Python Pool Multiprocessing Poor CPU Usage

Question

I have a bunch of independent N body sims I want to run in parallel in python. The walltime for individual sims is going to vary dramatically depending on the parameters of the bodies in the sims. It seemed like the best way to do this would be to build pool of processes with the multiprocessing module, give them the sim jobs with the starmap() function, and have them save the results to separate files based on the process ID. However, I've getting awful parallel performance. There is no speedup between 2 and 4 processes (I have 4 CPU on my laptop) and the unix time utility seems to think that the CPU usage percentage is ~150% which is terrible. Below is my code:

import rebound
import numpy as np
import multiprocessing as mp

def two_orbits_one_pool(orbit1, orbit2):
    #######################################
    print('process number', mp.current_process().name)
    #######################################
    # build simulation
    sim = rebound.Simulation()
    # add sun
    sim.add(m=1.)
    # add two overlapping orbits
    sim.add(primary=sim.particles[0], m=orbit1['m'], a=orbit1['a'], e=orbit1['e'], inc=orbit1['i'], \
            pomega=orbit1['lop'], Omega=orbit1['lan'], M=orbit1['M'])
    sim.add(primary=sim.particles[0], m=orbit2['m'], a=orbit2['a'], e=orbit2['e'], inc=orbit2['i'], \
            pomega=orbit2['lop'], Omega=orbit2['lan'], M=orbit2['M'])
    sim.move_to_com()
    # integrate for 10 orbits of orbit1
    P = 2.*np.pi * np.sqrt(orbit1['a']**3)
    sim.automateSimulationArchive("archive-{}.bin".format(mp.current_process().name), interval=P)
    sim.integrate(10.*P)

if __name__ == "__main__":
    # orbit definitions
    N_M = 10
    N_lop = 10
    m = 1e-6
    a, e = 1., 0.3
    inc, lop, lan = 0., 0., 0.
    M = np.linspace(0., 2*np.pi, endpoint=False, num=N_M)
    dlop = np.linspace(0., 0.05, num=N_lop)
    # orbit dictionaries
    args = []
    for i in range(dlop.shape[0]):
        for j in range(M.shape[0]):
            for k in range(M.shape[0]):
                args.append( ( {'m':m, 'a':a, 'e':e, 'i':inc, \
                                'lop':lop, 'lan':lan, 'M':M[j]},
                               {'m':m, 'a':a, 'e':e, 'i':inc, \
                                'lop':lop+dlop[i], 'lan':lan, 'M':M[k]} ) )
    # fill the pool with orbit jobs
    with mp.Pool() as pool:
        pool.starmap(two_orbits_one_pool, args)

Could someone explain why this is performing so poorly? I'm much more used to OpenMP and MPI; I'm not that familiar with parallel programming in Python. Overall, I've been quite disappointed in the multiprocessing module. I think maybe I should try using the numba module instead?

EDIT: In response to Roland Smith's response, I profiled the integration and save time for my code. Here is a stripplot showing the results. As you can see, both Roland Smith's and J_H's suggestions were true. There is a subset of initial conditions that result in extremely long integration times due to close encounters between the bodes. However, in general, the save time was about 5 times longer than the integration time. The job suffers from stragglers and is disk i/o bound.

Answer 1

If there is no discernable speedup, then probably your code is not CPU-bound.

In general, writing to a disk (even an SSD) is much slower than running code on the CPU. If several worker processes are writing significant amounts of data to disk, that might be the bottleneck.

To diagnose the problem, you have to measure.

You should separate the calculations from the saving of the data; e.g. run sim.integrate() followed by sim.simulationarchive_snapshot() 10 times, and sandwich each of those calls between time.monotonic() calls. Then return the average time of the integration step and the snapshot steps as shown below.

import time

def two_orbits_one_pool(orbit1, orbit2):
    #######################################
    print('process number', mp.current_process().name)
    #######################################
    # build simulation
    sim = rebound.Simulation()
    # add sun
    sim.add(m=1.)
    # add two overlapping orbits
    sim.add(primary=sim.particles[0], m=orbit1['m'], a=orbit1['a'], e=orbit1['e'], inc=orbit1['i'], \
            pomega=orbit1['lop'], Omega=orbit1['lan'], M=orbit1['M'])
    sim.add(primary=sim.particles[0], m=orbit2['m'], a=orbit2['a'], e=orbit2['e'], inc=orbit2['i'], \
            pomega=orbit2['lop'], Omega=orbit2['lan'], M=orbit2['M'])
    sim.move_to_com()
    # integrate for 10 orbits of orbit1
    P = 2.*np.pi * np.sqrt(orbit1['a']**3)
    arname = "archive-{}.bin".format(mp.current_process().name)
    itime, stime = 0.0, 0.0
    for k in range(10):
        start = time.monotonic()
        sim.integrate(P)
        itime += time.monotonic() - start
        start = time.monotonic()
        sim.simulationarchive_snapshot(arname)
        stime += time.monotonic() - start
    return (mp.current_process().name, itime/10, stime/10)

# Run the calculations
with mp.Pool() as pool:
    data = pool.starmap(two_orbits_one_pool, args)
# Print the times that it took.
for name, itime, stime in data:
   print(f"worker {name}: itime {itime} s, stime {stime} s")

That should tell you what the bottleneck is.

Possible solutions if writing to disk is the bottleneck;

• Use an SSD to store the simulation results.
• Use a RAM-disk to store the simulation results. (Although compared to an SSD not a huge performance boost.)
• Check if you can tune your OS for maximum write performance.

Edit1: Given your measurement result, the obvious performance improvement is to save less often.

Another option that might be worth looking at is staggering the writes. That only makes sense if there is significant overlap between the writes from different processes, and if those concurrent writes can saturate the disk I/O subsystem. So you'd have to measure that first. If there is overlap, create a Lock object in the parent process. Then acquire the lock before (explicitly) saving, and release it after. This won't work with automateSimulationArchive.

A last option is to write your own save function using mmap. Using mmap is somewhat clunky compared to normal file handling in Python. But it can significantly improve performance. However I am unsure that the gains justify the effort in this case.

Answer 2

The straggler effect can have a big impact on such jobs.

straggler effect

Suppose you have N tasks for N cores, and each task has a different duration. Order by duration to find min_time and max_time. All N cores will be busy up through min_time, but then they go idle, one by one. Just before max_time, only a single "straggler" core is being used.

predictions

If you can make a decent guess about task duration beforehand, use that to sort them in descending order. For T tasks > N cores, schedule the long tasks first. Then N tasks run for a while, the shortest of those completes, and the now-idle core picks up a task of "medium" duration. By the time we get to the T-th task, each core has a random amount of work still to do, and we're scheduling a "short" task. So cores are mostly busy doing useful work, right up till near the end.

logging

If you cannot make a useful duration estimate a priori, at least record the start times and durations. Use that to figure out whether stragglers are causing you grief, or if it's something else like L3 cache thrashing.