src/vector/fnvec/ring.rs
//! **The dependency ring's function vectors** — once per `crate@version`.
//!
//! The ring (every dependency crate the driver dumped, `@universe` in the
//! event log) does not move with a commit: it moves when the lockfile does.
//! So it is not embedded inside a commit's consume. It is ONE background job
//! ([`kick`]): started by the `embeddings` feed when the ring's manifest has
//! moved since the last ring run, one at a time per process, written to the
//! job ledger like every other long-running operation.
//!
//! * Unit = one ring crate (`name@version`). A crate that has its ledger row
//! for the current model space is done for good; a lockfile bump embeds only
//! the crate versions that are new.
//! * Source text: the directory of the crate's `Cargo.toml`, asked of `cargo
//! metadata --offline` in the consumer's clone of each member — the registry
//! checkout the driver compiled. The span comes from the lake's
//! `code_span_events`, the same driver spans as everywhere. A crate whose
//! source this process cannot see (a driver that runs inside a toolchain
//! image keeps its registry there) gets a `skipped` ledger row naming it.
//! * One embedding run at a time per process ([`super::embed_turn`]): the job
//! is started when the feed's own run is over, and takes its turn per crate.
//! * The whole job is priced ONCE with the embed budget
//! ([`crate::vector::store::embed_budget`]) before the first crate: on a CPU
//! box a ring is hours to days, and the job then refuses by the numbers —
//! a ledger row, not a silent non-run.
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use anyhow::{Context, Result, anyhow};
use arrow::array::{Array, Int64Array, StringArray, UInt32Array};
use mimir_ingest::{EventLog, FoldScope, Kind};
use mimir_room::{EdgeSel, MaskSpec, Room, Scope};
use mimir_schema::{EdgeKind, Origin, SPAN_NODE};
use super::{At, Embedding, FnChunk, FnNode, RING, Run, links};
use crate::warehouse::iceberg::IcebergWarehouse;
/// One ring crate and its function nodes.
#[derive(Debug, Clone, Default)]
pub struct RingCrate {
/// `name@version`.
pub krate: String,
/// Where its source is (`None` = no manifest path known for it).
pub dir: Option<PathBuf>,
pub functions: Vec<FnNode>,
/// Function nodes a macro wrote, or without a span.
pub without_chunk: usize,
}
/// What a ring run would do.
#[derive(Debug, Default)]
pub struct Plan {
/// The crates that have no ledger row for this model space yet.
pub todo: Vec<RingCrate>,
/// Ring crates that already have one.
pub done: usize,
/// The ring's manifest snapshot this plan is of (the lake's own marker:
/// it moves exactly when the ring does).
pub manifest: Option<i64>,
/// `fn` + `method` nodes defined in ring crates.
pub callables: usize,
}
impl Plan {
pub fn functions(&self) -> usize {
self.todo.iter().map(|c| c.functions.len()).sum()
}
pub fn unlocated(&self) -> Vec<&str> {
self.todo.iter().filter(|c| c.dir.is_none()).map(|c| c.krate.as_str()).collect()
}
pub fn line(&self) -> String {
format!(
"ring: {} crate(s) to embed ({} function(s); {} without a source directory), {} already done; {} function-like node(s) in the ring",
self.todo.len(),
self.functions(),
self.unlocated().len(),
self.done,
self.callables
)
}
}
/// `name@version` → the directory of its manifest, out of a `cargo metadata`
/// document (only directories that exist).
pub fn dirs_from_metadata(meta: &serde_json::Value) -> BTreeMap<String, PathBuf> {
let mut out = BTreeMap::new();
for p in meta["packages"].as_array().into_iter().flatten() {
let (Some(name), Some(version), Some(manifest)) = (p["name"].as_str(), p["version"].as_str(), p["manifest_path"].as_str()) else { continue };
if let Some(dir) = Path::new(manifest).parent().filter(|d| !manifest.is_empty() && d.is_dir()) {
out.entry(format!("{name}@{version}")).or_insert_with(|| dir.to_path_buf());
}
}
out
}
/// Where the source of every crate of the consumed members' builds is:
/// `cargo metadata --offline` in each member's consumer clone
/// (`<builds>/mimir_v2/clones/<member>`, the tree the driver compiled, at the
/// consumed commit). Never the network. The second value names the members it
/// could not ask (no cargo on this box, a workspace that does not load): their
/// crates then have no source directory, and the run says so per crate.
pub fn crate_dirs(log: &EventLog, clones_root: &Path) -> (BTreeMap<String, PathBuf>, Vec<String>) {
let members: Vec<String> = on_lake(log, || async { Ok(log.offsets().await?.into_keys().collect()) }).unwrap_or_default();
let (mut out, mut failed) = (BTreeMap::new(), Vec::new());
for member in members {
let dir = clones_root.join(&member);
let ran = std::process::Command::new("cargo").args(["metadata", "--offline", "--format-version", "1"]).current_dir(&dir).output();
match ran {
Ok(o) if o.status.success() => match serde_json::from_slice::<serde_json::Value>(&o.stdout) {
Ok(meta) => {
for (k, v) in dirs_from_metadata(&meta) {
out.entry(k).or_insert(v);
}
}
Err(e) => failed.push(format!("{member}: cargo metadata printed no JSON ({e})")),
},
Ok(o) => failed.push(format!("{member}: cargo metadata failed in {}: {}", dir.display(), String::from_utf8_lossy(&o.stderr).lines().next().unwrap_or(""))),
Err(e) => failed.push(format!("{member}: cargo could not be started ({e})")),
}
}
(out, failed)
}
fn on_lake<T: Send, F: Future<Output = Result<T>>>(log: &EventLog, make: impl FnOnce() -> F + Send) -> Result<T> {
// on a thread of its own, on the lake's runtime: callable from any thread
std::thread::scope(|s| s.spawn(|| log.block_on(make())).join().unwrap_or_else(|p| std::panic::resume_unwind(p)))
}
/// The spans of the ring's nodes named in `wanted` (moniker → file, start,
/// end), from the lake's span events. The room does not fold them: dependency
/// sources are not a workspace's files.
fn ring_spans(log: &EventLog, wanted: &HashSet<&str>) -> Result<HashMap<String, (String, u32, u32)>> {
let scope = FoldScope { repos: BTreeMap::new(), universe: true };
let rows = on_lake(log, || log.fold_rows(Kind::Spans, &scope, &["target", "key", "file", "start_line", "end_line"]))?;
let line = |b: &arrow::array::RecordBatch, name: &str, r: usize| -> Result<u32> {
let col = b.column_by_name(name).ok_or_else(|| anyhow!("code_span_events has no `{name}`"))?;
if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
return Ok(a.value(r) as u32);
}
col.as_any().downcast_ref::<UInt32Array>().map(|a| a.value(r)).ok_or_else(|| anyhow!("code_span_events.{name} is neither int64 nor uint32"))
};
let text = |b: &arrow::array::RecordBatch, name: &str| -> Result<StringArray> {
b.column_by_name(name).and_then(|c| c.as_any().downcast_ref::<StringArray>()).cloned().ok_or_else(|| anyhow!("code_span_events has no string column `{name}`"))
};
let cols: Vec<(StringArray, StringArray, StringArray)> = rows.batches.iter().map(|b| Ok((text(b, "target")?, text(b, "key")?, text(b, "file")?))).collect::<Result<_>>()?;
let mut out: HashMap<String, (String, u32, u32)> = HashMap::new();
for &(b, r) in &rows.open {
let (b, r) = (b as usize, r as usize);
let (target, key, file) = &cols[b];
if target.value(r) != SPAN_NODE || !wanted.contains(key.value(r)) {
continue;
}
let span = (file.value(r).to_string(), line(&rows.batches[b], "start_line", r)?, line(&rows.batches[b], "end_line", r)?);
// several units may position one id: the first in (file, line) order, as the room does
match out.get(key.value(r)) {
Some(prev) if *prev <= span => {}
_ => {
out.insert(key.value(r).to_string(), span);
}
}
}
Ok(out)
}
/// The ring crates that have no ledger row yet under `model_profile`, with
/// their function nodes positioned and their source directories resolved.
/// `retry_skipped`: also the crates a former run skipped (no source directory).
pub fn plan(wh: &IcebergWarehouse, log: &EventLog, dirs: &BTreeMap<String, PathBuf>, model_profile: &str, retry_skipped: bool) -> Result<Plan> {
let manifest = on_lake(log, || log.manifest_snapshot())?;
let settled: BTreeSet<String> = links::read_runs(wh)?
.into_iter()
.filter(|r| r.repo == RING && !r.unit.is_empty() && r.model_profile == model_profile)
.filter(|r| r.status == links::STATUS_OK || (r.status == links::STATUS_SKIPPED && !retry_skipped))
.map(|r| r.unit)
.collect();
// the ring's own nodes: the room with the universe folded in
let scope = Scope::current().with_universe(true);
let room = on_lake(log, || Room::build_uncached(log, &scope))?;
let mut ords: Vec<u64> = Vec::new();
for k in super::room::KINDS {
let spec = MaskSpec { node_kinds: vec![k.to_string()], origins: vec![Origin::Library, Origin::Lang], defined_only: true, ..Default::default() };
ords.extend_from_slice(room.mask(&spec)?.value.ordinals());
}
ords.sort_unstable();
ords.dedup();
let mut owners: BTreeSet<u64> = ords.iter().copied().collect();
owners.extend(ords.iter().filter_map(|o| room.node(*o).parent));
let owners = room.mask_of(&owners.into_iter().collect::<Vec<_>>());
let generated: BTreeSet<u64> =
room.edges(&EdgeSel { kinds: vec![EdgeKind::Expands], from: Some(&owners), to: None, cap: None, heaviest_first: false })?.value.edges.into_iter().map(|e| e.from).collect();
let unit_of = |o: u64| {
let n = room.node(o);
format!("{}@{}", n.krate.unwrap_or("?"), n.version.unwrap_or("?"))
};
let mut by_crate: BTreeMap<String, RingCrate> = BTreeMap::new();
let mut own: Vec<u64> = Vec::new();
let mut done: BTreeSet<String> = BTreeSet::new();
for &o in &ords {
let unit = unit_of(o);
if settled.contains(&unit) {
done.insert(unit);
continue;
}
let c = by_crate.entry(unit.clone()).or_insert_with(|| RingCrate { dir: dirs.get(&unit).cloned(), krate: unit, ..Default::default() });
if generated.contains(&o) || room.node(o).parent.is_some_and(|p| generated.contains(&p)) {
c.without_chunk += 1;
} else {
own.push(o);
}
}
let wanted: HashSet<&str> = own.iter().map(|o| room.moniker(*o)).collect();
let spans = if wanted.is_empty() { HashMap::new() } else { ring_spans(log, &wanted)? };
for &o in &own {
let c = by_crate.get_mut(&unit_of(o)).expect("grouped above");
match spans.get(room.moniker(o)) {
Some((file, start, end)) => c.functions.push(FnNode { moniker: room.moniker(o).to_string(), file: file.clone(), start_line: *start, end_line: *end }),
None => c.without_chunk += 1,
}
}
Ok(Plan { todo: by_crate.into_values().collect(), done: done.len(), manifest, callables: ords.len() })
}
/// Cut a ring crate's functions out of its source directory. Only a file path
/// the driver wrote relative to the manifest is read: a file outside the
/// package (build-script output in a target directory) is not its source.
fn cut_crate(c: &RingCrate) -> (Vec<FnChunk>, usize) {
let Some(dir) = &c.dir else { return (Vec::new(), c.functions.len()) };
let (chunks, unreadable) = super::chunks(c.functions.clone(), |f| super::room::read_under(dir, f));
(chunks, unreadable.len())
}
/// What a ring run did.
#[derive(Debug, Default, Clone)]
pub struct Report {
pub crates_embedded: usize,
pub crates_skipped: usize,
pub functions: usize,
pub new_vectors: usize,
pub embed_ms: u64,
pub ms: u64,
pub backend: String,
}
impl Report {
pub fn line(&self) -> String {
let rate = if self.embed_ms > 0 { format!(" at {:.2} chunk/s on {}", self.new_vectors as f64 * 1000.0 / self.embed_ms as f64, self.backend) } else { String::new() };
format!(
"ring: {} crate(s) embedded, {} skipped (no source); {} function(s), {} new vector(s){rate}; {:.1}s",
self.crates_embedded,
self.crates_skipped,
self.functions,
self.new_vectors,
self.ms as f64 / 1000.0
)
}
}
/// Run `plan`: price the whole job, then crate by crate — embed what has no
/// vector, write the links, write the crate's ledger row. A crate without a
/// source directory gets a `skipped` row naming it. When every crate is
/// settled, the ring-level row records the manifest snapshot the ring is of.
pub fn run(wh: &IcebergWarehouse, plan: &Plan, embedding: &mut dyn Embedding, on: &mut dyn FnMut(&str, &Run)) -> Result<Report> {
let t0 = std::time::Instant::now();
let model = embedding.profile();
let profile = model.id();
let mut stored = wh.block_on(crate::vector::store::existing_content_hashes_in(wh, crate::warehouse::iceberg::TABLE_EMBEDDINGS_RING, &profile))?;
let events = links::read_events_of(wh, &[RING])?;
let state = links::open(&events, RING, i64::MAX);
drop(events);
// ── price the WHOLE job once: a stride sample of what a forward pass will be
// handed (the pieces of the texts that need a vector) ──
let mut total_pieces = 0usize;
let mut sample: Vec<String> = Vec::new();
{
let mut seen: HashSet<String> = HashSet::new();
for c in &plan.todo {
for ch in cut_crate(c).0 {
if !stored.contains(&ch.content_hash) && seen.insert(ch.content_hash.clone()) {
for piece in super::pieces(&ch.text) {
if total_pieces % 97 == 0 && sample.len() < 64 {
sample.push(piece.to_string());
}
total_pieces += 1;
}
}
}
}
}
if total_pieces > 0 {
let _turn = super::embed_turn();
embedding.price_sampled(&sample.iter().map(String::as_str).collect::<Vec<_>>(), total_pieces)?;
}
let mut rep = Report { backend: embedding.backend(), ..Default::default() };
for c in &plan.todo {
let at = At::ring(&c.krate);
if c.dir.is_none() {
super::record_not_run(wh, &at, links::STATUS_SKIPPED, &format!("no source directory known for {} ({} function(s)): `cargo metadata` of the consumed members does not name its manifest on this box", c.krate, c.functions.len()))?;
rep.crates_skipped += 1;
continue;
}
// one crate per turn: a consume's feed gets in between two crates
let _turn = super::embed_turn();
let (chunks, unreadable) = cut_crate(c);
let mine: BTreeMap<String, String> = chunks.iter().filter_map(|ch| state.get(&ch.node.moniker).map(|h| (ch.node.moniker.clone(), h.clone()))).collect();
let run = super::embed_against(wh, &at, &chunks, c.without_chunk + unreadable, embedding, &mine, &mut stored)?;
rep.crates_embedded += 1;
rep.functions += run.functions;
rep.new_vectors += run.new_vectors;
rep.embed_ms += run.embed_ms;
on(&c.krate, &run);
}
rep.ms = t0.elapsed().as_millis() as u64;
// the ring-level row: the ring is embedded as of this manifest snapshot
let ring_run = Run { functions: rep.functions, new_vectors: rep.new_vectors, embed_ms: rep.embed_ms, ms: rep.ms, backend: rep.backend.clone(), ..Default::default() };
let at = At { repo: RING.to_string(), unit: String::new(), seq: 0, sha: manifest_mark(plan.manifest) };
links::append_run(wh, &links::RunRow::of(&at, &profile, &embedding.model_id(), &ring_run, links::STATUS_OK, Some(&rep.line())))?;
Ok(rep)
}
fn manifest_mark(manifest: Option<i64>) -> String {
manifest.map(|m| m.to_string()).unwrap_or_default()
}
/// Has the ring moved since the last ring run under `model_profile`? Cheap:
/// the lake's manifest snapshot against the newest ring-level ledger row.
pub fn due(wh: &IcebergWarehouse, log: &EventLog, model_profile: &str) -> Result<bool> {
let manifest = on_lake(log, || log.manifest_snapshot())?;
if manifest.is_none() {
return Ok(false); // nothing dumped into the ring yet
}
let mark = manifest_mark(manifest);
let last = links::read_runs(wh)?.into_iter().filter(|r| r.repo == RING && r.unit.is_empty() && r.model_profile == model_profile).max_by_key(|r| r.ts_micros);
// a ring run that was refused (over budget) or failed for THIS manifest is not asked again by itself
Ok(last.is_none_or(|r| r.sha != mark))
}
/// Where the background job reads and writes.
pub struct Host {
pub workspace: String,
pub warehouse: Arc<IcebergWarehouse>,
pub log: EventLog,
/// `<builds>/mimir_v2/clones`: the consumer's clones of the members.
pub clones_root: PathBuf,
/// The job ledger (`None`: the run is announced on stderr only).
pub jobs: Option<crate::jobs::JobSink>,
/// The workspace's jera lane (a server): the ring run is QUEUED there
/// ([`run_queued`] is its runner), never on a thread of its own. `None`
/// (the CLI, a test): the run takes a thread, as before.
pub queue: Option<Arc<crate::jobs::chain::ChainQueue>>,
}
/// Hosts whose ring run is queued on a jera lane, by workspace — what
/// [`run_queued`] picks up when the lane gets to the unit.
static QUEUED: std::sync::Mutex<std::collections::BTreeMap<String, Host>> = std::sync::Mutex::new(std::collections::BTreeMap::new());
/// **The ring run, as a queued jera job's body** — called by the lane runner
/// for the `vector_embed` unit `ring:<workspace>` that [`kick`] submitted. The
/// chain writes the job's ledger row; this returns what the row says.
pub fn run_queued(workspace: &str) -> Result<String> {
let Some(host) = QUEUED.lock().unwrap_or_else(|e| e.into_inner()).remove(workspace) else {
return Ok("nothing queued for this workspace (an earlier unit took it)".into());
};
if RUNNING.swap(true, Ordering::AcqRel) {
anyhow::bail!("a ring run is already going in this process");
}
let out = background(&host);
RUNNING.store(false, Ordering::Release);
out
}
/// What [`kick`] did.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Kick {
/// A ring job was started on a thread of its own.
Started,
/// A ring job was queued on the workspace's jera lane.
Queued,
/// One is already running in this process.
Running,
/// The ring has not moved since the last ring run.
NotDue,
/// This build cannot embed.
NoEmbedder(String),
}
static RUNNING: AtomicBool = AtomicBool::new(false);
/// **Start the ring job if the ring moved** — called by the `embeddings` feed
/// after a member's functions are in. Returns at once: the job runs on its own
/// thread, outside the consume and outside the workspace gate (it reads the
/// lake at one committed release and dependency sources, never a member's
/// checkout), and appends through the one vector writer lock.
pub fn kick(host: Host) -> Kick {
if let Some(why) = super::no_embedder() {
return Kick::NoEmbedder(why.to_string());
}
if RUNNING.swap(true, Ordering::AcqRel) {
return Kick::Running;
}
match due_now(&host) {
Ok(true) => {}
Ok(false) => {
RUNNING.store(false, Ordering::Release);
return Kick::NotDue;
}
Err(e) => eprintln!("nornir-embed: {}: ring check failed, trying the run: {e:#}", host.workspace),
}
// a server: the run is a QUEUED jera job (low priority, one unit per
// workspace — a second kick while it waits is the same unit)
if let Some(queue) = host.queue.clone() {
RUNNING.store(false, Ordering::Release);
let ws = host.workspace.clone();
QUEUED.lock().unwrap_or_else(|e| e.into_inner()).insert(ws.clone(), host);
let key = format!("ring:{ws}");
let spec = crate::jobs::chain::ChainSpec::new(crate::jobs::kind::VECTOR_EMBED, key.clone(), key, 0, &ws, serde_json::json!({ "scope": "ring" }))
.with_target("ring")
.with_priority(crate::jobs::chain::Priority::Low);
queue.submit(spec);
return Kick::Queued;
}
let spawned = std::thread::Builder::new().name("nornir-ring-embed".into()).spawn(move || {
let job = host.jobs.clone().map(|sink| crate::jobs::JobHandle::start(sink, crate::jobs::kind::VECTOR_EMBED, "ring", &host.workspace, serde_json::json!({ "scope": "dependency ring", "unit": "crate@version" })));
let out = background(&host);
RUNNING.store(false, Ordering::Release);
match (out, job) {
(Ok(line), Some(job)) => job.finish(serde_json::json!({ "result": line }), "embedding_node_runs:@universe"),
(Err(e), Some(job)) => job.fail_with_detail(serde_json::json!({ "error": format!("{e:#}") })),
(Ok(line), None) => eprintln!("nornir-embed: {}: {line}", host.workspace),
(Err(e), None) => eprintln!("nornir-embed: {}: ring job failed: {e:#}", host.workspace),
}
});
match spawned {
Ok(_) => Kick::Started,
Err(e) => {
RUNNING.store(false, Ordering::Release);
eprintln!("nornir-embed: the ring job could not start: {e}");
Kick::NotDue
}
}
}
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
fn due_now(host: &Host) -> Result<bool> {
due(&host.warehouse, &host.log, &crate::vector::embed_support::profile().id())
}
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
fn due_now(_: &Host) -> Result<bool> {
Ok(false)
}
/// The job's body: plan, run, and — when the run is refused or fails — the
/// ring-level ledger row that says so for this manifest (so the next commit
/// does not start the same refused job again).
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
fn background(host: &Host) -> Result<String> {
let mut embedding = super::Compiled::default();
let profile = embedding.profile().id();
let (dirs, unasked) = crate_dirs(&host.log, &host.clones_root);
for why in &unasked {
eprintln!("nornir-embed: {}: ring sources: {why}", host.workspace);
}
let plan = plan(&host.warehouse, &host.log, &dirs, &profile, false).context("plan the ring")?;
eprintln!("nornir-embed: {}: {}", host.workspace, plan.line());
let ws = host.workspace.clone();
let ran = run(&host.warehouse, &plan, &mut embedding, &mut |krate, run| eprintln!("nornir-embed: {ws}: ring {krate}: {}", run.line()));
match ran {
Ok(rep) => Ok(rep.line()),
Err(e) => {
let at = At { repo: RING.to_string(), unit: String::new(), seq: 0, sha: manifest_mark(plan.manifest) };
let mut row = links::RunRow::of(&at, &profile, &embedding.model_id(), &Run { backend: embedding.backend(), ..Default::default() }, links::STATUS_ERROR, Some(&format!("{e:#}")));
row.functions = plan.functions() as i64;
let _ = links::append_run(&host.warehouse, &row);
Err(e)
}
}
}
#[cfg(not(any(feature = "embed-tract", feature = "embed-ort")))]
fn background(_: &Host) -> Result<String> {
Err(anyhow!("{}", super::no_embedder().unwrap_or("no embedder")))
}