Context manager for mpi4py that selectively executes its body on the root node

mpiutils.py

import sys
import inspect

def mpi_rank(comm=None):
    """Get the rank of the curent MPI node

    Parameters
    ----------
    comm : mpi communicator, optional
        The MPI communicator. By default, we use the COMM_WORLD

    Returns
    -------
    rank : int
        The rank of the current mpi process. The return value is 0
        (the root node) if MPI is not running
    """
    if comm is None:
        try:
            from mpi4py import MPI
            comm = MPI.COMM_WORLD
            rank = comm.Get_rank()
        except ImportError:
            rank = 0
    else:
        rank = comm.Get_rank()

    return rank

class SelectiveExecution(object):
    """Contect manager that executes body only under a certain
    condition.

    http://stackoverflow.com/questions/12594148/skipping-execution-of-with-block
    """

    def __init__(self, skip=False):
        self.skip = skip

    def __enter__(self):
        if self.skip:
            # Do some magic
            sys.settrace(lambda *args, **keys: None)
            frame = inspect.currentframe(1)
            frame.f_trace = self.trace

    def trace(self, frame, event, arg):
        raise

    def __exit__(self, type, value, traceback):
        return True


class mpi_root(SelectiveExecution):
    """Context manager that selectively executes its body on the root node"""
    def __init__(self, comm=None):
        skip_if = (mpi_rank() != 0)
        super(mpi_root, self).__init__(skip_if)

def main():
    from mpi4py import MPI
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()

    a = None
    with mpi_root():
        print 'EXECUTING ON RANK', rank
        a = [1,2,3,4]
    a = comm.bcast(a)

    print 'RANK %s, a=%s' % (rank, a)


if __name__ == '__main__':
    main()