File size: 2,479 Bytes
b84549f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import subprocess
import time
import traceback
from xml.dom import minidom
def collect_gpu_usage(node_id):
cmd = 'nvidia-smi -q -x'.split()
info = None
try:
smi_output = subprocess.check_output(cmd)
info = parse_nvidia_smi_result(smi_output)
except Exception:
traceback.print_exc()
info = gen_empty_gpu_metric()
return info
def parse_nvidia_smi_result(smi):
try:
output = {}
xmldoc = minidom.parseString(smi)
gpuList = xmldoc.getElementsByTagName('gpu')
output["Timestamp"] = time.asctime(time.localtime())
output["gpuCount"] = len(gpuList)
output["gpuInfos"] = []
for gpuIndex, gpu in enumerate(gpuList):
gpuInfo = {}
gpuInfo['index'] = gpuIndex
gpuInfo['gpuUtil'] = gpu.getElementsByTagName('utilization')[0]\
.getElementsByTagName('gpu_util')[0]\
.childNodes[0].data.replace("%", "").strip()
gpuInfo['gpuMemUtil'] = gpu.getElementsByTagName('utilization')[0]\
.getElementsByTagName('memory_util')[0]\
.childNodes[0].data.replace("%", "").strip()
processes = gpu.getElementsByTagName('processes')
runningProNumber = len(processes[0].getElementsByTagName('process_info'))
gpuInfo['activeProcessNum'] = runningProNumber
gpuInfo['gpuType'] = gpu.getElementsByTagName('product_name')[0]\
.childNodes[0].data
memUsage = gpu.getElementsByTagName('fb_memory_usage')[0]
gpuInfo['gpuMemTotal'] = memUsage.getElementsByTagName('total')[0]\
.childNodes[0].data.replace("MiB", "").strip()
gpuInfo['gpuMemUsed'] = memUsage.getElementsByTagName('used')[0]\
.childNodes[0].data.replace("MiB", "").strip()
gpuInfo['gpuMemFree'] = memUsage.getElementsByTagName('free')[0]\
.childNodes[0].data.replace("MiB", "").strip()
output["gpuInfos"].append(gpuInfo)
except Exception:
traceback.print_exc()
output = {}
return output
def gen_empty_gpu_metric():
try:
output = {}
output["Timestamp"] = time.asctime(time.localtime())
output["gpuCount"] = 0
output["gpuInfos"] = []
except Exception:
traceback.print_exc()
output = {}
return output
|