Add noexcept funcitons gpu_count and is_available

Author: XuehaiPan

.gitignore

@@ -12,6 +12,7 @@ __pycache__/
 # Distribution / packaging
 .Python
 env/
+venv/
 build/
 develop-eggs/
 dist/

gpustat/__init__.py

@@ -14,13 +14,17 @@
     ) from ex
 
 from .core import GPUStat, GPUStatCollection
-from .core import new_query
+from .core import new_query, gpu_count, is_available
 from .cli import print_gpustat, main
 
 
 __all__ = (
     '__version__',
-    'GPUStat', 'GPUStatCollection',
+    'GPUStat',
+    'GPUStatCollection',
     'new_query',
-    'print_gpustat', 'main',
+    'gpu_count',
+    'is_available',
+    'print_gpustat',
+    'main',
 )

gpustat/core.py

@@ -740,3 +740,26 @@ def new_query():
     to get the list of GPUs and running process information.
     '''
     return GPUStatCollection.new_query()
+
+
+def gpu_count() -> int:
+    '''
+    Return the number of GPUs in the system.
+    '''
+    try:
+        N.nvmlInit()
+        return N.nvmlDeviceGetCount()
+    except N.NVMLError:
+        return 0  # fallback
+    finally:
+        try:
+            N.nvmlShutdown()
+        except N.NVMLError:
+            pass
+
+
+def is_available() -> bool:
+    '''
+    Return True if the NVML library and GPU devices are available.
+    '''
+    return gpu_count() > 0