hpcflow-new2 0.2.0a221__tar.gz → 0.2.0a223__tar.gz

{hpcflow_new2-0.2.0a221 → hpcflow_new2-0.2.0a223}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hpcflow-new2
-Version: 0.2.0a221
+Version: 0.2.0a223
 Summary: Computational workflow management
 License: MPL-2.0
 Author: aplowman

hpcflow_new2-0.2.0a223/hpcflow/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0a223"

{hpcflow_new2-0.2.0a221 → hpcflow_new2-0.2.0a223}/hpcflow/sdk/persistence/zarr.py

@@ -18,6 +18,7 @@ from numpy.ma.core import MaskedArray
 import zarr  # type: ignore
 from zarr.errors import BoundsCheckError  # type: ignore
 from zarr.storage import DirectoryStore, FSStore  # type: ignore
+from zarr.util import guess_chunks  # type: ignore
 from fsspec.implementations.zip import ZipFileSystem  # type: ignore
 from rich.console import Console
 from numcodecs import MsgPack, VLenArray, blosc, Blosc, Zstd  # type: ignore
@@ -59,6 +60,7 @@ from hpcflow.sdk.submission.submission import (
     SUBMISSION_SUBMIT_TIME_KEYS,
 )
 from hpcflow.sdk.utils.arrays import get_2D_idx, split_arr
+from hpcflow.sdk.utils.patches import override_module_attrs
 from hpcflow.sdk.utils.strings import shorten_list_str
 
 if TYPE_CHECKING:
@@ -89,6 +91,11 @@ if TYPE_CHECKING:
 ListAny: TypeAlias = "list[Any]"
 #: Zarr attribute mapping context.
 ZarrAttrs: TypeAlias = "dict[str, Any]"
+#: Soft lower limit for the number of bytes in an array chunk
+_ARRAY_CHUNK_MIN: int = 500 * 1024 * 1024  # 500 MiB
+#: Hard upper limit for the number of bytes in an array chunk. Should be lower than the
+#: maximum buffer size of the blosc encoder, if we're using it (2 GiB)
+_ARRAY_CHUNK_MAX: int = 1024 * 1024 * 1024  # 1 GiB
 _JS: TypeAlias = "dict[str, list[dict[str, dict]]]"
 
 
