Как использовать многопроцессорную очередь Python для доступа к GPU (через PyOpenCL)?

Question

У меня есть код, выполнение которого занимает много времени, и поэтому я исследовал многопроцессорную библиотеку Python, чтобы ускорить процесс. Мой код также имеет несколько шагов, которые используют GPU через PyOpenCL. Проблема заключается в том, что если я настрою одновременное выполнение нескольких процессов, все они в конечном итоге будут пытаться использовать графический процессор одновременно, и это часто приводит к тому, что один или несколько процессов выдают исключение и завершают работу.

Чтобы обойти это, я пошатнул начало каждого процесса, чтобы они не сталкивались друг с другом:

process_list = []
num_procs = 4

# break data into chunks so each process gets it's own chunk of the data
data_chunks = chunks(data,num_procs)
for chunk in data_chunks:
    if len(chunk) == 0:
        continue
    # Instantiates the process
    p = multiprocessing.Process(target=test, args=(arg1,arg2))
    # Sticks the thread in a list so that it remains accessible
    process_list.append(p)

# Start threads
j = 1
for process in process_list:
    print('\nStarting process %i' % j)
    process.start()
    time.sleep(5)
    j += 1

for process in process_list:
    process.join()

Я также обернул цикл try иначе вокруг функции, которая вызывает GPU, чтобы два процесса пытались получить к нему доступ одновременно, а тот, кто не получил доступа, подождет пару секунд и попробует снова:

wait = 2
n = 0
while True:
    try:
        gpu_out = GPU_Obj.GPU_fn(params)
    except:
        time.sleep(wait)
        print('\n Waiting for GPU memory...')
        n += 1
        if n == 5:
            raise Exception('Tried and failed %i times to allocate memory for opencl kernel.' % n)
        continue
    break

Этот обходной путь очень неуклюжий, и хотя он работает большую часть времени, процессы иногда выдают исключения, и я чувствую, что должно быть более эффективное / элегантное решение, использующее multiprocessing.queue или что-то подобное. Тем не менее, я не уверен, как интегрировать его с PyOpenCL для доступа к GPU.

Answer 1

Похоже, вы могли бы использовать multiprocessing.Lock для синхронизации доступа к графическому процессору:

data_chunks = chunks(data,num_procs)
lock = multiprocessing.Lock()
for chunk in data_chunks:
    if len(chunk) == 0:
        continue
    # Instantiates the process
    p = multiprocessing.Process(target=test, args=(arg1,arg2, lock))
    ...

Тогда внутри test где вы получаете доступ к GPU:

with lock:  # Only one process will be allowed in this block at a time.
    gpu_out = GPU_Obj.GPU_fn(params)

Редактировать:

Чтобы сделать это с пулом, вы должны сделать это:

# At global scope
lock = None

def init(_lock):
    global lock
    lock = _lock

data_chunks = chunks(data,num_procs)
lock = multiprocessing.Lock()
for chunk in data_chunks:
    if len(chunk) == 0:
        continue
    # Instantiates the process
    p = multiprocessing.Pool(initializer=init, initargs=(lock,))
    p.apply(test, args=(arg1, arg2))
    ...

Или же:

data_chunks = chunks(data,num_procs)
m = multiprocessing.Manager()
lock = m.Lock()
for chunk in data_chunks:
    if len(chunk) == 0:
        continue
    # Instantiates the process
    p = multiprocessing.Pool()
    p.apply(test, args=(arg1, arg2, lock))