feat(prof): add aggregate metric output in Bencher JSON format (#1315)

crates/prof/src/aggregate.rs

@@ -1,9 +1,10 @@
 use std::{collections::HashMap, io::Write};
 
 use eyre::Result;
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
-use crate::types::{Labels, MdTableCell, MetricDb};
+use crate::types::{BencherValue, Labels, MdTableCell, MetricDb};
 
 type MetricName = String;
 type MetricsByName = HashMap<MetricName, Vec<(f64, Labels)>>;
@@ -20,6 +21,21 @@ pub struct AggregateMetrics {
     /// "group" label => metric aggregate statitics
     #[serde(flatten)]
     pub by_group: HashMap<String, HashMap<MetricName, Stats>>,
+    /// In seconds
+    pub total_proof_time: MdTableCell,
+    /// In seconds
+    pub total_par_proof_time: MdTableCell,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct BencherAggregateMetrics {
+    // BMF max depth is 2
+    #[serde(flatten)]
+    pub by_group: HashMap<String, BencherValue>,
+    /// In seconds
+    pub total_proof_time: BencherValue,
+    /// In seconds
+    pub total_par_proof_time: BencherValue,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -28,6 +44,7 @@ pub struct Stats {
     pub max: MdTableCell,
     pub min: MdTableCell,
     pub avg: MdTableCell,
+    #[serde(skip)]
     pub count: usize,
 }
 
@@ -121,7 +138,12 @@ impl GroupedMetrics {
                 (group_name.clone(), group_summaries)
             })
             .collect();
-        AggregateMetrics { by_group }
+        let mut metrics = AggregateMetrics {
+            by_group,
+            ..Default::default()
+        };
+        metrics.compute_total();
+        metrics
     }
 }
 
@@ -140,6 +162,38 @@ pub(crate) fn group_weight(name: &str) -> usize {
 }
 
 impl AggregateMetrics {
+    pub fn compute_total(&mut self) {
+        let mut total_proof_time = MdTableCell::new(0.0, Some(0.0));
+        let mut total_par_proof_time = MdTableCell::new(0.0, Some(0.0));
+        for (group_name, metrics) in &self.by_group {
+            let stats = metrics.get(PROOF_TIME_LABEL);
+            if stats.is_none() {
+                continue;
+            }
+            let stats = stats.unwrap();
+            let mut sum = stats.sum;
+            let mut max = stats.max;
+            // convert ms to s
+            sum.val /= 1000.0;
+            max.val /= 1000.0;
+            if let Some(diff) = &mut sum.diff {
+                *diff /= 1000.0;
+            }
+            if let Some(diff) = &mut max.diff {
+                *diff /= 1000.0;
+            }
+            if !group_name.contains("keygen") {
+                // Proving time in keygen group is dummy and not part of total.
+                total_proof_time.val += sum.val;
+                *total_proof_time.diff.as_mut().unwrap() += sum.diff.unwrap_or(0.0);
+                total_par_proof_time.val += max.val;
+                *total_par_proof_time.diff.as_mut().unwrap() += max.diff.unwrap_or(0.0);
+            }
+        }
+        self.total_proof_time = total_proof_time;
+        self.total_par_proof_time = total_par_proof_time;
+    }
+
     pub fn set_diff(&mut self, prev: &Self) {
         for (group_name, metrics) in self.by_group.iter_mut() {
             if let Some(prev_metrics) = prev.by_group.get(group_name) {
@@ -150,6 +204,7 @@ impl AggregateMetrics {
                 }
             }
         }
+        self.compute_total();
     }
 
     pub fn to_vec(&self) -> Vec<(String, HashMap<MetricName, Stats>)> {
@@ -173,6 +228,41 @@ impl AggregateMetrics {
             .collect()
     }
 
+    pub fn to_bencher_metrics(&self) -> BencherAggregateMetrics {
+        let by_group = self
+            .by_group
+            .iter()
+            .flat_map(|(group_name, metrics)| {
+                metrics
+                    .iter()
+                    .flat_map(|(metric_name, stats)| {
+                        [
+                            (
+                                format!("{group_name}::{metric_name}::sum"),
+                                stats.sum.into(),
+                            ),
+                            (
+                                format!("{group_name}::{metric_name}"),
+                                BencherValue {
+                                    value: stats.avg.val,
+                                    lower_value: Some(stats.min.val),
+                                    upper_value: Some(stats.max.val),
+                                },
+                            ),
+                        ]
+                    })
+                    .collect_vec()
+            })
+            .collect();
+        let total_proof_time = self.total_proof_time.into();
+        let total_par_proof_time = self.total_par_proof_time.into();
+        BencherAggregateMetrics {
+            by_group,
+            total_proof_time,
+            total_par_proof_time,
+        }
+    }
+
     pub fn write_markdown(&self, writer: &mut impl Write, metric_names: &[&str]) -> Result<()> {
         self.write_summary_markdown(writer)?;
         writeln!(writer)?;
@@ -204,15 +294,13 @@ impl AggregateMetrics {
         Ok(())
     }
 
-    pub fn write_summary_markdown(&self, writer: &mut impl Write) -> Result<()> {
+    fn write_summary_markdown(&self, writer: &mut impl Write) -> Result<()> {
         writeln!(
             writer,
             "| Summary | Proof Time (s) | Parallel Proof Time (s) |"
         )?;
         writeln!(writer, "|:---|---:|---:|")?;
         let mut rows = Vec::new();
-        let mut total_proof_time = MdTableCell::new(0.0, Some(0.0));
-        let mut total_par_proof_time = MdTableCell::new(0.0, Some(0.0));
         for (group_name, summaries) in self.to_vec() {
             let stats = summaries.get(PROOF_TIME_LABEL);
             if stats.is_none() {
@@ -230,25 +318,28 @@ impl AggregateMetrics {
             if let Some(diff) = &mut max.diff {
                 *diff /= 1000.0;
             }
-            if !group_name.contains("keygen") {
-                // Proving time in keygen group is dummy and not part of total.
-                total_proof_time.val += sum.val;
-                *total_proof_time.diff.as_mut().unwrap() += sum.diff.unwrap_or(0.0);
-                total_par_proof_time.val += max.val;
-                *total_par_proof_time.diff.as_mut().unwrap() += max.diff.unwrap_or(0.0);
-            }
             rows.push((group_name, sum, max));
         }
         writeln!(
             writer,
-            "| Total | {total_proof_time} | {total_par_proof_time} |"
+            "| Total | {} | {} |",
+            self.total_proof_time, self.total_par_proof_time
         )?;
         for (group_name, proof_time, par_proof_time) in rows {
             writeln!(writer, "| {group_name} | {proof_time} | {par_proof_time} |")?;
         }
         writeln!(writer)?;
         Ok(())
     }
+
+    pub fn name(&self) -> String {
+        // A hacky way to determine the app name
+        self.by_group
+            .keys()
+            .find(|k| group_weight(k) == 0)
+            .unwrap_or_else(|| self.by_group.keys().next().unwrap())
+            .clone()
+    }
 }
 
 pub const PROOF_TIME_LABEL: &str = "total_proof_time_ms";

crates/prof/src/main.rs

@@ -10,7 +10,7 @@ use itertools::Itertools;
 use openvm_prof::{
     aggregate::{GroupedMetrics, VM_METRIC_NAMES},
     summary::GithubSummary,
-    types::MetricDb,
+    types::{BenchmarkOutput, MetricDb},
 };
 
 #[derive(Parser, Debug)]
@@ -28,6 +28,10 @@ struct Cli {
     /// Some file paths may be passed in that do not exist to account for new benchmarks.
     #[arg(long, value_delimiter = ',')]
     prev_json_paths: Option<Vec<PathBuf>>,
+
+    /// Path to write the output JSON in BMF format
+    #[arg(long)]
+    output_json: Option<PathBuf>,
 }
 
 #[derive(Subcommand, Debug)]
@@ -53,6 +57,7 @@ fn main() -> Result<()> {
     };
     let mut aggregated_metrics = Vec::new();
     let mut md_paths = Vec::new();
+    let mut output = BenchmarkOutput::default();
     for (metrics_path, prev_metrics_path) in args.json_paths.into_iter().zip_eq(prev_json_paths) {
         let db = MetricDb::new(&metrics_path)?;
         let grouped = GroupedMetrics::new(&db, "group")?;
@@ -67,8 +72,12 @@ fn main() -> Result<()> {
             }
         }
 
+        output
+            .by_name
+            .insert(aggregated.name(), aggregated.to_bencher_metrics());
         let mut writer = Vec::new();
         aggregated.write_markdown(&mut writer, VM_METRIC_NAMES)?;
+
         let mut markdown_output = String::from_utf8(writer)?;
 
         // TODO: calculate diffs for detailed metrics
@@ -82,6 +91,9 @@ fn main() -> Result<()> {
         md_paths.push(md_path);
         aggregated_metrics.push((aggregated, prev_aggregated));
     }
+    if let Some(path) = args.output_json {
+        fs::write(&path, serde_json::to_string_pretty(&output)?)?;
+    }
     if let Some(command) = args.command {
         match command {
             Commands::Summary(cmd) => {

crates/prof/src/summary.rs

@@ -4,7 +4,7 @@ use eyre::Result;
 use itertools::Itertools;
 
 use crate::{
-    aggregate::{group_weight, AggregateMetrics, CELLS_USED_LABEL, CYCLES_LABEL, PROOF_TIME_LABEL},
+    aggregate::{AggregateMetrics, CELLS_USED_LABEL, CYCLES_LABEL, PROOF_TIME_LABEL},
     types::MdTableCell,
 };
 
@@ -166,9 +166,8 @@ impl AggregateMetrics {
 
     /// Returns `None` if no group for app is found.
     pub fn get_summary_row(&self, md_filename: &str) -> Option<SummaryRow> {
-        // A hacky way to determine the app name
-        let app_name = self.by_group.keys().find(|k| group_weight(k) == 0)?;
-        let app = self.get_single_summary(app_name)?;
+        let app_name = self.name();
+        let app = self.get_single_summary(&app_name)?;
         let leaf = self.get_single_summary("leaf");
         let mut internals = Vec::new();
         let mut hgt = 0;

crates/prof/src/types.rs

@@ -3,6 +3,8 @@ use std::collections::{BTreeMap, HashMap};
 use num_format::{Locale, ToFormattedString};
 use serde::{Deserialize, Serialize};
 
+use crate::aggregate::BencherAggregateMetrics;
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Metric {
     pub name: String,
@@ -21,9 +23,26 @@ pub struct Labels(pub Vec<(String, String)>);
 #[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
 pub struct MdTableCell {
     pub val: f64,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub diff: Option<f64>,
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct BencherValue {
+    pub value: f64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub lower_value: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub upper_value: Option<f64>,
+}
+
+/// Benchmark output in [Bencher Metric Format](https://bencher.dev/docs/reference/bencher-metric-format/).
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
+pub struct BenchmarkOutput {
+    #[serde(flatten)]
+    pub by_name: HashMap<String, BencherAggregateMetrics>,
+}
+
 impl Labels {
     pub fn get(&self, key: &str) -> Option<&str> {
         self.0
@@ -117,6 +136,22 @@ fn format_cell(div: &str, span: Option<&str>, span_color: Option<&str>) -> Strin
     ret
 }
 
+impl BencherValue {
+    pub fn new(value: f64) -> Self {
+        Self {
+            value,
+            lower_value: None,
+            upper_value: None,
+        }
+    }
+}
+
+impl From<MdTableCell> for BencherValue {
+    fn from(cell: MdTableCell) -> Self {
+        Self::new(cell.val)
+    }
+}
+
 // For serialization purposes
 #[derive(Debug, Serialize, Deserialize)]
 pub struct MetricsFile {