feat: add openvm metrics export to integration test

Makefile

@@ -56,6 +56,9 @@ test-execute-chunk-multi:
 test-cycle:
 	@cargo test --release -p scroll-zkvm-integration --test chunk_circuit test_cycle -- --exact --nocapture
 
+test-chunk-cell:
+	@cargo test --release -p scroll-zkvm-integration --test chunk_circuit test_cell -- --exact --nocapture
+
 test-execute-batch: $(TESTDATA_PATH)/proofs/chunk-$(CHUNK_PROOF).json
 	@cargo test --release -p scroll-zkvm-integration --test batch_circuit test_e2e_execute -- --exact --nocapture
 

crates/integration/Cargo.toml

@@ -11,11 +11,14 @@ sbv-primitives = { workspace = true }
 tracing.workspace = true
 rkyv.workspace = true
 tracing-subscriber.workspace = true
+metrics.workspace = true
+metrics-util.workspace = true
 metrics-tracing-context.workspace = true
 
 openvm-build = { workspace = true, default-features = false }
 openvm-circuit.workspace = true
 openvm-sdk = { workspace = true, default-features = false }
+openvm-stark-sdk = { workspace = true, default-features = false }
 openvm-native-circuit = { workspace = true, default-features = false }
 openvm-native-compiler = { workspace = true, default-features = false }
 openvm-native-recursion = { workspace = true, default-features = false }
@@ -33,7 +36,6 @@ vm-zstd = { workspace = true, features = ["zstd"] }
 chrono = "0.4"
 ff = "0.13"
 glob = "0.3"
-once_cell = "1.20"
 revm = { workspace = true }
 serde_json = "1.0"
 sha2 = "0.10"
@@ -47,6 +49,7 @@ regex = "1.11.1"
 [dev-dependencies]
 halo2curves-axiom = "0.7.0"
 glob = "0.3"
+thousands = "0.2"
 
 [features]
 default = []

crates/integration/src/lib.rs

@@ -1,9 +1,17 @@
+#![feature(exit_status_error)]
+
 use cargo_metadata::MetadataCommand;
-use once_cell::sync::OnceCell;
+use eyre::Context;
+use metrics_tracing_context::TracingContextLayer;
+use metrics_util::{
+    debugging::{DebuggingRecorder, Snapshotter},
+    layers::Layer,
+};
 use openvm_sdk::{
     F, Sdk,
     config::{AppConfig, SdkVmConfig},
 };
