Skip to content

Commit f4ca8ca

Browse files
brevansbrevans
committed
minor reformatting
1 parent c175b74 commit f4ca8ca

File tree

2 files changed

+217
-112
lines changed

2 files changed

+217
-112
lines changed

‎dSQ.py‎

Lines changed: 161 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from datetime import datetime
44
from os import path
55
from subprocess import call, check_output
6-
from textwrap import fill
6+
from textwrap import fill
77
import argparse
88
import itertools
99
import os
@@ -12,12 +12,14 @@
1212

1313
__version__ = 1.04
1414

15+
1516
def safe_fill(text, wrap_width):
1617
if sys.__stdin__.isatty():
1718
return fill(text, wrap_width)
1819
else:
1920
return text
2021

22+
2123
# Check if dSQ is being run interactively
2224
if sys.__stdin__.isatty():
2325
# get terminal columns for wrapping
@@ -28,12 +30,16 @@ def safe_fill(text, wrap_width):
2830
else:
2931
term_columns = 25
3032

31-
#get slurm info
33+
# get slurm info
3234
try:
33-
#get max configured array index
34-
slurm_conf = check_output(["scontrol", "show", "conf"], universal_newlines=True).split("\n")[:-1]
35-
max_array_size = [int(x.split("=")[1]) for x in slurm_conf if x.startswith("MaxArraySize")][0]
36-
35+
# get max configured array index
36+
slurm_conf = check_output(
37+
["scontrol", "show", "conf"], universal_newlines=True
38+
).split("\n")[:-1]
39+
max_array_size = [
40+
int(x.split("=")[1]) for x in slurm_conf if x.startswith("MaxArraySize")
41+
][0]
42+
3743
except FileNotFoundError as e:
3844
print("You don't appear to have slurm available. Exiting!")
3945
sys.exit(1)
@@ -55,91 +61,125 @@ def safe_fill(text, wrap_width):
5561
-c, --cpus-per-task=ncpus Number of cpu cores required per job.
5662
--mem-per-cpu=MiB Amount of memory per allocated cpu core per job.
5763
58-
""".format(__version__)
59-
desc = "\n".join([safe_fill(x, term_columns-1) for x in str.splitlines(desc)])
64+
""".format(
65+
__version__
66+
)
67+
desc = "\n".join([safe_fill(x, term_columns - 1) for x in str.splitlines(desc)])
6068

6169
# helper functions for array range formatting
6270
# collapse job numbers in job file to ranges
6371
def _collapse_ranges(jobnums):
6472
# takes a list of numbers, returns tuples of numbers that specify representative ranges
6573
# inclusive
66-
for i, t in itertools.groupby(enumerate(jobnums), lambda tx: tx[1]-tx[0]):
74+
for i, t in itertools.groupby(enumerate(jobnums), lambda tx: tx[1] - tx[0]):
6775
t = list(t)
6876
yield t[0][1], t[-1][1]
6977

7078

7179
# format job ranges
7280
def format_range(jobnums):
7381
ranges = list(_collapse_ranges(jobnums))
74-
return ",".join(["{}-{}".format(x[0],x[1]) if x[0]!=x[1] else str(x[0]) for x in ranges])
82+
return ",".join(
83+
["{}-{}".format(x[0], x[1]) if x[0] != x[1] else str(x[0]) for x in ranges]
84+
)
85+
7586

7687
# argument parsing
77-
parser = argparse.ArgumentParser(description=desc,
78-
add_help=False,
79-
usage="%(prog)s --job-file jobfile [dSQ args] [slurm args]",
80-
formatter_class=argparse.RawTextHelpFormatter,
81-
prog=path.basename(sys.argv[0]))
88+
parser = argparse.ArgumentParser(
89+
description=desc,
90+
add_help=False,
91+
usage="%(prog)s --job-file jobfile [dSQ args] [slurm args]",
92+
formatter_class=argparse.RawTextHelpFormatter,
93+
prog=path.basename(sys.argv[0]),
94+
)
8295

