@contextmanager
def MPIContext():
    """Wrap this around code that you want to globally fail if it fails
    on any MPI process in MPI_WORLD.
    """
    try:
        yield
    except:
        exc_type, exc_val, exc_tb = sys.exc_info()
        if exc_val is not None:
            fail = True
        else:
            fail = False

        fails = MPI.COMM_WORLD.allgather(fail)

        if fail or not any(fails):
            raise exc_type, exc_val, exc_tb
        else:
            for i,f in enumerate(fails):
                if f:
                    raise RuntimeError("a test failed in (at least) rank %d" % i)
dataset/ETHPy150Open OpenMDAO/OpenMDAO-Framework/openmdao.main/src/openmdao/main/mpiwrap.py/MPIContext/original