candle_graph/instrument/
span.rs

1//! RAII span guards for [`super::session::TraceSession`].
2
3use std::time::Instant;
4
5pub use crate::trace::schema::SpanKind;
6
7/// Opaque numeric span identity; formatted only when emitting JSONL.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct SpanId(pub(crate) u64);
10
11/// Fields for one timed op observation inside an open span.
12#[derive(Debug, Clone, Copy)]
13pub struct OpRecord<'a> {
14    pub op_name: &'a str,
15    pub inputs: &'a [String],
16    pub output: Option<&'a str>,
17    pub shape: &'a [usize],
18    pub dtype: &'a str,
19    pub device: &'a str,
20    pub duration_ns: u64,
21    pub timestamp_ns: u64,
22    pub storage_bytes: Option<u64>,
23    pub input_storage_bytes: u64,
24    pub category: Option<crate::trace::MemoryCategory>,
25}
26
27/// Fields for an explicit memory alloc/free event.
28#[derive(Debug, Clone, Copy)]
29pub struct MemoryRecord<'a> {
30    pub tensor_id: &'a str,
31    pub device: &'a str,
32    pub bytes: u64,
33    pub dtype: &'a str,
34    pub category: crate::trace::MemoryCategory,
35    pub timestamp_ns: Option<u64>,
36    pub op_name: Option<&'a str>,
37    pub shape: &'a [usize],
38}
39
40/// Fields for a tensor memory observation inside an open span.
41#[derive(Debug, Clone, Copy)]
42pub struct TensorRecord<'a> {
43    pub tensor_id: &'a str,
44    pub shape: &'a [usize],
45    pub dtype: &'a str,
46    pub device: &'a str,
47    pub requires_grad: bool,
48    pub storage_bytes: Option<u64>,
49    pub category: crate::trace::MemoryCategory,
50}
51
52impl SpanId {
53    pub fn raw(self) -> u64 {
54        self.0
55    }
56}
57
58impl SpanGuard<'_> {
59    pub fn id(&self) -> SpanId {
60        self.id
61    }
62
63    pub fn elapsed_ns(&self) -> u64 {
64        self.started.elapsed().as_nanos().min(u64::MAX as u128) as u64
65    }
66}
67
68/// RAII span: emits `span_start` on creation and `span_end` with wall duration on drop.
69pub struct SpanGuard<'a> {
70    pub(crate) session: &'a super::session::TraceSession,
71    pub(crate) id: SpanId,
72    pub(crate) started: Instant,
73}
74
75impl Drop for SpanGuard<'_> {
76    fn drop(&mut self) {
77        let duration_ns = self.started.elapsed().as_nanos().min(u64::MAX as u128) as u64;
78        let _ = self.session.end_span(self.id, duration_ns);
79    }
80}
candle_graph/instrument/span.rs

candle_graph/instrument/
span.rs
