From 0e005817f766be66edc68846ff82c9cdfff8015c Mon Sep 17 00:00:00 2001 From: Daniel Huang Date: Wed, 11 Mar 2026 21:08:34 -0700 Subject: [PATCH] [Enhancement] Patch OmniStage.try_collect() with ray alive checks (#1561) Signed-off-by: Daniel Huang Signed-off-by: Daniel Huang Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- vllm_omni/distributed/ray_utils/utils.py | 24 ++++++++++++++++++++++-- vllm_omni/entrypoints/omni_stage.py | 19 +++++++++++++++---- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/vllm_omni/distributed/ray_utils/utils.py b/vllm_omni/distributed/ray_utils/utils.py index 07513b660..519460256 100644 --- a/vllm_omni/distributed/ray_utils/utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 19091034a..3f5a14e97 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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,