hpcflow 0.1.15__py3-none-any.whl → 0.2.0a271__py3-none-any.whl

hpcflow/sdk/persistence/store_resource.py

@@ -0,0 +1,203 @@
+"""
+Models of data stores as resources.
+"""
+
+from __future__ import annotations
+from abc import ABC, abstractmethod
+import copy
+import json
+from typing import Any, Callable, TYPE_CHECKING
+
+from hpcflow.sdk.core.utils import get_md5_hash
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from logging import Logger
+    from pathlib import Path
+    import zarr  # type: ignore
+    from fsspec import AbstractFileSystem  # type: ignore
+    from ..app import BaseApp
+
+
+class StoreResource(ABC):
+    """Class to represent a persistent resource within which store data lives.
+
+    A `PersistentStore` maps workflow data across zero or more store resources. Updates to
+    persistent workflow data that live in the same store resource are performed together.
+
+    Parameters
+    ----------
+    app: App
+        The main application context.
+    name:
+        The store name.
+    """
+
+    def __init__(self, app: BaseApp, name: str) -> None:
+        self._app = app
+        self.name = name
+        self.data: dict[str, Any] = {"read": None, "update": None}
+        self.hash = None
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(name={self.name!r})"
+
+    @property
+    def logger(self) -> Logger:
+        """
+        The logger.
+        """
+        return self._app.persistence_logger
+
+    @abstractmethod
+    def _load(self) -> Any:
+        pass
+
+    @abstractmethod
+    def _dump(self, data: dict | list):
+        pass
+
+    def open(self, action: str):
+        """
+        Open the store.
+
+        Parameters
+        ----------
+        action:
+            What we are opening the store for; typically either ``read`` or ``update``.
+        """
+
+        # TODO: some tests?
+
+        if action == "read":
+            # reuse "update" data if set, rather than re-loading from disk -- but copy,
+            # so changes made in the "read" scope do not update!
+            update_data = self.data["update"]
+            rd_msg = " (using `update` data)" if update_data else ""
+            self.logger.debug(f"{self!r}: opening to read{rd_msg}.")
+            data = copy.deepcopy(update_data) if update_data else self._load()
+
+        elif action == "update":
+            # reuse "read" data if set, rather than re-loading from disk; this also means
+            # updates will be reflected in the "read" data as soon as they are made:
+            read_data = self.data["read"]
+            upd_msg = " (using `read` data)" if read_data else ""
+            self.logger.debug(f"{self!r}: opening to update{upd_msg}.")
+            data = read_data or self._load()
+
+        else:
+            self._check_action(action)
+
+        self.data[action] = data
+
+        try:
+            self.hash = get_md5_hash(data)  # type: ignore
+        except Exception:
+            pass
+
+    def close(self, action: str):
+        """
+        Close the store for a particular action.
+
+        Parameters
+        ----------
+        action:
+            What we are closing the store for.
+            Should match a previous call to :py:meth:`close`.
+        """
+        if action == "read":
+            self.logger.debug(f"{self!r}: closing read.")
+        elif action == "update":
+            if self.hash:
+                # check if it has changed:
+                new_hash = get_md5_hash(self.data[action])
+            if not self.hash or self.hash != new_hash:
+                self.logger.debug(f"{self!r}: data (hash) changed.")
+                self._dump(self.data[action])
+            self.logger.debug(f"{self!r}: closing update.")
+        else:
+            self._check_action(action)
+
+        # unset data for this action:
+        self.data[action] = None
+
+    def _check_action(self, action: str):
+        if action not in self.data:
+            raise ValueError(
+                f"Action {action!r} not known for {self.__class__.__name__!r}"
+            )
+
+
+class JSONFileStoreResource(StoreResource):
+    """
+    For caching reads and writes to a JSON file.
+
+    Parameters
+    ----------
+    app: App
+        The main application context.
+    name:
+        The store name.
+    filename:
+        The name of the JSON file.
+    path:
+        The path to the directory containing the JSON file.
+    fs:
+        The filesystem that the JSON file resides within.
+    """
+
+    def __init__(
+        self,
+        app: BaseApp,
+        name: str,
+        filename: str,
+        path: str | Path,
+        fs: AbstractFileSystem,
+    ):
+        self.filename = filename
+        self.path = path
+        self.fs = fs
+        super().__init__(app, name)
+
+    @property
+    def _full_path(self) -> str:
+        return f"{self.path}/{self.filename}"
+
+    def _load(self) -> Any:
+        self.logger.debug(f"{self!r}: loading JSON from file.")
+        with self.fs.open(self._full_path, mode="rt") as fp:
+            return json.load(fp)
+
+    def _dump(self, data: Mapping | list):
+        self.logger.debug(f"{self!r}: dumping JSON to file")
+        with self.fs.open(self._full_path, mode="wt") as fp:
+            json.dump(data, fp, indent=2)
+
+
+class ZarrAttrsStoreResource(StoreResource):
+    """
+    For caching reads and writes to Zarr attributes on groups and arrays.
+
+    Parameters
+    ----------
+    app: App
+        The main application context.
+    name:
+        The store name.
+    open_call:
+        How to actually perform an open on the underlying resource.
+    """
+
+    def __init__(self, app: BaseApp, name: str, open_call: Callable[..., zarr.Group]):
+        self.open_call = open_call
+        super().__init__(app, name)
+
+    def _load(self) -> Any:
+        self.logger.debug(f"{self!r}: loading Zarr attributes.")
+        item = self.open_call(mode="r")
+        return copy.deepcopy(item.attrs.asdict())
+
+    def _dump(self, data: dict | list):
+        self.logger.debug(f"{self!r}: dumping Zarr attributes.")
+        item = self.open_call(mode="r+")
+        item.attrs.put(data)

