cubecl_environment/records/
span.rs

1#[cfg(not(native_cache))]
2use super::disabled as session;
3#[cfg(native_cache)]
4use super::session;
5use super::{Record, RecordEffect, Stamp, base::write_stamped};
6use alloc::string::String;
7use core::time::Duration;
8use serde::{Deserialize, Serialize};
9
10/// A stretch of the session in progress that a record describes: stamped when
11/// it opens, timed on the session's clock, and recorded when it closes, so
12/// everything stamped meanwhile lies inside it.
13#[must_use = "a span records nothing until it is closed"]
14#[derive(Debug)]
15pub struct Span {
16    start: Stamp,
17}
18
19impl Span {
20    /// Opens a span now, starting the session if there is none. `None` when
21    /// recording is off.
22    pub fn new() -> Option<Self> {
23        session::stamp().map(|start| Self { start })
24    }
25
26    /// The time since the span opened, on the session's clock. `None` when
27    /// the environment switched sessions since: the span went with its own.
28    pub fn elapsed(&self) -> Option<Duration> {
29        session::offset(self.start.session).map(|now| now.saturating_sub(self.start.offset))
30    }
31
32    /// Records `record`, stamped when the span opened: written if the session
33    /// is kept, as [`write()`](super::write) is. `false` when the session the
34    /// span opened in is gone, or the write failed.
35    pub fn close<R: Record + Serialize>(self, effect: RecordEffect, record: &R) -> bool {
36        write_stamped(self.start, effect, record)
37    }
38}
39
40/// A named span of the session in progress — a phase of the caller's work,
41/// `load weights`, `plan turn 7` — recorded as a [`MarkRecord`] when it
42/// drops. Everything else stamped meanwhile lies inside it on the session's
43/// clock, which is what lets a reader say what a phase cost.
44#[must_use = "a mark records the span until it is dropped"]
45#[derive(Debug)]
46pub struct Mark {
47    label: String,
48    span: Option<Span>,
49}
50
51impl Mark {
52    /// Opens the span now. The label is converted only when the environment
53    /// records, so a `&str` on a path that records nothing is never copied.
54    pub fn new<S: Into<String>>(label: S) -> Self {
55        let span = Span::new();
56        let label = match span {
57            Some(_) => label.into(),
58            None => String::new(),
59        };
60        Self { label, span }
61    }
62}
63
64impl Drop for Mark {
65    fn drop(&mut self) {
66        let Some(span) = self.span.take() else {
67            return;
68        };
69        let Some(wall) = span.elapsed() else {
70            return;
71        };
72        let record = MarkRecord {
73            label: core::mem::take(&mut self.label),
74            wall,
75        };
76        span.close(RecordEffect::Observed, &record);
77    }
78}
79
80/// A named span of a session, as a [`Mark`] records it.
81#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
82pub struct MarkRecord {
83    /// What the span was.
84    pub label: String,
85    /// How long it lasted.
86    pub wall: Duration,
87}
88
89impl Record for MarkRecord {
90    const KIND: &'static str = "marks";
91}
cubecl_environment/records/span.rs

cubecl_environment/records/
span.rs
