PythonからGPUデバイスが使用可能かチェックする

TL;DR

Pythonの処理からデバイス番号を指定して，指定したGPUが使えるかどうかをチェックする方法です．
ここで使えるとは，デバイスが存在することを指します．

GPUデバイスが使用可能かチェックする

NVIDIAからライブラリ等が提供されているわけではないので，nvidia-smiコマンドを叩いた結果から確認します．
（ライブラリがあった場合はご教授お願い致します）

以下のようにします．

"""
GPUデバイスが使用可能か確認する
"""
import subprocess


def check_gpu_device(device: int) -> bool:
    """
    GPUデバイスが使用可能か確認する
    Args:
        device: デバイス番号
    """
    output = subprocess.check_output(["nvidia-smi", "-L"])
    lines = output.decode().split('\n')
    for line in lines:
        if line == "":
            continue
        gpu_device_id = line.split(":")[0]
        if gpu_device_id == f"GPU {device}":
            return True
    return False

-Lオプションを付けることで1行1デバイスで列挙できるので，Pythonから扱いやすくなります．
（--query-gpuでもOKです）

まとめ

PythonからGPUデバイスが使用可能かチェックする処理を紹介しました．
基本的には自分が使っているPCのデバイス番号は把握しているかと思いますが，何かしらのアプリケーションの形で提供している場合には使えるかと思います．
ご参考になれば幸いです．