diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 171a6278c..8bae6b2d9 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/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" diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index f8d68d95a..266900f70 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/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}, + } diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py index 0c1ddc6a5..31f238382 100644 --- a/tests/entrypoints/test_omni_base_profiler.py +++ b/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 diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 0aeb6158e..f0ebe2a3f 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/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)] diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 488378b40..c04d46cf0 100644 --- a/vllm_omni/diffusion/data.py +++ b/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() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index ff0f753b4..c40101f0d 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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. diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 5a6fb6371..f11d300f7 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e6373ec96..b2ce6ec70 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 68c072c2b..280f2333c 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index a7a02eded..f52580737 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 96df0591e..1a56a0244 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 5ccc7c3d2..390dd2919 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/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(