diff --git a/.github/workflows/asv-benchmarking-pr.yml b/.github/workflows/asv-benchmarking-pr.yml index cae4aa6f0..66eca2401 100644 --- a/.github/workflows/asv-benchmarking-pr.yml +++ b/.github/workflows/asv-benchmarking-pr.yml @@ -31,6 +31,11 @@ jobs: fetch-depth: 0 + - name: Record CPU topology + run: | + # A diagnostic; never fail the job over it. + lscpu | grep -E 'Model name|^CPU\(s\):|Thread\(s\) per core|Core\(s\) per socket|Socket\(s\)|CPU max MHz' || lscpu || true + - name: Set up Conda environment uses: mamba-org/setup-micromamba@v3 with: @@ -48,6 +53,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..5c52cd586 100644 --- a/.gitignore +++ b/.gitignore @@ -163,3 +163,4 @@ docs/user-guide/psi_healpix.nc benchmarks/env benchmarks/results benchmarks/html +benchmarks/_io_cache diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 31a43921d..9b3852f53 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -61,7 +61,18 @@ // defaults to 10 min "install_timeout": 600, - "benchmark_timeout": 360, + // Fork each benchmark from one interpreter that has already imported the + // suite, instead of starting a fresh one per benchmark. Saves the ~0.9s + // uxarray import and the numba kernel load on every one of ~160 benchmark + // processes, and lets a fixture loaded at import be inherited rather than + // re-read. asv leaves this at "spawn" by default because fork and threads + // mix badly; nothing in this suite runs a parallel kernel at import time, + // which is what makes it safe here -- keep it that way. + "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/", @@ -91,10 +102,27 @@ // will not be set for the current combination. // "matrix": { - "setuptools_scm": [""], - "xarray": [""], - "netcdf4": [""], - "pip+pyfma": [""] + "req": { + "setuptools_scm": [""], + "xarray": [""], + "netcdf4": [""], + "pip+pyfma": [""], + + // numba picks its threading layer by availability, and only ``tbb`` + // and ``workqueue`` install fork handlers. Left to itself on a runner + // without TBB it picks ``omp``, which on Linux is libgomp: a benchmark + // forked from a parent that holds a pool then dies with "Terminating: + // fork() called from a process already using GNU OpenMP". numba + // *prefers* TBB, so a developer machine that has it cannot reproduce + // that -- hence pinning it here rather than inheriting whatever the + // environment happens to offer. + "tbb": [""] + }, + + // Belt to that brace: if TBB is ever unavailable, pick the other + // fork-safe layer rather than quietly falling back to the one that + // breaks. ``forksafe`` raises if no fork-safe layer exists at all. + "env_nobuild": {"NUMBA_THREADING_LAYER": ["forksafe"]} }, @@ -165,7 +193,12 @@ // `asv` will cache results of the recent builds in each // environment, making them faster to install next time. This is // the number of builds to keep, per environment. - // "build_cache_size": 2, + // + // Raised from the default of 2, which is exactly the number of commits + // ``asv continuous`` builds and so leaves no headroom: with rounds + // interleaved it revisits each commit, and a wheel that has been evicted + // in between is built again from scratch. + "build_cache_size": 4, // The commits after which the regression search in `asv publish` // should start looking for regressions. Dictionary whose keys are 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..f5d58ea47 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -1,18 +1,15 @@ -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 +from .helpers._memsize import grid_nbytes +from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss -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 +21,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 +60,30 @@ 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" 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..9ab574197 --- /dev/null +++ b/benchmarks/helpers/_fixtures.py @@ -0,0 +1,374 @@ +"""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 + +__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_topologies", + "prime", +] + +from ._warmup import warm_in_parent + +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``.""" + return xr.open_dataset(artifact_path, mask_and_scale=False).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 + + _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. + + The spec is handed back rather than left to default. ``Grid.__init__`` records + it in ``_ds.attrs``, and the artifact is that ``_ds``, so the original label + survives the round trip -- an MPAS grid stays MPAS. Omitting it warns on every + construction ("Attempting to construct a Grid without passing in + source_grid_spec"), once per benchmark process, which buried the actual + benchmark output. + """ + grid_ds = _cached_grid_ds(grid_path).copy() + spec = grid_ds.attrs.get("source_grid_spec", "UGRID") # uxarray's documented default + return ux.Grid(grid_ds, source_grid_spec=spec) + + +def cached_dataset(grid_path, data_path): + """A ``UxDataset`` over ``data_path``, on the cached grid.""" + source = (Path(grid_path), Path(data_path)) + artifact_path = _ensure(source, "data", ".nc") + if artifact_path not in _loaded: + _loaded[artifact_path] = _read_dataset(artifact_path) + return ux.UxDataset(_loaded[artifact_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 + + +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) + + +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) + + +def warm_netcdf(): + """Brings the netCDF/HDF5 stack up here, so the forks inherit it loaded. + + The first ``open_dataset`` in a process spends ~65ms initializing the + library before it reads a byte, and nothing else in the parent pays that: + the topology fixtures go through numpy. So today every forked child that + touches a cached ``.nc`` artifact pays it again, one at a time. + + The file is closed again. What the children want is the loaded library, not + a handle -- an inherited HDF5 handle is shared rather than copied, and + parent and child would then read through one file offset. + """ + artifacts = sorted(cache_dir().glob("*.nc"), key=lambda path: path.stat().st_size) + if not artifacts: + return # cache not primed yet, so the first child pays it as before + with xr.open_dataset(artifacts[0], mask_and_scale=False): + pass + + +warm_in_parent(warm_netcdf, "the netCDF backend") 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/_warmup.py b/benchmarks/helpers/_warmup.py new file mode 100644 index 000000000..ae84c10a9 --- /dev/null +++ b/benchmarks/helpers/_warmup.py @@ -0,0 +1,89 @@ +"""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. + +Warming can leave a numba thread pool behind. *Compiling* a ``parallel=True`` +kernel initializes the threading layer -- before the kernel is ever run -- and +``cache=True`` only avoids that while the on-disk cache is warm, which it is +not on a fresh runner. Whether an inherited pool is safe then depends entirely +on which layer numba picked: + +* ``tbb`` and ``workqueue`` install fork handlers and come through it intact. +* ``omp`` does not. Where that is libgomp -- Linux, unless TBB is installed -- + every forked child that touches a parallel kernel is killed with + ``Terminating: fork() called from a process already using GNU OpenMP``. + +numba prefers TBB, so a machine that has it cannot reproduce the failure at +all; the benchmark environment pins it instead (see ``tbb`` and +``NUMBA_THREADING_LAYER`` in ``asv.conf.json``). This reports when that has not +taken effect rather than raising. Raising fails the module's import in the +parent *and* again in every child that re-imports it, so it converts "these +benchmarks are slower than they could be" into "this commit produced no +results" -- and the pool exists either way, and asv forks either way. +""" + +import os +import sys + +import numba + +__all__ = ["warm_in_parent", "will_run_benchmarks"] + +# ``benchmark.py ``: asv runs discovery, setup_cache and check as their own +# processes, and none of them goes on to run a benchmark, so warming in them is +# time bought for nobody. ``run_server`` -- the forkserver parent, the one whose +# state every benchmark inherits -- is deliberately absent. +_NON_RUNNING_MODES = frozenset({"discover", "setup_cache", "check"}) + +# Numba installs fork handlers for these two; ``omp`` is on its own. +_FORK_SAFE = frozenset({"tbb", "workqueue"}) + +_reported = False + + +def will_run_benchmarks(): + """Whether this interpreter is going to run a benchmark. + + Conservative by construction: an argv this does not recognize is assumed to + be a benchmark run, so a change in how asv invokes itself costs the warm + twice over rather than losing it where it counts. + """ + return not ( + os.path.basename(sys.argv[0]) == "benchmark.py" + and len(sys.argv) > 1 + and sys.argv[1] in _NON_RUNNING_MODES + ) + + +def warm_in_parent(warm, what): + """Runs ``warm``, then checks any pool it leaves behind survives a fork. + + ``what`` names the thing being warmed, for the report. A no-op in an + interpreter that will not run a benchmark, per :func:`will_run_benchmarks`. + """ + if not will_run_benchmarks(): + return + + warm() + + try: + layer = numba.threading_layer() + except ValueError: + return # nothing initialized a pool, so there is nothing to inherit + + if layer in _FORK_SAFE: + return + + global _reported + if _reported: + return + _reported = True + print( + f"asv: warming {what} left numba's {layer!r} thread pool behind, and " + f"{layer!r} does not survive fork(). Forked benchmarks that run a " + "parallel kernel will be killed by the OpenMP runtime. Install tbb in " + "the benchmark environment, or set NUMBA_THREADING_LAYER=forksafe.", + file=sys.stderr, + ) 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..2136b2612 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -8,54 +8,50 @@ import uxarray as ux from uxarray.grid.neighbors import Neighborhood, _get_element_coords +from .helpers._fixtures import ( + OQU_DATASETS, + OQU_GRIDS, + OQU_RESOLUTIONS, + CachedFixtures, +) 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 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) +# 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 -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]} - - -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 +63,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 +91,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 +103,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 +129,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 +199,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 +217,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 +248,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 +257,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 +273,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 +290,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 +321,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 +330,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 +347,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 +378,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): 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/test/test_dependencies.py b/test/test_dependencies.py index fb27bfbce..2be243d01 100644 --- a/test/test_dependencies.py +++ b/test/test_dependencies.py @@ -38,4 +38,34 @@ def test_hvplot_optional(): _assert_not_imported_after_import_uxarray("hvplot") +def test_no_numba_kernels_built_on_import(): + """Test that `import uxarray` does not build any numba kernel. + + ``guvectorize`` compiles at decoration time when it is given explicit + signatures, so a kernel assigned at module scope is built during the + import. This compilation can dominate the uxarray import, and building a + ``target="parallel"`` kernel starts numba's threading layer, which + leaves a thread pool running, making forks unsafe. + """ + code = ( + "import numba, uxarray\n" + "try:\n" + " layer = numba.threading_layer()\n" + "except ValueError:\n" + " pass\n" + "else:\n" + " raise AssertionError(\n" + " f'`import uxarray` started numba threading layer {layer!r}. '\n" + " 'Something it imports builds a parallel kernel at module '\n" + " 'scope; build it on first use instead.'\n" + " )\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + # TODO: similar tests for cartopy, holoviews, and other optional deps. diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 03a4c7cd5..753092f18 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,3 +1,4 @@ +import functools import warnings from typing import Callable @@ -1277,26 +1278,6 @@ def _median(window, _): return np.median(window) -# One compiled kernel per reduction. The methods on ``Neighborhood`` below name -# these directly, so there is no dispatch table between the public API and the -# gufuncs: a reduction is reachable only if a method exists for it, and a method -# can only reach the kernel it names. ``Neighborhood`` is the only class that -# 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)) - - def _as_quantile(q, scale: float): """Validates ``q`` on a 0-``scale`` scale and returns it as a 0-1 fraction.""" value = float(q) @@ -1507,47 +1488,107 @@ def __repr__(self) -> str: f"neighbors_per_element=[{self._counts.min()}, {self._counts.max()}]>" ) + # One compiled kernel per reduction. The methods below call into these + # directly. Non-compiled functions are only provided hooks through + # ``reduce``. If new compiled reductions are desired, they should follow + # this pattern. + # + # ``functools.cache`` defers each build to the first call. The deferred + # compilation ensures that these kernels will only be compiled individually + # and lazily. Further, the lazy compilation prevents gufuncs from spawning + # threadpools eagerly and disrupting threading and forking in other + # contexts. They are ``staticmethod``s rather than attributes for the same + # reason: a class body runs at import, so assigning them there would + # compile all nine during ``import uxarray``. + + @staticmethod + @functools.cache + def _mean_kernel(): + return _make_kernel(lambda window, _: np.mean(window)) + + @staticmethod + @functools.cache + def _sum_kernel(): + return _make_kernel(lambda window, _: np.sum(window)) + + @staticmethod + @functools.cache + def _min_kernel(): + return _make_kernel(lambda window, _: np.min(window)) + + @staticmethod + @functools.cache + def _max_kernel(): + return _make_kernel(lambda window, _: np.max(window)) + + @staticmethod + @functools.cache + def _ptp_kernel(): + return _make_kernel(lambda window, _: np.max(window) - np.min(window)) + + @staticmethod + @functools.cache + def _median_kernel(): + return _make_kernel(_median) + + @staticmethod + @functools.cache + def _var_kernel(): + return _make_kernel(_variance) + + @staticmethod + @functools.cache + def _std_kernel(): + return _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. + @staticmethod + @functools.cache + def _quantile_kernel(): + return _make_kernel(lambda window, q: np.quantile(window, q)) + def mean(self, uxda): """Mean of each neighborhood.""" - return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0) + return self._apply_kernel(uxda, self._mean_kernel, 0.0) def sum(self, uxda): """Sum of each neighborhood.""" - return self._apply_kernel(uxda, _SUM_KERNEL, 0.0) + return self._apply_kernel(uxda, self._sum_kernel, 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, self._min_kernel, 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, self._max_kernel, 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, self._ptp_kernel, 0.0) def median(self, uxda): """Median of each neighborhood.""" - return self._apply_kernel(uxda, _MEDIAN_KERNEL, 0.0) + return self._apply_kernel(uxda, self._median_kernel, 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, self._var_kernel, 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, self._std_kernel, 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, self._quantile_kernel, _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, self._quantile_kernel, _as_quantile(q, 100.0)) def reduce(self, uxda, func: Callable): """Reduces each neighborhood with an arbitrary callable. @@ -1592,7 +1633,7 @@ def run(block, arrays): # path does too by writing into a float64 output. if block.dtype not in (np.float64, np.float32): block = block.astype(np.float64) - return kernel(block, *arrays, param) + return kernel()(block, *arrays, param) return self._apply(uxda, run)