haematite-cli 0.21.1

//! The required embedding settings and the optional writer wait in serve.json.

use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use haem_frames::embedding::{Limits, PROTOCOL, PROTOCOL_WINDOWS, Start, Windows, check_start};

use crate::serve::query::Setting;
use crate::serve::storage::embedded::WINDOWS_MOST;

pub const NAMES: [&str; 14] = [
    "embed_worker",
    "embed_model_dir",
    "embed_model_sha256",
    "embed_memory_bytes",
    "embed_threads",
    "embed_batch_items",
    "embed_input_bytes",
    "embed_tokens",
    "embed_frame_bytes",
    "embed_ready_ms",
    "embed_call_ms",
    "embed_waiting",
    "embed_request_items",
    "embed_page_records",
];

/// The most records one page of the embedding walk reads and commits (EW1-fix).
///
/// A page holds its texts (each at most `embed_input_bytes`) and its vectors
/// in memory, reads under held snapshots, and is lost whole to a worker
/// failure: 1024 records of 16 KiB text and 768 members is 16 MiB of text
/// and 3 MiB of vectors, about 22 seconds of the worker's measured 47 a
/// second, and every shard of a store touched by one commit.
///
/// EMB-E2 (h): with windows on, a record may hold up to `max_windows`
/// vectors, so a page's vectors are at most `page_records` × `max_windows`;
/// each worker frame bounds them by `embed_frame_bytes`
/// (`check_model_limits_windowed`).
pub const PAGE_RECORDS_MOST: usize = 1024;

/// The keys as serve.json gives them; all absent turns embedding off.
#[derive(Clone, Debug, Default)]
pub struct Options {
    pub worker: Setting<PathBuf>,
    pub model_dir: Setting<PathBuf>,
    pub model_sha256: Setting<String>,
    pub memory_bytes: Setting<u64>,
    pub threads: Setting<u64>,
    pub batch_items: Setting<u64>,
    pub input_bytes: Setting<u64>,
    pub tokens: Setting<u64>,
    pub frame_bytes: Setting<u32>,
    pub ready_ms: Setting<u64>,
    pub call_ms: Setting<u64>,
    pub waiting: Setting<u64>,
    pub request_items: Setting<u64>,
    pub page_records: Setting<u64>,
    pub writer_wait_ms: Setting<u64>,
    /// EMB-E2: `embed_windows`; off when absent.
    pub windows: Setting<bool>,
    /// EMB-E2: `embed_max_windows`; [`MAX_WINDOWS`] when absent.
    pub max_windows: Setting<u64>,
    /// EMB-E2: `embed_window_overlap_tokens`; [`OVERLAP_TOKENS`] when absent.
    pub overlap_tokens: Setting<u64>,
}

/// EMB-E2 (ruled 6 Oct 21:19): the most windows one text is embedded as.
pub const MAX_WINDOWS: u64 = 32;

/// EMB-E2 (ruled 6 Oct 21:19): the tokens consecutive windows share.
pub const OVERLAP_TOKENS: u64 = 64;

/// The worker, the Start it is sent, and the service's own bounds on it.
#[derive(Clone, Debug)]
pub struct Settings {
    pub worker: PathBuf,
    pub model_dir: String,
    /// The model's manifest a worker's Ready must name, or it is refused.
    pub model_sha256: String,
    pub limits: Limits,
    /// `limits.batch_items` as a count of texts.
    pub batch: NonZeroUsize,
    /// From Start until Ready.
    pub ready: Duration,
    /// From each Embed until its Vectors.
    pub call: Duration,
    /// Callers that may wait for the one worker at once.
    pub waiting: usize,
    /// Texts one request may ask for.
    pub request_items: NonZeroUsize,
    /// Records one page of the walk reads and commits once, from 1 to
    /// [`PAGE_RECORDS_MOST`].
    pub page_records: NonZeroUsize,
    pub writer_wait: Duration,
    /// EMB-E2: how a text longer than the model reads is embedded whole;
    /// `None` keeps EMB1, where it gets no vector.
    pub windows: Option<Windows>,
}

impl Settings {
    /// EMB-E2 (h): the most one worker call of documents may take.
    ///
    /// With windows on, a text may answer up to `max_windows` vectors, so the
    /// call is given `embed_call_ms` for each: a page's calls stay within
    /// `embed_call_ms` × ⌈(texts + windows) / batch⌉ at their most.
    #[must_use]
    pub fn document_call(&self) -> Duration {
        let windows = self.windows.map_or(1, |windows| windows.max_windows);
        self.call
            .saturating_mul(u32::try_from(windows).unwrap_or(u32::MAX))
    }
}

fn required<T>(setting: Setting<T>, name: &str) -> Result<T, String> {
    match setting {
        Setting::Present(value) => Ok(value),
        Setting::Absent => Err(format!("missing embed setting {name}")),
    }
}

fn count(setting: Setting<u64>, name: &str) -> Result<NonZeroUsize, String> {
    let value = required(setting, name)?;
    usize::try_from(value)
        .ok()
        .and_then(NonZeroUsize::new)
        .ok_or_else(|| format!("embed setting {name} must be a positive count"))
}

