diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 43bcee12ff40..76842cd05031 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -128,6 +128,7 @@ class CompletionOutput: logprobs_diff (TokenLogprobs | SimpleTokenLogprobs): Logprobs of newly generated tokens. text_diff (str): Newly generated tokens. routed_experts (Optional[torch.Tensor]): Per-token pre-EPLB logical top-k MoE expert ids (Router Replay / R3), or None when not requested. + postprocess_result (Any): What the post-processor returned for this output when post-processing ran in a postproc worker; None otherwise. """ index: int text: str = "" @@ -216,6 +217,14 @@ def token_ids_diff(self) -> List[int]: def logprobs_diff(self) -> TokenLogprobs | SimpleTokenLogprobs: return self.logprobs[self._last_logprobs_len:] + @property + def postprocess_result(self) -> Any: + """The post-processor's return value stored by a postproc worker. + + None when post-processing did not run in a postproc worker. + """ + return self._postprocess_result + class GenerationResultBase: """This holds the core logic of the GenerationResult class.""" diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 5079b9b073e4..9a355bab6c4e 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -118,11 +118,11 @@ build_request_metrics_record) from tensorrt_llm.serve.postprocess_handlers import ( ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs, - ResponsesAPIPostprocArgs, chat_harmony_post_processor, + PostprocStrategy, ResponsesAPIPostprocArgs, chat_harmony_post_processor, chat_harmony_streaming_post_processor, chat_response_post_processor, chat_stream_post_processor, completion_response_post_processor, completion_stream_post_processor, responses_api_post_processor, - responses_api_streaming_post_processor) + responses_api_streaming_post_processor, select_postproc_strategy) from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, UnsupportedRenderError, mount_render_endpoints, @@ -648,6 +648,7 @@ def __init__( load_builtin_extensions() self.generator = generator self._is_visual_gen = _is_visual_gen_instance(generator) + self._postproc = select_postproc_strategy(self.postproc_worker_enabled) self._embedding_max_queue_delay = embedding_max_queue_delay self._embedding_max_queue_size = embedding_max_queue_size self.embedding_batcher: Optional[EncodeBatcher] = None @@ -1146,6 +1147,13 @@ def postproc_worker_enabled(self) -> bool: return True if self.generator.args.num_postprocess_workers > 0 else False + @functools.cached_property + def _postproc(self) -> PostprocStrategy: + # __init__ pre-caches this; computing on demand keeps instances + # built without __init__ (unit tests use OpenAIServer.__new__) + # working. + return select_postproc_strategy(self.postproc_worker_enabled) + @property def _vocab_size(self) -> Optional[int]: if self.tokenizer is not None and self.tokenizer.tokenizer is not None: @@ -1960,11 +1968,7 @@ async def _create_chat_response( await promise.aresult() if promise.error is not None: raise RuntimeError(f"Generation failed: {promise.error}") - if self.postproc_worker_enabled: - chat_response = promise.outputs[0]._postprocess_result - else: - post_processor, args = postproc_params.post_processor, postproc_params.postproc_args - chat_response = post_processor(promise, args) + chat_response = self._postproc.apply(promise, postproc_params) if disaggregated_params is not None and disaggregated_params.request_type == "context_only": chat_response.prompt_token_ids = promise.prompt_token_ids @@ -2077,22 +2081,17 @@ async def chat_stream_generator( promise: RequestOutput, postproc_params: PostprocParams) -> AsyncGenerator[str, None]: try: - if not self.postproc_worker_enabled: - post_processor, args = postproc_params.post_processor, postproc_params.postproc_args first_response = await anext(promise) raw_request.state.server_first_token_time = self._adjusted_steady_clock.now( ) - pp_results = first_response.outputs[ - 0]._postprocess_result if self.postproc_worker_enabled else post_processor( - first_response, args) + pp_results = self._postproc.apply(first_response, + postproc_params) for pp_res in pp_results: yield pp_res # Making sure we can handle the situation where there is only one response res = first_response async for res in promise: - pp_results = res.outputs[ - 0]._postprocess_result if self.postproc_worker_enabled else post_processor( - res, args) + pp_results = self._postproc.apply(res, postproc_params) for pp_res in pp_results: yield pp_res yield "data: [DONE]\n\n" @@ -2958,11 +2957,7 @@ async def completion_response( response = await promise if response.error is not None: raise RuntimeError(f"Generation failed: {response.error}") - if not self.postproc_worker_enabled: - post_processor, args = postproc_params.post_processor, postproc_params.postproc_args - pp_result = post_processor(response, args) - else: - pp_result = response.outputs[0]._postprocess_result + pp_result = self._postproc.apply(response, postproc_params) if disaggregated_params and disaggregated_params.request_type and disaggregated_params.request_type == "context_only": # Include prompt token ids for context-only requests pp_result.prompt_token_ids = response.prompt_token_ids @@ -3003,11 +2998,7 @@ async def completion_generator(promise: RequestOutput, params: Optional[PostprocParams]): try: async for output in promise: - if not self.postproc_worker_enabled: - post_processor, args = params.post_processor, params.postproc_args - pp_result = post_processor(output, args) - else: - pp_result = output.outputs[0]._postprocess_result + pp_result = self._postproc.apply(output, params) for pp_res in pp_result: yield pp_res await self._extract_metrics(output, raw_request) @@ -3182,21 +3173,16 @@ async def chat_harmony(self, request: ChatCompletionRequest, async def create_streaming_generator(promise: RequestOutput, postproc_params: PostprocParams): try: - if not self.postproc_worker_enabled: - post_processor, args = postproc_params.post_processor, postproc_params.postproc_args first_response = await anext(promise) raw_request.state.server_first_token_time = ( self._adjusted_steady_clock.now()) - pp_results = (first_response.outputs[0]._postprocess_result if - self.postproc_worker_enabled else post_processor( - first_response, args)) + pp_results = self._postproc.apply(first_response, + postproc_params) for pp_res in pp_results: yield pp_res res = first_response async for res in promise: - pp_results = (res.outputs[0]._postprocess_result - if self.postproc_worker_enabled else - post_processor(res, args)) + pp_results = self._postproc.apply(res, postproc_params) for pp_res in pp_results: yield pp_res yield "data: [DONE]\n\n" @@ -3343,11 +3329,10 @@ async def create_response( promise: RequestOutput, postproc_params: PostprocParams) -> ResponsesResponse: await promise.aresult() - if self.postproc_worker_enabled: - response = promise.outputs[0]._postprocess_result - else: - args = postproc_params.postproc_args - response = await responses_api_create_response( + args = postproc_params.postproc_args + response = await self._postproc.apply_async( + promise, + lambda: responses_api_create_response( generator=promise, request=request, sampling_params=args.sampling_params, @@ -3361,7 +3346,7 @@ async def create_response( tool_parser=args.tool_parser, num_prompt_tokens=args.num_prompt_tokens, tokenizer=self.tokenizer, - ) + )) # A postprocessing worker's result has no prompt token ids, and the # context-only handoff needs them, so take them from the promise as # the chat path does. @@ -3375,16 +3360,14 @@ async def create_response( async def create_streaming_generator(promise: RequestOutput, postproc_params: PostprocParams): - post_processor, args = postproc_params.post_processor, postproc_params.postproc_args + args = postproc_params.postproc_args streaming_processor = args.streaming_processor initial_responses = streaming_processor.get_initial_responses() for initial_response in initial_responses: yield initial_response async for res in promise: - pp_results = res.outputs[ - 0]._postprocess_result if self.postproc_worker_enabled else post_processor( - res, args) + pp_results = self._postproc.apply(res, postproc_params) for pp_res in pp_results: yield pp_res await self._extract_metrics(res, raw_request) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index b882a2b6e7f3..65c397936ef7 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -14,8 +14,10 @@ # limitations under the License. import json import time +from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, List, Literal, Optional, Tuple, Union +from typing import (Any, Awaitable, Callable, List, Literal, Optional, Tuple, + Union) from tensorrt_llm.logger import logger from tensorrt_llm.serve.responses_utils import ResponsesStreamingProcessor @@ -25,7 +27,7 @@ from .._utils import nvtx_range_debug from ..executor import (DetokenizedGenerationResultBase, GenerationResult, GenerationResultBase) -from ..executor.postproc_worker import PostprocArgs +from ..executor.postproc_worker import PostprocArgs, PostprocParams from ..executor.result import Logprob, TokenLogprobs from ..llmapi import SamplingParams from ..llmapi.disagg_utils import (get_usage_tokens_from_ctx, @@ -1106,3 +1108,53 @@ def responses_api_streaming_post_processor( args.streaming_processor.get_final_response_non_store( rsp, args.num_prompt_tokens, tokenizer=args.tokenizer)) return outputs + + +class PostprocStrategy(ABC): + """How a handler gets the post-processed result of one engine response. + + A server chooses one strategy at construction from + num_postprocess_workers. WorkerPostprocStrategy: a postproc worker + already ran the post-processor and the frontend stored its return value + on outputs[0]. InlinePostprocStrategy: run the post-processor in the + server process. + + With n > 1 the two are not equivalent in what the frontend records: in + worker mode only outputs[0] carries the result and the metrics arrive + merged, while inline mode records metrics per candidate. That difference + is known and deliberately not changed here. + """ + + @abstractmethod + def apply(self, result: GenerationResultBase, + postproc_params: Optional[PostprocParams]) -> Any: + """Return the post-processed value for one engine response.""" + + @abstractmethod + async def apply_async(self, result: GenerationResultBase, + inline: Callable[[], Awaitable[Any]]) -> Any: + """Like apply, for in-server post-processing that is a coroutine.""" + + +class WorkerPostprocStrategy(PostprocStrategy): + + def apply(self, result, postproc_params): + return result.outputs[0].postprocess_result + + async def apply_async(self, result, inline): + return result.outputs[0].postprocess_result + + +class InlinePostprocStrategy(PostprocStrategy): + + def apply(self, result, postproc_params): + return postproc_params.post_processor(result, + postproc_params.postproc_args) + + async def apply_async(self, result, inline): + return await inline() + + +def select_postproc_strategy(postproc_worker_enabled: bool) -> PostprocStrategy: + return (WorkerPostprocStrategy() + if postproc_worker_enabled else InlinePostprocStrategy()) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 550796e3fa0c..ef7bfe0af34d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -334,6 +334,7 @@ l0_cpu: - unittest/quantization/test_mode.py - unittest/scaffolding - unittest/scripts + - unittest/serve/test_postproc_strategy.py - unittest/tools - unittest/usage - unittest/utils/test_logger.py diff --git a/tests/unittest/api_stability/references/completion_output.yaml b/tests/unittest/api_stability/references/completion_output.yaml index e183ebe20d0e..0fd559a65b44 100644 --- a/tests/unittest/api_stability/references/completion_output.yaml +++ b/tests/unittest/api_stability/references/completion_output.yaml @@ -35,3 +35,6 @@ properties: routed_experts: annotation: Optional[torch.Tensor] default: inspect._empty + postprocess_result: + annotation: Any + default: inspect._empty diff --git a/tests/unittest/disaggregated/test_disagg_responses_api.py b/tests/unittest/disaggregated/test_disagg_responses_api.py index 52a8ad9dcdce..df156bccac85 100644 --- a/tests/unittest/disaggregated/test_disagg_responses_api.py +++ b/tests/unittest/disaggregated/test_disagg_responses_api.py @@ -327,7 +327,7 @@ async def nothing(*args, **kwargs): return None promise = SimpleNamespace( - outputs=[SimpleNamespace(_postprocess_result=postprocessed)], + outputs=[SimpleNamespace(postprocess_result=postprocessed)], prompt_token_ids=[5, 6, 7], aresult=nothing, ) diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 3e97cc47d4b5..16ac85f0dafe 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -80,6 +80,7 @@ def _args() -> SimpleNamespace: enable_iter_perf_stats=True, guided_decoding_backend=None, reasoning_parser=None, + num_postprocess_workers=0, max_seq_len=100, max_batch_size=8, max_num_tokens=100, diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index b7e6cf62a3d6..7200f8b7d371 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -310,7 +310,7 @@ async def __anext__(self): frames = self._steps.pop(0) return SimpleNamespace( finished=not self._steps, - outputs=[SimpleNamespace(_postprocess_result=frames)], + outputs=[SimpleNamespace(postprocess_result=frames)], ) if self._then is not None: raise self._then diff --git a/tests/unittest/llmapi/run_llm_with_postproc.py b/tests/unittest/llmapi/run_llm_with_postproc.py index d192cd9edcde..9d52bf02cd2c 100644 --- a/tests/unittest/llmapi/run_llm_with_postproc.py +++ b/tests/unittest/llmapi/run_llm_with_postproc.py @@ -56,7 +56,7 @@ async def generate_async(): _postproc_params=postproc_params, streaming=True): print(output) - outputs.append(output.outputs[0]._postprocess_result) + outputs.append(output.outputs[0].postprocess_result) asyncio.run(generate_async()) diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py index 6dbe016150c0..13849025f5b0 100644 --- a/tests/unittest/llmapi/test_executor.py +++ b/tests/unittest/llmapi/test_executor.py @@ -429,6 +429,7 @@ def test_handle_response_propagates_should_abort(): result._handle_response(output) assert result.aborted() assert result._outputs[0]._postprocess_result == "mock_sse_data" + assert result._outputs[0].postprocess_result == "mock_sse_data" def test_PostprocWorker_Output_tracing_fields(): diff --git a/tests/unittest/serve/test_postproc_strategy.py b/tests/unittest/serve/test_postproc_strategy.py new file mode 100644 index 000000000000..37047873f934 --- /dev/null +++ b/tests/unittest/serve/test_postproc_strategy.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The post-processing strategy a server chooses once at construction. + +Every handler used to decide per response whether the post-processed value +comes from a postproc worker (stored on ``outputs[0]``) or must be computed +in the server process. These tests pin the strategy objects that replaced +those decisions and the public ``CompletionOutput.postprocess_result`` +accessor the worker strategy reads. +""" + +import asyncio +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.executor.postproc_worker import PostprocParams +from tensorrt_llm.executor.result import CompletionOutput, GenerationResultBase +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.serve.openai_server import OpenAIServer +from tensorrt_llm.serve.postprocess_handlers import ( + InlinePostprocStrategy, + PostprocStrategy, + WorkerPostprocStrategy, + select_postproc_strategy, +) + +pytestmark = pytest.mark.cpu_only + + +def _result_with_stored_value(value): + """A real result object as the frontend leaves it in worker mode.""" + result = GenerationResultBase(id=1, sampling_params=SamplingParams()) + result._outputs[0]._postprocess_result = value + return result + + +def test_select_postproc_strategy_maps_the_flag(): + worker = select_postproc_strategy(True) + inline = select_postproc_strategy(False) + assert isinstance(worker, WorkerPostprocStrategy) + assert isinstance(inline, InlinePostprocStrategy) + assert isinstance(worker, PostprocStrategy) + assert isinstance(inline, PostprocStrategy) + + +def test_worker_strategy_reads_the_stored_result_and_accepts_no_params(): + # openai_completion passes postproc_params=None in worker mode. + result = _result_with_stored_value("worker-result") + assert WorkerPostprocStrategy().apply(result, None) == "worker-result" + + +def test_inline_strategy_runs_the_post_processor_exactly_once(): + result = GenerationResultBase(id=2, sampling_params=SamplingParams()) + calls = [] + + def post_processor(res, args): + calls.append((res, args)) + return "inline-result" + + params = PostprocParams(post_processor=post_processor, postproc_args="args-sentinel") + assert InlinePostprocStrategy().apply(result, params) == "inline-result" + assert calls == [(result, "args-sentinel")] + + +def test_worker_apply_async_returns_the_field_without_calling_inline(): + result = _result_with_stored_value("stored") + + def inline(): + raise AssertionError("inline must not run in worker mode") + + value = asyncio.run(WorkerPostprocStrategy().apply_async(result, inline)) + assert value == "stored" + + +def test_inline_apply_async_awaits_the_inline_coroutine_once(): + result = GenerationResultBase(id=3, sampling_params=SamplingParams()) + calls = [] + + async def inline(): + calls.append("called") + return "async-result" + + value = asyncio.run(InlinePostprocStrategy().apply_async(result, inline)) + assert value == "async-result" + assert calls == ["called"] + + +class _ServerStandIn: + """Just what choosing the strategy reads off a server. + + Reuses the real ``postproc_worker_enabled`` property so these tests pin + the production decision, not a copy of it. + """ + + postproc_worker_enabled = OpenAIServer.postproc_worker_enabled + + def __init__(self, num_postprocess_workers): + self.generator = SimpleNamespace( + args=SimpleNamespace(num_postprocess_workers=num_postprocess_workers) + ) + self._is_visual_gen = False + + +def test_server_strategy_is_chosen_from_num_postprocess_workers(): + assert isinstance(OpenAIServer._postproc.func(_ServerStandIn(0)), InlinePostprocStrategy) + assert isinstance(OpenAIServer._postproc.func(_ServerStandIn(2)), WorkerPostprocStrategy) + + +def test_server_strategy_is_the_same_object_on_every_access(): + stand_in = _ServerStandIn(2) + first = OpenAIServer._postproc.__get__(stand_in) + assert OpenAIServer._postproc.__get__(stand_in) is first + assert isinstance(first, WorkerPostprocStrategy) + + +def test_completion_output_postprocess_result_accessor(): + output = CompletionOutput(0) + assert output.postprocess_result is None + output._postprocess_result = "worker-result" + assert output.postprocess_result == "worker-result"