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

[docs]def mprint(st): "print with prefixing with the process rank" print("MPI process %d : %s"%(MPI_rank,st))
[docs]def mprint_debug(st): "as mprint() but only if DEBUG is True" if DEBUG: mprint(st)
[docs]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
[docs]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
[docs]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")
[docs]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 ########
[docs]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
[docs]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
[docs]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
[docs]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()