ci: run benchmark command

Author: AvivYossef-starkware

bench_tools/Cargo.toml

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

bench_tools/src/gcs.rs

@@ -58,7 +58,7 @@ pub async fn download_inputs(benchmark_name: &str, output_dir: &Path) {
         .args(["storage", "cp", "-r", &source, &dest])
         .output()
         .await
-        .expect("Failed to cp inputs from GCS");
+        .unwrap_or_else(|e| panic!("Failed to cp inputs from GCS: {}", e));
 
     if !output.status.success() {
         let stderr = String::from_utf8_lossy(&output.stderr);

bench_tools/src/lib.rs

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

bench_tools/src/main.rs

@@ -25,6 +25,10 @@ enum Commands {
         /// Output directory for results.
         #[arg(short, long)]
         out: String,
+        /// Optional: Local directory containing input files. If not provided, inputs will be
+        /// downloaded from GCS for benchmarks that require them.
+        #[arg(long)]
+        input_dir: Option<String>,
     },
     /// List benchmarks for a package.
     List {
@@ -47,8 +51,14 @@ enum Commands {
 async fn main() {
     let cli = Cli::parse();
     match cli.command {
-        Commands::Run { package: _, out: _ } => {
-            unimplemented!()
+        Commands::Run { package, out, input_dir } => {
+            let benchmarks = find_benchmarks_by_package(&package);
+
+            if benchmarks.is_empty() {
+                panic!("No benchmarks found for package: {}", package);
+            }
+
+            bench_tools::runner::run_benchmarks(&benchmarks, input_dir.as_deref(), &out).await;
         }
         Commands::List { package } => match package {
             Some(package_name) => {

bench_tools/src/runner.rs

@@ -0,0 +1,99 @@
+use std::path::PathBuf;
+
+use tokio::fs;
+
+use crate::gcs;
+use crate::types::benchmark_config::BenchmarkConfig;
+
+/// Prepares inputs for a benchmark.
+/// If the benchmark needs inputs and a local input directory is provided,
+/// it verifies the directory exists.
+/// If the benchmark needs inputs and no local input directory is provided,
+/// it downloads the inputs from GCS.
+async fn prepare_inputs(bench: &BenchmarkConfig, input_dir: Option<&str>) {
+    if !bench.needs_inputs() {
+        return;
+    }
+
+    if let Some(local_dir) = input_dir {
+        let local_path = PathBuf::from(local_dir);
+        if !local_path.exists() {
+            panic!("Input directory does not exist: {}", local_dir);
+        }
+    } else {
+        let target = PathBuf::from(bench.input_dir.unwrap());
+        gcs::download_inputs(bench.name, &target).await;
+        if !target.exists() {
+            panic!("Failed to prepare inputs for {}", bench.name);
+        }
+    }
+}
+
+/// Runs a single benchmark and panic if it fails.
+async fn run_single_benchmark(bench: &BenchmarkConfig) {
+    println!("Running: {}", bench.name);
+
+    let output = tokio::process::Command::new("cargo")
+        .args(bench.cmd_args)
+        .output()
+        .await
+        .unwrap_or_else(|e| panic!("Failed to execute {}: {}", bench.name, e));
+
+    if !output.status.success() {
+        panic!("\nBenchmark {} failed:\n{}", bench.name, String::from_utf8_lossy(&output.stderr));
+    }
+}
+
+/// Collects benchmark results from criterion output and saves them to the output directory.
+async fn save_benchmark_results(_bench: &BenchmarkConfig, output_dir: &str) {
+    let criterion_base = PathBuf::from("target/criterion");
+    let Ok(mut entries) = fs::read_dir(&criterion_base).await else { return };
+
+    // Collect all estimates files.
+    while let Ok(Some(entry)) = entries.next_entry().await {
+        let path = entry.path();
+        if !path.is_dir() {
+            continue;
+        }
+
+        // Save estimates file.
+        let estimates_path = path.join("new/estimates.json");
+        if let Ok(data) = fs::read_to_string(&estimates_path).await {
+            if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) {
+                if let Ok(pretty) = serde_json::to_string_pretty(&json) {
+                    let bench_name = path.file_name().unwrap().to_string_lossy();
+                    let dest =
+                        PathBuf::from(output_dir).join(format!("{}_estimates.json", bench_name));
+                    if fs::write(&dest, pretty).await.is_ok() {
+                        println!("Saved results: {}", dest.display());
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// Runs benchmarks for a given package, handling input downloads if needed.
+pub async fn run_benchmarks(
+    benchmarks: &[&BenchmarkConfig],
+    input_dir: Option<&str>,
+    output_dir: &str,
+) {
+    // Prepare inputs
+    for bench in benchmarks {
+        prepare_inputs(bench, input_dir).await;
+    }
+
+    // Create output directory
+    tokio::fs::create_dir_all(output_dir)
+        .await
+        .unwrap_or_else(|e| panic!("Failed to create output dir: {}", e));
+
+    // Run benchmarks
+    for bench in benchmarks {
+        run_single_benchmark(bench).await;
+        save_benchmark_results(bench, output_dir).await;
+    }
+
+    println!("\n✓ All benchmarks completed! Results saved to: {}", output_dir);
+}