Repository: cubed-dev/cubed (python)

Issue:
Issue warning if a computation actually used more memory than `allowed_mem`
We know the peak memory used in every task, so if this value ever exceeds the amount specified in `allowed_mem` (specified by the user) then we can issue a warning. We should encourage users to report these as bugs in cubed (though the immediate workaround would be to increase the memory available to the runtime).

Pull request description:
Warn if peak mem exceeds `allowed_mem`
Fixes #110 

Patch:
diff --git a/cubed/extensions/mem_warn.py b/cubed/extensions/mem_warn.py
new file mode 100644
index 00000000..03f5b2bf
--- /dev/null
+++ b/cubed/extensions/mem_warn.py
@@ -0,0 +1,35 @@
+import warnings
+from collections import Counter
+
+from cubed.runtime.pipeline import visit_nodes
+from cubed.runtime.types import Callback
+
+
+class MemoryWarningCallback(Callback):
+    def on_compute_start(self, event):
+        # store ops keyed by name
+        self.ops = {}
+        for name, node in visit_nodes(event.dag, event.resume):
+            primitive_op = node["primitive_op"]
+            self.ops[name] = primitive_op
+
+        # count number of times each op exceeds allowed mem
+        self.counter = Counter()
+
+    def on_task_end(self, event):
+        allowed_mem = self.ops[event.name].allowed_mem
+        if (
+            event.peak_measured_mem_end is not None
+            and event.peak_measured_mem_end > allowed_mem
+        ):
+            self.counter.update({event.name: 1})
+
+    def on_compute_end(self, event):
+        if sum(self.counter.values()) > 0:
+            exceeded = [
+                f"{k} ({v}/{self.ops[k].num_tasks})" for k, v in self.counter.items()
+            ]
+            warnings.warn(
+                f"Peak memory usage exceeded allowed_mem when running tasks: {', '.join(exceeded)}",
+                UserWarning,
+            )
diff --git a/cubed/runtime/executors/local.py b/cubed/runtime/executors/local.py
index 599702ec..ebb18aa3 100644
--- a/cubed/runtime/executors/local.py
+++ b/cubed/runtime/executors/local.py
@@ -6,6 +6,7 @@ from functools import partial
 from typing import Any, AsyncIterator, Callable, Iterable, Optional, Sequence
 
 import cloudpickle
+import psutil
 from aiostream import stream
 from aiostream.core import Stream
 from networkx import MultiDiGraph
@@ -154,6 +155,16 @@ def pipeline_to_stream(
     )
 
 
