/
template.rb
316 lines (268 loc) · 11.2 KB
/
template.rb
1
2
3
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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
require "ood_core/refinements/hash_extensions"
module OodCore
module BatchConnect
# A template class that renders a batch script designed to facilitate
# external connections to the running job
class Template
using Refinements::HashExtensions
using Refinements::ArrayExtensions
# The context used to render this template
# @return [Hash] context hash
attr_reader :context
# @param context [#to_h] the context used to render the template
# @option context [#to_s] :work_dir Working directory for batch script
# @option context [#to_s] :conn_file ("connection.yml") The file that
# holds connection information
# @option context [#to_sym, Array<#to_sym>] :conn_params ([]) A list of
# connection parameters added to the connection file (`:host`, `:port`,
# and `:password` will always exist)
# @option context [#to_s] :bash_helpers ("...") Bash helper methods
# @option context [#to_i] :min_port (2000) Minimum port used when looking
# for available port
# @option context [#to_i] :max_port (65535) Maximum port used when
# looking for available port
# @option context [#to_i] :password_size (32) Length of randomly generated
# password
# @option context [#to_s] :header ("") Shell code prepended at the top of
# the script body
# @option context [#to_s] :footer ("") Shell code appended at the bottom
# of the script body
# @option context [#to_s] :script_wrapper ("%s") Bash code that wraps
# around the body of the template script (use `%s` to interpolate the
# body)
# @option context [#to_s] :set_host ("host=$(hostname)") Bash code used
# to set the `host` environment variable used for connection
# information
# @option context [#to_s] :before_script ("...") Bash code run before the
# main script is forked off
# @option context [#to_s] :before_file ("before.sh") Path to script that
# is sourced before main script is forked (assumes you don't modify
# `:before_script`)
# @option context [#to_s] :run_script ("...") Bash code that is forked
# off and treated as the main script
# @option context [#to_s] :script_file ("./script.sh") Path to script
# that is forked as the main scripta (assumes you don't modify
# `:run_script`)
# @option context [#to_s] :timeout ("") Timeout the main script in
# seconds, if empty then let script run for full walltime (assumes you
# don't modify `:run_script`)
# @option context [#to_s] :clean_script ("...") Bash code run during
# clean up after job finishes
# @option context [#to_s] :clean_file ("clean.sh") Path to script that is
# sourced during clean up (assumes you don't modify `:clean_script`)
def initialize(context = {})
@context = context.to_h.compact.symbolize_keys
raise ArgumentError, "No work_dir specified. Missing argument: work_dir" unless context.include?(:work_dir)
end
# Render this template as string
# @return [String] rendered template
def to_s
<<-EOT.gsub(/^ {10}/, '')
#{header}
#{script_wrapper}
#{footer}
EOT
end
protected
def password_size
context.fetch(:password_size, 8).to_i
end
def min_port
context.fetch(:min_port, 2000).to_i
end
def max_port
context.fetch(:max_port, 65535).to_i
end
private
# Working directory that batch script runs in
def work_dir
context.fetch(:work_dir).to_s
end
# The file that holds the connection information in yaml format
def conn_file
context.fetch(:conn_file, "connection.yml").to_s
end
# The parameters to include in the connection file
def conn_params
conn_params = Array.wrap(context.fetch(:conn_params, [])).map(&:to_sym)
(conn_params + [:host, :port, :password]).uniq
end
# Bash script used to define the `host` environment variable
def set_host
context.fetch(:set_host, "host=$(hostname)").to_s
end
# Helper methods used in the bash scripts
def bash_helpers
context.fetch(:bash_helpers) do
<<-EOT.gsub(/^ {14}/, '')
# Source in all the helper functions
source_helpers () {
# Generate random integer in range [$1..$2]
random_number () {
shuf -i ${1}-${2} -n 1
}
export -f random_number
port_used_python() {
python -c "import socket; socket.socket().connect(('$1',$2))" >/dev/null 2>&1
}
port_used_python3() {
python3 -c "import socket; socket.socket().connect(('$1',$2))" >/dev/null 2>&1
}
port_used_nc(){
nc -w 2 "$1" "$2" < /dev/null > /dev/null 2>&1
}
port_used_lsof(){
lsof -i :"$2" >/dev/null 2>&1
}
port_used_bash(){
local bash_supported=$(strings /bin/bash 2>/dev/null | grep tcp)
if [ "$bash_supported" == "/dev/tcp/*/*" ]; then
(: < /dev/tcp/$1/$2) >/dev/null 2>&1
else
return 127
fi
}
# Check if port $1 is in use
port_used () {
local port="${1#*:}"
local host=$((expr "${1}" : '\\(.*\\):' || echo "localhost") | awk 'END{print $NF}')
local port_strategies=(port_used_nc port_used_lsof port_used_bash port_used_python port_used_python3)
for strategy in ${port_strategies[@]};
do
$strategy $host $port
status=$?
if [[ "$status" == "0" ]] || [[ "$status" == "1" ]]; then
return $status
fi
done
return 127
}
export -f port_used
# Find available port in range [$2..$3] for host $1
# Default: [#{min_port}..#{max_port}]
find_port () {
local host="${1:-localhost}"
local port=$(random_number "${2:-#{min_port}}" "${3:-#{max_port}}")
while port_used "${host}:${port}"; do
port=$(random_number "${2:-#{min_port}}" "${3:-#{max_port}}")
done
echo "${port}"
}
export -f find_port
# Wait $2 seconds until port $1 is in use
# Default: wait 30 seconds
wait_until_port_used () {
local port="${1}"
local time="${2:-30}"
for ((i=1; i<=time*2; i++)); do
port_used "${port}"
port_status=$?
if [ "$port_status" == "0" ]; then
return 0
elif [ "$port_status" == "127" ]; then
echo "commands to find port were either not found or inaccessible."
echo "command options are lsof, nc, bash's /dev/tcp, or python (or python3) with socket lib."
return 127
fi
sleep 0.5
done
return 1
}
export -f wait_until_port_used
# Generate random alphanumeric password with $1 (default: #{password_size}) characters
create_passwd () {
tr -cd 'a-zA-Z0-9' < /dev/urandom 2> /dev/null | head -c${1:-#{password_size}}
}
export -f create_passwd
}
export -f source_helpers
EOT
end.to_s
end
# Shell code that is prepended at the top of the script body
def header
context.fetch(:header, "").to_s
end
# Shell code that is appended at the bottom of the script body
def footer
context.fetch(:footer, "").to_s
end
# Bash code that wraps around the body of the template script (use `%s`
# to interpolate the body)
def script_wrapper
context.fetch(:script_wrapper, "%s").to_s % base_script
end
# Source in a developer defined script before running the main script
def before_script
context.fetch(:before_script) do
before_file = context.fetch(:before_file, "before.sh").to_s
"[[ -e \"#{before_file}\" ]] && source \"#{before_file}\""
end.to_s
end
# Fork off a developer defined main script and possibly time it out after
# a period of time
def run_script
context.fetch(:run_script) do
script_file = context.fetch(:script_file, "./script.sh").to_s
timeout = context.fetch(:timeout, "").to_s
timeout.empty? ? "\"#{script_file}\"" : "timeout #{timeout} \"#{script_file}\""
end.to_s
end
# Source in a developer defined script after running the main script
def after_script
context.fetch(:after_script) do
after_file = context.fetch(:after_file, "after.sh").to_s
"[[ -e \"#{after_file}\" ]] && source \"#{after_file}\""
end.to_s
end
# Source in a developer defined clean up script that is run during the
# clean up stage
def clean_script
context.fetch(:clean_script) do
clean_file = context.fetch(:clean_file, "clean.sh").to_s
"[[ -e \"#{clean_file}\" ]] && source \"#{clean_file}\""
end.to_s
end
# The base script template
def base_script
<<-EOT.gsub(/^ {12}/, '')
cd #{work_dir}
# Export useful connection variables
export host
export port
# Generate a connection yaml file with given parameters
create_yml () {
echo "Generating connection YAML file..."
(
umask 077
echo -e "#{conn_params.map { |p| "#{p}: $#{p}" }.join('\n')}" > "#{conn_file}"
)
}
# Cleanliness is next to Godliness
clean_up () {
echo "Cleaning up..."
#{clean_script.gsub(/\n(?=[^\s])/, "\n ")}
[[ ${SCRIPT_PID} ]] && pkill -P ${SCRIPT_PID} || :
pkill -P $$
exit ${1:-0}
}
#{bash_helpers}
source_helpers
# Set host of current machine
#{set_host}
#{before_script}
echo "Script starting..."
#{run_script} &
SCRIPT_PID=$!
#{after_script}
# Create the connection yaml file
create_yml
# Wait for script process to finish
wait ${SCRIPT_PID} || clean_up 1
# Exit cleanly
clean_up
EOT
end
end
end
end