Repository: tomwhite/cubed
Pull request #182 (2023-05-24 14:35:29)

Rename `peak_memory` to `peak_measured_mem` in line with other memory settings


Patch:
diff --git a/cubed/core/array.py b/cubed/core/array.py
index 7f497514..52cc98e4 100644
--- a/cubed/core/array.py
+++ b/cubed/core/array.py
@@ -280,11 +280,11 @@ class TaskEndEvent:
     task_result_tstamp: Optional[float] = None
     """Timestamp of when the result of the task was received by the client."""
 
-    peak_memory_start: Optional[int] = None
-    """Peak memory usage on the remote worker before the function starts executing."""
+    peak_measured_mem_start: Optional[int] = None
+    """Peak memory usage measured on the remote worker before the function starts executing."""
 
-    peak_memory_end: Optional[int] = None
-    """Peak memory usage on the remote worker after the function finishes executing."""
+    peak_measured_mem_end: Optional[int] = None
+    """Peak memory usage measured on the remote worker after the function finishes executing."""
 
 
 def check_array_specs(arrays):
@@ -353,9 +353,9 @@ def visualize(*arrays, filename="cubed", format=None, optimize_graph=True):
     )
 
 
-class PeakMemoryCallback(Callback):
+class PeakMeasuredMemoryCallback(Callback):
     def on_task_end(self, event):
-        self.peak_memory = event.peak_memory_end
+        self.peak_measured_mem = event.peak_measured_mem_end
 
 
 def measure_reserved_memory(executor):
@@ -386,6 +386,6 @@ def measure_reserved_memory(executor):
 
     a = xp.ones((1,))
     b = xp.negative(a)
-    peak_memory_callback = PeakMemoryCallback()
-    b.compute(executor=executor, callbacks=[peak_memory_callback])
-    return peak_memory_callback.peak_memory
+    peak_measured_mem_callback = PeakMeasuredMemoryCallback()
+    b.compute(executor=executor, callbacks=[peak_measured_mem_callback])
+    return peak_measured_mem_callback.peak_measured_mem
diff --git a/cubed/extensions/history.py b/cubed/extensions/history.py
index f27212f0..c855ac77 100644
--- a/cubed/extensions/history.py
+++ b/cubed/extensions/history.py
@@ -57,18 +57,22 @@ def analyze(plan_df, events_df):
             "num_tasks",
         ]
     ]
-    events_df["peak_mem_start_mb"] = events_df["peak_memory_start"] / 1_000_000
-    events_df["peak_mem_end_mb"] = events_df["peak_memory_end"] / 1_000_000
-    events_df["peak_mem_delta_mb"] = (
-        events_df["peak_mem_end_mb"] - events_df["peak_mem_start_mb"]
+    events_df["peak_measured_mem_start_mb"] = (
+        events_df["peak_measured_mem_start"] / 1_000_000
+    )
+    events_df["peak_measured_mem_end_mb"] = (
+        events_df["peak_measured_mem_end"] / 1_000_000
+    )
+    events_df["peak_measured_mem_delta_mb"] = (
+        events_df["peak_measured_mem_end_mb"] - events_df["peak_measured_mem_start_mb"]
     )
 
     # find per-array stats
     df = events_df.groupby("array_name", as_index=False).agg(
         {
-            "peak_mem_start_mb": ["min", "mean", "max"],
-            "peak_mem_end_mb": ["max"],
-            "peak_mem_delta_mb": ["min", "mean", "max"],
+            "peak_measured_mem_start_mb": ["min", "mean", "max"],
+            "peak_measured_mem_end_mb": ["max"],
+            "peak_measured_mem_delta_mb": ["min", "mean", "max"],
         }
     )
 
@@ -77,7 +81,7 @@ def analyze(plan_df, events_df):
     df = df.merge(plan_df, on="array_name")
 
     def projected_mem_utilization(row):
-        return row["peak_mem_end_mb_max"] / row["projected_mem_mb"]
+        return row["peak_measured_mem_end_mb_max"] / row["projected_mem_mb"]
 
     df["projected_mem_utilization"] = df.apply(
         lambda row: projected_mem_utilization(row), axis=1
@@ -87,9 +91,9 @@ def analyze(plan_df, events_df):
             "array_name",
             "op_name",
             "num_tasks",
-            "peak_mem_start_mb_max",
-            "peak_mem_end_mb_max",
-            "peak_mem_delta_mb_max",
+            "peak_measured_mem_start_mb_max",
+            "peak_measured_mem_end_mb_max",
+            "peak_measured_mem_delta_mb_max",
             "projected_mem_mb",
             "reserved_mem_mb",
             "projected_mem_utilization",
diff --git a/cubed/runtime/executors/lithops.py b/cubed/runtime/executors/lithops.py
index 57bc10e2..5c832094 100644
--- a/cubed/runtime/executors/lithops.py
+++ b/cubed/runtime/executors/lithops.py
@@ -200,8 +200,8 @@ def lithops_stats_to_task_end_event(name, stats):
         function_start_tstamp=stats["worker_func_start_tstamp"],
         function_end_tstamp=stats["worker_func_end_tstamp"],
         task_result_tstamp=stats["host_status_done_tstamp"],
-        peak_memory_start=stats["worker_peak_memory_start"],
-        peak_memory_end=stats["worker_peak_memory_end"],
+        peak_measured_mem_start=stats["worker_peak_memory_start"],
+        peak_measured_mem_end=stats["worker_peak_memory_end"],
     )
 
 
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
index 62dccfec..76186f18 100644
--- a/cubed/runtime/executors/modal.py
+++ b/cubed/runtime/executors/modal.py
@@ -9,7 +9,7 @@ from tenacity import retry, retry_if_exception_type, stop_after_attempt
 from cubed.core.array import TaskEndEvent
 from cubed.core.plan import visit_nodes
 from cubed.runtime.types import DagExecutor
-from cubed.utils import peak_memory
+from cubed.utils import peak_measured_mem
 
 stub = modal.Stub()
 
@@ -43,17 +43,17 @@ else:
 )
 def run_remotely(input, func=None, config=None):
     print(f"running remotely on {input}")
