ci: upload benchmark input files

Author: AvivYossef-starkware

Cargo.lock

@@ -383,6 +383,7 @@ url = "2.5.0"
 validator = "0.12"
 void = "1.0.2"
 waker-fn = "1.2.0"
+walkdir = "2.5"
 workspace_tests.path = "workspace_tests"
 zstd = "0.13.1"
 

Cargo.toml

@@ -10,6 +10,7 @@ workspace = true
 clap = { workspace = true, features = ["derive"] }
 criterion.workspace = true
 serde = { workspace = true, features = ["derive"] }
+tokio = { workspace = true, features = ["full"] }
 
 [dev-dependencies]
 glob.workspace = true

crates/bench_tools/Cargo.toml

@@ -0,0 +1,37 @@
+use std::path::Path;
+
+/// Default GCS bucket for benchmarks.
+pub const BENCHMARKS_BUCKET: &str = "benchmarks_ci";
+
+/// Uploads all files from a local directory to Google Cloud Storage.
+///
+/// Uses gcloud CLI to upload files. Before running, authenticate with:
+/// `gcloud auth application-default login`
+///
+/// Files are uploaded to: `gs://{BENCHMARKS_BUCKET}/{benchmark_name}/input/`
+pub async fn upload_inputs(benchmark_name: &str, input_dir: &Path) {
+    println!(
+        "Uploading inputs from {} to gs://{}/{}/input/",
+        input_dir.display(),
+        BENCHMARKS_BUCKET,
+        benchmark_name
+    );
+
+    let source = format!("{}/*", input_dir.display());
+    let dest = format!("gs://{}/{}/input/", BENCHMARKS_BUCKET, benchmark_name);
+
+    // Use gcloud storage cp command to upload files.
+    let output = tokio::process::Command::new("gcloud")
+        .args(["storage", "cp", "-r", &source, &dest])
+        .output()
+        .await
+        .expect("Failed to upload inputs to GCS");
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        panic!("Failed to upload inputs to GCS: {}", stderr);
+    }
+
+    println!("{}", String::from_utf8_lossy(&output.stdout).trim());
+    println!("Input files uploaded successfully!");
+}

crates/bench_tools/src/gcs.rs

@@ -1,3 +1,4 @@
 #[cfg(test)]
 pub(crate) mod benches;
+pub mod gcs;
 pub mod types;

crates/bench_tools/src/lib.rs

@@ -1,4 +1,11 @@
-use bench_tools::types::benchmark_config::{find_benchmarks_by_package, BENCHMARKS};
+use std::path::PathBuf;
+
+use bench_tools::gcs;
+use bench_tools::types::benchmark_config::{
+    find_benchmark_by_name,
+    find_benchmarks_by_package,
+    BENCHMARKS,
+};
 use clap::{Parser, Subcommand};
 
 #[derive(Parser)]
@@ -25,9 +32,19 @@ enum Commands {
         #[arg(short, long)]
         package: Option<String>,
     },
+    /// Upload benchmark input files to GCS.
+    UploadInputs {
+        /// Benchmark name.
+        #[arg(long)]
+        benchmark: String,
+        /// Local directory containing input files.
+        #[arg(long)]
+        input_dir: String,
+    },
 }
 
-fn main() {
+#[tokio::main]
+async fn main() {
     let cli = Cli::parse();
     match cli.command {
         Commands::Run { package: _, out: _ } => {
@@ -59,5 +76,16 @@ fn main() {
                 }
             }
         },
+        Commands::UploadInputs { benchmark, input_dir } => {
+            // Validate benchmark exists
+            if find_benchmark_by_name(&benchmark).is_none() {
+                panic!("Unknown benchmark: {}", benchmark);
+            }
+
+            let input_path = PathBuf::from(&input_dir);
+            gcs::upload_inputs(&benchmark, &input_path).await;
+
+            println!("Input files uploaded successfully!");
+        }
     }
 }