Repository navigation
feat(gateway): Smooth transition from regular->PD - #1445
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
Warning Rate limit exceeded
To keep reviews running without waiting, you can enable usage-based add-on for your organization. This allows additional reviews beyond the hourly cap. Account admins can enable it under billing. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughRouterManager removes snapshot-based router selection and per-request header-driven PD preference, replacing it with worker-capability-driven weighted random routing. The new logic computes PD eligibility and weights from available prefill/decode/regular worker counts, then selects routers proportionally. All RouterTrait request handlers are updated to use the simplified ChangesRouter Selection Refactor
Sequence DiagramsequenceDiagram
participant Client
participant RouterManager
participant WorkerRegistry
participant WeightLogic
participant Router
Client->>RouterManager: route_request(model_id)
RouterManager->>WorkerRegistry: get_workers(model_id or all)
WorkerRegistry-->>RouterManager: [worker list]
RouterManager->>WeightLogic: select_router_for_workers([workers])
WeightLogic->>WeightLogic: count prefill/decode/regular workers
WeightLogic->>WeightLogic: compute PD vs regular weights
WeightLogic->>WeightLogic: pick_router_by_weights(weights)
WeightLogic-->>RouterManager: selected_router_id
RouterManager->>Router: execute_request()
Router-->>Client: response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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 |
|
Hi @ekzhang, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 95706ac52b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let workers = if let Some(model) = model_id { | ||
| self.worker_registry.get_by_model(model).to_vec() | ||
| } else { | ||
| // ZERO-ALLOCATION Snapshot Iteration (Hot Path Optimization) | ||
| // Atomic load avoids heap allocations and DashMap shard locks per-request | ||
| let routers_snapshot = self.routers_snapshot.load(); | ||
| for router in routers_snapshot.iter() { | ||
| let mut score = 1.0; | ||
|
|
||
| let is_pd = router.is_pd_mode(); | ||
| if prefer_pd && is_pd { | ||
| score += 2.0; | ||
| } else if !prefer_pd && !is_pd { | ||
| score += 1.0; | ||
| } | ||
| // TODO: Once routers expose worker stats, we can evaluate: | ||
| // - Average worker priority vs priority_threshold | ||
| // - Average worker cost vs max_cost | ||
| // - Current load and health status | ||
|
|
||
| if score > best_score && is_router_valid(is_pd) { | ||
| best_score = score; | ||
| best_router = Some(Arc::clone(router)); | ||
| } | ||
| } | ||
| } | ||
| self.worker_registry.get_all() |
There was a problem hiding this comment.
When model_id is absent, this now routes by global worker weights (get_all()), which means PD routers can be selected for endpoints like cancel_response that call select_router_for_request(None). Neither HTTP nor gRPC PD routers implement cancel_response, so they fall back to the RouterTrait default 501 Not Implemented path in routers/mod.rs; in mixed regular+PD deployments, cancel calls will fail intermittently depending on random selection. The previous no-model path preferred non-PD routing, so this introduces a user-visible regression specifically for cancellation flows.
Useful? React with 👍 / 👎.
| let grpc_pd = if grpc_prefill > 0 && grpc_decode > 0 { | ||
| grpc_prefill + grpc_decode |
There was a problem hiding this comment.
PD routing weight is computed as prefill + decode once both are non-zero, but each PD request consumes one prefill and one decode worker, so effective capacity is bounded by the smaller side. In imbalanced rollouts (for example, many prefill workers but one decode worker), this formula over-allocates traffic to PD and can overload the scarce side, undermining the “smooth transition” goal of this change. Using paired capacity (for example based on min(prefill, decode)) would avoid this skew.
Useful? React with 👍 / 👎.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| let pick = ((rand::random::<f64>() * total as f64) as usize).min(total - 1); | ||
| let mut cum = 0usize; | ||
| for (weight, router_id) in &options { | ||
| cum += weight; | ||
| if pick < cum { | ||
| return self.routers.get(*router_id).map(|r| r.clone()); | ||
| } | ||
| } | ||
| None |
There was a problem hiding this comment.
🟡 Nit: If the randomly selected router ID has non-zero weight (workers exist) but the router itself isn't registered in self.routers, this returns None and falls through to the default router — even when other weighted router types are registered and would be valid.
For example: 3 gRPC-regular workers + 7 HTTP-regular workers, but only HTTP_REGULAR is registered. ~30% of requests would bypass HTTP_REGULAR and hit the default fallback instead of being routed proportionally.
In practice routers are all created at startup in from_config, so this should be rare, but a defensive approach would be to filter options to only routers present in self.routers before computing weights:
let options: Vec<(usize, &RouterId)> = [
(grpc_pd, &router_ids::GRPC_PD),
(http_pd, &router_ids::HTTP_PD),
(grpc_regular, &router_ids::GRPC_REGULAR),
(http_regular, &router_ids::HTTP_REGULAR),
]
.into_iter()
.filter(|(_, id)| self.routers.contains_key(*id))
.collect();
| let grpc_pd = if grpc_prefill > 0 && grpc_decode > 0 { | ||
| grpc_prefill + grpc_decode | ||
| } else { | ||
| None | ||
| } | ||
| 0 | ||
| }; | ||
| let http_pd = if http_prefill > 0 && http_decode > 0 { | ||
| http_prefill + http_decode | ||
| } else { | ||
| 0 | ||
| }; |
There was a problem hiding this comment.
🟡 Nit: The PD weight is prefill + decode, but the effective throughput of a PD pipeline is bottlenecked at min(prefill, decode). In an unbalanced deployment (e.g. 10 prefill + 1 decode + 6 regular), the PD weight would be 11 vs. 6 regular, routing ~65% of traffic to a pipeline bottlenecked at 1 decode worker.
Using 2 * min(prefill, decode) (or just min(prefill, decode)) would better reflect actual PD capacity relative to regular workers. For balanced deployments the two formulas are identical, so the change only affects unbalanced scaling.
Not blocking since the typical migration adds P and D workers in roughly equal numbers, but worth considering for robustness.
There was a problem hiding this comment.
hmm I think it's more like prefill + decode in our case since each of the workers is around the same size
There was a problem hiding this comment.
Clean refactor. The weighted routing logic is correct and the test provides good statistical validation. Two minor nits posted inline (neither blocking):
- 2 × 🟡 Nit: (1)
pick_router_by_weightscan returnNonewhen the randomly selected router isn't registered, even when other valid routers are available — filtering to registered routers before sampling would be more defensive. (2) PD weight usesprefill + decoderather than throughput-awaremin(prefill, decode), which could over-route in unbalanced scaling scenarios.
![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/routers/router_manager.rs`:
- Around line 1021-1025: The test currently builds workers with
BasicWorkerBuilder::new(...) and relies on its default connection mode;
explicitly set the connection mode to HTTP by calling the builder method that
sets ConnectionMode::Http (e.g., add .connection_mode(ConnectionMode::Http) into
the chain before .build()) so BasicWorkerBuilder (and the created worker) cannot
switch to gRPC if defaults change; update usages where workers are constructed
(BasicWorkerBuilder::new, .worker_type, .labels, .circuit_breaker_config,
.build) to include the explicit connection mode.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: abd36245-0cfb-4d79-b9dd-b95c98d16c31
📒 Files selected for processing (1)
model_gateway/src/routers/router_manager.rs
95706ac to
0fd0b52
Compare
This adds support for smoothly transitioning from regular routing to PD routing and back, as well as between HTTP and gRPC worker engines. Before if there was at least 1 P or D worker, all requests would be routed to PD mode, even if there were other regular workers. After this change, we weighted route requests between PD and regular workers based on the total worker count, and only send to prefill-decode disaggregated workers if min(P, D) >= 1. Motivation here is to gradually transition a production deployment from regular to PD disaggregated mode without overloading a single worker with a storm of traffic as soon as it starts up. Instead, traffic will gradually be routed to the PD workers as they start up. Also cleaned up the `x-prefer-pd` header which is only used when running in IGW mode without a model ID passed in, deprecated usage. Signed-off-by: Eric Zhang <ekzhang1@gmail.com>
0fd0b52 to
d5015ba
Compare
|
@coderabbitai help |
ChatThere are 3 ways to chat with CodeRabbit:
CodeRabbit commands
Other keywords and placeholders
CodeRabbit configuration file (
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
3 participants
Description
Problem
It is not possible to transition an SMG deployment from regular routing to PD workers without downtime, or sending a storm of traffic to the first PD worker that starts up.
Solution
This adds support for smoothly transitioning from regular routing to PD routing and back, as well as between HTTP and gRPC worker engines.
Before if there was at least 1 P or D worker, all requests would be routed to PD mode, even if there were other regular workers. After this change, we weighted route requests between PD and regular workers based on the total worker count, and only send to prefill-decode disaggregated workers if min(P, D) >= 1.
Motivation here is to gradually transition a production deployment from regular to PD disaggregated mode without overloading a single worker with a storm of traffic as soon as it starts up. Instead, traffic will gradually be routed to the PD workers as they start up.
Also cleaned up the
x-prefer-pdheader which is only used when running in IGW mode without a model ID passed in, deprecated usage.Changes
Test Plan
New unit test
weighted_routing_splits_40_pd_60_regularregisters 2 prefill + 2 decode + 6 regular workers against a model, runs 10,000 routing calls, and asserts that ~40% of requests are directed to the PD router (±5%). At n=10,000 the expected standard deviation of the ratio is ~0.005, so the ±5% band is ~10σ wide — false failure probability is negligible (~10⁻²⁴).Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Refactor
Tests