github-starred/nvitop
2
0
Fork 0
mirror of https://github.com/XuehaiPan/nvitop.git synced 2026-05-21 06:45:24 -06:00
Code Issues 29 Projects Packages Wiki Activity Actions

fix(core/libnvml): fix memory version 2 API arguments

Browse source 
Signed-off-by: Xuehai Pan <XuehaiPan@pku.edu.cn>
This commit is contained in:
Xuehai Pan 2022-10-17 15:43:11 +08:00
parent 50548495b3
commit 6c343a24b6
1 changed files with 23 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

39
nvitop/core/libnvml.py
View file

@ -570,28 +570,35 @@ def nvmlDeviceGetMemoryInfo(handle): # pylint: disable=function-redefined
global __memory_info_v2_available # pylint: disable=global-statement
if __memory_info_v2_available is None:
try:
# pylint: disable-next=unexpected-keyword-arg
retval = _pynvml.nvmlDeviceGetMemoryInfo(handle, version=2)
except TypeError as ex:
if 'unexpected keyword argument' in str(ex).lower():
if not hasattr(_pynvml, 'nvmlMemory_v2'):
with __lock:
__memory_info_v2_available = False
LOGGER.debug(
'NVML constant `nvmlMemory_v2` not found. NVML memory info version 2 is not available.'
)
else:
try:
# pylint: disable-next=unexpected-keyword-arg,no-member
retval = _pynvml.nvmlDeviceGetMemoryInfo(handle, version=_pynvml.nvmlMemory_v2)
except TypeError as ex:
if 'unexpected keyword argument' in str(ex).lower():
with __lock:
__memory_info_v2_available = False
LOGGER.debug('NVML memory info version 2 is not available.')
else:
raise
except (NVMLError_FunctionNotFound, NVMLError_Unknown):
with __lock:
__memory_info_v2_available = False
LOGGER.debug('NVML memory info version 2 is not available.')
else:
raise
except (NVMLError_FunctionNotFound, NVMLError_Unknown):
with __lock:
__memory_info_v2_available = False
LOGGER.debug('NVML memory info version 2 is not available.')
else:
with __lock:
__memory_info_v2_available = True
LOGGER.debug('NVML memory info version 2 is available.')
return retval
with __lock:
__memory_info_v2_available = True
LOGGER.debug('NVML memory info version 2 is available.')
return retval
elif __memory_info_v2_available:
# pylint: disable-next=unexpected-keyword-arg
return _pynvml.nvmlDeviceGetMemoryInfo(handle, version=2)
return _pynvml.nvmlDeviceGetMemoryInfo(handle, version=_pynvml.nvmlMemory_v2)
return _pynvml.nvmlDeviceGetMemoryInfo(handle)