[Bugfix][Frontend] Preserve abort finish_reason for scale-out token streams (#47933)

Signed-off-by: Ting Sun <suntcrick@gmail.com>
Signed-off-by: Flora Feng <4florafeng@gmail.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
This commit is contained in:
Ting SUN
2026-10-04 02:57:26 +00:00
committed by GitHub
co-authored by Flora Feng
parent ce49174247
commit 50b404e71e
3 changed files with 76 additions and 32 deletions
@@ -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]."""
@@ -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)
@@ -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: