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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

use super::cpu::CpuSnapshot;
use super::gpu::GpuSnapshot;
use super::memory::MemorySnapshot;
use super::processes::ProcessesSnapshot;
use super::system::System;

use super::{update_kind::DataUpdateKind, DataSnapshot};

pub struct Collector {
    pub system: System,
}

impl Collector {
    pub fn new() -> Collector {
        let mut system = System::new();
        system.refresh(&DataUpdateKind::all());

        Collector { system }
    }

    pub fn collect(&mut self, kind: &DataUpdateKind) -> DataSnapshot {
        self.system.refresh(kind);

        let cpu = if kind.cpu() {
            Some(CpuSnapshot::from_sysinfo(&self.system.sys))
        } else {
            None
        };

        let memory = if kind.memory() {
            Some(MemorySnapshot::from_sysinfo(&self.system.sys))
        } else {
            None
        };

        let gpu = if kind.gpu() {
            match &self.system.nvml {
                Some(n) => {
                    let result = GpuSnapshot::from_nvml(&n);
                    match result {
                        Ok(g) => Some(g),
                        Err(_) => None,
                    }
                }
                None => None,
            }
        } else {
            None
        };

        let processes = if kind.processes() {
            Some(ProcessesSnapshot::from_sysinfo_nvml(
                &self.system.sys,
                self.system.nvml.as_ref(),
            ))
        } else {
            None
        };
        DataSnapshot {
            cpu,
            memory,
            gpu,
            processes,
        }
    }

    pub fn can_read_gpu(&self) -> bool {
        self.system.nvml.is_some()
    }
}