getnode for SLURM

#!/bin/bash

VERSION="1.3"
AUTHOR="Garrett Mills"

NODECOUNT=1
NTASKS=1
PARTITION=sixhour
SHELL=bash
JOBNAME="getnode"
TIME="03:00:00"
X11=1

ADDONS=""

function usage()
{
	echo -e "Usage: getnode [params]..."
	echo -e "Drop to an interactive shell on a node in the SLURM cluster."
	echo -e ""
	echo -e "\t -h,--help \t\t\t display this usage information"
	echo -e "\t -N=n,--nodes=n \t\t request n nodes for the shell job (default=$NODECOUNT)"
	echo -e "\t -n=n,--ntasks-per-node=n \t allow a maximum of n processes per node (default=$NTASKS)"
	echo -e "\t -p=PART,--partition=PART \t request nodes from the PART partition (default=$PARTITION)"
	echo -e "\t -s=PROG,--shell=PROG \t\t drop to the PROG shell (default=$SHELL)"
	echo -e "\t -t=TIME,--time=TIME \t\t set the maximum time for the process to TIME (default=$TIME)"
	echo -e "\t --jobname=NAME \t\t name of the node-request job (default=$JOBNAME)"
	echo -e "\t --constraints=list \t\t SBATCH constraints for the shell job"
	echo -e "\t --nodelist=list \t\t specific list of nodes to request"
	echo -e "\t --exclusive \t\t\t force the shell job to be the only running job on the node"
	echo -e "\t -xX,--no-x11 \t\t\t disable X.org forwarding (enabled by default)"
	echo -e "\t -l,--license \t\t\t display license information"
	echo -e ""
	echo -e "getnode (slurm-getnode srun wrapper) version $VERSION by $AUTHOR"
}

function license(){
	echo -e "getnode (slurm-getnode srun wrapper) version $VERSION by $AUTHOR"
	echo -e "Copyright 2018 (C) Garrett Mills."
	echo -e ""
	echo -e "Permission is hereby granted, free of charge, to any person obtaining"
	echo -e "a copy of this software and associated documentation files (the \"Software\")"
	echo -e "to deal in the Software without restriction, including without limitation"
	echo -e "to the rights to use, copy, modify, merge, publish, distribute, sublicense"
	echo -e "and/or sell copies of the Software, and to permit persons to whom the"
	echo -e "Software is furnished to do so, subject to the following conditions:"
	echo -e ""
	echo -e "The above copyright notice and this permission notice shall be included in"
	echo -e "all copies or substantial portions of the Software."
	echo -e ""
	echo -e "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND,"
	echo -e "EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF"
	echo -e "MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT."
	echo -e "IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM"
	echo -e "DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR"
	echo -e "OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE"
	echo -e "USE OR OTHER DEALINGS IN THE SOFTWARE."
}

while [ "$1" != "" ]; do
    PARAM=`echo $1 | awk -F= '{print $1}'`
    VALUE=`echo $1 | awk -F= '{print $2}'`
    case $PARAM in
        -h | --help)
            usage
            exit
            ;;
        -N | --nodes)
            NODECOUNT=$VALUE
            ;;
        -n | --ntasks-per-node)
            NTASKS=$VALUE
            ;;
        -p | --partition)
            PARTITION=$VALUE
            ;;
        -s | --shell)
            SHELL=$VALUE
            ;;
        --jobname)
            JOBNAME=$VALUE
            ;;
	-t | --time)
	    TIME=$VALUE
	    ;;
        --constraints)
            CONSTRAINTS=$VALUE
	    ADDONS="$ADDONS --constraint=$VALUE"
            ;;
        --nodelist)
            NODELIST=$VALUE
	    ADDONS="$ADDONS --nodelist=$VALUE"
            ;;
        --exclusive)
            EXCLUSIVE=1
	    ADDONS="$ADDONS --exclusive"
            ;;
        -xX | --no-x11)
	    X11=0
            ;;
        -l | --license)
	    license
	    exit
            ;;
        *)
            echo "ERROR: unknown parameter \"$PARAM\""
            usage
            exit 1
            ;;
    esac
    shift
done

command -v srun >/dev/null 2>&1 || {
	echo -e "Not a SLURM environment, or missing srun. Aborting."
	echo -e ""	
	echo -e "getnode (slurm-getnode srun wrapper) version $VERSION by $AUTHOR"
	exit 1
}

if [ $X11 -eq 1 ]; then
	ADDONS="$ADDONS --x11"
fi

srun --nodes=$NODECOUNT --ntasks-per-node=$NTASKS --time=$TIME --partition=$PARTITION --job-name=$JOBNAME $ADDONS --pty $SHELL