diff --git a/crates/huginn-daemon/src/index/fakes.rs b/crates/huginn-daemon/src/index/fakes.rs index ffb8570..47502b8 100644 --- a/crates/huginn-daemon/src/index/fakes.rs +++ b/crates/huginn-daemon/src/index/fakes.rs @@ -4,6 +4,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::{Arc, Condvar, Mutex}; +use std::time::Duration; use anyhow::{Result, bail, ensure}; @@ -60,6 +61,11 @@ pub(crate) struct EmbedderState { pub(crate) down: bool, /// The texts of every completed call, in order. pub(crate) embedded: Vec>, + /// The bound each completed call was given, in order. + pub(crate) bounds: Vec, + /// The length of the vectors it answers with, where that is not the + /// identity's. + pub(crate) vector_length: Option, } #[derive(Clone)] @@ -70,7 +76,7 @@ pub(crate) struct FakeEmbedder { impl FakeEmbedder { pub(crate) fn new(digest: &str) -> Self { - let state = EmbedderState { identity: identity(digest), down: false, embedded: Vec::new() }; + let state = EmbedderState { identity: identity(digest), down: false, embedded: Vec::new(), bounds: Vec::new(), vector_length: None }; Self { state: Arc::new(Mutex::new(state)), gate: None } } @@ -93,7 +99,7 @@ impl Embedder for FakeEmbedder { Ok(state.identity.clone()) } - fn embed(&mut self, texts: &[String]) -> Result>> { + fn embed(&mut self, texts: &[String], within: Duration) -> Result>> { if let Some(gate) = &self.gate { gate.pass(); } @@ -102,7 +108,8 @@ impl Embedder for FakeEmbedder { bail!("the embedder is down"); } state.embedded.push(texts.to_vec()); - let dimensions = state.identity.dimensions as usize; + state.bounds.push(within); + let dimensions = state.vector_length.unwrap_or(state.identity.dimensions as usize); Ok(texts.iter().map(|text| (0..dimensions).map(|at| (text.len() + at) as f32).collect()).collect()) } } diff --git a/crates/huginn-daemon/src/index/mod.rs b/crates/huginn-daemon/src/index/mod.rs index c0980c8..3f9102c 100644 --- a/crates/huginn-daemon/src/index/mod.rs +++ b/crates/huginn-daemon/src/index/mod.rs @@ -65,7 +65,8 @@ //! (`ServeOptions::search_timeout` bounds the wait), and takes back a search //! it has stopped waiting for. At most `SEARCHES_QUEUED_MAX` wait; one more is //! refused. A search whose embed is already running cannot be recalled, so it -//! can hold the writes off for as long as Ollama's own timeout. +//! can hold the writes off for as long as its bound, `QUERY_EMBED_TIMEOUT`, +//! which is `SEARCH_DEADLINE`, the deadline the serve loop gives the search. //! - Only a current index answers a search. While the worker is `Reconciling` //! or `Behind` (a re-pulled model, a lost collection, entries not yet //! written), `Failing` or `Refused`, a search is refused at once with an @@ -99,6 +100,22 @@ pub const MANAGED_BY: &str = "huginn"; /// Documents embedded and written per call. pub const BATCH: usize = 32; +/// How long a search may take, queue wait included: the serve loop's default +/// `ServeOptions::search_timeout`. It is the one number the query embed's +/// bound derives from, so the two cannot drift apart. +pub const SEARCH_DEADLINE: Duration = Duration::from_secs(30); + +/// The bound on one query embed. A running embed cannot be recalled, so it +/// holds the index's writes off for as long as it lasts; it is held to the +/// search's own deadline, after which nobody is waiting for its vector. +pub const QUERY_EMBED_TIMEOUT: Duration = SEARCH_DEADLINE; + +/// The bound on a bulk index-flush embed, and on the embedder's identity +/// reads. A cold model takes seconds to load. It does not block serving: the +/// serve loop never waits on the worker, and a search queued behind a flush is +/// answered by its own deadline. +pub const BULK_EMBED_TIMEOUT: Duration = Duration::from_secs(300); + /// The most candidates one search asks the index for. pub const OVERSAMPLE_MAX: u32 = 200; @@ -163,8 +180,8 @@ pub struct Hit { pub trait Embedder { fn model_identity(&mut self) -> Result; - /// One vector per text, in order. - fn embed(&mut self, texts: &[String]) -> Result>>; + /// One vector per text, in order, or an error once `within` has passed. + fn embed(&mut self, texts: &[String], within: Duration) -> Result>>; } pub trait VectorIndex { @@ -319,9 +336,6 @@ impl Projector { if self.wanted.is_empty() { return Ok(Advance::Idle); } - if !self.model_unchanged()? { - return Ok(Advance::Progressed); - } self.flush_batch()?; Ok(Advance::Progressed) } @@ -408,7 +422,9 @@ impl Projector { return Ok(()); } let texts: Vec = batch.iter().map(|entry| entry.text.clone()).collect(); - let vectors = self.embedder.embed(&texts).context("embedding")?; + // A model that changed leaves nothing written and the reconciliation + // forgotten: the next step reconciles and rebuilds. + let Ok(vectors) = self.embed_verified(&texts, BULK_EMBED_TIMEOUT)? else { return Ok(()) }; ensure!(vectors.len() == batch.len(), "the embedder returned {} vectors for {} texts", vectors.len(), batch.len()); let mut points = Vec::with_capacity(batch.len()); for (entry, vector) in batch.iter().zip(vectors) { @@ -437,6 +453,33 @@ impl Projector { Ok(()) } + /// The one check-then-embed: reads the model's identity and the + /// collection's label, embeds, reads the identity again, and gives the + /// vectors back only if the model is still the one the collection was + /// reconciled under. A model that changed forgets the reconciliation and + /// answers `Err(when)`; a collection that is no longer this mind's, built + /// as reconciled, is an error and forgets it too. Search and flush both + /// embed through here, so neither compares or writes vectors made by a + /// model the collection cannot mix. + fn embed_verified(&mut self, texts: &[String], within: Duration) -> Result>, &'static str>> { + let (Some(mind), Some(reconciled)) = (self.mind.clone(), self.reconciled.clone()) else { + bail!("the mind has no identity yet, so it has no index"); + }; + if !self.model_unchanged().map_err(|error| anyhow::anyhow!("{error}"))? { + return Ok(Err("since the collection was built")); + } + let described = self.index.describe(&self.collection).context("describing the collection")?; + if described != Described::Labelled(self.collection_meta(&mind, &reconciled)) { + self.reconciled = None; + bail!("collection {} is no longer the one this mind's index was built as: {described:?}", self.collection); + } + let vectors = self.embedder.embed(texts, within).context("embedding")?; + if !self.model_unchanged().map_err(|error| anyhow::anyhow!("{error}"))? { + return Ok(Err("while the texts were embedded")); + } + Ok(Ok(vectors)) + } + /// The candidates a query text finds: the text is embedded with the model's /// query instruction and the collection is searched, with no filter and an /// oversample of `min(top_k * 4, 200)`, because the mind will drop most of @@ -465,31 +508,17 @@ impl Projector { } fn nearest(&mut self, text: &str, top_k: u32) -> Result { - let (Some(mind), Some(reconciled)) = (self.mind.clone(), self.reconciled.clone()) else { + let Some(reconciled) = self.reconciled.clone() else { bail!("the mind has no identity yet, so it has no index"); }; let dimensions = reconciled.dimensions as usize; - ensure!( - self.model_unchanged().map_err(|error| anyhow::anyhow!("{error}"))?, - "the embedding model changed since the collection was built; the index is being rebuilt for it" - ); - let described = self.index.describe(&self.collection).context("describing the collection")?; - ensure!( - described == Described::Labelled(self.collection_meta(&mind, &reconciled)), - "collection {} is no longer the one this mind's index was built as: {described:?}", - self.collection - ); let query = format!("Instruct: {QUERY_INSTRUCTION}\nQuery: {text}"); - let mut vectors = self.embedder.embed(&[query]).context("embedding the query")?; + let mut vectors = match self.embed_verified(&[query], QUERY_EMBED_TIMEOUT)? { + Ok(vectors) => vectors, + Err(when) => bail!("the embedding model changed {when}; the index is being rebuilt for it"), + }; ensure!(vectors.len() == 1, "the embedder returned {} vectors for one query", vectors.len()); let vector = vectors.remove(0); - // The model is read again after the query was embedded: only a vector - // made between two reads of the same model is compared with the - // collection. - ensure!( - self.model_unchanged().map_err(|error| anyhow::anyhow!("{error}"))?, - "the embedding model changed while the query was embedded; the index is being rebuilt for it" - ); ensure!(vector.len() == dimensions, "the embedder returned a query vector of {} where the collection holds {dimensions}", vector.len()); let found = self.index.search(&self.collection, &vector, top_k.saturating_mul(4).min(OVERSAMPLE_MAX))?; let mut hits = Vec::with_capacity(found.len()); diff --git a/crates/huginn-daemon/src/index/ollama.rs b/crates/huginn-daemon/src/index/ollama.rs index 850b276..603bd08 100644 --- a/crates/huginn-daemon/src/index/ollama.rs +++ b/crates/huginn-daemon/src/index/ollama.rs @@ -6,7 +6,7 @@ use std::time::Duration; use anyhow::{Context, Result, bail, ensure}; use serde_json::{Value, json}; -use super::{Embedder, Http, Method, ModelIdentity}; +use super::{BULK_EMBED_TIMEOUT, Embedder, Http, Method, ModelIdentity}; /// The instruction Qwen3 embedding models take in front of a query (never in /// front of a document). Cut 11b's semantic read prepends it as @@ -15,20 +15,22 @@ pub const QUERY_INSTRUCTION: &str = "Given a question about a project's pipeline state, retrieve the documents that answer it"; /// One model on one Ollama. A cold model takes seconds to load, so the -/// timeout is generous; the worker owns the schedule around it. +/// identity reads get `BULK_EMBED_TIMEOUT`; the worker owns the schedule +/// around it. Each embed call carries its own bound. pub struct OllamaEmbedder { http: Http, + base_url: String, model: String, } impl OllamaEmbedder { pub fn new(base_url: &str, model: &str) -> Self { - Self { http: Http::new(base_url, Duration::from_secs(300)), model: model.to_owned() } + Self { http: Http::new(base_url, BULK_EMBED_TIMEOUT), base_url: base_url.to_owned(), model: model.to_owned() } } - fn ask(&self, method: Method, path: &str, body: Option<&Value>) -> Result { + fn ask(&self, http: &Http, method: Method, path: &str, body: Option<&Value>) -> Result { let text = body.map(Value::to_string); - let (status, answer) = self.http.request(method, path, text.as_deref())?; + let (status, answer) = http.request(method, path, text.as_deref())?; ensure!(status == 200, "Ollama answered {status} to {path}: {}", answer.chars().take(300).collect::()); serde_json::from_str(&answer).with_context(|| format!("Ollama's answer to {path} is not JSON")) } @@ -38,7 +40,7 @@ impl Embedder for OllamaEmbedder { /// The configured name, the digest Ollama lists for it, and the model's /// embedding length. The address is not part of the identity. fn model_identity(&mut self) -> Result { - let tags = self.ask(Method::Get, "/api/tags", None)?; + let tags = self.ask(&self.http, Method::Get, "/api/tags", None)?; let listed = tags["models"].as_array().context("Ollama's tag list has no models")?; let digest = listed .iter() @@ -46,7 +48,7 @@ impl Embedder for OllamaEmbedder { .and_then(|entry| entry["digest"].as_str()) .with_context(|| format!("Ollama does not list the model {}", self.model))? .to_owned(); - let shown = self.ask(Method::Post, "/api/show", Some(&json!({ "model": self.model })))?; + let shown = self.ask(&self.http, Method::Post, "/api/show", Some(&json!({ "model": self.model })))?; let info = shown["model_info"].as_object().context("Ollama's model description has no model_info")?; let Some(length) = info.iter().find(|(key, _)| key.ends_with(".embedding_length")).and_then(|(_, value)| value.as_u64()) else { @@ -56,8 +58,9 @@ impl Embedder for OllamaEmbedder { Ok(ModelIdentity { name: self.model.clone(), digest, dimensions }) } - fn embed(&mut self, texts: &[String]) -> Result>> { - let answer = self.ask(Method::Post, "/api/embed", Some(&json!({ "model": self.model, "input": texts })))?; + fn embed(&mut self, texts: &[String], within: Duration) -> Result>> { + let bounded = Http::new(&self.base_url, within); + let answer = self.ask(&bounded, Method::Post, "/api/embed", Some(&json!({ "model": self.model, "input": texts })))?; let embeddings = answer["embeddings"].as_array().context("Ollama's answer has no embeddings")?; embeddings .iter() diff --git a/crates/huginn-daemon/src/index/tests.rs b/crates/huginn-daemon/src/index/tests.rs index 2e7387c..c36699b 100644 --- a/crates/huginn-daemon/src/index/tests.rs +++ b/crates/huginn-daemon/src/index/tests.rs @@ -235,8 +235,8 @@ fn a_collection_owned_by_another_instance_is_never_recreated() { } } -/// A vector of the wrong length is refused before anything is written, and the -/// documents stay wanted. +/// The vector's length is checked against the collection's, so a flush that +/// the model did not change is still refused when the embedder answers wrongly. #[test] fn a_vector_of_the_wrong_length_is_refused_and_nothing_is_dropped() { let (_root, mind) = mind_with_documents(); @@ -245,13 +245,111 @@ fn a_vector_of_the_wrong_length_is_refused_and_nothing_is_dropped() { let mut projector = Projector::new(embedder.clone(), index.clone(), &slug(INSTANCE), Some(MIND.into())); projector.want(entries.iter().cloned()); projector.reconcile().unwrap(); - embedder.state.lock().unwrap().identity.dimensions = 8; + embedder.state.lock().unwrap().vector_length = Some(8); let error = projector.flush_batch().expect_err("a longer vector is refused"); assert!(matches!(error, StepError::Unavailable(_)), "{error}"); assert_eq!(projector.pending(), 4); assert!(index.state.lock().unwrap().upserts.is_empty()); } +/// A flush embeds through the same check as a search: a model that changed +/// before the embed is never asked, a model that changed during it has its +/// vectors dropped, and a collection that is no longer this mind's is refused. +/// In each case nothing is written, the documents stay wanted, and the +/// reconciliation is forgotten. +#[test] +fn a_flush_is_refused_as_a_search_is_when_the_model_or_the_collection_changed() { + let (_root, mind) = mind_with_documents(); + let entries = mind.index_entries(None).unwrap(); + let reconciled = || { + let (embedder, index) = pair(); + let mut projector = Projector::new(embedder.clone(), index.clone(), &slug(INSTANCE), Some(MIND.into())); + projector.want(entries.iter().cloned()); + projector.reconcile().unwrap(); + (embedder, index, projector) + }; + + let (embedder, index, mut projector) = reconciled(); + embedder.state.lock().unwrap().identity.digest = "d2".into(); + projector.flush_batch().expect("a changed model is not a failure"); + assert!(embedder.texts().is_empty(), "before: the texts were not embedded under the new model"); + assert!(index.state.lock().unwrap().upserts.is_empty()); + assert_eq!((projector.pending(), projector.stale()), (4, true), "before"); + + let (embedder, index, mut projector) = reconciled(); + change_label(&index, |meta| meta.mind = "mind-commit-another".into()); + let error = projector.flush_batch().expect_err("another mind's label is refused"); + assert!(error.to_string().contains("no longer the one this mind's index was built as"), "{error}"); + assert!(embedder.texts().is_empty(), "label: nothing was embedded"); + assert!(index.state.lock().unwrap().upserts.is_empty()); + assert_eq!((projector.pending(), projector.stale()), (4, true), "label"); + + let gate = Gate::shut(); + let embedder = FakeEmbedder::new("d1").behind(&gate); + let index = FakeIndex::default(); + let mut projector = Projector::new(embedder.clone(), index.clone(), &slug(INSTANCE), Some(MIND.into())); + projector.want(entries.iter().cloned()); + projector.reconcile().unwrap(); + let flushing = std::thread::spawn(move || { + let flushed = projector.flush_batch(); + (projector, flushed) + }); + gate.wait_arrived(1); + embedder.state.lock().unwrap().identity.digest = "d2".into(); + gate.open(); + let (projector, flushed) = flushing.join().unwrap(); + flushed.expect("a changed model is not a failure"); + assert!(index.state.lock().unwrap().upserts.is_empty(), "during: the vectors were not written"); + assert_eq!((projector.pending(), projector.stale()), (4, true), "during"); +} + +fn change_label(index: &FakeIndex, change: impl FnOnce(&mut CollectionMeta)) { + let mut state = index.state.lock().unwrap(); + let Described::Labelled(meta) = &mut state.collections.get_mut(&collection()).unwrap().label else { panic!() }; + change(meta); +} + +/// The query embed is bounded by the search's deadline, one constant behind +/// both, so a query embed that cannot be recalled cannot hold the writes off +/// past the moment nobody is waiting for it. A bulk flush keeps its own bound. +#[test] +fn the_query_embed_is_bounded_by_the_search_deadline_and_the_flush_by_its_own() { + let (_root, mind) = mind_with_documents(); + let entries = mind.index_entries(None).unwrap(); + let (embedder, index) = pair(); + let mut projector = projected(&embedder, &index, &entries); + assert!(!embedder.state.lock().unwrap().bounds.is_empty()); + assert!(embedder.state.lock().unwrap().bounds.iter().all(|bound| *bound == BULK_EMBED_TIMEOUT), "the flushes"); + projector.search("anything", 3).unwrap(); + let bounds = embedder.state.lock().unwrap().bounds.clone(); + assert_eq!(bounds.last(), Some(&QUERY_EMBED_TIMEOUT), "the query"); + let served = crate::serve::ServeOptions::default().search_timeout; + assert_eq!(served, SEARCH_DEADLINE); + assert!(QUERY_EMBED_TIMEOUT <= served); + assert!(QUERY_EMBED_TIMEOUT < BULK_EMBED_TIMEOUT); +} + +/// The Ollama adapter enforces the bound it is given: an embed that is never +/// answered fails at the bound, not at the adapter's default. +#[test] +fn an_ollama_embed_that_is_never_answered_fails_at_its_bound() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let held = std::thread::spawn(move || { + let mut held = Vec::new(); + for stream in listener.incoming().take(1) { + held.push(stream.unwrap()); + std::thread::sleep(Duration::from_secs(8)); + } + }); + let (sent, received) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let mut embedder = OllamaEmbedder::new(&url, "m:1"); + sent.send(embedder.embed(&["a".into()], Duration::from_millis(300)).is_err()).ok(); + }); + assert_eq!(received.recv_timeout(Duration::from_secs(5)), Ok(true), "the embed did not fail at its bound"); + drop(held); +} fn whoami_index(daemon: &mut Daemon) -> IndexStatus { match daemon.handle(HuginnMindRequest::Whoami, now()).answered() { HuginnMindResponse::Whoami(status) => status.index, @@ -446,7 +544,7 @@ fn ollama_identity_is_the_listed_digest_and_the_models_embedding_length() { let ollama = stub(&[("GET /api/tags", 200, TAGS), ("POST /api/show", 200, SHOW), ("POST /api/embed", 200, EMBED)]); let mut embedder = OllamaEmbedder::new(&ollama.url, "m:1"); assert_eq!(embedder.model_identity().unwrap(), ModelIdentity { name: "m:1".into(), digest: "ac6d".into(), dimensions: 4 }); - let vectors = embedder.embed(&["a".into(), "b".into(), "c".into(), "d".into()]).unwrap(); + let vectors = embedder.embed(&["a".into(), "b".into(), "c".into(), "d".into()], Duration::from_secs(5)).unwrap(); assert_eq!(vectors.len(), 4); assert_eq!(vectors[0], vec![0.5, 1.0, 2.0, 3.0]); let seen = ollama.seen.lock().unwrap(); @@ -464,7 +562,7 @@ fn ollama_identity_is_the_listed_digest_and_the_models_embedding_length() { let mut absent = OllamaEmbedder::new(&ollama.url, "not-pulled:1"); assert!(format!("{:#}", absent.model_identity().unwrap_err()).contains("does not list the model")); let refusing = stub(&[("POST /api/embed", 500, r#"{"error":"boom"}"#)]); - assert!(OllamaEmbedder::new(&refusing.url, "m:1").embed(&["a".into()]).is_err()); + assert!(OllamaEmbedder::new(&refusing.url, "m:1").embed(&["a".into()], Duration::from_secs(5)).is_err()); } /// Compatibility never reads the address: the same model at another address @@ -1188,7 +1286,7 @@ fn a_model_that_changes_while_the_query_is_embedded_is_refused() { gate.open(); let (projector, found) = searching.join().unwrap(); let refusal = format!("{:#}", found.unwrap_err()); - assert!(refusal.contains("while the query was embedded"), "{refusal}"); + assert!(refusal.contains("while the texts were embedded"), "{refusal}"); assert!(index.state.lock().unwrap().searches.is_empty(), "the vector was not used"); assert!(projector.stale()); } diff --git a/crates/huginn-daemon/src/serve.rs b/crates/huginn-daemon/src/serve.rs index 628abb9..5af8589 100644 --- a/crates/huginn-daemon/src/serve.rs +++ b/crates/huginn-daemon/src/serve.rs @@ -100,7 +100,7 @@ impl Default for ServeOptions { session_timeout: Duration::from_secs(30), idle_sleep: Duration::from_millis(2), deferred_ttl: Duration::from_secs(60), - search_timeout: Duration::from_secs(30), + search_timeout: crate::index::SEARCH_DEADLINE, } } }