[Router] Add composable scoring and eligibility policies (#37731)
Co-authored-by: inkcherry <mingzhi.liu@amd.com>
This commit is contained in:
@@ -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(),
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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(¤t)
|
||||
.and_then(|n| n.children.get(&h).copied())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
let Some(child) = self.nodes.get(&child_id) else {
|
||||
break;
|
||||
};
|
||||
child.last_used.store(now, Ordering::Relaxed);
|
||||
current = child_id;
|
||||
reached += 1;
|
||||
if reached == 1 {
|
||||
// Absent here means never in `alive`: a tail held without
|
||||
// block 0 is reported as holding nothing.
|
||||
alive = child.workers.iter().collect();
|
||||
} else {
|
||||
let mut still = Vec::with_capacity(alive.len());
|
||||
for w in alive {
|
||||
if child.workers.contains(w) {
|
||||
still.push(w);
|
||||
} else {
|
||||
depths.insert(w, reached - 1);
|
||||
}
|
||||
}
|
||||
alive = still;
|
||||
}
|
||||
if alive.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Whoever is still tracked held every level the walk reached.
|
||||
for w in alive {
|
||||
depths.insert(w, reached);
|
||||
}
|
||||
depths.into_iter().map(|(w, d)| (w.clone(), d)).collect()
|
||||
}
|
||||
|
||||
/// Approximate count of *non-root* nodes in the tree.
|
||||
fn node_count(&self) -> usize {
|
||||
// Subtract one for the root sentinel.
|
||||
@@ -652,6 +720,18 @@ impl HashTree {
|
||||
state.match_prefix(parent_hash, block_hashes)
|
||||
}
|
||||
|
||||
/// How many leading blocks of `block_hashes` each worker holds contiguously,
|
||||
/// in one descent under one read lock. [`Self::match_prefix`] names only the
|
||||
/// deepest matched node's holders, so it cannot answer this. Absent = none.
|
||||
pub fn prefix_depths(
|
||||
&self,
|
||||
parent_hash: Option<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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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()],
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user