[Router] Expose the KV storage-tier stream and tree occupancy on /metrics (2/4) (#39109)
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Kangyan Zhou
Claude Opus 5
Kangyan Zhou
Claude Opus 5
parent
24874f90a3
commit
3c48c1e967
@@ -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
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -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<u32> {
|
||||
.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<HashTree>,
|
||||
pub(crate) tally: Arc<EventTally>,
|
||||
}
|
||||
|
||||
impl KvIndexMetrics {
|
||||
pub fn new(tree: Arc<HashTree>, tally: Arc<EventTally>) -> 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<Mutex<HashMap<KvWorkerId, i64>>>,
|
||||
/// Applied events by kind and storage medium, for the `/metrics` scrape.
|
||||
/// Written only by the pump.
|
||||
tally: Arc<EventTally>,
|
||||
/// 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<Mutex<HashMap<KvWorkerId, i64>>> = Arc::new(Mutex::new(HashMap::new()));
|
||||
let live_workers: Arc<Mutex<HashSet<KvWorkerId>>> = Arc::new(Mutex::new(HashSet::new()));
|
||||
let pump_cancel = CancellationToken::new();
|
||||
let tally = Arc::new(EventTally::new());
|
||||
let pump = tokio::spawn(pump_loop(
|
||||
tree.clone(),
|
||||
engine_load.clone(),
|
||||
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<KvIndexMetrics> {
|
||||
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<EngineLoadTable>,
|
||||
cursors: Arc<Mutex<HashMap<KvWorkerId, i64>>>,
|
||||
live_workers: Arc<Mutex<HashSet<KvWorkerId>>>,
|
||||
tally: Arc<EventTally>,
|
||||
cancel: CancellationToken,
|
||||
mut rx: mpsc::Receiver<WorkerEvent>,
|
||||
) {
|
||||
@@ -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<HashTree>,
|
||||
engine_load: Arc<EngineLoadTable>,
|
||||
cursors: Arc<Mutex<HashMap<KvWorkerId, i64>>>,
|
||||
tally: Arc<EventTally>,
|
||||
#[allow(dead_code)]
|
||||
live_set: Arc<Mutex<HashSet<KvWorkerId>>>,
|
||||
#[allow(dead_code)]
|
||||
@@ -557,12 +627,14 @@ mod tests {
|
||||
let live_set: Arc<Mutex<HashSet<KvWorkerId>>> =
|
||||
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<i64>| {
|
||||
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]
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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<HashSet<String>>,
|
||||
}
|
||||
|
||||
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<TallyRow> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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<KvWorkerId, Tiers>, 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<KvWorkerId, Tiers>,
|
||||
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<NodeId, Node>,
|
||||
by_hash: HashMap<i64, HashSet<NodeId>>,
|
||||
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<KvWorkerId, TierCounts>,
|
||||
/// 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<NodeId> = 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<KvWorkerId, TierCounts> = 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
|
||||
|
||||
@@ -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<Arc<dyn sgl_kv_indexer::PrefixIndex>>,
|
||||
pub radix_tree_prefix_provider: Option<RadixTreePrefixProvider>,
|
||||
pub block_size_oracle: Arc<BlockSizeOracle>,
|
||||
/// Read-only handles `/metrics` pulls the KV storage-tier series from on
|
||||
/// scrape. `None` when this router maintains no local tree (external
|
||||
/// Indexer), where those series would all be a structural zero — see
|
||||
/// [`crate::policies::kv_events::KvEventIndex::metrics_source`].
|
||||
pub kv_metrics: Option<KvIndexMetrics>,
|
||||
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),
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<Arc<AppContext>>) -> 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<Arc<AppContext>>) -> 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<String> = [("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());
|
||||
|
||||
Reference in New Issue
Block a user