haematite-cli 0.21.3

//! The service's one embedding worker, kept between requests behind one owner.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Mutex, PoisonError};

use haem_frames::embedding::{Distance, Kind, input_refusal};

use super::Settings;
use super::worker::{Batch, Ended, Ending, Worker};
use crate::serve::storage::Failure;

/// The model that made a set of vectors, as its worker's Ready named it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ModelId {
    pub name: String,
    pub revision: String,
    pub manifest_sha256: String,
    pub dimensions: u64,
    pub distance: Distance,
    pub normalised: bool,
}

/// What the worker made of one text.
#[derive(Debug, Clone, PartialEq)]
pub enum Answer {
    Vector(Vec<f32>),
    /// EMB1: the text has more tokens than the model reads. No vector is
    /// made of it: a vector of part of a text is never kept.
    Longer(Longer),
    /// EMB-E2: a text longer than the model reads, embedded whole as
    /// overlapping windows in text order.
    Windows(Vec<WindowVector>),
    /// EMB-E2: a text whose whole cover needs more windows than the cap. No
    /// vector is made of it, as for [`Answer::Longer`].
    OverCap(Longer),
}

/// EMB-E2: one window of a text, its bytes `[start, end)`, and its vector.
#[derive(Debug, Clone, PartialEq)]
pub struct WindowVector {
    pub start: u64,
    pub end: u64,
    pub vector: Vec<f32>,
}

/// A text longer than the model reads: its content tokens, without special
/// or prefix tokens, and the model's limit as its Ready declared it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Longer {
    pub tokens_total: u64,
    pub tokens_limit: u64,
}

/// What the worker made of each text, in the order the texts were given,
/// and the model that made them.
#[derive(Debug)]
pub struct Embedded {
    pub model: ModelId,
    pub answers: Vec<Answer>,
}

/// Starts one worker on first use and keeps it; a second caller waits for it.
pub struct Embedder {
    settings: Option<Settings>,
    worker: Mutex<Option<Worker>>,
    waiting: AtomicUsize,
    /// The last worker ending, for the operator; taken only with the worker's
    /// lock released.
    last: Mutex<Option<Ending>>,
}

/// One caller counted among those waiting for the worker, until dropped.
struct Waiting<'a>(&'a AtomicUsize);

impl Drop for Waiting<'_> {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::AcqRel);
    }
}

impl Embedder {
    /// An embedder with `settings`; `None` refuses every request by name.
    /// Nothing is started until the first request.
    #[must_use]
    pub const fn new(settings: Option<Settings>) -> Self {
        Self {
            settings,
            worker: Mutex::new(None),
            waiting: AtomicUsize::new(0),
            last: Mutex::new(None),
        }
    }

    /// The vectors of `texts`, all of them in the order given, or a refusal.
    ///
    /// A kept worker found dead before any byte of a batch reached it is
    /// replaced, once, and the batch goes to the new one: a worker that died
    /// between requests costs no request. One that dies with a batch answers
    /// by name, and that batch is never sent again.
    ///
    /// # Errors
    ///
    /// `embed_unavailable` with nothing configured; `embed_request_empty`,
    /// `embed_request_too_large`, `check_embed`'s own name, or a text's own
    /// refusal before anything is sent, with any kept worker left running;
    /// `embed_busy` past the waiting bound; a start's refusal; any ending of
    /// the worker, by name, after which the next request starts another.
    pub fn embed(&self, kind: Kind, texts: &[String]) -> Result<Embedded, Failure> {
        let settings = self.settings.as_ref().ok_or_else(|| {
            Failure::new(
                "embed_unavailable",
                "no embedding worker is configured for this service",
            )
        })?;
        asked(settings, texts)?;
        let batches = texts
            .chunks(settings.batch.get())
            .enumerate()
            .map(|(number, batch)| Batch::new(kind, batch, number * settings.batch.get(), settings))
            .collect::<Result<Vec<_>, _>>()?;
        let waiting = self.wait(settings)?;
        let mut worker = self.worker.lock().map_err(|_| {
            Failure::new(
                "embed_worker_owner_failed",
                "a request panicked while it held the worker",
            )
        })?;
        drop(waiting);
        let mut answers = Vec::with_capacity(texts.len());
        let mut ending = None;
        for batch in &batches {
            match call(&mut worker, batch, settings, &mut ending) {
                Ok(made) => answers.extend(made),
                Err(ended) => {
                    drop(worker);
                    self.remember(ending);
                    return Err(ended);
                }
            }
        }
        let Some(kept) = worker.as_ref() else {
            return Err(not_started());
        };
        let model = &kept.ready.model;
        let model = ModelId {
            name: model.name.clone(),
            revision: model.revision.clone(),
            manifest_sha256: model.manifest_sha256.clone(),
            dimensions: model.dimensions,
            distance: model.distance,
            normalised: model.normalised,
        };
        drop(worker);
        self.remember(ending);
        Ok(Embedded { model, answers })
    }

