#!/usr/bin/env bash

export MASTER_PORT=$((12000 + $RANDOM % 20000))

set -x

PARTITION=$1

JOB_NAME=$2

CONFIG=$3

GPUS=${GPUS:-8}

GPUS_PER_NODE=${GPUS_PER_NODE:-8}

CPUS_PER_TASK=${CPUS_PER_TASK:-5}

SRUN_ARGS=${SRUN_ARGS:-""}

PY_ARGS=${@:4}  # Any arguments from the forth one are captured by this

PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \

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}