jlab-hobart

#!/bin/bash
set -e

#-- set defaults
# expected environment variables
walltime=08:00:00
queue=medium
notebook_dir=${HOME}
jlab_port=8888
#dash_port=8787

NODES=1
PPN=48

#-- define usage
usage () {
cat << EOF
Usage: launch dask
Possible options are:
 -w,--walltime: walltime [default: ${walltime}]
 -q,--queue: queue [default: ${queue}]
 -d,--directory: notebook directory
 -p,--port: [default: ${jlab_port}]
EOF
exit
}

#-- parse input args
while [[ $# -gt 0 ]]; do
  key="${1}"
  case ${key} in
    -h|--help)
      usage ;;
    -w|--walltime)
      walltime="${2}"
      shift 2
      ;;
    -q|--queue)
      queue="${2}"
      shift 2
      ;;
    -d|--directory)
      notebook_dir="${2}"
      shift 2
      ;;
    -p|--port)
      jlab_port="${2}"
      shift 2
      ;;
    *)
      echo "ERROR: unknown argument: ${key}"
      usage
      ;;
  esac
done

#-- check inputs
if [ -z "${walltime}" ]; then
  echo "ERROR: walltime not set."
  exit 1
fi
if [ -z "${queue}" ]; then
  echo "ERROR: queue not set."
  exit 1
fi
if [[ -z ${WORKDIR} ]]; then
    WORKDIR=/scratch/cluster/${USER}/tmp
fi

#-- make sure the working directory exists
if [[ ! -d ${WORKDIR} ]]; then
  mkdir -vp ${WORKDIR}
fi
LOGFILE=${WORKDIR}/jlab-hobart.jlab-log.$(date +%Y%m%dT%H%M%S)



echo "Launching notebook server"

echo "  queue = ${queue}"
echo "  nodes = ${NODES}"
echo "  walltime = ${walltime}"
echo "  port = ${jlab_port}"
echo

s=$(qsub << EOF
#!/bin/bash
#PBS -N jlab-hobart
#PBS -r n
#PBS -q ${queue}
#PBS -l nodes=${NODES}:ppn=${PPN}
#PBS -l walltime=${walltime}
#PBS -o ${WORKDIR}/
#PBS -e ${WORKDIR}/
#PBS -j oe

# setup environment

# run lab server
cd ${notebook_dir}
jupyter lab --no-browser --ip=\$(hostname) --notebook-dir ${notebook_dir} >& ${LOGFILE}
EOF
)
sjob=${s%%.*}
echo "submitted job: ${sjob} to queue ${queue}"

trap ctrl_c INT
function ctrl_c() {
  echo "killing ${sjob}"
  qdel ${sjob}
  exit 0
}


#-- wait for job to start
echo -n "waiting for job to run"
while [ 1 ]; do
  #-- get job status
  job_state=$(qstat -f ${sjob} | grep job_state | awk -F= '{print $2}')
  exec_host=$(qstat -f ${sjob} | grep exec_host | awk -F= '{print $2}')

  if [[ -n "${exec_host}" && ${job_state} =~ R ]]; then
    exec_host=$(echo -e "${exec_host%/*}" | tr -d '[:space:]')
    exec_host=$(echo -e "${exec_host%%.*}")
    echo .
    echo "Job ID: ${sjob}"
    echo "Job host: ${exec_host}"
    break
  elif [[ ${job_state} =~ F ]]; then
    echo "Job failed."
    qstat -x -f ${sjob}
    exit 1
  else
    echo -n ..
  fi
  sleep 1
done

#-- wait until log file is there
echo "waiting for job log to appear: ${LOGFILE}"
while [ ! -f "${LOGFILE}" ]; do
  sleep 1
done

#-- wait until address shows up in log
ADDRESS=
ELAPSED=0
while [ -z "${ADDRESS}" ]; do
  ADDRESS=$(grep -e '^\[.*\]\s*http://.*:' ${LOGFILE} | head -n 1)
  PORT=$(echo ${ADDRESS#*http://} | awk -F':' '{print $2}' | awk -F'/' '{print $1}')
  sleep 1
  ((ELAPSED+=1))
  if [[ ${ELAPSED} -gt 120 ]]; then
      echo -e "something went wrong\n---"
      cat ${LOGFILE}
      echo "---"
      ctrl_c
  fi
done

echo
echo "----------------------------------------------------------------------"
echo "Execute on local machine:"
echo "ssh -N -L ${jlab_port}:${exec_host}:${PORT} ${USER}@$(hostname)" # -L ${dash_port}:${exec_host}:8787
echo
echo "Open a browser on your local machine and type in the address bar:"
echo "http://localhost:${jlab_port}"
echo "----------------------------------------------------------------------"
echo

echo "Job log: ${LOGFILE}"
tail -f ${LOGFILE}