python监控显卡显存

内存监控

在工作组中，经常会遇到多人争抢显卡的情况，僧多肉少，所以可以做个监控显卡的头部程序，然后再调取主程序训练模型。
我们借用pynvml，没有安装的，可以使用下面命令直接安装：

pip install pynvml

下面使用pynvml可以查看显卡总共多少显存，剩余多少显存，以及已用多少显存。

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 这里的0表示第一张显卡
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(meminfo.total/1024/1024/1024) # 输出第一张显卡的显存
print(meminfo.free/1024/1024/1024) # 输出第一张显卡的可用显存
print(meminfo.used/1024/1024/1024) # 输出第一张显卡的已用显存

我们如果想对显卡进行监控，可以使用下面的命令：

import pynvml

# 定义一个读取gpu显卡的函数，gpu为第几张显卡，free表示需要多少剩余显存
def test_gpu(gpu, free):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)

    return meminfo.free/1024/1024/1024 < free

然后用while循环对显卡进行监控：

while test_gpu(1, 40):
	print("GPU memory < %dGB, wait 1 minut"%40)
    time.sleep(60)

当显卡的剩余显存大于40GB的时候，才可以执行下面的程序，否则将等待1分钟（可以根据自己情况设置）。
如果想要执行python程序文件，可以使用os库运行：

import os
os.system('python run.py')

进程监控

获取进程pid

既然要进程监控，就需要获得当前显卡上运行程序的进程pid，一张显卡上的进程个数可能是0，也可能是1，还有可能是N。没人用卡进程就是0，一个人用时1（tensorflow一般会将显存占完，不让其它新用户使用，pytorch则不是，用多少，占用多少），多人在用显卡则有N个进程。

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 这里的0表示第一张显卡
gpu_ps = pynvml.nvmlDeviceGetComputeRunningProcesses(handle) # 获取该显卡上运行的程序，他是一个list

# 输出该显卡一共多少进程
print(len(gpu_ps))
# 输出第一个进程的pid
print(gpu_ps[0].pid)
# 将所有进程的pid
for gpu_p in gpu_ps:
	print(gpu_p.pid)

管理进程pid的死活

杀死进程可以借助psutil，当然，需要保证的时该pid是你自己的进程，不然无法管理。

import os, psutil

os.system("kill gpu_ps[0].pid")
print(psutil.pid_exists(gpu_ps[0].pid)) # 还在活着，输出True，死掉了输出False