From 3fc7a12e9282fe2fe9c095c0480cd01a15c54898 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Thu, 1 Oct 2026 13:38:01 +0200
Subject: [PATCH 1/6] feat(orchestration): give services on a shared node their
 own GPUs

When two or more vLLM services share a node, each now gets its own
CUDA_VISIBLE_DEVICES (TP x PP x instances per node, in declaration order).
A ray service there gets the leftover GPUs. A node with one GPU service
renders nothing new.

A manual CUDA_VISIBLE_DEVICES index list wins and reserves its GPUs.
Overlaps and over-subscription fail validation. `share_gpus: true` opts
a service out.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py                 | 112 +++++++++++
 .../orchestration/executors/slurm_script.py   |  20 +-
 tests/unit_tests/test_slurm_script.py         | 190 ++++++++++++++++++
 3 files changed, 318 insertions(+), 4 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index e7f1fe2eaf5..ac0e1a59c97 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -86,6 +86,9 @@ class BaseServiceConfig(_StrictModel):
     # that cannot share a GPU with the policy, or a prefill/decode split. Pools take
     # contiguous node ranges in declaration order. None means the whole allocation.
     node_pool: str | None = None
+    # Services sharing a node get disjoint GPUs by default (see plan_gpus). True lets this
+    # one see every GPU instead, e.g. two small models each capped by --gpu-memory-utilization.
+    share_gpus: bool = False
     health_check: HealthCheckConfig | None = None
     # Values may be prefixed `lit:` (literal), `host:VAR` (read from the submitting
     # machine's env), or `runtime:VAR` (resolved from the job's own env at run time).
@@ -603,6 +606,7 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
                     # vLLM doesn't require auth; dummy key satisfies clients that require the header.
                     benchmark.run["policy_api_key"] = "dummy"  # pragma: allowlist secret
 
+        plan_gpus(self)
         return self
 
     def _validate_pd_services(self, compute: "ComputeConfig") -> None:
@@ -698,3 +702,111 @@ def _validate_vllm_gpu_footprint(
                 "allocation.",
                 stacklevel=2,
             )
