Skip to main content

burn_std/config/
remote.rs

1use cubecl_environment::config::logger::{LogLevel, LoggerConfig};
2
3/// Configuration for the remote backend.
4#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
5pub struct RemoteConfig {
6    /// Logger configuration for remote-backend logs (e.g. the bytes-saved metric of the
7    /// fusion / op-graph caching feature).
8    #[serde(default)]
9    pub logger: LoggerConfig<RemoteLogLevel>,
10
11    /// Flush the outgoing task buffer once this many tasks have accumulated.
12    ///
13    /// Wire-level batching only — every task keeps its own stream/request id, so the server sees
14    /// the same per-task semantics. Larger values batch more aggressively (fewer, bigger frames);
15    /// smaller values cut latency for chains of fire-and-forget submits.
16    #[serde(default = "default_flush_threshold")]
17    pub flush_threshold: usize,
18
19    /// Flush once this many bytes of buffered tensor data accumulate, independent of
20    /// [`flush_threshold`](Self::flush_threshold).
21    ///
22    /// Bounds how much tensor data sits unsent so large uploads go out promptly while small ops keep
23    /// batching. Larger batches fewer/bigger frames; smaller cuts latency for data-heavy streams.
24    #[serde(default = "default_flush_bytes_threshold")]
25    pub flush_bytes_threshold: usize,
26
27    /// Let the Iroh endpoints Burn binds send segmentation-offloaded (GSO) batches.
28    ///
29    /// Off by default: Iroh keeps sending them after the kernel refuses one, which stalls every
30    /// open connection ([iroh#4555](https://github.com/n0-computer/iroh/issues/4555)).
31    #[serde(default)]
32    pub iroh_segmentation_offload: bool,
33}
34
35impl Default for RemoteConfig {
36    fn default() -> Self {
37        Self {
38            logger: LoggerConfig::default(),
39            flush_threshold: default_flush_threshold(),
40            flush_bytes_threshold: default_flush_bytes_threshold(),
41            iroh_segmentation_offload: false,
42        }
43    }
44}
45
46fn default_flush_threshold() -> usize {
47    4
48}
49
50fn default_flush_bytes_threshold() -> usize {
51    1024 * 1024
52}
53
54/// Log levels for remote-backend logging.
55#[derive(
56    Default,
57    Clone,
58    Copy,
59    Debug,
60    PartialEq,
61    Eq,
62    PartialOrd,
63    Ord,
64    serde::Serialize,
65    serde::Deserialize,
66)]
67pub enum RemoteLogLevel {
68    /// Remote logging is disabled.
69    #[default]
70    #[serde(rename = "disabled")]
71    Disabled,
72
73    /// Log periodic summaries — e.g. the cumulative network bytes saved by op-graph caching,
74    /// emitted once per additional mebibyte saved.
75    #[serde(rename = "basic")]
76    Basic,
77
78    /// Log every optimization registration and replay, with per-message sizes.
79    #[serde(rename = "full")]
80    Full,
81}
82
83impl LogLevel for RemoteLogLevel {}