diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index b249e7b99..8f1e41416 100644 --- a/benchmarks/face_bounds.py +++ b/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 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/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 e53515f2a..32c87a27d 100644 --- a/benchmarks/import.py +++ b/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" diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 2d47627c1..f8f69b0b0 100644 --- a/benchmarks/mpas_ocean.py +++ b/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) diff --git a/benchmarks/quad_hexagon.py b/benchmarks/quad_hexagon.py index 4364f39b0..5b03eec6f 100644 --- a/benchmarks/quad_hexagon.py +++ b/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"