[cipheraxat]

Contributor intent (required human answers)

1. What outcome do you want to achieve?

The weekly BAGEL multi-replica e2e test (tests/e2e/online_serving/test_bagel_multi_replicas.py) passes again instead of OOMing.

2. What is the highest-leverage variable for that outcome?

Which physical GPU each headless DiT replica lands on.

3. How will this work move that variable?

The stage-CLI test harness launches one headless process per stage with --omni-dp-size-local N, so devices: "1,2,3" is split one replica per card.

4. Whose customer or user needs does this address, and what are those needs?

Maintainers / CI: a trustworthy weekly multi-replica diffusion signal without false OOM failures.

Purpose

Fixes #6843.

Before: OmniServerStageCli._start_server started one headless process per replica, each
with --replica-id i. serve ignores --replica-id (deprecated in #3569), so each process
ran with omni_dp_size_local=1. get_headless_replica_devices →
split_devices_for_replicas("1,2,3", 1, 1, 1) returned ["1,2,3"], so all three DiT replicas
used local device 0 (physical GPU 1). With gpu_memory_utilization: 0.7 the third replica
OOMs. The OOM line in the issue shows three 22–28 GiB processes on one GPU.

After: the harness starts one headless process per stage and adds
--omni-dp-size-local <num_replicas> when the stage has more than one replica. The supported
headless path then splits the pool to ["1", "2", "3"] and starts one replica per card.
--replica-id is no longer passed.

There is now one harness process per stage, so the per-replica bookkeeping is gone too:
stage_procs and the stage log maps are keyed by stage_id, and stage logs are named
omni_stage_<id>_<port>.log. The only other reader, test_gepard_tts.py, iterates
_stage_log_paths.values() and is unaffected.

Scope: bagel_multi_replicas_4gpu is the only stage-CLI overlay with num_replicas > 1.
test_qwen3_omni_multi_replicas.py uses the single-runtime server and is unaffected.

Docs: no change. docs/design/module/stage_runtime.md owns
stage_engine_startup.py (headless device splitting), which this PR does not change. It only
aligns the test harness with it. No production code changes.

Test Plan

vLLM Version: 0.31.0 (repo pin). See the note on the local environment below.

vLLM-Omni Commit: base b164fea79

pytest tests/helpers/tests/test_runtime_stage_cli.py tests/helpers/tests/test_runtime_startup_timeout.py -q

Regression tests (CPU, core_model + cpu):

• test_stage_cli_builds_headless_replica_cmd: stage 1 command carries
  --omni-dp-size-local 3 and no --replica-id.
• test_stage_cli_omits_dp_size_local_for_single_replica_stage
• test_stage_cli_launches_one_process_per_stage: _start_server launches stage 1 once,
  not once per replica.

Test Result

• Before the fix: the 3 new tests fail (stage 1 launched 3×, --replica-id present).
  After the fix: 22 passed (test_runtime_stage_cli.py + test_runtime_startup_timeout.py).
• split_devices_for_replicas (real source): ("1,2,3", 1, 1, 1) → ['1,2,3'] before;
  ("1,2,3", 3, 1, 1) → ['1', '2', '3'] after.
• pre-commit run --files on the 3 changed files: all hooks pass (ruff, mypy-3.10, test
  marks, SPDX, forbidden imports, typos).

Not run:

• The local runs were on macOS, where vllm 0.31.0 has no wheel, so vllm / vllm_omni imports
  were stubbed for these pure-Python harness tests. The tests have not been run in the real
  Linux CPU-test environment.
• The 4×H100 e2e test_bagel_multi_replicas.py was not run (no hardware). GPU placement is
  verified at source level only. A maintainer rerun of that weekly test would confirm the fix.

AI assistance: Used Claude Code (Claude Opus 5.5) to trace the root cause, write the harness fix and regression tests, and draft this description and the issue comment.
Validation: see Test Result above. CPU harness tests and pre-commit ran locally on macOS with vllm imports stubbed; the 4×H100 e2e was not run.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-17T12:49:46.375662Z	5dd5f2b	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR matches CODEOWNERS paths: /tests/.

Code owners: @NickCao @yenuo26

Routing: @NickCao via CODEOWNERS; @yenuo26 via CODEOWNERS

@cipheraxat, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 28c7d7e67245 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[alex-jw-brooks]

Looks good to me, thanks for the fix :)

Validated on my local machine that on main only 2/4 GPUs are used (which would push over capacity of an h100), and that on this branch things do indeed get distributed over the 4 GPUs.