[None][feat] Serve should support AGSI middlewares#13378
Conversation
Signed-off-by: Nikita Sokolov <faucct@nebius.com>
📝 WalkthroughWalkthroughIntroduces dynamic FastAPI middleware support to the serve CLI by adding Changes
Sequence Diagram(s)sequenceDiagram
participant User as User/CLI
participant Serve as serve_main()
participant Parser as Click Parser
participant Launch as launch_server()
participant Apply as _apply_fastapi_middlewares()
participant Loader as Dynamic Importer
participant App as FastAPI App
User->>Serve: invoke with --middleware
Serve->>Parser: parse CLI arguments
Parser->>Serve: extract middleware tuple
Serve->>Launch: call with middleware list
Launch->>Apply: pass middleware specs
Apply->>Loader: import module.object path
Loader->>Apply: return middleware class/function
Apply->>App: register via add_middleware/middleware()
App->>Apply: registration complete
Apply->>Launch: middlewares registered
Launch->>Serve: server ready
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/commands/serve.py (1)
1029-1036:⚠️ Potential issue | 🟠 Major
--grpcis effectively bypassed on the visual-gen path.When
is_visual_genis true, the code goes through_serve_visual_gen()and never enforces gRPC constraints, so HTTP middleware is still applied even if--grpcis provided. That makes CLI behavior inconsistent and can silently launch HTTP when users requested gRPC.🛠️ Proposed fix
is_visual_gen = extra_visual_gen_options is not None or get_is_diffusion_model( model) if is_visual_gen: + if grpc: + raise ValueError( + "Argument '--grpc' is not supported for visual generation serve mode." + ) _serve_visual_gen() else: _serve_llm()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/commands/serve.py` around lines 1029 - 1036, The visual generation path bypasses the CLI --grpc constraint because is_visual_gen triggers _serve_visual_gen() which never enforces gRPC/middleware rules; update the control flow so that when is_visual_gen is true you still honor the --grpc flag: either make _serve_visual_gen accept and enforce the grpc flag (disable HTTP middleware or initialize gRPC handlers when grpc is true) or dispatch to a visual-gen gRPC entrypoint (e.g., _serve_visual_gen_grpc) and keep existing _serve_visual_gen for HTTP; ensure the same grpc/middleware checks used by _serve_llm are applied in the visual-gen code paths.
🧹 Nitpick comments (1)
tests/unittest/llmapi/test_config_database.py (1)
242-304: Add negative tests for invalid middleware specs.Great happy-path coverage, but this change still needs failure-mode tests for:
- invalid import path/module/object resolution
- import target that is neither middleware class nor async function
Also, since this PR only adds unit tests under
tests/unittest/, QA integration test-list updates are unnecessary.As per coding guidelines: “Coverage expectations: Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix,” and for unit-scope changes, QA list updates are unnecessary/optional.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unittest/llmapi/test_config_database.py` around lines 242 - 304, Add negative unit tests that exercise failure modes when middleware specs cannot be resolved or are of the wrong type: create tests that call _apply_fastapi_middlewares and serve_main with (1) a non-existent import path (e.g., "non.existent.module.Missing") and assert the function raises the expected ImportError/ValueError, and (2) a valid import that resolves to an object that is neither a class inheriting ASGI middleware nor an async function (e.g., a plain int or regular sync function), asserting the same rejection behavior; reference the existing helpers _apply_fastapi_middlewares and serve_main and reuse the pattern in test_apply_fastapi_middlewares_supports_class_and_function and test_serve_cli_rejects_middleware_with_grpc to mock environment and check for raised exceptions and error messages.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/commands/serve.py`:
- Around line 3-4: Add the required NVIDIA SPDX copyright/license header at the
very top of the modified module (tensorrt_llm.commands.serve / serve.py) so the
file begins with the standard NVIDIA header block including the latest
modification year and the SPDX identifier; insert the header before the existing
imports (import importlib, import inspect) ensuring the full multi-line header
matches the project's copyright template and includes the correct year of last
meaningful change.
---
Outside diff comments:
In `@tensorrt_llm/commands/serve.py`:
- Around line 1029-1036: The visual generation path bypasses the CLI --grpc
constraint because is_visual_gen triggers _serve_visual_gen() which never
enforces gRPC/middleware rules; update the control flow so that when
is_visual_gen is true you still honor the --grpc flag: either make
_serve_visual_gen accept and enforce the grpc flag (disable HTTP middleware or
initialize gRPC handlers when grpc is true) or dispatch to a visual-gen gRPC
entrypoint (e.g., _serve_visual_gen_grpc) and keep existing _serve_visual_gen
for HTTP; ensure the same grpc/middleware checks used by _serve_llm are applied
in the visual-gen code paths.
---
Nitpick comments:
In `@tests/unittest/llmapi/test_config_database.py`:
- Around line 242-304: Add negative unit tests that exercise failure modes when
middleware specs cannot be resolved or are of the wrong type: create tests that
call _apply_fastapi_middlewares and serve_main with (1) a non-existent import
path (e.g., "non.existent.module.Missing") and assert the function raises the
expected ImportError/ValueError, and (2) a valid import that resolves to an
object that is neither a class inheriting ASGI middleware nor an async function
(e.g., a plain int or regular sync function), asserting the same rejection
behavior; reference the existing helpers _apply_fastapi_middlewares and
serve_main and reuse the pattern in
test_apply_fastapi_middlewares_supports_class_and_function and
test_serve_cli_rejects_middleware_with_grpc to mock environment and check for
raised exceptions and error messages.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 87f26fbd-35a9-4d3f-95f9-1d5f141cbe6c
📒 Files selected for processing (2)
tensorrt_llm/commands/serve.pytests/unittest/llmapi/test_config_database.py
svc-trtllm-gh-bot
added
the
Community want to contribute
|
/bot run |
|
PR_Github #46178 [ run ] triggered by Bot. Commit: |
venkywonka
left a comment
venkywonka
left a comment
There was a problem hiding this comment.
Lgtm! It would be great if a line about this can documented somewhere in the docs on serving!
|
PR_Github #46178 [ run ] completed with state
|
pcastonguay
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #46377 [ run ] triggered by Bot. Commit: |
|
PR_Github #46377 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #46693 [ run ] triggered by Bot. Commit: |
|
PR_Github #46693 [ run ] completed with state |
Signed-off-by: Nikita Sokolov <faucct@nebius.com>
6 participants
Summary by CodeRabbit
Release Notes
--middlewareoption to serve CLI for registering custom middleware with HTTP serversDescription
serveshould allow configuring AGSI middlewares, like vLLM does – https://github.com/vllm-project/vllm/pull/1106/changesTest Coverage
added a new test, but failed to run those because of environment issues. will try again later.
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
To see a list of available CI bot commands, please comment
/bot help.