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.
This commit is contained in:
Charlie Marsh
2025-11-02 16:21:44 -05:00
committed by GitHub
parent 040521802c
commit 6da135a66a
+22 -3
View File
@@ -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());
}
}
}