#!/bin/bash
#
# File:     slurm_submit.sh
# Author:   Jaime Frey (jfrey@cs.wisc.edu)
# Based on code by David Rebatto (david.rebatto@mi.infn.it)
#
# Description:
#   Submission script for SLURM, to be invoked by blahpd server.
#   Usage:
#     slurm_submit.sh -c <command> [-i <stdin>] [-o <stdout>] [-e <stderr>] [-w working dir] [-- command's arguments]
#
# Copyright (c) Members of the EGEE Collaboration. 2004. 
# Copyright (c) HTCondor Team, Computer Sciences Department,
#   University of Wisconsin-Madison, WI. 2015.
# 
# Licensed under the Apache License, Version 2.0 (the "License"); 
# you may not use this file except in compliance with the License. 
# You may obtain a copy of the License at 
# 
#     http://www.apache.org/licenses/LICENSE-2.0 
# 
# Unless required by applicable law or agreed to in writing, software 
# distributed under the License is distributed on an "AS IS" BASIS, 
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
# See the License for the specific language governing permissions and 
# limitations under the License.
#

. `dirname $0`/blah_common_submit_functions.sh

# Default values for configuration variables
slurm_std_storage=${slurm_std_storage:-/dev/null}
slurm_binpath=${slurm_binpath:-/usr/bin}

bls_submit_args_prefix="#SBATCH"

bls_parse_submit_options "$@"

bls_setup_all_files

# Write wrapper preamble
cat > $bls_tmp_file << end_of_preamble
#!/bin/bash
# SLURM job wrapper generated by `basename $0`
# on `/bin/date`
#
# stgcmd = $bls_opt_stgcmd
# proxy_string = $bls_opt_proxy_string
# proxy_local_file = $bls_proxy_local_file
#
# SLURM directives:
#SBATCH -o $slurm_std_storage
#SBATCH -e $slurm_std_storage
end_of_preamble

if [ "x$bls_opt_project" != "x" ] ; then
  echo "#SBATCH -A $bls_opt_project" >> $bls_tmp_file
fi

if [ "x$bls_opt_runtime" != "x" ] ; then
  echo "#SBATCH -t $((bls_opt_runtime / 60))" >> $bls_tmp_file
fi

if [ "x$bls_opt_req_mem" != "x" ]
then
  # Different schedulers require different memory checks
  echo "#SBATCH --mem=${bls_opt_req_mem}" >> $bls_tmp_file
fi

# Write SLURM directives according to command line options
# Map the queue option to slurm's partition option
# handle queue/partition overriding
cluster_name=`echo "$bls_opt_queue" | cut -s -f2 -d@`
if [ "$cluster_name" != "" ] ; then
  bls_opt_queue=`echo "$bls_opt_queue" | cut -f1 -d@`
fi
[ -z "$bls_opt_queue" ] || grep -q "^#SBATCH -p" $bls_tmp_file || echo "#SBATCH -p $bls_opt_queue" >> $bls_tmp_file
[ -z "$cluster_name" ] || grep -q "^#SBATCH -M" $bls_tmp_file || echo "#SBATCH -M $cluster_name" >> $bls_tmp_file

# Extended support for MPI attributes
if [ "x$bls_opt_wholenodes" == "xyes" ] ; then
  bls_opt_hostsmpsize=${bls_opt_hostsmpsize:-1}
  if [[ ! -z "$bls_opt_smpgranularity" ]] ; then
    if [[ -z "$bls_opt_hostnumber" ]] ; then
      echo "#SBATCH -N 1" >> $bls_tmp_file
      echo "#SBATCH --ntasks-per-node $bls_opt_smpgranularity" >> $bls_tmp_file
    else
      echo "#SBATCH -N $bls_opt_hostnumber" >> $bls_tmp_file
      echo "#SBATCH --ntasks-per-node $bls_opt_smpgranularity" >> $bls_tmp_file
    fi
    echo "#SBATCH --exclusive" >> $bls_tmp_file
  else
    if [[ ! -z "$bls_opt_hostnumber" ]] ; then
      if [[ $bls_opt_mpinodes -gt 0 ]] ; then
        echo "#SBATCH -N $bls_opt_hostnumber" >> $bls_tmp_file
        echo "#SBATCH -n $bls_opt_mpinodes" >> $bls_tmp_file
      else
        echo "#SBATCH -N $bls_opt_hostnumber" >> $bls_tmp_file
      fi
      echo "#SBATCH --exclusive" >> $bls_tmp_file
    fi
  fi
else
  if [[ ! -z "$bls_opt_mpinodes" ]] ; then
    echo "#SBATCH -n $bls_opt_mpinodes" >> $bls_tmp_file
  fi
  if [[ ! -z "$bls_opt_smpgranularity" ]] ; then
    echo "#SBATCH --ntasks-per-node $bls_opt_smpgranularity" >> $bls_tmp_file
  elif [[ ! -z "$bls_opt_hostnumber" ]] ; then
    echo "#SBATCH -N $bls_opt_hostnumber" >> $bls_tmp_file
  fi
fi
# --- End of MPI directives

# add GPU support
[ -z "$bls_opt_gpunumber" ] || [ -z "$bls_opt_gpumodel" ] || echo "#SBATCH --gres=gpu:${bls_opt_gpumodel}:${bls_opt_gpunumber}" >> $bls_tmp_file
[ -z "$bls_opt_gpunumber" ] || [ ! -z "$bls_opt_gpumodel" ] || echo "#SBATCH --gres=gpu:${bls_opt_gpunumber}" >> $bls_tmp_file
[ ! -z "$bls_opt_gpunumber" ] || [ -z "$bls_opt_gpumodel" ] || echo "#SBATCH --gres=gpu:${bls_opt_gpumodel}:1" >> $bls_tmp_file

#add MIC support
[ -z "$bls_opt_micnumber" ] || echo "#SBATCH --gres=mic:${bls_opt_micnumber}" >> $bls_tmp_file

# Do the local and extra args after all #SBATCH commands, otherwise slurm ignores anything
# after a non-#SBATCH command
bls_set_up_local_and_extra_args

# Input and output sandbox setup.
# Assume all filesystems are shared.

bls_add_job_wrapper
bls_save_submit

###############################################################
# Submit the script
###############################################################

datenow=`date +%Y%m%d`
jobID=`${slurm_binpath}/sbatch $bls_tmp_file` # actual submission
retcode=$?

if [ "$retcode" != "0" ] ; then
	rm -f $bls_tmp_file
	echo "Error from sbatch: $jobID" >&2
	exit 1
fi

# The job id is given by the line "Submitted batch job ###".
# Some installations print additional lines before or after this one.
jobID=`echo "$jobID" | awk '/^Submitted batch job [0-9]+/ {print $4}'`
if [ "X$jobID" == "X" ]; then
	rm -f $bls_tmp_file
	echo "Error: job id missing" >&2
	echo Error # for the sake of waiting fgets in blahpd
	exit 1
fi

# Compose the blahp jobID ("slurm/" + datenow + jobid [+ cluster])
blahp_jobID="slurm/`basename $datenow`/$jobID"
if [ "$cluster_name" != "" ] ; then
	blahp_jobID="${blahp_jobID}@${cluster_name}"
fi

echo "BLAHP_JOBID_PREFIX$blahp_jobID"
  
bls_wrap_up_submit

exit $retcode
