from pynvml import nvmlDeviceGetCount, nvmlInit, nvmlShutdown, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo


def get_free_gpu_number():
    free_list = []
    nvmlInit()
    for i in range(nvmlDeviceGetCount()):
        handle = nvmlDeviceGetHandleByIndex(i)
        mem_info = nvmlDeviceGetMemoryInfo(handle)
        free_list.append(mem_info.free / mem_info.total)
    nvmlShutdown()
    max_free = max(free_list)
    if max_free < 0.3:
        raise Exception('GPU计算资源不足')
    return free_list.index(max_free)


if __name__ == '__main__':
    print(get_free_gpu_number())