cubecl_environment/records/
base.rs

1#[cfg(not(native_cache))]
2use super::disabled as session;
3#[cfg(native_cache)]
4use super::session;
5use alloc::string::String;
6use core::time::Duration;
7use serde::{Deserialize, Serialize};
8
9/// The namespace root of every record, versioned so a reader selects the
10/// layout it knows and an export can drop them all by one prefix.
11pub const ROOT: &str = "records/v1";
12
13/// The namespace sessions are written to.
14pub const SESSIONS: &str = "records/v1/sessions";
15
16/// How much an environment records.
17#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19#[repr(u8)]
20pub enum RecordLevel {
21    /// Nothing.
22    Off = 0,
23    /// Every record, without the heavy parts: what a build cost, not the
24    /// artifacts it produced.
25    #[default]
26    Basic = 1,
27    /// Every record in full: a compiled kernel's IR and source.
28    Full = 2,
29}
30
31/// How much an environment records of its own build, and how long it keeps
32/// it: what [`configure`] takes, and the `[environment.records]` table of the
33/// runtime's configuration.
34#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
35pub struct RecordsConfig {
36    /// `basic` by default: every record, without the heavy parts.
37    #[serde(default)]
38    pub level: RecordLevel,
39
40    /// How many sessions an environment keeps; the oldest beyond it are
41    /// pruned with their records when a session is kept, at its first change
42    /// to the environment. A session that changes nothing prunes nothing, and
43    /// the session being kept always survives, so `0` keeps it alone, as `1`
44    /// does. Every session is kept when unset.
45    #[serde(default)]
46    pub keep_sessions: Option<u32>,
47}
48
49/// Whether a record accompanies a change to the environment.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub enum RecordEffect {
52    /// The environment's caches changed: a key was tuned, a kernel compiled.
53    /// The session is kept, with everything it recorded before.
54    Changed,
55    /// Something happened that changed nothing: a kernel loaded from the
56    /// store, a span marked, a snapshot taken. Kept only if the session
57    /// changes something.
58    Observed,
59}
60
61/// A type written to the environment as a record: its kind names the
62/// namespace, `records/v1/<KIND>`, it is written to and read back from.
63pub trait Record {
64    /// The kind records of this type are written under. Any but `sessions`,
65    /// which names the [namespace sessions are written to](SESSIONS).
66    const KIND: &'static str;
67
68    /// The namespace records of this type are written to.
69    fn namespace() -> String {
70        alloc::format!("{ROOT}/{}", Self::KIND.trim_matches('/'))
71    }
72}
73
74/// Names a [`Session`]: unique among the sessions of an environment. Taken
75/// from the session's start in nanoseconds, so ids sort sessions by start.
76#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
77#[serde(transparent)]
78pub struct SessionId(pub u64);
79
80impl core::fmt::Display for SessionId {
81    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
82        write!(f, "{}", self.0)
83    }
84}
85
86/// One process's use of one environment: what every record of it shares.
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
88pub struct Session {
89    /// Unique among the sessions of an environment.
90    pub id: SessionId,
91    /// Wall-clock start, in milliseconds since the Unix epoch.
92    pub started_unix_ms: u64,
93    /// The cubecl that wrote it.
94    pub cubecl_version: String,
95    /// What the process said it was doing, e.g. `models build qwen3-8b`.
96    pub label: Option<String>,
97    /// The operating system's process id.
98    pub process: u32,
99    /// The operating system, e.g. `linux`.
100    pub os: String,
101    /// The CPU architecture, e.g. `x86_64`.
102    pub arch: String,
103}
104
105/// Where a record sits in its session.
106#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
107pub struct Stamp {
108    /// The [`Session::id`] of the session the record belongs to.
109    pub session: SessionId,
110    /// The record's position among its session's, across every namespace.
111    pub seq: u64,
112    /// Time since the session started, on a monotonic clock.
113    pub offset: Duration,
114}
115
116/// A record as it is stored: its stamp, then its body. The key is the stamp's
117/// `(session, seq)`, which is unique and sorts a session's records in order.
118#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct Stamped<V> {
120    /// Where the record sits in its session.
121    pub stamp: Stamp,
122    /// The record itself.
123    pub record: V,
124}
125
126/// Sets how much is recorded, and how many sessions an environment keeps.
127/// Called by the runtime once its configuration is loaded.
128pub fn configure(config: RecordsConfig) {
129    session::configure(config);
130}
131
132/// The level in force.
133pub fn level() -> RecordLevel {
134    session::level()
135}
136
137/// Whether anything is recorded.
138pub fn enabled() -> bool {
139    level() != RecordLevel::Off
140}
141
142/// Names what this process is doing, for the session in progress and every
143/// later one: `models build qwen3-8b`, `serve`.
144pub fn label<S: Into<String>>(label: S) {
145    session::label(label.into());
146}
147
148/// Records `record` under its [namespace](Record::namespace) in the active
149/// environment, stamped now: written if the session is kept, which an
150/// [`Observed`](RecordEffect::Observed) record alone does not decide. `None`
151/// when recording is off or there is no database to write to.
152pub fn write<R: Record + Serialize>(effect: RecordEffect, record: &R) -> Option<Stamp> {
153    let stamp = session::stamp()?;
154    write_stamped(stamp, effect, record).then_some(stamp)
155}
156
157/// Records `record` with a stamp taken earlier. `false` when the session it
158/// was stamped in is gone, or the write failed.
159pub(crate) fn write_stamped<R: Record + Serialize>(
160    stamp: Stamp,
161    effect: RecordEffect,
162    record: &R,
163) -> bool {
164    let namespace = R::namespace();
165    // A record under the sessions' namespace would be read as a session, and
166    // skipped by every prune.
167    debug_assert_ne!(namespace, SESSIONS, "`sessions` is not a record kind");
168    session::write(&namespace, effect, &Stamped { stamp, record })
169}
cubecl_environment/records/base.rs

cubecl_environment/records/
base.rs
