python小工具开发专题：gpu监控工具

需求来源

在压测显卡时能看到gpu的波动情况，并输出波动的范围、gpu卡的index、显存占用、显存总量。
直接在linux 无gui的命令行展示上述信息

实现思路

使用nvitop的gpu信息接口获取信息
借助python的三方库asciichart画图

代码

import asciichartpy as acp
import argparse
import time
from nvitop import Device
import sys

parser = argparse.ArgumentParser()
parser.add_argument('-i', type=int, default=0, help="指定显卡的index")
args = parser.parse_args()

# 获取服务器上的Device信息，返回是一个list
devices = Device.all()
target_device = None
for device in devices:
    if device.index == args.i:
        target_device = device
        break

if target_device is None:
    print(f"找不到指定 ID 为 {args.i} 的设备")
    exit()
# 打印显卡的名称
print(f'  - device name :    {target_device.name()}')
# 打印显卡的序号
print(f'  - device index:    {args.i}')
# 打印显卡温度
print(f'  - Temperature :    {target_device.temperature()} \u2103')
# 打印显卡的显存大小
print(f'  - Total memory:    {target_device.memory_total_human()}')
# 打印显卡的显存占用大小
print(f'  - Used memory :    {target_device.memory_used_human()}')
print('-' * 120)

series = []

# 通过device.gpu_utilization 获取gpu利用率
try:
    while True:
        utilization = target_device.gpu_utilization()
        series.append(utilization)
        time.sleep(1)  # 每秒获取一次利用率
except KeyboardInterrupt:
    # plot gpu info
    print(acp.plot(series, {'height': 10, 'min': 0, 'max': 100, 'color': '\033[32m'}))
    # remove 0 from series
    series = [x for x in series if x != 0]
    # series = list(filter(lambda x: x != 0, series))
    if len(series) > 0:
        print("\n gpu useage range: {}%-{}%".format(min(series), max(series)))
    else:
        print("\n 检测到这张gpu上没有任务运行 \n gpu useage range: 0%-0%\n")
    sys.exit(0)

效果

开始测试gpu时，执行python plot_gpuinfo.py 运行脚本，测试结束时通过键盘ctrl+c退出监控

调研过的画图工具

asciichart:基于node js，也可以pip直接安装，无其他依赖
termplotlib：需要依赖gnuplot
bashplotlib:没有纵坐标的展示，只能画出散点图
plotille:体验不佳，每次刷新都重新绘图，图不直观，配置有点复杂
plotext：可以画图，纵坐标需要调整，图片样式有点突兀