Равномерное распределение заданий по нескольким графическим процессорам с помощью `multiprocessing.Pool` - PullRequest
Новая
       13

Равномерное распределение заданий по нескольким графическим процессорам с помощью `multiprocessing.Pool`

0 голосов
/ 22 ноября 2018

Допустим, у меня есть следующее:

  • Система с 4 графическими процессорами.
  • Функция foo, которая может запускаться до 2 раз одновременно на каждомGPU.
  • Список files, который необходимо обработать, используя foo в любом порядке.Однако для обработки каждого файла требуется непредсказуемое количество времени.

Я хотел бы обработать все файлы, поддерживая все графические процессоры как можно более занятыми, гарантируя, что всегда есть 8 экземпляров foo запуск в любое время (2 экземпляра на каждом графическом процессоре) до тех пор, пока не останется менее 8 файлов.

Фактические детали вызова графического процессора не являются моей проблемой.Я пытаюсь понять, как записать распараллеливание, чтобы я мог сохранить 8 экземпляров foo в рабочем состоянии, но каким-то образом убедиться, что ровно 2 из каждого идентификатора GPU используются постоянно.

IМы придумали один способ решения этой проблемы, используя multiprocessing.Pool, но решение довольно хрупкое и опирается на (AFAIK) недокументированные функции.Он основан на том факте, что процессы в Pool имеют имена в формате FormPoolWorker-%d, где %d - это число между одним и числом процессов в пуле.Я беру это значение и модифицирую его количеством графических процессоров, и это дает мне действительный идентификатор графического процессора.Однако было бы гораздо приятнее, если бы я мог как-то дать идентификатор GPU непосредственно каждому процессу, возможно, при инициализации, вместо того, чтобы полагаться на строковый формат имен процессов.

Одна вещь, которую я рассмотрел, заключается в том, что если The initializer и initargs параметры Pool.__init__ разрешенный список initargs так, что каждый процесс может быть инициализирован с другим набором аргументов, то проблема была бы спорной.К сожалению, это не работает.

Кто-нибудь может порекомендовать более надежное или питонное решение этой проблемы?

Хакерское решение (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()

1 Ответ

0 голосов
/ 22 ноября 2018

Я понял это.Это на самом деле довольно просто.Все, что нам нужно сделать, это использовать multiprocessing.Queue для управления доступными идентификаторами GPU.Начните с инициализации Queue для содержания 2 каждого идентификатора графического процессора, затем get идентификатор графического процессора из queue в начале 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()
...