Repository navigation
Conversation
Part of #1692. Per-request `select_worker` was O(healthy): it recomputed the fleet-mean throughput by scanning every healthy worker's cached load, then scored all healthy workers to take the argmin. Both scale with fleet size and contribute to the routing hot path's O(workers) cost. Two changes, both confined to `least_load.rs`: 1. Cache the fleet-mean throughput off the hot path. The mean of positive per-worker `total_gen_throughput()` now lives in an `RwLock<f64>`, recomputed in `update_loads` (and `remove_worker`) over the whole load cache and read O(1) in `select_worker`. `update_loads` runs per model group and extends the shared cache, so the mean is taken after the merge to preserve fleet-wide semantics. The `default_throughput` fallback (no positive reports) is unchanged. 2. Power-of-two-choices selection. For pools with <= 2 healthy workers we still score the whole set (exact argmin); for larger pools we sample two distinct random healthy workers and pick the lower-scored, the standard near-optimal load-balancing approximation. RNG mirrors `power_of_two.rs` (`rand::rng()` + offset for a distinct second pick). Behavior preserved: empty/single-worker fast paths, healthy-only filtering, `fleet_has_loads` semantics, in-flight token credit to the chosen worker, and `increment_processed()`. The only intended behavioral change is min-over-sample instead of min-over-all (power-of-two) plus the cached mean. Tradeoff: power-of-two is an approximation — a given request may not land on the global argmin, but expected max-load is near-optimal and the in-flight credit still water-fills across dispatches within a poll interval, so steady-state balance is preserved while per-request cost drops from O(healthy) to O(1). Tests: power-of-two picks the lower of the sampled pair; healthy-only preserved; in-flight credit applied to the chosen worker; cached mean updates on `update_loads`/`remove_worker`, falls back with no positive rates, and is used to score a missing-snapshot worker; degenerate sizes (0,1,2). Existing least_load tests stay green. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthrough
ChangesLeastLoadPolicy: Cached Nominal Throughput and Power-of-Two Selection
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![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 optimizes the LeastLoadPolicy by caching the fleet-nominal throughput off the hot path and implementing a power-of-two-choices selection strategy to reduce worker selection complexity to O(1). The reviewer feedback suggests further optimizing performance by replacing the RwLock<f64> used for nominal_throughput with a lock-free AtomicU64 (storing the float as bits), which eliminates lock contention and synchronization overhead during request routing.
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.
| // No loads yet -> start at the fallback; recomputed on first poll. | ||
| nominal_throughput: RwLock::new(default_throughput), |
There was a problem hiding this comment.
Initialize the nominal_throughput as a lock-free AtomicU64 using f64::to_bits.
| // No loads yet -> start at the fallback; recomputed on first poll. | |
| nominal_throughput: RwLock::new(default_throughput), | |
| // No loads yet -> start at the fallback; recomputed on first poll. | |
| nominal_throughput: std::sync::atomic::AtomicU64::new(default_throughput.to_bits()), |
| fn nominal_throughput(&self) -> f64 { | ||
| self.nominal_throughput | ||
| .read() | ||
| .map(|t| *t) | ||
| .unwrap_or(self.default_throughput) |
There was a problem hiding this comment.
Load the cached throughput atomically using Ordering::Relaxed and convert it back to f64 with f64::from_bits. This is lock-free and extremely fast.
fn nominal_throughput(&self) -> f64 {
f64::from_bits(self.nominal_throughput.load(std::sync::atomic::Ordering::Relaxed))
}| let nominal = self.compute_nominal_throughput(&cached); | ||
| if let Ok(mut tp) = self.nominal_throughput.write() { | ||
| *tp = nominal; | ||
| } |
There was a problem hiding this comment.
Store the newly computed nominal throughput atomically using Ordering::Relaxed and f64::to_bits.
let nominal = self.compute_nominal_throughput(&cached);
self.nominal_throughput.store(nominal.to_bits(), std::sync::atomic::Ordering::Relaxed);| let nominal = self.compute_nominal_throughput(&cached); | ||
| if let Ok(mut tp) = self.nominal_throughput.write() { | ||
| *tp = nominal; | ||
| } |
There was a problem hiding this comment.
Store the newly computed nominal throughput atomically using Ordering::Relaxed and f64::to_bits.
let nominal = self.compute_nominal_throughput(&cached);
self.nominal_throughput.store(nominal.to_bits(), std::sync::atomic::Ordering::Relaxed);![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: a6b5318ba6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| let (tp_sum, tp_count) = loads | ||
| .values() | ||
| .map(|l| l.total_gen_throughput()) | ||
| .filter(|t| *t > 0.0) | ||
| .fold((0.0, 0u32), |(s, n), t| (s + t, n + 1)); |
There was a problem hiding this comment.
When a worker transitions away from Ready but is not removed, WorkerMonitor only evicts it from the watch/DP caches on the status-change path, and remove_worker_from_load_aware is only called by the removal workflow. This new cached nominal mean is computed from every entry left in cached_loads, so a NotReady worker's last throughput continues to affect scoring for healthy workers that are missing a fresh snapshot; the previous per-select calculation filtered through the current healthy worker slice, so those stale entries did not skew the fallback drain-time estimate.
Useful? React with 👍 / 👎.
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you! |
1 participant
Part of #1692 (routing hot path is O(total workers) per request).
Problem
LeastLoadPolicy::select_workerwas O(healthy) per request on two counts:nominal_throughput(mean of positivetotal_gen_throughput()) by scanning every healthy worker's cached load oneach call.
Both scale with fleet size and feed the O(workers) routing hot path #1692 targets.
Change (confined to
model_gateway/src/policies/least_load.rs)1. Cache the fleet-mean throughput off the hot path. The mean now lives in an
RwLock<f64>, recomputed inupdate_loads(andremove_worker) over the wholeload cache and read O(1) in
select_worker.update_loadsruns per modelgroup and
extends the shared cache, so the mean is taken after the merge tokeep fleet-wide semantics across groups. The
default_throughputfallback (nopositive reports) is unchanged.
2. Power-of-two-choices selection. For pools with
<= 2healthy workers westill score the whole set (exact argmin); for larger pools we sample two distinct
random healthy workers and pick the lower-scored — the standard near-optimal
load-balancing approximation. The RNG mirrors
power_of_two.rs(rand::rng()+an offset that guarantees a distinct second pick in O(1)). The existing
score(...)fn is reused unchanged.
Preserved: empty/single-worker fast paths, healthy-only filtering,
fleet_has_loadssemantics, the in-flight token credit to the chosen worker, andincrement_processed(). The only intended behavioral change is min-over-sampleinstead of min-over-all (power-of-two) plus the cached mean.
Load-balance tradeoff
Power-of-two is an approximation: a given request may not land on the global
argmin. But its expected maximum load is near-optimal (exponentially better than
random, close to full-scan), and the in-flight credit still water-fills load
across dispatches within a poll interval — so steady-state balance is preserved
while per-request cost drops from O(healthy) to O(1). The exact-scan path is kept
for tiny pools where sampling buys nothing.
Caveat
The cached mean is computed over the entire load cache rather than only the
currently-healthy subset (
update_loads/remove_workerdon't have the liveworker/health list). Stale entries are pruned by
remove_worker, so the meantracks the reporting fleet; this is a negligible, off-hot-path approximation of
the previous per-call healthy-only mean.
Tests
New (all green): power-of-two picks the lower of the sampled pair (4-worker
sampling-branch distribution + deterministic pairwise rule); healthy-only
preserved; in-flight credit applied to the chosen worker; cached mean updates on
update_loads, recomputes onremove_worker, falls back with no positive rates,and is used to score a missing-snapshot worker; degenerate sizes (0, 1, 2).
Existing
least_loadtests stay green.Verification
cargo +nightly fmt --all— cleancargo clippy -p smg --all-targets -- -D warnings— cleancargo test -p smg least_load(lib) — 19 passed, 0 failedcargo test -p smg --lib— 1085 passed, 0 failed, 4 ignoredSummary by CodeRabbit
Refactor
Tests