Skip navigation
NASA Logo, National Aeronautics and Space Administration
Currently Being Moderated

Using mpi4py

VERSION 7  Click to view document history
Created on: Aug 27, 2013 9:32 AM by Jules Kouatchou - Last Modified:  Aug 16, 2018 4:12 PM by Jules Kouatchou

MPI for Python (mpi4py) is a package that enables applications exploit multiple processors using standard MPI “look and feel” in Python scripts. It is constructed on top of the MPI-1/2 specifications and provides an object oriented interface which closely follows MPI-2 C++ bindings. It supports point-to-point (sends, receives) and collective (broadcasts, scatters, gathers) communications of any picklable Python object, as well as optimized communications of Python object exposing the single-segment buffer interface (NumPy arrays, builtin bytes/string/array objects).

 

Sample mpi4py Scripts

Simple example of mpi4py scripts are available at http://mpi4py.scipy.org/docs/usrman/tutorial.html

 

Consider a script that implements the scatter-gather procedure:

 

#!/usr/bin/env python
 
#--------------
# Loaded Modules
#--------------
import numpy as np
from mpi4py import MPI
 
#-------------
# Communicator
#-------------
comm = MPI.COMM_WORLD
 
my_N = 4
N = my_N * comm.size
 
if comm.rank == 0:
    A = np.arange(N, dtype=np.float64)
else:
    #Note that if I am not the root processor A is an empty array
    A = np.empty(N, dtype=np.float64)
 
my_A = np.empty(my_N, dtype=np.float64)
 
#-------------------------
# Scatter data into my_A arrays
#-------------------------
comm.Scatter( [A, MPI.DOUBLE], [my_A, MPI.DOUBLE] )
 
if comm.rank == 0:
   print "After Scatter:"
 
for r in xrange(comm.size):
    if comm.rank == r:
        print "[%d] %s" % (comm.rank, my_A)
    comm.Barrier()
 
#-------------------------
# Everybody is multiplying by 2
#-------------------------
my_A *= 2
 
#-----------------------
# Allgather data into A again
#-----------------------
comm.Allgather( [my_A, MPI.DOUBLE], [A, MPI.DOUBLE] )
 
if comm.rank == 0:
   print "After Allgather:"
 
for r in xrange(comm.size):
    if comm.rank == r:
        print "[%d] %s" % (comm.rank, A)
    comm.Barrier()

 

 

Submitting the Script

Assume that you name the above script scatter_gather.py, you can run it on discover in an interactive session or using PBS. You will need to load the modules:

 

module load other/comp/gcc-4.9.2-sp3
module load other/mpi/mvapich2-2.1/gcc-4.9.2-sp3
module load lib/mkl-15.0.2.164
module load other/SIVO-PyD/spd_1.24.0_gcc-4.9.2-sp3_mkl-15.0.2.164_mvapich2-2.1

 

and submit the script as:

 

mpirun -np 24 scatter_gather.py

 

Here is an example of PBS script:

 

#!/bin/csh -f

#SBATCH --time=00:10:00
#SBATCH --job-name=ptest_mpi4py
#SBATCH --nodes=4
#SBATCH --constraint=hasw
#SBATCH -A XXXXX

source /usr/share/modules/init/csh
module purge
module load other/comp/gcc-4.9.2-sp3
module load other/mpi/mvapich2-2.1/gcc-4.9.2-sp3
module load lib/mkl-15.0.2.164
module load other/SIVO-PyD/spd_1.24.0_gcc-4.9.2-sp3_mkl-15.0.2.164_mvapich2-2.1

mpirun -np 24 scatter_gather.py

 

 

mpi4py and F2Py

Assume that you have a Fortran subroutine (named myFortSubroutine) that has MPI calls. You want to use that subroutine in a mpi4py script. Here are the steps:

 

Step 1: Using F2Py to create a module

We can assume that the subroutine is in the file myFortSubroutine.F90. To generate the Python module from the Fortran source code, type:

 

module purge
module load other/comp/gcc-4.9.2-sp3
module load other/mpi/mvapich2-2.1/gcc-4.9.2-sp3
module load lib/mkl-15.0.2.164
module load other/SIVO-PyD/spd_1.24.0_gcc-4.9.2-sp3_mkl-15.0.2.164_mvapich2-2.1

f2py -c --fcompiler=gnu95 --include-paths /path/to/MPI/include/ myFortSubroutine.F90 -m myroutine -L/path/to/MPI/lib/ -lmpi


 

Step 2: Calling the Fortran subroutine from the Python script

 

#!/usr/bin/env python

#--------------
# Loaded Modules
#--------------
from mpi4py import MPI
import myroutine

#------------------------------------
# Set the MPI communicator
#-------------------------------------
fcomm = MPI.COMM_WORLD.py2f()

#-----------------------------------------
# Calling the Fortran subroutine
#-----------------------------------------
myroutine.myfortsubroutine(fcomm)


 

Additional Resources:


Comments (0)
USAGov logo NASA Logo - nasa.gov