PythonからGPU使用率などの情報を取得する(nvidia-smiコマンド)

サーバのどのGPUが空いてるか調べたい→プログラムからGPU使用率とか取れないのかな〜と思ってちょっと調べたらいつも使っているnvidia-smiコマンドはオプション次第でマシンリーダブルな出力を出してくれそうなことがわかりました。

--query-gpuオプション

--query-gpuオプションは欲しい情報を並べるとそれをとってきて--formatオプションで指定した出力スタイルに加工して出力してくれます。

たとえば以下のような情報を指定することができるようです。

• index: 0からはじまるハードウェアのマシン上の番号, "0"とか"1"とか
• memory.free: 空きGPUメモリ量, "5123 MiB" とか
• memory.used: 使用GPUメモリ量, "3451 MiB" とか
• memory.total: GPUに搭載されているメモリ量, "8113 MiB" とか
• utilization.gpu: GPU使用率, "12 %" とか
• utilization.memory: GPUメモリ使用率, "49 %" とか
• timestamp: nvidia-smiコマンドが実行された時間, "2016/12/16 22:05:43.771" とか
• uuid: グローバルにユニークな変化しないID。GPUに貼られてる番号とかは関係ないらしい, "GPU-92275dc5-b7cc-1a5f-0348-f389e3040f2b" とか

他にもいろいろキーがあるのでnvidia-smi --help-query-gpuで確かめてみてください。

--formatオプション

--query-gpuで指定した情報をCSVに吐き出すために使用します。--format=csvが一番シンプルな呼び出し方ですが、noheaderとnounitsも指定することができます。

• noheader: CSVの一番最初の行としてヘッダ行が出てこなくなる
• nounits: たとえばGPU利用率(utlization.gpu)を取得したとき "10 %" となるのが "10" になる

これらを組み合わせて --format=csv,noheader とか --format=csv,noheader,nounits とできます。

Pythonでnvidia-smiの情報を取得する関数

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())

こんな感じで動きました。

• 第一引数( nvidia_smi_path ): nvidia-smiコマンドのパスを絶対パスなどで指定したい場合にセットしてください。デフォルトでは'nvidia-smi'になっています。
• 第二引数( keys ): 取得したい情報を各要素が文字列のlistまたはtupleオブジェクトで指定してください。デフォルトではコード中のDEFAULT_ATTRIBUTESにリストされている内容になっています。
• 第三引数( no_units ) : 各キーの値に単位を含めたいかどうかをboolでセットしてください。デフォルトではTrueになっており、単位は省かれます。
• 返り値:
  • GPUユニット個数分のdict要素を持つlistオブジェクト
  • dictの各キーは第二引数で指定したもの
  • dictの各値はバイト文字列(indexやutilization.gpuを数値として扱いたい場合はint()かけてください)

これでslackなどでbot gpu statusとかたずねるとGPU利用状況を教えてくれるbotが簡単に作れそうですね。
Enjoy!