haematite-cli 0.21.1

//! One embedding worker process: started, checked, and asked one batch at a
//! time. Dropping it kills and reaps the process.
//!
//! Every ending after Ready is named (box 42): the owner is finished and its
//! outcome read, so the answer carries the status and the end of the
//! worker's stderr, and a call past its deadline is a timeout whatever the
//! read then saw. A batch is encoded before anything is sent, and the bytes
//! the worker's input accepted are counted, so a worker found dead before a
//! single byte reached it is told apart from one that died with the batch.

use std::io::{self, Write};
use std::num::NonZeroU32;
use std::process::{ChildStdin, ChildStdout, Command, Stdio};
use std::time::Instant;

use haem_frames::embedding::{
    Embed, Input, ItemCode, Kind, MINIMUM_FRAME_BYTES, Message, Outcome, PROTOCOL,
    PROTOCOL_WINDOWS, Ready, Start, Window, check_embed, check_model_limits_windowed,
    check_vectors_windowed, decode_message, decode_vector, read_message, write_message,
};
use haem_frames::frame::{self, Fault};
use haematite_query::worker_owner::{Control, OwnedChild};

use super::Settings;
use super::embedder::{Answer, Longer, WindowVector};
use crate::serve::storage::Failure;

/// A started worker that sent a Ready the service accepts.
pub struct Worker {
    // Dropped first: closing the worker's input before its owner kills it.
    input: ChildStdin,
    output: ChildStdout,
    control: Control,
    pub ready: Ready,
    frame: NonZeroU32,
}

fn failure(code: &str, message: impl Into<String>) -> Failure {
    Failure::new(code, message.into())
}

/// The most of the worker's stderr an answer carries, from its end.
const TAIL_BYTES: usize = 512;

/// One batch, checked and encoded before anything is sent.
pub struct Batch {
    embed: Embed,
    frame: Vec<u8>,
}

/// The frame limit both sides use.
fn frame_limit(settings: &Settings) -> Result<NonZeroU32, Failure> {
    NonZeroU32::new(settings.limits.frame_bytes.max(MINIMUM_FRAME_BYTES))
        .ok_or_else(|| failure("embed_worker_not_ready", "no frame bytes"))
}

impl Batch {
    /// `texts` as one Embed whose item ids are their places in the whole
    /// request, from `first`.
    ///
    /// # Errors
    ///
    /// `check_embed`'s own name for an Embed it refuses, and
    /// `embed_request_too_large` for one the frame cannot carry. Nothing is
    /// sent either way, and no worker is touched.
    pub fn new(
        kind: Kind,
        texts: &[String],
        first: usize,
        settings: &Settings,
    ) -> Result<Self, Failure> {
        let items = texts
            .iter()
            .enumerate()
            .map(|(offset, text)| Input {
                id: (first + offset).to_string(),
                text: text.clone(),
            })
            .collect();
        let embed = Embed { kind, items };
        check_embed(&embed, &settings.limits)
            .map_err(|error| failure(error.code.as_str(), error.message))?;
        let message = Message::Embed(embed);
        let mut frame = Vec::new();
        write_message(&mut frame, &message, frame_limit(settings)?).map_err(|error| {
            failure(
                "embed_request_too_large",
                format!("the batch from text {first} does not fit one frame: {error}"),
            )
        })?;
        let Message::Embed(embed) = message else {
            return Err(failure("embed_worker_not_ready", "the Embed was not kept"));
        };
        Ok(Self { embed, frame })
    }
}

/// How a worker ended, for the operator.
///
/// Its name, its exit status, and the end of its stderr. Never put in a
/// caller's answer: a kept worker serves every caller, so its stderr may hold
/// what it wrote for another.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Ending {
    pub code: String,
    pub status: String,
    pub tail: String,
}

/// A call's ending: what the caller is told, and what the operator keeps.
pub struct Lost {
    pub failure: Failure,
    pub ending: Ending,
}

/// How a call ended without its vectors. Boxed, so a call's `Result` stays
/// small.
pub enum Ended {
    /// The worker was found dead before any byte of the batch reached it, so
    /// the batch may go to another.
    Unsent(Box<Lost>),
    /// The worker had the batch, or some of it; it is never sent again.
    Sent(Box<Lost>),
}

impl Ended {
    pub fn lost(self) -> Lost {
        match self {
            Self::Unsent(lost) | Self::Sent(lost) => *lost,
        }
    }
}

/// What a call's ending was, before the owner says whether the deadline
/// passed.
#[derive(Clone, Copy)]
enum Why {
    Exited,
    Protocol,
    VectorsInvalid,
    RefusedLate,
}

impl Why {
    const fn code(self) -> &'static str {
        match self {
            Self::Exited => "embed_worker_exited",
            Self::Protocol => "embed_worker_protocol",
            Self::VectorsInvalid => "embed_worker_vectors_invalid",
            Self::RefusedLate => "embed_worker_refused_late",
        }
    }
}

