1use std::collections::BTreeMap;
4use std::fmt;
5
6use crate::phase::ExecutionPhase;
7use crate::phase::ExecutionStep;
8
9use serde::{Deserialize, Serialize};
10
11pub const SCHEMA: &str = "candle-graph/trace/6";
13
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct TraceRunMeta {
17 pub run_id: String,
19 pub correlation_id: String,
21 pub entrypoint: String,
23 pub phase: ExecutionPhase,
25 pub timestamp: String,
27 pub capture_step: u64,
29 pub warmup_steps: u64,
31 pub device: String,
33 pub timing_mode: TimingMode,
34 pub tags: BTreeMap<String, String>,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub candle_version: Option<String>,
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
42#[serde(rename_all = "snake_case")]
43pub enum TimingMode {
44 Host,
45 DeviceSynchronized,
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
50#[serde(rename_all = "snake_case")]
51pub enum SpanKind {
52 Function,
53 Op,
54 Module,
55}
56
57impl fmt::Display for SpanKind {
58 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59 match self {
60 Self::Function => write!(f, "function"),
61 Self::Op => write!(f, "op"),
62 Self::Module => write!(f, "module"),
63 }
64 }
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum GradientState {
71 Present,
72 Missing,
73 Zero,
74 NonFinite,
75}
76
77impl fmt::Display for GradientState {
78 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79 match self {
80 Self::Present => write!(f, "present"),
81 Self::Missing => write!(f, "missing"),
82 Self::Zero => write!(f, "zero"),
83 Self::NonFinite => write!(f, "non_finite"),
84 }
85 }
86}
87
88#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
90pub struct TraceSummary {
91 pub op_count: usize,
92 pub total_ns: u64,
93 pub span_count: usize,
94 pub root_span_count: usize,
95 pub max_depth: usize,
96 pub alloc_count: usize,
97 pub free_count: usize,
98 pub peak_bytes: u64,
99}
100
101#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
103pub struct SpanRecord {
104 pub id: String,
105 #[serde(default, skip_serializing_if = "Option::is_none")]
106 pub parent_id: Option<String>,
107 pub name: String,
108 pub kind: SpanKind,
109 pub measured: bool,
111 pub start_ns: u64,
113 #[serde(default)]
115 pub closed: bool,
116 #[serde(default)]
118 pub duration_ns: u64,
119 #[serde(default, skip_serializing_if = "Option::is_none")]
120 pub step: Option<ExecutionStep>,
121}