#!/bin/bash
#SBATCH --job-name=conda-container
#SBATCH --partition=test
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem-per-cpu=4000
#SBATCH --time=00:10:00

# Usage: sbatch run-environment.sbatch \
#        "/project/${project}" "/project/${project}/images/environment-v1.sif" \
#        python simulation.py
set -euo pipefail
if (( $# < 3 )); then
    echo "Usage: PROJECT_DIRECTORY IMAGE COMMAND [ARG...]" >&2
    exit 1
fi
PROJECT_DIR=$(realpath -e -- "$1")
IMAGE=$(realpath -e -- "$2")
shift 2
module load apptainer/1.4.4

JOB_WORK=$(mktemp -d /tmp/conda-container.XXXXXXXX)
trap 'rm -rf -- "$JOB_WORK"' EXIT
mkdir -p "$JOB_WORK/cache" "$JOB_WORK/mamba"

# The image is shared read-only; each job has its own writable scratch space.
srun apptainer run --cleanenv --contain \
    --bind "$PROJECT_DIR:/work" --bind "$JOB_WORK:/jobtmp" --pwd /work \
    --env TMPDIR=/jobtmp,XDG_CACHE_HOME=/jobtmp/cache,MAMBA_ROOT_PREFIX=/jobtmp/mamba \
    --env NUMBA_CACHE_DIR=/jobtmp/cache/numba,MPLCONFIGDIR=/jobtmp/cache/matplotlib \
    --env OMP_NUM_THREADS="$SLURM_CPUS_PER_TASK",OPENBLAS_NUM_THREADS="$SLURM_CPUS_PER_TASK" \
    --env NUMBA_NUM_THREADS="$SLURM_CPUS_PER_TASK" \
    "$IMAGE" "$@"