diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 4bc3e05fe..9ec06e810 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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 diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py index ef3cfab3b..47b3d5e9f 100644 --- a/tests/engine/test_async_omni_engine_outputs.py +++ b/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 diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py new file mode 100644 index 000000000..6131c5e9b --- /dev/null +++ b/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) diff --git a/tests/engine/test_stage_engine_core_client.py b/tests/engine/test_stage_engine_core_client.py new file mode 100644 index 000000000..dde0927af --- /dev/null +++ b/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 diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 607b3eaa8..8a373f74d 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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") diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 6157d8231..1cc561665 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/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() diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 3cffcd37d..289cf2673 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/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() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index fe940d623..abaf59895 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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( diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 6325d4a16..dcb35cfde 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/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 diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 7876b0094..d1216e2f2 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/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 diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 480d113d1..ff2d9f1e8 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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: diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index eced444fd..d36c5c644 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/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() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 7b98c13b8..a26102d05 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e86388a38..ca60c7531 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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: diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 66bae22a5..f6683c1e0 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 6759fa5a6..67425ea2e 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 8ef7e2ee5..223c208af 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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( diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 98b5c7cf7..3180c9c80 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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], diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 5d8cc4ea2..1def10e64 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index 7c2c3164d..d46c8d43d 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/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.") diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index aee6c3674..06b739e3b 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index ba8292f0c..c275c7795 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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: diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index c02c0c142..23930f358 100644 --- a/vllm_omni/outputs.py +++ b/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,