GameCult/Huginn
Watch
9
0
Fork
You've already forked Huginn
 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

pool: identities are read off the batch's clock and held to the collection's model

Browse source 
pool-blackhole-stalls: each due endpoint's identity is read on its own thread
under a 5 s bound, and a batch waits for reads only while no healthy primary
exists to start on, so a dark endpoint never delays work a healthy one can do.
pool-pins-first-identity: the pool is told the model the collection was built
with (Embedder::built_under, called by reconcile and refresh before the
identity read) and never writes with an endpoint serving another; it adopts a
model only when there is no collection, in list order.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Meta Vens 2026-10-31 07:54:57 +01:00
commit 13647f9220
3 changed files with 315 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

15
crates/huginn-daemon/src/index/mod.rs
View file

@ -155,6 +155,13 @@ pub struct CollectionMeta {
pub index_text_version: u32,
}
impl CollectionMeta {
/// The model the collection was built with.
pub fn model_identity(&self) -> ModelIdentity {
ModelIdentity { name: self.model.clone(), digest: self.model_digest.clone(), dimensions: self.dimensions }
}
}
/// What the vector store holds under a collection's name.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Described {
@ -192,6 +199,10 @@ pub struct Hit {
}
pub trait Embedder {
/// The model the collection about to be served was built with, or `None`
/// when there is no collection yet. Called before `model_identity`. One
/// endpoint has nothing to do with it; a pool holds its endpoints to it.
fn built_under(&mut self, _built: Option<&ModelIdentity>) {}
/// The model's identity, or an error once `within` has passed.
fn model_identity(&mut self, within: Duration) -> Result<ModelIdentity>;
/// One vector per text, in order, or an error once `within` has passed.
@ -417,10 +428,12 @@ impl<E: Embedder, V: VectorIndex> Projector<E, V> {
let Some(mind) = self.mind.clone() else {
return Err(StepError::Unavailable(anyhow::anyhow!("the mind has no identity yet")));
};
let described = self.index.describe(&self.collection, INDEX_CALL_TIMEOUT).context("describing the collection")?;
self.embedder.built_under(if let Described::Labelled(found) = &described { Some(found.model_identity()) } else { None }.as_ref());
let identity =
self.embedder.model_identity(BULK_EMBED_TIMEOUT).context("reading the embedding model's identity")?;
let wanted = self.collection_meta(&mind, &identity);
match self.index.describe(&self.collection, INDEX_CALL_TIMEOUT).context("describing the collection")? {
match described {
Described::Absent => {
// A collection lost while the worker ran was an index; one that
// was never made is a first build.

368
crates/huginn-daemon/src/index/pool.rs
View file

@ -10,11 +10,21 @@
//! `retry_after` has passed. A piece that fails on one endpoint is embedded
//! again on another, so a batch fails only when no endpoint is left to take it.
//!
//! The model of the pool is the first identity any endpoint reports, primaries
//! first. An endpoint's identity is read when it is admitted, not on every
//! batch: a model re-pulled under a running pool is noticed at the next
//! admission, and a writer is a short-lived job.
//! The model of the pool is the one the collection was built with: the caller
//! hands it over with `built_under` before it reads the identity, and an
//! endpoint serving anything else never takes work, even when it is the only
//! one up, so a run cannot rebuild an index around a foreign model. Only when
//! there is no collection yet does the pool adopt a model: the first one
//! reported in list order, primaries first, once every endpoint has answered.
//!
//! An endpoint's identity is read off the batch's clock. Each read runs on its
//! own thread under `IDENTITY_WITHIN`, and a batch waits for the reads only
//! while no healthy primary exists to start on: a primary being read keeps the
//! reserve waiting, so a returning GPU is used and not passed over for the CPU,
//! but an endpoint that drops packets never delays work a healthy one can do.
use std::panic::AssertUnwindSafe;
use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError};
use std::time::{Duration, Instant};
use anyhow::{Result, bail};
@ -24,81 +34,93 @@ use super::{Embedder, ModelIdentity};
/// How long an endpoint that dropped out is left alone by default.
pub const RETRY_AFTER: Duration = Duration::from_secs(60);
/// The bound on one identity read, connect and answer together. Reading a
/// model's identity loads nothing, so an endpoint that needs longer is dark.
pub const IDENTITY_WITHIN: Duration = Duration::from_secs(5);
/// One piece's outcome: the member that took it, where it starts, how long it is,
/// and what came back.
type Answer = (usize, usize, usize, Result<Vec<Vec<f32>>>);
struct Member<E> {
struct Health {
name: String,
embedder: E,
reserve: bool,
/// `None` while the endpoint is in the pool; otherwise when it may be read
/// again, and why it left.
out: Option<(Instant, String)>,
probing: bool,
/// What the last read found, until the roster has judged it.
seen: Option<Result<ModelIdentity, String>>,
}
pub struct EmbedderPool<E: Embedder + Send> {
members: Vec<Member<E>>,
struct Roster {
pinned: Option<ModelIdentity>,
retry_after: Duration,
members: Vec<Health>,
}
impl<E: Embedder + Send> EmbedderPool<E> {
/// `primaries` share the work; `reserve` serves only when none of them can.
/// Every endpoint starts out, to be admitted by its first identity read.
pub fn new(primaries: Vec<(String, E)>, reserve: Vec<(String, E)>) -> Self {
let member = |reserve: bool| {
move |(name, embedder): (String, E)| Member { name, embedder, reserve, out: Some((Instant::now(), "not yet read".into())) }
};
let members = primaries.into_iter().map(member(false)).chain(reserve.into_iter().map(member(true))).collect();
Self { members, pinned: None, retry_after: RETRY_AFTER }
}
struct Shared {
roster: Mutex<Roster>,
changed: Condvar,
}
pub fn retrying_after(mut self, retry_after: Duration) -> Self {
self.retry_after = retry_after;
self
}
fn locked<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}
/// Reads the identity of every endpoint that is out and due, in list order,
/// and lets in those that serve the pool's model.
fn admit(&mut self, within: Duration) {
let now = Instant::now();
for at in 0..self.members.len() {
if self.members[at].out.as_ref().is_none_or(|(due, _)| *due > now) {
continue;
}
let name = self.members[at].name.clone();
let outcome = match self.members[at].embedder.model_identity(within) {
Ok(identity) => {
let pinned = self.pinned.get_or_insert_with(|| identity.clone());
if *pinned == identity {
None
} else {
Some(format!("serves {identity:?}, not {pinned:?}"))
}
}
Err(error) => Some(format!("{error:#}")),
};
if let Some(why) = &outcome {
eprintln!("huginn-sources: embedder {name} is out of the pool: {why}");
}
self.members[at].out = outcome.map(|why| (Instant::now() + self.retry_after, why));
}
}
fn drop_out(&mut self, at: usize, why: String) {
eprintln!("huginn-sources: embedder {} is out of the pool: {why}", self.members[at].name);
self.members[at].out = Some((Instant::now() + self.retry_after, why));
impl Roster {
fn healthy(&self, reserve: bool) -> Vec<usize> {
(0..self.members.len()).filter(|at| self.members[*at].reserve == reserve && self.members[*at].out.is_none()).collect()
}
/// The members to work with now: the healthy primaries, or if none, the
/// healthy reserve.
fn active(&self) -> Vec<usize> {
let healthy = |reserve: bool| -> Vec<usize> {
(0..self.members.len()).filter(|at| self.members[*at].reserve == reserve && self.members[*at].out.is_none()).collect()
};
let primaries = healthy(false);
if primaries.is_empty() { healthy(true) } else { primaries }
let primaries = self.healthy(false);
if primaries.is_empty() { self.healthy(true) } else { primaries }
}
fn reading(&self, reserve: bool) -> bool {
self.members.iter().any(|member| member.reserve == reserve && member.probing)
}
/// Whether a batch may start: a healthy primary is there to start on; else
/// no primary is still being read, and the reserve is healthy or nothing is
/// left being read. Without a model to hold endpoints to, it waits for every
/// read, so the model adopted does not depend on who answered first.
fn ready(&self) -> bool {
if self.pinned.is_none() && self.members.iter().any(|member| member.probing) {
return false;
}
if !self.healthy(false).is_empty() {
return true;
}
if self.reading(false) {
return false;
}
!self.healthy(true).is_empty() || !self.reading(true)
}
/// Adopts a model when there is none and every read is in, then lets in the
/// endpoints that serve it and puts out the rest.
fn judge(&mut self, retry_after: Duration) {
if self.pinned.is_none() && !self.members.iter().any(|member| member.probing) {
self.pinned =
self.members.iter().find_map(|member| member.seen.as_ref().and_then(|seen| seen.as_ref().ok())).cloned();
}
let pinned = self.pinned.clone();
for member in &mut self.members {
let why = match (&pinned, member.seen.take()) {
(_, None) | (None, Some(Ok(_))) => continue,
(_, Some(Err(why))) => why,
(Some(pinned), Some(Ok(identity))) if *pinned == identity => {
member.out = None;
continue;
}
(Some(pinned), Some(Ok(identity))) => format!("serves {identity:?}, not {pinned:?}"),
};
eprintln!("huginn embedder pool: {} is out of the pool: {why}", member.name);
member.out = Some((Instant::now() + retry_after, why));
}
}
fn unavailable(&self) -> anyhow::Error {
@ -111,12 +133,108 @@ impl<E: Embedder + Send> EmbedderPool<E> {
}
}
impl<E: Embedder + Send> Embedder for EmbedderPool<E> {
/// Reads one endpoint's identity and tells the roster. Runs on its own thread.
fn probe<E: Embedder>(shared: &Shared, embedder: &Mutex<E>, at: usize, retry_after: Duration) {
let read = std::panic::catch_unwind(AssertUnwindSafe(|| locked(embedder).model_identity(IDENTITY_WITHIN)))
.unwrap_or_else(|_| Err(anyhow::anyhow!("the embedder panicked")));
let mut roster = locked(&shared.roster);
roster.members[at].probing = false;
roster.members[at].seen = Some(read.map_err(|error| format!("{error:#}")));
roster.judge(retry_after);
drop(roster);
shared.changed.notify_all();
}
pub struct EmbedderPool<E: Embedder + Send + 'static> {
embedders: Vec<Arc<Mutex<E>>>,
shared: Arc<Shared>,
retry_after: Duration,
}
impl<E: Embedder + Send + 'static> EmbedderPool<E> {
/// `primaries` share the work; `reserve` serves only when none of them can.
/// Every endpoint starts out, to be admitted by its first identity read.
pub fn new(primaries: Vec<(String, E)>, reserve: Vec<(String, E)>) -> Self {
let mut embedders = Vec::new();
let mut members = Vec::new();
let tiers = primaries.into_iter().map(|primary| (false, primary)).chain(reserve.into_iter().map(|spare| (true, spare)));
for (is_reserve, (name, embedder)) in tiers {
embedders.push(Arc::new(Mutex::new(embedder)));
members.push(Health { name, reserve: is_reserve, out: Some((Instant::now(), "not yet read".into())), probing: false, seen: None });
}
let shared = Arc::new(Shared { roster: Mutex::new(Roster { pinned: None, members }), changed: Condvar::new() });
Self { embedders, shared, retry_after: RETRY_AFTER }
}
pub fn retrying_after(mut self, retry_after: Duration) -> Self {
self.retry_after = retry_after;
self
}
/// Starts a read of every endpoint that is out and due, each on its own
/// thread, and waits for them only until a batch can start (`Roster::ready`)
/// or `within` has passed.
fn admit(&self, within: Duration) {
let started = Instant::now();
let mut roster = locked(&self.shared.roster);
for at in 0..self.embedders.len() {
let member = &mut roster.members[at];
if member.probing || member.out.as_ref().is_none_or(|(due, _)| *due > started) {
continue;
}
member.probing = true;
let (shared, embedder, retry_after) = (Arc::clone(&self.shared), Arc::clone(&self.embedders[at]), self.retry_after);
std::thread::spawn(move || probe(&shared, &embedder, at, retry_after));
}
while !roster.ready() {
let left = (started + within).saturating_duration_since(Instant::now());
if left.is_zero() {
break;
}
roster = self.shared.changed.wait_timeout(roster, left).unwrap_or_else(PoisonError::into_inner).0;
}
}
fn drop_out(&self, at: usize, why: String) {
let mut roster = locked(&self.shared.roster);
eprintln!("huginn embedder pool: {} is out of the pool: {why}", roster.members[at].name);
roster.members[at].out = Some((Instant::now() + self.retry_after, why));
}
fn active(&self) -> Vec<usize> {
locked(&self.shared.roster).active()
}
/// Waits until no read is running: a test's way to say the pool has heard
/// every endpoint.
#[cfg(test)]
fn settle(&self) {
let mut roster = locked(&self.shared.roster);
while roster.members.iter().any(|member| member.probing) {
roster = self.shared.changed.wait_timeout(roster, Duration::from_millis(50)).unwrap_or_else(PoisonError::into_inner).0;
}
}
}
impl<E: Embedder + Send + 'static> Embedder for EmbedderPool<E> {
fn built_under(&mut self, built: Option<&ModelIdentity>) {
let mut roster = locked(&self.shared.roster);
if roster.pinned.as_ref() == built {
return;
}
roster.pinned = built.cloned();
for member in &mut roster.members {
member.seen = None;
member.out = Some((Instant::now(), "not yet read under the collection's model".into()));
}
}
fn model_identity(&mut self, within: Duration) -> Result<ModelIdentity> {
self.admit(within);
match &self.pinned {
Some(pinned) if !self.active().is_empty() => Ok(pinned.clone()),
_ => Err(self.unavailable()),
let roster = locked(&self.shared.roster);
match &roster.pinned {
Some(pinned) if !roster.active().is_empty() => Ok(pinned.clone()),
_ => Err(roster.unavailable()),
}
}
@ -126,9 +244,15 @@ impl<E: Embedder + Send> Embedder for EmbedderPool<E> {
// Pieces still owed, as (first index, length) over `texts`.
let mut owed = vec![(0, texts.len())];
while !owed.is_empty() {
let active = self.active();
let mut active = self.active();
if active.is_empty() {
bail!("{}", self.unavailable());
// Every member that took the last round failed; the reserve may
// still be on its first read.
self.admit(within);
active = self.active();
}
if active.is_empty() {
bail!("{}", locked(&self.shared.roster).unavailable());
}
// Each owed piece is cut evenly over the active members.
let mut jobs: Vec<(usize, usize, usize)> = Vec::new(); // (member, first, length)
@ -140,15 +264,15 @@ impl<E: Embedder + Send> Embedder for EmbedderPool<E> {
}
let mut answers: Vec<Answer> = Vec::new();
std::thread::scope(|scope| {
let mut taken: Vec<Option<&mut Member<E>>> = self.members.iter_mut().map(Some).collect();
let mut busy = vec![false; self.embedders.len()];
let mut running = Vec::new();
for (member, first, length) in jobs {
// A member holds one job per round; later pieces for it wait for the next.
if let Some(held) = taken[member].take() {
let piece = &texts[first..first + length];
running.push((member, first, length, scope.spawn(move || held.embedder.embed(piece, within))));
} else {
if std::mem::replace(&mut busy[member], true) {
owed.push((first, length));
} else {
let (piece, embedder) = (&texts[first..first + length], &self.embedders[member]);
running.push((member, first, length, scope.spawn(move || locked(embedder).embed(piece, within))));
}
}
for (member, first, length, handle) in running {
@ -181,7 +305,7 @@ impl<E: Embedder + Send> Embedder for EmbedderPool<E> {
#[cfg(test)]
mod tests {
use super::*;
use crate::index::fakes::{FakeEmbedder, Gate};
use crate::index::fakes::{FakeEmbedder, Gate, identity};
const WITHIN: Duration = Duration::from_secs(1);
@ -194,6 +318,12 @@ mod tests {
(EmbedderPool::new(primaries, reserve).retrying_after(Duration::ZERO), handles)
}
/// Hears every endpoint once, so a test starts from a pool that knows who is up.
fn warm(pool: &EmbedderPool<FakeEmbedder>) {
pool.admit(WITHIN);
pool.settle();
}
fn texts(count: usize) -> Vec<String> {
(0..count).map(|n| "x".repeat(n + 1)).collect()
}
@ -215,6 +345,7 @@ mod tests {
let b = FakeEmbedder::new("d1").behind(&gate);
let reserve = FakeEmbedder::new("d1");
let mut pool = EmbedderPool::new(vec![("raven".into(), a.clone()), ("starfire".into(), b.clone())], vec![("yggdrasil".into(), reserve.clone())]);
warm(&pool);
let batch = texts(8);
let opener = {
let gate = gate.clone();
@ -232,12 +363,14 @@ mod tests {
}
}
/// The reserve takes work only while no primary is in the pool, and gives
/// it back when a primary returns.
#[test]
fn the_reserve_serves_only_while_both_primaries_are_out() {
let (mut pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
down(&fakes[0], true);
warm(&pool);
pool.embed(&texts(4), WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (0, 4, 0), "one GPU left takes all of it");
down(&fakes[1], true);
@ -254,7 +387,7 @@ mod tests {
#[test]
fn a_failing_endpoint_drops_out_its_piece_moves_on_and_it_is_retried_later() {
let (mut pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
pool.model_identity(WITHIN).unwrap();
warm(&pool);
down(&fakes[0], true);
let batch = texts(6);
let vectors = pool.embed(&batch, WITHIN).unwrap();
@ -262,12 +395,31 @@ mod tests {
for (text, vector) in batch.iter().zip(&vectors) {
assert_eq!(vector[0], text.len() as f32);
}
let mut waiting = EmbedderPool::new(vec![("raven".into(), fakes[0].clone()), ("starfire".into(), fakes[1].clone())], vec![]);
down(&fakes[0], false);
waiting.embed(&texts(2), WITHIN).unwrap();
warm(&pool);
pool.embed(&texts(2), WITHIN).unwrap();
assert_eq!(used(&fakes[0]), 1, "back in the pool at the next admission");
}
/// With a retry time that has not come, a dropped endpoint is not read and
/// not used; once it has come, it is.
#[test]
fn a_dropped_endpoint_stays_out_until_its_retry_time() {
let (pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[]);
let mut pool = pool.retrying_after(Duration::from_millis(400));
warm(&pool);
down(&fakes[0], true);
pool.embed(&texts(2), WITHIN).unwrap();
down(&fakes[0], false);
let reads = fakes[0].state.lock().unwrap().identity_bounds.len();
pool.embed(&texts(2), WITHIN).unwrap();
assert_eq!((used(&fakes[0]), fakes[0].state.lock().unwrap().identity_bounds.len()), (0, reads), "not read, not used");
std::thread::sleep(Duration::from_millis(450));
warm(&pool);
pool.embed(&texts(2), WITHIN).unwrap();
assert_eq!(used(&fakes[0]), 1, "due, read and used");
}
/// An endpoint serving another digest or another length is never written
/// with, even when it is the only one answering; the log names it.
#[test]
@ -303,4 +455,74 @@ mod tests {
down(&cold_fakes[1], true);
assert!(format!("{:#}", cold.model_identity(WITHIN).unwrap_err()).contains("raven"));
}
/// An endpoint that never answers an identity read costs a batch nothing:
/// across several retry windows every batch finishes at the healthy
/// endpoint's pace, however long the caller's budget is, and the dark
/// endpoint is read under the pool's own short bound, never the caller's.
#[test]
fn a_dark_endpoint_never_delays_a_batch_a_healthy_one_can_take() {
let (pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d1")], &[("yggdrasil", "d1")]);
let mut pool = pool.retrying_after(Duration::from_millis(100));
fakes[1].state.lock().unwrap().hang_identity = true;
pool.built_under(Some(&identity("d1")));
let budget = Duration::from_secs(300);
for _ in 0..4 {
let started = Instant::now();
pool.embed(&texts(4), budget).unwrap();
assert!(started.elapsed() < Duration::from_millis(1500), "a batch waited {:?} on a dark endpoint", started.elapsed());
std::thread::sleep(Duration::from_millis(150));
}
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (16, 0, 0));
let bounds = fakes[1].state.lock().unwrap().identity_bounds.clone();
assert!(!bounds.is_empty() && bounds.iter().all(|bound| *bound == IDENTITY_WITHIN), "{bounds:?}");
}
/// The model of the pool is the collection's: an endpoint serving another
/// is not written with, however alone it is, and the identity is not given.
#[test]
fn the_pool_holds_endpoints_to_the_collections_model() {
let (mut pool, fakes) = pool(&[("raven", "d2"), ("starfire", "d1")], &[("yggdrasil", "d2")]);
pool.built_under(Some(&identity("d1")));
pool.embed(&texts(4), WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (0, 4, 0), "only the endpoint serving the collection's model works");
down(&fakes[1], true);
assert!(pool.embed(&texts(1), WITHIN).is_err(), "a different model is not a substitute");
assert!(pool.model_identity(WITHIN).is_err(), "and the pool does not claim the model of another");
assert_eq!((used(&fakes[0]), used(&fakes[2])), (0, 0));
down(&fakes[1], false);
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1");
}
/// With no collection to hold to, the model is the first in list order that
/// answers, whichever answers first; handing the pool a collection's model
/// changes it, and handing it the same one again reads nothing.
#[test]
fn the_pool_adopts_a_model_only_without_a_collection_and_a_new_pin_rereads() {
let (slow, quick) = (FakeEmbedder::new("d1"), FakeEmbedder::new("d2"));
slow.state.lock().unwrap().identity_delay = Duration::from_millis(200);
let mut pool = EmbedderPool::new(vec![("raven".into(), slow.clone()), ("starfire".into(), quick.clone())], vec![]);
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1", "list order, not the quickest");
pool.built_under(Some(&identity("d1")));
let reads = quick.state.lock().unwrap().identity_bounds.len();
pool.model_identity(WITHIN).unwrap();
assert_eq!(quick.state.lock().unwrap().identity_bounds.len(), reads, "the same model again reads nothing");
pool.built_under(Some(&identity("d2")));
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d2");
pool.embed(&texts(2), WITHIN).unwrap();
assert_eq!((used(&slow), used(&quick)), (0, 2));
pool.built_under(None);
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1", "no collection, the pool adopts again");
}
/// A primary that is still being read keeps the reserve waiting, so the
/// first batch goes to the GPU and not to the CPU that answered first.
#[test]
fn a_primary_still_being_read_is_waited_for_before_the_reserve_is_used() {
let (mut pool, fakes) = pool(&[("raven", "d1")], &[("yggdrasil", "d1")]);
fakes[0].state.lock().unwrap().identity_delay = Duration::from_millis(200);
pool.built_under(Some(&identity("d1")));
pool.embed(&texts(3), WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[1])), (3, 0));
}
}

15
crates/huginn-daemon/src/sources/mod.rs
View file

@ -29,7 +29,7 @@ use std::time::{Duration, Instant};
use anyhow::{Context, Result, bail, ensure};
use huginn_mind::{SourceHit, SourceHits, SourceModel, SourceQuery};
use crate::index::{BATCH, BULK_EMBED_TIMEOUT, Budget, Embedder, INDEX_CALL_TIMEOUT, MANAGED_BY, point_id};
use crate::index::{BATCH, BULK_EMBED_TIMEOUT, Budget, Embedder, INDEX_CALL_TIMEOUT, MANAGED_BY, ModelIdentity, point_id};
use chunk::{CHUNKER_VERSION, embedded_text};
use git::{BLOB_BYTES_MAX, Mirror, TreeFile};
@ -59,6 +59,13 @@ pub struct SourcesMeta {
pub chunker_version: u32,
}
impl SourcesMeta {
/// The model the collection was built with.
pub fn model_identity(&self) -> ModelIdentity {
ModelIdentity { name: self.model.clone(), digest: self.model_digest.clone(), dimensions: self.dimensions }
}
}
/// What the vector store holds under `COLLECTION`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum DescribedSources {
@ -187,6 +194,8 @@ pub fn refresh<E: Embedder, S: SourceIndex, F: Forge>(
state_root: &Path,
on_repo: &mut dyn FnMut(&Report),
) -> Result<Refreshed> {
let described = index.describe_sources(INDEX_CALL_TIMEOUT).context("describing the source collection")?;
embedder.built_under(if let DescribedSources::Labelled { meta, .. } = &described { Some(meta.model_identity()) } else { None }.as_ref());
let identity = embedder.model_identity(BULK_EMBED_TIMEOUT).context("reading the embedding model's identity")?;
let wanted = SourcesMeta {
managed_by: MANAGED_BY.into(),
@ -196,7 +205,7 @@ pub fn refresh<E: Embedder, S: SourceIndex, F: Forge>(
dimensions: identity.dimensions,
chunker_version: CHUNKER_VERSION,
};
let mut cursors = match index.describe_sources(INDEX_CALL_TIMEOUT).context("describing the source collection")? {
let mut cursors = match described {
DescribedSources::Absent => {
index.recreate_sources(&wanted).context("creating the source collection")?;
BTreeMap::new()
@ -368,12 +377,12 @@ impl Batch {
/// was built under another model, because vectors of two models do not compare.
pub fn search<E: Embedder, S: SourceIndex>(embedder: &mut E, index: &mut S, query: &SourceQuery, within: Duration) -> Result<SourceHits> {
let budget = Budget::new(within);
let identity = embedder.model_identity(budget.left()?).context("reading the embedding model's identity")?;
let built = match index.describe_sources(budget.left()?).context("describing the source collection")? {
DescribedSources::Absent => bail!("source index not built"),
DescribedSources::Unlabelled => bail!("source index not built: {COLLECTION} carries no Huginn metadata"),
DescribedSources::Labelled { meta, .. } => meta,
};
let identity = embedder.model_identity(budget.left()?).context("reading the embedding model's identity")?;
ensure!(
built.model == identity.name && built.model_digest == identity.digest && built.dimensions == identity.dimensions,
"the source index was built under {} {} ({} dimensions) and this daemon embeds under {} {} ({} dimensions)",