haematite_cli/serve/embed/
workers.rs

1//! Box 43 (ruled 5 Oct 09:17): a page's texts handed to several kept
2//! workers of the same build at once.
3//!
4//! The texts are cut into contiguous shares of near-equal size, one for each
5//! worker, or one for each text when there are fewer texts than workers.
6//! Each share goes to its own worker on its own scoped thread, and no thread
7//! outlives the call. The vectors come back in the order the texts came.
8//!
9//! The answer is whole or nothing. A share that fails fails the call, by its
10//! worker's code, with the worker named in the message; the shares that
11//! succeeded are thrown away, and every worker is left as it is. Every share
12//! must answer the same model.
13//!
14//! With one worker, or no texts, the call is the first worker's own, unchanged.
15
16use std::num::NonZeroUsize;
17use std::thread;
18
19use super::job::Embed;
20use super::{Embedded, Embedder, Settings};
21use crate::serve::storage::Failure;
22
23/// Several workers behind the job's trait, each its own `Embed`.
24pub struct Workers<E = Embedder> {
25    workers: Vec<E>,
26}
27
28impl Workers<Embedder> {
29    /// `count` embedders made from the same settings; none is started until
30    /// it is first given texts.
31    #[must_use]
32    pub fn new(settings: Option<&Settings>, count: NonZeroUsize) -> Self {
33        let workers = (0..count.get())
34            .map(|_worker| Embedder::new(settings.cloned()))
35            .collect();
36        Self { workers }
37    }
38}
39
40impl<E: Embed + Sync> Workers<E> {
41    /// Workers from `workers`, in the order their numbers are given.
42    ///
43    /// # Errors
44    /// `embed_workers_none` when there is none.
45    pub fn of(workers: Vec<E>) -> Result<Self, Failure> {
46        if workers.is_empty() {
47            return Err(Failure::new(
48                "embed_workers_none",
49                "at least one worker is needed",
50            ));
51        }
52        Ok(Self { workers })
53    }
54}
55
56impl<E: Embed + Sync> Embed for Workers<E> {
57    fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
58        let total = self.workers.len();
59        let Some(first) = self.workers.first() else {
60            return Err(Failure::new(
61                "embed_workers_none",
62                "at least one worker is needed",
63            ));
64        };
65        if texts.is_empty() || total == 1 {
66            return first.documents(texts);
67        }
68        let shares = total.min(texts.len());
69        let answers: Vec<Result<Embedded, Failure>> = thread::scope(|scope| {
70            let mut running = Vec::with_capacity(shares);
71            for (share, worker) in shares_of(texts, shares).into_iter().zip(&self.workers) {
72                running.push(scope.spawn(move || worker.documents(share)));
73            }
74            running
75                .into_iter()
76                .map(|handle| handle.join().unwrap_or_else(|_panic| Err(panicked())))
77                .collect()
78        });
79        let mut model = None;
80        let mut answered = Vec::with_capacity(texts.len());
81        for (index, answer) in answers.into_iter().enumerate() {
82            let made = answer.map_err(|failure| Failure {
83                message: format!("worker {} of {total}: {}", index + 1, failure.message),
84                ..failure
85            })?;
86            if model.as_ref().is_some_and(|known| *known != made.model) {
87                return Err(Failure::new(
88                    "embed_model_mixed",
89                    format!("worker {} of {total} answered another model", index + 1),
90                ));
91            }
92            if model.is_none() {
93                model = Some(made.model);
94            }
95            answered.extend(made.answers);
96        }
97        let model = model.ok_or_else(panicked)?;
98        Ok(Embedded {
99            model,
100            answers: answered,
101        })
102    }
103}
104
105/// `texts` cut into `shares` contiguous runs whose lengths differ by at most
106/// one, the longer first.
107fn shares_of(texts: &[String], shares: usize) -> Vec<&[String]> {
108    let base = texts.len() / shares;
109    let extra = texts.len() % shares;
110    let mut cut = Vec::with_capacity(shares);
111    let mut rest = texts;
112    for index in 0..shares {
113        let size = base + usize::from(index < extra);
114        let (share, after) = rest.split_at(size.min(rest.len()));
115        cut.push(share);
116        rest = after;
117    }
118    cut
119}
120
121fn panicked() -> Failure {
122    Failure::new(
123        "embed_worker_owner_failed",
124        "a share's thread panicked before it answered",
125    )
126}
127
128#[cfg(test)]
129#[path = "workers_tests.rs"]
130mod tests;
haematite_cli/serve/embed/workers.rs

haematite_cli/serve/embed/
workers.rs