fn page_records(setting: Setting<u64>) -> Result<NonZeroUsize, String> {
    let records = count(setting, NAMES[13])?;
    if records.get() > PAGE_RECORDS_MOST {
        return Err(format!(
            "embed setting {} must be from 1 to {PAGE_RECORDS_MOST}",
            NAMES[13]
        ));
    }
    Ok(records)
}

fn millis(setting: Setting<u64>, name: &str) -> Result<Duration, String> {
    let duration = Duration::from_millis(required(setting, name)?);
    if duration.is_zero() || Instant::now().checked_add(duration).is_none() {
        return Err(format!(
            "embed setting {name} must be positive and within the clock's range"
        ));
    }
    Ok(duration)
}

impl Options {
    /// The settings, `None` when no key is given.
    ///
    /// # Errors
    ///
    /// Names every missing key when only some are given, and any key whose
    /// value the worker's Start or the service cannot take.
    pub fn resolve(self) -> Result<Option<Settings>, String> {
        let present = [
            self.worker.is_set(),
            self.model_dir.is_set(),
            self.model_sha256.is_set(),
            self.memory_bytes.is_set(),
            self.threads.is_set(),
            self.batch_items.is_set(),
            self.input_bytes.is_set(),
            self.tokens.is_set(),
            self.frame_bytes.is_set(),
            self.ready_ms.is_set(),
            self.call_ms.is_set(),
            self.waiting.is_set(),
            self.request_items.is_set(),
            self.page_records.is_set(),
        ];
        // A window setting alone is named missing settings, not ignored.
        let windowing = [
            self.windows.is_set(),
            self.max_windows.is_set(),
            self.overlap_tokens.is_set(),
        ];
        if present.iter().chain(&windowing).all(|set| !set) && !self.writer_wait_ms.is_set() {
            return Ok(None);
        }
        let missing: Vec<_> = NAMES
            .into_iter()
            .zip(present)
            .filter_map(|(name, set)| (!set).then_some(name))
            .collect();
        if !missing.is_empty() {
            return Err(format!("missing embed settings: {}", missing.join(", ")));
        }
        let worker = required(self.worker, NAMES[0])?;
        if !worker.is_absolute() {
            return Err(format!(
                "embed setting {} must be an absolute path",
                NAMES[0]
            ));
        }
        let model_dir = required(self.model_dir, NAMES[1])?
            .into_os_string()
            .into_string()
            .map_err(|_| format!("embed setting {} must be UTF-8", NAMES[1]))?;
        let limits = Limits {
            memory_bytes: required(self.memory_bytes, NAMES[3])?,
            threads: required(self.threads, NAMES[4])?,
            batch_items: required(self.batch_items, NAMES[5])?,
            input_bytes: required(self.input_bytes, NAMES[6])?,
            tokens: required(self.tokens, NAMES[7])?,
            frame_bytes: required(self.frame_bytes, NAMES[8])?,
        };
        let windows = windows(self.windows, self.max_windows, self.overlap_tokens)?;
        let start = Start {
            protocol: if windows.is_some() {
                PROTOCOL_WINDOWS
            } else {
                PROTOCOL
            },
            model_dir,
            limits,
            windows,
        };
        check_start(&start).map_err(|error| format!("embed settings: {error}"))?;
        Ok(Some(Settings {
            worker,
            model_dir: start.model_dir,
            model_sha256: required(self.model_sha256, NAMES[2])?,
            limits,
            batch: count(Setting::Present(limits.batch_items), NAMES[5])?,
            ready: millis(self.ready_ms, NAMES[9])?,
            call: millis(self.call_ms, NAMES[10])?,
            waiting: count(self.waiting, NAMES[11])?.get(),
            request_items: count(self.request_items, NAMES[12])?,
            page_records: page_records(self.page_records)?,
            writer_wait: writer_wait(self.writer_wait_ms)?,
            windows,
        }))
    }
}

/// EMB-E2: windows when `embed_windows` is true, with the cap and overlap
/// given or their ruled defaults; a cap or overlap given with windows off is
/// refused by name, so a setting is never silently ignored.
fn windows(
    on: Setting<bool>,
    max_windows: Setting<u64>,
    overlap_tokens: Setting<u64>,
) -> Result<Option<Windows>, String> {
    let given = |setting: Setting<u64>, default: u64| match setting {
        Setting::Present(value) => value,
        Setting::Absent => default,
    };
    match on {
        Setting::Present(true) => {
            let max_windows = given(max_windows, MAX_WINDOWS);
            // A record keeps at most `WINDOWS_MOST` window keys.
            if !(2..=u64::from(WINDOWS_MOST)).contains(&max_windows) {
                return Err(format!(
                    "embed setting embed_max_windows is {max_windows}; it must be 2 to \
                     {WINDOWS_MOST}"
                ));
            }
            Ok(Some(Windows {
                overlap_tokens: given(overlap_tokens, OVERLAP_TOKENS),
                max_windows,
            }))
        }
        Setting::Present(false) | Setting::Absent
            if max_windows.is_set() || overlap_tokens.is_set() =>
        {
            Err(
                "embed settings embed_max_windows and embed_window_overlap_tokens need \
                 embed_windows true"
                    .to_owned(),
            )
        }
        Setting::Present(false) | Setting::Absent => Ok(None),
    }
}

