Fix Auto-select GPU (#1474)

* Fix Auto-select GPU * Format file * Add variable in init * Format files * Add small test to ensure environment variable is set * Make linter happy --------- Co-authored-by: roomrys <[email protected]>
talmolab · Sep 6, 2023 · 64655d6 · 64655d6
1 parent 0ef52cd
commit 64655d6
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/sleap/nn/__init__.py b/sleap/nn/__init__.py
@@ -14,3 +14,6 @@
 import sleap.nn.tracking
 import sleap.nn.viz
 import sleap.nn.identity
+import os
+
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
diff --git a/sleap/nn/system.py b/sleap/nn/system.py
@@ -195,6 +195,7 @@ def get_gpu_memory() -> List[int]:
         A list of the available memory on each GPU in MiB.
 
     """
+
     if shutil.which("nvidia-smi") is None:
         return []
 

diff --git a/tests/nn/test_system.py b/tests/nn/test_system.py
@@ -87,3 +87,9 @@ def test_gpu_order_and_length():
 
     # Assert that the order and length of GPU indices match
     assert sleap_indices == nvidia_indices
+
+
+def test_gpu_device_order():
+    """Indirectly tests GPU device order by ensuring environment variable is set."""
+
+    assert os.environ["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"
-Original file line number
+Diff line change
@@ Expand Up / @@ -195,6 +195,7 @@ def get_gpu_memory() -> List[int]: @@
             A list of the available memory on each GPU in MiB.
         """
         if shutil.which("nvidia-smi") is None:
             return []
@@ Expand Down @@