chore: bench_cache (#10028)

This commit is contained in:
xzfc
2026-07-30 11:58:06 +00:00
committed by GitHub
parent af4e9e5d77
commit 39547e3a67
9 changed files with 162 additions and 135 deletions
+11 -37
View File
@@ -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<Path>, 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<W: Weight>(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::<W, MmapFile>::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::<W, _>::from_ram_index(&MmapFs, index, path).unwrap();
})
}
/// Compute hash of the given [`InvertedIndexRam`].