Ray jobs on Amazon SageMaker HyperPod: scalable and resilient distributed AI

Basis mannequin (FM) coaching and inference has led to a major improve in computational wants throughout the business. These fashions require huge quantities of accelerated compute to coach and function successfully, pushing the boundaries of conventional computing infrastructure. They require environment friendly methods for distributing workloads throughout a number of GPU accelerated servers, and optimizing developer velocity in addition to efficiency.

Ray is an open supply framework that makes it easy to create, deploy, and optimize distributed Python jobs. At its core, Ray presents a unified programming mannequin that permits builders to seamlessly scale their functions from a single machine to a distributed cluster. It offers a set of high-level APIs for duties, actors, and knowledge that summary away the complexities of distributed computing, enabling builders to give attention to the core logic of their functions. Ray promotes the identical coding patterns for each a easy machine studying (ML) experiment and a scalable, resilient manufacturing utility. Ray’s key options embrace environment friendly activity scheduling, fault tolerance, and computerized useful resource administration, making it a robust instrument for constructing a variety of distributed functions, from ML fashions to real-time knowledge processing pipelines. With its rising ecosystem of libraries and instruments, Ray has grow to be a well-liked selection for organizations wanting to make use of the ability of distributed computing to deal with complicated and data-intensive issues.

Amazon SageMaker HyperPod is a purpose-built infrastructure to develop and deploy large-scale FMs. SageMaker HyperPod not solely offers the flexibleness to create and use your individual software program stack, but in addition offers optimum efficiency via identical backbone placement of cases, in addition to built-in resiliency. Combining the resiliency of SageMaker HyperPod and the effectivity of Ray offers a robust framework to scale up your generative AI workloads.

On this put up, we reveal the steps concerned in operating Ray jobs on SageMaker HyperPod.

Overview of Ray

This part offers a high-level overview of the Ray instruments and frameworks for AI/ML workloads. We primarily give attention to ML coaching use instances.

Ray is an open-source distributed computing framework designed to run extremely scalable and parallel Python functions. Ray manages, executes, and optimizes compute wants throughout AI workloads. It unifies infrastructure via a single, versatile framework—enabling AI workloads from knowledge processing, to mannequin coaching, to mannequin serving and past.

For distributed jobs, Ray offers intuitive instruments for parallelizing and scaling ML workflows. It permits builders to give attention to their coaching logic with out the complexities of useful resource allocation, activity scheduling, and inter-node communication.

At a excessive degree, Ray is made up of three layers:

• Ray Core: The inspiration of Ray, offering primitives for parallel and distributed computing
• Ray AI libraries:
  • Ray Practice – A library that simplifies distributed coaching by providing built-in assist for standard ML frameworks like PyTorch, TensorFlow, and Hugging Face
  • Ray Tune – A library for scalable hyperparameter tuning
  • Ray Serve – A library for distributed mannequin deployment and serving
• Ray clusters: A distributed computing platform the place employee nodes run person code as Ray duties and actors, usually within the cloud

On this put up, we dive deep into operating Ray clusters on SageMaker HyperPod. A Ray cluster consists of a single head node and numerous linked worker nodes. The pinnacle node orchestrates activity scheduling, useful resource allocation, and communication between nodes. The ray employee nodes execute the distributed workloads utilizing Ray duties and actors, comparable to mannequin coaching or knowledge preprocessing.

Ray clusters and Kubernetes clusters pair properly collectively. By operating a Ray cluster on Kubernetes utilizing the KubeRay operator, each Ray customers and Kubernetes directors profit from the graceful path from growth to manufacturing. For this use case, we use a SageMaker HyperPod cluster orchestrated via Amazon Elastic Kubernetes Service (Amazon EKS).

The KubeRay operator lets you run a Ray cluster on a Kubernetes cluster. KubeRay creates the next {custom} useful resource definitions (CRDs):

• RayCluster – The first useful resource for managing Ray cases on Kubernetes. The nodes in a Ray cluster manifest as pods within the Kubernetes cluster.
• RayJob – A single executable job designed to run on an ephemeral Ray cluster. It serves as a higher-level abstraction for submitting duties or batches of duties to be executed by the Ray cluster. A RayJob additionally manages the lifecycle of the Ray cluster, making it ephemeral by robotically spinning up the cluster when the job is submitted and shutting it down when the job is full.
• RayService – A Ray cluster and a Serve utility that runs on high of it right into a single Kubernetes manifest. It permits for the deployment of Ray functions that have to be uncovered for exterior communication, sometimes via a service endpoint.

For the rest of this put up, we don’t give attention to RayJob or RayService; we give attention to making a persistent Ray cluster to run distributed ML coaching jobs.

When Ray clusters are paired with SageMaker HyperPod clusters, Ray clusters unlock enhanced resiliency and auto-resume capabilities, which we are going to dive deeper into later on this put up. This mixture offers an answer for dealing with dynamic workloads, sustaining excessive availability, and offering seamless restoration from node failures, which is essential for long-running jobs.

Overview of SageMaker HyperPod

On this part, we introduce SageMaker HyperPod and its built-in resiliency options to supply infrastructure stability.

