diff --git a/crates/eureka-state/tests/cli.rs b/crates/eureka-state/tests/cli.rs index 4f17b46..3cbff32 100644 --- a/crates/eureka-state/tests/cli.rs +++ b/crates/eureka-state/tests/cli.rs @@ -1235,7 +1235,7 @@ impl huginn_daemon::sources::SourceIndex for Desk { anyhow::bail!("a daemon never writes the source index") } - fn delete_by(&mut self, _repo: &str, _paths: Option<&[String]>) -> anyhow::Result<()> { + fn delete_by(&mut self, _repo: &str, _paths: Option<&[String]>, _keeping: Option<&str>) -> anyhow::Result<()> { anyhow::bail!("a daemon never writes the source index") } diff --git a/crates/huginn-daemon/src/index/fakes.rs b/crates/huginn-daemon/src/index/fakes.rs index bf8833f..705090b 100644 --- a/crates/huginn-daemon/src/index/fakes.rs +++ b/crates/huginn-daemon/src/index/fakes.rs @@ -61,6 +61,8 @@ impl Gate { pub(crate) struct EmbedderState { pub(crate) identity: ModelIdentity, pub(crate) down: bool, + /// Whether embedding fails while the identity read still answers. + pub(crate) embed_fails: bool, /// The texts of every completed call, in order. pub(crate) embedded: Vec>, /// The bound each completed call was given, in order. @@ -85,7 +87,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(), bounds: Vec::new(), vector_length: None, identity_bounds: Vec::new(), hang_identity: false, identity_delay: Duration::ZERO }; + let state = EmbedderState { identity: identity(digest), down: false, embed_fails: false, embedded: Vec::new(), bounds: Vec::new(), vector_length: None, identity_bounds: Vec::new(), hang_identity: false, identity_delay: Duration::ZERO }; Self { state: Arc::new(Mutex::new(state)), gate: None } } @@ -122,7 +124,7 @@ impl Embedder for FakeEmbedder { gate.pass(); } let mut state = self.state.lock().unwrap(); - if state.down { + if state.down || state.embed_fails { bail!("the embedder is down"); } state.embedded.push(texts.to_vec()); @@ -327,14 +329,16 @@ impl SourceIndex for FakeSources { Ok(()) } - fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()> { + fn delete_by(&mut self, repo: &str, paths: Option<&[String]>, keeping: Option<&str>) -> Result<()> { let mut state = self.state.lock().unwrap(); if state.fail_deletes_of.as_deref() == Some(repo) { bail!("the vector store refused the delete"); } state.deletes.push((repo.to_owned(), paths.map(<[String]>::to_vec))); state.points.retain(|_, point| { - !(point.payload.repo == repo && paths.is_none_or(|paths| paths.contains(&point.payload.path))) + !(point.payload.repo == repo + && paths.is_none_or(|paths| paths.contains(&point.payload.path)) + && keeping.is_none_or(|commit| point.payload.commit != commit)) }); Ok(()) } @@ -376,8 +380,8 @@ impl SourceIndex for FakeIndex { fn set_cursors(&mut self, cursors: &BTreeMap) -> Result<()> { self.sources.set_cursors(cursors) } - fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()> { - self.sources.delete_by(repo, paths) + fn delete_by(&mut self, repo: &str, paths: Option<&[String]>, keeping: Option<&str>) -> Result<()> { + self.sources.delete_by(repo, paths, keeping) } fn upsert_sources(&mut self, points: &[SourcePoint]) -> Result<()> { self.sources.upsert_sources(points) diff --git a/crates/huginn-daemon/src/index/ollama.rs b/crates/huginn-daemon/src/index/ollama.rs index fc79872..afb624e 100644 --- a/crates/huginn-daemon/src/index/ollama.rs +++ b/crates/huginn-daemon/src/index/ollama.rs @@ -31,6 +31,12 @@ impl EmbedBounds { pub const DEFAULT: Self = Self { connect: Duration::from_secs(3), floor: Duration::from_secs(20), per_text: Duration::from_millis(500) }; + /// A CPU host. It embeds one text in 16 to 27 s idle and 52 s under load, + /// so a text is given 30 s and a batch 60 s to load; only the caller's + /// budget, which `for_batch` never exceeds, ends a slow reserve early. + pub const RESERVE: Self = + Self { connect: Duration::from_secs(3), floor: Duration::from_secs(60), per_text: Duration::from_secs(30) }; + fn for_batch(&self, texts: usize, within: Duration) -> Duration { (self.floor + self.per_text.saturating_mul(u32::try_from(texts).unwrap_or(u32::MAX))).min(within) } @@ -105,10 +111,24 @@ impl Embedder for OllamaEmbedder { /// One pool over Ollama endpoints serving `model`: the `primaries` (GPU hosts) /// share the work and the `reserve` (CPU hosts) serve only when none of them -/// can. The one embedder client of every Huginn process that embeds. +/// can, each tier under the bounds of its kind of host. The one embedder client +/// of every Huginn process that embeds. pub fn pool(primaries: &[String], reserve: &[String], model: &str) -> super::pool::EmbedderPool { - let endpoint = |url: &String| (url.clone(), OllamaEmbedder::new(url, model)); - super::pool::EmbedderPool::new(primaries.iter().map(endpoint).collect(), reserve.iter().map(endpoint).collect()) + pool_bounded(primaries, reserve, model, EmbedBounds::DEFAULT, EmbedBounds::RESERVE) +} + +fn pool_bounded( + primaries: &[String], + reserve: &[String], + model: &str, + primary: EmbedBounds, + reserve_bounds: EmbedBounds, +) -> super::pool::EmbedderPool { + let endpoint = |bounds: EmbedBounds| move |url: &String| (url.clone(), OllamaEmbedder::new(url, model).bounded_by(bounds)); + super::pool::EmbedderPool::new( + primaries.iter().map(endpoint(primary)).collect(), + reserve.iter().map(endpoint(reserve_bounds)).collect(), + ) } #[cfg(test)] @@ -121,4 +141,81 @@ mod tests { assert_eq!(bounds.for_batch(3, Duration::from_secs(100)), Duration::from_secs(16)); assert_eq!(bounds.for_batch(3, Duration::from_secs(12)), Duration::from_secs(12)); } + + /// A CPU host that takes 52 s for one text under load, or 16 to 27 s idle, + /// must not be given up on at a GPU's 20.5 s, and a full bulk batch of it is + /// bounded only by the caller's budget; a GPU's bound stays short. + #[test] + fn the_reserve_is_given_what_a_cpu_takes_and_a_gpu_is_not() { + let bulk = Duration::from_secs(300); + assert!(EmbedBounds::RESERVE.for_batch(1, bulk) >= Duration::from_secs(60)); + assert_eq!(EmbedBounds::RESERVE.for_batch(32, bulk), bulk); + assert!(EmbedBounds::DEFAULT.for_batch(1, bulk) < Duration::from_secs(25)); + } + + /// An Ollama that answers an embed only after `delay`. + fn slow_ollama(delay: Duration) -> String { + use std::io::{Read, Write}; + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + std::thread::spawn(move || { + for stream in listener.incoming().flatten() { + std::thread::spawn(move || { + let mut stream = stream; + let mut buffer = Vec::new(); + let (head, length) = loop { + let mut chunk = [0_u8; 4096]; + match stream.read(&mut chunk) { + Ok(0) | Err(_) => return, + Ok(read) => buffer.extend_from_slice(&chunk[..read]), + } + let text = String::from_utf8_lossy(&buffer).into_owned(); + if let Some(end) = text.find("\r\n\r\n") { + let length = text.lines().find_map(|line| line.to_ascii_lowercase().strip_prefix("content-length:").map(|v| v.trim().parse().unwrap())).unwrap_or(0); + break (end + 4, length); + } + }; + while buffer.len() < head + length { + let mut chunk = [0_u8; 4096]; + match stream.read(&mut chunk) { + Ok(0) | Err(_) => return, + Ok(read) => buffer.extend_from_slice(&chunk[..read]), + } + } + let request = String::from_utf8_lossy(&buffer[..head]).into_owned(); + let body = if request.starts_with("GET /api/tags") { + r#"{"models":[{"name":"m:1","model":"m:1","digest":"d1"}]}"# + } else if request.starts_with("POST /api/show") { + r#"{"model_info":{"qwen3.embedding_length":1}}"# + } else { + std::thread::sleep(delay); + r#"{"embeddings":[[1.0]]}"# + }; + let _ = write!(stream, "HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n{body}", body.len()); + }); + } + }); + url + } + + /// With every GPU dark, the CPU reserve answers a text slower than a GPU's + /// bound allows and the pool takes the answer; under the GPU's bound the + /// same host is given up on. The tier decides the bound. + #[test] + fn a_slow_reserve_answers_under_its_own_bound_where_a_gpus_would_have_failed() { + let dark = { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + [format!("http://{}", listener.local_addr().unwrap())] + }; + let cpu = [slow_ollama(Duration::from_millis(900))]; + let quick = EmbedBounds { connect: Duration::from_secs(1), floor: Duration::from_millis(300), per_text: Duration::ZERO }; + let patient = EmbedBounds { floor: Duration::from_secs(5), ..quick }; + + let mut as_reserve = pool_bounded(&dark, &cpu, "m:1", quick, patient); + let vectors = as_reserve.embed(&["a".to_owned()], Duration::from_secs(20)).unwrap(); + assert_eq!(vectors, vec![vec![1.0_f32]]); + + let mut as_a_gpu = pool_bounded(&dark, &cpu, "m:1", quick, quick); + assert!(as_a_gpu.embed(&["a".to_owned()], Duration::from_secs(20)).is_err()); + } } diff --git a/crates/huginn-daemon/src/index/qdrant.rs b/crates/huginn-daemon/src/index/qdrant.rs index b7c1d9a..bc8c6dc 100644 --- a/crates/huginn-daemon/src/index/qdrant.rs +++ b/crates/huginn-daemon/src/index/qdrant.rs @@ -273,13 +273,17 @@ impl SourceIndex for QdrantIndex { Ok(()) } - fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()> { + fn delete_by(&mut self, repo: &str, paths: Option<&[String]>, keeping: Option<&str>) -> Result<()> { let collection = self.sources.clone(); let mut must = vec![json!({ "key": "repo", "match": { "value": repo } })]; if let Some(paths) = paths { must.push(json!({ "key": "path", "match": { "any": paths } })); } - let body = json!({ "filter": { "must": must } }); + let mut filter = json!({ "must": must }); + if let Some(commit) = keeping { + filter["must_not"] = json!([{ "key": "commit", "match": { "value": commit } }]); + } + let body = json!({ "filter": filter }); self.ask(Method::Post, &format!("/collections/{collection}/points/delete?wait=true"), Some(&body), &[200])?; Ok(()) } diff --git a/crates/huginn-daemon/src/index/tests.rs b/crates/huginn-daemon/src/index/tests.rs index 2382a98..16af983 100644 --- a/crates/huginn-daemon/src/index/tests.rs +++ b/crates/huginn-daemon/src/index/tests.rs @@ -1079,6 +1079,29 @@ fn a_search_refuses_when_the_index_is_not_reconciled_with_the_model() { assert!(projector.stale(), "and the reconciliation is to be redone"); } +/// The source delete names the repo, the paths when given, and, when it keeps a +/// commit, excludes that commit's points: the filter is what spares the chunks +/// just written. +#[test] +fn qdrant_source_delete_filters_by_repo_paths_and_the_commit_it_keeps() { + use crate::sources::SourceIndex; + let deleting = stub(&[("POST /collections/", 200, r#"{"result":{"status":"completed"},"status":"ok"}"#)]); + let mut store = QdrantIndex::new(&deleting.url); + store.delete_by("Demo", Some(&["a.rs".to_owned()]), Some("abc123")).unwrap(); + store.delete_by("Gone", None, None).unwrap(); + let seen = deleting.seen.lock().unwrap(); + let bodies: Vec = seen.iter().map(|(_, _, body)| serde_json::from_str(body).unwrap()).collect(); + assert!(seen.iter().all(|(method, path, _)| method == "POST" && path.ends_with("/points/delete?wait=true"))); + assert_eq!( + bodies[0]["filter"], + serde_json::json!({ + "must": [{ "key": "repo", "match": { "value": "Demo" } }, { "key": "path", "match": { "any": ["a.rs"] } }], + "must_not": [{ "key": "commit", "match": { "value": "abc123" } }] + }) + ); + assert_eq!(bodies[1]["filter"], serde_json::json!({ "must": [{ "key": "repo", "match": { "value": "Gone" } }] })); +} + /// The search scores exactly the points it is handed (a has_id filter, exact /// search; which documents are members is the mind's), and reads back the /// document id, the kind and the score, in the order Qdrant gave them. A hit diff --git a/crates/huginn-daemon/src/sources/mod.rs b/crates/huginn-daemon/src/sources/mod.rs index 0a45872..ea88095 100644 --- a/crates/huginn-daemon/src/sources/mod.rs +++ b/crates/huginn-daemon/src/sources/mod.rs @@ -115,9 +115,10 @@ pub trait SourceIndex { fn recreate_sources(&mut self, meta: &SourcesMeta) -> Result<()>; /// Replaces the cursors object whole. fn set_cursors(&mut self, cursors: &BTreeMap) -> Result<()>; - /// Deletes a repo's chunks: those at `paths`, or all of them when `None`. - /// Acknowledged before it returns. - fn delete_by(&mut self, repo: &str, paths: Option<&[String]>) -> Result<()>; + /// Deletes a repo's chunks: those at `paths`, or all of them when `None`; + /// with `keeping`, only those not stamped with that commit. Acknowledged + /// before it returns. + fn delete_by(&mut self, repo: &str, paths: Option<&[String]>, keeping: Option<&str>) -> Result<()>; /// Writes chunks, acknowledged before it returns. fn upsert_sources(&mut self, points: &[SourcePoint]) -> Result<()>; /// The `limit` nearest chunks, best first, in `repo` alone when given. @@ -309,7 +310,7 @@ fn retire( } fn retire_one(index: &mut S, name: &str, mirrors: &Path, cursors: &mut BTreeMap) -> Result<()> { - index.delete_by(name, None).with_context(|| format!("deleting the chunks of retired {name}"))?; + index.delete_by(name, None, None).with_context(|| format!("deleting the chunks of retired {name}"))?; if cursors.remove(name).is_some() { index.set_cursors(cursors).context("dropping a retired repo's cursor")?; } @@ -342,23 +343,16 @@ fn refresh_repo( Some(old) if mirror.is_ancestor(old, &head) => Some(mirror.changed_paths(old, &head)?.into_iter().collect()), _ => None, }; - // From the first delete on, the collection holds no single commit's chunks + // From the first upsert on, the collection holds no single commit's chunks // for this repo until the run finishes. Say so before changing anything: // the cursor is dropped, so a run that dies half way leaves a repo the next // run re-ingests whole, whatever head it finds then. A cursor left at the // old commit would let a later head that reverts the files hide the hole. + // Nothing is deleted until the new chunks are written (below), so a failed + // embed leaves what the index held. if cursors.remove(&repo.name).is_some() { index.set_cursors(cursors).context("dropping the cursor before the repo's chunks change")?; } - match &changed { - Some(paths) => { - let paths: Vec = paths.iter().cloned().collect(); - for slice in paths.chunks(DELETE_PATHS_PER_CALL) { - index.delete_by(&repo.name, Some(slice)).context("deleting the changed paths' chunks")?; - } - } - None => index.delete_by(&repo.name, None).context("deleting the repo's chunks")?, - } let files: Vec = mirror .files(&head)? .into_iter() @@ -398,6 +392,18 @@ fn refresh_repo( } } batch.write(embedder, index, *dimensions)?; + // Every chunk at this head now carries it (a chunk id is its content, so a + // rewritten one replaced its old self). What is left in scope with another + // commit is stale: edited or deleted paths' chunks, or a repo's old self. + match &changed { + Some(paths) => { + let paths: Vec = paths.iter().cloned().collect(); + for slice in paths.chunks(DELETE_PATHS_PER_CALL) { + index.delete_by(&repo.name, Some(slice), Some(&head)).context("deleting the changed paths' stale chunks")?; + } + } + None => index.delete_by(&repo.name, None, Some(&head)).context("deleting the repo's stale chunks")?, + } cursors.insert(repo.name.clone(), head); index.set_cursors(cursors).context("advancing the cursor")?; Ok(report) diff --git a/crates/huginn-daemon/src/sources/tests.rs b/crates/huginn-daemon/src/sources/tests.rs index 698f828..7e62905 100644 --- a/crates/huginn-daemon/src/sources/tests.rs +++ b/crates/huginn-daemon/src/sources/tests.rs @@ -217,10 +217,10 @@ fn cursor_never_leads() { assert_eq!(rig.index.ids(), clean.index.ids()); } -/// A run that fails between a repo's deletes and its upserts leaves chunks of -/// no single commit. A head that then reverts the files to the old content -/// makes an empty diff from the old commit, which would keep the deleted -/// chunks deleted; the repo is ingested whole instead. +/// A run that fails part way through a repo's upserts leaves chunks of two +/// commits. A head that then reverts the files to the old content makes an +/// empty diff from the old commit, which would keep the half-written new chunks; +/// the repo is ingested whole instead and only the head's chunks remain. #[test] fn a_failed_run_then_a_revert_converges() { let remote = Remote::new(); @@ -242,7 +242,8 @@ fn a_failed_run_then_a_revert_converges() { } let (refreshed, _) = rig.run(&mut forge); assert_eq!(refreshed.unwrap().failed, ["Demo"]); - assert!(rig.index.held().len() < 40, "the failed run left a partial repo"); + let commits: BTreeSet = rig.index.held().into_iter().map(|(.., commit)| commit).collect(); + assert_eq!(commits.len(), 2, "the failed run left chunks of two commits"); rig.index.state.lock().unwrap().fail_upserts_after = None; let c = remote.commit(&original, &[]); @@ -255,6 +256,41 @@ fn a_failed_run_then_a_revert_converges() { assert_eq!(rig.index.ids(), clean.index.ids()); } +/// While embedding fails and the identity read still answers (every GPU dark, +/// the CPU too slow), a pushed repo keeps every chunk it held, run after run, +/// and the first run after the outage converges to a clean ingest. +#[test] +fn an_embed_outage_deletes_nothing_the_index_held() { + let remote = Remote::new(); + remote.commit(&[("a.rs", b"one +"), ("b.rs", b"two +"), ("c.rs", b"three +")], &[]); + let mut rig = Rig::new(); + let mut forge = FakeForge::of(&[("Demo", &remote)]); + rig.converged(&mut forge); + let before = rig.index.held(); + let ids_before = rig.index.ids(); + + let b = remote.commit(&[("a.rs", b"one edited +")], &["b.rs"]); + rig.embedder.state.lock().unwrap().embed_fails = true; + for _ in 0..2 { + let (refreshed, _) = rig.run(&mut forge); + assert_eq!(refreshed.unwrap().failed, ["Demo"]); + assert_eq!(rig.index.held(), before, "a failed embed left the index as it was"); + assert_eq!(rig.index.ids(), ids_before); + } + + rig.embedder.state.lock().unwrap().embed_fails = false; + rig.converged(&mut forge); + assert_eq!(rig.index.cursors().get("Demo"), Some(&b)); + let mut clean = Rig::new(); + clean.converged(&mut forge); + assert_eq!(rig.index.held(), clean.index.held()); + assert_eq!(rig.index.ids(), clean.index.ids()); +} + /// One repo failing is logged and the others go on. #[test] fn one_repos_failure_does_not_stop_the_others() {