Repository: tomwhite/cubed (python)

Issue:
Bugs with passing allowed_mem as a string
There are some bugs with #188

```python
In [2]: cubed.Spec(allowed_mem='0.8MB')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 1
----> 1 cubed.Spec(allowed_mem='0.8MB')

File ~/Documents/Work/Code/cubed/cubed/core/array.py:261, in Spec.__init__(self, work_dir, max_mem, allowed_mem, reserved_mem, executor, storage_options)
    259     self._allowed_mem = (max_mem or 0) + self.reserved_mem
    260 else:
--> 261     self._allowed_mem = convert_to_bytes(allowed_mem)
    263 self._executor = executor
    264 self._storage_options = storage_options

File ~/Documents/Work/Code/cubed/cubed/utils.py:216, in convert_to_bytes(size)
    214     value = size[:-2]
    215 else:
--> 216     raise ValueError(
    217         f"Invalid value: {size}. Expected a string ending with an SI prefix."
    218     )
    220 if unit in units and value.isdigit():
    221     # convert to bytes
    222     return int(value) * (1000 ** units[unit])

ValueError: Invalid value: 0.8MB. Expected a string ending with an SI prefix.

In [10]: cubed.Spec(allowed_mem='8 MB')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[10], line 1
----> 1 cubed.Spec(allowed_mem='8 MB')

File ~/Documents/Work/Code/cubed/cubed/core/array.py:261, in Spec.__init__(self, work_dir, max_mem, allowed_mem, reserved_mem, executor, storage_options)
    259     self._allowed_mem = (max_mem or 0) + self.reserved_mem
    260 else:
--> 261     self._allowed_mem = convert_to_bytes(allowed_mem)
    263 self._executor = executor
    264 self._storage_options = storage_options

File ~/Documents/Work/Code/cubed/cubed/utils.py:216, in convert_to_bytes(size)
    214     value = size[:-2]
    215 else:
--> 216     raise ValueError(
    217         f"Invalid value: {size}. Expected a string ending with an SI prefix."
    218     )
    220 if unit in units and value.isdigit():
    221     # convert to bytes
    222     return int(value) * (1000 ** units[unit])

ValueError: Invalid value: 8 MB. Expected a string ending with an SI prefix.
```

Pull request:
Fixes for converting strings to bytes
Closes #285 

Patch:
diff --git a/cubed/runtime/executors/dask.py b/cubed/runtime/executors/dask.py
deleted file mode 100644
index 0de8f99d..00000000
--- a/cubed/runtime/executors/dask.py
+++ /dev/null
@@ -1,34 +0,0 @@
-import dask
-
-from cubed.core.plan import visit_nodes
-from cubed.runtime.types import DagExecutor
-
-
-def exec_stage_func(func, *args, **kwargs):
-    # TODO would be good to give the dask tasks useful names
-    return dask.delayed(func)(*args, **kwargs)  # should we add pure=True?
-
-
-class DaskDelayedExecutor(DagExecutor):
-    """Executes each stage using dask.Delayed functions."""
-
-    def execute_dag(
-        self,
-        dag,
-        callbacks=None,
-        array_names=None,
-        resume=None,
-        spec=None,
-        **compute_kwargs,
-    ):
-        # Note this currently only builds the task graph for each stage once it gets to that stage in computation
-        for name, node in visit_nodes(dag, resume=resume):
-            pipeline = node["pipeline"]
-            stage_delayed_funcs = []
-            for m in pipeline.mappable:
-                delayed_func = exec_stage_func(
-                    pipeline.function, m, config=pipeline.config
-                )
-                stage_delayed_funcs.append(delayed_func)
-
-            dask.compute(*stage_delayed_funcs, **compute_kwargs)
diff --git a/cubed/runtime/executors/dask_distributed_async.py b/cubed/runtime/executors/dask_distributed_async.py
index 7de39827..5b35673f 100644
--- a/cubed/runtime/executors/dask_distributed_async.py
+++ b/cubed/runtime/executors/dask_distributed_async.py
@@ -92,6 +92,20 @@ def pipeline_to_stream(
     )
 
 
