A simple framework of multiprocessing with Pool. The following code with calculate the square of 1, 2, and 3 in three different parallel threads. If the function takes multiple arguments (function(x, y, z)), arg should be list of list ([[x1, y1, z1],[x2, y2, z2],...]). product (from itertools) may be useful to generate the list of list (or tuple).
Note: To run a python multiprocessing on a cluster, note that python multiprocessing doesn't handle jobs on distinct nodes. So, jobs (those multiprocesses by python) need to run in different CPUs in the SAME NODE. To make sure this, in SLURM system, use:
(Example of using 10 CPUs to handle multiple processes at once, while allocating each CPU 32GB)
#SBATCH --cpus-per-task=10
#SBATCH --mem-per-cpu=32GB
Note: When the function requires multiple arguments, use Pool.starmap, instead of Pool.map, with a list of tuples, each containing multiple arguments.
from multiprocessing import Pool
def function(x):
y = x**2
print(y)
def multiprocess(func,arg):
nt = 3 # the number of threads
p = Pool(processes = nt)
p.map(func,arg)
if __name__ == '__main__':
multiprocess(function, [1,2,3])
