src/vector/memory.rs
//! **The memory ceiling for `nornir vector index`** — the lever that used to be
//! `taskset`.
//!
//! ## The allocation shape (INVENTORY-semantic-twins-2026-09-10 §2, §5 #5)
//!
//! `vector index monetize` was OOM-killed twice at **7 528 584 kB anon-rss** on a
//! box carrying three other lanes' builds. The shape is NOT "all embeddings in
//! memory at once" — the output vectors are `n × dim × 4` bytes (25 700 gunnar
//! chunks × 768 × 4 = 79 MB) and the tokens about the same. The shape is the
//! **fan-out**: [`super::embed::JinaEmbedder::embed`] plans length-bucketed
//! batches of up to `NORNIR_EMBED_BATCH_TOKENS` (8192) padded tokens and runs
//! `available_parallelism()` of them CONCURRENTLY through tract, and every
//! in-flight forward owns its own activations — hidden states `[B, L, 768]`
//! plus the attention scores `[B, heads, L, L]` per layer. Peak RSS is therefore
//!
//! ```text
//! peak ≈ BASELINE (weights + graph) + workers × tokens_per_batch × BYTES_PER_TOKEN
//! ```
//!
//! and the only pre-existing knobs were the per-batch caps and the core count
//! (`taskset -c 0-5`, honoured by `available_parallelism`, took peak from 5.6 GB
//! to under 1.6 GB together with `NORNIR_EMBED_BATCH_ROWS=8`).
//!
//! ## The two constants are MEASURED, from those two points
//!
//! * 16 workers × ≤8192 padded tokens → 5.6 GB peak (the run that got killed at
//! 7.5 GB under memory pressure);
//! * 6 workers × ≤4096 effective padded tokens (rows capped at 8 halves the
//! typical batch) → <1.6 GB.
//!
//! With a ~700 MiB resident baseline (612 MiB of `model.onnx` weights + the
//! optimised graph + tokenizer) both points give **≈37 KiB per padded token**;
//! [`EMBED_BYTES_PER_PADDED_TOKEN`] rounds that UP to 40 KiB so the estimate
//! errs toward refusing a worker rather than admitting an OOM. Attention memory
//! is `L²` per row, so a batch of few long rows costs more per token than many
//! short ones; within the default 8192-token cap the linear model held to
//! within the two measurements. It is an estimate — the plan SAYS its estimated
//! peak so the operator can compare it with what `ps` reports and correct the
//! constant here. A real-warehouse re-measurement is still owed (oden offline).
//!
//! ## The lever
//!
//! `--max-memory 2G` (or `NORNIR_EMBED_MAX_MEMORY_MB=2048`) → [`EmbedLimits::bounded`]
//! spends the budget on WORKERS first (that is where the throughput is), and
//! only when not even one full-width batch fits does it shrink the per-batch
//! token cap, never below [`MIN_BATCH_TOKENS`]. When even the floor plan is
//! estimated over budget the plan still runs (a refusal would leave the operator
//! with `taskset` again) but [`EmbedLimits::over_budget`] is set and the
//! announcement says so. Zero silent fallbacks.
//!
//! The row-side ceiling is separate and already existed for the batched
//! multi-repo path: `store::EMBED_SLICE_CHUNKS` (2048 rows per embed call, each
//! slice appended durably). Gap 3 made the per-repo `index_repo` path stream
//! through the same slice door (`store::index_snapshot_streamed`), so the
//! vectors-in-flight are bounded by the slice on BOTH paths and a kill mid-run
//! resumes from what already landed.
//!
//! Compiled with the `vector` feature (no embedder needed — the CLI flag parser
//! and the budget are backend-independent); the backends read the plan.
use std::sync::atomic::{AtomicU64, Ordering};
/// Bytes of in-flight tract activations per PADDED token of one batch — MEASURED
/// (see the module note), rounded UP from ≈37 KiB.
pub const EMBED_BYTES_PER_PADDED_TOKEN: u64 = 40 * 1024;
/// Resident baseline that exists before the first batch runs: `model.onnx`
/// weights (612 MiB for jina-v2-base-code) + the optimised graph + tokenizer.
pub const EMBED_BASELINE_BYTES: u64 = 700 * 1024 * 1024;
/// The smallest padded-token cap the planner shrinks to — one 256-token row.
/// Below this the model's own per-row minimum dominates and the estimate is
/// meaningless, so the plan stops here and flags `over_budget` instead.
pub const MIN_BATCH_TOKENS: usize = 256;
/// Environment spelling of the budget. MiB when bare (`2048`); suffixes
/// accepted as in [`parse_memory_size`].
pub const MEMORY_BUDGET_ENV: &str = "NORNIR_EMBED_MAX_MEMORY_MB";
/// Process-wide budget set by the CLI (`--max-memory`); `0` = unset, fall back
/// to the env var. An atomic rather than `std::env::set_var` because the
/// embedders read it from worker threads and `set_var` is `unsafe` in edition
/// 2024 for exactly that reason.
static MEMORY_BUDGET: AtomicU64 = AtomicU64::new(0);
/// Declare the embed memory ceiling for this process (bytes). `0` clears it.
pub fn set_memory_budget_bytes(bytes: u64) {
MEMORY_BUDGET.store(bytes, Ordering::SeqCst);
}
/// The active ceiling: the CLI-set value first, else [`MEMORY_BUDGET_ENV`],
/// else `None` (= unbounded, the pre-gap-3 behaviour).
pub fn memory_budget_bytes() -> Option<u64> {
let set = MEMORY_BUDGET.load(Ordering::SeqCst);
if set > 0 {
return Some(set);
}
std::env::var(MEMORY_BUDGET_ENV)
.ok()
.and_then(|v| parse_memory_size(&v).ok())
.filter(|&b| b > 0)
}
/// Parse a memory size: `2G`, `1536M`, `768MiB`, `4GB`, `512K`, or a bare
/// number which is **MiB** (the env var's name says `_MB`). Binary units
/// throughout (`G` = GiB) — RSS is what the kernel kills on, and it counts in
/// KiB. Whitespace-tolerant, case-insensitive.
pub fn parse_memory_size(s: &str) -> Result<u64, String> {
let s = s.trim();
if s.is_empty() {
return Err("empty memory size".into());
}
let split = s
.find(|c: char| !c.is_ascii_digit() && c != '.')
.unwrap_or(s.len());
let (num, unit) = s.split_at(split);
let value: f64 = num
.trim()
.parse()
.map_err(|_| format!("`{s}`: not a number before the unit"))?;
if !(value >= 0.0) || !value.is_finite() {
return Err(format!("`{s}`: negative or non-finite"));
}
let lowered = unit.trim().to_ascii_lowercase();
let unit: &str = lowered
.strip_suffix("ib")
.or_else(|| lowered.strip_suffix('b'))
.unwrap_or(lowered.as_str());
let mult: f64 = match unit {
"" | "m" => 1024.0 * 1024.0,
"k" => 1024.0,
"g" => 1024.0 * 1024.0 * 1024.0,
"t" => 1024.0 * 1024.0 * 1024.0 * 1024.0,
other => return Err(format!("`{s}`: unknown unit `{other}` (use K, M, G, T; bare = MiB)")),
};
Ok((value * mult).round() as u64)
}
/// The embed plan the backends run: how many batches in flight and how big
/// each may be. Built by [`EmbedLimits::from_env`] (the unbounded legacy knobs)
/// or [`EmbedLimits::bounded`] (a memory ceiling).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EmbedLimits {
/// Concurrent forwards (the tract fan-out width). `1` on the ort backend.
pub workers: usize,
/// Row cap per batch (`NORNIR_EMBED_BATCH_ROWS`).
pub max_rows: usize,
/// Padded-token cap per batch (`NORNIR_EMBED_BATCH_TOKENS`) — the number
/// the memory model is linear in.
pub max_tokens: usize,
/// The ceiling this plan was fitted to, if any.
pub budget_bytes: Option<u64>,
/// `true` when even the floor plan (1 worker × [`MIN_BATCH_TOKENS`]) is
/// estimated over the budget. The plan still runs; the announcement says so.
pub over_budget: bool,
}
impl EmbedLimits {
/// No ceiling: every core, the given caps.
pub fn unbounded(workers_available: usize, max_rows: usize, max_tokens: usize) -> Self {
Self {
workers: workers_available.max(1),
max_rows: max_rows.max(1),
max_tokens: max_tokens.max(1),
budget_bytes: None,
over_budget: false,
}
}
/// The memory model: baseline + one activation set per in-flight batch.
pub fn estimated_peak_bytes(&self) -> u64 {
EMBED_BASELINE_BYTES
+ self.workers as u64 * self.max_tokens as u64 * EMBED_BYTES_PER_PADDED_TOKEN
}
/// Fit a plan under `budget`: workers first (throughput lives there), then
/// — only if not even one full-width batch fits — halve the token cap down
/// to [`MIN_BATCH_TOKENS`]. Never zero workers; never a silent overshoot.
pub fn bounded(
budget: u64,
workers_available: usize,
max_rows: usize,
max_tokens: usize,
) -> Self {
let per_batch = |tokens: usize| (tokens as u64 * EMBED_BYTES_PER_PADDED_TOKEN).max(1);
let avail = budget.saturating_sub(EMBED_BASELINE_BYTES);
let mut tokens = max_tokens.max(MIN_BATCH_TOKENS);
let fit = (avail / per_batch(tokens)) as usize;
let workers = if fit >= 1 {
fit.min(workers_available.max(1))
} else {
while tokens > MIN_BATCH_TOKENS && per_batch(tokens) > avail {
tokens = (tokens / 2).max(MIN_BATCH_TOKENS);
}
1
};
let mut lim = Self {
workers,
max_rows: max_rows.max(1),
max_tokens: tokens,
budget_bytes: Some(budget),
over_budget: false,
};
lim.over_budget = lim.estimated_peak_bytes() > budget;
lim
}
/// The plan for THIS process: the legacy caps from `NORNIR_EMBED_BATCH_ROWS`
/// / `NORNIR_EMBED_BATCH_TOKENS` (falling back to the given defaults), bounded
/// by [`memory_budget_bytes`] when a ceiling is declared.
pub fn from_env(workers_available: usize, default_rows: usize, default_tokens: usize) -> Self {
let max_rows = env_usize("NORNIR_EMBED_BATCH_ROWS", default_rows);
let max_tokens = env_usize("NORNIR_EMBED_BATCH_TOKENS", default_tokens);
match memory_budget_bytes() {
Some(b) => Self::bounded(b, workers_available, max_rows, max_tokens),
None => Self::unbounded(workers_available, max_rows, max_tokens),
}
}
/// One line for stderr: the plan and its estimated peak, and the budget it
/// was fitted to. Says `OVER BUDGET` when the floor plan does not fit.
pub fn summary(&self) -> String {
let peak = nornir_humanize::bytes(self.estimated_peak_bytes());
match self.budget_bytes {
None => format!(
"embed plan: {} worker(s) × ≤{} padded tokens (≤{} rows) — est. peak ~{peak}, \
NO ceiling (pass --max-memory / {MEMORY_BUDGET_ENV} to bound it)",
self.workers, self.max_tokens, self.max_rows
),
Some(b) => format!(
"embed plan: {} worker(s) × ≤{} padded tokens (≤{} rows) — est. peak ~{peak} \
under a {} ceiling{}",
self.workers,
self.max_tokens,
self.max_rows,
nornir_humanize::bytes(b),
if self.over_budget {
" — OVER BUDGET even at the floor plan (1 worker × 256 tokens); running \
anyway, expect the kernel to disagree"
} else {
""
}
),
}
}
/// Print [`Self::summary`] to stderr ONCE per process — only when a ceiling
/// is in force (the unbounded plan is the pre-existing behaviour and stays
/// quiet). The plan cannot change within a process, so once is exact.
pub fn announce_once(&self) {
static ONCE: std::sync::Once = std::sync::Once::new();
if self.budget_bytes.is_some() {
ONCE.call_once(|| eprintln!("nornir-embed: {}", self.summary()));
}
}
}
/// `$var` as a positive usize, else `default`. Local (std-only) twin of
/// `embed_support::env_usize`, which is compiled only with an embedder feature.
fn env_usize(var: &str, default: usize) -> usize {
std::env::var(var)
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.filter(|&x| x > 0)
.unwrap_or(default)
}
#[cfg(test)]
mod tests {
use super::*;
const ROWS: usize = 16;
const TOKENS: usize = 8192;
const GIB: u64 = 1 << 30;
/// THE ASSERTION OF GAP 3: for every budget that can hold the baseline plus
/// one batch, the plan's estimated peak stays UNDER the budget, uses at
/// least one worker, never more than the box has, and more budget never
/// buys fewer workers.
#[test]
fn bounded_plan_keeps_estimated_peak_under_budget() {
let mut prev_workers = 0usize;
for gib in [1u64, 2, 3, 4, 6, 8, 12, 16, 32] {
let b = gib * GIB;
let lim = EmbedLimits::bounded(b, 16, ROWS, TOKENS);
assert!(lim.workers >= 1 && lim.workers <= 16, "{gib} GiB: {lim:?}");
assert!(
lim.estimated_peak_bytes() <= b,
"{gib} GiB: est. {} > budget — {lim:?}",
lim.estimated_peak_bytes()
);
assert!(!lim.over_budget, "{gib} GiB: {lim:?}");
assert!(lim.workers >= prev_workers, "monotone in budget: {lim:?}");
prev_workers = lim.workers;
assert_eq!(lim.max_rows, ROWS, "rows are not the lever");
}
// Enough budget ⇒ the full box, full-width batches: same as unbounded.
let big = EmbedLimits::bounded(32 * GIB, 16, ROWS, TOKENS);
assert_eq!(big.workers, 16);
assert_eq!(big.max_tokens, TOKENS);
}
/// The two measured points from the inventory, through the model: the
/// default plan on 16 cores prices out near the 5.6 GB that was observed,
/// and a 2 GiB ceiling on the same box gives a plan in the range the
/// `taskset -c 0-5` + `BATCH_ROWS=8` workaround measured (<1.6 GB).
#[test]
fn the_model_reproduces_the_two_measured_points() {
let unb = EmbedLimits::unbounded(16, ROWS, TOKENS);
let est = unb.estimated_peak_bytes();
assert!(est > 5 * GIB && est < 7 * GIB, "16×8192: est {est} bytes, measured 5.6 GB / killed at 7.5 GB");
let two = EmbedLimits::bounded(2 * GIB, 16, ROWS, TOKENS);
assert_eq!(two.workers, 4, "{two:?}");
assert!(two.estimated_peak_bytes() <= 2 * GIB);
assert!(two.estimated_peak_bytes() > GIB, "the budget is SPENT, not left idle: {two:?}");
}
/// Below one full-width batch the planner shrinks the batch, not the worker
/// count (1 is the floor); below the baseline it still returns the floor
/// plan and SAYS it is over budget.
#[test]
fn small_budgets_shrink_the_batch_and_the_floor_is_flagged_not_hidden() {
// 1 GiB: baseline 700 MiB leaves 324 MiB; one 8192-token batch is 320 MiB.
let one = EmbedLimits::bounded(GIB, 16, ROWS, TOKENS);
assert_eq!((one.workers, one.max_tokens), (1, TOKENS), "{one:?}");
assert!(!one.over_budget);
// 800 MiB: 100 MiB of headroom ⇒ halve 8192 → 4096 → 2048 (80 MiB) fits.
let tight = EmbedLimits::bounded(800 << 20, 16, ROWS, TOKENS);
assert_eq!((tight.workers, tight.max_tokens), (1, 2048), "{tight:?}");
assert!(!tight.over_budget);
assert!(tight.estimated_peak_bytes() <= 800 << 20);
// 256 MiB: under the baseline itself ⇒ floor plan, flagged.
let floor = EmbedLimits::bounded(256 << 20, 16, ROWS, TOKENS);
assert_eq!((floor.workers, floor.max_tokens), (1, MIN_BATCH_TOKENS), "{floor:?}");
assert!(floor.over_budget);
assert!(floor.summary().contains("OVER BUDGET"), "{}", floor.summary());
// The summaries name the ceiling and the flag that sets one.
assert!(one.summary().contains("1.00 GiB ceiling") || one.summary().contains("GiB ceiling"), "{}", one.summary());
assert!(EmbedLimits::unbounded(4, ROWS, TOKENS).summary().contains("--max-memory"));
}
#[test]
fn parse_memory_size_accepts_the_usual_spellings() {
assert_eq!(parse_memory_size("2G").unwrap(), 2 * GIB);
assert_eq!(parse_memory_size("2g").unwrap(), 2 * GIB);
assert_eq!(parse_memory_size("2GiB").unwrap(), 2 * GIB);
assert_eq!(parse_memory_size("2GB").unwrap(), 2 * GIB);
assert_eq!(parse_memory_size("1536M").unwrap(), 1536 << 20);
assert_eq!(parse_memory_size("1536").unwrap(), 1536 << 20, "bare = MiB");
assert_eq!(parse_memory_size(" 512k ").unwrap(), 512 << 10);
assert_eq!(parse_memory_size("0.5G").unwrap(), GIB / 2);
assert_eq!(parse_memory_size("1T").unwrap(), 1 << 40);
assert!(parse_memory_size("").is_err());
assert!(parse_memory_size("lots").is_err());
assert!(parse_memory_size("2X").unwrap_err().contains("unknown unit"));
assert!(parse_memory_size("-1G").is_err());
}
/// The CLI-set budget outranks the env var, and `0` clears it. Serialised
/// on the process-global atomic with the env var left untouched.
#[test]
fn process_budget_outranks_env_and_zero_clears() {
set_memory_budget_bytes(3 * GIB);
assert_eq!(memory_budget_bytes(), Some(3 * GIB));
let lim = EmbedLimits::from_env(16, ROWS, TOKENS);
assert_eq!(lim.budget_bytes, Some(3 * GIB));
assert!(lim.estimated_peak_bytes() <= 3 * GIB);
set_memory_budget_bytes(0);
// With nothing set and (in this test process) no env var, the plan is unbounded.
if std::env::var_os(MEMORY_BUDGET_ENV).is_none() {
assert_eq!(EmbedLimits::from_env(16, ROWS, TOKENS).budget_bytes, None);
}
}
}