[Refactor]: clean up diffusion executor lifecycle (#5215)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -57,10 +57,8 @@ def _make_executor(num_gpus: int = 1):
     Returns ``(executor, request_queue, result_queue)``.
     """
     od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False)
-    monkeypatch = pytest.MonkeyPatch()
-    monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None)
-    executor = MultiprocDiffusionExecutor(od_cfg)
-    monkeypatch.undo()
+    executor = object.__new__(MultiprocDiffusionExecutor)
+    executor.od_config = od_cfg
 
     req_q: queue.Queue = queue.Queue()
     res_q: queue.Queue = queue.Queue()
@@ -73,7 +71,7 @@ def _make_executor(num_gpus: int = 1):
     executor._result_mq = mock_rmq
     executor._closed = False
     executor._processes = []
-    executor.is_failed = False
+    executor._is_failed = False
     executor._failure_callbacks = []
     return executor, req_q, res_q
 
@@ -622,7 +620,7 @@ class TestMultiprocExecutorRaisesEngineDeadError:
         executor._broadcast_mq = MagicMock()
         executor._result_mq = MagicMock()
         executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError)
-        executor.is_failed = True
+        executor._is_failed = True
 
         with pytest.raises(EngineDeadError):
             executor.collective_rpc(
@@ -643,7 +641,7 @@ class TestMultiprocExecutorRaisesEngineDeadError:
             nonlocal call_count
             call_count += 1
             if call_count == 1:
-                executor.is_failed = True
+                executor._is_failed = True
                 raise TimeoutError
             return orig_dequeue(timeout=timeout)
 
@@ -920,7 +918,7 @@ def _make_short_lived_process() -> mp.Process:
 
 
 class TestMultiprocExecutorWorkerMonitor:
-    """Integration tests for ``start_worker_monitor``.
+    """Integration tests for ``_start_worker_monitor``.
 
     Uses real short-lived subprocesses so that OS-level sentinel fd
     readiness is exercised end-to-end.
@@ -928,17 +926,17 @@ class TestMultiprocExecutorWorkerMonitor:
 
     def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
         """When a worker process dies, the monitor thread must:
-        1. Set ``is_failed = True``
+        1. Set ``_is_failed = True``
         2. Call ``shutdown()`` (which sets ``_closed = True``)
         3. Invoke all registered failure callbacks
         """
         executor = object.__new__(MultiprocDiffusionExecutor)
         executor._closed = False
-        executor.is_failed = False
+        executor._is_failed = False
         executor._failure_callbacks = []
         executor._broadcast_mq = None
         executor._result_mq = None
-        executor.resources = None
+        executor._shutdown_cleaner = None
         # Use a no-op so shutdown() doesn't crash on None resources.
         executor._finalizer = lambda: None
 
@@ -948,35 +946,38 @@ class TestMultiprocExecutorWorkerMonitor:
         callback_called = threading.Event()
         executor.register_failure_callback(callback_called.set)
 
-        executor.start_worker_monitor()
+        executor._start_worker_monitor()
 
         # Wait for the process to exit and the monitor to react.
         proc.join(5)
-        assert _poll_flag(lambda: executor.is_failed), "is_failed was not set"
+        assert _poll_flag(lambda: executor._is_failed), "_is_failed was not set"
         assert executor._closed, "shutdown() was not called"
         assert callback_called.wait(timeout=2), "failure callback was not invoked"
+        assert executor.is_dead
 
     def test_worker_monitor_noop_when_already_closed(self):
         """If ``_closed`` is already True when the process dies (orderly
-        shutdown), the monitor must *not* set ``is_failed``."""
+        shutdown), the monitor must *not* set ``_is_failed``."""
         executor = object.__new__(MultiprocDiffusionExecutor)
         executor._closed = True  # already shut down
-        executor.is_failed = False
+        executor._is_failed = False
         executor._failure_callbacks = []
         executor._broadcast_mq = None
         executor._result_mq = None
-        executor.resources = None
+        executor._shutdown_cleaner = None
         executor._finalizer = lambda: None
 
         proc = _make_short_lived_process()
         executor._processes = [proc]
 
-        executor.start_worker_monitor()
+        executor._start_worker_monitor()
         proc.join(5)
 
         # Give the monitor thread a chance to run (it should early-return).
         time.sleep(0.3)
-        assert not executor.is_failed, "is_failed should remain False on orderly shutdown"
+        assert not executor._is_failed, "_is_failed should remain False on orderly shutdown"
+        # Orderly close still reports dead via the public accessor.
+        assert executor.is_dead
 
 
 class TestStageDiffusionClientProcMonitor:

vllm_omni/diffusion/executor/abstract.py

@@ -59,7 +59,7 @@ class DiffusionExecutor(ABC):
             raise ValueError(f"Unknown distributed executor backend: {distributed_executor_backend}")
         return executor_class
 
-    def __init__(self, od_config: OmniDiffusionConfig):
+    def __init__(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
         self._init_executor()
 
@@ -68,6 +68,12 @@ class DiffusionExecutor(ABC):
         """Initialize the executor (e.g., launch workers, setup IPC)."""
         pass
 
+    @property
+    @abstractmethod
+    def is_dead(self) -> bool:
+        """Whether the executor is shut down or has failed fatally."""
+        pass
+
     @abstractmethod
     def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute request-mode work from a scheduler output."""

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -7,10 +7,11 @@ import time
 import weakref
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from multiprocessing.synchronize import Event
+from typing import TYPE_CHECKING, Any, cast
 
 import zmq
-from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
+from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
 
@@ -29,29 +30,21 @@ _DEQUEUE_TIMEOUT_S = 5.0
 
 
 @dataclass
-class BackgroundResources:
-    """
-    Used as a finalizer for clean shutdown.
-    """
+class _ExecutorShutdownCleaner:
+    """Finalizer that shuts down executor worker processes."""
 
     broadcast_mq: MessageQueue | None = None
-    result_mq: MessageQueue | None = None
     num_workers: int = 0
     processes: list[mp.Process] | None = None
 
-    def __call__(self):
+    def __call__(self) -> None:
         """Clean up background resources."""
-        if hasattr(self, "wake_events") and self.wake_events:
-            for ev in self.wake_events:
-                ev.set()
-
         if self.broadcast_mq is not None:
             try:
                 for _ in range(self.num_workers):
                     self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0)
 
                 self.broadcast_mq = None
-                self.result_mq = None
             except Exception as exc:
                 logger.warning("Failed to send shutdown signal: %s", exc)
 
@@ -72,10 +65,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     def _init_executor(self) -> None:
         self._processes: list[mp.Process] = []
         self._closed = False
-        self.is_failed = False
+        self._is_failed = False
         self._failure_callbacks: list[Callable[[], None]] = []
+        self._result_mq: MessageQueue | None = None
 
-        num_workers = self.od_config.num_gpus
+        num_workers = cast(int, self.od_config.num_gpus)
         self.wake_events = [mp.Event() for _ in range(num_workers)]
 
         self._broadcast_mq = self._init_broadcast_queue(num_workers)
@@ -83,18 +77,23 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
         # Launch workers
         processes, result_handle = self._launch_workers(broadcast_handle, self.wake_events)
-        self._result_mq = self._init_result_queue(result_handle)
         self._processes = processes
 
-        self.resources = BackgroundResources(
+        shutdown_cleaner = _ExecutorShutdownCleaner(
             broadcast_mq=self._broadcast_mq,
-            result_mq=self._result_mq,
             num_workers=num_workers,
             processes=self._processes,
         )
-        self._finalizer = weakref.finalize(self, self.resources)
+        self._shutdown_cleaner: _ExecutorShutdownCleaner | None = shutdown_cleaner
+        self._finalizer = weakref.finalize(self, shutdown_cleaner)
 
-        self.start_worker_monitor()
+        try:
+            self._result_mq = self._init_result_queue(result_handle)
+        except Exception:
+            self.shutdown()
+            raise
+
+        self._start_worker_monitor()
 
     def _init_broadcast_queue(self, num_workers: int) -> MessageQueue:
         return MessageQueue(
@@ -103,20 +102,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             local_reader_ranks=list(range(num_workers)),
         )
 
-    def _init_result_queue(self, result_handle) -> MessageQueue | None:
+    def _init_result_queue(self, result_handle: Handle | None) -> MessageQueue:
         if result_handle is None:
-            logger.error("Failed to get result queue handle from workers")
-            return None
+            raise RuntimeError("Failed to get result queue handle from workers")
         return MessageQueue.create_from_handle(result_handle, 0)
 
     def _ensure_open(self) -> None:
         if self._closed:
             raise RuntimeError("DiffusionExecutor is closed.")
         if self._result_mq is None:
-            raise RuntimeError("Result queue not initialized")
+            raise RuntimeError("Result queue is closed")
+        if self._broadcast_mq is None:
+            raise RuntimeError("Broadcast queue is closed")
 
     def _dequeue_one_with_failure_polling(self, deadline: float | None, method: str) -> Any:
-        """Block until one result message, polling ``is_failed`` between chunk timeouts."""
+        """Block until one result message, polling ``_is_failed`` between chunk timeouts."""
         while True:
             if deadline is None:
                 chunk_timeout = _DEQUEUE_TIMEOUT_S
@@ -126,9 +126,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     raise TimeoutError(f"RPC call to {method} timed out.")
                 chunk_timeout = min(_DEQUEUE_TIMEOUT_S, remaining)
             try:
-                return self._result_mq.dequeue(timeout=chunk_timeout)
+                return self._result_mq.dequeue(timeout=chunk_timeout)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
             except (TimeoutError, zmq.error.Again):
-                if self.is_failed:
+                if self._is_failed:
                     raise EngineDeadError()
                 continue
 
@@ -183,11 +183,15 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
         return response
 
-    def _launch_workers(self, broadcast_handle, wake_events):
+    def _launch_workers(
+        self,
+        broadcast_handle: Handle,
+        wake_events: list[Event],
+    ) -> tuple[list[mp.Process], Handle | None]:
         od_config = self.od_config
         logger.info("Starting server...")
 
-        num_gpus = od_config.num_gpus
+        num_gpus = cast(int, od_config.num_gpus)
         mp.set_start_method("spawn", force=True)
         processes = []
 
@@ -221,7 +225,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             processes.append(process)
 
         # Wait for all workers to be ready
-        scheduler_infos = []
         result_handle = None
         for writer in scheduler_pipe_writers:
             writer.close()
@@ -241,14 +244,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             if i == 0:
                 result_handle = data.get("result_handle")
 
-            scheduler_infos.append(data)
             reader.close()
 
         logger.debug("All workers are ready")
 
         return processes, result_handle
 
-    def start_worker_monitor(self) -> None:
+    @property
+    def is_dead(self) -> bool:
+        """Whether the executor is shut down or a worker has failed fatally."""
+        return self._closed or self._is_failed
+
+    def _start_worker_monitor(self) -> None:
         # Monitors worker process liveness. If any die unexpectedly,
         # logs an error, shuts down the executor and invokes the failure
         # callback to inform the engine.
@@ -287,7 +294,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     "Diffusion worker(s) died unexpectedly: %s",
                     details,
                 )
-                self.is_failed = True
+                self._is_failed = True
 
             self.shutdown()
 
@@ -416,7 +423,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
         try:
             # Broadcast RPC request to all workers via unified message queue
-            self._broadcast_mq.enqueue(rpc_request)
+            self._broadcast_mq.enqueue(rpc_request)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown
 
             # Only rank 0 has a result_mq, so we always expect exactly 1 response
             num_responses = 1
@@ -440,12 +447,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             raise
 
     def check_health(self) -> None:
-        if self.is_failed:
+        if self._is_failed:
             raise EngineDeadError()
         self._ensure_open()
         for p in self._processes:
             if not p.is_alive():
-                self.is_failed = True
+                self._is_failed = True
                 raise EngineDeadError(f"Worker process {p.name} is dead")
 
     def shutdown(self) -> None:
@@ -455,5 +462,5 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         finally:
             self._broadcast_mq = None
             self._result_mq = None
-            self.resources = None
+            self._shutdown_cleaner = None
             self._processes = []

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -40,6 +40,7 @@ from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 
 logger = init_logger(__name__)
 
@@ -92,20 +93,19 @@ class StageDiffusionProc:
     def _is_executor_dead(self) -> bool:
         """True iff the multiproc executor has been closed or marked failed.
 
-        Detects the "workers died but the diffusion proc is still pulling
-        requests" case: ``MultiprocDiffusionExecutor`` sets ``_closed = True``
-        and ``is_failed = True`` from its worker-monitor thread the moment any
-        worker process exits; every subsequent ``execute_request`` /
-        ``collective_rpc`` then raises ``RuntimeError("DiffusionExecutor is
-        closed.")`` inside the engine. Callers in ``run_loop`` use this to
-        decide whether a per-request failure is recoverable or fatal.
+        requests" case: ``MultiprocDiffusionExecutor`` sets ``is_dead`` from its
+        worker-monitor thread the moment any worker process exits; every
+        subsequent ``execute_request`` / ``collective_rpc`` then raises
+        ``RuntimeError("DiffusionExecutor is closed.")`` inside the engine.
+        Callers in ``run_loop`` use this to decide whether a per-request
+        failure is recoverable or fatal.
         """
         if self._engine is None:
             return False
-        executor = getattr(self._engine, "executor", None)
+        executor: DiffusionExecutor | None = getattr(self._engine, "executor", None)
         if executor is None:
             return False
-        return bool(getattr(executor, "_closed", False) or getattr(executor, "is_failed", False))
+        return executor.is_dead
 
     def _signal_fatal_engine_failure(self, reason: str) -> None:
         """Idempotently signal ``run_loop`` to tear down on a fatal engine error."""