diff --git a/experimental/sgl-router/monitoring/README.md b/experimental/sgl-router/monitoring/README.md index 1f4f9e65a13f..a8d997c75102 100644 --- a/experimental/sgl-router/monitoring/README.md +++ b/experimental/sgl-router/monitoring/README.md @@ -42,6 +42,8 @@ Families the router emits. The dashboard graphs all of them except the | `sgl_router_kv_event_batches_lost_total` | Counter | KV-event batches dropped in transit, from gaps in each publisher's sequence | | `sgl_router_kv_tree_accounting_errors_total` | Counter | Occupancy-bookkeeping contradictions, by `reason`. Always 0 on a correct tree | | `sgl_router_kv_tree_maintained` | Gauge | 1 when this router maintains its own KV tree, 0 under an external Indexer | +| `sgl_router_kv_bootstrap_peers` | Gauge | Ready sibling router replicas peer bootstrap could pull a tree snapshot from. Emitted only with `--kv-peer-selector` | +| `sgl_router_kv_bootstrap_peers_synced` | Gauge | 1 once peer discovery has reported at least once; a 0 that never becomes 1 means the peer watch is not delivering (check EndpointSlice RBAC). Emitted only with `--kv-peer-selector` | The legacy `sgl_router_overlap_blocks` metric was removed with the `cache_aware_zmq` policy and has no direct replacement. Remove queries, alerts, diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index a8981ea41a12..beb4e96831a4 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -3,11 +3,12 @@ //! Grouped CLI options and conversion into a validated [`Config`]. -use anyhow::{anyhow, ensure, Context, Result}; +use anyhow::{anyhow, bail, ensure, Context, Result}; use clap::Parser; use std::num::NonZeroU32; use crate::config::sampling::{parse_sampling_overrides, ConflictPolicy}; +use crate::config::types::is_selector_empty; 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, AffinityConfig, @@ -16,6 +17,7 @@ use crate::config::{ InflightLoadConfig, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE, + DEFAULT_KV_BOOTSTRAP_TIMEOUT_MS, }; const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100; @@ -164,6 +166,14 @@ pub struct DiscoveryArgs { /// Decode equality selector terms (key=value or key==value). Requires --prefill-selector. #[arg(long, num_args = 1..)] pub decode_selector: Vec, + + /// Label selector matching the EndpointSlices of this router's OWN Service + /// (the Service's labels, not the pods'), so a booting replica can find + /// siblings to pull a cache-aware tree snapshot from, e.g. + /// `kubernetes.io/service-name=sgl-router`. Unset disables peer bootstrap + /// and every replica starts cold. + #[arg(long)] + pub kv_peer_selector: Option, } #[derive(clap::Args, Debug)] @@ -238,6 +248,11 @@ pub struct CacheArgs { #[arg(long)] pub kv_indexer_query_max_inflight: Option, + /// Peer-bootstrap deadline in milliseconds. Requires --kv-peer-selector. + /// Defaults to 600000 (10 minutes). + #[arg(long)] + pub kv_bootstrap_timeout_ms: Option, + /// Minimum cache-hit tokens for a candidate. Defaults to 1024. #[arg(long)] pub cache_affinity_min_matched_tokens: Option, @@ -371,7 +386,27 @@ impl Cli { .map(load_bucket_config) .transpose()?; let circuit_breaker = self.routing.build_circuit_breaker()?; + let kv_bootstrap_timeout_ms = self.cache.kv_bootstrap_timeout_ms; let cache_aware = self.cache.into_config(self.routing.policy)?; + // Peer bootstrap grafts into this router's own radix tree, so it needs + // cache-aware over a local tree; checked here because it spans groups. + let has_peer_selector = discovery.peer_selector().is_some(); + ensure!( + has_peer_selector || kv_bootstrap_timeout_ms.is_none(), + "--kv-bootstrap-timeout-ms requires --kv-peer-selector, which is what \ + enables peer bootstrap" + ); + if has_peer_selector { + match cache_aware.as_ref().map(|c| c.prefix_provider) { + None => bail!("--kv-peer-selector requires --policy cache_aware"), + Some(p) if p != CachePrefixProvider::RadixTree => bail!( + "--kv-peer-selector requires --cache-prefix-provider radix_tree \ + (there is no local tree to bootstrap when an external Indexer is \ + the prefix source)" + ), + _ => {} + } + } let fused = self.routing.build_fused()?; let eligibility = self.routing.build_eligibility()?; let sticky = self.affinity.into_sticky_config(self.routing.policy)?; @@ -467,6 +502,14 @@ impl DiscoveryArgs { "--service-discovery-namespace / --selector / --prefill-selector / \ --decode-selector require --service-discovery" ); + // Peer replicas are found through EndpointSlices too, so with + // any other backend the flag would be accepted and then + // silently ignored. + ensure!( + self.kv_peer_selector.is_none(), + "--kv-peer-selector requires --service-discovery (peer replicas are \ + found via Kubernetes EndpointSlices)" + ); DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { urls: self.worker_urls, }) @@ -477,9 +520,19 @@ impl DiscoveryArgs { join_selector(&self.prefill_selector).as_deref(), join_selector(&self.decode_selector).as_deref(), )?; + // An empty selector matches every EndpointSlice in the watched + // namespace(s), turning unrelated Services into "siblings". + ensure!( + self.kv_peer_selector + .as_deref() + .is_none_or(|s| !is_selector_empty(s)), + "--kv-peer-selector must not be empty: an empty selector matches \ + every EndpointSlice in the watched namespace" + ); DiscoveryBackend::K8s(K8sDiscoveryConfig { namespace: self.service_discovery_namespace.unwrap_or_default(), mode, + peer_selector: self.kv_peer_selector, }) } }; @@ -630,6 +683,9 @@ impl CacheArgs { Ok(Some(CacheAwareConfig { prefix_provider: cache_prefix_provider, kv_indexer_endpoint, + bootstrap_timeout_ms: self + .kv_bootstrap_timeout_ms + .unwrap_or(DEFAULT_KV_BOOTSTRAP_TIMEOUT_MS), })) } } @@ -860,7 +916,9 @@ fn join_selector(terms: &[String]) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::config::{DiscoveryBackend, K8sDiscoveryMode, ScoreTermKind}; + use crate::config::{ + DiscoveryBackend, K8sDiscoveryMode, ScoreTermKind, MAX_KV_BOOTSTRAP_TIMEOUT_MS, + }; #[test] fn chat_routing_selects_policy_implementation_without_another_config() { @@ -1166,6 +1224,145 @@ mod tests { assert!(err.contains("require --service-discovery"), "got: {err}"); } + // ---- peer bootstrap ---- + + /// k8s discovery + cache-aware, the base every peer-bootstrap config needs. + fn peer_cfg(extra: &[&str]) -> Result { + let mut args = vec![ + "--service-discovery", + "--selector", + "app=sglang", + "--policy", + "cache_aware", + ]; + args.extend_from_slice(extra); + into_config_owned(with_model(&args)) + } + + #[test] + fn peer_selector_rides_on_the_k8s_backend() { + let c = peer_cfg(&[ + "--kv-peer-selector", + "app=sgl-router", + "--kv-bootstrap-timeout-ms", + "12000", + ]) + .unwrap(); + match &c.discovery { + DiscoveryBackend::K8s(k) => { + assert_eq!(k.peer_selector.as_deref(), Some("app=sgl-router")) + } + _ => panic!("expected k8s backend"), + } + assert_eq!( + c.model.cache_aware.as_ref().unwrap().bootstrap_timeout_ms, + 12_000, + ); + } + + #[test] + fn peer_bootstrap_is_off_by_default() { + let c = peer_cfg(&[]).unwrap(); + match &c.discovery { + DiscoveryBackend::K8s(k) => assert!(k.peer_selector.is_none()), + _ => panic!("expected k8s backend"), + } + assert_eq!( + c.model.cache_aware.as_ref().unwrap().bootstrap_timeout_ms, + DEFAULT_KV_BOOTSTRAP_TIMEOUT_MS, + ); + } + + /// Peers are found through EndpointSlices, so without k8s discovery the + /// flag would be accepted and then silently ignored. + #[test] + fn rejects_peer_selector_without_service_discovery() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://w:30000", + "--policy", + "cache_aware", + "--kv-peer-selector", + "app=sgl-router", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("requires --service-discovery"), "got: {err}"); + } + + #[test] + fn rejects_peer_selector_without_cache_aware_policy() { + let err = into_config_owned(with_model(&[ + "--service-discovery", + "--selector", + "app=sglang", + "--kv-peer-selector", + "app=sgl-router", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("--policy cache_aware"), "got: {err}"); + } + + /// With an external Indexer there is no local tree to graft into, so the + /// flag would be accepted and then do nothing. + #[test] + fn rejects_peer_selector_with_an_external_indexer() { + let err = peer_cfg(&[ + "--cache-prefix-provider", + "indexer", + "--kv-indexer-endpoint", + "http://indexer:50051", + "--kv-peer-selector", + "app=sgl-router", + ]) + .unwrap_err() + .to_string(); + assert!( + err.contains("--cache-prefix-provider radix_tree"), + "got: {err}", + ); + } + + #[test] + fn kv_bootstrap_timeout_bounds() { + let parse = |ms: u64| { + peer_cfg(&[ + "--kv-peer-selector", + "app=sgl-router", + "--kv-bootstrap-timeout-ms", + &ms.to_string(), + ]) + }; + assert!(parse(MAX_KV_BOOTSTRAP_TIMEOUT_MS).is_ok()); + let err = parse(MAX_KV_BOOTSTRAP_TIMEOUT_MS + 1) + .unwrap_err() + .to_string(); + assert!(err.contains("ceiling"), "got: {err}"); + let err = parse(0).unwrap_err().to_string(); + assert!(err.contains("greater than zero"), "got: {err}"); + } + + /// Without a selector the timeout tunes a bootstrap that never runs. + #[test] + fn rejects_bootstrap_timeout_without_a_peer_selector() { + let err = peer_cfg(&["--kv-bootstrap-timeout-ms", "20000"]) + .unwrap_err() + .to_string(); + assert!(err.contains("requires --kv-peer-selector"), "got: {err}"); + } + + /// An empty selector matches every EndpointSlice in the namespace. + #[test] + fn rejects_an_empty_peer_selector() { + for selector in ["", " , "] { + let err = peer_cfg(&["--kv-peer-selector", selector]) + .unwrap_err() + .to_string(); + assert!(err.contains("must not be empty"), "{selector:?} got: {err}"); + } + } + #[test] fn rejects_static_urls_duplicate() { let err = into_config_owned(with_model(&[ diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index bebf3c88c637..957a3f8a6b0a 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -60,6 +60,18 @@ impl Config { can check it against the pod's real budget", self.server.shutdown_drain_secs, ); + if let Some(cache) = self.model.cache_aware.as_ref() { + let ms = cache.bootstrap_timeout_ms; + ensure!( + ms > 0, + "--kv-bootstrap-timeout-ms must be greater than zero; to boot cold, \ + leave --kv-peer-selector unset" + ); + ensure!( + ms <= MAX_KV_BOOTSTRAP_TIMEOUT_MS, + "--kv-bootstrap-timeout-ms {ms} exceeds the {MAX_KV_BOOTSTRAP_TIMEOUT_MS}ms ceiling" + ); + } match &self.discovery { DiscoveryBackend::StaticUrls(s) => { ensure!( diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index bdda3164ab43..934d2bfb000a 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -401,14 +401,38 @@ pub enum CachePrefixProvider { } /// Per-model Cache-Aware configuration. -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub struct CacheAwareConfig { /// Prefix-match source for native Cache-Aware. pub prefix_provider: CachePrefixProvider, /// External Indexer configuration when `prefix_provider = indexer`. pub kv_indexer_endpoint: Option, + /// Deadline for peer bootstrap, validated by `Config::validate`. Only + /// meaningful when a peer selector is set; see + /// [`K8sDiscoveryConfig::peer_selector`]. + pub bootstrap_timeout_ms: u64, +} + +impl Default for CacheAwareConfig { + fn default() -> Self { + Self { + prefix_provider: CachePrefixProvider::default(), + kv_indexer_endpoint: None, + bootstrap_timeout_ms: DEFAULT_KV_BOOTSTRAP_TIMEOUT_MS, + } + } } +/// Budget for the whole peer-bootstrap sweep: 10 minutes, room for several +/// attempts at a fleet-sized snapshot (the producer's export build, the +/// transfer and the graft). Readiness waits on it, so a pod's startup or +/// readiness probe must tolerate a replica that stays unready this long. +pub const DEFAULT_KV_BOOTSTRAP_TIMEOUT_MS: u64 = 600_000; + +/// Ceiling on `--kv-bootstrap-timeout-ms` (1 hour); past it, `Instant + +/// Duration` can overflow and panic. +pub const MAX_KV_BOOTSTRAP_TIMEOUT_MS: u64 = 3_600_000; + /// Default request header for sticky routing. pub const DEFAULT_STICKY_HEADER: &str = "x-sgl-routing-key"; @@ -568,6 +592,17 @@ pub enum DiscoveryBackend { K8s(K8sDiscoveryConfig), } +impl DiscoveryBackend { + /// The peer-bootstrap selector, which only the Kubernetes backend carries. + /// `Some` is what enables peer bootstrap. + pub fn peer_selector(&self) -> Option<&str> { + match self { + Self::K8s(k) => k.peer_selector.as_deref(), + Self::StaticUrls(_) => None, + } + } +} + /// Workers registered at startup. Roles, models, and bootstrap ports come from /// `/server_info`; topology changes require a restart. #[derive(Debug, Clone)] @@ -582,6 +617,24 @@ pub struct K8sDiscoveryConfig { pub namespace: String, /// Resolved + validated selector mode (plain vs PD). pub mode: K8sDiscoveryMode, + /// Label selector matching the EndpointSlices of this router's OWN + /// Service, so a booting replica can find sibling replicas to pull a + /// cache-aware tree snapshot from. + /// + /// Matched against EndpointSlice labels, exactly like the worker + /// selectors: those are the Service's labels plus + /// `kubernetes.io/service-name`, never the pods' labels. A pod-template + /// label matches nothing, which reads as "no siblings" and boots every + /// replica cold; `kubernetes.io/service-name=` is the + /// unambiguous choice. + /// + /// `None` disables peer bootstrap. Scoped to the same namespace as + /// `namespace` (so an empty `namespace` means every namespace), which is + /// why this lives here rather than on [`crate::config::CacheAwareConfig`]. + /// + /// Requires the router's ServiceAccount to have `list`/`watch` on + /// EndpointSlices in that namespace. + pub peer_selector: Option, } /// Validated selector mode. Plain selectors run server-side; PD selectors @@ -639,7 +692,7 @@ pub enum ConfigError { } /// An empty selector matches every slice, which is invalid for a PD role. -fn is_selector_empty(selector: &str) -> bool { +pub(crate) fn is_selector_empty(selector: &str) -> bool { selector.split(',').all(|t| t.trim().is_empty()) } diff --git a/experimental/sgl-router/src/discovery/k8s.rs b/experimental/sgl-router/src/discovery/k8s.rs index 561efd5759b3..d4c9ede9a50d 100644 --- a/experimental/sgl-router/src/discovery/k8s.rs +++ b/experimental/sgl-router/src/discovery/k8s.rs @@ -335,7 +335,12 @@ pub async fn spawn( ) -> Result> { // The mode was resolved + validated at construction (`resolve_mode` in // `Cli::build_discovery`); just destructure it here. - let K8sDiscoveryConfig { namespace, mode } = cfg; + // `peer_selector` names sibling routers, not the workers streamed here. + let K8sDiscoveryConfig { + namespace, + mode, + peer_selector: _, + } = cfg; let client = Client::try_default() .await diff --git a/experimental/sgl-router/src/server/routes/metrics.rs b/experimental/sgl-router/src/server/routes/metrics.rs index 3a647ea027b3..4c3e2c860c1d 100644 --- a/experimental/sgl-router/src/server/routes/metrics.rs +++ b/experimental/sgl-router/src/server/routes/metrics.rs @@ -12,6 +12,7 @@ use crate::discovery::WorkerMode; use crate::server::app_context::AppContext; use crate::server::metrics::{escape_label, WorkerSnapshot}; +use crate::state::kv_events::bootstrap::PeerRegistry; use crate::state::kv_events::{KvIndexMetrics, Tiers, ACCOUNTING_REASONS}; use axum::extract::State; use axum::http::header::CONTENT_TYPE; @@ -71,6 +72,10 @@ pub async fn metrics(State(ctx): State>) -> impl IntoResponse { ctx.block_size_oracle.get().unwrap_or(0), )); } + // A local tree alone does not mean peer bootstrap is on; the selector does. + if let (Some(index), Some(_)) = (ctx.kv_index.as_ref(), ctx.config.discovery.peer_selector()) { + body.push_str(&render_kv_peers(&index.peers())); + } ( StatusCode::OK, [(CONTENT_TYPE, PROMETHEUS_CONTENT_TYPE)], @@ -78,6 +83,27 @@ pub async fn metrics(State(ctx): State>) -> impl IntoResponse { ) } +/// Render the peer-discovery series, emitted only with a peer selector and a +/// local tree. `synced` separates "not told yet" from "told and alone". +fn render_kv_peers(peers: &PeerRegistry) -> String { + let (len, synced) = peers.len_and_synced(); + let mut out = String::new(); + out.push_str( + "# HELP sgl_router_kv_bootstrap_peers Ready sibling router replicas this replica could pull a cache-aware tree snapshot from. 0 with sgl_router_kv_bootstrap_peers_synced=1 means the fleet genuinely has no other ready replica; 0 with synced=0 means peer discovery has not reported yet (check RBAC for endpointslices on the router's own Service).\n", + ); + out.push_str("# TYPE sgl_router_kv_bootstrap_peers gauge\n"); + out.push_str(&format!("sgl_router_kv_bootstrap_peers {len}\n")); + out.push_str( + "# HELP sgl_router_kv_bootstrap_peers_synced 1 once peer discovery has reported at least once, even with an empty result. Until then an empty peer set is not evidence of anything.\n", + ); + out.push_str("# TYPE sgl_router_kv_bootstrap_peers_synced gauge\n"); + out.push_str(&format!( + "sgl_router_kv_bootstrap_peers_synced {}\n", + u8::from(synced), + )); + out +} + /// Render the storage-tier series: what the tree holds per worker and tier, /// the block size to convert it to tokens, and the tagged event stream it /// consumed. @@ -182,6 +208,90 @@ mod tests { use http_body_util::BodyExt; use tower::ServiceExt; + /// The two readings of a zero peer count an operator has to tell apart: + /// "discovery has not reported yet" and "this replica is genuinely alone". + /// Only `synced` separates them, so both series are contract. + #[tokio::test] + async fn kv_peer_series_distinguish_unsynced_from_alone() { + let peers = PeerRegistry::new(); + let before = render_kv_peers(&peers); + assert!(before.contains("sgl_router_kv_bootstrap_peers 0\n")); + assert!(before.contains("sgl_router_kv_bootstrap_peers_synced 0\n")); + + peers.replace(vec![]); + let alone = render_kv_peers(&peers); + assert!(alone.contains("sgl_router_kv_bootstrap_peers 0\n")); + assert!( + alone.contains("sgl_router_kv_bootstrap_peers_synced 1\n"), + "a synced empty set is a different fact from an unsynced one", + ); + + peers.replace(vec!["http://a:30000".into(), "http://b:30000".into()]); + let warm = render_kv_peers(&peers); + assert!(warm.contains("sgl_router_kv_bootstrap_peers 2\n")); + assert!(warm.contains("sgl_router_kv_bootstrap_peers_synced 1\n")); + } + + /// A router that cannot peer-bootstrap — no local tree, or no selector — + /// emits no series rather than a confidently wrong 0. The live wiring is + /// part of the claim, hence the full scrape. + #[tokio::test] + async fn kv_peer_series_require_a_peer_selector() { + use crate::config::{DiscoveryBackend, K8sDiscoveryConfig, K8sDiscoveryMode}; + use crate::state::kv_events::KvEventIndex; + + let scrape = |ctx: Arc| async move { + let res = crate::server::app::build_router(ctx) + .oneshot( + Request::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = res.into_body().collect().await.unwrap().to_bytes(); + String::from_utf8(body.to_vec()).unwrap() + }; + let with_selector = || { + DiscoveryBackend::K8s(K8sDiscoveryConfig { + namespace: "ns".into(), + mode: K8sDiscoveryMode::Plain { + label_selector: "app=sglang".into(), + }, + peer_selector: Some("kubernetes.io/service-name=sgl-router".into()), + }) + }; + + let mut ctx = AppContext::stub(); + ctx.config.discovery = with_selector(); + assert!( + !scrape(Arc::new(ctx)) + .await + .contains("sgl_router_kv_bootstrap_peers"), + "a router with no local tree emits no peer series at all", + ); + + let index = KvEventIndex::new(); + index.peers().replace(vec!["http://sibling:30000".into()]); + + let mut ctx = AppContext::stub(); + ctx.kv_index = index.snapshot_source(); + assert!( + !scrape(Arc::new(ctx)) + .await + .contains("sgl_router_kv_bootstrap_peers"), + "a router without --kv-peer-selector emits no peer series", + ); + + let mut ctx = AppContext::stub(); + ctx.kv_index = index.snapshot_source(); + ctx.config.discovery = with_selector(); + let body = scrape(Arc::new(ctx)).await; + assert!(body.contains("sgl_router_kv_bootstrap_peers 1\n")); + assert!(body.contains("sgl_router_kv_bootstrap_peers_synced 1\n")); + } + /// The tier series are what a coverage dashboard joins on, so their names /// and label keys are contract: per-worker blocks by tier with zeros /// emitted, the block size to convert them, and every (event, medium) diff --git a/experimental/sgl-router/src/state/kv_events/bootstrap.rs b/experimental/sgl-router/src/state/kv_events/bootstrap.rs index 94bae78b3a59..4308a032e318 100644 --- a/experimental/sgl-router/src/state/kv_events/bootstrap.rs +++ b/experimental/sgl-router/src/state/kv_events/bootstrap.rs @@ -2,7 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 //! Peer-snapshot bootstrap for the KV-event tree: the wire shape -//! ([`PeerSnapshot`]) and its constants. +//! ([`PeerSnapshot`]), its constants, and the registry of peers a snapshot +//! may be pulled from. //! //! A freshly started replica subscribes to each worker's KV topic mid-stream: //! ZMQ SUB delivers deltas from whatever sequence the publisher has reached, @@ -29,6 +30,10 @@ use serde::{Deserialize, Serialize}; use super::tree::{KvWorkerId, SnapshotNode}; +mod peers; + +pub use peers::PeerRegistry; + /// Wire-format version. Bump on any incompatible change to [`PeerSnapshot`]. pub const SNAPSHOT_FORMAT: u32 = 1; diff --git a/experimental/sgl-router/src/state/kv_events/bootstrap/peers.rs b/experimental/sgl-router/src/state/kv_events/bootstrap/peers.rs new file mode 100644 index 000000000000..2ef1ca6bc673 --- /dev/null +++ b/experimental/sgl-router/src/state/kv_events/bootstrap/peers.rs @@ -0,0 +1,130 @@ +//! The sibling-replica set a snapshot may be pulled from. + +use parking_lot::Mutex; +use rand::seq::SliceRandom; +use tracing::info; + +/// Sibling replicas a snapshot may be pulled from. An empty set is conclusive +/// only after a sync and only if siblings were never seen: before the first +/// delivery it means "not told yet", and after siblings were seen it is a +/// transient dip (slice repack, rolling update). All state is behind one lock +/// so a reader never pairs fields from different writes. +#[derive(Debug, Default)] +pub struct PeerRegistry { + state: Mutex, +} + +#[derive(Debug, Default)] +struct PeerState { + peers: Vec, + /// Whether a peer set has been delivered at least once, even an empty one. + synced: bool, + /// Whether a non-empty peer set has ever been delivered. + ever_had_peers: bool, +} + +impl PeerRegistry { + pub fn new() -> Self { + Self::default() + } + + /// Replace the peer set wholesale, as a relist or a slice event produces. + pub fn replace(&self, peers: Vec) { + let changed = { + let mut state = self.state.lock(); + state.synced = true; + state.ever_had_peers |= !peers.is_empty(); + let changed = (state.peers != peers).then(|| peers.clone()); + state.peers = peers; + changed + }; + if let Some(peers) = changed { + info!(count = peers.len(), peers = ?peers, "kv-bootstrap: peer set updated"); + } + } + + pub fn len(&self) -> usize { + self.state.lock().peers.len() + } + + pub fn is_empty(&self) -> bool { + self.state.lock().peers.is_empty() + } + + /// Whether peer discovery has reported at least once. An empty peer set is + /// only conclusive once this is true. + pub fn synced(&self) -> bool { + self.state.lock().synced + } + + /// Peer count and `synced`, read together. + pub fn len_and_synced(&self) -> (usize, bool) { + let state = self.state.lock(); + (state.peers.len(), state.synced) + } + + /// True when discovery has confirmed this replica has no siblings, so + /// there is no point waiting for one. + pub fn known_to_have_no_peers(&self) -> bool { + let state = self.state.lock(); + state.synced && state.peers.is_empty() && !state.ever_had_peers + } + + /// Candidate peers in shuffled order, so simultaneous boots spread their + /// snapshot fetches instead of stampeding whichever peer sorts first. + pub fn candidates(&self) -> Vec { + let mut peers = self.state.lock().peers.clone(); + peers.shuffle(&mut rand::thread_rng()); + peers + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_transient_empty_peer_set_is_not_conclusive() { + let reg = PeerRegistry::new(); + assert!( + !reg.known_to_have_no_peers(), + "before any sync, empty means the watch has not delivered", + ); + + reg.replace(vec!["http://a:30000".into()]); + assert!(reg.synced()); + assert!(!reg.known_to_have_no_peers()); + + // An EndpointSlice repack, or a rolling update with every sibling + // notReady, empties the set for an instant. Having once seen siblings + // is what keeps that from reading as "I am alone". + reg.replace(vec![]); + assert!(reg.is_empty()); + assert!( + !reg.known_to_have_no_peers(), + "a dip after siblings were seen is transient, not conclusive", + ); + } + + #[test] + fn a_synced_empty_set_with_no_history_means_no_siblings() { + let reg = PeerRegistry::new(); + reg.replace(vec![]); + assert!(reg.synced()); + assert!(reg.known_to_have_no_peers()); + } + + #[test] + fn candidates_shuffle_without_losing_entries() { + let reg = PeerRegistry::new(); + let peers: Vec = (0..16).map(|i| format!("http://p{i}:30000")).collect(); + reg.replace(peers.clone()); + + let mut got = reg.candidates(); + assert_eq!(got.len(), peers.len()); + got.sort(); + let mut want = peers; + want.sort(); + assert_eq!(got, want, "shuffling must not drop or duplicate a peer"); + } +} diff --git a/experimental/sgl-router/src/state/kv_events/index.rs b/experimental/sgl-router/src/state/kv_events/index.rs index 8d8f2c46ad8f..5759aa61149d 100644 --- a/experimental/sgl-router/src/state/kv_events/index.rs +++ b/experimental/sgl-router/src/state/kv_events/index.rs @@ -37,6 +37,7 @@ use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; use super::block_size_oracle::BlockSizeOracle; +use super::bootstrap::PeerRegistry; use super::discovery::{fetch_event_config, EventConfig}; use super::subscriber::{KvEventSubscriberRegistry, SubKind, WorkerEvent}; use super::tally::{EventKind, EventTally}; @@ -130,6 +131,8 @@ pub struct KvEventIndex { /// Applied events by kind and storage medium, for the `/metrics` scrape. /// Written only by the pump. tally: Arc, + /// Sibling replicas a snapshot may be pulled from. See [`PeerRegistry`]. + peers: Arc, /// Last built snapshot; see [`KvEventIndex::peer_snapshot_body`]. snapshot_cache: Arc>>, /// Worker-sourced `page_size` shared with prefix providers. @@ -214,6 +217,7 @@ impl KvEventIndex { live_workers, cursors, tally, + peers: Arc::new(PeerRegistry::new()), snapshot_cache: Arc::new(AsyncMutex::new(None)), block_size_oracle, }) @@ -247,6 +251,11 @@ impl KvEventIndex { }) } + /// Shared handle to the peer registry. + pub fn peers(&self) -> Arc { + Arc::clone(&self.peers) + } + /// 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. diff --git a/experimental/sgl-router/tests/proxy/bucket_routing.rs b/experimental/sgl-router/tests/proxy/bucket_routing.rs index 214498eb84b0..1eb9cbca5ece 100644 --- a/experimental/sgl-router/tests/proxy/bucket_routing.rs +++ b/experimental/sgl-router/tests/proxy/bucket_routing.rs @@ -205,6 +205,7 @@ fn build_cache_ctx_with_affinity( query_timeout_ms: 100, query_max_inflight: 32, }), + ..Default::default() }); context.prefix_index = Some(prefix_index); context.block_size_oracle.try_set(1).unwrap();