Generative AI workloads comparable to coaching, inference, and fine-tuning contain constructing, sustaining, and optimizing massive clusters of hundreds of GPU accelerated cases. For distributed coaching, the purpose is to effectively parallelize workloads throughout these cases to be able to maximize cluster utilization and decrease time to coach. For giant-scale inference, it’s necessary to reduce latency, maximize throughput, and seamlessly scale throughout these cases for one of the best person expertise. SageMaker HyperPod is a purpose-built infrastructure to deal with these wants. It removes the undifferentiated heavy lifting concerned in constructing, sustaining, and optimizing a big GPU accelerated cluster. It additionally offers flexibility to completely customise your coaching or inference setting and compose your individual software program stack. You should utilize both Slurm or Amazon EKS for orchestration with SageMaker HyperPod.

As a result of their huge dimension and the necessity to practice on massive quantities of information, FMs are sometimes skilled and deployed on massive compute clusters composed of hundreds of AI accelerators comparable to GPUs and AWS Trainium. A single failure in one in all these thousand accelerators can interrupt all the coaching course of, requiring guide intervention to determine, isolate, debug, restore, and get better the defective node within the cluster. This workflow can take a number of hours for every failure and because the scale of the cluster grows, it’s widespread to see a failure each few days and even each few hours. SageMaker HyperPod offers resiliency towards infrastructure failures by making use of brokers that repeatedly run well being checks on cluster cases, repair the dangerous cases, reload the final legitimate checkpoint, and resume the coaching—with out person intervention. Consequently, you possibly can practice your fashions as much as 40% quicker. You may as well SSH into an occasion within the cluster for debugging and collect insights on hardware-level optimization throughout multi-node coaching. Orchestrators like Slurm or Amazon EKS facilitate environment friendly allocation and administration of assets, present optimum job scheduling, monitor useful resource utilization, and automate fault tolerance.

Resolution overview

This part offers an outline of find out how to run Ray jobs for multi-node distributed coaching on SageMaker HyperPod. We go over the structure and the method of making a SageMaker HyperPod cluster, putting in the KubeRay operator, and deploying a Ray coaching job.

Though this put up offers a step-by-step information to manually create the cluster, be at liberty to take a look at the aws-do-ray mission, which goals to simplify the deployment and scaling of distributed Python utility utilizing Ray on Amazon EKS or SageMaker HyperPod. It makes use of Docker to containerize the instruments essential to deploy and handle Ray clusters, jobs, and providers. Along with the aws-do-ray mission, we’d like to focus on the Amazon SageMaker Hyperpod EKS workshop, which presents an end-to-end expertise for operating numerous workloads on SageMaker Hyperpod clusters. There are a number of examples of coaching and inference workloads from the GitHub repository awsome-distributed-training.

As launched earlier on this put up, KubeRay simplifies the deployment and administration of Ray functions on Kubernetes. The next diagram illustrates the answer structure.

Create a SageMaker HyperPod cluster

Stipulations

Earlier than deploying Ray on SageMaker HyperPod, you want a HyperPod cluster:

In case you choose to deploy HyperPod on an present EKS cluster, please observe the directions right here which embrace:

• EKS cluster – You’ll be able to affiliate SageMaker HyperPod compute to an present EKS cluster that satisfies the set of conditions. Alternatively and really helpful, you possibly can deploy a ready-made EKS cluster with a single AWS CloudFormation template. Consult with the GitHub repo for directions on establishing an EKS cluster.
• Customized assets – Working multi-node distributed coaching requires numerous assets, comparable to machine plugins, Container Storage Interface (CSI) drivers, and coaching operators, to be pre-deployed on the EKS cluster. You additionally must deploy further assets for the well being monitoring agent and deep well being verify. HyperPodHelmCharts simplify the method utilizing Helm, one in all mostly used bundle mangers for Kubernetes. Consult with Set up packages on the Amazon EKS cluster utilizing Helm for set up directions.

The next present an instance workflow for making a HyperPod cluster on an present EKS Cluster after deploying conditions. That is for reference solely and never required for the short deploy possibility.

