fix unittest process crash

where the tests for tests?..
This commit is contained in:
Sergey Krashevich 2023-05-03 14:51:50 +03:00
parent e33c87e9fe
commit b56265e377
No known key found for this signature in database
GPG Key ID: 625171324E7D3856

View File

@ -925,30 +925,32 @@ def try_get_info(f, h, default="N/A"):
def get_nvidia_gpu_stats() -> dict[int, dict]: def get_nvidia_gpu_stats() -> dict[int, dict]:
nvml.nvmlInit()
deviceCount = nvml.nvmlDeviceGetCount()
results = {} results = {}
for i in range(deviceCount): try:
handle = nvml.nvmlDeviceGetHandleByIndex(i) nvml.nvmlInit()
meminfo = try_get_info(nvml.nvmlDeviceGetMemoryInfo, handle) deviceCount = nvml.nvmlDeviceGetCount()
util = try_get_info(nvml.nvmlDeviceGetUtilizationRates, handle) for i in range(deviceCount):
if util != "N/A": handle = nvml.nvmlDeviceGetHandleByIndex(i)
gpu_util = util.gpu meminfo = try_get_info(nvml.nvmlDeviceGetMemoryInfo, handle)
else: util = try_get_info(nvml.nvmlDeviceGetUtilizationRates, handle)
gpu_util = 0 if util != "N/A":
gpu_util = util.gpu
else:
gpu_util = 0
if meminfo != "N/A": if meminfo != "N/A":
gpu_mem_util = meminfo.used / meminfo.total * 100 gpu_mem_util = meminfo.used / meminfo.total * 100
else: else:
gpu_mem_util = -1 gpu_mem_util = -1
results[i] = {
"name": nvml.nvmlDeviceGetName(handle),
"gpu": gpu_util,
"mem": gpu_mem_util,
}
results[i] = {
"name": nvml.nvmlDeviceGetName(handle),
"gpu": gpu_util,
"mem": gpu_mem_util,
}
except:
return results return results
return results
def ffprobe_stream(path: str) -> sp.CompletedProcess: def ffprobe_stream(path: str) -> sp.CompletedProcess: