MPI4py
MPI for Python provides Python bindings for the Message Passing Interface (MPI) standard, allowing Python applications to exploit multiple processors on workstations, clusters and supercomputers.
Available versions
mpi4py is available as a module, and not from the wheelhouse as typical Python packages do.
One can find available version using:
[name@server ~]$ module spider mpi4py
and look for more information on a specific version by using:
[name@server ~]$ module spider mpi4py/X.Y.Z
where X.Y.Z is the exact desired version, for instance 4.0.0.
Famous first words: Hello World
1. Run a short interactive job :
[name@server ~]$ salloc --account=<your account> --ntasks=5
2. Load the module:
[name@server ~]$ module load mpi4py/4.0.0 python/3.12
3. Run a Hello World test:
[name@server ~]$ srun python -m mpi4py.bench helloworld
Hello, World! I am process 0 of 5 on node1.
Hello, World! I am process 1 of 5 on node1.
Hello, World! I am process 2 of 5 on node3.
Hello, World! I am process 3 of 5 on node3.
Hello, World! I am process 4 of 5 on node3.
In the case above, two nodes (node1 and node3) were allocated, and the tasks were distributed accross the available ressources.
mpi4py as a package dependency
Often mpi4py is a dependency for another package. In order to fulfill this dependency :
1. Deactivate any Python virtual environment:
[name@server ~]$ test $VIRTUAL_ENV && deactivate
Note: If you had a virtual environment activated, it is important to deactivate it first, then load the module, before re-activating your virtual environment.
2. Load the module:
[name@server ~]$ module load mpi4py/4.0.0 python/3.12
3. Check that it is visible by pip:
[name@server ~]$ pip list | grep mpi4py
mpi4py 4.0.0
and is accessible for your current python module loaded:
[name@server ~]$ python -c 'import mpi4py'
If no errors are raised, then everything is ok!
4. Create a virtual env. and install your packages
Running jobs with mpi4py
You can run mpi jobs distributed across multiple nodes or cores. For efficient MPI scheduling, please see:
CPU
1. Write your python code, for instance broadcasting a numpy array:
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if rank == 0:
data = np.arange(100, dtype='i')
else:
data = np.empty(100, dtype='i')
comm.Bcast(data, root=0)
for i in range(100):
assert data[i] == i
The example above is based on mpi4py tutorial.
2. Write your submission script:
#!/bin/bash
#SBATCH --account=def-someprof # adjust this to match the accounting group you are using to submit jobs
#SBATCH --time=08:00:00 # adjust this to match the walltime of your job
#SBATCH --ntasks=4 # adjust this to match the number of tasks/processes to run
#SBATCH --mem-per-cpu=4G # adjust this according to the memory you need per process
# Run on cores across the system : https://docs.alliancecan.ca/wiki/Advanced_MPI_scheduling#Few_cores,_any_number_of_nodes
# Load modules dependencies.
module load StdEnv/2023 gcc mpi4py/4.0.0 python/3.12
# create the virtual environment on each allocated node:
srun --ntasks $SLURM_NNODES --tasks-per-node=1 bash << EOF
virtualenv --no-download $SLURM_TMPDIR/env
source $SLURM_TMPDIR/env/bin/activate
pip install --no-index --upgrade pip
pip install --no-index numpy==2.1.1
EOF
# activate only on main node
source $SLURM_TMPDIR/env/bin/activate;
# srun exports the current env, which contains $VIRTUAL_ENV and $PATH variables
srun python mpi4py-np-bc.py;
#!/bin/bash
#SBATCH --account=def-someprof # adjust this to match the accounting group you are using to submit jobs
#SBATCH --time=01:00:00 # adjust this to match the walltime of your job
#SBATCH --nodes=2 # adjust this to match the number of whole node
#SBATCH --ntasks-per-node=40 # adjust this to match the number of tasks/processes to run per node
#SBATCH --mem-per-cpu=1G # adjust this according to the memory you need per process
# Run on N whole nodes : https://docs.alliancecan.ca/wiki/Advanced_MPI_scheduling#Whole_nodes
# Load modules dependencies.
module load StdEnv/2023 gcc openmpi mpi4py/4.0.0 python/3.12
# create the virtual environment on each allocated node:
srun --ntasks $SLURM_NNODES --tasks-per-node=1 bash << EOF
virtualenv --no-download $SLURM_TMPDIR/env
source $SLURM_TMPDIR/env/bin/activate
pip install --no-index --upgrade pip
pip install --no-index numpy==2.1.1
EOF
# activate only on main node
source $SLURM_TMPDIR/env/bin/activate;
# srun exports the current env, which contains $VIRTUAL_ENV and $PATH variables
srun python mpi4py-np-bc.py;
2. Submit your job to the scheduler.
2.1 Test your script.
Before submitting your job, it is important to test that your submission script will start without errors. You can do a quick test in an interactive job.
2.2 Submit your job
[name@server ~]$ sbatch submit-mpi4py-distributed.sh
GPU
1. From a login node, download the demo example:
[name@server ~]$ wget https://raw.githubusercontent.com/mpi4py/mpi4py/refs/heads/master/demo/cuda-aware-mpi/use_cupy.py
The example above and others, can be found in the demo folder.
2. Write your submission script:
#!/bin/bash
#SBATCH --account=def-someprof # adjust this to match the accounting group you are using to submit jobs
#SBATCH --time=08:00:00 # adjust this to match the walltime of your job
#SBATCH --ntasks=2 # adjust this to match the number of tasks/processes to run
#SBATCH --mem-per-cpu=2G # adjust this according to the memory you need per process
#SBATCH --gpus=1
# Load modules dependencies.
module load StdEnv/2023 gcc cuda/12 mpi4py/4.0.0 python/3.11
# create the virtual environment on each allocated node:
virtualenv --no-download $SLURM_TMPDIR/env
source $SLURM_TMPDIR/env/bin/activate
pip install --no-index --upgrade pip
pip install --no-index cupy numba
srun python use_cupy.py;
2. Submit your job to the scheduler.
2.1 Test your script.
Before submitting your job, it is important to test that your submission script will start without errors. You can do a quick test in an interactive job.
2.2 Submit your job
[name@server ~]$ sbatch submit-mpi4py-gpu.sh
Troubleshooting
ModuleNotFoundError: No module named 'mpi4py'
If mpi4py is not accessible, one may get the following error when importing it:
ModuleNotFoundError: No module named 'mpi4py'
Possible solution to fix this error:
- check compatible python version with
module spider mpi4py/X.Y.Zand thatpython -c 'import mpi4py'works. - load the module before activating a virtual environment: please see the above section MPI4py#mpi4py as a package dependency
And also see Python#ModuleNotFoundError:_No_module_named_'X'