qsm_core/models/
download.rs

1//! Native weight download and on-disk cache (`download` feature).
2//!
3//! Fetches a model's weight files over HTTP, verifies their SHA-256, and stores
4//! them under [`super::cache_dir`]. This is the "download on use" path for native
5//! hosts (e.g. QSMxT). It is not compiled for `wasm32-unknown-unknown`, where hosts fetch
6//! weights themselves and hand them over with [`super::provide_weights`].
7
8use std::fs;
9use std::io::{Read, Write};
10use std::path::PathBuf;
11use std::sync::atomic::{AtomicU64, Ordering};
12
13use super::{cache_dir, cache_path, resolve_local, sha256_hex, ModelSpec, WeightFile};
14
15/// Error from resolving or downloading model weights.
16#[derive(Debug)]
17pub enum DownloadError {
18    /// The model has no hosted URL yet (Pending) and no local copy was found.
19    NotHosted { model: String, file: String },
20    /// Network/transport failure.
21    Http(String),
22    /// Filesystem failure.
23    Io(std::io::Error),
24    /// Downloaded bytes did not match the expected SHA-256.
25    ChecksumMismatch { file: String, expected: String, got: String },
26}
27
28impl std::fmt::Display for DownloadError {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        match self {
31            Self::NotHosted { model, file } => write!(
32                f,
33                "weights for '{model}' are not hosted yet (file '{file}'). \
34                 Supply it locally via $QSM_MODEL_DIR."
35            ),
36            Self::Http(msg) => write!(f, "download failed: {msg}"),
37            Self::Io(e) => write!(f, "io error: {e}"),
38            Self::ChecksumMismatch { file, expected, got } => write!(
39                f,
40                "checksum mismatch for '{file}': expected {expected}, got {got}"
41            ),
42        }
43    }
44}
45
46impl std::error::Error for DownloadError {}
47
48impl From<std::io::Error> for DownloadError {
49    fn from(e: std::io::Error) -> Self {
50        Self::Io(e)
51    }
52}
53
54/// Ensure a single weight file is present locally, downloading it if needed, and
55/// return its path. See [`ensure_file_with_progress`].
56pub fn ensure_file(model_id: &str, file: &WeightFile) -> Result<PathBuf, DownloadError> {
57    ensure_file_with_progress(model_id, file, &mut |_, _| {})
58}
59
60/// Like [`ensure_file`], but reports download progress via `on_progress(downloaded,
61/// total)` (bytes). `total` is the registry-declared size ([`WeightFile::bytes`]),
62/// falling back to the response `Content-Length`. Not called on a cache hit.
63///
64/// Resolution: `$QSM_MODEL_DIR` / cache hit (checksum-validated) → return;
65/// otherwise download from [`WeightFile::url`], verify, and cache.
66pub fn ensure_file_with_progress(
67    model_id: &str,
68    file: &WeightFile,
69    on_progress: &mut dyn FnMut(u64, u64),
70) -> Result<PathBuf, DownloadError> {
71    // 1. Local override or a good cache entry.
72    if let Some(path) = resolve_local(file) {
73        if file.sha256.is_empty() || checksum_ok(&path, file.sha256)? {
74            return Ok(path);
75        }
76        // A cached file with the wrong hash is stale — re-fetch it.
77    }
78
79    // 2. Must have a URL to fetch.
80    if file.url.is_empty() {
81        return Err(DownloadError::NotHosted {
82            model: model_id.to_string(),
83            file: file.name.to_string(),
84        });
85    }
86
87    let bytes = http_get(file.url, file.bytes, on_progress)?;
88
89    // 3. Verify before trusting.
90    if !file.sha256.is_empty() {
91        let got = sha256_hex(&bytes);
92        if got != file.sha256 {
93            return Err(DownloadError::ChecksumMismatch {
94                file: file.name.to_string(),
95                expected: file.sha256.to_string(),
96                got,
97            });
98        }
99    }
100
101    // 4. Write atomically into the cache.
102    write_cached(file, &bytes)
103}
104
105/// Write `bytes` into the cache as `file`, atomically: a temp file unique to this writer,
106/// then a rename over the final name.
107///
108/// The temp name must be unique per writer, not just per file. Several threads or processes
109/// can fetch the same weights at once (a test binary runs its tests in parallel; CI jobs share
110/// a cache), and with one shared `<name>.part` a second writer's `File::create` truncated the
111/// first writer's temp file mid-write: the first then renamed a partial model into place
112/// ("model proto does not contain a graph") and the second's rename found no temp file
113/// ("No such file or directory"). With a unique temp file each rename installs a complete,
114/// verified file; the last one wins, and all of them hold the same bytes.
115fn write_cached(file: &WeightFile, bytes: &[u8]) -> Result<PathBuf, DownloadError> {
116    let path = cache_path(file);
117    write_atomically(&cache_dir(), &path, file.name, bytes)?;
118    Ok(path)
119}
120
121fn write_atomically(dir: &std::path::Path, final_path: &std::path::Path, name: &str, bytes: &[u8]) -> Result<(), DownloadError> {
122    static WRITER: AtomicU64 = AtomicU64::new(0);
123    fs::create_dir_all(dir)?;
124    let tmp = dir.join(format!(
125        "{}.{}.{}.part",
126        name,
127        std::process::id(),
128        WRITER.fetch_add(1, Ordering::Relaxed)
129    ));
130    let written = (|| {
131        let mut f = fs::File::create(&tmp)?;
132        f.write_all(bytes)?;
133        f.sync_all()?;
134        fs::rename(&tmp, final_path)
135    })();
136    if let Err(e) = written {
137        let _ = fs::remove_file(&tmp);
138        return Err(e.into());
139    }
140    Ok(())
141}
142
143
144/// Ensure every weight file for a model is present, returning their paths in
145/// [`ModelSpec::files`] order.
146pub fn ensure_model(spec: &ModelSpec) -> Result<Vec<PathBuf>, DownloadError> {
147    spec.files.iter().map(|f| ensure_file(spec.id, f)).collect()
148}
149
150fn checksum_ok(path: &PathBuf, expected: &str) -> Result<bool, DownloadError> {
151    let bytes = fs::read(path)?;
152    Ok(sha256_hex(&bytes) == expected)
153}
154
155/// Max HTTP attempts for a single weight file (1 initial + 4 retries).
156const MAX_ATTEMPTS: u32 = 5;
157
158/// Fetch a URL with exponential-backoff retries, streaming the body and reporting
159/// `on_progress(downloaded, total)` as it goes (`total` from `expected_total`, else the
160/// response `Content-Length`, else 0 = unknown). Progress resets to 0 at the start of each
161/// attempt (a retried download restarts from the beginning).
162///
163/// Weights are hosted on OSF, which throttles bursts of anonymous downloads with a
164/// **403** (not 429) — so when many jobs fetch weights at once (e.g. the CI model matrix)
165/// individual requests fail spuriously. We retry on 403/408/425/429/5xx and transport
166/// errors with backoff (0.5s, 1s, 2s, 4s) plus URL-derived jitter so concurrent fetchers
167/// of different files desynchronise. Permanent failures (404, 401, other 4xx) fail fast.
168fn http_get(
169    url: &str,
170    expected_total: u64,
171    on_progress: &mut dyn FnMut(u64, u64),
172) -> Result<Vec<u8>, DownloadError> {
173    let mut attempt = 0;
174    loop {
175        attempt += 1;
176        match http_get_once(url, expected_total, on_progress) {
177            Ok(bytes) => return Ok(bytes),
178            Err((msg, retryable)) => {
179                if !retryable || attempt >= MAX_ATTEMPTS {
180                    return Err(DownloadError::Http(format!("{msg} (after {attempt} attempt(s))")));
181                }
182                // Backoff: 0.5s · 2^(n-1), plus up to ~0.5s of URL-derived jitter.
183                let base = 500u64 << (attempt - 1);
184                let jitter = (url.bytes().map(u64::from).sum::<u64>() % 500) + u64::from(attempt) * 37;
185                std::thread::sleep(std::time::Duration::from_millis(base + jitter));
186            }
187        }
188    }
189}
190
191/// One HTTP GET, streamed in chunks. Returns `(message, retryable)` on failure.
192fn http_get_once(
193    url: &str,
194    expected_total: u64,
195    on_progress: &mut dyn FnMut(u64, u64),
196) -> Result<Vec<u8>, (String, bool)> {
197    match ureq::get(url).call() {
198        Ok(resp) => {
199            let total = if expected_total > 0 {
200                expected_total
201            } else {
202                resp.header("Content-Length").and_then(|s| s.parse().ok()).unwrap_or(0)
203            };
204            let mut reader = resp.into_reader();
205            let mut bytes: Vec<u8> = Vec::with_capacity(total as usize);
206            let mut buf = [0u8; 64 * 1024];
207            on_progress(0, total);
208            loop {
209                match reader.read(&mut buf) {
210                    Ok(0) => break,
211                    Ok(n) => {
212                        bytes.extend_from_slice(&buf[..n]);
213                        let done = bytes.len() as u64;
214                        on_progress(done, total.max(done));
215                    }
216                    // A truncated body mid-stream is transient — worth retrying.
217                    Err(e) => return Err((e.to_string(), true)),
218                }
219            }
220            Ok(bytes)
221        }
222        // OSF signals throttling with 403; treat the usual transient statuses as retryable.
223        Err(ureq::Error::Status(code, _)) => {
224            let retryable = matches!(code, 403 | 408 | 425 | 429 | 500 | 502 | 503 | 504);
225            Err((format!("status code {code}"), retryable))
226        }
227        // Connection resets / timeouts / DNS blips — retry.
228        Err(e @ ureq::Error::Transport(_)) => Err((e.to_string(), true)),
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    /// Concurrent writers of one cache entry must each succeed and leave a complete file:
237    /// with a shared `<name>.part`, one writer truncated another's temp file mid-write.
238    #[test]
239    fn concurrent_writers_leave_a_complete_file() {
240        let dir = std::env::temp_dir().join(format!("qsm-cache-race-{}", std::process::id()));
241        let _ = fs::remove_dir_all(&dir);
242        let final_path = dir.join("model.onnx");
243        let payload: Vec<u8> = (0..4_000_000u32).map(|i| (i % 251) as u8).collect();
244        std::thread::scope(|s| {
245            for _ in 0..8 {
246                s.spawn(|| {
247                    for _ in 0..5 {
248                        write_atomically(&dir, &final_path, "model.onnx", &payload).unwrap();
249                        assert_eq!(fs::read(&final_path).unwrap(), payload);
250                    }
251                });
252            }
253        });
254        let leftovers: Vec<_> = fs::read_dir(&dir).unwrap().filter_map(|e| e.ok())
255            .filter(|e| e.file_name().to_string_lossy().ends_with(".part")).collect();
256        assert!(leftovers.is_empty(), "temp files left behind: {leftovers:?}");
257        fs::remove_dir_all(&dir).unwrap();
258    }
259}
qsm_core/models/download.rs

qsm_core/models/
download.rs
