diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index c5cba1ea9..0553c03ff 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -161,3 +161,113 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) - finally: orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_diffusion_client_error_output_propagates_status_code(orchestrator_factory) -> None: + """A client-error OmniRequestOutput (e.g. a guardrail 4xx) from a diffusion + stage must surface as an ErrorMessage that carries status_code/error_type, + so the frontend can map it to the correct HTTP response instead of a 500. + """ + 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-blocked", + prompt={"prompt": "blocked"}, + original_prompt={"prompt": "blocked"}, + sampling_params_list=[params], + final_stage_id=0, + ) + + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + # Push a client-error output from the diffusion stage. + stage0.push_diffusion_output( + OmniRequestOutput.from_error( + "req-blocked", + "Input was blocked by Cosmos3 guardrails.", + status_code=400, + error_type="BadRequestError", + ) + ) + + msg = await _get_any_output_message(orchestrator_fixture) + + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "req-blocked" + assert msg.stage_id == 0 + assert msg.fatal is False + assert msg.status_code == 400 + assert msg.error_type == "BadRequestError" + assert msg.error == "Input was blocked by Cosmos3 guardrails." + + # Request state should be cleaned up. + await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states) + finally: + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status_code", "error_type"), + [ + pytest.param(429, "RateLimitError", id="429-too-many-requests"), + pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"), + pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"), + pytest.param(403, "PermissionDeniedError", id="403-forbidden"), + ], +) +async def test_diffusion_client_error_output_propagates_non_400_status( + orchestrator_factory, + status_code: int, + error_type: str, +) -> None: + """A diffusion-stage client error with a non-400 4xx status must be serialized + into an ErrorMessage that preserves that exact status_code/error_type. + """ + 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-blocked", + prompt={"prompt": "blocked"}, + original_prompt={"prompt": "blocked"}, + sampling_params_list=[params], + final_stage_id=0, + ) + + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + # Push a non-400 client-error output from the diffusion stage. + stage0.push_diffusion_output( + OmniRequestOutput.from_error( + "req-blocked", + "client error from diffusion stage", + status_code=status_code, + error_type=error_type, + ) + ) + + msg = await _get_any_output_message(orchestrator_fixture) + + assert isinstance(msg, ErrorMessage) + assert msg.request_id == "req-blocked" + assert msg.stage_id == 0 + assert msg.fatal is False + assert msg.status_code == status_code + assert msg.error_type == error_type + assert msg.error == "client error from diffusion stage" + + # Request state should be cleaned up. + await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states) + finally: + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 38cb86c6e..d108b9323 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import ( parse_size, ) from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.errors import GuardrailViolationError from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -587,6 +588,24 @@ def test_generate_single_image(test_client): assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"] +def test_generate_images_guardrail_error_returns_400(test_client, mock_async_diffusion): + async def blocked_generate(**kwargs): + raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.") + yield MockGenerationResult([]) # pragma: no cover + + mock_async_diffusion.generate = blocked_generate + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "blocked prompt", + "n": 1, + "size": "1024x1024", + }, + ) + assert response.status_code == 400 + assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails." + + def test_generate_images_async_omni_sampling_params(async_omni_test_client): """Test AsyncOmni path uses per-stage sampling params.""" response = async_omni_test_client.post( @@ -891,6 +910,31 @@ def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_ assert payloads[2] == "[DONE]" +def test_image_edits_streaming_guardrail_error_uses_400(streaming_image_edit_client): + async def generate_guardrail_error(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs): + raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.") + yield MockStageResult(stage_id=1, final_output_type="image", images=[]) # pragma: no cover + + streaming_image_edit_client.app.state.engine_client.generate = generate_guardrail_error + img_bytes = make_test_image_bytes((16, 16)) + response = streaming_image_edit_client.post( + "/v1/images/edits", + files=[("image", img_bytes)], + data={ + "prompt": "blocked prompt", + "stream": "true", + }, + ) + + assert response.status_code == 200 + payloads = _parse_sse_payloads(response.text) + assert payloads[0]["object"] == "error" + assert payloads[0]["error"]["message"] == "Input was blocked by Cosmos3 guardrails." + assert payloads[0]["error"]["type"] == "BadRequestError" + assert payloads[0]["error"]["code"] == 400 + assert payloads[1] == "[DONE]" + + def test_image_edits_streaming_rejects_single_stage(test_client): img_bytes = make_test_image_bytes((16, 16)) response = test_client.post( diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 1bd353841..b54afd54a 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -31,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.errors import GuardrailViolationError from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -1019,6 +1020,25 @@ def test_invalid_lora_returns_400(test_client): assert "lora object" in failed["error"]["message"].lower() +def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture): + mocker.patch.object( + OmniOpenAIServingVideo, + "generate_video_bytes", + side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."), + ) + response = test_client.post("/v1/videos", data={"prompt": "blocked prompt"}) + 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"] == 400 + assert failed["error"]["message"] == "Input was blocked by Cosmos3 guardrails." + + retrieve = test_client.get(f"/v1/videos/{video_id}") + assert retrieve.status_code == 400 + assert retrieve.json()["error"]["code"] == 400 + + def test_unsupported_image_reference_file_id_returns_400(test_client): response = test_client.post( "/v1/videos", @@ -1521,6 +1541,20 @@ def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture): assert "GPU exploded" in response.json()["detail"] +def test_sync_guardrail_error_returns_400(test_client, mocker: MockerFixture): + mocker.patch.object( + OmniOpenAIServingVideo, + "generate_video_bytes", + side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."), + ) + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "blocked prompt"}, + ) + assert response.status_code == 400 + assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails." + + def test_sync_does_not_create_store_entry(test_client, mocker: MockerFixture): """The sync endpoint should NOT leave any record in VIDEO_STORE — it is stateless by design.""" diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 28717d28c..8f647138b 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -18,6 +18,12 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni import Omni from vllm_omni.entrypoints.omni_base import OmniEngineDeadError +from vllm_omni.errors import ( + OmniClientError, + client_error_from_metadata, + client_error_metadata, + is_client_error_status, +) from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -604,6 +610,71 @@ async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.Mon assert getattr(exc_info.value, "error_stage_id") == 2 +def _enqueue_client_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + ErrorMessage( + request_id=msg["request_id"], + stage_id=2, + error="Input was blocked by Cosmos3 guardrails.", + status_code=400, + error_type="BadRequestError", + ) + ) + + +@pytest.mark.asyncio +async def test_async_omni_propagates_client_error_status(monkeypatch: pytest.MonkeyPatch): + """A non-fatal client error from the orchestrator must be routed to the + requesting generate() call (not raised in the shared dispatcher) and + surface as an OmniClientError carrying status_code/error_type.""" + engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_client_error_message) + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + with pytest.raises(OmniClientError) as exc_info: + async for _ in app.generate(prompt="blocked", request_id="req-1"): + pass + finally: + app.shutdown() + + assert exc_info.value.status_code == 400 + assert exc_info.value.error_type == "BadRequestError" + assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails." + + +def _enqueue_server_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + ErrorMessage( + request_id=msg["request_id"], + stage_id=2, + error="GPU exploded", + ) + ) + + +@pytest.mark.asyncio +async def test_async_omni_propagates_server_error_as_runtime(monkeypatch: pytest.MonkeyPatch): + """A non-fatal error WITHOUT a 4xx status_code is a server fault: it must + surface as a RuntimeError (-> HTTP 500), not an OmniClientError (-> 400). + Guards against the fix over-broadly mapping every error to a client error.""" + engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_server_error_message) + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + # OmniClientError subclasses ValueError, not RuntimeError, so matching + # RuntimeError here also proves it was NOT raised as a client error. + with pytest.raises(RuntimeError) as exc_info: + async for _ in app.generate(prompt="boom", request_id="req-1"): + pass + finally: + app.shutdown() + + assert not isinstance(exc_info.value, OmniClientError) + assert str(exc_info.value) == "GPU exploded" + + 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)] engine = FakeAsyncOmniEngine( @@ -761,6 +832,140 @@ def test_non_fatal_error_raises_runtime(): base._handle_output_message(msg) +def test_non_fatal_client_error_raises_omni_client_error(): + """A non-fatal ErrorMessage carrying a 4xx status_code (e.g. a guardrail + block) must surface as an OmniClientError with the metadata intact, not a + bare RuntimeError. This covers the offline/sync Omni consumer path.""" + base = _make_base() + msg = ErrorMessage( + error="Input was blocked by Cosmos3 guardrails.", + status_code=400, + error_type="BadRequestError", + ) + + with pytest.raises(OmniClientError) as exc_info: + base._handle_output_message(msg) + + assert exc_info.value.status_code == 400 + assert exc_info.value.error_type == "BadRequestError" + assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails." + + +_NON_400_CLIENT_ERRORS = [ + pytest.param(429, "RateLimitError", id="429-too-many-requests"), + pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"), + pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"), + pytest.param(403, "PermissionDeniedError", id="403-forbidden"), +] + + +@pytest.mark.parametrize( + ("status_code", "error_type"), [pytest.param(400, "BadRequestError", id="400")] + _NON_400_CLIENT_ERRORS +) +def test_client_error_metadata_round_trip_preserves_4xx(status_code: int, error_type: str): + original = OmniClientError("blocked", status_code=status_code, error_type=error_type) + + # Outbound: what the broad `except Exception` handlers serialize. + carried_status, carried_type = client_error_metadata(original) + assert carried_status == status_code + assert carried_type == error_type + assert is_client_error_status(carried_status) + + # Inbound: reconstruction at the consuming end. + rebuilt = client_error_from_metadata("blocked", status_code=carried_status, error_type=carried_type) + assert isinstance(rebuilt, OmniClientError) + assert rebuilt.status_code == status_code + assert rebuilt.error_type == error_type + assert str(rebuilt) == "blocked" + + +@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS) +def test_non_fatal_client_error_preserves_non_400_status(status_code: int, error_type: str): + base = _make_base() + msg = ErrorMessage( + error="client side failure", + status_code=status_code, + error_type=error_type, + ) + + with pytest.raises(OmniClientError) as exc_info: + base._handle_output_message(msg) + + assert exc_info.value.status_code == status_code + assert exc_info.value.error_type == error_type + assert str(exc_info.value) == "client side failure" + + +@pytest.mark.parametrize( + "status_code", + [ + pytest.param(None, id="no-status"), + pytest.param(399, id="399-below-4xx"), + pytest.param(500, id="500-server-error"), + pytest.param(503, id="503-service-unavailable"), + ], +) +def test_non_fatal_non_4xx_status_raises_runtime(status_code: int | None): + base = _make_base() + msg = ErrorMessage( + error="server side failure", + status_code=status_code, + error_type="ShouldBeIgnoredForNon4xx", + ) + + with pytest.raises(RuntimeError) as exc_info: + base._handle_output_message(msg) + + assert not isinstance(exc_info.value, OmniClientError) + assert str(exc_info.value) == "server side failure" + + +def _make_enqueue_client_error( + status_code: int, + error_type: str, + error_text: str, +) -> Callable[[FakeAsyncOmniEngine, dict[str, Any]], None]: + def _enqueue(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + ErrorMessage( + request_id=msg["request_id"], + stage_id=2, + error=error_text, + status_code=status_code, + error_type=error_type, + ) + ) + + return _enqueue + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS) +async def test_async_omni_propagates_non_400_client_error_status( + monkeypatch: pytest.MonkeyPatch, + status_code: int, + error_type: str, +): + error_text = f"blocked with {status_code}" + engine = FakeAsyncOmniEngine( + stage_metadata=THREE_STAGE_META, + on_add_request=_make_enqueue_client_error(status_code, error_type, error_text), + ) + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + with pytest.raises(OmniClientError) as exc_info: + async for _ in app.generate(prompt="blocked", request_id="req-1"): + pass + finally: + app.shutdown() + + assert exc_info.value.status_code == status_code + assert exc_info.value.error_type == error_type + assert str(exc_info.value) == error_text + + def test_async_omni_errored_property_alive(): omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9887a0e50..4017ac984 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.base_config import ( from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available +from vllm_omni.errors import client_error_metadata from vllm_omni.quantization import build_quant_config if TYPE_CHECKING: @@ -1062,6 +1063,8 @@ class DiffusionOutput: trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None trajectory_decoded: list[Image.Image] | None = None error: str | None = None + error_status_code: int | None = None + error_type: str | None = None aborted: bool = False abort_message: str | None = None @@ -1102,6 +1105,15 @@ class DiffusionOutput: if self.custom_output: self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()} + @classmethod + def from_exception(cls, exc: BaseException) -> "DiffusionOutput": + status_code, error_type = client_error_metadata(exc) + return cls( + error=str(exc), + error_status_code=status_code, + error_type=error_type, + ) + class DiffusionRequestAbortedError(RuntimeError): """Raised when a diffusion request ends via user-visible abort.""" diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 4e3853a2f..ce1185c66 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusio from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput +from vllm_omni.errors import client_error_from_metadata, is_client_error_status from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.outputs import OmniRequestOutput @@ -231,6 +232,12 @@ class DiffusionEngine: if output.aborted: raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.") if output.error: + if is_client_error_status(output.error_status_code): + raise client_error_from_metadata( + output.error, + status_code=output.error_status_code, + error_type=output.error_type, + ) raise RuntimeError(output.error) logger.debug("Generation completed successfully.") @@ -507,7 +514,7 @@ class DiffusionEngine: request_id=request_id, step_index=None, finished=True, - result=DiffusionOutput(error=str(exc)), + result=DiffusionOutput.from_exception(exc), ) for request_id in sched_output.scheduled_request_ids ] @@ -664,7 +671,7 @@ class DiffusionEngine: request_id=request_id, step_index=None, finished=True, - result=DiffusionOutput(error=str(exc)), + result=DiffusionOutput.from_exception(exc), ) self._process_aborts_queue() diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index a2717bcc8..137f2b887 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -22,6 +22,7 @@ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -143,11 +144,13 @@ class InlineStageDiffusionClient(StageClientBase): logger.info("request_id: %s aborted: %s", request_id, str(e)) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) - error_output = OmniRequestOutput.from_diffusion( + status_code, error_type = client_error_metadata(e) + error_output = OmniRequestOutput.from_error( request_id=request_id, - images=[], + error_message=str(e), + status_code=status_code, + error_type=error_type, ) - error_output.error = str(e) self._output_queue.put_nowait(error_output) finally: self._tasks.pop(request_id, None) @@ -248,11 +251,13 @@ class InlineStageDiffusionClient(StageClientBase): logger.info("request_id: %s aborted: %s", request_id, str(e)) except Exception as e: logger.exception("Batch diffusion request %s failed: %s", request_id, e) - error_output = OmniRequestOutput.from_diffusion( + status_code, error_type = client_error_metadata(e) + error_output = OmniRequestOutput.from_error( request_id=request_id, - images=[], + error_message=str(e), + status_code=status_code, + error_type=error_type, ) - error_output.error = str(e) self._output_queue.put_nowait(error_output) finally: self._tasks.pop(request_id, None) diff --git a/vllm_omni/diffusion/models/cosmos3/guardrails.py b/vllm_omni/diffusion/models/cosmos3/guardrails.py index 3458c2cc3..e249d2826 100644 --- a/vllm_omni/diffusion/models/cosmos3/guardrails.py +++ b/vllm_omni/diffusion/models/cosmos3/guardrails.py @@ -20,6 +20,7 @@ import torch from vllm.logger import init_logger from vllm_omni.diffusion.models.progress_bar import _is_rank_zero +from vllm_omni.errors import GuardrailViolationError if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -67,7 +68,7 @@ def _build_text_guardrail(checker: Any) -> TextGuardrailFn: def text_guardrail(prompt: str) -> None: if not checker.check_text_safety(prompt): # CosmosSafetyChecker logs the specific reason at CRITICAL. - raise ValueError("Guardrail blocked prompt") + raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.") return text_guardrail diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index c8da748eb..613fb64c7 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -228,7 +228,9 @@ class StageDiffusionClient(StageClientBase): elif msg_type == "error": req_id = msg.get("request_id") rpc_id = msg.get("rpc_id") - error_msg = msg.get("error") + error_msg = msg.get("error") or "Unknown diffusion subprocess error." + status_code = msg.get("status_code") + error_type = msg.get("error_type") logger.error( "[StageDiffusionClient] stage-%s [rep-%s] subprocess error for %s: %s", self.stage_id, @@ -245,7 +247,14 @@ class StageDiffusionClient(StageClientBase): # 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)) + self._output_queue.put_nowait( + OmniRequestOutput.from_error( + req_id, + error_msg, + status_code=status_code, + error_type=error_type, + ) + ) # Fields that are subprocess-local and cannot be serialized across # process boundaries. They are recreated in the subprocess with diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 07489d4c9..27c6f16a0 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -33,6 +33,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( ) from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage from vllm_omni.engine.stage_init_utils import set_death_signal +from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -400,12 +401,15 @@ class StageDiffusionProc: ) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) + status_code, error_type = client_error_metadata(e) await response_socket.send( encoder.encode( { "type": "error", "request_id": request_id, "error": str(e), + "status_code": status_code, + "error_type": error_type, } ) ) @@ -483,12 +487,15 @@ class StageDiffusionProc: ) except Exception as e: logger.exception("Batch diffusion request %s failed: %s", rid, e) + status_code, error_type = client_error_metadata(e) await response_socket.send( encoder.encode( { "type": "error", "request_id": rid, "error": str(e), + "status_code": status_code, + "error_type": error_type, } ) ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 3ade1f64d..cd9576437 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -801,7 +801,7 @@ class WorkerProc: except Exception as e: logger.error(f"Error processing RPC: {e}", exc_info=True) if self.result_mq is not None: - self.return_result(DiffusionOutput(error=str(e))) + self.return_result(DiffusionOutput.from_exception(e)) elif isinstance(msg, dict) and msg.get("type") == "shutdown": logger.info("Worker %s: Received shutdown message", self.gpu_id) @@ -817,7 +817,7 @@ class WorkerProc: f"Error executing forward in event loop: {e}", exc_info=True, ) - output = DiffusionOutput(error=str(e)) + output = DiffusionOutput.from_exception(e) try: self.return_result(output) diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 536f8d525..bb73f74a3 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -73,6 +73,8 @@ class UnregisterRemoteReplicaMessage(EngineQueueMessage, kw_only=True): class ErrorMessage(EngineQueueMessage, kw_only=True): type: Literal["error"] = "error" error: str + status_code: int | None = None + error_type: str | None = None fatal: bool = False request_id: str | None = None stage_id: int | None = None diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ae4d2bbcd..f8ff84de8 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -809,6 +809,8 @@ class Orchestrator: request_id=parent_id, stage_id=stage_id, error=output.error, + status_code=getattr(output, "error_status_code", None), + error_type=getattr(output, "error_type", None), ) ) await self._cleanup_request_ids( diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 93f2dadff..6a1f505e4 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -33,6 +33,7 @@ from vllm_omni.entrypoints.omni_base import ( OmniBase, OmniEngineDeadError, ) +from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput @@ -483,10 +484,13 @@ class AsyncOmni(EngineClient, OmniBase): except (asyncio.CancelledError, GeneratorExit): cancelled = True except Exception as error: + status_code, error_type = client_error_metadata(error) await req_state.queue.put( ErrorMessage( request_id=request_id, error=str(error), + status_code=status_code, + error_type=error_type, ) ) finally: @@ -589,7 +593,7 @@ class AsyncOmni(EngineClient, OmniBase): result.error, error_stage_id=result.stage_id, ) - raise RuntimeError(result.error) + self._raise_nonfatal_error_message(result) if not isinstance(result, OutputMessage): logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) @@ -662,6 +666,18 @@ class AsyncOmni(EngineClient, OmniBase): await self.event_resolver.resolve(msg) continue + if isinstance(msg, ErrorMessage) and not msg.fatal: + req_state = self.request_states.get(msg.request_id) + if req_state is not None: + await req_state.queue.put(msg) + else: + logger.warning( + "[%s] dropping non-fatal error for unknown req %s", + self._name, + msg.request_id, + ) + continue + should_continue, _, stage_id, req_state = self._handle_output_message(msg) if should_continue: continue diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index d415c07ea..029665f59 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -20,6 +20,7 @@ from vllm_omni.engine.messages import ( from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e +from vllm_omni.errors import raise_client_error_or from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize from vllm_omni.metrics.prometheus import OmniPrometheusMetrics from vllm_omni.metrics.stats import OrchestratorAggregator @@ -372,7 +373,7 @@ class OmniBase(PDDisaggregationMixin): msg.error, error_stage_id=msg.stage_id, ) - raise RuntimeError(msg.error) + self._raise_nonfatal_error_message(msg) if not isinstance(msg, OutputMessage): logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg.type) @@ -408,6 +409,15 @@ class OmniBase(PDDisaggregationMixin): return False, req_id, stage_id, req_state + def _raise_nonfatal_error_message(self, msg: ErrorMessage) -> None: + """Raise the exception for a non-fatal, request-scoped error message.""" + raise_client_error_or( + msg.error, + status_code=msg.status_code, + error_type=msg.error_type, + fallback=RuntimeError, + ) + def _check_engine_output_error( self, result: OutputMessage, @@ -424,6 +434,8 @@ class OmniBase(PDDisaggregationMixin): error_text = getattr(engine_outputs, "error", None) if error_text is None: return + status_code = engine_outputs.error_status_code + error_type = engine_outputs.error_type logger.error( "[%s] Stage error for req=%s stage-%s: %s", self.__class__.__name__, @@ -437,7 +449,12 @@ class OmniBase(PDDisaggregationMixin): error_text, error_stage_id=stage_id, ) - raise EngineGenerateError(error_text) + raise_client_error_or( + error_text, + status_code=status_code, + error_type=error_type, + fallback=EngineGenerateError, + ) def _process_single_result( self, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0b419c8ed..6df57b70c 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -142,6 +142,7 @@ from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url, decode_input_reference from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI +from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace @@ -1812,6 +1813,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise + except OmniClientError as e: + logger.info("Client error during image generation: %s", e) + raise HTTPException(status_code=e.status_code, detail=e.message) from e except ValueError as e: logger.error(f"Validation error: {e}") raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) @@ -2167,6 +2171,9 @@ async def edit_images( return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise + except OmniClientError as e: + logger.info("Client error during image edit: %s", e) + raise HTTPException(status_code=e.status_code, detail=e.message) from e except ValueError as e: logger.error(f"Validation error: {e}") raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) @@ -2667,6 +2674,9 @@ def _video_error_from_exception(exc: Exception) -> VideoError: message = str(exc.detail) if exc.detail else str(exc) return VideoError(code=exc.status_code, message=message) + if isinstance(exc, OmniClientError): + return VideoError(code=exc.status_code, message=exc.message) + if isinstance(exc, (EngineGenerateError, EngineDeadError)): err = create_error_response(exc) return VideoError(code=err.error.code, message=err.error.message) @@ -3021,6 +3031,9 @@ async def create_video_sync( return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise + except OmniClientError as exc: + logger.info("Client error during sync video generation: %s", exc) + raise HTTPException(status_code=exc.status_code, detail=exc.message) from exc except Exception as exc: logger.exception("Sync video generation failed for request_id=%s", request_id) raise HTTPException( diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 5dec1768a..35cfafcff 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -113,6 +113,7 @@ from vllm_omni.entrypoints.openai.utils import ( resolve_diffusion_od_config, validate_requested_speaker, ) +from vllm_omni.errors import OmniClientError from vllm_omni.lora.request import LoRARequest from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate @@ -2869,6 +2870,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "[OmniOpenAIServingChat] Engine dead during streaming image edit, " "but cannot terminate server (no raw_request context)." ) + except OmniClientError as exc: + logger.info("Client error during streaming image edit: %s", exc) + chunk = ImageEditStreamError( + created=created, + model=model, + error={ + "message": exc.message, + "type": exc.error_type, + "code": exc.status_code, + }, + ) + yield f"data: {chunk.model_dump_json()}\n\n" except Exception as exc: logger.exception("Streaming image edit failed: %s", exc) chunk = ImageEditStreamError( @@ -3265,6 +3278,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return response + except OmniClientError as e: + logger.info("Client error during diffusion chat completion: %s", e) + return self._create_error_response( + e.message, + err_type=e.error_type, + status_code=e.status_code, + ) except Exception as e: logger.exception("Diffusion chat completion failed: %s", e) return self._create_error_response( diff --git a/vllm_omni/errors.py b/vllm_omni/errors.py new file mode 100644 index 000000000..1bc7f5913 --- /dev/null +++ b/vllm_omni/errors.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Request-scoped client error types shared across vLLM-Omni entrypoints.""" + +from __future__ import annotations + +from collections.abc import Callable +from http import HTTPStatus +from typing import NoReturn + +DEFAULT_CLIENT_ERROR_TYPE = "BadRequestError" + + +class OmniClientError(ValueError): + """ + Request-scoped error that should be surfaced as a 4xx response. + One example of using OmniClientError is GuardrailViolationError, which is captured and resurfaced + as HTTP 400 error code, instead of a generic 500. OmniClientError should be used for any exceptions + which need to be resurfaced as 4xx, as opposed to EngineDeadError/EngineGenerateError which are resurfaced as 500. + """ + + def __init__( + self, + message: str, + *, + status_code: int = HTTPStatus.BAD_REQUEST.value, + error_type: str = DEFAULT_CLIENT_ERROR_TYPE, + ) -> None: + super().__init__(message) + self.message = message + self.status_code = int(status_code) + self.error_type = error_type + + +class GuardrailViolationError(OmniClientError): + """Raised when a model guardrail rejects request content.""" + + +def client_error_metadata(exc: BaseException) -> tuple[int | None, str | None]: + if isinstance(exc, OmniClientError): + return exc.status_code, exc.error_type + return None, None + + +def client_error_from_metadata( + message: str, + *, + status_code: int | None, + error_type: str | None, +) -> OmniClientError: + return OmniClientError( + message, + status_code=status_code or HTTPStatus.BAD_REQUEST.value, + error_type=error_type or DEFAULT_CLIENT_ERROR_TYPE, + ) + + +def is_client_error_status(status_code: int | None) -> bool: + return status_code is not None and 400 <= int(status_code) < 500 + + +def raise_client_error_or( + message: str, + *, + status_code: int | None, + error_type: str | None, + fallback: Callable[[str], BaseException], +) -> NoReturn: + """Raise a client error for 4xx statuses, otherwise raise ``fallback(message)``. + + Centralizes the "client-error-or-fallback" decision shared by the engine + error paths so the status mapping lives in one place. + """ + if is_client_error_status(status_code): + raise client_error_from_metadata( + message, + status_code=status_code, + error_type=error_type, + ) + raise fallback(message) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index b4d308ebf..7b3aab2cb 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -102,12 +102,17 @@ class OmniRequestOutput: # error handling error: str | None = None + error_status_code: int | None = None + error_type: str | None = None @classmethod def from_error( cls, request_id: str, error_message: str, + *, + status_code: int | None = None, + error_type: str | None = None, ) -> "OmniRequestOutput": """Create a terminal error output. @@ -122,6 +127,8 @@ class OmniRequestOutput: request_id=request_id, finished=True, error=error_message, + error_status_code=status_code, + error_type=error_type, ) @classmethod