mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
36 lines
1022 B
Python
Executable File
36 lines
1022 B
Python
Executable File
#!/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)
|