Get the container name of each process running in each gpu

container_gpu.py

import subprocess
import collections

gpus = subprocess.check_output("nvidia-smi | grep '    C' | awk '{print $2}'", shell=True).split('\n')[:-1]

gpids = subprocess.check_output("nvidia-smi | grep '    C' | awk '{print $3}'", shell=True).split('\n')[:-1]

cnames = subprocess.check_output("docker ps -q | xargs docker inspect --format '{{.State.Pid}} {{.Name}}' | awk '{print $2}'", shell=True).split('\n')[:-1]

cpids = subprocess.check_output("docker ps -q | xargs docker inspect --format '{{.State.Pid}} {{.Name}}' | awk '{print $1}'", shell=True).split('\n')[:-1]

gpu_dict = {}

def get_pchilds(pid):
        pids = []
        try:
                childs = subprocess.check_output("pgrep -P {}".format(pid), shell=True).split('\n')[:-1]
        except:
                childs = []

        pids.extend(childs)
        for child in childs:
                pids.extend(get_pchilds(child))

        return pids


for cname, cpid in zip(cnames, cpids):
        child_c_pids = get_pchilds(cpid)
                
          
        for child in child_c_pids:
                for index, gpid in enumerate(gpids):
                        if child == gpid:
                                if gpus[index] in gpu_dict:
                                        gpu_dict[gpus[index]].append(cname)
                                else:
                                        gpu_dict[gpus[index]] = [cname]

for gpu, containers in collections.OrderedDict(sorted(gpu_dict.items())).items():
        for cname in containers:
                print "GPU: {}, Container name: {}".format(gpu, cname)