cat > cluster-config.json << EOL
{
    "ClusterName": "ml-cluster",
    "Orchestrator": {
        "Eks": {
            "ClusterArn": "${EKS_CLUSTER_ARN}"
        }
    },
    "InstanceGroups": [
        {
            "InstanceGroupName": "worker-group-1",
            "InstanceType": "ml.p5.48xlarge",
            "InstanceCount": 4,
            "LifeCycleConfig": {
                "SourceS3Uri": "s3://amzn-s3-demo-bucket",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "${EXECUTION_ROLE}",
            "ThreadsPerCore": 1,
            "OnStartDeepHealthChecks": [
                "InstanceStress",
                "InstanceConnectivity"
            ]
        },
        {
            "InstanceGroupName": "head-group",
            "InstanceType": "ml.m5.2xlarge",
            "InstanceCount": 1,
            "LifeCycleConfig": {
                "SourceS3Uri": "s3://amzn-s3-demo-bucket",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "${EXECUTION_ROLE}",
            "ThreadsPerCore": 1,
        }
    ],
    "VpcConfig": {
        "SecurityGroupIds": [
            "${SECURITY_GROUP_ID}"
        ],
        "Subnets": [
            "${SUBNET_ID}"
        ]
    },
    "NodeRecovery": "Automated"
}
EOL

The supplied configuration file accommodates two key highlights:

• “OnStartDeepHealthChecks”: [“InstanceStress”, “InstanceConnectivity”] – Instructs SageMaker HyperPod to conduct a deep well being verify every time new GPU or Trainium cases are added
• “NodeRecovery”: “Automated” – Permits SageMaker HyperPod automated node restoration

You’ll be able to create a SageMaker HyperPod compute with the next AWS Command Line Interface (AWS CLI) command (AWS CLI model 2.17.47 or newer is required):

aws sagemaker create-cluster 
    --cli-input-json file://cluster-config.json
{
"ClusterArn": "arn:aws:sagemaker:us-east-2:xxxxxxxxxx:cluster/wccy5z4n4m49"
}

To confirm the cluster standing, you should use the next command:

aws sagemaker list-clusters --output desk

This command shows the cluster particulars, together with the cluster identify, standing, and creation time:

------------------------------------------------------------------------------------------------------------------------------------------------------
|                                                                    ListClusters                                                                    |
+----------------------------------------------------------------------------------------------------------------------------------------------------+
||                                                                 ClusterSummaries                                                                 ||
|+----------------------------------------------------------------+---------------------------+----------------+------------------------------------+|
||                           ClusterArn                           |        ClusterName        | ClusterStatus  |           CreationTime             ||
|+----------------------------------------------------------------+---------------------------+----------------+------------------------------------+|
||  arn:aws:sagemaker:us-west-2:xxxxxxxxxxxx:cluster/zsmyi57puczf |         ml-cluster        |   InService     |  2025-03-03T16:45:05.320000+00:00  ||
|+----------------------------------------------------------------+---------------------------+----------------+------------------------------------+|

Alternatively, you possibly can confirm the cluster standing on the SageMaker console. After a short interval, you possibly can observe that the standing for the nodes transitions to Working.

Create an FSx for Lustre shared file system

For us to deploy the Ray cluster, we’d like the SageMaker HyperPod cluster to be up and operating, and moreover we’d like a shared storage quantity (for instance, an Amazon FSx for Lustre file system). It is a shared file system that the SageMaker HyperPod nodes can entry. This file system could be provisioned statically earlier than launching your SageMaker HyperPod cluster or dynamically afterwards.

Specifying a shared storage location (comparable to cloud storage or NFS) is non-compulsory for single-node clusters, however it’s required for multi-node clusters. Utilizing an area path will raise an error throughout checkpointing for multi-node clusters.

The Amazon FSx for Lustre CSI driver makes use of IAM roles for service accounts (IRSA) to authenticate AWS API calls. To make use of IRSA, an IAM OpenID Join (OIDC) supplier must be related to the OIDC issuer URL that comes provisioned your EKS cluster.

Create an IAM OIDC id supplier in your cluster with the next command:

eksctl utils associate-iam-oidc-provider --cluster $EKS_CLUSTER_NAME --approve

Deploy the FSx for Lustre CSI driver:

helm repo add aws-fsx-csi-driver https://kubernetes-sigs.github.io/aws-fsx-csi-driver
helm repo replace
helm improve --install aws-fsx-csi-driver aws-fsx-csi-driver/aws-fsx-csi-driver
  --namespace kube-system 

This Helm chart features a service account named fsx-csi-controller-sa that will get deployed within the kube-system namespace.

Use the eksctl CLI to create an AWS Identification and Entry Administration (IAM) position certain to the service account utilized by the driving force, attaching the AmazonFSxFullAccess AWS managed coverage:

eksctl create iamserviceaccount 
  --name fsx-csi-controller-sa 
  --override-existing-serviceaccounts 
  --namespace kube-system 
  --cluster $EKS_CLUSTER_NAME 
  --attach-policy-arn arn:aws:iam::aws:coverage/AmazonFSxFullAccess 
  --approve 
  --role-name AmazonEKSFSxLustreCSIDriverFullAccess 
  --region $AWS_REGION

The --override-existing-serviceaccounts flag lets eksctl know that the fsx-csi-controller-sa service account already exists on the EKS cluster, so it skips creating a brand new one and updates the metadata of the present service account as an alternative.

Annotate the driving force’s service account with the Amazon Useful resource Identify (ARN) of the AmazonEKSFSxLustreCSIDriverFullAccess IAM position that was created:

SA_ROLE_ARN=$(aws iam get-role --role-name AmazonEKSFSxLustreCSIDriverFullAccess --query 'Position.Arn' --output textual content)

kubectl annotate serviceaccount -n kube-system fsx-csi-controller-sa 
  eks.amazonaws.com/role-arn=${SA_ROLE_ARN} --overwrite=true

This annotation lets the driving force know what IAM position it ought to use to work together with the FSx for Lustre service in your behalf.

Confirm that the service account has been correctly annotated:

kubectl get serviceaccount -n kube-system fsx-csi-controller-sa -o yaml

Restart the fsx-csi-controller deployment for the adjustments to take impact:

kubectl rollout restart deployment fsx-csi-controller -n kube-system

The FSx for Lustre CSI driver presents you with two choices for provisioning a file system:

• Dynamic provisioning – This selection makes use of Persistent Quantity Claims (PVCs) in Kubernetes. You outline a PVC with desired storage specs. The CSI driver robotically provisions the FSx for Lustre file system for you primarily based on the PVC request. This enables for easy scaling and eliminates the necessity to manually create file methods.
• Static provisioning – On this technique, you manually create the FSx for Lustre file system earlier than utilizing the CSI driver. You will have to configure particulars like subnet ID and safety teams for the file system. Then, you should use the driving force to mount this pre-created file system inside your container as a quantity.

For this instance, we use dynamic provisioning. Begin by making a storage class that makes use of the fsx.csi.aws.com provisioner:

cat <<EOF > storageclass.yaml
form: StorageClass
apiVersion: storage.k8s.io/v1
metadata:
  identify: fsx-sc
provisioner: fsx.csi.aws.com
parameters:
  subnetId: ${SUBNET_ID}
  securityGroupIds: ${SECURITYGROUP_ID}
  deploymentType: PERSISTENT_2
  automaticBackupRetentionDays: "0"
  copyTagsToBackups: "true"
  perUnitStorageThroughput: "250"
  dataCompressionType: "LZ4"
  fileSystemTypeVersion: "2.12"
mountOptions:
  - flock
EOF

kubectl apply -f storageclass.yaml

• SUBNET_ID: The subnet ID that the FSx for Lustre filesystem. Ought to be the identical non-public subnet that was used for HyperPod creation.
• SECURITYGROUP_ID: The safety group IDs that will probably be hooked up to the file system. Ought to be the identical Safety Group ID that’s utilized in HyperPod and EKS.

Subsequent, create a PVC that makes use of the fsx-claim storage declare:

cat <<EOF > pvc.yaml
apiVersion: v1
form: PersistentVolumeClaim
metadata:
  identify: fsx-claim
spec:
  accessModes:
  - ReadWriteMany
  storageClassName: fsx-sc
  assets:
    requests:
      storage: 1200Gi
EOF

kubectl apply -f pvc.yaml

This PVC will begin the dynamic provisioning of an FSx for Lustre file system primarily based on the specs supplied within the storage class.

Create the Ray cluster

Now that we’ve got each the SageMaker HyperPod cluster and the FSx for Lustre file system created, we will arrange the Ray cluster:

1. Arrange dependencies. We’ll create a brand new namespace in our Kubernetes cluster and set up the KubeRay operator utilizing a Helm chart.

We advocate utilizing KubeRay operator model 1.2.0 or greater, which helps computerized Ray Pod eviction and alternative in case of failures (for instance, {hardware} points on EKS or SageMaker HyperPod nodes).

# Create KubeRay namespace
kubectl create namespace kuberay
# Deploy the KubeRay operator with the Helm chart repository
helm repo add kuberay https://ray-project.github.io/kuberay-helm/
helm repo replace
#Set up each CRDs and Kuberay operator v1.2.0
helm set up kuberay-operator kuberay/kuberay-operator --version 1.2.0 --namespace kuberay
# Kuberay operator pod will probably be deployed onto head pod
kubectl get pods --namespace kuberay

2. Create a Ray Container Picture for the Ray Cluster manifest. With the current deprecation of the `rayproject/ray-ml` photos ranging from Ray model 2.31.0, it’s essential to create a {custom} container picture for our Ray cluster. Due to this fact, we are going to construct on high of the `rayproject/ray:2.42.1-py310-gpu` picture, which has all vital Ray dependencies, and embrace our coaching dependencies to construct our personal {custom} picture. Please be at liberty to change this Dockerfile as you want.

First, create a Dockerfile that builds upon the bottom Ray GPU picture and contains solely the mandatory dependencies:

cat <<EOF > Dockerfile
 
FROM rayproject/ray:2.42.1-py310-gpu
# Set up Python dependencies for PyTorch, Ray, Hugging Face, and extra
RUN pip set up --no-cache-dir 
    torch torchvision torchaudio 
    numpy 
    pytorch-lightning 
    transformers datasets consider tqdm click on 
    ray[train] ray[air] 
    ray[train-torch] ray[train-lightning] 
    torchdata 
    torchmetrics 
    torch_optimizer 
    speed up 
    scikit-learn 
    Pillow==9.5.0 
    protobuf==3.20.3
 
RUN pip set up --upgrade datasets transformers
 
# Set the person
USER ray
WORKDIR /house/ray
 
# Confirm ray set up
RUN which ray && 
    ray –-version
  
# Default command
CMD [ "/bin/bash" ]
 
EOF

Then, construct and push the picture to your container registry (Amazon ECR) utilizing the supplied script:

export AWS_REGION=$(aws configure get area)
export ACCOUNT=$(aws sts get-caller-identity --query Account --output textual content)
export REGISTRY=${ACCOUNT}.dkr.ecr.${AWS_REGION}.amazonaws.com/
 
echo "This course of could take 10-Quarter-hour to finish..."
 
echo "Constructing picture..."
 
docker construct --platform linux/amd64 -t ${REGISTRY}aws-ray-custom:newest .
 
# Create registry if wanted
REGISTRY_COUNT=$(aws ecr describe-repositories | grep "aws-ray-custom" | wc -l)
if [ "$REGISTRY_COUNT" == "0" ]; then
    aws ecr create-repository --repository-name aws-ray-custom
fi
 
# Login to registry
echo "Logging in to $REGISTRY ..."
aws ecr get-login-password --region $AWS_REGION| docker login --username AWS --password-stdin $REGISTRY
 
echo "Pushing picture to $REGISTRY ..."
 
# Push picture to registry
docker picture push ${REGISTRY}aws-ray-custom:newest 

Now, our Ray container picture is in Amazon ECR with all vital Ray dependencies, in addition to code library dependencies.

3. Create a Ray cluster manifest. We use a Ray cluster to host our coaching jobs. The Ray cluster is the first useful resource for managing Ray cases on Kubernetes. It represents a cluster of Ray nodes, together with a head node and a number of employee nodes. The Ray cluster CRD determines how the Ray nodes are arrange, how they impart, and the way assets are allotted amongst them. The nodes in a Ray cluster manifest as pods within the EKS or SageMaker HyperPod cluster.

Observe that there are two distinct sections within the cluster manifest. Whereas the `headGroupSpec` defines the top node of the Ray Cluster, the `workerGroupSpecs` outline the employee nodes of the Ray Cluster. Whereas a job might technically run on the Head node as properly, it’s common to separate the top node from the precise employee nodes the place jobs are executed. Due to this fact, the occasion for the top node can sometimes be a smaller occasion (i.e. we selected a m5.2xlarge). Because the head node additionally manages cluster-level metadata, it may be useful to have it run on a non-GPU node to reduce the chance of node failure (as GPU is usually a potential supply of node failure).

cat <<'EOF' > raycluster.yaml
apiVersion: ray.io/v1alpha1
form: RayCluster
metadata:
  identify: rayml
  labels:
    controller-tools.k8s.io: "1.0"
spec:
  # Ray head pod template
  headGroupSpec:
    # The `rayStartParams` are used to configure the `ray begin` command.
    # See https://github.com/ray-project/kuberay/blob/grasp/docs/steering/rayStartParams.md for the default settings of `rayStartParams` in KubeRay.
    # See https://docs.ray.io/en/newest/cluster/cli.html#ray-start for all obtainable choices in `rayStartParams`.
    rayStartParams:
      dashboard-host: '0.0.0.0'
    #pod template
    template:
      spec:
        #        nodeSelector:  
        #node.kubernetes.io/instance-type: "ml.m5.2xlarge"
        securityContext:
          runAsUser: 0
          runAsGroup: 0
          fsGroup: 0
        containers:
        - identify: ray-head
          picture: ${REGISTRY}aws-ray-custom:newest     ## IMAGE: Right here it's possible you'll select which picture your head pod will run
          env:                                ## ENV: Right here is the place you possibly can ship stuff to the top pod
            - identify: RAY_GRAFANA_IFRAME_HOST   ## PROMETHEUS AND GRAFANA
              worth: http://localhost:3000
            - identify: RAY_GRAFANA_HOST
              worth: http://prometheus-grafana.prometheus-system.svc:80
            - identify: RAY_PROMETHEUS_HOST
              worth: http://prometheus-kube-prometheus-prometheus.prometheus-system.svc:9090
          lifecycle:
            preStop:
              exec:
                command: ["/bin/sh","-c","ray stop"]
          assets:
            limits:                                    ## LIMITS: Set useful resource limits in your head pod
              cpu: 1
              reminiscence: 8Gi
            requests:                                    ## REQUESTS: Set useful resource requests in your head pod
              cpu: 1
              reminiscence: 8Gi
          ports:
          - containerPort: 6379
            identify: gcs-server
          - containerPort: 8265 # Ray dashboard
            identify: dashboard
          - containerPort: 10001
            identify: shopper
          - containerPort: 8000
            identify: serve
          volumeMounts:                                    ## VOLUMEMOUNTS
          - identify: fsx-storage
            mountPath: /fsx
          - identify: ray-logs
            mountPath: /tmp/ray
        volumes:
          - identify: ray-logs
            emptyDir: {}
          - identify: fsx-storage
            persistentVolumeClaim:
              claimName: fsx-claim
  workerGroupSpecs:
  # the pod replicas on this group typed employee
  - replicas: 4                                    ## REPLICAS: What number of employee pods you need 
    minReplicas: 1
    maxReplicas: 10
    # logical group identify, for this known as small-group, additionally could be purposeful
    groupName: gpu-group
    rayStartParams:
      num-gpus: "8"
    #pod template
    template:
      spec:
        #nodeSelector:
        # node.kubernetes.io/instance-type: "ml.p5.48xlarge"
        securityContext:
          runAsUser: 0
          runAsGroup: 0
          fsGroup: 0
        containers:
        - identify: ray-worker
          picture: ${REGISTRY}aws-ray-custom:newest             ## IMAGE: Right here it's possible you'll select which picture your head node will run
          env:
          lifecycle:
            preStop:
              exec:
                command: ["/bin/sh","-c","ray stop"]
          assets:
            limits:                                    ## LIMITS: Set useful resource limits in your employee pods
              nvidia.com/gpu: 8
              #vpc.amazonaws.com/efa: 32  
            requests:                                    ## REQUESTS: Set useful resource requests in your employee pods
              nvidia.com/gpu: 8
              #vpc.amazonaws.com/efa: 32
          volumeMounts:                                    ## VOLUMEMOUNTS
          - identify: ray-logs
            mountPath: /tmp/ray
          - identify: fsx-storage
            mountPath: /fsx
        volumes:
        - identify: fsx-storage
          persistentVolumeClaim:
            claimName: fsx-claim
        - identify: ray-logs
          emptyDir: {}
EOF

4. Deploy the Ray cluster:

envsubst < raycluster.yaml | kubectl apply -f -

5. Optionally, expose the Ray dashboard utilizing port forwarding:

# Will get identify of kubectl service that runs the top pod
export SERVICEHEAD=$(kubectl get service | grep head-svc | awk '{print $1}' | head -n 1)
# Port forwards the dashboard from the top pod service
kubectl port-forward --address 0.0.0.0 service/${SERVICEHEAD} 8265:8265 > /dev/null 2>&1 &

Now, you possibly can go to http://localhost:8265/ to go to the Ray Dashboard.

6. To launch a coaching job, there are just a few choices:
   1. Use the Ray jobs submission SDK, the place you possibly can submit jobs to the Ray cluster via the Ray dashboard port (8265 by default) the place Ray listens for job requests. To be taught extra, see Quickstart using the Ray Jobs CLI.
   2. Execute a Ray job within the head pod the place you exec instantly into the top pod after which submit your job. To be taught extra, see RayCluster Quickstart.

For this instance, we use the primary technique and submit the job via the SDK. Due to this fact, we merely run from an area setting the place the coaching code is offered in --working-dir. Relative to this path, we specify the principle coaching Python script situated at --train.py
Throughout the working-dir folder, we will additionally embrace further scripts we would must run the coaching.

The fsdp-ray.py instance is situated in aws-do-ray/Container-Root/ray/raycluster/jobs/fsdp-ray/fsdp-ray.py within the aws-do-ray GitHub repo.

# Inside jobs/ folder
ray job submit --address http://localhost:8265 --working-dir "fsdp-ray" -- python3 fsdp-ray.py

For our Python coaching script to run, we’d like to verify our coaching scripts are appropriately arrange to make use of Ray. This contains the next steps:

• Configure a mannequin to run distributed and on the proper CPU/GPU machine
• Configure a knowledge loader to shard knowledge throughout the workers and place knowledge on the proper CPU or GPU machine
• Configure a training function to report metrics and save checkpoints
• Configure scaling and CPU or GPU useful resource necessities for a coaching job
• Launch a distributed coaching job with a TorchTrainer class

For additional particulars on find out how to regulate your present coaching script to get probably the most out of Ray, seek advice from the Ray documentation.

The next diagram illustrates the entire structure you will have constructed after finishing these steps.

Implement coaching job resiliency with the job auto resume performance

Ray is designed with strong fault tolerance mechanisms to supply resilience in distributed methods the place failures are inevitable. These failures usually fall into two classes: application-level failures, which stem from bugs in person code or exterior system points, and system-level failures, attributable to node crashes, community disruptions, or inner bugs in Ray. To handle these challenges, Ray offers instruments and methods that allow functions to detect, get better, and adapt seamlessly, offering reliability and efficiency in distributed environments. On this part, we take a look at two of the most typical forms of failures, and find out how to implement fault tolerance in them that SageMaker HyperPod compliments: Ray Practice employee failures and Ray employee node failures.

• Ray Practice employee – It is a employee course of particularly used for coaching duties inside Ray Practice, Ray’s distributed coaching library. These staff deal with particular person duties or shards of a distributed coaching job. Every employee is chargeable for processing a portion of the information, coaching a subset of the mannequin, or performing computation throughout distributed coaching. They’re coordinated by the Ray Practice orchestration logic to collectively practice a mannequin.
• Ray employee node – On the Ray degree, this can be a Ray node in a Ray cluster. It’s a part of the Ray cluster infrastructure and is chargeable for operating duties, actors, and different processes as orchestrated by the Ray head node. Every employee node can host a number of Ray processes that execute duties or handle distributed objects. On the Kubernetes degree, a Ray employee node is a Kubernetes pod that’s managed by a KubeRay operator. For this put up, we will probably be speaking concerning the Ray employee nodes on the Kubernetes degree, so we are going to seek advice from them as pods.

On the time of writing, there are not any official updates relating to head pod fault tolerance and auto resume capabilities. Although head pod failures are uncommon, within the unlikely occasion of such a failure, you will have to manually restart your coaching job. Nevertheless, you possibly can nonetheless resume progress from the final saved checkpoint. To attenuate the chance of hardware-related head pod failures, it’s suggested to position the top pod on a devoted, CPU-only SageMaker HyperPod node, as a result of GPU failures are a typical coaching job failure level.

Ray Practice employee failures

Ray Practice is designed with fault tolerance to deal with employee failures, comparable to RayActorErrors. When a failure happens, the affected staff are stopped, and new ones are robotically began to take care of operations. Nevertheless, for coaching progress to proceed seamlessly after a failure, saving and loading checkpoints is crucial. With out correct checkpointing, the coaching script will restart, however all progress will probably be misplaced. Checkpointing is due to this fact a essential element of Ray Practice’s fault tolerance mechanism and must be carried out in your code.

Automated restoration

When a failure is detected, Ray shuts down failed staff and provisions new ones. Though this occurs, we will inform the coaching operate to all the time hold retrying till coaching can proceed. Every occasion of restoration from a employee failure is taken into account a retry. We will set the variety of retries via the max_failures attribute of the FailureConfig, which is about within the RunConfig handed to the Coach (for instance, TorchTrainer). See the next code:

from ray.practice import RunConfig, FailureConfig
# Tries to get better a run as much as this many occasions.
run_config = RunConfig(failure_config=FailureConfig(max_failures=2))
# No restrict on the variety of retries.
run_config = RunConfig(failure_config=FailureConfig(max_failures=-1))

For extra data, see Handling Failures and Node Preemption.

Checkpoints

A checkpoint in Ray Practice is a light-weight interface representing a listing saved both regionally or remotely. For instance, a cloud-based checkpoint would possibly level to s3://my-bucket/checkpoint-dir, and an area checkpoint would possibly level to /tmp/checkpoint-dir. To be taught extra, see Saving checkpoints during training.

To save lots of a checkpoint within the coaching loop, you first want to write down your checkpoint to an area listing, which could be short-term. When saving, you should use checkpoint utilities from different frameworks like torch.save, pl.Coach.save_checkpoint, accelerator.save_model, save_pretrained, tf.keras.Mannequin.save, and extra. You then create a checkpoint from the listing utilizing Checkpoint.from_directory. Lastly, report the checkpoint to Ray Practice utilizing ray.practice.report(metrics, checkpoint=...). The metrics reported alongside the checkpoint are used to maintain monitor of the best-performing checkpoints. Reporting will add the checkpoint to persistent storage.

In case you save checkpoints with ray.train.report(..., checkpoint=...) and run on a multi-node cluster, Ray Practice will elevate an error if NFS or cloud storage is just not arrange. It’s because Ray Practice expects all staff to have the ability to write the checkpoint to the identical persistent storage location.

Lastly, clear up the native short-term listing to liberate disk house (for instance, by exiting the tempfile.TemporaryDirectory context). We will save a checkpoint each epoch or each few iterations.

The next diagram illustrates this setup.

The next code is an instance of saving checkpoints utilizing native PyTorch:

import os
import tempfile

import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam

import ray.practice.torch
from ray import practice
from ray.practice import Checkpoint, ScalingConfig
from ray.practice.torch import TorchTrainer


def train_func(config):
    n = 100
    # create a toy dataset
    # knowledge   : X - dim = (n, 4)
    # goal : Y - dim = (n, 1)
    X = torch.Tensor(np.random.regular(0, 1, dimension=(n, 4)))
    Y = torch.Tensor(np.random.uniform(0, 1, dimension=(n, 1)))
    # toy neural community : 1-layer
    # Wrap the mannequin in DDP
    mannequin = ray.practice.torch.prepare_model(nn.Linear(4, 1))
    criterion = nn.MSELoss()

    optimizer = Adam(mannequin.parameters(), lr=3e-4)
    for epoch in vary(config["num_epochs"]):
        y = mannequin.ahead(X)
        loss = criterion(y, Y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        metrics = {"loss": loss.merchandise()}

        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            checkpoint = None

            should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
            # In normal DDP coaching, the place the mannequin is identical throughout all ranks,
            # solely the worldwide rank 0 employee wants to save lots of and report the checkpoint
            if practice.get_context().get_world_rank() == 0 and should_checkpoint:
                torch.save(
                    mannequin.module.state_dict(),  # NOTE: Unwrap the mannequin.
                    os.path.be a part of(temp_checkpoint_dir, "mannequin.pt"),
                )
                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)

            practice.report(metrics, checkpoint=checkpoint)


coach = TorchTrainer(
    train_func,
    train_loop_config={"num_epochs": 5},
    scaling_config=ScalingConfig(num_workers=2),
)
consequence = coach.match()

Ray Practice additionally comes with CheckpointConfig, a technique to configure checkpointing choices:

from ray.practice import RunConfig, CheckpointConfig
# Instance 1: Solely hold the two *most up-to-date* checkpoints and delete the others.
run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=2))
# Instance 2: Solely hold the two *finest* checkpoints and delete the others.
run_config = RunConfig(
    checkpoint_config=CheckpointConfig(
        num_to_keep=2,
        # *Greatest* checkpoints are decided by these params:
        checkpoint_score_attribute="mean_accuracy",
        checkpoint_score_order="max",
    ),
    # This can retailer checkpoints on S3.
    storage_path="s3://remote-bucket/location",
)

