Split off some functions

startvnc

@@ -14,55 +14,57 @@
 # The vncviewer client has to support the -via option. TurboVNC,
 # TigerVNC and TightVNC are known to work.
 #
-# Set SSH_MUX=true to share one ssh connection. This can speed up the
-# script if ssh login is slow. If you have to use password
+# Set SVNC_SSH_MUX="true" to share one ssh connection. This can speed
+# up the script if ssh login is slow. If you have to use password
 # authentication, you will only have to give your password twice.
 
 set -e
 
 # Default settings
-readonly USERNAME=                      # LDAP username
-readonly ACCTCODE=                      # project account code
-readonly TIME=240                       # session run time in minutes
-readonly QOS=normal                     # quality of service
-readonly GEOMETRY="1920x1200"           # resolution of remote desktop window
-readonly VNCVIEWER=/usr/bin/vncviewer   # path to local vncviewer
-readonly PASSWORD=true                  # copy vncpassword temporarily to
-                                        # the local workstation
-readonly SSH_MUX=true                   # use sharing of multiple ssh sessions
-                                        # over a single network connection
-#readonly OPTIONS="-extssh"             # uncomment this if you use
-                                        # TurboVNC as vncviewer
+
+SVNC_USERNAME=""                     # LDAP username
+SVNC_ACCTCODE=""                     # project account code
+SVNC_RUNTIME=240                     # session run time in minutes
+SVNC_QOS="normal"                    # quality of service
+SVNC_GEOMETRY="1920x1200"            # resolution of remote desktop window
+SVNC_VNCVIEWER="/usr/bin/vncviewer"  # path to local vncviewer
+
+readonly SVNC_PASSWORD="true"        # copy vncpassword temporarily to
+                                     # the local workstation
+readonly SVNC_SSH_MUX="true"         # use sharing of multiple ssh sessions
+                                     # over a single network connection
+#readonly SVNC_VNC_OPTIONS="-extssh" # uncomment this if you use
+                                     # TurboVNC as vncviewer
 
 clean_up () {
     trap - ERR EXIT
     set +e
-    local tmp_files=${job_script}
+    local tmp_files="${job_script_name}"
 
     if [[ ${vnc_host} ]]; then
-        ${ssh_mistral} "ssh ${vnc_host} \
-            \"/opt/TurboVNC/bin/vncserver -kill :${vnc_display}\""
+        ssh_mistral "ssh ${vnc_host} \"/opt/TurboVNC/bin/vncserver \
+                    -kill :${vnc_display}\""
     fi
 
     echo "Removing job and temporary files."
     if [[ ${job_id} ]]; then
         tmp_files="${tmp_files} out.${job_id} gethost.${job_id}"
-        ${ssh_mistral} "scancel -Q ${job_id}; sleep 1"
+        ssh_mistral "scancel -Q ${job_id}; sleep 1"
     else
         echo "Job ID not available. Make sure the vncjob is not running!"
-        ${ssh_mistral} "squeue -u ${username}"
+        ssh_mistral "squeue -u ${SVNC_USERNAME}"
     fi
 
     if [[ ${tmp_files} ]]; then
-        ${ssh_mistral} "cd .startvnc && rm -f ${tmp_files}"
+        ssh_mistral "cd .startvnc && rm -f ${tmp_files}"
     fi
 
-    if [[ ${SSH_MUX} = true ]]; then
-        ssh ${ssh_options} -O exit ${mistral}
+    if [[ ${SVNC_SSH_MUX} = "true" ]]; then
+        ssh_mistral "" "-O exit"
     fi
 
     # Remove local vnc PasswordFile
-    if [[ ${copy_password} = true ]]; then
+    if [[ ${SVNC_PASSWORD} = "true" ]]; then
         rm -f ${HOME}/vnc_passwd
     fi
     exit
@@ -74,84 +76,88 @@ usage () {
     printf -- "-v path_to_vncviewer\n-x vncviewer_options\n" >&2
 }
 
-while getopts 'A:g:q:t:u:v:x:' OPTION
-do
-    case $OPTION in
-        A) acctcode="$OPTARG"
-           ;;
-        g) geometry="$OPTARG"
-           ;;
-        q) qos="$OPTARG"
-           ;;
-        t) runtime="$OPTARG"
-           ;;
-        u) username="$OPTARG"
-           ;;
-        v) vncviewer="$OPTARG"
-           ;;
-        x) vnc_options="$OPTARG"
-           ;;
-        ?) usage
+ssh_mistral () {
+    local ssh_options="$2 -o ForwardX11=no"
+    if [[ ${SVNC_SSH_MUX} = "true" ]]; then
+        ssh_options+=" -o ControlPath=~/.ssh/control:%h:%p:%r"
+    fi
+    ssh ${ssh_options} "${SVNC_USERNAME}@mistral00.dkrz.de" "$1"
+}
+
+parse_options () {
+    local option
+    while getopts 'A:g:q:t:u:v:x:' option
+    do
+        case ${option} in
+            A) SVNC_ACCTCODE="$OPTARG"
+               ;;
+            g) SVNC_GEOMETRY="$OPTARG"
+               ;;
+            q) SVNC_QOS="$OPTARG"
+               ;;
+            t) SVNC_RUNTIME="$OPTARG"
+               ;;
+            u) SVNC_USERNAME="$OPTARG"
+               ;;
+            v) SVNC_VNCVIEWER="$OPTARG"
+               ;;
+            x) SVNC_VNC_OPTIONS="$OPTARG"
+               ;;
+            ?) usage
+            exit 2
+            ;;
+        esac
+    done
+
+    readonly SVNC_ACCTCODE
+    readonly SVNC_GEOMETRY
+    readonly SVNC_QOS
+    readonly SVNC_RUNTIME
+    readonly SVNC_USERNAME
+    readonly SVNC_VNCVIEWER
+    readonly SVNC_VNC_OPTIONS
+
+    if [[ -z ${SVNC_ACCTCODE} ]] || [[ -z ${SVNC_USERNAME} ]]; then
+        printf "ERROR: Please specify at least username and acctcode.\n\n" >&2
+        usage
         exit 2
