embedder pool in the shared index module (GPUs concurrent, CPU reserve, per-endpoint identity); huginn-sources uses it; the failover chain is deleted

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Meta Vens 2026-10-31 06:51:57 +01:00
commit bc8e071bf7
5 changed files with 314 additions and 161 deletions

View file

@ -8,8 +8,8 @@ use std::process::ExitCode;
use anyhow::{Context, Result, ensure};
use huginn_daemon::index::ollama::OllamaEmbedder;
use huginn_daemon::index::pool::EmbedderPool;
use huginn_daemon::index::qdrant::QdrantIndex;
use huginn_daemon::sources::failover::FailoverEmbedder;
use huginn_daemon::sources::forge::HttpForge;
use huginn_daemon::sources::refresh;
@ -20,8 +20,10 @@ const DEFAULT_FORGE_URL: &str = "http://127.0.0.1:3310";
struct Options {
state_root: PathBuf,
qdrant_url: String,
/// In order of preference; the first that serves the run's model embeds.
/// The GPU hosts: all take batches at once.
ollama_urls: Vec<String>,
/// The CPU host: takes batches only while no GPU host can.
reserve_urls: Vec<String>,
embedding_model: String,
forge_url: String,
}
@ -30,16 +32,17 @@ fn parse(args: impl Iterator<Item = String>) -> Result<Options> {
let mut args = args;
let mut values: BTreeMap<String, String> = BTreeMap::new();
let mut ollama_urls = Vec::new();
let mut reserve_urls = Vec::new();
while let Some(name) = args.next() {
let name = name.strip_prefix("--").with_context(|| format!("expected --option, got {name:?}"))?;
ensure!(
matches!(name, "state-root" | "qdrant-url" | "ollama-url" | "embedding-model" | "forge-url"),
matches!(name, "state-root" | "qdrant-url" | "ollama-url" | "reserve-ollama-url" | "embedding-model" | "forge-url"),
"unsupported huginn-sources option --{name}"
);
let value = args.next().with_context(|| format!("missing value for --{name}"))?;
if name == "ollama-url" {
ensure!(value.starts_with("http://"), "--ollama-url must be a plain http:// address, got {value:?}");
ollama_urls.push(value);
if name == "ollama-url" || name == "reserve-ollama-url" {
ensure!(value.starts_with("http://"), "--{name} must be a plain http:// address, got {value:?}");
if name == "ollama-url" { ollama_urls.push(value) } else { reserve_urls.push(value) }
} else {
ensure!(values.insert(name.to_owned(), value).is_none(), "duplicate option --{name}");
}
@ -57,9 +60,10 @@ fn parse(args: impl Iterator<Item = String>) -> Result<Options> {
state_root,
qdrant_url: url("qdrant-url", take("qdrant-url")?)?,
ollama_urls: {
ensure!(!ollama_urls.is_empty(), "at least one --ollama-url is required; repeat it in order of preference");
ensure!(!ollama_urls.is_empty(), "at least one --ollama-url is required; repeat it for each GPU host");
ollama_urls
},
reserve_urls,
embedding_model,
forge_url: url("forge-url", values.get("forge-url").cloned().unwrap_or_else(|| DEFAULT_FORGE_URL.to_owned()))?,
})
@ -73,9 +77,8 @@ fn main() -> ExitCode {
return ExitCode::from(2);
}
};
let mut embedder = FailoverEmbedder::new(
options.ollama_urls.iter().map(|url| (url.clone(), OllamaEmbedder::new(url, &options.embedding_model))).collect(),
);
let endpoint = |url: &String| (url.clone(), OllamaEmbedder::new(url, &options.embedding_model));
let mut embedder = EmbedderPool::new(options.ollama_urls.iter().map(endpoint).collect(), options.reserve_urls.iter().map(endpoint).collect());
let mut index = QdrantIndex::new(&options.qdrant_url);
let mut forge = HttpForge::new(&options.forge_url);
let mut journal = |report: &huginn_daemon::sources::Report| println!("{report}");

View file

@ -89,6 +89,7 @@
//! model, the collection's label) forgets it.
pub mod ollama;
pub mod pool;
pub mod qdrant;
#[cfg(test)]

View file

@ -0,0 +1,300 @@
//! A pool of embedder endpoints serving one model, as one `Embedder`.
//!
//! The endpoints are two tiers. The primaries (the GPU hosts) all take work at
//! once: a batch is cut into contiguous pieces, one per healthy primary, and
//! the pieces are embedded concurrently and joined back in order. The reserve
//! (the CPU host) takes work only while no primary is healthy. An endpoint is
//! admitted only after its model identity (name, digest, dimensions) equals the
//! pool's; one that is unreachable, failing or serving another model drops out,
//! is named in the log, is never written with, and is read again once
//! `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.
use std::time::{Duration, Instant};
use anyhow::{Result, bail};
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);
struct Member<E> {
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)>,
}
pub struct EmbedderPool<E: Embedder + Send> {
members: Vec<Member<E>>,
pinned: Option<ModelIdentity>,
retry_after: Duration,
}
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 }
}
pub fn retrying_after(mut self, retry_after: Duration) -> Self {
self.retry_after = retry_after;
self
}
/// 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_some_and(|(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));
}
/// 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 }
}
fn unavailable(&self) -> anyhow::Error {
let named: Vec<String> = self
.members
.iter()
.map(|member| format!("{}: {}", member.name, member.out.as_ref().map_or("failed this batch", |(_, why)| why.as_str())))
.collect();
anyhow::anyhow!("no embedder endpoint can take the batch: {}", named.join("; "))
}
}
impl<E: Embedder + Send> Embedder for EmbedderPool<E> {
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()),
}
}
fn embed(&mut self, texts: &[String], within: Duration) -> Result<Vec<Vec<f32>>> {
self.admit(within);
let mut vectors: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
// Pieces still owed, as (first index, length) over `texts`.
let mut owed = vec![(0, texts.len())];
while !owed.is_empty() {
let active = self.active();
if active.is_empty() {
bail!("{}", self.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(..) {
let share = length.div_ceil(active.len()).max(1);
for (n, start) in (first..first + length).step_by(share).enumerate() {
jobs.push((active[n % active.len()], start, share.min(first + length - start)));
}
}
let mut answers: Vec<(usize, usize, usize, Result<Vec<Vec<f32>>>)> = Vec::new();
std::thread::scope(|scope| {
let mut taken: Vec<Option<&mut Member<E>>> = self.members.iter_mut().map(Some).collect();
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 {
owed.push((first, length));
}
}
for (member, first, length, handle) in running {
let answer = handle.join().unwrap_or_else(|_| Err(anyhow::anyhow!("the embedder thread panicked")));
answers.push((member, first, length, answer));
}
});
for (member, first, length, answer) in answers {
match answer {
Ok(found) if found.len() == length => {
for (offset, vector) in found.into_iter().enumerate() {
vectors[first + offset] = Some(vector);
}
}
Ok(found) => {
self.drop_out(member, format!("returned {} vectors for {length} texts", found.len()));
owed.push((first, length));
}
Err(error) => {
self.drop_out(member, format!("failed an embed: {error:#}"));
owed.push((first, length));
}
}
}
}
Ok(vectors.into_iter().map(|vector| vector.expect("every piece was embedded")).collect())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::index::fakes::{FakeEmbedder, Gate};
const WITHIN: Duration = Duration::from_secs(1);
fn pool(primaries: &[(&str, &str)], reserve: &[(&str, &str)]) -> (EmbedderPool<FakeEmbedder>, Vec<FakeEmbedder>) {
let build = |spec: &[(&str, &str)]| -> Vec<(String, FakeEmbedder)> {
spec.iter().map(|(name, digest)| ((*name).to_owned(), FakeEmbedder::new(digest))).collect()
};
let (primaries, reserve) = (build(primaries), build(reserve));
let handles = primaries.iter().chain(&reserve).map(|(_, fake)| fake.clone()).collect();
(EmbedderPool::new(primaries, reserve).retrying_after(Duration::ZERO), handles)
}
fn texts(count: usize) -> Vec<String> {
(0..count).map(|n| "x".repeat(n + 1)).collect()
}
fn used(fake: &FakeEmbedder) -> usize {
fake.texts().len()
}
fn down(fake: &FakeEmbedder, is_down: bool) {
fake.state.lock().unwrap().down = is_down;
}
/// Both primaries take a batch at the same moment: each is inside its embed
/// before either is let go, and the vectors come back in the texts' order.
#[test]
fn both_primaries_embed_one_batch_concurrently_and_the_order_is_kept() {
let gate = Gate::shut();
let a = FakeEmbedder::new("d1").behind(&gate);
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())]);
let batch = texts(8);
let opener = {
let gate = gate.clone();
std::thread::spawn(move || {
gate.wait_arrived(2);
gate.open();
})
};
let vectors = pool.embed(&batch, WITHIN).unwrap();
opener.join().unwrap();
assert_eq!((used(&a), used(&b), used(&reserve)), (4, 4, 0), "an even share each, the reserve idle");
assert_eq!(vectors.len(), 8);
for (text, vector) in batch.iter().zip(&vectors) {
assert_eq!(vector[0], text.len() as f32, "the vector at each position is that text's");
}
}
/// 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);
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);
pool.embed(&texts(3), WITHIN).unwrap();
assert_eq!(used(&fakes[2]), 3, "both GPUs out, the CPU takes it");
down(&fakes[0], false);
pool.embed(&texts(2), WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[2])), (2, 3), "a GPU back in the pool, the CPU is left alone");
}
/// A piece that fails on one endpoint is embedded on another within the
/// same call, in order, and the failed endpoint is not asked again until
/// its retry time.
#[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();
down(&fakes[0], true);
let batch = texts(6);
let vectors = pool.embed(&batch, WITHIN).unwrap();
assert_eq!(used(&fakes[1]), 6, "starfire took raven's piece as well as its own");
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();
assert_eq!(used(&fakes[0]), 1, "back in the pool at the next admission");
}
/// 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]
fn a_mismatched_endpoint_is_never_written_with() {
let (mut pool, fakes) = pool(&[("raven", "d1"), ("starfire", "d2")], &[("yggdrasil", "d1")]);
pool.embed(&texts(4), WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (4, 0, 0));
down(&fakes[0], true);
pool.embed(&texts(2), WITHIN).unwrap();
assert_eq!((used(&fakes[1]), used(&fakes[2])), (0, 2), "starfire is skipped for the CPU, which serves the pool's model");
fakes[2].state.lock().unwrap().identity.dimensions = 8;
down(&fakes[2], false);
let error = format!("{:#}", pool.embed(&texts(1), WITHIN).unwrap_err());
assert!(error.contains("raven") && error.contains("starfire") && error.contains("yggdrasil"), "{error}");
assert_eq!(used(&fakes[1]), 0);
}
/// With every endpoint down, the identity read and the embed both fail
/// naming each; the pool's model is what the first endpoint reported.
#[test]
fn with_every_endpoint_down_the_pool_fails_naming_them_all() {
let (mut pool, fakes) = pool(&[("raven", "d1")], &[("yggdrasil", "d1")]);
assert_eq!(pool.model_identity(WITHIN).unwrap().digest, "d1");
for fake in &fakes {
down(fake, true);
}
let error = format!("{:#}", pool.embed(&texts(2), WITHIN).unwrap_err());
assert!(error.contains("raven") && error.contains("yggdrasil"), "{error}");
let (mut cold, cold_fakes) = self::pool(&[("raven", "d1")], &[("yggdrasil", "d1")]);
down(&cold_fakes[0], true);
down(&cold_fakes[1], true);
assert!(format!("{:#}", cold.model_identity(WITHIN).unwrap_err()).contains("raven"));
}
}

