✨ error-handling: recover Eliashberg calculations from walltime limits

src/aiida_epw/common/types.py

@@ -4,9 +4,43 @@
 class RestartType(enum.Enum):
     """Enumeration of EPW run/restart modes."""
 
-    NONE = "none"
+    FROM_SCRATCH = "from_scratch"
+    FROM_EPB = "from_epb"
+    FROM_EPMATWP = "from_epmatwp"
     EPHWRITE = "ephwrite"
-    EPHREAD = "ephread"
-    EPHWRITE_RESTART = "ephwrite_restart"
-    EPHREAD_RESTART = "ephread_restart"
-    EPWREAD = "epwread"
+    FROM_EPH = "from_eph"
+
+
+RESTART_TYPE_DEFAULTS = {
+    RestartType.FROM_SCRATCH: {
+        "epwread": False,
+        "epwwrite": True,
+        "epbwrite": True,
+        "epbread": False,
+    },
+    RestartType.FROM_EPB: {
+        "epbread": True,
+        "epbwrite": False,
+        "epwread": False,
+        "epwwrite": True,
+    },
+    RestartType.FROM_EPMATWP: {
+        "epwread": True,
+        "epwwrite": False,
+        "epbwrite": False,
+        "epbread": False,
+    },
+    RestartType.EPHWRITE: {
+        "epwread": True,
+        "ep_coupling": True,
+        "elph": True,
+        "ephwrite": True,
+    },
+    RestartType.FROM_EPH: {
+        "epwread": True,
+        "ep_coupling": False,
+        "elph": False,
+        "ephwrite": False,
+        "restart": False,
+    },
+}

src/aiida_epw/parsers/epw.py

@@ -43,12 +43,43 @@ class EpwParser(BaseParser):
     class_error_map = {
         r"internal error,\s*cannot bracket Ef": "ERROR_CANNOT_BRACKET_EF",
     }
+    _scheduler_walltime_patterns = (
+        r"(?i)due to time limit",
+        r"(?i)time limit reached",
+        r"(?i)walltime\s+(?:limit\s+)?exceeded",
+        r"(?i)cputime\s+(?:limit\s+)?exceeded",
+        r"(?i)term_runlimit",
+        r"(?i)exceeded.*wallclock time limit",
+        r"(?i)job.*cancelled.*due to time limit",
+    )
 
     @staticmethod
     def get_parser_settings_key():
         """Return the settings key reserved for parser-specific options."""
         return "parser_options"
 
+    def check_scheduler_errors(self, logs):
+        """Check scheduler stderr/stdout or node exit status for scheduler-level aborts."""
+        scheduler_walltime_status = getattr(
+            EpwCalculation.exit_codes, "ERROR_SCHEDULER_OUT_OF_WALLTIME", None
+        )
+        if (
+            scheduler_walltime_status is not None
+            and self.node.exit_status == scheduler_walltime_status.status
+        ):
+            logs.error.append("ERROR_OUT_OF_WALLTIME")
+            return
+
+        scheduler_stderr = self.get_retrieved_content("_scheduler-stderr.txt") or ""
+        scheduler_stdout = self.get_retrieved_content("_scheduler-stdout.txt") or ""
+        scheduler_output = f"{scheduler_stderr}\n{scheduler_stdout}"
+
+        if scheduler_output.strip():
+            for pattern in self._scheduler_walltime_patterns:
+                if re.search(pattern, scheduler_output):
+                    logs.error.append("ERROR_OUT_OF_WALLTIME")
+                    return
+
     def get_retrieved_content(self, *filenames):
         """Return the content of the first retrieved file that exists."""
         for filename in filenames:
@@ -88,8 +119,10 @@ def parse(self, **kwargs):
         parsed_epw, logs = self.parse_stdout(stdout, logs, code_version=code_version)
         parsed_data.update(parsed_epw)
 
