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 output_dense_bytes: Option<u64>,
23    pub input_dense_bytes: u64,
24}
25
26/// Fields for an explicit memory alloc/free event.
27#[derive(Debug, Clone, Copy)]
28pub struct MemoryRecord<'a> {
29    pub storage_id: &'a str,
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/// One already-resolved interval from a device timing adapter.
41#[derive(Debug, Clone, Copy)]
42pub struct DeviceIntervalRecord<'a> {
43    pub device: &'a str,
44    pub stream_id: &'a str,
45    pub clock_id: &'a str,
46    pub backend: &'a str,
47    pub start_ns: u64,
48    pub duration_ns: u64,
49}
50
51/// Independently observed physical-memory checkpoint values.
52#[derive(Debug, Clone, Copy)]
53pub struct DeviceMemoryRecord<'a> {
54    pub device: &'a str,
55    pub used_bytes: Option<u64>,
56    pub free_bytes: Option<u64>,
57    pub reserved_bytes: Option<u64>,
58    pub capacity_bytes: Option<u64>,
59    pub timestamp_ns: Option<u64>,
60}
61
62/// Fields for a tensor memory observation inside an open span.
63#[derive(Debug, Clone, Copy)]
64pub struct TensorRecord<'a> {
65    pub tensor_id: &'a str,
66    pub label: Option<&'a str>,
67    pub shape: &'a [usize],
68    pub dtype: &'a str,
69    pub device: &'a str,
70    pub requires_grad: bool,
71    pub dense_bytes: Option<u64>,
72    pub category: crate::trace::MemoryCategory,
73}
74
75impl SpanId {
76    pub fn raw(self) -> u64 {
77        self.0
78    }
79}
80
81impl SpanGuard<'_> {
82    pub fn id(&self) -> SpanId {
83        self.id
84    }
85
86    pub fn elapsed_ns(&self) -> u64 {
87        self.started.elapsed().as_nanos().min(u64::MAX as u128) as u64
88    }
89}
90
91/// RAII span: emits `span_start` on creation and `span_end` with wall duration on drop.
92pub struct SpanGuard<'a> {
93    pub(crate) session: &'a super::session::TraceSession,
94    pub(crate) id: SpanId,
95    pub(crate) started: Instant,
96}
97
98impl Drop for SpanGuard<'_> {
99    fn drop(&mut self) {
100        let duration_ns = self.started.elapsed().as_nanos().min(u64::MAX as u128) as u64;
101        let _ = self.session.end_span(self.id, duration_ns);
102    }
103}
candle_graph/instrument/span.rs

candle_graph/instrument/
span.rs
