Multiprocessing in python

parallel.py

import multiprocessing
from multiprocessing import Pool
from tqdm import tqdm
import functools
import time


def square_and_add(x, y=5):
    # function to multiprocess
    time.sleep(0.1)
    return x*x + y


def multi_arg(x, y, z, w=1, a=0.01):
    time.sleep(0.1)
    return w*x + w*y + z*a


def z_first_multi_arg(z, x, y,  **kwargs):
    # stupid hack to get it to work for partial 
    return multi_arg(x, y, z, **kwargs)


def y_first_multi_arg(y, x, z, **kwargs):
    return multi_arg(x, y, z, **kwargs)


def cpu_count():
    return multiprocessing.cpu_count()


def parallel(function, args, n_cores):
    with Pool(processes=n_cores) as p:
        results = list(tqdm(p.imap(function, args), total=len(args)))

    return results
        

if __name__ == "__main__":
    n_cores = cpu_count() - 1 or 1 
    # run first task
    xs = [1, 2, 3, 4, 5, 6, 7]

    print("calc square and add...")
    results = parallel(square_and_add, xs, n_cores)
    print(results)

    # now modify keyword arg 
    new_y = 3
    m_square_and_add = functools.partial(square_and_add, y=new_y)

    print("calc modified square and add...")
    m_results = parallel(m_square_and_add, xs, n_cores)
    print(m_results)

    # now many args and kwargs but only want to parallelise one 
    ys = [x*2 for x in xs]
    zs = [x*1.5 for x in xs]

    # first fix y, z 
    vary_x_multi_arg = functools.partial(multi_arg, y=ys[0], z=zs[0])
    print("calc multi arg for fixed y, z...")
    vary_x_results = parallel(vary_x_multi_arg, xs, n_cores)
    print(vary_x_results)


    # THESE ONES INVOLVE FIXING THE FIRST PARAM WHICH SEEMS TO BE A PROBLEM
    # fix x, y 
    vary_z_multi_arg = functools.partial(z_first_multi_arg, x=xs[0], y=ys[0])
    print("calc multi arg for fixed x, y...")
    vary_z_results = parallel(vary_z_multi_arg, zs, n_cores)
    print(vary_z_results)

    # fix x, z
    vary_y_multi_arg = functools.partial(y_first_multi_arg, x=xs[0], z=zs[0])
    print("calc multi arg for fixed x, z...")
    vary_y_results = parallel(vary_y_multi_arg, ys, n_cores)
    print(vary_y_results)