Add RayTrain DLC (Ray distributed training — EKS + EC2)#6373
Open
Eren-Jeager123 wants to merge 7 commits into
Open
Add RayTrain DLC (Ray distributed training — EKS + EC2)#6373Eren-Jeager123 wants to merge 7 commits into
Eren-Jeager123 wants to merge 7 commits into
Conversation
New multi-node GPU DLC bundling Ray Train/Tune/Data on the PyTorch training stack (EFA/NCCL/GDRCopy/OpenMPI), built on the base CUDA DLC (cu130, py3.13). - docker/ray-train/Dockerfile.cuda: builder-base/gdrcopy/oss -> runtime-base -> eks (KubeRay head/worker) + ec2 (manual) stages - docker/ray-train/2.56/cuda: pyproject.toml + uv.lock (ray[train,tune,data] 2.56.0, torch 2.11.0+cu130); training-scoped, no ray[serve] - configs: framework=ray, container_type=train-ml, platform=eks -> tags ray:train-ml-cuda and ray:train-ml-eks-cuda (no release-logic changes) - workflows: ray-train pipeline + PR + autorelease (sanity/security/telemetry) - basic sanity test (imports, versions, CUDA-torch build, serve-absent) EFA/multi-node/KubeRay-on-EKS functional tests deferred (see test/ray-train/README.md).
- Move pyproject.toml/uv.lock to docker/ray-train/ (flat, latest-only — mirrors the Ray Serve DLC; no version subfolder) - Bump torch 2.11.0 -> 2.13.0 (+cu130), torchvision 0.26.0 -> 0.28.0 - Drop torchaudio (no cu130 wheel past 2.11.0; not needed for training) - Update Dockerfile COPY path + TORCH_VERSION default, configs torch_version
…test Build fix (CI failure on PR): - runtime-base: install cuda-nvcc + cuda-cudart-devel (base runtime image has no CUDA compiler; nccl-tests build needs nvcc). Add global CUDA_VERSION ARG. - register ray-train.autorelease-gpu.yml cron in release-schedule.yml (moved to a free slot 00 23 * * 2,4 to avoid colliding with vllm-omni) Sanity test reuse (per repo convention — sanity scripts are shared by category, not per-framework): - add 'ray' to training_cluster_only in test_sanity_training.py so RayTrain runs the shared training contract (env/PATH/EFA/NCCL/CUDA/cuDNN/SSH/venv/OSS) - gate OpenMPI double-wrap + entrypoint.sh checks to pt_tf_only (RayTrain uses EFA's bundled OpenMPI and a passive/KubeRay entrypoint) - add ray-gated TestRayTrain class (ray version, extras import, serve absent) - delete standalone test_sanity_ray_train.py; wire the shared script in the workflow
…d one Revert changes to test_sanity_training.py (leave it exactly as on main) and give RayTrain its own test_sanity_ray_train.py, matching repo convention — the Ray Serve image likewise keeps its own suite rather than reusing another framework's sanity file, and RayTrain diverges from the PT/TF contract (EFA-bundled OpenMPI, no from-source double-wrap; passive/KubeRay entrypoint, no entrypoint.sh). The dedicated script covers the shared training-cluster contract RayTrain honors (env, PATH, EFA/NCCL, CUDA, SSH, venv, OSS, nccl-tests binary) plus Ray specifics (ray version, ray[train,tune,data] imports, CUDA torch build, serve extra absent). Workflow invokes it on framework==ray && job_type==training.
- pipeline: add run-efa-test input + efa-test job (uses _reusable.efa-tests.yml, 2x EFA GPU instances, all_reduce_perf across nodes). Add ci-config job to read platform; gate EFA to the ec2 config only (interconnect is identical across eks/ec2 — run the expensive 2x p4d test once). release-gate now needs efa-test. - pr-gpu: add test/efa/** path + efa-test-change filter; run EFA on build or EFA test changes (matches pytorch convention) - autorelease-gpu: run-efa-test: true so releases are gated on multi-node EFA
Every Dockerfile ARG must come from the config; the Dockerfile default is only a fallback for bare 'docker build'. Add the missing ones to both configs: - builder_base_image, base_image (were relying on Dockerfile defaults) - python_short_version (venv site-packages path; was Dockerfile-only) Drop redundant build.framework_version — the build-image action already injects FRAMEWORK_VERSION from metadata.framework_version (kept in metadata as required).
…tion DO NOT MERGE THIS COMMIT — revert with 'git revert <sha>' before merging the PR. Lets a PR run publish the RayTrain image to the GAMMA private ECR (account 028651357192, repo 'ray', tags train-ml-cuda / train-ml-eks-cuda) so it can be pulled onto a real EKS / HyperPod-EKS cluster and verified before production. Changes (all reverted by reverting this single commit): - ec2/eks configs: environment production->gamma, public_registry true->false - pipeline release-gate: drop the main-branch / autorelease-only / no-PR guards (safe — gamma is non-prod, private registry only) Verified gamma infra exists: 'ray' repo present in 028651357192 (already holds serve-ml-* tags from Ray Serve gamma releases).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds a new RayTrain DLC — a multi-node, GPU, distributed-training container bundling Ray Train/Tune/Data on the PyTorch training stack (EFA / NCCL / GDRCopy / OpenMPI), for HyperPod-EKS / plain EKS (KubeRay) and EC2.
Motivation: the official training-oriented Ray image (
rayproject/ray-ml) was deprecated at Ray 2.31.0, so customers hand-stitch NCCL + Ray + PyTorch + Lightning + HF today (see the HyperPod EKS Ray Train guide). This DLC replaces that with a validated, CVE-patched image. Distinct from the existing Ray Serve DLC (docker/ray, inference, single-node): this is training-scoped and multi-node.Contents
Image —
docker/ray-train/Dockerfile.cudaFROMour own base CUDA DLC (base:{devel,runtime}-cu130-amzn2023), reusing the base's Python + uv + OSS-compliance + patched CUDA. Interconnect layer (EFA / NCCL / GDRCopy / OpenMPI / nccl-tests) reused from the PyTorch DLC install scripts.builder-base→builder-gdrcopy/builder-oss→runtime-base→eks(KubeRay head/worker; CRD injectsray start, no baked entrypoint) +ec2(passive/manual entrypoint).runtime-baseaddscuda-nvcc+cuda-cudart-devel(runtime base ships no compiler; nccl-tests needs nvcc) and bakes EFA/NCCL runtime env (FI_PROVIDER=efa,NCCL_SOCKET_IFNAME=eth0, etc.).Dependencies —
docker/ray-train/{pyproject.toml,uv.lock}(flat, latest-only, mirrors Ray Serve)ray[train,tune,data]==2.56.0,torch==2.13.0+cu130,torchvision==0.28.0,pytorch-lightning,transformers,datasets,accelerate, py3.13.ray[serve](that is the inference DLC). torchaudio omitted (no cu130 wheel past 2.11). peft/trl noted as a TODO.Configs —
.github/config/image/ray-train/{ec2,eks}-gpu.ymlframework: ray,container_type: train-ml,platform: eks(ec2 has no platform → default variant). Renders toray:train-ml-cuda(ec2) andray:train-ml-eks-cuda(eks) via the existingrayframework entry — zero release-logic changes.docker buildfallbacks only.Workflows —
ray-train.{pipeline,pr-gpu,autorelease-gpu}.yml+release-schedule.ymlentry.Tests
test/sanity/scripts/test_sanity_ray_train.py(own script, like the Ray Serve image keeps its own suite; RayTrain diverges from the PT/TF contract — EFA-bundled OpenMPI, passive/KubeRay entrypoint). Covers the shared training-cluster contract it honors (env, PATH, EFA/NCCL, CUDA, SSH, venv, OSS, nccl-tests binary) plus Ray specifics (version,ray[train,tune,data]imports, CUDA torch build, serve-extra absent).ray/ECR scan allowlist (framework: ray)._reusable.efa-tests.yml(2x EFA GPU instances,all_reduce_perfacross nodes; runs once on the ec2 config since the interconnect layer is identical across stages; gates the release).Verified
Both images build green in CI (~9.5 min each); sanity, security, and telemetry pass. Lock resolution, config → build-args, release-spec, and tag rendering all validated. All pre-commit hooks pass.
Not in this PR (deferred)
Single-node multi-GPU Ray Train convergence and KubeRay-on-EKS functional tests — tracked in
test/ray-train/README.md. (No EKS-based test exists in this repo yet; being scoped separately, likely via gamma release + manual HyperPod verification.)