Some fixes for SLURM and SGE engine

sisyphus/engine.py

@@ -50,6 +50,10 @@ def submit_call(self, call, logpath, rqmt, name, task_name, task_ids):
     def get_default_rqmt(self, task):
         raise NotImplementedError
 
+    def get_task_termination_info(self, task, task_id, submit_info):
+        """Return engine-provided usage information for a terminated task."""
+        return {}
+
     def get_used_engine(self, engine_selector):
         return self
 
@@ -92,11 +96,19 @@ def get_rqmt(self, task, task_id, update=True):
         # find last requirements
         rqmt = self.add_defaults_to_rqmt(task, rqmt)
         rqmt_hist = self.get_submit_history(task)[task_id]
-        if rqmt_hist and rqmt_hist[0] == rqmt:
-            # job has been submitted before and the rqmt given by the recipe did not change
-            rqmt.update(rqmt_hist[-1])
+        if rqmt_hist:
+            # Keep previously submitted values for requirements that did not change in the recipe. Build the result
+            # from the currently known requirement keys so that submit metadata and removed requirements are ignored.
+            missing = object()
+            initial_rqmt = rqmt_hist[0]
+            last_rqmt = rqmt_hist[-1]
+            rqmt = {
+                key: last_rqmt.get(key, value) if initial_rqmt.get(key, missing) == value else value
+                for key, value in rqmt.items()
+            }
             if update:
-                rqmt = task.update_rqmt(rqmt, task_id)
+                termination_info = self.get_task_termination_info(task, task_id, last_rqmt)
+                rqmt = task.update_rqmt(rqmt, task_id, additional_usage=termination_info)
 
         if "mem" in rqmt:
             rqmt["mem"] = tools.str_to_GB(rqmt["mem"])
@@ -322,6 +334,9 @@ def submit_call(self, call, logpath, rqmt, name, task_name, task_ids):
     def get_default_rqmt(self, task):
         return self.get_used_engine_by_rqmt(task.rqmt()).get_default_rqmt(task)
 
