asyncmd 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

asyncmd/slurm.py

@@ -12,6 +12,18 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with asyncmd. If not, see <https://www.gnu.org/licenses/>.
+"""
+This module contains the implementation of the classes to interact with Slurm.
+
+The SlurmClusterMediator is a singleton class (handling all sacct calls in a
+coordinated fashion) for all SlurmProcess instances.
+The SlurmProcess is a drop-in replacement for asyncio.subprocess.Subprocess and
+in this spirit this module also contains the function create_slurmprocess_submit,
+which similarly to asyncio.create_subprocess_exec, creates a SlurmProcess and
+directly submits the job.
+Finally this module contains two functions to set the configuration of this module,
+set_all_slurm_settings and set_slurm_settings.
+"""
 import asyncio
 import collections
 import logging
@@ -28,6 +40,7 @@ from .tools import (ensure_executable_available,
                     remove_file_if_exist_async,
                     remove_file_if_exist,
                     )
+from .tools import attach_kwargs_to_object as _attach_kwargs_to_object
 from ._config import _SEMAPHORES
 
 
@@ -108,8 +121,6 @@ class SlurmClusterMediator:
     success_to_fail_ratio : int
         Number of successful jobs we need to observe per node to decrease the
         failed job counter by one.
-    exclude_nodes : list[str]
-        List of nodes to exclude in job submissions.
 
     """
 
@@ -130,38 +141,24 @@ class SlurmClusterMediator:
     #             (here forever means until we reinitialize SlurmClusterMediator)
 
     def __init__(self, **kwargs) -> None:
-        self._exclude_nodes = []
+        self._exclude_nodes: list[str] = []
         # make it possible to set any attribute via kwargs
         # check the type for attributes with default values
-        dval = object()
-        for kwarg, value in kwargs.items():
-            cval = getattr(self, kwarg, dval)
-            if cval is not dval:
-                if isinstance(value, type(cval)):
-                    # value is of same type as default so set it
-                    setattr(self, kwarg, value)
-                else:
-                    raise TypeError(f"Setting attribute {kwarg} with "
-                                    + f"mismatching type ({type(value)}). "
-                                    + f" Default type is {type(cval)}."
-                                    )
-            else:
-                # not previously defined, so warn that we ignore it
-                logger.warning("Ignoring unknown keyword-argument %s.", kwarg)
+        _attach_kwargs_to_object(obj=self, logger=logger, **kwargs)
         # this either checks for our defaults or whatever we just set via kwargs
         self.sacct_executable = ensure_executable_available(self.sacct_executable)
         self.sinfo_executable = ensure_executable_available(self.sinfo_executable)
-        self._node_job_fails = collections.Counter()
-        self._node_job_successes = collections.Counter()
+        self._node_job_fails: dict[str,int] = collections.Counter()
+        self._node_job_successes: dict[str,int] = collections.Counter()
         self._all_nodes = self.list_all_nodes()
-        self._jobids = []  # list of jobids of jobs we know about
-        self._jobids_sacct = []  # list of jobids we monitor actively via sacct
+        self._jobids: list[str] = []  # list of jobids of jobs we know about
+        self._jobids_sacct: list[str] = []  # list of jobids we monitor actively via sacct
         # we will store the info about jobs in a dict keys are jobids
         # values are dicts with key queried option and value the (parsed)
         # return value
         # currently queried options are: state, exitcode and nodelist
-        self._jobinfo = {}
-        self._last_sacct_call = 0  # make sure we dont call sacct too often
+        self._jobinfo: dict[str,dict] = {}
+        self._last_sacct_call = 0.  # make sure we dont call sacct too often
         # make sure we can only call sacct once at a time
         # (since there is only one ClusterMediator at a time we can create
         #  the semaphore here in __init__)
@@ -332,12 +329,12 @@ class SlurmClusterMediator:
         # (note that one semaphore counts for 3 files!)
         await _SEMAPHORES["MAX_FILES_OPEN"].acquire()
         try:
-            sacct_proc = await asyncio.subprocess.create_subprocess_exec(
+            sacct_proc = await asyncio.create_subprocess_exec(
                                                 *shlex.split(sacct_cmd),
                                                 stdout=asyncio.subprocess.PIPE,
                                                 stderr=asyncio.subprocess.PIPE,
                                                 close_fds=True,
-                                                                          )
+                                                              )
             stdout, stderr = await sacct_proc.communicate()
             sacct_return = stdout.decode()
         except asyncio.CancelledError as e:
@@ -383,17 +380,18 @@ class SlurmClusterMediator:
                 self._jobinfo[jobid]["nodelist"] = nodelist
                 self._jobinfo[jobid]["exitcode"] = exitcode
                 self._jobinfo[jobid]["state"] = state
-                logger.debug(f"Extracted from sacct output: jobid {jobid},"
-                             + f" state {state}, exitcode {exitcode} and "
-                             + f"nodelist {nodelist}.")
+                logger.debug("Extracted from sacct output: jobid %s, state %s, "
+                             "exitcode %s and nodelist %s.",
+                             jobid, state, exitcode, nodelist,
+                            )
                 parsed_ec = self._parse_exitcode_from_slurm_state(slurm_state=state)
                 self._jobinfo[jobid]["parsed_exitcode"] = parsed_ec
                 if parsed_ec is not None:
-                    logger.debug("Parsed slurm state %s for job %s"
-                                 " as returncode %s. Removing job"
-                                 "from sacct calls because its state will"
-                                 " not change anymore.",
-                                 state, jobid, parsed_ec,
+                    logger.debug("Parsed slurm state %s for job %s as "
+                                 "returncode %s. Removing job from sacct calls "
+                                 "because its state will not change anymore.",
+                                 state, jobid, str(parsed_ec) if parsed_ec is not None
+                                 else "not available",
                                  )
                     self._jobids_sacct.remove(jobid)
                     self._node_fail_heuristic(jobid=jobid,
@@ -434,8 +432,8 @@ class SlurmClusterMediator:
             # make the string a list of single node hostnames
             hostnameprefix, nums = nodelist.split("[")
             nums = nums.rstrip("]")
-            nums = nums.split(",")
-            return [f"{hostnameprefix}{num}" for num in nums]
+            nums_list = nums.split(",")
+            return [f"{hostnameprefix}{num}" for num in nums_list]
 
     def _parse_exitcode_from_slurm_state(self,
                                          slurm_state: str,
@@ -443,8 +441,9 @@ class SlurmClusterMediator:
         for ecode, regexp in self._ecode_for_slurmstate_regexps.items():
             if regexp.search(slurm_state):
                 # regexp matches the given slurm_state
-                logger.debug("Parsed SLURM state %s as exitcode %d.",
-                             slurm_state, ecode,
+                logger.debug("Parsed SLURM state %s as exitcode %s.",
+                             slurm_state, str(ecode) if ecode is not None
+                             else "not available",
                              )
                 return ecode
         # we should never finish the loop, it means we miss a slurm job state
@@ -522,11 +521,11 @@ class SlurmClusterMediator:
         all_nodes = len(self._all_nodes)
         exclude_nodes = len(self._exclude_nodes)
         if exclude_nodes >= all_nodes / 4:
-            logger.error("We already declared 1/4 of the cluster as broken."
-                         + "Houston, we might have a problem?")
+            logger.error("We already declared 1/4 of the cluster as broken. "
+                         "Houston, we might have a problem?")
             if exclude_nodes >= all_nodes / 2:
-                logger.error("In fact we declared 1/2 of the cluster as broken."
-                             + "Houston, we *do* have a problem!")
+                logger.error("In fact we declared 1/2 of the cluster as broken. "
+                             "Houston, we *do* have a problem!")
                 if exclude_nodes >= all_nodes * 0.75:
                     raise RuntimeError("Houston? 3/4 of the cluster is broken?")
 
@@ -581,9 +580,9 @@ class SlurmProcess:
         _slurm_cluster_mediator = None
         # we raise a ValueError if sacct/sinfo are not available
         logger.warning("Could not initialize SLURM cluster handling. "
-                       "If you are sure SLURM (sinfo/sacct/etc) is available"
-                       " try calling `asyncmd.config.set_slurm_settings()`"
-                       " with the appropriate arguments.")
+                       "If you are sure SLURM (sinfo/sacct/etc) is available "
+                       "try calling `asyncmd.config.set_slurm_settings()` "
+                       "with the appropriate arguments.")
     # we can not simply wait for the subprocess, since slurm exits directly
     # so we will sleep for this long between checks if slurm-job completed
     sleep_time = 15  # TODO: heuristic? dynamically adapt?
@@ -597,8 +596,9 @@ class SlurmProcess:
     scancel_executable = "scancel"
 
     def __init__(self, jobname: str, sbatch_script: str,
-                 workdir: typing.Optional[str] = None,
-                 time: typing.Optional[float] = None,
+                 workdir: str | None = None,
+                 time: float | None = None,
+                 sbatch_options: dict | None = None,
                  stdfiles_removal: str = "success",
                  **kwargs) -> None:
         """
