#!/bin/bash
#
# File:     flux_submit.sh
# Author:   Ian Lumsden (ilumsden@vols.utk.edu)
# Author:   Michela Taufer (mtaufer@utk.edu)
#
# Flux support was added as part of work supported by the US National Science Foundation
# under Grant No. 2530461, 2513101, 2331152, 2223704, 2138811, and 2103845.
# 
# Flux support was also added under the auspices of the U.S. Department of Energy by
# Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344
# and was supported by the LLNL-LDRD Program under Project No. 24-SI-005.
#
# Based on code by David Rebatto (david.rebatto@mi.infn.it)
#
# Description:
#   Submission script for Flux, to be invoked by blahpd server.
#   Usage:
#     flux_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.
#

# Source blah_common_submit_function.sh to get access to common functions
. `dirname $0`/blah_common_submit_functions.sh

. `dirname $0`/flux_utils.sh

# Set bls_submit_args_prefix so that any extra Flux directives are prefixed with
# "#FLUX: "
bls_submit_args_prefix="#FLUX:"

# Parse command-line options
bls_parse_submit_options "$@"

# Setup any necessary files
bls_setup_all_files

# Write wrapper preamble
# 
# This process considers the following BLAHP arguments
#   * Out: specifies the file for STDOUT
#   * Err: specifies the file for STDERR
#   * BatchProject: specifies the project/account for submission
#   * BatchRuntime: specifies the time limit for submission
#   * Queue: specifies the schedule queue/partition for submission (and optionally the cluster name)
#   * WholeNodes: if true, indicates that we want exclusive node allocation (i.e., no shared nodes)
#   * NodeNumber: specifies the total number of cores/tasks
#   * HostNumber: specifies the total number of nodes
#   * SMPGranularity: specifies the number of cores per node
#   * GPUNumber: specifies the number of GPUs per node
#   * BatchExtraSbumitArgs: specifies user-defined extra arguments.
#     * Note: this is not handled by this script. Instead, extra arguments are injected
#             by the bls_set_up_local_and_extra_args function from blah_common_submit_functions.sh
#
# As part of this, set Flux's "--output" flag to BLAHP's "Out" argument,
# and set Flux's "--error" flag to BLAHP's "Err" argument
cat > $bls_tmp_file << end_of_preamble
#!/bin/bash
# Flux 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
#
# Flux directives:
#FLUX: --output=$bls_wrapper_stdout
#FLUX: --error=$bls_wrapper_stderr
end_of_preamble

# Write Flux directives into preamble according to command-line options

# If a project/bank is specified (i.e., BLAHP BatchProject), add "--bank" to Flux
if [ "x$bls_opt_project" != "x" ] ; then
  echo "#FLUX: --bank=$bls_opt_project" >> $bls_tmp_file
fi

# If a time limit is specified (i.e., BLAHP BatchRuntime), add "--time-limit" to Flux
if [ "x$bls_opt_runtime" != "x" ] ; then
  echo "#FLUX: --time-limit=$((bls_opt_runtime / 60))" >> $bls_tmp_file
fi

# TODO if support for memory request is added to Flux, add support for RequestMemory

# If the queue/cluster name are specified (i.e., BLAHP Queue), do two things:
#   1. Split value on "@". The first part is the actual queue name. The second part
#      is the cluster name
#   2. If the queue name was provided, add "--queue" to Flux
# TODO if support for submission to other clusters is added to Flux (likely), add support for cluster_name
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 "^#FLUX: --queue" $bls_tmp_file || echo "#FLUX: --queue=$bls_opt_queue" >> $bls_tmp_file

# Process MPI/resource provisioning arguments

# If BLAHP WholeNodes is specified, add "--exclusive" to Flux
if [ "x$bls_opt_wholenodes" == "xyes" ]; then
  echo "#FLUX: --exclusive" >> $bls_tmp_file
