diff --git a/Cargo.lock b/Cargo.lock index e71d760..cb13b67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -527,6 +527,28 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e71406cd8807725f7ac2f999a4cdd32e98f829fdf65f528343cebf945e41df1e" +dependencies = [ + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-epoch", + "crossbeam-queue", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.7" @@ -546,6 +568,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03e8bd762f7479489c70ed6c768ddca99d7296857de437a68dcb2a94365b3fae" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.22" @@ -3058,6 +3089,19 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" +[[package]] +name = "tracel-artifact" +version = "0.11.0" +source = "git+https://github.com/tracel-ai/tracel?rev=f52ee9e5c76153abf2ca750c31427d44dd86cc0b#f52ee9e5c76153abf2ca750c31427d44dd86cc0b" +dependencies = [ + "crossbeam", + "reqwest", + "serde", + "sha2", + "tempfile", + "thiserror 2.0.20", +] + [[package]] name = "tracel-cli" version = "0.5.1" @@ -3081,6 +3125,7 @@ dependencies = [ "thiserror 2.0.20", "toml", "toml_edit", + "tracel-artifact", "tracel-client", "tracel-job", "url", @@ -3091,8 +3136,9 @@ dependencies = [ [[package]] name = "tracel-client" -version = "0.21.0" -source = "git+https://github.com/tracel-ai/burn-central-client?rev=8b909baa569effcd4bf336a1256e23be25036f01#8b909baa569effcd4bf336a1256e23be25036f01" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "456275c8a4a1fd65e147aac5e39a34de32070ad8efe03929cb7cf6104fe2f407" dependencies = [ "directories", "reqwest", @@ -3110,7 +3156,7 @@ dependencies = [ [[package]] name = "tracel-job" version = "0.11.0" -source = "git+https://github.com/tracel-ai/tracel?rev=84d5b9a3f6c9c1fee2c3401073a50c56bdc536b4#84d5b9a3f6c9c1fee2c3401073a50c56bdc536b4" +source = "git+https://github.com/tracel-ai/tracel?rev=f52ee9e5c76153abf2ca750c31427d44dd86cc0b#f52ee9e5c76153abf2ca750c31427d44dd86cc0b" dependencies = [ "clap", "clap_complete", diff --git a/Cargo.toml b/Cargo.toml index be58a01..c2e61ba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,8 +48,9 @@ rstest = { version = "0.27.0", default-features = false } ## External Crate -tracel-client = { git = "https://github.com/tracel-ai/burn-central-client", rev = "8b909baa569effcd4bf336a1256e23be25036f01" } -tracel-job = { git = "https://github.com/tracel-ai/tracel", rev = "84d5b9a3f6c9c1fee2c3401073a50c56bdc536b4" } +tracel-client = "0.22.0" +tracel-artifact = { git = "https://github.com/tracel-ai/tracel", rev = "f52ee9e5c76153abf2ca750c31427d44dd86cc0b" } +tracel-job = { git = "https://github.com/tracel-ai/tracel", rev = "f52ee9e5c76153abf2ca750c31427d44dd86cc0b" } ## Crate tracel-cli = { version = "0.5.0", path = "crates/tracel-cli" } diff --git a/crates/tracel-cli/Cargo.toml b/crates/tracel-cli/Cargo.toml index d75b3c4..10964ff 100644 --- a/crates/tracel-cli/Cargo.toml +++ b/crates/tracel-cli/Cargo.toml @@ -17,6 +17,7 @@ path = "src/bin/tracel.rs" [dependencies] tracel-client.workspace = true +tracel-artifact.workspace = true tracel-job = { workspace = true, features = ["cli"] } anyhow.workspace = true clap.workspace = true diff --git a/crates/tracel-cli/src/commands/artifacts.rs b/crates/tracel-cli/src/commands/artifacts.rs index 38f899b..40415da 100644 --- a/crates/tracel-cli/src/commands/artifacts.rs +++ b/crates/tracel-cli/src/commands/artifacts.rs @@ -192,7 +192,17 @@ pub fn handle_command(args: ArtifactsArgs, context: CliContext) -> anyhow::Resul }; let response = client .presign_artifact_download(namespace, name, args.experiment, &artifact.id) - .map_err(map_error)?; + .map_err(|error| { + map_resource_error( + error, + namespace, + name, + Resource::Artifact { + experiment: args.experiment, + artifact: &artifact.name, + }, + ) + })?; let files: Vec<_> = response .files .into_iter() @@ -206,8 +216,7 @@ pub fn handle_command(args: ArtifactsArgs, context: CliContext) -> anyhow::Resul } }) .collect(); - let files = - download_files(&client, context.terminal(), &directory, &files, args.force)?; + let files = download_files(context.terminal(), &directory, &files, args.force)?; let bytes = files.iter().map(|file| file.bytes).sum(); context.terminal().outro(&format!( "Downloaded artifact '{}' ({}) to {}.", diff --git a/crates/tracel-cli/src/commands/models.rs b/crates/tracel-cli/src/commands/models.rs index 55f78a6..6092b02 100644 --- a/crates/tracel-cli/src/commands/models.rs +++ b/crates/tracel-cli/src/commands/models.rs @@ -5,6 +5,7 @@ use std::path::{Path, PathBuf}; use clap::{Args, Subcommand}; use serde::Serialize; use serde_json::Value; +use tracel_artifact::upload::{MultipartUploadFile, MultipartUploadPart}; use tracel_client::console::model::request::{ ModelFileSpecRequest, ModelVersionListState, PromoteModelVersionRequest, RequestModelVersionUploadRequest, SetModelAliasRequest, @@ -19,8 +20,8 @@ use crate::commands::login::get_client_and_login_if_needed; use crate::context::CliContext; use crate::error::{CliError, ErrorKind}; use crate::helpers::{ - DownloadFile, DownloadResult, Resource, build_part_tasks, download_files, ensure_model_exists, - map_resource_error, parse_metadata, resolve_namespace_project, select_artifact, upload_parts, + DownloadFile, DownloadResult, Resource, download_files, ensure_model_exists, + map_resource_error, parse_metadata, resolve_namespace_project, select_artifact, upload_files, validate_auto_create, validate_rel_path, }; use crate::tools::fs::{build_file_specs, collect_files}; @@ -533,7 +534,7 @@ fn pull( checksum: Some(file.checksum), }) .collect(); - let files = download_files(&client, context.terminal(), &directory, &files, args.force)?; + let files = download_files(context.terminal(), &directory, &files, args.force)?; let bytes = files.iter().map(|file| file.bytes).sum(); context.terminal().outro(&format!( "Downloaded '{}' v{} ({}) to {}.", @@ -579,11 +580,6 @@ fn push( let file_specs = build_file_specs(&files)?; spinner.clear(); - let file_sizes: BTreeMap = file_specs - .iter() - .map(|f| (f.rel_path.clone(), f.size_bytes)) - .collect(); - let spinner = context.terminal().spinner("Starting the upload"); let upload_request = RequestModelVersionUploadRequest { files: file_specs @@ -601,8 +597,28 @@ fn push( .map_err(|e| map_resource_error(e, namespace, name, Resource::Model(&args.model_name)))?; spinner.clear(); - let tasks = build_part_tasks(&files, &file_sizes, &upload.files)?; - upload_parts(&client, tasks, context.terminal())?; + let plan: Vec = upload + .files + .into_iter() + .map(|file| MultipartUploadFile { + rel_path: file.rel_path, + parts: file + .urls + .parts + .into_iter() + .map(|part| MultipartUploadPart { + part: part.part, + url: part.url, + size_bytes: part.size_bytes, + }) + .collect(), + }) + .collect(); + let message = format!( + "Uploading {}", + counted(files.len() as u64, ["file", "files"]) + ); + upload_files(context.terminal(), message, &files, &plan)?; let map_error = |error| map_resource_error(error, namespace, name, Resource::Model(&args.model_name)); diff --git a/crates/tracel-cli/src/commands/package/code_version.rs b/crates/tracel-cli/src/commands/package/code_version.rs index 43f99d8..d8679aa 100644 --- a/crates/tracel-cli/src/commands/package/code_version.rs +++ b/crates/tracel-cli/src/commands/package/code_version.rs @@ -1,9 +1,10 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::sync::Arc; use anyhow::Context; use serde::Serialize; +use tracel_artifact::upload::{MultipartUploadFile, MultipartUploadPart}; use tracel_client::console::Client; use tracel_client::console::project::request::{ PublishArtifactRequest, PublishBinaryRequest, PublishProjectVersionRequest, @@ -12,6 +13,7 @@ use tracel_client::console::project::request::{ use super::{MODE_HINT, MODE_NEEDED, Mode, PackageArgs, check_targets_allowed}; use crate::context::CliContext; +use crate::helpers::upload_files; use crate::tools::build_driver::{self, BuildDriver}; use crate::tools::cargo::{self, Binary}; use crate::tools::fs::file_sha256_and_size; @@ -283,15 +285,28 @@ impl BuiltPackage { return Ok(version); } - let spinner = context.terminal().spinner("Uploading the code version"); + let mut files = BTreeMap::new(); + let mut plan = Vec::with_capacity(uploads.len()); for (key, url, path) in uploads { - let bytes = - std::fs::read(&path).with_context(|| format!("Can't read '{}'", path.display()))?; - client - .upload_bytes_to_url(url, bytes) - .with_context(|| format!("Can't upload '{key}'"))?; + let size_bytes = std::fs::metadata(&path) + .with_context(|| format!("Can't read '{}'", path.display()))? + .len(); + plan.push(MultipartUploadFile { + rel_path: key.clone(), + parts: vec![MultipartUploadPart { + part: 1, + url: url.to_string(), + size_bytes, + }], + }); + files.insert(key, path); } - spinner.clear(); + upload_files( + context.terminal(), + "Uploading the code version", + &files, + &plan, + )?; client .complete_project_version_upload(&linked.namespace, &linked.name, &version.version_id) diff --git a/crates/tracel-cli/src/helpers/download.rs b/crates/tracel-cli/src/helpers/download.rs index 68ae8ad..442ff32 100644 --- a/crates/tracel-cli/src/helpers/download.rs +++ b/crates/tracel-cli/src/helpers/download.rs @@ -1,14 +1,15 @@ -use std::fs::{self, File}; -use std::io::Write; use std::path::{Component, Path, PathBuf}; -use anyhow::Context; use serde::Serialize; -use sha2::{Digest, Sha256}; -use tracel_client::console::Client; +use tracel_artifact::ReqwestTransferClient; +use tracel_artifact::bundle::FsBundle; +use tracel_artifact::download::{ + ArtifactDownloadFile, download_artifacts_to_sink_with_client_and_observer, +}; +use super::transfer::{TransferProgress, download_failure}; use crate::error::{CliError, ErrorKind}; -use crate::ui::{Terminal, counted, short_digest}; +use crate::ui::{Terminal, counted}; /// A file to download into a directory, with the size and checksum to check it against. pub struct DownloadFile { @@ -61,133 +62,67 @@ fn check_destinations(directory: &Path, files: &[DownloadFile], force: bool) -> Ok(()) } -struct HashingWriter { - inner: W, - hasher: Sha256, - bytes: u64, -} - -impl HashingWriter { - fn new(inner: W) -> Self { - Self { - inner, - hasher: Sha256::new(), - bytes: 0, - } - } - - fn verify(&self, file: &DownloadFile) -> anyhow::Result<()> { - if let Some(size) = file.size_bytes.filter(|size| *size != 0) - && self.bytes != size - { - anyhow::bail!( - "Size mismatch for '{}': expected {size} bytes, received {}.", - file.rel_path, - self.bytes - ); - } - if let Some(checksum) = file - .checksum - .as_deref() - .filter(|checksum| !checksum.is_empty()) - { - let actual = format!("{:x}", self.hasher.clone().finalize()); - if !actual.eq_ignore_ascii_case(checksum) { - anyhow::bail!( - "Checksum mismatch for '{}': expected {}, received {}.", - file.rel_path, - short_digest(checksum), - short_digest(&actual) - ); - } - } - Ok(()) - } -} - -impl Write for HashingWriter { - fn write(&mut self, buffer: &[u8]) -> std::io::Result { - let written = self.inner.write(buffer)?; - self.hasher.update(&buffer[..written]); - self.bytes += written as u64; - Ok(written) - } - - fn flush(&mut self) -> std::io::Result<()> { - self.inner.flush() - } -} - -/// Download `files` into `directory` under a spinner, after checking every path and, +/// Download `files` into `directory` under a progress bar, after checking every path and, /// unless `force`, that none of them exists. +/// +/// The files are downloaded beside `directory`, and checked against their size and +/// checksum, before any of them moves into it, so a failed download leaves nothing in it. pub fn download_files( - client: &Client, terminal: &Terminal, directory: &Path, files: &[DownloadFile], force: bool, ) -> anyhow::Result> { check_destinations(directory, files, force)?; - let spinner = terminal.spinner(format!( + let staging_parent = match directory.parent() { + Some(parent) if !parent.as_os_str().is_empty() => parent, + _ => Path::new("."), + }; + let mut staging = FsBundle::temp_in(staging_parent, ".tracel-download-")?; + let downloads: Vec = files + .iter() + .map(|file| ArtifactDownloadFile { + rel_path: file.rel_path.clone(), + url: file.url.clone(), + size_bytes: file.size_bytes, + checksum: file.checksum.clone(), + }) + .collect(); + let total = files.iter().filter_map(|file| file.size_bytes).sum(); + let message = format!( "Downloading {}", counted(files.len() as u64, ["file", "files"]) - )); - let results = download_each(client, directory, files)?; - spinner.clear(); - Ok(results) -} - -fn download_each( - client: &Client, - directory: &Path, - files: &[DownloadFile], -) -> anyhow::Result> { - let mut results = Vec::with_capacity(files.len()); - for file in files { - let path = directory.join(&file.rel_path); - let parent = path.parent().expect("Validated relative file has a parent"); - fs::create_dir_all(parent) - .with_context(|| format!("Can't create '{}'", parent.display()))?; - let mut name = path - .file_name() - .expect("Validated relative file has a name") - .to_os_string(); - name.push(".part"); - let part = path.with_file_name(name); - let output = - File::create(&part).with_context(|| format!("Can't create '{}'", part.display()))?; - let result = (|| -> anyhow::Result { - let mut writer = HashingWriter::new(output); - client.download_from_url(&file.url, &mut writer)?; - writer.flush()?; - writer.verify(file)?; - let bytes = writer.bytes; - drop(writer); - fs::rename(&part, &path).with_context(|| format!("Can't save '{}'", path.display()))?; - Ok(bytes) - })(); - let bytes = result.inspect_err(|_| { - let _ = fs::remove_file(&part); - })?; - results.push(DownloadResult { + ); + let mut progress = TransferProgress::new(terminal.transfer(message, total)); + download_artifacts_to_sink_with_client_and_observer( + &ReqwestTransferClient::new(), + &mut staging, + &downloads, + &mut progress, + ) + .map_err(download_failure)?; + progress.finish(); + staging.move_into(directory)?; + + Ok(files + .iter() + .map(|file| DownloadResult { rel_path: file.rel_path.clone(), - path, - bytes, - }); - } - Ok(results) + path: directory.join(&file.rel_path), + bytes: progress.moved(&file.rel_path).unwrap_or_default(), + }) + .collect()) } #[cfg(test)] mod tests { + use std::fs; + use rstest::rstest; use super::*; use crate::error::classify; - const HELLO_WORLD_SHA256: &str = - "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"; - fn file(rel_path: &str, size_bytes: Option, checksum: Option<&str>) -> DownloadFile { DownloadFile { rel_path: rel_path.into(), @@ -197,26 +132,6 @@ mod tests { } } - fn written(contents: &[u8]) -> HashingWriter> { - let mut writer = HashingWriter::new(Vec::new()); - writer.write_all(contents).unwrap(); - writer - } - - struct TwoBytesAtATime(Vec); - - impl Write for TwoBytesAtATime { - fn write(&mut self, data: &[u8]) -> std::io::Result { - let count = data.len().min(2); - self.0.extend_from_slice(&data[..count]); - Ok(count) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - } - #[rstest] #[case::file("weights.bin")] #[case::nested("nested/weights.bin")] @@ -243,63 +158,6 @@ mod tests { assert!(validated.is_err()); } - #[test] - fn given_a_writer_that_takes_two_bytes_at_a_time_when_writing_then_only_written_bytes_are_hashed() - { - let mut writer = HashingWriter::new(TwoBytesAtATime(Vec::new())); - - writer.write_all(b"hello world").unwrap(); - - assert_eq!(writer.inner.0, b"hello world"); - assert_eq!(writer.bytes, 11); - assert_eq!( - format!("{:x}", writer.hasher.finalize()), - HELLO_WORLD_SHA256 - ); - } - - #[test] - fn given_the_declared_size_and_checksum_when_verifying_a_download_then_it_passes() { - let writer = written(b"hello world"); - let declared = file("hello", Some(11), Some(&HELLO_WORLD_SHA256.to_uppercase())); - - let verified = writer.verify(&declared); - - assert!(verified.is_ok()); - } - - #[rstest] - #[case::none(None, None)] - #[case::zero_and_empty(Some(0), Some(""))] - fn given_no_declared_size_or_checksum_when_verifying_a_download_then_it_passes( - #[case] size_bytes: Option, - #[case] checksum: Option<&str>, - ) { - let writer = written(b"hello world"); - - let verified = writer.verify(&file("hello", size_bytes, checksum)); - - assert!(verified.is_ok()); - } - - #[test] - fn given_another_size_when_verifying_a_download_then_it_fails() { - let writer = written(b"hello world"); - - let verified = writer.verify(&file("hello", Some(10), None)); - - assert!(verified.is_err()); - } - - #[test] - fn given_another_checksum_when_verifying_a_download_then_it_fails() { - let writer = written(b"hello world"); - - let verified = writer.verify(&file("hello", None, Some("0123abcd"))); - - assert!(verified.is_err()); - } - #[test] fn given_an_existing_file_when_checking_destinations_then_it_is_a_conflict() { let directory = tempfile::tempdir().unwrap(); diff --git a/crates/tracel-cli/src/helpers/mod.rs b/crates/tracel-cli/src/helpers/mod.rs index 2de94ac..3ab5797 100644 --- a/crates/tracel-cli/src/helpers/mod.rs +++ b/crates/tracel-cli/src/helpers/mod.rs @@ -1,15 +1,16 @@ mod download; mod jobs; -mod model_upload; +mod model; mod project; mod resources; +mod transfer; pub use download::{DownloadFile, DownloadResult, download_files, validate_rel_path}; pub use jobs::{ FinalStatus, final_status, follow_job, follow_logs, job_command, job_invocation, missing_job, poll_job, read_job, require_completed, }; -pub use model_upload::{build_part_tasks, ensure_model_exists, upload_parts}; +pub use model::ensure_model_exists; pub use project::{ ProjectSource, can_initialize_project, missing_project, parse_project, require_cargo_workspace, require_linked_project, require_workspace_project, resolve_namespace_project, resolve_project, @@ -18,3 +19,4 @@ pub use project::{ pub use resources::{ Resource, map_resource_error, parse_metadata, select_artifact, validate_auto_create, }; +pub use transfer::upload_files; diff --git a/crates/tracel-cli/src/helpers/model_upload/existence.rs b/crates/tracel-cli/src/helpers/model.rs similarity index 97% rename from crates/tracel-cli/src/helpers/model_upload/existence.rs rename to crates/tracel-cli/src/helpers/model.rs index a2e6634..f944658 100644 --- a/crates/tracel-cli/src/helpers/model_upload/existence.rs +++ b/crates/tracel-cli/src/helpers/model.rs @@ -1,3 +1,5 @@ +//! Making sure a model exists before a version is pushed to it. + use anyhow::Context; use tracel_client::console::Client; use tracel_client::console::model::request::CreateModelRequest; diff --git a/crates/tracel-cli/src/helpers/model_upload/mod.rs b/crates/tracel-cli/src/helpers/model_upload/mod.rs deleted file mode 100644 index c6c1b56..0000000 --- a/crates/tracel-cli/src/helpers/model_upload/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -//! Uploading a model version: making sure the model exists, and uploading its parts in -//! parallel. - -mod existence; -mod parts; -mod upload; - -pub use existence::ensure_model_exists; -pub use parts::build_part_tasks; -pub use upload::upload_parts; diff --git a/crates/tracel-cli/src/helpers/model_upload/parts.rs b/crates/tracel-cli/src/helpers/model_upload/parts.rs deleted file mode 100644 index 0791b21..0000000 --- a/crates/tracel-cli/src/helpers/model_upload/parts.rs +++ /dev/null @@ -1,168 +0,0 @@ -use std::collections::BTreeMap; -use std::path::PathBuf; -use tracel_client::console::artifact::response::PresignedUploadUrlResponse; -use tracel_client::console::model::response::PresignedModelFileUploadUrlsResponse; - -/// One part of a file to upload: where its bytes are in the file, and the URL to upload -/// them to. -#[derive(Clone, Debug)] -pub struct PartUploadTask { - pub rel_path: String, - pub absolute_path: PathBuf, - pub part: u32, - pub url: String, - pub offset: u64, - pub size_bytes: u64, -} - -fn validate_part_sequence( - rel_path: &str, - parts: &[PresignedUploadUrlResponse], -) -> anyhow::Result<()> { - for (index, part) in parts.iter().enumerate() { - let expected = index as u32 + 1; - if part.part != expected { - anyhow::bail!( - "The Tracel Console numbered the parts of '{rel_path}' wrongly: expected part {expected}, got {}.", - part.part - ); - } - } - Ok(()) -} - -fn validate_part_total_size( - rel_path: &str, - parts: &[PresignedUploadUrlResponse], - expected_size: u64, -) -> anyhow::Result<()> { - let total: u64 = parts.iter().map(|part| part.size_bytes).sum(); - if total != expected_size { - anyhow::bail!( - "The Tracel Console's parts of '{rel_path}' sum to {total} bytes, not {expected_size}." - ); - } - Ok(()) -} - -/// Split each file of an upload response into its parts, checking that they number from -/// 1 and cover the file. -pub fn build_part_tasks( - files: &BTreeMap, - file_sizes: &BTreeMap, - upload_files: &[PresignedModelFileUploadUrlsResponse], -) -> anyhow::Result> { - let mut tasks = Vec::new(); - - for upload_file in upload_files { - let absolute_path = files.get(&upload_file.rel_path).ok_or_else(|| { - anyhow::anyhow!( - "The Tracel Console sent parts of an unknown file '{}'.", - upload_file.rel_path - ) - })?; - let expected_size = *file_sizes.get(&upload_file.rel_path).ok_or_else(|| { - anyhow::anyhow!( - "The Tracel Console sent parts of '{}', whose size is unknown.", - upload_file.rel_path - ) - })?; - - if upload_file.urls.parts.is_empty() { - anyhow::bail!( - "The Tracel Console sent no parts of '{}'.", - upload_file.rel_path - ); - } - let mut parts = upload_file.urls.parts.clone(); - parts.sort_by_key(|part| part.part); - validate_part_sequence(&upload_file.rel_path, &parts)?; - validate_part_total_size(&upload_file.rel_path, &parts, expected_size)?; - - let mut offset = 0u64; - for part in parts { - tasks.push(PartUploadTask { - rel_path: upload_file.rel_path.clone(), - absolute_path: absolute_path.clone(), - part: part.part, - url: part.url, - offset, - size_bytes: part.size_bytes, - }); - offset += part.size_bytes; - } - } - - Ok(tasks) -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - use tracel_client::console::artifact::response::MultipartUploadResponse; - - use super::*; - - fn weights(size: u64) -> (BTreeMap, BTreeMap) { - let files = BTreeMap::from([("weights.bin".to_string(), PathBuf::from("weights.bin"))]); - let sizes = BTreeMap::from([("weights.bin".to_string(), size)]); - (files, sizes) - } - - fn parts_of(rel_path: &str, parts: &[(u32, u64)]) -> Vec { - vec![PresignedModelFileUploadUrlsResponse { - rel_path: rel_path.to_string(), - urls: MultipartUploadResponse { - id: "upload-1".to_string(), - parts: parts - .iter() - .map(|&(part, size_bytes)| PresignedUploadUrlResponse { - part, - url: format!("https://storage/part{part}"), - size_bytes, - }) - .collect(), - }, - }] - } - - #[test] - fn given_parts_out_of_order_when_splitting_a_file_then_each_part_starts_where_the_previous_ends() - { - let (files, sizes) = weights(300); - - let tasks = build_part_tasks( - &files, - &sizes, - &parts_of("weights.bin", &[(2, 100), (1, 200)]), - ) - .unwrap(); - - assert_eq!( - (tasks[0].part, tasks[0].offset, tasks[0].size_bytes), - (1, 0, 200) - ); - assert_eq!( - (tasks[1].part, tasks[1].offset, tasks[1].size_bytes), - (2, 200, 100) - ); - } - - #[rstest] - #[case::unknown_file("other.bin", &[(1, 300)], "'other.bin'")] - #[case::no_parts("weights.bin", &[], "no parts")] - #[case::too_small("weights.bin", &[(1, 200)], "200 bytes")] - #[case::repeated_part("weights.bin", &[(1, 150), (1, 150)], "got 1")] - #[case::missing_part("weights.bin", &[(1, 200), (3, 100)], "got 3")] - fn given_parts_that_do_not_cover_the_file_when_splitting_it_then_the_upload_is_refused_for_that_reason( - #[case] rel_path: &str, - #[case] parts: &[(u32, u64)], - #[case] reason: &str, - ) { - let (files, sizes) = weights(300); - - let error = build_part_tasks(&files, &sizes, &parts_of(rel_path, parts)).unwrap_err(); - - assert!(error.to_string().contains(reason), "{error}"); - } -} diff --git a/crates/tracel-cli/src/helpers/model_upload/upload.rs b/crates/tracel-cli/src/helpers/model_upload/upload.rs deleted file mode 100644 index 0c313f8..0000000 --- a/crates/tracel-cli/src/helpers/model_upload/upload.rs +++ /dev/null @@ -1,233 +0,0 @@ -use anyhow::Context; -use std::collections::BTreeSet; -use std::fs::File; -use std::io::{Read, Seek, SeekFrom}; -use std::path::Path; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::time::Duration; -use tracel_client::ClientError; - -use crate::ui::{Terminal, counted}; - -use super::parts::PartUploadTask; - -const UPLOAD_WORKER_COUNT: usize = 8; - -/// Uploads one part's bytes to its presigned URL. -pub trait PartUploader: Send + Sync { - fn upload(&self, url: &str, bytes: Vec) -> Result<(), ClientError>; -} - -impl PartUploader for tracel_client::console::Client { - fn upload(&self, url: &str, bytes: Vec) -> Result<(), ClientError> { - self.upload_bytes_to_url(url, bytes) - } -} - -fn read_part_bytes(path: &Path, offset: u64, size: u64) -> anyhow::Result> { - let mut file = File::open(path).with_context(|| format!("Can't open '{}'", path.display()))?; - file.seek(SeekFrom::Start(offset)) - .with_context(|| format!("Can't seek in '{}'", path.display()))?; - - let mut buffer = vec![0u8; size as usize]; - file.read_exact(&mut buffer) - .with_context(|| format!("Can't read '{}'", path.display()))?; - - Ok(buffer) -} - -/// Split `tasks` into up to `worker_count` contiguous chunks, one per worker. -fn chunk_tasks(tasks: Vec, worker_count: usize) -> Vec> { - if tasks.is_empty() { - return Vec::new(); - } - - let chunk_size = tasks.len().div_ceil(worker_count); - let mut iter = tasks.into_iter(); - let mut chunks = Vec::new(); - - loop { - let chunk: Vec<_> = iter.by_ref().take(chunk_size).collect(); - if chunk.is_empty() { - break; - } - chunks.push(chunk); - } - - chunks -} - -/// Upload a worker's parts in order, counting each in `completed`, until one fails or -/// another worker's failure sets `cancelled`. -fn upload_chunk( - chunk: Vec, - uploader: &impl PartUploader, - cancelled: &AtomicBool, - completed: &AtomicU64, -) -> anyhow::Result<()> { - for task in chunk { - if cancelled.load(Ordering::Relaxed) { - break; - } - - let result = - read_part_bytes(&task.absolute_path, task.offset, task.size_bytes).and_then(|bytes| { - uploader - .upload(&task.url, bytes) - .map_err(|e| anyhow::anyhow!(e)) - }); - - match result { - Ok(()) => { - completed.fetch_add(1, Ordering::Relaxed); - } - Err(e) => { - cancelled.store(true, Ordering::Relaxed); - return Err(e.context(format!( - "Can't upload '{}' (part {})", - task.rel_path, task.part - ))); - } - } - } - - Ok(()) -} - -/// Upload `tasks` with up to 8 workers under one progress bar, stopping at the first -/// failure. -pub fn upload_parts( - uploader: &impl PartUploader, - tasks: Vec, - terminal: &Terminal, -) -> anyhow::Result<()> { - let total_parts = tasks.len(); - let files: BTreeSet<&str> = tasks.iter().map(|task| task.rel_path.as_str()).collect(); - let progress = terminal.progress( - format!( - "Uploading {}", - counted(files.len() as u64, ["file", "files"]) - ), - total_parts as u64, - ); - let worker_count = UPLOAD_WORKER_COUNT.min(total_parts.max(1)); - let chunks = chunk_tasks(tasks, worker_count); - - let cancelled = AtomicBool::new(false); - let completed = AtomicU64::new(0); - - let failure = std::thread::scope(|scope| { - let handles: Vec<_> = chunks - .into_iter() - .map(|chunk| { - let cancelled = &cancelled; - let completed = &completed; - scope.spawn(move || upload_chunk(chunk, uploader, cancelled, completed)) - }) - .collect(); - - let mut reported = 0u64; - while handles.iter().any(|handle| !handle.is_finished()) { - let done = completed.load(Ordering::Relaxed); - if done > reported { - progress.set_position(done); - reported = done; - } - std::thread::sleep(Duration::from_millis(50)); - } - - // Join each worker here, so a panic becomes an error instead of unwinding. - handles - .into_iter() - .filter_map(|handle| match handle.join() { - Ok(Ok(())) => None, - Ok(Err(e)) => Some(e), - Err(_) => Some(anyhow::anyhow!("An upload worker stopped unexpectedly.")), - }) - .next() - }); - - match failure { - Some(error) => Err(error), - None => { - progress.clear(); - Ok(()) - } - } -} - -#[cfg(test)] -mod tests { - use std::fs; - use std::sync::Mutex; - - use super::*; - use crate::test_support::quiet_terminal; - - struct RecordingUploader { - urls: Mutex>, - failing_url: Option<&'static str>, - } - - impl RecordingUploader { - fn new(failing_url: Option<&'static str>) -> Self { - Self { - urls: Mutex::new(Vec::new()), - failing_url, - } - } - } - - impl PartUploader for RecordingUploader { - fn upload(&self, url: &str, _bytes: Vec) -> Result<(), ClientError> { - self.urls.lock().unwrap().push(url.to_string()); - if self.failing_url == Some(url) { - return Err(ClientError::NotFound); - } - Ok(()) - } - } - - fn task(directory: &Path, name: &str) -> PartUploadTask { - let path = directory.join(name); - fs::write(&path, b"data").unwrap(); - PartUploadTask { - rel_path: name.to_string(), - absolute_path: path, - part: 1, - url: format!("https://storage/{name}"), - offset: 0, - size_bytes: 4, - } - } - - #[test] - fn given_two_parts_when_uploading_then_each_is_uploaded_once() { - let directory = tempfile::tempdir().unwrap(); - let tasks = vec![ - task(directory.path(), "a.bin"), - task(directory.path(), "b.bin"), - ]; - let uploader = RecordingUploader::new(None); - - upload_parts(&uploader, tasks, &quiet_terminal()).unwrap(); - - let mut urls = uploader.urls.into_inner().unwrap(); - urls.sort(); - assert_eq!(urls, ["https://storage/a.bin", "https://storage/b.bin"]); - } - - #[test] - fn given_a_part_that_fails_when_uploading_then_the_error_names_it() { - let directory = tempfile::tempdir().unwrap(); - let tasks = vec![ - task(directory.path(), "a.bin"), - task(directory.path(), "b.bin"), - ]; - let uploader = RecordingUploader::new(Some("https://storage/b.bin")); - - let error = upload_parts(&uploader, tasks, &quiet_terminal()).unwrap_err(); - - assert!(error.to_string().contains("'b.bin' (part 1)"), "{error}"); - } -} diff --git a/crates/tracel-cli/src/helpers/resources.rs b/crates/tracel-cli/src/helpers/resources.rs index 375cf3d..5387d46 100644 --- a/crates/tracel-cli/src/helpers/resources.rs +++ b/crates/tracel-cli/src/helpers/resources.rs @@ -11,10 +11,34 @@ pub enum Resource<'a> { ModelVersion { model: &'a str, version: u32 }, ModelAlias { model: &'a str, alias: &'a str }, Experiment(i32), + Artifact { experiment: i32, artifact: &'a str }, Dataset(&'a str), Job(i32), } +impl Resource<'_> { + /// Name the resource in a sentence, such as "artifact 'weights' of experiment 42". + fn named(self) -> String { + match self { + Self::Model(model) => format!("model '{model}'"), + Self::ModelVersionRef { model, reference } => { + format!("version '{reference}' of model '{model}'") + } + Self::ModelVersion { model, version } => { + format!("version {version} of model '{model}'") + } + Self::ModelAlias { model, alias } => format!("alias '{alias}' of model '{model}'"), + Self::Experiment(experiment) => format!("experiment {experiment}"), + Self::Artifact { + experiment, + artifact, + } => format!("artifact '{artifact}' of experiment {experiment}"), + Self::Dataset(dataset) => format!("dataset '{dataset}'"), + Self::Job(job) => format!("job {job}"), + } + } +} + /// Turn the Tracel Console's "not found" for `resource` into an error that names it and /// how to list what exists. Other errors stay as they are. pub fn map_resource_error( @@ -24,6 +48,13 @@ pub fn map_resource_error( resource: Resource<'_>, ) -> anyhow::Error { let (message, hint) = match (&error, resource) { + (ClientError::NotFoundWithCode(ApiErrorCode::ObjectStorageFile), resource) => { + let message = format!( + "The files of {} are missing from the Tracel Console's storage.", + resource.named() + ); + return CliError::new(ErrorKind::NotFound, message).into(); + } ( ClientError::NotFoundWithCode(ApiErrorCode::Model), Resource::Model(model) @@ -62,12 +93,27 @@ pub fn map_resource_error( format!("List aliases with `tracel models alias list {model}`."), ), ( - ClientError::NotFound | ClientError::NotFoundWithCode(ApiErrorCode::Unknown), + ClientError::NotFound + | ClientError::NotFoundWithCode(ApiErrorCode::Unknown | ApiErrorCode::Experiment), Resource::Experiment(experiment), + ) + | ( + ClientError::NotFoundWithCode(ApiErrorCode::Experiment), + Resource::Artifact { experiment, .. }, ) => ( format!("No experiment {experiment} in {namespace}/{project}."), "List experiments with `tracel experiments list`.".to_string(), ), + ( + ClientError::NotFoundWithCode(ApiErrorCode::ExperimentArtifact), + Resource::Artifact { + experiment, + artifact, + }, + ) => ( + format!("No artifact '{artifact}' in experiment {experiment}."), + format!("List artifacts with `tracel artifacts list {experiment}`."), + ), (ClientError::NotFoundWithCode(ApiErrorCode::Dataset), Resource::Dataset(dataset)) => ( format!("No dataset '{dataset}' in {namespace}/{project}."), "List datasets with `tracel datasets list`.".to_string(), @@ -201,6 +247,24 @@ mod tests { "job 12", "tracel jobs list" )] + #[case::experiment( + ApiErrorCode::Experiment, + Resource::Experiment(9), + "experiment 9", + "tracel experiments list" + )] + #[case::experiment_of_an_artifact( + ApiErrorCode::Experiment, + Resource::Artifact { experiment: 9, artifact: "weights" }, + "experiment 9", + "tracel experiments list" + )] + #[case::artifact( + ApiErrorCode::ExperimentArtifact, + Resource::Artifact { experiment: 9, artifact: "weights" }, + "'weights'", + "tracel artifacts list 9" + )] fn given_the_console_says_what_is_missing_when_mapping_the_error_then_it_names_it_and_how_to_list_them( #[case] code: ApiErrorCode, #[case] resource: Resource<'static>, @@ -237,6 +301,27 @@ mod tests { ); } + #[test] + fn given_the_stored_files_are_missing_when_mapping_the_error_then_it_says_so_of_the_resource() { + let error = map_resource_error( + ClientError::NotFoundWithCode(ApiErrorCode::ObjectStorageFile), + "alice", + "demo", + Resource::Artifact { + experiment: 9, + artifact: "weights", + }, + ); + + assert_eq!(classify(&error), ErrorKind::NotFound); + assert!( + error + .to_string() + .contains("artifact 'weights' of experiment 9"), + "{error}" + ); + } + #[rstest] #[case::unauthenticated(ClientError::Unauthenticated)] #[case::server_error(ClientError::InternalServerError)] diff --git a/crates/tracel-cli/src/helpers/transfer.rs b/crates/tracel-cli/src/helpers/transfer.rs new file mode 100644 index 0000000..44229b7 --- /dev/null +++ b/crates/tracel-cli/src/helpers/transfer.rs @@ -0,0 +1,181 @@ +//! Moving files to and from presigned URLs with tracel-artifact, under one progress bar of +//! the bytes moved. + +use std::collections::{BTreeMap, HashMap}; +use std::fs::File; +use std::io::{Read, Seek, SeekFrom}; +use std::path::PathBuf; + +use tracel_artifact::TransferObserver; +use tracel_artifact::download::DownloadError; +use tracel_artifact::upload::{ + MultipartUploadFile, MultipartUploadSource, UploadError, + upload_bundle_multipart_with_client_and_observer, +}; + +use crate::error::{CliError, ErrorKind}; +use crate::ui::{Progress, Terminal}; + +/// Upload `files`, keyed by the path each `plan` entry names, to the URLs of `plan`, under a +/// progress bar saying `message`. +pub fn upload_files( + terminal: &Terminal, + message: impl std::fmt::Display, + files: &BTreeMap, + plan: &[MultipartUploadFile], +) -> anyhow::Result<()> { + let total = plan + .iter() + .flat_map(|file| &file.parts) + .map(|part| part.size_bytes) + .sum(); + let mut progress = TransferProgress::new(terminal.transfer(message, total)); + upload_bundle_multipart_with_client_and_observer( + &tracel_artifact::ReqwestTransferClient::new(), + &LocalFiles(files), + plan, + &mut progress, + ) + .map_err(upload_failure)?; + progress.finish(); + Ok(()) +} + +/// Shows on one bar the bytes every file of a transfer has moved. +pub struct TransferProgress { + progress: Progress, + moved: HashMap, +} + +impl TransferProgress { + pub fn new(progress: Progress) -> Self { + Self { + progress, + moved: HashMap::new(), + } + } + + /// The bytes `rel_path` moved, once it has. + pub fn moved(&self, rel_path: &str) -> Option { + self.moved.get(rel_path).copied() + } + + /// Take the bar off the screen. + pub fn finish(&self) { + self.progress.clear(); + } + + fn record(&mut self, rel_path: &str, bytes: u64) { + self.moved.insert(rel_path.to_string(), bytes); + self.progress.set_position(self.moved.values().sum()); + } +} + +impl TransferObserver for TransferProgress { + fn file_progress(&mut self, rel_path: &str, transferred_bytes: u64) { + self.record(rel_path, transferred_bytes); + } + + fn file_completed(&mut self, rel_path: &str, transferred_bytes: u64) { + self.record(rel_path, transferred_bytes); + } +} + +/// Local files, by the path the Tracel Console knows each by, read a part at a time. +struct LocalFiles<'a>(&'a BTreeMap); + +impl LocalFiles<'_> { + fn path(&self, rel_path: &str) -> Result<&PathBuf, UploadError> { + self.0.get(rel_path).ok_or_else(|| { + UploadError::InvalidMultipart(format!( + "The Tracel Console asked for '{rel_path}', which isn't being uploaded." + )) + }) + } +} + +impl MultipartUploadSource for LocalFiles<'_> { + fn file_len(&self, rel_path: &str) -> Result { + let path = self.path(rel_path)?; + let metadata = std::fs::metadata(path).map_err(|error| UploadError::MultipartReader { + rel_path: rel_path.to_string(), + source: Box::new(error), + })?; + Ok(metadata.len()) + } + + fn open_part( + &self, + rel_path: &str, + offset: u64, + size: u64, + ) -> Result, UploadError> { + let unreadable = |error: std::io::Error| UploadError::MultipartReader { + rel_path: rel_path.to_string(), + source: Box::new(error), + }; + let mut file = File::open(self.path(rel_path)?).map_err(unreadable)?; + file.seek(SeekFrom::Start(offset)).map_err(unreadable)?; + Ok(Box::new(file.take(size))) + } +} + +/// A failed upload: `UNAVAILABLE` when a part couldn't be sent. +fn upload_failure(error: UploadError) -> anyhow::Error { + match error { + UploadError::Transfer { + part_index, + total_parts, + rel_path, + source, + } => CliError::new( + ErrorKind::Unavailable, + format!("Can't upload '{rel_path}' (part {part_index} of {total_parts}): {source}"), + ) + .into(), + error => anyhow::Error::new(error).context("Can't upload the files"), + } +} + +/// A failed download: `UNAVAILABLE` when a file couldn't be fetched. +pub fn download_failure(error: DownloadError) -> anyhow::Error { + match error { + DownloadError::Transfer { rel_path, source } => CliError::new( + ErrorKind::Unavailable, + format!("Can't download '{rel_path}': {source}"), + ) + .into(), + error => anyhow::Error::new(error).context("Can't download the files"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn given_a_file_when_opening_a_part_then_it_reads_only_that_range() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("weights.bin"); + std::fs::write(&path, b"abcdef").unwrap(); + let files = BTreeMap::from([("weights.bin".to_string(), path)]); + let mut part = String::new(); + + LocalFiles(&files) + .open_part("weights.bin", 2, 3) + .unwrap() + .read_to_string(&mut part) + .unwrap(); + + assert_eq!(part, "cde"); + } + + #[test] + fn given_a_file_not_being_uploaded_when_the_console_asks_for_it_then_the_upload_is_refused() { + let files = BTreeMap::new(); + + let length = LocalFiles(&files).file_len("other.bin"); + + assert!(matches!(length, Err(UploadError::InvalidMultipart(_)))); + } +} diff --git a/crates/tracel-cli/src/ui/mod.rs b/crates/tracel-cli/src/ui/mod.rs index b471217..3744d44 100644 --- a/crates/tracel-cli/src/ui/mod.rs +++ b/crates/tracel-cli/src/ui/mod.rs @@ -18,7 +18,7 @@ use screen::{Screen, Terminals}; pub use format::{Format, FormatArg}; pub use human::{Details, Human, Table, counted, list_summary, number, short_digest, size}; pub use layout::{Layout, Place}; -pub use stderr::{Declined, Terminal, Unanswered}; +pub use stderr::{Declined, Progress, Terminal, Unanswered}; pub use stdout::{Outcome, Output, Render, StdoutClosed}; pub use time::{Timestamp, elapsed, log_time, when}; diff --git a/crates/tracel-cli/src/ui/stderr.rs b/crates/tracel-cli/src/ui/stderr.rs index e835693..1a0d8dc 100644 --- a/crates/tracel-cli/src/ui/stderr.rs +++ b/crates/tracel-cli/src/ui/stderr.rs @@ -193,9 +193,11 @@ impl Terminal { self.start(cliclack::spinner(), message.to_string()) } - /// Start a bar that counts the `total` steps of one task. - pub fn progress(&self, message: impl Display, total: u64) -> Progress { - self.start(cliclack::progress_bar(total), message.to_string()) + /// Start a bar that counts the bytes of a transfer of `total_bytes`, with its rate and + /// time left. + pub fn transfer(&self, message: impl Display, total_bytes: u64) -> Progress { + let bar = cliclack::progress_bar(total_bytes).with_download_template(); + self.start(bar, message.to_string()) } fn start(&self, bar: ProgressBar, message: String) -> Progress { @@ -368,7 +370,7 @@ impl Progress { } } - /// Set how many of a bar's steps are done. + /// Set how much of a bar is done, such as the bytes a transfer has moved. pub fn set_position(&self, done: u64) { if let Some(id) = self.id { self.screen.set_position(id, done); diff --git a/docs/cli.md b/docs/cli.md index f14d28f..0f138b2 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -169,7 +169,14 @@ false when it already had them all. into `./` when the artifact's name is a plain file name; `-d, --directory ` chooses another directory. Every path and existing file is checked before anything downloads, and `--force` overwrites existing files. -Each file's size and SHA-256 checksum are checked before it is saved. +The files download into a temporary directory beside it, a large file in +several ranges at once, and each file's size and SHA-256 checksum are checked +before any of them moves in, so a failed download leaves the directory +unchanged. + +`models push` and the code versions `package` uploads send up to eight parts or +files at once. A transfer on a terminal shows the bytes moved, the rate and the +time left. ## Models