How to use multiprocessing pool.map with multiple arguments

В multiprocessing библиотеке Python есть ли вариант pool.map, который поддерживает несколько аргументов?

import multiprocessing

text = "test"

def harvester(text, case):
    X = case[0]
    text + str(X)

if __name__ == '__main__':
    pool = multiprocessing.Pool(processes=6)
    case = RAW_DATASET
    pool.map(harvester(text, case), case, 1)
    pool.close()
    pool.join()

существует ли вариант pool.map, который поддерживает несколько аргументов?

Python 3.3 включает в себя pool.starmap() метод:

#!/usr/bin/env python3
from functools import partial
from itertools import repeat
from multiprocessing import Pool, freeze_support

def func(a, b):
    return a + b

def main():
    a_args = [1,2,3]
    second_arg = 1
    with Pool() as pool:
        L = pool.starmap(func, [(1, 1), (2, 1), (3, 1)])
        M = pool.starmap(func, zip(a_args, repeat(second_arg)))
        N = pool.map(partial(func, b=second_arg), a_args)
        assert L == M == N

if __name__=="__main__":
    freeze_support()
    main()

Для более старых версий:

#!/usr/bin/env python2
import itertools
from multiprocessing import Pool, freeze_support

def func(a, b):
    print a, b

def func_star(a_b):
    """Convert `f([1,2])` to `f(1,2)` call."""
    return func(*a_b)

def main():
    pool = Pool()
    a_args = [1,2,3]
    second_arg = 1
    pool.map(func_star, itertools.izip(a_args, itertools.repeat(second_arg)))

if __name__=="__main__":
    freeze_support()
    main()

Вывод

1 1
2 1
3 1

Обратите внимание, как здесь используются itertools.izip() и itertools.repeat().

Из-за ошибки, упомянутой @unutbu, вы не можете использовать functools.partial() или аналогичные возможности на Python 2.6, поэтому простая функция-оболочка func_star() должна быть определена явно. Смотрите также обходной путь предложенный uptimebox.

Ответ на этот вопрос зависит от версии и конкретной ситуации. Наиболее общий ответ для последних версий Python (начиная с 3.3) был впервые описан ниже Дж.Ф. Себастьяном.¹ Он использует Pool.starmap метод, который принимает последовательность кортежей аргументов. Затем он автоматически распаковывает аргументы из каждого кортежа и передает их данной функции:

import multiprocessing
from itertools import product

def merge_names(a, b):
    return '{} & {}'.format(a, b)

if __name__ == '__main__':
    names = ['Brown', 'Wilson', 'Bartlett', 'Rivera', 'Molloy', 'Opie']
    with multiprocessing.Pool(processes=3) as pool:
        results = pool.starmap(merge_names, product(names, repeat=2))
    print(results)

# Output: ['Brown & Brown', 'Brown & Wilson', 'Brown & Bartlett', ...

Для более ранних версий Python вам нужно будет написать вспомогательную функцию для явного распаковывания аргументов. Если вы хотите использовать with, вам также нужно будет написать оболочку для превращения Pool в контекстный менеджер. (Спасибо muon за указание на это.)

import multiprocessing
from itertools import product
from contextlib import contextmanager

def merge_names(a, b):
    return '{} & {}'.format(a, b)

def merge_names_unpack(args):
    return merge_names(*args)

@contextmanager
def poolcontext(*args, **kwargs):
    pool = multiprocessing.Pool(*args, **kwargs)
    yield pool
    pool.terminate()

if __name__ == '__main__':
    names = ['Brown', 'Wilson', 'Bartlett', 'Rivera', 'Molloy', 'Opie']
    with poolcontext(processes=3) as pool:
        results = pool.map(merge_names_unpack, product(names, repeat=2))
    print(results)

# Output: ['Brown & Brown', 'Brown & Wilson', 'Brown & Bartlett', ...

В более простых случаях, с фиксированным вторым аргументом, вы также можете использовать partial, но только в Python 2.7+.

import multiprocessing
from functools import partial
from contextlib import contextmanager

@contextmanager
def poolcontext(*args, **kwargs):
    pool = multiprocessing.Pool(*args, **kwargs)
    yield pool
    pool.terminate()

def merge_names(a, b):
    return '{} & {}'.format(a, b)

if __name__ == '__main__':
    names = ['Brown', 'Wilson', 'Bartlett', 'Rivera', 'Molloy', 'Opie']
    with poolcontext(processes=3) as pool:
        results = pool.map(partial(merge_names, b='Sons'), names)
    print(results)

# Output: ['Brown & Sons', 'Wilson & Sons', 'Bartlett & Sons', ...

^{1. Многое из этого было вдохновлено его ответом, который, вероятно, следовало принять вместо этого. Но поскольку этот застрял вверху, казалось, что лучше всего улучшить его для будущих читателей.}

Я думаю, что ниже будет лучше:

def multi_run_wrapper(args):
   return add(*args)

def add(x,y):
    return x+y

if __name__ == "__main__":
    from multiprocessing import Pool
    with Pool(4) as pool:
        results = pool.map(multi_run_wrapper,[(1,2),(2,3),(3,4)])
    print results

Вывод

[3, 5, 7]

Использование Python 3.3+ с pool.starmap():

from multiprocessing.dummy import Pool as ThreadPool 

def write_func(i, x):
    print(i, "---", x)

a = ["1","2","3"]
b = ["4","5","6"] 

pool = ThreadPool(2)
pool.starmap(write_func, zip(a,b)) 
pool.close() 
pool.join()

Результат:

1 --- 4
2 --- 5
3 --- 6

Вы также можете использовать zip() с дополнительными аргументами, если хотите: zip(a,b,c,d,e)

На случай, если вы хотите, чтобы в качестве аргумента передавалось постоянное значение:

import itertools

zip(itertools.repeat(constant), a)

На случай, если ваша функция должна вернуть что-то:

results = pool.starmap(write_func, zip(a,b))

Это выдает список с возвращаемыми значениями.