Commit 7065fab6 authored by Mathis Rosenhauer's avatar Mathis Rosenhauer

Script for jupyter on mistral

parent 5a0accae
#!/bin/bash
#
# Copyright 2018 Deutsches Klimarechenzentrum GmbH
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following
# disclaimer in the documentation and/or other materials provided
# with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
# start-jupyter
#
# This script is intended to be used on your local workstation running
# Ubuntu, Fedora Linux or macOS (tested). Other Unix flavors may work
# as well. It allows you to start jupyter notebook or lab on one of
# DKRZ's mistral nodes. The script then opens a ssh tunnel to connect
# your local browser to jupyter.
#
# If you indicate an account with the -A switch, the script will run
# jupyter in a job on dedicated resources. Otherwise jupyter uses a
# shared interactive node.
#
# In case of problems contact Mathis Rosenhauer <rosenhauer@dkrz.de>.
#
set -eufo pipefail
# Default settings
#
# You can change the settings here or override them with command line
# options.
# Project account code.
#
# Jupyter is started on the frontend if this is not set.
SJ_ACCTCODE=""
# LDAP username
#
# Specify your username on the frontend if it is not your local
# username.
SJ_USERNAME="$(id -un)"
# Jupyter command
#
# You could change this to lab for example
SJ_COMMAND="notebook"
# Generate debugging output if set to 1
SJ_DEBUG=0
# Session run time in minutes
SJ_RUNTIME=240
# Ntasks for job
SJ_NTASKS=1
# Partition for job
SJ_PARTITION=shared
# Incfile
#
# If indicated, this file will be sourced prior to running jupyter
# notebook. It has to be located on mistral. Set up the environment
# for starting the correct jupyter here.
SJ_INCFILE=""
# Frontend host
#
# Must be directly accessible from client. The frontend and the node
# where jupyter is running need a shared home file system.
readonly SJ_FRONTEND_HOST="mistralpp.dkrz.de"
function clean_up () {
trap - ERR EXIT
set +e
echo
if [[ -n ${SJ_ACCTCODE} ]]; then
if [[ -n ${jupyter_id:-} ]]; then
echo "Removing job ${jupyter_id}."
ssh_frontend "scancel -Q ${jupyter_id}"
else
echo "Job ID not available. Make sure the jupyter " \
"job is not running!"
ssh_frontend "squeue -u ${SJ_USERNAME}"
fi
else
if [[ -n ${jupyter_id:-} ]]; then
echo "Killing jupyter process ${jupyter_id}..."
ssh_frontend "kill ${jupyter_id}"
fi
fi
ssh_frontend "rm -f ${jupyter_log}"
ssh_frontend "" "-O exit"
rmdir "${ssh_socket_dir}"
exit
}
function usage () {
cat <<EOF
Usage: $(basename "$0") [OPTION]
Available values for OPTION include:
-A acctcode start a job with acctcode
-c command invoke jupyter with command
-d check for presence of jupyter
-i file source file prior to running jupyter
-n ntasks request ntasks tasks for job
-p partition run job on partition
-t time job runtime
-u username use username for login
EOF
}
function parse_options () {
local option
while getopts 'A:c:di:n:p:t:u:' option; do
case ${option} in
A) SJ_ACCTCODE="$OPTARG"
;;
c) SJ_COMMAND="$OPTARG"
;;
d) SJ_DEBUG=1
;;
i) SJ_INCFILE="$OPTARG"
;;
n) SJ_NTASKS="$OPTARG"
;;
p) SJ_PARTITION="$OPTARG"
;;
t) SJ_RUNTIME="$OPTARG"
;;
u) SJ_USERNAME="$OPTARG"
;;
?) usage
exit 1
;;
esac
done
readonly SJ_ACCTCODE
readonly SJ_COMMAND
readonly SJ_DEBUG
readonly SJ_INCFILE
readonly SJ_NTASKS
readonly SJ_PARTITION
readonly SJ_RUNTIME
readonly SJ_USERNAME
}
function ssh_frontend () {
# Run command on frontend with ssh
local command="$1"
local extra_options="${2:-}"
local options
options="${extra_options} -o ForwardX11=no \
-o ControlPath=${ssh_socket_dir}/control:%h:%p:%r"
ssh ${options} "${SJ_USERNAME}@${SJ_FRONTEND_HOST}" "${command}"
}
function source_incfile() {
# Add sourcing of incfile to commandline if an incfile was
# specified
local commandline="$1"
if [[ -n ${SJ_INCFILE} ]]; then
local incfile="${SJ_INCFILE}"
if [[ ${incfile:0:1} != "/" ]]; then
incfile="\${HOME}/${incfile}"
fi
commandline="source ${incfile}; ${commandline}"
fi
echo "${commandline}"
}
function which_jupyter() {
echo "Looking for Jupyter."
local which
which="$(source_incfile "which jupyter")"
ssh_frontend "/bin/bash -lc \"${which}\""
}
function assemble_commandline () {
local logfile="$1"
local commandline="jupyter ${SJ_COMMAND} --no-browser 2>> ${logfile}"
# If we are not running a job, we have to perform our own scheduling
if [[ -z ${SJ_ACCTCODE} ]]; then
commandline="nohup ${commandline} > /dev/null & echo \$!"
fi
commandline="$(source_incfile "${commandline}")"
echo "${commandline}"
}
function submit_jupyter_job () {
local commandline="$1"
local logfile="$2"
ssh_frontend "sbatch" <<EOF
#!/bin/bash -l
#SBATCH --job-name=Jupyter
#SBATCH --partition=${SJ_PARTITION}
#SBATCH --ntasks=${SJ_NTASKS}
#SBATCH --time=${SJ_RUNTIME}
#SBATCH --account=${SJ_ACCTCODE}
#SBATCH --output=/dev/null
#SBATCH --parsable
#SBATCH --dkrzepilog=0
cd \${HOME}
echo "NODE:\${SLURM_JOB_NODELIST}" > ${logfile}
${commandline}
EOF
}
function run_jupyter () {
local logfile="$1"
local commandline
commandline="$(assemble_commandline "${jupyter_log}")"
# Run commandline in job or directly on frontend
if [[ -n ${SJ_ACCTCODE} ]]; then
submit_jupyter_job "${commandline}" "${logfile}"
else
ssh_frontend "/bin/bash -ls" <<< "${commandline}"
fi
}
function extract_from_logs () {
local pattern="$1"
local haystack="$2"
ssh_frontend "/bin/bash -s" <<EOF
while [[ -z \${needle} ]]; do
sleep 1
if [[ -f ${haystack} ]]; then
needle="\$(grep -Pom 1 "${pattern}" "${haystack}")"
fi
printf "." >&2
done
printf "\n" >&2
echo "\${needle}"
EOF
}
function get_jupyter_node () {
local logfile="$1"
if [[ -n ${SJ_ACCTCODE} ]]; then
printf "Waiting for job to start" >&2
extract_from_logs "NODE:\K\w+" "${logfile}"
else
echo "${SJ_FRONTEND_HOST}"
fi
}
function open_tunnel () {
local node="$1"
local port="$2"
if [[ -n ${SJ_ACCTCODE} ]]; then
# Tunnel to notebook in job needs jump host since nodes
# usually have no direct external access.
# Unfortunately, doesn' seem to support connection sharing
# for the jump host
ssh -o ForwardX11=no \
-o StrictHostKeyChecking=accept-new \
-J "${SJ_USERNAME}@${SJ_FRONTEND_HOST}" \
-L "${port}:localhost:${port}" \
-Nf \
"${SJ_USERNAME}@${node}"
else
ssh_frontend "" "-O forward -L${port}:localhost:${port}"
fi
}
function show_url() {
local url="$1"
echo "Open the following URL with your browser"
echo "${url}"
}
function launch_browser() {
local url="$1"
case "$(uname -s)" in
Darwin)
open "${url}" || show_url "${url}"
;;
Linux)
xdg-open "${url}" || show_url "${url}"
;;
*)
show_url "${url}"
;;
esac
}
function main () {
parse_options "$@"
trap clean_up INT QUIT TERM ERR EXIT
echo "Establishing ssh master connection."
# Set up control master for connection sharing
ssh_socket_dir="$(mktemp -d "${HOME}/.ssh/socket.XXXXX")"
ssh_frontend "" "-MNf"
# Create unique output file for jupyter notebook
jupyter_log="$(ssh_frontend "mkdir -p \${HOME}/.jupyter \
&& mktemp \${HOME}/.jupyter/jupyter.XXXXX")"
ssh_frontend "chmod 600 ${jupyter_log}"
# Check for jupyter
[[ ${SJ_DEBUG} == 1 ]] && which_jupyter
jupyter_id="$(run_jupyter "${jupyter_log}")"
local node
node="$(get_jupyter_node "${jupyter_log}")"
# Get notebook url and token from output
printf "Starting jupyter server" >&2
local url
url="$(extract_from_logs "^.*\Khttp://localhost:.+" "${jupyter_log}")"
ssh_frontend "rm -f ${jupyter_log}"
local port
port=${url#*t:}
port=${port%%/*}
open_tunnel "${node}" "${port}"
echo "Established tunnel to ${node}:${port}."
launch_browser "${url}"
echo "Press Ctrl-C to stop jupyter and shut down tunnel."
sleep 604800
}
main "$@"
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment