[Enhancement] Patch OmniStage.try_collect() with ray alive checks (#1561)

Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
Signed-off-by: Daniel Huang <pilotflyer824@gmail.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

vllm_omni/distributed/ray_utils/utils.py

@@ -176,6 +176,26 @@ def start_ray_actor(
         runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}, "CUDA_LAUNCH_BLOCKING": "1"},
     ).remote()
 
-    worker_actor.run.remote(worker_entry_fn, *args, **kwargs)
+    task_ref = worker_actor.run.remote(worker_entry_fn, *args, **kwargs)
 
-    return worker_actor
+    return worker_actor, task_ref
+
+
+def is_ray_task_alive(task_ref: Any, **kwargs):
+    """Checks ray task status. Returns FALSE if ray task has exited for any reason."""
+    if not RAY_AVAILABLE:
+        raise ImportError("ray is required to query ray tasks")
+
+    ready, _ = ray.wait([task_ref], **kwargs)
+    return not bool(ready)
+
+
+def get_ray_task_error(task_ref: Any, **kwargs) -> Exception | None:
+    """Gets ray task. Returns RayTaskError if ray instance exited with any error, else None."""
+    if not RAY_AVAILABLE:
+        raise ImportError("ray is required to query ray tasks")
+
+    try:
+        ray.get(task_ref, **kwargs)
+    except Exception as e:
+        return e

vllm_omni/entrypoints/omni_stage.py

@@ -36,7 +36,12 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.distributed.omni_connectors import build_stage_connectors
 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
 from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase
-from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
+from vllm_omni.distributed.ray_utils.utils import (
+    get_ray_task_error,
+    is_ray_task_alive,
+    kill_ray_actor,
+    start_ray_actor,
+)
 from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs, OmniEngineArgs
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
@@ -303,6 +308,8 @@ class OmniStage:
         self._in_q: mp.queues.Queue | ZmqQueue | str | None = None
         self._out_q: mp.queues.Queue | ZmqQueue | str | None = None
         self._proc: mp.Process | None = None
+        self._ray_actor: Any | None = None
+        self._ray_task_ref: Any | None = None
         self._shm_threshold_bytes: int = 65536
         self._stage_init_timeout: int = stage_init_timeout
         # Callback used by the orchestrator's output handler to stash
@@ -544,7 +551,7 @@ class OmniStage:
             os.environ["VLLM_LOGGING_PREFIX"] = new_env
             if worker_backend == "ray":
                 if is_async:
-                    self._ray_actor = start_ray_actor(
+                    self._ray_actor, self._ray_task_ref = start_ray_actor(
                         _stage_worker_async_entry,
                         ray_placement_group,
                         self.stage_id,
@@ -556,7 +563,7 @@ class OmniStage:
                         stage_init_timeout=self._stage_init_timeout,
                     )
                 else:
-                    self._ray_actor = start_ray_actor(
+                    self._ray_actor, self._ray_task_ref = start_ray_actor(
                         _stage_worker,
                         ray_placement_group,
                         self.stage_id,
@@ -629,9 +636,10 @@ class OmniStage:
                     pass
             self._out_q = None
 
-        if hasattr(self, "_ray_actor") and self._ray_actor:
+        if self._ray_actor is not None:
             kill_ray_actor(self._ray_actor)
             self._ray_actor = None
+            self._ray_task_ref = None
         elif self._proc is not None:
             try:
                 self._proc.join(timeout=5)
@@ -701,6 +709,9 @@ class OmniStage:
             raise
         if self._proc is not None and not self._proc.is_alive():
             raise RuntimeError(f"OmniStage Worker process died unexpectedly with exit code {self._proc.exitcode}")
+        if self._ray_task_ref is not None and not is_ray_task_alive(self._ray_task_ref, timeout=0):
+            e = get_ray_task_error(self._ray_task_ref, timeout=0)
+            raise RuntimeError("OmniStage Ray actor died unexpectedly") from e
 
     def process_engine_inputs(
         self,