#!/bin/sh
#
# submission script for PMSCO calculations on the Ra cluster
#
# this version clones the current git repository at HEAD to the work directory.
# thus, version conflicts between jobs are avoided.
#

if [ $# -lt 1 ]; then
  echo "Usage: $0 [NOSUB] GIT_TAG DESTDIR JOBNAME NODES TASKS_PER_NODE WALLTIME:HOURS PROJECT [ARGS [ARGS [...]]]"
  echo ""
  echo "       NOSUB (optional): do not submit the script to the queue. default: submit."
  echo "       GIT_TAG: git tag or branch name of the code. HEAD for current code."
  echo "       DESTDIR: destination directory. must exist. a sub-dir \$JOBNAME is created."
  echo "       JOBNAME (text): name of job. use only alphanumeric characters, no spaces."
  echo "       NODES (integer): number of computing nodes. (1 node = 24 or 32 processors)."
  echo "          do not specify more than 2."
  echo "       TASKS_PER_NODE (integer): 1...24, or 32."
  echo "          24 or 32 for full-node allocation."
  echo "          1...23 for shared node allocation."
  echo "       WALLTIME:HOURS (integer): requested wall time."
  echo "          1...24 for day partition"
  echo "          24...192 for week partition"
  echo "          1...192 for shared partition"
  echo "       PROJECT: python module (file path) that declares the project and starts the calculation."
  echo "       ARGS (optional): any number of further PMSCO or project arguments (except time)."
  echo ""
  echo "the job script is written to \$DESTDIR/\$JOBNAME which is also the destination of calculation output."
  exit 1
fi

# location of the pmsco package is derived from the path of this script
SCRIPTDIR="$(dirname $(readlink -f $0))"
SOURCEDIR="$(readlink -f $SCRIPTDIR/..)"
PMSCO_SOURCE_DIR="$SOURCEDIR"

# read arguments
if [ "$1" == "NOSUB" ]; then
  NOSUB="true"
  shift
else
  NOSUB="false"
fi

if [ "$1" == "HEAD" ]; then
    BRANCH_ARG=""
else
    BRANCH_ARG="-b $1"
fi
shift

DEST_DIR="$1"
shift

PMSCO_JOBNAME=$1
shift

PMSCO_NODES=$1
PMSCO_TASKS_PER_NODE=$2
PMSCO_TASKS=$(expr $PMSCO_NODES \* $PMSCO_TASKS_PER_NODE)
shift 2

PMSCO_WALLTIME_HR=$1
PMSCO_WALLTIME_MIN=$(expr $PMSCO_WALLTIME_HR \* 60)
shift

# select partition
if [ $PMSCO_WALLTIME_HR -ge 25 ]; then
    PMSCO_PARTITION="week"
else
    PMSCO_PARTITION="day"
fi
if [ $PMSCO_TASKS_PER_NODE -lt 24 ]; then
    PMSCO_PARTITION="shared"
fi

PMSCO_PROJECT_FILE="$(readlink -f $1)"
shift

PMSCO_PROJECT_ARGS="$*"

# set up working directory
cd "$DEST_DIR"
if [ ! -d "$PMSCO_JOBNAME" ]; then
    mkdir "$PMSCO_JOBNAME"
fi
cd "$PMSCO_JOBNAME"
WORKDIR="$(pwd)"
PMSCO_WORK_DIR="$WORKDIR"

# copy code
PMSCO_SOURCE_REPO="file://$PMSCO_SOURCE_DIR"
echo "$PMSCO_SOURCE_REPO"

cd "$PMSCO_WORK_DIR"
git clone $BRANCH_ARG --single-branch --depth 1 $PMSCO_SOURCE_REPO pmsco || exit
cd pmsco
PMSCO_REV=$(git log --pretty=format:"%h, %ai" -1) || exit
cd "$WORKDIR"
echo "$PMSCO_REV" > revision.txt

# generate job script from template
sed -e "s:_PMSCO_WORK_DIR:$PMSCO_WORK_DIR:g" \
    -e "s:_PMSCO_JOBNAME:$PMSCO_JOBNAME:g" \
    -e "s:_PMSCO_NODES:$PMSCO_NODES:g" \
    -e "s:_PMSCO_WALLTIME_HR:$PMSCO_WALLTIME_HR:g" \
    -e "s:_PMSCO_PROJECT_FILE:$PMSCO_PROJECT_FILE:g" \
    -e "s:_PMSCO_PROJECT_ARGS:$PMSCO_PROJECT_ARGS:g" \
    "$SCRIPTDIR/pmsco.ra-git.template" > $PMSCO_JOBNAME.job

chmod u+x "$PMSCO_JOBNAME.job" || exit

# request nodes and tasks
#
# The option --ntasks-per-node is meant to be used with the --nodes option.
# (For the --ntasks option, the default is one task per node, use the --cpus-per-task option to change this default.)
#
# sbatch options
# --cores-per-socket=16
#   32 cores per node
# --partition=[shared|day|week]
# --time=8-00:00:00
#   override default time limit (2 days in long queue)
#   time formats: "minutes", "minutes:seconds", "hours:minutes:seconds", "days-hours", "days-hours:minutes", "days-hours:minutes:seconds"
# --mail-type=ALL
# --test-only
#   check script but do not submit
#
SLURM_ARGS="--nodes=$PMSCO_NODES --ntasks-per-node=$PMSCO_TASKS_PER_NODE"

if [ $PMSCO_TASKS_PER_NODE -gt 24 ]; then
    SLURM_ARGS="--cores-per-socket=16 $SLURM_ARGS"
fi

SLURM_ARGS="--partition=$PMSCO_PARTITION $SLURM_ARGS"

SLURM_ARGS="--time=$PMSCO_WALLTIME_HR:00:00 $SLURM_ARGS"

CMD="sbatch $SLURM_ARGS $PMSCO_JOBNAME.job"
echo $CMD
if [ "$NOSUB" != "true" ]; then
  $CMD
fi

exit 0