How to use multiprocessing for a Python for loop with a break statement?

Question

How to parallelize a for loop which has an if condition in it. If that condition is met, then there is no need to continue with the loop. It would be great if I could use multiprocessing for this.

for i in xrange(N):
    x = do_something_with()
    if x == 0:
        break

Can the above code be parallelized in Python?

Answer 1

You can terminate a multiprocessing.Pool while jobs are active and it will kill the child processes. If you can generate your parameters in advance, imap_unordered can be used to pull in results and terminate the pool when the condition is met. Subprocesses processing other jobs in the pool are killed so they will not return other results.

import multiprocessing as mp
import time

def worker(x):
    print('work item', x)
    time.sleep(x)
    result = x - 5
    if result == 0:
        print('termination condition')
    print('work item', x, 'done')
    return result

if __name__ == '__main__':
    p = mp.Pool(4)
    for result in p.imap_unordered(worker, range(20), chunksize=1):
        if result == 0:
            print('terminating')
            p.terminate()
            break
    print('done')

Results in

work item 0
work item 1
work item 0 done
work item 4
work item 3
work item 2
work item 1 done
work item 5
work item 2 done
work item 6
work item 3 done
work item 7
work item 4 done
work item 8
termination condition
work item 5 done
work item 9
terminating
done

Notice that some jobs were started but not completed.