furiosa_device/
proc.rs

1use std::fs;
2use std::io::prelude::*;
3
4use lazy_static::lazy_static;
5use memoize::memoize;
6use rayon::prelude::*;
7use regex::Regex;
8
9use crate::DeviceResult;
10
11lazy_static! {
12    static ref DEVICE_PATH_PATTERN: Regex = Regex::new(
13        r"^/dev/npu(?P<device_id>\d+)((?:pe)(?P<start_core>\d+)(-(?P<end_core>\d+))?)?$"
14    )
15    .unwrap();
16}
17
18pub struct NpuProcess {
19    pub dev_name: String,
20    pub pid: u32,
21    pub cmdline: String,
22}
23
24impl NpuProcess {
25    pub(crate) fn new(dev_name: String, pid: u32, cmdline: String) -> Self {
26        Self {
27            dev_name,
28            pid,
29            cmdline,
30        }
31    }
32
33    pub fn dev_name(&self) -> &str {
34        self.dev_name.as_str()
35    }
36
37    pub fn pid(&self) -> u32 {
38        self.pid
39    }
40
41    pub fn cmdline(&self) -> &str {
42        self.cmdline.as_str()
43    }
44}
45
46pub fn scan_processes() -> DeviceResult<Vec<NpuProcess>> {
47    let mut targets = Vec::new();
48
49    for entry in fs::read_dir("/proc")? {
50        if let Ok(pid) = entry?.file_name().to_string_lossy().parse::<u32>() {
51            let path = format!("/proc/{pid}/fd");
52            if let Ok(dirs) = fs::read_dir(&path) {
53                for entry in dirs {
54                    let entry = entry?;
55                    targets.push((
56                        pid,
57                        format!(
58                            "{}/{}",
59                            path,
60                            entry.file_name().as_os_str().to_string_lossy()
61                        ),
62                    ));
63                }
64            }
65        }
66    }
67
68    let mut results: Vec<NpuProcess> = targets
69        .into_par_iter()
70        .filter_map(|(pid, path)| {
71            if let Ok(link) = fs::read_link(path) {
72                let link = link.as_os_str().to_string_lossy();
73                if DEVICE_PATH_PATTERN.is_match(&link) {
74                    return Some((pid, link.replace("/dev/", "")));
75                }
76            }
77            None
78        })
79        .filter_map(|(pid, dev_path)| {
80            read_cmdline(pid).map(|cmdline| NpuProcess::new(dev_path, pid, cmdline))
81        })
82        .collect();
83
84    results.sort_by(|a, b| a.dev_name.cmp(&b.dev_name));
85
86    Ok(results)
87}
88
89#[memoize(TimeToLive: std::time::Duration::from_secs(1))]
90fn read_cmdline(pid: u32) -> Option<String> {
91    let path = format!("/proc/{pid}/cmdline");
92    let file = std::fs::File::open(path).ok()?;
93    let mut buf_reader = std::io::BufReader::new(file);
94    let mut contents = String::new();
95    buf_reader.read_to_string(&mut contents).ok()?;
96
97    Some(contents.replace('\0', " ").trim().to_string())
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn read_cmdline_test() {
106        let my_pid = std::process::id();
107        let res = read_cmdline(my_pid);
108        assert!(res.is_some());
109
110        let res = res.unwrap();
111        assert!(res.contains("/target/"));
112        // when "cargo test": ".../device-api/target/debug/deps/furiosa_device-5a7434828b2c179b"
113        // when module unit test: ".../device-api/target/debug/deps/furiosa_device-5a7434828b2c179b proc::tests --nocapture"
114        // when function unit test: ".../device-api/target/debug/deps/furiosa_device-96a8035dd957f8e4 proc::tests::read_cmdline_test --exact --nocapture"
115    }
116}
furiosa_device/proc.rs

furiosa_device/
proc.rs
