asyncmd 0.3.3__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.
 
     """
 
@@ -133,21 +144,7 @@ class SlurmClusterMediator:
         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)
@@ -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:
@@ -624,11 +621,11 @@ class SlurmProcess:
             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 correponding 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": ""}.
+            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
@@ -649,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)
@@ -892,13 +875,13 @@ class SlurmProcess:
         # 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:
@@ -1197,11 +1180,11 @@ async def create_slurmprocess_submit(jobname: str,
         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 correponding 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": ""}.
+        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

asyncmd/tools.py

@@ -12,8 +12,33 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with asyncmd. If not, see <https://www.gnu.org/licenses/>.
+"""
+This file contains functions and classes (re)used internally in asyncmd.
+
+These functions and classes are not (thought to be) exposed to the users but
+instead intended to be (re)used in newly added asyncmd code.
+This is also not the place for MD-related utility functions, for this see utils.py
+
+Currently in here are:
+
+- ensure_executable_available
+- remove_file_if_exist and remove_file_if_exist_async
+- attach_kwargs_to_object: a function to attach kwargs to an object as properties
+  or attributes. This does type checking and warns when previously unset things
+  are set. It is used, e.g., in the GmxEngine and SlurmProcess classes.
+- DescriptorWithDefaultOnInstanceAndClass and DescriptorOutputTrajType: two descriptor
+  classes to make default values accessible on the class level but still enable checks
+  when setting on the instance level (like a property), used in the GmxEngine classes
+  but could/should be useful for any MDEngine class
+- FlagChangeList (and its typed sibling): lists with some sugar to remember if
+  their content has changed
+
+"""
+import collections
 import os
 import shutil
+import logging
+import typing
 import aiofiles
 
 
@@ -45,11 +70,11 @@ def ensure_executable_available(executable: str) -> str:
         executable = os.path.abspath(executable)
         if not os.access(executable, os.X_OK):
             raise ValueError(f"{executable} must be executable.")
-    elif shutil.which(executable) is not None:
+    elif (which_exe := shutil.which(executable)) is not None:
         # see if we find it in $PATH
-        executable = shutil.which(executable)
+        executable = which_exe
     else:
-        raise ValueError(f"{executable} must be an existing path or accesible "
+        raise ValueError(f"{executable} must be an existing path or accessible "
                          + "via the $PATH environment variable.")
     return executable
 
@@ -66,7 +91,6 @@ def remove_file_if_exist(f: str):
     try:
         os.remove(f)
     except FileNotFoundError:
-        # TODO: should we info/warn if the file is not there?
         pass
 
 
@@ -82,5 +106,260 @@ async def remove_file_if_exist_async(f: str):
     try:
         await aiofiles.os.remove(f)
     except FileNotFoundError:
-        # TODO: should we info/warn if the file is not there?
         pass
+
+
+def attach_kwargs_to_object(obj, *, logger: logging.Logger,
+                            **kwargs
+                            ) -> None:
+    """
+    Set all kwargs as object attributes/properties, error on mismatching type.
+
+    Warn when we set an unknown (i.e. previously undefined attribute/property)
+
+    Parameters
+    ----------
+    obj : object
+        The object to attach the kwargs to.
+    logger: logging.Logger
+        The logger to use for logging.
+    **kwargs : dict
+        Zero to N keyword arguments.
+    """
+    dval = object()
+    for kwarg, value in kwargs.items():
+        if (cval := getattr(obj, kwarg, dval)) is not dval:
+            if isinstance(value, type(cval)):
+                # value is of same type as default so set it
+                setattr(obj, 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)
+
+
+class DescriptorWithDefaultOnInstanceAndClass:
+    """
+    A descriptor that makes the (default) value of the private attribute
+    ``_name`` of the class it is attached to accessible as ``name`` on both the
+    class and the instance level.
+    Accessing the default value works from the class-level, i.e. without
+    instantiating the object, but note that setting on the class level
+    overwrites the descriptor and does not call ``__set__``.
+    Setting from an instance calls __set__ and therefore only sets the attribute
+    for the given instance (and also runs potential checks done in ``__set__``).
+    Also see the python docs:
+    https://docs.python.org/3/howto/descriptor.html#customized-names
+    """
+    private_name: str
+
+    def __set_name__(self, owner, name: str) -> None:
+        self.private_name = "_" + name
+
+    def __get__(self, obj, objtype=None) -> typing.Any:
+        if obj is None:
+            # I (hejung) think if obj is None objtype will always be set
+            # to the class of the obj
+            obj = objtype
+        val = getattr(obj, self.private_name)
+        return val
+
+    def __set__(self, obj, val) -> None:
+        setattr(obj, self.private_name, val)
+
+
+class DescriptorOutputTrajType(DescriptorWithDefaultOnInstanceAndClass):
+    """
+    Check the value given is in the set of allowed values before setting.
+
+    Used to check ``output_traj_type`` of MDEngines for consistency when setting.
+    """
+    # set of allowed values, e.g., trajectory file endings (without "." and all lower case)
+    ALLOWED_VALUES: set[str] = set()
+
+    def __set_name__(self, owner, name: str) -> None:
+        if not self.ALLOWED_VALUES:
+            # make sure we can only instantiate with ALLOWED_VALUES set,
+            # i.e. make this class a sort of ABC :)
+            raise NotImplementedError(f"Can not instantiate {type(self)} "  # pragma: no cover
+                                      "without allowed trajectory types set. "
+                                      "Set ``ALLOWED_VALUES`` to a set of strings.")
+        super().__set_name__(owner, name)
+
+    def __set__(self, obj, val: str) -> None:
+        if (val := val.lower()) not in self.ALLOWED_VALUES:
+            raise ValueError("output_traj_type must be one of "
+                             + f"{self.ALLOWED_VALUES}, but was {val}."
+                             )
+        super().__set__(obj, val)
+
+    def __get__(self, obj, objtype=None) -> str:
+        return super().__get__(obj=obj, objtype=objtype)
+
+
+class FlagChangeList(collections.abc.MutableSequence):
+    """A list that knows if it has been changed after initializing."""
+
+    def __init__(self, data: collections.abc.Iterable) -> None:
+        """
+        Initialize a `FlagChangeList`.
+
+        Parameters
+        ----------
+        data : Iterable
+            The data this `FlagChangeList` will hold.
+
+        Raises
+        ------
+        TypeError
+            Raised when data is not an :class:`Iterable`.
+        """
+        self._data = list(data)
+        self._changed = False
+
+    @property
+    def changed(self) -> bool:
+        """
+        Whether this `FlagChangeList` has been modified since creation.
+
+        Returns
+        -------
+        bool
+        """
+        return self._changed
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return self._data.__repr__()
+
+    def __getitem__(self, index: int | slice) -> typing.Any:
+        return self._data.__getitem__(index)
+
+    def __len__(self) -> int:
+        return self._data.__len__()
+
+    def __setitem__(self, index: int | slice, value) -> None:
+        self._data.__setitem__(index, value)
+        self._changed = True
+
+    def __delitem__(self, index: int | slice) -> None:
+        self._data.__delitem__(index)
+        self._changed = True
+
+    def insert(self, index: int, value: typing.Any):
+        """
+        Insert `value` at position given by `index`.
+
+        Parameters
+        ----------
+        index : int
+            The index of the new value in the `FlagChangeList`.
+        value : typing.Any
+            The value to insert into this `FlagChangeList`.
+        """
+        self._data.insert(index, value)
+        self._changed = True
+
+    def __add__(self, other: collections.abc.Iterable):
+        return FlagChangeList(data=self._data + list(other))
+
+    def __iadd__(self, other: collections.abc.Iterable):
+        for val in other:
+            self.append(val)
+        return self
+
+
+class TypedFlagChangeList(FlagChangeList):
+    """
+    A :class:`FlagChangeList` with an ensured type for individual list items.
+
+    Note that single strings are not treated as Iterable, i.e. (as opposed to
+    a "normal" list) `TypedFlagChangeList(data="abc")` will result in
+    `data=["abc"]` (and not `data=["a", "b", "c"]`).
+    """
+
+    def __init__(self, data: collections.abc.Iterable, dtype: type) -> None:
+        """
+        Initialize a `TypedFlagChangeList`.
+
+        Parameters
+        ----------
+        data : Iterable
+            (Initial) data for this `TypedFlagChangeList`.
+        dtype : Callable datatype
+            The datatype for all entries in this `TypedFlagChangeList`. Will be
+            called on every value separately and is expected to convert to the
+            desired datatype.
+        """
+        self._dtype = dtype  # set first to use in _convert_type method
+        data = self._ensure_iterable(data)
+        typed_data = [self._convert_type(v, index=i)
+                      for i, v in enumerate(data)]
+        super().__init__(data=typed_data)
+
+    @property
+    def dtype(self) -> type:
+        """
+        All values in this `TypedFlagChangeList` are converted to dtype.
+
+        Returns
+        -------
+        type
+        """
+        return self._dtype
+
+    def _ensure_iterable(self, data) -> collections.abc.Iterable:
+        if getattr(data, '__iter__', None) is None:
+            # convenience for singular options,
+            # if it has no iter attribute we assume it is the only item
+            data = [data]
+        elif isinstance(data, str):
+            # strings have an iter but we still do not want to split them into
+            # single letters, so just put a list around
+            data = [data]
+        return data
+
+    def _convert_type(self, value,
+                      index: int | slice | None = None) -> list:
+        # here we ignore index, but passing it should in principal make it
+        # possible to use different dtypes for different indices
+        if isinstance(index, int):
+            return self._dtype(value)
+        return [self._dtype(v) for v in value]
+
+    def __setitem__(self, index: int | slice, value) -> None:
+        typed_value = self._convert_type(value, index=index)
+        self._data.__setitem__(index, typed_value)
+        self._changed = True
+
+    def insert(self, index: int, value) -> None:
+        """
+        Insert `value` at position given by `index`.
+
+        Parameters
+        ----------
+        index : int
+            The index of the new value in the `TypedFlagChangeList`.
+        value : typing.Any
+            The value to insert into this `TypedFlagChangeList`.
+        """
+        typed_value = self._convert_type(value, index=index)
+        self._data.insert(index, typed_value)
+        self._changed = True
+
+    def __add__(self, other: collections.abc.Iterable):
+        # cast other to an iterable as we expect it (excluding the strings)
+        other = self._ensure_iterable(other)
+        ret = TypedFlagChangeList(data=self._data + list(other),
+                                  dtype=self._dtype)
+        return ret
+
+    def __iadd__(self, other: collections.abc.Iterable):
+        # cast other to an iterable as we expect it (excluding the strings)
+        other = self._ensure_iterable(other)
+        for val in other:
+            self.append(val)
+        return self

asyncmd/trajectory/__init__.py

@@ -12,13 +12,31 @@
 #
 # 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 classes and functions for engine-agnostic but trajectory-related operations.
+
+All user-facing classes and functions are reexported here for convenience.
+This includes:
+
+- the TrajectoryFunctionWrapper classes for CV value calculation and caching,
+- the Conditional/InParts TrajectoryPropagator classes for propagation of MD in
+  parts and/or until a condition is reached (and related functions),
+- and classes for extracting and concatenating trajectories (FrameExtractors
+  and TrajectoryConcatenator)
+
+"""
+from .convert import (NoModificationFrameExtractor,
+                      InvertedVelocitiesFrameExtractor,
+                      RandomVelocitiesFrameExtractor,
+                      TrajectoryConcatenator,
+                      )
 from .functionwrapper import (PyTrajectoryFunctionWrapper,
                               SlurmTrajectoryFunctionWrapper,
                               )
 from .propagate import (ConditionalTrajectoryPropagator,
                         TrajectoryPropagatorUntilAnyState,
                         InPartsTrajectoryPropagator,
-                        construct_TP_from_plus_and_minus_traj_segments,
+                        construct_tp_from_plus_and_minus_traj_segments,
                         )
 from .trajectory import (_forget_trajectory,
                          _forget_all_trajectories,