cherenkov/model/index/
store.rs

1use super::{Artifact, ArtifactKind, Catalog, ModelEntry, ModelIndex};
2use anyhow::{Context, Result, ensure};
3use std::{
4    fs::File,
5    path::{Path, PathBuf},
6};
7
8/// Keeps owned files reachable while a reader, mmap, or GPU operation uses them.
9/// The caller must retain this handle until all derived views are finished.
10pub struct ArtifactLease {
11    /// Consumer directory, using the model root for legacy runtime layouts.
12    pub path: PathBuf,
13    _lock: Option<File>,
14}
15
16pub(super) struct Pending {
17    pub artifact: Artifact,
18    pub lease: ArtifactLease,
19}
20
21impl ModelIndex {
22    pub(super) fn artifact_root(&self, id: &str) -> PathBuf {
23        self.paths.data.join("artifacts").join(id)
24    }
25
26    pub(super) fn artifact_path(&self, artifact: &Artifact) -> PathBuf {
27        artifact
28            .external
29            .clone()
30            .unwrap_or_else(|| self.artifact_root(&artifact.id).join(&artifact.entry))
31    }
32
33    pub(super) fn lease_file(&self, id: &str) -> Result<File> {
34        let dir = self.paths.data.join("leases");
35
36        crate::storage::create_private_dir(&dir)?;
37
38        Ok(File::options()
39            .create(true)
40            .truncate(false)
41            .read(true)
42            .write(true)
43            .open(dir.join(id))?)
44    }
45
46    pub(super) fn lease(&self, artifact: &Artifact) -> Result<ArtifactLease> {
47        let lock = self.lease_file(&artifact.id)?;
48
49        lock.lock_shared()?;
50
51        let path = self.artifact_path(artifact);
52
53        ensure!(
54            path.exists(),
55            "artifact {} is missing at {}",
56            artifact.id,
57            path.display()
58        );
59
60        Ok(ArtifactLease {
61            path,
62            _lock: Some(lock),
63        })
64    }
65
66    /// Lease an existing, published prepared artifact without building or downloading.
67    /// Fail if the model is unprepared or its artifact path is missing.
68    pub fn acquire_prepared(&self, reference: &str) -> Result<ArtifactLease> {
69        let catalog = self.lock()?;
70        let model = super::lookup(&catalog.value, reference)?;
71        let id = model.prepared.as_ref().with_context(|| {
72            format!("model is not prepared; run `cherenkov prepare {reference}`")
73        })?;
74        let artifact = &catalog.value.artifacts[id];
75
76        ensure!(artifact.ready, "artifact is not ready");
77
78        self.lease(artifact)
79    }
80
81    pub(super) fn begin(&self, kind: ArtifactKind) -> Result<Pending> {
82        let mut catalog = self.lock()?;
83        let artifact = Artifact {
84            id: super::id(),
85            kind,
86            external: None,
87            entry: PathBuf::new(),
88            ready: false,
89        };
90        let root = self.artifact_root(&artifact.id);
91        let lock = self.lease_file(&artifact.id)?;
92
93        lock.lock_shared()?;
94        catalog
95            .value
96            .artifacts
97            .insert(artifact.id.clone(), artifact.clone());
98        catalog.save()?;
99        crate::storage::create_private_dir(&root)?;
100
101        Ok(Pending {
102            artifact,
103            lease: ArtifactLease {
104                path: root,
105                _lock: Some(lock),
106            },
107        })
108    }
109
110    pub(super) fn publish(
111        &self,
112        model: &ModelEntry,
113        mut pending: Artifact,
114        retained: Option<Artifact>,
115    ) -> Result<()> {
116        let mut catalog = self.lock()?;
117        let current = catalog
118            .value
119            .models
120            .get(&model.id)
121            .context("model was removed during import")?;
122
123        ensure!(
124            current.prepared == model.prepared && current.retained_source == model.retained_source,
125            "model changed during import; output was not published"
126        );
127
128        pending.ready = true;
129        let prepared_id = pending.id.clone();
130
131        catalog.value.artifacts.insert(prepared_id.clone(), pending);
132
133        let retained_id = retained.map(|mut artifact| {
134            artifact.ready = true;
135            let id = artifact.id.clone();
136
137            catalog.value.artifacts.insert(id.clone(), artifact);
138
139            id
140        });
141        let entry = catalog.value.models.get_mut(&model.id).unwrap();
142        entry.prepared = Some(prepared_id);
143
144        if let Some(id) = retained_id {
145            entry.retained_source = Some(id);
146        }
147
148        catalog.save()
149    }
150}
151
152impl ArtifactLease {
153    /// Select the model root for legacy `model/packed` stores with parent metadata.
154    /// The original artifact lock remains held; only the runner's input path changes.
155    pub(super) fn for_runtime(mut self) -> Self {
156        self.path = self.runtime_path().to_owned();
157
158        self
159    }
160
161    /// Locate legacy parent metadata without changing the artifact identity or lease.
162    pub(super) fn runtime_path(&self) -> &Path {
163        if self.path.join("config.json").is_file() && self.path.join("tokenizer.json").is_file() {
164            return &self.path;
165        }
166
167        if self.path.file_name().is_none_or(|name| name != "packed") {
168            return &self.path;
169        }
170
171        if let Some(parent) = self.path.parent()
172            && parent.join("config.json").is_file()
173            && parent.join("tokenizer.json").is_file()
174        {
175            return parent;
176        }
177
178        &self.path
179    }
180
181    /// Wrap a caller-owned path without checking it or acquiring a lock.
182    /// The caller must keep its files unchanged while readers use them.
183    pub fn external(path: PathBuf) -> Self {
184        Self { path, _lock: None }
185    }
186}
187
188pub(super) fn references(catalog: &Catalog, id: &str) -> usize {
189    catalog
190        .models
191        .values()
192        .filter(|m| m.prepared.as_deref() == Some(id) || m.retained_source.as_deref() == Some(id))
193        .count()
194}
195
196/// Do not follow links: external targets cannot become GC-owned by traversal.
197pub(super) fn bytes(path: &Path) -> Result<u64> {
198    let metadata = match std::fs::symlink_metadata(path) {
199        Ok(meta) => meta,
200        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(0),
201        Err(e) => return Err(e.into()),
202    };
203
204    if !metadata.is_dir() {
205        return Ok(metadata.len());
206    }
207
208    let mut total = 0_u64;
209
210    for entry in std::fs::read_dir(path)? {
211        total = total
212            .checked_add(bytes(&entry?.path())?)
213            .context("artifact size overflow")?;
214    }
215
216    Ok(total)
217}
218
219pub(super) fn precisions(path: &Path) -> Vec<u32> {
220    let Ok(manifest) = crate::qwen4_exp::Manifest::load(path) else {
221        return Vec::new();
222    };
223    let mut result = Vec::new();
224
225    if path.join("manifest.json").is_file() && path.join("experts.bin").is_file() {
226        result.push(4);
227    }
228
229    for bits in [3, 2] {
230        if crate::qwen4_exp::lowbit::is_usable(path, &manifest.experts, bits).unwrap_or(false) {
231            result.push(bits);
232        }
233    }
234
235    result
236}
237
238pub(super) fn rebuild_variants(path: &Path, experts: &[u32], force: bool) -> Vec<u32> {
239    let available = precisions(path);
240
241    experts
242        .iter()
243        .copied()
244        .filter(|&bits| bits != 4 && (force || !available.contains(&bits)))
245        .collect()
246}
cherenkov/model/index/store.rs

cherenkov/model/index/
store.rs
