diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 6bbf26d1a..1bcc6e08c 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/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.""" diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 1c27def99..63f13d725 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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())