-
Notifications
You must be signed in to change notification settings - Fork 0
/
sbatch_submit.py
executable file
·58 lines (48 loc) · 1.92 KB
/
sbatch_submit.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
#!/usr/bin/env python3
"""
Submit a batch job to SLURM on BlueBear.
"""
import os
import subprocess
import tempfile
import argparse
def main():
# Read in the template file
lib_dir = os.path.dirname(os.path.realpath(__file__))
with open(lib_dir + '/_sbatch_template.sh') as f:
template = f.readlines()
template = ''.join(template)
# Make the arg parser
desc = 'Submit code as a Bluebear job'
parser = argparse.ArgumentParser(description=desc)
# Arg specs: Short name, long name, description, default, other args
args = [['i', 'input', 'Shell cmd to run', None, {'required': True}],
['s', 'setup', 'Shell cmd that sets up the environment', ':', {}],
['n', 'ntasks', 'Number of tasks', '1', {}],
['o', 'nodes', 'Number of nodes', '1', {}],
['c', 'cpus', 'Number of CPUs per task', '1', {}],
['t', 'time', 'Maximum time to let the job run', '5:0:0', {}],
['m', 'mem', 'Memory allocation', '10G', {}],
['d', 'dir', 'Directory to save logfiles', '.', {}],
['q', 'qos', 'Select a QOS on BlueBear', 'bbdefault',
{'choices': ['bbdefault', 'bbshort']}]]
argnames = [a[1] for a in args]
for arg in args:
parser.add_argument('-' + arg[0],
'--' + arg[1],
help=arg[2],
default=arg[3],
**arg[4])
args = parser.parse_args()
# Replace the Slurm parameters in the sbatch script
for argname in argnames:
template = template.replace('_' + argname, getattr(args, argname))
# Write a temporary sbatch script
tf = tempfile.NamedTemporaryFile(mode='w')
tf.writelines(template)
tf.flush()
# Run the job
p = subprocess.call('sbatch ' + tf.name, shell=True)
#####p = subprocess.call('cat ' + tf.name, shell=True)
if __name__ == "__main__":
main()