From fa7531dc5f3b4574180b4da0381b62c5eacd0dae Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Wed, 23 Sep 2026 10:07:29 +0200
Subject: [PATCH 1/8] feat(orchestration): let a service claim its own nodes

A submit config could already declare several services, but every one of them
rendered as a bare `srun --overlap`, so they all landed on the same node and
competed for its GPUs. A scorer or judge that cannot share a GPU with the
policy had nowhere to go.

A service may now name a `node_pool`. Pools take contiguous slices of the
allocation in declaration order, and a pinned service is launched with
`srun --relative=<offset>`, which starts the step at that slice.

Sizing follows placement: a service pinned to a one-node pool inside a
ten-node job is a single-node deployment, so it is no longer built as a
multi-node Ray one, and its GPU footprint is checked against its own pool
instead of the allocation total.

This also fixes multiple pools being silently useless. The directives were
emitted per pool, so each --nodes/--partition overwrote the last and only the
final pool survived, while the executor sized itself on the sum. There is now
one directive block for the allocation, and pools that disagree on a
whole-allocation directive are refused instead of resolved by ordering.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py                 |  51 +++++---
 .../orchestration/executors/slurm_script.py   | 123 +++++++++++++++---
 tests/unit_tests/test_slurm_script.py         | 113 +++++++++++++++-
 3 files changed, 248 insertions(+), 39 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index ef653842ac3..b4306cbb93f 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -80,6 +80,12 @@ class BaseServiceConfig(_StrictModel):
     container: str
     # Resolved to the sole compute resource name at validation time when not set.
     placement: str | None = None
+    # Name of a node pool in the placed compute's `node_pools`. Pins this service to
+    # that pool's slice of the allocation instead of letting it land wherever srun
+    # starts, which is how two services get nodes of their own -- a scorer or judge
+    # 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
     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).
@@ -304,12 +310,8 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
         total_nodes = (
             sum(p.nodes for p in compute.node_pools.values()) if isinstance(compute, SlurmComputeConfig) else 1
         )
-        is_multi_node = total_nodes > 1
-        gpus_per_node_values = (
-            [p.gpus_per_node for p in compute.node_pools.values() if p.gpus_per_node is not None]
-            if isinstance(compute, SlurmComputeConfig)
-            else []
-        )
+
+        pool_names = set(compute.node_pools) if isinstance(compute, SlurmComputeConfig) else set()
 
         for service_name, service in self.services.items():
             if service.placement is None:
@@ -320,25 +322,42 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
                     f"({', '.join(sorted(compute_names))})."
                 )
 
+            if service.node_pool is not None and service.node_pool not in pool_names:
+                raise ValueError(
+                    f"Service '{service_name}' node_pool '{service.node_pool}' does not match any node pool of "
+                    f"compute '{service.placement}' ({', '.join(sorted(pool_names)) or 'none declared'})."
+                )
+
             if not isinstance(service, VllmServiceConfig):
                 continue
 
-            is_ray_serve = effective_ray_serve(service, total_nodes, gpus_per_node_values)
+            # A pinned service is sized against its own pool, not the whole job: one node of a
+            # ten-node allocation is a single-node deployment with that pool's GPUs, and judging
+            # it by the allocation total both mis-builds the command and mis-reports idle GPUs.
+            service_pools = (
+                {service.node_pool: compute.node_pools[service.node_pool]}
+                if service.node_pool is not None and isinstance(compute, SlurmComputeConfig)
+                else (compute.node_pools if isinstance(compute, SlurmComputeConfig) else {})
+            )
+            service_nodes = sum(p.nodes for p in service_pools.values()) or total_nodes
+            service_gpus = [p.gpus_per_node for p in service_pools.values() if p.gpus_per_node is not None]
+
+            is_ray_serve = effective_ray_serve(service, service_nodes, service_gpus)
 
             if (
-                is_multi_node
+                service_nodes > 1
                 and service.number_of_instances > 1
-                and service.number_of_instances % total_nodes != 0
+                and service.number_of_instances % service_nodes != 0
                 and not is_ray_serve
             ):
                 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 "
+                    f"be evenly divisible by the number of nodes ({service_nodes}) for multi-node data-parallel "
                     "deployment - each node hosts an equal share of the data-parallel replicas."
                 )
 
             self._validate_vllm_gpu_footprint(
-                service_name, service, total_nodes, compute, gpus_per_node_values, is_ray_serve
+                service_name, service, service_nodes, service_pools, service_gpus, is_ray_serve
             )
 
         if self.driver.policy_model is not None:
@@ -370,7 +389,7 @@ def _validate_vllm_gpu_footprint(
         service_name: str,
         service: "VllmServiceConfig",
         total_nodes: int,
-        compute: "SlurmComputeConfig",
+        node_pools: dict[str, "NodePool"],
         gpus_per_node_values: list[int],
         is_ray_serve: bool,
     ) -> None:
@@ -383,9 +402,7 @@ def _validate_vllm_gpu_footprint(
         if total_nodes > 1 and is_ray_serve:
             # Ray Serve's placement-group scheduler packs the aggregate footprint across the cluster.
             gpus_needed = tp_pp * service.number_of_instances
-            gpus_available = sum(
-                pool.nodes * pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node
-            )
+            gpus_available = sum(pool.nodes * pool.gpus_per_node for pool in node_pools.values() if pool.gpus_per_node)
             footprint = (
                 f"tensor_parallel_size={service.tensor_parallel_size} x "
                 f"pipeline_parallel_size={service.pipeline_parallel_size} x "
@@ -408,9 +425,7 @@ def _validate_vllm_gpu_footprint(
         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
-            )
+            gpus_available = sum(pool.nodes * pool.gpus_per_node for pool in 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}"
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index 399b899f539..7cebc623ce0 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -78,21 +78,54 @@ def _render_directives(compute: SlurmComputeConfig, remote_bench_dir: Path, benc
     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))
+    lines.extend(_render_pool_directives(compute.node_pools))
     return "\n".join(lines)
 
 
-def _render_pool_directives(pool_name: str, pool: NodePool) -> list[str]:
+def _pool_directive(pools: dict[str, NodePool], attribute: str) -> Any:
+    """The one value every pool agrees on for `attribute`.
+
+    A plain sbatch job takes a single --partition/--ntasks-per-node/--gpus-per-node
+    for the whole allocation, so pools that disagree cannot both be honoured. Slurm
+    would silently apply whichever directive came last; say so instead.
+    """
+    values = {getattr(pool, attribute) for pool in pools.values()}
+    if len(values) > 1:
+        named = ", ".join(f"{name}={getattr(pool, attribute)!r}" for name, pool in pools.items())
+        raise ValueError(
+            f"Node pools disagree on {attribute} ({named}). One Slurm job takes a single value for the whole "
+            "allocation; split the run or make the pools agree."
+        )
+    return next(iter(values))
+
+
+def _render_pool_directives(pools: dict[str, NodePool]) -> list[str]:
+    """One set of directives for the whole allocation, not one per pool.
+
+    Pools divide an allocation between services (see _pool_offsets); they are not
+    separate Slurm requests. Emitting --nodes per pool made every pool but the last
+    a no-op, so a two-pool job asked for one pool's nodes while the rest of the
+    executor sized itself on the sum.
+    """
+    if not pools:
+        return []
     lines = [
-        f"#SBATCH --partition={pool.partition}  # pool: {pool_name}",
-        f"#SBATCH --nodes={pool.nodes}",
-        f"#SBATCH --ntasks-per-node={pool.ntasks_per_node}",
+        f"#SBATCH --partition={_pool_directive(pools, 'partition')}",
+        f"#SBATCH --nodes={sum(pool.nodes for pool in pools.values())}",
+        f"#SBATCH --ntasks-per-node={_pool_directive(pools, '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}")
+    gpus_per_node = _pool_directive(pools, "gpus_per_node")
+    if gpus_per_node is not None:
+        lines.append(f"#SBATCH --gpus-per-node={gpus_per_node}")
+    extra_args: dict[str, str] = {}
+    for name, pool in pools.items():
+        for key, val in pool.extra_args.items():
+            if extra_args.setdefault(key, val) != val:
+                raise ValueError(
+                    f"Node pool {name!r} sets extra_args[{key!r}]={val!r}, which conflicts with another pool's "
+                    f"{extra_args[key]!r}. #SBATCH directives apply to the whole allocation."
+                )
+    lines.extend(f"#SBATCH --{key}={val}" for key, val in extra_args.items())
     return lines
 
 
@@ -133,6 +166,7 @@ def _render_service_command(
     pre_command: str = "",
     workdir: str | None = None,
     single_node: bool = False,
+    relative: int | None = None,
 ) -> str:
     var = bash_var(name)
     env_prefix = _resolve_env(env) if env else ""
@@ -141,6 +175,10 @@ def _render_service_command(
         node_flags = " --nodes=1 --ntasks=1"
     else:
         node_flags = f" --nodes={nodes} --ntasks={ntasks}" if (nodes is not None and nodes > 1) else ""
+    # --relative=N starts the step at node N of the allocation. With --overlap it is
+    # what gives a pinned service nodes of its own while other steps run elsewhere.
+    if relative is not None:
+        node_flags = f" --relative={relative} --nodes={nodes} --ntasks={ntasks}"
     mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else ""
     workdir_flag = f" --container-workdir={shlex.quote(workdir)}" if workdir else ""
     if pre_command:
@@ -403,6 +441,56 @@ def _render_collector_shutdown(config: SubmitConfig, remote_bench_dir: Path) ->
     )
 
 
+def _pool_offsets(compute: SlurmComputeConfig) -> dict[str, tuple[int, int]]:
+    """Each pool's (first node index, node count) within the allocation.
+
+    Pools are laid out contiguously in declaration order, which is the order the
+    single #SBATCH --nodes total is built from, so pool i owns the nodes after
+    every pool before it.
+    """
+    offsets: dict[str, tuple[int, int]] = {}
+    start = 0
+    for name, pool in compute.node_pools.items():
+        offsets[name] = (start, pool.nodes)
+        start += pool.nodes
+    return offsets
+
+
+def _service_nodes(
+    service: VllmServiceConfig | RayServiceConfig, compute: SlurmComputeConfig, total_nodes: int
+) -> int:
+    """How many nodes this service actually runs on.
+
+    A service pinned to a pool sees only that pool, so a single-node pool inside a
+    ten-node job is a single-node deployment and must not be built as a multi-node
+    Ray one.
+    """
+    if service.node_pool is None:
+        return total_nodes
+    return compute.node_pools[service.node_pool].nodes
+
+
+def _srun_nodes(
+    service: VllmServiceConfig | RayServiceConfig, compute: SlurmComputeConfig, total_nodes: int
+) -> int | None:
+    nodes = _service_nodes(service, compute, total_nodes)
+    if service.node_pool is not None:
+        return nodes
+    return total_nodes if _vllm_spans_multiple_nodes(service, total_nodes) else None
+
+
+def _srun_ntasks(
+    service: VllmServiceConfig | RayServiceConfig,
+    compute: SlurmComputeConfig,
+    total_nodes: int,
+    total_ntasks: int,
+) -> int | None:
+    if service.node_pool is not None:
+        pool = compute.node_pools[service.node_pool]
+        return pool.nodes * pool.ntasks_per_node
+    return total_ntasks if _vllm_spans_multiple_nodes(service, total_nodes) else None
+
+
 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())
