[Router] Add composable scoring and eligibility policies (#37731)

Co-authored-by: inkcherry <mingzhi.liu@amd.com>
This commit is contained in:
Vincent Gao
2026-09-04 00:01:41 +08:00
committed by GitHub
co-authored by inkcherry
parent 392841f47c
commit 54cadad151
36 changed files with 2256 additions and 227 deletions
+292 -31
View File
@@ -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<f32>,
/// 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<String>,
/// KV Indexer query timeout in milliseconds. Requires
/// `--kv-indexer-endpoint`; defaults to 100.
#[arg(long)]
pub kv_indexer_query_timeout_ms: Option<u64>,
/// 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<usize>,
/// Weighted terms for `--policy fused_score`.
#[arg(long, value_delimiter = ',')]
pub fuse: Vec<FusedTerm>,
/// Ordered hard constraints applied before policy selection.
#[arg(long, value_delimiter = ',')]
pub filter: Vec<FilterKind>,
/// Router-local in-flight limit for `--filter overloaded`.
#[arg(long)]
pub max_in_flight: Option<usize>,
/// Minimum cached prompt share for `--filter prefix_cache`.
#[arg(long)]
pub prefix_cache_min_share: Option<f32>,
// ---- 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<PolicyKind>,
pub sticky_fallback_policy: Option<StickyFallbackKind>,
/// 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<String> {
#[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<Config> {
@@ -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 <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<Config> {
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<Vec<(ScoreTermKind, Option<f32>)>> {
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<f32>)> {
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::<f32>` 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::<f32>().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}");
}
}
}
}
@@ -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(),
+135 -10
View File
@@ -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 = <Self as clap::ValueEnum>::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 = <Self as clap::ValueEnum>::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 = <Self as clap::ValueEnum>::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 = <Self as clap::ValueEnum>::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<CircuitBreakerConfig>,
/// 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<CacheAwareConfig>,
@@ -155,6 +231,10 @@ pub struct ModelConfig {
/// The chat handler reads `sticky.header_name` to populate
/// [`crate::policies::SelectionContext::routing_key`].
pub sticky: Option<StickyConfig>,
/// Terms for `policy = "fused_score"`.
pub fused: Option<Vec<FusedTerm>>,
/// Hard constraints applied before policy selection.
pub eligibility: Option<EligibilityConfig>,
}
/// 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<FilterKind>,
/// `overloaded`: in-flight count at which a worker stops being eligible.
pub max_in_flight: Option<usize>,
/// `prefix_cache` minimum cached prompt share.
pub min_prefix_share: Option<f32>,
}
/// 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<f32>,
}
impl std::str::FromStr for FusedTerm {
type Err = String;
fn from_str(s: &str) -> Result<Self, String> {
let (name, weight) = match s.split_once('=') {
Some((n, w)) => (n, Some(parse_fuse_weight(n, w)?)),
None => (s, None),
};
let kind = <ScoreTermKind as clap::ValueEnum>::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<f32, String> {
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<KvIndexerEndpointConfig>,
}
@@ -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(),
}
@@ -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";
+289 -62
View File
@@ -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<dyn Policy> {
/// Build a dependency-free policy for keyless sticky requests and new pins.
fn build_sticky_fallback(kind: StickyFallbackKind) -> Arc<dyn Policy> {
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<dyn Policy> {
))
}
/// 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<HashTree>,
tokenizers: Arc<TokenizerRegistry>,
block_size_oracle: Arc<BlockSizeOracle>,
) -> Arc<dyn Policy> {
match model.policy {
) -> Result<Arc<dyn Policy>> {
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<HashTree>,
tokenizers: &Arc<TokenizerRegistry>,
block_size_oracle: &Arc<BlockSizeOracle>,
) -> Result<Arc<dyn Policy>> {
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<HashTree>,
block_size_oracle: &Arc<BlockSizeOracle>,
) -> Arc<dyn Policy> {
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<dyn Policy> {
/// Builds one soft scoring term for `--policy fused_score`.
fn build_score(
kind: ScoreTermKind,
tree: &Arc<HashTree>,
block_size_oracle: &Arc<BlockSizeOracle>,
) -> Arc<dyn Policy> {
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<HashTree>,
oracle: &Arc<BlockSizeOracle>,
) -> Result<Arc<dyn Policy>> {
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<dyn Policy>, Option<f32>)> = 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<Arc<dyn Policy>> {
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<dyn Policy> {
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<PolicyRegistry> {
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<Worker> {
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:?}"
);
}
}
}
@@ -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<i64>,
block_hashes: &[i64],
) -> HashMap<KvWorkerId, usize> {
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(&current)
.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<i64>,
block_hashes: &[i64],
) -> HashMap<KvWorkerId, usize> {
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();
@@ -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<Worker>]) -> Option<Arc<Worker>> {
workers
.iter()
.min_by_key(|w| w.active_load())
.map(Arc::clone)
}
}
impl Policy for LoadBasedPolicy {
fn select(&self, workers: &[Arc<Worker>], _ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
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<Worker>], _ctx: &SelectionContext<'_>) -> Vec<f32> {
let loads: Vec<usize> = 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<Worker> {
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<usize> = spec.split(',').map(|s| s.parse().unwrap()).collect();
let ws: Vec<Arc<Worker>> = (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}");
}
}
}
+36 -78
View File
@@ -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<u32>,
/// 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<String> {
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<St
None
}
/// Selection input — carries the request body and the routing tokens
/// (computed once at ingress) so cache-aware policies can hash prefix
/// tokens without reshaping the [`Policy`] trait or re-tokenizing. Today's
/// load-only policies (round-robin, random, power-of-two, load-based) read
/// only `workers`; sticky reads `routing_key`.
///
/// Constructed via [`Self::new`] / [`Self::with_routing_key`]; the
/// ingress-computed tokens are attached with [`Self::with_request_tokens`].
/// Accessors expose immutable references so callers cannot mutate the model
/// id or swap in a different body without going through the constructor.
/// Immutable request data consumed by a routing policy.
pub struct SelectionContext<'a> {
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<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>>;
/// 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<MetricsRegistry>) {}
/// 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<MetricsRegistry>) {
for entry in self.by_model.iter() {
entry.value().attach_metrics(Arc::clone(&metrics));
+49 -5
View File
@@ -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<Worker>], _ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
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<Worker>], _ctx: &SelectionContext<'_>) -> Vec<f32> {
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<Worker> {
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<Arc<Worker>> = (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}");
}
}
}
@@ -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<Worker>], _ctx: &SelectionContext<'_>) -> Vec<bool> {
(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<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
let eligible: Vec<Arc<Worker>> = (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<Worker> {
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<Box<dyn EligibilityFilter>> = 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<Box<dyn EligibilityFilter>> = 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<Worker>], _: &SelectionContext<'_>) -> Vec<bool> {
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<Box<dyn EligibilityFilter>> =
vec![Box::new(Overloaded::new(2)), Box::new(AdmitAll)];
assert!(admit(refs(&chain), &ws, &ctx).is_none());
}
}
@@ -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<Worker>], scores: &[f32]) -> Option<usize>;
}
/// 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<Worker>], scores: &[f32]) -> Option<usize> {
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::<f32>, |acc, s| Some(acc.map_or(s, |b| b.max(s))));
let mut band: Vec<usize> = 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<usize> = 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<Worker> {
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<usize> = (0..3).filter_map(|_| sel.pick(&ws, &[1.0; 3])).collect();
assert_eq!(
picks.len(),
3,
"three tied picks must cover all three: {picks:?}"
);
}
}
@@ -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<Worker>], ctx: &SelectionContext<'_>) -> Vec<bool>;
/// 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<Worker>], ctx: &SelectionContext<'_>) -> Vec<f32>;
/// 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<Item = &'f dyn EligibilityFilter>,
workers: &[Arc<Worker>],
ctx: &SelectionContext<'_>,
) -> Option<Vec<Arc<Worker>>> {
let mut alive: Vec<Arc<Worker>> = 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<Arc<Worker>> = (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<T: ScoringPolicy> Policy for T {
fn select(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
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<dyn Policy>, Option<f32>)>,
}
fn view(t: &(Arc<dyn Policy>, Option<f32>)) -> (&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<dyn Policy>, Option<f32>)>) -> anyhow::Result<Self> {
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<Arc<dyn Policy>>,
inner: Arc<dyn Policy>,
}
impl Pipeline {
/// Reject policies that do not expose an eligibility view.
pub fn new(filters: Vec<Arc<dyn Policy>>, inner: Arc<dyn Policy>) -> anyhow::Result<Self> {
for f in &filters {
anyhow::ensure!(f.can_filter(), "policy {f:?} imposes no eligibility rule");
}
Ok(Self { filters, inner })
}
fn views(&self) -> impl Iterator<Item = &dyn EligibilityFilter> {
(self.filters.iter()).map(|p| p.as_filter().expect("checked by Pipeline::new"))
}
}
impl Policy for Pipeline {
fn select(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
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<crate::server::metrics::MetricsRegistry>) {
self.inner.attach_metrics(metrics);
}
}
impl ScoringPolicy for FusedScorePolicy {
fn scores(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Vec<f32> {
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<dyn EligibilityFilter>],
) -> impl Iterator<Item = &dyn EligibilityFilter> {
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<Worker> {
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<Arc<Worker>> {
vec![worker("a"), worker("b"), worker("c")]
}
fn urls(ws: &[Arc<Worker>]) -> Vec<String> {
ws.iter().map(|w| w.url.clone()).collect()
}
fn term(p: impl Policy + 'static, w: Option<f32>) -> (Arc<dyn Policy>, Option<f32>) {
(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<Worker>], _: &SelectionContext<'_>) -> Vec<f32> {
(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<Worker>], _: &SelectionContext<'_>) -> Vec<bool> {
(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<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
(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<dyn EligibilityFilter> {
Box::new(Keep(names.to_vec(), on_empty))
}
fn boxed(f: impl EligibilityFilter + 'static) -> Box<dyn EligibilityFilter> {
Box::new(f)
}
#[derive(Debug)]
struct PickFirst;
static PICK_FIRST: PickFirst = PickFirst;
impl Selector for PickFirst {
fn pick(&self, workers: &[Arc<Worker>], _: &[f32]) -> Option<usize> {
(!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<Arc<dyn Policy>> = 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<dyn Policy> = 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<Worker>], _: &SelectionContext<'_>) -> Vec<bool> {
vec![true; ws.len()]
}
fn needs_tokens(&self) -> bool {
true
}
}
impl Policy for Hungry {
fn select(&self, ws: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<Arc<Worker>> {
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<Worker>], _: &SelectionContext<'_>) -> Vec<bool> {
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<Worker>], _: &SelectionContext<'_>) -> Vec<bool> {
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<Worker>], _: &SelectionContext<'_>) -> Vec<bool> {
ws.iter().map(|w| !w.url.contains('c')).collect()
}
fn needs_tokens(&self) -> bool {
true
}
}
impl ScoringPolicy for Dual {
fn scores(&self, ws: &[Arc<Worker>], _: &SelectionContext<'_>) -> Vec<f32> {
vec![0.0; ws.len()]
}
fn as_filter(&self) -> Option<&dyn EligibilityFilter> {
Some(self)
}
}
let p: Arc<dyn Policy> = 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"
);
}
}
@@ -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<HashTree>,
block_size_oracle: Arc<BlockSizeOracle>,
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<HashTree>, block_size_oracle: Arc<BlockSizeOracle>, 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<Worker>], ctx: &SelectionContext<'_>) -> Vec<f32> {
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<Worker>], ctx: &SelectionContext<'_>) -> Vec<f32> {
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<Worker>], ctx: &SelectionContext<'_>) -> Vec<bool> {
(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<Worker> {
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<u32> {
(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<HashTree>) -> 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<Worker>], ctx: &SelectionContext<'_>) -> Vec<f32> {
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]);
}
}
@@ -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 {
@@ -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
@@ -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 {
@@ -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 {
@@ -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()],
@@ -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()],
@@ -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 {
@@ -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<Worker> {
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<u32> = (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<dyn Policy>, 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}");
}
}
@@ -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;
@@ -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<AppContext> {
let cfg = config();
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
@@ -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()],
@@ -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()],
@@ -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()],
@@ -47,6 +47,8 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
circuit_breaker: None,
cache_aware: None,
sticky: None,
fused: None,
eligibility: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
@@ -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()],
@@ -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()],
@@ -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()],
@@ -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()],
@@ -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()],
@@ -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<AppContext
cache_aware: None,
sticky: Some(StickyConfig {
header_name: header_name.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()],
@@ -41,6 +41,8 @@ fn config(_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()],
+15
View File
@@ -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
@@ -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()