nordisk/nornir

Runtime embedder selection — GPU/CPU/AMD in ONE binary (plan #9)

Internal design + guide for plan item #9: a plain cargo install nornir yields ONE binary that carries both embedder backends and picks the right one at runtime — NVIDIA GPU (CUDA) or AMD GPU (ROCm) when present, pure-Rust CPU otherwise — with no build failure and no panic on a GPU-less / library-less box. This file is a durable .nornir/ design note (not a rendered artifact).

Source: src/vector/select.rs (the selector), src/vector/cuda.rs (NVIDIA + onnxruntime discovery), src/vector/rocm.rs (AMD probe), src/vector/embed.rs (tract CPU), src/vector/embed_ort.rs (ort GPU/CPU). Companion: vector.md (the index), vector-rocm.md.

What changed vs the pre-#9 state

Before #9 the default features shipped only embed-tract (pure-Rust CPU). embed-ort (GPU) was opt-in, so the default binary never touched the GPU even on an RTX 4090. The runtime CUDA/ROCm probes (cuda.rs / rocm.rs) existed but were dead code in the default build (#[cfg(feature = "embed-ort")]).

#9 makes both backends part of default and adds a runtime selector that is the single place deciding which embedder to load on the running machine.

The selection (probe order → backend)

src/vector/select.rs::decide(&Probe) -> Backend is a pure function over a struct of booleans:

  probe onnxruntime present?  ── no ──▶ Backend::TractCpu  (CPU, pure Rust)
         │ yes
  probe NVIDIA CUDA ready?    ── yes ─▶ Backend::OrtCuda   (NVIDIA GPU)
         │ no
  probe AMD  ROCm  ready?     ── yes ─▶ Backend::OrtRocm   (AMD GPU)
         │ no
         └──────────────────────────▶ Backend::TractCpu  (CPU, pure Rust)

• onnxruntime is the gate for any ort path. NVIDIA (CUDA) precedes AMD (ROCm) because CUDA is the proven path. With no usable GPU we stay on tract CPU even if onnxruntime happens to be installed — there is no reason to pull the C++ runtime in to do CPU work tract already does well (pure-Rust, no FFI).

• live_probe() fills the Probe from the real probes:

  • onnxruntime ← cuda::arm_onnxruntime() — finds libonnxruntime.so and, if found, sets ORT_DYLIB_PATH to it so ort loads exactly that file.
  • cuda_ready ← cuda::preflight().0 — NVIDIA driver + cudart + cuDNN.
  • rocm_ready ← rocm::available() (only when embed-ort-rocm is compiled).

• load_embedder() (src/vector/mod.rs) routes to select::load() in the default both-backend build, and to the single compiled backend otherwise.

The build strategy: ort load-dynamic (why it can't break a CPU-only build)

The hard constraint: embed-ort must NOT break cargo build/cargo install on a box with no onnxruntime and no GPU. A normal embed-ort build links the onnxruntime C++ library (via ort's default download-binaries/static link), so a plain box would fail to build/link — that violates "works out of the box."

The fix is ort's load-dynamic feature (ort/load-dynamic → ort-sys/disable-linking, verified against ort 2.0.0-rc.12 in this repo):

• The binary links NO onnxruntime at build time. ort-sys' build script returns early under disable-linking, downloading and linking nothing. So cargo build --features embed-ort succeeds on any box, GPU or not, library or not. Proven here: the default build (now including embed-ort) compiles clean.
• At runtime, ort dlopens libonnxruntime.so from $ORT_DYLIB_PATH (or the bare soname) on the first call into any ort API. If the dylib is absent, ort PANICS (ort::api() is a get_or_init/expect).

That panic risk is exactly why the selector is probe-first: it never constructs an ort object unless arm_onnxruntime() confirmed a real libonnxruntime.so is loadable (the probe dlopens it and checks the OrtGetApiBase export, then drops the handle). No onnxruntime ⇒ decide() returns TractCpu and the ort module is never entered → no panic, no link.

Cargo.toml default change

default = [..., "embed-tract", "embed-ort", ...]   # both backends, was tract-only
embed-ort = ["vector", "dep:ort", "ort/load-dynamic", "dep:tokenizers", "dep:libloading"]
ort = { version = "2.0.0-rc.12", default-features = false,
        features = ["std", "ndarray", "cuda", "api-24"], optional = true }

default-features = false drops ort's download-binaries/tls-* (not needed under dynamic-load); load-dynamic brings ort-sys/disable-linking. The cuda feature only enables the CUDA EP bindings (no link under disable-linking); embed-ort-rocm adds ort/migraphx + ort/rocm for the AMD EP.

HONEST build-prerequisite note

• Default build needs NOTHING extra to compile or to run on CPU. No onnxruntime, no GPU, no CUDA libs — it builds clean and embeds via tract CPU.

• The GPU path additionally needs, at RUNTIME:

  1. libonnxruntime.so resolvable — via $ORT_DYLIB_PATH, dropped into /opt/nornir/cuda (the CUDA discovery dir), or on the system loader path. It must be ≥ the ORT version ort rc.12 targets (1.24.x / api-24); an older lib (e.g. 1.22) is rejected by ort's version check, and the selector then degrades to tract CPU.
  2. The NVIDIA driver + CUDA-12 runtime + cuDNN-9 (or AMD driver + ROCm/MIOpen), discovered by cuda.rs / rocm.rs (e.g. staged in /opt/nornir/cuda).

• When any of that is missing the selector silently uses tract CPU — never a build failure, never a panic. nornir vector doctor reports exactly which probe is missing and what to install.

So #9's "GPU by default" is honest: the binary is capable of GPU out of the box, and uses it the moment a real onnxruntime + GPU are present, with zero build prerequisites for the CPU case.

nornir vector doctor

In the default build, doctor prints the CUDA preflight, then the runtime probe booleans (onnxruntime / cuda_ready / rocm_ready), the selected backend, and a real test embed through it (dims + ms + backend label). It is the operator's single readout of the runtime choice.

Batched CPU embedding (2026-07-02) — ~2–4×, bit-identical

The tract CPU path (src/vector/embed.rs) used to run one forward per text. It now batches: tokenize all texts, length-bucket them (sort by token count so similar lengths pad together, minimizing wasted padding), pack up to MAX_BATCH_ROWS=16 rows / MAX_BATCH_TOKENS=8192 per batch (env NORNIR_EMBED_BATCH_ROWS / NORNIR_EMBED_BATCH_TOKENS), run one (B, pad_len) forward per batch, mask-aware mean-pool + normalize each row, distribute batches across cores via gatling_map_balanced, reassemble in original order.

• Feasible with no re-export: the jina ONNX has a dynamic batch dim (['batch','sequence']), so tract resolves (B, max_len) per run().
• Bit-identical: test batched_equals_one_at_a_time on the real model → max_abs_diff = 0e0 (every f32 equal by to_bits()). The attention mask makes real-token states independent of padding (masked scores underflow to exactly 0.0), and batching leaves the matmul contraction dim (model width) unchanged, so no float op is reordered. A ≤1e-5 hard cap guards regressions.
• This is EMBEDDING speed, not search speed. It is the transformer forward pass amortized over rows; it is unrelated to the int8/VNNI search kernels in vector-cpu-simd.md. Realistic ~2–4× on short similar-length code chunks.

Tests (LAW: inject-assert)

In src/vector/select.rs (#[cfg(test)]):

• picks_ort_cuda_when_gpu_ready — inject {onnxruntime, cuda_ready} ⇒ assert OrtCuda.
• cuda_precedes_rocm_but_rocm_chosen_alone — both ready ⇒ CUDA; only AMD ⇒ ROCm.
• falls_back_to_tract_without_gpu — onnxruntime present but no GPU ⇒ tract CPU.
• no_onnxruntime_never_chooses_ort — no onnxruntime ⇒ tract even with GPU flags (the panic-avoidance gate).
• live_probe_is_self_consistent — the live probe runs anywhere; on a no-onnxruntime box it lands on tract.
• selector_load_backend_embeds (#[ignore], needs the model) — forces TractCpu, asserts a real 768-dim L2-normalized vector and that two functions embed closer than fn-vs-prose. The CPU fallback genuinely embeds.

Plus the existing rocm::probe_degrades_to_cpu_without_rocm and the ort/tract #[ignore] load+embed tests.