Get information such as GPU usage from Python (nvidia-smi command)

I want to find out which GPU on the server is free → I wonder if I can get the GPU usage rate from the program ~ I did a little research and the nvidia-smi command I always use gives a machine-readable output depending on the option I found that.

--query-gpu option

The --query-gpu option fetches the desired information, processes it into the output style specified by the --format option, and outputs it.

For example, it seems that you can specify the following information.

--index: Hardware machine numbers starting with 0, such as "0" or "1" --memory.free: amount of free GPU memory, like "5123 MiB" --memory.used: GPU memory used, such as "3451 MiB" --memory.total: Amount of memory installed in GPU, such as "8113 MiB" --utilization.gpu: GPU utilization, like "12%" --utilization.memory: GPU memory usage, like "49%" --timestamp: time when the nvidia-smi command was executed, such as "2016/12/16 22: 05: 43.771" --uuid: A globally unique, unchanging ID. It seems that it has nothing to do with the number affixed to the GPU, such as "GPU-92275dc5-b7cc-1a5f-0348-f389e3040f2b"

There are many other keys, so check them out with nvidia-smi --help-query-gpu.

--format option

It is used to output the information specified by --query-gpu to CSV. --format = csv is the simplest way to call it, but you can also specify noheader and nounits.

--noheader: The header line does not appear as the first line of CSV --nounits: For example, when you get the GPU utilization (utlization.gpu)," 10% "becomes" 10 ".

You can combine these to make --format = csv, noheader or --format = csv, noheader, nounits.

Function to get nvidia-smi information in Python

import subprocess
import json

DEFAULT_ATTRIBUTES = (
    'index',
    'uuid',
    'name',
    'timestamp',
    'memory.total',
    'memory.free',
    'memory.used',
    'utilization.gpu',
    'utilization.memory'
)

def get_gpu_info(nvidia_smi_path='nvidia-smi', keys=DEFAULT_ATTRIBUTES, no_units=True):
    nu_opt = '' if not no_units else ',nounits'
    cmd = '%s --query-gpu=%s --format=csv,noheader%s' % (nvidia_smi_path, ','.join(keys), nu_opt)
    output = subprocess.check_output(cmd, shell=True)
    lines = output.decode().split('\n')
    lines = [ line.strip() for line in lines if line.strip() != '' ]

    return [ { k: v for k, v in zip(keys, line.split(', ')) } for line in lines ]


import pprint
pprint.pprint(get_gpu_info())

It worked like this.

--First argument (nvidia_smi_path): nvidia-smi Set this when you want to specify the command path as an absolute path. The default is 'nvidia-smi'. --Second argument (keys): Specify the information you want to get in the list or tuple object where each element is a character string. By default, it is listed in DEFAULT_ATTRIBUTES in the code. --Third argument (no_units): Set with bool whether you want to include units in the value of each key. The default is True and the unit is omitted. --Return value: --A list object with dict elements for the number of GPU units --Each key of dict is the one specified by the second argument --Each value of dict is a byte string (if you want to treat ʻindex or ʻutilization.gpu as a number, use ʻint ()`)

With this, if you ask bot gpu status with slack etc., it seems that you can easily create a bot that will tell you the GPU usage status. Enjoy!