Use donfig to control default spec

Author: tomwhite

cubed/__init__.py

@@ -12,6 +12,14 @@
 except Exception:  # pragma: no cover
     __version__ = "unknown"
 
+from donfig import Config
+
+config = Config(
+    "cubed",
+    # default spec is local temp dir and a modest amount of memory (200MB, of which 100MB is reserved)
+    defaults=[{"spec": {"allowed_mem": 200_000_000, "reserved_mem": 100_000_000}}],
+)
+
 from .array_api import Array
 from .core.array import compute, measure_reserved_mem, visualize
 from .core.gufunc import apply_gufunc
@@ -28,6 +36,7 @@
     "TaskEndEvent",
     "apply_gufunc",
     "compute",
+    "config",
     "from_array",
     "from_zarr",
     "map_blocks",

cubed/core/array.py

@@ -4,8 +4,9 @@
 import numpy as np
 from toolz import map, reduce
 
+from cubed import config
 from cubed.runtime.types import Callback, Executor
-from cubed.spec import Spec
+from cubed.spec import Spec, spec_from_config
 from cubed.storage.zarr import open_if_lazy_zarr_array
 from cubed.utils import chunk_memory
 from cubed.vendor.dask.array.core import normalize_chunks
@@ -41,11 +42,8 @@ def __init__(self, name, zarray, spec, plan):
         self._chunks = normalize_chunks(
             zarray.chunks, shape=self.shape, dtype=self.dtype
         )
-        # if no spec is supplied, use a default with local temp dir,
-        # and a modest amount of memory (200MB, of which 100MB is reserved)
-        self.spec = spec or Spec(
-            None, allowed_mem=200_000_000, reserved_mem=100_000_000
-        )
+        # get spec from config if not supplied
+        self.spec = spec or spec_from_config(config)
         self.plan = plan
 
     @property

cubed/runtime/executors/modal.py

@@ -32,6 +32,7 @@
     aws_image = modal.Image.debian_slim().pip_install(
         [
             "array-api-compat",
+            "donfig",
             "fsspec",
             "mypy_extensions",  # for rechunker
             "networkx",
@@ -46,6 +47,7 @@
     gcp_image = modal.Image.debian_slim().pip_install(
         [
             "array-api-compat",
+            "donfig",
             "fsspec",
             "mypy_extensions",  # for rechunker
             "networkx",

cubed/spec.py

@@ -1,5 +1,9 @@
+from functools import lru_cache
 from typing import Optional, Union
 
+import donfig
+from donfig.config_obj import expand_environment_variables
+
 from cubed.runtime.create import create_executor
 from cubed.runtime.types import Executor
 from cubed.utils import convert_to_bytes
@@ -109,3 +113,14 @@ def __eq__(self, other):
             )
         else:
             return False
+
+
+def spec_from_config(config):
+    return _spec_from_serialized_config(config.serialize())
+
+
+@lru_cache  # ensure arrays have the same Spec object for a given config
+def _spec_from_serialized_config(ser: str):
+    config = donfig.deserialize(ser)
+    spec_dict = expand_environment_variables(config["spec"])
+    return Spec(**spec_dict)

cubed/tests/runtime/test_modal_async.py

@@ -20,6 +20,7 @@
 image = modal.Image.debian_slim().pip_install(
     [
         "array-api-compat",
+        "donfig",
         "fsspec",
         "mypy_extensions",  # for rechunker
         "networkx",

cubed/tests/test_core.py

@@ -280,11 +280,21 @@ def test_default_spec(executor):
 
 def test_default_spec_allowed_mem_exceeded():
     # default spec fails for large computations
-    a = xp.ones((100000, 100000), chunks=(10000, 10000))
+    a = xp.ones((20000, 1000), chunks=(10000, 1000))
     with pytest.raises(ValueError):
         xp.negative(a)
 
 
+def test_default_spec_config_override():
+    # override default spec to increase allowed_mem
+    from cubed import config
+
+    with config.set({"spec.allowed_mem": "500MB"}):
+        a = xp.ones((20000, 1000), chunks=(10000, 1000))
+        b = xp.negative(a)
+        assert_array_equal(b.compute(), -np.ones((20000, 1000)))
+
+
 def test_different_specs(tmp_path):
     spec1 = cubed.Spec(tmp_path, allowed_mem=100000)
     spec2 = cubed.Spec(tmp_path, allowed_mem=200000)

pyproject.toml

@@ -23,6 +23,7 @@ requires-python = ">=3.9"
 dependencies = [
     "aiostream",
     "array-api-compat",
+    "donfig",
     "fsspec",
     "mypy_extensions", # for rechunker
     "networkx < 2.8.3",