MPI for Python (mpi4py)

MPI for Python is an implementation of language bindings for Python of the MPI parallel computing framework.

Homepage Version(s): 3.0.0

---

MPI for Python (mpi4py) requires an environment module

In order to use MPI for Python (mpi4py), you must first load the appropriate environment module:

module load gnu openmpi

This package allows access to MPI function calls and allows the user to leverage multiple processors to run their code natively in Python without having to use other bindings for MPI such as C or Fortran.

Running MPI for Python on the HPC

MPI for Python is a native Python library that allows access to the MPI framework for parallel processing. Once you have loaded the requisite environment modules (module load gnu openmpi), you can then run your Python script using the srun command:

Warning

mpi4py must use our centrally installed MPI. Otherwise, it is likely to fail with Slurm-related errors.

UV

$ module load gnu openmpi python-uv
$ uv venv ~/mpienv
$ source ~/mpienv/bin/activate
$ uv pip install --no-cache-dir mpi4py

$ srun -n 4 python MY_SCRIPT.py

Conda

$ module load gnu openmpi anaconda/3.12.4
$ conda create -n mpienv
$ conda activate mpienv
$ conda install mpi4py openmpi=4.1.*=external_*

$ srun -n 4 python MY_SCRIPT.py

PIP (Virtualenv)

$ module load gnu openmpi python
$ python -m venv ~/mpienv
$ source ~/mpienv/bin/activate
$ pip install --no-cache-dir mpi4py

$ srun -n 4 python MY_SCRIPT.py

Submitting a mpi4py Job to Slurm

Below is an example script that uses four processes:

#!/bin/bash

#SBATCH --job-name=MY_MPI4PY_JOB
#SBATCH -n 4
#SBATCH -t 01:00:00
# ..additional Slurm Params here..

# Load the MPI environment modules 
# (substitute python-uv with 'anaconda/3.12.4' if using Conda or 'python' if using Virtualenv)
module load gnu openmpi python-uv

# Activate your uv, Conda, or venv environment
source ~/mpienv/bin/activate

srun python MY_SCRIPT.py

For example Python scripts that you can use for testing, refer to the official tutorial.