Get gpu usage by pid about pynvml

[rjzamora]

Using nvmlDeviceGetAccountingStats might be solution, but I have not personally tested or experimented with this. I cannot immediately test this myself, because root/admin access is need to set the appropriate accounting mode for the target device (e.g. using nvmlDeviceSetAccountingMode).

If you are interested in getting a pid breakdown of gpu-memory consomption for a given device, yout can use nvmlDeviceGetComputeRunningProcesses:

import pynvml

pynvml.nvmlInit()
for dev_id in range(pynvml.nvmlDeviceGetCount()):
    handle = pynvml.nvmlDeviceGetHandleByIndex(dev_id)
    for proc in pynvml.nvmlDeviceGetComputeRunningProcesses(handle):
        print(
            "pid %d using %d bytes of memory on device %d."
            % (proc.pid, proc.usedGpuMemory, dev_id)
        )

I will try to take a harder look into this later and will update here if I find anything.

[cod3licious]

This is the code I'm using now. It only works for type "C" processes, not type "G" in the list when calling nvidia-smi directly in the terminal (e.g. my neural network in python was a type "C" process, Xorg type "G") but I think it should suffice:

import os
from pynvml.smi import nvidia_smi
from bs4 import BeautifulSoup as bs


def get_gpu_usage(pid=None):
    if pid is None:
        # get current process id
        pid = os.getpid()
    pid = str(pid)
    # find associated gpu usage
    soup = bs(nvidia_smi.getInstance().XmlDeviceQuery(), "lxml")
    for p in soup.find("processes").find_all("process_info"):
        if p.find("pid").get_text() == pid:
            return p.find("used_memory").get_text()
    # process not found
    return "0 MiB"

[cod3licious]

I'd love that feature as well! Calling .XmlDeviceQuery() gives an XML where the different processes are listed so it should be possible if something like beautiful soup could be an additional dependency.