diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index e2ad98f56..9c0533059 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/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() diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index d94996014..28083349a 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/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() diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index e74b9ea18..0c5bb217d 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/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 diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index d935949b3..50f8eef91 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/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() diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index a9702eeb3..7dc2255ec 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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, diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 3882fd42d..534019513 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e37111063..9a1d63170 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 56c7ccf46..5338761ad 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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()