From 7f867fb782f709a92776432ed47acaf3e4d5fec3 Mon Sep 17 00:00:00 2001 From: summer <128961079+zhang-keliang@users.noreply.github.com> Date: Thu, 24 Sep 2026 00:58:46 +0800 Subject: [PATCH] [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 --- .../test_diffusion_engine_cleanup.py | 33 ++ tests/diffusion/test_diffusion_ipc.py | 2 + .../test_multiproc_engine_concurrency.py | 5 + .../diffusion/test_multiproc_shutdown_race.py | 279 +++++++++++ tests/diffusion/test_result_pump.py | 448 +++++++++++++++++- vllm_omni/diffusion/diffusion_engine.py | 5 + vllm_omni/diffusion/executor/abstract.py | 10 + .../diffusion/executor/multiproc_executor.py | 197 +++++++- vllm_omni/diffusion/utils/future_utils.py | 15 +- 9 files changed, 969 insertions(+), 25 deletions(-) create mode 100644 tests/diffusion/test_multiproc_shutdown_race.py diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index eea749571..0d593fe73 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/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() diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py index aec062057..6caba1311 100644 --- a/tests/diffusion/test_diffusion_ipc.py +++ b/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() diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 50fca3ea8..22318f9f7 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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] diff --git a/tests/diffusion/test_multiproc_shutdown_race.py b/tests/diffusion/test_multiproc_shutdown_race.py new file mode 100644 index 000000000..e891e0678 --- /dev/null +++ b/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 == {} diff --git a/tests/diffusion/test_result_pump.py b/tests/diffusion/test_result_pump.py index fa172fd0c..2cadc8a3c 100644 --- a/tests/diffusion/test_result_pump.py +++ b/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 == {} diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 2f38fe10f..a8385d7c6 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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 diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index 296cb6de5..4bbef2f27 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/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.""" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 57b23656c..c58617499 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/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 diff --git a/vllm_omni/diffusion/utils/future_utils.py b/vllm_omni/diffusion/utils/future_utils.py index 5492c6a21..028c73c86 100644 --- a/vllm_omni/diffusion/utils/future_utils.py +++ b/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