chore(metrics): instrument total_proof_time_ms directly (#1952)

`total_proof_time_ms` can now be measured directly with a span instead
of by summing up individual parts, which is more accurate.

The increase in times from benchmarks includes some additional setup
time.

Author: jonathanpwang

crates/prof/src/lib.rs

@@ -46,6 +46,12 @@ impl MetricDb {
     pub fn apply_aggregations(&mut self) {
         for metrics in self.flat_dict.values_mut() {
             let get = |key: &str| metrics.iter().find(|m| m.name == key).map(|m| m.value);
+            let total_proof_time = get(PROOF_TIME_LABEL);
+            if total_proof_time.is_some() {
+                // We have instrumented total_proof_time_ms
+                continue;
+            }
+            // otherwise, calculate it from sub-components
             let execute_metered_time = get(EXECUTE_METERED_TIME_LABEL);
             let execute_preflight_time = get(EXECUTE_PREFLIGHT_TIME_LABEL);
             let trace_gen_time = get(TRACE_GEN_TIME_LABEL);

crates/vm/src/arch/vm.rs

@@ -870,7 +870,9 @@ where
         let mut proofs = Vec::with_capacity(segments.len());
         let mut state = Some(vm.create_initial_state(exe, input));
         for (seg_idx, segment) in segments.into_iter().enumerate() {
-            let _span = info_span!("prove_segment", segment = seg_idx).entered();
+            let _segment_span = info_span!("prove_segment", segment = seg_idx).entered();
+            // We need a separate span so the metric label includes "segment" from _segment_span
+            let _prove_span = info_span!("total_proof").entered();
             let Segment {
                 instret_start,
                 num_insns,
@@ -914,6 +916,7 @@ where
     <VB::VmConfig as VmExecutionConfig<Val<E::SC>>>::Executor:
         PreflightExecutor<Val<E::SC>, VB::RecordArena>,
 {
+    #[instrument(name = "total_proof", skip_all)]
     fn prove(
         &mut self,
         input: impl Into<Streams<Val<E::SC>>>,

docs/crates/metrics.md

@@ -16,9 +16,8 @@ For a segment proof, the following metrics are collected:
 - `trace_gen_time_ms` (gauge): The time to generate non-cached trace matrices from execution records.
   - If this is a segment in a VM with continuations enabled, a `segment: segment_idx` label is added to the metric.
 - All metrics collected by [`openvm-stark-backend`](https://github.com/openvm-org/stark-backend/blob/main/docs/metrics.md), in particular `stark_prove_excluding_trace_time_ms` (gauge).
-- The `total_proof_time_ms` of the proof is:
-  - The sum `execute_preflight_time_ms + trace_gen_time_ms + stark_prove_excluding_trace_time_ms` for app proofs. The `execute_metered_time_ms` is excluded for app proofs because it is not run on a per-segment basis.
-  - The sum `execute_metered_time_ms + execute_preflight_time_ms + trace_gen_time_ms + stark_prove_excluding_trace_time_ms` for non-app proofs.
+- The `total_proof_time_ms` of the proof is instrumented directly when possible. Otherwise, it is calculated as:
+  - The sum `execute_preflight_time_ms + trace_gen_time_ms + stark_prove_excluding_trace_time_ms`. The `execute_metered_time_ms` is excluded for app proofs because it is not run on a per-segment basis.
 - `insns` (counter): The total number of instructions executed in the segment.
 - `main_cells_used` (counter): The total number of main trace cells used by all chips in the segment. This does not include cells needed to pad rows to power-of-two matrix heights. Only main trace cells, not preprocessed or permutation trace cells, are counted.
 - `total_cells_used` (counter): The total number of preprocessed, main, and permutation trace cells used by all chips in the segment. This does not include cells needed to pad rows to power-of-two matrix heights.