Add Ray Serve gateway for multi-instance vLLM spanning multiple nodes per instance

examples/slurm_vllm_ray_multi_instance_multi_node.yaml

@@ -0,0 +1,42 @@
+compute:
+  cluster:
+    type: slurm
+    hostname: myslurmcluster
+    walltime: "01:00:00"
+    account: myaccount
+    node_pools:
+      compute:
+        partition: batch
+        nodes: 4
+        ntasks_per_node: 1
+        gpus_per_node: 8
+
+services:
+  vllm_model:
+    container: vllm/vllm-openai:v0.9.0
+    type: vllm
+    model: Qwen/Qwen2.5-72B-Instruct
+    trust_remote_code: true
+    tensor_parallel_size: 8      # TP is capped at gpus_per_node (8); PP crosses node boundaries.
+    pipeline_parallel_size: 2    # TP8 x PP2 = 16 GPUs/instance, spans 2 nodes/instance (16 / 8).
+    number_of_instances: 2       # 2 instances x 2 nodes/instance = 4 nodes = the whole allocation.
+    port: 8000
+    health_check:
+      timeout_seconds: 1200
+    mounts:
+      - /lustre/datasets:/data
+
+driver:
+  policy_model: vllm_model
+  container: python:3.13
+  gym_install:
+    ref: main
+  benchmarks:
+    gpqa:
+      run:
+        split: benchmark
+        config_paths:
+          - benchmarks/gpqa/config.yaml
+
+job:
+  output_path: /lustre/fsw/my-path

nemo_gym/orchestration/api.py

@@ -55,6 +55,11 @@ class VllmServiceConfig(BaseModelServiceConfig):
     pipeline_parallel_size: int = 1
     trust_remote_code: bool = False
     number_of_instances: int = 1
+    # Opt-in escape hatch: run multiple instances behind a Ray Serve gateway (which handles both
+    # instance creation and request routing) instead of vLLM's own --data-parallel-size/multi-node
+    # DP mechanism. Most users never need to set this - it's forced on regardless of this value
+    # when the topology requires it (see effective_ray_serve).
+    use_ray_serve: bool = False
 
     @field_validator("number_of_instances")
     @classmethod
@@ -74,6 +79,39 @@ def _default_health_check(self) -> "VllmServiceConfig":
         return self
 
 
+def effective_ray_serve(service: "VllmServiceConfig", total_nodes: int, gpus_per_node_values: list[int]) -> bool:
+    """Whether the Ray Serve gateway (ray_serve_gateway.py) manages instance creation and request
+    routing for this service, instead of vLLM's own --data-parallel-size/multi-node DP mechanism.
+
+    True either because the user opted in via `use_ray_serve`, or because the topology requires
+    it: an instance's own tensor/pipeline-parallel footprint would have to span multiple nodes,
+    which vLLM's own multi-node data-parallel mechanism cannot express. Shared between api.py's
+    validation and slurm_script.py's command building so both agree on the same decision.
+    """
+    if service.use_ray_serve:
+        return True
+    if not gpus_per_node_values:
+        return False
+    max_gpus_per_node = max(gpus_per_node_values)
+    tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size
+    return total_nodes > 1 and service.number_of_instances > 1 and tp_pp > max_gpus_per_node
+
+
+def gym_install_required_message(service_name: str | None = None) -> str:
+    """Shared explanation for why driver.gym_install must be set whenever effective_ray_serve is
+    true for a service - raised both by api.py's validation (fires first, at config-load time) and
+    slurm_script.py's command building (defense-in-depth for callers of _build_service_command/
+    build_sbatch_script directly, bypassing SubmitConfig validation). One shared message so the
+    two call sites can't drift out of sync."""
+    subject = f"Service '{service_name}'" if service_name else "Service"
+    return (
+        f"{subject} requires the Ray Serve gateway (use_ray_serve or an instance spanning multiple "
+        "nodes), but driver.gym_install is not set. The gateway script "
+        "(nemo_gym/orchestration/ray_serve_gateway.py) is fetched from that repo/ref into the vLLM "
+        "service's own container - set driver.gym_install.{repo,ref}."
+    )
+
+
 class RayServiceConfig(BaseServiceConfig):
     type: Literal["ray"]
 
@@ -164,6 +202,13 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
             sum(p.nodes for p in compute.node_pools.values()) if isinstance(compute, SlurmComputeConfig) else 1
         )
         is_multi_node = total_nodes > 1
