svisit

#!/bin/bash

usage() {
  svisit=$(basename $0)
  cat <<EOF
Usage:
  $svisit [-t] [[-j] job] [[-n] node] [[srun options] command]
    -t        Test instead - just show the srun command $svisit would use.
    -j job    A Slurm Job ID. If not provided one will be looked for via 'squeue --me'.
    -n node   A particular compute node. Defaults to the first node in the job.
    command   Defaults to your shell (${SHELL:-/bin/bash})
 
Description:
  $svisit starts a terminal session within your currently running Slurm job.
  It is a wrapper around 'srun --pty --overlap --jobid=...', and a substitute for
  most uses of ssh to reach compute nodes.

Example:
  svisit 123456_1 c001 htop
EOF
}

while [[ -n "$1" ]]; do
  if [[ $1 == '-h' ]]; then
    usage
    exit 0
  elif [[ $1 == '-t' ]]; then
    TEST="echo "
  elif [[ $1 == '-j' ]]; then
    shift
    JOB_ID=$1
  elif [[ -z $J && $1 =~ ^[0-9_.]+$ ]]; then
    JOB_ID=$1
  elif [[ $1 == '-n' ]]; then
    shift
    NODE_FILTER="-w $1"
  elif [[ -z $NODE_FILTER && $1 =~ ^[a-z]{1,2}[0-9]{2,3}$ ]]; then
    NODE_FILTER="-w $1"
  else
    break
  fi
  shift
done

if [[ -z $JOB_ID ]]; then
  JOB_ID=$(squeue --Format=JobID --me --noheader --states=R $NODE_FILTER | sort -nr | head -1)
  if [[ -z $JOB_ID ]]; then
    echo "No JobID provided and no running job found either. 'svisit -h' for help." 1>&2
    exit 1
  fi
fi

if [[ -n $1 ]]; then
  $TEST srun --pty --overlap --jobid=$JOB_ID $NODE_FILTER $@
else
  $TEST srun --pty --overlap --jobid=$JOB_ID $NODE_FILTER ${SHELL:-/bin/bash}
fi