haematite_cli/serve/embed/
embedder.rs

1//! The service's one embedding worker, kept between requests behind one owner.
2
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Mutex, PoisonError};
5
6use haem_frames::embedding::{Distance, Kind, input_refusal};
7
8use super::Settings;
9use super::worker::{Batch, Ended, Ending, Worker};
10use crate::serve::storage::Failure;
11
12/// The model that made a set of vectors, as its worker's Ready named it.
13#[derive(Clone, Debug, PartialEq, Eq)]
14pub struct ModelId {
15    pub name: String,
16    pub revision: String,
17    pub manifest_sha256: String,
18    pub dimensions: u64,
19    pub distance: Distance,
20    pub normalised: bool,
21}
22
23/// What the worker made of one text.
24#[derive(Debug, Clone, PartialEq)]
25pub enum Answer {
26    Vector(Vec<f32>),
27    /// EMB1: the text has more tokens than the model reads. No vector is
28    /// made of it: a vector of part of a text is never kept.
29    Longer(Longer),
30    /// EMB-E2: a text longer than the model reads, embedded whole as
31    /// overlapping windows in text order.
32    Windows(Vec<WindowVector>),
33    /// EMB-E2: a text whose whole cover needs more windows than the cap. No
34    /// vector is made of it, as for [`Answer::Longer`].
35    OverCap(Longer),
36}
37
38/// EMB-E2: one window of a text, its bytes `[start, end)`, and its vector.
39#[derive(Debug, Clone, PartialEq)]
40pub struct WindowVector {
41    pub start: u64,
42    pub end: u64,
43    pub vector: Vec<f32>,
44}
45
46/// A text longer than the model reads: its content tokens, without special
47/// or prefix tokens, and the model's limit as its Ready declared it.
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub struct Longer {
50    pub tokens_total: u64,
51    pub tokens_limit: u64,
52}
53
54/// What the worker made of each text, in the order the texts were given,
55/// and the model that made them.
56#[derive(Debug)]
57pub struct Embedded {
58    pub model: ModelId,
59    pub answers: Vec<Answer>,
60}
61
62/// Starts one worker on first use and keeps it; a second caller waits for it.
63pub struct Embedder {
64    settings: Option<Settings>,
65    worker: Mutex<Option<Worker>>,
66    waiting: AtomicUsize,
67    /// The last worker ending, for the operator; taken only with the worker's
68    /// lock released.
69    last: Mutex<Option<Ending>>,
70}
71
72/// One caller counted among those waiting for the worker, until dropped.
73struct Waiting<'a>(&'a AtomicUsize);
74
75impl Drop for Waiting<'_> {
76    fn drop(&mut self) {
77        self.0.fetch_sub(1, Ordering::AcqRel);
78    }
79}
80
81impl Embedder {
82    /// An embedder with `settings`; `None` refuses every request by name.
83    /// Nothing is started until the first request.
84    #[must_use]
85    pub const fn new(settings: Option<Settings>) -> Self {
86        Self {
87            settings,
88            worker: Mutex::new(None),
89            waiting: AtomicUsize::new(0),
90            last: Mutex::new(None),
91        }
92    }
93
94    /// The vectors of `texts`, all of them in the order given, or a refusal.
95    ///
96    /// A kept worker found dead before any byte of a batch reached it is
97    /// replaced, once, and the batch goes to the new one: a worker that died
98    /// between requests costs no request. One that dies with a batch answers
99    /// by name, and that batch is never sent again.
100    ///
101    /// # Errors
102    ///
103    /// `embed_unavailable` with nothing configured; `embed_request_empty`,
104    /// `embed_request_too_large`, `check_embed`'s own name, or a text's own
105    /// refusal before anything is sent, with any kept worker left running;
106    /// `embed_busy` past the waiting bound; a start's refusal; any ending of
107    /// the worker, by name, after which the next request starts another.
108    pub fn embed(&self, kind: Kind, texts: &[String]) -> Result<Embedded, Failure> {
109        let settings = self.settings.as_ref().ok_or_else(|| {
110            Failure::new(
111                "embed_unavailable",
112                "no embedding worker is configured for this service",
113            )
114        })?;
115        asked(settings, texts)?;
116        let batches = texts
117            .chunks(settings.batch.get())
118            .enumerate()
119            .map(|(number, batch)| Batch::new(kind, batch, number * settings.batch.get(), settings))
120            .collect::<Result<Vec<_>, _>>()?;
121        let waiting = self.wait(settings)?;
122        let mut worker = self.worker.lock().map_err(|_| {
123            Failure::new(
124                "embed_worker_owner_failed",
125                "a request panicked while it held the worker",
126            )
127        })?;
128        drop(waiting);
129        let mut answers = Vec::with_capacity(texts.len());
130        let mut ending = None;
131        for batch in &batches {
132            match call(&mut worker, batch, settings, &mut ending) {
133                Ok(made) => answers.extend(made),
134                Err(ended) => {
135                    drop(worker);
136                    self.remember(ending);
137                    return Err(ended);
138                }
139            }
140        }
141        let Some(kept) = worker.as_ref() else {
142            return Err(not_started());
143        };
144        let model = &kept.ready.model;
145        let model = ModelId {
146            name: model.name.clone(),
147            revision: model.revision.clone(),
148            manifest_sha256: model.manifest_sha256.clone(),
149            dimensions: model.dimensions,
150            distance: model.distance,
151            normalised: model.normalised,
152        };
153        drop(worker);
154        self.remember(ending);
155        Ok(Embedded { model, answers })
156    }
157
158    /// The last way a worker ended: its name, exit status and the end of its
159    /// stderr, for an operator. No caller's answer carries the stderr.
160    pub fn last_ending(&self) -> Option<Ending> {
161        let last = self.last.lock().unwrap_or_else(PoisonError::into_inner);
162        let ending = last.clone();
163        drop(last);
164        ending
165    }
166
167    /// Keep `ending`, if there is one, as the last. Called only once the
168    /// worker's lock is released, so the two locks are never held together.
169    fn remember(&self, ending: Option<Ending>) {
170        let Some(ending) = ending else {
171            return;
172        };
173        let mut last = self.last.lock().unwrap_or_else(PoisonError::into_inner);
174        *last = Some(ending);
175        drop(last);
176    }
177
178    fn wait(&self, settings: &Settings) -> Result<Waiting<'_>, Failure> {
179        let before = self.waiting.fetch_add(1, Ordering::AcqRel);
180        let waiting = Waiting(&self.waiting);
181        if before >= settings.waiting {
182            return Err(Failure::new(
183                "embed_busy",
184                format!("{before} requests already wait for the embedding worker"),
185            ));
186        }
187        Ok(waiting)
188    }
189}
190
191/// A request's texts, refused by name before anything is sent.
192fn asked(settings: &Settings, texts: &[String]) -> Result<(), Failure> {
193    if texts.is_empty() {
194        return Err(Failure::new("embed_request_empty", "no texts were given"));
195    }
196    if texts.len() > settings.request_items.get() {
197        return Err(Failure::new(
198            "embed_request_too_large",
199            format!(
200                "{} texts were given; at most {} may be",
201                texts.len(),
202                settings.request_items
203            ),
204        ));
205    }
206    for (index, text) in texts.iter().enumerate() {
207        if let Some(code) = input_refusal(text, settings.limits.input_bytes) {
208            let code = serde_json::to_value(code).map_err(Failure::storage)?;
209            return Err(Failure::new(
210                code.as_str().unwrap_or("embed_input_refused"),
211                format!("text {index} cannot be embedded"),
212            ));
213        }
214    }
215    Ok(())
216}
217
218/// One batch through the kept worker, started if none is kept. A kept worker
219/// found dead before the batch reached it is replaced once; any other ending
220/// drops the worker, already reaped, and answers by name.
221fn call(
222    worker: &mut Option<Worker>,
223    batch: &Batch,
224    settings: &Settings,
225    ending: &mut Option<Ending>,
226) -> Result<Vec<Answer>, Failure> {
227    for attempt in 0..2 {
228        if worker.is_none() {
229            *worker = Some(Worker::start(settings)?);
230        }
231        let Some(kept) = worker.as_mut() else {
232            return Err(not_started());
233        };
234        match kept.embed(batch, settings) {
235            Ok(made) => return Ok(made),
236            // Dead before the batch reached it: nothing of it was sent, so
237            // the new worker may have it.
238            Err(Ended::Unsent(unsent)) if attempt == 0 => {
239                *worker = None;
240                *ending = Some(unsent.ending);
241            }
242            Err(ended) => {
243                *worker = None;
244                let lost = ended.lost();
245                *ending = Some(lost.ending);
246                return Err(lost.failure);
247            }
248        }
249    }
250    Err(not_started())
251}
252
253fn not_started() -> Failure {
254    Failure::new("embed_worker_not_ready", "the worker was not kept")
255}
haematite_cli/serve/embed/embedder.rs

haematite_cli/serve/embed/
embedder.rs