To restore training state from a checkpoint in case your coaching job have been to fail and retry, it is best to modify your coaching loop to auto resume after which restore a Ray Practice job. By pointing to the trail of your saved checkpoints, you possibly can restore your coach and proceed coaching. Right here’s a fast instance:

from ray.practice.torch import TorchTrainer

restored_trainer = TorchTrainer.restore(
    path="~/ray_results/dl_trainer_restore",  # Will also be a cloud storage path like S3
    datasets=get_datasets(),
)
consequence = restored_trainer.match()

To streamline restoration, you possibly can add auto resume logic to your script. This checks if a sound experiment listing exists and restores the coach if obtainable. If not, it begins a brand new experiment:

experiment_path = "~/ray_results/dl_restore_autoresume"
if TorchTrainer.can_restore(experiment_path):
    coach = TorchTrainer.restore(experiment_path, datasets=get_datasets())
else:
    coach = TorchTrainer(
        train_loop_per_worker=train_loop_per_worker,
        datasets=get_datasets(),
        scaling_config=practice.ScalingConfig(num_workers=2),
        run_config=practice.RunConfig(
            storage_path="~/ray_results",
            identify="dl_restore_autoresume",
        ),
    )
consequence = coach.match()

To summarize, to supply fault tolerance and auto resume when utilizing Ray Practice libraries, set your max_failures parameter within the FailureConfig (we advocate setting it to -1 to verify it can hold retrying till the SageMaker HyperPod node is rebooted or changed), and be sure to have enabled checkpointing in your code.

