haematite_cli/serve/embed/

job.rs

//! V7a (ruled 5 Oct 08:48): one page of the embedding job as three steps, so
//! the service makes the worker call holding no storage admission.
//!
//! 1. `Storage::embed_pending`, a shared read: the next page of the walk.
//! 2. [`make`]: the worker's vectors for the page's texts, sent in batches of
//!    the worker's batch size. It touches no storage.
//! 3. `Storage::embed_commit`: the page and the job's position, one commit.
//!
//! [`step`] runs the three in order, for tests and for V7b to split across its
//! locks. A page with no text to send calls the worker not at all and still
//! commits, so the position moves past removals and records with no text. A
//! failed or refused worker call commits nothing: the position and the
//! companion are unchanged, the failure's name is returned, and the next step
//! takes the same page. There is no retry here; when to try again is the
//! service's (V7b).
//!
//! EMB1 (Tom: never trim below what the model reads, never throw data away):
//! a text the model reads only part of has no vector. The page names it,
//! with its tokens and the model's limit, commits the rest of its batch, and
//! the walk goes past it, so it is not sent again until its text changes.
//!
//! EMB-E1: a text the worker would refuse (`embed_input_empty`, or too long
//! in bytes) is never put in a batch, since one refused text refuses its
//! whole request. It is passed with no vector, named, and the walk goes on.
//!
//! EMB-E2: with windows on, a text longer than the model reads comes back as
//! its windows, each vector stored with its span; one past the window cap
//! comes back `windows_over_cap` and, as EMB1, gets no vector.

use std::num::NonZeroUsize;
use std::time::Duration;

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

use super::{Answer, Embedded, Embedder, Longer};
use crate::serve::protocol::{Selection, Version};
use crate::serve::storage::embed_job::{Made, MadeWindows, Page, Reach, Window};
use crate::serve::storage::embed_pending::{Pending, Waiting};
use crate::serve::storage::embedded::{Declaration, ModelName};
use crate::serve::storage::{Failure, Storage};

/// Texts in, vectors out: what the job needs of a worker.
pub trait Embed {
    /// The vectors of `texts` as documents, in the order given, and the
    /// model that made them.
    ///
    /// # Errors
    /// Whatever the worker answers.
    fn documents(&self, texts: &[String]) -> Result<Embedded, Failure>;
}

impl Embed for Embedder {
    fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
        self.embed(Kind::Document, texts)
    }
}

/// The job's bounds, from the service's settings.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Bounds {
    /// Records one page names; one commit each.
    pub page_records: NonZeroUsize,
    /// Texts one worker call takes (`embed_batch_items`).
    pub batch: NonZeroUsize,
    /// The longest text the worker takes (`embed_input_bytes`).
    pub input_bytes: NonZeroUsize,
    /// HANG1: the most one worker call may take (`embed_call_ms`).
    pub call: Duration,
}

/// A record whose text has more tokens than the model reads: no vector is
/// made of it (EMB1), as none is of a text too long in bytes (`too_long`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LongerThanModel {
    pub id: String,
    /// The text's content tokens, without special or prefix tokens.
    pub tokens_total: u64,
    /// The model's limit, as its Ready declared it.
    pub tokens_limit: u64,
}

/// A page made: what its commit writes, and the records it made no vector
/// of because the model reads only part of their text, or the worker would
/// refuse it.
#[derive(Debug, Clone, PartialEq)]
pub struct MadePage {
    pub page: Page,
    pub longer_than_model: Vec<LongerThanModel>,
    /// EMB-E1: records whose text is empty or only whitespace, never sent.
    pub empty_input: Vec<String>,
    /// Records whose text is longer in bytes than a worker takes, never sent.
    pub too_long: Vec<String>,
    /// EMB-E2: records whose whole cover needs more windows than the cap,
    /// given no vector.
    pub windows_over_cap: Vec<LongerThanModel>,
}

/// What one committed page held.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Counts {
    pub made: usize,
    pub removed: usize,
    pub no_text: usize,
    pub too_long: usize,
    pub longer_than_model: usize,
    /// EMB-E1: records passed because their text is empty or whitespace.
    pub empty_input: usize,
    /// EMB-E2: records embedded whole as windows.
    pub windowed: usize,
    /// EMB-E2: records past the window cap, given no vector.
    pub windows_over_cap: usize,
    /// The page ended its walk.
    pub ended: bool,
}

