Repository navigation
Conversation
… failures PositionalIndexer::intern_worker asserted id < MAX_WORKERS(2048). At fleets past 2048 workers — or any long-lived gateway with enough worker churn, since interned ids are monotonic and never recycled — the assert panicked inside the per-worker KV event subscription task, and the panic was discarded at await time, so cache-aware routing silently stopped seeing those workers' KV events. - Replace the fixed 2048-slot tree_sizes Vec with segmented, doubling, lazily-allocated storage (OnceLock<Box<[AtomicUsize]>> per segment): reads stay a lock-free array index on the query hot path, worker count is unbounded (full u32 id space), memory is 8 bytes per interned id and at most 2x the high-water id. - intern_worker returns Result instead of asserting; the only error is u32 id-space exhaustion, detected via a u64 counter instead of wrapping. Interning can no longer kill a subscription task. - Surface subscription failures instead of swallowing them: panics are caught inside the task (catch_unwind) and logged with worker context when they happen, JoinErrors are logged on the stop paths, and all failure modes increment the new smg_kv_event_subscription_failures_total counter (reasons: panic, join_error, intern_failed). - Tests: segment math across the full u32 space, interning past the old cap with matches probed on both sides of the 2048 boundary, id-lifecycle (no recycling) semantics, no-alloc reset for never written segments, exhaustion edge, and concurrent interning racing across a segment boundary. Refs #1685. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
… scoring At incident scale (#1693/#1685: 4000 workers, 3k req/s) two routing hot paths do O(workers) serialized or string-keyed work on every request. LeastLoad held a std RwLock WRITE guard across the whole scoring scan, serializing all routing for a model on one lock, plus a per-request url().to_string() allocation. The guard only protected the in-flight estimate's read-scan-credit atomicity, so the load report and the since-poll counter merge into one read-mostly map entry: selection now takes a shared read guard, resolves each worker's entry once (one URL probe per worker instead of three), and credits the chosen worker with a relaxed fetch_add, so concurrent credits are never lost. Credits to workers without a load entry are dropped: they were never observable, because scores only read the estimate when a report exists and the poll that installs one resets it. cache_aware's event-driven scoring probed the indexer's URL-keyed DashMap 2-3 times per worker per request (filter + max_by_key + debug log). PositionalIndexer now exposes a process-unique non-zero instance id (epoch), and workers cache their interned u32 id per epoch in one packed AtomicU64 in the shared WorkerRuntime (same-URL replace keeps it). The scoring loop is integer-keyed: one relaxed atomic load per worker, URL interning at most once per (worker, indexer instance) via intern_worker, and the epoch makes stale ids self-invalidating when a model's indexer is dropped and recreated. The request's per-block content hashes are memoized on SelectWorkerInfo keyed by block size: PD mode selects prefill and decode with the same info, so the full-prompt XXH3 pass runs once per request instead of once per select_worker call. Verified while in there: the second full-sequence hash flagged at cache_aware.rs:1012 (hash_token_path) is on the approximate-tree path, mutually exclusive with the event-driven hash within a request and gated behind mesh hash-index population, so it is not a duplicate and stays. One deliberate routing change: the approximate-tree fallback when a matched tenant is gone or unhealthy used to herd every such request onto healthy_indices.first(); it now picks the least-loaded healthy worker (cache_aware.rs:1020,1101). Selection is otherwise unchanged: same worker for the same state, including tie-breaks (fully equal score/load/tree-size keeps the later worker), guarded by an equivalence test that keeps the previous implementation as an oracle. Tests: LeastLoad lost-update concurrency test and selection-vs-writers safety test; score_overlap equivalence and epoch-invalidation tests; request-hash memoization tests; min-load fallback tests for both tree paths. Stacked on #1706 (intern_worker returning Result<WorkerId, WorkerIdExhausted>). Refs #1693 #1685 Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
grpc
![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 introduces significant performance optimizations and concurrency improvements to the routing policies in the model gateway. Key changes include: 1) Cache-Aware Policy Optimizations: Added process-unique epoch IDs (instance_id) to PositionalIndexer to safely cache interned worker IDs directly on worker objects, avoiding redundant URL-keyed map lookups during scoring. Additionally, full-prompt content hashing is now memoized within SelectWorkerInfo to prevent redundant hashing passes across prefill and decode selections. 2) Least-Load Policy Concurrency: Refactored LeastLoadPolicy to store in-flight token counts in atomics (AtomicU64) within a single read-mostly map, allowing concurrent worker selections to execute in parallel under a shared read lock and eliminating serialization bottlenecks. 3) Anti-Herding Fallbacks: Updated fallback logic in token-tree and string-tree routing paths to select the least-loaded healthy worker instead of herding all fallback requests onto the first healthy worker. Comprehensive unit and integration tests have been added to validate the correctness, concurrency safety, and equivalence of these optimizations. No review comments were provided, so there is no additional feedback to address.
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.
|
👋 The PR description doesn't fully follow PULL_REQUEST_TEMPLATE.md:
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}
There was a problem hiding this comment.
Reviewed all 8 changed files (event_tree.rs, cache_aware.rs, least_load.rs, mod.rs, worker_selection.rs, pd_router.rs, router.rs, worker.rs). This is a clean, well-tested performance optimization.
Findings: 0 🔴 Important · 0 🟡 Nit · 0 🟣 Pre-existing
Key things verified:
-
LeastLoadPolicy lock downgrade: The read-guard +
AtomicU64::fetch_addscheme is correct. Concurrent credits are never lost (fetch_addis atomic regardless of ordering), and the deliberate relaxation (two simultaneous selections may pick the same worker before either credit lands) is bounded by concurrency — the concurrent test with 8×250 selections proves exact credit accounting. -
Epoch-based cache invalidation: The
(epoch << 32) | worker_idpacking inWorkerRuntime::kv_worker_idis correct. Epoch 0 is reserved (sentinel for "unset"), epochs start at 1, and thecached_kv_worker_idguard rejects epoch 0. The cross-instance invalidation test demonstrates correctness even under deliberate raw-id collision. -
Content hash memoization:
OnceLockonSelectWorkerInfocorrectly memoizes per-block hashes. The block-size check prevents wrong-width reuse (Cow::Ownedfallback for mismatched sizes). PD mode's dual-selection reuse is tested. -
Herding fix:
healthy_indices.first()→min_by_key(load)in both tree fallbacks is a correctness improvement, not just perf. Both paths are covered by dedicated tests. -
resolve_worker_idside-effects: Interning via the scoring path (vs. the old read-onlyworker_id()) is safe —intern_workeris idempotent, the id space is shared with the KV-event path, and the worker_to_id map is bounded by fleet size. -
score_overlapbehavioral equivalence: The reference-oracle test validates that the integer-keyed loop matches the old URL-probingmax_by_keyon cold caches, warm caches, ties, and unknown workers.
ff20ff8 to
b901a2e
Compare
|
Hi @slin1237, 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
Summary
Final PR of the wave-1 dynamo-review batch (#1693/#1685): remove the two verified per-request serialization points in routing scoring. Stacked on #1706 (base branch is
fix/kv-index-worker-cap; this diff shows only the routing changes).K4 — LeastLoad: no more write-lock-across-scan
The old policy held a
RwLockwrite guard across its entire O(workers) scan on every request — a global serialization point (a prime suspect for the incident's 1.2-core ceiling) — plus aurl().to_string()allocation per request. What the lock actually protected: atomicity of read-scan-then-credit oninflight_tokens.Now: one read-mostly map
HashMap<String, WorkerLoadState>where the in-flight token credit is anAtomicU64inside the state. Selection takes a read guard (concurrent selections run fully parallel), resolves each worker once (1 string probe instead of 3 — the old code probed two maps plus a throughput re-probe), credits viafetch_add— no exclusive lock, no allocation, no lost updates (proven by an 8-thread × 250-selection test asserting the exact credit total). Merging the two maps also closes a pre-existing torn-update window between their separate locks. Deliberate relaxation: two truly simultaneous selections may pick the same worker before either credit lands — bounded by concurrency, dynamo-style; sequential semantics unchanged.K3 — cache_aware: de-stringed scoring + memoized hashing
AtomicU64on the sharedWorkerRuntime— same-URL replace inherits it, indexer replacement self-invalidates (regression-tested with a deliberate cross-instance raw-id collision). Tie-breaking preserved exactly, guarded by an oracle test that keeps the previous implementation as the reference.compute_request_content_hashes(full-prompt XXH3) ran once per select call — twice per PD request. Now memoized onSelectWorkerInfoviaOnceLock, block-size-keyed.:1012is on the mutually-exclusive approximate-tree path, not a per-request duplicate — left untouched.Herding fix (deliberate selection change)
Both dead/unmatched-tenant fallbacks routed to
healthy_indices.first()— herding traffic onto one worker exactly when a tenant dies. Now: least-loaded healthy worker, consistent with the file's other fallbacks. Covered by tests for both tree types.Tests
8 new: lost-update concurrency (exact credit accounting), concurrent update_loads safety, the score_overlap reference-oracle equivalence (cold/warm/ties/unknown workers), cross-instance id invalidation, hash memoization (per-policy and per-block-size), and both dead-tenant fallbacks.
Gates: workspace clippy
-D warningsclean, 3565 tests passed / 0 failed (115 policy + 197 kv-index among them).Refs #1693 #1685. After #1706 squash-merges, this branch gets a trivial rebase.