Ray employee pod failures

Along with the aforementioned mechanisms to get better from Ray Practice employee failures, Ray additionally offers fault tolerance on the employee pod degree. When a employee pod fails (this contains situations during which the raylet course of fails), the operating duties and actors on it can fail and the objects owned by employee processes of this pod will probably be misplaced. On this case, the tasks, actors, and objects fault tolerance mechanisms will begin and attempt to get better the failures utilizing different employee pods.

These mechanisms will probably be implicitly dealt with by the Ray Practice library. To be taught extra concerning the underlying fault tolerance on the duties, actors, objects (carried out on the Ray Core degree), see Fault Tolerance.

In apply, which means that in case of a employee pod failure, the next happens:

• If there’s a free employee pod within the Ray cluster, Ray will get better the failed employee pod by changing it with the free employee pod.
• If there isn’t any free employee pod, however within the underlying SageMaker HyperPod cluster there are free SageMaker HyperPod nodes, Ray will schedule a brand new employee pod onto one of many free SageMaker HyperPod nodes. This pod will be a part of the operating Ray cluster and the failure will probably be recovered utilizing this new employee pod.

Within the context of KubeRay, Ray employee nodes are represented by Kubernetes pods, and failures at this degree can embrace points comparable to pod eviction or preemption attributable to software-level elements.

