haematite_cli/serve/embed/
settings.rs

1//! The required embedding settings and the optional writer wait in serve.json.
2
3use std::num::NonZeroUsize;
4use std::path::PathBuf;
5use std::time::{Duration, Instant};
6
7use haem_frames::embedding::{Limits, PROTOCOL, PROTOCOL_WINDOWS, Start, Windows, check_start};
8
9use crate::serve::query::Setting;
10use crate::serve::storage::embedded::WINDOWS_MOST;
11
12pub const NAMES: [&str; 14] = [
13    "embed_worker",
14    "embed_model_dir",
15    "embed_model_sha256",
16    "embed_memory_bytes",
17    "embed_threads",
18    "embed_batch_items",
19    "embed_input_bytes",
20    "embed_tokens",
21    "embed_frame_bytes",
22    "embed_ready_ms",
23    "embed_call_ms",
24    "embed_waiting",
25    "embed_request_items",
26    "embed_page_records",
27];
28
29/// The most records one page of the embedding walk reads and commits (EW1-fix).
30///
31/// A page holds its texts (each at most `embed_input_bytes`) and its vectors
32/// in memory, reads under held snapshots, and is lost whole to a worker
33/// failure: 1024 records of 16 KiB text and 768 members is 16 MiB of text
34/// and 3 MiB of vectors, about 22 seconds of the worker's measured 47 a
35/// second, and every shard of a store touched by one commit.
36///
37/// EMB-E2 (h): with windows on, a record may hold up to `max_windows`
38/// vectors, so a page's vectors are at most `page_records` × `max_windows`;
39/// each worker frame bounds them by `embed_frame_bytes`
40/// (`check_model_limits_windowed`).
41pub const PAGE_RECORDS_MOST: usize = 1024;
42
43/// The keys as serve.json gives them; all absent turns embedding off.
44#[derive(Clone, Debug, Default)]
45pub struct Options {
46    pub worker: Setting<PathBuf>,
47    pub model_dir: Setting<PathBuf>,
48    pub model_sha256: Setting<String>,
49    pub memory_bytes: Setting<u64>,
50    pub threads: Setting<u64>,
51    pub batch_items: Setting<u64>,
52    pub input_bytes: Setting<u64>,
53    pub tokens: Setting<u64>,
54    pub frame_bytes: Setting<u32>,
55    pub ready_ms: Setting<u64>,
56    pub call_ms: Setting<u64>,
57    pub waiting: Setting<u64>,
58    pub request_items: Setting<u64>,
59    pub page_records: Setting<u64>,
60    pub writer_wait_ms: Setting<u64>,
61    /// EMB-E2: `embed_windows`; off when absent.
62    pub windows: Setting<bool>,
63    /// EMB-E2: `embed_max_windows`; [`MAX_WINDOWS`] when absent.
64    pub max_windows: Setting<u64>,
65    /// EMB-E2: `embed_window_overlap_tokens`; [`OVERLAP_TOKENS`] when absent.
66    pub overlap_tokens: Setting<u64>,
67}
68
69/// EMB-E2 (ruled 6 Oct 21:19): the most windows one text is embedded as.
70pub const MAX_WINDOWS: u64 = 32;
71
72/// EMB-E2 (ruled 6 Oct 21:19): the tokens consecutive windows share.
73pub const OVERLAP_TOKENS: u64 = 64;
74
75/// The worker, the Start it is sent, and the service's own bounds on it.
76#[derive(Clone, Debug)]
77pub struct Settings {
78    pub worker: PathBuf,
79    pub model_dir: String,
80    /// The model's manifest a worker's Ready must name, or it is refused.
81    pub model_sha256: String,
82    pub limits: Limits,
83    /// `limits.batch_items` as a count of texts.
84    pub batch: NonZeroUsize,
85    /// From Start until Ready.
86    pub ready: Duration,
87    /// From each Embed until its Vectors.
88    pub call: Duration,
89    /// Callers that may wait for the one worker at once.
90    pub waiting: usize,
91    /// Texts one request may ask for.
92    pub request_items: NonZeroUsize,
93    /// Records one page of the walk reads and commits once, from 1 to
94    /// [`PAGE_RECORDS_MOST`].
95    pub page_records: NonZeroUsize,
96    pub writer_wait: Duration,
97    /// EMB-E2: how a text longer than the model reads is embedded whole;
98    /// `None` keeps EMB1, where it gets no vector.
99    pub windows: Option<Windows>,
100}
101
102impl Settings {
103    /// EMB-E2 (h): the most one worker call of documents may take.
104    ///
105    /// With windows on, a text may answer up to `max_windows` vectors, so the
106    /// call is given `embed_call_ms` for each: a page's calls stay within
107    /// `embed_call_ms` × ⌈(texts + windows) / batch⌉ at their most.
108    #[must_use]
109    pub fn document_call(&self) -> Duration {
110        let windows = self.windows.map_or(1, |windows| windows.max_windows);
111        self.call
112            .saturating_mul(u32::try_from(windows).unwrap_or(u32::MAX))
113    }
114}
115
116fn required<T>(setting: Setting<T>, name: &str) -> Result<T, String> {
117    match setting {
118        Setting::Present(value) => Ok(value),
119        Setting::Absent => Err(format!("missing embed setting {name}")),
120    }
121}
122
123fn count(setting: Setting<u64>, name: &str) -> Result<NonZeroUsize, String> {
124    let value = required(setting, name)?;
125    usize::try_from(value)
126        .ok()
127        .and_then(NonZeroUsize::new)
128        .ok_or_else(|| format!("embed setting {name} must be a positive count"))
129}
130
131fn page_records(setting: Setting<u64>) -> Result<NonZeroUsize, String> {
132    let records = count(setting, NAMES[13])?;
133    if records.get() > PAGE_RECORDS_MOST {
134        return Err(format!(
135            "embed setting {} must be from 1 to {PAGE_RECORDS_MOST}",
136            NAMES[13]
137        ));
138    }
139    Ok(records)
140}
141
142fn millis(setting: Setting<u64>, name: &str) -> Result<Duration, String> {
143    let duration = Duration::from_millis(required(setting, name)?);
144    if duration.is_zero() || Instant::now().checked_add(duration).is_none() {
145        return Err(format!(
146            "embed setting {name} must be positive and within the clock's range"
147        ));
148    }
149    Ok(duration)
150}
151
152impl Options {
153    /// The settings, `None` when no key is given.
154    ///
155    /// # Errors
156    ///
157    /// Names every missing key when only some are given, and any key whose
158    /// value the worker's Start or the service cannot take.
159    pub fn resolve(self) -> Result<Option<Settings>, String> {
160        let present = [
161            self.worker.is_set(),
162            self.model_dir.is_set(),
163            self.model_sha256.is_set(),
164            self.memory_bytes.is_set(),
165            self.threads.is_set(),
166            self.batch_items.is_set(),
167            self.input_bytes.is_set(),
168            self.tokens.is_set(),
169            self.frame_bytes.is_set(),
170            self.ready_ms.is_set(),
171            self.call_ms.is_set(),
172            self.waiting.is_set(),
173            self.request_items.is_set(),
174            self.page_records.is_set(),
175        ];
176        // A window setting alone is named missing settings, not ignored.
177        let windowing = [
178            self.windows.is_set(),
179            self.max_windows.is_set(),
180            self.overlap_tokens.is_set(),
181        ];
182        if present.iter().chain(&windowing).all(|set| !set) && !self.writer_wait_ms.is_set() {
183            return Ok(None);
184        }
185        let missing: Vec<_> = NAMES
186            .into_iter()
187            .zip(present)
188            .filter_map(|(name, set)| (!set).then_some(name))
189            .collect();
190        if !missing.is_empty() {
191            return Err(format!("missing embed settings: {}", missing.join(", ")));
192        }
193        let worker = required(self.worker, NAMES[0])?;
194        if !worker.is_absolute() {
195            return Err(format!(
196                "embed setting {} must be an absolute path",
197                NAMES[0]
198            ));
199        }
200        let model_dir = required(self.model_dir, NAMES[1])?
201            .into_os_string()
202            .into_string()
203            .map_err(|_| format!("embed setting {} must be UTF-8", NAMES[1]))?;
204        let limits = Limits {
205            memory_bytes: required(self.memory_bytes, NAMES[3])?,
206            threads: required(self.threads, NAMES[4])?,
207            batch_items: required(self.batch_items, NAMES[5])?,
208            input_bytes: required(self.input_bytes, NAMES[6])?,
209            tokens: required(self.tokens, NAMES[7])?,
210            frame_bytes: required(self.frame_bytes, NAMES[8])?,
211        };
212        let windows = windows(self.windows, self.max_windows, self.overlap_tokens)?;
213        let start = Start {
214            protocol: if windows.is_some() {
215                PROTOCOL_WINDOWS
216            } else {
217                PROTOCOL
218            },
219            model_dir,
220            limits,
221            windows,
222        };
223        check_start(&start).map_err(|error| format!("embed settings: {error}"))?;
224        Ok(Some(Settings {
225            worker,
226            model_dir: start.model_dir,
227            model_sha256: required(self.model_sha256, NAMES[2])?,
228            limits,
229            batch: count(Setting::Present(limits.batch_items), NAMES[5])?,
230            ready: millis(self.ready_ms, NAMES[9])?,
231            call: millis(self.call_ms, NAMES[10])?,
232            waiting: count(self.waiting, NAMES[11])?.get(),
233            request_items: count(self.request_items, NAMES[12])?,
234            page_records: page_records(self.page_records)?,
235            writer_wait: writer_wait(self.writer_wait_ms)?,
236            windows,
237        }))
238    }
239}
240
241/// EMB-E2: windows when `embed_windows` is true, with the cap and overlap
242/// given or their ruled defaults; a cap or overlap given with windows off is
243/// refused by name, so a setting is never silently ignored.
244fn windows(
245    on: Setting<bool>,
246    max_windows: Setting<u64>,
247    overlap_tokens: Setting<u64>,
248) -> Result<Option<Windows>, String> {
249    let given = |setting: Setting<u64>, default: u64| match setting {
250        Setting::Present(value) => value,
251        Setting::Absent => default,
252    };
253    match on {
254        Setting::Present(true) => {
255            let max_windows = given(max_windows, MAX_WINDOWS);
256            // A record keeps at most `WINDOWS_MOST` window keys.
257            if !(2..=u64::from(WINDOWS_MOST)).contains(&max_windows) {
258                return Err(format!(
259                    "embed setting embed_max_windows is {max_windows}; it must be 2 to \
260                     {WINDOWS_MOST}"
261                ));
262            }
263            Ok(Some(Windows {
264                overlap_tokens: given(overlap_tokens, OVERLAP_TOKENS),
265                max_windows,
266            }))
267        }
268        Setting::Present(false) | Setting::Absent
269            if max_windows.is_set() || overlap_tokens.is_set() =>
270        {
271            Err(
272                "embed settings embed_max_windows and embed_window_overlap_tokens need \
273                 embed_windows true"
274                    .to_owned(),
275            )
276        }
277        Setting::Present(false) | Setting::Absent => Ok(None),
278    }
279}
280
281/// EWR1: the wait before the second start after a worker died or timed out;
282/// each later one doubles it, up to [`RESTART_MOST`]. The first start after
283/// a death is at once.
284pub const RESTART: &str = "embed_restart_ms";
285
286/// EWR1: the longest wait between two starts of the worker.
287pub const RESTART_MOST: &str = "embed_restart_most_ms";
288
289/// [`RESTART`] when serve.json does not give it: one second.
290pub const RESTART_DEFAULT_MS: u64 = 1_000;
291
292/// [`RESTART_MOST`] when serve.json does not give it: five minutes, so a
293/// ward's search comes back within five minutes of the cause clearing.
294pub const RESTART_MOST_DEFAULT_MS: u64 = 300_000;
295
296/// When the driver starts the worker again after it died.
297///
298/// It also holds how long a text search waits for a worker that is not ready
299/// (`embed_ready_ms`). There is no last start: the worker is started again for
300/// as long as the service runs.
301#[derive(Clone, Copy, Debug, PartialEq, Eq)]
302pub struct Restart {
303    pub first: Duration,
304    pub most: Duration,
305    pub ready: Duration,
306}
307
308impl Default for Restart {
309    /// The defaults, and a search that does not wait.
310    fn default() -> Self {
311        Self {
312            first: Duration::from_millis(RESTART_DEFAULT_MS),
313            most: Duration::from_millis(RESTART_MOST_DEFAULT_MS),
314            ready: Duration::ZERO,
315        }
316    }
317}
318
319impl Restart {
320    /// The two keys, each its default when absent, and `ready` from
321    /// `embed_ready_ms`.
322    ///
323    /// # Errors
324    ///
325    /// Names a key that is 0, out of the clock's range, or a most below the
326    /// first.
327    pub fn of(first: Setting<u64>, most: Setting<u64>, ready: Duration) -> Result<Self, String> {
328        let first = millis(given_or(first, RESTART_DEFAULT_MS), RESTART)?;
329        let most = millis(given_or(most, RESTART_MOST_DEFAULT_MS), RESTART_MOST)?;
330        if most < first {
331            return Err(format!(
332                "embed setting {RESTART_MOST} must be at least {RESTART}"
333            ));
334        }
335        Ok(Self { first, most, ready })
336    }
337
338    /// The wait before the start that follows the `deaths`th death in a row:
339    /// none after the first, then `first`, doubling each time, never past
340    /// `most`.
341    #[must_use]
342    pub fn wait(self, deaths: u32) -> Duration {
343        let Some(doublings) = deaths.checked_sub(2) else {
344            return Duration::ZERO;
345        };
346        1_u32
347            .checked_shl(doublings)
348            .and_then(|times| self.first.checked_mul(times))
349            .map_or(self.most, |wait| wait.min(self.most))
350    }
351}
352
353/// `setting`, or `default` when serve.json does not give it.
354const fn given_or(setting: Setting<u64>, default: u64) -> Setting<u64> {
355    match setting {
356        Setting::Present(value) => Setting::Present(value),
357        Setting::Absent => Setting::Present(default),
358    }
359}
360
361/// The name of the count of document workers (box 43).
362pub const WORKERS: &str = "embed_workers";
363
364/// The most document workers one service keeps: each holds its own copy of
365/// the model, and Tom's Mac was measured to eight.
366pub const MAX_WORKERS: usize = 16;
367
368/// How many document workers to keep: one when the key is absent, so a
369/// serve.json written before it still starts.
370///
371/// # Errors
372///
373/// Names the key when it is 0 or over [`MAX_WORKERS`].
374pub fn workers(setting: Setting<u64>) -> Result<NonZeroUsize, String> {
375    let Setting::Present(value) = setting else {
376        return Ok(NonZeroUsize::MIN);
377    };
378    usize::try_from(value)
379        .ok()
380        .filter(|count| *count <= MAX_WORKERS)
381        .and_then(NonZeroUsize::new)
382        .ok_or_else(|| format!("embed setting {WORKERS} must be from 1 to {MAX_WORKERS}"))
383}
384
385/// The writer wait when the optional setting is absent.
386pub const WRITER_WAIT_DEFAULT_MS: u64 = 1_000;
387
388/// How long foreground admission may wait for one embedding batch commit.
389///
390/// The absent setting uses [`WRITER_WAIT_DEFAULT_MS`]; an explicit value
391/// remains the operator's bound and is never replaced when invalid.
392///
393/// # Errors
394///
395/// Names `embed_writer_wait_ms` when it is zero or outside the clock's range.
396pub fn writer_wait(setting: Setting<u64>) -> Result<Duration, String> {
397    let millis = match setting {
398        Setting::Absent => WRITER_WAIT_DEFAULT_MS,
399        Setting::Present(millis) => millis,
400    };
401    let wait = Duration::from_millis(millis);
402    if wait.is_zero() || Instant::now().checked_add(wait).is_none() {
403        return Err(
404            "embed_writer_wait_ms must be positive and within the clock's range".to_owned(),
405        );
406    }
407    Ok(wait)
408}
409
410#[cfg(test)]
411#[path = "settings_tests.rs"]
412mod tests;
haematite_cli/serve/embed/settings.rs

haematite_cli/serve/embed/
settings.rs
