From 6da135a66a19c4bbd3b98b3822066850c696397e Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 2 Nov 2025 16:21:44 -0500 Subject: [PATCH] Respect multi-GPU outputs in nvidia-smi (#15460) ## Summary This initially included `NVIDIA_VISIBLE_DEVICES` masking, though it's now omitted for simplicity. Closes https://github.com/astral-sh/uv/issues/14647. --- crates/uv-torch/src/accelerator.rs | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/crates/uv-torch/src/accelerator.rs b/crates/uv-torch/src/accelerator.rs index b96b2b67b..2ed92f63a 100644 --- a/crates/uv-torch/src/accelerator.rs +++ b/crates/uv-torch/src/accelerator.rs @@ -132,9 +132,12 @@ impl Accelerator { .output() { if output.status.success() { - let driver_version = Version::from_str(&String::from_utf8(output.stdout)?)?; - debug!("Detected CUDA driver version from `nvidia-smi`: {driver_version}"); - return Ok(Some(Self::Cuda { driver_version })); + let stdout = String::from_utf8(output.stdout)?; + if let Some(first_line) = stdout.lines().next() { + let driver_version = Version::from_str(first_line.trim())?; + debug!("Detected CUDA driver version from `nvidia-smi`: {driver_version}"); + return Ok(Some(Self::Cuda { driver_version })); + } } debug!( @@ -354,4 +357,20 @@ mod tests { let result = parse_proc_driver_nvidia_version(content).unwrap(); assert_eq!(result, Some(Version::from_str("375.74").unwrap())); } + + #[test] + fn nvidia_smi_multi_gpu() { + // Test that we can parse nvidia-smi output with multiple GPUs (multiple lines) + let single_gpu = "572.60\n"; + if let Some(first_line) = single_gpu.lines().next() { + let version = Version::from_str(first_line.trim()).unwrap(); + assert_eq!(version, Version::from_str("572.60").unwrap()); + } + + let multi_gpu = "572.60\n572.60\n"; + if let Some(first_line) = multi_gpu.lines().next() { + let version = Version::from_str(first_line.trim()).unwrap(); + assert_eq!(version, Version::from_str("572.60").unwrap()); + } + } }