+        self.check_scheduler_errors(logs)
+
         elbands_contents = self.get_retrieved_content(
-            EpwCalculation._output_elbands_file
+            EpwCalculation._OUTPUT_ELBANDS_FILE
         )
         if elbands_contents is not None:
             self.out(
@@ -100,7 +133,7 @@ def parse(self, **kwargs):
             )
 
         phbands_contents = self.get_retrieved_content(
-            EpwCalculation._output_phbands_file
+            EpwCalculation._OUTPUT_PHBANDS_FILE
         )
         if phbands_contents is not None:
             self.out(
@@ -221,6 +254,9 @@ def parse(self, **kwargs):
         if "ERROR_FACTORIZATION" in logs.error:
             return self.exit(self.exit_codes.get("ERROR_FACTORIZATION"), logs)
 
+        if "ERROR_OUT_OF_WALLTIME" in logs.error:
+            return self.exit(self.exit_codes.get("ERROR_OUT_OF_WALLTIME"), logs)
+
         for exit_code in list(self.get_error_map().values()):
             if exit_code in logs.error:
                 return self.exit(self.exit_codes.get(exit_code), logs)

src/aiida_epw/tools/error_handling/supercon.py

@@ -1,4 +1,4 @@
-"""Pade error-recovery helpers for EPW workflows."""
+"""Pade and walltime error-recovery helpers for EPW workflows."""
 
 from copy import deepcopy
 
@@ -184,3 +184,99 @@ def get_temperature_out_of_range_action():
         "Temperature reached phase transition (delta converged to zero). "
         "Finishing workchain successfully."
     )
+
+
+def prepare_walltime_recovery(
+    parameters, output_parameters, restart_type, is_eliashberg
+):
+    """Return ``(parameters, restart_type, action)`` for a recoverable timeout."""
+    from aiida_epw.common.types import RestartType
+
+    if not is_eliashberg or restart_type is None:
+        return (
+            None,
+            None,
+            (
+                "Walltime reached but this calculation/restart type is not supported "
+                "for auto-recovery. Aborting."
+            ),
+        )
+
+    if restart_type is RestartType.EPHWRITE:
+        return (
+            parameters,
+            RestartType.EPHWRITE,
+            (
+                "Walltime reached during Eliashberg ephwrite calculation. "
+                "Restarting from the last checkpoint."
+            ),
+        )
+
+    if restart_type is RestartType.FROM_EPH:
+        outputs = output_parameters or {}
+        eliashberg_data = (
+            outputs.get("isotropic_eliashberg")
+            or outputs.get("anisotropic_eliashberg")
+            or {}
+        )
+        succeeded_temps = _get_succeeded_temperatures(eliashberg_data)
+        input_epw = parameters.get("INPUTEPW", {})
+        all_temps, is_linear_range = _get_input_temperatures(input_epw, eliashberg_data)
+        remaining_temps = [
+            temp
+            for temp in all_temps
+            if not any(abs(temp - succeeded) < 1.0e-4 for succeeded in succeeded_temps)
+        ]
+
+        if succeeded_temps and remaining_temps:
+            parameters = deepcopy(parameters)
+            input_epw_new = parameters.setdefault("INPUTEPW", {})
+            new_temps = (
+                [remaining_temps[0], remaining_temps[-1]]
+                if is_linear_range and len(remaining_temps) >= 2
+                else remaining_temps
+            )
+            input_epw_new["temps"] = (
+                " ".join(str(temp) for temp in new_temps)
+                if isinstance(input_epw.get("temps"), str)
+                else new_temps
+            )
+            input_epw_new["nstemp"] = len(remaining_temps)
+            input_epw_new.pop("tempsmin", None)
+            input_epw_new.pop("tempsmax", None)
+            return (
+                parameters,
+                RestartType.FROM_EPH,
+                (
+                    "Walltime reached during Eliashberg from_eph calculation. "
+                    f"Removed successfully calculated temperatures: {succeeded_temps}. "
+                    "Restarting."
+                ),
+            )
+        return (
+            None,
+            None,
+            (
+                "Walltime reached during Eliashberg from_eph calculation but no "
+                "temperatures finished. Aborting."
+            ),
+        )
+
+    if restart_type is RestartType.FROM_SCRATCH:
+        return (
+            None,
+            None,
+            (
+                "Walltime reached during a fresh Eliashberg run. Automatic recovery "
+                "requires an ephwrite checkpoint. Aborting."
+            ),
+        )
+
+    return (
+        None,
+        None,
+        (
+            "Walltime reached but this calculation/restart type is not supported "
+            "for auto-recovery. Aborting."
+        ),
+    )

