mltop 0.4.0

Resource monitor for ML engineers written in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

use crate::data::models::cpu::Cpu;
use crate::data::models::gpu::Gpu;
use crate::data::models::memory::Memory;
use crate::data::models::processes::Processes;
use nvml_wrapper::Nvml;
use sysinfo::{CpuRefreshKind, MemoryRefreshKind, Pid, ProcessRefreshKind, RefreshKind, System};

use super::update_kind::DataUpdateKind;

pub struct Collector {
    pub cpu: Cpu,
    pub memory: Memory,
    pub gpu: Option<Gpu>,
    pub processes: Processes,
    sys: System,
    nvml: Option<Nvml>,
}

impl Collector {
    pub fn new() -> Collector {
        let mut sys = System::new();

        Self::refresh_system(&mut sys);

        let mut cpu = Cpu::new();
        cpu.update(&sys);

        let nvml = match Nvml::init() {
            Ok(n) => Some(n),
            Err(_) => None,
        };

        let gpu = match Gpu::new(&nvml) {
            Err(_) => None,
            Ok(mut g) => {
                let _ = g.update(&nvml);
                Some(g)
            }
        };

        let mut processes = Processes::new();
        processes.update(&sys, &nvml);

        Collector {
            cpu,
            memory: Memory::read(&sys),
            gpu,
            processes,
            sys,
            nvml,
        }
    }

    pub fn update(&mut self, kind: &DataUpdateKind) {
        if kind.any() {
            Self::refresh_system(&mut self.sys);

            if kind.cpu() {
                self.cpu.update(&self.sys);
            }
            if kind.memory() {
                self.memory = Memory::read(&self.sys);
            }
            if kind.processes() {
                self.processes.update(&self.sys, &self.nvml);
            }

            if kind.gpu() {
                if let Some(gpu) = self.gpu.as_mut() {
                    let _ = gpu.update(&self.nvml);
                }
            }
        }
    }

    pub fn kill_process(&self, pid: usize) {
        if let Some(process) = self.sys.process(Pid::from(pid)) {
            process.kill();
        }
    }

    fn refresh_system(sys: &mut System) {
        std::thread::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
        let refresh_kind = RefreshKind::new()
            .with_cpu(CpuRefreshKind::new().with_cpu_usage())
            .with_memory(MemoryRefreshKind::everything())
            .with_processes(
                ProcessRefreshKind::everything(), // .with_cpu()
                                                  // .with_memory()
                                                  // .with_cmd(sysinfo::UpdateKind::OnlyIfNotSet),
            );
        sys.refresh_specifics(refresh_kind);
    }

    pub fn has_gpu(&self) -> bool {
        self.gpu.is_some()
    }
}