MPI4py

From Alliance Doc
Revision as of 14:51, 1 October 2024 by Coulombc (talk | contribs) (Updated title)
Jump to navigation Jump to search
Other languages:

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[edit]

mpi4py is available as a module, and not from the wheelhouse as typical Python packages do. One can find available version using:

Question.png
[name@server ~]$ module spider mpi4py

and look for more information on a specific version by using:

Question.png
[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[edit]

1. Run a short interactive job :

Question.png
[name@server ~]$ salloc --account=<your account> --ntasks=5

2. Load the module:

Question.png
[name@server ~]$ module load mpi4py/4.0.0 python/3.12

3. Run a Hello World test:

Question.png
[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[edit]

Often mpi4py is a dependency for another package. In order to fulfill this dependency :

1. Deactivate any Python virtual environment:

Question.png
[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:

Question.png
[name@server ~]$ module load mpi4py/4.0.0 python/3.12

3. Check that it is visible by pip:

Question.png
[name@server ~]$ pip list | grep mpi4py
mpi4py            4.0.0

and is accessible for your current python module loaded:

Question.png
[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[edit]

You can run mpi jobs distributed across multiple nodes or cores. For efficient MPI scheduling, please see:

CPU[edit]

1. Write your python code, for instance broadcasting a numpy array:

File : "mpi4py-np-bc.py"

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:

File : submit-mpi4py-distributed.sh

#!/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;


File : submit-mpi4py-whole-nodes.sh

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

Question.png
[name@server ~]$ sbatch submit-mpi4py-distributed.sh

GPU[edit]

1. From a login node, download the demo example:

Question.png
[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:

File : submit-mpi4py-gpu.sh

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

Question.png
[name@server ~]$ sbatch submit-mpi4py-gpu.sh

Troubleshooting[edit]

ModuleNotFoundError: No module named 'mpi4py'[edit]

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.Z and that python -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'