Pin nvidia-ml-py version to avoid API version issues

Requires nvidia-ml-py >= 11.450.129 and nvidia drivers >= 450.66.

We also pin the version of nvidia-ml-py < 11.510 against a breaking
change due to v3 APIs (from 510.39.01). The latest compatible version
would be nvidia-ml-py==11.495.46. See #107 for more details.

We no longer need to exclude 375.* versions as the wrong package
has been expelled from the pypi repository.

gpustat/nvml.py

@@ -2,21 +2,35 @@
 
 import textwrap
 
+
 try:
+    # Check pynvml version: we require 11.450.129 or newer.
+    # https://github.com/wookayin/gpustat/pull/107
     import pynvml
-    if not hasattr(pynvml, 'NVML_BRAND_TITAN_RTX'):
+    if not (
+        # Requires nvidia-ml-py >= 11.460.79
+        hasattr(pynvml, 'NVML_BRAND_NVIDIA_RTX') or
+        # Requires nvidia-ml-py >= 11.450.129, < 11.510.69
+        hasattr(pynvml, 'nvmlDeviceGetComputeRunningProcesses_v2')
+    ):
         raise RuntimeError("pynvml library is outdated.")
 except (ImportError, SyntaxError, RuntimeError) as e:
     raise ImportError(textwrap.dedent(
         """\
         pynvml is missing or an outdated version is installed.
 
+        We require nvidia-ml-py>=11.450.129; see GH-107 for more details.
+
+        -----------------------------------------------------------
         Please reinstall `gpustat`:
-        $ pip install -I gpustat
 
-        or manually fix the package:
+        $ pip install --force-reinstall gpustat
+
+        if it still does not fix the problem, manually fix nvidia-ml-py
+        installation:
+
         $ pip uninstall nvidia-ml-py3
-        $ pip install -I 'nvidia-ml-py!=375.*'
+        $ pip install --force-reinstall 'nvidia-ml-py<11.515'
         """)) from e
 
 

setup.py

@@ -77,7 +77,7 @@ def run(self):
 
 install_requires = [
     'six>=1.7',
-    'nvidia-ml-py!=375.*',  # see https://forums.developer.nvidia.com/t/mistake-with-nvidia-ml-py-on-pypi/112942
+    'nvidia-ml-py>=11.450.129,<=11.495.46',  # see #107
     'psutil>=5.6.0',    # GH-1447
     'blessed>=1.17.1',  # GH-126
 ]