一个小工具识别哪个docker占用gpu

我们经常会为了组内谁在占用某块gpu卡而不使用 烦恼，通过简单的代码就能快速识别到这块卡上面的进程是哪个容器的。
下面的代码会给出如下图这样的结果

#! /miniconda3/bin/python
# -*- coding: utf-8 -*-
import psutil
import subprocess as spdef pre():ans = sp.check_output('''nvidia-smi |grep MiB|grep -v Default|awk '{print $2"\t"$5}' ''',shell=True)ans = ans.strip().decode('utf-8')return ansdef pre1():id_name_di = {}ans = sp.check_output(''' docker ps|awk '{print $1"\t"$NF}' ''',shell=True)ans = ans.strip().decode('utf-8')for id_name in ans.strip().split('\n'):if 'CONTAINER' in id_name: continueid1_name = id_name.split()if len(id1_name) != 2: continueid1,name = id1_nameid_name_di[id1]=namereturn id_name_didef check(pid):p = psutil.Process(int(pid))for curpid in p.parents():if 'containerd-shim' in curpid.name():docker = curpid.cwd().strip().split('/')[-1]return dockerreturn Nonedef get(ans,id_name_di):id_docker = []for id_pid in ans.split('\n'):id1_pid = id_pid.strip().split()if len(id1_pid) != 2: continueid1,pid = id1_pidanstmp = check(pid)if not anstmp: continuedocker = anstmp[:12]if docker in id_name_di:id_docker.append({'id':id1,'pid':pid,'name': id_name_di[docker]})else:continueprint('=================================')#print('\033[31m=================================\033[0m')for id_docker in id_docker:#print('卡号:\033[31m[{}]\033[0m\tdocker名称:\033[31m[{}]\033[0m'.format(id1,docker))id1 = id_docker['id']pid = id_docker['pid']docker = id_docker['name']print('卡号:[{}]\tpid:{}\tdocker名称:[{}]'.format(id1,pid,docker))if __name__ == '__main__':ans = pre()id_name_di = pre1()get(ans,id_name_di)