HPC UNITe

MPI4py ping-pong demonstrates point-to-point communication between rank 0 (master) and rank 1 (worker). Rank 0 sends increasingly larger messages to rank 1, which echoes them back; timing measures bandwidth.

pingpong_mpi4py.py

from mpi4py import MPI
import numpy as np
import time
 
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
 
if size < 2:
    if rank == 0:
        print("Need at least 2 processes")
    exit()
 
N = 1000  # max message size
extent = 100  # message size steps
sbuf = np.zeros(1, dtype='d')
rbuf = np.zeros(1, dtype='d')
 
if rank == 0:
    t0 = time.time()
    for i in range(extent):
        size = (i + 1) * N
        sbuf = np.zeros(size, dtype='d') + rank
        t1 = time.time()
        comm.send(sbuf, dest=1, tag=i)
        comm.Recv(rbuf, source=1, tag=i)
        t2 = time.time()
        latency = (t2 - t1) * 1000  # ms
        bandwidth = (size * 8) / (t2 - t1) / 1e6  # MB/s
        print(f"Size {size}: latency {latency:.2f}ms, BW {bandwidth:.1f} MB/s")
    total_time = time.time() - t0
    print(f"Total time: {total_time:.2f}s")
elif rank == 1:
    for i in range(extent):
        rbuf = comm.recv(source=0, tag=i)
        comm.send(rbuf, dest=0, tag=i)

slurm_pingpong.job

#!/bin/bash
#SBATCH --job-name=pingpong_mpi4py
#SBATCH --partition=unite
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1
#SBATCH --time=00:01:00
#SBATCH --output=pingpong_%j.out
 
module purge
#Load necessary modules
module add unite/python/3.14/mpi4py
 
#Start the example
mpirun -np $SLURM_NTASKS python pingpong_mpi4py.py