feat(launcher): Pass CUDA_VISIBLE_DEVICES to the shard

cea6051e · OlivierDehaene · 427d7cc4 · cea6051e
Commit cea6051e authored Nov 04, 2022 by OlivierDehaene
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 0 deletions

launcher/src/main.rs launcher/src/main.rs +8 -0

No files found.
--- a/launcher/src/main.rs
+++ b/launcher/src/main.rs
@@ -306,6 +306,14 @@ fn shard_manager(
        ));
    };

+    // If the CUDA_VISIBLE_DEVICES env var is set, pass it to the shard
+    if let Ok(cuda_visible_devices) = env::var("CUDA_VISIBLE_DEVICES") {
+        env.push((
+            "CUDA_VISIBLE_DEVICES".parse().unwrap(),
+            cuda_visible_devices.parse().unwrap(),
+        ));
+    };
+
    // Start process
    tracing::info!("Starting shard {}", rank);
    let mut p = match Popen::create(