diff --git a/kernel-builder/Cargo.lock b/kernel-builder/Cargo.lock index b7926be7..9c593c7a 100644 --- a/kernel-builder/Cargo.lock +++ b/kernel-builder/Cargo.lock @@ -436,6 +436,19 @@ dependencies = [ "memchr", ] +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + [[package]] name = "const-str" version = "1.1.0" @@ -726,6 +739,12 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "encoding_rs" version = "0.8.35" @@ -1150,6 +1169,7 @@ dependencies = [ "eyre", "git2", "huggingface-hub", + "indicatif", "itertools 0.13.0", "kernels-data", "minijinja", @@ -1235,7 +1255,7 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "huggingface-hub" version = "0.1.0" -source = "git+https://github.com/huggingface/huggingface_hub_rust.git?rev=8cbc662035e04d4be8e829316272893e980f5926#8cbc662035e04d4be8e829316272893e980f5926" +source = "git+https://github.com/huggingface/huggingface_hub_rust.git?branch=support-commit-operation-callback#b73ef621e847fff79b4ac73d88d07530192e3679" dependencies = [ "async-trait", "base64", @@ -1478,6 +1498,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -2010,6 +2043,12 @@ dependencies = [ "libm", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "objc2-core-foundation" version = "0.3.2" @@ -2239,6 +2278,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + [[package]] name = "potential_utf" version = "0.1.4" @@ -4226,6 +4271,15 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + [[package]] name = "windows-sys" version = "0.60.2" diff --git a/kernel-builder/Cargo.toml b/kernel-builder/Cargo.toml index 4c172807..5c8ac065 100644 --- a/kernel-builder/Cargo.toml +++ b/kernel-builder/Cargo.toml @@ -15,7 +15,8 @@ clap = { version = "4", features = ["derive"] } 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"] } +huggingface-hub = { git = "https://github.com/huggingface/huggingface_hub_rust.git", branch = "support-commit-operation-callback", package = "huggingface-hub", features = ["blocking", "xet"] } +indicatif = "0.17" itertools = "0.13" minijinja = "2.5" minijinja-embed = "2.5" diff --git a/kernel-builder/src/main.rs b/kernel-builder/src/main.rs index 5d17132e..818e90a4 100644 --- a/kernel-builder/src/main.rs +++ b/kernel-builder/src/main.rs @@ -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 { diff --git a/kernel-builder/src/upload.rs b/kernel-builder/src/upload.rs index adc02c0b..0bc066f7 100644 --- a/kernel-builder/src/upload.rs +++ b/kernel-builder/src/upload.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 = 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(¶ms) .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(); diff --git a/nix-builder/pkgs/kernel-builder/default.nix b/nix-builder/pkgs/kernel-builder/default.nix index b0c08903..be2fcae0 100644 --- a/nix-builder/pkgs/kernel-builder/default.nix +++ b/nix-builder/pkgs/kernel-builder/default.nix @@ -54,7 +54,7 @@ rustPlatform.buildRustPackage { cargoLock = { lockFile = ../../../kernel-builder/Cargo.lock; outputHashes = { - "huggingface-hub-0.1.0" = "sha256-XgVrtujU7gPQ3XnUxeEVF9Kaf4+/EwLudKkwDPj44II="; + "huggingface-hub-0.1.0" = "sha256-xgKdI775XTb2gUQ+gBTPnhrCRKUnLiANk7hquXCYr9Q="; "hf-xet-1.4.0" = "sha256-/vvU8qy9U+suiH9MCcxrV3Ayw84yRV6EmW0yzB7Uvng="; }; };