通过显卡占用率和显存占用率获取空闲GPUs

创建idleGPUs.py，内容如下：

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from os import popen
from typing import List


def query_idle_gpus(threshold: int = 2) -> List[int]:
    CMD = (
        "nvidia-smi --format=csv,noheader,nounits"
        " --query-gpu=index,utilization.gpu,utilization.memory"
    )
    idle_gpus: List[int] = []
    with popen(CMD) as pipe:
        for line in pipe.readlines():
            strs = line.split(", ")
            if len(strs) != 3:  # nvidia-smi is not available
                return []
            idx, util, memory = int(strs[0]), int(strs[1]), int(strs[2])
            if util < threshold and memory < threshold:
                idle_gpus.append(int(idx))
    return idle_gpus


if __name__ == "__main__":
    idle_gpus = query_idle_gpus()
    print(*idle_gpus, sep=",")

这个程序利用了nvidia-smi命令，查询显卡的占用率和显存的占用率。当两者的占用率都低于阈值（这里是2%）的时候，则认为该显卡是空闲的。运行该程序，即可获得通过逗号分割的空闲显卡。当计算机未安装显卡或未正确安装显卡时，程序输出为空。

同时，您可以将下面的命令架在每个运行GPU程序的命令的开头，让程序仅使用空闲的显卡。

CUDA_VISIBLE_DEVICES="`python ./idleGPUs.py`"