#!/bin/bash
#
# File:     sfapi_submit.sh
# Author:   Karan Vahi (vahi@isi.edu)
#
# Description:
#   Submission script for SLURM via NERSC Superfacility API (SFAPI),
#   to be invoked by blahpd server. Submits an HTCondor job to a remote
#   Slurm cluster (Perlmutter at NERSC) using the SFAPI Python client
#   via functions defined in sfapi_helpers.py.
#   Usage:
#     sfapi_submit.sh -c <command> [-i <stdin>] [-o <stdout>] [-e <stderr>]
#                     [-w working_dir] [-q queue] [-A project] [-t runtime]
#                     [-n mpinodes] [-N hostsmpsize] [-h hostnumber]
#                     [-g gpunumber] [-M gpumodel] [-- command_arguments]
#
#   Credentials are read from ~/.superfacility/clientid.txt and
#   ~/.superfacility/priv_key.jwk by sfapi_helpers.load_sflapi_client_secret().
#
# 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
. `dirname $0`/sfapi_setup.sh

bls_submit_args_prefix="#SBATCH"

# transfer_executable value is not supported in blahp
# similar functionality by setting the variable
# _CONDOR_TRANSFER_EXECUTABLE=false in the
# environment variable in the condor submit file
bls_transfer_executable_env_regex="^_CONDOR_TRANSFER_EXECUTABLE=false$"

function bls_set_job_env ()
{
    # Karan: Copied from bls_start_job_wrapper() function in blah_common_submit_functions
    # TODO: Should be a separate function in common functions. Ask Jaime
    #
    # Set the required environment variables (escape values with double quotes)
    # We don't merge the existing PATH with the job-defined one if so configured
    if [ "x$bls_opt_environment" != "x" ] ; then
            echo ""
            echo "#1 Setting the environment:"
      eval "env_array=($bls_opt_environment)"
            for  env_var in "${env_array[@]}"; do
                     echo export \"$env_var\"
                     if [[ $env_var =~ $bls_transfer_executable_env_regex ]] ; then
                         bls_opt_stgcmd="no"
                     fi
            done
    else
            if [ "x$bls_opt_envir" != "x" ] ; then
                    echo ""
                    echo "#2 Setting the environment:"
                    echo "`echo ';'$bls_opt_envir | sed -e 's/;[^=]*;/;/g' -e 's/;[^=]*$//g' | sed -e 's/;\([^=]*\)=\([^;]*\)/;export \1=\"\2\"/g' | awk 'BEGIN { RS = ";" } ; { print $0 }'`"
            fi
    fi

    print_blahp_job_env
}

function bls_get_file_path ()
{
#
# Usage: bls_get_file_path file path
# updates the file path against the ${bls_opt_workdir} if file path is not
# an absolute path
# ${bls_opt_workdir} is the intialdir for the job or the submit directory
#
# Result is returned in $bls_get_file_path_result.
#
    local file_path=$1
    if [ "${file_path:0:1}" != "/" ] ;
        then file_path=${bls_opt_workdir}/${file_path} ;
    fi
    bls_get_file_path_result=${file_path}
}

bls_parse_submit_options "$@"

# karan: we are submitting job remotely. so this should be set to no to get
# full paths for stdout and stderr
#export bls_is_in_shared_dir="no"

# Use only temp file setup; full bls_setup_all_files is skipped because
# stdin/stdout/stderr are handled remotely via sfapi_helpers (not local sandbox).
bls_setup_temp_files


# Mirror the debug-dir setup from bls_setup_all_files so bls_save_submit works.
#bls_wrapper_stdout="/dev/null"
#bls_wrapper_stderr="/dev/null"
if [ -d "$blah_debug_save_submit_info" ] && [ -n "$bls_tmp_name" ]; then
    bls_info_dir="$blah_debug_save_submit_info/$bls_tmp_name.debug"
    mkdir "$bls_info_dir"
#    if [ -d "$bls_info_dir" ]; then
#        bls_wrapper_stdout="$bls_info_dir/wrapper.stdout"
#        bls_wrapper_stderr="$bls_info_dir/wrapper.stderr"
#    fi
fi

