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

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

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)]