utils/gpu: fix nvidia-smi parsing for new nvidia-smi that uses (N/A) when unsupported (#776)

This commit is contained in:
Mathis Felardos 2020-04-02 19:28:21 +02:00 committed by GitHub
parent db14a88fc9
commit 25f43621ed
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,11 +1,21 @@
import subprocess
from collections import namedtuple
from typing import Optional
GPUUsageInfo = namedtuple('GPUUsageInfo', ['total_mem', 'avail_mem', 'used_mem',
'temp', 'percent_fan',
'usage_gpu', 'usage_mem'])
def _convert_nvidia_smi_value(value) -> Optional[int]:
value = value.lower()
# If value contains 'not' or 'N/A' - it is not supported for this GPU
# (in fact, for now nvidia-smi returns '[Not Supported]' or '[N/A]' depending of its version)
if "not" in value or "n/a" in value:
return None
return int(value)
def query_nvidia_smi(gpu_number) -> GPUUsageInfo:
"""
:return:
@ -37,7 +47,6 @@ def query_nvidia_smi(gpu_number) -> GPUUsageInfo:
output = output.decode('utf-8').split("\n")[gpu_number].strip()
values = output.split(", ")
# If value contains 'not' - it is not supported for this GPU (in fact, for now nvidia-smi returns '[Not Supported]')
values = [None if ("not" in value.lower()) else int(value) for value in values]
values = [_convert_nvidia_smi_value(value) for value in values]
return GPUUsageInfo(*values)