ci: run benchmark command

Author: AvivYossef-starkware

Cargo.lock

@@ -253,6 +253,7 @@ ethnum = "1.5.0"
 expect-test = "1.5.1"
 flate2 = "1.0.24"
 fs2 = "0.4"
+fs_extra = "1.3.0"
 futures = "0.3.21"
 futures-channel = "0.3.21"
 futures-util = "0.3.21"

Cargo.toml

@@ -9,7 +9,9 @@ workspace = true
 [dependencies]
 clap = { workspace = true, features = ["derive"] }
 criterion.workspace = true
+fs_extra.workspace = true
 serde = { workspace = true, features = ["derive"] }
+serde_json.workspace = true
 
 [dev-dependencies]
 apollo_infra_utils.workspace = true

crates/bench_tools/Cargo.toml

@@ -57,7 +57,7 @@ pub fn download_inputs(benchmark_name: &str, local_input_dir: &Path) {
     let output = Command::new("gcloud")
         .args(["storage", "cp", "-r", &source, &dest])
         .output()
-        .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);

crates/bench_tools/src/gcs.rs

@@ -3,6 +3,7 @@ pub(crate) mod benches;
 pub mod gcs;
 #[cfg(test)]
 pub mod gcs_test;
+pub mod runner;
 #[cfg(test)]
 pub mod test_utils;
 pub mod types;

crates/bench_tools/src/lib.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 {
@@ -46,8 +50,14 @@ enum Commands {
 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);
         }
         Commands::List { package } => match package {
             Some(package_name) => {

crates/bench_tools/src/main.rs

@@ -0,0 +1,114 @@
+use std::fs;
+use std::path::PathBuf;
+
+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 copies the contents from the local directory to the expected input location.
+/// If the benchmark needs inputs and no local input directory is provided,
+/// it downloads the inputs from GCS.
+fn prepare_inputs(bench: &BenchmarkConfig, input_dir: Option<&str>) {
+    if !bench.needs_inputs() {
+        return;
+    }
+
+    let benchmark_input_dir = PathBuf::from(bench.input_dir.unwrap());
+
+    // Create the input directory if it doesn't exist.
+    fs::create_dir_all(&benchmark_input_dir).unwrap_or_else(|e| {
+        panic!("Failed to create directory {}: {}", benchmark_input_dir.display(), e)
+    });
+
+    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);
+        }
+
+        // Copy local directory contents to the benchmark input directory.
+        let options = fs_extra::dir::CopyOptions::new().content_only(true);
+        fs_extra::dir::copy(&local_path, &benchmark_input_dir, &options).unwrap_or_else(|e| {
+            panic!(
+                "Failed to copy inputs from {} to {}: {}",
+                local_dir,
+                benchmark_input_dir.display(),
+                e
+            )
+        });
+
+        println!("Copied inputs from {} to {}", local_dir, benchmark_input_dir.display());
+    } else {
+        gcs::download_inputs(bench.name, &benchmark_input_dir);
+        if !benchmark_input_dir.exists() {
+            panic!(
+                "Failed to download inputs for {}: {}",
+                bench.name,
+                benchmark_input_dir.display()
+            );
+        }
+    }
+}
+
+/// Runs a single benchmark and panic if it fails.
+fn run_single_benchmark(bench: &BenchmarkConfig) {
+    println!("Running: {}", bench.name);
+
+    let output = std::process::Command::new("cargo")
+        .args(bench.cmd_args)
+        .output()
+        .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.
+fn save_benchmark_results(_bench: &BenchmarkConfig, output_dir: &str) {
+    let criterion_base = PathBuf::from("target/criterion");
+    let Ok(entries) = fs::read_dir(&criterion_base) else { return };
+
+    // Collect all estimates files.
+    for entry in entries.flatten() {
+        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) {
+            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).is_ok() {
+                        println!("Saved results: {}", dest.display());
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// Runs benchmarks for a given package, handling input downloads if needed.
+pub fn run_benchmarks(benchmarks: &[&BenchmarkConfig], input_dir: Option<&str>, output_dir: &str) {
+    // Prepare inputs.
+    for bench in benchmarks {
+        prepare_inputs(bench, input_dir);
+    }
+
+    // Create output directory.
+    fs::create_dir_all(output_dir).unwrap_or_else(|e| panic!("Failed to create output dir: {}", e));
+
+    // Run benchmarks.
+    for bench in benchmarks {
+        run_single_benchmark(bench);
+        save_benchmark_results(bench, output_dir);
+    }
+
+    println!("\n✓ All benchmarks completed! Results saved to: {}", output_dir);
+}