candle_graph/
profile.rs

1//! Runtime profiler: attach during train or inference to emit timed operation/edge traces.
2//!
3//! Writes `candle-graph/runtime/3` JSONL compatible with [`crate::runtime`] and mergeable
4//! into the unified model IR via static `static_id` correlation.
5
6use std::collections::HashMap;
7use std::fs::OpenOptions;
8use std::io;
9use std::path::{Path, PathBuf};
10use std::time::Instant;
11
12use anyhow::{Context, Result};
13
14use crate::phase::ExecutionPhase;
15use crate::runtime::{
16    EdgeTimingObservation, OperationObservation, RunMetadata, RuntimeTraceWriter, SCHEMA_V3,
17};
18
19/// Configuration for a profiling session started alongside train or inference.
20#[derive(Debug, Clone)]
21pub struct ProfileConfig {
22    pub entrypoint: String,
23    pub phase: ExecutionPhase,
24    pub profile: String,
25    pub cargo_features: Vec<String>,
26    pub analysis_id: Option<String>,
27    pub build_id: Option<String>,
28}
29
30impl ProfileConfig {
31    pub fn train(entrypoint: impl Into<String>) -> Self {
32        Self {
33            entrypoint: entrypoint.into(),
34            phase: ExecutionPhase::Train,
35            profile: "release".into(),
36            cargo_features: Vec::new(),
37            analysis_id: None,
38            build_id: None,
39        }
40    }
41
42    pub fn infer(entrypoint: impl Into<String>) -> Self {
43        Self {
44            entrypoint: entrypoint.into(),
45            phase: ExecutionPhase::Infer,
46            profile: "release".into(),
47            cargo_features: Vec::new(),
48            analysis_id: None,
49            build_id: None,
50        }
51    }
52}
53
54struct ActiveOp {
55    static_id: String,
56    op: String,
57    inputs: Vec<String>,
58    started: Instant,
59}
60
61/// Streaming profiler session writing JSONL runtime v3 events.
62pub struct ProfileSession {
63    path: PathBuf,
64    writer: RuntimeTraceWriter<io::BufWriter<std::fs::File>>,
65    phase: ExecutionPhase,
66    step: u64,
67    active: HashMap<String, ActiveOp>,
68    event_counter: u64,
69}
70
71impl ProfileSession {
72    /// Open (or truncate) a JSONL profile trace at `path`.
73    pub fn open(path: impl AsRef<Path>, config: ProfileConfig) -> Result<Self> {
74        let path = path.as_ref().to_path_buf();
75        if let Some(parent) = path.parent() {
76            std::fs::create_dir_all(parent)
77                .with_context(|| format!("create profile dir {}", parent.display()))?;
78        }
79        let file = OpenOptions::new()
80            .create(true)
81            .write(true)
82            .truncate(true)
83            .open(&path)
84            .with_context(|| format!("open profile trace {}", path.display()))?;
85        let run = RunMetadata {
86            entrypoint: config.entrypoint,
87            profile: config.profile,
88            cargo_features: config.cargo_features,
89            cfg: Vec::new(),
90            analysis_id: config.analysis_id,
91            build_id: config.build_id,
92            phase: Some(config.phase.as_str().to_string()),
93        };
94        let writer = RuntimeTraceWriter::new_with_schema(io::BufWriter::new(file), SCHEMA_V3, run)?;
95        Ok(Self {
96            path,
97            writer,
98            phase: config.phase,
99            step: 0,
100            active: HashMap::new(),
101            event_counter: 0,
102        })
103    }
104
105    pub fn phase(&self) -> ExecutionPhase {
106        self.phase
107    }
108
109    pub fn step(&self) -> u64 {
110        self.step
111    }
112
113    pub fn set_step(&mut self, step: u64) {
114        self.step = step;
115    }
116
117    /// Begin timing an operation; returns an `event_id` for [`Self::end_operation`].
118    pub fn begin_operation(
119        &mut self,
120        static_id: impl Into<String>,
121        op: impl Into<String>,
122        inputs: &[String],
123    ) -> Result<String> {
124        self.event_counter += 1;
125        let event_id = format!(
126            "{}-op-{}-{}",
127            self.phase.as_str(),
128            self.step,
129            self.event_counter
130        );
131        self.active.insert(
132            event_id.clone(),
133            ActiveOp {
134                static_id: static_id.into(),
135                op: op.into(),
136                inputs: inputs.to_vec(),
137                started: Instant::now(),
138            },
139        );
140        Ok(event_id)
141    }
142
143    /// End a timed operation and emit a runtime v3 operation observation.
144    pub fn end_operation(&mut self, event_id: &str, output: Option<String>) -> Result<u64> {
145        let active = self
146            .active
147            .remove(event_id)
148            .with_context(|| format!("unknown profile operation `{event_id}`"))?;
149        let duration_ns = active.started.elapsed().as_nanos().min(u64::MAX as u128) as u64;
150        self.writer.operation(OperationObservation {
151            event_id: event_id.to_string(),
152            op: active.op,
153            static_id: Some(active.static_id),
154            source: None,
155            inputs: active.inputs,
156            output,
157            step: Some(self.step),
158            duration_ns: Some(duration_ns),
159        })?;
160        Ok(duration_ns)
161    }
162
163    /// Record average-friendly edge timing between two static tensor/operation ids.
164    pub fn record_edge_timing(
165        &mut self,
166        from_static_id: impl Into<String>,
167        to_static_id: impl Into<String>,
168        duration_ns: u64,
169    ) -> Result<()> {
170        self.event_counter += 1;
171        let event_id = format!(
172            "{}-edge-{}-{}",
173            self.phase.as_str(),
174            self.step,
175            self.event_counter
176        );
177        self.writer.edge_timing(EdgeTimingObservation {
178            event_id,
179            from_static_id: from_static_id.into(),
180            to_static_id: to_static_id.into(),
181            duration_ns,
182            step: Some(self.step),
183        })
184    }
185
186    pub fn flush(&mut self) -> Result<()> {
187        self.writer.flush()
188    }
189
190    pub fn finish(self) -> Result<PathBuf> {
191        self.writer.finish()?;
192        Ok(self.path)
193    }
194}
195
196/// Convenience RAII guard for [`ProfileSession::begin_operation`] / [`ProfileSession::end_operation`].
197pub struct TimedOperation<'a> {
198    session: &'a mut ProfileSession,
199    event_id: String,
200}
201
202impl<'a> TimedOperation<'a> {
203    pub fn begin(
204        session: &'a mut ProfileSession,
205        static_id: impl Into<String>,
206        op: impl Into<String>,
207        inputs: &[String],
208    ) -> Result<Self> {
209        let event_id = session.begin_operation(static_id, op, inputs)?;
210        Ok(Self { session, event_id })
211    }
212}
213
214impl Drop for TimedOperation<'_> {
215    fn drop(&mut self) {
216        let _ = self.session.end_operation(&self.event_id, None);
217    }
218}
candle_graph/profile.rs

candle_graph/
profile.rs
