Script to get list of docker container that used a particular GPU, based on nvitop – https://github.com/XuehaiPan/nvitop

ndtop.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
from typing import Dict
from nvitop import Device, GpuProcess, NA
import docker
from rich.console import Console
from rich.table import Table
from rich import box

docker_client = docker.from_env()
console = Console()

def get_container(searched_pid):
  """
  Get the container ID of a process with a given PID.

  Parameters:
  searched_pid (int): The PID of the process to search for.

  Returns:
  str: The ID of the container that contains the process with the given PID.
       Returns None if no container is found.
  """
  # get the list of all containers
  containers = docker_client.containers.list()
  # iterate trough the list of containers
  for container in containers:
    # get the list of processes running inside the container
    processes = docker_client.api.top(container.id)
    list_of_processes = processes["Processes"]
    # iterate through the list of processes
    for process in list_of_processes:
      pid = process[1]
      # check if the pid is the searched_pid
      if pid == searched_pid:
        return container.id
  return None
      
def add_value(list_var, value):
  """
  Adds a value to a list variable.

  Args:
    list_var (list): The list variable to which the value will be added.
    value: The value to be added to the list.

  Returns:
    list: The updated list variable with the value added.
  """
  list_var += f'\n{value}' if list_var else f'{value}'
  return list_var
      
table_docker = Table(show_header=True, header_style="cyan bold", show_lines=True, box=box.SQUARE_DOUBLE_HEAD)
table_docker.add_column("GPU #")
table_docker.add_column("%GPU")
table_docker.add_column("∑%GPU-Mem")
table_docker.add_column("PID")
table_docker.add_column("%CPU")
table_docker.add_column("%Mem")
table_docker.add_column("%GPU-Mem")
table_docker.add_column("Container ID")
table_docker.add_column("Container Name")
table_docker.add_column("Container Image")
# table_docker.add_column("Command")

# https://github.com/XuehaiPan/nvitop/blob/main/README.md#quick-start
devices = Device.all()  # or `Device.cuda.all()` to use CUDA ordinal instead

for device in devices:
  processes = device.processes()  # type: Dict[int, GpuProcess]
  sorted_pids = sorted(processes.keys())
  
  pids = containers = images = cpus = mems = cmds = container_ids = gpu_mems = ""
  for pid in sorted_pids:
    gpu_process = processes[pid]
    gpu_mem = gpu_process.gpu_memory_percent()
    container_id = get_container(str(pid))
    cpu_percent = processes[pid].cpu_percent()
    memory_percent = round(processes[pid].memory_percent(), 2)
    command = processes[pid].command()
    
    pids = add_value(pids, pid)
    cpus = add_value(cpus, f'{cpu_percent}%')
    mems = add_value(mems, f'{memory_percent}%')
    gpu_mems = add_value(gpu_mems, f'{gpu_mem}%')
    # cmds = add_value(cmds, command)

    if container_id:
      container = docker_client.containers.get(container_id)
      container_name = container.attrs["Name"].lstrip("/")
      container_image = container.attrs["Config"]["Image"]

      container_ids = add_value(container_ids, container_id[:12])
      containers = add_value(containers, container_name)
      images = add_value(images, container_image)
    else:
      container_ids = add_value(container_ids, NA)
      containers = add_value(containers, NA)
      images = add_value(images, NA)

  table_docker.add_row(
    str(device.index),
    str(device.gpu_utilization()) + "%",
    str(device.memory_percent()) + "%",
    pids if pids else NA,
    cpus if cpus else NA,
    mems if mems else NA,
    gpu_mems if gpu_mems else NA,
    container_ids if container_ids else NA,
    containers if containers else NA,
    images if images else NA,
    # cmds if cmds else NA
  )

console.print(table_docker)