[Bugfix][Diffusion] Drop aborted request's orphaned async output (fix #6413) (#6439)

Signed-off-by: summer <128961079+zhang-keliang@users.noreply.github.com>
Co-authored-by: multica-agent <github@multica.ai>

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -399,3 +399,36 @@ def test_fail_engine_does_not_reenter_when_executor_shutdown_fails() -> None:
     assert engine._closed is True
     assert engine._shutting_down is True
     assert engine._shutdown_complete is False
+
+
+def test_finalize_aborted_request_drops_pending_async_output() -> None:
+    """An aborted request with a pending async output must be drained (#6413).
+
+    The aborted branch of _finalize_finished_request returns before the
+    async_output_id branch, so without draining, the worker's late OUTPUT_READY
+    is cached in the executor's _completed_outputs forever. The engine must tell
+    the executor to drop that output id.
+    """
+    engine = _make_engine()
+    engine.executor = SimpleNamespace(drop_output=Mock())
+    request_id = engine.scheduler.add_request(_make_request("aborted-async"))
+    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+    runner_output = SimpleNamespace(result=None, async_output_id="aid-async-1")
+    output = engine._finalize_finished_request(request_id, runner_output)
+
+    assert output.aborted is True
+    engine.executor.drop_output.assert_called_once_with("aid-async-1")
+
+
+def test_finalize_aborted_request_without_async_output_skips_drop() -> None:
+    """No async output pending -> nothing to drop; drop_output not called."""
+    engine = _make_engine()
+    engine.executor = SimpleNamespace(drop_output=Mock())
+    request_id = engine.scheduler.add_request(_make_request("aborted-plain"))
+    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+    output = engine._finalize_finished_request(request_id)
+
+    assert output.aborted is True
+    engine.executor.drop_output.assert_not_called()

tests/diffusion/test_diffusion_ipc.py

@@ -6,6 +6,7 @@ import multiprocessing as mp
 import queue
 import threading
 import time
+from collections import OrderedDict
 from multiprocessing import shared_memory
 from multiprocessing.connection import Connection
 
@@ -108,6 +109,7 @@ def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> No
         executor._rpc_futures = {}
         executor._output_futures = {}
         executor._completed_outputs = {}
+        executor._dropped_output_ids = OrderedDict()
         executor._batch_split_map = {}
         executor._start_result_pump()
 

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -8,6 +8,7 @@ import signal
 import threading
 import time
 import weakref
+from collections import OrderedDict
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock
 
@@ -77,6 +78,8 @@ def _make_executor(num_gpus: int = 1):
     executor._processes = []
     executor._is_failed = False
     executor._failure_callbacks = []
+    executor._completed_outputs = {}
+    executor._dropped_output_ids = OrderedDict()
     return executor, req_q, res_q
 
 
@@ -1551,6 +1554,8 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._rpc_futures = {}
         executor._output_futures = {}
         executor._batch_split_map = {}
+        executor._completed_outputs = {}
+        executor._dropped_output_ids = OrderedDict()
 
         proc = _make_short_lived_process()
         executor._processes = [proc]

