多服务器显卡使用状态监控工具实现
一、前言
实验室是做计算机视觉的，拥有几台服务器，每台服务器上有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'].append({
'mem_usage': float(mem_usage),
'mem_total': float(mem_total),
})
lino += 3
lino = -1
while True:
lino -= 1
status = re.findall('\| +(\d+) +(\d+) +\w+ +([^ ]*) +(\d+)MiB \|', msg[lino])
if status == []: break
gpuid, pid, program, mem_usage = status[0]
username = get_owner(int(pid))
if username is None:
print('进程已经不存在')
continue
try:
p = psutil.Process(int(pid))
p.cpu_percent()
time.sleep(0.5)
cpu_percent = p.cpu_percent()
except psutil.NoSuchProcess:
print('进程已经不存在')
continue
info['process'].append({
'gpuid': int(gpuid),
'pid': int(pid),
'program': program,
'cpu_percent': cpu_percent,
'mem_usage': float(mem_usage),
'username': username,
})
info['process'].reverse()
return info
- 获取显卡信息后通过HTTP库request，把显卡信息发送给服务器，并sleep一定时间（每隔一定时间获取一次显卡信息并发送）
while True:
mean_info = get_info()
data = json.dumps(mean_info)
try:
response = requests.get(url, data = data)
print('HTTP状态码:', response.status_code)
except Exception as e:
print(e)
time.sleep(opt.persecond)
2.2 服务器
- 使用http.server中的HTTPServer类，开启一个HTTP服务器，在do_GET函数中对客户端的GET请求进行处理，主要是更新客户端的显卡信息和时间戳
class CustomHandler(BaseHTTPRequestHandler):
alert_record = { }
def do_GET(self):
length = int(self.headers['content-length'])
info = json.loads(self.rfile.read(length).decode())
slaver_address, _ = self.client_address
lock.acquire()
if slaver_address not in info_record:
info_record[slaver_address] = {}
info_record[slaver_address]['info'] = info
info_record[slaver_address]['timestamp'] = time.time()
lock.release()
report_user()
self.send_response(200)
self.end_headers()
- 在report_user函数中对显卡信息进行排序输出。遍历info_record（info_record是一个dict，用于保存每个客户端的显卡信息和时间戳，直接用客户IP地址字符串作为dict的键），判断客户端的显卡信息是否过期（默认100s不更新就算过期），若过期，直接continue；若不过期，则分别统计用户显存占用和用户显卡占用。处理完info_record后，把排序结果print到终端
def report_user():
usage_dict = { }
usage_num = { }
lock.acquire()
for slaver_address in sorted(info_record.keys()):
if isExpire(info_record[slaver_address]['timestamp']):
continue
pi_list = info_record[slaver_address]['info']['process']
for pi in pi_list:
username = pi['username']
mem_usage = pi['mem_usage']
gpu_id = pi['gpuid']
usage_dict[username] = usage_dict.get(username, 0) + mem_usage
if username in usage_num:
usage_num[username].add(slaver_address + ':' + str(gpu_id))
else:
s = set()
s.add(slaver_address + ':' + str(gpu_id))
usage_num[username] = s
lock.release()
print('=' * 60)
print(time.strftime('%Y-%m-%d %H:%M:%S\n', time.localtime(time.time())))
usage_list = sorted(usage_dict.items(), key = lambda x: x[1], reverse=True)
print('用户显存占用排序：')
print('<用户ID> : <占用显存(MB)>')
for username, usagememory in usage_list:
print('%s : %dM' % (username, usagememory))
# report1 = ['%s : %dM' % (n, u) for n, u in usage_list]
# report1 = '\n'.join(report1)
print()
usage_num_list = sorted(usage_num.items(), key = lambda x: len(x[1]), reverse=True)
print('用户显卡占用排序：')
print('<用户ID> : <占用显卡数量> : <显卡IP地址和ID>')
for username, gpuset in usage_num_list:
print('%s : %d : ' % (username, len(gpuset)), end='')
for gpuaddr in gpuset:
print('%s, ' % gpuaddr, end='')
print()
print('=' * 60)
return