/// The worker's input, counting the bytes it accepted.
struct Counted<'a> {
    input: &'a mut ChildStdin,
    written: usize,
}

impl Write for Counted<'_> {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        let count = self.input.write(bytes)?;
        self.written = self.written.saturating_add(count);
        Ok(count)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.input.flush()
    }
}

/// The end of the worker's stderr, as text.
fn tail(bytes: &[u8]) -> String {
    let from = bytes.len().saturating_sub(TAIL_BYTES);
    String::from_utf8_lossy(bytes.get(from..).unwrap_or_default())
        .trim()
        .to_owned()
}

impl Worker {
    /// Starts the configured worker and waits for its Ready.
    ///
    /// # Errors
    ///
    /// Named for each way a worker fails to become ready; the process is
    /// killed and reaped before any of them is returned.
    pub fn start(settings: &Settings) -> Result<Self, Failure> {
        let file = std::fs::metadata(&settings.worker)
            .map_err(|error| failure("embed_worker_missing", format!("embed worker: {error}")))?;
        if !file.is_file() {
            return Err(failure(
                "embed_worker_not_file",
                "embed worker is not a file",
            ));
        }
        let deadline = Instant::now()
            .checked_add(settings.ready)
            .ok_or_else(|| failure("embed_worker_start_timeout", "no deadline can be set"))?;
        let mut child = Command::new(&settings.worker)
            .env_clear()
            .current_dir("/")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|error| {
                failure(
                    "embed_worker_not_executable",
                    format!("embed worker: {error}"),
                )
            })?;
        let pipes = (child.stdin.take(), child.stdout.take(), child.stderr.take());
        let child = OwnedChild(child);
        let (Some(input), Some(output), Some(stderr)) = pipes else {
            drop(child);
            return Err(failure(
                "embed_worker_not_ready",
                "worker pipes were not opened",
            ));
        };
        let control = Control::start(child, stderr, deadline)
            .map_err(|error| failure("embed_worker_not_ready", format!("worker owner: {error}")))?;
        let frame = NonZeroU32::new(settings.limits.frame_bytes.max(MINIMUM_FRAME_BYTES))
            .ok_or_else(|| failure("embed_worker_not_ready", "no frame bytes"))?;
        let mut worker = Handshake {
            input,
            output,
            control,
            frame,
        };
        let ready = worker.ready(settings)?;
        let Handshake {
            input,
            output,
            control,
            frame,
        } = worker;
        control
            .arm(None)
            .map_err(|error| failure("embed_worker_exited_before_ready", error.to_string()))?;
        Ok(Self {
            input,
            output,
            control,
            ready,
            frame,
        })
    }

    /// What the worker made of each text of `batch`: its one Embed and the
    /// one Vectors answering it. EMB1: a vector of fewer tokens than its text
    /// holds, and a refusal as `text_longer_than_model`, are each the text's
    /// [`Answer::Longer`], never a vector, and leave the worker running.
    ///
    /// # Errors
    ///
    /// [`Ended::Unsent`] when the worker had ended before any byte of the
    /// batch reached it; [`Ended::Sent`] for any other ending, named. Either
    /// way the worker is finished and reaped before this returns, and the
    /// caller drops it.
    pub fn embed(&mut self, batch: &Batch, settings: &Settings) -> Result<Vec<Answer>, Ended> {
        // EMB-E2 (h): a call of documents may answer windows.
        let call = match batch.embed.kind {
            Kind::Document => settings.document_call(),
            Kind::Query => settings.call,
        };
        let Some(deadline) = Instant::now().checked_add(call) else {
            return Err(Ended::Sent(
                self.ending(Why::Exited, "no call deadline can be set"),
            ));
        };
        if let Err(error) = self.control.arm(Some(deadline)) {
            return Err(Ended::Unsent(self.ending(Why::Exited, error)));
        }
        let mut input = Counted {
            input: &mut self.input,
            written: 0,
        };
        let sent = input.write_all(&batch.frame).and_then(|()| input.flush());
        let written = input.written;
        if let Err(error) = sent {
            let ending = self.ending(Why::Exited, format!("the Embed was not sent: {error}"));
            return Err(if written == 0 {
                Ended::Unsent(ending)
            } else {
                Ended::Sent(ending)
            });
        }
        let vectors = match self.answer() {
            Ok(vectors) => vectors,
            Err((why, message)) => return Err(Ended::Sent(self.ending(why, message))),
        };
        // The service refused every text it would not have embedded before
        // sending, so a Refused here is the worker's own, named before the
        // check that would call it invalid output. EMB1: a text longer than
        // the model reads is the one refusal that is the text's own.
        let refused = vectors.items.iter().find_map(|outcome| match outcome {
            Outcome::Refused {
                code: ItemCode::TextLongerThanModel | ItemCode::WindowsOverCap,
                ..
            }
            | Outcome::Vector { .. }
            | Outcome::Windows { .. } => None,
            Outcome::Refused { id, .. } => Some(id.clone()),
        });
        if let Some(id) = refused {
            let why = format!("text {id} was refused after the service accepted it");
            return Err(Ended::Sent(self.ending(Why::RefusedLate, why)));
        }
        let checked = check_vectors_windowed(
            &vectors,
            &batch.embed,
            &self.ready,
            &settings.limits,
            settings.windows,
        );
        if let Err(error) = checked {
            return Err(Ended::Sent(self.ending(Why::VectorsInvalid, error)));
        }
        if let Err(error) = self.control.arm(None) {
            return Err(Ended::Sent(self.ending(Why::Exited, error)));
        }
        self.answers(&vectors)
    }

    /// The answers of checked `vectors`, one for each text: a vector, its
    /// windows, or why it has none.
    ///
    /// # Errors
    /// [`Ended::Sent`], named, for a vector or window that does not decode
    /// or a refusal no check let through.
    fn answers(&mut self, vectors: &haem_frames::embedding::Vectors) -> Result<Vec<Answer>, Ended> {
        let mut made = Vec::with_capacity(vectors.items.len());
        for outcome in &vectors.items {
            match outcome {
                Outcome::Vector {
                    tokens_read,
                    tokens_total,
                    ..
                } if tokens_read < tokens_total => made.push(Answer::Longer(Longer {
                    tokens_total: *tokens_total,
                    tokens_limit: self.ready.model.max_tokens,
                })),
                Outcome::Refused {
                    code: code @ (ItemCode::TextLongerThanModel | ItemCode::WindowsOverCap),
                    tokens_total: Some(tokens_total),
                    tokens_limit: Some(tokens_limit),
                    ..
                } => {
                    let longer = Longer {
                        tokens_total: *tokens_total,
                        tokens_limit: *tokens_limit,
                    };
                    made.push(if *code == ItemCode::WindowsOverCap {
                        Answer::OverCap(longer)
                    } else {
                        Answer::Longer(longer)
                    });
                }
                Outcome::Windows { windows, .. } => match decode_windows(windows) {
                    Ok(windows) => made.push(Answer::Windows(windows)),
                    Err(error) => {
                        return Err(Ended::Sent(self.ending(Why::VectorsInvalid, error)));
                    }
                },
                Outcome::Vector { vector, .. } => match decode_vector(vector) {
                    Ok(vector) => made.push(Answer::Vector(vector)),
                    Err(error) => {
                        return Err(Ended::Sent(self.ending(Why::VectorsInvalid, error)));
                    }
                },
                Outcome::Refused { id, .. } => {
                    let why = format!("text {id} was refused after the service accepted it");
                    return Err(Ended::Sent(self.ending(Why::RefusedLate, why)));
                }
            }
        }
        Ok(made)
    }

    /// The one Vectors the worker answers, or why there is none.
    fn answer(&mut self) -> Result<haem_frames::embedding::Vectors, (Why, String)> {
        let frame = match frame::read(&mut self.output, self.frame) {
            Ok(Some(frame)) => frame,
            Ok(None) => return Err((Why::Exited, "the worker closed its output".to_owned())),
            // The output ended inside a frame, or could not be read: the
            // worker went. A frame that is too large or empty is its fault.
            Err(error) => {
                let why = match Fault::of(&error) {
                    None | Some(Fault::Cut) => Why::Exited,
                    Some(Fault::TooLarge { .. } | Fault::Empty) => Why::Protocol,
                };
                return Err((why, error.to_string()));
            }
        };
        let Some((tag, payload)) = frame.split_first() else {
            return Err((Why::Protocol, "a frame with no tag".to_owned()));
        };
        match decode_message(*tag, payload) {
            Ok(Message::Vectors(vectors)) => Ok(vectors),
            Ok(Message::Failed(failed)) => Err((
                Why::Exited,
                format!(
                    "the worker sent Failed {}: {}",
                    failed.code.as_str(),
                    failed.message
                ),
            )),
            Ok(_) => Err((Why::Protocol, "a message other than Vectors".to_owned())),
            Err(error) => Err((Why::Protocol, error.to_string())),
        }
    }

    /// Finishes the worker and names how the call ended: a timeout when its
    /// deadline passed, whatever was read; otherwise `why`. The caller is
    /// told the name and the status; the end of the worker's stderr is kept
    /// for the operator only.
    fn ending(&mut self, why: Why, message: impl std::fmt::Display) -> Box<Lost> {
        match self.control.finish(true) {
            Ok(outcome) => {
                let code = if outcome.deadline {
                    "embed_worker_timeout"
                } else {
                    why.code()
                };
                Box::new(Lost {
                    failure: failure(code, format!("{message}; status {}", outcome.status)),
                    ending: Ending {
                        code: code.to_owned(),
                        status: outcome.status.to_string(),
                        tail: tail(&outcome.tail),
                    },
                })
            }
            Err(error) => {
                let status = format!("the worker could not be reaped: {error}");
                Box::new(Lost {
                    failure: failure(why.code(), format!("{message}; {status}")),
                    ending: Ending {
                        code: why.code().to_owned(),
                        status,
                        tail: String::new(),
                    },
                })
            }
        }
    }
}

