Only need first line of --query-gpu=count

common-workflow-language · Oct 16, 2023 · 239ce5a · 239ce5a
1 parent d3f89d3
commit 239ce5a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cwltool/cuda.py b/cwltool/cuda.py
@@ -19,7 +19,8 @@ def cuda_device_count() -> str:
     except Exception as e:
         _logger.warning("Error checking number of GPUs with nvidia-smi: %s", e)
         return "0"
-    return proc.stdout.decode("utf-8")
+    # NOTE: On a machine with N GPUs the query return N lines, each containing N.
+    return proc.stdout.decode("utf-8").split("\n")[0]
 
 
 def cuda_version_and_device_count() -> Tuple[str, int]: