From 4d88403ab25848e5ac6c462c01367b1be800a998 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 22 Jul 2026 14:03:07 +0200 Subject: [PATCH 01/66] feat(orchestration): gym eval submit command Signed-off-by: Tomasz Grzegorzek --- nemo_gym/cli/eval.py | 6 ++++++ nemo_gym/cli/main.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index 04ca1684f48..e65d360434d 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -467,3 +467,9 @@ def reward_profile(): # pragma: no cover Input rows: {completion_summary["total_input_rows"]} total; {completion_summary["complete_input_rows"]} complete; {completion_summary["partial_input_rows"]} partial; {completion_summary["missing_input_rows"]} without rollouts dropped from output. Reward profiling outputs: {reward_profiling_fpath} Agent-level metrics: {agent_level_metrics_fpath}""") + + +@exit_cleanly_on_config_error +def submit() -> None: # pragma: no cover + get_global_config_dict() + rich.print("[bold]gym eval submit[/bold]: not yet implemented.") diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 06d67ceb276..8b6863baeff 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -257,6 +257,7 @@ def _asset_selector(flag: str) -> Flag: ) + def _merge_config_paths(overrides: list[str]) -> list[str]: """Coalesce all `+config_paths=[...]` tokens (from --config and asset selectors) into one (Hydra rejects dupes).""" prefix = "+config_paths=[" @@ -536,6 +537,11 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: _value_flag("rollouts", "rollouts_jsonl_fpath", "Rollouts JSONL produced by collection."), ), ), + "eval submit": Command( + target="nemo_gym.cli.eval:submit", + summary="Submit a job.", + flags=(CONFIG,), + ), "dev test": Command(target="nemo_gym.cli.dev:dev_test", summary="Run NeMo Gym's unit tests."), } From 8f7ed77f600eaec751e2fed18c5a9d8e175c6fbb Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 10:59:59 +0200 Subject: [PATCH 02/66] move as much orchestration to a separate module Signed-off-by: Tomasz Grzegorzek --- nemo_gym/cli/eval.py | 7 ++----- nemo_gym/cli/main.py | 9 ++++++-- nemo_gym/orchestration/__init__.py | 0 nemo_gym/orchestration/api.py | 33 ++++++++++++++++++++++++++++++ nemo_gym/orchestration/submit.py | 33 ++++++++++++++++++++++++++++++ 5 files changed, 75 insertions(+), 7 deletions(-) create mode 100644 nemo_gym/orchestration/__init__.py create mode 100644 nemo_gym/orchestration/api.py create mode 100644 nemo_gym/orchestration/submit.py diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index e65d360434d..fdebef4ba13 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -17,7 +17,9 @@ import difflib import importlib import json +import sys from copy import deepcopy +from dataclasses import dataclass from glob import glob from multiprocessing import Pool from pathlib import Path @@ -468,8 +470,3 @@ def reward_profile(): # pragma: no cover Reward profiling outputs: {reward_profiling_fpath} Agent-level metrics: {agent_level_metrics_fpath}""") - -@exit_cleanly_on_config_error -def submit() -> None: # pragma: no cover - get_global_config_dict() - rich.print("[bold]gym eval submit[/bold]: not yet implemented.") diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 8b6863baeff..f865d0133a6 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -538,9 +538,14 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: ), ), "eval submit": Command( - target="nemo_gym.cli.eval:submit", + target="nemo_gym.orchestration.submit:submit", summary="Submit a job.", - flags=(CONFIG,), + flags=( + Flag( + register=lambda p: p.add_argument("--config", metavar="PATH", help="Config YAML file to load."), + translate_to_hydra=lambda args: [f"+_submit_config={args.config}"] if args.config else [], + ), + ), ), "dev test": Command(target="nemo_gym.cli.dev:dev_test", summary="Run NeMo Gym's unit tests."), } diff --git a/nemo_gym/orchestration/__init__.py b/nemo_gym/orchestration/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py new file mode 100644 index 00000000000..b1a6fedbd07 --- /dev/null +++ b/nemo_gym/orchestration/api.py @@ -0,0 +1,33 @@ +from typing import Annotated, Literal + +from pydantic import BaseModel, Discriminator, Tag + + +class VllmServiceConfig(BaseModel): + type: Literal["vllm"] + container: str + placement: str | None = None + + +class RayServiceConfig(BaseModel): + type: Literal["ray"] + container: str + placement: str | None = None + + +ServiceConfig = Annotated[ + Annotated[VllmServiceConfig, Tag("vllm")] | Annotated[RayServiceConfig, Tag("ray")], + Discriminator("type"), +] + + +class ComputeConfig(BaseModel): + type: str + hostname: str | None = None + walltime: str | None = None + node_pools: dict[str, dict] | None = None + + +class SubmitConfig(BaseModel): + services: dict[str, ServiceConfig] + compute: dict[str, ComputeConfig] diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py new file mode 100644 index 00000000000..ee251e25ec0 --- /dev/null +++ b/nemo_gym/orchestration/submit.py @@ -0,0 +1,33 @@ +import sys + +import rich +from omegaconf import OmegaConf + +from nemo_gym.orchestration.api import SubmitConfig + + +def _load_submit_config() -> SubmitConfig: + """Load the submit config YAML and apply +key=value overrides from the command line.""" + config_path: str | None = None + override_tokens: list[str] = [] + + for token in sys.argv[1:]: + stripped = token.lstrip("+") + if stripped.startswith("_submit_config="): + config_path = stripped[len("_submit_config="):] + else: + override_tokens.append(stripped) + + if config_path is None: + rich.print("[red]error:[/red] No config provided. Use [bold]--config [/bold] to specify a submit config YAML.") + sys.exit(1) + + base = OmegaConf.load(config_path) + overrides = OmegaConf.from_dotlist(override_tokens) if override_tokens else OmegaConf.create() + merged = OmegaConf.merge(base, overrides) + return SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)) + + +def submit() -> None: # pragma: no cover + _load_submit_config() + rich.print("[bold]gym eval submit[/bold]: not yet implemented.") From 748fd41622064db1c469b00584bde85e3532b9ec Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 11:06:23 +0200 Subject: [PATCH 03/66] base class for service type Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index b1a6fedbd07..974119f22f2 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -3,16 +3,17 @@ from pydantic import BaseModel, Discriminator, Tag -class VllmServiceConfig(BaseModel): - type: Literal["vllm"] +class BaseServiceConfig(BaseModel): container: str - placement: str | None = None + placement: str + +class VllmServiceConfig(BaseServiceConfig): + type: Literal["vllm"] -class RayServiceConfig(BaseModel): + +class RayServiceConfig(BaseServiceConfig): type: Literal["ray"] - container: str - placement: str | None = None ServiceConfig = Annotated[ From 864a876e7d0a1295400082e2f938a540a6b0cf39 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 11:12:00 +0200 Subject: [PATCH 04/66] mark experimental Signed-off-by: Tomasz Grzegorzek --- nemo_gym/decorators.py | 16 ++++++++++++++++ nemo_gym/orchestration/submit.py | 2 ++ 2 files changed, 18 insertions(+) create mode 100644 nemo_gym/decorators.py diff --git a/nemo_gym/decorators.py b/nemo_gym/decorators.py new file mode 100644 index 00000000000..e705ef18e2e --- /dev/null +++ b/nemo_gym/decorators.py @@ -0,0 +1,16 @@ +import functools + +import rich + + +def experimental(fn): + """Decorator that prints an experimental warning before the function runs.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + rich.print( + f"[yellow]Warning:[/yellow] [bold]{fn.__name__}[/bold] is experimental and may change or be removed without notice." + ) + return fn(*args, **kwargs) + + return wrapper diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index ee251e25ec0..8c269196051 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -3,6 +3,7 @@ import rich from omegaconf import OmegaConf +from nemo_gym.decorators import experimental from nemo_gym.orchestration.api import SubmitConfig @@ -28,6 +29,7 @@ def _load_submit_config() -> SubmitConfig: return SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)) +@experimental def submit() -> None: # pragma: no cover _load_submit_config() rich.print("[bold]gym eval submit[/bold]: not yet implemented.") From 356881c1e84b2e5a4d0954caf909ede228a21c23 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 11:21:26 +0200 Subject: [PATCH 05/66] validate placements Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 974119f22f2..9dab2fb99c6 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -1,11 +1,11 @@ from typing import Annotated, Literal -from pydantic import BaseModel, Discriminator, Tag +from pydantic import BaseModel, Discriminator, Tag, model_validator class BaseServiceConfig(BaseModel): container: str - placement: str + placement: str | None = None class VllmServiceConfig(BaseServiceConfig): @@ -32,3 +32,24 @@ class ComputeConfig(BaseModel): class SubmitConfig(BaseModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] + + @model_validator(mode="after") + def _resolve_and_validate_placements(self) -> "SubmitConfig": + compute_names = set(self.compute) + sole_compute = next(iter(compute_names)) if len(compute_names) == 1 else None + + for service_name, service in self.services.items(): + if service.placement is None: + if sole_compute is None: + raise ValueError( + f"Service '{service_name}' has no placement and there are multiple compute resources " + f"({', '.join(sorted(compute_names))}). Set placement explicitly." + ) + service.placement = sole_compute + elif service.placement not in compute_names: + raise ValueError( + f"Service '{service_name}' placement '{service.placement}' does not match any compute resource " + f"({', '.join(sorted(compute_names))})." + ) + + return self From 743ce418194a7801718fbfbc4ade11cc076505ea Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 11:32:25 +0200 Subject: [PATCH 06/66] only single compute with tests Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 13 +++--- tests/unit_tests/test_orchestration_api.py | 51 ++++++++++++++++++++++ 2 files changed, 58 insertions(+), 6 deletions(-) create mode 100644 tests/unit_tests/test_orchestration_api.py diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 9dab2fb99c6..96a34c86890 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -36,15 +36,16 @@ class SubmitConfig(BaseModel): @model_validator(mode="after") def _resolve_and_validate_placements(self) -> "SubmitConfig": compute_names = set(self.compute) - sole_compute = next(iter(compute_names)) if len(compute_names) == 1 else None + + if len(compute_names) > 1: + raise ValueError( + f"Multiple compute resources are not supported yet ({', '.join(sorted(compute_names))})." + ) + + sole_compute = next(iter(compute_names)) for service_name, service in self.services.items(): if service.placement is None: - if sole_compute is None: - raise ValueError( - f"Service '{service_name}' has no placement and there are multiple compute resources " - f"({', '.join(sorted(compute_names))}). Set placement explicitly." - ) service.placement = sole_compute elif service.placement not in compute_names: raise ValueError( diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py new file mode 100644 index 00000000000..ccb3bae1028 --- /dev/null +++ b/tests/unit_tests/test_orchestration_api.py @@ -0,0 +1,51 @@ +# 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 pydantic import ValidationError + +from nemo_gym.orchestration.api import SubmitConfig + +COMPUTE = {"cluster": {"type": "slurm", "hostname": "foo"}} +COMPUTE_TWO = { + "cluster_a": {"type": "slurm", "hostname": "foo"}, + "cluster_b": {"type": "slurm", "hostname": "bar"}, +} + +SERVICE = {"container": "gym:latest", "type": "vllm"} + + +def test_implicit_placement_single_compute(): + config = SubmitConfig.model_validate({"services": {"svc": SERVICE}, "compute": COMPUTE}) + assert config.services["svc"].placement == "cluster" + + +def test_explicit_valid_placement(): + config = SubmitConfig.model_validate( + {"services": {"svc": {**SERVICE, "placement": "cluster"}}, "compute": COMPUTE} + ) + assert config.services["svc"].placement == "cluster" + + +def test_multiple_compute_raises(): + with pytest.raises(ValidationError, match="Multiple compute resources are not supported yet"): + SubmitConfig.model_validate({"services": {"svc": SERVICE}, "compute": COMPUTE_TWO}) + + +def test_invalid_placement_raises(): + with pytest.raises(ValidationError, match="does not match any compute resource"): + SubmitConfig.model_validate( + {"services": {"svc": {**SERVICE, "placement": "nonexistent"}}, "compute": COMPUTE} + ) From b243b9ecd0465bf85f10364356981420f32b2446 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 13:56:04 +0200 Subject: [PATCH 07/66] builder draft; update services Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 22 ++++++++ nemo_gym/orchestration/executors/__init__.py | 0 nemo_gym/orchestration/executors/base.py | 8 +++ nemo_gym/orchestration/executors/slurm.py | 59 ++++++++++++++++++++ nemo_gym/orchestration/submit.py | 14 ++++- tests/unit_tests/test_orchestration_api.py | 19 ++++--- 6 files changed, 111 insertions(+), 11 deletions(-) create mode 100644 nemo_gym/orchestration/executors/__init__.py create mode 100644 nemo_gym/orchestration/executors/base.py create mode 100644 nemo_gym/orchestration/executors/slurm.py diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 96a34c86890..6fe3850bc91 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -3,13 +3,25 @@ from pydantic import BaseModel, Discriminator, Tag, model_validator +class HealthCheckConfig(BaseModel): + path: str = "/health" + port: int + timeout_seconds: int = 60 + + class BaseServiceConfig(BaseModel): container: str placement: str | None = None + health_check: HealthCheckConfig | None = None class VllmServiceConfig(BaseServiceConfig): type: Literal["vllm"] + model: str + port: int = 8000 + tensor_parallel_size: int = 1 + pipeline_parallel_size: int = 1 + trust_remote_code: bool = False class RayServiceConfig(BaseServiceConfig): @@ -29,9 +41,19 @@ class ComputeConfig(BaseModel): node_pools: dict[str, dict] | None = None +class BenchmarkRunConfig(BaseModel): + name: str + + +class DriverConfig(BaseModel): + container: str + benchmarks: list[BenchmarkRunConfig] + + class SubmitConfig(BaseModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] + driver: DriverConfig @model_validator(mode="after") def _resolve_and_validate_placements(self) -> "SubmitConfig": diff --git a/nemo_gym/orchestration/executors/__init__.py b/nemo_gym/orchestration/executors/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/nemo_gym/orchestration/executors/base.py b/nemo_gym/orchestration/executors/base.py new file mode 100644 index 00000000000..f77f20807e3 --- /dev/null +++ b/nemo_gym/orchestration/executors/base.py @@ -0,0 +1,8 @@ +from abc import ABC, abstractmethod + +from nemo_gym.orchestration.api import SubmitConfig + + +class BaseExecutor(ABC): + @abstractmethod + def run(self, config: SubmitConfig) -> None: ... diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py new file mode 100644 index 00000000000..15d019901c6 --- /dev/null +++ b/nemo_gym/orchestration/executors/slurm.py @@ -0,0 +1,59 @@ +from nemo_gym.orchestration.api import ( + BenchmarkRunConfig, + ComputeConfig, + RayServiceConfig, + SubmitConfig, + VllmServiceConfig, +) +from nemo_gym.orchestration.executors.base import BaseExecutor + + +def _build_vllm_command(service: VllmServiceConfig) -> str: + cmd = f"vllm serve {service.model} --port {service.port} --tensor-parallel-size {service.tensor_parallel_size}" + if service.trust_remote_code: + cmd += " --trust-remote-code" + return cmd + + +def _build_ray_command(service: RayServiceConfig) -> str: + return "ray start --head" + + +_BUILDERS = { + VllmServiceConfig: _build_vllm_command, + RayServiceConfig: _build_ray_command, +} + + +class SlurmExecutor(BaseExecutor): + def run(self, config: SubmitConfig) -> None: + compute = next(iter(config.compute.values())) + for benchmark in config.driver.benchmarks: + script = self._build_job_script(config, benchmark, compute) + self._sbatch(script, compute) + + def _build_job_script( + self, config: SubmitConfig, benchmark: BenchmarkRunConfig, compute: ComputeConfig + ) -> str: + lines = ["#!/bin/bash"] + + if compute.walltime: + lines.append(f"#SBATCH --time={compute.walltime}") + + for name, service in config.services.items(): + builder = _BUILDERS[type(service)] + lines.append(f"# service: {name}") + lines.append(builder(service) + " &") + + for name, service in config.services.items(): + if service.health_check: + hc = service.health_check + lines.append( + f"gym env wait --host localhost --port {hc.port} --path {hc.path} --timeout {hc.timeout_seconds}" + ) + + lines.append(f"gym eval run --benchmark {benchmark.name}") + return "\n".join(lines) + + def _sbatch(self, script: str, compute: ComputeConfig) -> None: + raise NotImplementedError # pragma: no cover diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index 8c269196051..0e38cce790c 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -5,6 +5,11 @@ from nemo_gym.decorators import experimental from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.executors.slurm import SlurmExecutor + +_EXECUTORS = { + "slurm": SlurmExecutor, +} def _load_submit_config() -> SubmitConfig: @@ -31,5 +36,10 @@ def _load_submit_config() -> SubmitConfig: @experimental def submit() -> None: # pragma: no cover - _load_submit_config() - rich.print("[bold]gym eval submit[/bold]: not yet implemented.") + config = _load_submit_config() + compute = next(iter(config.compute.values())) + executor_cls = _EXECUTORS.get(compute.type) + if executor_cls is None: + rich.print(f"[red]error:[/red] Unsupported compute type '{compute.type}'. Supported: {', '.join(_EXECUTORS)}.") + return + executor_cls().run(config) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index ccb3bae1028..33ba1d109ee 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -24,28 +24,29 @@ "cluster_b": {"type": "slurm", "hostname": "bar"}, } -SERVICE = {"container": "gym:latest", "type": "vllm"} +SERVICE = {"container": "gym:latest", "type": "vllm", "model": "org/model"} +DRIVER = {"container": "gym:latest", "benchmarks": [{"name": "gsm8k"}]} + + +def _config(**overrides): + return {"services": {"svc": SERVICE}, "compute": COMPUTE, "driver": DRIVER, **overrides} def test_implicit_placement_single_compute(): - config = SubmitConfig.model_validate({"services": {"svc": SERVICE}, "compute": COMPUTE}) + config = SubmitConfig.model_validate(_config()) assert config.services["svc"].placement == "cluster" def test_explicit_valid_placement(): - config = SubmitConfig.model_validate( - {"services": {"svc": {**SERVICE, "placement": "cluster"}}, "compute": COMPUTE} - ) + config = SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "placement": "cluster"}})) assert config.services["svc"].placement == "cluster" def test_multiple_compute_raises(): with pytest.raises(ValidationError, match="Multiple compute resources are not supported yet"): - SubmitConfig.model_validate({"services": {"svc": SERVICE}, "compute": COMPUTE_TWO}) + SubmitConfig.model_validate(_config(compute=COMPUTE_TWO)) def test_invalid_placement_raises(): with pytest.raises(ValidationError, match="does not match any compute resource"): - SubmitConfig.model_validate( - {"services": {"svc": {**SERVICE, "placement": "nonexistent"}}, "compute": COMPUTE} - ) + SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "placement": "nonexistent"}})) From 7e72d6f5abec1b6e4b3d23f5ae61cfc7742f20db Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 14:12:47 +0200 Subject: [PATCH 08/66] make ComputeConfig a desscriminated union Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 14 ++++++++++++-- nemo_gym/orchestration/executors/slurm.py | 8 ++++---- nemo_gym/orchestration/submit.py | 10 +++------- 3 files changed, 19 insertions(+), 13 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 6fe3850bc91..f0809f0fc25 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -34,13 +34,23 @@ class RayServiceConfig(BaseServiceConfig): ] -class ComputeConfig(BaseModel): - type: str +class BaseComputeConfig(BaseModel): + pass + + +class SlurmComputeConfig(BaseComputeConfig): + type: Literal["slurm"] hostname: str | None = None walltime: str | None = None node_pools: dict[str, dict] | None = None +ComputeConfig = Annotated[ + Annotated[SlurmComputeConfig, Tag("slurm")], + Discriminator("type"), +] + + class BenchmarkRunConfig(BaseModel): name: str diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 15d019901c6..cca3f0ca276 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -1,7 +1,7 @@ from nemo_gym.orchestration.api import ( BenchmarkRunConfig, - ComputeConfig, RayServiceConfig, + SlurmComputeConfig, SubmitConfig, VllmServiceConfig, ) @@ -15,7 +15,7 @@ def _build_vllm_command(service: VllmServiceConfig) -> str: return cmd -def _build_ray_command(service: RayServiceConfig) -> str: +def _build_ray_command(_service: RayServiceConfig) -> str: return "ray start --head" @@ -33,7 +33,7 @@ def run(self, config: SubmitConfig) -> None: self._sbatch(script, compute) def _build_job_script( - self, config: SubmitConfig, benchmark: BenchmarkRunConfig, compute: ComputeConfig + self, config: SubmitConfig, benchmark: BenchmarkRunConfig, compute: SlurmComputeConfig ) -> str: lines = ["#!/bin/bash"] @@ -55,5 +55,5 @@ def _build_job_script( lines.append(f"gym eval run --benchmark {benchmark.name}") return "\n".join(lines) - def _sbatch(self, script: str, compute: ComputeConfig) -> None: + def _sbatch(self, script: str, compute: SlurmComputeConfig) -> None: raise NotImplementedError # pragma: no cover diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index 0e38cce790c..24a77aab300 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -4,11 +4,11 @@ from omegaconf import OmegaConf from nemo_gym.decorators import experimental -from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.slurm import SlurmExecutor _EXECUTORS = { - "slurm": SlurmExecutor, + SlurmComputeConfig: SlurmExecutor, } @@ -38,8 +38,4 @@ def _load_submit_config() -> SubmitConfig: def submit() -> None: # pragma: no cover config = _load_submit_config() compute = next(iter(config.compute.values())) - executor_cls = _EXECUTORS.get(compute.type) - if executor_cls is None: - rich.print(f"[red]error:[/red] Unsupported compute type '{compute.type}'. Supported: {', '.join(_EXECUTORS)}.") - return - executor_cls().run(config) + _EXECUTORS[type(compute)]().run(config) From 65255333bfef3305af5dbf83c84354a4f5c4e1ee Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 16:37:55 +0200 Subject: [PATCH 09/66] connect to the host; replace gym env wait placeholder with a simple curl Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 5 + .../orchestration/executors/connection.py | 98 +++++++++++++++++++ nemo_gym/orchestration/executors/slurm.py | 33 +++++-- tests/unit_tests/test_orchestration_api.py | 3 +- 4 files changed, 131 insertions(+), 8 deletions(-) create mode 100644 nemo_gym/orchestration/executors/connection.py diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index f0809f0fc25..778336464f8 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -60,10 +60,15 @@ class DriverConfig(BaseModel): benchmarks: list[BenchmarkRunConfig] +class JobConfig(BaseModel): + output_path: str + + class SubmitConfig(BaseModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] driver: DriverConfig + job: JobConfig @model_validator(mode="after") def _resolve_and_validate_placements(self) -> "SubmitConfig": diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py new file mode 100644 index 00000000000..7958369a371 --- /dev/null +++ b/nemo_gym/orchestration/executors/connection.py @@ -0,0 +1,98 @@ +import shlex +import shutil +import socket +import subprocess +import tempfile +from abc import ABC, abstractmethod +from pathlib import Path + + +class Connection(ABC): + def __enter__(self) -> "Connection": + return self + + def __exit__(self, *_) -> None: + self.close() + + @abstractmethod + def copy(self, local: Path, remote: Path) -> None: ... + + @abstractmethod + def run(self, commands: list[str]) -> None: ... + + def close(self) -> None: + pass + + +class LocalConnection(Connection): + def copy(self, local: Path, remote: Path) -> None: + if remote.exists(): + shutil.rmtree(remote) + shutil.copytree(local, remote) + + def run(self, commands: list[str]) -> None: + for cmd in commands: + subprocess.run(shlex.split(cmd), check=True) + + +class SSHConnection(Connection): + """Single SSH master connection; copy and all commands reuse the same socket.""" + + _OPTS = ["-o", "StrictHostKeyChecking=accept-new", "-o", "BatchMode=yes"] + + def __init__(self, hostname: str) -> None: + self._hostname = hostname + self._socket = Path(tempfile.mktemp(prefix="gym-ssh-", suffix=".sock")) + self._master: subprocess.Popen | None = None + + def __enter__(self) -> "SSHConnection": + self._master = subprocess.Popen( + [ + "ssh", + *self._OPTS, + "-o", "ControlMaster=yes", + "-o", f"ControlPath={self._socket}", + "-o", "ControlPersist=yes", + "-N", + self._hostname, + ] + ) + return self + + def _ssh_opts(self) -> list[str]: + return [*self._OPTS, "-o", "ControlMaster=no", "-o", f"ControlPath={self._socket}"] + + def copy(self, local: Path, remote: Path) -> None: + subprocess.run( + [ + "rsync", "-az", "--delete", + "-e", f"ssh {' '.join(self._ssh_opts())}", + f"{local}/", + f"{self._hostname}:{remote}", + ], + check=True, + ) + + def run(self, commands: list[str]) -> None: + script = "\n".join(commands) + subprocess.run( + ["ssh", *self._ssh_opts(), self._hostname, "bash", "-s"], + input=script, + text=True, + check=True, + ) + + def close(self) -> None: + subprocess.run( + ["ssh", *self._ssh_opts(), "-O", "exit", self._hostname], + capture_output=True, + ) + self._socket.unlink(missing_ok=True) + if self._master: + self._master.wait() + + +def get_connection(hostname: str | None) -> Connection: + if hostname is None or hostname == socket.gethostname(): + return LocalConnection() + return SSHConnection(hostname) diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index cca3f0ca276..63677c990fe 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -1,3 +1,7 @@ +import shlex +import tempfile +from pathlib import Path + from nemo_gym.orchestration.api import ( BenchmarkRunConfig, RayServiceConfig, @@ -6,10 +10,11 @@ VllmServiceConfig, ) from nemo_gym.orchestration.executors.base import BaseExecutor +from nemo_gym.orchestration.executors.connection import get_connection def _build_vllm_command(service: VllmServiceConfig) -> str: - cmd = f"vllm serve {service.model} --port {service.port} --tensor-parallel-size {service.tensor_parallel_size}" + cmd = f"vllm serve {shlex.quote(service.model)} --port {service.port} --tensor-parallel-size {service.tensor_parallel_size}" if service.trust_remote_code: cmd += " --trust-remote-code" return cmd @@ -28,9 +33,25 @@ def _build_ray_command(_service: RayServiceConfig) -> str: class SlurmExecutor(BaseExecutor): def run(self, config: SubmitConfig) -> None: compute = next(iter(config.compute.values())) + output_path = Path(config.job.output_path) + + staging = self._stage(config, compute) + with get_connection(compute.hostname) as conn: + conn.copy(staging, output_path) + conn.run([ + f"sbatch {shlex.quote(str(output_path / b.name / 'job.sh'))}" + for b in config.driver.benchmarks + ]) + + def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig) -> Path: + staging = Path(tempfile.mkdtemp(prefix="gym-submit-")) for benchmark in config.driver.benchmarks: + bench_dir = staging / benchmark.name + bench_dir.mkdir() + (bench_dir / "logs").mkdir() script = self._build_job_script(config, benchmark, compute) - self._sbatch(script, compute) + (bench_dir / "job.sh").write_text(script) + return staging def _build_job_script( self, config: SubmitConfig, benchmark: BenchmarkRunConfig, compute: SlurmComputeConfig @@ -49,11 +70,9 @@ def _build_job_script( if service.health_check: hc = service.health_check lines.append( - f"gym env wait --host localhost --port {hc.port} --path {hc.path} --timeout {hc.timeout_seconds}" + f"until curl -sf http://localhost:{hc.port}{hc.path}; do sleep 2; done &" + f" # health check: {name} (timeout {hc.timeout_seconds}s)" ) - lines.append(f"gym eval run --benchmark {benchmark.name}") + lines.append(f"gym eval run --benchmark {shlex.quote(benchmark.name)}") return "\n".join(lines) - - def _sbatch(self, script: str, compute: SlurmComputeConfig) -> None: - raise NotImplementedError # pragma: no cover diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 33ba1d109ee..380c0b056b1 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -26,10 +26,11 @@ SERVICE = {"container": "gym:latest", "type": "vllm", "model": "org/model"} DRIVER = {"container": "gym:latest", "benchmarks": [{"name": "gsm8k"}]} +JOB = {"output_path": "/tmp/gym-jobs"} def _config(**overrides): - return {"services": {"svc": SERVICE}, "compute": COMPUTE, "driver": DRIVER, **overrides} + return {"services": {"svc": SERVICE}, "compute": COMPUTE, "driver": DRIVER, "job": JOB, **overrides} def test_implicit_placement_single_compute(): From 4c56c884a25b03912d431eb4225903f64ca22ba8 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 16:51:17 +0200 Subject: [PATCH 10/66] make SlurmExecutor explicitly pyxis based Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/executors/slurm.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 63677c990fe..9dc593bdec6 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -31,6 +31,13 @@ def _build_ray_command(_service: RayServiceConfig) -> str: class SlurmExecutor(BaseExecutor): + """Slurm executor for Pyxis-enabled clusters (https://github.com/NVIDIA/pyxis). + + Every service and the driver are launched via `srun --container-image` so they + run inside the container specified in their config. Health checks run as plain + bash inside the sbatch script (no container needed — they just poll HTTP). + """ + def run(self, config: SubmitConfig) -> None: compute = next(iter(config.compute.values())) output_path = Path(config.job.output_path) @@ -64,15 +71,22 @@ def _build_job_script( for name, service in config.services.items(): builder = _BUILDERS[type(service)] lines.append(f"# service: {name}") - lines.append(builder(service) + " &") + lines.append( + f"srun --container-image={shlex.quote(service.container)} " + f"{builder(service)} &" + ) for name, service in config.services.items(): if service.health_check: hc = service.health_check lines.append( - f"until curl -sf http://localhost:{hc.port}{hc.path}; do sleep 2; done &" - f" # health check: {name} (timeout {hc.timeout_seconds}s)" + f"# health check: {name} (timeout {hc.timeout_seconds}s)\n" + f"timeout {hc.timeout_seconds} bash -c " + f"'until curl -sf http://localhost:{hc.port}{shlex.quote(hc.path)}; do sleep 2; done'" ) - lines.append(f"gym eval run --benchmark {shlex.quote(benchmark.name)}") + lines.append( + f"srun --container-image={shlex.quote(config.driver.container)} " + f"gym eval run --benchmark {shlex.quote(benchmark.name)}" + ) return "\n".join(lines) From f1d968385cb89c2a0d63e6a30ca8f9fcbd45a38d Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Tue, 28 Jul 2026 18:17:21 +0200 Subject: [PATCH 11/66] timestamp rundir and redirect logs Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/executors/slurm.py | 28 +++++++++++++++-------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 9dc593bdec6..30db8e6fa57 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -1,5 +1,6 @@ import shlex import tempfile +from datetime import datetime from pathlib import Path from nemo_gym.orchestration.api import ( @@ -40,40 +41,48 @@ class SlurmExecutor(BaseExecutor): def run(self, config: SubmitConfig) -> None: compute = next(iter(config.compute.values())) - output_path = Path(config.job.output_path) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + remote_run_dir = Path(config.job.output_path) / timestamp - staging = self._stage(config, compute) + staging = self._stage(config, compute, remote_run_dir) with get_connection(compute.hostname) as conn: - conn.copy(staging, output_path) + conn.copy(staging, remote_run_dir) conn.run([ - f"sbatch {shlex.quote(str(output_path / b.name / 'job.sh'))}" + f"sbatch {shlex.quote(str(remote_run_dir / b.name / 'job.sh'))}" for b in config.driver.benchmarks ]) - def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig) -> Path: + def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> Path: staging = Path(tempfile.mkdtemp(prefix="gym-submit-")) for benchmark in config.driver.benchmarks: bench_dir = staging / benchmark.name bench_dir.mkdir() (bench_dir / "logs").mkdir() - script = self._build_job_script(config, benchmark, compute) + remote_bench_dir = remote_run_dir / benchmark.name + script = self._build_job_script(config, benchmark, compute, remote_bench_dir) (bench_dir / "job.sh").write_text(script) return staging def _build_job_script( - self, config: SubmitConfig, benchmark: BenchmarkRunConfig, compute: SlurmComputeConfig + self, + config: SubmitConfig, + benchmark: BenchmarkRunConfig, + compute: SlurmComputeConfig, + remote_bench_dir: Path, ) -> str: lines = ["#!/bin/bash"] if compute.walltime: lines.append(f"#SBATCH --time={compute.walltime}") + lines.append(f"#SBATCH --chdir={remote_bench_dir}") for name, service in config.services.items(): builder = _BUILDERS[type(service)] lines.append(f"# service: {name}") lines.append( f"srun --container-image={shlex.quote(service.container)} " - f"{builder(service)} &" + f"{builder(service)} " + f"> logs/{name}.log 2>&1 &" ) for name, service in config.services.items(): @@ -87,6 +96,7 @@ def _build_job_script( lines.append( f"srun --container-image={shlex.quote(config.driver.container)} " - f"gym eval run --benchmark {shlex.quote(benchmark.name)}" + f"gym eval run --benchmark {shlex.quote(benchmark.name)} " + f"> logs/driver.log 2>&1" ) return "\n".join(lines) From a247be9ad59eed33e5daea4fb26782e7933a0e34 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 00:56:12 +0200 Subject: [PATCH 12/66] str-format sbatch script Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 11 +- nemo_gym/orchestration/executors/slurm.py | 67 +---------- .../orchestration/executors/slurm_script.py | 111 ++++++++++++++++++ 3 files changed, 124 insertions(+), 65 deletions(-) create mode 100644 nemo_gym/orchestration/executors/slurm_script.py diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 778336464f8..5789c0ed62d 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -34,6 +34,14 @@ class RayServiceConfig(BaseServiceConfig): ] +class NodePool(BaseModel): + partition: str + nodes: int = 1 + ntasks_per_node: int = 1 + gpus_per_node: int | None = None + extra_args: dict[str, str] = {} + + class BaseComputeConfig(BaseModel): pass @@ -42,7 +50,8 @@ class SlurmComputeConfig(BaseComputeConfig): type: Literal["slurm"] hostname: str | None = None walltime: str | None = None - node_pools: dict[str, dict] | None = None + node_pools: dict[str, NodePool] = {} + extra_args: dict[str, str] = {} ComputeConfig = Annotated[ diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 30db8e6fa57..409403bc665 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -3,32 +3,10 @@ from datetime import datetime from pathlib import Path -from nemo_gym.orchestration.api import ( - BenchmarkRunConfig, - RayServiceConfig, - SlurmComputeConfig, - SubmitConfig, - VllmServiceConfig, -) +from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.base import BaseExecutor from nemo_gym.orchestration.executors.connection import get_connection - - -def _build_vllm_command(service: VllmServiceConfig) -> str: - cmd = f"vllm serve {shlex.quote(service.model)} --port {service.port} --tensor-parallel-size {service.tensor_parallel_size}" - if service.trust_remote_code: - cmd += " --trust-remote-code" - return cmd - - -def _build_ray_command(_service: RayServiceConfig) -> str: - return "ray start --head" - - -_BUILDERS = { - VllmServiceConfig: _build_vllm_command, - RayServiceConfig: _build_ray_command, -} +from nemo_gym.orchestration.executors.slurm_script import build_sbatch_script class SlurmExecutor(BaseExecutor): @@ -58,45 +36,6 @@ def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_d bench_dir = staging / benchmark.name bench_dir.mkdir() (bench_dir / "logs").mkdir() - remote_bench_dir = remote_run_dir / benchmark.name - script = self._build_job_script(config, benchmark, compute, remote_bench_dir) + script = build_sbatch_script(config, benchmark, compute, remote_run_dir / benchmark.name) (bench_dir / "job.sh").write_text(script) return staging - - def _build_job_script( - self, - config: SubmitConfig, - benchmark: BenchmarkRunConfig, - compute: SlurmComputeConfig, - remote_bench_dir: Path, - ) -> str: - lines = ["#!/bin/bash"] - - if compute.walltime: - lines.append(f"#SBATCH --time={compute.walltime}") - lines.append(f"#SBATCH --chdir={remote_bench_dir}") - - for name, service in config.services.items(): - builder = _BUILDERS[type(service)] - lines.append(f"# service: {name}") - lines.append( - f"srun --container-image={shlex.quote(service.container)} " - f"{builder(service)} " - f"> logs/{name}.log 2>&1 &" - ) - - for name, service in config.services.items(): - if service.health_check: - hc = service.health_check - lines.append( - f"# health check: {name} (timeout {hc.timeout_seconds}s)\n" - f"timeout {hc.timeout_seconds} bash -c " - f"'until curl -sf http://localhost:{hc.port}{shlex.quote(hc.path)}; do sleep 2; done'" - ) - - lines.append( - f"srun --container-image={shlex.quote(config.driver.container)} " - f"gym eval run --benchmark {shlex.quote(benchmark.name)} " - f"> logs/driver.log 2>&1" - ) - return "\n".join(lines) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py new file mode 100644 index 00000000000..9b55c3055a5 --- /dev/null +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -0,0 +1,111 @@ +import shlex +from pathlib import Path + +from nemo_gym.orchestration.api import ( + BenchmarkRunConfig, + NodePool, + RayServiceConfig, + SlurmComputeConfig, + SubmitConfig, + VllmServiceConfig, +) + +_SCRIPT_TEMPLATE = """\ +#!/bin/bash +{directives} + +{service_commands} + +{health_checks} + +{driver_command} +""" + + +def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path) -> str: + lines = [] + if compute.walltime: + lines.append(f"#SBATCH --time={compute.walltime}") + lines.append(f"#SBATCH --chdir={remote_bench_dir}") + for key, val in compute.extra_args.items(): + lines.append(f"#SBATCH --{key}={val}") + for pool_name, pool in compute.node_pools.items(): + lines.extend(_render_pool_directives(pool_name, pool)) + return "\n".join(lines) + + +def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: + lines = [ + f"#SBATCH --partition={pool.partition} # pool: {pool_name}", + f"#SBATCH --nodes={pool.nodes}", + f"#SBATCH --ntasks-per-node={pool.ntasks_per_node}", + ] + if pool.gpus_per_node is not None: + lines.append(f"#SBATCH --gpus-per-node={pool.gpus_per_node}") + for key, val in pool.extra_args.items(): + lines.append(f"#SBATCH --{key}={val}") + return lines + + +def _render_service_command(name: str, container: str, command: str) -> str: + return ( + f"# service: {name}\n" + f"srun --container-image={shlex.quote(container)} {command} > logs/{name}.log 2>&1 &" + ) + + +def _render_health_check(name: str, port: int, path: str, timeout: int) -> str: + return ( + f"# health check: {name} (timeout {timeout}s)\n" + f"timeout {timeout} bash -c 'until curl -sf http://localhost:{port}{shlex.quote(path)}; do sleep 2; done'" + ) + + +def _build_vllm_command(service: VllmServiceConfig) -> str: + cmd = f"vllm serve {shlex.quote(service.model)} --port {service.port} --tensor-parallel-size {service.tensor_parallel_size}" + if service.trust_remote_code: + cmd += " --trust-remote-code" + return cmd + + +def _build_ray_command(_service: RayServiceConfig) -> str: + return "ray start --head" + + +_BUILDERS = { + VllmServiceConfig: _build_vllm_command, + RayServiceConfig: _build_ray_command, +} + + +def build_sbatch_script( + config: SubmitConfig, + benchmark: BenchmarkRunConfig, + compute: SlurmComputeConfig, + remote_bench_dir: Path, +) -> str: + directives = _render_directives(compute, remote_bench_dir) + + service_commands = "\n\n".join( + _render_service_command(name, service.container, _BUILDERS[type(service)](service)) + for name, service in config.services.items() + ) + + health_checks = "\n\n".join( + _render_health_check(name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds) + for name, service in config.services.items() + if service.health_check + ) + + driver_command = ( + f"srun --container-image={shlex.quote(config.driver.container)} " + f"gym eval run --benchmark {shlex.quote(benchmark.name)} " + f"> logs/driver.log 2>&1" + ) + + return _SCRIPT_TEMPLATE.format( + directives=directives, + service_commands=service_commands, + health_checks=health_checks, + driver_command=driver_command, + ) From a7914f8e75c167632d7c778a1719b7cf2081dfb4 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 01:09:15 +0200 Subject: [PATCH 13/66] introduce a dry-run and improve socket handling Signed-off-by: Tomasz Grzegorzek --- nemo_gym/cli/main.py | 4 ++ nemo_gym/orchestration/executors/base.py | 2 +- .../orchestration/executors/connection.py | 43 ++++++++++++++++--- nemo_gym/orchestration/executors/slurm.py | 15 ++++++- nemo_gym/orchestration/submit.py | 11 +++-- 5 files changed, 62 insertions(+), 13 deletions(-) diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index f865d0133a6..a1b7261686e 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -545,6 +545,10 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: register=lambda p: p.add_argument("--config", metavar="PATH", help="Config YAML file to load."), translate_to_hydra=lambda args: [f"+_submit_config={args.config}"] if args.config else [], ), + Flag( + register=lambda p: p.add_argument("--dry-run", action="store_true", help="Print generated job scripts without submitting."), + translate_to_hydra=lambda args: ["+_dry_run=true"] if args.dry_run else [], + ), ), ), "dev test": Command(target="nemo_gym.cli.dev:dev_test", summary="Run NeMo Gym's unit tests."), diff --git a/nemo_gym/orchestration/executors/base.py b/nemo_gym/orchestration/executors/base.py index f77f20807e3..0ec4a44c672 100644 --- a/nemo_gym/orchestration/executors/base.py +++ b/nemo_gym/orchestration/executors/base.py @@ -5,4 +5,4 @@ class BaseExecutor(ABC): @abstractmethod - def run(self, config: SubmitConfig) -> None: ... + def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: ... diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py index 7958369a371..cbbeae2b5b8 100644 --- a/nemo_gym/orchestration/executors/connection.py +++ b/nemo_gym/orchestration/executors/connection.py @@ -3,6 +3,7 @@ import socket import subprocess import tempfile +import time from abc import ABC, abstractmethod from pathlib import Path @@ -32,13 +33,14 @@ def copy(self, local: Path, remote: Path) -> None: def run(self, commands: list[str]) -> None: for cmd in commands: - subprocess.run(shlex.split(cmd), check=True) + _checked(shlex.split(cmd)) class SSHConnection(Connection): """Single SSH master connection; copy and all commands reuse the same socket.""" _OPTS = ["-o", "StrictHostKeyChecking=accept-new", "-o", "BatchMode=yes"] + _MASTER_TIMEOUT = 10 # seconds to wait for the control socket to appear def __init__(self, hostname: str) -> None: self._hostname = hostname @@ -55,31 +57,47 @@ def __enter__(self) -> "SSHConnection": "-o", "ControlPersist=yes", "-N", self._hostname, - ] + ], + stderr=subprocess.PIPE, ) + self._wait_for_socket() return self + def _wait_for_socket(self) -> None: + deadline = time.monotonic() + self._MASTER_TIMEOUT + while time.monotonic() < deadline: + if self._socket.exists(): + return + if self._master and self._master.poll() is not None: + stderr = (self._master.stderr.read().decode(errors="replace") if self._master.stderr else "") + raise RuntimeError( + f"SSH master to '{self._hostname}' exited with code {self._master.returncode}.\n{stderr.strip()}" + ) + time.sleep(0.2) + raise RuntimeError( + f"SSH master to '{self._hostname}' did not produce a control socket within {self._MASTER_TIMEOUT}s." + ) + def _ssh_opts(self) -> list[str]: return [*self._OPTS, "-o", "ControlMaster=no", "-o", f"ControlPath={self._socket}"] def copy(self, local: Path, remote: Path) -> None: - subprocess.run( + _checked( [ "rsync", "-az", "--delete", "-e", f"ssh {' '.join(self._ssh_opts())}", f"{local}/", f"{self._hostname}:{remote}", ], - check=True, + context=f"rsync to {self._hostname}:{remote}", ) def run(self, commands: list[str]) -> None: script = "\n".join(commands) - subprocess.run( + _checked( ["ssh", *self._ssh_opts(), self._hostname, "bash", "-s"], input=script, - text=True, - check=True, + context=f"ssh commands on {self._hostname}", ) def close(self) -> None: @@ -92,6 +110,17 @@ def close(self) -> None: self._master.wait() +def _checked(cmd: list[str], *, input: str | None = None, context: str = "") -> None: + result = subprocess.run(cmd, input=input, text=True, capture_output=True) + if result.returncode != 0: + label = f" ({context})" if context else "" + raise RuntimeError( + f"Command failed{label} with exit code {result.returncode}:\n" + f" {' '.join(cmd)}\n" + f"{result.stderr.strip()}" + ) + + def get_connection(hostname: str | None) -> Connection: if hostname is None or hostname == socket.gethostname(): return LocalConnection() diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 409403bc665..d7034d192a7 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -17,11 +17,15 @@ class SlurmExecutor(BaseExecutor): bash inside the sbatch script (no container needed — they just poll HTTP). """ - def run(self, config: SubmitConfig) -> None: + def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: compute = next(iter(config.compute.values())) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") remote_run_dir = Path(config.job.output_path) / timestamp + if dry_run: + self._dry_run(config, compute, remote_run_dir) + return + staging = self._stage(config, compute, remote_run_dir) with get_connection(compute.hostname) as conn: conn.copy(staging, remote_run_dir) @@ -30,6 +34,15 @@ def run(self, config: SubmitConfig) -> None: for b in config.driver.benchmarks ]) + def _dry_run(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> None: + print(f"[dry-run] remote run dir: {remote_run_dir}") + for benchmark in config.driver.benchmarks: + script = build_sbatch_script(config, benchmark, compute, remote_run_dir / benchmark.name) + print(f"\n{'='*60}") + print(f"[dry-run] sbatch script for benchmark: {benchmark.name}") + print(f"{'='*60}") + print(script) + def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> Path: staging = Path(tempfile.mkdtemp(prefix="gym-submit-")) for benchmark in config.driver.benchmarks: diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index 24a77aab300..01b382f4394 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -12,15 +12,18 @@ } -def _load_submit_config() -> SubmitConfig: +def _load_submit_config() -> tuple[SubmitConfig, bool]: """Load the submit config YAML and apply +key=value overrides from the command line.""" config_path: str | None = None + dry_run: bool = False override_tokens: list[str] = [] for token in sys.argv[1:]: stripped = token.lstrip("+") if stripped.startswith("_submit_config="): config_path = stripped[len("_submit_config="):] + elif stripped == "_dry_run=true": + dry_run = True else: override_tokens.append(stripped) @@ -31,11 +34,11 @@ def _load_submit_config() -> SubmitConfig: base = OmegaConf.load(config_path) overrides = OmegaConf.from_dotlist(override_tokens) if override_tokens else OmegaConf.create() merged = OmegaConf.merge(base, overrides) - return SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)) + return SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)), dry_run @experimental def submit() -> None: # pragma: no cover - config = _load_submit_config() + config, dry_run = _load_submit_config() compute = next(iter(config.compute.values())) - _EXECUTORS[type(compute)]().run(config) + _EXECUTORS[type(compute)]().run(config, dry_run=dry_run) From 8f63fbe5df722c5f67a1954145641432eb0e5171 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 01:12:27 +0200 Subject: [PATCH 14/66] add account to slurm config; prefix remote dir Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 1 + nemo_gym/orchestration/executors/slurm.py | 2 +- nemo_gym/orchestration/executors/slurm_script.py | 1 + tests/unit_tests/test_orchestration_api.py | 6 +++--- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 5789c0ed62d..7b265f0d9b2 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -48,6 +48,7 @@ class BaseComputeConfig(BaseModel): class SlurmComputeConfig(BaseComputeConfig): type: Literal["slurm"] + account: str hostname: str | None = None walltime: str | None = None node_pools: dict[str, NodePool] = {} diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index d7034d192a7..a786aa2958e 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -20,7 +20,7 @@ class SlurmExecutor(BaseExecutor): def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: compute = next(iter(config.compute.values())) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - remote_run_dir = Path(config.job.output_path) / timestamp + remote_run_dir = Path(config.job.output_path) / f"gym-job-{timestamp}" if dry_run: self._dry_run(config, compute, remote_run_dir) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 9b55c3055a5..0b4cc1cae1b 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -24,6 +24,7 @@ def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path) -> str: lines = [] + lines.append(f"#SBATCH --account={compute.account}") if compute.walltime: lines.append(f"#SBATCH --time={compute.walltime}") lines.append(f"#SBATCH --chdir={remote_bench_dir}") diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 380c0b056b1..f2eb4fbcc9c 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -18,10 +18,10 @@ from nemo_gym.orchestration.api import SubmitConfig -COMPUTE = {"cluster": {"type": "slurm", "hostname": "foo"}} +COMPUTE = {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}} COMPUTE_TWO = { - "cluster_a": {"type": "slurm", "hostname": "foo"}, - "cluster_b": {"type": "slurm", "hostname": "bar"}, + "cluster_a": {"type": "slurm", "account": "my-account", "hostname": "foo"}, + "cluster_b": {"type": "slurm", "account": "my-account", "hostname": "bar"}, } SERVICE = {"container": "gym:latest", "type": "vllm", "model": "org/model"} From d5c9b8a0a2feaab0832991a1d892f31c252305ea Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 01:28:57 +0200 Subject: [PATCH 15/66] get submit confirmation Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/executors/connection.py | 18 +++++++++++------- nemo_gym/orchestration/executors/slurm.py | 15 ++++++++++++++- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py index cbbeae2b5b8..a6183cda43e 100644 --- a/nemo_gym/orchestration/executors/connection.py +++ b/nemo_gym/orchestration/executors/connection.py @@ -19,7 +19,7 @@ def __exit__(self, *_) -> None: def copy(self, local: Path, remote: Path) -> None: ... @abstractmethod - def run(self, commands: list[str]) -> None: ... + def run(self, commands: list[str]) -> str: ... def close(self) -> None: pass @@ -31,9 +31,8 @@ def copy(self, local: Path, remote: Path) -> None: shutil.rmtree(remote) shutil.copytree(local, remote) - def run(self, commands: list[str]) -> None: - for cmd in commands: - _checked(shlex.split(cmd)) + def run(self, commands: list[str]) -> str: + return "\n".join(_checked(shlex.split(cmd)) for cmd in commands) class SSHConnection(Connection): @@ -82,6 +81,10 @@ def _ssh_opts(self) -> list[str]: return [*self._OPTS, "-o", "ControlMaster=no", "-o", f"ControlPath={self._socket}"] def copy(self, local: Path, remote: Path) -> None: + _checked( + ["ssh", *self._ssh_opts(), self._hostname, "mkdir", "-p", str(remote)], + context=f"mkdir -p {self._hostname}:{remote}", + ) _checked( [ "rsync", "-az", "--delete", @@ -92,9 +95,9 @@ def copy(self, local: Path, remote: Path) -> None: context=f"rsync to {self._hostname}:{remote}", ) - def run(self, commands: list[str]) -> None: + def run(self, commands: list[str]) -> str: script = "\n".join(commands) - _checked( + return _checked( ["ssh", *self._ssh_opts(), self._hostname, "bash", "-s"], input=script, context=f"ssh commands on {self._hostname}", @@ -110,7 +113,7 @@ def close(self) -> None: self._master.wait() -def _checked(cmd: list[str], *, input: str | None = None, context: str = "") -> None: +def _checked(cmd: list[str], *, input: str | None = None, context: str = "") -> str: result = subprocess.run(cmd, input=input, text=True, capture_output=True) if result.returncode != 0: label = f" ({context})" if context else "" @@ -119,6 +122,7 @@ def _checked(cmd: list[str], *, input: str | None = None, context: str = "") -> f" {' '.join(cmd)}\n" f"{result.stderr.strip()}" ) + return result.stdout def get_connection(hostname: str | None) -> Connection: diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index a786aa2958e..5a2d4912d93 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -1,13 +1,18 @@ +import re import shlex import tempfile from datetime import datetime from pathlib import Path +import rich + from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.base import BaseExecutor from nemo_gym.orchestration.executors.connection import get_connection from nemo_gym.orchestration.executors.slurm_script import build_sbatch_script +_SBATCH_JOB_ID_RE = re.compile(r"Submitted batch job (\d+)") + class SlurmExecutor(BaseExecutor): """Slurm executor for Pyxis-enabled clusters (https://github.com/NVIDIA/pyxis). @@ -29,11 +34,19 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: staging = self._stage(config, compute, remote_run_dir) with get_connection(compute.hostname) as conn: conn.copy(staging, remote_run_dir) - conn.run([ + output = conn.run([ f"sbatch {shlex.quote(str(remote_run_dir / b.name / 'job.sh'))}" for b in config.driver.benchmarks ]) + benchmarks = [b.name for b in config.driver.benchmarks] + job_ids = _SBATCH_JOB_ID_RE.findall(output) + for benchmark, job_id in zip(benchmarks, job_ids): + rich.print(f"[green]submitted[/green] {benchmark} → Slurm job [bold]{job_id}[/bold]") + unmatched = benchmarks[len(job_ids):] + for benchmark in unmatched: + rich.print(f"[green]submitted[/green] {benchmark} (job ID unavailable)") + def _dry_run(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> None: print(f"[dry-run] remote run dir: {remote_run_dir}") for benchmark in config.driver.benchmarks: From ad69cd6732db02f13a8a0ab18fa4449c49499a59 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 01:42:55 +0200 Subject: [PATCH 16/66] slurm-native log handling; default container for the driver Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 2 +- nemo_gym/orchestration/executors/slurm_script.py | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 7b265f0d9b2..ad7990322c0 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -66,7 +66,7 @@ class BenchmarkRunConfig(BaseModel): class DriverConfig(BaseModel): - container: str + container: str = "python:3.12" benchmarks: list[BenchmarkRunConfig] diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 0b4cc1cae1b..23785082609 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -22,8 +22,9 @@ """ -def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path) -> str: +def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benchmark_name: str) -> str: lines = [] + lines.append(f"#SBATCH --job-name=gym-{benchmark_name}") lines.append(f"#SBATCH --account={compute.account}") if compute.walltime: lines.append(f"#SBATCH --time={compute.walltime}") @@ -51,7 +52,7 @@ def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: def _render_service_command(name: str, container: str, command: str) -> str: return ( f"# service: {name}\n" - f"srun --container-image={shlex.quote(container)} {command} > logs/{name}.log 2>&1 &" + f"srun --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &" ) @@ -85,7 +86,7 @@ def build_sbatch_script( compute: SlurmComputeConfig, remote_bench_dir: Path, ) -> str: - directives = _render_directives(compute, remote_bench_dir) + directives = _render_directives(compute, remote_bench_dir, benchmark.name) service_commands = "\n\n".join( _render_service_command(name, service.container, _BUILDERS[type(service)](service)) @@ -100,8 +101,8 @@ def build_sbatch_script( driver_command = ( f"srun --container-image={shlex.quote(config.driver.container)} " - f"gym eval run --benchmark {shlex.quote(benchmark.name)} " - f"> logs/driver.log 2>&1" + f"--output=logs/driver.log " + f"gym eval run --benchmark {shlex.quote(benchmark.name)}" ) return _SCRIPT_TEMPLATE.format( From 473a78f04d51a4970bf408920ef34a8afb8f01c4 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 10:50:40 +0200 Subject: [PATCH 17/66] add health check validation Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index ad7990322c0..232ce83a800 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -23,6 +23,12 @@ class VllmServiceConfig(BaseServiceConfig): pipeline_parallel_size: int = 1 trust_remote_code: bool = False + @model_validator(mode="after") + def _default_health_check(self) -> "VllmServiceConfig": + if self.health_check is None: + self.health_check = HealthCheckConfig(port=self.port) + return self + class RayServiceConfig(BaseServiceConfig): type: Literal["ray"] From ec63660965bb5bc63f26c0acad2015315c27c431 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 10:58:56 +0200 Subject: [PATCH 18/66] gym eval run parameter passing yaml > CLI Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 3 +- .../orchestration/executors/slurm_script.py | 17 +++++++ tests/unit_tests/test_slurm_script.py | 49 +++++++++++++++++++ 3 files changed, 68 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/test_slurm_script.py diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 232ce83a800..55f1da64901 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -1,4 +1,4 @@ -from typing import Annotated, Literal +from typing import Annotated, Any, Literal from pydantic import BaseModel, Discriminator, Tag, model_validator @@ -69,6 +69,7 @@ class SlurmComputeConfig(BaseComputeConfig): class BenchmarkRunConfig(BaseModel): name: str + run: dict[str, Any] = {} class DriverConfig(BaseModel): diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 23785082609..2ae97990b54 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -1,5 +1,6 @@ import shlex from pathlib import Path +from typing import Any from nemo_gym.orchestration.api import ( BenchmarkRunConfig, @@ -74,6 +75,20 @@ def _build_ray_command(_service: RayServiceConfig) -> str: return "ray start --head" +def _flatten_run_args(run: dict[str, Any], prefix: str = "") -> list[str]: + args = [] + for key, value in run.items(): + full_key = f"{prefix}.{key}" if prefix else key + if isinstance(value, dict): + args.extend(_flatten_run_args(value, full_key)) + elif isinstance(value, list): + items = ",".join(str(v) for v in value) + args.append(shlex.quote(f"++{full_key}=[{items}]")) + else: + args.append(shlex.quote(f"++{full_key}={value}")) + return args + + _BUILDERS = { VllmServiceConfig: _build_vllm_command, RayServiceConfig: _build_ray_command, @@ -99,10 +114,12 @@ def build_sbatch_script( if service.health_check ) + run_args = " ".join(_flatten_run_args(benchmark.run)) driver_command = ( f"srun --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log " f"gym eval run --benchmark {shlex.quote(benchmark.name)}" + + (f" {run_args}" if run_args else "") ) return _SCRIPT_TEMPLATE.format( diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py new file mode 100644 index 00000000000..bf5a79c1422 --- /dev/null +++ b/tests/unit_tests/test_slurm_script.py @@ -0,0 +1,49 @@ +# 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. + +from nemo_gym.orchestration.executors.slurm_script import _flatten_run_args + + +def test_scalar_values(): + assert _flatten_run_args({"temperature": 0.05, "top_p": 0.9}) == [ + "++temperature=0.05", + "++top_p=0.9", + ] + + +def test_nested_dict(): + assert _flatten_run_args({"responses_create_params": {"max_concurrent": 92, "temperature": 0.05}}) == [ + "++responses_create_params.max_concurrent=92", + "++responses_create_params.temperature=0.05", + ] + + +def test_list_value(): + assert _flatten_run_args({"config_paths": ["benchmarks/gsm8k/config.yaml", "benchmarks/foo/config.yaml"]}) == [ + "'++config_paths=[benchmarks/gsm8k/config.yaml,benchmarks/foo/config.yaml]'", + ] + + +def test_empty(): + assert _flatten_run_args({}) == [] + + +def test_value_with_spaces_is_quoted(): + result = _flatten_run_args({"name": "my model"}) + assert result == ["'++name=my model'"] + + +def test_deeply_nested(): + assert _flatten_run_args({"a": {"b": {"c": 1}}}) == ["++a.b.c=1"] From 5e818bd23dca4fb0c41b1748f2c96bdf8ee7773e Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 11:03:11 +0200 Subject: [PATCH 19/66] move argument passing out of SlurmExecutor Signed-off-by: Tomasz Grzegorzek --- .../orchestration/executors/slurm_script.py | 18 ++---------------- nemo_gym/orchestration/executors/utils.py | 17 +++++++++++++++++ tests/unit_tests/test_slurm_script.py | 2 +- 3 files changed, 20 insertions(+), 17 deletions(-) create mode 100644 nemo_gym/orchestration/executors/utils.py diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 2ae97990b54..2632f074d69 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -1,7 +1,7 @@ import shlex from pathlib import Path -from typing import Any +from nemo_gym.orchestration.executors.utils import flatten_run_args from nemo_gym.orchestration.api import ( BenchmarkRunConfig, NodePool, @@ -75,20 +75,6 @@ def _build_ray_command(_service: RayServiceConfig) -> str: return "ray start --head" -def _flatten_run_args(run: dict[str, Any], prefix: str = "") -> list[str]: - args = [] - for key, value in run.items(): - full_key = f"{prefix}.{key}" if prefix else key - if isinstance(value, dict): - args.extend(_flatten_run_args(value, full_key)) - elif isinstance(value, list): - items = ",".join(str(v) for v in value) - args.append(shlex.quote(f"++{full_key}=[{items}]")) - else: - args.append(shlex.quote(f"++{full_key}={value}")) - return args - - _BUILDERS = { VllmServiceConfig: _build_vllm_command, RayServiceConfig: _build_ray_command, @@ -114,7 +100,7 @@ def build_sbatch_script( if service.health_check ) - run_args = " ".join(_flatten_run_args(benchmark.run)) + run_args = " ".join(flatten_run_args(benchmark.run)) driver_command = ( f"srun --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log " diff --git a/nemo_gym/orchestration/executors/utils.py b/nemo_gym/orchestration/executors/utils.py new file mode 100644 index 00000000000..845c47c3cdd --- /dev/null +++ b/nemo_gym/orchestration/executors/utils.py @@ -0,0 +1,17 @@ +import shlex +from typing import Any + + +def flatten_run_args(run: dict[str, Any], prefix: str = "") -> list[str]: + """Flatten a nested run config dict into shell-quoted ++key.path=value Hydra override tokens.""" + args = [] + for key, value in run.items(): + full_key = f"{prefix}.{key}" if prefix else key + if isinstance(value, dict): + args.extend(flatten_run_args(value, full_key)) + elif isinstance(value, list): + items = ",".join(str(v) for v in value) + args.append(shlex.quote(f"++{full_key}=[{items}]")) + else: + args.append(shlex.quote(f"++{full_key}={value}")) + return args diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index bf5a79c1422..2ec20873866 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from nemo_gym.orchestration.executors.slurm_script import _flatten_run_args +from nemo_gym.orchestration.executors.utils import flatten_run_args as _flatten_run_args def test_scalar_values(): From 0ae886753b0b34729cb8e4ea353ae3853787f4d2 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 11:42:56 +0200 Subject: [PATCH 20/66] separate templating Signed-off-by: Tomasz Grzegorzek --- .../executors/script_templates.py | 36 +++++++++++++++++++ .../orchestration/executors/slurm_script.py | 14 +++----- 2 files changed, 41 insertions(+), 9 deletions(-) create mode 100644 nemo_gym/orchestration/executors/script_templates.py diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py new file mode 100644 index 00000000000..8364555a0b6 --- /dev/null +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -0,0 +1,36 @@ +import re + +_HEALTH_WAIT_MULTI = """\ +# Wait for {name} (try multiple health endpoints) +echo "Waiting for {name} at {url}..." +{name_upper}_READY=0 +for _i in $(seq 1 {max_attempts}); do + if curl -sf "{url}/health" > /dev/null 2>&1 || curl -sf "{url}/openapi.json" > /dev/null 2>&1; then + echo " {name} ready." + {name_upper}_READY=1 + break + fi + if [ -n "${{{name_upper}_PID:-}}" ] && ! kill -0 ${name_upper}_PID 2>/dev/null; then + echo " {name} died during startup." + exit 1 + fi + sleep 5 +done +if [ ${name_upper}_READY -eq 0 ]; then + echo "ERROR: {name} did not become healthy after {max_attempts} attempts." + exit 1 +fi +""" + + +def bash_var(name: str) -> str: + return re.sub(r"[^A-Z0-9]", "_", name.upper()) + + +def render_health_check(name: str, port: int, timeout: int) -> str: + return _HEALTH_WAIT_MULTI.format( + name=name, + name_upper=bash_var(name), + url=f"http://localhost:{port}", + max_attempts=timeout // 5, + ) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 2632f074d69..07e62a27c46 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -1,6 +1,7 @@ import shlex from pathlib import Path +from nemo_gym.orchestration.executors.script_templates import bash_var, render_health_check from nemo_gym.orchestration.executors.utils import flatten_run_args from nemo_gym.orchestration.api import ( BenchmarkRunConfig, @@ -51,16 +52,11 @@ def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: def _render_service_command(name: str, container: str, command: str) -> str: + var = bash_var(name) return ( f"# service: {name}\n" - f"srun --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &" - ) - - -def _render_health_check(name: str, port: int, path: str, timeout: int) -> str: - return ( - f"# health check: {name} (timeout {timeout}s)\n" - f"timeout {timeout} bash -c 'until curl -sf http://localhost:{port}{shlex.quote(path)}; do sleep 2; done'" + f"srun --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" + f"{var}_PID=$!" ) @@ -95,7 +91,7 @@ def build_sbatch_script( ) health_checks = "\n\n".join( - _render_health_check(name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds) + render_health_check(name, service.health_check.port, service.health_check.timeout_seconds) for name, service in config.services.items() if service.health_check ) From be265f37a1a844f388bbb8b425a2a84893d6551c Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 12:15:02 +0200 Subject: [PATCH 21/66] PolicyModel config and synthactic sugar for wiring policy model with duplication back-off Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 28 +++++++++++++++++-- .../orchestration/executors/slurm_script.py | 2 +- tests/unit_tests/test_orchestration_api.py | 28 +++++++++++++++++++ 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 55f1da64901..15405020562 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -15,10 +15,15 @@ class BaseServiceConfig(BaseModel): health_check: HealthCheckConfig | None = None -class VllmServiceConfig(BaseServiceConfig): - type: Literal["vllm"] +class BaseModelServiceConfig(BaseServiceConfig): + """Base for services that serve a model and can be wired as the policy model.""" + model: str port: int = 8000 + + +class VllmServiceConfig(BaseModelServiceConfig): + type: Literal["vllm"] tensor_parallel_size: int = 1 pipeline_parallel_size: int = 1 trust_remote_code: bool = False @@ -74,6 +79,7 @@ class BenchmarkRunConfig(BaseModel): class DriverConfig(BaseModel): container: str = "python:3.12" + policy_model: str | None = None benchmarks: list[BenchmarkRunConfig] @@ -107,4 +113,22 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": f"({', '.join(sorted(compute_names))})." ) + if self.driver.policy_model is not None: + if self.driver.policy_model not in self.services: + raise ValueError( + f"driver.policy_model '{self.driver.policy_model}' does not match any service " + f"({', '.join(sorted(self.services))})." + ) + service = self.services[self.driver.policy_model] + if isinstance(service, BaseModelServiceConfig): + for benchmark in self.driver.benchmarks: + conflicts = [k for k in ("policy_base_url", "policy_model_name") if k in benchmark.run] + if conflicts: + raise ValueError( + f"Benchmark '{benchmark.name}' run config already sets {conflicts} " + f"but driver.policy_model is also set. Remove one." + ) + benchmark.run["policy_base_url"] = f"http://localhost:{service.port}/v1" + benchmark.run["policy_model_name"] = service.model + return self diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 07e62a27c46..6d2074e37e7 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -9,7 +9,7 @@ RayServiceConfig, SlurmComputeConfig, SubmitConfig, - VllmServiceConfig, + VllmServiceConfig, # used in _BUILDERS dispatch table ) _SCRIPT_TEMPLATE = """\ diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index f2eb4fbcc9c..8b816f3db93 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -51,3 +51,31 @@ def test_multiple_compute_raises(): def test_invalid_placement_raises(): with pytest.raises(ValidationError, match="does not match any compute resource"): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "placement": "nonexistent"}})) + + +def test_valid_policy_model(): + config = SubmitConfig.model_validate(_config(driver={**DRIVER, "policy_model": "svc"})) + assert config.driver.policy_model == "svc" + + +def test_invalid_policy_model_raises(): + with pytest.raises(ValidationError, match="does not match any service"): + SubmitConfig.model_validate(_config(driver={**DRIVER, "policy_model": "nonexistent"})) + + +def test_no_policy_model(): + config = SubmitConfig.model_validate(_config()) + assert config.driver.policy_model is None + + +def test_policy_model_injects_run_args(): + config = SubmitConfig.model_validate(_config(driver={**DRIVER, "policy_model": "svc"})) + benchmark = config.driver.benchmarks[0] + assert benchmark.run["policy_base_url"] == "http://localhost:8000/v1" + assert benchmark.run["policy_model_name"] == "org/model" + + +def test_policy_model_conflict_raises(): + driver = {**DRIVER, "policy_model": "svc", "benchmarks": [{"name": "gsm8k", "run": {"policy_base_url": "http://other"}}]} + with pytest.raises(ValidationError, match="already sets"): + SubmitConfig.model_validate(_config(driver=driver)) From cf8ae940125ce66513ddda1663f65e13a12c3dbb Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 13:32:48 +0200 Subject: [PATCH 22/66] allow to install on the fly; do not allow extra across api Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 28 +++++++++++++------ .../executors/script_templates.py | 20 +++++++++++++ .../orchestration/executors/slurm_script.py | 11 ++++---- 3 files changed, 45 insertions(+), 14 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 15405020562..58801caa75f 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -1,15 +1,19 @@ from typing import Annotated, Any, Literal -from pydantic import BaseModel, Discriminator, Tag, model_validator +from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator -class HealthCheckConfig(BaseModel): +class _StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class HealthCheckConfig(_StrictModel): path: str = "/health" port: int timeout_seconds: int = 60 -class BaseServiceConfig(BaseModel): +class BaseServiceConfig(_StrictModel): container: str placement: str | None = None health_check: HealthCheckConfig | None = None @@ -45,7 +49,7 @@ class RayServiceConfig(BaseServiceConfig): ] -class NodePool(BaseModel): +class NodePool(_StrictModel): partition: str nodes: int = 1 ntasks_per_node: int = 1 @@ -53,7 +57,7 @@ class NodePool(BaseModel): extra_args: dict[str, str] = {} -class BaseComputeConfig(BaseModel): +class BaseComputeConfig(_StrictModel): pass @@ -72,22 +76,28 @@ class SlurmComputeConfig(BaseComputeConfig): ] -class BenchmarkRunConfig(BaseModel): +class BenchmarkRunConfig(_StrictModel): name: str run: dict[str, Any] = {} -class DriverConfig(BaseModel): +class GymInstallConfig(_StrictModel): + repo: str = "https://github.com/NVIDIA-NeMo/gym" + ref: str + + +class DriverConfig(_StrictModel): container: str = "python:3.12" + gym_install: GymInstallConfig | None = None policy_model: str | None = None benchmarks: list[BenchmarkRunConfig] -class JobConfig(BaseModel): +class JobConfig(_StrictModel): output_path: str -class SubmitConfig(BaseModel): +class SubmitConfig(_StrictModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] driver: DriverConfig diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index 8364555a0b6..91a54ba6d6b 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -1,4 +1,5 @@ import re +import shlex _HEALTH_WAIT_MULTI = """\ # Wait for {name} (try multiple health endpoints) @@ -34,3 +35,22 @@ def render_health_check(name: str, port: int, timeout: int) -> str: url=f"http://localhost:{port}", max_attempts=timeout // 5, ) + + +def render_gym_cmd(benchmark_name: str, run_args: list[str]) -> str: + """Render a GYM_CMD bash array with each argument on its own line.""" + args = ["gym eval run", f"--benchmark {shlex.quote(benchmark_name)}", *run_args] + return "GYM_CMD=(\n " + "\n ".join(args) + "\n)" + + +def render_gym_install(repo: str, ref: str) -> str: + """Render a bash -c wrapper that installs uv + gym before exec-ing GYM_CMD.""" + install_cmd = shlex.quote(f"git+{repo}@{ref}") + return ( + "bash -c '\n" + " curl -LsSf https://astral.sh/uv/install.sh | sh\n" + ' source "$HOME/.local/bin/env"\n' + f" uv pip install --system {install_cmd}\n" + " exec \"$@\"\n" + "' -- \"${GYM_CMD[@]}\"" + ) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 6d2074e37e7..73397b186f6 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -1,7 +1,7 @@ import shlex from pathlib import Path -from nemo_gym.orchestration.executors.script_templates import bash_var, render_health_check +from nemo_gym.orchestration.executors.script_templates import bash_var, render_gym_cmd, render_gym_install, render_health_check from nemo_gym.orchestration.executors.utils import flatten_run_args from nemo_gym.orchestration.api import ( BenchmarkRunConfig, @@ -96,12 +96,13 @@ def build_sbatch_script( if service.health_check ) - run_args = " ".join(flatten_run_args(benchmark.run)) + gym_cmd = render_gym_cmd(benchmark.name, flatten_run_args(benchmark.run)) + gi = config.driver.gym_install + entrypoint = render_gym_install(gi.repo, gi.ref) if gi else '"${GYM_CMD[@]}"' driver_command = ( + f"{gym_cmd}\n" f"srun --container-image={shlex.quote(config.driver.container)} " - f"--output=logs/driver.log " - f"gym eval run --benchmark {shlex.quote(benchmark.name)}" - + (f" {run_args}" if run_args else "") + f"--output=logs/driver.log {entrypoint}" ) return _SCRIPT_TEMPLATE.format( From 3d7fba1851a534bab62b9fc524a5b1bfe235579a Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 14:10:00 +0200 Subject: [PATCH 23/66] switch from list to dict for benchmarks; overlap and no container mount home in indiviudla sruns Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 11 ++++---- nemo_gym/orchestration/executors/slurm.py | 27 +++++++++---------- .../orchestration/executors/slurm_script.py | 9 ++++--- tests/unit_tests/test_orchestration_api.py | 6 ++--- 4 files changed, 27 insertions(+), 26 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 58801caa75f..e8911c89173 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -9,7 +9,7 @@ class _StrictModel(BaseModel): class HealthCheckConfig(_StrictModel): path: str = "/health" - port: int + port: int | None = None timeout_seconds: int = 60 @@ -36,6 +36,8 @@ class VllmServiceConfig(BaseModelServiceConfig): def _default_health_check(self) -> "VllmServiceConfig": if self.health_check is None: self.health_check = HealthCheckConfig(port=self.port) + elif self.health_check.port is None: + self.health_check.port = self.port return self @@ -77,7 +79,6 @@ class SlurmComputeConfig(BaseComputeConfig): class BenchmarkRunConfig(_StrictModel): - name: str run: dict[str, Any] = {} @@ -90,7 +91,7 @@ class DriverConfig(_StrictModel): container: str = "python:3.12" gym_install: GymInstallConfig | None = None policy_model: str | None = None - benchmarks: list[BenchmarkRunConfig] + benchmarks: dict[str, BenchmarkRunConfig] class JobConfig(_StrictModel): @@ -131,11 +132,11 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": ) service = self.services[self.driver.policy_model] if isinstance(service, BaseModelServiceConfig): - for benchmark in self.driver.benchmarks: + for bench_name, benchmark in self.driver.benchmarks.items(): conflicts = [k for k in ("policy_base_url", "policy_model_name") if k in benchmark.run] if conflicts: raise ValueError( - f"Benchmark '{benchmark.name}' run config already sets {conflicts} " + f"Benchmark '{bench_name}' run config already sets {conflicts} " f"but driver.policy_model is also set. Remove one." ) benchmark.run["policy_base_url"] = f"http://localhost:{service.port}/v1" diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 5a2d4912d93..888d3b9a563 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -35,33 +35,32 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: with get_connection(compute.hostname) as conn: conn.copy(staging, remote_run_dir) output = conn.run([ - f"sbatch {shlex.quote(str(remote_run_dir / b.name / 'job.sh'))}" - for b in config.driver.benchmarks + f"sbatch {shlex.quote(str(remote_run_dir / name / 'job.sh'))}" + for name in config.driver.benchmarks ]) - benchmarks = [b.name for b in config.driver.benchmarks] + benchmark_names = list(config.driver.benchmarks) job_ids = _SBATCH_JOB_ID_RE.findall(output) - for benchmark, job_id in zip(benchmarks, job_ids): - rich.print(f"[green]submitted[/green] {benchmark} → Slurm job [bold]{job_id}[/bold]") - unmatched = benchmarks[len(job_ids):] - for benchmark in unmatched: - rich.print(f"[green]submitted[/green] {benchmark} (job ID unavailable)") + for name, job_id in zip(benchmark_names, job_ids): + rich.print(f"[green]submitted[/green] {name} → Slurm job [bold]{job_id}[/bold]") + for name in benchmark_names[len(job_ids):]: + rich.print(f"[green]submitted[/green] {name} (job ID unavailable)") def _dry_run(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> None: print(f"[dry-run] remote run dir: {remote_run_dir}") - for benchmark in config.driver.benchmarks: - script = build_sbatch_script(config, benchmark, compute, remote_run_dir / benchmark.name) + for name, benchmark in config.driver.benchmarks.items(): + script = build_sbatch_script(config, name, benchmark, compute, remote_run_dir / name) print(f"\n{'='*60}") - print(f"[dry-run] sbatch script for benchmark: {benchmark.name}") + print(f"[dry-run] sbatch script for benchmark: {name}") print(f"{'='*60}") print(script) def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> Path: staging = Path(tempfile.mkdtemp(prefix="gym-submit-")) - for benchmark in config.driver.benchmarks: - bench_dir = staging / benchmark.name + for name, benchmark in config.driver.benchmarks.items(): + bench_dir = staging / name bench_dir.mkdir() (bench_dir / "logs").mkdir() - script = build_sbatch_script(config, benchmark, compute, remote_run_dir / benchmark.name) + script = build_sbatch_script(config, name, benchmark, compute, remote_run_dir / name) (bench_dir / "job.sh").write_text(script) return staging diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 73397b186f6..7983093161a 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -55,7 +55,7 @@ def _render_service_command(name: str, container: str, command: str) -> str: var = bash_var(name) return ( f"# service: {name}\n" - f"srun --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" + f"srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -79,11 +79,12 @@ def _build_ray_command(_service: RayServiceConfig) -> str: def build_sbatch_script( config: SubmitConfig, + benchmark_name: str, benchmark: BenchmarkRunConfig, compute: SlurmComputeConfig, remote_bench_dir: Path, ) -> str: - directives = _render_directives(compute, remote_bench_dir, benchmark.name) + directives = _render_directives(compute, remote_bench_dir, benchmark_name) service_commands = "\n\n".join( _render_service_command(name, service.container, _BUILDERS[type(service)](service)) @@ -96,12 +97,12 @@ def build_sbatch_script( if service.health_check ) - gym_cmd = render_gym_cmd(benchmark.name, flatten_run_args(benchmark.run)) + gym_cmd = render_gym_cmd(benchmark_name, flatten_run_args(benchmark.run)) gi = config.driver.gym_install entrypoint = render_gym_install(gi.repo, gi.ref) if gi else '"${GYM_CMD[@]}"' driver_command = ( f"{gym_cmd}\n" - f"srun --container-image={shlex.quote(config.driver.container)} " + f"srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 8b816f3db93..2ded3fde462 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -25,7 +25,7 @@ } SERVICE = {"container": "gym:latest", "type": "vllm", "model": "org/model"} -DRIVER = {"container": "gym:latest", "benchmarks": [{"name": "gsm8k"}]} +DRIVER = {"container": "gym:latest", "benchmarks": {"gsm8k": {}}} JOB = {"output_path": "/tmp/gym-jobs"} @@ -70,12 +70,12 @@ def test_no_policy_model(): def test_policy_model_injects_run_args(): config = SubmitConfig.model_validate(_config(driver={**DRIVER, "policy_model": "svc"})) - benchmark = config.driver.benchmarks[0] + benchmark = config.driver.benchmarks["gsm8k"] assert benchmark.run["policy_base_url"] == "http://localhost:8000/v1" assert benchmark.run["policy_model_name"] == "org/model" def test_policy_model_conflict_raises(): - driver = {**DRIVER, "policy_model": "svc", "benchmarks": [{"name": "gsm8k", "run": {"policy_base_url": "http://other"}}]} + driver = {**DRIVER, "policy_model": "svc", "benchmarks": {"gsm8k": {"run": {"policy_base_url": "http://other"}}}} with pytest.raises(ValidationError, match="already sets"): SubmitConfig.model_validate(_config(driver=driver)) From e4f14358eb19c95200718848270a467c8618bbef Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 18:20:23 +0200 Subject: [PATCH 24/66] add preparation stage as well Signed-off-by: Tomasz Grzegorzek --- .../data/gpqa_diamond_benchmark_metrics.json | 2 +- nemo_gym/orchestration/api.py | 1 + .../executors/script_templates.py | 52 +++++++++++++------ .../orchestration/executors/slurm_script.py | 24 +++++++-- 4 files changed, 58 insertions(+), 21 deletions(-) diff --git a/benchmarks/gpqa/data/gpqa_diamond_benchmark_metrics.json b/benchmarks/gpqa/data/gpqa_diamond_benchmark_metrics.json index 70c88882268..069b2c4b185 100644 --- a/benchmarks/gpqa/data/gpqa_diamond_benchmark_metrics.json +++ b/benchmarks/gpqa/data/gpqa_diamond_benchmark_metrics.json @@ -3,7 +3,7 @@ "type": "benchmark", "jsonl_fpath": "benchmarks/gpqa/data/gpqa_diamond_benchmark.jsonl", "prepare_script": "benchmarks/gpqa/prepare.py", - "prompt_config": "benchmarks/gpqa/prompts/default.yaml", + "prompt_config": "benchmarks/prompts/eval/aai/mcq-4choices.yaml", "num_repeats": 8, "Number of examples": 0, "Number of tools": { diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index e8911c89173..b3624d6ba52 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -79,6 +79,7 @@ class SlurmComputeConfig(BaseComputeConfig): class BenchmarkRunConfig(_StrictModel): + prepare: dict[str, Any] = {} run: dict[str, Any] = {} diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index 91a54ba6d6b..d08ff56bdf9 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -37,20 +37,38 @@ def render_health_check(name: str, port: int, timeout: int) -> str: ) -def render_gym_cmd(benchmark_name: str, run_args: list[str]) -> str: - """Render a GYM_CMD bash array with each argument on its own line.""" - args = ["gym eval run", f"--benchmark {shlex.quote(benchmark_name)}", *run_args] - return "GYM_CMD=(\n " + "\n ".join(args) + "\n)" - - -def render_gym_install(repo: str, ref: str) -> str: - """Render a bash -c wrapper that installs uv + gym before exec-ing GYM_CMD.""" - install_cmd = shlex.quote(f"git+{repo}@{ref}") - return ( - "bash -c '\n" - " curl -LsSf https://astral.sh/uv/install.sh | sh\n" - ' source "$HOME/.local/bin/env"\n' - f" uv pip install --system {install_cmd}\n" - " exec \"$@\"\n" - "' -- \"${GYM_CMD[@]}\"" - ) +def render_gym_cmd(subcommand: str, var_name: str, benchmark_name: str, args: list[str]) -> str: + """Render a bash array with each argument on its own line.""" + entries = [f"gym {subcommand}", f"--benchmark {shlex.quote(benchmark_name)}", *args] + return f"{var_name}=(\n " + "\n ".join(entries) + "\n)" + + +def render_driver_entrypoint( + repo: str | None, + ref: str | None, + prepare_cmd: str | None, +) -> str: + """Render the srun entrypoint for the driver step. + + When either gym_install or prepare is needed, wraps everything in a single + bash -c so prepare and run happen in the same srun step and container. + """ + preamble: list[str] = [] + + if repo and ref: + install_cmd = shlex.quote(f"git+{repo}@{ref}") + preamble += [ + "curl -LsSf https://astral.sh/uv/install.sh | sh", + 'source "$HOME/.local/bin/env"', + f"uv pip install --system {install_cmd}", + ] + + if prepare_cmd: + preamble.append(prepare_cmd) + + if not preamble: + return '"${GYM_CMD[@]}"' + + preamble.append('exec "$@"') + body = "\n ".join(preamble) + return f"bash -c '\n {body}\n' -- \"${{GYM_CMD[@]}}\"" diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 7983093161a..52e8948f4a2 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -1,7 +1,7 @@ import shlex from pathlib import Path -from nemo_gym.orchestration.executors.script_templates import bash_var, render_gym_cmd, render_gym_install, render_health_check +from nemo_gym.orchestration.executors.script_templates import bash_var, render_driver_entrypoint, render_gym_cmd, render_health_check from nemo_gym.orchestration.executors.utils import flatten_run_args from nemo_gym.orchestration.api import ( BenchmarkRunConfig, @@ -20,6 +20,8 @@ {health_checks} +{prepare_command} + {driver_command} """ @@ -97,9 +99,24 @@ def build_sbatch_script( if service.health_check ) - gym_cmd = render_gym_cmd(benchmark_name, flatten_run_args(benchmark.run)) gi = config.driver.gym_install - entrypoint = render_gym_install(gi.repo, gi.ref) if gi else '"${GYM_CMD[@]}"' + + prepare_cmd = None + if benchmark.prepare: + prepare_args = " ".join( + [f"gym eval prepare --benchmark {shlex.quote(benchmark_name)}"] + + flatten_run_args(benchmark.prepare) + ) + prepare_cmd = prepare_args + + extra_flags = ["--model-type openai_model"] if config.driver.policy_model else [] + gym_cmd = render_gym_cmd("eval run", "GYM_CMD", benchmark_name, extra_flags + flatten_run_args(benchmark.run)) + entrypoint = render_driver_entrypoint( + repo=gi.repo if gi else None, + ref=gi.ref if gi else None, + prepare_cmd=prepare_cmd, + ) + prepare_command = "" driver_command = ( f"{gym_cmd}\n" f"srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " @@ -110,5 +127,6 @@ def build_sbatch_script( directives=directives, service_commands=service_commands, health_checks=health_checks, + prepare_command=prepare_command, driver_command=driver_command, ) From a1f681e866ff45875268bf79d5a5b9b0ca54247a Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 18:54:28 +0200 Subject: [PATCH 25/66] fix hydra force-override Signed-off-by: Tomasz Grzegorzek --- .../orchestration/executors/script_templates.py | 4 ++-- nemo_gym/orchestration/executors/slurm.py | 1 + nemo_gym/orchestration/executors/slurm_script.py | 9 +++------ nemo_gym/orchestration/executors/utils.py | 4 ++-- tests/unit_tests/test_slurm_script.py | 14 +++++++------- 5 files changed, 15 insertions(+), 17 deletions(-) diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index d08ff56bdf9..c5eef66bf03 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -37,9 +37,9 @@ def render_health_check(name: str, port: int, timeout: int) -> str: ) -def render_gym_cmd(subcommand: str, var_name: str, benchmark_name: str, args: list[str]) -> str: +def render_gym_cmd(subcommand: str, var_name: str, args: list[str]) -> str: """Render a bash array with each argument on its own line.""" - entries = [f"gym {subcommand}", f"--benchmark {shlex.quote(benchmark_name)}", *args] + entries = [f"gym {subcommand}", *args] return f"{var_name}=(\n " + "\n ".join(entries) + "\n)" diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 888d3b9a563..24ba60174c4 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -61,6 +61,7 @@ def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_d bench_dir = staging / name bench_dir.mkdir() (bench_dir / "logs").mkdir() + (bench_dir / "artifacts").mkdir() script = build_sbatch_script(config, name, benchmark, compute, remote_run_dir / name) (bench_dir / "job.sh").write_text(script) return staging diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 52e8948f4a2..fe9b4822d57 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -103,14 +103,11 @@ def build_sbatch_script( prepare_cmd = None if benchmark.prepare: - prepare_args = " ".join( - [f"gym eval prepare --benchmark {shlex.quote(benchmark_name)}"] - + flatten_run_args(benchmark.prepare) - ) - prepare_cmd = prepare_args + prepare_cmd = "gym eval prepare " + " ".join(flatten_run_args(benchmark.prepare)) + output_path = "+output_jsonl_fpath=artifacts/rollouts.jsonl" extra_flags = ["--model-type openai_model"] if config.driver.policy_model else [] - gym_cmd = render_gym_cmd("eval run", "GYM_CMD", benchmark_name, extra_flags + flatten_run_args(benchmark.run)) + gym_cmd = render_gym_cmd("eval run", "GYM_CMD", [output_path] + extra_flags + flatten_run_args(benchmark.run)) entrypoint = render_driver_entrypoint( repo=gi.repo if gi else None, ref=gi.ref if gi else None, diff --git a/nemo_gym/orchestration/executors/utils.py b/nemo_gym/orchestration/executors/utils.py index 845c47c3cdd..36e257c2605 100644 --- a/nemo_gym/orchestration/executors/utils.py +++ b/nemo_gym/orchestration/executors/utils.py @@ -11,7 +11,7 @@ def flatten_run_args(run: dict[str, Any], prefix: str = "") -> list[str]: args.extend(flatten_run_args(value, full_key)) elif isinstance(value, list): items = ",".join(str(v) for v in value) - args.append(shlex.quote(f"++{full_key}=[{items}]")) + args.append(shlex.quote(f"+{full_key}=[{items}]")) else: - args.append(shlex.quote(f"++{full_key}={value}")) + args.append(shlex.quote(f"+{full_key}={value}")) return args diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 2ec20873866..bd4d78f8ed3 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -18,21 +18,21 @@ def test_scalar_values(): assert _flatten_run_args({"temperature": 0.05, "top_p": 0.9}) == [ - "++temperature=0.05", - "++top_p=0.9", + "+temperature=0.05", + "+top_p=0.9", ] def test_nested_dict(): assert _flatten_run_args({"responses_create_params": {"max_concurrent": 92, "temperature": 0.05}}) == [ - "++responses_create_params.max_concurrent=92", - "++responses_create_params.temperature=0.05", + "+responses_create_params.max_concurrent=92", + "+responses_create_params.temperature=0.05", ] def test_list_value(): assert _flatten_run_args({"config_paths": ["benchmarks/gsm8k/config.yaml", "benchmarks/foo/config.yaml"]}) == [ - "'++config_paths=[benchmarks/gsm8k/config.yaml,benchmarks/foo/config.yaml]'", + "'+config_paths=[benchmarks/gsm8k/config.yaml,benchmarks/foo/config.yaml]'", ] @@ -42,8 +42,8 @@ def test_empty(): def test_value_with_spaces_is_quoted(): result = _flatten_run_args({"name": "my model"}) - assert result == ["'++name=my model'"] + assert result == ["'+name=my model'"] def test_deeply_nested(): - assert _flatten_run_args({"a": {"b": {"c": 1}}}) == ["++a.b.c=1"] + assert _flatten_run_args({"a": {"b": {"c": 1}}}) == ["+a.b.c=1"] From 4d792070a27e6acb72c17cd0273c114b7c87bb7b Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 20:31:51 +0200 Subject: [PATCH 26/66] add a dummy policy_api_key; change to local copy because we cannot resolve configs properly (?) Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 3 ++- nemo_gym/orchestration/executors/script_templates.py | 7 +++++-- tests/unit_tests/test_orchestration_api.py | 1 + 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index b3624d6ba52..1ab9588aca5 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -134,7 +134,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": service = self.services[self.driver.policy_model] if isinstance(service, BaseModelServiceConfig): for bench_name, benchmark in self.driver.benchmarks.items(): - conflicts = [k for k in ("policy_base_url", "policy_model_name") if k in benchmark.run] + conflicts = [k for k in ("policy_base_url", "policy_model_name", "policy_api_key") if k in benchmark.run] if conflicts: raise ValueError( f"Benchmark '{bench_name}' run config already sets {conflicts} " @@ -142,5 +142,6 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": ) benchmark.run["policy_base_url"] = f"http://localhost:{service.port}/v1" benchmark.run["policy_model_name"] = service.model + benchmark.run["policy_api_key"] = "dummy" return self diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index c5eef66bf03..e2181ef7976 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -56,11 +56,14 @@ def render_driver_entrypoint( preamble: list[str] = [] if repo and ref: - install_cmd = shlex.quote(f"git+{repo}@{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"uv pip install --system {install_cmd}", + f"git clone {shlex.quote(repo)}", + f"cd {shlex.quote(repo_name)}", + f"git checkout {shlex.quote(ref)}", + "uv pip install -e . --system", ] if prepare_cmd: diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 2ded3fde462..b4a4cdb8d20 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -73,6 +73,7 @@ def test_policy_model_injects_run_args(): benchmark = config.driver.benchmarks["gsm8k"] assert benchmark.run["policy_base_url"] == "http://localhost:8000/v1" assert benchmark.run["policy_model_name"] == "org/model" + assert benchmark.run["policy_api_key"] == "dummy" def test_policy_model_conflict_raises(): From 8c27a4c23f0c85b111828b3b9feec554a0d65b21 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 23:44:40 +0200 Subject: [PATCH 27/66] add some more tests Signed-off-by: Tomasz Grzegorzek --- tests/unit_tests/test_slurm_script.py | 255 +++++++++++++++++++++++++- 1 file changed, 253 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index bd4d78f8ed3..30b81270e2c 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -13,8 +13,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +from pathlib import Path + +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_vllm_command, + _render_directives, + _render_pool_directives, + _render_service_command, + build_sbatch_script, +) from nemo_gym.orchestration.executors.utils import flatten_run_args as _flatten_run_args +# --------------------------------------------------------------------------- +# flatten_run_args +# --------------------------------------------------------------------------- def test_scalar_values(): assert _flatten_run_args({"temperature": 0.05, "top_p": 0.9}) == [ @@ -41,9 +55,246 @@ def test_empty(): def test_value_with_spaces_is_quoted(): - result = _flatten_run_args({"name": "my model"}) - assert result == ["'+name=my model'"] + assert _flatten_run_args({"name": "my model"}) == ["'+name=my model'"] def test_deeply_nested(): assert _flatten_run_args({"a": {"b": {"c": 1}}}) == ["+a.b.c=1"] + + +# --------------------------------------------------------------------------- +# _render_pool_directives +# --------------------------------------------------------------------------- + +def test_render_pool_directives_basic(pool): + lines = _render_pool_directives("main", pool) + assert "#SBATCH --partition=batch # pool: main" in lines + assert "#SBATCH --nodes=1" in lines + assert "#SBATCH --ntasks-per-node=4" in lines + + +def test_render_pool_directives_gpus(pool): + pool.gpus_per_node = 4 + lines = _render_pool_directives("main", pool) + assert "#SBATCH --gpus-per-node=4" in lines + + +def test_render_pool_directives_extra_args(pool): + pool.extra_args["gres"] = "shard:8" + lines = _render_pool_directives("main", pool) + assert "#SBATCH --gres=shard:8" in lines + + +# --------------------------------------------------------------------------- +# _render_directives +# --------------------------------------------------------------------------- + +def test_render_directives_job_name(compute, bench_dir): + out = _render_directives(compute, bench_dir, "gsm8k") + assert "#SBATCH --job-name=gym-gsm8k" in out + + +def test_render_directives_account(compute, bench_dir): + out = _render_directives(compute, bench_dir, "gsm8k") + assert "#SBATCH --account=my-account" in out + + +def test_render_directives_walltime(compute, bench_dir): + compute.walltime = "01:00:00" + out = _render_directives(compute, bench_dir, "gsm8k") + assert "#SBATCH --time=01:00:00" in out + + +def test_render_directives_no_walltime(compute, bench_dir): + compute.walltime = None + out = _render_directives(compute, bench_dir, "gsm8k") + assert "--time" not in out + + +def test_render_directives_chdir(compute, bench_dir): + out = _render_directives(compute, bench_dir, "gsm8k") + assert f"#SBATCH --chdir={bench_dir}" in out + + +# --------------------------------------------------------------------------- +# _render_service_command +# --------------------------------------------------------------------------- + +def test_render_service_command_contains_srun(): + out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model") + assert "srun --overlap --no-container-mount-home" in out + assert "--container-image=vllm:latest" in out + assert "vllm serve model" in out + assert out.endswith("VLLM_MODEL_PID=$!") + + +def test_render_service_command_backgrounded(): + out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model") + assert "& " in out or out.split("\n")[1].endswith(" &") + + +def test_render_service_command_log_file(): + out = _render_service_command("my_service", "img:latest", "cmd") + assert "--output=logs/my_service.log" in out + + +# --------------------------------------------------------------------------- +# _build_vllm_command +# --------------------------------------------------------------------------- + +def test_build_vllm_command_basic(vllm_service): + cmd = _build_vllm_command(vllm_service) + assert "vllm serve" in cmd + assert "--port 8000" in cmd + assert "--tensor-parallel-size 1" in cmd + + +def test_build_vllm_command_trust_remote_code(vllm_service): + vllm_service.trust_remote_code = True + cmd = _build_vllm_command(vllm_service) + assert "--trust-remote-code" in cmd + + +def test_build_vllm_command_no_trust_remote_code_by_default(vllm_service): + cmd = _build_vllm_command(vllm_service) + assert "--trust-remote-code" not in cmd + + +# --------------------------------------------------------------------------- +# render_gym_cmd +# --------------------------------------------------------------------------- + +def test_render_gym_cmd_subcommand(): + out = render_gym_cmd("eval run", "GYM_CMD", ["+foo=bar"]) + assert out.startswith("GYM_CMD=(") + assert "gym eval run" in out + assert "+foo=bar" in out + + +def test_render_gym_cmd_prepare(): + out = render_gym_cmd("eval prepare", "GYM_PREPARE_CMD", []) + assert "gym eval prepare" in out + assert "GYM_PREPARE_CMD=(" in out + + +# --------------------------------------------------------------------------- +# render_driver_entrypoint +# --------------------------------------------------------------------------- + +def test_render_driver_entrypoint_no_install_no_prepare(): + out = render_driver_entrypoint(None, None, None) + assert out == '"${GYM_CMD[@]}"' + + +def test_render_driver_entrypoint_with_gym_install(): + out = render_driver_entrypoint("https://github.com/NVIDIA-NeMo/gym", "main", None) + assert "git clone" in out + assert "git checkout main" in out + assert "uv pip install -e . --system" in out + assert 'exec "$@"' in out + assert '"${GYM_CMD[@]}"' in out + + +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 + assert 'exec "$@"' in out + + +def test_render_driver_entrypoint_install_and_prepare(): + out = render_driver_entrypoint("https://github.com/NVIDIA-NeMo/gym", "v1.0", "gym eval prepare") + assert "git clone" in out + assert "git checkout v1.0" in out + assert "gym eval prepare" in out + assert 'exec "$@"' in out + + +# --------------------------------------------------------------------------- +# build_sbatch_script (integration) +# --------------------------------------------------------------------------- + +def test_build_sbatch_script_contains_shebang(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert script.startswith("#!/bin/bash") + + +def test_build_sbatch_script_contains_vllm_srun(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "vllm serve" in script + assert "srun --overlap" in script + + +def test_build_sbatch_script_driver_output_flag(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "--output=logs/driver.log" in script + + +def test_build_sbatch_script_output_jsonl_fpath(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "+output_jsonl_fpath=artifacts/rollouts.jsonl" in script + + +def test_build_sbatch_script_policy_model_flags(submit_config_with_policy, bench_dir): + config = submit_config_with_policy + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "--model-type openai_model" in script + assert "+policy_base_url=" in script + assert "+policy_model_name=" in script + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +import pytest +from nemo_gym.orchestration.api import NodePool, SlurmComputeConfig, VllmServiceConfig + + +@pytest.fixture +def bench_dir(): + return Path("/remote/jobs/gym-job-20260729/gsm8k") + + +@pytest.fixture +def pool(): + return NodePool(partition="batch", nodes=1, ntasks_per_node=4) + + +@pytest.fixture +def compute(): + return SlurmComputeConfig(type="slurm", account="my-account", hostname="foo") + + +@pytest.fixture +def vllm_service(): + return VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model") + + +@pytest.fixture +def submit_config(): + return SubmitConfig.model_validate({ + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + }) + + +@pytest.fixture +def submit_config_with_policy(): + return SubmitConfig.model_validate({ + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "policy_model": "vllm_model", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + }) From 427346e1963862462b1145da3bb1c8e5e6b34768 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Wed, 29 Jul 2026 23:49:20 +0200 Subject: [PATCH 28/66] Add in-code doc comments Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 21 ++++++++++++++++--- .../orchestration/executors/connection.py | 6 +++++- .../orchestration/executors/slurm_script.py | 4 ++++ 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 1ab9588aca5..56cb0f3b3b0 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -3,18 +3,21 @@ from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator +# Reject unknown fields on all config models so typos in YAML surface immediately. class _StrictModel(BaseModel): model_config = ConfigDict(extra="forbid") class HealthCheckConfig(_StrictModel): path: str = "/health" + # port defaults to None so VllmServiceConfig can fill it from service.port when omitted. port: int | None = None timeout_seconds: int = 60 class BaseServiceConfig(_StrictModel): container: str + # Resolved to the sole compute resource name at validation time when not set. placement: str | None = None health_check: HealthCheckConfig | None = None @@ -34,6 +37,8 @@ class VllmServiceConfig(BaseModelServiceConfig): @model_validator(mode="after") def _default_health_check(self) -> "VllmServiceConfig": + # vLLM always exposes /health on its serving port; set it automatically + # so the sbatch script gets a health check without the user having to repeat the port. if self.health_check is None: self.health_check = HealthCheckConfig(port=self.port) elif self.health_check.port is None: @@ -45,6 +50,7 @@ class RayServiceConfig(BaseServiceConfig): type: Literal["ray"] +# Discriminated union keyed on `type`; Pydantic rejects unknown type values at parse time. ServiceConfig = Annotated[ Annotated[VllmServiceConfig, Tag("vllm")] | Annotated[RayServiceConfig, Tag("ray")], Discriminator("type"), @@ -55,7 +61,9 @@ class NodePool(_StrictModel): partition: str nodes: int = 1 ntasks_per_node: int = 1 + # Structured field the executor uses for smart deployment decisions (e.g. multi-instance vLLM). gpus_per_node: int | None = None + # Arbitrary #SBATCH directives forwarded verbatim for options we don't model explicitly. extra_args: dict[str, str] = {} @@ -66,10 +74,10 @@ class BaseComputeConfig(_StrictModel): class SlurmComputeConfig(BaseComputeConfig): type: Literal["slurm"] account: str - hostname: str | None = None + hostname: str | None = None # None means we're already on the login node; skip SSH. walltime: str | None = None node_pools: dict[str, NodePool] = {} - extra_args: dict[str, str] = {} + extra_args: dict[str, str] = {} # Job-level #SBATCH directives (e.g. --comment, --mail-user). ComputeConfig = Annotated[ @@ -79,23 +87,29 @@ class SlurmComputeConfig(BaseComputeConfig): class BenchmarkRunConfig(_StrictModel): + # Hydra overrides forwarded to `gym eval prepare`. Flattened to +key=value tokens. prepare: dict[str, Any] = {} + # Hydra overrides forwarded to `gym eval run`. policy_model wiring is injected here at + # validation time so all executors see it uniformly via flatten_run_args. run: dict[str, Any] = {} class GymInstallConfig(_StrictModel): repo: str = "https://github.com/NVIDIA-NeMo/gym" - ref: str + ref: str # Git tag or commit hash. class DriverConfig(_StrictModel): container: str = "python:3.12" gym_install: GymInstallConfig | None = None + # Name of a service in `services:` to use as the policy model. When set, injects + # policy_base_url/policy_model_name/policy_api_key into each benchmark's run config. policy_model: str | None = None benchmarks: dict[str, BenchmarkRunConfig] class JobConfig(_StrictModel): + # Remote base directory. Each submit creates a timestamped subdirectory here. output_path: str @@ -142,6 +156,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": ) benchmark.run["policy_base_url"] = f"http://localhost:{service.port}/v1" benchmark.run["policy_model_name"] = service.model + # vLLM doesn't require auth; dummy key satisfies clients that require the header. benchmark.run["policy_api_key"] = "dummy" return self diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py index a6183cda43e..82ffe33dbb4 100644 --- a/nemo_gym/orchestration/executors/connection.py +++ b/nemo_gym/orchestration/executors/connection.py @@ -47,6 +47,8 @@ def __init__(self, hostname: str) -> None: self._master: subprocess.Popen | None = None def __enter__(self) -> "SSHConnection": + # Open a persistent master connection. Subsequent ssh/rsync calls reuse + # the same socket via ControlMaster=no, avoiding per-command handshakes. self._master = subprocess.Popen( [ "ssh", @@ -54,7 +56,7 @@ def __enter__(self) -> "SSHConnection": "-o", "ControlMaster=yes", "-o", f"ControlPath={self._socket}", "-o", "ControlPersist=yes", - "-N", + "-N", # no remote command — just keep the tunnel open self._hostname, ], stderr=subprocess.PIPE, @@ -96,6 +98,8 @@ def copy(self, local: Path, remote: Path) -> None: ) def run(self, commands: list[str]) -> str: + # Send all commands as a single bash script over one SSH session to avoid + # repeated connection overhead for multi-benchmark submits. script = "\n".join(commands) return _checked( ["ssh", *self._ssh_opts(), self._hostname, "bash", "-s"], diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index fe9b4822d57..3874f4ceef7 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -32,6 +32,7 @@ def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benc lines.append(f"#SBATCH --account={compute.account}") if compute.walltime: lines.append(f"#SBATCH --time={compute.walltime}") + # --chdir makes relative paths (logs/, artifacts/) resolve correctly inside the job. lines.append(f"#SBATCH --chdir={remote_bench_dir}") for key, val in compute.extra_args.items(): lines.append(f"#SBATCH --{key}={val}") @@ -55,6 +56,9 @@ def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: def _render_service_command(name: str, container: str, command: str) -> str: var = bash_var(name) + # --overlap lets this step share the allocation with other concurrent steps (driver + services). + # --no-container-mount-home avoids polluting the container with host home directory contents. + # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" f"srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" From 9308727d037d75fa1799949df922064a4235a079 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:01:37 +0200 Subject: [PATCH 29/66] smarter entrypoint Signed-off-by: Tomasz Grzegorzek --- nemo_gym/cli/main.py | 19 ++++++++++++++---- nemo_gym/orchestration/submit.py | 33 +------------------------------- 2 files changed, 16 insertions(+), 36 deletions(-) diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index a1b7261686e..29f6aaa14ca 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -271,6 +271,19 @@ def _merge_config_paths(overrides: list[str]) -> list[str]: return ([f"+config_paths=[{','.join(paths)}]"] if paths else []) + rest +def _eval_submit(args: argparse.Namespace, overrides: list[str]) -> None: + from omegaconf import OmegaConf + from nemo_gym.orchestration.api import SubmitConfig + from nemo_gym.orchestration.submit import submit + + merged = OmegaConf.merge( + OmegaConf.load(args.config), + OmegaConf.from_dotlist([t.lstrip("+") for t in overrides]) if overrides else OmegaConf.create(), + ) + config = SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)) + submit(config, dry_run=args.dry_run) + + def _eval_run(args: argparse.Namespace, overrides: list[str]) -> None: target = "nemo_gym.cli.eval:collect_rollouts" if args.no_serve else "nemo_gym.cli.eval:e2e_rollout_collection" dispatch(target, overrides) @@ -538,16 +551,14 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: ), ), "eval submit": Command( - target="nemo_gym.orchestration.submit:submit", + target=_eval_submit, summary="Submit a job.", flags=( Flag( - register=lambda p: p.add_argument("--config", metavar="PATH", help="Config YAML file to load."), - translate_to_hydra=lambda args: [f"+_submit_config={args.config}"] if args.config else [], + register=lambda p: p.add_argument("--config", "-c", required=True, metavar="PATH", help="Submit config YAML file."), ), Flag( register=lambda p: p.add_argument("--dry-run", action="store_true", help="Print generated job scripts without submitting."), - translate_to_hydra=lambda args: ["+_dry_run=true"] if args.dry_run else [], ), ), ), diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index 01b382f4394..ff8bbb527e3 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -1,8 +1,3 @@ -import sys - -import rich -from omegaconf import OmegaConf - from nemo_gym.decorators import experimental from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.slurm import SlurmExecutor @@ -12,33 +7,7 @@ } -def _load_submit_config() -> tuple[SubmitConfig, bool]: - """Load the submit config YAML and apply +key=value overrides from the command line.""" - config_path: str | None = None - dry_run: bool = False - override_tokens: list[str] = [] - - for token in sys.argv[1:]: - stripped = token.lstrip("+") - if stripped.startswith("_submit_config="): - config_path = stripped[len("_submit_config="):] - elif stripped == "_dry_run=true": - dry_run = True - else: - override_tokens.append(stripped) - - if config_path is None: - rich.print("[red]error:[/red] No config provided. Use [bold]--config [/bold] to specify a submit config YAML.") - sys.exit(1) - - base = OmegaConf.load(config_path) - overrides = OmegaConf.from_dotlist(override_tokens) if override_tokens else OmegaConf.create() - merged = OmegaConf.merge(base, overrides) - return SubmitConfig.model_validate(OmegaConf.to_container(merged, resolve=True)), dry_run - - @experimental -def submit() -> None: # pragma: no cover - config, dry_run = _load_submit_config() +def submit(config: SubmitConfig, *, dry_run: bool = False) -> None: # pragma: no cover compute = next(iter(config.compute.values())) _EXECUTORS[type(compute)]().run(config, dry_run=dry_run) From c7487e9caa9a976a8f2a47687de30bbb80c432bf Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:04:28 +0200 Subject: [PATCH 30/66] cleanup tempdirs Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/executors/slurm.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 24ba60174c4..e3ed88c4649 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -31,13 +31,14 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: self._dry_run(config, compute, remote_run_dir) return - staging = self._stage(config, compute, remote_run_dir) - with get_connection(compute.hostname) as conn: - conn.copy(staging, remote_run_dir) - output = conn.run([ - f"sbatch {shlex.quote(str(remote_run_dir / name / 'job.sh'))}" - for name in config.driver.benchmarks - ]) + with tempfile.TemporaryDirectory(prefix="gym-submit-") as staging_str: + staging = self._stage(config, compute, remote_run_dir, Path(staging_str)) + with get_connection(compute.hostname) as conn: + conn.copy(staging, remote_run_dir) + output = conn.run([ + f"sbatch {shlex.quote(str(remote_run_dir / name / 'job.sh'))}" + for name in config.driver.benchmarks + ]) benchmark_names = list(config.driver.benchmarks) job_ids = _SBATCH_JOB_ID_RE.findall(output) @@ -55,8 +56,7 @@ def _dry_run(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run print(f"{'='*60}") print(script) - def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> Path: - staging = Path(tempfile.mkdtemp(prefix="gym-submit-")) + def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path, staging: Path) -> Path: for name, benchmark in config.driver.benchmarks.items(): bench_dir = staging / name bench_dir.mkdir() From 3d0a1efca994d8d1bed56eb612e86acc70f19d33 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:09:01 +0200 Subject: [PATCH 31/66] make CI happy Signed-off-by: Tomasz Grzegorzek --- tests/unit_tests/test_orchestration_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index b4a4cdb8d20..78168f1f964 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -73,7 +73,7 @@ def test_policy_model_injects_run_args(): benchmark = config.driver.benchmarks["gsm8k"] assert benchmark.run["policy_base_url"] == "http://localhost:8000/v1" assert benchmark.run["policy_model_name"] == "org/model" - assert benchmark.run["policy_api_key"] == "dummy" + assert benchmark.run["policy_api_key"] == "dummy" # pragma: allowlist secret def test_policy_model_conflict_raises(): From 8ec6fac350c4070aa5e1376cc2c29b943dda9312 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:11:38 +0200 Subject: [PATCH 32/66] add missing spdx headers Signed-off-by: Tomasz Grzegorzek --- nemo_gym/decorators.py | 15 +++++++++++++++ nemo_gym/orchestration/__init__.py | 16 ++++++++++++++++ nemo_gym/orchestration/api.py | 15 +++++++++++++++ nemo_gym/orchestration/executors/__init__.py | 16 ++++++++++++++++ nemo_gym/orchestration/executors/base.py | 15 +++++++++++++++ nemo_gym/orchestration/executors/connection.py | 15 +++++++++++++++ .../orchestration/executors/script_templates.py | 15 +++++++++++++++ nemo_gym/orchestration/executors/slurm.py | 15 +++++++++++++++ nemo_gym/orchestration/executors/slurm_script.py | 15 +++++++++++++++ nemo_gym/orchestration/executors/utils.py | 15 +++++++++++++++ nemo_gym/orchestration/submit.py | 15 +++++++++++++++ 11 files changed, 167 insertions(+) diff --git a/nemo_gym/decorators.py b/nemo_gym/decorators.py index e705ef18e2e..39a2bfcfaca 100644 --- a/nemo_gym/decorators.py +++ b/nemo_gym/decorators.py @@ -1,3 +1,18 @@ +# 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 functools import rich diff --git a/nemo_gym/orchestration/__init__.py b/nemo_gym/orchestration/__init__.py index e69de29bb2d..2c5a3a8a54b 100644 --- a/nemo_gym/orchestration/__init__.py +++ b/nemo_gym/orchestration/__init__.py @@ -0,0 +1,16 @@ +# 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. + + diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 56cb0f3b3b0..9bfdb6f7dc3 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -1,3 +1,18 @@ +# 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. + from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator diff --git a/nemo_gym/orchestration/executors/__init__.py b/nemo_gym/orchestration/executors/__init__.py index e69de29bb2d..2c5a3a8a54b 100644 --- a/nemo_gym/orchestration/executors/__init__.py +++ b/nemo_gym/orchestration/executors/__init__.py @@ -0,0 +1,16 @@ +# 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. + + diff --git a/nemo_gym/orchestration/executors/base.py b/nemo_gym/orchestration/executors/base.py index 0ec4a44c672..8eb95dce6de 100644 --- a/nemo_gym/orchestration/executors/base.py +++ b/nemo_gym/orchestration/executors/base.py @@ -1,3 +1,18 @@ +# 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. + from abc import ABC, abstractmethod from nemo_gym.orchestration.api import SubmitConfig diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py index 82ffe33dbb4..77949ee7308 100644 --- a/nemo_gym/orchestration/executors/connection.py +++ b/nemo_gym/orchestration/executors/connection.py @@ -1,3 +1,18 @@ +# 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 shlex import shutil import socket diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index e2181ef7976..a941c7b26c3 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -1,3 +1,18 @@ +# 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 re import shlex diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index e3ed88c4649..9213c6a1063 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -1,3 +1,18 @@ +# 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 re import shlex import tempfile diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 3874f4ceef7..34e39d74e74 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -1,3 +1,18 @@ +# 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 shlex from pathlib import Path diff --git a/nemo_gym/orchestration/executors/utils.py b/nemo_gym/orchestration/executors/utils.py index 36e257c2605..f4d779d7226 100644 --- a/nemo_gym/orchestration/executors/utils.py +++ b/nemo_gym/orchestration/executors/utils.py @@ -1,3 +1,18 @@ +# 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 shlex from typing import Any diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index ff8bbb527e3..e17bdc96e24 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -1,3 +1,18 @@ +# 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. + from nemo_gym.decorators import experimental from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.slurm import SlurmExecutor From 2daaa18a62098580b29b1e22cd6513c887448fb2 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:12:44 +0200 Subject: [PATCH 33/66] linting Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 8 ++-- .../orchestration/executors/connection.py | 22 ++++++----- .../executors/script_templates.py | 1 + nemo_gym/orchestration/executors/slurm.py | 17 +++++---- .../orchestration/executors/slurm_script.py | 10 ++++- nemo_gym/orchestration/submit.py | 1 + tests/unit_tests/test_orchestration_api.py | 3 +- tests/unit_tests/test_slurm_script.py | 38 +++++++++++++------ 8 files changed, 65 insertions(+), 35 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 9bfdb6f7dc3..933c7bc103a 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -139,9 +139,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": compute_names = set(self.compute) if len(compute_names) > 1: - raise ValueError( - f"Multiple compute resources are not supported yet ({', '.join(sorted(compute_names))})." - ) + raise ValueError(f"Multiple compute resources are not supported yet ({', '.join(sorted(compute_names))}).") sole_compute = next(iter(compute_names)) @@ -163,7 +161,9 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": service = self.services[self.driver.policy_model] if isinstance(service, BaseModelServiceConfig): for bench_name, benchmark in self.driver.benchmarks.items(): - conflicts = [k for k in ("policy_base_url", "policy_model_name", "policy_api_key") if k in benchmark.run] + conflicts = [ + k for k in ("policy_base_url", "policy_model_name", "policy_api_key") if k in benchmark.run + ] if conflicts: raise ValueError( f"Benchmark '{bench_name}' run config already sets {conflicts} " diff --git a/nemo_gym/orchestration/executors/connection.py b/nemo_gym/orchestration/executors/connection.py index 77949ee7308..3282cd28710 100644 --- a/nemo_gym/orchestration/executors/connection.py +++ b/nemo_gym/orchestration/executors/connection.py @@ -68,9 +68,12 @@ def __enter__(self) -> "SSHConnection": [ "ssh", *self._OPTS, - "-o", "ControlMaster=yes", - "-o", f"ControlPath={self._socket}", - "-o", "ControlPersist=yes", + "-o", + "ControlMaster=yes", + "-o", + f"ControlPath={self._socket}", + "-o", + "ControlPersist=yes", "-N", # no remote command — just keep the tunnel open self._hostname, ], @@ -85,7 +88,7 @@ def _wait_for_socket(self) -> None: if self._socket.exists(): return if self._master and self._master.poll() is not None: - stderr = (self._master.stderr.read().decode(errors="replace") if self._master.stderr else "") + stderr = self._master.stderr.read().decode(errors="replace") if self._master.stderr else "" raise RuntimeError( f"SSH master to '{self._hostname}' exited with code {self._master.returncode}.\n{stderr.strip()}" ) @@ -104,8 +107,11 @@ def copy(self, local: Path, remote: Path) -> None: ) _checked( [ - "rsync", "-az", "--delete", - "-e", f"ssh {' '.join(self._ssh_opts())}", + "rsync", + "-az", + "--delete", + "-e", + f"ssh {' '.join(self._ssh_opts())}", f"{local}/", f"{self._hostname}:{remote}", ], @@ -137,9 +143,7 @@ def _checked(cmd: list[str], *, input: str | None = None, context: str = "") -> if result.returncode != 0: label = f" ({context})" if context else "" raise RuntimeError( - f"Command failed{label} with exit code {result.returncode}:\n" - f" {' '.join(cmd)}\n" - f"{result.stderr.strip()}" + f"Command failed{label} with exit code {result.returncode}:\n {' '.join(cmd)}\n{result.stderr.strip()}" ) return result.stdout diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index a941c7b26c3..b9d62880330 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -16,6 +16,7 @@ import re import shlex + _HEALTH_WAIT_MULTI = """\ # Wait for {name} (try multiple health endpoints) echo "Waiting for {name} at {url}..." diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 9213c6a1063..88ebc00e8b7 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -26,6 +26,7 @@ from nemo_gym.orchestration.executors.connection import get_connection from nemo_gym.orchestration.executors.slurm_script import build_sbatch_script + _SBATCH_JOB_ID_RE = re.compile(r"Submitted batch job (\d+)") @@ -50,25 +51,27 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: staging = self._stage(config, compute, remote_run_dir, Path(staging_str)) with get_connection(compute.hostname) as conn: conn.copy(staging, remote_run_dir) - output = conn.run([ - f"sbatch {shlex.quote(str(remote_run_dir / name / 'job.sh'))}" - for name in config.driver.benchmarks - ]) + output = conn.run( + [ + f"sbatch {shlex.quote(str(remote_run_dir / name / 'job.sh'))}" + for name in config.driver.benchmarks + ] + ) benchmark_names = list(config.driver.benchmarks) job_ids = _SBATCH_JOB_ID_RE.findall(output) for name, job_id in zip(benchmark_names, job_ids): rich.print(f"[green]submitted[/green] {name} → Slurm job [bold]{job_id}[/bold]") - for name in benchmark_names[len(job_ids):]: + for name in benchmark_names[len(job_ids) :]: rich.print(f"[green]submitted[/green] {name} (job ID unavailable)") def _dry_run(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path) -> None: print(f"[dry-run] remote run dir: {remote_run_dir}") for name, benchmark in config.driver.benchmarks.items(): script = build_sbatch_script(config, name, benchmark, compute, remote_run_dir / name) - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"[dry-run] sbatch script for benchmark: {name}") - print(f"{'='*60}") + print(f"{'=' * 60}") print(script) def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_dir: Path, staging: Path) -> Path: diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 34e39d74e74..c724761f002 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -16,8 +16,6 @@ import shlex from pathlib import Path -from nemo_gym.orchestration.executors.script_templates import bash_var, render_driver_entrypoint, render_gym_cmd, render_health_check -from nemo_gym.orchestration.executors.utils import flatten_run_args from nemo_gym.orchestration.api import ( BenchmarkRunConfig, NodePool, @@ -26,6 +24,14 @@ SubmitConfig, VllmServiceConfig, # used in _BUILDERS dispatch table ) +from nemo_gym.orchestration.executors.script_templates import ( + bash_var, + render_driver_entrypoint, + render_gym_cmd, + render_health_check, +) +from nemo_gym.orchestration.executors.utils import flatten_run_args + _SCRIPT_TEMPLATE = """\ #!/bin/bash diff --git a/nemo_gym/orchestration/submit.py b/nemo_gym/orchestration/submit.py index e17bdc96e24..b15f778039d 100644 --- a/nemo_gym/orchestration/submit.py +++ b/nemo_gym/orchestration/submit.py @@ -17,6 +17,7 @@ from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.slurm import SlurmExecutor + _EXECUTORS = { SlurmComputeConfig: SlurmExecutor, } diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 78168f1f964..d51bd520aed 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -18,6 +18,7 @@ from nemo_gym.orchestration.api import SubmitConfig + COMPUTE = {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}} COMPUTE_TWO = { "cluster_a": {"type": "slurm", "account": "my-account", "hostname": "foo"}, @@ -73,7 +74,7 @@ def test_policy_model_injects_run_args(): benchmark = config.driver.benchmarks["gsm8k"] assert benchmark.run["policy_base_url"] == "http://localhost:8000/v1" assert benchmark.run["policy_model_name"] == "org/model" - assert benchmark.run["policy_api_key"] == "dummy" # pragma: allowlist secret + assert benchmark.run["policy_api_key"] == "dummy" # pragma: allowlist secret def test_policy_model_conflict_raises(): diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 30b81270e2c..8a5b09f48b1 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -26,10 +26,12 @@ ) from nemo_gym.orchestration.executors.utils import flatten_run_args as _flatten_run_args + # --------------------------------------------------------------------------- # flatten_run_args # --------------------------------------------------------------------------- + def test_scalar_values(): assert _flatten_run_args({"temperature": 0.05, "top_p": 0.9}) == [ "+temperature=0.05", @@ -66,6 +68,7 @@ def test_deeply_nested(): # _render_pool_directives # --------------------------------------------------------------------------- + def test_render_pool_directives_basic(pool): lines = _render_pool_directives("main", pool) assert "#SBATCH --partition=batch # pool: main" in lines @@ -89,6 +92,7 @@ def test_render_pool_directives_extra_args(pool): # _render_directives # --------------------------------------------------------------------------- + def test_render_directives_job_name(compute, bench_dir): out = _render_directives(compute, bench_dir, "gsm8k") assert "#SBATCH --job-name=gym-gsm8k" in out @@ -120,6 +124,7 @@ def test_render_directives_chdir(compute, bench_dir): # _render_service_command # --------------------------------------------------------------------------- + def test_render_service_command_contains_srun(): out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model") assert "srun --overlap --no-container-mount-home" in out @@ -142,6 +147,7 @@ def test_render_service_command_log_file(): # _build_vllm_command # --------------------------------------------------------------------------- + def test_build_vllm_command_basic(vllm_service): cmd = _build_vllm_command(vllm_service) assert "vllm serve" in cmd @@ -164,6 +170,7 @@ def test_build_vllm_command_no_trust_remote_code_by_default(vllm_service): # render_gym_cmd # --------------------------------------------------------------------------- + def test_render_gym_cmd_subcommand(): out = render_gym_cmd("eval run", "GYM_CMD", ["+foo=bar"]) assert out.startswith("GYM_CMD=(") @@ -181,6 +188,7 @@ def test_render_gym_cmd_prepare(): # render_driver_entrypoint # --------------------------------------------------------------------------- + def test_render_driver_entrypoint_no_install_no_prepare(): out = render_driver_entrypoint(None, None, None) assert out == '"${GYM_CMD[@]}"' @@ -213,6 +221,7 @@ def test_render_driver_entrypoint_install_and_prepare(): # build_sbatch_script (integration) # --------------------------------------------------------------------------- + def test_build_sbatch_script_contains_shebang(submit_config, bench_dir): benchmark = submit_config.driver.benchmarks["gsm8k"] compute = next(iter(submit_config.compute.values())) @@ -257,6 +266,7 @@ def test_build_sbatch_script_policy_model_flags(submit_config_with_policy, bench # --------------------------------------------------------------------------- import pytest + from nemo_gym.orchestration.api import NodePool, SlurmComputeConfig, VllmServiceConfig @@ -282,19 +292,23 @@ def vllm_service(): @pytest.fixture def submit_config(): - return SubmitConfig.model_validate({ - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - }) + return SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) @pytest.fixture def submit_config_with_policy(): - return SubmitConfig.model_validate({ - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "policy_model": "vllm_model", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - }) + return SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "policy_model": "vllm_model", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) From be3a1a0e4587f2ee729e915da809a54a6545e153 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:17:55 +0200 Subject: [PATCH 34/66] fix ignored health path Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/executors/script_templates.py | 5 +++-- nemo_gym/orchestration/executors/slurm_script.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index b9d62880330..bd1d2694e73 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -22,7 +22,7 @@ echo "Waiting for {name} at {url}..." {name_upper}_READY=0 for _i in $(seq 1 {max_attempts}); do - if curl -sf "{url}/health" > /dev/null 2>&1 || curl -sf "{url}/openapi.json" > /dev/null 2>&1; then + if curl -sf "{url}{path}" > /dev/null 2>&1; then echo " {name} ready." {name_upper}_READY=1 break @@ -44,11 +44,12 @@ def bash_var(name: str) -> str: return re.sub(r"[^A-Z0-9]", "_", name.upper()) -def render_health_check(name: str, port: int, timeout: int) -> str: +def render_health_check(name: str, port: int, path: str, timeout: int) -> str: return _HEALTH_WAIT_MULTI.format( name=name, name_upper=bash_var(name), url=f"http://localhost:{port}", + path=path, max_attempts=timeout // 5, ) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index c724761f002..f8b084cad93 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -119,7 +119,7 @@ def build_sbatch_script( ) health_checks = "\n\n".join( - render_health_check(name, service.health_check.port, service.health_check.timeout_seconds) + render_health_check(name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds) for name, service in config.services.items() if service.health_check ) From 2e263eac67d33572dea29417f703923968e269f8 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:20:42 +0200 Subject: [PATCH 35/66] ignore dummy api key Signed-off-by: Tomasz Grzegorzek --- nemo_gym/orchestration/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 933c7bc103a..0a06754d4e0 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -172,6 +172,6 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": benchmark.run["policy_base_url"] = f"http://localhost:{service.port}/v1" benchmark.run["policy_model_name"] = service.model # vLLM doesn't require auth; dummy key satisfies clients that require the header. - benchmark.run["policy_api_key"] = "dummy" + benchmark.run["policy_api_key"] = "dummy" # pragma: allowlist secret return self From dd335c32cbd25e8ecc39d282b5d419b81834d31d Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 00:22:50 +0200 Subject: [PATCH 36/66] linting Signed-off-by: Tomasz Grzegorzek --- nemo_gym/cli/eval.py | 4 ---- nemo_gym/cli/main.py | 10 +++++++--- nemo_gym/orchestration/__init__.py | 2 -- nemo_gym/orchestration/api.py | 2 +- nemo_gym/orchestration/executors/__init__.py | 2 -- nemo_gym/orchestration/executors/slurm_script.py | 4 +++- 6 files changed, 11 insertions(+), 13 deletions(-) diff --git a/nemo_gym/cli/eval.py b/nemo_gym/cli/eval.py index ef4c76b68e7..1c1f2987079 100644 --- a/nemo_gym/cli/eval.py +++ b/nemo_gym/cli/eval.py @@ -15,10 +15,7 @@ import asyncio import importlib import json -import sys from copy import deepcopy -from dataclasses import dataclass -from glob import glob from multiprocessing import Pool from pathlib import Path from typing import Any, Dict, List, Tuple @@ -479,4 +476,3 @@ def reward_profile(): # pragma: no cover Input rows: {completion_summary["total_input_rows"]} total; {completion_summary["complete_input_rows"]} complete; {completion_summary["partial_input_rows"]} partial; {completion_summary["missing_input_rows"]} without rollouts dropped from output. Reward profiling outputs: {reward_profiling_fpath} Agent-level metrics: {agent_level_metrics_fpath}""") - diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index ba6dfa71367..e194b314edd 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -306,7 +306,6 @@ def _asset_selector(flag: str) -> Flag: ) - def _merge_config_paths(overrides: list[str]) -> list[str]: """Coalesce all `+config_paths=[...]` tokens (from --config and asset selectors) into one (Hydra rejects dupes).""" prefix = "+config_paths=[" @@ -322,6 +321,7 @@ def _merge_config_paths(overrides: list[str]) -> list[str]: def _eval_submit(args: argparse.Namespace, overrides: list[str]) -> None: from omegaconf import OmegaConf + from nemo_gym.orchestration.api import SubmitConfig from nemo_gym.orchestration.submit import submit @@ -668,10 +668,14 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: summary="Submit a job.", flags=( Flag( - register=lambda p: p.add_argument("--config", "-c", required=True, metavar="PATH", help="Submit config YAML file."), + register=lambda p: p.add_argument( + "--config", "-c", required=True, metavar="PATH", help="Submit config YAML file." + ), ), Flag( - register=lambda p: p.add_argument("--dry-run", action="store_true", help="Print generated job scripts without submitting."), + register=lambda p: p.add_argument( + "--dry-run", action="store_true", help="Print generated job scripts without submitting." + ), ), ), ), diff --git a/nemo_gym/orchestration/__init__.py b/nemo_gym/orchestration/__init__.py index 2c5a3a8a54b..467079831e1 100644 --- a/nemo_gym/orchestration/__init__.py +++ b/nemo_gym/orchestration/__init__.py @@ -12,5 +12,3 @@ # 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. - - diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 0a06754d4e0..98d25f2937b 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -172,6 +172,6 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": benchmark.run["policy_base_url"] = f"http://localhost:{service.port}/v1" benchmark.run["policy_model_name"] = service.model # vLLM doesn't require auth; dummy key satisfies clients that require the header. - benchmark.run["policy_api_key"] = "dummy" # pragma: allowlist secret + benchmark.run["policy_api_key"] = "dummy" # pragma: allowlist secret return self diff --git a/nemo_gym/orchestration/executors/__init__.py b/nemo_gym/orchestration/executors/__init__.py index 2c5a3a8a54b..467079831e1 100644 --- a/nemo_gym/orchestration/executors/__init__.py +++ b/nemo_gym/orchestration/executors/__init__.py @@ -12,5 +12,3 @@ # 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. - - diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index f8b084cad93..07072821cab 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -119,7 +119,9 @@ def build_sbatch_script( ) health_checks = "\n\n".join( - render_health_check(name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds) + render_health_check( + name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds + ) for name, service in config.services.items() if service.health_check ) From 4afa47b32c3831c7259b6633966be4e4e101f1a5 Mon Sep 17 00:00:00 2001 From: Tomasz Grzegorzek Date: Thu, 30 Jul 2026 01:14:28 +0200 Subject: [PATCH 37/66] add exemplary config for orchestration Signed-off-by: Tomasz Grzegorzek --- examples/slurm_vllm_1IN.yaml | 46 ++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 examples/slurm_vllm_1IN.yaml diff --git a/examples/slurm_vllm_1IN.yaml b/examples/slurm_vllm_1IN.yaml new file mode 100644 index 00000000000..a593d4200d6 --- /dev/null +++ b/examples/slurm_vllm_1IN.yaml @@ -0,0 +1,46 @@ +compute: + cluster: + type: slurm + hostname: myslurmcluster + walltime: "00:30:00" + account: myaccount + node_pools: + compute: + partition: batch + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 4 + +services: + vllm_model: + container: vllm/vllm-openai:v0.26.0 # optional, defaults + type: vllm + model: Qwen/Qwen2.5-0.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + port: 8000 + health_check: + timeout_seconds: 1200 + +driver: + policy_model: vllm_model # synthactic sugar + container: python:3.12 # optional, defaults to the service container + gym_install: + ref: main + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/ifbench/config.yaml + run: # gym eval run + split: benchmark + overwrite_metrics_conflicts: true + responses_create_params: + temperature: 0.6 + top_p: 0.9 + config_paths: # if services include pre-defined service types, some configs might be created and added to the command automatically + - benchmarks/ifbench/config.yaml +job: + output_path: /lustre/fsw/my-path + From 4391efbac2352d830688726858f05a86cf9c031c Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 5 Aug 2026 13:58:11 -0400 Subject: [PATCH 38/66] Add env variable injection support Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_1IN.yaml | 8 +- nemo_gym/orchestration/api.py | 23 ++ .../orchestration/executors/slurm_script.py | 14 +- tests/unit_tests/test_connection.py | 111 ++++++++ tests/unit_tests/test_orchestration_api.py | 15 ++ tests/unit_tests/test_slurm_script.py | 243 ++++++++++++++++++ 6 files changed, 408 insertions(+), 6 deletions(-) create mode 100644 tests/unit_tests/test_connection.py diff --git a/examples/slurm_vllm_1IN.yaml b/examples/slurm_vllm_1IN.yaml index a593d4200d6..5a9c6b49cba 100644 --- a/examples/slurm_vllm_1IN.yaml +++ b/examples/slurm_vllm_1IN.yaml @@ -1,9 +1,9 @@ compute: cluster: type: slurm - hostname: myslurmcluster + hostname: cw-dfw-cs-001-login-01.cw-dfw-cs-001.hpc.nvidia.com walltime: "00:30:00" - account: myaccount + account: coreai_dlalgo_nemorl node_pools: compute: partition: batch @@ -17,7 +17,7 @@ services: type: vllm model: Qwen/Qwen2.5-0.5B-Instruct trust_remote_code: true - tensor_parallel_size: 1 + tensor_parallel_size: $TP pipeline_parallel_size: 1 port: 8000 health_check: @@ -42,5 +42,5 @@ driver: config_paths: # if services include pre-defined service types, some configs might be created and added to the command automatically - benchmarks/ifbench/config.yaml job: - output_path: /lustre/fsw/my-path + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/orch_tests diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 98d25f2937b..72502a0f6ed 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator @@ -35,6 +36,8 @@ class BaseServiceConfig(_StrictModel): # Resolved to the sole compute resource name at validation time when not set. placement: str | None = None health_check: HealthCheckConfig | None = None + # Values starting with "$" are resolved from the host environment at submit time. + env: dict[str, str] = {} class BaseModelServiceConfig(BaseServiceConfig): @@ -121,6 +124,8 @@ class DriverConfig(_StrictModel): # policy_base_url/policy_model_name/policy_api_key into each benchmark's run config. policy_model: str | None = None benchmarks: dict[str, BenchmarkRunConfig] + # Values starting with "$" are resolved from the host environment at submit time. + env: dict[str, str] = {} class JobConfig(_StrictModel): @@ -128,12 +133,30 @@ class JobConfig(_StrictModel): output_path: str +def _resolve_env_refs(data: Any) -> Any: + if isinstance(data, dict): + return {k: _resolve_env_refs(v) for k, v in data.items()} + if isinstance(data, list): + return [_resolve_env_refs(v) for v in data] + if isinstance(data, str) and data.startswith("$"): + host_var = data[1:] + if host_var not in os.environ: + raise ValueError(f"Host environment variable {host_var!r} is not set (referenced as {data!r} in config).") + return os.environ[host_var] + return data + + class SubmitConfig(_StrictModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] driver: DriverConfig job: JobConfig + @model_validator(mode="before") + @classmethod + def _resolve_host_env_vars(cls, data: Any) -> Any: + return _resolve_env_refs(data) + @model_validator(mode="after") def _resolve_and_validate_placements(self) -> "SubmitConfig": compute_names = set(self.compute) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 07072821cab..8643b90ebae 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -75,13 +75,20 @@ def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: return lines -def _render_service_command(name: str, container: str, command: str) -> str: +def _resolve_env(env: dict[str, str]) -> str: + return "\n".join(f"export {k}={shlex.quote(v)}" for k, v in env.items()) + + +def _render_service_command(name: str, container: str, command: str, env: dict[str, str] | None = None) -> str: var = bash_var(name) + env_block = _resolve_env(env) if env else "" + prefix = f"{env_block}\n" if env_block else "" # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" + f"{prefix}" f"srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -114,7 +121,7 @@ def build_sbatch_script( directives = _render_directives(compute, remote_bench_dir, benchmark_name) service_commands = "\n\n".join( - _render_service_command(name, service.container, _BUILDERS[type(service)](service)) + _render_service_command(name, service.container, _BUILDERS[type(service)](service), service.env or None) for name, service in config.services.items() ) @@ -141,8 +148,11 @@ def build_sbatch_script( prepare_cmd=prepare_cmd, ) prepare_command = "" + driver_env_block = _resolve_env(config.driver.env) if config.driver.env else "" + driver_env_prefix = f"{driver_env_block}\n" if driver_env_block else "" driver_command = ( f"{gym_cmd}\n" + f"{driver_env_prefix}" f"srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/tests/unit_tests/test_connection.py b/tests/unit_tests/test_connection.py new file mode 100644 index 00000000000..fe67bd96a5e --- /dev/null +++ b/tests/unit_tests/test_connection.py @@ -0,0 +1,111 @@ +# 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 socket +import subprocess +from unittest.mock import MagicMock, patch + +import pytest + +from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.executors.connection import LocalConnection, SSHConnection, get_connection +from nemo_gym.orchestration.executors.slurm import SlurmExecutor + + +# --------------------------------------------------------------------------- +# get_connection routing +# --------------------------------------------------------------------------- + + +def test_get_connection_none_returns_local(): + assert isinstance(get_connection(None), LocalConnection) + + +def test_get_connection_own_hostname_returns_local(): + assert isinstance(get_connection(socket.gethostname()), LocalConnection) + + +def test_get_connection_remote_hostname_returns_ssh(): + assert isinstance(get_connection("remote-login-node.example.com"), SSHConnection) + + +# --------------------------------------------------------------------------- +# SlurmExecutor on the login node makes no SSH calls +# --------------------------------------------------------------------------- + + +@pytest.fixture +def login_node_config(tmp_path): + """SubmitConfig with hostname=None (already on the login node).""" + return SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account"}}, # hostname omitted → None + "driver": { + "container": "python:3.12", + "policy_model": "vllm_model", + "benchmarks": {"gsm8k": {}}, + }, + "job": {"output_path": str(tmp_path / "gym-jobs")}, + } + ) + + +def test_slurm_executor_login_node_no_ssh(login_node_config): + """When hostname is None (login node), SlurmExecutor must not open any SSH connection.""" + ssh_calls = [] + + def fake_subprocess_run(cmd, **kwargs): + if isinstance(cmd, list) and cmd and cmd[0] == "ssh": + ssh_calls.append(cmd) + result = MagicMock(spec=subprocess.CompletedProcess) + result.returncode = 0 + result.stdout = "Submitted batch job 99999\n" + result.stderr = "" + return result + + with patch("nemo_gym.orchestration.executors.connection.subprocess.run", side_effect=fake_subprocess_run): + SlurmExecutor().run(login_node_config) + + assert ssh_calls == [], f"Unexpected SSH subprocess calls on login node: {ssh_calls}" + + +def test_slurm_executor_login_node_uses_local_connection(login_node_config): + """SlurmExecutor must resolve to LocalConnection (not SSHConnection) when hostname is None.""" + connections_created = [] + real_get_connection = get_connection + + def spy_get_connection(hostname): + conn = real_get_connection(hostname) + connections_created.append(conn) + return conn + + def fake_subprocess_run(cmd, **kwargs): + result = MagicMock(spec=subprocess.CompletedProcess) + result.returncode = 0 + result.stdout = "Submitted batch job 99999\n" + result.stderr = "" + return result + + with ( + patch("nemo_gym.orchestration.executors.slurm.get_connection", side_effect=spy_get_connection), + patch("nemo_gym.orchestration.executors.connection.subprocess.run", side_effect=fake_subprocess_run), + ): + SlurmExecutor().run(login_node_config) + + assert len(connections_created) == 1 + assert isinstance(connections_created[0], LocalConnection), ( + f"Expected LocalConnection for login node, got {type(connections_created[0]).__name__}" + ) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index d51bd520aed..2b284c71976 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -81,3 +81,18 @@ def test_policy_model_conflict_raises(): driver = {**DRIVER, "policy_model": "svc", "benchmarks": {"gsm8k": {"run": {"policy_base_url": "http://other"}}}} with pytest.raises(ValidationError, match="already sets"): SubmitConfig.model_validate(_config(driver=driver)) + + +def test_service_env_accepted(): + config = SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "env": {"FOO": "bar"}}})) + assert config.services["svc"].env == {"FOO": "bar"} + + +def test_driver_env_accepted(): + config = SubmitConfig.model_validate(_config(driver={**DRIVER, "env": {"KEY": "val"}})) + assert config.driver.env == {"KEY": "val"} + + +def test_service_unknown_field_raises(): + with pytest.raises(ValidationError): + SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "unknown_field": "x"}})) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 8a5b09f48b1..01b829bb02e 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -22,8 +22,10 @@ _render_directives, _render_pool_directives, _render_service_command, + _resolve_env, build_sbatch_script, ) +from nemo_gym.orchestration.api import _resolve_env_refs from nemo_gym.orchestration.executors.utils import flatten_run_args as _flatten_run_args @@ -261,6 +263,247 @@ def test_build_sbatch_script_policy_model_flags(submit_config_with_policy, bench assert "+policy_model_name=" in script +# --------------------------------------------------------------------------- +# _resolve_env +# --------------------------------------------------------------------------- + + +def test_resolve_env_literal(): + out = _resolve_env({"FOO": "bar", "BAZ": "qux"}) + assert "export FOO=bar" in out + assert "export BAZ=qux" in out + + +def test_resolve_env_value_with_spaces(): + out = _resolve_env({"MSG": "hello world"}) + assert "export MSG='hello world'" in out + + +def test_resolve_env_empty(): + assert _resolve_env({}) == "" + + +# --------------------------------------------------------------------------- +# _resolve_env_refs (pre-validation env var expansion) +# --------------------------------------------------------------------------- + + +def test_resolve_env_refs_literal(): + assert _resolve_env_refs({"key": "value"}) == {"key": "value"} + + +def test_resolve_env_refs_host_var(monkeypatch): + monkeypatch.setenv("MY_TOKEN", "secret123") + assert _resolve_env_refs({"TOKEN": "$MY_TOKEN"}) == {"TOKEN": "secret123"} + + +def test_resolve_env_refs_missing_host_var(monkeypatch): + monkeypatch.delenv("UNSET_VAR", raising=False) + with pytest.raises(ValueError, match="'UNSET_VAR' is not set"): + _resolve_env_refs({"KEY": "$UNSET_VAR"}) + + +def test_resolve_env_refs_nested(monkeypatch): + monkeypatch.setenv("TP", "4") + result = _resolve_env_refs({"service": {"tensor_parallel_size": "$TP"}}) + assert result == {"service": {"tensor_parallel_size": "4"}} + + +def test_resolve_env_refs_list(monkeypatch): + monkeypatch.setenv("VAL", "x") + assert _resolve_env_refs(["$VAL", "literal"]) == ["x", "literal"] + + +def test_resolve_env_refs_non_dollar_string_unchanged(): + assert _resolve_env_refs({"key": "http://localhost:8000/v1"}) == {"key": "http://localhost:8000/v1"} + + +def test_resolve_env_refs_dollar_not_at_start_unchanged(): + # Only leading $ triggers resolution — embedded $ is left as-is. + assert _resolve_env_refs({"key": "prefix_$VAR"}) == {"key": "prefix_$VAR"} + + +def test_resolve_env_refs_error_message_contains_ref(monkeypatch): + monkeypatch.delenv("MISSING", raising=False) + with pytest.raises(ValueError, match=r"\$MISSING"): + _resolve_env_refs({"key": "$MISSING"}) + + +def test_resolve_env_refs_non_string_values_unchanged(): + assert _resolve_env_refs({"n": 42, "b": True}) == {"n": 42, "b": True} + + +def test_submit_config_resolves_int_field(monkeypatch): + monkeypatch.setenv("TP", "2") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "tensor_parallel_size": "$TP"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + assert config.services["vllm_model"].tensor_parallel_size == 2 + + +def test_submit_config_resolves_string_fields(monkeypatch): + monkeypatch.setenv("MY_MODEL", "org/llama-3") + monkeypatch.setenv("MY_CONTAINER", "nvcr.io/my-vllm:latest") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "$MY_CONTAINER", "model": "$MY_MODEL"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + assert config.services["vllm_model"].model == "org/llama-3" + assert config.services["vllm_model"].container == "nvcr.io/my-vllm:latest" + + +def test_submit_config_resolves_compute_account(monkeypatch): + monkeypatch.setenv("SLURM_ACCOUNT", "my_team") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "$SLURM_ACCOUNT", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + assert config.compute["cluster"].account == "my_team" + + +def test_submit_config_resolves_output_path(monkeypatch): + monkeypatch.setenv("JOB_DIR", "/lustre/fsw/my-jobs") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "$JOB_DIR"}, + } + ) + assert config.job.output_path == "/lustre/fsw/my-jobs" + + +def test_submit_config_resolves_missing_var_raises(monkeypatch): + from pydantic import ValidationError + + monkeypatch.delenv("UNSET_MODEL", raising=False) + with pytest.raises(ValidationError, match="UNSET_MODEL"): + SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "$UNSET_MODEL"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + + +def test_submit_config_resolves_benchmark_run_args(monkeypatch): + monkeypatch.setenv("TEMP", "0.7") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {"run": {"temperature": "$TEMP"}}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + assert config.driver.benchmarks["gsm8k"].run["temperature"] == "0.7" + + +def test_build_sbatch_script_resolved_tp_in_vllm_cmd(bench_dir, monkeypatch): + monkeypatch.setenv("TP", "8") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "tensor_parallel_size": "$TP"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "--tensor-parallel-size 8" in script + + +def test_build_sbatch_script_service_env_before_driver_env(bench_dir, monkeypatch): + monkeypatch.setenv("SVC_TOKEN", "svc_val") + monkeypatch.setenv("DRV_TOKEN", "drv_val") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "env": {"SVC_KEY": "$SVC_TOKEN"}}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}, "env": {"DRV_KEY": "$DRV_TOKEN"}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + svc_export_idx = script.index("export SVC_KEY=svc_val") + drv_export_idx = script.index("export DRV_KEY=drv_val") + svc_srun_idx = script.index("srun --overlap --no-container-mount-home --container-image=vllm:latest") + drv_srun_idx = script.index("srun --overlap --no-container-mount-home --container-image=python:3.12") + # Service export appears before service srun; driver export appears before driver srun. + assert svc_export_idx < svc_srun_idx + assert drv_export_idx < drv_srun_idx + # Service export appears before driver export. + assert svc_export_idx < drv_export_idx + + +def test_render_service_command_with_env(): + out = _render_service_command("svc", "img:latest", "cmd", {"FOO": "bar"}) + lines = out.splitlines() + export_idx = next(i for i, l in enumerate(lines) if "export FOO=bar" in l) + srun_idx = next(i for i, l in enumerate(lines) if "srun" in l) + assert export_idx < srun_idx + + +def test_render_service_command_no_env(): + out = _render_service_command("svc", "img:latest", "cmd") + assert "export" not in out + + +def test_build_sbatch_script_service_env(bench_dir, monkeypatch): + monkeypatch.setenv("HF_TOKEN", "hf_test") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "env": {"HF_TOKEN": "$HF_TOKEN", "LIT": "val"}}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + # $HF_TOKEN resolved at parse time — env dict now has literal values + assert config.services["vllm_model"].env == {"HF_TOKEN": "hf_test", "LIT": "val"} + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "export HF_TOKEN=hf_test" in script + assert "export LIT=val" in script + + +def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): + monkeypatch.setenv("WANDB_KEY", "wb_secret") + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}, "env": {"WANDB_API_KEY": "$WANDB_KEY"}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "export WANDB_API_KEY=wb_secret" in script + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From b4455769b8a3edcd5494984ad6c05dbe41a0f3ad Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 5 Aug 2026 14:12:38 -0400 Subject: [PATCH 39/66] Fix merge conflicts Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 15 --------------- .../orchestration/executors/slurm_script.py | 17 ----------------- 2 files changed, 32 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 567c296c04b..72502a0f6ed 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -13,10 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -<<<<<<< HEAD import os -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator @@ -39,11 +36,8 @@ class BaseServiceConfig(_StrictModel): # Resolved to the sole compute resource name at validation time when not set. placement: str | None = None health_check: HealthCheckConfig | None = None -<<<<<<< HEAD # Values starting with "$" are resolved from the host environment at submit time. env: dict[str, str] = {} -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 class BaseModelServiceConfig(BaseServiceConfig): @@ -130,11 +124,8 @@ class DriverConfig(_StrictModel): # policy_base_url/policy_model_name/policy_api_key into each benchmark's run config. policy_model: str | None = None benchmarks: dict[str, BenchmarkRunConfig] -<<<<<<< HEAD # Values starting with "$" are resolved from the host environment at submit time. env: dict[str, str] = {} -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 class JobConfig(_StrictModel): @@ -142,7 +133,6 @@ class JobConfig(_StrictModel): output_path: str -<<<<<<< HEAD def _resolve_env_refs(data: Any) -> Any: if isinstance(data, dict): return {k: _resolve_env_refs(v) for k, v in data.items()} @@ -156,22 +146,17 @@ def _resolve_env_refs(data: Any) -> Any: return data -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 class SubmitConfig(_StrictModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] driver: DriverConfig job: JobConfig -<<<<<<< HEAD @model_validator(mode="before") @classmethod def _resolve_host_env_vars(cls, data: Any) -> Any: return _resolve_env_refs(data) -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 @model_validator(mode="after") def _resolve_and_validate_placements(self) -> "SubmitConfig": compute_names = set(self.compute) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 3d54c3b9003..8643b90ebae 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -75,7 +75,6 @@ def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: return lines -<<<<<<< HEAD def _resolve_env(env: dict[str, str]) -> str: return "\n".join(f"export {k}={shlex.quote(v)}" for k, v in env.items()) @@ -84,19 +83,12 @@ def _render_service_command(name: str, container: str, command: str, env: dict[s var = bash_var(name) env_block = _resolve_env(env) if env else "" prefix = f"{env_block}\n" if env_block else "" -======= -def _render_service_command(name: str, container: str, command: str) -> str: - var = bash_var(name) ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" -<<<<<<< HEAD f"{prefix}" -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 f"srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -129,11 +121,7 @@ def build_sbatch_script( directives = _render_directives(compute, remote_bench_dir, benchmark_name) service_commands = "\n\n".join( -<<<<<<< HEAD _render_service_command(name, service.container, _BUILDERS[type(service)](service), service.env or None) -======= - _render_service_command(name, service.container, _BUILDERS[type(service)](service)) ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 for name, service in config.services.items() ) @@ -160,16 +148,11 @@ def build_sbatch_script( prepare_cmd=prepare_cmd, ) prepare_command = "" -<<<<<<< HEAD driver_env_block = _resolve_env(config.driver.env) if config.driver.env else "" driver_env_prefix = f"{driver_env_block}\n" if driver_env_block else "" driver_command = ( f"{gym_cmd}\n" f"{driver_env_prefix}" -======= - driver_command = ( - f"{gym_cmd}\n" ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 f"srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) From b3aba0c1f71228d661af9bed8e0d09a45af407ee Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 5 Aug 2026 14:14:27 -0400 Subject: [PATCH 40/66] Fix merge conflicts Signed-off-by: Onur Yilmaz --- tests/unit_tests/test_orchestration_api.py | 3 --- tests/unit_tests/test_slurm_script.py | 8 -------- 2 files changed, 11 deletions(-) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index eaa31f5c869..2b284c71976 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -81,7 +81,6 @@ def test_policy_model_conflict_raises(): driver = {**DRIVER, "policy_model": "svc", "benchmarks": {"gsm8k": {"run": {"policy_base_url": "http://other"}}}} with pytest.raises(ValidationError, match="already sets"): SubmitConfig.model_validate(_config(driver=driver)) -<<<<<<< HEAD def test_service_env_accepted(): @@ -97,5 +96,3 @@ def test_driver_env_accepted(): def test_service_unknown_field_raises(): with pytest.raises(ValidationError): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "unknown_field": "x"}})) -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 91ae4671c6a..01b829bb02e 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -22,15 +22,10 @@ _render_directives, _render_pool_directives, _render_service_command, -<<<<<<< HEAD _resolve_env, build_sbatch_script, ) from nemo_gym.orchestration.api import _resolve_env_refs -======= - build_sbatch_script, -) ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 from nemo_gym.orchestration.executors.utils import flatten_run_args as _flatten_run_args @@ -269,7 +264,6 @@ def test_build_sbatch_script_policy_model_flags(submit_config_with_policy, bench # --------------------------------------------------------------------------- -<<<<<<< HEAD # _resolve_env # --------------------------------------------------------------------------- @@ -511,8 +505,6 @@ def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): # --------------------------------------------------------------------------- -======= ->>>>>>> fcca3a82814101d5234f5994b8d8c8d34e2d1749 # Fixtures # --------------------------------------------------------------------------- From 1479b60a7f79895f2f6d1c63981bfab1a3a0d5ca Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Fri, 7 Aug 2026 12:26:00 -0400 Subject: [PATCH 41/66] Fix pre-commit formatting and lint issues in test_slurm_script.py Reformat long dict literals per ruff line-length limit, consolidate duplicate import, and rename ambiguous loop variable `l` to `line`. Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Onur Yilmaz --- tests/unit_tests/test_slurm_script.py | 49 ++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 01b829bb02e..e5cd39376ac 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -15,7 +15,7 @@ from pathlib import Path -from nemo_gym.orchestration.api import SubmitConfig +from nemo_gym.orchestration.api import SubmitConfig, _resolve_env_refs from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd from nemo_gym.orchestration.executors.slurm_script import ( _build_vllm_command, @@ -25,7 +25,6 @@ _resolve_env, build_sbatch_script, ) -from nemo_gym.orchestration.api import _resolve_env_refs from nemo_gym.orchestration.executors.utils import flatten_run_args as _flatten_run_args @@ -337,7 +336,14 @@ def test_submit_config_resolves_int_field(monkeypatch): monkeypatch.setenv("TP", "2") config = SubmitConfig.model_validate( { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "tensor_parallel_size": "$TP"}}, + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "tensor_parallel_size": "$TP", + } + }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, "job": {"output_path": "/remote/jobs"}, @@ -419,7 +425,14 @@ def test_build_sbatch_script_resolved_tp_in_vllm_cmd(bench_dir, monkeypatch): monkeypatch.setenv("TP", "8") config = SubmitConfig.model_validate( { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "tensor_parallel_size": "$TP"}}, + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "tensor_parallel_size": "$TP", + } + }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, "job": {"output_path": "/remote/jobs"}, @@ -436,7 +449,14 @@ def test_build_sbatch_script_service_env_before_driver_env(bench_dir, monkeypatc monkeypatch.setenv("DRV_TOKEN", "drv_val") config = SubmitConfig.model_validate( { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "env": {"SVC_KEY": "$SVC_TOKEN"}}}, + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "env": {"SVC_KEY": "$SVC_TOKEN"}, + } + }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}, "env": {"DRV_KEY": "$DRV_TOKEN"}}, "job": {"output_path": "/remote/jobs"}, @@ -459,8 +479,8 @@ def test_build_sbatch_script_service_env_before_driver_env(bench_dir, monkeypatc def test_render_service_command_with_env(): out = _render_service_command("svc", "img:latest", "cmd", {"FOO": "bar"}) lines = out.splitlines() - export_idx = next(i for i, l in enumerate(lines) if "export FOO=bar" in l) - srun_idx = next(i for i, l in enumerate(lines) if "srun" in l) + export_idx = next(i for i, line in enumerate(lines) if "export FOO=bar" in line) + srun_idx = next(i for i, line in enumerate(lines) if "srun" in line) assert export_idx < srun_idx @@ -473,7 +493,14 @@ def test_build_sbatch_script_service_env(bench_dir, monkeypatch): monkeypatch.setenv("HF_TOKEN", "hf_test") config = SubmitConfig.model_validate( { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model", "env": {"HF_TOKEN": "$HF_TOKEN", "LIT": "val"}}}, + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "env": {"HF_TOKEN": "$HF_TOKEN", "LIT": "val"}, + } + }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, "job": {"output_path": "/remote/jobs"}, @@ -494,7 +521,11 @@ def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): { "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}, "env": {"WANDB_API_KEY": "$WANDB_KEY"}}, + "driver": { + "container": "python:3.12", + "benchmarks": {"gsm8k": {}}, + "env": {"WANDB_API_KEY": "$WANDB_KEY"}, + }, "job": {"output_path": "/remote/jobs"}, } ) From 2d58049964d91ad8f3645ba63efeade4f3359057 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Fri, 7 Aug 2026 12:51:19 -0400 Subject: [PATCH 42/66] Fix security related issues Signed-off-by: Onur Yilmaz --- .../orchestration/executors/slurm_script.py | 29 +++++++---- tests/unit_tests/test_slurm_script.py | 48 ++++++++++++------- 2 files changed, 51 insertions(+), 26 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 8643b90ebae..54e48ae3f70 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re import shlex from pathlib import Path @@ -75,21 +76,33 @@ def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]: return lines +_VALID_ENV_KEY = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _validate_env_key(key: str) -> None: + if not _VALID_ENV_KEY.match(key): + raise ValueError(f"Invalid environment variable name: {key!r}") + + def _resolve_env(env: dict[str, str]) -> str: - return "\n".join(f"export {k}={shlex.quote(v)}" for k, v in env.items()) + """Return an 'env K=V ...' prefix string (trailing space) scoped to a single command, or '' if empty.""" + if not env: + return "" + for k in env: + _validate_env_key(k) + pairs = " ".join(f"{k}={shlex.quote(v)}" for k, v in env.items()) + return f"env {pairs} " def _render_service_command(name: str, container: str, command: str, env: dict[str, str] | None = None) -> str: var = bash_var(name) - env_block = _resolve_env(env) if env else "" - prefix = f"{env_block}\n" if env_block else "" + env_prefix = _resolve_env(env) if env else "" # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" - f"{prefix}" - f"srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" + f"{env_prefix}srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -148,12 +161,10 @@ def build_sbatch_script( prepare_cmd=prepare_cmd, ) prepare_command = "" - driver_env_block = _resolve_env(config.driver.env) if config.driver.env else "" - driver_env_prefix = f"{driver_env_block}\n" if driver_env_block else "" + driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" driver_command = ( f"{gym_cmd}\n" - f"{driver_env_prefix}" - f"srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " + f"{driver_env_prefix}srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index e5cd39376ac..f06ebea0e94 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -269,19 +269,32 @@ def test_build_sbatch_script_policy_model_flags(submit_config_with_policy, bench def test_resolve_env_literal(): out = _resolve_env({"FOO": "bar", "BAZ": "qux"}) - assert "export FOO=bar" in out - assert "export BAZ=qux" in out + assert "FOO=bar" in out + assert "BAZ=qux" in out + assert out.startswith("env ") def test_resolve_env_value_with_spaces(): out = _resolve_env({"MSG": "hello world"}) - assert "export MSG='hello world'" in out + assert "MSG='hello world'" in out def test_resolve_env_empty(): assert _resolve_env({}) == "" +def test_resolve_env_invalid_key_raises(): + with pytest.raises(ValueError, match="Invalid environment variable name"): + _resolve_env({"X; rm -rf /": "v"}) + + +def test_resolve_env_valid_keys(): + out = _resolve_env({"_VALID_KEY": "a", "key1": "b", "KEY_123": "c"}) + assert "_VALID_KEY=a" in out + assert "key1=b" in out + assert "KEY_123=c" in out + + # --------------------------------------------------------------------------- # _resolve_env_refs (pre-validation env var expansion) # --------------------------------------------------------------------------- @@ -465,23 +478,24 @@ def test_build_sbatch_script_service_env_before_driver_env(bench_dir, monkeypatc benchmark = config.driver.benchmarks["gsm8k"] compute = next(iter(config.compute.values())) script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - svc_export_idx = script.index("export SVC_KEY=svc_val") - drv_export_idx = script.index("export DRV_KEY=drv_val") + svc_env_idx = script.index("SVC_KEY=svc_val") + drv_env_idx = script.index("DRV_KEY=drv_val") svc_srun_idx = script.index("srun --overlap --no-container-mount-home --container-image=vllm:latest") drv_srun_idx = script.index("srun --overlap --no-container-mount-home --container-image=python:3.12") - # Service export appears before service srun; driver export appears before driver srun. - assert svc_export_idx < svc_srun_idx - assert drv_export_idx < drv_srun_idx - # Service export appears before driver export. - assert svc_export_idx < drv_export_idx + # Service env prefix appears before service srun; driver env prefix appears before driver srun. + assert svc_env_idx < svc_srun_idx + assert drv_env_idx < drv_srun_idx + # Service env prefix appears before driver env prefix. + assert svc_env_idx < drv_env_idx def test_render_service_command_with_env(): out = _render_service_command("svc", "img:latest", "cmd", {"FOO": "bar"}) - lines = out.splitlines() - export_idx = next(i for i, line in enumerate(lines) if "export FOO=bar" in line) - srun_idx = next(i for i, line in enumerate(lines) if "srun" in line) - assert export_idx < srun_idx + assert "FOO=bar" in out + # env prefix must appear before srun on the same line or earlier + foo_idx = out.index("FOO=bar") + srun_idx = out.index("srun") + assert foo_idx < srun_idx def test_render_service_command_no_env(): @@ -511,8 +525,8 @@ def test_build_sbatch_script_service_env(bench_dir, monkeypatch): benchmark = config.driver.benchmarks["gsm8k"] compute = next(iter(config.compute.values())) script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - assert "export HF_TOKEN=hf_test" in script - assert "export LIT=val" in script + assert "HF_TOKEN=hf_test" in script + assert "LIT=val" in script def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): @@ -532,7 +546,7 @@ def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): benchmark = config.driver.benchmarks["gsm8k"] compute = next(iter(config.compute.values())) script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - assert "export WANDB_API_KEY=wb_secret" in script + assert "WANDB_API_KEY=wb_secret" in script # --------------------------------------------------------------------------- From 9aa6bbcbd065f3b35712b3e5f0453743c244b4e5 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Fri, 7 Aug 2026 14:43:18 -0400 Subject: [PATCH 43/66] Add output and mount support Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_1IN.yaml | 7 +- nemo_gym/orchestration/api.py | 7 ++ nemo_gym/orchestration/executors/slurm.py | 2 +- .../orchestration/executors/slurm_script.py | 20 +++- tests/unit_tests/test_slurm_script.py | 113 ++++++++++++++++++ 5 files changed, 142 insertions(+), 7 deletions(-) diff --git a/examples/slurm_vllm_1IN.yaml b/examples/slurm_vllm_1IN.yaml index 278a39e36c2..d89eb1a83a6 100644 --- a/examples/slurm_vllm_1IN.yaml +++ b/examples/slurm_vllm_1IN.yaml @@ -22,15 +22,20 @@ services: port: 8000 health_check: timeout_seconds: 1200 + mounts: + - /lustre/datasets:/data driver: policy_model: vllm_model # synthactic sugar container: python:3.12 # optional, defaults to the service container gym_install: ref: main + # mounts: + # - /lustre/datasets:/data benchmarks: gpqa: - prepare: + output_jsonl_fpath: artifacts/rollouts_gpqa.jsonl # default; relative to the job's working directory + prepare: config_paths: - benchmarks/ifbench/config.yaml run: # gym eval run diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 72502a0f6ed..1764508a414 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -38,6 +38,8 @@ class BaseServiceConfig(_StrictModel): health_check: HealthCheckConfig | None = None # Values starting with "$" are resolved from the host environment at submit time. env: dict[str, str] = {} + # Pyxis-style bind mounts passed as --container-mounts. Each entry is "src:dst" or "src". + mounts: list[str] = [] class BaseModelServiceConfig(BaseServiceConfig): @@ -110,6 +112,9 @@ class BenchmarkRunConfig(_StrictModel): # Hydra overrides forwarded to `gym eval run`. policy_model wiring is injected here at # validation time so all executors see it uniformly via flatten_run_args. run: dict[str, Any] = {} + # Path (relative to the job's working directory) where the rollout JSONL is written. + # The parent directory is pre-created in the staging area before job submission. + output_jsonl_fpath: str = "artifacts/rollouts.jsonl" class GymInstallConfig(_StrictModel): @@ -126,6 +131,8 @@ class DriverConfig(_StrictModel): benchmarks: dict[str, BenchmarkRunConfig] # Values starting with "$" are resolved from the host environment at submit time. env: dict[str, str] = {} + # Pyxis-style bind mounts passed as --container-mounts. Each entry is "src:dst" or "src". + mounts: list[str] = [] class JobConfig(_StrictModel): diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 88ebc00e8b7..bb52b3d6ee2 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -79,7 +79,7 @@ def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_d bench_dir = staging / name bench_dir.mkdir() (bench_dir / "logs").mkdir() - (bench_dir / "artifacts").mkdir() + (bench_dir / Path(benchmark.output_jsonl_fpath).parent).mkdir(parents=True, exist_ok=True) script = build_sbatch_script(config, name, benchmark, compute, remote_run_dir / name) (bench_dir / "job.sh").write_text(script) return staging diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 54e48ae3f70..ef2c1c08990 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -94,15 +94,22 @@ def _resolve_env(env: dict[str, str]) -> str: return f"env {pairs} " -def _render_service_command(name: str, container: str, command: str, env: dict[str, str] | None = None) -> str: +def _render_service_command( + name: str, + container: str, + command: str, + env: dict[str, str] | None = None, + mounts: list[str] | None = None, +) -> str: var = bash_var(name) env_prefix = _resolve_env(env) if env else "" + mounts_flag = f" --container-mounts={','.join(mounts)}" if mounts else "" # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" - f"{env_prefix}srun --overlap --no-container-mount-home --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" + f"{env_prefix}srun --overlap --no-container-mount-home{mounts_flag} --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -134,7 +141,9 @@ def build_sbatch_script( directives = _render_directives(compute, remote_bench_dir, benchmark_name) service_commands = "\n\n".join( - _render_service_command(name, service.container, _BUILDERS[type(service)](service), service.env or None) + _render_service_command( + name, service.container, _BUILDERS[type(service)](service), service.env or None, service.mounts or None + ) for name, service in config.services.items() ) @@ -152,7 +161,7 @@ def build_sbatch_script( if benchmark.prepare: prepare_cmd = "gym eval prepare " + " ".join(flatten_run_args(benchmark.prepare)) - output_path = "+output_jsonl_fpath=artifacts/rollouts.jsonl" + output_path = f"+output_jsonl_fpath={benchmark.output_jsonl_fpath}" extra_flags = ["--model-type openai_model"] if config.driver.policy_model else [] gym_cmd = render_gym_cmd("eval run", "GYM_CMD", [output_path] + extra_flags + flatten_run_args(benchmark.run)) entrypoint = render_driver_entrypoint( @@ -162,9 +171,10 @@ def build_sbatch_script( ) prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" + driver_mounts_flag = f" --container-mounts={','.join(config.driver.mounts)}" if config.driver.mounts else "" driver_command = ( f"{gym_cmd}\n" - f"{driver_env_prefix}srun --overlap --no-container-mount-home --container-image={shlex.quote(config.driver.container)} " + f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index f06ebea0e94..4029fbe73ba 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -549,6 +549,119 @@ def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): assert "WANDB_API_KEY=wb_secret" in script +# --------------------------------------------------------------------------- +# mounts: service and driver +# --------------------------------------------------------------------------- + + +def test_render_service_command_with_mounts(): + out = _render_service_command("svc", "img:latest", "cmd", mounts=["/src:/dst", "/data"]) + assert "--container-mounts=/src:/dst,/data" in out + + +def test_render_service_command_no_mounts_by_default(): + out = _render_service_command("svc", "img:latest", "cmd") + assert "--container-mounts" not in out + + +def test_render_service_command_empty_mounts_omits_flag(): + out = _render_service_command("svc", "img:latest", "cmd", mounts=[]) + assert "--container-mounts" not in out + + +def test_build_sbatch_script_service_mounts(bench_dir): + config = SubmitConfig.model_validate( + { + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "mounts": ["/lustre/datasets:/data", "/tmp/cache"], + } + }, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "--container-mounts=/lustre/datasets:/data,/tmp/cache" in script + + +def test_build_sbatch_script_driver_mounts(bench_dir): + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": { + "container": "python:3.12", + "benchmarks": {"gsm8k": {}}, + "mounts": ["/lustre/checkpoints:/ckpts"], + }, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + # mounts flag appears on the driver srun line + driver_srun_line = next(line for line in script.splitlines() if "python:3.12" in line) + assert "--container-mounts=/lustre/checkpoints:/ckpts" in driver_srun_line + + +def test_build_sbatch_script_no_mounts_by_default(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "--container-mounts" not in script + + +# --------------------------------------------------------------------------- +# output_jsonl_fpath: configurable per benchmark +# --------------------------------------------------------------------------- + + +def test_benchmark_run_config_default_output_jsonl_fpath(): + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + assert config.driver.benchmarks["gsm8k"].output_jsonl_fpath == "artifacts/rollouts.jsonl" + + +def test_build_sbatch_script_output_jsonl_fpath_default(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "+output_jsonl_fpath=artifacts/rollouts.jsonl" in script + + +def test_build_sbatch_script_output_jsonl_fpath_custom(bench_dir): + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, + "driver": { + "container": "python:3.12", + "benchmarks": {"gsm8k": {"output_jsonl_fpath": "out/results/rollouts.jsonl"}}, + }, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "+output_jsonl_fpath=out/results/rollouts.jsonl" in script + assert "artifacts/rollouts.jsonl" not in script + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From 39fb54fa18f78792773fbecbed69c58fc9dd6562 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 10 Aug 2026 13:28:20 -0400 Subject: [PATCH 44/66] Addressing the reviews Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_1IN.yaml | 2 +- nemo_gym/orchestration/api.py | 21 --- tests/unit_tests/test_slurm_script.py | 178 ++++++-------------------- 3 files changed, 37 insertions(+), 164 deletions(-) diff --git a/examples/slurm_vllm_1IN.yaml b/examples/slurm_vllm_1IN.yaml index 278a39e36c2..7cd7477d11b 100644 --- a/examples/slurm_vllm_1IN.yaml +++ b/examples/slurm_vllm_1IN.yaml @@ -17,7 +17,7 @@ services: type: vllm model: Qwen/Qwen2.5-0.5B-Instruct trust_remote_code: true - tensor_parallel_size: $TP + tensor_parallel_size: ${oc.env:TP} # set via: export TP=4 pipeline_parallel_size: 1 port: 8000 health_check: diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 72502a0f6ed..0d21c5ed6ac 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator @@ -36,7 +35,6 @@ class BaseServiceConfig(_StrictModel): # Resolved to the sole compute resource name at validation time when not set. placement: str | None = None health_check: HealthCheckConfig | None = None - # Values starting with "$" are resolved from the host environment at submit time. env: dict[str, str] = {} @@ -124,7 +122,6 @@ class DriverConfig(_StrictModel): # policy_base_url/policy_model_name/policy_api_key into each benchmark's run config. policy_model: str | None = None benchmarks: dict[str, BenchmarkRunConfig] - # Values starting with "$" are resolved from the host environment at submit time. env: dict[str, str] = {} @@ -133,30 +130,12 @@ class JobConfig(_StrictModel): output_path: str -def _resolve_env_refs(data: Any) -> Any: - if isinstance(data, dict): - return {k: _resolve_env_refs(v) for k, v in data.items()} - if isinstance(data, list): - return [_resolve_env_refs(v) for v in data] - if isinstance(data, str) and data.startswith("$"): - host_var = data[1:] - if host_var not in os.environ: - raise ValueError(f"Host environment variable {host_var!r} is not set (referenced as {data!r} in config).") - return os.environ[host_var] - return data - - class SubmitConfig(_StrictModel): services: dict[str, ServiceConfig] compute: dict[str, ComputeConfig] driver: DriverConfig job: JobConfig - @model_validator(mode="before") - @classmethod - def _resolve_host_env_vars(cls, data: Any) -> Any: - return _resolve_env_refs(data) - @model_validator(mode="after") def _resolve_and_validate_placements(self) -> "SubmitConfig": compute_names = set(self.compute) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index f06ebea0e94..0a1731b4dc7 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -15,7 +15,9 @@ from pathlib import Path -from nemo_gym.orchestration.api import SubmitConfig, _resolve_env_refs +import pytest + +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_vllm_command, @@ -295,147 +297,47 @@ def test_resolve_env_valid_keys(): assert "KEY_123=c" in out -# --------------------------------------------------------------------------- -# _resolve_env_refs (pre-validation env var expansion) -# --------------------------------------------------------------------------- - - -def test_resolve_env_refs_literal(): - assert _resolve_env_refs({"key": "value"}) == {"key": "value"} - - -def test_resolve_env_refs_host_var(monkeypatch): - monkeypatch.setenv("MY_TOKEN", "secret123") - assert _resolve_env_refs({"TOKEN": "$MY_TOKEN"}) == {"TOKEN": "secret123"} - - -def test_resolve_env_refs_missing_host_var(monkeypatch): - monkeypatch.delenv("UNSET_VAR", raising=False) - with pytest.raises(ValueError, match="'UNSET_VAR' is not set"): - _resolve_env_refs({"KEY": "$UNSET_VAR"}) - - -def test_resolve_env_refs_nested(monkeypatch): - monkeypatch.setenv("TP", "4") - result = _resolve_env_refs({"service": {"tensor_parallel_size": "$TP"}}) - assert result == {"service": {"tensor_parallel_size": "4"}} - - -def test_resolve_env_refs_list(monkeypatch): - monkeypatch.setenv("VAL", "x") - assert _resolve_env_refs(["$VAL", "literal"]) == ["x", "literal"] - - -def test_resolve_env_refs_non_dollar_string_unchanged(): - assert _resolve_env_refs({"key": "http://localhost:8000/v1"}) == {"key": "http://localhost:8000/v1"} - - -def test_resolve_env_refs_dollar_not_at_start_unchanged(): - # Only leading $ triggers resolution — embedded $ is left as-is. - assert _resolve_env_refs({"key": "prefix_$VAR"}) == {"key": "prefix_$VAR"} - - -def test_resolve_env_refs_error_message_contains_ref(monkeypatch): - monkeypatch.delenv("MISSING", raising=False) - with pytest.raises(ValueError, match=r"\$MISSING"): - _resolve_env_refs({"key": "$MISSING"}) - +def test_resolve_env_invalid_key_with_spaces(): + with pytest.raises(ValueError, match="Invalid environment variable name"): + _resolve_env({"KEY WITH SPACE": "v"}) -def test_resolve_env_refs_non_string_values_unchanged(): - assert _resolve_env_refs({"n": 42, "b": True}) == {"n": 42, "b": True} +def test_resolve_env_invalid_key_with_hyphen(): + with pytest.raises(ValueError, match="Invalid environment variable name"): + _resolve_env({"KEY-NAME": "v"}) -def test_submit_config_resolves_int_field(monkeypatch): - monkeypatch.setenv("TP", "2") - config = SubmitConfig.model_validate( - { - "services": { - "vllm_model": { - "type": "vllm", - "container": "vllm:latest", - "model": "org/model", - "tensor_parallel_size": "$TP", - } - }, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - assert config.services["vllm_model"].tensor_parallel_size == 2 +def test_resolve_env_invalid_key_starts_with_digit(): + with pytest.raises(ValueError, match="Invalid environment variable name"): + _resolve_env({"1KEY": "v"}) -def test_submit_config_resolves_string_fields(monkeypatch): - monkeypatch.setenv("MY_MODEL", "org/llama-3") - monkeypatch.setenv("MY_CONTAINER", "nvcr.io/my-vllm:latest") - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "$MY_CONTAINER", "model": "$MY_MODEL"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - assert config.services["vllm_model"].model == "org/llama-3" - assert config.services["vllm_model"].container == "nvcr.io/my-vllm:latest" +def test_resolve_env_invalid_key_with_equals(): + with pytest.raises(ValueError, match="Invalid environment variable name"): + _resolve_env({"KEY=BAD": "v"}) -def test_submit_config_resolves_compute_account(monkeypatch): - monkeypatch.setenv("SLURM_ACCOUNT", "my_team") - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "$SLURM_ACCOUNT", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - assert config.compute["cluster"].account == "my_team" +def test_resolve_env_invalid_key_with_dollar(): + with pytest.raises(ValueError, match="Invalid environment variable name"): + _resolve_env({"$KEY": "v"}) -def test_submit_config_resolves_output_path(monkeypatch): - monkeypatch.setenv("JOB_DIR", "/lustre/fsw/my-jobs") - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "$JOB_DIR"}, - } - ) - assert config.job.output_path == "/lustre/fsw/my-jobs" +def test_resolve_env_value_with_semicolons_is_quoted(): + out = _resolve_env({"KEY": "val;rm -rf /"}) + assert "KEY='val;rm -rf /'" in out -def test_submit_config_resolves_missing_var_raises(monkeypatch): - from pydantic import ValidationError - monkeypatch.delenv("UNSET_MODEL", raising=False) - with pytest.raises(ValidationError, match="UNSET_MODEL"): - SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "$UNSET_MODEL"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) +def test_resolve_env_value_with_newline_is_quoted(): + out = _resolve_env({"KEY": "line1\nline2"}) + assert "KEY='line1\nline2'" in out -def test_submit_config_resolves_benchmark_run_args(monkeypatch): - monkeypatch.setenv("TEMP", "0.7") - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {"run": {"temperature": "$TEMP"}}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - assert config.driver.benchmarks["gsm8k"].run["temperature"] == "0.7" +# --------------------------------------------------------------------------- +# build_sbatch_script — env injection +# --------------------------------------------------------------------------- -def test_build_sbatch_script_resolved_tp_in_vllm_cmd(bench_dir, monkeypatch): - monkeypatch.setenv("TP", "8") +def test_build_sbatch_script_resolved_tp_in_vllm_cmd(bench_dir): config = SubmitConfig.model_validate( { "services": { @@ -443,7 +345,7 @@ def test_build_sbatch_script_resolved_tp_in_vllm_cmd(bench_dir, monkeypatch): "type": "vllm", "container": "vllm:latest", "model": "org/model", - "tensor_parallel_size": "$TP", + "tensor_parallel_size": 8, } }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, @@ -457,9 +359,7 @@ def test_build_sbatch_script_resolved_tp_in_vllm_cmd(bench_dir, monkeypatch): assert "--tensor-parallel-size 8" in script -def test_build_sbatch_script_service_env_before_driver_env(bench_dir, monkeypatch): - monkeypatch.setenv("SVC_TOKEN", "svc_val") - monkeypatch.setenv("DRV_TOKEN", "drv_val") +def test_build_sbatch_script_service_env_before_driver_env(bench_dir): config = SubmitConfig.model_validate( { "services": { @@ -467,11 +367,11 @@ def test_build_sbatch_script_service_env_before_driver_env(bench_dir, monkeypatc "type": "vllm", "container": "vllm:latest", "model": "org/model", - "env": {"SVC_KEY": "$SVC_TOKEN"}, + "env": {"SVC_KEY": "svc_val"}, } }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}, "env": {"DRV_KEY": "$DRV_TOKEN"}}, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}, "env": {"DRV_KEY": "drv_val"}}, "job": {"output_path": "/remote/jobs"}, } ) @@ -503,8 +403,7 @@ def test_render_service_command_no_env(): assert "export" not in out -def test_build_sbatch_script_service_env(bench_dir, monkeypatch): - monkeypatch.setenv("HF_TOKEN", "hf_test") +def test_build_sbatch_script_service_env(bench_dir): config = SubmitConfig.model_validate( { "services": { @@ -512,7 +411,7 @@ def test_build_sbatch_script_service_env(bench_dir, monkeypatch): "type": "vllm", "container": "vllm:latest", "model": "org/model", - "env": {"HF_TOKEN": "$HF_TOKEN", "LIT": "val"}, + "env": {"HF_TOKEN": "hf_test", "LIT": "val"}, } }, "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, @@ -520,8 +419,6 @@ def test_build_sbatch_script_service_env(bench_dir, monkeypatch): "job": {"output_path": "/remote/jobs"}, } ) - # $HF_TOKEN resolved at parse time — env dict now has literal values - assert config.services["vllm_model"].env == {"HF_TOKEN": "hf_test", "LIT": "val"} benchmark = config.driver.benchmarks["gsm8k"] compute = next(iter(config.compute.values())) script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) @@ -529,8 +426,7 @@ def test_build_sbatch_script_service_env(bench_dir, monkeypatch): assert "LIT=val" in script -def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): - monkeypatch.setenv("WANDB_KEY", "wb_secret") +def test_build_sbatch_script_driver_env(bench_dir): config = SubmitConfig.model_validate( { "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, @@ -538,7 +434,7 @@ def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): "driver": { "container": "python:3.12", "benchmarks": {"gsm8k": {}}, - "env": {"WANDB_API_KEY": "$WANDB_KEY"}, + "env": {"WANDB_API_KEY": "wb_secret"}, }, "job": {"output_path": "/remote/jobs"}, } @@ -553,8 +449,6 @@ def test_build_sbatch_script_driver_env(bench_dir, monkeypatch): # Fixtures # --------------------------------------------------------------------------- -import pytest - from nemo_gym.orchestration.api import NodePool, SlurmComputeConfig, VllmServiceConfig From c688eb386eb3f49c72fe03a54ce313e9a86f7d48 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 10 Aug 2026 13:32:10 -0400 Subject: [PATCH 45/66] Fix CI secret issue Signed-off-by: Onur Yilmaz --- tests/unit_tests/test_slurm_script.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 0a1731b4dc7..33a056add26 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -434,7 +434,7 @@ def test_build_sbatch_script_driver_env(bench_dir): "driver": { "container": "python:3.12", "benchmarks": {"gsm8k": {}}, - "env": {"WANDB_API_KEY": "wb_secret"}, + "env": {"WANDB_API_KEY": "wb_secret"}, # pragma: allowlist secret }, "job": {"output_path": "/remote/jobs"}, } From e1ced0dd4d18244119e5a904721f4f3b60c9c83e Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 10 Aug 2026 15:32:15 -0400 Subject: [PATCH 46/66] Fix test issue Signed-off-by: Onur Yilmaz --- responses_api_agents/aviary_agent/tests/test_app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/responses_api_agents/aviary_agent/tests/test_app.py b/responses_api_agents/aviary_agent/tests/test_app.py index 28f35dbd57f..21fb2eced11 100644 --- a/responses_api_agents/aviary_agent/tests/test_app.py +++ b/responses_api_agents/aviary_agent/tests/test_app.py @@ -551,7 +551,7 @@ async def test_run_workflow(self) -> None: "env_id": env_id, "group_id": "0", "contains_transitions": True, - "output": [[{"role": "user", "content": "obs"}]], + "output": [[{"type": "message", "role": "user", "content": "obs"}]], "parallel_tool_calls": True, "tool_choice": "auto", "tools": [], From 6f7587481d3ec0cd1133873c7f3998f8f29cfe7e Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 10 Aug 2026 15:50:57 -0400 Subject: [PATCH 47/66] Revert output file changes Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_1IN.yaml | 1 - nemo_gym/orchestration/api.py | 3 -- nemo_gym/orchestration/executors/slurm.py | 2 +- .../orchestration/executors/slurm_script.py | 2 +- tests/unit_tests/test_slurm_script.py | 43 ------------------- 5 files changed, 2 insertions(+), 49 deletions(-) diff --git a/examples/slurm_vllm_1IN.yaml b/examples/slurm_vllm_1IN.yaml index dc5f9594510..0ea44df923f 100644 --- a/examples/slurm_vllm_1IN.yaml +++ b/examples/slurm_vllm_1IN.yaml @@ -34,7 +34,6 @@ driver: # - /lustre/datasets:/data benchmarks: gpqa: - output_jsonl_fpath: artifacts/rollouts_gpqa.jsonl # default; relative to the job's working directory prepare: config_paths: - benchmarks/ifbench/config.yaml diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index b0d87e7b1b4..a7cef27f7b5 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -110,9 +110,6 @@ class BenchmarkRunConfig(_StrictModel): # Hydra overrides forwarded to `gym eval run`. policy_model wiring is injected here at # validation time so all executors see it uniformly via flatten_run_args. run: dict[str, Any] = {} - # Path (relative to the job's working directory) where the rollout JSONL is written. - # The parent directory is pre-created in the staging area before job submission. - output_jsonl_fpath: str = "artifacts/rollouts.jsonl" class GymInstallConfig(_StrictModel): diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index bb52b3d6ee2..88ebc00e8b7 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -79,7 +79,7 @@ def _stage(self, config: SubmitConfig, compute: SlurmComputeConfig, remote_run_d bench_dir = staging / name bench_dir.mkdir() (bench_dir / "logs").mkdir() - (bench_dir / Path(benchmark.output_jsonl_fpath).parent).mkdir(parents=True, exist_ok=True) + (bench_dir / "artifacts").mkdir() script = build_sbatch_script(config, name, benchmark, compute, remote_run_dir / name) (bench_dir / "job.sh").write_text(script) return staging diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index ef2c1c08990..1508d135f0a 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -161,7 +161,7 @@ def build_sbatch_script( if benchmark.prepare: prepare_cmd = "gym eval prepare " + " ".join(flatten_run_args(benchmark.prepare)) - output_path = f"+output_jsonl_fpath={benchmark.output_jsonl_fpath}" + output_path = "+output_jsonl_fpath=artifacts/rollouts.jsonl" extra_flags = ["--model-type openai_model"] if config.driver.policy_model else [] gym_cmd = render_gym_cmd("eval run", "GYM_CMD", [output_path] + extra_flags + flatten_run_args(benchmark.run)) entrypoint = render_driver_entrypoint( diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 00fd120b68f..376c42fd7c8 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -515,49 +515,6 @@ def test_build_sbatch_script_no_mounts_by_default(submit_config, bench_dir): assert "--container-mounts" not in script -# --------------------------------------------------------------------------- -# output_jsonl_fpath: configurable per benchmark -# --------------------------------------------------------------------------- - - -def test_benchmark_run_config_default_output_jsonl_fpath(): - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - assert config.driver.benchmarks["gsm8k"].output_jsonl_fpath == "artifacts/rollouts.jsonl" - - -def test_build_sbatch_script_output_jsonl_fpath_default(submit_config, bench_dir): - benchmark = submit_config.driver.benchmarks["gsm8k"] - compute = next(iter(submit_config.compute.values())) - script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) - assert "+output_jsonl_fpath=artifacts/rollouts.jsonl" in script - - -def test_build_sbatch_script_output_jsonl_fpath_custom(bench_dir): - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": {"cluster": {"type": "slurm", "account": "my-account", "hostname": "foo"}}, - "driver": { - "container": "python:3.12", - "benchmarks": {"gsm8k": {"output_jsonl_fpath": "out/results/rollouts.jsonl"}}, - }, - "job": {"output_path": "/remote/jobs"}, - } - ) - benchmark = config.driver.benchmarks["gsm8k"] - compute = next(iter(config.compute.values())) - script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - assert "+output_jsonl_fpath=out/results/rollouts.jsonl" in script - assert "artifacts/rollouts.jsonl" not in script - - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From c2f6dc9d9ee303d05e3f7a1955c3b577d9eada73 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 10 Aug 2026 18:01:45 -0400 Subject: [PATCH 48/66] Check the mounts exist or not Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 6 +- nemo_gym/orchestration/executors/slurm.py | 25 ++++++- .../orchestration/executors/slurm_script.py | 8 +- tests/unit_tests/test_slurm_script.py | 73 +++++++++++++++++++ 4 files changed, 107 insertions(+), 5 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index a7cef27f7b5..c55c571a540 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -36,7 +36,8 @@ class BaseServiceConfig(_StrictModel): placement: str | None = None health_check: HealthCheckConfig | None = None env: dict[str, str] = {} - # Pyxis-style bind mounts passed as --container-mounts. Each entry is "src:dst" or "src". + # Pyxis-style bind mounts passed as --container-mounts. + # Each entry is "src", "src:dst", or "src:dst:flags" (e.g. "/data:/data:ro"). mounts: list[str] = [] @@ -125,7 +126,8 @@ class DriverConfig(_StrictModel): policy_model: str | None = None benchmarks: dict[str, BenchmarkRunConfig] env: dict[str, str] = {} - # Pyxis-style bind mounts passed as --container-mounts. Each entry is "src:dst" or "src". + # Pyxis-style bind mounts passed as --container-mounts. + # Each entry is "src", "src:dst", or "src:dst:flags" (e.g. "/data:/data:ro"). mounts: list[str] = [] diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index 88ebc00e8b7..ce7d15d7ad1 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -23,13 +23,35 @@ from nemo_gym.orchestration.api import SlurmComputeConfig, SubmitConfig from nemo_gym.orchestration.executors.base import BaseExecutor -from nemo_gym.orchestration.executors.connection import get_connection +from nemo_gym.orchestration.executors.connection import Connection, LocalConnection, get_connection from nemo_gym.orchestration.executors.slurm_script import build_sbatch_script _SBATCH_JOB_ID_RE = re.compile(r"Submitted batch job (\d+)") +def _validate_mounts(config: SubmitConfig, conn: Connection) -> None: + entries = [("driver", m) for m in config.driver.mounts] + for name, service in config.services.items(): + entries += [(f"services.{name}", m) for m in service.mounts] + srcs_by_label = [(label, mount.split(":")[0]) for label, mount in entries] + if not srcs_by_label: + return + + if isinstance(conn, LocalConnection): + missing = {src for _, src in srcs_by_label if not Path(src).exists()} + else: + # SSHConnection.run() pipes commands as a bash script, so || works fine. + checks = [f'test -e {shlex.quote(src)} || echo "__GYM_MISSING:{src}"' for _, src in srcs_by_label] + output = conn.run(checks) + missing = {line[len("__GYM_MISSING:"):] for line in output.splitlines() if line.startswith("__GYM_MISSING:")} + + if missing: + bad = [(label, src) for label, src in srcs_by_label if src in missing] + details = "\n".join(f" {label}: {src!r}" for label, src in bad) + raise ValueError(f"Mount src paths do not exist:\n{details}") + + class SlurmExecutor(BaseExecutor): """Slurm executor for Pyxis-enabled clusters (https://github.com/NVIDIA/pyxis). @@ -50,6 +72,7 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: with tempfile.TemporaryDirectory(prefix="gym-submit-") as staging_str: staging = self._stage(config, compute, remote_run_dir, Path(staging_str)) with get_connection(compute.hostname) as conn: + _validate_mounts(config, conn) conn.copy(staging, remote_run_dir) output = conn.run( [ diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 1508d135f0a..e9400a95a69 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -103,7 +103,7 @@ def _render_service_command( ) -> str: var = bash_var(name) env_prefix = _resolve_env(env) if env else "" - mounts_flag = f" --container-mounts={','.join(mounts)}" if mounts else "" + mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else "" # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. @@ -171,7 +171,11 @@ def build_sbatch_script( ) prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" - driver_mounts_flag = f" --container-mounts={','.join(config.driver.mounts)}" if config.driver.mounts else "" + driver_mounts_flag = ( + f" --container-mounts={','.join(shlex.quote(m) for m in config.driver.mounts)}" + if config.driver.mounts + else "" + ) driver_command = ( f"{gym_cmd}\n" f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 376c42fd7c8..7af6582392f 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -515,6 +515,79 @@ def test_build_sbatch_script_no_mounts_by_default(submit_config, bench_dir): assert "--container-mounts" not in script +# --------------------------------------------------------------------------- +# _validate_mounts +# --------------------------------------------------------------------------- + +from unittest.mock import MagicMock + +from nemo_gym.orchestration.executors.connection import LocalConnection +from nemo_gym.orchestration.executors.slurm import _validate_mounts + + +def _make_submit_config_with_mounts(driver_mounts=None, service_mounts=None): + return SubmitConfig.model_validate( + { + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + **({"mounts": service_mounts} if service_mounts is not None else {}), + } + }, + "compute": {"cluster": {"type": "slurm", "account": "acct", "hostname": "foo"}}, + "driver": { + "container": "python:3.12", + "benchmarks": {"gsm8k": {}}, + **({"mounts": driver_mounts} if driver_mounts is not None else {}), + }, + "job": {"output_path": "/remote/jobs"}, + } + ) + + +def test_validate_mounts_local_passes_when_src_exists(tmp_path): + src = str(tmp_path) + config = _make_submit_config_with_mounts(driver_mounts=[f"{src}:/data"]) + _validate_mounts(config, LocalConnection()) + + +def test_validate_mounts_local_raises_for_missing_src(tmp_path): + src = str(tmp_path / "nonexistent") + config = _make_submit_config_with_mounts(driver_mounts=[f"{src}:/data"]) + with pytest.raises(ValueError, match="driver"): + _validate_mounts(config, LocalConnection()) + + +def test_validate_mounts_local_parses_flags_format(tmp_path): + src = str(tmp_path) + config = _make_submit_config_with_mounts(driver_mounts=[f"{src}:/data:ro"]) + _validate_mounts(config, LocalConnection()) # should check src only, not src:ro + + +def test_validate_mounts_remote_passes_when_all_exist(): + config = _make_submit_config_with_mounts(driver_mounts=["/lustre/data:/data"]) + conn = MagicMock() + conn.run.return_value = "" # no __GYM_MISSING lines + _validate_mounts(config, conn) + + +def test_validate_mounts_remote_raises_for_missing_src(): + config = _make_submit_config_with_mounts(service_mounts=["/lustre/missing:/data"]) + conn = MagicMock() + conn.run.return_value = "__GYM_MISSING:/lustre/missing" + with pytest.raises(ValueError, match="services\\.vllm_model"): + _validate_mounts(config, conn) + + +def test_validate_mounts_no_mounts_passes(): + config = _make_submit_config_with_mounts() + conn = MagicMock() + _validate_mounts(config, conn) + conn.run.assert_not_called() + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From a0a5525e5db3168481555b804c8cd1b12b909965 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 11 Aug 2026 12:09:16 -0400 Subject: [PATCH 49/66] Add assertions Signed-off-by: Onur Yilmaz --- tests/unit_tests/test_slurm_script.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 7af6582392f..a623005aa7d 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -550,35 +550,42 @@ def _make_submit_config_with_mounts(driver_mounts=None, service_mounts=None): def test_validate_mounts_local_passes_when_src_exists(tmp_path): src = str(tmp_path) config = _make_submit_config_with_mounts(driver_mounts=[f"{src}:/data"]) + # No exception means all srcs were found — implicitly asserted by the call completing. _validate_mounts(config, LocalConnection()) def test_validate_mounts_local_raises_for_missing_src(tmp_path): src = str(tmp_path / "nonexistent") config = _make_submit_config_with_mounts(driver_mounts=[f"{src}:/data"]) - with pytest.raises(ValueError, match="driver"): + with pytest.raises(ValueError, match="driver") as exc_info: _validate_mounts(config, LocalConnection()) + assert src in str(exc_info.value) def test_validate_mounts_local_parses_flags_format(tmp_path): src = str(tmp_path) + # Passes because src exists; would raise if the code checked "src:ro" as a path instead of "src". config = _make_submit_config_with_mounts(driver_mounts=[f"{src}:/data:ro"]) - _validate_mounts(config, LocalConnection()) # should check src only, not src:ro + _validate_mounts(config, LocalConnection()) def test_validate_mounts_remote_passes_when_all_exist(): config = _make_submit_config_with_mounts(driver_mounts=["/lustre/data:/data"]) conn = MagicMock() - conn.run.return_value = "" # no __GYM_MISSING lines + conn.run.return_value = "" # no __GYM_MISSING lines → all exist _validate_mounts(config, conn) + # Confirms conn.run was called with a check for the correct src. + (commands,), _ = conn.run.call_args + assert any("/lustre/data" in cmd for cmd in commands) def test_validate_mounts_remote_raises_for_missing_src(): config = _make_submit_config_with_mounts(service_mounts=["/lustre/missing:/data"]) conn = MagicMock() conn.run.return_value = "__GYM_MISSING:/lustre/missing" - with pytest.raises(ValueError, match="services\\.vllm_model"): + with pytest.raises(ValueError, match="services\\.vllm_model") as exc_info: _validate_mounts(config, conn) + assert "/lustre/missing" in str(exc_info.value) def test_validate_mounts_no_mounts_passes(): From 1bdf41f09a1337e725284021d420b4b3a09c22ae Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 11 Aug 2026 12:11:17 -0400 Subject: [PATCH 50/66] Fix linting issues Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/executors/slurm.py | 2 +- nemo_gym/orchestration/executors/slurm_script.py | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm.py b/nemo_gym/orchestration/executors/slurm.py index ce7d15d7ad1..ae96686cf3d 100644 --- a/nemo_gym/orchestration/executors/slurm.py +++ b/nemo_gym/orchestration/executors/slurm.py @@ -44,7 +44,7 @@ def _validate_mounts(config: SubmitConfig, conn: Connection) -> None: # SSHConnection.run() pipes commands as a bash script, so || works fine. checks = [f'test -e {shlex.quote(src)} || echo "__GYM_MISSING:{src}"' for _, src in srcs_by_label] output = conn.run(checks) - missing = {line[len("__GYM_MISSING:"):] for line in output.splitlines() if line.startswith("__GYM_MISSING:")} + missing = {line[len("__GYM_MISSING:") :] for line in output.splitlines() if line.startswith("__GYM_MISSING:")} if missing: bad = [(label, src) for label, src in srcs_by_label if src in missing] diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index e9400a95a69..b056991803c 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -172,9 +172,7 @@ def build_sbatch_script( prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" driver_mounts_flag = ( - f" --container-mounts={','.join(shlex.quote(m) for m in config.driver.mounts)}" - if config.driver.mounts - else "" + f" --container-mounts={','.join(shlex.quote(m) for m in config.driver.mounts)}" if config.driver.mounts else "" ) driver_command = ( f"{gym_cmd}\n" From 62a6d41575989591e81ed2a7f3fe13cbff66b4ff Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 11 Aug 2026 14:51:49 -0400 Subject: [PATCH 51/66] Adding multi instance support using vllm serivce Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance.yaml | 43 +++++++++++++++++++ nemo_gym/orchestration/api.py | 35 ++++++++++++++- .../orchestration/executors/slurm_script.py | 10 ++++- tests/unit_tests/test_orchestration_api.py | 38 ++++++++++++++++ tests/unit_tests/test_slurm_script.py | 28 ++++++++++++ 5 files changed, 152 insertions(+), 2 deletions(-) create mode 100644 examples/slurm_vllm_multi_instance.yaml diff --git a/examples/slurm_vllm_multi_instance.yaml b/examples/slurm_vllm_multi_instance.yaml new file mode 100644 index 00000000000..22b1091c8cf --- /dev/null +++ b/examples/slurm_vllm_multi_instance.yaml @@ -0,0 +1,43 @@ +compute: + cluster: + type: slurm + hostname: myslurmcluster + walltime: "01:00:00" + account: myaccount + node_pools: + compute: + partition: batch + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:v0.9.0 + type: vllm + model: Qwen/Qwen2.5-7B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 + number_of_instances: 4 # 4 engine replicas, total 8 GPUs + distributed_backend: + type: vllm_service # vLLM native --data-parallel-size; ray_serve/dynamo coming later + port: 8000 + health_check: + timeout_seconds: 1200 + mounts: + - /lustre/datasets:/data + +driver: + policy_model: vllm_model + container: python:3.12 + 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 c55c571a540..f8e2dd5ab35 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -15,7 +15,7 @@ from typing import Annotated, Any, Literal -from pydantic import BaseModel, ConfigDict, Discriminator, Tag, model_validator +from pydantic import BaseModel, ConfigDict, Discriminator, Tag, field_validator, model_validator # Reject unknown fields on all config models so typos in YAML surface immediately. @@ -48,11 +48,44 @@ class BaseModelServiceConfig(BaseServiceConfig): port: int = 8000 +class VllmServiceDistributedBackend(_StrictModel): + """Use vLLM's native data-parallel multi-instance (--data-parallel-size N).""" + + type: Literal["vllm_service"] = "vllm_service" + + +# Future backends: add Annotated[RayServeDistributedBackend, Tag("ray_serve")], etc. +DistributedBackendConfig = Annotated[ + Annotated[VllmServiceDistributedBackend, Tag("vllm_service")], + Discriminator("type"), +] + + class VllmServiceConfig(BaseModelServiceConfig): type: Literal["vllm"] tensor_parallel_size: int = 1 pipeline_parallel_size: int = 1 trust_remote_code: bool = False + number_of_instances: int = 1 + distributed_backend: DistributedBackendConfig | None = None + + @field_validator("number_of_instances") + @classmethod + def _validate_number_of_instances(cls, v: int) -> int: + if v < 1: + raise ValueError(f"number_of_instances must be >= 1, got {v}") + return v + + @model_validator(mode="after") + def _validate_distributed_backend(self) -> "VllmServiceConfig": + if self.number_of_instances > 1 and self.distributed_backend is None: + raise ValueError( + f"distributed_backend must be set when number_of_instances > 1 (got {self.number_of_instances}). " + "Use: distributed_backend: {type: vllm_service}" + ) + if self.number_of_instances == 1 and self.distributed_backend is not None: + raise ValueError("distributed_backend should not be set when number_of_instances == 1") + return self @model_validator(mode="after") def _default_health_check(self) -> "VllmServiceConfig": diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index b056991803c..982f7d8924e 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -115,7 +115,15 @@ def _render_service_command( def _build_vllm_command(service: VllmServiceConfig) -> str: - cmd = f"vllm serve {shlex.quote(service.model)} --port {service.port} --tensor-parallel-size {service.tensor_parallel_size}" + cmd = ( + f"vllm serve {shlex.quote(service.model)}" + f" --port {service.port}" + f" --tensor-parallel-size {service.tensor_parallel_size}" + ) + if service.pipeline_parallel_size > 1: + cmd += f" --pipeline-parallel-size {service.pipeline_parallel_size}" + if service.number_of_instances > 1: + cmd += f" --data-parallel-size {service.number_of_instances}" if service.trust_remote_code: cmd += " --trust-remote-code" return cmd diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 2b284c71976..5343a21f3a6 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -96,3 +96,41 @@ def test_driver_env_accepted(): def test_service_unknown_field_raises(): with pytest.raises(ValidationError): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "unknown_field": "x"}})) + + +# --------------------------------------------------------------------------- +# number_of_instances / distributed_backend +# --------------------------------------------------------------------------- + +_MULTI_SERVICE = {**SERVICE, "number_of_instances": 4, "distributed_backend": {"type": "vllm_service"}} + + +def test_number_of_instances_with_backend_accepted(): + config = SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE})) + svc = config.services["svc"] + assert svc.number_of_instances == 4 + assert svc.distributed_backend is not None + assert svc.distributed_backend.type == "vllm_service" + + +def test_number_of_instances_defaults_to_1(): + config = SubmitConfig.model_validate(_config()) + assert config.services["svc"].number_of_instances == 1 + assert config.services["svc"].distributed_backend is None + + +def test_multi_instance_without_backend_raises(): + with pytest.raises(ValidationError, match="distributed_backend must be set"): + SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 4}})) + + +def test_single_instance_with_backend_raises(): + with pytest.raises(ValidationError, match="should not be set"): + SubmitConfig.model_validate( + _config(services={"svc": {**SERVICE, "distributed_backend": {"type": "vllm_service"}}}) + ) + + +def test_number_of_instances_zero_raises(): + with pytest.raises(ValidationError): + SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 0}})) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index a623005aa7d..30396e81491 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -169,6 +169,34 @@ def test_build_vllm_command_no_trust_remote_code_by_default(vllm_service): assert "--trust-remote-code" not in cmd +def test_build_vllm_command_multi_instance(): + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + number_of_instances=4, + distributed_backend={"type": "vllm_service"}, + ) + cmd = _build_vllm_command(service) + assert "--data-parallel-size 4" in cmd + + +def test_build_vllm_command_single_instance_omits_dp_flag(vllm_service): + cmd = _build_vllm_command(vllm_service) + assert "--data-parallel-size" not in cmd + + +def test_build_vllm_command_pipeline_parallel(): + service = VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model", pipeline_parallel_size=2) + cmd = _build_vllm_command(service) + assert "--pipeline-parallel-size 2" in cmd + + +def test_build_vllm_command_pipeline_parallel_1_omits_flag(vllm_service): + cmd = _build_vllm_command(vllm_service) + assert "--pipeline-parallel-size" not in cmd + + # --------------------------------------------------------------------------- # render_gym_cmd # --------------------------------------------------------------------------- From 641efc2bfb26ad66d3b29ab0a7111e7fdaea3aa2 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 12 Aug 2026 16:36:19 -0400 Subject: [PATCH 52/66] Adding multi node srun support Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance.yaml | 2 +- .../orchestration/executors/slurm_script.py | 25 ++- tests/unit_tests/test_slurm_script.py | 144 ++++++++++++++++++ 3 files changed, 167 insertions(+), 4 deletions(-) diff --git a/examples/slurm_vllm_multi_instance.yaml b/examples/slurm_vllm_multi_instance.yaml index 22b1091c8cf..4a19a1b22da 100644 --- a/examples/slurm_vllm_multi_instance.yaml +++ b/examples/slurm_vllm_multi_instance.yaml @@ -29,7 +29,7 @@ services: driver: policy_model: vllm_model - container: python:3.12 + container: python:3.13 gym_install: ref: main benchmarks: diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 982f7d8924e..5396893b2c2 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -100,16 +100,19 @@ def _render_service_command( command: str, env: dict[str, str] | None = None, mounts: list[str] | None = None, + nodes: int | None = None, + ntasks: int | None = None, ) -> str: var = bash_var(name) env_prefix = _resolve_env(env) if env else "" + node_flags = f" --nodes={nodes} --ntasks={ntasks}" if (nodes is not None and nodes > 1) else "" mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else "" # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" - f"{env_prefix}srun --overlap --no-container-mount-home{mounts_flag} --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" + f"{env_prefix}srun --overlap --no-container-mount-home{node_flags}{mounts_flag} --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -139,6 +142,12 @@ def _build_ray_command(_service: RayServiceConfig) -> str: } +def _node_totals(compute: SlurmComputeConfig) -> tuple[int, int]: + total_nodes = sum(pool.nodes for pool in compute.node_pools.values()) + total_ntasks = sum(pool.nodes * pool.ntasks_per_node for pool in compute.node_pools.values()) + return total_nodes, total_ntasks + + def build_sbatch_script( config: SubmitConfig, benchmark_name: str, @@ -148,9 +157,18 @@ def build_sbatch_script( ) -> str: directives = _render_directives(compute, remote_bench_dir, benchmark_name) + total_nodes, total_ntasks = _node_totals(compute) + is_multi_node = total_nodes > 1 + service_commands = "\n\n".join( _render_service_command( - name, service.container, _BUILDERS[type(service)](service), service.env or None, service.mounts or None + name, + service.container, + _BUILDERS[type(service)](service), + service.env or None, + service.mounts or None, + nodes=total_nodes if is_multi_node else None, + ntasks=total_ntasks if is_multi_node else None, ) for name, service in config.services.items() ) @@ -179,12 +197,13 @@ def build_sbatch_script( ) prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" + driver_node_flags = " --nodes=1 --ntasks=1" if is_multi_node else "" driver_mounts_flag = ( f" --container-mounts={','.join(shlex.quote(m) for m in config.driver.mounts)}" if config.driver.mounts else "" ) driver_command = ( f"{gym_cmd}\n" - f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " + f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_node_flags}{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 30396e81491..11b8961cebf 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -21,6 +21,7 @@ from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd from nemo_gym.orchestration.executors.slurm_script import ( _build_vllm_command, + _node_totals, _render_directives, _render_pool_directives, _render_service_command, @@ -623,6 +624,149 @@ def test_validate_mounts_no_mounts_passes(): conn.run.assert_not_called() +# --------------------------------------------------------------------------- +# _render_service_command — multi-node flags +# --------------------------------------------------------------------------- + + +def test_render_service_command_multi_node_adds_nodes_and_ntasks(): + out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model", nodes=4, ntasks=4) + assert "--nodes=4" in out + assert "--ntasks=4" in out + + +def test_render_service_command_multi_node_flags_before_container_image(): + out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model", nodes=4, ntasks=4) + assert out.index("--nodes=4") < out.index("--container-image=") + + +def test_render_service_command_single_node_omits_node_flags(): + out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model", nodes=1, ntasks=1) + assert "--nodes=" not in out + assert "--ntasks=" not in out + + +def test_render_service_command_no_nodes_kwarg_omits_node_flags(): + out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model") + assert "--nodes=" not in out + assert "--ntasks=" not in out + + +# --------------------------------------------------------------------------- +# _node_totals +# --------------------------------------------------------------------------- + + +def test_node_totals_empty_pools(): + compute = SlurmComputeConfig(type="slurm", account="acct") + assert _node_totals(compute) == (0, 0) + + +def test_node_totals_single_pool(): + compute = SlurmComputeConfig( + type="slurm", + account="acct", + node_pools={"main": NodePool(partition="gpu", nodes=4, ntasks_per_node=2)}, + ) + assert _node_totals(compute) == (4, 8) + + +def test_node_totals_multiple_pools(): + compute = SlurmComputeConfig( + type="slurm", + account="acct", + node_pools={ + "gpu": NodePool(partition="gpu", nodes=4, ntasks_per_node=1), + "cpu": NodePool(partition="cpu", nodes=2, ntasks_per_node=2), + }, + ) + assert _node_totals(compute) == (6, 8) + + +# --------------------------------------------------------------------------- +# build_sbatch_script — multi-node srun flags +# --------------------------------------------------------------------------- + + +def _multi_node_config(): + return SubmitConfig.model_validate( + { + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "tensor_parallel_size": 8, + } + }, + "compute": { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"main": {"partition": "gpu", "nodes": 4, "ntasks_per_node": 1}}, + } + }, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + + +def test_build_sbatch_script_multi_node_vllm_srun_gets_node_flags(bench_dir): + config = _multi_node_config() + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) + assert "--nodes=4" in vllm_line + assert "--ntasks=4" in vllm_line + + +def test_build_sbatch_script_multi_node_driver_srun_gets_nodes_1(bench_dir): + config = _multi_node_config() + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + driver_line = next(line for line in script.splitlines() if "python:3.12" in line) + assert "--nodes=1" in driver_line + assert "--ntasks=1" in driver_line + + +def test_build_sbatch_script_multi_node_node_flags_before_container_image(bench_dir): + config = _multi_node_config() + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) + assert vllm_line.index("--nodes=4") < vllm_line.index("--container-image=") + + +def test_build_sbatch_script_single_node_pool_omits_node_flags_from_srun(bench_dir): + config = SubmitConfig.model_validate( + { + "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, + "compute": { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"main": {"partition": "gpu", "nodes": 1, "ntasks_per_node": 4}}, + } + }, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) + driver_line = next(line for line in script.splitlines() if "python:3.12" in line) + assert "--nodes=" not in vllm_line + assert "--nodes=" not in driver_line + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From 83577b50b037302d59fdc9fcae85fdd084645fab Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 17 Aug 2026 16:55:59 -0400 Subject: [PATCH 53/66] Add ray based multi node Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance.yaml | 3 +- examples/slurm_vllm_ray_multi_node.yaml | 45 ++++++++++++ nemo_gym/orchestration/api.py | 33 +++++++-- .../orchestration/executors/slurm_script.py | 51 +++++++++++++- tests/unit_tests/test_orchestration_api.py | 50 ++++++++++++-- tests/unit_tests/test_slurm_script.py | 69 +++++++++++++++++++ 6 files changed, 237 insertions(+), 14 deletions(-) create mode 100644 examples/slurm_vllm_ray_multi_node.yaml diff --git a/examples/slurm_vllm_multi_instance.yaml b/examples/slurm_vllm_multi_instance.yaml index 4a19a1b22da..0e074ee45ec 100644 --- a/examples/slurm_vllm_multi_instance.yaml +++ b/examples/slurm_vllm_multi_instance.yaml @@ -20,7 +20,8 @@ services: tensor_parallel_size: 2 number_of_instances: 4 # 4 engine replicas, total 8 GPUs distributed_backend: - type: vllm_service # vLLM native --data-parallel-size; ray_serve/dynamo coming later + type: vllm_service # vLLM native --data-parallel-size, single node only. For multi-node, + # use distributed_backend: {type: ray} - see slurm_vllm_ray_multi_node.yaml port: 8000 health_check: timeout_seconds: 1200 diff --git a/examples/slurm_vllm_ray_multi_node.yaml b/examples/slurm_vllm_ray_multi_node.yaml new file mode 100644 index 00000000000..d0457558075 --- /dev/null +++ b/examples/slurm_vllm_ray_multi_node.yaml @@ -0,0 +1,45 @@ +compute: + cluster: + type: slurm + hostname: myslurmcluster + walltime: "01:00:00" + account: myaccount + node_pools: + compute: + partition: batch + nodes: 2 + 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: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 8 # needs the ray backend since the footprint exceeds one node + distributed_backend: + type: ray # vLLM's Ray core executor (--distributed-executor-backend ray). + # Not the ray.serve library - just lets vLLM's instances span nodes. + # A separate ray_serve backend may be added later. + 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 f8e2dd5ab35..f36380054d1 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -49,14 +49,24 @@ class BaseModelServiceConfig(BaseServiceConfig): class VllmServiceDistributedBackend(_StrictModel): - """Use vLLM's native data-parallel multi-instance (--data-parallel-size N).""" + """Use vLLM's native data-parallel multi-instance (--data-parallel-size N). Single node only.""" type: Literal["vllm_service"] = "vllm_service" -# Future backends: add Annotated[RayServeDistributedBackend, Tag("ray_serve")], etc. +class RayDistributedBackend(_StrictModel): + """Use vLLM's Ray executor (--distributed-executor-backend ray) so a single vLLM service can + span multiple physical Slurm nodes. This does not use the ray.serve library. + + Named "ray" for now; a separate "ray_serve" backend (using the actual ray.serve library) may + be added later. + """ + + type: Literal["ray"] = "ray" + + DistributedBackendConfig = Annotated[ - Annotated[VllmServiceDistributedBackend, Tag("vllm_service")], + Annotated[VllmServiceDistributedBackend, Tag("vllm_service")] | Annotated[RayDistributedBackend, Tag("ray")], Discriminator("type"), ] @@ -83,8 +93,6 @@ def _validate_distributed_backend(self) -> "VllmServiceConfig": f"distributed_backend must be set when number_of_instances > 1 (got {self.number_of_instances}). " "Use: distributed_backend: {type: vllm_service}" ) - if self.number_of_instances == 1 and self.distributed_backend is not None: - raise ValueError("distributed_backend should not be set when number_of_instances == 1") return self @model_validator(mode="after") @@ -183,6 +191,10 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": raise ValueError(f"Multiple compute resources are not supported yet ({', '.join(sorted(compute_names))}).") sole_compute = next(iter(compute_names)) + compute = self.compute[sole_compute] + is_multi_node = ( + isinstance(compute, SlurmComputeConfig) and sum(p.nodes for p in compute.node_pools.values()) > 1 + ) for service_name, service in self.services.items(): if service.placement is None: @@ -193,6 +205,17 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": f"({', '.join(sorted(compute_names))})." ) + if ( + is_multi_node + and isinstance(service, VllmServiceConfig) + and (service.distributed_backend is None or service.distributed_backend.type != "ray") + ): + raise ValueError( + f"Service '{service_name}' is placed on compute '{sole_compute}', which spans multiple nodes " + f"({sum(p.nodes for p in compute.node_pools.values())} total). vLLM services on multi-node " + "compute must use distributed_backend: {type: ray}." + ) + if self.driver.policy_model is not None: if self.driver.policy_model not in self.services: raise ValueError( diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 5396893b2c2..38475823147 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -38,6 +38,8 @@ #!/bin/bash {directives} +{ray_prelude} + {service_commands} {health_checks} @@ -48,6 +50,18 @@ """ +# Resolves the Ray head node from Slurm's node list, mirroring scripts/sbatch_base.sh. +# Runs once in the main sbatch shell (before any srun step) so $SLURM_JOB_NODELIST is available. +_RAY_PRELUDE = """\ +# Resolve Ray head node IP for multi-node vLLM services (distributed_backend: ray). +nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +nodes_array=($nodes) +head_node_hostname=${nodes_array[0]} +head_node_ip=$(getent hosts "$head_node_hostname" | awk '{print $1}') +RAY_HEAD_NODE_IP="$head_node_ip:6379" +echo "Ray head node IP address: $RAY_HEAD_NODE_IP\"""" + + def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benchmark_name: str) -> str: lines = [] lines.append(f"#SBATCH --job-name=gym-{benchmark_name}") @@ -132,6 +146,26 @@ def _build_vllm_command(service: VllmServiceConfig) -> str: return cmd +def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str: + # Uses vLLM's own Ray *core* executor to span nodes (--distributed-executor-backend ray). + # This is not the ray.serve library - no Serve deployment/ingress/HTTP proxy is involved. + inner_cmd = _build_vllm_command(service) + " --distributed-executor-backend ray" + if service.number_of_instances > 1: + inner_cmd += " --data-parallel-backend ray" + # ray symmetric-run starts/joins a Ray cluster across every task and runs the entrypoint + # only on the elected head node, mirroring scripts/sbatch_base.sh. + return ( + "bash -lc '\n" + " ray symmetric-run \\\n" + ' --address "$RAY_HEAD_NODE_IP" \\\n' + f" --min-nodes {total_nodes} \\\n" + " --num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} \\\n" + " --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE} \\\n" + f" -- {inner_cmd}\n" + "'" + ) + + def _build_ray_command(_service: RayServiceConfig) -> str: return "ray start --head" @@ -142,6 +176,18 @@ def _build_ray_command(_service: RayServiceConfig) -> str: } +def _uses_ray_distributed_backend(service: VllmServiceConfig | RayServiceConfig) -> bool: + return isinstance(service, VllmServiceConfig) and ( + service.distributed_backend is not None and service.distributed_backend.type == "ray" + ) + + +def _build_service_command(service: VllmServiceConfig | RayServiceConfig, total_nodes: int) -> str: + if _uses_ray_distributed_backend(service): + return _build_vllm_ray_command(service, total_nodes) + return _BUILDERS[type(service)](service) + + def _node_totals(compute: SlurmComputeConfig) -> tuple[int, int]: total_nodes = sum(pool.nodes for pool in compute.node_pools.values()) total_ntasks = sum(pool.nodes * pool.ntasks_per_node for pool in compute.node_pools.values()) @@ -160,11 +206,13 @@ def build_sbatch_script( total_nodes, total_ntasks = _node_totals(compute) is_multi_node = total_nodes > 1 + ray_prelude = _RAY_PRELUDE if any(_uses_ray_distributed_backend(s) for s in config.services.values()) else "" + service_commands = "\n\n".join( _render_service_command( name, service.container, - _BUILDERS[type(service)](service), + _build_service_command(service, total_nodes), service.env or None, service.mounts or None, nodes=total_nodes if is_multi_node else None, @@ -209,6 +257,7 @@ def build_sbatch_script( return _SCRIPT_TEMPLATE.format( directives=directives, + ray_prelude=ray_prelude, service_commands=service_commands, health_checks=health_checks, prepare_command=prepare_command, diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 5343a21f3a6..128d5b6eff2 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -124,13 +124,49 @@ def test_multi_instance_without_backend_raises(): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 4}})) -def test_single_instance_with_backend_raises(): - with pytest.raises(ValidationError, match="should not be set"): - SubmitConfig.model_validate( - _config(services={"svc": {**SERVICE, "distributed_backend": {"type": "vllm_service"}}}) - ) - - def test_number_of_instances_zero_raises(): with pytest.raises(ValidationError): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 0}})) + + +# --------------------------------------------------------------------------- +# ray distributed_backend (multi-node) +# --------------------------------------------------------------------------- + +COMPUTE_MULTI_NODE = { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"compute": {"partition": "batch", "nodes": 2, "gpus_per_node": 4}}, + } +} + + +def test_ray_backend_with_single_instance_accepted(): + # ray is valid even with number_of_instances == 1: it's needed to span nodes for + # a single instance's tensor/pipeline-parallel footprint, not just for data-parallel replicas. + config = SubmitConfig.model_validate( + _config(services={"svc": {**SERVICE, "distributed_backend": {"type": "ray"}}}) + ) + assert config.services["svc"].distributed_backend.type == "ray" + + +def test_vllm_service_backend_on_multi_node_compute_raises(): + with pytest.raises(ValidationError, match="must use distributed_backend"): + SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE}, compute=COMPUTE_MULTI_NODE)) + + +def test_vllm_service_no_backend_on_multi_node_compute_raises(): + with pytest.raises(ValidationError, match="must use distributed_backend"): + SubmitConfig.model_validate(_config(compute=COMPUTE_MULTI_NODE)) + + +def test_ray_backend_on_multi_node_compute_accepted(): + config = SubmitConfig.model_validate( + _config( + services={"svc": {**SERVICE, "distributed_backend": {"type": "ray"}}}, + compute=COMPUTE_MULTI_NODE, + ) + ) + assert config.services["svc"].distributed_backend.type == "ray" diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 11b8961cebf..fe6be3fff3e 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -21,6 +21,7 @@ from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd from nemo_gym.orchestration.executors.slurm_script import ( _build_vllm_command, + _build_vllm_ray_command, _node_totals, _render_directives, _render_pool_directives, @@ -198,6 +199,50 @@ def test_build_vllm_command_pipeline_parallel_1_omits_flag(vllm_service): assert "--pipeline-parallel-size" not in cmd +# --------------------------------------------------------------------------- +# _build_vllm_ray_command +# --------------------------------------------------------------------------- + + +def test_build_vllm_ray_command_uses_ray_distributed_executor(vllm_service): + cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) + assert "--distributed-executor-backend ray" in cmd + assert "vllm serve" in cmd + + +def test_build_vllm_ray_command_wraps_in_symmetric_run(vllm_service): + cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) + assert "ray symmetric-run" in cmd + assert "--min-nodes 2" in cmd + assert '--address "$RAY_HEAD_NODE_IP"' in cmd + + +def test_build_vllm_ray_command_multi_instance_adds_dp_backend(): + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + number_of_instances=4, + distributed_backend={"type": "ray"}, + ) + cmd = _build_vllm_ray_command(service, total_nodes=2) + assert "--data-parallel-size 4" in cmd + assert "--data-parallel-backend ray" in cmd + + +def test_build_vllm_ray_command_single_instance_omits_dp_backend(vllm_service): + cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) + assert "--data-parallel-backend" not in cmd + + +def test_build_vllm_ray_command_not_ray_serve_library(): + # Sanity check the plan constraint: this must not shell out to `serve` / ray.serve. + service = VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model") + cmd = _build_vllm_ray_command(service, total_nodes=2) + assert "ray.serve" not in cmd + assert "serve.run" not in cmd + + # --------------------------------------------------------------------------- # render_gym_cmd # --------------------------------------------------------------------------- @@ -697,6 +742,7 @@ def _multi_node_config(): "container": "vllm:latest", "model": "org/model", "tensor_parallel_size": 8, + "distributed_backend": {"type": "ray"}, } }, "compute": { @@ -767,6 +813,29 @@ def test_build_sbatch_script_single_node_pool_omits_node_flags_from_srun(bench_d assert "--nodes=" not in driver_line +# --------------------------------------------------------------------------- +# build_sbatch_script — ray prelude +# --------------------------------------------------------------------------- + + +def test_build_sbatch_script_ray_backend_adds_head_node_prelude(bench_dir): + config = _multi_node_config() + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + assert "scontrol show hostnames" in script + assert "RAY_HEAD_NODE_IP" in script + assert "ray symmetric-run" in script + + +def test_build_sbatch_script_vllm_service_backend_omits_ray_prelude(submit_config, bench_dir): + benchmark = submit_config.driver.benchmarks["gsm8k"] + compute = next(iter(submit_config.compute.values())) + script = build_sbatch_script(submit_config, "gsm8k", benchmark, compute, bench_dir) + assert "scontrol show hostnames" not in script + assert "ray symmetric-run" not in script + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From 81a26632b15dd3620cb657df1c3b7377d3afc954 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 17 Aug 2026 20:42:56 -0400 Subject: [PATCH 54/66] Support old ray command Signed-off-by: Onur Yilmaz --- .../orchestration/executors/slurm_script.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 38475823147..77b82e76597 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -152,16 +152,28 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str inner_cmd = _build_vllm_command(service) + " --distributed-executor-backend ray" if service.number_of_instances > 1: inner_cmd += " --data-parallel-backend ray" + resource_flags = ( + "--num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE}" + ) # ray symmetric-run starts/joins a Ray cluster across every task and runs the entrypoint - # only on the elected head node, mirroring scripts/sbatch_base.sh. + # only on the elected head node, mirroring scripts/sbatch_base.sh. It requires Ray >= 2.50, so + # containers with an older pin fall back to manually starting head/worker Ray processes, keyed + # on Slurm's per-node task rank ($SLURM_NODEID). vLLM's Ray executor blocks on placement-group + # scheduling until every node's GPUs join, so the fallback needs no separate cluster-ready wait. return ( "bash -lc '\n" - " ray symmetric-run \\\n" - ' --address "$RAY_HEAD_NODE_IP" \\\n' - f" --min-nodes {total_nodes} \\\n" - " --num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} \\\n" - " --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE} \\\n" - f" -- {inner_cmd}\n" + " if ray symmetric-run --help >/dev/null 2>&1; then\n" + " ray symmetric-run \\\n" + ' --address "$RAY_HEAD_NODE_IP" \\\n' + f" --min-nodes {total_nodes} \\\n" + f" {resource_flags} \\\n" + f" -- {inner_cmd}\n" + ' elif [ "$SLURM_NODEID" = "0" ]; then\n' + f" ray start --head --port=6379 {resource_flags}\n" + f" {inner_cmd}\n" + " else\n" + f' ray start --address="$RAY_HEAD_NODE_IP" {resource_flags} --block\n' + " fi\n" "'" ) From 70a90894ba12f99f3bcacfb9921d92f5e58f53bc Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 18 Aug 2026 12:45:34 -0400 Subject: [PATCH 55/66] Fix the deployment issue Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/executors/slurm_script.py | 4 ++++ tests/unit_tests/test_slurm_script.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 77b82e76597..a43f22b81ff 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -160,8 +160,12 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str # containers with an older pin fall back to manually starting head/worker Ray processes, keyed # on Slurm's per-node task rank ($SLURM_NODEID). vLLM's Ray executor blocks on placement-group # scheduling until every node's GPUs join, so the fallback needs no separate cluster-ready wait. + # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle the ray CLI - vLLM only + # needs ray as a runtime dependency when the ray executor backend is actually selected - so + # install it on the fly if it's missing before relying on either code path above. return ( "bash -lc '\n" + ' command -v ray >/dev/null 2>&1 || pip install -q "ray[default]"\n' " if ray symmetric-run --help >/dev/null 2>&1; then\n" " ray symmetric-run \\\n" ' --address "$RAY_HEAD_NODE_IP" \\\n' diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index fe6be3fff3e..3c8fcc2143b 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -243,6 +243,12 @@ def test_build_vllm_ray_command_not_ray_serve_library(): assert "serve.run" not in cmd +def test_build_vllm_ray_command_installs_ray_if_missing(vllm_service): + # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle the ray CLI. + cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) + assert 'command -v ray >/dev/null 2>&1 || pip install -q "ray[default]"' in cmd + + # --------------------------------------------------------------------------- # render_gym_cmd # --------------------------------------------------------------------------- From 878001badf7e3aac303c4956c31d4cb2c7bd4722 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 18 Aug 2026 13:20:55 -0400 Subject: [PATCH 56/66] Fix the issue Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/executors/slurm_script.py | 2 +- tests/unit_tests/test_slurm_script.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index a43f22b81ff..10118d228f9 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -58,7 +58,7 @@ nodes_array=($nodes) head_node_hostname=${nodes_array[0]} head_node_ip=$(getent hosts "$head_node_hostname" | awk '{print $1}') -RAY_HEAD_NODE_IP="$head_node_ip:6379" +export RAY_HEAD_NODE_IP="$head_node_ip:6379" echo "Ray head node IP address: $RAY_HEAD_NODE_IP\"""" diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 3c8fcc2143b..9fd7726a35b 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -830,8 +830,10 @@ def test_build_sbatch_script_ray_backend_adds_head_node_prelude(bench_dir): compute = next(iter(config.compute.values())) script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) assert "scontrol show hostnames" in script - assert "RAY_HEAD_NODE_IP" in script assert "ray symmetric-run" in script + # Must be exported: it's read inside a separate `srun ... bash -lc` subprocess, which only + # inherits *exported* environment variables, not plain shell variables from the parent script. + assert 'export RAY_HEAD_NODE_IP="$head_node_ip:6379"' in script def test_build_sbatch_script_vllm_service_backend_omits_ray_prelude(submit_config, bench_dir): From 73e5cbbb562bb21d1c0c362a1ce21471b569af92 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 18 Aug 2026 13:39:45 -0400 Subject: [PATCH 57/66] Add DP support Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 21 +++++- .../orchestration/executors/slurm_script.py | 73 ++++++++++++++++--- tests/unit_tests/test_orchestration_api.py | 25 +++++++ tests/unit_tests/test_slurm_script.py | 67 ++++++++++++----- 4 files changed, 155 insertions(+), 31 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index f36380054d1..a9bcacf8d32 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -192,9 +192,10 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": sole_compute = next(iter(compute_names)) compute = self.compute[sole_compute] - is_multi_node = ( - isinstance(compute, SlurmComputeConfig) and sum(p.nodes for p in compute.node_pools.values()) > 1 + total_nodes = ( + sum(p.nodes for p in compute.node_pools.values()) if isinstance(compute, SlurmComputeConfig) else 1 ) + is_multi_node = total_nodes > 1 for service_name, service in self.services.items(): if service.placement is None: @@ -212,8 +213,20 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": ): raise ValueError( f"Service '{service_name}' is placed on compute '{sole_compute}', which spans multiple nodes " - f"({sum(p.nodes for p in compute.node_pools.values())} total). vLLM services on multi-node " - "compute must use distributed_backend: {type: ray}." + f"({total_nodes} total). vLLM services on multi-node compute must use " + "distributed_backend: {type: ray}." + ) + + if ( + is_multi_node + and isinstance(service, VllmServiceConfig) + and service.number_of_instances > 1 + and service.number_of_instances % total_nodes != 0 + ): + raise ValueError( + f"Service '{service_name}' has number_of_instances={service.number_of_instances}, which must " + f"be evenly divisible by the number of nodes ({total_nodes}) for multi-node data-parallel " + "deployment - each node hosts an equal share of the data-parallel replicas." ) if self.driver.policy_model is not None: diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 10118d228f9..79d5330af7b 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -50,16 +50,20 @@ """ -# Resolves the Ray head node from Slurm's node list, mirroring scripts/sbatch_base.sh. +# Resolves the head node from Slurm's node list, mirroring scripts/sbatch_base.sh. # Runs once in the main sbatch shell (before any srun step) so $SLURM_JOB_NODELIST is available. +# Used by both multi-node code paths below: HEAD_NODE_IP for vLLM's own --data-parallel-address +# (data-parallel spanning nodes), RAY_HEAD_NODE_IP for `ray start`/`ray symmetric-run` (tensor/ +# pipeline-parallel spanning nodes). _RAY_PRELUDE = """\ -# Resolve Ray head node IP for multi-node vLLM services (distributed_backend: ray). +# Resolve the head node IP for multi-node vLLM services (distributed_backend: ray). nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") nodes_array=($nodes) head_node_hostname=${nodes_array[0]} head_node_ip=$(getent hosts "$head_node_hostname" | awk '{print $1}') +export HEAD_NODE_IP="$head_node_ip" export RAY_HEAD_NODE_IP="$head_node_ip:6379" -echo "Ray head node IP address: $RAY_HEAD_NODE_IP\"""" +echo "Head node IP address: $HEAD_NODE_IP\"""" def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benchmark_name: str) -> str: @@ -131,7 +135,7 @@ def _render_service_command( ) -def _build_vllm_command(service: VllmServiceConfig) -> str: +def _vllm_base_flags(service: VllmServiceConfig) -> str: cmd = ( f"vllm serve {shlex.quote(service.model)}" f" --port {service.port}" @@ -139,6 +143,11 @@ def _build_vllm_command(service: VllmServiceConfig) -> str: ) if service.pipeline_parallel_size > 1: cmd += f" --pipeline-parallel-size {service.pipeline_parallel_size}" + return cmd + + +def _build_vllm_command(service: VllmServiceConfig) -> str: + cmd = _vllm_base_flags(service) if service.number_of_instances > 1: cmd += f" --data-parallel-size {service.number_of_instances}" if service.trust_remote_code: @@ -146,12 +155,11 @@ def _build_vllm_command(service: VllmServiceConfig) -> str: return cmd -def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str: - # Uses vLLM's own Ray *core* executor to span nodes (--distributed-executor-backend ray). - # This is not the ray.serve library - no Serve deployment/ingress/HTTP proxy is involved. +def _build_vllm_multi_node_tp_command(service: VllmServiceConfig, total_nodes: int) -> str: + # A single instance's tensor/pipeline-parallel footprint spans nodes. Uses vLLM's own Ray + # *core* executor (--distributed-executor-backend ray) - not the ray.serve library, no Serve + # deployment/ingress/HTTP proxy is involved. inner_cmd = _build_vllm_command(service) + " --distributed-executor-backend ray" - if service.number_of_instances > 1: - inner_cmd += " --data-parallel-backend ray" resource_flags = ( "--num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE}" ) @@ -182,6 +190,53 @@ def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str ) +def _build_vllm_multi_node_dp_command(service: VllmServiceConfig, total_nodes: int) -> str: + # Data-parallel replicas span nodes. vLLM's Ray-based DP auto-placement doesn't spread ranks + # across physical nodes - launching a single `vllm serve --data-parallel-size N` from one node + # only sees that node's own GPUs when placing DP ranks. Real multi-node DP instead needs one + # `vllm serve` invocation per node: the head node's serves the OpenAI API and coordinates, + # worker nodes run `--headless` with a --data-parallel-start-rank offset. This is vLLM's + # documented multi-node data-parallel deployment pattern and doesn't use Ray at all - each + # node's tensor-parallel ranks stay local via vLLM's default (mp) executor backend. + if service.number_of_instances % total_nodes != 0: + raise ValueError( + f"number_of_instances ({service.number_of_instances}) must be evenly divisible by the number of " + f"nodes ({total_nodes}) for multi-node data-parallel deployment." + ) + dp_size_local = service.number_of_instances // total_nodes + common = _vllm_base_flags(service) + dp_flags = ( + f" --data-parallel-size {service.number_of_instances}" + f" --data-parallel-size-local {dp_size_local}" + ' --data-parallel-address "$HEAD_NODE_IP"' + " --data-parallel-rpc-port 13345" + ) + trust_flag = " --trust-remote-code" if service.trust_remote_code else "" + head_cmd = common + dp_flags + trust_flag + worker_cmd = ( + common + + dp_flags + + trust_flag + + " --headless" + + f" --data-parallel-start-rank $(( SLURM_NODEID * {dp_size_local} ))" + ) + return ( + "bash -lc '\n" + ' if [ "$SLURM_NODEID" = "0" ]; then\n' + f" {head_cmd}\n" + " else\n" + f" {worker_cmd}\n" + " fi\n" + "'" + ) + + +def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str: + if service.number_of_instances > 1: + return _build_vllm_multi_node_dp_command(service, total_nodes) + return _build_vllm_multi_node_tp_command(service, total_nodes) + + def _build_ray_command(_service: RayServiceConfig) -> str: return "ray start --head" diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 128d5b6eff2..0480149cab5 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -170,3 +170,28 @@ def test_ray_backend_on_multi_node_compute_accepted(): ) ) assert config.services["svc"].distributed_backend.type == "ray" + + +def test_ray_backend_multi_node_dp_evenly_divisible_accepted(): + # COMPUTE_MULTI_NODE has 2 nodes; 4 instances split evenly (2 per node). + config = SubmitConfig.model_validate( + _config( + services={ + "svc": {**SERVICE, "number_of_instances": 4, "distributed_backend": {"type": "ray"}}, + }, + compute=COMPUTE_MULTI_NODE, + ) + ) + assert config.services["svc"].number_of_instances == 4 + + +def test_ray_backend_multi_node_dp_uneven_split_raises(): + with pytest.raises(ValidationError, match="evenly divisible"): + SubmitConfig.model_validate( + _config( + services={ + "svc": {**SERVICE, "number_of_instances": 3, "distributed_backend": {"type": "ray"}}, + }, + compute=COMPUTE_MULTI_NODE, + ) + ) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 9fd7726a35b..88551c72220 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -200,7 +200,7 @@ def test_build_vllm_command_pipeline_parallel_1_omits_flag(vllm_service): # --------------------------------------------------------------------------- -# _build_vllm_ray_command +# _build_vllm_ray_command - single instance, TP/PP spans nodes (uses Ray core) # --------------------------------------------------------------------------- @@ -217,7 +217,27 @@ def test_build_vllm_ray_command_wraps_in_symmetric_run(vllm_service): assert '--address "$RAY_HEAD_NODE_IP"' in cmd -def test_build_vllm_ray_command_multi_instance_adds_dp_backend(): +def test_build_vllm_ray_command_not_ray_serve_library(): + # Sanity check the plan constraint: this must not shell out to `serve` / ray.serve. + service = VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model") + cmd = _build_vllm_ray_command(service, total_nodes=2) + assert "ray.serve" not in cmd + assert "serve.run" not in cmd + + +def test_build_vllm_ray_command_installs_ray_if_missing(vllm_service): + # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle the ray CLI. + cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) + assert 'command -v ray >/dev/null 2>&1 || pip install -q "ray[default]"' in cmd + + +# --------------------------------------------------------------------------- +# _build_vllm_ray_command - multiple instances (data parallel) span nodes +# --------------------------------------------------------------------------- + + +def test_build_vllm_ray_command_dp_does_not_use_ray(): + # Multi-node DP uses vLLM's own --data-parallel-address/--headless coordination, not Ray. service = VllmServiceConfig( type="vllm", container="vllm:latest", @@ -226,27 +246,38 @@ def test_build_vllm_ray_command_multi_instance_adds_dp_backend(): distributed_backend={"type": "ray"}, ) cmd = _build_vllm_ray_command(service, total_nodes=2) - assert "--data-parallel-size 4" in cmd - assert "--data-parallel-backend ray" in cmd - + assert "ray" not in cmd + assert "symmetric-run" not in cmd -def test_build_vllm_ray_command_single_instance_omits_dp_backend(vllm_service): - cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) - assert "--data-parallel-backend" not in cmd - -def test_build_vllm_ray_command_not_ray_serve_library(): - # Sanity check the plan constraint: this must not shell out to `serve` / ray.serve. - service = VllmServiceConfig(type="vllm", container="vllm:latest", model="org/model") +def test_build_vllm_ray_command_dp_head_and_worker_branches(): + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + number_of_instances=4, + distributed_backend={"type": "ray"}, + ) cmd = _build_vllm_ray_command(service, total_nodes=2) - assert "ray.serve" not in cmd - assert "serve.run" not in cmd + assert 'if [ "$SLURM_NODEID" = "0" ]; then' in cmd + assert "--headless" in cmd + assert "--data-parallel-size 4" in cmd + assert "--data-parallel-size-local 2" in cmd + assert '--data-parallel-address "$HEAD_NODE_IP"' in cmd + assert "--data-parallel-rpc-port 13345" in cmd + assert "--data-parallel-start-rank $(( SLURM_NODEID * 2 ))" in cmd -def test_build_vllm_ray_command_installs_ray_if_missing(vllm_service): - # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle the ray CLI. - cmd = _build_vllm_ray_command(vllm_service, total_nodes=2) - assert 'command -v ray >/dev/null 2>&1 || pip install -q "ray[default]"' in cmd +def test_build_vllm_ray_command_dp_uneven_split_raises(): + service = VllmServiceConfig( + type="vllm", + container="vllm:latest", + model="org/model", + number_of_instances=3, + distributed_backend={"type": "ray"}, + ) + with pytest.raises(ValueError, match="evenly divisible"): + _build_vllm_ray_command(service, total_nodes=2) # --------------------------------------------------------------------------- From 5d5a327e8524bdd5f9f7d771252dff0826dcfc65 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Tue, 18 Aug 2026 16:35:48 -0400 Subject: [PATCH 58/66] Add number of GPUs check Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 39 ++++++++++++++++ tests/unit_tests/test_orchestration_api.py | 54 ++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index f8e2dd5ab35..38b426723e4 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings from typing import Annotated, Any, Literal from pydantic import BaseModel, ConfigDict, Discriminator, Tag, field_validator, model_validator @@ -193,6 +194,9 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": f"({', '.join(sorted(compute_names))})." ) + if isinstance(service, VllmServiceConfig): + self._validate_vllm_gpu_footprint(service_name, service) + if self.driver.policy_model is not None: if self.driver.policy_model not in self.services: raise ValueError( @@ -216,3 +220,38 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": benchmark.run["policy_api_key"] = "dummy" # pragma: allowlist secret return self + + def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceConfig") -> 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 + ] + if not gpus_per_node_values: + return + + max_gpus_per_node = max(gpus_per_node_values) + gpus_needed = service.tensor_parallel_size * service.pipeline_parallel_size * service.number_of_instances + + if gpus_needed > max_gpus_per_node: + raise ValueError( + f"Service '{service_name}' requires {gpus_needed} GPUs " + 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}), which exceeds the largest available " + f"node pool's gpus_per_node ({max_gpus_per_node}) on compute '{service.placement}'. " + "Multi-node vLLM services are not supported yet by the 'vllm_service' distributed backend; " + "reduce number_of_instances/tensor_parallel_size/pipeline_parallel_size to fit on a single node." + ) + elif gpus_needed < max_gpus_per_node: + warnings.warn( + f"Service '{service_name}' requires {gpus_needed} GPUs " + 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}) but compute '{service.placement}' allocates " + f"nodes with {max_gpus_per_node} GPUs each, leaving {max_gpus_per_node - gpus_needed} GPU(s) idle. " + "Increase number_of_instances/tensor_parallel_size or reduce gpus_per_node to use the full node.", + stacklevel=2, + ) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 5343a21f3a6..496bd99b736 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -134,3 +134,57 @@ def test_single_instance_with_backend_raises(): def test_number_of_instances_zero_raises(): with pytest.raises(ValidationError): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 0}})) + + +# --------------------------------------------------------------------------- +# GPU footprint vs node pool capacity +# --------------------------------------------------------------------------- + +COMPUTE_8_GPUS_PER_NODE = { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"compute": {"partition": "batch", "gpus_per_node": 8}}, + } +} + + +def test_gpu_footprint_exact_fit_accepted(): + service = { + **SERVICE, + "tensor_parallel_size": 2, + "number_of_instances": 4, + "distributed_backend": {"type": "vllm_service"}, + } + config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) + assert config.services["svc"].number_of_instances == 4 + + +def test_gpu_footprint_exceeds_node_raises(): + service = { + **SERVICE, + "tensor_parallel_size": 2, + "number_of_instances": 8, + "distributed_backend": {"type": "vllm_service"}, + } + with pytest.raises(ValidationError, match="exceeds the largest available"): + SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) + + +def test_gpu_footprint_underutilized_warns(): + service = { + **SERVICE, + "tensor_parallel_size": 2, + "number_of_instances": 2, + "distributed_backend": {"type": "vllm_service"}, + } + with pytest.warns(UserWarning, match="leaving 4 GPU"): + config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) + assert config.services["svc"].number_of_instances == 2 + + +def test_gpu_footprint_no_node_pools_skips_validation(): + # Default COMPUTE fixture has no node_pools, so nothing to validate against. + config = SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE})) + assert config.services["svc"].number_of_instances == 4 From 5b5dddf04cbb7feecf828e15ed9c3f89fbbd3594 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 18:41:31 -0400 Subject: [PATCH 59/66] Remove the multi node pieces from this PR Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance.yaml | 2 +- examples/slurm_vllm_multi_instance_onur.yaml | 48 ++++++ examples/slurm_vllm_ray_multi_node_onur.yaml | 50 ++++++ .../slurm_vllm_ray_serve_multi_node_onur.yaml | 50 ++++++ .../orchestration/executors/slurm_script.py | 25 +-- tests/unit_tests/test_slurm_script.py | 144 ------------------ 6 files changed, 152 insertions(+), 167 deletions(-) create mode 100644 examples/slurm_vllm_multi_instance_onur.yaml create mode 100644 examples/slurm_vllm_ray_multi_node_onur.yaml create mode 100644 examples/slurm_vllm_ray_serve_multi_node_onur.yaml diff --git a/examples/slurm_vllm_multi_instance.yaml b/examples/slurm_vllm_multi_instance.yaml index 4a19a1b22da..22b1091c8cf 100644 --- a/examples/slurm_vllm_multi_instance.yaml +++ b/examples/slurm_vllm_multi_instance.yaml @@ -29,7 +29,7 @@ services: driver: policy_model: vllm_model - container: python:3.13 + container: python:3.12 gym_install: ref: main benchmarks: diff --git a/examples/slurm_vllm_multi_instance_onur.yaml b/examples/slurm_vllm_multi_instance_onur.yaml new file mode 100644 index 00000000000..c1c1b865f3f --- /dev/null +++ b/examples/slurm_vllm_multi_instance_onur.yaml @@ -0,0 +1,48 @@ +compute: + cluster: + type: slurm + hostname: cw-dfw-cs-001-login-01 + walltime: "01:00:00" + account: coreai_dlalgo_nemorl + node_pools: + compute: + partition: batch + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:latest + type: vllm + model: Qwen/Qwen2.5-1.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 4 # needs the ray backend since the footprint exceeds one node + port: 8000 + distributed_backend: + type: vllm_service + health_check: + timeout_seconds: 1200 + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + +driver: + policy_model: vllm_model + container: python:3.13 + gym_install: + ref: onur/multi-instance-vllm-service + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/gpqa/config.yaml + run: + split: benchmark + config_paths: + - benchmarks/gpqa/config.yaml + +job: + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_multi_node_onur.yaml b/examples/slurm_vllm_ray_multi_node_onur.yaml new file mode 100644 index 00000000000..200f0af0eb2 --- /dev/null +++ b/examples/slurm_vllm_ray_multi_node_onur.yaml @@ -0,0 +1,50 @@ +compute: + cluster: + type: slurm + hostname: cw-dfw-cs-001-login-01 + walltime: "01:00:00" + account: coreai_dlalgo_nemorl + node_pools: + compute: + partition: batch + nodes: 3 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:latest + type: vllm + model: Qwen/Qwen2.5-1.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 12 # needs the ray backend since the footprint exceeds one node + distributed_backend: + type: ray # vLLM's Ray core executor (--distributed-executor-backend ray). + # Not the ray.serve library - just lets vLLM's instances span nodes. + # A separate ray_serve backend may be added later. + port: 8000 + health_check: + timeout_seconds: 1200 + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + +driver: + policy_model: vllm_model + container: python:3.13 + gym_install: + ref: onur/ray-service-multi-gpu + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/gpqa/config.yaml + run: + split: benchmark + config_paths: + - benchmarks/gpqa/config.yaml + +job: + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_serve_multi_node_onur.yaml b/examples/slurm_vllm_ray_serve_multi_node_onur.yaml new file mode 100644 index 00000000000..22f554ca65b --- /dev/null +++ b/examples/slurm_vllm_ray_serve_multi_node_onur.yaml @@ -0,0 +1,50 @@ +compute: + cluster: + type: slurm + hostname: cw-dfw-cs-001-login-01 + walltime: "01:00:00" + account: coreai_dlalgo_nemorl + node_pools: + compute: + partition: batch + nodes: 2 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:latest + type: vllm + model: Qwen/Qwen2.5-1.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 8 # needs the ray backend since the footprint exceeds one node + distributed_backend: + type: ray_serve # vLLM's Ray core executor (--distributed-executor-backend ray). + # Not the ray.serve library - just lets vLLM's instances span nodes. + # A separate ray_serve backend may be added later. + port: 8000 + health_check: + timeout_seconds: 1200 + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + +driver: + policy_model: vllm_model + container: python:3.13 + gym_install: + ref: onur/add-ray-serve-multi-gpu + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/gpqa/config.yaml + run: + split: benchmark + config_paths: + - benchmarks/gpqa/config.yaml + +job: + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 5396893b2c2..982f7d8924e 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -100,19 +100,16 @@ def _render_service_command( command: str, env: dict[str, str] | None = None, mounts: list[str] | None = None, - nodes: int | None = None, - ntasks: int | None = None, ) -> str: var = bash_var(name) env_prefix = _resolve_env(env) if env else "" - node_flags = f" --nodes={nodes} --ntasks={ntasks}" if (nodes is not None and nodes > 1) else "" mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else "" # --overlap lets this step share the allocation with other concurrent steps (driver + services). # --no-container-mount-home avoids polluting the container with host home directory contents. # PID is captured so the health check can detect early service death. return ( f"# service: {name}\n" - f"{env_prefix}srun --overlap --no-container-mount-home{node_flags}{mounts_flag} --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" + f"{env_prefix}srun --overlap --no-container-mount-home{mounts_flag} --container-image={shlex.quote(container)} --output=logs/{name}.log {command} &\n" f"{var}_PID=$!" ) @@ -142,12 +139,6 @@ def _build_ray_command(_service: RayServiceConfig) -> str: } -def _node_totals(compute: SlurmComputeConfig) -> tuple[int, int]: - total_nodes = sum(pool.nodes for pool in compute.node_pools.values()) - total_ntasks = sum(pool.nodes * pool.ntasks_per_node for pool in compute.node_pools.values()) - return total_nodes, total_ntasks - - def build_sbatch_script( config: SubmitConfig, benchmark_name: str, @@ -157,18 +148,9 @@ def build_sbatch_script( ) -> str: directives = _render_directives(compute, remote_bench_dir, benchmark_name) - total_nodes, total_ntasks = _node_totals(compute) - is_multi_node = total_nodes > 1 - service_commands = "\n\n".join( _render_service_command( - name, - service.container, - _BUILDERS[type(service)](service), - service.env or None, - service.mounts or None, - nodes=total_nodes if is_multi_node else None, - ntasks=total_ntasks if is_multi_node else None, + name, service.container, _BUILDERS[type(service)](service), service.env or None, service.mounts or None ) for name, service in config.services.items() ) @@ -197,13 +179,12 @@ def build_sbatch_script( ) prepare_command = "" driver_env_prefix = _resolve_env(config.driver.env) if config.driver.env else "" - driver_node_flags = " --nodes=1 --ntasks=1" if is_multi_node else "" driver_mounts_flag = ( f" --container-mounts={','.join(shlex.quote(m) for m in config.driver.mounts)}" if config.driver.mounts else "" ) driver_command = ( f"{gym_cmd}\n" - f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_node_flags}{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " + f"{driver_env_prefix}srun --overlap --no-container-mount-home{driver_mounts_flag} --container-image={shlex.quote(config.driver.container)} " f"--output=logs/driver.log {entrypoint}" ) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 11b8961cebf..30396e81491 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -21,7 +21,6 @@ from nemo_gym.orchestration.executors.script_templates import render_driver_entrypoint, render_gym_cmd from nemo_gym.orchestration.executors.slurm_script import ( _build_vllm_command, - _node_totals, _render_directives, _render_pool_directives, _render_service_command, @@ -624,149 +623,6 @@ def test_validate_mounts_no_mounts_passes(): conn.run.assert_not_called() -# --------------------------------------------------------------------------- -# _render_service_command — multi-node flags -# --------------------------------------------------------------------------- - - -def test_render_service_command_multi_node_adds_nodes_and_ntasks(): - out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model", nodes=4, ntasks=4) - assert "--nodes=4" in out - assert "--ntasks=4" in out - - -def test_render_service_command_multi_node_flags_before_container_image(): - out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model", nodes=4, ntasks=4) - assert out.index("--nodes=4") < out.index("--container-image=") - - -def test_render_service_command_single_node_omits_node_flags(): - out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model", nodes=1, ntasks=1) - assert "--nodes=" not in out - assert "--ntasks=" not in out - - -def test_render_service_command_no_nodes_kwarg_omits_node_flags(): - out = _render_service_command("vllm_model", "vllm:latest", "vllm serve model") - assert "--nodes=" not in out - assert "--ntasks=" not in out - - -# --------------------------------------------------------------------------- -# _node_totals -# --------------------------------------------------------------------------- - - -def test_node_totals_empty_pools(): - compute = SlurmComputeConfig(type="slurm", account="acct") - assert _node_totals(compute) == (0, 0) - - -def test_node_totals_single_pool(): - compute = SlurmComputeConfig( - type="slurm", - account="acct", - node_pools={"main": NodePool(partition="gpu", nodes=4, ntasks_per_node=2)}, - ) - assert _node_totals(compute) == (4, 8) - - -def test_node_totals_multiple_pools(): - compute = SlurmComputeConfig( - type="slurm", - account="acct", - node_pools={ - "gpu": NodePool(partition="gpu", nodes=4, ntasks_per_node=1), - "cpu": NodePool(partition="cpu", nodes=2, ntasks_per_node=2), - }, - ) - assert _node_totals(compute) == (6, 8) - - -# --------------------------------------------------------------------------- -# build_sbatch_script — multi-node srun flags -# --------------------------------------------------------------------------- - - -def _multi_node_config(): - return SubmitConfig.model_validate( - { - "services": { - "vllm_model": { - "type": "vllm", - "container": "vllm:latest", - "model": "org/model", - "tensor_parallel_size": 8, - } - }, - "compute": { - "cluster": { - "type": "slurm", - "account": "my-account", - "hostname": "foo", - "node_pools": {"main": {"partition": "gpu", "nodes": 4, "ntasks_per_node": 1}}, - } - }, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - - -def test_build_sbatch_script_multi_node_vllm_srun_gets_node_flags(bench_dir): - config = _multi_node_config() - benchmark = config.driver.benchmarks["gsm8k"] - compute = next(iter(config.compute.values())) - script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) - assert "--nodes=4" in vllm_line - assert "--ntasks=4" in vllm_line - - -def test_build_sbatch_script_multi_node_driver_srun_gets_nodes_1(bench_dir): - config = _multi_node_config() - benchmark = config.driver.benchmarks["gsm8k"] - compute = next(iter(config.compute.values())) - script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - driver_line = next(line for line in script.splitlines() if "python:3.12" in line) - assert "--nodes=1" in driver_line - assert "--ntasks=1" in driver_line - - -def test_build_sbatch_script_multi_node_node_flags_before_container_image(bench_dir): - config = _multi_node_config() - benchmark = config.driver.benchmarks["gsm8k"] - compute = next(iter(config.compute.values())) - script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) - assert vllm_line.index("--nodes=4") < vllm_line.index("--container-image=") - - -def test_build_sbatch_script_single_node_pool_omits_node_flags_from_srun(bench_dir): - config = SubmitConfig.model_validate( - { - "services": {"vllm_model": {"type": "vllm", "container": "vllm:latest", "model": "org/model"}}, - "compute": { - "cluster": { - "type": "slurm", - "account": "my-account", - "hostname": "foo", - "node_pools": {"main": {"partition": "gpu", "nodes": 1, "ntasks_per_node": 4}}, - } - }, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - benchmark = config.driver.benchmarks["gsm8k"] - compute = next(iter(config.compute.values())) - script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) - vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) - driver_line = next(line for line in script.splitlines() if "python:3.12" in line) - assert "--nodes=" not in vllm_line - assert "--nodes=" not in driver_line - - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- From d82c18d0a00c0317a8a0c692bdb623e950a771f2 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 18:51:17 -0400 Subject: [PATCH 60/66] Remove the vllm_service type Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance.yaml | 2 -- nemo_gym/orchestration/api.py | 13 +++++-------- tests/unit_tests/test_orchestration_api.py | 22 +++++++++++----------- tests/unit_tests/test_slurm_script.py | 2 +- 4 files changed, 17 insertions(+), 22 deletions(-) diff --git a/examples/slurm_vllm_multi_instance.yaml b/examples/slurm_vllm_multi_instance.yaml index 22b1091c8cf..179fb4dadab 100644 --- a/examples/slurm_vllm_multi_instance.yaml +++ b/examples/slurm_vllm_multi_instance.yaml @@ -19,8 +19,6 @@ services: trust_remote_code: true tensor_parallel_size: 2 number_of_instances: 4 # 4 engine replicas, total 8 GPUs - distributed_backend: - type: vllm_service # vLLM native --data-parallel-size; ray_serve/dynamo coming later port: 8000 health_check: timeout_seconds: 1200 diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 38b426723e4..21d4e85311f 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -52,12 +52,12 @@ class BaseModelServiceConfig(BaseServiceConfig): class VllmServiceDistributedBackend(_StrictModel): """Use vLLM's native data-parallel multi-instance (--data-parallel-size N).""" - type: Literal["vllm_service"] = "vllm_service" + type: Literal["mp"] = "mp" # Future backends: add Annotated[RayServeDistributedBackend, Tag("ray_serve")], etc. DistributedBackendConfig = Annotated[ - Annotated[VllmServiceDistributedBackend, Tag("vllm_service")], + Annotated[VllmServiceDistributedBackend, Tag("mp")], Discriminator("type"), ] @@ -79,13 +79,10 @@ def _validate_number_of_instances(cls, v: int) -> int: @model_validator(mode="after") def _validate_distributed_backend(self) -> "VllmServiceConfig": - if self.number_of_instances > 1 and self.distributed_backend is None: - raise ValueError( - f"distributed_backend must be set when number_of_instances > 1 (got {self.number_of_instances}). " - "Use: distributed_backend: {type: vllm_service}" - ) if self.number_of_instances == 1 and self.distributed_backend is not None: raise ValueError("distributed_backend should not be set when number_of_instances == 1") + if self.number_of_instances > 1 and self.distributed_backend is None: + self.distributed_backend = VllmServiceDistributedBackend() return self @model_validator(mode="after") @@ -242,7 +239,7 @@ def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceC f"pipeline_parallel_size={service.pipeline_parallel_size} x " f"number_of_instances={service.number_of_instances}), which exceeds the largest available " f"node pool's gpus_per_node ({max_gpus_per_node}) on compute '{service.placement}'. " - "Multi-node vLLM services are not supported yet by the 'vllm_service' distributed backend; " + "Multi-node vLLM services are not supported yet by the 'mp' distributed backend; " "reduce number_of_instances/tensor_parallel_size/pipeline_parallel_size to fit on a single node." ) elif gpus_needed < max_gpus_per_node: diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 496bd99b736..2efde0532af 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -102,7 +102,7 @@ def test_service_unknown_field_raises(): # number_of_instances / distributed_backend # --------------------------------------------------------------------------- -_MULTI_SERVICE = {**SERVICE, "number_of_instances": 4, "distributed_backend": {"type": "vllm_service"}} +_MULTI_SERVICE = {**SERVICE, "number_of_instances": 4, "distributed_backend": {"type": "mp"}} def test_number_of_instances_with_backend_accepted(): @@ -110,7 +110,7 @@ def test_number_of_instances_with_backend_accepted(): svc = config.services["svc"] assert svc.number_of_instances == 4 assert svc.distributed_backend is not None - assert svc.distributed_backend.type == "vllm_service" + assert svc.distributed_backend.type == "mp" def test_number_of_instances_defaults_to_1(): @@ -119,16 +119,16 @@ def test_number_of_instances_defaults_to_1(): assert config.services["svc"].distributed_backend is None -def test_multi_instance_without_backend_raises(): - with pytest.raises(ValidationError, match="distributed_backend must be set"): - SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 4}})) +def test_multi_instance_without_backend_defaults_to_mp(): + config = SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 4}})) + svc = config.services["svc"] + assert svc.distributed_backend is not None + assert svc.distributed_backend.type == "mp" def test_single_instance_with_backend_raises(): with pytest.raises(ValidationError, match="should not be set"): - SubmitConfig.model_validate( - _config(services={"svc": {**SERVICE, "distributed_backend": {"type": "vllm_service"}}}) - ) + SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "distributed_backend": {"type": "mp"}}})) def test_number_of_instances_zero_raises(): @@ -155,7 +155,7 @@ def test_gpu_footprint_exact_fit_accepted(): **SERVICE, "tensor_parallel_size": 2, "number_of_instances": 4, - "distributed_backend": {"type": "vllm_service"}, + "distributed_backend": {"type": "mp"}, } config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) assert config.services["svc"].number_of_instances == 4 @@ -166,7 +166,7 @@ def test_gpu_footprint_exceeds_node_raises(): **SERVICE, "tensor_parallel_size": 2, "number_of_instances": 8, - "distributed_backend": {"type": "vllm_service"}, + "distributed_backend": {"type": "mp"}, } with pytest.raises(ValidationError, match="exceeds the largest available"): SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) @@ -177,7 +177,7 @@ def test_gpu_footprint_underutilized_warns(): **SERVICE, "tensor_parallel_size": 2, "number_of_instances": 2, - "distributed_backend": {"type": "vllm_service"}, + "distributed_backend": {"type": "mp"}, } with pytest.warns(UserWarning, match="leaving 4 GPU"): config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 30396e81491..f139139c47e 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -175,7 +175,7 @@ def test_build_vllm_command_multi_instance(): container="vllm:latest", model="org/model", number_of_instances=4, - distributed_backend={"type": "vllm_service"}, + distributed_backend={"type": "mp"}, ) cmd = _build_vllm_command(service) assert "--data-parallel-size 4" in cmd From 25f6e2f957542fb9d8d04ff57ec877419c2b71af Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 18:54:55 -0400 Subject: [PATCH 61/66] Untrack ray_serve multi-node scratch example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep the file locally but drop it from the branch — it belongs to future multi-node work, not this single-node multi-instance PR. Signed-off-by: Onur Yilmaz --- .../slurm_vllm_ray_serve_multi_node_onur.yaml | 50 ------------------- 1 file changed, 50 deletions(-) delete mode 100644 examples/slurm_vllm_ray_serve_multi_node_onur.yaml diff --git a/examples/slurm_vllm_ray_serve_multi_node_onur.yaml b/examples/slurm_vllm_ray_serve_multi_node_onur.yaml deleted file mode 100644 index 22f554ca65b..00000000000 --- a/examples/slurm_vllm_ray_serve_multi_node_onur.yaml +++ /dev/null @@ -1,50 +0,0 @@ -compute: - cluster: - type: slurm - hostname: cw-dfw-cs-001-login-01 - walltime: "01:00:00" - account: coreai_dlalgo_nemorl - node_pools: - compute: - partition: batch - nodes: 2 - ntasks_per_node: 1 - gpus_per_node: 8 - -services: - vllm_model: - container: vllm/vllm-openai:latest - type: vllm - model: Qwen/Qwen2.5-1.5B-Instruct - trust_remote_code: true - tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 8 # needs the ray backend since the footprint exceeds one node - distributed_backend: - type: ray_serve # vLLM's Ray core executor (--distributed-executor-backend ray). - # Not the ray.serve library - just lets vLLM's instances span nodes. - # A separate ray_serve backend may be added later. - port: 8000 - health_check: - timeout_seconds: 1200 - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - -driver: - policy_model: vllm_model - container: python:3.13 - gym_install: - ref: onur/add-ray-serve-multi-gpu - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - benchmarks: - gpqa: - prepare: - config_paths: - - benchmarks/gpqa/config.yaml - run: - split: benchmark - config_paths: - - benchmarks/gpqa/config.yaml - -job: - output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output From 06c6b4f09477679738a8dc5677b973671a6b70ab Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 18:55:16 -0400 Subject: [PATCH 62/66] Untrack remaining onur scratch examples MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep these locally but drop them from the branch — they're personal scratch configs, not part of this single-node multi-instance PR. Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance_onur.yaml | 48 ------------------- examples/slurm_vllm_ray_multi_node_onur.yaml | 50 -------------------- 2 files changed, 98 deletions(-) delete mode 100644 examples/slurm_vllm_multi_instance_onur.yaml delete mode 100644 examples/slurm_vllm_ray_multi_node_onur.yaml diff --git a/examples/slurm_vllm_multi_instance_onur.yaml b/examples/slurm_vllm_multi_instance_onur.yaml deleted file mode 100644 index c1c1b865f3f..00000000000 --- a/examples/slurm_vllm_multi_instance_onur.yaml +++ /dev/null @@ -1,48 +0,0 @@ -compute: - cluster: - type: slurm - hostname: cw-dfw-cs-001-login-01 - walltime: "01:00:00" - account: coreai_dlalgo_nemorl - node_pools: - compute: - partition: batch - nodes: 1 - ntasks_per_node: 1 - gpus_per_node: 8 - -services: - vllm_model: - container: vllm/vllm-openai:latest - type: vllm - model: Qwen/Qwen2.5-1.5B-Instruct - trust_remote_code: true - tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 4 # needs the ray backend since the footprint exceeds one node - port: 8000 - distributed_backend: - type: vllm_service - health_check: - timeout_seconds: 1200 - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - -driver: - policy_model: vllm_model - container: python:3.13 - gym_install: - ref: onur/multi-instance-vllm-service - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - benchmarks: - gpqa: - prepare: - config_paths: - - benchmarks/gpqa/config.yaml - run: - split: benchmark - config_paths: - - benchmarks/gpqa/config.yaml - -job: - output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_multi_node_onur.yaml b/examples/slurm_vllm_ray_multi_node_onur.yaml deleted file mode 100644 index 200f0af0eb2..00000000000 --- a/examples/slurm_vllm_ray_multi_node_onur.yaml +++ /dev/null @@ -1,50 +0,0 @@ -compute: - cluster: - type: slurm - hostname: cw-dfw-cs-001-login-01 - walltime: "01:00:00" - account: coreai_dlalgo_nemorl - node_pools: - compute: - partition: batch - nodes: 3 - ntasks_per_node: 1 - gpus_per_node: 8 - -services: - vllm_model: - container: vllm/vllm-openai:latest - type: vllm - model: Qwen/Qwen2.5-1.5B-Instruct - trust_remote_code: true - tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 12 # needs the ray backend since the footprint exceeds one node - distributed_backend: - type: ray # vLLM's Ray core executor (--distributed-executor-backend ray). - # Not the ray.serve library - just lets vLLM's instances span nodes. - # A separate ray_serve backend may be added later. - port: 8000 - health_check: - timeout_seconds: 1200 - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - -driver: - policy_model: vllm_model - container: python:3.13 - gym_install: - ref: onur/ray-service-multi-gpu - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - benchmarks: - gpqa: - prepare: - config_paths: - - benchmarks/gpqa/config.yaml - run: - split: benchmark - config_paths: - - benchmarks/gpqa/config.yaml - -job: - output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output From a76b461d3e93ffba54d1c21317372e411e0f2314 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 19:18:46 -0400 Subject: [PATCH 63/66] Remove the distributed param Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance_onur.yaml | 46 +++++++++++++++++ examples/slurm_vllm_ray_multi_node.yaml | 6 +-- examples/slurm_vllm_ray_multi_node_onur.yaml | 46 +++++++++++++++++ .../slurm_vllm_ray_serve_multi_node_onur.yaml | 50 +++++++++++++++++++ nemo_gym/orchestration/api.py | 14 ++---- tests/unit_tests/test_orchestration_api.py | 14 +++--- tests/unit_tests/test_slurm_script.py | 33 ++++++++++++ 7 files changed, 188 insertions(+), 21 deletions(-) create mode 100644 examples/slurm_vllm_multi_instance_onur.yaml create mode 100644 examples/slurm_vllm_ray_multi_node_onur.yaml create mode 100644 examples/slurm_vllm_ray_serve_multi_node_onur.yaml diff --git a/examples/slurm_vllm_multi_instance_onur.yaml b/examples/slurm_vllm_multi_instance_onur.yaml new file mode 100644 index 00000000000..0cb65fb995e --- /dev/null +++ b/examples/slurm_vllm_multi_instance_onur.yaml @@ -0,0 +1,46 @@ +compute: + cluster: + type: slurm + hostname: cw-dfw-cs-001-login-01 + walltime: "01:00:00" + account: coreai_dlalgo_nemorl + node_pools: + compute: + partition: batch + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:latest + type: vllm + model: Qwen/Qwen2.5-1.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 4 # needs the ray backend since the footprint exceeds one node + port: 8000 + health_check: + timeout_seconds: 1200 + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + +driver: + policy_model: vllm_model + container: python:3.13 + gym_install: + ref: onur/multi-instance-vllm-service + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/gpqa/config.yaml + run: + split: benchmark + config_paths: + - benchmarks/gpqa/config.yaml + +job: + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_multi_node.yaml b/examples/slurm_vllm_ray_multi_node.yaml index d0457558075..12a3db873e7 100644 --- a/examples/slurm_vllm_ray_multi_node.yaml +++ b/examples/slurm_vllm_ray_multi_node.yaml @@ -18,11 +18,7 @@ services: model: Qwen/Qwen2.5-72B-Instruct trust_remote_code: true tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 8 # needs the ray backend since the footprint exceeds one node - distributed_backend: - type: ray # vLLM's Ray core executor (--distributed-executor-backend ray). - # Not the ray.serve library - just lets vLLM's instances span nodes. - # A separate ray_serve backend may be added later. + number_of_instances: 8 # footprint exceeds one node, so the ray backend is used automatically port: 8000 health_check: timeout_seconds: 1200 diff --git a/examples/slurm_vllm_ray_multi_node_onur.yaml b/examples/slurm_vllm_ray_multi_node_onur.yaml new file mode 100644 index 00000000000..66fe7a69af1 --- /dev/null +++ b/examples/slurm_vllm_ray_multi_node_onur.yaml @@ -0,0 +1,46 @@ +compute: + cluster: + type: slurm + hostname: cw-dfw-cs-001-login-01 + walltime: "01:00:00" + account: coreai_dlalgo_nemorl + node_pools: + compute: + partition: batch + nodes: 2 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:latest + type: vllm + model: Qwen/Qwen2.5-1.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 8 # footprint exceeds one node, so the ray backend is used automatically + port: 8000 + health_check: + timeout_seconds: 1200 + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + +driver: + policy_model: vllm_model + container: python:3.13 + gym_install: + ref: onur/ray-service-multi-gpu + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/gpqa/config.yaml + run: + split: benchmark + config_paths: + - benchmarks/gpqa/config.yaml + +job: + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_serve_multi_node_onur.yaml b/examples/slurm_vllm_ray_serve_multi_node_onur.yaml new file mode 100644 index 00000000000..22f554ca65b --- /dev/null +++ b/examples/slurm_vllm_ray_serve_multi_node_onur.yaml @@ -0,0 +1,50 @@ +compute: + cluster: + type: slurm + hostname: cw-dfw-cs-001-login-01 + walltime: "01:00:00" + account: coreai_dlalgo_nemorl + node_pools: + compute: + partition: batch + nodes: 2 + ntasks_per_node: 1 + gpus_per_node: 8 + +services: + vllm_model: + container: vllm/vllm-openai:latest + type: vllm + model: Qwen/Qwen2.5-1.5B-Instruct + trust_remote_code: true + tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) + number_of_instances: 8 # needs the ray backend since the footprint exceeds one node + distributed_backend: + type: ray_serve # vLLM's Ray core executor (--distributed-executor-backend ray). + # Not the ray.serve library - just lets vLLM's instances span nodes. + # A separate ray_serve backend may be added later. + port: 8000 + health_check: + timeout_seconds: 1200 + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + +driver: + policy_model: vllm_model + container: python:3.13 + gym_install: + ref: onur/add-ray-serve-multi-gpu + env: + HF_TOKEN: ${oc.env:HF_TOKEN} + benchmarks: + gpqa: + prepare: + config_paths: + - benchmarks/gpqa/config.yaml + run: + split: benchmark + config_paths: + - benchmarks/gpqa/config.yaml + +job: + output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 14d03c8839d..be846bebbc5 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -206,16 +206,10 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": f"({', '.join(sorted(compute_names))})." ) - if ( - is_multi_node - and isinstance(service, VllmServiceConfig) - and (service.distributed_backend is None or service.distributed_backend.type != "ray") - ): - raise ValueError( - f"Service '{service_name}' is placed on compute '{sole_compute}', which spans multiple nodes " - f"({total_nodes} total). vLLM services on multi-node compute must use " - "distributed_backend: {type: ray}." - ) + if isinstance(service, VllmServiceConfig) and is_multi_node: + # Node count is authoritative: multi-node compute always needs the `ray` backend to + # span nodes, regardless of what (if anything) was set/defaulted at the service level. + service.distributed_backend = RayDistributedBackend() if ( is_multi_node diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index 6faf978cb71..a13273940eb 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -154,14 +154,16 @@ def test_ray_backend_with_single_instance_accepted(): assert config.services["svc"].distributed_backend.type == "ray" -def test_vllm_service_backend_on_multi_node_compute_raises(): - with pytest.raises(ValidationError, match="must use distributed_backend"): - SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE}, compute=COMPUTE_MULTI_NODE)) +def test_mp_backend_on_multi_node_compute_overridden_to_ray(): + # Node count is authoritative: an explicit (or auto-defaulted) `mp` backend is overridden to + # `ray` on multi-node compute instead of raising, so users don't have to spell it out. + config = SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE}, compute=COMPUTE_MULTI_NODE)) + assert config.services["svc"].distributed_backend.type == "ray" -def test_vllm_service_no_backend_on_multi_node_compute_raises(): - with pytest.raises(ValidationError, match="must use distributed_backend"): - SubmitConfig.model_validate(_config(compute=COMPUTE_MULTI_NODE)) +def test_no_backend_on_multi_node_compute_defaults_to_ray(): + config = SubmitConfig.model_validate(_config(compute=COMPUTE_MULTI_NODE)) + assert config.services["svc"].distributed_backend.type == "ray" def test_ray_backend_on_multi_node_compute_accepted(): diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index 05fc6d60017..ec1f1985aa0 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -796,6 +796,39 @@ def _multi_node_config(): ) +def test_build_sbatch_script_multi_node_no_backend_config_defaults_to_ray(bench_dir): + # Node count alone should be enough to select the ray backend - no distributed_backend needed. + config = SubmitConfig.model_validate( + { + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "tensor_parallel_size": 8, + } + }, + "compute": { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"main": {"partition": "gpu", "nodes": 4, "ntasks_per_node": 1}}, + } + }, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + assert config.services["vllm_model"].distributed_backend.type == "ray" + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) + assert "--nodes=4" in vllm_line + assert "ray symmetric-run" in script + + def test_build_sbatch_script_multi_node_vllm_srun_gets_node_flags(bench_dir): config = _multi_node_config() benchmark = config.driver.benchmarks["gsm8k"] From 098c42cf34660aa0fee9b8052a14f5f3f70eede4 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 19:24:44 -0400 Subject: [PATCH 64/66] Untrack onur scratch examples MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep these locally but drop them from the branch — personal scratch configs, not part of this PR. Signed-off-by: Onur Yilmaz --- examples/slurm_vllm_multi_instance_onur.yaml | 46 ----------------- examples/slurm_vllm_ray_multi_node_onur.yaml | 46 ----------------- .../slurm_vllm_ray_serve_multi_node_onur.yaml | 50 ------------------- 3 files changed, 142 deletions(-) delete mode 100644 examples/slurm_vllm_multi_instance_onur.yaml delete mode 100644 examples/slurm_vllm_ray_multi_node_onur.yaml delete mode 100644 examples/slurm_vllm_ray_serve_multi_node_onur.yaml diff --git a/examples/slurm_vllm_multi_instance_onur.yaml b/examples/slurm_vllm_multi_instance_onur.yaml deleted file mode 100644 index 0cb65fb995e..00000000000 --- a/examples/slurm_vllm_multi_instance_onur.yaml +++ /dev/null @@ -1,46 +0,0 @@ -compute: - cluster: - type: slurm - hostname: cw-dfw-cs-001-login-01 - walltime: "01:00:00" - account: coreai_dlalgo_nemorl - node_pools: - compute: - partition: batch - nodes: 1 - ntasks_per_node: 1 - gpus_per_node: 8 - -services: - vllm_model: - container: vllm/vllm-openai:latest - type: vllm - model: Qwen/Qwen2.5-1.5B-Instruct - trust_remote_code: true - tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 4 # needs the ray backend since the footprint exceeds one node - port: 8000 - health_check: - timeout_seconds: 1200 - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - -driver: - policy_model: vllm_model - container: python:3.13 - gym_install: - ref: onur/multi-instance-vllm-service - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - benchmarks: - gpqa: - prepare: - config_paths: - - benchmarks/gpqa/config.yaml - run: - split: benchmark - config_paths: - - benchmarks/gpqa/config.yaml - -job: - output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_multi_node_onur.yaml b/examples/slurm_vllm_ray_multi_node_onur.yaml deleted file mode 100644 index 66fe7a69af1..00000000000 --- a/examples/slurm_vllm_ray_multi_node_onur.yaml +++ /dev/null @@ -1,46 +0,0 @@ -compute: - cluster: - type: slurm - hostname: cw-dfw-cs-001-login-01 - walltime: "01:00:00" - account: coreai_dlalgo_nemorl - node_pools: - compute: - partition: batch - nodes: 2 - ntasks_per_node: 1 - gpus_per_node: 8 - -services: - vllm_model: - container: vllm/vllm-openai:latest - type: vllm - model: Qwen/Qwen2.5-1.5B-Instruct - trust_remote_code: true - tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 8 # footprint exceeds one node, so the ray backend is used automatically - port: 8000 - health_check: - timeout_seconds: 1200 - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - -driver: - policy_model: vllm_model - container: python:3.13 - gym_install: - ref: onur/ray-service-multi-gpu - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - benchmarks: - gpqa: - prepare: - config_paths: - - benchmarks/gpqa/config.yaml - run: - split: benchmark - config_paths: - - benchmarks/gpqa/config.yaml - -job: - output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output diff --git a/examples/slurm_vllm_ray_serve_multi_node_onur.yaml b/examples/slurm_vllm_ray_serve_multi_node_onur.yaml deleted file mode 100644 index 22f554ca65b..00000000000 --- a/examples/slurm_vllm_ray_serve_multi_node_onur.yaml +++ /dev/null @@ -1,50 +0,0 @@ -compute: - cluster: - type: slurm - hostname: cw-dfw-cs-001-login-01 - walltime: "01:00:00" - account: coreai_dlalgo_nemorl - node_pools: - compute: - partition: batch - nodes: 2 - ntasks_per_node: 1 - gpus_per_node: 8 - -services: - vllm_model: - container: vllm/vllm-openai:latest - type: vllm - model: Qwen/Qwen2.5-1.5B-Instruct - trust_remote_code: true - tensor_parallel_size: 2 # 8 instances x TP2 = 16 GPUs total, spans both nodes (2 x 8 GPUs) - number_of_instances: 8 # needs the ray backend since the footprint exceeds one node - distributed_backend: - type: ray_serve # vLLM's Ray core executor (--distributed-executor-backend ray). - # Not the ray.serve library - just lets vLLM's instances span nodes. - # A separate ray_serve backend may be added later. - port: 8000 - health_check: - timeout_seconds: 1200 - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - -driver: - policy_model: vllm_model - container: python:3.13 - gym_install: - ref: onur/add-ray-serve-multi-gpu - env: - HF_TOKEN: ${oc.env:HF_TOKEN} - benchmarks: - gpqa: - prepare: - config_paths: - - benchmarks/gpqa/config.yaml - run: - split: benchmark - config_paths: - - benchmarks/gpqa/config.yaml - -job: - output_path: /lustre/fsw/portfolios/coreai/users/oyilmaz/Gym_output From 0cdda1b3c2d4a2c6e23992fe4bebed1059bec185 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Wed, 19 Aug 2026 19:31:09 -0400 Subject: [PATCH 65/66] Add GPU number check Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 71 ++++++++++++++-------- tests/unit_tests/test_orchestration_api.py | 58 +++++++++++++++++- 2 files changed, 104 insertions(+), 25 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index be846bebbc5..81b289057cd 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -223,13 +223,8 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": "deployment - each node hosts an equal share of the data-parallel replicas." ) - is_ray_backend = ( - isinstance(service, VllmServiceConfig) - and service.distributed_backend is not None - and service.distributed_backend.type == "ray" - ) - if isinstance(service, VllmServiceConfig) and not (is_multi_node and is_ray_backend): - self._validate_vllm_gpu_footprint(service_name, service) + if isinstance(service, VllmServiceConfig): + self._validate_vllm_gpu_footprint(service_name, service, total_nodes) if self.driver.policy_model is not None: if self.driver.policy_model not in self.services: @@ -255,7 +250,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": return self - def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceConfig") -> None: + 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 @@ -267,25 +262,53 @@ def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceC return max_gpus_per_node = max(gpus_per_node_values) - gpus_needed = service.tensor_parallel_size * service.pipeline_parallel_size * service.number_of_instances + tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size + + if 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_node_dp_command); the + # per-node share, not the total footprint, has to fit in that node's GPU count. + instances_per_node = service.number_of_instances // total_nodes + gpus_needed = tp_pp * instances_per_node + gpus_available = max_gpus_per_node + footprint = ( + f"{instances_per_node} local replica(s) per node (number_of_instances=" + f"{service.number_of_instances} / {total_nodes} nodes) x tensor_parallel_size=" + f"{service.tensor_parallel_size} x pipeline_parallel_size={service.pipeline_parallel_size}" + ) + scope = f"a single node's gpus_per_node ({max_gpus_per_node})" + elif total_nodes > 1: + # Single instance's TP/PP footprint spans the whole allocation via the ray backend. + gpus_needed = tp_pp + 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}" + ) + scope = f"the total GPUs across all nodes ({gpus_available})" + else: + gpus_needed = tp_pp * service.number_of_instances + gpus_available = max_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}" + ) + scope = f"the node pool's gpus_per_node ({max_gpus_per_node})" - if gpus_needed > max_gpus_per_node: + if gpus_needed > gpus_available: raise ValueError( - f"Service '{service_name}' requires {gpus_needed} GPUs " - 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}), which exceeds the largest available " - f"node pool's gpus_per_node ({max_gpus_per_node}) on compute '{service.placement}'. " - "Multi-node vLLM services are not supported yet by the 'mp' distributed backend; " - "reduce number_of_instances/tensor_parallel_size/pipeline_parallel_size to fit on a single node." + f"Service '{service_name}' requires {gpus_needed} GPUs ({footprint}), which exceeds {scope} " + f"on compute '{service.placement}'. Reduce number_of_instances/tensor_parallel_size/" + "pipeline_parallel_size, or add more nodes/GPUs." ) - elif gpus_needed < max_gpus_per_node: + elif gpus_needed < gpus_available: warnings.warn( - f"Service '{service_name}' requires {gpus_needed} GPUs " - 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}) but compute '{service.placement}' allocates " - f"nodes with {max_gpus_per_node} GPUs each, leaving {max_gpus_per_node - gpus_needed} GPU(s) idle. " - "Increase number_of_instances/tensor_parallel_size or reduce gpus_per_node to use the full node.", + f"Service '{service_name}' requires {gpus_needed} GPUs ({footprint}) but compute " + f"'{service.placement}' provides {scope}, leaving {gpus_available - gpus_needed} GPU(s) idle. " + "Increase number_of_instances/tensor_parallel_size or reduce gpus_per_node to use the full " + "allocation.", stacklevel=2, ) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index a13273940eb..c50322946aa 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -201,6 +201,62 @@ def test_ray_backend_multi_node_dp_uneven_split_raises(): ) +# --------------------------------------------------------------------------- +# GPU footprint vs node pool capacity - multi-node +# --------------------------------------------------------------------------- + + +def test_multi_node_single_instance_footprint_exact_fit_accepted(): + # COMPUTE_MULTI_NODE has 2 nodes x 4 GPUs = 8 total; TP8 spans the whole allocation via ray. + config = SubmitConfig.model_validate( + _config(services={"svc": {**SERVICE, "tensor_parallel_size": 8}}, compute=COMPUTE_MULTI_NODE) + ) + assert config.services["svc"].tensor_parallel_size == 8 + + +def test_multi_node_single_instance_footprint_exceeds_total_raises(): + with pytest.raises(ValidationError, match="exceeds the total GPUs across all nodes"): + SubmitConfig.model_validate( + _config(services={"svc": {**SERVICE, "tensor_parallel_size": 9}}, compute=COMPUTE_MULTI_NODE) + ) + + +def test_multi_node_dp_per_node_footprint_exact_fit_accepted(): + # 4 instances / 2 nodes = 2 local replicas/node; TP2 x 2 local replicas = 4 GPUs, matches gpus_per_node. + config = SubmitConfig.model_validate( + _config( + services={ + "svc": { + **SERVICE, + "tensor_parallel_size": 2, + "number_of_instances": 4, + "distributed_backend": {"type": "ray"}, + } + }, + compute=COMPUTE_MULTI_NODE, + ) + ) + assert config.services["svc"].number_of_instances == 4 + + +def test_multi_node_dp_per_node_footprint_exceeds_raises(): + # 4 instances / 2 nodes = 2 local replicas/node; TP3 x 2 local replicas = 6 GPUs > gpus_per_node (4). + with pytest.raises(ValidationError, match="exceeds a single node's gpus_per_node"): + SubmitConfig.model_validate( + _config( + services={ + "svc": { + **SERVICE, + "tensor_parallel_size": 3, + "number_of_instances": 4, + "distributed_backend": {"type": "ray"}, + } + }, + compute=COMPUTE_MULTI_NODE, + ) + ) + + # --------------------------------------------------------------------------- # GPU footprint vs node pool capacity # --------------------------------------------------------------------------- @@ -233,7 +289,7 @@ def test_gpu_footprint_exceeds_node_raises(): "number_of_instances": 8, "distributed_backend": {"type": "mp"}, } - with pytest.raises(ValidationError, match="exceeds the largest available"): + with pytest.raises(ValidationError, match="exceeds the node pool's gpus_per_node"): SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) From d1ccb007b00bb76e0ef158f377af92cf6f979688 Mon Sep 17 00:00:00 2001 From: Onur Yilmaz Date: Mon, 24 Aug 2026 18:23:43 -0400 Subject: [PATCH 66/66] Address PR reviews Signed-off-by: Onur Yilmaz --- nemo_gym/orchestration/api.py | 56 +++----- .../executors/script_templates.py | 45 ++++++ .../orchestration/executors/slurm_script.py | 84 ++++-------- tests/unit_tests/test_orchestration_api.py | 129 +++++------------- tests/unit_tests/test_slurm_script.py | 80 ++++++----- 5 files changed, 166 insertions(+), 228 deletions(-) diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py index 81b289057cd..4bfddb8a0d9 100644 --- a/nemo_gym/orchestration/api.py +++ b/nemo_gym/orchestration/api.py @@ -49,36 +49,12 @@ class BaseModelServiceConfig(BaseServiceConfig): port: int = 8000 -class VllmServiceDistributedBackend(_StrictModel): - """Use vLLM's native data-parallel multi-instance (--data-parallel-size N). Single node only.""" - - type: Literal["mp"] = "mp" - - -class RayDistributedBackend(_StrictModel): - """Use vLLM's Ray executor (--distributed-executor-backend ray) so a single vLLM service can - span multiple physical Slurm nodes. This does not use the ray.serve library. - - Named "ray" for now; a separate "ray_serve" backend (using the actual ray.serve library) may - be added later. - """ - - type: Literal["ray"] = "ray" - - -DistributedBackendConfig = Annotated[ - Annotated[VllmServiceDistributedBackend, Tag("mp")] | Annotated[RayDistributedBackend, Tag("ray")], - Discriminator("type"), -] - - class VllmServiceConfig(BaseModelServiceConfig): type: Literal["vllm"] tensor_parallel_size: int = 1 pipeline_parallel_size: int = 1 trust_remote_code: bool = False number_of_instances: int = 1 - distributed_backend: DistributedBackendConfig | None = None @field_validator("number_of_instances") @classmethod @@ -87,14 +63,6 @@ def _validate_number_of_instances(cls, v: int) -> int: raise ValueError(f"number_of_instances must be >= 1, got {v}") return v - @model_validator(mode="after") - def _validate_distributed_backend(self) -> "VllmServiceConfig": - # distributed_backend may still be set explicitly with number_of_instances == 1 (e.g. `ray`, - # to span a single instance's tensor/pipeline-parallel footprint across multiple nodes). - if self.number_of_instances > 1 and self.distributed_backend is None: - self.distributed_backend = VllmServiceDistributedBackend() - return self - @model_validator(mode="after") def _default_health_check(self) -> "VllmServiceConfig": # vLLM always exposes /health on its serving port; set it automatically @@ -206,11 +174,6 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig": f"({', '.join(sorted(compute_names))})." ) - if isinstance(service, VllmServiceConfig) and is_multi_node: - # Node count is authoritative: multi-node compute always needs the `ray` backend to - # span nodes, regardless of what (if anything) was set/defaulted at the service level. - service.distributed_backend = RayDistributedBackend() - if ( is_multi_node and isinstance(service, VllmServiceConfig) @@ -265,9 +228,24 @@ def _validate_vllm_gpu_footprint(self, service_name: str, service: "VllmServiceC tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size if total_nodes > 1 and service.number_of_instances > 1: + if tp_pp > max_gpus_per_node: + # Each instance's own TP/PP footprint already exceeds a single node's GPU count, so + # spreading multiple such instances across nodes would require every instance to + # itself span multiple nodes. That's not supported: multi-node data-parallel only + # distributes whole instances across nodes with tensor/pipeline parallelism kept + # local to each node (see _build_vllm_multi_instance_multi_node_command). + raise ValueError( + f"Service '{service_name}' sets number_of_instances={service.number_of_instances} with " + f"tensor_parallel_size={service.tensor_parallel_size} x " + f"pipeline_parallel_size={service.pipeline_parallel_size}={tp_pp}, which exceeds a single " + f"node's gpus_per_node ({max_gpus_per_node}). Multiple instances where each instance's own " + "tensor/pipeline-parallel footprint spans multiple nodes is not supported - reduce " + "tensor_parallel_size/pipeline_parallel_size to fit within one node, or set " + "number_of_instances=1 to let a single instance span nodes." + ) # Multi-node data-parallel: each node runs its own equal share of the replicas with - # local tensor/pipeline parallelism (see _build_vllm_multi_node_dp_command); the - # per-node share, not the total footprint, has to fit in that node's GPU count. + # 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. 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/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py index bd1d2694e73..02d5e77ce82 100644 --- a/nemo_gym/orchestration/executors/script_templates.py +++ b/nemo_gym/orchestration/executors/script_templates.py @@ -17,6 +17,35 @@ import shlex +_RAY_PRELUDE = """\ +# Resolve the head node IP for multi-node vLLM services (spanning nodes via Ray). +nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +nodes_array=($nodes) +head_node_hostname=${nodes_array[0]} +head_node_ip=$(getent hosts "$head_node_hostname" | awk '{print $1}') +export HEAD_NODE_IP="$head_node_ip" +export RAY_HEAD_NODE_IP="$head_node_ip:6379" +echo "Head node IP address: $HEAD_NODE_IP\"""" + + +_VLLM_RAY_SYMMETRIC_RUN = """\ +bash -lc ' + command -v ray >/dev/null 2>&1 || pip install -q "ray[default]" + if ray symmetric-run --help >/dev/null 2>&1; then + ray symmetric-run \\ + --address "$RAY_HEAD_NODE_IP" \\ + --min-nodes {total_nodes} \\ + {resource_flags} \\ + -- {inner_cmd} + elif [ "$SLURM_NODEID" = "0" ]; then + ray start --head --port=6379 {resource_flags} + {inner_cmd} + else + ray start --address="$RAY_HEAD_NODE_IP" {resource_flags} --block + fi +'""" + + _HEALTH_WAIT_MULTI = """\ # Wait for {name} (try multiple health endpoints) echo "Waiting for {name} at {url}..." @@ -44,6 +73,22 @@ def bash_var(name: str) -> str: return re.sub(r"[^A-Z0-9]", "_", name.upper()) +def render_ray_prelude() -> str: + return _RAY_PRELUDE + + +def render_vllm_ray_symmetric_run(inner_cmd: str, total_nodes: int, resource_flags: str) -> str: + """Render the Ray head/worker bootstrap that wraps a single vLLM instance's TP/PP command so + it spans multiple Slurm nodes. + + Uses `ray symmetric-run` when available (Ray >= 2.50), which starts/joins a Ray cluster across + every task and runs the entrypoint only on the elected head node. Containers with an older Ray + pin fall back to manually starting head/worker Ray processes, keyed on Slurm's per-node task + rank ($SLURM_NODEID). + """ + return _VLLM_RAY_SYMMETRIC_RUN.format(total_nodes=total_nodes, resource_flags=resource_flags, inner_cmd=inner_cmd) + + def render_health_check(name: str, port: int, path: str, timeout: int) -> str: return _HEALTH_WAIT_MULTI.format( name=name, diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py index 79d5330af7b..030e8c49531 100644 --- a/nemo_gym/orchestration/executors/slurm_script.py +++ b/nemo_gym/orchestration/executors/slurm_script.py @@ -30,6 +30,8 @@ render_driver_entrypoint, render_gym_cmd, render_health_check, + render_ray_prelude, + render_vllm_ray_symmetric_run, ) from nemo_gym.orchestration.executors.utils import flatten_run_args @@ -50,22 +52,6 @@ """ -# Resolves the head node from Slurm's node list, mirroring scripts/sbatch_base.sh. -# Runs once in the main sbatch shell (before any srun step) so $SLURM_JOB_NODELIST is available. -# Used by both multi-node code paths below: HEAD_NODE_IP for vLLM's own --data-parallel-address -# (data-parallel spanning nodes), RAY_HEAD_NODE_IP for `ray start`/`ray symmetric-run` (tensor/ -# pipeline-parallel spanning nodes). -_RAY_PRELUDE = """\ -# Resolve the head node IP for multi-node vLLM services (distributed_backend: ray). -nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") -nodes_array=($nodes) -head_node_hostname=${nodes_array[0]} -head_node_ip=$(getent hosts "$head_node_hostname" | awk '{print $1}') -export HEAD_NODE_IP="$head_node_ip" -export RAY_HEAD_NODE_IP="$head_node_ip:6379" -echo "Head node IP address: $HEAD_NODE_IP\"""" - - def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benchmark_name: str) -> str: lines = [] lines.append(f"#SBATCH --job-name=gym-{benchmark_name}") @@ -155,7 +141,7 @@ def _build_vllm_command(service: VllmServiceConfig) -> str: return cmd -def _build_vllm_multi_node_tp_command(service: VllmServiceConfig, total_nodes: int) -> str: +def _build_vllm_single_instance_multi_node_command(service: VllmServiceConfig, total_nodes: int) -> str: # A single instance's tensor/pipeline-parallel footprint spans nodes. Uses vLLM's own Ray # *core* executor (--distributed-executor-backend ray) - not the ray.serve library, no Serve # deployment/ingress/HTTP proxy is involved. @@ -163,34 +149,15 @@ def _build_vllm_multi_node_tp_command(service: VllmServiceConfig, total_nodes: i resource_flags = ( "--num-cpus=${SLURM_CPUS_PER_TASK:-$SLURM_CPUS_ON_NODE} --num-gpus=${SLURM_GPUS_PER_TASK:-$SLURM_GPUS_ON_NODE}" ) - # ray symmetric-run starts/joins a Ray cluster across every task and runs the entrypoint - # only on the elected head node, mirroring scripts/sbatch_base.sh. It requires Ray >= 2.50, so - # containers with an older pin fall back to manually starting head/worker Ray processes, keyed - # on Slurm's per-node task rank ($SLURM_NODEID). vLLM's Ray executor blocks on placement-group - # scheduling until every node's GPUs join, so the fallback needs no separate cluster-ready wait. # Model-serving images (e.g. vllm/vllm-openai) don't necessarily bundle the ray CLI - vLLM only # needs ray as a runtime dependency when the ray executor backend is actually selected - so - # install it on the fly if it's missing before relying on either code path above. - return ( - "bash -lc '\n" - ' command -v ray >/dev/null 2>&1 || pip install -q "ray[default]"\n' - " if ray symmetric-run --help >/dev/null 2>&1; then\n" - " ray symmetric-run \\\n" - ' --address "$RAY_HEAD_NODE_IP" \\\n' - f" --min-nodes {total_nodes} \\\n" - f" {resource_flags} \\\n" - f" -- {inner_cmd}\n" - ' elif [ "$SLURM_NODEID" = "0" ]; then\n' - f" ray start --head --port=6379 {resource_flags}\n" - f" {inner_cmd}\n" - " else\n" - f' ray start --address="$RAY_HEAD_NODE_IP" {resource_flags} --block\n' - " fi\n" - "'" - ) + # render_vllm_ray_symmetric_run installs it on the fly if it's missing. vLLM's Ray executor + # blocks on placement-group scheduling until every node's GPUs join, so the fallback path there + # needs no separate cluster-ready wait. + return render_vllm_ray_symmetric_run(inner_cmd, total_nodes, resource_flags) -def _build_vllm_multi_node_dp_command(service: VllmServiceConfig, total_nodes: int) -> str: +def _build_vllm_multi_instance_multi_node_command(service: VllmServiceConfig, total_nodes: int) -> str: # Data-parallel replicas span nodes. vLLM's Ray-based DP auto-placement doesn't spread ranks # across physical nodes - launching a single `vllm serve --data-parallel-size N` from one node # only sees that node's own GPUs when placing DP ranks. Real multi-node DP instead needs one @@ -198,11 +165,8 @@ def _build_vllm_multi_node_dp_command(service: VllmServiceConfig, total_nodes: i # worker nodes run `--headless` with a --data-parallel-start-rank offset. This is vLLM's # documented multi-node data-parallel deployment pattern and doesn't use Ray at all - each # node's tensor-parallel ranks stay local via vLLM's default (mp) executor backend. - if service.number_of_instances % total_nodes != 0: - raise ValueError( - f"number_of_instances ({service.number_of_instances}) must be evenly divisible by the number of " - f"nodes ({total_nodes}) for multi-node data-parallel deployment." - ) + # number_of_instances is guaranteed evenly divisible by total_nodes here - api.py's + # SubmitConfig validation enforces this before build_sbatch_script is ever called. dp_size_local = service.number_of_instances // total_nodes common = _vllm_base_flags(service) dp_flags = ( @@ -233,8 +197,8 @@ def _build_vllm_multi_node_dp_command(service: VllmServiceConfig, total_nodes: i def _build_vllm_ray_command(service: VllmServiceConfig, total_nodes: int) -> str: if service.number_of_instances > 1: - return _build_vllm_multi_node_dp_command(service, total_nodes) - return _build_vllm_multi_node_tp_command(service, total_nodes) + return _build_vllm_multi_instance_multi_node_command(service, total_nodes) + return _build_vllm_single_instance_multi_node_command(service, total_nodes) def _build_ray_command(_service: RayServiceConfig) -> str: @@ -247,14 +211,15 @@ def _build_ray_command(_service: RayServiceConfig) -> str: } -def _uses_ray_distributed_backend(service: VllmServiceConfig | RayServiceConfig) -> bool: - return isinstance(service, VllmServiceConfig) and ( - service.distributed_backend is not None and service.distributed_backend.type == "ray" - ) +def _vllm_spans_multiple_nodes(service: VllmServiceConfig | RayServiceConfig, total_nodes: int) -> bool: + # Node count alone determines this: multi-node compute always spans a vLLM service across + # nodes via Ray, regardless of number_of_instances (single instance's TP/PP, or DP replicas). + # Non-vLLM services (e.g. a plain Ray head) never span nodes this way. + return isinstance(service, VllmServiceConfig) and total_nodes > 1 def _build_service_command(service: VllmServiceConfig | RayServiceConfig, total_nodes: int) -> str: - if _uses_ray_distributed_backend(service): + if _vllm_spans_multiple_nodes(service, total_nodes): return _build_vllm_ray_command(service, total_nodes) return _BUILDERS[type(service)](service) @@ -277,7 +242,11 @@ def build_sbatch_script( total_nodes, total_ntasks = _node_totals(compute) is_multi_node = total_nodes > 1 - ray_prelude = _RAY_PRELUDE if any(_uses_ray_distributed_backend(s) for s in config.services.values()) else "" + ray_prelude = ( + render_ray_prelude() + if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values()) + else "" + ) service_commands = "\n\n".join( _render_service_command( @@ -286,8 +255,11 @@ def build_sbatch_script( _build_service_command(service, total_nodes), service.env or None, service.mounts or None, - nodes=total_nodes if is_multi_node else None, - ntasks=total_ntasks if is_multi_node else None, + # Only services that actually span multiple nodes need the whole allocation's --nodes/ + # --ntasks - not every service in a multi-node job (e.g. a plain Ray head service runs + # on a single node regardless of how many nodes the overall job spans). + nodes=total_nodes if _vllm_spans_multiple_nodes(service, total_nodes) else None, + ntasks=total_ntasks if _vllm_spans_multiple_nodes(service, total_nodes) else None, ) for name, service in config.services.items() ) diff --git a/tests/unit_tests/test_orchestration_api.py b/tests/unit_tests/test_orchestration_api.py index c50322946aa..8cfa854a1e9 100644 --- a/tests/unit_tests/test_orchestration_api.py +++ b/tests/unit_tests/test_orchestration_api.py @@ -99,31 +99,20 @@ def test_service_unknown_field_raises(): # --------------------------------------------------------------------------- -# number_of_instances / distributed_backend +# number_of_instances # --------------------------------------------------------------------------- -_MULTI_SERVICE = {**SERVICE, "number_of_instances": 4, "distributed_backend": {"type": "mp"}} +_MULTI_SERVICE = {**SERVICE, "number_of_instances": 4} -def test_number_of_instances_with_backend_accepted(): +def test_number_of_instances_accepted(): config = SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE})) - svc = config.services["svc"] - assert svc.number_of_instances == 4 - assert svc.distributed_backend is not None - assert svc.distributed_backend.type == "mp" + assert config.services["svc"].number_of_instances == 4 def test_number_of_instances_defaults_to_1(): config = SubmitConfig.model_validate(_config()) assert config.services["svc"].number_of_instances == 1 - assert config.services["svc"].distributed_backend is None - - -def test_multi_instance_without_backend_defaults_to_mp(): - config = SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 4}})) - svc = config.services["svc"] - assert svc.distributed_backend is not None - assert svc.distributed_backend.type == "mp" def test_number_of_instances_zero_raises(): @@ -131,8 +120,15 @@ def test_number_of_instances_zero_raises(): SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "number_of_instances": 0}})) +def test_distributed_backend_is_not_a_settable_field(): + # distributed_backend (ray vs single-node mp) is an internal implementation detail, fully + # derived from node count / number_of_instances - not part of the public schema. + with pytest.raises(ValidationError): + SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "distributed_backend": {"type": "ray"}}})) + + # --------------------------------------------------------------------------- -# ray distributed_backend (multi-node) +# multi-node (ray-backed) deployment # --------------------------------------------------------------------------- COMPUTE_MULTI_NODE = { @@ -145,59 +141,24 @@ def test_number_of_instances_zero_raises(): } -def test_ray_backend_with_single_instance_accepted(): - # ray is valid even with number_of_instances == 1: it's needed to span nodes for - # a single instance's tensor/pipeline-parallel footprint, not just for data-parallel replicas. - config = SubmitConfig.model_validate( - _config(services={"svc": {**SERVICE, "distributed_backend": {"type": "ray"}}}) - ) - assert config.services["svc"].distributed_backend.type == "ray" - - -def test_mp_backend_on_multi_node_compute_overridden_to_ray(): - # Node count is authoritative: an explicit (or auto-defaulted) `mp` backend is overridden to - # `ray` on multi-node compute instead of raising, so users don't have to spell it out. - config = SubmitConfig.model_validate(_config(services={"svc": _MULTI_SERVICE}, compute=COMPUTE_MULTI_NODE)) - assert config.services["svc"].distributed_backend.type == "ray" - - -def test_no_backend_on_multi_node_compute_defaults_to_ray(): +def test_multi_node_single_instance_accepted(): + # Node count alone is enough to span a single instance's TP/PP footprint across nodes. config = SubmitConfig.model_validate(_config(compute=COMPUTE_MULTI_NODE)) - assert config.services["svc"].distributed_backend.type == "ray" - - -def test_ray_backend_on_multi_node_compute_accepted(): - config = SubmitConfig.model_validate( - _config( - services={"svc": {**SERVICE, "distributed_backend": {"type": "ray"}}}, - compute=COMPUTE_MULTI_NODE, - ) - ) - assert config.services["svc"].distributed_backend.type == "ray" + assert config.services["svc"].number_of_instances == 1 -def test_ray_backend_multi_node_dp_evenly_divisible_accepted(): +def test_multi_node_dp_evenly_divisible_accepted(): # COMPUTE_MULTI_NODE has 2 nodes; 4 instances split evenly (2 per node). config = SubmitConfig.model_validate( - _config( - services={ - "svc": {**SERVICE, "number_of_instances": 4, "distributed_backend": {"type": "ray"}}, - }, - compute=COMPUTE_MULTI_NODE, - ) + _config(services={"svc": {**SERVICE, "number_of_instances": 4}}, compute=COMPUTE_MULTI_NODE) ) assert config.services["svc"].number_of_instances == 4 -def test_ray_backend_multi_node_dp_uneven_split_raises(): +def test_multi_node_dp_uneven_split_raises(): with pytest.raises(ValidationError, match="evenly divisible"): SubmitConfig.model_validate( - _config( - services={ - "svc": {**SERVICE, "number_of_instances": 3, "distributed_backend": {"type": "ray"}}, - }, - compute=COMPUTE_MULTI_NODE, - ) + _config(services={"svc": {**SERVICE, "number_of_instances": 3}}, compute=COMPUTE_MULTI_NODE) ) @@ -225,14 +186,7 @@ def test_multi_node_dp_per_node_footprint_exact_fit_accepted(): # 4 instances / 2 nodes = 2 local replicas/node; TP2 x 2 local replicas = 4 GPUs, matches gpus_per_node. config = SubmitConfig.model_validate( _config( - services={ - "svc": { - **SERVICE, - "tensor_parallel_size": 2, - "number_of_instances": 4, - "distributed_backend": {"type": "ray"}, - } - }, + services={"svc": {**SERVICE, "tensor_parallel_size": 2, "number_of_instances": 4}}, compute=COMPUTE_MULTI_NODE, ) ) @@ -244,14 +198,20 @@ def test_multi_node_dp_per_node_footprint_exceeds_raises(): with pytest.raises(ValidationError, match="exceeds a single node's gpus_per_node"): SubmitConfig.model_validate( _config( - services={ - "svc": { - **SERVICE, - "tensor_parallel_size": 3, - "number_of_instances": 4, - "distributed_backend": {"type": "ray"}, - } - }, + services={"svc": {**SERVICE, "tensor_parallel_size": 3, "number_of_instances": 4}}, + compute=COMPUTE_MULTI_NODE, + ) + ) + + +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"): + SubmitConfig.model_validate( + _config( + services={"svc": {**SERVICE, "tensor_parallel_size": 5, "number_of_instances": 2}}, compute=COMPUTE_MULTI_NODE, ) ) @@ -272,34 +232,19 @@ def test_multi_node_dp_per_node_footprint_exceeds_raises(): def test_gpu_footprint_exact_fit_accepted(): - service = { - **SERVICE, - "tensor_parallel_size": 2, - "number_of_instances": 4, - "distributed_backend": {"type": "mp"}, - } + service = {**SERVICE, "tensor_parallel_size": 2, "number_of_instances": 4} config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) assert config.services["svc"].number_of_instances == 4 def test_gpu_footprint_exceeds_node_raises(): - service = { - **SERVICE, - "tensor_parallel_size": 2, - "number_of_instances": 8, - "distributed_backend": {"type": "mp"}, - } + service = {**SERVICE, "tensor_parallel_size": 2, "number_of_instances": 8} with pytest.raises(ValidationError, match="exceeds the node pool's gpus_per_node"): SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) def test_gpu_footprint_underutilized_warns(): - service = { - **SERVICE, - "tensor_parallel_size": 2, - "number_of_instances": 2, - "distributed_backend": {"type": "mp"}, - } + service = {**SERVICE, "tensor_parallel_size": 2, "number_of_instances": 2} with pytest.warns(UserWarning, match="leaving 4 GPU"): config = SubmitConfig.model_validate(_config(services={"svc": service}, compute=COMPUTE_8_GPUS_PER_NODE)) assert config.services["svc"].number_of_instances == 2 diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py index ec1f1985aa0..fcc30878a4c 100644 --- a/tests/unit_tests/test_slurm_script.py +++ b/tests/unit_tests/test_slurm_script.py @@ -177,7 +177,6 @@ def test_build_vllm_command_multi_instance(): container="vllm:latest", model="org/model", number_of_instances=4, - distributed_backend={"type": "mp"}, ) cmd = _build_vllm_command(service) assert "--data-parallel-size 4" in cmd @@ -243,7 +242,6 @@ def test_build_vllm_ray_command_dp_does_not_use_ray(): container="vllm:latest", model="org/model", number_of_instances=4, - distributed_backend={"type": "ray"}, ) cmd = _build_vllm_ray_command(service, total_nodes=2) assert "ray" not in cmd @@ -256,7 +254,6 @@ def test_build_vllm_ray_command_dp_head_and_worker_branches(): container="vllm:latest", model="org/model", number_of_instances=4, - distributed_backend={"type": "ray"}, ) cmd = _build_vllm_ray_command(service, total_nodes=2) assert 'if [ "$SLURM_NODEID" = "0" ]; then' in cmd @@ -268,18 +265,6 @@ def test_build_vllm_ray_command_dp_head_and_worker_branches(): assert "--data-parallel-start-rank $(( SLURM_NODEID * 2 ))" in cmd -def test_build_vllm_ray_command_dp_uneven_split_raises(): - service = VllmServiceConfig( - type="vllm", - container="vllm:latest", - model="org/model", - number_of_instances=3, - distributed_backend={"type": "ray"}, - ) - with pytest.raises(ValueError, match="evenly divisible"): - _build_vllm_ray_command(service, total_nodes=2) - - # --------------------------------------------------------------------------- # render_gym_cmd # --------------------------------------------------------------------------- @@ -779,7 +764,6 @@ def _multi_node_config(): "container": "vllm:latest", "model": "org/model", "tensor_parallel_size": 8, - "distributed_backend": {"type": "ray"}, } }, "compute": { @@ -796,31 +780,9 @@ def _multi_node_config(): ) -def test_build_sbatch_script_multi_node_no_backend_config_defaults_to_ray(bench_dir): - # Node count alone should be enough to select the ray backend - no distributed_backend needed. - config = SubmitConfig.model_validate( - { - "services": { - "vllm_model": { - "type": "vllm", - "container": "vllm:latest", - "model": "org/model", - "tensor_parallel_size": 8, - } - }, - "compute": { - "cluster": { - "type": "slurm", - "account": "my-account", - "hostname": "foo", - "node_pools": {"main": {"partition": "gpu", "nodes": 4, "ntasks_per_node": 1}}, - } - }, - "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, - "job": {"output_path": "/remote/jobs"}, - } - ) - assert config.services["vllm_model"].distributed_backend.type == "ray" +def test_build_sbatch_script_multi_node_selects_ray_by_node_count(bench_dir): + # Node count alone is enough to span the vLLM service via ray - no explicit config needed. + config = _multi_node_config() benchmark = config.driver.benchmarks["gsm8k"] compute = next(iter(config.compute.values())) script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) @@ -883,6 +845,42 @@ def test_build_sbatch_script_single_node_pool_omits_node_flags_from_srun(bench_d assert "--nodes=" not in driver_line +def test_build_sbatch_script_non_vllm_service_omits_node_flags_in_multi_node_job(bench_dir): + # A plain Ray head service doesn't span nodes itself, even when the vLLM service alongside it + # does - it must not get the whole allocation's --nodes/--ntasks (that would launch + # `ray start --head` once per node instead of once). + config = SubmitConfig.model_validate( + { + "services": { + "vllm_model": { + "type": "vllm", + "container": "vllm:latest", + "model": "org/model", + "tensor_parallel_size": 8, + }, + "ray_head": {"type": "ray", "container": "ray:latest"}, + }, + "compute": { + "cluster": { + "type": "slurm", + "account": "my-account", + "hostname": "foo", + "node_pools": {"main": {"partition": "gpu", "nodes": 4, "ntasks_per_node": 1}}, + } + }, + "driver": {"container": "python:3.12", "benchmarks": {"gsm8k": {}}}, + "job": {"output_path": "/remote/jobs"}, + } + ) + benchmark = config.driver.benchmarks["gsm8k"] + compute = next(iter(config.compute.values())) + script = build_sbatch_script(config, "gsm8k", benchmark, compute, bench_dir) + vllm_line = next(line for line in script.splitlines() if "vllm:latest" in line) + ray_line = next(line for line in script.splitlines() if "ray:latest" in line) + assert "--nodes=4" in vllm_line + assert "--nodes=" not in ray_line + + # --------------------------------------------------------------------------- # build_sbatch_script — ray prelude # ---------------------------------------------------------------------------