Python multiprocessing: reduce during map?

Question

Is there a way to reduce memory consumption when working with Python's pool.map ?

To give a short example: worker() does some heavy lifting and returns a larger array...

def worker():
    # cpu time intensive tasks
    return large_array

...and a Pool maps over some large sequence:

with mp.Pool(mp.cpu_count()) as p:
    result = p.map(worker, large_sequence)

Considering this setup, obviously, result will allocate a large portion of the system's memory. However, the final operation on the result is:

    final_result = np.sum(result, axis=0)

Thus, NumPy effectively does nothing else than reducing with a sum operation on the iterable:

    final_result = reduce(lambda x, y: x + y, result)

This, of course, would make it possible to consume results of pool.map as they come in and garbage-collecting them after reducing to eliminate the need of storing all the values first.

I could write some mp.queue now where results go into and then write some queue-consuming worker that sums up the results but this would (1) require significantly more lines of code and (2) feel like a (potentially slower) hack-around to me rather than clean code.

Is there a way to reduce results returned by a mp.Pool operation directly as they come in?

Answer 1

The iterator mappers imap and imap_unordered seem to do the trick:

#!/usr/bin/env python3

import multiprocessing
import numpy as np

def worker( a ):
    # cpu time intensive tasks
    large_array = np.ones((20,30))+a
    return large_array


if __name__ == '__main__':
    
    arraysum = np.zeros((20,30))
    large_sequence = range(20)
    num_cpus = multiprocessing.cpu_count()    
    
    with multiprocessing.Pool( processes=num_cpus ) as p:
        for large_array in p.imap_unordered( worker, large_sequence ):
            arraysum += large_array