#!/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()