+def check_runtime_memory(spec, max_workers):
+    allowed_mem = spec.allowed_mem if spec is not None else None
+    total_mem = psutil.virtual_memory().total
+    if allowed_mem is not None:
+        if total_mem < allowed_mem * max_workers:
+            raise ValueError(
+                f"Total memory on machine ({total_mem}) is less than allowed_mem * max_workers ({allowed_mem} * {max_workers} = {allowed_mem * max_workers})"
+            )
+
+
 async def async_execute_dag(
     dag: MultiDiGraph,
     callbacks: Optional[Sequence[Callback]] = None,
@@ -163,16 +174,18 @@ async def async_execute_dag(
     **kwargs,
 ) -> None:
     concurrent_executor: Executor
+    max_workers = kwargs.pop("max_workers", os.cpu_count())
     use_processes = kwargs.pop("use_processes", False)
+    if spec is not None:
+        check_runtime_memory(spec, max_workers)
     if use_processes:
-        max_workers = kwargs.pop("max_workers", None)
         context = multiprocessing.get_context("spawn")
         # max_tasks_per_child is only supported from Python 3.11
         concurrent_executor = ProcessPoolExecutor(
             max_workers=max_workers, mp_context=context, max_tasks_per_child=1
         )
     else:
-        concurrent_executor = ThreadPoolExecutor()
+        concurrent_executor = ThreadPoolExecutor(max_workers=max_workers)
     try:
         if not compute_arrays_in_parallel:
             # run one pipeline at a time
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
index 5bdb101c..98be5510 100644
--- a/cubed/runtime/executors/modal.py
+++ b/cubed/runtime/executors/modal.py
@@ -40,6 +40,7 @@ else:
             "mypy_extensions",  # for rechunker
             "ndindex",
             "networkx",
+            "psutil",
             "pytest-mock",  # TODO: only needed for tests
             "s3fs",
             "tenacity",
@@ -55,6 +56,7 @@ else:
             "mypy_extensions",  # for rechunker
             "ndindex",
             "networkx",
+            "psutil",
             "pytest-mock",  # TODO: only needed for tests
             "gcsfs",
             "tenacity",
diff --git a/pyproject.toml b/pyproject.toml
index 9c28bc33..97e1be3c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,6 +33,7 @@ dependencies = [
     "ndindex",
     "networkx != 2.8.3, != 2.8.4, != 2.8.5, != 2.8.6, != 2.8.7, != 2.8.8, != 3.0.*, != 3.1.*, != 3.2.*",
     "numpy >= 1.22",
+    "psutil",
     "tenacity",
     "toolz",
     "zarr",
diff --git a/requirements.txt b/requirements.txt
index 6f84c347..88f56304 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,6 +9,7 @@ mypy_extensions # for rechunker
 ndindex
 networkx != 2.8.3, != 2.8.4, != 2.8.5, != 2.8.6, != 2.8.7, != 2.8.8, != 3.0.*, != 3.1.*, != 3.2.*
 numpy >= 1.22
+psutil
 tenacity
 toolz
 zarr
diff --git a/setup.cfg b/setup.cfg
index 76b6a455..e59de308 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -52,6 +52,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-pandas.*]
 ignore_missing_imports = True
+[mypy-psutil.*]
+ignore_missing_imports = True
 [mypy-pylab.*]
 ignore_missing_imports = True
 [mypy-pytest.*]


Test patch:
diff --git a/cubed/tests/test_executor_features.py b/cubed/tests/test_executor_features.py
index 60bb397b..22fce7bf 100644
--- a/cubed/tests/test_executor_features.py
+++ b/cubed/tests/test_executor_features.py
@@ -1,8 +1,11 @@
 import contextlib
+import os
 import platform
+import re
 
 import fsspec
 import numpy as np
+import psutil
 import pytest
 from numpy.testing import assert_array_equal
 
@@ -10,6 +13,7 @@ import cubed
 import cubed.array_api as xp
 import cubed.random
 from cubed.extensions.history import HistoryCallback
+from cubed.extensions.mem_warn import MemoryWarningCallback
 from cubed.extensions.rich import RichProgressBar
 from cubed.extensions.timeline import TimelineVisualizationCallback
 from cubed.extensions.tqdm import TqdmProgressBar
@@ -145,6 +149,28 @@ def test_callbacks_modal(spec, modal_executor):
         fs.rm(tmp_path, recursive=True)
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="measuring memory does not run on windows"
+)
+def test_mem_warn(tmp_path, executor):
+    if executor.name not in ("processes", "lithops"):
+        pytest.skip(f"{executor.name} executor does not support MemoryWarningCallback")
+
+    spec = cubed.Spec(tmp_path, allowed_mem=200_000_000, reserved_mem=100_000_000)
+    mem_warn = MemoryWarningCallback()
+
+    def func(a):
+        np.ones(100_000_000)  # blow memory
+        return a
+
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
+    b = cubed.map_blocks(func, a, dtype=a.dtype)
+    with pytest.raises(
+        UserWarning, match="Peak memory usage exceeded allowed_mem when running tasks"
+    ):
+        b.compute(executor=executor, callbacks=[mem_warn])
+
+
 def test_resume(spec, executor):
     if executor.name == "beam":
         pytest.skip(f"{executor.name} executor does not support resume")
@@ -264,3 +290,28 @@ def test_check_runtime_memory_modal(spec, modal_executor):
         match=r"Runtime memory \(2097152000\) is less than allowed_mem \(4000000000\)",
     ):
         c.compute(executor=modal_executor)
+
+
+def test_check_runtime_memory_processes(spec, executor):
+    if executor.name != "processes":
+        pytest.skip(f"{executor.name} executor does not support check_runtime_memory")
+
+    total_mem = psutil.virtual_memory().total
+    max_workers = os.cpu_count()
+    mem_per_worker = total_mem // max_workers
+    allowed_mem = mem_per_worker * 2  # larger than will fit
+
+    spec = cubed.Spec(spec.work_dir, allowed_mem=allowed_mem)
+    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=re.escape(
+            f"Total memory on machine ({total_mem}) is less than allowed_mem * max_workers ({allowed_mem} * {max_workers} = {allowed_mem * max_workers})"
+        ),
+    ):
+        c.compute(executor=executor)
+
+    # OK if we use fewer workers
+    c.compute(executor=executor, max_workers=max_workers // 2)
diff --git a/cubed/tests/test_mem_utilization.py b/cubed/tests/test_mem_utilization.py
index 853f993f..85f067c8 100644
--- a/cubed/tests/test_mem_utilization.py
+++ b/cubed/tests/test_mem_utilization.py
@@ -14,6 +14,7 @@ import cubed.array_api as xp
 import cubed.random
 from cubed.backend_array_api import namespace as nxp
 from cubed.extensions.history import HistoryCallback
+from cubed.extensions.mem_warn import MemoryWarningCallback
 from cubed.runtime.executors.lithops import LithopsExecutor
 from cubed.tests.utils import LITHOPS_LOCAL_CONFIG
 
@@ -277,12 +278,13 @@ def run_operation(tmp_path, name, result_array, *, optimize_function=None):
     # result_array.visualize(f"cubed-{name}", optimize_function=optimize_function)
     executor = LithopsExecutor(config=LITHOPS_LOCAL_CONFIG)
     hist = HistoryCallback()
+    mem_warn = MemoryWarningCallback()
     # use store=None to write to temporary zarr
     cubed.to_zarr(
         result_array,
         store=None,
         executor=executor,
-        callbacks=[hist],
+        callbacks=[hist, mem_warn],
         optimize_function=optimize_function,
     )