@@ -447,21 +535,24 @@ def build_sbatch_script(
 
     observed = otel_active(config)
 
+    offsets = _pool_offsets(compute)
     service_commands = "\n\n".join(
         ([_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node)] if observed else [])
         + [
             _render_service_command(
                 name,
                 service.container,
-                _build_service_command(service, total_nodes, gpus_per_node_values),
+                _build_service_command(service, _service_nodes(service, compute, total_nodes), gpus_per_node_values),
                 service.env or None,
                 service.mounts or None,
-                # Only services that actually span multiple nodes need the whole allocation's --nodes/
-                # --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,
+                # 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
+                # regardless of how many nodes the overall job spans). A pinned service always gets
+                # them, since --relative is meaningless without a node count.
+                nodes=_srun_nodes(service, compute, total_nodes),
+                ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks),
                 pre_command=service.pre_command,
+                relative=offsets[service.node_pool][0] if service.node_pool else None,
             )
             for name, service in config.services.items()
         ]
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index ac78c03f865..328b5e25d4d 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -21,7 +21,7 @@
 
 import pytest
 
-from nemo_gym.orchestration.api import SubmitConfig
+from nemo_gym.orchestration.api import NodePool, SubmitConfig
 from nemo_gym.orchestration.executors.script_templates import (
     render_driver_entrypoint,
     render_gym_cmd,
@@ -87,24 +87,46 @@ def test_deeply_nested():
 
 
 def test_render_pool_directives_basic(pool):
-    lines = _render_pool_directives("main", pool)
-    assert "#SBATCH --partition=batch  # pool: main" in lines
+    lines = _render_pool_directives({"main": pool})
+    assert "#SBATCH --partition=batch" 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)
+    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)
+    lines = _render_pool_directives({"main": pool})
     assert "#SBATCH --gres=shard:8" in lines
 
 
+def test_pools_ask_for_the_sum_of_their_nodes_once(pool):
+    # One #SBATCH --nodes for the allocation, not one per pool. Emitting them per
+    # pool made every pool but the last a no-op while the rest of the executor
+    # sized itself on the sum, so a two-pool job asked for one pool's nodes.
+    aux = NodePool(partition="batch", nodes=3, ntasks_per_node=4)
+    lines = _render_pool_directives({"main": pool, "aux": aux})
+    assert [line for line in lines if line.startswith("#SBATCH --nodes")] == ["#SBATCH --nodes=4"]
+
+
+def test_pools_that_disagree_on_a_whole_allocation_directive_are_refused(pool):
+    other = NodePool(partition="interactive", nodes=1, ntasks_per_node=4)
+    with pytest.raises(ValueError, match="disagree on partition"):
+        _render_pool_directives({"main": pool, "other": other})
+
+
+def test_pools_that_disagree_on_extra_args_are_refused(pool):
+    pool.extra_args["gres"] = "shard:8"
+    other = NodePool(partition="batch", nodes=1, ntasks_per_node=4, extra_args={"gres": "shard:4"})
+    with pytest.raises(ValueError, match="conflicts with another pool"):
+        _render_pool_directives({"main": pool, "other": other})
+
+
 # ---------------------------------------------------------------------------
 # _render_directives
 # ---------------------------------------------------------------------------
@@ -1570,3 +1592,84 @@ def test_gym_install_runs_from_the_install_root():
     assert entrypoint.index('cd "$GYM_SRC/gym"') < entrypoint.index('exec "$@"')
     # The only `cd` is into the clone -- nothing else may move cwd.
     assert entrypoint.count("cd ") == entrypoint.count('cd "$GYM_SRC/gym"')
+
+
+# ---------------------------------------------------------------------------
+# node_pool placement
+# ---------------------------------------------------------------------------
+
+
+def _placement_config(tmp_path, services, pools):
+    return SubmitConfig.model_validate(
+        {
+            "services": services,
+            "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": pools}},
+            "driver": {"container": "gym:latest", "benchmarks": {"b": {"run": {}}}},
+            "job": {"output_path": str(tmp_path / "jobs")},
+        }
+    )
+
+
+def _vllm(port, pool, **extra):
+    return {
+        "type": "vllm",
+        "container": "img",
+        "model": "/ckpt",
+        "port": port,
+        "node_pool": pool,
+        **extra,
+    }
+
+
+_TWO_POOLS = {
+    "gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4},
+    "aux": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4},
+}
+
+
+def _render(tmp_path, services, pools=None):
+    config = _placement_config(tmp_path, services, pools or _TWO_POOLS)
+    return build_sbatch_script(
+        config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "jobs" / "b"
+    )
+
+
+def test_pinned_services_take_contiguous_node_ranges(tmp_path):
+    # --relative is what gives a service nodes of its own; without it both steps
+    # start at node 0 and the second one shares the first one's GPUs.
+    script = _render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, "gpu", tensor_parallel_size=4),
+            "scorer": _vllm(8001, "aux", tensor_parallel_size=4),
+        },
+    )
+    assert "--relative=0 --nodes=1 --ntasks=1" in script.split("# service: policy")[1]
+    assert "--relative=1 --nodes=1 --ntasks=1" in script.split("# service: scorer")[1]
+
+
+def test_a_pinned_single_node_service_is_not_built_as_multi_node(tmp_path):
+    # The job spans two nodes, but each service owns one. Sizing them by the
+    # allocation total would build both as multi-node Ray deployments.
+    script = _render(
+        tmp_path,
+        {
+            "policy": _vllm(8000, "gpu", tensor_parallel_size=4),
+            "scorer": _vllm(8001, "aux", tensor_parallel_size=4),
+        },
+    )
+    assert "ray start" not in script
+    assert script.count("vllm serve") == 2
+
+
+def test_an_unpinned_service_keeps_the_whole_allocation(tmp_path):
+    script = _render(
+        tmp_path,
+        {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 8}},
+    )
+    assert "--relative=" not in script
+
+
+def test_an_unknown_node_pool_is_named(tmp_path):
+    with pytest.raises(ValueError, match="node_pool 'nope' does not match any node pool"):
+        _placement_config(tmp_path, {"policy": _vllm(8000, "nope", tensor_parallel_size=4)}, _TWO_POOLS)

From 92baa09078ee66e0a3221fa2c5047044ba882f46 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Wed, 23 Sep 2026 10:09:21 +0200
Subject: [PATCH 2/8] feat(orchestration): make the ray service usable

`_build_ray_command` returned the fixed string "ray start --head". That starts
a cluster nobody can join, on whichever node srun picked, and returns
immediately -- ray daemonises, so the srun step exited and Slurm tore the
service down again.

A ray service now takes a mode. A head starts a cluster on a chosen port; a
worker joins one at `address`. Either may advertise custom Ray resources
(`{"extra_gpu": 4}`), which is how a benchmark asks for a node whose GPUs it
places actors on itself rather than letting Ray count them. --block keeps the
step alive for as long as the node should be up.

