diff --git a/crates/huginn-daemon/src/index/mod.rs b/crates/huginn-daemon/src/index/mod.rs index 2df23fa..db0e961 100644 --- a/crates/huginn-daemon/src/index/mod.rs +++ b/crates/huginn-daemon/src/index/mod.rs @@ -1142,8 +1142,18 @@ pub(crate) enum Method { impl Http { pub(crate) fn new(base_url: &str, timeout: Duration) -> Self { - let agent: ureq::Agent = - ureq::Agent::config_builder().timeout_global(Some(timeout)).http_status_as_error(false).build().into(); + Self::connecting_within(base_url, timeout, None) + } + + /// As `new`, with a bound on the connect alone, so a host that never + /// answers the connection is not waited on for the whole `timeout`. + pub(crate) fn connecting_within(base_url: &str, timeout: Duration, connect: Option) -> Self { + let agent: ureq::Agent = ureq::Agent::config_builder() + .timeout_global(Some(timeout)) + .timeout_connect(connect) + .http_status_as_error(false) + .build() + .into(); Self { agent, base: base_url.trim_end_matches('/').to_owned() } } diff --git a/crates/huginn-daemon/src/index/ollama.rs b/crates/huginn-daemon/src/index/ollama.rs index da7d356..1781aa2 100644 --- a/crates/huginn-daemon/src/index/ollama.rs +++ b/crates/huginn-daemon/src/index/ollama.rs @@ -14,15 +14,43 @@ use super::{Embedder, Http, Method, ModelIdentity}; pub const QUERY_INSTRUCTION: &str = "Given a question about a project's pipeline state, retrieve the documents that answer it"; +/// What an embed call may take, sized to the batch and never beyond the +/// caller's budget: a host that froze after accepting the batch is given up on +/// after `floor + per_text * texts`, not after the caller's whole `within`. +#[derive(Clone, Copy)] +pub struct EmbedBounds { + /// The bound on connecting alone. + pub connect: Duration, + /// What any batch is allowed, a model load included. + pub floor: Duration, + /// What each text adds. + pub per_text: Duration, +} + +impl EmbedBounds { + pub const DEFAULT: Self = + Self { connect: Duration::from_secs(3), floor: Duration::from_secs(20), per_text: Duration::from_millis(500) }; + + 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) + } +} + /// One model on one Ollama. Every call carries its own bound. pub struct OllamaEmbedder { base_url: String, model: String, + bounds: EmbedBounds, } impl OllamaEmbedder { pub fn new(base_url: &str, model: &str) -> Self { - Self { base_url: base_url.to_owned(), model: model.to_owned() } + Self { base_url: base_url.to_owned(), model: model.to_owned(), bounds: EmbedBounds::DEFAULT } + } + + pub fn bounded_by(mut self, bounds: EmbedBounds) -> Self { + self.bounds = bounds; + self } fn ask(&self, http: &Http, method: Method, path: &str, body: Option<&Value>) -> Result { @@ -57,7 +85,8 @@ impl Embedder for OllamaEmbedder { } fn embed(&mut self, texts: &[String], within: Duration) -> Result>> { - let bounded = Http::new(&self.base_url, within); + let allowed = self.bounds.for_batch(texts.len(), within); + let bounded = Http::connecting_within(&self.base_url, allowed, Some(self.bounds.connect.min(allowed))); 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 diff --git a/crates/huginn-daemon/src/index/pool.rs b/crates/huginn-daemon/src/index/pool.rs index de390cc..904425f 100644 --- a/crates/huginn-daemon/src/index/pool.rs +++ b/crates/huginn-daemon/src/index/pool.rs @@ -241,6 +241,9 @@ impl Embedder for EmbedderPool { } fn embed(&mut self, texts: &[String], within: Duration) -> Result>> { + // One budget for the batch: a piece moved off a failed endpoint is + // embedded under what is left of it, not under a fresh `within`. + let deadline = Instant::now() + within; self.admit(within); let mut vectors: Vec>> = vec![None; texts.len()]; // Pieces still owed, as (first index, length) over `texts`. @@ -250,12 +253,16 @@ impl Embedder for EmbedderPool { if active.is_empty() { // Every member that took the last round failed; the reserve may // still be on its first read. - self.admit(within); + self.admit(deadline.saturating_duration_since(Instant::now())); active = self.active(); } if active.is_empty() { bail!("{}", locked(&self.shared.roster).unavailable()); } + let left = deadline.saturating_duration_since(Instant::now()); + if left.is_zero() { + bail!("the batch's budget of {within:?} is spent; {}", 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) for (first, length) in owed.drain(..) { @@ -274,7 +281,7 @@ impl Embedder for EmbedderPool { 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)))); + running.push((member, first, length, scope.spawn(move || locked(embedder).embed(piece, left)))); } } for (member, first, length, handle) in running { @@ -560,4 +567,148 @@ mod tests { pool.embed(&texts(3), WITHIN).unwrap(); assert_eq!((used(&fakes[0]), used(&fakes[1])), (3, 0)); } + + // Real Ollama clients against sockets that go wrong after the pool has + // admitted them. + + use std::io::{Read, Write}; + use std::net::{TcpListener, TcpStream}; + use std::sync::atomic::{AtomicBool, Ordering}; + + use crate::index::ollama::{EmbedBounds, OllamaEmbedder}; + + const TAGS: &str = r#"{"models":[{"name":"m:1","model":"m:1","digest":"d1"}]}"#; + const SHOW: &str = r#"{"model_info":{"qwen3.embedding_length":1}}"#; + const QUICK: EmbedBounds = + EmbedBounds { connect: Duration::from_millis(200), floor: Duration::from_millis(300), per_text: Duration::ZERO }; + + /// An Ollama on a local port. `stalls` makes it take an embed request and + /// never answer; `gone` closes the listener, so the next connect is refused. + /// An embedding is `[text length]`, so order is checkable. + struct Endpoint { + url: String, + stalls: Arc, + gone: Arc, + listening: std::thread::JoinHandle<()>, + } + + fn endpoint() -> Endpoint { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let (stalls, gone) = (Arc::new(AtomicBool::new(false)), Arc::new(AtomicBool::new(false))); + let (stalling, leaving) = (Arc::clone(&stalls), Arc::clone(&gone)); + let listening = std::thread::spawn(move || { + while !leaving.load(Ordering::SeqCst) { + match listener.accept() { + Ok((stream, _)) => { + let stalling = Arc::clone(&stalling); + std::thread::spawn(move || answer(stream, &stalling)); + } + Err(_) => std::thread::sleep(Duration::from_millis(5)), + } + } + }); + Endpoint { url, stalls, gone, listening } + } + + fn answer(mut stream: TcpStream, stalls: &AtomicBool) { + stream.set_nonblocking(false).unwrap(); + 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") { + TAGS.to_owned() + } else if request.starts_with("POST /api/show") { + SHOW.to_owned() + } else { + if stalls.load(Ordering::SeqCst) { + std::thread::sleep(Duration::from_secs(30)); + return; + } + let asked: serde_json::Value = serde_json::from_slice(&buffer[head..head + length]).unwrap(); + let vectors: Vec> = asked["input"].as_array().unwrap().iter().map(|text| vec![text.as_str().unwrap().len()]).collect(); + serde_json::json!({ "embeddings": vectors }).to_string() + }; + let _ = write!(stream, "HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n{body}", body.len()); + } + + fn real_pool(endpoints: &[&Endpoint], bounds: EmbedBounds) -> EmbedderPool { + let members = endpoints.iter().map(|at| (at.url.clone(), OllamaEmbedder::new(&at.url, "m:1").bounded_by(bounds))).collect(); + let pool = EmbedderPool::new(members, Vec::new()).retrying_after(Duration::from_secs(3600)); + warm_real(&pool); + pool + } + + fn warm_real(pool: &EmbedderPool) { + pool.admit(Duration::from_secs(5)); + pool.settle(); + } + + fn lengths(vectors: &[Vec]) -> Vec { + vectors.iter().map(|vector| vector[0]).collect() + } + + /// An endpoint that takes the connection and then says nothing costs the + /// batch its own bound, not the caller's budget; its share goes to the + /// healthy endpoint, in order, and it stays out until its retry time. + #[test] + fn an_endpoint_that_stalls_mid_batch_costs_the_batch_its_bound_not_the_budget() { + let (healthy, stalling) = (endpoint(), endpoint()); + let mut pool = real_pool(&[&healthy, &stalling], QUICK); + stalling.stalls.store(true, Ordering::SeqCst); + let started = Instant::now(); + let vectors = pool.embed(&texts(8), Duration::from_secs(20)).unwrap(); + assert!(started.elapsed() < Duration::from_secs(3), "took {:?}", started.elapsed()); + assert_eq!(lengths(&vectors), (1..=8).map(|n| n as f32).collect::>()); + assert!(locked(&pool.shared.roster).members[1].out.is_some()); + let again = Instant::now(); + pool.embed(&texts(4), Duration::from_secs(20)).unwrap(); + assert!(again.elapsed() < Duration::from_millis(250), "the dropped endpoint was tried again: {:?}", again.elapsed()); + } + + /// An endpoint that stops listening after it was admitted: the same. + #[test] + fn an_endpoint_that_refuses_mid_batch_costs_the_batch_nothing_the_healthy_one_could_not_do() { + let (healthy, leaving) = (endpoint(), endpoint()); + let mut pool = real_pool(&[&healthy, &leaving], QUICK); + leaving.gone.store(true, Ordering::SeqCst); + leaving.listening.join().unwrap(); + let started = Instant::now(); + let vectors = pool.embed(&texts(6), Duration::from_secs(20)).unwrap(); + assert!(started.elapsed() < Duration::from_secs(3), "took {:?}", started.elapsed()); + assert_eq!(lengths(&vectors), (1..=6).map(|n| n as f32).collect::>()); + assert!(locked(&pool.shared.roster).members[1].out.is_some()); + } + + /// The bound never exceeds what the caller gave: with one endpoint, stalled + /// and a floor far above the budget, the batch fails at the budget. + #[test] + fn a_stall_is_never_waited_on_beyond_the_callers_budget() { + let only = endpoint(); + let generous = EmbedBounds { floor: Duration::from_secs(60), ..QUICK }; + let mut pool = real_pool(&[&only], generous); + only.stalls.store(true, Ordering::SeqCst); + let started = Instant::now(); + assert!(pool.embed(&texts(2), Duration::from_millis(400)).is_err()); + assert!(started.elapsed() < Duration::from_secs(3), "took {:?}", started.elapsed()); + } }