add --metrics to prove command (#2925)

allows exporting metrics to a json file

Cargo.toml

@@ -121,10 +121,12 @@ openvm-stark-sdk = { git = "https://github.com/powdr-labs/stark-backend.git", re
   "parallel",
   "jemalloc",
   "nightly-features",
+  "bench-metrics",
 ] }
 openvm-stark-backend = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "37c22d2", default-features = false, features = [
   "parallel",
   "jemalloc",
+  "bench-metrics",
 ] }
 
 [profile.pr-tests]

cli-openvm/Cargo.toml

@@ -19,11 +19,18 @@ openvm-stark-backend.workspace = true
 powdr-openvm.workspace = true
 
 eyre = "0.6.12"
-tracing = "0.1.40"
 
 clap = { version = "^4.3", features = ["derive"] }
 
 serde_cbor = "0.11.2"
 
+tracing = "0.1.40"
+tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }
+tracing-forest = "0.1"
+metrics = "0.23.0"
+metrics-tracing-context = "0.16.0"
+metrics-util = "0.17.0"
+serde_json = "1.0.117"
+
 [lints]
 workspace = true

cli-openvm/src/main.rs

@@ -1,11 +1,15 @@
 use eyre::Result;
+use metrics_tracing_context::{MetricsLayer, TracingContextLayer};
+use metrics_util::{debugging::DebuggingRecorder, layers::Layer};
 use openvm_sdk::StdIn;
-use openvm_stark_sdk::config::setup_tracing_with_log_level;
+use openvm_stark_sdk::bench::serialize_metric_snapshot;
 use powdr_openvm::{CompiledProgram, GuestOptions, PgoConfig, PgoType, PowdrConfig};
 
 use clap::{CommandFactory, Parser, Subcommand};
-use std::io;
+use std::{io, path::PathBuf};
 use tracing::Level;
+use tracing_forest::ForestLayer;
+use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
 
 #[derive(Parser)]
 #[command(name = "powdr-openvm", author, version, about, long_about = None)]
@@ -77,13 +81,16 @@ enum Commands {
 
         #[arg(long)]
         input: Option<u32>,
+
+        #[arg(long)]
+        metrics: Option<PathBuf>,
     },
 }
 
 fn main() -> Result<(), io::Error> {
     let args = Cli::parse();
 
-    setup_tracing_with_log_level(Level::WARN);
+    setup_tracing_with_log_level(Level::INFO);
 
     if let Some(command) = args.command {
         run_command(command);
@@ -132,12 +139,22 @@ fn run_command(command: Commands) {
             recursion,
             pgo,
             input,
+            metrics,
         } => {
             let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
             let pgo_config = get_pgo_config(guest.clone(), guest_opts.clone(), pgo, input);
             let program =
                 powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap();
-            powdr_openvm::prove(&program, mock, recursion, stdin_from(input), None).unwrap();
+            let prove =
+                || powdr_openvm::prove(&program, mock, recursion, stdin_from(input), None).unwrap();
+            if let Some(metrics_path) = metrics {
+                run_with_metric_collection_to_file(
+                    std::fs::File::create(metrics_path).expect("Failed to create metrics file"),
+                    prove,
+                );
+            } else {
+                prove()
+            }
         }
 
         // Run Pgo on the original openvm program (without powdr extension)
@@ -186,3 +203,26 @@ fn get_pgo_config(
         PgoType::None => PgoConfig::None,
     }
 }
+
+fn setup_tracing_with_log_level(level: Level) {
+    let env_filter = EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| EnvFilter::new(format!("{level},p3_=warn")));
+    let subscriber = Registry::default()
+        .with(env_filter)
+        .with(ForestLayer::default())
+        .with(MetricsLayer::new());
+    tracing::subscriber::set_global_default(subscriber).unwrap();
+}
+
+/// export stark-backend metrics to the given file
+pub fn run_with_metric_collection_to_file<R>(file: std::fs::File, f: impl FnOnce() -> R) -> R {
+    let recorder = DebuggingRecorder::new();
+    let snapshotter = recorder.snapshotter();
+    let recorder = TracingContextLayer::all().layer(recorder);
+    metrics::set_global_recorder(recorder).unwrap();
+    let res = f();
+
+    serde_json::to_writer_pretty(&file, &serialize_metric_snapshot(snapshotter.snapshot()))
+        .unwrap();
+    res
+}