diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index b7daa2e95f3..99b2326be1e 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -869,9 +869,45 @@ 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]) + output = DiffusionOutput(stage_durations={"denoise": 10.0}, output_ready_wait_time=output_wait) + 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"]) -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 +927,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] @@ -903,6 +949,8 @@ async def test_async_add_req_and_stream_response(): engine = object.__new__(DiffusionEngine) 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() @@ -924,7 +972,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 3f138e72107..bc5e6d7ab0b 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -2,8 +2,10 @@ # 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 @@ -13,8 +15,10 @@ 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 from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -33,10 +37,16 @@ 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._shutdown_output_futures = {} engine._closed = False engine._shutting_down = False engine._shutdown_complete = False @@ -47,6 +57,83 @@ 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._dropped_output_ids = OrderedDict() + executor._closed = False + 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="^Executor shut down$"): + 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() @@ -64,6 +151,291 @@ 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() + + 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_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() + + +@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 +async def test_closing_stream_discards_unclaimed_async_output() -> None: + engine = _make_engine() + request_id = "abandoned" + 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) + + 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 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 +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", 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 + output_queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + output_queue.put_nowait(pending_output) + engine._out_streams[request_id] = output_queue + 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.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") + + await stream.aclose() + 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[output_queue] = {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 not engine._unclaimed_async_outputs + + +@pytest.mark.asyncio +@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) + engine._out_streams[request_id] = output_queue + 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) + + next_output.cancel() + with pytest.raises(asyncio.CancelledError): + await next_output + + 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 not 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[output_queue] = {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 not 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 +451,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/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index d26f20c1d18..7986328a1fd 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -178,6 +178,8 @@ 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._shutdown_output_futures = {} engine._closed = False engine.abort_queue = queue.Queue() engine._rpc_queue = queue.Queue() @@ -521,14 +523,8 @@ 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) + 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 assert engine.worker_thread.is_alive() 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 57dddc88904..4adea89cb09 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 @@ -409,6 +410,11 @@ 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]] = {} + # 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 @@ -473,26 +479,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 = 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 + 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 postprocess_start_time = time.perf_counter() scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) try: @@ -1164,6 +1154,45 @@ async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOu try: while True: output: DiffusionOutput = await queue.get() + async_output_id = output.async_output_id + if async_output_id is not None: + 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: + output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) + 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. + if not fut.done(): + self.executor.drop_output(async_output_id) + raise + except (TimeoutError, asyncio.TimeoutError): + 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 " + "to allow slower steps. Executor state: %s", + timeout, + _ASYNC_OUTPUT_TIMEOUT_ENV, + describe(async_output_id) if describe else "unavailable", + ) + raise yield output if output.finished: break @@ -1174,10 +1203,20 @@ 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(queue, set()) + self._shutdown_output_futures.pop(queue, None) + 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) - return self.get_output_stream(request_id) + + 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()``. @@ -1544,14 +1583,20 @@ def _put_queue_output( queue.put_nowait(output) def _put_output(self, request_id: str, output: DiffusionOutput) -> 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: + 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) 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 +1608,25 @@ def close(self) -> None: self.stop_event.set() pending_streams = list(self._out_streams.values()) self._out_streams.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)) + 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: + 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 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