diff --git a/asset/config/config.toml.tmpl b/asset/config/config.toml.tmpl index f67ab90..0e84203 100644 --- a/asset/config/config.toml.tmpl +++ b/asset/config/config.toml.tmpl @@ -22,6 +22,17 @@ client_stall_timeout = "60s" # converges: systemd's TimeoutStopSec expires and escalates to SIGABRT. Cutting # a download deliberately is the better failure -- the client retries. shutdown_grace = "15s" +# Where an in-flight transfer's bytes are spooled so a second request for the +# same blob can be served from it while it runs, rather than waiting out the +# whole transfer with nothing on the wire. This is what makes single-flight's +# promise real. It holds no state: losing it loses nothing, because a blob is +# not recorded until it is durable in the bucket. PrivateTmp=true gives the unit +# its own /tmp, so this path is not shared with anything else on the host. +spool_dir = "/tmp/rustingface-spool" +# How long a request that joined an in-flight fetch waits for that fetch's spool +# to appear before giving up and returning 503. Not a wait for the transfer -- +# only for it to start. +flight_wait_timeout = "5s" [storage] endpoint = "{{S3_ENDPOINT}}" diff --git a/crates/rustingface-api/tests/sovereignty.rs b/crates/rustingface-api/tests/sovereignty.rs index 3746215..db11b12 100644 --- a/crates/rustingface-api/tests/sovereignty.rs +++ b/crates/rustingface-api/tests/sovereignty.rs @@ -253,6 +253,101 @@ async fn a_sealed_instance_refuses_loudly_for_anything_it_never_stored() { ); } +#[tokio::test] +async fn a_follower_streams_from_the_transfer_instead_of_waiting_it_out() { + // Spec §7 says concurrent callers "subscribe to the same broadcast". They + // could not: the leader's live stream cannot be replayed to someone who + // arrived mid-body, and the multipart upload is unreadable until it + // completes, so a follower simply waited -- with nothing on the wire -- + // for the whole transfer. On a multi-gigabyte blob that is minutes against + // a client that gives up in ten seconds. + // + // The spool is the missing buffer: the leader writes each chunk to local + // disk as it passes, and a follower reads that from byte zero, following + // it as it grows. + let hub = FakeHub::start().await; + let weights = publish_model(&hub).await; + // Slow enough that the leader is still going well after the follower has + // started reading; the whole point is what happens during that window. + hub.throttle(2 * 1024, Duration::from_millis(20)).await; + + let bucket = bucket_dir("follower-streams"); + let rf = instance(&bucket, Some(&hub), |_| {}); + hub.reset_counts(); + + let leader = { + let app = rf.app.clone(); + tokio::spawn(async move { + call( + &app, + "GET", + "/Qwen/Qwen3-32B/resolve/main/model.safetensors", + &[], + ) + .await + }) + }; + + // Join once the leader's fetch is genuinely under way. + for _ in 0..400 { + if hub.file_gets() >= 1 { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + + let response = rf + .app + .clone() + .oneshot( + Request::builder() + .uri("/Qwen/Qwen3-32B/resolve/main/model.safetensors") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response.status(), + StatusCode::OK, + "a follower is served, not turned away" + ); + + use futures::StreamExt; + let mut body = response.into_body().into_data_stream(); + let first = body + .next() + .await + .expect("a chunk while the leader is still transferring") + .unwrap(); + + // The assertion that matters. Chunked delivery proves nothing -- the old + // wait-it-out path also returned the file in chunks, just all of them + // after the transfer had finished. What distinguishes streaming from + // waiting is *when* the first byte lands: here, while the leader is still + // reading from upstream. + assert!( + !leader.is_finished(), + "the follower's first bytes arrived while the leader was still transferring" + ); + assert!(!first.is_empty()); + + let mut served = first.to_vec(); + while let Some(chunk) = body.next().await { + served.extend_from_slice(&chunk.unwrap()); + } + assert_eq!(served, weights, "and it receives the whole, correct file"); + + let (status, _, leader_body) = leader.await.unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(leader_body, weights); + assert_eq!( + hub.file_gets(), + 1, + "and it cost no second upstream fetch: single-flight still holds" + ); +} + #[tokio::test] async fn eight_concurrent_clients_produce_exactly_one_upstream_fetch() { let hub = FakeHub::start().await; diff --git a/crates/rustingface-core/src/fetch.rs b/crates/rustingface-core/src/fetch.rs index 556b28b..3768c37 100644 --- a/crates/rustingface-core/src/fetch.rs +++ b/crates/rustingface-core/src/fetch.rs @@ -22,6 +22,7 @@ use crate::digest::{self, Digest}; use crate::flight::Outcome; use crate::ports::{Store, UpstreamHead}; use crate::registry::{FetchTask, LeadResult, Registry, TEE_DEPTH}; +use crate::spool::Spool; /// Where a transfer's bytes are being written while the digest is unknown. enum Sink { @@ -76,6 +77,27 @@ impl Registry { let (client_tx, client_rx) = mpsc::channel::>(TEE_DEPTH); let (outcome_tx, outcome_rx) = oneshot::channel::(); + // Spool the bytes as they pass so a second request for the same blob + // can be served from this transfer rather than waiting it out. Every + // transfer is spooled, not just the large ones: the cost on a small + // file is a create and a delete, and making it conditional would mean + // followers behave one way for a config file and another for a shard, + // which is the kind of size-dependent behaviour this service has + // already been bitten by once. + let total = head.linked_size.or(head.content_length).unwrap_or(0); + let spool = match &self.config().server.spool_dir { + Some(dir) => match Spool::create(dir, total, head.clone()) { + Ok(spool) => Some(spool), + Err(err) => { + // Not fatal: without a spool, followers are told to retry. + // Losing the transfer over it would be worse. + warn!(%err, dir = %dir.display(), "could not open a spool for this transfer; followers will be told to retry"); + None + } + }, + None => None, + }; + let transfer = Transfer { store: Arc::clone(self.store()), registry_repo_type: repo_type, @@ -87,6 +109,7 @@ impl Registry { part_size: self.config().storage.multipart_part_size, stall_timeout: self.config().server.client_stall_timeout, client_range: client_range.clone(), + spool: spool.clone(), }; let registry = self.clone(); @@ -104,6 +127,7 @@ impl Registry { rx: client_rx, task: FetchTask::new(outcome_rx), range: client_range, + spool, }) } @@ -137,6 +161,9 @@ struct Transfer { stall_timeout: std::time::Duration, /// Which bytes of the file the client asked for. `None` is the whole file. client_range: Option>, + /// Where the passing bytes are also written so other requests for the same + /// blob can be served from them while this transfer runs. + spool: Option>, } impl Transfer { @@ -183,6 +210,15 @@ impl Transfer { None => buffered.push(chunk.clone()), } + // Publish only after the write lands: a follower reading ahead + // of the bytes actually on disk would get zeroes or a short + // read, which is the silent corruption this service exists to + // prevent. + if let Some(spool) = &self.spool { + spool.append(&chunk).await?; + spool.published(offset); + } + // 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. diff --git a/crates/rustingface-core/src/flight.rs b/crates/rustingface-core/src/flight.rs index 1cd8de0..3b4c273 100644 --- a/crates/rustingface-core/src/flight.rs +++ b/crates/rustingface-core/src/flight.rs @@ -19,6 +19,8 @@ use rustingface_entities::manifest::Entry; use rustingface_entities::repo::{RepoId, RepoType}; use tokio::sync::watch; +use crate::spool::Spool; + /// Identifies one uncached blob. #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct FlightKey { @@ -90,6 +92,30 @@ impl From<&Error> for Outcome { #[derive(Debug)] pub struct Flight { rx: watch::Receiver>, + spool: watch::Receiver>>, +} + +impl Flight { + /// A receiver for the flight's outcome, for a reader that needs to know + /// when the blob became durable. + pub fn outcome(&self) -> watch::Receiver> { + self.rx.clone() + } + + /// The leader's spool, once it has one. + /// + /// A follower arrives before the leader has read a byte, so this waits + /// briefly for the spool to be published rather than concluding there is + /// none. `None` means the leader is not spooling — a small buffered file, + /// or a transfer that failed before it started — and the caller falls back + /// to waiting for the outcome. + pub async fn spool(&mut self, within: std::time::Duration) -> Option> { + if let Some(spool) = self.spool.borrow_and_update().clone() { + return Some(spool); + } + let _ = tokio::time::timeout(within, self.spool.changed()).await; + self.spool.borrow().clone() + } } impl Flight { @@ -117,9 +143,10 @@ impl Flight { /// and closes the channel, so a panicking leader cannot wedge later requests /// on a key that will never complete. pub struct Lead { - map: Arc>>>, + map: Arc>, key: FlightKey, tx: watch::Sender>, + spool: watch::Sender>>, } impl Lead { @@ -128,6 +155,14 @@ impl Lead { self.map.remove(&self.key); let _ = self.tx.send(Some(outcome)); } + + /// Offer this transfer's spool to everyone waiting on the flight. + /// + /// Called as soon as the transfer starts, so a follower that arrived + /// moments earlier can begin streaming from byte zero instead of waiting. + pub fn publish_spool(&self, spool: Arc) { + let _ = self.spool.send(Some(spool)); + } } impl Drop for Lead { @@ -136,6 +171,13 @@ impl Drop for Lead { } } +/// What a flight publishes to its followers. +#[derive(Clone, Debug)] +struct Shared { + outcome: watch::Receiver>, + spool: watch::Receiver>>, +} + /// Either this caller leads the fetch, or it waits on one already running. pub enum Role { Leader(Lead), @@ -145,7 +187,7 @@ pub enum Role { /// The in-flight registry. #[derive(Debug, Default)] pub struct Flights { - map: Arc>>>, + map: Arc>, } impl Flights { @@ -162,15 +204,21 @@ impl Flights { use dashmap::mapref::entry::Entry as MapEntry; match self.map.entry(key.clone()) { MapEntry::Occupied(occupied) => Role::Follower(Flight { - rx: occupied.get().clone(), + rx: occupied.get().outcome.clone(), + spool: occupied.get().spool.clone(), }), MapEntry::Vacant(vacant) => { let (tx, rx) = watch::channel(None); - vacant.insert(rx); + let (spool_tx, spool_rx) = watch::channel(None); + vacant.insert(Shared { + outcome: rx, + spool: spool_rx, + }); Role::Leader(Lead { map: Arc::clone(&self.map), key, tx, + spool: spool_tx, }) } } diff --git a/crates/rustingface-core/src/lib.rs b/crates/rustingface-core/src/lib.rs index 340a649..92816eb 100644 --- a/crates/rustingface-core/src/lib.rs +++ b/crates/rustingface-core/src/lib.rs @@ -13,6 +13,7 @@ pub mod flight; pub mod policy; pub mod ports; pub mod registry; +pub mod spool; pub use policy::Policy; pub use registry::{Body, Registry, ResolvedFile, Source}; diff --git a/crates/rustingface-core/src/registry.rs b/crates/rustingface-core/src/registry.rs index 73c5f60..e0ae872 100644 --- a/crates/rustingface-core/src/registry.rs +++ b/crates/rustingface-core/src/registry.rs @@ -719,8 +719,56 @@ impl Registry { let key = FlightKey::new(repo_type, repo_id.clone(), &commit, path); match self.inner.flights.join(key) { - Role::Follower(flight) => { + Role::Follower(mut flight) => { debug!(%repo_id, commit, path, "joining an in-flight upstream fetch"); + + // Spec §7: followers subscribe to the same broadcast rather + // than issuing their own request. The spool is what makes that + // possible -- the leader's live stream cannot be replayed, but + // its bytes on disk can. Without one there is nothing to read + // and the only honest answer is to come back later. + let grace = self.inner.config.server.flight_wait_timeout; + if let Some(spool) = flight.spool(grace).await { + let total = spool.total(); + let asked = match &range { + // Same rule as a cold ranged lead: a range that stops + // short of the end cannot be completed before the whole + // file is verified. + Some(asked) if asked.end != total => { + return Err(Error::FetchInProgress(format!( + "{repo_id}@{commit}/{path} is being fetched; a range that stops \ + short of the end can only be served once it is stored. Retry \ + shortly." + ))); + } + Some(asked) => asked.clone(), + None => 0..total, + }; + let head = spool.head().clone(); + let len = asked.end - asked.start; + let ranged = range.clone(); + let stream = crate::spool::read_while_written( + Arc::clone(&spool), + asked, + flight.outcome(), + ); + return Ok(( + ResolvedFile { + commit, + size: total, + etag: head.etag, + linked_etag: head.linked_etag, + linked_size: head.linked_size, + body: Body::Stream { + stream, + len, + range: ranged, + }, + }, + Source::Upstream, + )); + } + let entry = flight.wait().await?; let file = self .serve_stored(&entry, &commit, range) @@ -758,7 +806,13 @@ impl Registry { rx, task, range, + spool, }) => { + // Offer the spool before handing the flight over, so a + // follower already waiting can start reading at once. + if let Some(spool) = spool { + lead.publish_spool(spool); + } // The upload task owns the flight from here: it // publishes the outcome whether or not this client // stays connected. @@ -939,6 +993,8 @@ pub(crate) enum LeadResult { rx: mpsc::Receiver>, /// The slice of the file the tee is narrowed to, if any. range: Option>, + /// Where the transfer's bytes are also landing, for followers to read. + spool: Option>, task: FetchTask, }, } diff --git a/crates/rustingface-core/src/spool.rs b/crates/rustingface-core/src/spool.rs new file mode 100644 index 0000000..53bdaf4 --- /dev/null +++ b/crates/rustingface-core/src/spool.rs @@ -0,0 +1,244 @@ +//! The spool: an in-flight transfer's bytes, readable before it finishes. +//! +//! Single-flight (spec §7) says concurrent requests for the same uncached blob +//! must produce exactly one upstream fetch, and that "subsequent callers +//! subscribe to the same broadcast". Subscribing is the part that was missing: +//! a follower cannot be handed the leader's live stream, because the bytes +//! that already passed are gone from it and the multipart upload is not +//! readable until it completes. So it waited instead — with nothing on the +//! wire — for the whole transfer, which on a multi-gigabyte blob is minutes +//! against a client that gives up in ten seconds. +//! +//! The spool is the missing buffer. The leader writes every chunk to a local +//! file as it passes and publishes how much is safe to read; a follower opens +//! that file and streams it from the start, waiting for growth when it catches +//! up. One upstream fetch, every caller served, nobody in silence. +//! +//! It holds no state. Losing it loses nothing: the blob is not recorded until +//! it is durable in the bucket, so a crash mid-transfer leaves the file +//! unreferenced and the next request refetches. It is the same role the +//! in-memory buffer already plays for small files, on disk so it can hold a +//! large one. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use tokio::sync::watch; + +use crate::ports::UpstreamHead; + +/// A transfer's bytes on local disk, growing as they arrive. +#[derive(Debug)] +pub struct Spool { + path: PathBuf, + total: u64, + head: UpstreamHead, + /// Bytes durably written to `path` and safe for a follower to read. + written: watch::Sender, +} + +impl Spool { + /// Create a spool for a transfer of `total` bytes. + pub fn create(dir: &Path, total: u64, head: UpstreamHead) -> std::io::Result> { + std::fs::create_dir_all(dir)?; + let path = dir.join(format!("spool-{}", uuid())); + // Created here so a follower can open it the moment it is published. + std::fs::File::create(&path)?; + Ok(Arc::new(Self { + path, + total, + head, + written: watch::channel(0).0, + })) + } + + pub fn path(&self) -> &Path { + &self.path + } + + pub fn total(&self) -> u64 { + self.total + } + + pub fn head(&self) -> &UpstreamHead { + &self.head + } + + /// Watch how much of the file is readable. + pub fn progress(&self) -> watch::Receiver { + self.written.subscribe() + } + + /// Append a chunk. The file is opened per call rather than held, so the + /// leader's handle cannot outlive a failed transfer and strand the file. + pub async fn append(&self, chunk: &[u8]) -> rustingface_entities::error::Result<()> { + use tokio::io::AsyncWriteExt; + let mut file = tokio::fs::OpenOptions::new() + .append(true) + .open(&self.path) + .await + .map_err(|e| spool_error(&self.path, e))?; + file.write_all(chunk) + .await + .map_err(|e| spool_error(&self.path, e))?; + Ok(()) + } + + /// Announce that `n` bytes are on disk. Call only after the write lands. + pub fn published(&self, n: u64) { + let _ = self.written.send(n); + } +} + +impl Drop for Spool { + fn drop(&mut self) { + // Every reader holds an Arc, so this runs once the last one is done. + let _ = std::fs::remove_file(&self.path); + } +} + +/// A spool failure is a storage failure: the bytes did not land where the +/// next reader expects them. +fn spool_error(path: &Path, err: std::io::Error) -> rustingface_entities::error::Error { + rustingface_entities::error::Error::Storage(format!("spool {}: {err}", path.display())) +} + +/// A unique-enough name; the directory is ours alone. +fn uuid() -> String { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::{SystemTime, UNIX_EPOCH}; + static SEQ: AtomicU64 = AtomicU64::new(0); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + format!("{nanos:x}-{:x}", SEQ.fetch_add(1, Ordering::Relaxed)) +} + +/// Stream a spool's bytes from `range.start`, following it as it grows. +/// +/// The final chunk is withheld until the flight reports the blob stored, for +/// the same reason the tee runs one chunk behind: a response that completes +/// before the manifest write would tell the client its bytes are recorded when +/// they are not. A failed flight ends the stream with that failure rather than +/// a short read. +pub fn read_while_written( + spool: Arc, + range: std::ops::Range, + outcome: watch::Receiver>, +) -> futures::stream::BoxStream<'static, rustingface_entities::error::Result> { + use futures::StreamExt; + use tokio::io::{AsyncReadExt, AsyncSeekExt}; + + const CHUNK: usize = 1024 * 1024; + + let progress = spool.progress(); + let stream = futures::stream::unfold( + (spool, range.start, range.end, None::, false), + move |(spool, mut pos, end, held, mut done)| { + let mut progress = progress.clone(); + let mut outcome = outcome.clone(); + async move { + loop { + // Everything asked for has been read; release the withheld + // tail once the blob is durable and referenced. + if pos >= end { + let tail = held?; + loop { + // Bound the borrow before awaiting: holding a + // watch guard across an await makes the stream + // non-Send. + let current = outcome.borrow_and_update().clone(); + match current { + Some(crate::flight::Outcome::Stored(_)) => break, + Some(failed) => { + let err = failed.into_result().unwrap_err(); + return Some((Err(err), (spool, pos, end, None, true))); + } + None => { + if outcome.changed().await.is_err() { + return Some(( + Err(rustingface_entities::error::Error::Storage( + "the transfer this response was reading from ended without reporting an outcome".into(), + )), + (spool, pos, end, None, true), + )); + } + } + } + } + return Some((Ok(tail), (spool, pos, end, None, true))); + } + + let available = *progress.borrow_and_update(); + if available > pos { + let want = (available.min(end) - pos).min(CHUNK as u64) as usize; + let mut buf = vec![0u8; want]; + let read = async { + let mut file = tokio::fs::File::open(spool.path()).await?; + file.seek(std::io::SeekFrom::Start(pos)).await?; + file.read_exact(&mut buf).await?; + Ok::<_, std::io::Error>(()) + } + .await; + if let Err(err) = read { + return Some(( + Err(spool_error(spool.path(), err)), + (spool, pos, end, None, true), + )); + } + pos += want as u64; + let chunk = bytes::Bytes::from(buf); + // One chunk behind, exactly like the tee. + let previous = held; + return match previous { + Some(previous) => { + Some((Ok(previous), (spool, pos, end, Some(chunk), done))) + } + None => { + // Nothing to emit yet; go round again holding + // this one. + let state = (spool, pos, end, Some(chunk), done); + Some((Ok(bytes::Bytes::new()), state)) + } + }; + } + + // Caught up with the writer. Either more is coming, or the + // transfer has ended and this is all there will be. + if done { + return None; + } + tokio::select! { + changed = progress.changed() => { + if changed.is_err() { + done = true; + } + } + changed = outcome.changed() => { + if changed.is_err() { + done = true; + } else { + let current = outcome.borrow().clone(); + if let Some(failed @ crate::flight::Outcome::Failed { .. }) = current + { + let err = failed.into_result().unwrap_err(); + return Some((Err(err), (spool, pos, end, None, true))); + } + } + } + } + } + } + }, + ); + + // The empty frames emitted while priming the holdback are noise on the + // wire; drop them rather than making the client parse zero-length chunks. + stream + .filter(|item| { + let keep = !matches!(item, Ok(bytes) if bytes.is_empty()); + async move { keep } + }) + .boxed() +} diff --git a/crates/rustingface-entities/src/config.rs b/crates/rustingface-entities/src/config.rs index 9b29fe6..78814f0 100644 --- a/crates/rustingface-entities/src/config.rs +++ b/crates/rustingface-entities/src/config.rs @@ -51,6 +51,24 @@ pub struct Server { /// detached and finished without it. #[serde(with = "humantime_serde")] pub client_stall_timeout: Duration, + /// How long a request that joined another request's in-flight fetch waits + /// for that fetch's spool to appear before giving up on it. + /// + /// A follower arrives before the leader has read a byte, so this is not a + /// wait for the *transfer* but for it to start. It must stay well under a + /// client's read timeout -- `huggingface_hub` allows 10s. + #[serde(with = "humantime_serde")] + pub flight_wait_timeout: Duration, + /// Where an in-flight transfer's bytes are spooled so that other requests + /// for the same blob can be served from it while it runs. + /// + /// This is what makes single-flight's promise real: without it a second + /// request for the same uncached blob has nothing to read and can only + /// wait for the whole transfer. `None` disables spooling, at the cost of + /// returning those requests a 503 instead. Holds no state — losing it + /// loses nothing, because a blob is not recorded until it is durable in + /// the bucket. + pub spool_dir: Option, /// How long to let in-flight requests drain on SIGTERM before stopping /// anyway. /// @@ -69,6 +87,8 @@ impl Default for Server { listen: "127.0.0.1:20482".into(), external_url: None, client_stall_timeout: Duration::from_secs(60), + flight_wait_timeout: Duration::from_secs(5), + spool_dir: Some(PathBuf::from("/tmp/rustingface-spool")), shutdown_grace: Duration::from_secs(15), } } diff --git a/doc/spec.md b/doc/spec.md index 4184b00..db7955e 100644 --- a/doc/spec.md +++ b/doc/spec.md @@ -227,6 +227,8 @@ listen = "127.0.0.1:8080" external_url = "https://models.internal.example" # used to build redirect URLs client_stall_timeout = "60s" # idle guard on the tee; see below shutdown_grace = "15s" # bounded drain on SIGTERM; see below +spool_dir = "/tmp/rustingface-spool" # in-flight bytes, readable by followers +flight_wait_timeout = "5s" # how long a follower waits for a spool [storage] endpoint = "https://s3.internal.example:9000" @@ -410,6 +412,22 @@ subsequent callers subscribe to the same broadcast rather than issuing their own request. Without this, a fleet-wide rollout of a new model produces N concurrent 40GB pulls. +Subscribing is what the spool is for. A follower cannot be handed the leader's +live stream — the bytes that already passed are gone from it, and the multipart +upload is not readable until it completes — so without one it can only wait for +the entire transfer, with nothing on the wire, which a client's read timeout +ends long before. The leader therefore writes every chunk to `server.spool_dir` +as it passes and publishes how much is safe to read; a follower opens that file +and streams it from the start, following it as it grows. One upstream fetch, +every caller served, nobody in silence. + +The spool is not state. Losing it loses nothing, because a blob is not recorded +until it is durable in the bucket, so a crash mid-transfer leaves an +unreferenced blob and the next request refetches. It is the role the in-memory +buffer already plays for small files, on disk so it can hold a large one. With +`spool_dir` unset there is nothing for a follower to read and it is answered +`503` instead. + ### Streaming tee The upstream response body is consumed once and written to two sinks