CODEWITHSUNDEEP
python-3.xnested-loopspython-multiprocessing
Spinor8
Spinor8

Reputation: 1607

Converting nested loop Python code to use multiprocessing pool

I have some code that has the following structure: a function that takes in nested range of parameters. For MVCE purposes, it is sufficient to consider the following

def func(x, y):
    return x*x, x+(5*y)

results = []
for a in range(3):
    for b in range(2,5):
        temp_results = func(a,b)
        results.append(temp_results)

sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
print(sorted_results)

I am trying to use a multiprocessing pool to run the func(a, b) part. I want to keep the above structure as much as possible.

Here's what I have so far based on multiprocessing pool examples I have seen. The map arguments don't work as they expect chunk-size in there. Also, how do I collect the pool results in a list to be sorted later? Is there a way to do it in a very similar way to the above single-threaded case?

a_params = []
b_params = []
for a in range(3):
    for b in range(2,5):
        a_params.append(a)
        b_params.append(b)

from multiprocessing import Pool

def pool_handler():
    p = Pool(2)
    p.map(func, a_params, b_params)

if __name__ == '__main__':
    pool_handler()

Upvotes: 1

Views: 1515

Answers (2)

hhsm95
hhsm95

Reputation: 311

You only needed to save the results and finally order them.

import multiprocessing as mp


def func(vals):
    x, y = vals
    return x*x, x+(5*y)


def main():
    input_data = []
    for a in range(3):
        for b in range(2, 5):
            input_data.append([a, b])

    n_threads = 4
    pool = mp.Pool
    with pool(n_threads) as p:
        results = p.map(func, input_data)

    sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
    print(sorted_results)


if __name__ == '__main__':
    main()

Upvotes: 1

Tim Peters
Tim Peters

Reputation: 70582

map is for a function with a single argument. If you want to pass multiple arguments, use starmap (or rewrite func to take a tuple of arguments and pass argument tuples to map).

There's nothing extra you need to do to put the results in a list. map/starmap automatically accumulate the results into a list.

So, in all, this works:

import multiprocessing as mp

def func(x, y):
    print(x*x, x+(5*y))
    return x*x, x+(5*y)

def main():
    with mp.Pool(2) as pool:
        results = pool.starmap(func,
                               [(a, b) for a in range(3)
                                       for b in range(2, 5)])
    sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
    print(sorted_results)

if __name__ == '__main__':
    main()

and displays:

0 10
0 15
0 20
1 11
1 16
1 21
4 12
4 17
4 22
[(4, 22), (1, 21), (0, 20), (4, 17), (1, 16), (0, 15), (4, 12), (1, 11), (0, 10)]

Upvotes: 1

Related Questions