superfluid_engine/
artifact.rs

1//! What a runtime reads and what it runs on, as plain data.
2
3use std::path::Path;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct Format {
7    pub id: String,
8    pub describe: String,
9    pub rule: Rule,
10}
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub enum Rule {
14    File { extensions: Vec<String>, magic: Option<Vec<u8>> },
15    Directory { files: Vec<String>, extensions: Vec<String> },
16}
17
18impl Format {
19    pub fn file(id: &str, describe: &str, extensions: &[&str], magic: Option<&[u8]>) -> Format {
20        Format {
21            id: id.into(),
22            describe: describe.into(),
23            rule: Rule::File {
24                extensions: extensions.iter().map(|e| e.to_string()).collect(),
25                magic: magic.map(<[u8]>::to_vec),
26            },
27        }
28    }
29
30    pub fn directory(id: &str, describe: &str, files: &[&str], extensions: &[&str]) -> Format {
31        Format {
32            id: id.into(),
33            describe: describe.into(),
34            rule: Rule::Directory {
35                files: files.iter().map(|f| f.to_string()).collect(),
36                extensions: extensions.iter().map(|e| e.to_string()).collect(),
37            },
38        }
39    }
40
41    pub fn matches(&self, path: &Path) -> bool {
42        match &self.rule {
43            Rule::File { extensions, magic } => {
44                let named = path
45                    .extension()
46                    .and_then(|e| e.to_str())
47                    .is_some_and(|e| extensions.iter().any(|x| x.eq_ignore_ascii_case(e)));
48                named || magic.as_deref().is_some_and(|m| starts_with(path, m))
49            }
50            Rule::Directory { files, extensions } => {
51                path.is_dir()
52                    && files.iter().all(|f| path.join(f).is_file())
53                    && std::fs::read_dir(path)
54                        .map(|rd| {
55                            rd.flatten().any(|e| {
56                                e.path()
57                                    .extension()
58                                    .and_then(|x| x.to_str())
59                                    .is_some_and(|x| extensions.iter().any(|want| want == x))
60                            })
61                        })
62                        .unwrap_or(false)
63            }
64        }
65    }
66}
67
68fn starts_with(path: &Path, magic: &[u8]) -> bool {
69    if magic.is_empty() || !std::fs::metadata(path).is_ok_and(|m| m.is_file()) {
70        return false;
71    }
72    let mut head = vec![0u8; magic.len()];
73    std::fs::File::open(path)
74        .and_then(|mut f| std::io::Read::read_exact(&mut f, &mut head))
75        .map(|_| head == magic)
76        .unwrap_or(false)
77}
78
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub enum TokenizerSource {
81    Library(String),
82    HuggingFace,
83    Engine,
84}
85
86pub fn valid_id(s: &str) -> bool {
87    !s.is_empty()
88        && !s.starts_with('.')
89        && s.chars().all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || matches!(c, '.' | '-' | '_'))
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct RuntimeInfo {
94    pub id: &'static str,
95    pub aliases: &'static [&'static str],
96    pub formats: Vec<Format>,
97    pub tokenizer: TokenizerSource,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct Device {
102    pub backend: String,
103    pub name: String,
104    /// Its memory as the runtime reads it: llama.cpp's device total, MLX's recommended working set, baseRT's own budget.
105    pub memory: u64,
106    /// What the runtime plans a model and its cache against here, the same rule for every runtime: 0 when it plans nothing on it.
107    pub budget: u64,
108}
109
110impl Device {
111    /// The device and its budget, as `superfluid runtimes` prints it.
112    pub fn summary(&self) -> String {
113        let mut s = format!("{}: {}", self.backend, self.name);
114        if self.budget > 0 {
115            s.push_str(&format!(", {:.1} GB", self.budget as f64 / 1e9));
116        }
117        s
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use std::path::PathBuf;
125
126    fn dir(tag: &str) -> PathBuf {
127        let d = std::env::temp_dir().join(format!("superfluid-artifact-{tag}-{}", std::process::id()));
128        let _ = std::fs::remove_dir_all(&d);
129        std::fs::create_dir_all(&d).unwrap();
130        d
131    }
132
133    #[test]
134    fn a_file_format_matches_by_extension_in_any_case_or_by_magic() {
135        let d = dir("file");
136        let gguf = Format::file("gguf", "a GGUF file", &["gguf"], Some(b"GGUF"));
137        std::fs::write(d.join("a.gguf"), b"").unwrap();
138        std::fs::write(d.join("a.GGUF"), b"").unwrap();
139        std::fs::write(d.join("weights.bin"), b"GGUF\x03\x00\x00\x00").unwrap();
140        std::fs::write(d.join("a.base"), b"BASE").unwrap();
141        assert!(gguf.matches(&d.join("a.gguf")));
142        assert!(gguf.matches(&d.join("a.GGUF")));
143        assert!(gguf.matches(&d.join("weights.bin")), "the magic names a file its name does not");
144        assert!(!gguf.matches(&d.join("a.base")));
145        assert!(!gguf.matches(&d.join("missing.bin")));
146        assert!(!gguf.matches(&d), "a directory has no magic");
147        let _ = std::fs::remove_dir_all(&d);
148    }
149
150    #[test]
151    fn a_directory_format_needs_its_files_and_one_weight() {
152        let d = dir("dir");
153        let mlx = Format::directory("mlx", "an MLX model directory", &["config.json"], &["safetensors"]);
154        assert!(!mlx.matches(&d));
155        std::fs::write(d.join("config.json"), b"{}").unwrap();
156        assert!(!mlx.matches(&d), "no weights");
157        std::fs::write(d.join("model.safetensors"), b"").unwrap();
158        assert!(mlx.matches(&d));
159        assert!(!mlx.matches(&d.join("config.json")), "a file is not a directory");
160        let _ = std::fs::remove_dir_all(&d);
161    }
162
163    #[test]
164    fn the_magic_probe_opens_only_regular_files() {
165        let d = dir("fifo");
166        let fifo = d.join("pipe");
167        let c = std::ffi::CString::new(fifo.to_string_lossy().as_bytes()).unwrap();
168        extern "C" {
169            fn mkfifo(path: *const std::ffi::c_char, mode: u32) -> i32;
170        }
171        // SAFETY: a valid NUL-terminated path; mkfifo creates the node and
172        // touches nothing else.
173        assert_eq!(unsafe { mkfifo(c.as_ptr(), 0o600) }, 0, "mkfifo");
174        let gguf = Format::file("gguf", "a GGUF file", &["gguf"], Some(b"GGUF"));
175        assert!(!gguf.matches(&fifo));
176        let _ = std::fs::remove_dir_all(&d);
177    }
178
179    #[test]
180    fn a_device_sums_itself_up() {
181        let d = Device { backend: "Metal".into(), name: "Apple M5 Pro".into(), memory: 40_000_000_000, budget: 34_000_000_000 };
182        assert_eq!(d.summary(), "Metal: Apple M5 Pro, 34.0 GB", "the budget, not what the device has");
183        let cpu = Device { backend: "CPU".into(), name: "Apple M5 Pro".into(), memory: 51_000_000_000, budget: 0 };
184        assert_eq!(cpu.summary(), "CPU: Apple M5 Pro", "nothing is planned on it");
185    }
186}
superfluid_engine/artifact.rs

superfluid_engine/
artifact.rs
