Skip to main content

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}