diff --git a/server/src/models.rs b/server/src/models.rs index acd3c0d..1bf4ecc 100644 --- a/server/src/models.rs +++ b/server/src/models.rs @@ -65,6 +65,11 @@ pub struct LocalModel { #[serde(rename_all = "snake_case")] pub enum DownloadState { Running, + /// Reading the finished file back to check it against the hash + /// HuggingFace publishes. Its own state because it takes real time on + /// a multi-gigabyte file and "still working" is the honest thing to + /// show, rather than a bar sitting at 100% for half a minute. + Verifying, Finished, Failed, Cancelled, @@ -301,26 +306,49 @@ impl ModelStore { /// The download loop: resume where a partial left off, write, report. fn fetch(&self, run: &Run, target: &Path) -> Result<()> { let partial = partial_of(target); + let identity = identity_of(target); if let Some(parent) = target.parent() { private::create_dir(parent)?; } - let have = partial.metadata().map(|m| m.len()).unwrap_or(0); + + // What we have, and what it was part of. A partial with no + // recorded identity is not resumable -- it could be a fragment of + // any revision -- so it is refetched rather than guessed at. + let known = std::fs::read_to_string(&identity) + .ok() + .map(|s| s.trim().to_string()); + let have = match known { + Some(_) => partial.metadata().map(|m| m.len()).unwrap_or(0), + None => 0, + }; let url = format!( "https://huggingface.co/{}/resolve/main/{}", run.repo, run.file.replace(' ', "%20") ); - let mut request = ureq::get(&url).header("User-Agent", USER_AGENT); - if have > 0 { - request = request.header("Range", &format!("bytes={have}-")); + let (mut response, mut resumed) = request(&url, have)?; + let mut etag = etag_of(&response); + + // HuggingFace's CDN ignores `If-Range` -- probed 2026-08-28: a + // deliberately stale validator still answers 206 with the ranged + // bytes rather than 200 with the whole file. So the header cannot + // be relied on to restart us, and the check is done here instead: + // if what arrived is not the revision our partial belongs to, + // resuming would splice two files into something of exactly the + // right length and the wrong contents. Throw the partial away and + // ask again from zero. + if resumed && etag.is_some() && etag != known { + tracing::info!( + "{} changed upstream since the partial was written -- starting again", + run.key, + ); + let (fresh, fresh_resumed) = request(&url, 0)?; + response = fresh; + resumed = fresh_resumed; + etag = etag_of(&response); } - let mut response = request.call().with_context(|| format!("GET {url}"))?; - // A server that ignores Range answers 200 with the whole file, and - // appending to what we have would corrupt it -- so trust the - // status, not the request. - let resumed = response.status() == 206; let length: Option = response .headers() .get("content-length") @@ -353,7 +381,13 @@ impl ModelStore { .context("seek to resume point")?; } else { file.set_len(0) - .context("truncate a partial the server would not resume")?; + .context("truncate a partial we cannot resume onto")?; + } + // Written before the body, so an interrupted download leaves a + // partial that can still say which revision it belongs to. That is + // what makes it safe to keep one across a restart of this server. + if let Some(etag) = &etag { + std::fs::write(&identity, etag).ok(); } let mut reader = response.body_mut().as_reader(); @@ -377,14 +411,95 @@ impl ModelStore { file.flush().context("flushing the model file")?; drop(file); + // Checked before the rename, so a file that fails never gets the + // real name and `list` never offers it. With the identity check + // above this should not fire; it is here because a download of + // this size has too many ways to go subtly wrong to take on + // trust, and because a wrong model is the kind of failure that + // surfaces as bad output rather than as an error. + if let Some(expected) = published_sha256(&run.repo, &run.file) { + run.progress.lock().unwrap().state = DownloadState::Verifying; + let actual = sha256_of(&partial)?; + if actual != expected { + std::fs::remove_file(&partial).ok(); + std::fs::remove_file(&identity).ok(); + bail!( + "{} arrived corrupted -- HuggingFace publishes sha256 {expected}, what \ + arrived hashes to {actual}. It has been deleted; downloading again \ + starts clean.", + run.key, + ); + } + } + // Renamed only once complete, so a file at its real name is always // a whole model -- `list` needs no other way to tell. std::fs::rename(&partial, target) .with_context(|| format!("finish {}", target.display()))?; + std::fs::remove_file(&identity).ok(); Ok(()) } } +/// The sha256 of a file, read in chunks -- these are gigabytes, and +/// reading one into memory to hash it would be the largest allocation this +/// server ever makes. +fn sha256_of(path: &Path) -> Result { + use sha2::{Digest, Sha256}; + let mut file = + std::fs::File::open(path).with_context(|| format!("reopen {}", path.display()))?; + let mut hasher = Sha256::new(); + let mut buffer = vec![0u8; CHUNK]; + loop { + let read = file.read(&mut buffer).context("reading back to verify")?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + // Hex the same way auth.rs does, since this sha2 version's output + // type does not implement LowerHex. + Ok(hasher + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect()) +} + +/// One GET, ranged when there is something to resume onto. +fn request(url: &str, from: u64) -> Result<(ureq::http::Response, bool)> { + let mut get = ureq::get(url).header("User-Agent", USER_AGENT); + if from > 0 { + get = get.header("Range", &format!("bytes={from}-")); + } + let response = get.call().with_context(|| format!("GET {url}"))?; + // Trust the status, not the request: a server that ignores Range + // answers 200 with the whole file, and appending to that would + // corrupt it. + let resumed = response.status() == 206; + Ok((response, resumed)) +} + +fn etag_of(response: &ureq::http::Response) -> Option { + Some( + response + .headers() + .get("etag")? + .to_str() + .ok()? + .trim() + .to_string(), + ) +} + +/// `x.gguf` -> `x.gguf.part.etag`, holding which revision the partial +/// beside it is a piece of. +fn identity_of(target: &Path) -> PathBuf { + let mut name = target.as_os_str().to_os_string(); + name.push(".part.etag"); + PathBuf::from(name) +} + /// `x.gguf` -> `x.gguf.part`, the in-progress name. fn partial_of(target: &Path) -> PathBuf { let mut name = target.as_os_str().to_os_string(); @@ -477,6 +592,20 @@ pub fn search(query: &str) -> Result> { .collect()) } +/// The sha256 HuggingFace publishes for one file, if it publishes one. +/// +/// It is the LFS object id, which for these repositories is the sha256 of +/// the content -- so it is a free integrity check on a download rather +/// than something we would have to compute a second source of truth for. +fn published_sha256(repo: &str, file: &str) -> Option { + let url = format!("https://huggingface.co/api/models/{repo}/tree/main?expand=true"); + let body = get_json(&url).ok()?; + body.as_array()?.iter().find_map(|f| { + (f.get("path")?.as_str()? == file) + .then(|| f.get("lfs")?.get("oid")?.as_str().map(str::to_string))? + }) +} + /// The GGUF files in one repository, largest last, with the ones already /// downloaded marked. pub fn files(repo: &str, store: &ModelStore) -> Result> {