Distributing jobs evenly across multiple GPUs with `multiprocessing.Pool`

Let's say that I have the following:

A system with 4 GPUs.
A function, foo, which may be run up to 2 times simultaneously on each GPU.
A list of files that need to be processed using foo in any order. However, each file takes an unpredictable amount of time to be processed.

I would like to process all the files, keeping all the GPUs as busy as possible by ensuring there are always 8 instances of foo running at any given time (2 instance on each GPU) until less than 8 files remain.

The actual details of invoking the GPU are not my issue. What I'm trying to figure out is how to write the parallelization so that I can keep 8 instances of foo running but somehow making sure that exactly 2 of each GPU ID are used at all times.

I've come up with one way to solve this problem using multiprocessing.Pool, but the solution is quite brittle and relies on (AFAIK) undocumented features. It relies on the fact that the processes within the Pool are named in the format FormPoolWorker-%d where %d is a number between one and the number of processes in the pool. I take this value and mod it with the number of GPUs and that gives me a valid GPU id. However, it would be much nicer if I could somehow give the GPU id directly to each process, perhaps on initialization, instead of relying on the string format of the process names.

One thing I considered is that if the initializer and initargs parameters of Pool.__init__ allowed for a list of initargs so that each process could be initialized with a different set of arguments then the problem would be moot. Unfortunately that doesn't appear to work.

Can anybody recommend a more robust or pythonic solution to this problem?

Hacky solution (Python 3.7):

from multiprocessing import Pool, current_process

def foo(filename):
    # Hacky way to get a GPU id using process name (format "ForkPoolWorker-%d")
    gpu_id = (int(current_process().name.split('-')[-1]) - 1) % 4

    # 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))

pool = Pool(processes=4*2)

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

Solution

I figured it out. It's actually quite simple. All we need to do is use a multiprocessing.Queue to manage the available GPU IDs. Start by initializing the Queue to contain 2 of each GPU ID, then get the GPU ID from the queue at the beginning of foo and put it back at the end.

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()