startvnc 8.37 KB
Newer Older
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
1
2
#!/bin/bash
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
3
# Copyright 2017 Deutsches Klimarechenzentrum GmbH
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
#
# 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.
#
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
32
# startvnc
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
33
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
34
35
36
37
38
39
40
41
42
# This script is intended to be used on your local workstation running
# Ubuntu or Fedora Linux or macOS (tested). Other Unix flavors may
# work as well. It allows you to connect to one of DKRZ's
# visualization nodes to work remotely on your visualizations.
#
# Technically, it starts a VNC server on one of the GPU nodes by
# submitting a generated job script. A local vncviewer client is used
# to connect to the server over an encypted ssh tunnel. After the
# client shuts down, the server job is terminated.
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
43
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
44
45
# Your local vncviewer client has to support the -via
# option. TigerVNC, TightVNC, and TurboVNC  are known to work.
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
46
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
47
# In case of problems contact Mathis Rosenhauer <rosenhauer@dkrz.de>.
48
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
49

Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
50
set -eu
51

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
52
# Default settings
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
53
54
55
56
#
# You can change the settings here or override them with command line
# options. SVNC_ACCTCODE is the only parameter you have to provide if
# your local username matches your username on the frontend and if
57
58
# vncviewer is installed in the search path (tight/tigervnc on Linux
# for example).
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
59
60
61
62

# Project account code
SVNC_ACCTCODE=""

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
63
64
65
66
67
# LDAP username
#
# Specify your username on the login node if it is not your local
# username.
SVNC_USERNAME=$(id -un)
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
68

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
69
70
# Path to local vncviewer
#
71
72
73
# If your local vnc client is not in the search path or named
# differently (e.g. TurboVNC, macOS), then change this parameter.
SVNC_VNCVIEWER="vncviewer"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
74
75

# Resolution of remote desktop window
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
76
77
78
#
# This is just the initial resolution. Resizing the vncviewer window
# later should also work.
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
79
80
SVNC_GEOMETRY="1920x1200"

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
81
82
# Session run time in minutes
SVNC_RUNTIME=240
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
83

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
84
85
86
87
88
# Number of GPU nodes reserved. Default is the special value "half"
# which allocates 24 cores of a node. To allocate one or more nodes,
# set to the integer number of nodes.
SVNC_NODES="half"

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
89
90
# Quality of service
SVNC_QOS="normal"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
91

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
92
93
94
95
96
97
98
99
100
101
102
103
104
105
# Uncomment this if you use TurboVNC as vncviewer.
#
# TurboVNC will work without this option but then it will use its own
# implementation of ssh which dosen't support public key auth among other
# things.
#
# readonly SVNC_VNC_OPTIONS="-extssh"

# Frontend host
#
# Must be directly accessible from client. The frontend and the node
# 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).
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
106
107
108
readonly SVNC_FRONTEND_HOST="mistral.dkrz.de"

# Copy vncpassword temporarily to the local workstation
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
109
110
111
#
# If your vncviewer client has trouble with TurboVNC's password format
# then disable this option (set to "false" or comment).
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
112
readonly SVNC_PASSWORD="true"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
113

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
114

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
115
clean_up () {
116
    trap - ERR EXIT
117
    set +e
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
118

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
119
    if [[ -n ${vnc_host:-} ]]; then
120
        echo "Killing vncserver :${vnc_display} on ${vnc_host}."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
121
122
        ssh_frontend "ssh ${vnc_host} \"/opt/TurboVNC/bin/vncserver \
                     -kill :${vnc_display}\""
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
123
124
    fi

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
125
    if [[ -n ${job_id:-} ]]; then
126
        echo "Removing job ${job_id}."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
127
        ssh_frontend "scancel -Q ${job_id}; sleep 1; \
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
128
                     cd .startvnc && rm -f out.${job_id}"
129
130
    else
        echo "Job ID not available. Make sure the vncjob is not running!"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
131
        ssh_frontend "squeue -u ${SVNC_USERNAME}"
132
133
    fi

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
134
135
    ssh_frontend "" "-O exit"
    rmdir "${ssh_socket_dir}"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
136

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
137
    # Remove local vnc PasswordFile
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
138
    if [[ ${SVNC_PASSWORD} = "true" ]]; then
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
139
        rm -f "${HOME}/vnc_passwd"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
140
    fi
141
    exit
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
142
143
}

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
144
145
usage () {
    printf "Usage: %s [OPTIONS]\n" $(basename $0) >&2
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
146
147
148
    printf -- "-u username\n-A acctcode\n-g geometry\n-n nodes\n" >&2
    printf -- "-q qos\n-t time\n-v path_to_vncviewer\n" >&2
    printf -- "-x vncviewer_options\n" >&2
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
149
150
}

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
151
ssh_frontend () {
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
152
153
154
    local command="$1"
    local extra_options="${2:-}"
    local options="${extra_options} -o ForwardX11=no \
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
155
          -o ControlPath=${ssh_socket_dir}/control:%h:%p:%r"
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
156
    ssh ${options} "${SVNC_USERNAME}@${SVNC_FRONTEND_HOST}" "${command}"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
157
158
159
160
}

