candle_graph/graph/
model.rs

1//! Execution graph model (`candle-graph/graph/5`).
2
3use serde::{Deserialize, Serialize};
4
5use crate::trace::memory::MemoryCategory;
6
7/// Schema identifier for [`ExecutionGraph`] documents.
8pub const SCHEMA: &str = "candle-graph/graph/5";
9
10/// Hierarchical execution graph built from a trace document.
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub struct ExecutionGraph {
13    #[serde(deserialize_with = "deserialize_schema")]
14    pub schema: String,
15    /// Flat list of span and op nodes (parent links encode the tree).
16    pub spans: Vec<GraphNode>,
17    /// Call hierarchy and tensor data-flow edges with timing.
18    pub edges: Vec<GraphEdge>,
19    /// Semantic tensor checkpoints captured by the application.
20    pub tensors: Vec<TensorRecord>,
21    pub gradients: Vec<GradientRecord>,
22    pub summary: GraphSummary,
23}
24
25fn deserialize_schema<'de, D>(deserializer: D) -> std::result::Result<String, D::Error>
26where
27    D: serde::Deserializer<'de>,
28{
29    let schema = String::deserialize(deserializer)?;
30    if schema != SCHEMA {
31        return Err(serde::de::Error::custom(format_args!(
32            "unsupported graph schema {schema:?}; expected {SCHEMA:?}"
33        )));
34    }
35    Ok(schema)
36}
37
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct TensorRecord {
40    pub span_id: String,
41    pub tensor_id: String,
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub label: Option<String>,
44    pub shape: Vec<usize>,
45    pub dtype: String,
46    pub device: String,
47    pub requires_grad: bool,
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub dense_bytes: Option<u64>,
50    pub category: MemoryCategory,
51}
52
53/// One span or attached operation in the execution tree.
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
55pub struct GraphNode {
56    pub id: String,
57    pub parent_id: Option<String>,
58    pub name: String,
59    pub kind: GraphNodeKind,
60    pub start_ns: u64,
61    /// Host wall time excluding nested spans/ops.
62    pub host_self_time_ns: u64,
63    /// Inclusive host wall time for this node.
64    pub host_total_time_ns: u64,
65    /// Overlap-safe timings kept separate for every incomparable device clock.
66    #[serde(default, skip_serializing_if = "Vec::is_empty")]
67    pub device_timings: Vec<DeviceNodeTiming>,
68    #[serde(default, skip_serializing_if = "Option::is_none")]
69    pub shape: Option<Vec<usize>>,
70    #[serde(default, skip_serializing_if = "Option::is_none")]
71    pub dtype: Option<String>,
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    pub device: Option<String>,
74    /// Logical storage bytes directly allocated by this node.
75    #[serde(default, skip_serializing_if = "Option::is_none")]
76    pub allocated_bytes: Option<u64>,
77    /// Peak logical live bytes in this node's subtree.
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub peak_live_bytes: Option<u64>,
80    /// Logical bytes still live when this node finishes.
81    #[serde(default, skip_serializing_if = "Option::is_none")]
82    pub residual_bytes: Option<u64>,
83    /// Dense tensor footprint (derived from shape×dtype), not backing allocation size.
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub dense_bytes: Option<u64>,
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
89#[serde(rename_all = "snake_case")]
90pub enum GraphNodeKind {
91    Root,
92    Function,
93    Module,
94    Op,
95    Tensor,
96    #[serde(other)]
97    Other,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
101pub struct DeviceNodeTiming {
102    pub device: String,
103    pub clock_id: String,
104    pub busy_ns: u64,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
108#[serde(tag = "kind", rename_all = "snake_case")]
109pub enum GraphEdge {
110    Call {
111        from_span: String,
112        to_span: String,
113        host_duration_ns: u64,
114        #[serde(default, skip_serializing_if = "Option::is_none")]
115        label: Option<String>,
116    },
117    Data {
118        from_tensor: String,
119        to_tensor: String,
120        #[serde(default, skip_serializing_if = "Option::is_none")]
121        label: Option<String>,
122    },
123}
124
125#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
126pub struct GradientRecord {
127    pub root: String,
128    pub key: String,
129    pub state: GradientRecordState,
130    #[serde(default, skip_serializing_if = "Option::is_none")]
131    pub norm: Option<f64>,
132}
133
134#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
135#[serde(rename_all = "snake_case")]
136pub enum GradientRecordState {
137    Present,
138    Missing,
139    Zero,
140    NonFinite,
141}
142
143#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
144pub struct GraphSummary {
145    pub entrypoint: String,
146    pub outer_wall_time_ns: u64,
147    pub slowest_host_spans: Vec<HostSpanCost>,
148    pub slowest_device_spans: Vec<DeviceSpanCost>,
149    /// Top spans by known logical allocation bytes.
150    pub heaviest_spans: Vec<HeavySpan>,
151}
152
153#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
154pub struct HostSpanCost {
155    pub id: String,
156    pub name: String,
157    /// Structural relation to the measured region. Concurrent spans remain outside its tree.
158    pub scope: MeasuredHostScope,
159    /// Full host self time across the span's complete interval.
160    pub host_self_time_ns: u64,
161    /// Host self time clipped to the measured interval.
162    pub measured_overlap_self_time_ns: u64,
163    /// Inclusive duration across the span's complete interval.
164    pub full_duration_ns: u64,
165    /// Inclusive duration intersected with the measured interval.
166    pub measured_overlap_duration_ns: u64,
167}
168
169/// How a host span entered measured-scope headline analysis.
170#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
171#[serde(rename_all = "snake_case")]
172pub enum MeasuredHostScope {
173    /// The span is the measured region or one of its structural descendants.
174    #[default]
175    MeasuredSubtree,
176    /// The span is structurally outside the measured subtree but its host interval overlaps it.
177    ConcurrentOverlap,
178}
179
180impl MeasuredHostScope {
181    pub const fn as_str(self) -> &'static str {
182        match self {
183            Self::MeasuredSubtree => "measured_subtree",
184            Self::ConcurrentOverlap => "concurrent_overlap",
185        }
186    }
187}
188
189#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
190pub struct DeviceSpanCost {
191    pub id: String,
192    pub name: String,
193    pub device: String,
194    pub clock_id: String,
195    pub device_busy_ns: u64,
196}
197
198#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
199pub struct HeavySpan {
200    pub id: String,
201    pub name: String,
202    pub allocated_bytes: u64,
203    pub peak_live_bytes: Option<u64>,
204}
205
206impl ExecutionGraph {
207    pub fn node(&self, id: &str) -> Option<&GraphNode> {
208        self.spans.iter().find(|n| n.id == id)
209    }
210
211    pub fn children(&self, parent_id: &str) -> Vec<&GraphNode> {
212        self.spans
213            .iter()
214            .filter(|n| n.parent_id.as_deref() == Some(parent_id))
215            .collect()
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn measured_scope_fields_are_required() {
225        let error = serde_json::from_value::<HostSpanCost>(serde_json::json!({
226            "id": "span",
227            "name": "work",
228            "host_self_time_ns": 17
229        }))
230        .unwrap_err();
231
232        assert!(error.to_string().contains("scope"));
233    }
234
235    #[test]
236    fn measured_host_scope_has_stable_wire_names() {
237        assert_eq!(
238            serde_json::to_value(MeasuredHostScope::MeasuredSubtree).unwrap(),
239            serde_json::json!("measured_subtree")
240        );
241        assert_eq!(
242            serde_json::to_value(MeasuredHostScope::ConcurrentOverlap).unwrap(),
243            serde_json::json!("concurrent_overlap")
244        );
245    }
246}
candle_graph/graph/model.rs

candle_graph/graph/
model.rs
