trustformers_debug/export/
perfetto.rs

1//! Perfetto/Chrome Trace Event Format export
2//!
3//! Exports profiling data in Perfetto JSON format compatible with
4//! `chrome://tracing` and the Perfetto UI at <https://ui.perfetto.dev>.
5//!
6//! # Example
7//!
8//! ```no_run
9//! use trustformers_debug::export::perfetto::{PerfettoTrace, PerfettoEvent, PerfettoPhase};
10//! use std::collections::HashMap;
11//!
12//! let mut trace = PerfettoTrace::new();
13//! let mut args = HashMap::new();
14//! args.insert("layer".to_string(), serde_json::json!("attention"));
15//! trace.add_event(PerfettoEvent {
16//!     name: "layer_forward".to_string(),
17//!     phase: PerfettoPhase::Complete,
18//!     timestamp_us: 12345,
19//!     duration_us: Some(1500),
20//!     pid: 1,
21//!     tid: 1,
22//!     args,
23//! });
24//! let json = trace.export_to_string().unwrap();
25//! println!("{}", json);
26//! ```
27
28use std::collections::HashMap;
29use std::io::Write;
30
31use anyhow::Result;
32use serde::{Deserialize, Serialize};
33use serde_json::Value;
34
35use crate::ProfilerReport;
36
37// ─────────────────────────────────────────────────────────────
38// Public types
39// ─────────────────────────────────────────────────────────────
40
41/// Phase codes used in the Perfetto/Chrome trace-event format.
42///
43/// Each variant corresponds to the `ph` field in a trace-event object.
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45pub enum PerfettoPhase {
46    /// `B` — duration begin.
47    Begin,
48    /// `E` — duration end.
49    End,
50    /// `X` — complete (begin + duration).
51    Complete,
52    /// `i` — instant event.
53    Instant,
54    /// `C` — counter.
55    Counter,
56}
57
58impl PerfettoPhase {
59    /// Returns the single-character phase code.
60    pub fn as_code(&self) -> &'static str {
61        match self {
62            Self::Begin => "B",
63            Self::End => "E",
64            Self::Complete => "X",
65            Self::Instant => "i",
66            Self::Counter => "C",
67        }
68    }
69}
70
71/// A single trace event in Perfetto/Chrome format.
72///
73/// # Example
74///
75/// ```
76/// use std::collections::HashMap;
77/// use trustformers_debug::export::perfetto::{PerfettoEvent, PerfettoPhase};
78///
79/// let event = PerfettoEvent {
80///     name: "forward".to_string(),
81///     phase: PerfettoPhase::Complete,
82///     timestamp_us: 0,
83///     duration_us: Some(500),
84///     pid: 1,
85///     tid: 1,
86///     args: HashMap::new(),
87/// };
88/// assert_eq!(event.phase.as_code(), "X");
89/// ```
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct PerfettoEvent {
92    /// Human-readable name shown in the trace viewer.
93    pub name: String,
94    /// Event phase (Begin, End, Complete, Instant, Counter).
95    pub phase: PerfettoPhase,
96    /// Timestamp in **microseconds** since the trace start.
97    pub timestamp_us: u64,
98    /// Duration in microseconds (required for `Complete` events).
99    pub duration_us: Option<u64>,
100    /// Process ID.
101    pub pid: u32,
102    /// Thread ID.
103    pub tid: u32,
104    /// Arbitrary key-value metadata.
105    pub args: HashMap<String, Value>,
106}
107
108/// An in-memory collection of [`PerfettoEvent`]s that can be serialised to
109/// the Chrome trace-event JSON format.
110///
111/// # Example
112///
113/// ```
114/// use trustformers_debug::export::perfetto::PerfettoTrace;
115///
116/// let trace = PerfettoTrace::new();
117/// assert_eq!(trace.len(), 0);
118/// ```
119#[derive(Debug, Default)]
120pub struct PerfettoTrace {
121    events: Vec<PerfettoEvent>,
122}
123
124impl PerfettoTrace {
125    /// Creates an empty trace.
126    pub fn new() -> Self {
127        Self::default()
128    }
129
130    /// Appends a single event.
131    pub fn add_event(&mut self, event: PerfettoEvent) {
132        self.events.push(event);
133    }
134
135    /// The events in the trace, in insertion order.
136    pub fn events(&self) -> &[PerfettoEvent] {
137        &self.events
138    }
139
140    /// Returns the number of events in the trace.
141    pub fn len(&self) -> usize {
142        self.events.len()
143    }
144
145    /// Returns `true` if the trace contains no events.
146    pub fn is_empty(&self) -> bool {
147        self.events.is_empty()
148    }
149
150    /// Serialises the trace to a JSON string in Perfetto format.
151    ///
152    /// # Errors
153    ///
154    /// Returns an error if JSON serialisation fails.
155    ///
156    /// # Example
157    ///
158    /// ```
159    /// use trustformers_debug::export::perfetto::PerfettoTrace;
160    ///
161    /// let trace = PerfettoTrace::new();
162    /// let json = trace.export_to_string().unwrap();
163    /// assert!(json.contains("traceEvents"));
164    /// ```
165    pub fn export_to_string(&self) -> Result<String> {
166        let doc = self.build_doc();
167        Ok(serde_json::to_string_pretty(&doc)?)
168    }
169
170    /// Writes the trace to a file at `path`.
171    ///
172    /// # Errors
173    ///
174    /// Returns an error if the file cannot be created or written to, or if
175    /// JSON serialisation fails.
176    ///
177    /// # Example
178    ///
179    /// ```no_run
180    /// use trustformers_debug::export::perfetto::PerfettoTrace;
181    ///
182    /// let trace = PerfettoTrace::new();
183    /// trace.export_to_file(std::path::Path::new("/tmp/trace.json")).unwrap();
184    /// ```
185    pub fn export_to_file(&self, path: &std::path::Path) -> Result<()> {
186        let json = self.export_to_string()?;
187        let mut file = std::fs::File::create(path)?;
188        file.write_all(json.as_bytes())?;
189        tracing::debug!("Perfetto trace written to {}", path.display());
190        Ok(())
191    }
192
193    // ── helpers ──────────────────────────────────────────────
194
195    fn build_doc(&self) -> Value {
196        let events: Vec<Value> = self.events.iter().map(event_to_value).collect();
197        serde_json::json!({
198            "traceEvents": events,
199            "displayTimeUnit": "ms",
200        })
201    }
202}
203
204// ─────────────────────────────────────────────────────────────
205// PerfettoExporter
206// ─────────────────────────────────────────────────────────────
207
208/// Converts a [`ProfilerReport`] to a [`PerfettoTrace`] and writes it to disk.
209///
210/// # Example
211///
212/// ```no_run
213/// use trustformers_debug::export::perfetto::PerfettoExporter;
214/// use trustformers_debug::ProfilerReport;
215/// use std::collections::HashMap;
216/// use std::time::Duration;
217///
218/// // Build a minimal report for demonstration.
219/// let report = ProfilerReport {
220///     total_events: 0,
221///     total_runtime: Duration::from_millis(0),
222///     statistics: HashMap::new(),
223///     bottlenecks: vec![],
224///     slowest_layers: vec![],
225///     memory_efficiency: Default::default(),
226///     recommendations: vec![],
227/// };
228/// PerfettoExporter::export_profiler_report(
229///     &report,
230///     std::path::Path::new("/tmp/report.json"),
231/// ).unwrap();
232/// ```
233pub struct PerfettoExporter;
234
235impl PerfettoExporter {
236    /// Converts a [`ProfilerReport`] into a Perfetto trace file.
237    ///
238    /// Each layer in [`ProfilerReport::slowest_layers`] becomes a `Complete`
239    /// trace event.  Bottlenecks are appended as `Instant` events.
240    ///
241    /// # Errors
242    ///
243    /// Returns an error if the file cannot be written.
244    pub fn export_profiler_report(report: &ProfilerReport, path: &std::path::Path) -> Result<()> {
245        let mut trace = PerfettoTrace::new();
246        let mut cursor_us: u64 = 0;
247
248        for (layer_name, duration) in &report.slowest_layers {
249            let dur_us = duration.as_micros() as u64;
250            let mut args = HashMap::new();
251            args.insert("layer_name".to_string(), Value::String(layer_name.clone()));
252            trace.add_event(PerfettoEvent {
253                name: layer_name.clone(),
254                phase: PerfettoPhase::Complete,
255                timestamp_us: cursor_us,
256                duration_us: Some(dur_us),
257                pid: 1,
258                tid: 1,
259                args,
260            });
261            cursor_us += dur_us;
262        }
263
264        for bottleneck in &report.bottlenecks {
265            let mut args = HashMap::new();
266            args.insert(
267                "description".to_string(),
268                Value::String(bottleneck.description.clone()),
269            );
270            args.insert(
271                "suggestion".to_string(),
272                Value::String(bottleneck.suggestion.clone()),
273            );
274            trace.add_event(PerfettoEvent {
275                name: format!("bottleneck:{}", bottleneck.location),
276                phase: PerfettoPhase::Instant,
277                timestamp_us: cursor_us,
278                duration_us: None,
279                pid: 1,
280                tid: 1,
281                args,
282            });
283        }
284
285        trace.export_to_file(path)
286    }
287}
288
289// ─────────────────────────────────────────────────────────────
290// Private helpers
291// ─────────────────────────────────────────────────────────────
292
293fn event_to_value(e: &PerfettoEvent) -> Value {
294    let mut obj = serde_json::json!({
295        "name": e.name,
296        "ph": e.phase.as_code(),
297        "ts": e.timestamp_us,
298        "pid": e.pid,
299        "tid": e.tid,
300    });
301
302    if let Some(dur) = e.duration_us {
303        obj["dur"] = Value::Number(dur.into());
304    }
305
306    if !e.args.is_empty() {
307        obj["args"] = Value::Object(e.args.iter().map(|(k, v)| (k.clone(), v.clone())).collect());
308    }
309
310    obj
311}
312
313// ─────────────────────────────────────────────────────────────
314// Tests
315// ─────────────────────────────────────────────────────────────
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320    use std::time::Duration;
321
322    fn make_event(name: &str, phase: PerfettoPhase, ts: u64, dur: Option<u64>) -> PerfettoEvent {
323        PerfettoEvent {
324            name: name.to_string(),
325            phase,
326            timestamp_us: ts,
327            duration_us: dur,
328            pid: 1,
329            tid: 1,
330            args: HashMap::new(),
331        }
332    }
333
334    #[test]
335    fn test_empty_trace_roundtrip() {
336        let trace = PerfettoTrace::new();
337        let json = trace.export_to_string().unwrap();
338        let parsed: Value = serde_json::from_str(&json).unwrap();
339        assert_eq!(parsed["traceEvents"].as_array().unwrap().len(), 0);
340        assert_eq!(parsed["displayTimeUnit"], "ms");
341    }
342
343    #[test]
344    fn test_add_complete_event() {
345        let mut trace = PerfettoTrace::new();
346        trace.add_event(make_event(
347            "forward",
348            PerfettoPhase::Complete,
349            1000,
350            Some(500),
351        ));
352        assert_eq!(trace.len(), 1);
353
354        let json = trace.export_to_string().unwrap();
355        let parsed: Value = serde_json::from_str(&json).unwrap();
356        let ev = &parsed["traceEvents"][0];
357        assert_eq!(ev["ph"], "X");
358        assert_eq!(ev["ts"], 1000_u64);
359        assert_eq!(ev["dur"], 500_u64);
360        assert_eq!(ev["name"], "forward");
361    }
362
363    #[test]
364    fn test_begin_end_phases() {
365        let mut trace = PerfettoTrace::new();
366        trace.add_event(make_event("op", PerfettoPhase::Begin, 0, None));
367        trace.add_event(make_event("op", PerfettoPhase::End, 200, None));
368        let json = trace.export_to_string().unwrap();
369        let parsed: Value = serde_json::from_str(&json).unwrap();
370        assert_eq!(parsed["traceEvents"][0]["ph"], "B");
371        assert_eq!(parsed["traceEvents"][1]["ph"], "E");
372    }
373
374    #[test]
375    fn test_export_to_file() {
376        let mut dir = std::env::temp_dir();
377        dir.push("perfetto_test_trace.json");
378
379        let mut trace = PerfettoTrace::new();
380        let mut args = HashMap::new();
381        args.insert("batch_size".to_string(), serde_json::json!(32));
382        trace.add_event(PerfettoEvent {
383            name: "attention_forward".to_string(),
384            phase: PerfettoPhase::Complete,
385            timestamp_us: 0,
386            duration_us: Some(1500),
387            pid: 1,
388            tid: 2,
389            args,
390        });
391        trace.export_to_file(&dir).unwrap();
392        assert!(dir.exists());
393
394        let content = std::fs::read_to_string(&dir).unwrap();
395        let parsed: Value = serde_json::from_str(&content).unwrap();
396        assert_eq!(parsed["traceEvents"].as_array().unwrap().len(), 1);
397
398        std::fs::remove_file(&dir).ok();
399    }
400
401    #[test]
402    fn test_exporter_from_profiler_report() {
403        use crate::profiler::{MemoryEfficiencyAnalysis, PerformanceBottleneck};
404
405        let mut dir = std::env::temp_dir();
406        dir.push("perfetto_profiler_report.json");
407
408        let report = ProfilerReport {
409            total_events: 2,
410            total_runtime: Duration::from_millis(100),
411            statistics: HashMap::new(),
412            bottlenecks: vec![PerformanceBottleneck {
413                bottleneck_type: crate::profiler::BottleneckType::CpuBound,
414                location: "attention".to_string(),
415                severity: crate::profiler::BottleneckSeverity::Medium,
416                description: "CPU saturated".to_string(),
417                suggestion: "Use flash attention".to_string(),
418                metrics: HashMap::new(),
419            }],
420            slowest_layers: vec![
421                ("attention".to_string(), Duration::from_millis(10)),
422                ("ffn".to_string(), Duration::from_millis(15)),
423            ],
424            memory_efficiency: MemoryEfficiencyAnalysis::default(),
425            recommendations: vec![],
426        };
427
428        PerfettoExporter::export_profiler_report(&report, &dir).unwrap();
429        assert!(dir.exists());
430
431        let content = std::fs::read_to_string(&dir).unwrap();
432        let parsed: Value = serde_json::from_str(&content).unwrap();
433        let events = parsed["traceEvents"].as_array().unwrap();
434        // 2 layer events + 1 bottleneck instant event
435        assert_eq!(events.len(), 3);
436        assert_eq!(events[2]["ph"], "i");
437
438        std::fs::remove_file(&dir).ok();
439    }
440
441    #[test]
442    fn test_instant_and_counter_phases() {
443        let mut trace = PerfettoTrace::new();
444        trace.add_event(make_event("checkpoint", PerfettoPhase::Instant, 500, None));
445        trace.add_event(make_event("loss", PerfettoPhase::Counter, 600, None));
446        let json = trace.export_to_string().unwrap();
447        let parsed: Value = serde_json::from_str(&json).unwrap();
448        assert_eq!(parsed["traceEvents"][0]["ph"], "i");
449        assert_eq!(parsed["traceEvents"][1]["ph"], "C");
450    }
451}
trustformers_debug/export/perfetto.rs

trustformers_debug/export/
perfetto.rs
