[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/diffusion/test_multiproc_engine_concurrency.py

@@ -1,17 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
+import multiprocessing as mp
 import queue
 import threading
+import time
 from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import torch
+import zmq
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.sched import RequestScheduler
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
@@ -51,6 +59,8 @@ def _make_executor(num_gpus: int = 1):
     executor._result_mq = mock_rmq
     executor._closed = False
     executor._processes = []
+    executor.is_failed = False
+    executor._failure_callbacks = []
     return executor, req_q, res_q
 
 
@@ -334,8 +344,9 @@ class TestSerialEngineOperations:
 
 class TestCollectiveRpcTimeoutWhileLockHeld:
     """``collective_rpc(timeout=...)`` must honour its timeout even when
-    another thread holds ``engine._rpc_lock`` indefinitely (e.g. a stalled
-    ``add_req`` waiting on an unresponsive worker).
+    another thread holds ``engine._rpc_lock`` indefinitely (e.g. request
+    execution stalled on ``add_req_and_wait_for_response`` → ``execute_fn``
+    → ``collective_rpc`` while blocked on an unresponsive worker).
     """
 
     def test_rpc_times_out_when_lock_held_directly(self):
@@ -359,10 +370,10 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         with pytest.raises(TimeoutError):
             engine.collective_rpc("health", timeout=0.5)
 
-    def test_rpc_times_out_when_add_req_stalled_on_worker(self):
+    def test_rpc_times_out_when_request_execution_stalled_on_worker(self):
         """Real-world scenario the bot flagged:
 
-        ``add_req`` holds ``_rpc_lock`` while blocked on
+        The scheduler/execute path holds ``_rpc_lock`` while blocked on
         ``executor._result_mq.dequeue()`` because the worker never replies.
         A concurrent ``collective_rpc(timeout=...)`` must still time out
         instead of hanging forever waiting for the lock.
@@ -428,3 +439,353 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         t.join(5)
 
         assert result.error == "ok"
+
+
+# ───────── error handling: EngineDeadError propagation through layers ─────
+
+
+class TestMultiprocExecutorRaisesEngineDeadError:
+    """``collective_rpc`` raises ``EngineDeadError`` when the engine is failed."""
+
+    def test_collective_rpc_raises_when_is_failed(self):
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = False
+        executor._broadcast_mq = MagicMock()
+        executor._result_mq = MagicMock()
+        executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError)
+        executor.is_failed = True
+
+        with pytest.raises(EngineDeadError):
+            executor.collective_rpc(
+                "generate",
+                args=(MagicMock(),),
+                unique_reply_rank=0,
+                exec_all_ranks=True,
+            )
+
+    def test_collective_rpc_raises_mid_dequeue_when_is_failed(self):
+        """Worker dies while we are polling the dequeue loop."""
+        executor, _, res_q = _make_executor()
+
+        call_count = 0
+        orig_dequeue = executor._result_mq.dequeue
+
+        def _dying_dequeue(timeout=None):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                executor.is_failed = True
+                raise TimeoutError
+            return orig_dequeue(timeout=timeout)
+
+        executor._result_mq.dequeue = _dying_dequeue
+
+        with pytest.raises(EngineDeadError):
+            executor.collective_rpc(
+                "generate",
+                args=(MagicMock(),),
+                unique_reply_rank=0,
+                exec_all_ranks=True,
+            )
+
+
+class TestDiffusionEngineDeadErrorPassthrough:
+    """``DiffusionEngine.add_req_and_wait_for_response`` re-raises
+    ``EngineDeadError`` from executor and wraps other errors."""
+
+    def test_engine_dead_error_propagates(self):
+        engine, executor, _, _ = _make_engine()
+        engine.execute_fn = Mock(side_effect=EngineDeadError())
+
+        with pytest.raises(EngineDeadError):
+            engine.add_req_and_wait_for_response(_mock_request("dead"))
+
+    def test_runtime_error_wrapped_in_output(self):
+        engine, executor, _, _ = _make_engine()
+        engine.execute_fn = Mock(side_effect=RuntimeError("gpu fault"))
+
+        out = engine.add_req_and_wait_for_response(_mock_request("fault"))
+        assert isinstance(out, DiffusionOutput)
+        assert "gpu fault" in out.error
+
+
+class TestStageDiffusionClientErrorPropagation:
+    """Error surface behaviour of ``StageDiffusionClient``.
+
+    Uses ``object.__new__`` to construct a client without spawning a real
+    subprocess, then manually sets the fields needed for each test.
+    """
+
+    def _make_client(self, *, engine_dead=False, proc_alive=True):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client.final_output = True
+        client.final_output_type = "image"
+        client.default_sampling_params = None
+        client.custom_process_input_func = None
+        client.engine_input_source = None
+
+        client._output_queue = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._tasks = {}
+        client._shutting_down = False
+        client._engine_dead = engine_dead
+        client._owns_process = True
+        client._proc = MagicMock(
+            is_alive=MagicMock(return_value=proc_alive),
+            exitcode=1,
+        )
+        client._request_socket = MagicMock()
+        client._response_socket = MagicMock()
+        client._encoder = MagicMock()
+        client._decoder = MagicMock()
+
+        return client
+
+    @pytest.mark.asyncio
+    async def test_add_request_raises_when_dead(self):
+        client = self._make_client(engine_dead=True)
+
+        with pytest.raises(EngineDeadError):
+            await client.add_request_async("req-3", "test prompt", None)
+
+    def test_check_health_raises_when_dead(self):
+        client = self._make_client(engine_dead=True)
+
+        with pytest.raises(EngineDeadError):
+            client.check_health()
+
+    def test_check_health_ok_when_alive(self):
+        client = self._make_client()
+        client.check_health()
+
+    def test_get_output_raises_engine_dead_when_dead(self):
+        """When ``_engine_dead`` is True and the output queue is empty,
+        ``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
+        client = self._make_client(engine_dead=True)
+        # Simulate _drain_responses as a no-op (no ZMQ socket)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        with pytest.raises(EngineDeadError):
+            client.get_diffusion_output_nowait()
+
+    def test_get_output_returns_none_when_alive_and_empty(self):
+        """When the engine is alive and the queue is empty, return None."""
+        client = self._make_client()
+        client._response_socket.recv.side_effect = zmq.Again
+
+        assert client.get_diffusion_output_nowait() is None
+
+    def test_check_health_raises_when_proc_dead(self):
+        """``check_health`` detects a dead subprocess via ``_proc.is_alive()``
+        and raises ``EngineDeadError``, setting ``_engine_dead`` as a
+        side effect."""
+        client = self._make_client(proc_alive=False)
+
+        with pytest.raises(EngineDeadError, match="not alive"):
+            client.check_health()
+
+        assert client._engine_dead is True
+
+    def test_get_output_raises_when_proc_dead(self):
+        """When the subprocess has died (non-signal exit) and the output
+        queue is empty, ``get_diffusion_output_nowait`` must raise
+        ``EngineDeadError`` with the exit code."""
+        client = self._make_client(proc_alive=False)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        with pytest.raises(EngineDeadError, match="exit code"):
+            client.get_diffusion_output_nowait()
+
+        assert client._engine_dead is True
+
+    def test_get_output_returns_none_on_signal_death(self):
+        """When the subprocess was killed by a signal (exit code > 128),
+        ``get_diffusion_output_nowait`` returns ``None`` and sets
+        ``_shutting_down`` instead of raising."""
+        client = self._make_client(proc_alive=False)
+        client._proc.exitcode = 137  # SIGKILL (128 + 9)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        result = client.get_diffusion_output_nowait()
+
+        assert result is None
+        assert client._shutting_down is True
+        assert client._engine_dead is True
+
+
+# ───────── monitor thread & death sentinel integration tests ─────────
+
+
+def _poll_flag(get_flag, *, timeout=5.0, interval=0.05) -> bool:
+    """Poll until ``get_flag()`` returns True or *timeout* elapses."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if get_flag():
+            return True
+        time.sleep(interval)
+    return False
+
+
+def _make_short_lived_process() -> mp.Process:
+    """Spawn a real subprocess that exits immediately.
+
+    The process must be started with ``"fork"`` (or the platform default)
+    so that it can use a plain ``lambda`` as its target — ``"spawn"`` would
+    fail to pickle it.
+    """
+    ctx = mp.get_context("fork")
+    p = ctx.Process(target=lambda: None, name="ShortLivedWorker-0")
+    p.start()
+    return p
+
+
+class TestMultiprocExecutorWorkerMonitor:
+    """Integration tests for ``start_worker_monitor``.
+
+    Uses real short-lived subprocesses so that OS-level sentinel fd
+    readiness is exercised end-to-end.
+    """
+
+    def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
+        """When a worker process dies, the monitor thread must:
+        1. Set ``is_failed = True``
+        2. Call ``shutdown()`` (which sets ``_closed = True``)
+        3. Invoke all registered failure callbacks
+        """
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = False
+        executor.is_failed = False
+        executor._failure_callbacks = []
+        executor._broadcast_mq = None
+        executor._result_mq = None
+        executor.resources = None
+        # Use a no-op so shutdown() doesn't crash on None resources.
+        executor._finalizer = lambda: None
+
+        proc = _make_short_lived_process()
+        executor._processes = [proc]
+
+        callback_called = threading.Event()
+        executor.register_failure_callback(callback_called.set)
+
+        executor.start_worker_monitor()
+
+        # Wait for the process to exit and the monitor to react.
+        proc.join(5)
+        assert _poll_flag(lambda: executor.is_failed), "is_failed was not set"
+        assert executor._closed, "shutdown() was not called"
+        assert callback_called.wait(timeout=2), "failure callback was not invoked"
+
+    def test_worker_monitor_noop_when_already_closed(self):
+        """If ``_closed`` is already True when the process dies (orderly
+        shutdown), the monitor must *not* set ``is_failed``."""
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = True  # already shut down
+        executor.is_failed = False
+        executor._failure_callbacks = []
+        executor._broadcast_mq = None
+        executor._result_mq = None
+        executor.resources = None
+        executor._finalizer = lambda: None
+
+        proc = _make_short_lived_process()
+        executor._processes = [proc]
+
+        executor.start_worker_monitor()
+        proc.join(5)
+
+        # Give the monitor thread a chance to run (it should early-return).
+        time.sleep(0.3)
+        assert not executor.is_failed, "is_failed should remain False on orderly shutdown"
+
+
+class TestStageDiffusionClientProcMonitor:
+    """Integration test for ``StageDiffusionClient._start_proc_monitor``.
+
+    Uses a real short-lived subprocess to verify the sentinel-based
+    detection pipeline.
+    """
+
+    def test_proc_monitor_sets_engine_dead_on_process_death(self):
+        """When the subprocess dies, the monitor thread must set
+        ``_engine_dead = True``."""
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client._shutting_down = False
+        client._engine_dead = False
+
+        proc = _make_short_lived_process()
+        client._proc = proc
+
+        client._start_proc_monitor()
+        proc.join(5)
+
+        assert _poll_flag(lambda: client._engine_dead), "_engine_dead was not set"
+
+
+class TestDrainResponsesDeathSentinel:
+    """Tests for death sentinel and error routing in
+    ``StageDiffusionClient._drain_responses()``.
+    """
+
+    def _make_client(self):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client._engine_dead = False
+        client._shutting_down = False
+        client._output_queue = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._response_socket = MagicMock()
+        client._decoder = MagicMock()
+        return client
+
+    def test_drain_responses_sets_engine_dead_on_death_sentinel(self):
+        """When ``_drain_responses`` receives the ``DIFFUSION_PROC_DEAD``
+        sentinel, it must set ``_engine_dead = True`` and stop draining
+        (decoder is never called)."""
+        client = self._make_client()
+
+        # First recv returns the death sentinel, second would be a normal
+        # message but should never be reached.
+        client._response_socket.recv.side_effect = [
+            StageDiffusionProc.DIFFUSION_PROC_DEAD,
+            b"should-not-be-reached",
+        ]
+
+        client._drain_responses()
+
+        assert client._engine_dead is True
+        client._decoder.decode.assert_not_called()
+
+    def test_drain_responses_routes_error_as_omni_request_output(self):
+        """When ``_drain_responses`` receives a ``{"type": "error"}`` message
+        with a ``request_id``, it must place an ``OmniRequestOutput`` with
+        the error on ``_output_queue``."""
+        client = self._make_client()
+
+        error_msg = {
+            "type": "error",
+            "request_id": "req-fail",
+            "error": "gpu fault",
+        }
+        # First recv returns the encoded error, second raises zmq.Again.
+        client._response_socket.recv.side_effect = [b"encoded-error", zmq.Again]
+        client._decoder.decode.return_value = error_msg
+
+        client._drain_responses()
+
+        assert not client._output_queue.empty()
+        output = client._output_queue.get_nowait()
+        assert isinstance(output, OmniRequestOutput)
+        assert output.request_id == "req-fail"
+        assert output.error == "gpu fault"
+        assert output.finished is True

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

tests/entrypoints/openai_api/test_image_server.py

@@ -407,6 +407,29 @@ def test_health_endpoint_no_engine():
     assert data["status"] == "unhealthy"
 
 
+def test_health_endpoint_dead_engine():
+    """Health returns 503 when the engine raises EngineDeadError."""
+    from unittest.mock import AsyncMock
+
+    from fastapi import FastAPI
+    from vllm.v1.engine.exceptions import EngineDeadError
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    dead_engine = AsyncMock()
+    dead_engine.check_health = AsyncMock(side_effect=EngineDeadError())
+    app.state.engine_client = dead_engine
+
+    client = TestClient(app)
+    response = client.get("/health")
+    assert response.status_code == 503
+    data = response.json()
+    assert data["status"] == "unhealthy"
+
+
 def test_models_endpoint(test_client):
     """Test /v1/models endpoint for diffusion mode"""
     response = test_client.get("/v1/models")

tests/entrypoints/openai_api/test_video_server.py

@@ -133,7 +133,6 @@ def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s:
     last_payload = None
     while time.time() < deadline:
         response = client.get(f"/v1/videos/{video_id}")
-        assert response.status_code == 200
         last_payload = response.json()
         if last_payload["status"] == status:
             return last_payload
@@ -644,7 +643,7 @@ def test_invalid_lora_returns_400(test_client):
     assert response.status_code == 200
     video_id = response.json()["id"]
     failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
-    assert failed["error"]["code"] == "HTTPException"
+    assert failed["error"]["code"] == 400
     assert "lora object" in failed["error"]["message"].lower()
 
 

tests/entrypoints/test_omni_entrypoints.py

@@ -4,11 +4,13 @@ import queue
 from collections.abc import Callable
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import MagicMock
 
 import pytest
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.omni import Omni
@@ -163,6 +165,15 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine)
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
 
 
+def _make_base():
+    from vllm_omni.entrypoints.omni_base import OmniBase
+
+    obj = object.__new__(OmniBase)
+    obj.engine = MagicMock()
+    obj.request_states = {}
+    return obj
+
+
 def _stage_spec(
     stage_id: int,
     *,
@@ -687,3 +698,255 @@ def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):
     assert submitted_params[1].output_kind == RequestOutputKind.FINAL_ONLY
     assert submitted_params[2].output_kind == original_diffusion_output_kind
     assert len(outputs) == 2
+
+
+def test_fatal_error_raises_engine_dead():
+    base = _make_base()
+    msg = {"type": "error", "error": "orchestrator crashed", "fatal": True}
+
+    with pytest.raises(EngineDeadError, match="orchestrator crashed"):
+        base._handle_output_message(msg)
+
+
+def test_non_fatal_error_raises_runtime():
+    base = _make_base()
+    msg = {"type": "error", "error": "something wrong"}
+
+    with pytest.raises(RuntimeError, match="something wrong"):
+        base._handle_output_message(msg)
+
+
+def test_async_omni_errored_property_alive():
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        is_alive=lambda: True,
+        stage_clients=[SimpleNamespace(is_comprehension=False)],
+    )
+
+    assert omni.errored is False
+
+
+def test_async_omni_errored_property_dead_engine():
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        is_alive=lambda: False,
+        stage_clients=[SimpleNamespace(is_comprehension=False)],
+    )
+
+    assert omni.errored is True
+
+
+def test_async_omni_errored_property_dead_stage():
+    omni = object.__new__(AsyncOmni)
+    dead_stage = SimpleNamespace(is_comprehension=False, _engine_dead=True)
+    omni.engine = SimpleNamespace(
+        is_alive=lambda: True,
+        stage_clients=[dead_stage],
+    )
+
+    assert omni.errored is True
+
+
+def _enqueue_stage_error(
+    engine: FakeAsyncOmniEngine,
+    msg,
+    *,
+    error_text: str,
+    kill_engine: bool = False,
+):
+    """Enqueue a stage error output, optionally killing the engine."""
+    if kill_engine:
+        engine._alive = False
+    engine.output_q.put_nowait(
+        {
+            "type": "output",
+            "request_id": msg["request_id"],
+            "stage_id": 0,
+            "engine_outputs": SimpleNamespace(
+                payload="",
+                finished=True,
+                images=[],
+                stage_durations={},
+                error=error_text,
+            ),
+            "finished": False,
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_omni_propagates_engine_dead_error(monkeypatch: pytest.MonkeyPatch):
+    """When the engine is dead and an error output arrives, ``generate()``
+    must raise ``EngineDeadError`` (not plain ``RuntimeError``)."""
+
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=lambda eng, msg: _enqueue_stage_error(eng, msg, error_text="worker OOM", kill_engine=True),
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        with pytest.raises(EngineDeadError, match="worker OOM"):
+            async for _ in app.generate(prompt="hello", request_id="req-dead"):
+                pass
+    finally:
+        app.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_async_omni_propagates_engine_generate_error(monkeypatch: pytest.MonkeyPatch):
+    """When the engine is alive but a stage error occurs, ``generate()``
+    must raise ``EngineGenerateError`` (recoverable, not ``EngineDeadError``)."""
+
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=lambda eng, msg: _enqueue_stage_error(eng, msg, error_text="diffusion step failed"),
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        with pytest.raises(EngineGenerateError):
+            async for _ in app.generate(prompt="hello", request_id="req-recover"):
+                pass
+    finally:
+        app.shutdown()
+
+
+# ───────── OmniBase.check_health() aggregation ─────────
+
+
+def test_check_health_passes_when_all_healthy():
+    base = _make_base()
+    healthy_stage = MagicMock()
+    healthy_stage.check_health = MagicMock()
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [healthy_stage]
+    base.check_health()  # should not raise
+
+
+def test_check_health_raises_when_stage_dead():
+    base = _make_base()
+    dead_stage = MagicMock()
+    dead_stage.check_health = MagicMock(side_effect=EngineDeadError("Stage-1 dead"))
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [dead_stage]
+    with pytest.raises(EngineDeadError, match="Stage-1 dead"):
+        base.check_health()
+
+
+def test_check_health_raises_when_orchestrator_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = False
+    base.engine.stage_clients = []
+    with pytest.raises(EngineDeadError, match="not alive"):
+        base.check_health()
+
+
+# ───────── OmniBase.errored property ─────────
+
+
+def test_omni_base_errored_false_when_alive():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [SimpleNamespace()]
+    assert base.errored is False
+
+
+def test_omni_base_errored_true_when_orchestrator_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = False
+    base.engine.stage_clients = []
+    assert base.errored is True
+
+
+def test_omni_base_errored_true_when_stage_engine_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    dead_stage = SimpleNamespace(_engine_dead=True)
+    base.engine.stage_clients = [dead_stage]
+    assert base.errored is True
+
+
+def test_omni_base_errored_true_when_stage_resources_engine_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    dead_stage = SimpleNamespace(resources=SimpleNamespace(engine_dead=True))
+    base.engine.stage_clients = [dead_stage]
+    assert base.errored is True
+
+
+# ───────── Omni (sync) EngineDeadError / EngineGenerateError ─────────
+
+
+def test_omni_propagates_engine_dead_error(monkeypatch: pytest.MonkeyPatch):
+    """When the engine is dead and a stage error output arrives,
+    ``Omni.generate()`` must raise ``EngineDeadError``."""
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=lambda eng, msg: _enqueue_stage_error(eng, msg, error_text="worker OOM", kill_engine=True),
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        with pytest.raises(EngineDeadError, match="worker OOM"):
+            list(app.generate(["hello"], py_generator=False, use_tqdm=False))
+    finally:
+        app.shutdown()
+
+
+def test_omni_propagates_engine_generate_error(monkeypatch: pytest.MonkeyPatch):
+    """When the engine is alive but a stage error occurs,
+    ``Omni.generate()`` must raise ``EngineGenerateError`` (recoverable)."""
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=lambda eng, msg: _enqueue_stage_error(eng, msg, error_text="diffusion step failed"),
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        with pytest.raises(EngineGenerateError):
+            list(app.generate(["hello"], py_generator=False, use_tqdm=False))
+    finally:
+        app.shutdown()
+
+
+def test_omni_errored_property_alive(monkeypatch: pytest.MonkeyPatch):
+    """Omni.errored (inherited from OmniBase) returns False when healthy."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        assert app.errored is False
+    finally:
+        app.shutdown()
+
+
+def test_omni_errored_property_dead_engine(monkeypatch: pytest.MonkeyPatch):
+    """Omni.errored returns True when the orchestrator is dead."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        engine._alive = False
+        assert app.errored is True
+    finally:
+        app.shutdown()
+
+
+def test_omni_errored_property_dead_stage(monkeypatch: pytest.MonkeyPatch):
+    """Omni.errored returns True when a stage client is marked dead."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        engine.stage_clients[0]._engine_dead = True
+        assert app.errored is True
+    finally:
+        app.shutdown()

vllm_omni/diffusion/diffusion_engine.py

@@ -14,6 +14,7 @@ import numpy as np
 import PIL.Image
 import torch
 from vllm.logger import init_logger
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
@@ -122,7 +123,7 @@ class DiffusionEngine:
         if output.aborted:
             raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
         if output.error:
-            raise RuntimeError(f"{output.error}")
+            raise RuntimeError(output.error)
         logger.info("Generation completed successfully.")
 
         if output.output is None:
@@ -358,6 +359,8 @@ class DiffusionEngine:
                 sched_req_id = sched_output.scheduled_req_ids[0]
                 try:
                     runner_output = self.execute_fn(sched_output)
+                except EngineDeadError:
+                    raise
                 except Exception as exc:
                     logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
                     runner_output = RunnerOutput(

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
 import multiprocessing as mp
+import multiprocessing.connection
+import threading
 import time
 import weakref
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 import zmq
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
@@ -22,6 +26,8 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+_DEQUEUE_TIMEOUT_S = 5.0
+
 
 @dataclass
 class BackgroundResources:
@@ -67,6 +73,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     def _init_executor(self) -> None:
         self._processes: list[mp.Process] = []
         self._closed = False
+        self.is_failed = False
+        self._failure_callbacks: list[Callable[[], None]] = []
 
         num_workers = self.od_config.num_gpus
         self.wake_events = [mp.Event() for _ in range(num_workers)]
@@ -87,6 +95,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         )
         self._finalizer = weakref.finalize(self, self.resources)
 
+        self.start_worker_monitor()
+
     def _init_broadcast_queue(self, num_workers: int) -> MessageQueue:
         return MessageQueue(
             n_reader=num_workers,
@@ -106,6 +116,23 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         if self._result_mq is None:
             raise RuntimeError("Result queue not initialized")
 
+    def _dequeue_one_with_failure_polling(self, deadline: float | None, method: str) -> Any:
+        """Block until one result message, polling ``is_failed`` between chunk timeouts."""
+        while True:
+            if deadline is None:
+                chunk_timeout = _DEQUEUE_TIMEOUT_S
+            else:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise TimeoutError(f"RPC call to {method} timed out.")
+                chunk_timeout = min(_DEQUEUE_TIMEOUT_S, remaining)
+            try:
+                return self._result_mq.dequeue(timeout=chunk_timeout)
+            except (TimeoutError, zmq.error.Again):
+                if self.is_failed:
+                    raise EngineDeadError()
+                continue
+
     def _launch_workers(self, broadcast_handle, wake_events):
         od_config = self.od_config
         logger.info("Starting server...")
@@ -171,6 +198,49 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
         return processes, result_handle
 
+    def start_worker_monitor(self) -> None:
+        # Monitors worker process liveness. If any die unexpectedly,
+        # logs an error, shuts down the executor and invokes the failure
+        # callback to inform the engine.
+        sentinels = [p.sentinel for p in self._processes]
+        if not sentinels:
+            return
+
+        def _monitor() -> None:
+            try:
+                finished = multiprocessing.connection.wait(sentinels)
+            except OSError:
+                return
+
+            if self._closed:
+                return
+
+            dead = [p.name for p in self._processes if p.sentinel in finished]
+            if dead:
+                logger.error(
+                    "Diffusion worker(s) died unexpectedly: %s",
+                    dead,
+                )
+                self.is_failed = True
+
+            self.shutdown()
+
+            for cb in self._failure_callbacks:
+                try:
+                    cb()
+                except Exception:
+                    logger.exception("failure_callback raised")
+
+        t = threading.Thread(target=_monitor, daemon=True, name="diffusion-worker-monitor")
+        t.start()
+
+    def register_failure_callback(
+        self,
+        callback: Callable[[], None],
+    ) -> None:
+        """Register a callback invoked when a worker process dies."""
+        self._failure_callbacks.append(callback)
+
     def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         self._ensure_open()
         rpc_request = {
@@ -293,27 +363,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
             responses = []
             for _ in range(num_responses):
-                dequeue_timeout = None if deadline is None else max(0, deadline - time.monotonic())
+                response = self._dequeue_one_with_failure_polling(deadline, method)
+
                 try:
-                    response = self._result_mq.dequeue(timeout=dequeue_timeout)
+                    unpack_diffusion_output_shm(response)
+                except Exception as e:
+                    logger.warning("SHM unpack failed (data may already be inline): %s", e)
 
-                    try:
-                        unpack_diffusion_output_shm(response)
-                    except Exception as e:
-                        logger.warning("SHM unpack failed (data may already be inline): %s", e)
+                # Check if response indicates an error
+                if isinstance(response, dict) and response.get("status") == "error":
+                    raise RuntimeError(
+                        f"Worker failed with error '{response.get('error')}', "
+                        "please check the stack trace above for the root cause"
+                    )
 
-                    # Check if response indicates an error
-                    if isinstance(response, dict) and response.get("status") == "error":
-                        raise RuntimeError(
-                            f"Worker failed with error '{response.get('error')}', "
-                            "please check the stack trace above for the root cause"
-                        )
-
-                    responses.append(response)
-                except zmq.error.Again as e:
-                    raise TimeoutError(f"RPC call to {method} timed out.") from e
-                except TimeoutError as e:
-                    raise TimeoutError(f"RPC call to {method} timed out.") from e
+                responses.append(response)
 
             return responses[0] if unique_reply_rank is not None else responses
         except Exception as e:
@@ -321,10 +385,13 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             raise
 
     def check_health(self) -> None:
-        # Simple check if processes are alive
+        self._ensure_open()
+        if self.is_failed:
+            raise EngineDeadError()
         for p in self._processes:
             if not p.is_alive():
-                raise RuntimeError(f"Worker process {p.name} is dead")
+                self.is_failed = True
+                raise EngineDeadError(f"Worker process {p.name} is dead")
 
     def shutdown(self) -> None:
         self._closed = True

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
 
 from itertools import chain
 from typing import Any

vllm_omni/diffusion/stage_diffusion_client.py

@@ -8,15 +8,20 @@ from any stage client.
 from __future__ import annotations
 
 import asyncio
+import multiprocessing.connection
 import time
 import uuid
+import weakref
 from dataclasses import fields, is_dataclass
+from threading import Thread
 from typing import TYPE_CHECKING, Any
 
 import zmq
 from vllm.logger import init_logger
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.stage_diffusion_proc import (
+    StageDiffusionProc,
     complete_diffusion_handshake,
     spawn_diffusion_proc,
 )
@@ -129,6 +134,12 @@ class StageDiffusionClient:
         self._pending_rpcs: set[str] = set()
         self._tasks: dict[str, asyncio.Task] = {}
         self._shutting_down = False
+        self._engine_dead: bool = False
+
+        # Background thread to detect silent process death (SIGKILL, segfault)
+        # where the subprocess cannot send the ZMQ death sentinel.
+        # Mirrors MPClient.start_engine_core_monitor() in vLLM.
+        self._start_proc_monitor()
 
         logger.info(
             "[StageDiffusionClient] Stage-%s initialized (owns_process=%s, batch_size=%d)",
@@ -137,6 +148,37 @@ class StageDiffusionClient:
             batch_size,
         )
 
+    # ------------------------------------------------------------------
+    # Process monitor (mirrors vLLM's MPClient.start_engine_core_monitor)
+    # ------------------------------------------------------------------
+
+    def _start_proc_monitor(self) -> None:
+        """Start a daemon thread that watches the subprocess sentinel.
+
+        When the subprocess dies without sending the ZMQ death sentinel
+        (e.g. SIGKILL, segfault), this thread sets ``_engine_dead`` so
+        subsequent calls raise ``EngineDeadError``.
+        """
+        proc = self._proc
+        self_ref = weakref.ref(self)
+
+        def _monitor() -> None:
+            try:
+                multiprocessing.connection.wait([proc.sentinel])
+            except Exception:
+                return
+            client = self_ref()
+            if client is None or client._shutting_down or client._engine_dead:
+                return
+            client._engine_dead = True
+            logger.error(
+                "[StageDiffusionClient] Stage-%s StageDiffusionProc died unexpectedly (exit code %s).",
+                client.stage_id,
+                proc.exitcode,
+            )
+
+        Thread(target=_monitor, daemon=True, name="DiffusionProcMonitor").start()
+
     # ------------------------------------------------------------------
     # Internal helpers
     # ------------------------------------------------------------------
@@ -149,6 +191,15 @@ class StageDiffusionClient:
             except zmq.Again:
                 break
 
+            # Check for the death sentinel (raw bytes, not msgpack-encoded).
+            if raw == StageDiffusionProc.DIFFUSION_PROC_DEAD:
+                self._engine_dead = True
+                logger.error(
+                    "[StageDiffusionClient] Stage-%s received DIFFUSION_PROC_DEAD sentinel from subprocess.",
+                    self.stage_id,
+                )
+                break
+
             msg = self._decoder.decode(raw)
             msg_type = msg.get("type")
 
@@ -172,13 +223,10 @@ class StageDiffusionClient:
                         "error": True,
                         "reason": error_msg,
                     }
-                elif req_id is not None:
-                    error_output = OmniRequestOutput.from_diffusion(
-                        request_id=req_id,
-                        images=[],
-                    )
-                    error_output.error = error_msg
-                    self._output_queue.put_nowait(error_output)
+                # Route request errors as error outputs so the Orchestrator
+                # sees the request complete (instead of hanging forever).
+                if req_id is not None:
+                    self._output_queue.put_nowait(OmniRequestOutput.from_error(req_id, error_msg))
 
     # Fields that are subprocess-local and cannot be serialized across
     # process boundaries.  They are recreated in the subprocess with
@@ -242,6 +290,8 @@ class StageDiffusionClient:
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
+        if self._engine_dead:
+            raise EngineDeadError()
         self._request_socket.send(
             self._encoder.encode(
                 {
@@ -269,6 +319,8 @@ class StageDiffusionClient:
         and the combined result is placed on the output queue with a single
         *request_id*.
         """
+        if self._engine_dead:
+            raise EngineDeadError()
         task = asyncio.create_task(
             self._run_batch(
                 request_id,
@@ -306,6 +358,7 @@ class StageDiffusionClient:
                 request_id,
                 e,
             )
+            await self._output_queue.put(OmniRequestOutput.from_error(request_id, str(e)))
         finally:
             self._tasks.pop(request_id, None)
 
@@ -314,7 +367,10 @@ class StageDiffusionClient:
         try:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
+            if self._engine_dead:
+                raise EngineDeadError()
             if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive():
+                self._engine_dead = True
                 exitcode = self._proc.exitcode
                 # One final drain – the last ZMQ frame may have arrived
                 # between the first drain and the is_alive() check.
@@ -328,7 +384,7 @@ class StageDiffusionClient:
                     logger.warning("StageDiffusionProc was killed by signal %d; treating as external shutdown.", sig)
                     self._shutting_down = True
                     return None
-                raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {exitcode})")
+                raise EngineDeadError(f"StageDiffusionProc died unexpectedly (exit code {exitcode})")
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
@@ -349,6 +405,9 @@ class StageDiffusionClient:
         kwargs: dict[str, Any] | None = None,
     ) -> Any:
         """Forward control RPCs to the diffusion subprocess."""
+        if self._engine_dead:
+            raise EngineDeadError()
+
         # Inject a default profile_prefix that includes stage_id when profiling.
         if method == "profile":
             args_list = list(args)
@@ -386,8 +445,9 @@ class StageDiffusionClient:
                 self._drain_responses()
                 if rpc_id in self._rpc_results:
                     return self._rpc_results.pop(rpc_id)
-                if self._owns_process and self._proc is not None and not self._proc.is_alive():
-                    raise RuntimeError(
+                if self._engine_dead or (self._owns_process and self._proc is not None and not self._proc.is_alive()):
+                    self._engine_dead = True
+                    raise EngineDeadError(
                         f"StageDiffusionProc died while waiting for "
                         f"collective_rpc '{method}' (exit code {self._proc.exitcode})"
                     )
@@ -397,6 +457,19 @@ class StageDiffusionClient:
         finally:
             self._pending_rpcs.discard(rpc_id)
 
+    def check_health(self) -> None:
+        """Raise ``EngineDeadError`` if the diffusion engine is dead.
+
+        Mirrors the ``check_health`` protocol on vLLM's ``EngineClient``.
+        """
+        if self._engine_dead:
+            raise EngineDeadError(f"Stage-{self.stage_id} diffusion subprocess is dead")
+        if self._proc is not None and not self._proc.is_alive():
+            self._engine_dead = True
+            raise EngineDeadError(
+                f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {self._proc.exitcode})."
+            )
+
     def shutdown(self) -> None:
         self._shutting_down = True
         try:

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -46,6 +46,8 @@ class StageDiffusionProc:
     and ZMQ-based communication with StageDiffusionClient.
     """
 
+    DIFFUSION_PROC_DEAD = b"DIFFUSION_PROC_DEAD"
+
     def __init__(self, model: str, od_config: OmniDiffusionConfig) -> None:
         self._model = model
         self._od_config = od_config
@@ -450,6 +452,16 @@ class StageDiffusionProc:
                 elif msg_type == "shutdown":
                     break
 
+        except Exception:
+            # Send the death sentinel so the client can detect the
+            # fatal failure promptly (mirrors EngineCoreProc._send_engine_dead).
+            try:
+                response_socket.setsockopt(zmq.LINGER, 4000)
+                await response_socket.send(StageDiffusionProc.DIFFUSION_PROC_DEAD)
+            except Exception:
+                logger.warning("Failed to send DIFFUSION_PROC_DEAD sentinel to client.")
+            raise
+
         finally:
             for task in tasks.values():
                 task.cancel()

vllm_omni/engine/async_omni_engine.py

@@ -94,6 +94,8 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+_STARTUP_POLL_INTERVAL_S = 1.0
+
 
 # ============================================================================
 # Parent-EngineArgs field-routing contracts (consumed by
@@ -341,22 +343,7 @@ class AsyncOmniEngine:
             name="orchestrator",
         )
         self.orchestrator_thread.start()
-
-        # Wait for stage/runtime initialization result from orchestrator thread.
-        try:
-            startup_future.result(timeout=startup_timeout)
-        except concurrent.futures.TimeoutError as e:
-            try:
-                self.shutdown()
-            except Exception:
-                logger.exception("[AsyncOmniEngine] Failed to cleanup after orchestrator startup timeout")
-            raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") from e
-        except Exception:
-            try:
-                self.shutdown()
-            except Exception:
-                logger.exception("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure")
-            raise
+        self._wait_for_orchestrator_init(startup_future, startup_timeout)
 
         # Stage runtime fields are assigned directly on self by the bootstrap thread.
         self._weak_finalizer = weakref.finalize(
@@ -946,13 +933,17 @@ class AsyncOmniEngine:
             loop.run_until_complete(_run_orchestrator())
         except Exception as e:
             if not startup_future.done():
-                startup_future.set_exception(RuntimeError(f"Orchestrator initialization failed: {e}"))
+                wrapped = RuntimeError(f"Orchestrator initialization failed: {e}")
+                wrapped.__cause__ = e
+                startup_future.set_exception(wrapped)
             logger.exception("[AsyncOmniEngine] Orchestrator thread crashed")
+            error_text = str(e) or "Orchestrator thread crashed"
             try:
+                error_msg = {"type": "error", "error": error_text, "fatal": True}
                 if self.output_queue is not None:
-                    self.output_queue.sync_q.put_nowait({"type": "error", "error": "Orchestrator thread crashed"})
+                    self.output_queue.sync_q.put_nowait(error_msg)
                 if self.rpc_output_queue is not None:
-                    self.rpc_output_queue.sync_q.put_nowait({"type": "error", "error": "Orchestrator thread crashed"})
+                    self.rpc_output_queue.sync_q.put_nowait(error_msg)
             except Exception:
                 pass
             raise
@@ -972,6 +963,31 @@ class AsyncOmniEngine:
                 asyncio.set_event_loop(None)
                 loop.close()
 
+    def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None:
+        """
+        Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process.
+        """
+        deadline = time.monotonic() + startup_timeout
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup timeout")
+                raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s")
+            try:
+                startup_future.result(
+                    timeout=min(remaining, _STARTUP_POLL_INTERVAL_S),
+                )
+                break
+            except concurrent.futures.TimeoutError:
+                if not self.orchestrator_thread.is_alive():
+                    self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure")
+                    if startup_future.done():
+                        startup_future.result()  # re-raises the real exception
+                    raise RuntimeError("Orchestrator thread died during startup")
+            except Exception:
+                self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure")
+                raise
+
     # ---- request helpers ----
 
     def _build_add_request_message(
@@ -1729,3 +1745,9 @@ class AsyncOmniEngine:
             except Exception:
                 logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during shutdown")
             self._omni_master_server = None
+
+    def _try_shutdown(self, *args, **kwargs) -> None:
+        try:
+            self.shutdown()
+        except Exception:
+            logger.exception(*args, **kwargs)

vllm_omni/engine/orchestrator.py

@@ -22,6 +22,7 @@ from vllm.outputs import RequestOutput
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length
 from vllm_omni.engine import (
@@ -179,6 +180,7 @@ class Orchestrator:
         # Shutdown coordination
         self._shutdown_event = asyncio.Event()
         self._stages_shutdown = False
+        self._fatal_error: str | None = None
 
     async def run(self) -> None:
         """Main entry point for the Orchestrator event loop."""
@@ -208,6 +210,12 @@ class Orchestrator:
             except Exception:
                 pass
 
+            # If a fatal error caused the shutdown, drain any pending
+            # add_request messages that were never processed and broadcast
+            # fatal error responses so callers are not left hanging.
+            if self._fatal_error is not None:
+                await self._drain_pending_requests_on_fatal()
+
             self._shutdown_stages()
 
             # Cancel any remaining tasks spawned by wait_for / gather so
@@ -271,6 +279,21 @@ class Orchestrator:
                     output = stage_client.get_diffusion_output_nowait()
                     if output is not None:
                         idle = False
+
+                        if getattr(output, "error", None) is not None:
+                            await self.output_async_queue.put(
+                                {
+                                    "type": "output",
+                                    "request_id": output.request_id,
+                                    "stage_id": stage_id,
+                                    "engine_outputs": output,
+                                    "metrics": None,
+                                    "finished": True,
+                                }
+                            )
+                            self.request_states.pop(output.request_id, None)
+                            continue
+
                         req_state = self.request_states.get(output.request_id)
                         if req_state is not None:
                             if getattr(output, "error", None) is not None:
@@ -301,6 +324,26 @@ class Orchestrator:
                     continue
                 except asyncio.CancelledError:
                     raise
+                except EngineDeadError as e:
+                    logger.error(
+                        "[Orchestrator] Stage-%s is dead: %s",
+                        stage_id,
+                        e,
+                    )
+                    self._fatal_error = str(e)
+                    for req_id, req_state in list(self.request_states.items()):
+                        if stage_id in req_state.stage_submit_ts:
+                            await self.output_async_queue.put(
+                                {
+                                    "type": "error",
+                                    "error": str(e),
+                                    "fatal": True,
+                                    "request_id": req_id,
+                                }
+                            )
+                            self.request_states.pop(req_id, None)
+                    self._shutdown_event.set()
+                    raise
                 except Exception:
                     if self._shutdown_event.is_set():
                         return
@@ -1064,6 +1107,53 @@ class Orchestrator:
             }
         )
 
+    async def _drain_pending_requests_on_fatal(self) -> None:
+        """Drain the request queue and broadcast fatal errors for any
+        pending add_request messages that were never processed.
+
+        Called from the ``run()`` finally block when a fatal error
+        (e.g. ``EngineDeadError``) caused the orchestrator to shut down
+        before the request handler could process all queued messages.
+        Also broadcasts for any already-tracked requests still in
+        ``request_states`` that were not yet notified.
+        """
+        assert self._fatal_error is not None
+
+        notified: set[str] = set()
+
+        # 1) Drain pending messages from the request queue.
+        while True:
+            try:
+                msg = self.request_async_queue.get_nowait()
+            except Exception:
+                break
+            if msg.get("type") == "add_request":
+                req_id = msg["request_id"]
+                await self.output_async_queue.put(
+                    {
+                        "type": "error",
+                        "error": self._fatal_error,
+                        "fatal": True,
+                        "request_id": req_id,
+                    }
+                )
+                notified.add(req_id)
+
+        # 2) Broadcast for any tracked requests not already notified
+        #    (e.g. request was registered but the EngineDeadError handler
+        #    missed it because it wasn't submitted to the dead stage yet).
+        for req_id in list(self.request_states):
+            if req_id not in notified:
+                await self.output_async_queue.put(
+                    {
+                        "type": "error",
+                        "error": self._fatal_error,
+                        "fatal": True,
+                        "request_id": req_id,
+                    }
+                )
+            self.request_states.pop(req_id, None)
+
     def _shutdown_stages(self) -> None:
         """Shutdown all stage clients."""
         if self._stages_shutdown:

vllm_omni/engine/stage_engine_core_client.py

@@ -6,13 +6,17 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
 
 from __future__ import annotations
 
+import multiprocessing.connection
 import socket
+import threading
+import weakref
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
     KV_TRANSFER_PORT_OFFSET,
@@ -158,13 +162,64 @@ class StageEngineCoreClientBase:
                     shutdown_error,
                 )
             raise
+
         self._initialize_kv_sender_endpoint()
+
+        if self._proc is not None:
+            self._start_proc_monitor()
+
         logger.info(
             "[%s] Stage-%s EngineCore running",
             client_name,
             self.stage_id,
         )
 
+    def _start_proc_monitor(self) -> None:
+        """Start a daemon thread that watches the subprocess sentinel.
+
+        When the subprocess dies without sending the ZMQ ``ENGINE_CORE_DEAD``
+        sentinel (e.g. SIGKILL, segfault, OOM-killer), this thread sets
+        ``resources.engine_dead`` so subsequent calls raise
+        ``EngineDeadError``.
+        """
+        proc = self._proc
+        resources_ref = weakref.ref(self.resources)
+        stage_id = self.stage_id
+
+        def _monitor() -> None:
+            try:
+                multiprocessing.connection.wait([proc.sentinel])
+            except Exception:
+                return
+            resources = resources_ref()
+            if resources is None or resources.engine_dead:
+                return
+            resources.engine_dead = True
+            logger.error(
+                "[StageEngineCoreClient] Stage-%s subprocess died unexpectedly (exit code %s).",
+                stage_id,
+                proc.exitcode,
+            )
+
+        t = threading.Thread(
+            target=_monitor,
+            daemon=True,
+            name=f"StageCoreProcMonitor-{stage_id}",
+        )
+        t.start()
+
+    def check_health(self) -> None:
+        """Raise ``EngineDeadError`` if the stage subprocess is dead.
+
+        Called by ``OmniBase.check_health()`` and transitively by the
+        ``/health`` HTTP endpoint.
+        """
+        if self.resources.engine_dead:
+            raise EngineDeadError(f"Stage-{self.stage_id} engine core is dead")
+        if self._proc is not None and not self._proc.is_alive():
+            self.resources.engine_dead = True
+            raise EngineDeadError(f"Stage-{self.stage_id} subprocess is not alive (exit code {self._proc.exitcode})")
+
     # ==================== Overrides ====================
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:

vllm_omni/entrypoints/async_omni.py

@@ -493,6 +493,9 @@ class AsyncOmni(EngineClient, OmniBase):
 
             stage_id = result.get("stage_id", 0)
 
+            if result.get("type") == "error" and result.get("fatal"):
+                raise EngineDeadError(result.get("error", ""))
+
             # Check for errors
             if "error" in result:
                 logger.error(
@@ -503,6 +506,8 @@ class AsyncOmni(EngineClient, OmniBase):
                 )
                 raise RuntimeError(result)
 
+            self._check_engine_output_error(result, request_id, stage_id)
+
             # Process the result (constructs OmniRequestOutput)
             output_to_yield = self._process_single_result(
                 result,
@@ -575,6 +580,16 @@ class AsyncOmni(EngineClient, OmniBase):
 
             except asyncio.CancelledError:
                 raise
+            except EngineDeadError as e:
+                logger.error("[AsyncOmni] Engine dead: %s", e)
+                for req_state in list(self.request_states.values()):
+                    error_msg = {
+                        "type": "error",
+                        "error": str(e),
+                        "fatal": True,
+                        "request_id": req_state.request_id,
+                    }
+                    await req_state.queue.put(error_msg)
             except Exception as e:
                 logger.exception("[AsyncOmni] final_output_loop failed.")
                 for req_state in list(self.request_states.values()):
@@ -841,12 +856,21 @@ class AsyncOmni(EngineClient, OmniBase):
     @property
     def is_running(self) -> bool:
         """Check if the engine is running."""
-        return self.final_output_task is not None and not self.final_output_task.done()
+        orchestrator_alive = self.engine.is_alive()
+        task_alive = self.final_output_task is not None and not self.final_output_task.done()
+        return orchestrator_alive and task_alive
 
     @property
     def errored(self) -> bool:
-        """Whether orchestrator thread has stopped unexpectedly."""
-        return not self.engine.is_alive()
+        """Whether the engine is in a non-recoverable error state.
+
+        Delegates to ``OmniBase.errored`` which checks the orchestrator
+        thread and all stage clients.  Redeclared here to satisfy the
+        ``EngineClient`` abstract-property requirement (Python's ABC
+        mechanism does not resolve abstract methods from sibling MRO
+        entries).
+        """
+        return OmniBase.errored.fget(self)  # type: ignore[union-attr]
 
     @property
     def _name(self) -> str:

vllm_omni/entrypoints/omni.py

@@ -166,6 +166,8 @@ class Omni(OmniBase):
                     logger.warning("[Omni] Received output for unknown/finished request_id=%s", req_id)
                     continue
 
+                self._check_engine_output_error(msg, req_id, stage_id)
+
                 if req_state.metrics is None:
                     continue
                 output_to_yield = self._process_single_result(

vllm_omni/entrypoints/omni_base.py

@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, Literal
 
 import huggingface_hub
 from vllm.logger import init_logger
-from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
@@ -198,9 +198,33 @@ class OmniBase(PDDisaggregationMixin):
     def is_running(self) -> bool:
         return self.engine.is_alive()
 
+    @property
+    def errored(self) -> bool:
+        """Whether the engine is in a non-recoverable error state.
+
+        True when the orchestrator thread is dead **or** any stage client
+        has been marked dead (e.g. diffusion worker OOM / process death).
+
+        Checks both ``_engine_dead`` (StageDiffusionClient) and
+        ``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
+
     def check_health(self) -> None:
         if not self.engine.is_alive():
             raise EngineDeadError("Orchestrator process is not alive")
+        for stage_client in self.engine.stage_clients:
+            if hasattr(stage_client, "check_health"):
+                stage_client.check_health()
 
     def resolve_sampling_params_list(
         self,
@@ -271,7 +295,10 @@ class OmniBase(PDDisaggregationMixin):
             return True, None, None, None
 
         if msg_type == "error":
-            raise RuntimeError(msg.get("error", "Orchestrator returned an error message"))
+            error_text = msg.get("error", "Orchestrator returned an error message")
+            if msg.get("fatal"):
+                raise EngineDeadError(error_text)
+            raise RuntimeError(error_text)
 
         if msg_type != "output":
             logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)
@@ -300,6 +327,34 @@ class OmniBase(PDDisaggregationMixin):
 
         return False, req_id, stage_id, req_state
 
+    def _check_engine_output_error(
+        self,
+        result: dict[str, Any],
+        request_id: str,
+        stage_id: int,
+    ) -> None:
+        """Raise if ``engine_outputs`` carries an error field.
+
+        Raises :class:`EngineDeadError` when ``self.errored`` indicates the
+        engine is unrecoverable, otherwise raises :class:`EngineGenerateError`
+        (recoverable, single-request failure).
+        """
+        engine_outputs = result.get("engine_outputs")
+        error_text = getattr(engine_outputs, "error", None)
+        if error_text is None:
+            return
+        logger.error(
+            "[%s] Stage error for req=%s stage-%s: %s",
+            self.__class__.__name__,
+            request_id,
+            stage_id,
+            error_text,
+        )
+        # NOTE: O(n_stages) check for every error.
+        if self.errored:
+            raise EngineDeadError(error_text)
+        raise EngineGenerateError(error_text)
+
     def _process_single_result(
         self,
         result: dict[str, Any],

vllm_omni/entrypoints/openai/api_server.py

@@ -31,7 +31,7 @@ from vllm import SamplingParams
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
 from vllm.entrypoints.chat_utils import load_chat_template
-from vllm.entrypoints.launcher import serve_http
+from vllm.entrypoints.launcher import serve_http, terminate_if_errored
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
@@ -50,6 +50,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     ModelCard,
     ModelList,
     ModelPermission,
+    RequestResponseMetadata,
 )
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
@@ -74,6 +75,7 @@ from vllm.entrypoints.serve.instrumentator.basic import base
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
 from vllm.entrypoints.utils import (
+    create_error_response,
     load_aware_call,
     process_lora_modules,
     with_cancellation,
@@ -83,6 +85,7 @@ from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
 from vllm.utils.system_utils import decorate_logs
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
@@ -199,6 +202,50 @@ def _remove_route_from_app(app, path: str, methods: set[str] | None = None):
         app.routes.remove(route)
 
 
+def _register_omni_exception_handlers(app) -> None:
+    """Override upstream vLLM exception handlers with Omni-aware versions.
+
+    The upstream ``engine_error_handler`` is designed for ``AsyncLLM`` (single
+    EngineCore process).  Omni uses a multi-stage orchestrator with different
+    health semantics, so we register our own handlers that:
+
+    - Log multi-stage diagnostic info (orchestrator liveness, per-stage health)
+      when an ``EngineDeadError`` is caught.
+    - Call ``terminate_if_errored``
+    - Return an OpenAI-compatible error JSON response.
+    """
+
+    async def omni_engine_error_handler(
+        req: Request,
+        exc: EngineDeadError | EngineGenerateError,
+    ):
+        request_id = req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None
+
+        if req.app.state.args.log_error_stack:
+            logger.exception("Engine Exception caught. Request id: %s", request_id)
+
+        engine = req.app.state.engine_client
+        if isinstance(exc, EngineDeadError):
+            # Log Omni-specific diagnostic information for dead engines.
+            orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A"
+            logger.error(
+                "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s",
+                orchestrator_alive,
+                engine.errored,
+                request_id,
+            )
+
+        terminate_if_errored(
+            server=req.app.state.server,
+            engine=engine,
+        )
+        err = create_error_response(exc)
+        return JSONResponse(err.model_dump(), status_code=err.error.code)
+
+    app.exception_handler(EngineGenerateError)(omni_engine_error_handler)
+    app.exception_handler(EngineDeadError)(omni_engine_error_handler)
+
+
 class _DiffusionServingModels:
     """Minimal OpenAIServingModels implementation for diffusion-only servers.
 
@@ -308,6 +355,10 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
         _remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
         app.include_router(router)
 
+        # OMNI: Override upstream exception handlers with Omni-aware versions
+        # that understand the multi-stage orchestrator lifecycle.
+        _register_omni_exception_handlers(app)
+
         await omni_init_app_state(engine_client, app.state, args)
 
         # Conditionally register profiler endpoints based on stage YAML configs
@@ -878,6 +929,8 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
         return base_server.create_error_response(message="The model does not support Chat Completions API")
     try:
         generator = await handler.create_chat_completion(request, raw_request)
+    except (EngineGenerateError, EngineDeadError):
+        raise  # Propagate to the global Omni exception handler
     except Exception as e:
         logger.exception("Chat completion failed: %s", e)
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
@@ -973,6 +1026,8 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
                 status_code=result.error.code if result.error else 400,
             )
         return result
+    except (EngineGenerateError, EngineDeadError):
+        raise  # Propagate to the global Omni exception handler
     except Exception as e:
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
 
@@ -1007,6 +1062,8 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request)
                 status_code=result.error.code if result.error else 400,
             )
         return JSONResponse(content=result.model_dump())
+    except (EngineGenerateError, EngineDeadError):
+        raise  # Propagate to the global Omni exception handler
     except ValueError as e:
         raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e
     except Exception as e:
@@ -1244,31 +1301,26 @@ _remove_route_from_router(router, "/health")
 async def health(raw_request: Request) -> JSONResponse:
     """Health check endpoint that works for both LLM and diffusion modes.
 
-    Returns 200 OK if the server is healthy.
-    For LLM mode: delegates to engine_client health check
-    For diffusion mode: checks if diffusion_engine is running
+    Returns 200 OK if the server is healthy, 503 if the engine is dead.
+    Mirrors vLLM upstream's /health which catches EngineDeadError -> 503.
     """
-    # Check if we're in diffusion mode
-    diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None)
-    if diffusion_engine is not None:
-        # Diffusion mode health check
-        if hasattr(diffusion_engine, "is_running") and diffusion_engine.is_running:
-            return JSONResponse(content={"status": "healthy"})
+    engine_client = getattr(raw_request.app.state, "engine_client", None) or getattr(
+        raw_request.app.state, "diffusion_engine", None
+    )
+    if engine_client is None:
         return JSONResponse(
-            content={"status": "unhealthy", "reason": "Diffusion engine is not running"},
+            content={"status": "unhealthy", "reason": "No engine initialized"},
             status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
         )
 
-    # LLM mode - delegate to engine_client
-    engine_client = getattr(raw_request.app.state, "engine_client", None)
-    if engine_client is not None:
+    try:
         await engine_client.check_health()
         return JSONResponse(content={"status": "healthy"})
-
-    return JSONResponse(
-        content={"status": "unhealthy", "reason": "No engine initialized"},
-        status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
-    )
+    except EngineDeadError:
+        return JSONResponse(
+            content={"status": "unhealthy"},
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+        )
 
 
 # Remove existing models endpoint if present (from vllm imports)
@@ -1437,6 +1489,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         _update_if_not_none(gen_params, "layers", request.layers)
 
         request_id = f"img_gen-{random_uuid()}"
+        raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
 
         logger.info(f"Generating {request.n} image(s) {size_str}")
 
@@ -1468,8 +1521,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             data=image_data,
         )
 
+    except (EngineGenerateError, EngineDeadError):
+        raise  # Propagate to the global Omni exception handler
     except HTTPException:
-        # Re-raise HTTPExceptions as-is
         raise
     except ValueError as e:
         logger.error(f"Validation error: {e}")
@@ -1672,6 +1726,7 @@ async def edit_images(
 
         # 4. Generate images using AsyncOmni (multi-stage mode)
         request_id = f"img_edit-{random_uuid()}"
+        raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
         logger.info(f"Generating {n} image(s) {size_str}")
         result = await _generate_with_async_omni(
             engine_client=engine_client,
@@ -1703,8 +1758,9 @@ async def edit_images(
             size=size_str,
         )
 
+    except (EngineGenerateError, EngineDeadError):
+        raise  # Propagate to the global Omni exception handler
     except HTTPException:
-        # Re-raise HTTPExceptions as-is
         raise
     except ValueError as e:
         logger.error(f"Validation error: {e}")
@@ -2110,6 +2166,48 @@ def video_response_from_request(model_name: str, req: VideoGenerationRequest) ->
     return resp
 
 
+def _status_code_for_video_failure(error: VideoError | None) -> int:
+    if error is None:
+        return HTTPStatus.INTERNAL_SERVER_ERROR.value
+
+    if isinstance(error.code, int):
+        if 400 <= error.code < 600:
+            return error.code
+        return HTTPStatus.INTERNAL_SERVER_ERROR.value
+
+    if error.code == "HTTPException":
+        status_text, _, _ = error.message.partition(":")
+        try:
+            status_code = int(status_text)
+        except ValueError:
+            return HTTPStatus.INTERNAL_SERVER_ERROR.value
+        if 400 <= status_code < 600:
+            return status_code
+        return HTTPStatus.INTERNAL_SERVER_ERROR.value
+
+    if error.code == "EngineDeadError":
+        return HTTPStatus.INTERNAL_SERVER_ERROR.value
+    if error.code == "EngineGenerateError":
+        return HTTPStatus.INTERNAL_SERVER_ERROR.value
+
+    return HTTPStatus.INTERNAL_SERVER_ERROR.value
+
+
+def _video_error_from_exception(exc: Exception) -> VideoError:
+    if isinstance(exc, HTTPException):
+        message = str(exc.detail) if exc.detail else str(exc)
+        return VideoError(code=exc.status_code, message=message)
+
+    if isinstance(exc, (EngineGenerateError, EngineDeadError)):
+        err = create_error_response(exc)
+        return VideoError(code=err.error.code, message=err.error.message)
+
+    return VideoError(
+        code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+        message=str(exc),
+    )
+
+
 def _cleanup_video(video_id: str, output_path: str | None):
     try:
         if output_path is not None:
@@ -2123,6 +2221,7 @@ async def _run_video_generation_job(
     request: VideoGenerationRequest,
     video_id: str,
     reference_image: ReferenceImage | None = None,
+    app_state: Any | None = None,
 ) -> None:
     job = await VIDEO_STORE.get(video_id)
     if job is None:
@@ -2153,17 +2252,36 @@ async def _run_video_generation_job(
                 "peak_memory_mb": peak_memory_mb,
             },
         )
-    except Exception as exc:
-        logger.exception("Video generation failed for id=%s", video_id)
+    except (EngineGenerateError, EngineDeadError) as exc:
+        logger.exception("Video generation failed (engine error) for id=%s", video_id)
 
         _cleanup_video(video_id, output_path)
-        # TODO: It would be better to have a finite collection of errors to return rather than the exception name
         await VIDEO_STORE.update_fields(
             video_id,
             {
                 "status": VideoGenerationStatus.FAILED,
                 "completed_at": int(time.time()),
-                "error": VideoError(code=type(exc).__name__, message=str(exc)),
+                "error": _video_error_from_exception(exc),
+                "inference_time_s": time.perf_counter() - started_at,
+            },
+        )
+        # Background tasks can't propagate exceptions to FastAPI handlers.
+        # Actively signal shutdown when the engine is dead.
+        if app_state is not None and isinstance(exc, EngineDeadError):
+            terminate_if_errored(
+                server=app_state.server,
+                engine=app_state.engine_client,
+            )
+    except Exception as exc:
+        logger.exception("Video generation failed for id=%s", video_id)
+
+        _cleanup_video(video_id, output_path)
+        await VIDEO_STORE.update_fields(
+            video_id,
+            {
+                "status": VideoGenerationStatus.FAILED,
+                "completed_at": int(time.time()),
+                "error": _video_error_from_exception(exc),
                 "inference_time_s": time.perf_counter() - started_at,
             },
         )
@@ -2294,6 +2412,7 @@ async def _parse_video_form(
     },
 )
 async def create_video(
+    raw_request: Request,
     ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
 ) -> VideoResponse:
     """Create an asynchronous video generation job.
@@ -2304,7 +2423,9 @@ async def create_video(
     request, handler, effective_model_name, reference_image = ctx
     ref = video_response_from_request(effective_model_name, request)
     await VIDEO_STORE.upsert(ref.id, ref)
-    task = asyncio.create_task(_run_video_generation_job(handler, request, ref.id, reference_image))
+    task = asyncio.create_task(
+        _run_video_generation_job(handler, request, ref.id, reference_image, app_state=raw_request.app.state)
+    )
     await VIDEO_TASKS.upsert(ref.id, task)
     return ref
 
@@ -2319,6 +2440,7 @@ async def create_video(
     },
 )
 async def create_video_sync(
+    raw_request: Request,
     ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
 ) -> Response:
     """Synchronous video generation endpoint.
@@ -2332,6 +2454,7 @@ async def create_video_sync(
     """
     request, handler, effective_model_name, reference_image = ctx
     request_id = f"video_sync-{random_uuid()}"
+    raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
     started_at = time.perf_counter()
     try:
         video_bytes, stage_durations, peak_memory_mb = await asyncio.wait_for(
@@ -2343,6 +2466,8 @@ async def create_video_sync(
             status_code=HTTPStatus.GATEWAY_TIMEOUT.value,
             detail=f"Video generation timed out after {VIDEO_SYNC_TIMEOUT_S}s.",
         )
+    except (EngineGenerateError, EngineDeadError):
+        raise  # Propagate to the global Omni exception handler
     except HTTPException:
         raise
     except Exception as exc:
@@ -2403,8 +2528,8 @@ async def list_videos(
     return VideoListResponse(data=jobs, has_more=has_more, first_id=first_id, last_id=last_id)
 
 
-@router.get("/v1/videos/{video_id}")
-async def retrieve_video(video_id: str) -> VideoResponse:
+@router.get("/v1/videos/{video_id}", response_model=None)
+async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse:
     """Retrieve metadata for a previously created video job.
 
     Args:
@@ -2419,6 +2544,15 @@ async def retrieve_video(video_id: str) -> VideoResponse:
     job = await VIDEO_STORE.get(video_id)
     if job is None:
         raise HTTPException(status_code=404, detail="Video not found")
+    if job.status == VideoGenerationStatus.FAILED:
+        status_code = _status_code_for_video_failure(job.error)
+        content = job.model_dump(mode="json")
+        if content.get("error") is not None:
+            content["error"]["code"] = status_code
+        return JSONResponse(
+            content=content,
+            status_code=status_code,
+        )
     return job
 
 

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -235,7 +235,7 @@ class VideoGenerationResponse(BaseModel):
 
 
 class VideoError(BaseModel):
-    code: str = Field(..., description="A machine-readable error code that was returned.")
+    code: int | str = Field(..., description="A machine-readable error code that was returned.")
     message: str = Field(..., description="A human-readable description of the error that was returned.")
 
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -32,6 +32,7 @@ from vllm.entrypoints.chat_utils import (
     get_history_tool_calls_cnt,
     make_tool_call_id,
 )
+from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
@@ -80,6 +81,7 @@ from vllm.tokenizers.mistral import (
 from vllm.tool_parsers import ToolParser
 from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.utils.collection_utils import as_list
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
@@ -435,6 +437,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tokenizer,
                 request_metadata,
                 reasoning_parser,
+                raw_request=raw_request,
             )
 
         try:
@@ -825,6 +828,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         tokenizer: AnyTokenizer,
         request_metadata: RequestResponseMetadata,
         reasoning_parser: ReasoningParser | None = None,
+        raw_request: Request | None = None,
     ):
         created_time = int(time.time())
         chunk_object_type: Final = "chat.completion.chunk"
@@ -1537,6 +1541,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         delta=False,
                     )
 