+def check_runtime_memory(spec, client):
+    allowed_mem = spec.allowed_mem if spec is not None else None
+    scheduler_info = client.scheduler_info()
+    workers = scheduler_info.get("workers")
+    if workers is None or len(workers) == 0:
+        raise ValueError("Cluster has no workers running")
+    runtime_memory = min(w["memory_limit"] // w["nthreads"] for w in workers.values())
+    if allowed_mem is not None:
+        if runtime_memory < allowed_mem:
+            raise ValueError(
+                f"Runtime memory ({runtime_memory}) is less than allowed_mem ({allowed_mem})"
+            )
+
+
 async def async_execute_dag(
     dag: MultiDiGraph,
     callbacks: Optional[Sequence[Callback]] = None,
@@ -104,6 +118,8 @@ async def async_execute_dag(
 ) -> None:
     compute_kwargs = compute_kwargs or {}
     async with Client(asynchronous=True, **compute_kwargs) as client:
+        if spec is not None:
+            check_runtime_memory(spec, client)
         if not compute_arrays_in_parallel:
             # run one pipeline at a time
             for name, node in visit_nodes(dag, resume=resume):
diff --git a/cubed/utils.py b/cubed/utils.py
index d608ca6e..82156bc8 100644
--- a/cubed/utils.py
+++ b/cubed/utils.py
@@ -9,7 +9,7 @@ from math import prod
 from operator import add
 from pathlib import Path
 from posixpath import join
-from typing import Tuple, Union
+from typing import Dict, Tuple, Union
 from urllib.parse import quote, unquote, urlsplit, urlunsplit
 
 import numpy as np
@@ -185,7 +185,7 @@ def extract_array_names(frame):
     return array_names_to_variable_names
 
 
-def convert_to_bytes(size: Union[int, str]) -> int:
+def convert_to_bytes(size: Union[int, float, str]) -> int:
     """
     Converts the input data size to bytes.
 
@@ -193,33 +193,53 @@ def convert_to_bytes(size: Union[int, str]) -> int:
 
     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.
+    size: int, float, or str:
+        Size of data. If numeric it should represent an integer >=0. If str it should be of form <value><unit> where unit can be B, 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"
+    units: Dict[str, int] = {"kB": 1, "MB": 2, "GB": 3, "TB": 4, "PB": 5}
+
+    def is_numeric_str(s: str) -> bool:
+        try:
+            float(s)
+            return True
+        except ValueError:
+            return False
+
+    if isinstance(size, str):
+        size = size.replace(" ", "")
+
+        # check if the format of the string is valid
+        if is_numeric_str(size):
+            unit_factor = 1.0
+            value = size
+        elif size[-1] == "B" and is_numeric_str(size[:-1]):
+            unit_factor = 1.0
             value = size[:-1]
-        elif size[-2:] in units and size[:-2].isdigit():
+        elif size[-2:] in units and is_numeric_str(size[:-2]):
             unit = size[-2:]
+            unit_factor = 1000 ** units[unit]
             value = size[:-2]
         else:
             raise ValueError(
-                f"Invalid value: {size}. Expected a string ending with an SI prefix."
+                f"Invalid value: {size}. Expected the string to be a numeric value ending with an SI prefix."
             )
 
-        if unit in units and value.isdigit():
-            # convert to bytes
-            return int(value) * (1000 ** units[unit])
-    raise ValueError(
-        f"Invalid value: {size}. Expected a positive integer or a string ending with an SI prefix."
-    )
+        # convert to float number of bytes
+        size = float(value) * unit_factor
+
+    if isinstance(size, float):
+        if size.is_integer():
+            size = int(size)
+        else:
+            raise ValueError(
+                f"Invalid value: {size}. Can't have a non-integer number of bytes"
+            )
+
+    if size >= 0:
+        return size
+    else:
+        raise ValueError(f"Invalid value: {size}. Must be a positive value")


Test patch:
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 5588167c..bd750146 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -283,12 +283,17 @@ class TestSpecMemArgTypes:
         [
             (500, 500),
             (100_000, 100_000),
+            (50.0, 50),
             ("500B", 500),
             ("1kB", 1000),
             ("1MB", 1000**2),
             ("1GB", 1000**3),
             ("1TB", 1000**4),
             ("1PB", 1000**5),
+            ("100_000", 100_000),
+            ("1.2MB", 1.2 * 1000**2),
+            ("1 MB", 1000**2),
+            ("1.2 MB", 1.2 * 1000**2),
         ],
     )
     def test_convert_to_bytes(self, input_value, expected_value):
@@ -302,7 +307,8 @@ class TestSpecMemArgTypes:
             "1kb",  # lower-case k is not valid
             "invalid",  # completely invalid input
             -512,  # negative integer
-            1000.0,  # invalid type
+            "kB",  # only unit, no value
+            "1.1B",  # can't have a fractional number of bytes
         ],
     )
     def test_convert_to_bytes_error(self, input_value):
diff --git a/cubed/tests/test_executor_features.py b/cubed/tests/test_executor_features.py
index 16add1a3..944640d4 100644
--- a/cubed/tests/test_executor_features.py
+++ b/cubed/tests/test_executor_features.py
@@ -77,14 +77,6 @@ def test_retries(mocker, spec):
     platform.system() == "Windows", reason="measuring memory does not run on windows"
 )
 def test_callbacks(spec, executor):
-    try:
-        from cubed.runtime.executors.dask import DaskDelayedExecutor
-
-        if isinstance(executor, DaskDelayedExecutor):
-            pytest.skip(f"{type(executor)} does not support callbacks")
-    except ImportError:
-        pass
-
     task_counter = TaskCounter()
     # test following indirectly by checking they don't cause a failure
     progress = TqdmProgressBar()
@@ -135,14 +127,6 @@ def test_resume(spec, executor):
     except ImportError:
         pass
 
-    try:
-        from cubed.runtime.executors.dask import DaskDelayedExecutor
-
-        if isinstance(executor, DaskDelayedExecutor):
-            pytest.skip(f"{type(executor)} does not support resume")
-    except ImportError:
-        pass
-
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
     c = xp.add(a, b)
@@ -215,6 +199,55 @@ def test_compute_arrays_in_parallel_modal(modal_executor, compute_arrays_in_para
         fs.rm(tmp_path, recursive=True)
 
 
+def test_check_runtime_memory_dask(spec, executor):
+    pytest.importorskip("dask.distributed")
+    try:
+        from cubed.runtime.executors.dask_distributed_async import (
+            AsyncDaskDistributedExecutor,
+        )
+
+        if not isinstance(executor, AsyncDaskDistributedExecutor):
+            pytest.skip(f"{type(executor)} does not support check_runtime_memory")
+    except ImportError:
+        pass
+
+    spec = cubed.Spec(spec.work_dir, allowed_mem="4GB")  # larger than runtime memory
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    with pytest.raises(
+        ValueError,
+        match=r"Runtime memory \(2000000000\) is less than allowed_mem \(4000000000\)",
+    ):
+        c.compute(
+            executor=executor,
+            compute_kwargs=dict(memory_limit="2 GB", threads_per_worker=1),
+        )
+
+
+def test_check_runtime_memory_dask_no_workers(spec, executor):
+    pytest.importorskip("dask.distributed")
+    try:
+        from cubed.runtime.executors.dask_distributed_async import (
+            AsyncDaskDistributedExecutor,
+        )
+
+        if not isinstance(executor, AsyncDaskDistributedExecutor):
+            pytest.skip(f"{type(executor)} does not support check_runtime_memory")
+    except ImportError:
+        pass
+
+    spec = cubed.Spec(spec.work_dir, allowed_mem=100000)
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
+    c = xp.add(a, b)
+    with pytest.raises(ValueError, match=r"Cluster has no workers running"):
+        c.compute(
+            executor=executor,
+            compute_kwargs=dict(n_workers=0),
+        )
+
+
 @pytest.mark.cloud
 def test_check_runtime_memory_modal(spec, modal_executor):
     tmp_path = "s3://cubed-unittest/check-runtime-memory"
diff --git a/cubed/tests/utils.py b/cubed/tests/utils.py
index ff62cb6d..aea525ab 100644
--- a/cubed/tests/utils.py
+++ b/cubed/tests/utils.py
@@ -31,15 +31,6 @@ try:
 except ImportError:
     pass
 
-try:
-    from cubed.runtime.executors.dask import DaskDelayedExecutor
-
-    ALL_EXECUTORS.append(DaskDelayedExecutor())
-
-    MAIN_EXECUTORS.append(DaskDelayedExecutor())
-except ImportError:
-    pass
-
 try:
     from cubed.runtime.executors.dask_distributed_async import (
         AsyncDaskDistributedExecutor,