From 578da6d45a03318a049425778e41b0872d18f161 Mon Sep 17 00:00:00 2001 From: Jongwook Choi Date: Sat, 30 Apr 2022 17:04:14 -0400 Subject: [PATCH] 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 | 22 ++++++++++++++++++---- setup.py | 2 +- 2 files changed, 19 insertions(+), 5 deletions(-) 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 ]