src/vector/select.rs
//! Runtime embedder backend selection (#9) — ONE binary, picked at runtime.
//!
//! A default `cargo install nornir` ships **both** embedders: the pure-Rust
//! **tract** CPU backend ([`super::embed`]) and the GPU-capable **ort** backend
//! ([`super::embed_ort`], CUDA/ROCm). This module is the single place that
//! decides which one to load on the *running* machine, gracefully:
//!
//! ```text
//! probe onnxruntime present? ── no ──▶ tract (CPU, pure Rust, needs NOTHING)
//! │ yes
//! probe NVIDIA CUDA ready? ── yes ─▶ ort CUDA (NVIDIA GPU)
//! │ no
//! probe AMD ROCm ready?
//! AND the model FITS in
//! DEDICATED VRAM? ── yes ─▶ ort ROCm (AMD GPU / APU)
//! │ no
//! └──────────────────────────▶ ort CPU (MLAS — 4.18× tract, measured)
//! ```
//!
//! **Best available, in order, and every rung is measured rather than assumed.**
//!
//! * NVIDIA is taken outright — `cuda_ready` means driver + cudart + cuDNN all resolved,
//! which in practice means a discrete card installed on purpose.
//! * AMD is taken **only if the device can hold the model in memory the CPU is not
//! already using**. An APU's GTT *is* system RAM, so a model that only fits there is
//! being streamed over the bus the CPU would have used anyway — see
//! [`Probe::amd_gpu_is_worth_it`] for the `dwarf` measurement that this encodes
//! (Radeon 680M: 1.7× the FLOPs at 1.0× the bandwidth, and it drove the DIMMs past
//! their 85 °C critical mark while the CPU package sat at 66 °C).
//! * Otherwise **ort's CPU EP**, not tract. This header used to end "…— there is no
//! reason to pull the C++ runtime in to do CPU work tract already does well". That was
//! never measured and is **false** on this hardware: 4.18× on batch throughput, 4.68×
//! on single-query latency, at a worst per-row cosine agreement of 1.0000 — the same
//! vectors, four times sooner. See [`Backend::OrtCpu`] and
//! `examples/embed_cpu_bakeoff.rs`.
//! * tract remains the guaranteed fallback for the case that genuinely needs it: **no
//! `libonnxruntime.so` on the box at all**. That is a real and sufficient reason to
//! keep it, and it is now the only one.
//!
//! `$NORNIR_EMBED_BACKEND` overrides the whole ladder when an operator wants a specific
//! rung (`ort-cuda | ort-rocm | ort-cpu | tract-cpu | static-table`).
//!
//! There is also a fifth backend that no probe can choose:
//! [`Backend::StaticTable`] is selected by the MODEL, because whether a forward pass
//! exists at all is a property of the weights ([`model_backend`], which outranks
//! [`decide`]).
//!
//! ## Why a probe-first gate (build/run safety)
//!
//! `embed-ort` is compiled with ort's **`load-dynamic`** feature, so the binary
//! LINKS NO onnxruntime — it `dlopen`s `libonnxruntime.so` at runtime. The
//! consequence: the *first* call into any ort API ([`ort::session::Session`])
//! triggers ort to load the dylib, and if it is absent ort **panics**. So the
//! selector must NEVER construct an ort object unless [`super::cuda::arm_onnxruntime`]
//! has confirmed a real onnxruntime is loadable. That is exactly what
//! [`decide`] + [`load`] enforce: a box with no onnxruntime (or no GPU) returns
//! [`Backend::TractCpu`] and the ort module is never entered.
//!
//! ## Testability (LAW: inject-assert)
//!
//! [`decide`] is a pure function over a [`Probe`] of booleans, so a test can
//! INJECT a fake probe (GPU-ready / not, onnxruntime present / not) and ASSERT
//! the chosen [`Backend`] — without needing a GPU or onnxruntime in CI. The live
//! path ([`live_probe`]) fills the same struct from the real cuda/rocm probes.
use anyhow::Result;
use super::store::Embedder;
/// The embedder backend the selector chose for this machine.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Backend {
/// ort (ONNX Runtime) with the **NVIDIA CUDA** execution provider.
OrtCuda,
/// ort (ONNX Runtime) with the **AMD ROCm / MIGraphX** execution provider.
OrtRocm,
/// ort (ONNX Runtime) on the **CPU** execution provider (MLAS).
///
/// Added 2026-08-14. This variant did not exist, which is *why* the CPU EP was
/// unreachable: the selector could not name the thing. The module header justified
/// that with "there is no reason to pull the C++ runtime in to do CPU work tract
/// already does well" — an assumption with no measurement behind it.
///
/// ── MEASURED 2026-08-14 on `dwarf` (Ryzen 5 7535HS, 6c/12t, **AVX2+FMA only — no
/// AVX-512, no VNNI**), 64 real chunks of this repo's `src/`, one process, same
/// weights, same tokenizer (`examples/embed_cpu_bakeoff.rs`):
///
/// ```text
/// batch-tract 91715 ms 1433.0 ms/chunk · 0.7 chunk/s
/// batch-ort 21931 ms 342.7 ms/chunk · 2.9 chunk/s → 4.18× FOR ORT
/// lat-tract 78.7 ms p50 single-text embed (ONE core, the n==1 path)
/// lat-ort 16.8 ms p50 → 4.68× FOR ORT
/// agreement 1.0000 worst per-row cosine between the two backends
/// ```
///
/// The claim is therefore **false on this hardware**, and the AVX-512/VNNI escape
/// does not apply: neither ISA is present here. Two details make it conclusive
/// rather than suggestive — the agreement is 1.0000 (identical vectors, so the
/// speed is not bought with a different answer), and ort's 342.7 ms/chunk
/// reproduces `.nornir/vector.md:435`'s 2026-06-04 figure of ~337 ms/chunk almost
/// exactly, on different silicon. tract did improve (1900 → 1433 ms/chunk, the
/// 2026-07-02 batching) — it is still 4× behind.
///
/// **This IS the default whenever a usable `libonnxruntime.so` is present and no
/// worthwhile GPU is** (the owner's ruling, 2026-08-14: "if none or crappy amd gpu
/// pick a cpu impl THAT WORKS and is highly functional"). It was briefly opt-in-only,
/// which was the wrong shape — an operator who has already installed the dylib has
/// asked for ort, and handing them the 4× slower backend anyway is not caution.
///
/// A box with NO dylib is unaffected and still gets [`Backend::TractCpu`], so
/// nothing silently changes for a deployment that never had onnxruntime.
OrtCpu,
/// Pure-Rust **tract** CPU embedder — the guaranteed-present fallback.
TractCpu,
/// A **static token→vector table** (model2vec / potion): no transformer graph at
/// all, so no attention and no feed-forward — the forward pass is a gather plus a
/// mean-pool. Two to three orders of magnitude less arithmetic than any runtime
/// executing the real graph, which is why it is not a "faster backend" so much as
/// a different, cheaper model.
///
/// Orthogonal to the hardware probe by construction: it is chosen by the MODEL
/// (`$NORNIR_EMBED_MODEL` resolving to a [`ModelKind::StaticTable`] entry), never
/// by what silicon is present — see [`model_backend`]. Its vectors carry their own
/// `model_profile`, so they can never be compared against transformer vectors.
///
/// [`ModelKind::StaticTable`]: super::embed_registry::ModelKind::StaticTable
StaticTable,
}
impl Backend {
/// Human-readable name, for `vector doctor` / diagnostics.
pub fn label(self) -> &'static str {
match self {
Backend::OrtCuda => "ort (ONNX Runtime, NVIDIA CUDA GPU)",
Backend::OrtRocm => "ort (ONNX Runtime, AMD ROCm GPU)",
Backend::OrtCpu => "ort (ONNX Runtime, CPU / MLAS)",
Backend::TractCpu => "tract-onnx (CPU, pure Rust)",
Backend::StaticTable => "static-table (model2vec, CPU, pure Rust)",
}
}
/// True if this backend runs on a GPU (ort CUDA or ROCm).
pub fn is_gpu(self) -> bool {
matches!(self, Backend::OrtCuda | Backend::OrtRocm)
}
/// True if this backend executes a real transformer graph. False for
/// [`Backend::StaticTable`], whose vectors are a lookup — the distinction that
/// matters when reporting *why* an embed was fast.
pub fn runs_a_graph(self) -> bool {
!matches!(self, Backend::StaticTable)
}
/// The `$NORNIR_EMBED_BACKEND` token that selects this backend explicitly, for the
/// diagnostics to print back the exact string an operator would have to set.
pub fn env_token(self) -> &'static str {
match self {
Backend::OrtCuda => "ort-cuda",
Backend::OrtRocm => "ort-rocm",
Backend::OrtCpu => "ort-cpu",
Backend::TractCpu => "tract-cpu",
Backend::StaticTable => "static-table",
}
}
/// Parse a `$NORNIR_EMBED_BACKEND` token. `None` for anything unrecognized —
/// callers must treat that as "no request", never as a silent default, so a typo
/// cannot quietly pin the slowest backend.
pub fn from_env_token(s: &str) -> Option<Self> {
match s.trim().to_ascii_lowercase().as_str() {
"ort-cuda" | "cuda" => Some(Backend::OrtCuda),
"ort-rocm" | "rocm" => Some(Backend::OrtRocm),
"ort-cpu" => Some(Backend::OrtCpu),
"tract-cpu" | "tract" => Some(Backend::TractCpu),
"static-table" | "static" => Some(Backend::StaticTable),
_ => None,
}
}
}
/// Env var that pins the backend explicitly, bypassing the hardware probe.
///
/// This exists so the ort-CPU claim is testable **in production**, not only from an
/// example: without it, `Backend::OrtCpu` would be a variant no configuration could
/// reach, which is the same unfalsifiable shape the old comment had.
pub const NORNIR_EMBED_BACKEND_ENV: &str = "NORNIR_EMBED_BACKEND";
/// Headroom an AMD device's DEDICATED VRAM must have over the model's weights before
/// the GPU is worth using: weights plus activations plus the runtime's own arena.
///
/// 1.5× is deliberately modest. It does not need to be precise, because the case it
/// exists to catch is not marginal — see [`Probe::amd_gpu_is_worth_it`].
pub const VRAM_HEADROOM: f64 = 1.5;
/// The runtime facts the selector decides on. A test injects this directly;
/// [`live_probe`] fills it from the real cuda/rocm/onnxruntime probes.
#[derive(Debug, Clone, Copy, Default)]
pub struct Probe {
/// `libonnxruntime.so` is present and loadable (the ort dynamic-load dylib).
pub onnxruntime: bool,
/// NVIDIA driver + CUDA runtime + cuDNN are all present (CUDA EP would run).
pub cuda_ready: bool,
/// AMD driver + HIP runtime + MIOpen are all present (ROCm EP would run).
pub rocm_ready: bool,
/// **DEDICATED** VRAM on the AMD device, in bytes (`mem_info_vram_total`) — NOT
/// GTT. `None` when unreadable.
///
/// GTT is deliberately excluded and that is the whole point: on an APU, GTT *is*
/// system RAM, so a model that only fits in GTT is being streamed over the same
/// memory bus the CPU would have used. That is not acceleration, it is the CPU's
/// bandwidth with extra steps.
pub amd_vram_bytes: Option<u64>,
/// Size of the selected model's weights file in bytes. `None` when unreadable.
pub model_bytes: Option<u64>,
}
impl Probe {
/// **Is this AMD device actually worth using for THIS model?**
///
/// ── MEASURED 2026-08-14 on `dwarf`, which is exactly the case this rejects ──
/// Radeon 680M (gfx1035, "yellow_carp"): 6 CU / 12 SIMD32 @ 1899 MHz ≈ 1.46 TFLOPS
/// fp32, against the host Ryzen 5 7535HS at ≈ 0.85 TFLOPS. Only **1.7×** on paper —
/// and it has **512 MB** of dedicated VRAM (`mem_info_vram_total`, dmesg-confirmed)
/// against 32 GB of GTT. jina-embeddings-v2-base-code is **612 MB** of f32 weights,
/// so it does not fit in VRAM and would stream out of GTT = the DDR5-4800 the CPU
/// already saturates. 1.7× the FLOPs at 1.0× the bandwidth, on a bandwidth-bound
/// workload, is not a win — it is a lateral move that also cooked the DIMMs to
/// 88.5 °C (crit 85) while the CPU package sat at 66 °C.
///
/// So the test is not "how fast is the GPU" — estimating FLOPS across GCN/RDNA/CDNA
/// SIMD widths is exactly the kind of arithmetic that is confidently wrong — it is
/// **"can the model live in memory the CPU is not already using?"** That question is
/// answerable from two numbers we can actually read, and it is the one that decides
/// whether offloading moves work or merely moves it around.
///
/// Unknowns are trusted, NOT rejected: a box where ROCm is installed and ready is a
/// deliberate act, and a discrete card whose VRAM we failed to read should not be
/// refused on the strength of a missing file.
pub fn amd_gpu_is_worth_it(&self) -> bool {
match (self.amd_vram_bytes, self.model_bytes) {
(Some(vram), Some(model)) => (vram as f64) >= (model as f64) * VRAM_HEADROOM,
_ => true,
}
}
}
/// Pure decision function (#9): given the runtime [`Probe`], choose the backend.
///
/// onnxruntime is the gate for *any* ort path; within it NVIDIA wins over AMD
/// (CUDA is the proven path), and absent any usable GPU we fall back to tract
/// CPU even when onnxruntime is present. No I/O, no globals — trivially testable.
pub fn decide(p: &Probe) -> Backend {
if p.onnxruntime {
// NVIDIA first, and WITHOUT a worth-it test. cuda_ready means driver + cudart +
// cuDNN all resolved, which in practice means a discrete card someone installed
// on purpose; there is no NVIDIA analogue of the shared-DRAM APU case below.
if p.cuda_ready {
return Backend::OrtCuda;
}
// AMD second, and only when the device can actually hold the model in memory
// the CPU is not already using. See `Probe::amd_gpu_is_worth_it` for the
// measurement this encodes.
if p.rocm_ready && p.amd_gpu_is_worth_it() {
return Backend::OrtRocm;
}
// ── ort IS here and no GPU is worth using: use ort's CPU EP, not tract. ──
//
// Until 2026-08-14 this fell through to `TractCpu` on the stated grounds that
// there was "no reason to pull the C++ runtime in to do CPU work tract already
// does well". MEASURED on `dwarf` (Ryzen 5 7535HS, AVX2+FMA only — no AVX-512,
// no VNNI, so MLAS's headline kernels are NOT what wins here): ort's CPU EP is
// **4.18× faster on batch throughput and 4.68× on single-query latency**, at a
// worst per-row cosine agreement of 1.0000 — the same vectors, four times
// sooner. `examples/embed_cpu_bakeoff.rs` is the harness.
//
// Reaching this arm means the dylib was already found and armed, so choosing
// tract here would decline a 4× speedup that costs nothing and changes no
// answer. That is not conservatism, it is leaving the win on the floor.
return Backend::OrtCpu;
}
// No onnxruntime anywhere: tract is the pure-Rust backend that needs NOTHING on the
// box, and it is why it stays the guaranteed fallback rather than the preference.
Backend::TractCpu
}
/// Run the real probes on this machine and, as a side effect, **arm** the ort
/// dynamic-load path (set `ORT_DYLIB_PATH` when onnxruntime is found) so a
/// subsequent ort load resolves the same dylib we probed. Pure-CPU boxes leave
/// the env untouched.
pub fn live_probe() -> Probe {
// Arming sets ORT_DYLIB_PATH iff a usable libonnxruntime.so exists; its
// bool answer IS the onnxruntime probe. On an AMD-only box the dylib lives
// in the ROCm dirs (/opt/nornir/rocm, /opt/rocm*), so when the CUDA-dir
// search misses, also try the ROCm discovery dirs.
#[cfg_attr(not(feature = "embed-ort-rocm"), allow(unused_mut))]
let mut onnxruntime = super::cuda::arm_onnxruntime();
#[cfg(feature = "embed-ort-rocm")]
if !onnxruntime {
onnxruntime = super::rocm::arm_onnxruntime();
}
// CUDA readiness: driver + cudart + cuDNN. `preflight` returns (ready, _).
let cuda_ready = super::cuda::preflight().0;
// ROCm readiness only when the AMD feature is compiled; otherwise false.
#[cfg(feature = "embed-ort-rocm")]
let rocm_ready = super::rocm::available();
#[cfg(not(feature = "embed-ort-rocm"))]
let rocm_ready = false;
Probe {
onnxruntime,
cuda_ready,
rocm_ready,
// Only read when an AMD device is actually in play — no point stat'ing sysfs on
// an NVIDIA or CPU-only box, and `amd_gpu_is_worth_it` trusts `None` anyway.
amd_vram_bytes: if rocm_ready {
amd_dedicated_vram()
} else {
None
},
model_bytes: selected_model_weight_bytes(),
}
}
/// Dedicated VRAM of the first AMD render node, in bytes, from
/// `/sys/class/drm/card*/device/mem_info_vram_total`.
///
/// ⚠ Deliberately NOT `mem_info_gtt_total`. On `dwarf` those read 512 MB and 32 GB
/// respectively, and the 32 GB is system RAM the CPU is already using — counting it
/// would make every APU look like it had plenty of memory and re-introduce exactly the
/// bad decision this function exists to prevent.
fn amd_dedicated_vram() -> Option<u64> {
let mut best: Option<u64> = None;
let entries = std::fs::read_dir("/sys/class/drm").ok()?;
for e in entries.flatten() {
let name = e.file_name();
let name = name.to_string_lossy();
// `card0`, not `card0-DP-1` (connectors have no device memory).
if !name.starts_with("card") || name.contains('-') {
continue;
}
let p = e.path().join("device/mem_info_vram_total");
if let Some(v) = std::fs::read_to_string(&p)
.ok()
.and_then(|s| s.trim().parse::<u64>().ok())
{
best = Some(best.map_or(v, |b: u64| b.max(v)));
}
}
best
}
/// Size of the selected model's weights file on disk, in bytes.
///
/// This is the number the VRAM test is against, and taking it from the FILE rather than
/// a constant matters: the registry carries several models whose weights differ by more
/// than an order of magnitude (jina 612 MB vs potion-base-8m 28.8 MB), so "does it fit"
/// is a question about the CURRENTLY SELECTED model, not about embedding in general.
fn selected_model_weight_bytes() -> Option<u64> {
let m = super::embed_support::active_model();
std::fs::metadata(super::embed_support::model_dir().join(m.weights_local_name))
.ok()
.map(|md| md.len())
}
/// **Which backend the MODEL forces, independent of hardware.**
///
/// A static-table model has no graph, so no execution provider can run it and no
/// probe result is relevant: `$NORNIR_EMBED_MODEL=potion-base-8m` on a 4090 still
/// means a gather and a mean-pool. Returning `Some` here therefore *outranks* the
/// hardware probe — and it must, because [`Backend::TractCpu`] would send a static
/// model into `tract_onnx().model_for_path(dir.join("model.onnx"))`, a file
/// `build.rs` never fetched for it.
///
/// Pure (over the registry's const table), so it is inject-assert testable in the
/// same style as [`decide`] — the two pure functions together are the whole policy.
pub fn model_backend(kind: super::embed_registry::ModelKind) -> Option<Backend> {
match kind {
super::embed_registry::ModelKind::StaticTable => Some(Backend::StaticTable),
// A graph model leaves the choice to the hardware probe.
super::embed_registry::ModelKind::Onnx => None,
}
}
/// The operator's explicit `$NORNIR_EMBED_BACKEND` request, if any and if legal.
///
/// A request for a **graph** backend cannot override a static model: the weights for
/// a graph do not exist on disk for that model, so honouring it would be a load
/// error dressed as a configuration. A static request is likewise refused for a graph
/// model. Both refusals are LOUD — a silently ignored pin is how an operator ends up
/// measuring the backend they thought they had turned off.
fn env_backend(kind: super::embed_registry::ModelKind) -> Option<Backend> {
let raw = std::env::var(NORNIR_EMBED_BACKEND_ENV).ok()?;
if raw.trim().is_empty() {
return None;
}
let Some(want) = Backend::from_env_token(&raw) else {
eprintln!(
"⚠ {NORNIR_EMBED_BACKEND_ENV}={raw:?} is not a known backend — IGNORED (no pin \
applied). Valid: ort-cuda | ort-rocm | ort-cpu | tract-cpu | static-table."
);
return None;
};
let forced = model_backend(kind);
match (forced, want) {
// Static model + static request, or graph model + graph request: coherent.
(Some(Backend::StaticTable), Backend::StaticTable) => Some(want),
(None, w) if w != Backend::StaticTable => Some(w),
_ => {
eprintln!(
"⚠ {NORNIR_EMBED_BACKEND_ENV}={} REFUSED — the selected model is a {} model \
and that backend cannot run it. Change $NORNIR_EMBED_MODEL instead; a pin \
that cannot be honoured is worse than none.",
want.env_token(),
if forced.is_some() {
"static-table"
} else {
"graph"
},
);
None
}
}
}
/// **Was a non-GPU backend explicitly ASKED for**, rather than fallen back to?
///
/// True when an honoured `$NORNIR_EMBED_BACKEND` pin names a non-GPU backend, or the
/// selected model is a static table (which no GPU can run). This is what lets
/// [`super::fallback::diagnose`] tell a deliberate CPU choice apart from a degradation:
/// `ORT_DYLIB_PATH` is evidence of GPU intent *and* the way one enables ort-CPU, so
/// intent alone can no longer decide.
pub fn cpu_was_requested() -> bool {
let kind = super::embed_support::active_model().kind;
if model_backend(kind).is_some_and(|b| !b.is_gpu()) {
return true;
}
env_backend(kind).is_some_and(|b| !b.is_gpu())
}
/// The backend [`load`] would choose on this machine right now (probes the real
/// hardware, and honours the model kind + any `$NORNIR_EMBED_BACKEND` pin). Cheap;
/// used by `vector doctor` and diagnostics to report the runtime choice without
/// actually loading a model.
pub fn chosen_backend() -> Backend {
let kind = super::embed_support::active_model().kind;
// Precedence: an honoured operator pin, then the model's own demand, then the
// hardware probe. The pin is checked first so the ort-CPU claim can actually be
// measured on a live server; `env_backend` has already refused any incoherent one.
if let Some(b) = env_backend(kind) {
return b;
}
if let Some(b) = model_backend(kind) {
return b;
}
decide(&live_probe())
}
/// Why a bulk embed NOBODY IS WAITING FOR should not run on this box — `None`
/// = run it. Pure: a test injects the three facts.
///
/// The ladder above always ends on a CPU backend that works, and for a query or
/// an embed somebody asked for that is the right answer. An unattended bulk
/// embed (the consumer's `embeddings` feed, once per member commit) is not that
/// case: on a CPU it is minutes per member, with the workspace gate held. It
/// runs where the selector found a GPU WORTH using — NVIDIA, or an AMD device
/// that holds the model in dedicated VRAM; an APU's shared memory does not
/// count ([`Probe::amd_gpu_is_worth_it`]) — or where CPU was ASKED for: a
/// `$NORNIR_EMBED_BACKEND` pin, a static-table model ([`cpu_was_requested`]).
pub fn unattended_embed_skip(chosen: Backend, cpu_requested: bool, p: &Probe) -> Option<String> {
if chosen.is_gpu() || cpu_requested {
return None;
}
let mb = |bytes: u64| bytes >> 20;
let found = match (p.amd_vram_bytes, p.model_bytes) {
_ if !p.onnxruntime => "no onnxruntime on this box, so no GPU backend can load".to_string(),
(Some(vram), Some(model)) if p.rocm_ready && !p.amd_gpu_is_worth_it() => format!(
"the AMD GPU has {} MB of dedicated VRAM and the model needs {} MB ({VRAM_HEADROOM}× its {} MB of weights) — \
the memory it shares with the CPU does not count",
mb(vram),
mb((model as f64 * VRAM_HEADROOM) as u64),
mb(model),
),
_ => "no NVIDIA CUDA runtime and no AMD ROCm runtime is ready".to_string(),
};
Some(format!(
"no embedding hardware here: {found}; the best this box has is {} — {NORNIR_EMBED_BACKEND_ENV}={} asks for it by name",
chosen.label(),
chosen.env_token(),
))
}
/// [`unattended_embed_skip`] for this machine, right now.
pub fn live_unattended_embed_skip() -> Option<String> {
unattended_embed_skip(chosen_backend(), cpu_was_requested(), &live_probe())
}
/// Load the embedder the runtime selector chose: pin → model kind → probe →
/// [`decide`] → construct.
///
/// GPU backends are best-effort *inside* ort too (the EP falls back to CPU if a
/// provider fails to register), but we only ever enter ort when [`decide`] said
/// onnxruntime + a GPU are present. Any ort load error degrades to tract CPU so
/// the binary never fails to produce an embedder when tract is available.
pub fn load() -> Result<Box<dyn Embedder>> {
load_backend(chosen_backend())
}
/// Construct the embedder for an explicitly chosen [`Backend`], degrading to
/// tract CPU if an ort GPU load fails (so selection is always graceful).
///
/// The GPU arm uses [`super::embed_ort::OrtEmbedder::load_gpu_strict`], NOT the
/// best-effort [`load`](super::embed_ort::OrtEmbedder::load). Reaching this arm
/// means [`decide`] already saw driver + cudart + cuDNN, so an EP that then fails
/// to register is a genuine fault; ort's default `fail_silently` would swallow it
/// and serve every `run` on the CPU EP from inside a session still labelled CUDA
/// — a second, deeper copy of the silent fallback that left oden's forge embedding
/// ~60× slow with nothing in the log naming a cause. We still degrade to tract so
/// the binary always produces an embedder, but never quietly.
pub fn load_backend(backend: Backend) -> Result<Box<dyn Embedder>> {
match backend {
Backend::OrtCuda | Backend::OrtRocm => {
match super::embed_ort::OrtEmbedder::load_gpu_strict() {
Ok(e) => Ok(Box::new(e)),
Err(e) => {
// GPU/ort path failed at runtime — fall back to the always-present
// pure-Rust CPU embedder rather than erroring out, but SAY SO with
// the cause and a tag a log grep can find.
eprintln!(
"⚠ GPU EMBEDDING SKIPPED [{}] — the GPU probes all passed \
(driver + CUDA runtime + cuDNN) yet the {} execution provider \
could not be registered: {e:#}. Falling back to tract CPU, which \
was measured ~60× slower on this constellation's RTX 4090. This is \
a real fault (usually onnxruntime built against a different CUDA \
major than the installed runtime), not a GPU-less box.",
super::fallback::EP_REGISTRATION_FAILED,
backend.label(),
);
Ok(Box::new(super::embed::JinaEmbedder::load()?))
}
}
}
// ort on the CPU EP (MLAS). NOT best-effort and NOT degraded to tract: reaching
// this arm means somebody explicitly pinned it, and silently serving a different
// backend than the one under measurement is the exact defect `fallback.rs` exists
// to stop. A missing/too-old libonnxruntime.so is therefore an ERROR here.
Backend::OrtCpu => Ok(Box::new(super::embed_ort::OrtEmbedder::load_cpu()?)),
Backend::TractCpu => Ok(Box::new(super::embed::JinaEmbedder::load()?)),
// Static table: no graph, no EP, no probe. Also not degraded — a static model's
// vectors are a DIFFERENT model_profile, so quietly substituting tract would
// write jina vectors under a run the operator believes was potion.
Backend::StaticTable => Ok(Box::new(super::embed_static::StaticEmbedder::load()?)),
}
}
#[cfg(test)]
mod tests {
use super::*;
/// LAW (inject-assert): when the probe says onnxruntime + a usable NVIDIA
/// GPU are present, the selector MUST pick the ort CUDA backend.
#[test]
fn picks_ort_cuda_when_gpu_ready() {
let p = Probe {
onnxruntime: true,
cuda_ready: true,
rocm_ready: false,
..Default::default()
};
assert_eq!(decide(&p), Backend::OrtCuda);
assert!(
decide(&p).is_gpu(),
"CUDA choice must report as a GPU backend"
);
assert!(decide(&p).label().contains("CUDA"));
}
/// LAW (inject-assert): NVIDIA wins over AMD when both probe ready (CUDA is
/// the proven path), and AMD is chosen when only ROCm is ready.
#[test]
fn cuda_precedes_rocm_but_rocm_chosen_alone() {
let both = Probe {
onnxruntime: true,
cuda_ready: true,
rocm_ready: true,
..Default::default()
};
assert_eq!(decide(&both), Backend::OrtCuda, "CUDA precedes ROCm");
let amd = Probe {
onnxruntime: true,
cuda_ready: false,
rocm_ready: true,
..Default::default()
};
assert_eq!(
decide(&amd),
Backend::OrtRocm,
"ROCm chosen when only AMD is ready"
);
assert!(decide(&amd).label().contains("ROCm"));
}
/// LAW (inject-assert): no GPU ⇒ **a CPU backend**, and specifically the FAST one
/// when ort is available.
///
/// CHANGED 2026-08-14. This asserted `TractCpu` with the rationale "we don't pull the
/// C++ runtime in to do CPU work tract already does" — an unmeasured claim, and the
/// measurement refutes it: ort's CPU EP is 4.18× faster on this very hardware at
/// cosine 1.0000 agreement. Reaching this arm means the dylib is already found and
/// armed, so choosing tract declines a 4× speedup that costs nothing and changes no
/// answer. tract stays the fallback for the case that actually needs it — no dylib
/// at all — which `no_onnxruntime_never_chooses_ort` below pins.
#[test]
fn no_gpu_with_ort_present_picks_the_fast_cpu_backend() {
let cpu_only = Probe {
onnxruntime: true,
cuda_ready: false,
rocm_ready: false,
..Default::default()
};
assert_eq!(decide(&cpu_only), Backend::OrtCpu);
assert!(!decide(&cpu_only).is_gpu(), "still not a GPU backend");
assert!(decide(&cpu_only).runs_a_graph(), "still a real transformer");
}
/// LAW (inject-assert): no onnxruntime ⇒ tract CPU regardless of GPU probe —
/// the gate that keeps ort (which would panic without its dylib) untouched.
#[test]
fn no_onnxruntime_never_chooses_ort() {
// Even with a GPU "ready", absent onnxruntime we must not pick ort.
let gpu_no_ort = Probe {
onnxruntime: false,
cuda_ready: true,
rocm_ready: true,
..Default::default()
};
assert_eq!(
decide(&gpu_no_ort),
Backend::TractCpu,
"no libonnxruntime.so ⇒ must stay on tract (ort would panic loading its dylib)"
);
let nothing = Probe::default();
assert_eq!(decide(¬hing), Backend::TractCpu);
}
/// The live probe must run on any box without panicking and agree with the
/// pure decision over its own result (the live path and decide() are
/// consistent). On a GPU-less CI box this asserts the CPU branch concretely.
#[test]
fn live_probe_is_self_consistent() {
let p = live_probe();
let live = chosen_backend();
assert_eq!(
live,
decide(&p),
"chosen_backend() must equal decide(live_probe())"
);
// A box with no onnxruntime must land on tract (the common CI case).
if !p.onnxruntime {
assert_eq!(live, Backend::TractCpu, "no onnxruntime ⇒ tract CPU");
}
}
/// LAW (inject-assert): when a *loadable* `libonnxruntime.so` is present
/// (pointed to by `$NORNIR_TEST_ORT_STUB`), `arm_onnxruntime()` must detect
/// it, set `ORT_DYLIB_PATH`, and the live selection must move OFF tract — to
/// `OrtCuda` on a CUDA box (this RTX-4090 dev box) or `OrtRocm` on an AMD one.
/// Proves the onnxruntime gate + GPU selection plumbing end-to-end without a
/// real GPU embed. Ignored by default; run with:
/// `NORNIR_TEST_ORT_STUB=/path/to/libonnxruntime.so \`
/// ` cargo test --lib -- --ignored selection_arms_ort_when_onnxruntime_present`
#[test]
#[ignore = "needs a loadable libonnxruntime.so via $NORNIR_TEST_ORT_STUB; run with --ignored"]
fn selection_arms_ort_when_onnxruntime_present() {
let stub = std::env::var("NORNIR_TEST_ORT_STUB")
.expect("set NORNIR_TEST_ORT_STUB to a loadable libonnxruntime.so");
// TODO: Audit that the environment access only happens in single-threaded code.
unsafe { std::env::set_var("ORT_DYLIB_PATH", &stub) };
// The onnxruntime probe must now report present.
assert!(
super::super::cuda::arm_onnxruntime(),
"loadable onnxruntime should arm"
);
assert!(
super::super::cuda::onnxruntime_dylib().is_some(),
"dylib must be discoverable"
);
let probe = live_probe();
assert!(
probe.onnxruntime,
"live probe must see the onnxruntime stub"
);
let backend = decide(&probe);
if probe.cuda_ready {
assert_eq!(
backend,
Backend::OrtCuda,
"CUDA box + onnxruntime ⇒ ort CUDA"
);
} else if probe.rocm_ready {
assert_eq!(
backend,
Backend::OrtRocm,
"AMD box + onnxruntime ⇒ ort ROCm"
);
} else {
// No GPU on this box: even with onnxruntime present we stay on tract.
assert_eq!(
backend,
Backend::TractCpu,
"no GPU ⇒ tract even with onnxruntime"
);
}
assert!(
backend.is_gpu() || !probe.cuda_ready,
"with CUDA ready the selected backend must be a GPU backend"
);
}
/// LAW (inject-assert): the selector's CPU-fallback path actually LOADS and
/// EMBEDS — not "didn't panic". Ignored by default (needs the model staged);
/// run with `cargo test --lib -- --ignored selector_load_backend_embeds`.
/// We force [`Backend::TractCpu`] (the universal fallback) so this runs on
/// any box, and assert a real 768-dim L2-normalized vector comes out and
/// that semantically-close code embeds nearer than unrelated prose.
#[test]
#[ignore = "needs the embedding model staged; loads + embeds for real; run with --ignored"]
fn selector_load_backend_embeds() {
let e = load_backend(Backend::TractCpu).expect("load tract CPU embedder");
let dim = super::super::embed_support::dim();
let v = e.embed(&["fn main() {}"]).expect("embed");
assert_eq!(v.len(), 1, "one vector per input text");
assert_eq!(v[0].len(), dim, "vector has the model's dimensionality");
let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 1e-3,
"vector is L2-normalized (norm {norm})"
);
// Real signal: two functions embed closer to each other than to prose.
let a = e
.embed(&["fn add(a: i32, b: i32) -> i32 { a + b }"])
.unwrap();
let b = e
.embed(&["pub fn sum(x: i32, y: i32) -> i32 { x + y }"])
.unwrap();
let c = e
.embed(&["the quick brown fox jumps over the lazy dog"])
.unwrap();
let dot = |x: &[f32], y: &[f32]| x.iter().zip(y).map(|(p, q)| p * q).sum::<f32>();
assert!(
dot(&a[0], &b[0]) > dot(&a[0], &c[0]),
"two fns ({}) should embed closer than fn-vs-prose ({})",
dot(&a[0], &b[0]),
dot(&a[0], &c[0])
);
}
// ── the 2026-08-14 additions: OrtCpu + StaticTable ───────────────────────────
/// THE FULL LADDER, injected: NVIDIA → worthwhile AMD → ort CPU → tract.
///
/// The owner's rule, 2026-08-14: "make sure NVIDIA is selected if we have nvidia, AMD
/// if we have an AMD GPU/APU, and if none or a crappy AMD gpu pick a CPU impl that
/// WORKS and is highly functional." Every rung asserted here, including the two the
/// old three-arm `decide` could not express.
#[test]
fn the_ladder_picks_nvidia_then_good_amd_then_ort_cpu_then_tract() {
let base = Probe {
onnxruntime: true,
..Default::default()
};
// NVIDIA wins outright, and is NOT subject to the VRAM test — a cuda_ready box
// is a discrete card by construction.
assert_eq!(
decide(&Probe {
cuda_ready: true,
rocm_ready: true,
amd_vram_bytes: Some(512 << 20),
model_bytes: Some(612 << 20),
..base
}),
Backend::OrtCuda,
"NVIDIA outranks AMD and skips the worth-it test"
);
// A REAL AMD card: 24 GB against a 612 MB model — take it.
assert_eq!(
decide(&Probe {
rocm_ready: true,
amd_vram_bytes: Some(24 << 30),
model_bytes: Some(612 << 20),
..base
}),
Backend::OrtRocm,
"a discrete AMD card that can hold the model is the right choice"
);
// THE `dwarf` CASE, with its real numbers: Radeon 680M, 512 MB dedicated VRAM,
// jina at 612 MB. It does not fit, so it would stream out of GTT = system RAM,
// which is the CPU's own bandwidth. Must fall through to ort CPU, NOT to ROCm.
assert_eq!(
decide(&Probe {
rocm_ready: true,
amd_vram_bytes: Some(536_870_912),
model_bytes: Some(641_517_466),
..base
}),
Backend::OrtCpu,
"an APU that cannot hold the model must NOT be preferred over the CPU"
);
// No GPU at all, but ort is present → the FAST cpu impl (measured 4.18x tract).
assert_eq!(
decide(&base),
Backend::OrtCpu,
"ort present and no usable GPU ⇒ MLAS, not tract"
);
// No onnxruntime anywhere → tract, which needs nothing on the box.
assert_eq!(
decide(&Probe {
onnxruntime: false,
cuda_ready: true,
rocm_ready: true,
..Default::default()
}),
Backend::TractCpu,
"without the dylib NO ort backend is reachable, GPU probes notwithstanding"
);
}
/// The worth-it test itself, including its deliberate trust of unknowns.
/// An unattended bulk embed runs on a GPU worth using, or where CPU was
/// asked for — and on nothing else. The `dwarf` numbers again: its APU is
/// "an AMD GPU", and it is not hardware for this.
#[test]
fn an_unattended_embed_needs_a_real_gpu_or_an_explicit_cpu() {
let ort = Probe { onnxruntime: true, ..Default::default() };
let apu = Probe { rocm_ready: true, amd_vram_bytes: Some(536_870_912), model_bytes: Some(641_517_466), ..ort };
let card = Probe { rocm_ready: true, amd_vram_bytes: Some(24 << 30), model_bytes: Some(641_517_466), ..ort };
let on = |p: &Probe, cpu_requested: bool| unattended_embed_skip(decide(p), cpu_requested, p);
// a GPU worth using: NVIDIA, a discrete AMD card
assert_eq!(on(&Probe { cuda_ready: true, ..ort }, false), None);
assert_eq!(on(&card, false), None);
// the on-chip AMD GPU is skipped, and the reason has its numbers
assert_eq!(decide(&apu), Backend::OrtCpu);
let why = on(&apu, false).expect("an APU that cannot hold the model is not embedding hardware");
assert!(why.starts_with("no embedding hardware here: the AMD GPU has 512 MB of dedicated VRAM and the model needs 917 MB"), "{why}");
assert!(why.contains("does not count") && why.ends_with("NORNIR_EMBED_BACKEND=ort-cpu asks for it by name"), "{why}");
// no GPU runtime at all, with and without onnxruntime
let why = on(&ort, false).unwrap();
assert!(why.contains("no NVIDIA CUDA runtime and no AMD ROCm runtime is ready") && why.contains(Backend::OrtCpu.label()), "{why}");
// … and the pin it names is the backend THIS box can load
let why = on(&Probe::default(), false).unwrap();
assert!(why.contains("no onnxruntime on this box") && why.contains(Backend::TractCpu.label()) && why.ends_with("NORNIR_EMBED_BACKEND=tract-cpu asks for it by name"), "{why}");
// CPU was ASKED for (a pin, a static-table model): it runs, on any box
assert_eq!(on(&apu, true), None);
assert_eq!(on(&Probe::default(), true), None);
assert_eq!(unattended_embed_skip(Backend::StaticTable, true, &apu), None);
// a GPU pin is taken at its word
assert_eq!(unattended_embed_skip(Backend::OrtCuda, false, &Probe::default()), None);
}
#[test]
fn vram_test_rejects_only_a_measured_shortfall() {
let p = |vram, model| Probe {
amd_vram_bytes: vram,
model_bytes: model,
..Default::default()
};
// Exactly at the 1.5x headroom → accepted; a hair under → rejected.
assert!(p(Some(900), Some(600)).amd_gpu_is_worth_it());
assert!(!p(Some(899), Some(600)).amd_gpu_is_worth_it());
// potion (28.8 MB) FITS even on the 680M's 512 MB — the rule is about the
// SELECTED model, not about embedding in general.
assert!(p(Some(536_870_912), Some(30_236_760)).amd_gpu_is_worth_it());
// jina (612 MB) does not.
assert!(!p(Some(536_870_912), Some(641_517_466)).amd_gpu_is_worth_it());
// Unknowns are TRUSTED: a ready ROCm install is deliberate, and a discrete card
// whose sysfs we failed to read must not be refused over a missing file.
assert!(p(None, Some(641_517_466)).amd_gpu_is_worth_it());
assert!(p(Some(536_870_912), None).amd_gpu_is_worth_it());
assert!(p(None, None).amd_gpu_is_worth_it());
}
/// PURE (LAW inject-assert): the MODEL's kind decides whether a graph backend is
/// even a candidate, and it is answered without touching hardware.
///
/// This is the guard that stops `$NORNIR_EMBED_MODEL=potion-base-8m` from being
/// routed to `TractCpu` and dying on a `model.onnx` that was never fetched for it.
#[test]
fn model_kind_outranks_the_hardware_probe() {
use super::super::embed_registry::ModelKind;
assert_eq!(
model_backend(ModelKind::StaticTable),
Some(Backend::StaticTable),
"a static table has no graph — no probe result can change that"
);
assert_eq!(
model_backend(ModelKind::Onnx),
None,
"a graph model defers to the hardware probe"
);
// The hardware ladder is asserted in full by
// `the_ladder_picks_nvidia_then_good_amd_then_ort_cpu_then_tract`; here we only
// pin that a graph model still DEFERS to it rather than forcing a backend.
let p = Probe {
onnxruntime: true,
cuda_ready: false,
rocm_ready: false,
..Default::default()
};
assert_eq!(
model_backend(ModelKind::Onnx),
None,
"a graph model must not preempt the probe"
);
assert!(
!decide(&p).is_gpu() && decide(&p).runs_a_graph(),
"and the probe then picks a real-graph CPU backend"
);
}
/// The `$NORNIR_EMBED_BACKEND` tokens round-trip, and an unknown token is `None`
/// (a "no request"), never a default. A typo that silently pinned a backend would
/// make every measurement taken afterwards a lie about which code ran.
#[test]
fn backend_env_tokens_round_trip_and_reject_junk() {
for b in [
Backend::OrtCuda,
Backend::OrtRocm,
Backend::OrtCpu,
Backend::TractCpu,
Backend::StaticTable,
] {
assert_eq!(
Backend::from_env_token(b.env_token()),
Some(b),
"{} must round-trip through its env token",
b.label()
);
// Case and whitespace are an operator's, not a contract's.
assert_eq!(
Backend::from_env_token(&format!(" {} ", b.env_token().to_uppercase())),
Some(b)
);
}
for junk in ["", "ortcpu", "ort_cpu", "gpu", "nvidia", "potion"] {
assert_eq!(
Backend::from_env_token(junk),
None,
"{junk:?} must not resolve to a backend"
);
}
}
/// `is_gpu` and `runs_a_graph` classify the five variants correctly — the two
/// questions the diagnostics ask, and the reason `StaticTable` needed the second
/// one: it is not a GPU backend AND not a slow-CPU backend, so `!is_gpu()` alone
/// would have made `fallback.rs` accuse it of being a silent GPU degradation.
#[test]
fn variant_classification_separates_gpu_from_graph() {
assert!(Backend::OrtCuda.is_gpu() && Backend::OrtRocm.is_gpu());
assert!(!Backend::OrtCpu.is_gpu());
assert!(!Backend::TractCpu.is_gpu());
assert!(!Backend::StaticTable.is_gpu());
assert!(Backend::OrtCuda.runs_a_graph());
assert!(Backend::OrtRocm.runs_a_graph());
assert!(Backend::OrtCpu.runs_a_graph());
assert!(Backend::TractCpu.runs_a_graph());
assert!(
!Backend::StaticTable.runs_a_graph(),
"the static table is a lookup — calling it an inference backend is the \
misreading this method exists to prevent"
);
// Every variant must have a distinct label; two backends sharing a name would
// make `vector doctor` and the fallback diagnostics unfalsifiable.
let labels = [
Backend::OrtCuda.label(),
Backend::OrtRocm.label(),
Backend::OrtCpu.label(),
Backend::TractCpu.label(),
Backend::StaticTable.label(),
];
let mut uniq = labels.to_vec();
uniq.sort_unstable();
uniq.dedup();
assert_eq!(
uniq.len(),
labels.len(),
"labels must be distinct: {labels:?}"
);
}
}