ASV better memory benchmarks

benchmarks/face_bounds.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 
 import uxarray as ux
+from .helpers._memsize import grid_nbytes
+from .helpers._peakmem import numba_threads, peak_allocated
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
 
@@ -14,8 +16,15 @@ class FaceBounds:
 
     params = [grid_quad_hex, grid_geoflow, grid_scrip, grid_mpas]
 
+    number = 1
+    warmup_time = 0
 
     def setup(self, grid_path):
+        # Warmed on the smallest grid in ``params`` so the njit kernels are
+        # compiled before anything is measured. ``track_peakmem_*`` would
+        # otherwise charge the first sample for loading them off numba's disk
+        # cache, which inflates the reported peak by ~3%.
+        ux.open_grid(grid_quad_hex).bounds
         self.uxgrid = ux.open_grid(grid_path)
 
     def teardown(self, n):
@@ -25,6 +34,50 @@ def time_face_bounds(self, grid_path):
         """Time to obtain ``Grid.face_bounds``"""
         self.uxgrid.bounds
 
-    def peakmem_face_bounds(self, grid_path):
-        """Peak memory usage obtain ``Grid.face_bounds."""
-        face_bounds = self.uxgrid.bounds
+    def track_nbytes_face_bounds(self, grid_path):
+        """Size of the materialized ``Grid.face_bounds`` array."""
+        return self.uxgrid.bounds.nbytes
+
+    track_nbytes_face_bounds.unit = "bytes"
+
+    def track_nbytes_grid_with_bounds(self, grid_path):
+        """Grid footprint after populating bounds -- catches cached arrays that
+        ``bounds`` adds to the ``Grid`` beyond the returned array itself."""
+        self.uxgrid.bounds
+        return grid_nbytes(self.uxgrid)
+
+    track_nbytes_grid_with_bounds.unit = "bytes"
+
+    def track_peakmem_face_bounds(self, grid_path):
+        """Transient high-water allocation of populating ``Grid.face_bounds``.
+
+        The kernel behind ``bounds`` is ``parallel=True``, hence the pinning --
+        see :func:`~benchmarks.helpers._peakmem.numba_threads`.
+        """
+        with numba_threads(1):
+            return peak_allocated(lambda: self.uxgrid.bounds)
+
+    track_peakmem_face_bounds.unit = "bytes"
+
+
+class FaceBoundsColdStartRss:
+    """Peak memory of a cold start: import uxarray, open a grid, get its bounds.
+
+    Whole-process ``ru_maxrss``, not tracemalloc -- the ~250MB uxarray import is
+    part of the number by design, because the cold start is the subject. For the
+    cost of ``bounds`` alone see ``FaceBounds.track_peakmem_face_bounds``, which
+    runs one to three orders of magnitude lower.
+    """
+
+    params = FaceBounds.params
+    param_names = ["grid_path"]
+
+    def setup_cache(self):
+        """Compile the njit kernels before anything is measured."""
+        for grid_path in self.params:
+            ux.open_grid(grid_path).bounds
+
+    setup_cache.timeout = 1800
+
+    def peakmem_open_and_bounds(self, grid_path):
+        ux.open_grid(grid_path).bounds

benchmarks/helpers/_memsize.py

@@ -0,0 +1,12 @@
+
+__all__ = ["grid_nbytes", "dataset_nbytes"]
+
+
+def grid_nbytes(uxgrid):
+    """Total size of the arrays a ``Grid`` currently holds, in bytes."""
+    return uxgrid._ds.nbytes
+
+
+def dataset_nbytes(uxds):
+    """Total size of a ``UxDataset``, in bytes, including its grid."""
+    return uxds.nbytes + grid_nbytes(uxds.uxgrid)

benchmarks/helpers/_peakmem.py

