candle_graph/graph/
model.rs

1//! Execution graph model (`candle-graph/graph/3`).
2
3use serde::{Deserialize, Serialize};
4
5use crate::trace::memory::{MemoryCategory, MemoryProfile, MemorySummary};
6
7/// Schema identifier for [`ExecutionGraph`] documents.
8pub const SCHEMA: &str = "candle-graph/graph/3";
9
10/// Hierarchical execution graph built from a trace document.
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub struct ExecutionGraph {
13    pub schema: String,
14    /// Flat list of span and op nodes (parent links encode the tree).
15    pub spans: Vec<GraphNode>,
16    /// Call hierarchy and tensor data-flow edges with timing.
17    pub edges: Vec<GraphEdge>,
18    /// Semantic tensor checkpoints captured by the application.
19    pub tensors: Vec<TensorRecord>,
20    pub gradients: Vec<GradientRecord>,
21    pub summary: GraphSummary,
22    pub memory: MemoryProfile,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26pub struct TensorRecord {
27    pub span_id: String,
28    pub tensor_id: String,
29    pub shape: Vec<usize>,
30    pub dtype: String,
31    pub device: String,
32    pub requires_grad: bool,
33    pub storage_bytes: u64,
34    pub category: MemoryCategory,
35}
36
37/// One span or attached operation in the execution tree.
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct GraphNode {
40    pub id: String,
41    pub parent_id: Option<String>,
42    pub name: String,
43    pub kind: GraphNodeKind,
44    pub start_ns: u64,
45    /// Wall time excluding nested spans/ops (TensorFlow profiler style).
46    pub self_time_ns: u64,
47    /// Inclusive wall time for this node.
48    pub total_time_ns: u64,
49    #[serde(default, skip_serializing_if = "Option::is_none")]
50    pub shape: Option<Vec<usize>>,
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub dtype: Option<String>,
53    #[serde(default, skip_serializing_if = "Option::is_none")]
54    pub device: Option<String>,
55    /// Total bytes requested by this node's ops (TF `bytes`).
56    #[serde(default)]
57    pub bytes: u64,
58    /// Peak live bytes in subtree (TF `peak_bytes`).
59    #[serde(default)]
60    pub peak_bytes: u64,
61    /// Bytes still live when node finishes (TF `residual_bytes`).
62    #[serde(default)]
63    pub residual_bytes: u64,
64    /// Op output storage (derived from shape×dtype).
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub storage_bytes: Option<u64>,
67}
68
69#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71pub enum GraphNodeKind {
72    Root,
73    Function,
74    Module,
75    Op,
76    #[serde(other)]
77    Other,
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81pub struct GraphEdge {
82    pub from: String,
83    pub to: String,
84    pub kind: GraphEdgeKind,
85    pub duration_ns: u64,
86    #[serde(default, skip_serializing_if = "Option::is_none")]
87    pub label: Option<String>,
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
91#[serde(rename_all = "snake_case")]
92pub enum GraphEdgeKind {
93    Call,
94    Data,
95}
96
97#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
98pub struct GradientRecord {
99    pub root: String,
100    pub key: String,
101    pub state: GradientRecordState,
102    #[serde(default, skip_serializing_if = "Option::is_none")]
103    pub norm: Option<f64>,
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
107#[serde(rename_all = "snake_case")]
108pub enum GradientRecordState {
109    Present,
110    Missing,
111    Zero,
112    NonFinite,
113}
114
115#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
116pub struct GraphSummary {
117    pub entrypoint: String,
118    /// Root span inclusive wall time in milliseconds.
119    pub total_ms: f64,
120    /// Top spans by self time (excluding op leaf nodes).
121    pub slowest_spans: Vec<SlowSpan>,
122    /// Top spans by requested bytes (TF scope -order_by bytes).
123    pub heaviest_spans: Vec<HeavySpan>,
124    pub memory: MemorySummary,
125}
126
127#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
128pub struct SlowSpan {
129    pub id: String,
130    pub name: String,
131    pub self_time_ns: u64,
132}
133
134#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
135pub struct HeavySpan {
136    pub id: String,
137    pub name: String,
138    pub bytes: u64,
139    pub peak_bytes: u64,
140}
141
142impl ExecutionGraph {
143    pub fn node(&self, id: &str) -> Option<&GraphNode> {
144        self.spans.iter().find(|n| n.id == id)
145    }
146
147    pub fn children(&self, parent_id: &str) -> Vec<&GraphNode> {
148        self.spans
149            .iter()
150            .filter(|n| n.parent_id.as_deref() == Some(parent_id))
151            .collect()
152    }
153}
candle_graph/graph/model.rs

candle_graph/graph/
model.rs