@@ -619,6 +619,15 @@ class SlurmProcess:
         time : float or None
             Timelimit for the job in hours. None will result in using the
             default as either specified in the sbatch script or the partition.
+        sbatch_options : dict or None
+            Dictionary of sbatch options, keys are long names for options,
+            values are the corresponding values. The keys/long names are given
+            without the dashes, e.g. to specify ``--mem=1024`` the dictionary
+            needs to be ``{"mem": "1024"}``. To specify options without values
+            use keys with empty strings as values, e.g. to specify
+            ``--contiguous`` the dictionary needs to be ``{"contiguous": ""}``.
+            See the SLURM documentation for a full list of sbatch options
+            (https://slurm.schedmd.com/sbatch.html).
         stdfiles_removal : str
             Whether to remove the stdout, stderr (and possibly stdin) files.
             Possible values are:
@@ -637,21 +646,7 @@ class SlurmProcess:
         # we expect sbatch_script to be a path to a file
         # make it possible to set any attribute via kwargs
         # check the type for attributes with default values
-        dval = object()
-        for kwarg, value in kwargs.items():
-            cval = getattr(self, kwarg, dval)
-            if cval is not dval:
-                if isinstance(value, type(cval)):
-                    # value is of same type as default so set it
-                    setattr(self, kwarg, value)
-                else:
-                    raise TypeError(f"Setting attribute {kwarg} with "
-                                    + f"mismatching type ({type(value)}). "
-                                    + f" Default type is {type(cval)}."
-                                    )
-            else:
-                # not previously defined, so warn that we ignore it
-                logger.warning("Ignoring unknown keyword-argument %s.", kwarg)
+        _attach_kwargs_to_object(obj=self, logger=logger, **kwargs)
         # this either checks for our defaults or whatever we just set via kwargs
         ensure_executable_available(self.sbatch_executable)
         ensure_executable_available(self.scancel_executable)
@@ -664,13 +659,136 @@ class SlurmProcess:
         if workdir is None:
             workdir = os.getcwd()
         self.workdir = os.path.abspath(workdir)
-        self.time = time
+        self._time = time
+        # Use the property to directly call _sanitize_sbatch_options when assigning
+        # Do this **after** setting self._time to ensure consistency
+        if sbatch_options is None:
+            sbatch_options = {}
+        self.sbatch_options = sbatch_options
         self.stdfiles_removal = stdfiles_removal
-        self._jobid = None
-        self._jobinfo = {}  # dict with jobinfo cached from slurm cluster mediator
-        self._stdout_data = None
-        self._stderr_data = None
-        self._stdin = None
+        self._jobid: None | str = None
+        # dict with jobinfo cached from slurm cluster mediator
+        self._jobinfo: dict[str,typing.Any] = {}
+        self._stdout_data: None | bytes = None
+        self._stderr_data: None | bytes = None
+        self._stdin: None | str = None
+
+    def _sanitize_sbatch_options(self, sbatch_options: dict) -> dict:
+        """
+        Return sane sbatch_options dictionary to be consistent (with self).
+
+        Parameters
+        ----------
+        sbatch_options : dict
+            Dictionary of sbatch options.
+
+        Returns
+        -------
+        dict
+            Dictionary with sanitized sbatch options.
+        """
+        # NOTE: this should be called every time we modify sbatch_options or self.time!
+        # This is the list of sbatch options we use ourself, they should not
+        # be in the dict to avoid unforseen effects. We treat 'time' special
+        # because we want to allow for it to be specified via sbtach_options if
+        # it is not set via the attribute self.time.
+        reserved_sbatch_options = ["job-name", "chdir", "output", "error",
+                                   "input", "exclude", "parsable"]
+        new_sbatch_options = sbatch_options.copy()
+        if "time" in sbatch_options:
+            if self._time is not None:
+                logger.warning("Removing sbatch option 'time' from 'sbatch_options'. "
+                               "Using the 'time' argument instead.")
+                del new_sbatch_options["time"]
+            else:
+                logger.debug("Using 'time' from 'sbatch_options' because "
+                             "self.time is None.")
+        for option in reserved_sbatch_options:
+            if option in sbatch_options:
+                logger.warning("Removing sbatch option '%s' from "
+                               "'sbatch_options' because it is used internaly "
+                               "by the `SlurmProcess`.", option)
+                del new_sbatch_options[option]
+
+        return new_sbatch_options
+
+    def _slurm_timelimit_from_time_in_hours(self, time: float) -> str:
+        """
+        Create timelimit in SLURM compatible format from time in hours.
+
+        Parameters
+        ----------
+        timelimit : float
+            Timelimit for job in hours
+
+        Returns
+        -------
+        str
+            Timelimit for job as SLURM compatible string.
+        """
+        timelimit = time * 60
+        timelimit_min = int(timelimit)  # take only the full minutes
+        timelimit_sec = round(60 * (timelimit - timelimit_min))
+        timelimit_str = f"{timelimit_min}:{timelimit_sec}"
+        return timelimit_str
+
+    @property
+    def time(self) -> float | None:
+        """
+        Timelimit for SLURM job in hours.
+
+        Can be a float or None (meaning do not specify a timelimit).
+        """
+        return self._time
+
+    @time.setter
+    def time(self, val: float | None) -> None:
+        self._time = val
+        self._sbatch_options: dict = self._sanitize_sbatch_options(self._sbatch_options)
+
+    @property
+    def sbatch_options(self) -> dict:
+        """
+        A copy of the sbatch_options dictionary.
+
+        Note that modifying single key, value pairs has no effect, to modify
+        (single) sbatch_options either use the `set_sbatch_option` and
+        `del_sbatch_option` methods or (re)assign a dictionary to
+        `sbatch_options`.
+        """
+        return self._sbatch_options.copy()
+
+    @sbatch_options.setter
+    def sbatch_options(self, val: dict) -> None:
+        self._sbatch_options = self._sanitize_sbatch_options(val)
+
+    def set_sbatch_option(self, key: str, value: str) -> None:
+        """
+        Set sbatch option with given key to value.
+
+        I.e. add/modify single key, value pair in sbatch_options.
+
+        Parameters
+        ----------
+        key : str
+            The name of the sbatch option.
+        value : str
+            The value for the sbatch option.
+        """
+        self._sbatch_options[key] = value
+        self._sbatch_options = self._sanitize_sbatch_options(self._sbatch_options)
+
+    def del_sbatch_option(self, key: str) -> None:
+        """
+        Delete sbatch option with given key from sbatch_options.
+
+        Parameters
+        ----------
+        key : str
+            The name of the sbatch option to delete.
+        """
+        del self._sbatch_options[key]
+        self._sbatch_options = self._sanitize_sbatch_options(self._sbatch_options)
 
     @property
     def stdfiles_removal(self) -> str:
@@ -734,10 +852,7 @@ class SlurmProcess:
         sbatch_cmd += f" --output=./{self._stdout_name(use_slurm_symbols=True)}"
         sbatch_cmd += f" --error=./{self._stderr_name(use_slurm_symbols=True)}"
         if self.time is not None:
-            timelimit = self.time * 60
-            timelimit_min = int(timelimit)  # take only the full minutes
-            timelimit_sec = round(60 * (timelimit - timelimit_min))
-            timelimit_str = f"{timelimit_min}:{timelimit_sec}"
+            timelimit_str = self._slurm_timelimit_from_time_in_hours(self.time)
             sbatch_cmd += f" --time={timelimit_str}"
         # keep a ref to the stdin value, we need it in communicate
         self._stdin = stdin
@@ -749,19 +864,24 @@ class SlurmProcess:
         exclude_nodes = self.slurm_cluster_mediator.exclude_nodes
         if len(exclude_nodes) > 0:
             sbatch_cmd += f" --exclude={','.join(exclude_nodes)}"
+        for key, val in self.sbatch_options.items():
+            if val == "":
+                sbatch_cmd += f" --{key}"
+            else:
+                sbatch_cmd += f" --{key}={val}"
         sbatch_cmd += f" --parsable {self.sbatch_script}"
         logger.debug("About to execute sbatch_cmd %s.", sbatch_cmd)
         # 3 file descriptors: stdin,stdout,stderr
         # Note: one semaphore counts for 3 open files!
         await _SEMAPHORES["MAX_FILES_OPEN"].acquire()
         try:
-            sbatch_proc = await asyncio.subprocess.create_subprocess_exec(
+            sbatch_proc = await asyncio.create_subprocess_exec(
                                                 *shlex.split(sbatch_cmd),
                                                 stdout=asyncio.subprocess.PIPE,
                                                 stderr=asyncio.subprocess.PIPE,
                                                 cwd=self.workdir,
                                                 close_fds=True,
-                                                                          )
+                                                               )
             stdout, stderr = await sbatch_proc.communicate()
             sbatch_return = stdout.decode()
         except asyncio.CancelledError as e:
@@ -908,7 +1028,7 @@ class SlurmProcess:
         RuntimeError
             If the job has never been submitted.
         """
-        if self._jobid is None:
+        if self.slurm_jobid is None:
             # make sure we can only wait after submitting, otherwise we would
             # wait indefinitively if we call wait() before submit()
             raise RuntimeError("Can only wait for submitted SLURM jobs with "
@@ -1012,8 +1132,10 @@ class SlurmProcess:
                         + f" and output {e.output}."
                         ) from e
             # if we got until here the job is successfuly canceled....
-            logger.debug(f"Canceled SLURM job with jobid {self.slurm_jobid}."
-                         + f"scancel returned {scancel_out}.")
+            logger.debug("Canceled SLURM job with jobid %s. "
+                         "scancel returned %s.",
+                         self.slurm_jobid, scancel_out,
+                         )
             # remove the job from the monitoring
             self.slurm_cluster_mediator.monitor_remove_job(jobid=self._jobid)
             if (self._stdfiles_removal == "yes"
@@ -1034,6 +1156,7 @@ async def create_slurmprocess_submit(jobname: str,
                                      sbatch_script: str,
                                      workdir: str,
                                      time: typing.Optional[float] = None,
+                                     sbatch_options: dict | None = None,
                                      stdfiles_removal: str = "success",
                                      stdin: typing.Optional[str] = None,
                                      **kwargs,
@@ -1055,6 +1178,15 @@ async def create_slurmprocess_submit(jobname: str,
     time : float or None
         Timelimit for the job in hours. None will result in using the
         default as either specified in the sbatch script or the partition.
+    sbatch_options : dict or None
+        Dictionary of sbatch options, keys are long names for options,
+        values are the corresponding values. The keys/long names are given
+        without the dashes, e.g. to specify ``--mem=1024`` the dictionary
+        needs to be ``{"mem": "1024"}``. To specify options without values
+        use keys with empty strings as values, e.g. to specify
+        ``--contiguous`` the dictionary needs to be ``{"contiguous": ""}``.
+        See the SLURM documentation for a full list of sbatch options
+        (https://slurm.schedmd.com/sbatch.html).
     stdfiles_removal : str
         Whether to remove the stdout, stderr (and possibly stdin) files.
         Possible values are:
@@ -1078,6 +1210,7 @@ async def create_slurmprocess_submit(jobname: str,
     """
     proc = SlurmProcess(jobname=jobname, sbatch_script=sbatch_script,
                         workdir=workdir, time=time,
+                        sbatch_options=sbatch_options,
                         stdfiles_removal=stdfiles_removal,
                         **kwargs)
     await proc.submit(stdin=stdin)