Rocky Python MPI Hello World

MPI4PY Jobs on Rocky

Supported versions

Rocky supports PMIx v3 for MPI process management. MPI applications should be launched via srun using --mpi=pmix_v3.

To see what versions of OpenMPI are available use the following command:

module spider OpenMPI

As of the writing of this document the versions available:

     Versions:
        OpenMPI/3.1.1-GCC-7.3.0-2.30
        OpenMPI/3.1.3-GCC-8.2.0-2.31.1
        OpenMPI/3.1.4-GCC-8.3.0
        OpenMPI/4.0.3-GCC-9.3.0
        OpenMPI/4.0.5-GCC-10.2.0
        OpenMPI/4.1.1-GCC-10.3.0
        OpenMPI/4.1.1-GCC-11.2.0
        OpenMPI/4.1.4-GCC-11.3.0
        OpenMPI/4.1.4-GCC-12.2.0
        OpenMPI/4.1.5-GCC-12.3.0
        OpenMPI/4.1.6-GCC-13.2.0
        OpenMPI/5.0.3-GCC-13.3.0
        OpenMPI/5.0.7-GCC-14.2.0
        OpenMPI/5.0.8-GCC-14.3.0
        OpenMPI/5.0.8-llvm-compilers-20.1.8

Similarly, to see what version of mpi4py is available use:

module spider mpi4py

     Versions:
        mpi4py/3.0.1 (E)
        mpi4py/3.0.2 (E)
        mpi4py/3.0.3 (E)
        mpi4py/3.1.1 (E)
        mpi4py/3.1.3 (E)
        mpi4py/3.1.4-gompi-2022b
        mpi4py/3.1.4-gompi-2023a
        mpi4py/3.1.4 (E)
        mpi4py/3.1.5-gompi-2023b
        mpi4py/3.1.5 (E)
        mpi4py/3.1.6 (E)
        mpi4py/4.0.1-gompi-2024a
        mpi4py/4.0.1 (E)

With the above given information we are going to use OpenMPI/5.0.8-GCC-14.3.0 and mpi4py/4.0.1.

hello_mpi4py.py

from mpi4py import MPI
import socket

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
hostname = socket.gethostname()

print(f"Hello from rank {rank} of {size} on {hostname}")

hello_mpi4py.sbatch

#!/usr/bin/env bash
#SBATCH --job-name=hello-mpi4py
#SBATCH --output=hello-mpi4py-%j.out
#SBATCH --error=hello-mpi4py-%j.err
#SBATCH --nodes=2
#SBATCH --ntasks=8
#SBATCH --time=00:05:00

module load OpenMPI/5.0.8-GCC-14.3.0
module load mpi4py/4.0.1

srun --mpi=pmix_v3 python hello_mpi4py.py

Submit the job:

sbatch hello_mpi4py.sbatch

Expected Output

Hello from rank 0 of 8 on rocky4.rocky.nimbios.org
Hello from rank 1 of 8 on rocky4.rocky.nimbios.org
Hello from rank 2 of 8 on rocky4.rocky.nimbios.org
Hello from rank 3 of 8 on rocky4.rocky.nimbios.org
Hello from rank 4 of 8 on rocky5.rocky.nimbios.org
Hello from rank 5 of 8 on rocky5.rocky.nimbios.org
Hello from rank 6 of 8 on rocky5.rocky.nimbios.org
Hello from rank 7 of 8 on rocky5.rocky.nimbios.org