MPI4PY
This page documents how to use the MPI for Python package within a Conda environment.
Using MPI4PY in a Python Script
To use MPI in a python script through mpi4py
, you must first import it using the following code:
from mpi4py import MPI
Example Script
Here is an example python script mpi4pytest.py
that uses MPI:
from mpi4py import MPI
import sys
def print_hello(rank, size, name):
msg = "Hello World! I am process {0} of {1} on {2}.\n"
sys.stdout.write(msg.format(rank, size, name))
if __name__ == "__main__":
size = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
name = MPI.Get_processor_name()
print_hello(rank, size, name)
Conda Environment
Start by creating and activating a conda environment:
Once you have activated your conda environment, run the following commands to install mpi4py
:
$ module load hpcx-mpi/4.1.5rc2-mt
$ pip install mpi4py
Here is an example batch job script mpi4pytest_conda.sh
that uses mpi4pytest.py
and the conda environment setup:
#!/bin/bash
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --mem=1G
module load miniconda3/23.11.0s
source /oscar/runtime/software/external/miniconda3/23.11.0/etc/profile.d/conda.sh
conda activate my_env
module hpcx-mpi/4.1.5rc2-mt
srun --mpi=pmix python mpi4pytest.py
Python Environment
Start by creating and activating a Python environment
Once you have activated your conda environment, run the following command to install mpi4py
:
$ python -m pip install mpi4py
$ deactivate
Below is an example batch job script mpi4pytest_env.sh
:
#!/bin/bash
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --mem=1G
module load hpcx-mpi/4.1.5rc2-mt
source my_env/bin/activate
srun --mpi=pmix python mpi4pytest.py
Last updated
Was this helpful?