Nevertheless, one other essential state of affairs to contemplate is {hardware} failures. If the underlying SageMaker HyperPod node turns into unavailable because of a {hardware} situation, comparable to a GPU error, it might inevitably trigger the Ray employee pod operating on that node to fail as properly. Now the fault tolerance and auto-healing mechanisms of your SageMaker HyperPod cluster begin and can reboot or change the defective node. After the brand new wholesome node is added into the SageMaker HyperPod cluster, Ray will schedule a brand new employee pod onto the SageMaker HyperPod node and get better the interrupted coaching. On this case, each the Ray fault tolerance mechanism and the SageMaker HyperPod resiliency options work collectively seamlessly and ensure that even in case of a {hardware} failure, your ML coaching workload can auto resume and choose up from the place it was interrupted.

As you will have seen, there are numerous built-in resiliency and fault-tolerance mechanisms that permit your Ray Practice workload on SageMaker HyperPod to get better and auto resume. As a result of these mechanisms will primarily get better by restarting the coaching job, it’s essential that checkpointing is carried out within the coaching script. Additionally it is usually suggested to save lots of the checkpoints on a shared and chronic path, comparable to an Amazon Easy Storage Service (Amazon S3) bucket or FSx for Lustre file system.

Clear up

To delete your SageMaker HyperPod cluster created on this put up, you possibly can both use the SageMaker AI console or use the next AWS CLI command:

