Levante

start-vnc

@@ -91,6 +91,9 @@ SVNC_PARTITION="gpu"
 # Quality of service
 SVNC_QOS="normal"
 
+# Use a reservation
+SVNC_RESERVATION=""
+
 # Constraint for feature selection (GPU and/or memory)
 SVNC_CONSTRAINT=""
 
@@ -108,7 +111,7 @@ SVNC_CONSTRAINT=""
 # where vncserver is running need a shared home file system. You will
 # have to change this for other sites than DKRZ (along with many other
 # things).
-readonly SVNC_FRONTEND_HOST="mistral.dkrz.de"
+readonly SVNC_FRONTEND_HOST="levante.dkrz.de"
 
 # Copy vncpassword temporarily to the local workstation
 #
@@ -125,22 +128,22 @@ clean_up () {
     set +e
 
     if [[ "${SVNC_KILLSERVER}" = true ]]; then
-        if [[ -n ${vnc_host:-} ]]; then
-            echo "Killing vncserver :${vnc_display} on ${vnc_host}."
-            ssh_frontend "ssh ${vnc_host} \"/opt/TurboVNC/bin/vncserver \
-                         -kill :${vnc_display}\""
+        if [[ -n ${SVNC_HOST:-} ]]; then
+            echo "Killing vncserver :${SVNC_DISPLAY} on ${SVNC_HOST}."
+            ssh_frontend "ssh ${SVNC_HOST} \"/opt/TurboVNC/bin/vncserver \
+                         -kill :${SVNC_DISPLAY}\""
         fi
 
-        if [[ -n ${job_id:-} ]]; then
-            echo "Removing job ${job_id}."
-            ssh_frontend "scancel -Q ${job_id}; sleep 1; \
-                         rm -f .startvnc/out.${job_id}"
+        if [[ -n ${SVNC_JOB_ID:-} ]]; then
+            echo "Removing job ${SVNC_JOB_ID}."
+            ssh_frontend "scancel -Q ${SVNC_JOB_ID}; sleep 1; \
+                         rm -f .startvnc/out.${SVNC_JOB_ID}"
         else
             echo "Job ID not available. Make sure the vncjob is not running!"
             ssh_frontend "squeue -u ${SVNC_USERNAME}"
         fi
     else
-        echo "Kept job ${job_id} running."
+        echo "Kept job ${SVNC_JOB_ID} running."
     fi
 
     ssh_frontend "" "-O exit"
@@ -159,16 +162,17 @@ Usage: $(basename "$0") [OPTION]
 
 Available values for OPTION include:
 
-  -A acctcode  acctcode of job
-  -C list      constraint for feature selection
-  -n nodes     number of nodes
-  -p partition run job on partition
-  -q qos       job qos
-  -s options   addidtional options for vncserver
-  -t time      job runtime
-  -u username  use username for login
-  -v path      path to vncviewer program
-  -x options   addidtional options for vncviewer
+  -A acctcode     acctcode of job
+  -C list         constraint for feature selection
+  -n nodes        number of nodes
+  -p partition    run job on partition
+  -q qos          job qos
+  -r reservation  use reservation
+  -s options      addidtional options for vncserver
+  -t time         job runtime
+  -u username     use username for login
+  -v path         path to vncviewer program
+  -x options      addidtional options for vncviewer
 
 EOF
 }
@@ -184,7 +188,7 @@ ssh_frontend () {
 
 parse_options () {
     local option
-    while getopts 'A:C:n:p:q:s:t:u:v:x:' option; do
+    while getopts 'A:C:n:p:q:r:s:t:u:v:x:' option; do
         case ${option} in
             A) SVNC_ACCTCODE="$OPTARG"
                ;;
@@ -196,6 +200,8 @@ parse_options () {
                ;;
             q) SVNC_QOS="$OPTARG"
                ;;
+            r) SVNC_RESERVATION="$OPTARG"
+               ;;
             s) SVNC_SERVER_OPTIONS="$OPTARG"
                ;;
             t) SVNC_RUNTIME="$OPTARG"
