Question

假设我具有以下条件：

具有4个GPU的系统。
foo函数，在每个GPU上最多可同时运行2次。
需要使用files以任意顺序处理的foo列表。但是，每个文件的处理时间都无法预测。

我想处理所有文件，通过确保在任何给定时间始终运行8个foo实例（每个GPU上有2个实例），直到少于8个文件，使所有GPU尽可能繁忙。留下来。

调用GPU的实际细节不是我的问题。我要弄清楚的是如何编写并行化，以便使foo的8个实例保持运行，但要确保始终使用每个GPU ID中的2个。

我已经提出了一种使用multiprocessing.Pool解决此问题的方法，但是该解决方案非常脆弱，并且依赖于（AFAIK）未记录的功能。它依赖于这样一个事实，Pool中的进程以FormPoolWorker-%d的格式命名，其中%d是一个与池中的进程数之间的数字。我采用此值并使用GPU的数量对其进行修改，这为我提供了有效的GPU ID。但是，如果我能以某种方式直接给每个进程GPU ID（也许在初始化时），而不是依赖进程名称的字符串格式，那就更好了。

我考虑的一件事是，如果如果 initializer的{{1}}和initargs参数允许使用Pool.__init__的列表，那么每个过程可以使用不同的参数集进行初始化，那么问题就不存在了。不幸的是，这似乎不起作用。

有人可以为这个问题推荐更健壮或Pythonic的解决方案吗？

棘手的解决方案（Python 3.7）：

initargs

Answer 1

我知道了。实际上很简单。我们需要做的就是使用multiprocessing.Queue来管理可用的GPU ID。首先初始化Queue以包含每个GPU ID中的2个，然后在get的开头queue来foo的GPU ID并将其返回put最后。

from multiprocessing import Pool, current_process, Queue

NUM_GPUS = 4
PROC_PER_GPU = 2    

queue = Queue()

def foo(filename):
    gpu_id = queue.get()
    try:
        # run processing on GPU <gpu_id>
        ident = current_process().ident
        print('{}: starting process on GPU {}'.format(ident, gpu_id))
        # ... process filename
        print('{}: finished'.format(ident))
    finally:
        queue.put(gpu_id)

# initialize the queue with the GPU ids
for gpu_ids in range(NUM_GPUS):
    for _ in range(PROC_PER_GPU):
        queue.put(gpu_ids)

pool = Pool(processes=PROC_PER_GPU * NUM_GPUS)
files = ['file{}.xyz'.format(x) for x in range(1000)]
for _ in pool.imap_unordered(foo, files):
    pass
pool.close()
pool.join()

使用`multiprocessing.Pool`在多个GPU之间平均分配作业

1 个答案: