feat: transfers through tracel-artifact, concurrent and with byte progress

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" }

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

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 {}.",

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<String, u64> = 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<MultipartUploadFile> = 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));

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)