8396
tmp = parser.add_argument_group("Required Arguments")
8497
required_dsq = tmp.add_mutually_exclusive_group(required=True)
85-
required_dsq.add_argument("--job-file",
86-
metavar="jobs.txt",
87-
nargs=1,
88-
type=argparse.FileType("r"),
89-
help="Job file, one self-contained job per line.")
98+
required_dsq.add_argument(
99+
"--job-file",
100+
metavar="jobs.txt",
101+
nargs=1,
102+
type=argparse.FileType("r"),
103+
help="Job file, one self-contained job per line.",
104+
)
90105
# add old option names, but keep them hidden
91-
required_dsq.add_argument("--taskfile",
92-
nargs=1,
93-
dest="job_file",
94-
type=argparse.FileType("r"),
95-
help=argparse.SUPPRESS)
96-
required_dsq.add_argument("--jobfile",
97-
nargs=1,
98-
dest="job_file",
99-
type=argparse.FileType("r"),
100-
help=argparse.SUPPRESS)
106+
required_dsq.add_argument(
107+
"--taskfile",
108+
nargs=1,
109+
dest="job_file",
110+
type=argparse.FileType("r"),
111+
help=argparse.SUPPRESS,
112+
)
113+
required_dsq.add_argument(
114+
"--jobfile",
115+
nargs=1,
116+
dest="job_file",
117+
type=argparse.FileType("r"),
118+
help=argparse.SUPPRESS,
119+
)
101120
# optional arguments
102121
optional_dsq = parser.add_argument_group("Optional Arguments")
103-
optional_dsq.add_argument("-h","--help",
104-
action="help",
105-
default=argparse.SUPPRESS,
106-
help="Show this help message and exit.")
107-
optional_dsq.add_argument("--version",
108-
action="version",
109-
version="%(prog)s {}".format(__version__))
110-
optional_dsq.add_argument("--batch-file",
111-
metavar="sub_script.sh",
112-
nargs=1,
113-
help=safe_fill("Name for batch script file. Defaults to dsq-jobfile-YYYY-MM-DD.sh", term_columns-24))
114-
optional_dsq.add_argument("-J", "--job-name",
115-
metavar="jobname",
116-
nargs=1,
117-
help="Name of your job array. Defaults to dsq-jobfile")
118-
optional_dsq.add_argument("--max-jobs",
119-
metavar="number",
120-
nargs=1,
121-
help="Maximum number of simultaneously running jobs from the job array.")
122-
optional_dsq.add_argument("-o", "--output",
123-
nargs=1,
124-
metavar="fmt_string",
125-
help=safe_fill("Slurm output file pattern. There will be one file per line in your job file. To suppress slurm out files, set this to /dev/null. Defaults to dsq-jobfile-%%A_%%a-%%N.out", term_columns-24))
126-
optional_dsq.add_argument("--status-dir",
127-
metavar="dir",
128-
nargs=1,
129-
help="Directory to save the job_jobid_status.tsv file to. Defaults to working directory.")
130-
optional_dsq.add_argument("--suppress-stats-file",
131-
action="store_true",
132-
help="Don't save job stats to job_jobid_status.tsv")
133-
optional_dsq.add_argument("--stdout",
134-
action="store_true",
135-
help=argparse.SUPPRESS)
136-
optional_dsq.add_argument("--submit",
137-
action="store_true",
138-
help="Submit the job array on the fly instead of creating a submission script.")
122+
optional_dsq.add_argument(
123+
"-h",
124+
"--help",
125+
action="help",
126+
default=argparse.SUPPRESS,
127+
help="Show this help message and exit.",
128+
)
129+
optional_dsq.add_argument(
130+
"--version", action="version", version="%(prog)s {}".format(__version__)
131+
)
132+
optional_dsq.add_argument(
133+
"--batch-file",
134+
metavar="sub_script.sh",
135+
nargs=1,
136+
help=safe_fill(
137+
"Name for batch script file. Defaults to dsq-jobfile-YYYY-MM-DD.sh",
138+
term_columns - 24,
139+
),
140+
)
141+
optional_dsq.add_argument(
142+
"-J",
143+
"--job-name",
144+
metavar="jobname",
145+
nargs=1,
146+
help="Name of your job array. Defaults to dsq-jobfile",
147+
)
148+
optional_dsq.add_argument(
149+
"--max-jobs",
150+
metavar="number",
151+
nargs=1,
152+
help="Maximum number of simultaneously running jobs from the job array.",
153+
)
154+
optional_dsq.add_argument(
155+
"-o",
156+
"--output",
157+
nargs=1,
158+
metavar="fmt_string",
159+
help=safe_fill(
160+
"Slurm output file pattern. There will be one file per line in your job file. To suppress slurm out files, set this to /dev/null. Defaults to dsq-jobfile-%%A_%%a-%%N.out",
161+
term_columns - 24,
162+
),
163+
)
164+
optional_dsq.add_argument(
165+
"--status-dir",
166+
metavar="dir",
167+
nargs=1,
168+
help="Directory to save the job_jobid_status.tsv file to. Defaults to working directory.",
169+
)
170+
optional_dsq.add_argument(
171+
"--suppress-stats-file",
172+
action="store_true",
173+
help="Don't save job stats to job_jobid_status.tsv",
174+
)
175+
optional_dsq.add_argument("--stdout", action="store_true", help=argparse.SUPPRESS)
176+
optional_dsq.add_argument(
177+
"--submit",
178+
action="store_true",
179+
help="Submit the job array on the fly instead of creating a submission script.",
180+
)
139181
# silently allow overriding --array, otherwise we calculate that
140-
optional_dsq.add_argument("-a", "--array",
141-
nargs=1,
142-
help=argparse.SUPPRESS)
182+
optional_dsq.add_argument("-a", "--array", nargs=1, help=argparse.SUPPRESS)
143183

144184
args, user_slurm_args = parser.parse_known_args()
145185

@@ -149,7 +189,9 @@ def format_range(jobnums):
149189
job_info["max_jobs"] = args.max_jobs
150190
job_info["num_jobs"] = 0
151191
job_info["job_id_list"] = []
152-
job_info["run_script"] = path.join(path.dirname(path.abspath(sys.argv[0])), "dSQBatch.py")
192+
job_info["run_script"] = path.join(
193+
path.dirname(path.abspath(sys.argv[0])), "dSQBatch.py"
194+
)
153195
job_info["job_file_name"] = path.abspath(args.job_file[0].name)
154196
job_info["job_file_arg"] = "--job-file {}".format(job_info["job_file_name"])
155197
job_info["slurm_args"] = {}
@@ -167,13 +209,20 @@ def format_range(jobnums):
167209
for i, line in enumerate(args.job_file[0]):
168210
if not (line.startswith("#") or line.rstrip() == ""):
169211
job_info["job_id_list"].append(i)
170-
job_info["num_jobs"]+=1
212+
job_info["num_jobs"] += 1
171213
job_info["max_array_idx"] = job_info["job_id_list"][-1]
172214
job_info["array_range"] = format_range(job_info["job_id_list"])
173215

174216
# quit if we have too many array jobs
175217
if job_info["max_array_idx"] > job_info["max_array_size"]:
176-
print(safe_fill("Your job file would result in a job array with a maximum index of {max_array_idx}. This exceeds allowed array size of {max_array_size}. Split the jobs into chunks that are smaller than {max_array_size}, or do more per job.".format(**job_info), term_columns-1))
218+
print(
219+
safe_fill(
220+
"Your job file would result in a job array with a maximum index of {max_array_idx}. This exceeds allowed array size of {max_array_size}. Split the jobs into chunks that are smaller than {max_array_size}, or do more per job.".format(
221+
**job_info
222+
),
223+
term_columns - 1,
224+
)
225+
)
177226
sys.exit(1)
178227
job_info["array_fmt_width"] = len(str(job_info["max_array_idx"]))
179228

