From 146350155321de06d178f50a5d5fff7afaad0128 Mon Sep 17 00:00:00 2001 From: David Cheung Date: Thu, 15 Oct 2026 06:32:19 -0700 Subject: [PATCH] [Bugfix] Terminate streaming sessions on ABORT/ERROR and prevent zombie sessions (#58942) Signed-off-by: David Cheung Signed-off-by: Nick Hill Co-authored-by: Benedikt Falk Co-authored-by: Nick Hill --- tests/v1/engine/test_engine_core_client.py | 37 ++++++++ tests/v1/engine/test_output_processor.py | 84 +++++++++++++++++++ .../test_async_llm_streaming.py | 50 +++++++++++ vllm/v1/engine/async_llm.py | 13 ++- vllm/v1/engine/core_client.py | 25 +++++- vllm/v1/engine/output_processor.py | 15 ++-- 6 files changed, 213 insertions(+), 11 deletions(-) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 9bcc9d91eb..10e9f1e92f 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -210,6 +210,7 @@ def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncM client = object.__new__(DPLBAsyncMPClient) client.client_count = client_count client.reqs_in_flight = {} + client.streaming_reqs = set() client.engine_inflight = Counter() client.core_engines = [bytes([i, 0]) for i in range(num_engines)] client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)] @@ -407,6 +408,42 @@ def test_dplb_finished_requests_release_inflight(): assert req.request_id not in client.reqs_in_flight +@pytest.mark.parametrize("final_sent", [False, True]) +def test_dplb_streaming_session_abort_goes_to_all_engines(final_sent: bool): + """A streaming-input chunk in flight when its session finished may + re-create it on any engine, so an open session's abort reaches every + engine. Other aborts go only to the routed engine.""" + client = _make_dplb_client(num_engines=3) + client.resources = SimpleNamespace(engine_dead=False) + sent: dict[bytes, list[str]] = {} + + async def abort_requests(request_ids, engine): + sent[engine] = request_ids + + client._abort_requests = abort_requests + + chunk = make_request(SamplingParams(max_tokens=1)) + chunk.resumable = True + client.get_core_engine_for_request(chunk) + if final_sent: + chunk.resumable = False + client.get_core_engine_for_request(chunk) + # The engine finished the session, dropping its routing entry. + outputs = EngineCoreOutputs(finished_requests={chunk.request_id}) + asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs)) + client.reqs_in_flight["live"] = client.core_engines[1] + + asyncio.run(client.abort_requests_async([chunk.request_id, "live"])) + + if final_sent: + expected = {client.core_engines[1]: ["live"]} + else: + expected = {engine: [chunk.request_id] for engine in client.core_engines} + expected[client.core_engines[1]].append("live") + assert sent == expected + assert not client.streaming_reqs + + @pytest.mark.parametrize( ("effective_size", "other_size"), [(None, None), (4224, 4224), (4224, 1056), (4224, None)], diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 3c186a4e73..eaa45e3320 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -1642,6 +1642,90 @@ async def test_cumulative_output_collector_n(): assert third[0].text == "c" +def _make_streaming_input_request( + request_id: str, prompt_token_ids: list[int] +) -> EngineCoreRequest: + """One chunk of a streaming-input (resumable) request.""" + return EngineCoreRequest( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + mm_features=None, + sampling_params=SamplingParams( + detokenize=False, output_kind=RequestOutputKind.DELTA + ), + pooling_params=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + resumable=True, + external_req_id=f"{request_id}-ext", + ) + + +def _add_streaming_request_with_queued_chunk( + processor: OutputProcessor, request_id: str +) -> RequestOutputCollector: + queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id=request_id) + processor.add_request( + _make_streaming_input_request(request_id, [1, 2, 3]), None, queue=queue + ) + processor.add_request(_make_streaming_input_request(request_id, [7, 8]), None) + return queue + + +def test_streaming_input_stop_is_chunk_boundary(): + """STOP is a per-chunk boundary: finished is forced False, the queued + chunk is applied, and the request state stays alive.""" + processor = OutputProcessor(None, log_stats=False) + queue = _add_streaming_request_with_queued_chunk(processor, "stream-req") + + processed = processor.process_outputs( + [ + EngineCoreOutput( + request_id="stream-req", + new_token_ids=[4], + finish_reason=FinishReason.STOP, + ) + ] + ) + + assert processed.reqs_to_abort == [] + assert processor.get_num_unfinished_requests() == 1 + output = queue.get_nowait() + assert output is not None + assert output.finished is False + assert output.outputs[0].finish_reason == "stop" + req_state = processor.request_states["stream-req"] + assert req_state.prompt_token_ids is not None + assert req_state.prompt_token_ids[-2:] == [7, 8] + + +@pytest.mark.parametrize("finish_reason", [FinishReason.ABORT, FinishReason.ERROR]) +def test_streaming_input_abort_error_is_terminal(finish_reason: FinishReason): + """ABORT/ERROR is terminal, not a chunk boundary: finished stays True, + the state is freed, and an abort is routed back to the engine core.""" + processor = OutputProcessor(None, log_stats=False) + queue = _add_streaming_request_with_queued_chunk(processor, "stream-req") + + processed = processor.process_outputs( + [ + EngineCoreOutput( + request_id="stream-req", + new_token_ids=[], + finish_reason=finish_reason, + ) + ] + ) + + assert processor.get_num_unfinished_requests() == 0 + assert processed.reqs_to_abort == ["stream-req"] + output = queue.get_nowait() + assert output is not None + assert output.finished is True + assert output.outputs[0].finish_reason == str(finish_reason) + + @pytest.mark.parametrize("runner", ["generate", "pooling"]) @pytest.mark.parametrize("abort_by", ["internal", "external"]) def test_abort_requests(runner: str, abort_by: str, dummy_test_vectors): diff --git a/tests/v1/streaming_input/test_async_llm_streaming.py b/tests/v1/streaming_input/test_async_llm_streaming.py index 21a49024ce..9dac5ee652 100644 --- a/tests/v1/streaming_input/test_async_llm_streaming.py +++ b/tests/v1/streaming_input/test_async_llm_streaming.py @@ -3,6 +3,7 @@ import asyncio from collections.abc import AsyncGenerator +from unittest import mock from unittest.mock import AsyncMock, MagicMock import pytest @@ -198,3 +199,52 @@ async def test_empty_input_stream_finishes_without_engine_request(): assert queue.get_nowait() is STREAM_FINISHED llm._add_request.assert_not_awaited() + + +def _make_streaming_llm(num_chunks: int, has_request: bool = True): + llm = AsyncLLM.__new__(AsyncLLM) + llm._validate_streaming_input_sampling_params = MagicMock() + llm.get_supported_tasks = AsyncMock(return_value=("generate",)) + llm._run_output_handler = MagicMock() + llm._add_request = AsyncMock() + llm.model_config = MagicMock() + llm.output_processor = MagicMock() + llm.output_processor.has_request.return_value = has_request + + # The first processed request is the final (end-of-stream) request. + reqs = [MagicMock(prompt_embeds=None) for _ in range(num_chunks + 1)] + reqs[0].request_id = "r1-int" + llm.input_processor = MagicMock() + llm.input_processor.process_inputs.side_effect = reqs + return llm, reqs + + +async def _stream_chunks(llm: AsyncLLM, num_chunks: int) -> None: + async def chunks(): + for i in range(num_chunks): + yield StreamingInput(prompt=f"frame-{i}") + + with mock.patch( + "vllm.v1.engine.async_llm.extract_prompt_components", + return_value=("", None, None), + ): + queue = await llm._add_streaming_input_request( + "r1", chunks(), SamplingParams(max_tokens=4) + ) + task = queue._input_stream_task + assert task is not None + await task + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_live", [True, False]) +async def test_handle_inputs_stops_after_session_finished(session_live: bool): + """Once the session's request state is gone (aborted or errored), no + further chunks or final request are sent: they would recreate a request + state that is never cleaned up.""" + llm, reqs = _make_streaming_llm(num_chunks=2, has_request=session_live) + await _stream_chunks(llm, num_chunks=2) + + sent = [call.args[0] for call in llm._add_request.await_args_list] + # Both chunks then the final request, or only the first chunk. + assert sent == ([reqs[1], reqs[2], reqs[0]] if session_live else [reqs[1]]) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index f2142af2b3..f282c1860b 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -628,6 +628,13 @@ class AsyncLLM(EngineClient): any_added = False try: async for input_chunk in input_stream: + # Stop once the session has finished (aborted or errored). + # There is no await between this check and the ADD send, + # so any chunk already sent precedes the engine abort. + if any_added and not self.output_processor.has_request( + internal_req_id + ): + break sp = input_chunk.sampling_params if sp: self._validate_streaming_input_sampling_params(sp) @@ -663,8 +670,10 @@ class AsyncLLM(EngineClient): if not cancelled: if any_added: # Send empty final request to indicate that inputs have - # finished. Don't send if cancelled (session was aborted). - await self._add_request(final_req, None, None, 0, queue) + # finished. Don't send if cancelled or if the session + # already finished (aborted or errored). + if self.output_processor.has_request(internal_req_id): + await self._add_request(final_req, None, None, 0, queue) elif not errored: queue.put(STREAM_FINISHED) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index dccc8cba53..ef6b269618 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1618,6 +1618,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # To route aborts to the correct engine. self.reqs_in_flight: dict[str, EngineIdentity] = {} + # Open streaming-input sessions, whose aborts go to all engines. + self.streaming_reqs: set[str] = set() # Exact per-engine count of this client's unfinished requests. self.engine_inflight: Counter[EngineIdentity] = Counter() @@ -1640,6 +1642,12 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) // client_count def get_core_engine_for_request(self, request: EngineCoreRequest) -> EngineIdentity: + if request.resumable: + self.streaming_reqs.add(request.request_id) + else: + # Also removes a session on its final (end-of-stream) request. + self.streaming_reqs.discard(request.request_id) + if (engine := self.reqs_in_flight.get(request.request_id)) is not None: # This is for streaming-input session affinity. return engine @@ -1759,15 +1767,24 @@ class DPLBAsyncMPClient(DPAsyncMPClient): if not request_ids or self.resources.engine_dead: return - if len(request_ids) == 1: + if ( + len(request_ids) == 1 + and (req_id := request_ids[0]) not in self.streaming_reqs + and (engine := self.reqs_in_flight.get(req_id)) + ): # Fast-path common case. - if engine := self.reqs_in_flight.get(request_ids[0]): - await self._abort_requests(request_ids, engine) + await self._abort_requests(request_ids, engine) return by_engine = defaultdict[EngineIdentity, list[str]](list) for req_id in request_ids: - if engine := self.reqs_in_flight.get(req_id): + if req_id in self.streaming_reqs: + # A chunk in flight when the session finished may have + # re-created it on any engine. + self.streaming_reqs.discard(req_id) + for engine in self.core_engines: + by_engine[engine].append(req_id) + elif engine := self.reqs_in_flight.get(req_id): by_engine[engine].append(req_id) for engine, req_ids in by_engine.items(): await self._abort_requests(req_ids, engine) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index eb499f6cc7..c5f25e4276 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -784,6 +784,10 @@ class OutputProcessor: # if required. req_state.logprobs_processor.update_from_output(engine_core_output) + streaming_chunk_boundary = req_state.streaming_input and ( + finish_reason not in (FinishReason.ABORT, FinishReason.ERROR) + ) + # 4) Create and handle RequestOutput objects. if request_output := req_state.make_request_output( new_token_ids, @@ -794,7 +798,7 @@ class OutputProcessor: ec_transfer_params, error=request_error, ): - if req_state.streaming_input: + if streaming_chunk_boundary: request_output.finished = False if req_state.queue is not None: @@ -806,7 +810,7 @@ class OutputProcessor: # Free completed requests. if finish_reason is not None: - if req_state.streaming_input: + if streaming_chunk_boundary: if req_state.input_chunk_queue: update = req_state.input_chunk_queue.popleft() req_state.apply_streaming_update(update) @@ -814,9 +818,10 @@ class OutputProcessor: req_state.input_chunk_queue = None else: self._finish_request(req_state) - if not engine_core_output.finished: - # If req not finished in EngineCore, but Detokenizer - # detected stop string, abort needed in EngineCore. + if req_state.streaming_input or not engine_core_output.finished: + # EngineCore may still have the request: a stop string + # was detected here, or an in-flight streaming-input + # chunk may re-create the session. reqs_to_abort.append(req_id) # Track per-request stats