在python multiprocess.spawn中传递对象

Question

我有50个进程，我想并行运行。我需要在gpu上运行进程。我的机器有8个gpus，我把设备编号传递给每个进程，这样它就知道在什么设备上运行。一旦该进程完成，我希望在该设备上运行另一个进程。使用下面的命令使用POpen将进程作为子进程运行

python special_process.py device

一个简单的方法就是

for group in groups:
    processes = [subprocess.POpen(f'python special_process.py {device}'.split()) for device in range(8)]
    [p.wait() for p in process]

其中的groups，是将50个进程分成8组。

这样做的缺点是，有些进程花费的时间比其他进程要长，所有进程都需要在转移到下一组之前完成。

我本来希望做一些类似multiprocess.spawn的事情，但我需要最后一个进程来返回设备编号，这样就可以清楚地知道哪个设备是打开的，可以在上面运行。我尝试从多进程中使用Queue和Process，但不能一次运行多个进程。

任何帮助都将不胜感激。谢谢

Answer 1

简单的while循环和构建您自己的队列是有效的。只是不要等到最后再用。

import subprocess

d = list(range(20))
num_gpus = 8
procs = []
gpus_free = set([j for j in range(num_gpus)])
gpus_used = set()

while len(d) > 0:
    for proc, gpu in procs:
        poll = proc.poll()
        if poll is None:
            # Proc still running
            continue
        else:
            # Proc complete - pop from list
            procs.remove((proc, gpu))
            gpus_free.add(gpu)

    # Submit new processes
    if len(procs) < num_gpus:
        this_process = d.pop()
        gpu_for_this_process = gpus_free.pop()
        command = f"python3 inner_function.py {gpu_for_this_process} {this_process}"
        proc = subprocess.Popen(command, shell= True)
        procs.append((proc, gpu_for_this_process))

[proc.wait() for proc, _ in procs]
print('DONE with all')