[Bugfix][Frontend] Surface streaming terminal submit failures (#7011)

Signed-off-by: BANANASJIM <bananasjim1@gmail.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/entrypoints/test_async_omni.py

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 
 import anyio
 import pytest
+from vllm.engine.protocol import StreamingInput
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
@@ -67,6 +68,7 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     omni._paused = False
     omni.engine = SimpleNamespace(
         num_stages=1,
+        stage_configs=[],
         add_request_async=fake_add_request,
         abort_async=fake_abort_request,
     )
@@ -473,6 +475,61 @@ def test_generate_accepts_request_after_repeated_cancellations():
     asyncio.run(run_test())
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("has_chunk", [False, True])
+@pytest.mark.parametrize("input_fails", [False, True])
+def test_streaming_input_terminal_submit_failure_reaches_generate(mocker, has_chunk, input_fails):
+    async def run_test():
+        submissions: list[tuple[str, bool]] = []
+
+        async def add_request_async(*, resumable, **kwargs):
+            del kwargs
+            submissions.append(("add", resumable))
+            if not resumable:
+                raise RuntimeError("terminal streaming submit failed")
+
+        async def add_streaming_update_async(*, resumable, **kwargs):
+            del kwargs
+            submissions.append(("update", resumable))
+            assert not resumable
+            raise RuntimeError("terminal streaming submit failed")
+
+        omni = get_async_omni_instance(fake_add_request=add_request_async)
+        omni.engine.add_streaming_update_async = add_streaming_update_async
+        omni.engine.stage_clients = []
+        omni.engine.model_config = object()
+        del omni._process_orchestrator_results
+        mocker.patch(
+            "vllm_omni.entrypoints.async_omni.extract_prompt_components",
+            return_value=(None, None, None),
+        )
+        error_metadata = mocker.spy(async_omni_mod, "client_error_metadata")
+
+        async def input_stream():
+            if has_chunk:
+                yield StreamingInput(prompt={"prompt": "chunk"})
+            if input_fails:
+                raise RuntimeError("streaming input failed")
+
+        async def consume():
+            async for _ in omni.generate(
+                prompt=input_stream(),
+                request_id="stream-submit-error",
+                sampling_params_list=[SamplingParams(output_kind=RequestOutputKind.DELTA)],
+                output_modalities=["text"],
+            ):
+                pass
+
+        expected_error = "streaming input failed" if input_fails else "terminal streaming submit failed"
+        with pytest.raises(RuntimeError, match=expected_error):
+            await asyncio.wait_for(consume(), timeout=1.0)
+        assert error_metadata.call_count == 1
+        assert submissions == ([("add", True), ("update", False)] if has_chunk else [("add", False)])
+        assert omni.request_states == {}
+
+    asyncio.run(run_test())
+
+
 @pytest.mark.cpu
 def test_generate_cancel_bounds_cleanup_abort(monkeypatch):
     """Cancelled generate() must not wait forever on a wedged orchestrator abort."""

vllm_omni/entrypoints/async_omni.py

@@ -371,6 +371,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         stage0_params = cast(SamplingParams, stage0_params)
         req_state = self.request_states[request_id]
         has_submitted_first_chunk = False
+        input_error_reported = False
 
         # NOTE: InputProcessor in vLLM should generally do this too, but for
         # now we do it defensively. TODO (Alex) ensure clones/copying are optimized
@@ -382,6 +383,21 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             if first_chunk_submitted is not None and not first_chunk_submitted.done():
                 first_chunk_submitted.set_result(None)
 
+        async def _report_input_error(error: Exception) -> None:
+            nonlocal input_error_reported
+            if input_error_reported:
+                return
+            input_error_reported = True
+            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,
+                )
+            )
+
         async def handle_inputs() -> None:
             nonlocal has_submitted_first_chunk
             cancelled = False
@@ -427,15 +443,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             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,
-                    )
-                )
+                await _report_input_error(error)
             finally:
                 try:
                     if not cancelled:
@@ -474,9 +482,11 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                                 )
                             )
                             has_submitted_first_chunk = True
+                except Exception as error:
+                    await _report_input_error(error)
                 finally:
                     # Unblock generate() even on cancel / empty stream / submit
-                    # failure so it can observe a terminal abort or empty result.
+                    # failure so it can observe the queued terminal result or error.
                     _mark_first_chunk_submitted()
 
         input_stream_task = asyncio.create_task(handle_inputs())