Repeating Tracemalloc Benchmark for accuracy (#5981)

* PoC custom benchmark.

* Working tracemalloc subclass.

* Docstrings and comments.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Finalised code.

* Make custom benchmarks installable.

* Experiment.

* Re-write following testing.

* Replace pyproject with setup.py.

* Testing.

* Make testing smaller.

* Better testing.

* Better testing 2.

* Better testing 3.

* Better testing 4.

* Better testing 5.

* Remove testing.

* Remove tracemalloc decorator.

* Docs tidy-up.

* Restructure and use custom install script.

* Temporary quick demo.

* Revert "Temporary quick demo."

This reverts commit 9dffa1b.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

benchmarks/README.md

@@ -118,6 +118,12 @@ repeats _between_ `setup()` calls using the `repeat` attribute.
 `warmup_time = 0` is also advisable since ASV performs independent re-runs to
 estimate run-time, and these will still be subject to the original problem.
 
+### Custom benchmarks
+
+Iris benchmarking implements custom benchmark types, such as a `tracemalloc`
+benchmark to measure memory growth. See [custom_bms/](./custom_bms) for more
+detail.
+
 ### Scaling / non-Scaling Performance Differences
 
 **(We no longer advocate the below for benchmarks run during CI, given the

benchmarks/asv.conf.json

@@ -53,9 +53,12 @@
     "command_comment": [
         "We know that the Nox command takes care of installation in each",
         "environment, and in the case of Iris no specialised uninstall or",
-        "build commands are needed to get it working."
+        "build commands are needed to get it working.",
+
+        "We do however need to install the custom benchmarks for them to be",
+        "usable."
     ],
     "install_command": [],
     "uninstall_command": [],
-    "build_command": []
+    "build_command": ["python {conf_dir}/custom_bms/install.py"]
 }

benchmarks/benchmarks/__init__.py

@@ -37,111 +37,6 @@ def disable_repeat_between_setup(benchmark_object):
     return benchmark_object
 
 
-class TrackAddedMemoryAllocation:
-    """Measures by how much process resident memory grew, during execution.
-
-    Context manager which measures by how much process resident memory grew,
-    during execution of its enclosed code block.
-
-    Obviously limited as to what it actually measures : Relies on the current
-    process not having significant unused (de-allocated) memory when the
-    tested codeblock runs, and only reliable when the code allocates a
-    significant amount of new memory.
-
-    Example:
-        with TrackAddedMemoryAllocation() as mb:
-            initial_call()
-            other_call()
-        result = mb.addedmem_mb()
-
-    Attributes
-    ----------
-    RESULT_MINIMUM_MB : float
-        The smallest result that should ever be returned, in Mb. Results
-        fluctuate from run to run (usually within 1Mb) so if a result is
-        sufficiently small this noise will produce a before-after ratio over
-        AVD's detection threshold and be treated as 'signal'. Results
-        smaller than this value will therefore be returned as equal to this
-        value, ensuring fractionally small noise / no noise at all.
-        Defaults to 1.0
-
-    RESULT_ROUND_DP : int
-        Number of decimal places of rounding on result values (in Mb).
-        Defaults to 1
-
-    """
-
-    RESULT_MINIMUM_MB = 0.2
-    RESULT_ROUND_DP = 1  # I.E. to nearest 0.1 Mb
-
-    def __enter__(self):
-        tracemalloc.start()
-        return self
-
-    def __exit__(self, *_):
-        _, peak_mem_bytes = tracemalloc.get_traced_memory()
-        tracemalloc.stop()
-        # Save peak-memory allocation, scaled from bytes to Mb.
-        self._peak_mb = peak_mem_bytes * (2.0**-20)
-
-    def addedmem_mb(self):
-        """Return measured memory growth, in Mb."""
-        result = self._peak_mb
-        # Small results are too vulnerable to noise being interpreted as signal.
-        result = max(self.RESULT_MINIMUM_MB, result)
-        # Rounding makes results easier to read.
-        result = np.round(result, self.RESULT_ROUND_DP)
-        return result
-
-    @staticmethod
-    def decorator(decorated_func):
-        """Benchmark to track growth in resident memory during execution.
-
-        Intended for use on ASV ``track_`` benchmarks. Applies the
-        :class:`TrackAddedMemoryAllocation` context manager to the benchmark
-        code, sets the benchmark ``unit`` attribute to ``Mb``.
-
-        """
-
-        def _wrapper(*args, **kwargs):
-            assert decorated_func.__name__[:6] == "track_"
-            # Run the decorated benchmark within the added memory context
-            # manager.
-            with TrackAddedMemoryAllocation() as mb:
-                decorated_func(*args, **kwargs)
-            return mb.addedmem_mb()
-
-        decorated_func.unit = "Mb"
-        return _wrapper
-
-    @staticmethod
-    def decorator_repeating(repeats=3):
-        """Benchmark to track growth in resident memory during execution.
-
-        Tracks memory for repeated calls of decorated function.
-
-        Intended for use on ASV ``track_`` benchmarks. Applies the
-        :class:`TrackAddedMemoryAllocation` context manager to the benchmark
-        code, sets the benchmark ``unit`` attribute to ``Mb``.
-
-        """
-
-        def decorator(decorated_func):
-            def _wrapper(*args, **kwargs):
-                assert decorated_func.__name__[:6] == "track_"
-                # Run the decorated benchmark within the added memory context
-                # manager.
-                with TrackAddedMemoryAllocation() as mb:
-                    for _ in range(repeats):
-                        decorated_func(*args, **kwargs)
-                return mb.addedmem_mb()
-
-            decorated_func.unit = "Mb"
-            return _wrapper
-
-        return decorator
-
-
 def on_demand_benchmark(benchmark_object):
     """Disable these benchmark(s) unless ON_DEMAND_BENCHARKS env var is set.
 

benchmarks/benchmarks/cperf/save.py

@@ -6,7 +6,7 @@
 
 from iris import save
 
-from .. import TrackAddedMemoryAllocation, on_demand_benchmark
+from .. import on_demand_benchmark
 from ..generate_data.ugrid import make_cube_like_2d_cubesphere, make_cube_like_umfield
 from . import _N_CUBESPHERE_UM_EQUIVALENT, _UM_DIMS_YX
 
@@ -36,6 +36,5 @@ def _save_data(self, cube):
     def time_save_data_netcdf(self, data_type):
         self._save_data(self.cube)
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_save_data_netcdf(self, data_type):
+    def tracemalloc_save_data_netcdf(self, data_type):
         self._save_data(self.cube)

benchmarks/benchmarks/merge_concat.py

@@ -8,7 +8,6 @@
 
 from iris.cube import CubeList
 
-from . import TrackAddedMemoryAllocation
 from .generate_data.stock import realistic_4d_w_everything
 
 
@@ -34,10 +33,11 @@ def setup(self):
     def time_merge(self):
         _ = self.cube_list.merge_cube()
 
-    @TrackAddedMemoryAllocation.decorator_repeating()
-    def track_mem_merge(self):
+    def tracemalloc_merge(self):
         _ = self.cube_list.merge_cube()
 
+    tracemalloc_merge.number = 3  # type: ignore[attr-defined]
+
 
 class Concatenate:
     # TODO: Improve coverage.
@@ -56,6 +56,7 @@ def setup(self):
     def time_concatenate(self):
         _ = self.cube_list.concatenate_cube()
 
-    @TrackAddedMemoryAllocation.decorator_repeating()
-    def track_mem_merge(self):
+    def tracemalloc_concatenate(self):
         _ = self.cube_list.concatenate_cube()
+
+    tracemalloc_concatenate.number = 3  # type: ignore[attr-defined]

benchmarks/benchmarks/mesh/utils/regions_combine.py

@@ -17,7 +17,6 @@
 from iris import load, load_cube, save
 from iris.mesh.utils import recombine_submeshes
 
-from ... import TrackAddedMemoryAllocation
 from ...generate_data.ugrid import make_cube_like_2d_cubesphere
 
 
@@ -169,8 +168,7 @@ def setup(self, n_cubesphere):
     def time_create_combined_cube(self, n_cubesphere):
         self.recombine()
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_create_combined_cube(self, n_cubesphere):
+    def tracemalloc_create_combined_cube(self, n_cubesphere):
         self.recombine()
 
 
@@ -180,8 +178,7 @@ class CombineRegionsComputeRealData(MixinCombineRegions):
     def time_compute_data(self, n_cubesphere):
         _ = self.recombined_cube.data
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_compute_data(self, n_cubesphere):
+    def tracemalloc_compute_data(self, n_cubesphere):
         _ = self.recombined_cube.data
 
 
@@ -199,8 +196,7 @@ def time_save(self, n_cubesphere):
         # Save to disk, which must compute data + stream it to file.
         save(self.recombined_cube, "tmp.nc")
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_save(self, n_cubesphere):
+    def tracemalloc_save(self, n_cubesphere):
         save(self.recombined_cube, "tmp.nc")
 
     def track_filesize_saved(self, n_cubesphere):
@@ -227,6 +223,5 @@ def time_stream_file2file(self, n_cubesphere):
         # Save to disk, which must compute data + stream it to file.
         save(self.recombined_cube, "tmp.nc")
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_stream_file2file(self, n_cubesphere):
+    def tracemalloc_stream_file2file(self, n_cubesphere):
         save(self.recombined_cube, "tmp.nc")

benchmarks/benchmarks/regridding.py

@@ -14,8 +14,6 @@
 from iris.analysis import AreaWeighted, PointInCell
 from iris.coords import AuxCoord
 
-from . import TrackAddedMemoryAllocation
-
 
 class HorizontalChunkedRegridding:
     def setup(self) -> None:
@@ -53,20 +51,22 @@ def time_regrid_area_w_new_grid(self) -> None:
         # Realise data
         out.data
 
-    @TrackAddedMemoryAllocation.decorator_repeating()
-    def track_mem_regrid_area_w(self) -> None:
+    def tracemalloc_regrid_area_w(self) -> None:
         # Regrid the chunked cube
         out = self.cube.regrid(self.template_cube, self.scheme_area_w)
         # Realise data
         out.data
 
-    @TrackAddedMemoryAllocation.decorator_repeating()
-    def track_mem_regrid_area_w_new_grid(self) -> None:
+    tracemalloc_regrid_area_w.number = 3  # type: ignore[attr-defined]
+
+    def tracemalloc_regrid_area_w_new_grid(self) -> None:
         # Regrid the chunked cube
         out = self.chunked_cube.regrid(self.template_cube, self.scheme_area_w)
         # Realise data
         out.data
 
+    tracemalloc_regrid_area_w_new_grid.number = 3  # type: ignore[attr-defined]
+
 
 class CurvilinearRegridding:
     def setup(self) -> None:
@@ -110,9 +110,10 @@ def time_regrid_pic(self) -> None:
         # Realise the data
         out.data
 
-    @TrackAddedMemoryAllocation.decorator_repeating()
-    def track_mem_regrid_pic(self) -> None:
+    def tracemalloc_regrid_pic(self) -> None:
         # Regrid the cube onto the template.
         out = self.cube.regrid(self.template_cube, self.scheme_pic)
         # Realise the data
         out.data
+
+    tracemalloc_regrid_pic.number = 3  # type: ignore[attr-defined]

benchmarks/benchmarks/save.py

@@ -7,7 +7,6 @@
 from iris import save
 from iris.mesh import save_mesh
 
-from . import TrackAddedMemoryAllocation, on_demand_benchmark
 from .generate_data.ugrid import make_cube_like_2d_cubesphere
 
 
@@ -38,8 +37,7 @@ def time_netcdf_save_mesh(self, n_cubesphere, is_unstructured):
         if is_unstructured:
             self._save_mesh(self.cube)
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_netcdf_save(self, n_cubesphere, is_unstructured):
+    def tracemalloc_netcdf_save(self, n_cubesphere, is_unstructured):
         # Don't need to copy the cube here since track_ benchmarks don't
         #  do repeats between self.setup() calls.
         self._save_data(self.cube, do_copy=False)

benchmarks/benchmarks/sperf/combine_regions.py

@@ -12,7 +12,7 @@
 from iris import load, load_cube, save
 from iris.mesh.utils import recombine_submeshes
 
-from .. import TrackAddedMemoryAllocation, on_demand_benchmark
+from .. import on_demand_benchmark
 from ..generate_data.ugrid import BENCHMARK_DATA, make_cube_like_2d_cubesphere
 
 
@@ -175,8 +175,7 @@ def setup(self, n_cubesphere, imaginary_data=True, create_result_cube=False):
     def time_create_combined_cube(self, n_cubesphere):
         self.recombine()
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_create_combined_cube(self, n_cubesphere):
+    def tracemalloc_create_combined_cube(self, n_cubesphere):
         self.recombine()
 
 
@@ -187,8 +186,7 @@ class ComputeRealData(Mixin):
     def time_compute_data(self, n_cubesphere):
         _ = self.recombined_cube.data
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_compute_data(self, n_cubesphere):
+    def tracemalloc_compute_data(self, n_cubesphere):
         _ = self.recombined_cube.data
 
 
@@ -206,8 +204,7 @@ def time_save(self, n_cubesphere):
         # Save to disk, which must compute data + stream it to file.
         self.save_recombined_cube()
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_save(self, n_cubesphere):
+    def tracemalloc_save(self, n_cubesphere):
         self.save_recombined_cube()
 
     def track_filesize_saved(self, n_cubesphere):
@@ -233,6 +230,5 @@ def time_stream_file2file(self, n_cubesphere):
         # Save to disk, which must compute data + stream it to file.
         self.save_recombined_cube()
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_stream_file2file(self, n_cubesphere):
+    def tracemalloc_stream_file2file(self, n_cubesphere):
         self.save_recombined_cube()

benchmarks/benchmarks/sperf/save.py

@@ -9,7 +9,7 @@
 from iris import save
 from iris.mesh import save_mesh
 
-from .. import TrackAddedMemoryAllocation, on_demand_benchmark
+from .. import on_demand_benchmark
 from ..generate_data.ugrid import make_cube_like_2d_cubesphere
 
 
@@ -36,8 +36,7 @@ def _save_mesh(self, cube):
     def time_save_cube(self, n_cubesphere, is_unstructured):
         self._save_cube(self.cube)
 
-    @TrackAddedMemoryAllocation.decorator
-    def track_addedmem_save_cube(self, n_cubesphere, is_unstructured):
+    def tracemalloc_save_cube(self, n_cubesphere, is_unstructured):
         self._save_cube(self.cube)
 
     def time_save_mesh(self, n_cubesphere, is_unstructured):