使用python多GPU任务分配

背景：在做地震波正演时需要逐炮进行正演，这里面涉及到多节点、多GPU的任务分发工作。前人之前写的版本静态分发任务，不能满足节点内部GPU算力不同（比如两张1080Ti两张780Ti）而导致的计算时间浪费。所以考虑使用动态任务分配进行GPU调度。

代码链接：https://github.com/sh39o/Acoustic3d-fd

目的：使用python将诸多任务动态分配到空闲的GPU上进行运算，直到任务全部分配完成。与CPU多线程计算不同，每个计算节点的GPU的数目不固定，并且假定每个GPU的计算能力不确定。本文考虑使用动态分配任务使得计算时间达到近似最优。

 

python版本及库：python3.6，subprocess，multiprocessing, time, numpy

 

• 利用nvidia-smi获得可用的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)

注意：subprocess.getoutput返回值为字符串，需要进行类型转换。

 

• 任务分发

实现逻辑：第一次分发nDevice个任务使用全部GPU，之后进行动态分配分发。multiprocessing库apply_async支持异步发射线程，使用Pool进行线程管理，对于GPU并行任务使用nDevice大小的线程池。

这里使用两个函数实现，其中第一个函数管理进程池，控制进程发射；第二个函数实现具体的GPU代码。当控制代码传递给指定GPU代号时使用指定GPU，否则动态分配GPU（传递空值None）。

    def parallel_run(self):
        from multiprocessing import freeze_support, Pool
        import subprocess
        import time
        nDevice = int(subprocess.getoutput("nvidia-smi -L | grep GPU |wc -l"))
        freeze_support()
        pool = Pool(nDevice)
        taskList = self.sg.nodetask[self.nodei]
        print("there are {} task on this node".format(len(taskList)))
        print("the nDevice is {}".format(nDevice))
        for device_i in range(nDevice):
            task = taskList.pop(0)
            pool.apply_async(func=self.run, args=(task, device_i))
            time.sleep(3)
        for task in taskList:
            pool.apply_async(func=self.run, args=(task, None))
            time.sleep(3)
        pool.close()
        pool.join()

注意事项：

freeze_support()在windows执行时需要加入

Pool分配进程池大小

time.sleep(3)给CUDA代码cudaFree，寻找可用GPU时空出的时间。否则可能会导致很多进程发到0号GPU上。

apply_async内的args=(msgs, )，如果只有一个参数的话，逗号和括号不可少，否则会出现typeError（实测>_<）

pool.join()需在pool.close()之后

 

• 任务执行部分，接收上一个函数的参数，并将参数打包传给动态链接库。如果接收的device为指定值，那么在指定GPU上计算，否则动态分配GPU。

    def run(self, ishot, device):
        from ctypes import c_char_p, c_int, c_float, c_bool, cdll
        try:
            self.kernel = cdll.LoadLibrary('./lib/acoustic1order.so')
        except OSError:
            print('cannot open ./lib/acoustic1order.so')
        if device is None:
            device = super().available_GPU()
        kernel(args..., c_int(device))
        pass

总结：程序的调用关系为parallel_run控制并调用run的执行，run去寻找合适的GPU。这样思路较为直接，实现比较容易。