feat: upload progress bar

Author: drbh

kernel-builder/Cargo.lock

@@ -16,6 +16,7 @@ clap_complete = "4"
 eyre = "0.6.12"
 git2 = "0.20"
 huggingface-hub = { git = "https://github.com/huggingface/huggingface_hub_rust.git", rev = "8cbc662035e04d4be8e829316272893e980f5926", package = "huggingface-hub", features = ["blocking", "xet"] }
+indicatif = "0.17"
 itertools = "0.13"
 minijinja = "2.5"
 minijinja-embed = "2.5"

kernel-builder/Cargo.toml

@@ -107,6 +107,10 @@ enum Commands {
         /// Repository type on Hugging Face Hub (`model` or `kernel`).
         #[arg(long, value_enum, default_value_t = RepoTypeArg::Model)]
         repo_type: RepoTypeArg,
+
+        /// Suppress progress output.
+        #[arg(long, short)]
+        quiet: bool,
     },
 
     /// Upload kernel build artifacts to the Hugging Face Hub.
@@ -222,6 +226,7 @@ fn main() -> Result<()> {
             branch,
             private,
             repo_type,
+            quiet,
         } => {
             run_build(
                 kernel_dir.clone(),
@@ -236,6 +241,7 @@ fn main() -> Result<()> {
                 branch,
                 private,
                 repo_type,
+                quiet,
             })
         }
         Commands::CreatePyproject {

kernel-builder/src/main.rs

@@ -2,14 +2,16 @@ use std::{
     collections::{BTreeMap, HashSet},
     fs,
     path::{Path, PathBuf},
+    sync::Arc,
 };
 
 use clap::Args;
 use eyre::{bail, Context, Result};
 use huggingface_hub::{
-    AddSource, CommitOperation, CreateBranchParams, CreateCommitParams, CreateRepoParams,
-    ListRepoFilesParams, ListRepoRefsParams, RepoType,
+    AddSource, CommitOperation, CommitProgressCallback, CreateBranchParams, CreateCommitParams,
+    CreateRepoParams, ListRepoFilesParams, ListRepoRefsParams, RepoType,
 };
+use indicatif::{ProgressBar, ProgressStyle};
 use walkdir::WalkDir;
 
 use crate::{
@@ -60,6 +62,10 @@ pub struct UploadArgs {
     /// Repository type on Hugging Face Hub (`model` or `kernel`).
     #[arg(long, value_enum, default_value_t = RepoTypeArg::Model)]
     pub repo_type: RepoTypeArg,
+
+    /// Suppress progress output.
+    #[arg(long, short)]
+    pub quiet: bool,
 }
 
 pub fn run_upload(args: UploadArgs) -> Result<()> {
@@ -185,20 +191,30 @@ pub fn run_upload(args: UploadArgs) -> Result<()> {
             continue;
         }
 
-        eprintln!(
-            "Uploading {} operations to branch `{}`...",
-            operations.len(),
-            branch
-        );
-
         let batch_count = operations.len().div_ceil(BUILD_COMMIT_BATCH_SIZE);
-        if batch_count > 1 {
-            eprintln!(
-                "Uploading in {} commits ({} operations).",
-                batch_count,
-                operations.len()
+        let progress = if args.quiet {
+            ProgressBar::hidden()
+        } else {
+            let pb = ProgressBar::new(operations.len() as u64);
+            pb.set_style(
+                ProgressStyle::with_template(
+                    "Uploading to `{msg}` [{bar:40.cyan/blue}] {pos}/{len} files",
+                )
+                .unwrap()
+                .progress_chars("=> "),
             );
-        }
+            pb.set_message(branch.clone());
+            pb
+        };
+
+        let progress_callback: Option<CommitProgressCallback> = if args.quiet {
+            None
+        } else {
+            let pb = progress.clone();
+            Some(Arc::new(move |_path: &str| {
+                pb.inc(1);
+            }))
+        };
 
         for (batch_index, chunk) in operations.chunks(BUILD_COMMIT_BATCH_SIZE).enumerate() {
             let commit_message = if batch_count > 1 {
@@ -219,14 +235,13 @@ pub fn run_upload(args: UploadArgs) -> Result<()> {
                 revision: Some(branch.clone()),
                 create_pr: None,
                 parent_commit: None,
+                progress_callback: progress_callback.clone(),
             };
             api.create_commit(&params)
                 .wrap_err_with(|| format!("Cannot create commit on branch `{branch}`"))?;
-
-            if batch_count > 1 {
-                eprintln!("  Uploaded batch {}/{batch_count}.", batch_index + 1);
-            }
         }
+
+        progress.finish_with_message(format!("Uploaded to `{branch}`"));
     }
 
     let total_ops: usize = operations_by_branch.values().map(|v| v.len()).sum();