File size: 695 Bytes
ca7299e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | """
Utilities for multiprocessing.
"""
import multiprocessing as mp
import numpy as np
def mp_map(func, args, n_cpu=None, verbose=True):
"""
Multiprocessing wrapper for map function.
"""
if n_cpu is None:
n_cpu = mp.cpu_count()
pool = mp.Pool(n_cpu)
if verbose:
print('Running on {} CPUs'.format(n_cpu))
result = pool.map(func, args)
pool.close()
pool.join()
return result
def parse_mp_result(result_list, sort_fn=None):
"""Parse output from mp_map"""
if sort_fn is None:
return result_list
result = sorted(result_list, key=sort_fn, ) # ascending
return result
# unit test
if __name__ == '__main__':
pass |