add check_health in inline client

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

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