List docker containers for each GPU

gpu_alloc.py

#!/usr/bin/env python

import re
import subprocess as sp
from operator import itemgetter

import docker


def map_gpu_ids():
    pattern = re.compile(r"^GPU\s+(\d+):.+UUID:\s+([A-Za-z0-9-]+)")
    rval = {}
    out = sp.check_output("nvidia-smi -L", shell=True, universal_newlines=True)
    for line in out.splitlines():
        m = pattern.match(line)
        if m:
            idx, id_ = m.groups()
            rval[id_] = idx
    return rval


def map_container_gpus(client, gpu_map):
    rval = {}
    for c in client.containers.list():
        for kv in c.attrs["Config"]["Env"]:
            try:
                k, v = kv.split("=", 1)
            except ValueError:
                continue
            if k == "NVIDIA_VISIBLE_DEVICES":
                gpus = v.split(",")
                if gpus == ["all"]:
                    gpus = gpu_map.values()
                for g in gpus:
                    g = gpu_map.get(g, g)  # map to idx if listed by UUID
                    rval.setdefault(g, []).append(c)
    return rval


def main():
    gpu_map = map_gpu_ids()
    container_gpus = map_container_gpus(docker.from_env(), gpu_map)
    for uuid, idx in sorted(gpu_map.items(), key=itemgetter(1)):
        containers = container_gpus.get(idx, "NONE")
        print(f"GPU {idx} ({uuid}):")
        for c in containers:
            print(f"  [{c.short_id}] {c.name}")


if __name__ == "__main__":
    main()