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

use nvml_wrapper::{error::NvmlError, Nvml};
use std::vec::IntoIter;
use sysinfo::{Pid, System};

#[derive(Clone)]
pub enum ProcessType {
    // TODO: Add cpu processes as well
    GpuGraphic,
    GpuCompute,
}

impl ToString for ProcessType {
    fn to_string(&self) -> String {
        match self {
            ProcessType::GpuGraphic => "GRAPHIC".to_string(),
            ProcessType::GpuCompute => "COMPUTE".to_string(),
        }
    }
}

#[derive(Clone)]
pub struct Process {
    pub pid: u32,
    pub type_: ProcessType,
    pub command: String,
    // percentage 0-100% (can be higher than 100% if uses more than one core)
    pub cpu_usage: f32,
    // bytes
    pub memory: u64,
    // percentage 0-100%
    pub memory_usage: f32,
}

#[derive(Clone)]
pub struct Processes(pub Vec<Process>);

impl Processes {
    pub fn read(sys: &System, nvml: &Option<Nvml>) -> Result<Processes, NvmlError> {
        if nvml.is_none() {
            return Ok(Processes(Vec::new()));
        }

        let nvml = nvml.as_ref().unwrap();
        let device = nvml.device_by_index(0)?;

        let mut processes = Vec::new();

        let total_memory = sys.total_memory();

        processes.extend(device.running_compute_processes()?.iter().filter_map(|x| {
            match sys.process(Pid::from(x.pid as usize)) {
                None => None,
                Some(p) => Some({
                    let memory = p.memory();
                    Process {
                        pid: x.pid,
                        type_: ProcessType::GpuCompute,
                        command: String::from(p.exe().unwrap().to_str().unwrap()),
                        memory,
                        memory_usage: (memory as f32 / total_memory as f32) * 100.0,
                        cpu_usage: p.cpu_usage(),
                    }
                }),
            }
        }));
        processes.extend(device.running_graphics_processes()?.iter().filter_map(|x| {
            match sys.process(Pid::from(x.pid as usize)) {
                None => None,
                Some(p) => Some({
                    let memory = p.memory();
                    Process {
                        pid: x.pid,
                        type_: ProcessType::GpuGraphic,
                        command: String::from(p.exe().unwrap().to_str().unwrap()),
                        memory,
                        memory_usage: (memory as f32 / total_memory as f32) * 100.0,
                        cpu_usage: p.cpu_usage(),
                    }
                }),
            }
        }));
        Ok(Processes(processes))
    }

    pub fn into_iter(&self) -> IntoIter<Process> {
        self.0.clone().into_iter()
    }
}