Rocky Python MPI Hello World

MPI Jobs on Rocky

Supported Stack

We currently recommend and support:

• OpenMPI 4.1.6
• PMIx 4.2.6
• mpi4py 4.0.1

MPI Example: Hello World in Python

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=1
#SBATCH --ntasks=3
#SBATCH --time=00:05:00
#SBATCH --partition=kivlin,compute_all

module load OpenMPI/4.1.6-GCC-13.2.0
module load mpi4py

echo "=== Running mpi4py test with Slurm PMIx ==="
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