With node_pool from the previous commit, this is enough to put such a worker
on a node of its own.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py                 | 22 +++++++++++
 .../orchestration/executors/slurm_script.py   | 23 ++++++++++-
 tests/unit_tests/test_slurm_script.py         | 39 ++++++++++++++++++-
 3 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index b4306cbb93f..a29a085547f 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -158,6 +158,28 @@ def effective_ray_serve(service: "VllmServiceConfig", total_nodes: int, gpus_per
 
 class RayServiceConfig(BaseServiceConfig):
     type: Literal["ray"]
+    # "head" starts a cluster; "worker" joins the one at `address`. A worker is how a
+    # second node joins the driver's Ray cluster and offers its GPUs to actors the
+    # benchmark schedules (e.g. a scorer that runs off the policy's node).
+    mode: Literal["head", "worker"] = "head"
+    # Required for mode="worker": the head's host:port.
+    address: str | None = None
+    # Head only; ignored by a worker, which takes the port from `address`.
+    port: int = 6379
+    # Custom Ray resources this node advertises, e.g. {"extra_gpu": 4}. A benchmark
+    # asks for these by name rather than by num_gpus when it manages device placement
+    # itself.
+    resources: dict[str, float] = {}
+    num_cpus: int | None = None
+    num_gpus: int | None = None
+
+    @model_validator(mode="after")
+    def _validate_mode(self) -> "RayServiceConfig":
+        if self.mode == "worker" and not self.address:
+            raise ValueError("A ray service with mode='worker' needs `address` set to the head's host:port.")
+        if self.mode == "head" and self.address:
+            raise ValueError("A ray service with mode='head' starts its own cluster; remove `address`.")
+        return self
 
 
 # Discriminated union keyed on `type`; Pydantic rejects unknown type values at parse time.
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index 7cebc623ce0..fdf9077c456 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 json
 import re
 import shlex
 from pathlib import Path
@@ -353,8 +354,26 @@ def _build_vllm_ray_serve_command(
     )
 
 
-def _build_ray_command(_service: RayServiceConfig) -> str:
-    return "ray start --head"
+def _build_ray_command(service: RayServiceConfig) -> str:
+    # --block keeps the srun step alive. `ray start` daemonises and returns, so
+    # without it the step exits the moment the node is up and Slurm tears the
+    # service down again.
+    cmd = "ray start --block"
+    if service.mode == "head":
+        cmd += f" --head --port {service.port}"
+    else:
+        cmd += f" --address {shlex.quote(str(service.address))}"
+    if service.num_cpus is not None:
+        cmd += f" --num-cpus {service.num_cpus}"
+    if service.num_gpus is not None:
+        cmd += f" --num-gpus {service.num_gpus}"
+    if service.resources:
+        # Ray takes fractional custom resources, so the field is float-typed, but a
+        # whole number is written as one: {"extra_gpu": 4}, not 4.0, so the rendered
+        # command reads the way the config does.
+        resources = {k: int(v) if v.is_integer() else v for k, v in service.resources.items()}
+        cmd += " --resources=" + shlex.quote(json.dumps(resources, sort_keys=True))
+    return cmd
 
 
 _BUILDERS = {
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 328b5e25d4d..9d36b216589 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -21,13 +21,14 @@
 
 import pytest
 
-from nemo_gym.orchestration.api import NodePool, SubmitConfig
+from nemo_gym.orchestration.api import NodePool, RayServiceConfig, SubmitConfig
 from nemo_gym.orchestration.executors.script_templates import (
     render_driver_entrypoint,
     render_gym_cmd,
 )
 from nemo_gym.orchestration.executors.slurm_script import (
     _RAY_SERVE_GATEWAY_SOURCE_PATH,
+    _build_ray_command,
     _build_service_command,
     _build_vllm_command,
     _build_vllm_multi_instance_multi_node_command,
@@ -1673,3 +1674,39 @@ def test_an_unpinned_service_keeps_the_whole_allocation(tmp_path):
 def test_an_unknown_node_pool_is_named(tmp_path):
     with pytest.raises(ValueError, match="node_pool 'nope' does not match any node pool"):
         _placement_config(tmp_path, {"policy": _vllm(8000, "nope", tensor_parallel_size=4)}, _TWO_POOLS)
+
+
+# ---------------------------------------------------------------------------
+# ray services
+# ---------------------------------------------------------------------------
+
+
+def test_a_ray_head_blocks_so_slurm_keeps_the_step_alive():
+    # `ray start` daemonises and returns; without --block the srun step exits the
+    # moment the node is up and the service is torn down again.
+    command = _build_ray_command(RayServiceConfig(type="ray", container="img"))
+    assert command == "ray start --block --head --port 6379"
+
+
+def test_a_ray_worker_joins_the_head_and_advertises_its_resources():
+    command = _build_ray_command(
+        RayServiceConfig(
+            type="ray",
+            container="img",
+            mode="worker",
+            address="10.0.0.1:6379",
+            num_gpus=4,
+            resources={"extra_gpu": 4},
+        )
+    )
+    assert command == "ray start --block --address 10.0.0.1:6379 --num-gpus 4 --resources='{\"extra_gpu\": 4}'"
+
+
+def test_a_ray_worker_without_an_address_is_refused():
+    with pytest.raises(ValueError, match="needs `address`"):
+        RayServiceConfig(type="ray", container="img", mode="worker")
+
+
+def test_a_ray_head_with_an_address_is_refused():
+    with pytest.raises(ValueError, match="starts its own cluster"):
+        RayServiceConfig(type="ray", container="img", address="10.0.0.1:6379")

From ee676cc447a15bdf494bf4107a71f413fb2a6a45 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Wed, 23 Sep 2026 16:42:32 +0200
Subject: [PATCH 3/8] feat(orchestration): let a ray worker join a head service
 by name

A worker could only join a fixed host:port, but the head's host is not known
until Slurm places the job. A worker may now name a ray head service in the
same config (`head: ray_head`); the job script exports that head's address as
GYM_RAY_ADDRESS_<SERVICE>, resolved from its pool's first node, and the worker
joins through it. The driver can read the same variable to connect.

Also adds `extra_args` to ray services for flags the schema does not model,
such as fixed ports that avoid another Ray instance on the same node.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py                 | 21 ++++++--
 .../orchestration/executors/slurm_script.py   | 42 +++++++++++++--
 tests/unit_tests/test_slurm_script.py         | 54 ++++++++++++++++++-
 3 files changed, 107 insertions(+), 10 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index a29a085547f..7c831a8fa6d 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -162,8 +162,10 @@ class RayServiceConfig(BaseServiceConfig):
     # second node joins the driver's Ray cluster and offers its GPUs to actors the
     # benchmark schedules (e.g. a scorer that runs off the policy's node).
     mode: Literal["head", "worker"] = "head"
-    # Required for mode="worker": the head's host:port.
+    # For mode="worker": the head's host:port, or `head`, the name of a ray head
+    # service in this config whose address is only known once the job is placed.
     address: str | None = None
+    head: str | None = None
     # Head only; ignored by a worker, which takes the port from `address`.
     port: int = 6379
     # Custom Ray resources this node advertises, e.g. {"extra_gpu": 4}. A benchmark
@@ -172,13 +174,15 @@ class RayServiceConfig(BaseServiceConfig):
     resources: dict[str, float] = {}
     num_cpus: int | None = None
     num_gpus: int | None = None
+    # Raw extra flags appended verbatim to `ray start` (e.g. fixed ports).
+    extra_args: str = ""
 
     @model_validator(mode="after")
     def _validate_mode(self) -> "RayServiceConfig":
-        if self.mode == "worker" and not self.address:
-            raise ValueError("A ray service with mode='worker' needs `address` set to the head's host:port.")
-        if self.mode == "head" and self.address:
-            raise ValueError("A ray service with mode='head' starts its own cluster; remove `address`.")
+        if self.mode == "worker" and bool(self.address) == bool(self.head):
+            raise ValueError("A ray service with mode='worker' needs exactly one of `address` or `head`.")
+        if self.mode == "head" and (self.address or self.head):
+            raise ValueError("A ray service with mode='head' starts its own cluster; remove `address`/`head`.")
         return self
 
 
@@ -344,6 +348,13 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
                     f"({', '.join(sorted(compute_names))})."
                 )
 
+            if isinstance(service, RayServiceConfig) and service.head is not None:
+                head = self.services.get(service.head)
+                if not isinstance(head, RayServiceConfig) or head.mode != "head":
+                    raise ValueError(
+                        f"Service '{service_name}' head '{service.head}' must name a ray service with mode='head'."
+                    )
+
             if service.node_pool is not None and service.node_pool not in pool_names:
                 raise ValueError(
                     f"Service '{service_name}' node_pool '{service.node_pool}' does not match any node pool of "
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index fdf9077c456..75bccf6f2a5 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -361,6 +361,10 @@ def _build_ray_command(service: RayServiceConfig) -> str:
     cmd = "ray start --block"
     if service.mode == "head":
         cmd += f" --head --port {service.port}"
+    elif service.head is not None:
+        # Double-quoted, not shlex-quoted: the head's address is only known once
+        # Slurm places the job, and build_sbatch_script exports it under this name.
+        cmd += f' --address "${ray_head_address_var(service.head)}"'
     else:
         cmd += f" --address {shlex.quote(str(service.address))}"
     if service.num_cpus is not None:
@@ -373,9 +377,34 @@ def _build_ray_command(service: RayServiceConfig) -> str:
         # command reads the way the config does.
         resources = {k: int(v) if v.is_integer() else v for k, v in service.resources.items()}
         cmd += " --resources=" + shlex.quote(json.dumps(resources, sort_keys=True))
+    if service.extra_args:
+        cmd += " " + service.extra_args
     return cmd
 
 
+def ray_head_address_var(head_service: str) -> str:
+    """The env var build_sbatch_script exports with a ray head service's host:port."""
+    return f"GYM_RAY_ADDRESS_{bash_var(head_service)}"
+
+
+def _render_ray_head_addresses(config: SubmitConfig, compute: SlurmComputeConfig) -> str:
+    """Export every ray head service's address, resolved from the node Slurm gave it.
+
+    A worker or the driver in another step can then join it without knowing,
+    when the config is written, which host the job will land on.
+    """
+    heads = {n: s for n, s in config.services.items() if isinstance(s, RayServiceConfig) and s.mode == "head"}
+    if not heads:
+        return ""
+    offsets = _pool_offsets(compute)
+    lines = ['gym_ray_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))']
+    for name, service in heads.items():
+        index = offsets[service.node_pool][0] if service.node_pool else 0
+        ip = f"$(getent hosts ${{gym_ray_nodes[{index}]}} | awk '{{print $1}}')"
+        lines.append(f'export {ray_head_address_var(name)}="{ip}:{service.port}"')
+    return "\n".join(lines)
+
+
 _BUILDERS = {
     VllmServiceConfig: _build_vllm_command,
     RayServiceConfig: _build_ray_command,
@@ -546,10 +575,15 @@ def build_sbatch_script(
         pool.gpus_per_node for pool in compute.node_pools.values() if pool.gpus_per_node is not None
     ]
 
-    ray_prelude = (
-        render_ray_prelude()
-        if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values())
-        else ""
+    ray_prelude = "\n".join(
+        block
+        for block in (
+            render_ray_prelude()
+            if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values())
+            else "",
+            _render_ray_head_addresses(config, compute),
+        )
+        if block
     )
 
     observed = otel_active(config)
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 9d36b216589..90ac4700b7a 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -1703,10 +1703,62 @@ def test_a_ray_worker_joins_the_head_and_advertises_its_resources():
 
 
 def test_a_ray_worker_without_an_address_is_refused():
-    with pytest.raises(ValueError, match="needs `address`"):
+    with pytest.raises(ValueError, match="exactly one of `address` or `head`"):
         RayServiceConfig(type="ray", container="img", mode="worker")
 
 
 def test_a_ray_head_with_an_address_is_refused():
     with pytest.raises(ValueError, match="starts its own cluster"):
         RayServiceConfig(type="ray", container="img", address="10.0.0.1:6379")
+
+
+def _head_worker_config(tmp_path, worker):
+    return _placement_config(
+        tmp_path,
+        {
+            "ray_head": {"type": "ray", "container": "img", "mode": "head", "port": 6380, "node_pool": "gpu"},
+            "scorer": {"type": "ray", "container": "img", "mode": "worker", "node_pool": "aux", **worker},
+        },
+        _TWO_POOLS,
+    )
+
+
+def test_a_worker_joins_a_head_service_by_its_placed_address(tmp_path):
+    # The head's host is only known once Slurm places the job, so the script
+    # exports it and the worker reads it rather than a hard-coded address.
+    config = _head_worker_config(tmp_path, {"head": "ray_head", "resources": {"extra_gpu": 4}})
+    script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
+    assert 'export GYM_RAY_ADDRESS_RAY_HEAD="$(getent hosts ${gym_ray_nodes[0]}' in script
+    assert ':6380"' in script
+    assert '--address "$GYM_RAY_ADDRESS_RAY_HEAD"' in script.split("# service: scorer")[1]
+
+
+def test_a_head_on_a_later_pool_exports_that_pools_node(tmp_path):
+    config = _placement_config(
+        tmp_path,
+        {"ray_head": {"type": "ray", "container": "img", "mode": "head", "node_pool": "aux"}},
+        _TWO_POOLS,
+    )
+    script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
+    assert "${gym_ray_nodes[1]}" in script
+
+
+def test_ray_extra_args_are_appended():
+    command = _build_ray_command(
+        RayServiceConfig(type="ray", container="img", extra_args="--include-dashboard=false --node-manager-port=8366")
+    )
+    assert command.endswith("--include-dashboard=false --node-manager-port=8366")
+
+
+def test_a_worker_naming_something_other_than_a_head_is_refused(tmp_path):
+    with pytest.raises(ValueError, match="must name a ray service with mode='head'"):
+        _placement_config(
+            tmp_path,
+            {"scorer": {"type": "ray", "container": "img", "mode": "worker", "head": "missing", "node_pool": "aux"}},
+            _TWO_POOLS,
+        )
+
+
+def test_a_worker_with_both_address_and_head_is_refused():
+    with pytest.raises(ValueError, match="exactly one of `address` or `head`"):
+        RayServiceConfig(type="ray", container="img", mode="worker", address="h:1", head="ray_head")

From 4552a418b8651b34688cca3090d1f3c64ac5df83 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Wed, 23 Sep 2026 17:35:56 +0200
Subject: [PATCH 4/8] fix(orchestration): pin a pooled service with --nodelist,
 not --relative

--relative is only where Slurm starts looking. On a real two-node job a ray
head with --relative=0 landed on node 1, next to the worker, because node 0's
resources were held by the policy step; it then tried to reach node 0's IP and
never came up.

The job script now exports each pool's hosts (GYM_POOL_<NAME>_NODES) from the
allocation's node list, and a pinned service passes them to --nodelist. The
ray head address reuses the same host list.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 .../orchestration/executors/slurm_script.py   | 40 ++++++++++++++-----
 tests/unit_tests/test_slurm_script.py         | 15 ++++---
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index 75bccf6f2a5..5b3dc4bc353 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -167,7 +167,7 @@ def _render_service_command(
     pre_command: str = "",
     workdir: str | None = None,
     single_node: bool = False,
-    relative: int | None = None,
+    nodelist: str | None = None,
 ) -> str:
     var = bash_var(name)
     env_prefix = _resolve_env(env) if env else ""
@@ -176,10 +176,10 @@ def _render_service_command(
         node_flags = " --nodes=1 --ntasks=1"
     else:
         node_flags = f" --nodes={nodes} --ntasks={ntasks}" if (nodes is not None and nodes > 1) else ""
-    # --relative=N starts the step at node N of the allocation. With --overlap it is
-    # what gives a pinned service nodes of its own while other steps run elsewhere.
-    if relative is not None:
-        node_flags = f" --relative={relative} --nodes={nodes} --ntasks={ntasks}"
+    # --nodelist names the exact hosts. --relative is only a starting point Slurm may
+    # move off when that node's resources are taken, landing a service on the wrong node.
+    if nodelist is not None:
+        node_flags = f' --nodelist="${nodelist}" --nodes={nodes} --ntasks={ntasks}'
     mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else ""
     workdir_flag = f" --container-workdir={shlex.quote(workdir)}" if workdir else ""
     if pre_command:
@@ -382,6 +382,25 @@ def _build_ray_command(service: RayServiceConfig) -> str:
     return cmd
 
 
+_NODE_ARRAY = 'gym_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))'
+
+
+def pool_nodes_var(pool: str) -> str:
+    """The env var build_sbatch_script exports with a node pool's comma-separated hosts."""
+    return f"GYM_POOL_{bash_var(pool)}_NODES"
+
+
+def _render_pool_nodes(config: SubmitConfig, compute: SlurmComputeConfig) -> str:
+    """Export each node pool's hosts, so a pinned service can name them in --nodelist."""
+    if not any(s.node_pool for s in config.services.values()):
+        return ""
+    lines = [_NODE_ARRAY]
+    for name, (start, count) in _pool_offsets(compute).items():
+        if count:
+            lines.append(f'export {pool_nodes_var(name)}="$(IFS=,; echo "${{gym_nodes[*]:{start}:{count}}}")"')
+    return "\n".join(lines)
+
+
 def ray_head_address_var(head_service: str) -> str:
     """The env var build_sbatch_script exports with a ray head service's host:port."""
     return f"GYM_RAY_ADDRESS_{bash_var(head_service)}"
@@ -397,10 +416,11 @@ def _render_ray_head_addresses(config: SubmitConfig, compute: SlurmComputeConfig
     if not heads:
         return ""
     offsets = _pool_offsets(compute)
-    lines = ['gym_ray_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))']
+    # _render_pool_nodes already declared the host array when any service is pinned.
+    lines = [] if any(s.node_pool for s in config.services.values()) else [_NODE_ARRAY]
     for name, service in heads.items():
         index = offsets[service.node_pool][0] if service.node_pool else 0
-        ip = f"$(getent hosts ${{gym_ray_nodes[{index}]}} | awk '{{print $1}}')"
+        ip = f"$(getent hosts ${{gym_nodes[{index}]}} | awk '{{print $1}}')"
         lines.append(f'export {ray_head_address_var(name)}="{ip}:{service.port}"')
     return "\n".join(lines)
 
@@ -581,6 +601,7 @@ def build_sbatch_script(
             render_ray_prelude()
             if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values())
             else "",
+            _render_pool_nodes(config, compute),
             _render_ray_head_addresses(config, compute),
         )
         if block
@@ -588,7 +609,6 @@ def build_sbatch_script(
 
     observed = otel_active(config)
 
-    offsets = _pool_offsets(compute)
     service_commands = "\n\n".join(
         ([_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node)] if observed else [])
         + [
@@ -601,11 +621,11 @@ def build_sbatch_script(
                 # 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
                 # regardless of how many nodes the overall job spans). A pinned service always gets
-                # them, since --relative is meaningless without a node count.
+                # 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,
-                relative=offsets[service.node_pool][0] if service.node_pool else None,
+                nodelist=pool_nodes_var(service.node_pool) if service.node_pool else None,
             )
             for name, service in config.services.items()
         ]
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 90ac4700b7a..114ca1412f0 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -1636,7 +1636,8 @@ def _render(tmp_path, services, pools=None):
 
 
 def test_pinned_services_take_contiguous_node_ranges(tmp_path):
-    # --relative is what gives a service nodes of its own; without it both steps
+    # --nodelist names each pool's hosts; --relative is only a hint Slurm can move a
+    # step off (seen on a cluster run: the head landed on the worker's node). Without it both steps
     # start at node 0 and the second one shares the first one's GPUs.
     script = _render(
         tmp_path,
@@ -1645,8 +1646,10 @@ def test_pinned_services_take_contiguous_node_ranges(tmp_path):
             "scorer": _vllm(8001, "aux", tensor_parallel_size=4),
         },
     )
-    assert "--relative=0 --nodes=1 --ntasks=1" in script.split("# service: policy")[1]
-    assert "--relative=1 --nodes=1 --ntasks=1" in script.split("# service: scorer")[1]
+    assert 'export GYM_POOL_GPU_NODES="$(IFS=,; echo "${gym_nodes[*]:0:1}")"' in script
+    assert 'export GYM_POOL_AUX_NODES="$(IFS=,; echo "${gym_nodes[*]:1:1}")"' in script
+    assert '--nodelist="$GYM_POOL_GPU_NODES" --nodes=1 --ntasks=1' in script.split("# service: policy")[1]
+    assert '--nodelist="$GYM_POOL_AUX_NODES" --nodes=1 --ntasks=1' in script.split("# service: scorer")[1]
 
 
 def test_a_pinned_single_node_service_is_not_built_as_multi_node(tmp_path):
@@ -1668,7 +1671,7 @@ def test_an_unpinned_service_keeps_the_whole_allocation(tmp_path):
         tmp_path,
         {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 8}},
     )
-    assert "--relative=" not in script
+    assert "--nodelist" not in script
 
 
 def test_an_unknown_node_pool_is_named(tmp_path):
@@ -1728,7 +1731,7 @@ def test_a_worker_joins_a_head_service_by_its_placed_address(tmp_path):
     # exports it and the worker reads it rather than a hard-coded address.
     config = _head_worker_config(tmp_path, {"head": "ray_head", "resources": {"extra_gpu": 4}})
     script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
-    assert 'export GYM_RAY_ADDRESS_RAY_HEAD="$(getent hosts ${gym_ray_nodes[0]}' in script
+    assert 'export GYM_RAY_ADDRESS_RAY_HEAD="$(getent hosts ${gym_nodes[0]}' in script
     assert ':6380"' in script
     assert '--address "$GYM_RAY_ADDRESS_RAY_HEAD"' in script.split("# service: scorer")[1]
 
@@ -1740,7 +1743,7 @@ def test_a_head_on_a_later_pool_exports_that_pools_node(tmp_path):
         _TWO_POOLS,
     )
     script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
-    assert "${gym_ray_nodes[1]}" in script
+    assert "${gym_nodes[1]}" in script
 
 
 def test_ray_extra_args_are_appended():

From 04acca9bef177c9ce3744d448801b7c921adda21 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Wed, 23 Sep 2026 17:50:09 +0200
Subject: [PATCH 5/8] fix(orchestration): keep the driver on the policy's node

With services pinned to pools, the driver's srun had no node constraint, and on
a real two-node job Slurm ran it on the aux node. It reaches the policy on
localhost and joins a Ray head through that node's raylet, so both failed. The
driver now goes to the first node of the policy's pool.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 .../orchestration/executors/slurm_script.py   |  6 ++++++
 tests/unit_tests/test_slurm_script.py         | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index 5b3dc4bc353..f57932dc719 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -667,6 +667,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 ""
+    if any(s.node_pool for s in config.services.values()):
+        # With pinned services Slurm may place the driver on any node. It reaches the
+        # policy on localhost, so it goes to the first node of the policy's pool.
+        policy = config.services.get(config.driver.policy_model or "")
+        index = _pool_offsets(compute)[policy.node_pool][0] if policy and policy.node_pool else 0
+        driver_node_flags = f' --nodelist="${{gym_nodes[{index}]}}" --nodes=1 --ntasks=1'
     # The driver writes everything relative to the job directory -- `output_path`
     # above is `artifacts/rollouts.jsonl`. `#SBATCH --chdir` sets the cwd of the
     # BATCH script on the host, but inside a Pyxis container the cwd is whatever
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 114ca1412f0..6b6b01ef296 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -1765,3 +1765,22 @@ def test_a_worker_naming_something_other_than_a_head_is_refused(tmp_path):
 def test_a_worker_with_both_address_and_head_is_refused():
     with pytest.raises(ValueError, match="exactly one of `address` or `head`"):
         RayServiceConfig(type="ray", container="img", mode="worker", address="h:1", head="ray_head")
+
+
+def test_the_driver_goes_to_the_policys_node_when_services_are_pinned(tmp_path):
+    # The driver reaches the policy on localhost. Unpinned, Slurm put it on the aux
+    # node of a real two-node job, where neither the policy nor a local raylet was.
+    config = SubmitConfig.model_validate(
+        {
+            "services": {
+                "scorer": _vllm(8001, "gpu", tensor_parallel_size=4),
+                "policy": _vllm(8000, "aux", tensor_parallel_size=4),
+            },
+            "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": _TWO_POOLS}},
+            "driver": {"container": "gym:latest", "policy_model": "policy", "benchmarks": {"b": {"run": {}}}},
+            "job": {"output_path": str(tmp_path / "jobs")},
+        }
+    )
+    script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
+    driver = next(line for line in script.splitlines() if "--output=logs/driver.log" in line)
+    assert '--nodelist="${gym_nodes[1]}" --nodes=1 --ntasks=1' in driver