+use openvm_stark_sdk::bench::serialize_metric_snapshot;
 use scroll_zkvm_prover::{
     ProverType, WrappedProof,
     setup::{read_app_config, read_app_exe},
@@ -12,7 +20,8 @@ use scroll_zkvm_prover::{
 use std::{
     path::{Path, PathBuf},
     process,
-    sync::LazyLock,
+    process::Command,
+    sync::{LazyLock, OnceLock},
 };
 use tracing::instrument;
 use tracing_subscriber::{fmt::format::FmtSpan, layer::SubscriberExt, util::SubscriberInitExt};
@@ -40,9 +49,20 @@ static DIR_OUTPUT: LazyLock<&Path> = LazyLock::new(|| {
     Box::leak(path.into_boxed_path())
 });
 
+pub static METRIC_SNAPSHOTTER: LazyLock<Snapshotter> = LazyLock::new(|| {
+    let recorder = DebuggingRecorder::new();
+    let snapshotter = recorder.snapshotter();
+    let recorder = TracingContextLayer::all().layer(recorder);
+    metrics::set_global_recorder(recorder).unwrap();
+    snapshotter
+});
+
 /// Directory to store proofs on disc.
 const DIR_PROOFS: &str = "proofs";
 
+/// Directory to store metrics on disc.
+const DIR_METRICS: &str = "metrics";
+
 /// File descriptor for app openvm config.
 const FD_APP_CONFIG: &str = "openvm.toml";
 
@@ -58,7 +78,7 @@ const ENV_OUTPUT_DIR: &str = "OUTPUT_DIR";
 /// - <DIR_OUTPUT>/chunk-tests-{timestamp}
 /// - <DIR_OUTPUT>/batch-tests-{timestamp}
 /// - <DIR_OUTPUT>/bundle-tests-{timestamp}
-static DIR_TESTRUN: OnceCell<PathBuf> = OnceCell::new();
+static DIR_TESTRUN: OnceLock<PathBuf> = OnceLock::new();
 
 /// Circuit that implements functionality required to run e2e tests.
 pub trait ProverTester {
@@ -102,6 +122,10 @@ pub trait ProverTester {
         Ok(())
     }
 
+    fn setup_metrics() {
+        LazyLock::force(&METRIC_SNAPSHOTTER);
+    }
+
     /// Load the app config.
     fn load_with_exe_fd(
         app_exe_fd: &str,
@@ -150,6 +174,36 @@ pub trait ProverTester {
     ) -> eyre::Result<Vec<F>> {
         Self::execute(app_config, &Self::gen_proving_task()?, exe_path)
     }
+
+    fn export_metrics() -> eyre::Result<()> {
+        Self::export_metrics_with_name("default")
+    }
+
+    fn export_metrics_with_name(name: &str) -> eyre::Result<()> {
+        let snapshot = METRIC_SNAPSHOTTER.snapshot();
+
+        let dir = DIR_TESTRUN
+            .get()
+            .ok_or(eyre::eyre!("missing assets dir"))?
+            .join(Self::DIR_ASSETS)
+            .join(DIR_METRICS);
+        std::fs::create_dir_all(&dir)?;
+
+        let path = dir.join(format!("metrics-{name}")).with_extension("json");
+        tracing::info!("exporting metrics to {}", path.display());
+        let f = std::fs::File::create(&path).context("failed to create metrics file")?;
+
+        serde_json::to_writer_pretty(f, &serialize_metric_snapshot(snapshot))
+            .context("failed to serialize metrics snapshot")?;
+
+        Command::new("openvm-prof")
+            .arg("--json-paths")
+            .arg(&path)
+            .status()
+            .context("failed to run openvm-prof")?
+            .exit_ok()
+            .context("openvm-prof failed")
+    }
 }
 
 /// The outcome of a successful prove-verify run.
@@ -203,6 +257,8 @@ fn setup_logger() -> eyre::Result<()> {
             .try_init()?;
     }
 
+    LazyLock::force(&METRIC_SNAPSHOTTER);
+
     Ok(())
 }
 

crates/integration/tests/chunk_circuit.rs

@@ -1,6 +1,7 @@
 use eyre::Ok;
+use metrics_util::{MetricKind, debugging::DebugValue};
 use scroll_zkvm_integration::{
-    ProverTester, prove_verify_multi, prove_verify_single,
+    METRIC_SNAPSHOTTER, ProverTester, prove_verify_multi, prove_verify_single,
     testers::chunk::{ChunkProverTester, MultiChunkProverTester, read_block_witness_from_testdata},
     utils::testing_hardfork,
 };
@@ -10,6 +11,7 @@ use scroll_zkvm_prover::{
     task::{ProvingTask, chunk::ChunkProvingTask},
     utils::{self, vm::ExecutionResult},
 };
+use thousands::Separable;
 
 fn exec_chunk(task: &ChunkProvingTask) -> eyre::Result<(ExecutionResult, u64)> {
     let (_path_app_config, app_config, path_exe) =
@@ -61,6 +63,64 @@ fn test_cycle() -> eyre::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_cell() -> eyre::Result<()> {
+    let task = ChunkProverTester::gen_proving_task()?;
+    let (exec_result, total_gas_used) = exec_chunk(&task)?;
+
+    ChunkProverTester::setup()?;
+    prove_verify_single::<ChunkProverTester>(None)?;
+
+    let snapshot = METRIC_SNAPSHOTTER.snapshot().into_hashmap();
+    let (total_cells, total_cycles) = snapshot
+        .iter()
+        .filter(|(ck, _)| {
+            matches!(ck.kind(), MetricKind::Counter)
+                && (ck.key().name() == "total_cells" || ck.key().name() == "total_cycles")
+                && ck
+                    .key()
+                    .labels()
+                    .find(|label| label.key() == "segment")
+                    .is_some() // filter only segment labels
+        })
+        .fold(
+            (0u64, 0u64),
+            |(cells_acc, cycles_acc), (ck, (_, _, value))| {
+                let DebugValue::Counter(value) = value else {
+                    panic!("Expected a counter value for total_cells or total_cycles");
+                };
+                if ck.key().name() == "total_cells" {
+                    (cells_acc + value, cycles_acc)
+                } else if ck.key().name() == "total_cycles" {
+                    (cells_acc, cycles_acc + value)
+                } else {
+                    unreachable!()
+                }
+            },
+        );
+
+    println!("Total cells: {}", total_cells.separate_with_commas());
+    println!("Total cycles: {}", total_cycles.separate_with_commas());
+    let cycles_per_gas = total_cycles as f64 / total_gas_used as f64;
+    let cells_per_gas = total_cells as f64 / total_gas_used as f64;
+    let cells_per_cycle = total_cells as f64 / total_cycles as f64;
+    println!(
+        "Cycles per gas: {}",
+        format!("{cycles_per_gas:.2}").separate_with_commas()
+    );
+    println!(
+        "Cells per gas: {}",
+        format!("{cells_per_gas:.2}").separate_with_commas()
+    );
+    println!(
+        "Cells per cycle: {}",
+        format!("{cells_per_cycle:.2}").separate_with_commas()
+    );
+    assert_eq!(exec_result.total_cycle, total_cycles);
+
+    Ok(())
+}
+
 #[test]
 fn test_execute() -> eyre::Result<()> {
     ChunkProverTester::setup()?;
@@ -69,7 +129,9 @@ fn test_execute() -> eyre::Result<()> {
     let (exec_result, total_gas_used) = exec_chunk(&task)?;
     let cycle_per_gas = exec_result.total_cycle / total_gas_used;
     assert_ne!(cycle_per_gas, 0);
-    assert!(cycle_per_gas <= 35);
+    // assert!(cycle_per_gas <= 35);
+    ChunkProverTester::export_metrics()?;
+
     Ok(())
 }
 
@@ -200,6 +262,7 @@ fn setup_prove_verify_single() -> eyre::Result<()> {
 
     prove_verify_single::<ChunkProverTester>(None)?;
 
+    ChunkProverTester::export_metrics()?;
     Ok(())
 }
 

crates/prover/Cargo.toml

@@ -33,7 +33,6 @@ base64 = "0.22"
 git-version = "0.3.5"
 hex = "0.4"
 munge = "=0.4.1"
-once_cell = "1.20"
 serde = "1.0"
 serde_json = "1.0"
 serde_stacker = "0.1"