-
Notifications
You must be signed in to change notification settings - Fork 0
/
zizibee.py
516 lines (458 loc) · 16.1 KB
/
zizibee.py
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
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
#!/usr/bin/env python3
import inspect
import os
import subprocess
import json
import re
import sys
import paramiko
import h5py
import numpy as np
import time
HOSTNAME = 'sshcampus.ccv.brown.edu'
def get_cell_content(notebook_path, cell_index):
'''
This function can be used to retrieve a specific cell
from a Jupyter notebook.
Parameters
----------
notebook_path (str): path to the notebook
cell_index (int): index of the cell to be retrieved
Returns
-------
cell_content (str): the content of the cell
'''
with open(notebook_path, 'r') as f:
data = json.load(f)
if cell_index < len(data['cells']):
cell = data['cells'][cell_index]
if cell['cell_type'] == 'code':
return ''.join(cell['source'])
else:
return "Selected cell is not a code cell"
else:
return "Cell index is out of range"
def get_all_fun(the_globals):
'''
This function receives a dictionary with the global
variables and returns a list with the definitions of all the
functions defined in the given namespace.
Parameters
----------
the_globals (dict): the dictionary with the global
variables, usually this is is simply the output of globals()
but it needs to be given explicitly because otherwise the
function would be using its own globals().
Returns
-------
func_defs (list): a list with the definitions the functions
defined in the_globals.
'''
avoid_funcs = ['open', 'get_all_fun']
globals_copy = the_globals.copy()
func_defs = []
for name, obj in globals_copy.items():
if (inspect.isfunction(obj)
and (name not in avoid_funcs)
and ('_' != name[0])):
try:
func_def = inspect.getsource(obj)
func_defs.append(func_def)
except TypeError:
print(f"Can't get source code for built-in function {name}.")
return func_defs
def execute_at_ccv(ccv_cmd, username='jlizaraz'):
'''
Execute a command at CCV and return its output. The command
can be multiline and include many statements. This is
accomplished by using the ssh command to execute the command
at CCV. This is done through a login mode (although it could
be modified) so in its present form it should't be used for
resource-intensive tasks.
Parameters
----------
ccv_cmd (str): a command (or many) to execute at CCV
Returns
-------
stdo (str): the standard output of the sent command.
Example
-------
>> long_cmd = \'\'\'
me="David"
for i in 1 2 3
do
echo "Hello $me."
echo "$(date)"
done\'\'\'
>> oh = execute_at_ccv(long_cmd)
>> print(oh)
>>> Hello David.
Tue Jul 25 17:25:08 EDT 2023
Hello David.
Tue Jul 25 17:25:08 EDT 2023
Hello David.
Tue Jul 25 17:25:08 EDT 2023
'''
# cmd = "ssh jlizaraz@compute_node '%s'" % (HOSTNAME, ccv_cmd)
cmd = "ssh %s@%s '%s'" % (username, HOSTNAME, ccv_cmd)
process = subprocess.Popen(cmd, shell=True,
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate()
if process.returncode != 0:
# An error occurred
raise Exception(stderr.decode())
else:
# Return command output
stdo = stdout.decode()
return stdo
def upload_to_ccv(filename, folder, username='jlizaraz', transferhost = HOSTNAME, verbose=False):
'''
This function can be used to upload local files to a folder
at CCV.
Parameters
----------
filename (str): path of the file to be uploaded, can be a
full or relative path.
folder (str): file will be uploaded to this folder
username (str): username at CCV
transferhost(str): the hostname of the transfer node at CCV
verbose (bool): whether to print out the command executed to
upload the file
Returns
-------
None
'''
root_name = os.path.split(filename)[-1]
command = """
echo put \"\\\"%s\\\"\" \\\"\"%s/%s\\\"\" | sftp %s@%s
""" % (filename, folder, root_name, username, transferhost)
if verbose:
print(command)
process = subprocess.Popen(command,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
# Wait for the command to complete
stdout, stderr = process.communicate()
if verbose:
# Print the output or the error message, if any
if process.returncode == 0:
print(f'Success:\n{stdout.decode()}')
else:
print(f'Error:\n{stderr.decode()}')
return None
def execute_shell_command(ssh_shell, command):
'''
Convenience function to execute shell commands at CCV.
Parameters
----------
ssh_shell (paramiko.Channel): a shell channel to CCV
command (str): the command as one would type it at the shell
Returns
-------
output (str): the output of the command, as one would see it
at the shell
'''
ssh_shell.send(command + '\n')
output = ''
while True:
while not ssh_shell.recv_ready():
pass
data = ssh_shell.recv(4096).decode()
output += data
if output.endswith('$ '):
break
return output
def run_at_ccv(job_config, verbose=False, closeSSH=False):
'''
Send a grid job to CCV.
Parameters
----------
job_config (dict): with at least the
following keys:
> username (str): the username at CCV.
> numCores (int): how many cores
required for each job.
> numJobs (int): how many evaluation
points will be run.
> memInGB (int): how much memory
required for each job.
> import_block (str): the import block
of the script to be run.
> extra_py (list): list of paths to
extra python files to be uploaded.
> theglobals (dict): the dictionary with
the global variables.
> fun_name (str): the name of the
function to be run at CCV.
> job_name (str): the name of the job.
verbose (bool): if True some debug mesages
are printed
closeSSH (bool): if True then the SSH
connection and sheel are closed at the end
Returns
-------
job_config (dict) with the following
additional keys
> data_dir_at_CCV (str): the path to the
data folder at CCV
> scratch_dir_at_CCV (str): the path to
the scratch folder at CCV
> data_dir_at_mac (str): the path to the
data folder at the mac
> scratch_dir_at_mac (str): the path to
the scratch folder at the mac
> sbatch (str): the text of the sbatch
script
> ccv_sbatch_cmds (list): with the
commands that were executed to run the
batch job
> script_text (str): the text of the
uploaded script
> (theglobals is deleted from the
job_config dictionary)
'''
numCores = job_config['numCores']
theglobals = job_config['theglobals']
del job_config['theglobals']
numJobs = job_config['numJobs']
memInGB = job_config['memInGB']
username = job_config['username']
job_name = job_config['job_name']
importblock = job_config['import_block']
extra_py = job_config['extra_py']
special_func = job_config['fun_name']
funs = get_all_fun(theglobals)
fire_bit = '''
def main():
fire.Fire(%s)
if __name__ == '__main__':
main()''' % special_func
data_dir = '/users/%s/data/%s/%s' % (username, username, job_name)
scratch_dir = '/users/%s/scratch/%s' % (username, job_name)
job_config['data_dir_at_CCV'] = data_dir
job_config['scratch_dir_at_CCV'] = scratch_dir
home_dir = os.path.expanduser('~')
data_dir_at_mac = os.path.join(home_dir, 'ccv', 'data', job_name)
scratch_dir_at_mac = os.path.join(home_dir, 'ccv', 'scratch', job_name)
job_config['data_dir_at_mac'] = data_dir_at_mac
job_config['scratch_dir_at_mac'] = scratch_dir_at_mac
zzbar_dict = {'data_dir':data_dir, 'scratch_dir':scratch_dir}
zzbars = []
for k, v in zzbar_dict.items():
if isinstance(v, str):
zzbars.append('%s = \'%s\'' % (k,v))
else:
zzbars.append('%s = %s' % (k,v))
zzvars = '\n'.join(zzbars)
pieces = [importblock] + [zzvars] + funs + [fire_bit]
script_text = '\n\n'.join(pieces)
job_config['script_text'] = script_text
print("Establishing an SSH connection to CCV and launching a shell ...")
ccv_ssh = paramiko.SSHClient()
ccv_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ccv_ssh.connect(HOSTNAME, username=username)
ssh_shell = ccv_ssh.invoke_shell()
# make sure that the relevant folders are created at CCV
for dirname in [data_dir, scratch_dir]:
execute_shell_command(ssh_shell, 'mkdir %s' % dirname)
# upload the code
print("Saving Python script to file ...")
with open(job_name + '.py', 'w') as f:
f.write(script_text)
# upload the script to CCV
print("Uploading script to CCV ...")
upload_to_ccv(job_name+'.py', data_dir, verbose=verbose)
if len(extra_py) > 0:
for extrap in extra_py:
if '.py' not in extrap:
extrap = extrap + '.py'
upload_to_ccv(extrap, data_dir)
print("Composing the sbatch script ...")
sbatch = '''#!/bin/bash
#SBATCH -n {numCores}
#SBATCH --mem={memInGB}GB
#SBATCH -t 1:00:00
#SBATCH --array=0-{numJobs}
#SBATCH -o {job_name}-%a.out
#SBATCH -e {job_name}-%a.out
cd {data_dir}
~/anaconda/foundation/bin/python {data_dir}/{job_name}.py $SLURM_ARRAY_TASK_ID
'''.format(numCores=numCores,
numJobs = numJobs - 1,
data_dir = data_dir,
memInGB = memInGB,
job_name = job_name
)
job_config['sbatch'] = sbatch
sbatch_fname = '%s-batch.sh' % job_name
print("Writing sbatch script ...")
open(sbatch_fname,'w').write(sbatch)
print("Sending sbatch script ...")
upload_to_ccv(sbatch_fname, data_dir)
ccv_sbatch_cmds = ['module load anaconda/3-5.2.0',
'conda activate foundation',
'cd',
'cd %s' % data_dir,
'sbatch ' + sbatch_fname]
job_config['ccv_sbatch_cmds'] = ccv_sbatch_cmds
outputs = [execute_shell_command(ssh_shell, cmd) for cmd in ccv_sbatch_cmds]
job_config['ccv_sbatch_cmd_outputs'] = outputs
if verbose:
print('\n'.join(outputs))
if closeSSH:
ssh_shell.close()
ccv_ssh.close()
ccv_ssh = None
return job_config
def execute_command(cmd):
'''
Execute a command in the local machine and return its output.
Parameters
----------
cmd (str): a command, just like one would type on a terminal.
Returns
-------
stdout (str): the standard output of the command.
'''
process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate()
if process.returncode != 0:
# An error occurred
raise Exception(stderr.decode())
else:
# Return command output
return stdout.decode()
def myq():
'''
Returns the output of running the myq command at CCV.
Parameters
----------
None
Returns
-------
None
'''
qu = execute_at_ccv('myq')
qu = re.sub(r'\n\s*\n', '\n', qu)
return qu
def pull_from_ccv_to_mac(ccv_folder, mac_folder):
'''
All the files from ccv_folder will be synced to mac_folder.
None of the changes at the mac_folder will be reflected at
the ccv_folder.
Parameters
----------
ccv_folder (str): path to a folder at CCV
mac_foler (str): path to a folder at the mac
Returns
-------
rsync_out (str): the stdout of the rsync command
'''
if not os.path.exists(mac_folder):
print("Mac folder does not exist, creating ...")
os.mkdir(mac_folder)
if mac_folder[-1] != '/':
mac_folder += '/'
if ccv_folder[-1] != '/':
ccv_folder += '/'
rsync_cmd = 'rsync -avz [email protected]:{ccv_folder} {mac_folder}'.format(ccv_folder=ccv_folder, mac_folder=mac_folder)
rsync_out = execute_command(rsync_cmd)
return rsync_out
def load_h5_data(fnames):
'''
This function takes the filenames from a bunch of .h5 files
and creates a function that can be used to explore the
return values for the corresponding inputs.
Parameters
----------
fnames (list): list of paths to the h5 files
Returns
-------
ccv_fun (function): function that takes the input parameters
and returns the data, this function has an attribute called
keys that contains the input parameters that it is defined
for.
'''
outs = {}
for fname in fnames:
with h5py.File(fname, 'r') as f:
the_index = int(fname.split('/')[-1].split('.')[0])
the_params = tuple(np.array(f['in']))
outs[the_params] = np.array(f['out'])
out_fun = lambda *params: outs.get(tuple(params), None)
out_fun.input_params = list(outs.keys())
return out_fun
def get_ccv_values(mac_folder, ccv_folder, numJobs):
'''
This function uses rsync to pull the data from CCV to the
Mac, it does this periodically until all the expected output
files have been downloaded.
Once these files have been all retrieved, they are then
loaded into a function that can be used to retrieve the data
with the same call signature as the original function (at
least in shape).
The returned function has an attribute called input_params
that contains the input parameters for which the function
was evaluated.
If an input is given for which the function was not
evaluated, then the function returns None.
Parameters
----------
mac_folder (str): path to the folder where the data will be
downloaded
ccv_folder (str): path to the folder where the data is at
CCV
numJobs (int): how many jobs are expected
Returns
-------
out_fun (function): function that takes the input parameters
and returns the data.
'''
out_fnames = []
wait_time = 1
num_done = 0
while len(out_fnames) < numJobs:
pull_from_ccv_to_mac(ccv_folder, mac_folder)
out_fnames = os.listdir(mac_folder)
out_fnames = list(filter(lambda x: x.endswith('.h5'), out_fnames))
# escalate the waiting time if subsequent checks show no progress
if num_done == len(out_fnames):
wait_time += 2
num_done = len(out_fnames)
# print('Waiting for %d jobs to finish' % (numJobs - num_done))
progress_bar(num_done, numJobs, prefix = 'Progress:', suffix = 'Complete', length = 30)
time.sleep(1)
out_fnames = [os.path.join(mac_folder, fname) for fname in out_fnames]
out_fun = load_h5_data(out_fnames)
return out_fun
def progress_bar(iteration, total, prefix = '', suffix = '', decimals = 1, length = 50, fill = '█'):
'''
A convenient progress bar.
Parameters
----------
iteration (int): current iteration
total (int): total iterations
prefix (str): prefix string
suffix (str): suffix string
decimals (int): positive number of decimals in percent complete
length (int): character length of bar
fill (str): bar fill character
Returns
-------
None
'''
# Calculate elapsed time and estimated remaining time
global start_time
if iteration == 0:
start_time = time.time()
elapsed_time = time.time() - start_time
remaining_time = (elapsed_time * (total / (iteration + 1))) - elapsed_time if iteration > 0 else 0
percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
filled_length = int(length * iteration // total)
bar = fill * filled_length + '-' * (length - filled_length)
print(f'\r{prefix} |{bar}| {percent}% {suffix} Elapsed Time: {elapsed_time:.2f}s Remaining Time: {remaining_time:.2f}s', end = '\r')
sys.stdout.flush() # add this line
# Print New Line on Complete
if iteration == total:
print()