nvidia_print.py

import subprocess
import re
import pprint
nvidia_ouput = subprocess.check_output(['nvidia-smi'])
part = re.findall('Usage(.+)', str(nvidia_ouput), re.DOTALL)
pid_memory = re.findall('\s(\d)\s+(\d+).+?\s(\d+MiB)', str(part), re.DOTALL)
user_dict = {}
print('='*50)
for gpu_num, pid, memory in pid_memory:
    print('GPU number: ', gpu_num)
    print('Pid: ', pid)
    print('GPU memory :', memory)

    cat_output = subprocess.check_output(['cat', '/proc/{}/cgroup'.format(pid)])
    try:
        docker_id = re.findall('/docker/([\d\w]+)', str(cat_output), re.DOTALL)[0]
    except:
        print('Pid {} not have docker'.format(pid))
        print('-'*50)
        continue
    print('container id: ', docker_id)
    user_name = subprocess.check_output(['docker', 'inspect', "--format", "'{{.Name}}'", "{}".format(docker_id)])
    print(str(user_name))
    print('-'*50)
    memory_int = int(re.findall('\d+', str(memory))[0])
    if user_dict.get(str(user_name), False):
        user_dict[str(user_name)] += memory_int
    else:
        user_dict[str(user_name)] = memory_int
print('='*50)
pprint.pprint(user_dict)