+        # Only one compute resource is supported, so every service's placement resolves to it -
+        # this is computed once here rather than re-derived per service.
+        gpus_per_node_values = (
+            [p.gpus_per_node for p in compute.node_pools.values() if p.gpus_per_node is not None]
+            if isinstance(compute, SlurmComputeConfig)
+            else []
+        )
 
         for service_name, service in self.services.items():
             if service.placement is None:
@@ -174,20 +219,29 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
                     f"({', '.join(sorted(compute_names))})."
                 )
 
+            if not isinstance(service, VllmServiceConfig):
+                continue
+
+            is_ray_serve = effective_ray_serve(service, total_nodes, gpus_per_node_values)
+
             if (
                 is_multi_node
-                and isinstance(service, VllmServiceConfig)
                 and service.number_of_instances > 1
                 and service.number_of_instances % total_nodes != 0
+                and not is_ray_serve
             ):
                 raise ValueError(
                     f"Service '{service_name}' has number_of_instances={service.number_of_instances}, which must "
                     f"be evenly divisible by the number of nodes ({total_nodes}) for multi-node data-parallel "
                     "deployment - each node hosts an equal share of the data-parallel replicas."
                 )
 
-            if isinstance(service, VllmServiceConfig):
-                self._validate_vllm_gpu_footprint(service_name, service, total_nodes)
+            self._validate_vllm_gpu_footprint(
+                service_name, service, total_nodes, compute, gpus_per_node_values, is_ray_serve
+            )
+
+            if is_ray_serve and self.driver.gym_install is None:
+                raise ValueError(gym_install_required_message(service_name))
 
         if self.driver.policy_model is not None:
             if self.driver.policy_model not in self.services:
@@ -213,39 +267,42 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
 
         return self
 
-    def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceConfig", total_nodes: int) -> None:
-        compute = self.compute[service.placement]
-        if not isinstance(compute, SlurmComputeConfig):
-            return
-
-        gpus_per_node_values = [
-            pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node is not None
-        ]
+    def _validate_vllm_gpu_footprint(
+        self,
+        service_name: str,
+        service: "VllmServiceConfig",
+        total_nodes: int,
+        compute: "SlurmComputeConfig",
+        gpus_per_node_values: list[int],
+        is_ray_serve: bool,
+    ) -> None:
         if not gpus_per_node_values:
             return
 
         max_gpus_per_node = max(gpus_per_node_values)
         tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size
 
-        if total_nodes > 1 and service.number_of_instances > 1:
-            if tp_pp > max_gpus_per_node:
-                # Each instance's own TP/PP footprint already exceeds a single node's GPU count, so
-                # spreading multiple such instances across nodes would require every instance to
-                # itself span multiple nodes. That's not supported: multi-node data-parallel only
-                # distributes whole instances across nodes with tensor/pipeline parallelism kept
-                # local to each node (see _build_vllm_multi_instance_multi_node_command).
-                raise ValueError(
-                    f"Service '{service_name}' sets number_of_instances={service.number_of_instances} with "
-                    f"tensor_parallel_size={service.tensor_parallel_size} x "
-                    f"pipeline_parallel_size={service.pipeline_parallel_size}={tp_pp}, which exceeds a single "
-                    f"node's gpus_per_node ({max_gpus_per_node}). Multiple instances where each instance's own "
-                    "tensor/pipeline-parallel footprint spans multiple nodes is not supported - reduce "
-                    "tensor_parallel_size/pipeline_parallel_size to fit within one node, or set "
-                    "number_of_instances=1 to let a single instance span nodes."
-                )
+        if total_nodes > 1 and is_ray_serve:
+            # Ray Serve gateway path: Ray's own placement-group scheduler packs each instance's
+            # TP*PP GPUs anywhere across the shared cluster, so an instance may itself span nodes
+            # and multiple instances may share a node. Only the aggregate footprint has to fit -
+            # see ray_serve_gateway.py and _build_vllm_ray_serve_command.
+            gpus_needed = tp_pp * service.number_of_instances
+            gpus_available = sum(
+                pool.nodes * pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node
+            )
+            footprint = (
+                f"tensor_parallel_size={service.tensor_parallel_size} x "
+                f"pipeline_parallel_size={service.pipeline_parallel_size} x "
+                f"number_of_instances={service.number_of_instances} (ray_serve gateway)"
+            )
+            scope = f"the total GPUs across all nodes ({gpus_available})"
+        elif total_nodes > 1 and service.number_of_instances > 1:
             # Multi-node data-parallel: each node runs its own equal share of the replicas with
             # local tensor/pipeline parallelism (see _build_vllm_multi_instance_multi_node_command);
             # the per-node share, not the total footprint, has to fit in that node's GPU count.
