[#15327][feat] Add per-request priority support to OpenAI chat/completions#15329
[#15327][feat] Add per-request priority support to OpenAI chat/completions#15329sopwg612 wants to merge 13 commits into
Conversation
…completions Expose the executor's per-request scheduling priority through the OpenAI frontend: accept 'priority' (float in [0,1], higher served first, default DEFAULT_REQUEST_PRIORITY) on ChatCompletionRequest and CompletionRequest, and forward it to generate_async in openai_chat and openai_completion. Only honored when scheduler_config.waiting_queue_policy=priority; the default FCFS waiting queue ignores it. Validated on a 4xB300 deployment at 160-way saturation: HIGH(1.0) p50 TTFT 0.947s vs LOW(0.0) 2.857s, with an equal-priority control showing no gap. Signed-off-by: Sihan Wang <sihan@deepinfra.com>
📝 WalkthroughWalkthroughTensorRT-LLM's OpenAI-compatible API now accepts an optional ChangesRequest Priority Feature
🎯 2 (Simple) | ⏱️ ~8 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/serve/openai_server.py (1)
1294-1307: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winForward
priorityon the Harmony chat path as well.
openai_chatandopenai_completionnow propagaterequest.priority, but the Harmony variant of/v1/chat/completions(chat_harmony, Line 1752 onward) still callsgenerate_asyncwithoutpriority. This makes scheduling behavior path-dependent for the same API surface.Suggested fix
promise = self.generator.generate_async( inputs=harmony_tokens, sampling_params=sampling_params, _postproc_params=postproc_params if self.postproc_worker_enabled else None, streaming=bool(request.stream), lora_request=request.lora_request, scheduling_params=scheduling_params, disaggregated_params=disaggregated_params, trace_headers=trace_headers, + priority=request.priority + if request.priority is not None else DEFAULT_REQUEST_PRIORITY, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_server.py` around lines 1294 - 1307, The Harmony chat path (chat_harmony) calls self.generator.generate_async without passing the priority, causing inconsistent scheduling; update chat_harmony to forward priority exactly like the other paths by adding the priority argument to the generate_async call (use request.priority if not None else DEFAULT_REQUEST_PRIORITY) so the call signature matches the openai_chat/openai_completion usage and scheduling is consistent across handlers.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 1294-1307: The Harmony chat path (chat_harmony) calls
self.generator.generate_async without passing the priority, causing inconsistent
scheduling; update chat_harmony to forward priority exactly like the other paths
by adding the priority argument to the generate_async call (use request.priority
if not None else DEFAULT_REQUEST_PRIORITY) so the call signature matches the
openai_chat/openai_completion usage and scheduling is consistent across
handlers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 86f2ada7-fd57-4dcc-99ed-dec6398c9bce
📒 Files selected for processing (2)
tensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/openai_server.py
Keep /v1/chat/completions behavior consistent between the standard and Harmony handlers: chat_harmony now forwards request.priority to generate_async the same way openai_chat and openai_completion do. Signed-off-by: Sihan Wang <sihan@deepinfra.com>
|
A note for the maintainer who triages this: this PR is api-compatible. It only adds an optional I don't have permission to set labels on this repo — could you please add the |
syuoni
left a comment
There was a problem hiding this comment.
LGTM, thanks for the contribution.
|
@syuoni could you launch the bot run for the blossom-ci check? |
Resolve conflicts in serve/openai_protocol.py and serve/openai_server.py: keep both the new upstream conversation_params field/kwarg and this PR's priority field/kwarg on CompletionRequest/ChatCompletionRequest and the generate_async call sites. Signed-off-by: Sihan Wang <sihan@deepinfra.com>
|
/bot run |
|
PR_Github #59356 [ run ] triggered by Bot. Commit: |
|
PR_Github #59356 [ run ] completed with state
|
|
@sopwg612 Please fix the following test:
|
…eference The api_stability test test_live_fields_are_in_yaml requires every live Pydantic field on a gated request model to be declared in trtllm_serve_api.yaml. Add the priority extension field to CompletionRequest and ChatCompletionRequest (kind: extension, status: prototype, default: null, required: false). Signed-off-by: Sihan Wang <sihan@deepinfra.com>
|
@tongyuantongyu Thanks! Fixed in |
|
/bot run |
|
PR_Github #59686 [ run ] triggered by Bot. Commit: |
|
PR_Github #59686 [ run ] completed with state
|
|
@tongyuantongyu Thanks for running CI again. The Unfortunately I can't view the failure details — both the CI report and the failure-analysis page are on NVIDIA-internal hosts I can't reach. Could you share the failing test name(s) and error output from pipeline #48118? Once I can see what's failing I'll fix it right away. If it looks like an unrelated flake, a re-run might also clear it. Thanks! |
|
/bot run |
|
PR_Github #59848 [ run ] triggered by Bot. Commit: |
|
PR_Github #59848 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59934 [ run ] triggered by Bot. Commit: |
|
PR_Github #59934 [ run ] completed with state
|
Description
Closes #15327.
The LLM API already supports per-request scheduling priority (
generate_async(..., priority=...), range[0, 1], defaultDEFAULT_REQUEST_PRIORITY), honored by the PyTorch backend when the engine runs withscheduler_config.waiting_queue_policy=priority. This PR exposes it through thetrtllm-serveOpenAI frontend:priority: floatextra param (ge=0.0, le=1.0, defaultNone) toChatCompletionRequestandCompletionRequestintensorrt_llm/serve/openai_protocol.py.generate_asyncinopenai_chatandopenai_completionintensorrt_llm/serve/openai_server.py, falling back toDEFAULT_REQUEST_PRIORITYwhen unset.Behavior is unchanged for clients that don't send the field, and under the default FCFS waiting-queue policy the value is accepted and ignored, so the change is backward-compatible. Out-of-range values are rejected with a 422 by Pydantic validation.
Validated on a 4xGB300 deployment at 160-way saturation: priority 1.0 requests saw p50 TTFT 0.947s vs 2.857s for priority 0.0, with an equal-priority control showing no gap.
Test Coverage
waiting_queue_policy=priority(numbers above). Happy to add an acceptance test modeled ontests/unittest/llmapi/apps/_test_openai_cache_salt.pyif reviewers want one.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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.
Summary by CodeRabbit