[Bugfix] Support diffusion worker dead detect when use inline engine (#3214)

Signed-off-by: wuhang <wuhang6@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -5,6 +5,7 @@ from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
@@ -19,6 +20,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def mock_engine():
     with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
         engine_instance = MagicMock()
+        engine_instance.executor = SimpleNamespace(
+            register_failure_callback=MagicMock(),
+            check_health=MagicMock(),
+        )
         mock.make_engine.return_value = engine_instance
         yield engine_instance
 
@@ -97,6 +102,41 @@ def test_inline_shutdown(client, mock_engine):
     mock_engine.close.assert_called_once()
 
 
+def test_inline_registers_executor_failure_callback(client, mock_engine):
+    mock_engine.executor.register_failure_callback.assert_called_once()
+
+
+def test_inline_executor_failure_marks_engine_dead(client, mock_engine):
+    callback = mock_engine.executor.register_failure_callback.call_args.args[0]
+
+    callback()
+
+    assert client._engine_dead is True
+    with pytest.raises(EngineDeadError, match="inline diffusion engine is dead"):
+        client.get_diffusion_output_nowait()
+
+
+def test_inline_returns_queued_output_before_engine_dead(client, mock_engine):
+    callback = mock_engine.executor.register_failure_callback.call_args.args[0]
+    output = OmniRequestOutput.from_diffusion(request_id="req-queued", images=[])
+    client._output_queue.put_nowait(output)
+
+    callback()
+
+    assert client.get_diffusion_output_nowait() is output
+    with pytest.raises(EngineDeadError, match="inline diffusion engine is dead"):
+        client.get_diffusion_output_nowait()
+
+
+def test_inline_check_health_marks_engine_dead(client, mock_engine):
+    mock_engine.executor.check_health.side_effect = EngineDeadError()
+
+    with pytest.raises(EngineDeadError):
+        client.check_health()
+
+    assert client._engine_dead is True
+
+
 def test_inline_client_requires_replica_id(mock_engine):
     metadata = SimpleNamespace(
         stage_id=0,

tests/entrypoints/test_omni_entrypoints.py

@@ -943,6 +943,20 @@ def test_omni_base_errored_false_when_alive():
     assert base.errored is False
 
 
+def test_omni_base_is_running_false_when_stage_engine_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [SimpleNamespace(_engine_dead=True)]
+    assert base.is_running is False
+
+
+def test_omni_base_is_running_false_when_stage_resources_engine_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [SimpleNamespace(resources=SimpleNamespace(engine_dead=True))]
+    assert base.is_running is False
+
+
 def test_omni_base_errored_true_when_orchestrator_dead():
     base = _make_base()
     base.engine.is_alive.return_value = False

vllm_omni/diffusion/executor/abstract.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -101,6 +102,14 @@ class DiffusionExecutor(ABC):
         """Check if the executor and workers are healthy."""
         pass
 
+    def register_failure_callback(self, callback: Callable[[], None]) -> None:
+        """Register a callback invoked when the executor fatally fails.
+
+        Executors without a background failure monitor can keep the default
+        no-op implementation.
+        """
+        return None
+
     @abstractmethod
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -64,8 +64,11 @@ class InlineStageDiffusionClient(StageClientBase):
 
         self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
         self._tasks: dict[str, asyncio.Task] = {}
+        self._engine_dead = False
         self._shutting_down = False
 
+        self._engine.executor.register_failure_callback(self._mark_engine_dead)
+
         logger.info(
             "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)",
             self.stage_id,
@@ -77,6 +80,16 @@ class InlineStageDiffusionClient(StageClientBase):
         """Load model metadata from HuggingFace and populate od_config fields."""
         self.od_config.enrich_config()
 
+    def _mark_engine_dead(self) -> None:
+        if self._engine_dead:
+            return
+        self._engine_dead = True
+        logger.error(
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] diffusion executor died unexpectedly.",
+            self.stage_id,
+            self.replica_id,
+        )
+
     # ------------------------------------------------------------------
     # Request processing
     # ------------------------------------------------------------------
@@ -249,6 +262,8 @@ class InlineStageDiffusionClient(StageClientBase):
         try:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
+            if self._engine_dead:
+                raise EngineDeadError(f"Stage-{self.stage_id} inline diffusion engine is dead")
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
@@ -351,7 +366,11 @@ class InlineStageDiffusionClient(StageClientBase):
         """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()
+        try:
+            self._engine.executor.check_health()
+        except EngineDeadError:
+            self._mark_engine_dead()
+            raise
 
     def shutdown(self) -> None:
         self._shutting_down = True

vllm_omni/entrypoints/omni_base.py

@@ -218,9 +218,18 @@ class OmniBase(PDDisaggregationMixin):
         """Expose engine stage configs for PD disaggregation detection and validation."""
         return self.engine.stage_configs
 
+    def _has_dead_stage(self) -> bool:
+        for stage_client in self.engine.stage_clients:
+            if getattr(stage_client, "_engine_dead", False):
+                return True
+            resources = getattr(stage_client, "resources", None)
+            if resources is not None and getattr(resources, "engine_dead", False):
+                return True
+        return False
+
     @property
     def is_running(self) -> bool:
-        return self.engine.is_alive()
+        return self.engine.is_alive() and not self._has_dead_stage()
 
     @property
     def errored(self) -> bool:
@@ -233,15 +242,7 @@ class OmniBase(PDDisaggregationMixin):
         ``resources.engine_dead`` (StageEngineCoreClient / AsyncMPClient)
         since the two client types store the flag differently.
         """
-        if not self.engine.is_alive():
-            return True
-        for stage_client in self.engine.stage_clients:
-            if getattr(stage_client, "_engine_dead", False):
-                return True
-            resources = getattr(stage_client, "resources", None)
-            if resources is not None and getattr(resources, "engine_dead", False):
-                return True
-        return False
+        return not self.engine.is_alive() or self._has_dead_stage()
 
     def check_health(self) -> None:
         if not self.engine.is_alive():