Streamline machine learning workflows using Skypilot on Amazon Sagemaker HyperPod

This post was co-authored with Zhanghao Wu, co-creator of Skypilot.

Rapid advances in generator AI and basic models (FMS) have significantly increased the computational resource requirements for machine learning (ML) workloads. Modern ML pipelines require efficient systems to distribute workloads across accelerated computational resources, ensuring developers are highly productive. Organizations need infrastructure solutions that are not only powerful, but also flexible, resilient, and easy to manage.

Skypilot is an open source framework that simplifies run ML workloads by providing an integrated layer of abstraction that helps ML engineers run workloads on a variety of computing resources without managing the complexity of the underlying infrastructure. It provides a simple, high-level interface for resource provisioning, job scheduling, and managing distributed training across multiple nodes.

Amazon Sagemaker HyperPod is a dedicated infrastructure for developing and deploying large-scale FMSs. Sagemaker HyperPod not only offers the flexibility to create and use your own software stack, but also provides optimal performance through the placement and built-in resilience of the same spine instance. The resilience of the Sagemaker HyperPod with the efficiency of Skypilot provides a powerful framework to expand generative AI workloads.

In this post, we share how Sagemaker HyperPod works with Skypilot to streamline AI development workflows. This integration makes the sophisticated GPU infrastructure more accessible to ML engineers, increasing productivity and resource utilization.

Challenges for tuning machine learning workloads

Kubernetes is popular in ML workloads due to its scalability and extensive open source tools. Organized into the Amazon Elastic Kubernetes Service (Amazon EKS), Sagemaker HyperPod combines the power of Kubernetes with the resilient environment of Sagemaker HyperPods designed for training large-scale models. Amazon EKS support for SageMaker HyperPod enhances resilience through deep health checks, automated node recovery and automated job display capabilities, providing uninterrupted training for large, long-term jobs.

ML engineers migrating from traditional VM or on-premises environments often face a steep learning curve. The complexity of Kubernetes and cluster management can slow development cycles and resource utilization, creating important challenges.

Additionally, the AI infrastructure team faced the challenge of balancing the desire to provide a user-friendly experience for ML engineers with the need for sophisticated management tools. They needed a solution that could provide both high levels of control and ease of use for daily operations.

Sagemaker HyperPod and Skypilot

To address these challenges, we partnered with Skypilot to showcase solutions that use the strengths of both platforms. Sagemaker HyperPod excels in managing underlying computational resources and instances, providing the robust infrastructure needed to request AI workloads. Skypilot complements this by providing an intuitive layer for job management, interactive development and team coordination.

Through this partnership, we can provide our customers with the best world of both. It combines the powerful and scalable infrastructure of SageMaker HyperPod with a user-friendly interface that significantly reduces the learning curve of ML engineers. For AI infrastructure teams, this integration provides sophisticated management capabilities, simplifies the experience of ML engineers and creates a favorable situation for all stakeholders.

Skypilot helps AI teams run their workloads on a variety of infrastructures with unified high-level interfaces and powerful management of resources and jobs. AI engineers can bring in AI frameworks and specify resource requirements for jobs. Skypilot intelligently schedules workloads with the best infrastructure. Find available GPUs, provision GPUs, run jobs, and manage the lifecycle.

Solution overview

Implementing this solution is easy whether you are using an existing Sagemaker HyperPod cluster or configuring a new deployment. For existing clusters, you can connect using AWS Command Line Interface (AWS CLI) commands to update Kubeconfig to verify setup. For new deployments, it guides you through the configuration of high-performance network options such as API server setup, cluster creation, and Elastic Fabric Adapter (EFA).

The following diagram illustrates the solution architecture.

The next section shows how to run a Skypilot job for multi-node distributed training on a Sagemaker HyperPod. Consider the process of creating a Sagemaker HyperPod cluster, installing Skypilot, creating a Skypilot cluster, and deploying a Skypilot training job.

Prerequisites

The following prerequisites are required:

