From 3c48c1e967f64a6b2b116b8f6719d51eb360eb55 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Tue, 15 Sep 2026 13:30:05 -0700 Subject: [PATCH] [Router] Expose the KV storage-tier stream and tree occupancy on /metrics (2/4) (#39109) Co-authored-by: Kangyan Zhou Co-authored-by: Claude Opus 5 (1M context) --- experimental/sgl-router/monitoring/README.md | 35 +- experimental/sgl-router/src/main.rs | 1 + .../src/policies/kv_events/index.rs | 171 ++++++++ .../sgl-router/src/policies/kv_events/mod.rs | 9 +- .../src/policies/kv_events/tally.rs | 282 ++++++++++++++ .../sgl-router/src/policies/kv_events/tree.rs | 364 ++++++++++++++++-- .../sgl-router/src/server/app_context.rs | 9 +- experimental/sgl-router/src/server/metrics.rs | 2 +- .../sgl-router/src/server/routes/metrics.rs | 253 +++++++++++- 9 files changed, 1085 insertions(+), 41 deletions(-) create mode 100644 experimental/sgl-router/src/policies/kv_events/tally.rs diff --git a/experimental/sgl-router/monitoring/README.md b/experimental/sgl-router/monitoring/README.md index b001b6fe1..073266540 100644 --- a/experimental/sgl-router/monitoring/README.md +++ b/experimental/sgl-router/monitoring/README.md @@ -16,7 +16,8 @@ on `/metrics` (text/plain, version 0.0.4) on the router's serving port ## Metrics covered -The dashboard graphs every family the router emits: +Families the router emits. The dashboard graphs all of them except the +`sgl_router_kv_*` series, whose panels ship separately: | Metric | Type | What it shows | |---|---|---| @@ -34,6 +35,13 @@ The dashboard graphs every family the router emits: | `sgl_router_stale_requests_total` | Counter | Stale-request cancellations | | `sgl_router_decode_affinity_total` | Counter | PD decode-affinity outcomes | | `sgl_router_sticky_total` | Counter | Sticky-session selection outcomes | +| `sgl_router_kv_events_total` | Counter | KV-cache events the pump consumed, by `event` and storage `medium` | +| `sgl_router_kv_event_blocks_total` | Counter | Block hashes those events carried, by `event` and `medium` | +| `sgl_router_kv_tree_blocks` | Gauge | Blocks the tree attributes to a `worker_url` / `dp_rank`, by storage `tier` | +| `sgl_router_kv_block_size` | Gauge | Tokens per block hash, as established from the fleet (0 until a worker reports) | +| `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 | The legacy `sgl_router_overlap_blocks` metric was removed with the `cache_aware_zmq` policy and has no direct replacement. Remove queries, alerts, @@ -41,7 +49,30 @@ and dashboard panels that depend on this metric before upgrading. The `sgl_router_workers` / `sgl_router_worker_*` gauges are sampled from the live worker registry on every scrape, so a removed worker stops emitting -series immediately rather than leaving a stale value. +series immediately rather than leaving a stale value. The `sgl_router_kv_*` +series are pulled from the KV-event index the same way. + +`sgl_router_kv_tree_blocks * sgl_router_kv_block_size` for one worker and +tier, divided by that pod's own occupancy of the tier (device: +`sglang_kv_used_tokens + sglang_kv_evictable_tokens`; host: +`sglang_hicache_host_used_tokens`; `tp_rank="0"`), is the tree's coverage of +that tier. Scope both sides to the same deployment before dividing — block +size and fleet membership both vary between them, and an unscoped ratio +divides one fleet's tree by another's occupancy. + +Read it as: about 1, the tree mirrors the engine; about 0, the engine holds a +tier routing cannot see; **above 1, the tree holds tiers a worker has already +released** — check `sgl_router_kv_event_batches_lost_total`, because a tagged +removal clears only its own tier and a lost batch strands the rest. + +`sgl_router_kv_events_total` renders every `(event, medium)` cell including +zeros, so a `CPU_PINNED` row pinned at 0 on a hierarchical-cache fleet is +visible rather than absent. A nonzero `block_stored/unknown` row is the +upgrade signal: the engine is publishing a storage tier this build cannot +rank, so the tree drops those stores rather than filing them under a guess. Comparing `sgl_router_kv_event_blocks_total` for +`block_stored/CPU_PINNED` against the engine's `sglang_hicache_backup_tokens_total` +needs `sum without(pool)` on the engine side, and the two are not equal +anyway: the engine also evicts device blocks it never backed up. ## Prometheus scrape config diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index 3d384168d..dcaf35fea 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -210,6 +210,7 @@ async fn main() -> Result<()> { }); app_ctx.block_size_oracle = block_size_oracle; app_ctx.engine_load = kv_index.engine_load(); + app_ctx.kv_metrics = kv_index.metrics_source(); let ctx = Arc::new(app_ctx); ctx.mark_ready(); diff --git a/experimental/sgl-router/src/policies/kv_events/index.rs b/experimental/sgl-router/src/policies/kv_events/index.rs index 83c52507d..4cc15c146 100644 --- a/experimental/sgl-router/src/policies/kv_events/index.rs +++ b/experimental/sgl-router/src/policies/kv_events/index.rs @@ -38,6 +38,7 @@ use tracing::{debug, info, warn}; use super::block_size_oracle::BlockSizeOracle; use super::discovery::{fetch_event_config, EventConfig}; 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; @@ -70,6 +71,21 @@ fn subscribable_ranks(port_base: u16, dp_size: u32) -> Vec { .collect() } +/// The read-only handles the `/metrics` scrape pulls the KV storage-tier +/// series from. Narrower than an [`KvEventIndex`] handle on purpose: a route +/// has no business calling `add_worker` / `remove_worker` / `shutdown`. +#[derive(Clone)] +pub struct KvIndexMetrics { + pub(crate) tree: Arc, + pub(crate) tally: Arc, +} + +impl KvIndexMetrics { + pub fn new(tree: Arc, tally: Arc) -> Self { + Self { tree, tally } + } +} + /// Bundle of `HashTree` + `KvEventSubscriberRegistry` + pump task. /// /// Construct one instance per router process and hand it to the worker @@ -103,6 +119,9 @@ pub struct KvEventIndex { /// may legitimately have a fresh publisher whose sequence numbers /// restart from 1. cursors: Arc>>, + /// Applied events by kind and storage medium, for the `/metrics` scrape. + /// Written only by the pump. + tally: Arc, /// Worker-sourced `page_size` shared with prefix providers. /// `add_worker` calls `try_set(cfg.block_size)` so the first worker /// establishes the value; subsequent workers that disagree are @@ -162,11 +181,13 @@ impl KvEventIndex { 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(), cursors.clone(), live_workers.clone(), + Arc::clone(&tally), pump_cancel.clone(), rx, )); @@ -182,6 +203,7 @@ impl KvEventIndex { http, live_workers, cursors, + tally, block_size_oracle, }) } @@ -198,6 +220,22 @@ impl KvEventIndex { self.tree.clone() } + /// Handles for the `/metrics` storage-tier series, or `None` when this + /// router does not maintain a local tree. + /// + /// In metadata-only mode (an external Indexer is the routing signal) no KV + /// subscription is opened, so every tier series would be a structural + /// zero — while their own HELP text tells the operator to read a zero + /// `CPU_PINNED` row as "the tier stream is not reaching the router". That + /// is a different fault with a different fix, so emit nothing rather than + /// a confidently wrong zero. + pub fn metrics_source(&self) -> Option { + self.maintain_tree.then(|| KvIndexMetrics { + tree: Arc::clone(&self.tree), + tally: Arc::clone(&self.tally), + }) + } + /// 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. @@ -422,6 +460,7 @@ async fn pump_loop( engine_load: Arc, cursors: Arc>>, live_workers: Arc>>, + tally: Arc, cancel: CancellationToken, mut rx: mpsc::Receiver, ) { @@ -480,6 +519,25 @@ async fn pump_loop( ); continue; } + // The publisher's seq is dense, so a jump is exactly the + // batches ZMQ dropped at its high-water mark. This became + // worth counting with tier-tagged removals: a removal now + // clears only its own tier, so losing the batch carrying a + // block's LAST removal leaves the worker owning it until + // the next AllBlocksCleared or teardown. The tree cannot + // see that happened — only the sequence can. The + // operator-visible signature is tree coverage above 1. + let lost = (seq - p - 1) as u64; + if lost > 0 { + tally.record_lost_batches(lost); + warn!( + worker = ?worker, + seq, + last_applied = p, + lost, + "kv-events pump: sequence gap; batches were dropped in transit and the tree may hold stale tiers for this worker", + ); + } } for event in &batch.events { // The `medium` tag decides which tier a store lands on and @@ -488,6 +546,11 @@ async fn pump_loop( // — see the tree's "Storage tiers" docs. match event { KvCacheEvent::BlockStored(b) => { + tally.record( + EventKind::BlockStored, + b.medium.as_deref(), + b.block_hashes.len(), + ); tree.insert_tiered( &worker, b.parent_block_hash, @@ -496,6 +559,11 @@ async fn pump_loop( ); } KvCacheEvent::BlockRemoved(b) => { + tally.record( + EventKind::BlockRemoved, + b.medium.as_deref(), + b.block_hashes.len(), + ); tree.remove_tiered( &worker, &b.block_hashes, @@ -503,6 +571,7 @@ async fn pump_loop( ); } KvCacheEvent::AllBlocksCleared => { + tally.record(EventKind::AllBlocksCleared, None, 0); tree.clear_worker(&worker); } } @@ -540,6 +609,7 @@ mod tests { tree: Arc, engine_load: Arc, cursors: Arc>>, + tally: Arc, #[allow(dead_code)] live_set: Arc>>, #[allow(dead_code)] @@ -557,12 +627,14 @@ mod tests { let live_set: Arc>> = Arc::new(Mutex::new(live.iter().cloned().collect())); let cancel = CancellationToken::new(); + let tally = Arc::new(EventTally::new()); let (tx, rx) = mpsc::channel(4); let pump = tokio::spawn(pump_loop( tree.clone(), engine_load.clone(), cursors.clone(), live_set.clone(), + Arc::clone(&tally), cancel.clone(), rx, )); @@ -570,6 +642,7 @@ mod tests { tree, engine_load, cursors, + tally, live_set, cancel, tx, @@ -656,6 +729,104 @@ mod tests { assert!(!m.device_workers().contains(&id), "device copy is gone"); } + /// The metadata-only gate. Its whole justification is that a structural + /// zero would be read as "the tier stream is not reaching the router" — a + /// different fault with a different fix — so the gate itself needs pinning: + /// inverting it leaves every test green while `/metrics` starts lying. + #[tokio::test] + async fn metrics_source_is_none_only_without_a_local_tree() { + let http = reqwest::Client::builder().build().unwrap(); + let with_tree = + KvEventIndex::new_with_http_and_oracle(http.clone(), BlockSizeOracle::new()); + assert!( + with_tree.metrics_source().is_some(), + "a router maintaining its own tree must publish the tier series", + ); + let metadata_only = + KvEventIndex::new_metadata_only_with_http_and_oracle(http, BlockSizeOracle::new()); + assert!( + metadata_only.metrics_source().is_none(), + "an external-Indexer router must emit nothing rather than a structural zero", + ); + } + + /// Every applied event is tallied by kind and medium, blocks included, so + /// the scrape can show the tier stream the tree is consuming. An + /// out-of-order batch is filtered before the tally and must not count. + #[tokio::test] + async fn pump_tallies_applied_events_by_medium() { + let id = worker_id("http://w1", 0); + let h = spawn_pump(std::slice::from_ref(&id)); + let (tally, tx, pump) = (h.tally, h.tx, h.pump); + + let stored = |medium: Option<&str>, hashes: Vec| { + KvCacheEvent::BlockStored(BlockStored { + parent_block_hash: None, + block_hashes: hashes, + token_ids: vec![], + block_size: 64, + lora_id: None, + medium: medium.map(str::to_owned), + }) + }; + tx.send(WorkerEvent::Batch { + worker: id.clone(), + seq: 2, + batch: batch(vec![ + stored(Some("GPU"), vec![10, 20, 30]), + stored(Some("CPU_PINNED"), vec![10, 20, 30]), + KvCacheEvent::BlockRemoved(BlockRemoved { + block_hashes: vec![30], + medium: Some("GPU".into()), + }), + ]), + }) + .await + .unwrap(); + // Out of order: filtered, must not be tallied. + tx.send(WorkerEvent::Batch { + worker: id.clone(), + seq: 1, + batch: batch(vec![stored(None, vec![99])]), + }) + .await + .unwrap(); + // A gap: seq 3 and 4 were dropped in transit. Counted, because a + // tagged removal now clears only its own tier, so a lost batch can + // strand a tier the tree will never clear on its own. + tx.send(WorkerEvent::Batch { + worker: id.clone(), + seq: 5, + batch: batch(vec![KvCacheEvent::AllBlocksCleared]), + }) + .await + .unwrap(); + drop(tx); + pump.await.unwrap(); + + assert_eq!(tally.batches_lost(), 2, "seq 3 and 4 never arrived"); + let rows = tally.snapshot(); + let cell = |event: &str, medium: &str| { + rows.iter() + .find(|r| r.event == event && r.medium == medium) + .cloned() + .expect("cell rendered") + }; + assert_eq!(cell("block_stored", "GPU").blocks, 3); + assert_eq!(cell("block_stored", "CPU_PINNED").blocks, 3); + assert_eq!(cell("block_removed", "GPU").events, 1); + assert_eq!( + cell("block_stored", "untagged").events, + 0, + "the out-of-order batch was filtered before the tally", + ); + assert_eq!( + cell("all_blocks_cleared", "untagged").events, + 1, + "a clear carries no medium and lands on the untagged row", + ); + } + /// A `WorkerEvent::Load` lands in the engine-load table (gauge, no /// cursor) keyed by the worker URL, and does not touch the tree. #[tokio::test] diff --git a/experimental/sgl-router/src/policies/kv_events/mod.rs b/experimental/sgl-router/src/policies/kv_events/mod.rs index b6a0e1786..266766fc4 100644 --- a/experimental/sgl-router/src/policies/kv_events/mod.rs +++ b/experimental/sgl-router/src/policies/kv_events/mod.rs @@ -12,6 +12,7 @@ //! - [`hash`] — block-hash compute mirroring SGLang `RadixKey.hash_page`. //! - [`tree`] — hash-keyed radix tree consumed by the routing path, //! tracking the storage tier each worker holds a block on. +//! - [`tally`] — per-(kind, medium) counters of the events the pump applied. //! - [`subscriber`] — per-worker ZMQ SUB tasks. //! - [`discovery`] — `/server_info` parse → publisher endpoint. //! - [`index`] — public façade bundling the tree + subscribers + pump. @@ -21,6 +22,7 @@ pub mod discovery; pub mod hash; pub mod index; pub mod subscriber; +pub mod tally; pub mod tree; pub mod wire; @@ -28,9 +30,12 @@ pub use block_size_oracle::BlockSizeOracle; pub(crate) use discovery::classify_bigram; pub use discovery::{fetch_event_config, EventConfig}; pub use hash::{compute_block_hashes, compute_block_hashes_bigram, sha256_to_i64}; -pub use index::KvEventIndex; +pub use index::{KvEventIndex, KvIndexMetrics}; pub use subscriber::{KvEventSubscriberRegistry, SubKind, WorkerEvent}; -pub use tree::{HashTree, KvWorkerId, MatchResult, Tiers}; +pub use tally::{EventKind, EventTally, TallyRow}; +pub use tree::{ + HashTree, KvWorkerId, MatchResult, TierCounts, Tiers, ACCOUNTING_REASONS, TIER_SLOT_COUNT, +}; pub use wire::{ decode_event_batch, BlockRemoved, BlockStored, DecodeError, KvCacheEvent, KvEventBatch, }; diff --git a/experimental/sgl-router/src/policies/kv_events/tally.rs b/experimental/sgl-router/src/policies/kv_events/tally.rs new file mode 100644 index 000000000..cd4a98202 --- /dev/null +++ b/experimental/sgl-router/src/policies/kv_events/tally.rs @@ -0,0 +1,282 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Counters for the KV-cache event stream the pump consumes: events by kind +//! and by the storage `medium` tag each carried, plus the batches the +//! transport lost. Rendered as `sgl_router_kv_events_total`, +//! `sgl_router_kv_event_blocks_total` and `sgl_router_kv_event_batches_lost_total`. +//! +//! WHY this exists: an engine running a hierarchical cache publishes a +//! host-tier store for every block it backs up and a device-tier removal when +//! the device copy goes. Whether those tagged events reach the router, and at +//! what volume, was not observable anywhere — the tree consumed them and +//! nothing counted them — so a router discarding the tag looked identical to +//! an engine never sending it. Counting by medium makes the tier stream a +//! time series: `block_stored/CPU_PINNED` tracks the engine's D2H backup +//! volume and `block_removed/GPU` its device eviction volume. The two are NOT +//! equal — the engine also evicts device blocks that were never backed up, +//! and counts those itself as `sglang_hicache_dropped_tokens_total` — but a +//! `CPU_PINNED` row pinned at zero with hicache enabled points at the +//! publisher or the subscription, not the tree. +//! +//! Every cell is rendered, zeros included: the zero IS the finding. +//! +//! Label cardinality is fixed: the medium label is folded to the values the +//! tree can rank plus `untagged` (no `medium` field) and `unknown` (a string +//! this build does not recognise), so a misbehaving publisher cannot mint +//! series. +//! +//! WHY lost batches are counted here rather than left to the tree: ZMQ drops +//! at the publisher's high-water mark, and since a tagged removal now clears +//! only its own tier, losing the batch that carried a block's LAST removal +//! leaves the worker owning that block until the next `AllBlocksCleared` or +//! worker teardown. The tree cannot see that it happened; only the sequence +//! numbers can. + +use std::collections::HashSet; +use std::sync::atomic::{AtomicU64, Ordering}; + +use parking_lot::Mutex; +use tracing::warn; + +use super::tree::Tiers; + +/// Medium labels, in the order [`EventTally`] stores them: the wire strings +/// the tree ranks, read off [`Tiers::WIRE_MEDIA`] so the two can never +/// disagree, then the two folds. +pub const MEDIUM_LABELS: [&str; 6] = [ + Tiers::WIRE_MEDIA[0].0, + Tiers::WIRE_MEDIA[1].0, + Tiers::WIRE_MEDIA[2].0, + Tiers::WIRE_MEDIA[3].0, + "untagged", + "unknown", +]; +const UNTAGGED: usize = 4; +const UNKNOWN: usize = 5; +const _: () = assert!( + Tiers::WIRE_MEDIA.len() == UNTAGGED, + "MEDIUM_LABELS lists every WIRE_MEDIA entry before the folds", +); + +/// Cap on the distinct unrecognised `medium` strings remembered for +/// warn-once. A publisher cannot grow router memory by inventing media; past +/// the cap the warning simply repeats. +const MAX_REMEMBERED_UNKNOWN_MEDIA: usize = 16; + +/// Which event a tally entry describes. +/// +/// [`Self::slot`] and [`Self::label`] are exhaustive matches rather than a +/// discriminant cast into a parallel `&[&str]` table: a cast plus a +/// length assertion still lets an APPENDED variant compile and then panic on +/// an out-of-range index inside the pump task, which would take the whole +/// cache-aware path down with no restart. A new variant here is two compile +/// errors instead. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EventKind { + BlockStored, + BlockRemoved, + AllBlocksCleared, +} + +impl EventKind { + /// Every kind, in storage order. + pub const ALL: [EventKind; 3] = [ + Self::BlockStored, + Self::BlockRemoved, + Self::AllBlocksCleared, + ]; + + const fn slot(self) -> usize { + match self { + Self::BlockStored => 0, + Self::BlockRemoved => 1, + Self::AllBlocksCleared => 2, + } + } + + const fn label(self) -> &'static str { + match self { + Self::BlockStored => "block_stored", + Self::BlockRemoved => "block_removed", + Self::AllBlocksCleared => "all_blocks_cleared", + } + } +} + +// `slot()` indexes the counter arrays, so it must be a permutation of +// `0..ALL.len()`. Pinned here rather than trusted. +const _: () = { + let mut i = 0; + while i < EventKind::ALL.len() { + assert!( + EventKind::ALL[i].slot() == i, + "EventKind::slot must match position in EventKind::ALL", + ); + i += 1; + } +}; + +/// One rendered cell of the tally. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TallyRow { + pub event: &'static str, + pub medium: &'static str, + /// Events applied. + pub events: u64, + /// Block hashes those events carried (0 for `all_blocks_cleared`, whose + /// wire type carries no hashes). + pub blocks: u64, +} + +/// Lock-free counters, written by the single pump task and read on scrape. +#[derive(Debug, Default)] +pub struct EventTally { + events: [[AtomicU64; MEDIUM_LABELS.len()]; EventKind::ALL.len()], + blocks: [[AtomicU64; MEDIUM_LABELS.len()]; EventKind::ALL.len()], + /// Batches the transport lost, inferred from gaps in the publisher's + /// dense sequence. + batches_lost: AtomicU64, + /// Unrecognised `medium` strings already warned about, so an engine that + /// adds a tier logs once per string rather than once per event. + warned_unknown_media: Mutex>, +} + +impl EventTally { + pub fn new() -> Self { + Self::default() + } + + fn medium_slot(&self, medium: Option<&str>) -> usize { + let Some(m) = medium else { + return UNTAGGED; + }; + match MEDIUM_LABELS[..UNTAGGED] + .iter() + .position(|known| *known == m) + { + Some(slot) => slot, + None => { + // Only an unrecognised medium takes the lock, so the common + // path stays allocation- and lock-free. + self.warn_unknown_medium(m); + UNKNOWN + } + } + } + + /// Log the first sighting of each unrecognised `medium`. This is the + /// "the engine added a storage tier, upgrade the router" line: the tree + /// drops a store tagged with one ([`Tiers::for_store`]) and would + /// otherwise do it in silence, and silence about a discarded tag is the + /// bug this module exists to prevent recurring. + fn warn_unknown_medium(&self, medium: &str) { + let mut seen = self.warned_unknown_media.lock(); + if seen.contains(medium) { + return; + } + if seen.len() < MAX_REMEMBERED_UNKNOWN_MEDIA { + seen.insert(medium.to_owned()); + } + drop(seen); + warn!( + medium, + known = ?MEDIUM_LABELS[..UNTAGGED], + "kv-events: unrecognised storage medium; blocks stored on it are routable but rank below every known tier, and a removal tagged with it clears every tier", + ); + } + + /// Book one applied event carrying `blocks` block hashes. + pub fn record(&self, event: EventKind, medium: Option<&str>, blocks: usize) { + let (e, m) = (event.slot(), self.medium_slot(medium)); + self.events[e][m].fetch_add(1, Ordering::Relaxed); + self.blocks[e][m].fetch_add(blocks as u64, Ordering::Relaxed); + } + + /// Book `count` batches the transport dropped between two applied + /// sequence numbers. + pub fn record_lost_batches(&self, count: u64) { + self.batches_lost.fetch_add(count, Ordering::Relaxed); + } + + pub fn batches_lost(&self) -> u64 { + self.batches_lost.load(Ordering::Relaxed) + } + + /// Every cell in (event, medium) order, zeros included. + pub fn snapshot(&self) -> Vec { + let mut rows = Vec::with_capacity(EventKind::ALL.len() * MEDIUM_LABELS.len()); + for event in EventKind::ALL { + for (m, medium) in MEDIUM_LABELS.iter().enumerate() { + rows.push(TallyRow { + event: event.label(), + medium, + events: self.events[event.slot()][m].load(Ordering::Relaxed), + blocks: self.blocks[event.slot()][m].load(Ordering::Relaxed), + }); + } + } + rows + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cell<'a>(rows: &'a [TallyRow], event: &str, medium: &str) -> &'a TallyRow { + rows.iter() + .find(|r| r.event == event && r.medium == medium) + .expect("every (event, medium) cell is rendered") + } + + #[test] + fn records_by_kind_and_medium_and_folds_the_rest() { + let t = EventTally::new(); + t.record(EventKind::BlockStored, Some("GPU"), 3); + t.record(EventKind::BlockStored, Some("CPU_PINNED"), 3); + t.record(EventKind::BlockRemoved, Some("GPU"), 1); + t.record(EventKind::BlockRemoved, None, 2); + t.record(EventKind::BlockStored, Some("NVLINK_PEER"), 5); + t.record(EventKind::AllBlocksCleared, None, 0); + + let rows = t.snapshot(); + assert_eq!(rows.len(), EventKind::ALL.len() * MEDIUM_LABELS.len()); + assert_eq!(cell(&rows, "block_stored", "GPU").blocks, 3); + assert_eq!(cell(&rows, "block_stored", "CPU_PINNED").events, 1); + assert_eq!(cell(&rows, "block_removed", "GPU").blocks, 1); + assert_eq!(cell(&rows, "block_removed", "untagged").blocks, 2); + assert_eq!(cell(&rows, "block_stored", "unknown").blocks, 5); + assert_eq!(cell(&rows, "all_blocks_cleared", "untagged").events, 1); + // Zero cells are present, not omitted. + assert_eq!(cell(&rows, "block_removed", "CPU_PINNED").events, 0); + } + + /// A publisher inventing media must not grow router memory without bound; + /// past the cap the warning repeats instead. + #[test] + fn remembered_unknown_media_are_bounded() { + let t = EventTally::new(); + for i in 0..MAX_REMEMBERED_UNKNOWN_MEDIA * 3 { + t.record(EventKind::BlockStored, Some(&format!("MEDIUM_{i}")), 1); + } + assert_eq!( + t.warned_unknown_media.lock().len(), + MAX_REMEMBERED_UNKNOWN_MEDIA, + ); + assert_eq!( + cell(&t.snapshot(), "block_stored", "unknown").events, + (MAX_REMEMBERED_UNKNOWN_MEDIA * 3) as u64, + "every unknown medium is still counted, only the warning is capped", + ); + } + + #[test] + fn lost_batches_accumulate() { + let t = EventTally::new(); + assert_eq!(t.batches_lost(), 0); + t.record_lost_batches(3); + t.record_lost_batches(1); + assert_eq!(t.batches_lost(), 4); + } +} diff --git a/experimental/sgl-router/src/policies/kv_events/tree.rs b/experimental/sgl-router/src/policies/kv_events/tree.rs index d5acb4f10..79f8da63e 100644 --- a/experimental/sgl-router/src/policies/kv_events/tree.rs +++ b/experimental/sgl-router/src/policies/kv_events/tree.rs @@ -161,9 +161,22 @@ impl Tiers { .union(Self::HOST) .union(Self::DISK) .union(Self::EXTERNAL); + /// Every tier with its metric label, in bit order. [`TierCounts`] is + /// indexed by position in this table, and the order is preference order: + /// a device copy is served in place, a host copy by load-back from host + /// memory, local disk slower still, a remote pool slower again. + pub const SLOTS: [(Tiers, &'static str); TIER_SLOT_COUNT] = [ + (Self::DEVICE, "device"), + (Self::HOST, "host"), + (Self::DISK, "disk"), + (Self::EXTERNAL, "external"), + ]; + /// The `StorageMedium` strings SGLang puts on the wire /// (`python/sglang/srt/disaggregation/kv_events.py`) and the tier each - /// lands on. The single source for the tree's ranking. + /// lands on. The single source for both the tree's ranking and the event + /// tally's medium labels, so a medium the tree ranks can never be one the + /// tally reports as unknown. pub const WIRE_MEDIA: [(&'static str, Tiers); 4] = [ ("GPU", Self::DEVICE), ("CPU_PINNED", Self::HOST), @@ -229,6 +242,16 @@ impl Tiers { self.0 &= !other.0; } + /// The bits set in both. + pub const fn intersection(self, other: Tiers) -> Tiers { + Tiers(self.0 & other.0) + } + + /// The bits set here and not in `other`. + pub const fn difference(self, other: Tiers) -> Tiers { + Tiers(self.0 & !other.0) + } + /// The bits set in either. `const` so the tier tables can be built from /// the individual tiers rather than from raw bit arithmetic. pub const fn union(self, other: Tiers) -> Tiers { @@ -236,31 +259,73 @@ impl Tiers { } } -// `ALL` has to stay in step with the individual tiers, and nothing about -// adding a `pub const` tier would otherwise force it: a bit missing from -// `ALL` is a bit an untagged `BlockRemoved` never clears — a permanently -// stale owner, the one failure `for_remove` exists to prevent. +/// Number of entries in [`Tiers::SLOTS`]. +pub const TIER_SLOT_COUNT: usize = 4; + +// The three tier tables have to stay in step, and nothing about adding a +// `pub const` tier would otherwise force it: +// +// * a bit missing from `ALL` is a bit an untagged `BlockRemoved` never +// clears — a permanently stale owner, the one failure `for_remove` exists +// to prevent; +// * a bit missing from `SLOTS` is never counted by `tally_tiers` nor +// decremented by `account_remove`, so the tier is invisible in +// `sgl_router_kv_tree_blocks` — this bug class, one tier later. The +// `debug_recount_occupancy` oracle cannot catch it either, because it +// walks the same `SLOTS`. const _: () = { + let mut union = Tiers(0); + let mut i = 0; + while i < TIER_SLOT_COUNT { + union = union.union(Tiers::SLOTS[i].0); + i += 1; + } + assert!( + union.0 == Tiers::ALL.0, + "Tiers::ALL must be exactly the union of Tiers::SLOTS", + ); let mut i = 0; while i < Tiers::WIRE_MEDIA.len() { assert!( Tiers::ALL.contains(Tiers::WIRE_MEDIA[i].1), - "every wire medium must map to a tier `ALL` clears", + "every wire medium must map to a tier that SLOTS ranks", ); i += 1; } }; +/// How many nodes one carrier holds on each tier, indexed like +/// [`Tiers::SLOTS`]. A node held on device and host counts under both. +pub type TierCounts = [u64; TIER_SLOT_COUNT]; + +/// Count one node's worth of `bits` into `counts`. +fn tally_tiers(counts: &mut TierCounts, bits: Tiers) { + for (slot, (tier, _)) in Tiers::SLOTS.iter().enumerate() { + if bits.contains(*tier) { + counts[slot] += 1; + } + } +} + /// Add `tiers` to `worker`'s hold in `carriers`, creating the entry on first -/// sight. One lookup on the re-store path — under a hierarchical cache the -/// host backup of a chain the worker already holds on device, the common -/// case — and two on first sight. Never leaves an entry with no bits, which -/// [`TreeState::remove`] relies on. -fn add_tiers(carriers: &mut HashMap, worker: &KvWorkerId, tiers: Tiers) { +/// sight, and return the bits that were newly set. One lookup on the re-store +/// path — under a hierarchical cache the host backup of a chain the worker +/// already holds on device, the common case — and two on first sight. Never +/// leaves an entry with no bits, which [`TreeState::remove`] relies on. +fn add_tiers( + carriers: &mut HashMap, + worker: &KvWorkerId, + tiers: Tiers, +) -> Tiers { match carriers.get_mut(worker) { - Some(held) => held.insert(tiers), + Some(held) => { + let added = tiers.difference(*held); + held.insert(tiers); + added + } None => { carriers.insert(worker.clone(), tiers); + tiers } } } @@ -371,8 +436,25 @@ struct TreeState { nodes: HashMap, by_hash: HashMap>, next_id: NodeId, + /// Nodes each carrier holds, per tier. Booked at every site where a + /// carrier's tier bits on a node change (`account_add` / + /// `account_remove`), so a scrape reads it without walking the tree. A + /// carrier holding nothing has no row. + occupancy: HashMap, + /// Times the occupancy bookkeeping contradicted itself, by reason. Both + /// reasons mean the same class of bug — bits added without being booked — + /// but the symptom the operator sees is a worker's + /// `sgl_router_kv_tree_blocks` rows vanishing while it still holds nodes, + /// which the metric's own HELP text says means "this worker publishes + /// nothing". A counter is the only thing separating those. + accounting_errors: [u64; ACCOUNTING_REASONS.len()], } +/// Reasons in [`TreeState::accounting_errors`] order. +pub const ACCOUNTING_REASONS: [&str; 2] = ["missing_row", "underflow"]; +const REASON_MISSING_ROW: usize = 0; +const REASON_UNDERFLOW: usize = 1; + const ROOT_ID: NodeId = 0; /// Sentinel block_hash for the root. Real workers can in principle emit /// `i64::MIN`, but the root is never looked up via `by_hash` so collisions @@ -397,6 +479,8 @@ impl TreeState { nodes, by_hash: HashMap::new(), next_id: 1, + occupancy: HashMap::new(), + accounting_errors: [0; ACCOUNTING_REASONS.len()], } } @@ -406,6 +490,74 @@ impl TreeState { id } + /// Book `delta` nodes-per-tier newly held by `worker` — one row lookup + /// for a whole chain, which is how `insert` uses it. + fn account_add(&mut self, worker: &KvWorkerId, delta: TierCounts) { + if delta.iter().all(|&c| c == 0) { + return; + } + match self.occupancy.get_mut(worker) { + Some(counts) => { + for (acc, d) in counts.iter_mut().zip(delta) { + *acc += d; + } + } + None => { + self.occupancy.insert(worker.clone(), delta); + } + } + } + + /// Book one node's worth of `removed` tier bits `worker` no longer holds. + fn account_remove(&mut self, worker: &KvWorkerId, removed: Tiers) { + let mut delta = TierCounts::default(); + tally_tiers(&mut delta, removed); + self.account_remove_counts(worker, delta); + } + + /// Book `delta` nodes-per-tier `worker` no longer holds, dropping the + /// carrier's row once it holds nothing. + /// + /// Takes a whole-carrier delta rather than one node's bits so + /// `clear_worker` can book a fleet-sized tree in ONE call: booking per + /// node would emit one `error!` per node on the failure path below, which + /// is hundreds of thousands of lines written while holding the tree's + /// write lock — every routing decision blocked behind a log flood. + fn account_remove_counts(&mut self, worker: &KvWorkerId, delta: TierCounts) { + if delta.iter().all(|&c| c == 0) { + return; + } + let Some(counts) = self.occupancy.get_mut(worker) else { + self.accounting_errors[REASON_MISSING_ROW] += 1; + error!( + worker = %worker.url, + dp_rank = worker.dp_rank, + "tree invariant violation: releasing tiers for a carrier with no occupancy row", + ); + return; + }; + let mut underflowed = false; + for (acc, d) in counts.iter_mut().zip(delta) { + underflowed |= *acc < d; + *acc = acc.saturating_sub(d); + } + if underflowed { + // Saturating so a release can never wrap the gauge — but then the + // all-zero test below would drop a row for a carrier that still + // holds nodes, and a missing series reads as "this worker + // publishes nothing". Count it so a release build is not blind; + // assert so a debug run stops at the release that exposed it. + self.accounting_errors[REASON_UNDERFLOW] += 1; + debug_assert!( + false, + "occupancy underflow: a tier was released that was never booked", + ); + } + if counts.iter().all(|&c| c == 0) { + self.occupancy.remove(worker); + } + } + /// Insert a brand-new child under `parent_id` and wire up the reverse /// index. Caller is responsible for ensuring `parent_id`'s child slot /// for `block_hash` is empty (else this overwrites it). @@ -501,6 +653,8 @@ impl TreeState { let mut current = self.resolve_parent(worker, parent_hash); let mut prev_hash = parent_hash; let now = now_millis(); + // Occupancy is booked once for the whole chain, not per block. + let mut delta = TierCounts::default(); for &h in block_hashes { let child_id = match self .nodes @@ -508,9 +662,13 @@ impl TreeState { .and_then(|n| n.children.get(&h).copied()) { Some(id) => id, + // `break`, not `return`: the tier bits are already written + // into the nodes visited so far, so bailing without reaching + // `account_add` below would leave the occupancy gauge + // permanently short by exactly those blocks. None => match self.create_child(current, h, prev_hash) { Some(id) => id, - None => return, + None => break, }, }; let Some(child) = self.nodes.get_mut(&child_id) else { @@ -519,13 +677,15 @@ impl TreeState { block_hash = h, "tree invariant violation: child node missing immediately after fetch/create; aborting chain", ); - return; + break; }; - add_tiers(&mut child.workers, worker, tiers); + let added = add_tiers(&mut child.workers, worker, tiers); child.last_used.store(now, Ordering::Relaxed); + tally_tiers(&mut delta, added); current = child_id; prev_hash = Some(h); } + self.account_add(worker, delta); } /// Clear `tiers` from `worker`'s hold on every node carrying any hash in @@ -543,18 +703,24 @@ impl TreeState { for id in targets { // Node may already be gone if a previous prune in this batch // cascaded through it — skip silently. - let prunable = match self.nodes.get_mut(&id) { + let (prunable, released) = match self.nodes.get_mut(&id) { Some(node) => { + let mut released = Tiers::default(); if let Some(held) = node.workers.get_mut(worker) { + released = held.intersection(tiers); held.remove(tiers); if held.is_empty() { node.workers.remove(worker); } } - node.workers.is_empty() && node.children.is_empty() + ( + node.workers.is_empty() && node.children.is_empty(), + released, + ) } - None => false, + None => (false, Tiers::default()), }; + self.account_remove(worker, released); if prunable { self.prune_cascade(id); } @@ -570,16 +736,23 @@ impl TreeState { .filter(|&id| id != ROOT_ID) .collect(); let mut prune_candidates: Vec = Vec::new(); + // Accumulated over every node, then booked once — see + // `account_remove_counts`. + let mut delta = TierCounts::default(); for id in ids { - if let Some(node) = self.nodes.get_mut(&id) { - if node.workers.remove(worker).is_some() - && node.workers.is_empty() - && node.children.is_empty() - { + let removed = self.nodes.get_mut(&id).and_then(|node| { + node.workers + .remove(worker) + .map(|held| (held, node.workers.is_empty() && node.children.is_empty())) + }); + if let Some((held, prunable)) = removed { + tally_tiers(&mut delta, held); + if prunable { prune_candidates.push(id); } } } + self.account_remove_counts(worker, delta); for id in prune_candidates { // Re-check: cascading prune from a sibling may have already // removed this id. @@ -844,9 +1017,14 @@ impl TreeState { }; // Force-prune even if the leaf still holds workers — eviction // intentionally evicts. We clear workers first so the cascade - // precondition holds. - if let Some(node) = self.nodes.get_mut(&victim) { - node.workers.clear(); + // precondition holds, releasing each carrier's tiers as we go so + // the occupancy counters stay in step. + let carriers: Vec<(KvWorkerId, Tiers)> = match self.nodes.get_mut(&victim) { + Some(node) => node.workers.drain().collect(), + None => Vec::new(), + }; + for (worker, held) in &carriers { + self.account_remove(worker, *held); } self.prune_cascade(victim); } @@ -980,6 +1158,33 @@ impl HashTree { self.state.read().by_hash.len() } + /// Times the occupancy bookkeeping contradicted itself, by + /// [`ACCOUNTING_REASONS`]. Always zero on a correct tree; rendered so a + /// release build, where the debug assertion is compiled out, still says + /// so out loud. + pub fn accounting_errors(&self) -> [u64; ACCOUNTING_REASONS.len()] { + self.state.read().accounting_errors + } + + /// Nodes each carrier holds on each tier, sorted by carrier. Rendered as + /// `sgl_router_kv_tree_blocks`; against the engine's own per-tier + /// occupancy for the same pod it is the tree's coverage of that tier — + /// the number that says whether a tier the engine holds is visible to + /// routing at all (module docs, "Storage tiers"). + /// + /// Read off the incrementally maintained accounting rather than walked, + /// so a scrape costs one read lock. A carrier holding nothing has no row. + pub fn tier_occupancy(&self) -> Vec<(KvWorkerId, TierCounts)> { + let state = self.state.read(); + let mut rows: Vec<(KvWorkerId, TierCounts)> = state + .occupancy + .iter() + .map(|(w, counts)| (w.clone(), *counts)) + .collect(); + rows.sort_by(|a, b| (&a.0.url, a.0.dp_rank).cmp(&(&b.0.url, b.0.dp_rank))); + rows + } + /// Evict least-recently-used nodes until `node_count() <= max_size`. /// /// Strategy: @@ -1012,6 +1217,32 @@ impl HashTree { .values() .all(|n| n.workers.values().all(|t| !t.is_empty())) } + + /// Recompute [`Self::tier_occupancy`] by walking every node — the oracle + /// the incrementally maintained counters are checked against. + /// + /// Its ceiling: it shares `tally_tiers` and [`Tiers::SLOTS`] with the code + /// it checks, so it proves only that incremental booking agrees with a + /// full walk. It is blind by construction to WHICH tier is right — a + /// SLOTS/bit mismatch would mis-tally identically on both sides. The + /// example-based tier tests and the compile-time coupling asserts own + /// that half. + fn debug_recount_occupancy(&self) -> Vec<(KvWorkerId, TierCounts)> { + let state = self.state.read(); + let mut total: HashMap = HashMap::new(); + for (&id, node) in &state.nodes { + if id == ROOT_ID { + continue; + } + for (worker, held) in &node.workers { + let acc = total.entry(worker.clone()).or_default(); + tally_tiers(acc, *held); + } + } + let mut rows: Vec<(KvWorkerId, TierCounts)> = total.into_iter().collect(); + rows.sort_by(|a, b| (&a.0.url, a.0.dp_rank).cmp(&(&b.0.url, b.0.dp_rank))); + rows + } } #[cfg(test)] @@ -1228,6 +1459,59 @@ mod tests { assert_eq!(tree.node_count(), 0); } + /// The per-tier occupancy is maintained incrementally at every mutation + /// site, so it is checked against a full recount after a sequence that + /// exercises all of them: tiered stores, partial and full removals, + /// `clear_worker`, and LRU eviction. + #[test] + fn tier_occupancy_matches_a_full_recount_after_mixed_mutations() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + let b = worker("http://b", 1); + let c = worker("http://c", 0); + + tree.insert_tiered(&a, None, &[1, 2, 3, 4], Tiers::DEVICE); + tree.insert_tiered(&a, None, &[1, 2], Tiers::HOST); + tree.insert_tiered(&b, None, &[1, 2, 3], Tiers::HOST); + tree.insert_tiered(&b, None, &[1, 2, 5, 6], Tiers::DEVICE); + tree.insert_tiered(&c, None, &[7], Tiers::for_store(Some("EXTERNAL"))); + tree.insert_tiered(&c, None, &[8], Tiers::for_store(Some("NVLINK_PEER"))); + for r in 0..40i64 { + tree.insert(&c, None, &[r * 4096 + 11, r * 4096 + 12]); + } + assert_eq!(tree.tier_occupancy(), tree.debug_recount_occupancy()); + + // Spot-check the shape: a holds 4 device nodes and 2 host nodes. + let rows = tree.tier_occupancy(); + let (_, a_counts) = rows.iter().find(|(w, _)| *w == a).unwrap(); + assert_eq!(a_counts[0], 4, "device"); + assert_eq!(a_counts[1], 2, "host"); + assert_eq!(a_counts[2], 0, "disk"); + let (_, c_counts) = rows.iter().find(|(w, _)| *w == c).unwrap(); + assert_eq!(c_counts[3], 1, "external"); + assert_eq!( + tree.match_prefix(None, &[8]).matched_blocks, + 0, + "a store on an unrankable medium is booked nowhere because it is never applied", + ); + + // Partial removal (device only) on a node held on both tiers, then a + // removal that clears the last tier and prunes. + tree.remove_tiered(&a, &[2], Tiers::DEVICE); + tree.remove_tiered(&a, &[4], Tiers::ALL); + assert_eq!(tree.tier_occupancy(), tree.debug_recount_occupancy()); + + // Whole-worker clear, then LRU eviction down to a small cap. + tree.clear_worker(&b); + assert_eq!(tree.tier_occupancy(), tree.debug_recount_occupancy()); + assert!(tree.evict_lru(10) > 0); + assert_eq!(tree.tier_occupancy(), tree.debug_recount_occupancy()); + assert!( + tree.tier_occupancy().iter().all(|(w, _)| *w != b), + "a cleared worker must leave no occupancy row", + ); + } + /// `DISK` (L3) and `EXTERNAL` (L4) are distinct tiers in the engine's own /// `StorageMedium`, so they must not share a bit: on a fleet running both, /// folding them would make an L3 eviction erase the router's knowledge of @@ -1250,13 +1534,16 @@ mod tests { assert_eq!(tree.match_prefix(None, &[1]).matched_blocks, 0); } - /// `remove` and `prune_cascade` rely on "no bits ⇒ no entry" to know - /// when a node is prunable, and every mutation site has to preserve it — - /// exactly the shape that rots under a later refactor. A deterministic - /// random walk over every operation and every medium, with the invariant - /// asserted after EVERY step. + /// The occupancy counters are maintained incrementally at four mutation + /// sites, which is exactly the shape that rots under a later refactor. A + /// deterministic random walk over every operation and every medium, with + /// both invariants asserted after EVERY step: + /// + /// * `tier_occupancy()` equals a full node-walk recount, and + /// * no carrier entry with empty tiers exists (`remove` relies on + /// "no bits ⇒ no entry" to know when a node is prunable). #[test] - fn carrier_invariant_holds_under_a_random_walk() { + fn occupancy_and_carrier_invariants_hold_under_a_random_walk() { // xorshift64*, so the walk is reproducible without a dev-dependency. struct Rng(u64); impl Rng { @@ -1307,6 +1594,11 @@ mod tests { _ => tree.insert_tiered(w, parent, &hashes, Tiers::for_store(medium)), } + assert_eq!( + tree.tier_occupancy(), + tree.debug_recount_occupancy(), + "seed {seed} step {step}: incremental occupancy diverged from a full recount", + ); assert!( tree.debug_no_empty_carrier(), "seed {seed} step {step}: a carrier is present holding no tier", @@ -1318,7 +1610,8 @@ mod tests { /// `AllBlocksCleared` is the pod-restart / scale-down path /// (`remove_worker` clears every rank), so it must drop a carrier /// regardless of which tiers it held — a hold on any lower tier is still - /// a hold. + /// a hold. Asserted at the tree, not only through the rendered metric, so + /// a metrics refactor cannot take the coverage with it. #[test] fn clear_worker_drops_carriers_on_every_tier() { let tree = HashTree::new(); @@ -1335,6 +1628,11 @@ mod tests { "a host-only carrier must be cleared like any other", ); assert_eq!(tree.match_prefix(None, &[3]).matched_blocks, 0); + assert!( + tree.tier_occupancy().iter().all(|(w, _)| *w != a), + "a cleared worker must leave no occupancy row on any tier", + ); + assert_eq!(tree.tier_occupancy(), tree.debug_recount_occupancy()); } /// `resolve_parent` disambiguates a shared hash by preferring a candidate diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 909039545..8c6e84e2b 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -6,7 +6,7 @@ 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; +use crate::policies::kv_events::{BlockSizeOracle, KvIndexMetrics}; use crate::policies::prefix_provider::RadixTreePrefixProvider; use crate::policies::PolicyRegistry; use crate::proxy::Proxy; @@ -36,6 +36,11 @@ pub struct AppContext { 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`]. + pub kv_metrics: Option, ready: AtomicBool, } @@ -91,6 +96,7 @@ impl AppContext { prefix_index: None, radix_tree_prefix_provider: None, block_size_oracle: BlockSizeOracle::new(), + kv_metrics: None, engine_load: EngineLoadTable::new(), ready: AtomicBool::new(false), } @@ -146,6 +152,7 @@ impl AppContext { prefix_index: None, radix_tree_prefix_provider: None, block_size_oracle: BlockSizeOracle::new(), + kv_metrics: None, engine_load: EngineLoadTable::new(), ready: AtomicBool::new(false), } diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index 0c2c86e08..2a8b4777e 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -1074,7 +1074,7 @@ fn render_histogram(out: &mut String, name: &str, label_body: &str, hist: &Histo /// https://prometheus.io/docs/instrumenting/exposition_formats/. /// We only escape `\`, `"`, and newline — the three characters the /// reference parser rejects unescaped. -fn escape_label(s: &str) -> String { +pub(crate) fn escape_label(s: &str) -> String { let mut out = String::with_capacity(s.len()); for c in s.chars() { match c { diff --git a/experimental/sgl-router/src/server/routes/metrics.rs b/experimental/sgl-router/src/server/routes/metrics.rs index 11583a0cc..23e09cfcc 100644 --- a/experimental/sgl-router/src/server/routes/metrics.rs +++ b/experimental/sgl-router/src/server/routes/metrics.rs @@ -10,8 +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::WorkerSnapshot; +use crate::server::metrics::{escape_label, WorkerSnapshot}; use axum::extract::State; use axum::http::header::CONTENT_TYPE; use axum::http::StatusCode; @@ -49,7 +50,27 @@ pub async fn metrics(State(ctx): State>) -> impl IntoResponse { } }) .collect(); - let body = ctx.metrics.render_with_workers(&workers); + let mut body = ctx.metrics.render_with_workers(&workers); + // Pull-on-scrape, like the worker gauges above: the tree and the tally + // own the numbers, so a worker that goes away stops emitting series + // without anything having to reset a pushed counter. + // Emitted unconditionally: without it, "no kv series" is indistinguishable + // between the intended metadata-only mode, a broken `kv_metrics` wiring, + // and a regressed endpoint. + body.push_str( + "# HELP sgl_router_kv_tree_maintained 1 when this router maintains its own cache-aware KV tree and therefore emits the sgl_router_kv_* series; 0 when placement comes from an external Indexer, where those series would be a structural zero and are omitted rather than reported as an empty tier stream.\n", + ); + body.push_str("# TYPE sgl_router_kv_tree_maintained gauge\n"); + body.push_str(&format!( + "sgl_router_kv_tree_maintained {}\n", + u8::from(ctx.kv_metrics.is_some()), + )); + if let Some(kv) = ctx.kv_metrics.as_ref() { + body.push_str(&render_kv_tiers( + kv, + ctx.block_size_oracle.get().unwrap_or(0), + )); + } ( StatusCode::OK, [(CONTENT_TYPE, PROMETHEUS_CONTENT_TYPE)], @@ -57,6 +78,101 @@ pub async fn metrics(State(ctx): State>) -> impl IntoResponse { ) } +/// 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. +/// +/// These exist to make a router-vs-engine tier mismatch a number instead of +/// an inference. `sgl_router_kv_tree_blocks * sgl_router_kv_block_size` +/// for a worker and tier, divided by that pod's own occupancy of the tier +/// (device: `sglang_kv_used_tokens + sglang_kv_evictable_tokens`; host: +/// `sglang_hicache_host_used_tokens`; `tp_rank="0"`), is the tree's coverage +/// of the tier. About 1 means the tree mirrors the engine; about 0 means the +/// engine holds a tier that routing cannot see; a missing series means the +/// worker publishes nothing. The event counters show whether the tagged +/// stream that should feed the tree is arriving at all. +/// Emitted only when the router maintains a local tree; in metadata-only mode +/// (external Indexer) the tier and event series here would be a structural +/// zero, which the HELP text below would have the operator read as a missing +/// tier stream. See `KvEventIndex::metrics_source`. +/// +/// `block_size` is 0 until the first worker reports, and a coverage panel +/// multiplies by it, so such a panel reads 0 rather than NaN on a fleet that +/// has not registered yet. +fn render_kv_tiers(kv: &KvIndexMetrics, block_size: u32) -> String { + let mut out = String::new(); + + out.push_str( + "# HELP sgl_router_kv_block_size Tokens per KV block hash, as established from the fleet (0 until a worker reports). Multiply sgl_router_kv_tree_blocks by this to compare with the engine's token gauges.\n", + ); + out.push_str("# TYPE sgl_router_kv_block_size gauge\n"); + out.push_str(&format!("sgl_router_kv_block_size {block_size}\n")); + + // Every tier is emitted per carrier, zeros included: a host row at 0 next + // to a device row in the millions is the mismatch signature, and an + // absent series cannot be told from a tier the tree never tracked. + out.push_str( + "# HELP sgl_router_kv_tree_blocks Blocks the cache-aware tree attributes to a worker rank, by the storage tier the worker holds them on (a block held on device and host counts under both). Times sgl_router_kv_block_size, and divided by the engine's own occupancy of that tier for the same pod (device: sglang_kv_used_tokens + sglang_kv_evictable_tokens; host: sglang_hicache_host_used_tokens; tp_rank=\"0\"), this is the tree's coverage of the tier: ~1 mirrors the engine, ~0 means the engine holds a tier routing cannot see.\n", + ); + out.push_str("# TYPE sgl_router_kv_tree_blocks gauge\n"); + for (id, counts) in kv.tree.tier_occupancy() { + for (slot, (_, tier)) in Tiers::SLOTS.iter().enumerate() { + out.push_str(&format!( + "sgl_router_kv_tree_blocks{{worker_url=\"{}\",dp_rank=\"{}\",tier=\"{}\"}} {}\n", + escape_label(&id.url), + id.dp_rank, + tier, + counts[slot], + )); + } + } + + let rows = kv.tally.snapshot(); + out.push_str( + "# HELP sgl_router_kv_events_total KV-cache events the pump consumed, by kind and the storage medium tag they carried (untagged = no medium field; unknown = a medium this build does not recognise). On a hierarchical-cache fleet block_stored/CPU_PINNED runs at about the block_removed/GPU rate; a CPU_PINNED row pinned at 0 with hicache enabled means the tier stream is not reaching the router. A nonzero block_stored/unknown row means the engine publishes a tier this build cannot rank and the tree is dropping those stores: upgrade the router.\n", + ); + out.push_str("# TYPE sgl_router_kv_events_total counter\n"); + for r in &rows { + out.push_str(&format!( + "sgl_router_kv_events_total{{event=\"{}\",medium=\"{}\"}} {}\n", + r.event, r.medium, r.events, + )); + } + out.push_str( + "# HELP sgl_router_kv_event_blocks_total Block hashes carried by the KV-cache events the pump consumed, by kind and storage medium tag. Times sgl_router_kv_block_size this is comparable to the engine's device eviction volume (block_removed/GPU) and, summed without its pool label, to sglang_hicache_backup_tokens_total (block_stored/CPU_PINNED). The two are not equal: the engine also evicts device blocks it never backed up.\n", + ); + out.push_str("# TYPE sgl_router_kv_event_blocks_total counter\n"); + for r in &rows { + out.push_str(&format!( + "sgl_router_kv_event_blocks_total{{event=\"{}\",medium=\"{}\"}} {}\n", + r.event, r.medium, r.blocks, + )); + } + + // A tagged removal clears only its own tier, so a batch lost in transit + // can strand a tier bit the tree will never clear on its own. Nonzero + // here is the explanation for tree coverage drifting above 1. + out.push_str( + "# HELP sgl_router_kv_event_batches_lost_total KV-event batches dropped in transit, inferred from gaps in each publisher's dense sequence number (ZMQ drops at the publisher's high-water mark). Nonzero means the tree may hold tiers a worker has already released, which shows up as sgl_router_kv_tree_blocks exceeding the engine's own occupancy of that tier.\n", + ); + out.push_str("# TYPE sgl_router_kv_event_batches_lost_total counter\n"); + out.push_str(&format!( + "sgl_router_kv_event_batches_lost_total {}\n", + kv.tally.batches_lost(), + )); + + out.push_str( + "# HELP sgl_router_kv_tree_accounting_errors_total Times the tree's per-tier occupancy bookkeeping contradicted itself. Always 0 on a correct tree. Nonzero means sgl_router_kv_tree_blocks understates what the tree holds, and can drop a worker's series entirely — which the gauge's own HELP would have you read as a worker that publishes nothing.\n", + ); + out.push_str("# TYPE sgl_router_kv_tree_accounting_errors_total counter\n"); + for (reason, count) in ACCOUNTING_REASONS.iter().zip(kv.tree.accounting_errors()) { + out.push_str(&format!( + "sgl_router_kv_tree_accounting_errors_total{{reason=\"{reason}\"}} {count}\n", + )); + } + out +} + #[cfg(test)] mod tests { use super::*; @@ -66,6 +182,139 @@ mod tests { use http_body_util::BodyExt; use tower::ServiceExt; + /// 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) + /// 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}; + + let kv = KvIndexMetrics::new(Arc::new(HashTree::new()), Arc::new(EventTally::new())); + let w = KvWorkerId::new("http://w0:30000".into(), 0); + kv.tree.insert_tiered(&w, None, &[1, 2, 3], Tiers::DEVICE); + kv.tree.insert_tiered(&w, None, &[1, 2], Tiers::HOST); + kv.tree.insert_tiered(&w, None, &[1], Tiers::EXTERNAL); + kv.tally + .record(EventKind::BlockStored, Some("CPU_PINNED"), 2); + + let out = render_kv_tiers(&kv, 64); + let blocks = |tier: &str, n: u32| { + format!( + r#"sgl_router_kv_tree_blocks{{worker_url="http://w0:30000",dp_rank="0",tier="{tier}"}} {n}"# + ) + }; + // Every tier is asserted, the zero rows included: those are the ones + // a later edit to `Tiers::SLOTS` would silently drop. + let mut want: Vec = [("device", 3), ("host", 2), ("disk", 0), ("external", 1)] + .iter() + .map(|(t, n)| blocks(t, *n)) + .collect(); + want.extend( + [ + "sgl_router_kv_block_size 64\n", + r#"sgl_router_kv_events_total{event="block_stored",medium="CPU_PINNED"} 1"#, + r#"sgl_router_kv_event_blocks_total{event="block_stored",medium="CPU_PINNED"} 2"#, + r#"sgl_router_kv_events_total{event="block_removed",medium="GPU"} 0"#, + ] + .iter() + .map(|s| (*s).to_owned()), + ); + for w in want { + assert!(out.contains(&w), "missing {w:?}; got:\n{out}"); + } + } + + /// The series are pulled from the tree on every scrape, so dropping a + /// worker must make its rows disappear rather than freeze at their last + /// value. This is what `KvEventIndex::remove_worker` relies on when it + /// calls `clear_worker`. + #[tokio::test] + async fn kv_tree_blocks_drop_with_the_worker() { + use crate::policies::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); + kv.tree.insert_tiered(&w, None, &[1, 2], Tiers::HOST); + assert!(render_kv_tiers(&kv, 64).contains("http://w0:30000")); + + kv.tree.clear_worker(&w); + let out = render_kv_tiers(&kv, 64); + assert!( + !out.contains("http://w0:30000"), + "a cleared worker must stop emitting series; got:\n{out}" + ); + // The family itself stays declared so the scrape shape is stable. + assert!(out.contains("# TYPE sgl_router_kv_tree_blocks gauge")); + } + + /// The route wiring itself: `render_kv_tiers` had two direct unit tests + /// but nothing exercised `ctx.kv_metrics`, so deleting the `if let` in the + /// handler left the suite green while `/metrics` silently stopped emitting + /// 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}; + + let mut ctx = AppContext::stub(); + let tree = Arc::new(HashTree::new()); + tree.insert_tiered( + &KvWorkerId::new("http://w0:30000".into(), 0), + None, + &[1, 2], + Tiers::HOST, + ); + ctx.kv_metrics = Some(KvIndexMetrics::new(tree, Arc::new(EventTally::new()))); + let ctx = Arc::new(ctx); + + let app = crate::server::app::build_router(ctx.clone()); + let res = app + .oneshot( + Request::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = res.into_body().collect().await.unwrap().to_bytes(); + let body = std::str::from_utf8(&body).unwrap(); + assert!(body.contains("sgl_router_kv_tree_maintained 1")); + assert!(body.contains( + r#"sgl_router_kv_tree_blocks{worker_url="http://w0:30000",dp_rank="0",tier="host"} 2"# + )); + assert!(body.contains("sgl_router_kv_event_batches_lost_total 0")); + assert!( + body.contains(r#"sgl_router_kv_tree_accounting_errors_total{reason="underflow"} 0"#) + ); + } + + /// The other half of the gate: with no local tree the families are absent, + /// but the mode gauge still says so rather than leaving the operator to + /// guess whether the endpoint regressed. + #[tokio::test] + async fn metrics_endpoint_reports_the_mode_when_no_tree_is_maintained() { + let ctx = Arc::new(AppContext::stub()); + assert!(ctx.kv_metrics.is_none()); + let app = crate::server::app::build_router(ctx.clone()); + let res = app + .oneshot( + Request::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = res.into_body().collect().await.unwrap().to_bytes(); + let body = std::str::from_utf8(&body).unwrap(); + assert!(body.contains("sgl_router_kv_tree_maintained 0")); + assert!( + !body.contains("sgl_router_kv_tree_blocks{"), + "a structural zero must not be emitted as if it were a reading", + ); + } + #[tokio::test] async fn metrics_endpoint_returns_prometheus_text() { let ctx = Arc::new(AppContext::stub());