At present, we’re asserting new Ray capabilities on Amazon SageMaker HyperPod that combine Ray with the HyperPod purpose-built infrastructure for basis mannequin coaching and serving.
Ray is an open-source framework that information scientists use to scale distributed Python workloads throughout clusters of GPUs, from distributed coaching with Ray Prepare to mannequin serving with Ray Serve. On Kubernetes, Ray clusters are managed by KubeRay, an open-source operator that handles cluster lifecycle by means of customized sources (RayCluster, RayJob, RayService). SageMaker HyperPod offers purpose-built infrastructure for large-scale machine studying (ML) on Amazon Elastic Kubernetes Service (Amazon EKS) with built-in node well being monitoring and automated restoration. Till now, working Ray on Kubernetes required information scientists to write down YAML manifests, handle Docker picture rebuilds for each dependency change, arrange kubectl port-forward to entry their Ray Dashboard, and configure Prometheus and Grafana manually for observability.
With this launch, information scientists can create Ray clusters, open the Ray Dashboard and Amazon Managed Grafana observability dashboards, join a JupyterLab or Code Editor workspace to their cluster, submit distributed jobs, and configure hung job detection, all from SageMaker Studio. On the software stage, Ray coaching jobs acquire automated fault tolerance by means of HyperPod node well being monitoring and restoration, plus tiered checkpointing for quicker resume by means of HyperPod distributed tiered storage. SageMaker JumpStart integration masses mannequin weights instantly into Ray Serve endpoints, with KV cache offloading to tiered storage for serving long-context requests. These capabilities work with open-source KubeRay and normal Ray APIs, so present scripts and workflows run with out modification.
On this submit, we stroll by means of tips on how to get began with every of those capabilities.
Conditions
To observe together with this submit, you want an Amazon SageMaker HyperPod cluster with Amazon EKS orchestration and the next parts put in in your cluster:
- SageMaker Areas EKS add-on, which allows JupyterLab and Code Editor workspaces in Amazon SageMaker HyperPod that connect to Ray clusters for interactive improvement.
- HyperPod Observability EKS add-on, which collects metrics from Ray workloads and provisions Grafana dashboards in Amazon Managed Grafana.
- KubeRay operator, which manages RayCluster, RayJob, and RayService as native Kubernetes sources.
- HyperPod Ray Endpoint Operator (Helm chart), which generates authenticated public endpoints for dashboard entry and distant job submission.
You additionally want a SageMaker Studio area, which offers the console interface for creating Ray clusters, viewing workloads, opening dashboards, and managing HyperPod Areas. For full setup directions, see the Ray on HyperPod getting began information.
Objective-built information scientist expertise
SageMaker Studio now offers an entire Ray improvement setting. Information scientists can create, handle, and monitor Ray clusters instantly from the console with out writing Kubernetes manifests or working kubectl instructions. You may try the total expertise on this interactive demo.
To get began, navigate to SageMaker Studio and select HyperPod. Choose your HyperPod cluster, then go to the Duties tab. From the duty sort listing, select RayCluster. The console shows your Ray clusters with their standing, occasion varieties, and out there actions. To create a brand new cluster, select Create Ray Cluster.
The creation kind prompts you for a cluster identify, head and employee occasion varieties, employee depend, and container picture. By default, clusters use the SageMaker Distribution picture, which comes with Ray pre-installed and is managed by AWS with common vulnerability patching and software program upgrades. You may also specify a customized container picture in case your workload requires further dependencies.
For purchasers preferring kubectl or want superior customization, an inline YAML editor in Studio exposes the total Kubernetes manifest. The KubeRay operator additionally integrates with HyperPod job governance, so directors can set compute quotas and scheduling priorities for Ray workloads alongside different coaching jobs.
Throughout creation, allow distant endpoints with the intention to entry the Ray Dashboard, submit jobs, and retrieve logs from wherever with web entry, securely and with out native kubectl port-forwarding. To be taught extra, see Ray Dashboard entry.
After the create cluster kind is submitted, return to the duties tab and it’ll present the cluster like within the screenshot under. From the Actions menu, you’ll be able to open the Ray Dashboard, view Grafana metrics, edit the cluster configuration, or delete the cluster.
After the cluster reaches Operating standing, select Open Ray Dashboard from the Actions menu. The system generates a short-lived, IAM-authenticated URL scoped to the cluster creator. The Ray Dashboard opens in a brand new tab, exhibiting cluster well being, working jobs, and node standing.
Distant job submission
For manufacturing workloads, you’ll be able to submit jobs remotely to Ray clusters from Studio, your laptop computer, or steady integration and steady supply (CI/CD) pipelines utilizing the toolkit-for-ray-on-sagemaker-ai Python bundle. The bundle handles endpoint decision and EKS API credential era by means of IAM authentication, so you employ normal Ray job submission APIs with a SageMaker-aware handle resolver:
Interactive improvement with SageMaker Areas
Information scientists can connect a Ray cluster to a HyperPod JupyterLab or Code Editor house from Studio. The house joins the cluster as a zero-compute employee node, giving the pocket book full native Ray driver entry. You choose your cluster from an inventory throughout house creation and start working instantly.
After making a HyperPod JupyterLab or Code Editor house, a Ray cluster integration choice seems within the configuration panel.
Choose the Ray cluster you need your workspace to hook up with.
Attaching a Ray cluster restarts the house. As soon as it restarts, open JupyterLab or Code Editor in your browser and name ray.init(handle="auto"). You’re related to your Ray cluster and might run distributed workloads as in the event you had been on the pinnacle node. Ray’s runtime_env parameter helps you to inject Python dependencies at runtime with out rebuilding container photos, and you may scale employees up or down with out recreating the cluster. For instance, an information scientist coaching a mannequin can begin prototyping in a pocket book with a single employee, then scale to 4 GPU employees by altering one line in ScalingConfig. The coaching runs distributed throughout the hooked up Ray cluster whereas the pocket book stays interactive for monitoring progress, adjusting hyperparameters, or inspecting intermediate outcomes. The whole workflow stays throughout the pocket book.
To be taught extra, see IDEs and Notebooks with Ray.
Out-of-the-box observability
Establishing Ray observability on Kubernetes with Amazon Managed Service for Prometheus and Amazon Managed Grafana is a multi-step course of that includes putting in Helm charts, creating PodMonitors and ServiceMonitors, configuring IAM roles for SigV4 signing, and manually importing dashboard JSON information.
The HyperPod Observability EKS add-on now handles all of this. It routinely discovers Ray head and employee pods, scrapes their metrics endpoints, and provisions 4 pre-built Grafana dashboards in Amazon Managed Grafana: Ray Core, Ray Information, Ray Prepare, and Ray Serve. You don’t want to create PodMonitors, configure scrape targets, or import dashboard JSON information.
All 4 dashboards are organized below a Ray folder in Amazon Managed Grafana and assist filtering metrics by particular Ray cluster. The Open Grafana motion from the cluster listing (proven within the first screenshot) takes you on to the metrics in your particular cluster in a single click on. The dashboards seem alongside present HyperPod infrastructure dashboards (GPU, EFA, job governance), so operations groups see Ray workload metrics and cluster well being in a single place.
Resilient coaching
SageMaker HyperPod offers three layers of resilience for Ray coaching workloads: automated node restoration when {hardware} fails, hung job detection when coaching stalls, and tiered checkpointing for quick restoration after both occasion.
Automated node restoration
SageMaker HyperPod repeatedly displays node well being and might routinely change defective nodes so your Ray coaching jobs can run for so long as wanted with out guide restarts. When a node is changed, Ray reschedules employee pods onto the brand new wholesome node. In case your coaching code saves checkpoints periodically and contains logic to renew from the most recent checkpoint, the job picks up the place it left off. This requires no modifications to your present Ray coaching code. You solely must configure a ample variety of retries in your RayJob’s FailureConfig in order that jobs get well routinely reasonably than failing completely on the primary interruption. To be taught extra, see Automated node restoration with Ray.
Hung job detection
Distributed coaching jobs can hold with out producing errors. A single pod fails (due to a community partition, storage mount concern, or {hardware} fault), and each different pod blocks on the subsequent collective operation, ready indefinitely. GPUs keep allotted with reminiscence loaded however produce no helpful compute. As a result of there isn’t a error message or crash, information scientists typically uncover the issue hours later after they test job progress manually. At scale, a number of hours of undetected hold time throughout dozens of GPUs represents important wasted value.
SageMaker HyperPod EKS now features a per-node Job Monitoring Agent that detects these situations for Ray Prepare workloads routinely, with no modifications to your code. The agent displays a number of node-level and job-level alerts to find out when a coaching job has stalled, and notifies the consumer by means of the cluster’s Amazon CloudWatch log group and the Ray Prepare Grafana dashboard provisioned by the HyperPod Observability add-on.
For customized detection guidelines, information scientists can use the toolkit-for-ray-on-sagemaker-ai library to outline log patterns and timeout thresholds. When the configured motion is cancel, HyperPod terminates the hung employee course of and Ray Prepare’s built-in FailureConfig restarts employees from the final checkpoint. To be taught extra, see HyperPod Hung Job Detection on Ray.
Tiered checkpointing
Quick checkpoint restoration is vital for each node alternative and hung job restart situations. The amzn-sagemaker-checkpointing library integrates with HyperPod managed tiered checkpointing, which writes checkpoints to native disk and asynchronously uploads them to Amazon Easy Storage Service (Amazon S3). When a job restarts, the library checks HyperPod Tiered Storage first. If the checkpoint remains to be out there there, restoration is quicker than restoring from Amazon S3. For big fashions, this could scale back restoration time in comparison with restoring instantly from Amazon S3. To be taught extra, see HyperPod Tiered Storage on Ray.
Accelerated inference
Ray Serve is Ray’s framework for deploying ML fashions as scalable, production-ready endpoints. It helps multi-model composition, autoscaling, and works with serving engines like vLLM. SageMaker HyperPod helps Ray Serve workloads on EKS, so you’ll be able to deploy and scale inference endpoints in your HyperPod cluster.
SageMaker JumpStart offers a catalog of pre-trained fashions. With this launch, the toolkit-for-ray-on-sagemaker-ai library now features a JumpStart mannequin loader that downloads mannequin weights from the JumpStart catalog and deploys them instantly into Ray Serve on HyperPod with out guide weight obtain, mannequin configuration, or container setup.
One widespread problem with massive language mannequin (LLM) serving is that inference latency grows with context size. Every new token requires recalculating consideration over all earlier tokens, making lengthy paperwork and multi-turn conversations sluggish and costly. SageMaker HyperPod addresses this with Managed Tiered KV Cache. The tiered cache shops consideration key-value vectors in CPU reminiscence on every node (L1) and on HyperPod Tiered Storage for cross-instance sharing (L2). Your Ray Serve deployments can reap the benefits of the Tiered Storage KV caching capabilities in SageMaker HyperPod with minimal code modifications, decreasing time-to-first-token for multi-turn conversations and long-document workloads. To be taught extra, see Accelerated Ray Inference on SageMaker HyperPod.
Clear up
To liberate compute capability in your cluster, delete any Ray clusters you created throughout this walkthrough. From the SageMaker Studio Duties tab, choose your Ray cluster and select Delete from the Actions menu. In the event you created a HyperPod cluster, you’ll be able to delete the cluster from the Amazon SageMaker AI console to cease incurring fees for the underlying compute cases. In the event you put in any prerequisite add-ons for testing, uninstall them from EKS console to liberate compute capability.
Conclusion
On this submit, we walked by means of how Amazon SageMaker HyperPod now offers an entire Ray expertise on EKS, from cluster creation and interactive notebooks to resilient coaching and accelerated inference. All of this works with open-source KubeRay and normal Ray APIs, so present scripts run with out modification. We’d wish to thank Dhawal Parkar, Pradeep Cruz, Mark Vinciguerra, and Giuseppe Angelo Porcelli for his or her contributions to this submit.
This integration is accessible as we speak in all AWS Areas the place SageMaker HyperPod EKS is supported. To get began, see the Amazon SageMaker HyperPod documentation and the Ray on HyperPod getting began information. You may also discover the total workflow on this interactive demo. In the event you’d like to debate how Ray on HyperPod can assist your workloads, contact an AWS consultant.
In regards to the authors
