[codex] Fix diffusion worker timeout for large broadcast payloads

[hotTeaFun]

Summary

Make the diffusion worker wait indefinitely for each complete shared-memory
broadcast message instead of abandoning a receive after the hard-coded
1-second timeout.

This fixes long video/diffusion requests whose RPC payloads go through the
vLLM MessageQueue overflow path. Once an overflow marker has been consumed,
the worker must wait for the corresponding socket payload; retrying
dequeue after a timeout can lose or reorder requests.

Fixes #4844.

Root cause

For large video transfer requests, the serialized RPC request can exceed 2 GB.
The current worker loop calls:

self.mq.dequeue(timeout=1.0)

If the overflow socket message is not ready within 1 second, dequeue raises,
the worker loop catches the exception and continues, and the request is lost
from the worker's perspective. The API server then waits for a result that never
comes back.

Change

Add WorkerProc.recv_message(), which calls
self.mq.dequeue(indefinite=True), and use it from the worker busy loop so a
consumed overflow marker is always paired with its complete multipart payload.

Add regression coverage for indefinite receive and for preserving request order
when an overflow payload is delayed.

Validation

Ran:

PYTHONPYCACHEPREFIX=/private/tmp/pycache python3 -m py_compile \
  vllm_omni/diffusion/worker/diffusion_worker.py

Runtime validation on a Cosmos3 V2V/transfer request:

• Before: 189-frame request created a 2,081,495,018-byte broadcast payload;
  all workers timed out at recv_poll_start timeout_ms=1000.
• After setting the effective timeout to 30 seconds: all workers reached
  recv_poll_ready, received the 2.08 GB multipart payload, entered
  pipeline_forward_start, and the request returned HTTP 200 with 189 frames.

[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.

[MaciejBalaNV]

@david6666666 @hsliuustc0106
Can we please merge it? It also impacts Cosmos models with v2v requests.

[david6666666]

Can we specify it without modifying environment variables?

[hotTeaFun]

Can we specify it without modifying environment variables?

Yes. This requires a slightly broader change because the timeout needs to flow through the configuration stack instead of an environment variable. I’ll add it as an OmniDiffusionConfig field, wire it through the stage config/CLI, and have the worker consume it directly. The default remains 30 seconds, with -1 meaning an indefinite wait. Would this approach be okay with you?

[Copilot]

Pull request overview

This PR makes the diffusion worker’s shared-memory broadcast receive loop more robust for multi-GB diffusion/video RPC payloads by replacing the hard-coded 1s MessageQueue.dequeue timeout with a configurable, longer default.

Changes:

• Add VLLM_OMNI_DIFFUSION_WORKER_BROADCAST_TIMEOUT_S to configure the diffusion worker’s broadcast dequeue timeout (default 30s, with fallback behavior for invalid values).
• Log the effective broadcast dequeue timeout at worker startup.
• Use the configured timeout value in the worker busy loop’s MessageQueue.dequeue(...) call.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[hsliuustc0106]

At vllm_omni/diffusion/worker/diffusion_worker.py:923, increasing the finite timeout only moves the loss boundary. vLLM acknowledges and advances the SHM overflow marker before polling the socket; after a timeout this loop calls dequeue again, so a late multipart can be paired with the next overflow marker. I reproduced TimeoutError -> inline-2 -> stale overflow-1 by delaying the first overflow send by 100 ms with a 20 ms timeout. Please make overflow receipt lossless (for example, wait indefinitely here or fix MessageQueue upstream) and add a regression test.

[hotTeaFun]

At vllm_omni/diffusion/worker/diffusion_worker.py:923, increasing the finite timeout only moves the loss boundary. vLLM acknowledges and advances the SHM overflow marker before polling the socket; after a timeout this loop calls dequeue again, so a late multipart can be paired with the next overflow marker. I reproduced TimeoutError -> inline-2 -> stale overflow-1 by delaying the first overflow send by 100 ms with a 20 ms timeout. Please make overflow receipt lossless (for example, wait indefinitely here or fix MessageQueue upstream) and add a regression test.

Thanks for pointing this out. I’ve switched to indefinite waiting and added a regression test for delayed overflow payloads. Does this look okay?

[david6666666]

The indefinite receive plus delayed-overflow regression test addresses the ordering loss. Looks good now.

[MaciejBalaNV]

@david6666666 Can we merge then?

[MaciejBalaNV]

@david6666666 @hsliuustc0106
Bumping up for the merge

[hsliuustc0106]

https://buildkite.com/vllm/vllm-omni/builds/13418/canvas a few ci failed

[hotTeaFun]

Collaborator

@hsliuustc0106 Fixed in 99bc28e. Could you please retrigger CI? Thanks!

[hsliuustc0106]

This regression doesn't bite on the pre-fix code: the delayed release is 0.1s (line 155), well inside the old dequeue(timeout=1.0) bound, so main passes this test as-is. Also _BROADCAST_DEQUEUE_TIMEOUT_S doesn't exist on either main or this branch — raising=False patches an attribute nothing reads (leftover from the env-var revision?).

Since the SHM marker is consumed before the socket read, a finite timeout only loses the message when the delay exceeds it. Can you drop the monkeypatch and extend the delay past the old 1s bound (e.g. 1.5s) so reintroducing a finite timeout actually fails this test? Good timing to fold it in before the Buildkite retrigger — the new test has no green run on 8b07c43 yet.

[hsliuustc0106]

The PR description still documents the earlier env-var revision — VLLM_OMNI_DIFFUSION_WORKER_BROADCAST_TIMEOUT_S, 30s default, invalid-value fallback, startup log — none of which is in the final diff. Can you update the Change section to match the indefinite dequeue? The root-cause and validation sections are still accurate.

[hotTeaFun]

The PR description still documents the earlier env-var revision — VLLM_OMNI_DIFFUSION_WORKER_BROADCAST_TIMEOUT_S, 30s default, invalid-value fallback, startup log — none of which is in the final diff. Can you update the Change section to match the indefinite dequeue? The root-cause and validation sections are still accurate.

Thanks, updated the PR description.