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    /// Whether the single measured region is bounded by device synchronizations.
34    /// Nested span durations remain governed by `timing_mode`.
35    #[serde(default)]
36    pub measured_region_device_synchronized: bool,
37    pub timing_mode: TimingMode,
38    /// Workload-specific provenance such as lesson, batch size, and source revision.
39    pub tags: BTreeMap<String, String>,
40    /// Optional candle crate version observed at probe time.
41    #[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/// Span classification — mirrors profiler function / op / module hierarchy.
53#[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/// Observed gradient presence / quality.
72#[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/// Aggregated profiler statistics derived from a [`super::document::TraceDocument`].
93#[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/// Resolved span node assembled from `span_start` / `span_end` events.
106#[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    /// True only for the caller-controlled region used as the run's performance total.
114    pub measured: bool,
115    /// Monotonic timestamp in nanoseconds since the profile run started.
116    pub start_ns: u64,
117    /// True when a matching `span_end` was observed.
118    #[serde(default)]
119    pub closed: bool,
120    /// Wall duration from `span_end` (nanoseconds).
121    #[serde(default)]
122    pub duration_ns: u64,
123    #[serde(default, skip_serializing_if = "Option::is_none")]
124    pub step: Option<ExecutionStep>,
125}