sinteractive

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

# defaults
readonly DEF_NUM_NODES=1
readonly DEF_NUM_TASKS=1
readonly DEF_NUM_CPUS=1
readonly DEF_PARTITION=debug
if [[ "$HOSTNAME" = "a0"* ]] || [[ "$HOSTNAME" = "ascend"* ]]; then
	readonly DEF_GPU_PARTITION=debug
else
	readonly DEF_GPU_PARTITION=gpudebug
fi
readonly DEF_JOB_NAME=interactive

######

readonly PROGNAME=$(basename $0)
readonly ARGS="$@"

function usage() {
    cat <<-EOF
Usage: $PROGNAME [-p] [-N] [-n] [-c] [-m] [-M] [-g] [-G] [-L] [-t] [-J] [-A] [-w]

Optional arguments:
    -p: partition of where to run job (default: $DEF_PARTITION)
    -N: number of nodes to request (default: $DEF_NUM_NODES)
    -n: number of tasks to request (default: $DEF_NUM_TASKS)
    -c: number of CPU cores to request (default: $DEF_NUM_CPUS)
    -m: memory per CPU (default: Partition default)
    -M: memory per node (default: Partition default)
    -g: number of GPUs to request (default: None)
    -G: GRES to request (default: None)
    -L: Licenses to request (default: None)
    -t: Time limit (default: Partition default)
    -J: job name (default: $DEF_JOB_NAME)
    -A: Account
    -w: node name

Additional help:
  https://www.osc.edu/supercomputing/batch-processing-at-osc/job-submission
EOF

    exit 0
}

function parse_options() {
    while getopts ":p:N:n:c:m:M:g:G:L:t:J:A:w:" opt; do
        case $opt in
            p)
                PARTITION=$OPTARG
                ;;
            N)
                # make sure -N is passed a valid integer
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    usage
                fi

                NUM_NODES=$OPTARG
                ;;
            n)
                # make sure -n is passed a valid integer
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    usage
                fi

                NUM_TASKS=$OPTARG
                ;;
            c)
                # make sure -c is passed a valid integer
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    usage
                fi

                NUM_CPUS=$OPTARG
                ;;
            m)
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    usage
                fi

                MEM_PER_CPU=$OPTARG
                ;;
            M)
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    usage
                fi

                MEM_PER_NODE=$OPTARG
                ;;
            g)
                if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
                    usage
                fi

                GPUS=$OPTARG
                ;;
            G)
                GRES=$OPTARG
                ;;
            L)
                LICENSE=$OPTARG
                ;;
            t)
                TIME_LIMIT=$OPTARG
                ;;
            J)
                JOB_NAME=$OPTARG
                ;;
            A)
                ACCOUNT=$OPTARG
                ;;
            w)
                NODE_NAME=$OPTARG
                ;;
            \?|:)
                usage
                ;;
        esac
    done
}

function envsetup() {
    readonly SALLOC_BIN=/usr/bin/salloc
    readonly SRUN_BIN=/usr/bin/srun
    readonly INTERACTIVE_SHELL=/bin/bash

    # default is to request 1 node, unless more are requested
    SALLOC_OPTS="-N ${NUM_NODES:-$DEF_NUM_NODES}"

    # default is to request 1 task, unless more are requested
    SALLOC_OPTS="$SALLOC_OPTS -n ${NUM_TASKS:-$DEF_NUM_TASKS}"

    # default is to request 1 CPU per task, unless more are requested
    SALLOC_OPTS="$SALLOC_OPTS -c ${NUM_CPUS:-$DEF_NUM_CPUS}"

    # check for a job name, otherwise use default
    SALLOC_OPTS="$SALLOC_OPTS -J ${JOB_NAME:-$DEF_JOB_NAME}"

    # see if the user specified a partition, otherwise use default
    if [[ -n "$GPUS" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS -p ${PARTITION:-$DEF_GPU_PARTITION}"
        SALLOC_OPTS="$SALLOC_OPTS --gpus ${GPUS}"
    else
        SALLOC_OPTS="$SALLOC_OPTS -p ${PARTITION:-$DEF_PARTITION}"
    fi

    # set mem-per-cpu if defined as option
    if [[ -n "$MEM_PER_CPU" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --mem-per-cpu ${MEM_PER_CPU}"
    fi

    # set mem if defined as option
    if [[ -n "$MEM_PER_NODE" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --mem ${MEM_PER_NODE}"
    fi

    # set gres if defined as option
    if [[ -n "$GRES" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --gres ${GRES}"
    fi

    # set license if defined as option
    if [[ -n "$LICENSE" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --licenses ${LICENSE}"
    fi

    # set time only if defined as option
    if [[ -n "$TIME_LIMIT" ]];then
        SALLOC_OPTS="$SALLOC_OPTS -t ${TIME_LIMIT}"
    fi

    # set account only if defined as option
    if [[ -n "$ACCOUNT" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS -A ${ACCOUNT}"
    fi

    # see if the user specified a nodename
    if [[ -n "$NODE_NAME" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS -w ${NODE_NAME} --ntasks-per-node=1"
    fi

    # check if we want X11
    if [[ -n "$DISPLAY" ]]; then
        SALLOC_OPTS="$SALLOC_OPTS --x11"
    fi
    SRUN_OPTS="--interactive --pty"
}

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

# setup the defaults
envsetup

$SALLOC_BIN $SALLOC_OPTS $SRUN_BIN $SRUN_OPTS $INTERACTIVE_SHELL