Compare commits
6 Commits
fix/bounde
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
3747af6ccc
|
|||
|
32c8a3c2b5
|
|||
|
93cc64baf3
|
|||
|
d1b1cdf3de
|
|||
|
14f8342cf8
|
|||
|
a88bbd1ee2
|
@@ -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}}"
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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::<Result<Bytes>>(TEE_DEPTH);
|
||||
let (outcome_tx, outcome_rx) = oneshot::channel::<Outcome>();
|
||||
|
||||
// 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<Range<u64>>,
|
||||
/// 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<Arc<Spool>>,
|
||||
}
|
||||
|
||||
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.
|
||||
|
||||
@@ -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<Option<Outcome>>,
|
||||
spool: watch::Receiver<Option<Arc<Spool>>>,
|
||||
}
|
||||
|
||||
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<Option<Outcome>> {
|
||||
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<Arc<Spool>> {
|
||||
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<DashMap<FlightKey, watch::Receiver<Option<Outcome>>>>,
|
||||
map: Arc<DashMap<FlightKey, Shared>>,
|
||||
key: FlightKey,
|
||||
tx: watch::Sender<Option<Outcome>>,
|
||||
spool: watch::Sender<Option<Arc<Spool>>>,
|
||||
}
|
||||
|
||||
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<Spool>) {
|
||||
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<Option<Outcome>>,
|
||||
spool: watch::Receiver<Option<Arc<Spool>>>,
|
||||
}
|
||||
|
||||
/// 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<DashMap<FlightKey, watch::Receiver<Option<Outcome>>>>,
|
||||
map: Arc<DashMap<FlightKey, Shared>>,
|
||||
}
|
||||
|
||||
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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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<Result<Bytes>>,
|
||||
/// The slice of the file the tee is narrowed to, if any.
|
||||
range: Option<Range<u64>>,
|
||||
/// Where the transfer's bytes are also landing, for followers to read.
|
||||
spool: Option<Arc<crate::spool::Spool>>,
|
||||
task: FetchTask,
|
||||
},
|
||||
}
|
||||
|
||||
244
crates/rustingface-core/src/spool.rs
Normal file
244
crates/rustingface-core/src/spool.rs
Normal file
@@ -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<u64>,
|
||||
}
|
||||
|
||||
impl Spool {
|
||||
/// Create a spool for a transfer of `total` bytes.
|
||||
pub fn create(dir: &Path, total: u64, head: UpstreamHead) -> std::io::Result<Arc<Self>> {
|
||||
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<u64> {
|
||||
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<Spool>,
|
||||
range: std::ops::Range<u64>,
|
||||
outcome: watch::Receiver<Option<crate::flight::Outcome>>,
|
||||
) -> futures::stream::BoxStream<'static, rustingface_entities::error::Result<bytes::Bytes>> {
|
||||
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::<bytes::Bytes>, 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()
|
||||
}
|
||||
@@ -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<PathBuf>,
|
||||
/// 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),
|
||||
}
|
||||
}
|
||||
|
||||
24
doc/spec.md
24
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
|
||||
@@ -443,7 +461,11 @@ file reconstruction information from a CAS rather than a URL. rustingface
|
||||
implements none of it.
|
||||
|
||||
- **Downstream:** never advertise Xet capability. Clients then take the ordinary
|
||||
resolve-and-redirect path.
|
||||
resolve-and-redirect path. This is the half that has to keep working as the
|
||||
Hub migrates, so the conformance suite runs a *Xet-capable* client — `hf_xet`
|
||||
installs as a `huggingface_hub` dependency whether or not a user asks for it
|
||||
— and asserts it still chooses HTTP. Testing it with Xet disabled client-side
|
||||
would prove nothing about the configuration people actually run.
|
||||
- **Upstream:** set `HF_HUB_DISABLE_XET=1` semantics on outbound fetches so
|
||||
rustingface receives whole-file URLs via the Git LFS bridge, which reconstructs
|
||||
the file and returns a single resource URL for legacy clients. Storing whole
|
||||
|
||||
@@ -208,9 +208,14 @@ def client_env(endpoint: str, home: Path) -> dict:
|
||||
"HF_ENDPOINT": endpoint,
|
||||
"HF_HOME": str(home),
|
||||
"HF_HUB_CACHE": str(home / "hub"),
|
||||
# rustingface implements no Xet client and never advertises the
|
||||
# capability; setting this makes the client's side explicit too.
|
||||
"HF_HUB_DISABLE_XET": "1",
|
||||
# Xet is deliberately *not* disabled here. `hf_xet` installs as a
|
||||
# huggingface_hub dependency, so a real user's client is Xet-capable
|
||||
# whether or not they asked for it, and the whole claim rustingface
|
||||
# makes is that pointing HF_ENDPOINT at it works with the client
|
||||
# people actually have. Setting HF_HUB_DISABLE_XET would test a
|
||||
# configuration nobody runs, and would hide the one failure mode
|
||||
# that matters here: a client choosing Xet against a service that
|
||||
# implements none of it (spec §8).
|
||||
"HF_HUB_DISABLE_TELEMETRY": "1",
|
||||
"HF_HUB_DISABLE_PROGRESS_BARS": "1",
|
||||
}
|
||||
@@ -270,6 +275,32 @@ meta = get_hf_file_metadata(url)
|
||||
print(json.dumps({{"commit": meta.commit_hash, "etag": meta.etag, "size": meta.size}}))
|
||||
"""
|
||||
|
||||
# The Xet negotiation, from the client's own side. `hf_xet` ships as a
|
||||
# huggingface_hub dependency, so a real user's client is Xet-capable whether or
|
||||
# not they asked for it. What keeps rustingface compatible is that the *server*
|
||||
# never advertises Xet: `parse_xet_file_data_from_response` returns None, and
|
||||
# `hf_hub_download` dispatches to plain HTTP (file_download.py, "if
|
||||
# xet_file_data is not None and is_xet_available()").
|
||||
#
|
||||
# Asserted from inside the client rather than by inspecting headers ourselves,
|
||||
# because the question is not "which headers did we send" but "what did the
|
||||
# client conclude from them".
|
||||
XET_FALLBACK = """
|
||||
import json
|
||||
from huggingface_hub import get_session, hf_hub_url
|
||||
from huggingface_hub.utils._runtime import is_xet_available
|
||||
from huggingface_hub.utils._xet import parse_xet_file_data_from_response
|
||||
|
||||
url = hf_hub_url({repo!r}, "config.json", revision={revision!r})
|
||||
response = get_session().get(url, headers={{"Accept-Encoding": "identity"}}, allow_redirects=False)
|
||||
response.raise_for_status()
|
||||
print(json.dumps({{
|
||||
"xet_capable": bool(is_xet_available()),
|
||||
"xet_offered": parse_xet_file_data_from_response(response) is not None,
|
||||
"xet_headers": sorted(h for h in response.headers if h.lower().startswith("x-xet-")),
|
||||
}}))
|
||||
"""
|
||||
|
||||
MISSING_FILE = """
|
||||
import json
|
||||
from huggingface_hub import hf_hub_download
|
||||
@@ -451,6 +482,37 @@ def main() -> int:
|
||||
|
||||
results.check("model_info answers through rustingface", model_info)
|
||||
|
||||
def xet_capable_client_falls_back_to_http() -> None:
|
||||
"""The central compatibility claim: point HF_ENDPOINT here and it works.
|
||||
|
||||
Xet is the thing most likely to break that quietly. The Hub is moving to
|
||||
it, `hf_xet` is installed alongside `huggingface_hub` whether or not the
|
||||
user asked for it, and rustingface implements none of it (spec §8) --
|
||||
deliberately, because a Xet-backed fetch would leave the bucket
|
||||
unreadable without a Xet client. That only holds while the client
|
||||
*chooses* HTTP, so assert the choice rather than assume it.
|
||||
"""
|
||||
with Instance(args.binary, workdir, storage, open_upstream) as rf:
|
||||
seen = json.loads(
|
||||
run_client(XET_FALLBACK.format(repo=REPO, revision=REVISION), rf.endpoint, hf_home)
|
||||
)
|
||||
if not seen["xet_capable"]:
|
||||
raise Failure(
|
||||
"hf_xet is not installed, so this check proves nothing: it would pass against a "
|
||||
"client that could not use Xet even if rustingface offered it. Install the pinned "
|
||||
"requirements."
|
||||
)
|
||||
if seen["xet_offered"] or seen["xet_headers"]:
|
||||
raise Failure(
|
||||
"rustingface advertised Xet to the client "
|
||||
f"(headers {seen['xet_headers']}). It implements none of it, and a client that "
|
||||
"took that path would write a bucket nobody can read back without a Xet client, "
|
||||
"which forfeits state portability (spec §8)."
|
||||
)
|
||||
print(" client is Xet-capable, rustingface offered no Xet, download stays on HTTP")
|
||||
|
||||
results.check("a Xet-capable client falls back to plain HTTP", xet_capable_client_falls_back_to_http)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Failures are loud.
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user