diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index b249e7b99..d860bfdb1 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -2,6 +2,7 @@ from pathlib import Path import uxarray as ux +from .helpers._memsize import grid_nbytes current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0] @@ -25,6 +26,31 @@ 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" + + +class FaceBoundsPeakMem: + """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.""" + 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..98296ba6a 100644 --- a/benchmarks/import.py +++ b/benchmarks/import.py @@ -3,3 +3,7 @@ class Imports: 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 diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 659c19014..d9befabe7 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__))) @@ -60,16 +61,44 @@ 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.""" + self.uxds[data_var].integrate() + return grid_nbytes(self.uxds.uxgrid) + + 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.""" + 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): diff --git a/benchmarks/quad_hexagon.py b/benchmarks/quad_hexagon.py index 4364f39b0..1e9dd50ad 100644 --- a/benchmarks/quad_hexagon.py +++ b/benchmarks/quad_hexagon.py @@ -2,6 +2,7 @@ from pathlib import Path import uxarray as ux +from .helpers._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"