-    peak_memory_start = peak_memory()
+    peak_measured_mem_start = peak_measured_mem()
     function_start_tstamp = time.time()
     result = func(input, config=config)
     function_end_tstamp = time.time()
-    peak_memory_end = peak_memory()
+    peak_measured_mem_end = peak_measured_mem()
     yield (
         result,
         function_start_tstamp,
         function_end_tstamp,
-        peak_memory_start,
-        peak_memory_end,
+        peak_measured_mem_start,
+        peak_measured_mem_end,
     )
 
 
@@ -86,16 +86,16 @@ def handle_callbacks(callbacks, array_name, result, task_create_tstamp):
             res,
             function_start_tstamp,
             function_end_tstamp,
-            peak_memory_start,
-            peak_memory_end,
+            peak_measured_mem_start,
+            peak_measured_mem_end,
         ) = result
         task_stats = dict(
             task_create_tstamp=task_create_tstamp,
             function_start_tstamp=function_start_tstamp,
             function_end_tstamp=function_end_tstamp,
             task_result_tstamp=task_result_tstamp,
-            peak_memory_start=peak_memory_start,
-            peak_memory_end=peak_memory_end,
+            peak_measured_mem_start=peak_measured_mem_start,
+            peak_measured_mem_end=peak_measured_mem_end,
         )
         event = TaskEndEvent(array_name=array_name, **task_stats)
         [callback.on_task_end(event) for callback in callbacks]
diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
index c6fa86c1..8935de43 100644
--- a/cubed/runtime/executors/modal_async.py
+++ b/cubed/runtime/executors/modal_async.py
@@ -13,7 +13,7 @@ from cubed.core.array import TaskEndEvent
 from cubed.core.plan import visit_nodes
 from cubed.runtime.backup import should_launch_backup
 from cubed.runtime.types import DagExecutor
-from cubed.utils import peak_memory
+from cubed.utils import peak_measured_mem
 
 async_stub = modal.aio.AioStub()
 
@@ -45,17 +45,17 @@ else:
 )
 async def async_run_remotely(input, func=None, config=None):
     print(f"running remotely on {input}")
-    peak_memory_start = peak_memory()
+    peak_measured_mem_start = peak_measured_mem()
     function_start_tstamp = time.time()
     result = func(input, config=config)
     function_end_tstamp = time.time()
-    peak_memory_end = peak_memory()
+    peak_measured_mem_end = peak_measured_mem()
     yield (
         result,
         function_start_tstamp,
         function_end_tstamp,
-        peak_memory_start,
-        peak_memory_end,
+        peak_measured_mem_start,
+        peak_measured_mem_end,
     )
 
 