/// A worker between Start and Ready.
/// EMB-E2: a checked text's windows, each decoded.
fn decode_windows(windows: &[Window]) -> Result<Vec<WindowVector>, String> {
    windows
        .iter()
        .map(|window| {
            decode_vector(&window.vector)
                .map(|vector| WindowVector {
                    start: window.start,
                    end: window.end,
                    vector,
                })
                .map_err(|error| error.to_string())
        })
        .collect()
}

struct Handshake {
    input: ChildStdin,
    output: ChildStdout,
    control: Control,
    frame: NonZeroU32,
}

impl Handshake {
    fn ready(&mut self, settings: &Settings) -> Result<Ready, Failure> {
        let protocol = if settings.windows.is_some() {
            PROTOCOL_WINDOWS
        } else {
            PROTOCOL
        };
        let start = Message::Start(Start {
            protocol,
            model_dir: settings.model_dir.clone(),
            limits: settings.limits,
            windows: settings.windows,
        });
        let sent = write_message(&mut self.input, &start, self.frame)
            .map_err(|error| error.to_string())
            .and_then(|()| self.input.flush().map_err(|error| error.to_string()));
        let read = sent.and_then(|()| {
            read_message(&mut self.output, self.frame).map_err(|error| error.to_string())
        });
        let ready = match read {
            Ok(Some(Message::Ready(ready))) => ready,
            Ok(Some(_)) => {
                return Err(self.refuse("embed_worker_not_ready", "the first frame was not Ready"));
            }
            Ok(None) => return Err(self.before_ready(settings, None)),
            Err(why) => return Err(self.before_ready(settings, Some(why))),
        };
        // EMB-E2: a worker answers the protocol its Start named; one that
        // asked for windows and was answered by a worker without them is
        // refused by name rather than run without windows.
        if ready.protocol != protocol {
            let code = if protocol == PROTOCOL_WINDOWS {
                "embed_protocol_old"
            } else {
                "embed_worker_not_ready"
            };
            let why = format!(
                "Start named protocol {protocol}; Ready names {}",
                ready.protocol
            );
            return Err(self.refuse(code, &why));
        }
        let limits = check_model_limits_windowed(&settings.limits, &ready.model, settings.windows);
        if let Err(error) = limits {
            return Err(self.refuse("embed_worker_not_ready", &error.to_string()));
        }
        if ready.core_limit != 0 || ready.environment != 0 {
            return Err(self.refuse(
                "embed_worker_not_ready",
                "the worker kept a core limit or an environment",
            ));
        }
        if ready.model.manifest_sha256 != settings.model_sha256 {
            return Err(self.refuse(
                "embed_model_pinned_mismatch",
                &format!(
                    "the worker's model {} is not the configured {}",
                    ready.model.manifest_sha256, settings.model_sha256
                ),
            ));
        }
        Ok(ready)
    }

