[release/v0.18.0.post1][Bugfix] fix oom no response to user (#2638)

Signed-off-by: bjf-frz <frz123db@gmail.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -6,7 +6,7 @@ from unittest.mock import Mock, patch
 
 import pytest
 
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, OmniRequestError, normalize_omni_error
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
@@ -245,6 +245,39 @@ class TestDiffusionEngine:
         assert output is expected
         engine.executor.add_req.assert_called_once_with(request)
 
+    def test_add_req_and_wait_for_response_wraps_executor_errors(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+        engine.executor = Mock()
+        engine.executor.add_req.side_effect = RuntimeError("worker oom")
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine-error")
+
+        with pytest.raises(OmniRequestError, match="worker oom") as exc_info:
+            engine.add_req_and_wait_for_response(request)
+
+        assert exc_info.value.error_type == "RuntimeError"
+
+    def test_add_req_and_wait_for_response_raises_scheduler_error_when_empty_without_requests(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = Mock()
+        engine.scheduler.add_request.return_value = "sched-1"
+        engine.scheduler.schedule.return_value = Mock(is_empty=True)
+        engine.scheduler.has_requests.return_value = False
+        engine.scheduler.pop_request_state.return_value = None
+        engine.executor = Mock()
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine-empty")
+
+        with pytest.raises(OmniRequestError, match="no runnable requests") as exc_info:
+            engine.add_req_and_wait_for_response(request)
+
+        assert exc_info.value.error_type == "SchedulerError"
+        engine.scheduler.pop_request_state.assert_called_once_with("sched-1")
+
     def test_supports_scheduler_interface_injection(self) -> None:
         request = _make_request("engine_iface")
         expected = DiffusionOutput(output=None)
@@ -297,3 +330,15 @@ class TestDiffusionEngine:
 
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
+
+
+def test_normalize_omni_error_preserves_existing_metadata() -> None:
+    err = OmniRequestError("boom", status_code=409, error_type="Conflict")
+
+    normalized = normalize_omni_error(err, request_id="req-1", stage_id=3)
+
+    assert normalized is err
+    assert normalized.request_id == "req-1"
+    assert normalized.stage_id == 3
+    assert normalized.status_code == 409
+    assert normalized.error_type == "Conflict"

tests/entrypoints/openai_api/test_video_server.py

@@ -15,10 +15,12 @@ from types import SimpleNamespace
 
 import pytest
 from fastapi import FastAPI
+from fastapi import HTTPException
 from fastapi.testclient import TestClient
 from PIL import Image
 from pytest_mock import MockerFixture
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.openai import api_server
 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -29,6 +31,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
 from vllm_omni.entrypoints.openai.storage import LocalStorageManager
 from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -56,6 +59,18 @@ class FakeAsyncOmni:
         yield MockVideoResult(videos)
 
 
+class ErrorAsyncOmni(FakeAsyncOmni):
+    def __init__(self, exc: Exception):
+        super().__init__()
+        self._exc = exc
+
+    async def generate(self, prompt, request_id, sampling_params_list):
+        self.captured_prompt = prompt
+        self.captured_sampling_params_list = sampling_params_list
+        raise self._exc
+        yield  # pragma: no cover
+
+
 class BlockingVideoHandler:
     def __init__(self):
         self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -735,3 +750,37 @@ def test_extra_params_merged_with_existing_extra_args(test_client, mocker: Mocke
     assert captured.extra_args["flow_shift"] == 0.5
     assert captured.extra_args["use_zero_init"] is True
     assert captured.extra_args["zero_steps"] == 2
+
+
+@pytest.mark.asyncio
+async def test_run_generation_maps_omni_request_error_to_http_exception():
+    serving = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=ErrorAsyncOmni(
+            OmniRequestError(
+                "oom",
+                status_code=507,
+                request_id="req-oom",
+                stage_id=2,
+                error_type="OutOfMemoryError",
+                detail={"retryable": False},
+            )
+        ),
+        model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await serving._run_generation(
+            {"prompt": "test"},
+            OmniDiffusionSamplingParams(),
+            "req-oom",
+        )
+
+    exc = exc_info.value
+    assert exc.status_code == 507
+    assert exc.detail == {
+        "message": "oom",
+        "request_id": "req-oom",
+        "stage_id": 2,
+        "error_type": "OutOfMemoryError",
+        "detail": {"retryable": False},
+    }

tests/entrypoints/test_omni_base_profiler.py

@@ -4,6 +4,8 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
+from vllm_omni.diffusion.data import OmniRequestError
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
@@ -204,6 +206,45 @@ class TestOmniBaseProfilerSignatureConsistency:
         # stages should default to None
         assert sig.parameters["stages"].default is None
 
+
+class TestOmniBaseOutputHandling:
+    def test_handle_output_message_raises_structured_omni_error(self):
+        mock_engine = MagicMock()
+        mock_engine.num_stages = 1
+        mock_engine.default_sampling_params_list = [MagicMock()]
+        mock_engine.get_stage_metadata.return_value = {
+            "final_output_type": "image",
+            "final_output": True,
+        }
+
+        with (
+            patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine),
+            patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x),
+            patch("vllm_omni.entrypoints.omni_base.weakref.finalize"),
+        ):
+            from vllm_omni.entrypoints.omni_base import OmniBase
+
+            omni_base = OmniBase(model="test-model")
+
+        with pytest.raises(OmniRequestError, match="oom") as exc_info:
+            omni_base._handle_output_message(
+                {
+                    "type": "error",
+                    "error": "oom",
+                    "status_code": 507,
+                    "request_id": "req-1",
+                    "stage_id": 2,
+                    "error_type": "OutOfMemoryError",
+                    "detail": {"retryable": False},
+                }
+            )
+
+        assert exc_info.value.status_code == 507
+        assert exc_info.value.request_id == "req-1"
+        assert exc_info.value.stage_id == 2
+        assert exc_info.value.error_type == "OutOfMemoryError"
+        assert exc_info.value.detail == {"retryable": False}
+
     def test_stop_profile_default_values(self):
         """Verify stop_profile has correct default parameter values."""
         import inspect

tests/entrypoints/test_omni_entrypoints.py

@@ -10,6 +10,7 @@ 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_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.omni import Omni
 
@@ -540,12 +541,16 @@ async def test_async_omni_propagates_engine_error(monkeypatch: pytest.MonkeyPatc
 
     app = AsyncOmni("dummy-model")
     try:
-        with pytest.raises(RuntimeError, match="engine boom"):
+        with pytest.raises(OmniRequestError, match="engine boom") as exc_info:
             async for _ in app.generate(prompt="hello", request_id="req-1"):
                 pass
     finally:
         app.shutdown()
 
+    assert exc_info.value.request_id == "req-1"
+    assert exc_info.value.stage_id == 0
+    assert exc_info.value.status_code == 500
+
 
 def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
     sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]

vllm_omni/diffusion/data.py

@@ -675,6 +675,49 @@ class DiffusionOutput:
     peak_memory_mb: float = 0.0
 
 
+@dataclass
+class OmniRequestError(RuntimeError):
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: int = 500,
+        request_id: str | None = None,
+        stage_id: int | None = None,
+        error_type: str | None = None,
+        detail: dict[str, Any] | None = None,
+    ):
+        super().__init__(message)
+        self.status_code = status_code
+        self.request_id = request_id
+        self.stage_id = stage_id
+        self.error_type = error_type or self.__class__.__name__
+        self.detail = detail or {}
+
+
+def normalize_omni_error(
+    exc: Exception,
+    *,
+    status_code: int = 500,
+    request_id: str | None = None,
+    stage_id: int | None = None,
+) -> OmniRequestError:
+    if isinstance(exc, OmniRequestError):
+        if exc.request_id is None:
+            exc.request_id = request_id
+        if exc.stage_id is None:
+            exc.stage_id = stage_id
+        return exc
+
+    return OmniRequestError(
+        str(exc),
+        status_code=status_code,
+        request_id=request_id,
+        stage_id=stage_id,
+        error_type=type(exc).__name__,
+    )
+
+
 class AttentionBackendEnum(enum.Enum):
     FA = enum.auto()
     SLIDING_TILE_ATTN = enum.auto()

vllm_omni/diffusion/diffusion_engine.py

@@ -11,7 +11,7 @@ import PIL.Image
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, OmniRequestError, normalize_omni_error
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
@@ -99,7 +99,12 @@ class DiffusionEngine:
         exec_total_time = time.perf_counter() - exec_start_time
 
         if output.error:
-            raise Exception(f"{output.error}")
+            # raise Exception(f"{output.error}")
+            raise OmniRequestError(
+                output.error,
+                status_code=500,
+                error_type="DiffusionExecutionError",
+            )
         logger.info("Generation completed successfully.")
 
         if output.output is None:
@@ -280,33 +285,52 @@ class DiffusionEngine:
             target_sched_req_id = self.scheduler.add_request(request)
 
             # keep scheduling and executing until the target request is finished
-            while True:
-                sched_output = self.scheduler.schedule()
-                if sched_output.is_empty:
-                    if not self.scheduler.has_requests():
-                        raise RuntimeError("Diffusion scheduler has no runnable requests.")
-                    continue
+            try:
+                while True:
+                    sched_output = self.scheduler.schedule()
+                    if sched_output.is_empty:
+                        if not self.scheduler.has_requests():
+                            # raise RuntimeError("Diffusion scheduler has no runnable requests.")
+                            raise OmniRequestError(
+                                "Diffusion scheduler has no runnable requests.",
+                                status_code=500,
+                                error_type="SchedulerError",
+                            )
+                        continue
 
-                # NOTE: add_req_and_wait_for_response() is synchronous, and
-                # the scheduler currently enforces _max_batch_size = 1 (see
-                # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
-                # take the single scheduled request here.
-                sched_req_id = sched_output.scheduled_req_ids[0]
-                req = sched_output.scheduled_new_reqs[0].req
+                    # NOTE: add_req_and_wait_for_response() is synchronous, and
+                    # the scheduler currently enforces _max_batch_size = 1 (see
+                    # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
+                    # take the single scheduled request here.
+                    sched_req_id = sched_output.scheduled_req_ids[0]
+                    req = sched_output.scheduled_new_reqs[0].req
+                    try:
+                        output = self.executor.add_req(req)
+                    except Exception as exc:
+                        logger.error(
+                            "Execution failed for diffusion request %s",
+                            sched_req_id,
+                            exc_info=True,
+                        )
+                        # output = DiffusionOutput(error=str(exc))
+                        raise normalize_omni_error(exc) from exc
+
+                    finished_req_ids = self.scheduler.update_from_output(sched_output, output)
+
+                    if output.error:
+                        raise OmniRequestError(
+                            output.error,
+                            status_code=500,
+                            error_type="DiffusionExecutionError",
+                        )
+                    if target_sched_req_id in finished_req_ids:
+                        # self.scheduler.pop_request_state(target_sched_req_id)
+                        return output
+            finally:
                 try:
-                    output = self.executor.add_req(req)
-                except Exception as exc:
-                    logger.error(
-                        "Execution failed for diffusion request %s",
-                        sched_req_id,
-                        exc_info=True,
-                    )
-                    output = DiffusionOutput(error=str(exc))
-
-                finished_req_ids = self.scheduler.update_from_output(sched_output, output)
-                if target_sched_req_id in finished_req_ids:
                     self.scheduler.pop_request_state(target_sched_req_id)
-                    return output
+                except Exception:
+                    logger.debug("Request state already removed: %s", target_sched_req_id, exc_info=True)
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         """Start or stop torch profiling on all diffusion workers.

vllm_omni/diffusion/stage_diffusion_client.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import normalize_omni_error
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.outputs import OmniRequestOutput
@@ -75,11 +76,24 @@ class StageDiffusionClient:
             result = await self._engine.generate(prompt, sampling_params, request_id)
             await self._output_queue.put(result)
         except Exception as e:
+            err = normalize_omni_error(e, request_id=request_id, stage_id=self.stage_id)
             logger.exception(
                 "[StageDiffusionClient] Stage-%s req=%s failed: %s",
                 self.stage_id,
                 request_id,
-                e,
+                err,
+            )
+            await self._output_queue.put(
+                {
+                    "type": "error",
+                    "request_id": request_id,
+                    "stage_id": self.stage_id,
+                    "status_code": err.status_code,
+                    "error": str(err),
+                    "error_type": err.error_type,
+                    "detail": err.detail,
+                    "finished": True,
+                }
             )
         finally:
             self._tasks.pop(request_id, None)

vllm_omni/engine/orchestrator.py

@@ -242,6 +242,29 @@ class Orchestrator:
                     output = stage_client.get_diffusion_output_async()
                     if output is not None:
                         idle = False
+
+                        if isinstance(output, dict) and output.get("type") == "error":
+                            req_id = output.get("request_id")
+                            req_state = self.request_states.get(req_id)
+
+                            await self.output_async_queue.put(
+                                {
+                                    "type": "error",
+                                    "request_id": req_id,
+                                    "stage_id": stage_id,
+                                    "status_code": output.get("status_code", 500),
+                                    "error": output.get("error", "Unknown diffusion stage error"),
+                                    "error_type": output.get("error_type"),
+                                    "detail": output.get("detail") or {},
+                                    "finished": True,
+                                    "stage_submit_ts": req_state.stage_submit_ts.get(stage_id) if req_state else None,
+                                }
+                            )
+
+                            self._cleanup_companion_state(req_id)
+                            self.request_states.pop(req_id, None)
+                            continue
+
                         req_state = self.request_states.get(output.request_id)
                         if req_state is not None:
                             stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state)

vllm_omni/entrypoints/async_omni.py

@@ -21,6 +21,7 @@ from vllm.pooling_params import PoolingParams
 from vllm.tasks import SupportedTask
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import OmniBase
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
@@ -290,14 +291,15 @@ class AsyncOmni(EngineClient, OmniBase):
             stage_id = result.get("stage_id", 0)
 
             # Check for errors
-            if "error" in result:
-                logger.error(
-                    "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s",
-                    request_id,
-                    stage_id,
-                    result["error"],
+            if isinstance(result, dict) and result.get("type") == "error":
+                raise OmniRequestError(
+                    result.get("error", "Unknown orchestrator error"),
+                    status_code=result.get("status_code", 500),
+                    request_id=result.get("request_id", request_id),
+                    stage_id=result.get("stage_id", stage_id),
+                    error_type=result.get("error_type"),
+                    detail=result.get("detail") or {},
                 )
-                raise RuntimeError(result)
 
             # Process the result (constructs OmniRequestOutput)
             output_to_yield = self._process_single_result(
@@ -344,6 +346,13 @@ class AsyncOmni(EngineClient, OmniBase):
                         await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S)
                         continue
 
+                    if isinstance(msg, dict) and msg.get("type") == "error":
+                        req_id = msg.get("request_id")
+                        req_state = self.request_states.get(req_id)
+                        if req_state is not None:
+                            await req_state.queue.put(msg)
+                        continue
+
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
                         continue
@@ -358,7 +367,16 @@ class AsyncOmni(EngineClient, OmniBase):
             except Exception as e:
                 logger.exception("[AsyncOmni] final_output_loop failed.")
                 for req_state in list(self.request_states.values()):
-                    error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    # error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    error_msg = {
+                        "type": "error",
+                        "request_id": req_state.request_id,
+                        "status_code": 500,
+                        "error": str(e),
+                        "error_type": type(e).__name__,
+                        "detail": {},
+                        "finished": True,
+                    }
                     await req_state.queue.put(error_msg)
                 self.final_output_task = None
 

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -18,7 +18,7 @@ from typing import Any
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig, normalize_omni_error
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -237,7 +237,8 @@ class AsyncOmniDiffusion:
             )
         except Exception as e:
             logger.error("Batch generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
+            # raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
+            raise normalize_omni_error(e, request_id=request_id) from e
 
         # Combine all per-prompt results into a single OmniRequestOutput
         all_images = []
@@ -310,7 +311,8 @@ class AsyncOmniDiffusion:
             result = result[0]
         except Exception as e:
             logger.error("Generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion generation failed: {e}") from e
+            # raise RuntimeError(f"Diffusion generation failed: {e}") from e
+            raise normalize_omni_error(e, request_id=request_id) from e
 
         if not result.request_id:
             result.request_id = request_id

vllm_omni/entrypoints/omni_base.py

@@ -12,6 +12,7 @@ import huggingface_hub
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
@@ -206,7 +207,14 @@ class OmniBase:
             return True, None, None, None
 
         if msg_type == "error":
-            raise RuntimeError(msg.get("error", "Orchestrator returned an error message"))
+            raise OmniRequestError(
+                msg.get("error", "Orchestrator returned an error message"),
+                status_code=msg.get("status_code", 500),
+                request_id=msg.get("request_id"),
+                stage_id=msg.get("stage_id"),
+                error_type=msg.get("error_type"),
+                detail=msg.get("detail") or {},
+            )
 
         if msg_type != "output":
             logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)

vllm_omni/entrypoints/openai/serving_video.py

@@ -13,6 +13,7 @@ from PIL import Image
 from vllm.engine.protocol import EngineClient
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoData,
@@ -214,12 +215,24 @@ class OmniOpenAIServingVideo:
         sampling_params_list: list[OmniSamplingParams] = [gen_params for _ in stage_configs]
 
         result = None
-        async for output in engine_client.generate(
-            prompt=prompt,
-            request_id=request_id,
-            sampling_params_list=sampling_params_list,
-        ):
-            result = output
+        try:
+            async for output in engine_client.generate(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params_list=sampling_params_list,
+            ):
+                result = output
+        except OmniRequestError as e:
+            raise HTTPException(
+                status_code=e.status_code,
+                detail={
+                    "message": str(e),
+                    "request_id": e.request_id,
+                    "stage_id": e.stage_id,
+                    "error_type": e.error_type,
+                    "detail": e.detail,
+                },
+            ) from e
 
         if result is None:
             raise HTTPException(