diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index dcb35cfde1f..72436c79b40 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -385,9 +385,9 @@ def collective_rpc( raise def check_health(self) -> None: - self._ensure_open() if self.is_failed: raise EngineDeadError() + self._ensure_open() for p in self._processes: if not p.is_alive(): self.is_failed = True diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index a33a3e95619..1c251867ae9 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -15,6 +15,7 @@ import torch from PIL import Image from vllm.logger import init_logger +from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine @@ -329,6 +330,12 @@ async def collective_rpc_async( None, ) + def check_health(self) -> None: + """Check if the inline diffusion engine and its workers are healthy.""" + if self._shutting_down: + raise EngineDeadError("InlineStageDiffusionClient is shutting down") + self._engine.executor.check_health() + def shutdown(self) -> None: self._shutting_down = True