examples/embed_cpu_bakeoff.rs
//! **tract CPU vs onnxruntime CPU — the measurement the design decision never had.**
//!
//! `src/vector/select.rs` refuses ort for CPU work by construction, on the stated
//! grounds that "there is no reason to pull the C++ runtime in to do CPU work tract
//! already does well". That is an *assumption*: as of 2026-08-14 there is no test and
//! no bench in this repo that times the real embedder on either backend. The only
//! datapoint that exists is prose in `.nornir/vector.md:435` — tract ~1900 ms/chunk vs
//! ort CPU ~337 ms/chunk (**5.6× for ort**) — recorded 2026-06-04 on a Zen 2
//! Threadripper 3975WX, *before* tract's batching landed (`fef7c70`, 2026-07-02,
//! "~2-4x CPU"), and never re-measured. It also cannot be reproduced by the harness
//! that produced it: `examples/vector_timetravel.rs` goes through `load_embedder()`,
//! which now routes via `select::load()` and so can no longer reach ort on CPU.
//!
//! This example settles it, in ONE process, on the SAME chunks, with the SAME
//! tokenizer and model. It bypasses the selector rather than weakening it:
//! [`OrtEmbedder::load_cpu`] is already `pub` and passes `prefer_gpu = false`, so no
//! change to `select.rs` is needed and the production gate stays exactly as it is.
//!
//! ## Two arms, because they stress different things
//!
//! * **batch throughput** — the populate path. tract parallelises ACROSS batches via
//! gatling (`embed.rs:231`); ort parallelises INSIDE one `run` via intra-op threads.
//! * **single-query latency** — the interactive path, and the asymmetry that matters
//! most: `embed.rs:203` short-circuits `n == 1` to `embed_one`, and tract-linalg
//! defaults to `SingleThread` (the `multithread-mm` feature is never enabled here),
//! so a one-text embed runs on ONE core against ort's whole intra-op pool.
//!
//! It also asserts the two backends AGREE (cosine ≈ 1.0 per row), so a speed number
//! can never be bought with a wrong answer.
//!
//! ## Running it
//!
//! ort is `load-dynamic`, so no install and no root — just point it at a CPU-only
//! onnxruntime ≥ 1.24 (ort 2.0.0-rc.13 requests `api-24`):
//!
//! ```text
//! ORT_DYLIB_PATH=/path/to/libonnxruntime.so \
//! cargo run --release --example embed_cpu_bakeoff -- [n_chunks]
//! ```
//!
//! Needs `--features embed-tract,embed-ort` (both are in `default`). Use `--release`:
//! a debug tract build is not the thing anyone ships and would flatter ort enormously.
#[cfg(not(all(feature = "embed-tract", feature = "embed-ort")))]
fn main() {
eprintln!(
"embed_cpu_bakeoff needs BOTH backends compiled: \
--features embed-tract,embed-ort (both are in `default`)"
);
std::process::exit(2);
}
#[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
fn main() -> anyhow::Result<()> {
use nornir::vector::chunk::{ChunkOptions, chunk_file};
use nornir::vector::embed::JinaEmbedder;
use nornir::vector::embed_ort::OrtEmbedder;
use nornir::vector::embed_support as es;
use nornir::vector::store::Embedder;
use std::time::Instant;
let want: usize = std::env::args()
.nth(1)
.and_then(|s| s.parse().ok())
.unwrap_or(64);
// ── the corpus: this repo's own sources ─────────────────────────────────────
// Real code with real prose in it. Deliberately NOT `bench-corpus/`, which is
// generated filler (`pub fn f0() -> u32 { .. }`) and would make any embedder look
// identical to any other.
let mut chunks = Vec::new();
let opts = ChunkOptions::default();
let mut files: Vec<_> = walk_rs("src")?;
files.sort();
for path in files {
if chunks.len() >= want {
break;
}
let Ok(text) = std::fs::read_to_string(&path) else {
continue;
};
chunks.extend(chunk_file(&path, &text, &opts));
}
chunks.truncate(want);
anyhow::ensure!(!chunks.is_empty(), "no chunks collected from src/");
let texts: Vec<&str> = chunks.iter().map(|c| c.text.as_str()).collect();
let total_bytes: usize = texts.iter().map(|t| t.len()).sum();
println!("model : {} ({} dims)", es::model_name(), es::dim());
println!("model dir : {}", es::model_dir().display());
println!(
"corpus : {} chunks / {} KiB from this repo's src/",
texts.len(),
total_bytes / 1024
);
match std::env::var("ORT_DYLIB_PATH") {
Ok(p) if !p.is_empty() => println!("ORT_DYLIB : {p}"),
_ => println!("ORT_DYLIB : (unset — the ort arm will fail; see this file's header)"),
}
println!(
"threads : {} logical CPUs\n",
std::thread::available_parallelism().map_or(0, |n| n.get())
);
// ── load both, timed (a one-off cost, but an honest one to report) ──────────
let t = Instant::now();
let tract = JinaEmbedder::load()?;
let tract_load = t.elapsed();
let t = Instant::now();
let ort = OrtEmbedder::load_cpu()?;
let ort_load = t.elapsed();
// Both backends must claim the SAME profile — they are the same weights, and the
// warehouse dedup key depends on it. If this ever differs, vectors would silently
// land under two profiles for one model.
anyhow::ensure!(
tract.profile().id() == ort.profile().id(),
"backends disagree on model_profile: tract {} vs ort {} — vectors would split",
tract.profile().id(),
ort.profile().id()
);
// ── arm 1: batch throughput (the populate path) ─────────────────────────────
let t = Instant::now();
let v_tract = tract.embed(&texts)?;
let tract_batch = t.elapsed();
let t = Instant::now();
let v_ort = ort.embed(&texts)?;
let ort_batch = t.elapsed();
// ── arm 2: single-query latency (the interactive path) ──────────────────────
// Median of a few, because this is the number a user feels and one sample of it is
// noise. `embed.rs:203` sends n == 1 down `embed_one` — one core, one GEMM.
let probe = "a function that batches texts by token length so padding is not wasted";
let mut lat_tract = Vec::new();
let mut lat_ort = Vec::new();
for _ in 0..5 {
let t = Instant::now();
let _ = tract.embed(&[probe])?;
lat_tract.push(t.elapsed().as_secs_f64() * 1e3);
let t = Instant::now();
let _ = ort.embed(&[probe])?;
lat_ort.push(t.elapsed().as_secs_f64() * 1e3);
}
lat_tract.sort_by(f64::total_cmp);
lat_ort.sort_by(f64::total_cmp);
let (p50_tract, p50_ort) = (lat_tract[lat_tract.len() / 2], lat_ort[lat_ort.len() / 2]);
// ── agreement: a speed win must not be a different answer ───────────────────
let mut worst = f32::INFINITY;
for (a, b) in v_tract.iter().zip(v_ort.iter()) {
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
worst = worst.min(dot);
}
// ── report (col-1 id · col-2 status · col-3 detail) ─────────────────────────
let n = texts.len() as f64;
let ms = |d: std::time::Duration| d.as_secs_f64() * 1e3;
let speedup = ms(tract_batch) / ms(ort_batch);
let lat_speedup = p50_tract / p50_ort;
println!(
"load-tract {:>8.0} ms one-off model load (tract-onnx graph build)",
ms(tract_load)
);
println!(
"load-ort {:>8.0} ms one-off model load (onnxruntime session commit)",
ms(ort_load)
);
println!(
"batch-tract {:>8.0} ms {:.1} ms/chunk · {:.1} chunk/s (gatling across batches)",
ms(tract_batch),
ms(tract_batch) / n,
n / tract_batch.as_secs_f64()
);
println!(
"batch-ort {:>8.0} ms {:.1} ms/chunk · {:.1} chunk/s (intra-op threads in one run)",
ms(ort_batch),
ms(ort_batch) / n,
n / ort_batch.as_secs_f64()
);
println!(
"BATCH {:>8.2}x {} on batch throughput",
if speedup >= 1.0 {
speedup
} else {
1.0 / speedup
},
if speedup >= 1.0 {
"ORT WINS"
} else {
"TRACT WINS"
}
);
println!(
"lat-tract {:>8.1} ms p50 single-text embed (ONE core — embed.rs:203 n==1 path)",
p50_tract
);
println!("lat-ort {:>8.1} ms p50 single-text embed", p50_ort);
println!(
"LATENCY {:>8.2}x {} on single-query latency",
if lat_speedup >= 1.0 {
lat_speedup
} else {
1.0 / lat_speedup
},
if lat_speedup >= 1.0 {
"ORT WINS"
} else {
"TRACT WINS"
}
);
println!(
"agreement {:>8.4} worst per-row cosine between the two backends (want ~1.0)",
worst
);
// Same weights, same pooling, both L2-normalized: disagreement beyond float noise
// means one of them is wrong, and then no timing above means anything.
anyhow::ensure!(
worst > 0.999,
"backends DISAGREE (worst cosine {worst:.6}) — the timings above are meaningless"
);
Ok(())
}
/// Collect `.rs` files under `dir`, recursively. Std-only; no walkdir dep for an example.
#[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
fn walk_rs(dir: &str) -> anyhow::Result<Vec<String>> {
let mut out = Vec::new();
let mut stack = vec![std::path::PathBuf::from(dir)];
while let Some(d) = stack.pop() {
for e in std::fs::read_dir(&d)?.flatten() {
let p = e.path();
if p.is_dir() {
stack.push(p);
} else if p.extension().is_some_and(|x| x == "rs") {
out.push(p.to_string_lossy().into_owned());
}
}
}
Ok(out)
}