ci: upload benchmark input files

Author: AvivYossef-starkware

crates/bench_tools/src/gcs.rs

@@ -0,0 +1,37 @@
+use std::path::Path;
+use std::process::Command;
+
+/// Default GCS bucket for benchmarks.
+pub const BENCHMARKS_BUCKET: &str = "apollo_benchmarks";
+
+/// 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 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 = Command::new("gcloud")
+        .args(["storage", "cp", "-r", &source, &dest])
+        .output()
+        .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/lib.rs

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

crates/bench_tools/src/main.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,6 +32,15 @@ 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() {
@@ -59,5 +75,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);
+
+            println!("Input files uploaded successfully!");
+        }
     }
 }