Permalink
Branch: master
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
executable file 125 lines (101 sloc) 3.23 KB
#!/bin/bash
# Copyright (C) 2014 Alan Orth
#
# 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
# defaults
readonly DEF_NUM_CPUS=1
readonly DEF_PARTITION=batch
readonly DEF_JOB_NAME=interactive
######
readonly PROGNAME=$(basename $0)
readonly ARGS="$@"
function usage() {
cat <<-EOF
Usage: $PROGNAME [-c] [-p] [-J] [-w]
Optional arguments:
-c: number of CPU cores to request (default: $DEF_NUM_CPUS)
-p: partition to run job in (default: $DEF_PARTITION)
-J: job name (default: $DEF_JOB_NAME)
-w: node name
NB: interactive jobs have a time limit of 8 hours.
Written by: Alan Orth <a.orth@cgiar.org>
EOF
exit 0
}
function parse_options() {
while getopts ":c:p:J:w:" opt; do
case $opt in
c)
# make sure -c is passed a valid integer
if ! [[ "$OPTARG" =~ ^[0-9]+$ ]]; then
usage
fi
NUM_CPUS=$OPTARG
;;
p)
PARTITION=$OPTARG
;;
J)
JOB_NAME=$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
# force environment to be proper login shells
readonly INTERACTIVE_SHELL_OPTS="-l "
# check to see if user requested a specific number of CPUs
# ... otherwise use the default
if [[ -z $NUM_CPUS ]]; then
NUM_CPUS=$DEF_NUM_CPUS
fi
# set the CPU allocation
SALLOC_OPTS="-n $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
SALLOC_OPTS="$SALLOC_OPTS -p ${PARTITION:-$DEF_PARTITION}"
# default time for interactive jobs is 8 hours, ~1 working day
SALLOC_OPTS="$SALLOC_OPTS -t 8:00:00"
# if user specifies a node name, run all the tasks in the specified node
if [[ -n "$NODE_NAME" ]]; then
SALLOC_OPTS="$SALLOC_OPTS -w ${NODE_NAME} --ntasks-per-node=$NUM_CPUS"
fi
# check if we want X11
if [[ -n "$DISPLAY" ]]; then
SRUN_OPTS="--x11 --pty"
else
SRUN_OPTS="--pty"
fi
}
# 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 $INTERACTIVE_SHELL_OPTS