hpcflow/sdk/persistence/types.py

@@ -0,0 +1,309 @@
+"""
+Types used in type-checking the persistence subsystem.
+"""
+
+from __future__ import annotations
+from typing import Any, Generic, TypeVar, TYPE_CHECKING, DefaultDict
+from typing_extensions import TypedDict, NotRequired, TypeAlias
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from .base import StoreTask, StoreElement, StoreElementIter, StoreEAR, StoreParameter
+    from ..core.json_like import JSONed
+    from ..core.parameters import ParameterValue
+    from ..core.types import IterableParam
+    from ..typing import DataIndex, ParamSource
+
+#: Bound type variable: :class:`StoreTask`.
+AnySTask = TypeVar("AnySTask", bound="StoreTask")
+#: Bound type variable: :class:`StoreElement`.
+AnySElement = TypeVar("AnySElement", bound="StoreElement")
+#: Bound type variable: :class:`StoreElementITer`.
+AnySElementIter = TypeVar("AnySElementIter", bound="StoreElementIter")
+#: Bound type variable: :class:`StoreEAR`.
+AnySEAR = TypeVar("AnySEAR", bound="StoreEAR")
+#: Bound type variable: :class:`StoreParameter`.
+AnySParameter = TypeVar("AnySParameter", bound="StoreParameter")
+#: Type of possible stored parameters.
+ParameterTypes: TypeAlias = (
+    "ParameterValue | list | tuple | set | dict | int | float | str | None | Any"
+)
+
+
+class File(TypedDict):
+    """
+    Descriptor for file metadata.
+    """
+
+    #: Whether to store the contents.
+    store_contents: bool
+    #: The path to the file.
+    path: str
+
+
+class FileDescriptor(TypedDict):
+    """
+    Descriptor for file metadata.
+    """
+
+    #: Whether this is an input file.
+    is_input: bool
+    #: Whether to store the contents.
+    store_contents: bool
+    #: Where the file will go.
+    dst_path: str
+    #: The path to the file.
+    path: str | None
+    #: Whether to delete the file after processing.
+    clean_up: bool
+    # The contents of the file.
+    contents: NotRequired[str]
+
+
+class LoopDescriptor(TypedDict):
+    """
+    Descriptor for loop metadata.
+    """
+
+    #: The parameters iterated over by the loop.
+    iterable_parameters: dict[str, IterableParam]
+    #: The parameters output by the loop, and the final task insert ID from which they
+    #: are output.
+    output_parameters: dict[str, int]
+    #: The template data from which the loop was created.
+    loop_template: NotRequired[dict[str, Any]]
+    #: The number of iterations generated by a loop.
+    #: Note that the type is really ``list[tuple[tuple[int, ...], int]]``
+    #: but the persistence implementations don't handle tuples usefully.
+    num_added_iterations: list[list[list[int] | int]]
+    #: The parents of the loop.
+    parents: list[str]
+
+
+# TODO: This type looks familiar...
+class StoreCreationInfo(TypedDict):
+    """
+    Information about the creation of the persistence store.
+    """
+
+    #: Information about the application.
+    app_info: dict[str, Any]
+    #: When the persistence store was created.
+    create_time: str
+    #: The unique identifier for for the store/workflow.
+    id: str
+    #: User's name.
+    user_name: str
+    #: User's ORCID.
+    user_orcid: str
+    #: User's affiliations.
+    user_affiliations: list[str]
+
+
+class ElemMeta(TypedDict):
+    """
+    The kwargs supported for a StoreElement.
+    """
+
+    #: The ID of the element.
+    id_: int
+    #: The index of the element.
+    index: int
+    #: The index of the element in its element set.
+    es_idx: int
+    #: The indices of the element in the sequences that contain it.
+    seq_idx: dict[str, int]
+    #: The indices of the element's sources.
+    src_idx: dict[str, int]
+    #: The task associated with the element.
+    task_ID: int
+    #: The iteration IDs.
+    iteration_IDs: list[int]
+
+
+class IterMeta(TypedDict):
+    """
+    The kwargs supported for a StoreElementIter.
+    """
+
+    #: The index of the iteration.
+    data_idx: DataIndex
+    #: The EARs associated with the iteration.
+    EAR_IDs: dict[int, list[int]]
+    #: Whether the EARs have been initialised.
+    EARs_initialised: bool
+    #: The ID of the element.
+    element_ID: int
+    #: The loops containing the iteration.
+    loop_idx: dict[str, int]
+    #: The schema parameters being iterated over.
+    schema_parameters: list[str]
+
+
+class RunMeta(TypedDict):
+    """
+    The kwargs supported for StoreEAR.
+    """
+
+    #: The ID of the EAR.
+    id_: int
+    #: The ID of the element iteration containing the EAR.
+    elem_iter_ID: int
+    #: The index of the action that generated the EAR.
+    action_idx: int
+    #: The commands that the EAR will run.
+    commands_idx: list[int]
+    #: The data handled by the EAR.
+    data_idx: DataIndex
+    #: Metadata about the EAR.
+    metadata: Metadata | None
+    #: When the EAR ended, if known.
+    end_time: NotRequired[str | None]
+    #: The exit code of the EAR, if known.
+    exit_code: int | None
+    #: When the EAR started, if known.
+    start_time: NotRequired[str | None]
+    #: Working directory snapshot at start.
+    snapshot_start: dict[str, Any] | None
+    #: Working directory snapshot at end.
+    snapshot_end: dict[str, Any] | None
+    #: The index of the EAR in the submissions.
+    submission_idx: int | None
+    #: Where the EAR is set to run.
+    run_hostname: str | None
+    #: Whether the EAR succeeded, if known.
+    success: bool | None
+    #: The skip reason ID, if EAR was skipped.
+    skip: int
+    #: Port number used by ZeroMQ during execution.
+    port_number: int | None
+    #: Run ID whose command file can be used for this run (may be this run's ID).
+    commands_file_ID: int | None
+
+
+class TaskMeta(TypedDict):
+    """
+    Information about a task.
+    """
+
+    #: The ID of the task.
+    id_: int
+    #: The index of the task in the workflow.
+    index: int
+    #: The elements in the task.
+    element_IDs: list[int]
+
+
+class TemplateMeta(TypedDict):  # FIXME: Incomplete, see WorkflowTemplate
+    """
+    Metadata about a workflow template.
+    """
+
+    #: Descriptors for loops.
+    loops: list[dict]
+    #: Descriptors for tasks.
+    tasks: list[dict]
+
+
+class Metadata(TypedDict):
+    """
+    Workflow metadata.
+    """
+
+    #: Information about the store's creation.
+    creation_info: NotRequired[StoreCreationInfo]
+    #: Elements in the workflow.
+    elements: NotRequired[list[ElemMeta]]
+    #: Iterations in the workflow.
+    iters: NotRequired[list[IterMeta]]
+    #: Loops in the workflow.
+    loops: NotRequired[list[LoopDescriptor]]
+    #: The name of the workflow.
+    name: NotRequired[str]
+    #: The number of added tasks.
+    num_added_tasks: NotRequired[int]
+    #: The replacement workflow, if any.
+    replaced_workflow: NotRequired[str]
+    #: Element Action Runs in the workflow.
+    runs: NotRequired[list[RunMeta]]
+    #: Tasks in the workflow.
+    tasks: NotRequired[list[TaskMeta]]
+    #: The template that generated the workflow.
+    template: NotRequired[TemplateMeta]
+    #: Custom template components used.
+    template_components: NotRequired[dict[str, Any]]
+    #: Format for timestamps.
+    ts_fmt: NotRequired[str]
+    #: Format for timestamps used in naming.
+    ts_name_fmt: NotRequired[str]
+
+
+#: Information for looking up the type of a parameter.
+TypeLookup = DefaultDict[str, list[list]]
+
+
+class EncodedStoreParameter(TypedDict):
+    """
+    The encoding of a :class:`StoreParameter`.
+    """
+
+    #: The parameter data.
+    data: Any
+    #: Information for looking up the type.
+    type_lookup: TypeLookup
+
+
+class PersistenceCache(
+    TypedDict, Generic[AnySTask, AnySElement, AnySElementIter, AnySEAR, AnySParameter]
+):
+    """
+    Cache used internally by the persistence engine.
+    """
+
+    #: Tasks.
+    tasks: dict[int, AnySTask]
+    #: Elements.
+    elements: dict[int, AnySElement]
+    #: Element iterations.
+    element_iters: dict[int, AnySElementIter]
+    #: Element action runs.
+    EARs: dict[int, AnySEAR]
+    #: Parameter sources.
+    param_sources: dict[int, ParamSource]
+    #: Number of tasks.
+    num_tasks: int | None
+    #: Parameters.
+    parameters: dict[int, AnySParameter]
+    #: Number of element action runs.
+    num_EARs: int | None
+    #: Number of parameters.
+    num_params: int | None
+
+
+class ZarrAttrsDict(TypedDict):
+    """
+    Zarr workflow attributes descriptor.
+    """
+
+    #: Workflow name.
+    name: str
+    #: Timestamp format.
+    ts_fmt: str
+    #: Timestamp format for names.
+    ts_name_fmt: str
+    #: Information about the creation of the workflow and persistent store.
+    creation_info: StoreCreationInfo
+    #: The template used to build the workflow.
+    template: TemplateMeta
+    #: Custom components used to build the workflow.
+    template_components: dict[str, Any]
+    #: Number of tasks added.
+    num_added_tasks: int
+    #: Tasks in the workflow.
+    tasks: list[dict[str, Any]]
+    #: Loops in the workflow.
+    loops: list[dict[str, Any]]
+    #: Submissions by the workflow.
+    submissions: list[Mapping[str, JSONed]]
+    #: Replacement workflow, if any.
+    replaced_workflow: NotRequired[str]

