diff --git a/experimental/sgl-router/src/policies/kv_events/index.rs b/experimental/sgl-router/src/policies/kv_events/index.rs index 4f3cf20fe..83c52507d 100644 --- a/experimental/sgl-router/src/policies/kv_events/index.rs +++ b/experimental/sgl-router/src/policies/kv_events/index.rs @@ -38,7 +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::tree::{HashTree, KvWorkerId}; +use super::tree::{HashTree, KvWorkerId, Tiers}; use super::wire::KvCacheEvent; use crate::policies::engine_load::EngineLoadTable; @@ -482,12 +482,25 @@ async fn pump_loop( } } for event in &batch.events { + // The `medium` tag decides which tier a store lands on and + // which tier a removal clears, so a device eviction leaves + // a worker that still holds the block on host as an owner + // — see the tree's "Storage tiers" docs. match event { KvCacheEvent::BlockStored(b) => { - tree.insert(&worker, b.parent_block_hash, &b.block_hashes); + tree.insert_tiered( + &worker, + b.parent_block_hash, + &b.block_hashes, + Tiers::for_store(b.medium.as_deref()), + ); } KvCacheEvent::BlockRemoved(b) => { - tree.remove(&worker, &b.block_hashes); + tree.remove_tiered( + &worker, + &b.block_hashes, + Tiers::for_remove(b.medium.as_deref()), + ); } KvCacheEvent::AllBlocksCleared => { tree.clear_worker(&worker); @@ -593,7 +606,54 @@ mod tests { let m = tree.match_prefix(None, &[10, 20, 30]); assert_eq!(m.matched_blocks, 3); - assert!(m.workers.contains(&id), "tree must hold the worker"); + assert!(m.workers().contains(&id), "tree must hold the worker"); + } + + /// The pump must carry each event's `medium` into the tree. The engine's + /// write-back sequence for a backed-up block is a host-tagged store + /// followed by a device-tagged removal; applied tier-blind, the removal + /// erased the worker and every repeat of that prefix routed cold for the + /// whole host retention horizon. + #[tokio::test] + async fn pump_keeps_host_backed_block_owned_across_device_eviction() { + let id = worker_id("http://w1", 0); + let h = spawn_pump(std::slice::from_ref(&id)); + let (tree, tx, pump) = (h.tree, h.tx, h.pump); + + let stored = |medium: Option<&str>| { + KvCacheEvent::BlockStored(BlockStored { + parent_block_hash: None, + block_hashes: vec![10, 20], + token_ids: vec![], + block_size: 64, + lora_id: None, + medium: medium.map(str::to_owned), + }) + }; + let removed = |medium: Option<&str>| { + KvCacheEvent::BlockRemoved(BlockRemoved { + block_hashes: vec![20], + medium: medium.map(str::to_owned), + }) + }; + tx.send(WorkerEvent::Batch { + worker: id.clone(), + seq: 1, + batch: batch(vec![ + stored(Some("GPU")), + stored(Some("CPU_PINNED")), + removed(Some("GPU")), + ]), + }) + .await + .unwrap(); + drop(tx); + pump.await.unwrap(); + + let m = tree.match_prefix(None, &[10, 20]); + assert_eq!(m.matched_blocks, 2, "host copy keeps the block routable"); + assert!(m.workers().contains(&id)); + assert!(!m.device_workers().contains(&id), "device copy is gone"); } /// A `WorkerEvent::Load` lands in the engine-load table (gauge, no diff --git a/experimental/sgl-router/src/policies/kv_events/mod.rs b/experimental/sgl-router/src/policies/kv_events/mod.rs index 86b9c5de7..b6a0e1786 100644 --- a/experimental/sgl-router/src/policies/kv_events/mod.rs +++ b/experimental/sgl-router/src/policies/kv_events/mod.rs @@ -10,7 +10,8 @@ //! - [`wire`] — msgpack types and [`decode_event_batch`]; the contract //! with the SGLang publisher. Pure decoding; no I/O. //! - [`hash`] — block-hash compute mirroring SGLang `RadixKey.hash_page`. -//! - [`tree`] — hash-keyed radix tree consumed by the routing path. +//! - [`tree`] — hash-keyed radix tree consumed by the routing path, +//! tracking the storage tier each worker holds a block on. //! - [`subscriber`] — per-worker ZMQ SUB tasks. //! - [`discovery`] — `/server_info` parse → publisher endpoint. //! - [`index`] — public façade bundling the tree + subscribers + pump. @@ -29,7 +30,7 @@ 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 subscriber::{KvEventSubscriberRegistry, SubKind, WorkerEvent}; -pub use tree::{HashTree, KvWorkerId, MatchResult}; +pub use tree::{HashTree, KvWorkerId, MatchResult, Tiers}; pub use wire::{ decode_event_batch, BlockRemoved, BlockStored, DecodeError, KvCacheEvent, KvEventBatch, }; diff --git a/experimental/sgl-router/src/policies/kv_events/tree.rs b/experimental/sgl-router/src/policies/kv_events/tree.rs index ba1b1d60f..d5acb4f10 100644 --- a/experimental/sgl-router/src/policies/kv_events/tree.rs +++ b/experimental/sgl-router/src/policies/kv_events/tree.rs @@ -3,8 +3,8 @@ //! Each non-root node represents one block hash (`i64`). A node's children //! are keyed by the *next* block hash in a chain, so a path from the root //! down to depth `n` represents a chain of `n` block hashes. Every node -//! tracks the set of [`KvWorkerId`]s that hold the chain ending at that -//! node. +//! tracks the [`KvWorkerId`]s that hold the chain ending at that node, and +//! on which storage [`Tiers`] each of them holds it. //! //! The tree is fed by `BlockStored` / `BlockRemoved` / `AllBlocksCleared` //! events from SGLang workers (decoded by [`super::wire`]) and is queried @@ -39,6 +39,41 @@ //! from the reverse index. Pruning cascades upward iteratively (chains //! can be deep — the recursive form would risk stack-overflow for //! pathological inputs). +//! +//! # Storage tiers +//! +//! An engine running a hierarchical cache holds a block on device and, once +//! it has been backed up, on host pinned memory (or a storage backend) as +//! well. It publishes every tier transition as its own event, tagged with a +//! `medium`: a host-tier `BlockStored` when the backup lands, a device-tier +//! `BlockRemoved` when the device copy is evicted, a host-tier `BlockRemoved` +//! when the host copy goes. +//! +//! The two orderings differ by write policy, and neither may be assumed. +//! Under write-through the pending D2H copy holds a lock ref, so the host +//! store is published before the device eviction. Under write-back +//! `_detach_backuped` publishes the device removal as soon as host slots are +//! reserved and the host store follows only when the copy lands — the +//! inverse. The tree therefore has to converge either way, and it does: a +//! removal clears only its own tier, and a later store re-adds the chain. +//! +//! WHY the tree keeps the tiers apart instead of treating any `BlockRemoved` +//! as "the worker lost the block": a device eviction that leaves a host copy +//! behind does not end the worker's ability to serve the prefix — it loads +//! the copy back at memory speed instead of recomputing it. Dropping the +//! worker on that event makes every prefix unroutable after one device +//! turnover, even though the fleet still holds it for the whole host +//! retention horizon, and the repeat request lands on a random worker that +//! then prefills it cold. Ownership here is "any tier", which is what makes +//! the prefix routable again. [`MatchResult::tiers`] reports which tier each +//! owner holds it on so a policy can price a load-back against an in-place +//! hit; no policy consumes it yet — the routing path reads +//! [`HashTree::prefix_depths`], which is tier-blind by design. +//! +//! Untagged events keep their pre-tiering meaning: an untagged store is a +//! device store, an untagged remove clears every tier. A store tagged with a +//! `medium` this build cannot rank is dropped rather than filed under a +//! guess, while a removal tagged with one clears every tier. See [`Tiers`]. use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, Ordering}; @@ -96,15 +131,173 @@ impl KvWorkerId { } } -/// Result of [`HashTree::match_prefix`]. -#[derive(Debug, Clone)] +/// The storage tiers on which one worker holds one block — a bitset, because +/// a backed-up block sits on device AND host at once (module docs, "Storage +/// tiers"). +/// +/// A worker owns a block, and is a routing candidate for it, while ANY bit is +/// set. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)] +pub struct Tiers(u8); + +impl Tiers { + /// Device HBM: `medium = "GPU"`, or an untagged event from a publisher + /// that predates tiering. + pub const DEVICE: Tiers = Tiers(1); + /// Host pinned memory: `medium = "CPU_PINNED"`. + pub const HOST: Tiers = Tiers(1 << 1); + /// L3 local SSD / NVMe: `medium = "DISK"`. + pub const DISK: Tiers = Tiers(1 << 2); + /// L4 shared / remote pool, e.g. Mooncake: `medium = "EXTERNAL"`. + /// + /// Separate from [`Self::DISK`] rather than folded into one "storage" bit + /// because the engine treats them as distinct tiers + /// (`StorageMedium` in `python/sglang/srt/disaggregation/kv_events.py`). + /// Sharing a bit would make an L3 eviction erase the router's knowledge of + /// an L4 copy on a fleet running both. + pub const EXTERNAL: Tiers = Tiers(1 << 3); + /// Every tier — what an untagged `BlockRemoved` clears. + pub const ALL: Tiers = Self::DEVICE + .union(Self::HOST) + .union(Self::DISK) + .union(Self::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. + pub const WIRE_MEDIA: [(&'static str, Tiers); 4] = [ + ("GPU", Self::DEVICE), + ("CPU_PINNED", Self::HOST), + ("DISK", Self::DISK), + ("EXTERNAL", Self::EXTERNAL), + ]; + + /// The tiers a `BlockStored` tagged `medium` lands on. Untagged reads as + /// device, which is what the event meant before tiers existed. An unknown + /// tag lands on NO tier: the store is dropped rather than filed under a + /// guess. + /// + /// Filing it on device would turn every store on a future medium into a + /// device hit. Filing it on an "unranked" catch-all tier is no better — it + /// makes the worker a routing candidate for a tier whose cost this build + /// cannot price, and a load-back the router assumes is cheap may be a + /// remote fetch dearer than recomputing the prefix elsewhere. Dropping + /// costs at most one cold prefill and is visible: the event still counts + /// on the tally's `unknown` medium row, and the string is logged once. + pub fn for_store(medium: Option<&str>) -> Tiers { + match medium { + None => Self::DEVICE, + Some(m) => Self::known(m).unwrap_or_default(), + } + } + + /// The tiers a `BlockRemoved` tagged `medium` clears: its own, or every + /// tier when the tag is absent or unrecognised. + /// + /// Deliberately asymmetric with [`Self::for_store`], which drops an + /// unknown tag instead of widening it. A removal is the one direction + /// where guessing wide is safe and guessing narrow is not: over-removal + /// costs at most one cold prefill, while a permanently stale owner is the + /// one failure a routing index must never manufacture. An engine that + /// stores under a `medium` this build knows but frees the block under one + /// it does not would otherwise strand that bit until the worker is + /// dropped entirely. + pub fn for_remove(medium: Option<&str>) -> Tiers { + medium.and_then(Self::known).unwrap_or(Self::ALL) + } + + fn known(medium: &str) -> Option { + Self::WIRE_MEDIA + .iter() + .find(|(name, _)| *name == medium) + .map(|(_, tier)| *tier) + } + + pub const fn is_empty(self) -> bool { + self.0 == 0 + } + + /// Whether every bit of `other` is set here. + pub const fn contains(self, other: Tiers) -> bool { + self.0 & other.0 == other.0 + } + + pub fn insert(&mut self, other: Tiers) { + self.0 |= other.0; + } + + pub fn remove(&mut self, other: 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 { + Tiers(self.0 | other.0) + } +} + +// `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. +const _: () = { + 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", + ); + i += 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) { + match carriers.get_mut(worker) { + Some(held) => held.insert(tiers), + None => { + carriers.insert(worker.clone(), tiers); + } + } +} + +/// Result of [`HashTree::match_prefix`]. `Default` is the no-match result: +/// zero depth, no carriers on any tier. +#[derive(Debug, Clone, Default)] pub struct MatchResult { /// Number of leading block hashes from the input slice that matched a /// path from the root. pub matched_blocks: usize, - /// Workers holding the deepest matched node. Empty when - /// `matched_blocks == 0`. - pub workers: HashSet, + /// Every worker holding the deepest matched node, and the tiers it holds + /// it on. A worker without the device bit serves the prefix by loading it + /// back from a lower tier — cheaper than a cold prefill, dearer than + /// serving in place — which is the ordering a policy would prefer on. + /// Empty when `matched_blocks == 0`. + /// + /// The single carrier list: [`Self::workers`] and [`Self::device_workers`] + /// are derived from it, so no two views of one match can disagree. + pub tiers: HashMap, +} + +impl MatchResult { + /// Workers holding the deepest matched node on ANY tier. + pub fn workers(&self) -> HashSet { + self.tiers.keys().cloned().collect() + } + + /// The subset of [`Self::workers`] holding the deepest matched node on + /// device. + pub fn device_workers(&self) -> HashSet { + self.tiers + .iter() + .filter(|(_, tiers)| tiers.contains(Tiers::DEVICE)) + .map(|(w, _)| w.clone()) + .collect() + } } /// Internal stable handle to a tree node. @@ -140,7 +333,10 @@ struct Node { parent_block_hash: Option, /// `None` only for the root sentinel. parent: Option, - workers: HashSet, + /// Carriers of the chain ending here, each with the tiers it holds the + /// block on. A carrier is dropped the moment its last tier bit clears; + /// an entry with empty tiers never exists (see [`TreeState::insert`]). + workers: HashMap, /// Children keyed by next-block hash. children: HashMap, last_used: AtomicU64, @@ -152,7 +348,7 @@ impl Node { block_hash, parent_block_hash, parent: Some(parent), - workers: HashSet::new(), + workers: HashMap::new(), children: HashMap::new(), last_used: AtomicU64::new(now_millis()), } @@ -192,7 +388,7 @@ impl TreeState { block_hash: ROOT_HASH_SENTINEL, parent_block_hash: None, parent: None, - workers: HashSet::new(), + workers: HashMap::new(), children: HashMap::new(), last_used: AtomicU64::new(now_millis()), }, @@ -273,7 +469,7 @@ impl TreeState { if self .nodes .get(&cand) - .is_some_and(|n| n.workers.contains(worker)) + .is_some_and(|n| n.workers.contains_key(worker)) { return cand; } @@ -288,8 +484,18 @@ impl TreeState { ROOT_ID } - fn insert(&mut self, worker: &KvWorkerId, parent_hash: Option, block_hashes: &[i64]) { - if block_hashes.is_empty() { + /// Mark every node along `block_hashes` as held by `worker` on `tiers`, + /// adding to whatever tiers it already holds there. Empty `tiers` is a + /// no-op: a carrier entry with no tier would be an owner of nothing, and + /// `remove` relies on "no bits ⇒ no entry" to know when to prune. + fn insert( + &mut self, + worker: &KvWorkerId, + parent_hash: Option, + block_hashes: &[i64], + tiers: Tiers, + ) { + if block_hashes.is_empty() || tiers.is_empty() { return; } let mut current = self.resolve_parent(worker, parent_hash); @@ -315,14 +521,17 @@ impl TreeState { ); return; }; - child.workers.insert(worker.clone()); + add_tiers(&mut child.workers, worker, tiers); child.last_used.store(now, Ordering::Relaxed); current = child_id; prev_hash = Some(h); } } - fn remove(&mut self, worker: &KvWorkerId, block_hashes: &[i64]) { + /// Clear `tiers` from `worker`'s hold on every node carrying any hash in + /// `block_hashes`. The worker leaves a node once no tier bit remains, and + /// a node that becomes empty + childless is pruned. + fn remove(&mut self, worker: &KvWorkerId, block_hashes: &[i64], tiers: Tiers) { // Collect all node ids to touch (fixed snapshot — avoids iterator // invalidation when pruning mutates `by_hash`). let mut targets: Vec = Vec::new(); @@ -334,14 +543,19 @@ impl TreeState { for id in targets { // Node may already be gone if a previous prune in this batch // cascaded through it — skip silently. - let still_present = match self.nodes.get_mut(&id) { + let prunable = match self.nodes.get_mut(&id) { Some(node) => { - node.workers.remove(worker); + if let Some(held) = node.workers.get_mut(worker) { + held.remove(tiers); + if held.is_empty() { + node.workers.remove(worker); + } + } node.workers.is_empty() && node.children.is_empty() } None => false, }; - if still_present { + if prunable { self.prune_cascade(id); } } @@ -358,7 +572,7 @@ impl TreeState { let mut prune_candidates: Vec = Vec::new(); for id in ids { if let Some(node) = self.nodes.get_mut(&id) { - if node.workers.remove(worker) + if node.workers.remove(worker).is_some() && node.workers.is_empty() && node.children.is_empty() { @@ -448,10 +662,7 @@ impl TreeState { /// [`HashTree::match_prefix`] documents the policy for callers. fn match_prefix(&self, parent_hash: Option, block_hashes: &[i64]) -> MatchResult { if block_hashes.is_empty() { - return MatchResult { - matched_blocks: 0, - workers: HashSet::new(), - }; + return MatchResult::default(); } // Determine starting node: root, or the unique node carrying // `parent_hash`. Multiple matches: bail to root (caller should @@ -487,17 +698,13 @@ impl TreeState { None => break, } } - let workers = match last_match_node { - Some(id) => self - .nodes - .get(&id) - .map(|n| n.workers.clone()) - .unwrap_or_default(), - None => HashSet::new(), - }; + let tiers: HashMap = last_match_node + .and_then(|id| self.nodes.get(&id)) + .map(|n| n.workers.clone()) + .unwrap_or_default(); MatchResult { matched_blocks: matched, - workers, + tiers, } } @@ -545,12 +752,14 @@ impl TreeState { reached += 1; if reached == 1 { // Absent here means never in `alive`: a tail held without - // block 0 is reported as holding nothing. - alive = child.workers.iter().collect(); + // block 0 is reported as holding nothing. Presence is "on any + // tier" — a worker whose device copy was evicted but whose + // host backup remains still holds the level. + alive = child.workers.keys().collect(); } else { let mut still = Vec::with_capacity(alive.len()); for w in alive { - if child.workers.contains(w) { + if child.workers.contains_key(w) { still.push(w); } else { depths.insert(w, reached - 1); @@ -665,27 +874,51 @@ impl HashTree { } } - /// Apply a `BlockStored` event. - /// - /// Walks from `parent_hash`'s node (or root) and descends along - /// `block_hashes`, marking every visited node as held by `worker`. - /// Empty `block_hashes` is a no-op. + /// Apply an untagged `BlockStored` event: a device store. See + /// [`Self::insert_tiered`]. pub fn insert(&self, worker: &KvWorkerId, parent_hash: Option, block_hashes: &[i64]) { - let mut state = self.state.write(); - state.insert(worker, parent_hash, block_hashes); + self.insert_tiered(worker, parent_hash, block_hashes, Tiers::DEVICE); } - /// Apply a `BlockRemoved` event. + /// Apply a `BlockStored` event on `tiers` (from its `medium`, via + /// [`Tiers::for_store`]). /// - /// For every node carrying any hash in `block_hashes`, drop `worker` - /// from that node's worker set. Nodes that become empty AND childless - /// are pruned (cascading upward). + /// Walks from `parent_hash`'s node (or root) and descends along + /// `block_hashes`, marking every visited node as held by `worker` on + /// `tiers` in addition to any tier it already holds there. Empty + /// `block_hashes` or empty `tiers` is a no-op. + pub fn insert_tiered( + &self, + worker: &KvWorkerId, + parent_hash: Option, + block_hashes: &[i64], + tiers: Tiers, + ) { + let mut state = self.state.write(); + state.insert(worker, parent_hash, block_hashes, tiers); + } + + /// Apply an untagged `BlockRemoved` event: the worker loses the blocks on + /// every tier. See [`Self::remove_tiered`]. + pub fn remove(&self, worker: &KvWorkerId, block_hashes: &[i64]) { + self.remove_tiered(worker, block_hashes, Tiers::ALL); + } + + /// Apply a `BlockRemoved` event for `tiers` (from its `medium`, via + /// [`Tiers::for_remove`]). + /// + /// For every node carrying any hash in `block_hashes`, clear `tiers` + /// from `worker`'s hold on it. The worker stays an owner of the node + /// while it holds the block on any other tier — a device eviction after a + /// host backup leaves the worker a host-tier owner. Once no tier remains + /// the worker is dropped, and nodes that become empty AND childless are + /// pruned (cascading upward). /// /// Removing the worker from a node does NOT remove the node if other /// workers still hold it. - pub fn remove(&self, worker: &KvWorkerId, block_hashes: &[i64]) { + pub fn remove_tiered(&self, worker: &KvWorkerId, block_hashes: &[i64], tiers: Tiers) { let mut state = self.state.write(); - state.remove(worker, block_hashes); + state.remove(worker, block_hashes, tiers); } /// Apply an `AllBlocksCleared` event for `worker`. @@ -767,6 +1000,20 @@ impl HashTree { // Tests // --------------------------------------------------------------------------- +#[cfg(test)] +impl HashTree { + /// Whether every carrier on every node holds at least one tier. `remove` + /// and `prune_cascade` rely on "no bits ⇒ no entry"; a violation means a + /// node can never be pruned and a worker never dropped. + fn debug_no_empty_carrier(&self) -> bool { + let state = self.state.read(); + state + .nodes + .values() + .all(|n| n.workers.values().all(|t| !t.is_empty())) + } +} + #[cfg(test)] mod tests { use super::*; @@ -809,7 +1056,7 @@ mod tests { // so `match_prefix` credits it with the full chain scored here at 1. let m = tree.match_prefix(None, &chain); assert_eq!(m.matched_blocks, 4); - assert_eq!(m.workers, workers(&[&deep, &holed])); + assert_eq!(m.workers(), workers(&[&deep, &holed])); } #[test] @@ -817,11 +1064,302 @@ mod tests { let tree = HashTree::new(); let m = tree.match_prefix(None, &[]); assert_eq!(m.matched_blocks, 0); - assert!(m.workers.is_empty()); + assert!(m.workers().is_empty()); let m2 = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m2.matched_blocks, 0); - assert!(m2.workers.is_empty()); + assert!(m2.workers().is_empty()); + } + + /// The write-back sequence the engine publishes for a backed-up block: + /// device store, host store once the D2H copy lands, then a DEVICE-tagged + /// removal when the device copy is evicted. The worker still holds the + /// block on host, so it must stay an owner — just no longer a device one. + /// Only the host-tagged removal ends ownership. + #[test] + fn device_eviction_after_host_backup_keeps_the_worker_as_host_owner() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + tree.insert_tiered(&a, None, &[1, 2, 3], Tiers::for_store(Some("GPU"))); + tree.insert_tiered(&a, None, &[1, 2, 3], Tiers::for_store(Some("CPU_PINNED"))); + + let m = tree.match_prefix(None, &[1, 2, 3]); + assert_eq!(m.matched_blocks, 3); + assert_eq!(m.workers(), workers(&[&a])); + assert_eq!(m.device_workers(), workers(&[&a]), "held on device too"); + + tree.remove_tiered(&a, &[3], Tiers::for_remove(Some("GPU"))); + let m = tree.match_prefix(None, &[1, 2, 3]); + assert_eq!(m.matched_blocks, 3, "host copy keeps the chain matchable"); + assert_eq!( + m.workers(), + workers(&[&a]), + "host-only holder is still an owner" + ); + assert!( + m.device_workers().is_empty(), + "but no longer a device owner" + ); + + tree.remove_tiered(&a, &[3], Tiers::for_remove(Some("CPU_PINNED"))); + let m = tree.match_prefix(None, &[1, 2, 3]); + assert_eq!(m.matched_blocks, 2, "last tier gone: node pruned"); + } + + /// The routing path reads `prefix_depths`, not `match_prefix`, so the + /// tier fix has to show up there: a chain whose device copy was evicted + /// after a host backup still scores at its full depth. The untagged + /// removal is the control — it clears every tier and the depth collapses, + /// which is what the tier-blind tree did for the tagged case too. + #[test] + fn prefix_depths_survive_a_device_eviction_with_a_host_backup() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + tree.insert_tiered(&a, None, &[1, 2, 3], Tiers::for_store(Some("GPU"))); + tree.insert_tiered(&a, None, &[1, 2, 3], Tiers::for_store(Some("CPU_PINNED"))); + + // Block 1 included on purpose. `prefix_depths` SEEDS its live set from + // level 1 and narrows from there, so evicting only from block 2 onward + // would leave that seed exercised against device-held state alone — and + // a seed narrowed to device owners would send every host-only worker to + // depth 0, restoring this bug with the suite still green. + tree.remove_tiered(&a, &[1, 2, 3], Tiers::for_remove(Some("GPU"))); + assert_eq!( + tree.prefix_depths(None, &[1, 2, 3]).get(&a).copied(), + Some(3), + "host backup keeps every level attributed to the worker", + ); + + tree.remove_tiered(&a, &[1, 2, 3], Tiers::for_remove(None)); + assert_eq!( + tree.prefix_depths(None, &[1, 2, 3]).get(&a).copied(), + None, + "an untagged removal still clears every tier, host backup included", + ); + } + + /// Untagged events keep the pre-tiering contract: a bare `BlockStored` is + /// a device store, a bare `BlockRemoved` clears every tier at once. A + /// publisher that never tags must see exactly the behaviour it always had. + #[test] + fn untagged_events_keep_the_legacy_meaning() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + tree.insert(&a, None, &[1, 2]); + let m = tree.match_prefix(None, &[1, 2]); + assert_eq!( + m.device_workers(), + workers(&[&a]), + "untagged store is device" + ); + + // Add a host copy, then an UNTAGGED removal: must clear both. + tree.insert_tiered(&a, None, &[1, 2], Tiers::HOST); + tree.remove(&a, &[2]); + assert_eq!(tree.match_prefix(None, &[1, 2]).matched_blocks, 1); + } + + /// An unknown medium is asymmetric on purpose: the store is dropped (no + /// tier at all), while a removal tagged with it clears every tier. Filing + /// the store under a guess would make the worker a routing candidate on a + /// tier this build cannot price; narrowing the removal to a bit that was + /// never set would leave a stale owner forever. + #[test] + fn unknown_medium_drops_the_store_and_removes_everything() { + assert_eq!(Tiers::for_store(Some("NVLINK_PEER")), Tiers::default()); + assert_eq!(Tiers::for_remove(Some("NVLINK_PEER")), Tiers::ALL); + assert_eq!(Tiers::for_store(None), Tiers::DEVICE); + assert_eq!(Tiers::for_remove(None), Tiers::ALL); + assert_eq!(Tiers::for_store(Some("CPU_PINNED")), Tiers::HOST); + assert_eq!(Tiers::for_store(Some("DISK")), Tiers::DISK); + assert_eq!(Tiers::for_store(Some("EXTERNAL")), Tiers::EXTERNAL); + // A known tag clears its own tier and nothing else. + assert_eq!(Tiers::for_remove(Some("DISK")), Tiers::DISK); + assert_eq!(Tiers::for_remove(Some("EXTERNAL")), Tiers::EXTERNAL); + + let tree = HashTree::new(); + let a = worker("http://a", 0); + let b = worker("http://b", 0); + tree.insert_tiered(&a, None, &[1], Tiers::DEVICE); + tree.insert_tiered(&a, None, &[1], Tiers::HOST); + tree.insert_tiered(&b, None, &[1], Tiers::for_store(Some("NVLINK_PEER"))); + let m = tree.match_prefix(None, &[1]); + assert_eq!( + m.workers(), + workers(&[&a]), + "a store on an unrankable tier must not make the worker an owner" + ); + tree.remove_tiered(&a, &[1], Tiers::for_remove(Some("NVLINK_PEER"))); + assert_eq!( + tree.match_prefix(None, &[1]).matched_blocks, + 0, + "an unknown-medium removal clears every tier the worker held", + ); + } + + /// Tiers are per (node, worker): a device removal by one worker must not + /// touch another worker's hold on the same node, and a node whose + /// carriers differ by tier reports the device subset exactly. + #[test] + fn tiers_are_tracked_per_worker() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + let b = worker("http://b", 0); + tree.insert_tiered(&a, None, &[1, 2], Tiers::DEVICE); + tree.insert_tiered(&b, None, &[1, 2], Tiers::HOST); + + let m = tree.match_prefix(None, &[1, 2]); + assert_eq!(m.workers(), workers(&[&a, &b])); + assert_eq!(m.device_workers(), workers(&[&a])); + + tree.remove_tiered(&a, &[2], Tiers::DEVICE); + let m = tree.match_prefix(None, &[1, 2]); + assert_eq!(m.workers(), workers(&[&b]), "a dropped, b untouched"); + assert!(m.device_workers().is_empty()); + } + + /// A store on no tier must not create a carrier: `remove` relies on + /// "no bits ⇒ no entry" to know when a node is prunable. + #[test] + fn empty_tier_insert_is_noop() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + tree.insert_tiered(&a, None, &[1], Tiers::default()); + assert_eq!(tree.node_count(), 0); + } + + /// `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 + /// the L4 copy. + #[test] + fn disk_and_external_are_independent_tiers() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + tree.insert_tiered(&a, None, &[1], Tiers::for_store(Some("EXTERNAL"))); + tree.insert_tiered(&a, None, &[1], Tiers::for_store(Some("DISK"))); + + tree.remove_tiered(&a, &[1], Tiers::for_remove(Some("DISK"))); + assert_eq!( + tree.match_prefix(None, &[1]).workers(), + workers(&[&a]), + "an L3 eviction must not take the L4 copy with it", + ); + + tree.remove_tiered(&a, &[1], Tiers::for_remove(Some("EXTERNAL"))); + 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. + #[test] + fn carrier_invariant_holds_under_a_random_walk() { + // xorshift64*, so the walk is reproducible without a dev-dependency. + struct Rng(u64); + impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } + } + + const MEDIA: [Option<&str>; 6] = [ + Some("GPU"), + Some("CPU_PINNED"), + Some("DISK"), + Some("EXTERNAL"), + Some("NVLINK_PEER"), + None, + ]; + + for seed in 1..=16u64 { + let tree = HashTree::new(); + let ws: Vec = (0..4) + .map(|i| worker(&format!("http://w{i}"), i % 2)) + .collect(); + let mut rng = Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + + for step in 0..400 { + let w = &ws[rng.below(ws.len() as u64) as usize]; + let medium = MEDIA[rng.below(MEDIA.len() as u64) as usize]; + // A small hash space so chains collide and share nodes. + let hashes: Vec = (0..1 + rng.below(4)) + .map(|_| rng.below(12) as i64) + .collect(); + // Sometimes anchor to a hash that may or may not be in the + // tree, to exercise `resolve_parent`'s fallbacks. + let parent = (rng.below(4) == 0).then(|| rng.below(12) as i64); + + match rng.below(16) { + 0 => tree.clear_worker(w), + 1 => { + tree.evict_lru(rng.below(20) as usize); + } + 2..=6 => tree.remove_tiered(w, &hashes, Tiers::for_remove(medium)), + _ => tree.insert_tiered(w, parent, &hashes, Tiers::for_store(medium)), + } + + assert!( + tree.debug_no_empty_carrier(), + "seed {seed} step {step}: a carrier is present holding no tier", + ); + } + } + } + + /// `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. + #[test] + fn clear_worker_drops_carriers_on_every_tier() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + let b = worker("http://b", 0); + tree.insert_tiered(&a, None, &[1, 2], Tiers::HOST); + tree.insert_tiered(&a, None, &[3], Tiers::for_store(Some("EXTERNAL"))); + tree.insert_tiered(&b, None, &[1, 2], Tiers::DEVICE); + + tree.clear_worker(&a); + assert_eq!( + tree.match_prefix(None, &[1, 2]).workers(), + workers(&[&b]), + "a host-only carrier must be cleared like any other", + ); + assert_eq!(tree.match_prefix(None, &[3]).matched_blocks, 0); + } + + /// `resolve_parent` disambiguates a shared hash by preferring a candidate + /// the worker already holds — on ANY tier. Narrowing that to device owners + /// would re-anchor a continuation at the root once the worker's device + /// copy was evicted, fragmenting the very prefix the host tier still + /// serves. The shape is a shared system prompt whose block hash also + /// appears in another chain. + #[test] + fn ambiguous_parent_resolves_through_a_host_only_hold() { + let tree = HashTree::new(); + let a = worker("http://a", 0); + let b = worker("http://b", 0); + tree.insert_tiered(&a, None, &[5, 7], Tiers::DEVICE); + tree.insert_tiered(&a, None, &[5, 7], Tiers::HOST); + tree.remove_tiered(&a, &[5, 7], Tiers::for_remove(Some("GPU"))); + // A second chain carrying hash 7, so `parent_hash = 7` is ambiguous. + tree.insert_tiered(&b, None, &[6, 7], Tiers::DEVICE); + + tree.insert_tiered(&a, Some(7), &[8], Tiers::DEVICE); + assert_eq!( + tree.prefix_depths(None, &[5, 7, 8]).get(&a).copied(), + Some(3), + "the continuation must attach under the host-only hold, not at root", + ); } #[test] @@ -832,21 +1370,21 @@ mod tests { let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); let m = tree.match_prefix(None, &[1, 2]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); // Diverges at depth 3 (input asks for 4, tree has 3). let m = tree.match_prefix(None, &[1, 2, 4]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); // No match at root. let m = tree.match_prefix(None, &[9, 9]); assert_eq!(m.matched_blocks, 0); - assert!(m.workers.is_empty()); + assert!(m.workers().is_empty()); } #[test] @@ -860,16 +1398,16 @@ mod tests { // Common prefix node carries both. let m = tree.match_prefix(None, &[1, 2]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&a, &b])); + assert_eq!(m.workers(), workers(&[&a, &b])); // Divergent leaf carries only the matching worker. let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); let m = tree.match_prefix(None, &[1, 2, 4]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&b])); + assert_eq!(m.workers(), workers(&[&b])); } #[test] @@ -881,7 +1419,7 @@ mod tests { let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); } #[test] @@ -902,12 +1440,12 @@ mod tests { // Match length 2 lands on node 2 (workers empty), so workers={}. let m = tree.match_prefix(None, &[1, 2]); assert_eq!(m.matched_blocks, 2); - assert!(m.workers.is_empty()); + assert!(m.workers().is_empty()); // Match length 3 lands on node 3 (workers still has A). let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); // Reverse-index sanity for hash 2: still present (node holds it). { @@ -932,13 +1470,13 @@ mod tests { let m = tree.match_prefix(None, &[1, 2, 3]); // Node 3 was pruned, so only 2 levels match. assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&b])); + assert_eq!(m.workers(), workers(&[&b])); // [1,2] now has only B (A was the only other holder of node 2; // wait — actually A held 1 and 2 too. But B also holds 1 and 2.) let m = tree.match_prefix(None, &[1, 2]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&b])); + assert_eq!(m.workers(), workers(&[&b])); // Node count: root + 1 + 2 + 4 (no 3) = 3 non-root nodes. assert_eq!(tree.node_count(), 3); @@ -986,11 +1524,11 @@ mod tests { // Both chains exist independently. let m = tree.match_prefix(None, &[1, 5]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); let m = tree.match_prefix(None, &[2, 5]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); // Reverse index for hash 5 has 2 distinct nodes. { @@ -1005,10 +1543,10 @@ mod tests { assert_eq!(tree.node_count(), 2); let m = tree.match_prefix(None, &[1, 5]); assert_eq!(m.matched_blocks, 1); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); let m = tree.match_prefix(None, &[2, 5]); assert_eq!(m.matched_blocks, 1); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); } #[test] @@ -1022,16 +1560,16 @@ mod tests { // Common prefix has both ranks. let m = tree.match_prefix(None, &[1, 2]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&w0, &w1])); + assert_eq!(m.workers(), workers(&[&w0, &w1])); // Divergent leaves: each rank on its own. let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&w0])); + assert_eq!(m.workers(), workers(&[&w0])); let m = tree.match_prefix(None, &[1, 2, 4]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&w1])); + assert_eq!(m.workers(), workers(&[&w1])); } #[test] @@ -1049,12 +1587,12 @@ mod tests { // The chain 1->5->7 must exist with A. let m = tree.match_prefix(None, &[1, 5, 7]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); // The chain 2->5 should NOT have a 7-child (we routed to A's branch). let m = tree.match_prefix(None, &[2, 5, 7]); assert_eq!(m.matched_blocks, 2); - assert_eq!(m.workers, workers(&[&b])); + assert_eq!(m.workers(), workers(&[&b])); } #[test] @@ -1073,7 +1611,7 @@ mod tests { // C is reachable as a fresh root child at hash=9. let m = tree.match_prefix(None, &[9]); assert_eq!(m.matched_blocks, 1); - assert_eq!(m.workers, workers(&[&c])); + assert_eq!(m.workers(), workers(&[&c])); } #[test] @@ -1086,7 +1624,7 @@ mod tests { assert_eq!(tree.node_count(), 3); let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); } #[test] @@ -1158,7 +1696,7 @@ mod tests { // The newer chain should still match fully. let m = tree.match_prefix(None, &[200, 201, 202]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); } #[test] @@ -1171,7 +1709,7 @@ mod tests { let m = tree.match_prefix(None, &[10, 20, 30]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&a])); + assert_eq!(m.workers(), workers(&[&a])); // Confirm parent_block_hash chain: node carrying 30 should record // parent_block_hash = Some(20). @@ -1199,6 +1737,6 @@ mod tests { let m = tree.match_prefix(None, &[1, 2, 3]); assert_eq!(m.matched_blocks, 3); - assert_eq!(m.workers, workers(&[&b])); + assert_eq!(m.workers(), workers(&[&b])); } } diff --git a/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs b/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs index b0afee7cb..fa4d0c3d6 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs @@ -155,9 +155,9 @@ fn match_prefix_is_consistent_with_concurrent_clear() { // consistent. if m.matched_blocks == chain.len() { assert!( - m.workers.contains(&w), + m.workers().contains(&w), "full match must include worker; got {:?}", - m.workers, + m.workers(), ); } } diff --git a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs index 8a1973812..ea488ee51 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs @@ -84,8 +84,8 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { let mb = router_b.tree().match_prefix(None, &hashes); let converged = ma.matched_blocks == target && mb.matched_blocks == target - && ma.workers.contains(&key) - && mb.workers.contains(&key); + && ma.workers().contains(&key) + && mb.workers().contains(&key); if converged { // Both trees agree on count AND on the worker that holds the // prefix. This is what the Radix Tree provider reads to @@ -96,7 +96,8 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { "subscribers disagreed on matched_blocks", ); assert_eq!( - ma.workers, mb.workers, + ma.workers(), + mb.workers(), "subscribers disagreed on worker set", ); break; @@ -106,7 +107,10 @@ async fn two_independent_subscribers_converge_to_same_tree_state() { "subscribers did not converge within 3s: \ router_a={{matched={}, workers={:?}}}, \ router_b={{matched={}, workers={:?}}}, target={target}", - ma.matched_blocks, ma.workers, mb.matched_blocks, mb.workers, + ma.matched_blocks, + ma.workers(), + mb.matched_blocks, + mb.workers(), ); } sequence += 1; @@ -249,34 +253,34 @@ async fn two_subscribers_merge_events_from_two_publishers() { && ay.matched_blocks == target_y && bx.matched_blocks == target_x && by.matched_blocks == target_y - && ax.workers.contains(&key_x) - && ay.workers.contains(&key_y) - && bx.workers.contains(&key_x) - && by.workers.contains(&key_y); + && ax.workers().contains(&key_x) + && ay.workers().contains(&key_y) + && bx.workers().contains(&key_x) + && by.workers().contains(&key_y); if converged { // Negative attribution: prefix X must not be attributed to // worker_y in either tree, and vice versa. A regression that // keyed events by arriving socket rather than announced // worker URL would set BOTH worker keys on each prefix. assert!( - !ax.workers.contains(&key_y), + !ax.workers().contains(&key_y), "router_a cross-attributed worker_y to prefix X: {:?}", - ax.workers, + ax.workers(), ); assert!( - !ay.workers.contains(&key_x), + !ay.workers().contains(&key_x), "router_a cross-attributed worker_x to prefix Y: {:?}", - ay.workers, + ay.workers(), ); assert!( - !bx.workers.contains(&key_y), + !bx.workers().contains(&key_y), "router_b cross-attributed worker_y to prefix X: {:?}", - bx.workers, + bx.workers(), ); assert!( - !by.workers.contains(&key_x), + !by.workers().contains(&key_x), "router_b cross-attributed worker_x to prefix Y: {:?}", - by.workers, + by.workers(), ); break; } @@ -287,13 +291,13 @@ async fn two_subscribers_merge_events_from_two_publishers() { router_b: X={{matched={}, workers={:?}}}, Y={{matched={}, workers={:?}}}\n \ targets: X={target_x}, Y={target_y}", ax.matched_blocks, - ax.workers, + ax.workers(), ay.matched_blocks, - ay.workers, + ay.workers(), bx.matched_blocks, - bx.workers, + bx.workers(), by.matched_blocks, - by.workers, + by.workers(), ); } sequence += 1;