• An existing Sagemaker HyperPod cluster with Amazon eks (see Deploying a HyperPod Cluster to create one). In the next section, you need to provision a single ml.p5.48xlarge instance in the code sample.
• Accessing the AWS CLI kubectl Command line tool.
• A Python environment for installing Skypilot.

Create a Sagemaker HyperPod cluster

You can create an EKS cluster with a single AWS CloudFormation stack according to instructions using CloudFormation, consisting of a Virtual Private Cloud (VPC) and storage resources.

You can use the AWS Management Console or the AWS CLI to create and manage a Sagemaker HyperPod cluster. If you are using the AWS CLI, specify the cluster configuration for the JSON file and select the EKS cluster created from the CloudFormation stack as an orchestrator for the Sagemaker HyperPod cluster. Next, create a cluster worker node NodeRecovery Set to Automatic Enable automatic node recovery OnStartDeepHealthChecksaddition InstanceStress and InstanceConnectivity To enable deep health checks. See the following code:

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

You can add InstanceStorageConfigs to provide and mount additional Amazon Elastic Block Store (Amazon EBS) volumes on your SageMaker HyperPod nodes.

To create a cluster using the Sagemaker HyperPod API, run the following AWS CLI command:

aws sagemaker create-cluster \ 
--cli-input-json file://cluster-config.json

You are now ready to set up Skypilot in your Sagemaker HyperPod cluster.

Connect to Sagemaker HyperPod EKS cluster

From the AWS CLI environment, run the AWS EKS Update-Kubeconfig command and find it in the local Kube configuration file ( ~/.kube/config) Use the required qualifications and configuration to connect to an eks cluster kubectl Command (provides a specific EKS cluster name):

aws eks update-kubeconfig --name $EKS_CLUSTER_NAME

You can verify that you are connected to your EKS cluster by running the following command:

kubectl config current-context

Install Skypilot with Kubernetes support

Use the following code to install Skypilot with Kubernetes support using PIP:

pip install skypilot[kubernetes]

This will install the latest build of Skypilot, which includes the required Kubernetes integrations.

Check Skypilot's connection to EKS cluster

Check if Skypilot can connect to your Kubernetes cluster.

sky check k8s

The output should look similar to the following code:

Checking credentials to enable clouds for SkyPilot.
Kubernetes: enabled [compute]

To enable a cloud, follow the hints above and rerun: sky check
If any problems remain, refer to detailed docs at: https://docs.skypilot.co/en/latest/getting-started/installation.html

🎉 Enabled clouds 🎉
Kubernetes [compute]
Active context: arn:aws:eks:us-east-2:XXXXXXXXXXXXX:cluster/sagemaker-hyperpod-eks-cluster

Using SkyPilot API server: http://127.0.0.1:46580

If you are new to using Skypilot on this Kubernetes cluster, you may be prompted to create a GPU label for the node. Run the following code and follow the instructions:

python -m sky.utils.kubernetes.gpu_labeler --context

This script helps Skypilot identify the available GPU resources for each node in the cluster. A GPU labeling job can take several minutes depending on the number of GPU resources in the cluster.

Discover the available GPUs in your cluster

To check the available GPU resources in your Sagemaker HyperPod cluster, use the following code:

sky show-gpus --cloud k8s

This lists the available GPU types and their counts. There are two p5.48xlarge instances, each equipped with eight nvidia H100 GPUs.

 Kubernetes GPUs
GPU REQUESTABLE_QTY_PER_NODE TOTAL_GPUS TOTAL_FREE_GPUS
H100 1, 2, 4, 8 16 16

Kubernetes per node accelerator availability
NODE_NAME GPU_NAME TOTAL_GPUS FREE_GPUS
hyperpod-i-00baa178bc31afde3 H100 8 8
hyperpod-i-038beefa954efab84 H100 8 8

Start an interactive development environment

Skypilot allows you to launch a Skypilot cluster for interactive development.

sky launch -c dev --gpus H100

This command creates an interactive development environment (IDE) with a single H100 GPU and synchronizes the local working directory to the cluster. Skypilot handles pod creation, resource allocation, and setup for your IDE.

