From 7982b5c339f16bd5d0346bec884084e449775326 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 24 Aug 2026 08:45:22 -0700 Subject: [PATCH 01/14] initial commit Signed-off-by: Rahul Steiger --- .../test_diffusion_engine_cleanup.py | 100 +++++++++++++++++- vllm_omni/diffusion/diffusion_engine.py | 69 ++++++++---- vllm_omni/diffusion/executor/abstract.py | 3 +- 3 files changed, 147 insertions(+), 25 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index 3f138e72107..c3f3b5fd712 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio +import concurrent.futures import queue import threading from types import SimpleNamespace @@ -15,6 +16,7 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler +from vllm_omni.diffusion.worker.utils import RunnerOutput from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -33,10 +35,15 @@ def _make_engine() -> DiffusionEngine: engine = DiffusionEngine.__new__(DiffusionEngine) engine.scheduler = RequestScheduler() engine.scheduler.initialize(SimpleNamespace()) - engine.executor = SimpleNamespace(shutdown=Mock()) + engine.executor = SimpleNamespace( + drop_output=Mock(), + shutdown=Mock(), + wait_output_ready=Mock(), + ) engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) engine._out_streams = {} + engine._unclaimed_async_outputs = {} engine._closed = False engine._shutting_down = False engine._shutdown_complete = False @@ -64,6 +71,81 @@ def test_close_completes_pending_output_streams() -> None: event_loop.close() +def test_put_output_discards_async_output_when_stream_is_missing() -> None: + engine = _make_engine() + + engine._put_output("abandoned", DiffusionOutput(async_output_id="aid-missing")) + + engine.executor.drop_output.assert_called_once_with("aid-missing") + + +@pytest.mark.asyncio +async def test_closing_stream_discards_unclaimed_async_output() -> None: + engine = _make_engine() + request_id = "abandoned" + engine._out_streams[request_id] = asyncio.Queue() + engine._unclaimed_async_outputs[request_id] = {"aid-queued"} + stream = engine.get_output_stream(request_id) + next_output = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + + next_output.cancel() + with pytest.raises(asyncio.CancelledError): + await next_output + + engine.executor.drop_output.assert_called_once_with("aid-queued") + assert request_id not in engine._out_streams + assert request_id not in engine._unclaimed_async_outputs + + +@pytest.mark.asyncio +async def test_async_output_is_claimed_after_materialization() -> None: + engine = _make_engine() + request_id = "claimed" + pending_output = DiffusionOutput(async_output_id="aid-claimed") + materialized_output = DiffusionOutput(output="materialized") + ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + ready.set_result(materialized_output) + engine.executor.wait_output_ready.return_value = ready + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + output_queue.put_nowait(pending_output) + engine._out_streams[request_id] = output_queue + engine._unclaimed_async_outputs[request_id] = {"aid-claimed"} + stream = engine.get_output_stream(request_id) + + assert await anext(stream) is materialized_output + assert request_id not in engine._unclaimed_async_outputs + engine.executor.wait_output_ready.assert_called_once_with("aid-claimed") + + await stream.aclose() + engine.executor.drop_output.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancelling_materialization_discards_async_output() -> None: + engine = _make_engine() + request_id = "cancelled" + pending_output = DiffusionOutput(async_output_id="aid-cancelled") + ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + engine.executor.wait_output_ready.return_value = ready + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + output_queue.put_nowait(pending_output) + engine._out_streams[request_id] = output_queue + engine._unclaimed_async_outputs[request_id] = {"aid-cancelled"} + stream = engine.get_output_stream(request_id) + next_output = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + + next_output.cancel() + with pytest.raises(asyncio.CancelledError): + await next_output + + engine.executor.drop_output.assert_called_with("aid-cancelled") + assert ready.cancelled() + assert request_id not in engine._out_streams + assert request_id not in engine._unclaimed_async_outputs + + def test_emit_finished_outputs_finalizes_already_drained_waiter() -> None: class RacingOutQueue(dict): def get(self, key, default=None): @@ -79,6 +161,22 @@ def get(self, key, default=None): assert engine.scheduler.get_request_state(request_id) is None +def test_emit_finished_outputs_discards_async_output_without_stream() -> None: + engine = _make_engine() + request_id = engine.scheduler.add_request(_make_request("completed-async")) + engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_COMPLETED) + runner_output = RunnerOutput( + request_id=request_id, + finished=True, + async_output_id="aid-completed", + ) + + engine._emit_finished_outputs({request_id}, runner_output) + + engine.executor.drop_output.assert_called_once_with("aid-completed") + assert engine.scheduler.get_request_state(request_id) is None + + def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None: engine = _make_engine() engine.execution_mode = DiffusionExecutionMode.STEP_BATCH diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 57dddc88904..54a64577b53 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -409,6 +409,7 @@ def _init_runtime_state(self) -> None: self._rpc_lock = threading.RLock() self._cv = threading.Condition(self._rpc_lock) self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {} + self._unclaimed_async_outputs: dict[str, set[str]] = {} self._closed = False self._shutting_down = False self._shutdown_complete = False @@ -474,25 +475,6 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ generator = self.get_output_stream(request_id) async for output in generator: exec_total_time = time.perf_counter() - exec_start_time - output_ready_wait_time = 0.0 - # Async mode: wait for background D2H/SHM to complete. - if output.async_output_id: - output_ready_wait_start_time = time.perf_counter() - fut = self.executor.wait_output_ready(output.async_output_id) - timeout = _async_output_timeout() - try: - output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) - except (TimeoutError, asyncio.TimeoutError): - describe = getattr(self.executor, "describe_pending_state", None) - logger.error( - "Timed out after %.1fs waiting for async output; set %s to a larger value " - "to allow slower steps. Executor state: %s", - timeout, - _ASYNC_OUTPUT_TIMEOUT_ENV, - describe(output.async_output_id) if describe else "unavailable", - ) - raise - output_ready_wait_time = time.perf_counter() - output_ready_wait_start_time postprocess_start_time = time.perf_counter() scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) try: @@ -506,11 +488,10 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000 logger.debug( "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, " - "add_req_and_wait=%.2f ms, output_ready_wait=%.2f ms, " + "add_req_and_wait=%.2f ms, " "postprocess=%.2f ms, total=%.2f ms", preprocess_time * 1000, exec_total_time * 1000, - output_ready_wait_time * 1000, postprocess_time * 1000, step_total_ms, ) @@ -518,7 +499,6 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ metrics_update = { "preprocess_time_ms": preprocess_time * 1000, "diffusion_engine_exec_time_ms": exec_total_time * 1000, - "output_ready_wait_time_ms": output_ready_wait_time * 1000, "postprocess_time_ms": postprocess_time * 1000, **scheduler_metrics, } @@ -1164,6 +1144,33 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu try: while True: output: DiffusionOutput = await queue.get() + async_output_id = getattr(output, "async_output_id", None) + if async_output_id is not None: + fut = self.executor.wait_output_ready(async_output_id) + timeout = _async_output_timeout() + try: + output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) + except asyncio.CancelledError: + self.executor.drop_output(async_output_id) + raise + except (TimeoutError, asyncio.TimeoutError): + self.executor.drop_output(async_output_id) + describe = getattr(self.executor, "describe_pending_state", None) + logger.error( + "Timed out after %.1fs waiting for async output; set %s to a larger value " + "to allow slower steps. Executor state: %s", + timeout, + _ASYNC_OUTPUT_TIMEOUT_ENV, + describe(async_output_id) if describe else "unavailable", + ) + raise + with self._cv: + pending_outputs = getattr(self, "_unclaimed_async_outputs", None) + pending_ids = pending_outputs.get(request_id) if pending_outputs is not None else None + if pending_ids is not None: + pending_ids.discard(async_output_id) + if not pending_ids and pending_outputs is not None: + pending_outputs.pop(request_id, None) yield output if output.finished: break @@ -1174,6 +1181,10 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu with self._cv: if self._out_streams.get(request_id) is queue: self._out_streams.pop(request_id, None) + pending_outputs = getattr(self, "_unclaimed_async_outputs", None) + abandoned_ids = pending_outputs.pop(request_id, set()) if pending_outputs is not None else set() + for async_output_id in abandoned_ids: + self.executor.drop_output(async_output_id) def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]: request_id = self.add_request(request) @@ -1544,14 +1555,20 @@ def _put_queue_output( queue.put_nowait(output) def _put_output(self, request_id: str, output: DiffusionOutput) -> None: + async_output_id = getattr(output, "async_output_id", None) with self._cv: queue = self._out_streams.get(request_id) + if queue is not None and async_output_id is not None: + self._unclaimed_async_outputs.setdefault(request_id, set()).add(async_output_id) if queue is None: + if async_output_id is not None: + self.executor.drop_output(async_output_id) return self._put_queue_output(queue, output) def close(self) -> None: pending_streams: list[asyncio.Queue[DiffusionOutput]] = [] + abandoned_ids: set[str] = set() with self._cv: if self._closed and self._shutdown_complete: return @@ -1563,8 +1580,16 @@ def close(self) -> None: self.stop_event.set() pending_streams = list(self._out_streams.values()) self._out_streams.clear() + pending_outputs = getattr(self, "_unclaimed_async_outputs", None) + if pending_outputs is not None: + for async_output_ids in pending_outputs.values(): + abandoned_ids.update(async_output_ids) + pending_outputs.clear() self._cv.notify_all() + for async_output_id in abandoned_ids: + self.executor.drop_output(async_output_id) + closed_output = DiffusionOutput(error="DiffusionEngine is closed.") for stream in pending_streams: self._put_queue_output(stream, closed_output) diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index 4bbef2f270b..97d5114eede 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -138,8 +138,7 @@ def wait_output_ready(self, async_output_id: str) -> Future[DiffusionOutput]: 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). + """Reclaim an async output whose consumer is no longer available. Only executors with an async output path (result pump) cache outputs that a consumer must later claim; executors without one have nothing to From 17480a9f34e5c1058066cd88ff9f71ec710380a7 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Tue, 25 Aug 2026 00:55:22 -0700 Subject: [PATCH 02/14] removed unnecesary getattr Signed-off-by: Rahul Steiger --- vllm_omni/diffusion/diffusion_engine.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 54a64577b53..de4ca2e455d 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations @@ -1144,7 +1144,7 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu try: while True: output: DiffusionOutput = await queue.get() - async_output_id = getattr(output, "async_output_id", None) + async_output_id = output.async_output_id if async_output_id is not None: fut = self.executor.wait_output_ready(async_output_id) timeout = _async_output_timeout() @@ -1165,12 +1165,11 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu ) raise with self._cv: - pending_outputs = getattr(self, "_unclaimed_async_outputs", None) - pending_ids = pending_outputs.get(request_id) if pending_outputs is not None else None + pending_ids = self._unclaimed_async_outputs.get(request_id) if pending_ids is not None: pending_ids.discard(async_output_id) - if not pending_ids and pending_outputs is not None: - pending_outputs.pop(request_id, None) + if not pending_ids: + self._unclaimed_async_outputs.pop(request_id, None) yield output if output.finished: break @@ -1181,8 +1180,7 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu with self._cv: if self._out_streams.get(request_id) is queue: self._out_streams.pop(request_id, None) - pending_outputs = getattr(self, "_unclaimed_async_outputs", None) - abandoned_ids = pending_outputs.pop(request_id, set()) if pending_outputs is not None else set() + abandoned_ids = self._unclaimed_async_outputs.pop(request_id, set()) for async_output_id in abandoned_ids: self.executor.drop_output(async_output_id) @@ -1555,7 +1553,7 @@ def _put_queue_output( queue.put_nowait(output) def _put_output(self, request_id: str, output: DiffusionOutput) -> None: - async_output_id = getattr(output, "async_output_id", None) + async_output_id = output.async_output_id with self._cv: queue = self._out_streams.get(request_id) if queue is not None and async_output_id is not None: @@ -1580,11 +1578,9 @@ def close(self) -> None: self.stop_event.set() pending_streams = list(self._out_streams.values()) self._out_streams.clear() - pending_outputs = getattr(self, "_unclaimed_async_outputs", None) - if pending_outputs is not None: - for async_output_ids in pending_outputs.values(): - abandoned_ids.update(async_output_ids) - pending_outputs.clear() + for async_output_ids in self._unclaimed_async_outputs.values(): + abandoned_ids.update(async_output_ids) + self._unclaimed_async_outputs.clear() self._cv.notify_all() for async_output_id in abandoned_ids: From 7e31ad58ac27520b4307202768b7b3f6cd969a5f Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Fri, 28 Aug 2026 04:57:53 -0700 Subject: [PATCH 03/14] fix review comments about edge cases Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine.py | 1 + .../test_diffusion_engine_cleanup.py | 49 +++++++++++++++++++ .../test_diffusion_engine_rpc_routing.py | 10 +--- vllm_omni/diffusion/diffusion_engine.py | 24 ++++++--- 4 files changed, 68 insertions(+), 16 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index b7daa2e95f3..c98e82d0c47 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -903,6 +903,7 @@ async def test_async_add_req_and_stream_response(): engine = object.__new__(DiffusionEngine) engine.scheduler = MockScheduler() engine._out_streams = {} + engine._unclaimed_async_outputs = {} engine.abort_queue: queue.Queue[str] = queue.Queue() engine._rpc_queue = queue.Queue() engine._rpc_lock = threading.RLock() diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index c3f3b5fd712..7ffca35fa8a 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler from vllm_omni.diffusion.worker.utils import RunnerOutput @@ -54,6 +55,15 @@ def _make_engine() -> DiffusionEngine: return engine +def _make_output_executor() -> MultiprocDiffusionExecutor: + executor = MultiprocDiffusionExecutor.__new__(MultiprocDiffusionExecutor) + executor._futures_lock = threading.RLock() + executor._output_futures = {} + executor._completed_outputs = {} + executor._closed = False + return executor + + def test_close_completes_pending_output_streams() -> None: engine = _make_engine() event_loop = asyncio.new_event_loop() @@ -79,6 +89,18 @@ def test_put_output_discards_async_output_when_stream_is_missing() -> None: engine.executor.drop_output.assert_called_once_with("aid-missing") +@pytest.mark.asyncio +async def test_unstarted_stream_does_not_admit_request() -> None: + engine = _make_engine() + request = _make_request("unstarted") + + stream = engine.async_add_req_and_stream_response(request) + await stream.aclose() + + assert engine.scheduler.get_request_state(request.request_id) is None + assert request.request_id not in engine._out_streams + + @pytest.mark.asyncio async def test_closing_stream_discards_unclaimed_async_output() -> None: engine = _make_engine() @@ -121,6 +143,33 @@ async def test_async_output_is_claimed_after_materialization() -> None: engine.executor.drop_output.assert_not_called() +@pytest.mark.asyncio +async def test_exceptional_materialization_retires_async_output() -> None: + engine = _make_engine() + executor = _make_output_executor() + engine.executor = executor + request_id = "failed" + async_output_id = "aid-failed" + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + output_queue.put_nowait(DiffusionOutput(async_output_id=async_output_id)) + engine._out_streams[request_id] = output_queue + engine._unclaimed_async_outputs[request_id] = {async_output_id} + stream = engine.get_output_stream(request_id) + next_output = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + + with executor._futures_lock: + ready = executor._output_futures.pop(async_output_id) + ready.set_exception(RuntimeError("materialization failed")) + + with pytest.raises(RuntimeError, match="materialization failed"): + await next_output + + assert executor._output_futures == {} + assert executor._completed_outputs == {} + assert request_id not in engine._unclaimed_async_outputs + + @pytest.mark.asyncio async def test_cancelling_materialization_discards_async_output() -> None: engine = _make_engine() diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index d26f20c1d18..79e2b04055d 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -178,6 +178,7 @@ def _make_engine_with_loop( engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) engine._out_streams = {} + engine._unclaimed_async_outputs = {} engine._closed = False engine.abort_queue = queue.Queue() engine._rpc_queue = queue.Queue() @@ -521,14 +522,7 @@ def fail_reservation(*args, **kwargs): request.diffusion_kv_requests = (object(),) try: response_stream = engine.async_add_req_and_stream_response(request) - output_queue = engine._out_streams[request.request_id] - for _ in range(300): - if not output_queue.empty(): - break - await asyncio.sleep(0.01) - assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream" - - output = await anext(response_stream) + output = await asyncio.wait_for(anext(response_stream), timeout=3.0) assert output.error == "native allocation bug" assert output.finished assert engine.worker_thread.is_alive() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index de4ca2e455d..b1710ab37ef 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1164,12 +1164,13 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu describe(async_output_id) if describe else "unavailable", ) raise - with self._cv: - pending_ids = self._unclaimed_async_outputs.get(request_id) - if pending_ids is not None: - pending_ids.discard(async_output_id) - if not pending_ids: - self._unclaimed_async_outputs.pop(request_id, None) + finally: + with self._cv: + pending_ids = self._unclaimed_async_outputs.get(request_id) + if pending_ids is not None: + pending_ids.discard(async_output_id) + if not pending_ids: + self._unclaimed_async_outputs.pop(request_id, None) yield output if output.finished: break @@ -1184,9 +1185,16 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu for async_output_id in abandoned_ids: self.executor.drop_output(async_output_id) - def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]: + async def async_add_req_and_stream_response( + self, request: OmniDiffusionRequest + ) -> AsyncGenerator[DiffusionOutput, None]: request_id = self.add_request(request) - return self.get_output_stream(request_id) + stream = self.get_output_stream(request_id) + try: + async for output in stream: + yield output + finally: + await stream.aclose() async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: """Deprecated compatibility wrapper over ``async_add_req_and_stream_response()``. From 8a6395a83875f08623b96e414598335e9d879ffa Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Fri, 28 Aug 2026 07:46:36 -0700 Subject: [PATCH 04/14] fixed pre-commits tests Signed-off-by: Rahul Steiger --- vllm_omni/diffusion/diffusion_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b1710ab37ef..b745a8b007f 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations From deddb9784f27156eb831b24ff01f9a0aa770101d Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Fri, 28 Aug 2026 07:52:36 -0700 Subject: [PATCH 05/14] fix headers to usptream Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine.py | 2 +- tests/diffusion/test_diffusion_engine_rpc_routing.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index c98e82d0c47..ddd16c882b7 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import queue diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 79e2b04055d..8c21ee5d7f8 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Regression tests for the busy-loop-based RPC routing in DiffusionEngine. These tests guard the invariants established when ``collective_rpc`` was From 9b3fcea5f8217828216e046d1c8c82211e1c4aa4 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Fri, 28 Aug 2026 07:54:21 -0700 Subject: [PATCH 06/14] even more header issues Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine_cleanup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index 7ffca35fa8a..c49027c16fa 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import concurrent.futures From 318707d470fa37330f5aad277eb103e766224645 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Thu, 17 Sep 2026 01:14:03 -0700 Subject: [PATCH 07/14] Fix async output cancellation cleanup and admission tests Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine.py | 19 +++++-- .../test_diffusion_engine_cleanup.py | 53 +++++++++++++++++-- .../test_diffusion_engine_rpc_routing.py | 2 +- vllm_omni/diffusion/diffusion_engine.py | 15 ++++-- 4 files changed, 77 insertions(+), 12 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index ddd16c882b7..88e0077e983 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import queue @@ -871,7 +871,8 @@ async def _consume_final_output(generator): @pytest.mark.cpu @pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"]) -def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None: +@pytest.mark.asyncio +async def test_engine_admission_preprocesses_request_once(entrypoint: str, mocker: MockerFixture) -> None: raw_request = OmniDiffusionRequest( prompt="raw", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), @@ -891,7 +892,17 @@ def preprocess(request): engine = _make_admission_engine(preprocess) - getattr(engine, entrypoint)(raw_request) + if entrypoint == "async_add_req_and_stream_response": + output = DiffusionOutput(output="prepared", finished=True) + + async def output_stream(request_id): + assert request_id == prepared_request.request_id + yield output + + mocker.patch.object(engine, "get_output_stream", side_effect=output_stream) + assert [result async for result in engine.async_add_req_and_stream_response(raw_request)] == [output] + else: + engine.add_request(raw_request) assert preprocess_calls == [raw_request] assert engine.scheduler._waiting_queue == [prepared_request] @@ -925,7 +936,7 @@ async def test_async_add_req_and_stream_response(): def _finalize(rid, out, err=None, **kwargs): # Stream consumers stop on ``finished``; keep result_data for assertions. - return SimpleNamespace(result_data=out.result.result_data, finished=True) + return SimpleNamespace(result_data=out.result.result_data, finished=True, async_output_id=None) engine._finalize_finished_request = _finalize diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index c49027c16fa..46deb9227e0 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import concurrent.futures import queue import threading +from collections import OrderedDict from types import SimpleNamespace from unittest.mock import Mock, patch @@ -60,6 +61,7 @@ def _make_output_executor() -> MultiprocDiffusionExecutor: executor._futures_lock = threading.RLock() executor._output_futures = {} executor._completed_outputs = {} + executor._dropped_output_ids = OrderedDict() executor._closed = False return executor @@ -136,6 +138,7 @@ async def test_async_output_is_claimed_after_materialization() -> None: stream = engine.get_output_stream(request_id) assert await anext(stream) is materialized_output + assert materialized_output.stage_durations["output_ready_wait"] >= 0.0 assert request_id not in engine._unclaimed_async_outputs engine.executor.wait_output_ready.assert_called_once_with("aid-claimed") @@ -171,11 +174,14 @@ async def test_exceptional_materialization_retires_async_output() -> None: @pytest.mark.asyncio -async def test_cancelling_materialization_discards_async_output() -> None: +@pytest.mark.parametrize("running", [False, True]) +async def test_cancelling_materialization_discards_async_output(running: bool) -> None: engine = _make_engine() request_id = "cancelled" pending_output = DiffusionOutput(async_output_id="aid-cancelled") ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + if running: + ready.set_running_or_notify_cancel() engine.executor.wait_output_ready.return_value = ready output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() output_queue.put_nowait(pending_output) @@ -189,8 +195,47 @@ async def test_cancelling_materialization_discards_async_output() -> None: with pytest.raises(asyncio.CancelledError): await next_output - engine.executor.drop_output.assert_called_with("aid-cancelled") - assert ready.cancelled() + if running: + engine.executor.drop_output.assert_called_once_with("aid-cancelled") + assert not ready.done() + else: + # The cancelled waiter remains registered until the executor drains it. + engine.executor.drop_output.assert_not_called() + assert ready.cancelled() + assert request_id not in engine._out_streams + assert request_id not in engine._unclaimed_async_outputs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failed", [False, True]) +async def test_cancellation_after_delivery_does_not_recreate_waiter(failed: bool) -> None: + engine = _make_engine() + executor = _make_output_executor() + engine.executor = executor + request_id = "delivered" + async_output_id = "aid-delivered" + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + output_queue.put_nowait(DiffusionOutput(async_output_id=async_output_id)) + engine._out_streams[request_id] = output_queue + engine._unclaimed_async_outputs[request_id] = {async_output_id} + stream = engine.get_output_stream(request_id) + next_output = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + + # Match delivery: remove the waiter and resolve it before cancellation + # reaches the coroutine awaiting its wrapped future. + with executor._futures_lock: + ready = executor._output_futures.pop(async_output_id) + if failed: + ready.set_exception(RuntimeError("materialization failed")) + else: + ready.set_result(DiffusionOutput(output="materialized", finished=True)) + next_output.cancel() + with pytest.raises(asyncio.CancelledError): + await next_output + + assert executor._output_futures == {} + assert executor._completed_outputs == {} assert request_id not in engine._out_streams assert request_id not in engine._unclaimed_async_outputs diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 8c21ee5d7f8..79e2b04055d 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for the busy-loop-based RPC routing in DiffusionEngine. These tests guard the invariants established when ``collective_rpc`` was diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b745a8b007f..68812f0178c 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -475,6 +475,7 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ generator = self.get_output_stream(request_id) async for output in generator: exec_total_time = time.perf_counter() - exec_start_time + output_ready_wait_time = getattr(output, "stage_durations", {}).get("output_ready_wait", 0.0) postprocess_start_time = time.perf_counter() scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) try: @@ -488,10 +489,11 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000 logger.debug( "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, " - "add_req_and_wait=%.2f ms, " + "add_req_and_wait=%.2f ms, output_ready_wait=%.2f ms, " "postprocess=%.2f ms, total=%.2f ms", preprocess_time * 1000, exec_total_time * 1000, + output_ready_wait_time * 1000, postprocess_time * 1000, step_total_ms, ) @@ -499,6 +501,7 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ metrics_update = { "preprocess_time_ms": preprocess_time * 1000, "diffusion_engine_exec_time_ms": exec_total_time * 1000, + "output_ready_wait_time_ms": output_ready_wait_time * 1000, "postprocess_time_ms": postprocess_time * 1000, **scheduler_metrics, } @@ -1148,13 +1151,19 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu if async_output_id is not None: fut = self.executor.wait_output_ready(async_output_id) timeout = _async_output_timeout() + output_ready_wait_start_time = time.perf_counter() try: output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) + output.stage_durations["output_ready_wait"] = time.perf_counter() - output_ready_wait_start_time except asyncio.CancelledError: - self.executor.drop_output(async_output_id) + # Delivery already retires completed futures; dropping + # them again would create an orphaned executor waiter. + if not fut.done(): + self.executor.drop_output(async_output_id) raise except (TimeoutError, asyncio.TimeoutError): - self.executor.drop_output(async_output_id) + if not fut.done(): + self.executor.drop_output(async_output_id) describe = getattr(self.executor, "describe_pending_state", None) logger.error( "Timed out after %.1fs waiting for async output; set %s to a larger value " From 0213b6909b487d2e87a8865f64cf6063b668a060 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Sun, 27 Sep 2026 01:26:31 -0700 Subject: [PATCH 08/14] Fix diffusion execution timing after async output materialization Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine.py | 36 ++++++++++++++++++++++++ vllm_omni/diffusion/diffusion_engine.py | 4 ++- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 88e0077e983..df343e00118 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -869,6 +869,42 @@ async def _consume_final_output(generator): return final_output +@pytest.mark.cpu +@pytest.mark.asyncio +@pytest.mark.parametrize("output_wait", [0.0, 4.0]) +async def test_step_streaming_excludes_output_wait_from_execution_time( + output_wait: float, mocker: MockerFixture +) -> None: + clock = [0.0] + mocker.patch.object(diffusion_engine_module.time, "perf_counter", side_effect=lambda: clock[0]) + stage_durations = {"output_ready_wait": output_wait} if output_wait else {} + output = DiffusionOutput(stage_durations=stage_durations) + formatted_output = SimpleNamespace(metrics={}) + request = SimpleNamespace(scheduler_queue_wait_ms=None) + + async def output_stream(request_id): + # Ten seconds of execution followed by output materialization. + clock[0] = 10.0 + output_wait + yield output + + engine = mocker.Mock( + pre_process_func=None, + _scheduler_num_waiting_reqs=0, + _check_and_start_background_loop=mocker.AsyncMock(), + _prepare_request_for_admission=mocker.Mock(return_value=request), + _add_prepared_request=mocker.Mock(return_value="timed-request"), + get_output_stream=output_stream, + postprocess_output=mocker.Mock(return_value=[formatted_output]), + ) + + results = [batch async for batch in DiffusionEngine.step_streaming(engine, request)] + + assert results == [[formatted_output]] + assert formatted_output.metrics["diffusion_engine_exec_time_ms"] == pytest.approx(10_000.0) + assert formatted_output.metrics["output_ready_wait_time_ms"] == pytest.approx(output_wait * 1000) + assert formatted_output.metrics["postprocess_time_ms"] == 0.0 + + @pytest.mark.cpu @pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"]) @pytest.mark.asyncio diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 68812f0178c..aba6de540df 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -474,8 +474,10 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ request_id = self._add_prepared_request(request) generator = self.get_output_stream(request_id) async for output in generator: - exec_total_time = time.perf_counter() - exec_start_time output_ready_wait_time = getattr(output, "stage_durations", {}).get("output_ready_wait", 0.0) + # The stream now materializes async outputs before yielding. Keep + # that wait separate from execution time in the timing breakdown. + exec_total_time = time.perf_counter() - exec_start_time - output_ready_wait_time postprocess_start_time = time.perf_counter() scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) try: From f05aa480f02fcd36198b81f5adbf15b51af86a78 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 5 Oct 2026 07:08:37 -0700 Subject: [PATCH 09/14] key unclaimed sets per stream Signed-off-by: Rahul Steiger --- .../test_diffusion_engine_cleanup.py | 59 +++++++++++++++---- vllm_omni/diffusion/diffusion_engine.py | 11 ++-- 2 files changed, 54 insertions(+), 16 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index 46deb9227e0..a414b4c1872 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -107,8 +107,9 @@ async def test_unstarted_stream_does_not_admit_request() -> None: async def test_closing_stream_discards_unclaimed_async_output() -> None: engine = _make_engine() request_id = "abandoned" - engine._out_streams[request_id] = asyncio.Queue() - engine._unclaimed_async_outputs[request_id] = {"aid-queued"} + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_streams[request_id] = output_queue + engine._unclaimed_async_outputs[output_queue] = {"aid-queued"} stream = engine.get_output_stream(request_id) next_output = asyncio.create_task(anext(stream)) await asyncio.sleep(0) @@ -119,7 +120,43 @@ async def test_closing_stream_discards_unclaimed_async_output() -> None: engine.executor.drop_output.assert_called_once_with("aid-queued") assert request_id not in engine._out_streams - assert request_id not in engine._unclaimed_async_outputs + assert not engine._unclaimed_async_outputs + + +@pytest.mark.asyncio +async def test_old_stream_cleanup_preserves_reused_request_outputs() -> None: + engine = _make_engine() + engine.main_loop = None + request_id = "reused" + old_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_streams[request_id] = old_queue + engine._put_output(request_id, DiffusionOutput(finished=True)) + old_stream = engine.get_output_stream(request_id) + await anext(old_stream) + # Leave an unclaimed output owned by the old stream to exercise cleanup. + engine._put_output(request_id, DiffusionOutput(async_output_id="aid-old")) + + new_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_streams[request_id] = new_queue + engine._put_output(request_id, DiffusionOutput(async_output_id="aid-new")) + + await old_stream.aclose() + + engine.executor.drop_output.assert_called_once_with("aid-old") + assert engine._out_streams[request_id] is new_queue + assert engine._unclaimed_async_outputs == {new_queue: {"aid-new"}} + + materialized_output = DiffusionOutput(output="new", finished=True) + ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + ready.set_result(materialized_output) + engine.executor.wait_output_ready.return_value = ready + new_stream = engine.get_output_stream(request_id) + assert await anext(new_stream) is materialized_output + engine.executor.wait_output_ready.assert_called_once_with("aid-new") + await new_stream.aclose() + assert not engine._out_streams + assert not engine._unclaimed_async_outputs + engine.executor.drop_output.assert_called_once_with("aid-old") @pytest.mark.asyncio @@ -134,12 +171,12 @@ async def test_async_output_is_claimed_after_materialization() -> None: output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() output_queue.put_nowait(pending_output) engine._out_streams[request_id] = output_queue - engine._unclaimed_async_outputs[request_id] = {"aid-claimed"} + engine._unclaimed_async_outputs[output_queue] = {"aid-claimed"} stream = engine.get_output_stream(request_id) assert await anext(stream) is materialized_output assert materialized_output.stage_durations["output_ready_wait"] >= 0.0 - assert request_id not in engine._unclaimed_async_outputs + assert not engine._unclaimed_async_outputs engine.executor.wait_output_ready.assert_called_once_with("aid-claimed") await stream.aclose() @@ -156,7 +193,7 @@ async def test_exceptional_materialization_retires_async_output() -> None: output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() output_queue.put_nowait(DiffusionOutput(async_output_id=async_output_id)) engine._out_streams[request_id] = output_queue - engine._unclaimed_async_outputs[request_id] = {async_output_id} + engine._unclaimed_async_outputs[output_queue] = {async_output_id} stream = engine.get_output_stream(request_id) next_output = asyncio.create_task(anext(stream)) await asyncio.sleep(0) @@ -170,7 +207,7 @@ async def test_exceptional_materialization_retires_async_output() -> None: assert executor._output_futures == {} assert executor._completed_outputs == {} - assert request_id not in engine._unclaimed_async_outputs + assert not engine._unclaimed_async_outputs @pytest.mark.asyncio @@ -186,7 +223,7 @@ async def test_cancelling_materialization_discards_async_output(running: bool) - output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() output_queue.put_nowait(pending_output) engine._out_streams[request_id] = output_queue - engine._unclaimed_async_outputs[request_id] = {"aid-cancelled"} + engine._unclaimed_async_outputs[output_queue] = {"aid-cancelled"} stream = engine.get_output_stream(request_id) next_output = asyncio.create_task(anext(stream)) await asyncio.sleep(0) @@ -203,7 +240,7 @@ async def test_cancelling_materialization_discards_async_output(running: bool) - engine.executor.drop_output.assert_not_called() assert ready.cancelled() assert request_id not in engine._out_streams - assert request_id not in engine._unclaimed_async_outputs + assert not engine._unclaimed_async_outputs @pytest.mark.asyncio @@ -217,7 +254,7 @@ async def test_cancellation_after_delivery_does_not_recreate_waiter(failed: bool output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() output_queue.put_nowait(DiffusionOutput(async_output_id=async_output_id)) engine._out_streams[request_id] = output_queue - engine._unclaimed_async_outputs[request_id] = {async_output_id} + engine._unclaimed_async_outputs[output_queue] = {async_output_id} stream = engine.get_output_stream(request_id) next_output = asyncio.create_task(anext(stream)) await asyncio.sleep(0) @@ -237,7 +274,7 @@ async def test_cancellation_after_delivery_does_not_recreate_waiter(failed: bool assert executor._output_futures == {} assert executor._completed_outputs == {} assert request_id not in engine._out_streams - assert request_id not in engine._unclaimed_async_outputs + assert not engine._unclaimed_async_outputs def test_emit_finished_outputs_finalizes_already_drained_waiter() -> None: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index aba6de540df..f0fd094574e 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -409,7 +409,8 @@ def _init_runtime_state(self) -> None: self._rpc_lock = threading.RLock() self._cv = threading.Condition(self._rpc_lock) self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {} - self._unclaimed_async_outputs: dict[str, set[str]] = {} + # Track ownership by stream because request IDs can be reused. + self._unclaimed_async_outputs: dict[asyncio.Queue[DiffusionOutput], set[str]] = {} self._closed = False self._shutting_down = False self._shutdown_complete = False @@ -1177,11 +1178,11 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu raise finally: with self._cv: - pending_ids = self._unclaimed_async_outputs.get(request_id) + pending_ids = self._unclaimed_async_outputs.get(queue) if pending_ids is not None: pending_ids.discard(async_output_id) if not pending_ids: - self._unclaimed_async_outputs.pop(request_id, None) + self._unclaimed_async_outputs.pop(queue, None) yield output if output.finished: break @@ -1192,7 +1193,7 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu with self._cv: if self._out_streams.get(request_id) is queue: self._out_streams.pop(request_id, None) - abandoned_ids = self._unclaimed_async_outputs.pop(request_id, set()) + abandoned_ids = self._unclaimed_async_outputs.pop(queue, set()) for async_output_id in abandoned_ids: self.executor.drop_output(async_output_id) @@ -1576,7 +1577,7 @@ def _put_output(self, request_id: str, output: DiffusionOutput) -> None: with self._cv: queue = self._out_streams.get(request_id) if queue is not None and async_output_id is not None: - self._unclaimed_async_outputs.setdefault(request_id, set()).add(async_output_id) + self._unclaimed_async_outputs.setdefault(queue, set()).add(async_output_id) if queue is None: if async_output_id is not None: self.executor.drop_output(async_output_id) From 7c566fb2163a2f234a32b2a2e28929b519774881 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 5 Oct 2026 07:16:34 -0700 Subject: [PATCH 10/14] added live stream check Signed-off-by: Rahul Steiger --- .../test_diffusion_engine_cleanup.py | 54 +++++++++++++++++++ vllm_omni/diffusion/diffusion_engine.py | 9 ++-- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index a414b4c1872..a0f834cd63c 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -83,6 +83,60 @@ def test_close_completes_pending_output_streams() -> None: event_loop.close() +@pytest.mark.asyncio +@pytest.mark.parametrize("finished", [False, True]) +async def test_close_preserves_cached_output_for_live_consumer(finished: bool) -> None: + engine = _make_engine() + engine.main_loop = None + executor = _make_output_executor() + # Isolate stream ownership from executor teardown, which clears its cache. + executor.shutdown = Mock() + engine.executor = executor + request_id = "live" + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_streams[request_id] = output_queue + stream = engine.get_output_stream(request_id) + next_output = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + + materialized_output = DiffusionOutput(output="cached", finished=finished) + ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + ready.set_result(materialized_output) + executor._completed_outputs["aid-live"] = ready + engine._put_output(request_id, DiffusionOutput(async_output_id="aid-live", finished=finished)) + + engine.close() + + assert engine._unclaimed_async_outputs == {output_queue: {"aid-live"}} + assert await next_output is materialized_output + if not finished: + closed_output = await anext(stream) + assert closed_output.error == "DiffusionEngine is closed." + with pytest.raises(StopAsyncIteration): + await anext(stream) + assert not engine._unclaimed_async_outputs + assert not executor._completed_outputs + assert not executor._dropped_output_ids + + +def test_close_discards_only_outputs_without_registered_streams() -> None: + engine = _make_engine() + engine.main_loop = None + live_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + abandoned_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_streams["reused"] = live_queue + engine._unclaimed_async_outputs = { + live_queue: {"aid-live"}, + abandoned_queue: {"aid-abandoned"}, + } + + engine.close() + engine.close() + + engine.executor.drop_output.assert_called_once_with("aid-abandoned") + assert engine._unclaimed_async_outputs == {live_queue: {"aid-live"}} + + def test_put_output_discards_async_output_when_stream_is_missing() -> None: engine = _make_engine() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index f0fd094574e..8ea6869699c 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1598,9 +1598,12 @@ def close(self) -> None: self.stop_event.set() pending_streams = list(self._out_streams.values()) self._out_streams.clear() - for async_output_ids in self._unclaimed_async_outputs.values(): - abandoned_ids.update(async_output_ids) - self._unclaimed_async_outputs.clear() + # Registered streams retain ownership of queued outputs until + # their consumers finish draining or close their generators. + live_streams = set(pending_streams) + for stream in list(self._unclaimed_async_outputs): + if stream not in live_streams: + abandoned_ids.update(self._unclaimed_async_outputs.pop(stream)) self._cv.notify_all() for async_output_id in abandoned_ids: From f53946f661b8c97f63161b81d20882fa249b251a Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 5 Oct 2026 07:23:10 -0700 Subject: [PATCH 11/14] admit requests eagerly, and return generator with contextlib.asclosing Signed-off-by: Rahul Steiger --- .../test_diffusion_engine_cleanup.py | 42 +++++++++++++++++-- .../test_diffusion_engine_rpc_routing.py | 1 + vllm_omni/diffusion/diffusion_engine.py | 18 ++++---- 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index a0f834cd63c..d1d589c6ac8 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -146,15 +146,51 @@ def test_put_output_discards_async_output_when_stream_is_missing() -> None: @pytest.mark.asyncio -async def test_unstarted_stream_does_not_admit_request() -> None: +async def test_stream_admits_request_before_iteration() -> None: engine = _make_engine() request = _make_request("unstarted") stream = engine.async_add_req_and_stream_response(request) + + assert engine.scheduler.get_request_state(request.request_id) is not None + assert request.request_id in engine._out_streams + engine.abort(request.request_id) + assert engine.abort_queue.get_nowait() == request.request_id await stream.aclose() - assert engine.scheduler.get_request_state(request.request_id) is None - assert request.request_id not in engine._out_streams + +@pytest.mark.parametrize("closed", [False, True]) +def test_stream_admission_errors_are_raised_at_call_site(closed: bool) -> None: + engine = _make_engine() + engine._closed = closed + if closed: + error_type = RuntimeError + message = "DiffusionEngine is closed." + else: + error_type = ValueError + message = "invalid request" + engine.pre_process_func = Mock(side_effect=ValueError(message)) + + with pytest.raises(error_type, match=message): + engine.async_add_req_and_stream_response(_make_request("invalid")) + assert not engine._out_streams + + +@pytest.mark.asyncio +async def test_closing_response_stream_closes_inner_stream() -> None: + engine = _make_engine() + engine.main_loop = None + request = _make_request("close-wrapper") + stream = engine.async_add_req_and_stream_response(request) + engine._put_output(request.request_id, DiffusionOutput(finished=False)) + await anext(stream) + engine._put_output(request.request_id, DiffusionOutput(async_output_id="aid-unclaimed")) + + await stream.aclose() + + engine.executor.drop_output.assert_called_once_with("aid-unclaimed") + assert not engine._out_streams + assert not engine._unclaimed_async_outputs @pytest.mark.asyncio diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 79e2b04055d..e3e9045294f 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -522,6 +522,7 @@ def fail_reservation(*args, **kwargs): request.diffusion_kv_requests = (object(),) try: response_stream = engine.async_add_req_and_stream_response(request) + assert request.request_id in engine._out_streams output = await asyncio.wait_for(anext(response_stream), timeout=3.0) assert output.error == "native allocation bug" assert output.finished diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 8ea6869699c..44504ec7011 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -12,6 +12,7 @@ import threading import time from collections.abc import AsyncGenerator, Iterable +from contextlib import aclosing from dataclasses import dataclass, field from enum import Enum from typing import TYPE_CHECKING, Any @@ -1197,16 +1198,15 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu for async_output_id in abandoned_ids: self.executor.drop_output(async_output_id) - async def async_add_req_and_stream_response( - self, request: OmniDiffusionRequest - ) -> AsyncGenerator[DiffusionOutput, None]: + def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]: request_id = self.add_request(request) - stream = self.get_output_stream(request_id) - try: - async for output in stream: - yield output - finally: - await stream.aclose() + + async def stream_response() -> AsyncGenerator[DiffusionOutput, None]: + async with aclosing(self.get_output_stream(request_id)) as stream: + async for output in stream: + yield output + + return stream_response() async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: """Deprecated compatibility wrapper over ``async_add_req_and_stream_response()``. From df84ac6007c49b32bf9e1fee7c7e7b2ce792ed94 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 5 Oct 2026 07:36:32 -0700 Subject: [PATCH 12/14] separate stage durations and output_ready_wait_time_ms Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine.py | 3 +-- tests/diffusion/test_diffusion_engine_cleanup.py | 5 +++-- vllm_omni/diffusion/data.py | 3 +++ vllm_omni/diffusion/diffusion_engine.py | 4 ++-- 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index df343e00118..a1d970bc12a 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -877,8 +877,7 @@ async def test_step_streaming_excludes_output_wait_from_execution_time( ) -> None: clock = [0.0] mocker.patch.object(diffusion_engine_module.time, "perf_counter", side_effect=lambda: clock[0]) - stage_durations = {"output_ready_wait": output_wait} if output_wait else {} - output = DiffusionOutput(stage_durations=stage_durations) + output = DiffusionOutput(stage_durations={"denoise": 10.0}, output_ready_wait_time=output_wait) formatted_output = SimpleNamespace(metrics={}) request = SimpleNamespace(scheduler_queue_wait_ms=None) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index d1d589c6ac8..a877f648c31 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -254,7 +254,7 @@ async def test_async_output_is_claimed_after_materialization() -> None: engine = _make_engine() request_id = "claimed" pending_output = DiffusionOutput(async_output_id="aid-claimed") - materialized_output = DiffusionOutput(output="materialized") + materialized_output = DiffusionOutput(output="materialized", stage_durations={"denoise": 1.5}) ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() ready.set_result(materialized_output) engine.executor.wait_output_ready.return_value = ready @@ -265,7 +265,8 @@ async def test_async_output_is_claimed_after_materialization() -> None: stream = engine.get_output_stream(request_id) assert await anext(stream) is materialized_output - assert materialized_output.stage_durations["output_ready_wait"] >= 0.0 + assert materialized_output.output_ready_wait_time >= 0.0 + assert materialized_output.stage_durations == {"denoise": 1.5} assert not engine._unclaimed_async_outputs engine.executor.wait_output_ready.assert_called_once_with("aid-claimed") diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9bf32a6bc64..ec8993696a6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1812,6 +1812,9 @@ class DiffusionOutput: # logged duration of stages stage_durations: dict[str, float] = field(default_factory=dict) + # Engine-side async materialization wait (seconds), separate from model stages. + output_ready_wait_time: float = 0.0 + # memory usage info peak_memory_mb: float = 0.0 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 44504ec7011..b009a276f09 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -476,7 +476,7 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ request_id = self._add_prepared_request(request) generator = self.get_output_stream(request_id) async for output in generator: - output_ready_wait_time = getattr(output, "stage_durations", {}).get("output_ready_wait", 0.0) + output_ready_wait_time = output.output_ready_wait_time # The stream now materializes async outputs before yielding. Keep # that wait separate from execution time in the timing breakdown. exec_total_time = time.perf_counter() - exec_start_time - output_ready_wait_time @@ -1158,7 +1158,7 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu output_ready_wait_start_time = time.perf_counter() try: output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) - output.stage_durations["output_ready_wait"] = time.perf_counter() - output_ready_wait_start_time + output.output_ready_wait_time = time.perf_counter() - output_ready_wait_start_time except asyncio.CancelledError: # Delivery already retires completed futures; dropping # them again would create an orphaned executor waiter. From 2d0f44b93f5608e6a0ec8b68581a64a709fca0a4 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 5 Oct 2026 08:04:13 -0700 Subject: [PATCH 13/14] Fix async output delivery during diffusion shutdown Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine.py | 1 + .../test_diffusion_engine_cleanup.py | 68 +++++++++++++++++++ .../test_diffusion_engine_rpc_routing.py | 1 + vllm_omni/diffusion/diffusion_engine.py | 34 +++++++--- .../diffusion/executor/multiproc_executor.py | 4 ++ 5 files changed, 100 insertions(+), 8 deletions(-) diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index a1d970bc12a..99b2326be1e 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -950,6 +950,7 @@ async def test_async_add_req_and_stream_response(): engine.scheduler = MockScheduler() engine._out_streams = {} engine._unclaimed_async_outputs = {} + engine._shutdown_output_futures = {} engine.abort_queue: queue.Queue[str] = queue.Queue() engine._rpc_queue = queue.Queue() engine._rpc_lock = threading.RLock() diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index a877f648c31..66364b354d5 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -46,6 +46,7 @@ def _make_engine() -> DiffusionEngine: engine._cv = threading.Condition(engine._rpc_lock) engine._out_streams = {} engine._unclaimed_async_outputs = {} + engine._shutdown_output_futures = {} engine._closed = False engine._shutting_down = False engine._shutdown_complete = False @@ -66,6 +67,73 @@ def _make_output_executor() -> MultiprocDiffusionExecutor: return executor +def _make_shutdown_output_executor() -> MultiprocDiffusionExecutor: + executor = _make_output_executor() + # Stub process infrastructure, but exercise the real shutdown/cache cleanup. + executor._pump_stop = threading.Event() + executor._shutdown_cleaner = None + executor._finalizer = SimpleNamespace(alive=False) + executor._result_pump_threads = [] + executor._rpc_futures = {} + executor._batch_split_map = {} + return executor + + +def test_wait_output_ready_after_shutdown_fails_immediately() -> None: + executor = _make_shutdown_output_executor() + executor.shutdown() + + ready = executor.wait_output_ready("aid-after-shutdown") + + try: + assert ready.done(), "A stopped executor must not create a waiter that cannot complete" + with pytest.raises(RuntimeError, match="(?i)shut down|shutdown|closed"): + ready.result() + assert not executor._output_futures + finally: + ready.cancel() + executor.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ready_before_close", [False, True]) +async def test_close_with_real_executor_resolves_queued_output(monkeypatch, ready_before_close: bool) -> None: + engine = _make_engine() + engine.main_loop = None + executor = _make_shutdown_output_executor() + engine.executor = executor + monkeypatch.setattr(diffusion_engine_module, "_async_output_timeout", lambda: 0.05) + engine._out_streams["live"] = asyncio.Queue() + stream = engine.get_output_stream("live") + next_output = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + + materialized_output = DiffusionOutput(output="cached", finished=True) + ready: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + ready.set_result(materialized_output) + if ready_before_close: + executor._completed_outputs["aid-live"] = ready + # Wake the consumer, then close before it gets another event-loop turn. + engine._put_output("live", DiffusionOutput(async_output_id="aid-live", finished=True)) + engine.close() + + try: + if ready_before_close: + assert await next_output is materialized_output + with pytest.raises(StopAsyncIteration): + await anext(stream) + else: + with pytest.raises(RuntimeError, match="Executor shut down"): + await next_output + assert not engine._unclaimed_async_outputs + assert not engine._shutdown_output_futures + assert not executor._completed_outputs + assert not executor._output_futures + finally: + await stream.aclose() + executor.shutdown() + + def test_close_completes_pending_output_streams() -> None: engine = _make_engine() event_loop = asyncio.new_event_loop() diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index e3e9045294f..7986328a1fd 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -179,6 +179,7 @@ def _make_engine_with_loop( engine._cv = threading.Condition(engine._rpc_lock) engine._out_streams = {} engine._unclaimed_async_outputs = {} + engine._shutdown_output_futures = {} engine._closed = False engine.abort_queue = queue.Queue() engine._rpc_queue = queue.Queue() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b009a276f09..4adea89cb09 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -412,6 +412,9 @@ def _init_runtime_state(self) -> None: self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {} # Track ownership by stream because request IDs can be reused. self._unclaimed_async_outputs: dict[asyncio.Queue[DiffusionOutput], set[str]] = {} + self._shutdown_output_futures: dict[ + asyncio.Queue[DiffusionOutput], dict[str, concurrent.futures.Future[DiffusionOutput]] + ] = {} self._closed = False self._shutting_down = False self._shutdown_complete = False @@ -1153,7 +1156,20 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu output: DiffusionOutput = await queue.get() async_output_id = output.async_output_id if async_output_id is not None: - fut = self.executor.wait_output_ready(async_output_id) + with self._cv: + saved_futures = self._shutdown_output_futures.get(queue, {}) + fut = saved_futures.pop(async_output_id, None) + if not saved_futures: + self._shutdown_output_futures.pop(queue, None) + if fut is None: + fut = self.executor.wait_output_ready(async_output_id) + # The consumer now owns this future. Shutdown must not + # claim it again after the executor retires delivery. + pending_ids = self._unclaimed_async_outputs.get(queue) + if pending_ids is not None: + pending_ids.discard(async_output_id) + if not pending_ids: + self._unclaimed_async_outputs.pop(queue, None) timeout = _async_output_timeout() output_ready_wait_start_time = time.perf_counter() try: @@ -1177,13 +1193,6 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu describe(async_output_id) if describe else "unavailable", ) raise - finally: - with self._cv: - pending_ids = self._unclaimed_async_outputs.get(queue) - if pending_ids is not None: - pending_ids.discard(async_output_id) - if not pending_ids: - self._unclaimed_async_outputs.pop(queue, None) yield output if output.finished: break @@ -1195,6 +1204,7 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu if self._out_streams.get(request_id) is queue: self._out_streams.pop(request_id, None) abandoned_ids = self._unclaimed_async_outputs.pop(queue, set()) + self._shutdown_output_futures.pop(queue, None) for async_output_id in abandoned_ids: self.executor.drop_output(async_output_id) @@ -1604,6 +1614,14 @@ def close(self) -> None: for stream in list(self._unclaimed_async_outputs): if stream not in live_streams: abandoned_ids.update(self._unclaimed_async_outputs.pop(stream)) + else: + # Transfer cached results out of the executor before + # shutdown clears its cache. Pending futures receive + # the executor's shutdown error instead. + self._shutdown_output_futures[stream] = { + output_id: self.executor.wait_output_ready(output_id) + for output_id in self._unclaimed_async_outputs[stream] + } self._cv.notify_all() for async_output_id in abandoned_ids: diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 40ab555525f..1d7f6f0453a 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -1130,6 +1130,10 @@ def wait_output_ready(self, async_output_id: str) -> concurrent.futures.Future[D fresh one that would never complete. """ with self._futures_lock: + if self._closed: + closed: concurrent.futures.Future[DiffusionOutput] = concurrent.futures.Future() + closed.set_exception(RuntimeError("Executor shut down")) + return closed cached = self._completed_outputs.pop(async_output_id, None) if cached is not None: return cached From 496d546c3b983465c6090f20268aa1cc3d9505b0 Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Mon, 5 Oct 2026 08:08:08 -0700 Subject: [PATCH 14/14] test cleanup Signed-off-by: Rahul Steiger --- tests/diffusion/test_diffusion_engine_cleanup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index 66364b354d5..bc5e6d7ab0b 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -87,7 +87,7 @@ def test_wait_output_ready_after_shutdown_fails_immediately() -> None: try: assert ready.done(), "A stopped executor must not create a waiter that cannot complete" - with pytest.raises(RuntimeError, match="(?i)shut down|shutdown|closed"): + with pytest.raises(RuntimeError, match="^Executor shut down$"): ready.result() assert not executor._output_futures finally: