feat(ray): staged Ray scheduler-executor support (RayJobClient + RayExecutor + get_executor branch) (#1435)

Signed-off-by: Nick Gupta <nicgupta@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: George Armstrong <georgea@nvidia.com>

Author: 3 people

cluster_configs/example-ray.yaml

@@ -0,0 +1,80 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Ray executor cluster config — initial Ray support for SFT (and eventually GRPO).
+#
+# Use this for either:
+#   (a) standalone Ray clusters (e.g., an existing self-managed Ray deployment), or
+#   (b) Ray-on-Slurm setups where Slurm allocates the nodes and Ray runs the
+#       job-submission layer inside the allocation.
+#
+# This is DISTINCT from the existing `with_ray=True` flag (Ray inside a
+# heterogeneous Slurm allocation) and from `nemo_run.core.execution.kuberay`
+# (Kubernetes-managed Ray clusters via the KubeRay operator).
+#
+# Supported in this release:
+#   - Single-command Ray jobs (SFT, eventually GRPO)
+#   - Dependency chaining via Ray submission IDs
+#   - Shared-FS runtime/code visibility (head + workers see the same paths)
+#
+# Out of scope (raises NotImplementedError):
+#   - Sandbox judge containers
+#   - Server co-scheduling (vLLM/SGLang/TRT-LLM alongside the main job)
+#   - Heterogeneous tasks
+#   - Multi-command task groups
+
+executor: ray
+
+# Ray cluster connection.
+ray:
+  # "auto" attaches to a Ray cluster started in the current environment
+  # (e.g., `ray start --head` in a Ray-on-Slurm setup, or RAY_ADDRESS env var).
+  # Use a "ray://host:10001" URI for a remote Ray client connection.
+  address: auto
+  # Namespace for job isolation across users/jobs on a shared cluster.
+  namespace: nemo
+  # Default *per-node* CPU allocation. NeMo-Skills multiplies this by the
+  # workflow's `num_nodes` to compute the per-job total CPU request, and Ray's
+  # `entrypoint_num_cpus` is then derived as total / num_nodes (= this value).
+  # GPU count is derived separately from the workflow's `num_gpus` parameter.
+  default_num_cpus: 8
+
+# Where Ray submission metadata + per-job logs are written. Should be on a
+# shared filesystem visible to head + workers so both sides see the same paths.
+jobs:
+  log_dir: /workspace/ray_jobs
+
+containers:
+  # Ray jobs typically use the same NeMo-RL / NeMo-Skills container as Slurm —
+  # specify the image refs here. For air-gapped deployments, use locally-built
+  # .sqsh / pre-staged images; runtime pulls are not required for the Ray path.
+  # Example (uncomment and fill in):
+  # nemo-rl: <local-path-or-registry-ref>
+  # nemo-skills: <local-path-or-registry-ref>
+  # vllm: <local-path-or-registry-ref>
+
+# Mounts visible to Ray workers. Code is auto-mounted at /nemo_run/code by
+# nemo-run for the Slurm path; for Ray, packaging happens via Ray's runtime_env
+# `working_dir`. If your shared FS already has the code staged, you do not need
+# to define a code mount here.
+#
+# mounts:
+#   - <shared-fs-path>/data:/data
+#   - <shared-fs-path>/models:/models
+
+# Environment variables for Ray jobs. HF_HOME is required by default — must be a
+# mounted path (or a path on the shared FS visible to Ray workers).
+# env_vars:
+#   - HF_HOME=/models/hf-cache
+#   - TOKENIZERS_PARALLELISM=false

nemo_skills/pipeline/utils/exp.py

@@ -47,6 +47,11 @@
     get_packager,
     get_registered_external_repo,
 )
+from nemo_skills.pipeline.utils.ray_executor import (
+    RayExecutor,
+    RayJobConfig,
+    get_ray_client,
+)
 from nemo_skills.pipeline.utils.server import get_free_port, get_server_command
 from nemo_skills.utils import get_logger_name, remove_handlers
 
@@ -270,6 +275,35 @@ def get_executor(
             additional_kwargs={"entrypoint": ""},
         )
 
+    if cluster_config["executor"] == "ray":
+        # Ray top-level scheduler executor (standalone Ray clusters or Ray-on-Slurm).
+        # Distinct from the existing `with_ray=True` flag, which is "Ray inside a Slurm
+        # allocation" via heterogeneous Slurm jobs. Here, Ray IS the scheduler.
+        # Actual job submission is performed in `add_task()` via `RayJobClient`; this
+        # branch returns a `RayExecutor` config object for callers that introspect the
+        # executor type (e.g., `Pipeline._create_executor`).
+        ray_config = cluster_config.get("ray", {})
+        # gpus_per_node convention (per get_executor docstring): 0 or None for
+        # CPU-only jobs. Treat None as "caller didn't say" → default 1, but
+        # respect an explicit 0 so CPU-only Ray jobs don't silently get a GPU.
+        if gpus_per_node is None:
+            ray_num_gpus = 1
+        else:
+            ray_num_gpus = int(gpus_per_node) * num_nodes
+        return RayExecutor(
+            ray_address=ray_config.get("address", "auto"),
+            ray_namespace=ray_config.get("namespace", "nemo"),
+            num_gpus=ray_num_gpus,
+            # cluster_config.ray.default_num_cpus is per-node; multiply by
+            # num_nodes to get the per-job total RayExecutor expects.
+            num_cpus=ray_config.get("default_num_cpus", 8) * num_nodes,
+            num_nodes=num_nodes,
+            ntasks_per_node=tasks_per_node,
+            log_dir=cluster_config.get("jobs", {}).get("log_dir", "/tmp/ray_jobs"),
+            env_vars=env_vars,
+            packager=packager,
+        )
+
     if not heterogeneous:
         env_vars["SLURM_MASTER_NODE"] = "$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n1)"
     else:
