-
Notifications
You must be signed in to change notification settings - Fork 8
/
SlurmEasy
executable file
·192 lines (155 loc) · 7.99 KB
/
SlurmEasy
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
#!/package/python/3.4.3/bin/python3
__description__="""
SlurmEasy-0.1.1 - A wrapper around sbatch for submitting jobs to slurm.
I've attempted as much as possible to copy over the functionality and options from SGEeasy.
TODO: dependency, start delay/time, time limit
Devon Ryan
"""
import argparse
import os
import datetime
import subprocess
import sys
now = datetime.datetime.now()
NTHREADS = 6 # -c or -S?
QUEUE = "bioinfo" # -p
EMAIL_TYPE = "ALL" # --mail-type=...
WORKING_DIR = os.getcwd() # -D
NAME = os.getlogin() # -J
MEM = "5781" # --mem-per-cpu=... 3541 5781
CMD_TO_LOG = False
QUEUES = ['bioinfo', 'defqueue', 'rapidus','mono']
ENV_STRING = "export LC_ALL=en_US.UTF-8; export LANG=en_US.UTF-8; export LANGUAGE=en_US.UTF-8;"
MAINTEMP = "/data/extended/"
# Could set a custom stdout and stderr with -o and -e
def parse_args():
parser = argparse.ArgumentParser(prog="SlurmEasy")
#These are optional and most have defaults
parser.add_argument("-t", "--threads", dest="threads", metavar="INT", help="maximum number of threads (default: '{}')".format(NTHREADS), type=int, default=NTHREADS)
parser.add_argument("-l", "--log", dest="log_dir", metavar="logdir", help="directory for writing stdout and stderr log files (default: '{}')".format(WORKING_DIR), type=str, default=WORKING_DIR)
parser.add_argument("--no_log", dest="no_log", action="store_true", default=False, help="Do not write log files.")
parser.add_argument("-D", "--working_dir", dest="working_dir", metavar="workingdir", help="The working directory, in case you're not already in it and your commands use relative paths (default: '{}')".format(WORKING_DIR), type=str, default=WORKING_DIR)
parser.add_argument("--tempdir", dest="tempdir", metavar="tempdir", help="User temp directory (default: '{}/tmp.XXXXXXXXXX')".format(MAINTEMP), type=str, default=MAINTEMP)
parser.add_argument("-d", "--dependency", dest="dependency", metavar="INT", help="A job number to use as a dependency. Slurm will then wait until it's completed before starting this one. N.B., an integer is only labeled as being optional here to make integration with snakemake simpler.", type=int, nargs='?', const=False)
parser.add_argument("-c", "--cmd_to_log", dest="cmd_to_log", action="store_true", help="echo bash commands to stdout (*.o) LOG file (default: '{}')".format(CMD_TO_LOG), default=CMD_TO_LOG)
parser.add_argument("-n", "--name", dest="name", help="job name, which will also affect the log file names (default: '{}')".format(NAME), type=str, default=NAME)
parser.add_argument("-m", "--mem-per-cpu", dest="mem", metavar="MEMORY", help="maximum memory per core in MB if not otherwise specified (e.g. '6G', default: '{}')".format(MEM), type=str, default=MEM)
parser.add_argument("-M", "--email", dest="email", help="Send an email to this address once the job has finished (see also --emailType) (default: ''). Note that you can specify this without an email address and {}@ie-freiburg.mpg.de will automatically be used.".format(NAME), type=str, const=True, default=False, nargs='?')
parser.add_argument("--env_string", dest="env_string", help="Environment settings (default: '{}')".format(ENV_STRING), type=str, default=ENV_STRING)
parser.add_argument("--emailType", dest="emailType", help="send email once the job has hit this state: BEGIN, END, FAIL or ALL finished (default: '{}')".format(EMAIL_TYPE), type=str, default=EMAIL_TYPE, choices=['BEGIN', 'END', 'FAIL', 'ALL'])
parser.add_argument("-q", "--queuename", dest="partition", help="The queue/partition to use (default: '{}')".format(QUEUE), type=str, default=QUEUE, choices=QUEUES)
parser.add_argument("-k", "--keepscript", dest="keep_script", action="store_true", default=False, help="keep auto-generated job script (*.slurm.sh)")
parser.add_argument("-w", "--nodelist", default=None, help="If specified, a comma separated list of nodes that can be used, e.g. 'deep14,deep15'")
parser.add_argument("-x", "--exclude", default=None, help="If specified, a comma separated list of nodes that can not be used. For example: deep9,deep10")
parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", default=False, help="verbose output")
#Required arguments
parser.add_argument('cmd', metavar="COMMANDS", help="Commands (in quotes!) to be run. This can also be a script.")
args = parser.parse_args()
return args
def main():
args = parse_args()
#Generate a shell script, line by line
script = ["#!/bin/bash", "#SBATCH --ntasks-per-node=1"]
# -w
script.append("#SBATCH -D {}".format(args.working_dir))
# -t
if args.threads > 1:
script.append("#SBATCH -c {}".format(args.threads))
# -n
if args.name:
script.append("#SBATCH -J {}".format(args.name))
# -w
if args.nodelist:
script.append("#SBATCH --nodelist={}".format(args.nodelist))
# -x
if args.exclude:
script.append("#SBATCH --exclude={}".format(args.exclude))
# -m
if args.mem:
if args.partition == 'mono' and args.mem == "5781":
# The mono queue only uses the newer nodes, so have it default to a lower per-core memory allocation so more jobs will fit
args.mem = "3541"
if args.threads > 64 and args.mem == "5781":
args.mem = "3541"
script.append("#SBATCH --mem-per-cpu={}".format(args.mem))
# -M and --emailType
if args.email is not False:
if args.email is True:
script.append("#SBATCH --mail-user={}@ie-freiburg.mpg.de".format(NAME))
else:
script.append("#SBATCH --mail-user={}".format(args.email))
script.append("#SBATCH --mail-type={}".format(args.emailType))
# -q
script.append("#SBATCH -p {}".format(args.partition))
# -l
if args.no_log:
script.append("#SBATCH -o /dev/null")
script.append("#SBATCH -e /dev/null")
else:
script.append("#SBATCH -o {}/{}.{}.%j.out".format(args.log_dir, args.name, now.strftime('%Y%m%d_%H%M%S')))
script.append("#SBATCH -e {}/{}.{}.%j.err".format(args.log_dir, args.name, now.strftime('%Y%m%d_%H%M%S')))
try:
os.makedirs(args.log_dir)
except OSError as exc:
# This will only throw an error if the user lacks permissions or there's a similar error
if type(exc) == FileExistsError:
pass
else:
raise
# -d
if args.dependency:
script.append("#SBATCH -d {}".format(args.dependency))
# Environment variables
if MAINTEMP == args.tempdir:
script.append("""
TMPDIR=%s/
scratch=$(mktemp -p $TMPDIR -d -t tmp.XXXXXXXXXX)""" % MAINTEMP)
else:
script.append("""
TMPDIR=%s/
mkdir $TMPDIR
scratch=$TMPDIR""" % args.tempdir)
script.append("""
export TMPDIR=$scratch/
export TMP=$scratch/
export TEMP=$scratch/
%s
function cleanup {
cd /
rm -rf $scratch
}
trap cleanup SIGTERM SIGKILL SIGUSR2""" % args.env_string)
# -c echo the command
if args.cmd_to_log:
script.append("echo \"{}\"".format(args.cmd))
# Add the command
script.append("srun {}".format(args.cmd))
script.append("cleanup")
# -v
if args.verbose:
print("\n".join(script))
ofname="{}/{}.{}.slurm.sh".format(args.log_dir, args.name, now.strftime('%Y%m%d_%H%M%S'))
outfile = open(ofname, "w")
for line in script:
outfile.write("{}\n".format(line))
outfile.close()
cmd2 = ["sbatch", ofname]
sp = subprocess.Popen(cmd2, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = sp.communicate()
# This is so snakemake can figure out dependencies
stdout = stdout.decode()
stderr = stderr.decode()
if stdout.startswith("Submitted batch"):
stdout = stdout[20:]
# In python3, this will print b'' if empty!
if len(stderr):
sys.stderr.write(stderr)
if len(stdout):
sys.stdout.write(stdout)
if not args.keep_script:
os.remove(ofname)
if args.verbose:
print("Removing sbatch script ({})".format(ofname))
return(sp.returncode)
if __name__ == "__main__":
main()