/// What one step did.
#[derive(Debug, Clone)]
pub enum Step {
    /// Nothing pending: the job's position is the head.
    Idle,
    /// One page committed.
    Committed(Counts),
    /// The worker failed or was refused, by this name; nothing committed.
    Failed(Failure),
}

/// The page of `pending`: its texts' vectors and its removals.
///
/// The vectors come from `worker` in batches of `bounds.batch`. Its texts the
/// model reads only part of, or the worker would refuse, are named and given
/// no vector.
///
/// # Errors
/// The worker's failure; `embed_vector_count` for an answer holding another
/// number of vectors than texts sent; `embed_vector_length` for a vector of
/// another length than declared; `embed_model_mismatch` for vectors of
/// another model than the declared one. Each is the worker's failing, so
/// [`step`] gives it as [`Step::Failed`], never as a storage error.
pub fn make(
    pending: &Pending,
    collection: &str,
    field: &str,
    declaration: &Declaration,
    worker: &dyn Embed,
    bounds: Bounds,
) -> Result<MadePage, Failure> {
    let Sorted {
        texts,
        removed,
        empty_input,
        too_long,
    } = sorted(pending, bounds.input_bytes);
    // RED4: a record too long in bytes, whose source is no longer a string
    // (P2a), or whose text is now empty (EMB-E1) keeps no vector of an
    // earlier text.
    let cleared: Vec<String> = too_long
        .iter()
        .chain(&pending.no_text)
        .chain(&empty_input)
        .cloned()
        .collect();
    let mut collected = Collected {
        cleared,
        ..Collected::default()
    };
    for chunk in texts.chunks(bounds.batch.get()) {
        let sent: Vec<String> = chunk.iter().map(|(_, _, text)| (*text).clone()).collect();
        let answers = checked(worker.documents(&sent)?, sent.len(), declaration)?;
        for ((id, made_from, _), answer) in chunk.iter().zip(answers) {
            collected.take(id, *made_from, answer);
        }
    }
    let page = Page {
        collection: collection.to_owned(),
        field: field.to_owned(),
        model: declaration.model.clone(),
        made: collected.made,
        windowed: collected.windowed,
        removed,
        cleared: collected.cleared,
        toward: pending.toward,
        reach: pending.reach.clone(),
    };
    Ok(MadePage {
        page,
        longer_than_model: collected.longer_than_model,
        empty_input,
        too_long,
        windows_over_cap: collected.windows_over_cap,
    })
}

/// One call's answers, checked: one for each text sent, every vector and
/// window of the declared length, made by the declared model.
///
/// # Errors
/// `embed_vector_count`, `embed_vector_length` or `embed_model_mismatch`.
fn checked(
    embedded: Embedded,
    sent: usize,
    declaration: &Declaration,
) -> Result<Vec<Answer>, Failure> {
    if embedded.answers.len() != sent {
        return Err(Failure::new(
            "embed_vector_count",
            format!(
                "{sent} texts were sent and {} vectors came back",
                embedded.answers.len()
            ),
        ));
    }
    let dimensions = declaration.dimensions.get();
    let wrong = embedded
        .answers
        .iter()
        .flat_map(|answer| match answer {
            Answer::Vector(vector) => vec![vector.len()],
            Answer::Windows(windows) => windows.iter().map(|window| window.vector.len()).collect(),
            Answer::Longer(_) | Answer::OverCap(_) => Vec::new(),
        })
        .find(|length| *length != dimensions);
    if let Some(wrong) = wrong {
        return Err(Failure::new(
            "embed_vector_length",
            format!("a vector of {wrong} members came back; {dimensions} are declared"),
        ));
    }
    let model = ModelName {
        name: embedded.model.name,
        revision: embedded.model.revision,
        manifest_sha256: embedded.model.manifest_sha256,
    };
    if model != declaration.model {
        return Err(Failure::new(
            "embed_model_mismatch",
            format!(
                "the worker runs {}; the field is declared for {}",
                model.name, declaration.model.name
            ),
        ));
    }
    Ok(embedded.answers)
}

/// What a page's answers made: its vectors, its windowed records (EMB-E2),
/// and the records given no vector, each named.
#[derive(Default)]
struct Collected {
    made: Vec<Made>,
    windowed: Vec<MadeWindows>,
    cleared: Vec<String>,
    longer_than_model: Vec<LongerThanModel>,
    windows_over_cap: Vec<LongerThanModel>,
}