+    def get_task_termination_info(self, task, task_id, submit_info):
+        return self.get_used_engine_by_rqmt(submit_info).get_task_termination_info(task, task_id, submit_info)
+
     def get_job_node_hostnames(self):
         raise Exception(
             f"{self.__class__.__name__} is never an active engine at job runtime, "

sisyphus/global_settings.py

@@ -123,13 +123,15 @@ def update_engine_rqmt(last_rqmt: Dict, last_usage: Dict):
     # Did we run out of time?
     requested_time = last_rqmt.get("time")
     used_time = last_usage.get("used_time", 0)
-    if requested_time and requested_time - used_time < 0.1:
+    out_of_time = last_usage.get("out_of_time")
+    if requested_time and (out_of_time or requested_time - used_time < 0.1):
         out["time"] = requested_time * 2
 
     # Did it (nearly) break the memory limits?
     requested_memory = last_rqmt.get("mem")
     used_memory = last_usage.get("max", {}).get("rss", 0)
-    if requested_memory and last_usage.get("out_of_memory") or requested_memory - used_memory < 0.25:
+    out_of_memory = last_usage.get("out_of_memory")
+    if requested_memory and (out_of_memory or requested_memory - used_memory < 0.25):
         out["mem"] = requested_memory * 2
 
     return out

sisyphus/simple_linux_utility_for_resource_management_engine.py

@@ -1,5 +1,5 @@
 # Author: Wilfried Michel <michel@cs.rwth-aachen.de>
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 from collections import defaultdict, namedtuple
 from enum import Enum
 import getpass  # used to get username
@@ -63,6 +63,7 @@ def __init__(
                                           running jobs anymore.
         """
         self._task_info_cache_last_update = 0
+        self._task_info_cache = defaultdict(list)
         self.gateway = gateway
         self.default_rqmt = default_rqmt
         self.has_memory_resource = has_memory_resource
@@ -229,14 +230,14 @@ def submit_call(self, call, logpath, rqmt, name, task_name, task_ids):
             else:
                 # this id doesn't fit pattern, this should only happen if only parts of the jobs are restarted
                 job_id = self.submit_helper(call, logpath, rqmt, name, task_name, start_id, end_id, step_size)
-                submitted.append((list(range(start_id, end_id, step_size)), job_id))
+                submitted.append((list(range(start_id, end_id + 1, step_size)), job_id))
                 start_id, end_id, step_size = (task_id, None, None)
         assert start_id is not None
         if end_id is None:
             end_id = start_id
             step_size = 1
         job_id = self.submit_helper(call, logpath, rqmt, name, task_name, start_id, end_id, step_size)
-        submitted.append((list(range(start_id, end_id, step_size)), job_id))
+        submitted.append((list(range(start_id, end_id + 1, step_size)), job_id))
         return ENGINE_NAME, submitted
 
     def submit_helper(self, call, logpath, rqmt, name, task_name, start_id, end_id, step_size):
@@ -296,10 +297,10 @@ def submit_helper(self, call, logpath, rqmt, name, task_name, start_id, end_id,
                 # reset cache, after error
                 self.reset_cache()
             else:
-                job_id = sout[3].decode().split(".")
+                job_id = sout[3].decode().split(".")[0]
 
                 logging.info("Submitted with job_id: %s %s" % (job_id, name))
-                for task_id in range(start_id, end_id, step_size):
+                for task_id in range(start_id, end_id + 1, step_size):
                     self._task_info_cache[(name, task_id)].append((job_id, "PENDING"))
 
                 if err:
@@ -325,6 +326,38 @@ def submit_helper(self, call, logpath, rqmt, name, task_name, start_id, end_id,
     def reset_cache(self):
         self._task_info_cache_last_update = -10
 
+    def get_task_termination_info(self, task, task_id, submit_info) -> Dict[str, bool]:
+        job_id = next(
+            (job_id for task_ids, job_id in submit_info.get("engine_info", []) if task_id in task_ids),
+            None,
+        )
+        if job_id is None:
+            return {}
+
+        slurm_task_id = f"{job_id}_{task_id}"
+        command = ["sacct", "-n", "-P", "-j", slurm_task_id, "--format=State%30"]
+        try:
+            out, err, retval = self.system_call(command)
+        except (OSError, subprocess.TimeoutExpired):
+            logging.warning(self._system_call_error_warn_msg(command))
+            return {}
+        if retval != 0:
+            logging.warning(self._system_call_error_warn_msg(command, err=err))
+            return {}
+
+        termination_info = {}
+        for raw_line in out:
+            state = raw_line.decode("utf-8").strip()
+            if not state:
+                logging.warning("Failed to parse sacct output: %s" % raw_line.decode("utf-8", errors="replace"))
+                continue
+            if state == "OUT_OF_MEMORY":
+                termination_info["out_of_memory"] = True
+            elif state == "TIMEOUT":
+                termination_info["out_of_time"] = True
+
+        return termination_info
+
     def queue_state(self):
         """Returns list with all currently running tasks in this queue"""
 

sisyphus/son_of_grid_engine.py

@@ -66,6 +66,7 @@ def __init__(self, default_rqmt, gateway=None, auto_clean_eqw=True, ignore_jobs=
                             The default "mpi" is somewhat arbitrarily chosen as we have it in our environment.
         """
         self._task_info_cache_last_update = 0
+        self._task_info_cache = defaultdict(list)
         self.gateway = gateway
         self.default_rqmt = default_rqmt
         self.auto_clean_eqw = auto_clean_eqw
@@ -225,14 +226,14 @@ def submit_call(self, call, logpath, rqmt, name, task_name, task_ids):
             else:
                 # this id doesn't fit pattern, this should only happen if only parts of the jobs are restarted
                 job_id = self.submit_helper(call, logpath, rqmt, name, task_name, start_id, end_id, step_size)
-                submitted.append((list(range(start_id, end_id, step_size)), job_id))
+                submitted.append((list(range(start_id, end_id + 1, step_size)), job_id))
                 start_id, end_id, step_size = (task_id, None, None)
         assert start_id is not None
         if end_id is None:
             end_id = start_id
             step_size = 1
         job_id = self.submit_helper(call, logpath, rqmt, name, task_name, start_id, end_id, step_size)
-        submitted.append((list(range(start_id, end_id, step_size)), job_id))
+        submitted.append((list(range(start_id, end_id + 1, step_size)), job_id))
         return ENGINE_NAME, submitted
 
     def submit_helper(self, call, logpath, rqmt, name, task_name, start_id, end_id, step_size):
@@ -289,7 +290,7 @@ def submit_helper(self, call, logpath, rqmt, name, task_name, start_id, end_id,
                 job_id = sjob_id[0]
 
                 logging.info("Submitted with job_id: %s %s" % (job_id, name))
-                for task_id in range(start_id, end_id, step_size):
+                for task_id in range(start_id, end_id + 1, step_size):
                     self._task_info_cache[(name, task_id)].append((job_id, "qw"))
 
                 if False:  # for debugging

sisyphus/task.py

@@ -452,7 +452,7 @@ def _get_arg_idx_for_task_id(self, task_id):
             start = (chunk_size + 1) * overflow + chunk_size * (task_id - 1 - overflow)
             return range(start, start + chunk_size)
 
-    def update_rqmt(self, last_rqmt, task_id):
+    def update_rqmt(self, last_rqmt, task_id, additional_usage=None):
         """Update task requirements of interrupted job"""
         last_rqmt = last_rqmt.copy()
         # Make sure mem and time are numbers and not str
@@ -461,9 +461,15 @@ def update_rqmt(self, last_rqmt, task_id):
         usage_file = self._job._sis_path(gs.PLOGGING_FILE + "." + self.name(), task_id, abspath=True)
 
         try:
-            last_usage = literal_eval(open(usage_file).read())
+            with open(usage_file) as usage_file_handle:
+                last_usage = literal_eval(usage_file_handle.read())
         except (SyntaxError, IOError):
-            # we don't know anything if no usage file is writen or is invalid, just reuse last rqmts
+            last_usage = {}
+
+        if additional_usage:
+            last_usage.update(additional_usage)
+        if not last_usage:
+            # We don't know anything if no usage information is available, so just reuse the last requirements.
             return last_rqmt
         return self._update_rqmt(last_rqmt=last_rqmt, last_usage=last_usage)