#!/bin/bash
# This wrapper script is intended to be submitted to Slurm to support
# communicating jobs.
#
# This script uses the following environment variables set by the submit MATLAB code:
# MDCE_MATLAB_EXE     - the MATLAB executable to use
# MDCE_MATLAB_ARGS    - the MATLAB args to use
# MDCE_TOTAL_TASKS    - the number of workers to start
#
# The following environment variables are forwarded through mpiexec:
# MDCE_DECODE_FUNCTION     - the decode function to use
# MDCE_STORAGE_LOCATION    - used by decode function
# MDCE_STORAGE_CONSTRUCTOR - used by decode function
# MDCE_JOB_LOCATION        - used by decode function
#
# For backward compatability, this wrapper script uses the following Slurm environment variables:
# SLURM_NODELIST           - instead of SLURM_JOB_NODELIST

# Copyright 2015-2017 The MathWorks, Inc.

# Intel MPI 4.1 implements MPI-2.2
# R2013b ships MPICH2 version 1.4.1p1
# Compatable through at least R2014b (and 15a?)

# Prep module load
source /etc/profile.d/modules.sh
module load profile/eng
module load autoload
# matlab module generally loads intel MPI, but due to the ssh connection do it explicitly
# module load intel/pe-xe-2018--binary
# module load intelmpi/2018--binary
module load matlab

# This gives MATLAB the best performance
export KMP_AFFINITY='granularity=thread,none'

# Test to make sure IB is being used
export I_MPI_FABRICS=shm:tmi

# Environment variables needed to use omnipath with MDCS
# Not sure which PSM version was used, so using both env. variables.
export IPATH_NO_BACKTRACE=1
export HFI_NO_BACKTRACE=1

# MPIEXEC=srun
MPIEXEC='srun --mpi=pmi2'
MPIEXEC_CODE=0
# export I_MPI_ROOT=
# For srun, we need to ensure that we're using the PMI library
# I_MPI_PMI_LIBRARY must be defined
export I_MPI_PMI_LIBRARY=/usr/lib64/libpmi.so

# Set stack size to unlimited
ulimit -s unlimited

export TZ=Europe/Zurich

# export I_MPI_FABRICS_LIST=shm:tcp
# export I_MPI_FABRICS=shm:tcp

# Work out how many processes to launch - set MACHINE_ARG
chooseMachineArg() {
    echo Node list: ${SLURM_NODELIST:?"Node list undefined"}
    MACHINE_ARG="-l -n ${MDCE_TOTAL_TASKS}"
}

runMpiexec() {

    ###################################################
    # DEBUGGING
    ###################################################
    # cmd="strace -o $STRACE_DIR/mdcs_test -tt -ff $cmd"
    # STRACE_LAUNCHER_PFX="strace -o $STRACE_DIR/mpiexec_a -tt -f"
    ## THIS PROBABLY ONLY WORKS WELL FOR 1 WORKER
    # STRACE_LAUNCHEE_PFX="strace -o $STRACE_DIR/worker_a -tt -f"
    # CMD="$STRACE_LAUNCHER_PFX $MPIEXEC -l $MACHINE_ARG $STRACE_LAUNCHEE_PFX $MDCE_MATLAB_EXE $MDCE_MATLAB_ARGS"
    ###################################################
    
    if [ "X${MDCE_DEBUG}X" == "XtrueX" ]; then
        # SRUN verbosity can go as high as -vvvvv
        MPI_VERBOSE=-vv

        # For each debug message, date and time, log level.  Don't have
        #   MPI display the rank level since we're already doing it
        # Could also add source file name (file), source line number (line)
        #   scope (scope), pid (pid), & host (host)
        export I_MPI_DEBUG=5,datetime,level,norank
    else
        MPI_VERBOSE=
    fi

    CMD="$MPIEXEC $MPI_VERBOSE $MACHINE_ARG $MDCE_MATLAB_EXE $MDCE_MATLAB_ARGS"

    # As a debug stage: echo the command line...
    echo $CMD

    # ...and then execute it
    eval $CMD

    MPIEXEC_CODE=${?}
}

###################################################
# DEBUGGING
###################################################
# STRACE_DIR=$HOME/matlab/stracelogs
# test -d $STRACE_DIR || mkdir -p $STRACE_DIR
###################################################

chooseMachineArg
runMpiexec

echo "Exiting with code: ${MPIEXEC_CODE}"
exit ${MPIEXEC_CODE}
