json formatting of nvidia-settings

gpu_stat.py

#!/usr/bin/env python
# gpu_stat.py [DELAY [COUNT]]
# dump gpu stats as a json and plot in file
# {"cpu":5.9,"gpu":{"used_mem":7959,"util":{"graphics":91,"memory":56,"video":0,"PCIe":2}},"time":1505480296.0436406}

import sys
import time
import json
from pathlib import Path
# import socket
import psutil
import subprocess
import numpy as np
import matplotlib.pylab as plt


if '__main__' == __name__:

    try:
        delay = int(sys.argv[1])
    except IndexError:
        delay = 0.25

    try:
        count = int(sys.argv[2])
    except IndexError:
        count = None

    conv_val = lambda t: (t[0], int(t[1]))
    start_time = time.time()

    json_path = Path(__file__).parent.joinpath(f'gpu_stats_{int(start_time)}.json')
    plot_path = str(Path(__file__).parent.joinpath(f'gpu_stats_{int(start_time)}.png'))
    print(f'JSON: {json_path}, PLOT: {plot_path}')

    stats = []
    fig = plt.figure()
    i = int(0)
    while True:
        # GPU = socket.gethostname() + ':0[gpu:0]'
        # cmd = ['nvidia-settings', '-t', '-q', GPU + '/GPUUtilization']
        cmd = ['nvidia-settings', '-t', '-q', 'GPUUtilization']

        gpu_util = subprocess.check_output(cmd).decode('utf-8').strip().split(',')
        gpu_util = dict(conv_val(f.strip().split('=')) for f in gpu_util)

        # cmd[-1] = GPU + '/UsedDedicatedGPUMemory'
        cmd[-1] = 'UsedDedicatedGPUMemory'
        gpu_used_mem = subprocess.check_output(cmd).decode('utf-8').strip()

        current_stats = dict(cpu=psutil.cpu_percent(),
                             gpu=dict(used_mem=int(gpu_used_mem), util=gpu_util),
                             time=time.time())
        stats.append(current_stats)

        # print(json.dumps(current_stats, separators=(',', ':')))
        if i % 100 > 5:
            with json_path.open('w+') as fp:
                json.dump(stats, fp, separators=(',', ':'))

        if count is not None and i == count:
            exit(0)
        i += 1

        if i % 10 == 9:
            latest_stats = stats[-200:]

            time_intervals = [s['time'] - start_time for s in latest_stats]
            cpu_ = [s['cpu'] for s in latest_stats]
            graphics_ = [s['gpu']['util']['graphics'] for s in latest_stats]
            memory_ = [s['gpu']['util']['memory'] for s in latest_stats]

            latest_10_percent = int(20) if len(latest_stats) == 200 else int(len(cpu_) * 0.9)

            fig.clf()
            plt.plot(time_intervals, cpu_, label=f'CPU: {np.mean(cpu_[latest_10_percent:]):.1f}')
            plt.plot(time_intervals, graphics_, label=f'GPU: {np.mean(graphics_[latest_10_percent:]):.1f}')
            plt.plot(time_intervals, memory_, label=f'GPU_IO: {np.mean(memory_[latest_10_percent:]):.1f}')
            plt.legend(loc='lower left')
            plt.savefig(plot_path, bbox_inches='tight', pad_inches=0)
            # plt.show(block=False)

        time.sleep(delay)