###############################################################
# Build the SBATCH job script
###############################################################
# Note: #SBATCH --output, --error, --chdir are prepended by
# sfapi_helpers.submit_remote_slurm_job() with remote Perlmutter paths,
# so we do NOT set them here.

cat > $bls_tmp_file << end_of_preamble
#!/bin/bash
# SLURM job wrapper for SFAPI (Perlmutter) generated by `basename $0`
# on `/bin/date`
#
# stgcmd = $bls_opt_stgcmd
# proxy_string = $bls_opt_proxy_string
#
# the paths below are the local paths where
# sfapi_status.sh will download the stdout and stderr
# the submit command of sfapi_helpers.py sets the
# remote stdout and stderr
# local_stdout = ${bls_opt_workdir}/$bls_opt_stdout
# local_stderr = ${bls_opt_workdir}/$bls_opt_stderr
#
# SLURM directives:
#
#
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
    echo "#SBATCH --mem=${bls_opt_req_mem}" >> $bls_tmp_file
fi

# Map the queue option to slurm's partition; handle cluster override
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

# 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

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

# Append local submit attributes and any extra #SBATCH args.
# Must come after all other #SBATCH directives or Slurm ignores them.
bls_set_up_local_and_extra_args

# set the job environment
bls_set_job_env >> $bls_tmp_file

###############################################################
# Job body: run the uploaded executable in the remote directory.
# sfapi_helpers sets --chdir to the remote job directory so all
# uploaded files are accessible by their basename.
###############################################################

sfapi_input_files=()

bls_command_basename="`basename $bls_opt_the_command`"

{
    echo ""
    echo "# Execute the  job command"
    # Karan TODO: Only if transfer_executable is set
    # we use the basename
    run_cmd=${bls_opt_the_command}
    if [ "x${bls_opt_stgcmd}" == "xyes" ] ; then
        run_cmd="./${bls_command_basename}"
        echo "chmod u+x \"${run_cmd}\""

        # figure out the executable path to use for transfer
        bls_get_file_path ${bls_opt_the_command}
        bls_opt_the_command=$bls_get_file_path_result
        if [ -f "$bls_opt_the_command" ] ; then
            sfapi_input_files+=("${bls_opt_the_command}")
        else
            echo "Error: command not found locally: $bls_opt_the_command" >&2
            rm -f $bls_tmp_file
            exit 1
        fi
    fi


    # If stdin was specified and is not /dev/null,
    # it will be uploaded; reference it by basename.
    if [[ -n "$bls_opt_stdin" && x"$bls_opt_stdin" != "x/dev/null" ]] ; then
        run_cmd="${run_cmd} < \"./$( basename "$bls_opt_stdin" )\""
    fi

    # stdout/stderr flow to #SBATCH --output/--error set by sfapi_helpers.py.
    # bls_arguments holds only the raw positional args after '--'.
    echo "${run_cmd} ${bls_arguments}"
    echo "exit \$?"
} >> $bls_tmp_file

bls_save_submit


###############################################################
# Collect input files to upload to Perlmutter
###############################################################


# pick up transfer_input_files specified by the user
# the filenames are in a file identified by $bls_opt_inputflstring
# e.g. /tmp/condor_g_scratch.0x55d2fcb281c0.542/InputFileList_3_17781864805412
if [ ! -z "$bls_opt_inputflstring" ] ; then
    for file in `cat $bls_opt_inputflstring`; do
        bls_get_file_path $file
        file=$bls_get_file_path_result
        sfapi_input_files+=("$file")
    done
fi

###############################################################
# Collect output files to download from Perlmutter
###############################################################
sfapi_output_files=()

# pick up transfer_input_files specified by the user
# the filenames are in a file identified by $bls_opt_outputflstring
if [ ! -z "$bls_opt_outputflstring" ] ; then
    for file in `cat $bls_opt_outputflstring`; do
        bls_get_file_path $file
        file=$bls_get_file_path_result
        sfapi_output_files+=("$file")
    done
fi


