Code:
#!/usr/bin/env python
import argparse
import subprocess
import sys
def nvidia_smi(verbose):
cmd = ['nvidia-smi','-q']
try:
output = subprocess.check_output(cmd).decode("utf-8")
except subprocess.CalledProcessError:
if verbose:
sys.stdout.write('ERROR: nvidia-smi failed with a nonzero exitcode')
sys.exit(1)
for line in output.splitlines():
yield line
def fan_percent(smi_output,gpu):
percentages = []
for line in smi_output:
if 'Fan Speed' in line:
percentages.append(line.split()[2])
return percentages[gpu]
def temperature(smi_output,gpu):
temps = []
for line in smi_output:
if 'Gpu' in line and 'C' in line:
temps.append(line.split()[2])
return temps[gpu]
def vram_free(smi_output,gpu):
free = []
for line in smi_output:
if 'Free' in line and 'MB' in line:
free.append(line.split()[2])
return free[gpu]
def vram_used(smi_output,gpu):
usage = []
for line in smi_output:
if 'Used' in line and 'MB' in line:
usage.append(line.split()[2])
return usage[gpu]
def vram_total(smi_output,gpu):
totals = []
for line in smi_output:
if 'Total' in line and 'MB' in line:
totals.append(line.split()[2])
return totals[gpu]
def main(args):
verbose = args.verbose
gpu = args.gpu[0]
try:
smi_output = [line for line in nvidia_smi(verbose) if 'N/A' not in line]
except (OSError,FileNotFoundError):
if verbose:
sys.stdout.write('ERROR: nvidia-smi binary not found in path')
sys.exit(1)
try:
argarray = {
'fanpct' : fan_percent(smi_output,gpu),
'temp' : temperature(smi_output,gpu),
'vramfree' : vram_free(smi_output,gpu),
'vramused' : vram_used(smi_output,gpu),
'vramtotal' : vram_total(smi_output,gpu)
}
except IndexError:
if verbose:
sys.stdout.write('ERROR: Invalid GPU specified')
sys.exit(1)
del smi_output
for arg in vars(args):
try:
sys.stdout.write(argarray[arg])
sys.exit()
except KeyError:
pass
if __name__ == '__main__':
desc = 'Parse the output of the nvidia-smi utility for simple Conky integration'
use = '%(prog)s [option] gpu'
ep = 'Example: python %(prog)s --temperature 0'
parser = argparse.ArgumentParser(usage=use,description=desc,epilog=ep)
parser.add_argument('-v','--verbose',action='store_true',default=False,dest='verbose',
help='display error messages')
parser.add_argument('--version',action='version',version='%(prog)s 1.1')
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('--fan-speed',action='store_true',default=argparse.SUPPRESS,dest='fanpct',
help='fan speed as a percentage')
group.add_argument('--temperature',action='store_true',default=argparse.SUPPRESS,dest='temp',
help='GPU temperature in degrees Celsius')
group.add_argument('--vram-free',action='store_true',default=argparse.SUPPRESS,dest='vramfree',
help='free VRAM in megabytes')
group.add_argument('--vram-used',action='store_true',default=argparse.SUPPRESS,dest='vramused',
help='used VRAM in megabytes')
group.add_argument('--vram-total',action='store_true',default=argparse.SUPPRESS,dest='vramtotal',
help='total VRAM in megabytes')
parser.add_argument('gpu',action='store',type=int,nargs=1,
help='GPU number, first GPU is 0')
args = parser.parse_args()
main(args)