@@ -186,7 +235,9 @@ def format_range(jobnums):
186235
if args.output is not None:
187236
job_info["slurm_args"]["--output"] = args.output[0]
188237
else:
189-
job_info["slurm_args"]["--output"] = "dsq-{job_file_no_ext}-%A_%{array_fmt_width}a-%N.out".format(**job_info)
238+
job_info["slurm_args"][
239+
"--output"
240+
] = "dsq-{job_file_no_ext}-%A_%{array_fmt_width}a-%N.out".format(**job_info)
190241

191242
# set ouput directory
192243
if args.suppress_stats_file:
@@ -197,7 +248,12 @@ def format_range(jobnums):
197248
else:
198249
job_info["status_dir"] = path.abspath("./")
199250
if not os.access(job_info["status_dir"], os.W_OK | os.X_OK):
200-
print("{status_dir} does not appear to be a writeable directory.".format(**job_info), file=sys.stderr)
251+
print(
252+
"{status_dir} does not appear to be a writeable directory.".format(
253+
**job_info
254+
),
255+
file=sys.stderr,
256+
)
201257
sys.exit(1)
202258
job_info["status_dir_arg"] = "--status-dir {}".format(job_info["status_dir"])
203259

@@ -222,7 +278,9 @@ def format_range(jobnums):
222278
for option, value in job_info["slurm_args"].items():
223279
job_info["cli_args"] += " %s=%s" % (option, value)
224280

225-
cmd = "sbatch {cli_args} {user_slurm_args} {run_script} {job_file_arg} {status_dir_arg}".format(**job_info)
281+
cmd = "sbatch {cli_args} {user_slurm_args} {run_script} {job_file_arg} {status_dir_arg}".format(
282+
**job_info
283+
)
226284
print("submitting:\n {}".format(cmd))
227285
ret = call(cmd, shell=True)
228286
sys.exit(ret)
@@ -236,17 +294,36 @@ def format_range(jobnums):
236294
if args.batch_file is not None:
237295
job_info["batch_script_out"] = open(args.batch_file[0], "w")
238296
else:
239-
job_info["batch_script_out"] = open("dsq-{job_file_no_ext}-{today}.sh".format(**job_info), "w")
297+
job_info["batch_script_out"] = open(
298+
"dsq-{job_file_no_ext}-{today}.sh".format(**job_info), "w"
299+
)
240300
except Exception as e:
241-
print("Error: Couldn't open {batch_script_out} for writing. ".format(**job_info), e)
242-
243-
print("#!/bin/bash", file=job_info["batch_script_out"])
301+
print(
302+
"Error: Couldn't open {batch_script_out} for writing. ".format(
303+
**job_info
304+
),
305+
e,
306+
)
307+
308+
print("#!/bin/bash", file=job_info["batch_script_out"])
244309
for option, value in job_info["slurm_args"].items():
245310
print("#SBATCH {} {}".format(option, value), file=job_info["batch_script_out"])
246311
if len(job_info["user_slurm_args"]) > 0:
247-
print("#SBATCH {user_slurm_args}".format(**job_info), file=job_info["batch_script_out"])
248-
print("\n# DO NOT EDIT LINE BELOW".format(**job_info), file=job_info["batch_script_out"])
249-
print("{run_script} {job_file_arg} {status_dir_arg}\n".format(**job_info), file=job_info["batch_script_out"])
312+
print(
313+
"#SBATCH {user_slurm_args}".format(**job_info),
314+
file=job_info["batch_script_out"],
315+
)
316+
print(
317+
"\n# DO NOT EDIT LINE BELOW".format(**job_info),
318+
file=job_info["batch_script_out"],
319+
)
320+
print(
321+
"{run_script} {job_file_arg} {status_dir_arg}\n".format(**job_info),
322+
file=job_info["batch_script_out"],
323+
)
250324
if not args.stdout:
251-
print("Batch script generated. To submit your jobs, run:\n sbatch {}".format(job_info["batch_script_out"].name))
252-
325+
print(
326+
"Batch script generated. To submit your jobs, run:\n sbatch {}".format(
327+
job_info["batch_script_out"].name
328+
)
329+
)

0 commit comments

Comments
 (0)