Example 5: Containerized MPI Multi-Process Job

This example shows how to run a multi-process MPI job with Slurm and Apptainer. It uses a small Python program with mpi4py to estimate pi using a distributed Monte Carlo calculation.

The example follows a hybrid launch model: the host MPI launcher starts one Apptainer container per MPI rank. The container provides Python and mpi4py; the host MPI stack coordinates the ranks across the allocated nodes.

1. Connect to the CoSTAR cluster

Open a terminal and connect to the CoSTAR login node:

ssh ab1234@costar-login01

Replace ab1234 with your University username.

2. Create the Apptainer definition file

Download the definition file, or create a file called example-05-mpi.def in your working directory and copy in the recipe below.

• example-05-mpi.def

example-05-mpi.def (click to collapse / view)

Bootstrap: docker
From: ubuntu:22.04

%post
    set -e
    export DEBIAN_FRONTEND=noninteractive

    apt-get update
    apt-get install -y --no-install-recommends \
        python3 python3-pip python3-dev build-essential \
        libopenmpi-dev openmpi-bin

    python3 -m pip install --upgrade pip
    python3 -m pip install --no-cache-dir mpi4py

    apt-get clean
    rm -rf /var/lib/apt/lists/* /root/.cache/pip || true

%environment
    export PATH=/usr/bin:$PATH
    export PYTHONUNBUFFERED=1

%labels
    Author CoSTAR
    Description "Python mpi4py example for Slurm multi-process jobs"

3. Create the MPI Python script

Download the Python script, or create a file called example-05-mpi.py in your working directory and copy in the script below.

• example-05-mpi.py

example-05-mpi.py (click to collapse / view)

from mpi4py import MPI

import random
import socket


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

if rank == 0:
    print(f"[INFO] World size = {size}")

comm.Barrier()
print(f"Rank {rank:02d} on host {host}")

total_samples = 2_000_000
local_samples = total_samples // size
random.seed(1234 + rank)

hits = 0
for _ in range(local_samples):
    x = random.random()
    y = random.random()
    if x * x + y * y <= 1.0:
        hits += 1

total_hits = comm.reduce(hits, op=MPI.SUM, root=0)

if rank == 0:
    samples_used = local_samples * size
    pi_estimate = 4.0 * total_hits / samples_used
    print(f"[RESULT] pi ~= {pi_estimate:.6f} using {samples_used} samples")

4. Build the Apptainer image

Note

On CoSTAR, Apptainer is available on compute nodes, not on the login node. If you want to build the image on the cluster, start an interactive Slurm session on a compute node first.

Start an interactive Slurm session:

srun -N 1 --time=00-01:00:00 --pty bash

Once the session starts on a compute node, build the image:

apptainer build example-05-mpi.sif example-05-mpi.def

When the build finishes, exit the interactive session:

exit

5. Run interactively

For a quick interactive test, first request an allocation across 2 nodes with 1 MPI rank per node:

salloc -N 2 --ntasks-per-node=1 -t 02:00:00

When the allocation starts, run the MPI program inside the Apptainer image:

mpirun -np 2 apptainer exec --cleanenv example-05-mpi.sif python3 example-05-mpi.py

This launches 2 MPI ranks in total. When you finish testing, leave the interactive allocation with:

exit

6. Create the batch submission script

Download the example submission script, or create a file called example-05-mpi.sh in your working directory and copy in the script below.

• example-05-mpi.sh

example-05-mpi.sh (click to collapse / view)

#!/bin/bash

#SBATCH --job-name=mpi-container
#SBATCH --partition=main
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --time=00:10:00
#SBATCH --hint=nomultithread
#SBATCH --output=slurm.%N.%j.out
#SBATCH --error=slurm.%N.%j.err

module purge
module load NVHPC/26.3-CUDA-13.1.0
module load UCX-CUDA/1.19.0-GCCcore-14.3.0-CUDA-13.1.0

IMAGE=$PWD/example-05-mpi.sif
SCRIPT=$PWD/example-05-mpi.py

echo "Job started on $(date)"
echo "Nodes allocated:"
scontrol show hostnames "$SLURM_NODELIST"

mpirun -np "$SLURM_NTASKS" apptainer exec --cleanenv "$IMAGE" python3 "$SCRIPT"

echo "Job finished on $(date)"

The job requests 2 nodes with 1 MPI rank per node, for 2 MPI ranks in total. The mpirun command launches those ranks and runs the Python script inside the Apptainer image.

7. Submit as a batch job

Submit the script to the Slurm scheduler:

sbatch example-05-mpi.sh

You can check the status of your job by running:

squeue --me

8. Check the output

When the job finishes, Slurm writes the standard output and error streams to files in the submission directory:

• slurm.<node>.<job_id>.out contains the allocated node list, MPI rank placement, and the final pi estimate.

• slurm.<node>.<job_id>.err contains error messages, if any were produced.

The output file should include lines similar to:

Nodes allocated:
costar01
costar02
[INFO] World size = 2
Rank 00 on host costar01
Rank 01 on host costar02
[RESULT] pi ~= 3.141820 using 2000000 samples

The rank order may differ between runs, but the world size should match the number of MPI ranks requested by Slurm.

9. What to change next

For your own MPI workload, adjust:

• --nodes for the number of compute nodes.

• --ntasks-per-node for the number of MPI ranks per node.

• IMAGE to point to your Apptainer image.

• SCRIPT to point to your MPI application or Python script.

• The MPI modules to match the MPI version needed by your application.