+        except EngineDeadError as e:
+            logger.error(
+                "EngineDeadError during streaming for request %s: %s",
+                request_id,
+                e,
+            )
+            data = self.create_streaming_error_response(e)
+            yield f"data: {data}\n\n"
+            # Actively signal shutdown instead of waiting for the watchdog
+            # (5s polling interval).
+            if raw_request is not None:
+                terminate_if_errored(
+                    server=raw_request.app.state.server,
+                    engine=self.engine_client,
+                )
         except Exception as e:
             logger.exception("Error in chat completion stream generator.")
             data = self.create_streaming_error_response(e)

vllm_omni/entrypoints/openai/serving_speech.py

@@ -17,12 +17,17 @@ import torch
 from fastapi import Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
 from transformers.utils.hub import cached_file
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.launcher import terminate_if_errored
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorResponse,
+    RequestResponseMetadata,
+)
 from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
 from vllm.utils import random_uuid
 from vllm.utils.async_utils import make_async
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.protocol.audio import (
@@ -1151,7 +1156,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             )
         return wav_np.tolist(), sr
 
-    async def _generate_audio_chunks(self, generator, request_id: str, response_format: str = "pcm"):
+    async def _generate_audio_chunks(
+        self,
+        generator,
+        request_id: str,
+        response_format: str = "pcm",
+        raw_request: Request | None = None,
+    ):
         """Generate audio chunks for streaming response.
 
         Handles two audio output modes from the engine:
@@ -1226,6 +1237,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         except asyncio.CancelledError:
             logger.info("Streaming request %s cancelled by client", request_id)
             raise
+        except EngineDeadError as e:
+            logger.error(
+                "EngineDeadError during streaming speech for %s: %s",
+                request_id,
+                e,
+            )
+            # Actively signal shutdown rather than relying on the watchdog.
+            if raw_request is not None:
+                terminate_if_errored(
+                    server=raw_request.app.state.server,
+                    engine=self.engine_client,
+                )
+            raise
         except Exception as e:
             logger.exception("Streaming speech generation failed for %s: %s", request_id, e)
             raise
@@ -1500,6 +1524,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     async def _prepare_speech_generation(
         self,
         request: OpenAICreateSpeechRequest,
+        request_id: str | None = None,
     ) -> tuple[str, Any, dict[str, Any]]:
         if self.engine_client.errored:
             raise self.engine_client.dead_error
@@ -1592,7 +1617,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             tts_params = {}
             prompt = {"prompt": request.input}
 
-        request_id = f"speech-{random_uuid()}"
+        request_id = request_id or f"speech-{random_uuid()}"
         if self._is_fish_speech:
             model_type = "fish_speech"
         elif self._tts_model_type == "voxtral_tts":
@@ -1674,8 +1699,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         self,
         request: OpenAICreateSpeechRequest,
         base64_encode: bool = False,
+        request_id: str | None = None,
     ) -> tuple[bytes | str, str]:
-        request_id, generator, _ = await self._prepare_speech_generation(request)
+        request_id, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
 
         final_output: OmniRequestOutput | None = None
         async for res in generator:
@@ -1799,6 +1825,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         except asyncio.CancelledError:
             return self._diffusion_error_response("Client disconnected")
+        except (EngineGenerateError, EngineDeadError):
+            raise  # Propagate to the global Omni exception handler
         except ValueError as e:
             return self._diffusion_error_response(str(e))
         except Exception as e:
@@ -1844,6 +1872,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             logger.error("Error with model %s", error_check_ret)
             return error_check_ret
 
+        request_id = f"speech-{random_uuid()}"
+        if raw_request:
+            raw_request.state.request_metadata = RequestResponseMetadata(
+                request_id=request_id,
+            )
+
         try:
             if request.stream:
                 # Determine response format and media type for streaming
@@ -1864,17 +1898,24 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     )
 
                 media_type = "audio/wav" if response_format == "wav" else "audio/pcm"
-                request_id, generator, _ = await self._prepare_speech_generation(request)
+                _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
                 return StreamingResponse(
-                    self._generate_audio_chunks(generator, request_id, response_format),
+                    self._generate_audio_chunks(
+                        generator,
+                        request_id,
+                        response_format,
+                        raw_request=raw_request,
+                    ),
                     media_type=media_type,
                 )
 
-            audio_bytes, media_type = await self._generate_audio_bytes(request)
+            audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id)
             return Response(content=audio_bytes, media_type=media_type)
 
         except asyncio.CancelledError:
             return self.create_error_response("Client disconnected")
+        except (EngineGenerateError, EngineDeadError):
+            raise  # Propagate to the global Omni exception handler
         except ValueError as e:
             return self.create_error_response(e)
         except Exception as e:

vllm_omni/outputs.py

@@ -100,9 +100,22 @@ class OmniRequestOutput:
     # memory usage info
     peak_memory_mb: float = 0.0
 
-    # error handling
+    # Error information -- set when the output represents a failed request.
     error: str | None = None
 
+    @classmethod
+    def from_error(
+        cls,
+        request_id: str,
+        error: str,
+    ) -> "OmniRequestOutput":
+        """Create an error output for a request that failed during generation."""
+        return cls(
+            request_id=request_id,
+            finished=True,
+            error=error,
+        )
+
     @classmethod
     def from_pipeline(
         cls,