diff --git a/diskann-disk/Cargo.toml b/diskann-disk/Cargo.toml index 693614381..e4928d0fb 100644 --- a/diskann-disk/Cargo.toml +++ b/diskann-disk/Cargo.toml @@ -82,6 +82,7 @@ proptest.workspace = true [features] default = [] perf_test = ["dep:opentelemetry"] +benchmark-search-tracing = [] experimental_diversity_search = [ "diskann/experimental_diversity_search", "diskann-providers/experimental_diversity_search", diff --git a/diskann-disk/README.md b/diskann-disk/README.md index f21ef9434..fa163b564 100644 --- a/diskann-disk/README.md +++ b/diskann-disk/README.md @@ -48,6 +48,12 @@ This crate has been populated with the core disk index functionality from the ma - Disk-specific partitioning utilities +## Features + +- `benchmark-search-tracing`: Adds coarse, opt-in search spans and aggregate traversal fields for + diagnostic benchmarks. The feature is disabled by default and is not intended for production + search builds. + ## Dependencies This crate depends on: diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 37a3171df..f1f6ad33a 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -47,9 +47,51 @@ use diskann_vector::{distance::Metric, DistanceFunction}; use tokio::runtime::Runtime; use tracing::debug; +#[cfg(feature = "benchmark-search-tracing")] +macro_rules! benchmark_trace_span { + ($($fields:tt)*) => { + tracing::trace_span!($($fields)*) + }; +} + +#[cfg(not(feature = "benchmark-search-tracing"))] +macro_rules! benchmark_trace_span { + ($($fields:tt)*) => { + BenchmarkSpan + }; +} + +#[cfg(feature = "benchmark-search-tracing")] +type BenchmarkSpan = tracing::Span; + +#[cfg(not(feature = "benchmark-search-tracing"))] +struct BenchmarkSpan; + +#[cfg(not(feature = "benchmark-search-tracing"))] +struct BenchmarkSpanGuard; + +#[cfg(not(feature = "benchmark-search-tracing"))] +impl BenchmarkSpan { + fn enter(&self) -> BenchmarkSpanGuard { + BenchmarkSpanGuard + } + + fn record(&self, _field: &str, _value: T) {} +} + +#[cfg(feature = "benchmark-search-tracing")] +fn benchmark_tracing_enabled() -> bool { + tracing::enabled!(tracing::Level::TRACE) +} + +#[cfg(not(feature = "benchmark-search-tracing"))] +fn benchmark_tracing_enabled() -> bool { + false +} + use crate::{ data_model::{CachingStrategy, GraphHeader}, - error::{diskann_error, ErrorKind}, + error::{diskann_error, Error as DiskError, ErrorKind}, search::{ provider::{ aligned_file_reader::AlignedFileReaderFactory, @@ -254,10 +296,111 @@ where // Struct to track IO. This is used by single thread, but needs to be Atomic as the Strategy has "Send" trait bound. // There should be minimal to no overhead compared to using a raw reference. +#[cfg(feature = "benchmark-search-tracing")] +#[derive(Default)] +struct BenchmarkSearchMetrics { + expand_time_us: AtomicU64, + pq_time_us: AtomicU64, + read_time_us: AtomicU64, + materialize_time_us: AtomicU64, + beam_count: AtomicUsize, + neighbor_count: AtomicUsize, +} + +#[cfg(not(feature = "benchmark-search-tracing"))] +#[derive(Default)] +struct BenchmarkSearchMetrics; + +#[cfg(feature = "benchmark-search-tracing")] +impl BenchmarkSearchMetrics { + fn record_expand(&self, expand_time_us: u64, neighbor_count: usize) { + IOTracker::add_time(&self.expand_time_us, expand_time_us); + self.beam_count + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.neighbor_count + .fetch_add(neighbor_count, std::sync::atomic::Ordering::Relaxed); + } + + fn record_pq_time(&self, pq_time_us: u64) { + IOTracker::add_time(&self.pq_time_us, pq_time_us); + } + + fn record_vertex_load(&self, timings: VertexLoadTimings) { + IOTracker::add_time(&self.read_time_us, timings.read_us); + IOTracker::add_time(&self.materialize_time_us, timings.materialize_us); + } + + fn time(category: &AtomicU64) -> u64 { + category.load(std::sync::atomic::Ordering::Relaxed) + } + + fn count(category: &AtomicUsize) -> usize { + category.load(std::sync::atomic::Ordering::Relaxed) + } + + fn beam_count(&self) -> usize { + Self::count(&self.beam_count) + } + + fn neighbor_count(&self) -> usize { + Self::count(&self.neighbor_count) + } + + fn expand_time_us(&self) -> u64 { + Self::time(&self.expand_time_us) + } + + fn pq_time_us(&self) -> u64 { + Self::time(&self.pq_time_us) + } + + fn read_time_us(&self) -> u64 { + Self::time(&self.read_time_us) + } + + fn materialize_time_us(&self) -> u64 { + Self::time(&self.materialize_time_us) + } +} + +#[cfg(not(feature = "benchmark-search-tracing"))] +impl BenchmarkSearchMetrics { + fn record_expand(&self, _expand_time_us: u64, _neighbor_count: usize) {} + + fn record_pq_time(&self, _pq_time_us: u64) {} + + fn record_vertex_load(&self, _timings: VertexLoadTimings) {} + + fn beam_count(&self) -> usize { + 0 + } + + fn neighbor_count(&self) -> usize { + 0 + } + + fn expand_time_us(&self) -> u64 { + 0 + } + + fn pq_time_us(&self) -> u64 { + 0 + } + + fn read_time_us(&self) -> u64 { + 0 + } + + fn materialize_time_us(&self) -> u64 { + 0 + } +} + struct IOTracker { io_time_us: AtomicU64, preprocess_time_us: AtomicU64, io_count: AtomicUsize, + benchmark: BenchmarkSearchMetrics, } impl Default for IOTracker { @@ -266,6 +409,7 @@ impl Default for IOTracker { io_time_us: AtomicU64::new(0), preprocess_time_us: AtomicU64::new(0), io_count: AtomicUsize::new(0), + benchmark: BenchmarkSearchMetrics::default(), } } } @@ -469,11 +613,26 @@ where > + Send + ?Sized, { + let trace_enabled = benchmark_tracing_enabled(); + let rerank_span = benchmark_trace_span!( + "diskann_rerank", + candidates = tracing::field::Empty, + uncached_candidates = tracing::field::Empty, + result_count = tracing::field::Empty, + vertex_read_time_us = tracing::field::Empty, + vertex_materialize_time_us = tracing::field::Empty, + ); + let _rerank_guard = rerank_span.enter(); let provider = accessor.provider; + let mut candidate_count = 0usize; let mut uncached_ids = Vec::new(); + let mut rerank_load_timings = VertexLoadTimings::default(); let mut reranked: Vec<_> = { let mut process = |n: u32| { + if trace_enabled { + candidate_count += 1; + } if let Some(entry) = accessor.scratch.distance_cache.get(&n) { Some(Neighbor::new((n, entry.1), entry.0)) } else { @@ -494,7 +653,8 @@ where } }; if !uncached_ids.is_empty() { - ensure_vertex_loaded(&mut accessor.scratch.vertex_provider, &uncached_ids)?; + rerank_load_timings = + ensure_vertex_loaded(&mut accessor.scratch.vertex_provider, &uncached_ids)?; for n in &uncached_ids { let v = accessor.scratch.vertex_provider.get_vector(n)?; let d = provider.distance_comparer.evaluate_similarity(query, v); @@ -507,7 +667,18 @@ where reranked.sort_unstable_by(neighbor::ord::fast_distance); // Store the reranked results. - extend_output(accessor, reranked, output) + let result_count = extend_output(accessor, reranked, output)?; + if trace_enabled { + rerank_span.record("candidates", candidate_count as u64); + rerank_span.record("uncached_candidates", uncached_ids.len() as u64); + rerank_span.record("result_count", result_count as u64); + rerank_span.record("vertex_read_time_us", rerank_load_timings.read_us); + rerank_span.record( + "vertex_materialize_time_us", + rerank_load_timings.materialize_us, + ); + } + Ok(result_count) } } @@ -758,6 +929,7 @@ where F: FnMut(f32, u32), { let pq_scratch = &mut self.scratch.pq_scratch; + let pq_timer = benchmark_tracing_enabled().then(Instant::now); compute_pq_distance( ids, self.provider.pq_data.get_num_chunks(), @@ -766,6 +938,11 @@ where &mut pq_scratch.aligned_pq_coord_scratch, &mut pq_scratch.aligned_dist_scratch, )?; + if let Some(pq_timer) = pq_timer { + self.io_tracker + .benchmark + .record_pq_time(pq_timer.elapsed().as_micros() as u64); + } for (i, id) in ids.iter().enumerate() { let distance = self.scratch.pq_scratch.aligned_dist_scratch[i]; @@ -816,9 +993,12 @@ where let result = (|| { let io_limit = self.provider.search_io_limit - self.io_tracker.io_count(); let load_ids: Box<[_]> = ids.take(io_limit).collect(); + let trace_enabled = benchmark_tracing_enabled(); + let expand_timer = trace_enabled.then(Instant::now); self.ensure_loaded(&load_ids)?; let mut ids = Vec::new(); + let mut neighbor_count = 0usize; for i in load_ids { ids.clear(); ids.extend( @@ -830,8 +1010,15 @@ where .filter(|id| pred.eval_mut(id)), ); + neighbor_count += ids.len(); self.pq_distances(&ids, &mut |dist, id| f(id, dist))?; } + if trace_enabled { + self.io_tracker.benchmark.record_expand( + expand_timer.map_or(0, |timer| timer.elapsed().as_micros() as u64), + neighbor_count, + ); + } Ok(()) })(); @@ -860,6 +1047,14 @@ where where VPF: VertexProviderFactory, { + let trace_enabled = benchmark_tracing_enabled(); + let prepare_span = benchmark_trace_span!( + "diskann_prepare_query", + query_dimensions = query.len() as u64, + pq_chunks = provider.pq_data.get_num_chunks() as u64, + preprocess_us = tracing::field::Empty, + ); + let _prepare_guard = prepare_span.enter(); let mut scratch = PoolOption::try_pooled( scratch_pool, &DiskSearchScratchArgs { @@ -884,10 +1079,11 @@ where provider.metric, &[start_vertex_id], )?; - IOTracker::add_time( - &io_tracker.preprocess_time_us, - timer.elapsed().as_micros() as u64, - ); + let preprocess_us = timer.elapsed().as_micros() as u64; + IOTracker::add_time(&io_tracker.preprocess_time_us, preprocess_us); + if trace_enabled { + prepare_span.record("preprocess_us", preprocess_us); + } Ok(Self { provider, @@ -902,14 +1098,10 @@ where if ids.is_empty() { return Ok(()); } - let scratch = &mut self.scratch; let timer = Instant::now(); - ensure_vertex_loaded(&mut scratch.vertex_provider, ids)?; - IOTracker::add_time( - &self.io_tracker.io_time_us, - timer.elapsed().as_micros() as u64, - ); - self.io_tracker.add_io_count(ids.len()); + let timings = ensure_vertex_loaded(&mut self.scratch.vertex_provider, ids)?; + self.record_vertex_load(ids.len(), timer.elapsed().as_micros() as u64, timings); + let scratch = &mut self.scratch; for id in ids { let vector = scratch.vertex_provider.get_vector(id)?; let distance = self @@ -924,6 +1116,12 @@ where } Ok(()) } + + fn record_vertex_load(&self, vertex_count: usize, total_us: u64, timings: VertexLoadTimings) { + IOTracker::add_time(&self.io_tracker.io_time_us, total_us); + self.io_tracker.add_io_count(vertex_count); + self.io_tracker.benchmark.record_vertex_load(timings); + } } impl Drop for DiskAccessor<'_, Data, VP> @@ -1351,6 +1549,79 @@ where associated_data: &mut [Data::AssociatedDataType], indexed_vectors: IndexedVectorOutput<'_, Data::VectorDataType>, mode: &SearchMode<'_>, + ) -> ANNResult { + let trace_enabled = benchmark_tracing_enabled(); + let search_span = benchmark_trace_span!( + "diskann_search_internal", + query_dimensions = query.len() as u64, + k = k_value as u64, + l = search_list_size as u64, + beam_width = beam_width.unwrap_or(1) as u64, + io_limit = self.index.provider().search_io_limit as u64, + result_count = tracing::field::Empty, + comparisons = tracing::field::Empty, + hops = tracing::field::Empty, + traversal_io_operations = tracing::field::Empty, + traversal_vertex_load_time_us = tracing::field::Empty, + preprocess_time_us = tracing::field::Empty, + traversal_non_vertex_load_wall_time_us = tracing::field::Empty, + beam_expansions = tracing::field::Empty, + neighbors_considered = tracing::field::Empty, + expand_wall_time_us = tracing::field::Empty, + pq_time_us = tracing::field::Empty, + traversal_vertex_read_time_us = tracing::field::Empty, + traversal_vertex_materialize_time_us = tracing::field::Empty, + outcome = tracing::field::Empty, + error_category = tracing::field::Empty, + ); + let _search_guard = search_span.enter(); + let result = self.search_internal_impl_inner( + query, + k_value, + search_list_size, + beam_width, + query_stats, + indices, + distances, + associated_data, + indexed_vectors, + mode, + trace_enabled, + &search_span, + ); + if trace_enabled { + match &result { + Ok(_) => { + search_span.record("outcome", "success"); + } + Err(error) => { + search_span.record("outcome", "error"); + match error.downcast_ref::() { + Some(error) => search_span + .record("error_category", tracing::field::debug(error.kind())), + None => search_span.record("error_category", "unknown"), + }; + } + } + } + result + } + + #[allow(clippy::too_many_arguments)] + fn search_internal_impl_inner( + &self, + query: &[Data::VectorDataType], + k_value: usize, + search_list_size: u32, + beam_width: Option, + query_stats: &mut QueryStatistics, + indices: &mut [u32], + distances: &mut [f32], + associated_data: &mut [Data::AssociatedDataType], + indexed_vectors: IndexedVectorOutput<'_, Data::VectorDataType>, + mode: &SearchMode<'_>, + trace_enabled: bool, + search_span: &BenchmarkSpan, ) -> ANNResult { let l = search_list_size as usize; @@ -1370,7 +1641,6 @@ where )?; let timer = Instant::now(); - let io_tracker = IOTracker::default(); // * `FlatScan` — `flat_search` filters the scan iterator at @@ -1469,6 +1739,42 @@ where query_stats.cpu_time_us = query_stats.total_execution_time_us - query_stats.io_time_us - query_stats.query_pq_preprocess_time_us; + if trace_enabled { + search_span.record("result_count", stats.result_count as u64); + search_span.record("comparisons", stats.cmps as u64); + search_span.record("hops", stats.hops as u64); + search_span.record( + "traversal_io_operations", + query_stats.total_io_operations as u64, + ); + search_span.record( + "traversal_vertex_load_time_us", + query_stats.io_time_us as u64, + ); + search_span.record( + "preprocess_time_us", + query_stats.query_pq_preprocess_time_us as u64, + ); + search_span.record( + "traversal_non_vertex_load_wall_time_us", + query_stats.cpu_time_us as u64, + ); + search_span.record("beam_expansions", io_tracker.benchmark.beam_count() as u64); + search_span.record( + "neighbors_considered", + io_tracker.benchmark.neighbor_count() as u64, + ); + search_span.record("expand_wall_time_us", io_tracker.benchmark.expand_time_us()); + search_span.record("pq_time_us", io_tracker.benchmark.pq_time_us()); + search_span.record( + "traversal_vertex_read_time_us", + io_tracker.benchmark.read_time_us(), + ); + search_span.record( + "traversal_vertex_materialize_time_us", + io_tracker.benchmark.materialize_time_us(), + ); + } Ok(SearchResultStats { cmps: query_stats.total_comparisons, result_count: stats.result_count, @@ -1482,15 +1788,29 @@ where /// This is a convenience function that combines `load_vertices` and `process_loaded_node` /// for each vertex ID. It first loads all the vertices in batch, then processes each /// loaded node. +#[derive(Default)] +struct VertexLoadTimings { + read_us: u64, + materialize_us: u64, +} + fn ensure_vertex_loaded>( vertex_provider: &mut V, ids: &[Data::VectorIdType], -) -> ANNResult<()> { +) -> ANNResult { + let trace_enabled = benchmark_tracing_enabled(); + let read_timer = trace_enabled.then(Instant::now); vertex_provider.load_vertices(ids)?; + let read_us = read_timer.map_or(0, |timer| timer.elapsed().as_micros() as u64); + let materialize_timer = trace_enabled.then(Instant::now); for (idx, id) in ids.iter().enumerate() { vertex_provider.process_loaded_node(id, idx)?; } - Ok(()) + let materialize_us = materialize_timer.map_or(0, |timer| timer.elapsed().as_micros() as u64); + Ok(VertexLoadTimings { + read_us, + materialize_us, + }) } #[cfg(test)] @@ -1522,6 +1842,26 @@ mod disk_provider_tests { "/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_truth_search"; const TEST_INDEX_128DIM: &str = "/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_truth_search_disk.index"; + + #[cfg(feature = "benchmark-search-tracing")] + #[test] + fn benchmark_search_metrics_aggregate_phase_totals() { + let metrics = BenchmarkSearchMetrics::default(); + metrics.record_expand(11, 17); + metrics.record_expand(13, 19); + metrics.record_pq_time(7); + metrics.record_vertex_load(VertexLoadTimings { + read_us: 23, + materialize_us: 5, + }); + + assert_eq!(metrics.beam_count(), 2); + assert_eq!(metrics.neighbor_count(), 36); + assert_eq!(metrics.expand_time_us(), 24); + assert_eq!(metrics.pq_time_us(), 7); + assert_eq!(metrics.read_time_us(), 23); + assert_eq!(metrics.materialize_time_us(), 5); + } const TEST_PQ_PIVOT_128DIM: &str = "/disk_index_search/disk_index_sift_learn_R4_L50_A1.2_truth_search_pq_pivots.bin"; const TEST_PQ_COMPRESSED_128DIM: &str =