From 1de2c3ce39e5a7ce78e80f5dabd81f53e2646d7d Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 7 Aug 2026 15:57:18 -0500 Subject: [PATCH 1/5] Peakmem benchmarks; chunked not working yet --- benchmarks/bench_connectivity.py | 322 ++++++++++++++++++++++++++++--- 1 file changed, 299 insertions(+), 23 deletions(-) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 5e3f53d02..4eb70410d 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -1,7 +1,10 @@ import os +import tracemalloc import urllib.request from pathlib import Path +import numpy as np + import uxarray as ux current_path = Path(os.path.dirname(os.path.realpath(__file__))) @@ -63,6 +66,45 @@ def teardown(self, resolution, *args, **kwargs): "node_face_connectivity", ] +# What each variable needs in place before its own construction routine can run, +# read off the ``_populate_*`` functions in ``uxarray/grid/connectivity.py``. +# Only direct prerequisites are listed: accessing one builds its own in turn. +CONNECTIVITY_PREREQUISITES = { + "n_nodes_per_face": (), + "face_node_connectivity": (), + "edge_node_connectivity": ("n_nodes_per_face",), + # ``_populate_edge_node_connectivity`` writes ``face_edge_connectivity`` out + # alongside its own variable, so this one costs nothing once it has run. + "face_edge_connectivity": ("edge_node_connectivity",), + "node_edge_connectivity": ("edge_node_connectivity",), + "face_face_connectivity": ("edge_face_connectivity",), + "edge_face_connectivity": ("face_edge_connectivity",), + "node_face_connectivity": (), +} + +def _apply_chunking(uxgrid, chunk_size): + """Chunks every grid variable in place, when ``chunk_size`` asks for it.""" + if chunk_size is not None: + # Chunks in place and returns None, so it cannot be chained. + uxgrid.chunk(n_node=chunk_size, n_edge=chunk_size, n_face=chunk_size) + return uxgrid + + +def _build_prerequisites(uxgrid, connectivity, chunk_size): + """Puts ``connectivity``'s prerequisites in place, leaving the grid chunked. + + Numba hands its results back as NumPy, so building a prerequisite on a + chunked grid quietly un-chunks the very input the measured routine is about + to read -- without the re-chunk, ``chunk_size`` only ever reached the + variables sitting at the root of a chain, and the three routines fed + entirely by Numba output showed no response to it at all. Re-chunking here + is what makes ``chunk_size`` mean "this routine is fed chunked input". + """ + for prerequisite in CONNECTIVITY_PREREQUISITES[connectivity]: + getattr(uxgrid, prerequisite) + return _apply_chunking(uxgrid, chunk_size) + + _numba_warmed_up = False def _warmup(uxgrid): @@ -80,62 +122,296 @@ def _warmup(uxgrid): _numba_warmed_up = True -class Connectivity(GridBenchmark): - # Each connectivity variable is cached in ``Grid._ds`` once constructed, so a - # sample may only contain a single call; otherwise every call but the first - # would time a dictionary lookup. - number = 1 +_topology_cache = {} - def setup(self, resolution, *args, **kwargs): - # The benchmark grids are MPAS meshes, which carry every connectivity - # variable on disk. Reading one would time the MPAS parser rather than - # the construction routines, so reduce the grid down to the minimal - # UGRID topology and let each variable be built on demand. + +def _source_topology(resolution): + """The minimal UGRID topology for ``resolution``, read once per process. + + The benchmark grids are MPAS meshes, which carry every connectivity variable + on disk. Reading one would measure the MPAS parser rather than the + construction routines, so the grid is reduced to the minimal UGRID topology + and each variable is left to be built on demand. + + Cached because asv re-runs ``setup`` between timing repeats: at the dyamond + resolutions re-reading the source grid costs orders of magnitude more than + the sample it precedes. + """ + if resolution not in _topology_cache: source_grid = ux.open_grid(file_path_dict[resolution]) - self.topology = ( + _topology_cache[resolution] = ( source_grid.node_lon.data, source_grid.node_lat.data, source_grid.face_node_connectivity.data, ) + return _topology_cache[resolution] + + +class MinimalGridBenchmark(GridBenchmark): + """Template for benchmarks that construct connectivity variables on demand. + + Holds a ``Grid`` carrying nothing but the minimal UGRID topology, plus the + topology needed to mint further ones, and leaves the Numba kernels compiled. + """ + + # ``None`` leaves the topology as NumPy; "auto" is what ``Grid.chunk`` + # itself defaults to, and unlike a fixed blocksize it stays sensible across + # the whole resolution range rather than degenerating into one chunk at + # 480km and tens of thousands at 3.75km. Add explicit sizes here to force + # multi-chunk behaviour at the smaller resolutions. + param_names = GridBenchmark.param_names + ["chunk_size"] + params = GridBenchmark.params + [[None, 8]] + + # Handover slot for ``_prerequisite_setup``; see its docstring. + active_grid = None + + # The default ``benchmark_timeout`` is not enough for a 3.75km grid once + # ``_warmup`` has to build every variable on it. + timeout = 1200 + + def setup(self, resolution, chunk_size=None, *args, **kwargs): + self.topology = _source_topology(resolution) - _warmup(self.minimal_grid()) - self.uxgrid = self.minimal_grid() + _warmup(self.minimal_grid(chunk_size)) + self.uxgrid = self.minimal_grid(chunk_size) + MinimalGridBenchmark.active_grid = self.uxgrid - def minimal_grid(self): - return ux.Grid.from_topology(*self.topology) + def minimal_grid(self, chunk_size=None): + """Mints a ``Grid`` holding nothing beyond the minimal UGRID topology. + + ``chunk_size`` is a dask blocksize rather than a number of chunks: that + is what ``Grid.chunk`` takes, and a count could not be converted into + one for ``n_edge`` anyway, whose length is unknown until + ``edge_node_connectivity`` has been built. ``None`` leaves the arrays as + NumPy. + """ + return _apply_chunking(ux.Grid.from_topology(*self.topology), chunk_size) def teardown(self, resolution, *args, **kwargs): + # Cleared so a per-benchmark setup can never reach a stale grid: it + # would quietly measure the wrong thing, where this raises instead. + MinimalGridBenchmark.active_grid = None del self.uxgrid del self.topology - def time_n_nodes_per_face(self, resolution): + +def _prerequisite_setup(connectivity): + """Builds a per-benchmark ``setup`` that puts ``connectivity``'s + prerequisites in place before the clock starts. + + asv collects ``setup`` from the benchmark function as well as from the + class, and runs the class one first, but it calls neither with the instance + -- only with the parameters. Hence the handover through + ``MinimalGridBenchmark.active_grid``, which the class ``setup`` has just + filled in. One benchmark runs per process, so there is nothing to collide + with. + """ + + def setup(resolution, chunk_size=None, *args, **kwargs): + _build_prerequisites( + MinimalGridBenchmark.active_grid, connectivity, chunk_size + ) + + return setup + + +class Connectivity(MinimalGridBenchmark): + """Time to construct each connectivity variable. + + Each variable's prerequisites are built during ``setup``, so a sample times + the one construction routine that produces that variable rather than the + whole chain rooted at it -- matching how + :class:`ConnectivityPeakAlloc` attributes memory. + """ + + # Each connectivity variable is cached in ``Grid._ds`` once constructed, so a + # sample may only contain a single call; otherwise every call but the first + # would time a dictionary lookup. + number = 1 + + def time_n_nodes_per_face(self, resolution, chunk_size): _ = self.uxgrid.n_nodes_per_face.compute() - def time_face_node(self, resolution): + time_n_nodes_per_face.setup = _prerequisite_setup("n_nodes_per_face") + + def time_face_node(self, resolution, chunk_size): _ = self.uxgrid.face_node_connectivity.compute() - def time_edge_node(self, resolution): + time_face_node.setup = _prerequisite_setup("face_node_connectivity") + + def time_edge_node(self, resolution, chunk_size): _ = self.uxgrid.edge_node_connectivity.compute() + time_edge_node.setup = _prerequisite_setup("edge_node_connectivity") + # TODO: Not yet supported? # def time_node_node(self, resolution): # _ = self.uxgrid.node_node_connectivity - def time_face_edge(self, resolution): + def time_face_edge(self, resolution, chunk_size): _ = self.uxgrid.face_edge_connectivity.compute() + time_face_edge.setup = _prerequisite_setup("face_edge_connectivity") + # TODO: Not yet supported? # def time_edge_edge(self, resolution): # _ = self.uxgrid.edge_edge_connectivity - def time_node_edge(self, resolution): + def time_node_edge(self, resolution, chunk_size): + _ = self.uxgrid.node_edge_connectivity.compute() + + time_node_edge.setup = _prerequisite_setup("node_edge_connectivity") + + def time_face_face(self, resolution, chunk_size): + _ = self.uxgrid.face_face_connectivity.compute() + + time_face_face.setup = _prerequisite_setup("face_face_connectivity") + + def time_edge_face(self, resolution, chunk_size): + _ = self.uxgrid.edge_face_connectivity.compute() + + time_edge_face.setup = _prerequisite_setup("edge_face_connectivity") + + def time_node_face(self, resolution, chunk_size): + _ = self.uxgrid.node_face_connectivity.compute() + + time_node_face.setup = _prerequisite_setup("node_face_connectivity") + + +def _peak_allocated(build): + """Bytes held at the high-water point of ``build``, counting only what it + allocated itself.""" + tracemalloc.start() + try: + tracemalloc.reset_peak() + build() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + +class ConnectivityPeakAlloc(MinimalGridBenchmark): + """Peak memory of each connectivity routine on its own. + + Reports the transient high-water allocation of the construction routine, + with the ~245MB the process is already holding subtracted out. + """ + + # Applies to every ``track_*`` in the class; asv resolves benchmark + # attributes from the instance when the function does not carry them. + unit = "bytes" + + def _peak_building(self, name, chunk_size): + """Peak allocation of ``name``'s own construction routine.""" + uxgrid = _build_prerequisites(self.minimal_grid(chunk_size), name, chunk_size) + return _peak_allocated(lambda: getattr(uxgrid, name).compute()) + + def track_peakmem_n_nodes_per_face(self, resolution, chunk_size): + return self._peak_building("n_nodes_per_face", chunk_size) + + def track_peakmem_face_node(self, resolution, chunk_size): + return self._peak_building("face_node_connectivity", chunk_size) + + def track_peakmem_edge_node(self, resolution, chunk_size): + return self._peak_building("edge_node_connectivity", chunk_size) + + def track_peakmem_face_edge(self, resolution, chunk_size): + return self._peak_building("face_edge_connectivity", chunk_size) + + def track_peakmem_node_edge(self, resolution, chunk_size): + return self._peak_building("node_edge_connectivity", chunk_size) + + def track_peakmem_face_face(self, resolution, chunk_size): + return self._peak_building("face_face_connectivity", chunk_size) + + def track_peakmem_edge_face(self, resolution, chunk_size): + return self._peak_building("edge_face_connectivity", chunk_size) + + def track_peakmem_node_face(self, resolution, chunk_size): + return self._peak_building("node_face_connectivity", chunk_size) + + +def _save_topology(uxgrid, npz_path): + """Writes the minimal UGRID topology of ``uxgrid`` out to ``npz_path``.""" + np.savez( + npz_path, + node_lon=uxgrid.node_lon.data, + node_lat=uxgrid.node_lat.data, + face_node_connectivity=uxgrid.face_node_connectivity.data, + ) + + +def _load_topology(npz_path, chunk_size=None): + """Builds a ``Grid`` holding nothing beyond the minimal UGRID topology.""" + with np.load(npz_path) as topology: + uxgrid = ux.Grid.from_topology( + topology["node_lon"], + topology["node_lat"], + topology["face_node_connectivity"], + ) + return _apply_chunking(uxgrid, chunk_size) + + +class ConnectivityPeakMem: + """Peak resident memory of the process while constructing each connectivity + variable.""" + + # Declared rather than inherited from ``MinimalGridBenchmark`` -- only the + # parameterization is shared, not the ``setup`` that opens a grid in the + # process being measured. + param_names = MinimalGridBenchmark.param_names + params = MinimalGridBenchmark.params + timeout = 1200 + + def setup_cache(self): + # asv runs this in its own process and passes the return value back as + # the leading argument of ``setup`` and of each benchmark, so nothing + # allocated here counts towards the samples. + topology_paths = {} + for resolution in self.params[0]: + npz_path = os.path.abspath(f"topology_{resolution}.npz") + _save_topology(ux.open_grid(file_path_dict[resolution]), npz_path) + topology_paths[resolution] = npz_path + + # Resolution only affects how long the kernels run, not which + # signatures get compiled, so warming up on the coarsest grid is enough. + _warmup(_load_topology(topology_paths[self.params[0][0]])) + + return topology_paths + + # Reading every grid in ``file_path_dict`` exceeds the default + # ``benchmark_timeout`` once the Glade paths are available. + setup_cache.timeout = 1800 + + def setup(self, topology_paths, resolution, chunk_size): + # Each connectivity variable is cached in ``Grid._ds`` once constructed, + # so the measured call needs a ``Grid`` that does not hold it yet. + self.uxgrid = _load_topology(topology_paths[resolution], chunk_size) + + def teardown(self, topology_paths, resolution, chunk_size): + del self.uxgrid + + def peakmem_n_nodes_per_face(self, topology_paths, resolution, chunk_size): + _ = self.uxgrid.n_nodes_per_face.compute() + + def peakmem_face_node(self, topology_paths, resolution, chunk_size): + _ = self.uxgrid.face_node_connectivity.compute() + + def peakmem_edge_node(self, topology_paths, resolution, chunk_size): + _ = self.uxgrid.edge_node_connectivity.compute() + + def peakmem_face_edge(self, topology_paths, resolution, chunk_size): + _ = self.uxgrid.face_edge_connectivity.compute() + + def peakmem_node_edge(self, topology_paths, resolution, chunk_size): _ = self.uxgrid.node_edge_connectivity.compute() - def time_face_face(self, resolution): + def peakmem_face_face(self, topology_paths, resolution, chunk_size): _ = self.uxgrid.face_face_connectivity.compute() - def time_edge_face(self, resolution): + def peakmem_edge_face(self, topology_paths, resolution, chunk_size): _ = self.uxgrid.edge_face_connectivity.compute() - def time_node_face(self, resolution): + def peakmem_node_face(self, topology_paths, resolution, chunk_size): _ = self.uxgrid.node_face_connectivity.compute() From e5878102df806aa8595d1f6255b0fdbc00f416ea Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 7 Aug 2026 19:22:42 -0500 Subject: [PATCH 2/5] Connectivity peakmem: remove chunking logic --- benchmarks/bench_connectivity.py | 162 +++++++++++++------------------ 1 file changed, 70 insertions(+), 92 deletions(-) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 4eb70410d..781193db2 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -82,32 +82,19 @@ def teardown(self, resolution, *args, **kwargs): "node_face_connectivity": (), } -def _apply_chunking(uxgrid, chunk_size): - """Chunks every grid variable in place, when ``chunk_size`` asks for it.""" - if chunk_size is not None: - # Chunks in place and returns None, so it cannot be chained. - uxgrid.chunk(n_node=chunk_size, n_edge=chunk_size, n_face=chunk_size) - return uxgrid - - -def _build_prerequisites(uxgrid, connectivity, chunk_size): - """Puts ``connectivity``'s prerequisites in place, leaving the grid chunked. - Numba hands its results back as NumPy, so building a prerequisite on a - chunked grid quietly un-chunks the very input the measured routine is about - to read -- without the re-chunk, ``chunk_size`` only ever reached the - variables sitting at the root of a chain, and the three routines fed - entirely by Numba output showed no response to it at all. Re-chunking here - is what makes ``chunk_size`` mean "this routine is fed chunked input". - """ +def _build_prerequisites(uxgrid, connectivity): + """Puts ``connectivity``'s prerequisites in place, so that what follows + measures the one construction routine rather than the whole chain rooted at + it.""" for prerequisite in CONNECTIVITY_PREREQUISITES[connectivity]: getattr(uxgrid, prerequisite) - return _apply_chunking(uxgrid, chunk_size) + return uxgrid _numba_warmed_up = False -def _warmup(uxgrid): +def _warmup(): """Compiles the Numba kernels backing each connectivity variable. ``_build_node_edge_connectivity`` is not disk-cached, so a fresh benchmark @@ -117,6 +104,9 @@ def _warmup(uxgrid): global _numba_warmed_up if _numba_warmed_up: return + # The resolution only affects how long the kernels run, not which signatures get + # compiled, so the coarsest grid is enough to warm any of them. + uxgrid = ux.Grid.from_topology(*_source_topology(GridBenchmark.params[0][0])) for name in CONNECTIVITY_NAMES: getattr(uxgrid, name) _numba_warmed_up = True @@ -154,38 +144,23 @@ class MinimalGridBenchmark(GridBenchmark): topology needed to mint further ones, and leaves the Numba kernels compiled. """ - # ``None`` leaves the topology as NumPy; "auto" is what ``Grid.chunk`` - # itself defaults to, and unlike a fixed blocksize it stays sensible across - # the whole resolution range rather than degenerating into one chunk at - # 480km and tens of thousands at 3.75km. Add explicit sizes here to force - # multi-chunk behaviour at the smaller resolutions. - param_names = GridBenchmark.param_names + ["chunk_size"] - params = GridBenchmark.params + [[None, 8]] - # Handover slot for ``_prerequisite_setup``; see its docstring. active_grid = None - # The default ``benchmark_timeout`` is not enough for a 3.75km grid once - # ``_warmup`` has to build every variable on it. + # The default ``benchmark_timeout`` is not enough to build a connectivity + # variable on a 3.75km grid. timeout = 1200 - def setup(self, resolution, chunk_size=None, *args, **kwargs): + def setup(self, resolution, *args, **kwargs): self.topology = _source_topology(resolution) - _warmup(self.minimal_grid(chunk_size)) - self.uxgrid = self.minimal_grid(chunk_size) + _warmup() + self.uxgrid = self.minimal_grid() MinimalGridBenchmark.active_grid = self.uxgrid - def minimal_grid(self, chunk_size=None): - """Mints a ``Grid`` holding nothing beyond the minimal UGRID topology. - - ``chunk_size`` is a dask blocksize rather than a number of chunks: that - is what ``Grid.chunk`` takes, and a count could not be converted into - one for ``n_edge`` anyway, whose length is unknown until - ``edge_node_connectivity`` has been built. ``None`` leaves the arrays as - NumPy. - """ - return _apply_chunking(ux.Grid.from_topology(*self.topology), chunk_size) + def minimal_grid(self): + """Mints a ``Grid`` holding nothing beyond the minimal UGRID topology.""" + return ux.Grid.from_topology(*self.topology) def teardown(self, resolution, *args, **kwargs): # Cleared so a per-benchmark setup can never reach a stale grid: it @@ -207,10 +182,8 @@ def _prerequisite_setup(connectivity): with. """ - def setup(resolution, chunk_size=None, *args, **kwargs): - _build_prerequisites( - MinimalGridBenchmark.active_grid, connectivity, chunk_size - ) + def setup(resolution, *args, **kwargs): + _build_prerequisites(MinimalGridBenchmark.active_grid, connectivity) return setup @@ -229,17 +202,17 @@ class Connectivity(MinimalGridBenchmark): # would time a dictionary lookup. number = 1 - def time_n_nodes_per_face(self, resolution, chunk_size): + def time_n_nodes_per_face(self, resolution): _ = self.uxgrid.n_nodes_per_face.compute() time_n_nodes_per_face.setup = _prerequisite_setup("n_nodes_per_face") - def time_face_node(self, resolution, chunk_size): + def time_face_node(self, resolution): _ = self.uxgrid.face_node_connectivity.compute() time_face_node.setup = _prerequisite_setup("face_node_connectivity") - def time_edge_node(self, resolution, chunk_size): + def time_edge_node(self, resolution): _ = self.uxgrid.edge_node_connectivity.compute() time_edge_node.setup = _prerequisite_setup("edge_node_connectivity") @@ -248,7 +221,7 @@ def time_edge_node(self, resolution, chunk_size): # def time_node_node(self, resolution): # _ = self.uxgrid.node_node_connectivity - def time_face_edge(self, resolution, chunk_size): + def time_face_edge(self, resolution): _ = self.uxgrid.face_edge_connectivity.compute() time_face_edge.setup = _prerequisite_setup("face_edge_connectivity") @@ -257,22 +230,22 @@ def time_face_edge(self, resolution, chunk_size): # def time_edge_edge(self, resolution): # _ = self.uxgrid.edge_edge_connectivity - def time_node_edge(self, resolution, chunk_size): + def time_node_edge(self, resolution): _ = self.uxgrid.node_edge_connectivity.compute() time_node_edge.setup = _prerequisite_setup("node_edge_connectivity") - def time_face_face(self, resolution, chunk_size): + def time_face_face(self, resolution): _ = self.uxgrid.face_face_connectivity.compute() time_face_face.setup = _prerequisite_setup("face_face_connectivity") - def time_edge_face(self, resolution, chunk_size): + def time_edge_face(self, resolution): _ = self.uxgrid.edge_face_connectivity.compute() time_edge_face.setup = _prerequisite_setup("edge_face_connectivity") - def time_node_face(self, resolution, chunk_size): + def time_node_face(self, resolution): _ = self.uxgrid.node_face_connectivity.compute() time_node_face.setup = _prerequisite_setup("node_face_connectivity") @@ -302,34 +275,34 @@ class ConnectivityPeakAlloc(MinimalGridBenchmark): # attributes from the instance when the function does not carry them. unit = "bytes" - def _peak_building(self, name, chunk_size): + def _peak_building(self, name): """Peak allocation of ``name``'s own construction routine.""" - uxgrid = _build_prerequisites(self.minimal_grid(chunk_size), name, chunk_size) + uxgrid = _build_prerequisites(self.minimal_grid(), name) return _peak_allocated(lambda: getattr(uxgrid, name).compute()) - def track_peakmem_n_nodes_per_face(self, resolution, chunk_size): - return self._peak_building("n_nodes_per_face", chunk_size) + def track_peakmem_n_nodes_per_face(self, resolution): + return self._peak_building("n_nodes_per_face") - def track_peakmem_face_node(self, resolution, chunk_size): - return self._peak_building("face_node_connectivity", chunk_size) + def track_peakmem_face_node(self, resolution): + return self._peak_building("face_node_connectivity") - def track_peakmem_edge_node(self, resolution, chunk_size): - return self._peak_building("edge_node_connectivity", chunk_size) + def track_peakmem_edge_node(self, resolution): + return self._peak_building("edge_node_connectivity") - def track_peakmem_face_edge(self, resolution, chunk_size): - return self._peak_building("face_edge_connectivity", chunk_size) + def track_peakmem_face_edge(self, resolution): + return self._peak_building("face_edge_connectivity") - def track_peakmem_node_edge(self, resolution, chunk_size): - return self._peak_building("node_edge_connectivity", chunk_size) + def track_peakmem_node_edge(self, resolution): + return self._peak_building("node_edge_connectivity") - def track_peakmem_face_face(self, resolution, chunk_size): - return self._peak_building("face_face_connectivity", chunk_size) + def track_peakmem_face_face(self, resolution): + return self._peak_building("face_face_connectivity") - def track_peakmem_edge_face(self, resolution, chunk_size): - return self._peak_building("edge_face_connectivity", chunk_size) + def track_peakmem_edge_face(self, resolution): + return self._peak_building("edge_face_connectivity") - def track_peakmem_node_face(self, resolution, chunk_size): - return self._peak_building("node_face_connectivity", chunk_size) + def track_peakmem_node_face(self, resolution): + return self._peak_building("node_face_connectivity") def _save_topology(uxgrid, npz_path): @@ -342,26 +315,30 @@ def _save_topology(uxgrid, npz_path): ) -def _load_topology(npz_path, chunk_size=None): +def _load_topology(npz_path): """Builds a ``Grid`` holding nothing beyond the minimal UGRID topology.""" with np.load(npz_path) as topology: - uxgrid = ux.Grid.from_topology( + return ux.Grid.from_topology( topology["node_lon"], topology["node_lat"], topology["face_node_connectivity"], ) - return _apply_chunking(uxgrid, chunk_size) class ConnectivityPeakMem: """Peak resident memory of the process while constructing each connectivity - variable.""" + variable. + + Unlike :class:`ConnectivityPeakAlloc`, no prerequisites are built during + ``setup``, so a sample covers the whole chain rooted at that variable -- + which is what the resident footprint of asking for it actually costs. + """ # Declared rather than inherited from ``MinimalGridBenchmark`` -- only the # parameterization is shared, not the ``setup`` that opens a grid in the # process being measured. - param_names = MinimalGridBenchmark.param_names - params = MinimalGridBenchmark.params + param_names = GridBenchmark.param_names + params = GridBenchmark.params timeout = 1200 def setup_cache(self): @@ -374,9 +351,10 @@ def setup_cache(self): _save_topology(ux.open_grid(file_path_dict[resolution]), npz_path) topology_paths[resolution] = npz_path - # Resolution only affects how long the kernels run, not which - # signatures get compiled, so warming up on the coarsest grid is enough. - _warmup(_load_topology(topology_paths[self.params[0][0]])) + # Being a separate process, this only reaches the benchmarks through + # Numba's on-disk cache -- which is the point, since a compilation it + # saves is one that would otherwise land inside a measured region. + _warmup() return topology_paths @@ -384,34 +362,34 @@ def setup_cache(self): # ``benchmark_timeout`` once the Glade paths are available. setup_cache.timeout = 1800 - def setup(self, topology_paths, resolution, chunk_size): + def setup(self, topology_paths, resolution): # Each connectivity variable is cached in ``Grid._ds`` once constructed, # so the measured call needs a ``Grid`` that does not hold it yet. - self.uxgrid = _load_topology(topology_paths[resolution], chunk_size) + self.uxgrid = _load_topology(topology_paths[resolution]) - def teardown(self, topology_paths, resolution, chunk_size): + def teardown(self, topology_paths, resolution): del self.uxgrid - def peakmem_n_nodes_per_face(self, topology_paths, resolution, chunk_size): + def peakmem_n_nodes_per_face(self, topology_paths, resolution): _ = self.uxgrid.n_nodes_per_face.compute() - def peakmem_face_node(self, topology_paths, resolution, chunk_size): + def peakmem_face_node(self, topology_paths, resolution): _ = self.uxgrid.face_node_connectivity.compute() - def peakmem_edge_node(self, topology_paths, resolution, chunk_size): + def peakmem_edge_node(self, topology_paths, resolution): _ = self.uxgrid.edge_node_connectivity.compute() - def peakmem_face_edge(self, topology_paths, resolution, chunk_size): + def peakmem_face_edge(self, topology_paths, resolution): _ = self.uxgrid.face_edge_connectivity.compute() - def peakmem_node_edge(self, topology_paths, resolution, chunk_size): + def peakmem_node_edge(self, topology_paths, resolution): _ = self.uxgrid.node_edge_connectivity.compute() - def peakmem_face_face(self, topology_paths, resolution, chunk_size): + def peakmem_face_face(self, topology_paths, resolution): _ = self.uxgrid.face_face_connectivity.compute() - def peakmem_edge_face(self, topology_paths, resolution, chunk_size): + def peakmem_edge_face(self, topology_paths, resolution): _ = self.uxgrid.edge_face_connectivity.compute() - def peakmem_node_face(self, topology_paths, resolution, chunk_size): + def peakmem_node_face(self, topology_paths, resolution): _ = self.uxgrid.node_face_connectivity.compute() From 4a2d24411b1557504c057f1c3040976c8cce4ef9 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 7 Aug 2026 19:53:06 -0500 Subject: [PATCH 3/5] missing node_edge caching --- uxarray/grid/connectivity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index ac9658979..0fa67b384 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -469,7 +469,7 @@ def _populate_node_edge_connectivity(grid): ) -@njit +@njit(cache=True) def _build_node_edge_connectivity(edge_nodes, n_node): """Constructs the Node Edge Connectivity, which stores the indices of the edges that are shared by each node.""" n_edge, nodes_per_edge = edge_nodes.shape From ad294a20198399528469b7b7458303a99e29f07c Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 14 Aug 2026 16:54:32 -0500 Subject: [PATCH 4/5] connectivity benchmarks comment cleanup --- benchmarks/bench_connectivity.py | 120 ++++++++++++------------------- 1 file changed, 44 insertions(+), 76 deletions(-) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 781193db2..be3f1571e 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -1,5 +1,4 @@ import os -import tracemalloc import urllib.request from pathlib import Path @@ -7,6 +6,8 @@ import uxarray as ux +from .helpers._peakmem import peak_allocated + current_path = Path(os.path.dirname(os.path.realpath(__file__))) grid_filename_480 = "oQU480.grid.nc" @@ -66,15 +67,14 @@ def teardown(self, resolution, *args, **kwargs): "node_face_connectivity", ] -# What each variable needs in place before its own construction routine can run, -# read off the ``_populate_*`` functions in ``uxarray/grid/connectivity.py``. -# Only direct prerequisites are listed: accessing one builds its own in turn. +# Direct prerequisites only, read off the ``_populate_*`` functions in +# ``uxarray/grid/connectivity.py``; accessing one builds its own in turn. CONNECTIVITY_PREREQUISITES = { "n_nodes_per_face": (), "face_node_connectivity": (), "edge_node_connectivity": ("n_nodes_per_face",), - # ``_populate_edge_node_connectivity`` writes ``face_edge_connectivity`` out - # alongside its own variable, so this one costs nothing once it has run. + # ``_populate_edge_node_connectivity`` writes this one out too, so it costs + # nothing once that has run. "face_edge_connectivity": ("edge_node_connectivity",), "node_edge_connectivity": ("edge_node_connectivity",), "face_face_connectivity": ("edge_face_connectivity",), @@ -84,9 +84,8 @@ def teardown(self, resolution, *args, **kwargs): def _build_prerequisites(uxgrid, connectivity): - """Puts ``connectivity``'s prerequisites in place, so that what follows - measures the one construction routine rather than the whole chain rooted at - it.""" + """Builds ``connectivity``'s prerequisites, so what follows measures one + construction routine rather than the whole chain rooted at it.""" for prerequisite in CONNECTIVITY_PREREQUISITES[connectivity]: getattr(uxgrid, prerequisite) return uxgrid @@ -97,15 +96,15 @@ def _build_prerequisites(uxgrid, connectivity): def _warmup(): """Compiles the Numba kernels backing each connectivity variable. - ``_build_node_edge_connectivity`` is not disk-cached, so a fresh benchmark - process would otherwise charge ~240ms of JIT compilation to whichever sample - happened to touch it first. + Every kernel in ``uxarray/grid/connectivity.py`` is ``@njit(cache=True)``, so + this carries across processes through Numba's on-disk cache -- what makes it + usable from ``setup_cache``. Loading from that cache still allocates, so it + matters for ``track_peakmem_*`` too, not just timing. """ global _numba_warmed_up if _numba_warmed_up: return - # The resolution only affects how long the kernels run, not which signatures get - # compiled, so the coarsest grid is enough to warm any of them. + # Resolution affects how long the kernels run, not which signatures compile. uxgrid = ux.Grid.from_topology(*_source_topology(GridBenchmark.params[0][0])) for name in CONNECTIVITY_NAMES: getattr(uxgrid, name) @@ -118,14 +117,12 @@ def _warmup(): def _source_topology(resolution): """The minimal UGRID topology for ``resolution``, read once per process. - The benchmark grids are MPAS meshes, which carry every connectivity variable - on disk. Reading one would measure the MPAS parser rather than the - construction routines, so the grid is reduced to the minimal UGRID topology - and each variable is left to be built on demand. + The benchmark grids are MPAS meshes carrying every connectivity variable on + disk; reading one would measure the MPAS parser rather than the construction + routines, so each variable is left to be built on demand. - Cached because asv re-runs ``setup`` between timing repeats: at the dyamond - resolutions re-reading the source grid costs orders of magnitude more than - the sample it precedes. + Cached because asv re-runs ``setup`` between repeats, and at dyamond + resolutions re-reading the source grid dwarfs the sample it precedes. """ if resolution not in _topology_cache: source_grid = ux.open_grid(file_path_dict[resolution]) @@ -147,8 +144,7 @@ class MinimalGridBenchmark(GridBenchmark): # Handover slot for ``_prerequisite_setup``; see its docstring. active_grid = None - # The default ``benchmark_timeout`` is not enough to build a connectivity - # variable on a 3.75km grid. + # asv's 60s default is not enough to build a connectivity variable at 3.75km. timeout = 1200 def setup(self, resolution, *args, **kwargs): @@ -163,8 +159,8 @@ def minimal_grid(self): return ux.Grid.from_topology(*self.topology) def teardown(self, resolution, *args, **kwargs): - # Cleared so a per-benchmark setup can never reach a stale grid: it - # would quietly measure the wrong thing, where this raises instead. + # Cleared so a per-benchmark setup raises rather than quietly measuring + # a stale grid. MinimalGridBenchmark.active_grid = None del self.uxgrid del self.topology @@ -174,12 +170,9 @@ def _prerequisite_setup(connectivity): """Builds a per-benchmark ``setup`` that puts ``connectivity``'s prerequisites in place before the clock starts. - asv collects ``setup`` from the benchmark function as well as from the - class, and runs the class one first, but it calls neither with the instance - -- only with the parameters. Hence the handover through - ``MinimalGridBenchmark.active_grid``, which the class ``setup`` has just - filled in. One benchmark runs per process, so there is nothing to collide - with. + asv collects ``setup`` from the benchmark function as well as the class and + runs the class one first, but passes neither the instance, hence the handover + through ``MinimalGridBenchmark.active_grid``. """ def setup(resolution, *args, **kwargs): @@ -191,16 +184,13 @@ def setup(resolution, *args, **kwargs): class Connectivity(MinimalGridBenchmark): """Time to construct each connectivity variable. - Each variable's prerequisites are built during ``setup``, so a sample times - the one construction routine that produces that variable rather than the - whole chain rooted at it -- matching how - :class:`ConnectivityPeakAlloc` attributes memory. + Prerequisites are built during ``setup``, so a sample times the one routine + that produces that variable rather than the whole chain rooted at it -- + matching how :class:`ConnectivityTracemalloc` attributes memory. """ - # Each connectivity variable is cached in ``Grid._ds`` once constructed, so a - # sample may only contain a single call; otherwise every call but the first - # would time a dictionary lookup. number = 1 + warmup_time = 0 def time_n_nodes_per_face(self, resolution): _ = self.uxgrid.n_nodes_per_face.compute() @@ -251,34 +241,19 @@ def time_node_face(self, resolution): time_node_face.setup = _prerequisite_setup("node_face_connectivity") -def _peak_allocated(build): - """Bytes held at the high-water point of ``build``, counting only what it - allocated itself.""" - tracemalloc.start() - try: - tracemalloc.reset_peak() - build() - _, peak = tracemalloc.get_traced_memory() - finally: - tracemalloc.stop() - return peak - - -class ConnectivityPeakAlloc(MinimalGridBenchmark): +class ConnectivityTracemalloc(MinimalGridBenchmark): """Peak memory of each connectivity routine on its own. - Reports the transient high-water allocation of the construction routine, - with the ~245MB the process is already holding subtracted out. + The transient high-water allocation of the construction routine, with the + ~245MB the process already holds excluded. """ - # Applies to every ``track_*`` in the class; asv resolves benchmark - # attributes from the instance when the function does not carry them. unit = "bytes" def _peak_building(self, name): """Peak allocation of ``name``'s own construction routine.""" uxgrid = _build_prerequisites(self.minimal_grid(), name) - return _peak_allocated(lambda: getattr(uxgrid, name).compute()) + return peak_allocated(lambda: getattr(uxgrid, name).compute()) def track_peakmem_n_nodes_per_face(self, resolution): return self._peak_building("n_nodes_per_face") @@ -306,8 +281,12 @@ def track_peakmem_node_face(self, resolution): def _save_topology(uxgrid, npz_path): - """Writes the minimal UGRID topology of ``uxgrid`` out to ``npz_path``.""" - np.savez( + """Writes the minimal UGRID topology of ``uxgrid`` out to ``npz_path``. + + Compressed because ``setup_cache`` writes into a ``tempfile.mkdtemp()`` and + ``face_node_connectivity`` alone is multi-GB at 3.75km. + """ + np.savez_compressed( npz_path, node_lon=uxgrid.node_lon.data, node_lat=uxgrid.node_lat.data, @@ -325,46 +304,35 @@ def _load_topology(npz_path): ) -class ConnectivityPeakMem: +class ConnectivityChainRss: """Peak resident memory of the process while constructing each connectivity variable. - Unlike :class:`ConnectivityPeakAlloc`, no prerequisites are built during - ``setup``, so a sample covers the whole chain rooted at that variable -- - which is what the resident footprint of asking for it actually costs. + Differs from :class:`ConnectivityTracemalloc` on two axes. Scope: no + prerequisites are built in ``setup``, so a sample covers + the whole chain rooted at that variable. Instrument: ``ru_maxrss`` for the + whole process, so the ~250MB import. """ - # Declared rather than inherited from ``MinimalGridBenchmark`` -- only the - # parameterization is shared, not the ``setup`` that opens a grid in the - # process being measured. + # Only the parameterization is shared with ``MinimalGridBenchmark`` param_names = GridBenchmark.param_names params = GridBenchmark.params timeout = 1200 def setup_cache(self): - # asv runs this in its own process and passes the return value back as - # the leading argument of ``setup`` and of each benchmark, so nothing - # allocated here counts towards the samples. topology_paths = {} for resolution in self.params[0]: npz_path = os.path.abspath(f"topology_{resolution}.npz") _save_topology(ux.open_grid(file_path_dict[resolution]), npz_path) topology_paths[resolution] = npz_path - # Being a separate process, this only reaches the benchmarks through - # Numba's on-disk cache -- which is the point, since a compilation it - # saves is one that would otherwise land inside a measured region. _warmup() return topology_paths - # Reading every grid in ``file_path_dict`` exceeds the default - # ``benchmark_timeout`` once the Glade paths are available. setup_cache.timeout = 1800 def setup(self, topology_paths, resolution): - # Each connectivity variable is cached in ``Grid._ds`` once constructed, - # so the measured call needs a ``Grid`` that does not hold it yet. self.uxgrid = _load_topology(topology_paths[resolution]) def teardown(self, topology_paths, resolution): From 8334943e1aaff99f31713ccfb566ad50bdcdd86a Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 18 Aug 2026 14:36:02 -0500 Subject: [PATCH 5/5] connectivity peakmem: whole chain benchmarks (may remove later) --- benchmarks/bench_connectivity.py | 100 ++++++++++--------------------- 1 file changed, 30 insertions(+), 70 deletions(-) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index be3f1571e..bdd5f4243 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -2,8 +2,6 @@ import urllib.request from pathlib import Path -import numpy as np - import uxarray as ux from .helpers._peakmem import peak_allocated @@ -280,84 +278,46 @@ def track_peakmem_node_face(self, resolution): return self._peak_building("node_face_connectivity") -def _save_topology(uxgrid, npz_path): - """Writes the minimal UGRID topology of ``uxgrid`` out to ``npz_path``. - - Compressed because ``setup_cache`` writes into a ``tempfile.mkdtemp()`` and - ``face_node_connectivity`` alone is multi-GB at 3.75km. - """ - np.savez_compressed( - npz_path, - node_lon=uxgrid.node_lon.data, - node_lat=uxgrid.node_lat.data, - face_node_connectivity=uxgrid.face_node_connectivity.data, - ) - - -def _load_topology(npz_path): - """Builds a ``Grid`` holding nothing beyond the minimal UGRID topology.""" - with np.load(npz_path) as topology: - return ux.Grid.from_topology( - topology["node_lon"], - topology["node_lat"], - topology["face_node_connectivity"], - ) - +class ConnectivityChainTracemalloc(MinimalGridBenchmark): + """Peak memory of the whole chain rooted at each connectivity variable. -class ConnectivityChainRss: - """Peak resident memory of the process while constructing each connectivity - variable. + Same instrument as :class:`ConnectivityTracemalloc` -- what the build + allocates, with the ~245MB the process already holds excluded -- but wider + in scope: no prerequisites are put in place beforehand, so a sample covers + everything the variable pulls in, not just the routine that produces it. - Differs from :class:`ConnectivityTracemalloc` on two axes. Scope: no - prerequisites are built in ``setup``, so a sample covers - the whole chain rooted at that variable. Instrument: ``ru_maxrss`` for the - whole process, so the ~250MB import. + The two series coincide for ``n_nodes_per_face``, ``face_node_connectivity`` + and ``node_face_connectivity``, which build straight off the minimal + topology; elsewhere the gap between them is what the prerequisites cost. """ - # Only the parameterization is shared with ``MinimalGridBenchmark`` - param_names = GridBenchmark.param_names - params = GridBenchmark.params - timeout = 1200 - - def setup_cache(self): - topology_paths = {} - for resolution in self.params[0]: - npz_path = os.path.abspath(f"topology_{resolution}.npz") - _save_topology(ux.open_grid(file_path_dict[resolution]), npz_path) - topology_paths[resolution] = npz_path - - _warmup() - - return topology_paths - - setup_cache.timeout = 1800 - - def setup(self, topology_paths, resolution): - self.uxgrid = _load_topology(topology_paths[resolution]) + unit = "bytes" - def teardown(self, topology_paths, resolution): - del self.uxgrid + def _peak_chain(self, name): + """Peak allocation of building ``name`` and everything it rests on.""" + uxgrid = self.minimal_grid() + return peak_allocated(lambda: getattr(uxgrid, name).compute()) - def peakmem_n_nodes_per_face(self, topology_paths, resolution): - _ = self.uxgrid.n_nodes_per_face.compute() + def track_peakmem_n_nodes_per_face(self, resolution): + return self._peak_chain("n_nodes_per_face") - def peakmem_face_node(self, topology_paths, resolution): - _ = self.uxgrid.face_node_connectivity.compute() + def track_peakmem_face_node(self, resolution): + return self._peak_chain("face_node_connectivity") - def peakmem_edge_node(self, topology_paths, resolution): - _ = self.uxgrid.edge_node_connectivity.compute() + def track_peakmem_edge_node(self, resolution): + return self._peak_chain("edge_node_connectivity") - def peakmem_face_edge(self, topology_paths, resolution): - _ = self.uxgrid.face_edge_connectivity.compute() + def track_peakmem_face_edge(self, resolution): + return self._peak_chain("face_edge_connectivity") - def peakmem_node_edge(self, topology_paths, resolution): - _ = self.uxgrid.node_edge_connectivity.compute() + def track_peakmem_node_edge(self, resolution): + return self._peak_chain("node_edge_connectivity") - def peakmem_face_face(self, topology_paths, resolution): - _ = self.uxgrid.face_face_connectivity.compute() + def track_peakmem_face_face(self, resolution): + return self._peak_chain("face_face_connectivity") - def peakmem_edge_face(self, topology_paths, resolution): - _ = self.uxgrid.edge_face_connectivity.compute() + def track_peakmem_edge_face(self, resolution): + return self._peak_chain("edge_face_connectivity") - def peakmem_node_face(self, topology_paths, resolution): - _ = self.uxgrid.node_face_connectivity.compute() + def track_peakmem_node_face(self, resolution): + return self._peak_chain("node_face_connectivity")