[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::{
|
use crate::config::{
|
||||||
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
|
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
|
||||||
resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config,
|
resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config,
|
||||||
DiscoveryBackend, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig,
|
DiscoveryBackend, EligibilityConfig, FilterKind, FusedTerm, K8sDiscoveryConfig,
|
||||||
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
KvIndexerEndpointConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig,
|
||||||
StickyConfig,
|
ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE,
|
||||||
};
|
};
|
||||||
|
|
||||||
const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100;
|
const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100;
|
||||||
@@ -73,19 +73,32 @@ pub struct Cli {
|
|||||||
/// Multiplicative load spread gating the absolute balance check.
|
/// Multiplicative load spread gating the absolute balance check.
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub balance_rel_threshold: Option<f32>,
|
pub balance_rel_threshold: Option<f32>,
|
||||||
/// External KV indexer gRPC endpoint used as the authoritative cache signal.
|
/// External KV indexer gRPC endpoint used as the cache signal.
|
||||||
/// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`.
|
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub kv_indexer_endpoint: Option<String>,
|
pub kv_indexer_endpoint: Option<String>,
|
||||||
/// KV Indexer query timeout in milliseconds. Requires
|
/// KV Indexer query timeout in milliseconds. Requires
|
||||||
/// `--kv-indexer-endpoint`; defaults to 100.
|
/// `--kv-indexer-endpoint`; defaults to 100.
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub kv_indexer_query_timeout_ms: Option<u64>,
|
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.
|
/// `--kv-indexer-endpoint`; defaults to 32.
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub kv_indexer_query_max_inflight: Option<usize>,
|
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`) ----
|
// ---- sticky-session policy (only used by `--policy sticky`) ----
|
||||||
/// Request header carrying the routing key for sticky-session routing.
|
/// Request header carrying the routing key for sticky-session routing.
|
||||||
/// Defaults to `x-sgl-routing-key` when `--policy sticky` is set.
|
/// 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
|
/// `round_robin` / `random` / `power_of_two` / `load_based`. Defaults
|
||||||
/// to `round_robin`.
|
/// to `round_robin`.
|
||||||
#[arg(long, value_enum)]
|
#[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
|
/// Evict a sticky assignment after it has been idle (unreferenced) this
|
||||||
/// many seconds. Defaults to 600.
|
/// many seconds. Defaults to 600.
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
@@ -179,7 +192,8 @@ impl Cli {
|
|||||||
|| self.kv_indexer_query_max_inflight.is_some();
|
|| self.kv_indexer_query_max_inflight.is_some();
|
||||||
if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq {
|
if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq {
|
||||||
return Err(anyhow!(
|
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) {
|
if self.kv_indexer_query_timeout_ms == Some(0) {
|
||||||
@@ -202,6 +216,57 @@ impl Cli {
|
|||||||
"--kv-indexer-query-max-inflight requires --kv-indexer-endpoint"
|
"--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()
|
let tuned_sticky = self.routing_key_header.is_some()
|
||||||
|| self.sticky_fallback_policy.is_some()
|
|| self.sticky_fallback_policy.is_some()
|
||||||
|| self.sticky_idle_secs.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
|
// policy is selected. The header name must parse as an HTTP header
|
||||||
// name so a typo fails at startup rather than silently never
|
// name so a typo fails at startup rather than silently never
|
||||||
// matching any request header; the fallback must be a
|
// matching any request header.
|
||||||
// dependency-free policy the factory can build standalone.
|
|
||||||
let sticky = if self.policy == PolicyKind::Sticky {
|
let sticky = if self.policy == PolicyKind::Sticky {
|
||||||
let d = StickyConfig::default();
|
let d = StickyConfig::default();
|
||||||
let header_name = self.routing_key_header.unwrap_or(d.header_name);
|
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}")
|
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);
|
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 idle_secs = self.sticky_idle_secs.unwrap_or(d.idle_secs);
|
||||||
let eviction_interval_secs = self
|
let eviction_interval_secs = self
|
||||||
.sticky_eviction_interval_secs
|
.sticky_eviction_interval_secs
|
||||||
@@ -315,6 +370,8 @@ impl Cli {
|
|||||||
circuit_breaker,
|
circuit_breaker,
|
||||||
cache_aware,
|
cache_aware,
|
||||||
sticky,
|
sticky,
|
||||||
|
fused,
|
||||||
|
eligibility,
|
||||||
},
|
},
|
||||||
discovery,
|
discovery,
|
||||||
proxy: ProxyConfig {
|
proxy: ProxyConfig {
|
||||||
@@ -401,7 +458,7 @@ fn join_selector(terms: &[String]) -> Option<String> {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::config::{DiscoveryBackend, K8sDiscoveryMode};
|
use crate::config::{DiscoveryBackend, K8sDiscoveryMode, ScoreTermKind};
|
||||||
|
|
||||||
/// Parse argv (without the leading binary name) into a `Config`.
|
/// Parse argv (without the leading binary name) into a `Config`.
|
||||||
fn into_config(args: &[&str]) -> Result<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.
|
/// `--policy load_based` parses to the load-based selector.
|
||||||
#[test]
|
#[test]
|
||||||
fn parses_load_based_policy() {
|
fn parses_load_based_policy() {
|
||||||
@@ -967,11 +1058,35 @@ mod tests {
|
|||||||
assert_eq!(c.model.policy, PolicyKind::Sticky);
|
assert_eq!(c.model.policy, PolicyKind::Sticky);
|
||||||
let s = c.model.sticky.expect("sticky config built");
|
let s = c.model.sticky.expect("sticky config built");
|
||||||
assert_eq!(s.header_name, "x-sgl-routing-key");
|
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.idle_secs, 600);
|
||||||
assert_eq!(s.eviction_interval_secs, 60);
|
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]
|
#[test]
|
||||||
fn sticky_flags_override_defaults() {
|
fn sticky_flags_override_defaults() {
|
||||||
let c = into_config_owned(with_model(&[
|
let c = into_config_owned(with_model(&[
|
||||||
@@ -991,11 +1106,74 @@ mod tests {
|
|||||||
.unwrap();
|
.unwrap();
|
||||||
let s = c.model.sticky.expect("sticky config built");
|
let s = c.model.sticky.expect("sticky config built");
|
||||||
assert_eq!(s.header_name, "x-session-id");
|
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.idle_secs, 120);
|
||||||
assert_eq!(s.eviction_interval_secs, 15);
|
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]
|
#[test]
|
||||||
fn non_sticky_policy_leaves_sticky_none() {
|
fn non_sticky_policy_leaves_sticky_none() {
|
||||||
let c = into_config_owned(with_model(&[
|
let c = into_config_owned(with_model(&[
|
||||||
@@ -1049,10 +1227,7 @@ mod tests {
|
|||||||
]))
|
]))
|
||||||
.unwrap_err()
|
.unwrap_err()
|
||||||
.to_string();
|
.to_string();
|
||||||
assert!(
|
assert!(err.contains("invalid value"), "got: {err}");
|
||||||
err.contains("--sticky-fallback-policy must be one of"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -1067,10 +1242,7 @@ mod tests {
|
|||||||
]))
|
]))
|
||||||
.unwrap_err()
|
.unwrap_err()
|
||||||
.to_string();
|
.to_string();
|
||||||
assert!(
|
assert!(err.contains("invalid value"), "got: {err}");
|
||||||
err.contains("--sticky-fallback-policy must be one of"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A zero eviction interval would panic `tokio::time::interval` at
|
/// A zero eviction interval would panic `tokio::time::interval` at
|
||||||
@@ -1110,4 +1282,93 @@ mod tests {
|
|||||||
"got: {err}"
|
"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,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: urls.iter().map(|s| s.to_string()).collect(),
|
urls: urls.iter().map(|s| s.to_string()).collect(),
|
||||||
|
|||||||
@@ -70,7 +70,8 @@ impl Default for ActiveLoadConfig {
|
|||||||
/// policy factory.
|
/// policy factory.
|
||||||
///
|
///
|
||||||
/// Accepted on the CLI (`--policy`) as `round_robin` / `random` /
|
/// 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)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)]
|
||||||
pub enum PolicyKind {
|
pub enum PolicyKind {
|
||||||
#[default]
|
#[default]
|
||||||
@@ -83,6 +84,9 @@ pub enum PolicyKind {
|
|||||||
/// Selects the currently least-loaded worker.
|
/// Selects the currently least-loaded worker.
|
||||||
#[value(name = "load_based")]
|
#[value(name = "load_based")]
|
||||||
LoadBased,
|
LoadBased,
|
||||||
|
/// Weighted sum of `--fuse` terms.
|
||||||
|
#[value(name = "fused_score")]
|
||||||
|
FusedScore,
|
||||||
/// Cache-aware routing fed by SGLang's ZMQ KV-cache event publisher.
|
/// Cache-aware routing fed by SGLang's ZMQ KV-cache event publisher.
|
||||||
/// Requires the model to have a tokenizer loaded; cache_aware tuning
|
/// Requires the model to have a tokenizer loaded; cache_aware tuning
|
||||||
/// lives on `ModelConfig::cache_aware`.
|
/// lives on `ModelConfig::cache_aware`.
|
||||||
@@ -97,6 +101,78 @@ pub enum PolicyKind {
|
|||||||
Sticky,
|
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)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct ServerConfig {
|
pub struct ServerConfig {
|
||||||
pub host: String,
|
pub host: String,
|
||||||
@@ -146,7 +222,7 @@ pub struct ModelConfig {
|
|||||||
pub tokenizer_path: String,
|
pub tokenizer_path: String,
|
||||||
pub policy: PolicyKind,
|
pub policy: PolicyKind,
|
||||||
pub circuit_breaker: Option<CircuitBreakerConfig>,
|
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 = "cache_aware_zmq"`. `None` falls back to defaults at
|
||||||
/// policy construction time.
|
/// policy construction time.
|
||||||
pub cache_aware: Option<CacheAwareConfig>,
|
pub cache_aware: Option<CacheAwareConfig>,
|
||||||
@@ -155,6 +231,10 @@ pub struct ModelConfig {
|
|||||||
/// The chat handler reads `sticky.header_name` to populate
|
/// The chat handler reads `sticky.header_name` to populate
|
||||||
/// [`crate::policies::SelectionContext::routing_key`].
|
/// [`crate::policies::SelectionContext::routing_key`].
|
||||||
pub sticky: Option<StickyConfig>,
|
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.
|
/// External KV Indexer client settings.
|
||||||
@@ -165,6 +245,54 @@ pub struct KvIndexerEndpointConfig {
|
|||||||
pub query_max_inflight: usize,
|
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.
|
/// Per-model cache-aware tuning.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct CacheAwareConfig {
|
pub struct CacheAwareConfig {
|
||||||
@@ -182,8 +310,7 @@ pub struct CacheAwareConfig {
|
|||||||
/// that the absolute check is gated on. Default 1.1 — 10 % relative
|
/// that the absolute check is gated on. Default 1.1 — 10 % relative
|
||||||
/// difference triggers re-balancing.
|
/// difference triggers re-balancing.
|
||||||
pub balance_rel_threshold: f32,
|
pub balance_rel_threshold: f32,
|
||||||
/// Optional external KV Indexer client configuration. When configured, it
|
/// Optional external KV Indexer client configuration.
|
||||||
/// replaces the local ZMQ radix tree as the cache signal.
|
|
||||||
pub kv_indexer_endpoint: Option<KvIndexerEndpointConfig>,
|
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`
|
/// Per-model sticky-session tuning. Built from the `--routing-key-header`
|
||||||
/// / `--sticky-*` flags by [`crate::config::cli::Cli::into_config`], which
|
/// / `--sticky-*` flags by [`crate::config::cli::Cli::into_config`], which
|
||||||
/// also validates that `header_name` parses as an HTTP header name and
|
/// also validates that `header_name` parses as an HTTP header name.
|
||||||
/// that `fallback_policy` is one of the dependency-free policies.
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct StickyConfig {
|
pub struct StickyConfig {
|
||||||
/// Request header carrying the routing key. Validated to parse as a
|
/// 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
|
/// to pick the initial worker when a new key is first seen. One of
|
||||||
/// `round_robin` / `random` / `power_of_two` / `load_based` — the
|
/// `round_robin` / `random` / `power_of_two` / `load_based` — the
|
||||||
/// dependency-free policies the factory can build standalone (no
|
/// dependency-free policies the factory can build standalone (no
|
||||||
/// `HashTree` / tokenizer / ZMQ feed). `cache_aware_zmq` and `sticky`
|
/// `HashTree` / tokenizer / ZMQ feed).
|
||||||
/// are rejected at config-build time.
|
pub fallback_policy: StickyFallbackKind,
|
||||||
pub fallback_policy: PolicyKind,
|
|
||||||
/// Evict an assignment after it has been idle (unreferenced) this many
|
/// Evict an assignment after it has been idle (unreferenced) this many
|
||||||
/// seconds. Bounds the map against unbounded routing-key cardinality.
|
/// seconds. Bounds the map against unbounded routing-key cardinality.
|
||||||
pub idle_secs: u64,
|
pub idle_secs: u64,
|
||||||
@@ -247,7 +372,7 @@ impl Default for StickyConfig {
|
|||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
header_name: DEFAULT_STICKY_HEADER.to_string(),
|
header_name: DEFAULT_STICKY_HEADER.to_string(),
|
||||||
fallback_policy: PolicyKind::RoundRobin,
|
fallback_policy: StickyFallbackKind::RoundRobin,
|
||||||
idle_secs: default_sticky_idle_secs(),
|
idle_secs: default_sticky_idle_secs(),
|
||||||
eviction_interval_secs: default_sticky_eviction_interval_secs(),
|
eviction_interval_secs: default_sticky_eviction_interval_secs(),
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -144,14 +144,12 @@ impl CacheAwareZmqPolicy {
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
// The index may include unhealthy workers or workers in another pool.
|
|
||||||
let best_routable_blocks = matches
|
let best_routable_blocks = matches
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|m| workers.iter().any(|worker| worker.url == m.address))
|
.filter(|m| workers.iter().any(|worker| worker.url == m.address))
|
||||||
.map(|m| m.matched_prefix_blocks)
|
.map(|m| m.matched_prefix_blocks)
|
||||||
.max()
|
.max()
|
||||||
.unwrap_or(0);
|
.unwrap_or(0);
|
||||||
|
|
||||||
let match_rate = best_routable_blocks as f32 / signal.query_blocks as f32;
|
let match_rate = best_routable_blocks as f32 / signal.query_blocks as f32;
|
||||||
if let Some(metrics) = self.metrics.get() {
|
if let Some(metrics) = self.metrics.get() {
|
||||||
metrics.observe_overlap_blocks(ctx.model().0.as_str(), best_routable_blocks as u64);
|
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);
|
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() {
|
if let Some(signal) = ctx.external_prefix() {
|
||||||
return self
|
return self
|
||||||
.select_external(workers, ctx, signal)
|
.select_external(workers, ctx, signal)
|
||||||
@@ -349,6 +345,8 @@ mod tests {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
||||||
crate::config::StaticUrlsDiscoveryConfig {
|
crate::config::StaticUrlsDiscoveryConfig {
|
||||||
@@ -400,7 +398,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn external_prefix_signal_skips_unroutable_best_match() {
|
fn external_prefix_signal_selects_the_best_routable_match() {
|
||||||
let mut config = cfg_default();
|
let mut config = cfg_default();
|
||||||
config.cache_threshold = 0.0;
|
config.cache_threshold = 0.0;
|
||||||
let policy = CacheAwareZmqPolicy::new(
|
let policy = CacheAwareZmqPolicy::new(
|
||||||
@@ -438,7 +436,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[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 tree = Arc::new(HashTree::new());
|
||||||
let registry = tokenizer_registry_with_tiny();
|
let registry = tokenizer_registry_with_tiny();
|
||||||
let text = "hello world hello world hello world";
|
let text = "hello world hello world hello world";
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
// SPDX-License-Identifier: Apache-2.0
|
// 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::discovery::ModelId;
|
||||||
use crate::policies::{
|
use crate::policies::{
|
||||||
cache_aware_zmq::CacheAwareZmqPolicy,
|
cache_aware_zmq::CacheAwareZmqPolicy,
|
||||||
@@ -10,27 +12,25 @@ use crate::policies::{
|
|||||||
power_of_two::PowerOfTwoChoicesPolicy,
|
power_of_two::PowerOfTwoChoicesPolicy,
|
||||||
random::RandomPolicy,
|
random::RandomPolicy,
|
||||||
round_robin::RoundRobinPolicy,
|
round_robin::RoundRobinPolicy,
|
||||||
|
scoring::{
|
||||||
|
admission::Overloaded, prefix_cache, prefix_cache::PrefixCachePolicy, FusedScorePolicy,
|
||||||
|
Pipeline,
|
||||||
|
},
|
||||||
sticky::StickyPolicy,
|
sticky::StickyPolicy,
|
||||||
Policy, PolicyRegistry,
|
Policy, PolicyRegistry,
|
||||||
};
|
};
|
||||||
use crate::tokenizer::TokenizerRegistry;
|
use crate::tokenizer::TokenizerRegistry;
|
||||||
use anyhow::Result;
|
use anyhow::{anyhow, Result};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
/// Build a dependency-free policy for use as the sticky-session fallback
|
/// Build a dependency-free policy for keyless sticky requests and new pins.
|
||||||
/// (keyless requests + initial pin of a new key). `Cli::into_config`
|
fn build_sticky_fallback(kind: StickyFallbackKind) -> Arc<dyn Policy> {
|
||||||
/// 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> {
|
|
||||||
match kind {
|
match kind {
|
||||||
PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()),
|
StickyFallbackKind::RoundRobin => Arc::new(RoundRobinPolicy::new()),
|
||||||
PolicyKind::Random => Arc::new(RandomPolicy::new()),
|
StickyFallbackKind::Random => Arc::new(RandomPolicy::new()),
|
||||||
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
StickyFallbackKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
||||||
PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
|
StickyFallbackKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
|
||||||
PolicyKind::CacheAwareZmq | PolicyKind::Sticky => {
|
|
||||||
unreachable!("sticky fallback is validated to be dependency-free in Cli::into_config")
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -46,20 +46,38 @@ fn build_sticky(model: &ModelConfig) -> Arc<dyn Policy> {
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Construct a policy for a single model from its [`ModelConfig`] and the
|
/// Constructs a policy for one model.
|
||||||
/// 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.
|
|
||||||
pub fn build_policy(
|
pub fn build_policy(
|
||||||
model: &ModelConfig,
|
model: &ModelConfig,
|
||||||
tree: Arc<HashTree>,
|
tree: Arc<HashTree>,
|
||||||
tokenizers: Arc<TokenizerRegistry>,
|
tokenizers: Arc<TokenizerRegistry>,
|
||||||
block_size_oracle: Arc<BlockSizeOracle>,
|
block_size_oracle: Arc<BlockSizeOracle>,
|
||||||
) -> Arc<dyn Policy> {
|
) -> Result<Arc<dyn Policy>> {
|
||||||
match model.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::RoundRobin => Arc::new(RoundRobinPolicy::new()),
|
||||||
PolicyKind::Random => Arc::new(RandomPolicy::new()),
|
PolicyKind::Random => Arc::new(RandomPolicy::new()),
|
||||||
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
||||||
@@ -74,32 +92,90 @@ pub fn build_policy(
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
PolicyKind::Sticky => build_sticky(model),
|
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
|
/// Builds one soft scoring term for `--policy fused_score`.
|
||||||
/// a policy without wiring the cache-aware dependencies; rejects
|
fn build_score(
|
||||||
/// `CacheAwareZmq` to keep the call sites that don't have a `HashTree` /
|
kind: ScoreTermKind,
|
||||||
/// `TokenizerRegistry` to hand from accidentally compiling.
|
tree: &Arc<HashTree>,
|
||||||
#[cfg(test)]
|
block_size_oracle: &Arc<BlockSizeOracle>,
|
||||||
pub fn build_policy_kind_only(kind: PolicyKind) -> Arc<dyn Policy> {
|
) -> Arc<dyn Policy> {
|
||||||
match kind {
|
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::RoundRobin => Arc::new(RoundRobinPolicy::new()),
|
||||||
PolicyKind::Random => Arc::new(RandomPolicy::new()),
|
PolicyKind::Random => Arc::new(RandomPolicy::new()),
|
||||||
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
||||||
PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
|
PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
|
||||||
PolicyKind::CacheAwareZmq => {
|
PolicyKind::CacheAwareZmq => Arc::new(CacheAwareZmqPolicy::new(
|
||||||
// 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(),
|
crate::config::CacheAwareConfig::default(),
|
||||||
Arc::new(HashTree::new()),
|
Arc::new(HashTree::new()),
|
||||||
Arc::new(TokenizerRegistry::default()),
|
Arc::new(TokenizerRegistry::default()),
|
||||||
BlockSizeOracle::new(),
|
BlockSizeOracle::new(),
|
||||||
))
|
)),
|
||||||
}
|
|
||||||
PolicyKind::Sticky => {
|
PolicyKind::Sticky => {
|
||||||
let s = crate::config::StickyConfig::default();
|
let s = crate::config::StickyConfig::default();
|
||||||
Arc::new(StickyPolicy::new(
|
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),
|
build_sticky_fallback(s.fallback_policy),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
PolicyKind::FusedScore => {
|
||||||
|
return Err(anyhow!("--policy {kind} needs --fuse terms from the model"))
|
||||||
}
|
}
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn build_registry(
|
pub fn build_registry(
|
||||||
@@ -126,19 +205,12 @@ pub fn build_registry(
|
|||||||
Arc::clone(&tree),
|
Arc::clone(&tree),
|
||||||
Arc::clone(&tokenizers),
|
Arc::clone(&tokenizers),
|
||||||
Arc::clone(&block_size_oracle),
|
Arc::clone(&block_size_oracle),
|
||||||
),
|
)?,
|
||||||
);
|
);
|
||||||
Ok(reg)
|
Ok(reg)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Convenience for tests + non-cache-aware callers: builds a registry with
|
/// Builds a registry with empty cache-aware dependencies.
|
||||||
/// 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.
|
|
||||||
pub fn build_registry_with_defaults(cfg: &Config) -> Result<PolicyRegistry> {
|
pub fn build_registry_with_defaults(cfg: &Config) -> Result<PolicyRegistry> {
|
||||||
build_registry(
|
build_registry(
|
||||||
cfg,
|
cfg,
|
||||||
@@ -156,7 +228,90 @@ mod tests {
|
|||||||
StaticUrlsDiscoveryConfig,
|
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 {
|
fn cfg_with_model(id: &str, policy: PolicyKind) -> Config {
|
||||||
Config {
|
Config {
|
||||||
@@ -172,6 +327,8 @@ mod tests {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
@@ -183,13 +340,68 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn build_policy_kind_only_covers_all_variants() {
|
fn build_policy_kind_only_covers_all_variants() {
|
||||||
// Trivially total — the match is exhaustive over `PolicyKind`.
|
for kind in [
|
||||||
let _ = build_policy_kind_only(PolicyKind::RoundRobin);
|
PolicyKind::RoundRobin,
|
||||||
let _ = build_policy_kind_only(PolicyKind::Random);
|
PolicyKind::Random,
|
||||||
let _ = build_policy_kind_only(PolicyKind::PowerOfTwo);
|
PolicyKind::PowerOfTwo,
|
||||||
let _ = build_policy_kind_only(PolicyKind::LoadBased);
|
PolicyKind::LoadBased,
|
||||||
let _ = build_policy_kind_only(PolicyKind::CacheAwareZmq);
|
PolicyKind::CacheAwareZmq,
|
||||||
let _ = build_policy_kind_only(PolicyKind::Sticky);
|
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]
|
#[test]
|
||||||
@@ -209,9 +421,6 @@ mod tests {
|
|||||||
let tokenizers = Arc::new(TokenizerRegistry::default());
|
let tokenizers = Arc::new(TokenizerRegistry::default());
|
||||||
let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap();
|
let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap();
|
||||||
let p = reg.get(&ModelId("modelA".into())).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:?}");
|
let dbg = format!("{p:?}");
|
||||||
assert!(
|
assert!(
|
||||||
dbg.contains("CacheAwareZmqPolicy"),
|
dbg.contains("CacheAwareZmqPolicy"),
|
||||||
@@ -246,4 +455,22 @@ mod tests {
|
|||||||
"expected StickyPolicy debug repr, got: {dbg}",
|
"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.
|
/// Approximate count of *non-root* nodes in the tree.
|
||||||
fn node_count(&self) -> usize {
|
fn node_count(&self) -> usize {
|
||||||
// Subtract one for the root sentinel.
|
// Subtract one for the root sentinel.
|
||||||
@@ -652,6 +720,18 @@ impl HashTree {
|
|||||||
state.match_prefix(parent_hash, block_hashes)
|
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
|
/// Approximate number of non-root nodes in the tree (the root sentinel
|
||||||
/// is not counted). Useful for metrics and to decide when to call
|
/// is not counted). Useful for metrics and to decide when to call
|
||||||
/// [`HashTree::evict_lru`].
|
/// [`HashTree::evict_lru`].
|
||||||
@@ -702,6 +782,36 @@ mod tests {
|
|||||||
ids.iter().map(|w| (*w).clone()).collect()
|
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]
|
#[test]
|
||||||
fn empty_match_returns_zero_no_workers() {
|
fn empty_match_returns_zero_no_workers() {
|
||||||
let tree = HashTree::new();
|
let tree = HashTree::new();
|
||||||
|
|||||||
@@ -1,14 +1,12 @@
|
|||||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
// SPDX-License-Identifier: Apache-2.0
|
// 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 crate::workers::Worker;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
/// Deterministic load-based policy.
|
/// Prefers the least-loaded candidate; `select()` is the blanket impl's.
|
||||||
///
|
|
||||||
/// Chooses the candidate with the lowest current `Worker::active_load`.
|
|
||||||
/// Ties follow the candidate slice order, which is registry-dependent.
|
|
||||||
#[derive(Debug, Default)]
|
#[derive(Debug, Default)]
|
||||||
pub struct LoadBasedPolicy;
|
pub struct LoadBasedPolicy;
|
||||||
|
|
||||||
@@ -16,18 +14,23 @@ impl LoadBasedPolicy {
|
|||||||
pub fn new() -> Self {
|
pub fn new() -> Self {
|
||||||
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 {
|
impl ScoringPolicy for LoadBasedPolicy {
|
||||||
fn select(&self, workers: &[Arc<Worker>], _ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
|
/// `1.0` for the least loaded down to `0.0` for the most, min-max scaled to
|
||||||
Self::pick_min_load(workers)
|
/// 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 {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||||
|
use crate::policies::scoring::argmax::TIE_EPSILON;
|
||||||
|
use crate::policies::Policy;
|
||||||
|
|
||||||
fn worker(id: &str) -> Arc<Worker> {
|
fn worker(id: &str) -> Arc<Worker> {
|
||||||
Arc::new(Worker::new(WorkerSpec {
|
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() {
|
fn empty_returns_none() {
|
||||||
let policy = LoadBasedPolicy::new();
|
let m = ModelId("tiny".into());
|
||||||
let model = ModelId("tiny".into());
|
let ctx = SelectionContext::new(&m, None);
|
||||||
let ctx = SelectionContext::new(&model, None);
|
assert!(LoadBasedPolicy::new().select(&[], &ctx).is_none());
|
||||||
assert!(policy.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]
|
#[test]
|
||||||
fn picks_lowest_active_load() {
|
fn scores_rank_strictly_by_load_and_the_choice_lands_on_the_minimum() {
|
||||||
let policy = LoadBasedPolicy::new();
|
|
||||||
let model = ModelId("tiny".into());
|
let model = ModelId("tiny".into());
|
||||||
let ctx = SelectionContext::new(&model, None);
|
let (ctx, p) = (SelectionContext::new(&model, None), LoadBasedPolicy::new());
|
||||||
let w0 = worker("w0");
|
for spec in ["0,1", "2,1,0", "1,0,1", "0,0", "5,2,9,2", "300,900"] {
|
||||||
let w1 = worker("w1");
|
let loads: Vec<usize> = spec.split(',').map(|s| s.parse().unwrap()).collect();
|
||||||
let _g0 = w0.load_guard();
|
let ws: Vec<Arc<Worker>> = (0..loads.len()).map(|i| worker(&format!("w{i}"))).collect();
|
||||||
assert_eq!(
|
let _held: Vec<_> = (ws.iter().zip(&loads))
|
||||||
policy.select(&[w0, Arc::clone(&w1)], &ctx).unwrap().id,
|
.flat_map(|(w, n)| (0..*n).map(move |_| w.load_guard()))
|
||||||
w1.id
|
.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 random;
|
||||||
pub mod registry;
|
pub mod registry;
|
||||||
pub mod round_robin;
|
pub mod round_robin;
|
||||||
|
pub mod scoring;
|
||||||
pub mod sticky;
|
pub mod sticky;
|
||||||
|
|
||||||
use crate::discovery::ModelId;
|
use crate::discovery::ModelId;
|
||||||
|
use crate::policies::scoring::{EligibilityFilter, ScoringPolicy};
|
||||||
use crate::server::metrics::MetricsRegistry;
|
use crate::server::metrics::MetricsRegistry;
|
||||||
use crate::tokenizer::{adapter, TokenizerRegistry};
|
use crate::tokenizer::{adapter, TokenizerRegistry};
|
||||||
use crate::workers::Worker;
|
use crate::workers::Worker;
|
||||||
use dashmap::DashMap;
|
use dashmap::DashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
/// Tokens produced once at ingress for a request. Consumed by the
|
/// Tokens produced at ingress for routing and optional engine forwarding.
|
||||||
/// 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).
|
|
||||||
pub struct RequestTokens {
|
pub struct RequestTokens {
|
||||||
/// The prompt token ids.
|
/// The prompt token ids.
|
||||||
pub ids: Vec<u32>,
|
pub ids: Vec<u32>,
|
||||||
/// True only when the ids were produced via the model's chat encoder —
|
/// Whether the token ids are safe to forward as engine `input_ids`.
|
||||||
/// 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.
|
|
||||||
pub engine_equivalent: bool,
|
pub engine_equivalent: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// External indexer answer prepared by the async ingress path for the
|
/// External indexer answer prepared by the async ingress path for a
|
||||||
/// synchronous cache-aware policy.
|
/// cache-aware policy.
|
||||||
pub struct ExternalPrefixSignal {
|
pub struct ExternalPrefixSignal {
|
||||||
pub outcome: sgl_kv_indexer::PrefixOutcome,
|
pub outcome: sgl_kv_indexer::PrefixOutcome,
|
||||||
pub query_blocks: usize,
|
pub query_blocks: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Produce the routing tokens — and whether they are engine-equivalent —
|
/// Tokenizes a request for routing. Chat-encoder tokens are engine-equivalent;
|
||||||
/// from an already-parsed request body, using the shared tokenizer registry.
|
/// raw prompt tokens are used only for routing.
|
||||||
///
|
|
||||||
/// 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.
|
|
||||||
pub fn request_tokens_for(
|
pub fn request_tokens_for(
|
||||||
tokenizers: &TokenizerRegistry,
|
tokenizers: &TokenizerRegistry,
|
||||||
model_id: &ModelId,
|
model_id: &ModelId,
|
||||||
@@ -82,11 +61,7 @@ pub fn request_tokens_for(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tokenize `text` for `model_id` via the shared registry. Returns `None` if
|
/// Tokenizes text with the model tokenizer.
|
||||||
/// 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.
|
|
||||||
fn tokenize_text(
|
fn tokenize_text(
|
||||||
tokenizers: &TokenizerRegistry,
|
tokenizers: &TokenizerRegistry,
|
||||||
model_id: &ModelId,
|
model_id: &ModelId,
|
||||||
@@ -97,13 +72,6 @@ fn tokenize_text(
|
|||||||
Ok(ids) if !ids.is_empty() => Some(ids),
|
Ok(ids) if !ids.is_empty() => Some(ids),
|
||||||
Ok(_) => None,
|
Ok(_) => None,
|
||||||
Err(e) => {
|
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!(
|
tracing::warn!(
|
||||||
model = %model_id,
|
model = %model_id,
|
||||||
error = %e,
|
error = %e,
|
||||||
@@ -114,10 +82,7 @@ fn tokenize_text(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Extract a raw prompt-text candidate from an already-parsed JSON request
|
/// Extracts raw prompt text from a parsed request body.
|
||||||
/// 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`]).
|
|
||||||
///
|
///
|
||||||
/// Supported shapes (in priority order):
|
/// Supported shapes (in priority order):
|
||||||
/// 1. `"prompt": "..."` — `/v1/completions`-style.
|
/// 1. `"prompt": "..."` — `/v1/completions`-style.
|
||||||
@@ -129,7 +94,6 @@ fn tokenize_text(
|
|||||||
/// multimodal content blocks; text-only blocks concatenated.
|
/// multimodal content blocks; text-only blocks concatenated.
|
||||||
/// 5. `"text": "..."` — SGLang `/generate` native form.
|
/// 5. `"text": "..."` — SGLang `/generate` native form.
|
||||||
///
|
///
|
||||||
/// Anything else yields `None`.
|
|
||||||
pub(crate) fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option<String> {
|
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()) {
|
if let Some(s) = v.get("prompt").and_then(|p| p.as_str()) {
|
||||||
return Some(s.to_string());
|
return Some(s.to_string());
|
||||||
@@ -173,16 +137,7 @@ pub(crate) fn extract_prompt_text_from_value(v: &serde_json::Value) -> Option<St
|
|||||||
None
|
None
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Selection input — carries the request body and the routing tokens
|
/// Immutable request data consumed by a routing policy.
|
||||||
/// (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.
|
|
||||||
pub struct SelectionContext<'a> {
|
pub struct SelectionContext<'a> {
|
||||||
model: &'a ModelId,
|
model: &'a ModelId,
|
||||||
request_body: Option<&'a [u8]>,
|
request_body: Option<&'a [u8]>,
|
||||||
@@ -216,9 +171,7 @@ impl<'a> SelectionContext<'a> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Attach the ingress-computed routing tokens. When present, the
|
/// Attaches ingress-computed routing tokens.
|
||||||
/// cache-aware policy consumes these instead of re-parsing and
|
|
||||||
/// re-tokenizing the body.
|
|
||||||
pub fn with_request_tokens(mut self, request_tokens: Option<&'a [u32]>) -> Self {
|
pub fn with_request_tokens(mut self, request_tokens: Option<&'a [u32]>) -> Self {
|
||||||
self.request_tokens = request_tokens;
|
self.request_tokens = request_tokens;
|
||||||
self
|
self
|
||||||
@@ -244,8 +197,7 @@ impl<'a> SelectionContext<'a> {
|
|||||||
self.routing_key
|
self.routing_key
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Ingress-precomputed routing tokens, if any. `None` means the policy
|
/// Returns ingress-computed routing tokens.
|
||||||
/// must derive tokens itself (e.g. a caller that didn't pre-tokenize).
|
|
||||||
pub fn request_tokens(&self) -> Option<&[u32]> {
|
pub fn request_tokens(&self) -> Option<&[u32]> {
|
||||||
self.request_tokens
|
self.request_tokens
|
||||||
}
|
}
|
||||||
@@ -258,25 +210,33 @@ impl<'a> SelectionContext<'a> {
|
|||||||
pub trait Policy: Send + Sync + std::fmt::Debug {
|
pub trait Policy: Send + Sync + std::fmt::Debug {
|
||||||
fn select(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>>;
|
fn select(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>>;
|
||||||
|
|
||||||
/// Whether this policy's ROUTING decision needs the request tokens (i.e.
|
/// Whether policy selection needs request tokens.
|
||||||
/// 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.
|
|
||||||
fn needs_request_tokens(&self) -> bool {
|
fn needs_request_tokens(&self) -> bool {
|
||||||
false
|
false
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Attach the process metrics registry after construction. Default is a
|
/// Attaches the process metrics registry after construction.
|
||||||
/// 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.
|
|
||||||
fn attach_metrics(&self, _metrics: Arc<MetricsRegistry>) {}
|
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)]
|
#[derive(Debug, Default)]
|
||||||
@@ -293,9 +253,7 @@ impl PolicyRegistry {
|
|||||||
self.by_model.get(model).map(|p| p.clone())
|
self.by_model.get(model).map(|p| p.clone())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Inject the metrics registry into every registered policy. Called once
|
/// Attaches metrics to each registered policy.
|
||||||
/// at startup (after the registry is built) so metrics-emitting policies
|
|
||||||
/// can record into the shared registry.
|
|
||||||
pub fn attach_metrics(&self, metrics: Arc<MetricsRegistry>) {
|
pub fn attach_metrics(&self, metrics: Arc<MetricsRegistry>) {
|
||||||
for entry in self.by_model.iter() {
|
for entry in self.by_model.iter() {
|
||||||
entry.value().attach_metrics(Arc::clone(&metrics));
|
entry.value().attach_metrics(Arc::clone(&metrics));
|
||||||
|
|||||||
@@ -1,9 +1,9 @@
|
|||||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
|
||||||
use crate::policies::{Policy, SelectionContext};
|
use crate::policies::{scoring::ScoringPolicy, SelectionContext};
|
||||||
use crate::workers::Worker;
|
use crate::workers::Worker;
|
||||||
use rand::seq::SliceRandom;
|
use rand::Rng;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
#[derive(Debug, Default)]
|
#[derive(Debug, Default)]
|
||||||
@@ -15,8 +15,52 @@ impl RandomPolicy {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Policy for RandomPolicy {
|
impl ScoringPolicy for RandomPolicy {
|
||||||
fn select(&self, workers: &[Arc<Worker>], _ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
|
/// Argmax of n iid uniforms IS a uniform choice: exactly the old `choose`.
|
||||||
workers.choose(&mut rand::thread_rng()).cloned()
|
/// 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,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
||||||
crate::config::StaticUrlsDiscoveryConfig {
|
crate::config::StaticUrlsDiscoveryConfig {
|
||||||
|
|||||||
@@ -54,6 +54,8 @@ mod tests {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
};
|
};
|
||||||
let app = crate::server::app::build_router(std::sync::Arc::new(ctx));
|
let app = crate::server::app::build_router(std::sync::Arc::new(ctx));
|
||||||
let res = app
|
let res = app
|
||||||
|
|||||||
@@ -121,6 +121,8 @@ mod tests {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
||||||
crate::config::StaticUrlsDiscoveryConfig {
|
crate::config::StaticUrlsDiscoveryConfig {
|
||||||
|
|||||||
@@ -235,6 +235,8 @@ mod tests {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
discovery: crate::config::DiscoveryBackend::StaticUrls(
|
||||||
crate::config::StaticUrlsDiscoveryConfig {
|
crate::config::StaticUrlsDiscoveryConfig {
|
||||||
|
|||||||
@@ -481,6 +481,8 @@ mod tests {
|
|||||||
}),
|
}),
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://test:30000".into()],
|
urls: vec!["http://test:30000".into()],
|
||||||
|
|||||||
@@ -133,6 +133,8 @@ async fn static_urls_pd_role_resolved_end_to_end() {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec![url.clone()],
|
urls: vec![url.clone()],
|
||||||
|
|||||||
@@ -72,6 +72,8 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: sgl_router::config::DiscoveryBackend::StaticUrls(
|
discovery: sgl_router::config::DiscoveryBackend::StaticUrls(
|
||||||
sgl_router::config::StaticUrlsDiscoveryConfig {
|
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 zmq_helpers;
|
||||||
|
|
||||||
mod cache_aware_zmq;
|
mod cache_aware_zmq;
|
||||||
|
mod fused_score;
|
||||||
mod kv_events_hash_parity;
|
mod kv_events_hash_parity;
|
||||||
mod kv_events_tree_concurrent;
|
mod kv_events_tree_concurrent;
|
||||||
mod kv_events_two_subscribers;
|
mod kv_events_two_subscribers;
|
||||||
|
|||||||
@@ -12,10 +12,18 @@
|
|||||||
//! doesn't render tool schemas, so its ids would diverge from the engine).
|
//! doesn't render tool schemas, so its ids would diverge from the engine).
|
||||||
//! * A request with multimodal (array) content → `input_ids` omitted (a text
|
//! * A request with multimodal (array) content → `input_ids` omitted (a text
|
||||||
//! tokenizer can't represent image content).
|
//! 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::body::Body;
|
||||||
use axum::http::{Request, StatusCode};
|
use axum::http::{Request, StatusCode};
|
||||||
use serde_json::{json, Value};
|
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::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||||
use sgl_router::policies::factory::build_registry;
|
use sgl_router::policies::factory::build_registry;
|
||||||
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
||||||
@@ -28,9 +36,35 @@ use std::sync::Arc;
|
|||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use tower::ServiceExt;
|
use tower::ServiceExt;
|
||||||
|
|
||||||
use crate::common::cache_aware_fixture::{config, MODEL};
|
|
||||||
use crate::common::mock_worker::MockWorker;
|
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> {
|
fn build_ctx(url: String) -> Arc<AppContext> {
|
||||||
let cfg = config();
|
let cfg = config();
|
||||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||||
|
|||||||
@@ -37,6 +37,8 @@ fn config_for(_worker_url: &str) -> Config {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -29,6 +29,8 @@ pub fn config() -> Config {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: Some(CacheAwareConfig::default()),
|
cache_aware: Some(CacheAwareConfig::default()),
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -41,6 +41,8 @@ async fn failover_when_one_worker_dies() {
|
|||||||
}),
|
}),
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
|
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,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -34,6 +34,8 @@ async fn forwards_whitelisted_headers_strips_others() {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -49,6 +49,8 @@ fn config() -> Config {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -48,6 +48,8 @@ fn config() -> Config {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -45,6 +45,8 @@ fn config() -> Config {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -25,7 +25,7 @@ use axum::http::{Request, StatusCode};
|
|||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use sgl_router::config::{
|
use sgl_router::config::{
|
||||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
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::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||||
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
|
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.
|
// mid-test; round-robin fallback for the initial pin of a key.
|
||||||
sticky: Some(StickyConfig {
|
sticky: Some(StickyConfig {
|
||||||
header_name: HEADER.to_string(),
|
header_name: HEADER.to_string(),
|
||||||
fallback_policy: PolicyKind::RoundRobin,
|
fallback_policy: StickyFallbackKind::RoundRobin,
|
||||||
idle_secs: 3600,
|
idle_secs: 3600,
|
||||||
eviction_interval_secs: 3600,
|
eviction_interval_secs: 3600,
|
||||||
}),
|
}),
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -8,7 +8,7 @@
|
|||||||
|
|
||||||
use sgl_router::config::{
|
use sgl_router::config::{
|
||||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
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::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||||
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
|
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,
|
cache_aware: None,
|
||||||
sticky: Some(StickyConfig {
|
sticky: Some(StickyConfig {
|
||||||
header_name: header_name.to_string(),
|
header_name: header_name.to_string(),
|
||||||
fallback_policy: PolicyKind::RoundRobin,
|
fallback_policy: StickyFallbackKind::RoundRobin,
|
||||||
idle_secs: 3600,
|
idle_secs: 3600,
|
||||||
eviction_interval_secs: 3600,
|
eviction_interval_secs: 3600,
|
||||||
}),
|
}),
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -41,6 +41,8 @@ fn config(_worker_url: &str) -> Config {
|
|||||||
circuit_breaker: None,
|
circuit_breaker: None,
|
||||||
cache_aware: None,
|
cache_aware: None,
|
||||||
sticky: None,
|
sticky: None,
|
||||||
|
fused: None,
|
||||||
|
eligibility: None,
|
||||||
},
|
},
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
urls: vec!["http://placeholder:0".into()],
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
|||||||
@@ -324,6 +324,20 @@ setup_cargo_cache() {
|
|||||||
mark_step_done "${FUNCNAME[0]}"
|
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() {
|
release_cargo_cache_lock() {
|
||||||
if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then
|
if [ "${CARGO_TARGET_LOCK_HELD:-0}" = "1" ]; then
|
||||||
flock --unlock 9
|
flock --unlock 9
|
||||||
@@ -904,6 +918,7 @@ main() {
|
|||||||
install_pytorch_stack
|
install_pytorch_stack
|
||||||
install_cuda12_deepep_wheel
|
install_cuda12_deepep_wheel
|
||||||
setup_cargo_cache
|
setup_cargo_cache
|
||||||
|
invalidate_torch_rust_cache
|
||||||
install_sglang
|
install_sglang
|
||||||
release_cargo_cache_lock
|
release_cargo_cache_lock
|
||||||
# Diffusion B200 CI imports torch inside install_sglang_kernel after removing
|
# 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