+
+
+_CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES"
+_GPU_INDEX_LIST_RE = re.compile(r"^\d+(,\d+)*$")
+
+
+def _manual_gpus(service: BaseServiceConfig) -> list[int] | None:
+    value = service.env.get(_CUDA_VISIBLE_DEVICES)
+    if value is None or not _GPU_INDEX_LIST_RE.match(value):
+        return None
+    return [int(index) for index in value.split(",")]
+
+
+def _takes_own_gpus(service: BaseServiceConfig) -> bool:
+    # A CUDA_VISIBLE_DEVICES that is not an index list (e.g. runtime:VAR) names unknown GPUs.
+    if service.share_gpus:
+        return False
+    return _CUDA_VISIBLE_DEVICES not in service.env or _manual_gpus(service) is not None
+
+
+def _gpus_needed_per_node(service: VllmServiceConfig, nodes: int, gpus_per_node: int) -> int | None:
+    """GPUs this service uses on each of its nodes, or None when Ray places it across nodes."""
+    tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size
+    if nodes <= 1:
+        return tp_pp * service.number_of_instances
+    if service.number_of_instances > 1 and not effective_ray_serve(service, nodes, [gpus_per_node]):
+        return tp_pp * (service.number_of_instances // nodes)
+    return None
+
+
+def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
+    """The CUDA_VISIBLE_DEVICES to render, keyed by node pool and then service.
+
+    A pool is planned only when two or more vLLM services on it want their own GPUs. Each
+    takes the lowest free indices in declaration order; a ray service there gets the rest.
+    Every other service keeps seeing every GPU, as it always has.
+    """
+    compute = next(iter(config.compute.values()))
+    pools = {name: pool for name, pool in compute.node_pools.items() if pool.nodes}
+    if not pools:
+        return {}
+    single_node = sum(pool.nodes for pool in pools.values()) == 1
+    policy = config.services.get(config.driver.policy_model or "")
+    # A ray head runs beside the driver, which runs on the policy's first node.
+    driver_pool = policy.node_pool if policy is not None and policy.node_pool else next(iter(pools))
+
+    plan: dict[str, dict[str, str]] = {}
+    for pool_name, pool in pools.items():
+        if pool.gpus_per_node is None:
+            continue
+        on_pool = {
+            name: service
+            for name, service in config.services.items()
+            if single_node
+            or (
+                pool_name in (*service.node_pools, driver_pool)
+                if isinstance(service, RayServiceConfig)
+                else service.node_pool == pool_name
+            )
+        }
+        vllms = {n: s for n, s in on_pool.items() if isinstance(s, VllmServiceConfig) and _takes_own_gpus(s)}
+        if len(vllms) < 2:
+            continue
+        needs = {n: _gpus_needed_per_node(s, pool.nodes, pool.gpus_per_node) for n, s in vllms.items()}
+        if None in needs.values():
+            continue
+        manual = {n: gpus for n, s in vllms.items() if (gpus := _manual_gpus(s)) is not None}
+        needs.update({n: len(gpus) for n, gpus in manual.items()})
+
+        total = sum(needs.values())
+        if total > pool.gpus_per_node:
+            breakdown = ", ".join(f"{n} {k}" for n, k in needs.items())
+            raise ValueError(
+                f"The services sharing node pool '{pool_name}' need {total} GPUs per node ({breakdown}), but it "
+                f"has {pool.gpus_per_node}. Pin a service to its own node_pool, add a node, or set "
+                "share_gpus: true on services that may share GPUs."
+            )
+
+        owner: dict[int, str] = {}
+        for name, gpus in manual.items():
+            for index in gpus:
+                if index >= pool.gpus_per_node:
+                    raise ValueError(
+                        f"Service '{name}' sets CUDA_VISIBLE_DEVICES={','.join(map(str, gpus))}, but node pool "
+                        f"'{pool_name}' has only GPUs 0-{pool.gpus_per_node - 1}."
+                    )
+                if owner.get(index, name) != name:
+                    raise ValueError(
+                        f"Services '{owner[index]}' and '{name}' both set GPU {index} in CUDA_VISIBLE_DEVICES on "
+                        f"node pool '{pool_name}'. Give them disjoint GPUs, or set share_gpus: true on one."
+                    )
+                owner[index] = name
+
+        free = [index for index in range(pool.gpus_per_node) if index not in owner]
+        assigned: dict[str, str] = {}
+        for name in vllms:
+            if name not in manual:
+                mine, free = free[: needs[name]], free[needs[name] :]
+                assigned[name] = ",".join(map(str, mine))
+        for name, service in on_pool.items():
+            if (
+                isinstance(service, RayServiceConfig)
+                and not service.share_gpus
+                and _CUDA_VISIBLE_DEVICES not in service.env
+            ):
+                assigned[name] = ",".join(map(str, free))
+        plan[pool_name] = assigned
+    return plan
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index cc5c56b9895..c6c41496c46 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -32,6 +32,7 @@
     VllmPDTierConfig,
     VllmServiceConfig,  # used in _BUILDERS dispatch table
     effective_ray_serve,
+    plan_gpus,
 )
 from nemo_gym.orchestration.executors.otel import (
     COLLECTOR_HEALTH_PORT,
@@ -551,12 +552,20 @@ def _pool_of(compute: SlurmComputeConfig, node: int) -> str | None:
     return None
 
 
+def _with_gpus(env: dict[str, str], gpus: str | None) -> dict[str, str] | None:
+    """The service's env plus its planned CUDA_VISIBLE_DEVICES, if plan_gpus gave it one."""
+    if gpus is not None:
+        env = {**env, "CUDA_VISIBLE_DEVICES": gpus}
+    return env or None
+
+
 def _render_ray_service(
     name: str,
     service: RayServiceConfig,
     config: SubmitConfig,
     compute: SlurmComputeConfig,
     driver_node: int | None,
+    gpu_plan: dict[str, dict[str, str]],
 ) -> str:
     """A ray service's srun steps: the head beside the driver, then one step per spanned pool's workers.
 
@@ -573,7 +582,7 @@ def _render_ray_service(
             _build_ray_command(
                 service, pool=head_pool if head_pool in service.node_pools else None, address_var=address_var
             ),
-            service.env or None,
+            _with_gpus(service.env, gpu_plan.get(head_pool, {}).get(name)),
             service.mounts or None,
             nodes=_srun_nodes(service, compute, total_nodes),
             ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks),
@@ -589,7 +598,7 @@ def _render_ray_service(
                 f"{name}_{pool}_workers",
                 service.container,
                 _build_ray_command(service, pool=pool, address_var=address_var, worker=True),
-                service.env or None,
+                _with_gpus(service.env, gpu_plan.get(pool, {}).get(name)),
                 service.mounts or None,
                 nodes=count,
                 ntasks=count,
@@ -909,6 +918,9 @@ def build_sbatch_script(
     instrumented = gym_telemetry_active(config)
 
     driver_node = _driver_node(config, compute) if is_multi_node else None
+    gpu_plan = plan_gpus(config)
+    # A vLLM service sits on one pool, so its name is unique across the plan.
+    vllm_gpus = {name: gpus for planned in gpu_plan.values() for name, gpus in planned.items()}
     service_commands = "\n\n".join(
         (
             [_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node, driver_node=driver_node)]
@@ -925,7 +937,7 @@ def build_sbatch_script(
                     gpus_per_node_values,
                     offsets,
                 ),
-                service.env or None,
+                _with_gpus(service.env, vllm_gpus.get(name)),
                 service.mounts or None,
                 # Only services that actually span multiple nodes need --nodes/--ntasks - not every
                 # service in a multi-node job (e.g. a plain Ray head service runs on a single node
@@ -940,7 +952,7 @@ def build_sbatch_script(
             if not isinstance(service, RayServiceConfig)
         ]
         + [
-            _render_ray_service(name, service, config, compute, driver_node)
+            _render_ray_service(name, service, config, compute, driver_node, gpu_plan)
             for name, service in config.deployed_services.items()
             if isinstance(service, RayServiceConfig)
         ]
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index a09297adbe7..a111c90d79f 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -2539,3 +2539,193 @@ def test_a_ray_head_on_node_0_is_probed_locally(tmp_path):
     }
     script = _render(tmp_path, {"policy": _vllm(8000, "gpu", tensor_parallel_size=4), "ray": ray})
     assert "Waiting for ray at http://localhost:8011" in script
+
+
+# ---------------------------------------------------------------------------
+# GPU packing: services sharing a node get disjoint GPUs
+# ---------------------------------------------------------------------------
+
+
+_ONE_NODE = {"gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}}
+
+
+def _step(script, name):
+    return script.split(f"# service: {name}\n")[1].split("\n\n")[0]
+
+
+def _gpus(step):
+    match = re.search(r"CUDA_VISIBLE_DEVICES=(\S*)", step)
+    return match and match.group(1)
+
+
+def test_a_judge_and_a_policy_on_one_node_get_disjoint_gpus(tmp_path):
+    script = _render(
+        tmp_path,
+        {"judge": _vllm(8001, None), "policy": _vllm(8000, None, tensor_parallel_size=2)},
+        _ONE_NODE,
+    )
+    assert _gpus(_step(script, "judge")) == "0"
+    assert _gpus(_step(script, "policy")) == "1,2"
+
+
+def test_data_parallel_replicas_take_tp_times_instances(tmp_path):
+    eight = {"gpu": {**_ONE_NODE["gpu"], "gpus_per_node": 8}}
+    script = _render(
+        tmp_path,
+        {"policy": _vllm(8000, None, tensor_parallel_size=2, number_of_instances=2), "judge": _vllm(8001, None)},
+        eight,
+    )
+    # One `vllm serve --data-parallel-size 2` puts rank r on visible devices [2r, 2r+2).
+    assert _gpus(_step(script, "policy")) == "0,1,2,3"
+    assert _gpus(_step(script, "judge")) == "4"
+
+
+def test_a_lone_service_renders_no_cuda_visible_devices(tmp_path):
+    script = _render(tmp_path, {"policy": _vllm(8000, None, tensor_parallel_size=2)}, _ONE_NODE)
+    assert "CUDA_VISIBLE_DEVICES" not in script
+
+
+def test_a_lone_service_next_to_a_ray_service_renders_nothing_new(tmp_path):
+    script = _render(
+        tmp_path,
+        {"policy": _vllm(8000, None, tensor_parallel_size=2), "ray": {"type": "ray", "container": "img"}},
+        _ONE_NODE,
+    )
+    assert "CUDA_VISIBLE_DEVICES" not in script
+
+
+def test_a_manual_cuda_visible_devices_wins_and_reserves_its_gpus(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None),
+            "policy": _vllm(8000, None, tensor_parallel_size=2, env={"CUDA_VISIBLE_DEVICES": "lit:0,1"}),
+        },
+        _ONE_NODE,
+    )
+    assert _gpus(_step(script, "policy")) == "0,1"
+    assert _gpus(_step(script, "judge")) == "2"
+
+
+def test_overlapping_manual_gpus_are_refused(tmp_path):
+    with pytest.raises(ValueError, match="Services 'a' and 'b' both set GPU 1"):
+        _placement_config(
+            tmp_path,
+            {
+                "a": _vllm(8000, None, env={"CUDA_VISIBLE_DEVICES": "lit:1"}),
+                "b": _vllm(8001, None, tensor_parallel_size=2, env={"CUDA_VISIBLE_DEVICES": "lit:1,2"}),
+            },
+            _ONE_NODE,
+        )
+
+
+def test_a_manual_gpu_past_the_node_is_refused(tmp_path):
+    with pytest.raises(ValueError, match="has only GPUs 0-3"):
+        _placement_config(
+            tmp_path,
+            {"a": _vllm(8000, None, env={"CUDA_VISIBLE_DEVICES": "lit:4"}), "b": _vllm(8001, None)},
+            _ONE_NODE,
+        )
+
+
+def test_services_that_need_more_gpus_than_the_node_has_are_refused(tmp_path):
+    with pytest.raises(ValueError, match=r"need 5 GPUs per node \(judge 1, policy 4\), but it has 4") as error:
+        _placement_config(
+            tmp_path, {"judge": _vllm(8001, None), "policy": _vllm(8000, None, tensor_parallel_size=4)}, _ONE_NODE
+        )
+    assert "node_pool" in str(error.value)
+    assert "share_gpus" in str(error.value)
+
+
+def test_share_gpus_opts_a_service_out(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None, share_gpus=True),
+            "policy": _vllm(8000, None, tensor_parallel_size=4),
+            "scorer": _vllm(8002, None, tensor_parallel_size=4, share_gpus=True),
+        },
+        _ONE_NODE,
+    )
+    assert "CUDA_VISIBLE_DEVICES" not in script
+
+
+def test_an_unreadable_manual_value_is_kept_and_shares(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None, env={"CUDA_VISIBLE_DEVICES": "runtime:CUDA_VISIBLE_DEVICES"}),
+            "policy": _vllm(8000, None, tensor_parallel_size=4),
+        },
+        _ONE_NODE,
+    )
+    assert "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" in _step(script, "judge")
+    assert "CUDA_VISIBLE_DEVICES" not in _step(script, "policy")
+
+
+def test_a_ray_service_gets_the_leftover_gpus(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None),
+            "policy": _vllm(8000, None, tensor_parallel_size=2),
+            "ray": {"type": "ray", "container": "img"},
+        },
+        _ONE_NODE,
+    )
+    assert _gpus(_step(script, "ray")) == "3"
+
+
+def test_a_ray_service_with_its_own_gpus_keeps_them(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None),
+            "policy": _vllm(8000, None, tensor_parallel_size=2),
+            "ray": {"type": "ray", "container": "img", "env": {"CUDA_VISIBLE_DEVICES": "lit:0,1,2,3"}},
+        },
+        _ONE_NODE,
+    )
+    assert _gpus(_step(script, "ray")) == "0,1,2,3"
+
+
+def test_pinned_pools_are_packed_per_pool(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, "gpu", tensor_parallel_size=4),
+            "judge": _vllm(8001, "aux"),
+            "scorer": _vllm(8002, "aux", tensor_parallel_size=2),
+            "ray": {"type": "ray", "container": "img", "node_pools": ["gpu", "aux"]},
+        },
+    )
+    # The policy is alone on its pool, so it and the ray head beside it are untouched.
+    assert "CUDA_VISIBLE_DEVICES" not in _step(script, "policy")
+    assert "CUDA_VISIBLE_DEVICES" not in _step(script, "ray")
+    assert _gpus(_step(script, "judge")) == "0"
+    assert _gpus(_step(script, "scorer")) == "1,2"
+    assert _gpus(_step(script, "ray_aux_workers")) == "3"
+
+
+def test_multi_node_data_parallel_is_packed_per_node(tmp_path):
+    pools = {"gpu": {"partition": "batch", "nodes": 2, "ntasks_per_node": 1, "gpus_per_node": 4}}
+    script = _render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, "gpu", tensor_parallel_size=1, number_of_instances=4),
+            "judge": _vllm(8001, "gpu", number_of_instances=2),
+        },
+        pools,
+    )
+    assert _gpus(_step(script, "policy")) == "0,1"
+    assert _gpus(_step(script, "judge")) == "2"
+
+
+def test_a_ray_placed_multi_node_service_leaves_its_pool_as_today(tmp_path):
+    pools = {"gpu": {"partition": "batch", "nodes": 2, "ntasks_per_node": 1, "gpus_per_node": 4}}
+    script = _render(
+        tmp_path,
+        {"policy": _vllm(8000, "gpu", tensor_parallel_size=8), "judge": _vllm(8001, "gpu", number_of_instances=2)},
+        pools,
+    )
+    assert "CUDA_VISIBLE_DEVICES" not in script

From 820824f55efe1eebe950315c99f32f3a4b65fefe Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Thu, 1 Oct 2026 13:38:31 +0200
Subject: [PATCH 2/6] fix(orchestration): export planned GPUs inside the step;
 pack a lone vLLM beside Ray

Slurm resets CUDA_VISIBLE_DEVICES to the step's GPUs, so a value passed as
`env K=V srun` never reaches the task. Checked on a 4-GPU GB200 node: the
env-prefix value came back as 0,1,2,3, while `export` inside the step's
command left CUDA seeing only the planned GPUs. Planned and manual values on
a packed node are now exported in the step.

A lone vLLM service beside a GPU-using ray service is now packed too, so
Ray no longer advertises the model's GPUs. A ray service with num_gpus: 0,
or a single GPU service on a node, still renders nothing new.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py                 | 26 ++++----
 .../orchestration/executors/slurm_script.py   | 35 +++++++----
 tests/unit_tests/test_slurm_script.py         | 62 ++++++++++++++++---
 3 files changed, 91 insertions(+), 32 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index ac0e1a59c97..a93b11bfc03 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -717,7 +717,7 @@ def _manual_gpus(service: BaseServiceConfig) -> list[int] | None:
 
 def _takes_own_gpus(service: BaseServiceConfig) -> bool:
     # A CUDA_VISIBLE_DEVICES that is not an index list (e.g. runtime:VAR) names unknown GPUs.
-    if service.share_gpus:
+    if service.share_gpus or (isinstance(service, RayServiceConfig) and service.num_gpus == 0):
         return False
     return _CUDA_VISIBLE_DEVICES not in service.env or _manual_gpus(service) is not None
 
@@ -735,9 +735,9 @@ def _gpus_needed_per_node(service: VllmServiceConfig, nodes: int, gpus_per_node:
 def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
     """The CUDA_VISIBLE_DEVICES to render, keyed by node pool and then service.
 
-    A pool is planned only when two or more vLLM services on it want their own GPUs. Each
-    takes the lowest free indices in declaration order; a ray service there gets the rest.
-    Every other service keeps seeing every GPU, as it always has.
+    A pool is planned only when a vLLM service on it shares it with another service that wants
+    its own GPUs (a vLLM or a ray service). Each vLLM takes the lowest free indices in declaration
+    order; a ray service there gets the rest. Every other service keeps seeing every GPU.
     """
     compute = next(iter(config.compute.values()))
     pools = {name: pool for name, pool in compute.node_pools.items() if pool.nodes}
@@ -762,13 +762,14 @@ def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
                 else service.node_pool == pool_name
             )
         }
-        vllms = {n: s for n, s in on_pool.items() if isinstance(s, VllmServiceConfig) and _takes_own_gpus(s)}
-        if len(vllms) < 2:
+        claimants = {n: s for n, s in on_pool.items() if _takes_own_gpus(s)}
+        vllms = {n: s for n, s in claimants.items() if isinstance(s, VllmServiceConfig)}
+        if not vllms or len(claimants) < 2:
             continue
         needs = {n: _gpus_needed_per_node(s, pool.nodes, pool.gpus_per_node) for n, s in vllms.items()}
         if None in needs.values():
             continue
-        manual = {n: gpus for n, s in vllms.items() if (gpus := _manual_gpus(s)) is not None}
+        manual = {n: gpus for n, s in claimants.items() if (gpus := _manual_gpus(s)) is not None}
         needs.update({n: len(gpus) for n, gpus in manual.items()})
 
         total = sum(needs.values())
@@ -801,12 +802,9 @@ def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
             if name not in manual:
                 mine, free = free[: needs[name]], free[needs[name] :]
                 assigned[name] = ",".join(map(str, mine))
-        for name, service in on_pool.items():
-            if (
-                isinstance(service, RayServiceConfig)
-                and not service.share_gpus
-                and _CUDA_VISIBLE_DEVICES not in service.env
-            ):
-                assigned[name] = ",".join(map(str, free))
+        for name in claimants.keys() - vllms.keys() - manual.keys():
+            assigned[name] = ",".join(map(str, free))
+        # Manual values are re-rendered too, where Slurm cannot reset them (see slurm_script._with_gpus).
+        assigned.update({name: ",".join(map(str, gpus)) for name, gpus in manual.items()})
         plan[pool_name] = assigned
     return plan
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index c6c41496c46..94e8971af2a 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -552,11 +552,16 @@ def _pool_of(compute: SlurmComputeConfig, node: int) -> str | None:
     return None
 
 
-def _with_gpus(env: dict[str, str], gpus: str | None) -> dict[str, str] | None:
-    """The service's env plus its planned CUDA_VISIBLE_DEVICES, if plan_gpus gave it one."""
-    if gpus is not None:
-        env = {**env, "CUDA_VISIBLE_DEVICES": gpus}
-    return env or None
+def _with_gpus(env: dict[str, str], pre_command: str, gpus: str | None) -> tuple[dict[str, str] | None, str]:
+    """The step's env and pre_command, exporting the CUDA_VISIBLE_DEVICES plan_gpus gave it.
+
+    Exported inside the step: Slurm resets CUDA_VISIBLE_DEVICES to the step's GPUs, so an
+    `env K=V srun` value never reaches the task.
+    """
+    if gpus is None:
+        return env or None, pre_command
+    env = {k: v for k, v in env.items() if k != "CUDA_VISIBLE_DEVICES"}
+    return env or None, f"export CUDA_VISIBLE_DEVICES={shlex.quote(gpus)}\n{pre_command}".rstrip()
 
 
 def _render_ray_service(
@@ -575,6 +580,7 @@ def _render_ray_service(
     head_node = _driver_node(config, compute)
     head_pool = _pool_of(compute, head_node)
     address_var = ray_head_address_var(name)
+    head_env, head_pre = _with_gpus(service.env, service.pre_command, gpu_plan.get(head_pool, {}).get(name))
     steps = [
         _render_service_command(
             name,
@@ -582,27 +588,30 @@ def _render_ray_service(
             _build_ray_command(
                 service, pool=head_pool if head_pool in service.node_pools else None, address_var=address_var
             ),
-            _with_gpus(service.env, gpu_plan.get(head_pool, {}).get(name)),
+            head_env,
             service.mounts or None,
             nodes=_srun_nodes(service, compute, total_nodes),
             ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks),
-            pre_command=service.pre_command,
+            pre_command=head_pre,
             nodelist=_service_nodelist(service, driver_node, total_nodes),
         )
     ]
     # A worker only joins a running head, and the head may still be installing.
     wait_for_head = f'until ray status --address "${address_var}" >/dev/null 2>&1; do sleep 5; done'
     for pool, (_, count) in _ray_worker_ranges(service, compute, head_node).items():
+        worker_env, worker_pre = _with_gpus(
+            service.env, f"{service.pre_command.rstrip()}\n{wait_for_head}".lstrip(), gpu_plan.get(pool, {}).get(name)
+        )
         steps.append(
             _render_service_command(
                 f"{name}_{pool}_workers",
                 service.container,
                 _build_ray_command(service, pool=pool, address_var=address_var, worker=True),
-                _with_gpus(service.env, gpu_plan.get(pool, {}).get(name)),
+                worker_env,
                 service.mounts or None,
                 nodes=count,
                 ntasks=count,
-                pre_command=f"{service.pre_command.rstrip()}\n{wait_for_head}".lstrip(),
+                pre_command=worker_pre,
                 nodelist=ray_workers_var(name, pool),
             )
         )
@@ -921,6 +930,10 @@ def build_sbatch_script(
     gpu_plan = plan_gpus(config)
     # A vLLM service sits on one pool, so its name is unique across the plan.
     vllm_gpus = {name: gpus for planned in gpu_plan.values() for name, gpus in planned.items()}
+    env_and_pre = {
+        name: _with_gpus(s.env, _service_pre_command(s), vllm_gpus.get(name))
+        for name, s in config.deployed_services.items()
+    }
     service_commands = "\n\n".join(
         (
             [_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node, driver_node=driver_node)]
@@ -937,7 +950,7 @@ def build_sbatch_script(
                     gpus_per_node_values,
                     offsets,
                 ),
-                _with_gpus(service.env, vllm_gpus.get(name)),
+                env_and_pre[name][0],
                 service.mounts or None,
                 # Only services that actually span multiple nodes need --nodes/--ntasks - not every
                 # service in a multi-node job (e.g. a plain Ray head service runs on a single node
@@ -945,7 +958,7 @@ def build_sbatch_script(
                 # them, so the node list and the step size agree.
                 nodes=_srun_nodes(service, compute, total_nodes),
                 ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks),
-                pre_command=_service_pre_command(service),
+                pre_command=env_and_pre[name][1],
                 nodelist=_service_nodelist(service, driver_node, total_nodes),
             )
             for name, service in config.deployed_services.items()
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index a111c90d79f..cc43b4c1fc1 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -2554,7 +2554,12 @@ def _step(script, name):
 
 
 def _gpus(step):
-    match = re.search(r"CUDA_VISIBLE_DEVICES=(\S*)", step)
+    # Slurm resets an `env K=V srun` value, so the plan is exported inside the step's command.
+    match = (
+        re.search(r"export CUDA_VISIBLE_DEVICES=([^\n]*)", shlex.split(step.split("bash -c ", 1)[1])[0])
+        if "bash -c " in step
+        else None
+    )
     return match and match.group(1)
 
 
@@ -2585,12 +2590,31 @@ def test_a_lone_service_renders_no_cuda_visible_devices(tmp_path):
     assert "CUDA_VISIBLE_DEVICES" not in script
 
 
-def test_a_lone_service_next_to_a_ray_service_renders_nothing_new(tmp_path):
+def test_a_lone_vllm_next_to_a_gpu_ray_service_is_packed(tmp_path):
+    # Otherwise Ray advertises the policy's GPUs and schedules work onto them.
     script = _render(
         tmp_path,
         {"policy": _vllm(8000, None, tensor_parallel_size=2), "ray": {"type": "ray", "container": "img"}},
         _ONE_NODE,
     )
+    assert _gpus(_step(script, "policy")) == "0,1"
+    assert _gpus(_step(script, "ray")) == "2,3"
+
+
+def test_a_lone_vllm_next_to_a_cpu_only_ray_service_renders_nothing_new(tmp_path):
+    script = _render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, None, tensor_parallel_size=2),
+            "ray": {"type": "ray", "container": "img", "num_gpus": 0},
+        },
+        _ONE_NODE,
+    )
+    assert "CUDA_VISIBLE_DEVICES" not in script
+
+
+def test_a_lone_ray_service_renders_nothing_new(tmp_path):
+    script = _render(tmp_path, {"ray": {"type": "ray", "container": "img"}}, _ONE_NODE)
     assert "CUDA_VISIBLE_DEVICES" not in script
 
 
@@ -2682,11 +2706,13 @@ def test_a_ray_service_with_its_own_gpus_keeps_them(tmp_path):
         {
             "judge": _vllm(8001, None),
             "policy": _vllm(8000, None, tensor_parallel_size=2),
-            "ray": {"type": "ray", "container": "img", "env": {"CUDA_VISIBLE_DEVICES": "lit:0,1,2,3"}},
+            "ray": {"type": "ray", "container": "img", "env": {"CUDA_VISIBLE_DEVICES": "lit:0"}},
         },
         _ONE_NODE,
     )
