openvm_prof/
lib.rs

Help
1use std::{collections::HashMap, fs::File, path::Path};
2
3use aggregate::{
4    EXECUTE_TIME_LABEL, PROOF_TIME_LABEL, PROVE_EXCL_TRACE_TIME_LABEL, TRACE_GEN_TIME_LABEL,
5};
6use eyre::Result;
7
8use crate::types::{Labels, Metric, MetricDb, MetricsFile};
9
10pub mod aggregate;
11pub mod summary;
12pub mod types;
13
14impl MetricDb {
15    pub fn new(metrics_file: impl AsRef<Path>) -> Result<Self> {
16        let file = File::open(metrics_file)?;
17        let metrics: MetricsFile = serde_json::from_reader(file)?;
18
19        let mut db = MetricDb::default();
20
21        // Process counters
22        for entry in metrics.counter {
23            if entry.value == 0.0 {
24                continue;
25            }
26            let labels = Labels::from(entry.labels);
27            db.add_to_flat_dict(labels, entry.metric, entry.value);
28        }
29
30        // Process gauges
31        for entry in metrics.gauge {
32            let labels = Labels::from(entry.labels);
33            db.add_to_flat_dict(labels, entry.metric, entry.value);
34        }
35
36        db.apply_aggregations();
37        db.separate_by_label_types();
38
39        Ok(db)
40    }
41
42    // Currently hardcoding aggregations
43    pub fn apply_aggregations(&mut self) {
44        for metrics in self.flat_dict.values_mut() {
45            let get = |key: &str| metrics.iter().find(|m| m.name == key).map(|m| m.value);
46            let execute_time = get(EXECUTE_TIME_LABEL);
47            let trace_gen_time = get(TRACE_GEN_TIME_LABEL);
48            let prove_excl_trace_time = get(PROVE_EXCL_TRACE_TIME_LABEL);
49            if let (Some(execute_time), Some(trace_gen_time), Some(prove_excl_trace_time)) =
50                (execute_time, trace_gen_time, prove_excl_trace_time)
51            {
52                let total_time = execute_time + trace_gen_time + prove_excl_trace_time;
53                metrics.push(Metric::new(PROOF_TIME_LABEL.to_string(), total_time));
54            }
55        }
56    }
57
58    pub fn add_to_flat_dict(&mut self, labels: Labels, metric: String, value: f64) {
59        self.flat_dict
60            .entry(labels)
61            .or_default()
62            .push(Metric::new(metric, value));
63    }
64
65    // Custom sorting function that ensures 'group' comes first.
66    // Other keys are sorted alphabetically.
67    pub fn custom_sort_label_keys(label_keys: &mut [String]) {
68        // Prioritize 'group' by giving it the lowest possible sort value
69        label_keys.sort_by_key(|key| {
70            if key == "group" {
71                (0, key.clone()) // Lowest priority for 'group'
72            } else {
73                (1, key.clone()) // Normal priority for other keys
74            }
75        });
76    }
77
78    pub fn separate_by_label_types(&mut self) {
79        self.dict_by_label_types.clear();
80
81        for (labels, metrics) in &self.flat_dict {
82            // Get sorted label keys
83            let mut label_keys: Vec<String> = labels.0.iter().map(|(key, _)| key.clone()).collect();
84            Self::custom_sort_label_keys(&mut label_keys);
85
86            // Create label_values based on sorted keys
87            let label_dict: HashMap<String, String> = labels.0.iter().cloned().collect();
88
89            let label_values: Vec<String> = label_keys
90                .iter()
91                .map(|key| label_dict.get(key).unwrap().clone())
92                .collect();
93
94            // Add to dict_by_label_types
95            self.dict_by_label_types
96                .entry(label_keys)
97                .or_default()
98                .entry(label_values)
99                .or_default()
100                .extend(metrics.clone());
101        }
102    }
103
104    pub fn generate_markdown_tables(&self) -> String {
105        let mut markdown_output = String::new();
106        // Get sorted keys to iterate in consistent order
107        let mut sorted_keys: Vec<_> = self.dict_by_label_types.keys().cloned().collect();
108        sorted_keys.sort();
109
110        for label_keys in sorted_keys {
111            if label_keys.contains(&"cycle_tracker_span".to_string()) {
112                // Skip cycle_tracker_span as it is too long for markdown and visualized in flamegraphs
113                continue;
114            }
115            let metrics_dict = &self.dict_by_label_types[&label_keys];
116            let mut metric_names: Vec<String> = metrics_dict
117                .values()
118                .flat_map(|metrics| metrics.iter().map(|m| m.name.clone()))
119                .collect::<std::collections::HashSet<_>>()
120                .into_iter()
121                .collect();
122            metric_names.sort_by(|a, b| b.cmp(a));
123
124            // Create table header
125            let header = format!(
126                "| {} | {} |",
127                label_keys.join(" | "),
128                metric_names.join(" | ")
129            );
130
131            let separator = "| ".to_string()
132                + &vec!["---"; label_keys.len() + metric_names.len()].join(" | ")
133                + " |";
134
135            markdown_output.push_str(&header);
136            markdown_output.push('\n');
137            markdown_output.push_str(&separator);
138            markdown_output.push('\n');
139
140            // Fill table rows
141            for (label_values, metrics) in metrics_dict {
142                let mut row = String::new();
143                row.push_str("| ");
144                row.push_str(&label_values.join(" | "));
145                row.push_str(" | ");
146
147                // Add metric values
148                for metric_name in &metric_names {
149                    let metric_value = metrics
150                        .iter()
151                        .find(|m| &m.name == metric_name)
152                        .map(|m| Self::format_number(m.value))
153                        .unwrap_or_default();
154
155                    row.push_str(&format!("{} | ", metric_value));
156                }
157
158                markdown_output.push_str(&row);
159                markdown_output.push('\n');
160            }
161
162            markdown_output.push('\n');
163        }
164
165        markdown_output
166    }
167}
openvm_prof/lib.rs

openvm_prof/
lib.rs
