diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index f360f78d1..1bca943c4 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -11,10 +11,11 @@ use std::num::NonZeroU32; use crate::config::{ default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs, - resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config, - DiscoveryBackend, EligibilityConfig, FilterKind, FusedTerm, K8sDiscoveryConfig, - KvIndexerEndpointConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig, - ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE, + resolve_mode, ActiveLoadConfig, AffinityConfig, AffinityMode, CacheAwareConfig, + CircuitBreakerConfig, Config, DiscoveryBackend, EligibilityConfig, FilterKind, FusedTerm, + K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig, ObservabilityConfig, + PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode, StaticUrlsDiscoveryConfig, + StickyConfig, StickyFallbackKind, DEFAULT_FUSE, }; const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100; @@ -63,7 +64,7 @@ pub struct Cli { #[arg(long)] pub cb_cool_down_secs: Option, - // ---- cache-aware-zmq tuning (only used by that policy) ---- + // ---- legacy cache-aware-zmq tuning ---- /// Min `matched_blocks / total_blocks` for a cache match to win. #[arg(long)] pub cache_threshold: Option, @@ -73,19 +74,63 @@ pub struct Cli { /// Multiplicative load spread gating the absolute balance check. #[arg(long)] pub balance_rel_threshold: Option, - /// External KV indexer gRPC endpoint used as the cache signal. + /// External KV indexer gRPC endpoint used as the authoritative cache signal. + /// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`. #[arg(long)] pub kv_indexer_endpoint: Option, /// KV Indexer query timeout in milliseconds. Requires /// `--kv-indexer-endpoint`; defaults to 100. #[arg(long)] pub kv_indexer_query_timeout_ms: Option, - /// Maximum concurrent KV Indexer queries. Requires + /// Maximum concurrent KV Indexer queries issued by this Router. Requires /// `--kv-indexer-endpoint`; defaults to 32. #[arg(long)] pub kv_indexer_query_max_inflight: Option, - /// Weighted terms for `--policy fused_score`. + // ---- session-affinity tuning ---- + /// Header carrying the session ID for `--policy session_aware`. + #[arg(long)] + pub session_id_header: Option, + /// Idle timeout for a session assignment, in seconds. + #[arg(long)] + pub session_idle_secs: Option, + /// Session-assignment eviction cadence, in seconds. + #[arg(long)] + pub session_eviction_interval_secs: Option, + /// Use a deterministic backup for the affinity key and candidate range. + #[arg(long)] + pub stable_pair: bool, + /// Session-affinity admission mode. + #[arg(long, value_enum)] + pub affinity_mode: Option, + /// Session-affinity primary lookup and fallback behavior. + #[arg(long, value_enum)] + pub session_affinity_mode: Option, + /// Minimum cache-hit tokens for a cache-aware candidate. + #[arg(long)] + pub cache_affinity_min_matched_tokens: Option, + /// Minimum cache-hit ratio for a cache-aware candidate. + #[arg(long)] + pub cache_affinity_min_match_ratio: Option, + /// Minimum number of cache-aware candidates to try. + #[arg(long)] + pub cache_candidate_min_workers: Option, + /// Fraction of healthy prefill workers considered as cache-aware candidates. + #[arg(long)] + pub cache_candidate_ratio: Option, + /// Maximum number of cache-aware candidates to try. + #[arg(long)] + pub cache_candidate_max_workers: Option, + /// Maximum uncached-work difference that pressure may override. + #[arg(long)] + pub cache_switch_margin_tokens: Option, + + // ---- score composition ---- + /// Policies to sum, spelled exactly as `--policy` spells them and each + /// optionally weighted: `--fuse prefix_cache=2.0,load_based=0.3`. An + /// omitted weight keeps that policy's own default. Requires `--policy + /// score_policy` or `fused_score`; when either policy is set and this flag + /// is omitted, the terms default to `prefix_cache,load_based`. #[arg(long, value_delimiter = ',')] pub fuse: Vec, @@ -184,16 +229,12 @@ impl Cli { enabled by --cb-threshold)" )); } - let tuned_cache_aware = self.cache_threshold.is_some() + let tuned_legacy_cache_aware = self.cache_threshold.is_some() || self.balance_abs_threshold.is_some() - || self.balance_rel_threshold.is_some() - || self.kv_indexer_endpoint.is_some() - || self.kv_indexer_query_timeout_ms.is_some() - || self.kv_indexer_query_max_inflight.is_some(); - if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq { + || self.balance_rel_threshold.is_some(); + if tuned_legacy_cache_aware && self.policy != PolicyKind::CacheAwareZmq { return Err(anyhow!( - "--cache-threshold / --balance-abs-threshold / --balance-rel-threshold \ - require --policy cache_aware_zmq" + "cache-aware tuning flags require --policy cache_aware_zmq" )); } if self.kv_indexer_query_timeout_ms == Some(0) { @@ -216,11 +257,59 @@ impl Cli { "--kv-indexer-query-max-inflight requires --kv-indexer-endpoint" )); } - - if !self.fuse.is_empty() && self.policy != PolicyKind::FusedScore { - return Err(anyhow!("--fuse requires --policy fused_score")); + if self.kv_indexer_endpoint.is_some() + && !matches!( + self.policy, + PolicyKind::CacheAware | PolicyKind::CacheAwareZmq + ) + { + return Err(anyhow!( + "--kv-indexer-endpoint requires --policy cache_aware or cache_aware_zmq" + )); } - let fused = if self.policy == PolicyKind::FusedScore { + if self.policy == PolicyKind::CacheAware && self.kv_indexer_endpoint.is_none() { + return Err(anyhow!( + "--policy cache_aware requires --kv-indexer-endpoint" + )); + } + let tuned_cache_aware = tuned_legacy_cache_aware || self.kv_indexer_endpoint.is_some(); + let affinity_policy = matches!( + self.policy, + PolicyKind::SessionAware | PolicyKind::CacheAware + ); + let tuned_session_affinity = self.session_id_header.is_some() + || self.session_idle_secs.is_some() + || self.session_eviction_interval_secs.is_some() + || self.stable_pair + || self.affinity_mode.is_some() + || self.session_affinity_mode.is_some(); + if tuned_session_affinity && self.policy != PolicyKind::SessionAware { + return Err(anyhow!( + "--session-id-header, --session-*-secs, --stable-pair, --affinity-mode, and \ + --session-affinity-mode require --policy session_aware" + )); + } + let tuned_cache_candidates = self.cache_affinity_min_matched_tokens.is_some() + || self.cache_affinity_min_match_ratio.is_some() + || self.cache_candidate_min_workers.is_some() + || self.cache_candidate_ratio.is_some() + || self.cache_candidate_max_workers.is_some() + || self.cache_switch_margin_tokens.is_some(); + if tuned_cache_candidates && self.policy != PolicyKind::CacheAware { + return Err(anyhow!( + "cache candidate tuning flags require --policy cache_aware" + )); + } + let is_score_composition = matches!( + self.policy, + PolicyKind::FusedScore | PolicyKind::ScorePolicy + ); + if !self.fuse.is_empty() && !is_score_composition { + return Err(anyhow!( + "--fuse requires --policy score_policy or fused_score" + )); + } + let fused = if is_score_composition { let terms = if self.fuse.is_empty() { DEFAULT_FUSE .iter() @@ -250,6 +339,9 @@ impl Cli { "--max-in-flight and `--filter overloaded` require each other" )); } + if self.max_in_flight == Some(0) { + return Err(anyhow!("--max-in-flight must be greater than 0")); + } if self.prefix_cache_min_share.is_some() != has(FilterKind::PrefixCache) { return Err(anyhow!( "--prefix-cache-min-share and `--filter prefix_cache` require each other" @@ -261,6 +353,9 @@ impl Cli { { return Err(anyhow!("--prefix-cache-min-share must be in (0, 1]")); } + if self.policy == PolicyKind::Sticky && !self.filter.is_empty() { + return Err(anyhow!("--filter cannot be combined with --policy sticky")); + } let eligibility = (!self.filter.is_empty()).then(|| EligibilityConfig { filters: self.filter.clone(), max_in_flight: self.max_in_flight, @@ -318,6 +413,81 @@ impl Cli { None }; + let affinity = if affinity_policy { + let d = AffinityConfig::default(); + let session_id_header = self.session_id_header.unwrap_or(d.session_id_header); + axum::http::HeaderName::try_from(session_id_header.as_str()).map_err(|e| { + anyhow!("--session-id-header {session_id_header:?} is not a valid HTTP header name: {e}") + })?; + let cache_affinity_min_match_ratio = self + .cache_affinity_min_match_ratio + .or(d.cache_affinity_min_match_ratio); + if cache_affinity_min_match_ratio + .is_some_and(|ratio| !ratio.is_finite() || !(0.0..=1.0).contains(&ratio)) + { + return Err(anyhow!( + "--cache-affinity-min-match-ratio must be finite and in [0, 1]" + )); + } + let cache_candidate_ratio = self + .cache_candidate_ratio + .unwrap_or(d.cache_candidate_ratio); + if !cache_candidate_ratio.is_finite() || !(0.0..=1.0).contains(&cache_candidate_ratio) { + return Err(anyhow!( + "--cache-candidate-ratio must be finite and in [0, 1]" + )); + } + let cache_candidate_min_workers = self + .cache_candidate_min_workers + .unwrap_or(d.cache_candidate_min_workers); + let cache_candidate_max_workers = self + .cache_candidate_max_workers + .unwrap_or(d.cache_candidate_max_workers); + if cache_candidate_min_workers == 0 + || cache_candidate_max_workers == 0 + || cache_candidate_min_workers > cache_candidate_max_workers + { + return Err(anyhow!( + "--cache-candidate-min-workers and --cache-candidate-max-workers must be \ + positive and min must not exceed max" + )); + } + let session_idle_secs = self.session_idle_secs.unwrap_or(d.session_idle_secs); + let session_eviction_interval_secs = self + .session_eviction_interval_secs + .unwrap_or(d.session_eviction_interval_secs); + if session_idle_secs == 0 { + return Err(anyhow!("--session-idle-secs must be greater than 0")); + } + if session_eviction_interval_secs == 0 { + return Err(anyhow!( + "--session-eviction-interval-secs must be greater than 0" + )); + } + Some(AffinityConfig { + session_id_header, + session_idle_secs, + session_eviction_interval_secs, + stable_pair: self.stable_pair, + mode: self.affinity_mode.unwrap_or(d.mode), + session_affinity_mode: self + .session_affinity_mode + .unwrap_or(d.session_affinity_mode), + cache_affinity_min_matched_tokens: self + .cache_affinity_min_matched_tokens + .or(d.cache_affinity_min_matched_tokens), + cache_affinity_min_match_ratio, + cache_candidate_min_workers, + cache_candidate_ratio, + cache_candidate_max_workers, + cache_switch_margin_tokens: self + .cache_switch_margin_tokens + .unwrap_or(d.cache_switch_margin_tokens), + }) + } 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), @@ -326,14 +496,17 @@ impl Cli { // Only build a CacheAwareConfig when the operator tuned at least // one knob; otherwise leave it None so the policy uses its own // defaults. Unset knobs fall back to the per-field defaults. - let kv_indexer_query_timeout_ms = self - .kv_indexer_query_timeout_ms - .unwrap_or(DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS); - let kv_indexer_query_max_inflight = self - .kv_indexer_query_max_inflight - .unwrap_or(DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT); let cache_aware = if tuned_cache_aware { let d = CacheAwareConfig::default(); + let kv_indexer_endpoint = self.kv_indexer_endpoint.map(|url| KvIndexerEndpointConfig { + url, + query_timeout_ms: self + .kv_indexer_query_timeout_ms + .unwrap_or(DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS), + query_max_inflight: self + .kv_indexer_query_max_inflight + .unwrap_or(DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT), + }); Some(CacheAwareConfig { cache_threshold: self.cache_threshold.unwrap_or(d.cache_threshold), balance_abs_threshold: self @@ -342,11 +515,7 @@ impl Cli { balance_rel_threshold: self .balance_rel_threshold .unwrap_or(d.balance_rel_threshold), - kv_indexer_endpoint: self.kv_indexer_endpoint.map(|url| KvIndexerEndpointConfig { - url, - query_timeout_ms: kv_indexer_query_timeout_ms, - query_max_inflight: kv_indexer_query_max_inflight, - }), + kv_indexer_endpoint, }) } else { None @@ -370,6 +539,7 @@ impl Cli { circuit_breaker, cache_aware, sticky, + affinity, fused, eligibility, }, @@ -889,7 +1059,7 @@ mod tests { "--worker-urls", "http://x:30000", "--policy", - "cache_aware_zmq", + "cache_aware", "--kv-indexer-endpoint", "http://indexer:50051", "--kv-indexer-query-timeout-ms", @@ -907,6 +1077,30 @@ mod tests { #[test] fn kv_indexer_uses_safe_query_defaults() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "cache_aware", + "--kv-indexer-endpoint", + "http://indexer:50051", + ])) + .unwrap(); + let indexer = c + .model + .cache_aware + .expect("cache-aware config") + .kv_indexer_endpoint + .expect("Indexer config"); + assert_eq!( + indexer.query_timeout_ms, + DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS + ); + assert_eq!(indexer.query_max_inflight, 32); + } + + #[test] + fn kv_indexer_is_accepted_by_cache_aware_zmq() { let c = into_config_owned(with_model(&[ "--worker-urls", "http://x:30000", @@ -922,8 +1116,7 @@ mod tests { .expect("cache-aware config") .kv_indexer_endpoint .expect("Indexer config"); - assert_eq!(indexer.query_timeout_ms, 100); - assert_eq!(indexer.query_max_inflight, 32); + assert_eq!(indexer.url, "http://indexer:50051"); } #[test] @@ -936,7 +1129,7 @@ mod tests { ])) .unwrap_err() .to_string(); - assert!(err.contains("require --policy cache_aware_zmq")); + assert!(err.contains("requires --policy cache_aware"), "got: {err}"); } #[test] @@ -945,7 +1138,7 @@ mod tests { "--worker-urls", "http://x:30000", "--policy", - "cache_aware_zmq", + "cache_aware", "--kv-indexer-query-timeout-ms", "75", ])) @@ -960,7 +1153,7 @@ mod tests { "--worker-urls", "http://x:30000", "--policy", - "cache_aware_zmq", + "cache_aware", "--kv-indexer-query-max-inflight", "17", ])) @@ -975,7 +1168,7 @@ mod tests { "--worker-urls", "http://x:30000", "--policy", - "cache_aware_zmq", + "cache_aware", "--kv-indexer-endpoint", "http://indexer:50051", "--kv-indexer-query-max-inflight", @@ -1145,7 +1338,7 @@ mod tests { #[test] fn filter_misconfigurations_fail_at_startup() { - let cases: [(&[&str], &str); 6] = [ + let cases: [(&[&str], &str); 8] = [ (&["--filter", "overloaded"], "require each other"), (&["--max-in-flight", "64"], "require each other"), (&["--filter", "prefix_cache"], "require each other"), @@ -1163,6 +1356,21 @@ mod tests { ], "must be in (0, 1]", ), + ( + &["--filter", "overloaded", "--max-in-flight", "0"], + "must be greater than 0", + ), + ( + &[ + "--policy", + "sticky", + "--filter", + "overloaded", + "--max-in-flight", + "1", + ], + "cannot be combined with --policy sticky", + ), ]; for (extra, want) in cases { let mut args = vec!["--worker-urls", "http://x:30000"]; @@ -1307,8 +1515,31 @@ mod tests { fused_of(argv).expect("fused_score builds a term list") } - /// `--policy fused_score` alone composes the useful pair, and `--fuse` - /// overrides it with names + optional per-term weights. + /// `score_policy` is an independent top-level policy. + #[test] + fn score_policy_is_a_top_level_policy_with_its_own_cli_spelling() { + use PolicyKind::ScorePolicy; + use ScoreTermKind::{LoadBased, PrefixCache}; + let pair = [(PrefixCache, None), (LoadBased, None)]; + let config = cfg_of("--policy score_policy").unwrap(); + assert_eq!(config.model.policy, ScorePolicy); + assert_eq!( + config + .model + .fused + .expect("score_policy must resolve its score terms") + .iter() + .map(|t| (t.kind, t.weight)) + .collect::>(), + pair, + ); + assert_eq!( + fuse_ok("--policy score_policy --fuse prefix_cache=2.0,load_based=0.3"), + [(PrefixCache, Some(2.0)), (LoadBased, Some(0.3))], + ); + } + + /// `fused_score` keeps the compatibility entry point. #[test] fn fuse_defaults_to_the_useful_pair_and_parses_weights() { use ScoreTermKind::{LoadBased, PrefixCache, Random}; @@ -1345,7 +1576,7 @@ mod tests { #[test] fn fuse_rejects_malformed_compositions() { - let cases: [(&str, &[&str]); 5] = [ + let cases: [(&str, &[&str]); 6] = [ ("--fuse load_based", &["--fuse requires", "fused_score"]), ( "--policy fused_score --fuse fused_score,load_based", @@ -1355,6 +1586,10 @@ mod tests { "--policy fused_score --fuse load_based,load_based", &["load_based", "listed more than once"], ), + ( + "--policy score_policy --fuse score_policy,load_based", + &["score_policy", "not a score term"], + ), ( "--policy fused_score --fuse not_a_policy", &["not_a_policy", "is not a score term"], @@ -1371,4 +1606,199 @@ mod tests { } } } + + #[test] + fn session_aware_builds_affinity_config_from_its_cli_knobs() { + let config = cfg_of( + "--policy session_aware --session-id-header x-agent-session --stable-pair \ + --affinity-mode strict --session-affinity-mode global-rebind", + ) + .unwrap(); + let affinity = config + .model + .affinity + .expect("session policy needs affinity config"); + + assert_eq!(config.model.policy, PolicyKind::SessionAware); + assert_eq!(affinity.session_id_header, "x-agent-session"); + assert!(affinity.stable_pair); + assert_eq!(affinity.mode, AffinityMode::Strict); + assert_eq!( + affinity.session_affinity_mode, + SessionAffinityMode::GlobalRebind + ); + } + + #[test] + fn rejects_removed_token_pressure_flags() { + for flag in [ + "--disable-pressure-guard", + "--pressure-abs-threshold-tokens 2048", + "--pressure-rel-threshold 2.0", + ] { + let error = cfg_of(&format!("--policy session_aware {flag}")) + .unwrap_err() + .to_string(); + assert!(error.contains("unexpected argument"), "{flag}: {error}"); + } + } + + #[test] + fn rejects_removed_affinity_aware_range_flag() { + let error = cfg_of("--policy session_aware --affinity-aware-range global-first") + .unwrap_err() + .to_string(); + assert!(error.contains("unexpected argument '--affinity-aware-range'")); + } + + #[test] + fn session_affinity_mode_accepts_all_new_values() { + for (value, expected) in [ + ("bucket", SessionAffinityMode::Bucket), + ("global-rebind", SessionAffinityMode::GlobalRebind), + ("global-preserve", SessionAffinityMode::GlobalPreserve), + ] { + let config = cfg_of(&format!( + "--policy session_aware --session-affinity-mode {value}" + )) + .unwrap(); + assert_eq!( + config.model.affinity.unwrap().session_affinity_mode, + expected + ); + } + } + + #[test] + fn session_aware_configures_bounded_assignment_lifetime() { + let config = cfg_of( + "--policy session_aware --session-idle-secs 120 \ + --session-eviction-interval-secs 15", + ) + .unwrap(); + let affinity = config.model.affinity.expect("session affinity config"); + assert_eq!(affinity.session_idle_secs, 120); + assert_eq!(affinity.session_eviction_interval_secs, 15); + } + + #[test] + fn cache_aware_accepts_indexer_endpoint_and_rejects_affinity_knobs_elsewhere() { + let config = cfg_of( + "--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \ + --kv-indexer-query-timeout-ms 40 \ + --cache-affinity-min-matched-tokens 512 --cache-affinity-min-match-ratio 0.25 \ + --cache-candidate-min-workers 4 --cache-candidate-ratio 0.1 \ + --cache-candidate-max-workers 16 --cache-switch-margin-tokens 128", + ) + .unwrap(); + assert_eq!(config.model.policy, PolicyKind::CacheAware); + assert_eq!( + config + .model + .cache_aware + .as_ref() + .expect("cache-aware needs indexer config") + .kv_indexer_endpoint + .as_ref() + .map(|indexer| indexer.url.as_str()), + Some("http://indexer:50051"), + ); + let indexer_timeout_ms = config + .model + .cache_aware + .as_ref() + .and_then(|cache| cache.kv_indexer_endpoint.as_ref()) + .expect("cache-aware needs indexer config") + .query_timeout_ms; + let affinity = config + .model + .affinity + .expect("cache-aware needs candidate config"); + assert_eq!(affinity.cache_affinity_min_matched_tokens, Some(512)); + assert_eq!(affinity.cache_affinity_min_match_ratio, Some(0.25)); + assert_eq!(affinity.cache_candidate_min_workers, 4); + assert_eq!(affinity.cache_candidate_ratio, 0.1); + assert_eq!(affinity.cache_candidate_max_workers, 16); + assert_eq!(affinity.cache_switch_margin_tokens, 128); + assert_eq!(indexer_timeout_ms, 40); + + let defaults = + cfg_of("--policy cache_aware --kv-indexer-endpoint http://indexer:50051").unwrap(); + let defaults_indexer_timeout_ms = defaults + .model + .cache_aware + .as_ref() + .and_then(|cache| cache.kv_indexer_endpoint.as_ref()) + .expect("default indexer config") + .query_timeout_ms; + let defaults_affinity = defaults + .model + .affinity + .expect("default cache candidate config"); + assert_eq!( + defaults_affinity.cache_affinity_min_matched_tokens, + Some(1_024) + ); + assert_eq!(defaults_affinity.cache_affinity_min_match_ratio, None); + assert_eq!( + defaults_indexer_timeout_ms, + DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS + ); + + let err = cfg_of("--policy power_of_two --stable-pair") + .unwrap_err() + .to_string(); + assert!( + err.contains("--stable-pair") && err.contains("session_aware"), + "got: {err}" + ); + + let err = + cfg_of("--policy cache_aware --kv-indexer-endpoint http://indexer:50051 --stable-pair") + .expect_err("Cache-Aware has no stable backup") + .to_string(); + assert!(err.contains("--stable-pair"), "got: {err}"); + } + + #[test] + fn cache_candidate_cli_rejects_invalid_bounds() { + for (args, expected) in [ + ( + "--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \ + --cache-affinity-min-match-ratio 1.1", + "--cache-affinity-min-match-ratio", + ), + ( + "--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \ + --cache-candidate-min-workers 9 --cache-candidate-max-workers 8", + "--cache-candidate-min-workers", + ), + ( + "--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \ + --cache-candidate-ratio=-0.1", + "--cache-candidate-ratio", + ), + ( + "--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \ + --kv-indexer-query-timeout-ms 0", + "--kv-indexer-query-timeout-ms", + ), + ] { + let err = cfg_of(args) + .expect_err("invalid candidate bound") + .to_string(); + assert!(err.contains(expected), "got: {err}"); + } + } + + #[test] + fn rejects_affinity_options_that_cannot_affect_the_selected_policy() { + let missing_indexer = cfg_of("--policy cache_aware") + .expect_err("cache_aware without an indexer can only behave like P2") + .to_string(); + assert!( + missing_indexer.contains("--kv-indexer-endpoint"), + "got: {missing_indexer}" + ); + } } diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index 3dad456d6..b8fc86b01 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -88,6 +88,7 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 9e56a75ca..0577691b2 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -70,7 +70,8 @@ impl Default for ActiveLoadConfig { /// policy factory. /// /// Accepted on the CLI (`--policy`) as `round_robin` / `random` / -/// `power_of_two` / `load_based` / `fused_score` / `cache_aware_zmq` / +/// `power_of_two` / `load_based` / `fused_score` / `score_policy` / +/// `session_aware` / `cache_aware` / `cache_aware_zmq` / /// `sticky`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)] pub enum PolicyKind { @@ -87,6 +88,15 @@ pub enum PolicyKind { /// Weighted sum of `--fuse` terms. #[value(name = "fused_score")] FusedScore, + /// Composes compatible scoring terms into a single routing policy. + #[value(name = "score_policy")] + ScorePolicy, + /// Selects a worker from session affinity. + #[value(name = "session_aware")] + SessionAware, + /// Selects cache-affine prefill candidates from external indexer data. + #[value(name = "cache_aware")] + CacheAware, /// Cache-aware routing fed by SGLang's ZMQ KV-cache event publisher. /// Requires the model to have a tokenizer loaded; cache_aware tuning /// lives on `ModelConfig::cache_aware`. @@ -222,16 +232,19 @@ pub struct ModelConfig { pub tokenizer_path: String, pub policy: PolicyKind, pub circuit_breaker: Option, - /// Tuning for the cache-aware ZMQ policy. Ignored unless - /// `policy = "cache_aware_zmq"`. `None` falls back to defaults at - /// policy construction time. + /// Cache-Aware ZMQ tuning and optional external Indexer endpoint. 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, - /// Terms for `policy = "fused_score"`. + /// Session and cache-affinity tuning. + pub affinity: Option, + /// Terms the score-composition policy sums. `Some` exactly when + /// `policy = "fused_score"` or `policy = "score_policy"` (built by + /// [`crate::config::cli::Cli::into_config`]), defaulting to + /// [`DEFAULT_FUSE`] when `--fuse` is omitted. pub fused: Option>, /// Hard constraints applied before policy selection. pub eligibility: Option, @@ -340,6 +353,76 @@ fn default_balance_rel() -> f32 { /// (`x-sgl-decode-url`, `x-sgl-router-error-code`). pub const DEFAULT_STICKY_HEADER: &str = "x-sgl-routing-key"; +/// Default request header for session-aware routing. +pub const DEFAULT_SESSION_ID_HEADER: &str = "x-session-id"; + +/// Default external-indexer request limits. +pub const DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT: usize = 32; + +/// Controls whether admission may select a session-affinity backup. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)] +pub enum AffinityMode { + /// Keep the primary after it passes admission. + #[value(name = "strict")] + Strict, + /// Allow the admitted backup to relieve pressure. + #[default] + #[value(name = "soft")] + Soft, +} + +/// Controls the session-affinity lookup and fallback behavior. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)] +pub enum SessionAffinityMode { + /// Search only within the target bucket. + #[default] + #[value(name = "bucket")] + Bucket, + /// Rebind to a target-bucket fallback when the global primary is unavailable. + #[value(name = "global-rebind")] + GlobalRebind, + /// Keep a valid global assignment when a bucket fallback is used. + #[value(name = "global-preserve")] + GlobalPreserve, +} + +/// Shared session-aware and cache-aware settings. +#[derive(Debug, Clone)] +pub struct AffinityConfig { + pub session_id_header: String, + pub session_idle_secs: u64, + pub session_eviction_interval_secs: u64, + pub stable_pair: bool, + pub mode: AffinityMode, + pub session_affinity_mode: SessionAffinityMode, + pub cache_affinity_min_matched_tokens: Option, + pub cache_affinity_min_match_ratio: Option, + pub cache_candidate_min_workers: usize, + pub cache_candidate_ratio: f64, + pub cache_candidate_max_workers: usize, + pub cache_switch_margin_tokens: u64, +} + +impl Default for AffinityConfig { + fn default() -> Self { + Self { + session_id_header: DEFAULT_SESSION_ID_HEADER.to_string(), + session_idle_secs: default_sticky_idle_secs(), + session_eviction_interval_secs: default_sticky_eviction_interval_secs(), + stable_pair: false, + mode: AffinityMode::Soft, + session_affinity_mode: SessionAffinityMode::Bucket, + // Indexer prefix scans are truncated, so use an absolute token floor. + cache_affinity_min_matched_tokens: Some(1_024), + cache_affinity_min_match_ratio: None, + cache_candidate_min_workers: 8, + cache_candidate_ratio: 0.05, + cache_candidate_max_workers: 32, + cache_switch_margin_tokens: 1_024, + } + } +} + /// 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. diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index e1fb8b08f..d4d172178 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -141,6 +141,7 @@ async fn main() -> Result<()> { kv_index.tree(), Arc::clone(&tokenizers), Arc::clone(&block_size_oracle), + kv_index.engine_load(), ) .context("build policy registry")?, ); @@ -197,6 +198,7 @@ async fn main() -> Result<()> { ); app_ctx.prefix_index = prefix_index; app_ctx.block_size_oracle = block_size_oracle; + app_ctx.engine_load = kv_index.engine_load(); let ctx = Arc::new(app_ctx); ctx.mark_ready(); diff --git a/experimental/sgl-router/src/policies/admission.rs b/experimental/sgl-router/src/policies/admission.rs new file mode 100644 index 000000000..5673c8641 --- /dev/null +++ b/experimental/sgl-router/src/policies/admission.rs @@ -0,0 +1,672 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Shared admission and candidate comparison for Prefill and Decode. +//! +//! Decisions use only fields published in `LoadStat`. + +use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad}; +use crate::policies::power_of_two::select_with_snapshot; +use crate::policies::{CacheCandidate, CacheCandidateProposal, SelectionProposal}; +use crate::workers::Worker; +use std::cmp::Ordering; +use std::collections::HashMap; +use std::sync::Arc; + +/// A Prefill candidate range and its optional pending-token budget. +pub struct CandidateRange<'a> { + pub id: &'a str, + pub workers: &'a [Arc], + pub max_pending_prefill_tokens: Option, +} + +impl<'a> CandidateRange<'a> { + pub fn global(workers: &'a [Arc]) -> Self { + Self { + id: "global", + workers, + max_pending_prefill_tokens: None, + } + } +} + +/// A role-specific candidate domain resolved before policy evaluation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RoutingStage { + Prefill, + Decode, +} + +#[derive(Clone)] +pub struct CandidateDomain { + pub id: String, + pub stage: RoutingStage, + pub workers: Vec>, + pub max_pending_prefill_tokens: Option, +} + +impl CandidateDomain { + pub fn global_prefill(workers: &[Arc]) -> Self { + Self { + id: "global".to_string(), + stage: RoutingStage::Prefill, + workers: workers.to_vec(), + max_pending_prefill_tokens: None, + } + } + + pub fn global_decode(workers: &[Arc]) -> Self { + Self { + id: "global".to_string(), + stage: RoutingStage::Decode, + workers: workers.to_vec(), + max_pending_prefill_tokens: None, + } + } + + pub fn bucket_prefill( + id: impl Into, + workers: Vec>, + max_pending_prefill_tokens: Option, + ) -> Self { + Self { + id: id.into(), + stage: RoutingStage::Prefill, + workers, + max_pending_prefill_tokens, + } + } + + pub fn bucket_decode(id: impl Into, workers: Vec>) -> Self { + Self { + id: id.into(), + stage: RoutingStage::Decode, + workers, + max_pending_prefill_tokens: None, + } + } + + pub fn prefill_range(&self) -> Option> { + (self.stage == RoutingStage::Prefill).then(|| CandidateRange { + id: self.id.as_str(), + workers: &self.workers, + max_pending_prefill_tokens: self.max_pending_prefill_tokens, + }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DecisionReason { + Primary, + CacheCandidate, + BackupPrimaryAdmission, + /// Both Decode candidates were admitted; the lower-pressure backup won. + BackupLoadComparison, + RangeFallback, + CapacityFallbackPowerOfTwo, +} + +#[derive(Clone)] +pub struct FinalDecision { + pub selected: Arc, + pub primary: Arc, + pub backup: Option>, + pub reason: DecisionReason, + pub candidate_range_id: String, + pub load_snapshot_version: u64, +} + +/// Resolves a bounded cache-candidate set, using pressure to break near ties. +pub fn resolve_cache_candidates( + proposal: &CacheCandidateProposal, + request_input_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> Option { + let loads = FreshLoadLookup::new( + Some(snapshot), + proposal + .candidates + .iter() + .map(|candidate| &candidate.worker), + ); + let admitted: Vec<&CacheCandidate> = proposal + .candidates + .iter() + .filter(|candidate| is_cache_candidate_admitted(candidate, request_input_tokens, &loads)) + .collect(); + let work_floor = admitted + .iter() + .copied() + .min_by_key(|candidate| candidate.uncached_tokens)?; + let near_tie_ceiling = work_floor + .uncached_tokens + .saturating_add(proposal.cache_switch_margin_tokens); + let mut winner = work_floor; + for candidate in admitted { + if candidate.uncached_tokens <= near_tie_ceiling + && compare_cache_candidates(winner, candidate, &loads).is_gt() + { + winner = candidate; + } + } + Some(FinalDecision { + selected: Arc::clone(&winner.worker), + primary: Arc::clone(&winner.worker), + backup: None, + reason: DecisionReason::CacheCandidate, + candidate_range_id: winner.candidate_range_id.clone(), + load_snapshot_version: snapshot.version, + }) +} + +pub fn resolve_prefill( + range: &CandidateRange<'_>, + proposal: &SelectionProposal, + request_input_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> Option { + if !contains_worker(range, &proposal.primary) { + return None; + } + let backup = proposal + .backup + .as_ref() + .filter(|worker| contains_worker(range, worker)) + .cloned(); + let primary_admitted = is_proposal_worker_eligible(proposal, &proposal.primary) + && is_prefill_admitted(&proposal.primary, request_input_tokens, snapshot); + let backup_admitted = backup.as_ref().is_some_and(|worker| { + is_proposal_worker_eligible(proposal, worker) + && is_prefill_admitted(worker, request_input_tokens, snapshot) + }); + + let (selected, reason) = match (primary_admitted, backup.as_ref(), backup_admitted) { + (true, _, _) => (Arc::clone(&proposal.primary), DecisionReason::Primary), + (false, Some(backup), true) => (Arc::clone(backup), DecisionReason::BackupPrimaryAdmission), + _ => { + let legal = legal_prefill_candidates(range, proposal); + range_fallback(&legal, request_input_tokens, snapshot).or_else(|| { + select_with_snapshot(&legal, Some(snapshot)) + .map(|worker| (worker, DecisionReason::CapacityFallbackPowerOfTwo)) + })? + } + }; + Some(FinalDecision { + selected, + primary: Arc::clone(&proposal.primary), + backup, + reason, + candidate_range_id: range.id.to_string(), + load_snapshot_version: snapshot.version, + }) +} + +pub fn resolve_decode( + domain: &CandidateDomain, + proposal: &SelectionProposal, + request_kv_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> Option { + if domain.stage != RoutingStage::Decode || !contains_domain_worker(domain, &proposal.primary) { + return None; + } + let backup = proposal + .backup + .as_ref() + .filter(|worker| contains_domain_worker(domain, worker)) + .cloned(); + let primary_admitted = is_decode_admitted(&proposal.primary, request_kv_tokens, snapshot); + let backup_admitted = backup + .as_ref() + .is_some_and(|worker| is_decode_admitted(worker, request_kv_tokens, snapshot)); + let (selected, reason) = match (primary_admitted, backup.as_ref(), backup_admitted) { + (true, Some(backup), true) => { + if compare_decode_pressure(&proposal.primary, backup, Some(snapshot)).is_gt() { + (Arc::clone(backup), DecisionReason::BackupLoadComparison) + } else { + (Arc::clone(&proposal.primary), DecisionReason::Primary) + } + } + (true, _, _) => (Arc::clone(&proposal.primary), DecisionReason::Primary), + (false, Some(backup), true) => (Arc::clone(backup), DecisionReason::BackupPrimaryAdmission), + _ => decode_domain_fallback(domain, request_kv_tokens, snapshot)?, + }; + Some(FinalDecision { + selected, + primary: Arc::clone(&proposal.primary), + backup, + reason, + candidate_range_id: domain.id.clone(), + load_snapshot_version: snapshot.version, + }) +} + +fn contains_worker(range: &CandidateRange<'_>, candidate: &Arc) -> bool { + range.workers.iter().any(|worker| worker.id == candidate.id) +} + +fn contains_domain_worker(domain: &CandidateDomain, candidate: &Arc) -> bool { + domain + .workers + .iter() + .any(|worker| worker.id == candidate.id) +} + +fn is_proposal_worker_eligible(proposal: &SelectionProposal, candidate: &Arc) -> bool { + proposal + .eligible_workers + .as_ref() + .is_none_or(|workers| workers.iter().any(|worker| worker.id == candidate.id)) +} + +/// A zero LoadStat capacity is unknown and does not reject a candidate. +fn has_kv_capacity(load: Option<&EngineWorkerLoad>, requested_tokens: u64) -> bool { + let Some(load) = load else { + return true; + }; + load.max_total_num_tokens == 0 + || load.num_tokens.saturating_add(requested_tokens) <= load.max_total_num_tokens +} + +fn is_prefill_admitted( + worker: &Arc, + request_input_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> bool { + has_kv_capacity( + snapshot.fresh_load_for_url(&worker.url), + request_input_tokens, + ) +} + +fn is_decode_admitted( + worker: &Arc, + request_kv_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> bool { + has_kv_capacity(snapshot.fresh_load_for_url(&worker.url), request_kv_tokens) +} + +fn is_cache_candidate_admitted( + candidate: &CacheCandidate, + request_input_tokens: u64, + loads: &FreshLoadLookup<'_>, +) -> bool { + has_kv_capacity(loads.get(&candidate.worker.id), request_input_tokens) +} + +fn compare_cache_candidates( + left: &CacheCandidate, + right: &CacheCandidate, + loads: &FreshLoadLookup<'_>, +) -> Ordering { + left.uncached_tokens + .cmp(&right.uncached_tokens) + .then_with(|| loads.compare_prefill_pressure(&left.worker, &right.worker)) + .then_with(|| left.worker.id.0.cmp(&right.worker.id.0)) +} + +/// Per-request lookup that uses engine pressure only for a complete fresh set. +pub(crate) struct FreshLoadLookup<'a> { + by_worker_id: HashMap, + local_active_by_worker_id: HashMap, + compare_engine: bool, +} + +impl<'a> FreshLoadLookup<'a> { + pub(crate) fn new<'w>( + snapshot: Option<&'a EngineLoadSnapshot>, + workers: impl IntoIterator>, + ) -> Self { + let workers: Vec<&Arc> = workers.into_iter().collect(); + let local_active_by_worker_id: HashMap = workers + .iter() + .map(|worker| (worker.id.0.clone(), worker.active_load())) + .collect(); + let by_worker_id = snapshot + .into_iter() + .flat_map(|snapshot| { + workers.iter().filter_map(move |worker| { + snapshot + .fresh_load_for_url(&worker.url) + .map(|load| (worker.id.0.clone(), load)) + }) + }) + .collect::>(); + let compare_engine = !local_active_by_worker_id.is_empty() + && by_worker_id.len() == local_active_by_worker_id.len(); + Self { + by_worker_id, + local_active_by_worker_id, + compare_engine, + } + } + + pub(crate) fn get( + &self, + worker_id: &crate::discovery::WorkerId, + ) -> Option<&'a EngineWorkerLoad> { + self.by_worker_id.get(worker_id.0.as_str()).copied() + } + + fn comparable_get( + &self, + worker_id: &crate::discovery::WorkerId, + ) -> Option<&'a EngineWorkerLoad> { + self.compare_engine.then(|| self.get(worker_id)).flatten() + } + + fn pressure_key(&self, worker: &Arc) -> PressureKey<'a> { + PressureKey { + load: self.comparable_get(&worker.id), + local_active: self + .local_active_by_worker_id + .get(worker.id.0.as_str()) + .copied() + .unwrap_or(usize::MAX), + } + } + + fn compare_prefill_keys(&self, left: &PressureKey<'a>, right: &PressureKey<'a>) -> Ordering { + match (left.load, right.load) { + (Some(left_load), Some(right_load)) => prefill_pressure_key(left_load) + .cmp(&prefill_pressure_key(right_load)) + .then_with(|| left.local_active.cmp(&right.local_active)), + _ => left.local_active.cmp(&right.local_active), + } + } + + fn compare_decode_keys(&self, left: &PressureKey<'a>, right: &PressureKey<'a>) -> Ordering { + match (left.load, right.load) { + (Some(left_load), Some(right_load)) => compare_decode_load(left_load, right_load) + .then_with(|| left.local_active.cmp(&right.local_active)), + _ => left.local_active.cmp(&right.local_active), + } + } + + pub(crate) fn compare_prefill_pressure( + &self, + left: &Arc, + right: &Arc, + ) -> Ordering { + self.compare_prefill_keys(&self.pressure_key(left), &self.pressure_key(right)) + } + + /// Returns corrected engine queue depth for a complete fresh set, otherwise + /// local load. + pub(crate) fn score_load(&self, worker: &Arc) -> usize { + self.comparable_get(&worker.id) + .map(|load| { + let recent_dispatches = worker + .slots_acquired_since(load.captured_at) + .try_into() + .unwrap_or(u64::MAX); + load.num_waiting_reqs + .saturating_add(load.num_running_reqs) + .saturating_add(recent_dispatches) + .try_into() + .unwrap_or(usize::MAX) + }) + .unwrap_or_else(|| { + self.local_active_by_worker_id + .get(worker.id.0.as_str()) + .copied() + .unwrap_or(usize::MAX) + }) + } + + fn min_by_pressure_key( + &self, + candidates: Vec>, + compare: impl Fn(&Self, &PressureKey<'a>, &PressureKey<'a>) -> Ordering, + ) -> Option> { + let mut candidates = candidates.into_iter(); + let mut best = candidates.next()?; + let mut best_key = self.pressure_key(&best); + for candidate in candidates { + let key = self.pressure_key(&candidate); + if compare(self, &key, &best_key).is_lt() { + best = candidate; + best_key = key; + } + } + Some(best) + } +} + +struct PressureKey<'a> { + load: Option<&'a EngineWorkerLoad>, + local_active: usize, +} + +fn range_fallback( + legal: &[Arc], + request_input_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> Option<(Arc, DecisionReason)> { + let admitted = legal + .iter() + .filter(|worker| is_prefill_admitted(worker, request_input_tokens, snapshot)) + .cloned() + .collect::>(); + let loads = FreshLoadLookup::new(Some(snapshot), admitted.iter()); + loads + .min_by_pressure_key(admitted, FreshLoadLookup::compare_prefill_keys) + .map(|worker| (worker, DecisionReason::RangeFallback)) +} + +fn legal_prefill_candidates( + range: &CandidateRange<'_>, + proposal: &SelectionProposal, +) -> Vec> { + proposal + .eligible_workers + .as_deref() + .unwrap_or(range.workers) + .iter() + .filter(|worker| contains_worker(range, worker)) + .cloned() + .collect() +} + +fn decode_domain_fallback( + domain: &CandidateDomain, + request_kv_tokens: u64, + snapshot: &EngineLoadSnapshot, +) -> Option<(Arc, DecisionReason)> { + let admitted = domain + .workers + .iter() + .filter(|worker| is_decode_admitted(worker, request_kv_tokens, snapshot)) + .cloned() + .collect::>(); + let loads = FreshLoadLookup::new(Some(snapshot), admitted.iter()); + loads + .min_by_pressure_key(admitted, FreshLoadLookup::compare_decode_keys) + .map(|worker| (worker, DecisionReason::RangeFallback)) +} + +/// Compares Prefill pressure by waiting requests, running requests, and KV use. +pub(crate) fn compare_prefill_pressure( + left: &Arc, + right: &Arc, + snapshot: Option<&EngineLoadSnapshot>, +) -> Ordering { + match snapshot.and_then(|snapshot| { + Some(( + snapshot.fresh_load_for_url(&left.url)?, + snapshot.fresh_load_for_url(&right.url)?, + )) + }) { + Some((left_load, right_load)) => prefill_pressure_key(left_load) + .cmp(&prefill_pressure_key(right_load)) + .then_with(|| left.active_load().cmp(&right.active_load())), + None => left.active_load().cmp(&right.active_load()), + } +} + +fn prefill_pressure_key(load: &EngineWorkerLoad) -> (u64, u64, u64, u64) { + ( + load.num_waiting_reqs, + load.num_running_reqs, + load.num_tokens, + load.max_total_num_tokens, + ) +} + +/// Compares Decode pressure using only LoadStat values. +pub(crate) fn compare_decode_pressure( + left: &Arc, + right: &Arc, + snapshot: Option<&EngineLoadSnapshot>, +) -> Ordering { + match snapshot.and_then(|snapshot| { + Some(( + snapshot.fresh_load_for_url(&left.url)?, + snapshot.fresh_load_for_url(&right.url)?, + )) + }) { + Some((left_load, right_load)) => compare_decode_load(left_load, right_load) + .then_with(|| left.active_load().cmp(&right.active_load())), + None => left.active_load().cmp(&right.active_load()), + } +} + +fn compare_decode_load(left: &EngineWorkerLoad, right: &EngineWorkerLoad) -> Ordering { + let kv_usage = match (left.max_total_num_tokens, right.max_total_num_tokens) { + (left_cap, right_cap) if left_cap > 0 && right_cap > 0 => u128::from(left.num_tokens) + .saturating_mul(u128::from(right_cap)) + .cmp(&u128::from(right.num_tokens).saturating_mul(u128::from(left_cap))), + _ => Ordering::Equal, + }; + left.num_waiting_reqs + .cmp(&right.num_waiting_reqs) + .then_with(|| left.num_running_reqs.cmp(&right.num_running_reqs)) + .then(kv_usage) + .then_with(|| left.num_tokens.cmp(&right.num_tokens)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use std::time::Instant; + + 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("model".into())], + bootstrap_port: None, + })) + } + + fn snapshot(entries: &[(&Arc, u64, u64, u64, u64)]) -> EngineLoadSnapshot { + EngineLoadSnapshot::from_workers( + 7, + entries + .iter() + .map(|(worker, running, waiting, used, capacity)| { + ( + worker.url.clone(), + EngineWorkerLoad { + num_running_reqs: *running, + num_waiting_reqs: *waiting, + num_tokens: *used, + max_total_num_tokens: *capacity, + captured_at: Instant::now(), + }, + ) + }) + .collect(), + ) + } + + #[test] + fn capacity_rejects_only_when_the_published_capacity_is_exceeded() { + let full = worker("full"); + let unknown = worker("unknown"); + let workers = vec![Arc::clone(&full), Arc::clone(&unknown)]; + let range = CandidateRange::global(&workers); + let loads = snapshot(&[(&full, 0, 0, 90, 100), (&unknown, 0, 0, 90, 0)]); + + assert!(resolve_prefill( + &range, + &SelectionProposal::primary(Arc::clone(&full)), + 20, + &loads + ) + .is_some()); + assert_eq!( + resolve_prefill(&range, &SelectionProposal::primary(full), 20, &loads) + .expect("fallback selects unknown-capacity worker") + .selected + .id, + unknown.id + ); + } + + #[test] + fn all_capacity_rejected_falls_back_to_power_of_two_within_eligible_domain() { + let primary = worker("primary"); + let backup = worker("backup"); + let filtered = worker("filtered"); + let workers = vec![ + Arc::clone(&primary), + Arc::clone(&backup), + Arc::clone(&filtered), + ]; + let proposal = SelectionProposal::with_backup(Arc::clone(&primary), Arc::clone(&backup)) + .with_eligible_workers(vec![Arc::clone(&primary), Arc::clone(&backup)]); + let loads = snapshot(&[ + (&primary, 0, 0, 100, 100), + (&backup, 0, 0, 100, 100), + (&filtered, 0, 0, 0, 100), + ]); + + let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &loads) + .expect("capacity exhaustion must degrade within the legal domain"); + + assert!(matches!( + decision.selected.id.0.as_str(), + "primary" | "backup" + )); + assert_eq!(decision.reason, DecisionReason::CapacityFallbackPowerOfTwo); + } + + #[test] + fn capacity_fallback_uses_the_explicit_snapshot_for_power_of_two() { + let primary = worker("primary"); + let backup = worker("backup"); + let workers = vec![Arc::clone(&primary), Arc::clone(&backup)]; + let proposal = SelectionProposal::with_backup(Arc::clone(&primary), Arc::clone(&backup)); + let explicit = snapshot(&[(&primary, 0, 0, 100, 100), (&backup, 0, 10, 100, 100)]); + let opposite = snapshot(&[(&primary, 0, 10, 100, 100), (&backup, 0, 0, 100, 100)]); + let opposite_decision = select_with_snapshot(&workers, Some(&opposite)) + .expect("the opposite snapshot has the same legal workers"); + assert_eq!(opposite_decision.id, backup.id); + + let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &explicit) + .expect("capacity exhaustion must degrade to Power-of-Two"); + + assert_eq!(decision.selected.id, primary.id); + assert_eq!(decision.load_snapshot_version, explicit.version); + } + + #[test] + fn prefill_pressure_uses_waiting_then_running_requests() { + let busy = worker("busy"); + let idle = worker("idle"); + let loads = snapshot(&[(&busy, 1, 8, 10, 100), (&idle, 9, 2, 90, 100)]); + assert!(compare_prefill_pressure(&busy, &idle, Some(&loads)).is_gt()); + } + + #[test] + fn missing_snapshot_uses_local_active_load() { + let left = worker("left"); + let right = worker("right"); + let _guard = left.load_guard(); + assert!(compare_prefill_pressure(&left, &right, None).is_gt()); + } +} diff --git a/experimental/sgl-router/src/policies/cache_aware.rs b/experimental/sgl-router/src/policies/cache_aware.rs new file mode 100644 index 000000000..a6f10b0e3 --- /dev/null +++ b/experimental/sgl-router/src/policies/cache_aware.rs @@ -0,0 +1,197 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Builds bounded cache-aware candidates from ingress Indexer results. + +use crate::config::AffinityConfig; +use crate::policies::admission::FreshLoadLookup; +use crate::policies::power_of_two::PowerOfTwoChoicesPolicy; +use crate::policies::{ + CacheCandidate, CacheCandidateProposal, Policy, PrefillProposal, ProposalKind, + SelectionContext, SelectionProposal, +}; +use crate::workers::Worker; +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +#[derive(Debug)] +pub struct CacheAwarePolicy { + config: AffinityConfig, +} + +impl CacheAwarePolicy { + pub fn new(config: AffinityConfig) -> Self { + Self { config } + } + + fn cache_candidate_proposal( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + let input_tokens = ctx.input_tokens()?; + let signal = ctx.external_prefix()?; + let sgl_kv_indexer::PrefixOutcome::Matched { matches, .. } = &signal.outcome else { + return None; + }; + if signal.query_blocks == 0 || workers.is_empty() { + return None; + } + + // The #33370 indexer contract routes on the worker address (matched + // byte-for-byte against registered worker URLs); worker_id is for + // logs only. + let by_url: HashMap<&str, &Arc> = workers + .iter() + .map(|worker| (worker.url.as_str(), worker)) + .collect(); + let mut seen = HashSet::new(); + let mut candidates = Vec::new(); + for entry in matches { + let Some(worker) = by_url.get(entry.address.as_str()) else { + continue; + }; + if entry.matched_prefix_blocks == 0 || !seen.insert(worker.id.clone()) { + continue; + } + let matched_prefix_tokens = estimate_matched_prefix_tokens( + input_tokens, + signal.query_blocks, + entry.matched_prefix_blocks, + ); + if !self.passes_cache_gate(input_tokens, matched_prefix_tokens) { + continue; + } + candidates.push(CacheCandidate { + worker: Arc::clone(worker), + matched_prefix_tokens, + uncached_tokens: input_tokens.saturating_sub(matched_prefix_tokens), + candidate_range_id: ctx.candidate_range_id().to_string(), + max_pending_prefill_tokens: None, + }); + } + + let limit = self.candidate_limit(workers.len()); + if limit == 0 { + return None; + } + let loads = FreshLoadLookup::new( + ctx.load_snapshot(), + candidates.iter().map(|candidate| &candidate.worker), + ); + if candidates.len() > limit { + candidates.select_nth_unstable_by(limit, |left, right| { + compare_candidate_seed(left, right, &loads) + }); + candidates.truncate(limit); + } + candidates.sort_by(|left, right| compare_candidate_seed(left, right, &loads)); + if candidates.is_empty() { + return None; + } + Some(CacheCandidateProposal { + candidates, + cache_switch_margin_tokens: self.config.cache_switch_margin_tokens, + }) + } + + fn passes_cache_gate(&self, input_tokens: u64, matched_prefix_tokens: u64) -> bool { + self.config + .cache_affinity_min_matched_tokens + .is_none_or(|minimum| matched_prefix_tokens >= minimum) + && self + .config + .cache_affinity_min_match_ratio + .is_none_or(|minimum| { + input_tokens > 0 + && matched_prefix_tokens as f64 / input_tokens as f64 >= minimum + }) + } + + fn candidate_limit(&self, worker_count: usize) -> usize { + let proportional = (self.config.cache_candidate_ratio.clamp(0.0, 1.0) * worker_count as f64) + .ceil() as usize; + worker_count + .min(self.config.cache_candidate_max_workers) + .min(self.config.cache_candidate_min_workers.max(proportional)) + } +} + +fn compare_candidate_seed( + left: &CacheCandidate, + right: &CacheCandidate, + loads: &FreshLoadLookup<'_>, +) -> Ordering { + right + .matched_prefix_tokens + .cmp(&left.matched_prefix_tokens) + .then_with(|| loads.compare_prefill_pressure(&left.worker, &right.worker)) + .then_with(|| left.worker.id.0.cmp(&right.worker.id.0)) +} + +impl Policy for CacheAwarePolicy { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + self.propose(workers, ctx).map(|proposal| proposal.primary) + } + + fn propose( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + match self.propose_prefill(workers, ctx)? { + PrefillProposal::Pair(proposal) => Some(proposal), + PrefillProposal::CacheCandidates(proposal) => { + let candidate = proposal.candidates.into_iter().next()?; + Some( + SelectionProposal::primary(candidate.worker) + .with_kind(ProposalKind::CacheAffinity), + ) + } + } + } + + fn propose_prefill( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + if ctx.affinity_lookup_enabled() { + if let Some(proposal) = self.cache_candidate_proposal(workers, ctx) { + return Some(PrefillProposal::CacheCandidates(proposal)); + } + } + PowerOfTwoChoicesPolicy::new() + .propose(workers, ctx) + .map(PrefillProposal::Pair) + } + + fn needs_request_tokens(&self) -> bool { + true + } + + fn uses_shared_prefill_admission(&self) -> bool { + true + } +} + +fn estimate_matched_prefix_tokens( + input_tokens: u64, + query_blocks: usize, + matched_prefix_blocks: u32, +) -> u64 { + let query_blocks = u64::try_from(query_blocks).unwrap_or(u64::MAX).max(1); + let matched_prefix_blocks = u64::from(matched_prefix_blocks).min(query_blocks); + input_tokens.saturating_mul(matched_prefix_blocks) / query_blocks +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matched_token_estimate_caps_untrusted_block_count() { + assert_eq!(estimate_matched_prefix_tokens(80, 8, 99), 80); + } +} diff --git a/experimental/sgl-router/src/policies/cache_aware_zmq.rs b/experimental/sgl-router/src/policies/cache_aware_zmq.rs index 3332e8929..731960b3a 100644 --- a/experimental/sgl-router/src/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/src/policies/cache_aware_zmq.rs @@ -13,6 +13,9 @@ //! caller) and a `SelectionContext` carrying the JSON request body and the //! ingress-precomputed routing tokens: //! +//! Load comparisons use [`WorkerLoads::load_of`], which owns fresh-snapshot +//! correction. +//! //! 1. **Load-imbalance fast-path.** If `max_load - min_load > //! balance_abs_threshold` AND `max_load > min_load * //! balance_rel_threshold`, skip the cache lookup and pick the @@ -28,8 +31,7 @@ //! for the longest matching prefix. If `match_rate > cache_threshold`, //! pick the lowest-load worker whose `url` appears in the match result. //! Otherwise, fall through. -//! 4. **Min-load fallback.** Pick the lowest-load worker by -//! `Worker::active_load()`. +//! 4. **Min-load fallback.** Pick the lowest-load worker. //! //! The implementation never returns `None` for a non-empty `workers` slice; //! a misconfigured tree or tokenizer degrades to round-robin-with-load @@ -37,6 +39,7 @@ use crate::config::CacheAwareConfig; +use crate::policies::engine_load::{EngineLoadSnapshot, EngineLoadTable}; use crate::policies::kv_events::{ compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree, }; @@ -44,7 +47,9 @@ use crate::policies::{request_tokens_for, Policy, SelectionContext}; use crate::server::metrics::MetricsRegistry; use crate::tokenizer::TokenizerRegistry; use crate::workers::Worker; +use std::collections::HashMap; use std::sync::{Arc, OnceLock}; +use std::time::Instant; /// Selection policy that scores candidates by tree-overlap with the /// request's prefix and falls back to load-based picking when the tree @@ -63,6 +68,12 @@ pub struct CacheAwareZmqPolicy { /// degrades to min-load — the router cannot hash a prompt without /// a block size that matches what the worker publishes. block_size_oracle: Arc, + /// Engine-reported per-worker load (running + waiting), shared with the + /// `KvEventIndex` load subscriber. Read once per selection; a worker with + /// a fresh snapshot uses it in place of the router-side in-flight counter + /// (`Worker::active_load`), falling back to that counter when the snapshot + /// is stale or absent (cold start / worker predates load publishing). + engine_load: Arc, /// Optional metrics sink. Set via [`Self::with_metrics`] by the policy /// factory for the production policy; `None` in unit tests and /// non-cache-aware call sites. When set, each cache-aware selection @@ -82,18 +93,114 @@ impl std::fmt::Debug for CacheAwareZmqPolicy { } } +/// Snapshot of the load-imbalance check, carried out of +/// [`CacheAwareZmqPolicy::balance_check`] so the caller can log the +/// numbers behind a rebalance decision. +struct BalanceCheck { + min_load: usize, + max_load: usize, + abs_diff: usize, + imbalanced: bool, +} + +/// Per-selection load lookup. Built once per `select` from a single +/// [`EngineLoadTable::fresh_worker_state`] pass: a worker with a fresh +/// engine-reported snapshot uses its queue depth (`num_running + +/// num_waiting`) plus its own dispatches acquired since that snapshot's +/// timestamp (see [`Self::load_of`]); otherwise it falls back to the +/// router-side in-flight counter (`Worker::active_load`). Holding the +/// snapshot keeps every per-worker `load_of` an O(1) map lookup. +struct WorkerLoads { + /// url -> (engine-reported depth, that snapshot's oldest-rank timestamp). + fresh: HashMap, +} + +impl WorkerLoads { + /// Build the per-selection snapshot from one `fresh_worker_state` pass. + /// The single construction chokepoint guarantees every comparison in a + /// given `select` sees one consistent view of load. + fn from_engine(table: &EngineLoadTable, now: Instant) -> Self { + Self { + fresh: table.fresh_worker_state(now), + } + } + + /// Builds a load view from the ingress snapshot for the current candidates. + fn from_snapshot(snapshot: &EngineLoadSnapshot, workers: &[Arc]) -> Self { + let fresh = workers + .iter() + .filter_map(|worker| { + snapshot.fresh_load_for_url(&worker.url).map(|load| { + ( + worker.url.clone(), + ( + load.num_running_reqs + .saturating_add(load.num_waiting_reqs) + .try_into() + .unwrap_or(usize::MAX), + load.captured_at, + ), + ) + }) + }) + .collect(); + Self { fresh } + } + + /// A worker's current load: the engine-reported queue depth as of the + /// last fresh snapshot, plus this worker's own dispatches made *since* + /// that snapshot's timestamp — i.e. exactly the requests the engine + /// hasn't had a chance to report back on yet. This is deliberately not + /// the worker's full `active_load()`: that counter also includes + /// long-held slots from slow-draining streaming responses (see + /// `crate::proxy::Proxy::forward_streaming_to`'s `stream_guards` doc) + /// that the engine's own last report has likely already accounted for — + /// adding the full counter on top would bias selection away from workers + /// that are idle on the engine side but still slowly draining a finished + /// stream to a client. + /// + /// This correction is per-router-process: it only sees dispatches THIS + /// router pod made. It closes the single-pod stale-gauge herd, but does + /// not coordinate with other router replicas — two pods can still both + /// read the same stale engine number and independently pile onto the + /// same worker within one gauge-refresh window. Closing that would need + /// cross-replica state sharing, which this fix does not attempt. + fn load_of(&self, w: &Worker) -> usize { + match self.fresh.get(w.url.as_str()) { + // `saturating_add`, not an assertable invariant: both operands + // are bounded by real concurrency limits (a worker's in-flight + // count is bounded well below `usize::MAX` by connection and + // request-rate limits upstream of the router), so overflow here + // is unreachable from real traffic — reaching it would mean a + // problem (memory exhaustion, a corrupt engine payload) that is + // already symptomatic elsewhere, not something worth a panic on + // this per-request hot path. + Some(&(engine_load, at)) => engine_load.saturating_add(w.slots_acquired_since(at)), + None => w.active_load(), + } + } + + /// Number of workers whose load came from the engine (vs the router-side + /// fallback). Used only to annotate the rebalance log. + fn engine_worker_count(&self) -> usize { + self.fresh.len() + } +} + impl CacheAwareZmqPolicy { pub fn new( config: CacheAwareConfig, tree: Arc, tokenizers: Arc, block_size_oracle: Arc, + engine_load: Arc, ) -> Self { Self { config, tree, tokenizers, block_size_oracle, + engine_load, metrics: OnceLock::new(), } } @@ -107,28 +214,44 @@ impl CacheAwareZmqPolicy { self } - /// Lowest-load worker — ties broken by stable iteration order (which - /// is the order the registry returned, i.e. dashmap-undefined). For - /// production traffic the ties are rare; tests pin the load skew. - fn pick_min_load(workers: &[Arc]) -> Option> { + /// Lowest-load worker by the per-selection load lookup — ties broken by + /// stable iteration order (the order the registry returned, i.e. + /// dashmap-undefined). For production traffic the ties are rare; tests + /// pin the load skew. + fn pick_min_load(workers: &[Arc], loads: &WorkerLoads) -> Option> { workers .iter() - .min_by_key(|w| w.active_load()) + .min_by_key(|w| loads.load_of(w)) .map(Arc::clone) } - /// Detect load imbalance. Returns `true` when the spread between max + /// Detect load imbalance. Returns the min/max load snapshot together + /// with the `imbalanced` verdict — `true` when the spread between max /// and min load is large enough that cache-aware routing would dump - /// even more on the hot worker. - fn is_imbalanced(&self, workers: &[Arc]) -> bool { + /// even more on the hot worker. The caller logs these numbers so every + /// rebalance decision is visible in the logs. + /// + /// `min_load`/`max_load` are [`WorkerLoads::load_of`] values, i.e. for a + /// worker with a fresh engine snapshot this is the engine-reported depth + /// PLUS this router's own not-yet-reported dispatches — not the raw + /// engine number alone. An on-call reader comparing this log's + /// `max_load` against the engine's own `/metrics` queue depth during an + /// incident should expect them to differ by that correction. + fn balance_check(&self, workers: &[Arc], loads: &WorkerLoads) -> BalanceCheck { let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(mn, mx), w| { - let l = w.active_load(); + let l = loads.load_of(w); (mn.min(l), mx.max(l)) }); let min_load = if min_load == usize::MAX { 0 } else { min_load }; let abs_diff = max_load.saturating_sub(min_load); let rel_threshold = (min_load as f32 * self.config.balance_rel_threshold) as usize; - abs_diff > self.config.balance_abs_threshold && max_load > rel_threshold + let imbalanced = abs_diff > self.config.balance_abs_threshold && max_load > rel_threshold; + BalanceCheck { + min_load, + max_load, + abs_diff, + imbalanced, + } } fn select_external( @@ -136,6 +259,7 @@ impl CacheAwareZmqPolicy { workers: &[Arc], ctx: &SelectionContext<'_>, signal: &crate::policies::ExternalPrefixSignal, + loads: &WorkerLoads, ) -> Option> { let sgl_kv_indexer::PrefixOutcome::Matched { matches, .. } = &signal.outcome else { return None; @@ -144,6 +268,7 @@ impl CacheAwareZmqPolicy { return None; } + // The index may include unhealthy workers or workers in another pool. let best_routable_blocks = matches .iter() .filter(|m| workers.iter().any(|worker| worker.url == m.address)) @@ -165,27 +290,73 @@ impl CacheAwareZmqPolicy { m.matched_prefix_blocks == best_routable_blocks && m.address == worker.url }) }) - .min_by_key(|worker| worker.active_load()) + .min_by_key(|worker| loads.load_of(worker)) .cloned() } } impl Policy for CacheAwareZmqPolicy { + fn needs_load_snapshot(&self) -> bool { + true + } + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { if workers.is_empty() { return None; } + // Per-selection load lookup: engine-reported queue depth where fresh, + // else the router-side in-flight counter. One snapshot pass serves + // every comparison below (imbalance check, min-load fallback, + // matched-set tiebreak). + let loads = ctx + .load_snapshot() + .map(|snapshot| WorkerLoads::from_snapshot(snapshot, workers)) + .unwrap_or_else(|| WorkerLoads::from_engine(&self.engine_load, Instant::now())); + // 1. Load-imbalance fast-path: even the best cache hit gets - // dropped in favour of evening out load. - if self.is_imbalanced(workers) { - return Self::pick_min_load(workers); + // dropped in favour of evening out load. Logged on every + // request (debug) so the input to the decision is auditable; + // the actual rebalance is logged at info when it fires. + let balance = self.balance_check(workers, &loads); + tracing::debug!( + model = %ctx.model(), + min_load = balance.min_load, + max_load = balance.max_load, + abs_diff = balance.abs_diff, + balance_abs_threshold = self.config.balance_abs_threshold, + balance_rel_threshold = self.config.balance_rel_threshold, + imbalanced = balance.imbalanced, + engine_load_workers = loads.engine_worker_count(), + engine_load_expected = self.engine_load.expected_count(), + "cache-aware-zmq: load-balance check considered", + ); + if balance.imbalanced { + let chosen = Self::pick_min_load(workers, &loads); + if let Some(w) = &chosen { + tracing::info!( + model = %ctx.model(), + worker = %w.url, + worker_load = loads.load_of(w), + min_load = balance.min_load, + max_load = balance.max_load, + abs_diff = balance.abs_diff, + balance_abs_threshold = self.config.balance_abs_threshold, + balance_rel_threshold = self.config.balance_rel_threshold, + engine_load_workers = loads.engine_worker_count(), + engine_load_expected = self.engine_load.expected_count(), + "cache-aware-zmq: load imbalance detected — bypassing cache, routing to min-load worker", + ); + } + return chosen; } + // An external signal is authoritative; empty or unusable results + // fall back to min-load without consulting the local radix tree. if let Some(signal) = ctx.external_prefix() { return self - .select_external(workers, ctx, signal) - .or_else(|| Self::pick_min_load(workers)); + .select_external(workers, ctx, signal, &loads) + .or_else(|| Self::pick_min_load(workers, &loads)); } // 2. Routing tokens. Prefer the ids computed once at ingress; fall @@ -198,13 +369,13 @@ impl Policy for CacheAwareZmqPolicy { _ => { let body = match ctx.request_body() { Some(b) if !b.is_empty() => b, - _ => return Self::pick_min_load(workers), + _ => return Self::pick_min_load(workers, &loads), }; let Ok(value) = serde_json::from_slice::(body) else { - return Self::pick_min_load(workers); + return Self::pick_min_load(workers, &loads); }; let Some(rt) = request_tokens_for(&self.tokenizers, ctx.model(), &value) else { - return Self::pick_min_load(workers); + return Self::pick_min_load(workers, &loads); }; fallback_ids = rt.ids; &fallback_ids @@ -221,7 +392,7 @@ impl Policy for CacheAwareZmqPolicy { model = %ctx.model(), "cache-aware-zmq: block size unknown (no worker page_size yet), falling back to min-load", ); - return Self::pick_min_load(workers); + return Self::pick_min_load(workers, &loads); }; // EAGLE-family workers hash KV blocks over token bigrams; the query // hashes must match the worker's stored hashes or the tree lookup @@ -234,7 +405,7 @@ impl Policy for CacheAwareZmqPolicy { compute_block_hashes(tokens, block_size as usize) }; if block_hashes.is_empty() { - return Self::pick_min_load(workers); + return Self::pick_min_load(workers, &loads); } let matched = self.tree.match_prefix(None, &block_hashes); let match_rate = matched.matched_blocks as f32 / block_hashes.len() as f32; @@ -262,7 +433,7 @@ impl Policy for CacheAwareZmqPolicy { cache_threshold = self.config.cache_threshold, "cache-aware-zmq: overlap below threshold, falling back to min-load", ); - return Self::pick_min_load(workers); + return Self::pick_min_load(workers, &loads); } // Among workers in the matched set, pick the lowest-load one. let matched_urls: std::collections::HashSet<&str> = @@ -270,9 +441,9 @@ impl Policy for CacheAwareZmqPolicy { let best_matched: Option> = workers .iter() .filter(|w| matched_urls.contains(w.url.as_str())) - .min_by_key(|w| w.active_load()) + .min_by_key(|w| loads.load_of(w)) .map(Arc::clone); - let chosen = best_matched.or_else(|| Self::pick_min_load(workers)); + let chosen = best_matched.or_else(|| Self::pick_min_load(workers, &loads)); if let Some(w) = &chosen { tracing::debug!( model = %ctx.model(), @@ -298,9 +469,11 @@ mod tests { use super::*; use crate::config::CacheAwareConfig; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::engine_load::{EngineWorkerLoad, LoadStat}; use crate::policies::kv_events::tree::KvWorkerId; use crate::policies::kv_events::HashTree; use crate::tokenizer::adapter; + use std::time::Duration; fn cfg_default() -> CacheAwareConfig { CacheAwareConfig { @@ -331,6 +504,39 @@ mod tests { })) } + /// Build a policy with a fresh (empty) engine-load table, so selection + /// reads the router-side `active_load` counter — matching the + /// pre-load-aware behaviour these tests assert. + fn new_policy( + config: CacheAwareConfig, + tree: Arc, + tokenizers: Arc, + oracle: Arc, + ) -> CacheAwareZmqPolicy { + CacheAwareZmqPolicy::new(config, tree, tokenizers, oracle, EngineLoadTable::new()) + } + + /// Build a policy with an explicit engine-load table, for tests that + /// exercise engine-reported load overriding the router-side counter. + fn new_policy_with_load( + config: CacheAwareConfig, + tree: Arc, + tokenizers: Arc, + oracle: Arc, + engine_load: Arc, + ) -> CacheAwareZmqPolicy { + CacheAwareZmqPolicy::new(config, tree, tokenizers, oracle, engine_load) + } + + fn load_stat(running: u64, waiting: u64) -> LoadStat { + LoadStat { + num_running_reqs: running, + num_waiting_reqs: waiting, + num_tokens: 0, + max_total_num_tokens: 0, + } + } + fn tokenizer_registry_with_tiny() -> Arc { let cfg = crate::config::Config { server: crate::config::ServerConfig { @@ -345,6 +551,7 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, @@ -363,7 +570,7 @@ mod tests { #[test] fn empty_workers_returns_none() { let tree = Arc::new(HashTree::new()); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( cfg_default(), tree, tokenizer_registry_with_tiny(), @@ -378,7 +585,7 @@ mod tests { #[test] fn empty_tree_falls_back_to_min_load() { let tree = Arc::new(HashTree::new()); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( cfg_default(), tree, tokenizer_registry_with_tiny(), @@ -398,7 +605,7 @@ mod tests { } #[test] - fn external_prefix_signal_selects_the_best_routable_match() { + fn external_prefix_signal_skips_unroutable_best_match() { let mut config = cfg_default(); config.cache_threshold = 0.0; let policy = CacheAwareZmqPolicy::new( @@ -406,6 +613,7 @@ mod tests { Arc::new(HashTree::new()), tokenizer_registry_with_tiny(), oracle_for_tests(4), + EngineLoadTable::new(), ); let w0 = worker("http://w0:30000", "tiny"); let w1 = worker("http://w1:30000", "tiny"); @@ -436,7 +644,7 @@ mod tests { } #[test] - fn external_empty_result_uses_min_load() { + fn external_empty_result_uses_min_load_without_local_tree() { let tree = Arc::new(HashTree::new()); let registry = tokenizer_registry_with_tiny(); let text = "hello world hello world hello world"; @@ -445,7 +653,13 @@ mod tests { let hashes = compute_block_hashes(&ids, 4); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new(cfg_default(), tree, registry, oracle_for_tests(4)); + let policy = CacheAwareZmqPolicy::new( + cfg_default(), + tree, + registry, + oracle_for_tests(4), + EngineLoadTable::new(), + ); let w0 = worker("http://w0:30000", "tiny"); let w1 = worker("http://w1:30000", "tiny"); let _load = w0.load_guard(); @@ -461,6 +675,79 @@ mod tests { assert_eq!(chosen.url, w1.url); } + /// Equal external cache matches are resolved from the request snapshot, + /// not router-local load that changes after ingress. + #[test] + fn external_match_tiebreak_uses_the_request_snapshot() { + let mut config = cfg_default(); + config.cache_threshold = 0.0; + config.balance_abs_threshold = 100; + let policy = new_policy( + config, + Arc::new(HashTree::new()), + tokenizer_registry_with_tiny(), + oracle_for_tests(4), + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + // Local load disagrees with the snapshot and must not affect this choice. + let _after_snapshot: Vec<_> = (0..10).map(|_| w1.load_guard()).collect(); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let snapshot_at = Instant::now(); + let snapshot = EngineLoadSnapshot::from_workers( + 17, + HashMap::from([ + ( + w0.url.clone(), + EngineWorkerLoad { + num_running_reqs: 50, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: snapshot_at, + }, + ), + ( + w1.url.clone(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: snapshot_at, + }, + ), + ]), + ); + let signal = crate::policies::ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![ + sgl_kv_indexer::PrefixMatch { + address: w0.url.clone(), + matched_prefix_blocks: 4, + worker_id: w0.id.0.clone(), + }, + sgl_kv_indexer::PrefixMatch { + address: w1.url.clone(), + matched_prefix_blocks: 4, + worker_id: w1.id.0.clone(), + }, + ], + best_prefix_blocks: 4, + }, + query_blocks: 4, + }; + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None) + .with_external_prefix(Some(&signal)) + .with_load_snapshot(&snapshot); + + assert_eq!( + policy.select(&workers, &ctx).expect("must select").url, + w1.url, + "external-match tiebreak must use the request snapshot, not later active load" + ); + } /// Tree contains w0's prefix; cache-aware selection picks w0 even /// though w1 has lower load (the load skew is below the imbalance /// threshold, so cache wins). @@ -483,7 +770,7 @@ mod tests { ); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, // any match counts balance_abs_threshold: 32, @@ -521,7 +808,7 @@ mod tests { tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); let metrics = MetricsRegistry::new(); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -567,7 +854,7 @@ mod tests { assert!(!hashes.is_empty()); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -619,7 +906,7 @@ mod tests { tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); let metrics = MetricsRegistry::new(); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 1.0, // match_rate <= 1.0 always -> always fall back balance_abs_threshold: 32, @@ -703,7 +990,7 @@ mod tests { oracle.try_set(block_size).unwrap(); oracle.set_bigram(true); let metrics = MetricsRegistry::new(); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -743,7 +1030,7 @@ mod tests { let oracle = BlockSizeOracle::new(); oracle.try_set(block_size).unwrap(); let metrics = MetricsRegistry::new(); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -802,7 +1089,7 @@ mod tests { &templated_hashes, ); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -875,7 +1162,7 @@ mod tests { let tree = Arc::new(HashTree::new()); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -914,7 +1201,7 @@ mod tests { ); let tree = Arc::new(HashTree::new()); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -1018,7 +1305,7 @@ mod tests { tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); tree.insert(&KvWorkerId::new("http://w1:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -1054,7 +1341,7 @@ mod tests { let hashes = compute_block_hashes(&ids, block_size as usize); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, // would normally always match balance_abs_threshold: 5, @@ -1080,14 +1367,321 @@ mod tests { assert_eq!(chosen.url, "http://w1:30000", "imbalance must dominate"); } + /// Fresh engine-reported load drives the imbalance + min-load decision + /// instead of the router-side in-flight counter. Both workers hold the + /// prefix and have zero router-side load, so without engine load the + /// tiebreak would pick w0 (stable order). Engine load says w0 is hot + /// (50) and w1 is light (1) → the imbalance branch routes to w1. + #[test] + fn engine_load_overrides_active_load() { + let tree = Arc::new(HashTree::new()); + let registry = tokenizer_registry_with_tiny(); + let text = "hello world hello world hello world"; + let tok = registry.get("tiny").unwrap(); + let ids = adapter::encode(&tok, text).unwrap(); + let hashes = compute_block_hashes(&ids, 4); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + tree.insert(&KvWorkerId::new("http://w1:30000".into(), 0), None, &hashes); + + let engine_load = EngineLoadTable::new(); + let now = Instant::now(); + engine_load.set("http://w0:30000", 0, load_stat(50, 0), now); + engine_load.set("http://w1:30000", 0, load_stat(1, 0), now); + + let policy = new_policy_with_load( + CacheAwareConfig { + cache_threshold: 0.0, + balance_abs_threshold: 5, + balance_rel_threshold: 2.0, + kv_indexer_endpoint: None, + }, + tree, + registry, + oracle_for_tests(4), + engine_load, + ); + // Router-side counters are both 0 — only engine load is skewed. + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w1:30000", + "engine-reported load must drive selection", + ); + } + + /// A request snapshot must override load updates that arrive after ingress. + #[test] + fn request_snapshot_is_stable_after_new_load_stats_arrive() { + let table = EngineLoadTable::new(); + let snapshot_at = Instant::now(); + let snapshot = EngineLoadSnapshot::from_workers( + 9, + HashMap::from([ + ( + "http://w0:30000".to_string(), + EngineWorkerLoad { + num_running_reqs: 50, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: snapshot_at, + }, + ), + ( + "http://w1:30000".to_string(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: snapshot_at, + }, + ), + ]), + ); + // The later gauge disagrees with the captured view; this request still picks w1. + table.set("http://w0:30000", 0, load_stat(1, 0), Instant::now()); + table.set("http://w1:30000", 0, load_stat(50, 0), Instant::now()); + let policy = new_policy_with_load( + cfg_default(), + Arc::new(HashTree::new()), + tokenizer_registry_with_tiny(), + oracle_for_tests(4), + table, + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + + assert_eq!( + policy.select(&workers, &ctx).expect("must select").url, + w1.url, + "the request must use its frozen snapshot, not the newer table value" + ); + } + + /// When load is balanced enough that the imbalance branch does NOT fire, + /// the matched-set tiebreak still uses engine load: both workers hold the + /// prefix, engine load says w1 is lighter → w1 wins. (Guards against a + /// regression that reverted the tiebreak to `active_load()`.) + #[test] + fn matched_set_tiebreak_uses_engine_load() { + let tree = Arc::new(HashTree::new()); + let registry = tokenizer_registry_with_tiny(); + let text = "hello world hello world hello world"; + let tok = registry.get("tiny").unwrap(); + let ids = adapter::encode(&tok, text).unwrap(); + let hashes = compute_block_hashes(&ids, 4); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + tree.insert(&KvWorkerId::new("http://w1:30000".into(), 0), None, &hashes); + + let engine_load = EngineLoadTable::new(); + let now = Instant::now(); + engine_load.set("http://w0:30000", 0, load_stat(10, 0), now); + engine_load.set("http://w1:30000", 0, load_stat(2, 0), now); + + let policy = new_policy_with_load( + CacheAwareConfig { + cache_threshold: 0.0, + // High thresholds so the imbalance fast-path never fires (10 vs + // 2) and selection reaches the matched-set tiebreak. + balance_abs_threshold: 100, + balance_rel_threshold: 100.0, + kv_indexer_endpoint: None, + }, + tree, + registry, + oracle_for_tests(4), + engine_load, + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w1:30000", + "matched-set tiebreak must use engine load", + ); + } + + /// Recent dispatches made AFTER the engine's last snapshot are added on + /// top of the reported load. Without this, repeated `select` calls in + /// the same burst would all read the same "worker looks idle" engine + /// number and all pile onto it before the gauge catches up. w0 looks + /// lighter by the raw engine numbers alone (1 vs 3), but three slots + /// claimed on w0 after the snapshot flip the effective load in w1's + /// favor (1+3=4 > 3+0=3). + #[test] + fn recent_dispatches_are_added_on_top_of_engine_load() { + let tree = Arc::new(HashTree::new()); + let registry = tokenizer_registry_with_tiny(); + let text = "hello world hello world hello world"; + let tok = registry.get("tiny").unwrap(); + let ids = adapter::encode(&tok, text).unwrap(); + let hashes = compute_block_hashes(&ids, 4); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + tree.insert(&KvWorkerId::new("http://w1:30000".into(), 0), None, &hashes); + + let engine_load = EngineLoadTable::new(); + let snapshot_at = Instant::now(); + engine_load.set("http://w0:30000", 0, load_stat(1, 0), snapshot_at); + engine_load.set("http://w1:30000", 0, load_stat(3, 0), snapshot_at); + + let policy = new_policy_with_load( + CacheAwareConfig { + cache_threshold: 0.0, + // High thresholds so the imbalance fast-path never fires on + // the raw engine numbers (1 vs 3) and selection reaches the + // matched-set tiebreak, which also uses `load_of`. + balance_abs_threshold: 100, + balance_rel_threshold: 100.0, + kv_indexer_endpoint: None, + }, + tree, + registry, + oracle_for_tests(4), + engine_load, + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + // Three requests dispatched to w0 AFTER the engine's snapshot — + // exactly the "burst the engine hasn't reported back on yet" shape. + let _g1 = w0.timestamped_load_guard(); + let _g2 = w0.timestamped_load_guard(); + let _g3 = w0.timestamped_load_guard(); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w1:30000", + "w0's effective load (1 engine + 3 recent = 4) must exceed w1's \ + (3 engine + 0 recent = 3), even though the raw engine numbers \ + alone favor w0", + ); + } + + /// `load_of` must use the OLDEST rank's timestamp as the "since" cutoff + /// for a multi-rank worker, not the newest — this pins the end-to-end + /// wiring of the choice `EngineLoadTable::fresh_worker_state` makes (see + /// its doc comment). A regression to "newest" would silently treat the + /// dispatch below as already covered by rank1's later snapshot, even + /// though rank0's older snapshot doesn't reflect it. + #[test] + fn load_of_uses_oldest_rank_timestamp_for_multi_rank_worker() { + let engine_load = EngineLoadTable::new(); + let earlier = Instant::now(); + let w = worker("http://w:30000", "tiny"); + // Real sleeps, not synthetic `Instant` offsets: the dispatch's + // timestamp is captured internally by `timestamped_load_guard()` and isn't + // injectable (see `worker.rs`'s `slots_acquired_since` tests for the + // same reasoning). + std::thread::sleep(Duration::from_millis(5)); + let _g = w.timestamped_load_guard(); // dispatched strictly between earlier/later + std::thread::sleep(Duration::from_millis(5)); + let later = Instant::now(); + engine_load.set("http://w:30000", 0, load_stat(1, 0), earlier); + engine_load.set("http://w:30000", 1, load_stat(1, 0), later); + + let loads = WorkerLoads::from_engine(&engine_load, later); + assert_eq!( + loads.load_of(&w), + 3, + "depth (1+1=2) plus the one dispatch made after the OLDEST \ + rank's timestamp = 3; using the newest rank's timestamp \ + instead would exclude that dispatch and wrongly give 2", + ); + } + + /// A stale engine snapshot falls back to PURE `active_load()` — the + /// recent-dispatch correction only applies alongside a fresh snapshot + /// (see `load_of`'s `Some` branch). A regression that added + /// `slots_acquired_since` to the fallback branch too would double-count + /// this worker's own in-flight guards. + #[test] + fn load_of_fallback_does_not_add_recent_dispatches_on_top_of_active_load() { + let engine_load = EngineLoadTable::new(); + let stale = Instant::now() - Duration::from_secs(3600); + engine_load.set("http://w:30000", 0, load_stat(50, 0), stale); + let w = worker("http://w:30000", "tiny"); + let _g1 = w.load_guard(); + let _g2 = w.load_guard(); + + let loads = WorkerLoads::from_engine(&engine_load, Instant::now()); + assert_eq!( + loads.load_of(&w), + 2, + "must equal active_load() exactly (2) — not the stale depth \ + (50) plus anything, and not active_load() plus a second \ + correction", + ); + } + + /// A stale engine snapshot is ignored: selection falls back to the + /// router-side `active_load` counter. w0's (stale) engine load is high, + /// but w1 carries a router-side guard, so fallback picks w0. + #[test] + fn stale_engine_load_falls_back_to_active_load() { + let tree = Arc::new(HashTree::new()); + let registry = tokenizer_registry_with_tiny(); + let text = "hello world hello world hello world"; + let tok = registry.get("tiny").unwrap(); + let ids = adapter::encode(&tok, text).unwrap(); + let hashes = compute_block_hashes(&ids, 4); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + tree.insert(&KvWorkerId::new("http://w1:30000".into(), 0), None, &hashes); + + // Past the default freshness window; an hour ago is comfortably stale. + let engine_load = EngineLoadTable::new(); + let stale = Instant::now() - Duration::from_secs(3600); + engine_load.set("http://w0:30000", 0, load_stat(50, 0), stale); + + let policy = new_policy_with_load( + CacheAwareConfig { + cache_threshold: 0.0, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, + }, + tree, + registry, + oracle_for_tests(4), + engine_load, + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + // Router-side: w1 has one in-flight request, w0 has none. With the + // stale engine load ignored, the tiebreak picks w0 (load 0 < 1). + let _g = w1.load_guard(); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!( + chosen.url, "http://w0:30000", + "stale engine load must be ignored in favour of active_load", + ); + } + /// Tokenizer is missing for the requested model → fall back to /// min-load (no panic, no error). #[test] fn missing_tokenizer_falls_back_to_min_load() { let tree = Arc::new(HashTree::new()); let empty_registry = Arc::new(TokenizerRegistry::default()); - let policy = - CacheAwareZmqPolicy::new(cfg_default(), tree, empty_registry, oracle_for_tests(4)); + let policy = new_policy(cfg_default(), tree, empty_registry, oracle_for_tests(4)); let w0 = worker("http://w0:30000", "tiny"); let w1 = worker("http://w1:30000", "tiny"); let _g = w0.load_guard(); @@ -1104,7 +1698,7 @@ mod tests { #[test] fn missing_request_body_falls_back_to_min_load() { let tree = Arc::new(HashTree::new()); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( cfg_default(), tree, tokenizer_registry_with_tiny(), @@ -1125,7 +1719,7 @@ mod tests { #[test] fn body_without_prompt_field_falls_back_to_min_load() { let tree = Arc::new(HashTree::new()); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( cfg_default(), tree, tokenizer_registry_with_tiny(), @@ -1149,7 +1743,7 @@ mod tests { #[test] fn empty_text_falls_back_to_min_load() { let tree = Arc::new(HashTree::new()); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( cfg_default(), tree, tokenizer_registry_with_tiny(), @@ -1180,7 +1774,7 @@ mod tests { &[999, 998, 997], ); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.99, balance_abs_threshold: 32, @@ -1264,7 +1858,7 @@ mod tests { let kw0 = KvWorkerId::new("http://w0:30000".into(), 0); tree.insert(&kw0, None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, @@ -1362,7 +1956,7 @@ mod tests { let tree = Arc::new(HashTree::new()); tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); - let policy = CacheAwareZmqPolicy::new( + let policy = new_policy( CacheAwareConfig { cache_threshold: 0.0, balance_abs_threshold: 32, diff --git a/experimental/sgl-router/src/policies/engine_load.rs b/experimental/sgl-router/src/policies/engine_load.rs new file mode 100644 index 000000000..fb6804988 --- /dev/null +++ b/experimental/sgl-router/src/policies/engine_load.rs @@ -0,0 +1,489 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Engine-reported runtime load, fed by the load subscriber. +//! +//! Workers publish a [`LoadStat`] gauge on their dedicated load socket (see +//! `python/sglang/srt/managers/scheduler_components/load_publisher.py`). The +//! load subscriber routes those into this table, keyed per +//! `(worker_url, dp_rank)`; the +//! cache-aware-zmq policy reads the freshest aggregate per worker as a +//! truthful load signal, falling back to the router-side in-flight counter +//! when no fresh snapshot exists (cold start, stale publisher, or a worker +//! that predates load publishing). +//! +//! Load is a *gauge*, not a delta: last value wins, no sequence/replay +//! semantics. Entries older than [`EngineLoadTable::freshness`] are ignored. + +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use dashmap::{DashMap, DashSet}; +use serde::de::{self, Deserializer, IgnoredAny, SeqAccess, Visitor}; +use serde::Deserialize; + +/// Per-scheduler runtime load snapshot. Mirrors the Python `LoadStat` in +/// `managers/scheduler_components/load_publisher.py`, published on the +/// worker's dedicated load socket (separate from KV-cache events). +/// +/// Wire shape (msgspec `tag=True` + `array_like`): +/// `["LoadStat", num_running_reqs, num_waiting_reqs, num_tokens, +/// max_total_num_tokens, attn_dp_rank?]`. We read the four counts and ignore +/// any trailing fields (`attn_dp_rank` — the router keys load by the +/// subscriber's socket rank, not the payload). +#[derive(Debug, Clone, PartialEq)] +pub struct LoadStat { + /// Requests currently running on the engine. + pub num_running_reqs: u64, + /// Requests queued waiting to run. + pub num_waiting_reqs: u64, + /// KV tokens currently in use. + pub num_tokens: u64, + /// KV-cache token capacity; 0 when unknown. + pub max_total_num_tokens: u64, +} + +/// Aggregated, usable Engine load for one Worker at a fixed instant. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EngineWorkerLoad { + pub num_running_reqs: u64, + pub num_waiting_reqs: u64, + pub num_tokens: u64, + pub max_total_num_tokens: u64, + pub captured_at: Instant, +} + +/// Immutable Engine-load view captured once at request ingress. +#[derive(Debug, Clone, Default)] +pub struct EngineLoadSnapshot { + pub version: u64, + workers: HashMap, +} + +impl EngineLoadSnapshot { + pub fn fresh_load_for_url(&self, worker_url: &str) -> Option<&EngineWorkerLoad> { + self.workers.get(worker_url) + } + + /// Builds a view from already validated Worker data for tests and offline checks. + pub fn from_workers(version: u64, workers: HashMap) -> Self { + Self { version, workers } + } +} + +impl<'de> Deserialize<'de> for LoadStat { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct LoadStatVisitor; + + impl<'de> Visitor<'de> for LoadStatVisitor { + type Value = LoadStat; + + fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("a tagged msgpack array [\"LoadStat\", ...fields]") + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: SeqAccess<'de>, + { + let tag: String = seq + .next_element()? + .ok_or_else(|| de::Error::missing_field("event tag"))?; + if tag != "LoadStat" { + return Err(de::Error::custom(format!( + "expected \"LoadStat\" tag, got {tag:?}" + ))); + } + // The Python publisher always emits all four counts. Treat a + // shortened frame as malformed rather than inventing zeros: + // a partial gauge must fall back to router-local load, never + // make a worker appear artificially idle. + let num_running_reqs: u64 = seq + .next_element()? + .ok_or_else(|| de::Error::missing_field("num_running_reqs"))?; + let num_waiting_reqs: u64 = seq + .next_element()? + .ok_or_else(|| de::Error::missing_field("num_waiting_reqs"))?; + let num_tokens: u64 = seq + .next_element()? + .ok_or_else(|| de::Error::missing_field("num_tokens"))?; + let max_total_num_tokens: u64 = seq + .next_element()? + .ok_or_else(|| de::Error::missing_field("max_total_num_tokens"))?; + while seq.next_element::()?.is_some() {} + Ok(LoadStat { + num_running_reqs, + num_waiting_reqs, + num_tokens, + max_total_num_tokens, + }) + } + } + + deserializer.deserialize_seq(LoadStatVisitor) + } +} + +/// Decode a single load frame's msgpack payload into a [`LoadStat`]. +pub fn decode_load_stat(payload: &[u8]) -> Result { + rmp_serde::from_slice(payload) +} + +/// A per-rank load snapshot older than this is treated as stale, so a silent +/// or slow publisher degrades to the router-side load signal rather than +/// pinning a worker at its last reported value. +const DEFAULT_FRESHNESS: Duration = Duration::from_secs(5); + +#[derive(Debug, Clone)] +struct LoadEntry { + load: LoadStat, + at: Instant, +} + +/// Per-`(worker_url, dp_rank)` engine-reported load. Written by the load +/// subscriber pump, read by the cache-aware-zmq policy. Shared out of +/// [`super::kv_events::index::KvEventIndex`] the same way the hash tree is. +#[derive(Debug)] +pub struct EngineLoadTable { + by_rank: DashMap<(String, u32), LoadEntry>, + /// Per-rank publishers the worker advertised. A worker is usable only + /// when every advertised rank has a fresh value; accepting a partial + /// aggregate would make a silent rank look idle and attract traffic. + expected: DashSet<(String, u32)>, + freshness: Duration, + version: AtomicU64, +} + +impl EngineLoadTable { + pub fn new() -> Arc { + Arc::new(Self { + by_rank: DashMap::new(), + expected: DashSet::new(), + freshness: DEFAULT_FRESHNESS, + version: AtomicU64::new(0), + }) + } + + #[cfg(test)] + pub fn with_freshness(freshness: Duration) -> Arc { + Arc::new(Self { + by_rank: DashMap::new(), + expected: DashSet::new(), + freshness, + version: AtomicU64::new(0), + }) + } + + /// Record the latest load for one `(worker_url, dp_rank)`. + pub fn set(&self, url: &str, dp_rank: u32, load: LoadStat, at: Instant) { + self.by_rank + .insert((url.to_string(), dp_rank), LoadEntry { load, at }); + self.version.fetch_add(1, Ordering::Relaxed); + } + + /// Mark one advertised scheduler rank as expected to publish load. + pub fn mark_expected_rank(&self, url: &str, dp_rank: u32) { + if self.expected.insert((url.to_string(), dp_rank)) { + self.version.fetch_add(1, Ordering::Relaxed); + } + } + + /// Number of workers expected to publish load. Compared against the size + /// of [`Self::snapshot_fresh`] to surface a dead/misconfigured publisher + /// (expected > 0 but no fresh snapshots) in logs. + pub fn expected_count(&self) -> usize { + self.expected + .iter() + .map(|entry| entry.key().0.clone()) + .collect::>() + .len() + } + + /// Shared accumulation pass behind [`Self::snapshot_fresh`] and + /// [`Self::capture_snapshot`]. It produces the #34608 fields summed across + /// ranks and the OLDEST snapshot timestamp — **but only for workers whose + /// every advertised rank is present and fresh**. A missing or stale rank is + /// omitted, so the caller falls back to its own load signal. (Summing + /// only the fresh ranks would make a worker whose other ranks went silent + /// look misleadingly idle and draw *more* traffic.) Callers that never + /// registered expected ranks retain the legacy all-known-ranks rule. + /// `snapshot_fresh` and any other consumer walking this same pass can + /// never disagree with each other about which workers count as fresh. + /// + /// The oldest (not newest) rank's timestamp is deliberately what's kept + /// alongside the depth: a caller using it as a "dispatches not yet + /// reflected in this number" cutoff (see + /// `crate::policies::cache_aware_zmq::WorkerLoads::load_of`) needs a + /// bound that never treats an unreported dispatch as already-covered — + /// the freshest rank's timestamp could do exactly that for whichever + /// rank published less recently. This conservatism is one-sided, not + /// free: for a multi-rank worker with skewed publish times, a dispatch + /// that landed on (and was already reported by) the FRESHER rank can + /// get re-added by the caller's cutoff-based correction anyway, since + /// that correction has no way to attribute a dispatch to a specific + /// rank. That's an accepted, bounded over-count (it biases the wrong + /// direction relative to the under-count this method exists to avoid, + /// not a correctness hole) rather than something this method can close + /// on its own — closing it would require per-rank dispatch attribution, + /// which the router-side slot tracking below doesn't have. + fn fresh_worker_loads(&self, now: Instant) -> HashMap { + // url -> rank -> (reported load, fresh, timestamp). + let mut observed: HashMap> = HashMap::new(); + for entry in self.by_rank.iter() { + let at = entry.value().at; + let fresh = now.saturating_duration_since(at) <= self.freshness; + observed + .entry(entry.key().0.clone()) + .or_default() + .insert(entry.key().1, (entry.value().load.clone(), fresh, at)); + } + let mut expected: HashMap> = HashMap::new(); + for entry in self.expected.iter() { + expected + .entry(entry.key().0.clone()) + .or_default() + .insert(entry.key().1); + } + + let workers: HashSet = observed.keys().chain(expected.keys()).cloned().collect(); + workers + .into_iter() + .filter_map(|url| { + let ranks = observed.get(&url)?; + let required: Vec = match expected.get(&url) { + Some(expected_ranks) => expected_ranks.iter().copied().collect(), + None => ranks.keys().copied().collect(), + }; + let mut num_running_reqs = 0u64; + let mut num_waiting_reqs = 0u64; + let mut num_tokens = 0u64; + let mut max_total_num_tokens = 0u64; + let mut oldest_at = None; + for rank in required { + let (load, fresh, at) = ranks.get(&rank)?; + if !fresh { + return None; + } + num_running_reqs = num_running_reqs.saturating_add(load.num_running_reqs); + num_waiting_reqs = num_waiting_reqs.saturating_add(load.num_waiting_reqs); + num_tokens = num_tokens.saturating_add(load.num_tokens); + max_total_num_tokens = + max_total_num_tokens.saturating_add(load.max_total_num_tokens); + oldest_at = Some(oldest_at.map_or(*at, |oldest: Instant| oldest.min(*at))); + } + oldest_at.map(|captured_at| { + ( + url, + EngineWorkerLoad { + num_running_reqs, + num_waiting_reqs, + num_tokens, + max_total_num_tokens, + captured_at, + }, + ) + }) + }) + .collect() + } + + /// Captures the immutable view consumed by all routing decisions in one request. + pub fn capture_snapshot(&self, now: Instant) -> EngineLoadSnapshot { + EngineLoadSnapshot { + version: self.version.load(Ordering::Acquire), + workers: self.fresh_worker_loads(now), + } + } + + pub(crate) fn fresh_worker_state(&self, now: Instant) -> HashMap { + self.fresh_worker_loads(now) + .into_iter() + .map(|(url, load)| { + ( + url, + ( + load.num_running_reqs + .saturating_add(load.num_waiting_reqs) + .try_into() + .unwrap_or(usize::MAX), + load.captured_at, + ), + ) + }) + .collect() + } + + /// Per worker URL, the summed queue depth (`num_running_reqs + + /// num_waiting_reqs`) across that worker's ranks, for workers whose + /// every advertised rank is fresh. Computed once per selection so per-worker + /// lookups are O(1). See [`Self::fresh_worker_state`] for the freshness + /// gate behind this. + pub fn snapshot_fresh(&self, now: Instant) -> HashMap { + self.fresh_worker_state(now) + .into_iter() + .map(|(url, (depth, _))| (url, depth)) + .collect() + } + + /// Drop every rank entry (and the expected mark) for a worker. Called on + /// worker removal so a re-added worker does not leave stale load behind. + pub fn forget_worker(&self, url: &str) { + self.by_rank.retain(|k, _| k.0 != url); + self.expected.retain(|key| key.0 != url); + self.version.fetch_add(1, Ordering::Relaxed); + } + + #[cfg(test)] + pub fn entry_count(&self) -> usize { + self.by_rank.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn load(running: u64, waiting: u64) -> LoadStat { + LoadStat { + num_running_reqs: running, + num_waiting_reqs: waiting, + num_tokens: 0, + max_total_num_tokens: 0, + } + } + + #[test] + fn load_wire_rejects_wrong_tag_and_missing_counts() { + let mut wrong_tag = Vec::new(); + rmp::encode::write_array_len(&mut wrong_tag, 5).unwrap(); + rmp::encode::write_str(&mut wrong_tag, "OtherStat").unwrap(); + for value in [1, 2, 3, 4] { + rmp::encode::write_u64(&mut wrong_tag, value).unwrap(); + } + assert!(decode_load_stat(&wrong_tag).is_err()); + + let mut missing_count = Vec::new(); + rmp::encode::write_array_len(&mut missing_count, 4).unwrap(); + rmp::encode::write_str(&mut missing_count, "LoadStat").unwrap(); + for value in [1, 2, 3] { + rmp::encode::write_u64(&mut missing_count, value).unwrap(); + } + assert!(decode_load_stat(&missing_count).is_err()); + } + + #[test] + fn sums_queue_depth_across_ranks() { + let t = EngineLoadTable::new(); + let now = Instant::now(); + t.set("http://w:30000", 0, load(5, 1), now); + t.set("http://w:30000", 1, load(3, 2), now); + let fresh = t.snapshot_fresh(now); + // (5+1) + (3+2) = 11 + assert_eq!(fresh.get("http://w:30000").copied(), Some(11)); + } + + #[test] + fn stale_entries_are_dropped_from_snapshot() { + let t = EngineLoadTable::with_freshness(Duration::from_millis(10)); + let old = Instant::now(); + t.set("http://w:30000", 0, load(9, 9), old); + // A read far in the future sees the entry as stale -> worker absent. + let later = old + Duration::from_secs(60); + assert!(!t.snapshot_fresh(later).contains_key("http://w:30000")); + } + + #[test] + fn forget_worker_clears_all_ranks() { + let t = EngineLoadTable::new(); + let now = Instant::now(); + t.set("http://w:30000", 0, load(1, 0), now); + t.set("http://w:30000", 1, load(1, 0), now); + t.set("http://other:30000", 0, load(1, 0), now); + t.forget_worker("http://w:30000"); + assert_eq!(t.entry_count(), 1); + assert!(!t.snapshot_fresh(now).contains_key("http://w:30000")); + assert!(t.snapshot_fresh(now).contains_key("http://other:30000")); + } + + /// A worker with any stale rank is omitted entirely (not summed over only + /// its fresh ranks), so a partially-silent worker falls back to the + /// router-side counter instead of looking misleadingly idle. + #[test] + fn partial_freshness_excludes_worker() { + let t = EngineLoadTable::with_freshness(Duration::from_secs(5)); + let now = Instant::now(); + let stale = now - Duration::from_secs(3600); + t.set("http://w:30000", 0, load(5, 1), now); // fresh + t.set("http://w:30000", 1, load(9, 9), stale); // stale + assert!( + !t.snapshot_fresh(now).contains_key("http://w:30000"), + "any stale rank must drop the whole worker from the snapshot" + ); + } + + #[test] + fn missing_expected_rank_excludes_worker() { + let t = EngineLoadTable::new(); + let now = Instant::now(); + t.mark_expected_rank("http://w:30000", 0); + t.mark_expected_rank("http://w:30000", 1); + t.set("http://w:30000", 0, load(5, 1), now); + assert!( + !t.snapshot_fresh(now).contains_key("http://w:30000"), + "an advertised rank without a reading must not produce a partial aggregate" + ); + + t.set("http://w:30000", 1, load(3, 2), now); + assert_eq!(t.snapshot_fresh(now).get("http://w:30000"), Some(&11)); + } + + #[test] + fn fresh_worker_state_picks_the_earliest_rank_timestamp() { + let t = EngineLoadTable::new(); + let earlier = Instant::now() - Duration::from_secs(2); + let later = earlier + Duration::from_secs(1); + t.set("http://w:30000", 0, load(5, 1), later); + t.set("http://w:30000", 1, load(3, 2), earlier); + let now = later + Duration::from_millis(1); + assert_eq!( + t.fresh_worker_state(now).get("http://w:30000").copied(), + Some((11, earlier)), + "must expose the OLDEST rank's timestamp, not the newest" + ); + } + + #[test] + fn fresh_worker_state_agrees_with_snapshot_fresh_on_which_workers_are_present() { + let t = EngineLoadTable::with_freshness(Duration::from_secs(5)); + let now = Instant::now(); + let stale = now - Duration::from_secs(3600); + t.set("http://fresh:30000", 0, load(1, 0), now); + t.set("http://mixed:30000", 0, load(1, 0), now); + t.set("http://mixed:30000", 1, load(1, 0), stale); + + let depths = t.snapshot_fresh(now); + let state = t.fresh_worker_state(now); + assert!(depths.contains_key("http://fresh:30000")); + assert!(state.contains_key("http://fresh:30000")); + assert!(!depths.contains_key("http://mixed:30000")); + assert!(!state.contains_key("http://mixed:30000")); + } + + #[test] + fn expected_count_tracks_marked_workers_and_forget() { + let t = EngineLoadTable::new(); + assert_eq!(t.expected_count(), 0); + t.mark_expected_rank("http://w:30000", 0); + t.mark_expected_rank("http://w:30000", 1); // same worker + t.mark_expected_rank("http://other:30000", 0); + assert_eq!(t.expected_count(), 2); + t.forget_worker("http://w:30000"); + assert_eq!(t.expected_count(), 1); + } +} diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index e2963cc58..f0e734599 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -6,7 +6,9 @@ use crate::config::{ }; use crate::discovery::ModelId; use crate::policies::{ + cache_aware::CacheAwarePolicy, cache_aware_zmq::CacheAwareZmqPolicy, + engine_load::EngineLoadTable, kv_events::{BlockSizeOracle, HashTree}, load_based::LoadBasedPolicy, power_of_two::PowerOfTwoChoicesPolicy, @@ -14,8 +16,9 @@ use crate::policies::{ round_robin::RoundRobinPolicy, scoring::{ admission::Overloaded, prefix_cache, prefix_cache::PrefixCachePolicy, FusedScorePolicy, - Pipeline, + Pipeline, ScorePolicy, }, + session_aware::SessionAwarePolicy, sticky::StickyPolicy, Policy, PolicyRegistry, }; @@ -52,8 +55,17 @@ pub fn build_policy( tree: Arc, tokenizers: Arc, block_size_oracle: Arc, + engine_load: Arc, ) -> Result> { - let inner = build_kind(model.policy, model, &tree, &tokenizers, &block_size_oracle)?; + validate_eligibility(model)?; + let inner = build_kind( + model.policy, + model, + &tree, + &tokenizers, + &block_size_oracle, + &engine_load, + )?; let Some(elig) = model.eligibility.as_ref().filter(|e| !e.filters.is_empty()) else { return Ok(inner); }; @@ -64,6 +76,26 @@ pub fn build_policy( Ok(Arc::new(Pipeline::new(filters, inner)?)) } +/// Reject configurations that can bypass CLI validation when constructed in code. +fn validate_eligibility(model: &ModelConfig) -> Result<()> { + let Some(eligibility) = model.eligibility.as_ref().filter(|e| !e.filters.is_empty()) else { + return Ok(()); + }; + + if model.policy == PolicyKind::Sticky { + return Err(anyhow!( + "eligibility filters cannot be combined with sticky policy" + )); + } + if eligibility.filters.contains(&FilterKind::Overloaded) + && eligibility.max_in_flight.unwrap_or(0) == 0 + { + return Err(anyhow!("max_in_flight must be greater than 0")); + } + + Ok(()) +} + /// Build one policy kind with the shared model dependencies. fn build_kind( kind: PolicyKind, @@ -71,6 +103,7 @@ fn build_kind( tree: &Arc, tokenizers: &Arc, block_size_oracle: &Arc, + engine_load: &Arc, ) -> Result> { let (tree, tokenizers, block_size_oracle) = ( Arc::clone(tree), @@ -89,10 +122,18 @@ fn build_kind( tree, tokenizers, block_size_oracle, + Arc::clone(engine_load), )) } + PolicyKind::SessionAware => Arc::new(SessionAwarePolicy::new( + model.affinity.clone().unwrap_or_default(), + )), + PolicyKind::CacheAware => Arc::new(CacheAwarePolicy::new( + model.affinity.clone().unwrap_or_default(), + )), PolicyKind::Sticky => build_sticky(model), PolicyKind::FusedScore => build_fused(model, &tree, &block_size_oracle)?, + PolicyKind::ScorePolicy => build_score_policy(model, &tree, &block_size_oracle)?, }) } @@ -162,6 +203,17 @@ fn build_fused( Ok(Arc::new(FusedScorePolicy::new(terms)?)) } +/// Builds the top-level `score_policy`. +fn build_score_policy( + model: &ModelConfig, + tree: &Arc, + oracle: &Arc, +) -> Result> { + Ok(Arc::new(ScorePolicy::new(build_fused( + model, tree, oracle, + )?))) +} + /// Builds a policy with test defaults. #[cfg(test)] pub fn build_policy_kind_only(kind: PolicyKind) -> Result> { @@ -175,6 +227,13 @@ pub fn build_policy_kind_only(kind: PolicyKind) -> Result> { Arc::new(HashTree::new()), Arc::new(TokenizerRegistry::default()), BlockSizeOracle::new(), + EngineLoadTable::new(), + )), + PolicyKind::SessionAware => Arc::new(SessionAwarePolicy::new( + crate::config::AffinityConfig::default(), + )), + PolicyKind::CacheAware => Arc::new(CacheAwarePolicy::new( + crate::config::AffinityConfig::default(), )), PolicyKind::Sticky => { let s = crate::config::StickyConfig::default(); @@ -184,7 +243,7 @@ pub fn build_policy_kind_only(kind: PolicyKind) -> Result> { build_sticky_fallback(s.fallback_policy), )) } - PolicyKind::FusedScore => { + PolicyKind::FusedScore | PolicyKind::ScorePolicy => { return Err(anyhow!("--policy {kind} needs --fuse terms from the model")) } }) @@ -195,6 +254,7 @@ pub fn build_registry( tree: Arc, tokenizers: Arc, block_size_oracle: Arc, + engine_load: Arc, ) -> Result { let reg = PolicyRegistry::default(); let m = &cfg.model; @@ -205,6 +265,7 @@ pub fn build_registry( Arc::clone(&tree), Arc::clone(&tokenizers), Arc::clone(&block_size_oracle), + Arc::clone(&engine_load), )?, ); Ok(reg) @@ -217,6 +278,7 @@ pub fn build_registry_with_defaults(cfg: &Config) -> Result { Arc::new(HashTree::new()), Arc::new(TokenizerRegistry::default()), BlockSizeOracle::new(), + EngineLoadTable::new(), ) } @@ -327,6 +389,7 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, @@ -340,17 +403,25 @@ mod tests { #[test] fn build_policy_kind_only_covers_all_variants() { - for kind in [ - PolicyKind::RoundRobin, - PolicyKind::Random, - PolicyKind::PowerOfTwo, - PolicyKind::LoadBased, - PolicyKind::CacheAwareZmq, - PolicyKind::Sticky, + for (kind, needs_load_snapshot) in [ + (PolicyKind::RoundRobin, false), + (PolicyKind::Random, false), + (PolicyKind::PowerOfTwo, true), + (PolicyKind::LoadBased, true), + (PolicyKind::CacheAwareZmq, true), + (PolicyKind::SessionAware, true), + (PolicyKind::CacheAware, true), + (PolicyKind::Sticky, false), ] { - assert!(build_policy_kind_only(kind).is_ok(), "{kind:?}"); + let policy = build_policy_kind_only(kind).unwrap(); + assert_eq!( + policy.needs_load_snapshot(), + needs_load_snapshot, + "{kind:?}" + ); } assert!(build_policy_kind_only(PolicyKind::FusedScore).is_err()); + assert!(build_policy_kind_only(PolicyKind::ScorePolicy).is_err()); } #[test] @@ -361,6 +432,7 @@ mod tests { &BlockSizeOracle::new(), ); assert!(p.can_fuse(), "prefix_cache must be usable as a --fuse term"); + assert!(!p.needs_load_snapshot()); } #[test] @@ -380,28 +452,84 @@ mod tests { .contains("at least one --fuse term")); } + /// `score_policy` uses its own top-level factory branch. #[test] - fn fused_score_accepts_an_outer_eligibility_pipeline() { - let mut cfg = cfg_with_model("modelA", PolicyKind::FusedScore); + fn score_policy_builds_via_its_own_factory_branch() { + let mut cfg = cfg_with_model("modelA", PolicyKind::ScorePolicy); + cfg.model.fused = Some(vec![crate::config::FusedTerm { + kind: ScoreTermKind::PrefixCache, + weight: Some(2.0), + }]); + + let registry = build_registry_with_defaults(&cfg).expect("score policy builds"); + let policy = registry + .get(&ModelId("modelA".into())) + .expect("configured model has a policy"); + assert!( + policy.can_fuse(), + "the score policy exposes score semantics" + ); + assert!( + policy.uses_shared_prefill_admission(), + "top-level score_policy participates in the shared hard admission layer" + ); + } + + #[test] + fn score_policy_with_filter_builds_one_outer_pipeline() { + let mut cfg = cfg_with_model("modelA", PolicyKind::ScorePolicy); cfg.model.fused = Some(vec![crate::config::FusedTerm { kind: ScoreTermKind::LoadBased, - weight: None, + weight: Some(1.0), }]); cfg.model.eligibility = Some(EligibilityConfig { filters: vec![FilterKind::Overloaded], - max_in_flight: Some(0), + max_in_flight: Some(2), min_prefix_share: None, }); - let policy = build_registry_with_defaults(&cfg) - .expect("an outer filter must not make fused terms non-fusable") + let registry = build_registry_with_defaults(&cfg) + .expect("a score policy keeps eligibility outside its scoring terms"); + let policy = registry .get(&ModelId("modelA".into())) - .unwrap(); - let workers = vec![worker("w0"), worker("w1")]; - let model = ModelId("modelA".into()); - assert!(policy - .select(&workers, &SelectionContext::new(&model, None)) - .is_none()); + .expect("configured model has a policy"); + assert!(policy.uses_shared_prefill_admission()); + } + + #[test] + fn factory_rejects_missing_or_zero_overloaded_capacity() { + for max_in_flight in [None, Some(0)] { + let mut cfg = cfg_with_model("modelA", PolicyKind::FusedScore); + cfg.model.fused = Some(vec![crate::config::FusedTerm { + kind: ScoreTermKind::LoadBased, + weight: None, + }]); + cfg.model.eligibility = Some(EligibilityConfig { + filters: vec![FilterKind::Overloaded], + max_in_flight, + min_prefix_share: None, + }); + + let error = build_registry_with_defaults(&cfg) + .expect_err("overloaded capacity must be positive at construction"); + let message = error.to_string(); + assert!(message.contains("max_in_flight must be greater than 0")); + } + } + + #[test] + fn factory_rejects_sticky_eligibility_pipeline() { + let mut cfg = cfg_with_model("modelA", PolicyKind::Sticky); + cfg.model.eligibility = Some(EligibilityConfig { + filters: vec![FilterKind::Overloaded], + max_in_flight: Some(2), + min_prefix_share: None, + }); + + let error = build_registry_with_defaults(&cfg) + .expect_err("sticky assignments cannot be wrapped by eligibility filters"); + let message = error.to_string(); + assert!(message.contains("eligibility filters cannot be combined with sticky")); } #[test] @@ -409,7 +537,14 @@ mod tests { let cfg = cfg_with_model("qwen", PolicyKind::RoundRobin); let tree = Arc::new(HashTree::new()); let tokenizers = Arc::new(TokenizerRegistry::default()); - let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap(); + let reg = build_registry( + &cfg, + tree, + tokenizers, + BlockSizeOracle::new(), + EngineLoadTable::new(), + ) + .unwrap(); assert!(reg.get(&ModelId("qwen".into())).is_some()); assert!(reg.get(&ModelId("missing".into())).is_none()); } @@ -419,13 +554,21 @@ mod tests { let cfg = cfg_with_model("modelA", PolicyKind::CacheAwareZmq); let tree = Arc::new(HashTree::new()); let tokenizers = Arc::new(TokenizerRegistry::default()); - let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap(); + let reg = build_registry( + &cfg, + tree, + tokenizers, + BlockSizeOracle::new(), + EngineLoadTable::new(), + ) + .unwrap(); let p = reg.get(&ModelId("modelA".into())).unwrap(); let dbg = format!("{p:?}"); assert!( dbg.contains("CacheAwareZmqPolicy"), "expected CacheAwareZmqPolicy debug repr, got: {dbg}", ); + assert!(p.needs_load_snapshot()); } #[test] @@ -433,13 +576,21 @@ mod tests { let cfg = cfg_with_model("modelA", PolicyKind::LoadBased); let tree = Arc::new(HashTree::new()); let tokenizers = Arc::new(TokenizerRegistry::default()); - let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap(); + let reg = build_registry( + &cfg, + tree, + tokenizers, + BlockSizeOracle::new(), + EngineLoadTable::new(), + ) + .unwrap(); let p = reg.get(&ModelId("modelA".into())).unwrap(); let dbg = format!("{p:?}"); assert!( dbg.contains("LoadBasedPolicy"), "expected LoadBasedPolicy debug repr, got: {dbg}", ); + assert!(p.needs_load_snapshot()); } #[test] @@ -447,7 +598,14 @@ mod tests { 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 reg = build_registry( + &cfg, + tree, + tokenizers, + BlockSizeOracle::new(), + EngineLoadTable::new(), + ) + .unwrap(); let p = reg.get(&ModelId("modelA".into())).unwrap(); let dbg = format!("{p:?}"); assert!( diff --git a/experimental/sgl-router/src/policies/kv_events/discovery.rs b/experimental/sgl-router/src/policies/kv_events/discovery.rs index 217af1dbb..a179efb35 100644 --- a/experimental/sgl-router/src/policies/kv_events/discovery.rs +++ b/experimental/sgl-router/src/policies/kv_events/discovery.rs @@ -39,6 +39,14 @@ pub struct EventConfig { pub port_base: u16, /// ZMQ topic prefix the gateway should SUBSCRIBE to. pub topic: String, + /// Base port of the worker's dedicated load-snapshot socket range + /// (per-rank load port = `load_port_base + dp_rank`). `None` when the + /// worker predates load publishing — the load subscriber is then skipped + /// and selection falls back to the router-side in-flight counter. + pub load_port_base: Option, + /// Topic prefix for the dedicated load socket. The publisher advertises + /// it with `load_port_base`; both fields are required before subscribing. + pub load_topic: Option, /// Worker-reported `page_size`. Callers MUST compare against their /// own configured `block_size`; a mismatch produces silent /// miscompute since [`super::hash::compute_block_hashes`] is keyed @@ -127,6 +135,8 @@ pub async fn fetch_event_config( host, port_base: block.endpoint_port_base, topic: block.topic, + load_port_base: block.load_endpoint_port_base, + load_topic: block.load_topic, block_size: block.block_size, dp_size: block.dp_size, is_bigram, @@ -252,6 +262,12 @@ struct KvEventsBlock { endpoint_port_base: u16, #[serde(default)] topic: String, + /// Base port of the dedicated load-snapshot socket range. Absent on + /// workers that predate load publishing (`None` ⇒ no load subscriber). + #[serde(default)] + load_endpoint_port_base: Option, + #[serde(default)] + load_topic: Option, block_size: u32, dp_size: u32, } @@ -306,6 +322,8 @@ mod tests { "endpoint_host": "*", "endpoint_port_base": 5557, "topic": "kv", + "load_endpoint_port_base": 5559, + "load_topic": "load", "block_size": 64, "dp_size": 2, } @@ -318,6 +336,8 @@ mod tests { host: "127.0.0.1".to_string(), port_base: 5557, topic: "kv".to_string(), + load_port_base: Some(5559), + load_topic: Some("load".to_string()), block_size: 64, dp_size: 2, is_bigram: false, diff --git a/experimental/sgl-router/src/policies/kv_events/index.rs b/experimental/sgl-router/src/policies/kv_events/index.rs index a91e1d205..a40ffa3c4 100644 --- a/experimental/sgl-router/src/policies/kv_events/index.rs +++ b/experimental/sgl-router/src/policies/kv_events/index.rs @@ -7,10 +7,11 @@ //! always operate together in production: //! //! - [`HashTree`] — the cache-aware routing index keyed by SGLang block hash. -//! - [`KvEventSubscriberRegistry`] — one ZMQ SUB connection per `(worker_url, -//! dp_rank)`. -//! - A pump task that drains [`WorkerEvent`]s from the subscriber and applies -//! them to the tree. +//! - [`EngineLoadTable`] — engine-reported per-worker load. +//! - Two [`KvEventSubscriberRegistry`]s — one per `(worker_url, dp_rank)` on +//! the cache topic, one on the load topic. +//! - A pump task that drains [`WorkerEvent`]s and applies KV batches to the +//! tree and `Load` snapshots to the engine-load table. //! //! `add_worker` / `remove_worker` are driven from the worker manager on every //! `DiscoveryEvent::Added` / `DiscoveryEvent::Removed`. @@ -26,7 +27,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use parking_lot::Mutex; use tokio::sync::mpsc; @@ -36,9 +37,10 @@ use tracing::{debug, info, warn}; use super::block_size_oracle::BlockSizeOracle; use super::discovery::{fetch_event_config, EventConfig}; -use super::subscriber::{KvEventSubscriberRegistry, WorkerEvent}; +use super::subscriber::{KvEventSubscriberRegistry, SubKind, WorkerEvent}; use super::tree::{HashTree, KvWorkerId}; use super::wire::KvCacheEvent; +use crate::policies::engine_load::EngineLoadTable; /// Channel buffer between the subscriber registry and the pump task. /// @@ -58,6 +60,16 @@ struct WorkerEntry { dp_ranks: Vec, } +/// Ranks whose socket port is representable for a publisher range. This is +/// shared by lifecycle bookkeeping and the subscriber registry contract so an +/// expected load rank always has a corresponding SUB socket. +fn subscribable_ranks(port_base: u16, dp_size: u32) -> Vec { + let port_base = u32::from(port_base); + (0..dp_size) + .filter(|rank| port_base.saturating_add(*rank) <= u32::from(u16::MAX)) + .collect() +} + /// Bundle of `HashTree` + `KvEventSubscriberRegistry` + pump task. /// /// Construct one instance per router process and hand it to the worker @@ -67,6 +79,14 @@ pub struct KvEventIndex { tree: Arc, maintain_tree: bool, subscribers: Arc, + /// Second registry subscribing to the load topic (one per worker rank), + /// feeding `LoadStat` snapshots into `engine_load`. Shares the pump + /// channel with `subscribers`; keyed independently so KV and load + /// subscribers for the same worker don't collide. + load_subscribers: Arc, + /// Engine-reported per-worker load, written by the pump from + /// `WorkerEvent::Load` and read by the cache-aware-zmq policy. + engine_load: Arc, pump: Mutex>>, pump_cancel: CancellationToken, workers: Mutex>, @@ -136,12 +156,15 @@ impl KvEventIndex { ) -> Arc { let tree = Arc::new(HashTree::new()); let (tx, rx) = mpsc::channel::(EVENT_CHANNEL_BUFFER); - let subscribers = Arc::new(KvEventSubscriberRegistry::new(tx)); + let subscribers = Arc::new(KvEventSubscriberRegistry::new(tx.clone())); + let load_subscribers = Arc::new(KvEventSubscriberRegistry::with_kind(tx, SubKind::Load)); + let engine_load = EngineLoadTable::new(); let cursors: Arc>> = Arc::new(Mutex::new(HashMap::new())); let live_workers: Arc>> = Arc::new(Mutex::new(HashSet::new())); let pump_cancel = CancellationToken::new(); let pump = tokio::spawn(pump_loop( tree.clone(), + engine_load.clone(), cursors.clone(), live_workers.clone(), pump_cancel.clone(), @@ -151,6 +174,8 @@ impl KvEventIndex { tree, maintain_tree, subscribers, + load_subscribers, + engine_load, pump: Mutex::new(Some(pump)), pump_cancel, workers: Mutex::new(HashMap::new()), @@ -176,15 +201,26 @@ impl KvEventIndex { self.tree.clone() } + /// Shared accessor for the engine-load table. The `CacheAwareZmqPolicy` + /// (via [`crate::policies::factory`]) holds the same `Arc` and only reads + /// it at selection time. Load *values* are written solely by the pump + /// (from `LoadStat` events); `add_worker` / `remove_worker` here manage + /// the expected set and per-worker eviction. + pub fn engine_load(&self) -> Arc { + Arc::clone(&self.engine_load) + } + /// Register a worker. If `preresolved` is `Some`, the caller has /// already fetched `/server_info` (worker manager path) and we skip /// the internal HTTP round-trip; otherwise (standalone callers, /// e.g. integration tests) we fall back to `fetch_event_config`. /// - /// Opens one ZMQ SUB per advertised DP rank. If the worker is not - /// publishing KV events (older SGLang, opt-out config), this is a - /// logged no-op — the worker still routes via the non-cache-aware - /// policies. + /// Opens one ZMQ SUB per advertised DP rank for each usable stream. In + /// metadata-only mode KV subscriptions remain disabled, but the separate + /// #34608 load stream is still attached when its full descriptor exists. + /// If the worker is not publishing KV events (older SGLang, opt-out + /// config), this is a logged no-op — the worker still routes via the + /// non-cache-aware policies. pub async fn add_worker(&self, worker_url: &str, preresolved: Option) { let cfg: EventConfig = match preresolved { Some(c) => c, @@ -228,39 +264,58 @@ impl KvEventIndex { // hash KV blocks over token bigrams, so the policy must use the bigram // hasher for its query hashes to match the worker's stored hashes. self.block_size_oracle.set_bigram(cfg.is_bigram); - if !self.maintain_tree { - info!( + let kv_dp_ranks = if self.maintain_tree { + subscribable_ranks(cfg.port_base, cfg.dp_size) + } else { + Vec::new() + }; + let load_descriptor_complete = cfg.load_port_base.is_some() && cfg.load_topic.is_some(); + if cfg.load_port_base.is_some() != cfg.load_topic.is_some() { + warn!( worker_url = %worker_url, - block_size = cfg.block_size, - is_bigram = cfg.is_bigram, - "kv-events: external Indexer configured; discovered hash metadata without subscribing" + load_port_base = ?cfg.load_port_base, + load_topic = ?cfg.load_topic, + "kv-events: incomplete load descriptor; refusing load subscription" ); - return; } - info!( - worker_url = %worker_url, - dp_size = cfg.dp_size, - port_base = cfg.port_base, - block_size = cfg.block_size, - is_bigram = cfg.is_bigram, - "kv-events: subscribing", - ); - // Compute the DP ranks that will actually be subscribed (skip - // ranks whose port overflows u16; the subscriber will warn on - // each skipped rank). - let port_base_u32 = u32::from(cfg.port_base); - let dp_ranks: Vec = (0..cfg.dp_size) - .filter(|rank| (port_base_u32 + rank) <= u32::from(u16::MAX)) - .collect(); + let load_dp_ranks = cfg + .load_port_base + .filter(|_| load_descriptor_complete) + .map(|port_base| subscribable_ranks(port_base, cfg.dp_size)) + .unwrap_or_default(); + let mut dp_ranks = kv_dp_ranks.clone(); + dp_ranks.extend(load_dp_ranks.iter().copied()); + dp_ranks.sort_unstable(); + dp_ranks.dedup(); if dp_ranks.is_empty() { warn!( worker_url = %worker_url, port_base = cfg.port_base, dp_size = cfg.dp_size, - "kv-events: every advertised rank's port overflows u16; skipping worker", + "kv-events: no usable KV or load publisher ranks; skipping worker", ); return; } + if self.maintain_tree { + info!( + worker_url = %worker_url, + dp_size = cfg.dp_size, + port_base = cfg.port_base, + load_port_base = ?cfg.load_port_base, + block_size = cfg.block_size, + is_bigram = cfg.is_bigram, + "kv-events: subscribing", + ); + } else { + info!( + worker_url = %worker_url, + dp_size = cfg.dp_size, + load_port_base = ?cfg.load_port_base, + block_size = cfg.block_size, + is_bigram = cfg.is_bigram, + "kv-events: external Indexer configured; subscribing only to engine load", + ); + } // Mark every rank live BEFORE the subscriber starts so any event // it queues is accepted by the pump. { @@ -278,7 +333,17 @@ impl KvEventIndex { dp_ranks: dp_ranks.clone(), }, ); - self.subscribers.add_worker(worker_url, &cfg).await; + if self.maintain_tree && !kv_dp_ranks.is_empty() { + self.subscribers.add_worker(worker_url, &cfg).await; + } + // Mark only the ranks that have an actual SUB socket. `EngineLoadTable` + // then rejects missing or stale advertised ranks as a whole worker. + if !load_dp_ranks.is_empty() { + for rank in &load_dp_ranks { + self.engine_load.mark_expected_rank(worker_url, *rank); + } + self.load_subscribers.add_worker(worker_url, &cfg).await; + } } /// Tear down a worker's subscribers and clear it from the tree. @@ -307,12 +372,14 @@ impl KvEventIndex { live.remove(id); } } - // 2. Cancel and join the per-rank subscriber tasks. No further - // events for these ranks will be queued after this returns. + // 2. Cancel and join the per-rank subscriber tasks (KV + load). No + // further events for these ranks will be queued after this returns. self.subscribers.remove_worker(worker_url).await; - // 3. Drop each rank's tree state and cursor. Any event already in - // the mpsc buffer at this point will be filtered by the - // live-set check inside the pump. + self.load_subscribers.remove_worker(worker_url).await; + // 3. Drop each rank's tree state and cursor, and the worker's engine + // load. Any event already in the mpsc buffer at this point will be + // filtered by the live-set check inside the pump. + self.engine_load.forget_worker(worker_url); let mut cursors = self.cursors.lock(); for id in &ids { self.tree.clear_worker(id); @@ -334,6 +401,7 @@ impl KvEventIndex { /// events are discarded and the task exits promptly. pub async fn shutdown(&self) { self.subscribers.shutdown().await; + self.load_subscribers.shutdown().await; self.pump_cancel.cancel(); let handle = self.pump.lock().take(); if let Some(h) = handle { @@ -349,12 +417,14 @@ impl KvEventIndex { } } -/// Drain `WorkerEvent`s and apply each batch to the tree. Out-of-order -/// (seq ≤ last_applied) and stale (worker not in `live_workers`) batches -/// are skipped. `PublisherReset` events clear the cursor so a publisher +/// Drain `WorkerEvent`s: apply KV `Batch`es to the tree and `Load` snapshots +/// to the engine-load table. Out-of-order (seq ≤ last_applied) and stale +/// (worker not in `live_workers`) KV batches are skipped; `Load` is a gauge +/// with no seq. `PublisherReset` events clear the cursor so a publisher /// restarting from seq=1 (after sending END_SEQ) is not filtered. async fn pump_loop( tree: Arc, + engine_load: Arc, cursors: Arc>>, live_workers: Arc>>, cancel: CancellationToken, @@ -390,6 +460,11 @@ async fn pump_loop( } match ev { + WorkerEvent::Load { worker, load } => { + // Gauge: last value wins, no sequence/dedup. The live-worker + // filter above already dropped load from detached workers. + engine_load.set(&worker.url, worker.dp_rank, load, Instant::now()); + } WorkerEvent::PublisherReset { worker } => { if cursors.lock().remove(&worker).is_some() { info!( @@ -433,6 +508,7 @@ async fn pump_loop( #[cfg(test)] mod tests { use super::*; + use crate::policies::engine_load::LoadStat; use crate::policies::kv_events::wire::{BlockRemoved, BlockStored, KvEventBatch}; fn worker_id(url: &str, rank: u32) -> KvWorkerId { @@ -454,6 +530,7 @@ mod tests { /// can destructure just the bits they need. struct PumpHarness { tree: Arc, + engine_load: Arc, cursors: Arc>>, #[allow(dead_code)] live_set: Arc>>, @@ -467,6 +544,7 @@ mod tests { /// the given workers pre-marked live. fn spawn_pump(live: &[KvWorkerId]) -> PumpHarness { let tree = Arc::new(HashTree::new()); + let engine_load = EngineLoadTable::new(); let cursors = Arc::new(Mutex::new(HashMap::new())); let live_set: Arc>> = Arc::new(Mutex::new(live.iter().cloned().collect())); @@ -474,6 +552,7 @@ mod tests { let (tx, rx) = mpsc::channel(4); let pump = tokio::spawn(pump_loop( tree.clone(), + engine_load.clone(), cursors.clone(), live_set.clone(), cancel.clone(), @@ -481,6 +560,7 @@ mod tests { )); PumpHarness { tree, + engine_load, cursors, live_set, cancel, @@ -521,6 +601,34 @@ mod tests { assert!(m.workers.contains(&id), "tree must hold the worker"); } + /// A `WorkerEvent::Load` lands in the engine-load table (gauge, no + /// cursor) keyed by the worker URL, and does not touch the tree. + #[tokio::test] + async fn pump_applies_load_to_engine_load_table() { + let id = worker_id("http://w1", 0); + let h = spawn_pump(std::slice::from_ref(&id)); + let (tree, engine_load, tx, pump) = (h.tree, h.engine_load, h.tx, h.pump); + + tx.send(WorkerEvent::Load { + worker: id.clone(), + load: LoadStat { + num_running_reqs: 8, + num_waiting_reqs: 4, + num_tokens: 0, + max_total_num_tokens: 0, + }, + }) + .await + .unwrap(); + drop(tx); + pump.await.unwrap(); + + let fresh = engine_load.snapshot_fresh(Instant::now()); + assert_eq!(fresh.get("http://w1").copied(), Some(12)); // 8 + 4 + // Load events must not pollute the cache tree. + assert_eq!(tree.node_count(), 0); + } + /// Out-of-order seq is filtered: a batch with seq <= last_applied is /// dropped silently and does not mutate the tree. #[tokio::test] @@ -678,6 +786,8 @@ mod tests { host: "127.0.0.1".into(), port_base: 30100, topic: String::new(), + load_port_base: None, + load_topic: None, block_size: 128, dp_size: 1, is_bigram: false, @@ -707,6 +817,8 @@ mod tests { host: "127.0.0.1".into(), port_base: 30200, topic: String::new(), + load_port_base: None, + load_topic: None, block_size: 64, dp_size: 0, is_bigram: false, @@ -728,6 +840,8 @@ mod tests { host: "127.0.0.1".into(), port_base: 30300, topic: String::new(), + load_port_base: None, + load_topic: None, block_size: 64, dp_size: 0, is_bigram: true, @@ -741,7 +855,7 @@ mod tests { } #[tokio::test] - async fn metadata_only_mode_seeds_oracle_without_registering_subscribers() { + async fn metadata_only_mode_keeps_the_load_subscriber() { let oracle = BlockSizeOracle::new(); let index = KvEventIndex::new_metadata_only_with_http_and_oracle( reqwest::Client::new(), @@ -751,6 +865,8 @@ mod tests { host: "127.0.0.1".into(), port_base: 30400, topic: "kv-events".into(), + load_port_base: Some(30410), + load_topic: Some("load".into()), block_size: 64, dp_size: 2, is_bigram: true, @@ -760,7 +876,55 @@ mod tests { assert_eq!(oracle.get(), Some(64)); assert!(oracle.is_bigram()); - assert_eq!(index.known_worker_count(), 0); + assert_eq!(index.known_worker_count(), 1); + assert_eq!(index.engine_load().expected_count(), 1); + index.shutdown().await; + } + + /// `remove_worker` clears the worker's engine load and its expected mark, + /// so a re-added worker does not inherit stale load. The worker advertises + /// a load port (no publisher there; the subscriber just retries in the + /// background and is cancelled on remove). + #[tokio::test] + async fn remove_worker_clears_engine_load() { + let index = KvEventIndex::new(); + let url = "http://127.0.0.1:59123"; + let cfg = EventConfig { + host: "127.0.0.1".into(), + port_base: 59123, + topic: String::new(), + load_port_base: Some(59223), + load_topic: Some("load".into()), + block_size: 64, + dp_size: 1, + is_bigram: false, + }; + index.add_worker(url, Some(cfg)).await; + assert_eq!(index.engine_load().expected_count(), 1); + + let now = Instant::now(); + index.engine_load().set( + url, + 0, + LoadStat { + num_running_reqs: 3, + num_waiting_reqs: 1, + num_tokens: 0, + max_total_num_tokens: 0, + }, + now, + ); + assert!(index.engine_load().snapshot_fresh(now).contains_key(url)); + + index.remove_worker(url).await; + assert!( + !index + .engine_load() + .snapshot_fresh(Instant::now()) + .contains_key(url), + "remove_worker must clear engine load" + ); + assert_eq!(index.engine_load().expected_count(), 0); index.shutdown().await; } } diff --git a/experimental/sgl-router/src/policies/kv_events/mod.rs b/experimental/sgl-router/src/policies/kv_events/mod.rs index 1923816f8..86b9c5de7 100644 --- a/experimental/sgl-router/src/policies/kv_events/mod.rs +++ b/experimental/sgl-router/src/policies/kv_events/mod.rs @@ -1,7 +1,8 @@ //! ZMQ-based KV-cache event indexer for cache-aware routing. //! //! Decodes the msgpack wire format emitted by SGLang's `ZmqEventPublisher` -//! (see `python/sglang/srt/disaggregation/kv_events.py`) and maintains the +//! (`python/sglang/srt/utils/event_publisher.py`; KV event types in +//! `python/sglang/srt/disaggregation/kv_events.py`) and maintains the //! router-side index used for cache-aware request routing. //! //! # Submodules @@ -27,7 +28,7 @@ pub(crate) use discovery::classify_bigram; pub use discovery::{fetch_event_config, EventConfig}; pub use hash::{compute_block_hashes, compute_block_hashes_bigram, sha256_to_i64}; pub use index::KvEventIndex; -pub use subscriber::{KvEventSubscriberRegistry, WorkerEvent}; +pub use subscriber::{KvEventSubscriberRegistry, SubKind, WorkerEvent}; pub use tree::{HashTree, KvWorkerId, MatchResult}; pub use wire::{ decode_event_batch, BlockRemoved, BlockStored, DecodeError, KvCacheEvent, KvEventBatch, diff --git a/experimental/sgl-router/src/policies/kv_events/subscriber.rs b/experimental/sgl-router/src/policies/kv_events/subscriber.rs index d564f9087..376861a54 100644 --- a/experimental/sgl-router/src/policies/kv_events/subscriber.rs +++ b/experimental/sgl-router/src/policies/kv_events/subscriber.rs @@ -1,12 +1,14 @@ //! Per-worker, per-DP-rank ZMQ subscriber for SGLang's `ZmqEventPublisher`. //! -//! This module owns the I/O plumbing between SGLang workers (which publish -//! KV-cache events on a PUB socket — see -//! `python/sglang/srt/disaggregation/kv_events.py`) and the in-memory hash -//! tree consumed by [`super::index::KvEventIndex`]. Each `(worker_url, -//! dp_rank)` pair gets its own SUB socket on its own tokio task, decodes -//! msgpack batches via [`super::wire`], and forwards [`WorkerEvent`]s to -//! a shared mpsc channel. +//! This module owns the I/O plumbing between SGLang workers (which publish on +//! a PUB socket via `python/sglang/srt/utils/event_publisher.py` — +//! KV-cache events from `disaggregation/kv_events.py`, load gauges from +//! `managers/scheduler_components/load_publisher.py`) and the in-memory state +//! consumed by [`super::index::KvEventIndex`]. Each `(worker_url, dp_rank)` +//! pair gets its own SUB socket on its own tokio task, decodes msgpack frames +//! by [`SubKind`] (KV batches via [`super::wire`], load via +//! [`crate::policies::engine_load`]), and forwards [`WorkerEvent`]s to a +//! shared mpsc channel. //! //! # Wire format (3-frame multipart) //! @@ -70,6 +72,7 @@ use zeromq::{Socket, SocketRecv, SubSocket, ZmqMessage}; use super::discovery::EventConfig; use super::tree::KvWorkerId; use super::wire::{decode_event_batch, KvEventBatch}; +use crate::policies::engine_load::{decode_load_stat, LoadStat}; /// Maximum number of consecutive `recv()` errors before the subscriber /// gives up and exits its task. ZMQ's internal reconnect handles transient @@ -92,7 +95,25 @@ const CONNECT_BACKOFF_CAP: Duration = Duration::from_secs(2); /// signal is handled correctly. const END_SEQ_SENTINEL: i64 = -1; +/// Which topic a subscriber task listens on, and therefore what kind of +/// [`WorkerEvent`] it produces. KV-cache events feed the hash tree (with +/// sequence-ordered dedup); load snapshots feed the engine-load table (a +/// gauge — no sequence semantics). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SubKind { + /// Cache-delta topic (`BlockStored` / `BlockRemoved` / `AllBlocksCleared`). + Kv, + /// Load-snapshot topic (`LoadStat`). + Load, +} + /// Message forwarded from a per-worker subscriber task to the pump. +/// +/// The variants partition by subscriber [`SubKind`]: `Batch` and +/// `PublisherReset` come only from a `SubKind::Kv` subscriber and carry the +/// cache stream's sequence/replay semantics; `Load` comes only from a +/// `SubKind::Load` subscriber and is a seqless gauge. A given subscriber +/// never emits both families. #[derive(Debug)] pub enum WorkerEvent { /// A normal decoded event batch. @@ -105,6 +126,14 @@ pub enum WorkerEvent { /// Decoded batch payload. batch: KvEventBatch, }, + /// A runtime load snapshot from the load topic. Carries no sequence + /// number: load is a gauge, applied last-value-wins with no dedup. + Load { + /// Identity of the SGLang worker (DP rank) that produced this load. + worker: KvWorkerId, + /// Latest load snapshot for this `(worker, dp_rank)`. + load: LoadStat, + }, /// The publisher emitted its `END_SEQ` (-1) sentinel, signalling /// shutdown. A re-connecting publisher will restart its sequence /// counter from 1; the pump uses this to reset the cursor so those @@ -117,6 +146,7 @@ impl WorkerEvent { pub fn worker(&self) -> &KvWorkerId { match self { Self::Batch { worker, .. } => worker, + Self::Load { worker, .. } => worker, Self::PublisherReset { worker } => worker, } } @@ -142,19 +172,32 @@ struct Inner { /// Owns one ZMQ SUB connection per `(worker_url, dp_rank)`. Forwards /// decoded batches to a tokio mpsc channel supplied at construction time. +/// +/// A registry is single-kind: a [`SubKind::Kv`] registry subscribes to the +/// cache topic and emits [`WorkerEvent::Batch`]; a [`SubKind::Load`] registry +/// subscribes to the load topic and emits [`WorkerEvent::Load`]. The index +/// runs one of each, both feeding the same pump channel, so KV and load +/// subscribers for the same worker never collide in the handle map. pub struct KvEventSubscriberRegistry { inner: Arc, + kind: SubKind, } impl KvEventSubscriberRegistry { - /// Build an empty registry. `tx` is where decoded events flow out; - /// the channel buffer capacity is the caller's choice. + /// Build an empty KV-cache registry. `tx` is where decoded events flow + /// out; the channel buffer capacity is the caller's choice. pub fn new(tx: mpsc::Sender) -> Self { + Self::with_kind(tx, SubKind::Kv) + } + + /// Build an empty registry of the given kind. + pub fn with_kind(tx: mpsc::Sender, kind: SubKind) -> Self { Self { inner: Arc::new(Inner { tx, handles: Mutex::new(HashMap::new()), }), + kind, } } @@ -174,6 +217,24 @@ impl KvEventSubscriberRegistry { /// If `cfg.port_base + dp_rank` overflows `u16`, that rank is skipped /// with a `warn!` log and the remaining ranks proceed. pub async fn add_worker(&self, worker_url: &str, cfg: &EventConfig) { + // (port_base, topic) depend on this registry's kind. KV uses the cache + // socket + configured topic; Load uses its own advertised socket + + // topic. Refuse an incomplete load descriptor rather than subscribe-all: + // the load wire is a distinct contract and a future mixed-use socket + // must not feed unrelated payloads into the load decoder. + let (port_base, topic) = match self.kind { + SubKind::Kv => (cfg.port_base, cfg.topic.clone()), + SubKind::Load => match (&cfg.load_port_base, &cfg.load_topic) { + (Some(port), Some(topic)) => (*port, topic.clone()), + _ => { + debug!( + worker_url = %worker_url, + "kv-events: worker lacks a complete load descriptor; skipping load subscribers" + ); + return; + } + }, + }; let mut handles = self.inner.handles.lock().await; for dp_rank in 0..cfg.dp_size { let id = KvWorkerId { @@ -188,13 +249,14 @@ impl KvEventSubscriberRegistry { ); continue; } - let port = match u16::try_from(cfg.port_base as u32 + dp_rank) { + let port = match u16::try_from(port_base as u32 + dp_rank) { Ok(p) => p, Err(_) => { warn!( worker_url = %worker_url, dp_rank, - port_base = cfg.port_base, + port_base, + kind = ?self.kind, "ZMQ event port overflows u16; skipping this rank" ); continue; @@ -205,7 +267,8 @@ impl KvEventSubscriberRegistry { let join = spawn_subscriber_task( id.clone(), endpoint, - cfg.topic.clone(), + topic.clone(), + self.kind, self.inner.tx.clone(), cancel.clone(), ); @@ -286,11 +349,12 @@ fn spawn_subscriber_task( id: KvWorkerId, endpoint: String, topic: String, + kind: SubKind, tx: mpsc::Sender, cancel: CancellationToken, ) -> JoinHandle<()> { tokio::spawn(async move { - run_subscriber(id, endpoint, topic, tx, cancel).await; + run_subscriber(id, endpoint, topic, kind, tx, cancel).await; }) } @@ -305,6 +369,7 @@ async fn run_subscriber( id: KvWorkerId, endpoint: String, topic: String, + kind: SubKind, tx: mpsc::Sender, cancel: CancellationToken, ) { @@ -313,6 +378,7 @@ async fn run_subscriber( dp_rank = id.dp_rank, endpoint = %endpoint, topic = %topic, + kind = ?kind, "starting kv-event subscriber" ); @@ -337,7 +403,7 @@ async fn run_subscriber( match res { Ok(msg) => { errors_in_a_row = 0; - if let Some(event) = decode_message(&id, msg) { + if let Some(event) = decode_message(&id, msg, kind) { if tx.send(event).await.is_err() { // The pump (or the entire index) is gone. // This is unexpected mid-stream; warn so @@ -462,8 +528,9 @@ async fn connect_with_backoff( /// Validate, parse, and decode a single 3-frame multipart ZMQ message. /// Returns `None` (with logging) for any non-event input (bad frame -/// count, sentinel sequence, or msgpack decode error). -fn decode_message(id: &KvWorkerId, msg: ZmqMessage) -> Option { +/// count, sentinel sequence, or msgpack decode error). `kind` selects +/// whether to emit a KV [`WorkerEvent::Batch`] or a [`WorkerEvent::Load`]. +fn decode_message(id: &KvWorkerId, msg: ZmqMessage, kind: SubKind) -> Option { if msg.len() != 3 { warn!( worker_url = %id.url, @@ -498,41 +565,77 @@ fn decode_message(id: &KvWorkerId, msg: ZmqMessage) -> Option { let seq = i64::from_be_bytes(seq_bytes); if seq == END_SEQ_SENTINEL { - info!( - worker_url = %id.url, - dp_rank = id.dp_rank, - "publisher signalled shutdown (END_SEQ); forwarding cursor reset" - ); - return Some(WorkerEvent::PublisherReset { worker: id.clone() }); + match kind { + SubKind::Kv => { + info!( + worker_url = %id.url, + dp_rank = id.dp_rank, + "publisher signalled shutdown (END_SEQ); forwarding cursor reset" + ); + return Some(WorkerEvent::PublisherReset { worker: id.clone() }); + } + // Load has no cursor / replay state to reset — just drop. + SubKind::Load => return None, + } } - let batch = match decode_event_batch(payload.as_ref()) { - Ok(b) => b, - Err(e) => { - warn!( + // Decode by kind: the cache topic carries `KvEventBatch`es, the load + // topic carries bare `LoadStat` snapshots — two independent wire formats + // on two independent sockets. + match kind { + SubKind::Kv => { + let batch = match decode_event_batch(payload.as_ref()) { + Ok(b) => b, + Err(e) => { + warn!( + worker_url = %id.url, + dp_rank = id.dp_rank, + seq, + error = %e, + "failed to decode KV event batch payload; dropping" + ); + return None; + } + }; + trace!( worker_url = %id.url, dp_rank = id.dp_rank, seq, - error = %e, - "failed to decode KV event batch payload; dropping" + n_events = batch.events.len(), + "decoded KV event batch" ); - return None; + Some(WorkerEvent::Batch { + worker: id.clone(), + seq, + batch, + }) } - }; - - trace!( - worker_url = %id.url, - dp_rank = id.dp_rank, - seq, - n_events = batch.events.len(), - "decoded KV event batch" - ); - - Some(WorkerEvent::Batch { - worker: id.clone(), - seq, - batch, - }) + SubKind::Load => { + let load = match decode_load_stat(payload.as_ref()) { + Ok(l) => l, + Err(e) => { + warn!( + worker_url = %id.url, + dp_rank = id.dp_rank, + seq, + error = %e, + "failed to decode load snapshot payload; dropping" + ); + return None; + } + }; + trace!( + worker_url = %id.url, + dp_rank = id.dp_rank, + seq, + "decoded load snapshot" + ); + Some(WorkerEvent::Load { + worker: id.clone(), + load, + }) + } + } } /// Pull the host out of a routing URL like `http://10.0.0.1:30000` or @@ -591,6 +694,8 @@ mod tests { host: extract_host(worker_url).unwrap_or_else(|| "127.0.0.1".to_string()), port_base, topic: String::new(), + load_port_base: None, + load_topic: None, block_size: 64, dp_size, is_bigram: false, @@ -618,6 +723,29 @@ mod tests { buf } + /// Encode a LoadStat batch `[ts, [["LoadStat", running, waiting, + /// num_tokens, max_total]], dp_rank?]` in msgspec's array layout. + /// Encode a bare LoadStat msgpack array `["LoadStat", running, waiting, + /// num_tokens, max_total, attn_dp_rank]` — the payload on the load + /// socket (no EventBatch envelope). + pub fn encode_load_stat( + running: u64, + waiting: u64, + num_tokens: u64, + max_total: u64, + attn_dp_rank: u32, + ) -> Vec { + let mut buf = Vec::new(); + mp::write_array_len(&mut buf, 6).unwrap(); + mp::write_str(&mut buf, "LoadStat").unwrap(); + mp::write_uint(&mut buf, running).unwrap(); + mp::write_uint(&mut buf, waiting).unwrap(); + mp::write_uint(&mut buf, num_tokens).unwrap(); + mp::write_uint(&mut buf, max_total).unwrap(); + mp::write_uint(&mut buf, attn_dp_rank as u64).unwrap(); + buf + } + /// Build a 3-frame multipart with topic="", the given seq (BE i64), /// and the given payload bytes. pub fn build_multipart(seq: i64, payload: Vec) -> ZmqMessage { @@ -644,6 +772,9 @@ mod tests { pub fn expect_batch(ev: WorkerEvent) -> (KvWorkerId, i64, KvEventBatch) { match ev { WorkerEvent::Batch { worker, seq, batch } => (worker, seq, batch), + WorkerEvent::Load { worker, .. } => { + panic!("expected Batch, got Load for {worker:?}") + } WorkerEvent::PublisherReset { worker } => { panic!("expected Batch, got PublisherReset for {worker:?}") } @@ -761,6 +892,56 @@ mod tests { registry.shutdown().await; } + /// The #34608 load stream has its own advertised topic. It must use that + /// filter too: accepting every frame on the socket would make a future + /// colocated publisher influence routing through a coincidentally + /// decodable payload. + #[tokio::test] + async fn load_subscriber_filters_by_advertised_topic() { + let (mut pub_sock, port) = helpers::make_pub_bound().await; + let (tx, mut rx) = mpsc::channel::(8); + let registry = KvEventSubscriberRegistry::with_kind(tx, SubKind::Load); + + let worker_url = "http://127.0.0.1:30101"; + let mut cfg = helpers::cfg_for(worker_url, port, 1); + cfg.load_port_base = Some(port); + cfg.load_topic = Some("load".into()); + registry.add_worker(worker_url, &cfg).await; + helpers::settle().await; + + let payload = helpers::encode_load_stat(5, 2, 100, 1000, 0); + pub_sock + .send(helpers::build_multipart_with_topic(b"load", 3, payload)) + .await + .unwrap(); + let other_payload = helpers::encode_load_stat(99, 0, 0, 0, 0); + pub_sock + .send(helpers::build_multipart_with_topic( + b"other", + 4, + other_payload, + )) + .await + .unwrap(); + + let event = timeout(Duration::from_millis(500), rx.recv()) + .await + .expect("timed out waiting for load event") + .expect("channel closed"); + match event { + WorkerEvent::Load { load, .. } => assert_eq!(load.num_running_reqs, 5), + other => panic!("expected Load, got {other:?}"), + } + assert!( + timeout(Duration::from_millis(200), rx.recv()) + .await + .is_err(), + "unmatched load topic must not reach the subscriber" + ); + + registry.shutdown().await; + } + /// DP rank fan-out: 3 PUB sockets, 3 distinct events, all delivered. #[tokio::test] async fn dp_rank_fan_out() { @@ -1252,34 +1433,63 @@ mod tests { // Wrong frame count. let one_frame = ZmqMessage::from(Bytes::from_static(b"only")); - assert!(decode_message(&id, one_frame).is_none()); + assert!(decode_message(&id, one_frame, SubKind::Kv).is_none()); // Sentinel seq = -1 now surfaces as PublisherReset (not None) so // the downstream pump can clear its cursor before a reconnecting // publisher restarts from seq=1. let sentinel = helpers::build_multipart(-1, b"ignored".to_vec()); - let reset = decode_message(&id, sentinel).expect("END_SEQ forwards"); + let reset = decode_message(&id, sentinel, SubKind::Kv).expect("END_SEQ forwards"); assert!(matches!(reset, WorkerEvent::PublisherReset { .. })); // Bad seq frame length. let mut bad_seq = ZmqMessage::from(Bytes::new()); bad_seq.push_back(Bytes::from_static(b"abc")); // 3 bytes, not 8 bad_seq.push_back(Bytes::from_static(b"")); - assert!(decode_message(&id, bad_seq).is_none()); + assert!(decode_message(&id, bad_seq, SubKind::Kv).is_none()); // Bad payload. let bad_payload = helpers::build_multipart(1, vec![0xff, 0xfe]); - assert!(decode_message(&id, bad_payload).is_none()); + assert!(decode_message(&id, bad_payload, SubKind::Kv).is_none()); // Happy path. let payload = helpers::encode_all_blocks_cleared_batch(0.0, None); let good = helpers::build_multipart(7, payload); - let event = decode_message(&id, good).expect("should decode"); + let event = decode_message(&id, good, SubKind::Kv).expect("should decode"); let (worker, seq, _batch) = helpers::expect_batch(event); assert_eq!(seq, 7); assert_eq!(worker, id); } + /// A `SubKind::Load` subscriber decodes a bare LoadStat frame into + /// `WorkerEvent::Load`, and drops the END_SEQ sentinel (no cursor state). + #[test] + fn decode_message_load_kind() { + let id = KvWorkerId { + url: "http://x".to_string(), + dp_rank: 1, + }; + + // END_SEQ is dropped for the load topic. + let sentinel = helpers::build_multipart(-1, b"ignored".to_vec()); + assert!(decode_message(&id, sentinel, SubKind::Load).is_none()); + + // A bare LoadStat frame becomes WorkerEvent::Load. + let payload = helpers::encode_load_stat(5, 2, 100, 1000, 1); + let msg = helpers::build_multipart(3, payload); + let event = decode_message(&id, msg, SubKind::Load).expect("should decode load"); + match event { + WorkerEvent::Load { worker, load } => { + assert_eq!(worker, id); + assert_eq!(load.num_running_reqs, 5); + assert_eq!(load.num_waiting_reqs, 2); + assert_eq!(load.num_tokens, 100); + assert_eq!(load.max_total_num_tokens, 1000); + } + other => panic!("expected Load, got {other:?}"), + } + } + /// Restart-resume contract: after a worker is removed and then re-added /// to the same endpoint, the new subscriber must connect and forward /// fresh events. Confirms that `remove_worker` releases the SUB socket diff --git a/experimental/sgl-router/src/policies/load_based.rs b/experimental/sgl-router/src/policies/load_based.rs index ccf3aeee6..998bc0874 100644 --- a/experimental/sgl-router/src/policies/load_based.rs +++ b/experimental/sgl-router/src/policies/load_based.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 +use crate::policies::admission::FreshLoadLookup; use crate::policies::scoring::ScoringPolicy; use crate::policies::SelectionContext; use crate::workers::Worker; @@ -17,6 +18,10 @@ impl LoadBasedPolicy { } impl ScoringPolicy for LoadBasedPolicy { + fn needs_load_snapshot(&self) -> bool { + true + } + /// `1.0` for the least loaded down to `0.0` for the most, min-max scaled to /// the CURRENT fleet -- relative, not absolute, so it cannot saturate: /// `1 - load/256` reads a busy fleet as all-`0.0`, tied inside @@ -24,8 +29,9 @@ impl ScoringPolicy for LoadBasedPolicy { /// /// Purely a preference: "everybody is busy" is not a reason to refuse to /// route, so this term never constrains. Capacity is `--filter`'s job. - fn scores(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Vec { - let loads: Vec = workers.iter().map(|w| w.active_load()).collect(); + fn scores(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec { + let lookup = FreshLoadLookup::new(ctx.load_snapshot(), workers.iter()); + let loads: Vec = workers.iter().map(|w| lookup.score_load(w)).collect(); let lo = loads.iter().min().copied().unwrap_or(0); let span = (loads.iter().max().copied().unwrap_or(0) - lo) as f32; // `max(1.0)` is exact: a zero span means every `l - lo` is zero too. @@ -38,8 +44,11 @@ impl ScoringPolicy for LoadBasedPolicy { mod tests { use super::*; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad}; use crate::policies::scoring::argmax::TIE_EPSILON; use crate::policies::Policy; + use std::collections::HashMap; + use std::time::Instant; fn worker(id: &str) -> Arc { Arc::new(Worker::new(WorkerSpec { @@ -85,4 +94,161 @@ mod tests { assert_eq!(got, *loads.iter().min().expect("non-empty"), "{spec}"); } } + + #[test] + fn request_snapshot_overrides_later_router_active_load() { + let model = ModelId("tiny".into()); + let w0 = worker("w0"); + let w1 = worker("w1"); + // After the request snapshot, local counters say w0 is lighter. + // The policy must still preserve the frozen Engine Load ordering. + let _after_snapshot: Vec<_> = (0..10).map(|_| w1.load_guard()).collect(); + let snapshot = EngineLoadSnapshot::from_workers( + 23, + HashMap::from([ + ( + w0.url.clone(), + EngineWorkerLoad { + num_running_reqs: 50, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + ), + ( + w1.url.clone(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + ), + ]), + ); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + + assert_eq!( + LoadBasedPolicy::new().select(&workers, &ctx).unwrap().id, + w1.id, + "load-based scoring must use the request snapshot before local active-load" + ); + } + + #[test] + fn recent_dispatches_after_snapshot_change_load_based_choice() { + let model = ModelId("tiny".into()); + let w0 = worker("w0"); + let w1 = worker("w1"); + let captured_at = Instant::now(); + let snapshot = EngineLoadSnapshot::from_workers( + 37, + HashMap::from([ + ( + w0.url.clone(), + EngineWorkerLoad { + num_running_reqs: 0, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at, + }, + ), + ( + w1.url.clone(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at, + }, + ), + ]), + ); + let _after_snapshot = [w0.timestamped_load_guard(), w0.timestamped_load_guard()]; + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + + assert_eq!( + LoadBasedPolicy::new().select(&workers, &ctx).unwrap().id, + w1.id, + "dispatches newer than Engine Load must correct its queue depth" + ); + } + + #[test] + fn dispatches_before_snapshot_are_not_double_counted() { + let model = ModelId("tiny".into()); + let w0 = worker("w0"); + let w1 = worker("w1"); + let _before_snapshot = [w0.timestamped_load_guard(), w0.timestamped_load_guard()]; + std::thread::sleep(std::time::Duration::from_millis(5)); + let captured_at = Instant::now(); + let snapshot = EngineLoadSnapshot::from_workers( + 41, + HashMap::from([ + ( + w0.url.clone(), + EngineWorkerLoad { + num_running_reqs: 0, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at, + }, + ), + ( + w1.url.clone(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at, + }, + ), + ]), + ); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + + assert_eq!( + LoadBasedPolicy::new().select(&workers, &ctx).unwrap().id, + w0.id, + "slots already covered by the snapshot must not be added again" + ); + } + + #[test] + fn incomplete_snapshot_uses_frozen_local_active_fallback() { + let model = ModelId("tiny".into()); + let w0 = worker("w0"); + let w1 = worker("w1"); + let _local_load = [w0.load_guard(), w0.load_guard()]; + let snapshot = EngineLoadSnapshot::from_workers( + 43, + HashMap::from([( + w0.url.clone(), + EngineWorkerLoad { + num_running_reqs: 0, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + )]), + ); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + + assert_eq!( + LoadBasedPolicy::new().select(&workers, &ctx).unwrap().id, + w1.id, + "a partial Engine Load set must not mix engine and local gauges" + ); + } } diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index 6bdda22f7..5e93f8656 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -2,7 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 pub mod active_load; +pub mod admission; +pub mod cache_aware; pub mod cache_aware_zmq; +pub mod engine_load; pub mod factory; pub mod kv_events; pub mod load_based; @@ -11,9 +14,11 @@ pub mod random; pub mod registry; pub mod round_robin; pub mod scoring; +pub mod session_aware; pub mod sticky; use crate::discovery::ModelId; +use crate::policies::engine_load::EngineLoadSnapshot; use crate::policies::scoring::{EligibilityFilter, ScoringPolicy}; use crate::server::metrics::MetricsRegistry; use crate::tokenizer::{adapter, TokenizerRegistry}; @@ -29,8 +34,8 @@ pub struct RequestTokens { pub engine_equivalent: bool, } -/// External indexer answer prepared by the async ingress path for a -/// cache-aware policy. +/// External indexer answer prepared by the async ingress path for the +/// synchronous cache-aware policy. pub struct ExternalPrefixSignal { pub outcome: sgl_kv_indexer::PrefixOutcome, pub query_blocks: usize, @@ -138,12 +143,19 @@ pub(crate) fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option { model: &'a ModelId, request_body: Option<&'a [u8]>, routing_key: Option<&'a str>, + session_id: Option<&'a str>, + candidate_range_id: &'a str, + input_tokens: Option, request_tokens: Option<&'a [u32]>, external_prefix: Option<&'a ExternalPrefixSignal>, + load_snapshot: Option<&'a EngineLoadSnapshot>, + affinity_lookup_enabled: bool, + affinity_assignment_enabled: bool, } impl<'a> SelectionContext<'a> { @@ -152,8 +164,14 @@ impl<'a> SelectionContext<'a> { model, request_body, routing_key: None, + session_id: None, + candidate_range_id: "global", + input_tokens: None, request_tokens: None, external_prefix: None, + load_snapshot: None, + affinity_lookup_enabled: true, + affinity_assignment_enabled: true, } } @@ -166,8 +184,14 @@ impl<'a> SelectionContext<'a> { model, request_body, routing_key, + session_id: None, + candidate_range_id: "global", + input_tokens: None, request_tokens: None, external_prefix: None, + load_snapshot: None, + affinity_lookup_enabled: true, + affinity_assignment_enabled: true, } } @@ -177,6 +201,24 @@ impl<'a> SelectionContext<'a> { self } + /// Attaches the Session-Aware session ID. + pub fn with_session_id(mut self, session_id: Option<&'a str>) -> Self { + self.session_id = session_id; + self + } + + /// Attaches this policy evaluation's candidate range ID. + pub fn with_candidate_range_id(mut self, candidate_range_id: &'a str) -> Self { + self.candidate_range_id = candidate_range_id; + self + } + + /// Attaches the request input-token count. + pub fn with_input_tokens(mut self, input_tokens: u64) -> Self { + self.input_tokens = Some(input_tokens); + self + } + pub fn with_external_prefix( mut self, external_prefix: Option<&'a ExternalPrefixSignal>, @@ -185,6 +227,25 @@ impl<'a> SelectionContext<'a> { self } + /// Attaches the Engine Load snapshot captured at request start. + pub fn with_load_snapshot(mut self, load_snapshot: &'a EngineLoadSnapshot) -> Self { + self.load_snapshot = Some(load_snapshot); + self + } + + /// Disables affinity lookup and assignment. + pub fn without_affinity_lookup(mut self) -> Self { + self.affinity_lookup_enabled = false; + self.affinity_assignment_enabled = false; + self + } + + /// Keeps affinity lookup but disables assignment writes. + pub fn without_affinity_assignment(mut self) -> Self { + self.affinity_assignment_enabled = false; + self + } + pub fn model(&self) -> &ModelId { self.model } @@ -197,6 +258,17 @@ impl<'a> SelectionContext<'a> { self.routing_key } + pub fn session_id(&self) -> Option<&str> { + self.session_id + } + + pub fn candidate_range_id(&self) -> &str { + self.candidate_range_id + } + pub fn input_tokens(&self) -> Option { + self.input_tokens + } + /// Returns ingress-computed routing tokens. pub fn request_tokens(&self) -> Option<&[u32]> { self.request_tokens @@ -205,12 +277,169 @@ impl<'a> SelectionContext<'a> { pub fn external_prefix(&self) -> Option<&ExternalPrefixSignal> { self.external_prefix } + + pub fn load_snapshot(&self) -> Option<&EngineLoadSnapshot> { + self.load_snapshot + } + + pub fn affinity_lookup_enabled(&self) -> bool { + self.affinity_lookup_enabled + } + + pub fn affinity_assignment_enabled(&self) -> bool { + self.affinity_assignment_enabled + } +} + +/// A policy's primary/backup proposal. +#[derive(Clone)] +pub struct SelectionProposal { + pub primary: Arc, + pub backup: Option>, + pub kind: ProposalKind, + /// Workers still eligible for fallback after filtering. + pub eligible_workers: Option>>, +} + +/// A Cache-Aware Prefill candidate with `E = L - H`. +#[derive(Clone)] +pub struct CacheCandidate { + pub worker: Arc, + pub matched_prefix_tokens: u64, + pub uncached_tokens: u64, + /// Candidate domain. + pub candidate_range_id: String, + /// Optional pending-Prefill limit checked with `E`. + pub max_pending_prefill_tokens: Option, +} + +/// A bounded Cache-Aware candidate set. +#[derive(Clone)] +pub struct CacheCandidateProposal { + pub candidates: Vec, + pub cache_switch_margin_tokens: u64, +} + +/// A Prefill policy result: a pair or Cache-Aware candidates. +#[derive(Clone)] +pub enum PrefillProposal { + Pair(SelectionProposal), + CacheCandidates(CacheCandidateProposal), +} + +impl PrefillProposal { + /// Applies EligibilityFilter results to either proposal form. + pub fn with_eligible_workers(self, workers: Vec>) -> Self { + match self { + Self::Pair(proposal) => Self::Pair(proposal.with_eligible_workers(workers)), + Self::CacheCandidates(mut proposal) => { + proposal.candidates.retain(|candidate| { + workers + .iter() + .any(|worker| worker.id == candidate.worker.id) + }); + Self::CacheCandidates(proposal) + } + } + } +} + +impl SelectionProposal { + /// Creates a proposal without a backup. + pub fn primary(primary: Arc) -> Self { + Self { + primary, + backup: None, + kind: ProposalKind::Generic, + eligible_workers: None, + } + } + + /// Creates a primary/backup proposal. + pub fn with_backup(primary: Arc, backup: Arc) -> Self { + Self { + primary, + backup: Some(backup), + kind: ProposalKind::PowerOfTwo, + eligible_workers: None, + } + } + + pub fn with_kind(mut self, kind: ProposalKind) -> Self { + self.kind = kind; + self + } + + pub fn with_eligible_workers(mut self, workers: Vec>) -> Self { + self.eligible_workers = Some(workers); + self + } +} + +/// The source of a primary/backup proposal. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProposalKind { + Generic, + PowerOfTwo, + SessionAffinity, + CacheAffinity, + Score, } pub trait Policy: Send + Sync + std::fmt::Debug { fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option>; - /// Whether policy selection needs request tokens. + /// Produces a primary worker and an optional backup. + fn propose( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + self.select(workers, ctx).map(SelectionProposal::primary) + } + + /// Produces a prefill proposal, including cache-aware candidate sets. + fn propose_prefill( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + self.propose(workers, ctx).map(PrefillProposal::Pair) + } + + /// Commits policy-owned affinity state after choosing the final prefill worker. + fn commit_prefill_selection( + &self, + _ctx: &SelectionContext<'_>, + _proposal_kind: ProposalKind, + _selected: &Arc, + ) { + } + + /// Indicates whether this policy uses shared prefill admission and guards. + fn uses_shared_prefill_admission(&self) -> bool { + false + } + + /// Whether routing needs one request-scoped Engine Load snapshot. + fn needs_load_snapshot(&self) -> bool { + self.uses_shared_prefill_admission() + } + + /// Whether this policy resolves an affinity primary within the candidate range. + fn is_bucket_affinity_policy(&self) -> bool { + false + } + + /// Whether this policy's routing decision needs request tokens (i.e. + /// it routes by prompt prefix). Ingress tokenization itself is no longer + /// gated on this — that is a model property (`has_chat_encoder`) decided at + /// ingress via [`request_tokens_for`]. This flag is the EXTRA gate that + /// keeps the cache-aware policy's RAW-prompt routing path alive: a + /// cache-aware model with no chat encoder still wants its `/v1/completions` + /// /`text` prompt tokenized for tree matching, which `has_chat_encoder` + /// alone would not trigger. Default `false` for load-only and sticky + /// routes; only the cache-aware policy overrides it. fn needs_request_tokens(&self) -> bool { false } @@ -260,3 +489,1087 @@ impl PolicyRegistry { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{AffinityConfig, SessionAffinityMode}; + use crate::discovery::{WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::admission::{ + resolve_cache_candidates, resolve_prefill, CandidateRange, DecisionReason, FreshLoadLookup, + }; + use crate::policies::cache_aware::CacheAwarePolicy; + use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad}; + use crate::policies::power_of_two::PowerOfTwoChoicesPolicy; + use crate::policies::round_robin::RoundRobinPolicy; + use crate::policies::session_aware::SessionAwarePolicy; + use std::collections::HashMap; + use std::time::Instant; + + /// Aggregated `LoadStat` values used only by policy tests. + #[derive(Clone, Default)] + struct TestEngineLoad { + num_running_reqs: u64, + num_waiting_reqs: u64, + num_tokens: u64, + max_total_num_tokens: u64, + } + + 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("model".into())], + bootstrap_port: None, + })) + } + + #[test] + fn default_proposal_preserves_legacy_single_worker_selection() { + let model = ModelId("model".into()); + let ctx = SelectionContext::new(&model, None); + let only = worker("only"); + let policy = PowerOfTwoChoicesPolicy::new(); + + let proposal = policy + .propose(&[Arc::clone(&only)], &ctx) + .expect("one candidate must produce a proposal"); + assert_eq!(proposal.primary.id, only.id); + assert!(proposal.backup.is_none()); + } + + #[test] + fn only_step_one_policies_opt_into_shared_prefill_admission() { + assert!(PowerOfTwoChoicesPolicy::new().uses_shared_prefill_admission()); + assert!(SessionAwarePolicy::new(AffinityConfig::default()).uses_shared_prefill_admission()); + assert!(CacheAwarePolicy::new(AffinityConfig::default()).uses_shared_prefill_admission()); + assert!(!RoundRobinPolicy::new().uses_shared_prefill_admission()); + } + + #[test] + fn shared_prefill_admission_policies_need_a_load_snapshot() { + assert!(PowerOfTwoChoicesPolicy::new().needs_load_snapshot()); + assert!(SessionAwarePolicy::new(AffinityConfig::default()).needs_load_snapshot()); + assert!(CacheAwarePolicy::new(AffinityConfig::default()).needs_load_snapshot()); + assert!(!RoundRobinPolicy::new().needs_load_snapshot()); + } + + #[test] + fn power_of_two_proposal_keeps_the_other_sample_as_backup() { + let model = ModelId("model".into()); + let ctx = SelectionContext::new(&model, None); + let workers = vec![worker("first"), worker("second")]; + let policy = PowerOfTwoChoicesPolicy::new(); + + let proposal = policy + .propose(&workers, &ctx) + .expect("two candidates must produce a proposal"); + let backup = proposal.backup.expect("P2 must retain its second sample"); + + assert_ne!(proposal.primary.id, backup.id); + assert_eq!(proposal.kind, ProposalKind::PowerOfTwo); + } + + #[test] + fn prefill_proposal_adapter_keeps_existing_pair_semantics() { + let model = ModelId("model".into()); + let workers = vec![worker("first"), worker("second")]; + let policy = PowerOfTwoChoicesPolicy::new(); + let ctx = SelectionContext::new(&model, None); + + let proposal = policy + .propose_prefill(&workers, &ctx) + .expect("P2 must produce a prefill proposal"); + + let PrefillProposal::Pair(pair) = proposal else { + panic!("existing policies must use the pair adapter"); + }; + assert_eq!(pair.kind, ProposalKind::PowerOfTwo); + assert!(pair.backup.is_some()); + } + + #[test] + fn cache_candidate_proposal_carries_target_specific_work() { + let hot = worker("hot"); + let proposal = CacheCandidateProposal { + candidates: vec![CacheCandidate { + worker: Arc::clone(&hot), + matched_prefix_tokens: 75, + uncached_tokens: 25, + candidate_range_id: "global".into(), + max_pending_prefill_tokens: None, + }], + cache_switch_margin_tokens: 8, + }; + + assert_eq!(proposal.candidates[0].worker.id, hot.id); + assert_eq!(proposal.candidates[0].matched_prefix_tokens, 75); + assert_eq!(proposal.candidates[0].uncached_tokens, 25); + } + + #[test] + fn power_of_two_orders_its_sample_with_fresh_engine_load_snapshot() { + let model = ModelId("model".into()); + let busy = worker("busy"); + let idle = worker("idle"); + let workers = vec![Arc::clone(&busy), Arc::clone(&idle)]; + let load_snapshot = snapshot(&[ + ( + &busy, + TestEngineLoad { + num_waiting_reqs: 512, + max_total_num_tokens: 4_096, + ..Default::default() + }, + ), + ( + &idle, + TestEngineLoad { + num_waiting_reqs: 16, + max_total_num_tokens: 4_096, + ..Default::default() + }, + ), + ]); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&load_snapshot); + + let proposal = PowerOfTwoChoicesPolicy::new() + .propose(&workers, &ctx) + .expect("two candidates must produce a proposal"); + + assert_eq!(proposal.primary.id, idle.id); + assert_eq!( + proposal.backup.expect("P2 keeps its other sample").id, + busy.id + ); + } + + #[test] + fn session_affinity_reuses_primary_and_stable_backup_without_remapping() { + let model = ModelId("model".into()); + let workers = vec![worker("first"), worker("second"), worker("third")]; + let policy = SessionAwarePolicy::new(AffinityConfig { + stable_pair: true, + ..Default::default() + }); + let ctx = SelectionContext::new(&model, None).with_session_id(Some("session-a")); + + let first = policy + .propose(&workers, &ctx) + .expect("a new session must get an initial P2 proposal"); + policy.commit_prefill_selection(&ctx, first.kind, &first.primary); + let second = policy + .propose(&workers, &ctx) + .expect("a mapped session must produce an affinity proposal"); + let third = policy + .propose(&workers, &ctx) + .expect("the session assignment must remain stable"); + + assert_eq!(second.kind, ProposalKind::SessionAffinity); + assert_eq!(second.primary.id, first.primary.id); + assert_eq!(third.primary.id, second.primary.id); + assert_eq!( + third.backup.expect("stable pair has backup").id, + second.backup.expect("stable pair has backup").id, + ); + } + + #[test] + fn new_session_commits_the_final_capacity_admitted_worker() { + let model = ModelId("model".into()); + let workers = vec![worker("first"), worker("second")]; + let policy = SessionAwarePolicy::new(AffinityConfig::default()); + let ctx = SelectionContext::new(&model, None).with_session_id(Some("session-a")); + let proposal = policy + .propose(&workers, &ctx) + .expect("a new session produces a P2 proposal"); + let backup = proposal + .backup + .clone() + .expect("two workers retain a backup"); + let loads = snapshot(&[ + ( + &proposal.primary, + TestEngineLoad { + num_running_reqs: 1, + num_tokens: 4_090, + max_total_num_tokens: 4_096, + ..Default::default() + }, + ), + ( + &backup, + TestEngineLoad { + max_total_num_tokens: 4_096, + ..Default::default() + }, + ), + ]); + let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &loads) + .expect("the admitted backup must become Final P"); + assert_eq!(decision.selected.id, backup.id); + policy.commit_prefill_selection(&ctx, proposal.kind, &decision.selected); + + let mapped = policy + .propose(&workers, &ctx) + .expect("the next turn must reuse the actual first-turn worker"); + assert_eq!(mapped.kind, ProposalKind::SessionAffinity); + assert_eq!(mapped.primary.id, backup.id); + } + + #[test] + fn read_only_affinity_probe_does_not_create_a_session_assignment() { + let model = ModelId("model".into()); + let workers = vec![worker("first"), worker("second")]; + let policy = SessionAwarePolicy::new(AffinityConfig::default()); + let probe = SelectionContext::new(&model, None) + .with_session_id(Some("session-a")) + .without_affinity_assignment(); + + let first = policy + .propose(&workers, &probe) + .expect("read-only probe still gets a P2 candidate"); + assert_eq!(first.kind, ProposalKind::PowerOfTwo); + + let normal = SelectionContext::new(&model, None).with_session_id(Some("session-a")); + let second = policy + .propose(&workers, &normal) + .expect("first admitted route creates the session assignment"); + assert_eq!(second.kind, ProposalKind::PowerOfTwo); + policy.commit_prefill_selection(&normal, second.kind, &second.primary); + + let mapped = policy + .propose(&workers, &normal) + .expect("subsequent route resolves the admitted assignment"); + assert_eq!(mapped.kind, ProposalKind::SessionAffinity); + } + + #[test] + fn bucket_scoped_session_affinity_remembers_each_bucket_independently() { + let model = ModelId("model".into()); + let short = worker("short"); + let long = worker("long"); + let policy = SessionAwarePolicy::new(AffinityConfig { + session_affinity_mode: SessionAffinityMode::Bucket, + ..Default::default() + }); + let short_ctx = SelectionContext::new(&model, None) + .with_session_id(Some("session-a")) + .with_candidate_range_id("p-short"); + let long_ctx = SelectionContext::new(&model, None) + .with_session_id(Some("session-a")) + .with_candidate_range_id("p-long"); + + let short_proposal = policy + .propose(&[Arc::clone(&short)], &short_ctx) + .expect("short bucket creates its assignment"); + policy.commit_prefill_selection(&short_ctx, short_proposal.kind, &short_proposal.primary); + let long_proposal = policy + .propose(&[Arc::clone(&long)], &long_ctx) + .expect("long bucket creates an independent assignment"); + policy.commit_prefill_selection(&long_ctx, long_proposal.kind, &long_proposal.primary); + let returned = policy + .propose(&[short], &short_ctx) + .expect("returning to short bucket reuses its assignment"); + + assert_eq!(returned.kind, ProposalKind::SessionAffinity); + } + + #[test] + fn decode_pressure_tie_is_not_broken_by_worker_id() { + let a = worker("a"); + let z = worker("z"); + assert_eq!( + admission::compare_decode_pressure(&a, &z, None), + std::cmp::Ordering::Equal, + "P2 must preserve random sampling when observable pressure is equal" + ); + } + + #[test] + fn cache_affinity_uses_longest_routable_prefix_holder() { + let model = ModelId("model".into()); + let hot = worker("hot"); + let other = worker("other"); + let workers = vec![Arc::clone(&hot), Arc::clone(&other)]; + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![ + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 8, + worker_id: "gone".into(), + address: "http://gone:30000".into(), + }, + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 6, + worker_id: "hot".into(), + address: "http://hot:30000".into(), + }, + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 4, + worker_id: "other".into(), + address: "http://other:30000".into(), + }, + ], + best_prefix_blocks: 8, + }, + query_blocks: 8, + }; + let ctx = SelectionContext::new(&model, None) + .with_request_tokens(Some(&[1, 2, 3, 4, 5, 6, 7, 8])) + .with_input_tokens(8_000) + .with_external_prefix(Some(&signal)); + let policy = CacheAwarePolicy::new(AffinityConfig::default()); + + let proposal = policy + .propose(&workers, &ctx) + .expect("a routable indexer hit must propose a worker"); + + assert_eq!(proposal.kind, ProposalKind::CacheAffinity); + assert_eq!(proposal.primary.id, hot.id); + } + + #[test] + fn cache_candidates_keep_bounded_target_specific_uncached_work() { + let model = ModelId("model".into()); + let hot = worker("hot"); + let warm = worker("warm"); + let workers = vec![Arc::clone(&hot), Arc::clone(&warm)]; + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![ + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 8, + worker_id: "gone".into(), + address: "http://gone:30000".into(), + }, + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 6, + worker_id: "hot".into(), + address: "http://hot:30000".into(), + }, + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 4, + worker_id: "warm".into(), + address: "http://warm:30000".into(), + }, + ], + best_prefix_blocks: 8, + }, + query_blocks: 8, + }; + let ctx = SelectionContext::new(&model, None) + .with_input_tokens(8_000) + .with_external_prefix(Some(&signal)); + let config = AffinityConfig { + cache_candidate_min_workers: 2, + cache_candidate_ratio: 0.0, + cache_candidate_max_workers: 2, + ..Default::default() + }; + let policy = CacheAwarePolicy::new(config); + + let PrefillProposal::CacheCandidates(proposal) = policy + .propose_prefill(&workers, &ctx) + .expect("routable matches must produce cache candidates") + else { + panic!("cache hits must not be collapsed to a primary/backup pair"); + }; + + assert_eq!(proposal.candidates.len(), 2); + assert_eq!(proposal.candidates[0].worker.id, hot.id); + assert_eq!(proposal.candidates[0].matched_prefix_tokens, 6_000); + assert_eq!(proposal.candidates[0].uncached_tokens, 2_000); + assert_eq!(proposal.candidates[1].worker.id, warm.id); + assert_eq!(proposal.candidates[1].matched_prefix_tokens, 4_000); + assert_eq!(proposal.candidates[1].uncached_tokens, 4_000); + } + + #[test] + fn cache_candidate_bound_keeps_the_best_k_from_a_large_match_set() { + let model = ModelId("model".into()); + let workers: Vec> = (0..64) + .map(|index| worker(&format!("w{index:02}"))) + .collect(); + let matches = workers + .iter() + .enumerate() + .map(|(index, worker)| sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: (index + 1) as u32, + worker_id: worker.id.0.clone(), + address: worker.url.clone(), + }) + .collect(); + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches, + best_prefix_blocks: workers.len() as u32, + }, + query_blocks: 64, + }; + let ctx = SelectionContext::new(&model, None) + .with_input_tokens(64_000) + .with_external_prefix(Some(&signal)); + let policy = CacheAwarePolicy::new(AffinityConfig { + cache_affinity_min_matched_tokens: Some(0), + cache_candidate_min_workers: 4, + cache_candidate_ratio: 0.0, + cache_candidate_max_workers: 4, + ..Default::default() + }); + + let PrefillProposal::CacheCandidates(proposal) = policy + .propose_prefill(&workers, &ctx) + .expect("the bounded best candidates must survive") + else { + panic!("cache hits must retain candidate-set semantics"); + }; + + assert_eq!(proposal.candidates.len(), 4); + assert_eq!( + proposal + .candidates + .iter() + .map(|candidate| candidate.matched_prefix_tokens) + .collect::>(), + vec![64_000, 63_000, 62_000, 61_000] + ); + } + + #[test] + fn equal_cache_hits_bound_by_the_captured_local_load_before_worker_id() { + let model = ModelId("model".into()); + let workers: Vec> = (0..8) + .map(|index| { + let worker = worker(&format!("w{index}")); + worker + .active_requests + .store(8 - index, std::sync::atomic::Ordering::Relaxed); + worker + }) + .collect(); + let matches = workers + .iter() + .map(|worker| sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 4, + worker_id: worker.id.0.clone(), + address: worker.url.clone(), + }) + .collect(); + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches, + best_prefix_blocks: 4, + }, + query_blocks: 4, + }; + let ctx = SelectionContext::new(&model, None) + .with_input_tokens(4_000) + .with_external_prefix(Some(&signal)); + let policy = CacheAwarePolicy::new(AffinityConfig { + cache_candidate_min_workers: 2, + cache_candidate_ratio: 0.0, + cache_candidate_max_workers: 2, + ..Default::default() + }); + + let PrefillProposal::CacheCandidates(proposal) = policy + .propose_prefill(&workers, &ctx) + .expect("equal hits must retain the least-loaded replicas") + else { + panic!("cache hits must retain candidate-set semantics"); + }; + + assert_eq!( + proposal + .candidates + .iter() + .map(|candidate| candidate.worker.id.0.as_str()) + .collect::>(), + vec!["w7", "w6"] + ); + } + + #[test] + fn cache_candidate_gates_are_configurable_lower_bounds_with_and_semantics() { + let model = ModelId("model".into()); + let half = worker("half"); + let below_ratio = worker("below-ratio"); + let workers = vec![Arc::clone(&half), Arc::clone(&below_ratio)]; + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![ + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 4, + worker_id: "half".into(), + address: "http://half:30000".into(), + }, + sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 3, + worker_id: "below-ratio".into(), + address: "http://below-ratio:30000".into(), + }, + ], + best_prefix_blocks: 4, + }, + query_blocks: 8, + }; + let ctx = SelectionContext::new(&model, None) + .with_input_tokens(80) + .with_external_prefix(Some(&signal)); + let config = AffinityConfig { + cache_affinity_min_matched_tokens: Some(30), + cache_affinity_min_match_ratio: Some(0.5), + cache_candidate_min_workers: 8, + cache_candidate_max_workers: 8, + ..Default::default() + }; + let policy = CacheAwarePolicy::new(config); + + let PrefillProposal::CacheCandidates(proposal) = policy + .propose_prefill(&workers, &ctx) + .expect("one candidate satisfies both lower bounds") + else { + panic!("the admitted cache candidate must retain H/E"); + }; + + assert_eq!(proposal.candidates.len(), 1); + assert_eq!(proposal.candidates[0].worker.id, half.id); + } + + #[test] + fn default_cache_gate_rejects_a_prefix_below_the_absolute_floor() { + let model = ModelId("model".into()); + let weak = worker("weak"); + let workers = vec![Arc::clone(&weak)]; + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 3, + worker_id: "weak".into(), + address: "http://weak:30000".into(), + }], + best_prefix_blocks: 3, + }, + query_blocks: 8, + }; + let ctx = SelectionContext::new(&model, None) + .with_input_tokens(80) + .with_external_prefix(Some(&signal)); + let policy = CacheAwarePolicy::new(AffinityConfig::default()); + + let proposal = policy + .propose_prefill(&workers, &ctx) + .expect("a weak hit must degrade to no-hit P2, not fail selection"); + + assert!( + matches!(proposal, PrefillProposal::Pair(_)), + "the default gate must keep a tiny hit from forcing cache affinity" + ); + } + + #[test] + fn default_cache_gate_accepts_the_indexer_scan_cap_for_a_long_prompt() { + let model = ModelId("model".into()); + let holder = worker("holder"); + let workers = vec![Arc::clone(&holder)]; + let signal = ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![sgl_kv_indexer::PrefixMatch { + matched_prefix_blocks: 2_048, + worker_id: "holder".into(), + address: "http://holder:30000".into(), + }], + best_prefix_blocks: 2, + }, + query_blocks: 4_125, + }; + let ctx = SelectionContext::new(&model, None) + .with_input_tokens(4_125) + .with_external_prefix(Some(&signal)); + let policy = CacheAwarePolicy::new(AffinityConfig::default()); + + let PrefillProposal::CacheCandidates(proposal) = policy + .propose_prefill(&workers, &ctx) + .expect("the default absolute gate must accept a 2048-token lower bound") + else { + panic!("a server-truncated long-prefix hit must not degrade to P2"); + }; + + assert_eq!(proposal.candidates[0].worker.id, holder.id); + assert_eq!(proposal.candidates[0].matched_prefix_tokens, 2_048); + assert_eq!(proposal.candidates[0].uncached_tokens, 2_077); + } + + #[test] + fn cache_affinity_without_signal_degrades_to_a_plain_p2_proposal() { + let model = ModelId("model".into()); + let workers = vec![worker("first"), worker("second")]; + let policy = CacheAwarePolicy::new(AffinityConfig::default()); + let ctx = SelectionContext::new(&model, None); + + let proposal = policy + .propose(&workers, &ctx) + .expect("cache miss must still route through P2"); + + assert_eq!(proposal.kind, ProposalKind::PowerOfTwo); + assert!(proposal.backup.is_some()); + } + + fn snapshot(entries: &[(&Arc, TestEngineLoad)]) -> EngineLoadSnapshot { + EngineLoadSnapshot::from_workers( + 1, + entries + .iter() + .map(|(worker, aggregate)| { + ( + worker.url.clone(), + EngineWorkerLoad { + num_running_reqs: aggregate.num_running_reqs, + num_waiting_reqs: aggregate.num_waiting_reqs, + num_tokens: aggregate.num_tokens, + max_total_num_tokens: aggregate.max_total_num_tokens, + captured_at: Instant::now(), + }, + ) + }) + .collect::>(), + ) + } + + #[test] + fn mixed_freshness_uses_one_captured_local_level_for_the_candidate_set() { + let aggregate_idle = worker("aggregate-idle"); + let aggregate_busy = worker("aggregate-busy"); + let stale = worker("stale"); + aggregate_idle + .active_requests + .store(5, std::sync::atomic::Ordering::Relaxed); + aggregate_busy + .active_requests + .store(1, std::sync::atomic::Ordering::Relaxed); + let snapshot = snapshot(&[ + ( + &aggregate_idle, + TestEngineLoad { + num_waiting_reqs: 0, + ..TestEngineLoad::default() + }, + ), + ( + &aggregate_busy, + TestEngineLoad { + num_waiting_reqs: 1_000, + ..TestEngineLoad::default() + }, + ), + ]); + + let lookup = + FreshLoadLookup::new(Some(&snapshot), [&aggregate_idle, &aggregate_busy, &stale]); + assert!(lookup.get(&aggregate_idle.id).is_some()); + assert!(lookup.get(&stale.id).is_none()); + assert_eq!( + lookup.compare_prefill_pressure(&aggregate_idle, &aggregate_busy), + std::cmp::Ordering::Greater, + "one stale member makes the complete candidate set compare by the captured local level" + ); + } + + fn cache_candidate( + worker: &Arc, + matched_prefix_tokens: u64, + uncached_tokens: u64, + max_pending_prefill_tokens: Option, + ) -> CacheCandidate { + CacheCandidate { + worker: Arc::clone(worker), + matched_prefix_tokens, + uncached_tokens, + candidate_range_id: "global".into(), + max_pending_prefill_tokens, + } + } + + #[test] + fn cache_tournament_skips_capacity_exhausted_matches_and_returns_no_backup() { + let full = worker("full"); + let winner = worker("winner"); + let proposal = CacheCandidateProposal { + candidates: vec![ + cache_candidate(&full, 90, 10, None), + cache_candidate(&winner, 70, 30, None), + ], + cache_switch_margin_tokens: 16, + }; + let loads = snapshot(&[ + ( + &full, + TestEngineLoad { + num_running_reqs: 8, + num_tokens: 9_950, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &winner, + TestEngineLoad { + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ]); + + let decision = resolve_cache_candidates(&proposal, 100, &loads) + .expect("a later admitted cache match must survive"); + + assert_eq!(decision.selected.id, winner.id); + assert_eq!(decision.primary.id, winner.id); + assert!(decision.backup.is_none()); + assert_eq!(decision.reason, DecisionReason::CacheCandidate); + } + + #[test] + fn cache_tournament_compares_every_admitted_challenger_before_finalizing() { + let first = worker("first"); + let second = worker("second"); + let final_winner = worker("final-winner"); + let proposal = CacheCandidateProposal { + candidates: vec![ + cache_candidate(&first, 40, 60, None), + cache_candidate(&second, 60, 40, None), + cache_candidate(&final_winner, 80, 20, None), + ], + cache_switch_margin_tokens: 0, + }; + let loads = snapshot(&[ + ( + &first, + TestEngineLoad { + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &second, + TestEngineLoad { + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &final_winner, + TestEngineLoad { + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ]); + + let decision = resolve_cache_candidates(&proposal, 100, &loads) + .expect("all admitted candidates must participate in the tournament"); + + assert_eq!(decision.selected.id, final_winner.id); + assert_eq!(decision.primary.id, final_winner.id); + assert!(decision.backup.is_none()); + } + + #[test] + fn cache_tournament_uses_uncached_work_for_pending_but_full_input_for_kv() { + let candidate = worker("candidate"); + let proposal = CacheCandidateProposal { + candidates: vec![cache_candidate(&candidate, 80, 20, Some(30))], + cache_switch_margin_tokens: 16, + }; + let pending_allows = snapshot(&[( + &candidate, + TestEngineLoad { + num_waiting_reqs: 5, + max_total_num_tokens: 1_000, + ..TestEngineLoad::default() + }, + )]); + assert!( + resolve_cache_candidates(&proposal, 100, &pending_allows).is_some(), + "pending admission must project E=20, not L=100" + ); + + let kv_rejects = snapshot(&[( + &candidate, + TestEngineLoad { + num_tokens: 30, + num_waiting_reqs: 5, + max_total_num_tokens: 100, + ..TestEngineLoad::default() + }, + )]); + assert!( + resolve_cache_candidates(&proposal, 100, &kv_rejects).is_none(), + "KV safety must conservatively project the complete input L=100" + ); + } + + #[test] + fn cache_tournament_keeps_cache_gain_when_legacy_token_guard_is_unavailable() { + let congested = worker("congested"); + let idle = worker("idle"); + let proposal = CacheCandidateProposal { + candidates: vec![ + cache_candidate(&congested, 90, 10, None), + cache_candidate(&idle, 80, 20, None), + ], + cache_switch_margin_tokens: 32, + }; + let loads = snapshot(&[ + ( + &congested, + TestEngineLoad { + num_waiting_reqs: 1_000, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &idle, + TestEngineLoad { + num_waiting_reqs: 10, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ]); + + let decision = resolve_cache_candidates(&proposal, 100, &loads).unwrap(); + assert_eq!(decision.selected.id, congested.id); + } + + #[test] + fn cache_tournament_keeps_a_material_cache_gain_despite_pressure() { + let hot = worker("hot"); + let idle = worker("idle"); + let proposal = CacheCandidateProposal { + candidates: vec![ + cache_candidate(&hot, 90, 10, None), + cache_candidate(&idle, 20, 80, None), + ], + cache_switch_margin_tokens: 32, + }; + let loads = snapshot(&[ + ( + &hot, + TestEngineLoad { + num_waiting_reqs: 1_000, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &idle, + TestEngineLoad { + num_waiting_reqs: 10, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ]); + + let decision = resolve_cache_candidates(&proposal, 100, &loads).unwrap(); + assert_eq!( + decision.selected.id, hot.id, + "pressure may break a near tie, but must not erase a material cache-work gain" + ); + } + + #[test] + fn cache_tournament_uses_work_order_when_legacy_token_guard_is_unavailable() { + let best_work = worker("best-work"); + let near_tie = worker("near-tie"); + let beyond_margin = worker("beyond-margin"); + let proposal = CacheCandidateProposal { + // The policy supplies candidates in increasing E order. Each + // adjacent pair is a near tie, but the last candidate is more + // than one configured margin away from the global work minimum. + candidates: vec![ + cache_candidate(&best_work, 100, 0, None), + cache_candidate(&near_tie, 80, 20, None), + cache_candidate(&beyond_margin, 60, 40, None), + ], + cache_switch_margin_tokens: 32, + }; + let loads = snapshot(&[ + ( + &best_work, + TestEngineLoad { + num_waiting_reqs: 10_000, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &near_tie, + TestEngineLoad { + num_waiting_reqs: 1_000, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ( + &beyond_margin, + TestEngineLoad { + num_waiting_reqs: 0, + max_total_num_tokens: 10_000, + ..TestEngineLoad::default() + }, + ), + ]); + + let decision = resolve_cache_candidates(&proposal, 100, &loads).unwrap(); + assert_eq!( + decision.selected.id, best_work.id, + "without a unit-compatible token-pressure signal, cache work remains authoritative" + ); + } + + #[test] + fn admission_uses_admitted_backup_before_scanning_candidate_range() { + let primary = worker("primary"); + let backup = worker("backup"); + let fallback = worker("fallback"); + let workers = vec![ + Arc::clone(&primary), + Arc::clone(&backup), + Arc::clone(&fallback), + ]; + let snapshot = snapshot(&[ + ( + &primary, + TestEngineLoad { + num_running_reqs: 4, + num_tokens: 990, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ( + &backup, + TestEngineLoad { + num_tokens: 10, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ( + &fallback, + TestEngineLoad { + num_tokens: 10, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ]); + let range = CandidateRange::global(&workers); + let proposal = SelectionProposal::with_backup(Arc::clone(&primary), Arc::clone(&backup)); + + let decision = resolve_prefill(&range, &proposal, 32, &snapshot) + .expect("an admitted backup must be selected"); + + assert_eq!(decision.selected.id, backup.id); + assert_eq!(decision.reason, DecisionReason::BackupPrimaryAdmission); + } + + #[test] + fn missing_engine_snapshot_does_not_hard_reject_a_registry_healthy_primary() { + let primary = worker("primary"); + let workers = vec![Arc::clone(&primary)]; + let snapshot = EngineLoadSnapshot::default(); + + let decision = resolve_prefill( + &CandidateRange::global(&workers), + &SelectionProposal::primary(Arc::clone(&primary)), + 1_000_000, + &snapshot, + ) + .expect("disabled reporting must preserve the healthy registry candidate"); + + assert_eq!(decision.selected.id, primary.id); + assert_eq!(decision.reason, DecisionReason::Primary); + } + + #[test] + fn prefill_pair_keeps_primary_when_both_workers_fit_capacity() { + let primary = worker("primary"); + let backup = worker("backup"); + let workers = vec![Arc::clone(&primary), Arc::clone(&backup)]; + let snapshot = snapshot(&[ + ( + &primary, + TestEngineLoad { + num_waiting_reqs: 200, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ( + &backup, + TestEngineLoad { + num_waiting_reqs: 20, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ]); + let proposal = SelectionProposal::with_backup(primary, backup); + + let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 80, &snapshot) + .expect("both candidates fit capacity"); + + assert_eq!(decision.reason, DecisionReason::Primary); + } + + #[test] + fn admission_scans_range_only_after_primary_and_backup_both_fail() { + let primary = worker("primary"); + let backup = worker("backup"); + let fallback = worker("fallback"); + let workers = vec![ + Arc::clone(&primary), + Arc::clone(&backup), + Arc::clone(&fallback), + ]; + let snapshot = snapshot(&[ + ( + &primary, + TestEngineLoad { + num_running_reqs: 4, + num_tokens: 990, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ( + &backup, + TestEngineLoad { + num_tokens: 990, + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ( + &fallback, + TestEngineLoad { + max_total_num_tokens: 1_000, + ..Default::default() + }, + ), + ]); + let proposal = SelectionProposal::with_backup(primary, backup); + + let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &snapshot) + .expect("an admitted range fallback must be selected"); + + assert_eq!(decision.selected.id, fallback.id); + assert_eq!(decision.reason, DecisionReason::RangeFallback); + } +} diff --git a/experimental/sgl-router/src/policies/power_of_two.rs b/experimental/sgl-router/src/policies/power_of_two.rs index 554f70511..b840157ca 100644 --- a/experimental/sgl-router/src/policies/power_of_two.rs +++ b/experimental/sgl-router/src/policies/power_of_two.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -use crate::policies::{Policy, SelectionContext}; +use crate::policies::admission::compare_prefill_pressure; +use crate::policies::engine_load::EngineLoadSnapshot; +use crate::policies::{Policy, ProposalKind, SelectionContext, SelectionProposal}; use crate::workers::Worker; use rand::Rng; use std::sync::Arc; @@ -16,10 +18,21 @@ impl PowerOfTwoChoicesPolicy { } impl Policy for PowerOfTwoChoicesPolicy { - fn select(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Option> { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + select_with_snapshot(workers, ctx.load_snapshot()) + } + + /// Returns the primary and backup from one sample. + fn propose( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { match workers.len() { 0 => None, - 1 => Some(workers[0].clone()), + 1 => Some( + SelectionProposal::primary(workers[0].clone()).with_kind(ProposalKind::PowerOfTwo), + ), len => { let mut rng = rand::thread_rng(); let i = rng.gen_range(0..len); @@ -27,8 +40,60 @@ impl Policy for PowerOfTwoChoicesPolicy { if j >= i { j += 1; } - Some(std::cmp::min_by_key(&workers[i], &workers[j], |w| w.active_load()).clone()) + let (primary, backup) = ordered_pair(&workers[i], &workers[j], ctx); + Some(SelectionProposal::with_backup(primary, backup)) } } } + + fn uses_shared_prefill_admission(&self) -> bool { + true + } +} + +pub(crate) fn select_with_snapshot( + workers: &[Arc], + snapshot: Option<&EngineLoadSnapshot>, +) -> Option> { + match workers.len() { + 0 => None, + 1 => Some(workers[0].clone()), + len => { + let mut rng = rand::thread_rng(); + let i = rng.gen_range(0..len); + let mut j = rng.gen_range(0..len - 1); + if j >= i { + j += 1; + } + Some(select_lower_pressure(&workers[i], &workers[j], snapshot)) + } + } +} + +fn select_lower_pressure( + left: &Arc, + right: &Arc, + snapshot: Option<&EngineLoadSnapshot>, +) -> Arc { + ordered_pair_with_snapshot(left, right, snapshot).0 +} + +fn ordered_pair( + left: &Arc, + right: &Arc, + ctx: &SelectionContext<'_>, +) -> (Arc, Arc) { + ordered_pair_with_snapshot(left, right, ctx.load_snapshot()) +} + +fn ordered_pair_with_snapshot( + left: &Arc, + right: &Arc, + snapshot: Option<&EngineLoadSnapshot>, +) -> (Arc, Arc) { + if compare_prefill_pressure(left, right, snapshot).is_gt() { + (Arc::clone(right), Arc::clone(left)) + } else { + (Arc::clone(left), Arc::clone(right)) + } } diff --git a/experimental/sgl-router/src/policies/scoring/admission.rs b/experimental/sgl-router/src/policies/scoring/admission.rs index de492d941..900d7047a 100644 --- a/experimental/sgl-router/src/policies/scoring/admission.rs +++ b/experimental/sgl-router/src/policies/scoring/admission.rs @@ -69,6 +69,7 @@ mod tests { #[test] fn the_cap_is_a_strict_ceiling() { + assert!(!Overloaded::new(3).needs_load_snapshot()); let ws = vec![worker("idle"), worker("under"), worker("at")]; let _under: Vec<_> = (0..2).map(|_| ws[1].load_guard()).collect(); let _at: Vec<_> = (0..3).map(|_| ws[2].load_guard()).collect(); diff --git a/experimental/sgl-router/src/policies/scoring/mod.rs b/experimental/sgl-router/src/policies/scoring/mod.rs index 31fbc9ddb..42587c2f2 100644 --- a/experimental/sgl-router/src/policies/scoring/mod.rs +++ b/experimental/sgl-router/src/policies/scoring/mod.rs @@ -7,7 +7,7 @@ pub mod admission; pub mod argmax; pub mod prefix_cache; -use crate::policies::{Policy, SelectionContext}; +use crate::policies::{Policy, PrefillProposal, SelectionContext, SelectionProposal}; use crate::workers::Worker; use argmax::{Selector, ARGMAX}; use std::sync::Arc; @@ -52,6 +52,11 @@ pub trait ScoringPolicy: Send + Sync + std::fmt::Debug { false } + /// Whether scoring reads the request-scoped Engine Load snapshot. + fn needs_load_snapshot(&self) -> bool { + false + } + /// Optional eligibility view for policies that provide both signals. fn as_filter(&self) -> Option<&dyn EligibilityFilter> { None @@ -99,11 +104,10 @@ pub fn admit<'f>( match on_empty { OnEmpty::Hold => return None, OnEmpty::Abstain if untouched => { - tracing::warn!( + tracing::debug!( filter = ?filter, n_workers = workers.len(), - "eligibility filter rejected every worker on its own; a filter \ - with no signal should abstain (all true), not veto the fleet", + "eligibility filter has no eligible workers; falling back to the full candidate set", ); continue; } @@ -135,6 +139,10 @@ impl Policy for T { ScoringPolicy::needs_tokens(self) || self.as_filter().is_some_and(|f| f.needs_tokens()) } + fn needs_load_snapshot(&self) -> bool { + ScoringPolicy::needs_load_snapshot(self) + } + fn as_scoring(&self) -> Option<&dyn ScoringPolicy> { Some(self) } @@ -185,12 +193,107 @@ impl Pipeline { fn views(&self) -> impl Iterator { (self.filters.iter()).map(|p| p.as_filter().expect("checked by Pipeline::new")) } + + /// Apply eligibility without rewriting an existing Session assignment. + fn propose_prefill_filtered( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + let eligible = admit(self.views(), workers, ctx)?; + if self.inner.is_bucket_affinity_policy() && ctx.affinity_lookup_enabled() { + let probe_ctx = (*ctx).clone().without_affinity_assignment(); + if let Some( + proposal @ PrefillProposal::Pair(SelectionProposal { + kind: crate::policies::ProposalKind::SessionAffinity, + .. + }), + ) = self.inner.propose_prefill(workers, &probe_ctx) + { + return Some(proposal.with_eligible_workers(eligible)); + } + } + self.inner + .propose_prefill(&eligible, ctx) + .map(|proposal| proposal.with_eligible_workers(eligible)) + } } impl Policy for Pipeline { fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { - let eligible = admit(self.views(), workers, ctx)?; - self.inner.select(&eligible, ctx) + let (proposal_kind, selected) = match self.propose_prefill_filtered(workers, ctx)? { + PrefillProposal::Pair(proposal) => { + let eligible = proposal.eligible_workers.as_deref().unwrap_or(workers); + if eligible + .iter() + .any(|worker| worker.id == proposal.primary.id) + { + (proposal.kind, proposal.primary) + } else { + let selected = proposal + .backup + .filter(|backup| eligible.iter().any(|worker| worker.id == backup.id)) + .or_else(|| eligible.first().cloned())?; + (proposal.kind, selected) + } + } + PrefillProposal::CacheCandidates(proposal) => { + let selected = proposal.candidates.into_iter().next()?.worker; + (crate::policies::ProposalKind::CacheAffinity, selected) + } + }; + self.inner + .commit_prefill_selection(ctx, proposal_kind, &selected); + Some(selected) + } + + /// Preserves the inner policy's complete proposal. + fn propose( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + match self.propose_prefill_filtered(workers, ctx)? { + PrefillProposal::Pair(proposal) => Some(proposal), + PrefillProposal::CacheCandidates(proposal) => { + let candidate = proposal.candidates.into_iter().next()?; + Some( + SelectionProposal::primary(candidate.worker) + .with_kind(crate::policies::ProposalKind::CacheAffinity), + ) + } + } + } + + fn propose_prefill( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + self.propose_prefill_filtered(workers, ctx) + } + + fn uses_shared_prefill_admission(&self) -> bool { + self.inner.uses_shared_prefill_admission() + } + + fn needs_load_snapshot(&self) -> bool { + self.inner.needs_load_snapshot() || self.filters.iter().any(|p| p.needs_load_snapshot()) + } + + fn commit_prefill_selection( + &self, + ctx: &SelectionContext<'_>, + proposal_kind: crate::policies::ProposalKind, + selected: &Arc, + ) { + self.inner + .commit_prefill_selection(ctx, proposal_kind, selected); + } + + /// Preserves the inner policy's Bucket-affinity semantics. + fn is_bucket_affinity_policy(&self) -> bool { + self.inner.is_bucket_affinity_policy() } fn needs_request_tokens(&self) -> bool { @@ -202,6 +305,54 @@ impl Policy for Pipeline { } } +/// Top-level score policy that enters shared Prefill admission. +#[derive(Debug)] +pub struct ScorePolicy { + inner: Arc, +} + +impl ScorePolicy { + pub fn new(inner: Arc) -> Self { + Self { inner } + } +} + +impl Policy for ScorePolicy { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + self.inner.select(workers, ctx) + } + + fn propose( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + self.inner + .propose(workers, ctx) + .map(|proposal| proposal.with_kind(crate::policies::ProposalKind::Score)) + } + + fn uses_shared_prefill_admission(&self) -> bool { + true + } + + fn needs_request_tokens(&self) -> bool { + self.inner.needs_request_tokens() + } + + fn attach_metrics(&self, metrics: Arc) { + self.inner.attach_metrics(metrics); + } + + fn as_scoring(&self) -> Option<&dyn ScoringPolicy> { + self.inner.as_scoring() + } + + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + self.inner.as_filter() + } +} + impl ScoringPolicy for FusedScorePolicy { fn scores(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec { let mut total = vec![0.0f32; workers.len()]; @@ -216,6 +367,12 @@ impl ScoringPolicy for FusedScorePolicy { fn needs_tokens(&self) -> bool { self.terms.iter().map(view).any(|(t, _)| t.needs_tokens()) } + + fn needs_load_snapshot(&self) -> bool { + self.terms + .iter() + .any(|(policy, _)| policy.needs_load_snapshot()) + } } /// Owned boxes as the borrowed views [`admit`] consumes. Shared by the tests @@ -230,8 +387,15 @@ pub(crate) fn refs( #[cfg(test)] mod tests { use super::*; + use crate::config::AffinityConfig; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::admission::{resolve_prefill, CandidateRange}; + use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad}; + use crate::policies::power_of_two::PowerOfTwoChoicesPolicy; use crate::policies::round_robin::RoundRobinPolicy; + use crate::policies::session_aware::SessionAwarePolicy; + use std::collections::HashMap; + use std::time::Instant; fn worker(id: &str) -> Arc { Arc::new(Worker::new(WorkerSpec { @@ -247,6 +411,27 @@ mod tests { vec![worker("a"), worker("b"), worker("c")] } + fn snapshot(entries: &[(&Arc, u64, u64, u64, u64)]) -> EngineLoadSnapshot { + EngineLoadSnapshot::from_workers( + 1, + entries + .iter() + .map(|(worker, running, waiting, used, capacity)| { + ( + worker.url.clone(), + EngineWorkerLoad { + num_running_reqs: *running, + num_waiting_reqs: *waiting, + num_tokens: *used, + max_total_num_tokens: *capacity, + captured_at: Instant::now(), + }, + ) + }) + .collect::>(), + ) + } + fn urls(ws: &[Arc]) -> Vec { ws.iter().map(|w| w.url.clone()).collect() } @@ -406,6 +591,39 @@ mod tests { assert!(filtered.needs_request_tokens(), "the filter is hungry"); } + #[test] + fn composer_propagates_load_snapshot_capability() { + #[derive(Debug)] + struct LoadHungry; + impl ScoringPolicy for LoadHungry { + fn scores(&self, workers: &[Arc], _: &SelectionContext<'_>) -> Vec { + vec![0.0; workers.len()] + } + fn needs_load_snapshot(&self) -> bool { + true + } + } + + let plain = FusedScorePolicy::new(vec![term(by(1.0), None)]).unwrap(); + assert!(!Policy::needs_load_snapshot(&plain)); + let fused = + FusedScorePolicy::new(vec![term(by(1.0), None), term(LoadHungry, None)]).unwrap(); + assert!(Policy::needs_load_snapshot(&fused)); + + let pipeline = Pipeline::new( + vec![Arc::new(Keep(vec!["a"], OnEmpty::Abstain))], + Arc::new(fused), + ) + .unwrap(); + assert!(pipeline.needs_load_snapshot()); + + let score = ScorePolicy::new(Arc::new(by(1.0))); + assert!( + score.needs_load_snapshot(), + "shared admission requires a snapshot" + ); + } + #[test] fn a_rejected_worker_cannot_be_out_weighed() { let ws = fleet(); @@ -431,6 +649,137 @@ mod tests { assert_eq!(open.select(&ws, &ctx).unwrap().url, ws[2].url); } + #[test] + fn pipeline_preserves_the_inner_step_one_proposal_and_admission_opt_in() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + let pipeline = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b", "c"], OnEmpty::Abstain))], + Arc::new(PowerOfTwoChoicesPolicy::new()), + ) + .expect("valid filter and inner policy"); + + let proposal = pipeline + .propose(&ws, &ctx) + .expect("eligible P2 must retain a pair"); + + assert!( + proposal.backup.is_some(), + "Pipeline must not collapse P2 to one primary" + ); + assert!(pipeline.uses_shared_prefill_admission()); + + let session_pipeline = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b", "c"], OnEmpty::Abstain))], + Arc::new(SessionAwarePolicy::new(AffinityConfig::default())), + ) + .expect("valid filter and inner session policy"); + assert!( + session_pipeline.is_bucket_affinity_policy(), + "Pipeline must forward the inner Session affinity range capability" + ); + } + + #[test] + fn shared_admission_fallback_cannot_reintroduce_a_filtered_worker() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + let pipeline = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b"], OnEmpty::Abstain))], + Arc::new(PowerOfTwoChoicesPolicy::new()), + ) + .expect("valid filter and inner policy"); + let proposal = pipeline + .propose(&ws, &ctx) + .expect("the two eligible workers produce a P2 proposal"); + let snapshot = snapshot(&[ + (&ws[0], 0, 0, 4_090, 4_096), + (&ws[1], 0, 0, 4_090, 4_096), + (&ws[2], 0, 0, 0, 4_096), + ]); + + let decision = resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot) + .expect("capacity exhaustion must degrade inside the filtered domain"); + assert!(matches!(decision.selected.id.0.as_str(), "a" | "b")); + } + + #[test] + fn eligibility_escape_does_not_rewrite_an_existing_session_assignment() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_session_id(Some("session-a")); + let session = Arc::new(SessionAwarePolicy::new(AffinityConfig::default())); + + let initial = session + .propose(&ws[2..], &ctx) + .expect("one-worker domain establishes c"); + assert_eq!(initial.primary.id, ws[2].id); + session.commit_prefill_selection(&ctx, initial.kind, &initial.primary); + + let pipeline = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b"], OnEmpty::Abstain))], + session.clone(), + ) + .expect("valid filter and session policy"); + let PrefillProposal::Pair(proposal) = pipeline + .propose_prefill(&ws, &ctx) + .expect("filtered session proposal") + else { + panic!("Session-Aware must retain pair semantics"); + }; + assert_eq!( + proposal.kind, + crate::policies::ProposalKind::SessionAffinity + ); + assert_eq!(proposal.primary.id, ws[2].id); + + let snapshot = EngineLoadSnapshot::default(); + let decision = resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot) + .expect("an eligible escape worker exists"); + assert_ne!(decision.selected.id, ws[2].id); + assert!(matches!(decision.selected.id.0.as_str(), "a" | "b")); + + let after = session + .propose(&ws, &ctx) + .expect("the original assignment remains readable"); + assert_eq!(after.kind, crate::policies::ProposalKind::SessionAffinity); + assert_eq!(after.primary.id, ws[2].id); + } + + #[test] + fn new_session_assignment_is_created_inside_the_eligible_set() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_session_id(Some("session-new")); + let session = Arc::new(SessionAwarePolicy::new(AffinityConfig::default())); + let pipeline = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b"], OnEmpty::Abstain))], + session.clone(), + ) + .expect("valid filter and session policy"); + + let PrefillProposal::Pair(proposal) = pipeline + .propose_prefill(&ws, &ctx) + .expect("eligible workers establish the session") + else { + panic!("Session-Aware must retain pair semantics"); + }; + assert!(matches!(proposal.primary.id.0.as_str(), "a" | "b")); + pipeline.commit_prefill_selection(&ctx, proposal.kind, &proposal.primary); + + let mapped = session + .propose(&ws, &ctx) + .expect("the assignment is stored by the inner policy"); + assert_eq!(mapped.kind, crate::policies::ProposalKind::SessionAffinity); + assert_eq!(mapped.primary.id, proposal.primary.id); + } + + /// Order is priority: the LOWER-priority filter yields, and what the + /// higher-priority one narrowed to is kept. Asserted on the surviving set + /// rather than on the winner, because with three workers a wrong rule can + /// still land on the right one by luck. #[test] fn a_conflict_yields_the_later_filter_and_keeps_the_earlier_narrowing() { let ws = fleet(); diff --git a/experimental/sgl-router/src/policies/session_aware.rs b/experimental/sgl-router/src/policies/session_aware.rs new file mode 100644 index 000000000..9b8cdad09 --- /dev/null +++ b/experimental/sgl-router/src/policies/session_aware.rs @@ -0,0 +1,353 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Session-aware prefill policy. Admission does not rewrite session assignments. + +use crate::config::{AffinityConfig, SessionAffinityMode}; +use crate::discovery::WorkerId; +use crate::policies::active_load::{spawn_sweeper, Clock, JanitorHandle, SystemTimeClock}; +use crate::policies::admission::compare_prefill_pressure; +use crate::policies::power_of_two::PowerOfTwoChoicesPolicy; +use crate::policies::{Policy, ProposalKind, SelectionContext, SelectionProposal}; +use crate::workers::Worker; +use dashmap::DashMap; +use rand::Rng; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +#[derive(Debug)] +struct Assignment { + worker_id: WorkerId, + last_seen: Instant, +} + +#[derive(Debug)] +struct SessionState { + assignments: DashMap, + clock: Arc, + idle: Duration, +} + +impl SessionState { + fn sweep_expired(&self) -> usize { + let now = self.clock.now(); + let mut removed = 0; + self.assignments.retain(|_, assignment| { + let keep = now.saturating_duration_since(assignment.last_seen) <= self.idle; + if !keep { + removed += 1; + } + keep + }); + removed + } +} + +pub struct SessionAwarePolicy { + state: Arc, + config: AffinityConfig, + _janitor: Option, +} + +impl SessionAwarePolicy { + pub fn new(config: AffinityConfig) -> Self { + let state = Arc::new(SessionState { + assignments: DashMap::new(), + clock: Arc::new(SystemTimeClock), + idle: Duration::from_secs(config.session_idle_secs), + }); + let _janitor = if tokio::runtime::Handle::try_current().is_ok() { + let swept = Arc::clone(&state); + Some(spawn_sweeper( + move || swept.sweep_expired(), + Duration::from_secs(config.session_eviction_interval_secs), + "session-affinity-eviction", + )) + } else { + tracing::debug!( + "SessionAwarePolicy constructed outside a Tokio runtime; idle eviction is disabled" + ); + None + }; + Self { + state, + config, + _janitor, + } + } + + #[cfg(test)] + fn with_clock(config: AffinityConfig, clock: Arc) -> Self { + Self { + state: Arc::new(SessionState { + assignments: DashMap::new(), + clock, + idle: Duration::from_secs(config.session_idle_secs), + }), + config, + _janitor: None, + } + } + + #[cfg(test)] + fn sweep_expired(&self) -> usize { + self.state.sweep_expired() + } + + #[cfg(test)] + fn assignment_count(&self) -> usize { + self.state.assignments.len() + } + + fn assignment_key(&self, session_id: &str, ctx: &SelectionContext<'_>) -> String { + match self.config.session_affinity_mode { + SessionAffinityMode::Bucket => { + format!("{}\0{}", ctx.candidate_range_id(), session_id) + } + SessionAffinityMode::GlobalRebind | SessionAffinityMode::GlobalPreserve => { + session_id.to_string() + } + } + } + + fn initial_proposal( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + PowerOfTwoChoicesPolicy::new().propose(workers, ctx) + } + + fn affinity_proposal( + &self, + primary: Arc, + workers: &[Arc], + ctx: &SelectionContext<'_>, + session_id: &str, + ) -> SelectionProposal { + let backup = affinity_backup( + workers, + &primary, + session_id, + ctx.candidate_range_id(), + self.config.stable_pair, + ctx, + ); + let proposal = match backup { + Some(backup) => SelectionProposal::with_backup(primary, backup), + None => SelectionProposal::primary(primary), + }; + proposal.with_kind(ProposalKind::SessionAffinity) + } +} + +impl Policy for SessionAwarePolicy { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + let proposal = self.propose(workers, ctx)?; + let proposal_kind = proposal.kind; + let selected = proposal.primary; + self.commit_prefill_selection(ctx, proposal_kind, &selected); + Some(selected) + } + + fn propose( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + ) -> Option { + if !ctx.affinity_lookup_enabled() { + return PowerOfTwoChoicesPolicy::new().propose(workers, ctx); + } + let Some(session_id) = ctx.session_id().filter(|id| !id.is_empty()) else { + return self.initial_proposal(workers, ctx); + }; + + let assignment_key = self.assignment_key(session_id, ctx); + let assigned = self + .state + .assignments + .get_mut(&assignment_key) + .map(|mut assignment| { + assignment.last_seen = self.state.clock.now(); + assignment.worker_id.clone() + }); + if let Some(assigned) = assigned { + if let Some(primary) = workers.iter().find(|worker| worker.id == assigned).cloned() { + return Some(self.affinity_proposal(primary, workers, ctx, session_id)); + } + } + + // Persist new assignments only after selecting the final prefill worker. + self.initial_proposal(workers, ctx) + } + + fn commit_prefill_selection( + &self, + ctx: &SelectionContext<'_>, + proposal_kind: ProposalKind, + selected: &Arc, + ) { + if proposal_kind != ProposalKind::PowerOfTwo || !ctx.affinity_assignment_enabled() { + return; + } + let Some(session_id) = ctx.session_id().filter(|id| !id.is_empty()) else { + return; + }; + self.state.assignments.insert( + self.assignment_key(session_id, ctx), + Assignment { + worker_id: selected.id.clone(), + last_seen: self.state.clock.now(), + }, + ); + } + + fn uses_shared_prefill_admission(&self) -> bool { + true + } + + fn is_bucket_affinity_policy(&self) -> bool { + true + } +} + +impl std::fmt::Debug for SessionAwarePolicy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SessionAwarePolicy") + .field("config", &self.config) + .field("assignments", &self.state.assignments.len()) + .finish_non_exhaustive() + } +} + +pub(crate) fn affinity_backup( + workers: &[Arc], + primary: &Arc, + affinity_key: &str, + candidate_range_id: &str, + stable_pair: bool, + ctx: &SelectionContext<'_>, +) -> Option> { + if stable_pair { + return stable_backup(workers, primary, affinity_key, candidate_range_id); + } + sampled_backup_excluding(workers, primary, ctx) +} + +fn sampled_backup_excluding( + workers: &[Arc], + primary: &Arc, + ctx: &SelectionContext<'_>, +) -> Option> { + let mut rng = rand::thread_rng(); + let first = sample_index_excluding(workers, &primary.id, None, &mut rng)?; + let Some(second) = sample_index_excluding(workers, &primary.id, Some(first), &mut rng) else { + return Some(Arc::clone(&workers[first])); + }; + let left = &workers[first]; + let right = &workers[second]; + if compare_prefill_pressure(left, right, ctx.load_snapshot()).is_gt() { + Some(Arc::clone(right)) + } else { + Some(Arc::clone(left)) + } +} + +fn sample_index_excluding( + workers: &[Arc], + primary_id: &WorkerId, + other_index: Option, + rng: &mut impl Rng, +) -> Option { + if workers.is_empty() { + return None; + } + for _ in 0..32 { + let index = rng.gen_range(0..workers.len()); + if Some(index) != other_index && workers[index].id != *primary_id { + return Some(index); + } + } + workers.iter().enumerate().find_map(|(index, worker)| { + (Some(index) != other_index && worker.id != *primary_id).then_some(index) + }) +} + +fn stable_backup( + workers: &[Arc], + primary: &Arc, + session_id: &str, + candidate_range_id: &str, +) -> Option> { + let mut others: Vec> = workers + .iter() + .filter(|worker| worker.id != primary.id) + .cloned() + .collect(); + others.sort_by(|left, right| left.id.0.cmp(&right.id.0)); + if others.is_empty() { + return None; + } + let mut hasher = DefaultHasher::new(); + session_id.hash(&mut hasher); + candidate_range_id.hash(&mut hasher); + Some(others[(hasher.finish() as usize) % others.len()].clone()) +} + +#[cfg(test)] +mod lifecycle_tests { + use super::*; + use crate::discovery::{ModelId, WorkerMode, WorkerSpec}; + use crate::policies::active_load::MockClock; + use std::sync::atomic::Ordering; + use std::time::{Duration, Instant}; + + 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("model".into())], + bootstrap_port: None, + })) + } + + #[test] + fn idle_session_assignments_are_evicted() { + let clock = Arc::new(MockClock::new(Instant::now())); + let policy = SessionAwarePolicy::with_clock( + AffinityConfig { + session_idle_secs: 10, + ..Default::default() + }, + clock.clone(), + ); + let model = ModelId("model".into()); + let ctx = SelectionContext::new(&model, None).with_session_id(Some("session-a")); + let proposal = policy.propose(&[worker("w")], &ctx).unwrap(); + policy.commit_prefill_selection(&ctx, proposal.kind, &proposal.primary); + assert_eq!(policy.assignment_count(), 1); + + clock.advance(Duration::from_secs(11)); + assert_eq!(policy.sweep_expired(), 1); + assert_eq!(policy.assignment_count(), 0); + } + + #[test] + fn sampled_backup_excludes_primary_without_materializing_the_remaining_fleet() { + let primary = worker("primary"); + let busy = worker("busy"); + let idle = worker("idle"); + busy.active_requests.store(8, Ordering::Relaxed); + idle.active_requests.store(1, Ordering::Relaxed); + let workers = vec![Arc::clone(&primary), busy, Arc::clone(&idle)]; + let model = ModelId("model".into()); + let ctx = SelectionContext::new(&model, None); + + let backup = sampled_backup_excluding(&workers, &primary, &ctx) + .expect("two non-primary workers are available"); + assert_eq!(backup.id, idle.id); + } +} diff --git a/experimental/sgl-router/src/policies/sticky.rs b/experimental/sgl-router/src/policies/sticky.rs index 541ea7314..265a59996 100644 --- a/experimental/sgl-router/src/policies/sticky.rs +++ b/experimental/sgl-router/src/policies/sticky.rs @@ -212,6 +212,10 @@ impl Policy for StickyPolicy { fn attach_metrics(&self, metrics: Arc) { let _ = self.state.metrics.set(metrics); } + + fn needs_load_snapshot(&self) -> bool { + self.fallback.needs_load_snapshot() + } } impl std::fmt::Debug for StickyPolicy { @@ -228,6 +232,16 @@ impl std::fmt::Debug for StickyPolicy { mod tests { use super::*; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + + #[test] + fn sticky_propagates_fallback_load_snapshot_capability() { + let policy = StickyPolicy::new( + Duration::from_secs(60), + Duration::from_secs(10), + Arc::new(crate::policies::power_of_two::PowerOfTwoChoicesPolicy::new()), + ); + assert!(policy.needs_load_snapshot()); + } use crate::policies::round_robin::RoundRobinPolicy; fn worker(id: &str) -> Arc { diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 5645f35b0..dd70239d8 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -4,6 +4,7 @@ use crate::config::Config; use crate::policies::active_load::ActiveLoadRegistry; +use crate::policies::engine_load::EngineLoadTable; use crate::policies::kv_events::BlockSizeOracle; use crate::policies::PolicyRegistry; use crate::proxy::Proxy; @@ -19,17 +20,16 @@ pub struct AppContext { pub proxy: Arc, pub registry: Arc, pub policies: Arc, - /// Per-worker active-load bookkeeping. Shared between the proxy - /// (which mints guards on the request hot path), the cache-aware - /// policy (which reads per-worker load when scoring candidates), and - /// the stale-request janitor (which sweeps expired entries). + /// Per-worker active-load bookkeeping shared by the proxy, policies, + /// timeout janitor, and metrics. pub active_load: Arc, /// Lightweight Prometheus-format metrics registry served via /// `/metrics`. Shared with the chat handler (requests_total), /// cache-aware-zmq policy (overlap_blocks), active-load registry - /// (active_load gauge + stale_requests_total), and PD resolver - /// (decode_affinity_total). + /// (active_load gauge + stale_requests_total), and PD dispatch. pub metrics: Arc, + /// Shared Engine LoadStat table; ingress captures one immutable snapshot per request. + pub engine_load: Arc, pub prefix_index: Option>, pub block_size_oracle: Arc, ready: AtomicBool, @@ -86,6 +86,7 @@ impl AppContext { metrics, prefix_index: None, block_size_oracle: BlockSizeOracle::new(), + engine_load: EngineLoadTable::new(), ready: AtomicBool::new(false), } } @@ -116,6 +117,7 @@ impl AppContext { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, @@ -135,6 +137,7 @@ impl AppContext { metrics: MetricsRegistry::new(), prefix_index: None, block_size_oracle: BlockSizeOracle::new(), + engine_load: EngineLoadTable::new(), ready: AtomicBool::new(false), } } diff --git a/experimental/sgl-router/src/server/error.rs b/experimental/sgl-router/src/server/error.rs index de4676916..253ba5e2c 100644 --- a/experimental/sgl-router/src/server/error.rs +++ b/experimental/sgl-router/src/server/error.rs @@ -229,10 +229,7 @@ impl IntoResponse for ApiError { ); "request expired before completion".to_string() } - ApiError::PolicySelectionFailed { model } => { - tracing::warn!(model = %model, reason = "policy_selection_failed", "service unavailable"); - "service unavailable".to_string() - } + ApiError::PolicySelectionFailed { .. } => "service unavailable".to_string(), ApiError::BreakerOpen { worker } => { tracing::warn!(upstream = %worker, reason = "breaker_open", "service unavailable"); "service unavailable".to_string() @@ -370,6 +367,21 @@ mod tests { ); } + #[test] + fn policy_selection_failed_envelope_is_unchanged() { + let resp = ApiError::PolicySelectionFailed { + model: "tiny".into(), + } + .into_response(); + let (status, code_header, env) = parse_envelope(resp); + + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(code_header.as_deref(), Some("policy_selection_failed")); + assert_eq!(env.error.typ, "server_error"); + assert_eq!(env.error.code, "policy_selection_failed"); + assert_eq!(env.error.message, "service unavailable"); + } + #[test] fn bad_request_envelope_has_expected_shape() { let msg = "invalid_request: body must be an object"; diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index cb1c8c5c6..89b2d48a4 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -32,6 +32,8 @@ //! | `sgl_router_stale_requests_total` | Counter | `outcome` | //! | `sgl_router_decode_affinity_total` | Counter | `outcome` | //! | `sgl_router_sticky_total` | Counter | `outcome` | +//! | `sgl_router_policy_decisions_total` | Counter | `policy`, `reason` | +//! | `sgl_router_policy_selection_failures_total` | Counter | `policy`, `reason` | //! | `sgl_router_ingress_tokenize_errors_total` | Counter | `model_id` | //! //! The four `sgl_router_worker*` gauges and `sgl_router_workers` are sampled @@ -42,6 +44,7 @@ //! //! The exposition is text/plain; version=0.0.4 per the Prometheus spec. +use crate::config::PolicyKind; use parking_lot::Mutex; use std::collections::HashMap; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; @@ -184,6 +187,23 @@ impl StaleRequestOutcome { } } +#[derive(Debug, Clone, Copy)] +pub(crate) enum PolicySelectionFailureReason { + PrefillAdmissionExhausted, + CacheCandidatesExhausted, + ProposalEmpty, +} + +impl PolicySelectionFailureReason { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::PrefillAdmissionExhausted => "prefill_admission_exhausted", + Self::CacheCandidatesExhausted => "cache_candidates_exhausted", + Self::ProposalEmpty => "proposal_empty", + } + } +} + /// Active-load kind label — separates the two axes of per-worker load. #[derive(Debug, Clone, Copy)] pub enum ActiveLoadKind { @@ -224,6 +244,8 @@ pub struct MetricsRegistry { stale_requests_total: Mutex>>, decode_affinity_total: Mutex>>, sticky_total: Mutex>>, + policy_decisions_total: Mutex>>, + policy_selection_failures_total: Mutex>>, ingress_tokenize_errors_total: Mutex>>, } @@ -274,6 +296,12 @@ struct ActiveLoadKey { kind: &'static str, } +#[derive(Debug, Hash, Eq, PartialEq, Clone)] +struct PolicyDecisionKey { + policy: String, + reason: String, +} + #[derive(Debug)] struct Histogram { /// Bucket upper bounds this histogram observes against (e.g. @@ -481,6 +509,39 @@ impl MetricsRegistry { counter.fetch_add(1, Ordering::Relaxed); } + /// Record the final Prefill policy decision. + pub fn record_policy_decision(&self, policy: &str, reason: &str) { + let key = PolicyDecisionKey { + policy: policy.to_owned(), + reason: reason.to_owned(), + }; + let mut guard = self.policy_decisions_total.lock(); + let counter = guard + .entry(key) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + drop(guard); + counter.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn record_policy_selection_failure( + &self, + policy: PolicyKind, + reason: PolicySelectionFailureReason, + ) { + let key = PolicyDecisionKey { + policy: policy.to_string(), + reason: reason.as_str().to_owned(), + }; + let mut guard = self.policy_selection_failures_total.lock(); + let counter = guard + .entry(key) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + drop(guard); + counter.fetch_add(1, Ordering::Relaxed); + } + /// Bump `sgl_router_ingress_tokenize_errors_total{model_id}`. /// /// Recorded ONLY when the tokenization offload SHOULD have fired but the @@ -785,6 +846,48 @@ impl MetricsRegistry { } drop(guard); + // policy_decisions_total + out.push_str( + "# HELP sgl_router_policy_decisions_total Final Prefill policy decisions by policy and bounded reason.\n", + ); + out.push_str("# TYPE sgl_router_policy_decisions_total counter\n"); + let guard = self.policy_decisions_total.lock(); + let mut entries: Vec<(&PolicyDecisionKey, u64)> = guard + .iter() + .map(|(key, value)| (key, value.load(Ordering::Relaxed))) + .collect(); + entries.sort_by(|a, b| (&a.0.policy, &a.0.reason).cmp(&(&b.0.policy, &b.0.reason))); + for (key, value) in entries { + out.push_str(&format!( + "sgl_router_policy_decisions_total{{policy=\"{}\",reason=\"{}\"}} {}\n", + escape_label(&key.policy), + escape_label(&key.reason), + value, + )); + } + drop(guard); + + // policy_selection_failures_total + out.push_str( + "# HELP sgl_router_policy_selection_failures_total Failed Prefill policy selections by policy and bounded reason.\n", + ); + out.push_str("# TYPE sgl_router_policy_selection_failures_total counter\n"); + let guard = self.policy_selection_failures_total.lock(); + let mut entries: Vec<(&PolicyDecisionKey, u64)> = guard + .iter() + .map(|(key, value)| (key, value.load(Ordering::Relaxed))) + .collect(); + entries.sort_by(|a, b| (&a.0.policy, &a.0.reason).cmp(&(&b.0.policy, &b.0.reason))); + for (key, value) in entries { + out.push_str(&format!( + "sgl_router_policy_selection_failures_total{{policy=\"{}\",reason=\"{}\"}} {}\n", + escape_label(&key.policy), + escape_label(&key.reason), + value, + )); + } + drop(guard); + // ingress_tokenize_errors_total out.push_str( "# HELP sgl_router_ingress_tokenize_errors_total Chat requests on a chat-encoder model whose ingress tokenization failed, silently falling back to engine-side tokenization (the input_ids offload was defeated).\n", @@ -870,6 +973,7 @@ mod tests { 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")); + assert!(out.contains("# TYPE sgl_router_policy_decisions_total counter")); assert!(out.contains("# TYPE sgl_router_ingress_tokenize_errors_total counter")); // Pool-size series exist (at 0) for all three modes even with no // workers, so dashboards have a stable series to graph. @@ -1172,6 +1276,50 @@ mod tests { assert!(out.contains(r#"sgl_router_sticky_total{outcome="no_routing_key"} 1"#)); } + #[test] + fn policy_decisions_are_keyed_by_policy_and_reason() { + let reg = MetricsRegistry::new(); + reg.record_policy_decision("session_aware", "session_primary"); + reg.record_policy_decision("session_aware", "session_primary"); + reg.record_policy_decision("cache_aware", "cache_candidate"); + + let out = reg.render(); + assert!(out.contains( + r#"sgl_router_policy_decisions_total{policy="cache_aware",reason="cache_candidate"} 1"# + )); + assert!(out.contains( + r#"sgl_router_policy_decisions_total{policy="session_aware",reason="session_primary"} 2"# + )); + } + + #[test] + fn policy_selection_failures_are_keyed_by_policy_and_reason() { + let reg = MetricsRegistry::new(); + reg.record_policy_selection_failure( + PolicyKind::SessionAware, + PolicySelectionFailureReason::PrefillAdmissionExhausted, + ); + reg.record_policy_selection_failure( + PolicyKind::CacheAware, + PolicySelectionFailureReason::CacheCandidatesExhausted, + ); + reg.record_policy_selection_failure( + PolicyKind::RoundRobin, + PolicySelectionFailureReason::ProposalEmpty, + ); + + let out = reg.render(); + assert!(out.contains( + r#"sgl_router_policy_selection_failures_total{policy="session_aware",reason="prefill_admission_exhausted"} 1"# + )); + assert!(out.contains( + r#"sgl_router_policy_selection_failures_total{policy="cache_aware",reason="cache_candidates_exhausted"} 1"# + )); + assert!(out.contains( + r#"sgl_router_policy_selection_failures_total{policy="round_robin",reason="proposal_empty"} 1"# + )); + } + #[test] fn ingress_tokenize_error_counter_increments_per_model() { 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 1fce3ec26..531a4aaa4 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -2,13 +2,18 @@ // SPDX-License-Identifier: Apache-2.0 use crate::discovery::{ModelId, WorkerMode}; +use crate::policies::admission::{resolve_cache_candidates, resolve_prefill, CandidateRange}; use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram}; use crate::policies::registry::{PdPoolResolver, PdResolveError}; -use crate::policies::{request_tokens_for, ExternalPrefixSignal, RequestTokens, SelectionContext}; +use crate::policies::{ + request_tokens_for, ExternalPrefixSignal, PrefillProposal, ProposalKind, RequestTokens, + SelectionContext, +}; use crate::server::app_context::AppContext; use crate::server::error::ApiError; use crate::server::metrics::{ - MetricsRegistry, RequestOutcome, StaleRequestOutcome, WorkerModeLabel, + MetricsRegistry, PolicySelectionFailureReason, RequestOutcome, StaleRequestOutcome, + WorkerModeLabel, }; use crate::workers::{LoadGuard, Worker}; use axum::body::Body; @@ -91,6 +96,24 @@ impl Drop for RecordDurationOnDrop { } } +fn policy_selection_failed( + ctx: &AppContext, + model: &str, + reason: PolicySelectionFailureReason, +) -> ApiError { + ctx.metrics + .record_policy_selection_failure(ctx.config.model.policy, reason); + tracing::warn!( + policy = %ctx.config.model.policy, + reason = reason.as_str(), + model, + "prefill policy selection failed" + ); + ApiError::PolicySelectionFailed { + model: model.to_owned(), + } +} + /// POST /v1/chat/completions — parse model from body, select a healthy /// worker via the per-model policy, then proxy the request. If the /// request opts into streaming (`stream: true`), we pipe SSE bytes back; @@ -198,6 +221,15 @@ pub async fn chat_completions( _ => None, }; + let prefill_load = request_tokens + .as_ref() + .map(|tokens| tokens.ids.len().max(1)) + .unwrap_or_else(|| estimate_prefill_tokens(&body)); + let request_input_tokens = prefill_load as u64; + let needs_load_snapshot = policy.needs_load_snapshot(); + let load_snapshot = + needs_load_snapshot.then(|| ctx.engine_load.capture_snapshot(std::time::Instant::now())); + // 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 @@ -210,15 +242,69 @@ pub async fn chat_completions( .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) - .with_request_tokens(request_tokens.as_ref().map(|t| t.ids.as_slice())) + let session_id = ctx + .config + .model + .affinity + .as_ref() + .and_then(|config| headers.get(config.session_id_header.as_str())) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.is_empty()); + let candidate_range = CandidateRange::global(&workers); + let mut selection_ctx = SelectionContext::with_routing_key(&model_id, Some(&body), routing_key) + .with_session_id(session_id) + .with_candidate_range_id(candidate_range.id) + .with_input_tokens(request_input_tokens) + .with_request_tokens(request_tokens.as_ref().map(|tokens| tokens.ids.as_slice())) .with_external_prefix(external_prefix.as_ref()); - let worker = - policy - .select(&workers, &selection_ctx) - .ok_or_else(|| ApiError::PolicySelectionFailed { - model: model_str.clone(), - })?; + if let Some(snapshot) = load_snapshot.as_ref() { + selection_ctx = selection_ctx.with_load_snapshot(snapshot); + } + let worker = match policy.propose_prefill(candidate_range.workers, &selection_ctx) { + Some(PrefillProposal::Pair(proposal)) if policy.uses_shared_prefill_admission() => { + let snapshot = load_snapshot + .as_ref() + .expect("shared prefill admission requires a load snapshot"); + let decision = + resolve_prefill(&candidate_range, &proposal, request_input_tokens, snapshot) + .ok_or_else(|| { + policy_selection_failed( + &ctx, + &model_str, + PolicySelectionFailureReason::PrefillAdmissionExhausted, + ) + })?; + policy.commit_prefill_selection(&selection_ctx, proposal.kind, &decision.selected); + decision.selected + } + Some(PrefillProposal::Pair(proposal)) => proposal.primary, + Some(PrefillProposal::CacheCandidates(proposal)) => { + let snapshot = load_snapshot + .as_ref() + .expect("cache candidate resolution requires a load snapshot"); + let decision = resolve_cache_candidates(&proposal, request_input_tokens, snapshot) + .ok_or_else(|| { + policy_selection_failed( + &ctx, + &model_str, + PolicySelectionFailureReason::CacheCandidatesExhausted, + ) + })?; + policy.commit_prefill_selection( + &selection_ctx, + ProposalKind::CacheAffinity, + &decision.selected, + ); + decision.selected + } + None => { + return Err(policy_selection_failed( + &ctx, + &model_str, + PolicySelectionFailureReason::ProposalEmpty, + )); + } + }; // PD-mode decoder affinity. When the selected prefill worker is // part of a PD-disagg deployment, also resolve the matching decode @@ -291,15 +377,11 @@ pub async fn chat_completions( // 0 here: the active-load registry's decode axis is reserved for a // future decode-side scheduler — current decode selection is // host-affinity only. - let guard = worker.load_guard(); - // Use the exact token count from the ingress tokenization when available; - // fall back to the byte-count heuristic for load-only policies that don't - // tokenize. The exact count makes the cache-aware load-imbalance fast-path - // accurate rather than off by the char/token ratio. - let prefill_load = request_tokens - .as_ref() - .map(|t| t.ids.len().max(1)) - .unwrap_or_else(|| estimate_prefill_tokens(&body)); + let guard = if needs_load_snapshot { + worker.timestamped_load_guard() + } else { + worker.load_guard() + }; let active_guard = ctx.active_load .register(worker.id.clone(), worker.url.clone(), prefill_load, 0); @@ -469,8 +551,9 @@ pub async fn chat_completions( // Synchronously await the decode worker. Its response is what // the client sees. The decode side gets its own LoadGuard so - // per-worker `active_requests` reflects decode-pool load for - // cache-aware-zmq decisions on the decode side. + // per-worker `active_requests` reflects decode-pool load. Decode + // selection reads that atomic counter directly, so it does not need + // the prefill policy's timestamp registry. let decode_guard = decode_worker.load_guard(); if streaming { let stream_guards: Box = diff --git a/experimental/sgl-router/src/server/routes/models.rs b/experimental/sgl-router/src/server/routes/models.rs index 51127f843..e46075346 100644 --- a/experimental/sgl-router/src/server/routes/models.rs +++ b/experimental/sgl-router/src/server/routes/models.rs @@ -54,6 +54,7 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }; diff --git a/experimental/sgl-router/src/server/routes/tokenize.rs b/experimental/sgl-router/src/server/routes/tokenize.rs index 5e05b76bb..d07715b14 100644 --- a/experimental/sgl-router/src/server/routes/tokenize.rs +++ b/experimental/sgl-router/src/server/routes/tokenize.rs @@ -121,6 +121,7 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/src/tokenizer/mod.rs b/experimental/sgl-router/src/tokenizer/mod.rs index 9b600a24c..0646e71bf 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -235,6 +235,7 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/src/workers/introspect.rs b/experimental/sgl-router/src/workers/introspect.rs index 1748b678a..2bfeccdba 100644 --- a/experimental/sgl-router/src/workers/introspect.rs +++ b/experimental/sgl-router/src/workers/introspect.rs @@ -318,6 +318,8 @@ pub(crate) fn resolve_event_config( host, port_base: block.endpoint_port_base, topic: block.topic, + load_port_base: block.load_endpoint_port_base, + load_topic: block.load_topic, block_size: block.block_size, dp_size: block.dp_size, is_bigram, @@ -374,6 +376,12 @@ pub(crate) struct KvEventsBlock { pub endpoint_port_base: u16, #[serde(default)] pub topic: String, + /// Base port of the dedicated load-snapshot socket range. Absent on + /// workers that predate load publishing (`None` ⇒ no load subscriber). + #[serde(default)] + pub load_endpoint_port_base: Option, + #[serde(default)] + pub load_topic: Option, pub block_size: u32, pub dp_size: u32, } diff --git a/experimental/sgl-router/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index e67769a7a..3336177f5 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -481,6 +481,7 @@ mod tests { }), cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/src/workers/worker.rs b/experimental/sgl-router/src/workers/worker.rs index a33a29a83..6bce68945 100644 --- a/experimental/sgl-router/src/workers/worker.rs +++ b/experimental/sgl-router/src/workers/worker.rs @@ -3,8 +3,10 @@ use crate::discovery::{ModelId, WorkerId, WorkerMode}; use crate::health::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; -use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; -use std::sync::Arc; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; /// Parse a host from a worker URL. Matches SMG's `worker_builder.rs` /// fallback chain: parse as-is, retry with `http://` prefix if missing, @@ -31,27 +33,75 @@ fn parse_bootstrap_host(url: &str) -> String { "localhost".to_string() } -/// RAII guard that increments `active_requests` on construction and -/// decrements on drop. Obtain via [`Worker::load_guard`]. +/// Tracks each in-flight slot with an acquisition timestamp so a routing +/// policy can ask how many slots were claimed recently +/// ([`count_acquired_since`](SlotRegistry::count_acquired_since)). The +/// registry is separate from [`Worker::active_requests`]: ordinary load +/// tracking stays lock-free, while policies that correct an engine snapshot +/// explicitly opt into timestamp tracking. +#[derive(Debug)] +pub struct SlotRegistry { + slots: Mutex>, + next_id: AtomicU64, +} + +impl SlotRegistry { + fn new() -> Arc { + Arc::new(Self { + slots: Mutex::new(HashMap::new()), + next_id: AtomicU64::new(0), + }) + } + + /// Records one timestamped slot and returns its identity. + fn claim(&self) -> u64 { + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + self.slots.lock().unwrap().insert(id, Instant::now()); + id + } + + /// Releases timestamped slot `id`. + fn release(&self, id: u64) { + self.slots.lock().unwrap().remove(&id); + } + + /// Count of currently-claimed slots acquired at or after `since`. Used to + /// bound how many of this worker's in-flight requests are dispatches the + /// engine hasn't reported back on yet (see + /// `crate::policies::cache_aware_zmq::WorkerLoads::load_of`), rather than + /// adding the full in-flight count — which would also include long-held + /// slots from slow-draining streaming responses (see + /// `crate::proxy::Proxy::forward_streaming_to`'s `stream_guards` doc) + /// that the engine's own last report likely already accounts for. + pub fn count_acquired_since(&self, since: Instant) -> usize { + self.slots + .lock() + .unwrap() + .values() + .filter(|&&t| t >= since) + .count() + } +} + +/// RAII guard that increments `active_requests` on construction and decrements +/// on drop. Obtain via [`Worker::load_guard`]. Policies that need to correct +/// an engine snapshot use the crate-private timestamped variant. /// /// `#[must_use]`: a statement-form call like `worker.load_guard();` would /// drop the guard on the same line, so the counter would never see the /// in-flight request. The compile-time warning catches that misuse. #[must_use = "LoadGuard must be held for the request's lifetime; dropping it immediately decrements active_requests"] pub struct LoadGuard { - counter: Arc, -} - -impl LoadGuard { - pub(crate) fn new(counter: Arc) -> Self { - counter.fetch_add(1, Ordering::Relaxed); - Self { counter } - } + active_requests: Arc, + tracked_slot: Option<(Arc, u64)>, } impl Drop for LoadGuard { fn drop(&mut self) { - self.counter.fetch_sub(1, Ordering::Relaxed); + if let Some((registry, id)) = &self.tracked_slot { + registry.release(*id); + } + self.active_requests.fetch_sub(1, Ordering::Relaxed); } } @@ -88,6 +138,9 @@ pub struct Worker { pub model_ids: Vec, pub breaker: Arc, pub active_requests: Arc, + /// Timestamped ledger for requests whose policy reads Engine Load; + /// answers [`Worker::slots_acquired_since`]. + slots: Arc, /// Hostname parsed from `url` at construction time and cached. /// Used as the `bootstrap_host` field on PD-disagg requests so the /// prefill engine can match incoming KV-transfer requests from @@ -117,13 +170,16 @@ impl Worker { None => Arc::new(CircuitBreaker::new()), }; let bootstrap_host = parse_bootstrap_host(&spec.url); + let active_requests = Arc::new(AtomicUsize::new(0)); + let slots = SlotRegistry::new(); Self { id: spec.id, url: spec.url, mode: AtomicU8::new(spec.mode.as_u8()), model_ids: spec.model_ids, breaker, - active_requests: Arc::new(AtomicUsize::new(0)), + active_requests, + slots, bootstrap_host, bootstrap_port: spec.bootstrap_port, } @@ -159,10 +215,30 @@ impl Worker { self.active_requests.load(Ordering::Relaxed) } + /// Number of this worker's currently in-flight requests dispatched at or + /// after `since`. See [`SlotRegistry::count_acquired_since`]. + pub fn slots_acquired_since(&self, since: Instant) -> usize { + self.slots.count_acquired_since(since) + } + /// Returns a RAII guard that increments `active_requests` now and /// decrements when the guard is dropped. pub fn load_guard(&self) -> LoadGuard { - LoadGuard::new(self.active_requests.clone()) + self.active_requests.fetch_add(1, Ordering::Relaxed); + LoadGuard { + active_requests: Arc::clone(&self.active_requests), + tracked_slot: None, + } + } + + /// Returns a load guard that also records when the request was dispatched. + pub(crate) fn timestamped_load_guard(&self) -> LoadGuard { + let slot_id = self.slots.claim(); + self.active_requests.fetch_add(1, Ordering::Relaxed); + LoadGuard { + active_requests: Arc::clone(&self.active_requests), + tracked_slot: Some((Arc::clone(&self.slots), slot_id)), + } } } @@ -181,6 +257,7 @@ impl std::fmt::Debug for Worker { mod tests { use super::*; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use std::time::Duration; #[test] fn load_guard_increments_and_decrements() { @@ -202,6 +279,33 @@ mod tests { assert_eq!(w.active_load(), 0); } + #[test] + fn plain_load_guard_does_not_track_a_timestamped_slot() { + let w = test_worker(); + let cutoff = Instant::now() - Duration::from_secs(1); + let guard = w.load_guard(); + + assert_eq!(w.active_load(), 1); + assert_eq!(w.slots_acquired_since(cutoff), 0); + + drop(guard); + assert_eq!(w.active_load(), 0); + } + + #[test] + fn timestamped_load_guard_tracks_and_releases_its_slot() { + let w = test_worker(); + let cutoff = Instant::now() - Duration::from_secs(1); + let guard = w.timestamped_load_guard(); + + assert_eq!(w.active_load(), 1); + assert_eq!(w.slots_acquired_since(cutoff), 1); + + drop(guard); + assert_eq!(w.active_load(), 0); + assert_eq!(w.slots_acquired_since(cutoff), 0); + } + #[test] fn mode_accessor_round_trips_all_variants() { for m in [WorkerMode::Plain, WorkerMode::Prefill, WorkerMode::Decode] { @@ -295,4 +399,53 @@ mod tests { }); assert_eq!(w.bootstrap_host(), "localhost"); } + + fn test_worker() -> Worker { + Worker::new(WorkerSpec { + id: WorkerId("w".into()), + url: "http://x".into(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + }) + } + + #[test] + fn slots_acquired_since_excludes_earlier_slots() { + let w = test_worker(); + let _g_old = w.timestamped_load_guard(); + // A real (small) sleep, not a synthetic `Instant` offset: the slot's + // acquisition time is captured internally by `claim()`, not + // injectable, so the ordering guarantee has to come from wall-clock + // separation wide enough to beat any platform's monotonic-clock + // resolution. + std::thread::sleep(Duration::from_millis(5)); + let cutoff = Instant::now(); + let _g_new1 = w.timestamped_load_guard(); + let _g_new2 = w.timestamped_load_guard(); + assert_eq!(w.active_load(), 3); + assert_eq!( + w.slots_acquired_since(cutoff), + 2, + "only slots claimed at/after cutoff should count" + ); + } + + #[test] + fn slots_acquired_since_counts_all_slots_for_a_cutoff_before_every_claim() { + let w = test_worker(); + let long_ago = Instant::now() - Duration::from_secs(3600); + let _g1 = w.timestamped_load_guard(); + let _g2 = w.timestamped_load_guard(); + assert_eq!(w.slots_acquired_since(long_ago), 2); + } + + #[test] + fn slots_acquired_since_is_zero_for_a_cutoff_after_every_claim() { + let w = test_worker(); + let _g = w.timestamped_load_guard(); + std::thread::sleep(Duration::from_millis(5)); + let cutoff = Instant::now(); + assert_eq!(w.slots_acquired_since(cutoff), 0); + } } diff --git a/experimental/sgl-router/tests/component/discovery/static_urls.rs b/experimental/sgl-router/tests/component/discovery/static_urls.rs index ecdf86259..423a7d6d6 100644 --- a/experimental/sgl-router/tests/component/discovery/static_urls.rs +++ b/experimental/sgl-router/tests/component/discovery/static_urls.rs @@ -133,6 +133,7 @@ async fn static_urls_pd_role_resolved_end_to_end() { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, 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 ae99b7297..38f3ce228 100644 --- a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs @@ -26,6 +26,7 @@ use sgl_router::config::{ActiveLoadConfig, ProxyConfig}; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::cache_aware_zmq::CacheAwareZmqPolicy; +use sgl_router::policies::engine_load::EngineLoadTable; use sgl_router::policies::kv_events::{compute_block_hashes, discovery::EventConfig, KvEventIndex}; use sgl_router::policies::{Policy, SelectionContext}; use sgl_router::tokenizer::TokenizerRegistry; @@ -72,6 +73,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, @@ -116,6 +118,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() { kv_index.tree(), Arc::clone(&tokenizers), block_size_oracle, + EngineLoadTable::new(), ); // 5. Register two workers. They share `127.0.0.1` so both @@ -129,6 +132,8 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() { topic: String::new(), block_size, dp_size: 1, + load_port_base: None, + load_topic: None, is_bigram: false, }; kv_index.add_worker(url_a, Some(preresolved.clone())).await; diff --git a/experimental/sgl-router/tests/component/policies/fused_score.rs b/experimental/sgl-router/tests/component/policies/fused_score.rs index 7a71e865d..2a5f24f11 100644 --- a/experimental/sgl-router/tests/component/policies/fused_score.rs +++ b/experimental/sgl-router/tests/component/policies/fused_score.rs @@ -12,14 +12,17 @@ //! and the idlest at 1.0, so every assertion below holds either way. use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; +use sgl_router::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad}; use sgl_router::policies::kv_events::{ compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId, }; use sgl_router::policies::load_based::LoadBasedPolicy; -use sgl_router::policies::scoring::{prefix_cache::PrefixCachePolicy, FusedScorePolicy}; +use sgl_router::policies::scoring::{ + prefix_cache::PrefixCachePolicy, FusedScorePolicy, ScorePolicy, +}; use sgl_router::policies::{Policy, SelectionContext}; use sgl_router::workers::Worker; -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc, time::Instant}; const BLOCK: usize = 4; @@ -73,3 +76,85 @@ fn the_weight_override_steers_a_two_term_fusion_past_either_term_alone() { assert_eq!(got.id, want.id, "--fuse load_based={load_weight}"); } } + +#[test] +fn fused_load_based_term_uses_the_request_snapshot() { + let ws = vec![worker("w0"), worker("w1")]; + // Local counters changed after the request snapshot and prefer w0. + let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect(); + let snapshot = EngineLoadSnapshot::from_workers( + 29, + HashMap::from([ + ( + ws[0].url.clone(), + EngineWorkerLoad { + num_running_reqs: 50, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + ), + ( + ws[1].url.clone(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + ), + ]), + ); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + let fused = FusedScorePolicy::new(vec![(Arc::new(LoadBasedPolicy::new()), None)]) + .expect("load-based is fusable"); + + assert_eq!( + fused.select(&ws, &ctx).expect("must route").id, + ws[1].id, + "fused score must pass the request snapshot to the load-based term" + ); +} + +#[test] +fn score_policy_forwards_the_request_snapshot_to_load_based() { + let ws = vec![worker("w0"), worker("w1")]; + let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect(); + let snapshot = EngineLoadSnapshot::from_workers( + 31, + HashMap::from([ + ( + ws[0].url.clone(), + EngineWorkerLoad { + num_running_reqs: 50, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + ), + ( + ws[1].url.clone(), + EngineWorkerLoad { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 0, + captured_at: Instant::now(), + }, + ), + ]), + ); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot); + let score = ScorePolicy::new(Arc::new(LoadBasedPolicy::new())); + + assert_eq!( + score.select(&ws, &ctx).expect("must route").id, + ws[1].id, + "ScorePolicy must preserve the load-based snapshot contract" + ); +} diff --git a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs index 87f4afe1c..b16a86da0 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs @@ -41,6 +41,8 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { topic: String::new(), block_size, dp_size: 1, + load_port_base: None, + load_topic: None, is_bigram: false, }; @@ -53,12 +55,7 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { router_a.add_worker(worker_url, Some(cfg.clone())).await; router_b.add_worker(worker_url, Some(cfg.clone())).await; - // SUB-side handshake settle. Publishing before the subscribers - // finish their initial connect loses messages in PUB/SUB semantics; - // the polling loop below would then never converge. - tokio::time::sleep(Duration::from_millis(200)).await; - - // 3. Publish a deterministic, multi-block event chain. + // 3. Build a deterministic, multi-block event chain. let tokens: Vec = (0..16).collect(); let hashes = compute_block_hashes(&tokens, block_size as usize); assert!( @@ -68,21 +65,21 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { ); let event_bytes = encode_block_stored_event(&hashes, None, &tokens, block_size); let payload = encode_event_batch(0.0, vec![event_bytes], Some(0)); - publisher - .send(build_multipart(1, payload)) - .await - .expect("publish BlockStored"); - - // 4. Poll both trees until both report the FULL chain matched. The - // SUB→mpsc→pump→tree pipeline is async; loopback delivery is - // reliable but not instantaneous. + // 4. Republish until both subscribers observe the event. PUB/SUB has + // no readiness acknowledgement, so a one-shot send can race a new + // subscriber's handshake under a parallel test load. let target = hashes.len(); let key = KvWorkerId { url: worker_url.into(), dp_rank: 0, }; let start = std::time::Instant::now(); + let mut sequence = 1i64; loop { + publisher + .send(build_multipart(sequence, payload.clone())) + .await + .expect("publish BlockStored"); let ma = router_a.tree().match_prefix(None, &hashes); let mb = router_b.tree().match_prefix(None, &hashes); let converged = ma.matched_blocks == target @@ -112,6 +109,7 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { ma.matched_blocks, ma.workers, mb.matched_blocks, mb.workers, ); } + sequence += 1; tokio::time::sleep(Duration::from_millis(20)).await; } @@ -174,6 +172,8 @@ async fn two_subscribers_merge_events_from_two_publishers() { topic: String::new(), block_size, dp_size: 1, + load_port_base: None, + load_topic: None, is_bigram: false, }; let cfg_y = EventConfig { @@ -182,6 +182,8 @@ async fn two_subscribers_merge_events_from_two_publishers() { topic: String::new(), block_size, dp_size: 1, + load_port_base: None, + load_topic: None, is_bigram: false, }; @@ -193,10 +195,6 @@ async fn two_subscribers_merge_events_from_two_publishers() { router_b.add_worker(worker_x, Some(cfg_x.clone())).await; router_b.add_worker(worker_y, Some(cfg_y.clone())).await; - // Four SUB→PUB handshakes need to settle before publishing; missed - // SUBSCRIBE frames lose messages forever in PUB/SUB semantics. - tokio::time::sleep(Duration::from_millis(200)).await; - // Two non-overlapping token streams → two distinct hash chains. The // gap between them (0..16 vs 1000..1016) keeps `compute_block_hashes` // outputs disjoint so a cross-attribution bug can't be masked by @@ -221,15 +219,6 @@ async fn two_subscribers_merge_events_from_two_publishers() { )], Some(0), ); - pub_x - .send(build_multipart(1, payload_x)) - .await - .expect("publish on pub_x"); - pub_y - .send(build_multipart(1, payload_y)) - .await - .expect("publish on pub_y"); - let key_x = KvWorkerId { url: worker_x.into(), dp_rank: 0, @@ -242,7 +231,16 @@ async fn two_subscribers_merge_events_from_two_publishers() { let target_y = hashes_y.len(); let start = std::time::Instant::now(); + let mut sequence = 1i64; loop { + pub_x + .send(build_multipart(sequence, payload_x.clone())) + .await + .expect("publish on pub_x"); + pub_y + .send(build_multipart(sequence, payload_y.clone())) + .await + .expect("publish on pub_y"); let ax = router_a.tree().match_prefix(None, &hashes_x); let ay = router_a.tree().match_prefix(None, &hashes_y); let bx = router_b.tree().match_prefix(None, &hashes_x); @@ -298,6 +296,7 @@ async fn two_subscribers_merge_events_from_two_publishers() { by.workers, ); } + sequence += 1; tokio::time::sleep(Duration::from_millis(20)).await; } diff --git a/experimental/sgl-router/tests/component/workers/manager.rs b/experimental/sgl-router/tests/component/workers/manager.rs index aeb62f5ee..27fbe4fce 100644 --- a/experimental/sgl-router/tests/component/workers/manager.rs +++ b/experimental/sgl-router/tests/component/workers/manager.rs @@ -7,9 +7,9 @@ use sgl_router::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, Worke use sgl_router::workers::{manager, WorkerRegistry}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; use tokio::net::TcpListener; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{mpsc, oneshot, Barrier}; /// Spin up a tiny fake worker that returns `body` on `GET /server_info`. /// Returns the worker base URL and a shutdown channel. @@ -48,6 +48,19 @@ fn spec_for(id: &str, url: &str, mode: WorkerMode) -> WorkerSpec { } } +async fn wait_until(condition: impl Fn() -> bool, description: &str) { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if condition() { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {description}")); +} + #[tokio::test] async fn manager_processes_added_then_removed() { let (url_a, _s_a) = spawn_fake_worker(json!({"served_model_name": "m"})).await; @@ -72,8 +85,11 @@ async fn manager_processes_added_then_removed() { .await .unwrap(); - // Give the manager time to drain. - tokio::time::sleep(Duration::from_millis(200)).await; + wait_until( + || registry.workers_for(&ModelId("m".into())).len() == 2, + "both workers to register", + ) + .await; assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 2); tx.send(DiscoveryEvent::Removed { @@ -81,7 +97,11 @@ async fn manager_processes_added_then_removed() { }) .await .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; + wait_until( + || registry.workers_for(&ModelId("m".into())).len() == 1, + "removed worker to leave the registry", + ) + .await; assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 1); drop(tx); @@ -103,7 +123,16 @@ async fn manager_handles_mode_changed() { ))) .await .unwrap(); - tokio::time::sleep(Duration::from_millis(200)).await; + wait_until( + || { + registry + .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill) + .len() + == 1 + }, + "prefill worker to register", + ) + .await; assert_eq!( registry .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill) @@ -117,7 +146,19 @@ async fn manager_handles_mode_changed() { }) .await .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; + wait_until( + || { + registry + .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill) + .is_empty() + && registry + .workers_for_mode(&ModelId("m".into()), WorkerMode::Decode) + .len() + == 1 + }, + "worker mode to change to decode", + ) + .await; assert_eq!( registry .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill) @@ -280,7 +321,15 @@ async fn manager_handles_duplicate_added_as_upsert() { ))) .await .unwrap(); - tokio::time::sleep(Duration::from_millis(300)).await; + wait_until( + || { + registry + .get(&WorkerId("w1".into())) + .is_some_and(|worker| worker.url == url_second) + }, + "replacement Added event to update the worker", + ) + .await; assert_eq!( registry.workers_for(&ModelId("m1".into())).len(), @@ -319,6 +368,32 @@ async fn spawn_slow_worker(body: Value, delay: Duration) -> (String, oneshot::Se (format!("http://127.0.0.1:{port}"), tx) } +async fn spawn_gated_worker(body: Value, gate: Arc) -> (String, oneshot::Sender<()>) { + let body = Arc::new(body); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let app = Router::new().route( + "/server_info", + get(move || { + let body = body.clone(); + let gate = Arc::clone(&gate); + async move { + gate.wait().await; + Json((*body).clone()) + } + }), + ); + let (tx, rx) = oneshot::channel::<()>(); + tokio::spawn(async move { + let _ = axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = rx.await; + }) + .await; + }); + (format!("http://127.0.0.1:{port}"), tx) +} + /// Spawn a fake worker that counts each `GET /server_info` hit in the /// returned `AtomicUsize`. Used to assert the manager makes exactly /// one round-trip per worker. @@ -350,25 +425,22 @@ async fn spawn_counting_worker(body: Value) -> (String, Arc, onesho (format!("http://127.0.0.1:{port}"), counter, tx) } -/// Registration must run in parallel across multiple `Added` events. -/// Each fake worker delays its `/server_info` by 200ms; with sequential -/// processing the manager would take ≥1000ms for 5 workers. We allow -/// up to 600ms (3x the per-fetch delay) as a generous bound that still -/// rejects the sequential implementation. +/// Registration must start every `/server_info` fetch before any response is +/// released. A sequential manager stalls at the first worker's barrier. #[tokio::test] async fn added_events_run_in_parallel() { - let delay = Duration::from_millis(200); let n = 5; + let gate = Arc::new(Barrier::new(n + 1)); let mut workers = Vec::new(); for _ in 0..n { - workers.push(spawn_slow_worker(json!({"served_model_name": "m"}), delay).await); + workers + .push(spawn_gated_worker(json!({"served_model_name": "m"}), Arc::clone(&gate)).await); } let (tx, rx) = mpsc::channel(16); let registry = Arc::new(WorkerRegistry::default()); let h = tokio::spawn(manager::run(rx, registry.clone())); - let start = Instant::now(); for (i, (url, _s)) in workers.iter().enumerate() { tx.send(DiscoveryEvent::Added(spec_for( &format!("w{i}"), @@ -378,22 +450,22 @@ async fn added_events_run_in_parallel() { .await .unwrap(); } - let registered = tokio::time::timeout(Duration::from_secs(5), async { + assert!( + tokio::time::timeout(Duration::from_secs(2), gate.wait()) + .await + .is_ok(), + "manager did not start all {n} /server_info requests concurrently" + ); + let registered = tokio::time::timeout(Duration::from_secs(2), async { loop { if registry.workers_for(&ModelId("m".into())).len() == n { - return true; + return; } tokio::time::sleep(Duration::from_millis(10)).await; } }) .await; - let elapsed = start.elapsed(); assert!(registered.is_ok(), "manager failed to register {n} workers"); - assert!( - elapsed < Duration::from_millis(600), - "registration of {n} workers took {elapsed:?}; sequential per-worker /server_info \ - fetches would take ≥1000ms — parallel spawn is required" - ); drop(tx); h.await.unwrap(); diff --git a/experimental/sgl-router/tests/e2e/infra/model_pool.py b/experimental/sgl-router/tests/e2e/infra/model_pool.py index c60bbee64..3541a8e9b 100644 --- a/experimental/sgl-router/tests/e2e/infra/model_pool.py +++ b/experimental/sgl-router/tests/e2e/infra/model_pool.py @@ -34,6 +34,18 @@ from .model_specs import get_model_spec logger = logging.getLogger(__name__) +def _wait_for_process_group_exit(pgid: int, timeout: float) -> bool: + deadline = time.monotonic() + timeout + while True: + try: + os.killpg(pgid, 0) + except ProcessLookupError: + return True + if time.monotonic() >= deadline: + return False + time.sleep(0.1) + + def _get_open_port() -> int: """Allocate an ephemeral TCP port in the range [20000, 55535]. @@ -72,6 +84,7 @@ class ModelInstance: model_id: str gpu_ids: list[int] = field(default_factory=list) kv_events_endpoint: str | None = None + _shutdown_started: bool = field(default=False, init=False, repr=False) def __enter__(self) -> "ModelInstance": return self @@ -80,16 +93,31 @@ class ModelInstance: self.shutdown() def shutdown(self) -> None: - if self.process is not None and self.process.poll() is None: - try: - self.process.send_signal(signal.SIGTERM) - try: - self.process.wait(timeout=60) - except subprocess.TimeoutExpired: - self.process.kill() - self.process.wait() - except ProcessLookupError: - pass + if self.process is None or self._shutdown_started: + return + self._shutdown_started = True + pgid = self.process.pid + + try: + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + return + + try: + self.process.wait(timeout=60) + except subprocess.TimeoutExpired: + pass + + if _wait_for_process_group_exit(pgid, timeout=30): + return + + try: + os.killpg(pgid, signal.SIGKILL) + except ProcessLookupError: + return + self.process.wait() + if not _wait_for_process_group_exit(pgid, timeout=5): + raise RuntimeError(f"worker process group {pgid} did not exit") def spawn_worker( diff --git a/experimental/sgl-router/tests/e2e/infra/test_model_pool.py b/experimental/sgl-router/tests/e2e/infra/test_model_pool.py new file mode 100644 index 000000000..4f1c0e208 --- /dev/null +++ b/experimental/sgl-router/tests/e2e/infra/test_model_pool.py @@ -0,0 +1,50 @@ +import signal + +from infra import model_pool + + +class _Process: + pid = 1234 + + def __init__(self): + self.wait_timeouts = [] + + def poll(self): + return None + + def send_signal(self, sig): + raise AssertionError(f"signaled only the parent process: {sig}") + + def wait(self, timeout=None): + self.wait_timeouts.append(timeout) + return 0 + + +def test_shutdown_waits_for_the_worker_process_group(monkeypatch): + process = _Process() + signals = [] + probes = iter([True, True, False]) + + def killpg(pgid, sig): + signals.append((pgid, sig)) + if sig == 0 and not next(probes): + raise ProcessLookupError + + monkeypatch.setattr(model_pool.os, "killpg", killpg) + monkeypatch.setattr(model_pool.time, "sleep", lambda _: None) + + instance = model_pool.ModelInstance( + url="http://127.0.0.1:30000", + port=30000, + process=process, + model_id="qwen3-0.6b", + ) + instance.shutdown() + + assert signals == [ + (process.pid, signal.SIGTERM), + (process.pid, 0), + (process.pid, 0), + (process.pid, 0), + ] + assert process.wait_timeouts == [60] diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py b/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py index e2795ac6e..f3c1e5171 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py @@ -11,6 +11,7 @@ Teardown: ./tests/e2e/k8s_integration/setup.sh teardown from __future__ import annotations +import json import logging import socket import subprocess @@ -97,6 +98,60 @@ def _wait_for_pod_ready( ) +def _wait_for_replacement_pod_ready( + old_pod: str, + selector: str, + namespace: str = NAMESPACE, + timeout: int = 120, + interval: float = 0.5, +) -> str: + deadline = time.time() + timeout + last_observed = "no pods" + + while time.time() < deadline: + result = _kubectl( + "get", + "pods", + "-n", + namespace, + "-l", + selector, + "-o", + "json", + check=False, + ) + if getattr(result, "returncode", 0) == 0: + pods = json.loads(result.stdout or "{}").get("items", []) + names = [pod.get("metadata", {}).get("name", "") for pod in pods] + last_observed = ", ".join(filter(None, names)) or "no pods" + + if old_pod not in names: + for pod in sorted( + pods, key=lambda item: item.get("metadata", {}).get("name", "") + ): + metadata = pod.get("metadata", {}) + status = pod.get("status", {}) + ready = any( + condition.get("type") == "Ready" + and condition.get("status") == "True" + for condition in status.get("conditions", []) + ) + if ( + metadata.get("name") != old_pod + and not metadata.get("deletionTimestamp") + and status.get("phase") == "Running" + and ready + ): + return metadata["name"] + + time.sleep(interval) + + raise TimeoutError( + f"No ready replacement for pod {old_pod!r} with selector {selector!r} " + f"after {timeout}s; last observed: {last_observed}" + ) + + def _wait_for_port(port: int, proc: subprocess.Popen, timeout: int = 15) -> None: """Poll until a TCP connection to localhost:port succeeds.""" deadline = time.time() + timeout diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py new file mode 100644 index 000000000..a89449860 --- /dev/null +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py @@ -0,0 +1,54 @@ +import json +from types import SimpleNamespace + +import conftest as k8s_conftest + + +def _pod(name: str, phase: str, ready: bool) -> dict: + return { + "metadata": {"name": name}, + "status": { + "phase": phase, + "conditions": [ + { + "type": "Ready", + "status": "True" if ready else "False", + } + ], + }, + } + + +def test_wait_for_replacement_pod_ignores_old_and_pending_pods(monkeypatch): + old_pod = "sgl-router-old" + new_pod = "sgl-router-new" + responses = iter( + [ + [_pod(old_pod, "Running", True)], + [ + _pod(old_pod, "Running", True), + _pod(new_pod, "Running", True), + ], + [_pod(new_pod, "Pending", False)], + [_pod(new_pod, "Running", True)], + ] + ) + calls = [] + + def fake_kubectl(*args, **kwargs): + calls.append((args, kwargs)) + return SimpleNamespace(stdout=json.dumps({"items": next(responses)})) + + monkeypatch.setattr(k8s_conftest, "_kubectl", fake_kubectl) + monkeypatch.setattr(k8s_conftest.time, "sleep", lambda _: None) + + replacement = k8s_conftest._wait_for_replacement_pod_ready( + old_pod, + "app=sgl-router", + timeout=5, + interval=0, + ) + + assert replacement == new_pod + assert len(calls) == 4 + assert all("-o" in args and "json" in args for args, _ in calls) diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_lifecycle.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_lifecycle.py index ab94b32c3..47ce7483d 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/test_lifecycle.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_lifecycle.py @@ -13,10 +13,7 @@ by driving the deployment scale. from __future__ import annotations -import logging - import httpx -import pytest from conftest import ( NAMESPACE, _cleanup_port_forward, @@ -24,7 +21,7 @@ from conftest import ( _poll_until, _port_forward_start, _wait_for_deployment_ready, - logger, + _wait_for_replacement_pod_ready, ) ROUTER_RESTART_PORT = 8092 @@ -132,7 +129,10 @@ class TestRouterRestart: _cleanup_port_forward("router-restart-pre-kill", pf_holder[0]) pf_holder[0] = None - _wait_for_deployment_ready("sgl-router") + if old_pod: + _wait_for_replacement_pod_ready(old_pod, "app=sgl-router") + else: + _wait_for_deployment_ready("sgl-router") pf_holder[0] = _port_forward_start( NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090 diff --git a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs index a307343df..155fa9f63 100644 --- a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs @@ -25,6 +25,7 @@ use sgl_router::config::{ PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; +use sgl_router::policies::engine_load::EngineLoadTable; use sgl_router::policies::factory::build_registry; use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree}; use sgl_router::proxy::Proxy; @@ -54,6 +55,7 @@ fn config() -> Config { circuit_breaker: None, cache_aware: Some(CacheAwareConfig::default()), sticky: None, + affinity: None, fused: None, eligibility: None, }, @@ -88,6 +90,7 @@ fn build_ctx(url: String) -> Arc { Arc::new(HashTree::new()), Arc::clone(&tokenizers), BlockSizeOracle::new(), + EngineLoadTable::new(), ) .unwrap(), ); diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index e91f9d808..7d9bedafb 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -37,6 +37,7 @@ fn config_for(_worker_url: &str) -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs index 07e404727..59deef3cc 100644 --- a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs +++ b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs @@ -29,6 +29,7 @@ pub fn config() -> Config { circuit_breaker: None, cache_aware: Some(CacheAwareConfig::default()), sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/external_indexer_routing.rs b/experimental/sgl-router/tests/proxy/external_indexer_routing.rs index 54501c9d2..7af27f0dd 100644 --- a/experimental/sgl-router/tests/proxy/external_indexer_routing.rs +++ b/experimental/sgl-router/tests/proxy/external_indexer_routing.rs @@ -17,6 +17,7 @@ use sgl_kv_indexer::{ server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerService, PrefixIndexConfig, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; +use sgl_router::policies::engine_load::EngineLoadTable; use sgl_router::policies::factory::build_registry; use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree}; use sgl_router::policies::request_tokens_for; @@ -94,6 +95,7 @@ async fn external_indexer_routes_to_the_cached_worker() { Arc::new(HashTree::new()), Arc::clone(&tokenizers), Arc::clone(&oracle), + EngineLoadTable::new(), ) .unwrap(), ); diff --git a/experimental/sgl-router/tests/proxy/failover.rs b/experimental/sgl-router/tests/proxy/failover.rs index a0a828282..acd6f11ca 100644 --- a/experimental/sgl-router/tests/proxy/failover.rs +++ b/experimental/sgl-router/tests/proxy/failover.rs @@ -41,6 +41,7 @@ async fn failover_when_one_worker_dies() { }), cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 041d5e9c3..9c5cd892c 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -47,6 +47,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/header_forwarding.rs b/experimental/sgl-router/tests/proxy/header_forwarding.rs index a7a8a701e..a76c978a6 100644 --- a/experimental/sgl-router/tests/proxy/header_forwarding.rs +++ b/experimental/sgl-router/tests/proxy/header_forwarding.rs @@ -34,6 +34,7 @@ async fn forwards_whitelisted_headers_strips_others() { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/main.rs b/experimental/sgl-router/tests/proxy/main.rs index 4b7d440ce..3a4002d0d 100644 --- a/experimental/sgl-router/tests/proxy/main.rs +++ b/experimental/sgl-router/tests/proxy/main.rs @@ -19,6 +19,7 @@ mod header_forwarding; mod pd_bootstrap_injection; mod pd_pool_isolation; mod roundrobin_input_ids; +mod shared_prefill_admission; mod sticky_input_ids; 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 5df65d763..a4c1d7147 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -49,6 +49,7 @@ fn config() -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index 4fd8b0cd4..99fccdbc1 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -48,6 +48,7 @@ fn config() -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs index 54fe93c0a..b0015205c 100644 --- a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs @@ -45,6 +45,7 @@ fn config() -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs b/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs new file mode 100644 index 000000000..f95210299 --- /dev/null +++ b/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs @@ -0,0 +1,427 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use sgl_router::config::{ + ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind, + ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, +}; +use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; +use sgl_router::policies::engine_load::LoadStat; +use sgl_router::policies::{ + CacheCandidate, CacheCandidateProposal, Policy, PolicyRegistry, PrefillProposal, ProposalKind, + SelectionContext, SelectionProposal, +}; +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::{Worker, WorkerRegistry}; +use tower::ServiceExt; + +use crate::common::mock_worker::MockWorker; + +#[derive(Debug)] +struct AdmissionProbePolicy { + primary: Arc, + backup: Arc, + committed: Arc>>, +} + +impl Policy for AdmissionProbePolicy { + fn select(&self, _: &[Arc], _: &SelectionContext<'_>) -> Option> { + panic!("chat routing must resolve the prefill proposal before selection") + } + + fn propose(&self, _: &[Arc], _: &SelectionContext<'_>) -> Option { + Some( + SelectionProposal::with_backup(Arc::clone(&self.primary), Arc::clone(&self.backup)) + .with_kind(ProposalKind::SessionAffinity), + ) + } + + fn uses_shared_prefill_admission(&self) -> bool { + true + } + + fn commit_prefill_selection( + &self, + _: &SelectionContext<'_>, + _: ProposalKind, + selected: &Arc, + ) { + *self.committed.lock().unwrap() = Some(selected.id.0.clone()); + } +} + +#[derive(Debug)] +struct EmptyPolicy; + +impl Policy for EmptyPolicy { + fn select(&self, _: &[Arc], _: &SelectionContext<'_>) -> Option> { + None + } +} + +#[derive(Debug)] +struct InvalidPairPolicy { + outsider: Arc, +} + +impl Policy for InvalidPairPolicy { + fn select(&self, _: &[Arc], _: &SelectionContext<'_>) -> Option> { + panic!("chat routing must use the invalid prefill proposal") + } + + fn propose(&self, _: &[Arc], _: &SelectionContext<'_>) -> Option { + Some(SelectionProposal::primary(Arc::clone(&self.outsider))) + } + + fn uses_shared_prefill_admission(&self) -> bool { + true + } +} + +#[derive(Debug)] +struct CacheCandidatesPolicy { + worker: Arc, +} + +#[derive(Debug)] +struct SnapshotProbePolicy { + worker: Arc, + needs_snapshot: bool, + observed_snapshot: Arc>>, +} + +impl Policy for SnapshotProbePolicy { + fn select(&self, _: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + *self.observed_snapshot.lock().unwrap() = Some(ctx.load_snapshot().is_some()); + Some(Arc::clone(&self.worker)) + } + + fn needs_load_snapshot(&self) -> bool { + self.needs_snapshot + } +} + +impl Policy for CacheCandidatesPolicy { + fn select(&self, _: &[Arc], _: &SelectionContext<'_>) -> Option> { + panic!("chat routing must use the cache-candidate proposal") + } + + fn propose_prefill( + &self, + _: &[Arc], + _: &SelectionContext<'_>, + ) -> Option { + Some(PrefillProposal::CacheCandidates(CacheCandidateProposal { + candidates: vec![CacheCandidate { + worker: Arc::clone(&self.worker), + matched_prefix_tokens: 1, + uncached_tokens: 1, + candidate_range_id: "global".into(), + max_pending_prefill_tokens: None, + }], + cache_switch_margin_tokens: 0, + })) + } + + fn needs_load_snapshot(&self) -> bool { + true + } +} + +fn config(policy: PolicyKind) -> Config { + 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, + circuit_breaker: None, + cache_aware: None, + sticky: None, + affinity: None, + fused: None, + eligibility: None, + }, + discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { + urls: vec!["http://placeholder:0".into()], + }), + proxy: ProxyConfig::default(), + active_load: ActiveLoadConfig::default(), + } +} + +struct TestFixture { + ctx: Arc, + backends: Vec, + workers: Vec>, +} + +async fn fixture( + policy_kind: PolicyKind, + build_policy: impl FnOnce(&[Arc]) -> Arc, +) -> TestFixture { + let backends = vec![ + MockWorker::start(vec![]).await, + MockWorker::start(vec![]).await, + ]; + let cfg = config(policy_kind); + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); + let registry = Arc::new(WorkerRegistry::default()); + for (index, backend) in backends.iter().enumerate() { + registry + .add(WorkerSpec { + id: WorkerId(if index == 0 { "primary" } else { "backup" }.into()), + url: backend.url.clone(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + }) + .unwrap(); + } + let registered = registry.workers_for(&ModelId("tiny".into())); + let workers = ["primary", "backup"] + .into_iter() + .map(|id| { + registered + .iter() + .find(|worker| worker.id.0 == id) + .cloned() + .unwrap() + }) + .collect::>(); + let policies = Arc::new(PolicyRegistry::default()); + policies.insert(ModelId("tiny".into()), build_policy(&workers)); + let ctx = Arc::new(AppContext::new( + cfg, + tokenizers, + Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()), + registry, + policies, + )); + TestFixture { + ctx, + backends, + workers, + } +} + +async fn send_chat(ctx: &Arc) -> StatusCode { + build_router(Arc::clone(ctx)) + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + serde_json::to_vec(&serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(), + )) + .unwrap(), + ) + .await + .unwrap() + .status() +} + +#[tokio::test] +async fn chat_attaches_load_snapshot_only_when_the_policy_needs_it() { + for needs_snapshot in [false, true] { + let observed_snapshot = Arc::new(Mutex::new(None)); + let fixture = fixture(PolicyKind::RoundRobin, |workers| { + Arc::new(SnapshotProbePolicy { + worker: Arc::clone(&workers[0]), + needs_snapshot, + observed_snapshot: Arc::clone(&observed_snapshot), + }) + }) + .await; + + assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK); + assert_eq!(*observed_snapshot.lock().unwrap(), Some(needs_snapshot)); + } +} + +fn assert_failure_metric(ctx: &AppContext, policy: &str, expected_reason: &str) { + let metrics = ctx.metrics.render(); + assert!(metrics.contains(&format!( + "sgl_router_policy_selection_failures_total{{policy=\"{policy}\",reason=\"{expected_reason}\"}} 1" + ))); + for other in [ + "prefill_admission_exhausted", + "cache_candidates_exhausted", + "proposal_empty", + ] { + if other != expected_reason { + assert!( + !metrics.contains(&format!("reason=\"{other}\"")), + "{metrics}" + ); + } + } +} + +#[tokio::test] +async fn chat_commits_the_admitted_prefill_backup() { + let committed = Arc::new(Mutex::new(None)); + let fixture = fixture(PolicyKind::SessionAware, |workers| { + Arc::new(AdmissionProbePolicy { + primary: Arc::clone(&workers[0]), + backup: Arc::clone(&workers[1]), + committed: Arc::clone(&committed), + }) + }) + .await; + fixture.ctx.engine_load.set( + &fixture.workers[0].url, + 0, + LoadStat { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 100, + max_total_num_tokens: 100, + }, + Instant::now(), + ); + + assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK); + assert!(fixture.backends[0] + .captured + .lock() + .unwrap() + .last_body + .is_none()); + assert!(fixture.backends[1] + .captured + .lock() + .unwrap() + .last_body + .is_some()); + assert_eq!(committed.lock().unwrap().as_deref(), Some("backup")); + assert!(!fixture + .ctx + .metrics + .render() + .contains("sgl_router_policy_selection_failures_total{")); +} + +#[tokio::test] +async fn capacity_exhaustion_does_not_return_503() { + let committed = Arc::new(Mutex::new(None)); + let fixture = fixture(PolicyKind::SessionAware, |workers| { + Arc::new(AdmissionProbePolicy { + primary: Arc::clone(&workers[0]), + backup: Arc::clone(&workers[1]), + committed: Arc::clone(&committed), + }) + }) + .await; + for worker in &fixture.workers { + fixture.ctx.engine_load.set( + &worker.url, + 0, + LoadStat { + num_running_reqs: 0, + num_waiting_reqs: 0, + num_tokens: 100, + max_total_num_tokens: 100, + }, + Instant::now(), + ); + } + + assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK); + assert_eq!( + fixture + .backends + .iter() + .filter(|backend| backend.captured.lock().unwrap().last_body.is_some()) + .count(), + 1, + "the request must be dispatched to exactly one legal backend" + ); + assert!(matches!( + committed.lock().unwrap().as_deref(), + Some("primary" | "backup") + )); + assert!(!fixture + .ctx + .metrics + .render() + .contains("sgl_router_policy_selection_failures_total{")); +} + +#[tokio::test] +async fn chat_records_proposal_empty() { + let fixture = fixture(PolicyKind::SessionAware, |_| Arc::new(EmptyPolicy)).await; + + assert_eq!( + send_chat(&fixture.ctx).await, + StatusCode::SERVICE_UNAVAILABLE + ); + assert_failure_metric(&fixture.ctx, "session_aware", "proposal_empty"); +} + +#[tokio::test] +async fn chat_records_prefill_admission_exhausted_for_out_of_range_primary() { + let outsider = Arc::new(Worker::new(WorkerSpec { + id: WorkerId("outsider".into()), + url: "http://outsider:30000".into(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })); + let fixture = fixture(PolicyKind::SessionAware, |_| { + Arc::new(InvalidPairPolicy { + outsider: Arc::clone(&outsider), + }) + }) + .await; + + assert_eq!( + send_chat(&fixture.ctx).await, + StatusCode::SERVICE_UNAVAILABLE + ); + assert_failure_metric(&fixture.ctx, "session_aware", "prefill_admission_exhausted"); +} + +#[tokio::test] +async fn chat_records_cache_candidates_exhausted() { + let fixture = fixture(PolicyKind::CacheAware, |workers| { + Arc::new(CacheCandidatesPolicy { + worker: Arc::clone(&workers[0]), + }) + }) + .await; + fixture.ctx.engine_load.set( + &fixture.workers[0].url, + 0, + LoadStat { + num_running_reqs: 1, + num_waiting_reqs: 0, + num_tokens: 100, + max_total_num_tokens: 100, + }, + Instant::now(), + ); + + assert_eq!( + send_chat(&fixture.ctx).await, + StatusCode::SERVICE_UNAVAILABLE + ); + assert_failure_metric(&fixture.ctx, "cache_aware", "cache_candidates_exhausted"); +} diff --git a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs index bbceb4599..63a4fefec 100644 --- a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs @@ -64,6 +64,7 @@ fn config() -> Config { idle_secs: 3600, eviction_interval_secs: 3600, }), + affinity: None, fused: None, eligibility: None, }, diff --git a/experimental/sgl-router/tests/proxy/sticky_routing.rs b/experimental/sgl-router/tests/proxy/sticky_routing.rs index 9a5f9b5a7..bd272c00c 100644 --- a/experimental/sgl-router/tests/proxy/sticky_routing.rs +++ b/experimental/sgl-router/tests/proxy/sticky_routing.rs @@ -50,6 +50,7 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc Config { circuit_breaker: None, cache_aware: None, sticky: None, + affinity: None, fused: None, eligibility: None, },