Add feature to run slurm restart script if the runtime seems to be not enough for the current job

Author: JanLuca

varipeps/optimization/optimizer.py

@@ -1,5 +1,6 @@
 import collections
 from collections import deque
+import datetime
 from functools import partial
 from os import PathLike
 import time
@@ -8,6 +9,8 @@
 
 from tqdm_loggable.auto import tqdm
 
+import numpy as np
+
 import jax
 from jax import jit
 import jax.numpy as jnp
@@ -22,6 +25,7 @@
 from varipeps.mapping import Map_To_PEPS_Model
 from varipeps.ctmrg import CTMRGNotConvergedError, CTMRGGradientNotConvergedError
 from varipeps.utils.random import PEPS_Random_Number_Generator
+from varipeps.utils.slurm import SlurmUtils
 
 from .inner_function import (
     calc_ctmrg_expectation,
@@ -283,6 +287,7 @@ def optimize_peps_network(
         [PathLike, Sequence[jnp.ndarray], PEPS_Unit_Cell], None
     ] = autosave_function,
     additional_input: Dict[str, jnp.ndarray] = {},
+    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.
@@ -830,6 +835,32 @@ def random_noise(a):
                 best_unitcell = working_unitcell
                 best_run = random_noise_retries
 
+            if (
+                slurm_restart_script is not None
+                and (slurm_data := SlurmUtils.get_own_job_data()) is not None
+            ):
+                flatten_runtime = [j for i in step_runtime for j in 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
+                    )
+
+                    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:
+                            tqdm.write(
+                                "Failed to start new Slurm job or parse its job id."
+                            )
+                        break
+
     if working_value < best_value:
         best_value = working_value
         best_tensors = working_tensors

varipeps/utils/slurm.py

@@ -91,3 +91,48 @@ 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(
+            ["scontrol", "show", "job", f"{job_id:d}"], capture_output=True, text=True
+        )
+
+        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
+
+        try:
+            job_id = int(p.stdout.split()[-1])
+        except ValueError:
+            job_id = None
+
+        return job_id