Repository: tomwhite/cubed
Pull request #190 (2023-05-30 11:33:12)

Improvements and fixes to Modal examples.
Add matplotlib/seaborn dependencies to diagnostics to make timeline vizualisations work.

tomwhite: Here's an example timeline when running `modal-matmul-random.py`
![1685445298_timeline](https://github.com/tomwhite/cubed/assets/85085/6f8ccdd8-6d38-4047-97ac-cf88e0874e5c)


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/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.")
diff --git a/examples/README.md b/examples/README.md
index f53f4dcb..42f05a13 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -2,14 +2,14 @@
 
 ## Which cloud service should I use?
 
-**Modal** is the easiest to get started with because it handles building a runtime environment for you automatically (note that it requires that you [sign up](https://modal.com/signup) for a free account). 
-It has been tested with ~50 workers.
+**Modal** is the easiest to get started with because it handles building a runtime environment for you automatically (note that it requires that you [sign up](https://modal.com/signup) for a free account).
+It has been tested with ~300 workers.
 
 **Lithops** requires slightly more work to get started since you have to build a runtime environment first.
 Lithops has support for many serverless services on various cloud providers, but has so far been tested on two:
- 
 
-- **AWS lambda** requires building a docker container first, but has been tested with hundreds of workers. 
+
+- **AWS lambda** requires building a docker container first, but has been tested with hundreds of workers.
 - **Google Cloud Functions** only requires building a Lithops runtime, which can be created from a pip-style `requirements.txt` without docker. Large-scale testing is ongoing.
 
 **Google Cloud Dataflow** is relatively straightforward to get started with. It has the highest overhead for worker startup (minutes compared to seconds for Modal or Lithops), and although it has only been tested with ~20 workers, it is the most mature service and therefore should be reliable for much larger computations.
diff --git a/examples/lithops/aws-lambda/README.md b/examples/lithops/aws-lambda/README.md
index 7470f354..d616fa7b 100644
--- a/examples/lithops/aws-lambda/README.md
+++ b/examples/lithops/aws-lambda/README.md
@@ -36,8 +36,38 @@ python lithops-add-asarray.py "s3://cubed-$USER-temp" cubed-runtime
 
 If successful it should print a 4x4 matrix.
 
-Run the other example in a similar way
+Run the other examples in a similar way
 
 ```shell
 python lithops-add-random.py "s3://cubed-$USER-temp" cubed-runtime
 ```
+
+and
+
+```shell
+python lithops-matmul-random.py "s3://cubed-$USER-temp" cubed-runtime
+```
+
+These will take longer to run as they operate on more data.
+
+
+The last two examples use `TimelineVisualizationCallback` which produce a plot showing the timeline of events in the task lifecycle.
+The plots are `png` files and are written in the `plots` directory with a timestamp. Open the latest one with
+
+```shell
+open plots/$(ls plots | tail -1)
+```
+
+## Cleaning up
+
+If you want to rebuild the Lithops runtime image you can delete the existing one by running
+
+```shell
+lithops runtime delete -b aws_lambda -d cubed-runtime
+```
+
+Or you can remove everything (except config files) with
+
+```shell
+lithops clean -b aws_lambda
+```
diff --git a/examples/lithops/aws-lambda/lithops-add-random.py b/examples/lithops/aws-lambda/lithops-add-random.py
new file mode 100644
index 00000000..46803dad
--- /dev/null
+++ b/examples/lithops/aws-lambda/lithops-add-random.py
@@ -0,0 +1,43 @@
+import logging
+import sys
+
+from tqdm.contrib.logging import logging_redirect_tqdm
+
+import cubed
+import cubed.array_api as xp
+import cubed.random
+from cubed.extensions.history import HistoryCallback
+from cubed.extensions.timeline import TimelineVisualizationCallback
+from cubed.extensions.tqdm import TqdmProgressBar
+from cubed.runtime.executors.lithops import LithopsDagExecutor
+
+logging.basicConfig(level=logging.INFO)
+# suppress harmless connection pool warnings
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+
+if __name__ == "__main__":
+    tmp_path = sys.argv[1]
+    runtime = sys.argv[2]
+    spec = cubed.Spec(tmp_path, allowed_mem=2_000_000_000)
+    executor = LithopsDagExecutor()
+
+    a = cubed.random.random(
+        (50000, 50000), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    b = cubed.random.random(
+        (50000, 50000), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    c = xp.add(a, b)
+    with logging_redirect_tqdm():
+        progress = TqdmProgressBar()
+        hist = HistoryCallback()
+        timeline_viz = TimelineVisualizationCallback()
+        # use store=None to write to temporary zarr
+        cubed.to_zarr(
+            c,
+            store=None,
+            executor=executor,
+            callbacks=[progress, hist, timeline_viz],
+            runtime=runtime,
+            runtime_memory=2000,
+        )
diff --git a/examples/lithops/aws-lambda/lithops-matmul-random.py b/examples/lithops/aws-lambda/lithops-matmul-random.py
new file mode 100644
index 00000000..fe2b73eb
--- /dev/null
+++ b/examples/lithops/aws-lambda/lithops-matmul-random.py
@@ -0,0 +1,45 @@
+import logging
+import sys
+
+from tqdm.contrib.logging import logging_redirect_tqdm
+
+import cubed
+import cubed.array_api as xp
+import cubed.random
+from cubed.extensions.history import HistoryCallback
+from cubed.extensions.timeline import TimelineVisualizationCallback
+from cubed.extensions.tqdm import TqdmProgressBar
+from cubed.runtime.executors.lithops import LithopsDagExecutor
+
+logging.basicConfig(level=logging.INFO)
+# suppress harmless connection pool warnings
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+
+if __name__ == "__main__":
+    tmp_path = sys.argv[1]
+    runtime = sys.argv[2]
+    spec = cubed.Spec(tmp_path, allowed_mem=2_000_000_000)
+    executor = LithopsDagExecutor()
+
+    a = cubed.random.random(
+        (50000, 50000), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    b = cubed.random.random(
+        (50000, 50000), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    c = xp.astype(a, xp.float32)
+    d = xp.astype(b, xp.float32)
+    e = xp.matmul(c, d)
+    with logging_redirect_tqdm():
+        progress = TqdmProgressBar()
+        hist = HistoryCallback()
+        timeline_viz = TimelineVisualizationCallback()
+        # use store=None to write to temporary zarr
+        cubed.to_zarr(
+            e,
+            store=None,
+            executor=executor,
+            callbacks=[progress, hist, timeline_viz],
+            runtime=runtime,
+            runtime_memory=2000,
+        )
diff --git a/examples/modal/README.md b/examples/modal/README.md
index df03d2c2..5d5c3d7e 100644
--- a/examples/modal/README.md
+++ b/examples/modal/README.md
@@ -41,3 +41,10 @@ python modal-matmul-random.py "s3://cubed-modal-$USER-temp"
 ```
 
 These will take longer to run as they operate on more data.
+
+The last two examples use `TimelineVisualizationCallback` which produce a plot showing the timeline of events in the task lifecycle.
+The plots are `png` files and are written in the `plots` directory with a timestamp. Open the latest one with
+
+```shell
+open plots/$(ls plots | tail -1)
+```
diff --git a/examples/modal/modal-add-random.py b/examples/modal/modal-add-random.py
index 43cd2eaa..57d4cef8 100644
--- a/examples/modal/modal-add-random.py
+++ b/examples/modal/modal-add-random.py
@@ -5,6 +5,7 @@ import cubed.array_api as xp
 import cubed.random
 from cubed.extensions.history import HistoryCallback
 from cubed.extensions.timeline import TimelineVisualizationCallback
+from cubed.extensions.tqdm import TqdmProgressBar, std_out_err_redirect_tqdm
 from cubed.runtime.executors.modal_async import AsyncModalDagExecutor
 
 if __name__ == "__main__":
@@ -18,7 +19,11 @@ if __name__ == "__main__":
         (50000, 50000), chunks=(5000, 5000), spec=spec
     )  # 200MB chunks
     c = xp.add(a, b)
-    hist = HistoryCallback()
-    timeline_viz = TimelineVisualizationCallback()
-    # use store=None to write to temporary zarr
-    cubed.to_zarr(c, store=None, executor=executor, callbacks=[timeline_viz, hist])
+    with std_out_err_redirect_tqdm() as orig_stdout:
+        progress = TqdmProgressBar(file=orig_stdout, dynamic_ncols=True)
+        hist = HistoryCallback()
+        timeline_viz = TimelineVisualizationCallback()
+        # use store=None to write to temporary zarr
+        cubed.to_zarr(
+            c, store=None, executor=executor, callbacks=[progress, hist, timeline_viz]
+        )
diff --git a/pyproject.toml b/pyproject.toml
index 1a3d20a4..e2afc818 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,7 +36,9 @@ diagnostics = [
     "tqdm",
     "graphviz",
     "pydot",
-    "pandas"
+    "pandas",
+    "matplotlib",
+    "seaborn",
 ]
 beam = ["apache-beam", "gcsfs"]
 lithops = ["lithops[aws] >= 2.7.0"]


Test patch:
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)