haematite-cli 0.21.1

//! Box 43 (ruled 5 Oct 09:17): a page's texts handed to several kept
//! workers of the same build at once.
//!
//! The texts are cut into contiguous shares of near-equal size, one for each
//! worker, or one for each text when there are fewer texts than workers.
//! Each share goes to its own worker on its own scoped thread, and no thread
//! outlives the call. The vectors come back in the order the texts came.
//!
//! The answer is whole or nothing. A share that fails fails the call, by its
//! worker's code, with the worker named in the message; the shares that
//! succeeded are thrown away, and every worker is left as it is. Every share
//! must answer the same model.
//!
//! With one worker, or no texts, the call is the first worker's own, unchanged.

use std::num::NonZeroUsize;
use std::thread;

use super::job::Embed;
use super::{Embedded, Embedder, Settings};
use crate::serve::storage::Failure;

/// Several workers behind the job's trait, each its own `Embed`.
pub struct Workers<E = Embedder> {
    workers: Vec<E>,
}

impl Workers<Embedder> {
    /// `count` embedders made from the same settings; none is started until
    /// it is first given texts.
    #[must_use]
    pub fn new(settings: Option<&Settings>, count: NonZeroUsize) -> Self {
        let workers = (0..count.get())
            .map(|_worker| Embedder::new(settings.cloned()))
            .collect();
        Self { workers }
    }
}

impl<E: Embed + Sync> Workers<E> {
    /// Workers from `workers`, in the order their numbers are given.
    ///
    /// # Errors
    /// `embed_workers_none` when there is none.
    pub fn of(workers: Vec<E>) -> Result<Self, Failure> {
        if workers.is_empty() {
            return Err(Failure::new(
                "embed_workers_none",
                "at least one worker is needed",
            ));
        }
        Ok(Self { workers })
    }
}

impl<E: Embed + Sync> Embed for Workers<E> {
    fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
        let total = self.workers.len();
        let Some(first) = self.workers.first() else {
            return Err(Failure::new(
                "embed_workers_none",
                "at least one worker is needed",
            ));
        };
        if texts.is_empty() || total == 1 {
            return first.documents(texts);
        }
        let shares = total.min(texts.len());
        let answers: Vec<Result<Embedded, Failure>> = thread::scope(|scope| {
            let mut running = Vec::with_capacity(shares);
            for (share, worker) in shares_of(texts, shares).into_iter().zip(&self.workers) {
                running.push(scope.spawn(move || worker.documents(share)));
            }
            running
                .into_iter()
                .map(|handle| handle.join().unwrap_or_else(|_panic| Err(panicked())))
                .collect()
        });
        let mut model = None;
        let mut answered = Vec::with_capacity(texts.len());
        for (index, answer) in answers.into_iter().enumerate() {
            let made = answer.map_err(|failure| Failure {
                message: format!("worker {} of {total}: {}", index + 1, failure.message),
                ..failure
            })?;
            if model.as_ref().is_some_and(|known| *known != made.model) {
                return Err(Failure::new(
                    "embed_model_mixed",
                    format!("worker {} of {total} answered another model", index + 1),
                ));
            }
            if model.is_none() {
                model = Some(made.model);
            }
            answered.extend(made.answers);
        }
        let model = model.ok_or_else(panicked)?;
        Ok(Embedded {
            model,
            answers: answered,
        })
    }
}

/// `texts` cut into `shares` contiguous runs whose lengths differ by at most
/// one, the longer first.
fn shares_of(texts: &[String], shares: usize) -> Vec<&[String]> {
    let base = texts.len() / shares;
    let extra = texts.len() % shares;
    let mut cut = Vec::with_capacity(shares);
    let mut rest = texts;
    for index in 0..shares {
        let size = base + usize::from(index < extra);
        let (share, after) = rest.split_at(size.min(rest.len()));
        cut.push(share);
        rest = after;
    }
    cut
}

fn panicked() -> Failure {
    Failure::new(
        "embed_worker_owner_failed",
        "a share's thread panicked before it answered",
    )
}

#[cfg(test)]
#[path = "workers_tests.rs"]
mod tests;