From 68bf8c2c4296fb851c36173e5dcdca55ea56500c Mon Sep 17 00:00:00 2001 From: Anurag <75563673+anurag12-webster@users.noreply.github.com> Date: Sat, 15 Aug 2026 21:49:45 +0530 Subject: [PATCH] [Bugfix] Prevent DiffusionResultPump crash on cancelled futures (#5793) (#5983) Signed-off-by: ANURAG KANADE --- tests/diffusion/test_result_pump.py | 74 +++++++++++++++++++ .../diffusion/executor/multiproc_executor.py | 15 ++-- vllm_omni/diffusion/utils/future_utils.py | 27 +++++++ 3 files changed, 109 insertions(+), 7 deletions(-) create mode 100644 vllm_omni/diffusion/utils/future_utils.py diff --git a/tests/diffusion/test_result_pump.py b/tests/diffusion/test_result_pump.py index 2f4a89ca1..fa172fd0c 100644 --- a/tests/diffusion/test_result_pump.py +++ b/tests/diffusion/test_result_pump.py @@ -418,3 +418,77 @@ class TestShutdownCleansUpFutures: assert len(executor._rpc_futures) == 0 assert len(executor._output_futures) == 0 + + +class _RacyFuture(concurrent.futures.Future): + # done() always lies and reports False, to deterministically force the + # pump's check-then-act race window without needing real thread timing. + def done(self) -> bool: + return False + + +def _racy_cancelled_future() -> concurrent.futures.Future: + fut = _RacyFuture() + fut.cancel() + assert fut.cancelled() + return fut + + +class TestResultPumpCancelledFutureRace: + """Regression for #5793: a future cancelled concurrently with the pump's + resolve call (e.g. an asyncio.wait_for timeout, or a request abort) must + be dropped, not raise InvalidStateError and kill the pump thread. + """ + + def test_compute_done_racing_cancel_does_not_crash_pump(self): + executor = _make_executor() + fut = _racy_cancelled_future() + with executor._futures_lock: + executor._rpc_futures["1"] = fut + + msg = AsyncDiffusionOutput(kind=AsyncOutputKind.COMPUTE_DONE, rpc_id="1", async_output_id="abc") + _feed_one_msg_to_pump(executor, msg) + + assert fut.cancelled() + + def test_output_ready_racing_cancel_does_not_crash_pump(self): + executor = _make_executor() + fut = _racy_cancelled_future() + with executor._futures_lock: + executor._output_futures["abc123"] = fut + + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id="abc123", + output=DiffusionOutput(output="late"), + ) + _feed_one_msg_to_pump(executor, msg) + + assert fut.cancelled() + + def test_batch_split_racing_cancel_does_not_crash_pump(self): + executor = _make_executor() + cancelled_fut = _racy_cancelled_future() + healthy_fut = concurrent.futures.Future() + with executor._futures_lock: + executor._output_futures["batch-1/r-aborted"] = cancelled_fut + executor._output_futures["batch-1/r-healthy"] = healthy_fut + executor._batch_split_map["batch-1"] = { + "batch-1/r-aborted": "r-aborted", + "batch-1/r-healthy": "r-healthy", + } + + outputs = { + "r-aborted": DiffusionOutput(output="late"), + "r-healthy": DiffusionOutput(output="ok"), + } + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id="batch-1", + output=_FakeBatchOutput(outputs), + ) + _feed_one_msg_to_pump(executor, msg) + + assert cancelled_fut.cancelled() + assert healthy_fut.done() + assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"] diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 74286dc94..6f0917e32 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -24,6 +24,7 @@ from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, Asy from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key +from vllm_omni.diffusion.utils.future_utils import try_set_exception, try_set_result from vllm_omni.diffusion.worker import WorkerProc if TYPE_CHECKING: @@ -870,9 +871,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): fut = self._rpc_futures.pop(msg.rpc_id, None) if msg.rpc_id else None if fut is not None and not fut.done(): if msg.error: - fut.set_exception(RuntimeError(msg.error)) + try_set_exception(fut, RuntimeError(msg.error)) else: - fut.set_result(msg) + try_set_result(fut, msg) elif msg.kind == AsyncOutputKind.OUTPUT_READY: batch_id = msg.async_output_id with self._futures_lock: @@ -903,9 +904,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): pending = self._output_futures.pop(batch_id, None) if pending is not None and not pending.done(): if exc is not None: - pending.set_exception(exc) + try_set_exception(pending, exc) else: - pending.set_result(output_result) + try_set_result(pending, output_result) else: fut = concurrent.futures.Future() if exc is not None: @@ -933,7 +934,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): with self._futures_lock: pending = self._output_futures.pop(per_req_id, None) if pending is not None and not pending.done(): - pending.set_result(per_req_result) + try_set_result(pending, per_req_result) else: fut: concurrent.futures.Future = concurrent.futures.Future() fut.set_result(per_req_result) @@ -999,10 +1000,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): with self._futures_lock: for fut in self._rpc_futures.values(): if not fut.done(): - fut.set_exception(RuntimeError("Executor shut down")) + try_set_exception(fut, RuntimeError("Executor shut down")) for fut in self._output_futures.values(): if not fut.done(): - fut.set_exception(RuntimeError("Executor shut down")) + try_set_exception(fut, RuntimeError("Executor shut down")) self._rpc_futures.clear() self._output_futures.clear() self._batch_split_map.clear() diff --git a/vllm_omni/diffusion/utils/future_utils.py b/vllm_omni/diffusion/utils/future_utils.py new file mode 100644 index 000000000..5492c6a21 --- /dev/null +++ b/vllm_omni/diffusion/utils/future_utils.py @@ -0,0 +1,27 @@ +"""Helpers for resolving ``concurrent.futures.Future`` objects shared across threads.""" + +import concurrent.futures +from typing import Any, TypeVar + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_ResultT = TypeVar("_ResultT") + + +def try_set_result(fut: concurrent.futures.Future[_ResultT], result: _ResultT) -> None: + # fut may be cancelled concurrently (e.g. asyncio.wait_for timeout) between + # the caller's fut.done() check and this call; drop the late result instead + # of crashing the pump thread. + try: + fut.set_result(result) + except concurrent.futures.InvalidStateError: + logger.debug("Dropping late result for already-resolved/cancelled future") + + +def try_set_exception(fut: concurrent.futures.Future[Any], exc: BaseException) -> None: + try: + fut.set_exception(exc) + except concurrent.futures.InvalidStateError: + logger.debug("Dropping late exception for already-resolved/cancelled future")