diff --git a/crates/tracel-artifact/src/transfer/mod.rs b/crates/tracel-artifact/src/transfer/mod.rs index e40ee744..5a45dbd0 100644 --- a/crates/tracel-artifact/src/transfer/mod.rs +++ b/crates/tracel-artifact/src/transfer/mod.rs @@ -38,8 +38,9 @@ mod ranged; /// Watches a transfer as it runs, and can stop it. /// /// Every method defaults to doing nothing, so an implementation only has to define the events it -/// cares about. Callbacks run on the transferring thread and block it, so an implementation that -/// does real work should hand it off. +/// cares about. Callbacks run on the thread that started the transfer and block it, so an +/// implementation that does real work should hand it off. An upload sends several parts at once, +/// so it can report on more than one file at a time; each call names its file. pub trait TransferObserver { /// Returns whether the active transfer should stop. /// @@ -55,8 +56,8 @@ pub trait TransferObserver { let _ = (rel_path, total_bytes); } - /// Bytes have moved for the file being transferred. `transferred_bytes` is the running total - /// for that file, not an increment. + /// Bytes have moved for the file `rel_path`. `transferred_bytes` is the running total for that + /// file, not an increment. fn file_progress(&mut self, rel_path: &str, transferred_bytes: u64) { let _ = (rel_path, transferred_bytes); } diff --git a/crates/tracel-artifact/src/upload.rs b/crates/tracel-artifact/src/upload.rs index cd506d31..1a1052e4 100644 --- a/crates/tracel-artifact/src/upload.rs +++ b/crates/tracel-artifact/src/upload.rs @@ -2,10 +2,16 @@ //! //! The upload process can be customized with any implementation of the FileTransferClient trait (e.g. for custom HTTP clients, authentication, retries, etc), and multipart file sources can be abstracted behind the MultipartUploadSource trait for maximum flexibility (e.g. to support streaming from large files without loading them fully into memory). +use std::collections::HashSet; +use std::io::{self, Read}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::thread; + +use crossbeam::channel::{self, Receiver, Select, Sender}; + use crate::transfer::{TransferError, TransferObserver}; use crate::{FileTransferClient, ReqwestTransferClient}; -use std::collections::HashSet; -use std::io::Read; /// Errors that can occur during artifact file uploads. #[derive(Debug, thiserror::Error)] @@ -98,6 +104,12 @@ pub fn upload_bundle_multipart_with_client Result<(), UploadError> { + let plan = plan_uploads(source, files)?; + let workers = UPLOAD_WORKERS.min(plan.iter().map(|file| file.parts.len()).sum::()); + let stop = Arc::new(AtomicBool::new(false)); + let mut run = UploadRun::new(&plan, observer, Arc::clone(&stop)); + + thread::scope(|scope| { + let (jobs, queue) = channel::bounded::>(0); + let (events, reports) = channel::unbounded(); + for _ in 0..workers { + let queue = queue.clone(); + let events = events.clone(); + let stop = Arc::clone(&stop); + scope.spawn(move || { + for job in queue { + let reader = CountingReader { + inner: job.reader, + file: job.file, + events: events.clone(), + stop: Arc::clone(&stop), + }; + let result = client.put_reader(job.url, reader, job.size); + let _ = events.send(PartEvent::Finished { + file: job.file, + part: job.part, + result, + }); + } + }); + } + drop(queue); + drop(events); + + 'files: for (file, planned) in plan.iter().enumerate() { + if !run.start(file) { + break; + } + for part in &planned.parts { + run.check_cancelled(file); + if run.stopped() { + break 'files; + } + let reader = match source.open_part(planned.rel_path, part.offset, part.size) { + Ok(reader) => reader, + Err(error) => { + run.fail(error); + break 'files; + } + }; + let job = PartJob { + file, + part: part.number, + url: part.url, + size: part.size, + reader, + }; + if !run.hand_over(job, &jobs, &reports) { + break 'files; + } + } + } + drop(jobs); + while run.in_flight() { + match reports.recv() { + Ok(event) => run.handle(event), + Err(_) => break, + } + } + }); + + run.finish() +} + +/// The most parts uploaded at once. +const UPLOAD_WORKERS: usize = 8; + +/// A file of the plan, checked against its source. +struct PlannedFile<'a> { + rel_path: &'a str, + len: u64, + parts: Vec>, +} + +/// A part of a planned file: its 1-based number, URL, and the bytes of the file it covers. +struct PlannedPart<'a> { + number: usize, + url: &'a str, + offset: u64, + size: u64, +} + +/// Checks that no file is described twice and that each file's parts are numbered from 1 and +/// cover it exactly, in part order. +fn plan_uploads<'a, S: MultipartUploadSource>( + source: &S, + files: &'a [MultipartUploadFile], +) -> Result>, UploadError> { let mut seen = HashSet::new(); + files + .iter() + .map(|file| { + if !seen.insert(file.rel_path.as_str()) { + return Err(UploadError::InvalidMultipart(format!( + "Duplicate multipart upload descriptor for {}", + file.rel_path + ))); + } + let len = source.file_len(&file.rel_path)?; + let mut parts: Vec<&MultipartUploadPart> = file.parts.iter().collect(); + parts.sort_by_key(|part| part.part); + + let mut offset = 0u64; + let mut planned = Vec::with_capacity(parts.len()); + for (index, part) in parts.into_iter().enumerate() { + let number = index + 1; + if part.part != number as u32 { + return Err(UploadError::InvalidMultipart(format!( + "Invalid part numbering for {}: expected {}, got {}", + file.rel_path, number, part.part + ))); + } + if offset + part.size_bytes > len { + return Err(UploadError::InvalidMultipart(format!( + "Part {} exceeds file length for {}", + number, file.rel_path + ))); + } + planned.push(PlannedPart { + number, + url: &part.url, + offset, + size: part.size_bytes, + }); + offset += part.size_bytes; + } + if offset != len { + return Err(UploadError::InvalidMultipart(format!( + "Multipart size mismatch for {} (parts cover {}, expected {})", + file.rel_path, offset, len + ))); + } + + Ok(PlannedFile { + rel_path: &file.rel_path, + len, + parts: planned, + }) + }) + .collect() +} - for file in files { - if !seen.insert(file.rel_path.clone()) { - return Err(UploadError::InvalidMultipart(format!( - "Duplicate multipart upload descriptor for {}", - file.rel_path - ))); - } +/// A part handed to a worker, with its open reader. +struct PartJob<'a> { + file: usize, + part: usize, + url: &'a str, + size: u64, + reader: Box, +} + +/// What a worker reports to the calling thread. +enum PartEvent { + /// Bytes of a part of `file` were read to be sent. + Sent { file: usize, bytes: u64 }, + /// A part of `file` was uploaded, or failed. + Finished { + file: usize, + part: usize, + result: Result<(), TransferError>, + }, +} + +/// A part's reader that reports the bytes it gives, and fails once the upload is stopped, which +/// aborts the request sending it. +struct CountingReader { + inner: Box, + file: usize, + events: Sender, + stop: Arc, +} - if observer.is_cancelled() { - return Err(UploadError::Cancelled { - rel_path: file.rel_path.clone(), +impl Read for CountingReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if self.stop.load(Ordering::Relaxed) { + return Err(io::Error::other("the upload was stopped")); + } + let read = self.inner.read(buf)?; + if read > 0 { + let _ = self.events.send(PartEvent::Sent { + file: self.file, + bytes: read as u64, }); } - - upload_source_file_multipart_streaming( - client, - source, - &file.rel_path, - &file.parts, - observer, - )?; + Ok(read) } - - Ok(()) } -fn upload_source_file_multipart_streaming< - FTC: FileTransferClient, - S: MultipartUploadSource, - O: TransferObserver, ->( - client: &FTC, - source: &S, - rel_path: &str, - parts: &[MultipartUploadPart], - observer: &mut O, -) -> Result<(), UploadError> { - let file_len = source.file_len(rel_path)?; - observer.file_started(rel_path, Some(file_len)); +/// The calling thread's view of an upload: what each file has sent, what remains, and the first +/// failure. +struct UploadRun<'p, 'a, O: TransferObserver> { + plan: &'p [PlannedFile<'a>], + observer: &'p mut O, + stop: Arc, + sent: Vec, + parts_left: Vec, + handed_over: usize, + finished: usize, + failure: Option, +} - let mut part_indices: Vec = (0..parts.len()).collect(); - part_indices.sort_by_key(|&i| parts[i].part); +impl<'p, 'a, O: TransferObserver> UploadRun<'p, 'a, O> { + fn new(plan: &'p [PlannedFile<'a>], observer: &'p mut O, stop: Arc) -> Self { + Self { + plan, + observer, + stop, + sent: vec![0; plan.len()], + parts_left: plan.iter().map(|file| file.parts.len()).collect(), + handed_over: 0, + finished: 0, + failure: None, + } + } - for (i, &part_idx) in part_indices.iter().enumerate() { - let part = &parts[part_idx]; - if part.part != (i as u32 + 1) { - return Err(UploadError::InvalidMultipart(format!( - "Invalid part numbering for {}: expected {}, got {}", - rel_path, - i + 1, - part.part - ))); + /// Announces `file`, completing it at once when it has no parts. False once stopped. + fn start(&mut self, file: usize) -> bool { + self.check_cancelled(file); + if self.stopped() { + return false; + } + let planned = &self.plan[file]; + self.observer + .file_started(planned.rel_path, Some(planned.len)); + if planned.parts.is_empty() { + self.observer.file_completed(planned.rel_path, 0); } + true } - let mut offset = 0u64; + /// Hands `job` to a worker, handling reports while all are busy. False once stopped. + fn hand_over( + &mut self, + job: PartJob<'a>, + jobs: &Sender>, + reports: &Receiver, + ) -> bool { + let mut job = Some(job); + while let Some(pending) = job.take() { + let mut select = Select::new(); + let send = select.send(jobs); + let receive = select.recv(reports); + let operation = select.select(); + match operation.index() { + index if index == send => { + if operation.send(jobs, pending).is_err() { + return false; + } + self.handed_over += 1; + } + index if index == receive => { + job = Some(pending); + match operation.recv(reports) { + Ok(event) => self.handle(event), + Err(_) => return false, + } + if self.stopped() { + return false; + } + } + _ => unreachable!("only two operations are selected"), + } + } + true + } - for (part_index, &part_idx) in part_indices.iter().enumerate() { - let part = &parts[part_idx]; - let size = part.size_bytes; + /// Whether a part handed to a worker has not finished yet. + fn in_flight(&self) -> bool { + self.finished < self.handed_over + } - if offset + size > file_len { - return Err(UploadError::InvalidMultipart(format!( - "Part {} exceeds file length for {}", - part_index + 1, - rel_path - ))); + /// Takes in a worker's report. + fn handle(&mut self, event: PartEvent) { + if let PartEvent::Finished { .. } = event { + self.finished += 1; } + let file = match event { + PartEvent::Sent { file, bytes } => { + self.sent[file] += bytes; + if self.failure.is_none() { + self.observer + .file_progress(self.plan[file].rel_path, self.sent[file]); + } + file + } + PartEvent::Finished { + file, + result: Ok(()), + .. + } => { + self.parts_left[file] -= 1; + if self.parts_left[file] == 0 && self.failure.is_none() { + let planned = &self.plan[file]; + self.observer.file_completed(planned.rel_path, planned.len); + } + file + } + PartEvent::Finished { + file, + part, + result: Err(source), + } => { + let planned = &self.plan[file]; + self.fail(UploadError::Transfer { + part_index: part, + total_parts: planned.parts.len(), + rel_path: planned.rel_path.to_string(), + source, + }); + file + } + }; + self.check_cancelled(file); + } - if observer.is_cancelled() { - return Err(UploadError::Cancelled { - rel_path: rel_path.to_string(), + fn check_cancelled(&mut self, file: usize) { + if self.failure.is_none() && self.observer.is_cancelled() { + self.fail(UploadError::Cancelled { + rel_path: self.plan[file].rel_path.to_string(), }); } + } - let reader = source.open_part(rel_path, offset, size)?; - client - .put_reader(&part.url, reader, size) - .map_err(|e| UploadError::Transfer { - part_index: part_index + 1, - total_parts: parts.len(), - rel_path: rel_path.to_string(), - source: e, - })?; - - offset += size; - observer.file_progress(rel_path, offset); + /// Stops the upload on `error`, unless it already stopped on another. + fn fail(&mut self, error: UploadError) { + if self.failure.is_none() { + self.failure = Some(error); + } + self.stop.store(true, Ordering::Relaxed); } - if offset != file_len { - return Err(UploadError::InvalidMultipart(format!( - "Multipart size mismatch for {} (uploaded {}, expected {})", - rel_path, offset, file_len - ))); + fn stopped(&self) -> bool { + self.failure.is_some() } - observer.file_completed(rel_path, offset); - Ok(()) + fn finish(self) -> Result<(), UploadError> { + match self.failure { + Some(error) => Err(error), + None => Ok(()), + } + } } #[cfg(test)] @@ -391,10 +652,230 @@ mod tests { upload_bundle_multipart_with_client(&client, &source, &files) .expect("valid multipart plan should upload"); - let puts = client.puts.lock().expect("lock puts"); - assert_eq!(puts.len(), 3); - assert_eq!(puts[0], ("u1".to_string(), 2, b"ab".to_vec())); - assert_eq!(puts[1], ("u2".to_string(), 2, b"cd".to_vec())); - assert_eq!(puts[2], ("u3".to_string(), 2, b"ef".to_vec())); + let mut puts = client.puts.lock().expect("lock puts").clone(); + puts.sort(); + assert_eq!( + puts, + [ + ("u1".to_string(), 2, b"ab".to_vec()), + ("u2".to_string(), 2, b"cd".to_vec()), + ("u3".to_string(), 2, b"ef".to_vec()), + ] + ); + } + + fn two_part_file(rel_path: &str) -> MultipartUploadFile { + MultipartUploadFile { + rel_path: rel_path.to_string(), + parts: vec![ + MultipartUploadPart { + part: 1, + url: format!("{rel_path}/1"), + size_bytes: 2, + }, + MultipartUploadPart { + part: 2, + url: format!("{rel_path}/2"), + size_bytes: 2, + }, + ], + } + } + + /// Holds each upload until another one is in flight, or a while has passed, and records the + /// most uploads in flight at once. + #[derive(Clone, Default)] + struct OverlapClient { + in_flight: Arc>, + most_in_flight: Arc>, + } + + impl FileTransferClient for OverlapClient { + fn put_reader( + &self, + _url: &str, + mut reader: R, + _size_bytes: u64, + ) -> Result<(), TransferError> { + let mut sink = Vec::new(); + reader.read_to_end(&mut sink).expect("read the part"); + *self.in_flight.lock().unwrap() += 1; + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let in_flight = *self.in_flight.lock().unwrap(); + let mut most = self.most_in_flight.lock().unwrap(); + *most = (*most).max(in_flight); + if *most > 1 || std::time::Instant::now() > deadline { + break; + } + drop(most); + std::thread::sleep(std::time::Duration::from_millis(5)); + } + *self.in_flight.lock().unwrap() -= 1; + Ok(()) + } + + fn get_reader( + &self, + _url: &str, + _expected_size_bytes: Option, + ) -> Result, TransferError> { + unreachable!("uploads only") + } + } + + #[test] + fn uploads_parts_concurrently() { + let client = OverlapClient::default(); + let source = MockSource::new(HashMap::from([( + "weights.bin".to_string(), + b"abcd".to_vec(), + )])); + + upload_bundle_multipart_with_client(&client, &source, &[two_part_file("weights.bin")]) + .expect("valid multipart plan should upload"); + + assert!(*client.most_in_flight.lock().unwrap() > 1); + } + + #[derive(Default)] + struct RecordingObserver { + started: Vec<(String, Option)>, + progress: HashMap, + completed: Vec<(String, u64)>, + cancel: bool, + } + + impl TransferObserver for RecordingObserver { + fn is_cancelled(&self) -> bool { + self.cancel + } + + fn file_started(&mut self, rel_path: &str, total_bytes: Option) { + self.started.push((rel_path.to_string(), total_bytes)); + } + + fn file_progress(&mut self, rel_path: &str, transferred_bytes: u64) { + self.progress + .insert(rel_path.to_string(), transferred_bytes); + } + + fn file_completed(&mut self, rel_path: &str, transferred_bytes: u64) { + self.completed + .push((rel_path.to_string(), transferred_bytes)); + } + } + + #[test] + fn reports_each_file_from_start_to_completion() { + let client = MockClient::default(); + let source = MockSource::new(HashMap::from([ + ("a.bin".to_string(), b"abcd".to_vec()), + ("b.bin".to_string(), b"efgh".to_vec()), + ])); + let mut observer = RecordingObserver::default(); + + upload_bundle_multipart_with_client_and_observer( + &client, + &source, + &[two_part_file("a.bin"), two_part_file("b.bin")], + &mut observer, + ) + .expect("valid multipart plan should upload"); + + let mut completed = observer.completed.clone(); + completed.sort(); + assert_eq!( + completed, + [("a.bin".to_string(), 4), ("b.bin".to_string(), 4)] + ); + assert_eq!( + observer.progress, + HashMap::from([("a.bin".to_string(), 4), ("b.bin".to_string(), 4)]) + ); + } + + /// Fails the upload to one URL. + #[derive(Clone)] + struct FailingClient { + failing_url: &'static str, + } + + impl FileTransferClient for FailingClient { + fn put_reader( + &self, + url: &str, + mut reader: R, + _size_bytes: u64, + ) -> Result<(), TransferError> { + let mut sink = Vec::new(); + if reader.read_to_end(&mut sink).is_err() { + return Err(TransferError::Transport("the part was aborted".to_string())); + } + if url == self.failing_url { + return Err(TransferError::Transport("connection reset".to_string())); + } + Ok(()) + } + + fn get_reader( + &self, + _url: &str, + _expected_size_bytes: Option, + ) -> Result, TransferError> { + unreachable!("uploads only") + } + } + + #[test] + fn a_failed_part_fails_the_upload_naming_the_part() { + let client = FailingClient { + failing_url: "weights.bin/2", + }; + let source = MockSource::new(HashMap::from([( + "weights.bin".to_string(), + b"abcd".to_vec(), + )])); + + let err = + upload_bundle_multipart_with_client(&client, &source, &[two_part_file("weights.bin")]) + .expect_err("a failed part fails the upload"); + + match err { + UploadError::Transfer { + part_index, + total_parts, + rel_path, + .. + } => assert_eq!( + (part_index, total_parts, rel_path.as_str()), + (2, 2, "weights.bin") + ), + other => panic!("unexpected error: {other:?}"), + } + } + + #[test] + fn a_cancelled_upload_sends_nothing() { + let client = MockClient::default(); + let source = MockSource::new(HashMap::from([( + "weights.bin".to_string(), + b"abcd".to_vec(), + )])); + let mut observer = RecordingObserver { + cancel: true, + ..RecordingObserver::default() + }; + + let err = upload_bundle_multipart_with_client_and_observer( + &client, + &source, + &[two_part_file("weights.bin")], + &mut observer, + ) + .expect_err("a cancelled upload fails"); + + assert!(matches!(err, UploadError::Cancelled { .. }), "{err:?}"); + assert!(client.puts.lock().unwrap().is_empty()); } }