[Bugfix] Prevent DiffusionResultPump crash on cancelled futures (#5793)

[anurag12-webster]

What broke

MultiprocDiffusionExecutor._result_pump is the sole reader of the diffusion worker result queue, running as a singleton background thread. Every dispatch branch resolves a shared future with:

if fut is not None and not fut.done():
    fut.set_result(result)  # or fut.set_exception(...)

The not fut.done() check and the set_result()/set_exception() call are not atomic. A consumer awaiting the future via asyncio.wrap_future() under asyncio.wait_for() (an inter-output timeout, or a request abort) can cancel it in the gap between the two. When that happens, set_result()/set_exception() raises concurrent.futures.InvalidStateError, which nothing catches — killing the pump thread permanently. Since the pump is a singleton, the engine then accepts and "runs" jobs that silently never complete (/health stays 200, GPU sits idle).

Repro

concurrent.futures._base.InvalidStateError: CANCELLED: <Future at 0x... state=cancelled>
  File "vllm_omni/diffusion/executor/multiproc_executor.py", line 883, in _result_pump
    fut.set_result(msg.output)
  File ".../concurrent/futures/_base.py", line 544, in set_result
    raise InvalidStateError(...)

Reported originally in #5793 via two real triggers: a worker-side generation error aborting a request, and a client-side inter-output timeout (_ASYNC_OUTPUT_TIMEOUT) firing while the worker was still running. Both land on the same race.

Root cause

Check-then-act race on concurrent.futures.Future state across threads: the pump's done() check and its resolve call aren't atomic, so a concurrent cancel() from the consumer side can land in between.

Fix

Added _try_set_result/_try_set_exception helpers that wrap set_result/set_exception in a narrow except concurrent.futures.InvalidStateError, dropping the late resolution instead of crashing the thread. Applied at every call site in _result_pump that resolves a future pulled from the shared _rpc_futures/_output_futures dicts, plus the same pattern in shutdown()'s cleanup loop for consistency. Left untouched the two set_result calls on freshly-created local Future() objects that aren't shared with any other thread yet — no race possible there.

Testing

Added tests/diffusion/test_multiproc_result_pump.py (CPU-only, no GPU needed): a _RacyFuture subclass whose done() deterministically lies (False) while genuinely cancelled underneath, reproducing the exact race window without relying on real thread timing.

Verified on an RTX 3090 (vLLM 0.26.0, vllm-omni installed per docs):

• With the fix: tests/diffusion/test_multiproc_result_pump.py — 3 passed.
• With the fix reverted: all 3 fail with the exact InvalidStateError traceback from the original report, at the same fut.set_result(msg.output) call site.
• Neighboring suite tests/diffusion/test_multiproc_engine_concurrency.py — 60 passed, no regressions.
• ruff check, ruff format --check, and the full pre-commit hook suite pass on both changed files.

Fixes #5793

[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 appears to belong to: docs/design/module/diffusion/continuous_batching.md.

Module owners: @Isotr0py @princepride

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

[anurag12-webster]

@vllm-omni-review-bot

[anurag12-webster]

@Isotr0py @princepride PTAL

[anurag12-webster]

hi @Isotr0py @princepride can you take a look at this

[anurag12-webster]

rebased onto main. #6023 rewrote _result_pump/_deliver_batch_split and merged before this pr, so the InvalidStateError fix here got dropped along with the old code shape it patched. the original #5793 crash was reproducible on main again.

reapplied the same fix on top of #6023's new code, and added the regression test to tests/diffusion/test_result_pump.py (the test file #6023 introduced) instead of a separate file.

verified on a real gpu: with the fix reverted, all 3 new tests fail with the exact InvalidStateError from the original report, at the 3 call sites in the new code. with the fix applied, all 21 tests in that file pass (3 new + 18 existing, no regressions).

[Gaohan123]

LGTM. Thanks

[anurag12-webster]

thankyouu!! @Gaohan123