Spaces:
Runtime error
Runtime error
File size: 1,022 Bytes
c19ca42 |
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 |
#!/usr/bin/env python
import os
import json
import shutil
import subprocess
import xmltodict
from rich import print # pylint: disable=redefined-builtin
from util import log, Map
def get_nvidia_smi(output='dict'):
smi = shutil.which('nvidia-smi')
if smi is None:
log.error("nvidia-smi not found")
return None
result = subprocess.run(f'"{smi}" -q -x', shell=True, check=False, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
xml = result.stdout.decode(encoding="utf8", errors="ignore")
d = xmltodict.parse(xml)
if 'nvidia_smi_log' in d:
d = d['nvidia_smi_log']
if 'gpu' in d and 'supported_clocks' in d['gpu']:
del d['gpu']['supported_clocks']
if output == 'dict':
return d
elif output == 'class' or output == 'map':
d = Map(d)
return d
elif output == 'json':
return json.dumps(d, indent=4)
return None
if __name__ == "__main__":
res = get_nvidia_smi(output='dict')
print(type(res), res)
|