From 2acee8729385bd4c21a03312fe9a9964a01a92d7 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Mon, 28 Sep 2026 09:21:14 +0200
Subject: [PATCH 6/8] fix(orchestration): place the driver in every multi-node
 job

The driver reaches the policy on localhost, but its srun only asked for one
node, so in a multi-node job Slurm could start it anywhere. A multi-node
policy serves its API from node 0, a pinned one from its pool's first node;
the driver now goes there in any multi-node job, not only when services are
pinned. An unpinned service that fits on one node follows it, for the same
reason.

Also brace the --nodelist expansion. Unbraced, $gym_nodes[0] expands to the
first host followed by a literal [0].

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 .../orchestration/executors/slurm_script.py   | 78 ++++++++++++++-----
 tests/unit_tests/test_orchestration_otel.py   |  2 +
 tests/unit_tests/test_slurm_script.py         | 32 +++++++-
 3 files changed, 90 insertions(+), 22 deletions(-)

diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index f57932dc719..763665e67ce 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -179,7 +179,7 @@ def _render_service_command(
     # --nodelist names the exact hosts. --relative is only a starting point Slurm may
     # move off when that node's resources are taken, landing a service on the wrong node.
     if nodelist is not None:
-        node_flags = f' --nodelist="${nodelist}" --nodes={nodes} --ntasks={ntasks}'
+        node_flags = f' --nodelist="${{{nodelist}}}" --nodes={nodes} --ntasks={ntasks}'
     mounts_flag = f" --container-mounts={','.join(shlex.quote(m) for m in mounts)}" if mounts else ""
     workdir_flag = f" --container-workdir={shlex.quote(workdir)}" if workdir else ""
     if pre_command:
@@ -390,9 +390,9 @@ def pool_nodes_var(pool: str) -> str:
     return f"GYM_POOL_{bash_var(pool)}_NODES"
 
 
-def _render_pool_nodes(config: SubmitConfig, compute: SlurmComputeConfig) -> str:
-    """Export each node pool's hosts, so a pinned service can name them in --nodelist."""
-    if not any(s.node_pool for s in config.services.values()):
+def _render_pool_nodes(config: SubmitConfig, compute: SlurmComputeConfig, is_multi_node: bool) -> str:
+    """Declare the allocation's hosts, and export each node pool's, for --nodelist."""
+    if not is_multi_node and not any(s.node_pool for s in config.services.values()):
         return ""
     lines = [_NODE_ARRAY]
     for name, (start, count) in _pool_offsets(compute).items():
@@ -416,8 +416,9 @@ def _render_ray_head_addresses(config: SubmitConfig, compute: SlurmComputeConfig
     if not heads:
         return ""
     offsets = _pool_offsets(compute)
-    # _render_pool_nodes already declared the host array when any service is pinned.
-    lines = [] if any(s.node_pool for s in config.services.values()) else [_NODE_ARRAY]
+    # _render_pool_nodes already declared the host array in a multi-node or pinned job.
+    declared = _node_totals(compute)[0] > 1 or any(s.node_pool for s in config.services.values())
+    lines = [] if declared else [_NODE_ARRAY]
     for name, service in heads.items():
         index = offsets[service.node_pool][0] if service.node_pool else 0
         ip = f"$(getent hosts ${{gym_nodes[{index}]}} | awk '{{print $1}}')"
@@ -450,7 +451,9 @@ def _build_service_command(
     return _BUILDERS[type(service)](service)
 
 
-def _render_collector_service(config: SubmitConfig, remote_bench_dir: Path, *, is_multi_node: bool) -> str:
+def _render_collector_service(
+    config: SubmitConfig, remote_bench_dir: Path, *, is_multi_node: bool, driver_node: int | None = None
+) -> str:
     """The collector's srun step. Started before the model services so the scrape covers their
     startup.
 
@@ -476,6 +479,10 @@ def _render_collector_service(config: SubmitConfig, remote_bench_dir: Path, *, i
             "SLURM_JOB_ID": f"{RUNTIME_ENV_PREFIX}SLURM_JOB_ID",
         },
         single_node=is_multi_node,
+        # Beside the driver: it scrapes the policy API on localhost, as the driver calls it.
+        nodelist=f"gym_nodes[{driver_node}]" if driver_node is not None else None,
+        nodes=1 if driver_node is not None else None,
+        ntasks=1 if driver_node is not None else None,
         **container_kwargs,
     )
 
@@ -538,13 +545,43 @@ def _service_nodes(
     return compute.node_pools[service.node_pool].nodes
 
 
+def _driver_node(config: SubmitConfig, compute: SlurmComputeConfig) -> int:
+    """The node the driver runs on in a multi-node job: the policy's first node.
+
+    The driver reaches the policy on localhost. A multi-node policy serves its API
+    from node 0, and a pinned one from its pool's first node.
+    """
+    policy = config.services.get(config.driver.policy_model or "")
+    if policy is not None and policy.node_pool is not None:
+        return _pool_offsets(compute)[policy.node_pool][0]
+    return 0
+
+
+def _service_nodelist(
+    service: VllmServiceConfig | RayServiceConfig, driver_node: int | None, total_nodes: int
+) -> str | None:
+    """What a service's --nodelist names, as the shell variable it expands.
+
+    A pinned service names its pool. In a multi-node job an unpinned service that
+    fits on one node joins the driver, which reaches it on localhost; otherwise
+    Slurm could start it on any node.
+    """
+    if service.node_pool is not None:
+        return pool_nodes_var(service.node_pool)
+    if driver_node is not None and not _vllm_spans_multiple_nodes(service, total_nodes):
+        return f"gym_nodes[{driver_node}]"
+    return None
+
+
 def _srun_nodes(
     service: VllmServiceConfig | RayServiceConfig, compute: SlurmComputeConfig, total_nodes: int
 ) -> int | None:
     nodes = _service_nodes(service, compute, total_nodes)
     if service.node_pool is not None:
         return nodes
-    return total_nodes if _vllm_spans_multiple_nodes(service, total_nodes) else None
+    if _vllm_spans_multiple_nodes(service, total_nodes):
+        return total_nodes
+    return 1 if total_nodes > 1 else None
 
 
 def _srun_ntasks(
@@ -556,7 +593,9 @@ def _srun_ntasks(
     if service.node_pool is not None:
         pool = compute.node_pools[service.node_pool]
         return pool.nodes * pool.ntasks_per_node
-    return total_ntasks if _vllm_spans_multiple_nodes(service, total_nodes) else None
+    if _vllm_spans_multiple_nodes(service, total_nodes):
+        return total_ntasks
+    return 1 if total_nodes > 1 else None
 
 
 def _node_totals(compute: SlurmComputeConfig) -> tuple[int, int]:
@@ -601,7 +640,7 @@ def build_sbatch_script(
             render_ray_prelude()
             if any(_vllm_spans_multiple_nodes(s, total_nodes) for s in config.services.values())
             else "",
-            _render_pool_nodes(config, compute),
+            _render_pool_nodes(config, compute, is_multi_node),
             _render_ray_head_addresses(config, compute),
         )
         if block
@@ -609,8 +648,13 @@ def build_sbatch_script(
 
     observed = otel_active(config)
 
+    driver_node = _driver_node(config, compute) if is_multi_node else None
     service_commands = "\n\n".join(
-        ([_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node)] if observed else [])
+        (
+            [_render_collector_service(config, remote_bench_dir, is_multi_node=is_multi_node, driver_node=driver_node)]
+            if observed
+            else []
+        )
         + [
             _render_service_command(
                 name,
@@ -625,7 +669,7 @@ def build_sbatch_script(
                 nodes=_srun_nodes(service, compute, total_nodes),
                 ntasks=_srun_ntasks(service, compute, total_nodes, total_ntasks),
                 pre_command=service.pre_command,
-                nodelist=pool_nodes_var(service.node_pool) if service.node_pool else None,
+                nodelist=_service_nodelist(service, driver_node, total_nodes),
             )
             for name, service in config.services.items()
         ]
@@ -666,13 +710,9 @@ 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 ""
-    if any(s.node_pool for s in config.services.values()):
-        # With pinned services Slurm may place the driver on any node. It reaches the
-        # policy on localhost, so it goes to the first node of the policy's pool.
-        policy = config.services.get(config.driver.policy_model or "")
-        index = _pool_offsets(compute)[policy.node_pool][0] if policy and policy.node_pool else 0
-        driver_node_flags = f' --nodelist="${{gym_nodes[{index}]}}" --nodes=1 --ntasks=1'
+    driver_node_flags = (
+        f' --nodelist="${{gym_nodes[{driver_node}]}}" --nodes=1 --ntasks=1' if driver_node is not None else ""
+    )
     # The driver writes everything relative to the job directory -- `output_path`
     # above is `artifacts/rollouts.jsonl`. `#SBATCH --chdir` sets the cwd of the
     # BATCH script on the host, but inside a Pyxis container the cwd is whatever
diff --git a/tests/unit_tests/test_orchestration_otel.py b/tests/unit_tests/test_orchestration_otel.py
index 3254f2d2115..c1334c09940 100644
--- a/tests/unit_tests/test_orchestration_otel.py
+++ b/tests/unit_tests/test_orchestration_otel.py
@@ -309,6 +309,8 @@ def test_script_pins_the_collector_to_one_node_of_a_multi_node_job():
     )
     line = _collector_line(_script(multi))
     assert " --nodes=1 --ntasks=1" in line
+    # Named, not left to Slurm: the same node the driver is placed on.
+    assert '--nodelist="${gym_nodes[0]}"' in line
     assert " --nodes=1 --ntasks=1" not in _collector_line(_script(_config()))
 
 
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 6b6b01ef296..050e1e61100 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -1441,7 +1441,8 @@ def test_build_sbatch_script_non_vllm_service_omits_node_flags_in_multi_node_job
     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
+    # Once, on the driver's node -- not once per node of the allocation.
+    assert '--nodelist="${gym_nodes[0]}" --nodes=1 --ntasks=1' in ray_line
 
 
 # ---------------------------------------------------------------------------
@@ -1648,8 +1649,8 @@ def test_pinned_services_take_contiguous_node_ranges(tmp_path):
     )
     assert 'export GYM_POOL_GPU_NODES="$(IFS=,; echo "${gym_nodes[*]:0:1}")"' in script
     assert 'export GYM_POOL_AUX_NODES="$(IFS=,; echo "${gym_nodes[*]:1:1}")"' in script
-    assert '--nodelist="$GYM_POOL_GPU_NODES" --nodes=1 --ntasks=1' in script.split("# service: policy")[1]
-    assert '--nodelist="$GYM_POOL_AUX_NODES" --nodes=1 --ntasks=1' in script.split("# service: scorer")[1]
+    assert '--nodelist="${GYM_POOL_GPU_NODES}" --nodes=1 --ntasks=1' in script.split("# service: policy")[1]
+    assert '--nodelist="${GYM_POOL_AUX_NODES}" --nodes=1 --ntasks=1' in script.split("# service: scorer")[1]
 
 
 def test_a_pinned_single_node_service_is_not_built_as_multi_node(tmp_path):
@@ -1671,7 +1672,32 @@ def test_an_unpinned_service_keeps_the_whole_allocation(tmp_path):
         tmp_path,
         {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 8}},
     )
+    policy_line = script.split("# service: policy")[1].splitlines()[1]
+    assert "--nodelist" not in policy_line
+    assert "--nodes=2" in policy_line
+
+
+def test_the_driver_runs_on_node_0_of_a_multi_node_job(tmp_path):
+    # A multi-node policy serves its API from node 0, and the driver reaches it on
+    # localhost. Left unpinned, Slurm may start the driver on any node.
+    script = _render(
+        tmp_path,
+        {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 8}},
+    )
+    driver_line = next(line for line in script.splitlines() if "logs/driver.log" in line)
+    assert '--nodelist="${gym_nodes[0]}" --nodes=1 --ntasks=1' in driver_line
+    assert 'gym_nodes=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))' in script
+
+
+def test_a_single_node_job_places_nothing(tmp_path):
+    one = {"gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}}
+    script = _render(
+        tmp_path,
+        {"policy": {"type": "vllm", "container": "img", "model": "/ckpt", "port": 8000, "tensor_parallel_size": 4}},
+        one,
+    )
     assert "--nodelist" not in script
+    assert "gym_nodes=" not in script
 
 
 def test_an_unknown_node_pool_is_named(tmp_path):

From 7581947aadbb7c50364b71aede6913130268b145 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Mon, 28 Sep 2026 11:43:22 +0200
Subject: [PATCH 7/8] refactor(orchestration): make a ray service one cluster
 over node pools

A ray service now names the node pools it spans instead of starting a head
and workers separately. The head starts beside the driver, every other node
of those pools joins it as a worker, and all of them use the service's one
container. `resources` is keyed by pool, so only the pools that need a custom
resource advertise it.

Removes `mode`, `address` and `head`. Workers wait for the head, and flags
Ray accepts only on a head are dropped from their command.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 nemo_gym/orchestration/api.py                 |  49 +++---
 .../orchestration/executors/slurm_script.py   | 158 ++++++++++++++----
 tests/unit_tests/test_slurm_script.py         | 154 +++++++++++------
 3 files changed, 258 insertions(+), 103 deletions(-)

diff --git a/nemo_gym/orchestration/api.py b/nemo_gym/orchestration/api.py
index 7c831a8fa6d..3b8bb47c387 100644
--- a/nemo_gym/orchestration/api.py
+++ b/nemo_gym/orchestration/api.py
@@ -158,31 +158,33 @@ def effective_ray_serve(service: "VllmServiceConfig", total_nodes: int, gpus_per
 
 class RayServiceConfig(BaseServiceConfig):
     type: Literal["ray"]
-    # "head" starts a cluster; "worker" joins the one at `address`. A worker is how a
-    # second node joins the driver's Ray cluster and offers its GPUs to actors the
-    # benchmark schedules (e.g. a scorer that runs off the policy's node).
-    mode: Literal["head", "worker"] = "head"
-    # For mode="worker": the head's host:port, or `head`, the name of a ray head
-    # service in this config whose address is only known once the job is placed.
-    address: str | None = None
-    head: str | None = None
-    # Head only; ignored by a worker, which takes the port from `address`.
+    # Node pools the cluster spans. The head starts on the driver's node and every
+    # other node of these pools joins it as a worker. Empty: a single-node cluster.
+    node_pools: list[str] = []
     port: int = 6379
-    # Custom Ray resources this node advertises, e.g. {"extra_gpu": 4}. A benchmark
-    # asks for these by name rather than by num_gpus when it manages device placement
-    # itself.
-    resources: dict[str, float] = {}
+    # Custom Ray resources advertised by each spanned pool's nodes, keyed by pool,
+    # e.g. {"aux": {"extra_gpu": 4}}. A benchmark asks for these by name when it
+    # manages device placement itself.
+    resources: dict[str, dict[str, float]] = {}
     num_cpus: int | None = None
     num_gpus: int | None = None
-    # Raw extra flags appended verbatim to `ray start` (e.g. fixed ports).
+    # Raw extra flags appended verbatim to `ray start` on every node (e.g. fixed ports).
     extra_args: str = ""
 
     @model_validator(mode="after")
-    def _validate_mode(self) -> "RayServiceConfig":
-        if self.mode == "worker" and bool(self.address) == bool(self.head):
-            raise ValueError("A ray service with mode='worker' needs exactly one of `address` or `head`.")
-        if self.mode == "head" and (self.address or self.head):
-            raise ValueError("A ray service with mode='head' starts its own cluster; remove `address`/`head`.")
+    def _validate_shape(self) -> "RayServiceConfig":
+        if self.node_pool is not None:
+            raise ValueError(
+                f"A ray service spans `node_pools`, not a single `node_pool`; use node_pools: [{self.node_pool}]."
+            )
+        if len(set(self.node_pools)) != len(self.node_pools):
+            raise ValueError(f"A ray service's node_pools lists a pool more than once: {self.node_pools}.")
+        unspanned = sorted(set(self.resources) - set(self.node_pools))
+        if unspanned:
+            raise ValueError(
+                f"A ray service sets resources for {', '.join(unspanned)}, which it does not span "
+                f"(node_pools: {self.node_pools}). Add the pool to node_pools or drop its resources."
+            )
         return self
 
 
@@ -348,11 +350,12 @@ def _resolve_and_validate_placements(self) -> "SubmitConfig":
                     f"({', '.join(sorted(compute_names))})."
                 )
 
-            if isinstance(service, RayServiceConfig) and service.head is not None:
-                head = self.services.get(service.head)
-                if not isinstance(head, RayServiceConfig) or head.mode != "head":
+            if isinstance(service, RayServiceConfig):
+                unknown = [pool for pool in service.node_pools if pool not in pool_names]
+                if unknown:
                     raise ValueError(
-                        f"Service '{service_name}' head '{service.head}' must name a ray service with mode='head'."
+                        f"Service '{service_name}' node_pools {unknown} do not match any node pool of compute "
+                        f"'{service.placement}' ({', '.join(sorted(pool_names)) or 'none declared'})."
                     )
 
             if service.node_pool is not None and service.node_pool not in pool_names:
diff --git a/nemo_gym/orchestration/executors/slurm_script.py b/nemo_gym/orchestration/executors/slurm_script.py
index 763665e67ce..5a1465221db 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -354,31 +354,42 @@ def _build_vllm_ray_serve_command(
     )
 
 
-def _build_ray_command(service: RayServiceConfig) -> str:
+# Flags `ray start` refuses on a worker. They reach a worker through the service's
+# shared extra_args; Gym decides which nodes are workers, so Gym drops them there.
+_RAY_HEAD_ONLY_FLAG = re.compile(r"\s--(?:port|redis-shard-ports|include-dashboard)(?:[= ]\S+)?")
+
+
+def _build_ray_command(
+    service: RayServiceConfig, *, pool: str | None = None, address_var: str | None = None, worker: bool = False
+) -> str:
+    """`ray start` for one node of a ray service: the head, or with `worker` a node joining it.
+
+    `address_var` names the env var holding the head's host:port (see ray_head_address_var).
+    """
     # --block keeps the srun step alive. `ray start` daemonises and returns, so
     # without it the step exits the moment the node is up and Slurm tears the
     # service down again.
     cmd = "ray start --block"
-    if service.mode == "head":
-        cmd += f" --head --port {service.port}"
-    elif service.head is not None:
-        # Double-quoted, not shlex-quoted: the head's address is only known once
-        # Slurm places the job, and build_sbatch_script exports it under this name.
-        cmd += f' --address "${ray_head_address_var(service.head)}"'
+    if worker:
+        cmd += f' --address "${address_var}"'
     else:
-        cmd += f" --address {shlex.quote(str(service.address))}"
+        cmd += f" --head --port {service.port}"
+        if address_var is not None:
+            # Bind the head to the address the workers and the driver dial.
+            cmd += f' --node-ip-address "${{{address_var}%:*}}"'
     if service.num_cpus is not None:
         cmd += f" --num-cpus {service.num_cpus}"
     if service.num_gpus is not None:
         cmd += f" --num-gpus {service.num_gpus}"
-    if service.resources:
+    resources = service.resources.get(pool, {}) if pool is not None else {}
+    if resources:
         # Ray takes fractional custom resources, so the field is float-typed, but a
-        # whole number is written as one: {"extra_gpu": 4}, not 4.0, so the rendered
-        # command reads the way the config does.
-        resources = {k: int(v) if v.is_integer() else v for k, v in service.resources.items()}
+        # whole number is written as one: {"extra_gpu": 4}, not 4.0.
+        resources = {k: int(v) if v.is_integer() else v for k, v in resources.items()}
         cmd += " --resources=" + shlex.quote(json.dumps(resources, sort_keys=True))
-    if service.extra_args:
-        cmd += " " + service.extra_args
+    extra_args = _RAY_HEAD_ONLY_FLAG.sub("", " " + service.extra_args).strip() if worker else service.extra_args
+    if extra_args:
+        cmd += " " + extra_args
     return cmd
 
 
@@ -390,9 +401,14 @@ def pool_nodes_var(pool: str) -> str:
     return f"GYM_POOL_{bash_var(pool)}_NODES"
 
 
+def _places_services(config: SubmitConfig, is_multi_node: bool) -> bool:
+    """Whether the script declares the allocation's host array to place services with."""
+    return is_multi_node or any(s.node_pool for s in config.services.values())
+
+
 def _render_pool_nodes(config: SubmitConfig, compute: SlurmComputeConfig, is_multi_node: bool) -> str:
     """Declare the allocation's hosts, and export each node pool's, for --nodelist."""
-    if not is_multi_node and not any(s.node_pool for s in config.services.values()):
+    if not _places_services(config, is_multi_node):
         return ""
     lines = [_NODE_ARRAY]
     for name, (start, count) in _pool_offsets(compute).items():
@@ -401,31 +417,109 @@ def _render_pool_nodes(config: SubmitConfig, compute: SlurmComputeConfig, is_mul
     return "\n".join(lines)
 
 
-def ray_head_address_var(head_service: str) -> str:
-    """The env var build_sbatch_script exports with a ray head service's host:port."""
-    return f"GYM_RAY_ADDRESS_{bash_var(head_service)}"
+def ray_head_address_var(ray_service: str) -> str:
+    """The env var build_sbatch_script exports with a ray service's head host:port."""
+    return f"GYM_RAY_ADDRESS_{bash_var(ray_service)}"
+
+
+def ray_workers_var(ray_service: str, pool: str) -> str:
+    """The env var build_sbatch_script exports with a ray service's worker hosts in one pool."""
+    return f"GYM_RAY_{bash_var(ray_service)}_{bash_var(pool)}_WORKERS"
+
+
+def _ray_worker_ranges(
+    service: RayServiceConfig, compute: SlurmComputeConfig, head_node: int
+) -> dict[str, tuple[int, int]]:
+    """Each spanned pool's worker nodes as (first index, count): every node but the head.
+
+    The head is the driver's node, which is always the first node of some pool.
+    """
+    offsets = _pool_offsets(compute)
+    ranges: dict[str, tuple[int, int]] = {}
+    for pool in service.node_pools:
+        start, count = offsets[pool]
+        if start == head_node:
+            start, count = start + 1, count - 1
+        if count > 0:
+            ranges[pool] = (start, count)
+    return ranges
 
 
 def _render_ray_head_addresses(config: SubmitConfig, compute: SlurmComputeConfig) -> str:
-    """Export every ray head service's address, resolved from the node Slurm gave it.
+    """Export every ray service's head address, and its worker hosts, from the nodes Slurm gave it.
 
-    A worker or the driver in another step can then join it without knowing,
+    The head runs beside the driver. Workers and the driver join it without knowing,
     when the config is written, which host the job will land on.
     """
-    heads = {n: s for n, s in config.services.items() if isinstance(s, RayServiceConfig) and s.mode == "head"}
-    if not heads:
+    services = {n: s for n, s in config.services.items() if isinstance(s, RayServiceConfig)}
+    if not services:
         return ""
-    offsets = _pool_offsets(compute)
-    # _render_pool_nodes already declared the host array in a multi-node or pinned job.
-    declared = _node_totals(compute)[0] > 1 or any(s.node_pool for s in config.services.values())
-    lines = [] if declared else [_NODE_ARRAY]
-    for name, service in heads.items():
-        index = offsets[service.node_pool][0] if service.node_pool else 0
-        ip = f"$(getent hosts ${{gym_nodes[{index}]}} | awk '{{print $1}}')"
+    head_node = _driver_node(config, compute)
+    lines = [] if _places_services(config, _node_totals(compute)[0] > 1) else [_NODE_ARRAY]
+    for name, service in services.items():
+        ip = f"$(getent hosts ${{gym_nodes[{head_node}]}} | awk '{{print $1}}')"
         lines.append(f'export {ray_head_address_var(name)}="{ip}:{service.port}"')
+        for pool, (start, count) in _ray_worker_ranges(service, compute, head_node).items():
+            lines.append(f'export {ray_workers_var(name, pool)}="$(IFS=,; echo "${{gym_nodes[*]:{start}:{count}}}")"')
     return "\n".join(lines)
 
 
+def _pool_of(compute: SlurmComputeConfig, node: int) -> str | None:
+    for name, (start, count) in _pool_offsets(compute).items():
+        if start <= node < start + count:
+            return name
+    return None
+
+
+def _render_ray_service(
+    name: str,
+    service: RayServiceConfig,
+    config: SubmitConfig,
+    compute: SlurmComputeConfig,
+    driver_node: int | None,
+) -> str:
+    """A ray service's srun steps: the head beside the driver, then one step per spanned pool's workers.
+
+    Every step uses the service's one container, env and mounts.
+    """
+    total_nodes, total_ntasks = _node_totals(compute)
+    head_node = _driver_node(config, compute)
+    head_pool = _pool_of(compute, head_node)
+    address_var = ray_head_address_var(name)
+    steps = [
+        _render_service_command(
+            name,
+            service.container,
+            _build_ray_command(
+                service, pool=head_pool if head_pool in service.node_pools else None, address_var=address_var
+            ),
+            service.env or None,
+            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,
+            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():
+        steps.append(
+            _render_service_command(
+                f"{name}_{pool}_workers",
+                service.container,
+                _build_ray_command(service, pool=pool, address_var=address_var, worker=True),
+                service.env or None,
+                service.mounts or None,
+                nodes=count,
+                ntasks=count,
+                pre_command=f"{service.pre_command.rstrip()}\n{wait_for_head}".lstrip(),
+                nodelist=ray_workers_var(name, pool),
+            )
+        )
+    return "\n\n".join(steps)
+
+
 _BUILDERS = {
     VllmServiceConfig: _build_vllm_command,
     RayServiceConfig: _build_ray_command,
@@ -672,6 +766,12 @@ def build_sbatch_script(
                 nodelist=_service_nodelist(service, driver_node, total_nodes),
             )
             for name, service in config.services.items()
+            if not isinstance(service, RayServiceConfig)
+        ]
+        + [
+            _render_ray_service(name, service, config, compute, driver_node)
+            for name, service in config.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 050e1e61100..9841ba40522 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -1717,80 +1717,132 @@ def test_a_ray_head_blocks_so_slurm_keeps_the_step_alive():
     assert command == "ray start --block --head --port 6379"
 
 
-def test_a_ray_worker_joins_the_head_and_advertises_its_resources():
-    command = _build_ray_command(
-        RayServiceConfig(
-            type="ray",
-            container="img",
-            mode="worker",
-            address="10.0.0.1:6379",
-            num_gpus=4,
-            resources={"extra_gpu": 4},
-        )
+def test_a_ray_worker_joins_the_head_and_advertises_its_pools_resources():
+    service = RayServiceConfig(
+        type="ray", container="img", node_pools=["aux"], num_gpus=0, resources={"aux": {"extra_gpu": 4}}
     )
-    assert command == "ray start --block --address 10.0.0.1:6379 --num-gpus 4 --resources='{\"extra_gpu\": 4}'"
+    command = _build_ray_command(service, pool="aux", address_var="GYM_RAY_ADDRESS_RAY", worker=True)
+    assert command == (
+        'ray start --block --address "$GYM_RAY_ADDRESS_RAY" --num-gpus 0 --resources=\'{"extra_gpu": 4}\''
+    )
+
+
+def test_a_ray_worker_drops_the_flags_ray_allows_only_on_a_head():
+    service = RayServiceConfig(
+        type="ray", container="img", extra_args="--include-dashboard=false --port 1 --node-manager-port=8366"
+    )
+    command = _build_ray_command(service, address_var="A", worker=True)
+    assert command == 'ray start --block --address "$A" --node-manager-port=8366'
+    assert _build_ray_command(service).endswith("--include-dashboard=false --port 1 --node-manager-port=8366")
+
+
+def test_the_old_head_and_worker_fields_are_gone():
+    with pytest.raises(ValueError, match="Extra inputs are not permitted"):
+        RayServiceConfig(type="ray", container="img", mode="worker", address="10.0.0.1:6379")
+
+
+def test_a_ray_service_names_node_pools_not_a_node_pool():
+    with pytest.raises(ValueError, match=r"use node_pools: \[aux\]"):
+        RayServiceConfig(type="ray", container="img", node_pool="aux")
 
 
-def test_a_ray_worker_without_an_address_is_refused():
-    with pytest.raises(ValueError, match="exactly one of `address` or `head`"):
-        RayServiceConfig(type="ray", container="img", mode="worker")
+def test_ray_resources_for_a_pool_the_service_does_not_span_are_refused():
+    with pytest.raises(ValueError, match="sets resources for aux, which it does not span"):
+        RayServiceConfig(type="ray", container="img", node_pools=["gpu"], resources={"aux": {"extra_gpu": 4}})
 
 
-def test_a_ray_head_with_an_address_is_refused():
-    with pytest.raises(ValueError, match="starts its own cluster"):
-        RayServiceConfig(type="ray", container="img", address="10.0.0.1:6379")
+def test_an_unknown_ray_node_pool_is_named(tmp_path):
+    with pytest.raises(ValueError, match=r"node_pools \['nope'\] do not match any node pool"):
+        _placement_config(tmp_path, {"ray": {"type": "ray", "container": "img", "node_pools": ["nope"]}}, _TWO_POOLS)
 
 
-def _head_worker_config(tmp_path, worker):
-    return _placement_config(
+def _comet_like(tmp_path, pools=None, **ray):
+    return _render(
         tmp_path,
         {
-            "ray_head": {"type": "ray", "container": "img", "mode": "head", "port": 6380, "node_pool": "gpu"},
-            "scorer": {"type": "ray", "container": "img", "mode": "worker", "node_pool": "aux", **worker},
+            "policy": _vllm(8000, "gpu", tensor_parallel_size=4),
+            "ray": {
+                "type": "ray",
+                "container": "img",
+                "node_pools": ["gpu", "aux"],
+                "port": 6380,
+                "resources": {"aux": {"extra_gpu": 4}},
+                **ray,
+            },
         },
-        _TWO_POOLS,
+        pools,
     )
 
 
-def test_a_worker_joins_a_head_service_by_its_placed_address(tmp_path):
-    # The head's host is only known once Slurm places the job, so the script
-    # exports it and the worker reads it rather than a hard-coded address.
-    config = _head_worker_config(tmp_path, {"head": "ray_head", "resources": {"extra_gpu": 4}})
-    script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
-    assert 'export GYM_RAY_ADDRESS_RAY_HEAD="$(getent hosts ${gym_nodes[0]}' in script
+def test_a_ray_service_starts_its_head_beside_the_driver(tmp_path):
+    script = _comet_like(tmp_path)
+    assert 'export GYM_RAY_ADDRESS_RAY="$(getent hosts ${gym_nodes[0]}' in script
     assert ':6380"' in script
-    assert '--address "$GYM_RAY_ADDRESS_RAY_HEAD"' in script.split("# service: scorer")[1]
+    head = script.split("# service: ray\n")[1].split("\n\n")[0]
+    assert '--nodelist="${gym_nodes[0]}" --nodes=1 --ntasks=1' in head
+    assert '--head --port 6380 --node-ip-address "${GYM_RAY_ADDRESS_RAY%:*}"' in head
+    # The head's own pool advertises nothing here, so it offers no extra_gpu.
+    assert "extra_gpu" not in head
 
 
-def test_a_head_on_a_later_pool_exports_that_pools_node(tmp_path):
-    config = _placement_config(
-        tmp_path,
-        {"ray_head": {"type": "ray", "container": "img", "mode": "head", "node_pool": "aux"}},
-        _TWO_POOLS,
-    )
-    script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
-    assert "${gym_nodes[1]}" in script
+def test_every_other_spanned_node_joins_as_a_worker(tmp_path):
+    pools = {**_TWO_POOLS, "aux": {**_TWO_POOLS["aux"], "nodes": 2}}
+    script = _comet_like(tmp_path, pools)
+    assert 'export GYM_RAY_RAY_AUX_WORKERS="$(IFS=,; echo "${gym_nodes[*]:1:2}")"' in script
+    # The gpu pool's only node is the head, so it gets no worker step.
+    assert "GYM_RAY_RAY_GPU_WORKERS" not in script
+    workers = script.split("# service: ray_aux_workers\n")[1].split("\n\n")[0]
+    assert '--nodelist="${GYM_RAY_RAY_AUX_WORKERS}" --nodes=2 --ntasks=2' in workers
+    assert '--address "$GYM_RAY_ADDRESS_RAY"' in workers
+    assert "extra_gpu" in workers
+
+
+def test_a_ray_service_runs_every_step_in_its_one_container(tmp_path):
+    script = _comet_like(tmp_path, container="ray-img:1", mounts=["/x:/x"], pre_command="setup")
+    for name in ("ray", "ray_aux_workers"):
+        step = script.split(f"# service: {name}\n")[1].split("\n\n")[0]
+        assert "--container-image=ray-img:1" in step
+        assert "--container-mounts=/x:/x" in step
+        assert "setup" in step
 
 
-def test_ray_extra_args_are_appended():
-    command = _build_ray_command(
-        RayServiceConfig(type="ray", container="img", extra_args="--include-dashboard=false --node-manager-port=8366")
+def test_a_worker_waits_for_the_head_after_its_own_setup(tmp_path):
+    script = _comet_like(tmp_path, pre_command="install-ray")
+    workers = script.split("# service: ray_aux_workers\n")[1].split("\n\n")[0]
+    assert workers.index("install-ray") < workers.index("until ray status") < workers.index("exec ray start")
+
+
+def test_the_ray_head_follows_the_driver_to_a_later_pool(tmp_path):
+    config = SubmitConfig.model_validate(
+        {
+            "services": {
+                "policy": _vllm(8000, "aux", tensor_parallel_size=4),
+                "ray": {"type": "ray", "container": "img", "node_pools": ["gpu", "aux"]},
+            },
+            "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": _TWO_POOLS}},
+            "driver": {"container": "gym:latest", "policy_model": "policy", "benchmarks": {"b": {"run": {}}}},
+            "job": {"output_path": str(tmp_path / "jobs")},
+        }
     )
-    assert command.endswith("--include-dashboard=false --node-manager-port=8366")
+    script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
+    assert '--nodelist="${gym_nodes[1]}"' in script.split("# service: ray\n")[1].split("\n\n")[0]
+    assert 'export GYM_RAY_ADDRESS_RAY="$(getent hosts ${gym_nodes[1]}' in script
+    assert 'export GYM_RAY_RAY_GPU_WORKERS="$(IFS=,; echo "${gym_nodes[*]:0:1}")"' in script
+    assert "GYM_RAY_RAY_AUX_WORKERS" not in script
 
 
-def test_a_worker_naming_something_other_than_a_head_is_refused(tmp_path):
-    with pytest.raises(ValueError, match="must name a ray service with mode='head'"):
-        _placement_config(
-            tmp_path,
-            {"scorer": {"type": "ray", "container": "img", "mode": "worker", "head": "missing", "node_pool": "aux"}},
-            _TWO_POOLS,
-        )
+def test_a_single_node_ray_service_is_just_a_head(tmp_path):
+    one = {"gpu": {"partition": "batch", "nodes": 1, "ntasks_per_node": 1, "gpus_per_node": 4}}
+    script = _render(tmp_path, {"ray": {"type": "ray", "container": "img", "node_pools": ["gpu"]}}, one)
+    assert "_workers" not in script
+    assert "--nodelist" not in script
+    assert 'export GYM_RAY_ADDRESS_RAY="$(getent hosts ${gym_nodes[0]}' in script
 
 
-def test_a_worker_with_both_address_and_head_is_refused():
-    with pytest.raises(ValueError, match="exactly one of `address` or `head`"):
-        RayServiceConfig(type="ray", container="img", mode="worker", address="h:1", head="ray_head")
+def test_a_rendered_ray_service_is_valid_bash(tmp_path):
+    script = _comet_like(tmp_path, pre_command="echo 'quoted'", extra_args="--temp-dir=/tmp/ray-$SLURM_JOB_ID")
+    result = subprocess.run(["bash", "-n"], input=script, text=True, capture_output=True)
+    assert result.returncode == 0, result.stderr
 
 
 def test_the_driver_goes_to_the_policys_node_when_services_are_pinned(tmp_path):

From 4808c74811f2d14e829926e70369e6f968a67ae0 Mon Sep 17 00:00:00 2001
From: Wojciech Prazuch <wprazuch@nvidia.com>
Date: Mon, 28 Sep 2026 13:46:44 +0200
Subject: [PATCH 8/8] fix(orchestration): probe each service on the node it
 runs on

The health probe runs from the batch script on node 0, but a pinned service
answers on its pool's first node, and a ray head and the collector follow the
driver's node. Probed on localhost, they waited out their timeout on a later
node. Probe them where _service_nodelist places them.

Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>
---
 .../executors/script_templates.py             |  6 ++-
 .../orchestration/executors/slurm_script.py   | 39 ++++++++++++--
 tests/unit_tests/test_slurm_script.py         | 54 +++++++++++++++++++
 3 files changed, 93 insertions(+), 6 deletions(-)

diff --git a/nemo_gym/orchestration/executors/script_templates.py b/nemo_gym/orchestration/executors/script_templates.py
index ce9fb5d52c1..f30029ddf26 100644
--- a/nemo_gym/orchestration/executors/script_templates.py
+++ b/nemo_gym/orchestration/executors/script_templates.py
@@ -134,11 +134,13 @@ def render_vllm_ray_symmetric_run(inner_cmd: str, total_nodes: int, resource_fla
     )
 
 
-def render_health_check(name: str, port: int, path: str, timeout: int) -> str:
+def render_health_check(name: str, port: int, path: str, timeout: int, host: str = "localhost") -> str:
+    """`host` is emitted verbatim so it may be a shell expansion: a service pinned to
+    a node pool answers on that pool's head, not on the node running this script."""
     return _HEALTH_WAIT_MULTI.format(
         name=name,
         name_upper=bash_var(name),
-        url=f"http://localhost:{port}",
+        url=f"http://{host}:{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 5a1465221db..1d6abfb5b69 100644
--- a/nemo_gym/orchestration/executors/slurm_script.py
+++ b/nemo_gym/orchestration/executors/slurm_script.py
@@ -581,9 +581,14 @@ def _render_collector_service(
     )
 
 
-def _render_collector_health_check(config: SubmitConfig) -> str:
+def _render_collector_health_check(config: SubmitConfig, driver_node: int | None = None) -> str:
+    # The collector runs beside the driver (see _render_collector_service).
     return render_health_check(
-        COLLECTOR_SERVICE_NAME, COLLECTOR_HEALTH_PORT, "/", config.otel.health_check_timeout_seconds
+        COLLECTOR_SERVICE_NAME,
+        COLLECTOR_HEALTH_PORT,
+        "/",
+        config.otel.health_check_timeout_seconds,
+        _probe_host(driver_node or 0),
     )
 
 
@@ -713,6 +718,28 @@ def _with_default_capture_dir(run: dict[str, Any], remote_bench_dir: Path) -> di
     return run
 
 
+def _probe_host(node: int) -> str:
+    """How the batch script, which runs on the allocation's first node, reaches `node`."""
+    return "localhost" if node == 0 else f"${{gym_nodes[{node}]}}"
+
+
+def _health_check_host(
+    service: VllmServiceConfig | RayServiceConfig,
+    config: SubmitConfig,
+    compute: SlurmComputeConfig,
+    driver_node: int | None,
+    total_nodes: int,
+) -> str:
+    """Where this service answers its health probe: the node it runs on, as _service_nodelist places it."""
+    if isinstance(service, RayServiceConfig):
+        return _probe_host(_driver_node(config, compute))
+    if service.node_pool is not None:
+        return _probe_host(_pool_offsets(compute)[service.node_pool][0])
+    if driver_node is not None and not _vllm_spans_multiple_nodes(service, total_nodes):
+        return _probe_host(driver_node)
+    return "localhost"
+
+
 def build_sbatch_script(
     config: SubmitConfig,
     benchmark_name: str,
@@ -776,10 +803,14 @@ def build_sbatch_script(
     )
 
     health_checks = "\n\n".join(
-        ([_render_collector_health_check(config)] if observed else [])
+        ([_render_collector_health_check(config, driver_node)] if observed else [])
         + [
             render_health_check(
-                name, service.health_check.port, service.health_check.path, service.health_check.timeout_seconds
+                name,
+                service.health_check.port,
+                service.health_check.path,
+                service.health_check.timeout_seconds,
+                _health_check_host(service, config, compute, driver_node, total_nodes),
             )
             for name, service in config.services.items()
             if service.health_check
diff --git a/tests/unit_tests/test_slurm_script.py b/tests/unit_tests/test_slurm_script.py
index 9841ba40522..74b0bfa245b 100644
--- a/tests/unit_tests/test_slurm_script.py
+++ b/tests/unit_tests/test_slurm_script.py
@@ -1862,3 +1862,57 @@ def test_the_driver_goes_to_the_policys_node_when_services_are_pinned(tmp_path):
     script = build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
     driver = next(line for line in script.splitlines() if "--output=logs/driver.log" in line)
     assert '--nodelist="${gym_nodes[1]}" --nodes=1 --ntasks=1' in driver
+
+
+# ---------------------------------------------------------------------------
+# health probes
+# ---------------------------------------------------------------------------
+
+
+def _driver_on_aux(tmp_path, services):
+    config = SubmitConfig.model_validate(
+        {
+            "services": {"policy": _vllm(8000, "aux", tensor_parallel_size=4), **services},
+            "compute": {"hsg": {"type": "slurm", "account": "acct", "node_pools": _TWO_POOLS}},
+            "driver": {"container": "gym:latest", "policy_model": "policy", "benchmarks": {"b": {"run": {}}}},
+            "job": {"output_path": str(tmp_path / "jobs")},
+        }
+    )
+    return build_sbatch_script(config, "b", config.driver.benchmarks["b"], config.compute["hsg"], tmp_path / "b")
+
+
+_RAY_WITH_PROBE = {
+    "type": "ray",
+    "container": "img",
+    "node_pools": ["gpu", "aux"],
+    "health_check": {"port": 8011, "path": "/"},
+}
+
+
+def test_a_pinned_service_on_a_later_pool_is_probed_where_it_runs(tmp_path):
+    # The probe runs on node 0; the aux pool's service answers on node 1.
+    script = _render(tmp_path, {"policy": _vllm(8000, "gpu"), "scorer": _vllm(8001, "aux")})
+    assert "Waiting for scorer at http://${gym_nodes[1]}:8001" in script
+    assert "Waiting for policy at http://localhost:8000" in script
+
+
+def test_a_ray_head_is_probed_beside_a_non_zero_driver_node(tmp_path):
+    script = _driver_on_aux(tmp_path, {"ray": _RAY_WITH_PROBE})
+    assert "Waiting for ray at http://${gym_nodes[1]}:8011" in script
+
+
+def test_the_collector_is_probed_beside_a_non_zero_driver_node(tmp_path):
+    script = _driver_on_aux(tmp_path, {})
+    assert "Waiting for otel_collector at http://${gym_nodes[1]}:13133" in script
+
+
+def test_services_on_node_0_are_probed_locally(tmp_path):
+    script = _render(tmp_path, {"policy": _vllm(8000, "gpu", tensor_parallel_size=4), "ray": _RAY_WITH_PROBE})
+    assert "Waiting for ray at http://localhost:8011" in script
+    assert "Waiting for policy at http://localhost:8000" in script
+
+
+def test_an_unpinned_multi_node_service_is_probed_on_node_0(tmp_path):
+    # It spans the allocation and serves its API from node 0, where the probe runs.
+    script = _driver_on_aux(tmp_path, {"judge": {"type": "vllm", "container": "img", "model": "/j", "port": 9000}})
+    assert "Waiting for judge at http://localhost:9000" in script