minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Router] Add load-aware prefill admission and bounded policy proposals (#37843)

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