Python 自动寻找可用的GPU

使用包含多个GPU的主机时，每次可用的GPU不是固定的，需要手动指定使用哪个GPU。但是，其实可以，自动寻找可用的GPU。

实现1：基于PyTorch

# 参照 lightning (原 pytorch-lightning)
import torch

"""
寻找空闲的 GPU
"""
def find_availabel_gpu(nb=1):
    for i in range(torch.cuda.device_count(), -1, -1):
        gpu_mem_size = torch.cuda.get_device_properties(0).total_memory  # 字节
        gpu_mem_size = int(gpu_mem_size * 0.95)
        # Try to allocate on device:
        device = torch.device(f"cuda:{i}")
        try:
            # torch.int8: 1B: 1个字节
            torch.ones(gpu_mem_size, dtype=torch.int8).to(device)
            torch.cuda.empty_cache()
        except RuntimeError:
            continue
        return i
    raise RuntimeError("No GPUs available.")

以上，参考lightning，使用lightning时，可以在Trainer(auto_select_gpus=True)启用。
但是，在lightning的实现中，有一个假设就是GPU启动独占模式，对没有开启独占模式的主机来说，自动寻找的可用GPU，实际上可能并不可用。

另外，上述代码有一个小的bug，就是会对每一个尝试的GPU，留下一个100M左右的显存占用，但是实际上程序并没有使用这部分显存，而是由PyTorch占用，没有找到可以释放这部分显存的代码，如果有人知道可以在评论区写出来。

实现2：基于Nvidia-smi

后台调用nvidia-smi查询GPU的使用信息，从中找出可以的GPU

import re
import subprocess


def execute_command(cmd):
    cmd_args = cmd.split()
    out = subprocess.check_output(cmd_args)
    return out.decode('utf-8')


"""
寻找空闲的 GPU
"""
def find_availabel_gpu():
    cmd = "nvidia-smi --query-gpu=memory.total,memory.free,utilization.gpu --format=csv,noheader,nounits"
    gpus_info = execute_command(cmd)
    # memory.total [MiB], memory.free [MiB], utilization.gpu [%]
    # 11264, 10393, 0
    # 11264, 11016, 0
    gpus_info = gpus_info.strip().split("\n")

    gpus_used_info = {}
    num_pattern = re.compile(r"(\d+)")
    for idx, gpu_info in enumerate(gpus_info):
        total_mem, free_mem, gpu_used = num_pattern.findall(gpu_info)
        mem_freed = int(free_mem) / int(total_mem)
        if mem_freed > 0.9:
            gpu_unit_freed = 100 - int(gpu_used)
            gpus_used_info[idx] = ([mem_freed, gpu_unit_freed])

    if len(gpus_used_info) > 1:
        gpus_used_info = sorted(
        	gpus_used_info.items(), key=lambda it: it[1], reverse=True
        )
        avail_gpu_idx = [gpu[0] for gpu in gpus_used_info]
        return avail_gpu_idx[0]
    raise RuntimeError("No GPUs available.")

如果需要监控其他的指标，可以用这个命令查看支持的属性
命令：nvidia-smi --help-query-gpu

参考: nvidia-smi 系列命令，查看gpu ，显存信息 - Wsnan - 博客园 (cnblogs.com)