#!/bin/sh -e

##########################################################################
#   Script description:
#       Run a shell under SLURM
#       
#   History:
#   Date        Name        Modification
#   2020-02-03  Jason Bacon Begin
##########################################################################

usage()
{
    cat << EOM

Usage: $0 cores memory [additional srun flags]

slurm-shell requires the number of cores and total amount of memory
(per node, not per core).  Memory specifications are in mebibytes unless
suffixed (e.g. with 'g' for gibibytes).

Be sure that the programs you run under slurm-shell do not use more cores
or memory than you specified!

Examples:

    slurm-shell 1 200               # 1 core, 200 mebibytes
    slurm-shell 5 3g                # 5 cores, 3 gibibytes

EOM
    exit 1
}


##########################################################################
#   Main
##########################################################################

if [ $# -lt 2 ]; then
    usage
fi

cluster-head-check $0
set +e

cat << EOM

********************************** NOTE ***********************************

You can use additional "srun" flags with slurm-shell to alter the
allocation in any other way supported by srun.

Run "man srun" for a complete list of options.

********************************* WARNING *********************************

Interactive jobs run in this manner may be terminated at any moment if
security updates on the login node require a reboot.

The ability to run interactive jobs is provided as a convenience with no
guarantees of persistence.

HPC clusters are intended for running batch jobs and are not well suited
for running interactive applications in general.

Users who require fast graphical sessions or reliable interactive sessions
are advised to use a desktop computer rather than an HPC cluster.

***************************************************************************

EOM

cores=$1
mem=$2
shift
shift

srun --mem=$mem --nodes=1 --cpus-per-task=$cores --pty --preserve-env $@ $SHELL -l
