feat: upload progress bar

kernel-builder/Cargo.toml

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

kernel-builder/src/main.rs

@@ -125,6 +125,10 @@ enum Commands {
         /// Repository type on Hugging Face Hub (`kernel` by default, or `model` for legacy repos).
         #[arg(long, value_enum, default_value_t = RepoTypeArg::Kernel)]
         repo_type: RepoTypeArg,
+
+        /// Suppress progress output.
+        #[arg(long, short)]
+        quiet: bool,
     },
 
     /// Upload kernel build artifacts to the Hugging Face Hub.
@@ -322,6 +326,7 @@ fn main() -> Result<()> {
             branch,
             private,
             repo_type,
+            quiet,
         } => {
             run_build(
                 kernel_dir.clone(),
@@ -336,6 +341,7 @@ fn main() -> Result<()> {
                 branch,
                 private,
                 repo_type,
+                quiet,
             })
         }
         Commands::CreatePyproject {

kernel-builder/src/upload.rs

@@ -2,14 +2,17 @@ use std::{
     collections::{BTreeMap, HashSet},
     fs,
     path::{Path, PathBuf},
+    sync::Arc,
 };
 
 use clap::Args;
 use eyre::{bail, Context, Result};
 use huggingface_hub::{
-    AddSource, CommitOperation, CreateRepoParams, RepoCreateBranchParams, RepoCreateCommitParams,
-    RepoListFilesParams, RepoListRefsParams, RepoType,
+    AddSource, CommitOperation, CreateRepoParams, ProgressEvent, ProgressHandler,
+    RepoCreateBranchParams, RepoCreateCommitParams, RepoListFilesParams, RepoListRefsParams,
+    RepoType, UploadEvent,
 };
+use indicatif::{ProgressBar, ProgressStyle};
 use walkdir::WalkDir;
 
 use crate::{
@@ -18,6 +21,17 @@ use crate::{
     util::{check_or_infer_kernel_dir, discover_variants, parse_build},
 };
 
+/// Bridges `ProgressHandler` events to an `indicatif::ProgressBar`.
+struct IndicatifProgress(ProgressBar);
+
+impl ProgressHandler for IndicatifProgress {
+    fn on_progress(&self, event: &ProgressEvent) {
+        if let ProgressEvent::Upload(UploadEvent::FileComplete { files, .. }) = event {
+            self.0.inc(files.len() as u64);
+        }
+    }
+}
+
 const MAIN_BRANCH: &str = "main";
 const BUILD_COMMIT_BATCH_SIZE: usize = 1_000;
 
@@ -59,6 +73,10 @@ pub struct UploadArgs {
     /// Repository type on Hugging Face Hub (`kernel` by default, or `model` for legacy repos).
     #[arg(long, value_enum, default_value_t = RepoTypeArg::Kernel)]
     pub repo_type: RepoTypeArg,
+
+    /// Suppress progress output.
+    #[arg(long, short)]
+    pub quiet: bool,
 }
 
 /// Get repository and branch from the given arguments, or fallback to
@@ -189,20 +207,28 @@ 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_handler: Option<Arc<dyn ProgressHandler>> = if args.quiet {
+            None
+        } else {
+            let pb = progress.clone();
+            Some(Arc::new(IndicatifProgress(pb)))
+        };
 
         for (batch_index, chunk) in operations.chunks(BUILD_COMMIT_BATCH_SIZE).enumerate() {
             let commit_message = if batch_count > 1 {
@@ -221,14 +247,13 @@ pub fn run_upload(args: UploadArgs) -> Result<()> {
                 revision: Some(branch.clone()),
                 create_pr: None,
                 parent_commit: None,
+                progress: progress_handler.clone(),
             };
             repo.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();

nix-builder/pkgs/kernel-abi-check/default.nix

@@ -33,7 +33,7 @@ rustPlatform.buildRustPackage {
   cargoLock = {
     lockFile = ../../../Cargo.lock;
     outputHashes = {
-      "huggingface-hub-0.0.1" = "sha256-By8b1NUPWu+XF3Om1NcEO+o2qdZUco+FxvrJGNRqxWs=";
+      "huggingface-hub-0.0.1" = "sha256-0HjQUEAH/EiG8/1snEWYRteidR+zJq0yct164+6K+IY=";
     };
   };
 

nix-builder/pkgs/kernel-builder/default.nix

@@ -52,7 +52,7 @@ rustPlatform.buildRustPackage {
   cargoLock = {
     lockFile = ../../../Cargo.lock;
     outputHashes = {
-      "huggingface-hub-0.0.1" = "sha256-By8b1NUPWu+XF3Om1NcEO+o2qdZUco+FxvrJGNRqxWs=";
+      "huggingface-hub-0.0.1" = "sha256-0HjQUEAH/EiG8/1snEWYRteidR+zJq0yct164+6K+IY=";
     };
   };
 

nix-builder/pkgs/python-modules/kernel-abi-check/default.nix

@@ -37,7 +37,7 @@ buildPythonPackage {
   cargoDeps = rustPlatform.importCargoLock {
     lockFile = ../../../../Cargo.lock;
     outputHashes = {
-      "huggingface-hub-0.0.1" = "sha256-By8b1NUPWu+XF3Om1NcEO+o2qdZUco+FxvrJGNRqxWs=";
+      "huggingface-hub-0.0.1" = "sha256-0HjQUEAH/EiG8/1snEWYRteidR+zJq0yct164+6K+IY=";
     };
   };
 

nix-builder/pkgs/python-modules/kernels-data/default.nix

@@ -35,7 +35,7 @@ buildPythonPackage {
   cargoDeps = rustPlatform.importCargoLock {
     lockFile = ../../../../Cargo.lock;
     outputHashes = {
-      "huggingface-hub-0.0.1" = "sha256-By8b1NUPWu+XF3Om1NcEO+o2qdZUco+FxvrJGNRqxWs=";
+      "huggingface-hub-0.0.1" = "sha256-0HjQUEAH/EiG8/1snEWYRteidR+zJq0yct164+6K+IY=";
     };
   };