fix(launcher): parse num cuda devices from CUDA_VISIBLE_DEVICES and NVIDIA_VISIBLE_DEVICES

49a6c8c1 · OlivierDehaene · 146e72c3 · 49a6c8c1
Commit 49a6c8c1 authored May 30, 2023 by OlivierDehaene
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 3 deletions

launcher/src/main.rs launcher/src/main.rs +4 -3

No files found.
--- a/launcher/src/main.rs
+++ b/launcher/src/main.rs
@@ -455,9 +455,10 @@ fn shutdown_shards(shutdown: Arc<Mutex<bool>>, shutdown_receiver: &mpsc::Receive
 }
 fn num_cuda_devices() -> Option<usize> {
-    let devices = env::var("CUDA_VISIBLE_DEVICES")
+    let devices = match env::var("CUDA_VISIBLE_DEVICES") {
-        .map_err(|_| env::var("NVIDIA_VISIBLE_DEVICES"))
+        Ok(devices) => devices,
-        .ok()?;
+        Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
+    };
    let n_devices = devices.split(',').count();
    Some(n_devices)
 }