kime_tensor/
blob.rs

1//! Read only bytes that are either a mapped file or owned memory.
2//!
3//! Weights are mapped rather than read (spec/07-engine.md), so opening a checkpoint costs page
4//! table entries and not a copy. Mapping needs `unsafe`, which is why this lives here and not in
5//! `kime-model`: the parsers there only ever see a `&[u8]`.
6
7use std::fs::File;
8use std::io;
9use std::ops::Deref;
10use std::path::Path;
11
12/// Bytes a checkpoint is parsed from.
13#[derive(Debug)]
14pub struct Blob(Inner);
15
16#[derive(Debug)]
17enum Inner {
18    Mapped(memmap2::Mmap),
19    Owned(Vec<u8>),
20}
21
22impl Blob {
23    /// Maps `path` read only. Empty files are returned as owned, because mapping zero bytes fails
24    /// on some platforms.
25    ///
26    /// # Errors
27    ///
28    /// Any error from opening or mapping the file.
29    pub fn map(path: impl AsRef<Path>) -> io::Result<Self> {
30        let file = File::open(path)?;
31        if file.metadata()?.len() == 0 {
32            return Ok(Self(Inner::Owned(Vec::new())));
33        }
34        // SAFETY: the mapping is read only and kime never writes through it. The remaining hazard
35        // is another process truncating or rewriting the file while it is mapped, which the OS
36        // turns into a signal or changed bytes rather than a Rust level data race on our side.
37        // Every loader treats the bytes as untrusted and checks bounds before each read, so changed
38        // bytes can give wrong weights but not an out of bounds access. This is the same contract
39        // safetensors, candle and llama.cpp use for weight files.
40        let map = unsafe { memmap2::Mmap::map(&file)? };
41        Ok(Self(Inner::Mapped(map)))
42    }
43
44    /// Wraps bytes already in memory, for tests and for files that arrive over the network.
45    #[must_use]
46    pub fn owned(bytes: Vec<u8>) -> Self {
47        Self(Inner::Owned(bytes))
48    }
49
50    /// Whether the bytes come from a mapping.
51    #[must_use]
52    pub fn is_mapped(&self) -> bool {
53        matches!(self.0, Inner::Mapped(_))
54    }
55}
56
57impl Deref for Blob {
58    type Target = [u8];
59
60    fn deref(&self) -> &[u8] {
61        match &self.0 {
62            Inner::Mapped(m) => m,
63            Inner::Owned(v) => v,
64        }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn maps_a_file() {
74        let path = std::env::temp_dir().join(format!("kime-blob-{}", std::process::id()));
75        std::fs::write(&path, b"hello").unwrap();
76        let b = Blob::map(&path).unwrap();
77        assert!(b.is_mapped());
78        assert_eq!(&*b, b"hello");
79        drop(b);
80        std::fs::write(&path, b"").unwrap();
81        assert_eq!(Blob::map(&path).unwrap().len(), 0);
82        std::fs::remove_file(path).unwrap();
83    }
84}
kime_tensor/blob.rs

kime_tensor/
blob.rs
