cubecl_runtime/logging/
server.rs

1use core::fmt::Display;
2
3use crate::config::{Logger, compilation::CompilationLogLevel, profiling::ProfilingLogLevel};
4use alloc::format;
5use alloc::string::String;
6use alloc::string::ToString;
7use async_channel::{Receiver, Sender};
8use cubecl_common::future::spawn_detached_fut;
9use cubecl_common::profile::ProfileDuration;
10
11use super::{ProfileLevel, Profiled};
12
13enum LogMessage {
14    Execution(String),
15    Compilation(String),
16    Profile(String, ProfileDuration),
17    ProfileSummary,
18}
19
20/// Server logger.
21#[derive(Debug)]
22pub struct ServerLogger {
23    profile_level: Option<ProfileLevel>,
24    log_compile_info: bool,
25    log_channel: Option<Sender<LogMessage>>,
26}
27
28impl Default for ServerLogger {
29    fn default() -> Self {
30        let logger = Logger::new();
31
32        let disabled = matches!(
33            logger.config.compilation.logger.level,
34            CompilationLogLevel::Disabled
35        ) && matches!(
36            logger.config.profiling.logger.level,
37            ProfilingLogLevel::Disabled
38        );
39
40        if disabled {
41            return Self {
42                profile_level: None,
43                log_compile_info: false,
44                log_channel: None,
45            };
46        }
47        let profile_level = match logger.config.profiling.logger.level {
48            ProfilingLogLevel::Disabled => None,
49            ProfilingLogLevel::Minimal => Some(ProfileLevel::ExecutionOnly),
50            ProfilingLogLevel::Basic => Some(ProfileLevel::Basic),
51            ProfilingLogLevel::Medium => Some(ProfileLevel::Medium),
52            ProfilingLogLevel::Full => Some(ProfileLevel::Full),
53        };
54
55        let log_compile_info = match logger.config.compilation.logger.level {
56            CompilationLogLevel::Disabled => false,
57            CompilationLogLevel::Basic => true,
58            CompilationLogLevel::Full => true,
59        };
60
61        let (send, rec) = async_channel::unbounded();
62
63        // Spawn the logger as a detached task.
64        let async_logger = AsyncLogger {
65            message: rec,
66            logger,
67            profiled: Default::default(),
68        };
69        // Spawn the future in the background to logs messages / durations.
70        spawn_detached_fut(async_logger.process());
71
72        Self {
73            profile_level,
74            log_compile_info,
75            log_channel: Some(send),
76        }
77    }
78}
79
80impl ServerLogger {
81    /// Returns the profile level, none if profiling is deactivated.
82    pub fn profile_level(&self) -> Option<ProfileLevel> {
83        self.profile_level
84    }
85
86    /// Returns true if compilation info should be logged.
87    pub fn compilation_activated(&self) -> bool {
88        self.log_compile_info
89    }
90
91    /// Log the argument to a file when the compilation logger is activated.
92    pub fn log_compilation<I>(&self, arg: &I)
93    where
94        I: Display,
95    {
96        if let Some(channel) = &self.log_channel {
97            if self.log_compile_info {
98                // Channel will never be full, don't care if it's closed.
99                let _ = channel.try_send(LogMessage::Compilation(arg.to_string()));
100            }
101        }
102    }
103
104    /// Register a profiled task without timing.
105    pub fn register_execution(&self, name: impl Display) {
106        if let Some(channel) = &self.log_channel {
107            if matches!(self.profile_level, Some(ProfileLevel::ExecutionOnly)) {
108                // Channel will never be full, don't care if it's closed.
109                let _ = channel.try_send(LogMessage::Execution(name.to_string()));
110            }
111        }
112    }
113
114    /// Register a profiled task.
115    pub fn register_profiled(&self, name: impl Display, duration: ProfileDuration) {
116        if let Some(channel) = &self.log_channel {
117            if self.profile_level.is_some() {
118                // Channel will never be full, don't care if it's closed.
119                let _ = channel.try_send(LogMessage::Profile(name.to_string(), duration));
120            }
121        }
122    }
123
124    /// Show the profiling summary if activated and reset its state.
125    pub fn profile_summary(&self) {
126        if let Some(channel) = &self.log_channel {
127            if self.profile_level.is_some() {
128                // Channel will never be full, don't care if it's closed.
129                let _ = channel.try_send(LogMessage::ProfileSummary);
130            }
131        }
132    }
133}
134
135struct AsyncLogger {
136    message: Receiver<LogMessage>,
137    logger: Logger,
138    profiled: Profiled,
139}
140
141impl AsyncLogger {
142    async fn process(mut self) {
143        while let Ok(msg) = self.message.recv().await {
144            match msg {
145                LogMessage::Compilation(msg) => {
146                    self.logger.log_compilation(&msg);
147                }
148                LogMessage::Profile(name, profile) => {
149                    let duration = profile.resolve().await.duration();
150                    self.profiled.update(&name, duration);
151                    self.logger
152                        .log_profiling(&format!("| {duration:<10?} | {name}"));
153                }
154                LogMessage::Execution(name) => {
155                    self.logger.log_profiling(&format!("Executing {name}"));
156                }
157                LogMessage::ProfileSummary => {
158                    if !self.profiled.is_empty() {
159                        self.logger.log_profiling(&self.profiled);
160                        self.profiled = Profiled::default();
161                    }
162                }
163            }
164        }
165    }
166}
cubecl_runtime/logging/server.rs

cubecl_runtime/logging/
server.rs