    /// The last way a worker ended: its name, exit status and the end of its
    /// stderr, for an operator. No caller's answer carries the stderr.
    pub fn last_ending(&self) -> Option<Ending> {
        let last = self.last.lock().unwrap_or_else(PoisonError::into_inner);
        let ending = last.clone();
        drop(last);
        ending
    }

    /// Keep `ending`, if there is one, as the last. Called only once the
    /// worker's lock is released, so the two locks are never held together.
    fn remember(&self, ending: Option<Ending>) {
        let Some(ending) = ending else {
            return;
        };
        let mut last = self.last.lock().unwrap_or_else(PoisonError::into_inner);
        *last = Some(ending);
        drop(last);
    }

    fn wait(&self, settings: &Settings) -> Result<Waiting<'_>, Failure> {
        let before = self.waiting.fetch_add(1, Ordering::AcqRel);
        let waiting = Waiting(&self.waiting);
        if before >= settings.waiting {
            return Err(Failure::new(
                "embed_busy",
                format!("{before} requests already wait for the embedding worker"),
            ));
        }
        Ok(waiting)
    }
}

/// A request's texts, refused by name before anything is sent.
fn asked(settings: &Settings, texts: &[String]) -> Result<(), Failure> {
    if texts.is_empty() {
        return Err(Failure::new("embed_request_empty", "no texts were given"));
    }
    if texts.len() > settings.request_items.get() {
        return Err(Failure::new(
            "embed_request_too_large",
            format!(
                "{} texts were given; at most {} may be",
                texts.len(),
                settings.request_items
            ),
        ));
    }
    for (index, text) in texts.iter().enumerate() {
        if let Some(code) = input_refusal(text, settings.limits.input_bytes) {
            let code = serde_json::to_value(code).map_err(Failure::storage)?;
            return Err(Failure::new(
                code.as_str().unwrap_or("embed_input_refused"),
                format!("text {index} cannot be embedded"),
            ));
        }
    }
    Ok(())
}

/// One batch through the kept worker, started if none is kept. A kept worker
/// found dead before the batch reached it is replaced once; any other ending
/// drops the worker, already reaped, and answers by name.
fn call(
    worker: &mut Option<Worker>,
    batch: &Batch,
    settings: &Settings,
    ending: &mut Option<Ending>,
) -> Result<Vec<Answer>, Failure> {
    for attempt in 0..2 {
        if worker.is_none() {
            *worker = Some(Worker::start(settings)?);
        }
        let Some(kept) = worker.as_mut() else {
            return Err(not_started());
        };
        match kept.embed(batch, settings) {
            Ok(made) => return Ok(made),
            // Dead before the batch reached it: nothing of it was sent, so
            // the new worker may have it.
            Err(Ended::Unsent(unsent)) if attempt == 0 => {
                *worker = None;
                *ending = Some(unsent.ending);
            }
            Err(ended) => {
                *worker = None;
                let lost = ended.lost();
                *ending = Some(lost.ending);
                return Err(lost.failure);
            }
        }
    }
    Err(not_started())
}

fn not_started() -> Failure {
    Failure::new("embed_worker_not_ready", "the worker was not kept")
}