cherenkov/model/index/
mod.rs

1//! Persistent model references and explicit ownership of prepared/source stores.
2mod catalog;
3mod disk;
4mod events;
5mod gc;
6mod hub;
7mod local;
8mod packing;
9mod records;
10mod reference;
11mod resolve;
12mod selector;
13pub use disk::{DiskLayout, DiskStore};
14pub use events::ModelEvent;
15pub use resolve::ResolveOptions;
16mod store;
17
18use crate::{
19    model::{Checkpoint, ModelDescription},
20    storage::Paths,
21};
22use anyhow::{Context, Result, ensure};
23pub use packing::PackOptions;
24use records::{Artifact, Catalog};
25pub use records::{
26    ArtifactDetails, ArtifactKind, GcReport, ModelDetails, ModelEntry, ModelSummary, Removal,
27    Source,
28};
29use reference::lookup;
30use std::path::Path;
31pub use store::ArtifactLease;
32
33/// Persistent model references and artifact ownership under one storage root.
34/// Operations lock the catalog as needed; this handle does not retain artifact leases.
35#[derive(Clone)]
36pub struct ModelIndex {
37    pub(super) paths: Paths,
38}
39
40impl ModelIndex {
41    /// Select the index location without opening or creating it.
42    pub fn new(paths: Paths) -> Self {
43        Self { paths }
44    }
45
46    /// Look up a source reference, alias, or model ID and return its catalog entry.
47    /// The returned metadata does not keep its artifacts alive.
48    pub fn resolve(&self, reference: &str) -> Result<ModelEntry> {
49        Ok(lookup(&self.lock()?.value, reference)?.clone())
50    }
51
52    /// Register a local checkpoint path or an `hf://owner/repo` source.
53    /// Hub registration pins a commit and reads metadata without downloading the
54    /// full weights. `revision` defaults to `main`; it and `token` apply only to HF.
55    /// Registration does not transfer file ownership. The same source reuses its ID.
56    pub fn add(
57        &self,
58        source: &str,
59        revision: Option<&str>,
60        name: Option<&str>,
61        token: Option<&str>,
62    ) -> Result<ModelDetails> {
63        self.select(
64            Path::new(source),
65            ResolveOptions {
66                name,
67                revision,
68                token,
69                ..ResolveOptions::default()
70            },
71        )
72    }
73
74    /// Inspect a local source and adopt existing prepared output without copying it.
75    fn add_local(&self, path: &Path, name: Option<&str>) -> Result<ModelDetails> {
76        let path = std::fs::canonicalize(path).context("opening local checkpoint")?;
77        let description = Checkpoint::open(&path)?.description;
78        let fingerprint = local::fingerprint(&path)?;
79        let prepared = if path.join("manifest.json").is_file() {
80            Some(path.clone())
81        } else {
82            path.join("packed/manifest.json")
83                .is_file()
84                .then(|| path.join("packed"))
85        };
86        let source = Source::Local { path, fingerprint };
87        let id = self.register(source, description, name, prepared)?;
88
89        self.show(&id)
90    }
91
92    fn register(
93        &self,
94        source: Source,
95        description: ModelDescription,
96        name: Option<&str>,
97        prepared: Option<std::path::PathBuf>,
98    ) -> Result<String> {
99        let mut catalog = self.lock()?;
100        let same = catalog
101            .value
102            .models
103            .values()
104            .find(|m| same_source(&m.source, &source))
105            .map(|m| m.id.clone());
106
107        if let Some(name) = name {
108            let owner = catalog
109                .value
110                .models
111                .values()
112                .find(|m| m.name.as_deref() == Some(name));
113
114            ensure!(
115                owner.is_none_or(|m| Some(&m.id) == same.as_ref()),
116                "model name {name:?} is already registered"
117            );
118        }
119
120        let mut model = match same {
121            Some(id) => catalog.value.models[&id].clone(),
122            None => ModelEntry {
123                id: id(),
124                name: None,
125                source,
126                description,
127                prepared: None,
128                retained_source: None,
129            },
130        };
131
132        if let Some(name) = name {
133            ensure!(
134                model.name.as_deref().is_none_or(|current| current == name),
135                "source already has name {:?}",
136                model.name
137            );
138
139            model.name = Some(name.to_owned());
140        }
141
142        if let Source::Local { path, .. } = &model.source {
143            self.require_registered_location(&catalog.value, path)?;
144
145            model.retained_source = model
146                .retained_source
147                .or_else(|| self.artifact_at(&catalog.value, ArtifactKind::Source, path));
148        }
149
150        if model.prepared.is_none()
151            && let Some(path) = prepared
152        {
153            model.prepared = Some(self.adopt_prepared(&mut catalog.value, &path)?);
154        }
155
156        let id = model.id.clone();
157
158        catalog.value.models.insert(id.clone(), model);
159        catalog.save()?;
160
161        Ok(id)
162    }
163
164    fn adopt_prepared(&self, catalog: &mut Catalog, path: &Path) -> Result<String> {
165        let path = path.canonicalize()?;
166
167        if let Some(id) = self.artifact_at(catalog, ArtifactKind::Prepared, &path) {
168            return Ok(id);
169        }
170
171        self.require_external_path(&path)?;
172
173        let artifact = Artifact {
174            id: id(),
175            kind: ArtifactKind::Prepared,
176            external: Some(path),
177            entry: Default::default(),
178            ready: true,
179        };
180        let id = artifact.id.clone();
181
182        catalog.artifacts.insert(id.clone(), artifact);
183
184        Ok(id)
185    }
186
187    /// An exact managed entry reuses its ownership record. Other paths inside
188    /// the artifact namespace cannot safely be recorded as external locations.
189    fn require_registered_location(&self, catalog: &Catalog, path: &Path) -> Result<()> {
190        if self
191            .artifact_at(catalog, ArtifactKind::Source, path)
192            .is_some()
193            || self
194                .artifact_at(catalog, ArtifactKind::Prepared, path)
195                .is_some()
196        {
197            return Ok(());
198        }
199
200        self.require_external_path(path)
201    }
202
203    fn require_external_path(&self, canonical_path: &Path) -> Result<()> {
204        let artifacts = self.paths.data.join("artifacts");
205        let root = match artifacts.canonicalize() {
206            Ok(root) => root,
207            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
208            Err(e) => return Err(e.into()),
209        };
210
211        ensure!(
212            !canonical_path.starts_with(root),
213            "path is inside managed artifacts; use an indexed model reference or a location outside the artifacts directory"
214        );
215
216        Ok(())
217    }
218
219    /// Registering a managed path adds a reference to its existing ownership record.
220    fn artifact_at(&self, catalog: &Catalog, kind: ArtifactKind, path: &Path) -> Option<String> {
221        catalog
222            .artifacts
223            .values()
224            .find(|artifact| {
225                artifact.ready
226                    && artifact.kind == kind
227                    && self.artifact_path(artifact).canonicalize().ok().as_deref() == Some(path)
228            })
229            .map(|artifact| artifact.id.clone())
230    }
231
232    /// Summarize registered models and inspect their current artifact availability.
233    /// An absent index returns an empty list without creating it.
234    pub fn list(&self) -> Result<Vec<ModelSummary>> {
235        // An empty index can be listed on a read-only root without creating it.
236        if !self.paths.data.join("index.json").exists() {
237            return Ok(Vec::new());
238        }
239
240        let catalog = self.lock()?;
241
242        catalog
243            .value
244            .models
245            .values()
246            .map(|m| self.details(&catalog.value, m).map(|d| d.summary))
247            .collect()
248    }
249
250    /// Report a model's source, preparation requirements, and referenced artifacts.
251    pub fn show(&self, reference: &str) -> Result<ModelDetails> {
252        let catalog = self.lock()?;
253
254        self.details(&catalog.value, lookup(&catalog.value, reference)?)
255    }
256
257    /// Inspect the referenced prepared artifact, or return registered source metadata.
258    /// A recorded prepared artifact that is missing or unreadable returns an error.
259    pub fn description(&self, reference: &str) -> Result<ModelDescription> {
260        let entry = self.resolve(reference)?;
261
262        if entry.prepared.is_some() {
263            let lease = self.acquire_prepared(reference)?;
264
265            return Ok(Checkpoint::open(&lease.path)?.description);
266        }
267
268        Ok(entry.description)
269    }
270
271    fn details(&self, catalog: &Catalog, model: &ModelEntry) -> Result<ModelDetails> {
272        let mut artifacts = Vec::new();
273        let mut owned_bytes = 0;
274        let mut external_bytes = 0;
275        let mut precisions = Vec::new();
276
277        for id in model.prepared.iter().chain(&model.retained_source) {
278            let artifact = &catalog.artifacts[id];
279            let path = self.artifact_path(artifact);
280            let owned = artifact.external.is_none();
281            let size_path = if owned {
282                self.artifact_root(id)
283            } else {
284                path.clone()
285            };
286            let bytes = store::bytes(&size_path)?;
287
288            if owned {
289                owned_bytes += bytes;
290            } else {
291                external_bytes += bytes;
292            }
293
294            if artifact.kind == ArtifactKind::Prepared {
295                precisions = store::precisions(&path);
296            }
297
298            artifacts.push(ArtifactDetails {
299                id: id.clone(),
300                kind: artifact.kind,
301                available: path.exists(),
302                path,
303                owned,
304                bytes,
305                references: store::references(catalog, id),
306            });
307        }
308
309        let source_local = match &model.source {
310            Source::Local { path, .. } => path.exists(),
311            Source::HuggingFace { .. } => artifacts
312                .iter()
313                .any(|a| a.kind == ArtifactKind::Source && a.available),
314        };
315        let prepared = artifacts
316            .iter()
317            .any(|a| a.kind == ArtifactKind::Prepared && a.available);
318
319        Ok(ModelDetails {
320            summary: ModelSummary {
321                id: model.id.clone(),
322                reference: reference::preferred(catalog, model),
323                name: model.name.clone(),
324                architecture: model.description.architecture.clone(),
325                source: model.source.clone(),
326                source_local,
327                prepared,
328                precisions,
329                owned_bytes,
330                external_bytes,
331            },
332            preparation: if prepared {
333                crate::model::Preparation::Direct
334            } else {
335                model.description.preparation()
336            },
337            artifacts,
338        })
339    }
340}
341
342/// Parse an alias, ID, or source URI without resolving it. Explicit paths return false.
343pub fn is_reference(value: &Path) -> bool {
344    selector::parse(value).is_ok_and(|selector| !matches!(selector, selector::Selector::Path(_)))
345}
346
347/// Anchor explicit relative paths while preserving parsed aliases and source URIs.
348/// Invalid locators return an error instead of becoming filesystem paths.
349pub fn anchor_selector(value: &Path, base: &Path) -> Result<std::path::PathBuf> {
350    Ok(match selector::parse(value)? {
351        selector::Selector::Path(path) => base.join(path),
352        _ => value.to_owned(),
353    })
354}
355
356/// Validate selector syntax without resolving an entry, reading files, or registering.
357pub fn validate_selector(value: &str) -> Result<()> {
358    selector::parse(Path::new(value)).map(|_| ())
359}
360
361/// List available expert precisions in descending order after loading the manifest.
362/// Q4 requires an existing base file. Q2/Q3 also check layout, size, and sample
363/// records. A manifest load failure returns an empty list.
364pub fn prepared_precisions(path: &Path) -> Vec<u32> {
365    store::precisions(path)
366}
367
368/// Resolve through the index and lease the prepared artifact.
369/// Explicit source paths are registered if needed; files remain externally owned.
370pub fn resolve_path(paths: Paths, path: &Path) -> Result<ArtifactLease> {
371    let index = ModelIndex::new(paths);
372    let selected = index.select(path, ResolveOptions::default())?;
373
374    index.acquire_prepared(&selected.summary.id)
375}
376
377/// Prepare missing variants through publication, never by mutating a leased view.
378/// Indexed models must already have a prepared base. If the options permit it,
379/// missing variants are built before returning a lease to the published store.
380/// On success, indexed loads disable later repacking and store creation in `options`.
381/// Legacy `model/packed` stores use their model root to find config and tokenizer.
382pub fn resolve_runtime(
383    paths: Paths,
384    path: &Path,
385    options: &mut crate::options::Options,
386) -> Result<ArtifactLease> {
387    let index = ModelIndex::new(paths);
388    let selected = index.select(path, ResolveOptions::default())?;
389    let reference = selected.summary.id.as_str();
390    let lease = index.acquire_prepared(reference)?;
391    let experts = [
392        options.experts,
393        options.miss_experts.unwrap_or(options.experts),
394    ];
395    let available = prepared_precisions(&lease.path);
396    let missing = experts.iter().any(|bits| !available.contains(bits));
397
398    if missing || options.repack {
399        ensure!(
400            options.build_missing_store,
401            "selected expert stores are missing; run `cherenkov prepare {reference} --experts {},{}`",
402            experts[0],
403            experts[1]
404        );
405        index.pack(
406            reference,
407            PackOptions {
408                output: None,
409                experts: &experts,
410                keep_source: false,
411                token: None,
412                repack: options.repack,
413            },
414        )?;
415    }
416
417    options.repack = false;
418    options.build_missing_store = false;
419
420    Ok(index.acquire_prepared(reference)?.for_runtime())
421}
422
423fn id() -> String {
424    format!("{:032x}", rand::random::<u128>())
425}
426
427fn validate_name(name: Option<&str>) -> Result<()> {
428    if let Some(name) = name {
429        ensure!(
430            crate::storage::safe_component(name) && name.len() <= 64,
431            "model names must contain 1-64 letters, digits, '.', '_' or '-'"
432        );
433        ensure!(
434            !(name.len() == 32 && name.bytes().all(|b| b.is_ascii_hexdigit())),
435            "model names cannot look like an ID"
436        );
437    }
438
439    Ok(())
440}
441
442fn same_source(a: &Source, b: &Source) -> bool {
443    match (a, b) {
444        (
445            Source::Local {
446                path: a,
447                fingerprint: x,
448            },
449            Source::Local {
450                path: b,
451                fingerprint: y,
452            },
453        ) => a == b && x == y,
454        (
455            Source::HuggingFace {
456                repo: a,
457                revision: x,
458                endpoint: p,
459            },
460            Source::HuggingFace {
461                repo: b,
462                revision: y,
463                endpoint: q,
464            },
465        ) => a == b && x == y && p == q,
466        _ => false,
467    }
468}
469
470#[cfg(test)]
471#[path = "../../../tests/unit/model/index/mod.rs"]
472mod tests;
cherenkov/model/index/mod.rs

cherenkov/model/index/
mod.rs
