Repository navigation
feat(policy): add event-driven cache-aware routing to CacheAwarePolicy - #571
Merged
Merged
Conversation
github-actions
Bot
added
python-bindings
|
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:
📝 WalkthroughWalkthroughAdds a Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router as CacheAwarePolicy
participant Monitor as KvEventMonitor
participant Indexer as PositionalIndexer
participant Workers as WorkerCluster
Client->>Router: route_request(model_id, content)
alt event-driven indexer exists
Router->>Monitor: has_event_indexer(model_id)?
Monitor-->>Router: indexer handle
Router->>Indexer: compute_hashes_and_positions(content)
Indexer-->>Router: matching positions
Router->>Router: score_overlap(tokens, block_size)
Router->>Router: choose best worker (tie-break: load/tree)
Router-->>Workers: dispatch to selected worker
else no indexer
Router->>Router: select_worker_token_or_string_tree()
Router-->>Workers: dispatch using traditional routing
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a significant feature: an event-driven, cache-aware routing policy. While the changes are well-structured and include comprehensive test coverage, a critical security issue was identified: the new block_size parameter in the configuration validation logic is not validated to be non-zero. If block_size is configured as 0, it can trigger a panic during request processing, leading to a Denial of Service. It is recommended to add a validation check to ensure block_size is always greater than 0. Furthermore, there are opportunities to enhance code readability by using more idiomatic Rust and improve the configurability of the new policy by exposing block_size as a command-line argument.
![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: dbc556b4ae
ℹ️ 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".
Comment on lines
+451
to
+452
| if self.has_event_indexer(model_id) { | ||
| self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id) |
There was a problem hiding this comment.
This lookup uses the normalized model key from routing, but KvEventMonitor indexes by raw worker.model_id(); for the backward-compat case where workers have an empty model id, routing normalizes to unknown while the monitor key stays empty, so has_event_indexer misses and event-driven routing is skipped even when KV events exist for that model.
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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
model_gateway/src/config/builder.rs (1)
111-126: 🧹 Nitpick | 🔵 TrivialUse a shared default source for
block_sizeinstead of a literal.Line 125 introduces another standalone
16; centralizing this default avoids future drift between config paths.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/config/builder.rs` around lines 111 - 126, The CacheAware policy builder uses a hardcoded block_size value (16) in cache_aware_policy, which duplicates a literal; replace that literal with a shared default constant to avoid drift: add or reuse a module-level constant (e.g., DEFAULT_BLOCK_SIZE) and set PolicyConfig::CacheAware.block_size to that constant in cache_aware_policy (and update any other places that currently use the literal 16 to use the same constant).bindings/python/src/lib.rs (1)
474-481:⚠️ Potential issue | 🟠 MajorExpose CacheAware
block_sizein Python instead of hard-coding 16.Line 480 hard-codes
block_size, so Python users cannot align routing with non-default backend KV block sizes.Proposed fix
struct Router { @@ + cache_block_size: usize, @@ } #[pymethods] impl Router { #[new] #[pyo3(signature = ( @@ + cache_block_size = 16, @@ ))] fn new( @@ + cache_block_size: usize, @@ ) -> PyResult<Self> { @@ Ok(Router { @@ + cache_block_size, @@ }) } } impl Router { @@ PolicyType::CacheAware => ConfigPolicyConfig::CacheAware { cache_threshold: self.cache_threshold, balance_abs_threshold: self.balance_abs_threshold, balance_rel_threshold: self.balance_rel_threshold, eviction_interval_secs: self.eviction_interval_secs, max_tree_size: self.max_tree_size, - block_size: 16, + block_size: self.cache_block_size.max(1), },🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@bindings/python/src/lib.rs` around lines 474 - 481, The CacheAware branch hard-codes block_size = 16; change it to use the value coming from the Rust-side Python-config struct (e.g., use self.block_size) and update the Python-facing config/constructor so block_size is exposed and populated from Python; specifically replace the literal 16 in PolicyType::CacheAware -> ConfigPolicyConfig::CacheAware with the corresponding field (block_size) on the config object and ensure the Python binding that builds this Rust config accepts and forwards a block_size parameter.model_gateway/src/main.rs (1)
751-758:⚠️ Potential issue | 🟠 MajorCLI parsing hard-codes CacheAware
block_size, making backend alignment impossible.At Line 757,
block_sizeis fixed to 16. Add a CLI option and pass it through so routing can match backend KV block size.Proposed fix
struct CliArgs { @@ + /// KV block size (tokens per block) for cache-aware event-driven routing + #[arg(long, default_value_t = 16, help_heading = "Routing Policy")] + cache_block_size: usize, @@ } fn parse_policy(&self, policy_str: &str) -> PolicyConfig { @@ "cache_aware" => PolicyConfig::CacheAware { cache_threshold: self.cache_threshold, balance_abs_threshold: self.balance_abs_threshold, balance_rel_threshold: self.balance_rel_threshold, eviction_interval_secs: self.eviction_interval, max_tree_size: self.max_tree_size, - block_size: 16, + block_size: self.cache_block_size.max(1), },🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/main.rs` around lines 751 - 758, The CLI currently hard-codes CacheAware block_size to 16; add a CLI flag (e.g., block_size) to the options struct used for argument parsing (the same struct that exposes cache_threshold, balance_abs_threshold, balance_rel_threshold, eviction_interval, max_tree_size), wire that flag into the runtime config struct on startup, and replace the literal 16 in PolicyConfig::CacheAware construction with the new field (self.block_size) so the chosen block size is passed through to routing/backend alignment; update any help text/validation (positive integer) for the new flag.model_gateway/src/config/validation.rs (1)
170-209:⚠️ Potential issue | 🟠 MajorAdd explicit validation for
CacheAware.block_size.Line 176 currently ignores
block_sizevia.., soblock_size = 0is accepted and can silently disable block-based overlap routing.Suggested fix
PolicyConfig::CacheAware { cache_threshold, balance_abs_threshold: _, balance_rel_threshold, eviction_interval_secs, max_tree_size, - .. + block_size, } => { if !(0.0..=1.0).contains(cache_threshold) { return Err(ConfigError::InvalidValue { field: "cache_threshold".to_string(), value: cache_threshold.to_string(), reason: "Must be between 0.0 and 1.0".to_string(), }); } if *balance_rel_threshold < 1.0 { return Err(ConfigError::InvalidValue { field: "balance_rel_threshold".to_string(), value: balance_rel_threshold.to_string(), reason: "Must be >= 1.0".to_string(), }); } if *eviction_interval_secs == 0 { return Err(ConfigError::InvalidValue { field: "eviction_interval_secs".to_string(), value: eviction_interval_secs.to_string(), reason: "Must be > 0".to_string(), }); } if *max_tree_size == 0 { return Err(ConfigError::InvalidValue { field: "max_tree_size".to_string(), value: max_tree_size.to_string(), reason: "Must be > 0".to_string(), }); } + + if *block_size == 0 { + return Err(ConfigError::InvalidValue { + field: "block_size".to_string(), + value: block_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/config/validation.rs` around lines 170 - 209, PolicyConfig::CacheAware currently ignores the block_size field (matched via ..) so a block_size of 0 is accepted; update the match arm for PolicyConfig::CacheAware to destructure block_size and add a validation that it is > 0 (similar to max_tree_size/eviction_interval_secs), returning ConfigError::InvalidValue with field "block_size" and an appropriate reason/value if the check fails so block-based overlap routing cannot be silently disabled.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/policies/cache_aware.rs`:
- Around line 451-453: Replace the two-step presence-then-use pattern with a
single atomic fetch of the indexer: call the getter that returns the indexer
(the same method used at the later check, e.g., get_event_indexer or similar)
once into a local variable, check that Option for Some and use that variable
when calling select_worker_event_driven instead of calling has_event_indexer
first; do the same refactor for the other occurrence around lines 510-511 so
both checks/uses use the same fetched indexer instance to avoid the TOCTOU race.
- Around line 551-574: The selection currently considers healthy_indices with
implicit score 0, so change the logic to first filter healthy_indices to only
include indices whose
overlap.scores.get(workers[idx].url()).copied().unwrap_or(0) > 0, return None
(or short-circuit to the no-overlap fallback) if that filtered list is empty,
and then run the existing max_by tie-breaker on that filtered iterator (use the
same score/load/tree_size comparisons). Reference symbols: healthy_indices,
overlap.scores, workers[..].url(), best_idx, and the max_by closure.
---
Outside diff comments:
In `@bindings/python/src/lib.rs`:
- Around line 474-481: The CacheAware branch hard-codes block_size = 16; change
it to use the value coming from the Rust-side Python-config struct (e.g., use
self.block_size) and update the Python-facing config/constructor so block_size
is exposed and populated from Python; specifically replace the literal 16 in
PolicyType::CacheAware -> ConfigPolicyConfig::CacheAware with the corresponding
field (block_size) on the config object and ensure the Python binding that
builds this Rust config accepts and forwards a block_size parameter.
In `@model_gateway/src/config/builder.rs`:
- Around line 111-126: The CacheAware policy builder uses a hardcoded block_size
value (16) in cache_aware_policy, which duplicates a literal; replace that
literal with a shared default constant to avoid drift: add or reuse a
module-level constant (e.g., DEFAULT_BLOCK_SIZE) and set
PolicyConfig::CacheAware.block_size to that constant in cache_aware_policy (and
update any other places that currently use the literal 16 to use the same
constant).
In `@model_gateway/src/config/validation.rs`:
- Around line 170-209: PolicyConfig::CacheAware currently ignores the block_size
field (matched via ..) so a block_size of 0 is accepted; update the match arm
for PolicyConfig::CacheAware to destructure block_size and add a validation that
it is > 0 (similar to max_tree_size/eviction_interval_secs), returning
ConfigError::InvalidValue with field "block_size" and an appropriate
reason/value if the check fails so block-based overlap routing cannot be
silently disabled.
In `@model_gateway/src/main.rs`:
- Around line 751-758: The CLI currently hard-codes CacheAware block_size to 16;
add a CLI flag (e.g., block_size) to the options struct used for argument
parsing (the same struct that exposes cache_threshold, balance_abs_threshold,
balance_rel_threshold, eviction_interval, max_tree_size), wire that flag into
the runtime config struct on startup, and replace the literal 16 in
PolicyConfig::CacheAware construction with the new field (self.block_size) so
the chosen block size is passed through to routing/backend alignment; update any
help text/validation (positive integer) for the new flag.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (11)
bindings/python/src/lib.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/config/validation.rsmodel_gateway/src/core/kv_event_monitor.rsmodel_gateway/src/main.rsmodel_gateway/src/policies/cache_aware.rsmodel_gateway/src/policies/factory.rsmodel_gateway/src/policies/mod.rsmodel_gateway/tests/routing/cache_aware_backward_compat_test.rsmodel_gateway/tests/routing/test_pd_routing.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 879be90684
ℹ️ 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".
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/core/kv_event_monitor.rs (1)
169-190:⚠️ Potential issue | 🟠 Major
stop()can leave stale model state when no subscriptions are active.Line 174 returns early before
indexers.clear()/block_sizes.clear(). If state exists without active handles (e.g., seeded viaset_block_size),stop()won’t actually clean up.Suggested fix
pub async fn stop(&self) { let subscriptions: HashMap<String, WorkerSubscription> = { let mut handles = self.worker_handles.lock().await; std::mem::take(&mut *handles) }; - if subscriptions.is_empty() { - return; - } - - info!( - count = subscriptions.len(), - "Stopping all KV event subscriptions" - ); - for (url, sub) in subscriptions { - debug!(worker_url = %url, "Aborting KV event subscription"); - sub.handle.abort(); - let _ = sub.handle.await; + if !subscriptions.is_empty() { + info!( + count = subscriptions.len(), + "Stopping all KV event subscriptions" + ); + for (url, sub) in subscriptions { + debug!(worker_url = %url, "Aborting KV event subscription"); + sub.handle.abort(); + let _ = sub.handle.await; + } } self.indexers.clear(); self.block_sizes.clear(); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/kv_event_monitor.rs` around lines 169 - 190, The stop() implementation returns early when subscriptions.is_empty(), which skips clearing in-memory state and can leave stale entries set via methods like set_block_size; modify stop() (the method that reads self.worker_handles) so it always calls self.indexers.clear() and self.block_sizes.clear() before returning when subscriptions.is_empty(), or move those clear() calls above the early return; ensure you still abort and await any existing WorkerSubscription handles (sub.handle.abort()/sub.handle.await) when subscriptions are present.
♻️ Duplicate comments (2)
model_gateway/src/policies/cache_aware.rs (2)
560-583:⚠️ Potential issue | 🟠 MajorRestrict overlap selection to healthy workers with positive overlap scores.
Current ranking includes healthy workers with implicit score
0. When overlap exists only for stale/unhealthy tenants, this still returns an “overlap match” path instead of no-overlap fallback.Suggested fix
- let best_idx = healthy_indices + let best_idx = healthy_indices .iter() + .filter(|&&idx| overlap.scores.get(workers[idx].url()).copied().unwrap_or(0) > 0) .max_by(|&&a, &&b| { let score_a = overlap.scores.get(workers[a].url()).copied().unwrap_or(0); let score_b = overlap.scores.get(workers[b].url()).copied().unwrap_or(0); score_a .cmp(&score_b) .then_with(|| workers[b].load().cmp(&workers[a].load())) .then_with(|| { let size_a = overlap .tree_sizes .get(workers[a].url()) .copied() .unwrap_or(0); let size_b = overlap .tree_sizes .get(workers[b].url()) .copied() .unwrap_or(0); size_b.cmp(&size_a) }) }) .copied()?;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/policies/cache_aware.rs` around lines 560 - 583, The selection currently considers healthy_indices with implicit score 0; restrict to only healthy workers that have a positive overlap by filtering healthy_indices using overlap.scores.get(workers[idx].url()).copied().unwrap_or(0) > 0 before calling max_by (or, alternatively, after computing best_idx verify that its score is > 0 and return None/trigger no-overlap fallback otherwise). Update the logic around best_idx / healthy_indices and reference overlap.scores, workers[..].url(), and best_idx so only workers with positive overlap are eligible.
451-453:⚠️ Potential issue | 🔴 CriticalFetch the event indexer once and pass it through to avoid TOCTOU.
Line 451 does a presence check, then Line 510-511 fetches again. If the map changes between those points, event-driven selection can return
Noneunexpectedly.Suggested fix
- if let Some(tokens) = request_tokens { - if self.has_event_indexer(model_id) { - self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id) - } else { - self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) - } + if let Some(tokens) = request_tokens { + if let Some(monitor) = self.kv_monitor.as_ref() { + if let Some(indexer) = monitor.get_indexer(model_id) { + self.select_worker_event_driven( + workers, + tokens, + &healthy_indices, + model_id, + monitor, + &indexer, + ) + } else { + self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) + } + } else { + self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) + } } else {- fn select_worker_event_driven( + fn select_worker_event_driven( &self, workers: &[Arc<dyn Worker>], tokens: &[u32], healthy_indices: &[usize], model_id: &str, + monitor: &KvEventMonitor, + indexer: &PositionalIndexer, ) -> Option<usize> { - let monitor = self.kv_monitor.as_ref()?; - let indexer = monitor.get_indexer(model_id)?;Also applies to: 510-512
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/policies/cache_aware.rs` around lines 451 - 453, Check for the event indexer once and pass that retrieved value into select_worker_event_driven instead of re-checking; replace the pattern that calls has_event_indexer(model_id) and later re-fetches the indexer (the code that currently fetches the event indexer at lines ~510-511) with a single retrieval (e.g., let indexer = ...get(model_id) or fetch_event_indexer(model_id)), then call select_worker_event_driven(workers, tokens, &healthy_indices, model_id, indexer) (or pass the Option/Reference) and update select_worker_event_driven’s signature to accept the indexer so the TOCTOU race is eliminated. Ensure both call sites that previously re-fetched the indexer (around has_event_indexer and the later fetch) are updated to use the single retrieved value.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/kv_event_monitor.rs`:
- Around line 161-164: There is a TOCTOU race between computing
should_remove_indexer and removing shared state; update the removal to re-check
model liveness just before mutating shared maps: after computing
should_remove_indexer, verify there are still no active workers/subscribers for
that model (re-read the subscribers collection or call a helper like
has_active_workers(model_id)) and only then call indexers.remove(&sub.model_id)
and block_sizes.remove(&sub.model_id). Reference the existing symbols
on_worker_added, should_remove_indexer, indexers.remove, block_sizes.remove and
the subscribers/workers collection to implement the final atomic liveness check.
---
Outside diff comments:
In `@model_gateway/src/core/kv_event_monitor.rs`:
- Around line 169-190: The stop() implementation returns early when
subscriptions.is_empty(), which skips clearing in-memory state and can leave
stale entries set via methods like set_block_size; modify stop() (the method
that reads self.worker_handles) so it always calls self.indexers.clear() and
self.block_sizes.clear() before returning when subscriptions.is_empty(), or move
those clear() calls above the early return; ensure you still abort and await any
existing WorkerSubscription handles (sub.handle.abort()/sub.handle.await) when
subscriptions are present.
---
Duplicate comments:
In `@model_gateway/src/policies/cache_aware.rs`:
- Around line 560-583: The selection currently considers healthy_indices with
implicit score 0; restrict to only healthy workers that have a positive overlap
by filtering healthy_indices using
overlap.scores.get(workers[idx].url()).copied().unwrap_or(0) > 0 before calling
max_by (or, alternatively, after computing best_idx verify that its score is > 0
and return None/trigger no-overlap fallback otherwise). Update the logic around
best_idx / healthy_indices and reference overlap.scores, workers[..].url(), and
best_idx so only workers with positive overlap are eligible.
- Around line 451-453: Check for the event indexer once and pass that retrieved
value into select_worker_event_driven instead of re-checking; replace the
pattern that calls has_event_indexer(model_id) and later re-fetches the indexer
(the code that currently fetches the event indexer at lines ~510-511) with a
single retrieval (e.g., let indexer = ...get(model_id) or
fetch_event_indexer(model_id)), then call select_worker_event_driven(workers,
tokens, &healthy_indices, model_id, indexer) (or pass the Option/Reference) and
update select_worker_event_driven’s signature to accept the indexer so the
TOCTOU race is eliminated. Ensure both call sites that previously re-fetched the
indexer (around has_event_indexer and the later fetch) are updated to use the
single retrieved value.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (3)
model_gateway/src/core/kv_event_monitor.rsmodel_gateway/src/policies/cache_aware.rsprotocols/src/worker.rs
slin1237
force-pushed
the
slin/cache-aware-event-policy
branch
from
March 2, 2026 15:48
879be90 to
9b97c93
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
model_gateway/src/policies/cache_aware.rs (2)
451-455:⚠️ Potential issue | 🟠 MajorTOCTOU window remains between indexer check and use.
has_event_indexer(model_id)at line 451 checks for indexer presence, butselect_worker_event_drivenat line 510-511 fetches it again. If the indexer is removed between these calls,select_worker_event_drivenreturnsNone, causingselect_workerto returnNoneeven when healthy workers exist.Consider fetching the indexer once and passing it to
select_worker_event_driven:Suggested fix
if let Some(tokens) = request_tokens { - if self.has_event_indexer(model_id) { - self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id) + if let Some(indexer) = self + .kv_monitor + .as_ref() + .and_then(|m| m.get_indexer(model_id)) + { + self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id, &indexer) } else { self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) }Then update
select_worker_event_drivento acceptindexer: &PositionalIndexeras a parameter instead of fetching it internally.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/policies/cache_aware.rs` around lines 451 - 455, The TOCTOU occurs because has_event_indexer(model_id) is called in select_worker and then select_worker_event_driven re-fetches the indexer, allowing it to be removed in between; change select_worker to fetch the indexer once (e.g. call get_positional_indexer or whatever lookup returns Option<PositionalIndexer>), branch on that Option, and pass a &PositionalIndexer into select_worker_event_driven instead of letting that function fetch it internally; update the select_worker_event_driven signature to accept indexer: &PositionalIndexer and remove its internal lookup, and if the initial fetch returns None fall back to calling select_worker_with_tokens (or return None consistent with current behaviour).
559-582:⚠️ Potential issue | 🟠 MajorHealthy workers with zero overlap should be excluded from selection.
If
overlap.scorescontains entries only for unhealthy workers (or workers not inhealthy_indices), the current logic will still select a healthy worker with implicit score 0 instead of returningNonefor min-load fallback.Suggested fix
// Select worker with best overlap; tie-break by load (lower) then tree size (smaller) let best_idx = healthy_indices .iter() + .filter(|&&idx| overlap.scores.get(workers[idx].url()).copied().unwrap_or(0) > 0) .max_by(|&&a, &&b| { let score_a = overlap.scores.get(workers[a].url()).copied().unwrap_or(0); let score_b = overlap.scores.get(workers[b].url()).copied().unwrap_or(0); score_a .cmp(&score_b) .then_with(|| workers[b].load().cmp(&workers[a].load())) .then_with(|| { let size_a = overlap .tree_sizes .get(workers[a].url()) .copied() .unwrap_or(0); let size_b = overlap .tree_sizes .get(workers[b].url()) .copied() .unwrap_or(0); size_b.cmp(&size_a) }) }) .copied()?;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/policies/cache_aware.rs` around lines 559 - 582, The selection currently treats missing overlap.scores as implicit zero and can pick a healthy worker with no explicit score; update the selection logic around best_idx to first filter healthy_indices to only those whose workers[idx].url() exists in overlap.scores (i.e., overlap.scores.contains_key), and if that filtered list is empty return None (preserve the min-load fallback behavior), otherwise run the existing max_by comparison on that filtered list using the same comparisons of score, load(), and tree_sizes; reference healthy_indices, overlap.scores, workers, and best_idx to locate and modify the code.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/policies/cache_aware.rs`:
- Around line 489-496: The helper method has_event_indexer checks
kv_monitor.as_ref().and_then(|m| m.get_indexer(model_id)).is_some() and is now
redundant if you apply the TOCTOU fix (fetch the indexer once at the call site);
remove the has_event_indexer function from cache_aware.rs and update any callers
to use the pre-fetched indexer value (the result of kv_monitor.get_indexer or
its Option) directly instead of calling has_event_indexer, ensuring callers
reference the fetched Option and check is_some()/is_none() or pattern-match as
appropriate.
---
Duplicate comments:
In `@model_gateway/src/policies/cache_aware.rs`:
- Around line 451-455: The TOCTOU occurs because has_event_indexer(model_id) is
called in select_worker and then select_worker_event_driven re-fetches the
indexer, allowing it to be removed in between; change select_worker to fetch the
indexer once (e.g. call get_positional_indexer or whatever lookup returns
Option<PositionalIndexer>), branch on that Option, and pass a &PositionalIndexer
into select_worker_event_driven instead of letting that function fetch it
internally; update the select_worker_event_driven signature to accept indexer:
&PositionalIndexer and remove its internal lookup, and if the initial fetch
returns None fall back to calling select_worker_with_tokens (or return None
consistent with current behaviour).
- Around line 559-582: The selection currently treats missing overlap.scores as
implicit zero and can pick a healthy worker with no explicit score; update the
selection logic around best_idx to first filter healthy_indices to only those
whose workers[idx].url() exists in overlap.scores (i.e.,
overlap.scores.contains_key), and if that filtered list is empty return None
(preserve the min-load fallback behavior), otherwise run the existing max_by
comparison on that filtered list using the same comparisons of score, load(),
and tree_sizes; reference healthy_indices, overlap.scores, workers, and best_idx
to locate and modify the code.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (3)
model_gateway/src/core/kv_event_monitor.rsmodel_gateway/src/policies/cache_aware.rsprotocols/src/worker.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9b97c931ca
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: df41861b24
ℹ️ 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".
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/config/validation.rs (1)
739-759: 🧹 Nitpick | 🔵 TrivialAdd one explicit negative test for
block_size: 0.The updated tests only exercise valid
block_sizevalues. Please add a dedicated case assertingConfigError::InvalidValueforblock_size == 0to lock this regression fix.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/config/validation.rs` around lines 739 - 759, The test suite is missing a negative case for block_size == 0 in PolicyConfig::CacheAware; add a new unit test (e.g., test_validate_cache_aware_block_size_zero) that constructs a RouterConfig with RoutingMode::Regular and PolicyConfig::CacheAware where block_size: 0 (keep other fields valid) and assert that ConfigValidator::validate(&config) returns an Err matching ConfigError::InvalidValue; reference the existing test_validate_cache_aware_thresholds as a template and reuse RouterConfig, PolicyConfig::CacheAware, and ConfigValidator::validate to locate where to add the new test.
♻️ Duplicate comments (1)
model_gateway/src/policies/cache_aware.rs (1)
450-453:⚠️ Potential issue | 🔴 CriticalFetch indexer once to avoid check/use race returning
None.Line 451 checks availability, but Lines 510-511 fetch again with
?. If the indexer disappears between those operations, routing can returnNoneeven when healthy workers exist. Fetch once inselect_workerand pass the resolved indexer/block_size intoselect_worker_event_driven.Suggested fix
- if let Some(tokens) = request_tokens { - if self.has_event_indexer(model_id) { - self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id) - } else { - self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) - } + if let Some(tokens) = request_tokens { + if let Some(monitor) = self.kv_monitor.as_ref() { + if let Some(indexer) = monitor.get_indexer(model_id) { + let block_size = + monitor.block_size(model_id).unwrap_or(self.config.block_size); + return self.select_worker_event_driven( + workers, + tokens, + &healthy_indices, + model_id, + &indexer, + block_size, + ); + } + } + self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) } else { let text = request_text.unwrap_or(""); self.select_worker_with_text(workers, text, &healthy_indices, model_id) } @@ fn select_worker_event_driven( &self, workers: &[Arc<dyn Worker>], tokens: &[u32], healthy_indices: &[usize], model_id: &str, + indexer: &PositionalIndexer, + block_size: usize, ) -> Option<usize> { - let monitor = self.kv_monitor.as_ref()?; - let indexer = monitor.get_indexer(model_id)?; - - // Per-model block_size: learned from events > config default - let block_size = monitor - .block_size(model_id) - .unwrap_or(self.config.block_size); - if let Some(idx) = - Self::score_overlap(workers, tokens, healthy_indices, &indexer, block_size) + Self::score_overlap(workers, tokens, healthy_indices, indexer, block_size) { return Some(idx); }Also applies to: 510-511
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/policies/cache_aware.rs` around lines 450 - 453, The code currently calls has_event_indexer(model_id) then later re-fetches the indexer with `?`, causing a race where the indexer may disappear and cause None to be returned; instead, in select_worker obtain the indexer and block_size once (replace the has_event_indexer check with a single fetch that returns Option<(indexer, block_size)>) and pass these resolved values into select_worker_event_driven (update its signature to accept the indexer/block_size), and apply the same change to the other call site that currently re-fetches the indexer so both places use the single fetched value to avoid the check/use race.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@model_gateway/src/config/validation.rs`:
- Around line 739-759: The test suite is missing a negative case for block_size
== 0 in PolicyConfig::CacheAware; add a new unit test (e.g.,
test_validate_cache_aware_block_size_zero) that constructs a RouterConfig with
RoutingMode::Regular and PolicyConfig::CacheAware where block_size: 0 (keep
other fields valid) and assert that ConfigValidator::validate(&config) returns
an Err matching ConfigError::InvalidValue; reference the existing
test_validate_cache_aware_thresholds as a template and reuse RouterConfig,
PolicyConfig::CacheAware, and ConfigValidator::validate to locate where to add
the new test.
---
Duplicate comments:
In `@model_gateway/src/policies/cache_aware.rs`:
- Around line 450-453: The code currently calls has_event_indexer(model_id) then
later re-fetches the indexer with `?`, causing a race where the indexer may
disappear and cause None to be returned; instead, in select_worker obtain the
indexer and block_size once (replace the has_event_indexer check with a single
fetch that returns Option<(indexer, block_size)>) and pass these resolved values
into select_worker_event_driven (update its signature to accept the
indexer/block_size), and apply the same change to the other call site that
currently re-fetches the indexer so both places use the single fetched value to
avoid the check/use race.
3 of 6 tasks
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
model_gateway/src/core/kv_event_monitor.rs (1)
161-164:⚠️ Potential issue | 🟠 MajorRe-check model liveness immediately before deleting shared model state.
There is still a TOCTOU window: after the initial check, a concurrent
on_worker_addedcan reintroduce the model before Line 161-Line 164 executes, causing removal of activeindexers/block_sizes.Suggested fix
- if should_remove_indexer { - self.indexers.remove(&sub.model_id); - self.block_sizes.remove(&sub.model_id); - } + if should_remove_indexer { + let still_last = { + let handles = self.worker_handles.lock().await; + !handles.values().any(|other| other.model_id == sub.model_id) + }; + if still_last { + self.indexers.remove(&sub.model_id); + self.block_sizes.remove(&sub.model_id); + } + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/kv_event_monitor.rs` around lines 161 - 164, There is a TOCTOU race: after deciding should_remove_indexer you must re-check model liveness immediately before mutating shared state so a concurrent on_worker_added cannot reintroduce the model between the check and removal; update the code around self.indexers.remove(&sub.model_id) / self.block_sizes.remove(&sub.model_id) to recompute the same liveness predicate (or call the same helper used earlier) while holding the same lock used by on_worker_added (or use an atomic/entry-style conditional remove) and only perform the removes if the fresh check still indicates the model is absent.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/app_context.rs`:
- Around line 568-571: The monitor is only created when the top-level flags
is_cache_aware and is_grpc are both true, which prevents dynamic or per-model
cache-aware routing; change the initialization so the monitor is created
whenever cache-aware routing may be needed (not only when config.policy is
globally CacheAware and config.connection_mode is statically Grpc). Concretely,
remove the strict conjunction of is_cache_aware && is_grpc and instead either
(a) initialize the monitor unconditionally and guard activation at use-time, or
(b) lazily create/enable the monitor when a model/worker reports a CacheAware
policy or a GRPC connection is detected; update the code paths around the
is_cache_aware, is_grpc checks and the monitor creation/registration logic so
per-model overrides and late-discovered gRPC workers correctly trigger
event-driven routing.
---
Duplicate comments:
In `@model_gateway/src/core/kv_event_monitor.rs`:
- Around line 161-164: There is a TOCTOU race: after deciding
should_remove_indexer you must re-check model liveness immediately before
mutating shared state so a concurrent on_worker_added cannot reintroduce the
model between the check and removal; update the code around
self.indexers.remove(&sub.model_id) / self.block_sizes.remove(&sub.model_id) to
recompute the same liveness predicate (or call the same helper used earlier)
while holding the same lock used by on_worker_added (or use an
atomic/entry-style conditional remove) and only perform the removes if the fresh
check still indicates the model is absent.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (9)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pymodel_gateway/src/app_context.rsmodel_gateway/src/core/kv_event_monitor.rsmodel_gateway/src/core/steps/worker/local/remove_from_policy_registry.rsmodel_gateway/src/core/steps/worker/shared/update_policies.rsmodel_gateway/src/main.rsmodel_gateway/src/policies/registry.rsmodel_gateway/src/service_discovery.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9cbc9d4f42
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 36bfbe291d
ℹ️ 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".
Comment on lines
+83
to
+87
| if let Some(p) = self.prefill_policy.get() { | ||
| Self::maybe_inject_monitor(p, monitor.as_ref()); | ||
| } | ||
| if let Some(p) = self.decode_policy.get() { | ||
| Self::maybe_inject_monitor(p, monitor.as_ref()); |
There was a problem hiding this comment.
set_kv_event_monitor only propagates to prefill_policy/decode_policy if those OnceLocks are already populated, but in the startup path AppContextBuilder::from_config wires the monitor before RouterFactory::create_pd_router/create_grpc_pd_router set PD policies. This leaves cache-aware PD policies without a monitor, so PD routing silently never enters the event-driven branch even though KV subscriptions may be active. Inject the monitor when setting prefill/decode policies (or run monitor propagation again after router policy initialization).
Useful? React with 👍 / 👎.
| fn with_kv_event_monitor(mut self, config: &RouterConfig) -> Self { | ||
| use crate::config::types::PolicyConfig; | ||
|
|
||
| let is_cache_aware = matches!(config.policy, PolicyConfig::CacheAware { .. }); |
There was a problem hiding this comment.
with_kv_event_monitor only checks config.policy, so configurations that use cache-aware through PD overrides (prefill_policy/decode_policy) or later per-model policy hints can run without any monitor at all. In those valid setups app_context.kv_event_monitor remains None, which prevents worker subscriptions and disables event-driven cache-aware routing for affected models. The monitor creation condition should account for all routes that can instantiate cache_aware, not just the main policy field.
Useful? React with 👍 / 👎.
slin1237
force-pushed
the
slin/cache-aware-event-policy
branch
from
March 2, 2026 20:22
36bfbe2 to
eddd62c
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Add a new highest-priority routing tier to CacheAwarePolicy that uses PositionalIndexer overlap scoring from KvEventMonitor when available. This enables precise KV cache hit routing based on real backend cache state rather than approximate prefix matching. Three-tier dispatch order (gRPC path): 1. Event-driven: PositionalIndexer overlap scoring (when monitor has data for the model) 2. Approximate token tree: existing TokenTree prefix matching (fallback) 3. Min-load: when workers are imbalanced (unchanged) HTTP path remains unchanged (string tree). What changed: - policies/cache_aware.rs: add kv_monitor field, set_kv_event_monitor setter, try_event_driven_routing and select_worker_with_events methods, update select_worker dispatch to 3-tier, add 9 new tests - policies/mod.rs: add block_size to CacheAwareConfig (default 16) - core/kv_event_monitor.rs: add manual Debug impl - config/types.rs: add block_size to PolicyConfig::CacheAware with serde default for backward compat - config/builder.rs, validation.rs, factory.rs, main.rs, bindings/python/src/lib.rs, test files: propagate block_size Refs: Phase 1, Task 6 of KV cache event-driven routing Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…ive type
Restructure CacheAwarePolicy from cascading tiers (event → token tree
fallback) to three mutually exclusive routing types selected by worker
connection mode and KV event availability:
1. Event-Driven (gRPC + KV events): PositionalIndexer overlap scoring
2. Approximate Token Tree (gRPC, no events): TokenTree prefix matching
3. Approximate String Tree (HTTP): Tree character prefix matching
What changed:
- model_gateway/src/policies/cache_aware.rs:
- Replace try_event_driven_routing + select_worker_with_events with
has_event_indexer, select_worker_event_driven, and score_overlap
- Dispatch logic uses if/else (mutually exclusive) instead of fallthrough
- Event-driven path is self-contained: overlap match or min-load fallback,
never falls through to approximate token tree
- Rewrite file header to document all three types
- Add 12 new tests: 6 score_overlap unit tests (best match, no match,
load tiebreak, tree size tiebreak, short request, partial match) and
6 select_worker_event_driven integration tests (overlap selects cached,
no overlap uses min-load, short request min-load, no monitor token tree,
set/clear monitor, imbalanced skips events)
- model_gateway/src/core/kv_event_monitor.rs:
- Change indexers field to pub(crate) for test access from cache_aware tests
Why:
The cascading tier design (try event-driven, fall back to token tree) was
incorrect — when event data exists for a model, approximate tree routing
would produce stale/conflicting decisions. The three types are mutually
exclusive by design: if we have real cache state from KV events, we should
not also consult an approximate tree. When event-driven has no overlap
(cold start, novel tokens), min-load is the correct fallback.
Signed-off-by: Simon Lin <simon@simonlin.me>
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…nitor Add kv_block_size as a configurable per-worker field in WorkerSpec, discoverable during worker registration. This replaces the hardcoded block_size=16 default with a per-model value that users can control via the workers API. What changed: - protocols/src/worker.rs: add kv_block_size: Option<usize> to WorkerSpec with serde support (default None, skip_serializing_if None) - model_gateway/src/core/kv_event_monitor.rs: add block_sizes DashMap per-model storage, block_size()/set_block_size() public API, seed from WorkerSpec during on_worker_added, cleanup on worker removal and stop - model_gateway/src/policies/cache_aware.rs: resolve per-model block_size from monitor (falling back to config default), pass explicit block_size to score_overlap instead of reading from config Why: Block size is a static property of the backend engine (e.g. 16 for SGLang page size). Different models or backends may use different block sizes. Hardcoding 16 in the config prevents multi-model deployments with heterogeneous block sizes. Following Dynamo's pattern where block_size is per-model via ModelDeploymentCard, discovered during registration. How: Discovery flow: WorkerSpec.kv_block_size (user sets via API) → on_worker_added seeds KvEventMonitor.block_sizes → CacheAwarePolicy reads monitor.block_size(model_id) at query time, falls back to CacheAwareConfig.block_size (default 16) if not set. Tests: 3 new tests (set_block_size, stop_clears_block_sizes, event_driven_uses_monitor_block_size), 6 existing tests updated for new score_overlap signature. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
- Add block_size > 0 validation in config validator to prevent panic from chunks_exact(0) in compute_request_content_hashes - Filter zero-overlap workers in score_overlap to avoid selecting a worker with no cache match as if it were an overlap hit - Simplify score_overlap ranking with max_by_key + tuple instead of manual max_by with then_with chains Signed-off-by: Simon Lin <simon@lightseek.ai> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Wire the KvEventMonitor into the production workflow so event-driven
cache-aware routing activates automatically for eligible workers.
AppContext:
- Add kv_event_monitor field (Option<Arc<KvEventMonitor>>)
- Create monitor in with_kv_event_monitor() when policy is cache_aware
and connection mode is gRPC
- Inject monitor into PolicyRegistry at startup
PolicyRegistry:
- Add kv_event_monitor field with RwLock for thread-safe late init
- Add set_kv_event_monitor() method (follows set_mesh_sync pattern)
- Inject monitor into new CacheAwarePolicy instances during
create_policy_from_type("cache_aware")
UpdatePoliciesStep:
- Call monitor.on_worker_added(worker) for gRPC workers with
cache_aware policy, starting KV event subscription
RemoveFromPolicyRegistryStep:
- Call monitor.on_worker_removed(worker_url) before policy cleanup,
stopping subscription and removing worker from indexer
block_size CLI parameter (replaces hardcoded 16):
- router_args.py: add block_size field + --block-size argument
- lib.rs (Python bindings): add block_size to Router struct, signature,
constructor, use self.block_size in CacheAware config
- main.rs: add --block-size CLI arg, use self.block_size in config
Bug fix:
- kv_event_monitor.rs: stop() now always clears indexers and
block_sizes, not only when subscriptions exist
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Interior mutability for kv_monitor (P1): - Change CacheAwarePolicy.kv_monitor from Option<Arc<KvEventMonitor>> to RwLock<Option<Arc<KvEventMonitor>>> so set_kv_event_monitor works via &self on policies already behind Arc<dyn LoadBalancingPolicy>. - PolicyRegistry.set_kv_event_monitor now propagates to all existing cache-aware policies (default, prefill, decode, per-model), fixing the issue where default_policy created before the monitor was available would never receive it. Empty indexer fallback (P1): - has_event_indexer now checks current_size() > 0, so an empty indexer (startup, reconnect) falls through to the approximate token tree instead of taking the event-driven path and landing on min-load. - Add test_empty_indexer_falls_through_to_token_tree test. Monitor creation scope (P1): - Remove gRPC connection_mode check from with_kv_event_monitor in AppContextBuilder. The monitor is cheap (empty DashMaps) and UpdatePoliciesStep already gates subscriptions on cache_aware && gRPC. This allows dynamic/per-model gRPC workers to use event-driven routing even when the global connection_mode is not gRPC. Model key normalization (P2): - Add normalize_model_id helper to KvEventMonitor that maps empty model_id to UNKNOWN_MODEL_ID, matching routing's normalize_model_key. Block size validation (P2): - Skip kv_block_size=0 from WorkerSpec instead of storing it, which would poison overlap scoring with zero-length chunks. TOCTOU fix in on_worker_removed (P2): - Re-acquire worker_handles lock after aborting the subscription task to check if the model's last worker was truly removed. Prevents race where concurrent on_worker_added for the same model could repopulate handles between lock release and indexer removal. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
slin1237
force-pushed
the
slin/cache-aware-event-policy
branch
from
March 2, 2026 22:23
eddd62c to
59b18a9
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
slin1237
added a commit
that referenced
this pull request
Mar 3, 2026
Add comprehensive benchmarks for the event-driven PositionalIndexer alongside the existing StringTree and TokenTree benchmarks in radix_tree_benchmark.rs. What changed: - model_gateway/benches/radix_tree_benchmark.rs: Added PositionalIndexer helper functions (generate_token_chunks, chunks_to_stored_blocks, flatten_tokens, build_populated_indexer), three benchmark macros (bench_indexer_store!, bench_indexer_match!, bench_indexer_concurrent!), benchmark invocations in bench_summary() covering STORE, MATCH, and CONCURRENT operations across worker counts [10, 100, 500], block sizes [16, 64], blocks per worker [64, 256, 1024], and query sizes [32, 128, 512], plus summary table output in print_summary(). - .github/workflows/benchmark-radix-tree.yml: Remove --exact flag from cargo bench invocation which prevented individual benchmarks within the benchmark_summary group from running (substring match is needed). Why: The PositionalIndexer (event-driven KV cache routing) was added in PR #571 but had no benchmarks. This covers the hot path (find_matches at 5-170µs depending on worker count) and background operations (apply_stored), matching the existing benchmark report format. How: Follows the same macro-based pattern as StringTree/TokenTree benchmarks with iter_custom, BTreeMap result collection, and formatted summary tables. The build_populated_indexer helper simulates realistic cache state with shared system prompt prefix blocks and unique per-worker blocks. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
3 tasks done
This was referenced Mar 4, 2026
3 of 4 tasks
2 of 3 tasks
This was referenced Aug 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds event-driven routing as a new cache-aware routing type in
CacheAwarePolicy, usingPositionalIndexeroverlap scoring fromKvEventMonitorto route requests based on actual backend KV cache state. This is Task 6 of the KV event-driven routing feature (Phase 1).Refs: #560, #565, #568
What changed
model_gateway/src/policies/cache_aware.rs:PositionalIndexeroverlap scoring with min-load fallbackTokenTreeprefix matchingTreecharacter prefix matchingkv_monitor: Option<Arc<KvEventMonitor>>field andset_kv_event_monitorsetterhas_event_indexer— type selection helper checking if monitor has an indexer for modelselect_worker_event_driven— self-contained event routing: overlap match or min-load fallbackscore_overlap— rank workers by overlap count (primary), load (secondary), tree size (tertiary)if/else(mutually exclusive), not fallthroughscore_overlapunit tests + 6select_worker_event_drivenintegration testsmodel_gateway/src/core/kv_event_monitor.rs:indexersfield from private topub(crate)for test access fromcache_awaretestsmodel_gateway/src/policies/mod.rs(from prior commit):block_size: usizefield toCacheAwareConfig(default: 16, SGLang page size)Why
When event data is available for a model, routing should use real cache state from KV events exclusively. The original cascading design (try event-driven, fall back to token tree) would produce conflicting decisions — if we have ground-truth cache state, an approximate tree built from routing history would be stale. The three types are mutually exclusive by design. When event-driven has no overlap (cold start, novel tokens), min-load is the correct fallback.
How
The dispatch logic selects exactly one routing type based on available data:
Event-driven routing:
compute_request_content_hashes(tokens, block_size)— chunk tokens into blocks, hash eachindexer.find_matches(content_hashes)— jump search for overlap scoresscore_overlap— select worker with best overlap; tie-break by load then tree sizeTest plan
cargo test -p kv-index— 158 tests pass (PositionalIndexer + hash functions + E2E)cargo clippy -p kv-index— cleancache_aware::tests:test_score_overlap_selects_best_match— w1 has blocks, w2 doesn't → selects w1test_score_overlap_no_match_returns_none— different tokens → Nonetest_score_overlap_load_tiebreak— equal overlap, w1 higher load → selects w2test_score_overlap_tree_size_tiebreak— equal overlap+load, w2 larger tree → selects w1test_score_overlap_short_request_returns_none— request < block_size → Nonetest_score_overlap_partial_match— w1 has 4 blocks, w2 has 2 → selects w1test_event_driven_overlap_selects_cached_worker— full dispatch viaselect_workertest_event_driven_no_overlap_uses_min_load— no overlap → min-load, not token treetest_event_driven_short_request_uses_min_load— short request → min-load fallbacktest_no_monitor_uses_token_tree— no monitor → falls back to token treetest_set_kv_event_monitor— set/clear monitor lifecycletest_imbalanced_skips_event_driven— load imbalance → shortest-queue regardless of eventsNote:
smgcrate has a pre-existing build error inproto_wrapper.rs(unrelatedimage_datafield removed from proto) that blockscargo test -p smg. Thekv-indexcrate tests and clippy pass clean.Summary by CodeRabbit
New Features
Internal Improvements