-    assert _gpus(_step(script, "ray")) == "0,1,2,3"
+    assert _gpus(_step(script, "ray")) == "0"
+    assert _gpus(_step(script, "judge")) == "1"
+    assert _gpus(_step(script, "policy")) == "2,3"
 
 
 def test_pinned_pools_are_packed_per_pool(tmp_path):
@@ -2699,9 +2725,9 @@ def test_pinned_pools_are_packed_per_pool(tmp_path):
             "ray": {"type": "ray", "container": "img", "node_pools": ["gpu", "aux"]},
         },
     )
-    # The policy is alone on its pool, so it and the ray head beside it are untouched.
-    assert "CUDA_VISIBLE_DEVICES" not in _step(script, "policy")
-    assert "CUDA_VISIBLE_DEVICES" not in _step(script, "ray")
+    # The policy fills its pool, so the ray head beside it sees no GPUs.
+    assert _gpus(_step(script, "policy")) == "0,1,2,3"
+    assert _gpus(_step(script, "ray")) == "''"
     assert _gpus(_step(script, "judge")) == "0"
     assert _gpus(_step(script, "scorer")) == "1,2"
     assert _gpus(_step(script, "ray_aux_workers")) == "3"
@@ -2729,3 +2755,25 @@ def test_a_ray_placed_multi_node_service_leaves_its_pool_as_today(tmp_path):
         pools,
     )
     assert "CUDA_VISIBLE_DEVICES" not in script
