Basic SLURM orchestration

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": {

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
+

nemo_gym/cli/main.py

@@ -319,6 +319,20 @@ 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)
@@ -661,6 +675,22 @@ 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=_eval_submit,
+        summary="Submit a job.",
+        flags=(
+            Flag(
+                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."
+                ),
+            ),
+        ),
+    ),
     "dev test": Command(target="nemo_gym.cli.dev:dev_test", summary="Run NeMo Gym's unit tests."),
 }
 

nemo_gym/decorators.py

@@ -0,0 +1,31 @@
+# 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
+
+
+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

nemo_gym/orchestration/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nemo_gym/orchestration/api.py

@@ -0,0 +1,177 @@
+# 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
+
+
+# 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
+
+
+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
+
+    @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:
+            self.health_check.port = self.port
+        return self
+
+
+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"),
+]
+
+
+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] = {}
+
+
+class BaseComputeConfig(_StrictModel):
+    pass
+
+
+class SlurmComputeConfig(BaseComputeConfig):
+    type: Literal["slurm"]
+    account: str
+    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] = {}  # Job-level #SBATCH directives (e.g. --comment, --mail-user).
+
+
+ComputeConfig = Annotated[
+    Annotated[SlurmComputeConfig, Tag("slurm")],
+    Discriminator("type"),
+]
+
+
+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  # 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
+
+
+class SubmitConfig(_StrictModel):
+    services: dict[str, ServiceConfig]
+    compute: dict[str, ComputeConfig]
+    driver: DriverConfig
+    job: JobConfig
+
+    @model_validator(mode="after")
+    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))}).")
+
+        sole_compute = next(iter(compute_names))
+
+        for service_name, service in self.services.items():
+            if service.placement is None:
+                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))})."
+                )
+
+        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 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
+                    ]
+                    if conflicts:
+                        raise ValueError(
+                            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"
+                    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
+
+        return self

nemo_gym/orchestration/executors/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nemo_gym/orchestration/executors/base.py

@@ -0,0 +1,23 @@
+# 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
+
+
+class BaseExecutor(ABC):
+    @abstractmethod
+    def run(self, config: SubmitConfig, *, dry_run: bool = False) -> None: ...