@@ -217,6 +223,7 @@ parse_options () {
     readonly SVNC_SERVER_OPTIONS
     readonly SVNC_PARTITION
     readonly SVNC_QOS
+    readonly SVNC_RESERVATION
     readonly SVNC_RUNTIME
     readonly SVNC_USERNAME
     readonly SVNC_CLIENT
@@ -234,9 +241,11 @@ prepare_vncserver () {
     if ! ssh_frontend "test -s .vnc/passwd"; then
         echo "No VNC password found. Please set now."
         echo "Do not use your LDAP password. Eight characters maximum."
-        ssh_frontend "salloc -Q -n1 -pgpu -A${SVNC_ACCTCODE} -- /bin/bash -c \
-                     'ssh -tt \$SLURM_JOB_NODELIST -- mkdir -p .vnc \
-                     && /opt/TurboVNC/bin/vncpasswd'" "-t"
+        echo "Wait for the prompt."
+        ssh_frontend "salloc -Q -n1 -p${SVNC_PARTITION} \
+                     --reservation=${SVNC_RESERVATION} \
+                     -A${SVNC_ACCTCODE} \
+                     srun --pty /opt/TurboVNC/bin/vncpasswd" "-t"
     fi
 }
 
@@ -255,25 +264,23 @@ submit_vnc_job () {
 #SBATCH --constraint=${SVNC_CONSTRAINT}
 ${sbatch_resources}
 #SBATCH --qos=${SVNC_QOS}
+#SBATCH --reservation=${SVNC_RESERVATION}
 #SBATCH --time=${SVNC_RUNTIME}
 #SBATCH --account=${SVNC_ACCTCODE}
 #SBATCH --output=out.%j
 #SBATCH --parsable
-#SBATCH --dkrzepilog=0
 cd \${HOME}
 /opt/TurboVNC/bin/vncserver -fg -localhost ${SVNC_SERVER_OPTIONS}
 EOF
 }
 
 get_vnc_host_and_display () {
-    local job_id="$1"
     ssh_frontend "/bin/bash -s" <<EOF
-logfile=\${HOME}/.startvnc/out.${job_id}
+logfile=\${HOME}/.startvnc/out.${SVNC_JOB_ID}
 while [[ -z \${host_and_display} ]]; do
     sleep 1
     if [[ -f \${logfile} ]]; then
-        host_and_display="\$(grep -Po "started on display \Kmg[0-9]+:[0-9]+" \
-                          \${logfile})"
+        host_and_display="\$(grep -Po "started on display \Kl[g0-9][0-9]+.atos.local:[0-9]+" \${logfile} | sed 's/.atos.local//g')"
     fi
     printf "." >&2
 done
@@ -282,47 +289,54 @@ echo \${host_and_display}
 EOF
 }
 
+get_jumphost_options () {
+    # Check for -J option introduced with OpenSSH 7.3
+    if (ssh || true) 2>&1 | grep -q -- "-J" ; then
+        echo "-J ${SVNC_USERNAME}@${SVNC_FRONTEND_HOST}"
+    else
+        echo "-o ProxyCommand=ssh ${SVNC_USERNAME}@${SVNC_FRONTEND_HOST} nc ${SVNC_HOST} 22"
+    fi
+}
+
 main () {
     parse_options "$@"
     trap clean_up INT QUIT TERM ERR EXIT
-
     mkdir -p "${HOME}/.ssh"
     ssh_socket_dir="$(mktemp -d "${HOME}/.ssh/socket.XXXXX")"
     ssh_frontend "" "-MNf"
     prepare_vncserver
 
     echo "Submitting vncserver job."
-    job_id="$(submit_vnc_job)"
+    SVNC_JOB_ID="$(submit_vnc_job)"
 
-    printf "Waiting for job ${job_id} to start" >&2
-    local host_and_display
-    host_and_display="$(get_vnc_host_and_display "$job_id")"
-    vnc_host=${host_and_display%:*}
-    vnc_display=${host_and_display#*:}
-    echo "Vncserver started on node ${vnc_host}.dkrz.de display \
-:${vnc_display}."
+    printf "Waiting for job ${SVNC_JOB_ID} to start" >&2
+    local host_and_display="$(get_vnc_host_and_display)"
+    SVNC_HOST=${host_and_display%:*}
+    SVNC_DISPLAY=${host_and_display#*:}
+    echo "Vncserver started on node ${SVNC_HOST}.dkrz.de display \
+:${SVNC_DISPLAY}."
 
     # WSL doesn't seem to detect used ports so randomize
     # local port to reduce risk of masking.
     local vnc_port_local=$((5900 + RANDOM % 100))
-    local vnc_port_remote=$(( 5900 + vnc_display ))
+    local vnc_port_remote=$(( 5900 + SVNC_DISPLAY ))
+    local jumphost_options="$(get_jumphost_options)"
     until ssh -o ForwardX11=no \
               -o StrictHostKeyChecking=ask \
-              -o ProxyCommand="ssh ${SVNC_USERNAME}@${SVNC_FRONTEND_HOST} nc ${vnc_host} 22" \
+              "${jumphost_options}" \
               -L "${vnc_port_local}:localhost:${vnc_port_remote}" \
               -Nf \
-              "${SVNC_USERNAME}@${vnc_host}"; do
+              "${SVNC_USERNAME}@${SVNC_HOST}"; do
         vnc_port_local=$((5900 + RANDOM % 100))
         echo "Trying local port ${vnc_port_local}."
     done
-    local client_options
-    client_options="${SVNC_CLIENT_OPTIONS:-}"
+    local client_options="${SVNC_CLIENT_OPTIONS:-}"
     if [[ "${SVNC_PASSWORD}" = true ]]; then
         echo "Fetching password from frontend."
         ssh_frontend "cat .vnc/passwd" > vnc_passwd
         client_options+=" -passwd vnc_passwd"
     fi
-    echo "Connecting vncviewer to ${vnc_host}.dkrz.de"
+    echo "Connecting vncviewer to ${SVNC_HOST}.dkrz.de"
     "${SVNC_CLIENT}" ${client_options} :$(( vnc_port_local - 5900 ))
 }