diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index c495f0832..9c10d46b4 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/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, diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 2360a68fc..adcdc3e97 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/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 diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index f8d2b364a..c5abaf59c 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/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.""" diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 62eff7ecc..b6f881d1c 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 78f2b7985..7bd27be81 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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():