trustformers_debug/export/
unified.rs

1//! Unified export interface for all supported trace formats.
2//!
3//! [`TraceExporter`] dispatches to the appropriate backend based on
4//! [`ExportConfig::format`].  New formats can be added by extending
5//! [`ExportFormat`] and [`TraceExporter::export_all`].
6
7use std::path::Path;
8
9use anyhow::Result;
10
11use super::perfetto::{PerfettoEvent, PerfettoExporter, PerfettoPhase, PerfettoTrace};
12use super::tracy::{TracyExporter, TracyTrace, TracyZone};
13
14// ─────────────────────────────────────────────────────────────────────────────
15// TimingEvent
16// ─────────────────────────────────────────────────────────────────────────────
17
18/// A generic timed event used by the CSV and JSON exporters.
19///
20/// Can be converted from a [`PerfettoEvent`] or a [`TracyZone`].
21///
22/// # Example
23///
24/// ```
25/// use trustformers_debug::export::unified::TimingEvent;
26///
27/// let ev = TimingEvent {
28///     timestamp_ns: 1_000_000,
29///     duration_ns: 500_000,
30///     thread_id: 0,
31///     name: "attention_forward".to_string(),
32/// };
33/// assert_eq!(ev.timestamp_ns, 1_000_000);
34/// ```
35#[derive(Debug, Clone, PartialEq)]
36pub struct TimingEvent {
37    /// Start timestamp in nanoseconds.
38    pub timestamp_ns: u64,
39    /// Duration in nanoseconds.
40    pub duration_ns: u64,
41    /// Thread / worker identifier.
42    pub thread_id: u32,
43    /// Human-readable operation name.
44    pub name: String,
45}
46
47impl From<&TracyZone> for TimingEvent {
48    fn from(z: &TracyZone) -> Self {
49        Self {
50            timestamp_ns: z.timestamp_ns,
51            duration_ns: z.duration_ns,
52            thread_id: z.thread_id,
53            name: z.name.clone(),
54        }
55    }
56}
57
58impl From<&PerfettoEvent> for TimingEvent {
59    fn from(e: &PerfettoEvent) -> Self {
60        Self {
61            timestamp_ns: e.timestamp_us * 1_000,
62            duration_ns: e.duration_us.unwrap_or(0) * 1_000,
63            thread_id: e.tid,
64            name: e.name.clone(),
65        }
66    }
67}
68
69// ─────────────────────────────────────────────────────────────────────────────
70// ExportFormat
71// ─────────────────────────────────────────────────────────────────────────────
72
73/// Supported export formats for profiling trace data.
74///
75/// # Example
76///
77/// ```
78/// use trustformers_debug::export::unified::ExportFormat;
79///
80/// let fmt = ExportFormat::Csv;
81/// assert_eq!(fmt.extension(), "csv");
82/// ```
83#[derive(Debug, Clone, PartialEq, Eq)]
84pub enum ExportFormat {
85    /// Perfetto/Chrome trace event JSON.
86    Perfetto,
87    /// Tracy offline CSV format.
88    Tracy,
89    /// `chrome://tracing` compatible JSON (alias for Perfetto).
90    ChromeTrace,
91    /// Generic timing CSV (timestamp_ns, duration_ns, thread_id, name).
92    Csv,
93    /// Simple JSON array of timing events.
94    Json,
95}
96
97impl ExportFormat {
98    /// Returns the conventional file extension for this format.
99    pub fn extension(&self) -> &str {
100        match self {
101            Self::Perfetto | Self::ChromeTrace => "json",
102            Self::Tracy => "csv",
103            Self::Csv => "csv",
104            Self::Json => "json",
105        }
106    }
107}
108
109// ─────────────────────────────────────────────────────────────────────────────
110// ExportConfig
111// ─────────────────────────────────────────────────────────────────────────────
112
113/// Configuration for a single export operation.
114///
115/// # Example
116///
117/// ```
118/// use trustformers_debug::export::unified::{ExportConfig, ExportFormat};
119///
120/// let cfg = ExportConfig {
121///     format: ExportFormat::Csv,
122///     output_path: "/tmp/trace.csv".to_string(),
123///     compress: false,
124/// };
125/// assert_eq!(cfg.format, ExportFormat::Csv);
126/// ```
127#[derive(Debug, Clone)]
128pub struct ExportConfig {
129    /// Target export format.
130    pub format: ExportFormat,
131    /// Filesystem path where the output will be written.
132    pub output_path: String,
133    /// Reserved for future compression support (currently ignored).
134    pub compress: bool,
135}
136
137// ─────────────────────────────────────────────────────────────────────────────
138// ExportError
139// ─────────────────────────────────────────────────────────────────────────────
140
141/// Errors returned by [`TraceExporter`].
142#[derive(Debug, Clone, PartialEq)]
143pub enum ExportError {
144    /// The requested format is not supported in this build.
145    UnsupportedFormat(String),
146    /// An I/O error occurred while writing the output file.
147    IoError(String),
148    /// The input trace/events slice was empty.
149    EmptyTrace,
150}
151
152impl std::fmt::Display for ExportError {
153    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154        match self {
155            Self::UnsupportedFormat(s) => write!(f, "unsupported export format: {s}"),
156            Self::IoError(s) => write!(f, "I/O error during export: {s}"),
157            Self::EmptyTrace => write!(f, "export trace is empty"),
158        }
159    }
160}
161
162impl std::error::Error for ExportError {}
163
164impl From<anyhow::Error> for ExportError {
165    fn from(e: anyhow::Error) -> Self {
166        Self::IoError(e.to_string())
167    }
168}
169
170// ─────────────────────────────────────────────────────────────────────────────
171// ProfilingTrace wrapper
172// ─────────────────────────────────────────────────────────────────────────────
173
174/// A unified view of a profiling trace that can be exported to multiple formats.
175///
176/// Build one from a [`PerfettoTrace`] or a [`TracyTrace`], then call
177/// [`TraceExporter::export_all`] to materialise it in the configured format.
178#[derive(Debug, Default)]
179pub struct ProfilingTrace {
180    events: Vec<TimingEvent>,
181}
182
183impl ProfilingTrace {
184    /// Creates an empty trace.
185    pub fn new() -> Self {
186        Self::default()
187    }
188
189    /// Appends a [`TimingEvent`].
190    pub fn add_event(&mut self, event: TimingEvent) {
191        self.events.push(event);
192    }
193
194    /// Returns a slice of all events.
195    pub fn events(&self) -> &[TimingEvent] {
196        &self.events
197    }
198
199    /// Returns the total number of events.
200    pub fn len(&self) -> usize {
201        self.events.len()
202    }
203
204    /// Returns `true` when the trace contains no events.
205    pub fn is_empty(&self) -> bool {
206        self.events.is_empty()
207    }
208
209    /// Builds a [`PerfettoTrace`] from this unified trace.
210    pub fn to_perfetto(&self) -> PerfettoTrace {
211        let mut trace = PerfettoTrace::new();
212        for ev in &self.events {
213            trace.add_event(PerfettoEvent {
214                name: ev.name.clone(),
215                phase: PerfettoPhase::Complete,
216                timestamp_us: ev.timestamp_ns / 1_000,
217                duration_us: Some(ev.duration_ns / 1_000),
218                pid: 1,
219                tid: ev.thread_id,
220                args: std::collections::HashMap::new(),
221            });
222        }
223        trace
224    }
225
226    /// Builds a [`TracyTrace`] from this unified trace.
227    pub fn to_tracy(&self) -> TracyTrace {
228        let mut trace = TracyTrace::new();
229        for ev in &self.events {
230            trace.add_zone(TracyZone {
231                name: ev.name.clone(),
232                timestamp_ns: ev.timestamp_ns,
233                duration_ns: ev.duration_ns,
234                thread_id: ev.thread_id,
235            });
236        }
237        trace
238    }
239}
240
241impl From<&TracyTrace> for ProfilingTrace {
242    fn from(t: &TracyTrace) -> Self {
243        let events = t.zones().iter().map(TimingEvent::from).collect();
244        Self { events }
245    }
246}
247
248impl From<&PerfettoTrace> for ProfilingTrace {
249    /// Convert every event in the Perfetto trace, in order, via
250    /// [`From<&PerfettoEvent> for TimingEvent`].
251    ///
252    /// This used to ignore `t` entirely and return an empty trace, so any
253    /// caller writing `let profiling: ProfilingTrace = (&perfetto).into();`
254    /// silently lost every event.
255    fn from(t: &PerfettoTrace) -> Self {
256        let mut trace = Self::new();
257        for event in t.events() {
258            trace.add_event(TimingEvent::from(event));
259        }
260        trace
261    }
262}
263
264// ─────────────────────────────────────────────────────────────────────────────
265// CsvExporter
266// ─────────────────────────────────────────────────────────────────────────────
267
268/// Exports a slice of [`TimingEvent`]s to a CSV string.
269///
270/// # Example
271///
272/// ```
273/// use trustformers_debug::export::unified::{CsvExporter, TimingEvent};
274///
275/// let events = vec![TimingEvent { timestamp_ns: 0, duration_ns: 1000, thread_id: 0, name: "op".to_string() }];
276/// let csv = CsvExporter::export_to_csv(&events);
277/// assert!(csv.contains("timestamp_ns,duration_ns,thread_id,name"));
278/// assert!(csv.contains("0,1000,0,op"));
279/// ```
280pub struct CsvExporter;
281
282impl CsvExporter {
283    /// Serialises `events` to a CSV string with header:
284    /// `timestamp_ns,duration_ns,thread_id,name`
285    pub fn export_to_csv(events: &[TimingEvent]) -> String {
286        let mut out = String::from("timestamp_ns,duration_ns,thread_id,name\n");
287        for ev in events {
288            let safe_name = ev.name.replace(',', "\\,");
289            use std::fmt::Write as _;
290            let _ = writeln!(
291                out,
292                "{},{},{},{}",
293                ev.timestamp_ns, ev.duration_ns, ev.thread_id, safe_name
294            );
295        }
296        out
297    }
298
299    /// Writes CSV to the file at `path`.
300    pub fn export_to_file(events: &[TimingEvent], path: &Path) -> Result<()> {
301        let csv = Self::export_to_csv(events);
302        std::fs::write(path, csv.as_bytes())?;
303        Ok(())
304    }
305}
306
307// ─────────────────────────────────────────────────────────────────────────────
308// JsonExporter
309// ─────────────────────────────────────────────────────────────────────────────
310
311/// Exports a slice of [`TimingEvent`]s to a JSON array string.
312///
313/// # Example
314///
315/// ```
316/// use trustformers_debug::export::unified::{JsonExporter, TimingEvent};
317///
318/// let events = vec![
319///     TimingEvent { timestamp_ns: 1000, duration_ns: 500, thread_id: 1, name: "ffn".to_string() },
320/// ];
321/// let json = JsonExporter::export_to_json(&events);
322/// assert!(json.starts_with('['));
323/// assert!(json.contains("\"name\":\"ffn\""));
324/// ```
325pub struct JsonExporter;
326
327impl JsonExporter {
328    /// Serialises `events` to a JSON array without external dependencies.
329    pub fn export_to_json(events: &[TimingEvent]) -> String {
330        use std::fmt::Write as _;
331        let mut out = String::from('[');
332        for (i, ev) in events.iter().enumerate() {
333            if i > 0 {
334                out.push(',');
335            }
336            let escaped_name = escape_json_string_local(&ev.name);
337            let _ = write!(
338                out,
339                r#"{{"timestamp_ns":{},"duration_ns":{},"thread_id":{},"name":"{}"}}"#,
340                ev.timestamp_ns, ev.duration_ns, ev.thread_id, escaped_name
341            );
342        }
343        out.push(']');
344        out
345    }
346
347    /// Writes JSON to the file at `path`.
348    pub fn export_to_file(events: &[TimingEvent], path: &Path) -> Result<()> {
349        let json = Self::export_to_json(events);
350        std::fs::write(path, json.as_bytes())?;
351        Ok(())
352    }
353}
354
355fn escape_json_string_local(s: &str) -> String {
356    use std::fmt::Write as _;
357    let mut out = String::with_capacity(s.len());
358    for c in s.chars() {
359        match c {
360            '"' => out.push_str("\\\""),
361            '\\' => out.push_str("\\\\"),
362            '\n' => out.push_str("\\n"),
363            '\r' => out.push_str("\\r"),
364            '\t' => out.push_str("\\t"),
365            c if (c as u32) < 0x20 => {
366                let _ = write!(out, "\\u{:04x}", c as u32);
367            },
368            c => out.push(c),
369        }
370    }
371    out
372}
373
374// ─────────────────────────────────────────────────────────────────────────────
375// TraceExporter
376// ─────────────────────────────────────────────────────────────────────────────
377
378/// Unified exporter that dispatches to the appropriate backend.
379///
380/// # Example
381///
382/// ```no_run
383/// use trustformers_debug::export::unified::{
384///     ExportConfig, ExportFormat, ProfilingTrace, TimingEvent, TraceExporter,
385/// };
386///
387/// let mut trace = ProfilingTrace::new();
388/// trace.add_event(TimingEvent {
389///     timestamp_ns: 0,
390///     duration_ns: 1_000_000,
391///     thread_id: 0,
392///     name: "attention".to_string(),
393/// });
394/// let config = ExportConfig {
395///     format: ExportFormat::Csv,
396///     output_path: "/tmp/trace.csv".to_string(),
397///     compress: false,
398/// };
399/// TraceExporter::export_all(&trace, &config).unwrap();
400/// ```
401pub struct TraceExporter;
402
403impl TraceExporter {
404    /// Exports `trace` according to `config`.
405    ///
406    /// # Errors
407    ///
408    /// Returns [`ExportError::EmptyTrace`] when the trace is empty.
409    /// Returns [`ExportError::IoError`] on filesystem errors.
410    pub fn export_all(trace: &ProfilingTrace, config: &ExportConfig) -> Result<(), ExportError> {
411        if trace.is_empty() {
412            return Err(ExportError::EmptyTrace);
413        }
414        let path = Path::new(&config.output_path);
415
416        match &config.format {
417            ExportFormat::Perfetto | ExportFormat::ChromeTrace => {
418                let perf = trace.to_perfetto();
419                perf.export_to_file(path).map_err(ExportError::from)?;
420            },
421            ExportFormat::Tracy => {
422                let tracy = trace.to_tracy();
423                tracy.export_to_file(path).map_err(ExportError::from)?;
424            },
425            ExportFormat::Csv => {
426                CsvExporter::export_to_file(trace.events(), path).map_err(ExportError::from)?;
427            },
428            ExportFormat::Json => {
429                JsonExporter::export_to_file(trace.events(), path).map_err(ExportError::from)?;
430            },
431        }
432        Ok(())
433    }
434
435    /// Exports from a [`crate::ProfilerReport`] via [`PerfettoExporter`] or
436    /// [`TracyExporter`] depending on the format.
437    ///
438    /// For CSV/JSON formats the slowest layers are converted to
439    /// [`TimingEvent`]s first.
440    pub fn export_profiler_report(
441        report: &crate::ProfilerReport,
442        config: &ExportConfig,
443    ) -> Result<(), ExportError> {
444        let path = Path::new(&config.output_path);
445        match &config.format {
446            ExportFormat::Perfetto | ExportFormat::ChromeTrace => {
447                PerfettoExporter::export_profiler_report(report, path)
448                    .map_err(ExportError::from)?;
449            },
450            ExportFormat::Tracy => {
451                TracyExporter::export_profiler_report(report, path).map_err(ExportError::from)?;
452            },
453            ExportFormat::Csv | ExportFormat::Json => {
454                let events: Vec<TimingEvent> = report
455                    .slowest_layers
456                    .iter()
457                    .enumerate()
458                    .map(|(i, (name, dur))| TimingEvent {
459                        timestamp_ns: i as u64 * 1_000_000,
460                        duration_ns: dur.as_nanos() as u64,
461                        thread_id: 0,
462                        name: name.clone(),
463                    })
464                    .collect();
465                if events.is_empty() {
466                    return Err(ExportError::EmptyTrace);
467                }
468                match &config.format {
469                    ExportFormat::Csv => {
470                        CsvExporter::export_to_file(&events, path).map_err(ExportError::from)?
471                    },
472                    _ => JsonExporter::export_to_file(&events, path).map_err(ExportError::from)?,
473                }
474            },
475        }
476        Ok(())
477    }
478}
479
480// ─────────────────────────────────────────────────────────────────────────────
481// Tests
482// ─────────────────────────────────────────────────────────────────────────────
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487
488    #[test]
489    fn perfetto_to_profiling_conversion_keeps_every_event() {
490        let mut perfetto = PerfettoTrace::new();
491        for (index, name) in ["load", "forward", "backward"].iter().enumerate() {
492            perfetto.add_event(PerfettoEvent {
493                name: (*name).to_string(),
494                phase: PerfettoPhase::Complete,
495                timestamp_us: 1_000 * (index as u64 + 1),
496                duration_us: Some(500),
497                pid: 1,
498                tid: 7,
499                args: std::collections::HashMap::new(),
500            });
501        }
502
503        let profiling = ProfilingTrace::from(&perfetto);
504        // The old impl discarded `t` and returned an empty trace.
505        assert_eq!(
506            profiling.len(),
507            3,
508            "every event must survive the conversion"
509        );
510        let names: Vec<&str> = profiling.events().iter().map(|e| e.name.as_str()).collect();
511        assert_eq!(
512            names,
513            vec!["load", "forward", "backward"],
514            "order must be preserved"
515        );
516        assert_eq!(profiling.events()[0].thread_id, 7);
517        // Microseconds become nanoseconds.
518        assert_eq!(profiling.events()[0].timestamp_ns, 1_000_000);
519        assert_eq!(profiling.events()[0].duration_ns, 500_000);
520    }
521
522    fn sample_events() -> Vec<TimingEvent> {
523        vec![
524            TimingEvent {
525                timestamp_ns: 0,
526                duration_ns: 1_000_000,
527                thread_id: 0,
528                name: "attention".to_string(),
529            },
530            TimingEvent {
531                timestamp_ns: 1_000_000,
532                duration_ns: 2_000_000,
533                thread_id: 1,
534                name: "ffn".to_string(),
535            },
536            TimingEvent {
537                timestamp_ns: 3_000_000,
538                duration_ns: 500_000,
539                thread_id: 0,
540                name: "layer_norm".to_string(),
541            },
542        ]
543    }
544
545    fn sample_trace() -> ProfilingTrace {
546        let mut t = ProfilingTrace::new();
547        for e in sample_events() {
548            t.add_event(e);
549        }
550        t
551    }
552
553    // ── CsvExporter ──────────────────────────────────────────────────────────
554
555    #[test]
556    fn test_csv_header() {
557        let csv = CsvExporter::export_to_csv(&[]);
558        assert_eq!(csv.trim(), "timestamp_ns,duration_ns,thread_id,name");
559    }
560
561    #[test]
562    fn test_csv_export_values() {
563        let csv = CsvExporter::export_to_csv(&sample_events());
564        assert!(csv.contains("0,1000000,0,attention"));
565        assert!(csv.contains("1000000,2000000,1,ffn"));
566    }
567
568    #[test]
569    fn test_csv_comma_escaping() {
570        let events = vec![TimingEvent {
571            timestamp_ns: 0,
572            duration_ns: 0,
573            thread_id: 0,
574            name: "op,with,commas".to_string(),
575        }];
576        let csv = CsvExporter::export_to_csv(&events);
577        assert!(csv.contains("op\\,with\\,commas"));
578    }
579
580    #[test]
581    fn test_csv_export_to_file() {
582        let path = std::env::temp_dir().join("csv_export_test.csv");
583        CsvExporter::export_to_file(&sample_events(), &path).unwrap();
584        assert!(path.exists());
585        let content = std::fs::read_to_string(&path).unwrap();
586        assert!(content.contains("attention"));
587        std::fs::remove_file(&path).ok();
588    }
589
590    // ── JsonExporter ─────────────────────────────────────────────────────────
591
592    #[test]
593    fn test_json_export_structure() {
594        let json = JsonExporter::export_to_json(&sample_events());
595        assert!(json.starts_with('['));
596        assert!(json.ends_with(']'));
597        assert!(json.contains("\"name\":\"attention\""));
598        assert!(json.contains("\"timestamp_ns\":0"));
599        assert!(json.contains("\"thread_id\":1"));
600    }
601
602    #[test]
603    fn test_json_export_empty() {
604        let json = JsonExporter::export_to_json(&[]);
605        assert_eq!(json, "[]");
606    }
607
608    #[test]
609    fn test_json_export_escaping() {
610        let events = vec![TimingEvent {
611            timestamp_ns: 0,
612            duration_ns: 0,
613            thread_id: 0,
614            name: "say \"hello\"".to_string(),
615        }];
616        let json = JsonExporter::export_to_json(&events);
617        assert!(json.contains("\\\"hello\\\""));
618    }
619
620    #[test]
621    fn test_json_export_to_file() {
622        let path = std::env::temp_dir().join("json_export_test.json");
623        JsonExporter::export_to_file(&sample_events(), &path).unwrap();
624        assert!(path.exists());
625        std::fs::remove_file(&path).ok();
626    }
627
628    // ── ExportFormat ─────────────────────────────────────────────────────────
629
630    #[test]
631    fn test_export_format_extension() {
632        assert_eq!(ExportFormat::Perfetto.extension(), "json");
633        assert_eq!(ExportFormat::ChromeTrace.extension(), "json");
634        assert_eq!(ExportFormat::Tracy.extension(), "csv");
635        assert_eq!(ExportFormat::Csv.extension(), "csv");
636        assert_eq!(ExportFormat::Json.extension(), "json");
637    }
638
639    // ── ProfilingTrace conversions ────────────────────────────────────────────
640
641    #[test]
642    fn test_profiling_trace_to_perfetto() {
643        let trace = sample_trace();
644        let perf = trace.to_perfetto();
645        assert_eq!(perf.len(), 3);
646    }
647
648    #[test]
649    fn test_profiling_trace_to_tracy() {
650        let trace = sample_trace();
651        let tracy = trace.to_tracy();
652        assert_eq!(tracy.zones().len(), 3);
653        assert_eq!(tracy.zones()[0].name, "attention");
654    }
655
656    #[test]
657    fn test_timing_event_from_tracy_zone() {
658        let zone = TracyZone {
659            name: "test".to_string(),
660            timestamp_ns: 5_000,
661            duration_ns: 1_000,
662            thread_id: 2,
663        };
664        let ev = TimingEvent::from(&zone);
665        assert_eq!(ev.timestamp_ns, 5_000);
666        assert_eq!(ev.duration_ns, 1_000);
667        assert_eq!(ev.thread_id, 2);
668        assert_eq!(ev.name, "test");
669    }
670
671    // ── TraceExporter ─────────────────────────────────────────────────────────
672
673    #[test]
674    fn test_trace_exporter_csv() {
675        let trace = sample_trace();
676        let path = std::env::temp_dir().join("unified_export_csv.csv");
677        let config = ExportConfig {
678            format: ExportFormat::Csv,
679            output_path: path.to_string_lossy().into_owned(),
680            compress: false,
681        };
682        TraceExporter::export_all(&trace, &config).unwrap();
683        assert!(path.exists());
684        std::fs::remove_file(&path).ok();
685    }
686
687    #[test]
688    fn test_trace_exporter_json() {
689        let trace = sample_trace();
690        let path = std::env::temp_dir().join("unified_export_json.json");
691        let config = ExportConfig {
692            format: ExportFormat::Json,
693            output_path: path.to_string_lossy().into_owned(),
694            compress: false,
695        };
696        TraceExporter::export_all(&trace, &config).unwrap();
697        assert!(path.exists());
698        let content = std::fs::read_to_string(&path).unwrap();
699        assert!(content.contains("attention"));
700        std::fs::remove_file(&path).ok();
701    }
702
703    #[test]
704    fn test_trace_exporter_empty_returns_error() {
705        let trace = ProfilingTrace::new();
706        let path = std::env::temp_dir().join("should_not_exist.csv");
707        let config = ExportConfig {
708            format: ExportFormat::Csv,
709            output_path: path.to_string_lossy().into_owned(),
710            compress: false,
711        };
712        let result = TraceExporter::export_all(&trace, &config);
713        assert!(matches!(result, Err(ExportError::EmptyTrace)));
714    }
715
716    #[test]
717    fn test_export_error_display() {
718        assert!(ExportError::EmptyTrace.to_string().contains("empty"));
719        assert!(ExportError::UnsupportedFormat("xyz".to_string()).to_string().contains("xyz"));
720        assert!(ExportError::IoError("perm denied".to_string())
721            .to_string()
722            .contains("perm denied"));
723    }
724}
trustformers_debug/export/unified.rs

trustformers_debug/export/
unified.rs
