candle_graph/trace/
events.rs

1//! JSONL event records for `candle-graph/trace/6`.
2
3use serde::{Deserialize, Serialize};
4
5use super::memory::{MemoryAction, MemoryCategory};
6use super::schema::{GradientState, 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: TraceRunMeta,
18    },
19    SpanStart(SpanStartEvent),
20    SpanEnd(SpanEndEvent),
21    Op(OpEvent),
22    Tensor(TensorEvent),
23    Memory(MemoryEvent),
24    DeviceMemory(DeviceMemoryEvent),
25    Gradient(GradientEvent),
26    Edge(EdgeEvent),
27}
28
29impl TraceEvent {
30    /// Convenience constructor for the required first meta line.
31    pub fn meta(run: TraceRunMeta) -> Self {
32        Self::Meta {
33            schema: SCHEMA.to_string(),
34            run,
35        }
36    }
37}
38
39/// Opens a span in the profiler hierarchy.
40#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
41pub struct SpanStartEvent {
42    pub id: String,
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub parent_id: Option<String>,
45    pub name: String,
46    /// Monotonic timestamp in nanoseconds since the profile run started.
47    pub start_ns: u64,
48    /// Span classification (`function` / `op` / `module`). Serialized as `span_kind` because
49    /// the JSONL event discriminator also uses the key `kind`.
50    #[serde(rename = "span_kind")]
51    pub kind: SpanKind,
52    #[serde(default)]
53    pub measured: bool,
54    /// PyTorch-style training step (`forward` / `backward` / `optimizer`) for memory categories.
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub step: Option<ExecutionStep>,
57}
58
59/// Closes a span opened by [`SpanStartEvent`].
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub struct SpanEndEvent {
62    pub id: String,
63    /// Wall duration in nanoseconds (from the probe's span guard).
64    #[serde(default)]
65    pub duration_ns: u64,
66}
67
68/// Timed operation inside a span (matmul, add, function body, …).
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
70pub struct OpEvent {
71    pub span_id: String,
72    pub op_name: String,
73    #[serde(default)]
74    pub inputs: Vec<String>,
75    #[serde(default, skip_serializing_if = "Option::is_none")]
76    pub output: Option<String>,
77    #[serde(default)]
78    pub shape: Vec<usize>,
79    pub dtype: String,
80    pub device: String,
81    pub duration_ns: u64,
82    /// Monotonic timestamp for memory timeline ordering (nanoseconds since probe start).
83    #[serde(default)]
84    pub timestamp_ns: u64,
85    /// Dense output storage bytes; derived from shape×dtype when omitted at parse time.
86    #[serde(default, skip_serializing_if = "Option::is_none")]
87    pub storage_bytes: Option<u64>,
88    /// Sum of input tensor storage (PyTorch `record_shapes` footprint).
89    #[serde(default)]
90    pub input_storage_bytes: u64,
91}
92
93/// Tensor snapshot associated with a span (create or metadata).
94#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
95pub struct TensorEvent {
96    pub span_id: String,
97    pub tensor_id: String,
98    #[serde(default)]
99    pub shape: Vec<usize>,
100    pub dtype: String,
101    pub device: String,
102    #[serde(default)]
103    pub requires_grad: bool,
104    #[serde(default, skip_serializing_if = "Option::is_none")]
105    pub storage_bytes: Option<u64>,
106    #[serde(default)]
107    pub category: MemoryCategory,
108}
109
110/// Tensor allocation or deallocation (TensorFlow Memory Profile timeline).
111#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
112pub struct MemoryEvent {
113    pub timestamp_ns: u64,
114    pub tensor_id: String,
115    pub span_id: String,
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub op_name: Option<String>,
118    pub device: String,
119    pub bytes: u64,
120    pub action: MemoryAction,
121    #[serde(default)]
122    pub shape: Vec<usize>,
123    pub dtype: String,
124    #[serde(default)]
125    pub category: MemoryCategory,
126}
127
128/// Optional device-level memory checkpoint (cudaMemGetInfo-style).
129#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
130pub struct DeviceMemoryEvent {
131    pub timestamp_ns: u64,
132    pub device: String,
133    pub used_bytes: u64,
134    pub free_bytes: u64,
135    /// Caching allocator reserved bytes (PyTorch `memory_reserved`); optional.
136    #[serde(default, skip_serializing_if = "Option::is_none")]
137    pub reserved_bytes: Option<u64>,
138}
139
140/// Parameter gradient fact recorded during a probe run.
141#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
142pub struct GradientEvent {
143    pub event_id: String,
144    pub root: String,
145    /// Parameter key under `root`.
146    pub key: String,
147    pub state: GradientState,
148    #[serde(default, skip_serializing_if = "Option::is_none")]
149    pub norm: Option<f64>,
150}
151
152impl GradientEvent {
153    pub fn param_key(&self) -> &str {
154        &self.key
155    }
156}
157
158/// Data-flow edge timing between two spans.
159#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
160pub struct EdgeEvent {
161    pub from_span: String,
162    pub to_span: String,
163    pub duration_ns: u64,
164}
candle_graph/trace/events.rs

candle_graph/trace/
events.rs
