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

qsm_core/models/
download.rs
