#!/bin/bash
#SBATCH --nodes=4                # num of dpnodes
#SBATCH --ntasks-per-node=3      # combination of num dpnodes and workers per node (should be one more than num workers per node)
#SBATCH --gres=gpu:2             # number of GPUs per node (should be one less than the num tasks per node)
#SBATCH --mem=60G                # mem per node
#SBATCH --time=0-00:20           # max time (DD-HH:MM)
#
# Usage: sbatch slurm-jobscript.sh N_file data_output_dir cudadl_path
# Example: sbatch slurm-jobscript.sh $HOME/N/N_70 $HOME/data $HOME/code/cudadl
#
TEMP_NFS_DIR="$HOME/controller_loc"
PROBLEM_FILE="$1"
SAVE_OUTPUT_DIR="$2"
CUDADL_PATH="$3"
#
##################################################################
#
NUM_DPNODES="$SLURM_JOB_NUM_NODES"
NUM_WORKERS="$(( SLURM_JOB_NUM_NODES*(SLURM_NTASKS_PER_NODE-1) ))"
MEM_PER_DPNODE="$(( SLURM_MEM_PER_NODE/1024 ))" # SLURM_MEM_PER_NODE will be in MB and we want GB
#
if [[ ! -e "$TEMP_NFS_DIR" ]]; then
	mkdir "$TEMP_NFS_DIR"
elif [[ ! -d "$TEMP_NFS_DIR" ]]; then
	echo "$TEMP_NFS_DIR already exists but is not a directory" 1>&2
	exit 1
fi
#
if [[ ! -e "$SAVE_OUTPUT_DIR" ]]; then
	mkdir "$SAVE_OUTPUT_DIR"
elif [[ ! -d "$SAVE_OUTPUT_DIR" ]]; then
	echo "$SAVE_OUTPUT_DIR already exists but is not a directory" 1>&2
	exit 1
fi
#
SAVE_OUTPUT_DIR="$SAVE_OUTPUT_DIR/$SLURM_JOB_ID"
mkdir "$SAVE_OUTPUT_DIR" || exit 1
#
TEMP_NFS_FILE=$(mktemp -u "${TEMP_NFS_DIR}/loc.XXXXXXXXXX")
#
echo "Starting with $NUM_DPNODES dpnodes ($MEM_PER_DPNODE GB memory) and $NUM_WORKERS workers."
#
srun ${CUDADL_PATH}/slurm-wrapper.sh $SAVE_OUTPUT_DIR $TEMP_NFS_FILE $CUDADL_PATH -w $NUM_WORKERS -n $NUM_DPNODES -m $MEM_PER_DPNODE $PROBLEM_FILE 1