多服务器显卡使用状态监控工具实现

一、前言

实验室是做计算机视觉的，拥有几台服务器，每台服务器上有1到8张数量不等的显卡，用于同学们跑深度学习算法。每次到了会议/期刊截稿前，服务器都人满为患，大家各种抢显卡用（同时占用多张显卡能提高程序的并行性，训练模型花费时间变少），旱的旱死涝旳涝死，有的同学占了十几张卡，而有的同学天天蹲守服务器也等不到一张空闲的卡。于是，产品经理（其实是实验室大师兄）提了个需求，让我写个显卡监控工具，用于监控所有服务器的显卡使用情况，并汇总到一起做个排名（每人用了多少张卡，占了多少显存等等）。

二、实现

• 我的想法基于C/S模型
  • 在每个需要监控的服务器上运行一个客户端，每隔一段时间就把当前服务器的显卡使用情况发送给服务器
  • 服务器负责接收这些信息，并做汇总排序
• 有了想法后，再参考网上大佬的代码，具体实现如下

2.1 客户端

• 用Popen函数执行nvidia-smi命令，获取返回信息。然后使用正则匹配的方式获取nvidia-smi返回信息中的GPU信息和进程信息。GPU信息包含每张显卡的显存大小、已被占用显存；进程信息包含显卡ID、进程ID、执行命令、占用显存。然后通过进程ID找到对应的用户名

def get_owner(pid):
    try:
        for line in open('/proc/%d/status' % pid):
            if line.startswith('Uid:'):
                uid = int(line.split()[1])
                return pwd.getpwuid(uid).pw_name
    except:
        return None

def get_info():
    info = {
    'gpu': [], 'process': [] }
    msg = subprocess.Popen('nvidia-smi', stdout = subprocess.PIPE).stdout.read().decode()
    msg = msg.strip().split('\n')

    lino = 8
    while True:
        status = re.findall('.*\d+%.*\d+C.*\d+W / +\d+W.* +(\d+)MiB / +(\d+)MiB.* +\d+%.*', msg[lino])
        if status == []: break
        mem_usage, mem_total = status[0]
        info['gpu'