burn_train/metric/
cpu_use.rs

1use super::{MetricMetadata, Numeric};
2use crate::metric::{Metric, MetricEntry, MetricName, NumericEntry};
3use std::{
4    sync::Arc,
5    time::{Duration, Instant},
6};
7use sysinfo::{CpuRefreshKind, RefreshKind, System};
8
9/// General CPU Usage metric
10pub struct CpuUse {
11    name: MetricName,
12    last_refresh: Instant,
13    refresh_frequency: Duration,
14    sys: System,
15    current: f64,
16}
17
18impl Clone for CpuUse {
19    fn clone(&self) -> Self {
20        Self {
21            name: self.name.clone(),
22            last_refresh: self.last_refresh,
23            refresh_frequency: self.refresh_frequency,
24            sys: System::new(),
25            current: self.current,
26        }
27    }
28}
29
30impl CpuUse {
31    /// Creates a new CPU metric
32    pub fn new() -> Self {
33        let mut sys = System::new();
34        let current = Self::refresh(&mut sys);
35        let name = "CPU Usage".to_string();
36
37        Self {
38            name: Arc::new(name),
39            last_refresh: Instant::now(),
40            refresh_frequency: Duration::from_millis(200),
41            sys,
42            current,
43        }
44    }
45
46    fn refresh(sys: &mut System) -> f64 {
47        sys.refresh_specifics(
48            RefreshKind::nothing().with_cpu(CpuRefreshKind::nothing().with_cpu_usage()),
49        );
50
51        let cpus = sys.cpus();
52        let num_cpus = cpus.len();
53        let use_percentage = cpus.iter().fold(0.0, |acc, cpu| acc + cpu.cpu_usage()) as f64;
54
55        use_percentage / num_cpus as f64
56    }
57}
58
59impl Default for CpuUse {
60    fn default() -> Self {
61        CpuUse::new()
62    }
63}
64
65impl Metric for CpuUse {
66    type Input = ();
67
68    fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> MetricEntry {
69        if self.last_refresh.elapsed() >= self.refresh_frequency {
70            self.current = Self::refresh(&mut self.sys);
71            self.last_refresh = Instant::now();
72        }
73
74        let formatted = format!("{}: {:.2} %", self.name(), self.current);
75        let raw = format!("{:.2}", self.current);
76
77        MetricEntry::new(self.name(), formatted, raw)
78    }
79
80    fn clear(&mut self) {}
81
82    fn name(&self) -> MetricName {
83        self.name.clone()
84    }
85}
86
87impl Numeric for CpuUse {
88    fn value(&self) -> NumericEntry {
89        NumericEntry::Value(self.current)
90    }
91}