fi
# Derive number of nodes, number of slots, and number of cores per slot from BLAHP's
# NodeNumber, HostNumber, and SMPGranularity arguments
flux_nnodes=
flux_nslots=
flux_ncores_per_slot=
# If number of nodes is specified or can be derived, get number of nodes and the larger possible number of cores/node
if [ ! -z "$bls_opt_hostnumber" ] || ([ ! -z "$bls_opt_smpgranularity" ] && [ ! -z "$bls_opt_mpinodes"]); then
  # If HostNumber is specified, we use it to get number of nodes
  if [ ! -z "$bls_opt_hostnumber" ]; then
    flux_nnodes=$bls_opt_hostnumber
  # If HostNumber is not specified, we get number of nodes from ceil(MpiNodes / SMPGranularity)
  else
    flux_nnodes=$(( (bls_opt_mpinodes + bls_opt_smpgranularity  - 1) / bls_opt_smpgranularity ))
  fi
  mpinodes_cores_per_slot=
  # To get the ncores_per_slot value associated with MpiNodes, we use ceil(MpiNodes / nnodes)
  if [ ! -z "$bls_opt_mpinodes" ]; then
    mpinodes_cores_per_slot=$(( (bls_opt_mpinodes + flux_nnodes - 1) / flux_nnodes ))
  fi
  # If we have cores_per_slot numbers associated with both MpiNodes and SMPGranularity,
  # we choose the larger of the two as flux_ncores_per_slot
  if [ ! -z "$mpinodes_cores_per_slot" ] && [ ! -z "$bls_opt_smpgranularity" ]; then
    if [ "$mpinodes_cores_per_slot" -ge "$bls_opt_smpgranularity" ]; then
      flux_ncores_per_slot=$mpinodes_cores_per_slot
    else
      flux_ncores_per_slot=$bls_opt_smpgranularity
    fi
  # If only the cores_per_slot number derived from MpiNodes exists, we set
  # flux_ncores_per_slot with that number
  elif [ ! -z "$mpinodes_cores_per_slot" ]; then 
    flux_ncores_per_slot=$mpinodes_cores_per_slot
  # If only SMPGranularity exists, we set flux_ncores_per_slot to SMPGranularity
  elif [ ! -z "$bls_opt_smpgranularity" ]; then
    flux_ncores_per_slot=$bls_opt_smpgranularity
  # If there is no provided or derived cores_per_slot value, we use the Flux default by not
  # providing a numbeer of cores per slot
  fi
  # Set the number of slots to equal the number of nodes so that the
  # --cores-per-slot and --gpus-per-slot flags act as if they are per-node
  flux_nslots=$flux_nnodes
# Else if only number of slots can be derived
elif [ ! -z "$bls_opt_mpinodes" ]; then
  flux_nslots=$bls_opt_mpinodes
  flux_ncores_per_slot=1
# Else we cannot derive the job shape, and this request is an error
# else
#   echo "Error: invalid resource specification"
# 	echo Error # for the sake of waiting fgets in blahpd
# 	exit 1
fi
# If flux_nnodes is defined, add "--nodes" to Flux
if [ ! -z "$flux_nnodes" ]; then
  echo "#FLUX: --nodes=$flux_nnodes" >> $bls_tmp_file
fi
# If flux_nslots is defined, add "--nslots" to Flux
if [ ! -z "$flux_nslots" ]; then
  echo "#FLUX: --nslots=$flux_nslots" >> $bls_tmp_file
fi
# If flux_ncores_per_slot is defined, add "--cores-per-slot" to Flux
if [ ! -z "$flux_ncores_per_slot" ]; then
  echo "#FLUX: --cores-per-slot=$flux_ncores_per_slot" >> $bls_tmp_file
fi

# If the number of GPUs per node (i.e., BLAHP GPUNumber) is specified and the number of
# nodes is derived above, add "--gpus-per-slot" to Flux
# TODO if support for adding GPU model information is added to Flux, add support for GPUModel
if [ ! -z "$flux_nnodes" ] && [ ! -z "$bls_opt_gpunumber" ]; then
  echo "#FLUX: --gpus-per-slot=$bls_opt_gpunumber" >> $bls_tmp_file
fi

# TODO if support for MIC is added to Flux (unlikely), add support for MICNumber

# Do the local and extra args after all #fLUX commands, otherwise Flux ignores anything
# after a non-#FLUX 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
###############################################################

jobID=`${flux_binpath}flux batch $bls_tmp_file` # actual submission
retcode=$?

# Verify that submission was successful
if [ "$retcode" != "0" ] ; then
	rm -f $bls_tmp_file
	echo "Error from flux batch: $jobID" >&2
	echo Error # for the sake of waiting fgets in blahpd
	exit 1
fi

# Ensure we properly got a job ID
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 ("flux/" + datenow + jobid [+ @cluster])
flux_utils_create_blahp_jobid "$jobID"

echo "BLAHP_JOBID_PREFIX$blahp_jobID"
  
bls_wrap_up_submit

exit $retcode