hpcflow/sdk/persistence/utils.py

@@ -0,0 +1,73 @@
+"""
+Miscellaneous persistence-related helpers.
+"""
+
+from __future__ import annotations
+from getpass import getpass
+from typing import TYPE_CHECKING
+
+from hpcflow.sdk.core.errors import WorkflowNotFoundError
+
+if TYPE_CHECKING:
+    from typing import Callable, TypeVar
+    from fsspec import AbstractFileSystem  # type: ignore
+
+    T = TypeVar("T")
+
+
+def ask_pw_on_auth_exc(
+    f: Callable[..., T], *args, add_pw_to: str | None = None, **kwargs
+) -> tuple[T, str | None]:
+    """
+    Run the given function on the given arguments and add a password if the function
+    fails with an SSHException.
+    """
+    from paramiko.ssh_exception import SSHException
+
+    try:
+        out = f(*args, **kwargs)
+        pw = None
+
+    except SSHException:
+        pw = getpass()
+
+        if not add_pw_to:
+            kwargs["password"] = pw
+        else:
+            kwargs[add_pw_to] = {**kwargs[add_pw_to], "password": pw}
+
+        out = f(*args, **kwargs)
+
+    return out, pw
+
+
+def infer_store(path: str, fs: AbstractFileSystem) -> str:
+    """Identify the store type using the path and file system parsed by fsspec.
+
+    Parameters
+    ----------
+    fs
+        fsspec file system
+
+    """
+
+    # TODO: raise WorkflowNotFoundError if the path does not exist
+    # TODO: raise MalformedWorkflowError if a known store type cannot be inferred
+
+    # try to identify store type just from the path string:
+    if path.endswith(".zip"):
+        store_fmt = "zip"
+
+    elif path.endswith(".json"):
+        store_fmt = "json-single"
+
+    else:
+        # look at the directory contents:
+        if fs.glob(f"{path}/.zattrs"):
+            store_fmt = "zarr"
+        elif fs.glob(f"{path}/metadata.json"):
+            store_fmt = "json"
+        else:
+            raise WorkflowNotFoundError(path, fs)
+
+    return store_fmt