GPU memory and user usage

lsgpu

#!/usr/bin/env /usr/bin/python
# location: ~/.local/bin/lsgpu

import subprocess
import re

# ps -aux
raw_stdout = subprocess.run(["ps", "-aux"], stdout=subprocess.PIPE).stdout.decode(
    "utf-8"
)
raw_stdout = raw_stdout.split("\n")[1:-1]

pid_to_user = {}

for x in raw_stdout:
    x = x.split()
    user = x[0]
    pid = int(x[1])
    pid_to_user[pid] = user

# nvidia-smi
nvidia_smi_output = subprocess.run(
    ["nvidia-smi"], stdout=subprocess.PIPE
).stdout.decode("utf-8")
raw_stdout = nvidia_smi_output.split("\n ")[0].split("===|")[1].split("--+\n")

gpu_ids = []
gpu_mem = {}

for x in raw_stdout:
    gpu_id = re.findall(r"\|\s+(\d+)\s", x)
    if len(gpu_id) == 1:
        gpu_id = int(gpu_id[0])

        gpu_ids += [gpu_id]

        mem = re.findall(r"\|\s+(\d+)MiB\s/\s(\d+)MiB\s+\|", x)
        if len(mem) == 1:
            mem = mem[0]
            mem = (int(mem[0]) / 1024, int(mem[1]) / 1024)
            mem = f"({mem[0]: 2.0f} /{mem[1]: 3.0f})GiB"
        else:
            mem = "(  N/A  )GiB"
        gpu_mem[gpu_id] = mem

gpu_to_user = {i: set() for i in gpu_ids}

if "No running processes" not in nvidia_smi_output:
    raw_stdout = (
        nvidia_smi_output.split("Processes")[1].split("=====|")[1].split("\n")[1:-2]
    )

    for x in raw_stdout:
        x = x.split()
        gpu_id = int(x[1])
        pid = int(x[4])
        gpu_to_user[gpu_id].add(pid_to_user[pid])

for gpu_id, users in gpu_to_user.items():
    users = " ".join(users)
    print(f"GPU {gpu_id} {gpu_mem[gpu_id]}: {users}")