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