List docker containers using GPU

show_containers_using_gpu.py

import os

def shell(command):
    stream = os.popen(command)
    out = stream.read().split("\n")[:-1]
    return out

def find_docker(pid: int):
    if pid == -1:
        return -1, "No Docker!"

    script = "ps -o ppid,cmd= {}".format(pid)
    pid, *com = shell(script)[1].split()
    com = " ".join(com)
    if "containerd-shim" in com:
        return pid, com
    return find_docker(pid)

if __name__ == "__main__":
    list_gpu_pids = "nvidia-smi | grep 'python' | awk '{ print $2,$5 }'"
    grep_docker = "docker ps | grep {}"

    pid_infos = shell(list_gpu_pids)

    pids = [i.split()[1] for i in pid_infos]
    gpu_ids = [i.split()[0] for i in pid_infos]

    print("GPU ID       CONTAINER INFO")
    for gpu_id, p in zip(gpu_ids, pids):
        pid, cmd = find_docker(p)
        if pid == -1:
            continue
        process = cmd.split()[4]
        container_id = process.split("/")[-1][:10]
        out = shell(grep_docker.format(container_id))
        print(gpu_id, out)