47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
import os
|
|
import sys
|
|
|
|
assert len(sys.argv) >= 2, f"Add a sbatch script as the first argument"
|
|
assert os.path.isfile(
|
|
sys.argv[1]
|
|
), f"First argument has to be an executable script (a file that exists), but got '{sys.argv[1]}'"
|
|
|
|
with open(sys.argv[1], "r") as f:
|
|
script = f.readlines()
|
|
|
|
script = [l.strip() for l in script if len(l.strip()) > 0]
|
|
|
|
# join lines ending with \
|
|
joined_script = []
|
|
current_line = ""
|
|
for line in script:
|
|
current_line += line
|
|
if current_line.endswith("\\"):
|
|
current_line = current_line[:-1]
|
|
else:
|
|
joined_script.append(current_line)
|
|
current_line = ""
|
|
script = joined_script
|
|
|
|
additional_srun_params = []
|
|
srun_lines = []
|
|
for line in script:
|
|
if line.upper().startswith("#SBATCH "):
|
|
param = line[len("#SBATCH ") :]
|
|
if param.startswith("--output="):
|
|
continue
|
|
additional_srun_params.append(param)
|
|
elif line.startswith("srun "):
|
|
srun_lines.append(line)
|
|
|
|
further_args = " ".join(sys.argv[2:])
|
|
|
|
sruns = [
|
|
line.replace("srun ", "srun " + " ".join(additional_srun_params) + " ").replace('"$@"', further_args)
|
|
for line in srun_lines
|
|
]
|
|
|
|
for srun_line in sruns:
|
|
print(f"I will run:\n{srun_line}", flush=True)
|
|
os.system(srun_line)
|