#!/usr/bin/env bash

ARGPARSE_DESCRIPTION="Train recognizer on slurm cluster"

source $(dirname $0)/argparse.bash || exit 1

argparse "$@" <<EOF || exit 1

parser.add_argument('-p', '--partition', default='${PARTITION:-""}', type=str,

                    help='partition (default: env PARTITION)')

parser.add_argument('--gpus', default=${GPUS:-8}, type=int,

                    help='number of gpus to use (default: env GPUS, or 8)')

parser.add_argument('--gpus_per_node', default=${GPUS_PER_NODE:-8}, type=int,

                    help='gpus per node (default: env GPUS_PER_NODE, or 8)')

parser.add_argument('--cpus_per_task', default=${CPUS_PER_TASK:-5}, type=int,

                    help='cpus per task (default: env CPUS_PER_TASK, or 5)')

parser.add_argument('-j', '--job_name', default='', type=str,

                    help='name of job (default: config filename)')

parser.add_argument('--srun_args', default='${SRUN_ARGS:-""}', type=str,

                    help='extra args for srun (default: env SRUN_ARGS)')

parser.add_argument('config', type=str,

                    help='config file')

parser.add_argument('py_args', nargs=argparse.REMAINDER,

                    help='extra args to tools/train.py (can be empty)')

EOF

BARECONFIG="${CONFIG##*/}"

JOB_NAME="${JOB_NAME:-${BARECONFIG%.*}}"

echo partition: "${PARTITION}"

echo jobname: "${JOB_NAME}"

echo config: "${CONFIG}"

echo gpus: "${GPUS}"

echo gpus per node: "${GPUS_PER_NODE}"

echo cpus per task: "${CPUS_PER_TASK}"

echo srun args: "${SRUN_ARGS[@]}"

echo py args: "${PY_ARGS[@]}"

echo command to run:

echo ----

echo srun -p "${PARTITION}" \

    --job-name="${JOB_NAME}" \

    --gres=gpu:${GPUS_PER_NODE} \

    --ntasks=${GPUS} \

    --ntasks-per-node=${GPUS_PER_NODE} \

    --cpus-per-task=${CPUS_PER_TASK} \

    --kill-on-bad-exit=1 \

    ${SRUN_ARGS} \

    python -u tools/train.py ${CONFIG} --launcher="slurm" ${PY_ARGS[@]}

echo ----

echo

srun -p "${PARTITION}" \

    --job-name="${JOB_NAME}" \

    --gres=gpu:${GPUS_PER_NODE} \

    --ntasks=${GPUS} \

    --ntasks-per-node=${GPUS_PER_NODE} \

    --cpus-per-task=${CPUS_PER_TASK} \

    --kill-on-bad-exit=1 \

    ${SRUN_ARGS} \

    python -u tools/train.py ${CONFIG} --launcher="slurm" ${PY_ARGS[@]}