fix(router): use independent P/D policies in gRPC PD selection - #2035
Conversation
Sharing one round_robin counter across prefill and decode made each role only see every other worker index. Signed-off-by: yangxurui <yangxurui@meituan.com> Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
github-actions
Bot
added
grpc
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change separates prefill and decode worker-selection policies in PD routing. It updates selection metrics to use each leg’s policy and adds regression coverage for independent round-robin state across worker pools. ChangesPD worker selection
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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
🤖 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.
Inline comments:
In `@model_gateway/src/policies/round_robin.rs`:
- Around line 169-216: Add an integration regression test for
WorkerSelectionStage::select_pd_pair using four prefill and four decode workers,
with independent round-robin policies registered through PolicyRegistry. Invoke
select_pd_pair 40 times and assert each worker in both pools is selected exactly
10 times; ensure the stage obtains policies via get_prefill_policy and
get_decode_policy rather than calling RoundRobinPolicy directly, then run the
pr-test-analyzer agent to verify coverage.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 774ff52c-bd54-4531-be87-1c99351ffa28
📒 Files selected for processing (2)
model_gateway/src/policies/round_robin.rsmodel_gateway/src/routers/grpc/common/stages/worker_selection.rs
| #[test] | ||
| fn test_independent_policies_cover_all_workers_across_two_pools() { | ||
| fn make_workers(prefix: &str, n: usize) -> Vec<Arc<dyn Worker>> { | ||
| (0..n) | ||
| .map(|i| { | ||
| Arc::new( | ||
| BasicWorkerBuilder::new(format!("http://{prefix}{i}:8000")) | ||
| .worker_type(WorkerType::Regular) | ||
| .health_config(no_health_check()) | ||
| .build(), | ||
| ) as Arc<dyn Worker> | ||
| }) | ||
| .collect() | ||
| } | ||
|
|
||
| let prefill_workers = make_workers("p", 4); | ||
| let decode_workers = make_workers("d", 4); | ||
| let info = SelectWorkerInfo::default(); | ||
|
|
||
| let shared = RoundRobinPolicy::new(); | ||
| let mut shared_prefill = [0usize; 4]; | ||
| let mut shared_decode = [0usize; 4]; | ||
| for _ in 0..40 { | ||
| let p = shared.select_worker(&prefill_workers, &info).unwrap(); | ||
| let d = shared.select_worker(&decode_workers, &info).unwrap(); | ||
| shared_prefill[p] += 1; | ||
| shared_decode[d] += 1; | ||
| } | ||
| assert_eq!(shared_prefill, [20, 0, 20, 0]); | ||
| assert_eq!(shared_decode, [0, 20, 0, 20]); | ||
|
|
||
| let prefill_policy = RoundRobinPolicy::new(); | ||
| let decode_policy = RoundRobinPolicy::new(); | ||
| let mut indep_prefill = [0usize; 4]; | ||
| let mut indep_decode = [0usize; 4]; | ||
| for _ in 0..40 { | ||
| let p = prefill_policy | ||
| .select_worker(&prefill_workers, &info) | ||
| .unwrap(); | ||
| let d = decode_policy | ||
| .select_worker(&decode_workers, &info) | ||
| .unwrap(); | ||
| indep_prefill[p] += 1; | ||
| indep_decode[d] += 1; | ||
| } | ||
| assert_eq!(indep_prefill, [10, 10, 10, 10]); | ||
| assert_eq!(indep_decode, [10, 10, 10, 10]); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🔴 Important Add a select_pd_pair integration regression test.
This test validates RoundRobinPolicy behavior only. It does not exercise the changed WorkerSelectionStage::select_pd_pair wiring.
Configure four prefill workers and four decode workers with round-robin policies. Call select_pd_pair 40 times. Assert that each prefill worker and each decode worker receives 10 selections. This test must use PolicyRegistry::get_prefill_policy and PolicyRegistry::get_decode_policy through the stage.
As per coding guidelines, “Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality.”
🤖 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 `@model_gateway/src/policies/round_robin.rs` around lines 169 - 216, Add an
integration regression test for WorkerSelectionStage::select_pd_pair using four
prefill and four decode workers, with independent round-robin policies
registered through PolicyRegistry. Invoke select_pd_pair 40 times and assert
each worker in both pools is selected exactly 10 times; ensure the stage obtains
policies via get_prefill_policy and get_decode_policy rather than calling
RoundRobinPolicy directly, then run the pr-test-analyzer agent to verify
coverage.
Source: Coding guidelines
Add a WorkerSelectionStage regression that selects 4P+4D with independent round_robin policies through get_prefill_policy / get_decode_policy, locking the PD wiring (not just RoundRobinPolicy). Signed-off-by: yangxurui <yangxurui@meituan.com> Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
Added an integration regression test in |
The test module cannot see the parent's private use import, which broke CI clippy/libtest compilation. Signed-off-by: yangxurui <yangxurui@meituan.com> Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Align formatting with nightly rustfmt so CI fmt --check passes. Signed-off-by: yangxurui <yangxurui@meituan.com> Co-authored-by: longcat-ia-team <312376714+longcat-ia-team@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
4 participants
Summary
select_pd_pairpreviously shared one policy instance for prefill and decode, so stateful policies likeround_robinadvanced the same counter twice per request and each role only hit every other worker.get_prefill_policy/get_decode_policy(same approach as the EPD path).Test plan
test_independent_policies_cover_all_workers_across_two_pools--prefill-policy/--decode-policy round_robin→ each of 4 prefill and 4 decode workers got equal trafficCo-authored-by: longcat-ia-team 312376714+longcat-ia-team@users.noreply.github.com
Made with Cursor