diff --git a/Cargo.lock b/Cargo.lock index 1229644a..2eb22bb4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1896,7 +1896,7 @@ checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" [[package]] name = "pegaflow-common" -version = "0.23.9" +version = "0.24.2" dependencies = [ "colored", "libc", @@ -1906,7 +1906,7 @@ dependencies = [ [[package]] name = "pegaflow-core" -version = "0.23.9" +version = "0.24.2" dependencies = [ "ahash", "bytesize", @@ -1939,7 +1939,7 @@ dependencies = [ [[package]] name = "pegaflow-metaserver" -version = "0.23.9" +version = "0.24.2" dependencies = [ "axum", "clap", @@ -1959,7 +1959,7 @@ dependencies = [ [[package]] name = "pegaflow-pd-wire" -version = "0.23.9" +version = "0.24.2" dependencies = [ "serde", "serde_json", @@ -1967,7 +1967,7 @@ dependencies = [ [[package]] name = "pegaflow-proto" -version = "0.23.9" +version = "0.24.2" dependencies = [ "prost", "tonic", @@ -1977,7 +1977,7 @@ dependencies = [ [[package]] name = "pegaflow-py" -version = "0.23.9" +version = "0.24.2" dependencies = [ "log", "mea", @@ -1995,7 +1995,7 @@ dependencies = [ [[package]] name = "pegaflow-server" -version = "0.23.9" +version = "0.24.2" dependencies = [ "axum", "clap", @@ -2027,7 +2027,7 @@ dependencies = [ [[package]] name = "pegaflow-transfer" -version = "0.23.9" +version = "0.24.2" dependencies = [ "anyhow", "bincode", diff --git a/Cargo.toml b/Cargo.toml index d74f9323..fb80fd5e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ members = [ resolver = "2" [workspace.package] -version = "0.23.9" +version = "0.24.2" edition = "2024" license = "Apache-2.0" diff --git a/docs/metrics.md b/docs/metrics.md index d2b1ec58..3b36e410 100644 --- a/docs/metrics.md +++ b/docs/metrics.md @@ -141,18 +141,25 @@ PegaFlow exposes the following metrics for monitoring KV cache operations: ### HLL Reuse Metrics - **pegaflow_hll_cardinality** (Gauge) - - Estimated distinct block hashes observed in a configured sliding window + - Estimated distinct `(namespace, block hash)` objects classified as misses in a configured sliding window - Labels: `window` (`15m`, `1h`, `1d` by default) - Use case: Derive approximate prefix reuse over longer windows without storing every block hash - **pegaflow_hll_total_requests** (Gauge) - - Total block hash observations in the same configured sliding window + - Total queried blocks in the same configured sliding window, including ready blocks and duplicates - Labels: `window` (`15m`, `1h`, `1d` by default) - - Use case: Denominator for HLL-based estimated hit-rate PromQL + - Use case: Denominator for HLL-based reference reuse rate -PegaFlow does not export a separate HLL hit-rate gauge. Use PromQL so the -ratio is computed from values in the same scrape: +- **pegaflow_hll_estimated_hit_rate** (Gauge) + - Server-computed miss-based infinite-cache reuse reference from the same + HLL snapshot as the two gauges above + - Labels: `window` + - Value is clamped to `[0, 1]` + +The existing cardinality and total metrics are retained. Existing PromQL +continues to work, but new dashboards should prefer the direct gauge because +it applies the same cardinality clamp as the tracker: ```promql 1 - ( @@ -162,6 +169,18 @@ ratio is computed from values in the same scrape: ) ``` +```promql +pegaflow_hll_estimated_hit_rate{window="1h"} +``` + +This is a metrics semantic update, not an `/metrics` protocol breaking change: +metric names, types, existing labels, and the HTTP endpoint are unchanged; +the new gauge is additive. The default HLL size changes from 16,384 registers +(`bucket_bits=14`, about 0.8% standard error) to 65,536 registers +(`bucket_bits=16`, about 0.4%). Three default windows use about 192 KiB of +register storage; sliding slots make the live tracker a few MiB per server. +The setting remains configurable with `--metric-hll-bucket-bits`. + ### Save Metrics (GPU → CPU) - **pegaflow_save_bytes_total** (Counter) - Total bytes saved from GPU to CPU storage @@ -245,15 +264,17 @@ tier counters. - Only used when `--metrics-otel-endpoint` is set - `--metric-hll-windows`: Comma-separated HLL sliding windows for estimated - prefix reuse (default: `15m,1h,24h`) + prefix reuse (default: `15m,1h,1d`) - Supported units: `s`, `m`, `h`, `d` - Each configured duration becomes a canonical `window` label. For example, the default config exports `window="15m"`, `window="1h"`, and `window="1d"`. - Empty entries such as `15m,,1h` and duplicate durations such as `1h,60m` are rejected at startup. -- `--metric-hll-bucket-bits`: HLL bucket index bits (default: `14`) - - Higher values use more memory and lower estimation error. +- `--metric-hll-bucket-bits`: HLL bucket index bits (default: `16`) + - `2^16 = 65,536` registers per window and about 0.4% standard error. + - Higher values use more memory and lower estimation error; `18` remains + the supported maximum. **Example: Prometheus Metrics** ```bash @@ -484,7 +505,10 @@ sum by (le) ( rate(pegaflow_cache_residence_duration_seconds_bucket{reason="pressure"}[5m]) ) -# HLL estimated hit rate for the 1h window +# HLL estimated hit rate for the 1h window (preferred) +pegaflow_hll_estimated_hit_rate{window="1h"} + +# Backward-compatible derivation from the retained gauges 1 - ( pegaflow_hll_cardinality{window="1h"} / diff --git a/docs/p2p.md b/docs/p2p.md index 8e54656e..ab3c6734 100644 --- a/docs/p2p.md +++ b/docs/p2p.md @@ -132,9 +132,11 @@ P2P-related Prometheus metrics (on `:9091/metrics` by default): | Metric | Type | Description | |---|---|---| -| `pegaflow_rdma_fetch_total` | Counter | Total RDMA fetch operations | +| `pegaflow_rdma_fetch_total` | Counter | Total per-segment RDMA fetch operations | | `pegaflow_rdma_fetch_duration` | Histogram | RDMA fetch latency distribution | | `pegaflow_rdma_fetch_bytes` | Counter | Total bytes fetched via RDMA | +| `pegaflow_rdma_fetch_plan_segments` | Histogram | Planned segment count per executed RDMA fetch plan | +| `pegaflow_rdma_fetch_plan_completed_segments` | Histogram | Completed segment count before a plan stops | | `pegaflow_rdma_qps` | Gauge | Active RDMA queue pairs | | `pegaflow_transfer_lock_active` | UpDownCounter | Currently held transfer locks | | `pegaflow_transfer_lock_timeouts_total` | Counter | Transfer lock timeout events | diff --git a/examples/ipc_sender.py b/examples/ipc_sender.py index f0be521d..12e2111e 100644 --- a/examples/ipc_sender.py +++ b/examples/ipc_sender.py @@ -18,7 +18,6 @@ class CudaIPCWrapper: """Wrapper for CUDA IPC handle with tensor metadata.""" def __init__(self, tensor: torch.Tensor): - assert tensor.storage_offset() == 0, "Tensor must have zero storage offset" assert tensor.is_contiguous(), "Tensor must be contiguous" storage = tensor.untyped_storage() diff --git a/pegaflow-common/src/block.rs b/pegaflow-common/src/block.rs index 6ff795b6..e899ba53 100644 --- a/pegaflow-common/src/block.rs +++ b/pegaflow-common/src/block.rs @@ -23,3 +23,51 @@ impl BlockKey { (self.namespace.capacity() + self.hash.capacity() + 48) as u64 } } + +/// Encode a raw content hash with a hybrid-cache group id. +/// +/// Group 0 keeps the raw hash byte-for-byte so existing single-group caches +/// (and every current connector) stay bit-identical. Groups >= 1 append the +/// big-endian group id, which cannot collide with a raw content hash because +/// every real hash family is fixed-length. +pub fn group_hash(hash: &[u8], group_id: u32) -> Vec { + if group_id == 0 { + return hash.to_vec(); + } + let mut encoded = Vec::with_capacity(hash.len() + 4); + encoded.extend_from_slice(hash); + encoded.extend_from_slice(&group_id.to_be_bytes()); + encoded +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn group_zero_hash_is_bit_identical_to_raw() { + // Backward-compat contract: existing single-group deployments must not + // observe any key change. + let hash = b"sha256-content-hash"; + assert_eq!(group_hash(hash, 0), hash.to_vec()); + assert_eq!(group_hash(b"", 0), Vec::::new()); + } + + #[test] + fn nonzero_group_appends_big_endian_group_id() { + let hash = [0xAA, 0xBB]; + assert_eq!(group_hash(&hash, 1), vec![0xAA, 0xBB, 0, 0, 0, 1]); + assert_eq!(group_hash(&hash, 0x01020304), vec![0xAA, 0xBB, 1, 2, 3, 4]); + } + + #[test] + fn distinct_groups_do_not_share_keys() { + // The same content hash in two groups must be two different keys, and + // an encoded key must never equal a raw hash of any length (length + // differs, so this holds even across hash families). + let hash = [1, 2, 3, 4]; + assert_ne!(group_hash(&hash, 0), group_hash(&hash, 1)); + assert_ne!(group_hash(&hash, 1), group_hash(&hash, 2)); + assert_ne!(group_hash(&hash, 1).len(), hash.len()); + } +} diff --git a/pegaflow-common/src/hll.rs b/pegaflow-common/src/hll.rs index 924d4182..cc817b19 100644 --- a/pegaflow-common/src/hll.rs +++ b/pegaflow-common/src/hll.rs @@ -22,6 +22,7 @@ pub const MAX_BUCKET_BITS: u8 = 18; /// Input hashes are expected to have good uniformity (e.g. SHA-256). /// The full hash is used: top `bucket_bits` bits select the register, /// remaining bits are scanned for leading zeros. +#[derive(Debug)] pub struct HyperLogLog { registers: Vec, bucket_bits: u8, @@ -34,8 +35,8 @@ impl HyperLogLog { /// Create a new HyperLogLog with the given bucket bits. /// /// `bucket_bits` determines the number of buckets (2^bucket_bits) and estimation - /// accuracy (~1.04 / sqrt(2^bucket_bits)). 14 gives 16384 buckets - /// and ~0.8% standard error. + /// accuracy (~1.04 / sqrt(2^bucket_bits)). 16 gives 65536 buckets + /// and ~0.4% standard error. pub fn new(bucket_bits: u8) -> Self { assert!( (MIN_BUCKET_BITS..=MAX_BUCKET_BITS).contains(&bucket_bits), @@ -205,7 +206,7 @@ fn alpha_m(m: usize) -> f64 { /// Metric snapshot returned by [`HllTracker::metric`]. #[derive(Debug, Clone)] pub struct HllMetric { - /// Estimated number of distinct block hashes in the window. + /// Estimated number of distinct miss block identities in the window. pub cardinality: f64, /// Total block requests (including duplicates) in the window. pub total_requests: u64, @@ -225,7 +226,7 @@ struct WindowSlot { /// /// Divides time into fixed-duration slots and maintains a ring of HLLs. /// The merged cardinality across all active slots approximates the number -/// of distinct blocks requested in the window. From this we derive: +/// of distinct miss blocks recorded in the window. From this we derive: /// /// ```text /// hit_rate = (total_requests - cardinality) / total_requests @@ -249,7 +250,7 @@ impl HllTracker { /// /// - `slot_duration`: how long each time slot lasts (e.g. 1 hour) /// - `window_duration`: total sliding window (e.g. 24 hours) - /// - `bucket_bits`: HLL bucket index bits (4..=18, default 14) + /// - `bucket_bits`: HLL bucket index bits (4..=18, default 16) pub fn new(slot_duration: Duration, window_duration: Duration, bucket_bits: u8) -> Self { Self { slots: VecDeque::new(), @@ -268,6 +269,20 @@ impl HllTracker { /// time drift. For example with 1h slots: if the first slot starts at 0:00 /// and the next request arrives at 1:30, the new slot starts at 1:00 (not 1:30). pub fn record(&mut self, hash: &[u8]) { + let hashes = [hash]; + self.record_hashes_with_total(&hashes, 1); + } + + /// Record a batch of distinct identities while accounting for a possibly + /// larger observation count. The identities are inserted into HLL, while + /// `total_requests` is used as the denominator. This is used by the + /// miss-only reference: only cache misses enter HLL, but every queried + /// block still contributes to the total observation count. + pub fn record_hashes_with_total>(&mut self, hashes: &[T], total_requests: u64) { + if total_requests == 0 { + return; + } + let now = Instant::now(); let need_new_slot = match self.slots.back() { @@ -295,15 +310,15 @@ impl HllTracker { } let slot = self.slots.back_mut().unwrap(); - slot.hll.insert(hash); - slot.request_count += 1; + for hash in hashes { + slot.hll.insert(hash.as_ref()); + } + slot.request_count += total_requests; } /// Record a batch of block hashes from a gRPC request. pub fn record_hashes(&mut self, hashes: &[Vec]) { - for hash in hashes { - self.record(hash); - } + self.record_hashes_with_total(hashes, hashes.len() as u64); } /// Compute and return the current metric snapshot. @@ -417,6 +432,30 @@ impl MultiWindowHllTracker { } } + /// Record all queried blocks in the denominator but insert only the + /// identities that were misses into HLL. Repeated misses remain deduped by + /// HLL, preserving the infinite-cache reuse reference without discarding + /// historical observations. + pub fn record_namespaced_misses( + &mut self, + namespace: &str, + total_requests: u64, + miss_hashes: &[Vec], + ) { + if total_requests == 0 { + return; + } + debug_assert!(miss_hashes.len() as u64 <= total_requests); + + let namespaced_hashes: Vec<[u8; 8]> = miss_hashes + .iter() + .map(|hash| namespaced_hash(namespace, hash)) + .collect(); + for (_, tracker) in &mut self.windows { + tracker.record_hashes_with_total(&namespaced_hashes, total_requests); + } + } + /// Snapshot every window. Returned in insertion order. pub fn metrics(&mut self) -> Vec<(String, HllMetric)> { self.windows @@ -426,6 +465,38 @@ impl MultiWindowHllTracker { } } +/// Stable identity hash for `(namespace, block_hash)`. +/// +/// Cluster aggregation unions raw HLL registers across nodes, so every node +/// must map the same object to the exact same bits regardless of platform, +/// architecture, or build. FNV-1a with a splitmix64 finalizer is fully +/// specified byte-by-byte; do not replace it with a hasher that does not +/// guarantee cross-platform stability (e.g. ahash, SipHash with random keys). +fn namespaced_hash(namespace: &str, block_hash: &[u8]) -> [u8; 8] { + const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325; + const FNV_PRIME: u64 = 0x0000_0100_0000_01b3; + let mut h = FNV_OFFSET; + // Length prefix keeps (namespace, hash) boundaries unambiguous. + for &byte in (namespace.len() as u64) + .to_le_bytes() + .iter() + .chain(namespace.as_bytes()) + .chain(block_hash) + { + h = (h ^ u64::from(byte)).wrapping_mul(FNV_PRIME); + } + splitmix64(h).to_be_bytes() +} + +/// splitmix64 finalizer: strengthens FNV-1a's avalanche so the top bits used +/// for HLL bucket selection are uniformly distributed. +fn splitmix64(mut x: u64) -> u64 { + x = x.wrapping_add(0x9e37_79b9_7f4a_7c15); + x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + x ^ (x >> 31) +} + fn derive_slot_duration(window: Duration) -> Duration { const MIN_SLOT: Duration = Duration::from_secs(60); const MAX_SLOT: Duration = Duration::from_secs(3600); @@ -766,7 +837,7 @@ mod tests { vec![ ("15m".into(), Duration::from_secs(15 * 60)), ("1h".into(), Duration::from_secs(3600)), - ("24h".into(), Duration::from_secs(86400)), + ("1d".into(), Duration::from_secs(86400)), ], 14, ); @@ -781,7 +852,7 @@ mod tests { assert_eq!(metrics.len(), 3); assert_eq!(metrics[0].0, "15m"); assert_eq!(metrics[1].0, "1h"); - assert_eq!(metrics[2].0, "24h"); + assert_eq!(metrics[2].0, "1d"); for (label, m) in metrics { assert_eq!(m.total_requests, 2000, "{label}: total"); assert!( @@ -792,6 +863,44 @@ mod tests { } } + #[test] + fn namespaced_misses_keep_all_observations_in_denominator() { + let mut tracker = + MultiWindowHllTracker::new(vec![("15m".into(), Duration::from_secs(15 * 60))], 14); + let miss = vec![sha256_like(1).to_vec(), sha256_like(2).to_vec()]; + tracker.record_namespaced_misses("model", 5, &miss); + + let metric = tracker.metrics().remove(0).1; + assert_eq!(metric.total_requests, 5); + assert!(metric.cardinality > 1.0 && metric.cardinality < 3.5); + assert!(metric.estimated_hit_rate > 0.3); + } + + #[test] + fn namespaced_misses_allow_all_hit_observation_without_hll_insert() { + let mut tracker = + MultiWindowHllTracker::new(vec![("15m".into(), Duration::from_secs(15 * 60))], 14); + tracker.record_namespaced_misses("model", 4, &[]); + + let metric = tracker.metrics().remove(0).1; + assert_eq!(metric.total_requests, 4); + assert_eq!(metric.cardinality, 0.0); + assert_eq!(metric.estimated_hit_rate, 1.0); + } + + #[test] + fn namespaced_misses_separate_equal_raw_hashes() { + let mut tracker = + MultiWindowHllTracker::new(vec![("15m".into(), Duration::from_secs(15 * 60))], 16); + let raw_hash = vec![42; 32]; + tracker.record_namespaced_misses("model-a", 1, std::slice::from_ref(&raw_hash)); + tracker.record_namespaced_misses("model-b", 1, std::slice::from_ref(&raw_hash)); + + let metric = tracker.metrics().remove(0).1; + assert_eq!(metric.total_requests, 2); + assert!((1.5..2.5).contains(&metric.cardinality)); + } + #[test] fn derive_slot_clamps() { assert_eq!( @@ -849,11 +958,4 @@ mod tests { hash[24..32].copy_from_slice(&m3.to_le_bytes()); hash } - - fn splitmix64(mut x: u64) -> u64 { - x = x.wrapping_add(0x9e3779b97f4a7c15); - x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); - x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb); - x ^ (x >> 31) - } } diff --git a/pegaflow-common/src/lib.rs b/pegaflow-common/src/lib.rs index abe43f30..2438e444 100644 --- a/pegaflow-common/src/lib.rs +++ b/pegaflow-common/src/lib.rs @@ -5,7 +5,7 @@ pub mod logging; #[cfg(target_os = "linux")] pub mod numa; -pub use block::BlockKey; +pub use block::{BlockKey, group_hash}; #[cfg(target_os = "linux")] pub use numa::{ NumaNode, NumaTopology, format_cpu_list, pin_thread_to_numa_node, query_pages_numa, diff --git a/pegaflow-core/src/backing/rdma_fetch.rs b/pegaflow-core/src/backing/rdma_fetch.rs index aacd92df..88e801ed 100644 --- a/pegaflow-core/src/backing/rdma_fetch.rs +++ b/pegaflow-core/src/backing/rdma_fetch.rs @@ -10,8 +10,8 @@ use log::{debug, info, warn}; use mea::singleflight::Group; use pegaflow_proto::proto::engine::engine_client::EngineClient; use pegaflow_proto::proto::engine::{ - QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, RdmaHandshakeRequest, - TransferBlockInfo, + FetchSegment, QueryBlocksForTransferRequest, QueryBlocksForTransferResponse, + RdmaHandshakeRequest, TransferBlockInfo, }; use pegaflow_transfer::{ConnectionStatus, HandshakeMetadata, TransferDesc, TransferOp}; use tonic::transport::{Channel, Endpoint}; @@ -58,6 +58,138 @@ pub(crate) struct RdmaFetchStore { connect_group: Arc>, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct FetchPlanSegment { + node: String, + start: usize, + end: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct FetchPlan { + segments: Vec, + block_count: usize, +} + +impl FetchPlan { + pub(crate) fn block_count(&self) -> usize { + self.block_count + } + + fn segment_blocks_summary(&self) -> String { + self.segments + .iter() + .map(|segment| (segment.end - segment.start).to_string()) + .collect::>() + .join(",") + } +} + +fn validate_fetch_plan( + segments: Vec, + hash_count: usize, + exclude_node: &str, +) -> Result, String> { + if segments.is_empty() { + return Ok(None); + } + + let mut validated = Vec::with_capacity(segments.len()); + let mut offset = 0usize; + for (index, segment) in segments.into_iter().enumerate() { + if segment.node.is_empty() { + return Err(format!("segment {index} has an empty node")); + } + if segment.node == exclude_node { + return Err(format!("segment {index} selects the excluded requester")); + } + if segment.block_count == 0 { + return Err(format!("segment {index} has zero blocks")); + } + if validated + .last() + .is_some_and(|previous: &FetchPlanSegment| previous.node == segment.node) + { + return Err(format!( + "segment {index} repeats the previous node instead of merging" + )); + } + + let count = usize::try_from(segment.block_count) + .map_err(|_| format!("segment {index} block count exceeds usize"))?; + let end = offset + .checked_add(count) + .ok_or_else(|| format!("segment {index} block count overflows"))?; + if end > hash_count { + return Err(format!( + "segment {index} ends at block {end}, beyond request length {hash_count}" + )); + } + validated.push(FetchPlanSegment { + node: segment.node, + start: offset, + end, + }); + offset = end; + } + + Ok(Some(FetchPlan { + segments: validated, + block_count: offset, + })) +} + +#[tonic::async_trait] +trait SegmentFetcher { + async fn fetch_segment(&self, remote_addr: &str, hashes: &[Vec]) -> PrefetchResult; +} + +struct RdmaSegmentFetcher<'a> { + store: &'a RdmaFetchStore, + req_id: &'a str, + namespace: &'a str, +} + +#[tonic::async_trait] +impl SegmentFetcher for RdmaSegmentFetcher<'_> { + async fn fetch_segment(&self, remote_addr: &str, hashes: &[Vec]) -> PrefetchResult { + self.store + .fetch_blocks(remote_addr, self.req_id, self.namespace, hashes) + .await + } +} + +async fn execute_fetch_plan( + fetcher: &F, + plan: &FetchPlan, + namespace: &str, + hashes: &[Vec], +) -> (PrefetchResult, usize, Option<(usize, usize, usize)>) { + let mut fetched = Vec::with_capacity(plan.block_count); + let mut completed_segments = 0usize; + let mut failed_segment = None; + + for (index, segment) in plan.segments.iter().enumerate() { + let expected = &hashes[segment.start..segment.end]; + let returned = fetcher.fetch_segment(&segment.node, expected).await; + let contiguous = returned + .iter() + .zip(expected) + .take_while(|((key, _), hash)| key.namespace == namespace && key.hash == **hash) + .count(); + let returned_count = returned.len(); + fetched.extend(returned.into_iter().take(contiguous)); + + if contiguous != expected.len() || returned_count != expected.len() { + failed_segment = Some((index, expected.len(), contiguous)); + break; + } + completed_segments += 1; + } + + (fetched, completed_segments, failed_segment) +} + impl RdmaFetchStore { pub(crate) fn new( metaserver_client: Arc, @@ -76,42 +208,89 @@ impl RdmaFetchStore { } } - /// Query MetaServer for the best remote node that holds a prefix of `hashes`. - /// Returns `(node_addr, prefix_len)`, or `None` if no remote node has any. - pub(crate) async fn query_prefix( + /// Query MetaServer for a validated ordered plan covering a prefix of `hashes`. + pub(crate) async fn query_plan( &self, namespace: &str, hashes: &[Vec], - ) -> Option<(String, usize)> { + ) -> Option { if hashes.is_empty() { return None; } - let nodes = match self.metaserver_client.query_prefix(namespace, hashes).await { - Ok(n) => n, + let segments = match self + .metaserver_client + .query_plan(namespace, hashes, &self.advertise_addr) + .await + { + Ok(segments) => segments, Err(e) => { warn!("MetaServer query failed for remote fetch: {e}"); return None; } }; - let best = nodes - .iter() - .filter(|n| n.node != self.advertise_addr) - .max_by_key(|n| n.prefix_len)?; - - let prefix_len = best.prefix_len as usize; - if prefix_len == 0 { - return None; - } + let plan = match validate_fetch_plan(segments, hashes.len(), &self.advertise_addr) { + Ok(plan) => plan?, + Err(error) => { + warn!("MetaServer returned invalid remote fetch plan: {error}"); + return None; + } + }; debug!( - "Remote prefix query: namespace={namespace} best_node={} prefix={prefix_len}/{}", - best.node, - hashes.len() + "Remote prefix query: segments={} prefix={}/{}", + plan.segments.len(), + plan.block_count, + hashes.len(), + ); + + Some(plan) + } + + pub(crate) async fn fetch_plan( + &self, + plan: &FetchPlan, + req_id: &str, + namespace: &str, + hashes: &[Vec], + ) -> PrefetchResult { + let started_at = Instant::now(); + let fetcher = RdmaSegmentFetcher { + store: self, + req_id, + namespace, + }; + let (fetched, completed_segments, failure) = + execute_fetch_plan(&fetcher, plan, namespace, hashes).await; + let metrics = core_metrics(); + metrics + .rdma_fetch_plan_segments + .record(plan.segments.len() as u64, &[]); + metrics + .rdma_fetch_plan_completed_segments + .record(completed_segments as u64, &[]); + let (failed_segment, failed_planned_blocks, failed_returned_blocks) = failure + .map(|(index, planned, returned)| { + (index.to_string(), planned.to_string(), returned.to_string()) + }) + .unwrap_or_else(|| ("none".into(), "none".into(), "none".into())); + + info!( + "RDMA multi-node fetch plan summary: req_id={} planned_segments={} completed_segments={} planned_blocks={} segment_blocks={} fetched_blocks={} failed_segment={} failed_segment_planned_blocks={} failed_segment_returned_blocks={} total_ms={:.2}", + req_id, + plan.segments.len(), + completed_segments, + plan.block_count, + plan.segment_blocks_summary(), + fetched.len(), + failed_segment, + failed_planned_blocks, + failed_returned_blocks, + started_at.elapsed().as_secs_f64() * 1000.0, ); - Some((best.node.clone(), prefix_len)) + fetched } /// Fetch `hashes` from `remote_addr`. @@ -735,6 +914,7 @@ fn transfer_timeout_from_server(lock_timeout_secs: u32) -> Duration { #[cfg(test)] mod tests { use super::*; + use std::collections::VecDeque; use std::num::NonZeroU64; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -757,6 +937,153 @@ mod tests { HashMap::from([(NumaNode(0), bytes)]) } + fn segment(node: &str, block_count: u32) -> FetchSegment { + FetchSegment { + node: node.to_string(), + block_count, + } + } + + fn fetched_block(hash: u8) -> (BlockKey, Arc) { + ( + BlockKey::new("ns".to_string(), vec![hash]), + Arc::new(SealedBlock::from_slots(Vec::new())), + ) + } + + #[derive(Default)] + struct FakeSegmentFetcher { + calls: Mutex>)>>, + responses: Mutex>, + } + + #[tonic::async_trait] + impl SegmentFetcher for FakeSegmentFetcher { + async fn fetch_segment(&self, remote_addr: &str, hashes: &[Vec]) -> PrefetchResult { + self.calls + .lock() + .unwrap() + .push((remote_addr.to_string(), hashes.to_vec())); + self.responses + .lock() + .unwrap() + .pop_front() + .unwrap_or_default() + } + } + + #[test] + fn validates_ordered_fetch_plan_offsets() { + let plan = validate_fetch_plan( + vec![segment("node-a", 2), segment("node-b", 1)], + 3, + "requester", + ) + .expect("plan should be valid") + .expect("plan should be non-empty"); + + assert_eq!(plan.block_count, 3); + assert_eq!(plan.segment_blocks_summary(), "2,1"); + assert_eq!( + plan.segments, + vec![ + FetchPlanSegment { + node: "node-a".into(), + start: 0, + end: 2, + }, + FetchPlanSegment { + node: "node-b".into(), + start: 2, + end: 3, + }, + ] + ); + } + + #[test] + fn rejects_invalid_fetch_plans() { + for (segments, expected) in [ + (vec![segment("", 1)], "empty node"), + (vec![segment("node-a", 0)], "zero blocks"), + (vec![segment("requester", 1)], "excluded requester"), + (vec![segment("node-a", 2)], "beyond request length"), + ( + vec![segment("node-a", 1), segment("node-a", 1)], + "repeats the previous node", + ), + ] { + let error = + validate_fetch_plan(segments, 1, "requester").expect_err("plan should be rejected"); + assert!(error.contains(expected), "unexpected error: {error}"); + } + } + + #[tokio::test] + async fn fetch_plan_executes_segments_in_order() { + let plan = validate_fetch_plan( + vec![segment("node-a", 2), segment("node-b", 1)], + 3, + "requester", + ) + .unwrap() + .unwrap(); + let fetcher = FakeSegmentFetcher { + calls: Mutex::new(Vec::new()), + responses: Mutex::new(VecDeque::from([ + vec![fetched_block(1), fetched_block(2)], + vec![fetched_block(3)], + ])), + }; + let hashes = vec![vec![1], vec![2], vec![3]]; + + let (fetched, completed, failure) = + execute_fetch_plan(&fetcher, &plan, "ns", &hashes).await; + + assert_eq!(fetched.len(), 3); + assert_eq!(completed, 2); + assert_eq!(failure, None); + assert_eq!( + *fetcher.calls.lock().unwrap(), + vec![ + ("node-a".into(), vec![vec![1], vec![2]]), + ("node-b".into(), vec![vec![3]]), + ] + ); + } + + #[tokio::test] + async fn fetch_plan_stops_after_first_short_segment() { + let plan = validate_fetch_plan( + vec![ + segment("node-a", 1), + segment("node-b", 1), + segment("node-c", 1), + ], + 3, + "requester", + ) + .unwrap() + .unwrap(); + let fetcher = FakeSegmentFetcher { + calls: Mutex::new(Vec::new()), + responses: Mutex::new(VecDeque::from([ + vec![fetched_block(1)], + Vec::new(), + vec![fetched_block(3)], + ])), + }; + let hashes = vec![vec![1], vec![2], vec![3]]; + + let (fetched, completed, failure) = + execute_fetch_plan(&fetcher, &plan, "ns", &hashes).await; + + assert_eq!(fetched.len(), 1); + assert_eq!(completed, 1); + assert_eq!(failure, Some((1, 1, 0))); + assert_eq!(fetcher.calls.lock().unwrap().len(), 2); + } + #[test] fn chunked_slabs_bump_within_chunk_then_refill() { let calls = Arc::new(AtomicUsize::new(0)); diff --git a/pegaflow-core/src/backing/ssd_cache.rs b/pegaflow-core/src/backing/ssd_cache.rs index db99ff8c..d84c4892 100644 --- a/pegaflow-core/src/backing/ssd_cache.rs +++ b/pegaflow-core/src/backing/ssd_cache.rs @@ -37,6 +37,11 @@ pub const DEFAULT_SSD_WRITE_INFLIGHT: usize = 2; /// Default max concurrent prefetches pub const DEFAULT_SSD_PREFETCH_INFLIGHT: usize = 16; +/// Upper bound for one pinned-pool allocation while staging an SSD prefetch. +/// Large contiguous requests can otherwise force disproportionate LRU reclaim +/// from a fragmented pool. +const SSD_PREFETCH_CHUNK_BYTES: u64 = 256 * 1024 * 1024; + /// Result of a single prefetch I/O. type SinglePrefetchResult = ( BlockKey, @@ -700,10 +705,21 @@ struct SlotRef { size: u64, } +struct SlotPlacement { + block_idx: usize, + slot_idx: usize, + offset: u64, +} + +struct PrefetchChunk { + capacity: u64, + size: u64, + slots: Vec, +} + /// Group all slots across all blocks by NUMA node. /// -/// Returns a map from NUMA key (None = global/unknown) to the list of slot -/// references that should share a single contiguous allocation. +/// Returns a map from NUMA key (None = global/unknown) to its slot references. fn group_slots_by_numa( is_numa: bool, requests: &[PrefetchRequest], @@ -727,6 +743,59 @@ fn group_slots_by_numa( groups } +fn chunk_slot_refs(refs: &[SlotRef], chunk_bytes: u64) -> Result, String> { + assert!(chunk_bytes > 0, "SSD prefetch chunk size must be non-zero"); + + let mut remaining = refs.iter().try_fold(0u64, |total, slot| { + total + .checked_add(slot.size) + .ok_or_else(|| "SSD prefetch allocation size overflow".to_string()) + })?; + let mut chunks = Vec::new(); + let mut current: Option = None; + + for slot in refs { + let needs_new_chunk = current.as_ref().is_none_or(|chunk| { + chunk + .size + .checked_add(slot.size) + .is_none_or(|end| end > chunk.capacity) + }); + + if needs_new_chunk { + if let Some(chunk) = current.take() { + chunks.push(chunk); + } + current = Some(PrefetchChunk { + capacity: remaining.min(chunk_bytes).max(slot.size), + size: 0, + slots: Vec::new(), + }); + } + + let chunk = current + .as_mut() + .expect("non-empty slot list must have an active prefetch chunk"); + chunk.slots.push(SlotPlacement { + block_idx: slot.block_idx, + slot_idx: slot.slot_idx, + offset: chunk.size, + }); + chunk.size = chunk + .size + .checked_add(slot.size) + .expect("new SSD prefetch chunk must fit its first slot"); + remaining = remaining + .checked_sub(slot.size) + .expect("SSD prefetch remaining bytes must cover every slot"); + } + + if let Some(chunk) = current { + chunks.push(chunk); + } + Ok(chunks) +} + /// Allocate per-slot memory grouped by NUMA, build PrefetchTasks, and enqueue. /// Returns false if the task channel is closed (should exit). async fn dispatch_prefetch_batch( @@ -739,37 +808,43 @@ async fn dispatch_prefetch_batch( // 1. Group all slots across all blocks by NUMA node let numa_groups = group_slots_by_numa(store.is_numa(), &requests); - // 2. Allocate per NUMA group, assign per-slot (allocation, offset). - // All blocks in a batch share the same slot layout, so a single NUMA - // group failure means every block is missing slots → fail the whole batch. + // 2. Allocate bounded chunks per NUMA group, assign per-slot offsets. + // A failure still fails the whole batch because every requested block + // must have all of its slots before it can be rebuilt. let mut slot_allocs: Vec>> = requests .iter() .map(|r| (0..r.entry.slots.len()).map(|_| None).collect()) .collect(); for (numa_node, refs) in &numa_groups { - let total_size: u64 = refs.iter().map(|r| r.size).sum(); - let allocation = match store.allocate_prefetch(total_size, *numa_node) { - Some(alloc) => alloc, - None => { - warn!( - "SSD prefetch dispatcher: alloc failed for {} bytes ({} slots) numa={:?}, failing entire batch", - total_size, - refs.len(), - numa_node - ); + let chunks = match chunk_slot_refs(refs, SSD_PREFETCH_CHUNK_BYTES) { + Ok(chunks) => chunks, + Err(err) => { + warn!("SSD prefetch dispatcher: {err}, failing entire batch"); let _ = done_tx.send(Vec::new()); return true; } }; - - let mut offset = 0usize; - for r in refs { - slot_allocs[r.block_idx][r.slot_idx] = Some(SlotAlloc { - allocation: allocation.clone(), - offset, - }); - offset += r.size as usize; + for chunk in chunks { + let allocation = match store.allocate_prefetch(chunk.size, *numa_node) { + Some(alloc) => alloc, + None => { + warn!( + "SSD prefetch dispatcher: alloc failed for {} bytes numa={:?}, failing entire batch", + chunk.size, numa_node + ); + let _ = done_tx.send(Vec::new()); + return true; + } + }; + for slot in chunk.slots { + let offset = + usize::try_from(slot.offset).expect("SSD prefetch chunk offset must fit usize"); + slot_allocs[slot.block_idx][slot.slot_idx] = Some(SlotAlloc { + allocation: Arc::clone(&allocation), + offset, + }); + } } } @@ -1294,4 +1369,77 @@ mod tests { assert_eq!(groups[&Some(NumaNode(0))].len(), 1); assert_eq!(groups[&None].len(), 1); // UNKNOWN → None } + + #[test] + fn test_k3_style_prefetch_uses_bounded_numa_chunks() { + const MIB: u64 = 1024 * 1024; + let slots = || { + (0..8) + .map(|slot| make_slot(NumaNode(slot / 4), 16 * MIB)) + .collect() + }; + let requests: Vec<_> = (1..=56) + .map(|block| make_prefetch_request(block, slots())) + .collect(); + + let groups = group_slots_by_numa(true, &requests); + for numa in [NumaNode(0), NumaNode(1)] { + let refs = &groups[&Some(numa)]; + let whole_batch_bytes: u64 = refs.iter().map(|slot| slot.size).sum(); + let chunks = chunk_slot_refs(refs, SSD_PREFETCH_CHUNK_BYTES).unwrap(); + + assert_eq!(whole_batch_bytes, 3584 * MIB); + assert_eq!(chunks.len(), 14); + assert_eq!( + chunks.iter().map(|chunk| chunk.size).max(), + Some(SSD_PREFETCH_CHUNK_BYTES) + ); + assert_eq!( + chunks.iter().map(|chunk| chunk.slots.len()).sum::(), + refs.len() + ); + } + } + + #[test] + fn test_oversized_prefetch_slot_gets_dedicated_chunk() { + let refs = vec![ + SlotRef { + block_idx: 0, + slot_idx: 0, + size: 300, + }, + SlotRef { + block_idx: 0, + slot_idx: 1, + size: 100, + }, + ]; + + let chunks = chunk_slot_refs(&refs, 256).unwrap(); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].size, 300); + assert_eq!(chunks[0].slots[0].offset, 0); + assert_eq!(chunks[1].size, 100); + assert_eq!(chunks[1].slots[0].offset, 0); + } + + #[test] + fn test_non_divisible_prefetch_chunks_allocate_only_slot_bytes() { + let refs = [200, 100, 100] + .into_iter() + .enumerate() + .map(|(slot_idx, size)| SlotRef { + block_idx: 0, + slot_idx, + size, + }) + .collect::>(); + + let chunks = chunk_slot_refs(&refs, 256).unwrap(); + let allocation_sizes = chunks.iter().map(|chunk| chunk.size).collect::>(); + + assert_eq!(allocation_sizes, vec![200, 200]); + assert_eq!(allocation_sizes.iter().sum::(), 400); + } } diff --git a/pegaflow-core/src/cache.rs b/pegaflow-core/src/cache.rs index 90e209e7..55620006 100644 --- a/pegaflow-core/src/cache.rs +++ b/pegaflow-core/src/cache.rs @@ -72,6 +72,16 @@ impl TinyLfuCache { self.lru.contains_key(key) } + /// Returns true when the cache is the only strong owner of the block. + /// + /// Weak references from fire-and-forget backing-store work do not pin the + /// resident block: they may fail to upgrade after pressure eviction. + pub(crate) fn is_cache_owned_only(&self, key: &BlockKey) -> bool { + self.lru + .peek(key) + .is_some_and(|block| Arc::strong_count(block) == 1) + } + /// Insert with TinyLFU admission. If the candidate is colder than the /// current LRU victim it is dropped. pub(crate) fn insert(&mut self, key: BlockKey, value: ArcSealedBlock) -> CacheInsertOutcome { diff --git a/pegaflow-core/src/instance.rs b/pegaflow-core/src/instance.rs index 263d0462..52c6da26 100644 --- a/pegaflow-core/src/instance.rs +++ b/pegaflow-core/src/instance.rs @@ -49,11 +49,25 @@ struct RegistrationState { /// /// Layer ids are the rank of the layer name in sorted order, so every /// instance that registers the same layer set derives the same ids — the -/// property block slot layout depends on. +/// property block slot layout depends on. On top of that, each layer belongs +/// to a hybrid-cache *storage group* (attention vs. recurrent state, mirroring +/// vLLM's `KVCacheGroupSpec`): slots are dense *within* a group, giving each +/// group its own seal domain. A group seals a block once every +/// `(group layer, tp_rank)` slot of that block is saved, regardless of other +/// groups — that is what lets a recurrent-state group seal the final block +/// even though attention groups save every block. #[derive(Debug)] pub(crate) struct LayerTopology { name_to_id: HashMap, tp_size: usize, + /// Storage group id per layer, indexed by layer_id. All zeros for + /// single-group (classic) instances. + layer_group: Vec, + /// Layer count per storage group, indexed by group id. + group_layer_count: Vec, + /// Dense rank of the layer within its storage group (0..group_layer_count), + /// indexed by layer_id. Within-group slots are `rank * tp_size + tp_rank`. + layer_group_rank: Vec, /// Page-first layout when `Some`: each block's layers collapse into /// contiguous per-shard pages, so `total_slots = num_shards`. `None` is the /// legacy layer-first layout (`total_slots = num_layers * tp_size`). @@ -167,6 +181,11 @@ impl LayerTopology { /// Total number of storage slots per block: `num_shards` page-first, else /// `num_layers * tp_size` (layer-first). + /// + /// This is the union across groups and remains the right answer for + /// single-group instances; multi-group callers must use + /// [`Self::group_total_slots`] — blocks seal per group, so one global + /// denominator would never fill. pub(crate) fn total_slots(&self) -> usize { match &self.page_layout { Some(p) => p.shard_page_sizes.len(), @@ -174,6 +193,34 @@ impl LayerTopology { } } + /// Number of distinct storage groups. Always >= 1. + pub(crate) fn num_groups(&self) -> usize { + self.group_layer_count.len() + } + + /// Storage group id of a layer. + pub(crate) fn group_of_layer(&self, layer_id: usize) -> u32 { + self.layer_group[layer_id] + } + + /// Number of storage slots per block within one group. A block of that + /// group seals when exactly this many slots are saved. + pub(crate) fn group_total_slots(&self, group_id: u32) -> Result { + let group_idx = group_id as usize; + if group_idx >= self.num_groups() { + return Err(EngineError::InvalidArgument(format!( + "storage group {group_id} out of range ({} groups)", + self.num_groups() + ))); + } + Ok(match &self.page_layout { + // Page-first is seal-validated to be single-group: slots follow + // shards, not the layer grid. + Some(p) => p.shard_page_sizes.len(), + None => self.group_layer_count[group_idx] * self.tp_size, + }) + } + /// Page-first only: contiguous page size in bytes of `shard` (one slot). pub(crate) fn shard_page_size(&self, shard: usize) -> Option { self.page_layout.as_ref().map(|p| p.shard_page_sizes[shard]) @@ -198,7 +245,9 @@ impl LayerTopology { /// /// Page-first collapses the layer dimension into per-shard pages, so the /// slot is the layer's shard; the layer's position is its page offset - /// ([`Self::page_placement`]). Layer-first keeps `[layer][tp_rank]`. + /// ([`Self::page_placement`]). Layer-first slots are dense *within* the + /// layer's storage group: `group_rank * tp_size + tp_rank`. Single-group + /// topologies degenerate to the historical `[layer][tp_rank]` grid. pub(crate) fn slot_index(&self, layer_id: usize, tp_rank: usize) -> Result { if layer_id >= self.num_layers() { return Err(EngineError::InvalidArgument(format!( @@ -215,7 +264,7 @@ impl LayerTopology { } Ok(match &self.page_layout { Some(p) => p.layer_shard[layer_id], - None => layer_id * self.tp_size + tp_rank, + None => self.layer_group_rank[layer_id] * self.tp_size + tp_rank, }) } } @@ -243,6 +292,9 @@ pub struct GpuContext { /// KV cache layouts by layer name. kv_caches: HashMap, + /// Hybrid-cache storage group id by layer name; absent = group 0. + layer_groups: HashMap, + /// CUDA context handle (kept alive for the lifetime of this context). _cuda_ctx: Arc, @@ -260,6 +312,10 @@ impl GpuContext { /// # Errors /// Returns `EngineError::CudaInit` if CUDA context creation or worker /// pool initialization fails. + #[allow( + clippy::too_many_arguments, + reason = "GPU context construction mirrors one registration payload" + )] fn new( cuda_ctx: Arc, device_id: i32, @@ -268,6 +324,7 @@ impl GpuContext { numa_node: NumaNode, transfer_mode: TransferMode, kv_caches: HashMap, + layer_groups: HashMap, ) -> Result { let worker_pool = GpuWorkerPool::spawn(device_id, numa_node, transfer_mode)?; @@ -277,6 +334,7 @@ impl GpuContext { pp_rank, preferred_numa: numa_node, kv_caches, + layer_groups, _cuda_ctx: cuda_ctx, worker_pool, }) @@ -311,6 +369,12 @@ impl GpuContext { pub(crate) fn worker_pool(&self) -> &GpuWorkerPool { &self.worker_pool } + + /// Hybrid-cache storage group of a layer; unregistered layers default to + /// group 0, preserving single-group behavior. + pub(crate) fn group_of_layer(&self, layer_name: &str) -> u32 { + self.layer_groups.get(layer_name).copied().unwrap_or(0) + } } pub(crate) struct GpuRegistration { @@ -320,6 +384,8 @@ pub(crate) struct GpuRegistration { pub(crate) numa_node: NumaNode, pub(crate) transfer_mode: TransferMode, pub(crate) kv_caches: HashMap, + /// Hybrid-cache storage group id by layer name; absent = group 0. + pub(crate) layer_groups: HashMap, } /// Instance context for a model inference process. @@ -483,6 +549,61 @@ impl InstanceContext { .map(|(id, name)| (name.clone(), id)) .collect(); + // Storage groups must agree across devices (same name ⇒ same group, + // checked while scanning) and be dense `0..N` so slot spaces are + // complete; a group nobody registers would never seal. + let mut group_by_name: HashMap<&str, u32> = HashMap::new(); + for gpu in gpus() { + for name in gpu.kv_caches.keys() { + let group = gpu.group_of_layer(name); + match group_by_name.insert(name, group) { + None => {} + Some(existing) if existing == group => {} + Some(existing) => { + return Err(EngineError::InvalidArgument(format!( + "layer {name} registered in storage group {existing} on one device \ + but group {group} on device {}", + gpu.device_id(), + ))); + } + } + } + } + let num_groups = 1 + group_by_name.values().copied().max().unwrap_or(0) as usize; + if group_by_name + .values() + .copied() + .collect::>() + .len() + != num_groups + { + return Err(EngineError::InvalidArgument(format!( + "storage groups must be dense 0..{num_groups}: an unregistered group \ + would never seal its blocks", + ))); + } + + let mut layer_group = vec![0u32; names.len()]; + let mut layer_group_rank = vec![0usize; names.len()]; + let mut group_layer_count = vec![0usize; num_groups]; + for name in &names { + let layer_id = name_to_id[name]; + let group = group_by_name[name.as_str()]; + layer_group[layer_id] = group; + layer_group_rank[layer_id] = group_layer_count[group as usize]; + group_layer_count[group as usize] += 1; + } + + // Page-first packs every layer of a block into one page slot per + // shard; splitting a block across storage groups contradicts that + // layout, so reject instead of silently picking one. + if self.page_first && num_groups > 1 { + return Err(EngineError::InvalidArgument(format!( + "page-first storage does not support multiple storage groups \ + ({num_groups} registered)", + ))); + } + // Validate registration completeness on the full `[layer][tp_rank]` // grid, independent of how slots are stored: every (layer, tp_rank) // pair needs an owner regardless of page-first collapsing. @@ -543,6 +664,9 @@ impl InstanceContext { Ok(LayerTopology { name_to_id, tp_size: self.tp_size, + layer_group, + group_layer_count, + layer_group_rank, page_layout, }) } @@ -555,6 +679,10 @@ impl InstanceContext { /// # Errors /// Returns `EngineError::InvalidArgument` for negative device IDs, /// or `EngineError::CudaInit` if CUDA context creation fails. + #[allow( + clippy::too_many_arguments, + reason = "GPU context construction mirrors one registration payload" + )] fn build_gpu_context( &self, device_id: i32, @@ -563,6 +691,7 @@ impl InstanceContext { numa_node: NumaNode, transfer_mode: TransferMode, kv_caches: HashMap, + layer_groups: HashMap, ) -> Result, EngineError> { if device_id < 0 { return Err(EngineError::InvalidArgument(format!( @@ -582,6 +711,7 @@ impl InstanceContext { numa_node, transfer_mode, kv_caches, + layer_groups, )?)) } @@ -635,6 +765,7 @@ impl InstanceContext { numa_node, transfer_mode, kv_caches, + layer_groups, } = registration; if tp_rank >= self.tp_size { @@ -661,6 +792,7 @@ impl InstanceContext { numa_node, transfer_mode, kv_caches, + layer_groups, )?; { diff --git a/pegaflow-core/src/instance/tests.rs b/pegaflow-core/src/instance/tests.rs index 6358dc93..9513ab21 100644 --- a/pegaflow-core/src/instance/tests.rs +++ b/pegaflow-core/src/instance/tests.rs @@ -22,6 +22,23 @@ fn gpu_registration(device_id: i32, tp_rank: usize, layers: &[&str]) -> GpuRegis gpu_registration_with_segment_bytes(device_id, tp_rank, layers, 1024) } +/// Build a `GpuRegistration` with explicit hybrid-cache group ids per layer. +/// `(name, group)` pairs; group 0 is the attention default, higher groups are +/// e.g. recurrent-state checkpoints in hybrid (mamba) models. +fn gpu_registration_with_groups( + device_id: i32, + tp_rank: usize, + layers_and_groups: &[(&str, u32)], +) -> GpuRegistration { + let layers: Vec<&str> = layers_and_groups.iter().map(|(name, _)| *name).collect(); + let mut registration = gpu_registration(device_id, tp_rank, &layers); + registration.layer_groups = layers_and_groups + .iter() + .map(|(name, group)| ((*name).to_string(), *group)) + .collect(); + registration +} + fn gpu_registration_with_segment_bytes( device_id: i32, tp_rank: usize, @@ -48,6 +65,7 @@ fn gpu_registration_with_segment_bytes( numa_node: NumaNode::UNKNOWN, transfer_mode: TransferMode::Direct, kv_caches, + layer_groups: HashMap::new(), } } @@ -448,3 +466,127 @@ fn seal_rejects_inconsistent_layer_geometry() { .expect_err("same name with different geometry must fail the seal"); assert!(err.to_string().contains("inconsistent geometry")); } + +/// Hybrid (mamba-style) topology: attention layers in group 0, recurrent-state +/// layers in group 1. Each group gets its own dense slot space, which is what +/// lets a group seal a block independently of the others — the precondition +/// for "recurrent saves only the final block" to ever become visible. +/// Commits device 0. +#[test] +fn hybrid_topology_seals_per_group_slot_spaces() { + // tp_size=1 keeps this single-device; the group-rank reset and per-group + // totals are what distinguish groups (tp_rank math is unchanged). + let instance = InstanceContext::new("hybrid".into(), "hybrid-ns".into(), 1, 1, false).unwrap(); + instance + .register_new_gpu(gpu_registration_with_groups( + 0, + 0, + &[ + ("attn_a", 0), + ("attn_b", 0), + ("recurrent_c", 1), + ("recurrent_d", 1), + ], + )) + .expect("register hybrid gpu"); + + let topology = instance.sealed_topology().expect("sealed"); + assert_eq!(topology.num_layers(), 4); + assert_eq!(topology.num_groups(), 2); + + // Sorted-name ids: attn_a=0, attn_b=1, recurrent_c=2, recurrent_d=3. + assert_eq!(topology.group_of_layer(0), 0); + assert_eq!(topology.group_of_layer(1), 0); + assert_eq!(topology.group_of_layer(2), 1); + assert_eq!(topology.group_of_layer(3), 1); + + // Per-group seal domains: group_slots = group_layers * tp_size. + assert_eq!(topology.group_total_slots(0).unwrap(), 2); + assert_eq!(topology.group_total_slots(1).unwrap(), 2); + // Union across groups stays the historical denominator. + assert_eq!(topology.total_slots(), 4); + + // Layer-first slots are dense WITHIN the group: group rank, not global id. + assert_eq!(topology.slot_index(0, 0).unwrap(), 0); // attn_a, group rank 0 + assert_eq!(topology.slot_index(1, 0).unwrap(), 1); // attn_b, group rank 1 + assert_eq!(topology.slot_index(2, 0).unwrap(), 0); // recurrent_c restarts at 0 + assert_eq!(topology.slot_index(3, 0).unwrap(), 1); +} + +/// With every layer in the default group 0 the within-group rank equals the +/// global layer id, so all existing slot math is bit-identical. This is the +/// backward-compat guard for current connectors. Commits device 0. +#[test] +fn default_groups_keep_global_slot_layout() { + let instance = + InstanceContext::new("classic".into(), "classic-ns".into(), 1, 1, false).unwrap(); + instance + .register_new_gpu(gpu_registration(0, 0, &["layer_a", "layer_b", "layer_c"])) + .expect("register classic gpu"); + + let topology = instance.sealed_topology().expect("sealed"); + assert_eq!(topology.num_groups(), 1); + assert_eq!(topology.group_total_slots(0).unwrap(), 3); + for layer_id in 0..3 { + assert_eq!(topology.group_of_layer(layer_id), 0); + assert_eq!(topology.slot_index(layer_id, 0).unwrap(), layer_id); + } + assert!(topology.group_total_slots(1).is_err()); +} + +/// A layer must live in the same storage group on every worker; disagreeing +/// devices would seal blocks with different slot counts per device silently. +/// Needs 2 CUDA devices. +#[test] +fn seal_rejects_inconsistent_layer_group_across_devices() { + if !has_cuda_devices(2) { + eprintln!( + "skipping seal_rejects_inconsistent_layer_group_across_devices: needs >= 2 CUDA devices" + ); + return; + } + + let instance = + InstanceContext::new("grp-conflict".into(), "grp-ns".into(), 1, 2, false).unwrap(); + instance + .register_new_gpu(gpu_registration_with_groups(0, 0, &[("layer_0", 0)])) + .expect("first worker"); + + let err = instance + .register_new_gpu(gpu_registration_with_groups(1, 0, &[("layer_0", 1)])) + .expect_err("same layer in different groups must fail the seal"); + assert!(err.to_string().contains("storage group"), "{err}"); +} + +/// Group ids must be dense `0..N`: a gap means a group nobody registers, whose +/// blocks would sit inflight forever waiting for slots that cannot arrive. +/// Commits device 0. +#[test] +fn seal_rejects_sparse_group_ids() { + let instance = InstanceContext::new("sparse".into(), "sparse-ns".into(), 1, 1, false).unwrap(); + let err = instance + .register_new_gpu(gpu_registration_with_groups( + 0, + 0, + &[("layer_a", 0), ("layer_b", 2)], + )) + .expect_err("group 1 with no layers must fail the seal"); + assert!(err.to_string().contains("dense"), "{err}"); +} + +/// Page-first storage packs all layers of a block into one page per shard; +/// splitting blocks across groups contradicts that. Reject instead of +/// silently picking one layout. Commits device 0. +#[test] +fn page_first_rejects_multiple_groups() { + let instance = + InstanceContext::new("page-hybrid".into(), "page-hybrid-ns".into(), 1, 1, true).unwrap(); + let err = instance + .register_new_gpu(gpu_registration_with_groups( + 0, + 0, + &[("layer_a", 0), ("layer_b", 1)], + )) + .expect_err("page-first with two storage groups must be rejected"); + assert!(err.to_string().contains("page-first"), "{err}"); +} diff --git a/pegaflow-core/src/internode/metaserver_client.rs b/pegaflow-core/src/internode/metaserver_client.rs index c5bf1c17..4193839a 100644 --- a/pegaflow-core/src/internode/metaserver_client.rs +++ b/pegaflow-core/src/internode/metaserver_client.rs @@ -4,11 +4,11 @@ use std::sync::Weak; use log::{debug, error, info, warn}; use pegaflow_common::grpc::{GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL, GRPC_CONNECT_TIMEOUT}; use pegaflow_proto::proto::engine::meta_server_client::MetaServerClient as MetaServerGrpcClient; +#[cfg(feature = "rdma")] +use pegaflow_proto::proto::engine::{FetchSegment, QueryPrefixBlocksRequest}; use pegaflow_proto::proto::engine::{ HeartbeatNodeRequest, InsertBlockHashesRequest, RemoveBlockHashesRequest, UnregisterNodeRequest, }; -#[cfg(feature = "rdma")] -use pegaflow_proto::proto::engine::{NodePrefixResult, QueryPrefixBlocksRequest}; use tokio::sync::{mpsc, oneshot}; use tokio::time::{Duration, Instant}; use tonic::Code; @@ -295,17 +295,18 @@ impl MetaServerClient { let _ = done_rx.await; } - /// Query MetaServer for the longest prefix of blocks that exist remotely. - /// Returns per-node prefix lengths. + /// Query MetaServer for an ordered remote fetch plan. #[cfg(feature = "rdma")] - pub(crate) async fn query_prefix( + pub(crate) async fn query_plan( &self, namespace: &str, hashes: &[Vec], - ) -> Result, ClientError> { + exclude_node: &str, + ) -> Result, ClientError> { let request = QueryPrefixBlocksRequest { namespace: namespace.to_string(), block_hashes: hashes.to_vec(), + exclude_node: exclude_node.to_string(), }; let response = self @@ -318,12 +319,12 @@ impl MetaServerClient { let resp = response.into_inner(); debug!( - "MetaServer query_prefix: namespace={} nodes={}", + "MetaServer query_plan: namespace={} segments={}", namespace, - resp.nodes.len() + resp.segments.len() ); - Ok(resp.nodes) + Ok(resp.segments) } } @@ -850,6 +851,7 @@ mod tests { reclaimable_hashes: Mutex>>, insert_requests: RequestLog, remove_requests: RequestLog, + query_requests: Mutex>, heartbeat_notify: Notify, insert_notify: Notify, remove_notify: Notify, @@ -949,9 +951,16 @@ mod tests { async fn query_prefix_blocks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Ok(Response::new(QueryPrefixBlocksResponse { nodes: vec![] })) + self.state + .query_requests + .lock() + .unwrap() + .push(request.into_inner()); + Ok(Response::new(QueryPrefixBlocksResponse { + segments: vec![], + })) } } @@ -1128,6 +1137,34 @@ mod tests { let _ = shutdown_tx.send(()); } + #[cfg(feature = "rdma")] + #[tokio::test] + async fn query_plan_sends_requester_as_excluded_node() { + let (addr, service, shutdown_tx) = start_fake_metaserver().await; + let client = MetaServerClient::new( + MetaServerClientConfig::new(addr, "node-a:50055".to_string()), + Weak::new(), + ); + let hashes = vec![vec![1], vec![2]]; + + let segments = client + .query_plan("ns", &hashes, "node-a:50055") + .await + .expect("query should succeed"); + + assert!(segments.is_empty()); + assert_eq!( + *service.query_requests.lock().unwrap(), + vec![QueryPrefixBlocksRequest { + namespace: "ns".to_string(), + block_hashes: hashes, + exclude_node: "node-a:50055".to_string(), + }] + ); + client.shutdown().await; + let _ = shutdown_tx.send(()); + } + #[test] fn unsent_after_failure_counts_only_unsent_tail() { let ns = |name: &str, n: usize| (name.to_string(), vec![vec![0u8]; n]); diff --git a/pegaflow-core/src/lib.rs b/pegaflow-core/src/lib.rs index 0dc1a3eb..7c15b49b 100644 --- a/pegaflow-core/src/lib.rs +++ b/pegaflow-core/src/lib.rs @@ -44,7 +44,7 @@ pub use internode::{ use layout::KVCacheLayout; pub use lease::QueryLeaseId; pub use pegaflow_common::NumaNode; -use pegaflow_common::NumaTopology; +use pegaflow_common::{NumaTopology, group_hash}; pub use pinned_pool::PinnedAllocation; pub use seal_offload::SlotMeta; pub use storage::{DEFAULT_RDMA_QPS_PER_PEER, MemoryCacheCleanupStats, StorageConfig}; @@ -281,6 +281,7 @@ impl PegaEngine { kv_stride_bytes_list, segments_list, None, + None, transfer_mode, page_first, ) @@ -315,6 +316,7 @@ impl PegaEngine { kv_stride_bytes_list: &[usize], segments_list: &[usize], block_stride_bytes_list: Option<&[usize]>, + layer_group_ids: Option<&[u32]>, transfer_mode: TransferMode, page_first: bool, ) -> Result<(), EngineError> { @@ -346,6 +348,14 @@ impl PegaEngine { strides.len() ))); } + if let Some(group_ids) = layer_group_ids + && group_ids.len() != batch_size + { + return Err(EngineError::InvalidArgument(format!( + "registration metadata length mismatch: layer_names={batch_size}, layer_group_ids={}", + group_ids.len() + ))); + } let mut kv_caches = HashMap::with_capacity(batch_size); for i in 0..batch_size { @@ -384,6 +394,15 @@ impl PegaEngine { } } + let layer_groups: HashMap = match layer_group_ids { + Some(group_ids) => layer_names + .iter() + .cloned() + .zip(group_ids.iter().copied()) + .collect(), + None => HashMap::new(), + }; + // Get or create instance let instance = self.get_or_create_instance(instance_id, namespace, tp_size, world_size, page_first)?; @@ -410,6 +429,7 @@ impl PegaEngine { numa_node, transfer_mode, kv_caches, + layer_groups, })?; info!( @@ -463,6 +483,11 @@ impl PegaEngine { .collect() } + /// Return the namespace associated with a registered instance. + pub fn instance_namespace(&self, instance_id: &str) -> Result { + Ok(self.get_instance(instance_id)?.namespace().to_string()) + } + /// Count prefix hit blocks with SSD prefetch support. /// /// Argument contract: @@ -507,6 +532,84 @@ impl PegaEngine { Ok(status) } + /// All-or-nothing membership fetch over one hybrid-cache storage group, + /// eligible for the same SSD prefetch and MetaServer + RDMA remote fetch + /// as prefix queries. + /// + /// Where [`Self::query_group_membership`] answers from the resident read + /// cache only, this treats `block_hashes` as an exact want-set: misses are + /// pulled from remote tiers, and the query stays `Loading` until the whole + /// set is fetchable (the prefix machinery over an explicit key list *is* + /// a set fetch once the full length is required). Use it when partial + /// state is useless — e.g. restoring a recurrent-state checkpoint on a + /// prefill/decode handoff, where the peer that saved the set holds every + /// member. A `Ready` result with fewer blocks than requested means the + /// set could not be completed anywhere; callers treat that as a miss. + pub async fn query_group_membership_with_fetch( + &self, + instance_id: &str, + req_id: &str, + group_id: u32, + block_hashes: &[Vec], + ) -> Result { + let instance = self.get_instance(instance_id)?; + let topology = instance.sealed_topology()?; + // Same contract as the local membership query: an unknown group is a + // bug in the caller, not an all-miss answer. + topology.group_total_slots(group_id)?; + + let namespace = instance.namespace(); + let encoded: Vec> = block_hashes + .iter() + .map(|hash| group_hash(hash, group_id)) + .collect(); + + let status = self + .storage + .check_prefix_and_prefetch(req_id, namespace, &encoded, true) + .await; + + if let PrefetchStatus::Ready { blocks, missing } = &status { + let metrics = core_metrics(); + metrics.cache_block_hits.add(blocks.len() as u64, &[]); + if *missing > 0 { + metrics.cache_block_misses.add(*missing as u64, &[]); + } + } + + Ok(status) + } + + /// Position-aligned membership query over one hybrid-cache storage group. + /// + /// Unlike prefix queries, every position reports independently: entry `i` + /// is the sealed block for `block_hashes[i]` in `group_id`, or `None` on + /// miss. Sparse hit patterns are the point — callers (e.g. the vLLM + /// connector's hybrid reconcile) pick the rightmost hit themselves. + /// Group 0 keeps raw-hash keys, so classic instances observe no change. + /// + /// The returned blocks hold plain `Arc` refs, not leases: pin what you + /// need via [`Self::create_query_lease`]. + pub fn query_group_membership( + &self, + instance_id: &str, + group_id: u32, + block_hashes: &[Vec], + ) -> Result>>, EngineError> { + let instance = self.get_instance(instance_id)?; + let topology = instance.sealed_topology()?; + // Validate the group against the sealed topology; an unknown group is + // a contract bug, not an answer of "all miss". + topology.group_total_slots(group_id)?; + + let namespace = instance.namespace(); + let encoded: Vec> = block_hashes + .iter() + .map(|hash| group_hash(hash, group_id)) + .collect(); + Ok(self.storage.get_membership(namespace, &encoded)) + } + /// Create an opaque lease that owns query-ready blocks. pub fn create_query_lease( &self, @@ -642,6 +745,32 @@ impl PegaEngine { )); } + // Resolve each load group's storage group and require homogeneity: a + // group's blocks seal against exactly one slot space. An empty load + // group owns no storage group; targets pointing at it load into + // nothing (it exists only to keep group indices aligned with the + // connector's cache-group layout). + let mut storage_group_of_load_group: Vec> = + Vec::with_capacity(layer_groups.len()); + for layer_names in layer_groups { + let mut group: Option = None; + for layer_name in layer_names { + let layer_id = topology.layer_id(layer_name)?; + let layer_group = topology.group_of_layer(layer_id); + match group { + None => group = Some(layer_group), + Some(existing) if existing == layer_group => {} + Some(existing) => { + return Err(EngineError::InvalidArgument(format!( + "load group mixes storage groups {existing} and {layer_group} \ + (layer {layer_name})", + ))); + } + } + } + storage_group_of_load_group.push(group); + } + // Consume query leases reserved for this load. trace_scope!("load.cache_lookup", _s); let mut block_targets_by_group = vec![Vec::new(); layer_groups.len()]; @@ -668,6 +797,28 @@ impl PegaEngine { group_index ))); } + // A stored block must carry exactly the slot layout of the + // storage group this target group loads into. A mismatch + // means the namespace is shared by instances with different + // layer sets (e.g. MTP enabled vs disabled) — loading would + // silently leave layers uninitialized, so fail loudly and + // let vLLM recompute. Empty load groups own no storage + // group and skip the check entirely. + if let Some(storage_group) = storage_group_of_load_group[group_index] { + let expected_slots = topology.group_total_slots(storage_group)?; + for (source_index, destination) in lease_block_targets.iter().enumerate() { + if destination.is_some() + && blocks[source_index].slots().len() != expected_slots + { + return Err(EngineError::InvalidArgument(format!( + "stored block has {} slots but storage group {storage_group} of \ + instance {instance_id} expects {expected_slots}: \ + namespace is shared by incompatible KV layouts", + blocks[source_index].slots().len(), + ))); + } + } + } block_targets_by_group[group_index].extend( lease_block_targets.iter().enumerate().filter_map( |(source_index, destination)| { @@ -676,21 +827,6 @@ impl PegaEngine { ), ); } - // A stored block must carry exactly this instance's slot layout. - // A mismatch means the namespace is shared by instances with - // different layer sets (e.g. MTP enabled vs disabled) — loading - // would silently leave layers uninitialized, so fail loudly and - // let vLLM recompute. - for block in &blocks { - if block.slots().len() != topology.total_slots() { - return Err(EngineError::InvalidArgument(format!( - "stored block has {} slots but instance {instance_id} expects {}: \ - namespace is shared by incompatible KV layouts", - block.slots().len(), - topology.total_slots() - ))); - } - } block_cache.extend(blocks); } trace_drop!(_s); diff --git a/pegaflow-core/src/metrics.rs b/pegaflow-core/src/metrics.rs index a6fd5362..4879499b 100644 --- a/pegaflow-core/src/metrics.rs +++ b/pegaflow-core/src/metrics.rs @@ -109,6 +109,10 @@ pub(crate) struct CoreMetrics { pub rdma_fetch_duration_seconds: Histogram, #[cfg(feature = "rdma")] pub rdma_fetch_bytes: Counter, + #[cfg(feature = "rdma")] + pub rdma_fetch_plan_segments: Histogram, + #[cfg(feature = "rdma")] + pub rdma_fetch_plan_completed_segments: Histogram, } fn init_meter() -> Meter { @@ -136,6 +140,14 @@ fn rdma_fetch_duration_boundaries() -> Vec { ] } +/// Histogram boundaries for the number of segments in one RDMA fetch plan. +#[cfg(feature = "rdma")] +fn rdma_fetch_plan_segment_boundaries() -> Vec { + vec![ + 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 16.0, 32.0, 64.0, 128.0, + ] +} + /// Tail-focused transfer duration boundaries in seconds. /// /// Load/save debugging cares more about sustained tail regressions than small @@ -480,6 +492,20 @@ pub(crate) fn core_metrics() -> &'static CoreMetrics { .with_unit("bytes") .with_description("Total bytes fetched via RDMA from remote nodes") .build(), + #[cfg(feature = "rdma")] + rdma_fetch_plan_segments: meter + .u64_histogram("pegaflow_rdma_fetch_plan_segments") + .with_description("Number of segments planned per executed RDMA fetch plan") + .with_boundaries(rdma_fetch_plan_segment_boundaries()) + .build(), + #[cfg(feature = "rdma")] + rdma_fetch_plan_completed_segments: meter + .u64_histogram("pegaflow_rdma_fetch_plan_completed_segments") + .with_description( + "Number of segments completed before an RDMA fetch plan stopped", + ) + .with_boundaries(rdma_fetch_plan_segment_boundaries()) + .build(), } }) } @@ -510,4 +536,15 @@ mod tests { &[KeyValue::new("reason", "cleanup")] ); } + + #[cfg(feature = "rdma")] + #[test] + fn rdma_fetch_plan_segment_boundaries_cover_failures_and_fragmented_plans() { + assert_eq!( + rdma_fetch_plan_segment_boundaries(), + vec![ + 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 16.0, 32.0, 64.0, 128.0, + ] + ); + } } diff --git a/pegaflow-core/src/offload.rs b/pegaflow-core/src/offload.rs index 7cf938a4..43106a2f 100644 --- a/pegaflow-core/src/offload.rs +++ b/pegaflow-core/src/offload.rs @@ -6,7 +6,7 @@ // is deferred to the storage insert worker via `RawSaveBatch`. // ============================================================================ -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::num::NonZeroU64; use std::sync::Arc; @@ -21,7 +21,7 @@ use crate::layout::KVCacheLayout; use crate::metrics::core_metrics; use crate::pinned_pool::PinnedAllocation; use crate::{EngineError, PegaEngine}; -use pegaflow_common::NumaNode; +use pegaflow_common::{NumaNode, group_hash}; // ============================================================================ // Types sent to the insert worker (deferred Phase 4) @@ -56,6 +56,8 @@ struct LayerContext { layer_name: String, layout: KVCacheLayout, slot_id: usize, + /// Hybrid-cache storage group of this layer; keyed and sealed per group. + group: u32, /// Blocks to save: (block_idx, hash). Filtered in Phase 1. blocks_to_save: Vec<(usize, Vec)>, /// Host-side block images, parallel to `blocks_to_save`. @@ -138,6 +140,35 @@ fn build_hashed_insert_entries(namespace: String, layers: Vec) -> .collect() } +/// Drop `(group, hash)` candidates whose group-encoded key already exists in +/// the read cache. Filtering is per group because groups key the same content +/// hash independently (e.g. attention block vs. recurrent checkpoint). +fn filter_new_hashes_per_group( + storage: &crate::storage::StorageEngine, + namespace: &str, + candidates: HashSet<(u32, Vec)>, +) -> HashSet<(u32, Vec)> { + let mut by_group: HashMap>> = HashMap::new(); + for (group, hash) in &candidates { + by_group.entry(*group).or_default().push(hash.clone()); + } + + let mut survivors = HashSet::with_capacity(candidates.len()); + for (group, raw_hashes) in by_group { + let mut encoded: HashSet> = raw_hashes + .iter() + .map(|hash| group_hash(hash, group)) + .collect(); + storage.filter_hashes_not_in_cache_inplace(namespace, &mut encoded); + for raw in raw_hashes { + if encoded.contains(&group_hash(&raw, group)) { + survivors.insert((group, raw)); + } + } + } + survivors +} + // ============================================================================ // Save methods on PegaEngine (moved from lib.rs) // ============================================================================ @@ -275,7 +306,6 @@ impl PegaEngine { let instance = self.get_instance(instance_id)?; let topology = instance.sealed_topology()?; let namespace = instance.namespace().to_string(); - let total_slots = topology.total_slots(); // ── Phase 0: Resolve per-layer metadata and build valid_blocks ── trace_scope!("save.resolve_metadata", _s); @@ -283,7 +313,9 @@ impl PegaEngine { let gpu_context = instance.get_gpu_for_save_group(device_id, tp_rank, pp_rank)?; let mut layer_contexts: Vec = Vec::with_capacity(saves.len()); - let mut hashes_to_save: HashSet> = HashSet::new(); + // Dedupe candidates per (group, hash): the same content hash may exist + // in the attention group while the recurrent group still needs it. + let mut hashes_to_save: HashSet<(u32, Vec)> = HashSet::new(); for LayerSave { layer_name, @@ -309,6 +341,7 @@ impl PegaEngine { })?; let slot_id = topology.slot_index(layer_id, tp_rank)?; + let group = topology.group_of_layer(layer_id); let num_blocks = layout.num_blocks(); @@ -331,13 +364,14 @@ impl PegaEngine { } for (_, hash) in &blocks_to_save { - hashes_to_save.insert(hash.clone()); + hashes_to_save.insert((group, hash.clone())); } layer_contexts.push(LayerContext { layer_name, layout, slot_id, + group, blocks_to_save, raw_blocks: Vec::new(), }); @@ -356,20 +390,23 @@ impl PegaEngine { trace_scope!("save.hash_filter", _s); - // Single in-place cache filter for all unique hashes - self.storage - .filter_hashes_not_in_cache_inplace(&namespace, &mut hashes_to_save); + let hashes_to_save = filter_new_hashes_per_group( + &self.storage, + &namespace, + std::mem::take(&mut hashes_to_save), + ); if hashes_to_save.is_empty() { trace_drop!(_s); return Ok(()); } - // Per-layer filter: keep only blocks whose hash needs saving + // Per-layer filter: keep only blocks whose (group, hash) needs saving let mut total_blocks_to_save = 0usize; for ctx in &mut layer_contexts { + let group = ctx.group; ctx.blocks_to_save - .retain(|(_, hash)| hashes_to_save.contains(hash.as_slice())); + .retain(|(_, hash)| hashes_to_save.contains(&(group, hash.clone()))); total_blocks_to_save += ctx.blocks_to_save.len(); } // Remove layers with no blocks to save @@ -550,10 +587,15 @@ impl PegaEngine { batch_start.elapsed().as_secs_f64() * 1000.0 ); - // Build RawSaveBatch and send to insert worker (fire-and-forget). - // Page-first collapses a shard's layers into one page slot per block, so - // a whole shard is a single RawSaveLayer regardless of its layer count. - let raw_layers: Vec = if page_first { + // Build per-group RawSaveBatches and send to the insert worker + // (fire-and-forget). Each group seals blocks against its own slot + // count and stores them under group-encoded hashes (group 0 keeps raw + // hashes, so single-group instances produce byte-identical keys). + if page_first { + // Page-first collapses a shard's layers into one page slot per + // block, so a whole shard is a single RawSaveLayer regardless of + // its layer count. Sealing rejects multi-group page-first, so the + // whole batch is group 0. let slot_id = layer_contexts[0].slot_id; let page_size = topology .shard_page_size(slot_id) @@ -570,37 +612,48 @@ impl PegaEngine { RawBlock::single_segment(Segment::new(ptr, page_size, page)) }) .collect(); - vec![RawSaveLayer { - slot_id, - padded_block_size: page_size, - blocks, - block_hashes, - }] + self.storage.send_raw_insert(RawSaveBatch { + namespace, + total_slots: topology.total_slots(), + numa_node: save_numa_node, + layers: vec![RawSaveLayer { + slot_id, + padded_block_size: page_size, + blocks, + block_hashes, + }], + }); } else { - layer_contexts - .into_iter() - .map(|ctx| { - let block_hashes: Vec> = ctx - .blocks_to_save - .into_iter() - .map(|(_, hash)| hash) - .collect(); - RawSaveLayer { - slot_id: ctx.slot_id, - padded_block_size: ctx.layout.padded_block_bytes(), - blocks: ctx.raw_blocks, - block_hashes, - } - }) - .collect() - }; - - self.storage.send_raw_insert(RawSaveBatch { - namespace, - total_slots, - numa_node: save_numa_node, - layers: raw_layers, - }); + let mut by_group: std::collections::BTreeMap> = + Default::default(); + for ctx in layer_contexts { + by_group.entry(ctx.group).or_default().push(ctx); + } + for (group, contexts) in by_group { + let raw_layers: Vec = contexts + .into_iter() + .map(|ctx| { + let block_hashes: Vec> = ctx + .blocks_to_save + .into_iter() + .map(|(_, hash)| group_hash(&hash, group)) + .collect(); + RawSaveLayer { + slot_id: ctx.slot_id, + padded_block_size: ctx.layout.padded_block_bytes(), + blocks: ctx.raw_blocks, + block_hashes, + } + }) + .collect(); + self.storage.send_raw_insert(RawSaveBatch { + namespace: namespace.clone(), + total_slots: topology.group_total_slots(group)?, + numa_node: save_numa_node, + layers: raw_layers, + }); + } + } Ok(()) } diff --git a/pegaflow-core/src/storage/mod.rs b/pegaflow-core/src/storage/mod.rs index 857ee8f2..e17aee96 100644 --- a/pegaflow-core/src/storage/mod.rs +++ b/pegaflow-core/src/storage/mod.rs @@ -5,8 +5,6 @@ pub(crate) mod transfer_lock; mod write_path; use bytesize::ByteSize; -#[cfg(not(feature = "rdma"))] -use log::warn; use log::{debug, info, warn}; use std::collections::HashSet; use std::num::NonZeroU64; @@ -380,6 +378,21 @@ impl StorageEngine { } } + /// Position-aligned membership lookup in the resident read cache: entry + /// `i` is the sealed block for `hashes[i]`, or `None` on miss. Hashes must + /// already carry any group encoding (see `group_hash`). + pub(crate) fn get_membership( + &self, + namespace: &str, + hashes: &[Vec], + ) -> Vec>> { + let keys: Vec = hashes + .iter() + .map(|hash| BlockKey::new(namespace.to_string(), hash.clone())) + .collect(); + self.read_cache.get_blocks_aligned(&keys) + } + /// Evict all blocks from the resident in-memory read cache. /// /// This preserves backing-store copies. Blocks with @@ -502,22 +515,12 @@ impl StorageEngine { } let mut batch_bytes = 0u64; - let mut still_referenced = 0u64; for (_key, block) in &evicted { let b = block.memory_footprint(); batch_bytes = batch_bytes.saturating_add(b); - if Arc::strong_count(block) > 1 { - still_referenced += 1; - } core_metrics().cache_resident_bytes.add(-(b as i64), &[]); } - if still_referenced > 0 { - core_metrics() - .cache_block_evictions_still_referenced - .add(still_referenced, &[]); - } - freed_bytes = freed_bytes.saturating_add(batch_bytes); freed_blocks += evicted.len(); diff --git a/pegaflow-core/src/storage/prefetch.rs b/pegaflow-core/src/storage/prefetch.rs index 5896c99f..29820210 100644 --- a/pegaflow-core/src/storage/prefetch.rs +++ b/pegaflow-core/src/storage/prefetch.rs @@ -44,23 +44,23 @@ impl RdmaFetch { remaining_hashes: &[Vec], require_full_prefix: bool, ) -> Option<(usize, PrefetchResult)> { - let (node, found) = self.0.query_prefix(namespace, remaining_hashes).await?; + let plan = self.0.query_plan(namespace, remaining_hashes).await?; + let found = plan.block_count(); if require_full_prefix && found != remaining_hashes.len() { return None; } let blocks = self .0 - .fetch_blocks(&node, req_id, namespace, &remaining_hashes[..found]) + .fetch_plan(&plan, req_id, namespace, remaining_hashes) .await; if require_full_prefix && blocks.len() != found { - // The advertised owner served fewer blocks than the MetaServer + // One planned segment served fewer blocks than the MetaServer // promised (stale advertisement or failed fetch). Keep the partial // result so poll_existing blacklists RDMA for this request instead // of the wait loop retrying the same fetch until timeout. warn!( - "RDMA fetch returned fewer blocks than advertised: req_id={} node={} returned={} advertised={}", + "RDMA fetch returned fewer blocks than planned: req_id={} returned={} planned={}", req_id, - node, blocks.len(), found ); diff --git a/pegaflow-core/src/storage/read_cache.rs b/pegaflow-core/src/storage/read_cache.rs index 1e3a6a01..5457bb26 100644 --- a/pegaflow-core/src/storage/read_cache.rs +++ b/pegaflow-core/src/storage/read_cache.rs @@ -140,17 +140,39 @@ impl ReadCache { found } + /// Position-aligned membership: entry `i` is the block for `keys[i]`, or + /// `None` on miss. Unlike [`Self::get_prefix_blocks`] this never stops at + /// the first gap — hybrid-cache checkpoint groups (recurrent state) have + /// sparse hit patterns by design, where the caller picks the rightmost + /// hit instead of a prefix. + pub(super) fn get_blocks_aligned(&self, keys: &[BlockKey]) -> Vec>> { + let mut inner = self.inner.lock(); + keys.iter() + .map(|key| { + inner.cache.get(key).inspect(|_| { + refresh_recency(&mut inner, key); + }) + }) + .collect() + } + pub(super) fn remove_lru_batch(&self, batch_size: usize) -> Vec<(BlockKey, Arc)> { let removed = { let mut inner = self.inner.lock(); let mut removed = Vec::with_capacity(batch_size); - while removed.len() < batch_size { - let next = remove_lru(&mut inner, ResidentClass::Reclaimable) - .or_else(|| remove_lru(&mut inner, ResidentClass::Retained)); - let Some(block) = next else { - break; - }; - removed.push(block); + remove_lru_batch_from_class( + &mut inner, + ResidentClass::Reclaimable, + batch_size, + &mut removed, + ); + if removed.len() < batch_size { + remove_lru_batch_from_class( + &mut inner, + ResidentClass::Retained, + batch_size, + &mut removed, + ); } removed }; @@ -337,6 +359,35 @@ fn remove_lru(inner: &mut ReadCacheInner, class: ResidentClass) -> Option, +) { + let candidates = class_lru(inner, class).len(); + for _ in 0..candidates { + if removed.len() == batch_size { + break; + } + + let Some(key) = class_lru(inner, class) + .iter() + .next() + .map(|(key, _)| key.clone()) + else { + break; + }; + if inner.cache.is_cache_owned_only(&key) { + let block = remove_lru(inner, class) + .expect("cache-owned LRU candidate must remain resident while locked"); + removed.push(block); + } else { + class_lru(inner, class).get(&key); + } + } +} + fn record_residence_durations( removed: Vec, attributes: &[opentelemetry::KeyValue], @@ -445,6 +496,35 @@ mod tests { ); } + #[test] + fn pressure_reclaim_ignores_weak_references() { + let cache = make_cache(); + let key = BlockKey::new("ns".into(), vec![1]); + let block = make_block(); + let weak = Arc::downgrade(&block); + cache.batch_insert(vec![(key.clone(), block)]); + + assert_eq!(cache.remove_lru_batch(1)[0].0, key); + assert!(weak.upgrade().is_none()); + } + + #[test] + fn pressure_reclaim_waits_for_external_strong_reference() { + let cache = make_cache(); + let key = BlockKey::new("ns".into(), vec![1]); + let block = make_block(); + let external = Arc::clone(&block); + let weak = Arc::downgrade(&block); + cache.batch_insert(vec![(key.clone(), block)]); + + assert!(cache.remove_lru_batch(1).is_empty()); + assert_class(&cache, &key, ResidentClass::Retained); + + drop(external); + assert_eq!(cache.remove_lru_batch(1)[0].0, key); + assert!(weak.upgrade().is_none()); + } + #[test] fn local_hit_refreshes_recency_without_changing_class() { let cache = make_cache(); diff --git a/pegaflow-core/tests/common/harness.rs b/pegaflow-core/tests/common/harness.rs index dfa2ef13..779b59a8 100644 --- a/pegaflow-core/tests/common/harness.rs +++ b/pegaflow-core/tests/common/harness.rs @@ -98,6 +98,11 @@ impl TestGpuData { assert_eq!(self.gpu.copy_to_host(), expected, "GPU data mismatch"); } + /// The deterministic host-side pattern backing this GPU buffer. + pub fn expected_bytes(&self) -> &[u8] { + &self.expected + } + pub fn ptr(&self) -> u64 { self.gpu.as_u64() } @@ -131,6 +136,7 @@ struct LayerSpec { block_size: usize, kv_stride: usize, segments: usize, + group: u32, } pub struct TestEnvBuilder { @@ -172,6 +178,27 @@ impl TestEnvBuilder { block_size, kv_stride: 0, segments: 1, + group: 0, + }); + self + } + + /// Add a contiguous layer in an explicit hybrid-cache storage group + /// (e.g. group 1 = recurrent-state checkpoints in mamba-style models). + pub fn grouped_layer( + mut self, + name: &'static str, + num_blocks: usize, + block_size: usize, + group: u32, + ) -> Self { + self.layers.push(LayerSpec { + name, + num_blocks, + block_size, + kv_stride: 0, + segments: 1, + group, }); self } @@ -198,6 +225,7 @@ impl TestEnvBuilder { block_size: segment_size, kv_stride, segments: 2, + group: 0, }); self } @@ -261,6 +289,7 @@ impl TestEnvBuilder { block_size: spec.block_size, kv_stride: spec.kv_stride, segments: spec.segments, + group: spec.group, }) .collect(); diff --git a/pegaflow-core/tests/common/helpers.rs b/pegaflow-core/tests/common/helpers.rs index 3b047bbd..3f24920a 100644 --- a/pegaflow-core/tests/common/helpers.rs +++ b/pegaflow-core/tests/common/helpers.rs @@ -21,6 +21,8 @@ pub struct LayerInfo { pub block_size: usize, pub kv_stride: usize, pub segments: usize, + /// Hybrid-cache storage group id; 0 is the default attention group. + pub group: u32, } #[allow( @@ -47,9 +49,10 @@ pub fn register_layers( let block_sizes: Vec = layers.iter().map(|l| l.block_size).collect(); let kv_strides: Vec = layers.iter().map(|l| l.kv_stride).collect(); let segments: Vec = layers.iter().map(|l| l.segments).collect(); + let group_ids: Vec = layers.iter().map(|l| l.group).collect(); engine - .register_context_layer_batch( + .register_context_layer_batch_strided( instance_id, namespace, device_id, @@ -64,6 +67,8 @@ pub fn register_layers( &block_sizes, &kv_strides, &segments, + None, + Some(&group_ids), transfer_mode, page_first, ) diff --git a/pegaflow-core/tests/eviction.rs b/pegaflow-core/tests/eviction.rs index bf5069e9..180b2a37 100644 --- a/pegaflow-core/tests/eviction.rs +++ b/pegaflow-core/tests/eviction.rs @@ -11,7 +11,7 @@ use pegaflow_core::StorageConfig; const BLOCK_SIZE: usize = 4096; const NUM_BLOCKS: usize = 4; -/// Pool fits one batch with headroom for metadata, but not two. +/// Pool fits two batches, so a third batch must reclaim one of them. const POOL_SIZE: usize = NUM_BLOCKS * BLOCK_SIZE * 2; fn eviction_storage_config() -> StorageConfig { @@ -61,11 +61,20 @@ async fn leased_blocks_survive_eviction_pressure() { env.save_and_wait(&hashes).await; let lease = env.assert_all_hit_lease(&hashes).await; - // Second batch creates eviction pressure while the first is leased. - let pressure = make_block_hashes(NUM_BLOCKS, 20); - env.save_layer(0, &pressure).await; + // Fill the remaining pool, then force a third batch to reclaim it while + // the first batch remains leased. + let reclaimed = make_block_hashes(NUM_BLOCKS, 20); + env.save_layer_and_flush(0, &reclaimed).await; + let pressure = make_block_hashes(NUM_BLOCKS, 21); + env.save_layer_and_flush(0, &pressure).await; + assert_eq!(env.count_hits_then_release(&reclaimed).await, 0); - // First batch is still leased — load must succeed. + // The lease pins both the allocation and its cache index, so an exact + // follow-up query must still see the same prefix after pressure. + let second_lease = env.assert_all_hit_lease(&hashes).await; + env.release(&second_lease); + + // The original lease must still load the pinned data. env.data().zero_gpu(); env.load_to_gpu(lease, hashes.len()).await; env.data().assert_gpu_matches_expected(); diff --git a/pegaflow-core/tests/hybrid_groups.rs b/pegaflow-core/tests/hybrid_groups.rs new file mode 100644 index 00000000..7fca17ad --- /dev/null +++ b/pegaflow-core/tests/hybrid_groups.rs @@ -0,0 +1,323 @@ +//! Hybrid-cache (mamba/HMA) storage groups. +//! +//! Attention groups save every block; the recurrent group saves only +//! checkpoint blocks (the state *after* a block's tokens). Each group seals +//! blocks against its own slot space, so a final-block-only recurrent save +//! becomes visible without waiting for attention's per-block cadence — and +//! one lease per group loads both halves back into the GPU. + +mod common; + +use common::*; +use pegaflow_core::*; + +/// The HMA bull's-eye: recurrent group seals a tail-only checkpoint save, +/// queries isolate groups by hash, and a two-lease load restores both the +/// attention prefix and the recurrent state block. +#[tokio::test] +async fn recurrent_group_seals_final_block_save() { + if !has_cuda_devices(1) { + eprintln!("skipping recurrent_group_seals_final_block_save: needs >= 1 CUDA device"); + return; + } + + const ATTN_BLOCK: usize = 1024; + const RECUR_BLOCK: usize = 2048; + + let env = TestEnvBuilder::new("hybrid-hma", "hybrid-hma-ns") + .grouped_layer("attn_0", 8, ATTN_BLOCK, 0) + .grouped_layer("attn_1", 8, ATTN_BLOCK, 0) + .grouped_layer("recurrent_state", 8, RECUR_BLOCK, 1) + .pool_size(64 << 20) + .build(); + + // Same hash namespace for both groups on purpose: group encoding, not the + // connector, must keep attention and recurrent blocks apart. + let hashes = make_block_hashes(8, 10); + let prefix_hashes = hashes[0..3].to_vec(); + + // Attention saves blocks 0..2 for both layers; the recurrent group saves + // ONLY the block-2 checkpoint (state after block 2's tokens), mirroring + // the connector's final-save. The recurrent checkpoint physically lives + // in GPU block 5 — state blocks are per-request slots, not prefix cells. + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![ + LayerSave { + layer_name: "attn_0".into(), + block_ids: vec![0, 1, 2], + block_hashes: prefix_hashes.clone(), + }, + LayerSave { + layer_name: "attn_1".into(), + block_ids: vec![0, 1, 2], + block_hashes: prefix_hashes.clone(), + }, + ], + ) + .await + .expect("save attention prefix"); + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![LayerSave { + layer_name: "recurrent_state".into(), + block_ids: vec![5], + block_hashes: vec![hashes[2].clone()], + }], + ) + .await + .expect("save recurrent checkpoint"); + env.engine.flush_saves().await; + + // Membership is per group: group 0 holds all three prefix blocks... + let attn_hits = env + .engine + .query_group_membership(&env.instance_id, 0, &prefix_hashes) + .expect("query group 0 membership"); + assert_eq!( + attn_hits.iter().map(|b| b.is_some()).collect::>(), + vec![true, true, true] + ); + + // ...while group 1 holds ONLY the block-2 checkpoint, even though the + // identical hash bytes for blocks 0/1 exist in group 0 (isolation). + let recur_hits = env + .engine + .query_group_membership(&env.instance_id, 1, &prefix_hashes) + .expect("query group 1 membership"); + assert_eq!( + recur_hits.iter().map(|b| b.is_some()).collect::>(), + vec![false, false, true], + "recurrent group must seal the tail-only checkpoint" + ); + + // The classic prefix query (group 0) is untouched by the recurrent save. + match env.query(&prefix_hashes).await { + PrefetchStatus::Ready { blocks, missing } => { + assert_eq!(blocks.len(), 3); + assert_eq!(missing, 0); + } + other => panic!("expected Ready, got {other:?}"), + } + + // Reconcile (connector-side): rightmost recurrent checkpoint within the + // attention prefix → hit = 2 + 1 = 3 blocks. + let attn_lease = match env.query(&prefix_hashes).await { + PrefetchStatus::Ready { blocks, .. } => env + .engine + .create_query_lease(&env.instance_id, blocks) + .expect("attention lease"), + other => panic!("expected Ready, got {other:?}"), + }; + let recur_block = recur_hits[2] + .as_ref() + .expect("rightmost recurrent checkpoint") + .clone(); + let recur_lease = env + .engine + .create_query_lease(&env.instance_id, vec![recur_block]) + .expect("recurrent lease"); + + // Snapshot expected bytes, then wipe GPU memory so the load proves the + // roundtrip instead of re-reading its own source. + let attn_expected: Vec> = env.layers[0..2] + .iter() + .map(|l| l.data.expected_bytes().to_vec()) + .collect(); + let recur_expected = env.layers[2].data.expected_bytes().to_vec(); + for layer in &env.layers { + layer.data.zero_gpu(); + } + + // One load, two leases, two layer groups: attention restores the [0,3) + // prefix; the recurrent group has exactly one physical target (the + // request's live state slot, block 7) and None elsewhere. + let layer_groups: Vec> = vec![vec!["attn_0", "attn_1"], vec!["recurrent_state"]]; + let load_state = LoadState::new().expect("create LoadState"); + let shm_name = load_state.shm_name().to_string(); + env.engine + .batch_load_kv_blocks_multi_layer( + &env.instance_id, + 0, + 0, + &shm_name, + &layer_groups, + &[ + ( + attn_lease, + vec![vec![Some(0), Some(1), Some(2)], vec![None, None, None]], + ), + (recur_lease, vec![vec![None], vec![Some(7)]]), + ], + ) + .expect("submit hybrid load"); + wait_for_load(&load_state, LOAD_WAIT_TIMEOUT).await; + + // Attention blocks 0..2 restored, the rest still zero. + for (layer_idx, expected) in attn_expected.iter().enumerate() { + let mut want = vec![0u8; expected.len()]; + want[..3 * ATTN_BLOCK].copy_from_slice(&expected[..3 * ATTN_BLOCK]); + env.layers[layer_idx].data.assert_gpu_matches(&want); + } + // Recurrent: only block 7 holds the checkpoint copied from block 5. + let mut want = vec![0u8; recur_expected.len()]; + want[7 * RECUR_BLOCK..8 * RECUR_BLOCK] + .copy_from_slice(&recur_expected[5 * RECUR_BLOCK..6 * RECUR_BLOCK]); + env.layers[2].data.assert_gpu_matches(&want); +} + +/// Group 0 must behave bit-identically to classic single-group storage when +/// nothing declares a nonzero group: prefix query, per-hash isolation, and +/// slot counts all follow the historical layout. +#[tokio::test] +async fn single_group_instances_are_unaffected() { + if !has_cuda_devices(1) { + eprintln!("skipping single_group_instances_are_unaffected: needs >= 1 CUDA device"); + return; + } + + let env = TestEnvBuilder::new("hybrid-compat", "hybrid-compat-ns") + .layer("layer_0", 4, 1024) + .layer("layer_1", 4, 1024) + .build(); + + let hashes = make_block_hashes(4, 7); + env.save_and_wait(&hashes[0..2]).await; + + let hits = env + .engine + .query_group_membership(&env.instance_id, 0, &hashes[0..4]) + .expect("membership"); + assert_eq!( + hits.iter().map(|b| b.is_some()).collect::>(), + vec![true, true, false, false] + ); + + // A nonsensical group on a single-group instance is a hard error, not an + // all-miss answer. + assert!( + env.engine + .query_group_membership(&env.instance_id, 7, &hashes[0..1]) + .is_err() + ); +} + +/// The prefill/decode handoff shape: `query_group_membership_with_fetch` +/// treats the hash list as an exact want-set with all-or-nothing semantics. +/// A fully saved recurrent checkpoint set answers Ready and complete; a +/// want-set with any absent member (no remote tier configured here) comes +/// back short, which callers must read as a miss. Group isolation holds: a +/// hash saved only under the attention group never satisfies the recurrent +/// want-set. +#[tokio::test] +async fn recurrent_membership_fetch_is_all_or_nothing() { + if !has_cuda_devices(1) { + eprintln!("skipping recurrent_membership_fetch_is_all_or_nothing: needs >= 1 CUDA device"); + return; + } + + const ATTN_BLOCK: usize = 1024; + const RECUR_BLOCK: usize = 2048; + + let env = TestEnvBuilder::new("hybrid-pd", "hybrid-pd-ns") + .grouped_layer("attn_0", 8, ATTN_BLOCK, 0) + .grouped_layer("recurrent_state", 8, RECUR_BLOCK, 1) + .pool_size(64 << 20) + .build(); + + let hashes = make_block_hashes(8, 21); + + // Attention saves blocks 0..2; the recurrent group saves the block-1 and + // block-2 checkpoints (a two-member set, like a multi-component state). + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![LayerSave { + layer_name: "attn_0".into(), + block_ids: vec![0, 1, 2], + block_hashes: hashes[0..3].to_vec(), + }], + ) + .await + .expect("save attention prefix"); + env.engine + .batch_save_kv_blocks_from_ipc( + &env.instance_id, + 0, + 0, + 0, + vec![LayerSave { + layer_name: "recurrent_state".into(), + block_ids: vec![4, 5], + block_hashes: vec![hashes[1].clone(), hashes[2].clone()], + }], + ) + .await + .expect("save recurrent checkpoints"); + env.engine.flush_saves().await; + + // The exact want-set resolves Ready and complete, in requested order. + let want = vec![hashes[1].clone(), hashes[2].clone()]; + match env + .engine + .query_group_membership_with_fetch(&env.instance_id, "pd-req-full", 1, &want) + .await + .expect("want-set query") + { + PrefetchStatus::Ready { blocks, missing } => { + assert_eq!(blocks.len(), 2, "full want-set must resolve completely"); + assert_eq!(missing, 0); + let _lease = env + .engine + .create_query_lease(&env.instance_id, blocks) + .expect("lease over the fetched set"); + } + other => panic!("expected Ready, got {other:?}"), + } + + // A want-set with an absent member comes back short (no remote tier is + // configured in this harness): the caller must treat that as a miss. + let short = vec![hashes[1].clone(), hashes[5].clone(), hashes[2].clone()]; + match env + .engine + .query_group_membership_with_fetch(&env.instance_id, "pd-req-short", 1, &short) + .await + .expect("short want-set query") + { + PrefetchStatus::Ready { blocks, missing } => { + assert!( + blocks.len() < short.len(), + "an incomplete want-set must not report as complete" + ); + assert!(missing > 0); + } + other => panic!("expected Ready, got {other:?}"), + } + + // Group isolation: hashes[0] exists in the attention group only, so the + // recurrent want-set containing it cannot complete. + let cross = vec![hashes[0].clone()]; + match env + .engine + .query_group_membership_with_fetch(&env.instance_id, "pd-req-cross", 1, &cross) + .await + .expect("cross-group query") + { + PrefetchStatus::Ready { blocks, missing } => { + assert_eq!(blocks.len(), 0); + assert_eq!(missing, 1); + } + other => panic!("expected Ready, got {other:?}"), + } +} diff --git a/pegaflow-metaserver/Cargo.toml b/pegaflow-metaserver/Cargo.toml index 1de7f6a5..44798aba 100644 --- a/pegaflow-metaserver/Cargo.toml +++ b/pegaflow-metaserver/Cargo.toml @@ -29,9 +29,11 @@ opentelemetry_sdk.workspace = true opentelemetry-prometheus.workspace = true prometheus.workspace = true axum.workspace = true +serde.workspace = true [dev-dependencies] criterion.workspace = true +tower = "0.5" [[bench]] name = "unregister_node" diff --git a/pegaflow-metaserver/README.md b/pegaflow-metaserver/README.md index f8afd910..1e91a1ee 100644 --- a/pegaflow-metaserver/README.md +++ b/pegaflow-metaserver/README.md @@ -164,28 +164,36 @@ message RemoveBlockHashesResponse { ### 5. QueryPrefixBlocks -Query the longest contiguous prefix of block hashes that exist, with per-node prefix lengths. +Build an ordered fetch plan for the longest contiguous prefix that is available +from live remote nodes. The requester is excluded from the owner candidates. **Request:** ```protobuf message QueryPrefixBlocksRequest { string namespace = 1; repeated bytes block_hashes = 2; // Ordered list of block hashes + string exclude_node = 3; // Requester address } ``` **Response:** ```protobuf -message NodePrefixResult { +message FetchSegment { string node = 1; - uint32 prefix_len = 2; // Consecutive hashes from h0 this node owns + uint32 block_count = 2; // Consecutive hashes fetched from this node } message QueryPrefixBlocksResponse { - repeated NodePrefixResult nodes = 1; + reserved 1; + reserved "nodes"; + repeated FetchSegment segments = 2; } ``` +Segments are ordered and contiguous. Their block counts are cumulative offsets +into the request's `block_hashes`; planning stops at the first hash with no live +remote owner. + ### 6. Health Health check endpoint. @@ -215,8 +223,8 @@ Graceful shutdown trigger. 2. **During server lifetime**: Call `HeartbeatNode` periodically with the same `node_id` 3. **On block save**: Call `InsertBlockHashes` with `{ node, node_id }` 4. **On cache eviction**: Call `RemoveBlockHashes` with `{ node, node_id }` -5. **On block query**: Call `QueryPrefixBlocks` to discover which live nodes hold a prefix -6. **On block load**: Query metaserver, then fetch from the best remote node via RDMA +5. **On block query**: Call `QueryPrefixBlocks` to build an ordered remote fetch plan +6. **On block load**: Fetch each plan segment in order via RDMA, stopping on the first failure ## Environment Variables diff --git a/pegaflow-metaserver/src/http_server.rs b/pegaflow-metaserver/src/http_server.rs index 5de06e18..e89fb1de 100644 --- a/pegaflow-metaserver/src/http_server.rs +++ b/pegaflow-metaserver/src/http_server.rs @@ -1,16 +1,29 @@ use axum::extract::State; use axum::http::StatusCode; use axum::response::IntoResponse; -use axum::{Router, routing::get}; +use axum::{ + Json, Router, + routing::{get, post}, +}; use log::{info, warn}; use prometheus::{Registry, TextEncoder}; +use serde::Serialize; use std::sync::Arc; use tokio::net::TcpListener; use tokio::sync::Notify; +use crate::store::{BlockHashStore, MANUAL_CLEANUP_AGE_SECS}; + #[derive(Clone)] struct AppState { prometheus_registry: Registry, + store: Arc, +} + +#[derive(Debug, Serialize)] +struct CleanupResponse { + removed_owners: usize, + removed_keys: usize, } async fn health_handler() -> &'static str { @@ -28,23 +41,46 @@ async fn metrics_handler(State(state): State) -> impl IntoResponse { ) } +async fn cleanup_expired_blocks_handler(State(state): State) -> Json { + let stats = state + .store + .remove_owners_older_than(std::time::Duration::from_secs(MANUAL_CLEANUP_AGE_SECS)); + Json(CleanupResponse { + removed_owners: stats.removed_owners, + removed_keys: stats.removed_keys, + }) +} + +fn app(state: AppState) -> Router { + Router::new() + .route("/health", get(health_handler)) + .route("/metrics", get(metrics_handler)) + .route( + "/admin/cleanup-expired-blocks", + post(cleanup_expired_blocks_handler), + ) + .with_state(state) +} + pub async fn start_http_server( addr: std::net::SocketAddr, prometheus_registry: Registry, + store: Arc, shutdown: Arc, ) -> Result, std::io::Error> { let listener = TcpListener::bind(addr).await?; let state = AppState { prometheus_registry, + store, }; - let app = Router::new() - .route("/health", get(health_handler)) - .route("/metrics", get(metrics_handler)) - .with_state(state); + let app = app(state); - info!("Starting HTTP server on {} (/health, /metrics)", addr); + info!( + "Starting HTTP server on {} (/health, /metrics, POST /admin/cleanup-expired-blocks)", + addr + ); let handle = tokio::spawn(async move { if let Err(err) = axum::serve(listener, app) @@ -59,3 +95,37 @@ pub async fn start_http_server( Ok(handle) } + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Body; + use axum::http::Request; + use tower::ServiceExt; + + #[tokio::test] + async fn cleanup_route_accepts_post_and_returns_stats() { + let response = app(AppState { + prometheus_registry: Registry::new(), + store: Arc::new(BlockHashStore::new()), + }) + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/cleanup-expired-blocks") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap() + .as_ref(), + br#"{"removed_owners":0,"removed_keys":0}"# + ); + } +} diff --git a/pegaflow-metaserver/src/lib.rs b/pegaflow-metaserver/src/lib.rs index ff5b137c..d2926be2 100644 --- a/pegaflow-metaserver/src/lib.rs +++ b/pegaflow-metaserver/src/lib.rs @@ -49,7 +49,7 @@ pub struct Cli { #[arg(long, default_value_t = store::DEFAULT_NODE_STALE_SECS)] pub node_stale_secs: u64, - /// Minutes before block ownership records are purged by the lifecycle sweep. + /// Deprecated compatibility setting. Background sweep no longer purges by owner age. #[arg(long, default_value_t = store::DEFAULT_TTL_MINUTES)] pub ttl_minutes: u64, @@ -113,33 +113,18 @@ pub async fn run() -> Result<(), Box> { info!("Starting PegaFlow MetaServer"); info!("Binding to address: {}", cli.addr); info!( - "Node lifecycle: stale_after={}s ttl={}m sweep_interval={}s", - cli.node_stale_secs, cli.ttl_minutes, cli.sweep_interval_secs + "Node lifecycle: stale_after={}s manual_cleanup_age=1h sweep_interval={}s (ttl_minutes={} retained for compatibility)", + cli.node_stale_secs, cli.sweep_interval_secs, cli.ttl_minutes ); - let ttl_secs = cli - .ttl_minutes - .checked_mul(60) - .ok_or("ttl-minutes is too large")?; - if cli.ttl_minutes == 0 { - return Err("ttl-minutes must be greater than 0".into()); - } if cli.sweep_interval_secs == 0 { return Err("sweep-interval-secs must be greater than 0".into()); } - if ttl_secs < cli.node_stale_secs { - return Err(format!( - "ttl-minutes ({}) must be >= node-stale-secs ({})", - cli.ttl_minutes, cli.node_stale_secs - ) - .into()); - } - // Initialize metrics let (meter_provider, prometheus_registry) = init_metrics()?; let store = Arc::new(BlockHashStore::with_config(store::StoreConfig { node_stale_after: Duration::from_secs(cli.node_stale_secs), - ttl: Duration::from_secs(ttl_secs), + ttl: Duration::MAX, })); // Register store observable gauges @@ -172,9 +157,13 @@ pub async fn run() -> Result<(), Box> { let shutdown = Arc::new(Notify::new()); // Start HTTP server for health check and metrics - let _http_handle = - http_server::start_http_server(cli.http_addr, prometheus_registry, Arc::clone(&shutdown)) - .await?; + let _http_handle = http_server::start_http_server( + cli.http_addr, + prometheus_registry, + Arc::clone(&store), + Arc::clone(&shutdown), + ) + .await?; // Create the gRPC service let service = GrpcMetaService::new(store.clone()); diff --git a/pegaflow-metaserver/src/service.rs b/pegaflow-metaserver/src/service.rs index 719eb59e..84ee9e14 100644 --- a/pegaflow-metaserver/src/service.rs +++ b/pegaflow-metaserver/src/service.rs @@ -1,18 +1,61 @@ use crate::metric::record_rpc_result; use crate::proto::engine::meta_server_server::MetaServer; use crate::proto::engine::{ - HeartbeatNodeRequest, HeartbeatNodeResponse, InsertBlockHashesRequest, - InsertBlockHashesResponse, NodePrefixResult, QueryPrefixBlocksRequest, - QueryPrefixBlocksResponse, RemoveBlockHashesRequest, RemoveBlockHashesResponse, ResponseStatus, - UnregisterNodeRequest, UnregisterNodeResponse, + FetchSegment, HeartbeatNodeRequest, HeartbeatNodeResponse, InsertBlockHashesRequest, + InsertBlockHashesResponse, QueryPrefixBlocksRequest, QueryPrefixBlocksResponse, + RemoveBlockHashesRequest, RemoveBlockHashesResponse, ResponseStatus, UnregisterNodeRequest, + UnregisterNodeResponse, }; -use crate::store::{BlockHashStore, StoreError}; +use crate::store::{BlockHashStore, PrefixEntry, StoreError}; use log::debug; use std::sync::Arc; use std::time::Instant; use tonic::{Request, Response, Status, async_trait}; use uuid::Uuid; +fn plan_fetch_segments( + entries: &[PrefixEntry], + exclude_node: &str, +) -> Result, &'static str> { + let mut segments = Vec::new(); + let mut offset = 0usize; + + while offset < entries.len() { + let mut best: Option<(&str, usize)> = None; + for candidate in &entries[offset].nodes { + let candidate = candidate.as_ref(); + if candidate == exclude_node { + continue; + } + + let end = entries[offset..] + .iter() + .take_while(|entry| entry.nodes.iter().any(|node| node.as_ref() == candidate)) + .count() + + offset; + + if best.is_none_or(|(best_node, best_end)| { + end > best_end || (end == best_end && candidate < best_node) + }) { + best = Some((candidate, end)); + } + } + + let Some((node, end)) = best else { + break; + }; + let block_count = + u32::try_from(end - offset).map_err(|_| "fetch segment block count exceeds uint32")?; + segments.push(FetchSegment { + node: node.to_string(), + block_count, + }); + offset = end; + } + + Ok(segments) +} + #[derive(Clone)] pub struct GrpcMetaService { store: Arc, @@ -236,9 +279,7 @@ impl MetaServer for GrpcMetaService { result } - /// Given an ordered list of block hashes, find the longest contiguous prefix - /// that exists in the store (stop at the first hash no node owns), then for - /// each node return how many consecutive hashes from h0 it holds. + /// Build an ordered plan that covers the longest remotely available prefix. async fn query_prefix_blocks( &self, request: Request, @@ -271,29 +312,10 @@ impl MetaServer for GrpcMetaService { req.namespace, prefix_len, total_queried, elapsed ); - // Per-node prefix: `existing[i]` maps to `block_hashes[i]`. - // A node's prefix = count of consecutive hashes from h0 it owns. - let mut node_prefix: std::collections::HashMap<&str, u32> = - std::collections::HashMap::new(); - for (i, entry) in existing.iter().enumerate() { - for node in &entry.nodes { - let count = node_prefix.entry(node.as_ref()).or_insert(0); - // Only extend if every previous hash (0..i) was also on this node - if *count == i as u32 { - *count += 1; - } - } - } - let nodes: Vec = node_prefix - .into_iter() - .filter(|(_, count)| *count > 0) - .map(|(node, count)| NodePrefixResult { - node: node.to_string(), - prefix_len: count, - }) - .collect(); + let segments = + plan_fetch_segments(&existing, &req.exclude_node).map_err(Status::invalid_argument)?; - let result = Ok(Response::new(QueryPrefixBlocksResponse { nodes })); + let result = Ok(Response::new(QueryPrefixBlocksResponse { segments })); record_rpc_result("query_prefix_blocks", &result, start); result } @@ -381,11 +403,12 @@ mod tests { .query_prefix_blocks(Request::new(QueryPrefixBlocksRequest { namespace: "ns".into(), block_hashes: vec![vec![1, 2, 3]], + exclude_node: String::new(), })) .await .unwrap() .into_inner(); - assert!(query_resp.nodes.is_empty()); + assert!(query_resp.segments.is_empty()); } #[tokio::test] @@ -423,11 +446,12 @@ mod tests { .query_prefix_blocks(Request::new(QueryPrefixBlocksRequest { namespace: "ns".into(), block_hashes: vec![vec![1, 2, 3]], + exclude_node: String::new(), })) .await .unwrap() .into_inner(); - assert_eq!(query_resp.nodes.len(), 1); + assert_eq!(query_resp.segments.len(), 1); } #[tokio::test] @@ -537,11 +561,12 @@ mod tests { .query_prefix_blocks(Request::new(QueryPrefixBlocksRequest { namespace: "ns".into(), block_hashes: vec![vec![1]], + exclude_node: String::new(), })) .await .unwrap() .into_inner(); - assert!(query_resp.nodes.is_empty()); + assert!(query_resp.segments.is_empty()); } #[tokio::test] @@ -582,22 +607,85 @@ mod tests { .query_prefix_blocks(Request::new(QueryPrefixBlocksRequest { namespace: namespace.into(), block_hashes: vec![h1, h2, h3, h4], + exclude_node: String::new(), })) .await .unwrap() .into_inner(); - let mut nodes: Vec<(String, u32)> = response - .nodes + let segments: Vec<(String, u32)> = response + .segments .into_iter() - .map(|entry| (entry.node, entry.prefix_len)) + .map(|entry| (entry.node, entry.block_count)) .collect(); - nodes.sort(); - // node-a owns h1..h4 (prefix 4), node-b owns h1..h3 (prefix 3) + assert_eq!(segments, vec![(node_a.to_string(), 4)]); + } + + fn prefix_entry(hash: u8, nodes: &[&str]) -> PrefixEntry { + PrefixEntry { + block_hash: vec![hash], + nodes: nodes.iter().map(|node| Arc::::from(*node)).collect(), + } + } + + fn planned(entries: &[PrefixEntry], exclude_node: &str) -> Vec<(String, u32)> { + plan_fetch_segments(entries, exclude_node) + .expect("small test plan should fit uint32") + .into_iter() + .map(|segment| (segment.node, segment.block_count)) + .collect() + } + + #[test] + fn planner_combines_fragmented_remote_prefix() { + let entries = vec![ + prefix_entry(1, &["node-a"]), + prefix_entry(2, &["node-a"]), + prefix_entry(3, &["node-b"]), + prefix_entry(4, &["node-b"]), + ]; + + assert_eq!( + planned(&entries, "requester"), + vec![("node-a".into(), 2), ("node-b".into(), 2)] + ); + } + + #[test] + fn planner_chooses_farthest_owner_and_stable_tie_break() { + let entries = vec![ + prefix_entry(1, &["node-c", "node-b", "node-a"]), + prefix_entry(2, &["node-c", "node-b", "node-a"]), + prefix_entry(3, &["node-c"]), + ]; + + assert_eq!(planned(&entries, "requester"), vec![("node-c".into(), 3)]); assert_eq!( - nodes, - vec![(node_a.to_string(), 4), (node_b.to_string(), 3),] + planned(&entries[..2], "requester"), + vec![("node-a".into(), 2)] ); } + + #[test] + fn planner_excludes_requester_and_stops_at_remote_gap() { + let entries = vec![ + prefix_entry(1, &["requester", "node-a"]), + prefix_entry(2, &["requester"]), + prefix_entry(3, &["node-b"]), + ]; + + assert_eq!(planned(&entries, "requester"), vec![("node-a".into(), 1)]); + } + + #[test] + fn planner_keeps_single_owner_prefix_in_one_segment() { + let entries = vec![ + prefix_entry(1, &["node-a"]), + prefix_entry(2, &["node-a"]), + prefix_entry(3, &["node-a"]), + ]; + + assert_eq!(planned(&entries, "requester"), vec![("node-a".into(), 3)]); + } } diff --git a/pegaflow-metaserver/src/store.rs b/pegaflow-metaserver/src/store.rs index 16d93470..7a4171a8 100644 --- a/pegaflow-metaserver/src/store.rs +++ b/pegaflow-metaserver/src/store.rs @@ -1,8 +1,11 @@ -use dashmap::{DashMap, mapref::entry::Entry}; +use dashmap::DashMap; use log::{info, warn}; use pegaflow_common::BlockKey; -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; +use std::collections::{HashMap, HashSet}; +use std::sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, +}; use std::time::{Duration, Instant}; use uuid::Uuid; @@ -10,6 +13,7 @@ const MIN_RECLAIMABLE_OWNER_COUNT: usize = 3; pub const DEFAULT_NODE_STALE_SECS: u64 = 30; pub const DEFAULT_TTL_MINUTES: u64 = 120; +pub const MANUAL_CLEANUP_AGE_SECS: u64 = 60 * 60; /// A prefix query result: one block hash and all live nodes that own it. #[derive(Debug, Clone)] @@ -46,11 +50,13 @@ impl SweepStats { } } -/// Live-owner redundancy distribution over block keys, recomputed each sweep and -/// cached so metric scrapes stay O(1). Keys are bucketed by their number of -/// query-visible owners (1, 2, 3, >=4); keys with zero visible owners are -/// excluded from every bucket. `copies` is the exact total of visible owners, so -/// average redundancy (the cache capacity shrink factor) is +/// Owner redundancy distribution over block keys, maintained incrementally so +/// metric scrapes stay O(1). Node liveness transitions are settled by the +/// periodic sweep, so a scrape before that sweep can briefly include a stale +/// node. Keys are bucketed by their number of accounted owners (1, 2, 3, >=4); +/// keys with zero accounted owners are excluded from every bucket. `copies` is +/// the exact total of accounted owners, so average redundancy (the cache +/// capacity shrink factor) is /// `copies / (keys_1 + keys_2 + keys_3 + keys_4plus)`. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct RedundancySnapshot { @@ -61,19 +67,6 @@ pub struct RedundancySnapshot { pub copies: u64, } -impl RedundancySnapshot { - fn record(&mut self, visible_owners: u64) { - match visible_owners { - 0 => {} - 1 => self.keys_1 += 1, - 2 => self.keys_2 += 1, - 3 => self.keys_3 += 1, - _ => self.keys_4plus += 1, - } - self.copies += visible_owners; - } -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum StoreError { UnknownNode, @@ -96,13 +89,61 @@ struct NodeRecord { last_seen: Instant, } -/// Both lifecycle decisions for one owner record, produced from a single `nodes` -/// lookup: `keep` (survives the TTL purge) and `visible` (query-visible: current -/// session and node still fresh). Lets the sweep tally live redundancy without -/// probing the node map twice per owner. -struct OwnerEval { - keep: bool, - visible: bool, +#[derive(Default)] +struct RedundancyCounters { + keys_1: AtomicU64, + keys_2: AtomicU64, + keys_3: AtomicU64, + keys_4plus: AtomicU64, + copies: AtomicU64, +} + +impl RedundancyCounters { + fn snapshot(&self) -> RedundancySnapshot { + RedundancySnapshot { + keys_1: self.keys_1.load(Ordering::Relaxed), + keys_2: self.keys_2.load(Ordering::Relaxed), + keys_3: self.keys_3.load(Ordering::Relaxed), + keys_4plus: self.keys_4plus.load(Ordering::Relaxed), + copies: self.copies.load(Ordering::Relaxed), + } + } + + fn adjust_bucket(&self, count: u64, delta: i64) { + let counter = match count { + 1 => &self.keys_1, + 2 => &self.keys_2, + 3 => &self.keys_3, + _ if count >= 4 => &self.keys_4plus, + _ => return, + }; + if delta > 0 { + counter.fetch_add(delta as u64, Ordering::Relaxed); + } else { + counter.fetch_sub((-delta) as u64, Ordering::Relaxed); + } + } + + fn adjust(&self, before: u64, after: u64) { + if before == after { + return; + } + self.adjust_bucket(before, -1); + self.adjust_bucket(after, 1); + if after > before { + self.copies.fetch_add(after - before, Ordering::Relaxed); + } else { + self.copies.fetch_sub(before - after, Ordering::Relaxed); + } + } + + fn reset(&self) { + self.keys_1.store(0, Ordering::Relaxed); + self.keys_2.store(0, Ordering::Relaxed); + self.keys_3.store(0, Ordering::Relaxed); + self.keys_4plus.store(0, Ordering::Relaxed); + self.copies.store(0, Ordering::Relaxed); + } } /// Async thread-safe block hash storage using DashMap. @@ -111,11 +152,14 @@ struct OwnerEval { /// the current MetaServer session and liveness for each node URL. pub struct BlockHashStore { blocks: DashMap, OwnerRecord>>, + node_blocks: DashMap, HashSet>, nodes: DashMap, NodeRecord>, config: StoreConfig, - /// Latest live-owner redundancy distribution, refreshed each sweep and read - /// by metric callbacks. Decouples the O(N) scan from the scrape path. - redundancy: Mutex, + /// Serializes metadata mutations so per-key visibility transitions and + /// aggregate redundancy counters stay consistent with one another. + mutation_lock: Mutex<()>, + /// Incremental live-owner redundancy counters read by metric callbacks. + redundancy: RedundancyCounters, } impl BlockHashStore { @@ -126,9 +170,11 @@ impl BlockHashStore { pub fn with_config(config: StoreConfig) -> Self { Self { blocks: DashMap::new(), + node_blocks: DashMap::new(), nodes: DashMap::new(), config, - redundancy: Mutex::new(RedundancySnapshot::default()), + mutation_lock: Mutex::new(()), + redundancy: RedundancyCounters::default(), } } @@ -144,61 +190,99 @@ impl BlockHashStore { } pub fn heartbeat_node(&self, node: &str, node_id: Uuid) -> Result<(), StoreError> { + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); let now = Instant::now(); - match self.nodes.entry(Arc::from(node)) { - Entry::Vacant(entry) => { - info!("MetaServer node registered: node={node} node_id={node_id}"); - entry.insert(NodeRecord { + let Some(current) = self.nodes.get(node).map(|record| record.clone()) else { + info!("MetaServer node registered: node={node} node_id={node_id}"); + self.nodes.insert( + Arc::from(node), + NodeRecord { node_id, last_seen: now, - }); - Ok(()) + }, + ); + return Ok(()); + }; + + let same_session = current.node_id == node_id; + let stale_session = now.duration_since(current.last_seen) > self.config.node_stale_after; + if same_session || stale_session { + if stale_session && !same_session { + info!( + "MetaServer node session takeover: node={} old_node_id={} new_node_id={}", + node, current.node_id, node_id + ); } - Entry::Occupied(mut entry) => { - let record = entry.get_mut(); - let same_session = record.node_id == node_id; - let stale_session = - now.duration_since(record.last_seen) > self.config.node_stale_after; - if same_session || stale_session { - if stale_session && !same_session { - info!( - "MetaServer node session takeover: node={} old_node_id={} new_node_id={}", - node, record.node_id, node_id - ); - } - record.node_id = node_id; - record.last_seen = now; - return Ok(()); + let keys = self + .node_blocks + .get(node) + .map(|keys| keys.iter().cloned().collect::>()) + .unwrap_or_default(); + let before: Vec<(BlockKey, u64)> = keys + .iter() + .filter_map(|key| { + self.blocks.get(key).map(|owners| { + ( + key.clone(), + self.visible_owner_count_before_heartbeat( + &owners, + node, + current.node_id, + now, + ), + ) + }) + }) + .collect(); + + let mut record = self + .nodes + .get_mut(node) + .expect("node disappeared under mutation lock"); + record.node_id = node_id; + record.last_seen = now; + drop(record); + + for (key, before_visible) in before { + if let Some(owners) = self.blocks.get(&key) { + let after_visible = self.accounted_owner_count(&owners); + self.redundancy.adjust(before_visible, after_visible); } - warn!( - "MetaServer heartbeat rejected stale session: node={} current_node_id={} rejected_node_id={}", - node, record.node_id, node_id - ); - Err(StoreError::StaleSession) } + return Ok(()); } + + warn!( + "MetaServer heartbeat rejected stale session: node={} current_node_id={} rejected_node_id={}", + node, current.node_id, node_id + ); + Err(StoreError::StaleSession) } pub fn unregister_node(&self, node: &str, node_id: Uuid) -> Result { - if self - .nodes - .remove_if(node, |_, record| record.node_id == node_id) - .is_none() - { - if self.nodes.contains_key(node) { - warn!( - "MetaServer unregister rejected stale session: node={} rejected_node_id={}", - node, node_id - ); - return Err(StoreError::StaleSession); - } + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); + let Some(record) = self.nodes.get(node) else { warn!( "MetaServer unregister rejected unknown node: node={} node_id={}", node, node_id ); return Err(StoreError::UnknownNode); + }; + if record.node_id != node_id { + warn!( + "MetaServer unregister rejected stale session: node={} rejected_node_id={}", + node, node_id + ); + return Err(StoreError::StaleSession); } - Ok(self.remove_node_owners(node, node_id)) + drop(record); + + // Remove owners while the node record is still present so the + // redundancy counters observe the same visibility as queries. + let removed = self.remove_node_owners(node, node_id); + self.nodes + .remove_if(node, |_, record| record.node_id == node_id); + Ok(removed) } pub fn insert_hashes( @@ -208,6 +292,7 @@ impl BlockHashStore { node: &str, node_id: Uuid, ) -> Result>, StoreError> { + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); self.touch_node_session(node, node_id)?; let node: Arc = Arc::from(node); let now = Instant::now(); @@ -215,6 +300,7 @@ impl BlockHashStore { for hash in hashes { let key = BlockKey::new(namespace.to_string(), hash.clone()); let mut owners = self.blocks.entry(key).or_default(); + let before_visible = self.accounted_owner_count(&owners); let previous = owners.insert( Arc::clone(&node), OwnerRecord { @@ -233,6 +319,14 @@ impl BlockHashStore { { reclaimable_hashes.push(hash.clone()); } + let after_visible = self.accounted_owner_count(&owners); + self.redundancy.adjust(before_visible, after_visible); + let key = BlockKey::new(namespace.to_string(), hash.clone()); + drop(owners); + self.node_blocks + .entry(Arc::clone(&node)) + .or_default() + .insert(key); } Ok(reclaimable_hashes) } @@ -244,23 +338,37 @@ impl BlockHashStore { node: &str, node_id: Uuid, ) -> Result { + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); self.touch_node_session(node, node_id)?; + let node_key: Arc = Arc::from(node); let mut removed = 0; for hash in hashes { let key = BlockKey::new(namespace.to_string(), hash.clone()); - let should_remove_key = if let Some(mut owners) = self.blocks.get_mut(&key) { + let change = if let Some(mut owners) = self.blocks.get_mut(&key) { + let before_visible = self.accounted_owner_count(&owners); if owners .get(node) .is_some_and(|owner| owner.node_id == node_id) { owners.remove(node); removed += 1; + let after_visible = self.accounted_owner_count(&owners); + Some((before_visible, after_visible)) + } else { + None } - owners.is_empty() } else { - false + None }; - if should_remove_key { + if let Some((before_visible, after_visible)) = change { + self.remove_reverse_index_key(&node_key, &key); + self.redundancy.adjust(before_visible, after_visible); + } + if self + .blocks + .get(&key) + .is_some_and(|owners| owners.is_empty()) + { self.blocks.remove_if(&key, |_, owners| owners.is_empty()); } } @@ -300,62 +408,118 @@ impl BlockHashStore { result } - /// Sweep owners whose node is missing or whose ownership TTL has expired, and - /// refresh the cached live-owner redundancy snapshot in the same walk. + /// Sweep owners belonging to nodes that are missing or no longer active. + /// The reverse index limits work to keys owned by those nodes. pub fn sweep_expired(&self) -> SweepStats { + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); let now = Instant::now(); let mut stats = SweepStats::default(); - let mut snapshot = RedundancySnapshot::default(); - - self.blocks.retain(|_, owners| { - let before = owners.len(); - let mut visible = 0u64; - owners.retain(|node, owner| { - let eval = self.eval_owner(node, owner, now); - if eval.keep && eval.visible { - visible += 1; + let stale_nodes: Vec<(Arc, Uuid, u64)> = self + .nodes + .iter() + .filter_map(|entry| { + let age = now.duration_since(entry.last_seen); + (age > self.config.node_stale_after) + .then(|| (Arc::clone(entry.key()), entry.node_id, age.as_secs())) + }) + .collect(); + + for (node, node_id, last_seen_age_secs) in stale_nodes { + let keys = self + .node_blocks + .get(&node) + .map(|keys| keys.iter().cloned().collect::>()) + .unwrap_or_default(); + for key in keys { + if let Some((before_visible, after_visible, empty)) = + self.remove_any_owner_for_key(&key, &node) + { + stats.removed_owners += 1; + stats.removed_keys += usize::from(empty); + self.redundancy.adjust(before_visible, after_visible); + self.remove_reverse_index_key(&node, &key); + } else if self.owner_for_node(&key, &node).is_none() { + // Clean an index entry left behind by an earlier remove. + // Keep it when a newer session has re-registered the owner. + self.remove_reverse_index_key(&node, &key); } - eval.keep - }); - stats.removed_owners += before.saturating_sub(owners.len()); - if owners.is_empty() { - stats.removed_keys += 1; - return false; } - snapshot.record(visible); - true - }); - let node_before = self.nodes.len(); - let ttl = self.config.ttl; - self.nodes.retain(|node, record| { - let keep = now.duration_since(record.last_seen) <= ttl; - if !keep { + if self + .nodes + .remove_if(&node, |_, record| record.node_id == node_id) + .is_some() + { + stats.removed_nodes += 1; info!( "MetaServer node swept: node={} node_id={} last_seen_age_secs={}", - node, - record.node_id, - now.duration_since(record.last_seen).as_secs() + node, node_id, last_seen_age_secs ); } - keep - }); - stats.removed_nodes = node_before.saturating_sub(self.nodes.len()); + } + + let empty_nodes: Vec> = self + .node_blocks + .iter() + .filter(|entry| entry.is_empty()) + .map(|entry| Arc::clone(entry.key())) + .collect(); + for node in empty_nodes { + self.node_blocks.remove_if(&node, |_, keys| keys.is_empty()); + } + + stats + } - *self - .redundancy - .lock() - .expect("redundancy snapshot mutex poisoned") = snapshot; + /// Remove ownership records older than `max_age`, regardless of node + /// liveness. This is reserved for explicit operator maintenance; the + /// periodic lifecycle sweep intentionally does not use owner age. + pub fn remove_owners_older_than(&self, max_age: Duration) -> SweepStats { + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); + let now = Instant::now(); + let mut stats = SweepStats::default(); + + let keys: Vec = self + .blocks + .iter() + .map(|entry| entry.key().clone()) + .collect(); + for key in keys { + let mut removed_nodes = Vec::new(); + if let Some(mut owners) = self.blocks.get_mut(&key) { + let before_visible = self.accounted_owner_count(&owners); + owners.retain(|node, owner| { + if now.duration_since(owner.key_register_time) > max_age { + removed_nodes.push(Arc::clone(node)); + false + } else { + true + } + }); + let after_visible = self.accounted_owner_count(&owners); + stats.removed_owners += removed_nodes.len(); + self.redundancy.adjust(before_visible, after_visible); + drop(owners); + } + for node in removed_nodes { + self.remove_reverse_index_key(&node, &key); + } + if self + .blocks + .get(&key) + .is_some_and(|owners| owners.is_empty()) + { + self.blocks.remove_if(&key, |_, owners| owners.is_empty()); + stats.removed_keys += 1; + } + } stats } - /// Latest cached live-owner redundancy distribution (refreshed each sweep). + /// Latest incrementally maintained live-owner redundancy distribution. pub fn redundancy_snapshot(&self) -> RedundancySnapshot { - *self - .redundancy - .lock() - .expect("redundancy snapshot mutex poisoned") + self.redundancy.snapshot() } pub fn entry_count(&self) -> u64 { @@ -377,7 +541,7 @@ impl BlockHashStore { let age = now.duration_since(node.last_seen); if age <= self.config.node_stale_after { active += 1; - } else if age <= self.config.ttl { + } else { stale += 1; } } @@ -389,12 +553,11 @@ impl BlockHashStore { reason = "maintenance API reserved for explicit store cleanup" )] pub fn invalidate_all(&self) { + let _mutation_guard = self.mutation_lock.lock().expect("mutation lock poisoned"); self.blocks.clear(); + self.node_blocks.clear(); self.nodes.clear(); - *self - .redundancy - .lock() - .expect("redundancy snapshot mutex poisoned") = RedundancySnapshot::default(); + self.redundancy.reset(); } fn touch_node_session(&self, node: &str, node_id: Uuid) -> Result<(), StoreError> { @@ -417,40 +580,121 @@ impl BlockHashStore { } fn remove_node_owners(&self, node: &str, node_id: Uuid) -> usize { + let node: Arc = Arc::from(node); + let keys = self + .node_blocks + .get(&node) + .map(|keys| keys.iter().cloned().collect::>()) + .unwrap_or_default(); let mut removed = 0; - self.blocks.retain(|_, owners| { - if owners - .get(node) - .is_some_and(|owner| owner.node_id == node_id) + for key in keys { + if let Some((before_visible, after_visible, _empty)) = + self.remove_owner_for_key(&key, &node, node_id) { - owners.remove(node); removed += 1; + self.redundancy.adjust(before_visible, after_visible); + self.remove_reverse_index_key(&node, &key); } - !owners.is_empty() - }); + } removed } - /// Evaluate one owner against the current node session and clocks with a - /// single `nodes` lookup. A missing node means the owner is neither kept nor - /// visible. - fn eval_owner(&self, node: &Arc, owner: &OwnerRecord, now: Instant) -> OwnerEval { - let Some(record) = self.nodes.get(node.as_ref()) else { - return OwnerEval { - keep: false, - visible: false, - }; - }; - let node_age = now.duration_since(record.last_seen); - OwnerEval { - keep: now.duration_since(owner.key_register_time) <= self.config.ttl - && node_age <= self.config.ttl, - visible: record.node_id == owner.node_id && node_age <= self.config.node_stale_after, + /// Counter visibility intentionally ignores node age. A stale node remains + /// represented until the lifecycle sweep removes it, so ordinary mutations + /// do not make aggregate counters disagree with the records they describe. + fn accounted_owner_count(&self, owners: &HashMap, OwnerRecord>) -> u64 { + owners + .iter() + .filter(|(node, owner)| { + self.nodes + .get(node.as_ref()) + .is_some_and(|record| record.node_id == owner.node_id) + }) + .count() as u64 + } + + fn visible_owner_count_before_heartbeat( + &self, + owners: &HashMap, OwnerRecord>, + node: &str, + node_id: Uuid, + now: Instant, + ) -> u64 { + owners + .iter() + .filter(|(owner_node, owner)| { + if owner_node.as_ref() == node { + owner.node_id == node_id + } else { + self.is_owner_visible(owner_node, owner, now) + } + }) + .count() as u64 + } + + fn remove_owner_for_key( + &self, + key: &BlockKey, + node: &Arc, + node_id: Uuid, + ) -> Option<(u64, u64, bool)> { + let mut owners = self.blocks.get_mut(key)?; + if !owners + .get(node) + .is_some_and(|owner| owner.node_id == node_id) + { + return None; + } + let before_visible = self.accounted_owner_count(&owners); + owners.remove(node); + let after_visible = self.accounted_owner_count(&owners); + let empty = owners.is_empty(); + drop(owners); + if empty { + self.blocks.remove_if(key, |_, owners| owners.is_empty()); + } + Some((before_visible, after_visible, empty)) + } + + fn remove_any_owner_for_key( + &self, + key: &BlockKey, + node: &Arc, + ) -> Option<(u64, u64, bool)> { + let mut owners = self.blocks.get_mut(key)?; + if !owners.contains_key(node) { + return None; + } + let before_visible = self.accounted_owner_count(&owners); + owners.remove(node); + let after_visible = self.accounted_owner_count(&owners); + let empty = owners.is_empty(); + drop(owners); + if empty { + self.blocks.remove_if(key, |_, owners| owners.is_empty()); } + Some((before_visible, after_visible, empty)) + } + + fn owner_for_node(&self, key: &BlockKey, node: &Arc) -> Option { + self.blocks + .get(key) + .and_then(|owners| owners.get(node).map(|owner| owner.node_id)) + } + + fn remove_reverse_index_key(&self, node: &Arc, key: &BlockKey) { + if let Some(mut keys) = self.node_blocks.get_mut(node) { + keys.remove(key); + } + self.node_blocks.remove_if(node, |_, keys| keys.is_empty()); } fn is_owner_visible(&self, node: &Arc, owner: &OwnerRecord, now: Instant) -> bool { - self.eval_owner(node, owner, now).visible + let Some(record) = self.nodes.get(node.as_ref()) else { + return false; + }; + let node_age = now.duration_since(record.last_seen); + record.node_id == owner.node_id && node_age <= self.config.node_stale_after } } @@ -917,6 +1161,58 @@ mod tests { assert_eq!(store.entry_count(), 2); } + #[test] + fn test_sweep_keeps_old_owner_on_active_node() { + let store = BlockHashStore::with_config(StoreConfig { + node_stale_after: Duration::from_secs(60), + ttl: Duration::from_secs(1), + }); + let node_id = heartbeat_node(&store, "node-a"); + store + .insert_hashes("ns", &[vec![1]], "node-a", node_id) + .unwrap(); + store + .blocks + .get_mut(&BlockKey::new("ns".to_string(), vec![1])) + .unwrap() + .get_mut("node-a") + .unwrap() + .key_register_time = Instant::now() - Duration::from_secs(2); + + let removed = store.sweep_expired(); + assert_eq!(removed, SweepStats::default()); + assert_eq!(store.owner_count(), 1); + assert_eq!(store.query_prefix("ns", &[vec![1]]).len(), 1); + } + + #[test] + fn test_remove_owners_older_than_removes_only_expired_owners() { + let store = BlockHashStore::new(); + let node_id = heartbeat_node(&store, "node-a"); + store + .insert_hashes("ns", &[vec![1], vec![2]], "node-a", node_id) + .unwrap(); + store + .blocks + .get_mut(&BlockKey::new("ns".to_string(), vec![1])) + .unwrap() + .get_mut("node-a") + .unwrap() + .key_register_time = Instant::now() - Duration::from_secs(3601); + + let removed = store.remove_owners_older_than(Duration::from_secs(3600)); + assert_eq!( + removed, + SweepStats { + removed_owners: 1, + removed_keys: 1, + removed_nodes: 0, + } + ); + assert_eq!(store.owner_count(), 1); + assert_eq!(store.entry_count(), 1); + } + #[test] fn test_concurrent_insert_and_remove() { use std::sync::Arc; @@ -1034,4 +1330,41 @@ mod tests { assert_eq!(store.owner_count(), 1, "raw record still present"); assert_eq!(store.redundancy_snapshot(), RedundancySnapshot::default()); } + + #[test] + fn test_redundancy_counters_follow_owner_mutations() { + let store = BlockHashStore::new(); + let node_a = heartbeat_node(&store, "node-a"); + let node_b = heartbeat_node(&store, "node-b"); + let hash = vec![1]; + + store + .insert_hashes("ns", std::slice::from_ref(&hash), "node-a", node_a) + .unwrap(); + assert_eq!(store.redundancy_snapshot().keys_1, 1); + + store + .insert_hashes("ns", std::slice::from_ref(&hash), "node-b", node_b) + .unwrap(); + assert_eq!(store.redundancy_snapshot().keys_2, 1); + assert_eq!(store.redundancy_snapshot().copies, 2); + + store + .remove_hashes("ns", std::slice::from_ref(&hash), "node-a", node_a) + .unwrap(); + assert_eq!(store.redundancy_snapshot().keys_1, 1); + assert_eq!(store.redundancy_snapshot().keys_2, 0); + assert_eq!(store.redundancy_snapshot().copies, 1); + + store + .blocks + .get_mut(&BlockKey::new("ns".to_string(), hash)) + .unwrap() + .get_mut("node-b") + .unwrap() + .key_register_time = Instant::now() - Duration::from_secs(3601); + let cleanup = store.remove_owners_older_than(Duration::from_secs(3600)); + assert_eq!(cleanup.removed_owners, 1); + assert_eq!(store.redundancy_snapshot(), RedundancySnapshot::default()); + } } diff --git a/pegaflow-proto/proto/engine.proto b/pegaflow-proto/proto/engine.proto index 3a039433..e3bfa8d5 100644 --- a/pegaflow-proto/proto/engine.proto +++ b/pegaflow-proto/proto/engine.proto @@ -46,6 +46,13 @@ message RegisterContextRequest { // Set by the connector for MLA; the engine derives per-layer page offsets // from the registered layouts. See spec.md. bool page_first = 18; + // Hybrid-cache storage group per layer, parallel to layer_names (e.g. + // vLLM's kv-cache group index: 0 = full attention, 1 = recurrent state). + // Empty means every layer belongs to group 0, preserving single-group + // behavior. Groups key blocks independently and seal against their own + // slot space, so a sparse (checkpoint-only) group can seal without + // attention's per-block cadence. Page-first stays single-group. + repeated uint32 layer_group_ids = 19; } message RegisterContextResponse { @@ -110,8 +117,20 @@ message QueryRequest { string req_id = 3; // Keep the query in Loading until the full remaining prefix is fetchable // from a remote node via MetaServer + RDMA. Does not observe saves landing - // in the local engine; no effect when RDMA is not configured. + // in the local engine; no effect when RDMA is not configured. On groups > 0 + // this switches the membership query to an all-or-nothing set fetch: the + // hash list is an exact want-set, misses are pulled from SSD / remote + // peers, and a Ready answer shorter than the request means the set could + // not be completed anywhere: callers treat it as a miss, and it carries no + // lease (the short hit count is reported for observability only). bool wait_for_full_prefix = 4; + // Hybrid-cache storage group to query. 0 (default) is the attention group + // and uses prefix semantics (hit count + lease over the contiguous prefix, + // SSD/RDMA prefetch eligible). Groups > 0 use membership semantics: every + // position reports independently (QueryReady.hit_positions), hits are not + // prefix-clamped, and the lease pins exactly the hit blocks; local-only + // unless wait_for_full_prefix requests the remote-fetching set form. + uint32 group_id = 5; } message QueryResponse { @@ -124,8 +143,18 @@ message QueryResponse { message QueryLoading {} message QueryReady { + // Prefix query (group_id = 0): number of leading hit blocks. + // Membership query (group_id > 0): number of hit positions in total. uint64 num_hit_blocks = 1; + // Lease over the hit blocks: the prefix for group 0, or exactly the blocks + // at hit_positions (in that order) for membership queries. Empty on 0 hits. bytes lease = 2; + // Membership queries only: indices into QueryRequest.block_hashes whose + // block is cached, in ascending order. The lease's i-th block corresponds + // to query position hit_positions[i]. Position-aligned so hybrid callers + // can pick e.g. the rightmost recurrent checkpoint within the attention + // prefix without a second round trip. + repeated uint32 hit_positions = 3; } message ReleaseRequest { @@ -276,15 +305,18 @@ message RemoveBlockHashesResponse { message QueryPrefixBlocksRequest { string namespace = 1; repeated bytes block_hashes = 2; // ordered by prefix position + string exclude_node = 3; // requester address; never return it as an owner } -message NodePrefixResult { +message FetchSegment { string node = 1; - uint32 prefix_len = 2; + uint32 block_count = 2; } message QueryPrefixBlocksResponse { - repeated NodePrefixResult nodes = 1; + reserved 1; + reserved "nodes"; + repeated FetchSegment segments = 2; } message HeartbeatNodeRequest { diff --git a/pegaflow-server/src/lib.rs b/pegaflow-server/src/lib.rs index e5117d20..97321513 100644 --- a/pegaflow-server/src/lib.rs +++ b/pegaflow-server/src/lib.rs @@ -174,11 +174,11 @@ pub struct Cli { /// HLL sliding-window list for hit-rate estimation. Comma-separated humantime /// durations; each becomes a canonical `window` label in metrics (e.g. `15m,1h,1d`). /// Slot duration is derived as `clamp(window/24, 1min, 1h)`. - #[arg(long, default_value = "15m,1h,24h", value_parser = parse_hll_windows_arg)] + #[arg(long, default_value = "15m,1h,1d", value_parser = parse_hll_windows_arg)] pub metric_hll_windows: String, - /// HLL bucket index bits 4–18 (default: 14 → 16384 buckets, ~0.8% error) - #[arg(long, default_value_t = 14, value_parser = parse_hll_bucket_bits)] + /// HLL bucket index bits 4–18 (default: 16 → 65536 buckets, ~0.4% error) + #[arg(long, default_value_t = 16, value_parser = parse_hll_bucket_bits)] pub metric_hll_bucket_bits: u8, /// Transfer lock timeout in seconds. Blocks held for cross-node RDMA transfer are @@ -211,7 +211,7 @@ fn parse_hll_windows_arg(s: &str) -> Result { Ok(s.to_string()) } -/// Parse a comma-separated list of humantime windows (e.g. `15m,1h,24h`). +/// Parse a comma-separated list of humantime windows (e.g. `15m,1h,1d`). /// Each entry becomes `(label, duration)` where label is canonicalized from /// the parsed duration. fn parse_hll_windows(s: &str) -> Result, String> { @@ -732,6 +732,7 @@ mod tests { parse_hll_windows(&cli.metric_hll_windows).unwrap(), expected_hll_windows() ); + assert_eq!(cli.metric_hll_bucket_bits, 16); } #[test] diff --git a/pegaflow-server/src/metric.rs b/pegaflow-server/src/metric.rs index fc4b8f8b..920edf39 100644 --- a/pegaflow-server/src/metric.rs +++ b/pegaflow-server/src/metric.rs @@ -8,6 +8,7 @@ use tonic::Status; struct HllGaugeHandles { _cardinality: ObservableGauge, _total_requests: ObservableGauge, + _estimated_hit_rate: ObservableGauge, } static HLL_GAUGES: OnceLock = OnceLock::new(); @@ -15,18 +16,20 @@ static HLL_GAUGES: OnceLock = OnceLock::new(); /// Register HLL observable gauges backed by the multi-window tracker. /// /// Emits one time series per configured window, labeled with `window=