Fix vLLM DP mode: one process per node instead of per GPU

[alec-flowers]

Summary

• vLLM handles spawning per-GPU DP engine cores internally, so srtslurm should launch one container per node with all GPUs visible
• The previous per-GPU approach launched separate containers per GPU, causing NVSHMEM IPC failures (DeepEP low-latency backend requires shared memory between GPU processes on the same node)
• endpoints_to_processes(): creates one Process per node with all GPUs instead of one Process per GPU
• build_worker_command(): replaces --data-parallel-rank with --data-parallel-size-local and --data-parallel-start-rank, adds --headless for non-leader nodes

Test plan

• All 285 existing tests pass
• Updated DP mode tests to verify per-node process creation and new CLI flags
• Deploy DeepSeek R1 recipe on EOS cluster to verify NVSHMEM IPC works with single-container-per-node

🤖 Generated with Claude Code

Summary by CodeRabbit

• Refactor

  • Changed distributed training spawning from per-GPU to per-node in DP+EP setups; per-node processes now manage all local GPUs.

• Tests

  • Updated distributed training tests to expect one process per node and adjusted per-process GPU counts and rank calculations.

• Documentation

  • Clarified DP+EP behavior and multi-node coordination in user-facing docs and comments.

• Chores

  • Job scripts now prefer a pre-staged config in the output directory when present and standardize config path passed to the orchestrator.

[coderabbitai]

📝 Walkthrough

Walkthrough

vLLM backend now spawns one process per node for DP+EP (vLLM spawns per-GPU DP engines inside each node process). Process gains per-process kv_events_port and nixl_port; command flags and config handling updated to use per-node DP sizing/start-rank. Job template copy behavior was guarded.

Changes

Cohort / File(s)	Summary
vLLM backend & process model src/srtctl/backends/vllm.py	Switched DP+EP spawning from per-GPU to one process per node; removed per-GPU process loop; added kv_events_port and nixl_port to Process; replaced per-GPU --data-parallel-rank usage with per-node --data-parallel-size-local and --data-parallel-start-rank; popped related config keys to avoid duplication.
Tests updated for per-node semantics tests/test_configs.py	Updated DP tests to expect one process per node (GPU count per process equals node GPU count); adjusted node_rank expectations and CLI assertions (removed per-GPU --data-parallel-rank, added per-node size/start-rank checks); preserved TP-mode checks.
Job template change src/srtctl/templates/job_script_minimal.j2	Guard config copy into OUTPUT_DIR/config.yaml (skip if already present) and update invocation to pass OUTPUT_DIR/config.yaml to the sweep entrypoint.

Sequence Diagram(s)

sequenceDiagram
    participant Orchestrator
    participant NodeProcess as "Node Process\n(vLLM master)"
    participant vLLM as "vLLM DP Engines\n(spawned per-GPU)"
    participant GPU

    Orchestrator->>NodeProcess: start process (includes kv_events_port, nixl_port, --data-parallel-size-local, --data-parallel-start-rank)
    NodeProcess->>vLLM: invoke vLLM runtime (internal spawn)
    vLLM->>GPU: spawn DP engine cores per GPU
    GPU-->>vLLM: report readiness/events via kv_events_port/nixl_port
    vLLM-->>NodeProcess: aggregate status
    NodeProcess-->>Orchestrator: process ready

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐰 I hopped from rank to node with glee,

One process tends each GPU tree,
vLLM spawns its cores inside,
KV and nixl ports open wide,
A nodal hop — compact and spry!

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically summarizes the main architectural change: replacing per-GPU process creation with per-node process creation in vLLM DP mode, which is the primary focus across the modified files.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch fix-vllm-dp-per-node

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/srtctl/backends/vllm.py (1)

174-181: ⚠️ Potential issue | 🟡 Minor

Stale docstring — still says "each GPU runs its own process".

This contradicts the new per-node model. The docstring in endpoints_to_processes was correctly updated (line 195–197), but this one was missed.

Proposed fix

     def _is_dp_mode(self, mode: WorkerMode) -> bool:
         """Check if this mode uses Data Parallel + Expert Parallel pattern.
 
-        DP+EP mode is detected when data-parallel-size is set in the mode's config.
-        In this mode, each GPU runs its own process (rather than TP across GPUs).
+        DP+EP mode is detected when data-parallel-size is set in the mode's config.
+        In this mode, vLLM spawns per-GPU DP engine cores internally within each
+        per-node process (rather than TP across GPUs).
         """

🤖 Fix all issues with AI agents

In `@src/srtctl/backends/vllm.py`:
- Line 322: The current assignment to dp_rpc_port pops "data-parallel-rpc-port"
and then uses or to fall back to popping "data_parallel_rpc_port", which can
leave the second key in config and produce duplicate CLI flags; change the logic
in the same block in _config_to_cli_args (where dp_rpc_port is set) to
explicitly pop both keys separately (e.g., first pop "data-parallel-rpc-port"
with a default None, then pop "data_parallel_rpc_port" with a default None and
use the first non-None value) so that both variants are removed from config and
only a single --data-parallel-rpc-port flag is emitted.

🧹 Nitpick comments (1)

tests/test_configs.py (1)

1002-1013: Fragile value assertions — "0" and "8" match unintended positions in cmd0.

assert "0" in cmd0 passes trivially because "10.0.0.1" is not "0", but e.g. "0" could appear as part of other flag values. More critically, if --data-parallel-start-rank ever changes from "0" to another value, this assertion wouldn't catch it precisely.

The non-leader test on lines 1026–1027 uses the better index-based pattern. Consider doing the same here for --data-parallel-size-local and --data-parallel-start-rank.

Proposed fix for precise assertions

         assert "--data-parallel-size-local" in cmd0
-        assert "8" in cmd0  # 8 GPUs per node
-        assert "--data-parallel-start-rank" in cmd0
-        assert "0" in cmd0  # start_rank = 0 * 8 = 0
+        idx_local = cmd0.index("--data-parallel-size-local")
+        assert cmd0[idx_local + 1] == "8"  # 8 GPUs per node
+        idx_start = cmd0.index("--data-parallel-start-rank")
+        assert cmd0[idx_start + 1] == "0"  # start_rank = 0 * 8 = 0