@@ -124,7 +131,13 @@ def _encode_numpy_array(
     new_idx = (
         max((int(i.removeprefix("arr_")) for i in param_arr_group.keys()), default=-1) + 1
     )
-    param_arr_group.create_dataset(name=f"arr_{new_idx}", data=obj, chunks=obj.shape)
+    with override_module_attrs(
+        "zarr.util", {"CHUNK_MIN": _ARRAY_CHUNK_MIN, "CHUNK_MAX": _ARRAY_CHUNK_MAX}
+    ):
+        # `guess_chunks` also ensures chunk shape is at least 1 in each dimension:
+        chunk_shape = guess_chunks(obj.shape, obj.dtype.itemsize)
+
+    param_arr_group.create_dataset(name=f"arr_{new_idx}", data=obj, chunks=chunk_shape)
     type_lookup["arrays"].append([path, new_idx])
 
     return len(type_lookup["arrays"]) - 1

hpcflow_new2-0.2.0a223/hpcflow/sdk/utils/patches.py

@@ -0,0 +1,31 @@
+from contextlib import contextmanager
+from pathlib import Path
+import sys
+from typing import Any
+
+
+def resolve_path(path):
+    """On Windows Python 3.8, 3.9, and 3.10, `Pathlib.resolve` does
+    not return an absolute path for non-existant paths, when it should.
+
+    See: https://github.com/python/cpython/issues/82852
+
+    """
+    # TODO: this only seems to be used in a test; remove?
+    return Path.cwd() / Path(path).resolve()  # cwd is ignored if already absolute
+
+
+@contextmanager
+def override_module_attrs(module_name: str, overrides: dict[str, Any]):
+    """Context manager to temporarily override module-level attributes. The module must be
+    imported (i.e. within `sys.modules`)."""
+
+    module = sys.modules[module_name]
+    original_values = {k: getattr(module, k) for k in overrides}
+    try:
+        for k, v in overrides.items():
+            setattr(module, k, v)
+        yield
+    finally:
+        for k, v in original_values.items():
+            setattr(module, k, v)

{hpcflow_new2-0.2.0a221 → hpcflow_new2-0.2.0a223}/hpcflow/tests/unit/test_persistence.py

@@ -6,7 +6,11 @@ from typing import cast, TYPE_CHECKING
 import numpy as np
 import zarr  # type: ignore
 import pytest
-from hpcflow.sdk.core.test_utils import make_test_data_YAML_workflow, make_workflow
+from hpcflow.sdk.core.test_utils import (
+    make_schemas,
+    make_test_data_YAML_workflow,
+    make_workflow,
+)
 from hpcflow.sdk.persistence.json import (
     JSONPersistentStore,
     JsonStoreElement,
@@ -551,3 +555,38 @@ def test_zarr_encode_decode_jobscript_block_dependencies_large_one_to_one():
     arr = ZarrPersistentStore._encode_jobscript_block_dependencies(deps_t)
     deps_rt = ZarrPersistentStore._decode_jobscript_block_dependencies(arr)
     assert deps_rt == deps
+
+
+@pytest.mark.parametrize(
+    "array",
+    (
+        np.array([]),
+        np.empty(0),
+        np.empty((0, 1, 2)),
+        np.array([1, 2, 3]),
+        np.array([[1, 2, 3], [4, 5, 6]]),
+    ),
+)
+def test_zarr_save_persistent_array_shape(null_config, tmp_path, array):
+    s1 = make_schemas(({"p1": None}, ()))
+    t1 = hf.Task(schema=s1, inputs={"p1": array})
+    wk = hf.Workflow.from_template_data(
+        template_name="test_save_empty_array",
+        tasks=[t1],
+        path=tmp_path,
+    )
+    assert array.shape == wk.tasks[0].elements[0].get("inputs.p1")[:].shape
+
+
+def test_zarr_single_chunk_threshold(null_config, tmp_path):
+    # test very large arrays (> ~1 GB) are saved using multiple chunks
+    array = np.zeros(
+        268_435_456
+    )  # ~ 2.147483647 GB; greater than blosc's max buffer size
+    s1 = make_schemas(({"p1": None}, ()))
+    t1 = hf.Task(schema=s1, inputs={"p1": array})
+    wk = hf.Workflow.from_template_data(
+        template_name="test_large_array",
+        tasks=[t1],
+        path=tmp_path,
+    )

{hpcflow_new2-0.2.0a221 → hpcflow_new2-0.2.0a223}/pyproject.toml

@@ -1,7 +1,7 @@
 
 [tool.poetry]
 name = "hpcflow-new2"
-version = "0.2.0a221"
+version = "0.2.0a223"
 
 description = "Computational workflow management"
 authors = ["aplowman <adam.plowman@manchester.ac.uk>"]
@@ -100,7 +100,7 @@ hook-dirs = "hpcflow.__pyinstaller:get_hook_dirs"
 
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.2.0a221"
+version = "0.2.0a223"
 tag_format = "v$version"
 version_files = [ 
     "pyproject.toml:version",

hpcflow_new2-0.2.0a221/hpcflow/_version.py

@@ -1 +0,0 @@
-__version__ = "0.2.0a221"

hpcflow_new2-0.2.0a221/hpcflow/sdk/utils/patches.py

@@ -1,12 +0,0 @@
-from pathlib import Path
-
-
-def resolve_path(path):
-    """On Windows Python 3.8, 3.9, and 3.10, `Pathlib.resolve` does
-    not return an absolute path for non-existant paths, when it should.
-
-    See: https://github.com/python/cpython/issues/82852
-
-    """
-    # TODO: this only seems to be used in a test; remove?
-    return Path.cwd() / Path(path).resolve()  # cwd is ignored if already absolute