+
+
+def test_planned_gpus_are_exported_inside_the_step_not_passed_through_srun(tmp_path):
+    # On a cluster, Slurm reset `env CUDA_VISIBLE_DEVICES=0 srun ...` to the step's four GPUs.
+    script = _render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None, env={"HF_HOME": "lit:/hf"}),
+            "policy": _vllm(8000, None, tensor_parallel_size=2, env={"CUDA_VISIBLE_DEVICES": "lit:2,3"}),
+        },
+        _ONE_NODE,
+    )
+    for name in ("judge", "policy"):
+        assert "CUDA_VISIBLE_DEVICES" not in _step(script, name).split("srun", 1)[0]
+    assert _step(script, "judge").startswith("env HF_HOME=/hf srun")
+    assert _gpus(_step(script, "policy")) == "2,3"
+
+
+def test_a_lone_services_manual_value_is_rendered_as_before(tmp_path):
+    script = _render(tmp_path, {"policy": _vllm(8000, None, env={"CUDA_VISIBLE_DEVICES": "lit:0"})}, _ONE_NODE)
+    assert _step(script, "policy").startswith("env CUDA_VISIBLE_DEVICES=0 srun")
+    assert "export CUDA_VISIBLE_DEVICES" not in script

From af57137f145e02544ee795aa0d36097c6c21d7f8 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Thu, 1 Oct 2026 13:38:57 +0200
Subject: [PATCH 3/6] feat(orchestration): size a node pool automatically with
 nodes: auto

