Add code which restarts slurm job from state file

JanLuca · JanLuca · commit 53b91fad40b4 · 2025-08-08T14:03:45.000+02:00
diff --git a/varipeps/config.py b/varipeps/config.py
@@ -39,6 +39,17 @@ class Wavevector_Type(IntEnum):
     TWO_PI_SYMMETRIC = auto()  #: Use interval [-2pi, 2pi) for q vectors
 
 
+@unique
+class Slurm_Restart_Mode(IntEnum):
+    DISABLED = (
+        auto()
+    )  #: Disable automatic restart of slurm job if maximal runtime limit is reached
+    WRITE_RESTART_SCRIPT = (
+        auto()
+    )  #: Write restart script but do not submit new slurm job
+    AUTOMATIC_RESTART = auto()  #: Write restart script and start new slurm job with it
+
+
 @dataclass
 @register_pytree_node_class
 class VariPEPS_Config:
@@ -190,6 +201,8 @@ class VariPEPS_Config:
         :obj:`scipy.optimize.basinhopping`. See this function for details.
       spiral_wavevector_type (:obj:`Wavevector_Type`):
         Type of wavevector to be used (only positive/symmetric interval/...).
+      slurm_restart_mode (:obj:`Slurm_Restart_Mode`):
+        Mode of operation to restart slurm job if maximal runtime is reached.
     """
 
     # AD config
@@ -265,6 +278,9 @@ class VariPEPS_Config:
     # Spiral PEPS
     spiral_wavevector_type: Wavevector_Type = Wavevector_Type.TWO_PI_POSITIVE_ONLY
 
+    # Slurm
+    slurm_restart_mode: Slurm_Restart_Mode = Slurm_Restart_Mode.DISABLED
+
     def update(self, name: str, value: Any) -> NoReturn:
         self.__setattr__(name, value)
 
@@ -349,6 +365,7 @@ class ConfigModuleWrapper:
         "Line_Search_Methods",
         "Projector_Method",
         "Wavevector_Type",
+        "Slurm_Restart_Mode",
         "VariPEPS_Config",
         "config",
     }
diff --git a/varipeps/optimization/optimizer.py b/varipeps/optimization/optimizer.py
@@ -21,7 +21,7 @@
 from jax.flatten_util import ravel_pytree
 
 from varipeps import varipeps_config, varipeps_global_state
-from varipeps.config import Optimizing_Methods
+from varipeps.config import Optimizing_Methods, Slurm_Restart_Mode
 from varipeps.peps import PEPS_Unit_Cell
 from varipeps.expectation import Expectation_Model
 from varipeps.config import Projector_Method
@@ -403,7 +403,6 @@ def optimize_peps_network(
     ] = autosave_function,
     additional_input: Dict[str, jnp.ndarray] = {},
     restart_state: Dict[str, Any] = {},
-    slurm_restart_script: Optional[PathLike] = None,
 ) -> Tuple[Sequence[jnp.ndarray], PEPS_Unit_Cell, Union[float, jnp.ndarray]]:
     """
     Optimize a PEPS unitcell using a variational method.
@@ -542,6 +541,9 @@ def random_noise(a):
     else:
         varipeps_global_state.ctmrg_projector_method = None
 
+    slurm_restart_written = False
+    slurm_new_job_id = None
+
     with tqdm(desc="Optimizing PEPS state", initial=count) as pbar:
         while count < varipeps_config.optimizer_max_steps:
             runtime_start = time.perf_counter()
@@ -1067,26 +1069,37 @@ def random_noise(a):
                 best_run = random_noise_retries
 
             if (
-                slurm_restart_script is not None
+                varipeps_config.slurm_restart_mode is not Slurm_Restart_Mode.DISABLED
                 and (slurm_data := SlurmUtils.get_own_job_data()) is not None
             ):
-                flatten_runtime = [j for i in step_runtime for j in i]
+                flatten_runtime = [j for i in step_runtime for j in step_runtime[i]]
                 runtime_mean = np.mean(flatten_runtime)
                 runtime_std = np.std(flatten_runtime)
 
