spike.util.mpiutil — spike 0.99.21-490 documentation

# Source code for spike.util.mpiutil

    #!/usr/bin/env python
    # encoding: utf-8

    """
    mpiutil.py

    Utilities for using MPI

    defines
    MPI_size : number of process in the pool
    MPI_rank : rank of the current process

    typical use is presented in test_server_worker()

    in main do :
        if MPI_rank == 0:
            compute()   # your code,     uses mpiutil.enum_imap(myfunction, xarg)
        else:
            mpiutil.slave()     # slave is a generic function used to compute myfunction( arg )

    create a function to do the processing
    def myfunction(arg)
        return whatever

    and in compute() :
        # create an iterator over arguments
        M = 1000            # size of the pb
        result = M*[0.0]    # create room for storing results (if needed)
        xarg = xrange(M)    # create arguments
        res = mpiutil.enum_imap(myfunction, xarg)    # enum_imap applies myfunction to the iterator xarg
        for i,r in res:         # collect results
             result[i] = r      # and store them (order is not garantied)

    Warning, (i,r) will not arrive in order

    Will import even in MPI is not installed, however slave() will not work
    The situation is reported as MPI_size = 1

    Created by Marc-Andre' on 2012-04-19.
    extended 2014-03-25

    Copyright (c) 2012-2014 IGBMC. All rights reserved.
    """

    from __future__ import print_function

    """
    Recv is for numpy buffers
    recv is for picklable python objects
    """

    try:
        from mpi4py import MPI
    except ImportError:
        print("no MPI available")
        MPI_rank = 0
        MPI_size = 1
    else:
        MPI_comm = MPI.COMM_WORLD
        MPI_rank = MPI_comm.Get_rank()
        MPI_size = MPI_comm.size

    # TAG for MPI
    WORK_TAG = 1    # tag work message
    DIE_TAG = 2     # tag die message - used for closing the workers
    INIT_TAG = 3    # tag init message - used for initializing the workers
    TAG_OFFSET = 100    # offset used to code the
    DEBUG = False

    def mprint(st):
        "print with prefixing with the process rank"
        print("MPI process %d : %s"%(MPI_rank,st))
    def mprint_debug(st):
        "as mprint() but only if DEBUG is True"
        if DEBUG:
            mprint(st)
    def enum_imap(function, iterator):
        """
        applies function() to each element of iterator
        creates an iterator imap(function, iterator) that returns (i, function(iterator_i))
        elements will be returned in an synchronous manner with no garanty on the order.

        similar to enumerate( pool.imap_async(function, iterator) )
        except for the order
        """
        if MPI_size == 1:
            raise Exception("program should be started with mpirun")
        mprint_debug('server')
        status = MPI.Status()
        #initialize
        initialize(function)
        todo = 0    # counts computation sent
        done = 0    # counts results returned

        # main loop
        for inp in iterator:
            result = False      # assume unproductive loop
            # first receive from slaves ! (at least ok from initialize)
            data = MPI_comm.recv( source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            worker = status.Get_source()    # who was it ?
            if status.Get_tag() != INIT_TAG:
                result = True   # was real data
                ind = status.Get_tag()-TAG_OFFSET   # tag codes for index
                mprint_debug('from proc %d, I get %s for index %d'%(worker, str(data), ind))
                done += 1
                yield (ind, data)
            mprint_debug( 'sending new stuff to %d'% worker)
            MPI_comm.send(inp, dest=worker, tag=todo+TAG_OFFSET )
            todo += 1
        # as loops starts with receiving, we need to clean-up by receiving all waiting jobs
        mprint_debug( 'clean-up')
        while done != todo:
            mprint_debug( "done %d over %d to do"%( done, todo))
            result = False      # assume unproductive loop
            data = MPI_comm.recv( source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            worker = status.Get_source()
            if status.Get_tag() != INIT_TAG:
                result = True   # was real data
                ind = status.Get_tag()-TAG_OFFSET   # tag codes for index
                mprint_debug('from proc %d, I get %s for index %d'%(worker, str(data), ind))
                done += 1
                yield (ind, data)
        return
    def slave():
        "to be called in conjunction with enum_imap"
        if MPI_size == 1:
            raise Exception("program should be started with mpirun")
        mprint_debug('worker starting')
        func = _nill # this is the working function, loaded at INIT
        status = MPI.Status()
        while True:
            # mprint_debug('waiting')
            data = MPI_comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
            if status.Get_tag() == DIE_TAG:
                mprint_debug('received DIE_TAG')
                break
            if status.Get_tag() == INIT_TAG:
                mprint_debug('received INIT_TAG')
                func = data
                MPI_comm.send(None, dest=0, tag=INIT_TAG)
                MPI_comm.Barrier()
                mprint_debug("synchronized")
            else:
                res = func(data)
                MPI_comm.send(res, dest=0, tag=status.Get_tag())
        mprint_debug('finished')

    def _nill(a):
        "empty function"
        return a
    def initialize(func=_nill):
        """
        bring all processes to a starting barrier
        to be called by master process (rank == 0)
        """
        if MPI_size == 1:
            raise Exception("program should be started with mpirun")
        status = MPI.Status()
        #initialize
        for d in range(1,MPI_size):
            MPI_comm.send(func, dest=d, tag=INIT_TAG)
        mprint_debug("Waiting for MPI synchronisation")
        MPI_comm.Barrier()
        mprint_debug("synchronized")

    def shutdown():
        """
        closes all MPI processes all processes
        to be called by master process (rank == 0)
        """
        if MPI_size == 1:
            raise Exception("program should be started with mpirun")
        mprint_debug("Closing all sub-processes")
        data = []
        for d in range(1,MPI_size):
            MPI_comm.send(data, dest=d, tag=DIE_TAG)
        mprint_debug("MPI closed")

    ###### Example of use - call the program itself with mpirun -np X ########
    def funct_array_simple1(a):
        """
        used for testing
        take an array and compute 2*sqr(a) +1
        plus a random nap to simulate a random processing time
        """
        import time
        import random
        import numpy
        time.sleep(1+random.random())
        return 2.0*numpy.sqrt(a)+1  # numpy operation
    def funct_array_simple2(a):
        """
        used for testing
        take an array and compute 3*sqr(a+1)
        plus a random nap to simulate a random processing time
        """
        import time
        import random
        import numpy
        time.sleep(random.random())
        return 3.0*numpy.sqrt(a)  # numpy operation
    def test_compute_simple(N, M):
        """run funct_array_simple over enum_imap"""
        import numpy
        import time
        import itertools as it
        # create dummy data
        data = numpy.arange(N, dtype=numpy.float64)
        print("""
    The program should print the value %f %d times
    then the value %f %d times
    with a process speed-up proportionnal to the number of (processes-1).
    """%(funct_array_simple1(data).sum(), M, funct_array_simple2(data).sum(), 2*M))
        # prepare an iterator over which enum_imap will be applied
        xarg = it.repeat(data, M)
        t0 = time.time()
        res = enum_imap(funct_array_simple1, xarg)    # apply it -
        for i,r in res:       # and get results
            print(i,r.sum())
        xarg2 = it.repeat(data, 2*M)
        res = enum_imap(funct_array_simple2, xarg2)    # apply it
        for i,r in res:       # and get results
            print(i,r.sum())
        cpu = time.time()-t0
        print("process time : %.2f sec"%(cpu))
        return cpu

    def test_server_worker():
        if MPI_rank == 0:
            import time
            t0 = time.time()
            # values used for the test
            N = 100000  # array size
            M = 20 # number of operation to apply
            cpu = test_compute_simple(N, M)
            elaps = time.time()-t0
            print('elapsed %.2f sec  MPI starting overhead is %.2f sec'%(elaps, elaps-cpu))
            spd = (1.5*M + M)/cpu    # funct_array_simple1 is ~ 1.5sec x M  and funct_array_simple2 ~0.5sec x 2M
            print("speed up is x %.2f  for a theoretical maximum speedup of x %d"%(spd, MPI_size-1))
            shutdown()
        else:
            slave()

    if __name__ == '__main__':
        # this should be called with mpirun to check the MPI behavior.
        # as the tested actions are mostly empty and just sleep, it can safely be tested with a large number of processes
        # even on a small machine.
        if MPI_size == 1:       # means no MPI
            print(__doc__)
        else:
            test_server_worker()