Search code examples
pythonmpimpi4py

Pre- and Post- Processing with mpi4py?


I am wanting to write a script that would look like the following:

Execute a pre-processing function (only 1 process should do this; all others wait until this is done)
--------
All processes then do the main function in parallel
--------
Finally, after all processes finish, execute a post-processing function

I tried:

from mpi4py import MPI
*pre-processing
MPI.Init()
*parallel script
MPI.Finalize()
*post-processing

But all cores are running all parts of the script


Solution

  • I eventually found a solution that worked for me. rank 0 is acting as a manager that pre-processes, then delegates the main script processing to all the other ranks, then post-processes. It works like this:

    if rank == 0:
        pre_processing()
        for i in range(1,size):
            comm.sendrecv('pre-processing finished; ask for work', dest=i)
        while (work_remains):
            source = comm.recv(source=MPI.ANY_SOURCE)
            comm.send(work_remains.pop(), dest=source)
        for i in range(1,size):
            comm.send('no more work', dest=i)
            comm.recv(source=i)
        post_processing()
    
    else:
        comm.recv(source=0) # Blocker; wait for rank 0 to finish
        comm.send(rank, dest=0) # Ask for work
        main_script_processing()