Repository navigation
feat(routing): honor X-SMG-Routing-Key on any policy (#1843) - #1848
Conversation
github-actions
Bot
added
documentation
|
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 (13)
📝 WalkthroughWalkthroughThe PR adds a routing-key override option through Python and Rust config paths, carries worker-leg metadata through HTTP and gRPC selection, and updates policy registry and manual routing to use ChangesRouting-key override rollout
Sequence Diagram(s)sequenceDiagram
participant AppContextBuilder
participant PolicyRegistry
participant ManualPolicy
AppContextBuilder->>PolicyRegistry: with_override(config.policy, config.routing_key_override)
PolicyRegistry->>PolicyRegistry: select_worker(request)
PolicyRegistry->>ManualPolicy: select_worker(info) when X-SMG-Routing-Key is present
ManualPolicy->>ManualPolicy: select_by_routing_id(info.leg.routing_id_prefix() + routing_key)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![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 implements an opt-in, per-request sticky-routing override via the X-SMG-Routing-Key header for any routing policy. It introduces a RoutingKeyOverride decorator policy, a WorkerLeg discriminator to support independent stickiness for prefill and decode legs in PD mode, and exposes the feature via CLI flags and Python bindings. The review feedback highlights critical integration issues in PolicyRegistry where direct downcasting and pointer comparisons (Arc::ptr_eq) will fail on wrapped policies, potentially breaking worker initialization and causing duplicate entries. Additionally, the feedback recommends defining explicit defaults during deserialization to prevent memory leaks from disabled evictions, and optimizing the hot path by avoiding heap allocations for the default WorkerLeg::Single case.
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.
| if let Some(routing_id) = extract_routing_key(info.headers) { | ||
| let (idx, branch) = self.select_by_routing_id(workers, routing_id, &healthy_indices); | ||
| let namespaced = format!("{}{}", info.leg.routing_id_prefix(), routing_id); | ||
| let (idx, branch) = self.select_by_routing_id(workers, &namespaced, &healthy_indices); | ||
| return (Some(idx), branch); | ||
| } |
There was a problem hiding this comment.
format!("{}{}", info.leg.routing_id_prefix(), routing_id) allocates a new String on the heap for every request.
Since WorkerLeg::Single is the default and most common case (where routing_id_prefix() is ""), we can completely avoid this heap allocation on the hot path by checking if info.leg == WorkerLeg::Single and using routing_id directly.
if let Some(routing_id) = extract_routing_key(info.headers) {
let (idx, branch) = if info.leg == crate::policies::WorkerLeg::Single {
self.select_by_routing_id(workers, routing_id, &healthy_indices)
} else {
let namespaced = format!("{}{}", info.leg.routing_id_prefix(), routing_id);
self.select_by_routing_id(workers, &namespaced, &healthy_indices)
};
return (Some(idx), branch);
}References
- Avoid heap allocations in hot or periodic paths. Keep allocation-heavy helper methods restricted to cold paths where the allocation overhead is negligible.
![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: d803a0a7c8
ℹ️ 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".
![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}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
model_gateway/src/policies/registry.rs (1)
373-396: 🩺 Stability & Availability | 🔴 CriticalUnwrap
RoutingKeyOverridebefore cache-aware/bucket init/cleanup.
set_prefill_policy/set_decode_policyandcreate_policy_from_typenow wrap eligible policies, but these helpers still checkpolicy.name()and downcast the wrapper directly:
init_cache_aware_policyremove_worker_from_cache_awareremove_worker_from_pd_cache_awareinit_pd_cache_aware_policiesinit_pd_bucket_policiesWhen
routing_key_overrideis enabled,cache_awareandbucketbecomerouting_key_override, so the downcast never reachesCacheAwarePolicy/BucketPolicy. That skips worker initialization and cleanup for wrapped policies. UseSelf::effective_policy(...)before the name check and downcast here, as in the monitor/load-receiver injection paths.🤖 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 373 - 396, The cache-aware and bucket policy init/cleanup paths in registry.rs are operating on the RoutingKeyOverride wrapper instead of the underlying policy. Update init_cache_aware_policy, remove_worker_from_cache_aware, remove_worker_from_pd_cache_aware, init_pd_cache_aware_policies, and init_pd_bucket_policies to call Self::effective_policy(...) before checking policy.name() and before any downcast, matching the existing monitor/load-receiver injection flow. This should ensure CacheAwarePolicy and BucketPolicy are reached correctly when routing_key_override is enabled.docs/superpowers/plans/2026-06-25-routing-key-override.md (1)
627-631: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRestore the blank lines around the final heading.
This trips MD022 and is likely to fail docs lint in CI.
🤖 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 `@docs/superpowers/plans/2026-06-25-routing-key-override.md` around lines 627 - 631, Restore the missing blank lines around the final markdown heading in the notes section so the heading is separated correctly from the surrounding list content and passes MD022 docs lint. Locate the final heading in the implementation notes block and adjust the surrounding spacing only; do not change the heading text or any other content.Source: Linters/SAST tools
model_gateway/src/routers/http/pd_router.rs (1)
882-917: 🗄️ Data Integrity & Integration | 🟠 MajorThread
WorkerLegthrough HTTP PD selection.
pick_worker_by_policy_archardcodesleg: WorkerLeg::Single, but it is used for both prefill and decode. WithRoutingKeyOverrideenabled,ManualPolicynamespacesX-SMG-Routing-Keybyinfo.leg.routing_id_prefix(), so the two pools share the same sticky key and can reuse the wrong entry across legs. PassWorkerLeg::PrefillandWorkerLeg::Decodefrom the two call sites instead ofSingle.🤖 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/routers/http/pd_router.rs` around lines 882 - 917, `pick_worker_by_policy_arc` is hardcoding `WorkerLeg::Single`, which breaks routing-key namespacing for both prefill and decode selections. Update the function to accept a `WorkerLeg` parameter and pass it into `SelectWorkerInfo` instead of `Single`. Then update the two HTTP PD call sites that select prefill and decode workers to pass `WorkerLeg::Prefill` and `WorkerLeg::Decode` respectively so `ManualPolicy` uses the correct sticky key namespace.
🤖 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 `@bindings/python/src/lib.rs`:
- Around line 760-769: The assignment_mode handling in Router policy
construction is inconsistent with convert_policy: invalid strings are currently
falling through to ManualAssignmentMode::Random instead of failing. Update the
mapping in the routing_key_override builder to reuse the same validated parsing
behavior as convert_policy, returning a ConfigError::InvalidValue (or shared
helper/closure) for unknown values rather than silently coercing them.
In `@docs/superpowers/specs/2026-06-25-routing-key-override-design.md`:
- Around line 101-107: The example block for RoutingId is missing a Rust code
fence label, which triggers markdownlint and removes syntax highlighting. Update
the fenced block in the spec so the snippet is marked as Rust, keeping the
example content the same while adding the proper language tag around the
RoutingId match info.leg section.
In `@model_gateway/src/config/types.rs`:
- Around line 309-323: The sticky-routing override config currently defaults
both eviction knobs to 0, which disables cleanup when `RoutingKeyOverrideConfig`
is enabled from config files. Update the serde defaults for
`eviction_interval_secs` and `max_idle_secs` in `RoutingKeyOverrideConfig` to
non-zero helper defaults that match the CLI path used in `to_router_config`, so
`ManualPolicy::with_config` will arm eviction automatically. Keep the fix
localized to the config defaults and verify the chosen values still satisfy the
`eviction_interval_secs > 0 && max_idle_secs > 0` check.
In `@model_gateway/src/policies/routing_key_override.rs`:
- Around line 50-52: The policy name returned by RoutingKeyOverride is masking
the inner routing policy, which causes router checks to miss cache_aware/manual
behavior. Update the name reporting in RoutingKeyOverride to expose the wrapped
policy’s effective name, and make router-side gating in the HTTP router use that
inner policy name before deciding whether to apply WorkerLoadGuard. Keep the
change focused around RoutingKeyOverride and the policy-name check used by the
request routing paths.
---
Outside diff comments:
In `@docs/superpowers/plans/2026-06-25-routing-key-override.md`:
- Around line 627-631: Restore the missing blank lines around the final markdown
heading in the notes section so the heading is separated correctly from the
surrounding list content and passes MD022 docs lint. Locate the final heading in
the implementation notes block and adjust the surrounding spacing only; do not
change the heading text or any other content.
In `@model_gateway/src/policies/registry.rs`:
- Around line 373-396: The cache-aware and bucket policy init/cleanup paths in
registry.rs are operating on the RoutingKeyOverride wrapper instead of the
underlying policy. Update init_cache_aware_policy,
remove_worker_from_cache_aware, remove_worker_from_pd_cache_aware,
init_pd_cache_aware_policies, and init_pd_bucket_policies to call
Self::effective_policy(...) before checking policy.name() and before any
downcast, matching the existing monitor/load-receiver injection flow. This
should ensure CacheAwarePolicy and BucketPolicy are reached correctly when
routing_key_override is enabled.
In `@model_gateway/src/routers/http/pd_router.rs`:
- Around line 882-917: `pick_worker_by_policy_arc` is hardcoding
`WorkerLeg::Single`, which breaks routing-key namespacing for both prefill and
decode selections. Update the function to accept a `WorkerLeg` parameter and
pass it into `SelectWorkerInfo` instead of `Single`. Then update the two HTTP PD
call sites that select prefill and decode workers to pass `WorkerLeg::Prefill`
and `WorkerLeg::Decode` respectively so `ManualPolicy` uses the correct sticky
key namespace.
🪄 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: ed8cf590-a3d3-4ad8-8e2d-471d65cf3748
📒 Files selected for processing (16)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pydocs/superpowers/plans/2026-06-25-routing-key-override.mddocs/superpowers/specs/2026-06-25-routing-key-override-design.mdmodel_gateway/src/app_context.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/policies/manual.rsmodel_gateway/src/policies/mod.rsmodel_gateway/src/policies/registry.rsmodel_gateway/src/policies/routing_key_override.rsmodel_gateway/src/routers/grpc/common/stages/worker_selection.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/service_discovery.rs
d803a0a to
4b49a2c
Compare
github-actions
Bot
removed
the
documentation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4b49a2c10d
ℹ️ 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/policies/registry.rs (1)
162-202: 🎯 Functional Correctness | 🔴 CriticalRoute cache-aware/bucket init and cleanup through
effective_policy(...)
init_cache_aware_policy,remove_worker_from_cache_aware,remove_worker_from_pd_cache_aware,init_pd_cache_aware_policies, andinit_pd_bucket_policiesstill inspect theRoutingKeyOverridewrapper directly. When the sticky override is enabled,policy.name()becomes"routing_key_override", so thesecache_aware/bucketchecks and downcasts miss the inner policy and the init/removal hooks never run. UseSelf::effective_policy(...)for both the name check and the downcast.🤖 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 162 - 202, The cache-aware/bucket setup and cleanup paths are still looking at the RoutingKeyOverride wrapper instead of the underlying policy, so the init/remove hooks can be skipped when sticky override is enabled. Update the affected registry methods such as init_cache_aware_policy, remove_worker_from_cache_aware, remove_worker_from_pd_cache_aware, init_pd_cache_aware_policies, and init_pd_bucket_policies to call Self::effective_policy(policy) before checking policy.name() and before any CacheAwarePolicy downcast, so the inner policy is used consistently.
♻️ Duplicate comments (3)
bindings/python/src/lib.rs (1)
760-769: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
assignment_modeparsing diverges fromconvert_policy— invalid values silently coerced toRandom.
convert_policy's Manual branch (lines 545-556) now returnsConfigError::InvalidValuefor unrecognized modes, but here the catch-all_ => config::ManualAssignmentMode::Randomsilently swallows typos. SinceRouter(...)can be constructed directly from Python with an arbitraryassignment_mode, a misspelled value (e.g.min_groupp) would route viaRandominstead of failing fast. Reuse the validated mapping (or a shared helper) here.🤖 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 `@bindings/python/src/lib.rs` around lines 760 - 769, The assignment_mode handling in the routing_key_override setup silently falls back to Random for unknown values, unlike convert_policy’s validated Manual branch. Update Router construction in lib.rs to reuse the same validation logic or shared helper used by convert_policy so invalid assignment_mode strings return ConfigError::InvalidValue instead of being coerced. Keep the mapping aligned for the assignment_mode field and the config::ManualAssignmentMode conversion.model_gateway/src/config/types.rs (1)
309-323: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftSerde-defaulted eviction knobs (
0) disable TTL eviction, risking unbounded sticky-map growth when enabled via config file.
eviction_interval_secsandmax_idle_secsboth default to0under#[serde(default)].ManualPolicy::with_configonly arms the eviction task when both are> 0, so a config-file user who sets onlyenabled: truegets a stickyDashMapthat never evicts — unbounded growth under per-session keys (the agentic/RL use case). This diverges from the CLI path (eviction_interval=120,max_idle_secs=14400).Consider non-zero
#[serde(default = "...")]helpers matching the CLI defaults.🤖 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/config/types.rs` around lines 309 - 323, The `RoutingKeyOverrideConfig` serde defaults leave `eviction_interval_secs` and `max_idle_secs` at `0`, so `ManualPolicy::with_config` never enables TTL eviction when config files only set `enabled`. Update the `RoutingKeyOverrideConfig` field defaults to non-zero helper defaults that match the CLI behavior, and keep the `enabled`, `eviction_interval_secs`, `max_idle_secs`, and `assignment_mode` wiring consistent so sticky-map eviction is armed automatically for config-file users.model_gateway/src/routers/http/router.rs (1)
321-323: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
WorkerLoadGuardgating misses wrappedcache_aware/manualpolicies.
policy.name()returns"routing_key_override"when the policy is wrapped by the override, so the["cache_aware", "manual"]check fails and load tracking is skipped for an otherwise-eligible inner policy. Gate on the effective (unwrapped) policy name here (the same applies to the multipart path at Lines 598-600).🤖 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/routers/http/router.rs` around lines 321 - 323, The `WorkerLoadGuard` gating in `router.rs` is checking only `policy.name()`, which can return `routing_key_override` for wrapped policies and cause eligible `cache_aware` or `manual` policies to skip load tracking. Update the guard condition to inspect the effective unwrapped policy name before deciding whether to create `WorkerLoadGuard::new`, and apply the same fix in the multipart routing path that uses the same `["cache_aware", "manual"]` check.
🤖 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 162-202: The cache-aware/bucket setup and cleanup paths are still
looking at the RoutingKeyOverride wrapper instead of the underlying policy, so
the init/remove hooks can be skipped when sticky override is enabled. Update the
affected registry methods such as init_cache_aware_policy,
remove_worker_from_cache_aware, remove_worker_from_pd_cache_aware,
init_pd_cache_aware_policies, and init_pd_bucket_policies to call
Self::effective_policy(policy) before checking policy.name() and before any
CacheAwarePolicy downcast, so the inner policy is used consistently.
---
Duplicate comments:
In `@bindings/python/src/lib.rs`:
- Around line 760-769: The assignment_mode handling in the routing_key_override
setup silently falls back to Random for unknown values, unlike convert_policy’s
validated Manual branch. Update Router construction in lib.rs to reuse the same
validation logic or shared helper used by convert_policy so invalid
assignment_mode strings return ConfigError::InvalidValue instead of being
coerced. Keep the mapping aligned for the assignment_mode field and the
config::ManualAssignmentMode conversion.
In `@model_gateway/src/config/types.rs`:
- Around line 309-323: The `RoutingKeyOverrideConfig` serde defaults leave
`eviction_interval_secs` and `max_idle_secs` at `0`, so
`ManualPolicy::with_config` never enables TTL eviction when config files only
set `enabled`. Update the `RoutingKeyOverrideConfig` field defaults to non-zero
helper defaults that match the CLI behavior, and keep the `enabled`,
`eviction_interval_secs`, `max_idle_secs`, and `assignment_mode` wiring
consistent so sticky-map eviction is armed automatically for config-file users.
In `@model_gateway/src/routers/http/router.rs`:
- Around line 321-323: The `WorkerLoadGuard` gating in `router.rs` is checking
only `policy.name()`, which can return `routing_key_override` for wrapped
policies and cause eligible `cache_aware` or `manual` policies to skip load
tracking. Update the guard condition to inspect the effective unwrapped policy
name before deciding whether to create `WorkerLoadGuard::new`, and apply the
same fix in the multipart routing path that uses the same `["cache_aware",
"manual"]` check.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 22583eab-e49e-4fac-b43e-b6cb8e6b713e
📒 Files selected for processing (14)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pymodel_gateway/src/app_context.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/policies/manual.rsmodel_gateway/src/policies/mod.rsmodel_gateway/src/policies/registry.rsmodel_gateway/src/policies/routing_key_override.rsmodel_gateway/src/routers/grpc/common/stages/worker_selection.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/service_discovery.rs
Opt-in per-request sticky routing: when a request carries X-SMG-Routing-Key, the selection path consults a shared manual sticky-map selector instead of the configured policy (cache_aware, etc.), pinning the key to a worker; without the header the configured policy decides. Off by default. - WorkerLeg (Single/Prefill/Decode) on SelectWorkerInfo; ManualPolicy namespaces its sticky entries per leg so PD prefill/decode stick independently (Single is unchanged and allocation-free). - PolicyRegistry owns the shared sticky selector and a select_worker(policy, ..) that applies the override, skipping policies that already honor the key (manual, consistent_hashing). Policies are never wrapped, so policy identity, downcast injection, and load-aware detection are untouched. All gRPC and HTTP selection sites route through it. - RoutingKeyOverrideConfig eviction knobs default to the manual values so config-file users with only `enabled: true` still get TTL eviction. Enabled via --routing-key-override (CLI), routing_key_override=True (Python), or config. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
4b49a2c to
e1d2125
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e1d212529c
ℹ️ 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 idx = self.policy_registry.select_worker( | ||
| &policy, |
There was a problem hiding this comment.
With --routing-key-override --assignment-mode min_load or min_group on the regular HTTP router and an underlying policy such as round_robin, this changed call can return a worker selected by the override's ManualPolicy, but route_typed_request_once still creates WorkerLoadGuard only when the underlying policy.name() is cache_aware or manual. Fresh evidence: this version moved override selection into PolicyRegistry::select_worker here, so the policy name remains round_robin; as a result Worker::load() and routing_key_load() are never incremented for keyed HTTP requests, making min_load/min_group decisions use stale counters and defeating the configured assignment mode.
Useful? React with 👍 / 👎.
| if let Some(sticky) = self.routing_key_sticky.as_ref() { | ||
| if Self::routing_key_override_applies(policy.name()) | ||
| && extract_routing_key(info.headers).is_some() | ||
| { | ||
| return sticky.select_worker(workers, info); |
There was a problem hiding this comment.
Because the override uses one shared ManualPolicy and calls it without any model namespace, all models share the same sticky entry for a bare X-SMG-Routing-Key (only the PD leg is namespaced in ManualPolicy). In a multi-model router, using the same session key against more than two models makes the manual policy's two-candidate URL list evict another model's worker URL, so a later request for that earlier model becomes an occupied miss and can move to a different healthy worker, breaking the sticky-routing guarantee; please include the model in the override key or keep per-model sticky selectors.
Useful? React with 👍 / 👎.
2 participants
Closes #1843.
What
Opt-in per-request sticky routing: when a request carries
X-SMG-Routing-Key, it is pinned to a worker using manual sticky-map semantics, overriding the model's configured policy (e.g.cache_aware). Without the header, the configured policy decides — unchanged behavior. Off by default.This lets operators keep
cache_awareas the default while consumers (RL samplers, agentic multi-turn rollouts) opt into session affinity per request, instead of committing each model tomanualvscache_awareahead of time.How
RoutingKeyOverridedecorator wraps the model's policy plus an internalManualPolicy; key present → sticky map, absent → inner policy. Non-selection trait methods delegate to the inner policy so a wrapped load-aware policy still works.WorkerLeg(Single/Prefill/Decode) added toSelectWorkerInfo;ManualPolicynamespaces its sticky entries per leg so PD prefill/decode stick independently (Singleis byte-identical to before). The gRPC PD path tags each leg.PolicyRegistry::with_overridewraps key-ignoring policies (cache_aware,least_load,power_of_two,round_robin,random,bucket) and skips key-native ones (manual,consistent_hashing,prefix_hash,passthrough). The decorator is made transparent to the registry'sCacheAwarePolicymonitor/load-receiver injection and load-aware detection (the one non-obvious part).--routing-key-override(CLI),routing_key_override=True(Python), orrouting_key_override: { enabled: true }(config). Reuses the existing manual eviction/idle/assignment knobs for the sticky map.Testing
WorkerLegdefault,ManualPolicyleg namespacing,RoutingKeyOverride(key→sticky / no-key→inner / per-leg independence), registry wrap/skip/disabled, and a load-aware transparency test (a wrappedpower_of_twois still detected + injected).cargo build/cargo +nightly fmt/cargo clippy --all-targetsclean acrosssmg,smg-python,smg-golang.Scope (v1)
Routing-key only (not
X-SMG-Target-Worker), a global flag (not per-model), manual semantics (no consistent-hash strategy option) — all additive later.Summary by CodeRabbit
New Features
X-SMG-Routing-Keyto stay consistent across supported routing policies.Bug Fixes