From ba6c9f53eda771c11e719d08c9a9c3bcfb430b7a Mon Sep 17 00:00:00 2001
From: Wirawan Purwanto <wpurwant@turing.hpc.odu.edu>
Date: Wed, 27 Mar 2019 13:28:33 -0400
Subject: [PATCH] * Added accommodation for Anaconda as well as site-provided
 python. * Added support for fully headless mode (connect via local browser).
 * Added some safeguards against job failing to start due to executable   not
 found, etc.

---
 turing/calculations/launch_jupyter | 119 ++++++++++++++++++++++++-----
 1 file changed, 99 insertions(+), 20 deletions(-)

diff --git a/turing/calculations/launch_jupyter b/turing/calculations/launch_jupyter
index 8acdb99..b0a25b9 100755
--- a/turing/calculations/launch_jupyter
+++ b/turing/calculations/launch_jupyter
@@ -6,6 +6,7 @@ clean_up() {
         job_id=$(awk '/srun: job/ {print $3;exit}' $jupyter_output)
         scancel $job_id
 	rm $jupyter_output
+        echo "Terminated SLURM job $job_id"
 }
 
 trap clean_up EXIT
@@ -13,40 +14,118 @@ trap clean_up EXIT
 set -m
 
 
-if [ x"$PYTHON_VER" = x ]; then
-        echo "Please load python and any needed python modules before launching juypter!"
+HAS_TIME=
+
+process_args() {
+        local arg args i
+        args=( "$@" )
+        for arg in "$@"; do
+                case "${args[i]}" in
+                    (-t|--time|--time=*)
+                        HAS_TIME=yes
+                        ;;
+                esac
+        done
+        # FIXME: Can do more thorough arg checking later.
+}
+
+# Allow several choices:
+# - anaconda3
+# - anaconda2
+# - python/2.7
+# - python/3.x
+
+enable_lmod
+if [ -n "$ANACONDA3_VER" ]; then
+        echo "Launching Jupyter for anaconda3 module version $ANACONDA3_VER"
+elif [ -n "$ANACONDA2_VER" ]; then
+        echo "Launching Jupyter for anaconda2 module version $ANACONDA2_VER"
+elif [ -n "$PYTHON_VER" ]; then
+        echo "Launching Jupyter for python module version $PYTHON_VER"
+        echo " * loading jupyter module"
+        if ! module is-loaded numpy; then module load numpy; fi
+        if ! module is-loaded ipython; then module load ipython; fi
+        if ! module is-loaded jupyter; then module load jupyter; fi
+elif [ x"$PYTHON_VER" = x ]; then
+        PYTHON_VER_DEFAULT=3.6
+        echo "Warning: Loading default Python module version $PYTHON_VER_DEFAULT"
+        echo "Warning: Please load python and any needed python modules before launching juypter!"
         echo "Example: "
         echo "  enable_lmod"
         echo "  module load python/3.6   # 2.7/3.7 is also avaiable"
-        exit 1
+        sleep 5s
+        #enable_lmod
+        module load "python/$PYTHON_VER_DEFAULT"
+        echo " * loading jupyter module"
+        module load jupyter
 fi
 
-echo " * loading jupyter module"
-
-enable_lmod
-module load python/$PYTHON_VER
-module load jupyter
+if [ -z "$(which jupyter-notebook)" ]; then
+        echo "Error: cannot find jupyter-notebook program" >&2
+        echo "Please contact site administrator with the following diagnosis:" >&2
+        echo
+        module avail
+        echo
+        echo "PATH=$PATH"
+        exit 1
+fi
 
-echo " * launching jupyter notebook server"
-(srun -n 1 -c 2 -t 240 -J juypter jupyter-notebook --no-browser --port=$((8000 + $RANDOM % 1000)) --ip=0.0.0.0 2>&1 | tee $jupyter_output > /dev/null) &
+# FIXME::
+SLURM_ARGS=( "$@" )
+process_args "${SLURM_ARGS[@]}"
 
 echo 
-echo "   jupyter server will be on for 4 hour"
-echo "   please do not close this window, or jupyter will be terminated immediately"
+if [ x"$HAS_TIME" != xyes ]; then
+        SLURM_ARGS=( -t 240 "${SLURM_ARGS[@]}" )
+        echo "   jupyter server will be TERMINATED after 4 hour"
+fi
+echo "   please do not close this terminal window, or jupyter will be terminated immediately"
 echo 
 
+module list
 
-sleep 10 
-
+echo " * launching jupyter notebook server"
+(srun -n 1 -c 2 -J juypter "${SLURM_ARGS[@]}" jupyter-notebook --no-browser --port=$((8000 + $RANDOM % 1000)) --ip=0.0.0.0 2>&1 | tee $jupyter_output > /dev/null) &
+
+echo " * waiting for Jupyter server to be ready ..."
+for ((i=0; i < 90; ++i)); do
+        tokens=$(awk -F: '/  http:/ {print $3}'        $jupyter_output)
+        if [ -n "$tokens" ]; then
+                break
+        fi
+        sleep 1s
+done
+# Don't use echo to avoid divulging the token (!!)
+port=${tokens%/*}
+token=${tokens#*token=}
+
+# Look for the hostname where the server is running:
 job_id=$(awk     '/srun: job/ {print $3;exit}' $jupyter_output)
-tokens=$(awk -F: '/  http:/ {print $3}'        $jupyter_output)
 host=$(squeue -j $job_id -O NodeList -h | sed 's/ //g')
 url="http://$host:$tokens"
 
-echo " * launching browser"
-(firefox $url 2>&1 )>/dev/null &
-
-echo "   if you closed browser, please use below to connect back"
-echo "      $url"
+echo " * SLURM job ID: $job_id"
+
+if [ -n "$DISPLAY" ]; then
+        echo " * launching browser"
+        (firefox $url 2>&1 )>/dev/null &
+        echo "   if you close your browser, please the link use below to connect back"
+        echo "      $url"
+else
+        echo
+        echo "NOTE: Jupyter is running in a headless mode!"
+        echo "To connect, please keep an open connection from your local computer using:"
+        echo
+        echo "    ssh -L 8888:$host:$port $USER@turing.hpc.odu.edu"
+        echo
+        echo "and point your browser to:"
+        echo
+        # https://stackoverflow.com/a/15229498/655885
+        cat <<EOF
+    http://localhost:8888/?token=${token}
+EOF
+        echo
+        echo "(If ssh fails because local port 8888 is taken, use another number.)"
+fi
 
 wait