Multi-Instance vLLM Support in gym eval submit (#2482)

## Multi-Instance vLLM Support in `gym eval submit`

### Summary

Adds support for running multiple vLLM engine replicas in `gym eval
submit` using vLLM's native data-parallel multi-instance mode.

Previously, `VllmServiceConfig` only supported a single vLLM engine.
This change introduces two new fields:

- **`number_of_instances`** — the number of engine replicas to run
(default `1`, preserving existing behavior)
- **`distributed_backend`** — a typed, extensible sub-config declaring
how replicas are coordinated. The only type available in this PR is
`vllm_service`, which maps to vLLM's `--data-parallel-size N` flag. vLLM
manages load balancing internally, so no external router or extra
processes are needed.

When `number_of_instances > 1`, the generated sbatch script passes
`--data-parallel-size N` to the single `vllm serve` command. The single
endpoint, health check, and driver URL wiring are all unchanged.

The `distributed_backend` field is designed to be extended in future PRs
— Ray Serve and Dynamo routing backends will each add a new type to the
discriminated union without touching existing configs.

### Validation rules

- `number_of_instances > 1` requires `distributed_backend` to be set
(fails loudly with a clear message)
- `number_of_instances == 1` with `distributed_backend` set is rejected
- `number_of_instances < 1` is rejected

### Also included

- `--pipeline-parallel-size` is now emitted in the generated `vllm
serve` command when `> 1` (it was in the config but was never passed to
the CLI)
- An example YAML (`examples/slurm_vllm_multi_instance.yaml`) showing a
4-instance setup with TP=2 across 8 GPUs

---------

Signed-off-by: Tomasz Grzegorzek <tgrzegorzek@nvidia.com>
Signed-off-by: prokotg <19536019+prokotg@users.noreply.github.com>
Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
Signed-off-by: Onur Yilmaz <35306097+oyilmaz-nvidia@users.noreply.github.com>
Co-authored-by: Tomasz Grzegorzek <tgrzegorzek@nvidia.com>
Co-authored-by: prokotg <19536019+prokotg@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: 3 people

examples/slurm_vllm_multi_instance.yaml

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

nemo_gym/orchestration/api.py

@@ -13,9 +13,10 @@
 # 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, 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 +49,41 @@ class BaseModelServiceConfig(BaseServiceConfig):
     port: int = 8000
 
 
+class VllmServiceDistributedBackend(_StrictModel):
+    """Use vLLM's native data-parallel multi-instance (--data-parallel-size N)."""
+
+    type: Literal["mp"] = "mp"
+
+
+# Future backends: add Annotated[RayServeDistributedBackend, Tag("ray_serve")], etc.
+DistributedBackendConfig = Annotated[
+    Annotated[VllmServiceDistributedBackend, Tag("mp")],
+    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 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")
     def _default_health_check(self) -> "VllmServiceConfig":
@@ -160,6 +191,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(
@@ -183,3 +217,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 '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:
+            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,
+            )

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

tests/unit_tests/test_orchestration_api.py

@@ -96,3 +96,95 @@ 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": "mp"}}
+
+
+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 == "mp"
+
+
+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_single_instance_with_backend_raises():
+    with pytest.raises(ValidationError, match="should not be set"):
+        SubmitConfig.model_validate(_config(services={"svc": {**SERVICE, "distributed_backend": {"type": "mp"}}}))
+
+
+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": "mp"},
+    }
+    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"},
+    }
+    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": "mp"},
+    }
+    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

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": "mp"},
+    )
+    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
 # ---------------------------------------------------------------------------