huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-05-23 12:02:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: count gpu uuids if NVIDIA_VISIBLE_DEVICES env set to all (#3230)

Browse Source
This commit is contained in:
drbh 2025-05-16 11:48:58 -04:00 committed by GitHub
parent 18cbecfb38
commit 58934c8b61
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 17 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
launcher/src/main.rs
View File

@ -1263,7 +1263,23 @@ fn num_cuda_devices() -> Option<usize> {
let devices = match env::var("CUDA_VISIBLE_DEVICES") {
Ok(devices) => devices,
Err(_) => match env::var("NVIDIA_VISIBLE_DEVICES") {
Ok(devices) => devices,
Ok(devices) => {
if devices.trim() == "all" {
// Count the number of all GPUs via nvidia-smi
let output = Command::new("nvidia-smi")
.args(["--query-gpu=uuid", "--format=csv,noheader"])
.output()
.ok()?;
String::from_utf8_lossy(&output.stdout)
.lines()
.filter(|line| !line.trim().is_empty())
.count()
.to_string()
} else {
devices
}
}
Err(_) => env::var("ZE_AFFINITY_MASK").ok()?,
},
};