#!/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)"