A pool with `nodes: auto` is sized by Gym. The policy goes on the pool's
first node, beside the driver. The other vLLM services go biggest first
into the first node with room, never split. A service bigger than one node
gets ceil(need / gpus_per_node) whole nodes of its own. More than one node
splits the pool into sub-pools `<pool>-<i>`, placed with the existing
node-pool --nodelist mechanism; a ray service spanning the pool spans them
all and gets each node's leftover GPUs. Explicit node counts are unchanged.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py         |  99 ++++++++++++++++-
 tests/unit_tests/test_slurm_script.py | 153 ++++++++++++++++++++++++++
 2 files changed, 251 insertions(+), 1 deletion(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index a93b11bfc03..736557d0958 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -296,7 +296,8 @@ def _validate_shape(self) -> "RayServiceConfig":
 
 class NodePool(_StrictModel):
     partition: str
-    nodes: int = 1
+    # "auto": Gym packs the pool's vLLM services onto as few nodes as fit (see _resolve_auto_pool).
+    nodes: int | Literal["auto"] = 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
@@ -498,6 +499,9 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
 
         sole_compute = next(iter(compute_names))
         compute = self.compute[sole_compute]
+        if isinstance(compute, SlurmComputeConfig):
+            for pool_name in [name for name, pool in compute.node_pools.items() if pool.nodes == "auto"]:
+                _resolve_auto_pool(self, compute, pool_name)
         total_nodes = (
             sum(p.nodes for p in compute.node_pools.values()) if isinstance(compute, SlurmComputeConfig) else 1
         )
@@ -808,3 +812,96 @@ def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
         assigned.update({name: ",".join(map(str, gpus)) for name, gpus in manual.items()})
         plan[pool_name] = assigned
     return plan
+
+
+def _resolve_auto_pool(config: SubmitConfig, compute: SlurmComputeConfig, pool_name: str) -> None:
+    """Size a `nodes: auto` pool and pin each of its vLLM services to a node of it.
+
+    The policy goes first, on the pool's first node, where the driver reaches it on localhost.
+    The rest go biggest first into the first node with room. A service bigger than one node
+    gets ceil(need / gpus_per_node) whole nodes to itself. More than one node splits the pool
+    into sub-pools named `<pool>-<i>`, which the rest of the executor places like any pool.
+    """
+    pool = compute.node_pools[pool_name]
+    gpus_per_node = pool.gpus_per_node
+    if gpus_per_node is None:
+        raise ValueError(
+            f"Node pool '{pool_name}' has nodes: auto but no gpus_per_node, so Gym cannot size it. "
+            "Set gpus_per_node, or give the pool an explicit node count."
+        )
+    sole_pool = len(compute.node_pools) == 1
+    members = {
+        name: service
+        for name, service in config.services.items()
+        if isinstance(service, VllmServiceConfig)
+        and (service.node_pool == pool_name or (sole_pool and service.node_pool is None))
+    }
+    if not members:
+        raise ValueError(
+            f"Node pool '{pool_name}' has nodes: auto but no vLLM service runs on it, so there is nothing to "
+            "size it by. Pin a service to it with node_pool, or give it an explicit node count."
+        )
+
+    def need(service: VllmServiceConfig) -> int:
+        if service.share_gpus:
+            return 0
+        manual = _manual_gpus(service)
+        if manual is not None:
+            return len(manual)
+        return service.tensor_parallel_size * service.pipeline_parallel_size * service.number_of_instances
+
+    policy = config.driver.policy_model
+    order = sorted(members, key=lambda name: (name != policy, -need(members[name])))
+    # One slot per sub-pool: the services on it, the manual indices they hold, the GPUs the
+    # rest take, and its node count (more than one only for a service bigger than a node).
+    bins: list[dict[str, Any]] = []
+    for name in order:
+        gpus = need(members[name])
+        manual = set() if members[name].share_gpus else set(_manual_gpus(members[name]) or [])
+        auto = 0 if manual else gpus
+        if gpus > gpus_per_node:
+            if manual:
+                raise ValueError(
+                    f"Service '{name}' sets {gpus} GPUs in CUDA_VISIBLE_DEVICES, but node pool '{pool_name}' has "
+                    f"{gpus_per_node} per node."
+                )
+            bins.append({"names": [name], "manual": set(), "auto": gpus, "nodes": -(-gpus // gpus_per_node)})
+            continue
+        for slot in bins:
+            if (
+                slot["nodes"] == 1
+                and not slot["manual"] & manual
+                and len(slot["manual"] | manual) + slot["auto"] + auto <= gpus_per_node
+            ):
+                slot["names"].append(name)
+                slot["manual"] |= manual
+                slot["auto"] += auto
+                break
+        else:
+            bins.append({"names": [name], "manual": manual, "auto": auto, "nodes": 1})
+
+    if len(bins) == 1 and bins[0]["nodes"] == 1:
+        compute.node_pools[pool_name] = pool.model_copy(update={"nodes": 1})
+        return
+
+    sub_pools = {f"{pool_name}-{index}": slot for index, slot in enumerate(bins)}
+    clashes = sorted(set(sub_pools) & set(compute.node_pools))
+    if clashes:
+        raise ValueError(f"Node pool '{pool_name}' has nodes: auto, but its sub-pool names {clashes} are taken.")
+    resolved: dict[str, NodePool] = {}
+    for name, existing in compute.node_pools.items():
+        if name != pool_name:
+            resolved[name] = existing
+            continue
+        for sub, slot in sub_pools.items():
+            resolved[sub] = pool.model_copy(update={"nodes": slot["nodes"]})
+            for member in slot["names"]:
+                members[member].node_pool = sub
+    compute.node_pools = resolved
+    for service in config.services.values():
+        if isinstance(service, RayServiceConfig) and pool_name in service.node_pools:
+            at = service.node_pools.index(pool_name)
+            service.node_pools[at : at + 1] = list(sub_pools)
+            if pool_name in service.resources:
+                resources = service.resources.pop(pool_name)
+                service.resources.update({sub: dict(resources) for sub in sub_pools})
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index cc43b4c1fc1..8a015febd9c 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -2777,3 +2777,156 @@ def test_a_lone_services_manual_value_is_rendered_as_before(tmp_path):
     script = _render(tmp_path, {"policy": _vllm(8000, None, env={"CUDA_VISIBLE_DEVICES": "lit:0"})}, _ONE_NODE)
     assert _step(script, "policy").startswith("env CUDA_VISIBLE_DEVICES=0 srun")
     assert "export CUDA_VISIBLE_DEVICES" not in script
+
+
+# ---------------------------------------------------------------------------
+# nodes: auto
+# ---------------------------------------------------------------------------
+
+
+def _auto_pool(gpus_per_node=4):
+    return {"gpu": {"partition": "batch", "nodes": "auto", "ntasks_per_node": 1, "gpus_per_node": gpus_per_node}}
+
+
+def _auto_config(tmp_path, services, pools=None, policy=None):
+    return SubmitConfig.model_validate(
+        {
+            "services": services,
+            "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": pools or _auto_pool()}},
+            "driver": {"container": "gym:latest", "policy_model": policy, "benchmarks": {"b": {"run": {}}}},
+            "job": {"output_path": str(tmp_path / "jobs")},
+        }
+    )
+
+
+def _auto_render(tmp_path, services, pools=None, policy=None):
+    config = _auto_config(tmp_path, services, pools, policy)
+    return config, build_sbatch_script(
+        config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "jobs" / "b"
+    )
+
+
+def test_auto_packs_a_judge_and_a_two_gpu_model_onto_one_node(tmp_path):
+    services = {"judge": _vllm(8001, None), "policy": _vllm(8000, None, tensor_parallel_size=2)}
+    config, script = _auto_render(tmp_path, services, policy="policy")
+    assert config.compute["hsg"].node_pools["gpu"].nodes == 1
+    assert "#SBATCH --nodes=1" in script
+    # GPU indices follow declaration order on a node.
+    assert _gpus(_step(script, "judge")) == "0"
+    assert _gpus(_step(script, "policy")) == "1,2"
+    # Exactly what the same config with an explicit `nodes: 1` renders.
+    _, explicit = _auto_render(tmp_path, services, {"gpu": {**_auto_pool()["gpu"], "nodes": 1}}, policy="policy")
+    assert script == explicit
+
+
+def test_auto_gives_a_service_bigger_than_a_node_whole_nodes_of_its_own(tmp_path):
+    # 6 GPUs on 4-GPU nodes: ceil(6 / 4) = 2 whole nodes, nothing else on them.
+    config, script = _auto_render(
+        tmp_path,
+        {
+            "judge": _vllm(8001, None, tensor_parallel_size=4),
+            "policy": _vllm(8000, None, number_of_instances=6),
+        },
+    )
+    pools = config.compute["hsg"].node_pools
+    assert {name: pool.nodes for name, pool in pools.items()} == {"gpu-0": 2, "gpu-1": 1}
+    assert config.services["policy"].node_pool == "gpu-0"
+    assert config.services["judge"].node_pool == "gpu-1"
+    assert "#SBATCH --nodes=3" in script
+    assert '--nodelist="${GYM_POOL_GPU_0_NODES}" --nodes=2 --ntasks=2' in _step(script, "policy")
+    assert '--nodelist="${GYM_POOL_GPU_1_NODES}" --nodes=1 --ntasks=1' in _step(script, "judge")
+    assert "--data-parallel-size-local 3" in _step(script, "policy")
+    # Each is alone on its nodes, so neither is narrowed.
+    assert "CUDA_VISIBLE_DEVICES" not in script
+
+
+def test_auto_packs_biggest_first_into_the_first_node_with_room(tmp_path):
+    config = _auto_config(
+        tmp_path,
+        {
+            "a": _vllm(8001, None),
+            "b": _vllm(8002, None, tensor_parallel_size=3),
+            "c": _vllm(8003, None, tensor_parallel_size=2),
+            "d": _vllm(8004, None, tensor_parallel_size=2),
+        },
+    )
+    placed = {name: service.node_pool for name, service in config.services.items()}
+    assert placed == {"b": "gpu-0", "c": "gpu-1", "d": "gpu-1", "a": "gpu-0"}
+
+
+def test_auto_puts_the_policy_on_node_0_beside_the_driver(tmp_path):
+    config, script = _auto_render(
+        tmp_path,
+        {"scorer": _vllm(8001, None, tensor_parallel_size=4), "policy": _vllm(8000, None)},
+        policy="policy",
+    )
+    assert config.services["policy"].node_pool == "gpu-0"
+    assert config.services["scorer"].node_pool == "gpu-1"
+    driver = next(line for line in script.splitlines() if "--output=logs/driver-$SLURM_JOB_ID.log" in line)
+    assert '--nodelist="${gym_nodes[0]}"' in driver
+
+
+def test_explicit_node_counts_are_not_touched(tmp_path):
+    config = _placement_config(
+        tmp_path,
+        {"policy": _vllm(8000, "gpu", tensor_parallel_size=4), "scorer": _vllm(8001, "aux", tensor_parallel_size=4)},
+        _TWO_POOLS,
+    )
+    assert {name: pool.nodes for name, pool in config.compute["hsg"].node_pools.items()} == {"gpu": 1, "aux": 1}
+    assert config.services["policy"].node_pool == "gpu"
+
+
+def test_auto_gives_the_leftovers_to_ray_on_every_node(tmp_path):
+    config, script = _auto_render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, None, tensor_parallel_size=3),
+            "judge": _vllm(8001, None, tensor_parallel_size=2),
+            "ray": {"type": "ray", "container": "img", "node_pools": ["gpu"], "resources": {"gpu": {"x": 1}}},
+        },
+        policy="policy",
+    )
+    ray = config.services["ray"]
+    assert ray.node_pools == ["gpu-0", "gpu-1"]
+    assert set(ray.resources) == {"gpu-0", "gpu-1"}
+    assert _gpus(_step(script, "policy")) == "0,1,2"
+    assert _gpus(_step(script, "ray")) == "3"
+    assert _gpus(_step(script, "judge")) == "0,1"
+    assert _gpus(_step(script, "ray_gpu-1_workers")) == "2,3"
+
+
+def test_auto_keeps_a_manual_cuda_visible_devices(tmp_path):
+    config, script = _auto_render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, None, tensor_parallel_size=2),
+            "judge": _vllm(8001, None, env={"CUDA_VISIBLE_DEVICES": "lit:0"}),
+        },
+        policy="policy",
+    )
+    assert config.compute["hsg"].node_pools["gpu"].nodes == 1
+    assert _gpus(_step(script, "judge")) == "0"
+    assert _gpus(_step(script, "policy")) == "1,2"
+
+
+def test_auto_without_gpus_per_node_is_refused(tmp_path):
+    pools = {"gpu": {"partition": "batch", "nodes": "auto"}}
+    with pytest.raises(ValueError, match="nodes: auto but no gpus_per_node"):
+        _auto_config(tmp_path, {"policy": _vllm(8000, None)}, pools)
+
+
+def test_auto_with_no_vllm_service_is_refused(tmp_path):
+    with pytest.raises(ValueError, match="no vLLM service runs on it"):
+        _auto_config(tmp_path, {"ray": {"type": "ray", "container": "img", "node_pools": ["gpu"]}})
+
+
+def test_auto_puts_services_with_clashing_manual_gpus_on_different_nodes(tmp_path):
+    config = _auto_config(
+        tmp_path,
+        {
+            "a": _vllm(8001, None, env={"CUDA_VISIBLE_DEVICES": "lit:0"}),
+            "b": _vllm(8002, None, env={"CUDA_VISIBLE_DEVICES": "lit:0"}),
+        },
+    )
+    assert config.services["a"].node_pool == "gpu-0"
+    assert config.services["b"].node_pool == "gpu-1"

