From 98370fb7817d64d312154d7dff493362b7e3472d Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Thu, 3 Sep 2026 13:34:56 -0400 Subject: [PATCH 1/8] Add ray serve for multi node span with multiple instances Signed-off-by: Onur Yilmaz --- ...rm_vllm_ray_multi_instance_multi_node.yaml | 42 ++++ nemo_gym/orchestration/api.py | 69 ++++-- .../orchestration/executors/slurm_script.py | 39 +++- nemo_gym/orchestration/ray_serve_gateway.py | 209 ++++++++++++++++++ tests/unit_tests/test_orchestration_api.py | 88 +++++++- tests/unit_tests/test_ray_serve_gateway.py | 146 ++++++++++++ tests/unit_tests/test_slurm_script.py | 84 +++++++ 7 files changed, 654 insertions(+), 23 deletions(-) create mode 100644 examples/slurm_vllm_ray_multi_instance_multi_node.yaml create mode 100644 nemo_gym/orchestration/ray_serve_gateway.py create mode 100644 tests/unit_tests/test_ray_serve_gateway.py diff --git a/examples/slurm_vllm_ray_multi_instance_multi_node.yaml b/examples/slurm_vllm_ray_multi_instance_multi_node.yaml new file mode 100644 index 0000000000..2fe8f55d86 --- /dev/null +++ b/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 diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 4bfddb8a0d..500135375a 100644 --- a/nemo_gym/orchestration/api.py +++ b/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 SubmitConfig._effective_ray_serve). + use_ray_serve: bool = False @field_validator("number_of_instances") @classmethod @@ -74,6 +79,24 @@ 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 + + class RayServiceConfig(BaseServiceConfig): type: Literal["ray"] @@ -179,6 +202,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": and isinstance(service, VllmServiceConfig) and service.number_of_instances > 1 and service.number_of_instances % total_nodes != 0 + and not self._effective_ray_serve(service, total_nodes) ): raise ValueError( f"Service '{service_name}' has number_of_instances={service.number_of_instances}, which must " @@ -213,6 +237,16 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": return self + def _effective_ray_serve(self, service: "VllmServiceConfig", total_nodes: int) -> bool: + """Whether the Ray Serve gateway path (see ray_serve_gateway.py) is used for this service.""" + compute = self.compute[service.placement] + if not isinstance(compute, SlurmComputeConfig): + return service.use_ray_serve + gpus_per_node_values = [ + pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node is not None + ] + return effective_ray_serve(service, total_nodes, gpus_per_node_values) + 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): @@ -226,26 +260,29 @@ def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceC max_gpus_per_node = max(gpus_per_node_values) tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size + effective_ray_serve = self._effective_ray_serve(service, total_nodes) - 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 effective_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 diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 030e8c4953..066ab3a014 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -24,6 +24,7 @@ SlurmComputeConfig, SubmitConfig, VllmServiceConfig, # used in _BUILDERS dispatch table + effective_ray_serve, ) from nemo_gym.orchestration.executors.script_templates import ( bash_var, @@ -201,6 +202,33 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str return _build_vllm_single_instance_multi_node_command(service, total_nodes) +def _build_vllm_ray_serve_command(service: VllmServiceConfig, total_nodes: 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 packs each instance's tensor/pipeline-parallel GPU footprint anywhere across the + # cluster, spanning nodes automatically when an instance's own footprint requires it. 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_cmd = ( + f"python -m nemo_gym.orchestration.ray_serve_gateway" + 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 service.trust_remote_code: + gateway_cmd += " --trust-remote-code" + 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. + return gateway_cmd + resource_flags = ( + "--num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE}" + ) + return render_vllm_ray_symmetric_run(gateway_cmd, total_nodes, resource_flags) + + def _build_ray_command(_service: RayServiceConfig) -> str: return "ray start --head" @@ -218,7 +246,11 @@ 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] +) -> str: + if isinstance(service, VllmServiceConfig) and effective_ray_serve(service, total_nodes, gpus_per_node_values): + return _build_vllm_ray_serve_command(service, total_nodes) if _vllm_spans_multiple_nodes(service, total_nodes): return _build_vllm_ray_command(service, total_nodes) return _BUILDERS[type(service)](service) @@ -241,6 +273,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 +287,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), service.env or None, service.mounts or None, # Only services that actually span multiple nodes need the whole allocation's --nodes/ diff --git a/nemo_gym/orchestration/ray_serve_gateway.py b/nemo_gym/orchestration/ray_serve_gateway.py new file mode 100644 index 0000000000..67eee89181 --- /dev/null +++ b/nemo_gym/orchestration/ray_serve_gateway.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 Serve gateway that launches multiple vLLM instances and routes requests across them. + +Selected automatically by `nemo_gym.orchestration` (see `api.effective_ray_serve` and +`executors/slurm_script._build_vllm_ray_serve_command`) whenever a vLLM service's +tensor/pipeline-parallel footprint would require an individual data-parallel instance to itself +span multiple Slurm nodes - something vLLM's own multi-node data-parallel mechanism can't express - +or whenever a user opts in via `use_ray_serve: true`. + +This process joins the (possibly multi-node) Ray cluster already bootstrapped by the sbatch script, +launches `number_of_instances` independent `vllm serve` subprocesses (each using vLLM's own Ray +core executor, `--distributed-executor-backend ray` - the same proven mechanism used for a single +instance spanning nodes), waits for each to become healthy, and then runs a Ray Serve HTTP ingress +that round-robins incoming requests across them. Ray's own placement-group scheduler decides where +each instance's tensor/pipeline-parallel GPU footprint lands, spanning nodes automatically when an +instance's own footprint requires it - this module never touches vLLM engine internals directly. +""" + +import argparse +import asyncio +import itertools +import logging +import os +import time + +import aiohttp +import ray +from fastapi import FastAPI, Request, Response +from ray import serve + + +logger = logging.getLogger(__name__) + +HEALTH_PATH = "/health" +HEALTH_POLL_INTERVAL_S = 5.0 +HEALTH_TIMEOUT_S = 900.0 + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--port", type=int, required=True, help="Port the gateway itself listens on.") + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--pipeline-parallel-size", type=int, default=1) + parser.add_argument("--number-of-instances", type=int, default=1) + parser.add_argument("--trust-remote-code", action="store_true") + return parser.parse_args(argv) + + +def instance_port(gateway_port: int, instance_index: int) -> int: + """Each backing vLLM instance listens on its own port, offset above the gateway's own port so + it never collides with the gateway's listening socket.""" + return gateway_port + 1 + instance_index + + +def build_instance_command(args: argparse.Namespace, instance_index: int) -> list[str]: + """Same flags as a single-instance-multi-node `vllm serve` invocation + (see `_build_vllm_single_instance_multi_node_command`), just run once per instance.""" + cmd = [ + "vllm", + "serve", + args.model, + "--port", + str(instance_port(args.port, instance_index)), + "--tensor-parallel-size", + str(args.tensor_parallel_size), + "--distributed-executor-backend", + "ray", + ] + if args.pipeline_parallel_size > 1: + cmd += ["--pipeline-parallel-size", str(args.pipeline_parallel_size)] + if args.trust_remote_code: + cmd.append("--trust-remote-code") + return cmd + + +class RoundRobinRouter: + """Cycles through backend instance URLs in order. Not load-aware - just spreads requests + evenly, mirroring what vLLM's own --data-parallel-size router would have done.""" + + def __init__(self, urls: list[str]): + if not urls: + raise ValueError("RoundRobinRouter requires at least one backend URL") + self._urls = list(urls) + self._cycle = itertools.cycle(self._urls) + + def next_url(self) -> str: + return next(self._cycle) + + +async def _wait_until_healthy( + session: aiohttp.ClientSession, url: str, proc: asyncio.subprocess.Process, instance_index: int, deadline: float +) -> None: + while True: + if proc.returncode is not None: + raise RuntimeError(f"vLLM instance {instance_index} ({url}) exited early with code {proc.returncode}") + try: + async with session.get(f"{url}{HEALTH_PATH}", timeout=aiohttp.ClientTimeout(total=5)) as resp: + if resp.status == 200: + logger.info("vLLM instance %d (%s) is healthy.", instance_index, url) + return + except (aiohttp.ClientError, asyncio.TimeoutError): + pass + if time.monotonic() > deadline: + raise TimeoutError(f"vLLM instance {instance_index} ({url}) did not become healthy in time") + await asyncio.sleep(HEALTH_POLL_INTERVAL_S) + + +async def launch_instances_and_wait(args: argparse.Namespace, gcs_address: str) -> list[str]: + """Launch every vLLM instance subprocess concurrently and block until all are healthy. + + Each `vllm serve` subprocess does its own internal `ray.init()` (inside its Ray distributed + executor) - without RAY_ADDRESS in its environment it can't discover the cluster this gateway + process already joined, and silently starts a separate, single-machine local Ray cluster of its + own instead. That defeats the whole point (Ray's placement-group scheduler no longer sees the + other instances, so multiple instances contend for the same GPUs). Passing RAY_ADDRESS + explicitly is what makes every instance join the one shared cluster. + + Returns the list of ready instance base URLs, in instance order. + """ + env = {**os.environ, "RAY_ADDRESS": gcs_address} + commands = [build_instance_command(args, i) for i in range(args.number_of_instances)] + procs = [await asyncio.create_subprocess_exec(*cmd, env=env) for cmd in commands] + urls = [f"http://localhost:{instance_port(args.port, i)}" for i in range(args.number_of_instances)] + + deadline = time.monotonic() + HEALTH_TIMEOUT_S + async with aiohttp.ClientSession() as session: + await asyncio.gather( + *(_wait_until_healthy(session, url, proc, i, deadline) for i, (url, proc) in enumerate(zip(urls, procs))) + ) + return urls + + +app = FastAPI() + + +@serve.deployment +@serve.ingress(app) +class VLLMGateway: + """Thin Ray Serve HTTP ingress that forwards every request to one of the ready vLLM instances. + + Ray Serve owns both instance creation (launch_instances_and_wait, called before serve.run) and + request routing (this class) - vLLM's own data-parallel mechanism is not used at all. + """ + + def __init__(self, instance_urls: list[str]): + self._router = RoundRobinRouter(instance_urls) + self._session = aiohttp.ClientSession() + + @app.get(HEALTH_PATH) + async def health(self) -> Response: + # By the time this deployment is serving traffic, every backing instance already passed + # its own health check in launch_instances_and_wait - nothing further to aggregate. + return Response(status_code=200) + + @app.api_route("/{path:path}", methods=["GET", "POST"]) + async def proxy(self, request: Request, path: str) -> Response: + target = self._router.next_url() + body = await request.body() + forward_headers = {k: v for k, v in request.headers.items() if k.lower() not in ("host", "content-length")} + async with self._session.request( + request.method, + f"{target}/{path}", + params=request.query_params, + data=body, + headers=forward_headers, + ) as resp: + content = await resp.read() + response_headers = {k: v for k, v in resp.headers.items() if k.lower() != "content-length"} + return Response(content=content, status_code=resp.status, headers=response_headers) + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + + try: + ray.init(address="auto") + except ConnectionError: + # No existing cluster to join (e.g. the single-node opt-in case, where the sbatch script + # skips the multi-node Ray bootstrap entirely) - start a local one. + ray.init() + gcs_address = ray.get_runtime_context().gcs_address + + instance_urls = asyncio.run(launch_instances_and_wait(args, gcs_address)) + + serve.start(http_options={"host": "0.0.0.0", "port": args.port}) + serve.run(VLLMGateway.bind(instance_urls)) + + logger.info("Ray Serve gateway ready on port %d, routing across %d instance(s).", args.port, len(instance_urls)) + while True: + time.sleep(3600) + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 8cfa854a1e..2e3eb0d615 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -204,11 +204,13 @@ def test_multi_node_dp_per_node_footprint_exceeds_raises(): ) -def test_multi_instance_per_instance_multi_node_tp_not_supported(): - # Each instance's own TP/PP footprint (TP5) already exceeds a single node's GPU count (4), so - # spreading 2 such instances across nodes would require each instance to itself span multiple - # nodes - that's not supported (only whole-instance placement across nodes is). - with pytest.raises(ValidationError, match="not supported"): +def test_multi_instance_per_instance_multi_node_tp_uses_ray_serve_gateway(): + # Each instance's own TP/PP footprint (TP5) exceeds a single node's GPU count (4): the Ray + # Serve gateway path is forced on automatically (no use_ray_serve needed) since only Ray's own + # placement-group scheduler - not vLLM's multi-node DP - can place such an instance. Aggregate + # footprint (5 x 2 = 10) exceeds COMPUTE_MULTI_NODE's total (2 nodes x 4 = 8), so this still + # raises, just against total cluster capacity rather than "not supported". + with pytest.raises(ValidationError, match="exceeds the total GPUs across all nodes"): SubmitConfig.model_validate( _config( services={"svc": {**SERVICE, "tensor_parallel_size": 5, "number_of_instances": 2}}, @@ -217,6 +219,82 @@ def test_multi_instance_per_instance_multi_node_tp_not_supported(): ) +COMPUTE_4_NODES_8_GPUS = { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"compute": {"partition": "batch", "nodes": 4, "gpus_per_node": 8}}, + } +} + + +def test_multi_instance_per_instance_multi_node_tp_accepted_when_footprint_fits(): + # TP8 x PP2 = 16 GPUs/instance > 8 gpus_per_node, so an instance must itself span nodes - the + # Ray Serve gateway path is forced on. 2 instances x 16 = 32 == 4 nodes x 8 gpus_per_node: fits + # exactly. This is the previously-forbidden topology that Ray Serve now supports. + config = SubmitConfig.model_validate( + _config( + services={ + "svc": {**SERVICE, "tensor_parallel_size": 8, "pipeline_parallel_size": 2, "number_of_instances": 2} + }, + compute=COMPUTE_4_NODES_8_GPUS, + ) + ) + assert config.services["svc"].number_of_instances == 2 + + +COMPUTE_4_NODES_6_GPUS = { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"compute": {"partition": "batch", "nodes": 4, "gpus_per_node": 6}}, + } +} + + +def test_multi_instance_per_instance_multi_node_tp_number_of_instances_need_not_divide_nodes(): + # 3 instances don't evenly divide 4 nodes, which would be rejected for vLLM's own multi-node DP + # - but Ray's placement-group scheduler packs flexibly, so the Ray Serve gateway path doesn't + # require this. tp_pp=8 (fits in one node) forces effective_ray_serve via use_ray_serve here so + # the mismatched node count is exercised without also tripping the footprint check (8 x 3 = 24 + # exactly matches 4 nodes x 6 gpus_per_node, so no idle-GPU warning either). + config = SubmitConfig.model_validate( + _config( + services={"svc": {**SERVICE, "tensor_parallel_size": 8, "number_of_instances": 3, "use_ray_serve": True}}, + compute=COMPUTE_4_NODES_6_GPUS, + ) + ) + assert config.services["svc"].number_of_instances == 3 + + +# --------------------------------------------------------------------------- +# use_ray_serve opt-in +# --------------------------------------------------------------------------- + + +def test_use_ray_serve_defaults_to_false(): + config = SubmitConfig.model_validate(_config()) + assert config.services["svc"].use_ray_serve is False + + +def test_use_ray_serve_opt_in_single_node_multi_instance_accepted(): + # Single-node, multi-instance: vLLM's own --data-parallel-size would normally handle this, but + # a user can still opt into the Ray Serve gateway for it. + service = {**SERVICE, "number_of_instances": 4, "use_ray_serve": True} + config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) + assert config.services["svc"].use_ray_serve is True + + +def test_use_ray_serve_opt_in_does_not_require_gpus_per_node(): + # No node_pools/gpus_per_node info at all - opting in still validates fine (nothing to check + # against), matching how the default path also skips footprint validation without that info. + service = {**SERVICE, "use_ray_serve": True} + config = SubmitConfig.model_validate(_config(services={"svc": service})) + assert config.services["svc"].use_ray_serve is True + + # --------------------------------------------------------------------------- # GPU footprint vs node pool capacity # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/test_ray_serve_gateway.py b/tests/unit_tests/test_ray_serve_gateway.py new file mode 100644 index 0000000000..db4ee72991 --- /dev/null +++ b/tests/unit_tests/test_ray_serve_gateway.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import pytest + +from nemo_gym.orchestration.ray_serve_gateway import ( + RoundRobinRouter, + build_instance_command, + instance_port, + parse_args, +) + + +# --------------------------------------------------------------------------- +# parse_args +# --------------------------------------------------------------------------- + + +def test_parse_args_required_fields(): + args = parse_args(["--model", "org/model", "--port", "8000"]) + assert args.model == "org/model" + assert args.port == 8000 + assert args.tensor_parallel_size == 1 + assert args.pipeline_parallel_size == 1 + assert args.number_of_instances == 1 + assert args.trust_remote_code is False + + +def test_parse_args_all_fields(): + args = parse_args( + [ + "--model", + "org/model", + "--port", + "9000", + "--tensor-parallel-size", + "8", + "--pipeline-parallel-size", + "2", + "--number-of-instances", + "4", + "--trust-remote-code", + ] + ) + assert args.tensor_parallel_size == 8 + assert args.pipeline_parallel_size == 2 + assert args.number_of_instances == 4 + assert args.trust_remote_code is True + + +def test_parse_args_missing_required_raises(): + with pytest.raises(SystemExit): + parse_args(["--port", "8000"]) + + +# --------------------------------------------------------------------------- +# instance_port +# --------------------------------------------------------------------------- + + +def test_instance_port_offsets_above_gateway_port(): + assert instance_port(8000, 0) == 8001 + assert instance_port(8000, 3) == 8004 + + +def test_instance_port_never_collides_with_gateway_port(): + for i in range(8): + assert instance_port(8000, i) != 8000 + + +# --------------------------------------------------------------------------- +# build_instance_command +# --------------------------------------------------------------------------- + + +def test_build_instance_command_basic(): + args = parse_args(["--model", "org/model", "--port", "8000"]) + cmd = build_instance_command(args, 0) + assert cmd[:3] == ["vllm", "serve", "org/model"] + assert "--port" in cmd and cmd[cmd.index("--port") + 1] == "8001" + assert "--tensor-parallel-size" in cmd + assert "--distributed-executor-backend" in cmd + assert cmd[cmd.index("--distributed-executor-backend") + 1] == "ray" + + +def test_build_instance_command_per_instance_port(): + args = parse_args(["--model", "org/model", "--port", "8000"]) + cmd0 = build_instance_command(args, 0) + cmd1 = build_instance_command(args, 1) + assert cmd0[cmd0.index("--port") + 1] == "8001" + assert cmd1[cmd1.index("--port") + 1] == "8002" + + +def test_build_instance_command_pipeline_parallel_flag_only_when_gt_1(): + args = parse_args(["--model", "org/model", "--port", "8000"]) + cmd = build_instance_command(args, 0) + assert "--pipeline-parallel-size" not in cmd + + args2 = parse_args(["--model", "org/model", "--port", "8000", "--pipeline-parallel-size", "2"]) + cmd2 = build_instance_command(args2, 0) + assert "--pipeline-parallel-size" in cmd2 + assert cmd2[cmd2.index("--pipeline-parallel-size") + 1] == "2" + + +def test_build_instance_command_trust_remote_code(): + args = parse_args(["--model", "org/model", "--port", "8000", "--trust-remote-code"]) + cmd = build_instance_command(args, 0) + assert "--trust-remote-code" in cmd + + +def test_build_instance_command_no_trust_remote_code_by_default(): + args = parse_args(["--model", "org/model", "--port", "8000"]) + cmd = build_instance_command(args, 0) + assert "--trust-remote-code" not in cmd + + +# --------------------------------------------------------------------------- +# RoundRobinRouter +# --------------------------------------------------------------------------- + + +def test_round_robin_router_cycles_in_order(): + router = RoundRobinRouter(["a", "b", "c"]) + assert [router.next_url() for _ in range(6)] == ["a", "b", "c", "a", "b", "c"] + + +def test_round_robin_router_single_url(): + router = RoundRobinRouter(["only"]) + assert [router.next_url() for _ in range(3)] == ["only", "only", "only"] + + +def test_round_robin_router_empty_raises(): + with pytest.raises(ValueError): + RoundRobinRouter([]) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index fcc30878a4..2853849b86 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -20,8 +20,10 @@ from nemo_gym.orchestration.api import SubmitConfig from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd from nemo_gym.orchestration.executors.slurm_script import ( + _build_service_command, _build_vllm_command, _build_vllm_ray_command, + _build_vllm_ray_serve_command, _node_totals, _render_directives, _render_pool_directives, @@ -265,6 +267,88 @@ def test_build_vllm_ray_command_dp_head_and_worker_branches(): assert "--data-parallel-start-rank $(( SLURM_NODEID * 2 ))" in cmd +# --------------------------------------------------------------------------- +# _build_vllm_ray_serve_command / Ray Serve gateway selection +# --------------------------------------------------------------------------- + + +def test_build_vllm_ray_serve_command_single_node_no_ray_bootstrap(vllm_service): + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1) + assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + assert "ray symmetric-run" not in cmd + assert "vllm serve" not in cmd # the gateway itself launches vllm serve, not this bash command + + +def test_build_vllm_ray_serve_command_multi_node_wraps_in_symmetric_run(vllm_service): + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=2) + assert "ray symmetric-run" in cmd + assert "--min-nodes 2" in cmd + assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + + +def test_build_vllm_ray_serve_command_passes_flags(): + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + port=9000, + tensor_parallel_size=8, + pipeline_parallel_size=2, + number_of_instances=2, + trust_remote_code=True, + ) + cmd = _build_vllm_ray_serve_command(service, total_nodes=4) + assert "--model org/model" in cmd + assert "--port 9000" in cmd + assert "--tensor-parallel-size 8" in cmd + assert "--pipeline-parallel-size 2" in cmd + assert "--number-of-instances 2" in cmd + assert "--trust-remote-code" in cmd + + +def test_build_service_command_uses_ray_serve_when_opted_in(vllm_service): + vllm_service.use_ray_serve = True + cmd = _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8]) + assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + + +def test_build_service_command_default_ignores_ray_serve_single_node(vllm_service): + cmd = _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8]) + assert "ray_serve_gateway" not in cmd + assert "vllm serve" in cmd + + +def test_build_service_command_mandatory_ray_serve_when_instance_spans_nodes(): + # TP*PP=16 > 8 GPUs/node with 2 instances on a 4-node/8-gpu allocation: an instance's own + # footprint must span nodes, which vLLM's own multi-node DP can't express - Ray Serve is + # forced on automatically, with no use_ray_serve set. + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + tensor_parallel_size=8, + pipeline_parallel_size=2, + number_of_instances=2, + ) + cmd = _build_service_command(service, total_nodes=4, gpus_per_node_values=[8]) + assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + + +def test_build_service_command_default_multi_instance_multi_node_unchanged(): + # tp_pp=8 fits within 8 gpus/node - stays on the existing (non-Ray) multi-node DP path even + # though it's multi-node and multi-instance. + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + tensor_parallel_size=8, + number_of_instances=4, + ) + cmd = _build_service_command(service, total_nodes=2, gpus_per_node_values=[8]) + assert "ray_serve_gateway" not in cmd + assert "--headless" in cmd + + # --------------------------------------------------------------------------- # render_gym_cmd # --------------------------------------------------------------------------- From 365b850860185e6df884399d9bd4cad0c1eef701 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Thu, 3 Sep 2026 22:00:42 -0400 Subject: [PATCH 2/8] Fix ray serve issues Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 7 +++ .../executors/script_templates.py | 9 ++++ .../orchestration/executors/slurm_script.py | 44 ++++++++++++++----- tests/unit_tests/test_orchestration_api.py | 37 +++++++++++++++- tests/unit_tests/test_slurm_script.py | 36 ++++++++++----- 5 files changed, 109 insertions(+), 24 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 500135375a..7111557430 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -212,6 +212,13 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": if isinstance(service, VllmServiceConfig): self._validate_vllm_gpu_footprint(service_name, service, total_nodes) + if self._effective_ray_serve(service, total_nodes) and self.driver.gym_install is None: + raise ValueError( + f"Service '{service_name}' 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}." + ) if self.driver.policy_model is not None: if self.driver.policy_model not in self.services: diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index 02d5e77ce8..8993129b6e 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -105,6 +105,15 @@ 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 `git clone {repo} && cd {name} && git checkout {ref}` as a single &&-chained + command, 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).""" + repo_name = repo.rstrip("/").split("/")[-1].removesuffix(".git") + return 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, diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 066ab3a014..1f1eb1541c 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -19,6 +19,7 @@ from nemo_gym.orchestration.api import ( BenchmarkRunConfig, + GymInstallConfig, NodePool, RayServiceConfig, SlurmComputeConfig, @@ -32,6 +33,7 @@ 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 @@ -202,31 +204,43 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str return _build_vllm_single_instance_multi_node_command(service, total_nodes) -def _build_vllm_ray_serve_command(service: VllmServiceConfig, total_nodes: int) -> str: +def _build_vllm_ray_serve_command(service: VllmServiceConfig, total_nodes: int, gym_install: GymInstallConfig) -> 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 packs each instance's tensor/pipeline-parallel GPU footprint anywhere across the # cluster, spanning nodes automatically when an instance's own footprint requires it. 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_cmd = ( - f"python -m nemo_gym.orchestration.ray_serve_gateway" - f" --model {shlex.quote(service.model)}" + 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 service.trust_remote_code: - gateway_cmd += " --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. - return gateway_cmd + return f"bash -lc '{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}" ) - return render_vllm_ray_symmetric_run(gateway_cmd, total_nodes, resource_flags) + return render_vllm_ray_symmetric_run(fetch_and_run, total_nodes, resource_flags) def _build_ray_command(_service: RayServiceConfig) -> str: @@ -247,10 +261,20 @@ def _vllm_spans_multiple_nodes(service: VllmServiceConfig | RayServiceConfig, to def _build_service_command( - service: VllmServiceConfig | RayServiceConfig, total_nodes: int, gpus_per_node_values: list[int] + 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): - return _build_vllm_ray_serve_command(service, total_nodes) + if gym_install is None: + raise ValueError( + "Service 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}." + ) + return _build_vllm_ray_serve_command(service, total_nodes, gym_install) if _vllm_spans_multiple_nodes(service, total_nodes): return _build_vllm_ray_command(service, total_nodes) return _BUILDERS[type(service)](service) @@ -287,7 +311,7 @@ def build_sbatch_script( _render_service_command( name, service.container, - _build_service_command(service, total_nodes, gpus_per_node_values), + _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/ diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 2e3eb0d615..1f0f954566 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -27,6 +27,9 @@ SERVICE = {"container": "gym:latest", "type": "vllm", "model": "org/model"} DRIVER = {"container": "gym:latest", "benchmarks": {"gsm8k": {}}} +# The Ray Serve gateway script is fetched from driver.gym_install's repo/ref into the vLLM +# service's own container - required whenever effective_ray_serve is true for a service. +DRIVER_WITH_GYM_INSTALL = {**DRIVER, "gym_install": {"ref": "main"}} JOB = {"output_path": "/tmp/gym-jobs"} @@ -239,11 +242,32 @@ def test_multi_instance_per_instance_multi_node_tp_accepted_when_footprint_fits( "svc": {**SERVICE, "tensor_parallel_size": 8, "pipeline_parallel_size": 2, "number_of_instances": 2} }, compute=COMPUTE_4_NODES_8_GPUS, + driver=DRIVER_WITH_GYM_INSTALL, ) ) assert config.services["svc"].number_of_instances == 2 +def test_multi_instance_per_instance_multi_node_tp_requires_gym_install(): + # Same topology as above but no driver.gym_install set - the Ray Serve gateway script has + # nowhere to be fetched from, so this must fail fast at config-validation time rather than + # only when the sbatch script is built. + with pytest.raises(ValidationError, match="gym_install"): + SubmitConfig.model_validate( + _config( + services={ + "svc": { + **SERVICE, + "tensor_parallel_size": 8, + "pipeline_parallel_size": 2, + "number_of_instances": 2, + } + }, + compute=COMPUTE_4_NODES_8_GPUS, + ) + ) + + COMPUTE_4_NODES_6_GPUS = { "cluster": { "type": "slurm", @@ -264,6 +288,7 @@ def test_multi_instance_per_instance_multi_node_tp_number_of_instances_need_not_ _config( services={"svc": {**SERVICE, "tensor_parallel_size": 8, "number_of_instances": 3, "use_ray_serve": True}}, compute=COMPUTE_4_NODES_6_GPUS, + driver=DRIVER_WITH_GYM_INSTALL, ) ) assert config.services["svc"].number_of_instances == 3 @@ -283,7 +308,9 @@ def test_use_ray_serve_opt_in_single_node_multi_instance_accepted(): # Single-node, multi-instance: vLLM's own --data-parallel-size would normally handle this, but # a user can still opt into the Ray Serve gateway for it. service = {**SERVICE, "number_of_instances": 4, "use_ray_serve": True} - config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) + config = SubmitConfig.model_validate( + _config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE, driver=DRIVER_WITH_GYM_INSTALL) + ) assert config.services["svc"].use_ray_serve is True @@ -291,10 +318,16 @@ def test_use_ray_serve_opt_in_does_not_require_gpus_per_node(): # No node_pools/gpus_per_node info at all - opting in still validates fine (nothing to check # against), matching how the default path also skips footprint validation without that info. service = {**SERVICE, "use_ray_serve": True} - config = SubmitConfig.model_validate(_config(services={"svc": service})) + config = SubmitConfig.model_validate(_config(services={"svc": service}, driver=DRIVER_WITH_GYM_INSTALL)) assert config.services["svc"].use_ray_serve is True +def test_use_ray_serve_opt_in_without_gym_install_raises(): + service = {**SERVICE, "use_ray_serve": True} + with pytest.raises(ValidationError, match="gym_install"): + SubmitConfig.model_validate(_config(services={"svc": service})) + + # --------------------------------------------------------------------------- # GPU footprint vs node pool capacity # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 2853849b86..05d657fc2a 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -17,7 +17,7 @@ import pytest -from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.api import GymInstallConfig, SubmitConfig from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd from nemo_gym.orchestration.executors.slurm_script import ( _build_service_command, @@ -272,18 +272,24 @@ def test_build_vllm_ray_command_dp_head_and_worker_branches(): # --------------------------------------------------------------------------- +_GYM_INSTALL = GymInstallConfig(repo="https://github.com/NVIDIA-NeMo/gym", ref="main") + + def test_build_vllm_ray_serve_command_single_node_no_ray_bootstrap(vllm_service): - cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1) - assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1, gym_install=_GYM_INSTALL) + assert "nemo_gym/orchestration/ray_serve_gateway.py" in cmd + assert "git clone" in cmd + assert "git checkout main" in cmd assert "ray symmetric-run" not in cmd assert "vllm serve" not in cmd # the gateway itself launches vllm serve, not this bash command def test_build_vllm_ray_serve_command_multi_node_wraps_in_symmetric_run(vllm_service): - cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=2) + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=2, gym_install=_GYM_INSTALL) assert "ray symmetric-run" in cmd assert "--min-nodes 2" in cmd - assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + assert "nemo_gym/orchestration/ray_serve_gateway.py" in cmd + assert "git clone" in cmd def test_build_vllm_ray_serve_command_passes_flags(): @@ -297,7 +303,7 @@ def test_build_vllm_ray_serve_command_passes_flags(): number_of_instances=2, trust_remote_code=True, ) - cmd = _build_vllm_ray_serve_command(service, total_nodes=4) + cmd = _build_vllm_ray_serve_command(service, total_nodes=4, gym_install=_GYM_INSTALL) assert "--model org/model" in cmd assert "--port 9000" in cmd assert "--tensor-parallel-size 8" in cmd @@ -308,12 +314,18 @@ def test_build_vllm_ray_serve_command_passes_flags(): def test_build_service_command_uses_ray_serve_when_opted_in(vllm_service): vllm_service.use_ray_serve = True - cmd = _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8]) - assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + cmd = _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8], gym_install=_GYM_INSTALL) + assert "nemo_gym/orchestration/ray_serve_gateway.py" in cmd + + +def test_build_service_command_missing_gym_install_raises(vllm_service): + vllm_service.use_ray_serve = True + with pytest.raises(ValueError, match="gym_install"): + _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8], gym_install=None) def test_build_service_command_default_ignores_ray_serve_single_node(vllm_service): - cmd = _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8]) + cmd = _build_service_command(vllm_service, total_nodes=1, gpus_per_node_values=[8], gym_install=None) assert "ray_serve_gateway" not in cmd assert "vllm serve" in cmd @@ -330,8 +342,8 @@ def test_build_service_command_mandatory_ray_serve_when_instance_spans_nodes(): pipeline_parallel_size=2, number_of_instances=2, ) - cmd = _build_service_command(service, total_nodes=4, gpus_per_node_values=[8]) - assert "python -m nemo_gym.orchestration.ray_serve_gateway" in cmd + cmd = _build_service_command(service, total_nodes=4, gpus_per_node_values=[8], gym_install=_GYM_INSTALL) + assert "nemo_gym/orchestration/ray_serve_gateway.py" in cmd def test_build_service_command_default_multi_instance_multi_node_unchanged(): @@ -344,7 +356,7 @@ def test_build_service_command_default_multi_instance_multi_node_unchanged(): tensor_parallel_size=8, number_of_instances=4, ) - cmd = _build_service_command(service, total_nodes=2, gpus_per_node_values=[8]) + cmd = _build_service_command(service, total_nodes=2, gpus_per_node_values=[8], gym_install=None) assert "ray_serve_gateway" not in cmd assert "--headless" in cmd From 370954f96727e804f82d0d24eb2338714dac3b7c Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Fri, 4 Sep 2026 11:02:12 -0400 Subject: [PATCH 3/8] Fixing ray serve issue Signed-off-by: Onur Yilmaz --- .../orchestration/executors/slurm_script.py | 19 ++-- nemo_gym/orchestration/ray_serve_gateway.py | 101 +++++++++++++++--- tests/unit_tests/test_ray_serve_gateway.py | 50 +++++++++ tests/unit_tests/test_slurm_script.py | 23 +++- 4 files changed, 170 insertions(+), 23 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 1f1eb1541c..a4f4def057 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -204,13 +204,16 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str return _build_vllm_single_instance_multi_node_command(service, total_nodes) -def _build_vllm_ray_serve_command(service: VllmServiceConfig, total_nodes: int, gym_install: GymInstallConfig) -> str: +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 packs each instance's tensor/pipeline-parallel GPU footprint anywhere across the - # cluster, spanning nodes automatically when an instance's own footprint requires it. 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). + # `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}" @@ -218,6 +221,8 @@ def _build_vllm_ray_serve_command(service: VllmServiceConfig, total_nodes: int, 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" @@ -274,7 +279,7 @@ def _build_service_command( "(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}." ) - return _build_vllm_ray_serve_command(service, total_nodes, gym_install) + 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) diff --git a/nemo_gym/orchestration/ray_serve_gateway.py b/nemo_gym/orchestration/ray_serve_gateway.py index 67eee89181..96ce1f6e81 100644 --- a/nemo_gym/orchestration/ray_serve_gateway.py +++ b/nemo_gym/orchestration/ray_serve_gateway.py @@ -26,15 +26,22 @@ core executor, `--distributed-executor-backend ray` - the same proven mechanism used for a single instance spanning nodes), waits for each to become healthy, and then runs a Ray Serve HTTP ingress that round-robins incoming requests across them. Ray's own placement-group scheduler decides where -each instance's tensor/pipeline-parallel GPU footprint lands, spanning nodes automatically when an -instance's own footprint requires it - this module never touches vLLM engine internals directly. +each instance's *worker* ranks land, spanning nodes automatically when an instance's own footprint +requires it - but it does NOT decide where the `vllm serve` driver process itself runs (that's +just wherever the OS process that launched it happens to execute). So each instance's driver is +explicitly pinned to a different node (round-robin, `nodes_per_instance` apart) via a small Ray +actor with node-affinity scheduling - otherwise every instance's driver would land on this +process's own node, and vLLM refuses to start once that node's local GPU share is exhausted by an +earlier instance, even though other nodes in the cluster are completely free. """ import argparse import asyncio import itertools import logging +import math import os +import subprocess import time import aiohttp @@ -48,6 +55,7 @@ HEALTH_PATH = "/health" HEALTH_POLL_INTERVAL_S = 5.0 HEALTH_TIMEOUT_S = 900.0 +NODE_AFFINITY_RESOURCE_WEIGHT = 0.001 def parse_args(argv: list[str] | None = None) -> argparse.Namespace: @@ -57,6 +65,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--tensor-parallel-size", type=int, default=1) parser.add_argument("--pipeline-parallel-size", type=int, default=1) parser.add_argument("--number-of-instances", type=int, default=1) + parser.add_argument( + "--gpus-per-node", + type=int, + default=None, + help="Used only to compute how many nodes a single instance's own TP*PP footprint needs " + "(for node-affinity assignment). Omit if every instance fits on one node.", + ) parser.add_argument("--trust-remote-code", action="store_true") return parser.parse_args(argv) @@ -67,6 +82,30 @@ def instance_port(gateway_port: int, instance_index: int) -> int: return gateway_port + 1 + instance_index +def nodes_per_instance(tensor_parallel_size: int, pipeline_parallel_size: int, gpus_per_node: int | None) -> int: + """How many physical nodes a single instance's own TP*PP footprint needs.""" + if not gpus_per_node: + return 1 + return max(1, math.ceil((tensor_parallel_size * pipeline_parallel_size) / gpus_per_node)) + + +def alive_node_ips() -> list[str]: + """Sorted, deduplicated IPs of every alive Ray node - deterministic so node assignment is + stable across the gateway process and any code that needs to reason about it.""" + ips = {n["NodeManagerAddress"] for n in ray.nodes() if n.get("Alive") and n.get("NodeManagerAddress")} + return sorted(ips) + + +def node_for_instance(instance_index: int, instance_nodes: int, node_ips: list[str]) -> str: + """Round-robin node assignment, `instance_nodes` apart, so each instance's driver (and, via + vLLM's own placement group, its worker ranks) lands on a distinct slice of the cluster instead + of every instance's driver piling onto this process's own node.""" + if not node_ips: + raise RuntimeError("No alive Ray nodes found - is the Ray cluster up?") + start = (instance_index * instance_nodes) % len(node_ips) + return node_ips[start] + + def build_instance_command(args: argparse.Namespace, instance_index: int) -> list[str]: """Same flags as a single-instance-multi-node `vllm serve` invocation (see `_build_vllm_single_instance_multi_node_command`), just run once per instance.""" @@ -88,6 +127,23 @@ def build_instance_command(args: argparse.Namespace, instance_index: int) -> lis return cmd +@ray.remote(num_cpus=0, num_gpus=0) +class _InstanceProcess: + """Supervises one `vllm serve` OS subprocess. Scheduled with a node-affinity resource (see + node_for_instance) so its physical node is chosen explicitly rather than left to chance - + Ray's placement-group scheduler only decides where vLLM's *worker* ranks land, not where this + driver process itself runs. num_gpus=0 here is intentional: this actor doesn't claim any GPU + itself, so it doesn't compete with vLLM's own internal placement-group GPU request for the + node it's pinned to. + """ + + def __init__(self, cmd: list[str], env: dict[str, str]) -> None: + self._proc = subprocess.Popen(cmd, env=env) + + def poll(self) -> int | None: + return self._proc.poll() + + class RoundRobinRouter: """Cycles through backend instance URLs in order. Not load-aware - just spreads requests evenly, mirroring what vLLM's own --data-parallel-size router would have done.""" @@ -103,11 +159,12 @@ def next_url(self) -> str: async def _wait_until_healthy( - session: aiohttp.ClientSession, url: str, proc: asyncio.subprocess.Process, instance_index: int, deadline: float + session: aiohttp.ClientSession, url: str, actor: "ray.actor.ActorHandle", instance_index: int, deadline: float ) -> None: while True: - if proc.returncode is not None: - raise RuntimeError(f"vLLM instance {instance_index} ({url}) exited early with code {proc.returncode}") + returncode = await actor.poll.remote() + if returncode is not None: + raise RuntimeError(f"vLLM instance {instance_index} ({url}) exited early with code {returncode}") try: async with session.get(f"{url}{HEALTH_PATH}", timeout=aiohttp.ClientTimeout(total=5)) as resp: if resp.status == 200: @@ -120,7 +177,9 @@ async def _wait_until_healthy( await asyncio.sleep(HEALTH_POLL_INTERVAL_S) -async def launch_instances_and_wait(args: argparse.Namespace, gcs_address: str) -> list[str]: +async def launch_instances_and_wait( + args: argparse.Namespace, gcs_address: str +) -> tuple[list[str], list["ray.actor.ActorHandle"]]: """Launch every vLLM instance subprocess concurrently and block until all are healthy. Each `vllm serve` subprocess does its own internal `ray.init()` (inside its Ray distributed @@ -130,19 +189,32 @@ async def launch_instances_and_wait(args: argparse.Namespace, gcs_address: str) other instances, so multiple instances contend for the same GPUs). Passing RAY_ADDRESS explicitly is what makes every instance join the one shared cluster. - Returns the list of ready instance base URLs, in instance order. + Returns (instance base URLs, actor handles) in instance order - the caller must keep the actor + handles alive for as long as the instances should keep running (Ray kills an actor once its + last handle is garbage collected). """ + instance_nodes = nodes_per_instance(args.tensor_parallel_size, args.pipeline_parallel_size, args.gpus_per_node) + node_ips = alive_node_ips() env = {**os.environ, "RAY_ADDRESS": gcs_address} - commands = [build_instance_command(args, i) for i in range(args.number_of_instances)] - procs = [await asyncio.create_subprocess_exec(*cmd, env=env) for cmd in commands] - urls = [f"http://localhost:{instance_port(args.port, i)}" for i in range(args.number_of_instances)] + + actors = [] + urls = [] + for i in range(args.number_of_instances): + node_ip = node_for_instance(i, instance_nodes, node_ips) + cmd = build_instance_command(args, i) + actor = _InstanceProcess.options(resources={f"node:{node_ip}": NODE_AFFINITY_RESOURCE_WEIGHT}).remote(cmd, env) + actors.append(actor) + urls.append(f"http://{node_ip}:{instance_port(args.port, i)}") deadline = time.monotonic() + HEALTH_TIMEOUT_S async with aiohttp.ClientSession() as session: await asyncio.gather( - *(_wait_until_healthy(session, url, proc, i, deadline) for i, (url, proc) in enumerate(zip(urls, procs))) + *( + _wait_until_healthy(session, url, actor, i, deadline) + for i, (url, actor) in enumerate(zip(urls, actors)) + ) ) - return urls + return urls, actors app = FastAPI() @@ -195,7 +267,10 @@ def main(argv: list[str] | None = None) -> None: ray.init() gcs_address = ray.get_runtime_context().gcs_address - instance_urls = asyncio.run(launch_instances_and_wait(args, gcs_address)) + # instance_actors must stay referenced for the rest of this (never-returning) function - Ray + # kills an actor once its last handle is garbage collected, and these back the running vLLM + # instances. Since main() never returns (it blocks forever below), this frame's locals live on. + instance_urls, instance_actors = asyncio.run(launch_instances_and_wait(args, gcs_address)) # noqa: F841 serve.start(http_options={"host": "0.0.0.0", "port": args.port}) serve.run(VLLMGateway.bind(instance_urls)) diff --git a/tests/unit_tests/test_ray_serve_gateway.py b/tests/unit_tests/test_ray_serve_gateway.py index db4ee72991..2e9ffaf865 100644 --- a/tests/unit_tests/test_ray_serve_gateway.py +++ b/tests/unit_tests/test_ray_serve_gateway.py @@ -19,6 +19,8 @@ RoundRobinRouter, build_instance_command, instance_port, + node_for_instance, + nodes_per_instance, parse_args, ) @@ -126,6 +128,54 @@ def test_build_instance_command_no_trust_remote_code_by_default(): assert "--trust-remote-code" not in cmd +# --------------------------------------------------------------------------- +# nodes_per_instance +# --------------------------------------------------------------------------- + + +def test_nodes_per_instance_no_gpus_per_node_defaults_to_one(): + assert nodes_per_instance(tensor_parallel_size=8, pipeline_parallel_size=2, gpus_per_node=None) == 1 + + +def test_nodes_per_instance_fits_in_one_node(): + assert nodes_per_instance(tensor_parallel_size=8, pipeline_parallel_size=1, gpus_per_node=8) == 1 + + +def test_nodes_per_instance_spans_two_nodes(): + assert nodes_per_instance(tensor_parallel_size=8, pipeline_parallel_size=2, gpus_per_node=8) == 2 + + +def test_nodes_per_instance_rounds_up(): + assert nodes_per_instance(tensor_parallel_size=5, pipeline_parallel_size=1, gpus_per_node=4) == 2 + + +# --------------------------------------------------------------------------- +# node_for_instance +# --------------------------------------------------------------------------- + + +def test_node_for_instance_single_node_per_instance_round_robins(): + nodes = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"] + assert [node_for_instance(i, 1, nodes) for i in range(4)] == nodes + + +def test_node_for_instance_wraps_around_when_more_instances_than_nodes(): + nodes = ["10.0.0.1", "10.0.0.2"] + assert [node_for_instance(i, 1, nodes) for i in range(4)] == nodes + nodes + + +def test_node_for_instance_multi_node_per_instance_uses_distinct_slices(): + # 4 nodes, 2 nodes/instance -> instance 0 anchors to node 0, instance 1 to node 2. + nodes = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"] + assert node_for_instance(0, 2, nodes) == "10.0.0.1" + assert node_for_instance(1, 2, nodes) == "10.0.0.3" + + +def test_node_for_instance_empty_nodes_raises(): + with pytest.raises(RuntimeError): + node_for_instance(0, 1, []) + + # --------------------------------------------------------------------------- # RoundRobinRouter # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 05d657fc2a..0436ec9d36 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -276,7 +276,7 @@ def test_build_vllm_ray_command_dp_head_and_worker_branches(): def test_build_vllm_ray_serve_command_single_node_no_ray_bootstrap(vllm_service): - cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1, gym_install=_GYM_INSTALL) + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1, gym_install=_GYM_INSTALL, gpus_per_node_values=[]) assert "nemo_gym/orchestration/ray_serve_gateway.py" in cmd assert "git clone" in cmd assert "git checkout main" in cmd @@ -285,13 +285,30 @@ def test_build_vllm_ray_serve_command_single_node_no_ray_bootstrap(vllm_service) def test_build_vllm_ray_serve_command_multi_node_wraps_in_symmetric_run(vllm_service): - cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=2, gym_install=_GYM_INSTALL) + cmd = _build_vllm_ray_serve_command( + vllm_service, total_nodes=2, gym_install=_GYM_INSTALL, gpus_per_node_values=[8] + ) assert "ray symmetric-run" in cmd assert "--min-nodes 2" in cmd assert "nemo_gym/orchestration/ray_serve_gateway.py" in cmd assert "git clone" in cmd +def test_build_vllm_ray_serve_command_passes_gpus_per_node(): + cmd = _build_vllm_ray_serve_command( + VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model"), + total_nodes=2, + gym_install=_GYM_INSTALL, + gpus_per_node_values=[8], + ) + assert "--gpus-per-node 8" in cmd + + +def test_build_vllm_ray_serve_command_omits_gpus_per_node_when_unknown(vllm_service): + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1, gym_install=_GYM_INSTALL, gpus_per_node_values=[]) + assert "--gpus-per-node" not in cmd + + def test_build_vllm_ray_serve_command_passes_flags(): service = VllmServiceConfig( type="vllm", @@ -303,7 +320,7 @@ def test_build_vllm_ray_serve_command_passes_flags(): number_of_instances=2, trust_remote_code=True, ) - cmd = _build_vllm_ray_serve_command(service, total_nodes=4, gym_install=_GYM_INSTALL) + cmd = _build_vllm_ray_serve_command(service, total_nodes=4, gym_install=_GYM_INSTALL, gpus_per_node_values=[8]) assert "--model org/model" in cmd assert "--port 9000" in cmd assert "--tensor-parallel-size 8" in cmd From 36dacb952e4b76da2fe977d2b16449cbf07ba300 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Sun, 6 Sep 2026 12:02:03 -0400 Subject: [PATCH 4/8] fix _build_vllm_ray_serve_command's Signed-off-by: Onur Yilmaz --- .../orchestration/executors/slurm_script.py | 13 ++++++++- tests/unit_tests/test_slurm_script.py | 28 +++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index a4f4def057..33df947f40 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -245,7 +245,18 @@ def _build_vllm_ray_serve_command( resource_flags = ( "--num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE}" ) - return render_vllm_ray_symmetric_run(fetch_and_run, total_nodes, resource_flags) + # 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 ""` 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. + escaped = fetch_and_run.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") + return render_vllm_ray_symmetric_run(f'bash -c "{escaped}"', total_nodes, resource_flags) def _build_ray_command(_service: RayServiceConfig) -> str: diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 0436ec9d36..1a87302adb 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import subprocess from pathlib import Path import pytest @@ -294,6 +295,33 @@ def test_build_vllm_ray_serve_command_multi_node_wraps_in_symmetric_run(vllm_ser assert "git clone" in cmd +def test_build_vllm_ray_serve_command_multi_node_chain_survives_symmetric_run_entrypoint(vllm_service): + # Regression test: `ray symmetric-run`'s entrypoint (everything after `--`) must be the whole + # git-clone-then-install-then-launch chain as ONE unit, not split by the outer bash -lc on the + # chain's own `&&` operators - a naive unquoted embedding lets that outer shell live-parse + # those operators, so `ray symmetric-run`'s entrypoint becomes just `git clone ...`, which + # succeeds and exits immediately, tearing the whole Ray cluster down before the gateway ever + # launches. Runs the *actual* generated bash through a stand-in `ray symmetric-run` to prove + # the whole chain lands as a single argv token, the same way it would for the real command. + cmd = _build_vllm_ray_serve_command( + vllm_service, total_nodes=2, gym_install=_GYM_INSTALL, gpus_per_node_values=[8] + ) + + script = cmd.replace( + "ray symmetric-run", + 'fake_symmetric_run() { for a in "$@"; do echo "ARG:$a"; done; }; fake_symmetric_run', + ).replace("if ray symmetric-run --help", "if true") + result = subprocess.run(["bash", "-c", script], capture_output=True, text=True, timeout=10) + + assert result.returncode == 0, result.stderr + args = [line.removeprefix("ARG:") for line in result.stdout.splitlines()] + assert args[-3:-1] == ["bash", "-c"] + chain = args[-1] + assert "git clone" in chain + assert "&&" in chain + assert "python3 nemo_gym/orchestration/ray_serve_gateway.py" in chain + + def test_build_vllm_ray_serve_command_passes_gpus_per_node(): cmd = _build_vllm_ray_serve_command( VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model"), From 71cc920dad61882b834574c31d6a2e9c7481dd73 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 7 Sep 2026 10:40:00 -0400 Subject: [PATCH 5/8] Fix git not found error Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/executors/slurm_script.py | 5 ++++- tests/unit_tests/test_slurm_script.py | 8 ++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 33df947f40..e7f87bf8a7 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -234,7 +234,10 @@ def _build_vllm_ray_serve_command( # 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)}" + # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle git - install it on + # the fly if missing, the same way render_vllm_ray_symmetric_run does for the ray CLI. + "(command -v git >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq git))" + 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}" ) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 1a87302adb..95d17933c9 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -285,6 +285,14 @@ def test_build_vllm_ray_serve_command_single_node_no_ray_bootstrap(vllm_service) assert "vllm serve" not in cmd # the gateway itself launches vllm serve, not this bash command +def test_build_vllm_ray_serve_command_installs_git_if_missing(vllm_service): + # Model-serving images (e.g. vllm/vllm-openai) don't always bundle git. + cmd = _build_vllm_ray_serve_command(vllm_service, total_nodes=1, gym_install=_GYM_INSTALL, gpus_per_node_values=[]) + assert "command -v git >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq git)" in cmd + # The git-install guard must come before the clone it's guarding. + assert cmd.index("command -v git") < cmd.index("git clone") + + def test_build_vllm_ray_serve_command_multi_node_wraps_in_symmetric_run(vllm_service): cmd = _build_vllm_ray_serve_command( vllm_service, total_nodes=2, gym_install=_GYM_INSTALL, gpus_per_node_values=[8] From 60f870e4afd943a4a7a7696429b8ef6a53d8628a Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 8 Sep 2026 16:02:29 -0400 Subject: [PATCH 6/8] Simplify the code Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 67 ++++++++++--------- .../executors/script_templates.py | 5 +- 2 files changed, 36 insertions(+), 36 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 7111557430..e1c9192590 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -58,7 +58,7 @@ class VllmServiceConfig(BaseModelServiceConfig): # 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 SubmitConfig._effective_ray_serve). + # when the topology requires it (see effective_ray_serve). use_ray_serve: bool = False @field_validator("number_of_instances") @@ -187,6 +187,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: @@ -197,12 +204,16 @@ 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 self._effective_ray_serve(service, total_nodes) + and not is_ray_serve ): raise ValueError( f"Service '{service_name}' has number_of_instances={service.number_of_instances}, which must " @@ -210,15 +221,17 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": "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) - if self._effective_ray_serve(service, total_nodes) and self.driver.gym_install is None: - raise ValueError( - f"Service '{service_name}' 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}." - ) + 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( + f"Service '{service_name}' 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}." + ) if self.driver.policy_model is not None: if self.driver.policy_model not in self.services: @@ -244,32 +257,22 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": return self - def _effective_ray_serve(self, service: "VllmServiceConfig", total_nodes: int) -> bool: - """Whether the Ray Serve gateway path (see ray_serve_gateway.py) is used for this service.""" - compute = self.compute[service.placement] - if not isinstance(compute, SlurmComputeConfig): - return service.use_ray_serve - gpus_per_node_values = [ - pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node is not None - ] - return effective_ray_serve(service, total_nodes, gpus_per_node_values) - - 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 - effective_ray_serve = self._effective_ray_serve(service, total_nodes) - if total_nodes > 1 and effective_ray_serve: + 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 - diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index 8993129b6e..92a2839ba2 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -127,13 +127,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", ] From 9b5190ef669910128f98e7e5e0594b3f7495d822 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 8 Sep 2026 16:58:23 -0400 Subject: [PATCH 7/8] Remove the second guard which is not needed Signed-off-by: Onur Yilmaz --- .../orchestration/executors/script_templates.py | 17 ++++++++++++----- .../orchestration/executors/slurm_script.py | 5 +---- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index 92a2839ba2..2a9a95b16f 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -106,12 +106,19 @@ def render_gym_cmd(subcommand: str, var_name: str, args: list[str]) -> str: def render_repo_checkout(repo: str, ref: str) -> str: - """Render `git clone {repo} && cd {name} && git checkout {ref}` as a single &&-chained - command, 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).""" + """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") - return f"git clone {shlex.quote(repo)} && cd {shlex.quote(repo_name)} && git checkout {shlex.quote(ref)}" + 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( diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index e7f87bf8a7..33df947f40 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -234,10 +234,7 @@ def _build_vllm_ray_serve_command( # 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 = ( - # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle git - install it on - # the fly if missing, the same way render_vllm_ray_symmetric_run does for the ray CLI. - "(command -v git >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq git))" - f" && {render_repo_checkout(gym_install.repo, gym_install.ref)}" + 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}" ) From 8b2ec7782cd6c234fac20781a818bbed53bcbe4f Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 8 Sep 2026 17:53:29 -0400 Subject: [PATCH 8/8] Fix the _escape_for_double_quoted_bash issue in the build command Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 22 +++++++--- .../orchestration/executors/slurm_script.py | 28 +++++++----- tests/unit_tests/test_slurm_script.py | 44 ++++++++++++++++++- 3 files changed, 77 insertions(+), 17 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index e1c9192590..e6903fb026 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -97,6 +97,21 @@ def effective_ray_serve(service: "VllmServiceConfig", total_nodes: int, gpus_per 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"] @@ -226,12 +241,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": ) if is_ray_serve and self.driver.gym_install is None: - raise ValueError( - f"Service '{service_name}' 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}." - ) + 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: diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 33df947f40..82a5d18e38 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -26,6 +26,7 @@ SubmitConfig, VllmServiceConfig, # used in _BUILDERS dispatch table effective_ray_serve, + gym_install_required_message, ) from nemo_gym.orchestration.executors.script_templates import ( bash_var, @@ -204,6 +205,14 @@ 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: @@ -240,8 +249,11 @@ def _build_vllm_ray_serve_command( ) 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. - return f"bash -lc '{fetch_and_run}'" + # 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}" ) @@ -255,8 +267,9 @@ def _build_vllm_ray_serve_command( # 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. - escaped = fetch_and_run.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`") - return render_vllm_ray_symmetric_run(f'bash -c "{escaped}"', total_nodes, resource_flags) + 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: @@ -284,12 +297,7 @@ def _build_service_command( ) -> str: if isinstance(service, VllmServiceConfig) and effective_ray_serve(service, total_nodes, gpus_per_node_values): if gym_install is None: - raise ValueError( - "Service 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}." - ) + 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) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 95d17933c9..4eef861bfa 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -19,7 +19,11 @@ import pytest from nemo_gym.orchestration.api import GymInstallConfig, SubmitConfig -from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd +from nemo_gym.orchestration.executors.script_templates import ( + render_driver_entrypoint, + render_gym_cmd, + render_repo_checkout, +) from nemo_gym.orchestration.executors.slurm_script import ( _build_service_command, _build_vllm_command, @@ -293,6 +297,36 @@ def test_build_vllm_ray_serve_command_installs_git_if_missing(vllm_service): assert cmd.index("command -v git") < cmd.index("git clone") +def test_build_vllm_ray_serve_command_single_node_model_with_space_survives_quoting(): + # Regression test for a real bug: shlex.quote() wraps a model name needing escaping (e.g. + # containing a space) in literal single quotes. Naively embedding that inside a Python-level + # single-quoted `bash -lc '...'` wrapper would let those literal `'` characters terminate the + # wrapper early, corrupting the command - everything from the space onward silently vanishes + # into inert extra positional args of the outer `bash -c` invocation instead of reaching the + # gateway. Runs the *actual* generated bash through a stand-in gateway to prove the model name + # survives intact, the same way test_..._multi_node_chain_survives_symmetric_run_entrypoint + # does for the multi-node path's `&&`-chain hazard. + service = VllmServiceConfig(type="vllm", container="vllm:latest", model="org/my model") + cmd = _build_vllm_ray_serve_command(service, total_nodes=1, gym_install=_GYM_INSTALL, gpus_per_node_values=[]) + + checkout = render_repo_checkout(_GYM_INSTALL.repo, _GYM_INSTALL.ref) + inner = ( + cmd.replace(checkout, "true") + .replace("pip install --quiet aiohttp", "true") + .replace("python3 nemo_gym/orchestration/ray_serve_gateway.py", "fake_gateway") + ) + # fake_gateway is defined and exported *before* the generated `bash -lc "..."` command, not + # spliced inline into it - inline injection of "$@" would itself get corrupted by the outer + # double-quoting the fix introduces (the same class of bug this test guards against). + script = 'fake_gateway() { for a in "$@"; do echo "ARG:$a"; done; }\nexport -f fake_gateway\n' + inner + "\n" + result = subprocess.run(["bash", "-c", script], capture_output=True, text=True, timeout=10) + + assert result.returncode == 0, result.stderr + args = [line.removeprefix("ARG:") for line in result.stdout.splitlines()] + assert "--model" in args, f"corrupted command, got args: {args}" + assert args[args.index("--model") + 1] == "org/my model" + + def test_build_vllm_ray_serve_command_multi_node_wraps_in_symmetric_run(vllm_service): cmd = _build_vllm_ray_serve_command( vllm_service, total_nodes=2, gym_install=_GYM_INSTALL, gpus_per_node_values=[8] @@ -451,6 +485,14 @@ def test_render_driver_entrypoint_with_gym_install(): assert '"${GYM_CMD[@]}"' in out +def test_render_driver_entrypoint_installs_git_if_missing(): + # The driver container (e.g. a minimal python image) may not bundle git any more than + # vllm/vllm-openai does - render_repo_checkout's guard protects both callers. + out = render_driver_entrypoint("https://github.com/NVIDIA-NeMo/gym", "main", None) + assert "command -v git >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq git)" in out + assert out.index("command -v git") < out.index("git clone") + + def test_render_driver_entrypoint_with_prepare(): out = render_driver_entrypoint(None, None, "gym eval prepare +foo=bar") assert "gym eval prepare +foo=bar" in out