candle_graph/instrument/
session.rs

1//! Representative-run profiler session — emits `candle-graph/trace/6` JSONL.
2
3use std::cell::RefCell;
4use std::collections::BTreeMap;
5use std::fs::{File, OpenOptions};
6use std::io::{self, Write};
7use std::path::{Path, PathBuf};
8use std::time::{Instant, SystemTime, UNIX_EPOCH};
9
10use anyhow::{Context, Result};
11use serde::Serialize;
12
13use crate::phase::ExecutionPhase;
14use crate::trace::events::{
15    DeviceMemoryEvent, GradientEvent, MemoryEvent, OpEvent, SpanEndEvent, SpanStartEvent,
16    TensorEvent, TraceEvent,
17};
18use crate::trace::memory::category_for_step;
19use crate::trace::memory::{resolve_storage_bytes, MemoryAction};
20use crate::trace::schema::{GradientState, TimingMode, TraceRunMeta};
21
22use super::span::{MemoryRecord, OpRecord, SpanGuard, SpanId, SpanKind, TensorRecord};
23
24/// Streaming trace session writing TensorFlow-Profiler-style span JSONL.
25pub struct TraceSession {
26    path: PathBuf,
27    inner: RefCell<SessionInner>,
28}
29
30/// Required provenance for one representative profile run.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct ProfileRun {
33    pub entrypoint: String,
34    pub correlation_id: String,
35    pub phase: ExecutionPhase,
36    /// One-based selected update or inference invocation.
37    pub capture_step: u64,
38    pub warmup_steps: u64,
39    pub device: String,
40    pub timing_mode: TimingMode,
41    pub tags: BTreeMap<String, String>,
42}
43
44impl ProfileRun {
45    pub fn training(
46        entrypoint: impl Into<String>,
47        capture_step: u64,
48        device: impl Into<String>,
49    ) -> Self {
50        let entrypoint = entrypoint.into();
51        Self {
52            correlation_id: format!("{entrypoint}/update-{capture_step}"),
53            entrypoint,
54            phase: ExecutionPhase::Train,
55            capture_step,
56            warmup_steps: capture_step.saturating_sub(1),
57            device: device.into(),
58            timing_mode: TimingMode::Host,
59            tags: BTreeMap::new(),
60        }
61    }
62
63    pub fn tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
64        self.tags.insert(key.into(), value.into());
65        self
66    }
67
68    pub fn correlation_id(mut self, value: impl Into<String>) -> Self {
69        self.correlation_id = value.into();
70        self
71    }
72
73    pub fn device_synchronized(mut self) -> Self {
74        self.timing_mode = TimingMode::DeviceSynchronized;
75        self
76    }
77}
78
79struct SessionInner {
80    writer: io::BufWriter<File>,
81    span_stack: Vec<u64>,
82    span_steps: Vec<Option<crate::phase::ExecutionStep>>,
83    next_span_id: u64,
84    next_event_id: u64,
85    id_buf: String,
86    probe_started: Instant,
87    sticky_error: Option<String>,
88}
89
90impl TraceSession {
91    /// Open a trace and own its single root span until [`Self::finish`].
92    pub fn open(path: impl AsRef<Path>, run: ProfileRun) -> Result<Self> {
93        anyhow::ensure!(
94            run.capture_step > 0,
95            "capture_step must be one-based and greater than zero"
96        );
97        let path = path.as_ref().to_path_buf();
98        if let Some(parent) = path.parent() {
99            std::fs::create_dir_all(parent)
100                .with_context(|| format!("create trace dir {}", parent.display()))?;
101        }
102        let file = OpenOptions::new()
103            .create(true)
104            .write(true)
105            .truncate(true)
106            .open(&path)
107            .with_context(|| format!("open trace {}", path.display()))?;
108        let mut writer = io::BufWriter::new(file);
109        let run_id = new_run_id();
110        let meta = TraceRunMeta {
111            run_id: run_id.clone(),
112            correlation_id: run.correlation_id,
113            entrypoint: run.entrypoint.clone(),
114            phase: run.phase,
115            timestamp: utc_iso8601_now(),
116            capture_step: run.capture_step,
117            warmup_steps: run.warmup_steps,
118            device: run.device,
119            timing_mode: run.timing_mode,
120            tags: run.tags,
121            candle_version: None,
122        };
123        write_event(&mut writer, &TraceEvent::meta(meta))?;
124        write_event(
125            &mut writer,
126            &TraceEvent::SpanStart(SpanStartEvent {
127                id: span_id_string(1),
128                parent_id: None,
129                name: run.entrypoint,
130                start_ns: 0,
131                kind: SpanKind::Function,
132                measured: false,
133                step: None,
134            }),
135        )?;
136        Ok(Self {
137            path,
138            inner: RefCell::new(SessionInner {
139                writer,
140                span_stack: vec![1],
141                span_steps: vec![None],
142                next_span_id: 1,
143                next_event_id: 0,
144                id_buf: String::with_capacity(24),
145                probe_started: Instant::now(),
146                sticky_error: None,
147            }),
148        })
149    }
150
151    /// Begin a nested span; parent is the top of the session span stack (TF Profiler call tree).
152    pub fn begin_span(&self, name: impl Into<String>, kind: SpanKind) -> SpanGuard<'_> {
153        self.begin_span_inner(name, kind, None, false)
154    }
155
156    /// Begin the single caller-controlled region used for total-time comparisons.
157    pub fn begin_measurement(&self, name: impl Into<String>) -> SpanGuard<'_> {
158        self.begin_span_inner(name, SpanKind::Function, None, true)
159    }
160
161    /// Begin a span tagged with a PyTorch-style training step (`forward` / `backward` / `optimizer`).
162    pub fn begin_step_span(
163        &self,
164        name: impl Into<String>,
165        step: crate::phase::ExecutionStep,
166        kind: SpanKind,
167    ) -> SpanGuard<'_> {
168        self.begin_span_inner(name, kind, Some(step), false)
169    }
170
171    fn begin_span_inner(
172        &self,
173        name: impl Into<String>,
174        kind: SpanKind,
175        step: Option<crate::phase::ExecutionStep>,
176        measured: bool,
177    ) -> SpanGuard<'_> {
178        let started = Instant::now();
179        let start_ns = self.elapsed_ns();
180        let mut inner = self.inner.borrow_mut();
181        inner.next_span_id += 1;
182        let span_id = inner.next_span_id;
183        let parent_id = inner.span_stack.last().copied().map(span_id_string);
184
185        format_span_id(&mut inner.id_buf, span_id);
186        let id_str = inner.id_buf.clone();
187
188        if let Err(error) = write_event(
189            &mut inner.writer,
190            &TraceEvent::SpanStart(SpanStartEvent {
191                id: id_str,
192                parent_id,
193                name: name.into(),
194                start_ns,
195                kind,
196                measured,
197                step,
198            }),
199        ) {
200            inner.sticky_error.get_or_insert_with(|| error.to_string());
201        }
202
203        inner.span_stack.push(span_id);
204        inner.span_steps.push(step);
205
206        SpanGuard {
207            session: self,
208            id: SpanId(span_id),
209            started,
210        }
211    }
212
213    fn current_step(&self) -> Option<crate::phase::ExecutionStep> {
214        self.inner
215            .borrow()
216            .span_steps
217            .iter()
218            .rev()
219            .find_map(|step| *step)
220    }
221
222    pub(crate) fn end_span(&self, id: SpanId, duration_ns: u64) -> Result<()> {
223        let mut inner = self.inner.borrow_mut();
224        let expected = inner
225            .span_stack
226            .last()
227            .copied()
228            .with_context(|| format!("span stack underflow closing span {}", id.0))?;
229        anyhow::ensure!(
230            expected == id.0,
231            "span_end id `{}` does not match open span `{}`",
232            id.0,
233            expected
234        );
235        inner.span_stack.pop();
236        inner.span_steps.pop();
237
238        format_span_id(&mut inner.id_buf, id.0);
239        let span_id = inner.id_buf.clone();
240        if let Err(error) = write_event(
241            &mut inner.writer,
242            &TraceEvent::SpanEnd(SpanEndEvent {
243                id: span_id,
244                duration_ns,
245            }),
246        ) {
247            inner.sticky_error.get_or_insert_with(|| error.to_string());
248            return Err(error);
249        }
250        Ok(())
251    }
252
253    pub fn elapsed_ns(&self) -> u64 {
254        self.inner
255            .borrow()
256            .probe_started
257            .elapsed()
258            .as_nanos()
259            .min(u64::MAX as u128) as u64
260    }
261
262    /// Record a timed op observation attached to `span_id`.
263    pub fn record_op(&self, span_id: SpanId, op: OpRecord<'_>) -> Result<()> {
264        let storage_bytes = resolve_storage_bytes(op.storage_bytes, op.shape, op.dtype);
265        let timestamp_ns = if op.timestamp_ns > 0 {
266            op.timestamp_ns
267        } else {
268            self.elapsed_ns()
269        };
270        let category = op
271            .category
272            .unwrap_or_else(|| category_for_step(self.current_step(), false));
273        {
274            let mut inner = self.inner.borrow_mut();
275            write_event(
276                &mut inner.writer,
277                &TraceEvent::Op(OpEvent {
278                    span_id: span_id_string(span_id.0),
279                    op_name: op.op_name.into(),
280                    inputs: op.inputs.to_vec(),
281                    output: op.output.map(str::to_string),
282                    shape: op.shape.to_vec(),
283                    dtype: op.dtype.into(),
284                    device: op.device.into(),
285                    duration_ns: op.duration_ns,
286                    timestamp_ns,
287                    storage_bytes: Some(storage_bytes),
288                    input_storage_bytes: op.input_storage_bytes,
289                }),
290            )?;
291        }
292
293        let _ = category;
294        Ok(())
295    }
296
297    /// Record tensor metadata and an allocation event.
298    pub fn record_tensor(&self, span_id: SpanId, tensor: TensorRecord<'_>) -> Result<()> {
299        let storage_bytes = resolve_storage_bytes(tensor.storage_bytes, tensor.shape, tensor.dtype);
300        let mut inner = self.inner.borrow_mut();
301        write_event(
302            &mut inner.writer,
303            &TraceEvent::Tensor(TensorEvent {
304                span_id: span_id_string(span_id.0),
305                tensor_id: tensor.tensor_id.into(),
306                shape: tensor.shape.to_vec(),
307                dtype: tensor.dtype.into(),
308                device: tensor.device.into(),
309                requires_grad: tensor.requires_grad,
310                storage_bytes: Some(storage_bytes),
311                category: tensor.category,
312            }),
313        )?;
314        Ok(())
315    }
316
317    /// Record an explicit tensor allocation (TensorFlow memory timeline).
318    pub fn record_memory_alloc(&self, span_id: SpanId, mem: MemoryRecord<'_>) -> Result<()> {
319        let timestamp_ns = mem.timestamp_ns.unwrap_or_else(|| self.elapsed_ns());
320        let mut inner = self.inner.borrow_mut();
321        write_event(
322            &mut inner.writer,
323            &TraceEvent::Memory(MemoryEvent {
324                timestamp_ns,
325                tensor_id: mem.tensor_id.into(),
326                span_id: span_id_string(span_id.0),
327                op_name: mem.op_name.map(str::to_string),
328                device: mem.device.into(),
329                bytes: mem.bytes,
330                action: MemoryAction::Alloc,
331                shape: mem.shape.to_vec(),
332                dtype: mem.dtype.into(),
333                category: mem.category,
334            }),
335        )
336    }
337
338    /// Record an explicit tensor deallocation.
339    pub fn record_memory_free(&self, span_id: SpanId, mem: MemoryRecord<'_>) -> Result<()> {
340        let timestamp_ns = mem.timestamp_ns.unwrap_or_else(|| self.elapsed_ns());
341        let mut inner = self.inner.borrow_mut();
342        write_event(
343            &mut inner.writer,
344            &TraceEvent::Memory(MemoryEvent {
345                timestamp_ns,
346                tensor_id: mem.tensor_id.into(),
347                span_id: span_id_string(span_id.0),
348                op_name: mem.op_name.map(str::to_string),
349                device: mem.device.into(),
350                bytes: mem.bytes,
351                action: MemoryAction::Free,
352                shape: mem.shape.to_vec(),
353                dtype: mem.dtype.into(),
354                category: mem.category,
355            }),
356        )
357    }
358
359    /// Record a device-level memory checkpoint (cudaMemGetInfo-style).
360    pub fn record_device_memory(
361        &self,
362        device: impl Into<String>,
363        used_bytes: u64,
364        free_bytes: u64,
365        timestamp_ns: Option<u64>,
366    ) -> Result<()> {
367        let timestamp_ns = timestamp_ns.unwrap_or_else(|| self.elapsed_ns());
368        let mut inner = self.inner.borrow_mut();
369        write_event(
370            &mut inner.writer,
371            &TraceEvent::DeviceMemory(DeviceMemoryEvent {
372                timestamp_ns,
373                device: device.into(),
374                used_bytes,
375                free_bytes,
376                reserved_bytes: None,
377            }),
378        )
379    }
380
381    /// Record one parameter gradient fact from a probe run.
382    pub fn record_gradient(
383        &self,
384        root: impl Into<String>,
385        key: impl Into<String>,
386        state: GradientState,
387        norm: Option<f64>,
388    ) -> Result<()> {
389        let mut inner = self.inner.borrow_mut();
390        inner.next_event_id += 1;
391        let event_id = format!("gradient-{}", inner.next_event_id);
392        write_event(
393            &mut inner.writer,
394            &TraceEvent::Gradient(GradientEvent {
395                event_id,
396                root: root.into(),
397                key: key.into(),
398                state,
399                norm,
400            }),
401        )
402    }
403
404    pub fn flush(&self) -> Result<()> {
405        let mut inner = self.inner.borrow_mut();
406        if let Some(error) = &inner.sticky_error {
407            anyhow::bail!("trace session previously failed: {error}");
408        }
409        inner.writer.flush().context("flushing trace JSONL")
410    }
411
412    /// Close the owned root span, flush, and return the trace path.
413    pub fn finish(self) -> Result<PathBuf> {
414        let duration_ns = self.elapsed_ns();
415        {
416            let mut inner = self.inner.borrow_mut();
417            if let Some(error) = &inner.sticky_error {
418                anyhow::bail!("trace session previously failed: {error}");
419            }
420            anyhow::ensure!(
421                inner.span_stack.as_slice() == [1],
422                "cannot finish trace with {} nested spans still open",
423                inner.span_stack.len().saturating_sub(1)
424            );
425            inner.span_stack.pop();
426            inner.span_steps.pop();
427            write_event(
428                &mut inner.writer,
429                &TraceEvent::SpanEnd(SpanEndEvent {
430                    id: span_id_string(1),
431                    duration_ns,
432                }),
433            )?;
434            inner.writer.flush().context("flushing trace JSONL")?;
435        }
436        Ok(self.path)
437    }
438}
439
440fn write_event<W: Write, T: Serialize>(writer: &mut W, event: &T) -> Result<()> {
441    let mut line = serde_json::to_vec(event).context("serializing trace JSONL event")?;
442    line.push(b'\n');
443    writer.write_all(&line).context("writing trace JSONL event")
444}
445
446fn format_span_id(buf: &mut String, id: u64) {
447    buf.clear();
448    use std::fmt::Write as _;
449    let _ = write!(buf, "s{id}");
450}
451
452fn span_id_string(id: u64) -> String {
453    format!("s{id}")
454}
455
456fn new_run_id() -> String {
457    let pid = std::process::id();
458    let nanos = SystemTime::now()
459        .duration_since(UNIX_EPOCH)
460        .map(|d| d.as_nanos())
461        .unwrap_or(0);
462    format!("run-{pid}-{nanos}")
463}
464
465fn utc_iso8601_now() -> String {
466    let now = SystemTime::now()
467        .duration_since(UNIX_EPOCH)
468        .expect("system clock before UNIX epoch");
469    let secs = now.as_secs();
470    let millis = now.subsec_millis();
471    let (year, month, day, hour, minute, second) = unix_secs_to_utc(secs);
472    format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}.{millis:03}Z")
473}
474
475/// Convert Unix seconds to UTC calendar components (Gregorian).
476fn unix_secs_to_utc(secs: u64) -> (u64, u64, u64, u64, u64, u64) {
477    const SECS_PER_DAY: u64 = 86_400;
478    let days = secs / SECS_PER_DAY;
479    let rem = secs % SECS_PER_DAY;
480    let hour = rem / 3600;
481    let minute = (rem % 3600) / 60;
482    let second = rem % 60;
483
484    let z = days + 719_468;
485    let era = z / 146_097;
486    let doe = z - era * 146_097;
487    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
488    let y = yoe + era * 400;
489    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
490    let mp = (5 * doy + 2) / 153;
491    let day = doy - (153 * mp + 2) / 5 + 1;
492    let month = if mp < 10 { mp + 3 } else { mp - 9 };
493    let year = if month <= 2 { y + 1 } else { y };
494    (year, month, day, hour, minute, second)
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500    use crate::trace::document::parse_trace;
501    use crate::trace::events::SpanStartEvent;
502    use serde_json::Value;
503    use std::io::{BufRead, BufReader};
504
505    fn temp_trace(name: &str) -> PathBuf {
506        std::env::temp_dir().join(format!(
507            "candle-graph-trace-{}-{}-{name}",
508            std::process::id(),
509            std::time::SystemTime::now()
510                .duration_since(UNIX_EPOCH)
511                .unwrap()
512                .as_nanos()
513        ))
514    }
515
516    fn span_end_durations(path: &Path) -> Vec<(String, u64)> {
517        let file = std::fs::File::open(path).unwrap();
518        let reader = BufReader::new(file);
519        reader
520            .lines()
521            .map(|line| line.unwrap())
522            .filter_map(|line| {
523                let value: Value = serde_json::from_str(&line).unwrap();
524                if value.get("kind")?.as_str()? != "span_end" {
525                    return None;
526                }
527                Some((
528                    value["id"].as_str().unwrap().to_string(),
529                    value["duration_ns"].as_u64().unwrap(),
530                ))
531            })
532            .collect()
533    }
534
535    fn read_events(path: &Path) -> Vec<TraceEvent> {
536        let file = std::fs::File::open(path).unwrap();
537        let reader = BufReader::new(file);
538        reader
539            .lines()
540            .map(|line| {
541                let line = line.unwrap();
542                serde_json::from_str(&line).unwrap_or_else(|err| {
543                    panic!("invalid JSONL line `{line}`: {err}");
544                })
545            })
546            .collect()
547    }
548
549    #[test]
550    fn nested_spans_emit_parent_hierarchy_and_durations() {
551        let path = temp_trace("nested");
552        let session =
553            TraceSession::open(&path, ProfileRun::training("model::forward", 1, "cpu")).unwrap();
554
555        let inner_id = {
556            let _outer = session.begin_measurement("Model::forward");
557            std::thread::sleep(std::time::Duration::from_micros(50));
558            let inner = session.begin_span("matmul", SpanKind::Op);
559            std::thread::sleep(std::time::Duration::from_micros(50));
560            inner.id
561        };
562
563        session
564            .record_op(
565                inner_id,
566                OpRecord {
567                    op_name: "matmul",
568                    inputs: &["a".into(), "b".into()],
569                    output: Some("c"),
570                    shape: &[8, 8],
571                    dtype: "f32",
572                    device: "cpu",
573                    duration_ns: 1200,
574                    timestamp_ns: 0,
575                    storage_bytes: None,
576                    input_storage_bytes: 0,
577                    category: None,
578                },
579            )
580            .unwrap();
581
582        session.finish().unwrap();
583
584        let events = read_events(&path);
585        assert!(matches!(events.first(), Some(TraceEvent::Meta { .. })));
586
587        let starts: Vec<&SpanStartEvent> = events
588            .iter()
589            .filter_map(|event| match event {
590                TraceEvent::SpanStart(start) => Some(start),
591                _ => None,
592            })
593            .collect();
594        assert_eq!(starts.len(), 3);
595        assert_eq!(starts[0].name, "model::forward");
596        assert!(starts[0].parent_id.is_none());
597        assert_eq!(starts[1].name, "Model::forward");
598        assert_eq!(starts[1].parent_id.as_deref(), Some("s1"));
599        assert_eq!(starts[2].name, "matmul");
600        assert_eq!(starts[2].parent_id.as_deref(), Some("s2"));
601
602        let ends = span_end_durations(&path);
603        assert_eq!(ends.len(), 3);
604        assert!(ends[0].1 > 0);
605
606        let doc = parse_trace(&path).unwrap();
607        assert_eq!(doc.run.entrypoint, "model::forward");
608        assert_eq!(doc.ops.len(), 1);
609        assert_eq!(doc.ops[0].storage_bytes, Some(8 * 8 * 4));
610        assert!(
611            doc.memory.is_empty(),
612            "op metadata must not fabricate tensor lifetime"
613        );
614    }
615
616    #[test]
617    fn record_gradient_round_trips_through_trace_parser() {
618        let path = temp_trace("gradient");
619        let session =
620            TraceSession::open(&path, ProfileRun::training("train::loss", 1, "cpu")).unwrap();
621        session
622            .record_gradient("vb", "encoder.weight", GradientState::Present, Some(0.42))
623            .unwrap();
624        session.finish().unwrap();
625
626        let doc = parse_trace(&path).unwrap();
627        assert_eq!(doc.gradients.len(), 1);
628        assert_eq!(doc.gradients[0].key, "encoder.weight");
629    }
630}
candle_graph/instrument/session.rs

candle_graph/instrument/
session.rs