From 232ef671dd00b83be6b975cb5c937fc549084d62 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Thu, 1 Oct 2026 13:39:09 +0200
Subject: [PATCH 4/6] fix(orchestration): make a hand-set CUDA_VISIBLE_DEVICES
 take effect

A service's own CUDA_VISIBLE_DEVICES was passed as `env K=V srun`, which
Slurm resets to the step's GPUs, so it was silently ignored. It is now
exported inside the step's command on every node, including a node with a
single service. A runtime:VAR value is relayed as GYM_CUDA_VISIBLE_DEVICES
so the step still exports the batch shell's $VAR.

Drop runtime:CUDA_VISIBLE_DEVICES from the multi-instance example: it
re-exported the job's full GPU set and did nothing.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 examples/slurm_vllm_multi_instance.yaml       |  3 +-
 .../orchestration/executors/slurm_script.py   | 17 ++++++++--
 tests/unit_tests/test_slurm_script.py         | 31 ++++++++++++++++---
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/examples/slurm_vllm_multi_instance.yaml b/examples/slurm_vllm_multi_instance.yaml
index f5af9d02a36..e3ea98732cf 100644
--- a/examples/slurm_vllm_multi_instance.yaml
+++ b/examples/slurm_vllm_multi_instance.yaml
@@ -32,8 +32,7 @@ services:
     env:
       HF_TOKEN: host:HF_TOKEN
       TRUST_REMOTE_CODE: lit:1
