Try to limit nvidia GPU queries to included GPUs (#5356)

* Try to limit nvidia GPU queries to included GPUs * ignore non digit GPU indexes * formatting * Formatting * Remove trailing spaces --------- Co-authored-by: Nicolas Mowen <[email protected]>
blakeblackshear · Feb 4, 2023 · 7083a5c · 7083a5c
1 parent db131d4
commit 7083a5c
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/frigate/util.py b/frigate/util.py
@@ -926,6 +926,17 @@ def get_nvidia_gpu_stats() -> dict[str, str]:
         "--format=csv",
     ]
 
+    if (
+        "CUDA_VISIBLE_DEVICES" in os.environ
+        and os.environ["CUDA_VISIBLE_DEVICES"].isdigit()
+    ):
+        nvidia_smi_command.extend(["--id", os.environ["CUDA_VISIBLE_DEVICES"]])
+    elif (
+        "NVIDIA_VISIBLE_DEVICES" in os.environ
+        and os.environ["NVIDIA_VISIBLE_DEVICES"].isdigit()
+    ):
+        nvidia_smi_command.extend(["--id", os.environ["NVIDIA_VISIBLE_DEVICES"]])
+
     p = sp.run(
         nvidia_smi_command,
         encoding="ascii",