From c4bb676a723e2e24d8f52d275404d2b84e5d543a Mon Sep 17 00:00:00 2001 From: Lars T Hansen Date: Wed, 10 Apr 2024 08:26:46 +0200 Subject: [PATCH] Fix #88 - spawn subprograms directly, not via shell --- src/amd.rs | 10 ++++++---- src/command.rs | 19 ++++++++++--------- src/nvidia.rs | 15 +++++++++------ 3 files changed, 25 insertions(+), 19 deletions(-) diff --git a/src/amd.rs b/src/amd.rs index 3de948c..bfaeea8 100644 --- a/src/amd.rs +++ b/src/amd.rs @@ -40,7 +40,7 @@ use crate::util::map; // too small. This is presumably all driver dependent.) pub fn get_amd_configuration() -> Option> { - match command::safe_command("rocm-smi --showproductname", TIMEOUT_SECONDS) { + match command::safe_command("rocm-smi", &["--showproductname"], TIMEOUT_SECONDS) { Ok(raw_text) => { let mut cards = vec![]; for l in raw_text.lines() { @@ -74,9 +74,9 @@ pub fn get_amd_information(user_by_pid: &UserTable) -> Result, // I've not been able to combine the two invocations of rocm-smi yet; we have to run the command // twice. Not a happy situation. - match command::safe_command(AMD_CONCISE_COMMAND, TIMEOUT_SECONDS) { + match command::safe_command(AMD_CONCISE_COMMAND, AMD_CONCISE_ARGS, TIMEOUT_SECONDS) { Ok(concise_raw_text) => { - match command::safe_command(AMD_SHOWPIDGPUS_COMMAND, TIMEOUT_SECONDS) { + match command::safe_command(AMD_SHOWPIDGPUS_COMMAND, AMD_SHOWPIDGPUS_ARGS, TIMEOUT_SECONDS) { Ok(showpidgpus_raw_text) => Ok(extract_amd_information( &concise_raw_text, &showpidgpus_raw_text, @@ -180,6 +180,7 @@ PID 28154 is using 1 DRM device(s): } const AMD_CONCISE_COMMAND: &str = "rocm-smi"; +const AMD_CONCISE_ARGS: &[&str] = &[]; // Return a vector of AMD GPU utilization indexed by device number: (gpu%, mem%) // @@ -238,7 +239,8 @@ GPU Temp (DieEdge) AvgPwr SCLK MCLK Fan Perf PwrCap VRAM% GPU% assert!(xs.eq(&vec![(99.0, 57.0), (63.0, 5.0)])); } -const AMD_SHOWPIDGPUS_COMMAND: &str = "rocm-smi --showpidgpus"; +const AMD_SHOWPIDGPUS_COMMAND: &str = "rocm-smi"; +const AMD_SHOWPIDGPUS_ARGS: &[&str] = &["--showpidgpus"]; // Return a vector of (PID, DEVICES) where DEVICES is a vector of the devices used by the PID. The // PID is a string, the devices are numbers. See test cases below for the various forms diff --git a/src/command.rs b/src/command.rs index 7302f3d..eab1834 100644 --- a/src/command.rs +++ b/src/command.rs @@ -22,8 +22,9 @@ pub enum CmdError { // especially at https://github.com/rust-lang/rust/issues/45572#issuecomment-860134955. See also // https://doc.rust-lang.org/std/process/index.html (second code blob under "Handling I/O"). -pub fn safe_command(command: &str, timeout_seconds: u64) -> Result { - let mut p = match Exec::shell(command) +pub fn safe_command(command: &str, args: &[&str], timeout_seconds: u64) -> Result { + let mut p = match Exec::cmd(command) + .args(args) .stdout(Redirection::Pipe) .stderr(Redirection::Pipe) .popen() @@ -162,35 +163,35 @@ fn format_failure(command: &str, stdout: &str, stderr: &str) -> String { #[test] fn test_safe_command() { // Should work, because we should be running this in the repo root. - match safe_command("ls Cargo.toml", 2) { + match safe_command("ls", &["Cargo.toml"], 2) { Ok(_) => {} Err(_) => assert!(false), } // This really needs to be the output - assert!(safe_command("grep '^name =' Cargo.toml", 2) == Ok("name = \"sonar\"\n".to_string())); + assert!(safe_command("grep", &["^name =", "Cargo.toml"], 2) == Ok("name = \"sonar\"\n".to_string())); // Not found - match safe_command("no-such-command-we-hope", 2) { + match safe_command("no-such-command-we-hope", &[], 2) { Err(CmdError::CouldNotStart(_)) => {} _ => { assert!(false) } } // Wrong permissions, not executable - match safe_command("/etc/passwd", 2) { + match safe_command("/etc/passwd", &[], 2) { Err(CmdError::CouldNotStart(_)) => {} _ => { assert!(false) } } // Should take too long - match safe_command("sleep 7", 2) { + match safe_command("sleep", &["7"], 2) { Err(CmdError::Hung(_)) => {} _ => { assert!(false) } } // Exited with error - match safe_command("ls /abracadabra", 2) { + match safe_command("ls", &["/abracadabra"], 2) { Err(CmdError::Failed(_)) => {} _ => { assert!(false) @@ -198,7 +199,7 @@ fn test_safe_command() { } // Should work even though output is large, but we want a file that is always present. The file // img/sonar.png is 1.2MB, which is probably good enough. - match safe_command("cat img/sonar.png", 5) { + match safe_command("cat", &["img/sonar.png"], 5) { Ok(_) => {} Err(_) => { assert!(false) diff --git a/src/nvidia.rs b/src/nvidia.rs index 736e11e..6066116 100644 --- a/src/nvidia.rs +++ b/src/nvidia.rs @@ -24,7 +24,7 @@ use crate::util::map; // Parsing all the output lines in order yields the information about all the cards. pub fn get_nvidia_configuration() -> Option> { - match command::safe_command("nvidia-smi -a", TIMEOUT_SECONDS) { + match command::safe_command("nvidia-smi", &["-a"], TIMEOUT_SECONDS) { Ok(raw_text) => { let mut cards = vec![]; let mut looking_for_total = false; @@ -77,10 +77,10 @@ pub fn get_nvidia_configuration() -> Option> { // command could not be found, we return Ok(vec![]). pub fn get_nvidia_information(user_by_pid: &UserTable) -> Result, String> { - match command::safe_command(NVIDIA_PMON_COMMAND, TIMEOUT_SECONDS) { + match command::safe_command(NVIDIA_PMON_COMMAND, NVIDIA_PMON_ARGS, TIMEOUT_SECONDS) { Ok(pmon_raw_text) => { let mut processes = parse_pmon_output(&pmon_raw_text, user_by_pid); - match command::safe_command(NVIDIA_QUERY_COMMAND, TIMEOUT_SECONDS) { + match command::safe_command(NVIDIA_QUERY_COMMAND, NVIDIA_QUERY_ARGS, TIMEOUT_SECONDS) { Ok(query_raw_text) => { processes.append(&mut parse_query_output(&query_raw_text, user_by_pid)); Ok(processes) @@ -107,7 +107,8 @@ pub fn get_nvidia_information(user_by_pid: &UserTable) -> Result (device-mask, gpu-util-pct, gpu-mem-pct, gpu-mem-size-in-kib) // where the values are summed across all devices and the device-mask is a bitmask for the @@ -154,8 +155,10 @@ fn parse_pmon_output(raw_text: &str, user_by_pid: &UserTable) -> Vec