使用“multiprocessing.Pool”在多个 GPU 上均匀分配作业

Question

假设我有以下内容：具有 4 个 GPU 的系统。一个函数，foo在每个 GPU 上最多可以同时运行 2 次。files需要foo按任何顺序处理的列表。但是，每个文件都需要不可预测的时间来处理。我想处理所有文件，通过确保foo在任何给定时间始终有 8 个运行实例（每个 GPU 上 2 个实例），直到剩下的文件少于 8 个，从而使所有 GPU 尽可能忙碌。调用 GPU 的实际细节不是我的问题。我想弄清楚的是如何编写并行化，以便我可以保持 8 个foo运行实例，但以某种方式确保始终使用每个 GPU ID 中的 2 个。我想出了一种使用 解决这个问题的方法multiprocessing.Pool，但该解决方案非常脆弱，并且依赖于（AFAIK）未记录的功能。它依赖于一个事实，即内部流程Pool的格式命名FormPoolWorker-%d，其中%d是一个与池中的进程数之间的数字。我取这个值并用 GPU 的数量修改它，这给了我一个有效的 GPU id。但是，如果我能以某种方式将 GPU id 直接提供给每个进程（可能是在初始化时），而不是依赖于进程名称的字符串格式，那就更好了。我考虑过的一件事是，如果 允许一个列表的initializer和initargs参数，以便每个进程都可以使用不同的参数集进行初始化，那么问题就没有实际意义了。不幸的是，这似乎不起作用。Pool.__init__initargs任何人都可以推荐一个更健壮或pythonic的解决方案来解决这个问题吗？

Answer 1

我想通了。其实很简单。我们需要做的就是使用 amultiprocessing.Queue来管理可用的 GPU ID。通过初始化开始Queue到含有2每个GPU ID的，然后get从所述GPU IDqueue之初foo和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()