tests/diffusion/test_multiproc_shutdown_race.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MultiprocDiffusionExecutor shutdown / delayed-pump race.
+
+Reproduces hsliuustc0106's #6439 review finding: `shutdown()` joins each result
+pump thread with a 2s timeout and only warns if it survives. A pump blocked
+inside `unpack_diffusion_output_shm` can resume AFTER `_completed_outputs` was
+cleared and repopulate the dictionary, reintroducing the leak this PR was meant
+to close.
+
+Both single-output and batch-split code paths must discard deliveries once
+`_closed` is set; `drop_output()` must be a no-op after shutdown. Discarding
+must NOT skip SHM unpack: `unpack_diffusion_output_shm()` is the only
+receive-side path that unlinks named SHM segments, so an OUTPUT_READY dequeued
+after shutdown still has to be unpacked (just never cached).
+"""
+
+import queue
+import threading
+import time
+from collections import OrderedDict
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.data import (
+    AsyncDiffusionOutput,
+    AsyncOutputKind,
+    DiffusionOutput,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_executor():
+    """Minimal MultiprocDiffusionExecutor-like object; mirrors test_result_pump helper."""
+    from vllm_omni.diffusion.executor.multiproc_executor import (
+        MultiprocDiffusionExecutor,
+    )
+
+    od_config = MagicMock()
+    od_config.step_execution = False
+
+    executor = object.__new__(MultiprocDiffusionExecutor)
+    executor.od_config = od_config
+    executor._rpc_id_counter = 0
+    executor._rpc_id_lock = threading.Lock()
+    executor._rpc_futures = {}
+    executor._output_futures = {}
+    executor._completed_outputs = {}
+    executor._dropped_output_ids = OrderedDict()
+    executor._batch_split_map = {}
+    executor._futures_lock = threading.RLock()
+    executor._pump_running = False
+    executor._pump_stop = threading.Event()
+    executor._sync_result_buffer = queue.Queue()
+    executor._result_mq = MagicMock()
+    executor._result_mqs = []
+    executor._broadcast_mq = MagicMock()
+    executor._closed = False
+    executor._is_failed = False
+    executor._finalizer = MagicMock()
+    executor._shutdown_cleaner = None
+    executor._processes = []
+    return executor
+
+
+class _FakeBatchOutput:
+    """Batch-level output exposing per-request results (mirrors test_result_pump helper)."""
+
+    def __init__(self, results):
+        self._results = results
+
+    def get_request_output(self, req_id):
+        result = self._results.get(req_id)
+        if result is None:
+            return None
+        return MagicMock(result=result)
+
+
+def _run_pump_with_one_msg(executor, msg):
+    """Drive the pump thread with one *msg*.
+
+    After the single msg, dequeue blocks until ``_pump_stop`` is set so the
+    pump loops out cleanly. Returns the pump thread. The caller finishes the
+    pump with ``_finish_pump``; the message's delivery completes before the
+    second dequeue call, so joining the thread is a deterministic barrier for
+    "delivery handled".
+    """
+    call_count = [0]
+
+    def mock_dequeue(timeout=None):
+        call_count[0] += 1
+        if call_count[0] == 1:
+            return msg
+        while not executor._pump_stop.is_set():
+            time.sleep(0.02)
+        raise TimeoutError
+
+    executor._result_mq.dequeue = mock_dequeue
+    t = threading.Thread(target=executor._result_pump, daemon=True)
+    t.start()
+    return t
+
+
+def _finish_pump(executor, pump_thread):
+    executor._pump_stop.set()
+    pump_thread.join(timeout=5.0)
+    assert not pump_thread.is_alive(), "pump thread did not exit"
+
+
+class TestResultPumpDelayedAfterShutdown:
+    """Reproduce hsliuustc0106's finding: a pump blocked inside SHM unpack can
+    resume after shutdown clears `_completed_outputs` and repopulate it.
+
+    Without gating, the assertions here fail (leak reintroduced through the
+    late-write path). With gating on `_closed`, deliveries after shutdown are
+    discarded. `entered_unpack` proves the pump is inside the mocked unpack
+    BEFORE `_closed` is set, so the post-unpack race path is really exercised
+    (a fixed sleep could let the pump take the early closed gate instead).
+    """
+
+    def test_result_pump_delayed_after_shutdown_single(self, mocker):
+        """Single-output pump path: late unpack must not repopulate the dict."""
+        executor = _make_executor()
+
+        entered_unpack = threading.Event()
+        unpack_gate = threading.Event()
+
+        def _blocked_unpack(*_args, **_kwargs):
+            entered_unpack.set()
+            unpack_gate.wait(timeout=10.0)
+
+        mocker.patch(
+            "vllm_omni.diffusion.executor.multiproc_executor.unpack_diffusion_output_shm",
+            side_effect=_blocked_unpack,
+        )
+
+        async_output_id = "abc-single"
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            output=DiffusionOutput(output="late"),
+        )
+
+        pump_thread = _run_pump_with_one_msg(executor, msg)
+
+        # The pump must be INSIDE unpack (past the early closed gate) before
+        # shutdown starts, otherwise the post-unpack race is not exercised.
+        assert entered_unpack.wait(timeout=5.0), "pump never reached SHM unpack"
+
+        # Simulate shutdown() clearing state while the pump is still stuck
+        # inside unpack. `_closed = True` mirrors shutdown()'s first line.
+        executor._closed = True
+        with executor._futures_lock:
+            executor._completed_outputs.clear()
+
+        # Release the pump: it resumes past unpack and tries to write into
+        # _completed_outputs. Without the gate, this reintroduces the leak.
+        unpack_gate.set()
+
+        # Joining the pump is a deterministic barrier: the delivery is fully
+        # handled before the pump blocks on the second dequeue and exits.
+        _finish_pump(executor, pump_thread)
+
+        # The core assertion: the delayed pump must not have repopulated
+        # `_completed_outputs` after shutdown cleared it.
+        assert executor._completed_outputs == {}, (
+            f"delayed pump repopulated dict after shutdown: {list(executor._completed_outputs)}"
+        )
+
+    def test_result_pump_delayed_after_shutdown_batch_split(self, mocker):
+        """Batch-split pump path: late unpack must not repopulate the dict."""
+        executor = _make_executor()
+
+        entered_unpack = threading.Event()
+        unpack_gate = threading.Event()
+
+        def _blocked_unpack(*_args, **_kwargs):
+            entered_unpack.set()
+            unpack_gate.wait(timeout=10.0)
+
+        mocker.patch(
+            "vllm_omni.diffusion.executor.multiproc_executor.unpack_diffusion_output_shm",
+            side_effect=_blocked_unpack,
+        )
+
+        batch_id = "batch-delayed"
+        # Pre-populate batch split map so pump goes down the batch path,
+        # WITHOUT registering per-request waiters (they are the ones aborted).
+        with executor._futures_lock:
+            executor._batch_split_map[batch_id] = {
+                f"{batch_id}/r-a": "r-a",
+                f"{batch_id}/r-b": "r-b",
+            }
+
+        outputs = {
+            "r-a": DiffusionOutput(output="late-a"),
+            "r-b": DiffusionOutput(output="late-b"),
+        }
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=batch_id,
+            output=_FakeBatchOutput(outputs),
+        )
+
+        pump_thread = _run_pump_with_one_msg(executor, msg)
+
+        assert entered_unpack.wait(timeout=5.0), "pump never reached SHM unpack"
+
+        # Simulate shutdown during unpack.
+        executor._closed = True
+        with executor._futures_lock:
+            executor._completed_outputs.clear()
+
+        # Release the pump so it enters _deliver_batch_split with no waiters
+        # registered; without gating, both per-request results land in
+        # _completed_outputs.
+        unpack_gate.set()
+
+        _finish_pump(executor, pump_thread)
+
+        assert executor._completed_outputs == {}, (
+            f"delayed batch-split pump repopulated dict after shutdown: {list(executor._completed_outputs)}"
+        )
+
+    def test_output_ready_after_shutdown_still_unpacks_shm(self, mocker):
+        """An OUTPUT_READY dequeued after `_closed` is set must still go
+        through `unpack_diffusion_output_shm` (the only receive-side path that
+        unlinks named SHM segments) while writing nothing into the caches.
+        Discarding it before unpack would leak the segments in /dev/shm for
+        the lifetime of the parent process.
+        """
+        executor = _make_executor()
+
+        unpack_calls = []
+        mocker.patch(
+            "vllm_omni.diffusion.executor.multiproc_executor.unpack_diffusion_output_shm",
+            side_effect=lambda output: unpack_calls.append(output),
+        )
+
+        # Shutdown completed BEFORE the pump dequeues the message.
+        executor._closed = True
+        with executor._futures_lock:
+            executor._completed_outputs.clear()
+
+        output = DiffusionOutput(output="late-shm")
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="abc-shm",
+            output=output,
+        )
+
+        pump_thread = _run_pump_with_one_msg(executor, msg)
+        _finish_pump(executor, pump_thread)
+
+        assert unpack_calls == [output], "closed-time discard skipped SHM unpack/unlink"
+        assert executor._completed_outputs == {}
+        assert executor._output_futures == {}
+
+
+class TestDropOutputAfterShutdown:
+    """`drop_output()` runs from the abort path and must be a no-op once the
+    executor has entered shutdown, so it does not race the pump into an
+    already-torn-down `_output_futures` / `_completed_outputs`.
+    """
+
+    def test_drop_output_after_shutdown_is_noop(self):
+        executor = _make_executor()
+        executor._closed = True
+
+        # Before the gate, `drop_output` would register a placeholder waiter
+        # in `_output_futures`. After shutdown, that would leak past teardown.
+        executor.drop_output("abc-abort")
+
+        assert executor._output_futures == {}, (
+            f"drop_output registered waiter after shutdown: {list(executor._output_futures)}"
+        )
+        assert executor._completed_outputs == {}

tests/diffusion/test_result_pump.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for MultiprocDiffusionExecutor async result pump and wait_output_ready."""
 
 import concurrent.futures
 import queue
 import threading
 import time