aws sagemaker delete-cluster --cluster-name <cluster_name>

Cluster deletion will take a couple of minutes. You’ll be able to affirm profitable deletion after you see no clusters on the SageMaker AI console.

In case you used the CloudFormation stack to create assets, you possibly can delete it utilizing the next command:

aws cloudformation delete-stack --stack-name <stack_name>

Conclusion

This put up demonstrated find out how to arrange and deploy Ray clusters on SageMaker HyperPod, highlighting key issues comparable to storage configuration and fault tolerance and auto resume mechanisms.

Working Ray jobs on SageMaker HyperPod presents a robust resolution for distributed AI/ML workloads, combining the flexibleness of Ray with the strong infrastructure of SageMaker HyperPod. This integration offers enhanced resiliency and auto resume capabilities, that are essential for long-running and resource-intensive duties. Through the use of Ray’s distributed computing framework and the built-in options of SageMaker HyperPod, you possibly can effectively handle complicated ML workflows, particularly coaching workloads as lined on this put up. As AI/ML workloads proceed to develop in scale and complexity, the mix of Ray and SageMaker HyperPod presents a scalable, resilient, and environment friendly platform for tackling probably the most demanding computational challenges in machine studying.

To get began with SageMaker HyperPod, seek advice from the Amazon EKS Support in Amazon SageMaker HyperPod workshop and the Amazon SageMaker HyperPod Developer Information. To be taught extra concerning the aws-do-ray framework, seek advice from the GitHub repo.

