From c97b5fc741ea44acaed37efe520ef8600d1c2abf Mon Sep 17 00:00:00 2001 From: shiyichuan <93317314+CarrotSwordsman@users.noreply.github.com> Date: Thu, 17 Sep 2026 02:38:45 +0800 Subject: [PATCH] [Bugfix][Frontend] Route synthetic aborts through final stage (#7006) Signed-off-by: mershi Co-authored-by: mershi --- tests/entrypoints/test_omni_entrypoints.py | 38 +++++++++++++++++++ vllm_omni/entrypoints/async_omni.py | 1 + vllm_omni/entrypoints/async_omni_base.py | 30 ++++++++++++--- vllm_omni/entrypoints/client_request_state.py | 5 +++ 4 files changed, 69 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 32d7d9988..ea57d07ab 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import gc import queue from collections.abc import Callable @@ -634,6 +635,43 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc assert "req-1" not in app.request_states +@pytest.mark.asyncio +async def test_async_omni_abort_yields_terminal_from_requested_final_stage(monkeypatch: pytest.MonkeyPatch): + stage_metadata = [ + _stage_meta(stage_type="llm", final_output=False, final_output_type=None), + _stage_meta(stage_type="llm", final_output=True, final_output_type="audio"), + ] + engine = FakeAsyncOmniEngine(stage_metadata=stage_metadata) + _patch_engine(monkeypatch, engine) + app = AsyncOmni("dummy-model") + + async def collect_outputs() -> list[OmniRequestOutput]: + return [ + output + async for output in app.generate( + prompt="hello", + request_id="req-1", + output_modalities=["audio"], + ) + ] + + try: + generate_task = asyncio.create_task(collect_outputs()) + while not engine.submitted: + await asyncio.sleep(0) + await app.abort("req-1") + outputs = await asyncio.wait_for(generate_task, timeout=1) + finally: + app.shutdown() + + assert len(outputs) == 1 + assert outputs[0].stage_id == 1 + assert outputs[0].final_output_type == "audio" + assert outputs[0].finished is True + assert outputs[0].outputs[0].finish_reason == "abort" + assert "req-1" not in app.request_states + + @pytest.mark.asyncio async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 6114bbea2..1291e68f0 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -236,6 +236,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): req_state = ClientRequestState( request_id=request_id, external_request_id=external_request_id, + final_stage_id=final_stage_id_for_e2e, ) req_state.metrics = metrics req_state.request_arrival_ts = wall_start_ts diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index e7c86de97..e9bf2c77e 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -479,7 +479,22 @@ class AsyncOmniBase(OmniBase): if state is not None and rid not in delivered: queue = getattr(state, "queue", None) if queue is not None: - await state.queue.put(self._synthetic_abort_output_message(rid)) + final_stage_id = getattr(state, "final_stage_id", None) + if final_stage_id is None: + final_stage_id = 0 + stage_metadata = getattr(self, "_stage_meta_list", ()) + final_output_type = ( + getattr(stage_metadata[final_stage_id], "final_output_type", None) + if final_stage_id < len(stage_metadata) + else None + ) or "text" + await state.queue.put( + self._synthetic_abort_output_message( + rid, + stage_id=final_stage_id, + final_output_type=final_output_type, + ) + ) delivered.add(rid) for rid in request_ids: self._record_request_failure_once(rid, reason="client_abort") @@ -491,13 +506,18 @@ class AsyncOmniBase(OmniBase): logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) @staticmethod - def _synthetic_abort_output_message(request_id: str) -> OutputMessage: + def _synthetic_abort_output_message( + request_id: str, + *, + stage_id: int, + final_output_type: str, + ) -> OutputMessage: """Terminal abort OutputMessage used when the engine returned none.""" engine_output = OmniRequestOutput( request_id=request_id, finished=True, - stage_id=0, - final_output_type="text", + stage_id=stage_id, + final_output_type=final_output_type, outputs=[ CompletionOutput( index=0, @@ -512,7 +532,7 @@ class AsyncOmniBase(OmniBase): ) return OutputMessage( request_id=request_id, - stage_id=0, + stage_id=stage_id, replica_id=None, engine_outputs=engine_output, metrics=None, diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 22f1c6d5e..ca9e96bab 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio from vllm_omni.metrics import OrchestratorAggregator @@ -11,10 +14,12 @@ class ClientRequestState: request_id: str, external_request_id: str | None = None, queue: asyncio.Queue | None = None, + final_stage_id: int | None = None, ): self.request_id = request_id self.external_request_id = external_request_id self.stage_id: int | None = None + self.final_stage_id: int | None = final_stage_id self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None # Request-scoped idempotency guard for Prometheus failure counters.