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

我们经常会为了组内谁在占用某块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)