parse_options () {
    local option
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
161
    while getopts 'A:g:n:q:t:u:v:x:' option; do
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
162
163
164
165
166
        case ${option} in
            A) SVNC_ACCTCODE="$OPTARG"
               ;;
            g) SVNC_GEOMETRY="$OPTARG"
               ;;
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
167
168
            n) SVNC_NODES="$OPTARG"
               ;;
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
169
170
171
172
173
174
175
176
177
178
179
            q) SVNC_QOS="$OPTARG"
               ;;
            t) SVNC_RUNTIME="$OPTARG"
               ;;
            u) SVNC_USERNAME="$OPTARG"
               ;;
            v) SVNC_VNCVIEWER="$OPTARG"
               ;;
            x) SVNC_VNC_OPTIONS="$OPTARG"
               ;;
            ?) usage
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
180
            exit 1
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
181
182
183
184
185
186
187
188
189
190
191
192
            ;;
        esac
    done

    readonly SVNC_ACCTCODE
    readonly SVNC_GEOMETRY
    readonly SVNC_QOS
    readonly SVNC_RUNTIME
    readonly SVNC_USERNAME
    readonly SVNC_VNCVIEWER
    readonly SVNC_VNC_OPTIONS

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
193
194
    if [[ -z ${SVNC_ACCTCODE} ]]; then
        printf "ERROR: Please specify an acctcode.\n\n" >&2
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
195
        usage
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
196
        exit 1
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
197
198
    fi
}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
199

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
200
prepare_vncserver () {
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
201
202
    ssh_frontend "mkdir -p .startvnc"
    if ! ssh_frontend "test -s .vnc/passwd"; then
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
203
        echo "No VNC password found. Please set now."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
204
205
206
        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"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
207
208
209
210
    fi
}

submit_vnc_job () {
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
211
212
213
214
    if [[ ${SVNC_NODES} = "half" ]]; then
        local sbatch_resources="#SBATCH --ntasks=24"
    else
        local sbatch_resources="#SBATCH --nodes=${SVNC_NODES}"
215
        sbatch_resources+=$'\n#SBATCH --exclusive'
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
216
    fi
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
217
    ssh_frontend "cd .startvnc && sbatch" <<EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
218
#!/bin/bash -l
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
219
#SBATCH --partition=gpu
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
220
${sbatch_resources}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
221
222
223
#SBATCH --qos=${SVNC_QOS}
#SBATCH --time=${SVNC_RUNTIME}
#SBATCH --account=${SVNC_ACCTCODE}
224
#SBATCH --output=out.%j
225
#SBATCH --parsable
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
226
#SBATCH --dkrzepilog=0
227
cd \${HOME}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
228
/opt/TurboVNC/bin/vncserver -fg -geometry "${SVNC_GEOMETRY}" -localhost
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
229
EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
230
}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
231

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
232
233
get_vnc_host_and_display () {
    local job_id="$1"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
234
    ssh_frontend "/bin/bash" <<EOF
235
logfile=\${HOME}/.startvnc/out.${job_id}
236
while [[ -z \${host_and_display} ]]; do
237
    sleep 1
238
    if [[ -f \${logfile} ]]; then
239
        host_and_display=\$(grep -Po "started on display \Kmg[0-9]+:[0-9]+" \
240
                          \${logfile})
241
    fi
242
done
243
echo \${host_and_display}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
244
EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
245
246
247
248
249
}

main () {
    parse_options "$@"
    trap clean_up INT QUIT TERM ERR EXIT
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
250

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
251
252
    ssh_socket_dir=$(mktemp -d "${HOME}/.ssh/socket.XXXXX")
    ssh_frontend "" "-MNf"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
253
    prepare_vncserver
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
254
255

    echo "Submitting vncserver job."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
256
    job_id=$(submit_vnc_job)
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
257

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
258
259
260
261
    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#*:}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
262
263
    echo "Vncserver started on node ${vnc_host}.dkrz.de display \
:${vnc_display}."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
264

Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
265
    local vnc_options="${SVNC_VNC_OPTIONS:-} \
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
266
267
                       -via ${SVNC_USERNAME}@${vnc_host}.dkrz.de"
    if [[ ${SVNC_PASSWORD} = "true" ]]; then
268
        echo "Fetching password from frontend."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
269
        ssh_frontend "cat .vnc/passwd" > ${HOME}/vnc_passwd
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
270
271
        vnc_options+=" -passwd ${HOME}/vnc_passwd"
    fi
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
272
    echo "Connecting vncviewer to ${vnc_host}.dkrz.de"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
273
    "${SVNC_VNCVIEWER}" ${vnc_options} :${vnc_display}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
274
}
275

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
276
main "$@"