Repository: tomwhite/cubed
Pull request #188 (2023-05-26 18:10:42)

Convert str to bytes
This is just a small quality-of-life improvement that allows passing memory arguments to `Spec` as strings with units, e.g. `allowed_mem='100MB'` rather than `allowed_mem=100_000_000`.

I went down a little rabbit hole learning about how there [isn't a particularly neat way](https://florimond.dev/en/posts/2018/10/reconciling-dataclasses-and-properties-in-python/) of doing input validation on a frozen `dataclass`, and ended up just removing the dataclass in favour of a normal class. 

The `Spec` object is now immutable. It could be mutable, like it technically was before, but I wasn't sure what the intended use was. If `Spec` is just intended as a container of kwargs then it should really be immutable (or even a `NamedTuple`, though that has similar problems with input validation).


Patch:
diff --git a/cubed/core/array.py b/cubed/core/array.py
index 52cc98e4..60aed17a 100644
--- a/cubed/core/array.py
+++ b/cubed/core/array.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from operator import mul
-from typing import Optional, TypeVar
+from typing import Optional, TypeVar, Union
 from warnings import warn
 
 import numpy as np
@@ -8,7 +8,7 @@ from toolz import map, reduce
 
 from cubed.runtime.pipeline import already_computed
 from cubed.runtime.types import Executor
-from cubed.utils import chunk_memory
+from cubed.utils import chunk_memory, convert_to_bytes
 from cubed.vendor.dask.array.core import normalize_chunks
 
 from .plan import arrays_to_plan
@@ -186,42 +186,107 @@ class CoreArray:
         return f"cubed.core.CoreArray<{self.name}, shape={self.shape}, dtype={self.dtype}, chunks={self.chunks}>"
 
 
-@dataclass
 class Spec:
     """Specification of resources available to run a computation."""
 
-    work_dir: Optional[str] = None
-    """The directory path (specified as an fsspec URL) used for storing intermediate data."""
+    def __init__(
+        self,
+        work_dir: Union[str, None] = None,
+        max_mem: Union[int, None] = None,
+        allowed_mem: Union[int, str, None] = None,
+        reserved_mem: Union[int, str, None] = 0,
+        executor: Union[Executor, None] = None,
+        storage_options: Union[dict, None] = None,
+    ):
+        """
+        Specify resources available to run a computation.
 
-    max_mem: Optional[int] = None
-    """**Deprecated**. The maximum memory available to a worker for data use for the computation, in bytes."""
+        Parameters
+        ----------
+        work_dir : str or None
+            The directory path (specified as an fsspec URL) used for storing intermediate data.
+        max_mem : int, optional
+            **Deprecated**. The maximum memory available to a worker for data use for the computation, in bytes.
+        allowed_mem : int or str, optional
+            The total memory available to a worker for running a task, in bytes.
+
+            If int it should be >=0. If str it should be of form <value><unit> where unit can be kB, MB, GB, TB etc.
+            This includes any ``reserved_mem`` that has been set.
+        reserved_mem : int or str, optional
+            The memory reserved on a worker for non-data use when running a task, in bytes.
+
+            If int it should be >=0. If str it should be of form <value><unit> where unit can be kB, MB, GB, TB etc.
+        executor : Executor, optional
+            The default executor for running computations.
+        storage_options : dict, optional
+            Storage options to be passed to fsspec.
+        """
 
-    allowed_mem: Optional[int] = None
-    """The total memory available to a worker for running a task, in bytes.
+        if max_mem is not None:
+            warn(
+                "`max_mem` is deprecated, please use `allowed_mem` instead",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
-    This includes any ``reserved_mem`` that has been set."""
+        self._work_dir = work_dir
 
-    reserved_mem: int = 0
-    """The memory reserved on a worker for non-data use when running a task, in bytes.
+        self._reserved_mem = convert_to_bytes(reserved_mem)
+        if allowed_mem is None:
+            self._allowed_mem = (max_mem or 0) + self.reserved_mem
+        else:
+            self._allowed_mem = convert_to_bytes(allowed_mem)
 
-    See ``cubed.measure_reserved_memory``.
-    """
+        self._executor = executor
+        self._storage_options = storage_options
 
-    executor: Executor = None
-    """The default executor for running computations."""
+    @property
+    def work_dir(self) -> str:
+        """The directory path (specified as an fsspec URL) used for storing intermediate data."""
+        return self._work_dir
 
-    storage_options: dict = None  # type: ignore[assignment]
-    """Storage options to be passed to fsspec"""
+    @property
+    def allowed_mem(self) -> int:
+        """
+        The total memory available to a worker for running a task, in bytes.
 
-    def __post_init__(self):
-        if self.max_mem is not None:
-            warn(
-                "`max_mem` is deprecated, please use `allowed_mem` instead",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-        if self.allowed_mem is None:
-            self.allowed_mem = (self.max_mem or 0) + self.reserved_mem
+        This includes any ``reserved_mem`` that has been set.
+        """
+        return self._allowed_mem
+
+    @property
+    def reserved_mem(self) -> int:
+        """
+        The memory reserved on a worker for non-data use when running a task, in bytes.
+
+        See Also
+        --------
+        cubed.measure_reserved_memory
+        """
+        return self._reserved_mem
+
+    @property
+    def executor(self) -> Executor:
+        """The default executor for running computations."""
+        return self._executor
+
+    @property
+    def storage_options(self) -> dict:
+        """Storage options to be passed to fsspec."""
+        return self._storage_options
+
+    def __repr__(self) -> str:
+        return f"cubed.Spec(work_dir={self._work_dir}, allowed_mem={self._allowed_mem}, " \
+               f"reserved_mem={self._reserved_mem}, executor={self._executor}, storage_options={self._storage_options})"
+
+    def __eq__(self, other):
+        if isinstance(other, Spec):
+            return self.work_dir == other.work_dir \
+                and self.allowed_mem == other.allowed_mem \
+                and self.reserved_mem == other.reserved_mem and self.executor == other.executor \
+                and self.storage_options == other.storage_options
+        else:
+            return False
 
 
 class Callback:
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
index 76186f18..9b30cbc2 100644
--- a/cubed/runtime/executors/modal.py
+++ b/cubed/runtime/executors/modal.py
@@ -11,7 +11,7 @@ from cubed.core.plan import visit_nodes
 from cubed.runtime.types import DagExecutor
 from cubed.utils import peak_measured_mem
 
-stub = modal.Stub()
+stub = modal.Stub("sync-stub")
 
 requirements_file = os.getenv("CUBED_MODAL_REQUIREMENTS_FILE")
 
diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
index 8935de43..e62ca318 100644
--- a/cubed/runtime/executors/modal_async.py
+++ b/cubed/runtime/executors/modal_async.py
@@ -15,7 +15,7 @@ from cubed.runtime.backup import should_launch_backup
 from cubed.runtime.types import DagExecutor
 from cubed.utils import peak_measured_mem
 
-async_stub = modal.aio.AioStub()
+async_stub = modal.aio.AioStub("async-stub")
 
 requirements_file = os.getenv("CUBED_MODAL_REQUIREMENTS_FILE")
 
diff --git a/cubed/utils.py b/cubed/utils.py
index 656f0b21..cb38d487 100644
--- a/cubed/utils.py
+++ b/cubed/utils.py
@@ -157,3 +157,41 @@ def extract_stack_summaries(frame, limit=None):
     stack_summaries.reverse()
 
     return stack_summaries
+
+
+def convert_to_bytes(size: Union[int, str]) -> int:
+    """
+    Converts the input data size to bytes.
+
+    The data size can be expressed as an integer or as a string with different SI prefixes such as '500kB', '2MB', or '1GB'.
+
+    Parameters
+    ----------
+    size: in or str:
+        Size of data. If int it should be >=0. If str it should be of form <value><unit> where unit can be kB, MB, GB, TB etc.
+
+    Returns
+    -------
+    int: The size in bytes
+    """
+    units = {"B": 0, "kB": 1, "MB": 2, "GB": 3, "TB": 4, "PB": 5}
+
+    if isinstance(size, int) and size >= 0:
+        return size
+    elif isinstance(size, str):
+        # check if the format is valid
+        if size[-1] == "B" and size[:-1].isdigit():
+            unit = "B"
+            value = size[:-1]
+        elif size[-2:] in units and size[:-2].isdigit():
+            unit = size[-2:]
+            value = size[:-2]
+        else:
+            raise ValueError(f"Invalid value: {size}. Expected a string ending with an SI prefix.")
+
+        if unit in units and value.isdigit():
+            value = int(value)
+            # convert to bytes
+            return value * (1000 ** units[unit])
+    else:
+        raise ValueError(f"Invalid value: {size}. Expected a positive integer or a string ending with an SI prefix.")


Test patch:
diff --git a/cubed/tests/runtime/test_modal_async.py b/cubed/tests/runtime/test_modal_async.py
index 1336dace..8101e65a 100644
--- a/cubed/tests/runtime/test_modal_async.py
+++ b/cubed/tests/runtime/test_modal_async.py
@@ -24,7 +24,7 @@ def write_int_to_file(path, i):
 tmp_path = "s3://cubed-unittest/map_unordered"
 
 
-stub = modal.aio.AioStub()
+stub = modal.aio.AioStub("async-test-stub")
 
 image = modal.Image.debian_slim().pip_install(
     [
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 036e8cf7..fd2e384f 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -248,6 +248,44 @@ def test_different_specs(tmp_path):
         xp.add(a, b)
 
 
+class TestSpecMemArgTypes:
+    def test_max_mem_deprecation_warning(self):
+        # Remove once max_mem fully deprecated in favour of allowed_mem
+        with pytest.warns(DeprecationWarning, match="`max_mem` is deprecated, please use `allowed_mem` instead"):
+            cubed.Spec(max_mem=100_000)
+
+    @pytest.mark.parametrize(
+        "input_value, expected_value",
+        [
+            (500, 500),
+            (100_000, 100_000),
+            ("500B", 500),
+            ("1kB", 1000),
+            ("1MB", 1000 ** 2),
+            ("1GB", 1000 ** 3),
+            ("1TB", 1000 ** 4),
+            ("1PB", 1000 ** 5)
+        ]
+    )
+    def test_convert_to_bytes(self, input_value, expected_value):
+        spec = cubed.Spec(allowed_mem=input_value)
+        assert spec.allowed_mem == expected_value
+
+    @pytest.mark.parametrize(
+        "input_value",
+        [
+            "1EB",  # EB is not a valid unit in this function
+            "1kb",  # lower-case k is not valid
+            "invalid",  # completely invalid input
+            -512,  # negative integer
+            1000.0,  # invalid type
+        ]
+    )
+    def test_convert_to_bytes_error(self, input_value):
+        with pytest.raises(ValueError):
+            cubed.Spec(allowed_mem=input_value)
+
+
 def test_reduction_multiple_rounds(tmp_path, executor):
     spec = cubed.Spec(tmp_path, allowed_mem=1000)
     a = xp.ones((100, 10), dtype=np.uint8, chunks=(1, 10), spec=spec)