ci: download bench input from gcs

Author: AvivYossef-starkware

Cargo.lock

@@ -16,6 +16,7 @@ apollo_infra_utils.workspace = true
 glob.workspace = true
 rstest.workspace = true
 serde_json.workspace = true
+tempfile.workspace = true
 
 [[bench]]
 harness = false

crates/bench_tools/Cargo.toml

@@ -7,7 +7,7 @@ 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`
+/// `gcloud auth login`
 ///
 /// Files are uploaded to: `gs://{BENCHMARKS_BUCKET}/{benchmark_name}/input/`
 pub fn upload_inputs(benchmark_name: &str, input_dir: &Path) {
@@ -35,3 +35,35 @@ pub fn upload_inputs(benchmark_name: &str, input_dir: &Path) {
     println!("{}", String::from_utf8_lossy(&output.stdout).trim());
     println!("Input files uploaded successfully!");
 }
+
+/// Downloads all input files for a benchmark from Google Cloud Storage.
+///
+/// Uses gcloud CLI to download files. Before running, authenticate with:
+/// `gcloud auth login`
+///
+/// Downloads from: `gs://{BENCHMARKS_BUCKET}/{benchmark_name}/input/` to the local input directory.
+pub fn download_inputs(benchmark_name: &str, local_input_dir: &Path) {
+    println!(
+        "Downloading inputs from gs://{}/{}/input/ to {}",
+        BENCHMARKS_BUCKET,
+        benchmark_name,
+        local_input_dir.display()
+    );
+
+    let source = format!("gs://{}/{}/input/*", BENCHMARKS_BUCKET, benchmark_name);
+    let dest = local_input_dir.display().to_string();
+
+    // Use gcloud storage cp command to download files.
+    let output = Command::new("gcloud")
+        .args(["storage", "cp", "-r", &source, &dest])
+        .output()
+        .expect("Failed to cp inputs from GCS");
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        panic!("Failed to download inputs from GCS: {}", stderr);
+    }
+
+    println!("{}", String::from_utf8_lossy(&output.stdout).trim());
+    println!("Input files downloaded successfully!");
+}

crates/bench_tools/src/gcs.rs

@@ -0,0 +1,48 @@
+use std::fs;
+
+use tempfile::TempDir;
+
+use crate::gcs::{download_inputs, upload_inputs};
+use crate::test_utils;
+
+#[test]
+#[ignore] // Run with: cargo test -p bench_tools -- --ignored
+fn test_upload_and_download_inputs() {
+    let benchmark_name = "dummy_benchmark";
+
+    // Get paths relative to bench_tools crate directory.
+    let source_dir = test_utils::bench_tools_crate_dir().join("data/dummy_bench_input");
+
+    // Ensure source files exist.
+    assert!(source_dir.exists(), "Source directory does not exist: {}", source_dir.display());
+
+    // Upload inputs.
+    println!("Testing upload...");
+    upload_inputs(benchmark_name, &source_dir);
+
+    // Create temp directory for download.
+    let temp_dir = TempDir::new().unwrap();
+    let download_dir = temp_dir.path();
+
+    println!("\nDownload directory: {}", download_dir.display());
+
+    // Download inputs to temp directory.
+    println!("\nTesting download...");
+    download_inputs(benchmark_name, download_dir);
+
+    // Verify files were downloaded.
+    let small_input = download_dir.join("small_input.json");
+    let large_input = download_dir.join("large_input.json");
+
+    assert!(small_input.exists(), "small_input.json was not downloaded");
+    assert!(large_input.exists(), "large_input.json was not downloaded");
+
+    // Verify content matches original.
+    let original_small = fs::read_to_string(source_dir.join("small_input.json")).unwrap();
+    let downloaded_small = fs::read_to_string(&small_input).unwrap();
+    assert_eq!(original_small, downloaded_small, "small_input.json content does not match");
+
+    let original_large = fs::read_to_string(source_dir.join("large_input.json")).unwrap();
+    let downloaded_large = fs::read_to_string(&large_input).unwrap();
+    assert_eq!(original_large, downloaded_large, "large_input.json content does not match");
+}

crates/bench_tools/src/gcs_test.rs

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

crates/bench_tools/src/lib.rs

@@ -0,0 +1,11 @@
+use std::path::PathBuf;
+
+use rstest::fixture;
+
+/// Returns the bench_tools crate directory.
+#[fixture]
+pub fn bench_tools_crate_dir() -> PathBuf {
+    std::env::var("CARGO_MANIFEST_DIR")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| std::env::current_dir().unwrap())
+}

crates/bench_tools/src/test_utils.rs

@@ -5,16 +5,9 @@ use std::process::Command;
 use apollo_infra_utils::path::project_path;
 use rstest::{fixture, rstest};
 
+use crate::test_utils::bench_tools_crate_dir;
 use crate::types::estimates::Estimates;
 
-/// Returns the bench_tools crate directory.
-#[fixture]
-fn bench_tools_crate_dir() -> PathBuf {
-    std::env::var("CARGO_MANIFEST_DIR")
-        .map(PathBuf::from)
-        .unwrap_or_else(|_| std::env::current_dir().unwrap())
-}
-
 /// Returns the directory where dummy benchmark estimate results are stored.
 #[fixture]
 fn dummy_bench_results_dir(bench_tools_crate_dir: PathBuf) -> PathBuf {