#!/usr/bin/env bash

HELP_MSG="\
Script for running a pytest test suite in parallel across multiple jobs.

Only the tests that use the given number of processors are selected from the suite. This list of tests is distributed between multiple jobs and each job outputs its own log file.

Usage:

    firedrake-run-split-tests <nprocs> <njobs> <pytest_args...>

  where:
    * <nprocs> is the number of ranks used in each test
    * <njobs> is the number of different jobs
    * <pytest_args...> are additional arguments that are passed to pytest

  The following environment variables can be used to configure the
  outer process-tree timeout for each split job:

    * FIREDRAKE_RUN_SPLIT_TESTS_TIMEOUT: maximum wall time for each job
      (default: 3600s)
    * FIREDRAKE_RUN_SPLIT_TESTS_KILL_AFTER: grace period before forcibly
      killing a timed-out job (default: 60s)

Example:

    firedrake-run-split-tests 3 4 tests/unit --verbose

  will run all of the parallel[3] tests inside tests/unit verbosely
  and split between 4 different jobs.

Run with [no arguments | -h | --help] to print this help message.

Requires:

  * pytest
  * pytest-split
  * mpi-pytest

Optional:

  * GNU parallel (if unavailable, the script falls back to bash job control)
  * GNU timeout or gtimeout (if unavailable, jobs run without an outer timeout)"

# Print out help message with no arguments or "-h" or "--help"
if [[ "$#" -eq "0" ]] || [[ "$1" == "-h" ]] || [[ "$1" == "--help" ]]; then
   echo -e "${HELP_MSG}"
   exit
fi

num_procs=$1
num_jobs=$2
extra_args=${@:3}
job_timeout=${FIREDRAKE_RUN_SPLIT_TESTS_TIMEOUT:-3600s}
kill_after=${FIREDRAKE_RUN_SPLIT_TESTS_KILL_AFTER:-60s}

# Callback to kill child processes if Ctrl-C hit
pids=()
cleanup() {
    for pid in "${pids[@]}"; do
        if ps -p "$pid" > /dev/null; then
            kill -INT "$pid"  # Send SIGINT to the sub-process
            wait "$pid"
        fi
    done
}
trap cleanup SIGINT

timeout_exec=""
if command -v timeout >/dev/null 2>&1; then
    timeout_exec="timeout"
elif command -v gtimeout >/dev/null 2>&1; then
    timeout_exec="gtimeout"
else
    echo "Warning: GNU timeout/gtimeout not found; running split jobs without an outer timeout" >&2
fi

if [ $num_procs = 1 ]; then
    # Cannot use mpiexec -n 1 because this can sometimes hang with
    # OpenMPI at MPI_Finalize
    pytest_exec="python3 -m pytest"
else
    # Set '--bind-to none' so we don't end up pinning MPI jobs to the same cores
    pytest_exec="mpiexec -n ${num_procs} --bind-to none python3 -m pytest"
fi
pytest_cmd="${pytest_exec} -v \
            --splits ${num_jobs} --group {#} \
            -m parallel[match] ${extra_args}"
if [ -n "${timeout_exec}" ]; then
    pytest_cmd="${timeout_exec} --kill-after=${kill_after} ${job_timeout} ${pytest_cmd}"
fi

log_file_prefix="pytest_nprocs${num_procs}_job"

# This incantation:
# * Runs pytest under GNU parallel using the right number of jobs
# * Applies an outer timeout to the whole pytest/mpiexec process tree, if available
# * Uses tee to pipe stdout+stderr to both stdout and a log file
# * Writes pytest's exit code to a file called jobN.errcode (for later inspection)
if command -v parallel >/dev/null 2>&1; then
    # Print the command
    set -x
    parallel --line-buffer --tag \
        "${pytest_cmd} 2>&1 | tee ${log_file_prefix}{#}.log; \
        echo \${PIPESTATUS[0]} > job{#}.errcode" \
        ::: $(seq ${num_jobs}) &
    set +x
    pid=$!
    pids+=($pid)
    wait ${pid} || true
else
    # Fallback when GNU parallel is not available: use bash job control.
    # We keep per-job logs and errcodes to match the GNU parallel behavior.
    # Print the commands
    set -x
    for i in $(seq ${num_jobs}); do
        (
            eval "${pytest_cmd/\{#\}/$i}" 2>&1 | tee ${log_file_prefix}${i}.log
            echo ${PIPESTATUS[0]} > job${i}.errcode
        ) &
        pids+=($!)
    done
    set +x
    for pid in "${pids[@]}"; do
        wait ${pid} || true
    done
fi

pass=true
for i in $(seq 1 ${num_jobs}); do
    # If we terminated early then this file doesn't exist
    if [ -f "job${i}.errcode" ]; then
        error_code=$(cat job${i}.errcode)
    else
        error_code=1
    fi
    # pytest uses exit code 5 if no tests were found, which we also treat that as a success
    # (see https://docs.pytest.org/en/7.1.x/reference/exit-codes.html)
    if [ ${error_code} = "0" ] || [ ${error_code} = "5" ]; then
        echo Job ${i} passed
    else
        echo Job ${i} failed, inspect the logs in ${log_file_prefix}${i}.log
        pass=false
    fi
done

echo Cleaning up
rm -f job*.errcode
echo Done

if $pass; then
    exit 0
else
    exit 1
fi
