metrale_grammar/compiler/
mask_snapshot.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// Cross-process persistence for the Tier-2 [`RuleLevelCache`] — issue #918.
4//
5// WHY THE *RULE*-LEVEL CACHE AND NOT THE COMPILED GRAMMAR
6// -------------------------------------------------------
7// #918 reports "~3.4 s cold grammar compile per new tool schema". The
8// CPU measurement that motivated this module (M5 Max, release, Qwen3
9// ByteLevel-BPE tokenizer, 151,669-token vocabulary, the coherency
10// gate's `get_weather` schema through `compile_qwen3_coder_tool_grammar`)
11// shows the cost is per *process*, not per schema:
12//
13//   grammar construction (EBNF + parse + normalize + optimize)   5.5 ms
14//   mask prewarm, first grammar in the process (123 masks)     621.2 ms
15//   mask prewarm, same schema again                              0.0 ms
16//   mask prewarm, a DIFFERENT schema (new tool + field names)    16.0 ms
17//
18// The second distinct schema costs ~2.5% of the first because
19// [`RuleLevelCache`] keys masks *structurally* — every JSON tool schema
20// reuses the same string/number/whitespace/punctuation sub-rules. So a
21// disk cache keyed by (tokenizer, schema) would only ever help a repeat
22// of the *same* schema, and would miss the reuse that already carries
23// 97% of the load. Persisting the rule-level entries instead makes a
24// fresh process warm for schemas it has never seen.
25//
26// Scaling that ~600 ms to the H100 host in the #918 report (248,320
27// tokens, ~1.64x the vocabulary, a slower host core than the M5) lands
28// on the ~3.2 s of cold mask generation the issue still has open.
29//
30// FORMAT
31// ------
32// A single little-endian binary file. The header pins everything a
33// cached mask depends on — format version, `usize` width (the accepted
34// bitset is written as raw `BitVec` words), the tokenizer fingerprint
35// and the vocabulary size — and a trailing FNV-1a digest covers the
36// whole body, so a truncated or corrupt file is a cache MISS rather
37// than a wrong mask. Any mismatch returns `Ok(0)`: the caller recomputes.
38
39use std::io::{Read, Write};
40use std::path::Path;
41use std::sync::Arc;
42
43use bitvec::order::Lsb0;
44use bitvec::vec::BitVec;
45
46use super::mask::{AdaptiveTokenMask, StoreType};
47use super::rule_cache::{RuleLevelCache, RuleMaskKey};
48
49/// File magic — `METRL` + "grammar masks", version 1.
50const MAGIC: &[u8; 8] = b"METRLGM1";
51/// Bumped whenever the encoding or the mask semantics change, so a
52/// snapshot written by an older build is a miss, never a mis-decode.
53const FORMAT_VERSION: u32 = 1;
54
55/// Errors from reading or writing a mask snapshot. A *stale* or
56/// corrupt snapshot is NOT an error — it is a miss (`Ok(0)`); only
57/// genuine I/O failures surface here.
58#[derive(Debug, thiserror::Error)]
59pub enum SnapshotError {
60    #[error("mask snapshot I/O failed: {0}")]
61    Io(String),
62}
63
64impl From<std::io::Error> for SnapshotError {
65    fn from(e: std::io::Error) -> Self {
66        SnapshotError::Io(e.to_string())
67    }
68}
69
70/// The identity a snapshot is only valid for: the tokenizer it was
71/// computed against, plus that tokenizer's vocabulary size.
72///
73/// A mask is a partition of the sorted decoded vocabulary, so a
74/// different tokenizer (or a different `vocab_size` cap over the same
75/// tokenizer) invalidates every entry.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub struct SnapshotIdentity {
78    /// [`crate::tokenizer::TokenizerInfo::fingerprint`].
79    pub tokenizer_fingerprint: u64,
80    /// [`crate::tokenizer::TokenizerInfo::vocab_size`].
81    pub vocab_size: usize,
82}
83
84/// Bits per `usize` on this target — part of the header because the
85/// accepted bitset is written as raw `BitVec<usize, Lsb0>` words.
86const fn usize_bits() -> u32 {
87    usize::BITS
88}
89
90// ── FNV-1a (64-bit) ────────────────────────────────────────────────
91//
92// A fixed-seed, dependency-free digest. `ahash`'s default state is
93// randomized per process, which is exactly wrong for a value that has
94// to be stable across restarts.
95
96const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
97const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
98
99/// Fold `bytes` into an FNV-1a accumulator.
100pub fn fnv1a(mut hash: u64, bytes: &[u8]) -> u64 {
101    for &b in bytes {
102        hash ^= b as u64;
103        hash = hash.wrapping_mul(FNV_PRIME);
104    }
105    hash
106}
107
108/// FNV-1a over `bytes`, starting from the standard offset basis.
109pub fn fnv1a_of(bytes: &[u8]) -> u64 {
110    fnv1a(FNV_OFFSET, bytes)
111}
112
113// ── Encoding ───────────────────────────────────────────────────────
114
115/// A byte sink that digests everything it is handed, so the trailing
116/// checksum never needs a second pass over the buffer.
117struct Writer {
118    buf: Vec<u8>,
119    digest: u64,
120}
121
122impl Writer {
123    fn new() -> Self {
124        Self {
125            buf: Vec::new(),
126            digest: FNV_OFFSET,
127        }
128    }
129    fn bytes(&mut self, b: &[u8]) {
130        self.digest = fnv1a(self.digest, b);
131        self.buf.extend_from_slice(b);
132    }
133    fn u8(&mut self, v: u8) {
134        self.bytes(&[v]);
135    }
136    fn u32(&mut self, v: u32) {
137        self.bytes(&v.to_le_bytes());
138    }
139    fn u64(&mut self, v: u64) {
140        self.bytes(&v.to_le_bytes());
141    }
142    fn i32(&mut self, v: i32) {
143        self.bytes(&v.to_le_bytes());
144    }
145    fn i32_slice(&mut self, v: &[i32]) {
146        self.u32(v.len() as u32);
147        for &x in v {
148            self.i32(x);
149        }
150    }
151}
152
153/// A checked byte source. Every read is bounds-guarded; a short read
154/// yields `None`, which the caller turns into a cache miss.
155struct Reader<'a> {
156    buf: &'a [u8],
157    pos: usize,
158}
159
160impl<'a> Reader<'a> {
161    fn new(buf: &'a [u8]) -> Self {
162        Self { buf, pos: 0 }
163    }
164    fn take(&mut self, n: usize) -> Option<&'a [u8]> {
165        let end = self.pos.checked_add(n)?;
166        let out = self.buf.get(self.pos..end)?;
167        self.pos = end;
168        Some(out)
169    }
170    fn u8(&mut self) -> Option<u8> {
171        Some(self.take(1)?[0])
172    }
173    fn u32(&mut self) -> Option<u32> {
174        Some(u32::from_le_bytes(self.take(4)?.try_into().ok()?))
175    }
176    fn u64(&mut self) -> Option<u64> {
177        Some(u64::from_le_bytes(self.take(8)?.try_into().ok()?))
178    }
179    fn i32(&mut self) -> Option<i32> {
180        Some(i32::from_le_bytes(self.take(4)?.try_into().ok()?))
181    }
182    fn i32_vec(&mut self) -> Option<Vec<i32>> {
183        let n = self.u32()? as usize;
184        let mut out = Vec::with_capacity(n.min(1 << 20));
185        for _ in 0..n {
186            out.push(self.i32()?);
187        }
188        Some(out)
189    }
190}
191
192fn store_tag(t: StoreType) -> u8 {
193    match t {
194        StoreType::Accepted => 0,
195        StoreType::Rejected => 1,
196        StoreType::AcceptedBitset => 2,
197    }
198}
199
200fn store_from_tag(tag: u8) -> Option<StoreType> {
201    match tag {
202        0 => Some(StoreType::Accepted),
203        1 => Some(StoreType::Rejected),
204        2 => Some(StoreType::AcceptedBitset),
205        _ => None,
206    }
207}
208
209/// Serialize `entries` into the snapshot byte format.
210pub fn encode(
211    identity: SnapshotIdentity,
212    entries: &[(RuleMaskKey, Arc<AdaptiveTokenMask>)],
213) -> Vec<u8> {
214    let mut w = Writer::new();
215    w.bytes(MAGIC);
216    w.u32(FORMAT_VERSION);
217    w.u32(usize_bits());
218    w.u64(identity.tokenizer_fingerprint);
219    w.u64(identity.vocab_size as u64);
220    w.u64(entries.len() as u64);
221    for (key, mask) in entries {
222        w.u64(key.fsm_hash);
223        w.i32(key.fsm_new_node_id);
224        w.i32(key.state_cnt);
225        w.i32(key.edge_cnt);
226        w.u8(store_tag(mask.store_type));
227        w.i32_slice(&mask.accepted_indices);
228        w.i32_slice(&mask.rejected_indices);
229        w.i32_slice(&mask.uncertain_indices);
230        let words = mask.accepted_bitset.as_raw_slice();
231        w.u64(mask.accepted_bitset.len() as u64);
232        w.u64(words.len() as u64);
233        for &word in words {
234            w.u64(word as u64);
235        }
236    }
237    let digest = w.digest;
238    w.buf.extend_from_slice(&digest.to_le_bytes());
239    w.buf
240}
241
242/// Parse a snapshot. Returns `None` for ANY reason the bytes are not
243/// usable under `identity` — wrong magic, older format, different
244/// `usize` width, different tokenizer, truncation, checksum mismatch.
245/// Those are cache misses, not failures.
246pub fn decode(
247    bytes: &[u8],
248    identity: SnapshotIdentity,
249) -> Option<Vec<(RuleMaskKey, Arc<AdaptiveTokenMask>)>> {
250    let body_len = bytes.len().checked_sub(8)?;
251    let (body, tail) = bytes.split_at(body_len);
252    if fnv1a_of(body) != u64::from_le_bytes(tail.try_into().ok()?) {
253        return None;
254    }
255    let mut r = Reader::new(body);
256    if r.take(8)? != MAGIC
257        || r.u32()? != FORMAT_VERSION
258        || r.u32()? != usize_bits()
259        || r.u64()? != identity.tokenizer_fingerprint
260        || r.u64()? != identity.vocab_size as u64
261    {
262        return None;
263    }
264    let count = r.u64()? as usize;
265    let mut out = Vec::with_capacity(count.min(1 << 16));
266    for _ in 0..count {
267        let key = RuleMaskKey {
268            fsm_hash: r.u64()?,
269            fsm_new_node_id: r.i32()?,
270            state_cnt: r.i32()?,
271            edge_cnt: r.i32()?,
272        };
273        let store_type = store_from_tag(r.u8()?)?;
274        let accepted_indices = r.i32_vec()?;
275        let rejected_indices = r.i32_vec()?;
276        let uncertain_indices = r.i32_vec()?;
277        let bit_len = r.u64()? as usize;
278        let word_count = r.u64()? as usize;
279        let mut words: Vec<usize> = Vec::with_capacity(word_count.min(1 << 20));
280        for _ in 0..word_count {
281            words.push(r.u64()? as usize);
282        }
283        let mut accepted_bitset: BitVec<usize, Lsb0> = BitVec::from_vec(words);
284        if accepted_bitset.len() < bit_len {
285            return None;
286        }
287        accepted_bitset.truncate(bit_len);
288        out.push((
289            key,
290            Arc::new(AdaptiveTokenMask {
291                store_type,
292                accepted_indices,
293                rejected_indices,
294                accepted_bitset,
295                uncertain_indices,
296            }),
297        ));
298    }
299    // Trailing bytes mean the writer and reader disagree about the
300    // format; refuse rather than half-trust it.
301    if r.pos != body.len() {
302        return None;
303    }
304    Some(out)
305}
306
307// ── File-level helpers ─────────────────────────────────────────────
308
309/// Write `cache`'s entries to `path`, atomically (temp file + rename)
310/// so a concurrent reader never observes a half-written snapshot.
311///
312/// At most `max_entries` are written, taken from the MOST-recently-used
313/// end of the cache's LRU order. The rule cache is bounded by a memory
314/// budget (~1/3 of a GiB by default), and a snapshot that large has no
315/// business sitting next to a checkpoint: measured mask footprint is
316/// ~19 KB/mask at a 151,669-token vocabulary, so the cap is what keeps
317/// the file in the tens of megabytes.
318///
319/// Returns the number of entries written.
320pub fn save_to_file(
321    cache: &RuleLevelCache,
322    identity: SnapshotIdentity,
323    path: &Path,
324    max_entries: usize,
325) -> Result<usize, SnapshotError> {
326    let all = cache.entries();
327    let entries = &all[all.len().saturating_sub(max_entries)..];
328    let bytes = encode(identity, entries);
329    if let Some(dir) = path.parent() {
330        std::fs::create_dir_all(dir)?;
331    }
332    // The pid keeps two servers starting at once from colliding on the
333    // temp name; the rename is what makes the visible file atomic.
334    let tmp = path.with_extension(format!("tmp{}", std::process::id()));
335    {
336        let mut f = std::fs::File::create(&tmp)?;
337        f.write_all(&bytes)?;
338        f.sync_all()?;
339    }
340    std::fs::rename(&tmp, path)?;
341    Ok(entries.len())
342}
343
344/// Load `path` into `cache`. Returns the number of entries imported —
345/// `0` when the file is absent, stale, or corrupt (a miss).
346///
347/// Entries already present in `cache` win: [`RuleLevelCache::add`]
348/// rejects a duplicate key, so loading never overwrites a mask this
349/// process computed itself.
350pub fn load_from_file(
351    cache: &RuleLevelCache,
352    identity: SnapshotIdentity,
353    path: &Path,
354) -> Result<usize, SnapshotError> {
355    let mut bytes = Vec::new();
356    match std::fs::File::open(path) {
357        Ok(mut f) => f.read_to_end(&mut bytes)?,
358        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(0),
359        Err(e) => return Err(e.into()),
360    };
361    let Some(entries) = decode(&bytes, identity) else {
362        return Ok(0);
363    };
364    let mut imported = 0usize;
365    for (key, mask) in entries {
366        if cache.add(key, mask) {
367            imported += 1;
368        }
369    }
370    Ok(imported)
371}
372
373#[cfg(test)]
374#[path = "mask_snapshot_tests.rs"]
375mod tests;
metrale_grammar/compiler/mask_snapshot.rs

metrale_grammar/compiler/
mask_snapshot.rs
