haematite_cli/serve/embed/
worker.rs

1//! One embedding worker process: started, checked, and asked one batch at a
2//! time. Dropping it kills and reaps the process.
3//!
4//! Every ending after Ready is named (box 42): the owner is finished and its
5//! outcome read, so the answer carries the status and the end of the
6//! worker's stderr, and a call past its deadline is a timeout whatever the
7//! read then saw. A batch is encoded before anything is sent, and the bytes
8//! the worker's input accepted are counted, so a worker found dead before a
9//! single byte reached it is told apart from one that died with the batch.
10
11use std::io::{self, Write};
12use std::num::NonZeroU32;
13use std::process::{ChildStdin, ChildStdout, Command, Stdio};
14use std::time::Instant;
15
16use haem_frames::embedding::{
17    Embed, Input, ItemCode, Kind, MINIMUM_FRAME_BYTES, Message, Outcome, PROTOCOL,
18    PROTOCOL_WINDOWS, Ready, Start, Window, check_embed, check_model_limits_windowed,
19    check_vectors_windowed, decode_message, decode_vector, read_message, write_message,
20};
21use haem_frames::frame::{self, Fault};
22use haematite_query::worker_owner::{Control, OwnedChild};
23
24use super::Settings;
25use super::embedder::{Answer, Longer, WindowVector};
26use crate::serve::storage::Failure;
27
28/// A started worker that sent a Ready the service accepts.
29pub struct Worker {
30    // Dropped first: closing the worker's input before its owner kills it.
31    input: ChildStdin,
32    output: ChildStdout,
33    control: Control,
34    pub ready: Ready,
35    frame: NonZeroU32,
36}
37
38fn failure(code: &str, message: impl Into<String>) -> Failure {
39    Failure::new(code, message.into())
40}
41
42/// The most of the worker's stderr an answer carries, from its end.
43const TAIL_BYTES: usize = 512;
44
45/// One batch, checked and encoded before anything is sent.
46pub struct Batch {
47    embed: Embed,
48    frame: Vec<u8>,
49}
50
51/// The frame limit both sides use.
52fn frame_limit(settings: &Settings) -> Result<NonZeroU32, Failure> {
53    NonZeroU32::new(settings.limits.frame_bytes.max(MINIMUM_FRAME_BYTES))
54        .ok_or_else(|| failure("embed_worker_not_ready", "no frame bytes"))
55}
56
57impl Batch {
58    /// `texts` as one Embed whose item ids are their places in the whole
59    /// request, from `first`.
60    ///
61    /// # Errors
62    ///
63    /// `check_embed`'s own name for an Embed it refuses, and
64    /// `embed_request_too_large` for one the frame cannot carry. Nothing is
65    /// sent either way, and no worker is touched.
66    pub fn new(
67        kind: Kind,
68        texts: &[String],
69        first: usize,
70        settings: &Settings,
71    ) -> Result<Self, Failure> {
72        let items = texts
73            .iter()
74            .enumerate()
75            .map(|(offset, text)| Input {
76                id: (first + offset).to_string(),
77                text: text.clone(),
78            })
79            .collect();
80        let embed = Embed { kind, items };
81        check_embed(&embed, &settings.limits)
82            .map_err(|error| failure(error.code.as_str(), error.message))?;
83        let message = Message::Embed(embed);
84        let mut frame = Vec::new();
85        write_message(&mut frame, &message, frame_limit(settings)?).map_err(|error| {
86            failure(
87                "embed_request_too_large",
88                format!("the batch from text {first} does not fit one frame: {error}"),
89            )
90        })?;
91        let Message::Embed(embed) = message else {
92            return Err(failure("embed_worker_not_ready", "the Embed was not kept"));
93        };
94        Ok(Self { embed, frame })
95    }
96}
97
98/// How a worker ended, for the operator.
99///
100/// Its name, its exit status, and the end of its stderr. Never put in a
101/// caller's answer: a kept worker serves every caller, so its stderr may hold
102/// what it wrote for another.
103#[derive(Clone, Debug, PartialEq, Eq)]
104pub struct Ending {
105    pub code: String,
106    pub status: String,
107    pub tail: String,
108}
109
110/// A call's ending: what the caller is told, and what the operator keeps.
111pub struct Lost {
112    pub failure: Failure,
113    pub ending: Ending,
114}
115
116/// How a call ended without its vectors. Boxed, so a call's `Result` stays
117/// small.
118pub enum Ended {
119    /// The worker was found dead before any byte of the batch reached it, so
120    /// the batch may go to another.
121    Unsent(Box<Lost>),
122    /// The worker had the batch, or some of it; it is never sent again.
123    Sent(Box<Lost>),
124}
125
126impl Ended {
127    pub fn lost(self) -> Lost {
128        match self {
129            Self::Unsent(lost) | Self::Sent(lost) => *lost,
130        }
131    }
132}
133
134/// What a call's ending was, before the owner says whether the deadline
135/// passed.
136#[derive(Clone, Copy)]
137enum Why {
138    Exited,
139    Protocol,
140    VectorsInvalid,
141    RefusedLate,
142}
143
144impl Why {
145    const fn code(self) -> &'static str {
146        match self {
147            Self::Exited => "embed_worker_exited",
148            Self::Protocol => "embed_worker_protocol",
149            Self::VectorsInvalid => "embed_worker_vectors_invalid",
150            Self::RefusedLate => "embed_worker_refused_late",
151        }
152    }
153}
154
155/// The worker's input, counting the bytes it accepted.
156struct Counted<'a> {
157    input: &'a mut ChildStdin,
158    written: usize,
159}
160
161impl Write for Counted<'_> {
162    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
163        let count = self.input.write(bytes)?;
164        self.written = self.written.saturating_add(count);
165        Ok(count)
166    }
167
168    fn flush(&mut self) -> io::Result<()> {
169        self.input.flush()
170    }
171}
172
173/// The end of the worker's stderr, as text.
174fn tail(bytes: &[u8]) -> String {
175    let from = bytes.len().saturating_sub(TAIL_BYTES);
176    String::from_utf8_lossy(bytes.get(from..).unwrap_or_default())
177        .trim()
178        .to_owned()
179}
180
181impl Worker {
182    /// Starts the configured worker and waits for its Ready.
183    ///
184    /// # Errors
185    ///
186    /// Named for each way a worker fails to become ready; the process is
187    /// killed and reaped before any of them is returned.
188    pub fn start(settings: &Settings) -> Result<Self, Failure> {
189        let file = std::fs::metadata(&settings.worker)
190            .map_err(|error| failure("embed_worker_missing", format!("embed worker: {error}")))?;
191        if !file.is_file() {
192            return Err(failure(
193                "embed_worker_not_file",
194                "embed worker is not a file",
195            ));
196        }
197        let deadline = Instant::now()
198            .checked_add(settings.ready)
199            .ok_or_else(|| failure("embed_worker_start_timeout", "no deadline can be set"))?;
200        let mut child = Command::new(&settings.worker)
201            .env_clear()
202            .current_dir("/")
203            .stdin(Stdio::piped())
204            .stdout(Stdio::piped())
205            .stderr(Stdio::piped())
206            .spawn()
207            .map_err(|error| {
208                failure(
209                    "embed_worker_not_executable",
210                    format!("embed worker: {error}"),
211                )
212            })?;
213        let pipes = (child.stdin.take(), child.stdout.take(), child.stderr.take());
214        let child = OwnedChild(child);
215        let (Some(input), Some(output), Some(stderr)) = pipes else {
216            drop(child);
217            return Err(failure(
218                "embed_worker_not_ready",
219                "worker pipes were not opened",
220            ));
221        };
222        let control = Control::start(child, stderr, deadline)
223            .map_err(|error| failure("embed_worker_not_ready", format!("worker owner: {error}")))?;
224        let frame = NonZeroU32::new(settings.limits.frame_bytes.max(MINIMUM_FRAME_BYTES))
225            .ok_or_else(|| failure("embed_worker_not_ready", "no frame bytes"))?;
226        let mut worker = Handshake {
227            input,
228            output,
229            control,
230            frame,
231        };
232        let ready = worker.ready(settings)?;
233        let Handshake {
234            input,
235            output,
236            control,
237            frame,
238        } = worker;
239        control
240            .arm(None)
241            .map_err(|error| failure("embed_worker_exited_before_ready", error.to_string()))?;
242        Ok(Self {
243            input,
244            output,
245            control,
246            ready,
247            frame,
248        })
249    }
250
251    /// What the worker made of each text of `batch`: its one Embed and the
252    /// one Vectors answering it. EMB1: a vector of fewer tokens than its text
253    /// holds, and a refusal as `text_longer_than_model`, are each the text's
254    /// [`Answer::Longer`], never a vector, and leave the worker running.
255    ///
256    /// # Errors
257    ///
258    /// [`Ended::Unsent`] when the worker had ended before any byte of the
259    /// batch reached it; [`Ended::Sent`] for any other ending, named. Either
260    /// way the worker is finished and reaped before this returns, and the
261    /// caller drops it.
262    pub fn embed(&mut self, batch: &Batch, settings: &Settings) -> Result<Vec<Answer>, Ended> {
263        // EMB-E2 (h): a call of documents may answer windows.
264        let call = match batch.embed.kind {
265            Kind::Document => settings.document_call(),
266            Kind::Query => settings.call,
267        };
268        let Some(deadline) = Instant::now().checked_add(call) else {
269            return Err(Ended::Sent(
270                self.ending(Why::Exited, "no call deadline can be set"),
271            ));
272        };
273        if let Err(error) = self.control.arm(Some(deadline)) {
274            return Err(Ended::Unsent(self.ending(Why::Exited, error)));
275        }
276        let mut input = Counted {
277            input: &mut self.input,
278            written: 0,
279        };
280        let sent = input.write_all(&batch.frame).and_then(|()| input.flush());
281        let written = input.written;
282        if let Err(error) = sent {
283            let ending = self.ending(Why::Exited, format!("the Embed was not sent: {error}"));
284            return Err(if written == 0 {
285                Ended::Unsent(ending)
286            } else {
287                Ended::Sent(ending)
288            });
289        }
290        let vectors = match self.answer() {
291            Ok(vectors) => vectors,
292            Err((why, message)) => return Err(Ended::Sent(self.ending(why, message))),
293        };
294        // The service refused every text it would not have embedded before
295        // sending, so a Refused here is the worker's own, named before the
296        // check that would call it invalid output. EMB1: a text longer than
297        // the model reads is the one refusal that is the text's own.
298        let refused = vectors.items.iter().find_map(|outcome| match outcome {
299            Outcome::Refused {
300                code: ItemCode::TextLongerThanModel | ItemCode::WindowsOverCap,
301                ..
302            }
303            | Outcome::Vector { .. }
304            | Outcome::Windows { .. } => None,
305            Outcome::Refused { id, .. } => Some(id.clone()),
306        });
307        if let Some(id) = refused {
308            let why = format!("text {id} was refused after the service accepted it");
309            return Err(Ended::Sent(self.ending(Why::RefusedLate, why)));
310        }
311        let checked = check_vectors_windowed(
312            &vectors,
313            &batch.embed,
314            &self.ready,
315            &settings.limits,
316            settings.windows,
317        );
318        if let Err(error) = checked {
319            return Err(Ended::Sent(self.ending(Why::VectorsInvalid, error)));
320        }
321        if let Err(error) = self.control.arm(None) {
322            return Err(Ended::Sent(self.ending(Why::Exited, error)));
323        }
324        self.answers(&vectors)
325    }
326
327    /// The answers of checked `vectors`, one for each text: a vector, its
328    /// windows, or why it has none.
329    ///
330    /// # Errors
331    /// [`Ended::Sent`], named, for a vector or window that does not decode
332    /// or a refusal no check let through.
333    fn answers(&mut self, vectors: &haem_frames::embedding::Vectors) -> Result<Vec<Answer>, Ended> {
334        let mut made = Vec::with_capacity(vectors.items.len());
335        for outcome in &vectors.items {
336            match outcome {
337                Outcome::Vector {
338                    tokens_read,
339                    tokens_total,
340                    ..
341                } if tokens_read < tokens_total => made.push(Answer::Longer(Longer {
342                    tokens_total: *tokens_total,
343                    tokens_limit: self.ready.model.max_tokens,
344                })),
345                Outcome::Refused {
346                    code: code @ (ItemCode::TextLongerThanModel | ItemCode::WindowsOverCap),
347                    tokens_total: Some(tokens_total),
348                    tokens_limit: Some(tokens_limit),
349                    ..
350                } => {
351                    let longer = Longer {
352                        tokens_total: *tokens_total,
353                        tokens_limit: *tokens_limit,
354                    };
355                    made.push(if *code == ItemCode::WindowsOverCap {
356                        Answer::OverCap(longer)
357                    } else {
358                        Answer::Longer(longer)
359                    });
360                }
361                Outcome::Windows { windows, .. } => match decode_windows(windows) {
362                    Ok(windows) => made.push(Answer::Windows(windows)),
363                    Err(error) => {
364                        return Err(Ended::Sent(self.ending(Why::VectorsInvalid, error)));
365                    }
366                },
367                Outcome::Vector { vector, .. } => match decode_vector(vector) {
368                    Ok(vector) => made.push(Answer::Vector(vector)),
369                    Err(error) => {
370                        return Err(Ended::Sent(self.ending(Why::VectorsInvalid, error)));
371                    }
372                },
373                Outcome::Refused { id, .. } => {
374                    let why = format!("text {id} was refused after the service accepted it");
375                    return Err(Ended::Sent(self.ending(Why::RefusedLate, why)));
376                }
377            }
378        }
379        Ok(made)
380    }
381
382    /// The one Vectors the worker answers, or why there is none.
383    fn answer(&mut self) -> Result<haem_frames::embedding::Vectors, (Why, String)> {
384        let frame = match frame::read(&mut self.output, self.frame) {
385            Ok(Some(frame)) => frame,
386            Ok(None) => return Err((Why::Exited, "the worker closed its output".to_owned())),
387            // The output ended inside a frame, or could not be read: the
388            // worker went. A frame that is too large or empty is its fault.
389            Err(error) => {
390                let why = match Fault::of(&error) {
391                    None | Some(Fault::Cut) => Why::Exited,
392                    Some(Fault::TooLarge { .. } | Fault::Empty) => Why::Protocol,
393                };
394                return Err((why, error.to_string()));
395            }
396        };
397        let Some((tag, payload)) = frame.split_first() else {
398            return Err((Why::Protocol, "a frame with no tag".to_owned()));
399        };
400        match decode_message(*tag, payload) {
401            Ok(Message::Vectors(vectors)) => Ok(vectors),
402            Ok(Message::Failed(failed)) => Err((
403                Why::Exited,
404                format!(
405                    "the worker sent Failed {}: {}",
406                    failed.code.as_str(),
407                    failed.message
408                ),
409            )),
410            Ok(_) => Err((Why::Protocol, "a message other than Vectors".to_owned())),
411            Err(error) => Err((Why::Protocol, error.to_string())),
412        }
413    }
414
415    /// Finishes the worker and names how the call ended: a timeout when its
416    /// deadline passed, whatever was read; otherwise `why`. The caller is
417    /// told the name and the status; the end of the worker's stderr is kept
418    /// for the operator only.
419    fn ending(&mut self, why: Why, message: impl std::fmt::Display) -> Box<Lost> {
420        match self.control.finish(true) {
421            Ok(outcome) => {
422                let code = if outcome.deadline {
423                    "embed_worker_timeout"
424                } else {
425                    why.code()
426                };
427                Box::new(Lost {
428                    failure: failure(code, format!("{message}; status {}", outcome.status)),
429                    ending: Ending {
430                        code: code.to_owned(),
431                        status: outcome.status.to_string(),
432                        tail: tail(&outcome.tail),
433                    },
434                })
435            }
436            Err(error) => {
437                let status = format!("the worker could not be reaped: {error}");
438                Box::new(Lost {
439                    failure: failure(why.code(), format!("{message}; {status}")),
440                    ending: Ending {
441                        code: why.code().to_owned(),
442                        status,
443                        tail: String::new(),
444                    },
445                })
446            }
447        }
448    }
449}
450
451/// A worker between Start and Ready.
452/// EMB-E2: a checked text's windows, each decoded.
453fn decode_windows(windows: &[Window]) -> Result<Vec<WindowVector>, String> {
454    windows
455        .iter()
456        .map(|window| {
457            decode_vector(&window.vector)
458                .map(|vector| WindowVector {
459                    start: window.start,
460                    end: window.end,
461                    vector,
462                })
463                .map_err(|error| error.to_string())
464        })
465        .collect()
466}
467
468struct Handshake {
469    input: ChildStdin,
470    output: ChildStdout,
471    control: Control,
472    frame: NonZeroU32,
473}
474
475impl Handshake {
476    fn ready(&mut self, settings: &Settings) -> Result<Ready, Failure> {
477        let protocol = if settings.windows.is_some() {
478            PROTOCOL_WINDOWS
479        } else {
480            PROTOCOL
481        };
482        let start = Message::Start(Start {
483            protocol,
484            model_dir: settings.model_dir.clone(),
485            limits: settings.limits,
486            windows: settings.windows,
487        });
488        let sent = write_message(&mut self.input, &start, self.frame)
489            .map_err(|error| error.to_string())
490            .and_then(|()| self.input.flush().map_err(|error| error.to_string()));
491        let read = sent.and_then(|()| {
492            read_message(&mut self.output, self.frame).map_err(|error| error.to_string())
493        });
494        let ready = match read {
495            Ok(Some(Message::Ready(ready))) => ready,
496            Ok(Some(_)) => {
497                return Err(self.refuse("embed_worker_not_ready", "the first frame was not Ready"));
498            }
499            Ok(None) => return Err(self.before_ready(settings, None)),
500            Err(why) => return Err(self.before_ready(settings, Some(why))),
501        };
502        // EMB-E2: a worker answers the protocol its Start named; one that
503        // asked for windows and was answered by a worker without them is
504        // refused by name rather than run without windows.
505        if ready.protocol != protocol {
506            let code = if protocol == PROTOCOL_WINDOWS {
507                "embed_protocol_old"
508            } else {
509                "embed_worker_not_ready"
510            };
511            let why = format!(
512                "Start named protocol {protocol}; Ready names {}",
513                ready.protocol
514            );
515            return Err(self.refuse(code, &why));
516        }
517        let limits = check_model_limits_windowed(&settings.limits, &ready.model, settings.windows);
518        if let Err(error) = limits {
519            return Err(self.refuse("embed_worker_not_ready", &error.to_string()));
520        }
521        if ready.core_limit != 0 || ready.environment != 0 {
522            return Err(self.refuse(
523                "embed_worker_not_ready",
524                "the worker kept a core limit or an environment",
525            ));
526        }
527        if ready.model.manifest_sha256 != settings.model_sha256 {
528            return Err(self.refuse(
529                "embed_model_pinned_mismatch",
530                &format!(
531                    "the worker's model {} is not the configured {}",
532                    ready.model.manifest_sha256, settings.model_sha256
533                ),
534            ));
535        }
536        Ok(ready)
537    }
538
539    /// Kills and reaps the worker, then names why it was refused.
540    fn refuse(&mut self, code: &str, why: &str) -> Failure {
541        let ending = match self.control.finish(true) {
542            Ok(_) => String::new(),
543            Err(error) => format!("; the worker could not be reaped: {error}"),
544        };
545        failure(code, format!("{why}{ending}"))
546    }
547
548    /// The worker ended, or stopped answering, before its Ready.
549    fn before_ready(&mut self, settings: &Settings, read: Option<String>) -> Failure {
550        let why = read.unwrap_or_else(|| "the worker closed its output".to_owned());
551        match self.control.finish(true) {
552            // STOP2: what the worker said before it was killed, which no
553            // caller's text can be in: none is sent before Ready.
554            Ok(outcome) if outcome.deadline => failure(
555                "embed_worker_start_timeout",
556                format!(
557                    "no Ready within {:?}; status {}; stderr: {}",
558                    settings.ready,
559                    outcome.status,
560                    tail(&outcome.tail)
561                ),
562            ),
563            Ok(outcome) => failure(
564                "embed_worker_exited_before_ready",
565                format!("{why}; status {}", outcome.status),
566            ),
567            Err(error) => failure(
568                "embed_worker_exited_before_ready",
569                format!("{why}; the worker could not be reaped: {error}"),
570            ),
571        }
572    }
573}
574
575#[cfg(all(test, unix))]
576#[path = "worker_start_tests.rs"]
577mod tests;
haematite_cli/serve/embed/worker.rs

haematite_cli/serve/embed/
worker.rs
