From 4ba8e14981bb80a1835a7956357ebf32011b0c27 Mon Sep 17 00:00:00 2001 From: wuhang Date: Mon, 18 May 2026 08:54:57 +0800 Subject: [PATCH] Fix diffusion engine cleanup lifecycle (#3494) Signed-off-by: wuhang Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- .../test_diffusion_engine_cleanup.py | 141 ++++++++++++++++++ tests/diffusion/test_diffusion_scheduler.py | 46 +++++- .../diffusion/test_diffusion_step_pipeline.py | 1 + .../test_multiproc_engine_concurrency.py | 1 + vllm_omni/diffusion/diffusion_engine.py | 102 ++++++++++--- vllm_omni/diffusion/sched/base_scheduler.py | 16 +- 6 files changed, 282 insertions(+), 25 deletions(-) create mode 100644 tests/diffusion/test_diffusion_engine_cleanup.py diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py new file mode 100644 index 000000000..a7b945eea --- /dev/null +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import queue +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_request( + request_id: str, + *, + request_ids: list[str] | None = None, +) -> OmniDiffusionRequest: + return OmniDiffusionRequest( + prompts=[f"prompt_{request_id}"], + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_ids=request_ids or [request_id], + request_id=request_id, + ) + + +def _make_engine() -> DiffusionEngine: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = RequestScheduler() + engine.scheduler.initialize(SimpleNamespace()) + engine.executor = SimpleNamespace(shutdown=Mock()) + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._out_queue = {} + engine._closed = False + engine._shutdown_complete = False + engine.abort_queue = queue.Queue() + engine._loop_started = False + engine.stop_event = None + engine.worker_thread = None + return engine + + +def test_close_completes_pending_async_waiters() -> None: + engine = _make_engine() + event_loop = asyncio.new_event_loop() + try: + future = event_loop.create_future() + engine._out_queue["pending-req"] = future + + engine.close() + + assert future.done() + assert future.result().error == "DiffusionEngine is closed." + finally: + event_loop.close() + + +def test_handle_finished_requests_ignores_already_drained_waiter() -> None: + class RacingOutQueue(dict): + def __contains__(self, key): + return True + + def pop(self, key, default=None): + return default + + engine = _make_engine() + engine._out_queue = RacingOutQueue() + engine._finalize_finished_request = Mock(side_effect=AssertionError("should not finalize drained waiters")) + + engine._handle_finished_requests({"pending-req"}) + + engine._finalize_finished_request.assert_not_called() + + +def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None: + engine = _make_engine() + request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"]) + sched_req_id = engine.scheduler.add_request(request) + + engine.abort("batch-parent") + engine._process_aborts_queue() + + state = engine.scheduler.get_request_state(sched_req_id) + assert state is not None + assert state.status == DiffusionRequestStatus.FINISHED_ABORTED + + +def test_close_rejects_late_async_requests() -> None: + engine = _make_engine() + event_loop = asyncio.new_event_loop() + try: + engine.main_loop = event_loop + engine.close() + + with pytest.raises(RuntimeError, match="closed"): + engine.add_request(_make_request("late-req")) + finally: + event_loop.close() + + +def test_close_resets_loop_started_for_dead_worker_thread() -> None: + engine = _make_engine() + engine._loop_started = True + engine.worker_thread = SimpleNamespace(is_alive=Mock(return_value=False)) + + engine.close() + + assert engine._loop_started is False + + +def test_close_defers_resource_shutdown_until_worker_thread_stops() -> None: + engine = _make_engine() + engine.scheduler.close = Mock() + engine._loop_started = True + worker_thread = SimpleNamespace( + is_alive=Mock(side_effect=[True, True, False, False]), + join=Mock(), + ) + engine.worker_thread = worker_thread + + engine.close() + + worker_thread.join.assert_called_once_with(timeout=10) + engine.scheduler.close.assert_not_called() + engine.executor.shutdown.assert_not_called() + assert engine._shutdown_complete is False + assert engine._loop_started is True + + engine.close() + + engine.scheduler.close.assert_called_once() + engine.executor.shutdown.assert_called_once() + assert engine._shutdown_complete is True + assert engine._loop_started is False diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index c242a0572..f0bc1da35 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -305,17 +305,55 @@ class TestRequestScheduler: prompts=["prompt_map_a", "prompt_map_b"], sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_ids=["map-a", "map-b"], + request_id="map-parent", ) sched_req_id = self.scheduler.add_request(request) assert self.scheduler.get_sched_req_id("map-a") == sched_req_id assert self.scheduler.get_sched_req_id("map-b") == sched_req_id + assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id self.scheduler.pop_request_state(sched_req_id) assert self.scheduler.get_sched_req_id("map-a") is None assert self.scheduler.get_sched_req_id("map-b") is None + assert self.scheduler.get_sched_req_id("map-parent") is None + + def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None: + self.scheduler.add_request( + OmniDiffusionRequest( + prompts=["prompt_existing"], + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_ids=["existing-child"], + request_id="duplicate-parent", + ) + ) + + colliding_request = OmniDiffusionRequest( + prompts=["prompt_unique"], + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_ids=["unique-child"], + request_id="duplicate-parent", + ) + + with pytest.raises(ValueError, match="duplicate-parent"): + self.scheduler.add_request(colliding_request) + + assert self.scheduler.get_request_state("unique-child") is None + assert self.scheduler.get_sched_req_id("unique-child") is None + + valid_request = OmniDiffusionRequest( + prompts=["prompt_unique"], + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_ids=["unique-child"], + request_id="unique-parent", + ) + sched_req_id = self.scheduler.add_request(valid_request) + + assert sched_req_id == "unique-child" + assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id + assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id class TestDiffusionEngine: @@ -324,6 +362,8 @@ class TestDiffusionEngine: engine.scheduler = RequestScheduler() engine.scheduler.initialize(SimpleNamespace()) engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._closed = False engine.abort_queue = queue.Queue() request = _make_request("engine") @@ -343,6 +383,8 @@ class TestDiffusionEngine: engine = DiffusionEngine.__new__(DiffusionEngine) engine.scheduler = scheduler engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._closed = False engine.abort_queue = queue.Queue() engine.execute_fn = mocker.Mock(return_value=runner_output) @@ -394,7 +436,8 @@ class TestDiffusionEngine: @pytest.mark.asyncio async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) - engine._loop_started = False + engine._closed = False + engine._loop_started = True engine._init_lock = asyncio.Lock() engine.main_loop = asyncio.get_running_loop() engine.stop_event = threading.Event() @@ -410,6 +453,7 @@ class TestDiffusionEngine: engine = DiffusionEngine.__new__(DiffusionEngine) engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) + engine._closed = False engine.scheduler = RequestScheduler() engine.scheduler.initialize(SimpleNamespace()) engine.abort_queue = queue.Queue() diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index a443363a9..1954f4a05 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -305,6 +305,7 @@ def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine: engine.execute_fn = execute_fn engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) + engine._closed = False engine.abort_queue = queue.Queue() return engine diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index a1e5c3456..aa2a74ac1 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -79,6 +79,7 @@ def _make_engine(num_gpus: int = 1): engine.executor = executor engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) + engine._closed = False engine._loop_started = False engine._rpc_queue = queue.Queue() engine.abort_queue = queue.Queue() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index fe4a4c77e..c13bd3c0c 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -148,6 +148,7 @@ class DiffusionEngine: logger.warning(f"Non-stepwise-execution does not support max-num-seqs={max_num_seqs}, set it to 1.") self.main_loop: asyncio.AbstractEventLoop | None = None self.stop_event: threading.Event | None = None + self.worker_thread: threading.Thread | None = None self._loop_started = False self._init_lock = asyncio.Lock() # _rpc_lock is retained solely as the underlying lock for self._cv, @@ -157,6 +158,8 @@ class DiffusionEngine: self._rpc_lock = threading.RLock() self._cv = threading.Condition(self._rpc_lock) self._out_queue: dict[str, asyncio.Future] = {} + self._closed = False + self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request @@ -169,11 +172,15 @@ class DiffusionEngine: raise e async def _check_and_start_background_loop(self): + if self._closed: + raise RuntimeError("DiffusionEngine is closed.") if self._loop_started: return async with self._init_lock: # double check, in case of lock queue issue + if self._closed: + raise RuntimeError("DiffusionEngine is closed.") if self._loop_started: return @@ -541,15 +548,16 @@ class DiffusionEngine: missing_result_error: str = "Diffusion execution finished without a final output", ): for rid in finished_ids: - if rid in self._out_queue: - fut = self._out_queue.pop(rid) - if runner_output is not None: - _output = runner_output.get_req_output(rid) - else: - _output = None - out = self._finalize_finished_request(rid, _output, missing_result_error) - if not fut.done(): - self.main_loop.call_soon_threadsafe(fut.set_result, out) + with self._cv: + fut = self._out_queue.pop(rid, None) + if fut is None: + continue + if runner_output is not None: + _output = runner_output.get_req_output(rid) + else: + _output = None + out = self._finalize_finished_request(rid, _output, missing_result_error) + self._complete_future(fut, out) @staticmethod def make_engine( @@ -567,8 +575,10 @@ class DiffusionEngine: return DiffusionEngine(config, scheduler=scheduler) def add_request(self, request: OmniDiffusionRequest) -> str: - fut = self.main_loop.create_future() with self._cv: + if self._closed: + raise RuntimeError("DiffusionEngine is closed.") + fut = self.main_loop.create_future() sched_req_id = self.scheduler.add_request(request) self._out_queue[sched_req_id] = fut self._cv.notify_all() @@ -594,6 +604,8 @@ class DiffusionEngine: def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: with self._rpc_lock: + if self._closed: + raise RuntimeError("DiffusionEngine is closed.") target_sched_req_id = self.scheduler.add_request(request) # keep scheduling and executing until the target request is finished @@ -801,31 +813,79 @@ class DiffusionEngine: task.future.cancel() raise TimeoutError(f"RPC call to {method} timed out.") from exc + def _complete_future(self, fut: asyncio.Future, output: DiffusionOutput) -> None: + if fut.done(): + return + + def _set_result() -> None: + if not fut.done(): + fut.set_result(output) + + try: + loop = fut.get_loop() + except AttributeError: + loop = self.main_loop + + try: + running_loop = asyncio.get_running_loop() + except RuntimeError: + running_loop = None + + if loop is not None and loop.is_running() and loop is not running_loop: + loop.call_soon_threadsafe(_set_result) + else: + _set_result() + def close(self) -> None: - if hasattr(self, "scheduler"): - self.scheduler.close() - if hasattr(self, "stop_event") and self.stop_event is not None: - with self._cv: - self.stop_event.set() + pending_futures: list[asyncio.Future] = [] + with self._cv: + if self._closed and self._shutdown_complete: + return + if not self._closed: + self._closed = True + if self.stop_event is not None: + self.stop_event.set() + pending_futures = list(self._out_queue.values()) + self._out_queue.clear() self._cv.notify_all() - if hasattr(self, "worker_thread") and self.worker_thread.is_alive(): - self.worker_thread.join(timeout=10) - if self.worker_thread.is_alive(): - logger.warning("Worker thread did not terminate within 10s, resources may not be released.") + + closed_output = DiffusionOutput(error="DiffusionEngine is closed.") + for fut in pending_futures: + self._complete_future(fut, closed_output) + + worker_thread = self.worker_thread + if worker_thread is not None: + if worker_thread.is_alive(): + worker_thread.join(timeout=10) + if worker_thread.is_alive(): + logger.warning( + "Worker thread did not terminate within 10s; scheduler and executor shutdown will be skipped." + ) + return else: self._loop_started = False - if hasattr(self, "executor"): - self.executor.shutdown() + else: + self._loop_started = False + + self.scheduler.close() + self.executor.shutdown() + self._shutdown_complete = True def abort(self, request_id: str | Iterable[str]) -> None: request_ids = [request_id] if isinstance(request_id, str) else list(request_id) with self._cv: + if self._closed: + return for req_id in request_ids: self.abort_queue.put(req_id) self._cv.notify_all() def _process_aborts_queue(self) -> None: + with self._cv: + self._drain_abort_queue() + + def _drain_abort_queue(self) -> None: if self.abort_queue.empty(): return diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index d2c0f1e75..f21c03e98 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -70,8 +70,9 @@ class _BaseScheduler(SchedulerInterface): def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str: state = self._make_request_state(sched_req_id, request) + request_ids = self._request_ids_for_mapping(request) + self._register_request_ids(request_ids, sched_req_id) self._request_states[sched_req_id] = state - self._register_request_ids(request.request_ids, sched_req_id) self._waiting.append(sched_req_id) logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, sched_req_id, len(self._waiting)) return sched_req_id @@ -134,7 +135,7 @@ class _BaseScheduler(SchedulerInterface): self._pop_extra_request_state(sched_req_id) state = self._request_states.pop(sched_req_id, None) if state is not None: - self._unregister_request_ids(state.req.request_ids, sched_req_id) + self._unregister_request_ids(self._request_ids_for_mapping(state.req), sched_req_id) return state def preempt_request(self, sched_req_id: str) -> bool: @@ -253,13 +254,22 @@ class _BaseScheduler(SchedulerInterface): return self._running_sampling_params_key def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None: - for request_id in request_ids: + unique_request_ids = list(dict.fromkeys(request_ids)) + for request_id in unique_request_ids: existing = self._request_id_to_sched_req_id.get(request_id) if existing is not None and existing != sched_req_id: raise ValueError(f"request_id {request_id!r} is already mapped to active sched_req_id {existing!r}.") + for request_id in unique_request_ids: self._request_id_to_sched_req_id[request_id] = sched_req_id def _unregister_request_ids(self, request_ids: list[str], sched_req_id: str) -> None: for request_id in request_ids: if self._request_id_to_sched_req_id.get(request_id) == sched_req_id: self._request_id_to_sched_req_id.pop(request_id, None) + + @staticmethod + def _request_ids_for_mapping(request: OmniDiffusionRequest) -> list[str]: + request_ids = list(request.request_ids) + if request_id := getattr(request, "request_id", None): + request_ids.append(request_id) + return list(dict.fromkeys(request_ids))