diff --git a/Cargo.toml b/Cargo.toml index 3bd623763..7880cd0cc 100644 --- a/Cargo.toml +++ b/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] diff --git a/cli-openvm/Cargo.toml b/cli-openvm/Cargo.toml index fc5e6127d..a82d899ff 100644 --- a/cli-openvm/Cargo.toml +++ b/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 diff --git a/cli-openvm/src/main.rs b/cli-openvm/src/main.rs index c8a42cc41..29a1fd5c9 100644 --- a/cli-openvm/src/main.rs +++ b/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, + + #[arg(long)] + metrics: Option, }, } 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(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 +}