Select PD workers in same roleset #1409

varungup90 · 2025-08-06T23:37:05Z

Pull Request Description

[Please provide a clear and concise description of your changes here]

Related Issues

Resolves: #[Insert issue number(s)]

Important: Before submitting, please complete the description above and review the checklist below.

Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

[Bug]: Corrections to existing functionality
[CI]: Changes to build process or CI pipeline
[Docs]: Updates or additions to documentation
[API]: Modifications to aibrix's API or interface
[CLI]: Changes or additions to the Command Line Interface
[Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

PR title includes appropriate prefix(es)
Changes are clearly explained in the PR description
New and existing tests pass successfully
Code adheres to project style and best practices
Documentation updated to reflect changes (if applicable)
Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

gemini-code-assist

Summary of Changes

Hello @varungup90, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refines the pod selection logic within the pdRouter to ensure that prefill (P) and decode (D) workers are selected from the same roleset. This change is crucial for maintaining consistency and potentially improving performance or resource utilization in disaggregated P/D architectures by ensuring related tasks are handled by co-located or logically grouped resources.

Highlights

Introduce RoleSet Constant: A new constant RoleSet ("roleset-name") has been added to pkg/plugins/gateway/algorithms/pd_disaggregation.go to identify pods belonging to a specific roleset, which is crucial for the new selection logic.
Refined Decode Pod Selection: The selectDecodePod function now takes the prefillPod as an argument and filters the available decode pods to only consider those that share the same RoleSet label as the prefillPod. This ensures that prefill and decode operations occur within the same logical grouping.
Improved Prefill Request Handling: The doPrefillRequest function has been updated to return the selected prefillPod along with any error. This allows the calling Route function to use this specific prefill pod's roleset for subsequent decode pod selection, ensuring consistency.
Enhanced Logging: Added klog.InfoS statements to log the names of the selected prefill and decode pods, as well as the IP of the prefill pod upon request completion. This enhances operational visibility and aids in debugging the P/D disaggregation routing.
Error Handling for Decode Pod: A new check has been added in the Route method to explicitly handle cases where no suitable decode pod is found after filtering by roleset, returning an error to prevent unexpected behavior.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a mechanism to select decode pods from the same RoleSet as the chosen prefill pod, which is a valuable feature for ensuring affinity between prefill and decode workers. The implementation correctly refactors doPrefillRequest to return the selected prefill pod and updates selectDecodePod to filter based on the new RoleSet label.

However, there are a couple of key areas for improvement. Critically, the new logic is not covered by unit tests, and existing tests are likely to fail due to the new RoleSet label requirement. Additionally, the selectDecodePod function could be refactored for better clarity, robustness, and traceability. My review includes a critical comment about the need for updated tests and a medium-severity suggestion to refactor selectDecodePod.

pkg/plugins/gateway/algorithms/pd_disaggregation.go

Jeffwan · 2025-08-07T17:28:38Z

pkg/plugins/gateway/algorithms/pd_disaggregation.go

@@ -90,12 +91,18 @@ func (r pdRouter) Route(ctx *types.RoutingContext, readyPodList types.PodList) (
 		return "", err
 	}

-	if err = r.doPrefillRequest(ctx, prefillPods, getLLMEngine(prefillPods[0], LLMEngineIdentifier, VLLMEngine)); err != nil {
+	prefillPod, err := r.doPrefillRequest(ctx, prefillPods, getLLMEngine(prefillPods[0], LLMEngineIdentifier, VLLMEngine))


This works but may not be optimal.

It selects the prefill role without considering the load on the decode nodes. It's possible that the decode nodes in the selected RoleSet are already heavily loaded.

It filters decode nodes on every request, which is slightly inefficient—though this is a minor issue given the current scale.

For the first issue, I believe we should improve it. We can either merge this PR and address it in a follow-up, or update this one directly. WDYT?

Yes, scoring will change as it needs to take prefill and decode pod both in consideration at the same time. I will need to think off a design, so will be in follow up PR.

Filtering is in-memory so not expensive, it is ok for this PR. In follow up PR, as prefill and decode pods both are scorered together, at that time this part can be addressed to pre-built cache index on roleset name.

Sounds good. Let's merge this one first.

Signed-off-by: Varun Gupta <varungup90@gmail.com>

* Select PD workers in same roleset * nit * update ut --------- Signed-off-by: Varun Gupta <varungup90@gmail.com>

…se-0.4 branch (#1468) * Select PD workers in same roleset (#1409) * Select PD workers in same roleset * nit * update ut --------- Signed-off-by: Varun Gupta <varungup90@gmail.com> * [Bug] fix webhook config output when using make manifests (#1412) fix webhook config output when using make manifests Signed-off-by: googs1025 <googs1025@gmail.com> * [Fix] Fix vLLM NIXL-based P/D samples (#1425) Signed-off-by: Haiyang Shi <haiyang.shi@bytedance.com> Co-authored-by: Haiyang Shi <haiyang.shi@bytedance.com> * [Fix] Disable GGA in NIXL samples (#1436) [Fix] Fix NIXL samples Explicitly set UCX_TLS to let UCX not use GGA (GPU Direct) transport Signed-off-by: Haiyang Shi <haiyang.shi@bytedance.com> Co-authored-by: Haiyang Shi <haiyang.shi@bytedance.com> * Fix P/D disaggregation router to follow Nixl kv_transfer_params (#1429) - Add kv_transfer_params configuration to prefill requests and decode requests Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com> * [Bug] Corrected naming convention for AIBRIX_MODEL_GPU_PROFILE_CACHING_FLAG (#1427) Corrected naming convention for AIBRIX_MODEL_GPU_PROFILE_CACHING_FLAG Signed-off-by: Jonathon Shea <sheajonathon0@gmail.com> * [Bug] stormservice's headless service not set ownerRef (#1442) * fix: stormservice's headless service not set ownerRef Signed-off-by: dajun.cui <dajun.cui@bytedance.com> * fix: patch ut test for service sync Signed-off-by: dajun.cui <dajun.cui@bytedance.com> --------- Signed-off-by: dajun.cui <dajun.cui@bytedance.com> * [Bug] stormservice's headless service need set PublishNotReadyAddresses (#1441) * fix: stormservice's headless service need set PublishNotReadyAddresses Signed-off-by: dajun.cui <dajun.cui@bytedance.com> * fix: isServiceEqual check PublishNotReadyAddresses Signed-off-by: dajun.cui <dajun.cui@bytedance.com> --------- Signed-off-by: dajun.cui <dajun.cui@bytedance.com> --------- Signed-off-by: Varun Gupta <varungup90@gmail.com> Signed-off-by: googs1025 <googs1025@gmail.com> Signed-off-by: Haiyang Shi <haiyang.shi@bytedance.com> Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com> Signed-off-by: Jonathon Shea <sheajonathon0@gmail.com> Signed-off-by: dajun.cui <dajun.cui@bytedance.com> Co-authored-by: Varun Gupta <varungup90@gmail.com> Co-authored-by: CYJiang <86391540+googs1025@users.noreply.github.com> Co-authored-by: Haiyang Shi <dwyane.shi@gmail.com> Co-authored-by: Haiyang Shi <haiyang.shi@bytedance.com> Co-authored-by: Jonathon Shea <35823125+JonathonShea@users.noreply.github.com> Co-authored-by: cuidajun <dajun.cui@bytedance.com>

gemini-code-assist bot reviewed Aug 6, 2025

View reviewed changes

pkg/plugins/gateway/algorithms/pd_disaggregation.go Show resolved Hide resolved

pkg/plugins/gateway/algorithms/pd_disaggregation.go Outdated Show resolved Hide resolved

varungup90 assigned varungup90 and Jeffwan and unassigned varungup90 Aug 7, 2025

Jeffwan reviewed Aug 7, 2025

View reviewed changes

varungup90 added 3 commits August 7, 2025 16:29

Select PD workers in same roleset

deebf05

Signed-off-by: Varun Gupta <varungup90@gmail.com>

nit

afb45e4

Signed-off-by: Varun Gupta <varungup90@gmail.com>

update ut

bbd67e8

Signed-off-by: Varun Gupta <varungup90@gmail.com>

Jeffwan force-pushed the roleset-based-selection branch from 70c9068 to bbd67e8 Compare August 7, 2025 23:29

Jeffwan merged commit 5d33437 into vllm-project:main Aug 8, 2025
14 checks passed

Jeffwan deleted the roleset-based-selection branch August 8, 2025 02:44

Jeffwan mentioned this pull request Aug 8, 2025

Support storm service roleset and roles in AIBrix routing cache #1273

Open

Jeffwan pushed a commit to Jeffwan/aibrix that referenced this pull request Aug 18, 2025

Select PD workers in same roleset (vllm-project#1409)

3cad1d0

* Select PD workers in same roleset * nit * update ut --------- Signed-off-by: Varun Gupta <varungup90@gmail.com>

This was referenced Aug 18, 2025

Cherry picks #1409 #1412 #1425 #1436 #1429 #1427 #1442 #1441 to release-0.4 branch #1468

Merged

v0.4.1 patch release #1452

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Select PD workers in same roleset #1409

Select PD workers in same roleset #1409

Uh oh!

varungup90 commented Aug 6, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Jeffwan Aug 7, 2025

Uh oh!

varungup90 Aug 7, 2025

Uh oh!

Jeffwan Aug 7, 2025

Uh oh!

Uh oh!

Uh oh!

Select PD workers in same roleset #1409

Select PD workers in same roleset #1409

Uh oh!

Conversation

varungup90 commented Aug 6, 2025

Pull Request Description

Related Issues

Pull Request Title Format

Submission Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Jeffwan Aug 7, 2025

Choose a reason for hiding this comment

Uh oh!

varungup90 Aug 7, 2025

Choose a reason for hiding this comment

Uh oh!

Jeffwan Aug 7, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!