interactive

#!/bin/bash

# Copyright (C) 2013 Alan Orth
# Modified by Jordi Blasco <jordi.blasco@hpcnow.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

#
# interactive: Obtain interactive shell with slurm allocation
#
# Basically simulates:
#
#    salloc -n1 srun --x11 --pty bash

# source /sNow/snow-utils/etc/snow.env

# The srun options have been defined in the slurm.conf with :
# SallocDefaultCommand="srun -n1 -N1 --mem-per-cpu=0 --pty --preserve-env --mpi=none $SHELL"
DEF_PARTITION=test
DEF_ARCH=hsw
DEF_MEM_PER_CPU=1
DEF_NUM_CPUS=1
DEF_NUM_TASKS=1
DEF_JOB_NAME=interactive
DEF_TIME=01:00:00
SRUN_BIN=/usr/bin/srun
#SRUN_OPTS="--pty"
SRUN_OPTS="-n1 -N1 --mem-per-cpu=0 --pty --preserve-env --mpi=none"
INTERACTIVE_SHELL=/bin/bash
INTERACTIVE_SHELL_OPTS=" "

######
 
function help() {
    echo -e "Usage: `basename $0` [-A] [-a] [-c] [-m] [-J] [-e]\n"
    echo "Optional arguments:"
    echo -e "\t -A: account (non-default account)"
    echo -e "\t -p: partition (default: $DEF_PARTITION)"
    echo -e "\t -a: architecture (default: $DEF_ARCH, values hsw=Haswell skl=SkyLake wsw=Warsaw)"
    echo -e "\t -n: number of tasks (default: $DEF_NUM_TASKS)"
    echo -e "\t -c: number of CPU cores (default: $DEF_NUM_CPUS)"
    echo -e "\t -m: amount of memory (GB) per core (default: $DEF_MEM_PER_CPU [GB])"
    echo -e "\t -e: email address to which the begin session notification is to be sent"
    echo -e "\t -r: specify a reservation name"
    echo -e "\t -w: target node"
    echo -e "\t -J: job name"
    echo -e "\t -x: binary that you want to run interactively"
    echo -e "example : `basename $0` -A snow -a hsw -c 4 -J MyFirstInteractiveJob"
    echo -e "example : `basename $0` -A snow -a hsw -c 4 -J MyFirstInteractiveJob -x \"MyBinary MyOptions\""

    echo -e "\nWritten by: Alan Orth <a.orth@cgiar.org>"
    echo -e "Modified by: Jordi Blasco <jordi.blasco@hpcnow.com>"

    exit 0
}

function parse_options() {
    while getopts ":A:p:a:n:c:m:e:r:w:J:x:" opt; do
        case $opt in
            a)
                # make sure -a is passed a valid architecture
                case $OPTARG in 
                hsw|skl|wsw|wm|sb|ib|kepler|fermi)
                    DEF_ARCH=$OPTARG
                    ;;
                *)
                    help   
                    ;;
               esac
               ;; 
            p)
                # make sure -a is passed a valid architecture
                case $OPTARG in 
                serial|debug|8_cores|12_cores|16_cores|20_cores|GPU)
                   DEF_PARTITION=$OPTARG
                  ;;
               *)
                  help   
                  ;;
               esac
               ;; 
            n)
                # make sure -n is passed a valid integer
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    help
                fi
                NUM_TASKS=$OPTARG
                ;;
            c)
                # make sure -c is passed a valid integer
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    help
                fi
                NUM_CPUS=$OPTARG
                ;;
            m)
                # make sure -m is passed a valid integer
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    help
                fi

                MEM_PER_CPU=$OPTARG
                ;;
            p)
                PARTITION=$OPTARG
                ;;
            e) 
                EMAIL_NOTIFY=$OPTARG
                ;;
            J)
                JOB_NAME=$OPTARG
                ;;
            A)
                ACCOUNT=$OPTARG
                ;;
            x)
                RUNEXEC=$OPTARG
                ;;
            r)
                RES=$OPTARG
                ;;
            w)
                NODELIST=$OPTARG
                ;;

            \?|:)
                help
                ;;
        esac
    done
}

function envsetup() {
    SALLOC_BIN=/usr/bin/salloc

    # default is to request 1 CPU, unless more are requested
    if [[ $NUM_TASKS -gt $DEF_NUM_TASKS ]]; then
        SALLOC_OPTS="-n $NUM_TASKS"
    else
        SALLOC_OPTS="-n $DEF_NUM_TASKS"
    fi

    # default is to request 1 CPU, unless more are requested
    if [[ $NUM_CPUS -gt $DEF_NUM_CPUS ]]; then
        SALLOC_OPTS="$SALLOC_OPTS -c $NUM_CPUS"
    else
        SALLOC_OPTS="$SALLOC_OPTS -c $DEF_NUM_CPUS"
    fi

    # default is to request 1GB of Memory / CPU, unless more are requested
    if [[ $MEM_PER_CPU -gt $DEF_MEM_PER_CPU ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --mem-per-cpu=$(($MEM_PER_CPU*1024))"
    else
        SALLOC_OPTS="$SALLOC_OPTS --mem-per-cpu=$(($DEF_MEM_PER_CPU*1024))"
    fi

    # check for a job name
    if [[ -n "$JOB_NAME" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS -J $JOB_NAME"
    else
        SALLOC_OPTS="$SALLOC_OPTS -J $DEF_JOB_NAME"
    fi

    if [[ -n "$EMAIL_NOTIFY" ]]; then
         SALLOC_OPTS="$SALLOC_OPTS --mail-type=BEGIN --mail-user=$EMAIL_NOTIFY"
    fi

    if [[ -n "$ACCOUNT" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS -A $ACCOUNT"
    else
        echo "Account is mandatory, please add : -A account_code"
        exit 0
    fi 

    # add the partition 
    SALLOC_OPTS="$SALLOC_OPTS -p $DEF_PARTITION"

    # add the architecture as a feature
    SALLOC_OPTS="$SALLOC_OPTS -C $DEF_ARCH"

    # add the default time limit
    SALLOC_OPTS="$SALLOC_OPTS -t $DEF_TIME"

    # check for executable
    if [[ -n "$RUNEXEC" ]]; then
        SRUN_OPTS=$RUNEXEC
        unset INTERACTIVE_SHELL
        unset INTERACTIVE_SHELL_OPTS 
    fi

    # Check for a reservation name
    if [[ -n "$RES" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --reservation=$RES"
    fi

    # Check for a node name
    if [[ -n "$NODELIST" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --nodelist=$NODELIST"
    fi

}

# pass the shell's argument array to the parsing function
parse_options "$@"

# Don't propogate microarchitecture of login node.
if [[ -z "$MODULE_LOAD_ALLOWED" ]]; then
    unset MODULEPATH MODULEPATH_ROOT module ml LMOD_CMD LMOD_SPIDER_CACHE_DIRS BASH_ENV
fi

# setup the defaults
envsetup

$SALLOC_BIN $SALLOC_OPTS $SRUN_BIN $SRUN_OPTS $INTERACTIVE_SHELL $INTERACTIVE_SHELL_OPTS