diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 4217e2156..5c7066e0c 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -1813,6 +1813,56 @@ async def test_stage_pool_submit_update_refreshes_output_processor_state() -> No @pytest.mark.asyncio +async def test_handle_streaming_update_unknown_request_is_dropped() -> None: + """A streaming update for an untracked request must not resurrect it. + + Updates always follow the first-chunk add_request through the same + ordered submission queue, so an unknown id can only mean the request + already finished or was aborted (e.g. the client disconnected). + Falling back to add_request created headless sessions that kept + cycling through the stages (issue #4271). + """ + + class RecordingPool: + def __init__(self) -> None: + self.calls: list[tuple[str, Any]] = [] + + async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int: + self.calls.append((request_id, prompt_text)) + return 0 + + pool = RecordingPool() + orchestrator = object.__new__(Orchestrator) + orchestrator.async_chunk = False + orchestrator.request_states = {} + orchestrator.stage_pools = [pool] + add_request_calls: list[str] = [] + + async def record_add_request(msg) -> None: + add_request_calls.append(msg.request_id) + + orchestrator._handle_add_request = record_add_request + + await orchestrator._handle_streaming_update( + StageSubmissionMessage( + type="streaming_update", + request_id="req-unknown", + prompt=SimpleNamespace(request_id="req-unknown", prompt_token_ids=[1], resumable=True), + original_prompt={"prompt": "segment-2"}, + output_prompt_text="segment-2", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + preprocess_ms=0.0, + request_timestamp=time.time(), + enqueue_ts=time.perf_counter(), + ) + ) + + assert pool.calls == [] + assert add_request_calls == [] + assert "req-unknown" not in orchestrator.request_states + + async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None: class RecordingPool: def __init__(self) -> None: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index affd0e452..fa10653be 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -681,24 +681,16 @@ class Orchestrator: final_stage_id = msg.final_stage_id req_state = self.request_states.get(request_id) if req_state is None: + # Streaming updates always follow the first-chunk add_request + # through the same ordered submission queue, so an unknown id + # here can only mean the request already finished or was + # aborted (e.g. the client disconnected). Re-adding it would + # resurrect a headless session that keeps cycling through the + # stages with nobody consuming its outputs (issue #4271). logger.warning( - "[Orchestrator] streaming_update for unknown req=%s, falling back to add_request", + "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)", request_id, ) - fallback_msg = StageSubmissionMessage( - type="add_request", - request_id=msg.request_id, - prompt=msg.prompt, - original_prompt=msg.original_prompt, - output_prompt_text=msg.output_prompt_text, - sampling_params_list=msg.sampling_params_list, - final_stage_id=msg.final_stage_id, - final_output_stage_ids=msg.final_output_stage_ids, - preprocess_ms=msg.preprocess_ms, - request_timestamp=msg.request_timestamp, - enqueue_ts=msg.enqueue_ts, - ) - await self._handle_add_request(fallback_msg) return if msg.sampling_params_list: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 332e5b921..1437cc680 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -628,6 +628,13 @@ class AsyncOmni(EngineClient, OmniBase): logger.debug(f"[AsyncOmni] Request {request_id} completed") + # The input pump can outlive a normally-completed generation + # (e.g. a realtime client that keeps streaming after the final + # output); a live pump keeps issuing streaming updates for a + # request the orchestrator no longer tracks. + if input_stream_task is not None and not input_stream_task.done(): + input_stream_task.cancel() + self._log_summary_and_cleanup(request_id) except (asyncio.CancelledError, GeneratorExit): @@ -1059,7 +1066,10 @@ class AsyncOmni(EngineClient, OmniBase): """Submit request IDs to be aborted to the engine.""" await self.engine.abort_async(request_ids) for rid in request_ids: - self.request_states.pop(rid, None) + state = self.request_states.pop(rid, None) + input_stream_task = getattr(state, "input_stream_task", None) + if input_stream_task is not None and not input_stream_task.done(): + input_stream_task.cancel() if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 0f0efbef8..aad7ea658 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -138,6 +138,7 @@ class RealtimeConnection(VllmRealtimeConnection): is_streaming=True, ) + result_gen = None try: result_gen = self.engine.generate( prompt=streaming_input_gen, @@ -196,6 +197,16 @@ class RealtimeConnection(VllmRealtimeConnection): logger.exception("Error in generation: %s", e) await self.send_error(str(e), "processing_error") finally: + # Close the generator explicitly so AsyncOmni.generate's cleanup + # (input-pump cancellation and engine-side abort) runs now rather + # than whenever the event loop garbage-collects the async + # generator; the delay window is where a disconnected session + # keeps cycling through the stages (issue #4271). + if result_gen is not None: + try: + await result_gen.aclose() + except Exception: + logger.exception("Failed to close realtime result generator") # Always send terminal event so clients don't hang forever. if self._is_connected and not audio_done_sent: try: @@ -206,4 +217,10 @@ class RealtimeConnection(VllmRealtimeConnection): self.audio_queue.get_nowait() async def send_json(self, payload: dict): - await self.websocket.send_text(json.dumps(payload)) + try: + await self.websocket.send_text(json.dumps(payload)) + except Exception: + # A failed send means the client is gone; flag it so the + # generation loop stops instead of retrying into a dead socket. + self._is_connected = False + raise