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

Signed-off-by: ANURAG KANADE <anuragkanade6@gmail.com>

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"]

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()

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")