Search code examples
pythonmultiprocessingpoolmap-function

Passing multiple parameters to pool.map() function in Python


I need some way to use a function within pool.map() that accepts more than one parameter. As per my understanding, the target function of pool.map() can only have one iterable as a parameter but is there a way that I can pass other parameters in as well? In this case, I need to pass in a few configuration variables, like my Lock() and logging information to the target function.

I have tried to do some research and I think that I may be able to use partial functions to get it to work? However I don't fully understand how these work. Any help would be greatly appreciated! Here is a simple example of what I want to do:

def target(items, lock):
    for item in items:
        # Do cool stuff
        if (... some condition here ...):
            lock.acquire()
            # Write to stdout or logfile, etc.
            lock.release()

def main():
    iterable = [1, 2, 3, 4, 5]
    pool = multiprocessing.Pool()
    pool.map(target(PASS PARAMS HERE), iterable)
    pool.close()
    pool.join()

Solution

  • You can use functools.partial for this (as you suspected):

    from functools import partial
    
    def target(lock, iterable_item):
        for item in iterable_item:
            # Do cool stuff
            if (... some condition here ...):
                lock.acquire()
                # Write to stdout or logfile, etc.
                lock.release()
    
    def main():
        iterable = [1, 2, 3, 4, 5]
        pool = multiprocessing.Pool()
        l = multiprocessing.Lock()
        func = partial(target, l)
        pool.map(func, iterable)
        pool.close()
        pool.join()
    

    Example:

    def f(a, b, c):
        print("{} {} {}".format(a, b, c))
    
    def main():
        iterable = [1, 2, 3, 4, 5]
        pool = multiprocessing.Pool()
        a = "hi"
        b = "there"
        func = partial(f, a, b)
        pool.map(func, iterable)
        pool.close()
        pool.join()
    
    if __name__ == "__main__":
        main()
    

    Output:

    hi there 1
    hi there 2
    hi there 3
    hi there 4
    hi there 5