diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 8d0e02b2d..f360f78d1 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -12,9 +12,9 @@ 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, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig, - ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, - StickyConfig, + DiscoveryBackend, EligibilityConfig, FilterKind, FusedTerm, K8sDiscoveryConfig, + KvIndexerEndpointConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig, + ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE, }; const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100; @@ -73,19 +73,32 @@ 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 authoritative cache signal. - /// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`. + /// External KV indexer gRPC endpoint used as the cache signal. #[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 issued by this Router. Requires + /// Maximum concurrent KV Indexer queries. Requires /// `--kv-indexer-endpoint`; defaults to 32. #[arg(long)] pub kv_indexer_query_max_inflight: Option, + /// Weighted terms for `--policy fused_score`. + #[arg(long, value_delimiter = ',')] + pub fuse: Vec, + + /// Ordered hard constraints applied before policy selection. + #[arg(long, value_delimiter = ',')] + pub filter: Vec, + /// Router-local in-flight limit for `--filter overloaded`. + #[arg(long)] + pub max_in_flight: Option, + /// Minimum cached prompt share for `--filter prefix_cache`. + #[arg(long)] + pub prefix_cache_min_share: Option, + // ---- sticky-session policy (only used by `--policy sticky`) ---- /// Request header carrying the routing key for sticky-session routing. /// Defaults to `x-sgl-routing-key` when `--policy sticky` is set. @@ -96,7 +109,7 @@ pub struct Cli { /// `round_robin` / `random` / `power_of_two` / `load_based`. Defaults /// to `round_robin`. #[arg(long, value_enum)] - pub sticky_fallback_policy: Option, + pub sticky_fallback_policy: Option, /// Evict a sticky assignment after it has been idle (unreferenced) this /// many seconds. Defaults to 600. #[arg(long)] @@ -179,7 +192,8 @@ impl Cli { || self.kv_indexer_query_max_inflight.is_some(); if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq { return Err(anyhow!( - "cache-aware tuning flags require --policy cache_aware_zmq" + "--cache-threshold / --balance-abs-threshold / --balance-rel-threshold \ + require --policy cache_aware_zmq" )); } if self.kv_indexer_query_timeout_ms == Some(0) { @@ -202,6 +216,57 @@ 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")); + } + let fused = if self.policy == PolicyKind::FusedScore { + let terms = if self.fuse.is_empty() { + DEFAULT_FUSE + .iter() + .map(|&kind| FusedTerm { kind, weight: None }) + .collect() + } else { + self.fuse.clone() + }; + for (i, t) in terms.iter().enumerate() { + if terms[..i].iter().any(|p| p.kind == t.kind) { + return Err(anyhow!("--fuse: `{}` is listed more than once", t.kind)); + } + } + Some(terms) + } else { + None + }; + + for (i, kind) in self.filter.iter().enumerate() { + if self.filter[..i].contains(kind) { + return Err(anyhow!("--filter: `{kind}` is listed more than once")); + } + } + let has = |k: FilterKind| self.filter.contains(&k); + if self.max_in_flight.is_some() != has(FilterKind::Overloaded) { + return Err(anyhow!( + "--max-in-flight and `--filter overloaded` require each other" + )); + } + 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" + )); + } + if self + .prefix_cache_min_share + .is_some_and(|s| !(s > 0.0 && s <= 1.0)) + { + return Err(anyhow!("--prefix-cache-min-share must be in (0, 1]")); + } + let eligibility = (!self.filter.is_empty()).then(|| EligibilityConfig { + filters: self.filter.clone(), + max_in_flight: self.max_in_flight, + min_prefix_share: self.prefix_cache_min_share, + }); + let tuned_sticky = self.routing_key_header.is_some() || self.sticky_fallback_policy.is_some() || self.sticky_idle_secs.is_some() @@ -213,11 +278,10 @@ impl Cli { )); } - // Build (and validate) the sticky config exactly when the sticky + // Build and validate the sticky config exactly when the sticky // policy is selected. The header name must parse as an HTTP header // name so a typo fails at startup rather than silently never - // matching any request header; the fallback must be a - // dependency-free policy the factory can build standalone. + // matching any request header. let sticky = if self.policy == PolicyKind::Sticky { let d = StickyConfig::default(); let header_name = self.routing_key_header.unwrap_or(d.header_name); @@ -225,15 +289,6 @@ impl Cli { anyhow!("--routing-key-header {header_name:?} is not a valid HTTP header name: {e}") })?; let fallback_policy = self.sticky_fallback_policy.unwrap_or(d.fallback_policy); - if matches!( - fallback_policy, - PolicyKind::Sticky | PolicyKind::CacheAwareZmq - ) { - return Err(anyhow!( - "--sticky-fallback-policy must be one of round_robin / random / \ - power_of_two / load_based; cache_aware_zmq and sticky are not allowed" - )); - } let idle_secs = self.sticky_idle_secs.unwrap_or(d.idle_secs); let eviction_interval_secs = self .sticky_eviction_interval_secs @@ -315,6 +370,8 @@ impl Cli { circuit_breaker, cache_aware, sticky, + fused, + eligibility, }, discovery, proxy: ProxyConfig { @@ -401,7 +458,7 @@ fn join_selector(terms: &[String]) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::config::{DiscoveryBackend, K8sDiscoveryMode}; + use crate::config::{DiscoveryBackend, K8sDiscoveryMode, ScoreTermKind}; /// Parse argv (without the leading binary name) into a `Config`. fn into_config(args: &[&str]) -> Result { @@ -703,6 +760,40 @@ mod tests { ); } + #[test] + fn policy_accepts_only_routing_strategies() { + for value in ["prefix_cache", "overloaded"] { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + value, + ])) + .expect_err("score terms and filters are not top-level policies") + .to_string(); + assert!(err.contains(value), "{value}: {err}"); + } + } + + #[test] + fn filters_and_fuse_terms_reject_non_members() { + let cases = [ + (vec!["--filter", "load_based"], "load_based"), + ( + vec!["--policy", "fused_score", "--fuse", "sticky"], + "sticky", + ), + ]; + for (args, value) in cases { + let err = into_config_owned(with_model( + &[&["--worker-urls", "http://x:30000"], &args[..]].concat(), + )) + .expect_err("the option must reject a kind from another layer") + .to_string(); + assert!(err.contains(value), "{value}: {err}"); + } + } + /// `--policy load_based` parses to the load-based selector. #[test] fn parses_load_based_policy() { @@ -967,11 +1058,35 @@ mod tests { assert_eq!(c.model.policy, PolicyKind::Sticky); let s = c.model.sticky.expect("sticky config built"); assert_eq!(s.header_name, "x-sgl-routing-key"); - assert_eq!(s.fallback_policy, PolicyKind::RoundRobin); + assert_eq!(s.fallback_policy, StickyFallbackKind::RoundRobin); assert_eq!(s.idle_secs, 600); assert_eq!(s.eviction_interval_secs, 60); } + #[test] + fn sticky_fallback_help_lists_only_dependency_free_policies() { + use clap::CommandFactory; + + let mut command = Cli::command(); + let mut help = Vec::new(); + command.write_long_help(&mut help).unwrap(); + let help = String::from_utf8(help).unwrap(); + let (_, after) = help + .split_once("--sticky-fallback-policy ") + .expect("sticky fallback option is documented"); + let choices = after + .split_once("--sticky-idle-secs") + .expect("sticky fallback precedes its tuning") + .0; + + for value in ["round_robin", "random", "power_of_two", "load_based"] { + assert!(choices.contains(value), "missing {value}: {choices}"); + } + for value in ["fused_score", "cache_aware_zmq", "sticky"] { + assert!(!choices.contains(value), "unexpected {value}: {choices}"); + } + } + #[test] fn sticky_flags_override_defaults() { let c = into_config_owned(with_model(&[ @@ -991,11 +1106,74 @@ mod tests { .unwrap(); let s = c.model.sticky.expect("sticky config built"); assert_eq!(s.header_name, "x-session-id"); - assert_eq!(s.fallback_policy, PolicyKind::LoadBased); + assert_eq!(s.fallback_policy, StickyFallbackKind::LoadBased); assert_eq!(s.idle_secs, 120); assert_eq!(s.eviction_interval_secs, 15); } + #[test] + fn filter_builds_the_eligibility_config_in_order_and_is_off_by_default() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "round_robin", + "--filter", + "overloaded,prefix_cache", + "--max-in-flight", + "64", + "--prefix-cache-min-share", + "0.6", + ])) + .unwrap(); + let e = c.model.eligibility.expect("--filter must build the config"); + assert_eq!( + e.filters, + vec![FilterKind::Overloaded, FilterKind::PrefixCache], + "order is priority, so it must survive parsing", + ); + assert_eq!((e.max_in_flight, e.min_prefix_share), (Some(64), Some(0.6))); + assert_eq!( + c.model.policy, + PolicyKind::RoundRobin, + "not gated on --policy" + ); + + let bare = into_config_owned(with_model(&["--worker-urls", "http://x:30000"])).unwrap(); + assert!(bare.model.eligibility.is_none(), "no --filter, no layer"); + } + + #[test] + fn filter_misconfigurations_fail_at_startup() { + let cases: [(&[&str], &str); 6] = [ + (&["--filter", "overloaded"], "require each other"), + (&["--max-in-flight", "64"], "require each other"), + (&["--filter", "prefix_cache"], "require each other"), + (&["--prefix-cache-min-share", "0.6"], "require each other"), + ( + &["--filter", "overloaded,overloaded", "--max-in-flight", "64"], + "listed more than once", + ), + ( + &[ + "--filter", + "prefix_cache", + "--prefix-cache-min-share", + "0.0", + ], + "must be in (0, 1]", + ), + ]; + for (extra, want) in cases { + let mut args = vec!["--worker-urls", "http://x:30000"]; + args.extend_from_slice(extra); + let err = into_config_owned(with_model(&args)) + .unwrap_err() + .to_string(); + assert!(err.contains(want), "for {extra:?} got: {err}"); + } + } + #[test] fn non_sticky_policy_leaves_sticky_none() { let c = into_config_owned(with_model(&[ @@ -1049,10 +1227,7 @@ mod tests { ])) .unwrap_err() .to_string(); - assert!( - err.contains("--sticky-fallback-policy must be one of"), - "got: {err}" - ); + assert!(err.contains("invalid value"), "got: {err}"); } #[test] @@ -1067,10 +1242,7 @@ mod tests { ])) .unwrap_err() .to_string(); - assert!( - err.contains("--sticky-fallback-policy must be one of"), - "got: {err}" - ); + assert!(err.contains("invalid value"), "got: {err}"); } /// A zero eviction interval would panic `tokio::time::interval` at @@ -1110,4 +1282,93 @@ mod tests { "got: {err}" ); } + + /// `argv` is space-split, so a case reads as the command line an operator + /// would type. Model + worker URL are supplied. + fn cfg_of(argv: &str) -> Result { + let extra: Vec<&str> = argv.split_whitespace().collect(); + into_config_owned(with_model( + &[&["--worker-urls", "http://10.0.0.1:30000"], &extra[..]].concat(), + )) + } + + fn fuse_err(argv: &str) -> String { + cfg_of(argv).unwrap_err().to_string() + } + + /// Resolved terms as `(kind, weight)` pairs; `None` when the policy is + /// not `fused_score` and so builds no term list at all. + fn fused_of(argv: &str) -> Option)>> { + let ts = cfg_of(argv).unwrap().model.fused?; + Some(ts.iter().map(|t| (t.kind, t.weight)).collect()) + } + + fn fuse_ok(argv: &str) -> Vec<(ScoreTermKind, Option)> { + 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. + #[test] + fn fuse_defaults_to_the_useful_pair_and_parses_weights() { + use ScoreTermKind::{LoadBased, PrefixCache, Random}; + let pair = [(PrefixCache, None), (LoadBased, None)]; + assert_eq!(fuse_ok("--policy fused_score"), pair); + // Comma-separated, order preserved, weight optional per term. + assert_eq!( + fuse_ok("--policy fused_score --fuse load_based=0.3,random"), + [(LoadBased, Some(0.3)), (Random, None)], + ); + assert!(fused_of("").is_none(), "round_robin builds no term list"); + } + + /// Non-finite and negative weights are refused, naming the term. + /// + /// `nan`/`inf` matter more than they look: `str::parse::` accepts + /// both, and a NaN weight makes every worker's fused total NaN, so argmax + /// discards them all and the router silently degrades to least-load. + #[test] + fn fuse_rejects_non_finite_and_negative_weights() { + for bad in ["nan", "NaN", "inf", "-inf", "-0.5", "banana"] { + let err = fuse_err(&format!("--policy fused_score --fuse load_based={bad}")); + assert!(err.contains("load_based"), "{bad}: names the term: {err}"); + assert!( + err.contains("must be finite and >= 0") || err.contains("is not a number"), + "{bad}: {err}", + ); + } + for good in ["0", "0.3", "2", "1e3"] { + let got = fuse_ok(&format!("--policy fused_score --fuse load_based={good}"))[0].1; + assert_eq!(got, Some(good.parse::().unwrap())); + } + } + + #[test] + fn fuse_rejects_malformed_compositions() { + let cases: [(&str, &[&str]); 5] = [ + ("--fuse load_based", &["--fuse requires", "fused_score"]), + ( + "--policy fused_score --fuse fused_score,load_based", + &["fused_score", "not a score term"], + ), + ( + "--policy fused_score --fuse load_based,load_based", + &["load_based", "listed more than once"], + ), + ( + "--policy fused_score --fuse not_a_policy", + &["not_a_policy", "is not a score term"], + ), + ( + "--policy sticky --sticky-fallback-policy prefix_cache", + &["prefix_cache", "invalid value"], + ), + ]; + for (argv, wants) in cases { + let err = fuse_err(argv); + for want in wants { + assert!(err.contains(want), "{argv}: want {want:?}, got: {err}"); + } + } + } } diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index df012c7c7..3dad456d6 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -88,6 +88,8 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: urls.iter().map(|s| s.to_string()).collect(), diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 519f78769..9e56a75ca 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` / `cache_aware_zmq` / `sticky`. +/// `power_of_two` / `load_based` / `fused_score` / `cache_aware_zmq` / +/// `sticky`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)] pub enum PolicyKind { #[default] @@ -83,6 +84,9 @@ pub enum PolicyKind { /// Selects the currently least-loaded worker. #[value(name = "load_based")] LoadBased, + /// Weighted sum of `--fuse` terms. + #[value(name = "fused_score")] + FusedScore, /// 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`. @@ -97,6 +101,78 @@ pub enum PolicyKind { Sticky, } +impl std::fmt::Display for PolicyKind { + /// The CLI spelling for this policy kind. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let v = ::to_possible_value(self) + .expect("PolicyKind skips no variants"); + f.write_str(v.get_name()) + } +} + +/// A hard admission constraint accepted by `--filter`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] +pub enum FilterKind { + /// Router-local in-flight capacity limit. + #[value(name = "overloaded")] + Overloaded, + /// Requires a minimum share of cached prompt blocks. + #[value(name = "prefix_cache")] + PrefixCache, +} + +impl std::fmt::Display for FilterKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let v = ::to_possible_value(self) + .expect("FilterKind skips no variants"); + f.write_str(v.get_name()) + } +} + +/// A soft scoring term accepted by `--fuse`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] +pub enum ScoreTermKind { + /// Independent uniform-random preference. + #[value(name = "random")] + Random, + /// Prefers the least router-local active load. + #[value(name = "load_based")] + LoadBased, + /// Prefers the largest local prefix-cache overlap. + #[value(name = "prefix_cache")] + PrefixCache, +} + +impl std::fmt::Display for ScoreTermKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let v = ::to_possible_value(self) + .expect("ScoreTermKind skips no variants"); + f.write_str(v.get_name()) + } +} + +/// Policy choices that can initialize or handle a keyless sticky request. +/// These policies have no request-scoped cache or sticky-state dependency. +#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] +pub enum StickyFallbackKind { + #[value(name = "round_robin")] + RoundRobin, + #[value(name = "random")] + Random, + #[value(name = "power_of_two")] + PowerOfTwo, + #[value(name = "load_based")] + LoadBased, +} + +impl std::fmt::Display for StickyFallbackKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let v = ::to_possible_value(self) + .expect("StickyFallbackKind skips no variants"); + f.write_str(v.get_name()) + } +} + #[derive(Debug, Clone)] pub struct ServerConfig { pub host: String, @@ -146,7 +222,7 @@ pub struct ModelConfig { pub tokenizer_path: String, pub policy: PolicyKind, pub circuit_breaker: Option, - /// Tuning for cache-aware routing. Ignored unless + /// Tuning for the cache-aware ZMQ policy. Ignored unless /// `policy = "cache_aware_zmq"`. `None` falls back to defaults at /// policy construction time. pub cache_aware: Option, @@ -155,6 +231,10 @@ pub struct ModelConfig { /// The chat handler reads `sticky.header_name` to populate /// [`crate::policies::SelectionContext::routing_key`]. pub sticky: Option, + /// Terms for `policy = "fused_score"`. + pub fused: Option>, + /// Hard constraints applied before policy selection. + pub eligibility: Option, } /// External KV Indexer client settings. @@ -165,6 +245,54 @@ pub struct KvIndexerEndpointConfig { pub query_max_inflight: usize, } +/// Eligibility filter configuration. +#[derive(Debug, Clone, Default)] +pub struct EligibilityConfig { + /// Filters in priority order. + pub filters: Vec, + /// `overloaded`: in-flight count at which a worker stops being eligible. + pub max_in_flight: Option, + /// `prefix_cache` minimum cached prompt share. + pub min_prefix_share: Option, +} + +/// Default `--policy fused_score` terms. +pub const DEFAULT_FUSE: [ScoreTermKind; 2] = [ScoreTermKind::PrefixCache, ScoreTermKind::LoadBased]; + +/// One `--fuse` policy and optional weight. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct FusedTerm { + pub kind: ScoreTermKind, + /// Weight override; `None` keeps the term's own `Criterion::weight()`. + pub weight: Option, +} + +impl std::str::FromStr for FusedTerm { + type Err = String; + fn from_str(s: &str) -> Result { + let (name, weight) = match s.split_once('=') { + Some((n, w)) => (n, Some(parse_fuse_weight(n, w)?)), + None => (s, None), + }; + let kind = ::from_str(name, false) + .map_err(|_| format!("--fuse: `{name}` is not a score term"))?; + Ok(FusedTerm { kind, weight }) + } +} + +/// Parses a finite, non-negative term weight. +fn parse_fuse_weight(name: &str, raw: &str) -> Result { + let w: f32 = raw + .parse() + .map_err(|_| format!("--fuse: `{name}` weight `{raw}` is not a number"))?; + if !w.is_finite() || w < 0.0 { + return Err(format!( + "--fuse: `{name}` weight `{raw}` must be finite and >= 0" + )); + } + Ok(w) +} + /// Per-model cache-aware tuning. #[derive(Debug, Clone)] pub struct CacheAwareConfig { @@ -182,8 +310,7 @@ pub struct CacheAwareConfig { /// that the absolute check is gated on. Default 1.1 — 10 % relative /// difference triggers re-balancing. pub balance_rel_threshold: f32, - /// Optional external KV Indexer client configuration. When configured, it - /// replaces the local ZMQ radix tree as the cache signal. + /// Optional external KV Indexer client configuration. pub kv_indexer_endpoint: Option, } @@ -215,8 +342,7 @@ pub const DEFAULT_STICKY_HEADER: &str = "x-sgl-routing-key"; /// Per-model sticky-session tuning. Built from the `--routing-key-header` /// / `--sticky-*` flags by [`crate::config::cli::Cli::into_config`], which -/// also validates that `header_name` parses as an HTTP header name and -/// that `fallback_policy` is one of the dependency-free policies. +/// also validates that `header_name` parses as an HTTP header name. #[derive(Debug, Clone)] pub struct StickyConfig { /// Request header carrying the routing key. Validated to parse as a @@ -226,9 +352,8 @@ pub struct StickyConfig { /// to pick the initial worker when a new key is first seen. One of /// `round_robin` / `random` / `power_of_two` / `load_based` — the /// dependency-free policies the factory can build standalone (no - /// `HashTree` / tokenizer / ZMQ feed). `cache_aware_zmq` and `sticky` - /// are rejected at config-build time. - pub fallback_policy: PolicyKind, + /// `HashTree` / tokenizer / ZMQ feed). + pub fallback_policy: StickyFallbackKind, /// Evict an assignment after it has been idle (unreferenced) this many /// seconds. Bounds the map against unbounded routing-key cardinality. pub idle_secs: u64, @@ -247,7 +372,7 @@ impl Default for StickyConfig { fn default() -> Self { Self { header_name: DEFAULT_STICKY_HEADER.to_string(), - fallback_policy: PolicyKind::RoundRobin, + fallback_policy: StickyFallbackKind::RoundRobin, idle_secs: default_sticky_idle_secs(), eviction_interval_secs: default_sticky_eviction_interval_secs(), } diff --git a/experimental/sgl-router/src/policies/cache_aware_zmq.rs b/experimental/sgl-router/src/policies/cache_aware_zmq.rs index f8430134b..3332e8929 100644 --- a/experimental/sgl-router/src/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/src/policies/cache_aware_zmq.rs @@ -144,14 +144,12 @@ 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)) .map(|m| m.matched_prefix_blocks) .max() .unwrap_or(0); - let match_rate = best_routable_blocks as f32 / signal.query_blocks as f32; if let Some(metrics) = self.metrics.get() { metrics.observe_overlap_blocks(ctx.model().0.as_str(), best_routable_blocks as u64); @@ -184,8 +182,6 @@ impl Policy for CacheAwareZmqPolicy { return Self::pick_min_load(workers); } - // An external signal is authoritative: an empty/unusable result - // degrades only to min-load and never consults the local radix tree. if let Some(signal) = ctx.external_prefix() { return self .select_external(workers, ctx, signal) @@ -349,6 +345,8 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { @@ -400,7 +398,7 @@ mod tests { } #[test] - fn external_prefix_signal_skips_unroutable_best_match() { + fn external_prefix_signal_selects_the_best_routable_match() { let mut config = cfg_default(); config.cache_threshold = 0.0; let policy = CacheAwareZmqPolicy::new( @@ -438,7 +436,7 @@ mod tests { } #[test] - fn external_empty_result_uses_min_load_without_local_tree() { + fn external_empty_result_uses_min_load() { let tree = Arc::new(HashTree::new()); let registry = tokenizer_registry_with_tiny(); let text = "hello world hello world hello world"; diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index f9d8ca663..e2963cc58 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -use crate::config::{Config, ModelConfig, PolicyKind}; +use crate::config::{ + Config, FilterKind, ModelConfig, PolicyKind, ScoreTermKind, StickyFallbackKind, +}; use crate::discovery::ModelId; use crate::policies::{ cache_aware_zmq::CacheAwareZmqPolicy, @@ -10,27 +12,25 @@ use crate::policies::{ power_of_two::PowerOfTwoChoicesPolicy, random::RandomPolicy, round_robin::RoundRobinPolicy, + scoring::{ + admission::Overloaded, prefix_cache, prefix_cache::PrefixCachePolicy, FusedScorePolicy, + Pipeline, + }, sticky::StickyPolicy, Policy, PolicyRegistry, }; use crate::tokenizer::TokenizerRegistry; -use anyhow::Result; +use anyhow::{anyhow, Result}; use std::sync::Arc; use std::time::Duration; -/// Build a dependency-free policy for use as the sticky-session fallback -/// (keyless requests + initial pin of a new key). `Cli::into_config` -/// validates `--sticky-fallback-policy` to one of these four, so the -/// `CacheAwareZmq`/`Sticky` arms are never reached in practice. -fn build_sticky_fallback(kind: PolicyKind) -> Arc { +/// Build a dependency-free policy for keyless sticky requests and new pins. +fn build_sticky_fallback(kind: StickyFallbackKind) -> Arc { match kind { - PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), - PolicyKind::Random => Arc::new(RandomPolicy::new()), - PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), - PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()), - PolicyKind::CacheAwareZmq | PolicyKind::Sticky => { - unreachable!("sticky fallback is validated to be dependency-free in Cli::into_config") - } + StickyFallbackKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), + StickyFallbackKind::Random => Arc::new(RandomPolicy::new()), + StickyFallbackKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), + StickyFallbackKind::LoadBased => Arc::new(LoadBasedPolicy::new()), } } @@ -46,20 +46,38 @@ fn build_sticky(model: &ModelConfig) -> Arc { )) } -/// Construct a policy for a single model from its [`ModelConfig`] and the -/// process-shared `HashTree` + `TokenizerRegistry` + `BlockSizeOracle`. -/// -/// The tree, tokenizer registry, and oracle are only consulted by the -/// cache-aware-zmq variant; other policies ignore them. Callers building -/// all policies for the same process pass the same instances to every -/// model. +/// Constructs a policy for one model. pub fn build_policy( model: &ModelConfig, tree: Arc, tokenizers: Arc, block_size_oracle: Arc, -) -> Arc { - match model.policy { +) -> Result> { + let inner = build_kind(model.policy, model, &tree, &tokenizers, &block_size_oracle)?; + let Some(elig) = model.eligibility.as_ref().filter(|e| !e.filters.is_empty()) else { + return Ok(inner); + }; + let mut filters = Vec::with_capacity(elig.filters.len()); + for &kind in &elig.filters { + filters.push(build_filter(kind, model, &tree, &block_size_oracle)); + } + Ok(Arc::new(Pipeline::new(filters, inner)?)) +} + +/// Build one policy kind with the shared model dependencies. +fn build_kind( + kind: PolicyKind, + model: &ModelConfig, + tree: &Arc, + tokenizers: &Arc, + block_size_oracle: &Arc, +) -> Result> { + let (tree, tokenizers, block_size_oracle) = ( + Arc::clone(tree), + Arc::clone(tokenizers), + Arc::clone(block_size_oracle), + ); + Ok(match kind { PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), PolicyKind::Random => Arc::new(RandomPolicy::new()), PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), @@ -74,32 +92,90 @@ pub fn build_policy( )) } PolicyKind::Sticky => build_sticky(model), + PolicyKind::FusedScore => build_fused(model, &tree, &block_size_oracle)?, + }) +} + +/// Builds one hard admission filter with the shared model dependencies. +fn build_filter( + kind: FilterKind, + model: &ModelConfig, + tree: &Arc, + block_size_oracle: &Arc, +) -> Arc { + match kind { + FilterKind::Overloaded => { + let cap = (model.eligibility.as_ref()) + .and_then(|e| e.max_in_flight) + .unwrap_or(usize::MAX); + Arc::new(Overloaded::new(cap)) + } + FilterKind::PrefixCache => { + let share = (model.eligibility.as_ref()) + .and_then(|e| e.min_prefix_share) + .unwrap_or(0.0); + Arc::new( + PrefixCachePolicy::new( + Arc::clone(tree), + Arc::clone(block_size_oracle), + prefix_cache::DEFAULT_WEIGHT, + ) + .with_min_share(share), + ) + } } } -/// Compatibility shim used by tests + non-cache-aware code paths. Builds -/// a policy without wiring the cache-aware dependencies; rejects -/// `CacheAwareZmq` to keep the call sites that don't have a `HashTree` / -/// `TokenizerRegistry` to hand from accidentally compiling. -#[cfg(test)] -pub fn build_policy_kind_only(kind: PolicyKind) -> Arc { +/// Builds one soft scoring term for `--policy fused_score`. +fn build_score( + kind: ScoreTermKind, + tree: &Arc, + block_size_oracle: &Arc, +) -> Arc { match kind { + ScoreTermKind::Random => Arc::new(RandomPolicy::new()), + ScoreTermKind::LoadBased => Arc::new(LoadBasedPolicy::new()), + ScoreTermKind::PrefixCache => Arc::new(PrefixCachePolicy::new( + Arc::clone(tree), + Arc::clone(block_size_oracle), + prefix_cache::DEFAULT_WEIGHT, + )), + } +} + +/// Builds `--policy fused_score`. +fn build_fused( + model: &ModelConfig, + tree: &Arc, + oracle: &Arc, +) -> Result> { + let spec = model.fused.as_deref().unwrap_or_default(); + if spec.is_empty() { + return Err(anyhow!( + "--policy fused_score needs at least one --fuse term" + )); + } + let mut terms: Vec<(Arc, Option)> = Vec::with_capacity(spec.len()); + for t in spec { + terms.push((build_score(t.kind, tree, oracle), t.weight)); + } + Ok(Arc::new(FusedScorePolicy::new(terms)?)) +} + +/// Builds a policy with test defaults. +#[cfg(test)] +pub fn build_policy_kind_only(kind: PolicyKind) -> Result> { + Ok(match kind { PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), PolicyKind::Random => Arc::new(RandomPolicy::new()), PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()), - PolicyKind::CacheAwareZmq => { - // Provide an empty tree + empty tokenizer registry + fresh - // oracle so the test policy is constructible. Production - // callers go through `build_policy` with the real - // process-shared instances. - Arc::new(CacheAwareZmqPolicy::new( - crate::config::CacheAwareConfig::default(), - Arc::new(HashTree::new()), - Arc::new(TokenizerRegistry::default()), - BlockSizeOracle::new(), - )) - } + PolicyKind::CacheAwareZmq => Arc::new(CacheAwareZmqPolicy::new( + crate::config::CacheAwareConfig::default(), + Arc::new(HashTree::new()), + Arc::new(TokenizerRegistry::default()), + BlockSizeOracle::new(), + )), PolicyKind::Sticky => { let s = crate::config::StickyConfig::default(); Arc::new(StickyPolicy::new( @@ -108,7 +184,10 @@ pub fn build_policy_kind_only(kind: PolicyKind) -> Arc { build_sticky_fallback(s.fallback_policy), )) } - } + PolicyKind::FusedScore => { + return Err(anyhow!("--policy {kind} needs --fuse terms from the model")) + } + }) } pub fn build_registry( @@ -126,19 +205,12 @@ pub fn build_registry( Arc::clone(&tree), Arc::clone(&tokenizers), Arc::clone(&block_size_oracle), - ), + )?, ); Ok(reg) } -/// Convenience for tests + non-cache-aware callers: builds a registry with -/// a fresh, empty `HashTree` and an empty `TokenizerRegistry`. The -/// cache-aware-zmq policy will then degrade to min-load (no tokenizer + -/// no worker-published block size → fallback) — which is exactly what -/// the legacy tests assume. -/// -/// Production callers go through [`build_registry`] with the real -/// process-shared instances. +/// Builds a registry with empty cache-aware dependencies. pub fn build_registry_with_defaults(cfg: &Config) -> Result { build_registry( cfg, @@ -156,7 +228,90 @@ mod tests { StaticUrlsDiscoveryConfig, }; - use crate::config::PolicyKind; + use crate::config::{ + EligibilityConfig, FilterKind, PolicyKind, ScoreTermKind, StickyFallbackKind, + }; + use crate::discovery::{WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::SelectionContext; + use crate::workers::Worker; + + 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("modelA".into())], + bootstrap_port: None, + })) + } + + #[test] + fn filter_overloaded_wires_through_the_factory() { + let mut cfg = cfg_with_model("modelA", PolicyKind::LoadBased); + cfg.model.eligibility = Some(EligibilityConfig { + filters: vec![FilterKind::Overloaded], + max_in_flight: Some(2), + min_prefix_share: None, + }); + let reg = build_registry_with_defaults(&cfg).unwrap(); + let p = reg.get(&ModelId("modelA".into())).unwrap(); + + let ws = vec![worker("w0"), worker("w1")]; + let model = ModelId("modelA".into()); + let ctx = SelectionContext::new(&model, None); + let _one = ws[1].load_guard(); + assert_eq!( + p.select(&ws, &ctx).unwrap().id, + ws[0].id, + "load still ranks" + ); + + let _fill: Vec<_> = (ws.iter()) + .flat_map(|w| (0..2).map(|_| w.load_guard())) + .collect(); + assert!( + p.select(&ws, &ctx).is_none(), + "every worker is over the cap" + ); + } + + #[test] + fn a_filter_must_actually_constrain() { + let mut cfg = cfg_with_model("modelA", PolicyKind::LoadBased); + let built = |c: &Config| { + build_registry_with_defaults(c).map(|r| r.get(&ModelId("modelA".into())).unwrap()) + }; + assert!( + !built(&cfg).unwrap().can_filter(), + "no floor, so the term stays a pure preference", + ); + + cfg.model.eligibility = Some(EligibilityConfig { + filters: vec![FilterKind::PrefixCache], + max_in_flight: None, + min_prefix_share: Some(0.6), + }); + assert!( + built(&cfg).unwrap().needs_request_tokens(), + "the floor reads the prompt" + ); + + cfg.model.eligibility = Some(EligibilityConfig { + filters: vec![FilterKind::PrefixCache], + max_in_flight: None, + min_prefix_share: Some(0.6), + }); + let filter = build_filter( + FilterKind::PrefixCache, + &cfg.model, + &Arc::new(HashTree::new()), + &BlockSizeOracle::new(), + ); + assert!( + filter.can_filter(), + "a configured prefix-cache floor is a filter", + ); + } fn cfg_with_model(id: &str, policy: PolicyKind) -> Config { Config { @@ -172,6 +327,8 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], @@ -183,13 +340,68 @@ mod tests { #[test] fn build_policy_kind_only_covers_all_variants() { - // Trivially total — the match is exhaustive over `PolicyKind`. - let _ = build_policy_kind_only(PolicyKind::RoundRobin); - let _ = build_policy_kind_only(PolicyKind::Random); - let _ = build_policy_kind_only(PolicyKind::PowerOfTwo); - let _ = build_policy_kind_only(PolicyKind::LoadBased); - let _ = build_policy_kind_only(PolicyKind::CacheAwareZmq); - let _ = build_policy_kind_only(PolicyKind::Sticky); + for kind in [ + PolicyKind::RoundRobin, + PolicyKind::Random, + PolicyKind::PowerOfTwo, + PolicyKind::LoadBased, + PolicyKind::CacheAwareZmq, + PolicyKind::Sticky, + ] { + assert!(build_policy_kind_only(kind).is_ok(), "{kind:?}"); + } + assert!(build_policy_kind_only(PolicyKind::FusedScore).is_err()); + } + + #[test] + fn prefix_cache_builds_as_a_score_term() { + let p = build_score( + ScoreTermKind::PrefixCache, + &Arc::new(HashTree::new()), + &BlockSizeOracle::new(), + ); + assert!(p.can_fuse(), "prefix_cache must be usable as a --fuse term"); + } + + #[test] + fn fused_score_builds_score_terms_and_rejects_an_empty_config() { + let mut cfg = cfg_with_model("m", PolicyKind::FusedScore); + let term = |kind, weight| crate::config::FusedTerm { kind, weight }; + + for weight in [None, Some(2.5)] { + cfg.model.fused = Some(vec![term(ScoreTermKind::PrefixCache, weight)]); + assert!(build_registry_with_defaults(&cfg).is_ok(), "{weight:?}"); + } + + cfg.model.fused = Some(vec![]); + assert!(build_registry_with_defaults(&cfg) + .unwrap_err() + .to_string() + .contains("at least one --fuse term")); + } + + #[test] + fn fused_score_accepts_an_outer_eligibility_pipeline() { + 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: Some(0), + min_prefix_share: None, + }); + + let policy = build_registry_with_defaults(&cfg) + .expect("an outer filter must not make fused terms non-fusable") + .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()); } #[test] @@ -209,9 +421,6 @@ mod tests { let tokenizers = Arc::new(TokenizerRegistry::default()); let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap(); let p = reg.get(&ModelId("modelA".into())).unwrap(); - // Down-cast probe via Debug — cheaper than carrying a type-tag - // on the trait. Pinning the debug repr is fine because the field - // name is part of the file's public test surface. let dbg = format!("{p:?}"); assert!( dbg.contains("CacheAwareZmqPolicy"), @@ -246,4 +455,22 @@ mod tests { "expected StickyPolicy debug repr, got: {dbg}", ); } + + #[test] + fn sticky_fallback_builder_covers_all_typed_choices() { + let workers = vec![worker("w0"), worker("w1")]; + let model = ModelId("modelA".into()); + let ctx = SelectionContext::new(&model, None); + for kind in [ + StickyFallbackKind::RoundRobin, + StickyFallbackKind::Random, + StickyFallbackKind::PowerOfTwo, + StickyFallbackKind::LoadBased, + ] { + assert!( + build_sticky_fallback(kind).select(&workers, &ctx).is_some(), + "{kind:?}" + ); + } + } } diff --git a/experimental/sgl-router/src/policies/kv_events/tree.rs b/experimental/sgl-router/src/policies/kv_events/tree.rs index 2409b6d10..ba1b1d60f 100644 --- a/experimental/sgl-router/src/policies/kv_events/tree.rs +++ b/experimental/sgl-router/src/policies/kv_events/tree.rs @@ -501,6 +501,74 @@ impl TreeState { } } + /// Contiguous prefix depth for every worker on the chain, in one descent. + /// `insert` marks a worker at every node it descends, so "holds the first + /// `d` blocks" means present at each of levels `1..=d`. A worker is frozen + /// at the first level that omits it, so a `remove`d interior node stops the + /// count at the hole instead of counting past it. + fn prefix_depths( + &self, + parent_hash: Option, + block_hashes: &[i64], + ) -> HashMap { + if block_hashes.is_empty() { + return HashMap::new(); + } + // Same start resolution as `match_prefix`: an ambiguous `parent_hash` + // has no worker context to disambiguate, so fall back to root. + let start = match parent_hash { + None => ROOT_ID, + Some(p) => match self.by_hash.get(&p) { + Some(set) if set.len() == 1 => *set.iter().next().unwrap(), + _ => ROOT_ID, + }, + }; + // Keys borrow the arena for the walk; ids are cloned once on the way out. + let mut depths: HashMap<&KvWorkerId, usize> = HashMap::new(); + let mut alive: Vec<&KvWorkerId> = Vec::new(); + let mut current = start; + let mut reached = 0usize; + let now = now_millis(); + for &h in block_hashes { + let Some(child_id) = self + .nodes + .get(¤t) + .and_then(|n| n.children.get(&h).copied()) + else { + break; + }; + let Some(child) = self.nodes.get(&child_id) else { + break; + }; + child.last_used.store(now, Ordering::Relaxed); + current = child_id; + reached += 1; + if reached == 1 { + // Absent here means never in `alive`: a tail held without + // block 0 is reported as holding nothing. + alive = child.workers.iter().collect(); + } else { + let mut still = Vec::with_capacity(alive.len()); + for w in alive { + if child.workers.contains(w) { + still.push(w); + } else { + depths.insert(w, reached - 1); + } + } + alive = still; + } + if alive.is_empty() { + break; + } + } + // Whoever is still tracked held every level the walk reached. + for w in alive { + depths.insert(w, reached); + } + depths.into_iter().map(|(w, d)| (w.clone(), d)).collect() + } + /// Approximate count of *non-root* nodes in the tree. fn node_count(&self) -> usize { // Subtract one for the root sentinel. @@ -652,6 +720,18 @@ impl HashTree { state.match_prefix(parent_hash, block_hashes) } + /// How many leading blocks of `block_hashes` each worker holds contiguously, + /// in one descent under one read lock. [`Self::match_prefix`] names only the + /// deepest matched node's holders, so it cannot answer this. Absent = none. + pub fn prefix_depths( + &self, + parent_hash: Option, + block_hashes: &[i64], + ) -> HashMap { + let state = self.state.read(); + state.prefix_depths(parent_hash, block_hashes) + } + /// Approximate number of non-root nodes in the tree (the root sentinel /// is not counted). Useful for metrics and to decide when to call /// [`HashTree::evict_lru`]. @@ -702,6 +782,36 @@ mod tests { ids.iter().map(|w| (*w).clone()).collect() } + /// One descent answers for every worker at its *own* depth, where + /// `match_prefix` names only whoever sits at the deepest matched node. + /// Also pins that a `remove`d interior block stops the count at the hole. + #[test] + fn prefix_depths_answers_every_worker_at_its_own_depth() { + let chain = [1i64, 2, 3, 4]; + let (deep, shallow, holed) = ( + worker("http://a", 0), + worker("http://b", 0), + worker("http://c", 0), + ); + let tree = HashTree::new(); + tree.insert(&deep, None, &chain); + tree.insert(&shallow, None, &chain[..2]); + tree.insert(&holed, None, &chain); + tree.remove(&holed, &chain[1..2]); + + let depths = tree.prefix_depths(None, &chain); + assert_eq!(depths.get(&deep), Some(&4), "holds the whole chain"); + assert_eq!(depths.get(&shallow), Some(&2), "holds two of four"); + assert_eq!(depths.get(&holed), Some(&1), "stops at the cleared block"); + assert_eq!(depths.get(&worker("http://d", 0)), None, "holds nothing"); + + // Contiguity matters: `remove` left `holed` listed at the deepest node, + // so `match_prefix` credits it with the full chain scored here at 1. + let m = tree.match_prefix(None, &chain); + assert_eq!(m.matched_blocks, 4); + assert_eq!(m.workers, workers(&[&deep, &holed])); + } + #[test] fn empty_match_returns_zero_no_workers() { let tree = HashTree::new(); diff --git a/experimental/sgl-router/src/policies/load_based.rs b/experimental/sgl-router/src/policies/load_based.rs index 6767ad12f..ccf3aeee6 100644 --- a/experimental/sgl-router/src/policies/load_based.rs +++ b/experimental/sgl-router/src/policies/load_based.rs @@ -1,14 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -use crate::policies::{Policy, SelectionContext}; +use crate::policies::scoring::ScoringPolicy; +use crate::policies::SelectionContext; use crate::workers::Worker; use std::sync::Arc; -/// Deterministic load-based policy. -/// -/// Chooses the candidate with the lowest current `Worker::active_load`. -/// Ties follow the candidate slice order, which is registry-dependent. +/// Prefers the least-loaded candidate; `select()` is the blanket impl's. #[derive(Debug, Default)] pub struct LoadBasedPolicy; @@ -16,18 +14,23 @@ impl LoadBasedPolicy { pub fn new() -> Self { Self } - - pub fn pick_min_load(workers: &[Arc]) -> Option> { - workers - .iter() - .min_by_key(|w| w.active_load()) - .map(Arc::clone) - } } -impl Policy for LoadBasedPolicy { - fn select(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Option> { - Self::pick_min_load(workers) +impl ScoringPolicy for LoadBasedPolicy { + /// `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 + /// `TIE_EPSILON`, so the term dies exactly when load matters most. + /// + /// 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(); + 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. + let score = |l: usize| 1.0 - (l - lo) as f32 / span.max(1.0); + loads.into_iter().map(score).collect() } } @@ -35,6 +38,8 @@ impl Policy for LoadBasedPolicy { mod tests { use super::*; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::scoring::argmax::TIE_EPSILON; + use crate::policies::Policy; fn worker(id: &str) -> Arc { Arc::new(Worker::new(WorkerSpec { @@ -46,25 +51,38 @@ mod tests { })) } - #[test] + #[test] // upstream's, retargeted from `pick_min_load` onto blanket `select` fn empty_returns_none() { - let policy = LoadBasedPolicy::new(); - let model = ModelId("tiny".into()); - let ctx = SelectionContext::new(&model, None); - assert!(policy.select(&[], &ctx).is_none()); + let m = ModelId("tiny".into()); + let ctx = SelectionContext::new(&m, None); + assert!(LoadBasedPolicy::new().select(&[], &ctx).is_none()); } + /// `select()` alone CANNOT detect a broken score: ARGMAX breaks a tie on + /// load, so a constant `scores()` still lands on the minimum and that arm + /// passes for the wrong reason. Ranking is therefore asserted on the vector + /// itself, strictly outside `TIE_EPSILON` so a saturating curve cannot hide + /// in the tie band -- what `300,900` is for. NaN needs its own arm because + /// no ORDERING sees it: it makes every comparison false, which on `0,0` is + /// the expected answer. Upstream's `picks_lowest_active_load` goes under + /// rule 4 -- the unique-minimum 2-worker case, which `0,1` subsumes. #[test] - fn picks_lowest_active_load() { - let policy = LoadBasedPolicy::new(); + fn scores_rank_strictly_by_load_and_the_choice_lands_on_the_minimum() { let model = ModelId("tiny".into()); - let ctx = SelectionContext::new(&model, None); - let w0 = worker("w0"); - let w1 = worker("w1"); - let _g0 = w0.load_guard(); - assert_eq!( - policy.select(&[w0, Arc::clone(&w1)], &ctx).unwrap().id, - w1.id - ); + let (ctx, p) = (SelectionContext::new(&model, None), LoadBasedPolicy::new()); + for spec in ["0,1", "2,1,0", "1,0,1", "0,0", "5,2,9,2", "300,900"] { + let loads: Vec = spec.split(',').map(|s| s.parse().unwrap()).collect(); + let ws: Vec> = (0..loads.len()).map(|i| worker(&format!("w{i}"))).collect(); + let _held: Vec<_> = (ws.iter().zip(&loads)) + .flat_map(|(w, n)| (0..*n).map(move |_| w.load_guard())) + .collect(); + let scores = p.scores(&ws, &ctx); + for (i, j) in (0..loads.len()).flat_map(|i| (0..loads.len()).map(move |j| (i, j))) { + let ok = (scores[i] > scores[j] + TIE_EPSILON, scores[i].is_nan()); + assert_eq!(ok, (loads[i] < loads[j], false), "{spec} scored {scores:?}"); + } + let got = p.select(&ws, &ctx).expect("non-empty").active_load(); + assert_eq!(got, *loads.iter().min().expect("non-empty"), "{spec}"); + } } } diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index ec2cd3c9b..6bdda22f7 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -10,55 +10,34 @@ pub mod power_of_two; pub mod random; pub mod registry; pub mod round_robin; +pub mod scoring; pub mod sticky; use crate::discovery::ModelId; +use crate::policies::scoring::{EligibilityFilter, ScoringPolicy}; use crate::server::metrics::MetricsRegistry; use crate::tokenizer::{adapter, TokenizerRegistry}; use crate::workers::Worker; use dashmap::DashMap; use std::sync::Arc; -/// Tokens produced once at ingress for a request. Consumed by the -/// cache-aware selection decision and, when `engine_equivalent`, forwarded -/// to the engine as `input_ids` so the engine skips its own prompt -/// tokenization (the router and engine would otherwise tokenize the same -/// prompt twice in the same cluster). +/// Tokens produced at ingress for routing and optional engine forwarding. pub struct RequestTokens { /// The prompt token ids. pub ids: Vec, - /// True only when the ids were produced via the model's chat encoder — - /// i.e. they match what the engine would tokenize from the chat - /// template. False for the raw-prompt fallback, where the engine must - /// tokenize the text itself, so the ids are NOT safe to forward. + /// Whether the token ids are safe to forward as engine `input_ids`. pub engine_equivalent: bool, } -/// External indexer answer prepared by the async ingress path for the -/// synchronous cache-aware policy. +/// External indexer answer prepared by the async ingress path for a +/// cache-aware policy. pub struct ExternalPrefixSignal { pub outcome: sgl_kv_indexer::PrefixOutcome, pub query_blocks: usize, } -/// Produce the routing tokens — and whether they are engine-equivalent — -/// from an already-parsed request body, using the shared tokenizer registry. -/// -/// Tokenization is a property of the MODEL (does it have a chat encoder?), -/// not of the routing policy, so this lives here as a free function the -/// ingress calls directly with `ctx.tokenizers` — every policy (sticky, -/// round-robin, cache-aware) shares one tokenize. The cache-aware policy also -/// calls it as a body-tokenize fallback for callers that didn't pre-tokenize. -/// -/// Chat requests (`messages`) on a model that has a chat encoder are rendered -/// through that encoder and tokenized the way the engine does, so the query -/// hashes match the engine's cached blocks (chat-templated tokens) AND the ids -/// are safe to hand the engine as `input_ids` (`engine_equivalent = true`). -/// Everything else — `/v1/completions` (`prompt`), `/generate` (`text`), or a -/// chat model without an encoder — tokenizes the raw extracted prompt text; -/// those ids only match the engine after it applies its own template, so they -/// are NOT engine-equivalent. A failed encoder render/encode falls through to -/// the raw path rather than failing the request. +/// Tokenizes a request for routing. Chat-encoder tokens are engine-equivalent; +/// raw prompt tokens are used only for routing. pub fn request_tokens_for( tokenizers: &TokenizerRegistry, model_id: &ModelId, @@ -82,11 +61,7 @@ pub fn request_tokens_for( }) } -/// Tokenize `text` for `model_id` via the shared registry. Returns `None` if -/// no tokenizer is loaded (the model_id may be misconfigured) or if encoding -/// fails / yields no tokens. An encode error logs at WARN (a loaded-but-erroring -/// tokenizer silently disables the offload); the no-text / empty-output paths -/// are expected and stay quiet. +/// Tokenizes text with the model tokenizer. fn tokenize_text( tokenizers: &TokenizerRegistry, model_id: &ModelId, @@ -97,13 +72,6 @@ fn tokenize_text( Ok(ids) if !ids.is_empty() => Some(ids), Ok(_) => None, Err(e) => { - // WARN, not DEBUG: a tokenizer that is loaded but consistently - // erroring silently turns the whole tokenization offload into a - // no-op, so the failure must be visible above DEBUG. Sustained - // failure logs once per request; the volume signal is the - // `sgl_router_ingress_tokenize_errors_total` counter (which the - // chat handler bumps on the chat-encode failure), so no - // rate-limiter here. tracing::warn!( model = %model_id, error = %e, @@ -114,10 +82,7 @@ fn tokenize_text( } } -/// Extract a raw prompt-text candidate from an already-parsed JSON request -/// body. Returns `None` when there's no routable text field; the caller then -/// skips tokenization. This is the raw path — chat requests on a model with a -/// chat encoder are rendered via the encoder instead (see [`request_tokens_for`]). +/// Extracts raw prompt text from a parsed request body. /// /// Supported shapes (in priority order): /// 1. `"prompt": "..."` — `/v1/completions`-style. @@ -129,7 +94,6 @@ fn tokenize_text( /// multimodal content blocks; text-only blocks concatenated. /// 5. `"text": "..."` — SGLang `/generate` native form. /// -/// Anything else yields `None`. pub(crate) fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option { if let Some(s) = v.get("prompt").and_then(|p| p.as_str()) { return Some(s.to_string()); @@ -173,16 +137,7 @@ pub(crate) fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option { model: &'a ModelId, request_body: Option<&'a [u8]>, @@ -216,9 +171,7 @@ impl<'a> SelectionContext<'a> { } } - /// Attach the ingress-computed routing tokens. When present, the - /// cache-aware policy consumes these instead of re-parsing and - /// re-tokenizing the body. + /// Attaches ingress-computed routing tokens. pub fn with_request_tokens(mut self, request_tokens: Option<&'a [u32]>) -> Self { self.request_tokens = request_tokens; self @@ -244,8 +197,7 @@ impl<'a> SelectionContext<'a> { self.routing_key } - /// Ingress-precomputed routing tokens, if any. `None` means the policy - /// must derive tokens itself (e.g. a caller that didn't pre-tokenize). + /// Returns ingress-computed routing tokens. pub fn request_tokens(&self) -> Option<&[u32]> { self.request_tokens } @@ -258,25 +210,33 @@ impl<'a> SelectionContext<'a> { pub trait Policy: Send + Sync + std::fmt::Debug { fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option>; - /// Whether this policy's ROUTING decision needs the 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` (load-only + sticky route - /// without prefix tokens); only the cache-aware policy overrides it. + /// Whether policy selection needs request tokens. fn needs_request_tokens(&self) -> bool { false } - /// Attach the process metrics registry after construction. Default is a - /// no-op — only policies that emit metrics (cache-aware-zmq's - /// `sgl_router_overlap_blocks`) override it. Mirrors - /// `ActiveLoadRegistry::attach_metrics`: the registry is built after the - /// policies, so it is injected here rather than passed to the constructor. + /// Attaches the process metrics registry after construction. fn attach_metrics(&self, _metrics: Arc) {} + + /// Returns the optional per-worker scoring view. + fn as_scoring(&self) -> Option<&dyn ScoringPolicy> { + None + } + + /// Returns the optional per-worker eligibility view. + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + None + } + + /// Whether this policy exposes scores for `--fuse`. + fn can_fuse(&self) -> bool { + self.as_scoring().is_some() + } + + /// Whether this policy exposes an eligibility filter. + fn can_filter(&self) -> bool { + self.as_filter().is_some() + } } #[derive(Debug, Default)] @@ -293,9 +253,7 @@ impl PolicyRegistry { self.by_model.get(model).map(|p| p.clone()) } - /// Inject the metrics registry into every registered policy. Called once - /// at startup (after the registry is built) so metrics-emitting policies - /// can record into the shared registry. + /// Attaches metrics to each registered policy. pub fn attach_metrics(&self, metrics: Arc) { for entry in self.by_model.iter() { entry.value().attach_metrics(Arc::clone(&metrics)); diff --git a/experimental/sgl-router/src/policies/random.rs b/experimental/sgl-router/src/policies/random.rs index 05547429b..3c73b2c89 100644 --- a/experimental/sgl-router/src/policies/random.rs +++ b/experimental/sgl-router/src/policies/random.rs @@ -1,9 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -use crate::policies::{Policy, SelectionContext}; +use crate::policies::{scoring::ScoringPolicy, SelectionContext}; use crate::workers::Worker; -use rand::seq::SliceRandom; +use rand::Rng; use std::sync::Arc; #[derive(Debug, Default)] @@ -15,8 +15,52 @@ impl RandomPolicy { } } -impl Policy for RandomPolicy { - fn select(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Option> { - workers.choose(&mut rand::thread_rng()).cloned() +impl ScoringPolicy for RandomPolicy { + /// Argmax of n iid uniforms IS a uniform choice: exactly the old `choose`. + /// Never constrains: a coin toss is not an eligibility rule. + fn scores(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Vec { + let mut rng = rand::thread_rng(); + (0..workers.len()).map(|_| rng.gen()).collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::Policy; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + /// Distributional: `select()` is not pure. Marginals alone are satisfied by + /// a ROTATION -- what a constant `scores()` becomes under ARGMAX's rotating + /// tiebreak -- so REPEATS share the band: P(pick==prev) is 1/N iid, 0 rotating. + #[test] + fn picks_uniformly_over_20k_draws_and_repeats_at_the_iid_rate() { + const MEAN: f64 = 5_000.0; // 20_000 draws over 4 workers; repeats too + const BAND: f64 = 5.0 * 61.237_244; // 5 sigma, sqrt(20_000 / 4 * 3 / 4) + let (policy, model) = (RandomPolicy::new(), ModelId("tiny".into())); + let ctx = SelectionContext::new(&model, None); + let ws: Vec> = (0..4).map(|i| worker(&format!("w{i}"))).collect(); + assert!(policy.select(&[], &ctx).is_none(), "empty fleet"); + let (mut counts, mut repeats, mut prev) = ([0usize; 4], 0usize, None); + for _ in 0..20_000 { + let got = policy.select(&ws, &ctx).expect("non-empty fleet"); + let i = ws.iter().position(|w| w.id == got.id).expect("a candidate"); + counts[i] += 1; + repeats += usize::from(prev.replace(i) == Some(i)); + } + for n in counts.iter().chain([&repeats]) { + let dev = (*n as f64 - MEAN).abs(); // `> 0` is no-starvation + assert!(*n > 0 && dev < BAND, "{counts:?} {repeats}"); + } } } diff --git a/experimental/sgl-router/src/policies/scoring/admission.rs b/experimental/sgl-router/src/policies/scoring/admission.rs new file mode 100644 index 000000000..de492d941 --- /dev/null +++ b/experimental/sgl-router/src/policies/scoring/admission.rs @@ -0,0 +1,123 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Router-local in-flight admission control. + +use super::{EligibilityFilter, OnEmpty}; +use crate::policies::{Policy, SelectionContext}; +use crate::workers::Worker; +use std::sync::Arc; + +/// Rejects workers at the router-local in-flight limit. +#[derive(Debug)] +pub struct Overloaded { + max_in_flight: usize, +} + +impl Overloaded { + pub fn new(max_in_flight: usize) -> Self { + Self { max_in_flight } + } +} + +impl EligibilityFilter for Overloaded { + fn keep(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Vec { + (workers.iter()) + .map(|w| w.active_load() < self.max_in_flight) + .collect() + } + + /// Do not route to an over-capacity worker. + fn on_empty(&self) -> OnEmpty { + OnEmpty::Hold + } +} + +impl Policy for Overloaded { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + let eligible: Vec> = (workers.iter()) + .zip(self.keep(workers, ctx)) + .filter(|(_, ok)| *ok) + .map(|(w, _)| Arc::clone(w)) + .collect(); + eligible + .iter() + .min_by_key(|w| w.active_load()) + .map(Arc::clone) + } + + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + Some(self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::scoring::{admit, refs}; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + #[test] + fn the_cap_is_a_strict_ceiling() { + 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(); + + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + assert_eq!( + Overloaded::new(3).keep(&ws, &ctx), + vec![true, true, false], + "load 3 against a cap of 3 is over", + ); + } + + #[test] + fn a_full_fleet_refuses_rather_than_picking_the_least_bad() { + let ws = vec![worker("a"), worker("b")]; + let _a: Vec<_> = (0..5).map(|_| ws[0].load_guard()).collect(); + let _b: Vec<_> = (0..9).map(|_| ws[1].load_guard()).collect(); + + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + let full: Vec> = vec![Box::new(Overloaded::new(4))]; + assert!( + admit(refs(&full), &ws, &ctx).is_none(), + "both over the cap, and the filter Holds", + ); + + let some: Vec> = vec![Box::new(Overloaded::new(6))]; + let out = admit(refs(&some), &ws, &ctx).expect("a is under the cap"); + assert_eq!(out.len(), 1); + assert_eq!(out[0].url, ws[0].url); + } + + #[test] + fn hold_does_not_yield_to_a_later_filter() { + #[derive(Debug)] + struct AdmitAll; + impl EligibilityFilter for AdmitAll { + fn keep(&self, ws: &[Arc], _: &SelectionContext<'_>) -> Vec { + vec![true; ws.len()] + } + } + let ws = vec![worker("a")]; + let _busy: Vec<_> = (0..9).map(|_| ws[0].load_guard()).collect(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let chain: Vec> = + vec![Box::new(Overloaded::new(2)), Box::new(AdmitAll)]; + assert!(admit(refs(&chain), &ws, &ctx).is_none()); + } +} diff --git a/experimental/sgl-router/src/policies/scoring/argmax.rs b/experimental/sgl-router/src/policies/scoring/argmax.rs new file mode 100644 index 000000000..1d1fa4ee5 --- /dev/null +++ b/experimental/sgl-router/src/policies/scoring/argmax.rs @@ -0,0 +1,113 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Selects one worker from per-worker scores. + +use crate::workers::Worker; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +/// Scores within this distance of the best are tied. +pub const TIE_EPSILON: f32 = 1e-6; + +pub trait Selector: Send + Sync + std::fmt::Debug { + /// Index into `workers` of the chosen candidate, or `None` when there is + /// nothing to choose from. `scores[i]` belongs to `workers[i]`. + fn pick(&self, workers: &[Arc], scores: &[f32]) -> Option; +} + +/// Highest score wins; ties choose the least-loaded candidate and rotate. +#[derive(Debug, Default)] +pub struct Argmax { + rotor: AtomicUsize, +} + +/// The default selector, shared by every scoring policy that does not override +/// [`super::ScoringPolicy::selector`]. +pub static ARGMAX: Argmax = Argmax { + rotor: AtomicUsize::new(0), +}; + +impl Selector for Argmax { + fn pick(&self, workers: &[Arc], scores: &[f32]) -> Option { + if workers.is_empty() { + return None; + } + let n = workers.len().min(scores.len()); + let best = (0..n) + .map(|i| scores[i]) + .filter(|s| !s.is_nan()) + .fold(None::, |acc, s| Some(acc.map_or(s, |b| b.max(s)))); + let mut band: Vec = match best { + Some(b) => (0..n) + .filter(|&i| !scores[i].is_nan() && scores[i] >= b - TIE_EPSILON) + .collect(), + None => Vec::new(), + }; + if band.is_empty() { + tracing::debug!( + n_workers = workers.len(), + n_scores = scores.len(), + "no usable score; falling back to load + rotation", + ); + band = (0..workers.len()).collect(); + } + let min_load = band.iter().map(|&i| workers[i].active_load()).min()?; + let tied: Vec = band + .into_iter() + .filter(|&i| workers[i].active_load() == min_load) + .collect(); + let k = self.rotor.fetch_add(1, Ordering::Relaxed) % tied.len(); + Some(tied[k]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use std::collections::HashSet; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + #[test] + fn score_wins_unless_the_gap_is_inside_the_tie_band() { + let ws = vec![worker("a"), worker("b")]; + let sel = Argmax::default(); + let _loaded = ws[1].load_guard(); + + assert_eq!(sel.pick(&ws, &[1.0, 1.0 - 1e-3]), Some(0), "clear winner"); + let tie = [1.0 - 5e-7, 1.0]; + assert_eq!(sel.pick(&ws, &tie), Some(0), "tie -> less load"); + assert_eq!(sel.pick(&[], &[]), None, "nothing to choose from"); + } + + #[test] + fn nan_never_wins_from_either_position() { + let ws = vec![worker("a"), worker("b")]; + let sel = Argmax::default(); + assert_eq!(sel.pick(&ws, &[f32::NAN, 0.0]), Some(1)); + assert_eq!(sel.pick(&ws, &[0.0, f32::NAN]), Some(0)); + assert!(sel.pick(&ws, &[f32::NAN, f32::NAN]).is_some()); + } + + #[test] + fn a_total_tie_rotates_over_every_candidate() { + let ws = vec![worker("a"), worker("b"), worker("c")]; + let sel = Argmax::default(); + let picks: HashSet = (0..3).filter_map(|_| sel.pick(&ws, &[1.0; 3])).collect(); + assert_eq!( + picks.len(), + 3, + "three tied picks must cover all three: {picks:?}" + ); + } +} diff --git a/experimental/sgl-router/src/policies/scoring/mod.rs b/experimental/sgl-router/src/policies/scoring/mod.rs new file mode 100644 index 000000000..31fbc9ddb --- /dev/null +++ b/experimental/sgl-router/src/policies/scoring/mod.rs @@ -0,0 +1,579 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Candidate eligibility and scoring policies. + +pub mod admission; +pub mod argmax; +pub mod prefix_cache; + +use crate::policies::{Policy, SelectionContext}; +use crate::workers::Worker; +use argmax::{Selector, ARGMAX}; +use std::sync::Arc; + +/// What a filter means when it has rejected every worker it was shown. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OnEmpty { + /// Ignore this filter when it rejects every candidate. + Abstain, + /// Keep the rejection: no worker is admissible. + Hold, +} + +/// A hard constraint applied before scoring. +pub trait EligibilityFilter: Send + Sync + std::fmt::Debug { + /// Returns one admission flag per worker; `true` keeps the candidate. + fn keep(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec; + + /// Whether this constraint reads `ctx.request_tokens()`. + fn needs_tokens(&self) -> bool { + false + } + + /// Controls the result when this filter rejects every candidate. + fn on_empty(&self) -> OnEmpty { + OnEmpty::Abstain + } +} + +/// A soft preference for eligible candidates. +pub trait ScoringPolicy: Send + Sync + std::fmt::Debug { + /// What each candidate is worth, parallel to `workers`; higher is better. + fn scores(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec; + + /// Default multiplier as a fused term; `--fuse name=weight` overrides it. + fn weight(&self) -> f32 { + 1.0 + } + + /// Whether scoring needs request tokens. + fn needs_tokens(&self) -> bool { + false + } + + /// Optional eligibility view for policies that provide both signals. + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + None + } + + /// Selects a winner from the score vector. + fn selector(&self) -> &dyn Selector { + &ARGMAX + } +} + +/// Applies ordered filters. `None` means a holding filter rejected all candidates. +pub fn admit<'f>( + filters: impl IntoIterator, + workers: &[Arc], + ctx: &SelectionContext<'_>, +) -> Option>> { + let mut alive: Vec> = workers.to_vec(); + + for filter in filters { + if alive.is_empty() { + break; + } + let flags = filter.keep(&alive, ctx); + let on_empty = filter.on_empty(); + if flags.len() != alive.len() { + tracing::debug!( + filter = ?filter, + n_workers = alive.len(), + n_flags = flags.len(), + "eligibility filter returned the wrong arity", + ); + if on_empty == OnEmpty::Hold { + return None; + } + } + let untouched = alive.len() == workers.len(); + + let next: Vec> = (alive.iter().enumerate()) + .filter(|(i, _)| flags.get(*i).copied().unwrap_or(true)) + .map(|(_, w)| Arc::clone(w)) + .collect(); + + if next.is_empty() { + match on_empty { + OnEmpty::Hold => return None, + OnEmpty::Abstain if untouched => { + tracing::warn!( + 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", + ); + continue; + } + OnEmpty::Abstain => { + tracing::debug!( + filter = ?filter, + n_alive = alive.len(), + "eligibility filter conflicts with a higher-priority one; yielding", + ); + continue; + } + } + } + alive = next; + } + + Some(alive) +} + +/// Selects the best-scoring worker. +impl Policy for T { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + let scores = self.scores(workers, ctx); + let i = self.selector().pick(workers, &scores)?; + workers.get(i).map(Arc::clone) + } + + fn needs_request_tokens(&self) -> bool { + ScoringPolicy::needs_tokens(self) || self.as_filter().is_some_and(|f| f.needs_tokens()) + } + + fn as_scoring(&self) -> Option<&dyn ScoringPolicy> { + Some(self) + } + + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + ScoringPolicy::as_filter(self) + } +} + +/// A weighted sum of scoring terms. +#[derive(Debug)] +pub struct FusedScorePolicy { + /// Terms and optional `--fuse name=weight` overrides. + terms: Vec<(Arc, Option)>, +} + +fn view(t: &(Arc, Option)) -> (&dyn ScoringPolicy, f32) { + let s = t.0.as_scoring().expect("checked by FusedScorePolicy::new"); + (s, t.1.unwrap_or_else(|| s.weight())) +} + +impl FusedScorePolicy { + /// Reject non-scoring terms during construction. + pub fn new(terms: Vec<(Arc, Option)>) -> anyhow::Result { + for (p, _) in &terms { + anyhow::ensure!(p.can_fuse(), "policy {p:?} does not support fusion"); + } + Ok(Self { terms }) + } +} + +/// Applies eligibility filters before an inner policy. +#[derive(Debug)] +pub struct Pipeline { + filters: Vec>, + inner: Arc, +} + +impl Pipeline { + /// Reject policies that do not expose an eligibility view. + pub fn new(filters: Vec>, inner: Arc) -> anyhow::Result { + for f in &filters { + anyhow::ensure!(f.can_filter(), "policy {f:?} imposes no eligibility rule"); + } + Ok(Self { filters, inner }) + } + + fn views(&self) -> impl Iterator { + (self.filters.iter()).map(|p| p.as_filter().expect("checked by Pipeline::new")) + } +} + +impl Policy for Pipeline { + fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + let eligible = admit(self.views(), workers, ctx)?; + self.inner.select(&eligible, ctx) + } + + fn needs_request_tokens(&self) -> bool { + self.inner.needs_request_tokens() || self.views().any(|f| f.needs_tokens()) + } + + fn attach_metrics(&self, metrics: Arc) { + self.inner.attach_metrics(metrics); + } +} + +impl ScoringPolicy for FusedScorePolicy { + fn scores(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec { + let mut total = vec![0.0f32; workers.len()]; + for (term, w) in self.terms.iter().map(view) { + for (acc, s) in total.iter_mut().zip(term.scores(workers, ctx)) { + *acc += w * s; + } + } + total + } + + fn needs_tokens(&self) -> bool { + self.terms.iter().map(view).any(|(t, _)| t.needs_tokens()) + } +} + +/// Owned boxes as the borrowed views [`admit`] consumes. Shared by the tests +/// in this module and its siblings. +#[cfg(test)] +pub(crate) fn refs( + fs: &[Box], +) -> impl Iterator { + fs.iter().map(|f| &**f) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::round_robin::RoundRobinPolicy; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + fn fleet() -> Vec> { + vec![worker("a"), worker("b"), worker("c")] + } + + fn urls(ws: &[Arc]) -> Vec { + ws.iter().map(|w| w.url.clone()).collect() + } + + fn term(p: impl Policy + 'static, w: Option) -> (Arc, Option) { + (Arc::new(p), w) + } + + #[derive(Debug)] + struct ByIndex(f32, bool, bool); + + fn by(w: f32) -> ByIndex { + ByIndex(w, false, false) + } + + impl ScoringPolicy for ByIndex { + fn scores(&self, workers: &[Arc], _: &SelectionContext<'_>) -> Vec { + (0..workers.len()).map(|i| i as f32).collect() + } + fn weight(&self) -> f32 { + self.0 + } + fn needs_tokens(&self) -> bool { + self.2 + } + fn selector(&self) -> &dyn Selector { + if self.1 { + &PICK_FIRST + } else { + &ARGMAX + } + } + } + + #[derive(Debug)] + struct Keep(Vec<&'static str>, OnEmpty); + + impl EligibilityFilter for Keep { + fn keep(&self, workers: &[Arc], _: &SelectionContext<'_>) -> Vec { + (workers.iter()) + .map(|w| self.0.iter().any(|n| w.url.contains(n))) + .collect() + } + fn on_empty(&self) -> OnEmpty { + self.1 + } + } + + impl Policy for Keep { + fn select(&self, ws: &[Arc], ctx: &SelectionContext<'_>) -> Option> { + (ws.iter().zip(self.keep(ws, ctx))) + .find(|(_, ok)| *ok) + .map(|(w, _)| Arc::clone(w)) + } + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + Some(self) + } + } + + fn keep(names: &[&'static str], on_empty: OnEmpty) -> Box { + Box::new(Keep(names.to_vec(), on_empty)) + } + + fn boxed(f: impl EligibilityFilter + 'static) -> Box { + Box::new(f) + } + + #[derive(Debug)] + struct PickFirst; + static PICK_FIRST: PickFirst = PickFirst; + impl Selector for PickFirst { + fn pick(&self, workers: &[Arc], _: &[f32]) -> Option { + (!workers.is_empty()).then_some(0) + } + } + + #[test] + fn selector_dispatch_uses_the_policys_own_selector() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let scores = by(1.0).scores(&ws, &ctx); + assert_eq!(ARGMAX.pick(&ws, &scores), Some(2)); + assert_eq!(PICK_FIRST.pick(&ws, &scores), Some(0)); + + assert_eq!(by(1.0).select(&ws, &ctx).unwrap().id, ws[2].id); + let first = ByIndex(1.0, true, false); + assert_eq!(first.select(&ws, &ctx).unwrap().id, ws[0].id); + } + + #[test] + fn can_fuse_is_derived_and_gates_construction() { + let fused = FusedScorePolicy::new(vec![term(by(1.0), None)]).unwrap(); + let fusable: Vec> = vec![Arc::new(by(1.0)), Arc::new(fused)]; + for p in &fusable { + assert!(p.can_fuse()); + assert!(p.as_scoring().is_some(), "the flag agrees with the view"); + } + let rr: Arc = Arc::new(RoundRobinPolicy::new()); + assert!(!rr.can_fuse()); + assert!(rr.as_scoring().is_none()); + let err = FusedScorePolicy::new(vec![term(RoundRobinPolicy::new(), None)]) + .expect_err("round_robin has no per-worker preference to contribute"); + assert!(err.to_string().contains("does not support fusion"), "{err}"); + } + + #[test] + fn fusion_nests_and_the_override_replaces_the_terms_own_weight() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let inner = FusedScorePolicy::new(vec![term(by(2.0), None)]).unwrap(); + assert_eq!(inner.scores(&ws, &ctx), vec![0.0, 2.0, 4.0], "its own 2i"); + + let outer = + FusedScorePolicy::new(vec![term(inner, None), term(by(3.0), Some(10.0))]).unwrap(); + assert_eq!(outer.scores(&ws, &ctx), vec![0.0, 12.0, 24.0], "2i + 10i"); + assert_eq!(outer.select(&ws, &ctx).unwrap().id, ws[2].id); + } + + #[test] + fn composer_propagates_needs_tokens_from_any_term() { + let plain = FusedScorePolicy::new(vec![term(by(1.0), None)]).unwrap(); + assert!(!plain.needs_request_tokens()); + let hungry = FusedScorePolicy::new(vec![ + term(by(1.0), None), + term(ByIndex(1.0, false, true), None), + ]) + .unwrap(); + assert!(hungry.needs_request_tokens(), "any one term is enough"); + + #[derive(Debug)] + struct Hungry; + impl EligibilityFilter for Hungry { + fn keep(&self, ws: &[Arc], _: &SelectionContext<'_>) -> Vec { + vec![true; ws.len()] + } + fn needs_tokens(&self) -> bool { + true + } + } + impl Policy for Hungry { + fn select(&self, ws: &[Arc], _: &SelectionContext<'_>) -> Option> { + ws.first().map(Arc::clone) + } + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + Some(self) + } + } + let filtered = Pipeline::new( + vec![Arc::new(Hungry)], + Arc::new(FusedScorePolicy::new(vec![term(by(1.0), None)]).unwrap()), + ) + .unwrap(); + assert!(filtered.needs_request_tokens(), "the filter is hungry"); + } + + #[test] + fn a_rejected_worker_cannot_be_out_weighed() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let fused = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b"], OnEmpty::Abstain))], + Arc::new(FusedScorePolicy::new(vec![term(by(1.0), Some(1e9))]).unwrap()), + ) + .unwrap(); + assert_eq!( + fused.select(&ws, &ctx).unwrap().url, + ws[1].url, + "the best ELIGIBLE, not the best" + ); + + let open = Pipeline::new( + vec![Arc::new(Keep(vec!["a", "b", "c"], OnEmpty::Abstain))], + Arc::new(FusedScorePolicy::new(vec![term(by(1.0), Some(1e9))]).unwrap()), + ) + .unwrap(); + assert_eq!(open.select(&ws, &ctx).unwrap().url, ws[2].url); + } + + #[test] + fn a_conflict_yields_the_later_filter_and_keeps_the_earlier_narrowing() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let chain = vec![ + keep(&["a", "b"], OnEmpty::Abstain), + keep(&["c"], OnEmpty::Abstain), + ]; + let out = admit(refs(&chain), &ws, &ctx).expect("Abstain never holds"); + assert_eq!( + urls(&out), + urls(&ws[..2]), + "the second filter yields; falling back to the raw fleet would read [a, b, c]", + ); + + let rev = vec![ + keep(&["c"], OnEmpty::Abstain), + keep(&["a", "b"], OnEmpty::Abstain), + ]; + assert_eq!(urls(&admit(refs(&rev), &ws, &ctx).unwrap()), urls(&ws[2..])); + } + + #[test] + fn a_filter_after_a_conflict_still_applies() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let chain = vec![ + keep(&["a", "b"], OnEmpty::Abstain), + keep(&["c"], OnEmpty::Abstain), + keep(&["b", "c"], OnEmpty::Abstain), + ]; + assert_eq!( + urls(&admit(refs(&chain), &ws, &ctx).unwrap()), + vec![ws[1].url.clone()] + ); + } + + #[test] + fn a_holding_filter_refuses_instead_of_yielding() { + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + + let held = vec![ + keep(&["a", "b"], OnEmpty::Abstain), + keep(&["c"], OnEmpty::Hold), + ]; + assert!( + admit(refs(&held), &ws, &ctx).is_none(), + "no eligible worker, and the filter said Hold", + ); + + let ok = vec![ + keep(&["a", "b"], OnEmpty::Abstain), + keep(&["b"], OnEmpty::Hold), + ]; + assert_eq!( + urls(&admit(refs(&ok), &ws, &ctx).unwrap()), + vec![ws[1].url.clone()] + ); + + let fused = Pipeline::new( + vec![ + Arc::new(Keep(vec!["a", "b"], OnEmpty::Abstain)), + Arc::new(Keep(vec!["c"], OnEmpty::Hold)), + ], + Arc::new(FusedScorePolicy::new(vec![term(by(1.0), None)]).unwrap()), + ) + .unwrap(); + assert!(fused.select(&ws, &ctx).is_none()); + } + + #[test] + fn a_short_flag_vector_degrades_instead_of_panicking() { + #[derive(Debug)] + struct Short; + impl EligibilityFilter for Short { + fn keep(&self, _: &[Arc], _: &SelectionContext<'_>) -> Vec { + vec![false] + } + } + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + let out = admit(refs(&[boxed(Short)]), &ws, &ctx).expect("the tail was admitted"); + assert_eq!(urls(&out), urls(&ws[1..]), "only index 0 rejected"); + } + + #[test] + fn a_short_hold_filter_fails_closed() { + #[derive(Debug)] + struct ShortHold; + impl EligibilityFilter for ShortHold { + fn keep(&self, _: &[Arc], _: &SelectionContext<'_>) -> Vec { + vec![true] + } + + fn on_empty(&self) -> OnEmpty { + OnEmpty::Hold + } + } + + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + assert!(admit(refs(&[boxed(ShortHold)]), &ws, &ctx).is_none()); + } + + #[test] + fn a_dual_role_term_exposes_its_filter_half_through_policy() { + #[derive(Debug)] + struct Dual; + impl EligibilityFilter for Dual { + fn keep(&self, ws: &[Arc], _: &SelectionContext<'_>) -> Vec { + ws.iter().map(|w| !w.url.contains('c')).collect() + } + fn needs_tokens(&self) -> bool { + true + } + } + impl ScoringPolicy for Dual { + fn scores(&self, ws: &[Arc], _: &SelectionContext<'_>) -> Vec { + vec![0.0; ws.len()] + } + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + Some(self) + } + } + + let p: Arc = Arc::new(Dual); + assert!(p.can_fuse(), "it still scores"); + let f = p.as_filter().expect("and it filters"); + let ws = fleet(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + assert_eq!(f.keep(&ws, &ctx), vec![true, true, false]); + assert!( + p.needs_request_tokens(), + "hunger comes from the filter half" + ); + } +} diff --git a/experimental/sgl-router/src/policies/scoring/prefix_cache.rs b/experimental/sgl-router/src/policies/scoring/prefix_cache.rs new file mode 100644 index 000000000..326802277 --- /dev/null +++ b/experimental/sgl-router/src/policies/scoring/prefix_cache.rs @@ -0,0 +1,264 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Prefix-cache scores from the KV-event [`HashTree`]. + +use super::{EligibilityFilter, ScoringPolicy}; +use crate::policies::kv_events::{ + compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree, +}; +use crate::policies::SelectionContext; +use crate::workers::Worker; +use std::sync::Arc; + +/// Score for a miss or unavailable prefix signal. +const NO_HOLDING: f32 = 0.0; + +/// Default fused-term weight. +pub const DEFAULT_WEIGHT: f32 = 1.0; + +pub struct PrefixCachePolicy { + tree: Arc, + block_size_oracle: Arc, + weight: f32, + /// Minimum cached share for eligibility; zero disables filtering. + min_share: f32, +} + +impl std::fmt::Debug for PrefixCachePolicy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrefixCachePolicy") + .field("weight", &self.weight) + .field("min_share", &self.min_share) + .field("tree_nodes", &self.tree.node_count()) + .finish() + } +} + +impl PrefixCachePolicy { + pub fn new(tree: Arc, block_size_oracle: Arc, weight: f32) -> Self { + Self { + tree, + block_size_oracle, + weight, + min_share: 0.0, + } + } + + /// Require a cached share for eligibility. + pub fn with_min_share(mut self, share: f32) -> Self { + self.min_share = share; + self + } +} + +impl PrefixCachePolicy { + /// Returns each worker's cached prompt share. + fn shares(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec { + let flat = || vec![NO_HOLDING; workers.len()]; + + let Some(tokens) = ctx.request_tokens().filter(|t| !t.is_empty()) else { + return flat(); + }; + let Some(block_size) = self.block_size_oracle.get() else { + return flat(); + }; + let hashes = if self.block_size_oracle.is_bigram() { + compute_block_hashes_bigram(tokens, block_size as usize) + } else { + compute_block_hashes(tokens, block_size as usize) + }; + if hashes.is_empty() { + return flat(); + } + + let depths = self.tree.prefix_depths(None, &hashes); + let total = hashes.len() as f32; + workers + .iter() + .map(|w| { + depths + .iter() + .filter(|(kw, _)| kw.url == w.url) + .map(|(_, &d)| d) + .max() + .map_or(NO_HOLDING, |d| d as f32 / total) + }) + .collect() + } +} + +impl ScoringPolicy for PrefixCachePolicy { + fn scores(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec { + self.shares(workers, ctx) + } + + fn weight(&self) -> f32 { + self.weight + } + + fn as_filter(&self) -> Option<&dyn EligibilityFilter> { + (self.min_share > 0.0).then_some(self as &dyn EligibilityFilter) + } + + fn needs_tokens(&self) -> bool { + true + } +} + +impl EligibilityFilter for PrefixCachePolicy { + fn keep(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Vec { + (self.shares(workers, ctx).into_iter()) + .map(|share| share >= self.min_share) + .collect() + } + + fn needs_tokens(&self) -> bool { + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::kv_events::KvWorkerId; + + const BLOCK: usize = 4; + + fn worker(url: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(url.into()), + url: url.into(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + fn tokens() -> Vec { + (0..(BLOCK as u32 * 4)).collect() + } + + fn insert(tree: &HashTree, url: &str, rank: u32, from: usize, blocks: usize) { + let all = compute_block_hashes(&tokens(), BLOCK); + let parent = if from == 0 { None } else { Some(all[from - 1]) }; + tree.insert( + &KvWorkerId::new(url.into(), rank), + parent, + &all[from..from + blocks], + ); + } + + fn policy(tree: Arc) -> PrefixCachePolicy { + let oracle = BlockSizeOracle::new(); + oracle + .try_set(BLOCK as u32) + .expect("a fresh oracle accepts the first block size"); + PrefixCachePolicy::new(tree, oracle, 1.0) + } + + fn shares(p: &PrefixCachePolicy, ws: &[Arc], ctx: &SelectionContext<'_>) -> Vec { + assert!( + ScoringPolicy::as_filter(p).is_none(), + "no floor configured, so this term must not be a filter at all", + ); + p.scores(ws, ctx) + } + + #[test] + fn depth_is_a_fraction_and_a_tail_without_block_zero_misses() { + let tree = Arc::new(HashTree::new()); + insert(&tree, "deep", 0, 0, 3); + insert(&tree, "tail", 0, 2, 2); + let ws = vec![worker("deep"), worker("tail"), worker("cold")]; + + let model = ModelId("tiny".into()); + let ids = tokens(); + let ctx = SelectionContext::new(&model, None).with_request_tokens(Some(&ids)); + + let scores = shares(&policy(tree), &ws, &ctx); + assert_eq!(scores[0], 0.75, "3 of 4 blocks held, not a neutral 1.0"); + assert_eq!(scores[1], 0.0, "tail without block 0 holds nothing"); + assert_eq!(scores[2], 0.0, "never seen"); + } + + #[test] + fn several_dp_ranks_of_one_worker_collapse_to_the_deepest() { + let tree = Arc::new(HashTree::new()); + insert(&tree, "dp", 0, 0, 1); + insert(&tree, "dp", 1, 0, 3); + let ws = vec![worker("dp")]; + + let model = ModelId("tiny".into()); + let ids = tokens(); + let ctx = SelectionContext::new(&model, None).with_request_tokens(Some(&ids)); + assert_eq!( + shares(&policy(tree), &ws, &ctx), + vec![0.75], + "3 of 4, not 1" + ); + } + + #[test] + fn without_tokens_every_worker_scores_the_same() { + let tree = Arc::new(HashTree::new()); + insert(&tree, "deep", 0, 0, 3); + let ws = vec![worker("deep"), worker("cold")]; + let model = ModelId("tiny".into()); + let policy = policy(tree); + + let ids = tokens(); + let with = SelectionContext::new(&model, None).with_request_tokens(Some(&ids)); + assert_eq!( + shares(&policy, &ws, &with), + vec![0.75, 0.0], + "signal is live" + ); + + let without = SelectionContext::new(&model, None); + assert_eq!( + shares(&policy, &ws, &without), + vec![0.0, 0.0], + "and inert here" + ); + } + + #[test] + fn without_a_block_size_no_worker_looks_like_a_hit() { + let tree = Arc::new(HashTree::new()); + insert(&tree, "deep", 0, 0, 3); + let ws = vec![worker("deep"), worker("cold")]; + + let model = ModelId("tiny".into()); + let ids = tokens(); + let ctx = SelectionContext::new(&model, None).with_request_tokens(Some(&ids)); + let cold = PrefixCachePolicy::new(tree, BlockSizeOracle::new(), 1.0); + assert_eq!(shares(&cold, &ws, &ctx), vec![0.0, 0.0]); + } + + #[test] + fn the_bigram_branch_queries_a_different_chain() { + let ids = tokens(); + let unigram = compute_block_hashes(&ids, BLOCK); + let bigram = compute_block_hashes_bigram(&ids, BLOCK); + assert_ne!(unigram, bigram, "the two hashers must disagree, else this"); + + let tree = Arc::new(HashTree::new()); + insert(&tree, "deep", 0, 0, 4); + let ws = vec![worker("deep")]; + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_request_tokens(Some(&ids)); + + let oracle = BlockSizeOracle::new(); + oracle.try_set(BLOCK as u32).unwrap(); + oracle.set_bigram(true); + let p = PrefixCachePolicy::new(Arc::clone(&tree), oracle, 1.0); + assert_eq!( + shares(&p, &ws, &ctx), + vec![0.0], + "bigram query, unigram tree" + ); + assert_eq!(shares(&policy(tree), &ws, &ctx), vec![1.0]); + } +} diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index a655913f0..5645f35b0 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -116,6 +116,8 @@ impl AppContext { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/server/routes/models.rs b/experimental/sgl-router/src/server/routes/models.rs index 10b635efc..51127f843 100644 --- a/experimental/sgl-router/src/server/routes/models.rs +++ b/experimental/sgl-router/src/server/routes/models.rs @@ -54,6 +54,8 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }; let app = crate::server::app::build_router(std::sync::Arc::new(ctx)); let res = app diff --git a/experimental/sgl-router/src/server/routes/tokenize.rs b/experimental/sgl-router/src/server/routes/tokenize.rs index f45f7bea4..5e05b76bb 100644 --- a/experimental/sgl-router/src/server/routes/tokenize.rs +++ b/experimental/sgl-router/src/server/routes/tokenize.rs @@ -121,6 +121,8 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/tokenizer/mod.rs b/experimental/sgl-router/src/tokenizer/mod.rs index 8df65fd57..9b600a24c 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -235,6 +235,8 @@ mod tests { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: crate::config::DiscoveryBackend::StaticUrls( crate::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index 71d8cfd49..e67769a7a 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -481,6 +481,8 @@ mod tests { }), cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://test:30000".into()], diff --git a/experimental/sgl-router/tests/component/discovery/static_urls.rs b/experimental/sgl-router/tests/component/discovery/static_urls.rs index 3364c98e5..ecdf86259 100644 --- a/experimental/sgl-router/tests/component/discovery/static_urls.rs +++ b/experimental/sgl-router/tests/component/discovery/static_urls.rs @@ -133,6 +133,8 @@ async fn static_urls_pd_role_resolved_end_to_end() { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec![url.clone()], diff --git a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs index 2f4e7e6fe..ae99b7297 100644 --- a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs @@ -72,6 +72,8 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: sgl_router::config::DiscoveryBackend::StaticUrls( sgl_router::config::StaticUrlsDiscoveryConfig { diff --git a/experimental/sgl-router/tests/component/policies/fused_score.rs b/experimental/sgl-router/tests/component/policies/fused_score.rs new file mode 100644 index 000000000..7a71e865d --- /dev/null +++ b/experimental/sgl-router/tests/component/policies/fused_score.rs @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Two REAL scoring policies whose terms DISAGREE, composed and routed. The +//! in-crate fusion tests sum `ByIndex` stubs that rank the fleet the SAME way, +//! so their `select()` half lands on ws[2] whichever term you read; this one's +//! half discriminates. (Their `scores()` half does catch a dropped term — +//! verified by mutation, so this file does not claim otherwise.) +//! +//! NOT pinned here: how `load_based` scales load — W2's min-max scale-free +//! defect is unruled, and both candidate curves put the busiest worker at 0.0 +//! 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::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::{Policy, SelectionContext}; +use sgl_router::workers::Worker; +use std::sync::Arc; + +const BLOCK: usize = 4; + +fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: id.into(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) +} + +#[test] +fn the_weight_override_steers_a_two_term_fusion_past_either_term_alone() { + let ids: Vec = (0..(BLOCK as u32 * 4)).collect(); + let tree = Arc::new(HashTree::new()); + tree.insert( + &KvWorkerId::new("hot".into(), 0), + None, + &compute_block_hashes(&ids, BLOCK), + ); + let oracle = BlockSizeOracle::new(); + oracle.try_set(BLOCK as u32).expect("fresh oracle"); + + // "hot" holds the whole prompt AND is the busiest: the two terms disagree. + let ws = vec![worker("hot"), worker("cold")]; + let _held: Vec<_> = (0..3).map(|_| ws[0].load_guard()).collect(); + + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_request_tokens(Some(&ids)); + let cache = || PrefixCachePolicy::new(Arc::clone(&tree), Arc::clone(&oracle), 1.0); + + // Vacuity guard: if the terms agreed, no weight could change the answer and + // everything below would pass against a composer that read only one of them. + assert_eq!(cache().select(&ws, &ctx).unwrap().id, ws[0].id, "cache→hot"); + assert_eq!( + LoadBasedPolicy::new().select(&ws, &ctx).unwrap().id, + ws[1].id, + "load→cold" + ); + + // Same two terms, same fleet, same request — only the override differs. + for (load_weight, want) in [(0.25, &ws[0]), (4.0, &ws[1])] { + let fused = FusedScorePolicy::new(vec![ + (Arc::new(cache()) as Arc, None), + (Arc::new(LoadBasedPolicy::new()), Some(load_weight)), + ]) + .expect("both terms are fusable"); + let got = fused.select(&ws, &ctx).expect("non-empty fleet"); + assert_eq!(got.id, want.id, "--fuse load_based={load_weight}"); + } +} diff --git a/experimental/sgl-router/tests/component/policies/mod.rs b/experimental/sgl-router/tests/component/policies/mod.rs index 6837a9b64..a2e0a3d85 100644 --- a/experimental/sgl-router/tests/component/policies/mod.rs +++ b/experimental/sgl-router/tests/component/policies/mod.rs @@ -4,6 +4,7 @@ mod zmq_helpers; mod cache_aware_zmq; +mod fused_score; mod kv_events_hash_parity; mod kv_events_tree_concurrent; mod kv_events_two_subscribers; 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 d311a5213..a307343df 100644 --- a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs @@ -12,10 +12,18 @@ //! doesn't render tool schemas, so its ids would diverge from the engine). //! * A request with multimodal (array) content → `input_ids` omitted (a text //! tokenizer can't represent image content). +//! +//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches +//! the built-in V4 chat encoder — the engine-equivalent path — without a +//! template fixture. use axum::body::Body; use axum::http::{Request, StatusCode}; use serde_json::{json, Value}; +use sgl_router::config::{ + ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, + PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, +}; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::factory::build_registry; use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree}; @@ -28,9 +36,35 @@ use std::sync::Arc; use std::time::Duration; use tower::ServiceExt; -use crate::common::cache_aware_fixture::{config, MODEL}; use crate::common::mock_worker::MockWorker; +const MODEL: &str = "deepseek-v4-tiny"; + +fn config() -> Config { + Config { + server: ServerConfig { + host: "0".into(), + port: 0, + }, + observability: ObservabilityConfig::default(), + model: ModelConfig { + id: MODEL.into(), + tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + policy: PolicyKind::CacheAwareZmq, + circuit_breaker: None, + cache_aware: Some(CacheAwareConfig::default()), + sticky: None, + fused: None, + eligibility: None, + }, + discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { + urls: vec!["http://placeholder:0".into()], + }), + proxy: ProxyConfig::default(), + active_load: ActiveLoadConfig::default(), + } +} + fn build_ctx(url: String) -> Arc { let cfg = config(); let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 6d4399147..e91f9d808 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -37,6 +37,8 @@ fn config_for(_worker_url: &str) -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], 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 456a4e659..07e404727 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,8 @@ pub fn config() -> Config { circuit_breaker: None, cache_aware: Some(CacheAwareConfig::default()), sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/failover.rs b/experimental/sgl-router/tests/proxy/failover.rs index 8ca7aeab4..a0a828282 100644 --- a/experimental/sgl-router/tests/proxy/failover.rs +++ b/experimental/sgl-router/tests/proxy/failover.rs @@ -41,6 +41,8 @@ async fn failover_when_one_worker_dies() { }), cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()], diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 2ec959778..041d5e9c3 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -47,6 +47,8 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/header_forwarding.rs b/experimental/sgl-router/tests/proxy/header_forwarding.rs index 06c3e1284..a7a8a701e 100644 --- a/experimental/sgl-router/tests/proxy/header_forwarding.rs +++ b/experimental/sgl-router/tests/proxy/header_forwarding.rs @@ -34,6 +34,8 @@ async fn forwards_whitelisted_headers_strips_others() { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index 9a1ce40db..5df65d763 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -49,6 +49,8 @@ fn config() -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index 0e67f3694..4fd8b0cd4 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -48,6 +48,8 @@ fn config() -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs index c14adf730..54fe93c0a 100644 --- a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs @@ -45,6 +45,8 @@ fn config() -> Config { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs index afe7d6ff7..bbceb4599 100644 --- a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs @@ -25,7 +25,7 @@ use axum::http::{Request, StatusCode}; use serde_json::{json, Value}; use sgl_router::config::{ ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind, - ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, + ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry; @@ -60,10 +60,12 @@ fn config() -> Config { // mid-test; round-robin fallback for the initial pin of a key. sticky: Some(StickyConfig { header_name: HEADER.to_string(), - fallback_policy: PolicyKind::RoundRobin, + fallback_policy: StickyFallbackKind::RoundRobin, idle_secs: 3600, eviction_interval_secs: 3600, }), + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/experimental/sgl-router/tests/proxy/sticky_routing.rs b/experimental/sgl-router/tests/proxy/sticky_routing.rs index 8d3fb2bff..9a5f9b5a7 100644 --- a/experimental/sgl-router/tests/proxy/sticky_routing.rs +++ b/experimental/sgl-router/tests/proxy/sticky_routing.rs @@ -8,7 +8,7 @@ use sgl_router::config::{ ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind, - ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, + ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry; @@ -46,10 +46,12 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc Config { circuit_breaker: None, cache_aware: None, sticky: None, + fused: None, + eligibility: None, }, discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: vec!["http://placeholder:0".into()], diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 067c64c0a..eb57a8a97 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -324,6 +324,20 @@ setup_cargo_cache() { mark_step_done "${FUNCNAME[0]}" } +invalidate_torch_rust_cache() { + if [ "${SGLANG_BUILD_RUST_EXTS:-}" = "none" ]; then + mark_step_done "${FUNCNAME[0]}" + return + fi + + # uv's editable build uses a temporary torch path. Rebuild these units + # under the lock so Cargo does not reuse that path in a later job. + cargo clean --release --manifest-path "${REPO_ROOT}/rust/sglang-radix-tree/Cargo.toml" \ + -p torch-sys -p sglang-radix-tree + + mark_step_done "${FUNCNAME[0]}" +} + release_cargo_cache_lock() { if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then flock --unlock 9 @@ -904,6 +918,7 @@ main() { install_pytorch_stack install_cuda12_deepep_wheel setup_cargo_cache + invalidate_torch_rust_cache install_sglang release_cargo_cache_lock # Diffusion B200 CI imports torch inside install_sglang_kernel after removing diff --git a/test/registered/unit/tools/test_cuda_ci_install_dependency.py b/test/registered/unit/tools/test_cuda_ci_install_dependency.py new file mode 100644 index 000000000..d6679c6b5 --- /dev/null +++ b/test/registered/unit/tools/test_cuda_ci_install_dependency.py @@ -0,0 +1,46 @@ +import importlib.util +import re +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[4] +CI_REGISTER_PATH = REPO_ROOT / "python" / "sglang" / "test" / "ci" / "ci_register.py" +INSTALL_SCRIPT = REPO_ROOT / "scripts" / "ci" / "cuda" / "ci_install_dependency.sh" + + +def _load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +register_cpu_ci = _load_module("ci_register", CI_REGISTER_PATH).register_cpu_ci +register_cpu_ci(est_time=0, suite="base-a-test-cpu") + + +class TestCudaCiInstallDependencyTorchCache(unittest.TestCase): + def test_rebuilds_torch_extensions_before_editable_install(self): + script = INSTALL_SCRIPT.read_text() + + self.assertRegex( + script, + re.compile( + r"setup_cargo_cache\s*\n" + r"\s*invalidate_torch_rust_cache\s*\n" + r"\s*install_sglang" + ), + ) + self.assertIn('"${SGLANG_BUILD_RUST_EXTS:-}" = "none"', script) + self.assertRegex( + script, + re.compile( + r"cargo clean --release --manifest-path " + r"\"\$\{REPO_ROOT\}/rust/sglang-radix-tree/Cargo\.toml\"\s*\\\n" + r"\s*-p torch-sys -p sglang-radix-tree" + ), + ) + + +if __name__ == "__main__": + unittest.main()