# Upload stdin file if specified
if [ -n "$bls_opt_stdin" ] && [ -f "$bls_opt_stdin" ] && [ x"$bls_opt_stdin" != "x/dev/null" ] ; then
    bls_get_file_path $bls_opt_stdin
    bls_opt_stdin=$bls_get_file_path_result
    sfapi_input_files+=("$bls_opt_stdin")
fi

###############################################################
# Submit via NERSC SFAPI
###############################################################

datenow=`date +%Y%m%d`
sfapi_helpers_dir=`dirname $0`

# Join the input-files array into a comma-separated string for --input-files
sfapi_input_files_csv=$(IFS=,; echo "${sfapi_input_files[*]}")
command="python3 "$sfapi_helpers_dir/sfapi_helpers.py" submit \
    --job-name    "$bls_tmp_name" \
    --input-files "$sfapi_input_files_csv" \
    --script      "$bls_tmp_file")"

sfapi_result=$(python3 "$sfapi_helpers_dir/sfapi_helpers.py" submit \
    --job-name    "$bls_tmp_name" \
    --input-files "$sfapi_input_files_csv" \
    --script      "$bls_tmp_file" 2>&1)
retcode=$?

if [ "$retcode" != "0" ] ; then
    rm -f $bls_tmp_file
    echo "Error submitting command $command"
    echo "Error submitting job via SFAPI: $sfapi_result" >&2
    exit 1
fi

# Parse "SFAPI_RESULT:<jobid>:<remote_dir>:<remote_stdout>:<remote_stderr>"
sfapi_line=$(echo "$sfapi_result" | grep "^SFAPI_RESULT:")
if [ -z "$sfapi_line" ] ; then
    rm -f $bls_tmp_file
    echo "Error: no SFAPI_RESULT in submission output" >&2
    echo "Output: $sfapi_result" >&2
    echo Error
    exit 1
fi

jobID=$(echo "$sfapi_line"        | cut -d: -f2)
remote_dir=$(echo "$sfapi_line"    | cut -d: -f3)
remote_stdout=$(echo "$sfapi_line" | cut -d: -f4)
remote_stderr=$(echo "$sfapi_line" | cut -d: -f5)

if [ "X$jobID" == "X" ] ; then
    rm -f $bls_tmp_file
    echo "Error: job id missing from SFAPI response" >&2
    echo Error
    exit 1
fi

# Compose the blahp job ID: "sfapi/" + datenow + "/" + jobID
blahp_jobID="sfapi/`basename $datenow`/$jobID"

# Persist the remote output paths so sfapi_status.sh can retrieve them later
sfapi_state_dir="${HOME}/.blah/sfapi_jobs"
mkdir -p "$sfapi_state_dir"

# resolve stdout and stderr path against bls_opt_workdir
bls_get_file_path $bls_opt_stdout
bls_opt_stdout=$bls_get_file_path_result

bls_get_file_path $bls_opt_stderr
bls_opt_stderr=$bls_get_file_path_result

# map local stdout and stderr to their remote counterparts on perlmutter
read -r -d '' fragment << EOM
# type::<local file on submit host>:<remote file to retrieve via sfapi>
remote_dir::${remote_dir}
stdout::${bls_opt_stdout}:${remote_stdout}
stderr::${bls_opt_stderr}:${remote_stderr}
EOM

jobstate_file=${sfapi_state_dir}/`basename $datenow`_$jobID
echo "${fragment}" > "${jobstate_file}"

# add output files that we need to download after job
# finishes into the jobstate file
for outfile in "${sfapi_output_files[@]}"
do
    echo "output::${outfile}":${remote_dir}/`basename ${outfile}` >> "${jobstate_file}"
done

# if we are saving the slurm submit file also copy the jobstate
# file that has the paths for remote stdout and stderr
if [ -d "${bls_info_dir}" ] && [ -n "${jobstate_file}" ]; then
    cp ${jobstate_file} ${bls_info_dir}
fi

# echo out the job id for BLAHP to pick up
echo "BLAHP_JOBID_PREFIX$blahp_jobID"

bls_wrap_up_submit

exit $retcode
