[wprazuch]

Follows up on review of #3654 (#3654 (comment)): automatic GPU placement, agreed to land separately. Stacked on it: base is wprazuch/service-node-placement.

In one sentence

Services that share a node now get their own GPUs automatically, unless you say otherwise.

The problem

Every service on a node sees every GPU. Put a 1-GPU judge and a 2-GPU policy on one 4-GPU node and both start on GPU 0. Today the only fix is writing CUDA_VISIBLE_DEVICES by hand in each service's env.

What changed

When two or more services that use GPUs share a node (same node_pool, or both in a one-node job), each vLLM service gets its own CUDA_VISIBLE_DEVICES. The count is TP x PP x instances on that node. Services are handed the lowest free GPUs in the order they are declared.

The value is exported inside the step's command, not passed as env K=V srun. On a 4-GPU GB200 node, Slurm reset an env-prefix value to 0,1,2,3 (with and without pyxis). An export inside the step survived, and CUDA saw only the planned GPUs:

[gym_judge_want_0]    CVD=0   cuda sees 1: ['0008:01:00.0']
[gym_policy_want_1,2] CVD=1,2 cuda sees 2: ['0009:01:00.0', '0018:01:00.0']
[gym_ray_want_3]      CVD=3   cuda sees 1: ['0019:01:00.0']

services:
  judge:  {type: vllm, model: /judge}                           # -> CUDA_VISIBLE_DEVICES=0
  policy: {type: vllm, model: /policy, tensor_parallel_size: 2} # -> CUDA_VISIBLE_DEVICES=1,2
  ray:    {type: ray}                                           # -> CUDA_VISIBLE_DEVICES=3 (the rest)
  shared: {type: vllm, model: /small, share_gpus: true}         # opt out: sees every GPU

• A CUDA_VISIBLE_DEVICES you set yourself wins, and its GPUs count as taken. It now actually takes effect: it is exported inside the step too, even for a single service on a node (under Slurm it was silently ignored before).
• Overlapping manual sets, or asking for more GPUs than the node has, fail at config validation. The error lists what each service needs and suggests node_pool, another node, or share_gpus.
• share_gpus: true opts a service out.
• A ray service on a packed node gets the GPUs nobody took, unless it sets its own. One vLLM service next to a GPU-using ray service is packed too, so Ray does not advertise the model's GPUs.

Automatic node choice

A pool with nodes: auto is sized by Gym. The policy goes on the pool's first node, beside the driver. The other services go biggest first into the first node with room, and are never split. A service bigger than one node gets whole nodes of its own: 6 GPUs on 4-GPU nodes takes 2 nodes. More than one node splits the pool into sub-pools <pool>-0, <pool>-1, ..., placed like any pool. A ray service spanning the pool spans all of them and gets each node's leftovers.

compute:
  cluster:
    node_pools:
      gpu: {partition: batch, nodes: auto, gpus_per_node: 4}
services:
  judge:  {type: vllm, model: /judge}                           # 1 GPU
  policy: {type: vllm, model: /policy, tensor_parallel_size: 2} # 2 GPUs
# -> one node: judge on GPU 0, policy on GPUs 1,2

Health checks reach a service on the node it runs on. That fix now comes from #3654. This PR only extends its coverage to auto-sized pools: a service nodes: auto moves to a later node is probed there.

What stays the same

A node with a single GPU service renders nothing new, and explicit node counts are unchanged. A ray service with num_gpus: 0 does not count as a GPU service. A manual value that is not an index list (e.g. runtime:CUDA_VISIBLE_DEVICES) is kept and treated as sharing.

Out of scope

• Multi-node services that Ray places (one instance spanning nodes, or Ray Serve across nodes). A pool holding one is left as today.
• An unpinned vLLM service in a multi-node job. It spans every node, so it is left as today.
• Moving services between explicitly sized pools. nodes: auto packs only within its own pool.
• With nodes: auto, a Ray-placed service bigger than a node can leave GPUs unused on its last node.

Testing

• New unit tests in tests/unit_tests/test_slurm_script.py cover: disjoint sets, a lone service renders nothing, manual wins and reserves, overlap and over-subscription errors, opt-out, Ray leftovers, per-pool packing, multi-node data parallel, the in-step export, and nodes: auto (1 GPU + 2 GPUs packs to one node and renders exactly like nodes: 1, 6 + 4 GPUs on 4-GPU nodes, the policy on node 0, explicit counts unchanged, Ray leftovers, manual values).
• All orchestration tests pass (399). ruff check and ruff format --check are clean.
• Mutation checks: breaking disjointness, the two-service trigger, first-fit sharing, policy-first order, whole-node sizing, or the manual-overlap check each fails at least one test.
• Cluster: the export shape was run on oci-hsg (output above).
• End to end on oci-hsg (4-GPU GB200 nodes, gym eval submit, AIME25 with a vLLM policy and a separate vLLM judge on one node):
  • nodes: 1: policy got CUDA_VISIBLE_DEVICES=0,1, judge got 2. nvidia-smi showed the policy's two workers on GPUs 0 and 1, the judge on GPU 2, and GPU 3 empty. 256/256 rollouts, 0 failures. The policy served 256 requests and the judge 70. No OOM.
  • nodes: auto, same services: resolved to 1 node and rendered the same script, with the same GPU split. 256/256 rollouts, the judge served 64 requests.
  • nodes: auto with 6 GPUs of need: resolved to 2 nodes. Policy (0,1) and judge (2,3) on node 0, the third model alone on node 1. It passed its health check at ${gym_nodes[1]} and answered a request there.
• A WMT24++ COMET config on this branch renders byte-identical to the base branch, and ran clean (8/8 rollouts, COMET and spBLEU in the aggregate metrics).

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.