fix(cache-aware): use backend global backlog for load-aware fallback#1863
fix(cache-aware): use backend global backlog for load-aware fallback#18632JooYeon wants to merge 2 commits into
Conversation
github-actions
Bot
added
protocols
📝 WalkthroughWalkthroughAdds a worker load aggregation helper, switches cache_aware routing to use snapshot-based load resolution with local fallback, and wires cache_aware policies to receive backend load snapshots through registry and app-context setup. ChangesGlobal load snapshot adoption in cache_aware
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AppContextBuilder
participant PolicyRegistry
participant WorkerMonitor
participant CacheAwarePolicy
participant load_rx
AppContextBuilder->>WorkerMonitor: create worker monitor
AppContextBuilder->>PolicyRegistry: set_load_receiver(Some(worker_monitor.subscribe()))
PolicyRegistry->>load_rx: store backend load snapshot receiver
PolicyRegistry->>CacheAwarePolicy: inject kv_event_monitor + load_rx
CacheAwarePolicy->>load_rx: read latest snapshot
load_rx-->>CacheAwarePolicy: WorkerLoadResponse map
CacheAwarePolicy->>CacheAwarePolicy: total_running_waiting_reqs() / effective_load()
CacheAwarePolicy->>CacheAwarePolicy: select_worker / score_overlap
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
![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: ef991b5dde
ℹ️ 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".
| snapshot | ||
| .and_then(|loads| loads.get(url)) | ||
| .map(|resp| resp.total_running_waiting_reqs().max(0) as usize) | ||
| .unwrap_or(local_load) |
There was a problem hiding this comment.
When a worker has any monitor entry, this path uses the polled backend count and ignores local_load entirely. Because the load monitor only refreshes periodically, a snapshot can still say a worker is idle while this gateway has since dispatched many in-flight requests to it; until the next poll, count-spread detection and shortest-queue selection continue treating that worker as idle and can stampede it. Consider taking at least max(snapshot_count, local_load) (or another since-poll correction) instead of falling back to local load only when the snapshot entry is missing.
Useful? React with 👍 / 👎.
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/policies/registry.rs (1)
448-458: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDoc comment now stale after including
cache_awarein the load-aware set.
get_all_load_aware_policieswas updated so the nestedis_load_awarepredicate includescache_aware, soremove_worker_from_load_awarenow also iteratescache_awarepolicies (harmlessly, since itsremove_workeris a no-op per the new comment above). The doc here still says "all load-aware policies (power_of_two,least_load)" — worth updating to avoid confusing future readers about which policies flow through this path.✏️ Suggested doc update
- /// Drop a removed worker's cached load report from all load-aware policies - /// (`power_of_two`, `least_load`). + /// Drop a removed worker's cached load report from all load-aware policies + /// (`power_of_two`, `least_load`; `cache_aware` is also included but its + /// `remove_worker` is a no-op).🤖 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/registry.rs` around lines 448 - 458, Update the stale doc comment on remove_worker_from_load_aware to match the current behavior of get_all_load_aware_policies, which now includes cache_aware in addition to power_of_two and least_load. Keep the comment consistent with the policy set actually iterated by remove_worker, and mention that cache_aware is included even though its remove_worker is a no-op.
🤖 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 `@crates/protocols/src/worker.rs`:
- Line 1158: The request count calculation in the worker routing logic is
summing two i32 counters before converting to i64, which can overflow first;
update the expression in the code path using the lambda over the load value so
each counter is cast to i64 before addition. Locate the mapping that uses
num_running_reqs and num_waiting_reqs in the worker selection logic and ensure
the sum is performed in i64 to preserve correct cache-aware routing.
---
Outside diff comments:
In `@model_gateway/src/policies/registry.rs`:
- Around line 448-458: Update the stale doc comment on
remove_worker_from_load_aware to match the current behavior of
get_all_load_aware_policies, which now includes cache_aware in addition to
power_of_two and least_load. Keep the comment consistent with the policy set
actually iterated by remove_worker, and mention that cache_aware is included
even though its remove_worker is a no-op.
🪄 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: 2fdb5764-9bd4-4025-8fe2-04f6bd6fe695
📒 Files selected for processing (3)
crates/protocols/src/worker.rsmodel_gateway/src/policies/cache_aware.rsmodel_gateway/src/policies/registry.rs
|
Hi @2JooYeon, 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:
|
There was a problem hiding this comment.
💡 Codex Review
smg/model_gateway/src/policies/registry.rs
Line 255 in ef991b5
Wire receivers for label-selected cache-aware policies
When cache_aware is selected from a worker policy label while the router default is not cache_aware, UpdatePoliciesStep still creates this policy via on_worker_added(..., policy_hint) (model_gateway/src/workflow/steps/shared/update_policies.rs:107-118), but self.load_rx remains None because AppContextBuilder only calls registry.set_load_receiver inside the default-policy CacheAware branch (model_gateway/src/app_context.rs:650-665). The policy is then included in get_all_load_aware_policies, so the monitor polls, but effective_load never receives the watch channel and silently falls back to per-gateway local counters, leaving the new global-backlog routing disabled for per-model cache-aware deployments.
ℹ️ 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".
2JooYeon
force-pushed
the
fix/cache-aware-global-load
branch
from
ef991b5 to
00baf82
Compare
|
Hi @2JooYeon, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 00baf82a30
ℹ️ 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".
Signed-off-by: 2JooYeon <lewndus7749@naver.com>
Signed-off-by: 2JooYeon <lewndus7749@naver.com>
2JooYeon
force-pushed
the
fix/cache-aware-global-load
branch
from
00baf82 to
b3f4758
Compare
mergify
Bot
removed
the
needs-rebase
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)
model_gateway/src/policies/registry.rs (1)
307-317: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winInject context after factory creation so cache-aware aliases are covered.
Line 308 only matches the exact
"cache_aware"string, butPolicyFactory::create_by_nameaccepts aliases/case variants like"cacheaware". Those supported names now createCacheAwarePolicyinstances without the global load receiver/KV monitor, silently reverting to local-load routing.Proposed fix
/// Create a policy from a type string (delegates to PolicyFactory) fn create_policy_from_type(&self, policy_type: &str) -> Arc<dyn LoadBalancingPolicy> { - if policy_type == "cache_aware" { - let policy: Arc<dyn LoadBalancingPolicy> = Arc::new(CacheAwarePolicy::new()); - self.inject_policy_context(&policy); - policy - } else { - PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| { - warn!("Unknown policy type '{}', using default", policy_type); - Arc::clone(&self.default_policy) - }) - } + let policy = PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| { + warn!("Unknown policy type '{}', using default", policy_type); + Arc::clone(&self.default_policy) + }); + self.inject_policy_context(&policy); + policy }🤖 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/registry.rs` around lines 307 - 317, Update create_policy_from_type in registry.rs so context injection happens for every CacheAwarePolicy returned by PolicyFactory::create_by_name, not just the exact "cache_aware" string. The current special-case branch misses aliases and case variants like "cacheaware", which can skip inject_policy_context and leave the policy without the global load receiver/KV monitor. After resolving the policy from PolicyFactory, detect CacheAwarePolicy instances and call inject_policy_context before returning them, while keeping the default-policy fallback behavior unchanged.
🤖 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 `@model_gateway/src/policies/registry.rs`:
- Around line 307-317: Update create_policy_from_type in registry.rs so context
injection happens for every CacheAwarePolicy returned by
PolicyFactory::create_by_name, not just the exact "cache_aware" string. The
current special-case branch misses aliases and case variants like "cacheaware",
which can skip inject_policy_context and leave the policy without the global
load receiver/KV monitor. After resolving the policy from PolicyFactory, detect
CacheAwarePolicy instances and call inject_policy_context before returning them,
while keeping the default-policy fallback behavior unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: a3c83d1c-5e61-46b7-a17c-82f43ed4c24d
📒 Files selected for processing (4)
crates/protocols/src/worker.rsmodel_gateway/src/app_context.rsmodel_gateway/src/policies/cache_aware.rsmodel_gateway/src/policies/registry.rs
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates the backend-reported global request count (running + waiting) into the cache_aware load balancing policy to ensure consistent imbalance detection and shortest-queue selection across multi-gateway deployments. It introduces an effective_load helper that prefers global loads from the WorkerMonitor snapshot and falls back to local in-flight counters when unavailable. Additionally, the cache_aware policy is registered as load-aware to trigger periodic polling of /v1/loads, and the policy registry is updated to properly inject context into prefill and decode policies. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
Hi @2JooYeon, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
1 participant
Description
Problem
cache_awarestill uses per-gatewayworker.load()for count-based imbalance detection and min-load fallback.Solution
WorkerMonitorwhen available.running + waitingrequest count for cache-aware load decisions.worker.load()when a worker has no fresh snapshot.Changes
WorkerLoadResponse::total_running_waiting_reqs().CacheAwarePolicy::effective_load()to resolve global-or-local load.cache_awareto use effective load for:cache_awareas load-aware so cache-aware-only deployments poll backend load snapshots.Test Plan
Added regression tests for cache-aware effective load selection:
cache_aware_uses_injected_global_load_for_selectionworker.load()cache_aware_falls_back_to_local_load_when_worker_unmonitoredworker.load()Manually validated with a multi-gateway mock setup:
mainand this branch with:cargo build --release -p smg -p mock-servercache_awarefallback decisions instead of each gateway's localworker.load()Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Summary by CodeRabbit
New Features
Bug Fixes
Tests