clawft_kernel/model_manifest/
registry.rs

1//! The local model registry: which attested models are adopted on this node,
2//! where their files lie, and whether they are usable right now.
3//!
4//! Adoption never copies. The registry stores the signed envelope and the
5//! absolute root (local only; never advertised), plus a `(size, mtime)`
6//! stamp per file recorded when its hash was last confirmed. Re-hashing is
7//! lazy: [`CheckMode::Lazy`] re-hashes a file only when its stamp changed or
8//! was never recorded, so a 45 GB shard is hashed once, not on every start.
9//! A file whose hash differs from the attestation refuses the model until an
10//! operator-run full check passes again.
11
12use std::collections::{BTreeMap, BTreeSet};
13use std::path::{Path, PathBuf};
14use std::sync::{Arc, Mutex};
15
16use serde::{Deserialize, Serialize};
17
18use super::adopt::{FileRole, ScannedModel, io_err};
19use super::resolve::{consistency, files_of_body};
20use super::check::{CheckMode, ModelCheck, ModelState, check_entry};
21use super::body::{ModelError, ModelFormat, ModelPackageBody, VerifiedModel, attest, verify_model};
22use crate::revocation::RevocationList;
23use crate::workload_pkg::{ManifestEnvelope, TrustAnchors};
24
25/// Registry file schema id.
26pub const REGISTRY_SCHEMA: &str = "weftos.model-registry.v1";
27/// Largest registry file read.
28pub const MAX_REGISTRY_BYTES: u64 = 16 * 1024 * 1024;
29
30/// A file of an adopted model and when its hash was last confirmed.
31#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
32pub struct LocalFile {
33    /// Path relative to the model root.
34    pub path: String,
35    /// Role.
36    pub role: FileRole,
37    /// Attested size.
38    pub size: u64,
39    /// Attested BLAKE3.
40    pub blake3: String,
41    /// `(size, mtime_ns)` when the hash was last confirmed on this node.
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub verified: Option<(u64, u64)>,
44}
45
46/// One adopted model.
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct ModelEntry {
49    /// The signed attestation.
50    pub envelope: ManifestEnvelope,
51    /// Canonical root directory (local only).
52    pub root: PathBuf,
53    /// Files with their verification stamps.
54    pub files: Vec<LocalFile>,
55    /// Set when a file failed its hash: the model is refused until a full
56    /// check passes.
57    #[serde(default, skip_serializing_if = "Option::is_none")]
58    pub refused: Option<String>,
59    /// Local opt-out: when true the model is not advertised in node facts
60    /// (shard hashes fingerprint which models a node holds). Not signed.
61    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
62    pub hidden: bool,
63}
64
65/// What a node trusts when it lets a model be used: pinned signers and the
66/// revocation list. Held by the registry; replaced with
67/// [`ModelRegistry::set_trust`] when keys or revocations change.
68#[derive(Clone)]
69pub struct ModelTrust {
70    /// Pinned signer keys.
71    pub anchors: TrustAnchors,
72    /// Revocation list consulted for the package, its signers and shard hashes.
73    pub revocations: Option<Arc<RevocationList>>,
74}
75
76impl std::fmt::Debug for ModelTrust {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("ModelTrust")
79            .field("signers", &self.anchors.signers.len())
80            .field("revocations", &self.revocations.is_some())
81            .finish()
82    }
83}
84
85impl ModelTrust {
86    /// Anchors with no revocation list.
87    pub fn new(anchors: TrustAnchors) -> Self {
88        Self { anchors, revocations: None }
89    }
90}
91
92impl ModelEntry {
93    /// The typed body (the envelope was verified before insertion).
94    pub fn body(&self) -> Result<ModelPackageBody, ModelError> {
95        serde_json::from_value(self.envelope.body.clone())
96            .map_err(|e| ModelError::Registry(format!("stored body: {e}")))
97    }
98
99    /// Package id of the attestation.
100    pub fn package_id(&self) -> Result<String, ModelError> {
101        Ok(self.envelope.package_id()?)
102    }
103}
104
105#[derive(Debug, Default, Serialize, Deserialize)]
106pub(super) struct RegistryFile {
107    pub(super) schema: String,
108    pub(super) models: BTreeMap<String, ModelEntry>,
109}
110
111/// A usable model, for an inference adapter to consume.
112#[derive(Debug, Clone)]
113pub struct ResolvedModel {
114    /// Package id (the manifest hash adapters put in `ModelRef.manifest`).
115    pub package_id: String,
116    /// Model name.
117    pub name: String,
118    /// Layout.
119    pub format: ModelFormat,
120    /// Directory holding the files (what `--model` takes for MLX/HF).
121    pub root: PathBuf,
122    /// Absolute path of each shard, in manifest order.
123    pub shards: Vec<PathBuf>,
124    /// Absolute tokenizer path, when the model has one.
125    pub tokenizer: Option<PathBuf>,
126    /// The attested body.
127    pub body: ModelPackageBody,
128}
129
130/// Result of [`ModelRegistry::adopt`] / [`ModelRegistry::attach`].
131#[derive(Debug, Clone)]
132pub struct AdoptedModel {
133    /// Package id.
134    pub package_id: String,
135    /// The signed envelope.
136    pub envelope: ManifestEnvelope,
137    /// Model root.
138    pub root: PathBuf,
139    /// The verified model.
140    pub verified: VerifiedModel,
141}
142
143/// Adopted models, persisted as JSON when opened on a path.
144pub struct ModelRegistry {
145    path: Option<PathBuf>,
146    pub(super) inner: Mutex<RegistryFile>,
147    pub(super) trust: Mutex<Option<ModelTrust>>,
148    /// Models whose every byte was hashed in this process. Not persisted:
149    /// the first [`ModelRegistry::resolve`] after a start does a full check.
150    pub(super) fully_verified: Mutex<BTreeSet<String>>,
151}
152
153impl std::fmt::Debug for ModelRegistry {
154    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
155        f.debug_struct("ModelRegistry").field("path", &self.path).finish()
156    }
157}
158
159impl ModelRegistry {
160    /// A registry that is not persisted.
161    pub fn in_memory() -> Self {
162        Self {
163            path: None,
164            inner: Mutex::new(RegistryFile {
165                schema: REGISTRY_SCHEMA.into(),
166                models: BTreeMap::new(),
167            }),
168            trust: Mutex::new(None),
169            fully_verified: Mutex::new(BTreeSet::new()),
170        }
171    }
172
173    /// Set what [`resolve`](Self::resolve) verifies against.
174    pub fn with_trust(self, trust: ModelTrust) -> Self {
175        *self.trust.lock().unwrap_or_else(|p| p.into_inner()) = Some(trust);
176        self
177    }
178
179    /// Replace the trust (a key was pinned or removed, a revocation list was
180    /// swapped). Takes effect on the next resolve.
181    pub fn set_trust(&self, trust: ModelTrust) {
182        *self.trust.lock().unwrap_or_else(|p| p.into_inner()) = Some(trust);
183    }
184
185    /// Open (or start) the registry file at `path`.
186    pub fn open(path: impl Into<PathBuf>) -> Result<Self, ModelError> {
187        let path = path.into();
188        let file = match std::fs::metadata(&path) {
189            Ok(m) if m.len() > MAX_REGISTRY_BYTES => {
190                return Err(ModelError::Registry("registry file too large".into()));
191            }
192            Ok(_) => {
193                let bytes = std::fs::read(&path).map_err(|e| io_err(&path, e))?;
194                let f: RegistryFile = serde_json::from_slice(&bytes)
195                    .map_err(|e| ModelError::Registry(format!("{}: {e}", path.display())))?;
196                if f.schema != REGISTRY_SCHEMA {
197                    return Err(ModelError::Registry(format!("unknown schema {:?}", f.schema)));
198                }
199                f
200            }
201            Err(e) if e.kind() == std::io::ErrorKind::NotFound => RegistryFile {
202                schema: REGISTRY_SCHEMA.into(),
203                models: BTreeMap::new(),
204            },
205            Err(e) => return Err(io_err(&path, e)),
206        };
207        for (id, e) in &file.models {
208            consistency(id, e).map_err(|why| {
209                ModelError::Registry(format!("{}: entry {id} is inconsistent: {why}", path.display()))
210            })?;
211        }
212        Ok(Self {
213            path: Some(path),
214            inner: Mutex::new(file),
215            trust: Mutex::new(None),
216            fully_verified: Mutex::new(BTreeSet::new()),
217        })
218    }
219
220    pub(super) fn lock(&self) -> std::sync::MutexGuard<'_, RegistryFile> {
221        self.inner.lock().unwrap_or_else(|p| p.into_inner())
222    }
223
224    pub(super) fn save(&self, f: &RegistryFile) -> Result<(), ModelError> {
225        let Some(path) = &self.path else { return Ok(()) };
226        if let Some(dir) = path.parent() {
227            std::fs::create_dir_all(dir).map_err(|e| io_err(dir, e))?;
228        }
229        let bytes = serde_json::to_vec_pretty(f).map_err(|e| ModelError::Registry(e.to_string()))?;
230        let tmp = path.with_extension("json.tmp");
231        std::fs::write(&tmp, &bytes).map_err(|e| io_err(&tmp, e))?;
232        #[cfg(unix)]
233        {
234            use std::os::unix::fs::PermissionsExt;
235            std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(0o600))
236                .map_err(|e| io_err(&tmp, e))?;
237        }
238        std::fs::rename(&tmp, path).map_err(|e| io_err(path, e))
239    }
240
241    fn insert(
242        &self,
243        verified: &VerifiedModel,
244        root: PathBuf,
245        files: Vec<LocalFile>,
246        replace: bool,
247    ) -> Result<AdoptedModel, ModelError> {
248        let entry = ModelEntry {
249            envelope: verified.envelope.clone(),
250            root: root.clone(),
251            files: files.clone(),
252            refused: None,
253            hidden: false,
254        };
255        consistency(&verified.package_id, &entry).map_err(ModelError::Registry)?;
256        let mut g = self.lock();
257        let name = &verified.body.name;
258        let clash: Vec<String> = g
259            .models
260            .iter()
261            .filter(|(id, e)| {
262                **id != verified.package_id
263                    && e.body().map(|b| &b.name == name).unwrap_or(false)
264            })
265            .map(|(id, _)| id.clone())
266            .collect();
267        if !clash.is_empty() && !replace {
268            return Err(ModelError::NameConflict(name.clone()));
269        }
270        for id in clash {
271            g.models.remove(&id);
272        }
273        g.models.insert(
274            verified.package_id.clone(),
275            ModelEntry {
276                envelope: verified.envelope.clone(),
277                root: root.clone(),
278                files,
279                refused: None,
280                hidden: false,
281            },
282        );
283        self.save(&g)?;
284        Ok(AdoptedModel {
285            package_id: verified.package_id.clone(),
286            envelope: verified.envelope.clone(),
287            root,
288            verified: verified.clone(),
289        })
290    }
291
292    /// Attest a scanned model with the operator key and register it. The
293    /// key must be pinned in `anchors`, otherwise the attestation does not
294    /// verify and nothing is registered. A different manifest under an
295    /// existing name is refused unless `replace` is set.
296    pub fn adopt(
297        &self,
298        scanned: ScannedModel,
299        key: &ed25519_dalek::SigningKey,
300        key_id: &str,
301        anchors: &TrustAnchors,
302        replace: bool,
303    ) -> Result<AdoptedModel, ModelError> {
304        let envelope = attest(&scanned.body, key, key_id)?;
305        let verified = verify_model(&envelope, anchors)?;
306        let files = scanned
307            .files
308            .iter()
309            .map(|f| LocalFile {
310                path: f.path.clone(),
311                role: f.role,
312                size: f.size,
313                blake3: f.blake3.clone(),
314                verified: Some((f.size, f.mtime_ns)),
315            })
316            .collect();
317        let adopted = self.insert(&verified, scanned.root, files, replace)?;
318        // The scan just hashed every byte in this process.
319        self.fully_verified
320            .lock()
321            .unwrap_or_else(|p| p.into_inner())
322            .insert(adopted.package_id.clone());
323        Ok(adopted)
324    }
325
326    /// Register an already-attested manifest (for example one fetched from
327    /// another node) whose files lie under `root`. Its signature is checked
328    /// against `anchors`; file hashes are checked lazily on first use.
329    pub fn attach(
330        &self,
331        envelope: &ManifestEnvelope,
332        root: &Path,
333        anchors: &TrustAnchors,
334        replace: bool,
335    ) -> Result<AdoptedModel, ModelError> {
336        let verified = verify_model(envelope, anchors)?;
337        let root = root.canonicalize().map_err(|e| io_err(root, e))?;
338        let files = files_of_body(&verified.body);
339        self.insert(&verified, root, files, replace)
340    }
341
342    /// Forget a model (files are untouched).
343    pub fn remove(&self, id_or_name: &str) -> Result<(), ModelError> {
344        let mut g = self.lock();
345        let id = Self::find(&g, id_or_name)?;
346        g.models.remove(&id);
347        self.save(&g)
348    }
349
350    pub(super) fn find(g: &RegistryFile, id_or_name: &str) -> Result<String, ModelError> {
351        if g.models.contains_key(id_or_name) {
352            return Ok(id_or_name.to_string());
353        }
354        g.models
355            .iter()
356            .find(|(_, e)| e.body().map(|b| b.name == id_or_name).unwrap_or(false))
357            .map(|(id, _)| id.clone())
358            .ok_or_else(|| ModelError::Unknown(id_or_name.chars().take(80).collect()))
359    }
360
361    /// Snapshot of every entry, keyed by package id.
362    pub fn entries(&self) -> Vec<(String, ModelEntry)> {
363        self.lock().models.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
364    }
365
366    /// One entry by package id or name.
367    pub fn get(&self, id_or_name: &str) -> Result<(String, ModelEntry), ModelError> {
368        let g = self.lock();
369        let id = Self::find(&g, id_or_name)?;
370        Ok((id.clone(), g.models[&id].clone()))
371    }
372
373    /// Check a model's files.
374    pub fn check(&self, id_or_name: &str, mode: CheckMode) -> Result<ModelCheck, ModelError> {
375        let (id, entry) = self.get(id_or_name)?;
376        let (check, stamps) = check_entry(&id, &entry, mode)?;
377        if mode != CheckMode::Stat {
378            let mut g = self.lock();
379            if let Some(e) = g.models.get_mut(&id) {
380                let mut changed = false;
381                for (f, s) in e.files.iter_mut().zip(stamps) {
382                    if f.verified != s {
383                        f.verified = s;
384                        changed = true;
385                    }
386                }
387                match &check.state {
388                    ModelState::Refused { reason } if e.refused.as_deref() != Some(reason) => {
389                        e.refused = Some(reason.clone());
390                        changed = true;
391                    }
392                    ModelState::Ready if mode == CheckMode::Full && e.refused.is_some() => {
393                        e.refused = None;
394                        changed = true;
395                    }
396                    _ => {}
397                }
398                if changed {
399                    self.save(&g)?;
400                }
401            }
402        }
403        Ok(check)
404    }
405}
clawft_kernel/model_manifest/registry.rs

clawft_kernel/model_manifest/
registry.rs
