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 #[serde(default)]
36 pub measured_region_device_synchronized: bool,
37 pub timing_mode: TimingMode,
38 pub tags: BTreeMap<String, String>,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub candle_version: Option<String>,
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
46#[serde(rename_all = "snake_case")]
47pub enum TimingMode {
48 Host,
49 DeviceSynchronized,
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
54#[serde(rename_all = "snake_case")]
55pub enum SpanKind {
56 Function,
57 Op,
58 Module,
59}
60
61impl fmt::Display for SpanKind {
62 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63 match self {
64 Self::Function => write!(f, "function"),
65 Self::Op => write!(f, "op"),
66 Self::Module => write!(f, "module"),
67 }
68 }
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
73#[serde(rename_all = "snake_case")]
74pub enum GradientState {
75 Present,
76 Missing,
77 Zero,
78 NonFinite,
79}
80
81impl fmt::Display for GradientState {
82 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
83 match self {
84 Self::Present => write!(f, "present"),
85 Self::Missing => write!(f, "missing"),
86 Self::Zero => write!(f, "zero"),
87 Self::NonFinite => write!(f, "non_finite"),
88 }
89 }
90}
91
92#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
94pub struct TraceSummary {
95 pub op_count: usize,
96 pub total_ns: u64,
97 pub span_count: usize,
98 pub root_span_count: usize,
99 pub max_depth: usize,
100 pub alloc_count: usize,
101 pub free_count: usize,
102 pub peak_bytes: u64,
103}
104
105#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
107pub struct SpanRecord {
108 pub id: String,
109 #[serde(default, skip_serializing_if = "Option::is_none")]
110 pub parent_id: Option<String>,
111 pub name: String,
112 pub kind: SpanKind,
113 pub measured: bool,
115 pub start_ns: u64,
117 #[serde(default)]
119 pub closed: bool,
120 #[serde(default)]
122 pub duration_ns: u64,
123 #[serde(default, skip_serializing_if = "Option::is_none")]
124 pub step: Option<ExecutionStep>,
125}