superfluid_worker/
check.rs

1//! `<worker> check [--model <path>]`.
2
3use std::path::PathBuf;
4
5use superfluid_engine::artifact::{Device, Format, Rule, RuntimeInfo, TokenizerSource};
6use serde_json::{json, Value};
7
8pub struct ModelCheck {
9    pub path: PathBuf,
10    pub reads: Result<(), String>,
11    pub facts: Option<Result<String, String>>,
12    pub sizing: Option<Result<String, String>>,
13}
14
15/// An option the runtime takes on `superfluid serve` as `--<runtime>-<name> <value>`, and its
16/// worker as `--<name> <value>`, as the worker's report declares it.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct OptionDecl {
19    pub name: String,
20    /// `int`, `float`, `on-off`, `choice` (one of `choices`) or `text`.
21    pub kind: String,
22    pub choices: Vec<String>,
23    /// The least an `int` or a `float` takes, when it has a floor.
24    pub least: Option<i64>,
25    /// The most an `int` takes, when it has a ceiling.
26    pub most: Option<i64>,
27    /// What the runtime does without it, as words for the help.
28    pub default: String,
29    pub about: String,
30}
31
32impl OptionDecl {
33    pub fn to_json(&self) -> Value {
34        let mut v = json!({"name": self.name, "kind": self.kind, "default": self.default, "about": self.about});
35        if !self.choices.is_empty() {
36            v["choices"] = json!(self.choices);
37        }
38        if let Some(least) = self.least {
39            v["least"] = json!(least);
40        }
41        if let Some(most) = self.most {
42            v["most"] = json!(most);
43        }
44        v
45    }
46}
47
48pub fn options_from(v: &Value) -> Vec<OptionDecl> {
49    let one = |o: &Value| {
50        Some(OptionDecl {
51            name: o["name"].as_str()?.to_string(),
52            kind: o["kind"].as_str().unwrap_or("text").to_string(),
53            choices: o["choices"].as_array().map(|a| a.iter().filter_map(|c| c.as_str().map(str::to_string)).collect()).unwrap_or_default(),
54            least: o["least"].as_i64(),
55            most: o["most"].as_i64(),
56            default: o["default"].as_str().unwrap_or_default().to_string(),
57            about: o["about"].as_str().unwrap_or_default().to_string(),
58        })
59    };
60    v.as_array().map(|a| a.iter().filter_map(one).collect()).unwrap_or_default()
61}
62
63pub struct Report {
64    pub info: RuntimeInfo,
65    pub version: Option<String>,
66    /// The contract number the adapter checked the runtime against, when it states one.
67    pub abi: Option<u64>,
68    pub available: Result<(), String>,
69    pub devices: Vec<Device>,
70    pub capabilities: String,
71    pub model: Option<ModelCheck>,
72    pub pull: Option<Vec<(String, bool, String)>>,
73    /// The options it takes on `serve`, as `--<runtime>-<name> <value>`.
74    pub options: Vec<OptionDecl>,
75    pub named: Option<String>,
76}
77
78fn verdict(r: &Result<(), String>) -> Value {
79    match r {
80        Ok(()) => Value::Bool(true),
81        Err(why) => Value::String(why.clone()),
82    }
83}
84
85fn parsed(text: &str) -> Value {
86    serde_json::from_str(text).unwrap_or(Value::Null)
87}
88
89fn strings(v: &Value) -> Vec<String> {
90    v.as_array().map(|a| a.iter().filter_map(|e| e.as_str().map(str::to_string)).collect()).unwrap_or_default()
91}
92
93pub fn format_json(f: &Format) -> Value {
94    let mut v = json!({"id": f.id, "describe": f.describe});
95    match &f.rule {
96        Rule::File { extensions, magic } => {
97            v["file"] = json!({"extensions": extensions});
98            if let Some(m) = magic {
99                v["file"]["magic"] = json!(String::from_utf8_lossy(m));
100            }
101        }
102        Rule::Directory { files, extensions } => v["directory"] = json!({"files": files, "extensions": extensions}),
103    }
104    v
105}
106
107pub fn formats_from(v: &Value) -> Vec<Format> {
108    let one = |f: &Value| -> Option<Format> {
109        let rule = if let Some(file) = f.get("file") {
110            Rule::File {
111                extensions: strings(&file["extensions"]),
112                magic: file["magic"].as_str().map(|m| m.as_bytes().to_vec()),
113            }
114        } else {
115            let d = f.get("directory")?;
116            Rule::Directory { files: strings(&d["files"]), extensions: strings(&d["extensions"]) }
117        };
118        Some(Format { id: f["id"].as_str()?.to_string(), describe: f["describe"].as_str()?.to_string(), rule })
119    };
120    v.as_array().map(|a| a.iter().filter_map(one).collect()).unwrap_or_default()
121}
122
123pub fn tokenizer_json(t: &TokenizerSource) -> Value {
124    match t {
125        TokenizerSource::Library(name) => json!({"from": "library", "name": name}),
126        TokenizerSource::HuggingFace => json!({"from": "huggingface"}),
127        TokenizerSource::Engine => json!({"from": "engine"}),
128    }
129}
130
131pub fn tokenizer_from(v: &Value) -> Option<TokenizerSource> {
132    match v["from"].as_str()? {
133        "library" => Some(TokenizerSource::Library(v["name"].as_str()?.to_string())),
134        "huggingface" => Some(TokenizerSource::HuggingFace),
135        "engine" => Some(TokenizerSource::Engine),
136        _ => None,
137    }
138}
139
140pub fn device_json(d: &Device) -> Value {
141    json!({"backend": d.backend, "name": d.name, "memory": d.memory})
142}
143
144pub fn devices_from(v: &Value) -> Vec<Device> {
145    let one = |d: &Value| -> Option<Device> {
146        Some(Device {
147            backend: d["backend"].as_str()?.to_string(),
148            name: d["name"].as_str().unwrap_or_default().to_string(),
149            memory: d["memory"].as_u64().unwrap_or(0),
150        })
151    };
152    v.as_array().map(|a| a.iter().filter_map(one).collect()).unwrap_or_default()
153}
154
155impl Report {
156    pub fn to_json(&self) -> Value {
157        let mut v = json!({
158            "runtime": {"id": self.info.id, "version": self.version, "aliases": self.info.aliases},
159            "link_w": [crate::server::PROTO_VERSION],
160            "available": verdict(&self.available),
161            "formats": self.info.formats.iter().map(format_json).collect::<Vec<_>>(),
162            "tokenizer": tokenizer_json(&self.info.tokenizer),
163            "devices": self.devices.iter().map(device_json).collect::<Vec<_>>(),
164            "capabilities": parsed(&self.capabilities),
165        });
166        if let Some(abi) = self.abi {
167            v["runtime"]["abi"] = json!(abi);
168        }
169        if let Some(options) = &self.pull {
170            v["pull"] = json!({
171                "options": options.iter().map(|(name, value, about)| json!({"name": name, "value": value, "about": about})).collect::<Vec<_>>(),
172            });
173        }
174        if !self.options.is_empty() {
175            v["options"] = Value::Array(self.options.iter().map(OptionDecl::to_json).collect());
176        }
177        if let Some(by) = &self.named {
178            v["named"] = json!(by);
179        }
180        if let Some(m) = &self.model {
181            v["model"] = json!({"path": m.path.to_string_lossy(), "reads": verdict(&m.reads)});
182            for (key, said) in [("facts", &m.facts), ("sizing", &m.sizing)] {
183                match said {
184                    Some(Ok(text)) => v["model"][key] = parsed(text),
185                    Some(Err(why)) => v["model"][key] = Value::String(why.clone()),
186                    None => {}
187                }
188            }
189        }
190        v
191    }
192
193    pub fn exit_code(&self) -> i32 {
194        let reads = self.model.as_ref().is_none_or(|m| m.reads.is_ok());
195        if self.available.is_ok() && reads {
196            0
197        } else {
198            1
199        }
200    }
201
202    pub fn finish(&self) -> ! {
203        println!("{}", self.to_json());
204        std::process::exit(self.exit_code())
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    fn report() -> Report {
213        Report {
214            info: RuntimeInfo {
215                id: "llamacpp",
216                aliases: &["llama.cpp"],
217                formats: vec![Format::file("gguf", "a GGUF file", &["gguf"], Some(b"GGUF"))],
218                tokenizer: TokenizerSource::Library("superfluid_tokenizer_llamacpp".into()),
219            },
220            version: Some("llama.cpp b11284".into()),
221            abi: None,
222            available: Ok(()),
223            devices: vec![Device { backend: "Metal".into(), name: "Apple M5 Pro".into(), memory: 1 << 30 }],
224            capabilities: r#"{"descriptor_version":1}"#.into(),
225            model: None,
226            pull: None,
227            options: Vec::new(),
228            named: None,
229        }
230    }
231
232    #[test]
233    fn a_usable_runtime_reports_itself_and_exits_zero() {
234        let r = report();
235        assert_eq!(
236            r.to_json(),
237            json!({
238                "runtime": {"id": "llamacpp", "version": "llama.cpp b11284", "aliases": ["llama.cpp"]},
239                "link_w": [crate::server::PROTO_VERSION],
240                "available": true,
241                "formats": [{"id": "gguf", "describe": "a GGUF file", "file": {"extensions": ["gguf"], "magic": "GGUF"}}],
242                "tokenizer": {"from": "library", "name": "superfluid_tokenizer_llamacpp"},
243                "devices": [{"backend": "Metal", "name": "Apple M5 Pro", "memory": 1u64 << 30}],
244                "capabilities": {"descriptor_version": 1},
245            })
246        );
247        assert_eq!(r.exit_code(), 0);
248    }
249
250    #[test]
251    fn what_a_report_declares_reads_back() {
252        let dir = Format::directory("mlx", "an MLX model directory", &["config.json"], &["safetensors"]);
253        let mut r = report();
254        r.info.formats.push(dir);
255        let v = r.to_json();
256        assert_eq!(formats_from(&v["formats"]), r.info.formats);
257        assert_eq!(devices_from(&v["devices"]), r.devices);
258        assert_eq!(tokenizer_from(&v["tokenizer"]), Some(r.info.tokenizer.clone()));
259        for t in [TokenizerSource::HuggingFace, TokenizerSource::Engine] {
260            assert_eq!(tokenizer_from(&tokenizer_json(&t)), Some(t));
261        }
262        assert_eq!(formats_from(&json!([{"id": "x", "describe": "x", "archive": {}}])), vec![]);
263        assert_eq!(tokenizer_from(&json!({"from": "sentencepiece"})), None);
264    }
265
266    #[test]
267    fn the_options_a_runtime_takes_read_back() {
268        let mut r = report();
269        assert!(r.to_json().get("options").is_none());
270        r.options = vec![
271            OptionDecl { name: "n-gpu-layers".into(), kind: "int".into(), choices: vec![], least: Some(-1), most: None, default: "-1".into(), about: "layers on the device".into() },
272            OptionDecl { name: "threads".into(), kind: "int".into(), choices: vec![], least: Some(0), most: Some(512), default: "0".into(), about: "threads".into() },
273            OptionDecl {
274                name: "flash-attn".into(),
275                kind: "choice".into(),
276                choices: vec!["auto".into(), "on".into(), "off".into()],
277                least: None,
278                most: None,
279                default: "auto".into(),
280                about: String::new(),
281            },
282        ];
283        let v = r.to_json();
284        assert_eq!(v["options"][0], json!({"name": "n-gpu-layers", "kind": "int", "least": -1, "default": "-1", "about": "layers on the device"}));
285        assert_eq!(v["options"][1]["most"], 512);
286        assert_eq!(v["options"][2]["choices"], json!(["auto", "on", "off"]));
287        assert_eq!(options_from(&v["options"]), r.options);
288        assert_eq!(options_from(&json!([{"kind": "int"}])), vec![], "an option needs a name");
289    }
290
291    #[test]
292    fn what_is_not_usable_says_why_and_exits_one() {
293        let mut r = report();
294        r.version = None;
295        r.available = Err("ModuleNotFoundError: No module named 'mlx'".into());
296        assert_eq!(r.to_json()["available"], "ModuleNotFoundError: No module named 'mlx'");
297        assert!(r.to_json().get("named").is_none());
298        r.named = Some("SUPERFLUID_LLAMA_LIB".into());
299        assert_eq!(r.to_json()["named"], "SUPERFLUID_LLAMA_LIB");
300        assert_eq!(r.to_json()["runtime"]["version"], Value::Null);
301        assert_eq!(r.exit_code(), 1);
302
303        let mut r = report();
304        r.model = Some(ModelCheck { path: "/m".into(), reads: Err("/m is not a GGUF file".into()), facts: None, sizing: None });
305        assert_eq!(r.to_json()["model"], json!({"path": "/m", "reads": "/m is not a GGUF file"}));
306        assert_eq!(r.exit_code(), 1);
307
308        r.model = Some(ModelCheck {
309            path: "/m.gguf".into(),
310            reads: Ok(()),
311            facts: Some(Ok(r#"{"architecture":"qwen3"}"#.into())),
312            sizing: Some(Ok(r#"{"trained_context":40960}"#.into())),
313        });
314        assert_eq!(r.to_json()["model"]["facts"], json!({"architecture": "qwen3"}));
315        assert_eq!(r.to_json()["model"]["sizing"], json!({"trained_context": 40960}));
316        assert_eq!(r.exit_code(), 0);
317        r.model.as_mut().unwrap().sizing = Some(Err("no GPU here to size a window against".into()));
318        assert_eq!(r.to_json()["model"]["sizing"], "no GPU here to size a window against");
319        assert_eq!(r.exit_code(), 0);
320    }
321}
superfluid_worker/check.rs

superfluid_worker/
check.rs