-                if runtime_std > 0:
-                    remaining_slurm_time = (
-                        slurm_data["TimeLimit"] - slurm_data["RunTime"]
-                    )
-                    time_of_one_step = datetime.timedelta(
-                        seconds=runtime_mean + 3 * runtime_std
+                remaining_slurm_time = slurm_data["TimeLimit"] - slurm_data["RunTime"]
+                time_of_one_step = datetime.timedelta(
+                    seconds=runtime_mean + 3 * runtime_std
+                )
+
+                if remaining_slurm_time < time_of_one_step:
+                    SlurmUtils.generate_restart_scripts(
+                        f"{str(autosave_filename)}.restart.slurm",
+                        f"{str(autosave_filename)}.restart.py",
+                        f"{str(autosave_filename)}.restartable",
+                        slurm_data,
                     )
 
-                    if remaining_slurm_time < time_of_one_step:
-                        new_job_id = SlurmUtils.run_slurm_script(slurm_restart_script)
-                        if new_job_id is not None:
-                            tqdm.write(f"Started new Slurm job with ID {new_job_id:d}.")
-                        else:
+                    slurm_restart_written = True
+
+                    if (
+                        varipeps_config.slurm_restart_mode
+                        is Slurm_Restart_Mode.AUTOMATIC_RESTART
+                    ):
+                        slurm_new_job_id = SlurmUtils.run_slurm_script(
+                            f"{str(autosave_filename)}.restart.slurm",
+                            slurm_data["WorkDir"],
+                        )
+                        if slurm_new_job_id is None:
                             tqdm.write(
                                 "Failed to start new Slurm job or parse its job id."
                             )
@@ -1124,6 +1137,11 @@ def random_noise(a):
 
     print(f"Best energy result found: {best_value}")
 
+    if slurm_restart_written:
+        print("Wrote script to restart optimizer job with Slurm.")
+    if slurm_new_job_id is not None:
+        print(f"Started new Slurm job with ID {slurm_new_job_id:d}.")
+
     return OptimizeResult(
         success=True,
         x=best_tensors,
@@ -1136,6 +1154,8 @@ def random_noise(a):
         step_conv=step_conv,
         step_runtime=step_runtime,
         best_run=best_run,
+        slurm_restart_written=slurm_restart_written,
+        slurm_new_job_id=slurm_new_job_id,
     )
 
 
diff --git a/varipeps/peps/unitcell.py b/varipeps/peps/unitcell.py
@@ -1230,6 +1230,10 @@ def load_from_group(
                 config_dict["spiral_wavevector_type"] = varipeps.config.Wavevector_Type(
                     config_dict["spiral_wavevector_type"]
                 )
+            if config_dict.get("slurm_restart_mode"):
+                config_dict["slurm_restart_mode"] = varipeps.config.Slurm_Restart_Mode(
+                    config_dict["slurm_restart_mode"]
+                )
 
             return cls(
                 data=data, real_ix=real_ix, real_iy=real_iy
diff --git a/varipeps/utils/slurm.py b/varipeps/utils/slurm.py
@@ -1,6 +1,10 @@
 import datetime
+import math
 import os
+import pathlib
+import sys
 import subprocess
+import textwrap
 
 
 class SlurmUtils:
@@ -15,8 +19,9 @@ def parse_special_fields(job_data):
             "StartTime",
             "EndTime",
             "LastSchedEval",
+            "PreemptEligibleTime",
         ):
-            if (entry := job_data.get(field)) is not None:
+            if (entry := job_data.get(field)) is not None and entry != "Unknown":
                 job_data[field] = datetime.datetime.fromisoformat(entry)
 
         for field in ("RunTime", "TimeLimit", "DelayBoot"):
@@ -37,10 +42,11 @@ def parse_special_fields(job_data):
                     seconds=int(seconds),
                 )
 
-        if (entry := job_data.get("TRES")) is not None:
-            entry = entry.split(",")
-            if len(entry) > 0:
-                job_data["TRES"] = dict(e.split("=", 1) for e in entry)
+        for field in ("TRES", "ReqTRES", "AllocTRES"):
+            if (entry := job_data.get(field)) is not None:
+                entry = entry.split(",")
+                if len(entry) > 0:
+                    job_data[field] = dict(e.split("=", 1) for e in entry)
 
         for field in job_data:
             try:
@@ -93,40 +99,13 @@ def get_own_job_data(cls):
         return None
 
     @staticmethod
-    def run_slurm_script(path):
+    def run_slurm_script(path, cwd=None):
+        cwd = pathlib.Path(cwd).resolve()
+
         p = subprocess.run(
-            ["scontrol", "show", "job", f"{job_id:d}"], capture_output=True, text=True
+            ["sbatch", str(path)], capture_output=True, text=True, cwd=cwd
         )
 
-        if p.returncode != 0:
-            return None
-
-        job_data = p.stdout.split()
-
-        slice_comb_list = []
-        for i, e in enumerate(job_data):
-            if "=" not in e:
-                slice_comb_list[-1] = slice(slice_comb_list[-1].start, i + 1)
-            else:
-                slice_comb_list.append(slice(i, i + 1))
-
-        job_data = ["".join(job_data[s]) for s in slice_comb_list]
-        job_data = dict(e.split("=", 1) for e in job_data)
-
-        job_data = cls.parse_special_fields(job_data)
-
-        return job_data
-
-    @classmethod
-    def get_own_job_data(cls):
-        if (job_id := os.environ.get("SLURM_JOB_ID")) is not None:
-            return cls.get_job_data(job_id)
-        return None
-
-    @staticmethod
-    def run_slurm_script(path):
-        p = subprocess.run(["sbatch", str(path)], capture_output=True, text=True)
-
         if p.returncode != 0:
             return None
 
@@ -136,3 +115,96 @@ def run_slurm_script(path):
             job_id = None
 
         return job_id
+
+    @staticmethod
+    def generate_restart_scripts(
+        slurm_script_path,
+        python_script_path,
+        restart_state_file,
+        slurm_data,
+        executable=None,
+    ):
+        TEMPLATE_PYTHON = textwrap.dedent(
+            """\
+        #!/usr/bin/env python3
+        import argparse
+        import pathlib
+        import varipeps
+
+        parser = argparse.ArgumentParser()
+        parser.add_argument('filename', type=pathlib.Path)
+        args = parser.parse_args()
+
+        varipeps.optimization.restart_from_state_file(args.filename)
+        """
+        )
+
+        TEMPLATE_SLURM = textwrap.dedent(
+            """\
+        #!/bin/bash
+
+        #SBATCH --partition={partition}
+        #SBATCH --qos={qos}
+        #SBATCH --job-name={job_name}
+        #SBATCH --ntasks={ntasks:d}
+        #SBATCH --nodes={nodes:d}
+        #SBATCH --cpus-per-task={ncpus:d}
+        #SBATCH --mem={mem}
+        #SBATCH --time={time_limit}
+        #SBATCH --mail-type=FAIL,END
+
+        "{executable}" "{python_script}" "{state_file}"
+        """
+        )
+
+        python_script_path = pathlib.Path(python_script_path).resolve()
+
+        restart_state_file = pathlib.Path(restart_state_file).resolve()
+
+        if executable is None:
+            executable = pathlib.Path(sys.executable).absolute()
+
+        if (tres := slurm_data.get("ReqTRES")) is not None:
+            mem = tres["mem"]
+        else:
+            tres = slurm_data["TRES"]
+            mem = tres["mem"]
+
+        time_limit = slurm_data["TimeLimit"]
+        if time_limit.days > 0:
+            time_limit_diff = time_limit - datetime.timedelta(days=time_limit.days)
+        else:
+            time_limit_diff = time_limit
+
+        time_limit_hours = math.floor(time_limit_diff / datetime.timedelta(hours=1))
+        time_limit_diff -= datetime.timedelta(hours=time_limit_hours)
+
+        time_limit_minutes = math.floor(time_limit_diff / datetime.timedelta(minutes=1))
+        time_limit_diff -= datetime.timedelta(minutes=time_limit_minutes)
+
+        time_limit_seconds = math.floor(time_limit_diff.total_seconds())
+
+        if time_limit.days > 0:
+            time_limit_str = f"{time_limit.days}-{time_limit_hours:02d}:{time_limit_minutes:02d}:{time_limit_seconds:02d}"
+        else:
+            time_limit_str = f"{time_limit_hours:02d}:{time_limit_minutes:02d}:{time_limit_seconds:02d}"
+
+        slurm_file_content = TEMPLATE_SLURM.format(
+            partition=slurm_data["Partition"],
+            qos=slurm_data["QOS"],
+            job_name=f"{slurm_data['JobName']}_restarted",
+            ntasks=slurm_data["NumTasks"],
+            nodes=slurm_data["NumNodes"],
+            ncpus=slurm_data["CPUs/Task"],
+            mem=mem,
+            time_limit=time_limit_str,
+            executable=str(executable),
+            python_script=str(python_script_path),
+            state_file=restart_state_file,
+        )
+
+        with python_script_path.open("w") as f:
+            f.write(TEMPLATE_PYTHON)
+
+        with pathlib.Path(slurm_script_path).open("w") as f:
+            f.write(slurm_file_content)
