Skip to main content

cubecl_runtime/config/
streaming.rs

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