-        ;;
-    esac
-done
+    fi
+}
 
-acctcode=${acctcode:-${ACCTCODE}}
-geometry=${geometry:-${GEOMETRY}}
-qos=${qos:-${QOS}}
-runtime=${runtime:-${TIME}}
-username=${username:-${USERNAME}}
-vncviewer=${vncviewer:-${VNCVIEWER}}
-vnc_options=${vnc_options:-${OPTIONS}}
-mistral="${username}@mistral00.dkrz.de"
-ssh_options="-o ForwardX11=no"
-
-if [[ ${SSH_MUX} = true ]]; then
-    ssh_options+=" -o ControlPath=~/.ssh/control:%h:%p:%r"
-fi
-
-ssh_mistral="ssh ${ssh_options} ${mistral}"
-
-if [[ -z ${acctcode} ]] || [[ -z ${username} ]]; then
-    printf "ERROR: Please specify at least username and acctcode.\n\n" >&2
-    usage
-    exit 2
-fi
-
-if [[ ${SSH_MUX} = true ]]; then
-    ssh -MNf ${ssh_options} ${mistral}
-fi
-
-if ! ${ssh_mistral} "test -s .vnc/passwd"; then
-    echo "No VNC password found. Please set now."
-    ${ssh_mistral} -t "salloc -Q -n1 -pgpu -A${acctcode} -- /bin/bash -c \
-        'ssh -tt \$SLURM_JOB_NODELIST -- mkdir -p .vnc \
-        && /opt/TurboVNC/bin/vncpasswd'"
-fi
-
-trap clean_up INT QUIT TERM ERR EXIT
-
-job_script=$(${ssh_mistral} \
-             "mkdir -p .startvnc && cd .startvnc && mktemp vncjob.XXXXX")
-${ssh_mistral} "cat > .startvnc/${job_script}" <<EOF
-#!/bin/bash
+check_vnc_password () {
+    if ! ssh_mistral "test -s .vnc/passwd"; then
+        echo "No VNC password found. Please set now."
+        ssh_mistral "salloc -Q -n1 -pgpu -A${SVNC_ACCTCODE} -- /bin/bash -c \
+                    'ssh -tt \$SLURM_JOB_NODELIST -- mkdir -p .vnc \
+                    && /opt/TurboVNC/bin/vncpasswd'" "-t"
+    fi
+}
+
+get_job_script_name () {
+    ssh_mistral "mkdir -p .startvnc && cd .startvnc && mktemp vncjob.XXXXX"
+}
+
+submit_vnc_job () {
+    local job_script_name="$1"
+    ssh_mistral "cat > .startvnc/${job_script_name}" <<EOF
+#!/bin/bash -l
 #SBATCH --partition=gpu
 #SBATCH --ntasks=12
-#SBATCH --qos=${qos}
-#SBATCH --time=${runtime}
-#SBATCH --account=${acctcode}
+#SBATCH --qos=${SVNC_QOS}
+#SBATCH --time=${SVNC_RUNTIME}
+#SBATCH --account=${SVNC_ACCTCODE}
 #SBATCH --output=out.%j
 #SBATCH --parsable
-##SBATCH --dkrzepilog=0
+#SBATCH --dkrzepilog=0
 cd \${HOME}
-/opt/TurboVNC/bin/vncserver -fg -geometry ${geometry} -localhost
+/opt/TurboVNC/bin/vncserver -fg -geometry "${SVNC_GEOMETRY}" -localhost
 EOF
+    ssh_mistral "cd .startvnc && sbatch ${job_script_name}"
+}
 