@@ -0,0 +1,80 @@
+
+import contextlib
+import subprocess
+import sys
+
+import numba
+
+__all__ = ["peak_allocated", "numba_threads", "subprocess_peak_rss"]
+
+
+def peak_allocated(build):
+    """Bytes held at the high-water allocation point of ``build``.
+
+    ``tracemalloc.start`` begins with an empty trace table, so whatever the
+    process is already holding -- including everything allocated in ``setup`` --
+    is excluded, and only what ``build`` allocates counts. A ``reset_peak()``
+    here would be a no-op for that reason.
+    """
+    # Imported here rather than at module scope: asv preimports every benchmark
+    # module under its default ``forkserver`` launch method
+    import tracemalloc
+
+    if tracemalloc.is_tracing():
+        raise RuntimeError("tracemalloc is already tracing")
+
+    # nframe=1: the reported peak is identical at any traceback depth, while the
+    # cost is not -- nframe=25 runs 26x slower.
+    tracemalloc.start(1)
+    try:
+        build()
+        _, peak = tracemalloc.get_traced_memory()
+    finally:
+        tracemalloc.stop()
+    return peak
+
+
+@contextlib.contextmanager
+def numba_threads(n):
+    """Runs the block with numba's thread pool held at ``n``.
+
+    Tracing serializes on tracemalloc's global allocator lock, so a
+    ``parallel=True`` kernel under :func:`peak_allocated` degrades badly as
+    threads contend for it.
+    """
+    restore = numba.get_num_threads()
+    numba.set_num_threads(n)
+    try:
+        yield
+    finally:
+        numba.set_num_threads(restore)
+
+
+# ``ru_maxrss`` is bytes on macOS and kilobytes elsewhere, mirroring
+# ``asv_runner/benchmarks/_maxrss.py:117,132``.
+_MAXRSS_TO_BYTES = 1 if sys.platform == "darwin" else 1024
+
+
+def subprocess_peak_rss(statement):
+    """Peak resident bytes of a fresh interpreter that has run ``statement``.
+
+    For the cases an in-process metric cannot reach: measuring an import, whose
+    cost is already paid before a ``peakmem_*`` body runs, and measuring
+    ``numba.typed`` containers, which tracemalloc does not see at all.
+
+    The child reports its own ``ru_maxrss`` rather than the parent reading
+    ``RUSAGE_CHILDREN``, which is a maximum over every child that has exited and
+    so would not isolate this one.
+    """
+    reporter = (
+        "import resource, sys\n"
+        f"exec({statement!r})\n"
+        "sys.stdout.write(str(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss))\n"
+    )
+    completed = subprocess.run(
+        [sys.executable, "-c", reporter],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return int(completed.stdout) * _MAXRSS_TO_BYTES

benchmarks/import.py

@@ -1,5 +1,14 @@
+from .helpers._peakmem import subprocess_peak_rss
+
+
 class Imports:
     """Benchmark importing uxarray."""
 
     def timeraw_import_uxarray(self):
         return "import uxarray"
+
+    def track_peakmem_import_uxarray(self):
+        """Peak resident memory of a process that has imported uxarray."""
+        return subprocess_peak_rss("import uxarray")
+
+    track_peakmem_import_uxarray.unit = "bytes"

benchmarks/mpas_ocean.py

@@ -6,6 +6,9 @@
 
 import uxarray as ux
 
+from .helpers._memsize import grid_nbytes
+from .helpers._peakmem import numba_threads, peak_allocated
+
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
 data_var = 'bottomDepth'
@@ -57,27 +60,92 @@ def teardown(self, resolution, *args, **kwargs):
 
 
 class FaceAreas(GridBenchmark):
-    def time_compute_face_areas(self, resolution):
-        self.uxgrid.compute_face_areas()
+    number = 1
+    warmup_time = 0
+
+    def setup(self, resolution, *args, **kwargs):
+        # The coarsest grid, purely to compile the njit kernel
+        warmup_grid = ux.open_grid(file_path_dict[self.params[0][0]][0])
+        _ = warmup_grid.face_areas
+        super().setup(resolution, *args, **kwargs)
+        self.uxgrid._ds = self.uxgrid._ds.drop_vars("face_areas", errors="ignore")
+
+    def time_face_areas(self, resolution):
+        _ = self.uxgrid.face_areas
+
+    def track_nbytes_face_areas(self, resolution):
+        """Size of the materialized ``Grid.face_areas`` array."""
+        return self.uxgrid.face_areas.nbytes
+
+    track_nbytes_face_areas.unit = "bytes"
 
-    def peakmem_compute_face_areas(self, resolution):
-        self.uxgrid.compute_face_areas()
+    def track_peakmem_face_areas(self, resolution):
+        """Transient high-water allocation of computing ``Grid.face_areas``."""
+        with numba_threads(1):
+            return peak_allocated(lambda: self.uxgrid.face_areas)
+
+    track_peakmem_face_areas.unit = "bytes"
 
 
 class Gradient(DatasetBenchmark):
+    def setup(self, resolution, *args, **kwargs):
+        super().setup(resolution, *args, **kwargs)
+        # Compiles the gradient kernels on the coarsest grid
+        grid, data = file_path_dict[self.params[0][0]]
+        _ = ux.open_dataset(grid, data)[data_var].gradient()
+
     def time_gradient(self, resolution):
         self.uxds[data_var].gradient()
 
-    def peakmem_gradient(self, resolution):
-        grad = self.uxds[data_var].gradient()
+    def track_nbytes_gradient(self, resolution):
+        """Size of the gradient result."""
+        return self.uxds[data_var].gradient().nbytes
+
+    track_nbytes_gradient.unit = "bytes"
+
+    def track_peakmem_gradient(self, resolution):
+        """Transient high-water allocation of taking a gradient."""
+        return peak_allocated(lambda: self.uxds[data_var].gradient())
+
+    track_peakmem_gradient.unit = "bytes"
 
 
 class Integrate(DatasetBenchmark):
+
     def time_integrate(self, resolution):
         self.uxds[data_var].integrate()
 
-    def peakmem_integrate(self, resolution):
-        integral = self.uxds[data_var].integrate()
+    def track_nbytes_integrate(self, resolution):
+        """Grid footprint after integrating."""
+        self.uxds[data_var].integrate()
+        return grid_nbytes(self.uxds.uxgrid)
+
+    track_nbytes_integrate.unit = "bytes"
+
+
+class GradientColdStartRss:
+    """Peak memory of a cold start: import uxarray, open a dataset, take a gradient.
+
+    Whole-process ``ru_maxrss``, not tracemalloc -- the ~250MB uxarray import is
+    part of the number by design, because the cold start is the subject. For the
+    gradient's own transient cost see ``Gradient.track_peakmem_gradient``, which
+    runs one to three orders of magnitude lower.
+    """
+
+    param_names = ["resolution"]
+    params = [["480km", "120km"]]
+
+    def setup_cache(self):
+        """Compile the njit kernels before anything is measured."""
+        for resolution in self.params[0]:
+            grid, data = file_path_dict[resolution]
+            ux.open_dataset(grid, data)[data_var].gradient()
+
+    setup_cache.timeout = 1800
+
+    def peakmem_gradient(self, resolution):
+        grid, data = file_path_dict[resolution]
+        ux.open_dataset(grid, data)[data_var].gradient()
 
 
 class GeoDataFrame(DatasetBenchmark):
@@ -90,7 +158,7 @@ def time_to_geodataframe(self, resolution, exclude_antimeridian):
 
 class ConnectivityConstruction(DatasetBenchmark):
     def time_n_nodes_per_face(self, resolution):
-        self.uxds.uxgrid.n_nodes_per_face
+        _ = self.uxds.uxgrid.n_nodes_per_face
 
     def time_face_face_connectivity(self, resolution):
         ux.grid.connectivity._populate_face_face_connectivity(self.uxds.uxgrid)

benchmarks/quad_hexagon.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 
 import uxarray as ux
+from .helpers._memsize import dataset_nbytes, grid_nbytes
+from .helpers._peakmem import peak_allocated
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
 
@@ -10,27 +12,40 @@
 
 
 class QuadHexagon:
+    def setup(self):
+        # Opening a grid for the first time in a process pulls in xarray's
+        # backend machinery and the netCDF library
+        ux.open_grid(grid_path)
+        ux.open_dataset(grid_path, data_path)
+
     def time_open_grid(self):
         """Time to open a `Grid`"""
         ux.open_grid(grid_path)
 
-    # def mem_open_grid(self):
-    #     """Memory Occupied by a `Grid`"""
-    #     return ux.open_grid(grid_path)
+    def track_nbytes_open_grid(self):
+        """Memory occupied by a `Grid`"""
+        return grid_nbytes(ux.open_grid(grid_path))
 
-    def peakmem_open_grid(self):
-        """Peak memory usage of a `Grid`"""
-        uxgrid = ux.open_grid(grid_path)
+    track_nbytes_open_grid.unit = "bytes"
 
+    def track_peakmem_open_grid(self):
+        """Transient high-water allocation of opening a `Grid`"""
+        return peak_allocated(lambda: ux.open_grid(grid_path))
+
+    track_peakmem_open_grid.unit = "bytes"
 
     def time_open_dataset(self):
         """Time to open a `UxDataset`"""
         ux.open_dataset(grid_path, data_path)
 
-    # def mem_open_dataset(self):
-    #     """Memory occupied by a `UxDataset`"""
-    #     return ux.open_dataset(grid_path, data_path)
+    def track_nbytes_open_dataset(self):
+        """Memory occupied by a `UxDataset`, including its grid"""
+        return dataset_nbytes(ux.open_dataset(grid_path, data_path))
+
+    track_nbytes_open_dataset.unit = "bytes"
+
+    def track_peakmem_open_dataset(self):
+        """Transient high-water allocation of opening a `UxDataset`"""
+        return peak_allocated(lambda: ux.open_dataset(grid_path, data_path))
 
-    def peakmem_open_dataset(self):
-        """Peak memory usage of a `UxDataset`"""
-        uxds = ux.open_dataset(grid_path, data_path)
+    track_peakmem_open_dataset.unit = "bytes"