examples/embed_budget_bench.rs
//! Measure-first harness for the embed budget gate defect (2026-08-26).
//!
//! The live journal on t14s shows every repo REFUSED: `0.14 chunk/s measured on
//! 7 sample chunks` → ~262 min projected for Skidbladnir's 2258 chunks. This
//! bench answers, with real numbers on the real corpus:
//!
//! 1. model LOAD time alone (tract graph optimize is minutes-scale);
//! 2. the gate's own probe (stride sample of 8, `Embedder::estimate`) — the
//! number the refusal quotes;
//! 3. sustained SINGLE-batch rate (what the probe actually times: one
//! `[7, L]` forward on ONE thread, because `threads = min(cores, batches)`);
//! 4. sustained MULTI-batch rate (N chunks spanning many batches, gatling
//! fan-out across all cores — what the real run gets);
//!
//! Run: `cargo run --release --features embed-tract --example embed_budget_bench -- <repo_root> [n_sustained]`
use std::time::Instant;
use nornir::vector::chunk::{ChunkOptions, chunk_file};
use nornir::vector::store::{Embedder, collect_rust_sources};
/// Same striding as `store::stride_sample` (private there): spread `want`
/// picks evenly across `texts`.
fn stride_sample<'a>(texts: &[&'a str], want: usize) -> Vec<&'a str> {
if texts.len() <= want {
return texts.to_vec();
}
let step = texts.len() as f64 / want as f64;
(0..want)
.map(|i| texts[((i as f64 * step) as usize).min(texts.len() - 1)])
.collect()
}
fn main() -> anyhow::Result<()> {
let mut args: Vec<String> = std::env::args().skip(1).collect();
let real_mode = args.iter().any(|a| a == "--real");
args.retain(|a| a != "--real");
let mut args = args.into_iter();
let root = args
.next()
.unwrap_or_else(|| format!("{}/git/Skidbladnir", std::env::var("HOME").unwrap_or_default()));
let n_sustained: usize = args
.next()
.and_then(|s| s.parse().ok())
.unwrap_or(128);
let cores = std::thread::available_parallelism().map(|x| x.get()).unwrap_or(1);
eprintln!("host: {cores} logical cores; corpus root {root}");
// ---- REAL mode: the whole gated write path against a temp warehouse ----
// `index_repos_batched` chunk-dedupes, runs the budget gate, embeds in
// durable slices, and commits the snapshot — exactly what the monitor's
// populate loop does per member. The before/after proof for the refused
// repos: before the fix the gate refused here; after, it prices honestly
// and the member commits (or refuses honestly if it truly cannot fit).
if real_mode {
let repo = std::path::Path::new(&root)
.file_name()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|| "repo".into());
let files = collect_rust_sources(std::path::Path::new(&root));
eprintln!("REAL : indexing {} ({} files) into a temp warehouse …", repo, files.len());
let dir = tempfile::tempdir()?;
let wh = nornir::warehouse::iceberg::IcebergWarehouse::open(dir.path())?;
let embedder = nornir::vector::embed::JinaEmbedder::load()?;
let input = nornir::vector::store::RepoInput {
workspace: "bench".into(),
repo: repo.clone(),
git_sha: "bench-sha".into(),
branch: "main".into(),
complete: true,
files,
};
let t = Instant::now();
let res = nornir::vector::store::index_repos_batched(
&wh,
&[input],
&ChunkOptions::default(),
&embedder,
)
.pop()
.expect("one result for one member");
let s = t.elapsed().as_secs_f64();
match res {
Ok(snap) => eprintln!(
"REAL : {} COMMITTED — {} occurrences, {} new vectors in {:.1}s ({:.1} min)",
repo, snap.occurrences, snap.new_vectors, s, s / 60.0
),
Err(e) => eprintln!("REAL : {} REFUSED/FAILED after {:.1}s: {e:#}", repo, s),
}
return Ok(());
}
// ---- corpus: chunk exactly as the monitor does ----
let files = collect_rust_sources(std::path::Path::new(&root));
let opts = ChunkOptions::default();
let chunks: Vec<_> = files
.iter()
.flat_map(|(p, c)| chunk_file(p, c, &opts))
.collect();
let texts: Vec<&str> = chunks.iter().map(|c| c.text.as_str()).collect();
eprintln!("corpus: {} files -> {} chunks", files.len(), chunks.len());
// ---- 1. LOAD time alone ----
let t = Instant::now();
let e = nornir::vector::embed::JinaEmbedder::load()?;
let load_s = t.elapsed().as_secs_f64();
eprintln!("LOAD : {load_s:.1}s (tract load + optimize + runnable)");
// ---- warm-up forward (first-touch page faults etc.) ----
let t = Instant::now();
let _ = e.embed(&texts[..1])?;
eprintln!("WARMUP: {:.1}s (first single-chunk forward)", t.elapsed().as_secs_f64());
// ---- 2. the gate's own probe, verbatim ----
let sample = stride_sample(&texts, 8);
// token lengths of the sample rows (what the probe actually pays for)
let t = Instant::now();
let eta = e.estimate(&sample, texts.len())?;
eprintln!(
"GATE : {:.1}s probe -> {:.2} chunk/s on {} timed -> eta {:.1} min for {} chunks",
t.elapsed().as_secs_f64(),
eta.rate,
eta.sampled,
eta.eta.as_secs_f64() / 60.0,
eta.total_chunks
);
// ---- 3. single-batch sustained (what the probe times) ----
let seven: Vec<&str> = sample[1..].to_vec();
let t = Instant::now();
let _ = e.embed(&seven)?;
let s = t.elapsed().as_secs_f64();
eprintln!(
"1BATCH: {:.1}s for {} chunks -> {:.2} chunk/s (threads limited by batch count)",
s,
seven.len(),
seven.len() as f64 / s
);
// ---- 4. multi-batch sustained (the real run's shape) ----
let many = stride_sample(&texts, n_sustained);
let t = Instant::now();
let _ = e.embed(&many)?;
let s = t.elapsed().as_secs_f64();
let rate = many.len() as f64 / s;
eprintln!(
"NBATCH: {:.1}s for {} chunks -> {:.2} chunk/s -> full corpus {} chunks in {:.1} min",
s,
many.len(),
rate,
texts.len(),
texts.len() as f64 / rate / 60.0
);
Ok(())
}