[Enhancement] Engine runtime errors (#2426)

Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
Signed-off-by: Daniel Huang <pilotflyer824@gmail.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Chenguang Zheng <645327136@qq.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: bjf-frz <frz123db@gmail.com>

tests/engine/test_async_omni_engine_outputs.py

@@ -63,3 +63,39 @@ async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: Mocke
 
     with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"):
         await engine.try_get_output_async()
+
+
+def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixture):
+    """When the orchestrator thread crashes, it enqueues a fatal error message.
+
+    ``try_get_output`` must return this message so the caller
+    (``OmniBase._handle_output_message``) can detect the fatal flag.
+    """
+    fatal_msg = {"type": "error", "error": "Orchestrator thread crashed", "fatal": True}
+
+    mock_queue = mocker.MagicMock()
+    mock_queue.sync_q.get.return_value = fatal_msg
+
+    engine = _make_engine(mock_queue, mocker, thread_alive=False)
+
+    msg = engine.try_get_output()
+    assert msg is not None
+    assert msg["type"] == "error"
+    assert msg["fatal"] is True
+    assert "crashed" in msg["error"]
+
+
+@pytest.mark.asyncio
+async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker: MockerFixture):
+    """Async variant of the fatal error message test."""
+    fatal_msg = {"type": "error", "error": "Orchestrator thread crashed", "fatal": True}
+
+    mock_queue = mocker.MagicMock()
+    mock_queue.sync_q.get_nowait.return_value = fatal_msg
+
+    engine = _make_engine(mock_queue, mocker, thread_alive=False)
+
+    msg = await engine.try_get_output_async()
+    assert msg is not None
+    assert msg["type"] == "error"
+    assert msg["fatal"] is True

tests/engine/test_orchestrator_error_handling.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for error propagation paths within the Orchestrator.
+
+Covers:
+- EngineDeadError from an LLM stage poll → fatal error broadcast + shutdown
+- Diffusion stage error output (OmniRequestOutput.from_error) → routed correctly
+"""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import time
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+from .test_orchestrator import (
+    FakeStageClient,
+    OrchestratorFixture,
+    _build_harness,
+    _enqueue_add_request,
+    _wait_for,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _sampling_params(max_tokens: int = 4):
+    from vllm.sampling_params import SamplingParams
+
+    return SamplingParams(max_tokens=max_tokens)
+
+
+async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+    """Like _get_output_message but returns any message type (including errors)."""
+    deadline = time.monotonic() + timeout
+    while True:
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for orchestrator output")
+        try:
+            return fixture.output_sync_q.get_nowait()
+        except queue.Empty:
+            await asyncio.sleep(0.01)
+
+
+@pytest.fixture
+def orchestrator_factory():
+    fixtures: list[OrchestratorFixture] = []
+
+    def _factory(*args, **kwargs) -> OrchestratorFixture:
+        fixture = _build_harness(*args, **kwargs)
+        fixtures.append(fixture)
+        return fixture
+
+    yield _factory
+
+    for fixture in fixtures:
+        if fixture.thread.is_alive():
+            fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            fixture.thread.join(timeout=5)
+        for q in fixture.queues:
+            q.close()
+
+
+# ───────── EngineDeadError from LLM stage poll ─────────
+
+
+class FakeDeadLLMStageClient(FakeStageClient):
+    """LLM stage client that raises EngineDeadError on get_output_async."""
+
+    async def get_output_async(self):
+        raise EngineDeadError("Stage-0 engine core is dead")
+
+
+@pytest.mark.asyncio
+async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_factory) -> None:
+    """When a stage raises EngineDeadError during poll, the orchestrator must:
+    1. Enqueue a fatal error message for each affected request
+    2. Shut itself down (thread exits)
+    """
+    stage0 = FakeDeadLLMStageClient(stage_type="llm", final_output=True)
+    orchestrator_fixture = orchestrator_factory([stage0])
+    request = SimpleNamespace(request_id="req-dead", prompt_token_ids=[1, 2])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-dead",
+            prompt=request,
+            original_prompt={"prompt": "hello"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+
+        # Collect the fatal error message.
+        msg = await _get_any_output_message(orchestrator_fixture)
+
+        assert msg["type"] == "error"
+        assert msg["fatal"] is True
+        assert msg["request_id"] == "req-dead"
+        assert "Stage-0 engine core is dead" in msg["error"]
+
+        # The orchestrator thread should exit after the fatal error.
+        orchestrator_fixture.thread.join(timeout=5)
+        assert not orchestrator_fixture.thread.is_alive()
+
+        # Request state should be cleaned up.
+        assert "req-dead" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        if orchestrator_fixture.thread.is_alive():
+            orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            orchestrator_fixture.thread.join(timeout=5)
+
+
+# ───────── Diffusion stage error output routing ─────────
+
+
+@pytest.mark.asyncio
+async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -> None:
+    """When a diffusion stage returns an OmniRequestOutput with a non-None
+    error, the orchestrator must route it as a finished output message and
+    clean up the request state.
+    """
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-err",
+            prompt={"prompt": "draw a cat"},
+            original_prompt={"prompt": "draw a cat"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        # Push an error output from the diffusion stage.
+        stage0.push_diffusion_output(OmniRequestOutput.from_error("req-err", "gpu fault"))
+
+        msg = await _get_any_output_message(orchestrator_fixture)
+
+        assert msg["type"] == "output"
+        assert msg["request_id"] == "req-err"
+        assert msg["finished"] is True
+        assert msg["engine_outputs"].error == "gpu fault"
+
+        # Request state should be cleaned up.
+        await _wait_for(lambda: "req-err" not in orchestrator_fixture.orchestrator.request_states)
+    finally:
+        orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+        orchestrator_fixture.thread.join(timeout=5)

tests/engine/test_stage_engine_core_client.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for StageEngineCoreClient.check_health().
+
+Uses object.__new__ to construct a minimal client — check_health only
+touches self.resources, self.stage_id, and self._proc.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_client(*, engine_dead=False, proc_alive=True):
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.resources = SimpleNamespace(engine_dead=engine_dead)
+    client._proc = MagicMock(is_alive=MagicMock(return_value=proc_alive), exitcode=1)
+    return client
+
+
+def test_check_health_passes_when_alive():
+    client = _make_client(engine_dead=False, proc_alive=True)
+    client.check_health()  # no exception
+
+
+def test_check_health_raises_when_resources_engine_dead():
+    client = _make_client(engine_dead=True, proc_alive=True)
+    with pytest.raises(EngineDeadError, match="engine core is dead"):
+        client.check_health()
+
+
+def test_check_health_raises_when_proc_not_alive():
+    client = _make_client(engine_dead=False, proc_alive=False)
+    with pytest.raises(EngineDeadError, match="not alive"):
+        client.check_health()
+    # Verify it set resources.engine_dead as a side effect
+    assert client.resources.engine_dead is True