cubecl_runtime/logging/observer.rs
1//! Watching kernel launches from the process that issues them.
2//!
3//! The profiling logger ([`ServerLogger`](super::ServerLogger)) already knows
4//! every kernel that runs, but it formats them into a sink: durations are
5//! aggregated by name into a private table and written out through a detached
6//! task. That is the right shape for reading a log and the wrong one for a
7//! caller that wants to *attribute* the launches — by the time a line is
8//! written, the context that issued it is gone, and nothing about the ordering
9//! is guaranteed against the caller's own state.
10//!
11//! An observer is the other half: a hook called **synchronously, on the thread
12//! that issued the launch, before it is submitted**. That is the only point
13//! where host-side context still exists, so a caller that keeps a stack of what
14//! it is currently doing can pair a kernel with it.
15//!
16//! ```
17//! use std::collections::HashMap;
18//! use std::sync::{Arc, Mutex};
19//!
20//! use cubecl_runtime::logging::{LaunchObservation, LaunchObserver};
21//!
22//! #[derive(Default)]
23//! struct CountThem(Mutex<HashMap<&'static str, usize>>);
24//!
25//! impl LaunchObserver for CountThem {
26//! fn launched(&self, kernel: &'static str) {
27//! *self.0.lock().unwrap().entry(kernel).or_default() += 1;
28//! }
29//! }
30//!
31//! let counts = Arc::new(CountThem::default());
32//! let watching = LaunchObservation::new(counts.clone());
33//! the_pass_to_attribute();
34//! drop(watching);
35//!
36//! for (kernel, count) in counts.0.lock().unwrap().iter() {
37//! println!("{count} × {kernel}");
38//! }
39//! # fn the_pass_to_attribute() {}
40//! ```
41//!
42//! # Cost
43//!
44//! One relaxed atomic load per launch when nothing is installed, which is every
45//! ordinary run. The kernel's name is a `&'static str` the kernel already
46//! carries, so an idle hook allocates and formats nothing. An observer that
47//! asks for timing is the expensive case, and pays per launch →
48//! [`timing`](LaunchObserver::timing).
49//!
50//! # What it reports
51//!
52//! A launch that was **issued**, not one that finished: [`launched`] arrives
53//! before the kernel reaches the server, and a measurement — when one was asked
54//! for — arrives afterwards, either unread through [`profiled`] or as a
55//! duration through [`timed`].
56//!
57//! Issued is not the same as executed. Under a
58//! [`DryRun`](crate::dry_run::DryRun) every launch is still compiled and still
59//! reported here, and is then dropped instead of reaching the device; a
60//! duration measured over one is the compile and the submit, with no kernel
61//! under it. An observer that cares about the difference checks
62//! [`dry_run`](crate::dry_run::dry_run).
63//!
64//! A replayed [`Graph`](crate::client::Graph) is the other direction: its
65//! kernels were observed once, when the capture window recorded them, and a
66//! replay re-executes the whole graph without issuing them again — so an
67//! observed benchmark of a graph-replayed pass reports the capture run and
68//! nothing per replay.
69//!
70//! [`launched`]: LaunchObserver::launched
71//! [`profiled`]: LaunchObserver::profiled
72//! [`timed`]: LaunchObserver::timed
73
74use alloc::sync::Arc;
75use core::sync::atomic::{AtomicBool, Ordering};
76
77/// Re-exported because [`LaunchObserver::timed`]'s signature names them: an
78/// implementor that reached this trait through `cubecl` cannot otherwise spell
79/// its own arguments, and `Duration` is not `core::time::Duration` on every
80/// target.
81pub use cubecl_common::profile::{Duration, ProfileDuration, ProfileTicks, TimingMethod};
82use cubecl_environment::sync::RwLock;
83
84/// What an observer asks be done with each launch's measurement.
85///
86/// Declared once, by [`timing`](LaunchObserver::timing): the launch path has to
87/// know where a measurement will go *before* it takes one.
88///
89/// **It must not change while an observation is installed.** A launch bracketed
90/// under one answer and delivered under another loses a measurement that has
91/// already been paid for. An observer that measures only part of a run keeps
92/// that region in its own state, rather than changing its answer here.
93#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
94pub enum TimingRequest {
95 /// Don't time launches. The default, because timing is not free.
96 #[default]
97 None,
98 /// Time each launch and read its measurement back, delivering the length
99 /// to [`timed`](LaunchObserver::timed).
100 ///
101 /// Reading blocks the issuing thread until the kernel has run, so kernels
102 /// run one at a time and their sum is not the pass's device time. Right for
103 /// each kernel's own cost, wrong for measuring a pipeline.
104 Resolved,
105 /// Time each launch and hand its measurement over **unread**, to
106 /// [`profiled`](LaunchObserver::profiled).
107 ///
108 /// Nothing waits, so the kernels around it keep running back to back and
109 /// the observer reads them once the pass is over. Right for measuring where
110 /// a pass spends its time.
111 ///
112 /// **Only while the profiling logger is off.** A measurement is read once,
113 /// and past [`ExecutionOnly`](super::ProfileLevel::ExecutionOnly) the logger
114 /// needs it too, so the launch path reads it there and delivers a length to
115 /// [`timed`](LaunchObserver::timed) instead — which runs the kernels one at
116 /// a time. An observer whose figures only mean something on a pipelined pass
117 /// should say so when timings arrive that way.
118 Deferred,
119}
120
121/// Notified of every kernel launch, on the thread that issued it.
122///
123/// Implementations must be cheap, must not launch, and must not install or drop
124/// a [`LaunchObservation`]: this runs inside the launch path, before the kernel
125/// reaches the server, and holds the lock that guards the installed observer.
126pub trait LaunchObserver: Send + Sync {
127 /// A kernel was issued, named as the kernel names itself. Pass it through
128 /// [`type_name_format`](crate::config::type_name_format) to shorten it the
129 /// way the profiling logger does.
130 fn launched(&self, kernel: &'static str);
131
132 /// Whether each launch should be timed, and which of
133 /// [`profiled`](Self::profiled) and [`timed`](Self::timed) its measurement
134 /// reaches.
135 ///
136 /// **[`None`](TimingRequest::None) by default, because timing is not free.**
137 /// Bracketing a launch with profile markers costs the issuing thread a round
138 /// trip to the server per kernel, and [`Resolved`](TimingRequest::Resolved)
139 /// also blocks until the kernel has run, removing the overlap between
140 /// kernels. An observer that only wants to know *which* kernels ran should
141 /// leave this alone; one measuring where a pass spends its time wants
142 /// [`Deferred`](TimingRequest::Deferred).
143 ///
144 /// Two situations refuse the measurement without refusing the launch:
145 ///
146 /// * A profile the server cannot take — a graph capture window refuses
147 /// them on the spot. The kernel is still launched, still reported to
148 /// [`launched`](Self::launched), and the measurement is skipped for it,
149 /// with a warning in the log.
150 /// * Don't ask for [`Resolved`](TimingRequest::Resolved) around **collective**
151 /// kernels. Reading blocks until the kernel completes, and a collective
152 /// completes only when its peers launch — a thread that issues more than
153 /// one side of a collective deadlocks waiting for the first.
154 fn timing(&self) -> TimingRequest {
155 TimingRequest::None
156 }
157
158 /// A kernel was timed, and this is its measurement — **not yet read
159 /// back**. Keep it, and read it once the work being measured is over.
160 ///
161 /// Only called under [`TimingRequest::Deferred`], on the thread that issued
162 /// the launch, right after it. Where the backend times on the device
163 /// without waiting — CUDA, HIP, and wgpu's timestamp queries — the
164 /// measurement is two events in the stream that nothing has waited for, and
165 /// the resolved [`ProfileTicks`] carry the window's start and end on one
166 /// clock, so an observer can also say where the device sat idle between
167 /// kernels. Metal waits for the window and places it at the moment it was
168 /// read: its lengths are device time, its starts and ends do not line up,
169 /// and keeping the measurement saves nothing there.
170 ///
171 /// Called under the same lock as every other method here, so keeping the
172 /// measurement must be cheap. Reading it — [`ProfileDuration::resolve`] —
173 /// blocks until the device has reached both events, so doing that here
174 /// would hold the lock for the length of the kernel, which is the thing
175 /// [`Deferred`](TimingRequest::Deferred) exists to avoid.
176 ///
177 /// The default drops it: an observer declaring
178 /// [`Deferred`](TimingRequest::Deferred) owes an implementation.
179 fn profiled(&self, _kernel: &'static str, _profile: ProfileDuration) {}
180
181 /// A kernel finished, and took this long.
182 ///
183 /// Called under [`TimingRequest::Resolved`] once the launch path has read the
184 /// measurement — and under [`TimingRequest::Deferred`] too, in place of
185 /// [`profiled`](Self::profiled), whenever the profiling logger is reading
186 /// measurements as well. **So an observer that asked to keep measurements
187 /// unread still has to implement this**, or it loses every timing whenever
188 /// the logger is on.
189 ///
190 /// It arrives *after* the launch rather than before it, so an observer
191 /// pairing kernels with its own state should do that in
192 /// [`launched`](Self::launched) and use this only for the duration. A
193 /// duration goes to the observer the launch was reported to, and only while
194 /// it is still installed: an observation that ends mid-read is not told.
195 ///
196 /// **`method` is not a detail.** A backend falls back to
197 /// [`System`](TimingMethod::System) where it cannot get a device
198 /// timestamp — wgpu does exactly that once the timestamp-query budget is
199 /// spent — and a system timing is host wall around a blocking submit,
200 /// which includes submission, sync, and the kernel's compilation on its
201 /// first launch, rather than the kernel. The two are not the same
202 /// measurement and an observer reporting them as one will show a number
203 /// that moves several-fold between runs.
204 fn timed(&self, _kernel: &'static str, _duration: Duration, _method: TimingMethod) {}
205}
206
207/// Watches every launch the process issues for as long as it lives, then puts
208/// back whatever it replaced.
209///
210/// Process-wide rather than per client: a caller attributing launches wants
211/// every one its work causes, and work reaches several clients on several
212/// streams. Filtering is the observer's to do, since only it knows what it is
213/// attributing to.
214///
215/// A guard rather than an install/stop pair so the scope being attributed is
216/// the scope the observer is installed for, with no restore step a caller can
217/// skip on an early return. There is one slot, so a second observation replaces
218/// the first for its lifetime; guards dropped in the order they were taken
219/// leave the process as they found it.
220#[must_use = "an observation stops as soon as it is dropped"]
221pub struct LaunchObservation {
222 previous: Option<Arc<dyn LaunchObserver>>,
223}
224
225impl LaunchObservation {
226 /// Installs `observer` until the guard drops.
227 pub fn new(observer: Arc<dyn LaunchObserver>) -> Self {
228 let previous = OBSERVER.write().replace(observer);
229 // Last, so the flag is never set over an empty slot.
230 OBSERVING.store(true, Ordering::Relaxed);
231 Self { previous }
232 }
233}
234
235impl Drop for LaunchObservation {
236 fn drop(&mut self) {
237 let previous = self.previous.take();
238 let still_observed = previous.is_some();
239 *OBSERVER.write() = previous;
240 // Last, so the flag is never cleared while an observer is still
241 // installed: the launches this guard covers are the ones it must not
242 // miss, and a notify that reads the flag in between finds an empty
243 // slot and does nothing.
244 OBSERVING.store(still_observed, Ordering::Relaxed);
245 }
246}
247
248impl core::fmt::Debug for LaunchObservation {
249 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
250 f.debug_struct("LaunchObservation")
251 .field("replaced_an_observer", &self.previous.is_some())
252 .finish()
253 }
254}
255
256/// Whether an observer is actively listening.
257pub(crate) fn is_observing() -> bool {
258 OBSERVING.load(Ordering::Relaxed)
259}
260
261/// Tell the installed observer, if there is one, that `kernel` was issued.
262pub(crate) fn notify_launch(kernel: &'static str) {
263 if !OBSERVING.load(Ordering::Relaxed) {
264 return;
265 }
266 if let Some(observer) = OBSERVER.read().as_ref() {
267 observer.launched(kernel);
268 }
269}
270
271/// Whether the installed observer asked for each launch to be timed.
272pub(crate) fn timing_wanted() -> bool {
273 timing_requested() != TimingRequest::None
274}
275
276/// Say once that the profiling logger is taking the measurements an observer
277/// asked to keep.
278///
279/// A measurement is read once, so with the logger set past `ExecutionOnly` the
280/// launch path reads it and the observer is told a duration instead: its
281/// [`profiled`](LaunchObserver::profiled) never runs and its kernels stop
282/// overlapping, which is a pass other than the one it asked to measure, with
283/// nothing in the numbers to say so.
284///
285/// Once per process, because it is a configuration mistake and not a per-launch
286/// event: at one line per kernel it would be the log.
287pub(crate) fn warn_logger_takes_deferred_measurements() {
288 static SAID: AtomicBool = AtomicBool::new(false);
289 if timing_requested() == TimingRequest::Deferred && !SAID.swap(true, Ordering::Relaxed) {
290 log::warn!(
291 "The profiling logger is reading every launch's measurement, so this run's \
292 launch observer is told durations instead of keeping them, and its kernels \
293 run one at a time. Turn the profile logging off to measure the pass as it runs."
294 );
295 }
296}
297
298/// What the installed observer asked be done with each measurement.
299fn timing_requested() -> TimingRequest {
300 if !OBSERVING.load(Ordering::Relaxed) {
301 return TimingRequest::None;
302 }
303 OBSERVER
304 .read()
305 .as_ref()
306 .map_or(TimingRequest::None, |observer| observer.timing())
307}
308
309/// Read `profile` for the observer, and hand the reading back for the logger.
310///
311/// Both want the same measurement and a measurement is read once, so the
312/// launch path reads it here and passes on what it got. Reading blocks for the
313/// length of the kernel, so it happens with the slot's lock released, and the
314/// duration reaches only the observer that was installed when the read started
315/// — the same guarantee [`notify_profiled`] gives.
316pub(crate) fn read_and_notify_timed(
317 kernel: &'static str,
318 profile: ProfileDuration,
319) -> ProfileDuration {
320 let method = profile.timing_method();
321 let observer = installed_observer();
322 let ticks = cubecl_environment::future::block_on(profile.resolve());
323
324 match (&ticks, &observer) {
325 (Some(ticks), Some(observer)) => deliver_timed(observer, kernel, ticks.duration(), method),
326 // Nothing to report: the window carried no measurement, and a zero
327 // would put a launch that was never timed in the timings.
328 (None, _) => log::warn!(
329 "Skipped timing a launch of `{kernel}` for its observer: \
330 the profiled window carried no measurement"
331 ),
332 (Some(_), None) => {}
333 }
334
335 ProfileDuration::new(alloc::boxed::Box::pin(async move { ticks }), method)
336}
337
338/// The observer installed right now, taken out from under the lock so a read
339/// can outlive holding it.
340fn installed_observer() -> Option<Arc<dyn LaunchObserver>> {
341 if !OBSERVING.load(Ordering::Relaxed) {
342 return None;
343 }
344 OBSERVER.read().as_ref().map(Arc::clone)
345}
346
347/// Report a duration to `observer`, if it is still the installed one.
348///
349/// A duration belongs to the observer the launch was reported to, and reading
350/// a measurement takes as long as the kernel: an observation that ended
351/// underneath the read is done receiving, and the one that replaced it never
352/// saw the launch.
353fn deliver_timed(
354 observer: &Arc<dyn LaunchObserver>,
355 kernel: &'static str,
356 duration: Duration,
357 method: TimingMethod,
358) {
359 let slot = OBSERVER.read();
360 if slot
361 .as_ref()
362 .is_some_and(|installed| Arc::ptr_eq(installed, observer))
363 {
364 observer.timed(kernel, duration, method);
365 }
366}
367
368/// Deliver a launch's measurement the way the observer asked for it.
369///
370/// [`Deferred`](TimingRequest::Deferred) hands it over unread and is done.
371/// [`Resolved`](TimingRequest::Resolved) reads it here with the slot's lock
372/// released, since holding the lock across a read would stall every other
373/// launching thread behind this one's kernel, and delivers through
374/// [`deliver_timed`] like every other duration.
375pub(crate) fn notify_profiled(kernel: &'static str, profile: ProfileDuration) {
376 if !OBSERVING.load(Ordering::Relaxed) {
377 return;
378 }
379 let (observer, profile) = {
380 let slot = OBSERVER.read();
381 let Some(observer) = slot.as_ref() else {
382 return;
383 };
384 match observer.timing() {
385 // Kept: the observer reads it once the work it is measuring is over.
386 TimingRequest::Deferred => {
387 observer.profiled(kernel, profile);
388 return;
389 }
390 // Asked for no timing between the launch being bracketed and this
391 // call — an observation that ended underneath it, or one answering
392 // differently at two moments, which `TimingRequest` forbids. Said
393 // out loud because the measurement has already been paid for, and a
394 // gap in a breakdown with nothing in the log is unattributable.
395 TimingRequest::None => {
396 // Once per process, like the logger's: a `timing()` that
397 // answers differently at two moments does so on every launch,
398 // and at one line each the warning would be the log.
399 static SAID: AtomicBool = AtomicBool::new(false);
400 if !SAID.swap(true, Ordering::Relaxed) {
401 log::warn!(
402 "Dropped a timing of `{kernel}`: its observer asked for none by the \
403 time the measurement arrived"
404 );
405 }
406 return;
407 }
408 TimingRequest::Resolved => (Arc::clone(observer), profile),
409 }
410 };
411
412 let method = profile.timing_method();
413 let Some(ticks) = cubecl_environment::future::block_on(profile.resolve()) else {
414 // Nothing to report: the window carried no measurement, and a zero
415 // would put a launch that was never timed in the timings.
416 log::warn!(
417 "Skipped timing a launch of `{kernel}` for its observer: \
418 the profiled window carried no measurement"
419 );
420 return;
421 };
422
423 deliver_timed(&observer, kernel, ticks.duration(), method);
424}
425
426/// Whether anything is watching. Separate from the observer itself so the
427/// unobserved path — every ordinary run — is one relaxed load rather than a
428/// lock acquisition on the launch path.
429static OBSERVING: AtomicBool = AtomicBool::new(false);
430
431static OBSERVER: RwLock<Option<Arc<dyn LaunchObserver>>> = RwLock::new(None);
432
433#[cfg(test)]
434mod tests {
435 use super::*;
436 use alloc::vec::Vec;
437 // `serial_test`'s macro expands to `vec!`, which a `no_std` crate has to
438 // bring in itself.
439 use alloc::vec;
440 use cubecl_environment::sync::Mutex;
441
442 /// The order launches arrive in, which is what makes attribution possible:
443 /// an observer is called before the launch is submitted, so whatever the
444 /// caller was doing when it issued the kernel is still true.
445 #[test]
446 #[serial_test::serial]
447 fn launches_arrive_in_issue_order() {
448 let recorder = Arc::new(Recorder::default());
449 let watching = LaunchObservation::new(recorder.clone());
450
451 notify_launch("first");
452 notify_launch("second");
453 assert_eq!(*recorder.0.lock(), ["first", "second"]);
454
455 drop(watching);
456 notify_launch("after");
457 assert_eq!(
458 recorder.0.lock().len(),
459 2,
460 "an observation that ended must not keep receiving"
461 );
462 }
463
464 /// Timing is opt-in, and the launch path asks before paying for it: an
465 /// observer that only wants the names must not make every launch blocking.
466 #[test]
467 #[serial_test::serial]
468 fn timing_is_off_unless_an_observer_asks() {
469 assert!(!timing_wanted(), "nothing installed, nothing to time");
470
471 let names_only = LaunchObservation::new(Arc::new(Recorder::default()));
472 assert!(!timing_wanted(), "names only, by default");
473 drop(names_only);
474
475 let timed = Arc::new(Timed::default());
476 let watching = LaunchObservation::new(timed.clone());
477 assert!(timing_wanted());
478 read_and_notify_timed("a_kernel", measured_on_device(7));
479 // The method travels with the duration: a backend that fell back to
480 // the system timer measured host wall around a blocking submit, and an
481 // observer that could not tell would report it as device time.
482 assert_eq!(
483 *timed.0.lock(),
484 [("a_kernel", Duration::from_micros(7), TimingMethod::Device)]
485 );
486
487 drop(watching);
488 assert!(!timing_wanted());
489 }
490
491 /// A nested observation puts back the one it replaced, so a caller that
492 /// watches a sub-pass does not silently take the process's only slot from
493 /// whoever was already watching.
494 #[test]
495 #[serial_test::serial]
496 fn an_observation_restores_the_one_it_replaced() {
497 let outer = Arc::new(Recorder::default());
498 let inner = Arc::new(Recorder::default());
499
500 let watching_outer = LaunchObservation::new(outer.clone());
501 {
502 let _watching_inner = LaunchObservation::new(inner.clone());
503 notify_launch("during_the_inner_pass");
504 }
505 notify_launch("after_the_inner_pass");
506 drop(watching_outer);
507 notify_launch("unobserved");
508
509 assert_eq!(*inner.0.lock(), ["during_the_inner_pass"]);
510 assert_eq!(*outer.0.lock(), ["after_the_inner_pass"]);
511 }
512
513 /// A measurement of `micros`, as a backend without device timestamps
514 /// hands one over: already known, so reading it back waits on nothing.
515 fn measured(micros: u64) -> ProfileDuration {
516 let start = cubecl_common::profile::Instant::now();
517 ProfileDuration::new_system_time(start, start + Duration::from_micros(micros))
518 }
519
520 /// The same, claiming the device timer, to check what an observer is told
521 /// about how a measurement was taken.
522 fn measured_on_device(micros: u64) -> ProfileDuration {
523 let start = cubecl_common::profile::Instant::now();
524 let ticks = ProfileTicks::from_start_end(start, start + Duration::from_micros(micros));
525 ProfileDuration::new(
526 alloc::boxed::Box::pin(async move { Some(ticks) }),
527 TimingMethod::Device,
528 )
529 }
530
531 /// A measurement whose read ends the observation that asked for it, the
532 /// way a pass that finishes while its last kernel is still running does.
533 fn measured_while(
534 micros: u64,
535 during_the_read: impl FnOnce() + Send + 'static,
536 ) -> ProfileDuration {
537 let start = cubecl_common::profile::Instant::now();
538 ProfileDuration::new(
539 alloc::boxed::Box::pin(async move {
540 during_the_read();
541 Some(ProfileTicks::from_start_end(
542 start,
543 start + Duration::from_micros(micros),
544 ))
545 }),
546 TimingMethod::System,
547 )
548 }
549
550 /// An observer asking for [`TimingRequest::Resolved`] is told the
551 /// duration, read back for it — the arm every observer written against
552 /// `timed` alone wants.
553 #[test]
554 #[serial_test::serial]
555 fn a_measurement_is_read_back_for_an_observer_that_only_wants_durations() {
556 let timed = Arc::new(Timed::default());
557 let watching = LaunchObservation::new(timed.clone());
558 notify_profiled("a_kernel", measured(7));
559 drop(watching);
560
561 assert_eq!(
562 *timed.0.lock(),
563 [("a_kernel", Duration::from_micros(7), TimingMethod::System)]
564 );
565 }
566
567 /// An observation dropped while its measurement is being read back is not
568 /// told the duration: the read-back runs outside the lock, so the guard
569 /// can drop mid-read, and the owner has already collected what it wanted.
570 #[test]
571 #[serial_test::serial]
572 fn an_observation_that_ended_mid_read_is_not_told_the_duration() {
573 let timed = Arc::new(Timed::default());
574 let watching = Arc::new(Mutex::new(Some(LaunchObservation::new(timed.clone()))));
575
576 // Stands in for a kernel still running when the owner's pass ends: the
577 // guard drops while the launch path waits on the measurement.
578 let ends_the_observation = watching.clone();
579 notify_profiled(
580 "a_kernel",
581 measured_while(7, move || drop(ends_the_observation.lock().take())),
582 );
583
584 assert!(watching.lock().is_none(), "the read-back ended it");
585 assert!(
586 timed.0.lock().is_empty(),
587 "an observation that ended must not keep receiving"
588 );
589 }
590
591 /// The same guarantee on the path the profiling logger takes, where the
592 /// launch path reads the measurement for both of them: the reading still
593 /// reaches the logger, and no observer is told a launch it never saw.
594 #[test]
595 #[serial_test::serial]
596 fn an_observation_that_ended_mid_read_is_not_told_the_loggers_reading() {
597 let timed = Arc::new(Timed::default());
598 let watching = Arc::new(Mutex::new(Some(LaunchObservation::new(timed.clone()))));
599
600 let ends_the_observation = watching.clone();
601 let for_the_logger = read_and_notify_timed(
602 "a_kernel",
603 measured_while(7, move || drop(ends_the_observation.lock().take())),
604 );
605
606 assert!(watching.lock().is_none(), "the read-back ended it");
607 assert!(
608 timed.0.lock().is_empty(),
609 "an observation that ended must not keep receiving"
610 );
611 let ticks = cubecl_environment::future::block_on(for_the_logger.resolve())
612 .expect("the logger still gets the reading");
613 assert_eq!(
614 ticks.duration(),
615 Duration::from_micros(7),
616 "read once, and handed on"
617 );
618 }
619
620 /// An observer that takes measurements unread is handed each one as it
621 /// was taken, and reads it back when it chooses — which is what keeps the
622 /// kernels around a timed launch running back to back.
623 #[test]
624 #[serial_test::serial]
625 fn an_observer_can_keep_a_measurement_unread() {
626 let kept = Arc::new(Kept::default());
627 let watching = LaunchObservation::new(kept.clone());
628 notify_profiled("first", measured(3));
629 notify_profiled("second", measured(5));
630 drop(watching);
631
632 let read: Vec<(&'static str, Duration)> = core::mem::take(&mut *kept.0.lock())
633 .into_iter()
634 .map(|(kernel, profile)| {
635 let ticks = cubecl_environment::future::block_on(profile.resolve())
636 .expect("a system measurement always carries its ticks");
637 (kernel, ticks.duration())
638 })
639 .collect();
640 assert_eq!(
641 read,
642 [
643 ("first", Duration::from_micros(3)),
644 ("second", Duration::from_micros(5))
645 ],
646 "in issue order, each still carrying its own window"
647 );
648 }
649
650 #[derive(Default)]
651 struct Kept(Mutex<Vec<(&'static str, ProfileDuration)>>);
652
653 impl LaunchObserver for Kept {
654 fn launched(&self, _kernel: &'static str) {}
655 fn timing(&self) -> TimingRequest {
656 TimingRequest::Deferred
657 }
658 fn profiled(&self, kernel: &'static str, profile: ProfileDuration) {
659 self.0.lock().push((kernel, profile));
660 }
661 }
662
663 #[derive(Default)]
664 struct Recorder(Mutex<Vec<&'static str>>);
665
666 impl LaunchObserver for Recorder {
667 fn launched(&self, kernel: &'static str) {
668 self.0.lock().push(kernel);
669 }
670 }
671
672 #[derive(Default)]
673 struct Timed(Mutex<Vec<(&'static str, Duration, TimingMethod)>>);
674
675 impl LaunchObserver for Timed {
676 fn launched(&self, _kernel: &'static str) {}
677 fn timing(&self) -> TimingRequest {
678 TimingRequest::Resolved
679 }
680 fn timed(&self, kernel: &'static str, duration: Duration, method: TimingMethod) {
681 self.0.lock().push((kernel, duration, method));
682 }
683 }
684}