examples/vector_timetravel.rs
//! Full-vectorization + time-travel benchmark.
//!
//! Indexes the **entire** Rust source of one or more repos with the real
//! jina-code embedder, simulates a never-pushed local change (a fresh,
//! distinct git_sha that only re-embeds the touched chunks), and **measures
//! the time to load a vector "blob"** — i.e. `store::reconstruct`, which
//! rebuilds the in-memory `VectorIndex` for a sha purely from the warehouse
//! (no re-embed, no git walk).
//!
//! Run (needs the build-time model cache + an embedder feature):
//! cargo run --release --features embed-tract --example vector_timetravel
//! cargo run --release --features embed-ort --example vector_timetravel
//! cargo run --release --features embed-tract --example vector_timetravel -- /path/a /path/b
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
fn main() {
eprintln!(
"vector_timetravel needs an embedder backend: \
re-run with `--features embed-tract` (CPU) or `--features embed-ort` (GPU/CPU)."
);
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
use real::main;
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
mod real {
use std::path::PathBuf;
use std::time::{Duration, Instant};
use anyhow::Result;
use nornir::vector::chunk::{ChunkOptions, chunk_text};
use nornir::vector::store::{RepoRef, collect_rust_sources, index_repo, reconstruct, search};
use nornir::warehouse::iceberg::IcebergWarehouse;
fn short(sha: &str) -> &str {
&sha[..sha.len().min(12)]
}
pub fn main() -> Result<()> {
let args: Vec<String> = std::env::args().skip(1).collect();
let repos: Vec<PathBuf> = if args.is_empty() {
vec![
PathBuf::from(format!("/run/media/{}/T9/git/holger", std::env::var("USER").unwrap_or_default())),
PathBuf::from(format!("/run/media/{}/T9/git/znippy", std::env::var("USER").unwrap_or_default())),
]
} else {
args.iter().map(PathBuf::from).collect()
};
// Fresh warehouse in a temp dir (no dev-dep tempfile in examples).
let whdir = std::env::temp_dir().join(format!("nornir-vecbench-{}", std::process::id()));
std::fs::create_dir_all(&whdir)?;
let wh = IcebergWarehouse::open(&whdir)?;
let t = Instant::now();
// Picks ort (GPU-capable) when `embed-ort` is on, else pure-Rust tract.
let embedder = nornir::vector::load_embedder()?;
println!(
"embedder backend: {} | model: {} (load {:.2?})",
nornir::vector::embedder_backend(),
nornir::vector::selected_model_desc(),
t.elapsed()
);
#[cfg(feature = "embed-ort")]
{
let c = nornir::vector::cuda::ensure();
println!(
"CUDA preload: {} libs from {} dir(s); cuDNN={} → GPU {}",
c.loaded.len(),
c.dirs.len(),
c.cudnn,
if c.cudnn {
"ENGAGED"
} else {
"unavailable (CPU fallback)"
}
);
}
println!(
"active SIMD for the in-memory index: {}",
nornir::vector::active_simd()
);
let mp = embedder.profile().id();
let opts = ChunkOptions::default();
for repo in &repos {
let name = repo
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| repo.display().to_string());
let files = collect_rust_sources(repo);
if files.is_empty() {
println!(
"\n=== {name}: no .rs files at {} — skipping ===",
repo.display()
);
continue;
}
let nchunks: usize = files
.iter()
.map(|(f, c)| chunk_text(f, c, &opts).len())
.sum();
let sha = nornir::gitio::head_sha(repo).unwrap_or_else(|_| "local-head".into());
println!(
"\n=== {name}: {} files, {nchunks} chunks, HEAD {} ===",
files.len(),
short(&sha)
);
// 1. FULL vectorization at HEAD.
let t = Instant::now();
let snap = index_repo(
&wh,
&RepoRef {
workspace: "bench",
repo: &name,
git_sha: &sha,
branch: "main",
complete: true,
},
&files,
&opts,
&*embedder,
)?;
let embed_dt = t.elapsed();
println!(
" full vectorize @HEAD: {:.2?} ({} occurrences, {} vectors embedded)",
embed_dt, snap.occurrences, snap.new_vectors
);
if snap.new_vectors > 0 {
println!(
" => {:.1} ms / embedded chunk",
embed_dt.as_secs_f64() * 1000.0 / snap.new_vectors as f64
);
}
// 2. Never-pushed local change: edit one file in-memory, index at a
// distinct (dirty) sha. Only the touched chunks should re-embed.
let mut dirty = files.clone();
dirty[0].1.push_str(
"\n// nornir: simulated uncommitted change\nfn __nornir_dirty_marker() {}\n",
);
let dsha = format!("{}-dirty", short(&sha));
let t = Instant::now();
let dsnap = index_repo(
&wh,
&RepoRef {
workspace: "bench",
repo: &name,
git_sha: &dsha,
branch: "main",
complete: true,
},
&dirty,
&opts,
&*embedder,
)?;
println!(
" index never-pushed change @{dsha}: {:.2?} ({} re-embedded of {} occurrences)",
t.elapsed(),
dsnap.new_vectors,
dsnap.occurrences
);
// 3. LOAD the vector blob (reconstruct) — the headline measurement.
let mut best = Duration::MAX;
let mut sum = Duration::ZERO;
let runs = 5;
let mut idx_len = 0;
for _ in 0..runs {
let t = Instant::now();
let r = reconstruct(&wh, &name, Some(&sha), &mp)?;
let e = t.elapsed();
idx_len = r.index.len();
best = best.min(e);
sum += e;
}
println!(
" >>> LOAD vector blob @HEAD (reconstruct {idx_len} vectors): best {:.2?}, avg {:.2?} over {runs}",
best,
sum / runs
);
let t = Instant::now();
let rd = reconstruct(&wh, &name, Some(&dsha), &mp)?;
println!(
" >>> LOAD vector blob @dirty (reconstruct {} vectors): {:.2?}",
rd.index.len(),
t.elapsed()
);
// 4. A sample semantic query at HEAD.
let query = "parallel compression and decompression of files";
// `Embedder::embed` takes `&[&str]` — a `&[String]` does not coerce, so
// this example failed `--all-targets` (lib/bins/tests never built it).
let q = embedder.embed(&[query])?;
let t = Instant::now();
let hits = search(&wh, &name, Some(&sha), &mp, &q[0], 3)?;
println!(" sample query {query:?} ({:.2?}):", t.elapsed());
for (s, o) in &hits {
println!(" {:.3} {}:{}-{}", s, o.file, o.start_line, o.end_line);
}
}
let _ = std::fs::remove_dir_all(&whdir);
Ok(())
}
}