使用“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, QueueNUM_GPUS = 4PROC_PER_GPU = 2&nbsp; &nbsp;&nbsp;queue = Queue()def foo(filename):&nbsp; &nbsp; gpu_id = queue.get()&nbsp; &nbsp; try:&nbsp; &nbsp; &nbsp; &nbsp; # run processing on GPU <gpu_id>&nbsp; &nbsp; &nbsp; &nbsp; ident = current_process().ident&nbsp; &nbsp; &nbsp; &nbsp; print('{}: starting process on GPU {}'.format(ident, gpu_id))&nbsp; &nbsp; &nbsp; &nbsp; # ... process filename&nbsp; &nbsp; &nbsp; &nbsp; print('{}: finished'.format(ident))&nbsp; &nbsp; finally:&nbsp; &nbsp; &nbsp; &nbsp; queue.put(gpu_id)# initialize the queue with the GPU idsfor gpu_ids in range(NUM_GPUS):&nbsp; &nbsp; for _ in range(PROC_PER_GPU):&nbsp; &nbsp; &nbsp; &nbsp; 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):&nbsp; &nbsp; passpool.close()pool.join()