[Bugfix][Frontend] Remove async-chunk guard from realtime API (#5383)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -57,13 +57,22 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
     )
 
 
-def _run_resumable_segment_stop(session: Request) -> None:
+def _run_resumable_segment_stop(
+    session: Request,
+    *,
+    session_finished: bool = False,
+):
     sched = MagicMock()
     sched.requests = {session.request_id: session}
     sched.perf_metrics = None
     sched.structured_output_manager.should_advance.return_value = False
-    sched._update_request_with_output.return_value = ([42], True)
-    sched._handle_stopped_request.return_value = False
+
+    def stop_request(request: Request, _token_ids: list[int]):
+        request.status = RequestStatus.FINISHED_STOPPED
+        return [42], True
+
+    sched._update_request_with_output.side_effect = stop_request
+    sched._handle_stopped_request.return_value = session_finished
     sched.chunk_transfer_adapter = None
     sched.running = [session]
     sched.waiting_for_transfer_free = set()
@@ -91,7 +100,7 @@ def _run_resumable_segment_stop(session: Request) -> None:
     model_runner_output.req_id_to_index = {session.request_id: 0}
     model_runner_output.routed_experts = None
 
-    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+    return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
 
 
 @pytest.mark.parametrize("outstanding_async_tokens", [0, 1, 2])
@@ -116,6 +125,18 @@ def test_resumable_segment_stop_reconciles_async_placeholders(
     assert session._output_token_ids == []
 
 
+def test_resumable_session_terminal_is_not_marked_as_segment_boundary() -> None:
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.resumable = True
+
+    outputs = _run_resumable_segment_stop(session, session_finished=True)
+
+    output = outputs[session.client_index].outputs[0]
+    assert output.finish_reason is not None
+    assert output.is_segment_finished is False
+
+
 def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
     sched = _make_scheduler(stage_id=0)
     session = _make_request()

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -9,6 +9,8 @@ update payloads early.
 from __future__ import annotations
 
 from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 
@@ -17,7 +19,9 @@ import pytest
 # isort: off
 import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
@@ -77,6 +81,58 @@ def _make_update(**kwargs) -> StreamingUpdate:
     return StreamingUpdate(**defaults)
 
 
+def test_resumable_generation_stop_marks_segment_boundary() -> None:
+    session = _make_request(request_id="req-generation-segment")
+    session.status = RequestStatus.RUNNING
+    session.resumable = True
+    session.num_computed_tokens = len(session.prompt_token_ids)
+
+    sched = MagicMock()
+    sched.requests = {session.request_id: session}
+    sched.perf_metrics = None
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        is_done_receiving_chunks=lambda _request_id: True,
+        segment_finished_requests={session.request_id},
+    )
+    sched._handle_stopped_request.return_value = False
+    sched.running = [session]
+    sched.waiting = MagicMock()
+    sched.skipped_waiting = MagicMock()
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._pending_finish_reqs = []
+    sched.recompute_kv_load_failures = False
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [[]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {session.request_id: 0}
+    model_runner_output.routed_experts = None
+
+    outputs = OmniGenerationScheduler.update_from_output(
+        sched,
+        scheduler_output,
+        model_runner_output,
+    )
+
+    output = outputs[session.client_index].outputs[0]
+    assert output.finish_reason is not None
+    assert output.is_segment_finished is True
+
+
 class TestReplaceSessionWithStreamingUpdate:
     def test_resets_tokens_and_prompt_from_update(self) -> None:
         sched = _SchedulerStub()

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import queue
 from types import SimpleNamespace
 from typing import Any
@@ -241,3 +242,38 @@ async def test_async_route_forwards_to_outgoing_only_stage() -> None:
     await orchestrator._route_output(0, 0, output, req_state, None)
 
     orchestrator._forward_to_next_stage.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_streaming_segment_does_not_complete_final_output_stage() -> None:
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = True
+    orchestrator._pd_pair = None
+    orchestrator._cfg_tracker = SimpleNamespace(
+        is_companion=lambda _request_id: False,
+        has_companions=lambda _request_id: False,
+        cleanup_parent=lambda _request_id: [],
+    )
+    orchestrator.stage_pools = [SimpleNamespace(final_output=True)]
+    orchestrator.output_async_queue = asyncio.Queue()
+    orchestrator._cleanup_request_ids = AsyncMock()
+
+    req_state = OrchestratorRequestState(
+        request_id="req-segment-final-output",
+        sampling_params_list=[SamplingParams(max_tokens=1)],
+        final_stage_id=0,
+        final_output_stage_ids={0},
+    )
+    req_state.streaming.enabled = True
+    req_state.streaming.segment_finished = True
+    output = SimpleNamespace(
+        request_id=req_state.request_id,
+        finished=True,
+    )
+
+    await orchestrator._route_output(0, 0, output, req_state, None)
+
+    assert req_state.finished_final_output_stage_ids == set()
+    orchestrator._cleanup_request_ids.assert_not_awaited()
+    routed = orchestrator.output_async_queue.get_nowait()
+    assert routed.finished is False

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -234,12 +234,6 @@ class TestQwen3OmniRealtimeWebSocket:
     @pytest.mark.omni
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
     @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True)
-    @pytest.mark.skip(
-        reason=(
-            "/v1/realtime rejects async_chunk since #3907 (server closes with 1000 after "
-            "unsupported error); tracked in #5363"
-        )
-    )
     def test_streaming_audio_input_pcm_output_async_chunk(self, omni_server) -> None:
         """Merge CI: async_chunk on, paced upload, full accuracy check."""
         pcm16 = _synthetic_pcm16_input()

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -429,8 +429,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
                 finish_reason = request.get_finished_reason()
-                is_segment_finished = True
                 finished = self._handle_stopped_request(request)
+                is_segment_finished = not finished
+                if finished:
+                    request.resumable = False
                 if not finished:
                     # for streaming input request only
                     if self.chunk_transfer_adapter:
@@ -494,7 +496,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
-            if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or is_segment_finished):
+            if self.chunk_transfer_adapter is not None and (
+                inter_stage_output is not None or is_segment_finished or finished
+            ):
                 self.chunk_transfer_adapter.save_async(
                     inter_stage_output,
                     request,

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -480,6 +480,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             mm_output = mm_outputs[req_index] if mm_outputs else None
             status_before_stop = request.status
             finish_reason = None
+            is_segment_finished = False
             routed_experts = None
 
             # One-shot generation request: finish after its current input unit
@@ -503,6 +504,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
+                is_segment_finished = not finished
                 if not finished:
                     # for streaming input request only
                     if self.chunk_transfer_adapter:
@@ -556,6 +558,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         trace_headers=request.trace_headers,
                         routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
+                        is_segment_finished=is_segment_finished,
                     )
                 )
             else:
@@ -570,6 +573,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             request.status = RequestStatus.FINISHED_STOPPED
             finish_reason = request.get_finished_reason()
             finished = self._handle_stopped_request(request)
+            is_segment_finished = not finished
             kv_transfer_params = None
             if finished:
                 kv_transfer_params = self._free_request(request)
@@ -587,6 +591,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     events=request.take_events(),
                     kv_transfer_params=kv_transfer_params,
                     trace_headers=request.trace_headers,
+                    is_segment_finished=is_segment_finished,
                 )
             )
             stopped_running_reqs.add(request)

vllm_omni/engine/orchestrator.py

@@ -1183,7 +1183,11 @@ class Orchestrator:
             return
 
         request_finished = False
-        if finished and self.stage_pools[stage_id].final_output:
+        if (
+            finished
+            and self.stage_pools[stage_id].final_output
+            and not (req_state.streaming.enabled and req_state.streaming.segment_finished)
+        ):
             req_state.finished_final_output_stage_ids.add(stage_id)
             final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
             request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)

vllm_omni/entrypoints/openai/api_server.py

@@ -1635,22 +1635,6 @@ async def realtime_websocket(websocket: WebSocket):
         await duplex_handler.handle_realtime_session(websocket)
         return
 
-    engine_client = getattr(websocket.app.state, "engine_client", None)
-    if engine_client is not None and getattr(engine_client, "async_chunk", False):
-        await websocket.accept()
-        await websocket.send_json(
-            {
-                "type": "error",
-                "error": (
-                    "The /v1/realtime API is not supported when async_chunk is enabled on the server. "
-                    "Use a stage configuration with async_chunk disabled and restart the server before using "
-                    "this endpoint."
-                ),
-                "code": "unsupported",
-            }
-        )
-        await websocket.close()
-        return
     serving = getattr(websocket.app.state, "openai_serving_realtime", None)
     if serving is None:
         await websocket.accept()