/// EWR1: the wait before the second start after a worker died or timed out;
/// each later one doubles it, up to [`RESTART_MOST`]. The first start after
/// a death is at once.
pub const RESTART: &str = "embed_restart_ms";

/// EWR1: the longest wait between two starts of the worker.
pub const RESTART_MOST: &str = "embed_restart_most_ms";

/// [`RESTART`] when serve.json does not give it: one second.
pub const RESTART_DEFAULT_MS: u64 = 1_000;

/// [`RESTART_MOST`] when serve.json does not give it: five minutes, so a
/// ward's search comes back within five minutes of the cause clearing.
pub const RESTART_MOST_DEFAULT_MS: u64 = 300_000;

/// When the driver starts the worker again after it died.
///
/// It also holds how long a text search waits for a worker that is not ready
/// (`embed_ready_ms`). There is no last start: the worker is started again for
/// as long as the service runs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Restart {
    pub first: Duration,
    pub most: Duration,
    pub ready: Duration,
}

impl Default for Restart {
    /// The defaults, and a search that does not wait.
    fn default() -> Self {
        Self {
            first: Duration::from_millis(RESTART_DEFAULT_MS),
            most: Duration::from_millis(RESTART_MOST_DEFAULT_MS),
            ready: Duration::ZERO,
        }
    }
}

impl Restart {
    /// The two keys, each its default when absent, and `ready` from
    /// `embed_ready_ms`.
    ///
    /// # Errors
    ///
    /// Names a key that is 0, out of the clock's range, or a most below the
    /// first.
    pub fn of(first: Setting<u64>, most: Setting<u64>, ready: Duration) -> Result<Self, String> {
        let first = millis(given_or(first, RESTART_DEFAULT_MS), RESTART)?;
        let most = millis(given_or(most, RESTART_MOST_DEFAULT_MS), RESTART_MOST)?;
        if most < first {
            return Err(format!(
                "embed setting {RESTART_MOST} must be at least {RESTART}"
            ));
        }
        Ok(Self { first, most, ready })
    }

    /// The wait before the start that follows the `deaths`th death in a row:
    /// none after the first, then `first`, doubling each time, never past
    /// `most`.
    #[must_use]
    pub fn wait(self, deaths: u32) -> Duration {
        let Some(doublings) = deaths.checked_sub(2) else {
            return Duration::ZERO;
        };
        1_u32
            .checked_shl(doublings)
            .and_then(|times| self.first.checked_mul(times))
            .map_or(self.most, |wait| wait.min(self.most))
    }
}

/// `setting`, or `default` when serve.json does not give it.
const fn given_or(setting: Setting<u64>, default: u64) -> Setting<u64> {
    match setting {
        Setting::Present(value) => Setting::Present(value),
        Setting::Absent => Setting::Present(default),
    }
}

/// The name of the count of document workers (box 43).
pub const WORKERS: &str = "embed_workers";

/// The most document workers one service keeps: each holds its own copy of
/// the model, and Tom's Mac was measured to eight.
pub const MAX_WORKERS: usize = 16;

/// How many document workers to keep: one when the key is absent, so a
/// serve.json written before it still starts.
///
/// # Errors
///
/// Names the key when it is 0 or over [`MAX_WORKERS`].
pub fn workers(setting: Setting<u64>) -> Result<NonZeroUsize, String> {
    let Setting::Present(value) = setting else {
        return Ok(NonZeroUsize::MIN);
    };
    usize::try_from(value)
        .ok()
        .filter(|count| *count <= MAX_WORKERS)
        .and_then(NonZeroUsize::new)
        .ok_or_else(|| format!("embed setting {WORKERS} must be from 1 to {MAX_WORKERS}"))
}

/// The writer wait when the optional setting is absent.
pub const WRITER_WAIT_DEFAULT_MS: u64 = 1_000;

/// How long foreground admission may wait for one embedding batch commit.
///
/// The absent setting uses [`WRITER_WAIT_DEFAULT_MS`]; an explicit value
/// remains the operator's bound and is never replaced when invalid.
///
/// # Errors
///
/// Names `embed_writer_wait_ms` when it is zero or outside the clock's range.
pub fn writer_wait(setting: Setting<u64>) -> Result<Duration, String> {
    let millis = match setting {
        Setting::Absent => WRITER_WAIT_DEFAULT_MS,
        Setting::Present(millis) => millis,
    };
    let wait = Duration::from_millis(millis);
    if wait.is_zero() || Instant::now().checked_add(wait).is_none() {
        return Err(
            "embed_writer_wait_ms must be positive and within the clock's range".to_owned(),
        );
    }
    Ok(wait)
}

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