utils/gpu: fix nvidia-smi parsing for new nvidia-smi that uses (N/A) when unsupported (#776)
This commit is contained in:
parent
db14a88fc9
commit
25f43621ed
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user