gadi_jupyter

#!/bin/bash

# Scott Wales 20190522

print_help() {
cat <<EOF

Run a Jupyter notebook on Gadi's compute nodes, presenting the interface in a
browser on the local machine

You can set a default username in your local ~/.ssh/config, e.g.

    Host gadi.nci.org.au
    User abc123

Your default project is set in the file on Gadi ~/.config/gadi-login.conf

General Options:
    -h:         Print help
    -l USER:    NCI username
    -L HOST:    NCI login node (default 'gadi.nci.org.au')
    -e ENVIRON: Conda environment
    -d:         Debug mode
    -D:         Profile mode

Queue Options:
    -q QUEUE:   Queue name
    -n NCPU:    Use NCPU cpus
    -g NGPU:    Use NGPU gpus (only on some queues)
    -m MEM:     Memory allocation (default 4*NCPU GB)
    -t TIME:    Walltime limit (default 1 hour)
    -J JOBFS:   Jobfs allocation (default 100 GB)
    -P PROJ:    Submit job under project PROJ
    -s STORAGE: PBS Storage flags (e.g. 'scratch/w35+gdata/v45')
                If not specified, all available storage is requested

EOF
}

set -eu

# Internal defaults
USER=''
PROJECT=''
LOGINNODE='gadi.nci.org.au'
QUEUE='normalbw'
NCPUS='1'
NGPUS=''
WALLTIME=1:00:00
JOBFS=100gb
MEM=''
CONDA_ENV=analysis3
DEBUG=""
STORAGE=""
STRACE=""

# Handle arguments
optspec="hl:L:q:n:m:t:J:P:e:ds:g:D"
while getopts "$optspec" optchar; do
    case "${optchar}" in
        h)
            print_help
            exit 2
            ;;
        l)
            USER="${OPTARG}"
            ;;
        L)
            LOGINNODE="${OPTARG}"
            ;;
        q)
            QUEUE="${OPTARG}"
            ;;
        n)
            NCPUS="${OPTARG}"
            ;;
        g)
            NGPUS="${OPTARG}"
            ;;
        t)
            WALLTIME="${OPTARG}"
            ;;
        J)
            JOBFS="${OPTARG}"
            ;;
        m)
            MEM="${OPTARG}"
            ;;
        P)
            PROJECT="${OPTARG}"
            ;;
        e)
            CONDA_ENV="${OPTARG}"
            ;;
        s)
            STORAGE="${OPTARG}"
            ;;
        d)
            DEBUG=true
            ;;
        D)
            DEBUG=true
            STRACE="strace --follow-forks -e trace=%file --syscall-times --absolute-timestamps=precision:us --output=\$WORKDIR/strace"
            ;;
        *)
            print_help
            exit 2
            ;;
    esac
done

# This gets evaluated on Gadi in the SSH script
WORKDIR=/scratch/${PROJECT:-\$PROJECT}/\$USER/tmp/runjp

# Check for agent
set +e
ssh-add -l &> /dev/null
AGENT=$?
set -e
if [ $AGENT -eq 2 ]; then
    # Restart with an agent
    ssh-agent "$0" "$@"
    exit
fi
if [ $AGENT -eq 1 ]; then
    # No keys saved
    ssh-add
fi

SSH='ssh -x -oBatchMode=yes'
if [ -n "$USER" ]; then
    SSH="${SSH} -l ${USER}"
fi

if [ $NCPUS -gt 48 ]; then
    echo "WARNING: Using more than one node with Dask needs extra setup and is not supported by this script"
fi
if [ $NCPUS -gt 8 ]; then
    echo "WARNING: Using a large number of CPUs in an interactive session can waste lots of NCI compute time. Try keeping the number of CPUs under 8."
    read -p "Proceed? [y/N] " -n 1 cpu_confirm
    echo
    if [[ ! "$cpu_confirm" =~ ^[Yy]$ ]]; then
        exit
    fi
fi
if [ -z "$MEM" ]; then
    MEM_PER_CPU=4


    MEM="$(( NCPUS * MEM_PER_CPU ))gb"
fi

SUBMITOPTS="-N jupyter-lab ${PROJECT:+-P '$PROJECT'} -q '$QUEUE' -l 'ncpus=${NCPUS},mem=${MEM},walltime=${WALLTIME},jobfs=${JOBFS}'"

if [ -n "$NGPUS" ]; then
    SUBMITOPTS="$SUBMITOPTS -l 'ngpus=$NGPUS'"
fi

echo "Starting notebook on ${LOGINNODE}..."
if [ -n "$DEBUG" ]; then
    echo "DEBUG: Connecting using: '$SSH $LOGINNODE'"
fi

# Check connection
$SSH "$LOGINNODE" echo "Working directory: $WORKDIR"

echo "qsub ${SUBMITOPTS}"

# Kill the job if this top-level script is cancelled while the job is still in the queue
trap "{ echo 'Stopping queued job... (Ctrl-C will leave job in the queue)' ; $SSH \"$LOGINNODE\" > /dev/null <<< 'qdel \$(cat $WORKDIR/jobid)' ; }" EXIT

