candle_graph/trace/
events.rs

1//! JSONL event records for `candle-graph/trace/10`.
2
3use serde::{Deserialize, Serialize};
4
5use super::memory::{MemoryAction, MemoryCategory};
6use super::schema::{GradientState, RunOutcome, SpanKind, TraceRunMeta, SCHEMA};
7use crate::phase::ExecutionStep;
8
9/// One JSONL record in a trace stream.
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11#[serde(tag = "kind", rename_all = "snake_case")]
12pub enum TraceEvent {
13    /// Declares schema + run metadata (exactly one per stream).
14    Meta {
15        schema: String,
16        #[serde(flatten)]
17        run: Box<TraceRunMeta>,
18    },
19    SpanStart(SpanStartEvent),
20    SpanEnd(SpanEndEvent),
21    Op(OpEvent),
22    Tensor(TensorEvent),
23    TensorStats(TensorStatsEvent),
24    Memory(MemoryEvent),
25    DeviceMemory(DeviceMemoryEvent),
26    DeviceInterval(DeviceIntervalEvent),
27    Gradient(GradientEvent),
28    Edge(EdgeEvent),
29    Terminal(TerminalEvent),
30}
31
32impl TraceEvent {
33    /// Convenience constructor for the required first meta line.
34    pub fn meta(run: TraceRunMeta) -> Self {
35        Self::Meta {
36            schema: SCHEMA.to_string(),
37            run: Box::new(run),
38        }
39    }
40}
41
42/// Opens a span in the profiler hierarchy.
43#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
44pub struct SpanStartEvent {
45    pub id: String,
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub parent_id: Option<String>,
48    pub name: String,
49    /// Monotonic timestamp in nanoseconds since the profile run started.
50    pub start_ns: u64,
51    /// Span classification (`function` / `op` / `module`). Serialized as `span_kind` because
52    /// the JSONL event discriminator also uses the key `kind`.
53    #[serde(rename = "span_kind")]
54    pub kind: SpanKind,
55    #[serde(default)]
56    pub measured: bool,
57    /// PyTorch-style training step (`forward` / `backward` / `optimizer`) for memory categories.
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub step: Option<ExecutionStep>,
60}
61
62/// Closes a span opened by [`SpanStartEvent`].
63#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
64pub struct SpanEndEvent {
65    pub id: String,
66    /// Wall duration in nanoseconds (from the probe's span guard).
67    #[serde(default)]
68    pub duration_ns: u64,
69}
70
71/// Timed operation inside a span (matmul, add, function body, …).
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73pub struct OpEvent {
74    pub span_id: String,
75    pub op_name: String,
76    #[serde(default)]
77    pub inputs: Vec<String>,
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub output: Option<String>,
80    pub shape: Vec<usize>,
81    pub dtype: String,
82    pub device: String,
83    pub duration_ns: u64,
84    /// Monotonic timestamp for memory timeline ordering (nanoseconds since probe start).
85    #[serde(default)]
86    pub timestamp_ns: u64,
87    /// Dense output tensor footprint; this is not backing-allocation size.
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub output_dense_bytes: Option<u64>,
90    /// Sum of dense input tensor footprints (PyTorch `record_shapes`-style metadata).
91    pub input_dense_bytes: u64,
92}
93
94/// Tensor snapshot associated with a span (create or metadata).
95#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
96pub struct TensorEvent {
97    pub span_id: String,
98    /// Backend tensor identity used for graph joins and deduplication.
99    pub tensor_id: String,
100    /// Optional caller-owned observation label; it is never used as tensor identity.
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub label: Option<String>,
103    pub shape: Vec<usize>,
104    pub dtype: String,
105    pub device: String,
106    #[serde(default)]
107    pub requires_grad: bool,
108    #[serde(default, skip_serializing_if = "Option::is_none")]
109    pub dense_bytes: Option<u64>,
110    #[serde(default)]
111    pub category: MemoryCategory,
112}
113
114/// Numerical summary of a caller-labeled tensor at capture time.
115#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
116pub struct TensorStatsEvent {
117    pub span_id: String,
118    /// Caller-owned semantic label (stable across runs; used for joins).
119    pub label: String,
120    pub shape: Vec<usize>,
121    pub dtype: String,
122    pub elements: u64,
123    /// Count of non-finite elements (NaN or +-inf).
124    pub non_finite: u64,
125    pub rms: f64,
126    pub abs_max: f64,
127    pub mean: f64,
128}
129
130/// Tensor allocation or deallocation (TensorFlow Memory Profile timeline).
131#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
132pub struct MemoryEvent {
133    pub timestamp_ns: u64,
134    /// Backend storage identity. Aliased tensor IDs share this identity.
135    pub storage_id: String,
136    pub tensor_id: String,
137    pub span_id: String,
138    #[serde(default, skip_serializing_if = "Option::is_none")]
139    pub op_name: Option<String>,
140    pub device: String,
141    pub bytes: u64,
142    pub action: MemoryAction,
143    #[serde(default)]
144    pub shape: Vec<usize>,
145    pub dtype: String,
146    #[serde(default)]
147    pub category: MemoryCategory,
148}
149
150/// Optional device-level memory checkpoint (cudaMemGetInfo-style).
151#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
152pub struct DeviceMemoryEvent {
153    pub timestamp_ns: u64,
154    pub device: String,
155    #[serde(default, skip_serializing_if = "Option::is_none")]
156    pub used_bytes: Option<u64>,
157    #[serde(default, skip_serializing_if = "Option::is_none")]
158    pub free_bytes: Option<u64>,
159    /// Caching allocator reserved bytes (PyTorch `memory_reserved`).
160    #[serde(default, skip_serializing_if = "Option::is_none")]
161    pub reserved_bytes: Option<u64>,
162    /// Independently observed device capacity; never derived from used + free.
163    #[serde(default, skip_serializing_if = "Option::is_none")]
164    pub capacity_bytes: Option<u64>,
165}
166
167/// Resolved device interval on one device clock and stream.
168#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
169pub struct DeviceIntervalEvent {
170    pub span_id: String,
171    pub device: String,
172    pub stream_id: String,
173    pub clock_id: String,
174    pub backend: String,
175    pub start_ns: u64,
176    pub duration_ns: u64,
177}
178
179/// Parameter gradient fact recorded during a probe run.
180#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
181pub struct GradientEvent {
182    pub event_id: String,
183    pub root: String,
184    /// Parameter key under `root`.
185    pub key: String,
186    pub state: GradientState,
187    #[serde(default, skip_serializing_if = "Option::is_none")]
188    pub norm: Option<f64>,
189}
190
191impl GradientEvent {
192    pub fn param_key(&self) -> &str {
193        &self.key
194    }
195}
196
197/// A typed call-hierarchy or tensor data-flow edge.
198#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
199#[serde(tag = "edge_kind", rename_all = "snake_case")]
200pub enum EdgeEvent {
201    Call {
202        from_span: String,
203        to_span: String,
204        host_duration_ns: u64,
205    },
206    Data {
207        from_tensor: String,
208        to_tensor: String,
209    },
210}
211
212#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
213pub struct TerminalEvent {
214    pub outcome: RunOutcome,
215    pub timestamp_ns: u64,
216    #[serde(default, skip_serializing_if = "Option::is_none")]
217    pub reason: Option<String>,
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn tensor_event_without_label_remains_readable() {
226        let event: TraceEvent = serde_json::from_str(
227            r#"{"kind":"tensor","span_id":"root","tensor_id":"backend:1","shape":[1],"dtype":"f32","device":"cpu"}"#,
228        )
229        .unwrap();
230        let TraceEvent::Tensor(tensor) = event else {
231            panic!("expected tensor event");
232        };
233        assert_eq!(tensor.tensor_id, "backend:1");
234        assert_eq!(tensor.label, None);
235    }
236}
candle_graph/trace/events.rs

candle_graph/trace/
events.rs
