diff --git a/.github/workflows/asv-benchmarking-pr.yml b/.github/workflows/asv-benchmarking-pr.yml index cae4aa6f0..49ec2cc27 100644 --- a/.github/workflows/asv-benchmarking-pr.yml +++ b/.github/workflows/asv-benchmarking-pr.yml @@ -48,6 +48,8 @@ jobs: id: benchmark run: | set -x + # Fill the fixture cache before asv preimports the suite, which would otherwise build it serially in the forkserver parent + (cd .. && python -m benchmarks.helpers._fixtures) # ID this runner asv machine --yes echo "Baseline: ${{ github.event.pull_request.base.sha }} (${{ github.event.pull_request.base.label }})" diff --git a/.github/workflows/asv-benchmarking.yml b/.github/workflows/asv-benchmarking.yml index aa16fc78a..f5f751cf0 100644 --- a/.github/workflows/asv-benchmarking.yml +++ b/.github/workflows/asv-benchmarking.yml @@ -77,6 +77,8 @@ jobs: shell: bash -l {0} id: benchmark run: | + # Fill the fixture cache before asv preimports the suite, which would otherwise build it serially in the forkserver parent + python -m benchmarks.helpers._fixtures cd benchmarks asv machine --machine GH-Actions --os ubuntu-latest --arch x64 --cpu "2-core unknown" --ram 7GB asv run v2024.02.0..main --skip-existing --parallel || true diff --git a/.gitignore b/.gitignore index c8309635c..ca0ea9bb8 100644 --- a/.gitignore +++ b/.gitignore @@ -163,3 +163,10 @@ docs/user-guide/psi_healpix.nc benchmarks/env benchmarks/results benchmarks/html +benchmarks/_io_cache +# Cookbook meshes, fetched on demand by helpers/_fixtures.py -- same four files +# already ignored above under docs/user-guide. oQU120.grid.nc alone is ~97MB. +benchmarks/oQU120.data.nc +benchmarks/oQU120.grid.nc +benchmarks/oQU480.data.nc +benchmarks/oQU480.grid.nc diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 31a43921d..a30740357 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -61,7 +61,39 @@ // defaults to 10 min "install_timeout": 600, - "benchmark_timeout": 360, + // Fork each benchmark from one interpreter that has already imported the + // suite, rather than starting a fresh one per benchmark. Saves the ~0.9s + // uxarray import and the numba kernel load on every one of ~250 benchmark + // processes, and lets a fixture loaded at import be inherited rather than + // re-read. Measured on bench_connectivity alone, 32 benchmark processes: + // 24s forked against 72s spawned, so roughly 1.5s per process, which over a + // full suite is minutes per commit. That is why this is not "spawn". + // + // The caveat, and it is a real one: forking is only safe while nothing has + // started numba's thread pool before the fork, and importing uxarray does + // exactly that on any commit whose neighbors.py builds its guvectorize + // kernels at module scope -- guvectorize with explicit signatures compiles + // at decoration time, and building a target="parallel" ufunc launches the + // pool. This branch fixes that, but `asv continuous` installs the baseline + // commit too, so until the fix reaches main every baseline, and every + // historical commit before it, forks from a parent that already has a pool. + // + // helpers/_warmup.py reports that instead of failing, so those runs + // complete: look for "asv: numba's thread pool was already running". A run + // with no such line forked from a clean parent and the caveat did not + // apply. Forking with a live pool and then running a parallel kernel in the + // child was checked by hand on macOS/libomp and works; it is not verified + // on the Linux runner, where a GNU OpenMP runtime would be entitled to + // abort the child instead. + // + // So if a benchmark process ever hangs, or dies complaining about fork() + // and OpenMP, set this to "spawn". That removes the hazard outright, at the + // ~1.5s per process above, and nothing else in the suite has to change. + "launch_method": "forkserver", + + // ``benchmark_timeout`` is not a key asv reads -- the one that sets the + // default is ``default_benchmark_timeout`` (asv/config.py) + "default_benchmark_timeout": 360, // the base URL to show a commit for the project. "show_commit_url": "https://github.com/UXARRAY/uxarray/commit/", diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 5e3f53d02..85ca61cb7 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -1,53 +1,25 @@ -import os -import urllib.request -from pathlib import Path - import uxarray as ux -current_path = Path(os.path.dirname(os.path.realpath(__file__))) - -grid_filename_480 = "oQU480.grid.nc" -grid_filename_120 = "oQU120.grid.nc" -filenames = [grid_filename_480, grid_filename_120] - -for filename in filenames: - if not os.path.isfile(current_path / filename): - # downloads the files from Cookbook repo, if they haven't been downloaded locally yet - url = f"https://github.com/ProjectPythia/unstructured-grid-viz-cookbook/raw/main/meshfiles/{filename}" - _, headers = urllib.request.urlretrieve(url, filename=current_path / filename) - -oQU_path_dict = {"480km": current_path / grid_filename_480, - "120km": current_path / grid_filename_120} - -# Paths to grid files on Glade -dyamond_path_dict = {"30km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/30km/grid.nc", - "15km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/15km/grid.nc", - "7.5km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/7.5km/grid.nc", - "3.75km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/3.75km/grid.nc"} - -# Determines if all file paths exist and are accesible -all_paths_exist = True -for file_path in dyamond_path_dict.values(): - all_paths_exist = all_paths_exist and os.path.exists(file_path) - -file_path_dict = oQU_path_dict -if all_paths_exist: - file_path_dict = file_path_dict | dyamond_path_dict +from .helpers._fixtures import ( + ALL_RESOLUTIONS, + GRIDS_BY_RESOLUTION, + CachedFixtures, + cached_topology, + preload_topologies, +) +from .helpers._warmup import warm_in_parent -class GridBenchmark: +class GridBenchmark(CachedFixtures): """Class used as a template for benchmarks requiring a ``Grid`` in this module across both resolutions.""" - param_names = ['resolution', ] - # Conditionally available; could get annoying if there are downstream tools relying on it. - if all_paths_exist: - params = [['480km', '120km', '30km', '15km', '7.5km', '3.75km'], ] - else: - params = [['480km', '120km'], ] + param_names = ['resolution', ] + params = [ALL_RESOLUTIONS, ] + timeout = 1200 def setup(self, resolution, *args, **kwargs): - self.uxgrid = ux.open_grid(file_path_dict[resolution]) + self.uxgrid = self.cached_grid(GRIDS_BY_RESOLUTION[resolution]) def teardown(self, resolution, *args, **kwargs): del self.uxgrid @@ -65,40 +37,36 @@ def teardown(self, resolution, *args, **kwargs): _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 - process would otherwise charge ~240ms of JIT compilation to whichever sample - happened to touch it first. + ``_build_node_edge_connectivity`` is ``@njit`` without ``cache=True``, so a + fresh benchmark process would otherwise charge ~240ms of JIT compilation to + whichever sample happened to touch it first. """ global _numba_warmed_up if _numba_warmed_up: return + uxgrid = ux.Grid.from_topology(*cached_topology(GRIDS_BY_RESOLUTION[ALL_RESOLUTIONS[0]])) for name in CONNECTIVITY_NAMES: getattr(uxgrid, name) _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. + # connectivity is cached in ``Grid._ds`` on construction, so only run them once number = 1 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. - source_grid = ux.open_grid(file_path_dict[resolution]) - self.topology = ( - source_grid.node_lon.data, - source_grid.node_lat.data, - source_grid.face_node_connectivity.data, - ) - - _warmup(self.minimal_grid()) + # the construction routines, so this takes the minimal UGRID topology + # fixture and lets each variable be built on demand. + self.topology = self.cached_topology(GRIDS_BY_RESOLUTION[resolution]) + + # A no-op once the module-level warm below has run; kept so the class is + # still correct if that ever goes away. + _warmup() self.uxgrid = self.minimal_grid() def minimal_grid(self): @@ -139,3 +107,15 @@ def time_edge_face(self, resolution): def time_node_face(self, resolution): _ = self.uxgrid.node_face_connectivity.compute() + + +# Compiled at import rather than in ``setup``. ASV imports the suite once and forks +# every benchmark from that parent, so kernels compiled here are inherited by all +# of them. Only safe while the connectivity kernels are serial +def _warm_parent(): + _warmup() + # And the topologies themselves... + preload_topologies(GRIDS_BY_RESOLUTION[res] for res in ALL_RESOLUTIONS) + + +warm_in_parent(_warm_parent, "the connectivity kernels") diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index 8f1e41416..0caf838df 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -1,18 +1,16 @@ -import os -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] +from .helpers._fixtures import GRIDS_BY_FORMAT, CachedFixtures, preload_grids +from .helpers._memsize import grid_nbytes +from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss +from .helpers._warmup import warm_in_parent -grid_quad_hex = current_path / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "grid.nc" -grid_geoflow = current_path / "test" / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc" -grid_scrip = current_path / "test" / "meshfiles" / "scrip" / "outCSne8" / "outCSne8.nc" -grid_mpas= current_path / "test" / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc" +grid_quad_hex = GRIDS_BY_FORMAT["ugrid-quad-hexagon"] +grid_geoflow = GRIDS_BY_FORMAT["ugrid-geoflow"] +grid_scrip = GRIDS_BY_FORMAT["scrip-outCSne8"] +grid_mpas = GRIDS_BY_FORMAT["mpas-oQU480"] -class FaceBounds: +class FaceBounds(CachedFixtures): params = [grid_quad_hex, grid_geoflow, grid_scrip, grid_mpas] @@ -24,8 +22,8 @@ def setup(self, grid_path): # 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) + self.cached_grid(grid_quad_hex).bounds + self.uxgrid = self.cached_grid(grid_path) def teardown(self, n): del self.uxgrid @@ -63,21 +61,35 @@ def track_peakmem_face_bounds(self, grid_path): 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. + Whole-process peak resident memory, 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. + + Measured in a subprocess of its own rather than through asv's ``peakmem_*``, + which reports ``ru_maxrss`` for the benchmark process. """ params = FaceBounds.params param_names = ["grid_path"] def setup_cache(self): - """Compile the njit kernels before anything is measured.""" + """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 + def track_peakmem_open_and_bounds(self, grid_path): + return subprocess_peak_rss( + f"import uxarray as ux; ux.open_grid({str(grid_path)!r}).bounds" + ) + + track_peakmem_open_and_bounds.unit = "bytes" + + +# Inherited by every forked benchmark rather than re-read in each; see the same +# note in ``mpas_ocean``. These four are the small ones, ~3.7MB all told. +warm_in_parent(lambda: preload_grids(GRIDS_BY_FORMAT.values()), "the source-format grids") diff --git a/benchmarks/geometry_kernels.py b/benchmarks/geometry_kernels.py index 08aea236e..ceb9f995c 100644 --- a/benchmarks/geometry_kernels.py +++ b/benchmarks/geometry_kernels.py @@ -15,6 +15,8 @@ import numpy as np +from .helpers._warmup import warm_in_parent + def _unit(v): return v / np.linalg.norm(v) @@ -205,3 +207,25 @@ def time_try_gca_const_lat_intersection(self): def time_gca_const_lat_intersection(self): """Layer 3: dispatcher (full public API).""" self.gca_const_lat_intersection(self.gca_cart, _CONST_Z) + + +def _warm_classes(): + """Compiles what each class's ``setup`` compiles, once per process. + + Runs the setups themselves rather than a copy of their warm calls, so this + cannot drift out of step with them. + """ + for cls in ( + EFTPrimitives, + AccucrossKernels, + OrientPredicates, + GCAGCAIntersection, + GCAConstLatIntersection, + ): + try: + cls().setup() + except Exception: + pass + + +warm_in_parent(_warm_classes, "the geometry kernels") diff --git a/benchmarks/geometry_samebody.py b/benchmarks/geometry_samebody.py index 7712d1c7a..bee8cdf4b 100644 --- a/benchmarks/geometry_samebody.py +++ b/benchmarks/geometry_samebody.py @@ -60,6 +60,8 @@ gca_const_lat_intersection, ) +from .helpers._warmup import warm_in_parent + # --------------------------------------------------------------------------- # L1 (FP64 body) — direct double-precision kernel, verbatim from # fp64_GCAconstLat.hh. Scalar in / scalar out so Numba keeps it in registers, @@ -392,16 +394,33 @@ def main(): # --------------------------------------------------------------------------- +_prepared = None + + +def _prepare(): + """The packed cases, with the batched drivers warmed, once per process. + + Hoisted out of ``setup`` so a forked benchmark inherits them: the seed is + fixed, so the arrays are the same ones ``setup`` used to build, and the + drivers are ``@njit(cache=True)`` either way. + """ + global _prepared + if _prepared is None: + packed = _pack(_make_cases(20_000, seed=20251104)) + A, B, Z, gcas = packed + _batch_fp64_kernel(A, B, Z) + _batch_accux_kernel(A, B, Z) + _batch_fp64_dispatch(gcas, Z) + _batch_accux_dispatch(gcas, Z) + _prepared = packed + return _prepared + + class SameBodyConstLat: """asv: same-body FP64 vs real AccuX at kernel (L1) and dispatch (L3) levels.""" def setup(self): - cases = _make_cases(20_000, seed=20251104) - self.A, self.B, self.Z, self.gcas = _pack(cases) - _batch_fp64_kernel(self.A, self.B, self.Z) - _batch_accux_kernel(self.A, self.B, self.Z) - _batch_fp64_dispatch(self.gcas, self.Z) - _batch_accux_dispatch(self.gcas, self.Z) + self.A, self.B, self.Z, self.gcas = _prepare() def time_fp64_kernel(self): _batch_fp64_kernel(self.A, self.B, self.Z) @@ -416,5 +435,10 @@ def time_accux_dispatch(self): _batch_accux_dispatch(self.gcas, self.Z) +# Prepared at import so every forked benchmark inherits it; see +# :mod:`benchmarks.helpers._warmup`. +warm_in_parent(_prepare, "the const-lat drivers") + + if __name__ == "__main__": main() diff --git a/benchmarks/geometry_samebody_gcagca.py b/benchmarks/geometry_samebody_gcagca.py index 348f55859..fd3494ed3 100644 --- a/benchmarks/geometry_samebody_gcagca.py +++ b/benchmarks/geometry_samebody_gcagca.py @@ -23,6 +23,8 @@ from uxarray.grid.arcs import on_minor_arc from uxarray.grid.intersections import _accux_gca, gca_gca_intersection +from .helpers._warmup import warm_in_parent + @njit(cache=True, inline="always") def _fp64_gca(w0, w1, v0, v1): @@ -307,6 +309,29 @@ def ns(t): print("=" * 70) +_prepared = None + + +def _prepare(): + """The packed cases, with the batched drivers warmed, once per process. + + Building the cases is 4.06s of the 4.27s this used to spend in every + ``setup``; compiling the drivers is 0.08s, since they are all + ``@njit(cache=True)``. This method allows for reuse of cases in forked + benchmarks to reduce time spent on case generation. + """ + global _prepared + if _prepared is None: + packed = _pack_gca(_make_gca_cases(100_000, seed=20251104)) + wa, wb, va, vb, ga, gb = packed + _batch_fp64_gca_kernel(wa, wb, va, vb) + _batch_accux_gca_kernel(wa, wb, va, vb) + _batch_fp64_gca_dispatch(ga, gb) + _batch_accux_gca_dispatch(ga, gb) + _prepared = packed + return _prepared + + class SameBodyGcaGca: """ asv timing class (Numba warmed in setup, distinct cases) @@ -314,12 +339,7 @@ class SameBodyGcaGca: """ def setup(self): - cases = _make_gca_cases(100_000, seed=20251104) - self.wa, self.wb, self.va, self.vb, self.ga, self.gb = _pack_gca(cases) - _batch_fp64_gca_kernel(self.wa, self.wb, self.va, self.vb) - _batch_accux_gca_kernel(self.wa, self.wb, self.va, self.vb) - _batch_fp64_gca_dispatch(self.ga, self.gb) - _batch_accux_gca_dispatch(self.ga, self.gb) + self.wa, self.wb, self.va, self.vb, self.ga, self.gb = _prepare() def time_fp64_kernel(self): _batch_fp64_gca_kernel(self.wa, self.wb, self.va, self.vb) @@ -334,5 +354,8 @@ def time_accux_dispatch(self): _batch_accux_gca_dispatch(self.ga, self.gb) +warm_in_parent(_prepare, "the gca-gca drivers") + + if __name__ == "__main__": main() diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py new file mode 100644 index 000000000..99852457f --- /dev/null +++ b/benchmarks/helpers/_fixtures.py @@ -0,0 +1,433 @@ +"""Cached inputs for benchmarks whose subject is not reading a file. + +ASV gives every benchmark its own process, so a grid opened in ``setup`` is +opened once per benchmark rather than once per run. This module provides flexible +access to files across benchmark runs by caching the needed files. + +Benchmarks that do measure opening a file behave as before. + +Two flavors: + +``topology`` + the three arrays ``Grid.from_topology`` needs and nothing else +``grid`` / ``dataset`` + everything the reader produced from ``Grid.open_grid`` and + ``Grid.open_dataset`` + +Artifacts are keyed on both the uxarray build and the files, because an +artifact is one version's reader output and ASV diffs commits. Likewise, there's a +fresh read per commit. ``prime`` covers every source that is readable here, +and there is a CLI (``python -m benchmarks.helpers._fixtures``) to fill the +cache from a batch script instead of from inside a benchmark. +""" + +import hashlib +import multiprocessing +import os +import urllib.request +import uuid +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import numpy as np +import xarray as xr + +import uxarray as ux + +from ._warmup import warm_in_parent + +__all__ = [ + "ALL_RESOLUTIONS", + "DYAMOND_AVAILABLE", + "DYAMOND_GRIDS", + "GRIDS_BY_FORMAT", + "GRIDS_BY_RESOLUTION", + "OQU_DATASETS", + "OQU_GRIDS", + "OQU_RESOLUTIONS", + "QUAD_HEXAGON_DATASET", + "CachedFixtures", + "cache_dir", + "cached_dataset", + "cached_grid", + "cached_topology", + "preload_datasets", + "preload_grids", + "preload_topologies", + "prime", + "warm_netcdf", +] + +BENCHMARK_DIR = Path(__file__).resolve().parents[1] +REPO_DIR = BENCHMARK_DIR.parent + +_COOKBOOK_MESH_URL = ( + "https://github.com/ProjectPythia/unstructured-grid-viz-cookbook/raw/main/meshfiles" +) + + +def _cookbook_mesh(filename): + """Path to a Cookbook mesh, fetched once if this checkout lacks it.""" + path = BENCHMARK_DIR / filename + if not path.is_file(): + urllib.request.urlretrieve(f"{_COOKBOOK_MESH_URL}/{filename}", filename=path) + return path + + +# Grids and grid/data pairs, by mesh resolution. +OQU_GRIDS = { + "480km": _cookbook_mesh("oQU480.grid.nc"), + "120km": _cookbook_mesh("oQU120.grid.nc"), +} +OQU_DATASETS = { + "480km": (OQU_GRIDS["480km"], _cookbook_mesh("oQU480.data.nc")), + "120km": (OQU_GRIDS["120km"], _cookbook_mesh("oQU120.data.nc")), +} + +DYAMOND_GRIDS = { + "30km": Path("/glade/campaign/cisl/vast/uxarray/data/dyamond/30km/grid.nc"), + "15km": Path("/glade/campaign/cisl/vast/uxarray/data/dyamond/15km/grid.nc"), + "7.5km": Path("/glade/campaign/cisl/vast/uxarray/data/dyamond/7.5km/grid.nc"), + "3.75km": Path("/glade/campaign/cisl/vast/uxarray/data/dyamond/3.75km/grid.nc"), +} + +# Find out which files are actually available +DYAMOND_AVAILABLE = all(path.exists() for path in DYAMOND_GRIDS.values()) + +GRIDS_BY_RESOLUTION = dict(OQU_GRIDS) +if DYAMOND_AVAILABLE: + GRIDS_BY_RESOLUTION |= DYAMOND_GRIDS + +OQU_RESOLUTIONS = list(OQU_GRIDS) +ALL_RESOLUTIONS = list(GRIDS_BY_RESOLUTION) + +# Grids by source format, for benchmarks whose axis is the reader rather than the +# mesh size. ``mpas-oQU480`` is the same 1,791-face mesh as ``480km`` above, +# reached through the copy in the repo instead of the Cookbook download. +GRIDS_BY_FORMAT = { + "ugrid-quad-hexagon": REPO_DIR / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "grid.nc", + "ugrid-geoflow": REPO_DIR / "test" / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc", + "scrip-outCSne8": REPO_DIR / "test" / "meshfiles" / "scrip" / "outCSne8" / "outCSne8.nc", + "mpas-oQU480": REPO_DIR / "test" / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc", +} + +QUAD_HEXAGON_DATASET = ( + GRIDS_BY_FORMAT["ugrid-quad-hexagon"], + REPO_DIR / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "data.nc", +) + +# Per path, which artifacts are actually loaded +_loaded = {} + + +def cache_dir(): + """Directory the cached artifacts live in: ``benchmarks/_io_cache``.""" + root = Path(os.environ.get("UXARRAY_BENCH_CACHE_DIR") or BENCHMARK_DIR) + cached = root / "_io_cache" + cached.mkdir(parents=True, exist_ok=True) + return cached + + +def _artifact(source, flavor, suffix): + """Where ``flavor`` of ``source`` is cached. + + ``source`` is a grid path, or a (grid, data) pair. Keyed on each file's size + and mtime as well as its path, so a replaced source misses rather than being + served something stale. + """ + parts = [] + for path in source: + stat = os.stat(path) + parts.append(f"{os.path.realpath(path)}:{stat.st_size}:{stat.st_mtime_ns}") + digest = hashlib.sha256("|".join(parts).encode()).hexdigest()[:16] + # The dyamond grids are all named ``grid.nc``; the directory above them is + # what tells one resolution from the next. + grid_path = Path(source[0]) + stem = f"{grid_path.parent.name}-{grid_path.stem}" + return cache_dir() / f"{stem}-{flavor}-{digest}{suffix}" + + +def _write(dataset, artifact_path, writer): + """Writes ``dataset`` to ``artifact_path`` atomically. + + Via a scratch name in the same directory, so a process racing this one sees + either no artifact or a complete one. + """ + if artifact_path.exists(): + # If the path exists, someone else is already working + return + scratch = artifact_path.with_name( + f"{artifact_path.stem}.{uuid.uuid4().hex}.tmp{artifact_path.suffix}" + ) + writer(dataset, scratch) + os.replace(scratch, artifact_path) + + +def _read_dataset(artifact_path): + """Reads back a cached ``xr.Dataset``, holding no file handle afterwards. + + ``load()`` brings every array into memory, but on its own it leaves the + netCDF handle behind it open in xarray's file cache. Under + ``launch_method: forkserver`` a handle still open when the suite finishes + importing is inherited by every benchmark process forked from that + interpreter, and HDF5 does not support being used across a fork. + """ + with xr.open_dataset(artifact_path, mask_and_scale=False) as dataset: + return dataset.load() + + +def _build(source): + """Reads ``source`` and writes every flavor of it, from the one read.""" + grid_path = source[0] + if len(source) == 1: + uxgrid = ux.open_grid(grid_path) + uxgrid._ds.load() + data_ds = None + else: + uxds = ux.open_dataset(*source) + uxds.load() + uxgrid = uxds.uxgrid + uxgrid._ds.load() + # A ``UxDataset`` grid is cached separately. + data_ds = uxds + uxds.close() + + # Everything is in memory by now, so the source files are done with. Closed + # rather than left to the garbage collector because on a cold cache this + # runs at import, in the interpreter every benchmark is forked from, and an + # inherited HDF5 handle is not something that library supports. + uxgrid._ds.close() + + _write( + {name: getattr(uxgrid, name).data for name in ["node_lon", "node_lat", "face_node_connectivity"]}, + _artifact(source[:1], "topology", ".npz"), + # Uncompressed: written once, read by every benchmark process after. + lambda arrays, path: np.savez(path, **arrays), + ) + _write(uxgrid._ds, _artifact(source[:1], "grid", ".nc"), lambda ds, path: ds.to_netcdf(path)) + if data_ds is not None: + _write(data_ds, _artifact(source, "data", ".nc"), lambda ds, path: ds.to_netcdf(path)) + + +def _ensure(source, flavor, suffix): + """Path to a cached artifact, building the source's artifacts if need be.""" + artifact_path = _artifact(source, flavor, suffix) + if not artifact_path.exists(): + _build(source if flavor == "data" else source[:1]) + return artifact_path + + +def cached_topology(grid_path): + """``(node_lon, node_lat, face_node_connectivity)`` for ``grid_path``. + + The arrays are shared rather than copied: ``Grid.from_topology`` wraps + ``node_lon`` and ``node_lat`` without copying them, and the construction + routines assign to ``Grid._ds`` rather than writing through their inputs, + which is what makes one copy per process safe. Treat them as read-only. + """ + artifact_path = _ensure((Path(grid_path),), "topology", ".npz") + if artifact_path not in _loaded: + with np.load(artifact_path) as cached: + _loaded[artifact_path] = tuple(cached[name] for name in ["node_lon", "node_lat", "face_node_connectivity"]) + return _loaded[artifact_path] + + +def _cached_grid_ds(grid_path): + """The cached internal dataset of ``grid_path``, held for this process.""" + artifact_path = _ensure((Path(grid_path),), "grid", ".nc") + if artifact_path not in _loaded: + _loaded[artifact_path] = _read_dataset(artifact_path) + return _loaded[artifact_path] + + +def cached_grid(grid_path): + """A fresh ``Grid`` carrying everything the reader found in ``grid_path`` via shallow copy.""" + return ux.Grid(_cached_grid_ds(grid_path).copy()) + + +def _cached_data_ds(grid_path, data_path): + """The cached dataset of a ``(grid, data)`` pair, held for this process.""" + artifact_path = _ensure((Path(grid_path), Path(data_path)), "data", ".nc") + if artifact_path not in _loaded: + _loaded[artifact_path] = _read_dataset(artifact_path) + return _loaded[artifact_path] + + +def cached_dataset(grid_path, data_path): + """A ``UxDataset`` over ``data_path``, on the cached grid.""" + return ux.UxDataset( + _cached_data_ds(grid_path, data_path).copy(), uxgrid=cached_grid(grid_path) + ) + + +def prime(workers=1): + """Fills the cache for every source the fixtures can serve. + + Returns the sources it had to read. Idempotent, and once warm costs a + ``stat`` per file, so it is cheap to call ahead of every run. + + ``workers`` reads that many sources at once, which is worth having when the + sources differ wildly in size and live somewhere slow: the small ones finish + while a larger grid is still being read, instead of queueing behind all the + file fetches. Only pays off when a read is slower than a process start due to + the new interpreter startup. Processes rather than threads because the standard + netCDF/HDF5 stack is not generally thread-safe for concurrent opens. + """ + sources = [(path,) for path in OQU_GRIDS.values()] + sources += [(path,) for path in GRIDS_BY_FORMAT.values()] + sources += list(OQU_DATASETS.values()) + if DYAMOND_AVAILABLE: + sources += [(path,) for path in DYAMOND_GRIDS.values()] + + missing = [] + for source in sources: + flavor, suffix = ("data", ".nc") if len(source) == 2 else ("grid", ".nc") + if not _artifact(source, flavor, suffix).exists(): + missing.append(source) + + # Reading a (grid, data) pair produces that grid's artifacts too + paired_grids = {source[0] for source in missing if len(source) == 2} + missing = [ + source for source in missing if len(source) == 2 or source[0] not in paired_grids + ] + + if workers > 1 and len(missing) > 1: + # Largest first: with a pool, the longest read should start earliest, or + # it lands last and everything waits on it. + missing.sort(key=lambda source: -sum(os.path.getsize(path) for path in source)) + + # Spawned. This process has the netCDF/HDF5 library loaded by the time it primes, + # and a fresh interpreter per worker keeps that state out of the children. + with ProcessPoolExecutor( + max_workers=min(workers, len(missing)), + mp_context=multiprocessing.get_context("spawn"), + ) as pool: + list(pool.map(_build, missing)) + else: + for source in missing: + _build(source) + return missing + + +def preload_topologies(grid_paths): + """Loads topologies here so forked benchmarks inherit them. + + Reading an artifact caches it on disk, not in the next process: under + ``launch_method: forkserver`` each benchmark is forked from the interpreter + that imported the suite, so it starts with whatever *that* process holds and + reads its own copy of everything else. + + Safe to call at import: reading arrays starts no numba thread pool, which is + what a forked child cannot inherit (see :mod:`benchmarks.helpers._warmup`). + """ + loaded = 0 + for grid_path in grid_paths: + cached_topology(grid_path) # held by the process-level memo from here on + loaded += 1 + return loaded + + +def preload_grids(grid_paths): + """Loads cached ``grid`` artifacts here so forked benchmarks inherit them. + + The counterpart of :func:`preload_topologies` for the flavor that carries + everything the reader produced, and it rests on the same thing: under + ``launch_method: forkserver`` a benchmark starts with whatever the + interpreter that imported the suite is holding, and reads its own copy of + anything else. + + Same caveat as well -- the arrays are shared rather than copied, so treat + them as read-only. ``cached_grid`` hands out a shallow copy of the dataset, + which is what keeps that safe. + + This costs the parent the size of every artifact named here for the length + of the run. Cheaper in total than the alternative, because the forks share + those pages copy-on-write instead of each holding a private read, but it is + resident rather than transient -- which is why the dyamond grids are left + out of the callers below and read per fork instead. + """ + loaded = 0 + for grid_path in grid_paths: + _cached_grid_ds(grid_path) # held by the process-level memo from here on + loaded += 1 + return loaded + + +def preload_datasets(sources): + """Loads cached ``(grid, data)`` artifacts, and their grids, for the same reason.""" + loaded = 0 + for grid_path, data_path in sources: + _cached_data_ds(grid_path, data_path) + _cached_grid_ds(grid_path) + loaded += 1 + return loaded + + +def warm_netcdf(): + """Pays netCDF4/HDF5's per-process start-up cost here, once, before any fork. + + The *first* ``xr.open_dataset`` in an interpreter spends ~164ms resolving + xarray's backend entry points and bringing up the netCDF library, and only + the first -- reads after it cost single-digit milliseconds even for the + 18MB oQU120 grid. Once the cache is warm nothing else in this module opens a + netCDF file at import (:func:`preload_topologies` goes through numpy), so + without this every forked benchmark that reaches ``cached_grid`` or + ``cached_dataset`` pays that start-up again on its own account. Roughly 134 + of the ~235 benchmark processes in a run do. + + Deliberately opens the smallest grid in the repository rather than a cached + artifact, so it still works before the cache exists, and closes it again so + nothing is inherited across the fork. + """ + try: + with xr.open_dataset(GRIDS_BY_FORMAT["ugrid-quad-hexagon"]) as dataset: + dataset.load() + except (KeyError, OSError, ValueError, RuntimeError): + # Nothing to warm is not a failure -- the forks just pay for themselves. + pass + + +class CachedFixtures: + """Mixin for benchmarks whose subject is not reading a file. + + Holds the one canonical ``setup_cache`` the suite shares. ASV keys ``setup_cache`` + on where it is defined and groups benchmarks by that key, so this single + definition runs once per ``asv run`` rather than once per class. It returns + ``None``, which asv reads as "no cache argument". + + The accessors are re-exported as methods so a ``setup`` reads as + ``self.cached_grid(...)`` instead of importing the module's functions + alongside its registries. + """ + + def setup_cache(self): + prime() + + # Reading grids off campaign storage is the point of the cache, and does not + # fit in a benchmark-sized timeout. + setup_cache.timeout = 7200 + + cached_topology = staticmethod(cached_topology) + cached_grid = staticmethod(cached_grid) + cached_dataset = staticmethod(cached_dataset) + + +# Run at import so the whole suite inherits an initialised netCDF stack. Through +# ``warm_in_parent`` for the two things it decides: skip this in the interpreters +# that only import the suite to enumerate it, and fail loudly if warming ever +# starts a numba thread pool a fork could not inherit. +warm_in_parent(warm_netcdf, "the netCDF backend") + + +if __name__ == "__main__": + # Fills the cache ahead of ``asv run``, so no benchmark -- and not even + # ``setup_cache`` -- pays for reading a source grid. Worth a line in a batch + # script whenever the dyamond grids are in play. + print(f"fixture cache: {cache_dir()}", flush=True) + + # Four at a time only where the dyamond grids are readable, since those are + # the reads worth overlapping: on the oQU pair alone, priming in parallel is + # slower than doing it sequentially (1.69s against 0.52s), because starting + # an interpreter costs more than reading a small local file. + for source in prime(workers=4 if DYAMOND_AVAILABLE else 1) or [None]: + print(f" read {' + '.join(Path(p).name for p in source)}" if source else " nothing to do", flush=True) diff --git a/benchmarks/helpers/_peakmem.py b/benchmarks/helpers/_peakmem.py index 0265d80f8..011e03034 100644 --- a/benchmarks/helpers/_peakmem.py +++ b/benchmarks/helpers/_peakmem.py @@ -1,7 +1,9 @@ import contextlib +import os import subprocess import sys +import tempfile import numba @@ -66,15 +68,22 @@ def subprocess_peak_rss(statement): ``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 + with tempfile.TemporaryDirectory() as scratch: + report_path = os.path.join(scratch, "peak_rss") + reporter = ( + "import resource\n" + f"exec({statement!r})\n" + "peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss\n" + # Reported through a file rather than stdout, which is not ours + # alone: uxarray prints coordinate warnings there for some grids, + # and the number would arrive with prose in front of it. + f"open({report_path!r}, 'w').write(str(peak))\n" + ) + subprocess.run( + [sys.executable, "-c", reporter], + capture_output=True, + text=True, + check=True, + ) + with open(report_path) as report: + return int(report.read()) * _MAXRSS_TO_BYTES diff --git a/benchmarks/helpers/_shard.py b/benchmarks/helpers/_shard.py new file mode 100644 index 000000000..151fc3331 --- /dev/null +++ b/benchmarks/helpers/_shard.py @@ -0,0 +1,660 @@ +"""Running the ``track_*`` half of the suite as concurrent ``asv run`` processes. + +asv never runs two benchmarks at once. ``--parallel`` builds environments in +parallel and does nothing else -- its own help text is "Build (but don't +benchmark) in parallel" (``asv/commands/common_args.py``) -- and the forkserver +is a ``listen(1)`` socket that forks one child and ``waitpid``s on it before it +will read another command (``asv_runner/server.py``). Concurrency has to come +from running several ``asv run`` processes side by side, which is what this +module does. + +Only the ``track_*`` benchmarks are safe to run that way, and they are also the +ones worth running that way: + +* A ``time_*`` result is a wall-clock measurement, and a co-tenant on the + machine corrupts it. Those stay serial. +* A ``track_*`` result is a count. ``track_nbytes_*`` returns an array size and + ``track_peakmem_*`` returns a tracemalloc high-water figure or a child's + ``ru_maxrss``; none of them move because another core is busy. +* Every traced benchmark runs under ``numba_threads(1)`` (see + :mod:`benchmarks.helpers._peakmem`), so today the whole ``track_`` pass leaves + every core but one idle while it runs the most expensive benchmarks in the + suite. + +What is *not* invariant is the thread count itself. The peaks measured through +``subprocess_peak_rss`` come from a child that runs numba's kernels at whatever +width the platform gives it, and numba follows CPU affinity -- so N shards +sharing a node would each see fewer threads and report a different number than +an unsharded run would. ``--threads-per-shard`` pins that width explicitly (1 by +default, which is what the traced benchmarks already do to themselves), so the +figure stays a property of the benchmark rather than of how the run happened to +be sharded. Expect a one-time step in the ``subprocess_peak_rss`` series the +first time this is used, and none after. + +A run goes:: + + prime fill the fixture cache once, here, rather than in N shards + discover one serial ``asv run --bench just-discover``, which builds and + installs the project at the commit and writes the full + ``benchmarks.json`` + shards N concurrent ``asv run --bench ...``, each with its own + ``results_dir`` so they cannot clobber one another's file + merge fold the shard result files into the real results dir + serial the remaining (``time_*``) benchmarks, one process, as usual + +The discover pass is what lets the shards start together: +``Environment.install_project`` returns immediately once ``installed_commit_hash`` +matches the commit, so the shards find the project built instead of N of them +racing to build it into the one shared environment. + +The serial pass runs last on purpose. ``asv run`` calls ``Results.load_data`` +before it records anything (``asv/commands/run.py``), so it reads the merged +track results off disk and writes them back out alongside the timings rather +than replacing the file. + +Usage, from the repository root:: + + python -m benchmarks.helpers._shard # branch head + python -m benchmarks.helpers._shard main^! --workers 8 + python -m benchmarks.helpers._shard v2024.02.0..main -- --skip-existing + +Anything after a bare ``--`` is handed to every ``asv run`` this launches. +""" + +import argparse +import copy +import itertools +import json +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +__all__ = ["main", "merge_results", "missing_results", "plan_shards"] + +BENCHMARK_DIR = Path(__file__).resolve().parents[1] +REPO_DIR = BENCHMARK_DIR.parent + +# Columns of an asv result row that are indexed by parameter combination, and so +# merge element-wise across shards. Everything else in the row describes the +# benchmark as a whole. Names come from ``Results.save`` (``asv/results.py``). +_PER_PARAM_COLUMNS = frozenset( + { + "result", + "samples", + "stats_ci_99_a", + "stats_ci_99_b", + "stats_q_25", + "stats_q_75", + "stats_number", + "stats_repeat", + } +) + +# ``setup_cache`` groups that may be split across shards. Every other group in +# this suite compiles kernels and computes at every resolution before anything is +# measured -- they carry ``setup_cache.timeout = 1800`` for that reason -- so +# splitting one would pay that cost once per shard. ``CachedFixtures.setup_cache`` +# is just ``prime()``, which costs a stat per file once the cache is warm, and +# ``None`` is no setup_cache at all. +_SPLITTABLE_PREFIX = "helpers._fixtures:" + + +def _splittable(setup_cache_key): + """Whether instances sharing this ``setup_cache_key`` may land in different shards.""" + return setup_cache_key is None or str(setup_cache_key).startswith(_SPLITTABLE_PREFIX) + + +def _asv(): + """Imports asv, with a message pointing at the environment if it is absent. + + Imported lazily: this module is a CLI that drives ``asv``, and asv itself is + not a dependency of the benchmarks it runs. The project environment asv + builds has no asv in it. + """ + try: + import asv.config + import asv.results + import asv.util + except ImportError as exc: # pragma: no cover - environment problem, not logic + raise SystemExit( + "benchmarks.helpers._shard needs asv importable in the interpreter " + f"running it (the one that provides the `asv` command): {exc}" + ) + return asv + + +def _instance_names(benchmark): + """Every selectable name of ``benchmark``, as asv's ``--bench`` matches them. + + asv filters parameterized benchmarks on ``funcname(param0, param1)`` built + from the parameter *reprs* it recorded at discovery, not on the values + (``Benchmarks.__init__`` in ``asv/benchmarks.py``), so this reproduces that + string exactly rather than re-deriving it from the benchmark's own params. + """ + params = benchmark.get("params") or [] + if not params: + return [benchmark["name"]] + return [ + f"{benchmark['name']}({', '.join(combo)})" + for combo in itertools.product(*params) + ] + + +def _weights(results_dir): + """Per-instance runtime hints, in seconds, from whatever results are on disk. + + Bin packing on measured cost beats packing on a count, because this suite is + lopsided: one traced gradient runs longer than every ``track_nbytes_*`` + benchmark put together. Missing entries fall back to a flat weight, so a + first run still shards, just less evenly. + + asv records one ``duration`` per benchmark rather than per parameter + combination, so this spreads it evenly over the combinations. + """ + asv = _asv() + hints = {} + root = Path(results_dir) + if not root.is_dir(): + return hints + for path in sorted(root.glob("*/*.json")): + if path.name in ("machine.json", "benchmarks.json"): + continue + try: + data = asv.util.load_json(path, api_version=asv.results.Results.api_version) + except Exception: + continue + columns = data.get("result_columns") or [] + if "duration" not in columns: + continue + index = columns.index("duration") + for name, row in data.get("results", {}).items(): + if len(row) <= index or row[index] is None: + continue + duration = row[index] + if isinstance(duration, list): + duration = sum(x for x in duration if x) + hints[name] = float(duration) + return hints + + +def plan_shards(benchmarks, selected, n_shards, weights=None): + """Partitions ``selected`` instance names into ``n_shards`` lists. + + ``benchmarks`` is the parsed ``benchmarks.json``. Units that share a + non-splittable ``setup_cache_key`` move together, so an expensive + ``setup_cache`` is paid once rather than once per shard. Longest-first onto + the lightest shard, which is the standard greedy fit and is good enough here: + the schedule is dominated by one benchmark that no packing can subdivide. + + Returns a list of ``n_shards`` lists of instance names, shortest lists last. + """ + weights = weights or {} + by_name = {b["name"]: b for b in benchmarks} + + # Build the units: either one instance, or a whole setup_cache group. + units = {} + for instance in selected: + benchmark_name = instance.split("(", 1)[0] + key = by_name.get(benchmark_name, {}).get("setup_cache_key") + unit = instance if _splittable(key) else f"setup_cache:{key}" + units.setdefault(unit, []).append(instance) + + def cost(instances): + total = 0.0 + for instance in instances: + benchmark_name = instance.split("(", 1)[0] + hint = weights.get(benchmark_name) + benchmark = by_name.get(benchmark_name) + if hint is None or benchmark is None: + total += 1.0 + else: + total += hint / max(1, len(_instance_names(benchmark))) + return total + + shards = [[] for _ in range(n_shards)] + loads = [0.0] * n_shards + for instances in sorted(units.values(), key=lambda i: -cost(i)): + target = loads.index(min(loads)) + shards[target].extend(instances) + loads[target] += cost(instances) + return [s for s in shards if s] + + +def _shard_home(scratch, index): + """A private ``HOME`` for one shard, seeded with the files asv reads from it. + + Every ``asv run`` ends ``Machine.load`` with an unconditional + ``MachineCollection.save`` back to ``~/.asv-machine.json``, and + ``util.write_json`` opens that path ``'w'`` -- truncating it before it + writes. Concurrent shards therefore race on one file, and a shard that + reads it inside another's truncate window dies with "Error parsing JSON in + file ... Expecting value: line 1 column 1". Observed, repeatedly: it takes + out that shard's whole share of the suite. + + asv resolves the path through ``os.path.expanduser``, so a private ``HOME`` + per shard gives each its own copy to rewrite and leaves the real one alone. + ``.condarc`` comes along because asv shells out to conda for the + environment, and the machine file because that is the point. + """ + home = os.path.join(scratch, f"shard{index}", "home") + os.makedirs(home, exist_ok=True) + for name in (".asv-machine.json", ".condarc"): + source = os.path.join(os.path.expanduser("~"), name) + if os.path.isfile(source): + shutil.copyfile(source, os.path.join(home, name)) + return home + + +def _shard_config(raw_conf, conf, results_dir, html_dir): + """A copy of the project's asv config that writes results somewhere else. + + Every path is absolutised. ``asv --config`` chdirs to the directory holding + the config file (``asv/main.py``), so a shard config in a scratch directory + would otherwise resolve ``repo``, ``benchmark_dir`` and ``env_dir`` against + that scratch directory instead of the project. + + ``env_dir`` deliberately stays the shared one: the shards run against the + environment the discover pass already installed into. + """ + shard = copy.deepcopy(raw_conf) + shard["repo"] = os.path.abspath(conf.repo) + shard["benchmark_dir"] = os.path.abspath(conf.benchmark_dir) + shard["env_dir"] = os.path.abspath(conf.env_dir) + shard["results_dir"] = os.path.abspath(results_dir) + shard["html_dir"] = os.path.abspath(html_dir) + return shard + + +def merge_results(shard_dirs, dest_dir): + """Folds each shard's result file into the matching file in ``dest_dir``. + + Shards run disjoint benchmarks, but two shards can hold different parameter + combinations of the *same* benchmark, and asv writes those as one row per + benchmark with ``None`` in the combinations it did not run. So rows merge + element-wise on the parameter-indexed columns and take the first value on + the rest. + + Returns the list of files written. + """ + asv = _asv() + api_version = asv.results.Results.api_version + written = [] + + # Group by the per-(machine, commit, env) filename asv chose. + files = {} + for shard_dir in shard_dirs: + for path in sorted(Path(shard_dir).glob("*/*.json")): + if path.name in ("machine.json", "benchmarks.json"): + continue + files.setdefault(str(Path(path.parent.name) / path.name), []).append(path) + + for relative, paths in sorted(files.items()): + dest = Path(dest_dir) / relative + merged = None + columns = None + if dest.is_file(): + merged = asv.util.load_json(dest, api_version=api_version) + columns = merged.get("result_columns") + + for path in paths: + data = asv.util.load_json(path, api_version=api_version) + if merged is None: + merged = data + columns = data.get("result_columns") + continue + if data.get("result_columns") != columns: + raise RuntimeError( + f"{path} has different result columns than the file it merges " + "into; the asv versions writing them differ" + ) + for name, row in data.get("results", {}).items(): + existing = merged["results"].get(name) + merged["results"][name] = ( + list(row) if existing is None else _merge_row(existing, row, columns) + ) + for key, value in (data.get("durations") or {}).items(): + merged.setdefault("durations", {}) + merged["durations"][key] = merged["durations"].get(key, 0) + value + + # asv drops trailing nulls from every row; keep the file in that shape. + for name, row in merged["results"].items(): + while row and row[-1] is None: + row.pop() + + dest.parent.mkdir(parents=True, exist_ok=True) + asv.util.write_json(dest, merged, api_version=api_version, compact=True) + written.append(dest) + + # The machine description is identical in every shard; carry one over. + machine = paths[0].parent / "machine.json" + if machine.is_file() and not (dest.parent / "machine.json").is_file(): + shutil.copyfile(machine, dest.parent / "machine.json") + + return written + + +def missing_results(written, selected, benchmarks): + """Selected instances that came back with no result in the merged files. + + A shard that dies -- a timeout, an OOM, an asv that exits non-zero -- takes + its share of the suite with it, and the merge of what survives is a + perfectly well-formed results file that is quietly missing benchmarks. That + is the worst outcome available here, because the run looks like it worked. + So the instances asked for are checked against the instances that came back. + + A benchmark that ran and genuinely failed also records no result, and shows + up here too. Both mean the same thing to a caller: this run did not produce + the number. + """ + asv = _asv() + by_name = {b["name"]: b for b in benchmarks} + wanted = {} + for instance in selected: + benchmark_name = instance.split("(", 1)[0] + benchmark = by_name.get(benchmark_name) + if benchmark is None: + continue + names = _instance_names(benchmark) + wanted.setdefault(benchmark_name, []).append((names.index(instance), instance)) + + absent = [] + for path in written: + data = asv.util.load_json(path, api_version=asv.results.Results.api_version) + columns = data.get("result_columns") or [] + if "result" not in columns: + continue + index = columns.index("result") + for benchmark_name, entries in wanted.items(): + row = data.get("results", {}).get(benchmark_name) + values = (row[index] if row and len(row) > index else None) or [] + for param_idx, instance in entries: + if param_idx >= len(values) or values[param_idx] is None: + absent.append(instance) + return absent + + +def _merge_row(a, b, columns): + """Merges two asv result rows for the same benchmark.""" + width = len(columns) + a = list(a) + [None] * (width - len(a)) + b = list(b) + [None] * (width - len(b)) + row = [] + for index, column in enumerate(columns): + x, y = a[index], b[index] + if column in _PER_PARAM_COLUMNS: + row.append(_merge_per_param(x, y)) + elif column == "duration": + row.append(None if x is None and y is None else (x or 0) + (y or 0)) + elif column == "started_at": + present = [v for v in (x, y) if v is not None] + row.append(min(present) if present else None) + else: + # params, version, profile: describe the benchmark, not one run of it. + row.append(x if x is not None else y) + return row + + +def _merge_per_param(x, y): + """Element-wise merge of two parameter-indexed columns, taking whichever ran.""" + if x is None: + return y + if y is None: + return x + width = max(len(x), len(y)) + x = list(x) + [None] * (width - len(x)) + y = list(y) + [None] * (width - len(y)) + return [xi if xi is not None else yi for xi, yi in zip(x, y)] + + +def _core_ranges(n_shards, threads_per_shard): + """``taskset`` core lists giving each shard its own cores, or ``None``. + + ``None`` where the platform has no ``taskset`` or there are not enough cores + to go round, in which case the shards are left to the scheduler -- which + costs nothing here, since the measurements are counts rather than times. + """ + if not shutil.which("taskset"): + return None + available = os.cpu_count() or 1 + if n_shards * threads_per_shard > available: + return None + return [ + ",".join( + str(c) + for c in range(i * threads_per_shard, (i + 1) * threads_per_shard) + ) + for i in range(n_shards) + ] + + +def _run(command, env=None, check=True): + print(f" $ {shlex.join(command)}", flush=True) + return subprocess.run(command, env=env, check=check) + + +def main(argv=None): + parser = argparse.ArgumentParser( + prog="python -m benchmarks.helpers._shard", + description="Run the track_* benchmarks as concurrent asv processes, then the rest serially.", + ) + parser.add_argument( + "range", + nargs="?", + default=None, + help="Commit range, as asv takes it. Defaults to the configured branch head.", + ) + parser.add_argument( + "--workers", type=int, default=4, help="Number of concurrent asv processes (default 4)." + ) + parser.add_argument( + "--threads-per-shard", + type=int, + default=1, + help="numba/OpenMP threads each shard may use. Pinned rather than inherited so a " + "subprocess_peak_rss figure does not depend on how the run was sharded (default 1).", + ) + parser.add_argument( + "--select", + default=r"\.track_", + help=r"Regex picking the benchmarks to run in parallel (default '\.track_').", + ) + parser.add_argument( + "--config", + default=str(BENCHMARK_DIR / "asv.conf.json"), + help="Path to asv.conf.json.", + ) + parser.add_argument( + "--track-only", + action="store_true", + help="Stop after the sharded pass instead of running the serial remainder.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the plan and the commands, run nothing.", + ) + parser.add_argument( + "asv_args", + nargs="*", + default=[], + help="Arguments after a bare -- are passed to every asv run (e.g. --skip-existing).", + ) + args = parser.parse_args(argv) + + asv = _asv() + config_path = os.path.abspath(args.config) + # Match ``asv --config``: paths in the file are relative to the file. + os.chdir(os.path.dirname(config_path)) + conf = asv.config.Config.load(config_path) + raw_conf = asv.util.load_json(config_path, js_comments=True) + results_dir = os.path.abspath(conf.results_dir) + + passthrough = list(args.asv_args) + commit = [args.range] if args.range else [] + + # 1. Fill the fixture cache here rather than letting N shards each find it + # cold. Also keeps it off the forkserver parent's import, which is where + # it otherwise lands -- asv preimports the suite before any setup_cache. + print("priming the fixture cache", flush=True) + if not args.dry_run: + # Not fatal: the shards fill the cache themselves if this cannot run, + # just once per shard instead of once here. Priming needs uxarray + # importable in *this* interpreter, which is the one that has asv -- + # true of the CI environment (``ci/asv.yml`` installs the project into + # it) but not of every environment asv can be driven from. + try: + from . import _fixtures + + read = _fixtures.prime(workers=args.workers) + print(f" built {len(read)} source(s)", flush=True) + except Exception as exc: + print(f" skipped ({type(exc).__name__}: {exc})", flush=True) + + # 2. One serial pass that discovers and, as a side effect, builds and + # installs the project at the commit -- so the shards do not race to. + print("discovering benchmarks (and installing the project)", flush=True) + discover = ["asv", "run", *commit, "--bench", "just-discover", *passthrough] + if not args.dry_run: + _run(discover) + else: + print(f" $ {shlex.join(discover)}", flush=True) + + benchmarks_json = Path(results_dir) / "benchmarks.json" + if not benchmarks_json.is_file(): + raise SystemExit(f"no benchmarks.json at {benchmarks_json}; discovery failed") + discovered = asv.util.load_json(benchmarks_json, api_version=2) + benchmarks = list(discovered.values()) + + pattern = re.compile(args.select) + selected = [ + instance + for benchmark in benchmarks + for instance in _instance_names(benchmark) + if pattern.search(instance) + ] + if not selected: + raise SystemExit(f"nothing matched --select {args.select!r}") + + weights = _weights(results_dir) + shards = plan_shards(benchmarks, selected, args.workers, weights) + print(f"{len(selected)} instances over {len(shards)} shard(s):", flush=True) + for i, shard in enumerate(shards): + load = sum( + weights.get(instance.split("(", 1)[0], 0.0) + / max(1, len(_instance_names(next(b for b in benchmarks if b["name"] == instance.split("(", 1)[0])))) + for instance in shard + ) + print(f" shard {i}: {len(shard):3} instances, ~{load:.1f}s of recorded work", flush=True) + for instance in shard: + print(f" {instance}", flush=True) + + cores = _core_ranges(len(shards), args.threads_per_shard) + if cores is None: + print(" (not pinning cores: no taskset, or too few cores)", flush=True) + + scratch = tempfile.mkdtemp(prefix="asv-shard-") + shard_dirs = [] + processes = [] + incomplete = False + try: + for i, shard in enumerate(shards): + shard_results = os.path.join(scratch, f"shard{i}", "results") + shard_dirs.append(shard_results) + + # Seed the shard with the results already on disk, so that + # ``--skip-existing`` and friends see the same history a plain + # ``asv run`` would. Without this a shard reads an empty results dir + # and re-runs everything. Re-merging the seeded rows afterwards is a + # no-op: they are the values already in the destination. + if os.path.isdir(results_dir): + shutil.copytree(results_dir, shard_results, dirs_exist_ok=True) + shard_conf_path = os.path.join(scratch, f"shard{i}", "asv.conf.json") + os.makedirs(os.path.dirname(shard_conf_path), exist_ok=True) + with open(shard_conf_path, "w") as handle: + json.dump( + _shard_config( + raw_conf, conf, shard_results, os.path.join(scratch, f"shard{i}", "html") + ), + handle, + ) + + # Results are per (machine, commit, env); the shards must agree on + # the machine or they would write to different subdirectories and + # never merge. + command = ["asv", "--config", shard_conf_path, "run", *commit] + for instance in shard: + command += ["--bench", f"^{re.escape(instance)}$"] + command += passthrough + + if cores is not None: + command = ["taskset", "-c", cores[i]] + command + + env = dict(os.environ) + env["NUMBA_NUM_THREADS"] = str(args.threads_per_shard) + env["OMP_NUM_THREADS"] = str(args.threads_per_shard) + env["HOME"] = _shard_home(scratch, i) + + print(f" $ {shlex.join(command)}", flush=True) + if not args.dry_run: + processes.append(subprocess.Popen(command, env=env)) + + codes = [p.wait() for p in processes] + if any(codes): + incomplete = True + print(f"error: shard exit codes {codes}", file=sys.stderr, flush=True) + + # 3. Fold the shards back together, then make sure everything asked + # for actually came back -- see ``missing_results``. + if not args.dry_run: + print("merging shard results", flush=True) + written = merge_results(shard_dirs, results_dir) + for path in written: + print(f" wrote {path}", flush=True) + + absent = missing_results(written, selected, benchmarks) + if absent: + incomplete = True + print( + f"error: {len(absent)} of {len(selected)} selected instances " + "came back with no result:", + file=sys.stderr, + flush=True, + ) + for instance in absent: + print(f" {instance}", file=sys.stderr, flush=True) + else: + print(f" all {len(selected)} instances accounted for", flush=True) + finally: + if not args.dry_run: + shutil.rmtree(scratch, ignore_errors=True) + + # 4. The remainder, serially, on a quiet machine. asv reads the merged file + # back in (Results.load_data) before it records, so this adds to it. + if not args.track_only: + print("running the remaining benchmarks serially", flush=True) + serial = [ + "asv", + "run", + *commit, + "--bench", + f"^(?!.*(?:{args.select}))", + *passthrough, + ] + if not args.dry_run: + _run(serial, check=False) + else: + print(f" $ {shlex.join(serial)}", flush=True) + + # Non-zero rather than a warning buried in the log: a sharded run that lost + # a shard writes a well-formed results file with benchmarks missing from it, + # which is indistinguishable from a clean run unless the caller is told. + return 1 if incomplete else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py new file mode 100644 index 000000000..96a63b3c2 --- /dev/null +++ b/benchmarks/helpers/_warmup.py @@ -0,0 +1,130 @@ +"""Warming benchmark state in the interpreter every benchmark is forked from. + +Under ``launch_method: forkserver`` ASV imports the suite once and forks each +benchmark from that interpreter, so whatever a module prepares at import is +inherited copy-on-write. + +Only tasks that leaves no numba thread pool behind may be warmed this way. +Running a ``parallel=True`` kernel or calling ``.compile()`` on one +launches the pool, and numba's OpenMP layer is not fork-safe. So this checks when a +warmed kernel goes parallel, it becomes an import error rather than a benchmark +that hangs on a cluster. + +Not every interpreter that imports the suite goes on to run something from it, +and the ones that do not should not pay to warm it -- see +:func:`will_run_benchmarks`. +""" + +import os +import sys + +import numba + +__all__ = ["warm_in_parent", "will_run_benchmarks"] + +# ASV drives one script (``asv/benchmark.py``) in several modes, named by its +# first argument. These are the modes that import the suite without ever calling +# anything in it. +_NON_RUNNING_MODES = frozenset({"discover", "setup_cache", "check"}) + + +def will_run_benchmarks(): + """Whether this interpreter is going to execute a benchmark body. + + ``discover`` is the mode this exists for. It imports every module in the + suite to enumerate what is there, calls none of it, and sits on the critical + path of every ``asv run`` -- once for the run itself and once more for each + concurrent shard (:mod:`benchmarks.helpers._shard`). Warming it is pure + cost: measured at ~5s of the ~6s that importing this suite takes, nearly all + of it numba loading kernels that discovery will never call. + + ``setup_cache`` and ``check`` are the same story, though cheaper -- ASV + imports one module for those rather than the tree. + + The modes that do run something (``run``, ``run_server``, ``timing``) still + warm, and so does anything this cannot identify: an unrecognised argv means + warm, which is what happens when the suite is imported from a REPL or a + script. This only skips where it is certain nothing will be measured. + """ + return not ( + os.path.basename(sys.argv[0]) == "benchmark.py" + and len(sys.argv) > 1 + and sys.argv[1] in _NON_RUNNING_MODES + ) + + +def _pool_running(): + """Whether numba's thread pool has been launched in this interpreter. + + ``threading_layer()`` raises until a ``parallel=True`` kernel has been + compiled or run, which is exactly the event that makes this process unsafe + to fork a numba benchmark from. + """ + try: + numba.threading_layer() + except ValueError: + return False + return True + + +_REPORTED_FOREIGN_POOL = False + + +def _report_foreign_pool(what): + """Says once that this interpreter had a thread pool before we warmed anything.""" + global _REPORTED_FOREIGN_POOL + if _REPORTED_FOREIGN_POOL: + return + _REPORTED_FOREIGN_POOL = True + print( + f"asv: numba's thread pool was already running before {what} was warmed. " + "Something in the import of uxarray itself compiled or ran a parallel=True " + "kernel, so this interpreter was already unsafe to fork a numba benchmark " + "from -- warming is not the cause and cannot be the fix. Benchmarking " + "continues. To remove it, stop that kernel being built at import, or set " + "launch_method to spawn.", + file=sys.stderr, + ) + + +def warm_in_parent(warm, what): + """Runs ``warm``, then fails if *warming* started a numba thread pool. + + ``what`` names the thing being warmed, for the error message. + + The pool is sampled before as well as after, and the two cases are not the + same problem. A pool that warming started is ours and is preventable: move + that work into ``setup`` and the parent stays forkable. A pool that was + already running when we got here belongs to something in uxarray's own + import, and refusing to proceed does not help -- the pool exists either way, + asv forks either way, and raising only means the suite cannot benchmark that + commit at all. Which is what it did: every module failed to import against a + uxarray whose import builds a ``parallel=True`` kernel, so every benchmark + on that commit failed, and a run comparing against it compared against + nothing. So that case reports itself and carries on. + + Warming still happens in that case, deliberately. Skipping it would leave + the forks on such a commit compiling kernels that the forks on a fixed + commit inherit, which is a difference between the two sides of a comparison + that has nothing to do with the code being measured. + + A no-op in an interpreter that will not run a benchmark, per + :func:`will_run_benchmarks`. + """ + if not will_run_benchmarks(): + return + + already_running = _pool_running() + warm() + if not _pool_running(): + return # nothing launched a pool, which is what makes this inheritable + + if already_running: + _report_foreign_pool(what) + return + + raise RuntimeError( + f"warming {what} started numba's thread pool, which a forked benchmark " + "cannot safely inherit -- warm it from setup() instead, now that these " + "kernels run in parallel" + ) diff --git a/benchmarks/mpas_dyamond.py b/benchmarks/mpas_dyamond.py index 3e4ecc8c0..8c9454a14 100644 --- a/benchmarks/mpas_dyamond.py +++ b/benchmarks/mpas_dyamond.py @@ -1,31 +1,23 @@ -import os - from asv_runner.benchmarks.mark import skip_benchmark_if, timeout_class_at import uxarray as ux -# Paths to grid files on Glade -grid_path_dict = {"30km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/30km/grid.nc", - "15km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/15km/grid.nc", - "7.5km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/7.5km/grid.nc", - "3.75km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/3.75km/grid.nc"} - - +from .helpers._fixtures import DYAMOND_AVAILABLE, DYAMOND_GRIDS, CachedFixtures -# Determines if all file paths exist and are accesible -all_paths_exist = True -for file_path in grid_path_dict.values(): - all_paths_exist = all_paths_exist and os.path.exists(file_path) +# Paths, and the question of whether this machine can see them, both come from +# ``helpers._fixtures`` -- ``bench_connectivity`` asks for the same four grids. +grid_path_dict = DYAMOND_GRIDS -class BaseGridBenchmark: +class BaseGridBenchmark(CachedFixtures): """Base class for Grid Benchmarks across the four supported resolutions (30km, 15km, 7.5km, 3.75km)""" param_names = ['resolution'] - params = [['30km', '15km', '7.5km', '3.75km'], ] + params = [list(DYAMOND_GRIDS), ] def setup(self, resolution, **kwargs): - self.uxgrid = ux.open_grid(grid_path_dict[resolution]) + # The cached grid, not a fresh read + self.uxgrid = self.cached_grid(grid_path_dict[resolution]) def teardown(self, resolution, **kwargs): del self.uxgrid @@ -33,21 +25,21 @@ def teardown(self, resolution, **kwargs): @timeout_class_at(1200) class OpenGrid: param_names = ['resolution'] - params = [['30km', '15km', '7.5km', '3.75km'], ] + params = [list(DYAMOND_GRIDS), ] - @skip_benchmark_if(not all_paths_exist) + @skip_benchmark_if(not DYAMOND_AVAILABLE) def time_open_grid(self, resolution): _ = ux.open_grid(grid_path_dict[resolution]) @timeout_class_at(1200) class Bounds(BaseGridBenchmark): - @skip_benchmark_if(not all_paths_exist) + @skip_benchmark_if(not DYAMOND_AVAILABLE) def time_bounds(self, resolution): _ = self.uxgrid.bounds @timeout_class_at(1200) class GeoDataFrame(BaseGridBenchmark): - @skip_benchmark_if(not all_paths_exist) + @skip_benchmark_if(not DYAMOND_AVAILABLE) def time_to_geodataframe(self, resolution): self.uxgrid.to_geodataframe(exclude_antimeridian=True) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index adf679a3a..bc3719895 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -8,54 +8,53 @@ import uxarray as ux from uxarray.grid.neighbors import Neighborhood, _get_element_coords +from .helpers._fixtures import ( + OQU_DATASETS, + OQU_GRIDS, + OQU_RESOLUTIONS, + CachedFixtures, + preload_datasets, + preload_grids, +) from .helpers._memsize import grid_nbytes -from .helpers._peakmem import numba_threads, peak_allocated - -current_path = Path(os.path.dirname(os.path.realpath(__file__))) +from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss +from .helpers._warmup import warm_in_parent data_var = 'bottomDepth' -grid_filename_480 = "oQU480.grid.nc" -data_filename_480 = "oQU480.data.nc" - -grid_filename_120 = "oQU120.grid.nc" -data_filename_120 = "oQU120.data.nc" - -filenames = [grid_filename_480, data_filename_480, grid_filename_120, data_filename_120] - -for filename in filenames: - if not os.path.isfile(current_path / filename): - # downloads the files from Cookbook repo, if they haven't been downloaded locally yet - url = f"https://github.com/ProjectPythia/unstructured-grid-viz-cookbook/raw/main/meshfiles/{filename}" - _, headers = urllib.request.urlretrieve(url, filename=current_path / filename) - +# Sample budget for the benchmarks with long single call runtimes. +# +# Only the classes whose slowest parameter clears ~0.25s carry this; on the rest +# the cap would never bind and would cost samples for nothing. +SLOW_CALL_REPEAT = (2, 3, 8.0) -file_path_dict = {"480km": [current_path / grid_filename_480, current_path / data_filename_480], - "120km": [current_path / grid_filename_120, current_path / data_filename_120]} +# Paths, and fetching the files in the first place, both live in +# ``helpers._fixtures`` now -- ``bench_connectivity`` draws the same grids from it. +file_path_dict = OQU_DATASETS - -class DatasetBenchmark: +class DatasetBenchmark(CachedFixtures): """Class used as a template for benchmarks requiring a ``UxDataset`` in - this module across both resolutions.""" + this module across both resolutions. + """ param_names = ['resolution', ] - params = [['480km', '120km'], ] + params = [OQU_RESOLUTIONS, ] def setup(self, resolution, *args, **kwargs): - self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1]) + self.uxds = self.cached_dataset(*file_path_dict[resolution]) def teardown(self, resolution, *args, **kwargs): del self.uxds -class GridBenchmark: +class GridBenchmark(CachedFixtures): """Class used as a template for benchmarks requiring a ``Grid`` in this module across both resolutions.""" param_names = ['resolution', ] - params = [['480km', '120km'], ] + params = [OQU_RESOLUTIONS, ] def setup(self, resolution, *args, **kwargs): - self.uxgrid = ux.open_grid(file_path_dict[resolution][0]) + self.uxgrid = self.cached_grid(file_path_dict[resolution][0]) def teardown(self, resolution, *args, **kwargs): del self.uxgrid @@ -67,9 +66,11 @@ class FaceAreas(GridBenchmark): 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 + _ = self.cached_grid(OQU_GRIDS[OQU_RESOLUTIONS[0]]).face_areas super().setup(resolution, *args, **kwargs) + + # MPAS meshes carry ``face_areas`` on disk and computation requires it to + # be dropped. Still safe, because each fixture ``Grid`` is a shallow copy. self.uxgrid._ds = self.uxgrid._ds.drop_vars("face_areas", errors="ignore") def time_face_areas(self, resolution): @@ -93,8 +94,7 @@ 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() + _ = self.cached_dataset(*file_path_dict[OQU_RESOLUTIONS[0]])[data_var].gradient() def time_gradient(self, resolution): self.uxds[data_var].gradient() @@ -106,8 +106,12 @@ 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()) + """Transient high-water allocation of taking a gradient. + + The kernel behind ``gradient`` is ``parallel=True``, hence the pinning + """ + with numba_threads(1): + return peak_allocated(lambda: self.uxds[data_var].gradient()) track_peakmem_gradient.unit = "bytes" @@ -128,30 +132,45 @@ def track_nbytes_integrate(self, resolution): 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. + Whole-process peak resident memory, not tracemalloc -- the ~226MB 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. + + Measured in a subprocess of its own rather than through asv's ``peakmem_*``, + which reports ``ru_maxrss`` for the benchmark process. """ param_names = ["resolution"] - params = [["480km", "120km"]] + params = [OQU_RESOLUTIONS] def setup_cache(self): - """Compile the njit kernels before anything is measured.""" + """Compile the njit kernels before anything is measured. + + The subprocess inherits numba's on-disk cache rather than this process's + memory, so this keeps compilation out of the measured cold start. + """ 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): + def track_peakmem_gradient(self, resolution): grid, data = file_path_dict[resolution] - ux.open_dataset(grid, data)[data_var].gradient() + return subprocess_peak_rss( + "import uxarray as ux\n" + f"uxds = ux.open_dataset({str(grid)!r}, {str(data)!r})\n" + f"uxds[{data_var!r}].gradient()\n" + ) + + track_peakmem_gradient.unit = "bytes" class GeoDataFrame(DatasetBenchmark): param_names = DatasetBenchmark.param_names + ['exclude_antimeridian'] + repeat = SLOW_CALL_REPEAT params = DatasetBenchmark.params + [[True, False]] def time_to_geodataframe(self, resolution, exclude_antimeridian): @@ -183,11 +202,11 @@ def time_ball_tree(self, resolution): self.uxds.uxgrid.get_ball_tree() -class RemapDownsample: +class RemapDownsample(CachedFixtures): def setup(self): - self.uxds_120 = ux.open_dataset(file_path_dict['120km'][0], file_path_dict['120km'][1]) - self.uxds_480 = ux.open_dataset(file_path_dict['480km'][0], file_path_dict['480km'][1]) + self.uxds_120 = self.cached_dataset(*file_path_dict['120km']) + self.uxds_480 = self.cached_dataset(*file_path_dict['480km']) def teardown(self): del self.uxds_120, self.uxds_480 @@ -201,11 +220,12 @@ def time_inverse_distance_weighted_remapping(self): def time_bilinear_remapping(self): self.uxds_120["bottomDepth"].remap.bilinear(self.uxds_480.uxgrid) -class RemapUpsample: +class RemapUpsample(CachedFixtures): + repeat = SLOW_CALL_REPEAT def setup(self): - self.uxds_120 = ux.open_dataset(file_path_dict['120km'][0], file_path_dict['120km'][1]) - self.uxds_480 = ux.open_dataset(file_path_dict['480km'][0], file_path_dict['480km'][1]) + self.uxds_120 = self.cached_dataset(*file_path_dict['120km']) + self.uxds_480 = self.cached_dataset(*file_path_dict['480km']) def teardown(self): del self.uxds_120, self.uxds_480 @@ -231,6 +251,8 @@ def time_dual_mesh_construction(self, resolution): class ConstructFaceLatLon(GridBenchmark): + repeat = SLOW_CALL_REPEAT + def time_welzl(self, resolution): self.uxgrid.construct_face_centers(method='welzl') @@ -238,12 +260,12 @@ def time_cartesian_averaging(self, resolution): self.uxgrid.construct_face_centers(method='cartesian average') -class CheckNorm: +class CheckNorm(CachedFixtures): param_names = ['resolution'] - params = ['480km', '120km'] + params = OQU_RESOLUTIONS def setup(self, resolution): - self.uxgrid = ux.open_grid(file_path_dict[resolution][0]) + self.uxgrid = self.cached_grid(file_path_dict[resolution][0]) def teardown(self, resolution): del self.uxgrid @@ -254,10 +276,11 @@ def time_check_norm(self, resolution): class CrossSections(DatasetBenchmark): param_names = DatasetBenchmark.param_names + ['n_lat'] + repeat = SLOW_CALL_REPEAT params = DatasetBenchmark.params + [[1, 2, 4]] def setup(self, resolution, lat_step): - self.uxgrid = ux.open_grid(file_path_dict[resolution][0]) + self.uxgrid = self.cached_grid(file_path_dict[resolution][0]) self.uxgrid.normalize_cartesian_coordinates() self.lats = np.arange(-45, 45, lat_step) _ = self.uxgrid.bounds @@ -270,12 +293,12 @@ def time_const_lat(self, resolution, lat_step): self.uxgrid.cross_section.constant_latitude(lat) -class PointInPolygon: +class PointInPolygon(CachedFixtures): param_names = ['resolution'] - params = ['480km', '120km'] + params = OQU_RESOLUTIONS def setup(self, resolution): - self.uxgrid = ux.open_grid(file_path_dict[resolution][0]) + self.uxgrid = self.cached_grid(file_path_dict[resolution][0]) self.uxgrid.normalize_cartesian_coordinates() # Construct variables needed to ensure that the benchmark doesn't measure construction time @@ -301,7 +324,7 @@ def time_face_search_lonlat(self, resolution): class ZonalAverage(DatasetBenchmark): def setup(self, resolution, *args, **kwargs): - self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1]) + super().setup(resolution, *args, **kwargs) bounds = self.uxds.uxgrid.bounds def time_zonal_average(self, resolution): @@ -310,10 +333,14 @@ def time_zonal_average(self, resolution): class ZonalAveragePeakMem: - """Peak memory of a cold-start non-conservative zonal-mean sweep.""" + """Peak memory of a cold-start non-conservative zonal-mean sweep. + + A fresh interpreter per sample. The cold start is the subject, and a forked + benchmark process no longer has one. + """ param_names = ["resolution"] - params = [["480km", "120km"]] + params = [OQU_RESOLUTIONS] def setup_cache(self): """Compile the njit kernels before anything is measured.""" @@ -323,18 +350,28 @@ def setup_cache(self): uxds.uxgrid.bounds uxds[data_var].zonal_mean(lat=(-45, 45, 10)) - def peakmem_zonal_average(self, resolution): + setup_cache.timeout = 1800 + + def track_peakmem_zonal_average(self, resolution): grid, data = file_path_dict[resolution] - uxds = ux.open_dataset(grid, data) - uxds.uxgrid.bounds - uxds[data_var].zonal_mean(lat=(-45, 45, 10)) + return subprocess_peak_rss( + "import uxarray as ux\n" + f"uxds = ux.open_dataset({str(grid)!r}, {str(data)!r})\n" + "uxds.uxgrid.bounds\n" + f"uxds[{data_var!r}].zonal_mean(lat=(-45, 45, 10))\n" + ) + + track_peakmem_zonal_average.unit = "bytes" class CrossSectionsPeakMem: - """Peak memory of a cold-start constant-latitude cross-section sweep.""" + """Peak memory of a cold-start constant-latitude cross-section sweep. + + A fresh interpreter per sample, + """ param_names = ["resolution", "lat_step"] - params = [["480km", "120km"], [1, 2, 4]] + params = [OQU_RESOLUTIONS, [1, 2, 4]] def setup_cache(self): """Compile the njit kernels before anything is measured.""" @@ -344,12 +381,20 @@ def setup_cache(self): uxgrid.bounds uxgrid.cross_section.constant_latitude(0.0) - def peakmem_const_lat(self, resolution, lat_step): - uxgrid = ux.open_grid(file_path_dict[resolution][0]) - uxgrid.normalize_cartesian_coordinates() - uxgrid.bounds - for lat in np.arange(-45, 45, lat_step): - uxgrid.cross_section.constant_latitude(lat) + setup_cache.timeout = 1800 + + def track_peakmem_const_lat(self, resolution, lat_step): + grid = file_path_dict[resolution][0] + return subprocess_peak_rss( + "import numpy as np, uxarray as ux\n" + f"uxgrid = ux.open_grid({str(grid)!r})\n" + "uxgrid.normalize_cartesian_coordinates()\n" + "uxgrid.bounds\n" + f"for lat in np.arange(-45, 45, {lat_step}):\n" + " uxgrid.cross_section.constant_latitude(lat)\n" + ) + + track_peakmem_const_lat.unit = "bytes" class NeighborhoodBuild(DatasetBenchmark): @@ -525,3 +570,16 @@ def track_peakmem_mean(self, resolution, chunking): return peak_allocated(lambda: self.nb.mean().compute()) track_peakmem_mean.unit = "bytes" + + +# Read at import rather than in ``setup``. ASV imports the suite once and forks +# every benchmark from that parent, so a fixture loaded here is inherited by all +# of them copy-on-write instead of being re-read by each -- and, more to the +# point, none of them has to bring up the netCDF stack for itself. Only the two +# oQU meshes, which come to ~22MB; the dyamond grids stay per-fork reads. +def _warm_parent(): + preload_grids(OQU_GRIDS.values()) + preload_datasets(OQU_DATASETS.values()) + + +warm_in_parent(_warm_parent, "the oQU fixtures") diff --git a/benchmarks/quad_hexagon.py b/benchmarks/quad_hexagon.py index 5b03eec6f..a3c5b95ef 100644 --- a/benchmarks/quad_hexagon.py +++ b/benchmarks/quad_hexagon.py @@ -1,14 +1,12 @@ -import os -from pathlib import Path - import uxarray as ux + +from .helpers._fixtures import QUAD_HEXAGON_DATASET 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] - -grid_path = current_path / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "grid.nc" -data_path = current_path / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "data.nc" +# Opening these files is what this module measures, so it reads them for real +# every time; only the paths come from the shared registry. +grid_path, data_path = QUAD_HEXAGON_DATASET class QuadHexagon: diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 03a4c7cd5..1ae97f085 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1284,17 +1284,35 @@ def _median(window, _): # names them -- the data-bound classes reach a kernel by naming the # ``Neighborhood`` method for it, so there is one place per reduction where its # kernel and parameter are chosen. -_MEAN_KERNEL = _make_kernel(lambda window, _: np.mean(window)) -_SUM_KERNEL = _make_kernel(lambda window, _: np.sum(window)) -_MIN_KERNEL = _make_kernel(lambda window, _: np.min(window)) -_MAX_KERNEL = _make_kernel(lambda window, _: np.max(window)) -_PTP_KERNEL = _make_kernel(lambda window, _: np.max(window) - np.min(window)) -_MEDIAN_KERNEL = _make_kernel(_median) -_VAR_KERNEL = _make_kernel(_variance) -_STD_KERNEL = _make_kernel(lambda window, ddof: np.sqrt(_variance(window, ddof))) -# ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale onto -# this one kernel rather than compiling a near-duplicate. -_QUANTILE_KERNEL = _make_kernel(lambda window, q: np.quantile(window, q)) +# +# Built on first use rather than at import. ``guvectorize`` with explicit +# signatures compiles at decoration time, and building a ``target="parallel"`` +# ufunc launches numba's OpenMP thread pool as part of that. Assigning +# these at module scope made ``import uxarray`` start a thread pool before any +# user code ran. + +_REDUCERS = { + "mean": lambda window, _: np.mean(window), + "sum": lambda window, _: np.sum(window), + "min": lambda window, _: np.min(window), + "max": lambda window, _: np.max(window), + "ptp": lambda window, _: np.max(window) - np.min(window), + "median": _median, + "var": _variance, + "std": lambda window, ddof: np.sqrt(_variance(window, ddof)), + # ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale + # onto this one kernel rather than compiling a near-duplicate. + "quantile": lambda window, q: np.quantile(window, q), +} + +_KERNELS = {} + + +def _kernel(reduction): + """The compiled kernel for ``reduction``, built once on first use.""" + if reduction not in _KERNELS: + _KERNELS[reduction] = _make_kernel(_REDUCERS[reduction]) + return _KERNELS[reduction] def _as_quantile(q, scale: float): @@ -1509,45 +1527,45 @@ def __repr__(self) -> str: def mean(self, uxda): """Mean of each neighborhood.""" - return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("mean"), 0.0) def sum(self, uxda): """Sum of each neighborhood.""" - return self._apply_kernel(uxda, _SUM_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("sum"), 0.0) def min(self, uxda): """Smallest value in each neighborhood.""" - return self._apply_kernel(uxda, _MIN_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("min"), 0.0) def max(self, uxda): """Largest value in each neighborhood.""" - return self._apply_kernel(uxda, _MAX_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("max"), 0.0) def ptp(self, uxda): """Peak-to-peak spread (``max - min``) of each neighborhood.""" - return self._apply_kernel(uxda, _PTP_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("ptp"), 0.0) def median(self, uxda): """Median of each neighborhood.""" - return self._apply_kernel(uxda, _MEDIAN_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("median"), 0.0) def var(self, uxda, ddof: int = 0): """Variance of each neighborhood, with ``ddof`` delta degrees of freedom.""" - return self._apply_kernel(uxda, _VAR_KERNEL, float(ddof)) + return self._apply_kernel(uxda, _kernel("var"), float(ddof)) def std(self, uxda, ddof: int = 0): """Standard deviation of each neighborhood, with ``ddof`` delta degrees of freedom.""" - return self._apply_kernel(uxda, _STD_KERNEL, float(ddof)) + return self._apply_kernel(uxda, _kernel("std"), float(ddof)) def quantile(self, uxda, q: float): """Quantile ``q`` (between 0 and 1) of each neighborhood.""" - return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 1.0)) + return self._apply_kernel(uxda, _kernel("quantile"), _as_quantile(q, 1.0)) def percentile(self, uxda, q: float): """Percentile ``q`` (between 0 and 100) of each neighborhood.""" - return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 100.0)) + return self._apply_kernel(uxda, _kernel("quantile"), _as_quantile(q, 100.0)) def reduce(self, uxda, func: Callable): """Reduces each neighborhood with an arbitrary callable.