    /// Kills and reaps the worker, then names why it was refused.
    fn refuse(&mut self, code: &str, why: &str) -> Failure {
        let ending = match self.control.finish(true) {
            Ok(_) => String::new(),
            Err(error) => format!("; the worker could not be reaped: {error}"),
        };
        failure(code, format!("{why}{ending}"))
    }

    /// The worker ended, or stopped answering, before its Ready.
    fn before_ready(&mut self, settings: &Settings, read: Option<String>) -> Failure {
        let why = read.unwrap_or_else(|| "the worker closed its output".to_owned());
        match self.control.finish(true) {
            // STOP2: what the worker said before it was killed, which no
            // caller's text can be in: none is sent before Ready.
            Ok(outcome) if outcome.deadline => failure(
                "embed_worker_start_timeout",
                format!(
                    "no Ready within {:?}; status {}; stderr: {}",
                    settings.ready,
                    outcome.status,
                    tail(&outcome.tail)
                ),
            ),
            Ok(outcome) => failure(
                "embed_worker_exited_before_ready",
                format!("{why}; status {}", outcome.status),
            ),
            Err(error) => failure(
                "embed_worker_exited_before_ready",
                format!("{why}; the worker could not be reaped: {error}"),
            ),
        }
    }
}

#[cfg(all(test, unix))]
#[path = "worker_start_tests.rs"]
mod tests;