Skip to main content

candle_graph/trace/
schema.rs

1//! Schema types for the current candle-graph execution-evidence stream.
2
3use std::collections::BTreeMap;
4use std::fmt;
5
6use crate::phase::ExecutionPhase;
7use crate::phase::ExecutionStep;
8
9use serde::{Deserialize, Serialize};
10
11/// Schema identifier written into every trace document and expected on parse.
12pub const SCHEMA: &str = "candle-graph/trace/6";
13
14/// Run metadata carried in the first JSONL `meta` event.
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct TraceRunMeta {
17    /// Stable identifier for this probe run (UUID, build id, or caller-defined key).
18    pub run_id: String,
19    /// Stable caller-owned key shared with NVTX range labels and external artifacts.
20    pub correlation_id: String,
21    /// Analyzed entrypoint, e.g. `train::leworld_loss` or `Model::forward`.
22    pub entrypoint: String,
23    /// Execution phase: `train` or `infer`.
24    pub phase: ExecutionPhase,
25    /// ISO-8601 timestamp when the trace was captured.
26    pub timestamp: String,
27    /// One-based optimizer update or inference invocation selected for capture.
28    pub capture_step: u64,
29    /// Completed invocations before the selected capture.
30    pub warmup_steps: u64,
31    /// Device requested by the caller (`cpu`, `cuda:0`, ...).
32    pub device: String,
33    pub timing_mode: TimingMode,
34    /// Workload-specific provenance such as lesson, batch size, and source revision.
35    pub tags: BTreeMap<String, String>,
36    /// Optional candle crate version observed at probe time.
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub candle_version: Option<String>,
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
42#[serde(rename_all = "snake_case")]
43pub enum TimingMode {
44    Host,
45    DeviceSynchronized,
46}
47
48/// Span classification — mirrors profiler function / op / module hierarchy.
49#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
50#[serde(rename_all = "snake_case")]
51pub enum SpanKind {
52    Function,
53    Op,
54    Module,
55}
56
57impl fmt::Display for SpanKind {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        match self {
60            Self::Function => write!(f, "function"),
61            Self::Op => write!(f, "op"),
62            Self::Module => write!(f, "module"),
63        }
64    }
65}
66
67/// Observed gradient presence / quality.
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum GradientState {
71    Present,
72    Missing,
73    Zero,
74    NonFinite,
75}
76
77impl fmt::Display for GradientState {
78    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79        match self {
80            Self::Present => write!(f, "present"),
81            Self::Missing => write!(f, "missing"),
82            Self::Zero => write!(f, "zero"),
83            Self::NonFinite => write!(f, "non_finite"),
84        }
85    }
86}
87
88/// Aggregated profiler statistics derived from a [`super::document::TraceDocument`].
89#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
90pub struct TraceSummary {
91    pub op_count: usize,
92    pub total_ns: u64,
93    pub span_count: usize,
94    pub root_span_count: usize,
95    pub max_depth: usize,
96    pub alloc_count: usize,
97    pub free_count: usize,
98    pub peak_bytes: u64,
99}
100
101/// Resolved span node assembled from `span_start` / `span_end` events.
102#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
103pub struct SpanRecord {
104    pub id: String,
105    #[serde(default, skip_serializing_if = "Option::is_none")]
106    pub parent_id: Option<String>,
107    pub name: String,
108    pub kind: SpanKind,
109    /// True only for the caller-controlled region used as the run's performance total.
110    pub measured: bool,
111    /// Monotonic timestamp in nanoseconds since the profile run started.
112    pub start_ns: u64,
113    /// True when a matching `span_end` was observed.
114    #[serde(default)]
115    pub closed: bool,
116    /// Wall duration from `span_end` (nanoseconds).
117    #[serde(default)]
118    pub duration_ns: u64,
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub step: Option<ExecutionStep>,
121}