impl Collected {
    /// One record's answer. A text the model reads only part of, or one past
    /// the window cap, gets no vector, and any earlier one is cleared.
    fn take(&mut self, id: &str, made_from: Version, answer: Answer) {
        let named = |longer: Longer| LongerThanModel {
            id: id.to_owned(),
            tokens_total: longer.tokens_total,
            tokens_limit: longer.tokens_limit,
        };
        match answer {
            Answer::Vector(vector) => self.made.push(Made {
                id: id.to_owned(),
                made_from,
                vector,
            }),
            Answer::Windows(windows) => self.windowed.push(MadeWindows {
                id: id.to_owned(),
                made_from,
                windows: windows
                    .into_iter()
                    .map(|window| Window {
                        start: window.start,
                        end: window.end,
                        vector: window.vector,
                    })
                    .collect(),
            }),
            Answer::Longer(longer) => {
                self.cleared.push(id.to_owned());
                self.longer_than_model.push(named(longer));
            }
            Answer::OverCap(longer) => {
                self.cleared.push(id.to_owned());
                self.windows_over_cap.push(named(longer));
            }
        }
    }
}

/// A page's texts to send, its removals, and the texts the worker would
/// refuse, each named once.
struct Sorted<'page> {
    texts: Vec<(&'page String, Version, &'page String)>,
    removed: Vec<String>,
    empty_input: Vec<String>,
    too_long: Vec<String>,
}

/// EMB-E1: `pending`'s texts, each checked as the worker checks a request
/// before any batch is built, so a text it would refuse is never sent and
/// never refuses its batch-mates. Those the walk already passed come first.
fn sorted(pending: &Pending, input_bytes: NonZeroUsize) -> Sorted<'_> {
    let limit = u64::try_from(input_bytes.get()).unwrap_or(u64::MAX);
    let mut sorted = Sorted {
        texts: Vec::new(),
        removed: Vec::new(),
        empty_input: pending.empty_input.clone(),
        too_long: pending.too_long.clone(),
    };
    for waiting in &pending.waiting {
        match waiting {
            Waiting::Text {
                id,
                made_from,
                text,
            } => match input_refusal(text, limit) {
                None => sorted.texts.push((id, *made_from, text)),
                Some(ItemCode::EmbedInputEmpty) => sorted.empty_input.push(id.clone()),
                Some(_) => {
                    sorted.too_long.push(id.clone());
                }
            },
            Waiting::Removed { id } => sorted.removed.push(id.clone()),
        }
    }
    sorted
}

/// One page of the job on `branch`'s `collection.field`: pending, made,
/// committed.
///
/// # Errors
/// Whatever reading the walk or committing the page answers; a worker's
/// failure is not an error but [`Step::Failed`].
pub fn step(
    storage: &mut Storage,
    branch: &str,
    collection: &str,
    field: &str,
    worker: &dyn Embed,
    bounds: Bounds,
) -> Result<Step, Failure> {
    let pending = storage.embed_pending(
        branch,
        collection,
        field,
        bounds.page_records,
        bounds.input_bytes,
    )?;
    let Some(pending) = pending else {
        return Ok(Step::Idle);
    };
    let at = storage.capture(&Selection {
        branch: branch.to_owned(),
        version: None,
    })?;
    let (declaration, _companion) = storage.embedded_field(&at, collection, field)?;
    let made = make(&pending, collection, field, &declaration, worker, bounds);
    let made = match made {
        Ok(made) => made,
        Err(failure) => return Ok(Step::Failed(failure)),
    };
    let page = &made.page;
    storage.embed_commit(branch, page, bounds.page_records)?;
    Ok(Step::Committed(Counts {
        made: page.made.len(),
        removed: page.removed.len(),
        no_text: pending.no_text.len(),
        too_long: made.too_long.len(),
        // EMB-E2 (i): a text past the cap is one the model reads part of.
        longer_than_model: made
            .longer_than_model
            .len()
            .saturating_add(made.windows_over_cap.len()),
        empty_input: made.empty_input.len(),
        windowed: page.windowed.len(),
        windows_over_cap: made.windows_over_cap.len(),
        ended: page.reach == Reach::End,
    }))
}

#[cfg(test)]
#[path = "job_tests.rs"]
mod tests;