[TRTLLM-7906][feat] Support multiple post process for Responses API - #9908
Conversation
|
/bot run |
|
/bot run |
📝 WalkthroughWalkthroughThis PR introduces API-level streaming and post-processing support for the Responses API, adding per-request streaming processor orchestration, new post-processing argument objects (ResponsesAPIPostprocArgs), and conditional store logic based on post-processing worker usage. It refactors streaming event handling into an object-oriented ResponsesStreamingProcessor and updates server-level request routing to support per-request post-processor selection. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Server as openai_server.py
participant Postproc as postprocess_handlers
participant Utils as responses_utils
participant Executor as LLM Executor
participant Store
Client->>Server: Request (streaming=true)
rect rgb(230, 245, 255)
Note over Server: Per-request Setup
Server->>Server: Check post_processor type
Server->>Utils: Create ResponsesStreamingProcessor
Utils-->>Server: Processor instance
Server->>Server: Build ResponsesAPIPostprocArgs
Server->>Server: Determine store enablement<br/>(consider workers + request flag)
end
Server->>Executor: Call generate with<br/>_postproc_params
rect rgb(240, 255, 240)
Note over Executor,Utils: Streaming Loop
loop For each GenerationResult
Executor-->>Utils: Yield output
Utils->>Utils: process_single_output()
Utils-->>Server: Streaming event
Server-->>Client: Send to client
end
end
rect rgb(255, 250, 240)
Note over Utils,Store: Final Response
Utils->>Utils: get_final_response_non_store()
Utils-->>Postproc: Final response data
Postproc->>Postproc: responses_api_streaming_post_processor
Postproc-->>Server: ResponsesResponse
alt store_enabled == true
Server->>Store: Persist result
Store-->>Server: Stored
else store_enabled == false or workers enabled
Note over Server: Bypass storage
end
end
Server-->>Client: Final response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45–75 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tensorrt_llm/serve/openai_server.py (2)
917-937: Minor formatting issue: missing space after=on Line 921.The code logic is correct, but there's a minor formatting inconsistency.
if self.postproc_worker_enabled: - response =promise.outputs[0]._postprocess_result + response = promise.outputs[0]._postprocess_result else:
1014-1017: Consider storing the task reference for proper lifecycle management.Static analysis flagged that
asyncio.create_taskreturns a task that isn't stored. While this pattern exists elsewhere in this file (Lines 556, 638, 785, 899) for theawait_disconnectedhelper, not keeping a reference means you cannot cancel the task if needed and any exceptions in the task will be silently lost.However, since this follows the established pattern in the codebase and the
await_disconnectedfunction handles its own lifecycle (checkingpromise.finished), this is consistent with existing code.tensorrt_llm/serve/responses_utils.py (2)
1111-1111: Use explicitOptional[int]instead of implicitNonedefault.PEP 484 prohibits implicit
Optional. The type hint should explicitly declareOptional[int].- create_time: int = None, + create_time: Optional[int] = None,
1149-1174: Use explicitOptional[int]forcreate_timeparameter.Same issue as
create_response- use explicitOptional[int]type hint.- create_time: int = None, + create_time: Optional[int] = None,
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
tensorrt_llm/serve/openai_server.py(4 hunks)tensorrt_llm/serve/postprocess_handlers.py(3 hunks)tensorrt_llm/serve/responses_utils.py(7 hunks)tests/unittest/llmapi/apps/_test_openai_responses.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces; do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used (e.g., usefrom package.subpackage import fooand thenfoo.SomeClass()instead offrom package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g.,some_file.py)
Python class names should use PascalCase (e.g.,class SomeClass)
Python function and method names should use snake_case (e.g.,def my_awesome_function():)
Python local variable names should use snake_case, with prefixkfor variable names that start with a number (e.g.,k_99th_percentile = ...)
Python global variables should use upper snake_case with prefixG(e.g.,G_MY_GLOBAL = ...)
Python constants should use upper snake_case (e.g.,MY_CONSTANT = ...)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g.,self.x = 5followed by"""<type>: Description of 'x'""")
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic
Files:
tests/unittest/llmapi/apps/_test_openai_responses.pytensorrt_llm/serve/postprocess_handlers.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/responses_utils.py
**/*.{cpp,h,cu,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top
Files:
tests/unittest/llmapi/apps/_test_openai_responses.pytensorrt_llm/serve/postprocess_handlers.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/responses_utils.py
🧠 Learnings (1)
📚 Learning: 2025-07-22T09:22:14.726Z
Learnt from: yechank-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Applied to files:
tensorrt_llm/serve/responses_utils.py
🧬 Code graph analysis (2)
tensorrt_llm/serve/postprocess_handlers.py (2)
tensorrt_llm/serve/responses_utils.py (2)
ResponsesStreamingProcessor(1708-1842)create_response_non_store(1149-1174)tensorrt_llm/serve/openai_protocol.py (4)
ResponsesRequest(759-855)ResponsesResponse(895-960)StreamOptions(76-78)ToolCall(382-386)
tensorrt_llm/serve/openai_server.py (3)
tensorrt_llm/serve/openai_protocol.py (2)
ResponsesRequest(759-855)ResponsesResponse(895-960)tensorrt_llm/serve/postprocess_handlers.py (1)
ResponsesAPIPostprocArgs(546-553)tensorrt_llm/serve/responses_utils.py (3)
ConversationHistoryStore(176-405)ResponsesStreamingProcessor(1708-1842)create_response(1102-1146)
🪛 Ruff (0.14.8)
tensorrt_llm/serve/openai_server.py
1017-1017: Store a reference to the return value of asyncio.create_task
(RUF006)
tensorrt_llm/serve/responses_utils.py
1127-1127: Avoid specifying long messages outside the exception class
(TRY003)
1155-1155: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
1840-1840: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (14)
tests/unittest/llmapi/apps/_test_openai_responses.py (2)
24-28: LGTM - New fixture for testing post-processing worker configurations.The
num_postprocess_workersfixture correctly parameterizes tests to run with both disabled (0) and enabled (2) post-processing worker configurations, providing good coverage for the new multi-postprocessing feature.
31-35: LGTM - Server fixture properly propagates the new parameter.The server fixture correctly accepts and passes the
num_postprocess_workersparameter to the server process via command-line arguments.tensorrt_llm/serve/postprocess_handlers.py (3)
4-6: LGTM - New imports for Responses API post-processing.The imports follow the coding guidelines by maintaining the namespace when importing, using
from tensorrt_llm.serve.responses_utils import ....
545-553: LGTM - Well-structured post-processing arguments dataclass.The
ResponsesAPIPostprocArgsdataclass follows the established pattern of other postproc args classes in this file, with appropriate fields for the Responses API workflow.
556-568: LGTM - Non-streaming post-processor correctly delegates to response creation.The function properly extracts fields from
argsand passes them toresponses_api_create_response_non_store.tensorrt_llm/serve/openai_server.py (4)
120-121: Verify the interaction betweenenable_storeandpostproc_worker_enabled.The
enable_storeflag now depends on both the environment variable andpostproc_worker_enabledstatus. This means storage is disabled when post-processing workers are enabled, which aligns with the warning at Lines 1014-1015. However, this creates a somewhat silent behavior change - ensure this is documented or users are aware that enabling post-processing workers disables response storage.
939-949: LGTM - Streaming generator properly yields initial responses and processes outputs.The generator correctly:
- Sends initial responses (created/in_progress events) before processing
- Handles both post-processing worker enabled and disabled paths
- Yields individual results from the processed output list
979-1006: LGTM - Proper conditional creation of streaming processor and post-processing arguments.The code correctly:
- Creates
streaming_processoronly when streaming is requested- Populates
ResponsesAPIPostprocArgswith all required fields- Selects the appropriate post-processor based on the
streamflag
1019-1026: LGTM - Return paths correctly handle streaming and non-streaming responses.Both paths properly return the expected response types.
tensorrt_llm/serve/responses_utils.py (5)
13-13: LGTM - AddedTupleto typing imports.The import is needed for the new function signatures that return tuples.
1061-1099: LGTM - Centralized response creation logic.The
_create_responsefunction cleanly abstracts the response creation logic, handling both harmony and non-harmony modes while returning output messages for optional storage.
1708-1742: LGTM - Well-structured streaming processor class.The
ResponsesStreamingProcessorclass properly encapsulates streaming state, event generation, and response finalization. The separation ofget_final_response(async, with store) andget_final_response_non_store(sync, without store) correctly handles the different use cases.
1815-1842: LGTM - Process single output correctly handles event generation.The method properly handles both harmony and non-harmony paths and formats events for streaming output.
1845-1881: LGTM - Streaming events function properly orchestrates the streaming flow.The function correctly:
- Creates a streaming processor
- Yields initial responses
- Processes each output and yields events
- Yields the final response with optional storage
|
PR_Github #27836 [ run ] triggered by Bot. Commit: |
50a5b87 to
e7a77b7
Compare
|
/bot run |
|
PR_Github #27837 [ run ] triggered by Bot. Commit: |
|
PR_Github #27836 [ run ] completed with state |
|
PR_Github #27838 [ run ] triggered by Bot. Commit: |
|
PR_Github #27837 [ run ] completed with state |
|
PR_Github #27838 [ run ] completed with state |
|
/bot run |
|
/bot kill |
e7a77b7 to
5d930be
Compare
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #28092 [ run ] triggered by Bot. Commit: |
|
PR_Github #28092 [ run ] completed with state |
6d4bfee to
302caeb
Compare
|
/bot run |
|
PR_Github #28325 [ run ] triggered by Bot. Commit: |
|
PR_Github #28325 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
/bot run |
|
/bot run |
|
PR_Github #29009 [ run ] triggered by Bot. Commit: |
|
PR_Github #29009 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29083 [ run ] triggered by Bot. Commit: |
|
PR_Github #29083 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29120 [ run ] triggered by Bot. Commit: |
|
PR_Github #29120 [ run ] completed with state
|
dda6848 to
4f30fd3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #29150 [ run ] triggered by Bot. Commit: |
|
PR_Github #29150 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29184 [ run ] triggered by Bot. Commit: |
|
PR_Github #29184 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29346 [ run ] triggered by Bot. Commit: |
Signed-off-by: Junyi Xu <219237550+JunyiXu-nv@users.noreply.github.com>
4f30fd3 to
d56f655
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #29357 [ run ] triggered by Bot. Commit: |
|
PR_Github #29357 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29438 [ run ] triggered by Bot. Commit: |
|
PR_Github #29438 [ run ] completed with state |
…VIDIA#9908) Signed-off-by: Junyi Xu <219237550+JunyiXu-nv@users.noreply.github.com>
…VIDIA#9908) Signed-off-by: Junyi Xu <219237550+JunyiXu-nv@users.noreply.github.com>
…VIDIA#9908) Signed-off-by: Junyi Xu <219237550+JunyiXu-nv@users.noreply.github.com> Signed-off-by: Daniil Kulko <kulkodaniil@gmail.com>
3 participants
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.
Description
Feature:
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.