Package mpi4py :: Module __main__
[hide private]
[frames] | no frames]

Source Code for Module mpi4py.__main__

  1  # Author:  Lisandro Dalcin 
  2  # Contact: dalcinl@gmail.com 
  3  """ 
  4  Run some benchmarks and tests 
  5  """ 
  6  import sys as _sys 
  7   
  8   
9 -def helloworld(comm, args=None, verbose=True):
10 """ 11 Hello, World! using MPI 12 """ 13 from . import MPI 14 from optparse import OptionParser 15 parser = OptionParser(prog="mpi4py helloworld") 16 parser.add_option("-q", "--quiet", action="store_false", 17 dest="verbose", default=verbose) 18 (options, args) = parser.parse_args(args) 19 20 size = comm.Get_size() 21 rank = comm.Get_rank() 22 name = MPI.Get_processor_name() 23 message = ("Hello, World! I am process %*d of %d on %s.\n" 24 % (len(str(size - 1)), rank, size, name)) 25 comm.Barrier() 26 if rank > 0: 27 comm.Recv([None, 'B'], rank - 1) 28 if options.verbose: 29 _sys.stdout.write(message) 30 _sys.stdout.flush() 31 if rank < size - 1: 32 comm.Send([None, 'B'], rank + 1) 33 comm.Barrier() 34 return message
35 36
37 -def ringtest(comm, args=None, verbose=True):
38 """ 39 Time a message going around the ring of processes 40 """ 41 # pylint: disable=too-many-locals 42 # pylint: disable=too-many-branches 43 from . import MPI 44 from array import array 45 from optparse import OptionParser 46 parser = OptionParser(prog="mpi4py ringtest") 47 parser.add_option("-q", "--quiet", action="store_false", 48 dest="verbose", default=verbose) 49 parser.add_option("-n", "--size", type="int", default=1, dest="size", 50 help="message size") 51 parser.add_option("-s", "--skip", type="int", default=0, dest="skip", 52 help="number of warm-up iterations") 53 parser.add_option("-l", "--loop", type="int", default=1, dest="loop", 54 help="number of iterations") 55 (options, args) = parser.parse_args(args) 56 57 def ring(comm, n=1, loop=1, skip=0): 58 # pylint: disable=invalid-name 59 # pylint: disable=bad-whitespace 60 # pylint: disable=missing-docstring 61 iterations = list(range((loop + skip))) 62 size = comm.Get_size() 63 rank = comm.Get_rank() 64 source = (rank - 1) % size 65 dest = (rank + 1) % size 66 Sendrecv = comm.Sendrecv 67 Send = comm.Send 68 Recv = comm.Recv 69 Wtime = MPI.Wtime 70 sendmsg = array('B', [+42]) * n 71 recvmsg = array('B', [0x0]) * n 72 if size == 1: 73 for i in iterations: 74 if i == skip: 75 tic = Wtime() 76 Sendrecv(sendmsg, dest, 0, 77 recvmsg, source, 0) 78 else: 79 if rank == 0: 80 for i in iterations: 81 if i == skip: 82 tic = Wtime() 83 Send(sendmsg, dest, 0) 84 Recv(recvmsg, source, 0) 85 else: 86 sendmsg = recvmsg 87 for i in iterations: 88 if i == skip: 89 tic = Wtime() 90 Recv(recvmsg, source, 0) 91 Send(sendmsg, dest, 0) 92 toc = Wtime() 93 if comm.rank == 0 and sendmsg != recvmsg: # pragma: no cover 94 import warnings 95 import traceback 96 try: 97 warnings.warn("received message does not match!") 98 except UserWarning: 99 traceback.print_exc() 100 comm.Abort(2) 101 return toc - tic
102 103 size = getattr(options, 'size', 1) 104 loop = getattr(options, 'loop', 1) 105 skip = getattr(options, 'skip', 0) 106 comm.Barrier() 107 elapsed = ring(comm, size, loop, skip) 108 if options.verbose and comm.rank == 0: 109 message = ("time for %d loops = %g seconds (%d processes, %d bytes)\n" 110 % (loop, elapsed, comm.size, size)) 111 _sys.stdout.write(message) 112 _sys.stdout.flush() 113 return elapsed 114 115
116 -def main(args=None):
117 "Entry-point for ``python -m mpi4py``" 118 from optparse import OptionParser 119 from . import __name__ as prog 120 from . import __version__ as version 121 parser = OptionParser(prog=prog, version='%prog ' + version, 122 usage="%prog [options] <command> [args]") 123 parser.add_option("--no-threads", 124 action="store_false", dest="threads", default=True, 125 help="initialize MPI without thread support") 126 parser.add_option("--thread-level", type="choice", metavar="LEVEL", 127 choices=["single", "funneled", "serialized", "multiple"], 128 action="store", dest="thread_level", default="multiple", 129 help="initialize MPI with required thread support") 130 parser.add_option("--mpe", 131 action="store_true", dest="mpe", default=False, 132 help="use MPE for MPI profiling") 133 parser.add_option("--vt", 134 action="store_true", dest="vt", default=False, 135 help="use VampirTrace for MPI profiling") 136 parser.disable_interspersed_args() 137 (options, args) = parser.parse_args(args) 138 139 from . import rc, profile 140 rc.threads = options.threads 141 rc.thread_level = options.thread_level 142 if options.mpe: 143 profile('mpe', logfile='mpi4py') 144 if options.vt: 145 profile('vt', logfile='mpi4py') 146 147 from . import MPI 148 comm = MPI.COMM_WORLD 149 if not args: 150 if comm.rank == 0: 151 parser.print_usage() 152 parser.exit() 153 command = args.pop(0) 154 if command not in main.commands: 155 if comm.rank == 0: 156 parser.error("unknown command '%s'" % command) 157 parser.exit(2) 158 command = main.commands[command] 159 command(comm, args=args) 160 parser.exit()
161 162 main.commands = { 163 'helloworld': helloworld, 164 'ringtest': ringtest, 165 } 166 167 if __name__ == '__main__': 168 main() 169