fix(api/libnvml): fix removal for process info v3 APIs on the upstrea…

…m 535.98 driver (#89)

CHANGELOG.md

@@ -21,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
--
+- Fix removal for process info v3 APIs on the upstream 535.98 driver by [@XuehaiPan](https://github.com/XuehaiPan) in [#89](https://github.com/XuehaiPan/nvitop/pull/89).
 
 ### Removed
 

nvitop/api/libnvml.py

@@ -594,8 +594,18 @@ def __determine_get_running_processes_version_suffix() -> str:
             # pylint: disable-next=protected-access,no-member
             _nvmlGetFunctionPointer = _pynvml._nvmlGetFunctionPointer
             __get_running_processes_version_suffix = '_v3'
+
+            def lookup(symbol: str) -> _Any:
+                try:
+                    ptr = _nvmlGetFunctionPointer(symbol)
+                except NVMLError_FunctionNotFound:
+                    LOGGER.debug('Failed to found symbol `%s`.', symbol)
+                    raise
+                LOGGER.debug('Found symbol `%s`.', symbol)
+                return ptr
+
             try:
-                _nvmlGetFunctionPointer('nvmlDeviceGetConfComputeMemSizeInfo')
+                lookup('nvmlDeviceGetConfComputeMemSizeInfo')
             except NVMLError_FunctionNotFound:
                 c_nvmlProcessInfo_t = c_nvmlProcessInfo_v2_t
                 LOGGER.debug(
@@ -604,7 +614,7 @@ def __determine_get_running_processes_version_suffix() -> str:
                     'version 3 API with v2 type struct.',
                 )
                 try:
-                    _nvmlGetFunctionPointer('nvmlDeviceGetComputeRunningProcesses_v3')
+                    lookup('nvmlDeviceGetComputeRunningProcesses_v3')
                 except NVMLError_FunctionNotFound:
                     __get_running_processes_version_suffix = '_v2'
                     LOGGER.debug(
@@ -613,7 +623,7 @@ def __determine_get_running_processes_version_suffix() -> str:
                         'process version 2 API with v2 type struct.',
                     )
                     try:
-                        _nvmlGetFunctionPointer('nvmlDeviceGetComputeRunningProcesses_v2')
+                        lookup('nvmlDeviceGetComputeRunningProcesses_v2')
                     except NVMLError_FunctionNotFound:
                         c_nvmlProcessInfo_t = c_nvmlProcessInfo_v1_t
                         __get_running_processes_version_suffix = ''
@@ -632,9 +642,20 @@ def __determine_get_running_processes_version_suffix() -> str:
                         'NVML get running process version 3 API with v2 type struct is available.',
                     )
             else:
-                LOGGER.debug(
-                    'NVML get running process version 3 API with v3 type struct is available.',
-                )
+                try:
+                    lookup('nvmlDeviceGetComputeRunningProcesses_v3')
+                except NVMLError_FunctionNotFound:
+                    c_nvmlProcessInfo_t = c_nvmlProcessInfo_v2_t
+                    __get_running_processes_version_suffix = '_v2'
+                    LOGGER.debug(
+                        'NVML get running process version 3 API with v3 type struct is not '
+                        'available due to incompatible NVIDIA driver. Fallback to use get running '
+                        'process version 2 API with v2 type struct.',
+                    )
+                else:
+                    LOGGER.debug(
+                        'NVML get running process version 3 API with v3 type struct is available.',
+                    )
 
         return __get_running_processes_version_suffix
 
@@ -820,13 +841,15 @@ def __determine_get_memory_info_version_suffix() -> str:
             try:
                 _nvmlGetFunctionPointer('nvmlDeviceGetMemoryInfo_v2')
             except NVMLError_FunctionNotFound:
+                LOGGER.debug('Failed to found symbol `nvmlDeviceGetMemoryInfo_v2`.')
                 c_nvmlMemory_t = c_nvmlMemory_v1_t
                 __get_memory_info_version_suffix = ''
                 LOGGER.debug(
                     'NVML get memory info version 2 API is not available due to incompatible '
                     'NVIDIA driver. Fallback to use NVML get memory info version 1 API.',
                 )
             else:
+                LOGGER.debug('Found symbol `nvmlDeviceGetMemoryInfo_v2`.')
                 LOGGER.debug('NVML get memory info version 2 is available.')
 
         return __get_memory_info_version_suffix