[None][fix] Bind explicit DP rank for new conversations - #16815
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
WalkthroughConversation-aware routing now assigns unbound conversations to explicit ranks when capacity allows, records their affinity, reuses it on later turns, and prevents conflicting explicit targets from overriding established bindings. ChangesConversation affinity routing
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py (1)
862-866: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrefer Python 3.10 built-in generic annotations.
Use
list[...]anddict[...]instead ofList/Dictin this new helper.Suggested change
- requests: List["RequestQueueItem"], - all_ranks_new_requests: Dict[int, List["RequestQueueItem"]], - all_ranks_num_active_requests: List[int], + requests: list["RequestQueueItem"], + all_ranks_new_requests: dict[int, list["RequestQueueItem"]], + all_ranks_num_active_requests: list[int], ... - ) -> List["RequestQueueItem"]: + ) -> list["RequestQueueItem"]:As per coding guidelines, Python 3.10+ code should prefer built-in generic types and
|.🤖 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/_torch/pyexecutor/scheduler/adp_router.py` around lines 862 - 866, Update the new helper signature around its request and rank parameters to use Python 3.10 built-in generics: replace List[...] with list[...] and Dict[...] with dict[...]. Remove any now-unused typing imports if applicable, without changing the helper’s behavior.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py`:
- Around line 862-866: Update the new helper signature around its request and
rank parameters to use Python 3.10 built-in generics: replace List[...] with
list[...] and Dict[...] with dict[...]. Remove any now-unused typing imports if
applicable, without changing the helper’s behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 91dda924-fc49-423e-ba5d-0a7a1bf75889
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.pytests/unittest/_torch/executor/test_adp_router.py
Signed-off-by: krishung5 <krish@nvidia.com>
5d67377 to
e83e196
Compare
|
Dynamo-side companion PR: ai-dynamo/dynamo#12225. It adds an opt-in Dynamo-owned initial DP-rank source while keeping TensorRT-LLM-owned placement as the default. |
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62330 [ run ] triggered by Bot. Commit: |
|
PR_Github #62330 [ run ] completed with state
|
|
@chienchunhung / @pcastonguay I see the CI fails due to some OOM/timeout/slurm issues, could you help rerun if possible? |
|
/bot run --disable-fail-fast |
|
PR_Github #62829 [ run ] triggered by Bot. Commit: |
|
PR_Github #62829 [ run ] completed with state
|
|
/bot skip --comment "Flaky tests" |
|
PR_Github #63085 [ skip ] triggered by Bot. Commit: |
|
PR_Github #63085 [ skip ] completed with state |
Signed-off-by: krishung5 <krish@nvidia.com>
4 participants
Dev Engineer Review
ConversationAwareADPRouterexplicit first-turn handling so that when a new conversation provides an explicitattention_dp_rank, the conversation→DP-rank affinity binding is recorded (only if the hinted rank has capacity) and then reused on subsequent turns._assign_new_conversation_explicit_dp_ranksand updatedroute_requests()to call it._conv_to_rankare skipped during explicit placement, so later conflicting explicitattention_dp_rankhints cannot override an existing binding.attention_dp_rankis at capacity, the request is returned for normal affinity/load-balanced routing and is not bound.conversation_idkeep the existing explicit-placement/fallback behavior but are not recorded into the conversation affinity map.QA Engineer Review
tests/unittest/_torch/executor/test_adp_router.pyunderTestConversationAwareADPRouter:test_explicit_target_dp_rank_establishes_binding: verifies an explicittarget_dp_rank/attention_dp_rankfor a new conversation creates and records a conversation→rank binding, and subsequent turns reuse it when later requests omit the explicit target.test_existing_binding_wins_over_later_explicit_target: verifies a later request with a conflicting explicit target does not override the already-recorded conversation binding.test_no_conversation_id_falls_back_and_is_not_recorded: verifies conversation-less requests are routed without being recorded in_conv_to_rank.tests/integration/test_lists/were identified.Description
ConversationAwareADPRoutercurrently honors an explicitattention_dp_rankbefore applying conversation-affinity routing. This placesthe first request on the requested rank but does not record the resulting
conversation_id -> DP rankbinding. Subsequent requests in the sameconversation can therefore be routed to a different DP rank.
This change makes explicit placement establish the conversation binding:
attention_dp_rankis routed to thatrank and records the affinity binding.
supplies a conflicting explicit rank.
conversation-affinity and load-balancing path.
behavior.
This enables external orchestrators such as Dynamo to select the initial DP
rank while allowing TensorRT-LLM to preserve that placement for later turns.
Context:
ai-dynamo/dynamo#11609 (comment)
There are no public API changes or new dependencies.
Test Coverage
Added unit coverage for:
explicit rank.
Local validation completed:
pre-commit run --files ...ruff checkruff format --checkpython3 -m compileallgit diff --checkconflicting later hints, and requests without conversation IDs.
The full pytest suite could not be collected in the local host environment
because its optional
mpi4pyandnvtxdependencies are unavailable. Theadded tests should be exercised by the TensorRT-LLM CI pipeline.
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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.