python 查看free GPU信息

还没测试

有问题请指出，谢谢！

 

def _get_free_gpu_id(free_size=0.7):
    free_command = "nvidia-smi -q -d Memory |grep -A4 GPU|grep Free"
    all_command = "nvidia-smi -q -d Memory |grep -A4 GPU|grep Total"

    free_result = subprocess.getoutput(free_command)
    all_result = subprocess.getoutput(all_command)

    # free
    free_data = [float(item.split(':')[1].strip('MiB').strip(' ')) for item in free_result.split('\n')]
    all_data = [float(item.split(':')[1].strip('MiB').strip(' ')) for item in all_result.split('\n')]

    free_gpu_id = []
    assert len(free_data) == len(all_data)
    for idx in range(len(free_data)):
        if free_data[idx] / all_data[idx] >= free_size:
            free_gpu_id.append(idx)

    free_gpu_id = [i for i in free_gpu_id if i in config.use_gpu]

    if free_gpu_id != []:
        return choice(free_gpu_id)
    else:
        return False