diff --git a/gpustat/nvml.py b/gpustat/nvml.py index 5d3db2f..4792ffc 100644 --- a/gpustat/nvml.py +++ b/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 diff --git a/setup.py b/setup.py index 6b7529c..9031f36 100644 --- a/setup.py +++ b/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 ]