[Bugfix] Terminate streaming sessions on ABORT/ERROR and prevent zombie sessions (#58942)

Signed-off-by: David Cheung <d7cheung@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Benedikt Falk <bfalk@nvidia.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

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)],

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):

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]])

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)
 

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)

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