In regards to the Authors

Mark Vinciguerra is an Affiliate Specialist Options Architect at Amazon Internet Companies (AWS) primarily based in New York. He focuses on the Automotive and Manufacturing sector, specializing in serving to organizations architect, optimize, and scale synthetic intelligence and machine studying options, with explicit experience in autonomous automobile applied sciences. Previous to AWS, he went to Boston College and graduated with a level in Laptop Engineering.

Florian Stahl is a Worldwide Specialist Options Architect at AWS, primarily based in Hamburg, Germany. He makes a speciality of Synthetic Intelligence, Machine Studying, and Generative AI options, serving to clients optimize and scale their AI/ML workloads on AWS. With a background as a Information Scientist, Florian focuses on working with clients within the Autonomous Automobile house, bringing deep technical experience to assist organizations design and implement refined machine studying options. He works carefully with clients worldwide to remodel their AI initiatives and maximize the worth of their machine studying investments on AWS.

Anoop Saha is a Sr GTM Specialist at Amazon Internet Companies (AWS) specializing in Gen AI mannequin coaching and inference. He’s partnering with high basis mannequin builders, strategic clients, and AWS service groups to allow distributed coaching and inference at scale on AWS and lead joint GTM motions. Earlier than AWS, Anoop has held a number of management roles at startups and enormous companies, primarily specializing in silicon and system structure of AI infrastructure.

Alex Iankoulski is a Principal Options Architect, ML/AI Frameworks, who focuses on serving to clients orchestrate their AI workloads utilizing containers and accelerated computing infrastructure on AWS. He’s additionally the writer of the open supply do framework and a Docker captain who loves making use of container applied sciences to speed up the tempo of innovation whereas fixing the world’s greatest challenges.