How to use multiprocessing pool.map with multiple arguments?

前端 未结 20 3448
-上瘾入骨i
-上瘾入骨i 2020-11-21 11:24

In the Python multiprocessing library, is there a variant of pool.map which supports multiple arguments?

text = "test"
def         


        
20条回答
  •  [愿得一人]
    2020-11-21 11:48

    The answer to this is version- and situation-dependent. The most general answer for recent versions of Python (since 3.3) was first described below by J.F. Sebastian.1 It uses the Pool.starmap method, which accepts a sequence of argument tuples. It then automatically unpacks the arguments from each tuple and passes them to the given function:

    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', ...
    

    For earlier versions of Python, you'll need to write a helper function to unpack the arguments explicitly. If you want to use with, you'll also need to write a wrapper to turn Pool into a context manager. (Thanks to muon for pointing this out.)

    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', ...
    

    In simpler cases, with a fixed second argument, you can also use partial, but only in 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. Much of this was inspired by his answer, which should probably have been accepted instead. But since this one is stuck at the top, it seemed best to improve it for future readers.

提交回复
热议问题