Add simple hardware switch functionalty.

gpustat/core.py

@@ -30,9 +30,15 @@
 from blessed import Terminal
 
 from gpustat import util
-from gpustat import rocml as nvml
-from gpustat import rocml as N
-from gpustat.rocml import check_driver_nvml_version
+
+if util.hasNvidia():
+    from gpustat import nvml
+    from gpustat.nvml import nvml as N
+    from gpustat.nvml import check_driver_nvml_version
+else:
+    from gpustat import rocml as nvml
+    from gpustat import rocml as N
+    from gpustat.rocml import check_driver_nvml_version
 
 NOT_SUPPORTED = 'Not Supported'
 MB = 1024 * 1024
@@ -555,7 +561,6 @@ def _wrapped(*args, **kwargs):
                 processes = []
                 nv_comp_processes = nv_comp_processes or []
                 nv_graphics_processes = nv_graphics_processes or []
-                print(nv_comp_processes)
                 # A single process might run in both of graphics and compute mode,
                 # However we will display the process only once
                 seen_pids = set()
@@ -611,10 +616,10 @@ def _wrapped(*args, **kwargs):
                 gpu_stat = GPUStat(gpu_info)
             except Exception as e:
                 gpu_stat = InvalidGPU(index, "((Unknown Error))", e)
-            #except N.NVMLError_Unknown as e:
-            #    gpu_stat = InvalidGPU(index, "((Unknown Error))", e)
-            #except N.NVMLError_GpuIsLost as e:
-            #    gpu_stat = InvalidGPU(index, "((GPU is lost))", e)
+            except N.NVMLError_Unknown as e:
+                gpu_stat = InvalidGPU(index, "((Unknown Error))", e)
+            except N.NVMLError_GpuIsLost as e:
+                gpu_stat = InvalidGPU(index, "((GPU is lost))", e)
 
             if isinstance(gpu_stat, InvalidGPU):
                 log.add_exception("GPU %d" % index, gpu_stat.exception)

gpustat/rocml.py

@@ -1,4 +1,4 @@
-"""Imports pynvml with sanity checks and custom patches."""
+"""Imports pyrsmi and wraps it in a pynvml compatible interface."""
 
 # pylint: disable=protected-access
 
@@ -16,10 +16,19 @@
 
 NVML_TEMPERATURE_GPU = 1
 
+class NVMLError_Unknown(Exception):
+    def __init__(self, message="An unknown ROCMLError has occurred"):
+        self.message = message
+        super().__init__(self.message)
+
+class NVMLError_GpuIsLost(Exception):
+    def __init__(self, message="ROCM Device is lost."):
+        self.message = message
+        super().__init__(self.message)
+
 def nvmlDeviceGetCount():
     return rocml.smi_get_device_count()
 
-
 def nvmlDeviceGetHandleByIndex(dev):
     return dev
 

gpustat/util.py

@@ -2,6 +2,7 @@
 
 import collections
 import os.path
+import subprocess
 import sys
 import traceback
 from typing import Callable, Tuple, Type, TypeVar, Union
@@ -101,3 +102,11 @@ def report_summary(self, concise=True):
                 self._write("{msg} -> Total {value} occurrences.".format(
                     msg=msg, value=value))
             self._write('')
+
+
+def hasNvidia():
+    try:
+        subprocess.check_output('nvidia-smi')
+        return True
+    except Exception:
+        return False