-      NEL_INVOCATION_ID: runtime:NEL_INVOCATION_ID      
-      CUDA_VISIBLE_DEVICES: runtime:CUDA_VISIBLE_DEVICES
+      NEL_INVOCATION_ID: runtime:NEL_INVOCATION_ID
 
 driver:
   policy_model: vllm_model
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index 94e8971af2a..60bce9d3634 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -552,16 +552,27 @@ def _pool_of(compute: SlurmComputeConfig, node: int) -> str | None:
     return None
 
 
+_GPU_RELAY_VAR = "GYM_CUDA_VISIBLE_DEVICES"
+
+
 def _with_gpus(env: dict[str, str], pre_command: str, gpus: str | None) -> tuple[dict[str, str] | None, str]:
-    """The step's env and pre_command, exporting the CUDA_VISIBLE_DEVICES plan_gpus gave it.
+    """The step's env and pre_command, exporting the CUDA_VISIBLE_DEVICES plan_gpus gave it,
+    else the service's own.
 
     Exported inside the step: Slurm resets CUDA_VISIBLE_DEVICES to the step's GPUs, so an
-    `env K=V srun` value never reaches the task.
+    `env K=V srun` value never reaches the task. A runtime:VAR value is relayed under another
+    name, so the step still exports the batch shell's $VAR.
     """
+    gpus = gpus if gpus is not None else env.get("CUDA_VISIBLE_DEVICES")
     if gpus is None:
         return env or None, pre_command
     env = {k: v for k, v in env.items() if k != "CUDA_VISIBLE_DEVICES"}
-    return env or None, f"export CUDA_VISIBLE_DEVICES={shlex.quote(gpus)}\n{pre_command}".rstrip()
+    if gpus.startswith(RUNTIME_ENV_PREFIX):
+        env[_GPU_RELAY_VAR] = gpus
+        export = f'export CUDA_VISIBLE_DEVICES="${_GPU_RELAY_VAR}"'
+    else:
+        export = f"export CUDA_VISIBLE_DEVICES={shlex.quote(gpus)}"
+    return env or None, f"{export}\n{pre_command}".rstrip()
 
 
 def _render_ray_service(
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 8a015febd9c..f1663ce8d77 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -20,6 +20,7 @@
 from pathlib import Path
 
 import pytest
+import yaml
 
 from nemo_gym.orchestration.api import NodePool, RayServiceConfig, SubmitConfig, VllmPDTierConfig
 from nemo_gym.orchestration.executors.script_templates import (
@@ -2683,7 +2684,7 @@ def test_an_unreadable_manual_value_is_kept_and_shares(tmp_path):
         },
         _ONE_NODE,
     )
-    assert "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" in _step(script, "judge")
+    assert _gpus(_step(script, "judge")) == '"$GYM_CUDA_VISIBLE_DEVICES"'
     assert "CUDA_VISIBLE_DEVICES" not in _step(script, "policy")
 
 
@@ -2773,10 +2774,32 @@ def test_planned_gpus_are_exported_inside_the_step_not_passed_through_srun(tmp_p
     assert _gpus(_step(script, "policy")) == "2,3"
 
 
-def test_a_lone_services_manual_value_is_rendered_as_before(tmp_path):
+def test_a_lone_services_manual_value_is_exported_inside_the_step(tmp_path):
+    # Left in the srun env, Slurm silently replaced it with every GPU of the step.
     script = _render(tmp_path, {"policy": _vllm(8000, None, env={"CUDA_VISIBLE_DEVICES": "lit:0"})}, _ONE_NODE)
-    assert _step(script, "policy").startswith("env CUDA_VISIBLE_DEVICES=0 srun")
-    assert "export CUDA_VISIBLE_DEVICES" not in script
+    assert "CUDA_VISIBLE_DEVICES" not in _step(script, "policy").split("srun", 1)[0]
+    assert _gpus(_step(script, "policy")) == "0"
+
+
+def test_a_runtime_manual_value_exports_the_batch_shells_value(tmp_path):
+    script = _render(
+        tmp_path, {"policy": _vllm(8000, None, env={"CUDA_VISIBLE_DEVICES": "runtime:MY_GPUS"})}, _ONE_NODE
+    )
+    step = _step(script, "policy")
+    assert step.startswith("env GYM_CUDA_VISIBLE_DEVICES=${MY_GPUS} srun")
+    assert _gpus(step) == '"$GYM_CUDA_VISIBLE_DEVICES"'
+
+
+def test_the_multi_instance_example_renders(tmp_path, monkeypatch):
+    raw = yaml.safe_load((Path(__file__).parents[2] / "examples" / "slurm_vllm_multi_instance.yaml").read_text())
+    raw["job"]["output_path"] = str(tmp_path)
+    monkeypatch.setenv("HF_TOKEN", "x")
+    config = SubmitConfig.model_validate(raw)
+    compute = next(iter(config.compute.values()))
+    script = build_sbatch_script(config, "gpqa", config.driver.benchmarks["gpqa"], compute, tmp_path / "gpqa")
+    assert "CUDA_VISIBLE_DEVICES" not in script
+    assert "--data-parallel-size 4" in _step(script, "vllm_model")
+    assert subprocess.run(["bash", "-n"], input=script, text=True).returncode == 0
 
 
 # ---------------------------------------------------------------------------

From 704663645534fe49ff4c6a0ed63307458e81fe63 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Thu, 1 Oct 2026 13:43:12 +0200
Subject: [PATCH 5/6] feat(orchestration): keep GPU packing and nodes: auto
 clear of vllm_pd

A prefill/decode tier owns its nodes, so a pool holding one is never
packed, and the vllm-router takes no GPUs. `nodes: auto` refuses a pool a
vllm_pd service uses: P/D pools need explicit node counts for now.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py         | 24 ++++++++++++++++---
 tests/unit_tests/test_slurm_script.py | 33 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index 736557d0958..f9ffc6a38be 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -741,7 +741,8 @@ def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
 
     A pool is planned only when a vLLM service on it shares it with another service that wants
     its own GPUs (a vLLM or a ray service). Each vLLM takes the lowest free indices in declaration
-    order; a ray service there gets the rest. Every other service keeps seeing every GPU.
+    order; a ray service there gets the rest. Every other service keeps seeing every GPU. A
+    prefill/decode tier owns its nodes, so a pool holding one is never planned.
     """
     compute = next(iter(config.compute.values()))
     pools = {name: pool for name, pool in compute.node_pools.items() if pool.nodes}
