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

# 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 [TESTSUITE_NAME]"
    exit 1
else
    CI_PROJECT_HOME="$1"
fi

# Validate optional TESTSUITE_NAME argument (case-insensitive) and set TEST_SUITE
if [ -z "$2" ]; then
    TEST_SUITE="all"
else
    case "$(echo "$2" | tr '[:upper:]' '[:lower:]')" in
        core|numba|numba_high_level_1|numba_high_level_2|all)
            TEST_SUITE="$(echo "$2" | tr '[:upper:]' '[:lower:]')"
            ;;
        *)
            echo "Invalid TESTSUITE_NAME: $2"
            echo "Allowed values: core, numba, numba_high_level_1, numba_high_level_2, all"
            exit 1
            ;;
    esac
fi

date

# Setup Python env
# HACK! Remove when container has this in it:
apt update && apt install -y librdmacm-dev rdma-core
apt update && apt install -y ibverbs-providers
#apt update && apt install -y --reinstall libc6-dev

date

python3 -m venv /workspace/nvshmem/nvshmem4py_venv
source /workspace/nvshmem/nvshmem4py_venv/bin/activate
pip install --force-reinstall $CI_PROJECT_DIR/nvshmem_pkg/build/dist/nvshmem4py_cu12-0.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_34_x86_64.whl
pip install cuda-python
pip install mpi4py
pip install cupy-cuda12x
pip install cffi
pip install "numba-cuda[cu12]==0.20.1"
pip install cuda.core==0.4.0
# Execute python tests
# Add more tests here if/when they exist
pip install torch
pip install nvidia-cuda-nvcc-cu12==12.8.*
pip install pytest pytest-mpi

date

# HACK! Overwrite the headers and libraries in libnvidia-nvshmem wheel that NVSHMEM4Py will search with the ones from this branch
cp -rf $CI_PROJECT_DIR/nvshmem_pkg/include/* /workspace/nvshmem/nvshmem4py_venv/lib/python3.12/site-packages/nvidia/nvshmem/include/
cp -rf $CI_PROJECT_DIR/nvshmem_pkg/lib/* /workspace/nvshmem/nvshmem4py_venv/lib/python3.12/site-packages/nvidia/nvshmem/lib/

date

echo "================================================"
echo "PIP LIST"
pip list
echo "================================================"


export NVSHMEM_HOME=$CI_PROJECT_DIR/nvshmem_pkg
export RDMA_CORE_HOME=/usr
export LD_LIBRARY_PATH=/usr/local/openmpi/lib:$LD_LIBRARY_PATH

EXIT_CODE=0
export PATH=/usr/local/openmpi/bin:$PATH

run_core_tests() {
python3 $CI_PROJECT_DIR/nvshmem4py/test/wheel_sanity_test.py
if [ $? -ne 0 ]; then
    echo "Test failed: wheel sanity test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -x NVSHMEM_BOOTSTRAP=MPI -np 2 -- python3 $CI_PROJECT_DIR/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 --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/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 --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/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 --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/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 --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/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 --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/rma_test.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: RMA Test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 8 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/teams_test.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Team Management Test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi
}

rm -rf $1/nvshmem4py_venv
echo Test wrapper is exiting with $EXIT_CODE...

run_numba_tests() {
# NUMBA TESTS:
echo "================================================"
echo "NUMBA TESTS"
numba -s
echo PATH=$PATH
echo LD_LIBRARY_PATH=$LD_LIBRARY_PATH
echo CPATH=$CPATH
echo "================================================"

unset NVSHMEM_BOOTSTRAP
mpirun --allow-run-as-root  -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/test_get_version.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Device get version test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

unset NVSHMEM_BOOTSTRAP
mpirun --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/test_npe.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Device number of processing elementstest"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/test_ring.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Simple int put on ring test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/test_collective.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Simple test that requires cooperative launch."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi


export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/test_highlevel_bindings.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: High-level bindings POC test"
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi


export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- python3 $CI_PROJECT_DIR/nvshmem4py/test/test_ring_allreduce.py -i mpi
if [ $? -ne 0 ]; then
    echo "Test failed: Ring all-reduce example test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

}

run_numba_high_level_tests_1() {
# NUMBA high level API tests:
echo "================================================"
echo "NUMBA high level API tests batch 1 (Collective/Sync/Barrier)"
echo "================================================"

# Change to test directory, pytest requires launching from the test directory to locate conftest
pushd $CI_PROJECT_DIR/nvshmem4py/test/device/numba/ || exit 1

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- pytest --init-type mpi --with-mpi test_device_coll.py -v -s
if [ $? -ne 0 ]; then
    echo "Test failed: High-level API Collective test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- pytest --init-type mpi --with-mpi test_device_sync.py -v -s
if [ $? -ne 0 ]; then
    echo "Test failed: High-level API Barrier All Sync test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- pytest --init-type mpi --with-mpi test_device_barrier.py -v -s
if [ $? -ne 0 ]; then
    echo "Test failed: High-level API Barrier test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun -x NUMBA_CUDA_ENABLE_NRT=1 --allow-run-as-root -np 2 -- pytest --init-type mpi --with-mpi test_device_mem.py -v -s
if [ $? -ne 0 ]; then
    echo "Test failed: High-level API Memory test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

popd || exit 1
}

run_numba_high_level_tests_2() {
# NUMBA high level API tests:
echo "================================================"
echo "NUMBA high level API tests batch 1 (RMA/AMO)"
echo "================================================"

# Change to test directory, pytest requires launching from the test directory to locate conftest
pushd $CI_PROJECT_DIR/nvshmem4py/test/device/numba/ || exit 1

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- pytest --init-type mpi --with-mpi test_device_rma.py -v -s
if [ $? -ne 0 ]; then
    echo "Test failed: High-level API RMA test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi

export NVSHMEM_BOOTSTRAP=MPI
mpirun --allow-run-as-root -np 2 -- pytest --init-type mpi --with-mpi test_device_amo.py -v -s
if [ $? -ne 0 ]; then
    echo "Test failed: High-level API AMO test."
    EXIT_CODE=`expr $EXIT_CODE + 1`
fi


popd || exit 1
}

if [ "$TEST_SUITE" = "core" ]; then
    run_core_tests
elif [ "$TEST_SUITE" = "numba" ]; then
    run_numba_tests
elif [ "$TEST_SUITE" = "numba_high_level_1" ]; then
    run_numba_high_level_tests_1
elif [ "$TEST_SUITE" = "numba_high_level_2" ]; then
    run_numba_high_level_tests_2
else
    run_core_tests
    run_numba_tests
    run_numba_high_level_tests_1
    run_numba_high_level_tests_2
fi

rm -rf $CI_PROJECT_DIR/nvshmem4py_venv
exit $EXIT_CODE