src/aiida_epw/tools/workchain.py

@@ -342,3 +342,105 @@ def get_default_target_basepath(computer):
     else:
         raise ValueError(f"Unsupported transport type: {computer.transport_type}")
     return target_basepath
+
+
+def pop_succeeded_temperatures(parameters, output_parameters):
+    """Analyze Eliashberg output parameters and pop successfully completed temperatures from input parameters.
+
+    :param parameters: Dict containing input namelist parameters
+    :param output_parameters: Dict containing output parameters of the calculation
+    :return: A tuple of (updated_parameters_dict, succeeded_temps, remaining_temps, eliashberg_data)
+    """
+    eliashberg_data = (
+        output_parameters.get("isotropic_eliashberg")
+        or output_parameters.get("anisotropic_eliashberg")
+        or {}
+    )
+
+    succeeded_temps = []
+    for temp_str, data in eliashberg_data.items():
+        temp = float(temp_str)
+        has_failed = False
+
+        iterations = data.get("iterations", {})
+        if iterations:
+            if not iterations.get("ethr") or any(
+                v is None for v in iterations.get("ethr", [])
+            ):
+                has_failed = True
+            elif any(v is None for v in iterations.get("znormi", [])) or any(
+                v is None for v in iterations.get("deltai", [])
+            ):
+                has_failed = True
+
+        pade = data.get("pade", {})
+        if pade:
+            if any(
+                pade.get(k) is None for k in ("delta", "znorm", "shift") if k in pade
+            ):
+                has_failed = True
+
+        if not iterations and not pade:
+            has_failed = True
+
+        if not has_failed:
+            succeeded_temps.append(temp)
+
+    input_epw = parameters.get("INPUTEPW", {})
+    all_temps = []
+    is_linear_range = False
+    if "temps" in input_epw:
+        temps_val = input_epw["temps"]
+        if isinstance(temps_val, str):
+            original_temps = [float(t) for t in temps_val.replace(",", " ").split()]
+        elif isinstance(temps_val, (int, float)):
+            original_temps = [float(temps_val)]
+        else:
+            original_temps = [float(t) for t in temps_val]
+
+        nstemp = input_epw.get("nstemp", len(original_temps))
+        if len(original_temps) == 2 and nstemp >= 2:
+            is_linear_range = True
+            t_min, t_max = original_temps[0], original_temps[1]
+            all_temps = [
+                t_min + i * (t_max - t_min) / (nstemp - 1) for i in range(nstemp)
+            ]
+        else:
+            all_temps = original_temps
+    elif "tempsmin" in input_epw and "tempsmax" in input_epw and "nstemp" in input_epw:
+        is_linear_range = True
+        t_min = float(input_epw["tempsmin"])
+        t_max = float(input_epw["tempsmax"])
+        n_temp = int(input_epw["nstemp"])
+        if n_temp > 1:
+            all_temps = [
+                t_min + i * (t_max - t_min) / (n_temp - 1) for i in range(n_temp)
+            ]
+        else:
+            all_temps = [t_min]
+    else:
+        all_temps = [float(k) for k in eliashberg_data.keys()]
+
+    remaining_temps = [
+        t for t in all_temps if not any(abs(t - st) < 1e-4 for st in succeeded_temps)
+    ]
+
+    new_temps_list = []
+    new_nstemp = len(remaining_temps)
+    if is_linear_range and new_nstemp >= 2:
+        new_temps_list = [remaining_temps[0], remaining_temps[-1]]
+    else:
+        new_temps_list = remaining_temps
+
+    updated_parameters = dict(parameters)
+    input_epw_new = updated_parameters.setdefault("INPUTEPW", {})
+
+    if isinstance(input_epw.get("temps"), str):
+        input_epw_new["temps"] = " ".join(str(t) for t in new_temps_list)
+    else:
+        input_epw_new["temps"] = new_temps_list
+    input_epw_new["nstemp"] = new_nstemp
+    input_epw_new.pop("tempsmin", None)
+    input_epw_new.pop("tempsmax", None)
+
+    return updated_parameters, succeeded_temps, remaining_temps, eliashberg_data