-echo "Submitting vncserver job."
-job_id=$(${ssh_mistral} "cd .startvnc && sbatch ${job_script}")
-
-${ssh_mistral} "cat > .startvnc/gethost.${job_id}" <<EOF
+get_vnc_host_and_display () {
+    local job_id="$1"
+    ssh_mistral "cat > .startvnc/gethost.${job_id}" <<EOF
 #!/bin/bash
 LOGFILE=\${HOME}/.startvnc/out.${job_id}
 while [[ -z \${host_and_display} ]]; do
@@ -163,18 +169,36 @@ while [[ -z \${host_and_display} ]]; do
 done
 echo \${host_and_display}
 EOF
+    ssh_mistral "/bin/bash .startvnc/gethost.${job_id}"
+}
+
+main () {
+    parse_options "$@"
+    trap clean_up INT QUIT TERM ERR EXIT
 
-echo "Waiting for job ${job_id} to start."
-host_and_display=$(${ssh_mistral} "/bin/bash .startvnc/gethost.${job_id}")
+    if [[ ${SVNC_SSH_MUX} = "true" ]]; then
+        ssh_mistral "" "-MNf"
+    fi
+
+    check_vnc_password
+
+    job_script_name=$(get_job_script_name)
+    echo "Submitting vncserver job."
+    job_id=$(submit_vnc_job "${job_script_name}")
 
-vnc_host=${host_and_display%:*}
-vnc_display=${host_and_display#*:}
-vnc_options+=" -via ${username}@${vnc_host}.dkrz.de"
+    echo "Waiting for job ${job_id} to start."
+    local host_and_display=$(get_vnc_host_and_display "$job_id")
+    vnc_host=${host_and_display%:*}
+    vnc_display=${host_and_display#*:}
 
-if [[ ${PASSWORD} = true ]]; then
-    scp -q ${ssh_options} ${mistral}:.vnc/passwd ${HOME}/vnc_passwd
-    vnc_options+=" -passwd ${HOME}/vnc_passwd"
-fi
+    local vnc_options="${SVNC_VNC_OPTIONS} \
+                       -via ${SVNC_USERNAME}@${vnc_host}.dkrz.de"
+    if [[ ${SVNC_PASSWORD} = "true" ]]; then
+        ssh_mistral "cat .vnc/passwd" > ${HOME}/vnc_passwd
+        vnc_options+=" -passwd ${HOME}/vnc_passwd"
+    fi
+
+    ${SVNC_VNCVIEWER} ${vnc_options} :${vnc_display}
+}
 
-# Start VNC viewer
-${vncviewer} ${vnc_options} :${vnc_display}
+main "$@"