perf(agent): store embeddings once, not twice
MemoryCache held every embedding in two places: a `Vec<Vec<f32>>` and a flattened copy for the batched kernels, kept in lock-step on every push, update and compaction. A store loaded from disk therefore carried the corpus twice, plus one heap allocation per entry. A new `cache::Embeddings` owns just the flat `[N x dim]` buffer and indexes into it, so `embeddings[i]` still reads as a `&[f32]` row. The batch kernels take a `VectorSet` (implemented for both `Embeddings` and `Vec<Vec<f32>>`) instead of `&[Vec<f32>]`, so their callers and tests are unchanged. Loading no longer unflattens what it just read. 100k 384-dim entries, reopened from disk: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors. Recall (1.0000 at ef=64) and query latency are unchanged. Rows are now always exactly `dim` long, shorter ones zero-padded. The old representation allowed ragged rows, which silently misaligned the flattened copy — every row after a wrong-length embedding — and `update` carried a comment about falling back to a rebuild to avoid exactly that. It is now unrepresentable. A record saved without an embedding holds a zero row and is told apart by its norm, which is what `total_embeddings` now counts. Measured with a counting allocator rather than RSS: freeing a structure returns its pages to the allocator's pool, not the OS, so an RSS reading from inside the process showed the two representations as identical. Breaking: MemoryCache::embeddings changes type, embeddings_flat is replaced by flat_embeddings(), rebuild_flat() is a deprecated no-op. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -198,6 +198,57 @@ fn summarize(mut samples: Vec<Duration>) -> Latency {
|
||||
}
|
||||
}
|
||||
|
||||
/// Counts live heap bytes, so a structure's cost can be measured by
|
||||
/// difference.
|
||||
///
|
||||
/// RSS cannot do this from inside one process: freeing a large structure
|
||||
/// returns its pages to the allocator's pool rather than to the OS, so
|
||||
/// allocating the next one shows no change. Measured that way, a store that
|
||||
/// holds the corpus twice and one that holds it once look identical.
|
||||
struct CountingAllocator;
|
||||
|
||||
static LIVE_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);
|
||||
|
||||
// SAFETY: every method forwards to the system allocator with the same layout
|
||||
// it was given, and only adds bookkeeping around it.
|
||||
unsafe impl std::alloc::GlobalAlloc for CountingAllocator {
|
||||
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
|
||||
let ptr = unsafe { std::alloc::System.alloc(layout) };
|
||||
if !ptr.is_null() {
|
||||
LIVE_BYTES.fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
ptr
|
||||
}
|
||||
|
||||
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
|
||||
LIVE_BYTES.fetch_sub(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
||||
unsafe { std::alloc::System.dealloc(ptr, layout) }
|
||||
}
|
||||
|
||||
unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
|
||||
let new_ptr = unsafe { std::alloc::System.realloc(ptr, layout, new_size) };
|
||||
if !new_ptr.is_null() {
|
||||
LIVE_BYTES.fetch_add(
|
||||
new_size as i64 - layout.size() as i64,
|
||||
std::sync::atomic::Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
new_ptr
|
||||
}
|
||||
}
|
||||
|
||||
#[global_allocator]
|
||||
static ALLOCATOR: CountingAllocator = CountingAllocator;
|
||||
|
||||
/// Live heap bytes right now.
|
||||
fn heap_bytes() -> u64 {
|
||||
LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
|
||||
}
|
||||
|
||||
fn mib(bytes: u64) -> f64 {
|
||||
bytes as f64 / (1 << 20) as f64
|
||||
}
|
||||
|
||||
fn micros(d: Duration) -> f64 {
|
||||
d.as_secs_f64() * 1e6
|
||||
}
|
||||
@@ -450,6 +501,62 @@ fn fusion_study(n: usize) {
|
||||
}
|
||||
}
|
||||
|
||||
/// What an in-memory store costs, stage by stage. The vectors are the floor:
|
||||
/// everything above it is bookkeeping that could in principle be shared.
|
||||
fn bench_footprint(n: usize) {
|
||||
let data = make_dataset(n, 0xF007 ^ n as u64);
|
||||
let mut rng = Rng(11);
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let path = dir.path().join("footprint.h5");
|
||||
|
||||
let base = heap_bytes();
|
||||
let entries: Vec<MemoryEntry> = data
|
||||
.vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| MemoryEntry {
|
||||
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||
embedding: v.clone(),
|
||||
source_channel: "bench".into(),
|
||||
timestamp: i as f64,
|
||||
session_id: format!("s{}", i % 50),
|
||||
tags: format!("t{i}"),
|
||||
})
|
||||
.collect();
|
||||
let after_entries = heap_bytes();
|
||||
|
||||
let mut mem = HDF5Memory::create(MemoryConfig::new(path, "bench", DIM)).unwrap();
|
||||
mem.save_batch(entries).unwrap();
|
||||
let after_store = heap_bytes();
|
||||
|
||||
// First query builds the vector and keyword indexes.
|
||||
std::hint::black_box(mem.hybrid_search(&data.queries[0], "record", 0.7, 0.3, K));
|
||||
let after_indexes = heap_bytes();
|
||||
|
||||
// Reopening is the figure that matters for a long-lived process, and the
|
||||
// only one RSS reports honestly: memory freed when the ingest buffers went
|
||||
// away stays in the allocator's pool, so the stage deltas above understate
|
||||
// what was given back.
|
||||
let path = mem.config().path.clone();
|
||||
drop(mem);
|
||||
let before_open = heap_bytes();
|
||||
let reopened = HDF5Memory::open(&path).unwrap();
|
||||
let after_open = heap_bytes();
|
||||
let loaded = after_open.saturating_sub(before_open);
|
||||
drop(reopened);
|
||||
|
||||
let raw = (n * DIM * 4) as u64;
|
||||
println!(
|
||||
"| {n} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.2}x |",
|
||||
mib(raw),
|
||||
mib(after_entries.saturating_sub(base)),
|
||||
mib(after_store.saturating_sub(after_entries)),
|
||||
mib(after_indexes.saturating_sub(after_store)),
|
||||
mib(loaded),
|
||||
loaded as f64 / raw as f64,
|
||||
);
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
let full = args.iter().any(|a| a == "--full");
|
||||
@@ -485,6 +592,18 @@ fn main() {
|
||||
|
||||
let mut json = Vec::new();
|
||||
println!("## Search harness");
|
||||
|
||||
if args.iter().any(|a| a == "--footprint") {
|
||||
println!("\n### Resident memory, {DIM}-dim f32\n");
|
||||
println!(
|
||||
"| N | vectors (raw) | entries MiB | store MiB | indexes MiB | reopened MiB | reopened / raw |"
|
||||
);
|
||||
println!("|---:|---:|---:|---:|---:|---:|---:|");
|
||||
for &n in sizes {
|
||||
bench_footprint(n);
|
||||
}
|
||||
return;
|
||||
}
|
||||
// `--e2e-only` skips the index benchmarks, so the end-to-end section runs
|
||||
// in a process that has not already spun up a thread pool.
|
||||
if !args.iter().any(|a| a == "--e2e-only") {
|
||||
|
||||
Reference in New Issue
Block a user