diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 145a0275dd0a..34f1e4ecc908 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -334,6 +334,11 @@ def error(self) -> Optional[str]: """Return the error message if this result completed with an error.""" return self._error_msg + @property + def finished(self) -> bool: + """Whether the whole request has finished. A read-only view of ``_done``.""" + return self._done + def _maybe_fill_spec_dec_perf_metrics( self, perf_metrics: "tllm.RequestPerfMetrics") -> None: """Backfill RequestPerfMetrics.speculative_decoding in the PyTorch flow. @@ -1145,10 +1150,6 @@ def abort(self) -> None: self._executor().abort_request(self.request_id) super().abort() - @property - def finished(self) -> bool: - return self._done - def clear_logprob_params(self) -> None: # Remove temporary attribute used in executor # for a cleaner external-facing output. diff --git a/tensorrt_llm/scaffolding/contrib/AsyncGeneration/__init__.py b/tensorrt_llm/scaffolding/contrib/AsyncGeneration/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/scaffolding/contrib/AsyncGeneration/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py b/tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py index 4e5ee87011c9..8e1f3e5f1d3c 100644 --- a/tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py +++ b/tensorrt_llm/scaffolding/contrib/AsyncGeneration/stream_generation.py @@ -43,12 +43,11 @@ async def get_step_or_more_tokens(task: StreamGenerationTask): return TaskStatus.SUCCESS for _ in range(task.streaming_step): - await task.request_handle._aresult_step() - if task.request_handle._done: + await anext(task.request_handle) + if task.request_handle.finished: break - while not task.request_handle._done: - async_task = asyncio.create_task( - task.request_handle._aresult_step()) + while not task.request_handle.finished: + async_task = asyncio.create_task(anext(task.request_handle)) if not async_task.done(): async_task.cancel() break @@ -58,7 +57,7 @@ async def get_step_or_more_tokens(task: StreamGenerationTask): task.cumulative_logprob = task.request_handle.outputs[ 0].cumulative_logprob task.logprobs = task.request_handle.outputs[0].logprobs - if task.request_handle._done: + if task.request_handle.finished: task.end_flag = True sampling_params = worker.convert_task_params(task) diff --git a/tensorrt_llm/scaffolding/worker.py b/tensorrt_llm/scaffolding/worker.py index da70945f81d9..cd24f7aee6c4 100644 --- a/tensorrt_llm/scaffolding/worker.py +++ b/tensorrt_llm/scaffolding/worker.py @@ -598,8 +598,8 @@ def convert_task_params(self, task: GenerationTask): async def streaming_generate_helper(self, generate_result, step_at_least, streaming_output_list): step = 0 - while not generate_result._done: - async_task = asyncio.create_task(generate_result._aresult_step()) + while not generate_result.finished: + async_task = asyncio.create_task(anext(generate_result)) if step_at_least and step >= step_at_least and not async_task.done( ): async_task.cancel() @@ -607,7 +607,7 @@ async def streaming_generate_helper(self, generate_result, step_at_least, await async_task step += 1 # do not put the last token to the streaming_output_list - if streaming_output_list is not None and not generate_result._done: + if streaming_output_list is not None and not generate_result.finished: streaming_output_list.append( ScaffoldingOutput( generate_result.outputs[0].text, @@ -656,7 +656,7 @@ async def stream_generation_handler( self.fill_task_with_result(task, task.request_handle) - if task.request_handle._done: + if task.request_handle.finished: task.end_flag = True return TaskStatus.SUCCESS diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index b882a2b6e7f3..e12ff9acb5c2 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -624,7 +624,7 @@ def yield_first_chat(num_tokens: int, data = chunk.model_dump_json(exclude_none=True) res.append(f"data: {data}\n\n") - if include_usage and rsp._done: + if include_usage and rsp.finished: completion_tokens = sum(output.length for output in rsp.outputs) final_usage = UsageInfo( prompt_tokens=prompt_tokens, @@ -915,7 +915,7 @@ def completion_stream_post_processor(rsp: DetokenizedGenerationResultBase, data = chunk.model_dump_json(exclude_unset=False) res.append(f"data: {data}\n\n") - if include_usage and rsp._done: + if include_usage and rsp.finished: completion_tokens = sum(output.length for output in rsp.outputs) final_usage = UsageInfo( prompt_tokens=prompt_tokens, @@ -1054,7 +1054,7 @@ def chat_harmony_streaming_post_processor( result=rsp, model=args.model, request_id=str(rsp.id), - done=rsp._done, + done=rsp.finished, num_prompt_tokens=prompt_tokens, first_iteration=args.first_iteration, stream_options=args.stream_options, @@ -1101,7 +1101,7 @@ def responses_api_streaming_post_processor( raise ValueError( "streaming_processor is required for streaming post-processing") outputs = args.streaming_processor.process_single_output(rsp) - if rsp._done: + if rsp.finished: outputs.append( args.streaming_processor.get_final_response_non_store( rsp, args.num_prompt_tokens, tokenizer=args.tokenizer)) diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index 9592b43b8b72..da3bdcad8641 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -3092,7 +3092,7 @@ def process_single_output(self, res: GenerationResult) -> list[str]: event_generator = _generate_streaming_event( output=output, request=self.request, - finished_generation=res._done, + finished_generation=res.finished, streaming_events_helper=self.streaming_events_helper, reasoning_parser_id=self.reasoning_parser, tool_parser_id=self.tool_parser, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 89ac19c97063..6727faddb512 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -295,6 +295,7 @@ l0_cpu: - unittest/quantization/test_mode.py - unittest/scaffolding - unittest/scripts + - unittest/serve/test_postprocess_finished.py - unittest/tools - unittest/usage - unittest/utils/test_logger.py diff --git a/tests/unittest/llmapi/apps/test_harmony_parsing.py b/tests/unittest/llmapi/apps/test_harmony_parsing.py index 13393bb0fbad..e53e011424c3 100644 --- a/tests/unittest/llmapi/apps/test_harmony_parsing.py +++ b/tests/unittest/llmapi/apps/test_harmony_parsing.py @@ -117,7 +117,7 @@ def _make_mock_result( result = Mock() result.outputs = [output] - result._done = True + result.finished = True result.cached_tokens = cached_tokens return result @@ -650,7 +650,7 @@ class TestStreamMetadata: def test_postprocessor_passes_cached_stream_metadata(self): mock_result = _make_mock_result(token_ids_diff=[1, 2, 3]) mock_result.id = 456 - mock_result._done = False + mock_result.finished = False args = ChatCompletionPostprocArgs( model="test-model", tools=[], diff --git a/tests/unittest/llmapi/apps/test_responses_stream_final_consistency.py b/tests/unittest/llmapi/apps/test_responses_stream_final_consistency.py index e7e2ceca84cb..15ea757d3624 100644 --- a/tests/unittest/llmapi/apps/test_responses_stream_final_consistency.py +++ b/tests/unittest/llmapi/apps/test_responses_stream_final_consistency.py @@ -81,7 +81,9 @@ def _generation_chunk(accumulated, diff, done, finish_reason="stop"): token_ids=[1, 2, 3], disaggregated_params=None, ) - return SimpleNamespace(outputs=[output], _done=done, prompt_token_ids=[1, 2], cached_tokens=0) + return SimpleNamespace( + outputs=[output], finished=done, prompt_token_ids=[1, 2], cached_tokens=0 + ) def _stream(processor, chunks): diff --git a/tests/unittest/llmapi/apps/test_responses_streaming_events.py b/tests/unittest/llmapi/apps/test_responses_streaming_events.py index b5b82e1eb59f..5bf82587ac32 100644 --- a/tests/unittest/llmapi/apps/test_responses_streaming_events.py +++ b/tests/unittest/llmapi/apps/test_responses_streaming_events.py @@ -523,7 +523,7 @@ def test_streamed_responses_use_wire_field_names(): token_ids=[1], disaggregated_params=None, ) - result = SimpleNamespace(outputs=[output], _done=True, prompt_token_ids=[1], cached_tokens=0) + result = SimpleNamespace(outputs=[output], finished=True, prompt_token_ids=[1], cached_tokens=0) frames = processor.get_initial_responses() + processor.process_single_output(result) frames.append(processor.get_final_response_non_store(result)) @@ -548,9 +548,7 @@ def _step(text, diff, done): token_ids=[1] * len(text), disaggregated_params=None, ) - return SimpleNamespace( - outputs=[output], _done=done, prompt_token_ids=[1], cached_tokens=0, finished=done - ) + return SimpleNamespace(outputs=[output], finished=done, prompt_token_ids=[1], cached_tokens=0) def _server_frames(steps, processor=None, fail=None, produced=None): @@ -573,7 +571,7 @@ async def frames(): text += diff result = _step(text, diff, done=fail is None and i == len(steps) - 1) step_frames = processor.process_single_output(result) - if result._done: + if result.finished: step_frames.append(processor.get_final_response_non_store(result)) for frame in step_frames: produced.append(frame) @@ -972,7 +970,7 @@ async def __aiter__(self): await asyncio.sleep(pause) text += diff result = _step(text, diff, done=i == len(steps) - 1) - self.finished = result._done + self.finished = result.finished yield result server.generator = SimpleNamespace( diff --git a/tests/unittest/llmapi/apps/test_responses_usage.py b/tests/unittest/llmapi/apps/test_responses_usage.py index 13c7490da625..d3a27db07f96 100644 --- a/tests/unittest/llmapi/apps/test_responses_usage.py +++ b/tests/unittest/llmapi/apps/test_responses_usage.py @@ -241,7 +241,7 @@ def test_a_streamed_response_counts_reasoning_tokens(): token_ids=[1, 2, 3], disaggregated_params=None, ) - result = SimpleNamespace(outputs=[output], _done=True, prompt_token_ids=[1], cached_tokens=0) + result = SimpleNamespace(outputs=[output], finished=True, prompt_token_ids=[1], cached_tokens=0) processor.process_single_output(result) frame = processor.get_final_response_non_store(result, tokenizer=_WordTokenizer()) usage = json.loads(frame.split("data: ", 1)[1])["response"]["usage"] diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 87b38e4053ca..3a5cfbb7a007 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -5841,7 +5841,7 @@ def _rsp(text, finish_reason): rsp = Mock() rsp.outputs = [output] - rsp._done = finish_reason is not None + rsp.finished = finish_reason is not None rsp.cached_tokens = 0 rsp.prompt_token_ids = [1, 2, 3] # Read via getattr() and fed straight into a pydantic model, so it must @@ -6370,7 +6370,7 @@ def _stream_deltas(args, chunks, finish_reason="stop"): rsp = SimpleNamespace(outputs=[output], cached_tokens=0, id="test-request", - _done=last) + finished=last) for line in chat_stream_post_processor(rsp, args): payloads.append(json.loads(line[len("data: "):])) return payloads diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py index 059c7bf27f2d..de54a9f95670 100644 --- a/tests/unittest/llmapi/test_executor.py +++ b/tests/unittest/llmapi/test_executor.py @@ -109,10 +109,16 @@ def test_GenerationResultBase(): ) result._handle_response(create_rsp(2, finished=False)) result._handle_response(create_rsp(3, finished=False)) + assert not result.finished result._handle_response(create_rsp(4, finished=True)) print(result.outputs[0]) assert len(result.outputs[0].token_ids) == 3 assert result._done + assert result.finished + # finished is a read-only view of the completion flag; only `_done` is + # ever written (by the response-handling paths), never the property. + with pytest.raises(AttributeError): + result.finished = False @pytest.mark.parametrize("pytorch", [False, True]) diff --git a/tests/unittest/scaffolding/test_worker_streaming_step.py b/tests/unittest/scaffolding/test_worker_streaming_step.py new file mode 100644 index 000000000000..9511732f41f1 --- /dev/null +++ b/tests/unittest/scaffolding/test_worker_streaming_step.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Scaffolding streaming helpers use only the public result surface. + +That surface is the ``finished`` property and async iteration (``anext``). +The fake handle below deliberately implements nothing else, so a regression +back to private members (``_done``, ``_aresult_step``) fails here. +""" + +import asyncio + +import pytest + +from tensorrt_llm.scaffolding import TaskStatus +from tensorrt_llm.scaffolding.contrib.AsyncGeneration.stream_generation import ( + StreamGenerationTask as ContribStreamGenerationTask, +) +from tensorrt_llm.scaffolding.contrib.AsyncGeneration.stream_generation import ( + stream_generation_handler, +) +from tensorrt_llm.scaffolding.task import StreamGenerationTask +from tensorrt_llm.scaffolding.worker import TRTLLMWorker + +pytestmark = pytest.mark.cpu_only + + +class _FakeOutput: + def __init__(self): + self.text = "" + self.token_ids = [] + self.cumulative_logprob = None + self.logprobs = [] + + +class _FakeStreamingHandle: + """A streaming generation handle exposing only the public surface. + + One generated token per async-iteration step; ``finished`` once + ``total_steps`` tokens are out. No ``_done``, no ``_aresult_step``. + """ + + def __init__(self, total_steps): + self.total_steps = total_steps + self.steps_taken = 0 + self.outputs = [_FakeOutput()] + self.context_logits = None + + @property + def finished(self) -> bool: + return self.steps_taken >= self.total_steps + + def __aiter__(self): + return self + + async def __anext__(self): + if self.finished: + raise StopAsyncIteration + await asyncio.sleep(0) + self.steps_taken += 1 + self.outputs[0].text += "t" + self.outputs[0].token_ids = self.outputs[0].token_ids + [self.steps_taken] + return self + + +def test_streaming_generate_helper_steps_until_finished(): + worker = TRTLLMWorker(llm=None, tokenizer=None) + handle = _FakeStreamingHandle(total_steps=3) + collected = [] + + asyncio.run(worker.streaming_generate_helper(handle, None, collected)) + + assert handle.finished + assert handle.steps_taken == 3 + # One ScaffoldingOutput per non-final step; the last token is not queued. + assert [out.text for out in collected] == ["t", "tt"] + assert [out.token_ids for out in collected] == [[1], [1, 2]] + + +def test_streaming_generate_helper_stops_early_at_step_at_least(): + worker = TRTLLMWorker(llm=None, tokenizer=None) + handle = _FakeStreamingHandle(total_steps=10) + collected = [] + + asyncio.run(worker.streaming_generate_helper(handle, 2, collected)) + + assert not handle.finished + assert handle.steps_taken == 2 + assert [out.text for out in collected] == ["t", "tt"] + + +def test_worker_stream_generation_handler_sets_end_flag_when_finished(): + worker = TRTLLMWorker(llm=None, tokenizer=None) + task = StreamGenerationTask() + task.streaming_step = 5 + task.request_handle = _FakeStreamingHandle(total_steps=2) + + status = asyncio.run(worker.stream_generation_handler(task)) + + assert status == TaskStatus.SUCCESS + assert task.end_flag + assert task.output_str == "tt" + assert task.output_tokens == [1, 2] + + +class _FakeContribWorker: + """Only what the AsyncGeneration handler calls on a worker.""" + + def convert_task_params(self, task): + return None + + +def test_contrib_stream_generation_handler_sets_end_flag_when_finished(): + task = ContribStreamGenerationTask() + task.streaming_step = 5 + task.request_handle = _FakeStreamingHandle(total_steps=3) + + status = asyncio.run(stream_generation_handler(_FakeContribWorker(), task)) + + assert status == TaskStatus.SUCCESS + assert task.end_flag + assert task.output_str == "ttt" + assert task.output_tokens == [1, 2, 3] + + +def test_contrib_stream_generation_handler_waits_for_more_steps(): + task = ContribStreamGenerationTask() + task.streaming_step = 2 + task.request_handle = _FakeStreamingHandle(total_steps=5) + + status = asyncio.run(stream_generation_handler(_FakeContribWorker(), task)) + + assert status == TaskStatus.SUCCESS + assert not task.end_flag + assert task.request_handle.steps_taken == 2 + assert task.output_str == "tt" + assert task.output_tokens == [1, 2] diff --git a/tests/unittest/serve/test_postprocess_finished.py b/tests/unittest/serve/test_postprocess_finished.py new file mode 100644 index 000000000000..390f3468f35d --- /dev/null +++ b/tests/unittest/serve/test_postprocess_finished.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Streaming post-processors read completion through ``finished``. + +``finished`` lives on ``GenerationResultBase`` so that the +``DetokenizedGenerationResultBase`` records the post-processing workers hand +to the post-processors (``num_postprocess_workers > 0``) expose it exactly +like the ``RequestOutput`` objects the in-process path uses. +""" + +import json + +import pytest + +from tensorrt_llm.executor import DetokenizedGenerationResultBase, GenerationResultBase +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.serve.openai_protocol import ResponsesRequest, StreamOptions +from tensorrt_llm.serve.postprocess_handlers import ( + ChatPostprocArgs, + CompletionPostprocArgs, + ResponsesAPIPostprocArgs, + chat_stream_post_processor, + completion_stream_post_processor, + responses_api_streaming_post_processor, +) + +pytestmark = pytest.mark.cpu_only + +RESULT_CLASSES = [GenerationResultBase, DetokenizedGenerationResultBase] + + +def _make_result(result_cls): + """A real result object mid-stream: one token out, not finished.""" + result = result_cls(1, SamplingParams()) + output = result._outputs[0] + output.text = "x" + output.token_ids = [11] + return result + + +def _finish(result): + """Advance the result by one token and mark the request finished.""" + output = result._outputs[0] + output._last_text_len = len(output.text) + output._last_token_ids_len = len(output.token_ids) + output.text += "y" + output.token_ids.append(12) + output.finish_reason = "stop" + result._done = True + + +def _payloads(chunks): + payloads = [] + for chunk in chunks: + for line in chunk.splitlines(): + if line.startswith("data: ") and not line.endswith("[DONE]"): + payloads.append(json.loads(line[len("data: ") :])) + return payloads + + +@pytest.mark.parametrize("result_cls", RESULT_CLASSES) +def test_finished_is_a_read_only_view_of_the_completion_flag(result_cls): + result = result_cls(1, SamplingParams()) + assert result.finished is False + result._done = True + assert result.finished is True + with pytest.raises(AttributeError): + result.finished = False + + +@pytest.mark.parametrize("result_cls", RESULT_CLASSES) +def test_chat_stream_emits_final_usage_only_once_finished(result_cls): + result = _make_result(result_cls) + args = ChatPostprocArgs( + role="assistant", + model="test-model", + num_prompt_tokens=3, + stream_options=StreamOptions(include_usage=True), + ) + + mid_payloads = _payloads(chat_stream_post_processor(result, args)) + assert mid_payloads, "the mid-stream call must emit delta chunks" + assert all(payload["choices"] for payload in mid_payloads), ( + "no usage-only chunk before the request finished" + ) + + _finish(result) + final_payloads = _payloads(chat_stream_post_processor(result, args)) + usage_chunk = final_payloads[-1] + assert usage_chunk["choices"] == [] + assert usage_chunk["usage"]["prompt_tokens"] == 3 + assert usage_chunk["usage"]["completion_tokens"] == 2 + assert usage_chunk["usage"]["total_tokens"] == 5 + + +@pytest.mark.parametrize("result_cls", RESULT_CLASSES) +def test_completion_stream_emits_final_usage_only_once_finished(result_cls): + result = _make_result(result_cls) + args = CompletionPostprocArgs( + model="test-model", num_prompt_tokens=3, stream_options=StreamOptions(include_usage=True) + ) + + mid_payloads = _payloads(completion_stream_post_processor(result, args)) + assert mid_payloads, "the mid-stream call must emit delta chunks" + assert all(payload["choices"] for payload in mid_payloads), ( + "no usage-only chunk before the request finished" + ) + + _finish(result) + final_payloads = _payloads(completion_stream_post_processor(result, args)) + usage_chunk = final_payloads[-1] + assert usage_chunk["choices"] == [] + assert usage_chunk["usage"]["prompt_tokens"] == 3 + assert usage_chunk["usage"]["completion_tokens"] == 2 + assert usage_chunk["usage"]["total_tokens"] == 5 + + +class _RecordingStreamingProcessor: + """Stands in for the Responses streaming processor; records the calls.""" + + def __init__(self): + self.single_output_calls = [] + self.final_response_calls = [] + + def process_single_output(self, rsp): + self.single_output_calls.append(rsp) + return ["data: step\n\n"] + + def get_final_response_non_store(self, rsp, num_prompt_tokens, tokenizer=None): + self.final_response_calls.append((rsp, num_prompt_tokens)) + return "data: final\n\n" + + +@pytest.mark.parametrize("result_cls", RESULT_CLASSES) +def test_responses_stream_requests_final_response_only_once_finished(result_cls): + processor = _RecordingStreamingProcessor() + args = ResponsesAPIPostprocArgs( + model="test-model", + request=ResponsesRequest(input="hi", model="test-model"), + sampling_params=SamplingParams(), + use_harmony=False, + streaming_processor=processor, + num_prompt_tokens=3, + ) + result = _make_result(result_cls) + + outputs = responses_api_streaming_post_processor(result, args) + assert outputs == ["data: step\n\n"] + assert processor.final_response_calls == [] + + _finish(result) + outputs = responses_api_streaming_post_processor(result, args) + assert outputs == ["data: step\n\n", "data: final\n\n"] + assert len(processor.final_response_calls) == 1 + assert processor.final_response_calls[0] == (result, 3)