From 50b404e71e7cfaec7ac85dcf8304bfc15a5a6bb6 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Sun, 4 Oct 2026 04:57:26 +0200 Subject: [PATCH] [Bugfix][Frontend] Preserve abort finish_reason for scale-out token streams (#47933) Signed-off-by: Ting Sun Signed-off-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_generate_stream.py | 62 ++++++++++++++++--- .../token_in_token_out/test_serving_tokens.py | 18 +++--- .../scale_out/token_in_token_out/serving.py | 28 ++++----- 3 files changed, 76 insertions(+), 32 deletions(-) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 6b6d732f07..7e7f74de9f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -456,11 +456,13 @@ async def test_stream_returns_spec_decode_metrics_on_empty_terminal_output( data_chunks = [chunk for chunk in _parse_sse_chunks(chunks) if chunk != "[DONE]"] assert "metrics" not in data_chunks[0] - assert len(data_chunks) == 1 + int(include_usage) + assert data_chunks[1]["choices"][0]["finish_reason"] == "stop" + assert "metrics" not in data_chunks[1] + assert len(data_chunks) == 2 + int(include_usage) if include_usage: - assert data_chunks[1]["choices"] == [] + assert data_chunks[2]["choices"] == [] assert ( - data_chunks[1]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3 + data_chunks[2]["metrics"]["speculative_decoding"]["num_draft_tokens"] == 3 ) @@ -823,30 +825,76 @@ async def test_stream_zero_token_completion_still_delivers_prompt_metadata(): @pytest.mark.asyncio -async def test_stream_zero_token_completion_emits_no_chunk_by_default(): +@pytest.mark.parametrize("finish_reason", ["stop", "length", "abort"]) +async def test_stream_emits_terminal_output_without_new_tokens(finish_reason): + """The final output may carry no new tokens (always for an abort); the + client must still see why the choice finished.""" engine = _mock_engine() async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) yield _make_request_output( - "req-1", token_ids=[], finish_reason="stop", finished=True + "req-1", token_ids=[], finish_reason=finish_reason, finished=True ) engine.generate = MagicMock(side_effect=mock_generate) serving = _build_serving_tokens(engine) request = GenerateRequest( token_ids=[1, 2, 3], - sampling_params=SamplingParams(max_tokens=1), + sampling_params=SamplingParams(max_tokens=10), model=MODEL_NAME, stream=True, ) response = await serving.serve_tokens(request) parsed = _parse_sse_chunks([chunk async for chunk in response]) + choices = [c["choices"][0] for c in parsed[:-1]] - assert not [c for c in parsed if isinstance(c, dict) and c.get("choices")] + assert [c["token_ids"] for c in choices] == [[10], []] + assert [c["finish_reason"] for c in choices] == [None, finish_reason] assert parsed[-1] == "[DONE]" +@pytest.mark.asyncio +async def test_stream_parallel_sampling_first_output_missing_a_choice(): + """With n=2 the first output may carry only one choice; the other must + still stream, and an abort finishes both.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10], index=0) + yield _make_request_output("req-1", token_ids=[11], index=1) + aborted = _make_request_output( + "req-1", token_ids=[], finish_reason="abort", finished=True + ) + aborted.outputs += _make_request_output( + "req-1", token_ids=[], finish_reason="abort", index=1 + ).outputs + yield aborted + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10, n=2), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=True), + ) + + response = await serving.serve_tokens(request) + parsed = _parse_sse_chunks([chunk async for chunk in response]) + choices = [c["choices"][0] for c in parsed[:-2]] + + assert [(c["index"], c["token_ids"], c["finish_reason"]) for c in choices] == [ + (0, [10], None), + (1, [11], None), + (0, [], "abort"), + (1, [], "abort"), + ] + assert parsed[-2]["usage"]["completion_tokens"] == 2 + + @pytest.mark.asyncio async def test_stream_error_mid_generation(): """finish_reason='error' mid-stream yields error chunk then [DONE].""" diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 2921422f01..2ee8f76e0f 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -710,21 +710,23 @@ async def test_text_mode_rejected_when_tokens_only(client): @pytest.mark.asyncio @pytest.mark.skipif( envs.VLLM_USE_RUST_FRONTEND, - reason="output_mode is not supported by the Rust frontend", + reason="the Rust frontend does not serve /inference/v1/abort_requests", ) -async def test_text_mode_stream_delivers_abort_finish_chunk( - client, tokenizer, messages +@pytest.mark.parametrize("output_mode", ["tokens", "text"]) +async def test_stream_delivers_abort_finish_chunk( + client, tokenizer, messages, output_mode ): - """The final output after an abort has no new token IDs, so it - reaches a text stream only because text mode emits finish-only chunks.""" + """The final output after an abort has no new token IDs, but the + stream must still deliver its finish reason.""" payload = { **_text_mode_payload( _chat_prompt_token_ids(tokenizer, messages), max_tokens=900, ignore_eos=True, ), + "output_mode": output_mode, "stream": True, - "request_id": "text-abort-e2e", + "request_id": f"{output_mode}-abort-e2e", } chunks = [] @@ -736,7 +738,7 @@ async def test_text_mode_stream_delivers_abort_finish_chunk( continue chunk = json.loads(line[len("data: ") :]) chunks.append(chunk) - if not aborted and chunk["choices"] and chunk["choices"][0]["text"]: + if not aborted and chunk["choices"] and chunk["choices"][0]["token_ids"]: abort = await client.post( ABORT_ENDPOINT, json={"request_ids": [chunk["request_id"]]}, @@ -745,7 +747,7 @@ async def test_text_mode_stream_delivers_abort_finish_chunk( aborted = True assert aborted - assert {chunk["output_mode"] for chunk in chunks} == {"text"} + assert {chunk["output_mode"] for chunk in chunks} == {output_mode} final_choice = chunks[-1]["choices"][0] assert final_choice["finish_reason"] == "abort" generated = sum(len(chunk["choices"][0]["token_ids"] or []) for chunk in chunks) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index a30e519cb2..6b0195b104 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -516,11 +516,12 @@ class ServingTokens(GenerateBaseServing): request_metadata: RequestResponseMetadata, ) -> AsyncGenerator[str, None]: num_prompt_tokens = 0 - num_generated_tokens: list[int] = [] first_iteration = True prompt_token_ids: list[int] | None = None num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params + # With n > 1, the first output may not include every choice yet. + num_generated_tokens = [0] * sampling_params.n last_res: RequestOutput | None = None text_mode = request.output_mode == "text" tokenizer = self._logprobs_tokenizer(text_mode) @@ -540,7 +541,6 @@ class ServingTokens(GenerateBaseServing): if res.encoder_prompt_token_ids is not None: num_prompt_tokens += len(res.encoder_prompt_token_ids) num_cached_tokens = res.num_cached_tokens - num_generated_tokens = [0] * len(res.outputs) first_iteration = False for output in res.outputs: @@ -551,21 +551,15 @@ class ServingTokens(GenerateBaseServing): finish_reason = output.finish_reason self._raise_if_error(finish_reason, request_id) - if text_mode: - # Text held back for stop string matching is flushed - # with the final output, and the abort output has no - # new token IDs, so text or a finish reason is enough. - emit = bool( - delta_token_ids or output.text or finish_reason is not None - ) - else: - # Still emit a terminal empty chunk while prompt - # metadata is pending, so zero-token completions - # deliver it. - emit = bool(delta_token_ids) or ( - finish_reason is not None and prompt_token_ids is not None - ) - if not emit: + # Terminal outputs are always emitted so the client sees + # the finish reason, e.g. an abort, which has no new + # tokens. Text mode also emits text held back for stop + # string matching, which can arrive without token IDs. + if not ( + delta_token_ids + or finish_reason is not None + or (text_mode and output.text) + ): continue if sampling_params.logprobs is not None: