From 8707c1673bc5beaae88d22a20b0794fc1d84e1fe Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 22 Jul 2026 12:40:01 -0500 Subject: [PATCH 1/8] ASV peakmem benchmark fix Three benchmarks reported a near-identical "improvement" on every PR regardless of what the PR touched: 578M -> 391M 0.68 face_bounds.FaceBounds.peakmem_face_bounds(geoflow-small) 708M -> 390M 0.55 face_bounds.FaceBounds.peakmem_face_bounds(quad-hexagon) 498M -> 384M 0.77 mpas_ocean.Gradient.peakmem_gradient('480km') They were not measuring the operation under test. asv's peakmem_* records the max RSS of the whole process, and per asv's docs it "also counts memory usage during the setup routine". Profiling the face_bounds params: grid import +open_grid +.bounds attributable to op quad-hexagon ( 24K) 236MB 265MB 301MB 36MB (12%) geoflow-small (1.1M) 236MB 263MB 487MB 224MB (46%) outCSne8 ( 48K) 235MB 281MB 317MB 35MB (11%) oQU480 (4.6M) 236MB 270MB 306MB 36MB (12%) `import uxarray` alone is ~226 MB and constant to within 1 MB. oQU480 is 190x larger than quad-hexagon yet both attributed ~36 MB to Grid.bounds -- the benchmark was nearly insensitive to its own workload. The part that did vary was numba: uxarray has 81 @njit(cache=True) kernels and both affected paths go through them (uxarray/grid/bounds.py, uxarray/core/gradient.py). With a cold JIT cache the quad-hexagon case peaked at 599 MB, with a warm one 298 MB -- a ratio of 0.50, matching the ratios seen in CI. Which side of an `asv continuous` comparison paid the compile cost depended on run order, not on the code under review. peakmem_* could not be repaired in place, because there was nothing to measure. Across every operation the five peakmem benchmarks covered, against every mesh in the suite including the 98 MB / 28,571-face oQU120: Grid.bounds ~1 MB open_grid 0-10 MB (lazy; mostly allocator noise) gradient 0-0.1 MB integrate 0.0 MB open_dataset 0.0 MB Two better instruments were evaluated and rejected: - Sampled peak RSS. Validates cleanly (a known 200 MB allocation measures as 200.0 MB) and is immune to process history. But the warm-up call needed to keep JIT out of the measurement leaves freed pages in the allocator, so RSS *growth* undercounts -- every operation measured 0.0-0.1 MB this way. - tracemalloc. Measures allocation volume rather than RSS growth, so it would sidestep page reuse, but it does not observe numba NRT allocations, which is where uxarray's array memory is allocated. What remains is deterministic size accounting via track_* benchmarks, which also sidesteps the setup confound entirely -- track_* does not count setup. Verified byte-identical across cold JIT, warm JIT, and an independent second cold JIT for all 14 parameter combinations, and it scales correctly with the mesh (quad-hexagon reports 128 bytes = 4 faces x 32). It does not capture transient peaks, but the measurements above show those are ~1 MB, far below anything worth gating on, and no available instrument captures them reliably here. The rationale is recorded in benchmarks/_memsize.py so the next person does not reintroduce peakmem_*. The commented-out mem_* stubs in quad_hexagon.py, an earlier attempt at the same thing, are removed. Co-Authored-By: Claude Opus 4.8 --- benchmarks/_memsize.py | 66 ++++++++++++++++++++++++++++++++++++++ benchmarks/face_bounds.py | 17 ++++++++-- benchmarks/mpas_ocean.py | 18 ++++++++--- benchmarks/quad_hexagon.py | 22 ++++++------- 4 files changed, 103 insertions(+), 20 deletions(-) create mode 100644 benchmarks/_memsize.py diff --git a/benchmarks/_memsize.py b/benchmarks/_memsize.py new file mode 100644 index 000000000..bc00a3b0e --- /dev/null +++ b/benchmarks/_memsize.py @@ -0,0 +1,66 @@ +"""Deterministic memory-size accounting for the benchmark suite. + +This module backs the ``track_nbytes_*`` benchmarks that replaced the former +``peakmem_*`` ones. + +Why ``peakmem_*`` was removed +----------------------------- +asv's ``peakmem_*`` records the maximum resident set size of the *whole +process*, and per asv's own docs it "also counts memory usage during the setup +routine". For uxarray that number is almost entirely fixed overhead: + + grid import +open_grid +.bounds attributable to op + quad-hexagon ( 24K) 236MB 265MB 301MB 36MB (12%) + geoflow-small (1.1M) 236MB 263MB 487MB 224MB (46%) + outCSne8 ( 48K) 235MB 281MB 317MB 35MB (11%) + oQU480 (4.6M) 236MB 270MB 306MB 36MB (12%) + +``import uxarray`` alone is ~226 MB and constant to within 1 MB. oQU480 is 190x +larger than quad-hexagon yet both attributed ~36 MB to ``Grid.bounds`` -- the +benchmark was nearly insensitive to its own workload. The part that did vary was +numba: with a cold JIT cache the quad-hexagon case peaked at 599 MB, with a warm +one 298 MB. Which side of an ``asv continuous`` comparison paid the compile cost +depended on run order, so unrelated PRs kept reporting identical ~0.55-0.77 +"improvements". + +Why not sampled peak RSS +------------------------ +Sampling current RSS around the call and reporting the delta is immune to +process history, and it validates cleanly (a known 200 MB allocation measures as +200.0 MB). But after a warm-up call -- which is required to keep JIT compilation +out of the measurement -- the allocator reuses the pages it just freed, so RSS +*growth* systematically undercounts. Every operation the old benchmarks covered +measured 0.0-0.1 MB that way, including on the 98 MB / 28,571-face oQU120 mesh. + +Why not tracemalloc +------------------- +It measures allocation volume rather than RSS growth, which would sidestep the +page-reuse problem, but it does not observe numba NRT allocations -- and that is +where uxarray's array memory is allocated. + +What is left +------------ +Deterministic size accounting. It is bit-reproducible across runs, platforms and +JIT states (verified), scales with the mesh, and catches the memory regressions +that actually matter in an array library: dtype widening, densified +connectivity, and newly cached arrays. It does not capture transient peaks -- +but no available instrument captures those reliably here, and the measurements +above show the transients are ~1 MB, far below anything worth gating on. +""" + +__all__ = ["grid_nbytes", "dataset_nbytes"] + + +def grid_nbytes(uxgrid): + """Total size of the arrays a ``Grid`` currently holds, in bytes. + + Counts whatever has been materialized so far, so calling this after an + operation that caches results onto the grid (``bounds``, ``face_areas``, + connectivity) reports that operation's contribution to the grid's footprint. + """ + 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) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index b249e7b99..6f2faf2eb 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -2,6 +2,7 @@ from pathlib import Path import uxarray as ux +from ._memsize import grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0] @@ -25,6 +26,16 @@ 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" diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 659c19014..39e413fca 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -5,6 +5,7 @@ import numpy as np import uxarray as ux +from ._memsize import grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))) @@ -60,16 +61,25 @@ class Gradient(DatasetBenchmark): 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" 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. ``integrate`` returns a scalar, so + the memory that matters is what it caches onto the ``Grid`` (face + areas) to get there.""" + self.uxds[data_var].integrate() + return grid_nbytes(self.uxds.uxgrid) + + track_nbytes_integrate.unit = "bytes" class GeoDataFrame(DatasetBenchmark): diff --git a/benchmarks/quad_hexagon.py b/benchmarks/quad_hexagon.py index 4364f39b0..683e13b1b 100644 --- a/benchmarks/quad_hexagon.py +++ b/benchmarks/quad_hexagon.py @@ -2,6 +2,7 @@ from pathlib import Path import uxarray as ux +from ._memsize import dataset_nbytes, grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0] @@ -14,23 +15,18 @@ 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 peakmem_open_grid(self): - """Peak memory usage of a `Grid`""" - uxgrid = ux.open_grid(grid_path) + def track_nbytes_open_grid(self): + """Memory occupied by a `Grid`""" + return grid_nbytes(ux.open_grid(grid_path)) + track_nbytes_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)) - def peakmem_open_dataset(self): - """Peak memory usage of a `UxDataset`""" - uxds = ux.open_dataset(grid_path, data_path) + track_nbytes_open_dataset.unit = "bytes" From 2113022a5e742e47b8f2f7beeb75f0fd3f7ea84c Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 22 Jul 2026 13:14:45 -0500 Subject: [PATCH 2/8] Add process-scope peakmem benchmarks The previous commit removed peakmem_* because the per-operation memory it claimed to measure was ~1 MB against a reported 400-700 MB. But the larger question those benchmarks were reaching for -- "how much memory does it take to open this grid and compute on it, from a cold start" -- is real, and asv can answer it with its own peakmem_* once two things are controlled: 1. Process history. ru_maxrss is a high-water mark that never falls, so imports, setup() and earlier work in the same process set a floor under the result. asv already spawns a separate process per benchmark *and per parameter* (runner._run_benchmark_single_param), so this is handled as long as the benchmark classes declare no setup() -- asv counts setup memory towards peakmem_*, which is exactly the confound its own docs warn about. 2. numba JIT cache warmth. Compiling uxarray's 81 @njit(cache=True) kernels costs a few hundred MB of transient RSS, so an otherwise identical run peaked at 599 MB cold and 298 MB warm. Whichever side of an `asv continuous` comparison happened to compile paid that cost. setup_cache handles (2). asv runs it in its own process (runner.Spawner.create_setup_cache), so LLVM's footprint never enters the high-water mark of the processes that do the measuring; they load the compiled kernels from numba's on-disk cache instead. It runs once per commit, so both sides of a comparison are warmed symmetrically. Every parameter is warmed, not just one -- the grids do not all reach the same njit signatures. Measured via `asv run --python=same --quick -b peakmem`, twice warm and once after deleting every .nbi/.nbc in the installed uxarray: warm warm COLD spread import uxarray 280M 282M 280M 0.7% open+bounds quad-hexagon 349M 349M 346M 0.9% open+bounds geoflow-small 348M 348M 348M 0.0% open+bounds outCSne8 365M 364M 370M 1.6% open+bounds oQU480 354M 356M 351M 1.4% gradient 480km 358M 355M 354M 1.1% gradient 120km 371M 370M 381M 2.9% The cold column is the condition that used to halve the result; the cache repopulated from 0 to 33 entries during that run, confirming setup_cache did the compiling. Everything sits inside 2.9%, against asv's default 10% factor. peakmem_import_uxarray is included deliberately. ~280 MB of every row above is just importing uxarray, so tracking it on its own means a heavy new top-level import shows up as itself instead of silently inflating everything else. Also moves _memsize into benchmarks/helpers/ and gives it an __init__.py, so the package structure is explicit rather than relying on namespace packages. Co-Authored-By: Claude Opus 4.8 --- benchmarks/_memsize.py | 66 ---------------------------------- benchmarks/face_bounds.py | 30 +++++++++++++++- benchmarks/helpers/__init__.py | 0 benchmarks/helpers/_memsize.py | 12 +++++++ benchmarks/import.py | 9 +++++ benchmarks/mpas_ocean.py | 27 +++++++++++++- benchmarks/quad_hexagon.py | 2 +- 7 files changed, 77 insertions(+), 69 deletions(-) delete mode 100644 benchmarks/_memsize.py create mode 100644 benchmarks/helpers/__init__.py create mode 100644 benchmarks/helpers/_memsize.py diff --git a/benchmarks/_memsize.py b/benchmarks/_memsize.py deleted file mode 100644 index bc00a3b0e..000000000 --- a/benchmarks/_memsize.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Deterministic memory-size accounting for the benchmark suite. - -This module backs the ``track_nbytes_*`` benchmarks that replaced the former -``peakmem_*`` ones. - -Why ``peakmem_*`` was removed ------------------------------ -asv's ``peakmem_*`` records the maximum resident set size of the *whole -process*, and per asv's own docs it "also counts memory usage during the setup -routine". For uxarray that number is almost entirely fixed overhead: - - grid import +open_grid +.bounds attributable to op - quad-hexagon ( 24K) 236MB 265MB 301MB 36MB (12%) - geoflow-small (1.1M) 236MB 263MB 487MB 224MB (46%) - outCSne8 ( 48K) 235MB 281MB 317MB 35MB (11%) - oQU480 (4.6M) 236MB 270MB 306MB 36MB (12%) - -``import uxarray`` alone is ~226 MB and constant to within 1 MB. oQU480 is 190x -larger than quad-hexagon yet both attributed ~36 MB to ``Grid.bounds`` -- the -benchmark was nearly insensitive to its own workload. The part that did vary was -numba: with a cold JIT cache the quad-hexagon case peaked at 599 MB, with a warm -one 298 MB. Which side of an ``asv continuous`` comparison paid the compile cost -depended on run order, so unrelated PRs kept reporting identical ~0.55-0.77 -"improvements". - -Why not sampled peak RSS ------------------------- -Sampling current RSS around the call and reporting the delta is immune to -process history, and it validates cleanly (a known 200 MB allocation measures as -200.0 MB). But after a warm-up call -- which is required to keep JIT compilation -out of the measurement -- the allocator reuses the pages it just freed, so RSS -*growth* systematically undercounts. Every operation the old benchmarks covered -measured 0.0-0.1 MB that way, including on the 98 MB / 28,571-face oQU120 mesh. - -Why not tracemalloc -------------------- -It measures allocation volume rather than RSS growth, which would sidestep the -page-reuse problem, but it does not observe numba NRT allocations -- and that is -where uxarray's array memory is allocated. - -What is left ------------- -Deterministic size accounting. It is bit-reproducible across runs, platforms and -JIT states (verified), scales with the mesh, and catches the memory regressions -that actually matter in an array library: dtype widening, densified -connectivity, and newly cached arrays. It does not capture transient peaks -- -but no available instrument captures those reliably here, and the measurements -above show the transients are ~1 MB, far below anything worth gating on. -""" - -__all__ = ["grid_nbytes", "dataset_nbytes"] - - -def grid_nbytes(uxgrid): - """Total size of the arrays a ``Grid`` currently holds, in bytes. - - Counts whatever has been materialized so far, so calling this after an - operation that caches results onto the grid (``bounds``, ``face_areas``, - connectivity) reports that operation's contribution to the grid's footprint. - """ - 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) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index 6f2faf2eb..9009755e8 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -2,7 +2,7 @@ from pathlib import Path import uxarray as ux -from ._memsize import grid_nbytes +from .helpers._memsize import grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0] @@ -39,3 +39,31 @@ def track_nbytes_grid_with_bounds(self, grid_path): return grid_nbytes(self.uxgrid) track_nbytes_grid_with_bounds.unit = "bytes" + + +class FaceBoundsPeakMem: + """Peak memory of a cold start: import uxarray, open a grid, get its bounds. + + Deliberately defines no ``setup``. asv counts setup memory towards + ``peakmem_*``, and asv gives each parameter its own process, so with the JIT + warmed in ``setup_cache`` the measured process does exactly the work named in + the benchmark and nothing else. + """ + + params = FaceBounds.params + param_names = ["grid_path"] + + def setup_cache(self): + """Compile the njit kernels before anything is measured. + + asv runs this in its own process, so LLVM's few hundred MB stays out of + the measured ones -- they load from numba's on-disk cache instead. + Without it the first measured process compiles and the rest do not, which + is what made the old ``peakmem_*`` swing ~2x with run order. All params + are warmed; the grids do not reach the same njit signatures. + """ + for grid_path in self.params: + ux.open_grid(grid_path).bounds + + def peakmem_open_and_bounds(self, grid_path): + ux.open_grid(grid_path).bounds diff --git a/benchmarks/helpers/__init__.py b/benchmarks/helpers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/benchmarks/helpers/_memsize.py b/benchmarks/helpers/_memsize.py new file mode 100644 index 000000000..cf13f70f8 --- /dev/null +++ b/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) diff --git a/benchmarks/import.py b/benchmarks/import.py index e53515f2a..291672cff 100644 --- a/benchmarks/import.py +++ b/benchmarks/import.py @@ -3,3 +3,12 @@ class Imports: def timeraw_import_uxarray(self): return "import uxarray" + + def peakmem_import_uxarray(self): + """Peak memory of a process that has imported uxarray. + + This is the floor under every other ``peakmem_*`` result, so it is worth + tracking on its own: a heavy new top-level import shows up here rather + than silently inflating everything else. + """ + import uxarray # noqa: F401 diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 39e413fca..221df4d57 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -5,7 +5,7 @@ import numpy as np import uxarray as ux -from ._memsize import grid_nbytes +from .helpers._memsize import grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))) @@ -82,6 +82,31 @@ def track_nbytes_integrate(self, resolution): track_nbytes_integrate.unit = "bytes" +class GradientPeakMem: + """Peak memory of a cold start: import uxarray, open a dataset, take a gradient. + + Not a :class:`DatasetBenchmark` subclass -- that would open the dataset in + ``setup``, and asv counts setup memory towards ``peakmem_*``. + """ + + param_names = ["resolution"] + params = [["480km", "120km"]] + + def setup_cache(self): + """Compile the njit kernels before anything is measured. + + See :meth:`face_bounds.FaceBoundsPeakMem.setup_cache` -- asv runs this in + its own process, keeping LLVM's footprint out of the measured ones. + """ + for resolution in self.params[0]: + grid, data = file_path_dict[resolution] + ux.open_dataset(grid, data)[data_var].gradient() + + def peakmem_gradient(self, resolution): + grid, data = file_path_dict[resolution] + ux.open_dataset(grid, data)[data_var].gradient() + + class GeoDataFrame(DatasetBenchmark): param_names = DatasetBenchmark.param_names + ['exclude_antimeridian'] params = DatasetBenchmark.params + [[True, False]] diff --git a/benchmarks/quad_hexagon.py b/benchmarks/quad_hexagon.py index 683e13b1b..1e9dd50ad 100644 --- a/benchmarks/quad_hexagon.py +++ b/benchmarks/quad_hexagon.py @@ -2,7 +2,7 @@ from pathlib import Path import uxarray as ux -from ._memsize import dataset_nbytes, grid_nbytes +from .helpers._memsize import dataset_nbytes, grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0] From 7a17f020c25588f10cbe4d917a2106189317a35e Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 22 Jul 2026 13:58:03 -0500 Subject: [PATCH 3/8] Clean up Claude comments --- benchmarks/face_bounds.py | 17 ++--------------- benchmarks/import.py | 7 +------ benchmarks/mpas_ocean.py | 10 ++-------- 3 files changed, 5 insertions(+), 29 deletions(-) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index 9009755e8..d860bfdb1 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -42,26 +42,13 @@ def track_nbytes_grid_with_bounds(self, grid_path): class FaceBoundsPeakMem: - """Peak memory of a cold start: import uxarray, open a grid, get its bounds. - - Deliberately defines no ``setup``. asv counts setup memory towards - ``peakmem_*``, and asv gives each parameter its own process, so with the JIT - warmed in ``setup_cache`` the measured process does exactly the work named in - the benchmark and nothing else. - """ + """Peak memory of a cold start: import uxarray, open a grid, get its bounds.""" params = FaceBounds.params param_names = ["grid_path"] def setup_cache(self): - """Compile the njit kernels before anything is measured. - - asv runs this in its own process, so LLVM's few hundred MB stays out of - the measured ones -- they load from numba's on-disk cache instead. - Without it the first measured process compiles and the rest do not, which - is what made the old ``peakmem_*`` swing ~2x with run order. All params - are warmed; the grids do not reach the same njit signatures. - """ + """Compile the njit kernels before anything is measured.""" for grid_path in self.params: ux.open_grid(grid_path).bounds diff --git a/benchmarks/import.py b/benchmarks/import.py index 291672cff..98296ba6a 100644 --- a/benchmarks/import.py +++ b/benchmarks/import.py @@ -5,10 +5,5 @@ def timeraw_import_uxarray(self): return "import uxarray" def peakmem_import_uxarray(self): - """Peak memory of a process that has imported uxarray. - - This is the floor under every other ``peakmem_*`` result, so it is worth - tracking on its own: a heavy new top-level import shows up here rather - than silently inflating everything else. - """ + """Peak memory of a process that has imported uxarray.""" import uxarray # noqa: F401 diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 221df4d57..d9befabe7 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -73,9 +73,7 @@ def time_integrate(self, resolution): self.uxds[data_var].integrate() def track_nbytes_integrate(self, resolution): - """Grid footprint after integrating. ``integrate`` returns a scalar, so - the memory that matters is what it caches onto the ``Grid`` (face - areas) to get there.""" + """Grid footprint after integrating.""" self.uxds[data_var].integrate() return grid_nbytes(self.uxds.uxgrid) @@ -93,11 +91,7 @@ class GradientPeakMem: params = [["480km", "120km"]] def setup_cache(self): - """Compile the njit kernels before anything is measured. - - See :meth:`face_bounds.FaceBoundsPeakMem.setup_cache` -- asv runs this in - its own process, keeping LLVM's footprint out of the measured ones. - """ + """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() From 4f9e0d1ed426b4cca197ddc799aa59b1a2ec247f Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 10 Aug 2026 16:17:27 -0500 Subject: [PATCH 4/8] update peakmem face_areas --- benchmarks/mpas_ocean.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index cfd7008ac..df2b4cf06 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -58,11 +58,11 @@ def teardown(self, resolution, *args, **kwargs): class FaceAreas(GridBenchmark): - def time_compute_face_areas(self, resolution): - self.uxgrid.compute_face_areas() + def time_face_areas(self, resolution): + _ = self.uxgrid.face_areas - def peakmem_compute_face_areas(self, resolution): - self.uxgrid.compute_face_areas() + def peakmem_face_areas(self, resolution): + _ = self.uxgrid.face_areas class Gradient(DatasetBenchmark): @@ -119,7 +119,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) From 1c3d176c434baf1bcdefdb1a060e8f2cc71dd719 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 10 Aug 2026 19:33:16 -0500 Subject: [PATCH 5/8] Fix face_areas benchmark --- benchmarks/mpas_ocean.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index df2b4cf06..9d70e86c7 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -5,6 +5,7 @@ import numpy as np import uxarray as ux + from .helpers._memsize import grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))) @@ -58,11 +59,20 @@ def teardown(self, resolution, *args, **kwargs): class FaceAreas(GridBenchmark): + number = 1 # face_areas only calculates once before being cached + + def setup(self, resolution, *args, **kwargs): + super().setup(resolution, *args, **kwargs) + del self.uxgrid._ds["face_areas"] # guarantee it is empty + def time_face_areas(self, resolution): _ = self.uxgrid.face_areas - def peakmem_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" class Gradient(DatasetBenchmark): From b96aade5489555426185f06cb3a2e17103c4375d Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Thu, 13 Aug 2026 19:01:35 -0500 Subject: [PATCH 6/8] Tracemalloc 'peakmem' benchmarks for faithful metrics --- benchmarks/face_bounds.py | 28 +++++++++++- benchmarks/helpers/_peakmem.py | 80 ++++++++++++++++++++++++++++++++++ benchmarks/import.py | 11 +++-- benchmarks/mpas_ocean.py | 36 +++++++++++++-- benchmarks/quad_hexagon.py | 19 ++++++++ 5 files changed, 166 insertions(+), 8 deletions(-) create mode 100644 benchmarks/helpers/_peakmem.py diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index d860bfdb1..5b79c9ca7 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -3,6 +3,7 @@ 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] @@ -15,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): @@ -40,9 +48,25 @@ def track_nbytes_grid_with_bounds(self, grid_path): 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 FaceBoundsPeakMem: - """Peak memory of a cold start: import uxarray, open a grid, get its bounds.""" + """Peak memory of a cold start: import uxarray, open a grid, get its bounds. + + Whole-process ``ru_maxrss``, so the ~250MB uxarray import is part of the + number by design -- this is the cold-start cost, not the cost of ``bounds``. + For that, see ``FaceBounds.track_peakmem_face_bounds``. + """ params = FaceBounds.params param_names = ["grid_path"] @@ -52,5 +76,7 @@ def setup_cache(self): 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 diff --git a/benchmarks/helpers/_peakmem.py b/benchmarks/helpers/_peakmem.py new file mode 100644 index 000000000..0265d80f8 --- /dev/null +++ b/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 diff --git a/benchmarks/import.py b/benchmarks/import.py index 98296ba6a..32c87a27d 100644 --- a/benchmarks/import.py +++ b/benchmarks/import.py @@ -1,9 +1,14 @@ +from .helpers._peakmem import subprocess_peak_rss + + class Imports: """Benchmark importing uxarray.""" def timeraw_import_uxarray(self): return "import uxarray" - def peakmem_import_uxarray(self): - """Peak memory of a process that has imported uxarray.""" - import uxarray # noqa: F401 + 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" diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 9d70e86c7..d8324c0b3 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -7,6 +7,7 @@ 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__))) @@ -59,11 +60,15 @@ def teardown(self, resolution, *args, **kwargs): class FaceAreas(GridBenchmark): - number = 1 # face_areas only calculates once before being cached + 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) - del self.uxgrid._ds["face_areas"] # guarantee it is empty + self.uxgrid._ds = self.uxgrid._ds.drop_vars("face_areas", errors="ignore") def time_face_areas(self, resolution): _ = self.uxgrid.face_areas @@ -74,8 +79,21 @@ def track_nbytes_face_areas(self, resolution): track_nbytes_face_areas.unit = "bytes" + 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() @@ -85,8 +103,15 @@ def track_nbytes_gradient(self, resolution): 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() @@ -101,8 +126,9 @@ def track_nbytes_integrate(self, resolution): class GradientPeakMem: """Peak memory of a cold start: import uxarray, open a dataset, take a gradient. - Not a :class:`DatasetBenchmark` subclass -- that would open the dataset in - ``setup``, and asv counts setup memory towards ``peakmem_*``. + Whole-process ``ru_maxrss``, so the ~250MB uxarray import is part of the + number by design. For the cost of the gradient alone, see + ``Gradient.track_peakmem_gradient``. """ param_names = ["resolution"] @@ -114,6 +140,8 @@ def setup_cache(self): 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() diff --git a/benchmarks/quad_hexagon.py b/benchmarks/quad_hexagon.py index 1e9dd50ad..5b03eec6f 100644 --- a/benchmarks/quad_hexagon.py +++ b/benchmarks/quad_hexagon.py @@ -3,6 +3,7 @@ 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] @@ -11,6 +12,12 @@ 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) @@ -21,6 +28,12 @@ def track_nbytes_open_grid(self): 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) @@ -30,3 +43,9 @@ def track_nbytes_open_dataset(self): 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)) + + track_peakmem_open_dataset.unit = "bytes" From 8bfd409ce5ba95cda4bb2b52f32cac88037ebce9 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 14 Aug 2026 14:11:00 -0500 Subject: [PATCH 7/8] rename 'peakmem' benchmark classes that use tracemalloc --- benchmarks/face_bounds.py | 2 +- benchmarks/mpas_ocean.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index 5b79c9ca7..5e7550e3f 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -60,7 +60,7 @@ def track_peakmem_face_bounds(self, grid_path): track_peakmem_face_bounds.unit = "bytes" -class FaceBoundsPeakMem: +class FaceBoundsTracemalloc: """Peak memory of a cold start: import uxarray, open a grid, get its bounds. Whole-process ``ru_maxrss``, so the ~250MB uxarray import is part of the diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index d8324c0b3..8b16425c0 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -123,7 +123,7 @@ def track_nbytes_integrate(self, resolution): track_nbytes_integrate.unit = "bytes" -class GradientPeakMem: +class GradientTracemalloc: """Peak memory of a cold start: import uxarray, open a dataset, take a gradient. Whole-process ``ru_maxrss``, so the ~250MB uxarray import is part of the From 88d25deb4498a4f2069acc40a14bd2c1e26beae1 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 14 Aug 2026 15:39:46 -0500 Subject: [PATCH 8/8] Name RSS peakmem classes for their instrument, not tracemalloc 8bfd409c renamed the two whole-process classes to *Tracemalloc, but they hold asv-native peakmem_* benchmarks, which measure ru_maxrss for the entire process. The tracemalloc benchmarks are the track_peakmem_* in FaceBounds and Gradient, so the labels pointed at the wrong mechanism. Renames them to *ColdStartRss, matching ConnectivityChainRss on cmd/connectivity_peakmem, and notes in each docstring that the series runs one to three orders of magnitude above its track_peakmem_* counterpart -- the two sit adjacent on the dashboard, both formatted as bytes, and the gap is the differing definitions rather than a regression. Co-Authored-By: Claude Opus 5 --- benchmarks/face_bounds.py | 9 +++++---- benchmarks/mpas_ocean.py | 9 +++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index 5e7550e3f..8f1e41416 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -60,12 +60,13 @@ def track_peakmem_face_bounds(self, grid_path): track_peakmem_face_bounds.unit = "bytes" -class FaceBoundsTracemalloc: +class FaceBoundsColdStartRss: """Peak memory of a cold start: import uxarray, open a grid, get its bounds. - Whole-process ``ru_maxrss``, so the ~250MB uxarray import is part of the - number by design -- this is the cold-start cost, not the cost of ``bounds``. - For that, see ``FaceBounds.track_peakmem_face_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 diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 8b16425c0..f8f69b0b0 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -123,12 +123,13 @@ def track_nbytes_integrate(self, resolution): track_nbytes_integrate.unit = "bytes" -class GradientTracemalloc: +class GradientColdStartRss: """Peak memory of a cold start: import uxarray, open a dataset, take a gradient. - Whole-process ``ru_maxrss``, so the ~250MB uxarray import is part of the - number by design. For the cost of the gradient alone, see - ``Gradient.track_peakmem_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"]