smux

#!/usr/bin/env bash
# shellcheck disable=SC2128

###############################################################################
# Launch a tmux server on a slurm node and interact with it from a login node #
#  (particularly useful when a slurm cluster has multiple login nodes). Has   #
# the same command line interface as tmux with an additional 'batch' command  #
#   which starts the tmux server using sbatch. Options are forwarded to the   #
#                         underlying sbatch command.                          #
###############################################################################

set -o errexit -o nounset -o pipefail

SBATCH_JOB_NAME="$(basename "$BASH_SOURCE")"
STATE_DIR="${XDG_STATE_HOME:-$HOME/.local/state}/$SBATCH_JOB_NAME"
JOB_ID_FILE="$STATE_DIR/jobid"

mkdir -p "$STATE_DIR"
touch "$JOB_ID_FILE"
export SBATCH_JOB_NAME

function job_id(){
  cat "$JOB_ID_FILE"
}

function job_state(){
  [[ -n $(squeue -ht "$1" -j "$(job_id)" 2>/dev/null) ]]
}

function cancel(){
  if job_state running,pending; then
    scancel "$(job_id)"
  fi
  truncate -s 0 "$JOB_ID_FILE"
}

function multiplex(){
  while job_state pending; do
    sleep 1
  done
  if job_state running; then
    srun --overlap --pty --jobid "$(job_id)" tmux "$@"
  else
    echo "No server running on any slurm node!" \
      "Run '$SBATCH_JOB_NAME batch [options]' first." >&2
    exit 1
  fi
}

if [[ "$BASH_SOURCE" == "$0" ]]; then
  case "${1:-}" in
    batch)
      sbatch --parsable "${@:2}" --wrap 'tmux -D' >"$JOB_ID_FILE"
      ;;
    switch)
      echo "${2:-$(job_id)}" >"$JOB_ID_FILE"
      ;;
    cancel)
      cancel
      ;;
    *)
      multiplex "$@"
      ;;
  esac
fi