@@ -533,6 +567,78 @@ def add_task(
         if not is_mounted_filepath(cluster_config, env_vars["HF_HOME"]):
             raise RuntimeError(f"Invalid cluster_config: HF_HOME={env_vars['HF_HOME']} is not a mounted path.")
 
+    # Ray executor path: bypass nemo_run.Experiment.add() and submit directly via
+    # RayJobClient. Scoped to single-command Ray jobs; unsupported modes raise
+    # NotImplementedError to fail fast. Distinct from `with_ray=True`, which runs
+    # Ray inside a Slurm allocation rather than as the top-level scheduler.
+    if cluster_config["executor"] == "ray":
+        if with_sandbox:
+            raise NotImplementedError(
+                "Ray executor does not support sandbox containers in this release. "
+                "Sandbox judge containers are out of scope for the initial Ray support. "
+                "Use cluster_config.executor='slurm' for sandbox workflows."
+            )
+        if server_config is not None:
+            raise NotImplementedError(
+                "Ray executor does not support co-scheduled servers (vLLM/SGLang/TRT-LLM) "
+                "alongside the main job in this release. For Ray-based workflows that "
+                "require an inference endpoint, use a separately-deployed external endpoint "
+                "and call it via server_type='openai' (OpenAI-compatible)."
+            )
+        if heterogeneous:
+            raise NotImplementedError("Ray executor does not support heterogeneous tasks in this release.")
+        if isinstance(cmd, list) and len(cmd) > 1:
+            raise NotImplementedError(
+                f"Ray executor only supports single-command tasks in this release. Got {len(cmd)} commands."
+            )
+        if with_ray:
+            raise NotImplementedError(
+                "with_ray=True (Ray-inside-Slurm-allocation, heterogeneous Slurm jobs) is "
+                "not compatible with cluster_config.executor='ray' (Ray-as-top-level-scheduler). "
+                "Use exactly one path per task."
+            )
+
+        # Validate task_name length already happened above. Build the Ray submission.
+        ray_cmd = cmd if isinstance(cmd, str) else cmd[0]
+        ray_dependencies = []
+        for dep in task_dependencies or []:
+            # task_dependencies for the Ray path must be Ray submission IDs (strings).
+            # Slurm callers commonly pass nemo-run task handles; silently dropping those
+            # would cause the new job to run without waiting for its dependency, defeating
+            # the contract without any signal. Fail fast instead.
+            if not isinstance(dep, str):
+                raise NotImplementedError(
+                    f"Ray executor task_dependencies must be Ray submission IDs (str); "
+                    f"got {type(dep).__name__}. If you copied a slurm pattern, replace the "
+                    f"nemo-run task handle with the string returned by the prior "
+                    f"add_task() call under cluster_config.executor='ray'."
+                )
+            ray_dependencies.append(dep)
+        ray_cluster_config = cluster_config.get("ray", {})
+        # default_num_cpus is per-node; multiply by num_nodes to get the per-job total.
+        ray_default_cpus_per_node = ray_cluster_config.get("default_num_cpus", 8)
+        ray_log_dir = log_dir or cluster_config.get("jobs", {}).get("log_dir", "/tmp/ray_jobs")
+
+        ray_job_config = RayJobConfig(
+            name=task_name,
+            command=ray_cmd,
+            num_gpus=num_gpus if num_gpus is not None else 1,
+            num_cpus=ray_default_cpus_per_node * num_nodes,
+            num_nodes=num_nodes,
+            env_vars=env_vars,
+            log_dir=ray_log_dir,
+            dependencies=ray_dependencies if ray_dependencies else None,
+        )
+
+        if dry_run:
+            LOG.info("Dry run mode: would submit Ray job %s with command: %s", task_name, ray_cmd)
+            return f"<dry-run-ray:{task_name}>"
+
+        ray_client = get_ray_client(cluster_config)
+        submission_id = ray_client.submit_job(ray_job_config)
+        LOG.info("Ray job submitted: task=%s submission_id=%s", task_name, submission_id)
+        return submission_id
+
     het_group = 0
     het_group_indices = []
     total_het_groups = (n_servers if server_config is not None else 0) + bool(cmd) + with_sandbox