+            # tp_pp > max_gpus_per_node can't reach here - that shape always sets
+            # effective_ray_serve above.
             instances_per_node = service.number_of_instances // total_nodes
             gpus_needed = tp_pp * instances_per_node
             gpus_available = max_gpus_per_node

nemo_gym/orchestration/executors/script_templates.py

@@ -105,6 +105,22 @@ def render_gym_cmd(subcommand: str, var_name: str, args: list[str]) -> str:
     return f"{var_name}=(\n    " + "\n    ".join(entries) + "\n)"
 
 
+def render_repo_checkout(repo: str, ref: str) -> str:
+    """Render an &&-chained command that installs git if missing, then does
+    `git clone {repo} && cd {name} && git checkout {ref}`, leaving the shell's cwd at the repo
+    root. Shared by the driver entrypoint (which additionally pip-installs the whole package) and
+    the Ray Serve gateway command (which just needs the raw ray_serve_gateway.py file - see
+    _build_vllm_ray_serve_command). The git-install guard lives here rather than per call site
+    because every caller needs git first, regardless of container image - e.g. vllm/vllm-openai
+    doesn't bundle it, and neither does every driver image."""
+    repo_name = repo.rstrip("/").split("/")[-1].removesuffix(".git")
+    ensure_git = "command -v git >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq git)"
+    return (
+        f"({ensure_git})"
+        f" && git clone {shlex.quote(repo)} && cd {shlex.quote(repo_name)} && git checkout {shlex.quote(ref)}"
+    )
+
+
 def render_driver_entrypoint(
     repo: str | None,
     ref: str | None,
@@ -118,13 +134,10 @@ def render_driver_entrypoint(
     preamble: list[str] = []
 
     if repo and ref:
-        repo_name = repo.rstrip("/").split("/")[-1].removesuffix(".git")
         preamble += [
             "curl -LsSf https://astral.sh/uv/install.sh | sh",
             'source "$HOME/.local/bin/env"',
-            f"git clone {shlex.quote(repo)}",
-            f"cd {shlex.quote(repo_name)}",
-            f"git checkout {shlex.quote(ref)}",
+            render_repo_checkout(repo, ref),
             "uv pip install -e . --system",
         ]
 

nemo_gym/orchestration/executors/slurm_script.py

@@ -19,18 +19,22 @@
 
 from nemo_gym.orchestration.api import (
     BenchmarkRunConfig,
+    GymInstallConfig,
     NodePool,
     RayServiceConfig,
     SlurmComputeConfig,
     SubmitConfig,
     VllmServiceConfig,  # used in _BUILDERS dispatch table
+    effective_ray_serve,
+    gym_install_required_message,
 )
 from nemo_gym.orchestration.executors.script_templates import (
     bash_var,
     render_driver_entrypoint,
     render_gym_cmd,
     render_health_check,
     render_ray_prelude,
+    render_repo_checkout,
     render_vllm_ray_symmetric_run,
 )
 from nemo_gym.orchestration.executors.utils import flatten_run_args
@@ -201,6 +205,73 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str
     return _build_vllm_single_instance_multi_node_command(service, total_nodes)
 
 
+def _escape_for_double_quoted_bash(text: str) -> str:
+    """Escape text for safe embedding inside a double-quoted bash string ("..."). Needed instead
+    of shlex.quote (which produces single-quote-wrapped output) whenever the text is substituted
+    inside an already-open single-quoted bash region - a literal `'` from shlex.quote would
+    terminate that outer quoting early and corrupt the command."""
+    return text.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`")
+
+
+def _build_vllm_ray_serve_command(
+    service: VllmServiceConfig, total_nodes: int, gym_install: GymInstallConfig, gpus_per_node_values: list[int]
+) -> str:
+    # Ray Serve (nemo_gym/orchestration/ray_serve_gateway.py) launches all number_of_instances
+    # `vllm serve` processes itself and routes requests across them. Ray's own placement-group
+    # scheduler only decides where each instance's *worker* ranks land - not where the `vllm serve`
+    # driver process itself runs - so the gateway needs gpus_per_node to compute how many nodes
+    # each instance's own footprint needs, to explicitly pin different instances' drivers to
+    # different nodes. This is the only place NeMo Gym uses the `ray.serve` library, as opposed to
+    # vLLM's own Ray core executor (see _build_vllm_single_instance_multi_node_command).
+    gateway_args = (
+        f"--model {shlex.quote(service.model)}"
+        f" --port {service.port}"
+        f" --tensor-parallel-size {service.tensor_parallel_size}"
+        f" --pipeline-parallel-size {service.pipeline_parallel_size}"
+        f" --number-of-instances {service.number_of_instances}"
+    )
+    if gpus_per_node_values:
+        gateway_args += f" --gpus-per-node {max(gpus_per_node_values)}"
+    if service.trust_remote_code:
+        gateway_args += " --trust-remote-code"
+
+    # ray_serve_gateway.py has no internal nemo_gym imports (stdlib + ray/fastapi/aiohttp only), so
+    # it's fetched and run as a standalone script rather than `pip install -e .`-ing the whole
+    # nemo_gym package - the vLLM service's own container (e.g. vllm/vllm-openai) never has
+    # nemo_gym installed, and a full package install there risks nemo_gym's own pinned deps (torch,
+    # ray, ...) clobbering vLLM's already-working install. Reuses driver.gym_install (the same
+    # repo/ref the driver itself checks out) instead of a separate per-service field, so the two
+    # can't drift out of sync.
+    fetch_and_run = (
+        f"{render_repo_checkout(gym_install.repo, gym_install.ref)}"
+        " && pip install --quiet aiohttp"
+        f" && python3 nemo_gym/orchestration/ray_serve_gateway.py {gateway_args}"
+    )
+    if total_nodes <= 1:
+        # No multi-node Ray cluster to join - the gateway starts its own local Ray instance and
+        # launches all instances on this one node. Still needs double-quote escaping: fetch_and_run
+        # contains shlex.quote(service.model), which wraps its output in literal single quotes
+        # whenever the model name needs escaping (e.g. contains a space) - those would otherwise
+        # terminate this bash -lc '...' wrapper early and corrupt the command.
+        return f'bash -lc "{_escape_for_double_quoted_bash(fetch_and_run)}"'
+    resource_flags = (
+        "--num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE}"
+    )
+    # render_vllm_ray_symmetric_run splices inner_cmd, unquoted, into a `ray symmetric-run ... --
+    # {inner_cmd}` statement that itself sits inside that template's own single-quoted `bash -lc
+    # '...'` wrapper. Without protection, fetch_and_run's `&&` chain gets live-parsed as bash
+    # operators once that outer bash -lc runs its script: `ray symmetric-run`'s entrypoint would
+    # become just the `git clone` (the first `&&`-segment), which succeeds and exits immediately,
+    # tearing down the whole Ray cluster it just stood up before the repo checkout/install/gateway
+    # launch ever run. Wrapping in `bash -c "<escaped>"` makes the whole chain one opaque token
+    # immune to that live-parsing - double quotes, not shlex.quote's single-quote style, because
+    # this text is substituted inside that template's own single-quoted region: a literal `'`
+    # (which shlex.quote would introduce) would terminate that outer quoting early.
+    return render_vllm_ray_symmetric_run(
+        f'bash -c "{_escape_for_double_quoted_bash(fetch_and_run)}"', total_nodes, resource_flags
+    )
+
+
 def _build_ray_command(_service: RayServiceConfig) -> str:
     return "ray start --head"
 
@@ -218,7 +289,16 @@ def _vllm_spans_multiple_nodes(service: VllmServiceConfig | RayServiceConfig, to
     return isinstance(service, VllmServiceConfig) and total_nodes > 1
 
 
-def _build_service_command(service: VllmServiceConfig | RayServiceConfig, total_nodes: int) -> str:
+def _build_service_command(
+    service: VllmServiceConfig | RayServiceConfig,
+    total_nodes: int,
+    gpus_per_node_values: list[int],
+    gym_install: GymInstallConfig | None,
+) -> str:
+    if isinstance(service, VllmServiceConfig) and effective_ray_serve(service, total_nodes, gpus_per_node_values):
+        if gym_install is None:
+            raise ValueError(gym_install_required_message())
+        return _build_vllm_ray_serve_command(service, total_nodes, gym_install, gpus_per_node_values)
     if _vllm_spans_multiple_nodes(service, total_nodes):
         return _build_vllm_ray_command(service, total_nodes)
     return _BUILDERS[type(service)](service)
@@ -241,6 +321,9 @@ def build_sbatch_script(
 
     total_nodes, total_ntasks = _node_totals(compute)
     is_multi_node = total_nodes > 1
+    gpus_per_node_values = [
+        pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node is not None
+    ]
 
     ray_prelude = (
         render_ray_prelude()
@@ -252,7 +335,7 @@ def build_sbatch_script(
         _render_service_command(
             name,
             service.container,
-            _build_service_command(service, total_nodes),
+            _build_service_command(service, total_nodes, gpus_per_node_values, config.driver.gym_install),
             service.env or None,
             service.mounts or None,
             # Only services that actually span multiple nodes need the whole allocation's --nodes/