我們經常會為了組內誰在佔用某塊gpu卡而不使用 煩惱,透過簡單的程式碼就能快速識別到這塊卡上面的程序是哪個容器的。
下面的程式碼會給出如下圖這樣的結果
#! /miniconda3/bin/python
# -*- coding: utf-8 -*-
import psutil
import subprocess as sp
def 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 ans
def 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: continue
id1_name = id_name.split()
if len(id1_name) != 2: continue
id1,name = id1_name
id_name_di[id1]=name
return id_name_di
def 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 docker
return None
def 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: continue
id1,pid = id1_pid
anstmp = check(pid)
if not anstmp: continue
docker = anstmp[:12]
if docker in id_name_di:
id_docker.append({'id':id1,'pid':pid,'name': id_name_di[docker]})
else:
continue
print('=================================')
#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)