diff --git a/Cargo.lock b/Cargo.lock index 279bc2630c..0a4a3991a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1534,6 +1534,7 @@ dependencies = [ "fs-err", "fs4 1.1.0", "fs_extra", + "humantime", "io-uring", "itertools 0.15.0", "log", diff --git a/lib/common/common/Cargo.toml b/lib/common/common/Cargo.toml index 33be40fabf..d87b5b1977 100644 --- a/lib/common/common/Cargo.toml +++ b/lib/common/common/Cargo.toml @@ -26,6 +26,7 @@ chrono = { workspace = true } fs-err = { workspace = true } fs4 = { workspace = true } fs_extra = { workspace = true } +humantime = { workspace = true } itertools = { workspace = true } log = { workspace = true } memmap2 = { workspace = true } diff --git a/lib/common/common/benches/mmap_hashmap.rs b/lib/common/common/benches/mmap_hashmap.rs index 1fa7e85eb5..ff0f38cfc3 100644 --- a/lib/common/common/benches/mmap_hashmap.rs +++ b/lib/common/common/benches/mmap_hashmap.rs @@ -1,11 +1,11 @@ #![cfg_attr(not(target_os = "linux"), expect(clippy::unit_arg))] use std::collections::BTreeMap; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; +use common::bench_cache::{build_once, cache_path}; use common::persisted_hashmap::{MmapHashMap, UniversalHashMap, serialize_hashmap}; use common::universal_io::{OpenOptions, UniversalIoError, UniversalReadFileOps}; use criterion::{Criterion, criterion_group, criterion_main}; -use fs_err as fs; use rand::rngs::SmallRng; use rand::{RngExt, SeedableRng}; @@ -147,28 +147,13 @@ fn bench_mmap_hashmap(c: &mut Criterion) { } fn make_serialized_hashmap(count: usize) -> PathBuf { - let path = Path::new(env!("CARGO_TARGET_TMPDIR")) - .join(env!("CARGO_PKG_NAME")) - .join(env!("CARGO_CRATE_NAME")) - .join(format!("hashmap-{count}.bin")); - - if !path.exists() { - eprintln!("Building serialized hashmap at {path:?}..."); - fs::create_dir_all(path.parent().unwrap()).unwrap(); - + build_once(cache_path!("hashmap-{count}"), |path| { let mut rng = SmallRng::seed_from_u64(42); let map = gen_map(&mut rng, count); - serialize_hashmap::( - &path, - map.iter().map(|(k, v)| (k.as_str(), v.iter().copied())), - ) - .unwrap(); - } - - let size = fs::metadata(&path).unwrap().len(); - eprintln!("Serialized hashmap at {path:?} ({size} bytes)"); - path + let it = map.iter().map(|(k, v)| (k.as_str(), v.iter().copied())); + serialize_hashmap::(path, it).unwrap(); + }) } fn gen_map(rng: &mut SmallRng, count: usize) -> BTreeMap> { diff --git a/lib/common/common/benches/universal_io.rs b/lib/common/common/benches/universal_io.rs index 404407f4fa..667988ae24 100644 --- a/lib/common/common/benches/universal_io.rs +++ b/lib/common/common/benches/universal_io.rs @@ -2,6 +2,7 @@ use std::hint::black_box; use std::io::Write as _; use std::path::{Path, PathBuf}; +use common::bench_cache::{build_once, cache_path}; use common::generic_consts::{Random, Sequential}; use common::mmap::AdviceSetting; #[cfg(target_os = "linux")] @@ -185,35 +186,21 @@ fn ranges_full_file() -> impl Iterator { } fn make_random_file() -> PathBuf { - let path = Path::new(env!("CARGO_TARGET_TMPDIR")) - .join(env!("CARGO_PKG_NAME")) - .join(env!("CARGO_CRATE_NAME")) - .join(format!("random-{FILE_SIZE_BYTES}.bin")); + build_once(cache_path!("random-{FILE_SIZE_BYTES}"), |path| { + let mut file = fs::File::create(path).unwrap(); + let mut rng = SmallRng::seed_from_u64(42); + let mut buffer = vec![0; 1024 * 1024]; + let mut bytes_left = FILE_SIZE_BYTES as usize; - if let Ok(metadata) = fs::metadata(&path) - && metadata.len() == FILE_SIZE_BYTES - { - return path; - } + while bytes_left > 0 { + let len = bytes_left.min(buffer.len()); + rng.fill_bytes(&mut buffer[..len]); + file.write_all(&buffer[..len]).unwrap(); + bytes_left -= len; + } - eprintln!("Building random benchmark file at {path:?}..."); - fs_err::create_dir_all(path.parent().unwrap()).unwrap(); - - let mut file = fs::File::create(&path).unwrap(); - let mut rng = SmallRng::seed_from_u64(42); - let mut buffer = vec![0; 1024 * 1024]; - let mut bytes_left = FILE_SIZE_BYTES as usize; - - while bytes_left > 0 { - let len = bytes_left.min(buffer.len()); - rng.fill_bytes(&mut buffer[..len]); - file.write_all(&buffer[..len]).unwrap(); - bytes_left -= len; - } - - file.flush().unwrap(); - eprintln!("Random benchmark file cached at {path:?}."); - path + file.flush().unwrap(); + }) } criterion_group! { diff --git a/lib/common/common/src/bench_cache.rs b/lib/common/common/src/bench_cache.rs new file mode 100644 index 0000000000..527003cc02 --- /dev/null +++ b/lib/common/common/src/bench_cache.rs @@ -0,0 +1,105 @@ +//! On-disk cache for expensive benchmark setup. + +use std::fmt::Debug; +use std::path::Path; +use std::sync::Once; +use std::time::{Duration, Instant}; + +use fs_err as fs; +use serde::Serialize; +use serde::de::DeserializeOwned; + +/// Where to store the cache. Accepts [`format!`]-style arguments. +/// +/// `cache_path!("foo")` -> `target/tmp/segment/hnsw_search_graph/foo` +#[macro_export] +#[doc(hidden)] +macro_rules! __cache_path { + ($($args:tt)*) => { + ::std::path::Path::new(env!("CARGO_TARGET_TMPDIR")) + .join(env!("CARGO_PKG_NAME")) + .join(env!("CARGO_CRATE_NAME")) + .join(format!($($args)*)) + }; +} +pub use __cache_path as cache_path; + +static HINT: Once = Once::new(); + +static WARN_MSG: &str = r" +note: $BENCH_CACHE is set => using the cache from previous runs. The stale cache + might screw the results. It's on you to delete the stale cache manually. +"; + +/// Preserve/reuse benchmark setup between runs. +/// +/// By default, cache is written to `path` but not reused. +/// Set `$BENCH_CACHE` to skip `build()` if the `path` already exists. +/// +/// `build(arg)` should create file/dir at `arg`, not `path`. +pub fn build_once + Debug>(path: P, build: impl FnOnce(&Path)) -> P { + let using_cache = std::env::var_os("BENCH_CACHE").is_some(); + let path_old = path.as_ref().with_added_extension("old"); + let path_tmp = path.as_ref().with_added_extension("tmp"); + + rm_rf(&path_tmp).unwrap(); + + if path.as_ref().exists() { + rm_rf(&path_old).unwrap(); + } else if path_old.exists() { + fs::rename(&path_old, path.as_ref()).unwrap(); + } + + if path.as_ref().exists() && using_cache { + let age = fs::metadata(path.as_ref()).and_then(|meta| meta.modified()); + let age = age.map_or(Duration::ZERO, |time| time.elapsed().unwrap_or_default()); + let age = humantime::format_duration(Duration::from_secs(age.as_secs() / 60 * 60)); + eprintln!("build_once: using cache {path:?}, built {age} ago."); + HINT.call_once(|| eprintln!("{}", WARN_MSG.trim())); + return path; + } + + eprintln!("build_once: Building {path:?}..."); + if !using_cache { + // Discoverability hint. + HINT.call_once(|| eprintln!("hint: Set BENCH_CACHE=1 to reuse caches between runs.")); + } + + fs::create_dir_all(path.as_ref().parent().unwrap()).unwrap(); + if path.as_ref().exists() { + // We won't use the existing cache this time (BENCH_CACHE is not set), + // but still preserve it in case the build fails/interrupted. + // Give the user the last chance to press ^C and re-run with BENCH_CACHE=1. + fs::rename(path.as_ref(), &path_old).unwrap(); + } + + let started = Instant::now(); + build(&path_tmp); + if path.as_ref().exists() { + // Catch misbehaving `build()`. + rm_rf(path.as_ref()).unwrap(); + panic!("build_once: build() should not create the final path"); + } + fs::rename(&path_tmp, path.as_ref()).unwrap(); + rm_rf(&path_old).unwrap(); + eprintln!("build_once: took {:?} to build {path:?}", started.elapsed()); + + path +} + +/// [`build_once`], for values stored as JSON. +pub fn cached_json(path: &Path, build: impl FnOnce() -> T) -> T { + build_once(path, |path| { + fs::write(path, serde_json::to_vec(&build()).unwrap()).unwrap(); + }); + crate::fs::read_json(path).unwrap() +} + +fn rm_rf(path: &Path) -> std::io::Result<()> { + match fs::symlink_metadata(path) { + Ok(meta) if meta.is_dir() => fs::remove_dir_all(path), + Ok(_) => fs::remove_file(path), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e), + } +} diff --git a/lib/common/common/src/lib.rs b/lib/common/common/src/lib.rs index de44c55b4d..accaf44a34 100644 --- a/lib/common/common/src/lib.rs +++ b/lib/common/common/src/lib.rs @@ -1,4 +1,6 @@ pub mod aligned_buf; +#[cfg(feature = "testing")] +pub mod bench_cache; pub mod binary_search; pub mod bitpacking; pub mod bitpacking_links; diff --git a/lib/segment/benches/fixture.rs b/lib/segment/benches/fixture.rs index 4e3137ff78..483f7a2daa 100644 --- a/lib/segment/benches/fixture.rs +++ b/lib/segment/benches/fixture.rs @@ -1,6 +1,4 @@ -use std::path::Path; -use std::time::Duration; - +use common::bench_cache::{build_once, cache_path}; use common::types::PointOffsetType; use fs_err as fs; use rand::SeedableRng as _; @@ -32,16 +30,6 @@ where { use indicatif::{ParallelProgressIterator as _, ProgressStyle}; - let path = Path::new(env!("CARGO_TARGET_TMPDIR")) - .join(env!("CARGO_PKG_NAME")) - .join(env!("CARGO_CRATE_NAME")) - .join(format!( - // The "smallrng" suffix keys the cache by RNG algorithm: the cached - // graph must match the vectors regenerated below. - "{num_vectors}-{dim}-{m}-{ef_construct}-{use_heuristic}-{:?}-smallrng", - METRIC::distance(), - )); - // Note: make sure that vector generation is deterministic. let vector_holder = TestRawScorerProducer::new( dim, @@ -51,13 +39,12 @@ where &mut SmallRng::seed_from_u64(42), ); - let graph_layers_path = GraphLayers::get_path(&path); - let graph_layers = if graph_layers_path.exists() { - let updated_ago = updated_ago(&graph_layers_path).unwrap_or_else(|_| "???".to_string()); - eprintln!("Loading cached links (built {updated_ago} ago) from {graph_layers_path:?}."); - eprintln!("Delete the directory above if code related to HNSW graph building is changed"); - GraphLayers::load(&path, GraphLinksResidency::Cached, false).unwrap() - } else { + let path = cache_path!( + "graph-{num_vectors}-{dim}-{m}-{ef_construct}-{use_heuristic}-{:?}", + METRIC::distance(), + ); + + build_once(&path, |path| { let mut graph_layers_builder = GraphLayersBuilder::new(num_vectors, HnswM::new2(m), ef_construct, 10, use_heuristic); @@ -80,17 +67,13 @@ where ) .for_each(add_point); - fs::create_dir_all(&path).unwrap(); + fs::create_dir_all(path).unwrap(); graph_layers_builder - .into_graph_layers(&path, GraphLinksFormatParam::Plain, false) - .unwrap() - }; + .into_graph_layers(path, GraphLinksFormatParam::Plain, false) + .unwrap(); + }); + + let graph_layers = GraphLayers::load(&path, GraphLinksResidency::Cached, false).unwrap(); (vector_holder, graph_layers) } - -fn updated_ago(path: &Path) -> Result> { - let elapsed = fs::metadata(path)?.modified()?.elapsed()?; - let secs_rounded = elapsed.as_secs().next_multiple_of(60); - Ok(humantime::format_duration(Duration::from_secs(secs_rounded)).to_string()) -} diff --git a/lib/segment/benches/hnsw_incremental_build.rs b/lib/segment/benches/hnsw_incremental_build.rs index 240b06577e..c21fbbab22 100644 --- a/lib/segment/benches/hnsw_incremental_build.rs +++ b/lib/segment/benches/hnsw_incremental_build.rs @@ -11,10 +11,10 @@ use std::sync::atomic::AtomicBool; use atomic_refcell::AtomicRefCell; use clap::Parser; +use common::bench_cache::{cache_path, cached_json}; use common::budget::ResourcePermit; use common::counter::hardware_counter::HardwareCounterCell; use common::flags::{FeatureFlags, feature_flags, init_feature_flags}; -use common::fs::{atomic_save_json, read_json}; use common::progress_tracker::ProgressTracker; use common::types::ScoredPointOffset; use fs_err as fs; @@ -157,11 +157,6 @@ fn main() { .tempdir() .unwrap(); - let cache_path = Path::new(env!("CARGO_TARGET_TMPDIR")) - .join(env!("CARGO_PKG_NAME")) - .join(env!("CARGO_CRATE_NAME")); - fs::create_dir_all(&cache_path).unwrap(); - // Load the dataset or generate random vectors. let (dataset_mmap, dataset); let vectors_mem; @@ -238,13 +233,13 @@ fn main() { args.distance, ); let initial_index_path = if args.cache { - cache_path.join(format!( + cache_path!( "initial-{dataset_hash}-{m}-{ef_construct}-{distance:?}", dataset_hash = dataset_hash(vectors[sliding_window.clone()].iter().copied()), m = args.m, ef_construct = args.ef_construct, distance = args.distance, - )) + ) } else { last_segment.data_path().join("hnsw_bench") }; @@ -281,10 +276,10 @@ fn main() { && (iteration % args.accuracy_check_period == 0 || iteration == args.iterations - 1) { let top = 10; - let exact_cache_path = cache_path.join(format!( + let exact_cache_path = cache_path!( "exact-{queries_hash}-{}-{top}", dataset_hash(sliding_window.clone().map(|i| vectors[i % vectors.len()])), - )); + ); let accuracy = measure_accuracy( &exact_cache_path, &segment, @@ -434,12 +429,8 @@ fn measure_accuracy( let id_tracker = segment.id_tracker.borrow(); // Exact search (aka full scan) is slow, so we cache the results. - let exact_search_results; - if exact_cache_path.exists() { - exact_search_results = read_json(exact_cache_path).unwrap() - } else { - let start = std::time::Instant::now(); - exact_search_results = query_vectors + let exact_search_results: Vec<_> = cached_json(exact_cache_path, || { + query_vectors .par_iter() .map(|query| { segment.vector_data[DEFAULT_VECTOR_NAME] @@ -448,10 +439,8 @@ fn measure_accuracy( .search(&[query], None, top, None, &Default::default()) .pipe(|results| process_search_results(&id_tracker, results)) }) - .collect::>(); - log::debug!("Exact search time = {:?}", start.elapsed()); - atomic_save_json(exact_cache_path, &exact_search_results).unwrap(); - } + .collect::>() + }); let sames: usize = query_vectors .par_iter() diff --git a/lib/sparse/benches/search.rs b/lib/sparse/benches/search.rs index dd1c57c196..2eca4d5e79 100644 --- a/lib/sparse/benches/search.rs +++ b/lib/sparse/benches/search.rs @@ -2,6 +2,7 @@ use std::borrow::Cow; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; +use common::bench_cache::{build_once, cache_path}; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use common::universal_io::{MmapFile, MmapFs}; @@ -235,55 +236,28 @@ fn load_csr_index(path: impl AsRef, ratio: f32) -> InvertedIndexRam { builder.build() } -fn cache_dir() -> PathBuf { - Path::new(env!("CARGO_TARGET_TMPDIR")) - .join(env!("CARGO_PKG_NAME")) - .join(env!("CARGO_CRATE_NAME")) - // Keyed by RNG algorithm: cached indexes must match freshly generated queries. - .join("smallrng") -} - /// Load an [`InvertedIndexRam`] from the cache. /// If not exists, calls `build()` to create it. fn cached_ram_index(name: &str, build: impl FnOnce() -> InvertedIndexRam) -> InvertedIndexRam { - let path = cache_dir().join(name); - if !path.exists() { - eprintln!("Building cache: {path:?}"); - let index = build(); - fs::create_dir_all(path.parent().unwrap()).unwrap(); - index.test_save(&path.with_extension("tmp")).unwrap(); - fs::rename(path.with_extension("tmp"), &path).unwrap(); - } - eprintln!("Using cache: {path:?}"); + let path = build_once(cache_path!("{name}"), |path| { + build().test_save(path).unwrap(); + }); InvertedIndexRam::test_load(&path).unwrap() } /// Loads [`InvertedIndexCompressedMmap`] from the cache. /// If not exists, converts it from the given [`InvertedIndexRam`]. fn cached_compressed_index(index: &InvertedIndexRam, name: &str) -> PathBuf { - let path = cache_dir().join(format!( + let path = cache_path!( "{name}-{}-{hash}", W::NAME, hash = inverted_index_partial_hash(index) - )); - - if !path.exists() { - eprintln!("Building cache: {path:?}"); - let tmp_path = path.with_extension("tmp"); - if tmp_path.exists() { - fs::remove_dir_all(&tmp_path).unwrap(); - } - fs::create_dir_all(&tmp_path).unwrap(); - InvertedIndexCompressedMmap::::from_ram_index( - &MmapFs, - Cow::Borrowed(index), - &tmp_path, - ) - .unwrap(); - fs::rename(tmp_path, &path).unwrap(); - } - eprintln!("Using cache: {path:?}"); - path + ); + build_once(path, |path| { + fs::create_dir_all(path).unwrap(); + let index = Cow::Borrowed(index); + InvertedIndexCompressedMmap::::from_ram_index(&MmapFs, index, path).unwrap(); + }) } /// Compute hash of the given [`InvertedIndexRam`].