src/vector/embed_ort.rs
//! ort (ONNX Runtime) embedder (GPU/CPU) — `jina-embeddings-v2-base-code`.
//!
//! Implements [`super::store::Embedder`] by running the model's ONNX export via
//! the `ort` crate (Microsoft ONNX Runtime). Requests a **GPU execution
//! provider** and falls back to CPU when none is available, so the same binary
//! works on a GPU box or a plain server. Two GPU vendors are supported:
//!
//! - **NVIDIA / CUDA** (always wired when `embed-ort` is on): the `CUDA` EP,
//! with [`super::cuda`] preloading the runtime libs.
//! - **AMD / ROCm** (extra feature `embed-ort-rocm`): the `MIGraphX` then
//! `ROCm` EPs, with [`super::rocm`] preloading the runtime libs and a real
//! runtime probe ([`super::rocm::available`]) so the AMD EP is only requested
//! when ROCm is actually installed. A non-AMD build never pulls this in; a
//! build with it but no AMD GPU degrades to the CUDA-then-CPU path unchanged.
//!
//! EP registration is **best-effort** (`fail_silently`): a missing GPU runtime
//! quietly drops to the next provider, ending at CPU — it never errors out of
//! `load`.
//!
//! This is the GPU-capable alternative to the pure-Rust tract backend
//! ([`super::embed`]). Both run the same model and share
//! [`super::embed_support`], so a vector produced here is identical to one from
//! tract — embeddings made by either backend share the same `model_profile`
//! dedup key in the warehouse. The model itself is selectable via the registry
//! ([`super::embed_registry`]); this backend runs whichever ONNX export
//! `build.rs` fetched (jina-v2-base-code, 768-dim, by default).
//!
//! ⚠ Heaviest no-C exception: `ort` links the onnxruntime C++ library (and,
//! with CUDA present, the CUDA runtime). A nornir build without `embed-ort`
//! pulls none of this. Accepted by the owner, 2026-06-04.
//!
//! onnxruntime parallelizes a single `run` internally (intra-op threads). The
//! embedder builds a real `[B, L]` batched tensor per length bucket and issues
//! ONE `session.run` per batch (not one per text), so the GPU sees big batches
//! instead of N lock-serialized `[1, seq]` singletons — the batched-submit fix.
//! GPU work serializes on the device anyway; batching amortizes the launch.
//!
//! Cargo feature: `embed-ort`.
use std::sync::Mutex;
use anyhow::{Context, Result};
use ort::session::Session;
use ort::value::Tensor;
use tokenizers::Tokenizer;
use super::embed_support as es;
use super::store::{Embedder, ModelProfile};
/// A loaded jina-v2-base-code ONNX model + tokenizer (ort / GPU-or-CPU).
///
/// `Session::run` takes `&mut self`, so the session is behind a `Mutex`; calls
/// serialize, but each `run` already uses many cores (CPU EP) or the GPU.
pub struct OrtEmbedder {
session: Mutex<Session>,
tokenizer: Tokenizer,
/// Whether a GPU EP was *requested* (it falls back to CPU silently if the
/// GPU runtime is absent).
gpu_requested: bool,
/// Which GPU vendor's EP was requested, for diagnostics: `"cuda"`,
/// `"rocm"`, or `"cpu"` (no GPU requested / available).
gpu_vendor: &'static str,
}
impl OrtEmbedder {
/// Load the model, requesting the best available GPU execution provider
/// (AMD ROCm when `embed-ort-rocm` is on and ROCm is present, else NVIDIA
/// CUDA), with a clean CPU fallback.
pub fn load() -> Result<Self> {
Self::load_inner(true, false)
}
/// Load DEMANDING the GPU execution provider actually register: registration
/// uses ort's `error_on_failure` instead of `fail_silently`, so a CUDA/ROCm EP
/// that cannot be registered returns an error here rather than quietly serving
/// every `run` on the CPU EP inside a session that still calls itself a GPU one.
///
/// This is for the caller that has ALREADY PROVEN the GPU is usable — the
/// runtime selector, which only reaches [`super::select::Backend::OrtCuda`]
/// after driver + cudart + cuDNN all dlopen'ed. Once readiness is established,
/// a registration failure is a real fault (a CUDA-major mismatch between
/// onnxruntime and the installed runtime is the usual one), not the ordinary
/// "this box has no GPU" case that `fail_silently` exists to absorb — and a
/// real fault must be named, not absorbed. Absorbing it is how a 4090 box ends
/// up embedding ~60× slower with nothing in the log saying so.
///
/// Callers that have NOT proven readiness should use [`load`], whose
/// best-effort registration is the correct behaviour for an unknown box.
pub fn load_gpu_strict() -> Result<Self> {
Self::load_inner(true, true)
}
/// Load the model on CPU only (do not request any GPU EP).
pub fn load_cpu() -> Result<Self> {
Self::load_inner(false, false)
}
/// Load the model. When `prefer_gpu`, register a GPU EP first; ort's
/// registration is best-effort (`fail_silently`), so a machine without that
/// GPU runtime quietly uses the CPU EP instead.
///
/// EP precedence: AMD ROCm (only when `embed-ort-rocm` is compiled **and**
/// [`super::rocm::available`] runtime-probes true) → NVIDIA CUDA → CPU.
/// MIGraphX is registered ahead of plain ROCm because it graph-compiles the
/// whole model; ort tries them in order and falls through on registration
/// failure.
pub fn load_with(prefer_gpu: bool) -> Result<Self> {
Self::load_inner(prefer_gpu, false)
}
/// Shared body of [`load`], [`load_gpu_strict`] and [`load_cpu`]. `strict`
/// selects ort's `error_on_failure` over `fail_silently` for the GPU EP —
/// see [`load_gpu_strict`] for when each is right.
fn load_inner(prefer_gpu: bool, strict: bool) -> Result<Self> {
let dir = es::model_dir();
let tokenizer = Tokenizer::from_file(dir.join("tokenizer.json"))
.map_err(|e| anyhow::anyhow!("load tokenizer.json: {e}"))?;
let mut builder = Session::builder().context("ort session builder")?;
// `$NORNIR_EMBED_CPU_THREADS` caps the CPU EP's threads: a GPU worker's
// CPU fallback (texts its device cannot hold) must not take a shared
// box's every core.
if !prefer_gpu {
if let Some(n) = std::env::var("NORNIR_EMBED_CPU_THREADS").ok().and_then(|v| v.trim().parse::<usize>().ok()).filter(|n| *n > 0) {
builder = builder.with_intra_threads(n).map_err(|e| anyhow::anyhow!("ort intra threads: {e}"))?;
}
}
let mut gpu_vendor = "cpu";
// One writer for the strict/best-effort choice, so the CUDA and ROCm arms
// below cannot drift apart on it.
let arm = |d: ort::ep::ExecutionProviderDispatch| {
if strict {
d.error_on_failure()
} else {
d.fail_silently()
}
};
if prefer_gpu {
// AMD ROCm path — only when the feature is built and the runtime
// probe says ROCm is actually installed. On any other box this
// branch is skipped entirely, leaving the CUDA path below untouched.
#[cfg(feature = "embed-ort-rocm")]
{
if super::rocm::available() {
// Make the ROCm libs resolvable before the provider
// dlopen's them, mirroring the CUDA preload.
let _ = super::rocm::ensure();
builder = builder
.with_execution_providers([
// MIGraphX compiles the graph; ROCm is the generic
// fallback. Best-effort → CPU on any miss, unless the
// caller already proved the GPU usable (`strict`).
arm(ort::ep::MIGraphX::default().build()),
arm(ort::ep::ROCm::default().build()),
])
.map_err(|e| anyhow::anyhow!("register ROCm/MIGraphX EP: {e}"))?;
gpu_vendor = "rocm";
}
}
// NVIDIA CUDA path — the original behaviour, used whenever the ROCm
// branch didn't claim the GPU (no AMD feature, or no ROCm present).
if gpu_vendor != "rocm" {
// Make CUDA libs resolvable before the provider dlopen's them,
// so the GPU EP works without a manual LD_LIBRARY_PATH.
let _ = super::cuda::ensure();
// the arena grows by what a run asks, not by the next power of
// two: a long text's attention bias is GBs, and doubling past it
// left a 16 GB card unable to place the next one (MEASURED on
// loki's RTX 5080, 2026-10-17: a 6.2 GB request refused with
// 7.8 GB in use)
builder = builder
.with_execution_providers([arm(ort::ep::CUDA::default().with_arena_extend_strategy(ort::ep::ArenaExtendStrategy::SameAsRequested).build())])
.map_err(|e| anyhow::anyhow!("register CUDA execution provider: {e}"))?;
gpu_vendor = "cuda";
}
}
let session = builder
.commit_from_file(dir.join("model.onnx"))
.context("commit onnx session")?;
Ok(Self {
session: Mutex::new(session),
tokenizer,
gpu_requested: prefer_gpu,
gpu_vendor,
})
}
/// True if a GPU EP was requested at load time. (ort doesn't expose which EP
/// actually serviced a run; GPU EPs fall back to CPU silently.)
pub fn gpu_requested(&self) -> bool {
self.gpu_requested
}
/// Which GPU vendor's EP was requested: `"cuda"`, `"rocm"`, or `"cpu"`.
pub fn gpu_vendor(&self) -> &'static str {
self.gpu_vendor
}
/// Whether the AMD ROCm runtime preload found a usable set (HIP + MIOpen).
/// Always `false` unless the `embed-ort-rocm` feature is compiled.
pub fn rocm_ready(&self) -> bool {
#[cfg(feature = "embed-ort-rocm")]
{
return super::rocm::available();
}
#[cfg(not(feature = "embed-ort-rocm"))]
{
false
}
}
/// Whether the runtime CUDA-lib preload found a usable set (cudart +
/// cuDNN). When false, the CUDA EP almost certainly fell back to CPU.
pub fn cuda_ready(&self) -> bool {
let p = super::cuda::ensure();
p.cudnn && p.loaded.iter().any(|s| s.starts_with("libcudart"))
}
/// Whether the active EP requires a CONSTANT input shape (MIGraphX/ROCm =
/// ONNX Runtime's AMD GPU path). Dynamic EPs (CPU / tract / CUDA) accept any
/// `[B, L]` per run; static EPs graph-compile per shape, so a batch must pad
/// to one of the fixed [`es::SHAPE_BUCKETS`] lengths.
fn static_shape(&self) -> bool {
self.gpu_vendor == "rocm"
}
/// Tokenize one text into truncated `i64` ids (length ∈ `[1, max_tokens]`) —
/// the same ids `embed_one` would feed, so a batched forward built from these
/// is equivalent to running each text one at a time.
fn tokenize(&self, text: &str) -> Result<Vec<i64>> {
let enc = self
.tokenizer
.encode(text, true)
.map_err(|e| anyhow::anyhow!("tokenize: {e}"))?;
Ok(es::prepare_ids(enc.get_ids(), es::max_tokens()))
}
/// Run ONE `session.run` over a batch of pre-tokenized texts (`toks[i]` for
/// each `i` in `batch`) at shape `[B, pad_len]`. `pad_len` is the batch's
/// longest row (dynamic EP) or the fixed bucket length `L` (static EP, passed
/// in `fixed_len`). Padding tokens carry mask `0`, so mask-aware pooling makes
/// each row's vector equal its one-at-a-time `embed_one` result. Returns
/// `(original_index, embedding)` per row. **One lock, one GPU pass per batch**
/// — the fix for the batch-1 Mutex-serialized submit (finding #1).
fn forward_batch(
&self,
batch: &[usize],
toks: &[Vec<i64>],
fixed_len: Option<usize>,
) -> Result<Vec<(usize, Vec<f32>)>> {
let b = batch.len();
debug_assert!(b > 0, "empty batch");
let dim = es::dim();
// pad_len: the bucket L on the static path, else the batch's own max row.
let pad_len = fixed_len
.unwrap_or_else(|| batch.iter().map(|&i| toks[i].len()).max().unwrap_or(1))
.max(1);
// Build the [B, pad_len] `input_ids` + `attention_mask` from front-packed
// real tokens (rest zero-padded, mask 0). One flat alloc each, no per-row
// Vec churn.
let mut ids = vec![0i64; b * pad_len];
let mut mask = vec![0i64; b * pad_len];
for (row, &i) in batch.iter().enumerate() {
let base = row * pad_len;
let n_real = toks[i].len().min(pad_len);
for j in 0..n_real {
ids[base + j] = toks[i][j];
mask[base + j] = 1;
}
}
let shape = [b as i64, pad_len as i64];
let id_tensor = Tensor::from_array((shape, ids)).context("build input_ids tensor")?;
let mask_tensor =
Tensor::from_array((shape, mask.clone())).context("build attention_mask tensor")?;
// ONE lock for the whole batch (not per text). `Session::run` needs
// `&mut self` and its `outputs` borrow the session, so the guard is held
// through extraction + pooling — the GPU serializes on the device anyway,
// so this is strictly fewer acquisitions than the old per-text lock storm.
// On the CUDA EP the arena gives back what a run left unused when the
// run ends. Without it the arena only grows: one long text's attention
// (GBs, quadratic in tokens) stayed held, the next long text found no
// room on a 16 GB card, and every later long text failed the device
// (MEASURED on loki's RTX 5080, 2026-10-17: 15.7 GB held at 0 % load
// while 31 texts of 5–11 kB were refused).
let mut opts = ort::session::RunOptions::new().context("ort run options")?;
if self.gpu_vendor == "cuda" {
opts.set("memory.enable_memory_arena_shrinkage", "gpu:0").context("arena shrinkage run option")?;
}
let mut session = self.session.lock().unwrap();
let outputs = session
.run_with_options(
ort::inputs![
"input_ids" => id_tensor,
"attention_mask" => mask_tensor,
],
&opts,
)
.context("onnx batched forward")?;
let (out_shape, data) = outputs[0]
.try_extract_tensor::<f32>()
.context("extract hidden state")?;
let dims = out_shape.as_ref();
anyhow::ensure!(
dims.len() == 3 && dims[2] as usize == dim,
"unexpected output shape {dims:?} (expected last dim {dim})"
);
let out_rows = dims[1] as usize;
if std::env::var_os("NORNIR_EMBED_DEBUG").is_some() {
eprintln!(
"embed: static={} B={b} pad_len={pad_len} out_rows={out_rows} out_dims={dims:?}",
self.static_shape()
);
}
// CORRECTNESS GUARD. The EP must return exactly the `pad_len` token rows we
// fed it: dynamic EPs always do; for the padded MIGraphX path a mismatch
// means it reused a DIFFERENT bucket's compiled program (its model cache is
// keyed without the runtime shape) — pooling that would mix the wrong
// tokens. Refuse rather than poison the warehouse with a silent-garbage
// vector.
anyhow::ensure!(
out_rows == pad_len && dims[0] as usize == b,
"embedder returned shape {dims:?} for a [{b}, {pad_len}] input — the \
active execution provider (MIGraphX/ROCm) ignored the requested shape \
(shape-locked to another bucket) and would produce WRONG embeddings. \
Use the CPU embedder (unset ORT_DYLIB_PATH / build without \
embed-ort-rocm) for correct vectors."
);
// Slice each row's `[pad_len, dim]` block and mask-pool its real tokens
// (padding rows contribute nothing), so a batched/padded embed equals the
// unpadded one-at-a-time result.
let mut out = Vec::with_capacity(b);
let stride = pad_len * dim;
for (row, &i) in batch.iter().enumerate() {
let start = row * stride;
let row_hidden = &data[start..start + stride];
let row_mask = &mask[row * pad_len..(row + 1) * pad_len];
out.push((i, es::pool_and_normalize_masked(row_hidden, row_mask, dim)));
}
Ok(out)
}
}
impl Embedder for OrtEmbedder {
fn profile(&self) -> ModelProfile {
es::profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let n = texts.len();
if n == 0 {
return Ok(Vec::new());
}
// Tokenize once (cheap vs the forward); lengths drive the bucketing.
let toks: Vec<Vec<i64>> = texts
.iter()
.map(|t| self.tokenize(t))
.collect::<Result<_>>()?;
let lens: Vec<usize> = toks.iter().map(|t| t.len()).collect();
// Same plan door as the tract backend (`vector::memory`): the legacy
// caps, shrunk under `--max-memory` when a ceiling is declared. ort runs
// ONE submit at a time (the device is the fan-out), so `workers` = 1 here
// and only the per-batch token cap can move.
let limits = super::memory::EmbedLimits::from_env(1, es::MAX_BATCH_ROWS, es::MAX_BATCH_TOKENS);
limits.announce_once();
let (max_rows, max_tokens) = (limits.max_rows, limits.max_tokens);
let max = es::max_tokens();
// Plan the `[B, L]` submits. Dynamic EP: length-bucket then pack (each
// batch pads to its own max row). Static EP (MIGraphX): first split by the
// fixed SHAPE bucket `L` so every row in a batch shares the compiled shape,
// then pack each bucket under the same row/token caps.
let mut plan: Vec<(Vec<usize>, Option<usize>)> = Vec::new();
if self.static_shape() {
// Group indices by their fixed bucket length.
let mut by_bucket: std::collections::BTreeMap<usize, Vec<usize>> =
std::collections::BTreeMap::new();
for i in 0..n {
let l = es::bucket_len(lens[i].min(max), max);
by_bucket.entry(l).or_default().push(i);
}
for (l, idxs) in by_bucket {
// Within a bucket all rows pad to `l`, so weight is rows*l — cap by
// both the row count and the token budget.
let rows_cap = max_rows.min((max_tokens / l.max(1)).max(1));
for chunk in idxs.chunks(rows_cap) {
plan.push((chunk.to_vec(), Some(l)));
}
}
} else {
for b in es::plan_batches(&lens, max_rows, max_tokens) {
plan.push((b, None));
}
}
// Run each planned batch as ONE forward. The session is behind a Mutex, so
// batches serialize on it (and on the single GPU) — but each acquisition
// now processes a whole `[B, L]` batch instead of one text, which is the
// point. Scatter each batch's rows back to their original positions.
let mut out: Vec<Vec<f32>> = vec![Vec::new(); n];
for (batch, fixed) in &plan {
for (i, v) in self.forward_batch(batch, &toks, *fixed)? {
out[i] = v;
}
}
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Loads the real ONNX model via ort (CPU EP — CUDA optional). Ignored by
/// default; run with `cargo test --features embed-ort -- --ignored embed`.
#[test]
#[ignore = "loads the real ONNX model via ort (CPU EP); run with --features embed-ort -- --ignored"]
fn loads_and_embeds_cpu() {
let e = OrtEmbedder::load_cpu().expect("load model");
let p = e.profile();
assert_eq!(p.dim, es::dim());
assert_eq!(p.model_name, es::model_name());
let v = e.embed(&["fn main() {}"]).unwrap();
assert_eq!(v.len(), 1);
assert_eq!(v[0].len(), es::dim());
let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
assert!((norm - 1.0).abs() < 1e-3, "norm {norm}");
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 ({}) closer than fn-vs-prose ({})",
dot(&a[0], &b[0]),
dot(&a[0], &c[0])
);
}
}