diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 41361eee8..4cb7be1ed 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -13,7 +13,7 @@ use crate::config::{ default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs, resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config, DiscoveryBackend, K8sDiscoveryConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, - ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, + ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, }; /// `sgl-router` — slim KV-aware OpenAI-compatible router for SGLang workers. @@ -70,6 +70,26 @@ pub struct Cli { #[arg(long)] pub balance_rel_threshold: Option, + // ---- sticky-session policy (only used by `--policy sticky`) ---- + /// Request header carrying the routing key for sticky-session routing. + /// Defaults to `x-sgl-routing-key` when `--policy sticky` is set. + #[arg(long)] + pub routing_key_header: Option, + /// Policy used to select a worker for requests with no routing key, and + /// to pick the initial worker when a new key is first seen. One of + /// `round_robin` / `random` / `power_of_two` / `load_based`. Defaults + /// to `round_robin`. + #[arg(long, value_enum)] + pub sticky_fallback_policy: Option, + /// Evict a sticky assignment after it has been idle (unreferenced) this + /// many seconds. Defaults to 600. + #[arg(long)] + pub sticky_idle_secs: Option, + /// Wall-clock cadence of the sticky idle-eviction sweep, in seconds. + /// Defaults to 60. + #[arg(long)] + pub sticky_eviction_interval_secs: Option, + // ---- discovery: static ---- /// Static worker URLs (space-separated or repeated). Mutually /// exclusive with `--service-discovery`. @@ -145,6 +165,67 @@ impl Cli { )); } + let tuned_sticky = self.routing_key_header.is_some() + || self.sticky_fallback_policy.is_some() + || self.sticky_idle_secs.is_some() + || self.sticky_eviction_interval_secs.is_some(); + if tuned_sticky && self.policy != PolicyKind::Sticky { + return Err(anyhow!( + "--routing-key-header / --sticky-fallback-policy / --sticky-idle-secs / \ + --sticky-eviction-interval-secs require --policy sticky" + )); + } + + // Build (and validate) the sticky config exactly when the sticky + // policy is selected. The header name must parse as an HTTP header + // name so a typo fails at startup rather than silently never + // matching any request header; the fallback must be a + // dependency-free policy the factory can build standalone. + let sticky = if self.policy == PolicyKind::Sticky { + let d = StickyConfig::default(); + let header_name = self.routing_key_header.unwrap_or(d.header_name); + axum::http::HeaderName::try_from(header_name.as_str()).map_err(|e| { + anyhow!("--routing-key-header {header_name:?} is not a valid HTTP header name: {e}") + })?; + let fallback_policy = self.sticky_fallback_policy.unwrap_or(d.fallback_policy); + if matches!( + fallback_policy, + PolicyKind::Sticky | PolicyKind::CacheAwareZmq + ) { + return Err(anyhow!( + "--sticky-fallback-policy must be one of round_robin / random / \ + power_of_two / load_based; cache_aware_zmq and sticky are not allowed" + )); + } + let idle_secs = self.sticky_idle_secs.unwrap_or(d.idle_secs); + let eviction_interval_secs = self + .sticky_eviction_interval_secs + .unwrap_or(d.eviction_interval_secs); + // Reject zero durations: `--sticky-eviction-interval-secs 0` would + // panic `tokio::time::interval` at startup, and `--sticky-idle-secs + // 0` would evict every assignment on the next sweep (defeating + // stickiness entirely). Fail fast with a clear message instead. + if eviction_interval_secs == 0 { + return Err(anyhow!( + "--sticky-eviction-interval-secs must be greater than 0" + )); + } + if idle_secs == 0 { + return Err(anyhow!( + "--sticky-idle-secs must be greater than 0 (0 would evict every \ + assignment immediately, defeating sticky routing)" + )); + } + Some(StickyConfig { + header_name, + fallback_policy, + idle_secs, + eviction_interval_secs, + }) + } else { + None + }; + let circuit_breaker = self.cb_threshold.map(|threshold| CircuitBreakerConfig { threshold, cool_down_secs: self.cb_cool_down_secs.unwrap_or_else(default_cb_cool_down), @@ -185,6 +266,7 @@ impl Cli { policy: self.policy, circuit_breaker, cache_aware, + sticky, }, discovery, proxy: ProxyConfig { @@ -721,4 +803,160 @@ mod tests { assert_eq!(c.proxy.request_timeout_secs, 120); assert_eq!(c.active_load.stale_request_timeout_secs, 240); } + + #[test] + fn sticky_policy_defaults_header_and_tuning() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + ])) + .unwrap(); + assert_eq!(c.model.policy, PolicyKind::Sticky); + let s = c.model.sticky.expect("sticky config built"); + assert_eq!(s.header_name, "x-sgl-routing-key"); + assert_eq!(s.fallback_policy, PolicyKind::RoundRobin); + assert_eq!(s.idle_secs, 600); + assert_eq!(s.eviction_interval_secs, 60); + } + + #[test] + fn sticky_flags_override_defaults() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + "--routing-key-header", + "x-session-id", + "--sticky-fallback-policy", + "load_based", + "--sticky-idle-secs", + "120", + "--sticky-eviction-interval-secs", + "15", + ])) + .unwrap(); + let s = c.model.sticky.expect("sticky config built"); + assert_eq!(s.header_name, "x-session-id"); + assert_eq!(s.fallback_policy, PolicyKind::LoadBased); + assert_eq!(s.idle_secs, 120); + assert_eq!(s.eviction_interval_secs, 15); + } + + #[test] + fn non_sticky_policy_leaves_sticky_none() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "round_robin", + ])) + .unwrap(); + assert!(c.model.sticky.is_none()); + } + + #[test] + fn rejects_sticky_flags_without_sticky_policy() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--routing-key-header", + "x-session-id", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("require --policy sticky"), "got: {err}"); + } + + #[test] + fn rejects_invalid_routing_key_header() { + // A space is not a legal HTTP header-name character. + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + "--routing-key-header", + "bad header", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("not a valid HTTP header name"), "got: {err}"); + } + + #[test] + fn rejects_cache_aware_zmq_as_sticky_fallback() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + "--sticky-fallback-policy", + "cache_aware_zmq", + ])) + .unwrap_err() + .to_string(); + assert!( + err.contains("--sticky-fallback-policy must be one of"), + "got: {err}" + ); + } + + #[test] + fn rejects_sticky_as_sticky_fallback() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + "--sticky-fallback-policy", + "sticky", + ])) + .unwrap_err() + .to_string(); + assert!( + err.contains("--sticky-fallback-policy must be one of"), + "got: {err}" + ); + } + + /// A zero eviction interval would panic `tokio::time::interval` at + /// startup — reject it at config-build time with a clear message. + #[test] + fn rejects_zero_sticky_eviction_interval() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + "--sticky-eviction-interval-secs", + "0", + ])) + .unwrap_err() + .to_string(); + assert!( + err.contains("--sticky-eviction-interval-secs must be greater than 0"), + "got: {err}" + ); + } + + #[test] + fn rejects_zero_sticky_idle() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "sticky", + "--sticky-idle-secs", + "0", + ])) + .unwrap_err() + .to_string(); + assert!( + err.contains("--sticky-idle-secs must be greater than 0"), + "got: {err}" + ); + } } diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index baef6ddc8..df012c7c7 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -87,6 +87,7 @@ mod tests { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: urls.iter().map(|s| s.to_string()).collect(), diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 0af45461f..f4c414be7 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -70,7 +70,7 @@ impl Default for ActiveLoadConfig { /// policy factory. /// /// Accepted on the CLI (`--policy`) as `round_robin` / `random` / -/// `power_of_two` / `load_based` / `cache_aware_zmq`. +/// `power_of_two` / `load_based` / `cache_aware_zmq` / `sticky`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)] pub enum PolicyKind { #[default] @@ -88,6 +88,13 @@ pub enum PolicyKind { /// lives on `ModelConfig::cache_aware`. #[value(name = "cache_aware_zmq")] CacheAwareZmq, + /// Sticky-session routing: pins a routing key (read from a + /// configurable request header) to a worker via an in-memory map, so + /// stateful sessions land on the same backend. Tuning — header name, + /// keyless-fallback policy, and TTL eviction — lives on + /// `ModelConfig::sticky`. + #[value(name = "sticky")] + Sticky, } #[derive(Debug, Clone)] @@ -143,6 +150,11 @@ pub struct ModelConfig { /// `policy = "cache_aware_zmq"`. `None` falls back to defaults at /// policy construction time. pub cache_aware: Option, + /// Tuning for the sticky-session policy. `Some` exactly when + /// `policy = "sticky"` (built by [`crate::config::cli::Cli::into_config`]). + /// The chat handler reads `sticky.header_name` to populate + /// [`crate::policies::SelectionContext::routing_key`]. + pub sticky: Option, } /// Per-model cache-aware-ZMQ tuning. @@ -184,6 +196,52 @@ fn default_balance_rel() -> f32 { 1.1 } +/// Default routing-key header for the sticky policy. The `x-sgl-` prefix +/// matches the router's other emitted/consumed metadata headers +/// (`x-sgl-decode-url`, `x-sgl-router-error-code`). +pub const DEFAULT_STICKY_HEADER: &str = "x-sgl-routing-key"; + +/// Per-model sticky-session tuning. Built from the `--routing-key-header` +/// / `--sticky-*` flags by [`crate::config::cli::Cli::into_config`], which +/// also validates that `header_name` parses as an HTTP header name and +/// that `fallback_policy` is one of the dependency-free policies. +#[derive(Debug, Clone)] +pub struct StickyConfig { + /// Request header carrying the routing key. Validated to parse as a + /// `http::HeaderName` at config-build time. + pub header_name: String, + /// Policy used to pick a worker when a request has no routing key, and + /// to pick the initial worker when a new key is first seen. One of + /// `round_robin` / `random` / `power_of_two` / `load_based` — the + /// dependency-free policies the factory can build standalone (no + /// `HashTree` / tokenizer / ZMQ feed). `cache_aware_zmq` and `sticky` + /// are rejected at config-build time. + pub fallback_policy: PolicyKind, + /// Evict an assignment after it has been idle (unreferenced) this many + /// seconds. Bounds the map against unbounded routing-key cardinality. + pub idle_secs: u64, + /// Wall-clock cadence of the background eviction sweep. + pub eviction_interval_secs: u64, +} + +pub fn default_sticky_idle_secs() -> u64 { + 600 +} +pub fn default_sticky_eviction_interval_secs() -> u64 { + 60 +} + +impl Default for StickyConfig { + fn default() -> Self { + Self { + header_name: DEFAULT_STICKY_HEADER.to_string(), + fallback_policy: PolicyKind::RoundRobin, + idle_secs: default_sticky_idle_secs(), + eviction_interval_secs: default_sticky_eviction_interval_secs(), + } + } +} + #[derive(Debug, Clone)] pub struct CircuitBreakerConfig { /// Consecutive failures required before the breaker opens. Encoded diff --git a/experimental/sgl-router/src/policies/active_load.rs b/experimental/sgl-router/src/policies/active_load.rs index f87c9782b..4428dfa6d 100644 --- a/experimental/sgl-router/src/policies/active_load.rs +++ b/experimental/sgl-router/src/policies/active_load.rs @@ -418,6 +418,24 @@ impl ActiveLoadRegistry { /// `Arc` (cloned from the shared one held in /// `AppContext`). pub fn spawn_janitor(registry: Arc, interval: Duration) -> JanitorHandle { + spawn_sweeper(move || registry.sweep_stale(), interval, "active-load") +} + +/// Spawn a background task that calls `sweep` on a fixed cadence until its +/// [`JanitorHandle`] is cancelled or dropped. +/// +/// `sweep` returns the number of entries it removed; a non-zero count is +/// logged at info under `{label} janitor`. This is the shared engine +/// behind [`spawn_janitor`] (active-load stale-request reaping) and the +/// sticky policy's idle-assignment eviction — both want the same +/// cancel-aware ticker loop, differing only in what they sweep. +/// +/// `interval` is the wall-clock cadence. Missed ticks are skipped (a long +/// sweep does not cause a catch-up burst). +pub fn spawn_sweeper(mut sweep: F, interval: Duration, label: &'static str) -> JanitorHandle +where + F: FnMut() -> usize + Send + 'static, +{ let cancel = CancellationToken::new(); let cancel_for_task = cancel.clone(); let join = tokio::spawn(async move { @@ -427,16 +445,13 @@ pub fn spawn_janitor(registry: Arc, interval: Duration) -> J tokio::select! { biased; _ = cancel_for_task.cancelled() => { - tracing::debug!("active-load janitor: shutdown requested"); + tracing::debug!("{label} janitor: shutdown requested"); return; } _ = ticker.tick() => { - let n = registry.sweep_stale(); + let n = sweep(); if n > 0 { - tracing::info!( - swept = n, - "active-load janitor: removed stale requests", - ); + tracing::info!(swept = n, "{label} janitor: removed entries"); } } } diff --git a/experimental/sgl-router/src/policies/cache_aware_zmq.rs b/experimental/sgl-router/src/policies/cache_aware_zmq.rs index 9c6ecb7e2..e9e38ea01 100644 --- a/experimental/sgl-router/src/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/src/policies/cache_aware_zmq.rs @@ -360,6 +360,7 @@ mod tests { policy: crate::config::PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index 271614313..498f892f9 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -10,11 +10,41 @@ use crate::policies::{ power_of_two::PowerOfTwoChoicesPolicy, random::RandomPolicy, round_robin::RoundRobinPolicy, + sticky::StickyPolicy, Policy, PolicyRegistry, }; use crate::tokenizer::TokenizerRegistry; use anyhow::Result; use std::sync::Arc; +use std::time::Duration; + +/// Build a dependency-free policy for use as the sticky-session fallback +/// (keyless requests + initial pin of a new key). `Cli::into_config` +/// validates `--sticky-fallback-policy` to one of these four, so the +/// `CacheAwareZmq`/`Sticky` arms are never reached in practice. +fn build_sticky_fallback(kind: PolicyKind) -> Arc { + match kind { + PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), + PolicyKind::Random => Arc::new(RandomPolicy::new()), + PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), + PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()), + PolicyKind::CacheAwareZmq | PolicyKind::Sticky => { + unreachable!("sticky fallback is validated to be dependency-free in Cli::into_config") + } + } +} + +/// Construct a [`StickyPolicy`] from a model's `sticky` config (or +/// defaults). Shared by `build_policy` and the test shim so the duration +/// conversion + fallback wiring live in one place. +fn build_sticky(model: &ModelConfig) -> Arc { + let s = model.sticky.clone().unwrap_or_default(); + Arc::new(StickyPolicy::new( + Duration::from_secs(s.idle_secs), + Duration::from_secs(s.eviction_interval_secs), + build_sticky_fallback(s.fallback_policy), + )) +} /// Construct a policy for a single model from its [`ModelConfig`] and the /// process-shared `HashTree` + `TokenizerRegistry` + `BlockSizeOracle`. @@ -43,6 +73,7 @@ pub fn build_policy( block_size_oracle, )) } + PolicyKind::Sticky => build_sticky(model), } } @@ -69,6 +100,14 @@ pub fn build_policy_kind_only(kind: PolicyKind) -> Arc { BlockSizeOracle::new(), )) } + PolicyKind::Sticky => { + let s = crate::config::StickyConfig::default(); + Arc::new(StickyPolicy::new( + Duration::from_secs(s.idle_secs), + Duration::from_secs(s.eviction_interval_secs), + build_sticky_fallback(s.fallback_policy), + )) + } } } @@ -132,6 +171,7 @@ mod tests { policy, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], @@ -149,6 +189,7 @@ mod tests { let _ = build_policy_kind_only(PolicyKind::PowerOfTwo); let _ = build_policy_kind_only(PolicyKind::LoadBased); let _ = build_policy_kind_only(PolicyKind::CacheAwareZmq); + let _ = build_policy_kind_only(PolicyKind::Sticky); } #[test] @@ -191,4 +232,18 @@ mod tests { "expected LoadBasedPolicy debug repr, got: {dbg}", ); } + + #[test] + fn sticky_builds_via_factory() { + let cfg = cfg_with_model("modelA", PolicyKind::Sticky); + let tree = Arc::new(HashTree::new()); + let tokenizers = Arc::new(TokenizerRegistry::default()); + let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap(); + let p = reg.get(&ModelId("modelA".into())).unwrap(); + let dbg = format!("{p:?}"); + assert!( + dbg.contains("StickyPolicy"), + "expected StickyPolicy debug repr, got: {dbg}", + ); + } } diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index c5fa0243a..d4ec58201 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -10,6 +10,7 @@ pub mod power_of_two; pub mod random; pub mod registry; pub mod round_robin; +pub mod sticky; use crate::discovery::ModelId; use crate::server::metrics::MetricsRegistry; diff --git a/experimental/sgl-router/src/policies/sticky.rs b/experimental/sgl-router/src/policies/sticky.rs new file mode 100644 index 000000000..541ea7314 --- /dev/null +++ b/experimental/sgl-router/src/policies/sticky.rs @@ -0,0 +1,462 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Sticky-session routing policy. +//! +//! Pins a request's routing key — read from a configurable header into +//! [`SelectionContext::routing_key`] by the chat handler — to a consistent +//! worker via an in-memory map, so stateful sessions land on the same +//! backend. Unlike consistent hashing, this policy never redistributes +//! existing keys when a worker is *added*; a key is only remapped when its +//! assigned worker leaves the healthy candidate set. +//! +//! # Behavior +//! - **No routing key** → delegate to the configured `fallback` policy (no +//! pinning). This lets clients that don't send the header still be served. +//! - **Known key, worker healthy** → return the pinned worker. +//! - **New key, or pinned worker unhealthy** → pick a worker via `fallback` +//! and record the assignment. +//! +//! Worker identity is the worker URL (stable across discovery events). +//! +//! # Eviction +//! A background sweeper (shared engine with the active-load janitor, see +//! [`super::active_load::spawn_sweeper`]) removes assignments idle longer +//! than `idle`, bounding the map against unbounded routing-key cardinality. +//! The sweeper is spawned only when constructed inside a Tokio runtime; +//! unit tests use [`StickyPolicy::with_clock`] and drive eviction +//! deterministically via a `MockClock` + direct `sweep_expired`. +//! +//! # HA +//! This map is per-router-instance state, so it is NOT consistent across +//! multiple router replicas or across a failover. HA sticky routing needs +//! a stateless deterministic scheme (rendezvous / consistent hashing) and +//! is intentionally out of scope here. + +use std::sync::{Arc, OnceLock}; +use std::time::{Duration, Instant}; + +use dashmap::DashMap; + +use crate::policies::active_load::{spawn_sweeper, Clock, JanitorHandle, SystemTimeClock}; +use crate::policies::{Policy, SelectionContext}; +use crate::server::metrics::{MetricsRegistry, StickyOutcome}; +use crate::workers::Worker; + +/// One routing-key → worker pin, with the last time it was referenced (used +/// by the idle-eviction sweep). +#[derive(Debug)] +struct Assignment { + worker_url: String, + last_seen: Instant, +} + +/// Shared inner state. Held behind an `Arc` so the background sweeper can +/// reference the same map the `select` hot path mutates. +#[derive(Debug)] +struct StickyState { + assignments: DashMap, + clock: Arc, + idle: Duration, + /// Metrics sink. Set once via the `Policy::attach_metrics` hook + /// (production) — `None` until then, in which case recording is a no-op. + metrics: OnceLock>, +} + +impl StickyState { + /// Remove every assignment idle longer than `idle`. Returns the count + /// removed. Called on a fixed cadence by the background sweeper. + fn sweep_expired(&self) -> usize { + let now = self.clock.now(); + let mut removed = 0; + self.assignments.retain(|_key, a| { + let keep = now.saturating_duration_since(a.last_seen) <= self.idle; + if !keep { + removed += 1; + } + keep + }); + removed + } + + fn record(&self, outcome: StickyOutcome) { + if let Some(m) = self.metrics.get() { + m.record_sticky(outcome); + } + } +} + +/// Sticky-session policy. See the module docs for behavior and limitations. +pub struct StickyPolicy { + state: Arc, + /// Selector for keyless requests and for the initial pin of a new key. + fallback: Arc, + /// Background idle-eviction sweeper. `None` when constructed outside a + /// Tokio runtime (unit tests). Dropping it cancels the task, so the + /// sweeper lives exactly as long as the policy. + _janitor: Option, +} + +impl StickyPolicy { + /// Production constructor: monotonic `SystemTimeClock`, with a + /// background eviction sweeper spawned on `eviction_interval` cadence + /// (only if called inside a Tokio runtime — the factory runs inside + /// `main`'s runtime). + pub fn new(idle: Duration, eviction_interval: Duration, fallback: Arc) -> Self { + let state = Arc::new(StickyState { + assignments: DashMap::new(), + clock: Arc::new(SystemTimeClock), + idle, + metrics: OnceLock::new(), + }); + // `spawn_sweeper` needs a runtime; the factory builds policies inside + // `main`'s Tokio runtime. Guard so sync constructions (e.g. the + // factory's `build_policy_kind_only` test helper) don't panic. + let _janitor = if tokio::runtime::Handle::try_current().is_ok() { + let swept = Arc::clone(&state); + Some(spawn_sweeper( + move || swept.sweep_expired(), + eviction_interval, + "sticky-eviction", + )) + } else { + // Only reached by sync construction (test helpers). In production + // the factory builds policies inside `main`'s runtime, so the + // sweeper always spawns. Log it so a future off-runtime + // construction that silently disables eviction is greppable. + tracing::debug!( + "StickyPolicy constructed outside a Tokio runtime; idle eviction is disabled" + ); + None + }; + Self { + state, + fallback, + _janitor, + } + } + + /// Test constructor: injectable clock, no background sweeper. Tests + /// advance a `MockClock` and call [`Self::sweep_expired`] directly for + /// deterministic eviction coverage. + #[cfg(test)] + fn with_clock(idle: Duration, fallback: Arc, clock: Arc) -> Self { + Self { + state: Arc::new(StickyState { + assignments: DashMap::new(), + clock, + idle, + metrics: OnceLock::new(), + }), + fallback, + _janitor: None, + } + } + + #[cfg(test)] + fn sweep_expired(&self) -> usize { + self.state.sweep_expired() + } + + #[cfg(test)] + fn assignment_count(&self) -> usize { + self.state.assignments.len() + } +} + +impl Policy for StickyPolicy { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + let Some(key) = ctx.routing_key().filter(|k| !k.is_empty()) else { + self.state.record(StickyOutcome::NoRoutingKey); + return self.fallback.select(workers, ctx); + }; + + // Fast path: an existing pin whose worker is still in the healthy set. + let mut existing = false; + if let Some(mut entry) = self.state.assignments.get_mut(key) { + existing = true; + if let Some(worker) = workers.iter().find(|w| w.url == entry.worker_url).cloned() { + entry.last_seen = self.state.clock.now(); + drop(entry); // release the shard lock before recording + self.state.record(StickyOutcome::Hit); + return Some(worker); + } + // Pinned worker is no longer healthy — fall through to reassign. + drop(entry); + } + + // Vacant key, or the pinned worker dropped out: (re)assign via the + // fallback. The read-miss above and this insert are intentionally NOT + // atomic — the shard lock is released before `fallback.select` (which + // may do real work, e.g. `load_based`) so it is never held across an + // unrelated computation. Two requests racing the *same* fresh key may + // therefore both assign (last-writer-wins in the map; both may record + // `Assigned`). The scatter is transient and self-heals: the next + // request for that key hits the surviving pin. + let chosen = self.fallback.select(workers, ctx)?; + self.state.assignments.insert( + key.to_string(), + Assignment { + worker_url: chosen.url.clone(), + last_seen: self.state.clock.now(), + }, + ); + self.state.record(if existing { + StickyOutcome::Remap + } else { + StickyOutcome::Assigned + }); + Some(chosen) + } + + fn attach_metrics(&self, metrics: Arc) { + let _ = self.state.metrics.set(metrics); + } +} + +impl std::fmt::Debug for StickyPolicy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("StickyPolicy") + .field("fallback", &self.fallback) + .field("idle", &self.state.idle) + .field("assignments", &self.state.assignments.len()) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::round_robin::RoundRobinPolicy; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + fn fallback() -> Arc { + Arc::new(RoundRobinPolicy::new()) + } + + fn policy(idle_secs: u64) -> StickyPolicy { + let clock = Arc::new(crate::policies::active_load::MockClock::new(Instant::now())); + StickyPolicy::with_clock(Duration::from_secs(idle_secs), fallback(), clock) + } + + #[test] + fn empty_workers_returns_none() { + let model = ModelId("tiny".into()); + let p = policy(600); + let ctx = SelectionContext::with_routing_key(&model, None, Some("u1")); + assert!(p.select(&[], &ctx).is_none()); + } + + #[test] + fn keyless_request_delegates_to_fallback_without_pinning() { + let model = ModelId("tiny".into()); + let p = policy(600); + let workers = vec![worker("w0"), worker("w1")]; + // No routing key on the context. + let ctx = SelectionContext::new(&model, None); + assert!(p.select(&workers, &ctx).is_some()); + assert_eq!(p.assignment_count(), 0, "keyless request must not pin"); + } + + #[test] + fn same_key_sticks_to_same_worker() { + let model = ModelId("tiny".into()); + let p = policy(600); + let workers = vec![worker("w0"), worker("w1")]; + let ctx = SelectionContext::with_routing_key(&model, None, Some("u1")); + + let first = p.select(&workers, &ctx).unwrap(); + // Many repeats must all return the same worker (the hit path never + // consults the fallback, so this is independent of round-robin). + for _ in 0..10 { + let again = p.select(&workers, &ctx).unwrap(); + assert_eq!(again.id, first.id); + } + assert_eq!(p.assignment_count(), 1); + } + + #[test] + fn distinct_keys_get_independent_pins() { + let model = ModelId("tiny".into()); + let p = policy(600); + let workers = vec![worker("w0"), worker("w1")]; + + let ctx_a = SelectionContext::with_routing_key(&model, None, Some("a")); + let ctx_b = SelectionContext::with_routing_key(&model, None, Some("b")); + let a = p.select(&workers, &ctx_a).unwrap(); + let b = p.select(&workers, &ctx_b).unwrap(); + // Two keys are tracked independently (two map entries), and the + // round-robin fallback hands the two fresh keys distinct workers. + assert_ne!(a.id, b.id); + assert_eq!(p.assignment_count(), 2); + // The core property: each key independently stays on its own pin. + for _ in 0..5 { + assert_eq!(p.select(&workers, &ctx_a).unwrap().id, a.id); + assert_eq!(p.select(&workers, &ctx_b).unwrap().id, b.id); + } + } + + #[test] + fn adding_a_worker_does_not_redistribute_existing_key() { + let model = ModelId("tiny".into()); + let p = policy(600); + let w0 = worker("w0"); + let w1 = worker("w1"); + let ctx = SelectionContext::with_routing_key(&model, None, Some("u1")); + + let pinned = p.select(&[Arc::clone(&w0), Arc::clone(&w1)], &ctx).unwrap(); + // Scale up: a third worker joins. The existing key must stay pinned. + let w2 = worker("w2"); + let after = p + .select(&[Arc::clone(&w0), Arc::clone(&w1), w2], &ctx) + .unwrap(); + assert_eq!(after.id, pinned.id, "true-sticky: no redistribution on add"); + } + + #[test] + fn remaps_when_pinned_worker_becomes_unhealthy() { + let model = ModelId("tiny".into()); + let p = policy(600); + let w0 = worker("w0"); + let w1 = worker("w1"); + let ctx = SelectionContext::with_routing_key(&model, None, Some("u1")); + + let pinned = p.select(&[Arc::clone(&w0), Arc::clone(&w1)], &ctx).unwrap(); + // Drop the pinned worker from the healthy set; only the other remains. + let survivor = if pinned.id == w0.id { + Arc::clone(&w1) + } else { + Arc::clone(&w0) + }; + let remapped = p.select(&[Arc::clone(&survivor)], &ctx).unwrap(); + assert_eq!(remapped.id, survivor.id); + // The new pin sticks across subsequent calls. + let again = p.select(&[Arc::clone(&survivor)], &ctx).unwrap(); + assert_eq!(again.id, survivor.id); + } + + #[test] + fn sweep_evicts_idle_entries_keeps_fresh() { + let model = ModelId("tiny".into()); + let clock = Arc::new(crate::policies::active_load::MockClock::new(Instant::now())); + let p = StickyPolicy::with_clock(Duration::from_secs(10), fallback(), clock.clone()); + let workers = vec![worker("w0"), worker("w1")]; + + // Pin key "old" at t0. + let ctx_old = SelectionContext::with_routing_key(&model, None, Some("old")); + p.select(&workers, &ctx_old).unwrap(); + + // Advance 6s, pin key "new" at t6. + clock.advance(Duration::from_secs(6)); + let ctx_new = SelectionContext::with_routing_key(&model, None, Some("new")); + p.select(&workers, &ctx_new).unwrap(); + assert_eq!(p.assignment_count(), 2); + + // Advance to t11: "old" has been idle 11s (> 10), "new" idle 5s. + clock.advance(Duration::from_secs(5)); + assert_eq!(p.sweep_expired(), 1); + assert_eq!(p.assignment_count(), 1); + + // "new" survived and is still pinned. + assert!(p.select(&workers, &ctx_new).is_some()); + assert_eq!(p.assignment_count(), 1); + } + + #[test] + fn hit_refreshes_last_seen_so_active_key_is_not_evicted() { + let model = ModelId("tiny".into()); + let clock = Arc::new(crate::policies::active_load::MockClock::new(Instant::now())); + let p = StickyPolicy::with_clock(Duration::from_secs(10), fallback(), clock.clone()); + let workers = vec![worker("w0")]; + let ctx = SelectionContext::with_routing_key(&model, None, Some("u1")); + + p.select(&workers, &ctx).unwrap(); + // Keep referencing the key just under the idle window each step. + for _ in 0..5 { + clock.advance(Duration::from_secs(8)); + p.select(&workers, &ctx).unwrap(); // hit → refreshes last_seen + assert_eq!( + p.sweep_expired(), + 0, + "an actively-used key must not be evicted" + ); + } + assert_eq!(p.assignment_count(), 1); + } + + /// Exercises the production path: `new` (not `with_clock`) spawns the + /// real background sweeper because we are inside a Tokio runtime. Uses + /// sub-second idle + interval so the sweep fires within the test's + /// wall-time, proving `StickyPolicy::new` correctly wires `sweep_expired` + /// into the runtime sweeper. + #[tokio::test] + async fn background_sweeper_evicts_idle_entry_in_runtime() { + let model = ModelId("tiny".into()); + let p = StickyPolicy::new( + Duration::from_millis(20), + Duration::from_millis(10), + fallback(), + ); + let workers = vec![worker("w0")]; + let ctx = SelectionContext::with_routing_key(&model, None, Some("u1")); + p.select(&workers, &ctx).unwrap(); + assert_eq!(p.assignment_count(), 1); + + // Idle window is 20ms; wait well past it plus several sweep ticks. + tokio::time::sleep(Duration::from_millis(400)).await; + assert_eq!( + p.assignment_count(), + 0, + "background sweeper should have evicted the idle assignment" + ); + } + + /// Many concurrent first-touch requests for the SAME fresh key converge: + /// the map ends with exactly one pin and every subsequent select agrees + /// on it (the documented self-heal after the benign assign race). + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_first_touch_converges_to_one_pin() { + let p = Arc::new(StickyPolicy::new( + Duration::from_secs(3600), + Duration::from_secs(3600), + fallback(), + )); + let workers = Arc::new(vec![worker("w0"), worker("w1"), worker("w2")]); + let model = ModelId("tiny".into()); + + let mut handles = Vec::new(); + for _ in 0..32 { + let p = Arc::clone(&p); + let workers = Arc::clone(&workers); + let model = model.clone(); + handles.push(tokio::spawn(async move { + let ctx = SelectionContext::with_routing_key(&model, None, Some("race")); + p.select(&workers[..], &ctx).map(|w| w.id.clone()) + })); + } + for h in handles { + h.await.unwrap().unwrap(); + } + + assert_eq!( + p.assignment_count(), + 1, + "concurrent first-touch must converge to a single pin" + ); + let ctx = SelectionContext::with_routing_key(&model, None, Some("race")); + let pinned = p.select(&workers[..], &ctx).unwrap().id.clone(); + for _ in 0..10 { + assert_eq!(p.select(&workers[..], &ctx).unwrap().id, pinned); + } + } +} diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 2ff31bf1a..04778aa04 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -111,6 +111,7 @@ impl AppContext { policy: crate::config::PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index 678b10d1e..2ecb22672 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -23,6 +23,7 @@ //! | `sgl_router_active_load` | Gauge | `worker_url`, `kind` | //! | `sgl_router_stale_requests_total` | Counter | `outcome` | //! | `sgl_router_decode_affinity_total` | Counter | `outcome` | +//! | `sgl_router_sticky_total` | Counter | `outcome` | //! //! The exposition is text/plain; version=0.0.4 per the Prometheus spec. @@ -97,6 +98,31 @@ impl DecodeAffinityOutcome { } } +/// Sticky-policy selection outcome — see `StickyPolicy::select` for the +/// four branches. +#[derive(Debug, Clone, Copy)] +pub enum StickyOutcome { + /// Routing key found and its assigned worker is still healthy. + Hit, + /// Routing key seen for the first time — a worker was assigned. + Assigned, + /// Routing key's assigned worker left the healthy set — remapped. + Remap, + /// Request carried no routing key — delegated to the fallback policy. + NoRoutingKey, +} + +impl StickyOutcome { + fn as_str(self) -> &'static str { + match self { + Self::Hit => "hit", + Self::Assigned => "assigned", + Self::Remap => "remap", + Self::NoRoutingKey => "no_routing_key", + } + } +} + /// Stale-request outcome label. #[derive(Debug, Clone, Copy)] pub enum StaleRequestOutcome { @@ -136,6 +162,7 @@ pub struct MetricsRegistry { active_load: Mutex>>, stale_requests_total: Mutex>>, decode_affinity_total: Mutex>>, + sticky_total: Mutex>>, } #[derive(Debug, Hash, Eq, PartialEq, Clone)] @@ -265,6 +292,17 @@ impl MetricsRegistry { counter.fetch_add(1, Ordering::Relaxed); } + /// Bump `sgl_router_sticky_total{outcome}`. + pub fn record_sticky(&self, outcome: StickyOutcome) { + let mut guard = self.sticky_total.lock(); + let counter = guard + .entry(outcome.as_str()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + drop(guard); + counter.fetch_add(1, Ordering::Relaxed); + } + /// Render the registry as a Prometheus 0.0.4 exposition-format string. pub fn render(&self) -> String { let mut out = String::new(); @@ -398,6 +436,25 @@ impl MetricsRegistry { } drop(guard); + // sticky_total + out.push_str( + "# HELP sgl_router_sticky_total Sticky-session selection outcomes from StickyPolicy.\n", + ); + out.push_str("# TYPE sgl_router_sticky_total counter\n"); + let guard = self.sticky_total.lock(); + let mut entries: Vec<(&&str, u64)> = guard + .iter() + .map(|(k, v)| (k, v.load(Ordering::Relaxed))) + .collect(); + entries.sort_by_key(|e| *e.0); + for (outcome, value) in entries { + out.push_str(&format!( + "sgl_router_sticky_total{{outcome=\"{}\"}} {}\n", + outcome, value, + )); + } + drop(guard); + out } } @@ -433,6 +490,7 @@ mod tests { assert!(out.contains("# TYPE sgl_router_active_load gauge")); assert!(out.contains("# TYPE sgl_router_stale_requests_total counter")); assert!(out.contains("# TYPE sgl_router_decode_affinity_total counter")); + assert!(out.contains("# TYPE sgl_router_sticky_total counter")); } #[test] @@ -518,6 +576,21 @@ mod tests { .contains(r#"sgl_router_decode_affinity_total{outcome="fallback_load_imbalance"} 1"#,)); } + #[test] + fn sticky_counter_emits_all_outcomes() { + let reg = MetricsRegistry::new(); + reg.record_sticky(StickyOutcome::Hit); + reg.record_sticky(StickyOutcome::Hit); + reg.record_sticky(StickyOutcome::Assigned); + reg.record_sticky(StickyOutcome::Remap); + reg.record_sticky(StickyOutcome::NoRoutingKey); + let out = reg.render(); + assert!(out.contains(r#"sgl_router_sticky_total{outcome="hit"} 2"#)); + assert!(out.contains(r#"sgl_router_sticky_total{outcome="assigned"} 1"#)); + assert!(out.contains(r#"sgl_router_sticky_total{outcome="remap"} 1"#)); + assert!(out.contains(r#"sgl_router_sticky_total{outcome="no_routing_key"} 1"#)); + } + #[test] fn label_values_escape_quotes_and_backslashes() { let reg = MetricsRegistry::new(); diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index 9fb5aa036..b3f7ab1ec 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -107,7 +107,19 @@ pub async fn chat_completions( .policies .get(&model_id) .ok_or_else(|| ApiError::ModelNotFound(model_str.clone()))?; - let selection_ctx = SelectionContext::new(&model_id, Some(&body)); + // Sticky-session routing key. When the sticky policy is configured, + // read the routing key from the operator-chosen header into the + // selection context; the policy pins it to a worker. Other policies + // leave `routing_key` `None` and ignore it. + let routing_key = ctx + .config + .model + .sticky + .as_ref() + .and_then(|s| headers.get(s.header_name.as_str())) + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()); + let selection_ctx = SelectionContext::with_routing_key(&model_id, Some(&body), routing_key); let worker = policy .select(&workers, &selection_ctx) diff --git a/experimental/sgl-router/src/server/routes/models.rs b/experimental/sgl-router/src/server/routes/models.rs index 5d66cb9ed..10b635efc 100644 --- a/experimental/sgl-router/src/server/routes/models.rs +++ b/experimental/sgl-router/src/server/routes/models.rs @@ -53,6 +53,7 @@ mod tests { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }; let app = crate::server::app::build_router(std::sync::Arc::new(ctx)); let res = app diff --git a/experimental/sgl-router/src/server/routes/tokenize.rs b/experimental/sgl-router/src/server/routes/tokenize.rs index 8be123bbd..f45f7bea4 100644 --- a/experimental/sgl-router/src/server/routes/tokenize.rs +++ b/experimental/sgl-router/src/server/routes/tokenize.rs @@ -120,6 +120,7 @@ mod tests { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/tokenizer/mod.rs b/experimental/sgl-router/src/tokenizer/mod.rs index fcb7bdeac..b862f98a0 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -59,6 +59,7 @@ mod tests { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index 1e8670057..89e2cd78f 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -304,6 +304,7 @@ mod tests { cool_down_secs, }), cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://test:30000".into()], diff --git a/experimental/sgl-router/tests/component/discovery/static_urls.rs b/experimental/sgl-router/tests/component/discovery/static_urls.rs index 8c7ae5fc5..3364c98e5 100644 --- a/experimental/sgl-router/tests/component/discovery/static_urls.rs +++ b/experimental/sgl-router/tests/component/discovery/static_urls.rs @@ -132,6 +132,7 @@ async fn static_urls_pd_role_resolved_end_to_end() { policy: sgl_router::config::PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec![url.clone()], diff --git a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs index a5585f819..7c6d81b82 100644 --- a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs @@ -71,6 +71,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() { policy: sgl_router::config::PolicyKind::CacheAwareZmq, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: sgl_router::config::DiscoveryBackend::StaticUrls( sgl_router::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 7347ee18b..7425052b4 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -35,6 +35,7 @@ fn config_for(_worker_url: &str) -> Config { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/failover.rs b/experimental/sgl-router/tests/proxy/failover.rs index 07920d806..8ca7aeab4 100644 --- a/experimental/sgl-router/tests/proxy/failover.rs +++ b/experimental/sgl-router/tests/proxy/failover.rs @@ -40,6 +40,7 @@ async fn failover_when_one_worker_dies() { cool_down_secs: 30, }), cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()], diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 8c5145299..2ec959778 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -46,6 +46,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/header_forwarding.rs b/experimental/sgl-router/tests/proxy/header_forwarding.rs index 79b86bc02..06c3e1284 100644 --- a/experimental/sgl-router/tests/proxy/header_forwarding.rs +++ b/experimental/sgl-router/tests/proxy/header_forwarding.rs @@ -33,6 +33,7 @@ async fn forwards_whitelisted_headers_strips_others() { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/main.rs b/experimental/sgl-router/tests/proxy/main.rs index 9507099bb..ff0da7fc9 100644 --- a/experimental/sgl-router/tests/proxy/main.rs +++ b/experimental/sgl-router/tests/proxy/main.rs @@ -16,4 +16,5 @@ mod graceful_shutdown; mod header_forwarding; mod pd_bootstrap_injection; mod pd_pool_isolation; +mod sticky_routing; mod timeout; diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index 059494ec4..9a1ce40db 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -48,6 +48,7 @@ fn config() -> Config { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index 5d6fdec99..f0e4a525b 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -47,6 +47,7 @@ fn config() -> Config { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/sticky_routing.rs b/experimental/sgl-router/tests/proxy/sticky_routing.rs new file mode 100644 index 000000000..647aa87af --- /dev/null +++ b/experimental/sgl-router/tests/proxy/sticky_routing.rs @@ -0,0 +1,255 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! End-to-end coverage for `policy = "sticky"`: a routing key read from the +//! operator-configured header pins a session to one worker, and the +//! `sgl_router_sticky_total` outcomes are recorded. Runs against two +//! `MockWorker` backends (CPU-only, no GPU). + +use sgl_router::config::{ + ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind, + ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, +}; +use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; +use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry; +use sgl_router::proxy::Proxy; +use sgl_router::server::app::build_router; +use sgl_router::server::app_context::AppContext; +use sgl_router::tokenizer::TokenizerRegistry; +use sgl_router::workers::WorkerRegistry; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use std::sync::Arc; +use std::time::Duration; +use tower::ServiceExt; + +use crate::common::mock_worker::MockWorker; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +/// Build an `AppContext` running the sticky policy over the given worker +/// URLs, reading the routing key from `header_name`. Eviction is pushed far +/// out so the background sweeper never fires mid-test. +fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc { + let cfg = Config { + server: ServerConfig { + host: "0".into(), + port: 0, + }, + observability: ObservabilityConfig::default(), + model: ModelConfig { + id: "tiny".into(), + tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + policy: PolicyKind::Sticky, + circuit_breaker: None, + cache_aware: None, + sticky: Some(StickyConfig { + header_name: header_name.to_string(), + fallback_policy: PolicyKind::RoundRobin, + idle_secs: 3600, + eviction_interval_secs: 3600, + }), + }, + discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { + urls: vec!["http://placeholder:0".into()], + }), + proxy: ProxyConfig::default(), + active_load: ActiveLoadConfig::default(), + }; + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); + let registry = Arc::new(WorkerRegistry::default()); + for (i, url) in worker_urls.iter().enumerate() { + let _ = registry.add(WorkerSpec { + id: WorkerId(format!("w{i}")), + url: url.clone(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + }); + } + let policies = Arc::new(build_policy_registry(&cfg).unwrap()); + let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()); + Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) +} + +fn chat_request(header: Option<(&str, &str)>) -> Request { + let mut builder = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json"); + if let Some((name, value)) = header { + builder = builder.header(name, value); + } + builder + .body(Body::from( + serde_json::to_vec(&serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": false + })) + .unwrap(), + )) + .unwrap() +} + +/// Parse `sgl_router_requests_total{...,outcome="success"} N` lines into a +/// map of worker_url -> success count. +fn success_counts(metrics: &str) -> std::collections::HashMap { + let mut counts = std::collections::HashMap::new(); + for line in metrics.lines() { + let Some(rest) = line.strip_prefix("sgl_router_requests_total{") else { + continue; + }; + if !rest.contains(r#"outcome="success""#) { + continue; + } + let Some(url_start) = rest.find(r#"worker_url=""#) else { + continue; + }; + let after = &rest[url_start + r#"worker_url=""#.len()..]; + let Some(url_end) = after.find('"') else { + continue; + }; + let url = after[..url_end].to_string(); + let value: u64 = line + .rsplit(' ') + .next() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + *counts.entry(url).or_insert(0) += value; + } + counts +} + +/// Read `sgl_router_sticky_total{outcome=""} N`. +fn sticky_count(metrics: &str, outcome: &str) -> u64 { + let needle = format!(r#"sgl_router_sticky_total{{outcome="{outcome}"}} "#); + metrics + .lines() + .find_map(|l| l.strip_prefix(&needle)) + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(0) +} + +#[tokio::test] +async fn same_routing_key_pins_to_one_worker() { + let w0 = MockWorker::start(vec![]).await; + let w1 = MockWorker::start(vec![]).await; + let ctx = build_sticky_ctx("x-sgl-routing-key", &[w0.url.clone(), w1.url.clone()]); + let app = build_router(ctx.clone()); + + const N: usize = 5; + for _ in 0..N { + let res = app + .clone() + .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + } + + let metrics = ctx.metrics.render(); + let counts = success_counts(&metrics); + let total: u64 = counts.values().sum(); + assert_eq!(total, N as u64, "all requests should succeed: {counts:?}"); + let pinned: Vec<_> = counts.iter().filter(|(_, &c)| c > 0).collect(); + assert_eq!( + pinned.len(), + 1, + "all same-key requests must hit exactly one worker: {counts:?}" + ); + assert_eq!(*pinned[0].1, N as u64); + + // One assignment, the rest hits. + assert_eq!(sticky_count(&metrics, "assigned"), 1, "metrics:\n{metrics}"); + assert_eq!(sticky_count(&metrics, "hit"), (N - 1) as u64); +} + +#[tokio::test] +async fn remaps_to_survivor_when_pinned_worker_is_removed() { + let w0 = MockWorker::start(vec![]).await; + let w1 = MockWorker::start(vec![]).await; + let worker_urls = vec![w0.url.clone(), w1.url.clone()]; + let ctx = build_sticky_ctx("x-sgl-routing-key", &worker_urls); + let app = build_router(ctx.clone()); + + // Pin the key, then discover which worker it landed on. + let res = app + .clone() + .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + let pinned_url = success_counts(&ctx.metrics.render()) + .into_iter() + .find(|(_, c)| *c > 0) + .map(|(url, _)| url) + .expect("a worker should have served the first request"); + let pinned_idx = worker_urls.iter().position(|u| *u == pinned_url).unwrap(); + let survivor_url = worker_urls[1 - pinned_idx].clone(); + + // Remove the pinned worker from the registry; the next same-key request + // must remap to the survivor (not fail). + ctx.registry.remove(&WorkerId(format!("w{pinned_idx}"))); + let res = app + .clone() + .oneshot(chat_request(Some(("x-sgl-routing-key", "alice")))) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + let metrics = ctx.metrics.render(); + assert_eq!(sticky_count(&metrics, "remap"), 1, "{metrics}"); + // The survivor served the second request. + let counts = success_counts(&metrics); + assert_eq!(counts.get(&survivor_url).copied().unwrap_or(0), 1); +} + +#[tokio::test] +async fn keyless_request_is_served_via_fallback() { + let w0 = MockWorker::start(vec![]).await; + let w1 = MockWorker::start(vec![]).await; + let ctx = build_sticky_ctx("x-sgl-routing-key", &[w0.url.clone(), w1.url.clone()]); + let app = build_router(ctx.clone()); + + // No routing-key header at all. + let res = app.clone().oneshot(chat_request(None)).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + let metrics = ctx.metrics.render(); + assert_eq!(sticky_count(&metrics, "no_routing_key"), 1, "{metrics}"); + assert_eq!(sticky_count(&metrics, "assigned"), 0); +} + +#[tokio::test] +async fn only_the_configured_header_name_is_honored() { + // Router configured to read the key from `x-session-id`. + let w0 = MockWorker::start(vec![]).await; + let w1 = MockWorker::start(vec![]).await; + let ctx = build_sticky_ctx("x-session-id", &[w0.url.clone(), w1.url.clone()]); + let app = build_router(ctx.clone()); + + // A request using the configured header pins (assigned), and a repeat hits. + for _ in 0..2 { + let res = app + .clone() + .oneshot(chat_request(Some(("x-session-id", "s-1")))) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + } + // A request carrying the *default* header (not the configured one) must + // be treated as keyless — proving the header name is dynamic, not baked in. + let res = app + .clone() + .oneshot(chat_request(Some(("x-sgl-routing-key", "s-1")))) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + let metrics = ctx.metrics.render(); + assert_eq!(sticky_count(&metrics, "assigned"), 1, "{metrics}"); + assert_eq!(sticky_count(&metrics, "hit"), 1); + assert_eq!(sticky_count(&metrics, "no_routing_key"), 1); +} diff --git a/experimental/sgl-router/tests/proxy/timeout.rs b/experimental/sgl-router/tests/proxy/timeout.rs index c5426d871..03851d626 100644 --- a/experimental/sgl-router/tests/proxy/timeout.rs +++ b/experimental/sgl-router/tests/proxy/timeout.rs @@ -40,6 +40,7 @@ fn config(_worker_url: &str) -> Config { policy: PolicyKind::RoundRobin, circuit_breaker: None, cache_aware: None, + sticky: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()],