map_parallel.py

import multiprocessing

def map_parallel(func, iterable, chunksize=1, *,
                 pool=None,
                 n_processes=None,
                 progress_bar_func=None,
                 pool_kwargs=None,
                 progress_bar_func_kwargs=None,
                 total=None
                ):
    if pool_kwargs is None:
        pool_kwargs = {}
    if progress_bar_func_kwargs is None:
        progress_bar_func_kwargs = {}

    own_pool = False
    if pool is None:
        pool = multiprocessing.Pool(n_processes, **pool_kwargs)
        own_pool = True

    try:
        if not callable(progress_bar_func):
            return pool.map(func, iterable, chunksize)
        else:
            pool_iter = pool.imap(func, iterable, chunksize)
            return [item for item in
                    progress_bar_func(pool_iter,
                                      total=total or len(iterable),
                                      **progress_bar_func_kwargs)]
    finally:
        if own_pool:
            pool.terminate()


if __name__ == '__main__':
    import time
    import sys

    import numpy as np
    from tqdm import tqdm


    def func(arg1):
        time.sleep(1.0)
        return arg1 * arg1


    arr = np.arange(0, 10)

    start_time = time.time()
    # res = [func(el) for el in arr]
    res = map_parallel(func, arr, n_processes=4, chunksize=1,
                       progress_bar_func=tqdm,
                       progress_bar_func_kwargs=dict(desc='myfunc'))

    print('{}'.format(res), file=sys.stderr)
    print('Elapsed: {}'.format(time.time() - start_time), file=sys.stderr)