start-vnc 9.5 KB
Newer Older
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
1
2
#!/bin/bash
#
3
# Copyright 2019 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.
#
#
32
# start-vnc
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
33
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
34
35
36
37
38
39
40
41
# 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
42
43
# client shuts down, the server job is terminated. TigerVNC, TightVNC,
# and TurboVNC are known to work.
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
44
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
45
# In case of problems contact Mathis Rosenhauer <rosenhauer@dkrz.de>.
46
#
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
47

48
set -eufo pipefail
49

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
50
# Default settings
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
51
52
53
54
#
# 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
55
56
# vncviewer is installed in the search path (tight/tigervnc on Linux
# for example).
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
57
58
59
60

# Project account code
SVNC_ACCTCODE=""

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

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

73
# Server options
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
74
#
75
76
77
# More options for the vncserver. TurboVNC on A Mac may produce
# a 'javax.net.ssl.SSLHandshakeException' in this case adding
# ' -securitytypes none' to SVNC_SERVER_OPTIONS may help.
78
SVNC_SERVER_OPTIONS="-geometry 1920x1200"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
79

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

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
83
84
85
86
87
# 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
88
89
# Quality of service
SVNC_QOS="normal"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
90

91
92
93
# Constraint for feature selection (GPU and/or memory)
SVNC_CONSTRAINT=""

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
94
95
96
97
98
99
# 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.
#
100
# readonly SVNC_CLIENT_OPTIONS="-extssh"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
101
102
103
104
105
106
107

# 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
108
109
110
readonly SVNC_FRONTEND_HOST="mistral.dkrz.de"

# Copy vncpassword temporarily to the local workstation
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
111
112
113
#
# 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
114
readonly SVNC_PASSWORD="true"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
115

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
116

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

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

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

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

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

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
146
usage () {
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
147
148
149
150
151
152
    cat <<EOF
Usage: $(basename "$0") [OPTION]

Available values for OPTION include:

  -A acctcode  acctcode of job
153
  -C list      constraint for feature selection
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
154
155
  -n nodes     number of nodes
  -q qos       job qos
156
  -s options   addidtional options for vncserver
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
157
158
159
160
161
162
  -t time      job runtime
  -u username  use username for login
  -v path      path to vncviewer program
  -x options   addidtional options for vncviewer

EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
163
164
}

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
165
ssh_frontend () {
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
166
167
    local command="$1"
    local extra_options="${2:-}"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
168
169
170
    local options
    options="${extra_options} -o ForwardX11=no \
            -o ControlPath=${ssh_socket_dir}/control:%h:%p:%r"
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
171
    ssh ${options} "${SVNC_USERNAME}@${SVNC_FRONTEND_HOST}" "${command}"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
172
173
174
175
}

parse_options () {
    local option
176
    while getopts 'A:C:n:q:s:t:u:v:x:' option; do
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
177
178
179
        case ${option} in
            A) SVNC_ACCTCODE="$OPTARG"
               ;;
180
181
            C) SVNC_CONSTRAINT="$OPTARG"
               ;;
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
182
183
            n) SVNC_NODES="$OPTARG"
               ;;
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
184
185
            q) SVNC_QOS="$OPTARG"
               ;;
186
187
            s) SVNC_SERVER_OPTIONS="$OPTARG"
               ;;
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
188
189
190
191
            t) SVNC_RUNTIME="$OPTARG"
               ;;
            u) SVNC_USERNAME="$OPTARG"
               ;;
192
            v) SVNC_CLIENT="$OPTARG"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
193
               ;;
194
            x) SVNC_CLIENT_OPTIONS="$OPTARG"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
195
196
               ;;
            ?) usage
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
197
            exit 1
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
198
199
200
201
202
            ;;
        esac
    done

    readonly SVNC_ACCTCODE
203
    readonly SVNC_CONSTRAINT
204
    readonly SVNC_SERVER_OPTIONS
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
205
206
207
    readonly SVNC_QOS
    readonly SVNC_RUNTIME
    readonly SVNC_USERNAME
208
209
    readonly SVNC_CLIENT
    readonly SVNC_CLIENT_OPTIONS
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
210

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
211
212
    if [[ -z ${SVNC_ACCTCODE} ]]; then
        printf "ERROR: Please specify an acctcode.\n\n" >&2
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
213
        usage
