mltop 0.5.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

use crate::data::update_kind::DataUpdateKind;
use nvml_wrapper::Nvml;
use sysinfo::{CpuRefreshKind, MemoryRefreshKind, Pid, ProcessRefreshKind, RefreshKind};
use sysinfo::{System as SysinfoSystem, UpdateKind};

pub struct System {
    pub sys: SysinfoSystem,
    pub nvml: Option<Nvml>,
}

impl System {
    pub fn new() -> System {
        let nvml = match Nvml::init() {
            Ok(n) => Some(n),
            Err(_) => None,
        };
        System {
            sys: SysinfoSystem::new(),
            nvml,
        }
    }

    pub fn refresh(&mut self, kind: &DataUpdateKind) {
        std::thread::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);

        let refresh_kind = if kind.processes() {
            RefreshKind::new()
                .with_cpu(CpuRefreshKind::new().with_cpu_usage())
                .with_memory(MemoryRefreshKind::everything())
                .with_processes(
                    ProcessRefreshKind::new()
                        .with_cpu()
                        .with_memory()
                        .with_cmd(UpdateKind::OnlyIfNotSet),
                )
        } else {
            RefreshKind::new()
                .with_cpu(CpuRefreshKind::new().with_cpu_usage())
                .with_memory(MemoryRefreshKind::everything())
        };
        self.sys.refresh_specifics(refresh_kind);
    }

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