reimplementation of gpu_count

nipy · effigies · Mar 19, 2025 · Mar 4, 2025 · Mar 5, 2025 · Mar 7, 2025
commit 4dfbbbee48fb92927dacb1905dc77bb5ed7382e4
diff --git a/nipype/pipeline/plugins/multiproc.py b/nipype/pipeline/plugins/multiproc.py
@@ -21,7 +21,7 @@
 from ...utils.profiler import get_system_total_memory_gb
 from ..engine import MapNode
 from .base import DistributedPluginBase
-from .tools import gpu_count
+from ...utils.gpu_count import gpu_count

 try:
    from textwrap import indent

diff --git a/nipype/pipeline/plugins/tools.py b/nipype/pipeline/plugins/tools.py
@@ -175,13 +175,3 @@ def create_pyscript(node, updatehash=False, store_exception=True):
    with open(pyscript, "w") as fp:
        fp.writelines(cmdstr)
    return pyscript
-
-
-def gpu_count():
-    n_gpus = 1
-    try:
-        import GPUtil
-    except ImportError:
-        return 1
-    else:
-        return len(GPUtil.getGPUs())
diff --git a/nipype/utils/gpu_count.py b/nipype/utils/gpu_count.py
@@ -0,0 +1,55 @@
+# -*- DISCLAIMER: this file contains code derived from gputil (https://github.com/anderskm/gputil)
+# and therefore is distributed under to the following license:
+#
+# MIT License
+#
+# Copyright (c) 2017 anderskm
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import platform
+import shutil
+from subprocess import Popen, PIPE
+import os
+
+
+def gpu_count():
+    try:
+        if platform.system() == "Windows":
+            nvidia_smi = shutil.which('nvidia-smi')
+            if nvidia_smi is None:
+                nvidia_smi = (
+                    "%s\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe"
+                    % os.environ['systemdrive']
+                )
+        else:
+            nvidia_smi = "nvidia-smi"
+
+        p = Popen(
+            [nvidia_smi, "--query-gpu=name", "--format=csv,noheader,nounits"],
+            stdout=PIPE,
+        )
+        stdout, stderror = p.communicate()
+
+        output = stdout.decode('UTF-8')
+        lines = output.split(os.linesep)
+        num_devices = len(lines) - 1
+        return num_devices
+    except:
+        return 0
-from subprocess import Popen, PIPE
-import os
-
-
-def gpu_count():
-    try:
-        if platform.system() == "Windows":
-            nvidia_smi = shutil.which('nvidia-smi')
-            if nvidia_smi is None:
-                nvidia_smi = (
-                    "%s\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe"
-                    % os.environ['systemdrive']
-                )
-        else:
-            nvidia_smi = "nvidia-smi"
-
-        p = Popen(
-            [nvidia_smi, "--query-gpu=name", "--format=csv,noheader,nounits"],
-            stdout=PIPE,
-        )
-        stdout, stderror = p.communicate()
-
-        output = stdout.decode('UTF-8')
-        lines = output.split(os.linesep)
-        num_devices = len(lines) - 1
-        return num_devices
-    except:
-        return 0
+import subprocess
+import os
+
+
+def gpu_count():
+    nvidia_smi = shutil.which('nvidia-smi')
+    if nvidia_smi is None and platform.system() == "Windows":
+        nvidia_smi = f'{os.environ["systemdrive"]}\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'
+
+    if nvidia_smi is None:
+        return 0
+
+    try:
+        p = subprocess.run(
+            [nvidia_smi, "--query-gpu=name", "--format=csv,noheader,nounits"],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+    except (OSError, UnicodeDecodeError):
+        return 0
+
+    return len(output.splitlines())
-from subprocess import Popen, PIPE
-import os
-
-
-def gpu_count():
-    try:
-        if platform.system() == "Windows":
-            nvidia_smi = shutil.which('nvidia-smi')
-            if nvidia_smi is None:
-                nvidia_smi = (
-                    "%s\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe"
-                    % os.environ['systemdrive']
-                )
-        else:
-            nvidia_smi = "nvidia-smi"
-
-        p = Popen(
-            [nvidia_smi, "--query-gpu=name", "--format=csv,noheader,nounits"],
-            stdout=PIPE,
-        )
-        stdout, stderror = p.communicate()
-
-        output = stdout.decode('UTF-8')
-        lines = output.split(os.linesep)
-        num_devices = len(lines) - 1
-        return num_devices
-    except:
-        return 0
+import subprocess
+import os
+
+
+def gpu_count():
+    nvidia_smi = shutil.which('nvidia-smi')
+    if nvidia_smi is None and platform.system() == "Windows":
+        nvidia_smi = f'{os.environ["systemdrive"]}\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'
+
+    if nvidia_smi is None:
+        return 0
+
+    try:
+        p = subprocess.run(
+            [nvidia_smi, "--query-gpu=name", "--format=csv,noheader,nounits"],
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+    except (OSError, UnicodeDecodeError):
+        return 0
+
+    return len(output.splitlines())