haematite-cli 0.21.3

//! Box 43: several workers behind the job's trait, with in-process stand-ins.

use std::error::Error;
use std::sync::{Mutex, PoisonError};

use haem_frames::embedding::Distance;

use super::Workers;
use crate::serve::embed::job::Embed;
use crate::serve::embed::{Answer, Embedded, ModelId};
use crate::serve::storage::Failure;

type TestResult = Result<(), Box<dyn Error>>;

/// A worker that answers each text `t-{n}` with `[n]`, keeps every call's
/// texts, and fails or names another model when told to.
struct StandIn {
    calls: Mutex<Vec<Vec<String>>>,
    fails: bool,
    model: &'static str,
}

impl StandIn {
    const fn new() -> Self {
        Self {
            calls: Mutex::new(Vec::new()),
            fails: false,
            model: "stand-in",
        }
    }

    fn calls(&self) -> Vec<Vec<String>> {
        self.calls
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .clone()
    }
}

impl Embed for StandIn {
    fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
        self.calls
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .push(texts.to_vec());
        if self.fails {
            return Err(Failure::new("embed_worker_exited", "the stand-in ended"));
        }
        let answers = texts
            .iter()
            .map(|text| {
                let number = text.strip_prefix("t-").and_then(|n| n.parse::<u16>().ok());
                Answer::Vector(vec![f32::from(number.unwrap_or(u16::MAX))])
            })
            .collect();
        Ok(Embedded {
            model: model(self.model),
            answers,
        })
    }
}

fn model(name: &str) -> ModelId {
    ModelId {
        name: name.to_owned(),
        revision: "1".to_owned(),
        manifest_sha256: "abc".to_owned(),
        dimensions: 1,
        distance: Distance::Cosine,
        normalised: true,
    }
}

fn texts(count: usize) -> Vec<String> {
    (0..count).map(|index| format!("t-{index}")).collect()
}

fn stand_ins(count: usize) -> Vec<StandIn> {
    (0..count).map(|_worker| StandIn::new()).collect()
}

/// The first number of each vector, in the order given. An answer that is
/// not one vector gives none, so the count no longer matches what is asked.
fn numbers(embedded: &Embedded) -> Vec<f32> {
    embedded
        .answers
        .iter()
        .filter_map(|answer| match answer {
            Answer::Vector(vector) => vector.first().copied(),
            Answer::Longer(_) | Answer::Windows(_) | Answer::OverCap(_) => None,
        })
        .collect()
}

fn expected(count: usize) -> Vec<f32> {
    texts(count)
        .iter()
        .filter_map(|text| text.strip_prefix("t-")?.parse::<u16>().ok())
        .map(f32::from)
        .collect()
}

/// Fails while the shares are not contiguous and near-equal, or the vectors
/// are not given back in the texts' order: 10 texts over 4 workers are
/// 3, 3, 2 and 2, each worker one call.
#[test]
fn ten_texts_over_four_workers_come_back_in_order() -> TestResult {
    let workers = Workers::of(stand_ins(4))?;
    let embedded = workers.documents(&texts(10))?;
    assert_eq!(numbers(&embedded), expected(10));
    let all = texts(10);
    let seen: Vec<Vec<Vec<String>>> = workers.workers.iter().map(StandIn::calls).collect();
    let shares = [
        vec![all[0..3].to_vec()],
        vec![all[3..6].to_vec()],
        vec![all[6..8].to_vec()],
        vec![all[8..10].to_vec()],
    ];
    assert_eq!(seen, shares);
    Ok(())
}

/// Fails while a worker is given an empty share: 2 texts over 4 workers go
/// to the first two, one text each, and the other two are never called.
#[test]
fn two_texts_over_four_workers_use_two() -> TestResult {
    let workers = Workers::of(stand_ins(4))?;
    let embedded = workers.documents(&texts(2))?;
    assert_eq!(numbers(&embedded), expected(2));
    let calls: Vec<usize> = workers
        .workers
        .iter()
        .map(|worker| worker.calls().len())
        .collect();
    assert_eq!(calls, [1, 1, 0, 0]);
    Ok(())
}

/// Fails while a share's failure is not the call's: worker 2 of 4 failing
/// fails the call by its code, named in the message, and nothing is given.
#[test]
fn one_failing_worker_fails_the_call() -> TestResult {
    let mut given = stand_ins(4);
    let second = given.get_mut(1).ok_or("no second worker")?;
    second.fails = true;
    let workers = Workers::of(given)?;
    let refused = workers
        .documents(&texts(10))
        .err()
        .ok_or("the call did not fail")?;
    assert_eq!(refused.code, "embed_worker_exited");
    assert_eq!(refused.message, "worker 2 of 4: the stand-in ended");
    Ok(())
}

/// Fails while the models are not compared: one worker naming another model
/// refuses the call by name.
#[test]
fn another_model_in_one_share_is_refused() -> TestResult {
    let mut given = stand_ins(4);
    let third = given.get_mut(2).ok_or("no third worker")?;
    third.model = "another";
    let workers = Workers::of(given)?;
    let refused = workers
        .documents(&texts(10))
        .err()
        .ok_or("the call did not fail")?;
    assert_eq!(refused.code, "embed_model_mixed");
    Ok(())
}

/// Fails while one worker is not the single worker's own call: all texts in
/// one call, and its failure given unchanged, with no worker number.
#[test]
fn one_worker_is_the_single_worker() -> TestResult {
    let workers = Workers::of(stand_ins(1))?;
    let embedded = workers.documents(&texts(5))?;
    assert_eq!(numbers(&embedded), expected(5));
    let seen: Vec<Vec<Vec<String>>> = workers.workers.iter().map(StandIn::calls).collect();
    assert_eq!(seen, [vec![texts(5)]]);
    let mut failing = StandIn::new();
    failing.fails = true;
    let workers = Workers::of(vec![failing])?;
    let refused = workers
        .documents(&texts(5))
        .err()
        .ok_or("the call did not fail")?;
    assert_eq!(refused.message, "the stand-in ended");
    Ok(())
}

/// Fails while no workers make a `Workers`: an empty list is refused
/// `embed_workers_none`.
#[test]
fn no_workers_are_refused() {
    let refused = Workers::<StandIn>::of(Vec::new()).err();
    let code = refused.map(|failure| failure.code);
    assert_eq!(code.as_deref(), Some("embed_workers_none"));
}