Repository: tomwhite/cubed (python)

Issue:
Error when trying to measure reserved memory with lithops
I tried to measure the reserved memory with lithops (set up to run on Google Cloud Functions) and this happened:

```python
IPython 8.13.2 -- An enhanced Interactive Python. Type '?' for help.

In [1]: from cubed import measure_reserved_memory

In [2]: from cubed.runtime.executors.lithops import LithopsDagExecutor

In [3]: import cubed

In [4]: measure_reserved_memory(LithopsDagExecutor())
2023-05-24 17:59:58,539 [INFO] config.py:134 -- Lithops v2.9.0
2023-05-24 17:59:58,878 [INFO] gcp_storage.py:45 -- Google Cloud Storage client created
2023-05-24 17:59:58,990 [INFO] gcp_functions.py:68 -- Google Cloud Functions client created - Region: us-east1 - Project: pangeo-forge-4967
2023-05-24 17:59:58,990 [INFO] invokers.py:108 -- ExecutorID 646b95-0 | JobID M000 - Selected Runtime: default-runtime-v39 - 256MB
2023-05-24 17:59:59,545 [INFO] invokers.py:172 -- ExecutorID 646b95-0 | JobID M000 - Starting function invocation: <lambda>() - Total: 1 activations
2023-05-24 17:59:59,548 [INFO] invokers.py:208 -- ExecutorID 646b95-0 | JobID M000 - View execution logs at /tmp/lithops-tom/logs/646b95-0-M000.log
2023-05-24 17:59:59,549 [INFO] wait.py:97 -- ExecutorID 646b95-0 - Waiting for -1% of 1 function activations to complete
2023-05-24 18:00:02,556 [WARNING] future.py:257 -- ExecutorID 646b95-0 | JobID M000 - There was an exception - Activation ID: 405f3d1778b148a5be84a124639af17b
2023-05-24 18:00:02,557 [WARNING] future.py:260 -- Exception: ModuleNotFoundError - No module named 'cubed'
2023-05-24 18:00:02,557 [INFO] invokers.py:108 -- ExecutorID 646b95-0 | JobID M001 - Selected Runtime: default-runtime-v39 - 256MB
2023-05-24 18:00:02,563 [INFO] invokers.py:172 -- ExecutorID 646b95-0 | JobID M001 - Starting function invocation: <lambda>() - Total: 1 activations
2023-05-24 18:00:02,563 [INFO] invokers.py:208 -- ExecutorID 646b95-0 | JobID M001 - View execution logs at /tmp/lithops-tom/logs/646b95-0-M001.log
2023-05-24 18:00:02,564 [INFO] wait.py:97 -- ExecutorID 646b95-0 - Waiting for -1% of 1 function activations to complete
2023-05-24 18:00:05,572 [WARNING] future.py:257 -- ExecutorID 646b95-0 | JobID M001 - There was an exception - Activation ID: ee879e0e92b242c4b8c57813a5376753
2023-05-24 18:00:05,572 [WARNING] future.py:260 -- Exception: ModuleNotFoundError - No module named 'cubed'
2023-05-24 18:00:05,573 [INFO] invokers.py:108 -- ExecutorID 646b95-0 | JobID M002 - Selected Runtime: default-runtime-v39 - 256MB
2023-05-24 18:00:05,577 [INFO] invokers.py:172 -- ExecutorID 646b95-0 | JobID M002 - Starting function invocation: <lambda>() - Total: 1 activations
2023-05-24 18:00:05,577 [INFO] invokers.py:208 -- ExecutorID 646b95-0 | JobID M002 - View execution logs at /tmp/lithops-tom/logs/646b95-0-M002.log
2023-05-24 18:00:05,578 [INFO] wait.py:97 -- ExecutorID 646b95-0 - Waiting for -1% of 1 function activations to complete
2023-05-24 18:00:08,584 [WARNING] future.py:257 -- ExecutorID 646b95-0 | JobID M002 - There was an exception - Activation ID: 4f8232940f6741ff9c80c7c12db8b3cf
2023-05-24 18:00:08,584 [WARNING] future.py:260 -- Exception: ModuleNotFoundError - No module named 'cubed'
2023-05-24 18:00:08,584 [INFO] invokers.py:108 -- ExecutorID 646b95-0 | JobID M003 - Selected Runtime: default-runtime-v39 - 256MB
2023-05-24 18:00:08,588 [INFO] invokers.py:172 -- ExecutorID 646b95-0 | JobID M003 - Starting function invocation: <lambda>() - Total: 1 activations
2023-05-24 18:00:08,589 [INFO] invokers.py:208 -- ExecutorID 646b95-0 | JobID M003 - View execution logs at /tmp/lithops-tom/logs/646b95-0-M003.log
2023-05-24 18:00:08,590 [INFO] wait.py:97 -- ExecutorID 646b95-0 - Waiting for -1% of 1 function activations to complete
2023-05-24 18:00:11,598 [WARNING] future.py:257 -- ExecutorID 646b95-0 | JobID M003 - There was an exception - Activation ID: ab403d5e96084537813e08e0c7ffd07a
2023-05-24 18:00:11,598 [WARNING] future.py:260 -- Exception: ModuleNotFoundError - No module named 'cubed'
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[4], line 1
----> 1 measure_reserved_memory(LithopsDagExecutor())

File ~/Documents/Work/Code/cubed/cubed/core/array.py:390, in measure_reserved_memory(executor)
    388 b = xp.negative(a)
    389 peak_memory_callback = PeakMemoryCallback()
--> 390 b.compute(executor=executor, callbacks=[peak_memory_callback])
    391 return peak_memory_callback.peak_memory

File ~/Documents/Work/Code/cubed/cubed/core/array.py:115, in CoreArray.compute(self, executor, callbacks, optimize_graph, **kwargs)
    106 def compute(
    107     self,
    108     *,
   (...)
    112     **kwargs,
    113 ):
    114     """Compute this array, and any arrays that it depends on."""
--> 115     result = compute(
    116         self,
    117         executor=executor,
    118         callbacks=callbacks,
    119         optimize_graph=optimize_graph,
    120         **kwargs,
    121     )
    122     if result:
    123         return result[0]

File ~/Documents/Work/Code/cubed/cubed/core/array.py:316, in compute(executor, callbacks, optimize_graph, *arrays, **kwargs)
    313         executor = PythonDagExecutor()
    315 _return_in_memory_array = kwargs.pop("_return_in_memory_array", True)
--> 316 plan.execute(
    317     executor=executor,
    318     callbacks=callbacks,
    319     optimize_graph=optimize_graph,
    320     array_names=[a.name for a in arrays],
    321     **kwargs,
    322 )
    324 if _return_in_memory_array:
    325     return tuple(a._read_stored() for a in arrays)

File ~/Documents/Work/Code/cubed/cubed/core/plan.py:156, in Plan.execute(self, executor, callbacks, optimize_graph, array_names, **kwargs)
    154 if callbacks is not None:
    155     [callback.on_compute_start(dag) for callback in callbacks]
--> 156 executor.execute_dag(
    157     dag, callbacks=callbacks, array_names=array_names, **kwargs
    158 )
    159 if callbacks is not None:
    160     [callback.on_compute_end(dag) for callback in callbacks]

File ~/Documents/Work/Code/cubed/cubed/runtime/executors/lithops.py:255, in LithopsDagExecutor.execute_dag(self, dag, callbacks, array_names, **kwargs)
    253 def execute_dag(self, dag, callbacks=None, array_names=None, **kwargs):
    254     merged_kwargs = {**self.kwargs, **kwargs}
--> 255     execute_dag(dag, callbacks=callbacks, array_names=array_names, **merged_kwargs)

File ~/Documents/Work/Code/cubed/cubed/runtime/executors/lithops.py:179, in execute_dag(dag, callbacks, array_names, **kwargs)
    177 for stage in pipeline.stages:
    178     if stage.mappable is not None:
--> 179         for _, stats in map_unordered(
    180             executor,
    181             run_func,
    182             list(stage.mappable),
    183             func=stage.function,
    184             config=pipeline.config,
    185             name=name,
    186             use_backups=use_backups,
    187             return_stats=True,
    188         ):
    189             if callbacks is not None:
    190                 event = lithops_stats_to_task_end_event(name, stats)

File ~/Documents/Work/Code/cubed/cubed/runtime/executors/lithops.py:118, in map_unordered(lithops_function_executor, map_function, map_iterdata, include_modules, max_failures, use_backups, return_stats, **kwargs)
    114     if failures > max_failures:
    115         # re-raise exception
    116         # TODO: why does calling status not raise the exception?
    117         future.status(throw_except=True)
--> 118         reraise(*future._exception)
    119     failed.append(future)
    120 else:

File ~/miniconda3/envs/cubed_xarray/lib/python3.9/site-packages/six.py:719, in reraise(tp, value, tb)
    717     if value.__traceback__ is not tb:
    718         raise value.with_traceback(tb)
--> 719     raise value
    720 finally:
    721     value = None

File /workspace/lithops/worker/jobrunner.py:209, in run()

ModuleNotFoundError: No module named 'cubed'
```

