#!/bin/bash
# Launcher script to avoid early-exit preventing output from being printed
PROJECT_DIR="$1"
TEST_TYPE="$2"

SBATCH_LOG="$PROJECT_DIR/nvshmem4py_sbatch.out"
TEST_LOG="$PROJECT_DIR/nvshmem4py_test.out"

source $PROJECT_DIR/docker/clusters/eos-config.sh
CONTAINER_IMAGE=$(get_build_tools_image $(get_build_image_version))

date

# launch sbatch
sbatch -W --container-image=$CONTAINER_IMAGE --container-mounts=$PROJECT_DIR/ \
    $PROJECT_DIR/nvshmem4py/test/python_test_wrapper.sh $PROJECT_DIR $TEST_TYPE > $SBATCH_LOG 2>&1 &
SBATCH_PID=$!

# TODO(vkhodel): Matrix uses a loop to wait for job id info
sleep 5

JOB_ID=$(cat $SBATCH_LOG | sed -e 's/^.* \([0-9]\+\)$/\1/')
echo JOB_ID is $JOB_ID

# TODO(vkhodel): CodeRabbit suggests to use a loop here to wait for Pending state
squeue -j $JOB_ID --start

# show output, wait for the file to appear, exit on sbatch exit
tail --retry --sleep-interval 5 --pid $SBATCH_PID -f $TEST_LOG &

wait $SBATCH_PID
SBATCH_EXIT_CODE=$?

date

# this might fail, let's skip for now
#EXIT_CODE=$(sacct -j "$JOB_ID" --format=ExitCode --noheader | awk -F: '{print $1}' | head -1)
#
#if [ -z "$EXIT_CODE" ]; then
#    echo "ERROR: No exit code for job $JOB_ID found"
#    exit 1
#fi
#
#echo "Slurm job exited with $EXIT_CODE"

echo "Sbatch exited with $SBATCH_EXIT_CODE"
exit $SBATCH_EXIT_CODE