Considered resources (1 node):
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
 CLOUD        INSTANCE            vCPUs   Mem(GB)   ACCELERATORS   REGION/ZONE                                                                 COST ($)   CHOSEN   
-------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Kubernetes   2CPU--8GB--H100:1   2       8         H100:1         arn:aws:eks:us-east-2:XXXXXXXXXX:cluster/sagemaker-hyperpod-eks-cluster   0.00          ✔     
------------------------------------------------------------------------------------------------------------------------------------------------------------------
Launching a new cluster 'dev'. Proceed? [Y/n]: Y
• Launching on Kubernetes.
Pod is up.
✔ Cluster launched: dev. View logs: sky api logs -1 sky-2025-05-05-15-28-47-523797/provision. log
• Syncing files.
Run commands not specified or empty.
Useful Commands
Cluster name: dey
To log into the head VM:   ssh dev
To submit a job:           sky exec dev yaml_file
To stop the cluster:       sky stop dev
To teardown the cluster:   sky down dev

Once booted you can connect to the IDE.

ssh dev

This provides an interactive shell to your IDE, allowing you to run code, install packages, and run ML experiments.

Run the training job

Skypilot allows you to run distributed training jobs on a Sagemaker HyperPod cluster. Below is an example of starting a distributed training job using a YAML configuration file:

First, create a file with the name train.yaml In the training job configuration:

resources:
    accelerators: H100

num_nodes: 1

setup: |
    git clone --depth 1 https://github.com/pytorch/examples || true
    cd examples
    git filter-branch --prune-empty --subdirectory-filter distributed/minGPT-ddp
    # SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
    uv venv --python 3.10
    source .venv/bin/activate
    uv pip install -r requirements.txt "numpy<2" "torch"

run: |
    cd examples
    source .venv/bin/activate
    cd mingpt
    export LOGLEVEL=INFO

    MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
    echo "Starting distributed training, head node: $MASTER_ADDR"

    torchrun \
    --nnodes=$SKYPILOT_NUM_NODES \
    --nproc_per_node=$SKYPILOT_NUM_GPUS_PER_NODE \
    --master_addr=$MASTER_ADDR \
    --master_port=8008 \
    --node_rank=${SKYPILOT_NODE_RANK} \
    main.py

After that, start the training job.

sky launch -c train train.yaml

This creates a training job on a single P5.48XLARGE node equipped with an 8 H100 NVIDIA GPU. You can monitor the output with the following command:

sky logs train

Run a multi-node training job with EFA

The Elastic Fabric Adapter (EFA) is the network interface for Amazon Elastic Compute Cloud (Amazon EC2) instances that can run applications requiring high-level internode communication on AWS via custom built operating system bypass hardware interfaces. This allows applications to communicate directly with network hardware while bypassing the operating system kernel, significantly reducing latency and CPU overhead. This direct hardware access is particularly beneficial for distributed ML workloads where internode communication can often become a bottleneck during gradient synchronization. By using EFA-enabled instances such as P5.48XLARGE and P6-B200.48XLARGE, data scientists scale training jobs across multiple nodes, maintaining low latency, high bandwidth communications essential for efficient distributed training, ultimately reducing training times and improving resources for large AI workloads.

The following code snippet shows how to incorporate this into a Skypilot job:

name: nccl-test-efa

resources:
  cloud: kubernetes
  accelerators: H100:8
  image_id: docker:public.ecr.aws/hpc-cloud/nccl-tests:latest

num_nodes: 2

envs:
  USE_EFA: "true"

