From 2cf5f57b3e0a5bc820f9f382ea88db68006915e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:30:59 +0800 Subject: [PATCH] [Refactor]: clean up diffusion executor lifecycle (#5215) Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> --- .../test_multiproc_engine_concurrency.py | 37 ++++----- vllm_omni/diffusion/executor/abstract.py | 8 +- .../diffusion/executor/multiproc_executor.py | 81 ++++++++++--------- vllm_omni/diffusion/stage_diffusion_proc.py | 18 ++--- 4 files changed, 79 insertions(+), 65 deletions(-) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 070d899e6..9b8e08a74 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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: diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index 8bfb5923a..17e5c0504 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/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.""" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 2e491ff69..d0ed76bb9 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/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 = [] diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index e74a34b59..0e00f8eea 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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."""