[wprazuch]

Same change as #3657, rebased onto main on a new branch; #3657's branch stays untouched for now.

In one sentence

Express prefill/decode disaggregated serving, so a policy can be split across two tiers behind a router.

Builds on #3654 (per-service node placement, now on main): prefill and decode need node pools of their own.

Why

A large policy served for an agentic benchmark spends most of its time in prefill. Splitting prefill and decode onto separate node tiers, with KV cache transferred between them and a router in front, is how that workload is served in practice today. SubmitConfig had no way to say it.

The change

One new service type, vllm_pd. It nests its two tiers and owns the router and KV settings:

services:
  policy:
    type: vllm_pd
    container: vllm/vllm-openai:latest
    model: /models/checkpoint
    served_model_name: my-model
    port: 8000                     # the router; driver.policy_model names this service
    prefill:
      node_pool: prefill
      tensor_parallel_size: 4
      number_of_instances: 4
    decode:
      node_pool: decode
      tensor_parallel_size: 4
      number_of_instances: 6
    # optional: server_per_node (false), nixl_side_channel_port (5600; decode +1),
    # prefill_policy, decode_policy, kv_connector (NixlConnector),
    # kv_load_failure_policy (fail), request_timeout_secs, log_level
driver:
  policy_model: policy

• prefill: and decode: are plain vLLM service settings. Each inherits container, model and served_model_name.
• Every P/D setting lives on the vllm_pd service. server_per_node applies to both tiers. nixl_side_channel_port is prefill's; decode gets it + 1. Setting either on a tier is refused.
• Each tier gets its KV role automatically.
• Tier ports default to router port + 1 and + 2. Decode's data-parallel port defaults to prefill's + 1. All can be set.

Why vllm_pd: the service is a vLLM prefill/decode deployment, not a general router. A different kind of router can have its own type later.

No logic is duplicated. At validation, the service expands into three internal services: policy-prefill, policy-decode and policy (the router). The existing multi-node rank builder, node placement and health checks run on those unchanged. The resolved config written to the job directory keeps the nested form, so it loads back.

Bugs found by generating a real job script

Before running anything, we generated the job script for SWE-bench Pro's real layout (4 prefill nodes, 6 decode nodes) and read it. That found three bugs, all still guarded:

• Both tiers used the same meeting point. Each tier's vLLM servers now coordinate on their own first node and their own port.
• Decode's health check looked in the wrong place. Each tier and the router are now probed on the node they run on.
• Nothing stopped the two tiers from picking the same ports. Validation refuses a data-parallel port clash and names both tiers. The NIXL ports are derived, so they cannot clash.

Found on a real cluster

A SWE-bench Pro prefill/decode run on a real cluster exposed four more problems. All are fixed:

• One server per node. A tier with server_per_node: true runs an independent vllm serve on every node, and the router lists each one.
• The router listens on 0.0.0.0. The driver and the health probe call it on localhost.
• The router and driver share node 0. This comes from feat(orchestration): let a service claim its own nodes #3654's placement rule.
• driver.policy_host. An agent in a remote sandbox cannot reach the driver's localhost, so the injected policy URL can name a routable host.

Testing

• Same job script as before. One commit renders SWE-bench Pro's layout with the earlier three-service form and pins the output as fixtures. There are two variants: data-parallel tiers and one server per node. The next commit switches to vllm_pd, and the new form renders both fixtures byte for byte. The collector's scrape config is also unchanged.
• Unit tests. Tests cover the expansion, inheritance and port defaults. They also check that the resolved config loads back. Refusal tests cover: a port clash, a tier without a pool or with an unknown pool, a missing tier, a name that clashes with a tier's name, KV fields or server_per_node / nixl_side_channel_port set on a tier, and those fields set on a plain vLLM service.
• Each new test was revert-checked. It fails when the piece it guards is removed.
• Non-P/D rendering is unchanged.

Verified end to end

SWE-bench Pro ran the full 4 prefill + 6 decode shape on a real cluster with this change. That was before the config was reshaped, and the job script is byte-identical. All 731 tasks x 3 = 2193 rollouts finished. Mean reward was 0.533, against a reference of 0.526 from the same checkpoint served another way.

A deployment still needs its fabric environment (UCX_*, NCCL_*) supplied through each tier's env or pre_command. This PR does not presume a set.

[prokotg]

Looks good but I think a change in API might simplify things. Let's discuss

[prokotg]

Thank you!

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/36722816677