import pynvml
threshold = 0.9
def get_gpu(need_gpu_id):
    used = get_gpu_threshold(need_gpu_id)
    #小于一定的
    if used > threshold:
        need_use = get_use_gpu()
        return need_use[0]
    else:
        return need_gpu_id

def get_use_gpu():
    use=[]
    for index in range(pynvml.nvmlDeviceGetCount()):
        used = get_gpu_threshold(index)
        if used > threshold:
            use.append(index)
    return use

def get_gpu_count():
    return pynvml.nvmlDeviceGetCount()

def get_gpu_threshold(index):
    handle = pynvml.nvmlDeviceGetHandleByIndex(index)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
    used = meminfo.used / meminfo.total
    return used