#!/bin/bash
#SBATCH --job-name=coreai_libraries_nvshmem-ci-jobs.nvshmem4py
#SBATCH --ntasks=2
#SBATCH --time=00:05:00         
#SBATCH --account coreai_libraries_nvshmem
#SBATCH --partition=batch
#SBATCH --output=nvshmem4py_test.out
#SBATCH --error=nvshmem4py_test.out

# Define log file based on job name and ID
LOG_FILE="$(basename $SLURM_JOB_NAME)_${SLURM_JOB_ID}.out"
set -x

# Ensure an argument (CI_PROJECT_HOME) is provided
if [ -z "$1" ]; then
    echo "Usage: sbatch python_test_wrapper.sh CI_PROJECT_HOME"
    exit 1
fi

# Setup Python env
source /home/benjaming/sourceme/eos.sh
pip install --force-reinstall $1/nvshmem_pkg/lib/python/dist/nvshmem4py_cu12-0.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_34_x86_64.whl
pip install cuda-python cuda.core
pip install mpi4py
pip install cupy-cuda12x
setrunpath
# Execute python tests
# Add more tests here if/when they exist
EXIT_CODE=0
python3 $1/nvshmem4py/test/wheel_sanity_test.py
if [ $? -ne 0 ]; then
    echo "Test failed: wheel sanity test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export LD_LIBRARY_PATH=$1/nvshmem_pkg/lib/:$LD_LIBRARY_PATH
export NVSHMEM_BOOTSTRAP=MPI
mpirun -x LD_PRELOAD=$1/nvshmem_pkg/lib/libnvshmem_host.so.3 -x NVSHMEM_BOOTSTRAP=MPI -np 2 -- python3 nvshmem4py/test/init_test.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: MPI init test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

unset NVSHMEM_BOOTSTRAP
mpirun -x LD_PRELOAD=$1/nvshmem_pkg/lib/libnvshmem_host.so.3 -np 2 -- python3 nvshmem4py/test/init_test.py -i uid
if [ $? -ne 0 ]; then
    echo "Test failed: UID init test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

unset NVSHMEM_BOOTSTRAP
mpirun -x LD_PRELOAD=$1/nvshmem_pkg/lib/libnvshmem_host.so.3 -np 2 -- python3 nvshmem4py/test/init_test.py -i emulated_mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Emulated MPI init test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun -x LD_PRELOAD=$1/nvshmem_pkg/lib/libnvshmem_host.so.3 -np 2 -- python3 nvshmem4py/test/memory_test.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Memory Test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun -x LD_PRELOAD=$1/nvshmem_pkg/lib/libnvshmem_host.so.3 -np 2 -- python3 nvshmem4py/test/collective_test.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Collectives Test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun -x LD_PRELOAD=$1/nvshmem_pkg/lib/libnvshmem_host.so.3 -np 2 -- python3 nvshmem4py/test/rma_test.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: RMA Test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

rm -rf $1/nvshmem4py_venv
exit $EXIT_CODE
