多GPU分配#

实现逻辑：寻找（可用显存 / 总显存）最大的的GPU，并优先安排任务

nvidia-smi可以很方便的获得GPU的各种详细信息。

首先获得可用的GPU数目，nvidia-smi -L | grep GPU |wc -l

然后获得GPU各自的总显存，nvidia-smi -q -d Memory | grep -A4 GPU | grep Total | grep -o ‘[0-9]+‘

最后获得GPU各自的可用显存，nvidia-smi -q -d Memory | grep -A4 GPU | grep Free | grep -o ‘[0-9]+‘

将（可用显存 / 总显存）另存为numpy数组，并使用np.argmax返回值即为可用GPU

def available_GPU(self):
    import subprocess
    import numpy as np
    nDevice = int(subprocess.getoutput("nvidia-smi -L | grep GPU |wc -l"))
    total_GPU_str = subprocess.getoutput("nvidia-smi -q -d Memory | grep -A4 GPU | grep Total | grep -o '[0-9]\+'")
    total_GPU = total_GPU_str.split('\n')
    total_GPU = np.array([int(device_i) for device_i in total_GPU])
    avail_GPU_str = subprocess.getoutput("nvidia-smi -q -d Memory | grep -A4 GPU | grep Free | grep -o '[0-9]\+'")
    avail_GPU = avail_GPU_str.split('\n')
    avail_GPU = np.array([int(device_i) for device_i in avail_GPU])
    avail_GPU = avail_GPU / total_GPU
    return np.argmax(avail_GPU)