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

[sgl-router] refactor - move policy-required states under src/state (#40272)

Browse Source 
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
2026-09-20 16:49:23 -07:00
committed by GitHub
co-authored by Claude Fable 5.1
parent acd20a516e
commit 4a9dc5c4af
78 changed files with 787 additions and 655 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/benches/tree_lookup.rs
+1 -1
View File
@@ -35,7 +35,7 @@ use std::thread;
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use sgl_router::policies::kv_events::tree::{HashTree, KvWorkerId};
use sgl_router::state::kv_events::tree::{HashTree, KvWorkerId};
fn build_tree(num_workers: usize, blocks_per_worker: usize, seed: u64) -> HashTree {
let tree = HashTree::new();
experimental/sgl-router/src/config/cli.rs
+7 -7
View File
@@ -10,10 +10,10 @@ use std::num::NonZeroU32;
use crate::config::sampling::{parse_sampling_overrides, ConflictPolicy};
use crate::config::{
default_cb_cool_down, default_host, default_port, default_proxy_request_timeout_secs,
default_shutdown_drain_secs, default_stale_request_timeout_secs, resolve_mode,
ActiveLoadConfig, AffinityConfig, AffinityMode, CacheAwareConfig, CachePrefixProvider,
CircuitBreakerConfig, Config, DecodePolicyKind, DiscoveryBackend, EligibilityConfig,
FilterKind, FusedTerm, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig,
default_shutdown_drain_secs, default_stale_request_timeout_secs, resolve_mode, AffinityConfig,
AffinityMode, CacheAwareConfig, CachePrefixProvider, CircuitBreakerConfig, Config,
DecodePolicyKind, DiscoveryBackend, EligibilityConfig, FilterKind, FusedTerm,
InflightLoadConfig, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig,
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode,
StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE,
};
@@ -390,7 +390,7 @@ impl Cli {
proxy: ProxyConfig {
request_timeout_secs: self.server.request_timeout_secs,
},
active_load: ActiveLoadConfig {
router_inflight_load: InflightLoadConfig {
stale_request_timeout_secs: self.server.stale_request_timeout_secs,
},
};
@@ -906,7 +906,7 @@ mod tests {
assert_eq!(c.model.policy, PolicyKind::RoundRobin);
assert_eq!(c.model.id, "qwen3-0.6b");
assert_eq!(c.proxy.request_timeout_secs, 300);
assert_eq!(c.active_load.stale_request_timeout_secs, 600);
assert_eq!(c.router_inflight_load.stale_request_timeout_secs, 600);
assert_eq!(c.server.shutdown_drain_secs, 30);
}
@@ -1473,7 +1473,7 @@ mod tests {
]))
.unwrap();
assert_eq!(c.proxy.request_timeout_secs, 120);
assert_eq!(c.active_load.stale_request_timeout_secs, 240);
assert_eq!(c.router_inflight_load.stale_request_timeout_secs, 240);
}
#[test]
experimental/sgl-router/src/config/mod.rs
+1 -1
View File
@@ -237,7 +237,7 @@ mod tests {
urls: urls.iter().map(|s| s.to_string()).collect(),
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/src/config/types.rs
+3 -3
View File
@@ -11,7 +11,7 @@ pub struct Config {
/// Discovery mode resolved from CLI options; static URLs are checked by [`Config::validate`].
pub discovery: DiscoveryBackend,
pub proxy: ProxyConfig,
pub active_load: ActiveLoadConfig,
pub router_inflight_load: InflightLoadConfig,
}
/// Outbound request timeout settings.
@@ -35,7 +35,7 @@ impl Default for ProxyConfig {
/// Request-tracking timeout; defaults above the proxy timeout.
#[derive(Debug, Clone, Copy)]
pub struct ActiveLoadConfig {
pub struct InflightLoadConfig {
/// Maximum request-entry lifetime before cancellation with 504 `stale_request_expired`.
pub stale_request_timeout_secs: u64,
}
@@ -44,7 +44,7 @@ pub fn default_stale_request_timeout_secs() -> u64 {
600
}
impl Default for ActiveLoadConfig {
impl Default for InflightLoadConfig {
fn default() -> Self {
Self {
stale_request_timeout_secs: default_stale_request_timeout_secs(),
experimental/sgl-router/src/discovery/k8s.rs
+1 -1
View File
@@ -1022,7 +1022,7 @@ mod tests {
}
/// Pod is replaced (same IP, different UID) — router must see this as
/// a Removed+Added cycle so the new pod gets fresh CB/active_load
/// a Removed+Added cycle so the new pod gets fresh CB/router_inflight_load
/// state. Without UID-keyed WorkerIds, two consecutive
/// `process_events` snapshots would dedup by `addr:port` and the
/// new pod would inherit the dead pod's state.
experimental/sgl-router/src/lib.rs
+1
View File
@@ -14,5 +14,6 @@ pub mod health;
pub mod policies;
pub mod proxy;
pub mod server;
pub mod state;
pub mod tokenizer;
pub mod workers;
experimental/sgl-router/src/main.rs
+20 -13
View File
@@ -8,14 +8,17 @@ use sgl_router::{
config::{CachePrefixProvider, Cli, Config, KvIndexerEndpointConfig, LogFormat, PolicyKind},
discovery::spawn_discovery,
policies::{
active_load::{spawn_janitor, ActiveLoadRegistry, JanitorHandle, SystemTimeClock},
factory::build_registry as build_policy_registry,
kv_events::{BlockSizeOracle, KvEventIndex},
prefix_provider::RadixTreePrefixProvider,
factory::build_registry as build_policy_registry, prefix_provider::RadixTreePrefixProvider,
PolicyRegistry,
},
proxy::Proxy,
server::{app::build_router, app_context::AppContext, shutdown::drain_for_termination},
state::{
kv_events::{BlockSizeOracle, KvEventIndex},
load_monitor::router_inflight_load::{
spawn_janitor, JanitorHandle, RouterInflightLoadRegistry, SystemTimeClock,
},
},
tokenizer::TokenizerRegistry,
workers::{manager, WorkerRegistry},
};
@@ -228,10 +231,14 @@ fn start_engine_state_monitor(use_external_indexer: bool) -> Arc<KvEventIndex> {
}
}
fn start_local_inflight_tracker(config: &Config) -> (Arc<ActiveLoadRegistry>, JanitorHandle) {
let timeout_secs = config.active_load.stale_request_timeout_secs;
let local_inflight_requests =
ActiveLoadRegistry::new(Arc::new(SystemTimeClock), Duration::from_secs(timeout_secs));
fn start_local_inflight_tracker(
config: &Config,
) -> (Arc<RouterInflightLoadRegistry>, JanitorHandle) {
let timeout_secs = config.router_inflight_load.stale_request_timeout_secs;
let local_inflight_requests = RouterInflightLoadRegistry::new(
Arc::new(SystemTimeClock),
Duration::from_secs(timeout_secs),
);
// Reap stale requests at one tenth of their timeout, bounded to 1–60 seconds.
let sweep_interval = Duration::from_secs((timeout_secs / 10).clamp(1, 60));
let inflight_cleanup = spawn_janitor(Arc::clone(&local_inflight_requests), sweep_interval);
@@ -242,7 +249,7 @@ async fn start_worker_discovery_and_manager(
config: &Config,
worker_registry: &Arc<WorkerRegistry>,
engine_state: &Arc<KvEventIndex>,
local_inflight_requests: &Arc<ActiveLoadRegistry>,
local_inflight_requests: &Arc<RouterInflightLoadRegistry>,
) -> Result<(JoinHandle<()>, JoinHandle<()>)> {
let (worker_events, discovery_handle) =
spawn_discovery(config).await.context("spawn discovery")?;
@@ -262,7 +269,7 @@ fn build_app_context(
tokenizers: Arc<TokenizerRegistry>,
worker_registry: Arc<WorkerRegistry>,
routing_policies: Arc<PolicyRegistry>,
local_inflight_requests: Arc<ActiveLoadRegistry>,
local_inflight_requests: Arc<RouterInflightLoadRegistry>,
engine_state: &KvEventIndex,
external_kv_indexer_client: Option<Arc<dyn PrefixIndex>>,
) -> Result<Arc<AppContext>> {
@@ -272,7 +279,7 @@ fn build_app_context(
.context("build proxy client")?,
);
let mut app_context = AppContext::with_active_load(
let mut app_context = AppContext::with_router_inflight_load(
config.clone(),
tokenizers,
proxy,
@@ -289,7 +296,7 @@ fn build_app_context(
.is_some_and(|cache| cache.prefix_provider == CachePrefixProvider::RadixTree))
.then(|| RadixTreePrefixProvider::new(engine_state.tree(), Arc::clone(&block_size_oracle)));
app_context.block_size_oracle = block_size_oracle;
app_context.engine_load = engine_state.engine_load();
app_context.engine_reported_load = engine_state.engine_reported_load();
app_context.kv_metrics = engine_state.metrics_source();
Ok(Arc::new(app_context))
}
@@ -350,7 +357,7 @@ async fn report_drain_progress(
tracing::$level!(
elapsed_secs = $elapsed,
inflight_http = app_context.inflight_http.count(),
inflight_proxied = app_context.active_load.inflight_count(),
inflight_proxied = app_context.router_inflight_load.inflight_count(),
"still draining in-flight requests; this phase is unbounded and ends at \
SIGKILL when terminationGracePeriodSeconds expires",
)
experimental/sgl-router/src/policies/admission.rs
+58 -38
View File
@@ -8,7 +8,7 @@
//! Router-local load.
//!
//! The optional queue gate (`--worker-queue-limit`) is the one criterion here
//! that reads [`EngineWorkerLoad::num_waiting_reqs`] rather than the native
//! that reads [`EngineReportedWorkerLoad::num_waiting_reqs`] rather than the native
//! monitor fields: a worker already making requests wait cannot win on cache
//! affinity. It fails open on a missing sample — see [`queue_gate_admits`].
//! The optional saturation floor (`--saturation-queue-floor`) cancels a
@@ -17,9 +17,11 @@
//! fleet reads below the floor, the request pins to the least-pressured
//! prefix owner instead of cold-prefilling on a non-owner.
use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad, NativeCacheWorkerLoad};
use crate::policies::power_of_two::select_k_with_snapshot;
use crate::policies::{CacheCandidate, CacheCandidateProposal, GuardHints, SelectionProposal};
use crate::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad, EngineReportedWorkerLoad,
};
use crate::workers::Worker;
use std::cmp::Ordering;
use std::collections::HashMap;
@@ -173,13 +175,13 @@ pub struct CacheCandidateResolution {
/// The queue gate, in one place. Both subtle decisions live here: the
/// boundary is `<` (a worker AT the limit is already making this request
/// wait), and an unknown queue ADMITS. The gate reads
/// [`EngineWorkerLoad::num_waiting_reqs`] because that is what the request
/// [`EngineReportedWorkerLoad::num_waiting_reqs`] because that is what the request
/// cares about, and the router-side in-flight counter cannot separate a
/// running request from a waiting one — so there is no honest substitute,
/// and the gate fails open rather than comparing the limit against a
/// different quantity.
pub(crate) fn queue_gate_admits(
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
worker: &Worker,
limit: Option<u64>,
) -> bool {
@@ -196,7 +198,7 @@ pub(crate) fn queue_gate_admits(
/// an empty fleet, or a single worker with no fresh sample all make this
/// false — an unknown queue is not a proven full one.
pub(crate) fn fleet_is_all_queued(
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
fleet: &[Arc<Worker>],
limit: Option<u64>,
) -> bool {
@@ -216,7 +218,7 @@ pub(crate) fn fleet_is_all_queued(
pub fn resolve_cache_candidates(
proposal: &CacheCandidateProposal,
request_input_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
fleet: &[Arc<Worker>],
) -> CacheCandidateResolution {
let queue_limit = proposal.worker_queue_limit;
@@ -397,7 +399,7 @@ pub fn resolve_prefill(
range: &CandidateRange<'_>,
proposal: &SelectionProposal,
request_input_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
queue_limit: Option<u64>,
min_load_choices: usize,
) -> Option<FinalDecision> {
@@ -431,7 +433,7 @@ pub fn resolve_prefill_admitted(
range: &CandidateRange<'_>,
proposal: &SelectionProposal,
request_input_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
queue_limit: Option<u64>,
) -> Option<FinalDecision> {
if !contains_worker(range, &proposal.primary) {
@@ -492,7 +494,7 @@ pub fn resolve_decode(
domain: &CandidateDomain,
proposal: &SelectionProposal,
request_kv_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
) -> Option<FinalDecision> {
if domain.stage != RoutingStage::Decode || !contains_domain_worker(domain, &proposal.primary) {
return None;
@@ -548,7 +550,7 @@ fn is_proposal_worker_eligible(proposal: &SelectionProposal, candidate: &Arc<Wor
/// Applies snapshot-backed capacity admission when native monitor data is complete.
/// Workers without monitor data remain eligible and use Router-local ordering.
fn has_kv_capacity(load: Option<&NativeCacheWorkerLoad>, requested_tokens: u64) -> bool {
fn has_kv_capacity(load: Option<&EngineReportedSchedulingLoad>, requested_tokens: u64) -> bool {
let Some(load) = load else {
return true;
};
@@ -560,7 +562,7 @@ fn is_prefill_admitted(
range: &CandidateRange<'_>,
worker: &Arc<Worker>,
request_input_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
) -> bool {
let load = snapshot.fresh_native_cache_load_for_url(&worker.url);
has_kv_capacity(load, request_input_tokens)
@@ -576,7 +578,7 @@ fn is_prefill_admitted(
fn is_decode_admitted(
worker: &Arc<Worker>,
request_kv_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
) -> bool {
has_kv_capacity(
snapshot.fresh_native_cache_load_for_url(&worker.url),
@@ -692,8 +694,8 @@ fn materially_more_pressured(
/// External values are compared only when every candidate is present. Mixed
/// candidate sets use Router-local active load to preserve ordering.
pub(crate) struct FreshLoadLookup<'a> {
by_worker_id: HashMap<String, &'a NativeCacheWorkerLoad>,
basic_by_worker_id: HashMap<String, &'a EngineWorkerLoad>,
by_worker_id: HashMap<String, &'a EngineReportedSchedulingLoad>,
basic_by_worker_id: HashMap<String, &'a EngineReportedWorkerLoad>,
local_active_by_worker_id: HashMap<String, usize>,
compare_engine: bool,
compare_basic_engine: bool,
@@ -701,13 +703,13 @@ pub(crate) struct FreshLoadLookup<'a> {
impl<'a> FreshLoadLookup<'a> {
pub(crate) fn new<'w>(
snapshot: Option<&'a EngineLoadSnapshot>,
snapshot: Option<&'a EngineReportedLoadSnapshot>,
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()))
.map(|worker| (worker.id.0.clone(), worker.router_inflight_load()))
.collect();
let by_worker_id = snapshot
.into_iter()
@@ -745,14 +747,14 @@ impl<'a> FreshLoadLookup<'a> {
pub(crate) fn get(
&self,
worker_id: &crate::discovery::WorkerId,
) -> Option<&'a NativeCacheWorkerLoad> {
) -> Option<&'a EngineReportedSchedulingLoad> {
self.by_worker_id.get(worker_id.0.as_str()).copied()
}
fn comparable_get(
&self,
worker_id: &crate::discovery::WorkerId,
) -> Option<&'a NativeCacheWorkerLoad> {
) -> Option<&'a EngineReportedSchedulingLoad> {
self.compare_engine.then(|| self.get(worker_id)).flatten()
}
@@ -853,7 +855,7 @@ impl<'a> FreshLoadLookup<'a> {
}
struct PressureKey<'a> {
load: Option<&'a NativeCacheWorkerLoad>,
load: Option<&'a EngineReportedSchedulingLoad>,
local_active: usize,
}
@@ -861,7 +863,7 @@ fn range_fallback(
range: &CandidateRange<'_>,
legal: &[Arc<Worker>],
request_input_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
queue_limit: Option<u64>,
) -> Option<(Arc<Worker>, DecisionReason)> {
let admitted = legal
@@ -922,7 +924,7 @@ fn legal_prefill_candidates(
fn decode_domain_fallback(
domain: &CandidateDomain,
request_kv_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
) -> Option<(Arc<Worker>, DecisionReason)> {
let admitted = domain
.workers
@@ -940,7 +942,7 @@ fn decode_domain_fallback(
pub(crate) fn compare_prefill_pressure(
left: &Arc<Worker>,
right: &Arc<Worker>,
snapshot: Option<&EngineLoadSnapshot>,
snapshot: Option<&EngineReportedLoadSnapshot>,
) -> Ordering {
match snapshot.and_then(|snapshot| {
Some((
@@ -948,13 +950,19 @@ pub(crate) fn compare_prefill_pressure(
snapshot.fresh_native_cache_load_for_url(&right.url)?,
))
}) {
Some((left_load, right_load)) => compare_prefill_load(left_load, right_load)
.then_with(|| left.active_load().cmp(&right.active_load())),
None => left.active_load().cmp(&right.active_load()),
Some((left_load, right_load)) => {
compare_prefill_load(left_load, right_load).then_with(|| {
left.router_inflight_load()
.cmp(&right.router_inflight_load())
})
}
None => left
.router_inflight_load()
.cmp(&right.router_inflight_load()),
}
}
fn prefill_pressure_key(load: &NativeCacheWorkerLoad) -> (u64, u64, u64) {
fn prefill_pressure_key(load: &EngineReportedSchedulingLoad) -> (u64, u64, u64) {
(
load.num_waiting_uncached_tokens,
load.num_waiting_reqs,
@@ -962,7 +970,10 @@ fn prefill_pressure_key(load: &NativeCacheWorkerLoad) -> (u64, u64, u64) {
)
}
fn compare_prefill_load(left: &NativeCacheWorkerLoad, right: &NativeCacheWorkerLoad) -> Ordering {
fn compare_prefill_load(
left: &EngineReportedSchedulingLoad,
right: &EngineReportedSchedulingLoad,
) -> Ordering {
match (
left.estimated_prefill_queue_ms,
right.estimated_prefill_queue_ms,
@@ -978,7 +989,7 @@ fn compare_prefill_load(left: &NativeCacheWorkerLoad, right: &NativeCacheWorkerL
pub(crate) fn compare_decode_pressure(
left: &Arc<Worker>,
right: &Arc<Worker>,
snapshot: Option<&EngineLoadSnapshot>,
snapshot: Option<&EngineReportedLoadSnapshot>,
) -> Ordering {
match snapshot.and_then(|snapshot| {
Some((
@@ -986,13 +997,22 @@ pub(crate) fn compare_decode_pressure(
snapshot.fresh_native_cache_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()),
Some((left_load, right_load)) => {
compare_decode_load(left_load, right_load).then_with(|| {
left.router_inflight_load()
.cmp(&right.router_inflight_load())
})
}
None => left
.router_inflight_load()
.cmp(&right.router_inflight_load()),
}
}
fn compare_decode_load(left: &NativeCacheWorkerLoad, right: &NativeCacheWorkerLoad) -> Ordering {
fn compare_decode_load(
left: &EngineReportedSchedulingLoad,
right: &EngineReportedSchedulingLoad,
) -> 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_used_tokens)
.saturating_mul(u128::from(right_cap))
@@ -1010,7 +1030,7 @@ fn pressure_guard_prefers_backup(
primary: &Arc<Worker>,
backup: &Arc<Worker>,
hints: &GuardHints,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
) -> bool {
if !hints.enable_pressure_guard {
return false;
@@ -1059,15 +1079,15 @@ mod tests {
}))
}
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
7,
entries
.iter()
.map(|(worker, running, waiting, used, capacity)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: *running,
num_waiting_reqs: *waiting,
num_waiting_uncached_tokens: *waiting,
@@ -1693,7 +1713,7 @@ mod tests {
];
let proposal = SelectionProposal::primary(Arc::clone(&primary));
let load = |waiting: u64, waiting_uncached: u64, total: u64, max_total: u64| {
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: 0,
num_waiting_reqs: waiting,
num_waiting_uncached_tokens: waiting_uncached,
@@ -1710,7 +1730,7 @@ mod tests {
// waits 5 (over the limit) behind 1 uncached token, busy_unqueued
// waits 3 (under the limit) behind 1000 uncached tokens. The primary
// is KV-full, so it is not admitted at all.
let loads = EngineLoadSnapshot::from_native_cache_workers(
let loads = EngineReportedLoadSnapshot::from_native_cache_workers(
7,
[
(primary.url.clone(), load(0, 0, 10_000, 10_000)),
experimental/sgl-router/src/policies/decode.rs
+5 -5
View File
@@ -8,15 +8,15 @@ use crate::policies::admission::{
compare_decode_pressure, resolve_decode, CandidateDomain, DecisionReason, FinalDecision,
RoutingStage,
};
use crate::policies::engine_load::EngineLoadSnapshot;
use crate::policies::registry::select_decode_with_affinity;
use crate::policies::{ProposalKind, SelectionProposal};
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadSnapshot;
use rand::Rng;
use std::sync::Arc;
#[derive(Debug, Default)]
pub struct DecodeSelectionContext<'a> {
load_snapshot: Option<&'a EngineLoadSnapshot>,
load_snapshot: Option<&'a EngineReportedLoadSnapshot>,
prefill_url: Option<&'a str>,
}
@@ -29,12 +29,12 @@ impl<'a> DecodeSelectionContext<'a> {
}
/// Engine load snapshot captured at request ingress.
pub fn with_load_snapshot(mut self, load_snapshot: &'a EngineLoadSnapshot) -> Self {
pub fn with_load_snapshot(mut self, load_snapshot: &'a EngineReportedLoadSnapshot) -> Self {
self.load_snapshot = Some(load_snapshot);
self
}
pub fn load_snapshot(&self) -> Option<&EngineLoadSnapshot> {
pub fn load_snapshot(&self) -> Option<&EngineReportedLoadSnapshot> {
self.load_snapshot
}
@@ -62,7 +62,7 @@ pub fn resolve_decode_with_capacity_fallback(
domain: &CandidateDomain,
proposal: &SelectionProposal,
request_kv_tokens: u64,
snapshot: &EngineLoadSnapshot,
snapshot: &EngineReportedLoadSnapshot,
) -> Option<FinalDecision> {
if let Some(decision) = resolve_decode(domain, proposal, request_kv_tokens, snapshot) {
return Some(decision);
experimental/sgl-router/src/policies/factory.rs
+3 -3
View File
@@ -7,7 +7,6 @@ use crate::config::{
use crate::discovery::ModelId;
use crate::policies::{
cache_aware::CacheAwarePolicy,
kv_events::{BlockSizeOracle, HashTree},
load_based::LoadBasedPolicy,
power_of_two::PowerOfTwoChoicesPolicy,
random::RandomPolicy,
@@ -20,6 +19,7 @@ use crate::policies::{
sticky::StickyPolicy,
Policy, PolicyRegistry,
};
use crate::state::kv_events::{BlockSizeOracle, HashTree};
use anyhow::{anyhow, Result};
use std::sync::Arc;
use std::time::Duration;
@@ -237,7 +237,7 @@ pub fn build_registry_with_defaults(cfg: &Config) -> Result<PolicyRegistry> {
mod tests {
use super::*;
use crate::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ProxyConfig, ServerConfig,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ProxyConfig, ServerConfig,
StaticUrlsDiscoveryConfig,
};
@@ -353,7 +353,7 @@ mod tests {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/src/policies/load_based.rs
+18 -13
View File
@@ -48,9 +48,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 crate::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedWorkerLoad,
};
use std::collections::HashMap;
use std::time::Instant;
@@ -94,7 +96,10 @@ mod tests {
let ok = (scores[i] > scores[j] + TIE_EPSILON, scores[i].is_nan());
assert_eq!(ok, (loads[i] < loads[j], false), "{spec} scored {scores:?}");
}
let got = p.select(&ws, &ctx).expect("non-empty").active_load();
let got = p
.select(&ws, &ctx)
.expect("non-empty")
.router_inflight_load();
assert_eq!(got, *loads.iter().min().expect("non-empty"), "{spec}");
}
}
@@ -107,12 +112,12 @@ mod tests {
// 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(
let snapshot = EngineReportedLoadSnapshot::from_workers(
23,
HashMap::from([
(
w0.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 50,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -122,7 +127,7 @@ mod tests {
),
(
w1.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -148,12 +153,12 @@ mod tests {
let w0 = worker("w0");
let w1 = worker("w1");
let captured_at = Instant::now();
let snapshot = EngineLoadSnapshot::from_workers(
let snapshot = EngineReportedLoadSnapshot::from_workers(
37,
HashMap::from([
(
w0.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -163,7 +168,7 @@ mod tests {
),
(
w1.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -192,12 +197,12 @@ mod tests {
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(
let snapshot = EngineReportedLoadSnapshot::from_workers(
41,
HashMap::from([
(
w0.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -207,7 +212,7 @@ mod tests {
),
(
w1.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -233,11 +238,11 @@ mod tests {
let w0 = worker("w0");
let w1 = worker("w1");
let _local_load = [w0.load_guard(), w0.load_guard()];
let snapshot = EngineLoadSnapshot::from_workers(
let snapshot = EngineReportedLoadSnapshot::from_workers(
43,
HashMap::from([(
w0.url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_tokens: 0,
experimental/sgl-router/src/policies/mod.rs
+11 -12
View File
@@ -1,14 +1,11 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
pub mod active_load;
pub mod admission;
pub mod buckets;
pub mod cache_aware;
pub mod decode;
pub mod engine_load;
pub mod factory;
pub mod kv_events;
pub mod load_based;
pub mod power_of_two;
pub mod prefix_provider;
@@ -22,9 +19,9 @@ pub mod sticky;
use crate::discovery::ModelId;
use crate::policies::buckets::{BucketRequest, BucketSelector};
use crate::policies::engine_load::EngineLoadSnapshot;
use crate::policies::scoring::{EligibilityFilter, ScoringPolicy};
use crate::server::metrics::MetricsRegistry;
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadSnapshot;
use crate::tokenizer::{adapter, TokenizerRegistry};
use crate::workers::Worker;
use dashmap::DashMap;
@@ -176,7 +173,7 @@ pub struct SelectionContext<'a> {
input_tokens: Option<u64>,
request_tokens: Option<&'a [u32]>,
external_prefix: Option<&'a ExternalPrefixSignal>,
load_snapshot: Option<&'a EngineLoadSnapshot>,
load_snapshot: Option<&'a EngineReportedLoadSnapshot>,
prefill_cache_bucket: Option<(&'a BucketSelector, BucketRequest)>,
affinity_lookup_enabled: bool,
affinity_assignment_enabled: bool,
@@ -254,7 +251,7 @@ impl<'a> SelectionContext<'a> {
}
/// Attaches the engine load snapshot captured at request ingress.
pub fn with_load_snapshot(mut self, load_snapshot: &'a EngineLoadSnapshot) -> Self {
pub fn with_load_snapshot(mut self, load_snapshot: &'a EngineReportedLoadSnapshot) -> Self {
self.load_snapshot = Some(load_snapshot);
self
}
@@ -315,7 +312,7 @@ impl<'a> SelectionContext<'a> {
self.external_prefix
}
pub fn load_snapshot(&self) -> Option<&EngineLoadSnapshot> {
pub fn load_snapshot(&self) -> Option<&EngineReportedLoadSnapshot> {
self.load_snapshot
}
@@ -592,10 +589,12 @@ mod tests {
resolve_cache_candidates, resolve_prefill, CandidateRange, DecisionReason, FreshLoadLookup,
};
use crate::policies::cache_aware::CacheAwarePolicy;
use crate::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
use crate::policies::power_of_two::PowerOfTwoChoicesPolicy;
use crate::policies::round_robin::RoundRobinPolicy;
use crate::policies::session_aware::SessionAwarePolicy;
use crate::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad,
};
use std::collections::HashMap;
use std::time::Instant;
@@ -1222,15 +1221,15 @@ mod tests {
assert!(proposal.backup.is_some());
}
fn snapshot(entries: &[(&Arc<Worker>, TestEngineLoad)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, TestEngineLoad)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
1,
entries
.iter()
.map(|(worker, aggregate)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: aggregate.num_running_reqs,
num_waiting_reqs: aggregate.num_waiting_reqs,
num_waiting_uncached_tokens: aggregate
@@ -1616,7 +1615,7 @@ mod tests {
fn missing_engine_snapshot_does_not_hard_reject_a_registry_healthy_primary() {
let primary = worker("primary");
let workers = vec![Arc::clone(&primary)];
let snapshot = EngineLoadSnapshot::default();
let snapshot = EngineReportedLoadSnapshot::default();
let decision = resolve_prefill(
&CandidateRange::global(&workers),
experimental/sgl-router/src/policies/power_of_two.rs
+11 -11
View File
@@ -15,8 +15,8 @@
use crate::config::DEFAULT_MIN_LOAD_CHOICES;
use crate::policies::admission::{compare_prefill_pressure, queue_gate_admits};
use crate::policies::engine_load::EngineLoadSnapshot;
use crate::policies::{Policy, ProposalKind, SelectionContext, SelectionProposal};
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadSnapshot;
use crate::workers::Worker;
use rand::seq::index::sample;
use rand::Rng;
@@ -81,7 +81,7 @@ impl Policy for PowerOfTwoChoicesPolicy {
pub(crate) fn select_k_with_snapshot(
workers: &[Arc<Worker>],
snapshot: Option<&EngineLoadSnapshot>,
snapshot: Option<&EngineReportedLoadSnapshot>,
choices: usize,
queue_limit: Option<u64>,
) -> Option<Arc<Worker>> {
@@ -95,7 +95,7 @@ pub(crate) fn select_k_with_snapshot(
/// a queueing worker while an unqueued one exists.
fn sample_pool<'w>(
workers: &'w [Arc<Worker>],
snapshot: Option<&EngineLoadSnapshot>,
snapshot: Option<&EngineReportedLoadSnapshot>,
queue_limit: Option<u64>,
) -> Cow<'w, [Arc<Worker>]> {
// Without a limit there is nothing to gate on, and without a snapshot
@@ -138,7 +138,7 @@ fn sample_pool<'w>(
/// unwinding whichever request task was selecting at the time.
fn best_two_of_sample(
pool: &[Arc<Worker>],
snapshot: Option<&EngineLoadSnapshot>,
snapshot: Option<&EngineReportedLoadSnapshot>,
choices: usize,
) -> Option<(Arc<Worker>, Option<Arc<Worker>>)> {
let len = pool.len();
@@ -186,7 +186,7 @@ fn best_two_of_sample(
mod tests {
use super::*;
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use crate::policies::engine_load::NativeCacheWorkerLoad;
use crate::state::load_monitor::engine_reported_load::EngineReportedSchedulingLoad;
use std::time::Instant;
fn worker(id: &str) -> Arc<Worker> {
@@ -202,15 +202,15 @@ mod tests {
/// Snapshot keyed on waiting depth; `waiting` sets both the queue-gate
/// reading (`num_waiting_reqs`) and the pressure ordering
/// (`num_waiting_uncached_tokens`), so one knob drives both.
fn snapshot(entries: &[(&Arc<Worker>, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
7,
entries
.iter()
.map(|(worker, waiting)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: 0,
num_waiting_reqs: *waiting,
num_waiting_uncached_tokens: *waiting,
@@ -234,8 +234,8 @@ mod tests {
/// `compare_prefill_pressure` intransitive: a slow worker with a shallow
/// queue loses to a fast worker with a deep one on the estimate, while
/// both are ordered against an estimate-less worker on waiting tokens.
fn mixed_estimate_snapshot(workers: &[Arc<Worker>]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn mixed_estimate_snapshot(workers: &[Arc<Worker>]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
11,
workers
.iter()
@@ -244,7 +244,7 @@ mod tests {
let waiting = (index as u64 * 7) % 13;
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_waiting_uncached_tokens: waiting,
experimental/sgl-router/src/policies/prefix_provider.rs
+1 -1
View File
@@ -2,7 +2,7 @@
// SPDX-License-Identifier: Apache-2.0
use super::ExternalPrefixSignal;
use crate::policies::kv_events::{
use crate::state::kv_events::{
compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree,
};
use sgl_kv_indexer::{PrefixMatch, PrefixOutcome};
experimental/sgl-router/src/policies/registry.rs
+8 -5
View File
@@ -211,7 +211,7 @@ impl PdPoolResolver {
///
/// 1. **Same-host preference.** Parse the host portion of both URLs
/// (`url::Url::host_str`). If any candidate shares the host AND has
/// a closed circuit breaker AND has `active_load <=
/// a closed circuit breaker AND has `router_inflight_load <=
/// AFFINITY_LOAD_TOLERANCE × median(decode_pool_load)`, return it.
/// 2. **Fallback: min-load among closed-breaker candidates.** No
/// same-host peer, or the same-host peer was filtered by rule 1's
@@ -257,7 +257,7 @@ pub fn select_decode_with_affinity(
let load_tolerance = if healthy.is_empty() {
0
} else {
let mut loads: Vec<usize> = healthy.iter().map(|w| w.active_load()).collect();
let mut loads: Vec<usize> = healthy.iter().map(|w| w.router_inflight_load()).collect();
loads.sort_unstable();
let median = loads[loads.len() / 2];
((median as f64) * AFFINITY_LOAD_TOLERANCE).ceil() as usize
@@ -267,7 +267,7 @@ pub fn select_decode_with_affinity(
if let Some(host) = prefill_host.as_deref() {
let affinity_peer = healthy.iter().find(|w| {
host_of(&w.url).as_deref() == Some(host)
&& (load_tolerance == 0 || w.active_load() <= load_tolerance)
&& (load_tolerance == 0 || w.router_inflight_load() <= load_tolerance)
});
if let Some(w) = affinity_peer {
return Some(Arc::clone(w));
@@ -275,14 +275,17 @@ pub fn select_decode_with_affinity(
}
// Rule 2: min-load among healthy.
if let Some(w) = healthy.iter().min_by_key(|w| w.active_load()) {
if let Some(w) = healthy.iter().min_by_key(|w| w.router_inflight_load()) {
return Some(Arc::clone(w));
}
// Rule 3: last-resort min-load over all candidates (every
// breaker is open). The caller's dispatch will likely fail and
// surface `BreakerOpen`, but the selection function stays total.
candidates.iter().min_by_key(|w| w.active_load()).cloned()
candidates
.iter()
.min_by_key(|w| w.router_inflight_load())
.cloned()
}
/// Parse the host portion of a worker URL. Returns `None` when the URL
experimental/sgl-router/src/policies/scoring/admission.rs
+2 -2
View File
@@ -23,7 +23,7 @@ impl Overloaded {
impl EligibilityFilter for Overloaded {
fn keep(&self, workers: &[Arc<Worker>], _ctx: &SelectionContext<'_>) -> Vec<bool> {
(workers.iter())
.map(|w| w.active_load() < self.max_in_flight)
.map(|w| w.router_inflight_load() < self.max_in_flight)
.collect()
}
@@ -42,7 +42,7 @@ impl Policy for Overloaded {
.collect();
eligible
.iter()
.min_by_key(|w| w.active_load())
.min_by_key(|w| w.router_inflight_load())
.map(Arc::clone)
}
experimental/sgl-router/src/policies/scoring/argmax.rs
+5 -2
View File
@@ -52,10 +52,13 @@ impl Selector for Argmax {
);
band = (0..workers.len()).collect();
}
let min_load = band.iter().map(|&i| workers[i].active_load()).min()?;
let min_load = band
.iter()
.map(|&i| workers[i].router_inflight_load())
.min()?;
let tied: Vec<usize> = band
.into_iter()
.filter(|&i| workers[i].active_load() == min_load)
.filter(|&i| workers[i].router_inflight_load() == min_load)
.collect();
let k = self.rotor.fetch_add(1, Ordering::Relaxed) % tied.len();
Some(tied[k])
experimental/sgl-router/src/policies/scoring/mod.rs
+7 -5
View File
@@ -417,11 +417,13 @@ mod tests {
use crate::config::AffinityConfig;
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use crate::policies::admission::{resolve_prefill, CandidateRange};
use crate::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
use crate::policies::load_based::LoadBasedPolicy;
use crate::policies::power_of_two::PowerOfTwoChoicesPolicy;
use crate::policies::round_robin::RoundRobinPolicy;
use crate::policies::session_aware::SessionAwarePolicy;
use crate::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad,
};
use std::collections::HashMap;
use std::time::Instant;
@@ -439,15 +441,15 @@ mod tests {
vec![worker("a"), worker("b"), worker("c")]
}
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
1,
entries
.iter()
.map(|(worker, running, waiting, used, capacity)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: *running,
num_waiting_reqs: *waiting,
num_waiting_uncached_tokens: *waiting,
@@ -789,7 +791,7 @@ mod tests {
);
assert_eq!(proposal.primary.id, ws[2].id);
let snapshot = EngineLoadSnapshot::default();
let snapshot = EngineReportedLoadSnapshot::default();
let decision = resolve_prefill(
&CandidateRange::global(&ws),
&proposal,
experimental/sgl-router/src/policies/scoring/prefix_cache.rs
+3 -3
View File
@@ -4,10 +4,10 @@
//! Prefix-cache scores from the KV-event [`HashTree`].
use super::{EligibilityFilter, ScoringPolicy};
use crate::policies::kv_events::{
use crate::policies::SelectionContext;
use crate::state::kv_events::{
compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree,
};
use crate::policies::SelectionContext;
use crate::workers::Worker;
use std::sync::Arc;
@@ -122,7 +122,7 @@ impl EligibilityFilter for PrefixCachePolicy {
mod tests {
use super::*;
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use crate::policies::kv_events::KvWorkerId;
use crate::state::kv_events::KvWorkerId;
const BLOCK: usize = 4;
experimental/sgl-router/src/policies/selection.rs
+14 -12
View File
@@ -35,11 +35,11 @@ use crate::policies::buckets::{BucketRequest, BucketSelector};
use crate::policies::decode::{
build_decode_policy, resolve_decode_with_capacity_fallback, DecodeSelectionContext,
};
use crate::policies::engine_load::EngineLoadSnapshot;
use crate::policies::{
ExternalPrefixSignal, Policy, PrefillProposal, ProposalKind, SelectionContext,
};
use crate::server::metrics::{CacheAwareDecision, MetricsRegistry, PolicySelectionFailureReason};
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadSnapshot;
use crate::workers::Worker;
/// Everything one prefill selection reads. Collaborators first, then the
@@ -61,7 +61,7 @@ pub(crate) struct PrefillSelectionInputs<'a> {
/// per-domain rung panics without it. `Policy::needs_load_snapshot`
/// defaults to `uses_shared_prefill_admission`, which is what keeps the
/// two in step for the ingress caller.
pub load_snapshot: Option<&'a EngineLoadSnapshot>,
pub load_snapshot: Option<&'a EngineReportedLoadSnapshot>,
pub workers: &'a [Arc<Worker>],
pub ttft_slo_ms: Option<u64>,
pub tps_slo: Option<f64>,
@@ -607,7 +607,7 @@ pub(crate) struct DecodeSelectionInputs<'a> {
/// Required: every rung resolves its proposal against the snapshot, so
/// without one the ladder reports no peer at all rather than picking one
/// blind.
pub load_snapshot: Option<&'a EngineLoadSnapshot>,
pub load_snapshot: Option<&'a EngineReportedLoadSnapshot>,
}
/// Runs the decode selection ladder.
@@ -699,12 +699,14 @@ mod tests {
use crate::policies::admission::{resolve_prefill_admitted, CandidateRange, DecisionReason};
use crate::policies::buckets::BucketSelector;
use crate::policies::cache_aware::CacheAwarePolicy;
use crate::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
use crate::policies::power_of_two::PowerOfTwoChoicesPolicy;
use crate::policies::{ExternalPrefixSignal, Policy, ProposalKind, SelectionProposal};
use crate::server::metrics::{
CacheAwareDecision, MetricsRegistry, PolicySelectionFailureReason,
};
use crate::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad,
};
use crate::workers::Worker;
use std::sync::Arc;
use std::time::Instant;
@@ -720,15 +722,15 @@ mod tests {
}
/// `(worker, tokens already held, published KV capacity)`.
fn snapshot(entries: &[(&Arc<Worker>, u64, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, u64, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
7,
entries
.iter()
.map(|(worker, used, capacity)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_waiting_uncached_tokens: 0,
@@ -749,15 +751,15 @@ mod tests {
/// `(worker, waiting requests, tokens already held, published KV
/// capacity)`. The queue gate reads `num_waiting_reqs`, which the plain
/// [`snapshot`] fixture pins at zero.
fn queued_snapshot(entries: &[(&Arc<Worker>, u64, u64, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn queued_snapshot(entries: &[(&Arc<Worker>, u64, u64, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
7,
entries
.iter()
.map(|(worker, waiting, used, capacity)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: 1,
num_waiting_reqs: *waiting,
num_waiting_uncached_tokens: *waiting,
@@ -811,7 +813,7 @@ mod tests {
metrics: &'a MetricsRegistry,
model_id: &'a ModelId,
workers: &'a [Arc<Worker>],
load_snapshot: Option<&'a EngineLoadSnapshot>,
load_snapshot: Option<&'a EngineReportedLoadSnapshot>,
request_input_tokens: u64,
) -> PrefillSelectionInputs<'a> {
PrefillSelectionInputs {
@@ -841,7 +843,7 @@ mod tests {
bucket_selector: &'a BucketSelector,
model_id: &'a ModelId,
decode_workers: &'a [Arc<Worker>],
load_snapshot: Option<&'a EngineLoadSnapshot>,
load_snapshot: Option<&'a EngineReportedLoadSnapshot>,
request_input_tokens: u64,
) -> DecodeSelectionInputs<'a> {
DecodeSelectionInputs {
experimental/sgl-router/src/policies/session_aware.rs
+23 -78
View File
@@ -5,87 +5,43 @@
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::{GuardHints, Policy, ProposalKind, SelectionContext, SelectionProposal};
use crate::state::load_monitor::router_inflight_load::JanitorHandle;
use crate::state::AffinityStore;
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
}
}
use std::time::Duration;
pub struct SessionAwarePolicy {
state: Arc<SessionState>,
store: Arc<AffinityStore>,
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
};
let store = AffinityStore::new(Duration::from_secs(config.session_idle_secs));
let _janitor =
store.spawn_sweeper(Duration::from_secs(config.session_eviction_interval_secs));
Self {
state,
store,
config,
_janitor,
}
}
#[cfg(test)]
fn with_clock(config: AffinityConfig, clock: Arc<dyn Clock>) -> Self {
fn with_clock(
config: AffinityConfig,
clock: Arc<dyn crate::state::load_monitor::router_inflight_load::Clock>,
) -> Self {
Self {
state: Arc::new(SessionState {
assignments: DashMap::new(),
clock,
idle: Duration::from_secs(config.session_idle_secs),
}),
store: AffinityStore::with_clock(Duration::from_secs(config.session_idle_secs), clock),
config,
_janitor: None,
}
@@ -93,12 +49,12 @@ impl SessionAwarePolicy {
#[cfg(test)]
fn sweep_expired(&self) -> usize {
self.state.sweep_expired()
self.store.sweep_expired()
}
#[cfg(test)]
fn assignment_count(&self) -> usize {
self.state.assignments.len()
self.store.len()
}
fn assignment_key(&self, session_id: &str, ctx: &SelectionContext<'_>) -> String {
@@ -173,18 +129,8 @@ impl Policy for SessionAwarePolicy {
};
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));
}
if let Some(primary) = self.store.bound(&assignment_key, workers) {
return Some(self.affinity_proposal(Arc::clone(primary), workers, ctx, session_id));
}
// Persist new assignments only after selecting the final prefill worker.
@@ -203,12 +149,11 @@ impl Policy for SessionAwarePolicy {
let Some(session_id) = ctx.session_id().filter(|id| !id.is_empty()) else {
return;
};
self.state.assignments.insert(
// Overwrites any previous binding for the key.
self.store.bind(
self.assignment_key(session_id, ctx),
Assignment {
worker_id: selected.id.clone(),
last_seen: self.state.clock.now(),
},
selected,
std::slice::from_ref(selected),
);
}
@@ -225,7 +170,7 @@ 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())
.field("assignments", &self.store.len())
.finish_non_exhaustive()
}
}
@@ -308,7 +253,7 @@ fn stable_backup(
mod lifecycle_tests {
use super::*;
use crate::discovery::{ModelId, WorkerMode, WorkerSpec};
use crate::policies::active_load::MockClock;
use crate::state::load_monitor::router_inflight_load::MockClock;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
experimental/sgl-router/src/policies/sticky.rs
+67 -151
View File
@@ -17,11 +17,11 @@
//! - **New key, or pinned worker unhealthy** → pick a worker via `fallback`
//! and record the assignment.
//!
//! Worker identity is the worker URL (stable across discovery events).
//! Worker identity is the worker ID supplied by discovery.
//!
//! # Eviction
//! A background sweeper (shared engine with the active-load janitor, see
//! [`super::active_load::spawn_sweeper`]) removes assignments idle longer
//! [`crate::state::load_monitor::router_inflight_load::spawn_sweeper`]) removes assignments idle longer
//! than `idle`, bounding the map against unbounded routing-key cardinality.
//! The sweeper is spawned only when constructed inside a Tokio runtime;
//! unit tests use [`StickyPolicy::with_clock`] and drive eviction
@@ -34,49 +34,60 @@
//! is intentionally out of scope here.
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use std::time::Duration;
use dashmap::DashMap;
use crate::policies::active_load::{spawn_sweeper, Clock, JanitorHandle, SystemTimeClock};
use crate::policies::{Policy, SelectionContext};
use crate::server::metrics::{MetricsRegistry, StickyOutcome};
use crate::state::load_monitor::router_inflight_load::JanitorHandle;
use crate::state::AffinityStore;
use crate::workers::Worker;
/// One routing-key → worker pin, with the last time it was referenced (used
/// by the idle-eviction sweep).
#[derive(Debug)]
struct Assignment {
worker_url: String,
last_seen: Instant,
}
/// Shared inner state. Held behind an `Arc` so the background sweeper can
/// reference the same map the `select` hot path mutates.
#[derive(Debug)]
struct StickyState {
assignments: DashMap<String, Assignment>,
clock: Arc<dyn Clock>,
idle: Duration,
/// Metrics sink. Set once via the `Policy::attach_metrics` hook
/// (production) — `None` until then, in which case recording is a no-op.
/// Sticky-session policy. See the module docs for behavior and limitations.
pub struct StickyPolicy {
store: Arc<AffinityStore>,
/// Selector for keyless requests and for the initial pin of a new key.
fallback: Arc<dyn Policy>,
/// Set once via `Policy::attach_metrics`; recording is a no-op until then.
metrics: OnceLock<Arc<MetricsRegistry>>,
/// Background idle-eviction sweeper; `None` outside a Tokio runtime.
_janitor: Option<JanitorHandle>,
}
impl StickyState {
/// Remove every assignment idle longer than `idle`. Returns the count
/// removed. Called on a fixed cadence by the background sweeper.
impl StickyPolicy {
pub fn new(idle: Duration, eviction_interval: Duration, fallback: Arc<dyn Policy>) -> Self {
let store = AffinityStore::new(idle);
let _janitor = store.spawn_sweeper(eviction_interval);
Self {
store,
fallback,
metrics: OnceLock::new(),
_janitor,
}
}
/// Test constructor: injectable clock, no background sweeper.
#[cfg(test)]
fn with_clock(
idle: Duration,
fallback: Arc<dyn Policy>,
clock: Arc<dyn crate::state::load_monitor::router_inflight_load::Clock>,
) -> Self {
Self {
store: AffinityStore::with_clock(idle, clock),
fallback,
metrics: OnceLock::new(),
_janitor: None,
}
}
#[cfg(test)]
fn sweep_expired(&self) -> usize {
let now = self.clock.now();
let mut removed = 0;
self.assignments.retain(|_key, a| {
let keep = now.saturating_duration_since(a.last_seen) <= self.idle;
if !keep {
removed += 1;
}
keep
});
removed
self.store.sweep_expired()
}
#[cfg(test)]
fn assignment_count(&self) -> usize {
self.store.len()
}
fn record(&self, outcome: StickyOutcome) {
@@ -86,122 +97,21 @@ impl StickyState {
}
}
/// Sticky-session policy. See the module docs for behavior and limitations.
pub struct StickyPolicy {
state: Arc<StickyState>,
/// Selector for keyless requests and for the initial pin of a new key.
fallback: Arc<dyn Policy>,
/// Background idle-eviction sweeper. `None` when constructed outside a
/// Tokio runtime (unit tests). Dropping it cancels the task, so the
/// sweeper lives exactly as long as the policy.
_janitor: Option<JanitorHandle>,
}
impl StickyPolicy {
/// Production constructor: monotonic `SystemTimeClock`, with a
/// background eviction sweeper spawned on `eviction_interval` cadence
/// (only if called inside a Tokio runtime — the factory runs inside
/// `main`'s runtime).
pub fn new(idle: Duration, eviction_interval: Duration, fallback: Arc<dyn Policy>) -> Self {
let state = Arc::new(StickyState {
assignments: DashMap::new(),
clock: Arc::new(SystemTimeClock),
idle,
metrics: OnceLock::new(),
});
// `spawn_sweeper` needs a runtime; the factory builds policies inside
// `main`'s Tokio runtime. Guard so sync constructions (e.g. the
// factory's `build_policy_kind_only` test helper) don't panic.
let _janitor = if tokio::runtime::Handle::try_current().is_ok() {
let swept = Arc::clone(&state);
Some(spawn_sweeper(
move || swept.sweep_expired(),
eviction_interval,
"sticky-eviction",
))
} else {
// Only reached by sync construction (test helpers). In production
// the factory builds policies inside `main`'s runtime, so the
// sweeper always spawns. Log it so a future off-runtime
// construction that silently disables eviction is greppable.
tracing::debug!(
"StickyPolicy constructed outside a Tokio runtime; idle eviction is disabled"
);
None
};
Self {
state,
fallback,
_janitor,
}
}
/// Test constructor: injectable clock, no background sweeper. Tests
/// advance a `MockClock` and call [`Self::sweep_expired`] directly for
/// deterministic eviction coverage.
#[cfg(test)]
fn with_clock(idle: Duration, fallback: Arc<dyn Policy>, clock: Arc<dyn Clock>) -> Self {
Self {
state: Arc::new(StickyState {
assignments: DashMap::new(),
clock,
idle,
metrics: OnceLock::new(),
}),
fallback,
_janitor: None,
}
}
#[cfg(test)]
fn sweep_expired(&self) -> usize {
self.state.sweep_expired()
}
#[cfg(test)]
fn assignment_count(&self) -> usize {
self.state.assignments.len()
}
}
impl Policy for StickyPolicy {
fn select(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
let Some(key) = ctx.routing_key().filter(|k| !k.is_empty()) else {
self.state.record(StickyOutcome::NoRoutingKey);
self.record(StickyOutcome::NoRoutingKey);
return self.fallback.select(workers, ctx);
};
// Fast path: an existing pin whose worker is still in the healthy set.
let mut existing = false;
if let Some(mut entry) = self.state.assignments.get_mut(key) {
existing = true;
if let Some(worker) = workers.iter().find(|w| w.url == entry.worker_url).cloned() {
entry.last_seen = self.state.clock.now();
drop(entry); // release the shard lock before recording
self.state.record(StickyOutcome::Hit);
return Some(worker);
}
// Pinned worker is no longer healthy — fall through to reassign.
drop(entry);
if let Some(worker) = self.store.bound(key, workers) {
self.record(StickyOutcome::Hit);
return Some(Arc::clone(worker));
}
// Vacant key, or the pinned worker dropped out: (re)assign via the
// fallback. The read-miss above and this insert are intentionally NOT
// atomic — the shard lock is released before `fallback.select` (which
// may do real work, e.g. `load_based`) so it is never held across an
// unrelated computation. Two requests racing the *same* fresh key may
// therefore both assign (last-writer-wins in the map; both may record
// `Assigned`). The scatter is transient and self-heals: the next
// request for that key hits the surviving pin.
// Vacant key, or the pinned worker dropped out: (re)pin the fallback's choice.
let remap = self.store.contains(key);
let chosen = self.fallback.select(workers, ctx)?;
self.state.assignments.insert(
key.to_string(),
Assignment {
worker_url: chosen.url.clone(),
last_seen: self.state.clock.now(),
},
);
self.state.record(if existing {
let chosen = Arc::clone(self.store.bind(key.to_string(), &chosen, workers));
self.record(if remap {
StickyOutcome::Remap
} else {
StickyOutcome::Assigned
@@ -210,7 +120,7 @@ impl Policy for StickyPolicy {
}
fn attach_metrics(&self, metrics: Arc<MetricsRegistry>) {
let _ = self.state.metrics.set(metrics);
let _ = self.metrics.set(metrics);
}
fn needs_load_snapshot(&self) -> bool {
@@ -226,8 +136,7 @@ impl std::fmt::Debug for StickyPolicy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StickyPolicy")
.field("fallback", &self.fallback)
.field("idle", &self.state.idle)
.field("assignments", &self.state.assignments.len())
.field("assignments", &self.store.len())
.finish_non_exhaustive()
}
}
@@ -236,6 +145,7 @@ impl std::fmt::Debug for StickyPolicy {
mod tests {
use super::*;
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use std::time::Instant;
#[test]
fn sticky_propagates_fallback_load_snapshot_capability() {
@@ -271,7 +181,9 @@ mod tests {
}
fn policy(idle_secs: u64) -> StickyPolicy {
let clock = Arc::new(crate::policies::active_load::MockClock::new(Instant::now()));
let clock = Arc::new(
crate::state::load_monitor::router_inflight_load::MockClock::new(Instant::now()),
);
StickyPolicy::with_clock(Duration::from_secs(idle_secs), fallback(), clock)
}
@@ -374,7 +286,9 @@ mod tests {
#[test]
fn sweep_evicts_idle_entries_keeps_fresh() {
let model = ModelId("tiny".into());
let clock = Arc::new(crate::policies::active_load::MockClock::new(Instant::now()));
let clock = Arc::new(
crate::state::load_monitor::router_inflight_load::MockClock::new(Instant::now()),
);
let p = StickyPolicy::with_clock(Duration::from_secs(10), fallback(), clock.clone());
let workers = vec![worker("w0"), worker("w1")];
@@ -401,7 +315,9 @@ mod tests {
#[test]
fn hit_refreshes_last_seen_so_active_key_is_not_evicted() {
let model = ModelId("tiny".into());
let clock = Arc::new(crate::policies::active_load::MockClock::new(Instant::now()));
let clock = Arc::new(
crate::state::load_monitor::router_inflight_load::MockClock::new(Instant::now()),
);
let p = StickyPolicy::with_clock(Duration::from_secs(10), fallback(), clock.clone());
let workers = vec![worker("w0")];
let ctx = SelectionContext::with_routing_key(&model, None, Some("u1"));
experimental/sgl-router/src/proxy/mod.rs
+1 -1
View File
@@ -262,7 +262,7 @@ impl Proxy {
/// pump task and held for the entire body lifetime (headers → last byte
/// / client disconnect). The proxy does not inspect the boxed value; it
/// relies entirely on `Drop` semantics, so callers typically pack
/// `(LoadGuard, ActiveLoadGuard)` here. This keeps both the per-worker
/// `(LoadGuard, RouterInflightLoadGuard)` here. This keeps both the per-worker
/// `active_requests` counter and the per-request active-load entry alive
/// for the full streaming lifetime — without which a long-running SSE
/// response would under-report load.
experimental/sgl-router/src/proxy/sse.rs
+1 -1
View File
@@ -88,7 +88,7 @@ impl ErrorEventScanner {
/// pump finishes (stream exhausted, client disconnects, or upstream errors).
/// The opaque `Box<dyn Send + 'static>` accepts any drop-only payload — most
/// commonly a tuple of [`crate::workers::LoadGuard`] and
/// [`crate::policies::active_load::ActiveLoadGuard`]. The proxy does not
/// [`crate::state::load_monitor::router_inflight_load::RouterInflightLoadGuard`]. The proxy does not
/// inspect the value; it relies entirely on `Drop` semantics, so callers can
/// pack arbitrary cleanup state in. Pass `None` for callers that manage the
/// guard externally (e.g. non-streaming paths where the handler itself is the
experimental/sgl-router/src/server/app_context.rs
+18 -18
View File
@@ -3,15 +3,15 @@
use crate::config::Config;
use crate::policies::active_load::ActiveLoadRegistry;
use crate::policies::buckets::BucketSelector;
use crate::policies::engine_load::EngineLoadTable;
use crate::policies::kv_events::{BlockSizeOracle, KvIndexMetrics};
use crate::policies::prefix_provider::RadixTreePrefixProvider;
use crate::policies::PolicyRegistry;
use crate::proxy::Proxy;
use crate::server::inflight::InflightHttp;
use crate::server::metrics::MetricsRegistry;
use crate::state::kv_events::{BlockSizeOracle, KvIndexMetrics};
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
use crate::state::load_monitor::router_inflight_load::RouterInflightLoadRegistry;
use crate::tokenizer::TokenizerRegistry;
use crate::workers::WorkerRegistry;
use std::sync::atomic::{AtomicU8, Ordering};
@@ -36,24 +36,24 @@ pub struct AppContext {
pub bucket_selector: Arc<BucketSelector>,
/// Per-worker active-load bookkeeping shared by the proxy, policies,
/// timeout janitor, and metrics.
pub active_load: Arc<ActiveLoadRegistry>,
pub router_inflight_load: Arc<RouterInflightLoadRegistry>,
/// Lightweight Prometheus-format metrics registry served via
/// `/metrics`. Shared with the edge middleware (requests_total /
/// responses_total), the chat handler (worker_requests_total), the
/// active-load registry, policy-specific counters, and PD dispatch.
pub metrics: Arc<MetricsRegistry>,
/// Shared Engine LoadStat table; ingress captures one immutable snapshot per request.
pub engine_load: Arc<EngineLoadTable>,
pub engine_reported_load: Arc<EngineReportedLoadTable>,
pub prefix_index: Option<Arc<dyn sgl_kv_indexer::PrefixIndex>>,
pub radix_tree_prefix_provider: Option<RadixTreePrefixProvider>,
pub block_size_oracle: Arc<BlockSizeOracle>,
/// Read-only handles `/metrics` pulls the KV storage-tier series from on
/// scrape. `None` when this router maintains no local tree (external
/// Indexer), where those series would all be a structural zero — see
/// [`crate::policies::kv_events::KvEventIndex::metrics_source`].
/// [`crate::state::kv_events::KvEventIndex::metrics_source`].
pub kv_metrics: Option<KvIndexMetrics>,
/// Open HTTP exchanges, on every route. What axum's graceful shutdown
/// is actually waiting on during the drain — `active_load` sees only the
/// is actually waiting on during the drain — `router_inflight_load` sees only the
/// proxied subset.
pub inflight_http: Arc<InflightHttp>,
readiness: AtomicU8,
@@ -67,34 +67,34 @@ impl AppContext {
registry: Arc<WorkerRegistry>,
policies: Arc<PolicyRegistry>,
) -> Self {
Self::with_active_load(
Self::with_router_inflight_load(
config,
tokenizers,
proxy,
registry,
policies,
ActiveLoadRegistry::with_defaults(),
RouterInflightLoadRegistry::with_defaults(),
)
}
/// Construct an [`AppContext`] with an explicit [`ActiveLoadRegistry`].
/// Construct an [`AppContext`] with an explicit [`RouterInflightLoadRegistry`].
/// Production wires the default (5-minute timeout, SystemTimeClock)
/// via [`Self::new`]; tests that exercise the janitor pass a registry
/// built with a `MockClock`.
pub fn with_active_load(
pub fn with_router_inflight_load(
config: Config,
tokenizers: Arc<TokenizerRegistry>,
proxy: Arc<Proxy>,
registry: Arc<WorkerRegistry>,
policies: Arc<PolicyRegistry>,
active_load: Arc<ActiveLoadRegistry>,
router_inflight_load: Arc<RouterInflightLoadRegistry>,
) -> Self {
let metrics = MetricsRegistry::new();
// Wire the per-worker active-load gauge so `sgl_router_active_load`
// mirrors the live counter on every register / drop / sweep.
// Without this, the metric is permanently 0 in production even
// though the chat handler is faithfully calling `register`.
active_load.attach_metrics(Arc::clone(&metrics));
router_inflight_load.attach_metrics(Arc::clone(&metrics));
// The metrics registry is built after the policy registry, so attach
// it here for policies that emit their own counters.
policies.attach_metrics(Arc::clone(&metrics));
@@ -106,13 +106,13 @@ impl AppContext {
registry,
policies,
bucket_selector,
active_load,
router_inflight_load,
metrics,
prefix_index: None,
radix_tree_prefix_provider: None,
block_size_oracle: BlockSizeOracle::new(),
kv_metrics: None,
engine_load: EngineLoadTable::new(),
engine_reported_load: EngineReportedLoadTable::new(),
inflight_http: InflightHttp::new(),
readiness: AtomicU8::new(READINESS_NOT_READY),
}
@@ -188,20 +188,20 @@ impl AppContext {
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
router_inflight_load: crate::config::InflightLoadConfig::default(),
},
tokenizers: Arc::new(TokenizerRegistry::default()),
proxy: Arc::new(Proxy::new(std::time::Duration::from_secs(60)).expect("stub proxy")),
registry: Arc::new(WorkerRegistry::default()),
policies: Arc::new(PolicyRegistry::default()),
bucket_selector: Arc::new(BucketSelector::new(None)),
active_load: ActiveLoadRegistry::with_defaults(),
router_inflight_load: RouterInflightLoadRegistry::with_defaults(),
metrics: MetricsRegistry::new(),
prefix_index: None,
radix_tree_prefix_provider: None,
block_size_oracle: BlockSizeOracle::new(),
kv_metrics: None,
engine_load: EngineLoadTable::new(),
engine_reported_load: EngineReportedLoadTable::new(),
inflight_http: InflightHttp::new(),
readiness: AtomicU8::new(READINESS_NOT_READY),
}
experimental/sgl-router/src/server/inflight.rs
+1 -1
View File
@@ -3,7 +3,7 @@
//! In-flight HTTP accounting, for the termination drain to report on.
//!
//! [`ActiveLoadRegistry`](crate::policies::active_load::ActiveLoadRegistry)
//! [`RouterInflightLoadRegistry`](crate::state::load_monitor::router_inflight_load::RouterInflightLoadRegistry)
//! counts *proxied* requests — what the workers are busy with. Axum's graceful
//! shutdown waits on something different and larger: every HTTP exchange still
//! open on an accepted connection, on any route, until its response body has
experimental/sgl-router/src/server/metrics.rs
+18 -13
View File
@@ -375,12 +375,12 @@ impl PolicySelectionFailureReason {
/// Active-load kind label — separates the two axes of per-worker load.
#[derive(Debug, Clone, Copy)]
pub enum ActiveLoadKind {
pub enum RouterInflightLoadKind {
PrefillTokens,
DecodeBlocks,
}
impl ActiveLoadKind {
impl RouterInflightLoadKind {
fn as_str(self) -> &'static str {
match self {
Self::PrefillTokens => "prefill_tokens",
@@ -408,7 +408,7 @@ pub struct MetricsRegistry {
request_duration: Mutex<HashMap<String, Histogram>>,
ttft_seconds: Mutex<HashMap<String, Histogram>>,
stream_outcome_total: Mutex<HashMap<StreamOutcomeKey, Arc<AtomicU64>>>,
active_load: Mutex<HashMap<ActiveLoadKey, Arc<AtomicI64>>>,
router_inflight_load: Mutex<HashMap<RouterInflightLoadKey, Arc<AtomicI64>>>,
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>>>,
@@ -470,12 +470,12 @@ pub struct WorkerSnapshot {
pub healthy: bool,
/// Circuit breaker state code: 0=closed, 1=open, 2=half_open.
pub cb_state: u8,
/// In-flight request count for this worker (`Worker::active_load`).
/// In-flight request count for this worker (`Worker::router_inflight_load`).
pub inflight: i64,
}
#[derive(Debug, Hash, Eq, PartialEq, Clone)]
struct ActiveLoadKey {
struct RouterInflightLoadKey {
worker_url: String,
kind: &'static str,
}
@@ -657,12 +657,17 @@ impl MetricsRegistry {
/// Set `sgl_router_active_load` for the given worker + kind. Replaces the
/// previous value (gauge semantics).
pub fn set_active_load(&self, worker_url: &str, kind: ActiveLoadKind, value: i64) {
let key = ActiveLoadKey {
pub fn set_router_inflight_load(
&self,
worker_url: &str,
kind: RouterInflightLoadKind,
value: i64,
) {
let key = RouterInflightLoadKey {
worker_url: worker_url.to_owned(),
kind: kind.as_str(),
};
let mut guard = self.active_load.lock();
let mut guard = self.router_inflight_load.lock();
let gauge = guard
.entry(key)
.or_insert_with(|| Arc::new(AtomicI64::new(0)))
@@ -984,13 +989,13 @@ impl MetricsRegistry {
}
drop(guard);
// active_load gauge
// router_inflight_load gauge
out.push_str(
"# HELP sgl_router_active_load Per-worker active load (prefill_tokens or decode_blocks).\n",
);
out.push_str("# TYPE sgl_router_active_load gauge\n");
let guard = self.active_load.lock();
let mut entries: Vec<(&ActiveLoadKey, i64)> = guard
let guard = self.router_inflight_load.lock();
let mut entries: Vec<(&RouterInflightLoadKey, i64)> = guard
.iter()
.map(|(k, v)| (k, v.load(Ordering::Relaxed)))
.collect();
@@ -1651,8 +1656,8 @@ mod tests {
#[test]
fn set_active_load_gauge_overwrites() {
let reg = MetricsRegistry::new();
reg.set_active_load("http://w:30000", ActiveLoadKind::PrefillTokens, 100);
reg.set_active_load("http://w:30000", ActiveLoadKind::PrefillTokens, 250);
reg.set_router_inflight_load("http://w:30000", RouterInflightLoadKind::PrefillTokens, 100);
reg.set_router_inflight_load("http://w:30000", RouterInflightLoadKind::PrefillTokens, 250);
let out = reg.render();
assert!(out.contains(
r#"sgl_router_active_load{worker_url="http://w:30000",kind="prefill_tokens"} 250"#,
experimental/sgl-router/src/server/routes/chat.rs
+5 -5
View File
@@ -6,8 +6,6 @@ mod preparation;
use crate::config::{SessionAffinityMode, DEFAULT_MIN_LOAD_CHOICES};
use crate::discovery::{ModelId, WorkerMode};
use crate::policies::engine_load::EngineLoadSnapshot;
use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram};
use crate::policies::registry::{PdPoolResolver, PdResolveError};
use crate::policies::selection::{
select_decode_peer, select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs,
@@ -16,6 +14,8 @@ use crate::policies::{ExternalPrefixSignal, Policy};
use crate::server::app_context::AppContext;
use crate::server::error::ApiError;
use crate::server::metrics::PolicySelectionFailureReason;
use crate::state::kv_events::{compute_block_hashes, compute_block_hashes_bigram};
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadSnapshot;
use crate::workers::Worker;
use axum::body::Body;
use axum::extract::State;
@@ -113,17 +113,17 @@ fn capture_load_snapshot(
ctx: &AppContext,
policy: &dyn Policy,
candidates: &[Arc<Worker>],
) -> Option<EngineLoadSnapshot> {
) -> Option<EngineReportedLoadSnapshot> {
let needed = policy.needs_load_snapshot()
|| candidates
.iter()
.any(|worker| worker.mode() == WorkerMode::Prefill);
needed.then(|| ctx.engine_load.capture_snapshot(Instant::now()))
needed.then(|| ctx.engine_reported_load.capture_snapshot(Instant::now()))
}
struct RoutingContext<'a> {
prefix_matches: Option<ExternalPrefixSignal>,
load_snapshot: Option<EngineLoadSnapshot>,
load_snapshot: Option<EngineReportedLoadSnapshot>,
ttft_slo_ms: Option<u64>,
tps_slo: Option<f64>,
routing_key: Option<&'a str>,
experimental/sgl-router/src/server/routes/chat/forward.rs
+4 -4
View File
@@ -5,7 +5,6 @@
use super::preparation::{generate_room_id, BootstrapFields, PreparedChatRequest};
use crate::discovery::WorkerMode;
use crate::policies::active_load::ActiveLoadGuard;
use crate::proxy::sse::StreamEnd;
use crate::server::app_context::AppContext;
use crate::server::error::ApiError;
@@ -13,6 +12,7 @@ use crate::server::metrics::{
classify_stream_end, outcome_from_status, MetricsRegistry, RequestLogContext, RequestOutcome,
StaleRequestOutcome, WorkerModeLabel,
};
use crate::state::load_monitor::router_inflight_load::RouterInflightLoadGuard;
use crate::workers::{LoadGuard, Worker};
use axum::body::Body;
use axum::http::{HeaderMap, HeaderName, HeaderValue, Response};
@@ -24,7 +24,7 @@ use std::time::Instant;
const CHAT_PATH: &str = "/v1/chat/completions";
// Expose the selected decode worker to both PD workers and the client.
const X_SGL_DECODE_URL: HeaderName = HeaderName::from_static("x-sgl-decode-url");
type LoadGuards = (LoadGuard, ActiveLoadGuard);
type LoadGuards = (LoadGuard, RouterInflightLoadGuard);
/// A plain worker, or a prefill worker paired with a decode worker for PD.
pub(super) struct SelectedWorkers {
@@ -58,7 +58,7 @@ pub(super) async fn forward_chat_request(
} else {
prefill.load_guard()
};
let active_request_guard = ctx.active_load.register(
let active_request_guard = ctx.router_inflight_load.register(
prefill.id.clone(),
prefill.url.clone(),
request.input_token_count,
@@ -97,7 +97,7 @@ pub(super) async fn forward_chat_request(
);
let decode_load_guards = (
decode.load_guard(),
ctx.active_load
ctx.router_inflight_load
.register(decode.id.clone(), decode.url.clone(), 0, 1),
);
(decode, decode_load_guards)
experimental/sgl-router/src/server/routes/metrics.rs
+5 -5
View File
@@ -10,9 +10,9 @@
//! discovered" failure mode is observable.
use crate::discovery::WorkerMode;
use crate::policies::kv_events::{KvIndexMetrics, Tiers, ACCOUNTING_REASONS};
use crate::server::app_context::AppContext;
use crate::server::metrics::{escape_label, WorkerSnapshot};
use crate::state::kv_events::{KvIndexMetrics, Tiers, ACCOUNTING_REASONS};
use axum::extract::State;
use axum::http::header::CONTENT_TYPE;
use axum::http::StatusCode;
@@ -46,7 +46,7 @@ pub async fn metrics(State(ctx): State<Arc<AppContext>>) -> impl IntoResponse {
// Saturating rather than `as i64`: a guard-accounting
// underflow would wrap usize and render as a nonsensical
// negative gauge; clamp to a large positive ceiling instead.
inflight: i64::try_from(w.active_load()).unwrap_or(i64::MAX),
inflight: i64::try_from(w.router_inflight_load()).unwrap_or(i64::MAX),
}
})
.collect();
@@ -188,7 +188,7 @@ mod tests {
/// cell of the tally.
#[tokio::test]
async fn kv_tier_series_render_per_worker_and_per_medium() {
use crate::policies::kv_events::{EventKind, EventTally, HashTree, KvWorkerId};
use crate::state::kv_events::{EventKind, EventTally, HashTree, KvWorkerId};
let kv = KvIndexMetrics::new(Arc::new(HashTree::new()), Arc::new(EventTally::new()));
let w = KvWorkerId::new("http://w0:30000".into(), 0);
@@ -231,7 +231,7 @@ mod tests {
/// calls `clear_worker`.
#[tokio::test]
async fn kv_tree_blocks_drop_with_the_worker() {
use crate::policies::kv_events::{EventTally, HashTree, KvWorkerId};
use crate::state::kv_events::{EventTally, HashTree, KvWorkerId};
let kv = KvIndexMetrics::new(Arc::new(HashTree::new()), Arc::new(EventTally::new()));
let w = KvWorkerId::new("http://w0:30000".into(), 0);
@@ -254,7 +254,7 @@ mod tests {
/// all four families.
#[tokio::test]
async fn metrics_endpoint_emits_kv_series_when_a_tree_is_maintained() {
use crate::policies::kv_events::{EventTally, HashTree, KvWorkerId};
use crate::state::kv_events::{EventTally, HashTree, KvWorkerId};
let mut ctx = AppContext::stub();
let tree = Arc::new(HashTree::new());
experimental/sgl-router/src/server/routes/tokenize.rs
+1 -1
View File
@@ -136,7 +136,7 @@ mod tests {
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
router_inflight_load: crate::config::InflightLoadConfig::default(),
};
let registry = crate::tokenizer::TokenizerRegistry::load_from_config(&cfg).unwrap();
let proxy = Arc::new(
experimental/sgl-router/src/state/affinity_store.rs
+195
View File
@@ -0,0 +1,195 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Session and routing-key assignments with idle expiry. The store keeps
//! bindings; deciding whether to reuse, create or replace one is the policy's.
use std::sync::Arc;
use std::time::{Duration, Instant};
use dashmap::DashMap;
use super::load_monitor::router_inflight_load::{
spawn_sweeper, Clock, JanitorHandle, SystemTimeClock,
};
use crate::discovery::WorkerId;
use crate::workers::Worker;
#[derive(Debug)]
struct Assignment {
engine: WorkerId,
last_seen: Instant,
}
#[derive(Debug)]
pub struct AffinityStore {
assignments: DashMap<String, Assignment>,
clock: Arc<dyn Clock>,
idle: Duration,
}
impl AffinityStore {
pub fn new(idle: Duration) -> Arc<Self> {
Self::with_clock(idle, Arc::new(SystemTimeClock))
}
pub fn with_clock(idle: Duration, clock: Arc<dyn Clock>) -> Arc<Self> {
Arc::new(Self {
assignments: DashMap::new(),
clock,
idle,
})
}
/// The bound engine when it is still among `engines`; refreshes the binding.
pub fn bound<'e>(&self, key: &str, engines: &'e [Arc<Worker>]) -> Option<&'e Arc<Worker>> {
let mut assignment = self.assignments.get_mut(key)?;
let engine = engines
.iter()
.find(|engine| engine.id == assignment.engine)?;
assignment.last_seen = self.clock.now();
Some(engine)
}
pub fn contains(&self, key: &str) -> bool {
self.assignments.contains_key(key)
}
/// The bound engine id, whether or not it is still a candidate.
pub fn binding(&self, key: &str) -> Option<WorkerId> {
self.assignments.get(key).map(|a| a.engine.clone())
}
/// Binds `engine`, unless a concurrent pick already bound another engine
/// that is still in `engines`; that binding wins so racing first touches
/// converge on one engine.
pub fn bind<'e>(
&self,
key: String,
engine: &'e Arc<Worker>,
engines: &'e [Arc<Worker>],
) -> &'e Arc<Worker> {
let now = self.clock.now();
let mut assignment = self.assignments.entry(key).or_insert_with(|| Assignment {
engine: engine.id.clone(),
last_seen: now,
});
match engines.iter().find(|bound| bound.id == assignment.engine) {
Some(bound) => {
assignment.last_seen = now;
bound
}
None => {
assignment.engine = engine.id.clone();
assignment.last_seen = now;
engine
}
}
}
pub fn len(&self) -> usize {
self.assignments.len()
}
pub fn is_empty(&self) -> bool {
self.assignments.is_empty()
}
pub 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
}
/// Periodic eviction; `None` outside a Tokio runtime.
pub fn spawn_sweeper(self: &Arc<Self>, interval: Duration) -> Option<JanitorHandle> {
let store = Arc::clone(self);
tokio::runtime::Handle::try_current()
.ok()
.map(|_| spawn_sweeper(move || store.sweep_expired(), interval, "affinity-eviction"))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::discovery::{ModelId, WorkerMode, WorkerSpec};
use crate::state::load_monitor::router_inflight_load::MockClock;
fn worker(id: &str) -> Arc<Worker> {
Arc::new(Worker::new(WorkerSpec {
id: WorkerId(id.into()),
url: format!("http://{id}"),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("m".into())],
bootstrap_port: None,
}))
}
#[test]
fn a_live_binding_wins_over_a_later_bind_and_a_dead_one_is_replaced() {
let store = AffinityStore::new(Duration::from_secs(60));
let (a, b) = (worker("a"), worker("b"));
let fleet = [a.clone(), b.clone()];
assert_eq!(store.bind("k".into(), &a, &fleet).id.0, "a");
assert_eq!(store.bind("k".into(), &b, &fleet).id.0, "a");
assert_eq!(store.bind("k".into(), &b, &fleet[1..]).id.0, "b");
assert_eq!(store.bound("k", &fleet).unwrap().id.0, "b");
assert!(store.bound("k", &fleet[..1]).is_none());
}
#[test]
fn idle_bindings_are_swept_and_hits_refresh() {
let clock = Arc::new(MockClock::new(Instant::now()));
let store = AffinityStore::with_clock(Duration::from_secs(10), clock.clone());
let fleet = [worker("a")];
store.bind("hot".into(), &fleet[0], &fleet);
store.bind("cold".into(), &fleet[0], &fleet);
clock.advance(Duration::from_secs(8));
store.bound("hot", &fleet);
clock.advance(Duration::from_secs(8));
assert_eq!(store.sweep_expired(), 1);
assert!(store.contains("hot") && !store.contains("cold"));
}
#[test]
fn concurrent_bindings_do_not_count_as_evictions() {
use std::sync::Barrier;
let clock = Arc::new(MockClock::new(Instant::now()));
let store = AffinityStore::with_clock(Duration::from_secs(60), clock);
let engine = worker("a");
let start = Arc::new(Barrier::new(5));
std::thread::scope(|scope| {
for writer in 0..4 {
let store = Arc::clone(&store);
let engine = Arc::clone(&engine);
let start = Arc::clone(&start);
scope.spawn(move || {
start.wait();
for key in 0..5000 {
store.bind(
format!("{writer}-{key}"),
&engine,
std::slice::from_ref(&engine),
);
}
});
}
start.wait();
for _ in 0..1000 {
// The clock never advances: every binding must survive,
// even when requests insert new keys during a sweep.
assert_eq!(store.sweep_expired(), 0);
}
});
assert_eq!(store.len(), 20_000);
}
}
experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs → experimental/sgl-router/src/state/kv_events/block_size_oracle.rs
+1 -1
View File
@@ -8,7 +8,7 @@
//! `compute_block_hashes` must hash with the **same** block size the
//! worker uses to publish KV-cache events; otherwise every cache-aware
//! lookup misses silently. The worker advertises its `page_size` via
//! `/server_info` (parsed into [`crate::policies::kv_events::EventConfig::block_size`]).
//! `/server_info` (parsed into [`crate::state::kv_events::EventConfig::block_size`]).
//! Dynamo's design treats `kv_cache_block_size` as a property of the
//! `ModelDeploymentCard` populated by the worker registrar (see
//! `~/dynamo/components/src/dynamo/sglang/register.py`); a mismatch
experimental/sgl-router/src/policies/kv_events/discovery.rs → experimental/sgl-router/src/state/kv_events/discovery.rs
View File
experimental/sgl-router/src/policies/kv_events/hash.rs → experimental/sgl-router/src/state/kv_events/hash.rs
View File
experimental/sgl-router/src/policies/kv_events/index.rs → experimental/sgl-router/src/state/kv_events/index.rs
+29 -28
View File
@@ -7,7 +7,7 @@
//! always operate together in production:
//!
//! - [`HashTree`] — the cache-aware routing index keyed by SGLang block hash.
//! - [`EngineLoadTable`] — engine-reported per-worker load.
//! - [`EngineReportedLoadTable`] — 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
@@ -41,7 +41,7 @@ use super::subscriber::{KvEventSubscriberRegistry, SubKind, WorkerEvent};
use super::tally::{EventKind, EventTally};
use super::tree::{HashTree, KvWorkerId, Tiers};
use super::wire::KvCacheEvent;
use crate::policies::engine_load::EngineLoadTable;
use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
/// Channel buffer between the subscriber registry and the pump task.
///
@@ -96,13 +96,13 @@ pub struct KvEventIndex {
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
/// feeding `LoadStat` snapshots into `engine_reported_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 captured at request ingress.
engine_load: Arc<EngineLoadTable>,
engine_reported_load: Arc<EngineReportedLoadTable>,
pump: Mutex<Option<JoinHandle<()>>>,
pump_cancel: CancellationToken,
workers: Mutex<HashMap<String, WorkerEntry>>,
@@ -177,14 +177,14 @@ impl KvEventIndex {
let (tx, rx) = mpsc::channel::<WorkerEvent>(EVENT_CHANNEL_BUFFER);
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 engine_reported_load = EngineReportedLoadTable::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 tally = Arc::new(EventTally::new());
let pump = tokio::spawn(pump_loop(
tree.clone(),
engine_load.clone(),
engine_reported_load.clone(),
cursors.clone(),
live_workers.clone(),
Arc::clone(&tally),
@@ -196,7 +196,7 @@ impl KvEventIndex {
maintain_tree,
subscribers,
load_subscribers,
engine_load,
engine_reported_load,
pump: Mutex::new(Some(pump)),
pump_cancel,
workers: Mutex::new(HashMap::new()),
@@ -239,8 +239,8 @@ impl KvEventIndex {
/// Shared accessor for the engine-load table. 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)
pub fn engine_reported_load(&self) -> Arc<EngineReportedLoadTable> {
Arc::clone(&self.engine_reported_load)
}
/// Register a worker. If `preresolved` is `Some`, the caller has
@@ -369,11 +369,12 @@ impl KvEventIndex {
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`
// Mark only the ranks that have an actual SUB socket. `EngineReportedLoadTable`
// 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.engine_reported_load
.mark_expected_rank(worker_url, *rank);
}
self.load_subscribers.add_worker(worker_url, &cfg).await;
}
@@ -412,7 +413,7 @@ impl KvEventIndex {
// 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);
self.engine_reported_load.forget_worker(worker_url);
let mut cursors = self.cursors.lock();
for id in &ids {
self.tree.clear_worker(id);
@@ -457,7 +458,7 @@ impl KvEventIndex {
/// restarting from seq=1 (after sending END_SEQ) is not filtered.
async fn pump_loop(
tree: Arc<HashTree>,
engine_load: Arc<EngineLoadTable>,
engine_reported_load: Arc<EngineReportedLoadTable>,
cursors: Arc<Mutex<HashMap<KvWorkerId, i64>>>,
live_workers: Arc<Mutex<HashSet<KvWorkerId>>>,
tally: Arc<EventTally>,
@@ -497,7 +498,7 @@ async fn pump_loop(
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());
engine_reported_load.set(&worker.url, worker.dp_rank, load, Instant::now());
}
WorkerEvent::PublisherReset { worker } => {
if cursors.lock().remove(&worker).is_some() {
@@ -585,8 +586,8 @@ 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};
use crate::state::kv_events::wire::{BlockRemoved, BlockStored, KvEventBatch};
use crate::state::load_monitor::engine_reported_load::LoadStat;
fn worker_id(url: &str, rank: u32) -> KvWorkerId {
KvWorkerId {
@@ -607,7 +608,7 @@ mod tests {
/// can destructure just the bits they need.
struct PumpHarness {
tree: Arc<HashTree>,
engine_load: Arc<EngineLoadTable>,
engine_reported_load: Arc<EngineReportedLoadTable>,
cursors: Arc<Mutex<HashMap<KvWorkerId, i64>>>,
tally: Arc<EventTally>,
#[allow(dead_code)]
@@ -622,7 +623,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 engine_reported_load = EngineReportedLoadTable::new();
let cursors = Arc::new(Mutex::new(HashMap::new()));
let live_set: Arc<Mutex<HashSet<KvWorkerId>>> =
Arc::new(Mutex::new(live.iter().cloned().collect()));
@@ -631,7 +632,7 @@ mod tests {
let (tx, rx) = mpsc::channel(4);
let pump = tokio::spawn(pump_loop(
tree.clone(),
engine_load.clone(),
engine_reported_load.clone(),
cursors.clone(),
live_set.clone(),
Arc::clone(&tally),
@@ -640,7 +641,7 @@ mod tests {
));
PumpHarness {
tree,
engine_load,
engine_reported_load,
cursors,
tally,
live_set,
@@ -833,7 +834,7 @@ mod tests {
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);
let (tree, engine_reported_load, tx, pump) = (h.tree, h.engine_reported_load, h.tx, h.pump);
tx.send(WorkerEvent::Load {
worker: id.clone(),
@@ -850,7 +851,7 @@ mod tests {
drop(tx);
pump.await.unwrap();
let snapshot = engine_load.capture_snapshot(Instant::now());
let snapshot = engine_reported_load.capture_snapshot(Instant::now());
let load = snapshot.fresh_load_for_url("http://w1").unwrap();
assert_eq!(load.num_running_reqs + load.num_waiting_reqs, 12);
// Load events must not pollute the cache tree.
@@ -1105,7 +1106,7 @@ mod tests {
assert_eq!(oracle.get(), Some(64));
assert!(oracle.is_bigram());
assert_eq!(index.known_worker_count(), 1);
assert_eq!(index.engine_load().expected_count(), 1);
assert_eq!(index.engine_reported_load().expected_count(), 1);
index.shutdown().await;
}
@@ -1128,10 +1129,10 @@ mod tests {
is_bigram: false,
};
index.add_worker(url, Some(cfg)).await;
assert_eq!(index.engine_load().expected_count(), 1);
assert_eq!(index.engine_reported_load().expected_count(), 1);
let now = Instant::now();
index.engine_load().set(
index.engine_reported_load().set(
url,
0,
LoadStat {
@@ -1144,7 +1145,7 @@ mod tests {
now,
);
assert!(index
.engine_load()
.engine_reported_load()
.capture_snapshot(now)
.fresh_load_for_url(url)
.is_some());
@@ -1152,13 +1153,13 @@ mod tests {
index.remove_worker(url).await;
assert!(
index
.engine_load()
.engine_reported_load()
.capture_snapshot(Instant::now())
.fresh_load_for_url(url)
.is_none(),
"remove_worker must clear engine load"
);
assert_eq!(index.engine_load().expected_count(), 0);
assert_eq!(index.engine_reported_load().expected_count(), 0);
index.shutdown().await;
}
experimental/sgl-router/src/policies/kv_events/mod.rs → experimental/sgl-router/src/state/kv_events/mod.rs
View File
experimental/sgl-router/src/policies/kv_events/subscriber.rs → experimental/sgl-router/src/state/kv_events/subscriber.rs
+3 -3
View File
@@ -7,7 +7,7 @@
//! 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
//! [`crate::state::load_monitor::engine_reported_load`]), and forwards [`WorkerEvent`]s to a
//! shared mpsc channel.
//!
//! # Wire format (3-frame multipart)
@@ -72,7 +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};
use crate::state::load_monitor::engine_reported_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
@@ -665,7 +665,7 @@ mod tests {
use tokio::time::timeout;
use zeromq::{Endpoint, PubSocket, Socket, SocketSend, ZmqMessage};
use crate::policies::kv_events::wire::KvCacheEvent;
use crate::state::kv_events::wire::KvCacheEvent;
mod helpers {
use super::*;
experimental/sgl-router/src/policies/kv_events/tally.rs → experimental/sgl-router/src/state/kv_events/tally.rs
View File
experimental/sgl-router/src/policies/kv_events/tree.rs → experimental/sgl-router/src/state/kv_events/tree.rs
View File
experimental/sgl-router/src/policies/kv_events/wire.rs → experimental/sgl-router/src/state/kv_events/wire.rs
View File
experimental/sgl-router/src/policies/engine_load.rs → experimental/sgl-router/src/state/load_monitor/engine_reported_load.rs
+35 -33
View File
@@ -10,7 +10,7 @@
//! aggregate and falls back to Router-local load when it is unavailable.
//!
//! Load is a *gauge*, not a delta: last value wins, no sequence/replay
//! semantics. Entries older than [`EngineLoadTable::freshness`] are ignored.
//! semantics. Entries older than [`EngineReportedLoadTable::freshness`] are ignored.
use std::collections::{HashMap, HashSet};
use std::fmt;
@@ -58,12 +58,12 @@ pub struct LoadStat {
pub native_cache: Option<NativeCacheRankLoad>,
}
/// Engine load for one worker captured at a fixed point in time.
/// Engine-reported request and KV counters for one worker, summed across DP ranks.
///
/// The four #34608 fields are summed across DP ranks. `captured_at` retains
/// the oldest rank timestamp so later local dispatches can be added.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EngineWorkerLoad {
pub struct EngineReportedWorkerLoad {
pub num_running_reqs: u64,
pub num_waiting_reqs: u64,
pub num_tokens: u64,
@@ -71,12 +71,14 @@ pub struct EngineWorkerLoad {
pub captured_at: Instant,
}
/// Complete ZMQ monitor aggregate used by native Cache-Aware.
/// Engine-reported per-worker scheduling load: capacity, queue pressure, and prefill estimates.
///
/// Shared by cache-aware, power-of-two, session-aware, and decode selection and admission.
///
/// Prefill throughput and queue time require two monotonic samples from every
/// DP rank. Initial samples and counter resets leave both values unavailable.
#[derive(Debug, Clone, PartialEq)]
pub struct NativeCacheWorkerLoad {
pub struct EngineReportedSchedulingLoad {
pub num_running_reqs: u64,
pub num_waiting_reqs: u64,
pub num_waiting_uncached_tokens: u64,
@@ -89,19 +91,19 @@ pub struct NativeCacheWorkerLoad {
pub captured_at: Instant,
}
/// Immutable engine load view captured once at request ingress.
/// Immutable fleet-wide view of engine-reported load, captured once at request ingress.
///
/// Keys are worker URLs used for dispatch. Missing, stale, or rank-incomplete
/// workers are omitted and must use Router-local active load.
#[derive(Debug, Clone, Default)]
pub struct EngineLoadSnapshot {
pub struct EngineReportedLoadSnapshot {
pub version: u64,
workers: HashMap<String, EngineWorkerLoad>,
native_cache_workers: HashMap<String, NativeCacheWorkerLoad>,
workers: HashMap<String, EngineReportedWorkerLoad>,
native_cache_workers: HashMap<String, EngineReportedSchedulingLoad>,
}
impl EngineLoadSnapshot {
pub fn fresh_load_for_url(&self, worker_url: &str) -> Option<&EngineWorkerLoad> {
impl EngineReportedLoadSnapshot {
pub fn fresh_load_for_url(&self, worker_url: &str) -> Option<&EngineReportedWorkerLoad> {
self.workers.get(worker_url)
}
@@ -136,13 +138,13 @@ impl EngineLoadSnapshot {
pub fn fresh_native_cache_load_for_url(
&self,
worker_url: &str,
) -> Option<&NativeCacheWorkerLoad> {
) -> Option<&EngineReportedSchedulingLoad> {
self.native_cache_workers.get(worker_url)
}
/// Builds a view from worker data that already passed freshness and rank checks.
/// Production requests should use [`EngineLoadTable::capture_snapshot`].
pub fn from_workers(version: u64, workers: HashMap<String, EngineWorkerLoad>) -> Self {
/// Production requests should use [`EngineReportedLoadTable::capture_snapshot`].
pub fn from_workers(version: u64, workers: HashMap<String, EngineReportedWorkerLoad>) -> Self {
Self {
version,
workers,
@@ -151,17 +153,17 @@ impl EngineLoadSnapshot {
}
/// Builds a test snapshot from complete native monitor data.
/// Production requests must use [`EngineLoadTable::capture_snapshot`].
/// Production requests must use [`EngineReportedLoadTable::capture_snapshot`].
pub fn from_native_cache_workers(
version: u64,
workers: HashMap<String, NativeCacheWorkerLoad>,
workers: HashMap<String, EngineReportedSchedulingLoad>,
) -> Self {
let basic = workers
.iter()
.map(|(url, load)| {
(
url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: load.num_running_reqs,
num_waiting_reqs: load.num_waiting_reqs,
num_tokens: load.num_used_tokens,
@@ -292,7 +294,7 @@ type NativeWorkerObservations = HashMap<u32, NativeRankObservation>;
/// Per-`(worker_url, dp_rank)` engine-reported load, written by the load
/// subscriber pump and captured once at request ingress.
#[derive(Debug)]
pub struct EngineLoadTable {
pub struct EngineReportedLoadTable {
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
@@ -302,7 +304,7 @@ pub struct EngineLoadTable {
version: AtomicU64,
}
impl EngineLoadTable {
impl EngineReportedLoadTable {
pub fn new() -> Arc<Self> {
Arc::new(Self {
by_rank: DashMap::new(),
@@ -364,7 +366,7 @@ impl EngineLoadTable {
/// look misleadingly idle and draw *more* traffic.) Callers that never
/// registered expected ranks retain the all-known-ranks rule. The oldest
/// timestamp represents the freshness of the complete aggregate.
fn fresh_worker_loads(&self, now: Instant) -> HashMap<String, EngineWorkerLoad> {
fn fresh_worker_loads(&self, now: Instant) -> HashMap<String, EngineReportedWorkerLoad> {
// 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() {
@@ -412,7 +414,7 @@ impl EngineLoadTable {
oldest_at.map(|captured_at| {
(
url,
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs,
num_waiting_reqs,
num_tokens,
@@ -432,7 +434,7 @@ impl EngineLoadTable {
fn fresh_native_cache_worker_loads(
&self,
now: Instant,
) -> HashMap<String, NativeCacheWorkerLoad> {
) -> HashMap<String, EngineReportedSchedulingLoad> {
let mut observed: HashMap<String, NativeWorkerObservations> = HashMap::new();
for entry in self.by_rank.iter() {
let at = entry.value().at;
@@ -521,7 +523,7 @@ impl EngineLoadTable {
oldest_at.map(|captured_at| {
(
url,
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs,
num_waiting_reqs,
num_waiting_uncached_tokens,
@@ -540,8 +542,8 @@ impl EngineLoadTable {
}
/// Captures one immutable view for all routing decisions in a request.
pub fn capture_snapshot(&self, now: Instant) -> EngineLoadSnapshot {
EngineLoadSnapshot {
pub fn capture_snapshot(&self, now: Instant) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot {
version: self.version.load(Ordering::Acquire),
workers: self.fresh_worker_loads(now),
native_cache_workers: self.fresh_native_cache_worker_loads(now),
@@ -635,7 +637,7 @@ mod tests {
#[test]
fn sums_queue_depth_across_ranks() {
let t = EngineLoadTable::new();
let t = EngineReportedLoadTable::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);
@@ -647,7 +649,7 @@ mod tests {
#[test]
fn stale_entries_are_dropped_from_snapshot() {
let t = EngineLoadTable::with_freshness(Duration::from_millis(10));
let t = EngineReportedLoadTable::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.
@@ -660,7 +662,7 @@ mod tests {
#[test]
fn forget_worker_clears_all_ranks() {
let t = EngineLoadTable::new();
let t = EngineReportedLoadTable::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);
@@ -677,7 +679,7 @@ mod tests {
/// router-side counter instead of looking misleadingly idle.
#[test]
fn partial_freshness_excludes_worker() {
let t = EngineLoadTable::with_freshness(Duration::from_secs(5));
let t = EngineReportedLoadTable::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
@@ -692,7 +694,7 @@ mod tests {
#[test]
fn missing_expected_rank_excludes_worker() {
let t = EngineLoadTable::new();
let t = EngineReportedLoadTable::new();
let now = Instant::now();
t.mark_expected_rank("http://w:30000", 0);
t.mark_expected_rank("http://w:30000", 1);
@@ -712,7 +714,7 @@ mod tests {
#[test]
fn capture_snapshot_uses_the_earliest_rank_timestamp() {
let t = EngineLoadTable::new();
let t = EngineReportedLoadTable::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);
@@ -726,7 +728,7 @@ mod tests {
#[test]
fn expected_count_tracks_marked_workers_and_forget() {
let t = EngineLoadTable::new();
let t = EngineReportedLoadTable::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
@@ -738,7 +740,7 @@ mod tests {
#[test]
fn complete_v3_semantic_samples_derive_prefill_queue_time() {
let t = EngineLoadTable::new();
let t = EngineReportedLoadTable::new();
let first = Instant::now();
let second = first + Duration::from_secs(2);
let mut old = load(2, 3);
experimental/sgl-router/src/state/load_monitor/mod.rs
+7
View File
@@ -0,0 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Load signals: engine-reported load and router-local in-flight accounting.
pub mod engine_reported_load;
pub mod router_inflight_load;
experimental/sgl-router/src/policies/active_load.rs → experimental/sgl-router/src/state/load_monitor/router_inflight_load.rs
+38 -32
View File
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Per-worker active-load tracking with RAII guards and a stale-request
//! Router-local per-worker in-flight load tracking with RAII guards and a stale-request
//! janitor.
//!
//! The per-worker `Worker::active_requests` counter tracks in-flight HTTP
@@ -18,7 +18,7 @@
//! 2. **Two-axis tracking** so PD-disaggregation can score prefill (token
//! count) separately from decode (block count). The two counters share
//! the same registry shape; we expose them as a single
//! [`ActiveLoadGuard`] holding both so the proxy's hot path mints one
//! [`RouterInflightLoadGuard`] holding both so the proxy's hot path mints one
//! guard per request rather than two.
//!
//! # Drop semantics
@@ -34,7 +34,7 @@
//!
//! # Clock injection
//!
//! [`ActiveLoadRegistry::new`] is generic over the clock so tests can drive
//! [`RouterInflightLoadRegistry::new`] is generic over the clock so tests can drive
//! the janitor deterministically. Production wires a `SystemTimeClock`;
//! tests use a `MockClock`. The `Instant`-based timestamp on registration
//! is sufficient for the timeout comparison (monotonic), so the clock
@@ -42,7 +42,7 @@
//! type whose `duration_since(other)` returns the wall-clock delta.
use crate::discovery::WorkerId;
use crate::server::metrics::{ActiveLoadKind, MetricsRegistry};
use crate::server::metrics::{MetricsRegistry, RouterInflightLoadKind};
use dashmap::DashMap;
use parking_lot::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
@@ -52,7 +52,7 @@ use tokio_util::sync::CancellationToken;
use uuid::Uuid;
/// Unique identifier for an in-flight request. Minted by
/// [`ActiveLoadRegistry::register`] and carried inside [`ActiveLoadGuard`]
/// [`RouterInflightLoadRegistry::register`] and carried inside [`RouterInflightLoadGuard`]
/// so the janitor can address one request at a time.
#[derive(Clone, Eq, Hash, PartialEq, Debug)]
pub struct RequestId(pub Uuid);
@@ -86,7 +86,7 @@ struct WorkerCounters {
///
/// `cancel` is a [`CancellationToken`] the janitor fires when the entry
/// is swept. The chat handler holds a clone (via
/// [`ActiveLoadGuard::cancel_token`]) and aborts its upstream fetch
/// [`RouterInflightLoadGuard::cancel_token`]) and aborts its upstream fetch
/// with `ApiError::StaleRequestExpired` when the token resolves —
/// surfacing the stale-request expiry as a 504 to the client instead
/// of leaving the handler hung on a long-lived upstream.
@@ -166,12 +166,12 @@ impl Clock for MockClock {
/// Registry of in-flight requests + per-worker active-load counters.
///
/// Constructed once per `AppContext`; the proxy holds an [`ActiveLoadGuard`]
/// Constructed once per `AppContext`; the proxy holds an [`RouterInflightLoadGuard`]
/// per request so counters decrement on drop. A background task periodically calls
/// [`Self::sweep_stale`] to evict requests that outlived
/// `stale_request_timeout`.
#[derive(Debug)]
pub struct ActiveLoadRegistry {
pub struct RouterInflightLoadRegistry {
workers: DashMap<WorkerId, Arc<WorkerCounters>>,
requests: DashMap<RequestId, RequestEntry>,
clock: Arc<dyn Clock>,
@@ -185,8 +185,8 @@ pub struct ActiveLoadRegistry {
metrics: Mutex<Option<Arc<MetricsRegistry>>>,
}
impl ActiveLoadRegistry {
/// Construct an [`ActiveLoadRegistry`] wrapped in an [`Arc`].
impl RouterInflightLoadRegistry {
/// Construct an [`RouterInflightLoadRegistry`] wrapped in an [`Arc`].
///
/// The registry is always shared (proxy + janitor + selector all hold
/// the same instance), so the public constructor mints the `Arc`
@@ -221,14 +221,14 @@ impl ActiveLoadRegistry {
let Some(metrics) = self.metrics.lock().clone() else {
return;
};
metrics.set_active_load(
metrics.set_router_inflight_load(
worker_url,
ActiveLoadKind::PrefillTokens,
RouterInflightLoadKind::PrefillTokens,
counters.prefill_load.load(Ordering::Relaxed) as i64,
);
metrics.set_active_load(
metrics.set_router_inflight_load(
worker_url,
ActiveLoadKind::DecodeBlocks,
RouterInflightLoadKind::DecodeBlocks,
counters.decode_load.load(Ordering::Relaxed) as i64,
);
}
@@ -262,7 +262,7 @@ impl ActiveLoadRegistry {
worker_url: impl Into<String>,
prefill_load: usize,
decode_load: usize,
) -> ActiveLoadGuard {
) -> RouterInflightLoadGuard {
let worker_url = worker_url.into();
let request_id = RequestId::new_v4();
let counters = self
@@ -291,7 +291,7 @@ impl ActiveLoadRegistry {
cancel: cancel.clone(),
},
);
ActiveLoadGuard {
RouterInflightLoadGuard {
registry: Some(Arc::clone(self)),
request_id: Some(request_id),
worker,
@@ -398,7 +398,7 @@ impl ActiveLoadRegistry {
}
/// Spawn a background janitor task that periodically calls
/// [`ActiveLoadRegistry::sweep_stale`].
/// [`RouterInflightLoadRegistry::sweep_stale`].
///
/// Returns a [`JanitorHandle`] that owns the join handle and a cancellation
/// token. Dropping the handle cancels the task; calling
@@ -407,9 +407,12 @@ impl ActiveLoadRegistry {
/// `interval` is the wall-clock cadence of the sweep. A sensible default
/// is half the configured `stale_request_timeout` so an expired entry is
/// reaped within 1.5× the timeout in the worst case. Pass a fresh
/// `Arc<ActiveLoadRegistry>` (cloned from the shared one held in
/// `Arc<RouterInflightLoadRegistry>` (cloned from the shared one held in
/// `AppContext`).
pub fn spawn_janitor(registry: Arc<ActiveLoadRegistry>, interval: Duration) -> JanitorHandle {
pub fn spawn_janitor(
registry: Arc<RouterInflightLoadRegistry>,
interval: Duration,
) -> JanitorHandle {
spawn_sweeper(move || registry.sweep_stale(), interval, "active-load")
}
@@ -481,16 +484,16 @@ impl Drop for JanitorHandle {
}
}
/// RAII guard returned by [`ActiveLoadRegistry::register`].
/// RAII guard returned by [`RouterInflightLoadRegistry::register`].
///
/// `#[must_use]`: a statement-form `registry.register(...)` would drop the
/// guard on the same line and decrement the counter before the request
/// actually executed, defeating the purpose. The compile-time warning
/// catches that misuse.
#[must_use = "ActiveLoadGuard must be held for the request's lifetime; dropping it immediately decrements counters"]
#[must_use = "RouterInflightLoadGuard must be held for the request's lifetime; dropping it immediately decrements counters"]
#[derive(Debug)]
pub struct ActiveLoadGuard {
registry: Option<Arc<ActiveLoadRegistry>>,
pub struct RouterInflightLoadGuard {
registry: Option<Arc<RouterInflightLoadRegistry>>,
/// `None` after the janitor expired this request — drop becomes a
/// no-op in that case. The guard keeps only the `RequestId`; the
/// per-axis amounts (and the captured `Arc<WorkerCounters>`) live
@@ -505,7 +508,7 @@ pub struct ActiveLoadGuard {
cancel: CancellationToken,
}
impl ActiveLoadGuard {
impl RouterInflightLoadGuard {
/// Read-only accessor (mainly for tests + diagnostic logging).
pub fn worker(&self) -> &WorkerId {
&self.worker
@@ -520,7 +523,7 @@ impl ActiveLoadGuard {
}
}
impl Drop for ActiveLoadGuard {
impl Drop for RouterInflightLoadGuard {
fn drop(&mut self) {
// If the janitor already expired this request (or `expire_now` was
// called explicitly), `request_id` is `None` and we skip — the
@@ -531,7 +534,7 @@ impl Drop for ActiveLoadGuard {
// `remove` returns `Some` exactly once; if the janitor races us
// and wins, we skip the decrement here. Decrement the **same**
// counters Arc the register call incremented (see
// `ActiveLoadGuard::counters`) — pinning the decrement to a
// `RouterInflightLoadGuard::counters`) — pinning the decrement to a
// specific WorkerCounters instance keeps the math correct
// across `forget_worker` + re-register cycles.
if let Some((_, entry)) = registry.requests.remove(&id) {
@@ -553,9 +556,12 @@ mod tests {
use super::*;
use std::sync::Arc;
fn registry_with_mock_clock(timeout: Duration) -> (Arc<ActiveLoadRegistry>, Arc<MockClock>) {
fn registry_with_mock_clock(
timeout: Duration,
) -> (Arc<RouterInflightLoadRegistry>, Arc<MockClock>) {
let clock = Arc::new(MockClock::new(Instant::now()));
let registry = ActiveLoadRegistry::new(Arc::clone(&clock) as Arc<dyn Clock>, timeout);
let registry =
RouterInflightLoadRegistry::new(Arc::clone(&clock) as Arc<dyn Clock>, timeout);
(registry, clock)
}
@@ -617,7 +623,7 @@ mod tests {
/// Gap closer #2: double-drop safety.
///
/// Rust's affine type system makes a literal double-drop of the same
/// `ActiveLoadGuard` value impossible — the compiler rejects
/// `RouterInflightLoadGuard` value impossible — the compiler rejects
/// `drop(g); drop(g);`. The interesting property is that the
/// registry's own bookkeeping never under-decrements, even if the
/// janitor and a guard's drop race. We assert that by simulating the
@@ -688,7 +694,7 @@ mod tests {
#[tokio::test]
async fn spawn_janitor_sweeps_stale_entries() {
let clock: Arc<dyn Clock> = Arc::new(SystemTimeClock);
let registry = ActiveLoadRegistry::new(clock, Duration::from_millis(30));
let registry = RouterInflightLoadRegistry::new(clock, Duration::from_millis(30));
let w = WorkerId("w0".into());
let _g = registry.register(w.clone(), "test://50-2", 50, 2);
assert_eq!(registry.inflight_count(), 1);
@@ -708,7 +714,7 @@ mod tests {
#[tokio::test]
async fn spawn_janitor_shutdown_is_clean() {
let clock: Arc<dyn Clock> = Arc::new(SystemTimeClock);
let registry = ActiveLoadRegistry::new(clock, Duration::from_secs(60));
let registry = RouterInflightLoadRegistry::new(clock, Duration::from_secs(60));
let handle = spawn_janitor(Arc::clone(&registry), Duration::from_millis(100));
// Verify shutdown completes within a generous bound.
let r = tokio::time::timeout(Duration::from_secs(2), handle.shutdown()).await;
@@ -840,7 +846,7 @@ mod tests {
/// When a [`MetricsRegistry`] is attached, the per-worker active-load
/// gauge mirrors the live counter on register / drop / sweep.
/// Regression: prior code exposed [`MetricsRegistry::set_active_load`]
/// Regression: prior code exposed [`MetricsRegistry::set_router_inflight_load`]
/// but nothing in the request hot path ever called it, leaving
/// `sgl_router_active_load` permanently at 0 in production.
#[test]
experimental/sgl-router/src/state/mod.rs
+12
View File
@@ -0,0 +1,12 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Shared mutable state that selection reads: the KV-event cache index, load
//! monitoring, and affinity assignments. Both the legacy policies and
//! `policies_reorg` read it here.
pub mod affinity_store;
pub mod kv_events;
pub mod load_monitor;
pub use affinity_store::AffinityStore;
experimental/sgl-router/src/tokenizer/mod.rs
+1 -1
View File
@@ -186,7 +186,7 @@ mod tests {
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
router_inflight_load: crate::config::InflightLoadConfig::default(),
}
}
experimental/sgl-router/src/workers/introspect.rs
+2 -2
View File
@@ -31,7 +31,7 @@ use serde::Deserialize;
use tracing::warn;
use url::Url;
use crate::policies::kv_events::EventConfig;
use crate::state::kv_events::EventConfig;
/// Default timeout for `/server_info`. Conservative for a small JSON
/// payload served by SGLang's HTTP server.
@@ -148,7 +148,7 @@ impl WorkerIntrospector {
// EAGLE-family speculative decoding ⇒ the worker hashes KV blocks over
// token bigrams; the router must mirror that on the selection side.
let is_bigram = crate::policies::kv_events::classify_bigram(
let is_bigram = crate::state::kv_events::classify_bigram(
parsed.speculative_algorithm.as_deref(),
worker_url,
);
experimental/sgl-router/src/workers/manager.rs
+20 -20
View File
@@ -4,8 +4,8 @@
use crate::config::Config;
use crate::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, WorkerSpec};
use crate::health::circuit_breaker::CircuitBreakerConfig;
use crate::policies::active_load::ActiveLoadRegistry;
use crate::policies::kv_events::KvEventIndex;
use crate::state::kv_events::KvEventIndex;
use crate::state::load_monitor::router_inflight_load::RouterInflightLoadRegistry;
use crate::workers::introspect::{DisaggregationRole, WorkerIntrospector};
use crate::workers::{WireProtocol, WorkerRegistry};
use std::collections::HashMap;
@@ -109,7 +109,7 @@ pub async fn run(rx: mpsc::Receiver<DiscoveryEvent>, registry: Arc<WorkerRegistr
/// manager does not need a handle to the proxy.
///
/// When `kv_index` is `None`, KV-event and load-subscriber state is disabled; when
/// `active_load` is `None` the active-load bookkeeping is not pruned
/// `router_inflight_load` is `None` the active-load bookkeeping is not pruned
/// on worker removal (leaks one `WorkerCounters` slot per departed
/// worker — fine for tests, but production passes `Some(...)`); when
/// `cfg` is `None` the default CB config is used for every worker
@@ -123,14 +123,14 @@ pub async fn run_with_config(
registry: Arc<WorkerRegistry>,
cfg: Option<Arc<Config>>,
kv_index: Option<Arc<KvEventIndex>>,
active_load: Option<Arc<ActiveLoadRegistry>>,
router_inflight_load: Option<Arc<RouterInflightLoadRegistry>>,
) {
run_with_introspector(
rx,
registry,
cfg,
kv_index,
active_load,
router_inflight_load,
Arc::new(WorkerIntrospector::default()),
)
.await
@@ -146,7 +146,7 @@ pub async fn run_with_introspector(
registry: Arc<WorkerRegistry>,
cfg: Option<Arc<Config>>,
kv_index: Option<Arc<KvEventIndex>>,
active_load: Option<Arc<ActiveLoadRegistry>>,
router_inflight_load: Option<Arc<RouterInflightLoadRegistry>>,
introspector: Arc<WorkerIntrospector>,
) {
run_with_introspector_and_reconcile(
@@ -154,7 +154,7 @@ pub async fn run_with_introspector(
registry,
cfg,
kv_index,
active_load,
router_inflight_load,
introspector,
RECONCILE_INTERVAL,
)
@@ -186,7 +186,7 @@ pub async fn run_with_introspector_and_reconcile(
registry: Arc<WorkerRegistry>,
cfg: Option<Arc<Config>>,
kv_index: Option<Arc<KvEventIndex>>,
active_load: Option<Arc<ActiveLoadRegistry>>,
router_inflight_load: Option<Arc<RouterInflightLoadRegistry>>,
introspector: Arc<WorkerIntrospector>,
reconcile_interval: Duration,
) {
@@ -226,7 +226,7 @@ pub async fn run_with_introspector_and_reconcile(
&registry,
&cfg,
&kv_index,
&active_load,
&router_inflight_load,
&introspector,
&mut pending,
)
@@ -264,7 +264,7 @@ async fn handle_discovery_event(
registry: &Arc<WorkerRegistry>,
cfg: &Option<Arc<Config>>,
kv_index: &Option<Arc<KvEventIndex>>,
active_load: &Option<Arc<ActiveLoadRegistry>>,
router_inflight_load: &Option<Arc<RouterInflightLoadRegistry>>,
introspector: &Arc<WorkerIntrospector>,
pending: &mut HashMap<WorkerId, JoinHandle<()>>,
) {
@@ -327,7 +327,7 @@ async fn handle_discovery_event(
// per-worker counters slot will not be re-created
// (selectors no longer see the worker, so no new
// requests can register against it).
if let Some(al) = active_load {
if let Some(al) = router_inflight_load {
al.forget_worker(&id);
}
}
@@ -570,7 +570,7 @@ async fn register_one(
mod tests {
use super::*;
use crate::config::{
ActiveLoadConfig, CircuitBreakerConfig as RawCbConfig, DiscoveryBackend, ModelConfig,
CircuitBreakerConfig as RawCbConfig, DiscoveryBackend, InflightLoadConfig, ModelConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use crate::discovery::{WorkerId, WorkerMode};
@@ -610,7 +610,7 @@ mod tests {
urls: vec!["http://test:30000".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
@@ -970,7 +970,7 @@ mod tests {
}
/// Task B: `DiscoveryEvent::Removed` calls
/// `ActiveLoadRegistry::forget_worker` so the per-worker counters
/// `RouterInflightLoadRegistry::forget_worker` so the per-worker counters
/// slot is reaped. Without this, a long-lived cluster with worker
/// churn would leak one `WorkerCounters` entry per departed worker.
#[tokio::test]
@@ -984,14 +984,14 @@ mod tests {
spawn_fake_server_info_worker(json!({"served_model_name": "m"})).await;
let registry = Arc::new(WorkerRegistry::default());
let active_load = ActiveLoadRegistry::with_defaults();
let router_inflight_load = RouterInflightLoadRegistry::with_defaults();
let (tx, rx) = mpsc::channel::<DiscoveryEvent>(8);
let manager_handle = tokio::spawn(run_with_introspector(
rx,
registry.clone(),
None,
None,
Some(Arc::clone(&active_load)),
Some(Arc::clone(&router_inflight_load)),
fast_introspector(),
));
@@ -1019,8 +1019,8 @@ mod tests {
// Mint a guard to force the active-load registry to create a
// per-worker counters slot for this id.
let _g = active_load.register(id.clone(), "test://", 10, 1);
assert!(active_load.is_known(&id));
let _g = router_inflight_load.register(id.clone(), "test://", 10, 1);
assert!(router_inflight_load.is_known(&id));
// Now drive the Removed event and assert the counters slot is
// gone. We tear down the guard last so the request entry is
@@ -1030,7 +1030,7 @@ mod tests {
.unwrap();
let removed = timeout(Duration::from_secs(2), async {
loop {
if !active_load.is_known(&id) && registry.get(&id).is_none() {
if !router_inflight_load.is_known(&id) && registry.get(&id).is_none() {
return true;
}
tokio::time::sleep(Duration::from_millis(20)).await;
@@ -1039,7 +1039,7 @@ mod tests {
.await;
assert!(
removed.is_ok(),
"manager must call active_load.forget_worker on Removed",
"manager must call router_inflight_load.forget_worker on Removed",
);
drop(tx);
experimental/sgl-router/src/workers/registry.rs
+2 -2
View File
@@ -77,7 +77,7 @@ impl WorkerRegistry {
/// spec carries an id that already has an entry, the prior entry
/// stays put — it's the caller's responsibility to decide whether
/// to evict it (and, importantly, to also clean up sidecar state
/// in `KvEventIndex` / `ActiveLoadRegistry` if so). Doing that
/// in `KvEventIndex` / `RouterInflightLoadRegistry` if so). Doing that
/// cleanup here would leak orphan state into those sidecars when
/// a caller actually wanted to keep the prior entry.
pub fn add_with_cb(
@@ -463,7 +463,7 @@ mod tests {
/// On a rejected upsert with `MixedPdAndPlain`, the registry is
/// **not** mutated — the prior entry for the rejected id stays
/// put. Eviction (with the matching `KvEventIndex` /
/// `ActiveLoadRegistry` cleanup) is the manager's responsibility;
/// `RouterInflightLoadRegistry` cleanup) is the manager's responsibility;
/// doing it here would leak orphan state in those sidecars.
#[test]
fn upsert_rejected_with_mixed_modes_leaves_registry_unchanged() {
experimental/sgl-router/src/workers/worker.rs
+12 -12
View File
@@ -246,7 +246,7 @@ impl Worker {
self.protocol
}
pub fn active_load(&self) -> usize {
pub fn router_inflight_load(&self) -> usize {
self.active_requests.load(Ordering::Relaxed)
}
@@ -284,7 +284,7 @@ impl std::fmt::Debug for Worker {
.field("url", &self.url)
.field("mode", &self.mode())
.field("protocol", &self.protocol)
.field("active_load", &self.active_load())
.field("router_inflight_load", &self.router_inflight_load())
.finish()
}
}
@@ -304,15 +304,15 @@ mod tests {
model_ids: vec![ModelId("m".into())],
bootstrap_port: None,
});
assert_eq!(w.active_load(), 0);
assert_eq!(w.router_inflight_load(), 0);
let g = w.load_guard();
assert_eq!(w.active_load(), 1);
assert_eq!(w.router_inflight_load(), 1);
let g2 = w.load_guard();
assert_eq!(w.active_load(), 2);
assert_eq!(w.router_inflight_load(), 2);
drop(g);
assert_eq!(w.active_load(), 1);
assert_eq!(w.router_inflight_load(), 1);
drop(g2);
assert_eq!(w.active_load(), 0);
assert_eq!(w.router_inflight_load(), 0);
}
#[test]
@@ -321,11 +321,11 @@ mod tests {
let cutoff = Instant::now() - Duration::from_secs(1);
let guard = w.load_guard();
assert_eq!(w.active_load(), 1);
assert_eq!(w.router_inflight_load(), 1);
assert_eq!(w.slots_acquired_since(cutoff), 0);
drop(guard);
assert_eq!(w.active_load(), 0);
assert_eq!(w.router_inflight_load(), 0);
}
#[test]
@@ -334,11 +334,11 @@ mod tests {
let cutoff = Instant::now() - Duration::from_secs(1);
let guard = w.timestamped_load_guard();
assert_eq!(w.active_load(), 1);
assert_eq!(w.router_inflight_load(), 1);
assert_eq!(w.slots_acquired_since(cutoff), 1);
drop(guard);
assert_eq!(w.active_load(), 0);
assert_eq!(w.router_inflight_load(), 0);
assert_eq!(w.slots_acquired_since(cutoff), 0);
}
@@ -477,7 +477,7 @@ mod tests {
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.router_inflight_load(), 3);
assert_eq!(
w.slots_acquired_since(cutoff),
2,
experimental/sgl-router/tests/component/discovery/static_urls.rs
+3 -2
View File
@@ -90,7 +90,8 @@ async fn static_urls_pd_role_resolved_end_to_end() {
use axum::{routing::get, Json, Router};
use serde_json::json;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ObservabilityConfig, ProxyConfig, ServerConfig,
Config, DiscoveryBackend, InflightLoadConfig, ObservabilityConfig, ProxyConfig,
ServerConfig,
};
use sgl_router::discovery::{spawn_discovery, WorkerId};
use sgl_router::workers::{manager, WorkerRegistry};
@@ -146,7 +147,7 @@ async fn static_urls_pd_role_resolved_end_to_end() {
urls: vec![url.clone()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
let registry = Arc::new(WorkerRegistry::default());
experimental/sgl-router/tests/component/policies/cache_prefix_provider.rs
+1 -3
View File
@@ -5,10 +5,8 @@ use std::collections::HashMap;
use std::sync::Arc;
use sgl_kv_indexer::PrefixOutcome;
use sgl_router::policies::kv_events::{
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
};
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
#[test]
fn radix_tree_reports_contiguous_prefix_depth_per_worker() {
experimental/sgl-router/tests/component/policies/decode.rs
+6 -4
View File
@@ -12,8 +12,10 @@ use sgl_router::policies::decode::{
resolve_decode_with_capacity_fallback, DecodePolicy, DecodePowerOfTwoPolicy,
DecodeSelectionContext, LegacyHostAffinityDecodePolicy,
};
use sgl_router::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
use sgl_router::policies::SelectionProposal;
use sgl_router::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad,
};
use sgl_router::workers::Worker;
use std::collections::HashMap;
use std::sync::atomic::Ordering;
@@ -30,15 +32,15 @@ fn worker(id: &str) -> Arc<Worker> {
}))
}
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
7,
entries
.iter()
.map(|(worker, running, waiting, used, capacity)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: *running,
num_waiting_reqs: *waiting,
num_waiting_uncached_tokens: *waiting,
experimental/sgl-router/tests/component/policies/fused_score.rs
+10 -10
View File
@@ -12,15 +12,15 @@
//! 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, ScorePolicy,
};
use sgl_router::policies::{Policy, SelectionContext};
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
use sgl_router::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedWorkerLoad,
};
use sgl_router::workers::Worker;
use std::{collections::HashMap, sync::Arc, time::Instant};
@@ -82,12 +82,12 @@ 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(
let snapshot = EngineReportedLoadSnapshot::from_workers(
29,
HashMap::from([
(
ws[0].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 50,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -97,7 +97,7 @@ fn fused_load_based_term_uses_the_request_snapshot() {
),
(
ws[1].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -123,12 +123,12 @@ fn fused_load_based_term_uses_the_request_snapshot() {
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(
let snapshot = EngineReportedLoadSnapshot::from_workers(
31,
HashMap::from([
(
ws[0].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 50,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -138,7 +138,7 @@ fn score_policy_forwards_the_request_snapshot_to_load_based() {
),
(
ws[1].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
experimental/sgl-router/tests/component/policies/kv_events_hash_parity.rs
+1 -1
View File
@@ -19,7 +19,7 @@
//! whatever fixture is checked in.
use serde::Deserialize;
use sgl_router::policies::kv_events::compute_block_hashes;
use sgl_router::state::kv_events::compute_block_hashes;
use std::path::PathBuf;
#[derive(Debug, Deserialize)]
experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs
+2 -2
View File
@@ -3,7 +3,7 @@
//! Concurrent-mutation stress test for `HashTree`.
//!
//! The inline tests in `policies::kv_events::tree` are all
//! The inline tests in `state::kv_events::tree` are all
//! single-threaded. Under production load, multiple worker subscribers
//! drive `insert` / `remove` / `clear_worker` against the same tree from
//! tokio worker threads while the chat handler simultaneously calls
@@ -25,7 +25,7 @@
use std::sync::Arc;
use std::thread;
use sgl_router::policies::kv_events::{HashTree, KvWorkerId};
use sgl_router::state::kv_events::{HashTree, KvWorkerId};
fn worker(i: usize) -> KvWorkerId {
KvWorkerId {
experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs
+2 -2
View File
@@ -22,8 +22,8 @@ use std::time::Duration;
use zeromq::SocketSend;
use sgl_router::policies::kv_events::discovery::EventConfig;
use sgl_router::policies::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
use sgl_router::state::kv_events::discovery::EventConfig;
use sgl_router::state::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
use super::zmq_helpers::{
build_multipart, encode_block_stored_event, encode_event_batch, make_pub_bound,
experimental/sgl-router/tests/component/policies/zmq_helpers.rs
+1 -1
View File
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Shared ZMQ wire-format helpers for the `policies::kv_events` component
//! Shared ZMQ wire-format helpers for the `state::kv_events` component
//! tests. Encodes events in the same msgspec layout SGLang emits, builds
//! the two-frame `[seq, payload]` ZMQ message a real publisher sends, and
//! binds a loopback PUB socket on an OS-assigned port.
experimental/sgl-router/tests/component/tokenizer/render_parity.rs
+2 -2
View File
@@ -6,7 +6,7 @@
use serde::Deserialize;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::ModelId;
@@ -128,7 +128,7 @@ fn registry(model_id: &str, tokenizer_path: PathBuf) -> TokenizerRegistry {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
TokenizerRegistry::load_from_config(&cfg).unwrap()
}
experimental/sgl-router/tests/component/workers/concurrent_state.rs
+3 -3
View File
@@ -131,17 +131,17 @@ fn load_guard_decrements_on_panic_unwind() {
model_ids: vec![ModelId("m".into())],
bootstrap_port: None,
}));
assert_eq!(w.active_load(), 0);
assert_eq!(w.router_inflight_load(), 0);
let w_inner = w.clone();
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
let _g = w_inner.load_guard();
assert_eq!(w_inner.active_load(), 1);
assert_eq!(w_inner.router_inflight_load(), 1);
panic!("synthetic panic to exercise Drop on unwind");
}));
assert!(result.is_err(), "the closure must have panicked");
assert_eq!(
w.active_load(),
w.router_inflight_load(),
0,
"LoadGuard's Drop must decrement even when the holder panics",
);
experimental/sgl-router/tests/component/workers/manager.rs
+1 -1
View File
@@ -740,7 +740,7 @@ async fn removed_awaits_pending_added() {
/// independently — 2N round-trips for N workers.
#[tokio::test]
async fn manager_emits_single_server_info_fetch_per_worker() {
use sgl_router::policies::kv_events::KvEventIndex;
use sgl_router::state::kv_events::KvEventIndex;
let body = json!({
"served_model_name": "m",
experimental/sgl-router/tests/proxy/bucket_routing.rs
+5 -5
View File
@@ -10,17 +10,17 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use sgl_kv_indexer::{PrefixIndex, PrefixIndexError, PrefixMatch, PrefixOutcome};
use sgl_router::config::{
ActiveLoadConfig, AffinityConfig, BucketConfig, BucketSpec, BucketStage, CacheAwareConfig,
CachePrefixProvider, Config, DiscoveryBackend, KvIndexerEndpointConfig, ModelConfig,
AffinityConfig, BucketConfig, BucketSpec, BucketStage, CacheAwareConfig, CachePrefixProvider,
Config, DiscoveryBackend, InflightLoadConfig, KvIndexerEndpointConfig, ModelConfig,
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode,
SloBucketPolicy, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::{LoadStat, NativeCacheRankLoad};
use sgl_router::policies::factory::build_registry_with_defaults;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::load_monitor::engine_reported_load::{LoadStat, NativeCacheRankLoad};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use std::sync::atomic::{AtomicUsize, Ordering};
@@ -77,7 +77,7 @@ fn build_app_context(
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&config).unwrap());
let registry = Arc::new(WorkerRegistry::default());
@@ -236,7 +236,7 @@ fn set_native_load_with_waiting(
max_total_num_tokens: u64,
num_waiting_uncached_tokens: u64,
) {
ctx.engine_load.set(
ctx.engine_reported_load.set(
worker_url,
0,
LoadStat {
experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs
+1 -1
View File
@@ -18,10 +18,10 @@ use axum::http::{Request, StatusCode};
use serde_json::{json, Value};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::kv_events::{BlockSizeOracle, HashTree};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use std::sync::Arc;
experimental/sgl-router/tests/proxy/chat_routing.rs
+34 -32
View File
@@ -2,7 +2,7 @@
// SPDX-License-Identifier: Apache-2.0
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -50,7 +50,7 @@ fn config_for(_worker_url: &str) -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
@@ -1322,7 +1322,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br
/// streaming response, not just for the handler lifetime.
///
/// Before the fix, the handler dropped `_guard` as soon as it returned
/// (which happens when headers arrive), so `active_load()` was 0 while
/// (which happens when headers arrive), so `router_inflight_load()` was 0 while
/// the SSE pump was still relaying bytes. This test catches that bug.
#[tokio::test]
async fn streaming_load_guard_persists_for_body_lifetime() {
@@ -1360,7 +1360,7 @@ async fn streaming_load_guard_persists_for_body_lifetime() {
));
let app = build_router(ctx);
// Grab the Worker handle so we can assert active_load().
// Grab the Worker handle so we can assert router_inflight_load().
let w_handle: Arc<Worker> = registry
.workers_for(&ModelId("tiny".into()))
.into_iter()
@@ -1386,9 +1386,9 @@ async fn streaming_load_guard_persists_for_body_lifetime() {
// first chunk's delay to pass, then assert load is still held.
tokio::time::sleep(Duration::from_millis(20)).await;
assert!(
w_handle.active_load() >= 1,
w_handle.router_inflight_load() >= 1,
"load should be >= 1 mid-stream, got {}",
w_handle.active_load()
w_handle.router_inflight_load()
);
// Drain the entire body — this drives the SSE pump to completion.
@@ -1398,25 +1398,25 @@ async fn streaming_load_guard_persists_for_body_lifetime() {
// released. Give the spawned task a brief moment to clean up.
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(
w_handle.active_load(),
w_handle.router_inflight_load(),
0,
"load should be 0 after stream completes"
);
}
/// Task A: the chat handler mints an `ActiveLoadGuard` from the shared
/// `ActiveLoadRegistry` and drops it when the request completes. The
/// Task A: the chat handler mints an `RouterInflightLoadGuard` from the shared
/// `RouterInflightLoadRegistry` and drops it when the request completes. The
/// non-streaming path drops the guard on handler exit; this test
/// asserts the round-trip increment → 0 across a single request.
#[tokio::test]
async fn non_streaming_active_load_increments_then_returns_to_zero() {
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
let ctx = build_ctx_with_worker(&worker.url);
let active_load = Arc::clone(&ctx.active_load);
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
let app = build_router(ctx);
assert_eq!(
active_load.inflight_count(),
router_inflight_load.inflight_count(),
0,
"registry must start with no in-flight requests",
);
@@ -1442,19 +1442,19 @@ async fn non_streaming_active_load_increments_then_returns_to_zero() {
// The handler has returned, so the active-load guard must have
// dropped — counters are back to zero.
assert_eq!(
active_load.inflight_count(),
router_inflight_load.inflight_count(),
0,
"active-load registry must be empty after non-streaming handler returns",
);
let w_id = WorkerId("w1".into());
assert_eq!(
active_load.prefill_load(&w_id),
router_inflight_load.prefill_load(&w_id),
0,
"prefill_load must decrement on response end",
);
}
/// Task A: the streaming path holds the `ActiveLoadGuard` until the
/// Task A: the streaming path holds the `RouterInflightLoadGuard` until the
/// SSE pump finishes. Mid-stream the registry shows `inflight_count >= 1`;
/// after the body drains it returns to 0. Counterpart to
/// `streaming_load_guard_persists_for_body_lifetime` — both guards must
@@ -1486,7 +1486,7 @@ async fn streaming_active_load_persists_for_body_lifetime() {
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
let ctx = Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies));
let active_load = Arc::clone(&ctx.active_load);
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
let app = build_router(ctx);
let req = Request::builder()
@@ -1508,15 +1508,15 @@ async fn streaming_active_load_persists_for_body_lifetime() {
// still running, so the registry's per-request entry must remain.
tokio::time::sleep(Duration::from_millis(20)).await;
assert!(
active_load.inflight_count() >= 1,
router_inflight_load.inflight_count() >= 1,
"registry inflight must be >= 1 mid-stream, got {}",
active_load.inflight_count(),
router_inflight_load.inflight_count(),
);
let w_id = WorkerId("w1".into());
assert!(
active_load.prefill_load(&w_id) >= 1,
router_inflight_load.prefill_load(&w_id) >= 1,
"prefill_load must be > 0 mid-stream, got {}",
active_load.prefill_load(&w_id),
router_inflight_load.prefill_load(&w_id),
);
// Drain the body — drives the SSE pump to completion.
@@ -1524,12 +1524,12 @@ async fn streaming_active_load_persists_for_body_lifetime() {
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(
active_load.inflight_count(),
router_inflight_load.inflight_count(),
0,
"registry must be empty after stream drains",
);
assert_eq!(
active_load.prefill_load(&w_id),
router_inflight_load.prefill_load(&w_id),
0,
"prefill_load must be 0 after stream drains",
);
@@ -1555,7 +1555,7 @@ async fn streaming_active_load_drops_on_client_disconnect() {
)
.await;
let (ctx, body) = stream_chat(&worker.url).await;
let active_load = Arc::clone(&ctx.active_load);
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
// Read one chunk to confirm the stream is live, then drop the body.
use futures::StreamExt;
@@ -1570,7 +1570,7 @@ async fn streaming_active_load_drops_on_client_disconnect() {
wait_for_metric(&ctx, &expected).await;
assert_eq!(
active_load.inflight_count(),
router_inflight_load.inflight_count(),
0,
"client disconnect must drop the streaming pump's guards within one tick",
);
@@ -1583,13 +1583,15 @@ async fn streaming_active_load_drops_on_client_disconnect() {
/// `ApiError::StaleRequestExpired`.
///
/// Wiring: build an `AppContext` with a short
/// `stale_request_timeout` `ActiveLoadRegistry` + spawn a janitor
/// `stale_request_timeout` `RouterInflightLoadRegistry` + spawn a janitor
/// with sub-second cadence + dispatch to a slow upstream that takes
/// longer than the timeout. The janitor sweeps before the upstream
/// returns; cancellation fires; handler returns 504.
#[tokio::test]
async fn janitor_expiry_returns_504_stale_request_expired() {
use sgl_router::policies::active_load::{spawn_janitor, ActiveLoadRegistry};
use sgl_router::state::load_monitor::router_inflight_load::{
spawn_janitor, RouterInflightLoadRegistry,
};
// Upstream that takes 2s to respond — longer than our 50ms
// stale_request_timeout.
let worker =
@@ -1610,18 +1612,18 @@ async fn janitor_expiry_returns_504_stale_request_expired() {
// Aggressive 50ms timeout: the janitor will sweep on the next
// tick (every 20ms) and fire the cancellation token before the
// upstream returns.
let active_load = ActiveLoadRegistry::new(
Arc::new(sgl_router::policies::active_load::SystemTimeClock),
let router_inflight_load = RouterInflightLoadRegistry::new(
Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock),
Duration::from_millis(50),
);
let _janitor = spawn_janitor(Arc::clone(&active_load), Duration::from_millis(20));
let ctx = Arc::new(AppContext::with_active_load(
let _janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20));
let ctx = Arc::new(AppContext::with_router_inflight_load(
cfg,
tokenizers,
proxy,
registry,
policies,
active_load,
router_inflight_load,
));
let app = build_router(ctx);
@@ -1671,7 +1673,7 @@ async fn non_streaming_error_path_drops_active_load_guard() {
drop(listener);
let ctx = build_ctx_with_worker(&dead_url);
let active_load = Arc::clone(&ctx.active_load);
let router_inflight_load = Arc::clone(&ctx.router_inflight_load);
let app = build_router(ctx);
let req = Request::builder()
@@ -1692,7 +1694,7 @@ async fn non_streaming_error_path_drops_active_load_guard() {
// Drain so any drop-on-body-end work runs.
let _ = res.into_body().collect().await.unwrap().to_bytes();
assert_eq!(
active_load.inflight_count(),
router_inflight_load.inflight_count(),
0,
"error path must drop the active-load guard",
);
experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs
+3 -3
View File
@@ -7,8 +7,8 @@
//! built-in V4 chat formatter — the engine-equivalent path — with no template fixture.
use sgl_router::config::{
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
CacheAwareConfig, Config, DiscoveryBackend, InflightLoadConfig, ModelConfig,
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
pub const MODEL: &str = "deepseek-v4-tiny";
@@ -42,6 +42,6 @@ pub fn config() -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/tests/proxy/external_indexer_routing.rs
+1 -1
View File
@@ -19,11 +19,11 @@ use sgl_kv_indexer::{
use sgl_router::config::{AffinityConfig, CachePrefixProvider, PolicyKind};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
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;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use tokio_stream::wrappers::TcpListenerStream;
experimental/sgl-router/tests/proxy/failover.rs
+1 -1
View File
@@ -54,7 +54,7 @@ async fn failover_when_one_worker_dies() {
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
experimental/sgl-router/tests/proxy/graceful_shutdown.rs
+4 -4
View File
@@ -25,7 +25,7 @@
use futures::future::join_all;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -69,7 +69,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
@@ -554,7 +554,7 @@ async fn wait_for_inflight_http(ctx: &Arc<AppContext>, want: usize) {
/// response BODY finishing, not the handler returning. A streaming completion
/// hands back its headers immediately, so a count released at handler exit
/// would read 0 for the entire window the heartbeat exists to explain — the
/// same blind spot `active_load.inflight_count()` has, reproduced in the
/// same blind spot `router_inflight_load.inflight_count()` has, reproduced in the
/// replacement.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn inflight_http_counts_a_streaming_response_until_its_body_finishes() {
@@ -623,7 +623,7 @@ async fn inflight_http_counts_a_streaming_response_until_its_body_finishes() {
/// Every route is instrumented, not only the proxied ones. `/metrics`,
/// `/readyz` and a 404 are exchanges axum's drain waits on too, and they are
/// exactly the traffic `active_load` cannot see — so a guard that leaked on a
/// exactly the traffic `router_inflight_load` cannot see — so a guard that leaked on a
/// non-proxied route would leave the heartbeat permanently busy and turn the
/// drain report back into noise.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
experimental/sgl-router/tests/proxy/header_forwarding.rs
+2 -2
View File
@@ -4,7 +4,7 @@
use axum::body::Body;
use axum::http::Request;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -47,7 +47,7 @@ async fn forwards_whitelisted_headers_strips_others() {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs
+3 -3
View File
@@ -21,7 +21,7 @@ use axum::http::{Request, StatusCode};
use bytes::Bytes;
use serde_json::{json, Value};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -62,7 +62,7 @@ fn config() -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
@@ -225,7 +225,7 @@ async fn round_robin_pd_prefill_does_not_track_dispatch_timestamps() {
let request = tokio::spawn(build_router(Arc::clone(&ctx)).oneshot(chat_request()));
await_captured_body(&prefill, Duration::from_secs(2), "prefill").await;
assert_eq!(prefill_worker.active_load(), 1);
assert_eq!(prefill_worker.router_inflight_load(), 1);
assert_eq!(prefill_worker.slots_acquired_since(cutoff), 0);
assert_eq!(request.await.unwrap().unwrap().status(), StatusCode::OK);
experimental/sgl-router/tests/proxy/pd_pool_isolation.rs
+2 -2
View File
@@ -20,7 +20,7 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -61,7 +61,7 @@ fn config() -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/tests/proxy/pd_protocol_binding.rs
+2 -2
View File
@@ -32,7 +32,7 @@ use hyper::service::service_fn;
use hyper::Response as HyperResponse;
use hyper_util::rt::{TokioExecutor, TokioIo};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -149,7 +149,7 @@ fn config() -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/tests/proxy/radix_tree_routing.rs
+1 -3
View File
@@ -10,14 +10,12 @@ use serde_json::json;
use sgl_router::config::{AffinityConfig, CachePrefixProvider, PolicyKind};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
};
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
use sgl_router::policies::request_tokens_for;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use tower::ServiceExt;
experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs
+2 -2
View File
@@ -11,7 +11,7 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::{json, Value};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -59,7 +59,7 @@ fn config() -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/tests/proxy/shared_prefill_admission.rs
+7 -7
View File
@@ -7,11 +7,10 @@ use std::time::{Duration, Instant};
use axum::body::Body;
use axum::http::{Request, StatusCode};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::{LoadStat, NativeCacheRankLoad};
use sgl_router::policies::{
CacheCandidate, CacheCandidateProposal, Policy, PolicyRegistry, PrefillProposal, ProposalKind,
SelectionContext, SelectionProposal,
@@ -19,6 +18,7 @@ use sgl_router::policies::{
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::load_monitor::engine_reported_load::{LoadStat, NativeCacheRankLoad};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::{Worker, WorkerRegistry};
use tower::ServiceExt;
@@ -165,7 +165,7 @@ fn config(policy: PolicyKind) -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
@@ -307,7 +307,7 @@ async fn chat_commits_the_admitted_prefill_backup() {
total_prefill_busy_us,
}),
};
fixture.ctx.engine_load.set(
fixture.ctx.engine_reported_load.set(
&fixture.workers[0].url,
0,
native_load(1, 1),
@@ -315,7 +315,7 @@ async fn chat_commits_the_admitted_prefill_backup() {
);
fixture
.ctx
.engine_load
.engine_reported_load
.set(&fixture.workers[0].url, 0, native_load(2, 2), now);
assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK);
@@ -351,7 +351,7 @@ async fn capacity_exhaustion_does_not_return_503() {
})
.await;
for worker in &fixture.workers {
fixture.ctx.engine_load.set(
fixture.ctx.engine_reported_load.set(
&worker.url,
0,
LoadStat {
@@ -434,7 +434,7 @@ async fn chat_records_cache_candidates_exhausted() {
})
})
.await;
fixture.ctx.engine_load.set(
fixture.ctx.engine_reported_load.set(
&fixture.workers[0].url,
0,
LoadStat {
experimental/sgl-router/tests/proxy/sticky_input_ids.rs
+2 -2
View File
@@ -24,7 +24,7 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::{json, Value};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -77,7 +77,7 @@ fn config() -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}
experimental/sgl-router/tests/proxy/sticky_routing.rs
+2 -2
View File
@@ -7,7 +7,7 @@
//! `MockWorker` backends (CPU-only, no GPU).
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -63,7 +63,7 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc<AppContext
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
experimental/sgl-router/tests/proxy/timeout.rs
+2 -2
View File
@@ -14,7 +14,7 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
Config, DiscoveryBackend, InflightLoadConfig, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
@@ -55,7 +55,7 @@ fn config(_worker_url: &str) -> Config {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
router_inflight_load: InflightLoadConfig::default(),
}
}