message=$(
$SSH -q "$LOGINNODE" <<EOF | tail -n 1

set -eu

WORKDIR="$WORKDIR"
mkdir -p "\$WORKDIR"

# Check if a hh5 member
if [ ! -d /g/data/hh5 ]; then
    echo "ERROR: \$USER is not a member of hh5. Join at https://my.nci.org.au/mancini/project/hh5/join" >&2
    echo "x x x x ERROR"
    exit
fi

# Check if already running
if [ -f "\$WORKDIR/jobid" ] && [ "\$(qstat -x \$(cat "\$WORKDIR/jobid") 2> /dev/null | sed -n '\$s/\S\+\s\+jupyter-lab\s.*\s[QR]\s\+\S\+\s*$/\0/p')" ] ; then
    while [ ! -f "\$WORKDIR/message" ]; do
        sleep 5
    done
    cat "\$WORKDIR/message" | sed 's/$/ RECONNECT/'
    exit
fi

rm -f "\$WORKDIR/message"

cat > "\$WORKDIR/runjp.sh" <<EOQ
#!/bin/bash

module purge

eval "\\\$(/g/data/hh5/public/apps/miniconda3/bin/conda shell.bash hook)"
conda activate "${CONDA_ENV}"

set -eu

# Jupyter security token
TOKEN=\\\$(uuidgen)

# Find a remote port https://unix.stackexchange.com/a/132524
PORT=\\\$(python -c 'import socket; s=socket.socket(); s.bind(("", 0)); print(s.getsockname()[1]); s.close()')

# Write message file with info for the local connection
echo "\\\$HOSTNAME \\\$TOKEN \\\$PBS_JOBID \\\$PORT" > "\$WORKDIR/message"

echo "runjp log dir \$WORKDIR"
cat "\$WORKDIR/message"

export DASK_LABEXTENSION__FACTORY__MODULE=dask.distributed
export DASK_LABEXTENSION__FACTORY__CLASS=LocalCluster
export DASK_LABEXTENSION__FACTORY__KWARGS__MEMORY_LIMIT=3900MB
export DASK_LABEXTENSION__FACTORY__KWARGS__LOCAL_DIRECTORY=\\\$PBS_JOBFS/dask-worker-space
export DASK_LABEXTENSION__DEFAULT__WORKERS=\\\$PBS_NCPUS
export DASK_DISTRIBUTED__DASHBOARD__LINK="/proxy/{port}/status"
export DASK_TEMPORARY_DIRECTORY="\\\$TMPDIR"

$STRACE jupyter lab --NotebookApp.token="\\\$TOKEN" --no-browser --ip="\\\$HOSTNAME" --port "\\\$PORT" --port-retries=0
EOQ

# Required for conda
storage="gdata/hh5"

if [ -n "$STORAGE" ]; then
    storage="$STORAGE+\$storage"
else
    for p in \$(id -nG); do
        if [ \$p = "access.dev" ]; then continue; fi
        if [ \$p = "access.admin" ]; then continue; fi
        if [ -d "/scratch/\$p" ]; then storage="scratch/\$p+\$storage"; fi
        if [ -d "/g/data/\$p" ]; then storage="gdata/\$p+\$storage"; fi
    done
fi

qsub $SUBMITOPTS -W umask=0022 -l "storage=\$storage" -j oe -o "\$WORKDIR/pbs.log" "\$WORKDIR/runjp.sh" > "\$WORKDIR/jobid"

# Wait for the message file to appear, then return to the local process
while [ ! -f "\$WORKDIR/message" ]; do
    sleep 5
done
cat "\$WORKDIR/message" | sed 's/$/ NEW/'
EOF
)

if [ -n "$DEBUG" ]; then
    echo "DEBUG: Remote Message: '$message'"
fi

# Grab info from the PBS job
read jobhost token jobid remote_port type <<< "$message"

if [ "$type" = "ERROR" ]; then
    echo
    echo "Stopping due to error"
    exit
fi
if [ "$type" = "RECONNECT" ]; then
    echo
    echo "Existing jupyterlab found, reconnecting to that instead"
fi

# Find a local port
for local_port in {8888..9000}; do
    if ! echo > /dev/tcp/127.0.0.1/${local_port} ; then
        break
    fi 2> /dev/null
done

echo
echo "Notebook running as PBS job ${jobid}"
echo
echo "Starting tunnel..."
if [ -n "$DEBUG" ]; then
    echo "DEBUG:" $SSH -N -L "${local_port}:$jobhost:${remote_port}" "$LOGINNODE"
fi
$SSH -q -N -L "${local_port}:$jobhost:${remote_port}" "$LOGINNODE" &
tunnelid=$!

# Shut everything down on exit
trap "{ echo 'Closing connections... (Ctrl-C will leave job in the queue)' ; kill $tunnelid ; $SSH "$LOGINNODE" qdel $jobid ; }" EXIT

URL="http://localhost:${local_port}/lab?token=${token}"

# Wait for startup then open a browser
if which curl > /dev/null; then
    until curl $URL &> /dev/null; do echo -n '.'; sleep 1; done
    echo
else
    sleep 5
fi

cat << EOF

Start a Dask cluster in your notebook using the Dask panel of Jupyterlab, or by
running (needs kernel analysis3-20.01 or later):

---------------------------------------------------------------
import climtas.nci

climtas.nci.GadiClient()
---------------------------------------------------------------

Opening ${URL}
EOF

set +e
if [ -f "/mnt/c/Windows/system32/cmd.exe" ]; then
    # Windows subsystem for linux
    /mnt/c/Windows/system32/cmd.exe /c start "$URL"
elif which xdg-open > /dev/null; then
    # Linux
    xdg-open "$URL"
elif which open > /dev/null; then
    # OSX
    open "$URL"
elif which explorer > /dev/null; then
    # Windows
    explorer "$URL"
else
    echo
    echo "----"
    echo "Can't find a way to open the URL automatically, please copy the URL above into your browser"
fi
set -e
echo

$SSH "$LOGINNODE" /g/data/hh5/public/apps/nci_scripts/qmonitor $jobid

# Move the cursor past the progress bars
echo
echo
echo