From daa449a457fc9a2b3c1c09cdec74e99d8087a3f7 Mon Sep 17 00:00:00 2001 From: Vlad Poenaru Date: Wed, 11 Nov 2026 09:06:58 -0800 Subject: [PATCH] lore: vectorize emails 64 at a time On a production-sized lore database, tokenizing at most 64 emails per model call cuts cold `--vectors` peak memory from 32-36 GB to 22 GB on main, and from 22-27 GB to 10.4-10.6 GB with #112 and #113 applied, with the same vectors and no added wall time. update_lore_vectors() starts one worker per CPU, and each worker passed its whole RecordBatch, up to 1,024 emails, to vectorize_batch() in one call. model2vec's chunk size defaults to num_cpus * 128, larger than any batch, so the tokenizer held the untruncated encoding of every email in the batch until it had pooled them all. A lance batch never spans fragments, so the peak follows fragment size. Main writes lore in small fragments, where the gain is modest (4.3-5.1 GB to 4.2 GB on two archives); #112's larger inserts, and any compaction of the lore table, write large ones, where 316 workers can hold the whole table at once. The 1,039- and 93-fragment databases measured below were built with earlier revisions of #112 and #113, the 632-fragment one by main. Each worker now encodes its batch 64 emails at a time. model2vec pools every email on its own, and the tokenizer of the model that scripts/nomic2vec.py produces adds no batch padding, so the vectors do not depend on the chunking. A call that returns the wrong number of vectors is an error rather than a silent shift of every later vector onto the wrong message_id. The global SEMCODE_BATCH_SIZE default stays as it is, because it also sets the chunk size for code vectors; it still applies inside each lore call, so it can lower the lore chunk but not raise it. Measured with `semcode-index --vectors -j 32` on reflinked copies, on a shared 316-CPU host (1-minute load 66-227). "#112 + #113" is main merged with the PR heads 0fad419 and 4d66a45. peak RSS (GB) wall (s) production-sized: 995,564 emails, 1,039 fragments main 35.85, 32.24 931, 967 main + this commit 22.09, 22.09 932, 936 main + #112 + #113 27.03, 22.21 167, 307 main + #112 + #113 + this 10.38, 10.56 154, 256 two archives: 104,189 emails, 93 fragments main 7.85, 8.09 25, 23 main + this commit 3.30, 3.38 23, 23 two archives: 104,189 emails, 632 fragments main 4.33, 5.09 38, 39 main + this commit 4.20, 4.23 33, 36 The sorted (message_id, vector) digest of lore_vectors matches one call per batch on the production-sized and 93-fragment databases. Main's 930 s production wall time comes from elsewhere in the run: with #112 and #113 the same run takes 154-307 s. Of the gap between 22 GB on main and 10.4-10.6 GB with both PRs, about 6 GB is #112's cap of 32 tokio workers: #113 alone with this commit peaks at 16.7-17.1 GB, and at 10.8 GB with TOKIO_WORKER_THREADS=32. --- src/database/search.rs | 117 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 115 insertions(+), 2 deletions(-) diff --git a/src/database/search.rs b/src/database/search.rs index 56458cf..57bf631 100644 --- a/src/database/search.rs +++ b/src/database/search.rs @@ -3044,8 +3044,11 @@ impl VectorSearchManager { .map(|(_, text)| text.as_str()) .collect(); - let vector_results = - vectorizer_for_batch.vectorize_batch(&texts)?; + let vector_results = vectorize_in_chunks( + &texts, + LORE_VECTORIZE_CHUNK, + |part| vectorizer_for_batch.vectorize_batch(part), + )?; // Combine with message IDs let entries: Vec = emails_to_vectorize @@ -3145,6 +3148,43 @@ impl VectorSearchManager { } } +/// Emails each `update_lore_vectors()` worker hands the model in one call. +/// +/// The tokenizer keeps the full, untruncated encoding of every email in a call +/// until it has pooled them all, and one worker runs per CPU, so peak memory +/// grows with workers times this number rather than with the size of the +/// RecordBatches the lore scan returns (up to 1,024 rows). model2vec pools +/// each email on its own, and the tokenizer of the model +/// `scripts/nomic2vec.py` produces adds no batch padding, so the vectors do +/// not depend on it. `SEMCODE_BATCH_SIZE` (also set by `-j threads:batch`) +/// still applies inside each call, so it can lower the effective chunk below +/// this but not raise it. +const LORE_VECTORIZE_CHUNK: usize = 64; + +/// Encode `texts` with `encode`, at most `chunk` texts per call, and return +/// one vector per text in input order. +/// +/// Fails if a call returns a different number of vectors than texts it was +/// given: the results are concatenated by position, so a short chunk would +/// pair every later vector with the wrong email. +fn vectorize_in_chunks(texts: &[&str], chunk: usize, mut encode: F) -> Result>> +where + F: FnMut(&[&str]) -> Result>>, +{ + let mut vectors = Vec::with_capacity(texts.len()); + for part in texts.chunks(chunk) { + let part_vectors = encode(part)?; + anyhow::ensure!( + part_vectors.len() == part.len(), + "vectorizer returned {} vectors for {} texts", + part_vectors.len(), + part.len() + ); + vectors.extend(part_vectors); + } + Ok(vectors) +} + /// Static helper for inserting lore vectors with a cached table handle async fn insert_lore_vectors_batch_with_table( lore_vectors_table: &lancedb::table::Table, @@ -3744,3 +3784,76 @@ impl VectorSearchManager { Ok(results) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// A deterministic stand-in for the model: each text's vector depends on + /// that text alone, as model2vec's mean pooling does. + fn stub_vector(text: &str) -> Vec { + vec![text.len() as f32, text.bytes().map(f32::from).sum()] + } + + fn texts(n: usize) -> Vec { + (0..n).map(|i| format!("email {i}")).collect() + } + + #[test] + fn lore_vectors_chunks_match_one_call_in_input_order() { + for n in [0, 1, 63, 64, 65, 128, 1000] { + let owned = texts(n); + let refs: Vec<&str> = owned.iter().map(String::as_str).collect(); + let mut call_sizes = Vec::new(); + let chunked = vectorize_in_chunks(&refs, 64, |part| { + call_sizes.push(part.len()); + Ok(part.iter().map(|t| stub_vector(t)).collect()) + }) + .unwrap(); + let whole: Vec> = refs.iter().map(|t| stub_vector(t)).collect(); + assert_eq!(chunked, whole, "n = {n}"); + assert_eq!(call_sizes.iter().sum::(), n, "n = {n}"); + assert_eq!(call_sizes.len(), n.div_ceil(64), "n = {n}"); + assert!(call_sizes.iter().all(|&s| (1..=64).contains(&s)), "n = {n}"); + } + } + + #[test] + fn lore_vectors_chunks_fail_on_short_result() { + let owned = texts(100); + let refs: Vec<&str> = owned.iter().map(String::as_str).collect(); + let err = vectorize_in_chunks(&refs, 64, |part| { + Ok(part.iter().skip(1).map(|t| stub_vector(t)).collect()) + }) + .unwrap_err(); + assert!(err.to_string().contains("63 vectors for 64 texts"), "{err}"); + } + + #[test] + fn lore_vectors_chunks_fail_on_long_result() { + let owned = texts(100); + let refs: Vec<&str> = owned.iter().map(String::as_str).collect(); + let err = vectorize_in_chunks(&refs, 64, |part| { + Ok((0..=part.len()).map(|_| stub_vector("extra")).collect()) + }) + .unwrap_err(); + assert!(err.to_string().contains("65 vectors for 64 texts"), "{err}"); + } + + #[test] + fn lore_vectors_chunks_propagate_encoder_error() { + let owned = texts(130); + let refs: Vec<&str> = owned.iter().map(String::as_str).collect(); + let mut calls = 0; + let err = vectorize_in_chunks(&refs, 64, |part| { + calls += 1; + if calls == 2 { + anyhow::bail!("encoder failed"); + } + Ok(part.iter().map(|t| stub_vector(t)).collect()) + }) + .unwrap_err(); + assert_eq!(err.to_string(), "encoder failed"); + assert_eq!(calls, 2); + } +}