From 4a9dc5c4af9631b658c6f9d61129f5d2ac32e500 Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Sun, 20 Sep 2026 16:49:23 -0700 Subject: [PATCH] [sgl-router] refactor - move policy-required states under src/state (#40272) Co-authored-by: Claude Fable 5.1 --- .../sgl-router/benches/tree_lookup.rs | 2 +- experimental/sgl-router/src/config/cli.rs | 14 +- experimental/sgl-router/src/config/mod.rs | 2 +- experimental/sgl-router/src/config/types.rs | 6 +- experimental/sgl-router/src/discovery/k8s.rs | 2 +- experimental/sgl-router/src/lib.rs | 1 + experimental/sgl-router/src/main.rs | 33 +-- .../sgl-router/src/policies/admission.rs | 96 +++++--- .../sgl-router/src/policies/decode.rs | 10 +- .../sgl-router/src/policies/factory.rs | 6 +- .../sgl-router/src/policies/load_based.rs | 31 +-- experimental/sgl-router/src/policies/mod.rs | 23 +- .../sgl-router/src/policies/power_of_two.rs | 22 +- .../src/policies/prefix_provider.rs | 2 +- .../sgl-router/src/policies/registry.rs | 13 +- .../src/policies/scoring/admission.rs | 4 +- .../sgl-router/src/policies/scoring/argmax.rs | 7 +- .../sgl-router/src/policies/scoring/mod.rs | 12 +- .../src/policies/scoring/prefix_cache.rs | 6 +- .../sgl-router/src/policies/selection.rs | 26 ++- .../sgl-router/src/policies/session_aware.rs | 101 ++------ .../sgl-router/src/policies/sticky.rs | 218 ++++++------------ experimental/sgl-router/src/proxy/mod.rs | 2 +- experimental/sgl-router/src/proxy/sse.rs | 2 +- .../sgl-router/src/server/app_context.rs | 36 +-- .../sgl-router/src/server/inflight.rs | 2 +- experimental/sgl-router/src/server/metrics.rs | 31 +-- .../sgl-router/src/server/routes/chat.rs | 10 +- .../src/server/routes/chat/forward.rs | 8 +- .../sgl-router/src/server/routes/metrics.rs | 10 +- .../sgl-router/src/server/routes/tokenize.rs | 2 +- .../sgl-router/src/state/affinity_store.rs | 195 ++++++++++++++++ .../kv_events/block_size_oracle.rs | 2 +- .../kv_events/discovery.rs | 0 .../src/{policies => state}/kv_events/hash.rs | 0 .../{policies => state}/kv_events/index.rs | 57 ++--- .../src/{policies => state}/kv_events/mod.rs | 0 .../kv_events/subscriber.rs | 6 +- .../{policies => state}/kv_events/tally.rs | 0 .../src/{policies => state}/kv_events/tree.rs | 0 .../src/{policies => state}/kv_events/wire.rs | 0 .../load_monitor/engine_reported_load.rs} | 68 +++--- .../sgl-router/src/state/load_monitor/mod.rs | 7 + .../load_monitor/router_inflight_load.rs} | 70 +++--- experimental/sgl-router/src/state/mod.rs | 12 + experimental/sgl-router/src/tokenizer/mod.rs | 2 +- .../sgl-router/src/workers/introspect.rs | 4 +- .../sgl-router/src/workers/manager.rs | 40 ++-- .../sgl-router/src/workers/registry.rs | 4 +- experimental/sgl-router/src/workers/worker.rs | 24 +- .../tests/component/discovery/static_urls.rs | 5 +- .../policies/cache_prefix_provider.rs | 4 +- .../tests/component/policies/decode.rs | 10 +- .../tests/component/policies/fused_score.rs | 20 +- .../policies/kv_events_hash_parity.rs | 2 +- .../policies/kv_events_tree_concurrent.rs | 4 +- .../policies/kv_events_two_subscribers.rs | 4 +- .../tests/component/policies/zmq_helpers.rs | 2 +- .../component/tokenizer/render_parity.rs | 4 +- .../component/workers/concurrent_state.rs | 6 +- .../tests/component/workers/manager.rs | 2 +- .../sgl-router/tests/proxy/bucket_routing.rs | 10 +- .../tests/proxy/cache_aware_input_ids.rs | 2 +- .../sgl-router/tests/proxy/chat_routing.rs | 66 +++--- .../tests/proxy/common/cache_aware_fixture.rs | 6 +- .../tests/proxy/external_indexer_routing.rs | 2 +- .../sgl-router/tests/proxy/failover.rs | 2 +- .../tests/proxy/graceful_shutdown.rs | 8 +- .../tests/proxy/header_forwarding.rs | 4 +- .../tests/proxy/pd_bootstrap_injection.rs | 6 +- .../tests/proxy/pd_pool_isolation.rs | 4 +- .../tests/proxy/pd_protocol_binding.rs | 4 +- .../tests/proxy/radix_tree_routing.rs | 4 +- .../tests/proxy/roundrobin_input_ids.rs | 4 +- .../tests/proxy/shared_prefill_admission.rs | 14 +- .../tests/proxy/sticky_input_ids.rs | 4 +- .../sgl-router/tests/proxy/sticky_routing.rs | 4 +- .../sgl-router/tests/proxy/timeout.rs | 4 +- 78 files changed, 787 insertions(+), 655 deletions(-) create mode 100644 experimental/sgl-router/src/state/affinity_store.rs rename experimental/sgl-router/src/{policies => state}/kv_events/block_size_oracle.rs (98%) rename experimental/sgl-router/src/{policies => state}/kv_events/discovery.rs (100%) rename experimental/sgl-router/src/{policies => state}/kv_events/hash.rs (100%) rename experimental/sgl-router/src/{policies => state}/kv_events/index.rs (96%) rename experimental/sgl-router/src/{policies => state}/kv_events/mod.rs (100%) rename experimental/sgl-router/src/{policies => state}/kv_events/subscriber.rs (99%) rename experimental/sgl-router/src/{policies => state}/kv_events/tally.rs (100%) rename experimental/sgl-router/src/{policies => state}/kv_events/tree.rs (100%) rename experimental/sgl-router/src/{policies => state}/kv_events/wire.rs (100%) rename experimental/sgl-router/src/{policies/engine_load.rs => state/load_monitor/engine_reported_load.rs} (94%) create mode 100644 experimental/sgl-router/src/state/load_monitor/mod.rs rename experimental/sgl-router/src/{policies/active_load.rs => state/load_monitor/router_inflight_load.rs} (94%) create mode 100644 experimental/sgl-router/src/state/mod.rs diff --git a/experimental/sgl-router/benches/tree_lookup.rs b/experimental/sgl-router/benches/tree_lookup.rs index 35f7b1159..dddcd486b 100644 --- a/experimental/sgl-router/benches/tree_lookup.rs +++ b/experimental/sgl-router/benches/tree_lookup.rs @@ -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(); diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 581072783..30d5c34cf 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -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] diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index f5b84fd62..f086f4607 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -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(), } } diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index ed704ec28..58bae48da 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -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(), diff --git a/experimental/sgl-router/src/discovery/k8s.rs b/experimental/sgl-router/src/discovery/k8s.rs index e0230fc74..ec9598b22 100644 --- a/experimental/sgl-router/src/discovery/k8s.rs +++ b/experimental/sgl-router/src/discovery/k8s.rs @@ -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. diff --git a/experimental/sgl-router/src/lib.rs b/experimental/sgl-router/src/lib.rs index 469e1fddc..bb7a1252e 100644 --- a/experimental/sgl-router/src/lib.rs +++ b/experimental/sgl-router/src/lib.rs @@ -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; diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index 71e6ecdab..dd2c132e9 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -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 { } } -fn start_local_inflight_tracker(config: &Config) -> (Arc, 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, 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, engine_state: &Arc, - local_inflight_requests: &Arc, + local_inflight_requests: &Arc, ) -> 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, worker_registry: Arc, routing_policies: Arc, - local_inflight_requests: Arc, + local_inflight_requests: Arc, engine_state: &KvEventIndex, external_kv_indexer_client: Option>, ) -> Result> { @@ -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", ) diff --git a/experimental/sgl-router/src/policies/admission.rs b/experimental/sgl-router/src/policies/admission.rs index 75dce9554..6807cdd2e 100644 --- a/experimental/sgl-router/src/policies/admission.rs +++ b/experimental/sgl-router/src/policies/admission.rs @@ -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, ) -> 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], limit: Option, ) -> 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], ) -> 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, min_load_choices: usize, ) -> Option { @@ -431,7 +433,7 @@ pub fn resolve_prefill_admitted( range: &CandidateRange<'_>, proposal: &SelectionProposal, request_input_tokens: u64, - snapshot: &EngineLoadSnapshot, + snapshot: &EngineReportedLoadSnapshot, queue_limit: Option, ) -> Option { 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 { 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, 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, 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, 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, - basic_by_worker_id: HashMap, + by_worker_id: HashMap, + basic_by_worker_id: HashMap, local_active_by_worker_id: HashMap, 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>, ) -> Self { let workers: Vec<&Arc> = workers.into_iter().collect(); let local_active_by_worker_id: HashMap = 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], request_input_tokens: u64, - snapshot: &EngineLoadSnapshot, + snapshot: &EngineReportedLoadSnapshot, queue_limit: Option, ) -> Option<(Arc, 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, DecisionReason)> { let admitted = domain .workers @@ -940,7 +942,7 @@ fn decode_domain_fallback( pub(crate) fn compare_prefill_pressure( left: &Arc, right: &Arc, - 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, right: &Arc, - 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, backup: &Arc, hints: &GuardHints, - snapshot: &EngineLoadSnapshot, + snapshot: &EngineReportedLoadSnapshot, ) -> bool { if !hints.enable_pressure_guard { return false; @@ -1059,15 +1079,15 @@ mod tests { })) } - fn snapshot(entries: &[(&Arc, u64, u64, u64, u64)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn snapshot(entries: &[(&Arc, 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)), diff --git a/experimental/sgl-router/src/policies/decode.rs b/experimental/sgl-router/src/policies/decode.rs index cbfe5dc77..2edb82877 100644 --- a/experimental/sgl-router/src/policies/decode.rs +++ b/experimental/sgl-router/src/policies/decode.rs @@ -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 { if let Some(decision) = resolve_decode(domain, proposal, request_kv_tokens, snapshot) { return Some(decision); diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index 132755cb9..82a69604e 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -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 { 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(), } } diff --git a/experimental/sgl-router/src/policies/load_based.rs b/experimental/sgl-router/src/policies/load_based.rs index 2a6b70578..c26fd81e6 100644 --- a/experimental/sgl-router/src/policies/load_based.rs +++ b/experimental/sgl-router/src/policies/load_based.rs @@ -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, diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index 35095ee30..a154cb16c 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -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, 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, TestEngineLoad)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn snapshot(entries: &[(&Arc, 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), diff --git a/experimental/sgl-router/src/policies/power_of_two.rs b/experimental/sgl-router/src/policies/power_of_two.rs index 35c44d841..8f178b991 100644 --- a/experimental/sgl-router/src/policies/power_of_two.rs +++ b/experimental/sgl-router/src/policies/power_of_two.rs @@ -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], - snapshot: Option<&EngineLoadSnapshot>, + snapshot: Option<&EngineReportedLoadSnapshot>, choices: usize, queue_limit: Option, ) -> Option> { @@ -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], - snapshot: Option<&EngineLoadSnapshot>, + snapshot: Option<&EngineReportedLoadSnapshot>, queue_limit: Option, ) -> Cow<'w, [Arc]> { // 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], - snapshot: Option<&EngineLoadSnapshot>, + snapshot: Option<&EngineReportedLoadSnapshot>, choices: usize, ) -> Option<(Arc, Option>)> { 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 { @@ -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, u64)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn snapshot(entries: &[(&Arc, 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]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn mixed_estimate_snapshot(workers: &[Arc]) -> 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, diff --git a/experimental/sgl-router/src/policies/prefix_provider.rs b/experimental/sgl-router/src/policies/prefix_provider.rs index 7ec544cf0..bdbf0b632 100644 --- a/experimental/sgl-router/src/policies/prefix_provider.rs +++ b/experimental/sgl-router/src/policies/prefix_provider.rs @@ -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}; diff --git a/experimental/sgl-router/src/policies/registry.rs b/experimental/sgl-router/src/policies/registry.rs index 44c4bb5e3..e01530ae5 100644 --- a/experimental/sgl-router/src/policies/registry.rs +++ b/experimental/sgl-router/src/policies/registry.rs @@ -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 = healthy.iter().map(|w| w.active_load()).collect(); + let mut loads: Vec = 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 diff --git a/experimental/sgl-router/src/policies/scoring/admission.rs b/experimental/sgl-router/src/policies/scoring/admission.rs index 900d7047a..73c647ed8 100644 --- a/experimental/sgl-router/src/policies/scoring/admission.rs +++ b/experimental/sgl-router/src/policies/scoring/admission.rs @@ -23,7 +23,7 @@ impl Overloaded { impl EligibilityFilter for Overloaded { fn keep(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Vec { (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) } diff --git a/experimental/sgl-router/src/policies/scoring/argmax.rs b/experimental/sgl-router/src/policies/scoring/argmax.rs index 1d1fa4ee5..f64ac44d2 100644 --- a/experimental/sgl-router/src/policies/scoring/argmax.rs +++ b/experimental/sgl-router/src/policies/scoring/argmax.rs @@ -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 = 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]) diff --git a/experimental/sgl-router/src/policies/scoring/mod.rs b/experimental/sgl-router/src/policies/scoring/mod.rs index 6332d1a9d..23df9fcbd 100644 --- a/experimental/sgl-router/src/policies/scoring/mod.rs +++ b/experimental/sgl-router/src/policies/scoring/mod.rs @@ -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, u64, u64, u64, u64)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn snapshot(entries: &[(&Arc, 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, diff --git a/experimental/sgl-router/src/policies/scoring/prefix_cache.rs b/experimental/sgl-router/src/policies/scoring/prefix_cache.rs index 326802277..ea33dc875 100644 --- a/experimental/sgl-router/src/policies/scoring/prefix_cache.rs +++ b/experimental/sgl-router/src/policies/scoring/prefix_cache.rs @@ -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; diff --git a/experimental/sgl-router/src/policies/selection.rs b/experimental/sgl-router/src/policies/selection.rs index 92fc274b3..a7a996046 100644 --- a/experimental/sgl-router/src/policies/selection.rs +++ b/experimental/sgl-router/src/policies/selection.rs @@ -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], pub ttft_slo_ms: Option, pub tps_slo: Option, @@ -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, u64, u64)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn snapshot(entries: &[(&Arc, 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, u64, u64, u64)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( + fn queued_snapshot(entries: &[(&Arc, 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], - 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], - load_snapshot: Option<&'a EngineLoadSnapshot>, + load_snapshot: Option<&'a EngineReportedLoadSnapshot>, request_input_tokens: u64, ) -> DecodeSelectionInputs<'a> { DecodeSelectionInputs { diff --git a/experimental/sgl-router/src/policies/session_aware.rs b/experimental/sgl-router/src/policies/session_aware.rs index 067f63e91..4fe6699af 100644 --- a/experimental/sgl-router/src/policies/session_aware.rs +++ b/experimental/sgl-router/src/policies/session_aware.rs @@ -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, - clock: Arc, - 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, + store: Arc, config: AffinityConfig, _janitor: Option, } 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) -> Self { + fn with_clock( + config: AffinityConfig, + clock: Arc, + ) -> 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}; diff --git a/experimental/sgl-router/src/policies/sticky.rs b/experimental/sgl-router/src/policies/sticky.rs index 1ab768215..7df3f4ece 100644 --- a/experimental/sgl-router/src/policies/sticky.rs +++ b/experimental/sgl-router/src/policies/sticky.rs @@ -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, - clock: Arc, - 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, + /// Selector for keyless requests and for the initial pin of a new key. + fallback: Arc, + /// Set once via `Policy::attach_metrics`; recording is a no-op until then. metrics: OnceLock>, + /// Background idle-eviction sweeper; `None` outside a Tokio runtime. + _janitor: Option, } -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) -> 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, + clock: Arc, + ) -> 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, - /// Selector for keyless requests and for the initial pin of a new key. - fallback: Arc, - /// 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, -} - -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) -> 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, clock: Arc) -> 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], ctx: &SelectionContext<'_>) -> Option> { 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) { - 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")); diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index 69542844d..227af7948 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -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. diff --git a/experimental/sgl-router/src/proxy/sse.rs b/experimental/sgl-router/src/proxy/sse.rs index 700c5836a..3ce5e8ea0 100644 --- a/experimental/sgl-router/src/proxy/sse.rs +++ b/experimental/sgl-router/src/proxy/sse.rs @@ -88,7 +88,7 @@ impl ErrorEventScanner { /// pump finishes (stream exhausted, client disconnects, or upstream errors). /// The opaque `Box` 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 diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 029e70b01..eaeb3da1a 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -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, /// Per-worker active-load bookkeeping shared by the proxy, policies, /// timeout janitor, and metrics. - pub active_load: Arc, + pub router_inflight_load: Arc, /// 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, /// Shared Engine LoadStat table; ingress captures one immutable snapshot per request. - pub engine_load: Arc, + pub engine_reported_load: Arc, pub prefix_index: Option>, pub radix_tree_prefix_provider: Option, pub block_size_oracle: Arc, /// 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, /// 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, readiness: AtomicU8, @@ -67,34 +67,34 @@ impl AppContext { registry: Arc, policies: Arc, ) -> 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, proxy: Arc, registry: Arc, policies: Arc, - active_load: Arc, + router_inflight_load: Arc, ) -> 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), } diff --git a/experimental/sgl-router/src/server/inflight.rs b/experimental/sgl-router/src/server/inflight.rs index e8a58a709..daf27d9cf 100644 --- a/experimental/sgl-router/src/server/inflight.rs +++ b/experimental/sgl-router/src/server/inflight.rs @@ -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 diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index 79d383726..d153c0b61 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -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>, ttft_seconds: Mutex>, stream_outcome_total: Mutex>>, - active_load: Mutex>>, + router_inflight_load: Mutex>>, stale_requests_total: Mutex>>, decode_affinity_total: Mutex>>, sticky_total: Mutex>>, @@ -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"#, diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index 4dc0a959c..7d4d72673 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -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], -) -> Option { +) -> Option { 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, - load_snapshot: Option, + load_snapshot: Option, ttft_slo_ms: Option, tps_slo: Option, routing_key: Option<&'a str>, diff --git a/experimental/sgl-router/src/server/routes/chat/forward.rs b/experimental/sgl-router/src/server/routes/chat/forward.rs index 7653a1b09..f07fa64e9 100644 --- a/experimental/sgl-router/src/server/routes/chat/forward.rs +++ b/experimental/sgl-router/src/server/routes/chat/forward.rs @@ -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) diff --git a/experimental/sgl-router/src/server/routes/metrics.rs b/experimental/sgl-router/src/server/routes/metrics.rs index 23e09cfcc..3a647ea02 100644 --- a/experimental/sgl-router/src/server/routes/metrics.rs +++ b/experimental/sgl-router/src/server/routes/metrics.rs @@ -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>) -> 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()); diff --git a/experimental/sgl-router/src/server/routes/tokenize.rs b/experimental/sgl-router/src/server/routes/tokenize.rs index 46131aaf7..9dd03714a 100644 --- a/experimental/sgl-router/src/server/routes/tokenize.rs +++ b/experimental/sgl-router/src/server/routes/tokenize.rs @@ -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( diff --git a/experimental/sgl-router/src/state/affinity_store.rs b/experimental/sgl-router/src/state/affinity_store.rs new file mode 100644 index 000000000..f51ec68f7 --- /dev/null +++ b/experimental/sgl-router/src/state/affinity_store.rs @@ -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, + clock: Arc, + idle: Duration, +} + +impl AffinityStore { + pub fn new(idle: Duration) -> Arc { + Self::with_clock(idle, Arc::new(SystemTimeClock)) + } + + pub fn with_clock(idle: Duration, clock: Arc) -> Arc { + 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]) -> Option<&'e Arc> { + 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 { + 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, + engines: &'e [Arc], + ) -> &'e Arc { + 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, interval: Duration) -> Option { + 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 { + 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); + } +} diff --git a/experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs b/experimental/sgl-router/src/state/kv_events/block_size_oracle.rs similarity index 98% rename from experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs rename to experimental/sgl-router/src/state/kv_events/block_size_oracle.rs index ef4c3ec41..d3346734b 100644 --- a/experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs +++ b/experimental/sgl-router/src/state/kv_events/block_size_oracle.rs @@ -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 diff --git a/experimental/sgl-router/src/policies/kv_events/discovery.rs b/experimental/sgl-router/src/state/kv_events/discovery.rs similarity index 100% rename from experimental/sgl-router/src/policies/kv_events/discovery.rs rename to experimental/sgl-router/src/state/kv_events/discovery.rs diff --git a/experimental/sgl-router/src/policies/kv_events/hash.rs b/experimental/sgl-router/src/state/kv_events/hash.rs similarity index 100% rename from experimental/sgl-router/src/policies/kv_events/hash.rs rename to experimental/sgl-router/src/state/kv_events/hash.rs diff --git a/experimental/sgl-router/src/policies/kv_events/index.rs b/experimental/sgl-router/src/state/kv_events/index.rs similarity index 96% rename from experimental/sgl-router/src/policies/kv_events/index.rs rename to experimental/sgl-router/src/state/kv_events/index.rs index f5936cf63..a6127e484 100644 --- a/experimental/sgl-router/src/policies/kv_events/index.rs +++ b/experimental/sgl-router/src/state/kv_events/index.rs @@ -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, /// 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, /// Engine-reported per-worker load, written by the pump from /// `WorkerEvent::Load` and captured at request ingress. - engine_load: Arc, + engine_reported_load: Arc, pump: Mutex>>, pump_cancel: CancellationToken, workers: Mutex>, @@ -177,14 +177,14 @@ impl KvEventIndex { let (tx, rx) = mpsc::channel::(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>> = Arc::new(Mutex::new(HashMap::new())); let live_workers: Arc>> = 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 { - Arc::clone(&self.engine_load) + pub fn engine_reported_load(&self) -> Arc { + 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, - engine_load: Arc, + engine_reported_load: Arc, cursors: Arc>>, live_workers: Arc>>, tally: Arc, @@ -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, - engine_load: Arc, + engine_reported_load: Arc, cursors: Arc>>, tally: Arc, #[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>> = 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; } diff --git a/experimental/sgl-router/src/policies/kv_events/mod.rs b/experimental/sgl-router/src/state/kv_events/mod.rs similarity index 100% rename from experimental/sgl-router/src/policies/kv_events/mod.rs rename to experimental/sgl-router/src/state/kv_events/mod.rs diff --git a/experimental/sgl-router/src/policies/kv_events/subscriber.rs b/experimental/sgl-router/src/state/kv_events/subscriber.rs similarity index 99% rename from experimental/sgl-router/src/policies/kv_events/subscriber.rs rename to experimental/sgl-router/src/state/kv_events/subscriber.rs index 7e18382c0..342050eac 100644 --- a/experimental/sgl-router/src/policies/kv_events/subscriber.rs +++ b/experimental/sgl-router/src/state/kv_events/subscriber.rs @@ -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::*; diff --git a/experimental/sgl-router/src/policies/kv_events/tally.rs b/experimental/sgl-router/src/state/kv_events/tally.rs similarity index 100% rename from experimental/sgl-router/src/policies/kv_events/tally.rs rename to experimental/sgl-router/src/state/kv_events/tally.rs diff --git a/experimental/sgl-router/src/policies/kv_events/tree.rs b/experimental/sgl-router/src/state/kv_events/tree.rs similarity index 100% rename from experimental/sgl-router/src/policies/kv_events/tree.rs rename to experimental/sgl-router/src/state/kv_events/tree.rs diff --git a/experimental/sgl-router/src/policies/kv_events/wire.rs b/experimental/sgl-router/src/state/kv_events/wire.rs similarity index 100% rename from experimental/sgl-router/src/policies/kv_events/wire.rs rename to experimental/sgl-router/src/state/kv_events/wire.rs diff --git a/experimental/sgl-router/src/policies/engine_load.rs b/experimental/sgl-router/src/state/load_monitor/engine_reported_load.rs similarity index 94% rename from experimental/sgl-router/src/policies/engine_load.rs rename to experimental/sgl-router/src/state/load_monitor/engine_reported_load.rs index 208014be8..effbb2744 100644 --- a/experimental/sgl-router/src/policies/engine_load.rs +++ b/experimental/sgl-router/src/state/load_monitor/engine_reported_load.rs @@ -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, } -/// 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, - native_cache_workers: HashMap, + workers: HashMap, + native_cache_workers: HashMap, } -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) -> Self { + /// Production requests should use [`EngineReportedLoadTable::capture_snapshot`]. + pub fn from_workers(version: u64, workers: HashMap) -> 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, + workers: HashMap, ) -> 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; /// 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 { 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 { + fn fresh_worker_loads(&self, now: Instant) -> HashMap { // url -> rank -> (reported load, fresh, timestamp). let mut observed: HashMap> = 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 { + ) -> HashMap { let mut observed: HashMap = 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); diff --git a/experimental/sgl-router/src/state/load_monitor/mod.rs b/experimental/sgl-router/src/state/load_monitor/mod.rs new file mode 100644 index 000000000..3d5aeee18 --- /dev/null +++ b/experimental/sgl-router/src/state/load_monitor/mod.rs @@ -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; diff --git a/experimental/sgl-router/src/policies/active_load.rs b/experimental/sgl-router/src/state/load_monitor/router_inflight_load.rs similarity index 94% rename from experimental/sgl-router/src/policies/active_load.rs rename to experimental/sgl-router/src/state/load_monitor/router_inflight_load.rs index 1f4a223b3..8d609a0d9 100644 --- a/experimental/sgl-router/src/policies/active_load.rs +++ b/experimental/sgl-router/src/state/load_monitor/router_inflight_load.rs @@ -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>, requests: DashMap, clock: Arc, @@ -185,8 +185,8 @@ pub struct ActiveLoadRegistry { metrics: Mutex>>, } -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, 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` (cloned from the shared one held in +/// `Arc` (cloned from the shared one held in /// `AppContext`). -pub fn spawn_janitor(registry: Arc, interval: Duration) -> JanitorHandle { +pub fn spawn_janitor( + registry: Arc, + 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>, +pub struct RouterInflightLoadGuard { + registry: Option>, /// `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`) 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, Arc) { + fn registry_with_mock_clock( + timeout: Duration, + ) -> (Arc, Arc) { let clock = Arc::new(MockClock::new(Instant::now())); - let registry = ActiveLoadRegistry::new(Arc::clone(&clock) as Arc, timeout); + let registry = + RouterInflightLoadRegistry::new(Arc::clone(&clock) as Arc, 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 = 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 = 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(®istry), 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] diff --git a/experimental/sgl-router/src/state/mod.rs b/experimental/sgl-router/src/state/mod.rs new file mode 100644 index 000000000..62491d493 --- /dev/null +++ b/experimental/sgl-router/src/state/mod.rs @@ -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; diff --git a/experimental/sgl-router/src/tokenizer/mod.rs b/experimental/sgl-router/src/tokenizer/mod.rs index 9dd594da8..73909b9ce 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -186,7 +186,7 @@ mod tests { }, ), proxy: crate::config::ProxyConfig::default(), - active_load: crate::config::ActiveLoadConfig::default(), + router_inflight_load: crate::config::InflightLoadConfig::default(), } } diff --git a/experimental/sgl-router/src/workers/introspect.rs b/experimental/sgl-router/src/workers/introspect.rs index 3ad6583c3..10e8bd950 100644 --- a/experimental/sgl-router/src/workers/introspect.rs +++ b/experimental/sgl-router/src/workers/introspect.rs @@ -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, ); diff --git a/experimental/sgl-router/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index 58ac77bb9..4c753e41e 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -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, registry: Arc, cfg: Option>, kv_index: Option>, - active_load: Option>, + router_inflight_load: Option>, ) { 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, cfg: Option>, kv_index: Option>, - active_load: Option>, + router_inflight_load: Option>, introspector: Arc, ) { 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, cfg: Option>, kv_index: Option>, - active_load: Option>, + router_inflight_load: Option>, introspector: Arc, reconcile_interval: Duration, ) { @@ -226,7 +226,7 @@ pub async fn run_with_introspector_and_reconcile( ®istry, &cfg, &kv_index, - &active_load, + &router_inflight_load, &introspector, &mut pending, ) @@ -264,7 +264,7 @@ async fn handle_discovery_event( registry: &Arc, cfg: &Option>, kv_index: &Option>, - active_load: &Option>, + router_inflight_load: &Option>, introspector: &Arc, pending: &mut HashMap>, ) { @@ -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::(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); diff --git a/experimental/sgl-router/src/workers/registry.rs b/experimental/sgl-router/src/workers/registry.rs index 0a89c42fa..2fd6b0c8f 100644 --- a/experimental/sgl-router/src/workers/registry.rs +++ b/experimental/sgl-router/src/workers/registry.rs @@ -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() { diff --git a/experimental/sgl-router/src/workers/worker.rs b/experimental/sgl-router/src/workers/worker.rs index 27583d0c8..bf494ad14 100644 --- a/experimental/sgl-router/src/workers/worker.rs +++ b/experimental/sgl-router/src/workers/worker.rs @@ -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, diff --git a/experimental/sgl-router/tests/component/discovery/static_urls.rs b/experimental/sgl-router/tests/component/discovery/static_urls.rs index 57068a9e8..852e020b1 100644 --- a/experimental/sgl-router/tests/component/discovery/static_urls.rs +++ b/experimental/sgl-router/tests/component/discovery/static_urls.rs @@ -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()); diff --git a/experimental/sgl-router/tests/component/policies/cache_prefix_provider.rs b/experimental/sgl-router/tests/component/policies/cache_prefix_provider.rs index 62e23a0cc..45756b67f 100644 --- a/experimental/sgl-router/tests/component/policies/cache_prefix_provider.rs +++ b/experimental/sgl-router/tests/component/policies/cache_prefix_provider.rs @@ -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() { diff --git a/experimental/sgl-router/tests/component/policies/decode.rs b/experimental/sgl-router/tests/component/policies/decode.rs index 733bd20ef..41ee4f394 100644 --- a/experimental/sgl-router/tests/component/policies/decode.rs +++ b/experimental/sgl-router/tests/component/policies/decode.rs @@ -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 { })) } -fn snapshot(entries: &[(&Arc, u64, u64, u64, u64)]) -> EngineLoadSnapshot { - EngineLoadSnapshot::from_native_cache_workers( +fn snapshot(entries: &[(&Arc, 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, diff --git a/experimental/sgl-router/tests/component/policies/fused_score.rs b/experimental/sgl-router/tests/component/policies/fused_score.rs index 2a5f24f11..e5e8bc046 100644 --- a/experimental/sgl-router/tests/component/policies/fused_score.rs +++ b/experimental/sgl-router/tests/component/policies/fused_score.rs @@ -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, diff --git a/experimental/sgl-router/tests/component/policies/kv_events_hash_parity.rs b/experimental/sgl-router/tests/component/policies/kv_events_hash_parity.rs index a0520d141..4a1c41a78 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_hash_parity.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_hash_parity.rs @@ -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)] diff --git a/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs b/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs index b365b9afc..04ed5eb0f 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs @@ -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 { diff --git a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs index 6eb627a2f..313fc3cc0 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs @@ -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, diff --git a/experimental/sgl-router/tests/component/policies/zmq_helpers.rs b/experimental/sgl-router/tests/component/policies/zmq_helpers.rs index 9f0e81624..3ca1ad521 100644 --- a/experimental/sgl-router/tests/component/policies/zmq_helpers.rs +++ b/experimental/sgl-router/tests/component/policies/zmq_helpers.rs @@ -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. diff --git a/experimental/sgl-router/tests/component/tokenizer/render_parity.rs b/experimental/sgl-router/tests/component/tokenizer/render_parity.rs index f69220729..fcc6e7dc9 100644 --- a/experimental/sgl-router/tests/component/tokenizer/render_parity.rs +++ b/experimental/sgl-router/tests/component/tokenizer/render_parity.rs @@ -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() } diff --git a/experimental/sgl-router/tests/component/workers/concurrent_state.rs b/experimental/sgl-router/tests/component/workers/concurrent_state.rs index 3cf6eeaf6..4d187d273 100644 --- a/experimental/sgl-router/tests/component/workers/concurrent_state.rs +++ b/experimental/sgl-router/tests/component/workers/concurrent_state.rs @@ -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", ); diff --git a/experimental/sgl-router/tests/component/workers/manager.rs b/experimental/sgl-router/tests/component/workers/manager.rs index a8d8cf3b4..6be1e5a50 100644 --- a/experimental/sgl-router/tests/component/workers/manager.rs +++ b/experimental/sgl-router/tests/component/workers/manager.rs @@ -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", diff --git a/experimental/sgl-router/tests/proxy/bucket_routing.rs b/experimental/sgl-router/tests/proxy/bucket_routing.rs index f35f5ef6e..554e35d22 100644 --- a/experimental/sgl-router/tests/proxy/bucket_routing.rs +++ b/experimental/sgl-router/tests/proxy/bucket_routing.rs @@ -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 { diff --git a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs index 59d0666ca..9e21e044b 100644 --- a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs @@ -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; diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 000704faa..32f11a73a 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -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 = 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", ); diff --git a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs index bfbe55e6e..a2883e5cb 100644 --- a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs +++ b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs @@ -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(), } } diff --git a/experimental/sgl-router/tests/proxy/external_indexer_routing.rs b/experimental/sgl-router/tests/proxy/external_indexer_routing.rs index 4abed4aef..984bf30b0 100644 --- a/experimental/sgl-router/tests/proxy/external_indexer_routing.rs +++ b/experimental/sgl-router/tests/proxy/external_indexer_routing.rs @@ -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; diff --git a/experimental/sgl-router/tests/proxy/failover.rs b/experimental/sgl-router/tests/proxy/failover.rs index 9a3a97789..d2d4d6e92 100644 --- a/experimental/sgl-router/tests/proxy/failover.rs +++ b/experimental/sgl-router/tests/proxy/failover.rs @@ -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()); diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 6eae3320c..01eb76805 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -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 { 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, 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)] diff --git a/experimental/sgl-router/tests/proxy/header_forwarding.rs b/experimental/sgl-router/tests/proxy/header_forwarding.rs index bad3360f4..e29bd57e4 100644 --- a/experimental/sgl-router/tests/proxy/header_forwarding.rs +++ b/experimental/sgl-router/tests/proxy/header_forwarding.rs @@ -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()); diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index d0fd95422..fcbc8738f 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -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); diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index 2d09f97f4..23ed74189 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -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(), } } diff --git a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs index f310fb296..60bec8f38 100644 --- a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs +++ b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs @@ -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(), } } diff --git a/experimental/sgl-router/tests/proxy/radix_tree_routing.rs b/experimental/sgl-router/tests/proxy/radix_tree_routing.rs index f88da3cd3..b9875d591 100644 --- a/experimental/sgl-router/tests/proxy/radix_tree_routing.rs +++ b/experimental/sgl-router/tests/proxy/radix_tree_routing.rs @@ -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; diff --git a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs index fbd2ff073..a05371dcd 100644 --- a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs @@ -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(), } } diff --git a/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs b/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs index 52414e1aa..a275e2b8a 100644 --- a/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs +++ b/experimental/sgl-router/tests/proxy/shared_prefill_admission.rs @@ -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 { diff --git a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs index 8c20010f7..e6de41117 100644 --- a/experimental/sgl-router/tests/proxy/sticky_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/sticky_input_ids.rs @@ -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(), } } diff --git a/experimental/sgl-router/tests/proxy/sticky_routing.rs b/experimental/sgl-router/tests/proxy/sticky_routing.rs index 96ddad1e0..db3a066f6 100644 --- a/experimental/sgl-router/tests/proxy/sticky_routing.rs +++ b/experimental/sgl-router/tests/proxy/sticky_routing.rs @@ -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 Config { urls: vec!["http://placeholder:0".into()], }), proxy: ProxyConfig::default(), - active_load: ActiveLoadConfig::default(), + router_inflight_load: InflightLoadConfig::default(), } }