A simple Nvidia GPU monitor parsing from nvidia-smi

gpu_monitor.py

import os
import sys
import json
import time
import argparse
import xmltodict
import subprocess as sp
from collections import namedtuple

def main(show_processes, length=None, driver=True):
    # Get Nvidia SMI log message
    log, _ = sp.Popen(['nvidia-smi', '-q', '-x'], stdout=sp.PIPE).communicate()
    log = log.decode('UTF-8')
    log = xmltodict.parse(log)['nvidia_smi_log']

    # Parse information
    driver_version = log['driver_version']
    gpu_name = log['gpu']['product_name']
    fan_speed = proc(log['gpu']['fan_speed'])
    cuda_version = log['cuda_version']
    used_memory = proc(log['gpu']['fb_memory_usage']['used'])
    total_memory = proc(log['gpu']['fb_memory_usage']['total'])
    memory_usage = used_memory / total_memory * 100
    gpu_usage = proc(log['gpu']['utilization']['gpu_util'])
    gpu_temperature = proc(log['gpu']['temperature']['gpu_temp'])

    # Make header and calculate length
    header = f'{gpu_name}'
    if driver:
        header = f'{header} | Driver v{driver_version} CUDA {cuda_version}'
    if length is None:
        length = len(header) - 21

    def draw_bar(val, base=100, front='=', back='-'):
        pt = int(val / base * length)
        return '[' + front * pt + back * (length - pt) + ']'

    # Display GPU information
    print(
        f'{header}\n'
        f'  GPU Usage    {draw_bar(gpu_usage)} {gpu_usage:.0f}%\n'
        f'  Memory Usage {draw_bar(memory_usage)} {memory_usage:.0f}%\n'
        f'  Fan Speed    {draw_bar(fan_speed)} {fan_speed:.0f}%\n'
        f'  Temperature  {draw_bar(gpu_temperature)} {gpu_temperature:.0f}C'
    )

    # Parse processes information
    if show_processes:
        info = [parse_proc(p) for p in log['gpu']['processes']['process_info']]
        info = sorted(info, key=lambda x: x.memory, reverse=True)
        pid_len = max(map(len, [x.pid for x in info])) + 2
        name_len = max(map(len, [x.name for x in info]))

        print('\nProcesses')
        print(f'{"PID":>{pid_len}s} {"Process":>{name_len}s} Memory')
        for p in info:
            pid = f'{p.pid:>{pid_len}s}'
            name = f'{p.name:>{name_len}s}'
            mem = f'{p.memory / total_memory * 100:5.0f}%'
            print(f'{pid} {name} {mem}')

Process = namedtuple('Process', ['name', 'pid', 'memory'])

def parse_proc(p):
    name = os.path.basename(p['process_name'])
    mem = proc(p['used_memory'])
    pid = p['pid']

    return Process(name, pid, mem)

def proc(x):
    return float(x.split(' ')[0])

def show_cursor():
    sys.stdout.write("\033[?25h")
    sys.stdout.flush()

def hide_cursor():
    sys.stdout.write("\033[?25l")
    sys.stdout.flush()

def init_screen():
    os.system('clear')
    hide_cursor()

def exit_screen(call):
    print('\033[0;0H')
    call()
    show_cursor()
    exit(0)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--interval', '-i', default=0.1)
    parser.add_argument('--processes', '-p', default=False, action='store_true')
    parser.add_argument('--query', '-q', default=False, action='store_true')
    parser.add_argument('--length', '-l', default=None, type=int)
    parser.add_argument('--driver', '-d', default=False, action='store_true')
    args = parser.parse_args()

    call = lambda: main(args.processes, args.length, args.driver)

    if args.query:
        call()
    else:
        init_screen()
        while True:
            try:
                print('\033[0;0H')
                call()
                time.sleep(args.interval)
            except KeyboardInterrupt:
                exit_screen(call)