[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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user