improvements

slurm_submission.py

@@ -108,6 +108,7 @@ def create_slurm_script(slurm_template:str,
                         job:str,
                         job_config_file:str,
                         job_log:str,
+                        array:tuple|None=None,
                         ):
     """
     Creates a SLURM submission script based on the provided template.
@@ -116,6 +117,12 @@ def create_slurm_script(slurm_template:str,
          - simbelmyne
          - scola
      """
+    
+    if array is not None and job != "scola":
+        raise ValueError(f"Array job range provided for job type {job}.")
+    if array is None and job == "scola":
+        raise ValueError(f"Array job range not provided for job type {job}.")
+    
     from os.path import isfile
     if not isfile(slurm_template):
         raise FileNotFoundError(f"SLURM template {slurm_template} does not exist.")
@@ -127,6 +134,10 @@ def create_slurm_script(slurm_template:str,
     # Create the script file
     with open(slurm_script, "w") as f:
         for line in template:
+            if array is not None and "--array" in line:
+                line = line.replace("%a-%b",f"{array[0]}-{array[1]}")
+            if array is not None and ("--output" in line  or "--error" in line):
+                line = line.replace("%j","%a_%A")
             f.write(line)
         
         # Add the job command
@@ -136,7 +147,7 @@ def create_slurm_script(slurm_template:str,
             case "simbelmyne":
                 f.write(f"{job} {job_config_file} {job_log}")
             case "scola":
-                f.write(f"{job} {job_config_file} {job_log} "+"${SLURM_ARRAY_TASK_ID}")
+                f.write(f"{job} {job_config_file} {job_log} "+"-b ${SLURM_ARRAY_TASK_ID}")
             case _:
                 raise ValueError(f"Job type {job} not recognized.")