hf_fetch_model/
header_cache.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Opt-in on-disk cache for parsed remote tensor-file headers.
4//!
5//! `inspect --cache-headers` persists the parsed header — the same
6//! [`SafetensorsHeaderInfo`] every format (`.safetensors` / `.gguf` /
7//! `.npz` / `.pth`) normalizes into —
8//! keyed on `(repo, revision, filename, etag)`, so repeat inspection of the
9//! same remote file across invocations (the natural pattern of iterative
10//! narrowing over a handful of quant candidates) is free on the second and
11//! third call.
12//!
13//! Off by default: a plain remote `inspect` never touches local disk
14//! without this flag. Entries live under
15//! [`cache_layout::header_cache_path`](crate::cache_layout::header_cache_path)
16//! — an `hf-fm`-private sidecar directory alongside the standard `hf-hub`
17//! layout, not inside it — following the same precedent as the
18//! `.hf-fm-snapshot.json` sidecar in [`crate::cache`].
19
20use std::path::Path;
21use std::time::SystemTime;
22
23use serde::{Deserialize, Serialize};
24
25use crate::error::FetchError;
26use crate::inspect::SafetensorsHeaderInfo;
27
28/// Cache schema version embedded in every entry.
29///
30/// Bumped whenever the on-disk JSON shape changes incompatibly, including
31/// transitively via [`SafetensorsHeaderInfo`]'s own fields. A mismatch is
32/// treated as a cache miss — see [`HeaderCacheEntry::is_compatible_with`].
33const SCHEMA_VERSION: u32 = 1;
34
35/// On-disk cache entry for one parsed remote header.
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct HeaderCacheEntry {
38    /// Cache schema version. Compared against `SCHEMA_VERSION`.
39    pub schema_version: u32,
40    /// The repository identifier this entry was parsed for.
41    pub repo: String,
42    /// The resolved revision (branch, tag, or commit SHA — `"main"` when
43    /// the caller did not pin one).
44    pub revision: String,
45    /// The filename within the repo.
46    pub filename: String,
47    /// The remote etag the header was fetched against. A different etag on
48    /// a later call means the upstream file changed, invalidating this entry.
49    pub etag: String,
50    /// When this entry was written — feeds the `Source: cached header (age:
51    /// ...)` display line.
52    pub cached_at: SystemTime,
53    /// The parsed header.
54    pub info: SafetensorsHeaderInfo,
55}
56
57impl HeaderCacheEntry {
58    /// Builds a fresh entry for a header just parsed remotely.
59    #[must_use]
60    pub fn new(
61        repo: String,
62        revision: String,
63        filename: String,
64        etag: String,
65        info: SafetensorsHeaderInfo,
66    ) -> Self {
67        Self {
68            schema_version: SCHEMA_VERSION,
69            repo,
70            revision,
71            filename,
72            etag,
73            cached_at: SystemTime::now(),
74            info,
75        }
76    }
77
78    /// Returns `true` when this entry is still valid for the given
79    /// request — every field but `cached_at`/`info` must match exactly.
80    #[must_use]
81    pub fn is_compatible_with(
82        &self,
83        repo: &str,
84        revision: &str,
85        filename: &str,
86        etag: &str,
87    ) -> bool {
88        self.schema_version == SCHEMA_VERSION
89            && self.repo == repo
90            && self.revision == revision
91            && self.filename == filename
92            && self.etag == etag
93    }
94
95    /// Reads and validates the cache entry at `path` against the given
96    /// request.
97    ///
98    /// Returns `None` on a cache miss for any reason — an absent file,
99    /// unparseable JSON, or a stale/mismatched entry — never an error; the
100    /// caller falls back to a normal remote fetch either way.
101    pub async fn load(
102        path: &Path,
103        repo: &str,
104        revision: &str,
105        filename: &str,
106        etag: &str,
107    ) -> Option<Self> {
108        let text = tokio::fs::read_to_string(path).await.ok()?;
109        let entry: Self = serde_json::from_str(&text).ok()?;
110        if entry.is_compatible_with(repo, revision, filename, etag) {
111            Some(entry)
112        } else {
113            None
114        }
115    }
116
117    /// Writes this entry to `path` atomically (write-tmp + rename), via the
118    /// shared `atomic_write::write_atomic` helper — the same
119    /// durability pattern `chunked_state::ChunkedState::save_atomic`
120    /// uses. Creates the parent directory (`.hf-fm-header-cache/`) if it
121    /// does not exist yet — the first `--cache-headers` call against a repo
122    /// that was never downloaded.
123    ///
124    /// # Errors
125    ///
126    /// Returns [`FetchError::Io`] on filesystem errors during the parent
127    /// directory creation, the temp write, or the rename.
128    pub async fn save_atomic(&self, path: &Path) -> Result<(), FetchError> {
129        let json = serde_json::to_string(self).map_err(|e| {
130            FetchError::Http(format!("failed to serialize header-cache entry: {e}"))
131        })?;
132        let tmp = path.with_extension("json.tmp");
133        crate::atomic_write::write_atomic(path, &tmp, json.as_bytes()).await
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    #![allow(clippy::panic, clippy::unwrap_used, clippy::expect_used)]
140
141    use super::*;
142    use crate::inspect::TensorInfo;
143
144    fn sample_info() -> SafetensorsHeaderInfo {
145        SafetensorsHeaderInfo::new(
146            vec![TensorInfo {
147                name: "model.embed.weight".to_owned(),
148                dtype: "BF16".to_owned(),
149                shape: vec![100, 100],
150                data_offsets: (0, 20_000),
151            }],
152            None,
153            64,
154            Some(20_064),
155            None,
156        )
157    }
158
159    #[test]
160    fn is_compatible_with_matches_every_key_field() {
161        let entry = HeaderCacheEntry::new(
162            "org/model".to_owned(),
163            "main".to_owned(),
164            "model.gguf".to_owned(),
165            "etag-1".to_owned(),
166            sample_info(),
167        );
168        assert!(entry.is_compatible_with("org/model", "main", "model.gguf", "etag-1"));
169        assert!(!entry.is_compatible_with("org/other", "main", "model.gguf", "etag-1"));
170        assert!(!entry.is_compatible_with("org/model", "v2", "model.gguf", "etag-1"));
171        assert!(!entry.is_compatible_with("org/model", "main", "other.gguf", "etag-1"));
172        assert!(!entry.is_compatible_with("org/model", "main", "model.gguf", "etag-2"));
173    }
174
175    #[test]
176    fn is_compatible_with_rejects_schema_version_mismatch() {
177        let mut entry = HeaderCacheEntry::new(
178            "org/model".to_owned(),
179            "main".to_owned(),
180            "model.gguf".to_owned(),
181            "etag-1".to_owned(),
182            sample_info(),
183        );
184        entry.schema_version = SCHEMA_VERSION + 1;
185        assert!(!entry.is_compatible_with("org/model", "main", "model.gguf", "etag-1"));
186    }
187
188    #[tokio::test]
189    async fn save_then_load_round_trips() {
190        let dir = tempfile::tempdir().expect("tempdir");
191        let path = dir.path().join(".hf-fm-header-cache").join("entry.json");
192        let entry = HeaderCacheEntry::new(
193            "org/model".to_owned(),
194            "main".to_owned(),
195            "model.gguf".to_owned(),
196            "etag-1".to_owned(),
197            sample_info(),
198        );
199
200        entry.save_atomic(&path).await.expect("save");
201        let loaded = HeaderCacheEntry::load(&path, "org/model", "main", "model.gguf", "etag-1")
202            .await
203            .expect("load should hit");
204
205        assert_eq!(loaded.repo, entry.repo);
206        assert_eq!(loaded.etag, entry.etag);
207        assert_eq!(loaded.info.tensors.len(), entry.info.tensors.len());
208        assert_eq!(
209            loaded.info.tensors.first().map(|t| t.name.as_str()),
210            Some("model.embed.weight")
211        );
212    }
213
214    #[tokio::test]
215    async fn load_returns_none_for_mismatched_etag() {
216        let dir = tempfile::tempdir().expect("tempdir");
217        let path = dir.path().join(".hf-fm-header-cache").join("entry.json");
218        let entry = HeaderCacheEntry::new(
219            "org/model".to_owned(),
220            "main".to_owned(),
221            "model.gguf".to_owned(),
222            "etag-1".to_owned(),
223            sample_info(),
224        );
225        entry.save_atomic(&path).await.expect("save");
226
227        let loaded =
228            HeaderCacheEntry::load(&path, "org/model", "main", "model.gguf", "etag-2").await;
229        assert!(loaded.is_none());
230    }
231
232    #[tokio::test]
233    async fn load_returns_none_for_missing_file() {
234        let dir = tempfile::tempdir().expect("tempdir");
235        let path = dir.path().join(".hf-fm-header-cache").join("missing.json");
236        let loaded =
237            HeaderCacheEntry::load(&path, "org/model", "main", "model.gguf", "etag-1").await;
238        assert!(loaded.is_none());
239    }
240
241    #[tokio::test]
242    async fn save_atomic_does_not_leave_tmp_behind() {
243        let dir = tempfile::tempdir().expect("tempdir");
244        let path = dir.path().join(".hf-fm-header-cache").join("entry.json");
245        let entry = HeaderCacheEntry::new(
246            "org/model".to_owned(),
247            "main".to_owned(),
248            "model.gguf".to_owned(),
249            "etag-1".to_owned(),
250            sample_info(),
251        );
252        entry.save_atomic(&path).await.expect("save");
253        assert!(!path.with_extension("json.tmp").exists());
254        assert!(path.exists());
255    }
256}
hf_fetch_model/header_cache.rs

hf_fetch_model/
header_cache.rs