run: |
  if [ "${SKYPILOT_NODE_RANK}" == "0" ]; then
    echo "Head node"

    # Total number of processes, NP should be the total number of GPUs in the cluster
    NP=$(($SKYPILOT_NUM_GPUS_PER_NODE * $SKYPILOT_NUM_NODES))

    # Append :${SKYPILOT_NUM_GPUS_PER_NODE} to each IP as slots
    nodes=""
    for ip in $SKYPILOT_NODE_IPS; do
      nodes="${nodes}${ip}:${SKYPILOT_NUM_GPUS_PER_NODE},"
    done
    nodes=${nodes::-1}
    echo "All nodes: ${nodes}"

    # Set environment variables
    export PATH=$PATH:/usr/local/cuda-12.2/bin:/opt/amazon/efa/bin:/usr/bin
    export LD_LIBRARY_PATH=/usr/local/cuda-12.2/lib64:/opt/amazon/openmpi/lib:/opt/nccl/build/lib:/opt/amazon/efa/lib:/opt/aws-ofi-nccl/install/lib:/usr/local/nvidia/lib:$LD_LIBRARY_PATH
    export NCCL_HOME=/opt/nccl
    export CUDA_HOME=/usr/local/cuda-12.2
    export NCCL_DEBUG=INFO
    export NCCL_BUFFSIZE=8388608
    export NCCL_P2P_NET_CHUNKSIZE=524288
    export NCCL_TUNER_PLUGIN=/opt/aws-ofi-nccl/install/lib/libnccl-ofi-tuner.so

    if [ "${USE_EFA}" == "true" ]; then
      export FI_PROVIDER="efa"
    else
      export FI_PROVIDER=""
    fi

    /opt/amazon/openmpi/bin/mpirun \
      --allow-run-as-root \
      --tag-output \
      -H $nodes \
      -np $NP \
      -N $SKYPILOT_NUM_GPUS_PER_NODE \
      --bind-to none \
      -x FI_PROVIDER \
      -x PATH \
      -x LD_LIBRARY_PATH \
      -x NCCL_DEBUG=INFO \
      -x NCCL_BUFFSIZE \
      -x NCCL_P2P_NET_CHUNKSIZE \
      -x NCCL_TUNER_PLUGIN \
      --mca pml ^cm,ucx \
      --mca btl tcp,self \
      --mca btl_tcp_if_exclude lo,docker0,veth_def_agent \
      /opt/nccl-tests/build/all_reduce_perf \
      -b 8 \
      -e 2G \
      -f 2 \
      -g 1 \
      -c 5 \
      -w 5 \
      -n 100
  else
    echo "Worker nodes"
  fi

config:
  kubernetes:
    pod_config:
      spec:
        containers:
        - resources:
            limits:
              
              vpc.amazonaws.com/efa: 32
            requests:
              
              vpc.amazonaws.com/efa: 32

cleaning

To delete a Skypilot cluster, run the following command:

sky down

To delete the Sagemaker HyperPod cluster created in this post, you can use the Sagemaker AI console or one of the following AWS CLI commands for users:

aws sagemaker delete-cluster --cluster-name

It takes a few minutes to delete a cluster. After the cluster does not appear in the Sagemaker AI console, you can check for successful deletions.

If you created a resource using the CloudFormation stack, you can delete it using the following command:

aws cloudformation delete-stack --stack-name

Conclusion

By combining the robust infrastructure capabilities of the Sagemaker HyperPod with Skypilot's user-friendly interface, we've introduced solutions that will help teams focus on innovation rather than infrastructure complexity. This approach not only simplifies operations, but also enhances productivity and resource utilization across organizations of all sizes. To get started, see Amazon EKS Support Skypilot in Amazon Sagemaker HyperPod Workshop.

About the author

Roy Arera I am AWS Senior AI/ML Specialist Solution Architect. He supports AWS customers (from small startups to large companies) and efficiently drives and deploys basic models on AWS. He is passionate about computational optimization problems and improving performance for AI workloads.

ZHANGHAO WU He is a co-creator of the Skypilot Open Source Project and holds a PhD in Computer Science from Berkeley, California. He is committed to improving the AI experience for Skypilot Core, client-server architecture, managed jobs, and a wide range of cloud infrastructures in general.

Anniversary I am a mobile market (GTM) specialist in AWS' senior foundation model. He partners with Top Generative AI Model Builders, Strategic Customers and AWS Services teams to enable the next generation of AI/ML workloads on AWS. Ankit's experience includes Amazon Alexa's product management expertise within the financial services industry in high frequency and low latency trading and business development.