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

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

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 = []