View file

@ -1,150 +0,0 @@
//! The writer's embedder: an ordered list of endpoints serving one model. The
//! first that answers with the run's model identity embeds; one that is
//! unreachable, failing, or serving another model (name, digest or length) is
//! skipped and named in the log, and is never written with. Daemons do not use
//! this: a query embeds on its daemon's own embedder.
use std::time::Duration;
use anyhow::{Result, bail};
use crate::index::{Embedder, ModelIdentity};
/// Endpoints in the order they are preferred, each with the name the log uses.
/// The first identity read of a run pins the model; every endpoint must then
/// report that identity before it is written with.
pub struct FailoverEmbedder<E: Embedder> {
endpoints: Vec<(String, E)>,
pinned: Option<ModelIdentity>,
}
impl<E: Embedder> FailoverEmbedder<E> {
pub fn new(endpoints: Vec<(String, E)>) -> Self {
Self { endpoints, pinned: None }
}
/// The first endpoint at or after `from` that answers with the pinned
/// identity. The others are named in the log and in the error if none does.
fn serving(&mut self, from: usize, within: Duration) -> Result<usize> {
let mut skipped = Vec::new();
for at in from..self.endpoints.len() {
let name = self.endpoints[at].0.clone();
match self.endpoints[at].1.model_identity(within) {
Ok(identity) => {
let pinned = self.pinned.get_or_insert_with(|| identity.clone());
if *pinned == identity {
return Ok(at);
}
eprintln!("huginn-sources: embedder {name} serves {identity:?}, not {pinned:?}; skipped");
skipped.push(format!("{name}: another model"));
}
Err(error) => {
eprintln!("huginn-sources: embedder {name} is unavailable: {error:#}; trying the next");
skipped.push(format!("{name}: {error:#}"));
}
}
}
bail!("no embedder endpoint can serve this run: {}", skipped.join("; "))
}
}
impl<E: Embedder> Embedder for FailoverEmbedder<E> {
fn model_identity(&mut self, within: Duration) -> Result<ModelIdentity> {
self.serving(0, within)?;
Ok(self.pinned.clone().expect("serving pins the model"))
}
fn embed(&mut self, texts: &[String], within: Duration) -> Result<Vec<Vec<f32>>> {
let mut failures = Vec::new();
let mut from = 0;
while from < self.endpoints.len() {
let at = match self.serving(from, within) {
Ok(at) => at,
Err(error) => bail!("no embedder endpoint embedded the batch: {}{error:#}", failures.join("; ")),
};
match self.endpoints[at].1.embed(texts, within) {
Ok(vectors) => return Ok(vectors),
Err(error) => {
let name = &self.endpoints[at].0;
eprintln!("huginn-sources: embedder {name} failed an embed: {error:#}; trying the next");
failures.push(format!("{name}: {error:#}; "));
from = at + 1;
}
}
}
bail!("no embedder endpoint embedded the batch: {}", failures.concat())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::index::fakes::FakeEmbedder;
const WITHIN: Duration = Duration::from_secs(1);
fn failover(spec: &[(&str, &str)]) -> (FailoverEmbedder<FakeEmbedder>, Vec<FakeEmbedder>) {
let fakes: Vec<FakeEmbedder> = spec.iter().map(|(_, digest)| FakeEmbedder::new(digest)).collect();
let endpoints = spec.iter().zip(&fakes).map(|((name, _), fake)| ((*name).to_owned(), fake.clone())).collect();
(FailoverEmbedder::new(endpoints), fakes)
}
fn used(fake: &FakeEmbedder) -> usize {
fake.texts().len()
}
/// The first healthy endpoint in order embeds; the others are not touched.
#[test]
fn the_first_healthy_endpoint_embeds() {
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("starfire", "d1"), ("yggdrasil", "d1")]);
embedder.embed(&["a".into()], WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (1, 0, 0));
}
/// A down endpoint hands the batch to the next, the CPU box embeds when
/// both desktops are down, and the preferred one is used again once it
/// answers.
#[test]
fn a_down_endpoint_hands_over_to_the_next_and_comes_back() {
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("starfire", "d1"), ("yggdrasil", "d1")]);
fakes[0].state.lock().unwrap().down = true;
embedder.embed(&["a".into()], WITHIN).unwrap();
assert_eq!((used(&fakes[0]), used(&fakes[1]), used(&fakes[2])), (0, 1, 0));
fakes[1].state.lock().unwrap().down = true;
embedder.embed(&["b".into()], WITHIN).unwrap();
assert_eq!(used(&fakes[2]), 1, "both desktops down, the CPU embeds");
fakes[0].state.lock().unwrap().down = false;
embedder.embed(&["c".into()], WITHIN).unwrap();
assert_eq!(used(&fakes[0]), 1, "the preferred endpoint is used again when it answers");
}
/// An endpoint serving another digest or another length is skipped and
/// never written with, even though it answers; with nothing left the run
/// fails naming each endpoint.
#[test]
fn a_mismatched_endpoint_is_skipped_and_never_embeds() {
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("starfire", "d2"), ("yggdrasil", "d1")]);
embedder.model_identity(WITHIN).unwrap();
fakes[0].state.lock().unwrap().down = true;
embedder.embed(&["a".into()], WITHIN).unwrap();
assert_eq!((used(&fakes[1]), used(&fakes[2])), (0, 1), "starfire serves d2 and is skipped");
fakes[2].state.lock().unwrap().identity.dimensions = 8;
let error = format!("{:#}", embedder.embed(&["b".into()], WITHIN).unwrap_err());
assert!(error.contains("raven") && error.contains("starfire: another model") && error.contains("yggdrasil: another model"), "{error}");
assert_eq!(used(&fakes[1]), 0);
}
/// With every endpoint down the identity read fails naming them all, and
/// the first one back up serves.
#[test]
fn with_every_endpoint_down_the_run_fails_naming_them_all() {
let (mut embedder, fakes) = failover(&[("raven", "d1"), ("yggdrasil", "d1")]);
for fake in &fakes {
fake.state.lock().unwrap().down = true;
}
let error = format!("{:#}", embedder.model_identity(WITHIN).unwrap_err());
assert!(error.contains("raven") && error.contains("yggdrasil"), "{error}");
fakes[1].state.lock().unwrap().down = false;
assert_eq!(embedder.model_identity(WITHIN).unwrap().digest, "d1");
}
}

View file

@ -16,7 +16,6 @@
//! The cursors object is written whole by this job alone, once per repo.
pub mod chunk;
pub mod failover;
pub mod forge;
pub mod git;