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

hpcflow/sdk/core/zarr_io.py

@@ -0,0 +1,206 @@
+"""
+Utilities for working with Zarr.
+"""
+
+from __future__ import annotations
+from typing import Any
+from typing_extensions import Self
+
+import zarr  # type: ignore
+import numpy as np
+
+from hpcflow.sdk.core.utils import get_in_container, get_relative_path, set_in_container
+
+
+#: The basic types that Zarr can handle directly with no special action.
+PRIMITIVES: tuple[type, ...] = (
+    int,
+    float,
+    str,
+    type(None),
+)
+
+#: Maximum nesting depth for encoding.
+MAX_DEPTH = 50
+
+
+def zarr_encode(data, zarr_group: zarr.Group, is_pending_add: bool, is_set: bool):
+    """
+    Encode data into a zarr group.
+    """
+
+    encoded: list[dict] = []
+
+    def encode(obj: Any, path: list) -> Any:
+        if len(path) > MAX_DEPTH:
+            raise RuntimeError("I'm in too deep!")
+
+        if isinstance(obj, ZarrEncodable):
+            return encode(obj.to_dict(), path)
+        elif isinstance(obj, (list, tuple, set)):
+            out = (encode(item, [*path, idx]) for idx, item in enumerate(obj))
+            if isinstance(obj, tuple):
+                return tuple(out)
+            elif isinstance(obj, set):
+                return set(out)
+            else:
+                return list(out)
+        elif isinstance(obj, dict):
+            return {
+                dct_key: encode(dct_val, [*path, dct_key])
+                for dct_key, dct_val in obj.items()
+            }
+        elif isinstance(obj, PRIMITIVES):
+            return obj
+        elif isinstance(obj, np.ndarray):
+            new_name = str(max((int(i) + 1 for i in zarr_group.keys()), default=0))
+            zarr_group.create_dataset(name=new_name, data=obj)
+            encoded.append(
+                {
+                    "path": path,
+                    "dataset": new_name,
+                }
+            )
+            return None
+        else:
+            raise ValueError(f"unserializable type: {type(obj)}")
+
+    zarr_group.attrs["data"] = encode(data, [])
+    zarr_group.attrs["encoded"] = encoded
+    zarr_group.attrs["is_set"] = is_set
+    if is_pending_add:
+        zarr_group.attrs["is_pending_add"] = is_pending_add
+
+
+# TODO: need to separate so this func doesn't need matflow
+
+
+def _zarr_encode_NEW(
+    obj: Any,
+    root_group: zarr.Group,
+    arr_path: str,
+) -> tuple[Any, list[list]]:
+    """
+    Save arbitrarily-nested Python-primitive, `ZarrEncodable` and numpy array objects into
+    Zarr.
+
+    Parameters
+    ----------
+    obj:
+        Object to encode.
+    root_group:
+        Parent Zarr group into which new Zarr arrays will be added (at `arr_path`).
+    arr_path:
+        Path relative to `root_group` into which new Zarr arrays will be added.
+
+    Returns
+    -------
+    data
+        The encoded data.
+    arr_lookup
+        How to look up where to rebuild Numpy arrays.
+    """
+
+    arr_lookup: list[list] = []
+
+    def encode(obj: Any, path: list) -> Any:
+        if len(path) > MAX_DEPTH:
+            raise RuntimeError("I'm in too deep!")
+
+        if isinstance(obj, ZarrEncodable):
+            return encode(obj.to_dict(), path)
+        elif isinstance(obj, (list, tuple, set)):
+            items = (encode(item, [*path, idx]) for idx, item in enumerate(obj))
+            if isinstance(obj, tuple):
+                return tuple(items)
+            elif isinstance(obj, set):
+                return set(items)
+            else:
+                return list(items)
+        elif isinstance(obj, dict):
+            return {key: encode(val, [*path, key]) for key, val in obj.items()}
+        elif isinstance(obj, PRIMITIVES):
+            return obj
+        elif isinstance(obj, np.ndarray):
+            # Might need to generate new group:
+            param_arr_group = root_group.require_group(arr_path)
+            new_idx = max((int(i) + 1 for i in param_arr_group.keys()), default=0)
+            param_arr_group.create_dataset(name=f"arr_{new_idx}", data=obj)
+            arr_lookup.append([path, new_idx])
+            return None
+        else:
+            raise ValueError(f"unserializable type: {type(obj)}")
+
+    return encode(obj, []), arr_lookup
+
+
+def zarr_decode(
+    param_data: None | dict,
+    arr_group: zarr.Group,
+    path: list | None = None,
+    dataset_copy: bool = False,
+):
+    """
+    Decode data from a zarr group.
+    """
+    if param_data is None:
+        return None
+
+    path = path or []
+
+    data = get_in_container(param_data["data"], path)
+    # data = copy.deepcopy(data)  # TODO: why did we copy?
+
+    for arr_path, arr_idx in param_data["arr_lookup"]:
+        try:
+            rel_path = get_relative_path(arr_path, path)
+        except ValueError:
+            continue
+
+        dataset = arr_group.get(f"arr_{arr_idx}")
+
+        if dataset_copy:
+            dataset = dataset[:]
+
+        if rel_path:
+            set_in_container(data, rel_path, dataset)
+        else:
+            data = dataset
+
+    return data
+
+
+class ZarrEncodable:
+    """
+    Base class of data that can be converted to and from zarr form.
+    """
+
+    _typ = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert this object to a dict.
+        """
+        if hasattr(self, "__dict__"):
+            return dict(self.__dict__)
+        elif hasattr(self, "__slots__"):
+            return {k: getattr(self, k) for k in self.__slots__}
+        else:
+            # Should be unreachable
+            return {}
+
+    def to_zarr(self, zarr_group: zarr.Group):
+        """
+        Save this object into the given zarr group.
+        """
+        zarr_encode(self.to_dict(), zarr_group, is_pending_add=False, is_set=False)
+
+    @classmethod
+    def from_zarr(cls, zarr_group: zarr.Group, dataset_copy: bool = False) -> Self:
+        """
+        Read an instance of this class from the given zarr group.
+        """
+        # FIXME: Do the read of the data!
+        param_data = None
+        data = zarr_decode(param_data, zarr_group, dataset_copy=dataset_copy)
+        return cls(**data)

hpcflow/sdk/data/__init__.py

@@ -0,0 +1,13 @@
+"""
+YAML schemas.
+
+Contents:
+
+* ``config_file_schema.yaml`` – Schema for configuration selection files.
+* ``config_schema.yaml`` – Schema for configuration files.
+* ``environments_spec_schema.yaml`` – Schema for execution environment definition files.
+* ``files_spec_schema.yaml`` – Schema for input/output specification files.
+* ``parameters_spec_schema.yaml`` – Schema for parameter specification files.
+* ``task_schema_spec_schema.yaml`` – Schema for task template specification files.
+* ``workflow_spec_schema.yaml`` – Schema for workflow files.
+"""

hpcflow/sdk/data/config_file_schema.yaml

@@ -0,0 +1,34 @@
+rules:
+  - path: []
+    condition:
+      and:
+        - value.required_keys:
+            - configs
+        - value.allowed_keys:
+            - configs
+
+  - path: [configs]
+    doc: >
+      A mapping of string config "invocation keys" to configuration settings.
+    condition:
+      and:
+        - value.type.equal_to: dict
+        - value.keys_is_instance: [str]
+
+  - path: [configs, { type: map_value }]
+    condition:
+      and:
+        - value.type.equal_to: dict
+        - value.required_keys:
+            - invocation
+            - config
+
+  - path: [configs, { type: map_value }, invocation]
+    condition:
+      value.allowed_keys:
+        - environment_setup
+        - match
+
+  - path: [configs, { type: map_value }, invocation, match]
+    condition:
+      value.type.equal_to: dict

hpcflow/sdk/data/config_schema.yaml

@@ -0,0 +1,260 @@
+rules:
+  - path: []
+    condition:
+      value.allowed_keys:
+        - machine
+        - user_name
+        - user_orcid
+        - user_affiliations
+        - linux_release_file
+        - log_file_path
+        - log_file_level
+        - log_console_level
+        - task_schema_sources
+        - parameter_sources
+        - command_file_sources
+        - environment_sources
+        - default_scheduler
+        - default_shell
+        - schedulers
+        - shells
+        - demo_data_dir
+        - demo_data_manifest_file
+        - default_workflow_path
+        - workflow_name_use_dir
+        - workflow_name_add_timestamp
+  - path: []
+    condition:
+      value.required_keys:
+        - machine
+        - task_schema_sources
+        - parameter_sources
+        - command_file_sources
+        - environment_sources
+        - default_scheduler
+        - default_shell
+        - schedulers
+  - path: []
+    condition:
+      value.forbidden_keys:
+        - config_directory # Config metadata
+        - config_file_name # Config metadata
+        - config_file_path # Config metadata
+        - config_file_contents # Config metadata
+        - config_key # Config metadata
+        - config_schemas # Config metadata
+        - invoking_user_id # Config metadata
+        - host_user_id # Config metadata
+        - host_user_id_file_path # Config metadata
+        - config_callback # Config method
+        - get_all # Config method
+        - get # Config method
+        - set # Config method
+        - _set # Config method
+        - unset # Config method
+        - save # Config method
+        - reload # Config method
+
+  - path: [machine]
+    doc: >
+      A label that references the current machine. By default, this uses the return from
+      the Python function `socket.gethostname()`.
+    condition:
+      value.type.equal_to: str
+
+  - path: [user_name]
+    condition:
+      value.type.equal_to: str
+
+  - path: [user_orcid]
+    condition:
+      value.type.equal_to: str
+
+  - path: [user_affiliations]
+    condition:
+      value.type.equal_to: list
+
+  - path: [user_affiliations, { type: list_value }]
+    condition:
+      value.is_instance: [str]
+
+  - path: [linux_release_file]
+    condition:
+      value.type.equal_to: str
+
+  - path: [log_file_path]
+    doc:
+      - >
+        File path to the app log file. The variables `<<app_name>>` and `<<app_version>>`
+        are available to be used in this option, which resolve to the app name and version,
+        respectively.
+      - >
+        The file path may be nested in non-existent directories, in which case
+        those directories will be generated. If specified as a relative path, the path will
+        be resolved relative to the config directory (the directory that contains the config
+        file).
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [log_file_level]
+    condition:
+      and:
+        - value.type.equal_to: str
+        - value.in: [notset, debug, info, warning, error, critical]
+
+  - path: [log_console_level]
+    condition:
+      and:
+        - value.type.equal_to: str
+        - value.in: [notset, debug, info, warning, error, critical]
+
+  - path: [task_schema_sources]
+    condition:
+      value.type.equal_to: list
+
+  - path: [environment_sources]
+    condition:
+      value.type.equal_to: list
+
+  - path: [parameter_sources]
+    condition:
+      value.type.equal_to: list
+
+  - path: [command_file_sources]
+    condition:
+      value.type.equal_to: list
+
+  - path: [task_schema_sources, { type: list_value }]
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [environment_sources, { type: list_value }]
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [parameter_sources, { type: list_value }]
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [command_file_sources, { type: list_value }]
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [default_scheduler]
+    condition:
+      value.in: [direct, sge, slurm]
+
+  - path: [default_shell]
+    condition:
+      value.in: [powershell, bash, wsl, wsl+bash]
+
+  - path: [schedulers]
+    condition:
+      value.allowed_keys: [direct, direct_posix, sge, slurm]
+
+  - path: [schedulers, { type: map_value }]
+    condition:
+      value.allowed_keys: [defaults, parallel_environments, partitions] # TODO split these up based on scheduler type
+
+  - path: [schedulers, { type: map_value }, defaults]
+    condition:
+      value.allowed_keys: # TODO: split these up based on scheduler type
+        - shebang_executable
+        - directives # TODO: only actually allowed in `QueuedScheduler` sub-classes
+        - options # deprecated, replaced by directives
+        - submit_cmd
+        - show_cmd
+        - del_cmd
+        - js_cmd
+        - array_switch
+        - array_item_var
+        - cwd_switch
+
+  - path: [schedulers, sge]
+    condition:
+      value.type.equal_to: dict
+
+  - path: [schedulers, sge, parallel_environments]
+    condition:
+      value.type.equal_to: dict
+
+  - path: [schedulers, sge, parallel_environments, { type: map_value }]
+    condition:
+      and:
+        - value.allowed_keys:
+            - num_cores
+            - num_nodes
+            - parallel_modes
+        - value.required_keys:
+            - num_cores
+
+  - path: [schedulers, slurm]
+    condition:
+      value.type.equal_to: dict
+
+  - path: [schedulers, slurm, partitions]
+    condition:
+      value.type.equal_to: dict
+
+  - path: [schedulers, slurm, partitions, { type: map_value }]
+    condition:
+      and:
+        - value.allowed_keys:
+            - num_cores
+            - num_nodes
+            - num_cores_per_node
+            - parallel_modes
+        - value.required_keys:
+            - num_cores
+
+  - path: [shells]
+    condition:
+      value.allowed_keys: [powershell, bash, wsl, wsl+bash]
+
+  - path: [shells, { type: map_value }]
+    condition:
+      value.allowed_keys: [defaults]
+
+  - path: [shells, { type: map_value }, defaults]
+    condition:
+      value.allowed_keys:
+        - executable
+        - executable_args
+        - os_args
+        - WSL_executable
+        - WSL_distribution
+        - WSL_user
+
+  - path: [demo_data_dir]
+    doc: >
+      If provided, this will be used as the path to the directory containing demo data
+      files. This can be a local path or an fsspec URL.
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [demo_data_manifest_file]
+    doc: >
+      If provided, this will be used as the path to the manifest for the built-in example 
+      data files. It should point to a JSON file whose top-level keys are example data
+      file names. Each value should be an empty map/dict or a map like: `{"in_zip": "text_file.zip"}`,
+      where `text_file.zip` is the name of a zip archive that contains the example data
+      file. This can be a local path or an fsspec URL.
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [default_workflow_path]
+    doc: Default directory path in which new workflows will be generated.
+    condition:
+      value.is_instance: [str, path]
+
+  - path: [workflow_name_add_timestamp]
+    doc: By default, whether to suffix the workflow name with a date-timestamp.
+    condition:
+      value.is_instance: [bool]
+
+  - path: [workflow_name_add_timestamp]
+    doc: >
+      By default, whether to use a directory named according to the workflow name, and 
+      a separate child directory named according to the date-timestamp if present.
+    condition:
+      value.is_instance: [bool]

hpcflow/sdk/data/environments_spec_schema.yaml

@@ -0,0 +1,21 @@
+rules:
+  - path: []
+    condition:
+      {
+        value.allowed_keys: [name, setup, executables, specifiers, doc, _hash_value],
+      }
+  - path: [setup]
+    condition:
+      or:
+        - value.type.equal_to: str
+        - value.null: null
+  - path: [executables]
+    condition: { value.type.equal_to: list }
+  - path: [executables, { type: list_value }]
+    condition: { value.allowed_keys: [label, instances] }
+  - path: [executables, { type: list_value }, instances]
+    condition: { value.type.equal_to: list }
+  - path: [executables, { type: list_value }, instances, { type: list_value }]
+    condition: { value.type.equal_to: dict }
+  - path: [executables, { type: list_value }, instances, { type: list_value }]
+    condition: { value.allowed_keys: [command, num_cores, parallel_mode] }

hpcflow/sdk/data/files_spec_schema.yaml

@@ -0,0 +1,5 @@
+rules:
+  - path: []
+    condition: { value.allowed_keys: [label, name, doc, _hash_value] }
+  - path: [name]
+    condition: { value.allowed_keys: [name, args, is_regex] }

hpcflow/sdk/data/parameters_spec_schema.yaml

@@ -0,0 +1,7 @@
+rules:
+  - path: []
+    condition:
+      {
+        value.allowed_keys:
+          [type, sub_parameters, is_file, _hash_value, _validation],
+      }

hpcflow/sdk/data/task_schema_spec_schema.yaml

@@ -0,0 +1,3 @@
+rules:
+  - path: [inputs, { type: list_value }]
+    condition: { value.type.equal_to: dict }

hpcflow/sdk/data/workflow_spec_schema.yaml

@@ -0,0 +1,22 @@
+rules:
+  - path: []
+    condition:
+      value.allowed_keys:
+        - doc
+        - name
+        - source_file
+        - resources
+        - config
+        - environments
+        - env_presets
+        - template_components
+        - tasks
+        - meta_tasks
+        - loops
+        - store_kwargs
+        - merge_resources
+        - merge_envs
+        - workflow
+
+  - path: [tasks]
+    condition: { value.type.equal_to: list }

hpcflow/sdk/demo/__init__.py

@@ -0,0 +1,3 @@
+"""
+Demonstration code.
+"""