diff --git a/crates/rustingface-api/tests/sovereignty.rs b/crates/rustingface-api/tests/sovereignty.rs index cf355cd..2c0752f 100644 --- a/crates/rustingface-api/tests/sovereignty.rs +++ b/crates/rustingface-api/tests/sovereignty.rs @@ -364,6 +364,75 @@ async fn a_client_that_disconnects_mid_transfer_still_leaves_the_file_stored() { assert_eq!(served, weights); } +#[tokio::test] +async fn a_cold_ranged_request_streams_while_the_file_is_still_being_fetched() { + // The failure this pins down: a ranged GET on a file the bucket does not + // hold used to fetch and store the whole thing before sending a single + // byte. A client asking for a 3GB shard therefore saw an open connection + // and total silence for the length of the transfer, which its read timeout + // ends long before -- and because `hf download` resumes with a Range + // header, *every retry after the first* took that path. The retry loop + // could never converge, and the server logged nothing but success. + // + // The whole file must still be stored (never a fragment); what must not + // happen is the client waiting on that in silence. + let hub = FakeHub::start().await; + let weights = publish_model(&hub).await; + hub.throttle(8 * 1024, Duration::from_millis(5)).await; + + let bucket = bucket_dir("cold-range-streams"); + let rf = instance(&bucket, Some(&hub), |_| {}); + + // Resume from an offset, open-ended, which is the shape huggingface_hub + // actually sends. + let response = rf + .app + .clone() + .oneshot( + Request::builder() + .uri("/Qwen/Qwen3-32B/resolve/main/model.safetensors") + .header("range", "bytes=1000-") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT); + assert_eq!( + response.headers()["content-range"], + format!("bytes 1000-{}/{}", weights.len() - 1, weights.len()) + ); + + use futures::StreamExt; + let mut body = response.into_body().into_data_stream(); + let first = body + .next() + .await + .expect("a chunk before the fetch ends") + .unwrap(); + assert!( + first.len() < weights.len() - 1000, + "the first chunk arrived while the transfer was still running, not after it" + ); + + let mut served = first.to_vec(); + while let Some(chunk) = body.next().await { + served.extend_from_slice(&chunk.unwrap()); + } + assert_eq!(served, weights[1000..], "and the range itself is intact"); + + // The bucket still holds the whole file, not the slice that was served. + let oid = { + use sha2::{Digest, Sha256}; + hex::encode(Sha256::digest(&weights)) + }; + assert_eq!( + std::fs::read(bucket.join(key::blob(&oid))).unwrap(), + weights, + "a ranged cold request stores the whole file, never a fragment" + ); +} + #[tokio::test] async fn a_ranged_request_resumes_from_the_offset_it_asks_for() { let hub = FakeHub::start().await; diff --git a/crates/rustingface-core/src/fetch.rs b/crates/rustingface-core/src/fetch.rs index b3a7086..556b28b 100644 --- a/crates/rustingface-core/src/fetch.rs +++ b/crates/rustingface-core/src/fetch.rs @@ -6,6 +6,7 @@ //! Ctrl-C at 30GB does not discard 30GB: the retention guarantee is better //! served by finishing (spec §7). +use std::ops::Range; use std::sync::Arc; use bytes::Bytes; @@ -36,16 +37,19 @@ enum Sink { impl Registry { /// Lead an upstream fetch for one file. /// - /// `blocking` forces the transfer to run to completion before returning, - /// which is what a ranged request on a cold file needs: a partial blob - /// must never reach the bucket. + /// `client_range` narrows what the *client* is sent; the whole file is + /// fetched and stored either way, because a partial blob must never reach + /// the bucket. Narrowing the tee rather than waiting for the store is what + /// keeps a ranged request on a cold file from sitting in silence for the + /// length of a multi-gigabyte transfer — which is what a resumed + /// `hf download` asks for on every retry after its first failure. pub(crate) async fn lead_fetch( &self, repo_type: RepoType, repo_id: &RepoId, commit: &str, path: &str, - blocking: bool, + client_range: Option>, ) -> Result { let body = self .upstream() @@ -82,16 +86,9 @@ impl Registry { sink, part_size: self.config().storage.multipart_part_size, stall_timeout: self.config().server.client_stall_timeout, + client_range: client_range.clone(), }; - if blocking { - // Nobody is reading the tee, so drop the client half up front and - // let the transfer run at full speed. - drop(client_rx); - let entry = transfer.run(self, body.stream, client_tx).await?; - return Ok(LeadResult::Completed(entry)); - } - let registry = self.clone(); tokio::spawn(async move { let result = transfer.run(®istry, body.stream, client_tx).await; @@ -106,6 +103,7 @@ impl Registry { head, rx: client_rx, task: FetchTask::new(outcome_rx), + range: client_range, }) } @@ -137,6 +135,8 @@ struct Transfer { sink: Sink, part_size: u64, stall_timeout: std::time::Duration, + /// Which bytes of the file the client asked for. `None` is the whole file. + client_range: Option>, } impl Transfer { @@ -166,19 +166,33 @@ impl Transfer { } }; + // Where the next chunk starts in the file, so a narrowed tee can work + // out which slice of it the client asked for. + let mut offset: u64 = 0; + let transfer = async { while let Some(chunk) = stream.next().await { let chunk = chunk?; hasher.update(&chunk); + let start = offset; + offset += chunk.len() as u64; + match &mut writer { Some(writer) => writer.write(chunk.clone()).await?, None => buffered.push(chunk.clone()), } + // The client's share of this chunk. The whole chunk when it + // asked for the whole file; the overlap with its range + // otherwise; nothing at all for a chunk outside it. + let Some(share) = self.client_share(&chunk, start, offset) else { + continue; + }; + // Tee to the client one chunk behind, dropping the client -- // not the transfer -- if it has gone away or stopped reading. - if let Some(previous) = holdback.replace(chunk) + if let Some(previous) = holdback.replace(share) && let Some(tx) = &client { match tokio::time::timeout(self.stall_timeout, tx.send(Ok(previous))).await { @@ -316,6 +330,24 @@ impl Transfer { } } +impl Transfer { + /// The part of `chunk` — which spans `[start, end)` of the file — that the + /// client asked for, or `None` when none of it is. + /// + /// Slicing is cheap: `Bytes` shares the buffer rather than copying it. + fn client_share(&self, chunk: &Bytes, start: u64, end: u64) -> Option { + let Some(range) = &self.client_range else { + return Some(chunk.clone()); + }; + let from = start.max(range.start); + let to = end.min(range.end); + if from >= to { + return None; + } + Some(chunk.slice((from - start) as usize..(to - start) as usize)) + } +} + /// Build the manifest entry for a completed transfer, carrying upstream's /// values through unchanged. fn entry_from(path: &str, size: u64, hex: &str, head: &UpstreamHead) -> Entry { diff --git a/crates/rustingface-core/src/registry.rs b/crates/rustingface-core/src/registry.rs index fa3f1c6..2b38d7d 100644 --- a/crates/rustingface-core/src/registry.rs +++ b/crates/rustingface-core/src/registry.rs @@ -734,7 +734,7 @@ impl Registry { } Role::Leader(lead) => { let outcome = self - .lead_fetch(repo_type, repo_id, &commit, path, range.is_some()) + .lead_fetch(repo_type, repo_id, &commit, path, range.clone()) .await; match outcome { Err(err) => { @@ -753,12 +753,25 @@ impl Registry { })?; Ok((file, Source::Bucket)) } - Ok(LeadResult::Streaming { head, rx, task }) => { + Ok(LeadResult::Streaming { + head, + rx, + task, + range, + }) => { // The upload task owns the flight from here: it // publishes the outcome whether or not this client // stays connected. task.attach(lead); let size = head.linked_size.or(head.content_length).unwrap_or(0); + // A ranged request is answered from the same tee, + // narrowed. The whole file is still fetched and stored + // -- a fragment must never reach the bucket -- but the + // client is not made to wait for that in silence. + let len = match &range { + Some(range) => range.end - range.start, + None => size, + }; Ok(( ResolvedFile { commit, @@ -768,25 +781,13 @@ impl Registry { linked_size: head.linked_size, body: Body::Stream { stream: Box::pin(tokio_stream_of(rx)), - len: size, - range: None, + len, + range, }, }, Source::Upstream, )) } - Ok(LeadResult::Completed(entry)) => { - lead.finish(Outcome::Stored(entry.clone())); - let file = self - .serve_stored(&entry, &commit, range) - .await? - .ok_or_else(|| { - Error::Storage(format!( - "the blob for {repo_id}@{commit}/{path} vanished after being stored" - )) - })?; - Ok((file, Source::Upstream)) - } } } } @@ -898,11 +899,10 @@ pub(crate) enum LeadResult { Streaming { head: UpstreamHead, rx: mpsc::Receiver>, + /// The slice of the file the tee is narrowed to, if any. + range: Option>, task: FetchTask, }, - /// The transfer ran to completion before the response was built, because - /// the client asked for a range and a partial blob may not be stored. - Completed(Entry), } /// Handle to a detached transfer, which owns the flight once attached. diff --git a/doc/spec.md b/doc/spec.md index 04aa970..dd6b77f 100644 --- a/doc/spec.md +++ b/doc/spec.md @@ -291,6 +291,13 @@ already on the wire and nothing logged. The adapter must disable that cap explicitly, or every proxied blob over `180s x bandwidth` fails deterministically and forever -- at 8MB/s, anything over about 1.4GB. +A ranged request on a file the bucket does not hold is subject to the same +principle from the other side. The whole file is still fetched and stored, since +a fragment must never be stored, but the client is sent its slice as those bytes +pass rather than being made to wait in silence for the transfer to end. A +resumed `hf download` asks for exactly that shape on every retry, and silence +there is indistinguishable to it from a dead server. + ### `policy.ref_resolution` - `freeze` (default) — the first resolution of a mutable ref is recorded and