@@ -758,7 +759,7 @@ def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
             continue
         on_pool = {
             name: service
-            for name, service in config.services.items()
+            for name, service in config.deployed_services.items()
             if single_node
             or (
                 pool_name in (*service.node_pools, driver_pool)
@@ -766,7 +767,13 @@ def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
                 else service.node_pool == pool_name
             )
         }
-        claimants = {n: s for n, s in on_pool.items() if _takes_own_gpus(s)}
+        if any(isinstance(s, VllmPDTierConfig) for s in on_pool.values()):
+            continue
+        claimants = {
+            n: s
+            for n, s in on_pool.items()
+            if isinstance(s, (VllmServiceConfig, RayServiceConfig)) and _takes_own_gpus(s)
+        }
         vllms = {n: s for n, s in claimants.items() if isinstance(s, VllmServiceConfig)}
         if not vllms or len(claimants) < 2:
             continue
@@ -829,6 +836,17 @@ def _resolve_auto_pool(config: SubmitConfig, compute: SlurmComputeConfig, pool_n
             f"Node pool '{pool_name}' has nodes: auto but no gpus_per_node, so Gym cannot size it. "
             "Set gpus_per_node, or give the pool an explicit node count."
         )
+    pd_users = [
+        name
+        for name, service in config.services.items()
+        if isinstance(service, VllmPDServiceConfig)
+        and pool_name in (service.node_pool, service.prefill.node_pool, service.decode.node_pool)
+    ]
+    if pd_users:
+        raise ValueError(
+            f"Node pool '{pool_name}' has nodes: auto, but vllm_pd service {', '.join(map(repr, pd_users))} "
+            "uses it. A prefill/decode pool needs an explicit node count."
+        )
     sole_pool = len(compute.node_pools) == 1
     members = {
         name: service
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index f1663ce8d77..3493c306cad 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -22,6 +22,7 @@
 import pytest
 import yaml
 
+import nemo_gym.orchestration.executors.slurm_script as slurm_script_module
 from nemo_gym.orchestration.api import NodePool, RayServiceConfig, SubmitConfig, VllmPDTierConfig
 from nemo_gym.orchestration.executors.script_templates import (
     render_driver_entrypoint,
@@ -2953,3 +2954,35 @@ def test_auto_puts_services_with_clashing_manual_gpus_on_different_nodes(tmp_pat
     )
     assert config.services["a"].node_pool == "gpu-0"
     assert config.services["b"].node_pool == "gpu-1"
+
+
+# ---------------------------------------------------------------------------
+# GPU packing and nodes: auto beside vllm_pd
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("server_per_node", [False, True])
+def test_a_pd_config_renders_the_same_with_and_without_gpu_packing(tmp_path, monkeypatch, server_per_node):
+    # A tier owns its nodes: nothing beside it is packed, whatever else shares the pool.
+    services = _pd_services(server_per_node=server_per_node)
+    if server_per_node:
+        for tier in ("prefill", "decode"):
+            services["policy"][tier]["number_of_instances"] = 1
+    services["judge"] = _vllm(8100, "decode")
+    services["ray"] = {"type": "ray", "container": "img", "node_pools": ["prefill"]}
+    packed = _pd_script(tmp_path, services=services)
+    monkeypatch.setattr(slurm_script_module, "plan_gpus", lambda config: {})
+    assert _pd_script(tmp_path, services=services) == packed
+    assert "CUDA_VISIBLE_DEVICES" not in packed
+
+
+def test_auto_refuses_a_pool_a_pd_tier_uses(tmp_path):
+    pools = {**_PD_POOLS, "decode": {**_PD_POOLS["decode"], "nodes": "auto"}}
+    with pytest.raises(ValueError, match="vllm_pd service 'policy' uses it"):
+        _pd_config(tmp_path, pools=pools)
+
+
+def test_a_pd_router_beside_a_lone_vllm_takes_no_gpus(tmp_path):
+    services = {**_pd_services(node_pool="front"), "scorer": _vllm(8200, "front")}
+    pools = {"front": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}, **_PD_POOLS}
+    assert "CUDA_VISIBLE_DEVICES" not in _pd_script(tmp_path, services=services, pools=pools)

From 62ecf0318f2576edcb8333c184fee673294a05a6 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Thu, 1 Oct 2026 13:43:46 +0200
Subject: [PATCH 6/6] test(orchestration): probe a service nodes: auto moves on
 the node it lands on

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 tests/unit_tests/test_slurm_script.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 3493c306cad..38d041d48c6 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -2986,3 +2986,13 @@ def test_a_pd_router_beside_a_lone_vllm_takes_no_gpus(tmp_path):
     services = {**_pd_services(node_pool="front"), "scorer": _vllm(8200, "front")}
     pools = {"front": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}, **_PD_POOLS}
     assert "CUDA_VISIBLE_DEVICES" not in _pd_script(tmp_path, services=services, pools=pools)
+
+
+def test_auto_probes_a_service_it_moves_to_a_second_node_there(tmp_path):
+    _, script = _auto_render(
+        tmp_path,
+        {"policy": _vllm(8000, None, tensor_parallel_size=4), "judge": _vllm(8001, None)},
+        policy="policy",
+    )
+    assert "Waiting for policy at http://localhost:8000" in script
+    assert "Waiting for judge at http://${gym_nodes[1]}:8001" in script