Toolkit for Ray on Amazon SageMaker AI

Amazon SageMaker Ray library — utilities that make running Ray workloads on Amazon SageMaker HyperPod simpler and more reliable.

Overview

The toolkit-for-ray-on-sagemaker-ai library adds SageMaker HyperPod capabilities to your Ray workflows using standard Ray APIs:

• Authenticated job submission — submit jobs to secured HyperPod Ray clusters through Ray's native JobSubmissionClient and ray job CLI, with authentication handled transparently.
• Hung job detection — define log-based rules that detect when a Ray Train job stops making progress, and optionally cancel it automatically.
• JumpStart model loading — download Amazon SageMaker JumpStart model artifacts onto Ray worker nodes for Ray Serve LLM deployments.

Installation

pip install toolkit-for-ray-on-sagemaker-ai

Requires Python 3.11+ and Ray 2.0.0+ (ray[default]).

Job Submission

The library registers a sagemaker_ray:// address scheme for Ray's job submission APIs. When you use this scheme, the library acquires credentials and establishes an authenticated session for you — no custom client classes and no explicit imports are required.

Prerequisites

• A SageMaker HyperPod cluster (orchestrated by Amazon EKS) running a Ray cluster with the public Ray dashboard endpoint enabled.
• kubectl configured for the cluster: aws eks update-kubeconfig --name <cluster>. The library reads your kubeconfig (~/.kube/config or $KUBECONFIG).
• The AWS CLI installed and AWS credentials configured — the EKS kubeconfig uses aws eks get-token to obtain a token.
• Job submission is restricted to the cluster's owner.

Python SDK

from ray.job_submission import JobSubmissionClient

# Address form: sagemaker_ray://<cluster-name>/<namespace>
client = JobSubmissionClient("sagemaker_ray://my-cluster/team-a")

job_id = client.submit_job(entrypoint="python train.py")
print(client.get_job_status(job_id))
print(client.get_job_logs(job_id))

CLI

ray job submit --address "sagemaker_ray://my-cluster/team-a" -- python train.py
ray job status --address "sagemaker_ray://my-cluster/team-a" <job_id>
ray job logs   --address "sagemaker_ray://my-cluster/team-a" <job_id>

Address formats

• Cluster name: sagemaker_ray://<cluster-name>/<namespace>. The namespace defaults to default if omitted: sagemaker_ray://my-cluster.
• Dashboard URL: sagemaker_ray://<dashboard-host> when you already have the fully qualified dashboard hostname.

Hung Job Detection

SageMakerLogMonitoring watches your training job's stdout for expected log patterns. If an expected pattern stops appearing within a configured window (or an error pattern appears), the job is flagged as hung and can be cancelled automatically. A rule can also define a stop_pattern that deactivates it once the job reaches a known terminal state (e.g. training completes), so normal shutdown isn't mistaken for a hang.

from toolkit_for_ray_on_sagemaker_ai.log_monitoring import SageMakerLogMonitoring, LogMonitorConfig

def train_func():
    SageMakerLogMonitoring(config=LogMonitorConfig(
        enabled=True,
        rules=[
            {
                "name": "training_progress",
                "type": "log_pattern",
                "enabled": True,
                "log_pattern": r"(Epoch|Step|Iteration) \d+",
                "timeout_minutes": 10,
                "start_timeout_minutes": 30,
                "stop_pattern": "Training complete",
                "metric_evaluation_data_points": 3,
                "fault_on_match": False,
            },
            {
                "name": "oom_detection",
                "type": "log_pattern",
                "enabled": True,
                "log_pattern": "CUDA out of memory|OutOfMemoryError|OOM",
                "fault_on_match": True,
            },
        ],
        action="cancel",
    )).start()

    # ... your training loop ...

# To opt out of hang detection entirely:
SageMakerLogMonitoring(config=LogMonitorConfig(enabled=False)).start()

Rule fields

LogMonitorConfig

Hung job detection relies on the monitoring capability provided by SageMaker HyperPod. On clusters where it is not available, calls are a no-op and your training job runs normally.

JumpStart Model Loading

JumpStartModelLoaderCallback downloads SageMaker JumpStart model artifacts to each Ray worker node before the serving engine initializes, for use with Ray Serve LLM deployments.

from toolkit_for_ray_on_sagemaker_ai.jumpstart import JumpStartModelLoaderCallback

callback_config = CallbackConfig(
    callback_class=JumpStartModelLoaderCallback,
    callback_kwargs={
        "jumpstart_model_id": "meta-textgeneration-llama-3-1-8b-instruct",
        "region": "us-east-1",
        "accept_eula": True,
    },
)

Troubleshooting

• "Access denied … Only the cluster creator can submit jobs." — Job submission is restricted to the cluster owner. Use credentials for the identity that created the cluster.
• "Kubeconfig not found" / authentication failed — Run aws eks update-kubeconfig --name <cluster> and confirm your AWS credentials are valid.
• "The public endpoint may not be enabled for this cluster." — Ensure the cluster's public Ray dashboard endpoint is enabled.
• Resource not found — Verify the cluster name and namespace in the address.

Requirements

• Python >= 3.11
• Ray >= 2.0.0 (ray[default])
• AWS CLI installed and credentials configured

Dependencies

requests, pyyaml, ray[default], aiohttp, boto3

License

Apache License 2.0. See LICENSE.txt in the package for details.