cubecl_runtime/config/
streaming.rs

1use super::logger::{LogLevel, LoggerConfig};
2
3pub use cubecl_environment::stream::StreamPolicy;
4
5/// Configuration for streaming settings in `CubeCL`.
6#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
7pub struct StreamingConfig {
8    /// Logger configuration for streaming logs, using binary log levels.
9    #[serde(default)]
10    pub logger: LoggerConfig<StreamingLogLevel>,
11    /// The maximum number of streams to be used.
12    #[serde(default = "default_max_streams")]
13    pub max_streams: u8,
14    /// Backend stream priority hint.
15    ///
16    /// Backends that expose stream priorities (e.g. CUDA via
17    /// `cuStreamCreateWithPriority`) map this to their native range. Backends
18    /// without a notion of stream priority ignore it. The default is
19    /// [`StreamPriority::Default`], which preserves existing behavior.
20    #[serde(default)]
21    pub priority: StreamPriority,
22    /// How the current stream is derived: `"per-thread"` (default),
23    /// `"per-task"` (stable stream per async task, requires the `tokio`
24    /// feature to take effect) or `"single"`.
25    ///
26    /// A programmatic `cubecl_environment::stream::set_policy` call always
27    /// wins over this setting.
28    #[serde(default)]
29    pub policy: StreamPolicy,
30}
31
32impl Default for StreamingConfig {
33    fn default() -> Self {
34        Self {
35            logger: Default::default(),
36            max_streams: default_max_streams(),
37            priority: StreamPriority::default(),
38            policy: StreamPolicy::default(),
39        }
40    }
41}
42
43fn default_max_streams() -> u8 {
44    128
45}
46
47/// Stream priority hint, mapped to the backend's native priority range.
48///
49/// CUDA convention is that lower numerical priorities run first; this enum is
50/// the portable abstraction so other runtimes can adopt it without changing
51/// their public surface. Backends clamp to their supported range — passing
52/// [`StreamPriority::Low`] on a device whose only priority bucket is the
53/// default is harmless.
54///
55/// The motivating use case is sharing a single GPU with a desktop compositor
56/// (WSL2, dev laptops, single-GPU workstations): running long compute batches
57/// on a low-priority stream lets the compositor preempt cleanly and keeps the
58/// UI responsive.
59#[derive(Default, Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
60pub enum StreamPriority {
61    /// Default backend priority — current behavior on every backend.
62    #[default]
63    #[serde(rename = "default")]
64    Default,
65    /// Lowest priority the backend supports. Useful when sharing the GPU with
66    /// an interactive workload such as a compositor; long-running compute
67    /// batches yield to higher-priority work like UI rendering.
68    #[serde(rename = "low")]
69    Low,
70    /// Highest priority the backend supports. Useful for latency-critical
71    /// foreground work that must not be preempted by background batches.
72    #[serde(rename = "high")]
73    High,
74}
75
76/// Log levels for streaming in `CubeCL`.
77#[derive(Default, Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
78pub enum StreamingLogLevel {
79    /// Compilation logging is disabled.
80    #[default]
81    #[serde(rename = "disabled")]
82    Disabled,
83
84    /// Basic streaming information is logged such as when streams are merged.
85    #[serde(rename = "basic")]
86    Basic,
87
88    /// Full streaming details are logged.
89    #[serde(rename = "full")]
90    Full,
91}
92
93impl LogLevel for StreamingLogLevel {}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn default_priority_is_default_variant() {
101        assert_eq!(StreamingConfig::default().priority, StreamPriority::Default);
102    }
103
104    #[cfg(feature = "std")]
105    #[test]
106    fn priority_omitted_in_toml_falls_back_to_default() {
107        // Existing config files written before this field was added must
108        // continue to deserialize unchanged.
109        let cfg: StreamingConfig = toml::from_str("max_streams = 64").unwrap();
110        assert_eq!(cfg.priority, StreamPriority::Default);
111        assert_eq!(cfg.max_streams, 64);
112    }
113
114    #[test]
115    fn priority_serde_roundtrip() {
116        for p in [
117            StreamPriority::Default,
118            StreamPriority::Low,
119            StreamPriority::High,
120        ] {
121            let s = serde_json::to_string(&p).unwrap();
122            let back: StreamPriority = serde_json::from_str(&s).unwrap();
123            assert_eq!(p, back);
124        }
125    }
126
127    #[test]
128    fn priority_serde_uses_lowercase_names() {
129        // Stable on-disk representation — don't break user configs by accident.
130        assert_eq!(
131            serde_json::to_string(&StreamPriority::Low).unwrap(),
132            "\"low\""
133        );
134        assert_eq!(
135            serde_json::to_string(&StreamPriority::High).unwrap(),
136            "\"high\""
137        );
138        assert_eq!(
139            serde_json::to_string(&StreamPriority::Default).unwrap(),
140            "\"default\""
141        );
142    }
143}
cubecl_runtime/config/streaming.rs

cubecl_runtime/config/
streaming.rs
