From 1c7665d9a5b417aac4ef9e3f6cd541a5595a35a9 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 12:07:04 -0500 Subject: [PATCH 01/18] Cached benchmark IO sketch --- benchmarks/asv.conf.json | 9 + benchmarks/bench_connectivity.py | 82 +++----- benchmarks/face_bounds.py | 55 +++-- benchmarks/helpers/_fixtures.py | 334 +++++++++++++++++++++++++++++++ benchmarks/mpas_dyamond.py | 34 ++-- benchmarks/mpas_ocean.py | 179 ++++++++++------- benchmarks/quad_hexagon.py | 12 +- 7 files changed, 533 insertions(+), 172 deletions(-) create mode 100644 benchmarks/helpers/_fixtures.py diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 31a43921d..41ee30072 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -61,6 +61,15 @@ // defaults to 10 min "install_timeout": 600, + // 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": 360, // the base URL to show a commit for the project. diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 5e3f53d02..3a5859ee9 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -1,53 +1,27 @@ -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"} +from .helpers._fixtures import ( + ALL_RESOLUTIONS, + GRIDS_BY_RESOLUTION, + CachedFixtures, + cached_topology, +) -# 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 - - -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'], ] + params = [ALL_RESOLUTIONS, ] + + # A single connectivity build at 3.75km does not fit in the 360s default + # from ``asv.conf.json``. + 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,16 +39,23 @@ 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. + + Warmed on the coarsest grid in ``params``, because resolution decides how + long the kernels run but not which signatures compile. Warming at the + benchmark's own resolution instead means eight full connectivity builds + before the sample -- 0.493s against a 0.106s sample at 120km, and the gap + only widens from there. """ 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 @@ -89,16 +70,11 @@ class Connectivity(GridBenchmark): 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]) + + _warmup() self.uxgrid = self.minimal_grid() def minimal_grid(self): diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index 8f1e41416..d406723d3 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -1,18 +1,17 @@ -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" +# One grid per reader, from the shared registry. ``mpas`` here is the same mesh +# the oQU ``480km`` benchmarks use, through the copy in the repo. +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 +23,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 +62,37 @@ 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. Under + ``launch_method: forkserver`` that process is forked from an interpreter that + has already imported the suite, so a ``peakmem_*`` here would be reporting a + warm start plus whatever the parent was holding. A fresh interpreter is the + only way to keep measuring what this benchmark is named for. """ 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. + + The subprocess inherits numba's on-disk cache, not this process's + memory, so this keeps compilation out of the measured cold start. + """ 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/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py new file mode 100644 index 000000000..a39975c14 --- /dev/null +++ b/benchmarks/helpers/_fixtures.py @@ -0,0 +1,334 @@ +"""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 -- around 160 times across +the suite once the dyamond resolutions are visible. For a mesh on local disk +that is a rounding error. For one on campaign storage it is most of the run, and +it is spent inside the benchmark's own timeout. + +Benchmarks that do measure opening a file keep reading the real thing: +``quad_hexagon``, ``OpenGrid`` in ``mpas_dyamond``, ``import``, and the +cold-start peak-memory benchmarks. They take their paths from the registries +here too, so the suite declares its inputs in one place either way. + +Two flavours, picked per benchmark: + +``topology`` + the three arrays ``Grid.from_topology`` needs and nothing else, for + benchmarks that mean to build the rest themselves -- 2.3MB of the 102MB + 120km MPAS file. +``grid`` / ``dataset`` + everything the reader produced, so a benchmark still gets the + ``face_areas`` and connectivity variables an MPAS file carries on disk + rather than silently measuring their construction. + +One source read produces both, so choosing between them costs nothing. + +Artifacts are keyed on the uxarray build as well as on the file, because an +artifact is one version's reader output and asv walks commits. That means a +fresh read per commit; ``prime`` therefore leaves the dyamond grids out unless +asked, 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 os +import tempfile +import urllib.request +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", + "prime", +] + +BENCHMARK_DIR = Path(__file__).resolve().parents[1] +REPO_DIR = BENCHMARK_DIR.parent + +_COOKBOOK_URL = ( + "https://github.com/ProjectPythia/unstructured-grid-viz-cookbook/raw/main/meshfiles" +) + + +def _cookbook(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_URL}/{filename}", filename=path) + return path + + +# Grids, and grid/data pairs, by mesh resolution. +OQU_GRIDS = { + "480km": _cookbook("oQU480.grid.nc"), + "120km": _cookbook("oQU120.grid.nc"), +} +OQU_DATASETS = { + "480km": (OQU_GRIDS["480km"], _cookbook("oQU480.data.nc")), + "120km": (OQU_GRIDS["120km"], _cookbook("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"), +} + +# Asked once here, rather than separately in each module that cares. +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 + +# Two ladders rather than one, because which of them a benchmark belongs on is a +# per-benchmark decision: an algorithm that does not care which model wrote the +# mesh can take the wide one, while anything tied to the oQU pair -- or too slow +# to run four dyamond resolutions of -- stays on the narrow one. +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", +) + +CACHE_DIR_VAR = "UXARRAY_BENCH_CACHE_DIR" +PRIME_VAR = "UXARRAY_BENCH_PRIME" + +# The arguments ``ux.Grid.from_topology`` takes, in order. +_TOPOLOGY_ARRAYS = ("node_lon", "node_lat", "face_node_connectivity") + +# Artifact path -> what it holds, for this process. +_loaded = {} + + +def cache_dir(): + """Directory the cached artifacts live in. + + ``UXARRAY_BENCH_CACHE_DIR`` overrides the default, and on a cluster it + should: the cache only pays off on a filesystem faster than the one holding + the source grids, and putting it somewhere that outlives the job means each + grid is read once per machine rather than once per job. + """ + root = Path(os.environ.get(CACHE_DIR_VAR) or tempfile.gettempdir()) + cached = root / "uxarray-bench-fixtures" + cached.mkdir(parents=True, exist_ok=True) + return cached + + +def _artifact(source, flavour, suffix): + """Where ``flavour`` 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, and on the uxarray version, because the artifact is + that version's reader output. + """ + parts = [ux.__version__] + 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}-{flavour}-{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, never a half-written file. + """ + scratch = artifact_path.with_name( + f"{artifact_path.stem}.{os.getpid()}.tmp{artifact_path.suffix}" + ) + writer(dataset, scratch) + os.replace(scratch, artifact_path) + + +def _read_dataset(artifact_path): + """Reads back a cached ``xr.Dataset``. + + ``mask_and_scale=False`` is load-bearing: the connectivity variables carry + ``_FillValue``, which xarray would otherwise consume, handing back float64 + where uxarray's njit kernels require int64 -- they fail to type rather than + returning something wrong, but they do fail. + """ + return xr.open_dataset(artifact_path, mask_and_scale=False).load() + + +def _build(source): + """Reads ``source`` and writes every flavour 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`` is an ``xr.Dataset`` subclass, so it writes itself; the + # grid half is cached separately just below. + data_ds = uxds + + _write( + {name: getattr(uxgrid, name).data for name in _TOPOLOGY_ARRAYS}, + _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, flavour, suffix): + """Path to a cached artifact, building the source's artifacts if need be.""" + artifact_path = _artifact(source, flavour, suffix) + if not artifact_path.exists(): + _build(source if flavour == "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 _TOPOLOGY_ARRAYS) + 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 ``Grid`` carrying everything the reader found in ``grid_path``. + + A fresh ``Grid`` over a shallow copy each call, so a benchmark that + populates or normalizes something does not hand its leftovers to the next + repeat: uxarray assigns new variables into ``_ds`` rather than writing + through the arrays, so the copy isolates that while the data stays shared. + """ + return ux.Grid(_cached_grid_ds(grid_path).copy()) + + +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(include_dyamond=None): + """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. + + The dyamond grids are left out unless ``UXARRAY_BENCH_PRIME=all`` (or + ``include_dyamond``) asks for them, because a filtered run should not pay + for reading four grids off campaign storage that it will never touch. On a + machine that does have them, prime from the CLI before ``asv run``. + """ + if include_dyamond is None: + include_dyamond = os.environ.get(PRIME_VAR, "").lower() == "all" + + sources = [(path,) for path in OQU_GRIDS.values()] + sources += [(path,) for path in GRIDS_BY_FORMAT.values()] + sources += list(OQU_DATASETS.values()) + if include_dyamond and DYAMOND_AVAILABLE: + sources += [(path,) for path in DYAMOND_GRIDS.values()] + + read = [] + for source in sources: + flavour, suffix = ("data", ".nc") if len(source) == 2 else ("grid", ".nc") + if not _artifact(source, flavour, suffix).exists(): + _build(source) + read.append(source) + return read + + +class CachedFixtures: + """Mixin for benchmarks whose subject is not reading a file. + + Holds the one ``setup_cache`` the suite shares. asv keys ``setup_cache`` on + where it is defined and groups benchmarks by that key, so this single + definition -- inherited by every such class in every module -- runs once per + ``asv run`` rather than once per class. It returns ``None``, which asv reads + as "no cache argument", so the benchmark signatures stay as they are. + + 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) + for source in prime(include_dyamond=True) or [None]: + # Unbuffered and one line per source, so a batch log shows how far the + # reading has got. + print(f" read {' + '.join(Path(p).name for p in source)}" if source else " nothing to do", flush=True) diff --git a/benchmarks/mpas_dyamond.py b/benchmarks/mpas_dyamond.py index 3e4ecc8c0..a91db9f8a 100644 --- a/benchmarks/mpas_dyamond.py +++ b/benchmarks/mpas_dyamond.py @@ -1,31 +1,25 @@ -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: what these benchmarks measure is + # ``bounds`` and ``to_geodataframe``, not the MPAS reader. ``OpenGrid`` + # below is the one that measures reading, and it opens the real file. + self.uxgrid = self.cached_grid(grid_path_dict[resolution]) def teardown(self, resolution, **kwargs): del self.uxgrid @@ -33,21 +27,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 cd6177f47..ba352db5f 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -1,59 +1,51 @@ -import os -import urllib.request -from pathlib import Path - import numpy as np import uxarray as ux +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) - - -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. + + The dataset comes from the fixture cache rather than a fresh + ``open_dataset``: every benchmark below measures an algorithm over the mesh, + not the reader that produced it. The fixture is what the reader produced, + connectivity and ``face_areas`` included, so nothing here silently starts + measuring construction that used to come off disk. + """ 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 @@ -65,9 +57,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 the fixture keeps it, so + # it is dropped here to leave the computation to be measured. Safe to do + # to a fixture: each handout is a fresh ``Grid`` over a shallow copy. self.uxgrid._ds = self.uxgrid._ds.drop_vars("face_areas", errors="ignore") def time_face_areas(self, resolution): @@ -91,8 +85,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() @@ -126,26 +119,44 @@ 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. Under + ``launch_method: forkserver`` that process is forked from an interpreter + that has already imported the suite, so ``peakmem_*`` would report a warm + start plus whatever the parent held. A fresh interpreter is the only way to + keep measuring the thing this benchmark is named for. """ 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): @@ -181,11 +192,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 @@ -199,11 +210,11 @@ 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): 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 @@ -236,12 +247,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 @@ -255,7 +266,7 @@ class CrossSections(DatasetBenchmark): 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 @@ -268,12 +279,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 @@ -299,7 +310,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): @@ -308,10 +319,15 @@ 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, for the reason spelled out in + :class:`GradientColdStartRss`: 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.""" @@ -321,18 +337,29 @@ 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, for the reason spelled out in + :class:`GradientColdStartRss`. + """ 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.""" @@ -342,9 +369,17 @@ 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" 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: From 787e2d48e2dc6ce45fef54c15e73aa9a8119109b Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 12:49:42 -0500 Subject: [PATCH 02/18] benchmark debugging --- benchmarks/asv.conf.json | 4 +++- benchmarks/helpers/_peakmem.py | 33 +++++++++++++++++++++------------ benchmarks/mpas_ocean.py | 8 ++++++-- 3 files changed, 30 insertions(+), 15 deletions(-) diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 41ee30072..535ec6b2a 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -70,7 +70,9 @@ // which is what makes it safe here -- keep it that way. "launch_method": "forkserver", - "benchmark_timeout": 360, + // ``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/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/mpas_ocean.py b/benchmarks/mpas_ocean.py index ba352db5f..b62620d83 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -97,8 +97,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" From 96ce618dfd32171614ffc9a784bb4f3900344b25 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 15:29:08 -0500 Subject: [PATCH 03/18] Sample large benchmarks fewer times --- benchmarks/mpas_ocean.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index b62620d83..cb3a178bf 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -13,6 +13,12 @@ data_var = 'bottomDepth' +# 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 @@ -165,6 +171,7 @@ def track_peakmem_gradient(self, resolution): 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): @@ -215,6 +222,7 @@ def time_bilinear_remapping(self): self.uxds_120["bottomDepth"].remap.bilinear(self.uxds_480.uxgrid) class RemapUpsample(CachedFixtures): + repeat = SLOW_CALL_REPEAT def setup(self): self.uxds_120 = self.cached_dataset(*file_path_dict['120km']) @@ -244,6 +252,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') @@ -267,6 +277,7 @@ 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): From 4060897d7ec66fd9fc00589b7f4e8e0d4b0d2df1 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 16:00:47 -0500 Subject: [PATCH 04/18] Temp fork/thread speedup for connectivity. REEVALUATE OR REMOVE WHEN CONNECTIVITY BECOMES MULTITHREADED --- benchmarks/bench_connectivity.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 3a5859ee9..6732c395c 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -1,3 +1,5 @@ +import numba + import uxarray as ux from .helpers._fixtures import ( @@ -74,6 +76,8 @@ def setup(self, resolution, *args, **kwargs): # 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() @@ -115,3 +119,29 @@ def time_edge_face(self, resolution): def time_node_face(self, resolution): _ = self.uxgrid.node_face_connectivity.compute() + + +# Compiled at import rather than only in ``setup``: under +# ``launch_method: forkserver`` asv imports the suite once and forks every +# benchmark from that interpreter, so kernels compiled here are inherited by all +# of them. A forked child otherwise spends 0.496s of its own on JIT and cache +# loading before it can build anything. +# +# Safe only while the connectivity kernels are serial. Warming a +# ``parallel=True`` kernel -- or even calling ``.compile()`` on one -- launches +# numba's thread pool in the parent, and its OpenMP layer is not fork-safe, with +# no at-fork handler to rebuild it in the child. So check, rather than trust: +# this turns the day chunked connectivity lands into a loud import error instead +# of a hung benchmark. +_warmup() + +try: + numba.threading_layer() +except ValueError: + pass # nothing launched a pool, which is what we want to inherit +else: + raise RuntimeError( + "warming the connectivity kernels started numba's thread pool, which a " + "forked benchmark cannot safely inherit -- move _warmup() back into " + "setup() now that these kernels run in parallel" + ) From 7ab239d6192b0e70c78fb36a11e14ee37e816905 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 16:13:31 -0500 Subject: [PATCH 05/18] Cache benchmark sampling across forks too --- benchmarks/bench_connectivity.py | 27 ++++-------------- benchmarks/geometry_kernels.py | 29 +++++++++++++++++++ benchmarks/geometry_samebody.py | 36 ++++++++++++++++++++---- benchmarks/geometry_samebody_gcagca.py | 39 ++++++++++++++++++++++---- benchmarks/helpers/_warmup.py | 37 ++++++++++++++++++++++++ 5 files changed, 134 insertions(+), 34 deletions(-) create mode 100644 benchmarks/helpers/_warmup.py diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 6732c395c..ca1b4e20e 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -1,5 +1,3 @@ -import numba - import uxarray as ux from .helpers._fixtures import ( @@ -8,6 +6,7 @@ CachedFixtures, cached_topology, ) +from .helpers._warmup import warm_in_parent class GridBenchmark(CachedFixtures): @@ -125,23 +124,7 @@ def time_node_face(self, resolution): # ``launch_method: forkserver`` asv imports the suite once and forks every # benchmark from that interpreter, so kernels compiled here are inherited by all # of them. A forked child otherwise spends 0.496s of its own on JIT and cache -# loading before it can build anything. -# -# Safe only while the connectivity kernels are serial. Warming a -# ``parallel=True`` kernel -- or even calling ``.compile()`` on one -- launches -# numba's thread pool in the parent, and its OpenMP layer is not fork-safe, with -# no at-fork handler to rebuild it in the child. So check, rather than trust: -# this turns the day chunked connectivity lands into a loud import error instead -# of a hung benchmark. -_warmup() - -try: - numba.threading_layer() -except ValueError: - pass # nothing launched a pool, which is what we want to inherit -else: - raise RuntimeError( - "warming the connectivity kernels started numba's thread pool, which a " - "forked benchmark cannot safely inherit -- move _warmup() back into " - "setup() now that these kernels run in parallel" - ) +# loading before it can build anything. Guarded, because this only stays safe +# while the connectivity kernels are serial -- see +# :mod:`benchmarks.helpers._warmup`. +warm_in_parent(_warmup, "the connectivity kernels") diff --git a/benchmarks/geometry_kernels.py b/benchmarks/geometry_kernels.py index 08aea236e..058d555ea 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,30 @@ 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. Failures are swallowed on purpose: this + module imports its kernels inside ``setup`` so that a commit missing a symbol + fails one benchmark rather than the whole module's collection, and warming + here must not take that away. + """ + for cls in ( + EFTPrimitives, + AccucrossKernels, + OrientPredicates, + GCAGCAIntersection, + GCAConstLatIntersection, + ): + try: + cls().setup() + except Exception: + pass + + +# Warmed at import so every forked benchmark inherits the compiled kernels; see +# :mod:`benchmarks.helpers._warmup`. +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..e93108025 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,30 @@ 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)``. The seed is fixed, so hoisting the arrays out of + ``setup`` changes what is measured not at all -- and lets a forked benchmark + inherit them rather than generate them again. + """ + 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 +340,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 +355,11 @@ def time_accux_dispatch(self): _batch_accux_gca_dispatch(self.ga, self.gb) +# Prepared at import so every forked benchmark inherits it; see +# :mod:`benchmarks.helpers._warmup` for why that is safe only while these +# kernels stay serial. +warm_in_parent(_prepare, "the gca-gca drivers") + + if __name__ == "__main__": main() diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py new file mode 100644 index 000000000..dbab5692b --- /dev/null +++ b/benchmarks/helpers/_warmup.py @@ -0,0 +1,37 @@ +"""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 instead of being paid for again by every benchmark +process. On this suite that is 0.5s of JIT and cache loading for the +connectivity kernels and 4.1s of case generation for the gca-gca drivers, each +of which was being repeated per benchmark. + +Only work that leaves no numba thread pool behind may be warmed this way. +Running a ``parallel=True`` kernel -- or merely calling ``.compile()`` on one -- +launches the pool, and numba's OpenMP layer is not fork-safe, with no at-fork +handler to rebuild it in the child. So this checks rather than trusts: the day a +warmed kernel goes parallel becomes a loud import error rather than a benchmark +that hangs on a cluster. +""" + +import numba + +__all__ = ["warm_in_parent"] + + +def warm_in_parent(warm, what): + """Runs ``warm``, then fails if it left a numba thread pool behind. + + ``what`` names the thing being warmed, for the error message. + """ + warm() + try: + numba.threading_layer() + except ValueError: + return # nothing launched a pool, which is what makes this inheritable + 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" + ) From c63676c29a3a9427ac7d2ec832348fc00de09ba1 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 17:37:52 -0500 Subject: [PATCH 06/18] Per-resolution async benchmarks --- benchmarks/helpers/_fixtures.py | 56 +++++++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 6 deletions(-) diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index a39975c14..50722d303 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -32,9 +32,12 @@ """ import hashlib +import multiprocessing import os import tempfile import urllib.request +import uuid +from concurrent.futures import ProcessPoolExecutor from pathlib import Path import numpy as np @@ -172,8 +175,13 @@ def _write(dataset, artifact_path, writer): Via a scratch name in the same directory, so a process racing this one sees either no artifact or a complete one, never a half-written file. """ + if artifact_path.exists(): + # A grid reached through both its own source and a (grid, data) pair + # would otherwise be written twice, and two writers racing on one + # scratch name is worse than wasteful. + return scratch = artifact_path.with_name( - f"{artifact_path.stem}.{os.getpid()}.tmp{artifact_path.suffix}" + f"{artifact_path.stem}.{uuid.uuid4().hex}.tmp{artifact_path.suffix}" ) writer(dataset, scratch) os.replace(scratch, artifact_path) @@ -268,7 +276,7 @@ def cached_dataset(grid_path, data_path): return ux.UxDataset(_loaded[artifact_path].copy(), uxgrid=cached_grid(grid_path)) -def prime(include_dyamond=None): +def prime(include_dyamond=None, 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 @@ -278,6 +286,14 @@ def prime(include_dyamond=None): ``include_dyamond``) asks for them, because a filtered run should not pay for reading four grids off campaign storage that it will never touch. On a machine that does have them, prime from the CLI before ``asv 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 dyamond grid is still being read, instead of queueing behind it. It + costs an interpreter per worker, so it only pays when a read is slower than + a process start -- the default stays sequential for that reason. Processes + rather than threads because the netCDF/HDF5 stack here is not thread-safe + for concurrent opens: threads produce HDF5 errors, measured, not assumed. """ if include_dyamond is None: include_dyamond = os.environ.get(PRIME_VAR, "").lower() == "all" @@ -288,13 +304,37 @@ def prime(include_dyamond=None): if include_dyamond and DYAMOND_AVAILABLE: sources += [(path,) for path in DYAMOND_GRIDS.values()] - read = [] + missing = [] for source in sources: flavour, suffix = ("data", ".nc") if len(source) == 2 else ("grid", ".nc") if not _artifact(source, flavour, suffix).exists(): + missing.append(source) + + # Reading a (grid, data) pair produces that grid's artifacts too, so a + # grid-only source covered by a pair here would just read the grid a second + # time to write nothing. + 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, not forked: 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. The extra second of startup is nothing against a + # read this is worth parallelizing. + 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) - read.append(source) - return read + return missing class CachedFixtures: @@ -328,7 +368,11 @@ def setup_cache(self): # ``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) - for source in prime(include_dyamond=True) or [None]: + # 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(include_dyamond=True, workers=4 if DYAMOND_AVAILABLE else 1) or [None]: # Unbuffered and one line per source, so a batch log shows how far the # reading has got. print(f" read {' + '.join(Path(p).name for p in source)}" if source else " nothing to do", flush=True) From 85f7c0c5a10d3520507f08182f72bc1d0e42a808 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 18:04:09 -0500 Subject: [PATCH 07/18] Benchmark IO caching fixes --- .gitignore | 1 + benchmarks/bench_connectivity.py | 10 +++++- benchmarks/helpers/_fixtures.py | 57 ++++++++++++++++++++++++-------- 3 files changed, 54 insertions(+), 14 deletions(-) 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/bench_connectivity.py b/benchmarks/bench_connectivity.py index ca1b4e20e..16dd3b590 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -5,6 +5,7 @@ GRIDS_BY_RESOLUTION, CachedFixtures, cached_topology, + preload_topologies, ) from .helpers._warmup import warm_in_parent @@ -127,4 +128,11 @@ def time_node_face(self, resolution): # loading before it can build anything. Guarded, because this only stays safe # while the connectivity kernels are serial -- see # :mod:`benchmarks.helpers._warmup`. -warm_in_parent(_warmup, "the connectivity kernels") +def _warm_parent(): + _warmup() + # And, if asked, the topologies themselves, so a forked benchmark inherits + # them instead of reading its resolution's artifact again. + preload_topologies(GRIDS_BY_RESOLUTION[res] for res in ALL_RESOLUTIONS) + + +warm_in_parent(_warm_parent, "the connectivity kernels") diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index 50722d303..cd357c439 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -34,7 +34,6 @@ import hashlib import multiprocessing import os -import tempfile import urllib.request import uuid from concurrent.futures import ProcessPoolExecutor @@ -60,6 +59,7 @@ "cached_dataset", "cached_grid", "cached_topology", + "preload_topologies", "prime", ] @@ -127,6 +127,7 @@ def _cookbook(filename): CACHE_DIR_VAR = "UXARRAY_BENCH_CACHE_DIR" PRIME_VAR = "UXARRAY_BENCH_PRIME" +PRELOAD_VAR = "UXARRAY_BENCH_PRELOAD" # The arguments ``ux.Grid.from_topology`` takes, in order. _TOPOLOGY_ARRAYS = ("node_lon", "node_lat", "face_node_connectivity") @@ -136,15 +137,9 @@ def _cookbook(filename): def cache_dir(): - """Directory the cached artifacts live in. - - ``UXARRAY_BENCH_CACHE_DIR`` overrides the default, and on a cluster it - should: the cache only pays off on a filesystem faster than the one holding - the source grids, and putting it somewhere that outlives the job means each - grid is read once per machine rather than once per job. - """ - root = Path(os.environ.get(CACHE_DIR_VAR) or tempfile.gettempdir()) - cached = root / "uxarray-bench-fixtures" + """Directory the cached artifacts live in: ``benchmarks/_io_cache``.""" + root = Path(os.environ.get(CACHE_DIR_VAR) or BENCHMARK_DIR) + cached = root / "_io_cache" cached.mkdir(parents=True, exist_ok=True) return cached @@ -154,10 +149,18 @@ def _artifact(source, flavour, suffix): ``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, and on the uxarray version, because the artifact is - that version's reader output. + served something stale. + + Deliberately not keyed on the uxarray version. These meshes are stable and + the artifacts are meant to persist -- across jobs, and across the commits asv + walks. Including the version cost a fresh read per commit and produced two + full sets of artifacts here, because a benchmark run imports the uxarray + installed in asv's environment while a direct run from the repo root imports + the working tree, and those report different versions. The tradeoff is that a + change to how a reader parses these files does not invalidate the cache on + its own: delete ``_io_cache`` when that happens. """ - parts = [ux.__version__] + parts = [] for path in source: stat = os.stat(path) parts.append(f"{os.path.realpath(path)}:{stat.st_size}:{stat.st_mtime_ns}") @@ -337,6 +340,34 @@ def prime(include_dyamond=None, workers=1): 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. Which is why a benchmark's memory + appears and then vanishes when it exits -- expected, but at 3.75km it means + re-reading gigabytes for every benchmark in the module. + + Loading them in the parent instead means one read per resolution per run, + shared copy-on-write by every child. Off unless ``UXARRAY_BENCH_PRELOAD`` + asks for it, because the parent then holds every resolution at once and asv's + discovery import pays for it too -- a poor trade unless the reads are slow, + which on campaign storage they are. + + 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`). + """ + if not os.environ.get(PRELOAD_VAR): + return 0 + 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. From 9cd8fa8c38e0f73f55d4edff30785a1da44fde2f Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 11:20:46 -0500 Subject: [PATCH 08/18] cached IO benchmarks deslopping --- benchmarks/bench_connectivity.py | 29 +---- benchmarks/face_bounds.py | 11 +- benchmarks/geometry_kernels.py | 7 +- benchmarks/geometry_samebody_gcagca.py | 8 +- benchmarks/helpers/_fixtures.py | 150 ++++++++----------------- benchmarks/helpers/_warmup.py | 16 +-- benchmarks/mpas_dyamond.py | 4 +- benchmarks/mpas_ocean.py | 24 +--- 8 files changed, 68 insertions(+), 181 deletions(-) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 16dd3b590..4796c1224 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -13,13 +13,9 @@ 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. + param_names = ['resolution', ] params = [ALL_RESOLUTIONS, ] - - # A single connectivity build at 3.75km does not fit in the 360s default - # from ``asv.conf.json``. timeout = 1200 def setup(self, resolution, *args, **kwargs): @@ -47,12 +43,6 @@ def _warmup(): ``_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. - - Warmed on the coarsest grid in ``params``, because resolution decides how - long the kernels run but not which signatures compile. Warming at the - benchmark's own resolution instead means eight full connectivity builds - before the sample -- 0.493s against a 0.106s sample at 120km, and the gap - only widens from there. """ global _numba_warmed_up if _numba_warmed_up: @@ -64,9 +54,7 @@ def _warmup(): 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): @@ -121,17 +109,12 @@ def time_node_face(self, resolution): _ = self.uxgrid.node_face_connectivity.compute() -# Compiled at import rather than only in ``setup``: under -# ``launch_method: forkserver`` asv imports the suite once and forks every -# benchmark from that interpreter, so kernels compiled here are inherited by all -# of them. A forked child otherwise spends 0.496s of its own on JIT and cache -# loading before it can build anything. Guarded, because this only stays safe -# while the connectivity kernels are serial -- see -# :mod:`benchmarks.helpers._warmup`. +# 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, if asked, the topologies themselves, so a forked benchmark inherits - # them instead of reading its resolution's artifact again. + # And, if asked, the topologies themselves... preload_topologies(GRIDS_BY_RESOLUTION[res] for res in ALL_RESOLUTIONS) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index d406723d3..f5d58ea47 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -4,8 +4,6 @@ from .helpers._memsize import grid_nbytes from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss -# One grid per reader, from the shared registry. ``mpas`` here is the same mesh -# the oQU ``480km`` benchmarks use, through the copy in the repo. grid_quad_hex = GRIDS_BY_FORMAT["ugrid-quad-hexagon"] grid_geoflow = GRIDS_BY_FORMAT["ugrid-geoflow"] grid_scrip = GRIDS_BY_FORMAT["scrip-outCSne8"] @@ -69,11 +67,7 @@ class FaceBoundsColdStartRss: magnitude lower. Measured in a subprocess of its own rather than through asv's ``peakmem_*``, - which reports ``ru_maxrss`` for the benchmark process. Under - ``launch_method: forkserver`` that process is forked from an interpreter that - has already imported the suite, so a ``peakmem_*`` here would be reporting a - warm start plus whatever the parent was holding. A fresh interpreter is the - only way to keep measuring what this benchmark is named for. + which reports ``ru_maxrss`` for the benchmark process. """ params = FaceBounds.params @@ -81,9 +75,6 @@ class FaceBoundsColdStartRss: def setup_cache(self): """Compile the njit kernels before anything is measured. - - The subprocess inherits numba's on-disk cache, not this process's - memory, so this keeps compilation out of the measured cold start. """ for grid_path in self.params: ux.open_grid(grid_path).bounds diff --git a/benchmarks/geometry_kernels.py b/benchmarks/geometry_kernels.py index 058d555ea..ceb9f995c 100644 --- a/benchmarks/geometry_kernels.py +++ b/benchmarks/geometry_kernels.py @@ -213,10 +213,7 @@ 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. Failures are swallowed on purpose: this - module imports its kernels inside ``setup`` so that a commit missing a symbol - fails one benchmark rather than the whole module's collection, and warming - here must not take that away. + cannot drift out of step with them. """ for cls in ( EFTPrimitives, @@ -231,6 +228,4 @@ def _warm_classes(): pass -# Warmed at import so every forked benchmark inherits the compiled kernels; see -# :mod:`benchmarks.helpers._warmup`. warm_in_parent(_warm_classes, "the geometry kernels") diff --git a/benchmarks/geometry_samebody_gcagca.py b/benchmarks/geometry_samebody_gcagca.py index e93108025..fd3494ed3 100644 --- a/benchmarks/geometry_samebody_gcagca.py +++ b/benchmarks/geometry_samebody_gcagca.py @@ -317,9 +317,8 @@ def _prepare(): 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)``. The seed is fixed, so hoisting the arrays out of - ``setup`` changes what is measured not at all -- and lets a forked benchmark - inherit them rather than generate them again. + ``@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: @@ -355,9 +354,6 @@ def time_accux_dispatch(self): _batch_accux_gca_dispatch(self.ga, self.gb) -# Prepared at import so every forked benchmark inherits it; see -# :mod:`benchmarks.helpers._warmup` for why that is safe only while these -# kernels stay serial. warm_in_parent(_prepare, "the gca-gca drivers") diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index cd357c439..b2c32df0b 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -1,32 +1,22 @@ """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 -- around 160 times across -the suite once the dyamond resolutions are visible. For a mesh on local disk -that is a rounding error. For one on campaign storage it is most of the run, and -it is spent inside the benchmark's own timeout. +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 keep reading the real thing: -``quad_hexagon``, ``OpenGrid`` in ``mpas_dyamond``, ``import``, and the -cold-start peak-memory benchmarks. They take their paths from the registries -here too, so the suite declares its inputs in one place either way. +Benchmarks that do measure opening a file behave as before. -Two flavours, picked per benchmark: +Two flavors: ``topology`` - the three arrays ``Grid.from_topology`` needs and nothing else, for - benchmarks that mean to build the rest themselves -- 2.3MB of the 102MB - 120km MPAS file. + the three arrays ``Grid.from_topology`` needs and nothing else ``grid`` / ``dataset`` - everything the reader produced, so a benchmark still gets the - ``face_areas`` and connectivity variables an MPAS file carries on disk - rather than silently measuring their construction. + everything the reader produced from ``Grid.open_grid`` and + ``Grid.open_dataset`` -One source read produces both, so choosing between them costs nothing. - -Artifacts are keyed on the uxarray build as well as on the file, because an -artifact is one version's reader output and asv walks commits. That means a -fresh read per commit; ``prime`` therefore leaves the dyamond grids out unless +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`` therefore leaves the dyamond grids out unless asked, and there is a CLI (``python -m benchmarks.helpers._fixtures``) to fill the cache from a batch script instead of from inside a benchmark. """ @@ -79,7 +69,7 @@ def _cookbook(filename): return path -# Grids, and grid/data pairs, by mesh resolution. +# Grids and grid/data pairs, by mesh resolution. OQU_GRIDS = { "480km": _cookbook("oQU480.grid.nc"), "120km": _cookbook("oQU120.grid.nc"), @@ -96,17 +86,13 @@ def _cookbook(filename): "3.75km": Path("/glade/campaign/cisl/vast/uxarray/data/dyamond/3.75km/grid.nc"), } -# Asked once here, rather than separately in each module that cares. +# 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 -# Two ladders rather than one, because which of them a benchmark belongs on is a -# per-benchmark decision: an algorithm that does not care which model wrote the -# mesh can take the wide one, while anything tied to the oQU pair -- or too slow -# to run four dyamond resolutions of -- stays on the narrow one. OQU_RESOLUTIONS = list(OQU_GRIDS) ALL_RESOLUTIONS = list(GRIDS_BY_RESOLUTION) @@ -129,10 +115,7 @@ def _cookbook(filename): PRIME_VAR = "UXARRAY_BENCH_PRIME" PRELOAD_VAR = "UXARRAY_BENCH_PRELOAD" -# The arguments ``ux.Grid.from_topology`` takes, in order. -_TOPOLOGY_ARRAYS = ("node_lon", "node_lat", "face_node_connectivity") - -# Artifact path -> what it holds, for this process. +# Per path, which artifacts are actually loaded _loaded = {} @@ -144,21 +127,12 @@ def cache_dir(): return cached -def _artifact(source, flavour, suffix): - """Where ``flavour`` of ``source`` is 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. - - Deliberately not keyed on the uxarray version. These meshes are stable and - the artifacts are meant to persist -- across jobs, and across the commits asv - walks. Including the version cost a fresh read per commit and produced two - full sets of artifacts here, because a benchmark run imports the uxarray - installed in asv's environment while a direct run from the repo root imports - the working tree, and those report different versions. The tradeoff is that a - change to how a reader parses these files does not invalidate the cache on - its own: delete ``_io_cache`` when that happens. """ parts = [] for path in source: @@ -169,19 +143,17 @@ def _artifact(source, flavour, suffix): # 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}-{flavour}-{digest}{suffix}" + return cache_dir() / f"{stem}-{flavor}-{digest}{suffix}" def _write(dataset, artifact_path, writer): - """Writes ``dataset`` to ``artifact_path``, atomically. + """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, never a half-written file. + either no artifact or a complete one. """ if artifact_path.exists(): - # A grid reached through both its own source and a (grid, data) pair - # would otherwise be written twice, and two writers racing on one - # scratch name is worse than wasteful. + # 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}" @@ -191,18 +163,12 @@ def _write(dataset, artifact_path, writer): def _read_dataset(artifact_path): - """Reads back a cached ``xr.Dataset``. - - ``mask_and_scale=False`` is load-bearing: the connectivity variables carry - ``_FillValue``, which xarray would otherwise consume, handing back float64 - where uxarray's njit kernels require int64 -- they fail to type rather than - returning something wrong, but they do fail. - """ + """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 flavour of it, from the one read.""" + """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) @@ -213,12 +179,11 @@ def _build(source): uxds.load() uxgrid = uxds.uxgrid uxgrid._ds.load() - # A ``UxDataset`` is an ``xr.Dataset`` subclass, so it writes itself; the - # grid half is cached separately just below. + # A ``UxDataset`` grid is cached separately. data_ds = uxds _write( - {name: getattr(uxgrid, name).data for name in _TOPOLOGY_ARRAYS}, + {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), @@ -228,11 +193,11 @@ def _build(source): _write(data_ds, _artifact(source, "data", ".nc"), lambda ds, path: ds.to_netcdf(path)) -def _ensure(source, flavour, suffix): +def _ensure(source, flavor, suffix): """Path to a cached artifact, building the source's artifacts if need be.""" - artifact_path = _artifact(source, flavour, suffix) + artifact_path = _artifact(source, flavor, suffix) if not artifact_path.exists(): - _build(source if flavour == "data" else source[:1]) + _build(source if flavor == "data" else source[:1]) return artifact_path @@ -247,7 +212,7 @@ def cached_topology(grid_path): 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 _TOPOLOGY_ARRAYS) + _loaded[artifact_path] = tuple(cached[name] for name in ["node_lon", "node_lat", "face_node_connectivity"]) return _loaded[artifact_path] @@ -260,13 +225,7 @@ def _cached_grid_ds(grid_path): def cached_grid(grid_path): - """A ``Grid`` carrying everything the reader found in ``grid_path``. - - A fresh ``Grid`` over a shallow copy each call, so a benchmark that - populates or normalizes something does not hand its leftovers to the next - repeat: uxarray assigns new variables into ``_ds`` rather than writing - through the arrays, so the copy isolates that while the data stays shared. - """ + """A fresh ``Grid`` carrying everything the reader found in ``grid_path`` via shallow copy.""" return ux.Grid(_cached_grid_ds(grid_path).copy()) @@ -285,18 +244,12 @@ def prime(include_dyamond=None, workers=1): 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. - The dyamond grids are left out unless ``UXARRAY_BENCH_PRIME=all`` (or - ``include_dyamond``) asks for them, because a filtered run should not pay - for reading four grids off campaign storage that it will never touch. On a - machine that does have them, prime from the CLI before ``asv 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 dyamond grid is still being read, instead of queueing behind it. It - costs an interpreter per worker, so it only pays when a read is slower than - a process start -- the default stays sequential for that reason. Processes - rather than threads because the netCDF/HDF5 stack here is not thread-safe - for concurrent opens: threads produce HDF5 errors, measured, not assumed. + 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. """ if include_dyamond is None: include_dyamond = os.environ.get(PRIME_VAR, "").lower() == "all" @@ -309,13 +262,11 @@ def prime(include_dyamond=None, workers=1): missing = [] for source in sources: - flavour, suffix = ("data", ".nc") if len(source) == 2 else ("grid", ".nc") - if not _artifact(source, flavour, suffix).exists(): + 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, so a - # grid-only source covered by a pair here would just read the grid a second - # time to write nothing. + # 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 @@ -325,10 +276,9 @@ def prime(include_dyamond=None, workers=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, not forked: 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. The extra second of startup is nothing against a - # read this is worth parallelizing. + + # 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"), @@ -346,15 +296,7 @@ def preload_topologies(grid_paths): 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. Which is why a benchmark's memory - appears and then vanishes when it exits -- expected, but at 3.75km it means - re-reading gigabytes for every benchmark in the module. - - Loading them in the parent instead means one read per resolution per run, - shared copy-on-write by every child. Off unless ``UXARRAY_BENCH_PRELOAD`` - asks for it, because the parent then holds every resolution at once and asv's - discovery import pays for it too -- a poor trade unless the reads are slow, - which on campaign storage they are. + 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`). @@ -371,11 +313,10 @@ def preload_topologies(grid_paths): class CachedFixtures: """Mixin for benchmarks whose subject is not reading a file. - Holds the one ``setup_cache`` the suite shares. asv keys ``setup_cache`` on - where it is defined and groups benchmarks by that key, so this single - definition -- inherited by every such class in every module -- runs once per - ``asv run`` rather than once per class. It returns ``None``, which asv reads - as "no cache argument", so the benchmark signatures stay as they are. + 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 @@ -399,11 +340,10 @@ def setup_cache(self): # ``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(include_dyamond=True, workers=4 if DYAMOND_AVAILABLE else 1) or [None]: - # Unbuffered and one line per source, so a batch log shows how far the - # reading has got. print(f" read {' + '.join(Path(p).name for p in source)}" if source else " nothing to do", flush=True) diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py index dbab5692b..3ad8a9c1e 100644 --- a/benchmarks/helpers/_warmup.py +++ b/benchmarks/helpers/_warmup.py @@ -1,17 +1,13 @@ """Warming benchmark state in the interpreter every benchmark is forked from. -Under ``launch_method: forkserver`` asv imports the suite once and forks each +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 instead of being paid for again by every benchmark -process. On this suite that is 0.5s of JIT and cache loading for the -connectivity kernels and 4.1s of case generation for the gca-gca drivers, each -of which was being repeated per benchmark. +inherited copy-on-write. -Only work that leaves no numba thread pool behind may be warmed this way. -Running a ``parallel=True`` kernel -- or merely calling ``.compile()`` on one -- -launches the pool, and numba's OpenMP layer is not fork-safe, with no at-fork -handler to rebuild it in the child. So this checks rather than trusts: the day a -warmed kernel goes parallel becomes a loud import error rather than a benchmark +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. """ diff --git a/benchmarks/mpas_dyamond.py b/benchmarks/mpas_dyamond.py index a91db9f8a..8c9454a14 100644 --- a/benchmarks/mpas_dyamond.py +++ b/benchmarks/mpas_dyamond.py @@ -16,9 +16,7 @@ class BaseGridBenchmark(CachedFixtures): params = [list(DYAMOND_GRIDS), ] def setup(self, resolution, **kwargs): - # The cached grid, not a fresh read: what these benchmarks measure is - # ``bounds`` and ``to_geodataframe``, not the MPAS reader. ``OpenGrid`` - # below is the one that measures reading, and it opens the real file. + # The cached grid, not a fresh read self.uxgrid = self.cached_grid(grid_path_dict[resolution]) def teardown(self, resolution, **kwargs): diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index cb3a178bf..749a3ab88 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -27,12 +27,6 @@ class DatasetBenchmark(CachedFixtures): """Class used as a template for benchmarks requiring a ``UxDataset`` in this module across both resolutions. - - The dataset comes from the fixture cache rather than a fresh - ``open_dataset``: every benchmark below measures an algorithm over the mesh, - not the reader that produced it. The fixture is what the reader produced, - connectivity and ``face_areas`` included, so nothing here silently starts - measuring construction that used to come off disk. """ param_names = ['resolution', ] params = [OQU_RESOLUTIONS, ] @@ -65,9 +59,9 @@ def setup(self, resolution, *args, **kwargs): # The coarsest grid, purely to compile the njit kernel _ = self.cached_grid(OQU_GRIDS[OQU_RESOLUTIONS[0]]).face_areas super().setup(resolution, *args, **kwargs) - # MPAS meshes carry ``face_areas`` on disk and the fixture keeps it, so - # it is dropped here to leave the computation to be measured. Safe to do - # to a fixture: each handout is a fresh ``Grid`` over a shallow copy. + + # 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): @@ -136,11 +130,7 @@ class GradientColdStartRss: magnitude lower. Measured in a subprocess of its own rather than through asv's ``peakmem_*``, - which reports ``ru_maxrss`` for the benchmark process. Under - ``launch_method: forkserver`` that process is forked from an interpreter - that has already imported the suite, so ``peakmem_*`` would report a warm - start plus whatever the parent held. A fresh interpreter is the only way to - keep measuring the thing this benchmark is named for. + which reports ``ru_maxrss`` for the benchmark process. """ param_names = ["resolution"] @@ -336,8 +326,7 @@ def time_zonal_average(self, resolution): class ZonalAveragePeakMem: """Peak memory of a cold-start non-conservative zonal-mean sweep. - A fresh interpreter per sample, for the reason spelled out in - :class:`GradientColdStartRss`: the cold start is the subject, and a forked + A fresh interpreter per sample. The cold start is the subject, and a forked benchmark process no longer has one. """ @@ -369,8 +358,7 @@ def track_peakmem_zonal_average(self, resolution): class CrossSectionsPeakMem: """Peak memory of a cold-start constant-latitude cross-section sweep. - A fresh interpreter per sample, for the reason spelled out in - :class:`GradientColdStartRss`. + A fresh interpreter per sample, """ param_names = ["resolution", "lat_step"] From b1a6be60ec13982c985ed25f8c08f12374c62e92 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 11:41:58 -0500 Subject: [PATCH 09/18] cached IO deslop global variables --- benchmarks/helpers/_fixtures.py | 36 ++++++++++++++++++++------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index b2c32df0b..e1da4d1db 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -19,6 +19,18 @@ fresh read per commit. ``prime`` therefore leaves the dyamond grids out unless asked, and there is a CLI (``python -m benchmarks.helpers._fixtures``) to fill the cache from a batch script instead of from inside a benchmark. + +Three environment variables tune all of this: + +``UXARRAY_BENCH_CACHE_DIR`` + root the ``_io_cache`` directory is created under, in place of + ``benchmarks/`` -- worth pointing at local scratch when the checkout itself + lives on a shared filesystem +``UXARRAY_BENCH_PRIME`` + ``all`` primes the dyamond grids as well, which ``prime`` skips by default +``UXARRAY_BENCH_PRELOAD`` + any non-empty value has ``preload_topologies`` do its work; unset, it is a + no-op, since preloading only pays off under ``launch_method: forkserver`` """ import hashlib @@ -56,27 +68,27 @@ BENCHMARK_DIR = Path(__file__).resolve().parents[1] REPO_DIR = BENCHMARK_DIR.parent -_COOKBOOK_URL = ( +_COOKBOOK_MESH_URL = ( "https://github.com/ProjectPythia/unstructured-grid-viz-cookbook/raw/main/meshfiles" ) -def _cookbook(filename): +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_URL}/{filename}", filename=path) + urllib.request.urlretrieve(f"{_COOKBOOK_MESH_URL}/{filename}", filename=path) return path # Grids and grid/data pairs, by mesh resolution. OQU_GRIDS = { - "480km": _cookbook("oQU480.grid.nc"), - "120km": _cookbook("oQU120.grid.nc"), + "480km": _cookbook_mesh("oQU480.grid.nc"), + "120km": _cookbook_mesh("oQU120.grid.nc"), } OQU_DATASETS = { - "480km": (OQU_GRIDS["480km"], _cookbook("oQU480.data.nc")), - "120km": (OQU_GRIDS["120km"], _cookbook("oQU120.data.nc")), + "480km": (OQU_GRIDS["480km"], _cookbook_mesh("oQU480.data.nc")), + "120km": (OQU_GRIDS["120km"], _cookbook_mesh("oQU120.data.nc")), } DYAMOND_GRIDS = { @@ -111,17 +123,13 @@ def _cookbook(filename): REPO_DIR / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "data.nc", ) -CACHE_DIR_VAR = "UXARRAY_BENCH_CACHE_DIR" -PRIME_VAR = "UXARRAY_BENCH_PRIME" -PRELOAD_VAR = "UXARRAY_BENCH_PRELOAD" - # 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(CACHE_DIR_VAR) or BENCHMARK_DIR) + 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 @@ -252,7 +260,7 @@ def prime(include_dyamond=None, workers=1): netCDF/HDF5 stack is not generally thread-safe for concurrent opens. """ if include_dyamond is None: - include_dyamond = os.environ.get(PRIME_VAR, "").lower() == "all" + include_dyamond = os.environ.get("UXARRAY_BENCH_PRIME", "").lower() == "all" sources = [(path,) for path in OQU_GRIDS.values()] sources += [(path,) for path in GRIDS_BY_FORMAT.values()] @@ -301,7 +309,7 @@ def preload_topologies(grid_paths): 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`). """ - if not os.environ.get(PRELOAD_VAR): + if not os.environ.get("UXARRAY_BENCH_PRELOAD"): return 0 loaded = 0 for grid_path in grid_paths: From 41f68b21eba1d37f81926f9ca80c8ebdeb515733 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 13:06:43 -0500 Subject: [PATCH 10/18] bench cached IO: more global variable deslopping --- benchmarks/bench_connectivity.py | 2 +- benchmarks/helpers/_fixtures.py | 29 ++++++----------------------- 2 files changed, 7 insertions(+), 24 deletions(-) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 4796c1224..85ca61cb7 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -114,7 +114,7 @@ def time_node_face(self, resolution): # of them. Only safe while the connectivity kernels are serial def _warm_parent(): _warmup() - # And, if asked, the topologies themselves... + # And the topologies themselves... preload_topologies(GRIDS_BY_RESOLUTION[res] for res in ALL_RESOLUTIONS) diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index e1da4d1db..2996be33a 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -16,21 +16,9 @@ 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`` therefore leaves the dyamond grids out unless -asked, and there is a CLI (``python -m benchmarks.helpers._fixtures``) to fill -the cache from a batch script instead of from inside a benchmark. - -Three environment variables tune all of this: - -``UXARRAY_BENCH_CACHE_DIR`` - root the ``_io_cache`` directory is created under, in place of - ``benchmarks/`` -- worth pointing at local scratch when the checkout itself - lives on a shared filesystem -``UXARRAY_BENCH_PRIME`` - ``all`` primes the dyamond grids as well, which ``prime`` skips by default -``UXARRAY_BENCH_PRELOAD`` - any non-empty value has ``preload_topologies`` do its work; unset, it is a - no-op, since preloading only pays off under ``launch_method: forkserver`` +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 @@ -246,7 +234,7 @@ def cached_dataset(grid_path, data_path): return ux.UxDataset(_loaded[artifact_path].copy(), uxgrid=cached_grid(grid_path)) -def prime(include_dyamond=None, workers=1): +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 @@ -259,13 +247,10 @@ def prime(include_dyamond=None, workers=1): the new interpreter startup. Processes rather than threads because the standard netCDF/HDF5 stack is not generally thread-safe for concurrent opens. """ - if include_dyamond is None: - include_dyamond = os.environ.get("UXARRAY_BENCH_PRIME", "").lower() == "all" - sources = [(path,) for path in OQU_GRIDS.values()] sources += [(path,) for path in GRIDS_BY_FORMAT.values()] sources += list(OQU_DATASETS.values()) - if include_dyamond and DYAMOND_AVAILABLE: + if DYAMOND_AVAILABLE: sources += [(path,) for path in DYAMOND_GRIDS.values()] missing = [] @@ -309,8 +294,6 @@ def preload_topologies(grid_paths): 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`). """ - if not os.environ.get("UXARRAY_BENCH_PRELOAD"): - return 0 loaded = 0 for grid_path in grid_paths: cached_topology(grid_path) # held by the process-level memo from here on @@ -353,5 +336,5 @@ def setup_cache(self): # 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(include_dyamond=True, workers=4 if DYAMOND_AVAILABLE else 1) or [None]: + 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) From 4dad004743e6405d77f114550ab7637ab3914e63 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 18:39:23 -0500 Subject: [PATCH 11/18] Prime the fixture cache before asv run, not during it asv preimports the benchmark suite before it runs any setup_cache (asv/runner.py, spawner.preimport() ahead of the run loop), so on a cold cache bench_connectivity's import-time preload_topologies is what fills it -- serially, in the forkserver parent, before a single benchmark starts. CachedFixtures.setup_cache then finds everything already built, and prime(workers=...) never runs on the path it was written for. Filling it from the CLI first puts those reads back in the parallel prime. Worth a second or two on the GitHub runners, which only see the oQU grids; worth rather more on a machine that can reach the four dyamond grids on campaign storage. Co-Authored-By: Claude Opus 5 --- .github/workflows/asv-benchmarking-pr.yml | 2 ++ .github/workflows/asv-benchmarking.yml | 2 ++ 2 files changed, 4 insertions(+) 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 From ec83767c568e9c96cda504b75ca9748c5da26599 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 25 Aug 2026 15:30:30 -0500 Subject: [PATCH 12/18] Lazy neighborhood filter kernel compilation --- test/test_dependencies.py | 30 +++++++++++++++++++ uxarray/grid/neighbors.py | 62 +++++++++++++++++++++++++++++++++------ 2 files changed, 83 insertions(+), 9 deletions(-) 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..a0c301b18 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,3 +1,4 @@ +import threading import warnings from typing import Callable @@ -1246,6 +1247,43 @@ def kernel(data, flat, starts, counts, param, out): return kernel +class _LazyKernel: + """A kernel that builds itself the first time it is called. + + ``guvectorize`` compiles at decoration time when it is given explicit + signatures, so calling :func:`_make_kernel` at module scope would compile + every reduction during ``import uxarray`` + """ + + __slots__ = ("_reduce_fn", "_kernel", "_lock") + + def __init__(self, reduce_fn): + self._reduce_fn = reduce_fn + self._kernel = None + self._lock = threading.Lock() + + def __call__(self, *args): + kernel = self._kernel + if kernel is None: + # Checked again under the lock: dask's threaded scheduler can call + # one reduction from several workers at once + with self._lock: + kernel = self._kernel + if kernel is None: + kernel = self._kernel = _make_kernel(self._reduce_fn) + return kernel(*args) + + def __getstate__(self): + # Only the reducer travels. A compiled gufunc is a ``numpy.ufunc``, + # which pickle cannot address by name and so refuses outright. + return self._reduce_fn + + def __setstate__(self, reduce_fn): + self._reduce_fn = reduce_fn + self._kernel = None + self._lock = threading.Lock() + + # Reducers take ``(window, param)``; those without a parameter ignore the # second argument. Numba keys its cache by code object rather than qualified # name, so the identically-named lambdas below do not collide. @@ -1284,17 +1322,23 @@ 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))) +# +# Naming a kernel is what binds a reduction to it; building it is deferred to +# the first call, for the reasons in :class:`_LazyKernel`. The two are +# separate on purpose -- each name below still resolves to its own object when +# this module is read, so a reduction that names no kernel is a ``NameError`` +# here rather than a lookup that fails once someone runs it. +_MEAN_KERNEL = _LazyKernel(lambda window, _: np.mean(window)) +_SUM_KERNEL = _LazyKernel(lambda window, _: np.sum(window)) +_MIN_KERNEL = _LazyKernel(lambda window, _: np.min(window)) +_MAX_KERNEL = _LazyKernel(lambda window, _: np.max(window)) +_PTP_KERNEL = _LazyKernel(lambda window, _: np.max(window) - np.min(window)) +_MEDIAN_KERNEL = _LazyKernel(_median) +_VAR_KERNEL = _LazyKernel(_variance) +_STD_KERNEL = _LazyKernel(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)) +_QUANTILE_KERNEL = _LazyKernel(lambda window, q: np.quantile(window, q)) def _as_quantile(q, scale: float): From 856441007de9366f08df5d64286302c32e94c136 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 25 Aug 2026 15:58:47 -0500 Subject: [PATCH 13/18] Lazy nb kernels with functools cache instead --- uxarray/grid/neighbors.py | 135 ++++++++++++++++++-------------------- 1 file changed, 65 insertions(+), 70 deletions(-) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index a0c301b18..d796f6d78 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1,4 +1,4 @@ -import threading +import functools import warnings from typing import Callable @@ -1247,43 +1247,6 @@ def kernel(data, flat, starts, counts, param, out): return kernel -class _LazyKernel: - """A kernel that builds itself the first time it is called. - - ``guvectorize`` compiles at decoration time when it is given explicit - signatures, so calling :func:`_make_kernel` at module scope would compile - every reduction during ``import uxarray`` - """ - - __slots__ = ("_reduce_fn", "_kernel", "_lock") - - def __init__(self, reduce_fn): - self._reduce_fn = reduce_fn - self._kernel = None - self._lock = threading.Lock() - - def __call__(self, *args): - kernel = self._kernel - if kernel is None: - # Checked again under the lock: dask's threaded scheduler can call - # one reduction from several workers at once - with self._lock: - kernel = self._kernel - if kernel is None: - kernel = self._kernel = _make_kernel(self._reduce_fn) - return kernel(*args) - - def __getstate__(self): - # Only the reducer travels. A compiled gufunc is a ``numpy.ufunc``, - # which pickle cannot address by name and so refuses outright. - return self._reduce_fn - - def __setstate__(self, reduce_fn): - self._reduce_fn = reduce_fn - self._kernel = None - self._lock = threading.Lock() - - # Reducers take ``(window, param)``; those without a parameter ignore the # second argument. Numba keys its cache by code object rather than qualified # name, so the identically-named lambdas below do not collide. @@ -1315,30 +1278,62 @@ 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. +# One compiled kernel per reduction. The methods on ``Neighborhood`` below call +# into these directly. Non-compiled functions are only provided hooks through +# ``Neighborhood.reduce``. If new compiled reductions are desired, they should +# follow this pattern. # -# Naming a kernel is what binds a reduction to it; building it is deferred to -# the first call, for the reasons in :class:`_LazyKernel`. The two are -# separate on purpose -- each name below still resolves to its own object when -# this module is read, so a reduction that names no kernel is a ``NameError`` -# here rather than a lookup that fails once someone runs it. -_MEAN_KERNEL = _LazyKernel(lambda window, _: np.mean(window)) -_SUM_KERNEL = _LazyKernel(lambda window, _: np.sum(window)) -_MIN_KERNEL = _LazyKernel(lambda window, _: np.min(window)) -_MAX_KERNEL = _LazyKernel(lambda window, _: np.max(window)) -_PTP_KERNEL = _LazyKernel(lambda window, _: np.max(window) - np.min(window)) -_MEDIAN_KERNEL = _LazyKernel(_median) -_VAR_KERNEL = _LazyKernel(_variance) -_STD_KERNEL = _LazyKernel(lambda window, ddof: np.sqrt(_variance(window, ddof))) +# ``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. + + +@functools.cache +def _mean_kernel(): + return _make_kernel(lambda window, _: np.mean(window)) + + +@functools.cache +def _sum_kernel(): + return _make_kernel(lambda window, _: np.sum(window)) + + +@functools.cache +def _min_kernel(): + return _make_kernel(lambda window, _: np.min(window)) + + +@functools.cache +def _max_kernel(): + return _make_kernel(lambda window, _: np.max(window)) + + +@functools.cache +def _ptp_kernel(): + return _make_kernel(lambda window, _: np.max(window) - np.min(window)) + + +@functools.cache +def _median_kernel(): + return _make_kernel(_median) + + +@functools.cache +def _var_kernel(): + return _make_kernel(_variance) + + +@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. -_QUANTILE_KERNEL = _LazyKernel(lambda window, q: np.quantile(window, q)) +@functools.cache +def _quantile_kernel(): + return _make_kernel(lambda window, q: np.quantile(window, q)) def _as_quantile(q, scale: float): @@ -1553,45 +1548,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, _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, _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, _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, _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, _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, _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, _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, _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, _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, _quantile_kernel, _as_quantile(q, 100.0)) def reduce(self, uxda, func: Callable): """Reduces each neighborhood with an arbitrary callable. @@ -1636,7 +1631,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) From 7127726beb6ec7a36740730b409df1ba23a62fd7 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 26 Aug 2026 12:34:36 -0500 Subject: [PATCH 14/18] Move kernels inside --- benchmarks/mpas_ocean.py | 3 +- uxarray/grid/neighbors.py | 138 +++++++++++++++++++------------------- 2 files changed, 72 insertions(+), 69 deletions(-) diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index adf679a3a..04325563f 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -107,7 +107,8 @@ def track_nbytes_gradient(self, resolution): def track_peakmem_gradient(self, resolution): """Transient high-water allocation of taking a gradient.""" - return peak_allocated(lambda: self.uxds[data_var].gradient()) + with numba_threads(1): + return peak_allocated(lambda: self.uxds[data_var].gradient()) track_peakmem_gradient.unit = "bytes" diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index d796f6d78..753092f18 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1278,64 +1278,6 @@ def _median(window, _): return np.median(window) -# One compiled kernel per reduction. The methods on ``Neighborhood`` below call -# into these directly. Non-compiled functions are only provided hooks through -# ``Neighborhood.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. - - -@functools.cache -def _mean_kernel(): - return _make_kernel(lambda window, _: np.mean(window)) - - -@functools.cache -def _sum_kernel(): - return _make_kernel(lambda window, _: np.sum(window)) - - -@functools.cache -def _min_kernel(): - return _make_kernel(lambda window, _: np.min(window)) - - -@functools.cache -def _max_kernel(): - return _make_kernel(lambda window, _: np.max(window)) - - -@functools.cache -def _ptp_kernel(): - return _make_kernel(lambda window, _: np.max(window) - np.min(window)) - - -@functools.cache -def _median_kernel(): - return _make_kernel(_median) - - -@functools.cache -def _var_kernel(): - return _make_kernel(_variance) - - -@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. -@functools.cache -def _quantile_kernel(): - return _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) @@ -1546,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. From c00fd1d67d4e08c1ef41310990e7fd77665cbbba Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 26 Aug 2026 13:32:30 -0500 Subject: [PATCH 15/18] Cached IO: try to impose more CI thread safety --- benchmarks/asv.conf.json | 25 ++++++++++++--- benchmarks/helpers/_warmup.py | 57 +++++++++++++++++++++++++++-------- 2 files changed, 65 insertions(+), 17 deletions(-) diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 535ec6b2a..770ea641d 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -102,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"]} }, diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py index 3ad8a9c1e..5479794de 100644 --- a/benchmarks/helpers/_warmup.py +++ b/benchmarks/helpers/_warmup.py @@ -4,30 +4,61 @@ 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. +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 sys + import numba __all__ = ["warm_in_parent"] +# Numba installs fork handlers for these two; ``omp`` is on its own. +_FORK_SAFE = frozenset({"tbb", "workqueue"}) + +_reported = False + def warm_in_parent(warm, what): - """Runs ``warm``, then fails if it left a numba thread pool behind. + """Runs ``warm``, then checks any pool it leaves behind survives a fork. - ``what`` names the thing being warmed, for the error message. + ``what`` names the thing being warmed, for the report. """ warm() + try: - numba.threading_layer() + layer = numba.threading_layer() except ValueError: - return # nothing launched a pool, which is what makes this inheritable - 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" + 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, ) From 526e3206a6f45cd64880b1cdf940f80e905a00f2 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 26 Aug 2026 14:29:27 -0500 Subject: [PATCH 16/18] Cached IO: lazy interpreter warming; NetCDF warming --- benchmarks/helpers/_fixtures.py | 38 +++++++++++++++++++++++++++++++-- benchmarks/helpers/_warmup.py | 29 +++++++++++++++++++++++-- 2 files changed, 63 insertions(+), 4 deletions(-) diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index 2996be33a..9ab574197 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -53,6 +53,8 @@ "prime", ] +from ._warmup import warm_in_parent + BENCHMARK_DIR = Path(__file__).resolve().parents[1] REPO_DIR = BENCHMARK_DIR.parent @@ -221,8 +223,18 @@ def _cached_grid_ds(grid_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()) + """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): @@ -338,3 +350,25 @@ def setup_cache(self): # 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/_warmup.py b/benchmarks/helpers/_warmup.py index 5479794de..ae84c10a9 100644 --- a/benchmarks/helpers/_warmup.py +++ b/benchmarks/helpers/_warmup.py @@ -24,11 +24,18 @@ results" -- and the pool exists either way, and asv forks either way. """ +import os import sys import numba -__all__ = ["warm_in_parent"] +__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"}) @@ -36,11 +43,29 @@ _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. + ``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: From a064b3046d180eebe8b08b7bbf49480bf5823237 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 26 Aug 2026 15:12:54 -0500 Subject: [PATCH 17/18] Cached IO: configuration tweaks --- .github/workflows/asv-benchmarking-pr.yml | 9 +++++++++ benchmarks/asv.conf.json | 7 ++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/.github/workflows/asv-benchmarking-pr.yml b/.github/workflows/asv-benchmarking-pr.yml index 49ec2cc27..c23abe74d 100644 --- a/.github/workflows/asv-benchmarking-pr.yml +++ b/.github/workflows/asv-benchmarking-pr.yml @@ -43,6 +43,15 @@ jobs: python-build mamba + # asv builds its own conda environment under ``benchmarks/env``, separate + # from the one above, and without this it solves and installs that from + # scratch on every run before a single benchmark starts. + - name: Cache asv's benchmark environment + uses: actions/cache@v6 + with: + path: ${{ env.ASV_DIR }}/env + key: asv-env-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('benchmarks/asv.conf.json', 'ci/environment.yml') }} + - name: Run Benchmarks shell: bash -l {0} id: benchmark diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 770ea641d..9b3852f53 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -193,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 From a0f4a80adb13e5c7399c85c8bc08e615db9d2693 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 26 Aug 2026 16:36:04 -0500 Subject: [PATCH 18/18] Partially revert a064b30, add cpu info dump --- .github/workflows/asv-benchmarking-pr.yml | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/.github/workflows/asv-benchmarking-pr.yml b/.github/workflows/asv-benchmarking-pr.yml index c23abe74d..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: @@ -43,15 +48,6 @@ jobs: python-build mamba - # asv builds its own conda environment under ``benchmarks/env``, separate - # from the one above, and without this it solves and installs that from - # scratch on every run before a single benchmark starts. - - name: Cache asv's benchmark environment - uses: actions/cache@v6 - with: - path: ${{ env.ASV_DIR }}/env - key: asv-env-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('benchmarks/asv.conf.json', 'ci/environment.yml') }} - - name: Run Benchmarks shell: bash -l {0} id: benchmark