mpi4py
mpi4py, which stands for "MPI for Python", is a library that provides a Python interface to the Message Passing Interface (MPI) and its implementation in C, allowing for interprocess communication. While programming with MPI is generally done at a rather low level, with explicit allocation of buffers and specification of data types being communicated between processes, mpi4py provides a somewhat less complicated, more Pythonic interface to the underlying functionality. In particular, mpi4py provides functions for:
- Communication of generic Python objects, using the Python pickle module to serialize data for communication among processes
- Fast, near C-speed communication of NumPy arrays and other buffer-provider objects
MPI for Python is predictably related to the standard MPI interface. The mpi4py library automatically initializes MPI when it is imported, and the size of the global communicator and the rank of the local process are easy to find. When entering the following into the Python interpreter, you will discover that your interactive MPI world consists of only one process:
For many scientific programs, a common pattern of data storage and exchange centers around NumPy arrays (which might be communicated among processors in a message-passing application). Since NumPy arrays have a fixed type and size which can be queried, MPI calls which communicate array data can be executed both easily and efficiently:
Note that for point-to-point calls, you must still supply the destination task ID and the tag that helps to distinguish between messages. In the above example, you should find that the initial zeros in data2 have been replaced by a sequence of non-zero numbers.
All of the standard MPI calls are implemented, including the usual collective communications. As noted above, it is possible to send Python objects of arbitrary type, which are serialized with the pickle module, communicated between processes as a flat, encoded, data string, and then unserialized at the destination process using pickle once more. The details of this communication are available in the online mpi4py documentation. A more involved example demonstrating the usage of mpi4py is contained in the Exercise at the end of this topic.
If you want to have multiple MPI tasks run your Python script, you must initialize the set of Python processes externally to the script using a method from the locally installed MPI software, e.g., the mpiexec
or ibrun
command. Be aware that this mpiexec
must be associated with the same MPI implementation that was used to build mpi4py. For example, on Frontera, you execute Python MPI code in the same way that you would execute any other MPI program, using the ibrun
command:
The ibrun command starts MPI tasks on every core of every allocated node in your job. By default, the SLURM batch scheduler will tell your job to start up with a copy of your current environment variables. That means you should use the module command to add the correct Python version before running sbatch to submit your job.