Chimie Théorique » QMX

import sys

import time

import atexit

def paropen(name, mode='r', buffering=0):

    """MPI-safe version of open function.

    In read mode, the file is opened on all nodes.  In write and

    append mode, the file is opened on the master only, and /dev/null

    is opened on all other nodes.

    """

    if rank > 0 and mode[0] != 'r':

        name = '/dev/null'

    return open(name, mode, buffering)

def parprint(*args, **kwargs):

    """MPI save print - prints only from master.

    Tries to adopt python 3 behaviour.

    """

    if rank > 0:

        return

    defaults = { 'end' : '\n',

                 'file' : sys.stdout }

    for key in defaults:

        if not key in kwargs:

            kwargs[key] = defaults[key]

    for arg in args[:-1]:

        print >> kwargs['file'], arg,

    if len(args):

        last = args[-1]

    else:

        last = ''

    if kwargs['end'] == '\n':

        print last

    else:

        print last,

# Check for special MPI-enabled Python interpreters:

if '_gpaw' in sys.modules:

    # http://wiki.fysik.dtu.dk/gpaw

    from gpaw.mpi import world

    rank = world.rank

    size = world.size

    barrier = world.barrier

elif 'asapparallel3' in sys.modules:

    # http://wiki.fysik.dtu.dk/Asap

    # We cannot import asap3.mpi here, as that creates an import deadlock

    #from asap3.mpi import world

    import asapparallel3

    world = asapparallel3.Communicator()

    rank = world.rank

    size = world.size

    barrier = world.barrier

elif 'Scientific_mpi' in sys.modules:

    #

    from Scientific.MPI import world

    rank = world.rank

    size = world.size

    barrier = world.barrier

else:

    # This is a standard Python interpreter:

    rank = 0

    size = 1

    world = None

    def barrier():

        pass

def register_parallel_cleanup_function():

    """Call MPI_Abort if python crashes.

    This will terminate the processes on the other nodes."""

    if size == 1:

        return

    def cleanup(sys=sys, time=time, world=world):

        error = getattr(sys, 'last_type', None)

        if error:

            sys.stdout.flush()

            sys.stderr.write(('ASE CLEANUP (node %d): %s occurred.  ' +

                              'Calling MPI_Abort!\n') % (world.rank, error))

            sys.stderr.flush()

            # Give other nodes a moment to crash by themselves (perhaps

            # producing helpful error messages):

            time.sleep(3)

            world.abort(42)

    atexit.register(cleanup)