AAAI Version

AAAI Supplementary Material/Model Training Code/srun-sbatch.py

@@ -0,0 +1,46 @@
+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)