I assume this means that lithops needs to be given a runtime environment containing the cubed package, but it's not clear how to pass this in from the existing documentation on `measure_reserved_memory`. Perhaps there should be a section in the docs on how to use `measure_reserved_memory`?

Pull request description:
Fix measure reserved
Fixes #184. Also renames the function to `measure_reserved_mem` to be in line with `reserved_mem` and all of the `*_mem` settings.

Patch:
diff --git a/cubed/__init__.py b/cubed/__init__.py
index 7c82da4a..63580391 100644
--- a/cubed/__init__.py
+++ b/cubed/__init__.py
@@ -18,6 +18,7 @@ from .core.array import (
     Spec,
     TaskEndEvent,
     compute,
+    measure_reserved_mem,
     measure_reserved_memory,
     visualize,
 )
@@ -35,6 +36,7 @@ __all__ = [
     "from_array",
     "from_zarr",
     "map_blocks",
+    "measure_reserved_mem",
     "measure_reserved_memory",
     "store",
     "to_zarr",
diff --git a/cubed/core/__init__.py b/cubed/core/__init__.py
index 43cced07..a4dc25c9 100644
--- a/cubed/core/__init__.py
+++ b/cubed/core/__init__.py
@@ -6,6 +6,7 @@ from .array import (
     TaskEndEvent,
     compute,
     gensym,
+    measure_reserved_mem,
     measure_reserved_memory,
     visualize,
 )
diff --git a/cubed/core/array.py b/cubed/core/array.py
index 53224735..2e1cc7f7 100644
--- a/cubed/core/array.py
+++ b/cubed/core/array.py
@@ -261,7 +261,7 @@ class Spec:
 
         See Also
         --------
-        cubed.measure_reserved_memory
+        cubed.measure_reserved_mem
         """
         return self._reserved_mem
 
@@ -428,7 +428,20 @@ class PeakMeasuredMemoryCallback(Callback):
         self.peak_measured_mem = event.peak_measured_mem_end
 
 
-def measure_reserved_memory(executor):
+def measure_reserved_memory(
+    executor: Executor, work_dir: Optional[str] = None, **kwargs
+) -> int:
+    warn(
+        "`measure_reserved_memory` is deprecated, please use `measure_reserved_mem` instead",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return measure_reserved_mem(executor, work_dir=work_dir, **kwargs)
+
+
+def measure_reserved_mem(
+    executor: Executor, work_dir: Optional[str] = None, **kwargs
+) -> int:
     """Measures the reserved memory use for a given executor runtime.
 
     This is the memory used by the Python process for running a task,
@@ -447,6 +460,13 @@ def measure_reserved_memory(executor):
         The executor to use to run the computation. It must be an executor that
         reports peak memory, such as Lithops or Modal.
 
+    work_dir : str or None, optional
+        The directory path (specified as an fsspec URL) used for storing intermediate data.
+        This is required when using a cloud runtime.
+
+    kwargs
+        Keyword arguments to pass to the ``compute`` function.
+
     Returns
     -------
     int
@@ -454,8 +474,9 @@ def measure_reserved_memory(executor):
     """
     import cubed.array_api as xp
 
-    a = xp.ones((1,))
+    # give a generous memory allowance
+    a = xp.ones((1,), spec=Spec(work_dir, allowed_mem="500MB"))
     b = xp.negative(a)
     peak_measured_mem_callback = PeakMeasuredMemoryCallback()
-    b.compute(executor=executor, callbacks=[peak_measured_mem_callback])
+    b.compute(executor=executor, callbacks=[peak_measured_mem_callback], **kwargs)
     return peak_measured_mem_callback.peak_measured_mem
diff --git a/docs/computation.md b/docs/computation.md
index 90469491..577ee74b 100644
--- a/docs/computation.md
+++ b/docs/computation.md
@@ -19,7 +19,7 @@ The maximum amount of memory that tasks are allowed to use must be specified by
 
 If the `projected_mem` calculated by Cubed is greater than the value of `allowed_mem` set by the user, an exception is raised during the planning phase. This check means that the user can have high confidence that the operation will run within its memory budget.
 
-It is also a good idea to set `reserved_mem`, the amount of memory reserved on a worker for non-data use - it's whatever is needed by the Python process for running a task, and can be estimated using the `measure_reserved_memory` function. Cubed will use `reserved_mem` as a baseline when calculating `projected_mem`, in order to more accurately estimate the upper bound on memory usage.
+It is also a good idea to set `reserved_mem`, the amount of memory reserved on a worker for non-data use - it's whatever is needed by the Python process for running a task, and can be estimated using the `measure_reserved_mem` function. Cubed will use `reserved_mem` as a baseline when calculating `projected_mem`, in order to more accurately estimate the upper bound on memory usage.
 
 The actual (peak) memory usage of each task in a computation can be measured at runtime and analysed to see how well the projected memory matched `peak_measured_mem`. See `HistoryCallback` for details.
 
diff --git a/examples/lithops/gcf/README.md b/examples/lithops/gcf/README.md
index 59ec848d..e25397c4 100644
--- a/examples/lithops/gcf/README.md
+++ b/examples/lithops/gcf/README.md
@@ -23,6 +23,7 @@ pip install -r requirements.txt  # use requirements file from same directory as
 
 ```shell
 lithops runtime build -f requirements.txt cubed-runtime -b gcp_functions
+lithops runtime deploy -b gcp_functions --memory 2048 cubed-runtime # optional, will be done automatically on first use
 ```
 
 ## Running
@@ -40,3 +41,33 @@ Run the other examples in a similar way
 ```shell
 python lithops-add-random.py "gs://cubed-$USER-temp" cubed-runtime
 ```
+
+and
+
+```shell
+python lithops-matmul-random.py "gs://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 gcp_functions -d cubed-runtime
+```
+
+Or you can remove everything (except config files) with
+
+```shell
+lithops clean -b gcp_functions
+```
diff --git a/examples/lithops/gcf/lithops-add-random.py b/examples/lithops/gcf/lithops-add-random.py
new file mode 100644
index 00000000..8c504bb9
--- /dev/null
+++ b/examples/lithops/gcf/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=2048,
+        )
diff --git a/examples/lithops/gcf/lithops-matmul-random.py b/examples/lithops/gcf/lithops-matmul-random.py
new file mode 100644
index 00000000..a9d1eb10
--- /dev/null
+++ b/examples/lithops/gcf/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=2048,
+        )
diff --git a/examples/lithops/gcf/requirements.txt b/examples/lithops/gcf/requirements.txt
index 87e72a85..10b767ef 100644
--- a/examples/lithops/gcf/requirements.txt
+++ b/examples/lithops/gcf/requirements.txt
@@ -1,4 +1,4 @@
 cubed
-lithops[gcf,gcp]
+lithops[gcp]
 gcsfs
 tqdm


Test patch:
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 7a8a14ce..a524c073 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -484,13 +484,13 @@ def test_already_computed(spec):
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")
-def test_measure_reserved_memory(executor):
+def test_measure_reserved_mem(executor):
     pytest.importorskip("lithops")
 
     from cubed.runtime.executors.lithops import LithopsDagExecutor
 
     if not isinstance(executor, LithopsDagExecutor):
-        pytest.skip(f"{type(executor)} does not support measure_reserved_memory")
+        pytest.skip(f"{type(executor)} does not support measure_reserved_mem")
 
-    reserved_memory = cubed.measure_reserved_memory(executor=executor)
+    reserved_memory = cubed.measure_reserved_mem(executor=executor)
     assert reserved_memory > 1_000_000  # over 1MB
diff --git a/cubed/tests/test_mem_utilization.py b/cubed/tests/test_mem_utilization.py
index 4a096086..f1be412c 100644
--- a/cubed/tests/test_mem_utilization.py
+++ b/cubed/tests/test_mem_utilization.py
@@ -20,7 +20,7 @@ def spec(tmp_path, reserved_mem):
 @pytest.fixture(scope="module")
 def reserved_mem():
     executor = LithopsDagExecutor(config=LITHOPS_LOCAL_CONFIG)
-    res = cubed.measure_reserved_memory(executor) * 1.05  # add some wiggle room
+    res = cubed.measure_reserved_mem(executor) * 1.05  # add some wiggle room
     return round_up_to_multiple(res, 10_000_000)  # round up to nearest multiple of 10MB