From 7083a5c9b6eca5a5f4f5cfdfc60a5d91f51c44bc Mon Sep 17 00:00:00 2001 From: jvrobert Date: Fri, 3 Feb 2023 18:34:07 -0700 Subject: [PATCH] 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 --- frigate/util.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/frigate/util.py b/frigate/util.py index 69ead2a7a..02433b358 100755 --- 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",