Fix diffusion engine cleanup lifecycle (#3494)

Signed-off-by: wuhang <wuhang6@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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()

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
 

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()

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
 

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))