diff --git a/Cargo.toml b/Cargo.toml index 284e0af..74379b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,3 +12,11 @@ nvml-wrapper = {version = "0.9.0", features = ["legacy-functions"]} pretty-bytes = "0.2.2" sysinfo = "0.27.7" users = "0.11.0" + +[[bin]] +name = "clobber" +path = "src/clobber.rs" + +[[bin]] +name = "clobberd" +path = "src/clobberd.rs" diff --git a/src/main.rs b/src/clobber.rs similarity index 100% rename from src/main.rs rename to src/clobber.rs diff --git a/src/clobberd.rs b/src/clobberd.rs new file mode 100644 index 0000000..57fb556 --- /dev/null +++ b/src/clobberd.rs @@ -0,0 +1,87 @@ +use nvml_wrapper::{Nvml,error::NvmlError}; +use users::get_user_by_uid; +use sysinfo::{ProcessRefreshKind, Pid, ProcessExt, System, SystemExt, Signal}; +use std::{thread, time::Duration}; + +pub struct GPUprocess { + name: String, + pid: usize, + device_number: usize, + user: User, +} + +#[derive(Clone, Eq)] +pub struct User { + name: String, + uid: usize, +} + +impl PartialEq for User { + fn eq(&self, other: &User) -> bool { + self.uid == other.uid + } +} + +fn get_processes(nvml: &Nvml, system: &mut System, device_index: u32) -> Result, NvmlError>{ + let mut gpu_processes = vec![]; + let device = nvml.device_by_index(device_index).unwrap(); + let nvml_processes = device.running_compute_processes_v2().unwrap(); + for proc in nvml_processes { + if let Some(process) = system.process(Pid::from(proc.pid as usize)) { + gpu_processes.push(GPUprocess { + name: process.name().to_string(), + pid: proc.pid as usize, + device_number: device_index as usize, + user: match process.user_id() { + Some(user_id) => { + let user = get_user_by_uid(**user_id).unwrap(); + User { + uid: **user_id as usize, + name: user.name().to_string_lossy().to_string() + } + }, + None => User { + uid: 0, + name: "Unknown".to_string() + } + } + }); + } + } + Ok(gpu_processes) +} + +fn kill_process(system: &mut System, pid: usize) -> bool { + system.process(Pid::from(pid)).map(|process| process.kill_with(Signal::Term)).is_some() +} + +fn main() { + let nvml = Nvml::init().unwrap(); + let mut system = System::new_all(); + + let nvml_device_count = nvml.device_count().unwrap(); + let mut locks: Vec> = vec![None; usize::try_from(nvml_device_count).unwrap()]; + + loop { + system.refresh_processes_specifics(ProcessRefreshKind::everything()); + system.refresh_users_list(); + for device_number in 0..nvml_device_count { + let processes = get_processes(&nvml, &mut system, device_number).unwrap_or(vec![]); + if processes.iter().filter(|p| p.user.uid != 0).count() == 0 { + locks[device_number as usize] = None; + } else { + for p in processes.iter() { + match &locks[device_number as usize] { + None => { + locks[device_number as usize] = Some(p.user.clone()); + }, + Some(current) => if current != &p.user { + kill_process(&mut system, p.pid); + } + } + } + } + } + thread::sleep(Duration::from_millis(250)); + } +}