@@ -156,16 +156,16 @@ def handle_callbacks(callbacks, array_name, result, task_create_tstamp):
             res,
             function_start_tstamp,
             function_end_tstamp,
-            peak_memory_start,
-            peak_memory_end,
+            peak_measured_mem_start,
+            peak_measured_mem_end,
         ) = result
         task_stats = dict(
             task_create_tstamp=task_create_tstamp,
             function_start_tstamp=function_start_tstamp,
             function_end_tstamp=function_end_tstamp,
             task_result_tstamp=task_result_tstamp,
-            peak_memory_start=peak_memory_start,
-            peak_memory_end=peak_memory_end,
+            peak_measured_mem_start=peak_measured_mem_start,
+            peak_measured_mem_end=peak_measured_mem_end,
         )
         event = TaskEndEvent(array_name=array_name, **task_stats)
         [callback.on_task_end(event) for callback in callbacks]
diff --git a/cubed/runtime/executors/python_async.py b/cubed/runtime/executors/python_async.py
index 13ce25bd..92fb5df5 100644
--- a/cubed/runtime/executors/python_async.py
+++ b/cubed/runtime/executors/python_async.py
@@ -9,23 +9,23 @@ from tenacity import retry, stop_after_attempt
 from cubed.core.array import TaskEndEvent
 from cubed.core.plan import visit_node_generations
 from cubed.runtime.types import DagExecutor
-from cubed.utils import peak_memory
+from cubed.utils import peak_measured_mem
 
 
 def execution_stats(func):
     """Measure timing information and peak memory usage of a function call."""
 
     def wrapper(*args, **kwargs):
-        peak_memory_start = peak_memory()
+        peak_measured_mem_start = peak_measured_mem()
         function_start_tstamp = time.time()
         result = func(*args, **kwargs)
         function_end_tstamp = time.time()
-        peak_memory_end = peak_memory()
+        peak_measured_mem_end = peak_measured_mem()
         return result, dict(
             function_start_tstamp=function_start_tstamp,
             function_end_tstamp=function_end_tstamp,
-            peak_memory_start=peak_memory_start,
-            peak_memory_end=peak_memory_end,
+            peak_measured_mem_start=peak_measured_mem_start,
+            peak_measured_mem_end=peak_measured_mem_end,
         )
 
     return wrapper
diff --git a/cubed/utils.py b/cubed/utils.py
index 96a506b7..656f0b21 100644
--- a/cubed/utils.py
+++ b/cubed/utils.py
@@ -70,14 +70,14 @@ def memory_repr(num):
         num /= 1000.0
 
 
-def peak_memory():
-    """Return the peak memory usage in bytes.
+def peak_measured_mem():
+    """Measures the peak memory usage in bytes.
 
     Note: this function currently doesn't work on Windows.
     """
 
     if platform.system() == "Windows":
-        raise NotImplementedError("`peak_memory` is not implemented on Windows")
+        raise NotImplementedError("`peak_measured_mem` is not implemented on Windows")
 
     from resource import RUSAGE_SELF, getrusage
 


Test patch:
diff --git a/cubed/tests/test_utils.py b/cubed/tests/test_utils.py
index 71595feb..30782341 100644
--- a/cubed/tests/test_utils.py
+++ b/cubed/tests/test_utils.py
@@ -9,7 +9,7 @@ from cubed.utils import (
     extract_stack_summaries,
     join_path,
     memory_repr,
-    peak_memory,
+    peak_measured_mem,
     to_chunksize,
 )
 
@@ -57,8 +57,8 @@ def test_memory_repr():
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")
-def test_peak_memory():
-    assert peak_memory() > 0
+def test_peak_measured_mem():
+    assert peak_measured_mem() > 0
 
 
 def test_extract_stack_summaries():
diff --git a/cubed/tests/utils.py b/cubed/tests/utils.py
index dac9b7cf..7976d155 100644
--- a/cubed/tests/utils.py
+++ b/cubed/tests/utils.py
@@ -21,7 +21,7 @@ MAIN_EXECUTORS = [PythonPipelineExecutor(), PythonDagExecutor()]
 
 
 if platform.system() != "Windows":
-    # AsyncPythonDagExecutor calls `peak_memory` which is not supported on Windows
+    # AsyncPythonDagExecutor calls `peak_measured_mem` which is not supported on Windows
     ALL_EXECUTORS.append(AsyncPythonDagExecutor())