Mathis Rosenhauer's avatar
set -u    
Mathis Rosenhauer committed
214
        exit 1
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
215
216
    fi
}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
217

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
218
prepare_vncserver () {
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
219
220
    ssh_frontend "mkdir -p .startvnc"
    if ! ssh_frontend "test -s .vnc/passwd"; then
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
221
        echo "No VNC password found. Please set now."
222
        echo "Do not use your LDAP password. Eight characters maximum."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
223
224
225
        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
226
227
228
229
    fi
}

submit_vnc_job () {
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
230
    local sbatch_resources
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
231
    if [[ ${SVNC_NODES} = "half" ]]; then
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
232
        sbatch_resources="#SBATCH --ntasks=24"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
233
    else
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
234
        sbatch_resources="#SBATCH --nodes=${SVNC_NODES}"
235
        sbatch_resources+=$'\n#SBATCH --exclusive'
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
236
    fi
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
237
    ssh_frontend "cd .startvnc && sbatch" <<EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
238
#!/bin/bash -l
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
239
#SBATCH --job-name=vncserver
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
240
#SBATCH --partition=gpu
241
#SBATCH --constraint=${SVNC_CONSTRAINT}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
242
${sbatch_resources}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
243
244
245
#SBATCH --qos=${SVNC_QOS}
#SBATCH --time=${SVNC_RUNTIME}
#SBATCH --account=${SVNC_ACCTCODE}
246
#SBATCH --output=out.%j
247
#SBATCH --parsable
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
248
#SBATCH --dkrzepilog=0
249
cd \${HOME}
250
/opt/TurboVNC/bin/vncserver -fg -localhost ${SVNC_SERVER_OPTIONS}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
251
EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
252
}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
253

Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
254
255
get_vnc_host_and_display () {
    local job_id="$1"
256
    ssh_frontend "/bin/bash -s" <<EOF
257
logfile=\${HOME}/.startvnc/out.${job_id}
258
while [[ -z \${host_and_display} ]]; do
259
    sleep 1
260
    if [[ -f \${logfile} ]]; then
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
261
262
        host_and_display="\$(grep -Po "started on display \Kmg[0-9]+:[0-9]+" \
                          \${logfile})"
263
    fi
264
    printf "." >&2
265
done
266
printf "\n" >&2
267
echo \${host_and_display}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
268
EOF
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
269
270
271
272
273
}

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

275
    mkdir -p "${HOME}/.ssh"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
276
    ssh_socket_dir="$(mktemp -d "${HOME}/.ssh/socket.XXXXX")"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
277
    ssh_frontend "" "-MNf"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
278
    prepare_vncserver
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
279
280

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

283
    printf "Waiting for job ${job_id} to start" >&2
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
284
285
    local host_and_display
    host_and_display="$(get_vnc_host_and_display "$job_id")"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
286
287
    vnc_host=${host_and_display%:*}
    vnc_display=${host_and_display#*:}
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
288
289
    echo "Vncserver started on node ${vnc_host}.dkrz.de display \
:${vnc_display}."
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
290

291
292
293
294
295
296
297
298
299
300
301
302
    local vnc_port_local=$((5900 + RANDOM % 100))
    # WSL doesn't seem to detect used ports so randomize
    # local port to reduce risk of masking.
    local vnc_port_remote=$(( 5900 + vnc_display ))
    while ! ssh -o ForwardX11=no \
            -o StrictHostKeyChecking=accept-new \
            -L "${vnc_port_local}:localhost:${vnc_port_remote}" \
            -Nf \
            "${SVNC_USERNAME}@${vnc_host}.dkrz.de"; do
        vnc_port_local=$((5900 + RANDOM % 100))
        echo "Trying local port ${vnc_port_local}."
    done
303
304
    local client_options
    client_options="${SVNC_CLIENT_OPTIONS:-}"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
305
    if [[ ${SVNC_PASSWORD} = "true" ]]; then
306
        echo "Fetching password from frontend."
307
        ssh_frontend "cat .vnc/passwd" > vnc_passwd
308
        client_options+=" -passwd vnc_passwd"
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
309
    fi
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
310
    echo "Connecting vncviewer to ${vnc_host}.dkrz.de"
311
    "${SVNC_CLIENT}" ${client_options} :$(( vnc_port_local - 5900 ))
Mathis Rosenhauer's avatar
Mathis Rosenhauer committed
312
}
313

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