+from collections import OrderedDict
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -30,6 +31,7 @@ def _make_executor(step_execution=False):
     executor._rpc_futures = {}
     executor._output_futures = {}
     executor._completed_outputs = {}
+    executor._dropped_output_ids = OrderedDict()
     executor._batch_split_map = {}
     executor._futures_lock = threading.RLock()
     executor._pump_running = False
@@ -465,6 +467,7 @@ class TestResultPumpCancelledFutureRace:
         _feed_one_msg_to_pump(executor, msg)
 
         assert fut.cancelled()
+        assert executor._completed_outputs == {}
 
     def test_batch_split_racing_cancel_does_not_crash_pump(self):
         executor = _make_executor()
@@ -492,3 +495,446 @@ class TestResultPumpCancelledFutureRace:
         assert cancelled_fut.cancelled()
         assert healthy_fut.done()
         assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"]
+        assert executor._completed_outputs == {}
+
+
+class TestResultPumpDiscardsDoneFuture:
+    """A registered future that is already cancelled/done must have its late
+    OUTPUT_READY result discarded, not re-cached into _completed_outputs.
+
+    rahul-steiger-nv (#6439 review): the old `else` branch cached the result
+    whenever `pending.done()` was true, so a genuinely cancelled waiter (the
+    common abort case) leaked the late result back into the cache. The pump
+    must distinguish "no waiter registered" (cache) from "waiter registered
+    but done/cancelled" (discard).
+    """
+
+    def test_single_output_discards_for_cancelled_future(self):
+        executor = _make_executor()
+        fut: concurrent.futures.Future = concurrent.futures.Future()
+        fut.cancel()
+        assert fut.cancelled() and fut.done()
+        with executor._futures_lock:
+            executor._output_futures["aid-x"] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="aid-x",
+            output=DiffusionOutput(output="late"),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert executor._completed_outputs == {}
+        assert "aid-x" not in executor._output_futures
+
+    def test_batch_split_discards_for_cancelled_member(self):
+        executor = _make_executor()
+        cancelled_fut: concurrent.futures.Future = concurrent.futures.Future()
+        cancelled_fut.cancel()
+        assert cancelled_fut.cancelled() and cancelled_fut.done()
+        healthy_fut: concurrent.futures.Future = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._output_futures["batch-9/r-aborted"] = cancelled_fut
+            executor._output_futures["batch-9/r-healthy"] = healthy_fut
+            executor._batch_split_map["batch-9"] = {
+                "batch-9/r-aborted": "r-aborted",
+                "batch-9/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-9",
+            output=_FakeBatchOutput(outputs),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        # Aborted member's result is discarded, healthy member still delivered.
+        assert healthy_fut.done()
+        assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"]
+        assert executor._completed_outputs == {}
+
+
+class TestDropOutput:
+    """drop_output drains an aborted request's async output so it cannot leak.
+
+    An aborted request never calls wait_output_ready, so a late OUTPUT_READY is
+    unpacked and cached in _completed_outputs forever (issue #6413). drop_output
+    evicts it, handling both arrival orderings.
+    """
+
+    def test_drops_already_cached_output(self):
+        executor = _make_executor()
+        fut = concurrent.futures.Future()
+        fut.set_result(DiffusionOutput(output="leaked"))
+        with executor._futures_lock:
+            executor._completed_outputs["aid-1"] = fut
+
+        executor.drop_output("aid-1")
+
+        with executor._futures_lock:
+            assert "aid-1" not in executor._completed_outputs
+
+    def test_prevents_caching_when_output_not_yet_arrived(self):
+        executor = _make_executor()
+
+        # Abort finalized before OUTPUT_READY: register a placeholder waiter.
+        executor.drop_output("aid-2")
+
+        # The late result now resolves into the placeholder instead of caching.
+        output = DiffusionOutput(output="late")
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="aid-2",
+            output=output,
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        with executor._futures_lock:
+            assert "aid-2" not in executor._completed_outputs
+            assert "aid-2" not in executor._output_futures
+
+    def test_leaves_real_waiter_untouched(self):
+        executor = _make_executor()
+        real = executor.wait_output_ready("aid-3")
+
+        executor.drop_output("aid-3")
+
+        with executor._futures_lock:
+            # A genuine waiter still drains via the normal path.
+            assert executor._output_futures.get("aid-3") is real
+
+    def test_drop_then_pump_then_late_wait_fails_fast_without_caching(self):
+        """drop_output → OUTPUT_READY → late wait_output_ready.
+
+        The aborted output's tensors must NOT be retained anywhere, and the
+        late waiter must get a completed (failed) Future rather than a fresh
+        one that never resolves — the verl-omni fully-async abort ordering.
+        """
+        executor = _make_executor()
+
+        executor.drop_output("aid-dpw")
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="aid-dpw",
+            output=DiffusionOutput(output="discarded"),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        with executor._futures_lock:
+            assert "aid-dpw" not in executor._completed_outputs
+            assert "aid-dpw" not in executor._output_futures
+
+        fut = executor.wait_output_ready("aid-dpw")
+        assert fut.done(), "late wait after drop returned a never-completing future"
+        with pytest.raises(RuntimeError, match="was dropped"):
+            fut.result(timeout=0)
+        # Terminal: the dropped-id LRU keeps the id so any later wait_output_ready
+        # on the same id also fails fast instead of hanging on a fresh Future.
+        with executor._futures_lock:
+            assert "aid-dpw" in executor._dropped_output_ids
+        second = executor.wait_output_ready("aid-dpw")
+        assert second.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            second.result(timeout=0)
+
+    def test_drop_then_wait_share_placeholder_and_fail_on_delivery(self):
+        """drop_output → wait_output_ready (before pump) share one Future, and
+        delivery terminates it with the dropped error instead of the tensors."""
+        executor = _make_executor()
+
+        executor.drop_output("aid-dw")
+        fut = executor.wait_output_ready("aid-dw")
+        with executor._futures_lock:
+            assert executor._output_futures.get("aid-dw") is fut
+        assert not fut.done()
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="aid-dw",
+            output=DiffusionOutput(output="discarded"),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            fut.result(timeout=0)
+        with executor._futures_lock:
+            assert "aid-dw" not in executor._completed_outputs
+            assert "aid-dw" not in executor._output_futures
+
+    def test_drop_after_cached_then_late_wait_fails_fast(self):
+        """Result already cached → drop_output → late wait must not hang."""
+        executor = _make_executor()
+        fut = concurrent.futures.Future()
+        fut.set_result(DiffusionOutput(output="cached"))
+        with executor._futures_lock:
+            executor._completed_outputs["aid-cd"] = fut
+
+        executor.drop_output("aid-cd")
+        late = executor.wait_output_ready("aid-cd")
+        assert late.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            late.result(timeout=0)
+
+    def test_dropped_ids_memory_is_bounded(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as mpe
+
+        cap = 8
+        monkeypatch.setattr(mpe, "_DROPPED_OUTPUT_IDS_MAX", cap)
+        executor = _make_executor()
+        n = cap + 5
+        for i in range(n):
+            executor.drop_output(f"aid-b{i}")
+            executor._pump_stop.clear()
+            _feed_one_msg_to_pump(
+                executor,
+                AsyncDiffusionOutput(
+                    kind=AsyncOutputKind.OUTPUT_READY,
+                    async_output_id=f"aid-b{i}",
+                    output=DiffusionOutput(output=i),
+                ),
+            )
+        with executor._futures_lock:
+            assert len(executor._dropped_output_ids) == cap
+            # Oldest evicted, newest kept; no tensors retained anywhere.
+            assert "aid-b0" not in executor._dropped_output_ids
+            assert f"aid-b{n - 1}" in executor._dropped_output_ids
+            assert executor._completed_outputs == {}
+            assert executor._output_futures == {}
+
+    def test_cancelled_waiter_then_pump_then_late_wait_fails_fast_twice(self):
+        """Fully-async abort overlap: step_streaming registers a real waiter,
+        the async abort cancels it, then the pump delivers OUTPUT_READY.
+
+        Reproduces SamitHuang's 09-14 P1 chain:
+          * _finish_output line 1029 discarded a cancelled waiter without
+            _remember_dropped → later wait_output_ready allocated a fresh
+            Future and hung.
+          * wait_output_ready line 1097 returned the cancelled existing future
+            as if it were live → abort-after-cancel never installed dropped
+            state.
+          * wait_output_ready line 1101 deleted the LRU entry on the first
+            failed wait → a second late wait on the same id missed the record
+            and hung.
+        """
+        executor = _make_executor()
+
+        # step_streaming registers a real waiter.
+        real = executor.wait_output_ready("aid-cx")
+        assert not real.done()
+
+        # asyncio.wrap_future(...).cancel() marks the underlying
+        # concurrent.futures.Future cancelled.
+        assert real.cancel()
+        assert real.cancelled()
+        assert real.done()
+
+        # Pump delivers OUTPUT_READY for the same id after the cancel.
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="aid-cx",
+            output=DiffusionOutput(output="discarded-after-cancel"),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        with executor._futures_lock:
+            # Tensors are not re-cached; the cancelled waiter is gone.
+            assert "aid-cx" not in executor._completed_outputs
+            assert "aid-cx" not in executor._output_futures
+            # The id was recorded as dropped so late waits fail fast.
+            assert "aid-cx" in executor._dropped_output_ids
+
+        # Late wait returns an already-failed Future (not a fresh never-completing one).
+        late = executor.wait_output_ready("aid-cx")
+        assert late.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            late.result(timeout=0)
+
+        # A second late wait must ALSO fail fast: the LRU entry must persist.
+        late2 = executor.wait_output_ready("aid-cx")
+        assert late2.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            late2.result(timeout=0)
+        with executor._futures_lock:
+            assert "aid-cx" in executor._dropped_output_ids
+            assert "aid-cx" not in executor._output_futures
+            assert "aid-cx" not in executor._completed_outputs
+
+    def test_wait_output_ready_replaces_stale_cancelled_existing_future(self):
+        """A cancelled non-placeholder Future in _output_futures is not a live
+        waiter: the next wait_output_ready must evict it and install dropped
+        state rather than returning the cancelled future.
+
+        Reproduces SamitHuang's wait_output_ready:1135 finding: without this,
+        abort-after-cancel skips _remember_dropped entirely and downstream
+        pump delivery falls into the discard branch that used to lose the id.
+        """
+        executor = _make_executor()
+
+        stale = executor.wait_output_ready("aid-stale")
+        assert stale.cancel()
+        assert stale.done() and stale.cancelled()
+
+        replacement = executor.wait_output_ready("aid-stale")
+        assert replacement is not stale
+        assert replacement.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            replacement.result(timeout=0)
+        with executor._futures_lock:
+            assert "aid-stale" not in executor._output_futures
+            assert "aid-stale" in executor._dropped_output_ids
+
+    def test_pump_after_cancelled_waiter_then_wait_stays_dropped(self):
+        """Timeline-B: the cancelled-waiter has already been observed by a
+        second wait_output_ready and evicted into the dropped-id LRU, THEN the
+        pump delivers OUTPUT_READY.
+
+        Without a dropped-id guard in _finish_output, the pump's ``pending``
+        is ``None`` (the second wait already popped the stale entry), so the
+        no-waiter branch caches the tensors in ``_completed_outputs``. A
+        subsequent wait_output_ready then pops that cache and returns a
+        successful Future — violating the "every late wait on a dropped id
+        fails the same way" contract from wait_output_ready.
+        """
+        executor = _make_executor()
+
+        # (1) step_streaming registers a real waiter.
+        real = executor.wait_output_ready("aid-tb")
+        assert not real.done()
+
+        # (2) asyncio cancels it.
+        assert real.cancel()
+
+        # (3) Second wait_output_ready observes the stale entry, evicts it into
+        # the dropped-id LRU and returns an already-failed Future.
+        first_late = executor.wait_output_ready("aid-tb")
+        assert first_late.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            first_late.result(timeout=0)
+        with executor._futures_lock:
+            assert "aid-tb" not in executor._output_futures
+            assert "aid-tb" in executor._dropped_output_ids
+
+        # (4) Pump delivers OUTPUT_READY for the dropped id AFTER the LRU
+        # record is already in place.
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="aid-tb",
+            output=DiffusionOutput(output="late-after-evict"),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        # (5) The tensors must NOT be cached, and the id must remain in the LRU.
+        with executor._futures_lock:
+            assert "aid-tb" not in executor._completed_outputs
+            assert "aid-tb" not in executor._output_futures
+            assert "aid-tb" in executor._dropped_output_ids
+
+        # (6) A subsequent wait_output_ready must still fail fast, NOT pop a
+        # successful cached Future.
+        second_late = executor.wait_output_ready("aid-tb")
+        assert second_late.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            second_late.result(timeout=0)
+
+    def test_cancel_between_done_check_and_set_result_records_dropped(self):
+        """Race: consumer cancels the Future between the pump's ``.done()``
+        check and the actual ``set_result``.
+
+        Reproduces Gaohan123's 09-19 P2: ``try_set_result`` catches the
+        ``InvalidStateError`` silently, but the waiter has already been
+        popped from ``_output_futures`` by ``_finish_output``. Without
+        recording the id as dropped, a subsequent ``wait_output_ready``
+        allocates a fresh Future that can never complete.
+        """
+
+        class _RaceCancellingFuture(concurrent.futures.Future):
+            """Future that mimics a consumer cancelling in the tiny window
+            between ``pending.done()`` and ``pending.set_result(...)``."""
+
+            def set_result(self, result):
+                self.cancel()  # race lands here — mid-branch cancel
+                super().set_result(result)  # raises InvalidStateError
+
+            def set_exception(self, exc):
+                self.cancel()
+                super().set_exception(exc)
+
+        executor = _make_executor()
+        racy = _RaceCancellingFuture()
+        with executor._futures_lock:
+            executor._output_futures["aid-race"] = racy
+
+        # Simulate the pump delivery: at entry .done() is False, then the
+        # concurrent cancel lands, then set_result raises InvalidStateError.
+        with executor._futures_lock:
+            executor._finish_output("aid-race", DiffusionOutput(output="lost"), None)
+
+        assert racy.cancelled()
+        with executor._futures_lock:
+            # The waiter has been popped and the id must be recorded as dropped.
+            assert "aid-race" not in executor._output_futures
+            assert "aid-race" not in executor._completed_outputs
+            assert "aid-race" in executor._dropped_output_ids
+
+        # Late wait_output_ready must fail fast, not allocate a fresh Future.
+        late = executor.wait_output_ready("aid-race")
+        assert late.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            late.result(timeout=0)
+        # And a second late wait must also fail fast (LRU stays populated).
+        late2 = executor.wait_output_ready("aid-race")
+        assert late2.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            late2.result(timeout=0)
+
+    def test_batch_split_drop_placeholder_discards_member(self):
+        """_deliver_batch_split applies the same contract: a dropped member is
+        discarded (never cached), a live member is resolved directly."""
+        executor = _make_executor()
+        live = executor.wait_output_ready("batch-d/r-live")
+        executor.drop_output("batch-d/r-dropped")
+        with executor._futures_lock:
+            executor._batch_split_map["batch-d"] = {
+                "batch-d/r-dropped": "r-dropped",
+                "batch-d/r-live": "r-live",
+            }
+        outputs = {
+            "r-dropped": DiffusionOutput(output="discarded"),
+            "r-live": DiffusionOutput(output="ok"),
+        }
+        _feed_one_msg_to_pump(
+            executor,
+            AsyncDiffusionOutput(
+                kind=AsyncOutputKind.OUTPUT_READY,
+                async_output_id="batch-d",
+                output=_FakeBatchOutput(outputs),
+            ),
+        )
+        assert live.result(timeout=1.0) is outputs["r-live"]
+        with executor._futures_lock:
+            assert executor._completed_outputs == {}
+            assert executor._output_futures == {}
+        late = executor.wait_output_ready("batch-d/r-dropped")
+        assert late.done()
+        with pytest.raises(RuntimeError, match="was dropped"):
+            late.result(timeout=0)
+
+
+class TestShutdownClearsCompletedOutputs:
+    """shutdown() must release cached async outputs (issue #6413)."""
+
+    def test_shutdown_clears_completed_outputs(self):
+        executor = _make_executor()
+        fut = concurrent.futures.Future()
+        fut.set_result(DiffusionOutput(output="cached"))
+        with executor._futures_lock:
+            executor._completed_outputs["aid-4"] = fut
+
+        executor.shutdown()
+
+        assert executor._completed_outputs == {}

vllm_omni/diffusion/diffusion_engine.py

@@ -1609,6 +1609,11 @@ class DiffusionEngine:
             raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")
 
         if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
+            # An aborted request is never waited on, so a pending async output
+            # would be cached forever by the executor result pump (issue #6413).
+            # Tell the executor to drop it before returning the aborted result.
+            if runner_output is not None and runner_output.async_output_id is not None:
+                self.executor.drop_output(runner_output.async_output_id)
             # Preserve runner-provided abort details when available.
             if runner_output is not None and runner_output.result is not None and runner_output.result.aborted:
                 return runner_output.result

vllm_omni/diffusion/executor/abstract.py

@@ -137,6 +137,16 @@ class DiffusionExecutor(ABC):
         """Resolve deferred output; only asynchronous executors implement this."""
         raise NotImplementedError(f"{type(self).__name__} does not support asynchronous output")
 
+    def drop_output(self, async_output_id: str) -> None:
+        """Reclaim an async output that will never be waited on (e.g. an
+        aborted request).
+
+        Only executors with an async output path (result pump) cache outputs
+        that a consumer must later claim; executors without one have nothing to
+        reclaim and can keep the default no-op implementation.
+        """
+        return None
+
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
         """Collect rank-local native specs after every Worker loads its model."""
 

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -12,6 +12,7 @@ import queue
 import threading
 import time
 import weakref
+from collections import OrderedDict
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from multiprocessing.synchronize import Event
@@ -47,6 +48,24 @@ _WORKER_SHUTDOWN_GRACE_S = 15.0
 _WORKER_TERMINATE_GRACE_S = 5.0
 _WORKER_KILL_GRACE_S = 5.0
 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0
+# Upper bound on remembered dropped async_output_ids (see drop_output).
+_DROPPED_OUTPUT_IDS_MAX = 4096
+
+
+class _DropPlaceholder(concurrent.futures.Future):
+    """Placeholder registered by :meth:`MultiprocDiffusionExecutor.drop_output`.
+
+    Distinguishes "abort said nobody wants this output" from a genuine
+    ``wait_output_ready`` waiter, so the result pump can discard the tensors
+    for the former while still resolving the latter directly.
+    """
+
+
+def _dropped_output_error(async_output_id: str) -> RuntimeError:
+    return RuntimeError(
+        f"async output {async_output_id} was dropped: the request was aborted "
+        "before its output was claimed; retry with a new request."
+    )
 
 
 def _is_empty_dp_prompt(prompt: object) -> bool:
@@ -190,6 +209,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._output_futures: dict[str, concurrent.futures.Future[DiffusionOutput]] = {}
         self._completed_outputs: dict[str, concurrent.futures.Future[DiffusionOutput]] = {}
         self._batch_split_map: dict[str, dict[str, str]] = {}  # batch_id -> {per_req_id: request_id}
+        # Bounded memory of ids drained by drop_output() so a late
+        # wait_output_ready() fails fast instead of hanging on a new Future.
+        self._dropped_output_ids: OrderedDict[str, None] = OrderedDict()
         self._futures_lock = threading.RLock()
         self._pump_running = False
         self._pump_stop = threading.Event()
@@ -922,6 +944,16 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 self._sync_result_buffer.put(msg)
                 continue
 
+            # If shutdown started while we were dequeuing, drop this delivery
+            # so it cannot repopulate _completed_outputs after shutdown()
+            # cleared it (issue #6413 / #6439 review). OUTPUT_READY must still
+            # flow through the dispatch below: unpack_diffusion_output_shm()
+            # is the only receive-side path that unlinks named SHM segments,
+            # and the closed-time re-checks under _futures_lock already
+            # prevent any cache write after unpack.
+            if self._closed and msg.kind != AsyncOutputKind.OUTPUT_READY:
+                continue
+
             if msg.kind in (AsyncOutputKind.RPC_RESULT, AsyncOutputKind.COMPUTE_DONE):
                 with self._futures_lock:
                     fut = self._rpc_futures.pop(msg.rpc_id, None) if msg.rpc_id else None
@@ -957,19 +989,82 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
                     if batch_id:
                         with self._futures_lock:
-                            pending = self._output_futures.pop(batch_id, None)
-                            if pending is not None and not pending.done():
-                                if exc is not None:
-                                    try_set_exception(pending, exc)
-                                else:
-                                    try_set_result(pending, output_result)
-                            else:
-                                fut = concurrent.futures.Future()
-                                if exc is not None:
-                                    fut.set_exception(exc)
-                                else:
-                                    fut.set_result(output_result)
-                                self._completed_outputs[batch_id] = fut
+                            if self._closed:
+                                # shutdown() cleared _completed_outputs while
+                                # we were unpacking; drop this delivery.
+                                continue
+                            self._finish_output(batch_id, output_result, exc)
+
+    def _finish_output(
+        self,
+        async_output_id: str,
+        result: DiffusionOutput | None,
+        exc: BaseException | None,
+    ) -> None:
+        """Hand one delivered async output to its waiter. Caller holds ``_futures_lock``.
+
+        * id already recorded as dropped -> discard the tensors: a
+          ``wait_output_ready`` observed a cancelled waiter and evicted it
+          into the dropped-id LRU before this delivery, so caching would
+          break the "every late wait on a dropped id fails the same way"
+          contract;
+        * no waiter registered -> cache for a later ``wait_output_ready()``;
+        * :class:`_DropPlaceholder` (request aborted) -> discard the tensors:
+          terminate the placeholder with an error so a caller that already
+          reused it via ``wait_output_ready()`` wakes up, and remember the id
+          (bounded) so a later wait fails fast instead of hanging;
+        * genuine pending waiter -> resolve it directly, never cache;
+        * waiter already cancelled/done -> discard the tensors and record the
+          id as dropped so a later wait fails fast instead of hanging.
+        """
+        pending = self._output_futures.pop(async_output_id, None)
+        if pending is None:
+            if async_output_id in self._dropped_output_ids:
+                # Timeline-B: wait_output_ready has already observed a
+                # cancelled/done stale entry and moved this id into the
+                # dropped-id LRU. Late OUTPUT_READY must not repopulate
+                # _completed_outputs — a subsequent wait_output_ready would
+                # otherwise pop a successful Future and violate the dropped
+                # contract.
+                self._dropped_output_ids.move_to_end(async_output_id)
+                return
+            fut: concurrent.futures.Future = concurrent.futures.Future()
+            if exc is not None:
+                fut.set_exception(exc)
+            else:
+                fut.set_result(result)
+            self._completed_outputs[async_output_id] = fut
+        elif isinstance(pending, _DropPlaceholder):
+            try_set_exception(pending, _dropped_output_error(async_output_id))
+            self._remember_dropped(async_output_id)
+        elif not pending.done():
+            # Concurrent cancellation can still race between ``.done()`` and
+            # ``set_result/set_exception``. The helpers swallow the resulting
+            # ``InvalidStateError`` and report ``False`` so we can still
+            # record the id as dropped — the waiter was popped from
+            # ``_output_futures`` above, so without this a subsequent
+            # ``wait_output_ready`` would allocate a fresh Future and hang.
+            if exc is not None:
+                delivered = try_set_exception(pending, exc)
+            else:
+                delivered = try_set_result(pending, result)
+            if not delivered:
+                self._remember_dropped(async_output_id)
+        else:
+            # Waiter already cancelled or resolved. Do not re-cache the delivered
+            # tensors, but remember the id so a later ``wait_output_ready`` on
+            # the same id fails fast instead of allocating a fresh Future that
+            # would never complete (fully-async abort overlap: ``step_streaming``
+            # took a live waiter, ``asyncio.wrap_future`` cancelled it, then
+            # ``OUTPUT_READY`` lands here).
+            self._remember_dropped(async_output_id)
+
+    def _remember_dropped(self, async_output_id: str) -> None:
+        dropped = self._dropped_output_ids
+        dropped[async_output_id] = None
+        dropped.move_to_end(async_output_id)
+        while len(dropped) > _DROPPED_OUTPUT_IDS_MAX:
+            dropped.popitem(last=False)
 
     def _deliver_batch_split(
         self,
@@ -978,6 +1073,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         error: str | None = None,
     ) -> None:
         """Resolve per-request futures from one batch-level output."""
+        if self._closed:
+            # shutdown() has taken over; do not touch _completed_outputs.
+            return
         for per_req_id, req_id in per_req_map.items():
             req_output = batch_output.get_request_output(req_id) if batch_output is not None else None
             per_req_result: DiffusionOutput
@@ -988,13 +1086,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             else:
                 per_req_result = DiffusionOutput(error="No output result for batch request")
             with self._futures_lock:
-                pending = self._output_futures.pop(per_req_id, None)
-                if pending is not None and not pending.done():
-                    try_set_result(pending, per_req_result)
-                else:
-                    fut: concurrent.futures.Future = concurrent.futures.Future()
-                    fut.set_result(per_req_result)
-                    self._completed_outputs[per_req_id] = fut
+                if self._closed:
+                    # Belt-and-braces: re-check under the lock so a shutdown
+                    # that started mid-loop cannot repopulate the cleared dict.
+                    return
+                self._finish_output(per_req_id, per_req_result, None)
 
     def describe_pending_state(self, async_output_id: str | None = None) -> str:
         """Summarize async-output bookkeeping for diagnosing stuck waits."""
@@ -1017,15 +1113,70 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             return str(self._rpc_id_counter)
 
     def wait_output_ready(self, async_output_id: str) -> concurrent.futures.Future[DiffusionOutput]:
-        """Return a Future that resolves when the async output is ready."""
+        """Return a Future that resolves when the async output is ready.
+
+        After :meth:`drop_output` has drained an id, the output is gone for
+        good: a wait on that id returns an already-failed Future instead of a
+        fresh one that would never complete.
+        """
         with self._futures_lock:
             cached = self._completed_outputs.pop(async_output_id, None)
             if cached is not None:
                 return cached
+            # Share an already-registered Future (a genuine pending waiter or a
+            # drop_output placeholder) rather than clobbering it. But if the
+            # registered entry is already cancelled or resolved, it is not a
+            # live waiter — evict it, remember the id, and fall through to the
+            # dropped path so abort-after-cancel installs terminal state
+            # instead of returning a Future that never completes.
+            existing = self._output_futures.get(async_output_id)
+            if existing is not None:
+                if isinstance(existing, _DropPlaceholder) or not existing.done():
+                    return existing
+                self._output_futures.pop(async_output_id, None)
+                self._remember_dropped(async_output_id)
+            if async_output_id in self._dropped_output_ids:
+                # Keep the id in the bounded LRU so *every* late wait on a
+                # dropped id fails the same way. Deleting here would make
+                # fail-fast one-shot and the next wait would hang on a fresh
+                # Future.
+                self._dropped_output_ids.move_to_end(async_output_id)
+                dropped: concurrent.futures.Future = concurrent.futures.Future()
+                dropped.set_exception(_dropped_output_error(async_output_id))
+                return dropped
             fut: concurrent.futures.Future = concurrent.futures.Future()
             self._output_futures[async_output_id] = fut
         return fut
 
+    def drop_output(self, async_output_id: str) -> None:
+        """Discard an async output that will never be waited on.
+
+        An aborted request never calls :meth:`wait_output_ready`, so a late
+        ``OUTPUT_READY`` would be unpacked and cached in ``_completed_outputs``
+        forever (issue #6413). Draining it here keeps engine-process memory
+        bounded under abort traffic. Handles both arrival orderings:
+
+        * result already arrived -> pop and drop the cached future;
+        * result not yet arrived -> register a :class:`_DropPlaceholder` so the
+          pump discards the tensors instead of caching them.
+
+        Either way the id is remembered (bounded) so a late
+        :meth:`wait_output_ready` fails fast rather than hanging. A genuine
+        waiter that is already registered is left untouched so it can drain
+        through the normal path.
+        """
+        with self._futures_lock:
+            if self._closed:
+                # Executor is torn down; abort path must not repopulate the
+                # cleared state (issue #6413 / #6439 review).
+                return
+            if self._completed_outputs.pop(async_output_id, None) is not None:
+                self._remember_dropped(async_output_id)
+                return
+            if async_output_id in self._output_futures:
+                return
+            self._output_futures[async_output_id] = _DropPlaceholder()
+
     def check_health(self) -> None:
         if self._is_failed:
             raise EngineDeadError()
@@ -1067,6 +1218,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 self._rpc_futures.clear()
                 self._output_futures.clear()
                 self._batch_split_map.clear()
+                # Cached async outputs hold unpacked tensors; drop them so they
+                # do not survive shutdown (issue #6413).
+                self._completed_outputs.clear()
+                self._dropped_output_ids.clear()
             self._processes = (cleaner.processes or []) if cleaner is not None else []
             if not self._processes:
                 self._shutdown_cleaner = None

vllm_omni/diffusion/utils/future_utils.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Helpers for resolving ``concurrent.futures.Future`` objects shared across threads."""
 
 import concurrent.futures
@@ -10,18 +13,24 @@ logger = init_logger(__name__)
 _ResultT = TypeVar("_ResultT")
 
 
-def try_set_result(fut: concurrent.futures.Future[_ResultT], result: _ResultT) -> None:
+def try_set_result(fut: concurrent.futures.Future[_ResultT], result: _ResultT) -> bool:
     # 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.
+    # of crashing the pump thread. Return whether the future was actually
+    # resolved so callers that need to mirror the terminal state elsewhere
+    # (e.g. record the id as dropped) can react to a lost delivery.
     try:
         fut.set_result(result)
     except concurrent.futures.InvalidStateError:
         logger.debug("Dropping late result for already-resolved/cancelled future")
+        return False
+    return True
 
 
-def try_set_exception(fut: concurrent.futures.Future[Any], exc: BaseException) -> None:
+def try_set_exception(fut: concurrent.futures.Future[Any], exc: BaseException) -> bool:
     try:
         fut.set_exception(exc)
     except concurrent.futures.InvalidStateError:
         logger.debug("Dropping late exception for already-resolved/cancelled future")
+        return False
+    return True