diff --git a/asap-precompute-rs/Cargo.lock b/asap-precompute-rs/Cargo.lock index aec0f19..79e6252 100644 --- a/asap-precompute-rs/Cargo.lock +++ b/asap-precompute-rs/Cargo.lock @@ -381,8 +381,8 @@ dependencies = [ [[package]] name = "asap_sketchlib" -version = "0.2.2" -source = "git+https://github.com/ProjectASAP/asap_sketchlib?rev=010457de46ba9ec5574a53476bbe8b4bdb02b965#010457de46ba9ec5574a53476bbe8b4bdb02b965" +version = "0.3.0" +source = "git+https://github.com/ProjectASAP/asap_sketchlib?rev=3fc3442c3275d374f79ba522f473de83912c6345#3fc3442c3275d374f79ba522f473de83912c6345" dependencies = [ "bytes", "prost 0.13.5", @@ -5443,7 +5443,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "petgraph", @@ -5464,7 +5464,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.118", @@ -5477,7 +5477,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.118", diff --git a/asap-precompute-rs/Cargo.toml b/asap-precompute-rs/Cargo.toml index 71e7512..302ff29 100644 --- a/asap-precompute-rs/Cargo.toml +++ b/asap-precompute-rs/Cargo.toml @@ -7,7 +7,7 @@ license = "MIT" # asap_sketchlib is a public dependency. [dependencies] -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "010457de46ba9ec5574a53476bbe8b4bdb02b965" } +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "3fc3442c3275d374f79ba522f473de83912c6345" } serde = { version = "1", features = ["derive"] } prost = "0.13" # OTAP's generated protobuf types implement the 0.14 Message trait. diff --git a/asap-precompute-rs/src/bin/asap_stream_bench.rs b/asap-precompute-rs/src/bin/asap_stream_bench.rs index 892eda4..958cd8c 100644 --- a/asap-precompute-rs/src/bin/asap_stream_bench.rs +++ b/asap-precompute-rs/src/bin/asap_stream_bench.rs @@ -4,6 +4,9 @@ mod model; #[path = "stream_bench/profile.rs"] mod profile; +#[path = "stream_bench/topk.rs"] +mod topk; +use asap_precompute_rs::otap::codec::decode_pdata_to_observations; use async_trait::async_trait; use linkme::distributed_slice; use otel_arrow_dfe_config::{ @@ -89,7 +92,7 @@ impl Options { signals_per_second: number("ASAP_SIGNALS_PER_SECOND", "100000")?, rolling_windows: number("ASAP_ROLLING_WINDOWS", "1")?, }; - if !["raw", "exact", "kll"].contains(&o.scenario.as_str()) + if !["raw", "exact", "kll", "topk-exact", "cmsheap"].contains(&o.scenario.as_str()) || !["generator", "branch", "merge", "estimate", "backend"].contains(&o.role.as_str()) || o.points == 0 || o.batch == 0 @@ -117,6 +120,8 @@ struct Stats { output_messages: u64, window_latency_ms: BTreeMap, latest_results: Vec, + latest_topk: Vec, + topk_accuracy: Vec, } struct Shared { stats: Mutex, @@ -223,8 +228,33 @@ fn start_control() -> Result<(), String> { }); Ok(()) } +/// Quantile scenarios (raw/exact/kll) or keyed top-k scenarios (topk-exact/cmsheap). +enum Stage { + Quantile(model::Processor), + TopK(topk::Processor), +} +impl Stage { + fn process(&mut self, pdata: OtapPdata) -> Result, String> { + match self { + Self::Quantile(p) => p.process(pdata), + Self::TopK(p) => p.process(pdata), + } + } + fn process_normalized( + &mut self, + batch: model::NormalizedBatch, + ) -> Result, String> { + match self { + Self::Quantile(p) => p.process_normalized(batch), + Self::TopK(p) => p.process_normalized(batch), + } + } +} +fn is_topk(scenario: &str) -> bool { + matches!(scenario, "topk-exact" | "cmsheap") +} struct StreamProcessor { - processor: model::Processor, + processor: Stage, normalizer: Option, } #[async_trait(?Send)] @@ -284,13 +314,23 @@ fn create_processor( let o = OPTIONS.get().unwrap(); Ok(ProcessorWrapper::local( StreamProcessor { - processor: model::Processor::new( - &o.scenario, - &o.role, - o.points, - o.batch, - o.rolling_windows, - ), + processor: if is_topk(&o.scenario) { + Stage::TopK(topk::Processor::new( + &o.scenario, + &o.role, + o.points, + o.batch, + o.rolling_windows, + )) + } else { + Stage::Quantile(model::Processor::new( + &o.scenario, + &o.role, + o.points, + o.batch, + o.rolling_windows, + )) + }, normalizer: (o.role == "branch") .then(|| model::InputNormalizer::new(o.source, o.points, o.batch)), }, @@ -306,7 +346,35 @@ static PROCESSOR_FACTORY: ProcessorFactory = ProcessorFactory { wiring_contract: otel_arrow_dfe_engine::wiring_contract::WiringContract::UNRESTRICTED, validate_config: otel_arrow_dfe_config::validation::no_config, }; -struct ValidatingBackend(model::Backend); +enum ValidatingBackend { + Quantile(model::Backend), + TopK(topk::Backend, u64), +} +impl ValidatingBackend { + /// Completed input signals and optional (window id, start timestamp). + fn receive(&mut self, pdata: OtapPdata) -> Result<(u64, Option<(u64, u64)>), String> { + match self { + Self::Quantile(backend) => { + let received = backend.receive(pdata)?; + if received.1.is_some() { + let mut stats = shared().stats.lock().unwrap(); + stats.latest_results.clone_from(&backend.latest); + } + Ok(received) + } + Self::TopK(backend, signals) => { + let obs = decode_pdata_to_observations(pdata) + .map_err(|e| e.to_string())? + .observations; + let completed = backend.receive(&obs)?; + let mut stats = shared().stats.lock().unwrap(); + stats.latest_topk.clone_from(&backend.latest); + stats.topk_accuracy.clone_from(&backend.accuracy); + Ok((*signals, Some(completed))) + } + } + } +} #[async_trait(?Send)] impl exporter::Exporter for ValidatingBackend { async fn start( @@ -319,7 +387,6 @@ impl exporter::Exporter for ValidatingBackend { Message::PData(pdata) => { // Validation failures terminate the benchmark, never increment goodput. let (count, completed) = self - .0 .receive(pdata.clone()) .expect("backend correctness failure"); { @@ -327,7 +394,6 @@ impl exporter::Exporter for ValidatingBackend { stats.received_signals += count; if let Some((window, start)) = completed { assert_eq!(window, stats.completed_windows, "backend window gap"); - stats.latest_results.clone_from(&self.0.latest); stats.completed_windows += 1; let elapsed_ms = now_ns().saturating_sub(start) / 1_000_000; *stats.window_latency_ms.entry(elapsed_ms).or_default() += 1; @@ -351,11 +417,18 @@ fn create_backend( ) -> Result, otel_arrow_dfe_config::error::Error> { let o = OPTIONS.get().unwrap(); Ok(ExporterWrapper::local( - ValidatingBackend(model::Backend::new( - &o.scenario, - o.points, - o.rolling_windows, - )), + if is_topk(&o.scenario) { + ValidatingBackend::TopK( + topk::Backend::new(&o.scenario, o.points, o.rolling_windows), + (o.points * 2) as u64, + ) + } else { + ValidatingBackend::Quantile(model::Backend::new( + &o.scenario, + o.points, + o.rolling_windows, + )) + }, node, config, runtime, diff --git a/asap-precompute-rs/src/bin/stream_bench/model.rs b/asap-precompute-rs/src/bin/stream_bench/model.rs index 0e03090..b125d14 100644 --- a/asap-precompute-rs/src/bin/stream_bench/model.rs +++ b/asap-precompute-rs/src/bin/stream_bench/model.rs @@ -53,7 +53,7 @@ impl Meta { started_ns: get("bench.started_ns")?, }) } - fn envelope(self, value: f64) -> SketchEnvelope { + pub fn envelope(self, value: f64) -> SketchEnvelope { SketchEnvelope { schema_version: 1, sketch_type: SketchType::Unspecified, diff --git a/asap-precompute-rs/src/bin/stream_bench/topk.rs b/asap-precompute-rs/src/bin/stream_bench/topk.rs new file mode 100644 index 0000000..3f5f138 --- /dev/null +++ b/asap-precompute-rs/src/bin/stream_bench/topk.rs @@ -0,0 +1,577 @@ +//! Keyed top-k over the same count-aligned windows as `model.rs`. +//! +//! `topk-exact` keeps a per-key count map per window, ships every entry and +//! sums the maps; `cmsheap` ships one ASAPv1 Count-Min-with-heap sketch per +//! source and merges sketches. Both answer the top `K` keys per window. +use super::model::{Meta, NormalizedBatch, MAX_PENDING}; +use super::profile::{self, Category}; +use asap_precompute_rs::{ + envelope::{Encoding, SketchEnvelope, SketchType}, + observation::{KeyValue, Observation}, + otap::codec::{decode_pdata_to_observations, encode_envelopes_to_pdata}, + precompute::FrequencySketch, + sketches::CMSHeapWrapper, +}; +use otel_arrow_dfe_otap::pdata::OtapPdata; +use serde::Serialize; +use std::collections::{BTreeMap, HashMap, VecDeque}; +use std::hash::{BuildHasherDefault, Hasher}; +use std::sync::OnceLock; + +pub const K: usize = 10; +const KEYS: usize = 1 << 20; +const ZIPF_S: f64 = 1.1; +const CMS_ROWS: usize = 4; +const CMS_COLS: usize = 2048; +const KEY_SEED: u64 = 0x2545_f491_4f6c_dd1d; +/// Minimum recall@K the sketch must reach in every window. +const MIN_RECALL: f64 = 0.9; + +fn mix(mut x: u64) -> u64 { + x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + x ^ (x >> 31) +} + +/// Deterministic Zipf(s=1.1) key over 2^20 ranks; ranks are scrambled so the +/// heaviest keys are not small integers. +pub fn key(index: u64) -> u64 { + static CDF: OnceLock> = OnceLock::new(); + let cdf = CDF.get_or_init(|| { + let mut total = 0.0; + let mut cdf = (1..=KEYS) + .map(|rank| { + total += (rank as f64).powf(-ZIPF_S); + total + }) + .collect::>(); + cdf.iter_mut().for_each(|c| *c /= total); + cdf + }); + let unit = (mix(index.wrapping_add(KEY_SEED)) >> 11) as f64 / (1_u64 << 53) as f64; + let rank = cdf.partition_point(|&c| c < unit).min(KEYS - 1) as u64; + mix(rank ^ KEY_SEED) +} + +/// Keys are already scrambled 64-bit values, so the exact baseline hashes +/// them with the identity rather than SipHash. +#[derive(Default)] +struct KeyHasher(u64); +impl Hasher for KeyHasher { + fn finish(&self) -> u64 { + self.0 + } + fn write(&mut self, _: &[u8]) { + unreachable!("only u64 keys are hashed") + } + fn write_u64(&mut self, key: u64) { + self.0 = key; + } +} +type Counts = HashMap>; + +/// Largest counts first; ties by key so exact answers are deterministic. +fn top(counts: &Counts) -> Vec<(u64, u64)> { + let mut entries = counts.iter().map(|(k, c)| (*k, *c)).collect::>(); + entries.sort_unstable_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0))); + entries.truncate(K); + entries +} + +fn label(o: &Observation, name: &str) -> Result { + o.labels + .iter() + .find(|kv| kv.key == name) + .ok_or_else(|| format!("missing {name}"))? + .value + .parse() + .map_err(|e| format!("{name}: {e}")) +} + +fn counts_pdata(entries: &[(u64, u64)], total: usize, meta: Meta) -> Result { + let _cpu = profile::scope(Category::PdataCodec); + let envelopes = entries + .iter() + .map(|(key, count)| { + let mut env = meta.envelope(*count as f64); + env.labels.push(KeyValue::new("topk.key", key.to_string())); + env.labels + .push(KeyValue::new("bench.total", total.to_string())); + env + }) + .collect::>(); + encode_envelopes_to_pdata(&envelopes).map_err(|e| e.to_string()) +} + +fn sketch_pdata(sketch: &CMSHeapWrapper, count: usize, meta: Meta) -> Result { + use asap_precompute_rs::precompute::Sketch; + let _cpu = profile::scope(Category::PdataCodec); + let mut env: SketchEnvelope = meta.envelope(0.0); + env.sketch_type = SketchType::CountMinSketchWithHeap; + env.encoding = Encoding::Msgpack; + env.payload = { + let _cpu = profile::scope(Category::SketchCodec); + sketch.snapshot().map_err(|e| e.to_string())? + }; + env.count = count as u64; + encode_envelopes_to_pdata(&[env]).map_err(|e| e.to_string()) +} + +fn decode_sketch(o: &Observation, count: usize) -> Result { + let _cpu = profile::scope(Category::SketchCodec); + let env = o.value.envelope.as_ref().ok_or("missing sketch envelope")?; + if env.count != count as u64 || env.sketch_type != SketchType::CountMinSketchWithHeap { + return Err("sketch count/type mismatch".into()); + } + CMSHeapWrapper::from_snapshot(&env.payload).map_err(|e| e.to_string()) +} + +fn sketch_top(sketch: &CMSHeapWrapper) -> Vec<(u64, u64)> { + sketch + .top_k(K) + .into_iter() + .map(|e| { + let key = u64::from_be_bytes(e.key.as_slice().try_into().expect("8-byte key")); + (key, e.count as u64) + }) + .collect() +} + +fn topk_pdata(results: &[(&str, Vec<(u64, u64)>)], meta: Meta) -> Result { + let _cpu = profile::scope(Category::PdataCodec); + let envelopes = results + .iter() + .flat_map(|(kind, entries)| { + entries.iter().enumerate().map(move |(rank, (key, count))| { + let mut env = meta.envelope(*count as f64); + env.metric_name = "http.server.request.topk".into(); + env.labels.push(KeyValue::new("topk.key", key.to_string())); + env.labels + .push(KeyValue::new("topk.rank", (rank + 1).to_string())); + env.labels.push(KeyValue::new("window_kind", *kind)); + env + }) + }) + .collect::>(); + encode_envelopes_to_pdata(&envelopes).map_err(|e| e.to_string()) +} + +struct Window { + /// Signals (branch) or entries (merge/estimate) received per source. + received: [usize; 3], + /// Entries expected per source, once its first chunk arrives. + total: [Option; 3], + counts: Counts, + sketch: Option, + started_ns: u64, +} +impl Window { + fn new(started_ns: u64) -> Self { + Self { + received: [0; 3], + total: [None; 3], + counts: Counts::default(), + sketch: None, + started_ns, + } + } +} + +pub struct Processor { + sketch: bool, + role: String, + points: usize, + batch: usize, + windows: BTreeMap, + next: [u64; 3], + rolling: usize, + count_history: VecDeque, + sketch_history: VecDeque, +} +impl Processor { + pub fn new(scenario: &str, role: &str, points: usize, batch: usize, rolling: usize) -> Self { + Self { + sketch: scenario == "cmsheap", + role: role.into(), + points, + batch, + windows: BTreeMap::new(), + next: [0; 3], + rolling, + count_history: VecDeque::new(), + sketch_history: VecDeque::new(), + } + } + + fn new_sketch() -> CMSHeapWrapper { + CMSHeapWrapper::new(CMS_ROWS, CMS_COLS, K) + } + + /// Branch input: the normalized corpus indices of one source. + pub fn process_normalized(&mut self, batch: NormalizedBatch) -> Result, String> { + let _cpu = profile::scope(Category::ProcessorBookkeeping); + let meta = batch.meta; + if self.role != "branch" || meta.source > 1 || meta.window != self.next[meta.source] { + return Err("duplicate or out-of-order window".into()); + } + let w = self + .windows + .entry(meta.window) + .or_insert_with(|| Window::new(meta.started_ns)); + if meta.offset != w.received[meta.source] || meta.offset + batch.count > self.points { + return Err("duplicate, missing, or overflowing batch".into()); + } + { + let _cpu = profile::scope(Category::Computation); + let keys = (0..batch.count as u64).map(|i| key(batch.first_index + i)); + if self.sketch { + let sketch = w.sketch.get_or_insert_with(Self::new_sketch); + keys.for_each(|k| sketch.update(&k.to_be_bytes(), 1)); + } else { + keys.for_each(|k| *w.counts.entry(k).or_default() += 1); + } + } + w.received[meta.source] += batch.count; + if w.received[meta.source] < self.points { + if self.windows.len() > MAX_PENDING as usize { + return Err("pending window bound exceeded".into()); + } + return Ok(vec![]); + } + self.next[meta.source] += 1; + let w = self.windows.remove(&meta.window).unwrap(); + let out = Meta { + offset: 0, + started_ns: w.started_ns, + ..meta + }; + self.emit(w, out, self.points) + } + + /// Ship a completed window: one sketch, or every count entry in chunks. + fn emit(&self, w: Window, meta: Meta, signals: usize) -> Result, String> { + if self.sketch { + return Ok(vec![sketch_pdata( + w.sketch.as_ref().unwrap(), + signals, + meta, + )?]); + } + let entries = w.counts.into_iter().collect::>(); + let total = entries.len(); + entries + .chunks(self.batch) + .enumerate() + .map(|(i, chunk)| { + counts_pdata( + chunk, + total, + Meta { + offset: i * self.batch, + ..meta + }, + ) + }) + .collect() + } + + /// Merge and estimate input. + pub fn process(&mut self, pdata: OtapPdata) -> Result, String> { + let _cpu = profile::scope(Category::ProcessorBookkeeping); + let obs = { + let _cpu = profile::scope(Category::PdataCodec); + decode_pdata_to_observations(pdata) + .map_err(|e| e.to_string())? + .observations + }; + let meta = Meta::read(obs.first().ok_or("empty batch")?)?; + let estimate = self.role == "estimate"; + if (estimate && meta.source != 2) || (!estimate && meta.source > 1) { + return Err("wrong source for role".into()); + } + if meta.window != self.next[meta.source] { + return Err("duplicate or out-of-order window".into()); + } + let signals = if estimate { + self.points * 2 + } else { + self.points + }; + let w = self + .windows + .entry(meta.window) + .or_insert_with(|| Window::new(meta.started_ns)); + w.started_ns = w.started_ns.min(meta.started_ns); + if meta.offset != w.received[meta.source] { + return Err("duplicate or missing batch".into()); + } + let _cpu = profile::scope(Category::Computation); + if self.sketch { + if obs.len() != 1 { + return Err("expected one sketch".into()); + } + let decoded = decode_sketch(&obs[0], signals)?; + match &mut w.sketch { + Some(sketch) => sketch.merge_from(&decoded).map_err(|e| e.to_string())?, + None => w.sketch = Some(decoded), + } + w.received[meta.source] = 1; + w.total[meta.source] = Some(1); + } else { + let total = label(&obs[0], "bench.total")? as usize; + if *w.total[meta.source].get_or_insert(total) != total + || w.received[meta.source] + obs.len() > total + { + return Err("count entry overflow".into()); + } + for o in &obs { + *w.counts.entry(label(o, "topk.key")?).or_default() += o.value.float as u64; + } + w.received[meta.source] += obs.len(); + } + let done = |s: usize| w.total[s] == Some(w.received[s]); + if done(meta.source) { + self.next[meta.source] += 1; + } + let complete = if estimate { + done(2) + } else { + done(0) && done(1) + }; + if !complete { + if self.windows.len() > MAX_PENDING as usize { + return Err("pending window bound exceeded".into()); + } + return Ok(vec![]); + } + let w = self.windows.remove(&meta.window).unwrap(); + let out = Meta { + source: 2, + offset: 0, + started_ns: w.started_ns, + ..meta + }; + if !estimate { + return self.emit(w, out, self.points * 2); + } + let mut results = vec![("tumbling", Vec::new())]; + if self.sketch { + let sketch = w.sketch.unwrap(); + results[0].1 = sketch_top(&sketch); + if self.rolling > 1 { + self.sketch_history.push_back(sketch); + if self.sketch_history.len() > self.rolling { + self.sketch_history.pop_front(); + } + let mut merged = Self::new_sketch(); + for s in &self.sketch_history { + merged.merge_from(s).map_err(|e| e.to_string())?; + } + results.push(("rolling", sketch_top(&merged))); + } + } else { + results[0].1 = top(&w.counts); + if self.rolling > 1 { + self.count_history.push_back(w.counts); + if self.count_history.len() > self.rolling { + self.count_history.pop_front(); + } + let mut merged = Counts::default(); + for counts in &self.count_history { + for (k, c) in counts { + *merged.entry(*k).or_default() += c; + } + } + results.push(("rolling", top(&merged))); + } + } + Ok(vec![topk_pdata(&results, out)?]) + } +} + +/// One reported top-k entry and its exact count, for dashboards. +#[derive(Serialize, Clone, Debug, Default)] +pub struct TopkPoint { + pub window_kind: String, + pub rank: usize, + pub key: String, + pub estimate: f64, + pub exact: f64, +} +/// Accuracy of one window's answer against the exact top-k. +#[derive(Serialize, Clone, Debug, Default)] +pub struct TopkAccuracy { + pub window_kind: String, + pub recall: f64, + pub mean_relative_error: f64, +} + +pub struct Backend { + sketch: bool, + rolling: usize, + /// Exact per-key counts of one window; every window replays the corpus. + counts: Counts, + reference: Vec<(u64, u64)>, + next: u64, + pub latest: Vec, + pub accuracy: Vec, +} +impl Backend { + pub fn new(scenario: &str, points: usize, rolling: usize) -> Self { + let mut counts = Counts::default(); + for index in 1..=points as u64 * 2 { + *counts.entry(key(index)).or_default() += 1; + } + Self { + sketch: scenario == "cmsheap", + rolling, + reference: top(&counts), + counts, + next: 0, + latest: Vec::new(), + accuracy: Vec::new(), + } + } + + /// Validate one window's answers; returns the window id and start time. + pub fn receive(&mut self, obs: &[Observation]) -> Result<(u64, u64), String> { + let meta = Meta::read(obs.first().ok_or("empty backend batch")?)?; + let mut kinds = vec![("tumbling", 1)]; + if self.rolling > 1 { + kinds.push(("rolling", (meta.window as usize + 1).min(self.rolling))); + } + if meta.window != self.next || meta.source != 2 || obs.len() != kinds.len() * K { + return Err("top-k window missing/duplicated".into()); + } + let (mut latest, mut accuracy) = (Vec::new(), Vec::new()); + for (kind, copies) in kinds { + let mut answer = obs + .iter() + .filter(|o| { + o.labels + .iter() + .any(|kv| kv.key == "window_kind" && kv.value == kind) + }) + .map(|o| Ok((label(o, "topk.rank")?, label(o, "topk.key")?, o.value.float))) + .collect::, String>>()?; + answer.sort_by_key(|(rank, _, _)| *rank); + if answer.len() != K { + return Err(format!("{kind}: expected {K} entries")); + } + let exact = + |key: u64| (self.counts.get(&key).copied().unwrap_or(0) * copies as u64) as f64; + let found = |key: u64| answer.iter().find(|(_, k, _)| *k == key).map(|e| e.2); + let hits = self + .reference + .iter() + .filter(|(k, _)| found(*k).is_some()) + .count(); + let recall = hits as f64 / K as f64; + let error = self + .reference + .iter() + .map(|(k, _)| (found(*k).unwrap_or(0.0) - exact(*k)).abs() / exact(*k)) + .sum::() + / K as f64; + let expected = self + .reference + .iter() + .map(|(k, c)| (*k, (c * copies as u64) as f64)) + .collect::>(); + let actual = answer.iter().map(|(_, k, c)| (*k, *c)).collect::>(); + if (!self.sketch && actual != expected) || recall < MIN_RECALL { + return Err(format!( + "{kind}: recall {recall}, answer {actual:?} != {expected:?}" + )); + } + for (rank, k, c) in &answer { + latest.push(TopkPoint { + window_kind: kind.into(), + rank: *rank as usize, + key: format!("{k:016x}"), + estimate: *c, + exact: exact(*k), + }); + } + accuracy.push(TopkAccuracy { + window_kind: kind.into(), + recall, + mean_relative_error: error, + }); + } + self.latest = latest; + self.accuracy = accuracy; + self.next += 1; + Ok((meta.window, meta.started_ns)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_precompute_rs::otap::codec::decode_pdata_to_observations; + + fn exercise(scenario: &str, rolling: usize) -> Backend { + let (n, batch) = (4096, 300); + let mut branches = [0, 1].map(|_| Processor::new(scenario, "branch", n, batch, rolling)); + let mut merge = Processor::new(scenario, "merge", n, batch, rolling); + let mut estimate = Processor::new(scenario, "estimate", n, batch, rolling); + let mut backend = Backend::new(scenario, n, rolling); + let mut completed = 0; + for window in 0..3 { + for offset in (0..n).step_by(batch) { + for (source, branch) in branches.iter_mut().enumerate() { + let input = NormalizedBatch { + first_index: (source * n + offset + 1) as u64, + count: (offset + batch).min(n) - offset, + meta: Meta { + window, + source, + offset, + started_ns: 1, + }, + }; + for a in branch.process_normalized(input).unwrap() { + for m in merge.process(a).unwrap() { + for e in estimate.process(m).unwrap() { + let obs = decode_pdata_to_observations(e).unwrap().observations; + assert_eq!(backend.receive(&obs).unwrap().0, completed); + completed += 1; + } + } + } + } + } + } + assert_eq!(completed, 3); + backend + } + + #[test] + fn exact_top_k_matches_reference() { + let backend = exercise("topk-exact", 2); + assert!(backend + .accuracy + .iter() + .all(|a| a.recall == 1.0 && a.mean_relative_error == 0.0)); + assert_eq!(backend.latest.len(), 2 * K); + } + + #[test] + fn sketch_top_k_meets_recall_bound() { + let backend = exercise("cmsheap", 2); + assert!(backend.accuracy.iter().all(|a| a.recall >= MIN_RECALL)); + } + + #[test] + fn keys_are_heavy_tailed() { + let mut counts = Counts::default(); + (1..=100_000).for_each(|i| *counts.entry(key(i)).or_default() += 1); + let heaviest = top(&counts)[0].1; + assert!( + heaviest > 5_000 && counts.len() > 10_000, + "{heaviest} {}", + counts.len() + ); + } +} diff --git a/asap-precompute-rs/src/config.rs b/asap-precompute-rs/src/config.rs index ba8e032..cb25410 100644 --- a/asap-precompute-rs/src/config.rs +++ b/asap-precompute-rs/src/config.rs @@ -170,13 +170,15 @@ pub fn sketch_size_string(sketch_type: SketchType, params: &SketchParams) -> Opt _ => None, }, }, - SketchType::CountMinSketch => match (params.get("rows"), params.get("columns")) { - (Some(r), Some(c)) => match (as_wire_u64(*r), as_wire_u64(*c)) { - (Some(r), Some(c)) => Some(format!("{r} x {c}")), + SketchType::CountMinSketch | SketchType::CountMinSketchWithHeap => { + match (params.get("rows"), params.get("columns")) { + (Some(r), Some(c)) => match (as_wire_u64(*r), as_wire_u64(*c)) { + (Some(r), Some(c)) => Some(format!("{r} x {c}")), + _ => None, + }, _ => None, - }, - _ => None, - }, + } + } SketchType::Unspecified => None, } } diff --git a/asap-precompute-rs/src/envelope.rs b/asap-precompute-rs/src/envelope.rs index 37706d1..cf6e69f 100644 --- a/asap-precompute-rs/src/envelope.rs +++ b/asap-precompute-rs/src/envelope.rs @@ -51,6 +51,8 @@ pub enum SketchType { CountSketch, /// Count-Min Sketch frequency sketch. CountMinSketch, + /// Count-Min Sketch plus a top-k heavy-hitter heap (ASAPv1 kind `0x03 0x00`). + CountMinSketchWithHeap, } impl SketchType { @@ -63,6 +65,7 @@ impl SketchType { Self::HLLSketch => "HLLSketch", Self::CountSketch => "CountSketch", Self::CountMinSketch => "CountMinSketch", + Self::CountMinSketchWithHeap => "CountMinSketchWithHeap", Self::Unspecified => "Unspecified", } } @@ -225,6 +228,10 @@ mod tests { assert_eq!(SketchType::HLLSketch.name(), "HLLSketch"); assert_eq!(SketchType::CountSketch.name(), "CountSketch"); assert_eq!(SketchType::CountMinSketch.name(), "CountMinSketch"); + assert_eq!( + SketchType::CountMinSketchWithHeap.name(), + "CountMinSketchWithHeap" + ); assert_eq!(SketchType::Unspecified.name(), "Unspecified"); } diff --git a/asap-precompute-rs/src/otap/codec.rs b/asap-precompute-rs/src/otap/codec.rs index a32d84a..7d8e2d0 100644 --- a/asap-precompute-rs/src/otap/codec.rs +++ b/asap-precompute-rs/src/otap/codec.rs @@ -185,7 +185,7 @@ impl OtapSketchEncoder { /// Encodes one flush while retaining series identity across later flushes. pub fn encode(&mut self, envelopes: &[SketchEnvelope]) -> Result { for env in envelopes { - validate_self_describing_payload(env.encoding, &env.payload) + validate_self_describing_payload(env.sketch_type, env.encoding, &env.payload) .map_err(CodecError::Encode)?; } let series_ids = envelopes @@ -223,7 +223,11 @@ impl OtapSketchEncoder { /// /// `sketch.envelope` contains the canonical sketchlib ASAPv1 envelope: /// magic, version, kind ID, length-prefixed metadata, and payload. -fn validate_self_describing_payload(encoding: Encoding, payload: &[u8]) -> Result<(), String> { +fn validate_self_describing_payload( + sketch_type: crate::envelope::SketchType, + encoding: Encoding, + payload: &[u8], +) -> Result<(), String> { if payload.is_empty() { return Ok(()); } @@ -233,6 +237,11 @@ fn validate_self_describing_payload(encoding: Encoding, payload: &[u8]) -> Resul encoding.name() )); } + if sketch_type == crate::envelope::SketchType::CountMinSketchWithHeap { + return crate::sketches::CMSHeapWrapper::from_snapshot(payload) + .map(|_| ()) + .map_err(|error| format!("invalid self-describing CMSHeap sketch.envelope: {error}")); + } asap_sketchlib::sketches::KLL::::deserialize_from_bytes(payload) .map(|_| ()) .map_err(|error| format!("invalid self-describing KLL sketch.envelope: {error}")) @@ -1579,7 +1588,8 @@ impl DecodeAccumulator { .transpose() .map_err(|error| CodecError::Decode(error.to_string()))? .unwrap_or(crate::envelope::Encoding::Unspecified); - validate_self_describing_payload(encoding, &payload).map_err(CodecError::Decode)?; + validate_self_describing_payload(sketch_type, encoding, &payload) + .map_err(CodecError::Decode)?; let timestamp_ms = dp.time_unix_nano() / 1_000_000; let envelope = SketchEnvelope { schema_version: inherited_schema.schema_version.unwrap_or(0), @@ -1881,6 +1891,26 @@ mod tests { assert_eq!(decoded_env.window_end_ms, 2_000); } + #[test] + fn summary_carrier_round_trips_a_cms_heap_envelope() { + use crate::sketches::CMSHeapWrapper; + let mut env = one_envelope("topk_stream", 0.0, "region", "us-east"); + env.sketch_type = SketchType::CountMinSketchWithHeap; + env.encoding = Encoding::Msgpack; + let mut sketch = CMSHeapWrapper::new(4, 64, 3); + sketch.update(b"checkout", 5); + env.payload = sketch.snapshot().expect("CMSHeap ASAPv1 envelope"); + let pdata = encode_envelopes_to_pdata(std::slice::from_ref(&env)).expect("encode"); + let outcome = decode_pdata_to_observations(pdata).expect("decode"); + let decoded = outcome.observations[0].value.envelope.as_ref().unwrap(); + assert_eq!(decoded.sketch_type, SketchType::CountMinSketchWithHeap); + assert_eq!(decoded.payload, env.payload); + + // A CMSHeap payload labelled as KLL is still rejected. + env.sketch_type = SketchType::KLLSketch; + assert!(encode_envelopes_to_pdata(std::slice::from_ref(&env)).is_err()); + } + #[test] fn asapv1_rejects_non_self_describing_envelope_bytes() { let mut env = one_envelope("sketch_stream", 0.0, "region", "us-east"); diff --git a/asap-precompute-rs/src/otap/decode.rs b/asap-precompute-rs/src/otap/decode.rs index b8c6958..b1e502b 100644 --- a/asap-precompute-rs/src/otap/decode.rs +++ b/asap-precompute-rs/src/otap/decode.rs @@ -235,6 +235,7 @@ pub(super) fn parse_sketch_type(row: usize, raw: &str) -> Result Ok(SketchType::HLLSketch), "CountSketch" => Ok(SketchType::CountSketch), "CountMinSketch" => Ok(SketchType::CountMinSketch), + "CountMinSketchWithHeap" => Ok(SketchType::CountMinSketchWithHeap), "Unspecified" | "" => Ok(SketchType::Unspecified), _ => Err(OtapDecodeError::UnknownSketchType { row, diff --git a/asap-precompute-rs/src/sketches/cms.rs b/asap-precompute-rs/src/sketches/cms.rs index 4ffee0e..932a1a3 100644 --- a/asap-precompute-rs/src/sketches/cms.rs +++ b/asap-precompute-rs/src/sketches/cms.rs @@ -25,7 +25,7 @@ const MAX_ROW_HASH_BITS: usize = 64; /// `log2(cols)` == the per-row bit width. When `bits_per_row == 0` /// (cols == 1) every row maps to column 0 regardless, so the slicing /// never overflows and `rows` is left as-is. Returns at least 1. -fn clamp_rows_for_hash_bits(rows: usize, cols: usize) -> usize { +pub(crate) fn clamp_rows_for_hash_bits(rows: usize, cols: usize) -> usize { let bits_per_row = cols.trailing_zeros() as usize; if bits_per_row == 0 { return rows.max(1); diff --git a/asap-precompute-rs/src/sketches/cms_heap.rs b/asap-precompute-rs/src/sketches/cms_heap.rs new file mode 100644 index 0000000..f9c4388 --- /dev/null +++ b/asap-precompute-rs/src/sketches/cms_heap.rs @@ -0,0 +1,171 @@ +//! Count-Min sketch with a top-k heap, over [`asap_sketchlib::CMSHeap`]. +//! +//! Implements [`Sketch`] + [`FrequencySketch`]. Snapshots are sketchlib's +//! ASAPv1 envelope (kind `0x03 0x00`): the counter matrix plus the heap's +//! keys and counts. There is no delta encoding; every payload is full state. + +use std::any::Any; + +use asap_sketchlib::{CMSHeap, DataInput, HeapItem}; + +use super::cms::clamp_rows_for_hash_bits; +use crate::precompute::{DeltaResult, FrequencyEntry, FrequencySketch, PrecomputeError, Sketch}; + +/// Count-Min sketch plus heavy-hitter heap. Keys are raw bytes. +pub struct CMSHeapWrapper { + sk: CMSHeap, +} + +impl CMSHeapWrapper { + /// Construct with `top_k` heap slots. `cols` is rounded up to a power of + /// two and `rows` clamped to the 64-bit hash budget, as for + /// [`super::CMSWrapper`]. + pub fn new(rows: usize, cols: usize, top_k: usize) -> Self { + let cols = cols.max(1).next_power_of_two(); + let rows = clamp_rows_for_hash_bits(rows, cols); + Self { + sk: CMSHeap::new(rows, cols, top_k.max(1)), + } + } + + /// Decode a full ASAPv1 snapshot. + pub fn from_snapshot(bytes: &[u8]) -> Result { + CMSHeap::deserialize_from_bytes(bytes) + .map(|sk| Self { sk }) + .map_err(|e| PrecomputeError::Other(format!("CMSHeapWrapper decode: {e}"))) + } + + /// Add `count` occurrences of `key` and refresh its heap entry. + pub fn update(&mut self, key: &[u8], count: i64) { + self.sk.insert_many(&DataInput::Bytes(key), count); + } + + /// Merge counters, then re-rank both heaps' candidates on the merged counters. + pub fn merge_from(&mut self, other: &Self) -> Result<(), PrecomputeError> { + if (self.sk.rows(), self.sk.cols()) != (other.sk.rows(), other.sk.cols()) { + return Err(PrecomputeError::Other( + "CMSHeapWrapper merge: dimension mismatch".into(), + )); + } + self.sk.merge(&other.sk); + Ok(()) + } + + /// Borrow the underlying sketch. + pub fn inner(&self) -> &CMSHeap { + &self.sk + } +} + +impl Sketch for CMSHeapWrapper { + fn snapshot(&self) -> Result, PrecomputeError> { + if self.sk.heap().is_empty() { + return Ok(Vec::new()); + } + self.sk + .serialize_to_bytes() + .map_err(|e| PrecomputeError::Other(format!("CMSHeapWrapper snapshot: {e}"))) + } + + fn compute_delta_against(&self, _: &[u8], _: u64) -> Result { + Ok(DeltaResult { + payload: self.snapshot()?, + is_full: true, + }) + } + + fn apply_delta(&mut self, payload: &[u8]) -> Result<(), PrecomputeError> { + if payload.is_empty() { + return Ok(()); + } + self.merge_from(&Self::from_snapshot(payload)?) + } + + fn merge(&mut self, other: &dyn Sketch) -> Result<(), PrecomputeError> { + self.apply_delta(&other.snapshot()?) + } + + fn reset(&mut self) { + self.sk = CMSHeap::new(self.sk.rows(), self.sk.cols(), self.sk.heap().capacity()); + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +impl FrequencySketch for CMSHeapWrapper { + fn estimate_count(&self, key: &[u8]) -> f64 { + self.sk.estimate(&DataInput::Bytes(key)) as f64 + } + + /// Heap entries, largest count first. + fn top_k(&self, k: usize) -> Vec { + let mut entries = self + .sk + .heap() + .heap() + .iter() + .filter_map(|item| match &item.key { + HeapItem::Bytes(key) => Some(FrequencyEntry { + key: key.clone(), + count: item.count as f64, + }), + _ => None, + }) + .collect::>(); + entries.sort_by(|a, b| b.count.total_cmp(&a.count).then_with(|| a.key.cmp(&b.key))); + entries.truncate(k); + entries + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn keyed(counts: &[(u64, i64)]) -> CMSHeapWrapper { + let mut sk = CMSHeapWrapper::new(4, 2048, 3); + for (key, count) in counts { + sk.update(&key.to_be_bytes(), *count); + } + sk + } + + fn keys(sk: &CMSHeapWrapper) -> Vec { + sk.top_k(3) + .iter() + .map(|e| u64::from_be_bytes(e.key.as_slice().try_into().unwrap())) + .collect() + } + + #[test] + fn top_k_orders_heavy_hitters() { + let sk = keyed(&[(1, 5), (2, 50), (3, 20), (4, 1), (5, 30)]); + assert_eq!(keys(&sk), vec![2, 5, 3]); + assert_eq!(sk.estimate_count(&2u64.to_be_bytes()), 50.0); + } + + #[test] + fn snapshot_round_trips_through_asapv1() { + let sk = keyed(&[(7, 9), (8, 4)]); + let bytes = sk.snapshot().unwrap(); + let decoded = CMSHeapWrapper::from_snapshot(&bytes).unwrap(); + assert_eq!(decoded.top_k(3), sk.top_k(3)); + assert_eq!(decoded.snapshot().unwrap(), bytes); + } + + #[test] + fn merge_reranks_candidates_from_both_heaps() { + let mut left = keyed(&[(1, 10), (2, 9), (3, 8), (4, 7)]); + let right = keyed(&[(4, 7), (5, 1)]); + left.merge(&right).unwrap(); + assert_eq!(keys(&left), vec![4, 1, 2]); + assert_eq!(left.estimate_count(&4u64.to_be_bytes()), 14.0); + } + + #[test] + fn empty_sketch_has_empty_snapshot() { + assert!(CMSHeapWrapper::new(4, 64, 2).snapshot().unwrap().is_empty()); + } +} diff --git a/asap-precompute-rs/src/sketches/ddsketch.rs b/asap-precompute-rs/src/sketches/ddsketch.rs index 49a8fec..6e96350 100644 --- a/asap-precompute-rs/src/sketches/ddsketch.rs +++ b/asap-precompute-rs/src/sketches/ddsketch.rs @@ -63,18 +63,9 @@ impl DDSketchWrapper { } fn build_state(&self) -> DdSketchState { - DdSketchState { - // Use the gamma-roundtripped alpha so the on-the-wire bytes - // match the portable serialization exactly. - // - // The DataPoint-level METRIC scalars (count/sum/min/max) were - // dropped from `DdSketchState`: the count is recoverable by - // summing `store_counts` and the others are bucket-estimated, - // so the wire format now carries only alpha + the bucket store. - alpha: self.sk.wire_alpha(), - store_counts: self.sk.store_counts.clone(), - store_offset: self.sk.store_offset, - } + // asap_sketchlib owns the proto mapping, including the negative-value + // and zero stores and the gamma-roundtripped alpha. + self.sk.to_proto() } fn encode_envelope(&self) -> Vec { @@ -107,17 +98,10 @@ impl DDSketchWrapper { state.alpha ))); } - // `DdSketchState` no longer carries the count/sum/min/max scalars. - // The total count is recovered by summing `store_counts` (see - // `DdSketch::total_count`); min/max are estimated from the extreme - // non-empty buckets inside `DdSketch::quantile`, within DDSketch's - // alpha relative-accuracy bound. So we reconstruct purely from the - // bucket store + alpha. - Ok(DdSketch::from_raw( - state.alpha, - state.store_counts, - state.store_offset, - )) + // `DdSketchState` carries no count/sum/min/max scalars: the count is + // the sum of all stores and min/max are bucket-estimated, within + // DDSketch's alpha relative-accuracy bound. + Ok(DdSketch::from_proto(state)) } } diff --git a/asap-precompute-rs/src/sketches/mod.rs b/asap-precompute-rs/src/sketches/mod.rs index 404c5e7..a539d4f 100644 --- a/asap-precompute-rs/src/sketches/mod.rs +++ b/asap-precompute-rs/src/sketches/mod.rs @@ -33,12 +33,14 @@ //! inner state proto. pub mod cms; +pub mod cms_heap; pub mod countsketch; pub mod ddsketch; pub mod hll; pub mod kll; pub use cms::{CMSObserver, CMSWrapper}; +pub use cms_heap::CMSHeapWrapper; pub use countsketch::{CountSketchObserver, CountSketchWrapper}; pub use ddsketch::{DDSketchObserver, DDSketchWrapper}; pub use hll::{HLLObserver, HLLWrapper}; diff --git a/asap-precompute-rs/tests/otap_codec.rs b/asap-precompute-rs/tests/otap_codec.rs index fcb661a..1a2eaef 100644 --- a/asap-precompute-rs/tests/otap_codec.rs +++ b/asap-precompute-rs/tests/otap_codec.rs @@ -61,6 +61,7 @@ fn batch_from_envelope_entries(entries: &[(SketchType, &[u8])]) -> RecordBatch { SketchType::HLLSketch => "metric_HLLSketch", SketchType::CountSketch => "metric_CountSketch", SketchType::CountMinSketch => "metric_CountMinSketch", + SketchType::CountMinSketchWithHeap => "metric_CountMinSketchWithHeap", SketchType::Unspecified => "metric_Unspecified", }) .collect(); @@ -183,6 +184,7 @@ fn per_sketch_type_round_trip_smoke() { (SketchType::HLLSketch, vec![0xcc; 16]), (SketchType::CountSketch, vec![0xdd; 128]), (SketchType::CountMinSketch, vec![0xee; 256]), + (SketchType::CountMinSketchWithHeap, vec![0xef; 96]), ] { let env = envelope_for(sketch_type, &payload); let batch = encode_batch(std::slice::from_ref(&env)).expect("encode"); diff --git a/benchmarks/STREAMING.md b/benchmarks/STREAMING.md index 059e266..0abd87f 100644 --- a/benchmarks/STREAMING.md +++ b/benchmarks/STREAMING.md @@ -44,6 +44,14 @@ wall-clock rotation, late-event handling, or window-watermark policies. Msgpack sketch per source/window, merges them, then returns p50/p99. - Backend validates every quantile result. Exact must match; KLL tolerance is `max(abs(reference) * 0.05, 0.001)`. This is value error, not rank error. +- Top-k scenarios use 64-bit keys drawn from Zipf(s=1.1) over + 2^20 ranks instead of durations. `topk-exact` counts keys per window in a + hash map, ships every (key, count) entry in bounded batches, sums the maps, + then returns the top 10. `cmsheap` updates a Count-Min sketch (4 x 2048) + with a 10-slot heavy-hitter heap, sends one ASAPv1 sketch per source/window, + merges sketches, then returns the heap's top 10. The backend requires exact + answers from `topk-exact` and recall@10 >= 0.9 from `cmsheap` in every + window, and reports recall and mean relative count error over the true top 10. The official generator emits generic synthetic metrics. Each branch maps the received item count to the deterministic long-tail HTTP-duration corpus used by diff --git a/benchmarks/run-streaming.py b/benchmarks/run-streaming.py index e563a31..802f75b 100755 --- a/benchmarks/run-streaming.py +++ b/benchmarks/run-streaming.py @@ -23,7 +23,7 @@ ROLES = ("generator_a", "generator_b", "branch_a", "branch_b", "merge", "estimate", "backend") PROCESSOR_ROLES = ("branch_a", "branch_b", "merge", "estimate") -SCENARIOS = ("raw", "exact", "kll") +SCENARIOS = ("raw", "exact", "kll", "topk-exact", "cmsheap") HTTP = urllib.request.build_opener(urllib.request.ProxyHandler({})) @@ -375,6 +375,12 @@ def backlog(s): "generated_signals_after_drain": drained["backend"]["stats"]["completed_windows"] * self.window_points * 2, "validated_signals_after_drain": drained["backend"]["stats"]["completed_windows"] * self.window_points * 2, "correctness": "passed", "resources": resources} + # Top-k answers are validated per window (exact must match, recall@k >= 0.9); + # report the last window's accuracy. + for accuracy in after["backend"]["stats"].get("topk_accuracy", []): + if accuracy["window_kind"] == "tumbling": + report["recall_at_k"] = accuracy["recall"] + report["mean_relative_count_error"] = accuracy["mean_relative_error"] if accounted: inputs = drained["backend"]["stats"]["received_signals"] - cpu_before["backend"]["stats"]["received_signals"] cpu_profile = {} @@ -415,7 +421,7 @@ def close(self): def draw_summary(rows, path): - colors = {"raw": "#64748b", "exact": "#dc2626", "kll": "#2563eb"} + colors = {"raw": "#64748b", "exact": "#dc2626", "kll": "#2563eb", "topk-exact": "#b45309", "cmsheap": "#059669"} maximum = max(r["median_signals_per_second"] for r in rows) height = 100 + len(rows) * 55 svg = [f'', @@ -504,6 +510,15 @@ def gauge(name, labels, value): gauge("asap_demo_reference", labels, point["reference"]) gauge("asap_demo_relative_error", labels, abs(point["value"] - point["reference"]) / abs(point["reference"]) if point["reference"] else 0) + for point in after["backend"]["stats"].get("latest_topk", []): + labels = {"scenario": scenario, "window_kind": point["window_kind"], + "rank": point["rank"], "key": point["key"]} + gauge("asap_demo_topk_estimate", labels, point["estimate"]) + gauge("asap_demo_topk_exact", labels, point["exact"]) + for accuracy in after["backend"]["stats"].get("topk_accuracy", []): + labels = {"scenario": scenario, "window_kind": accuracy["window_kind"]} + gauge("asap_demo_topk_recall", labels, accuracy["recall"]) + gauge("asap_demo_topk_mean_relative_error", labels, accuracy["mean_relative_error"]) return "\n".join(lines) + "\n" @@ -650,6 +665,10 @@ def main(): "min_signals_per_second": min(s["signals_per_second"] for s in samples), "max_signals_per_second": max(s["signals_per_second"] for s in samples), "median_window_latency_p99_ms": statistics.median(s["window_latency_p99_ms"] for s in samples)}) + if all("recall_at_k" in s for s in samples): + rows[-1]["min_recall_at_k"] = min(s["recall_at_k"] for s in samples) + rows[-1]["median_mean_relative_count_error"] = statistics.median( + s["mean_relative_count_error"] for s in samples) latency_limit_ms = max( args.sustainable_max_p99_ms, 2000 * window_points / traffic_rate, @@ -711,7 +730,7 @@ def main(): writer.writerows(capacities) (args.output / "summary.json").write_text(json.dumps(rows, indent=2) + "\n") with (args.output / "summary.csv").open("w", newline="") as f: - writer = csv.DictWriter(f, fieldnames=list(rows[0])) + writer = csv.DictWriter(f, fieldnames=list(dict.fromkeys(k for row in rows for k in row)), restval="") writer.writeheader() writer.writerows(rows) diff --git a/demos/kll-grafana/grafana/dashboards/topk-stream.json b/demos/kll-grafana/grafana/dashboards/topk-stream.json new file mode 100644 index 0000000..d7480e7 --- /dev/null +++ b/demos/kll-grafana/grafana/dashboards/topk-stream.json @@ -0,0 +1,304 @@ +{ + "uid": "asap-topk-stream", + "title": "ASAP CMS + heap vs exact top-k: multi-process streaming", + "schemaVersion": 39, + "version": 1, + "refresh": "5s", + "time": { + "from": "now-15m", + "to": "now" + }, + "panels": [ + { + "id": 1, + "type": "timeseries", + "title": "Backend-validated throughput", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 0 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_validated_signals_per_second{scenario=~\"topk-exact|cmsheap\"}", + "legendFormat": "{{scenario}}" + } + ], + "fieldConfig": { + "defaults": { + "unit": "short", + "min": 0 + }, + "overrides": [] + } + }, + { + "id": 2, + "type": "bargauge", + "title": "Processor CPU-seconds per million signals (branches + merge + estimate)", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 0 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_processor_cpu_seconds_per_million_signals{scenario=~\"topk-exact|cmsheap\"}", + "legendFormat": "{{scenario}}", + "instant": true + } + ], + "fieldConfig": { + "defaults": { + "decimals": 2 + }, + "overrides": [] + }, + "options": { + "orientation": "horizontal", + "displayMode": "basic", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + } + }, + { + "id": 3, + "type": "timeseries", + "title": "CPU cores per pinned processor", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 8 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_cpu_cores{scenario=~\"topk-exact|cmsheap\",component=~\"branch_a|branch_b|merge|estimate\"}", + "legendFormat": "{{scenario}} {{component}}" + } + ], + "fieldConfig": { + "defaults": { + "unit": "short" + }, + "overrides": [] + } + }, + { + "id": 4, + "type": "timeseries", + "title": "Data-network egress per process", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 8 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_data_tx_mbit_per_second{scenario=~\"topk-exact|cmsheap\",component=~\"branch_a|branch_b|merge|estimate\"}", + "legendFormat": "{{scenario}} {{component}}" + } + ], + "fieldConfig": { + "defaults": { + "unit": "Mbits" + }, + "overrides": [] + } + }, + { + "id": 5, + "type": "bargauge", + "title": "Top-10 keys, latest window: CMS + heap estimate", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 10, + "w": 12, + "x": 0, + "y": 16 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_topk_estimate{scenario=\"cmsheap\",window_kind=\"tumbling\"}", + "legendFormat": "#{{rank}} {{key}}", + "instant": true + } + ], + "fieldConfig": { + "defaults": { + "decimals": 0 + }, + "overrides": [] + }, + "options": { + "orientation": "horizontal", + "displayMode": "basic", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + } + }, + { + "id": 6, + "type": "bargauge", + "title": "Top-10 keys, latest window: exact counts", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 10, + "w": 12, + "x": 12, + "y": 16 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_topk_estimate{scenario=\"topk-exact\",window_kind=\"tumbling\"}", + "legendFormat": "#{{rank}} {{key}}", + "instant": true + } + ], + "fieldConfig": { + "defaults": { + "decimals": 0 + }, + "overrides": [] + }, + "options": { + "orientation": "horizontal", + "displayMode": "basic", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + } + } + }, + { + "id": 7, + "type": "timeseries", + "title": "Recall@10 vs exact", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 26 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_topk_recall", + "legendFormat": "{{scenario}} {{window_kind}}" + } + ], + "fieldConfig": { + "defaults": { + "unit": "percentunit", + "min": 0, + "max": 1 + }, + "overrides": [] + } + }, + { + "id": 8, + "type": "timeseries", + "title": "Mean relative count error over the true top-10", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 26 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_topk_mean_relative_error", + "legendFormat": "{{scenario}} {{window_kind}}" + } + ], + "fieldConfig": { + "defaults": { + "unit": "percentunit" + }, + "overrides": [] + } + }, + { + "id": 9, + "type": "timeseries", + "title": "Resident memory per process", + "datasource": { + "type": "prometheus", + "uid": "asap-prometheus" + }, + "gridPos": { + "h": 8, + "w": 24, + "x": 0, + "y": 34 + }, + "targets": [ + { + "refId": "A", + "expr": "asap_demo_rss_mib{scenario=~\"topk-exact|cmsheap\",component=~\"branch_a|branch_b|merge|estimate\"}", + "legendFormat": "{{scenario}} {{component}}" + } + ], + "fieldConfig": { + "defaults": { + "unit": "decmbytes" + }, + "overrides": [] + } + } + ] +} diff --git a/demos/kll-grafana/verify-streaming.py b/demos/kll-grafana/verify-streaming.py index 2940c2f..fdb016e 100644 --- a/demos/kll-grafana/verify-streaming.py +++ b/demos/kll-grafana/verify-streaming.py @@ -39,8 +39,11 @@ def main(): parser.add_argument('--docker-command', default='docker') parser.add_argument('--scenarios', nargs='+', default=['exact', 'kll']) parser.add_argument('--tolerance', type=float, default=0.05, - help='maximum relative error for sketch scenarios') + help='maximum relative (count) error for sketch scenarios') args, extra = parser.parse_known_args() + topk = set(args.scenarios) <= {'topk-exact', 'cmsheap'} + rolling_metric = ('asap_demo_topk_recall{window_kind="rolling"}' if topk + else 'asap_demo_estimate{window_kind="rolling"}') demo = subprocess.Popen([ sys.executable, str(REPO / 'benchmarks/run-streaming.py'), '--demo', '--docker-command', args.docker_command, '--output', str(REPO / 'benchmark-results/demo-verify'), @@ -52,7 +55,7 @@ def main(): assert demo.poll() is None, 'demo exited before metrics were verified' assert time.monotonic() < deadline, 'demo metrics did not appear' windows = values(args.prometheus, 'asap_demo_completed_windows', ['scenario']) - rolling = values(args.prometheus, 'asap_demo_estimate{window_kind="rolling"}', ['scenario']) + rolling = values(args.prometheus, rolling_metric, ['scenario']) if all(windows.get((s,), 0) >= 4 for s in args.scenarios) and len(rolling) == len(args.scenarios): break time.sleep(2) @@ -60,17 +63,27 @@ def main(): assert all(throughput[(s,)] > 0 for s in args.scenarios), throughput cpu = values(args.prometheus, 'asap_demo_cpu_cores', ['scenario', 'component']) assert all((s, c) in cpu for s in args.scenarios for c in PROCESSORS), cpu - errors = values(args.prometheus, 'asap_demo_relative_error', - ['scenario', 'window_kind', 'quantile']) - assert {(s, k, q) for s in args.scenarios for k in ('tumbling', 'rolling') - for q in ('0.5', '0.99')} <= set(errors), errors - for (scenario, _, _), error in errors.items(): - assert error <= (0.0 if scenario == 'exact' else args.tolerance), errors - expression = 'asap_demo_estimate' - keys = ['scenario', 'window_kind', 'quantile'] + if topk: + errors = values(args.prometheus, 'asap_demo_topk_mean_relative_error', + ['scenario', 'window_kind']) + recall = values(args.prometheus, 'asap_demo_topk_recall', ['scenario', 'window_kind']) + assert {(s, k) for s in args.scenarios for k in ('tumbling', 'rolling')} <= set(recall) + assert all(r == 1.0 for (s, _), r in recall.items() if s == 'topk-exact'), recall + assert all(r >= 0.9 for r in recall.values()), recall + print('recall@10:', {'/'.join(k): v for k, v in sorted(recall.items())}) + expression, keys = 'asap_demo_topk_estimate', ['scenario', 'window_kind', 'rank'] + else: + errors = values(args.prometheus, 'asap_demo_relative_error', + ['scenario', 'window_kind', 'quantile']) + assert {(s, k, q) for s in args.scenarios for k in ('tumbling', 'rolling') + for q in ('0.5', '0.99')} <= set(errors), errors + expression, keys = 'asap_demo_estimate', ['scenario', 'window_kind', 'quantile'] + for key, error in errors.items(): + assert error <= (0.0 if key[0] in ('exact', 'topk-exact') else args.tolerance), errors assert set(values(args.prometheus, expression, keys)) == set( values(args.grafana, expression, keys, grafana=True)) - with urllib.request.urlopen(args.grafana.rstrip('/') + '/api/dashboards/uid/asap-kll-stream', + uid = 'asap-topk-stream' if topk else 'asap-kll-stream' + with urllib.request.urlopen(args.grafana.rstrip('/') + '/api/dashboards/uid/' + uid, timeout=10) as response: dashboard = json.load(response)['dashboard'] assert any('asap_demo_validated_signals_per_second' in t['expr'] diff --git a/docs/images/topk-streaming-grafana.png b/docs/images/topk-streaming-grafana.png new file mode 100644 index 0000000..b6893c6 Binary files /dev/null and b/docs/images/topk-streaming-grafana.png differ diff --git a/docs/images/topk-streaming-throughput.png b/docs/images/topk-streaming-throughput.png new file mode 100644 index 0000000..54aaf9b Binary files /dev/null and b/docs/images/topk-streaming-throughput.png differ diff --git a/docs/topk-demo.md b/docs/topk-demo.md new file mode 100644 index 0000000..98e2144 --- /dev/null +++ b/docs/topk-demo.md @@ -0,0 +1,108 @@ +# Top-k heavy hitters: Count-Min sketch with heap vs exact counting + +The top-k demo runs the same multi-process topology as the +[KLL demo](./kll-windows.md): every processor is its own OS process in its own +container, pinned to one CPU, and every hop is OTLP/HTTP over TCP. Two +pipelines run side by side on identical keyed traffic: + +```mermaid +flowchart LR + G["generators A/B
Zipf(1.1) keys"] --> B["branch A/B
1 CPU each"] + B -->|"cmsheap: one ASAPv1 sketch
topk-exact: every (key, count)"| M["merge
1 CPU"] + M --> E["estimate
1 CPU
top 10 (tumbling + rolling)"] + E --> V["validating backend
recall@10, count error"] +``` + +- **`cmsheap`**: each branch inserts every key into a + [`CMSHeapWrapper`](../asap-precompute-rs/src/sketches/cms_heap.rs), which + wraps `asap_sketchlib::CMSHeap`: a 4 x 2048 Count-Min sketch plus a 10-slot + heavy-hitter heap. One ASAPv1 envelope (`CountMinSketchWithHeap`, kind + `0x03 0x00`) per source and window goes to merge. Merge adds the counters + and re-ranks both heaps' candidates on the merged counters; estimate reads + the top 10 from the heap. +- **`topk-exact`**: each branch counts keys in a hash map and ships every + (key, count) entry; merge sums the maps; estimate sorts them for the top 10. + +Keys are 64-bit values drawn from Zipf(s=1.1) over 2^20 ranks, deterministic +per corpus index, with disjoint index ranges per source. A 65,536-key window +has about 17,000 distinct keys per source and 30,000 after merging both. Every window replays the same +corpus, so the backend compares each answer with the exact top 10 computed +once at startup: `topk-exact` must match exactly, and `cmsheap` must reach +recall@10 of at least 0.9 in every window, or the run fails. With +`--rolling-windows N`, estimate also answers over the last N windows by +merging sketches or summing maps. + +## Run + +Linux with Docker and Docker Compose. From the repository root: + +```sh +demos/kll-grafana/run-demo.sh --scenarios topk-exact cmsheap # Ctrl+C stops +``` + +Open . The dashboard shows +backend-validated throughput, processor CPU-seconds per million signals, +per-process CPU, egress and memory, both pipelines' current top 10 with +counts, recall@10, and mean relative count error over the true top 10. + +![CMS with heap vs exact top-k dashboard](images/topk-streaming-grafana.png) + +*56-core x86 server, 50,000 signals/s per source. Both pipelines keep up; +CMS + heap uses 1.7 processor CPU-seconds per million signals, exact 13.0, +and its top 10 matches the exact keys and order with counts within 0.6%.* + +## Benchmark results + +`benchmarks/run-streaming.py --scenarios raw topk-exact cmsheap +--window-points 65536 --repetitions 3` on a 56-core x86 server, four generator CPUs per source and one dedicated CPU +per processor. Median backend-validated signals/s over three runs (both +sources combined): + +| Offered | Raw | Exact top-k | CMS + heap | +| ---: | ---: | ---: | ---: | +| 20,000 | 21,813 | 21,813 | 21,813 | +| 50,000 | 47,988 | 52,348 | 47,989 | +| 100,000 | 100,327 | 100,336 | 100,335 | +| 200,000 | 196,309 | 148,316 | 200,675 | +| 400,000 | 274,848 | 130,879 | 397,004 | +| 600,000 | 279,185 | 130,860 | 597,659 | +| 800,000 | 283,555 | 126,516 | 798,355 | + +(Short runs round to whole windows, so low-rate rows can exceed the offer.) + +![Validated throughput vs offered load](images/topk-streaming-throughput.png) + +By the [capacity criteria](../benchmarks/CAPACITY_DESIGN.md), the largest +sustainable offered rate per source is 50,000 for exact top-k, 100,000 for +raw, and at least 400,000 (the highest tested) for CMS + heap. Exact top-k +plateaus near 130,000 signals/s; raw forwarding near 280,000. + +Summed over branch A, branch B, merge and estimate, at the same offered load: + +| Offered per source | Exact CPU | CMS + heap CPU | Exact egress | CMS + heap egress | +| ---: | ---: | ---: | ---: | ---: | +| 10,000 | 0.29 cores | 0.039 cores | 43.9 Mbit/s | 0.1 Mbit/s | +| 25,000 | 0.66 cores | 0.084 cores | 102.9 Mbit/s | 0.3 Mbit/s | +| 50,000 | 1.29 cores | 0.165 cores | 201.1 Mbit/s | 0.5 Mbit/s | + +Accuracy, checked by the backend in every window of every run: exact top-k +matched exactly; CMS + heap had **recall@10 = 1.0** (same ten keys, same +order) and a **0.62% mean relative count error** over the true top 10, in +tumbling windows. Count-Min only overestimates, so every reported count is +at or above the true count. + +Verify the demo end to end (60 seconds): + +```sh +python3 demos/kll-grafana/verify-streaming.py --scenarios topk-exact cmsheap +``` + +Benchmark sweep (see [STREAMING.md](../benchmarks/STREAMING.md)): + +```sh +python3 benchmarks/run-streaming.py --scenarios raw topk-exact cmsheap \ + --window-points 65536 --repetitions 3 +``` + +`summary.csv` adds `min_recall_at_k` and `median_mean_relative_count_error` +for top-k rows.