From 1c7665d9a5b417aac4ef9e3f6cd541a5595a35a9 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 21 Aug 2026 12:07:04 -0500 Subject: [PATCH 01/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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 3cb8de062e047fe8744ec66b128c1a1b84619236 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 18:39:38 -0500 Subject: [PATCH 12/19] Shard the track_* benchmarks across concurrent asv processes asv never runs two benchmarks at once: --parallel builds environments and nothing else ("Build (but don't benchmark) in parallel", asv/commands/common_args.py), and the forkserver is a listen(1) socket that forks one child and waitpids on it before reading another command (asv_runner/server.py). Concurrency has to come from running several asv run processes, which is what this does. Only track_* is safe to run that way, and it is also the half worth running that way. A time_* result is a wall-clock measurement that a co-tenant corrupts, so those stay serial. A track_* result is a count -- an array size, a tracemalloc high-water mark, a child's ru_maxrss -- and none of them move because another core is busy. They are also the benchmarks that leave the machine idle: every traced one runs under numba_threads(1), so the track_ pass spends itself on one core. Measured on a 12-core box with the oQU grids only: the serial pass takes 89.3s and a 4-way shard takes 44.4s. Every track_nbytes_* and traced track_peakmem_* value comes back bit-identical; the subprocess_peak_rss figures land within 3%, mean -0.31%, which is the run-to-run spread those have anyway. It does not reach 4x because each extra asv run costs 13.1s of fixed overhead -- ~6.5s discovering and ~6s preimporting, both full suite imports -- and the 7.17s traced gradient at 120km cannot be subdivided. Past four to six shards there is nothing left to win. Details worth knowing: - setup_cache groups that compile kernels at every resolution move whole, so a shard cannot pay one twice. Only the shared helpers._fixtures key, which is just prime(), and benchmarks with no setup_cache are split. - Packing is per parameter combination, weighted by the durations in whatever results are already on disk, because the suite is lopsided. - Shard results dirs are seeded from the real one so --skip-existing behaves as it would under a plain asv run. - The serial remainder runs last on purpose: asv run calls Results.load_data before it records, so it reads the merged track results back and writes them out alongside the timings. Co-Authored-By: Claude Opus 5 --- benchmarks/helpers/_shard.py | 573 +++++++++++++++++++++++++++++++++++ 1 file changed, 573 insertions(+) create mode 100644 benchmarks/helpers/_shard.py diff --git a/benchmarks/helpers/_shard.py b/benchmarks/helpers/_shard.py new file mode 100644 index 000000000..7147d6b1e --- /dev/null +++ b/benchmarks/helpers/_shard.py @@ -0,0 +1,573 @@ +"""Running the ``track_*`` half of the suite as concurrent ``asv run`` processes. + +asv never runs two benchmarks at once. ``--parallel`` builds environments in +parallel and does nothing else -- its own help text is "Build (but don't +benchmark) in parallel" (``asv/commands/common_args.py``) -- and the forkserver +is a ``listen(1)`` socket that forks one child and ``waitpid``s on it before it +will read another command (``asv_runner/server.py``). Concurrency has to come +from running several ``asv run`` processes side by side, which is what this +module does. + +Only the ``track_*`` benchmarks are safe to run that way, and they are also the +ones worth running that way: + +* A ``time_*`` result is a wall-clock measurement, and a co-tenant on the + machine corrupts it. Those stay serial. +* A ``track_*`` result is a count. ``track_nbytes_*`` returns an array size and + ``track_peakmem_*`` returns a tracemalloc high-water figure or a child's + ``ru_maxrss``; none of them move because another core is busy. +* Every traced benchmark runs under ``numba_threads(1)`` (see + :mod:`benchmarks.helpers._peakmem`), so today the whole ``track_`` pass leaves + every core but one idle while it runs the most expensive benchmarks in the + suite. + +What is *not* invariant is the thread count itself. The peaks measured through +``subprocess_peak_rss`` come from a child that runs numba's kernels at whatever +width the platform gives it, and numba follows CPU affinity -- so N shards +sharing a node would each see fewer threads and report a different number than +an unsharded run would. ``--threads-per-shard`` pins that width explicitly (1 by +default, which is what the traced benchmarks already do to themselves), so the +figure stays a property of the benchmark rather than of how the run happened to +be sharded. Expect a one-time step in the ``subprocess_peak_rss`` series the +first time this is used, and none after. + +A run goes:: + + prime fill the fixture cache once, here, rather than in N shards + discover one serial ``asv run --bench just-discover``, which builds and + installs the project at the commit and writes the full + ``benchmarks.json`` + shards N concurrent ``asv run --bench ...``, each with its own + ``results_dir`` so they cannot clobber one another's file + merge fold the shard result files into the real results dir + serial the remaining (``time_*``) benchmarks, one process, as usual + +The discover pass is what lets the shards start together: +``Environment.install_project`` returns immediately once ``installed_commit_hash`` +matches the commit, so the shards find the project built instead of N of them +racing to build it into the one shared environment. + +The serial pass runs last on purpose. ``asv run`` calls ``Results.load_data`` +before it records anything (``asv/commands/run.py``), so it reads the merged +track results off disk and writes them back out alongside the timings rather +than replacing the file. + +Usage, from the repository root:: + + python -m benchmarks.helpers._shard # branch head + python -m benchmarks.helpers._shard main^! --workers 8 + python -m benchmarks.helpers._shard v2024.02.0..main -- --skip-existing + +Anything after a bare ``--`` is handed to every ``asv run`` this launches. +""" + +import argparse +import copy +import itertools +import json +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +__all__ = ["main", "merge_results", "plan_shards"] + +BENCHMARK_DIR = Path(__file__).resolve().parents[1] +REPO_DIR = BENCHMARK_DIR.parent + +# Columns of an asv result row that are indexed by parameter combination, and so +# merge element-wise across shards. Everything else in the row describes the +# benchmark as a whole. Names come from ``Results.save`` (``asv/results.py``). +_PER_PARAM_COLUMNS = frozenset( + { + "result", + "samples", + "stats_ci_99_a", + "stats_ci_99_b", + "stats_q_25", + "stats_q_75", + "stats_number", + "stats_repeat", + } +) + +# ``setup_cache`` groups that may be split across shards. Every other group in +# this suite compiles kernels and computes at every resolution before anything is +# measured -- they carry ``setup_cache.timeout = 1800`` for that reason -- so +# splitting one would pay that cost once per shard. ``CachedFixtures.setup_cache`` +# is just ``prime()``, which costs a stat per file once the cache is warm, and +# ``None`` is no setup_cache at all. +_SPLITTABLE_PREFIX = "helpers._fixtures:" + + +def _splittable(setup_cache_key): + """Whether instances sharing this ``setup_cache_key`` may land in different shards.""" + return setup_cache_key is None or str(setup_cache_key).startswith(_SPLITTABLE_PREFIX) + + +def _asv(): + """Imports asv, with a message pointing at the environment if it is absent. + + Imported lazily: this module is a CLI that drives ``asv``, and asv itself is + not a dependency of the benchmarks it runs. The project environment asv + builds has no asv in it. + """ + try: + import asv.config + import asv.results + import asv.util + except ImportError as exc: # pragma: no cover - environment problem, not logic + raise SystemExit( + "benchmarks.helpers._shard needs asv importable in the interpreter " + f"running it (the one that provides the `asv` command): {exc}" + ) + return asv + + +def _instance_names(benchmark): + """Every selectable name of ``benchmark``, as asv's ``--bench`` matches them. + + asv filters parameterized benchmarks on ``funcname(param0, param1)`` built + from the parameter *reprs* it recorded at discovery, not on the values + (``Benchmarks.__init__`` in ``asv/benchmarks.py``), so this reproduces that + string exactly rather than re-deriving it from the benchmark's own params. + """ + params = benchmark.get("params") or [] + if not params: + return [benchmark["name"]] + return [ + f"{benchmark['name']}({', '.join(combo)})" + for combo in itertools.product(*params) + ] + + +def _weights(results_dir): + """Per-instance runtime hints, in seconds, from whatever results are on disk. + + Bin packing on measured cost beats packing on a count, because this suite is + lopsided: one traced gradient runs longer than every ``track_nbytes_*`` + benchmark put together. Missing entries fall back to a flat weight, so a + first run still shards, just less evenly. + + asv records one ``duration`` per benchmark rather than per parameter + combination, so this spreads it evenly over the combinations. + """ + asv = _asv() + hints = {} + root = Path(results_dir) + if not root.is_dir(): + return hints + for path in sorted(root.glob("*/*.json")): + if path.name in ("machine.json", "benchmarks.json"): + continue + try: + data = asv.util.load_json(path, api_version=asv.results.Results.api_version) + except Exception: + continue + columns = data.get("result_columns") or [] + if "duration" not in columns: + continue + index = columns.index("duration") + for name, row in data.get("results", {}).items(): + if len(row) <= index or row[index] is None: + continue + duration = row[index] + if isinstance(duration, list): + duration = sum(x for x in duration if x) + hints[name] = float(duration) + return hints + + +def plan_shards(benchmarks, selected, n_shards, weights=None): + """Partitions ``selected`` instance names into ``n_shards`` lists. + + ``benchmarks`` is the parsed ``benchmarks.json``. Units that share a + non-splittable ``setup_cache_key`` move together, so an expensive + ``setup_cache`` is paid once rather than once per shard. Longest-first onto + the lightest shard, which is the standard greedy fit and is good enough here: + the schedule is dominated by one benchmark that no packing can subdivide. + + Returns a list of ``n_shards`` lists of instance names, shortest lists last. + """ + weights = weights or {} + by_name = {b["name"]: b for b in benchmarks} + + # Build the units: either one instance, or a whole setup_cache group. + units = {} + for instance in selected: + benchmark_name = instance.split("(", 1)[0] + key = by_name.get(benchmark_name, {}).get("setup_cache_key") + unit = instance if _splittable(key) else f"setup_cache:{key}" + units.setdefault(unit, []).append(instance) + + def cost(instances): + total = 0.0 + for instance in instances: + benchmark_name = instance.split("(", 1)[0] + hint = weights.get(benchmark_name) + benchmark = by_name.get(benchmark_name) + if hint is None or benchmark is None: + total += 1.0 + else: + total += hint / max(1, len(_instance_names(benchmark))) + return total + + shards = [[] for _ in range(n_shards)] + loads = [0.0] * n_shards + for instances in sorted(units.values(), key=lambda i: -cost(i)): + target = loads.index(min(loads)) + shards[target].extend(instances) + loads[target] += cost(instances) + return [s for s in shards if s] + + +def _shard_config(raw_conf, conf, results_dir, html_dir): + """A copy of the project's asv config that writes results somewhere else. + + Every path is absolutised. ``asv --config`` chdirs to the directory holding + the config file (``asv/main.py``), so a shard config in a scratch directory + would otherwise resolve ``repo``, ``benchmark_dir`` and ``env_dir`` against + that scratch directory instead of the project. + + ``env_dir`` deliberately stays the shared one: the shards run against the + environment the discover pass already installed into. + """ + shard = copy.deepcopy(raw_conf) + shard["repo"] = os.path.abspath(conf.repo) + shard["benchmark_dir"] = os.path.abspath(conf.benchmark_dir) + shard["env_dir"] = os.path.abspath(conf.env_dir) + shard["results_dir"] = os.path.abspath(results_dir) + shard["html_dir"] = os.path.abspath(html_dir) + return shard + + +def merge_results(shard_dirs, dest_dir): + """Folds each shard's result file into the matching file in ``dest_dir``. + + Shards run disjoint benchmarks, but two shards can hold different parameter + combinations of the *same* benchmark, and asv writes those as one row per + benchmark with ``None`` in the combinations it did not run. So rows merge + element-wise on the parameter-indexed columns and take the first value on + the rest. + + Returns the list of files written. + """ + asv = _asv() + api_version = asv.results.Results.api_version + written = [] + + # Group by the per-(machine, commit, env) filename asv chose. + files = {} + for shard_dir in shard_dirs: + for path in sorted(Path(shard_dir).glob("*/*.json")): + if path.name in ("machine.json", "benchmarks.json"): + continue + files.setdefault(str(Path(path.parent.name) / path.name), []).append(path) + + for relative, paths in sorted(files.items()): + dest = Path(dest_dir) / relative + merged = None + columns = None + if dest.is_file(): + merged = asv.util.load_json(dest, api_version=api_version) + columns = merged.get("result_columns") + + for path in paths: + data = asv.util.load_json(path, api_version=api_version) + if merged is None: + merged = data + columns = data.get("result_columns") + continue + if data.get("result_columns") != columns: + raise RuntimeError( + f"{path} has different result columns than the file it merges " + "into; the asv versions writing them differ" + ) + for name, row in data.get("results", {}).items(): + existing = merged["results"].get(name) + merged["results"][name] = ( + list(row) if existing is None else _merge_row(existing, row, columns) + ) + for key, value in (data.get("durations") or {}).items(): + merged.setdefault("durations", {}) + merged["durations"][key] = merged["durations"].get(key, 0) + value + + # asv drops trailing nulls from every row; keep the file in that shape. + for name, row in merged["results"].items(): + while row and row[-1] is None: + row.pop() + + dest.parent.mkdir(parents=True, exist_ok=True) + asv.util.write_json(dest, merged, api_version=api_version, compact=True) + written.append(dest) + + # The machine description is identical in every shard; carry one over. + machine = paths[0].parent / "machine.json" + if machine.is_file() and not (dest.parent / "machine.json").is_file(): + shutil.copyfile(machine, dest.parent / "machine.json") + + return written + + +def _merge_row(a, b, columns): + """Merges two asv result rows for the same benchmark.""" + width = len(columns) + a = list(a) + [None] * (width - len(a)) + b = list(b) + [None] * (width - len(b)) + row = [] + for index, column in enumerate(columns): + x, y = a[index], b[index] + if column in _PER_PARAM_COLUMNS: + row.append(_merge_per_param(x, y)) + elif column == "duration": + row.append(None if x is None and y is None else (x or 0) + (y or 0)) + elif column == "started_at": + present = [v for v in (x, y) if v is not None] + row.append(min(present) if present else None) + else: + # params, version, profile: describe the benchmark, not one run of it. + row.append(x if x is not None else y) + return row + + +def _merge_per_param(x, y): + """Element-wise merge of two parameter-indexed columns, taking whichever ran.""" + if x is None: + return y + if y is None: + return x + width = max(len(x), len(y)) + x = list(x) + [None] * (width - len(x)) + y = list(y) + [None] * (width - len(y)) + return [xi if xi is not None else yi for xi, yi in zip(x, y)] + + +def _core_ranges(n_shards, threads_per_shard): + """``taskset`` core lists giving each shard its own cores, or ``None``. + + ``None`` where the platform has no ``taskset`` or there are not enough cores + to go round, in which case the shards are left to the scheduler -- which + costs nothing here, since the measurements are counts rather than times. + """ + if not shutil.which("taskset"): + return None + available = os.cpu_count() or 1 + if n_shards * threads_per_shard > available: + return None + return [ + ",".join( + str(c) + for c in range(i * threads_per_shard, (i + 1) * threads_per_shard) + ) + for i in range(n_shards) + ] + + +def _run(command, env=None, check=True): + print(f" $ {shlex.join(command)}", flush=True) + return subprocess.run(command, env=env, check=check) + + +def main(argv=None): + parser = argparse.ArgumentParser( + prog="python -m benchmarks.helpers._shard", + description="Run the track_* benchmarks as concurrent asv processes, then the rest serially.", + ) + parser.add_argument( + "range", + nargs="?", + default=None, + help="Commit range, as asv takes it. Defaults to the configured branch head.", + ) + parser.add_argument( + "--workers", type=int, default=4, help="Number of concurrent asv processes (default 4)." + ) + parser.add_argument( + "--threads-per-shard", + type=int, + default=1, + help="numba/OpenMP threads each shard may use. Pinned rather than inherited so a " + "subprocess_peak_rss figure does not depend on how the run was sharded (default 1).", + ) + parser.add_argument( + "--select", + default=r"\.track_", + help=r"Regex picking the benchmarks to run in parallel (default '\.track_').", + ) + parser.add_argument( + "--config", + default=str(BENCHMARK_DIR / "asv.conf.json"), + help="Path to asv.conf.json.", + ) + parser.add_argument( + "--track-only", + action="store_true", + help="Stop after the sharded pass instead of running the serial remainder.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the plan and the commands, run nothing.", + ) + parser.add_argument( + "asv_args", + nargs="*", + default=[], + help="Arguments after a bare -- are passed to every asv run (e.g. --skip-existing).", + ) + args = parser.parse_args(argv) + + asv = _asv() + config_path = os.path.abspath(args.config) + # Match ``asv --config``: paths in the file are relative to the file. + os.chdir(os.path.dirname(config_path)) + conf = asv.config.Config.load(config_path) + raw_conf = asv.util.load_json(config_path, js_comments=True) + results_dir = os.path.abspath(conf.results_dir) + + passthrough = list(args.asv_args) + commit = [args.range] if args.range else [] + + # 1. Fill the fixture cache here rather than letting N shards each find it + # cold. Also keeps it off the forkserver parent's import, which is where + # it otherwise lands -- asv preimports the suite before any setup_cache. + print("priming the fixture cache", flush=True) + if not args.dry_run: + # Not fatal: the shards fill the cache themselves if this cannot run, + # just once per shard instead of once here. Priming needs uxarray + # importable in *this* interpreter, which is the one that has asv -- + # true of the CI environment (``ci/asv.yml`` installs the project into + # it) but not of every environment asv can be driven from. + try: + from . import _fixtures + + read = _fixtures.prime(workers=args.workers) + print(f" built {len(read)} source(s)", flush=True) + except Exception as exc: + print(f" skipped ({type(exc).__name__}: {exc})", flush=True) + + # 2. One serial pass that discovers and, as a side effect, builds and + # installs the project at the commit -- so the shards do not race to. + print("discovering benchmarks (and installing the project)", flush=True) + discover = ["asv", "run", *commit, "--bench", "just-discover", *passthrough] + if not args.dry_run: + _run(discover) + else: + print(f" $ {shlex.join(discover)}", flush=True) + + benchmarks_json = Path(results_dir) / "benchmarks.json" + if not benchmarks_json.is_file(): + raise SystemExit(f"no benchmarks.json at {benchmarks_json}; discovery failed") + discovered = asv.util.load_json(benchmarks_json, api_version=2) + benchmarks = list(discovered.values()) + + pattern = re.compile(args.select) + selected = [ + instance + for benchmark in benchmarks + for instance in _instance_names(benchmark) + if pattern.search(instance) + ] + if not selected: + raise SystemExit(f"nothing matched --select {args.select!r}") + + weights = _weights(results_dir) + shards = plan_shards(benchmarks, selected, args.workers, weights) + print(f"{len(selected)} instances over {len(shards)} shard(s):", flush=True) + for i, shard in enumerate(shards): + load = sum( + weights.get(instance.split("(", 1)[0], 0.0) + / max(1, len(_instance_names(next(b for b in benchmarks if b["name"] == instance.split("(", 1)[0])))) + for instance in shard + ) + print(f" shard {i}: {len(shard):3} instances, ~{load:.1f}s of recorded work", flush=True) + for instance in shard: + print(f" {instance}", flush=True) + + cores = _core_ranges(len(shards), args.threads_per_shard) + if cores is None: + print(" (not pinning cores: no taskset, or too few cores)", flush=True) + + scratch = tempfile.mkdtemp(prefix="asv-shard-") + shard_dirs = [] + processes = [] + try: + for i, shard in enumerate(shards): + shard_results = os.path.join(scratch, f"shard{i}", "results") + shard_dirs.append(shard_results) + + # Seed the shard with the results already on disk, so that + # ``--skip-existing`` and friends see the same history a plain + # ``asv run`` would. Without this a shard reads an empty results dir + # and re-runs everything. Re-merging the seeded rows afterwards is a + # no-op: they are the values already in the destination. + if os.path.isdir(results_dir): + shutil.copytree(results_dir, shard_results, dirs_exist_ok=True) + shard_conf_path = os.path.join(scratch, f"shard{i}", "asv.conf.json") + os.makedirs(os.path.dirname(shard_conf_path), exist_ok=True) + with open(shard_conf_path, "w") as handle: + json.dump( + _shard_config( + raw_conf, conf, shard_results, os.path.join(scratch, f"shard{i}", "html") + ), + handle, + ) + + # Results are per (machine, commit, env); the shards must agree on + # the machine or they would write to different subdirectories and + # never merge. + command = ["asv", "--config", shard_conf_path, "run", *commit] + for instance in shard: + command += ["--bench", f"^{re.escape(instance)}$"] + command += passthrough + + if cores is not None: + command = ["taskset", "-c", cores[i]] + command + + env = dict(os.environ) + env["NUMBA_NUM_THREADS"] = str(args.threads_per_shard) + env["OMP_NUM_THREADS"] = str(args.threads_per_shard) + + print(f" $ {shlex.join(command)}", flush=True) + if not args.dry_run: + processes.append(subprocess.Popen(command, env=env)) + + codes = [p.wait() for p in processes] + if any(codes): + print(f"warning: shard exit codes {codes}", file=sys.stderr, flush=True) + + # 3. Fold the shards back together. + if not args.dry_run: + print("merging shard results", flush=True) + for path in merge_results(shard_dirs, results_dir): + print(f" wrote {path}", flush=True) + finally: + if not args.dry_run: + shutil.rmtree(scratch, ignore_errors=True) + + # 4. The remainder, serially, on a quiet machine. asv reads the merged file + # back in (Results.load_data) before it records, so this adds to it. + if not args.track_only: + print("running the remaining benchmarks serially", flush=True) + serial = [ + "asv", + "run", + *commit, + "--bench", + f"^(?!.*(?:{args.select}))", + *passthrough, + ] + if not args.dry_run: + _run(serial, check=False) + else: + print(f" $ {shlex.join(serial)}", flush=True) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 658119956bc07cf5743f89a7650f11f13a5f6041 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:11:43 -0500 Subject: [PATCH 13/19] Build the neighborhood kernels on first use, not at import The PR benchmark job dies in asv's discover step: RuntimeError: warming the connectivity kernels started numba's thread pool, which a forked benchmark cannot safely inherit taking discovery with it, and with that the whole run, for the baseline as well as the contender -- asv takes the benchmark suite from the working tree and only the project from each commit. The connectivity kernels are innocent. neighbors.py built nine reduction kernels at module scope, and `guvectorize` with explicit signatures compiles at decoration time; building a `target="parallel"` ufunc calls `_launch_threads()`. So `import uxarray` started an OpenMP pool before any user code ran. numba's OpenMP layer installs no at-fork handler, so a child inherits a pool it cannot use -- which breaks anything that forks workers after importing uxarray, multiprocessing's "fork" start method included, and here asv's `forkserver`, which forks every benchmark from one interpreter that imported the suite. This branch only started failing when it merged main, where the kernels arrived with #941. They are now built the first time they are named, through a module `__getattr__`. Every call site looks them up as module globals, so nothing else changes. Importing uxarray no longer starts a thread pool, and no longer pays for nine compiles it will usually not need. warm_in_parent also sampled the pool only after warming, so it blamed whatever it had just warmed for a pool that anything earlier in the import could have started -- which is what sent this to the benchmarks rather than to neighbors.py. It now samples before as well, and the two cases raise different errors, because they need different fixes. Verified against this tree: discovery fails as CI reports it, and with the kernels built lazily it succeeds, 86 benchmarks, with the connectivity warm left exactly as it was. Co-Authored-By: Claude Opus 5 --- benchmarks/helpers/_warmup.py | 35 +++++++++++++++++++++++--- uxarray/grid/neighbors.py | 46 ++++++++++++++++++++++++++++------- 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py index 3ad8a9c1e..21b68af85 100644 --- a/benchmarks/helpers/_warmup.py +++ b/benchmarks/helpers/_warmup.py @@ -16,16 +16,43 @@ __all__ = ["warm_in_parent"] -def warm_in_parent(warm, what): - """Runs ``warm``, then fails if it left a numba thread pool behind. +def _pool_running(): + """Whether numba's thread pool has been launched in this interpreter. - ``what`` names the thing being warmed, for the error message. + ``threading_layer()`` raises until a ``parallel=True`` kernel has been + compiled or run, which is exactly the event that makes this process unsafe + to fork a numba benchmark from. """ - warm() try: numba.threading_layer() except ValueError: + return False + return True + + +def warm_in_parent(warm, what): + """Runs ``warm``, then fails if a numba thread pool is running afterwards. + + ``what`` names the thing being warmed, for the error message. + + The pool is sampled before as well as after. Without that baseline this + reports whatever it happened to warm as the cause of a pool that something + else started earlier in the import, which sends you looking in the wrong + module -- the two cases need different fixes, so they get different errors. + """ + already_running = _pool_running() + warm() + if not _pool_running(): return # nothing launched a pool, which is what makes this inheritable + + if already_running: + raise RuntimeError( + f"numba's thread pool was already running before {what} was warmed, " + "so this interpreter cannot safely fork a benchmark that uses numba. " + "Warming is not the cause: something earlier in the import of this " + "suite, or of uxarray itself, compiled or ran a parallel=True kernel. " + "Find that and move it, or drop launch_method: forkserver." + ) 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 " diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 03a4c7cd5..7cf3fdf8b 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1284,17 +1284,45 @@ 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))) +# +# Each is built the first time it is named rather than at import. ``guvectorize`` +# with explicit signatures compiles at decoration time, and building a +# ``target="parallel"`` ufunc launches numba's OpenMP thread pool as part of +# that -- so assigning these at module scope made ``import uxarray`` start a +# thread pool before any user code ran. Besides charging every import for nine +# compiles it will usually not need, that leaves the process unsafe to fork: +# numba's OpenMP layer installs no at-fork handler, so a child inherits a pool +# it cannot use. Anything that forks workers after importing uxarray is affected +# -- multiprocessing's "fork" start method, and asv's ``forkserver`` launch +# method, which forks every benchmark from one interpreter that imported the +# suite. +_KERNEL_BUILDERS = { + "_MEAN_KERNEL": lambda: _make_kernel(lambda window, _: np.mean(window)), + "_SUM_KERNEL": lambda: _make_kernel(lambda window, _: np.sum(window)), + "_MIN_KERNEL": lambda: _make_kernel(lambda window, _: np.min(window)), + "_MAX_KERNEL": lambda: _make_kernel(lambda window, _: np.max(window)), + "_PTP_KERNEL": lambda: _make_kernel(lambda window, _: np.max(window) - np.min(window)), + "_MEDIAN_KERNEL": lambda: _make_kernel(_median), + "_VAR_KERNEL": lambda: _make_kernel(_variance), + "_STD_KERNEL": lambda: _make_kernel( + lambda window, ddof: np.sqrt(_variance(window, ddof)) + ), # ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale onto # this one kernel rather than compiling a near-duplicate. -_QUANTILE_KERNEL = _make_kernel(lambda window, q: np.quantile(window, q)) + "_QUANTILE_KERNEL": lambda: _make_kernel(lambda window, q: np.quantile(window, q)), +} + +_KERNEL_CACHE = {} + + +def __getattr__(name): + """Builds a reduction kernel the first time it is named (PEP 562).""" + builder = _KERNEL_BUILDERS.get(name) + if builder is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + if name not in _KERNEL_CACHE: + _KERNEL_CACHE[name] = builder() + return _KERNEL_CACHE[name] def _as_quantile(q, scale: float): From ad1be8809ea9998e40527413d9242e0d451daf36 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:12:05 -0500 Subject: [PATCH 14/19] Give each benchmark shard its own HOME, and check nothing went missing Two bugs in the sharded track_ pass, both found by measuring it rather than by reading it. Every asv run finishes Machine.load with an unconditional MachineCollection.save back to ~/.asv-machine.json, and util.write_json opens that path 'w' -- truncating it before it writes. Concurrent shards race on the one file, and a shard that reads it inside another's truncate window dies with "Error parsing JSON in file ... Expecting value: line 1 column 1", taking its whole share of the suite with it. It cost two shards of four, in two runs out of three. asv resolves that path through os.path.expanduser, so each shard now gets a private HOME seeded with the machine file and .condarc, and the real one is left alone. Six-way runs go 3/3 clean after, against 1/3 before. The second is what made the first easy to miss: a shard that dies leaves a perfectly well-formed results file that is quietly missing benchmarks, and the run looks like it worked. The instances asked for are now checked against the instances that came back, a shortfall is an error rather than a warning, and the exit code follows. Tested against both a dropped parameter combination and a dropped benchmark. While measuring these: the 89.3s serial baseline quoted when this landed was a first run on cold numba and page caches. Warm, and on the same 12-core box, the serial track pass is 57.7s against 40s sharded four ways. Co-Authored-By: Claude Opus 5 --- benchmarks/helpers/_shard.py | 97 ++++++++++++++++++++++++++++++++++-- 1 file changed, 92 insertions(+), 5 deletions(-) diff --git a/benchmarks/helpers/_shard.py b/benchmarks/helpers/_shard.py index 7147d6b1e..151fc3331 100644 --- a/benchmarks/helpers/_shard.py +++ b/benchmarks/helpers/_shard.py @@ -74,7 +74,7 @@ import tempfile from pathlib import Path -__all__ = ["main", "merge_results", "plan_shards"] +__all__ = ["main", "merge_results", "missing_results", "plan_shards"] BENCHMARK_DIR = Path(__file__).resolve().parents[1] REPO_DIR = BENCHMARK_DIR.parent @@ -225,6 +225,31 @@ def cost(instances): return [s for s in shards if s] +def _shard_home(scratch, index): + """A private ``HOME`` for one shard, seeded with the files asv reads from it. + + Every ``asv run`` ends ``Machine.load`` with an unconditional + ``MachineCollection.save`` back to ``~/.asv-machine.json``, and + ``util.write_json`` opens that path ``'w'`` -- truncating it before it + writes. Concurrent shards therefore race on one file, and a shard that + reads it inside another's truncate window dies with "Error parsing JSON in + file ... Expecting value: line 1 column 1". Observed, repeatedly: it takes + out that shard's whole share of the suite. + + asv resolves the path through ``os.path.expanduser``, so a private ``HOME`` + per shard gives each its own copy to rewrite and leaves the real one alone. + ``.condarc`` comes along because asv shells out to conda for the + environment, and the machine file because that is the point. + """ + home = os.path.join(scratch, f"shard{index}", "home") + os.makedirs(home, exist_ok=True) + for name in (".asv-machine.json", ".condarc"): + source = os.path.join(os.path.expanduser("~"), name) + if os.path.isfile(source): + shutil.copyfile(source, os.path.join(home, name)) + return home + + def _shard_config(raw_conf, conf, results_dir, html_dir): """A copy of the project's asv config that writes results somewhere else. @@ -313,6 +338,46 @@ def merge_results(shard_dirs, dest_dir): return written +def missing_results(written, selected, benchmarks): + """Selected instances that came back with no result in the merged files. + + A shard that dies -- a timeout, an OOM, an asv that exits non-zero -- takes + its share of the suite with it, and the merge of what survives is a + perfectly well-formed results file that is quietly missing benchmarks. That + is the worst outcome available here, because the run looks like it worked. + So the instances asked for are checked against the instances that came back. + + A benchmark that ran and genuinely failed also records no result, and shows + up here too. Both mean the same thing to a caller: this run did not produce + the number. + """ + asv = _asv() + by_name = {b["name"]: b for b in benchmarks} + wanted = {} + for instance in selected: + benchmark_name = instance.split("(", 1)[0] + benchmark = by_name.get(benchmark_name) + if benchmark is None: + continue + names = _instance_names(benchmark) + wanted.setdefault(benchmark_name, []).append((names.index(instance), instance)) + + absent = [] + for path in written: + data = asv.util.load_json(path, api_version=asv.results.Results.api_version) + columns = data.get("result_columns") or [] + if "result" not in columns: + continue + index = columns.index("result") + for benchmark_name, entries in wanted.items(): + row = data.get("results", {}).get(benchmark_name) + values = (row[index] if row and len(row) > index else None) or [] + for param_idx, instance in entries: + if param_idx >= len(values) or values[param_idx] is None: + absent.append(instance) + return absent + + def _merge_row(a, b, columns): """Merges two asv result rows for the same benchmark.""" width = len(columns) @@ -495,6 +560,7 @@ def main(argv=None): scratch = tempfile.mkdtemp(prefix="asv-shard-") shard_dirs = [] processes = [] + incomplete = False try: for i, shard in enumerate(shards): shard_results = os.path.join(scratch, f"shard{i}", "results") @@ -531,6 +597,7 @@ def main(argv=None): env = dict(os.environ) env["NUMBA_NUM_THREADS"] = str(args.threads_per_shard) env["OMP_NUM_THREADS"] = str(args.threads_per_shard) + env["HOME"] = _shard_home(scratch, i) print(f" $ {shlex.join(command)}", flush=True) if not args.dry_run: @@ -538,13 +605,30 @@ def main(argv=None): codes = [p.wait() for p in processes] if any(codes): - print(f"warning: shard exit codes {codes}", file=sys.stderr, flush=True) + incomplete = True + print(f"error: shard exit codes {codes}", file=sys.stderr, flush=True) - # 3. Fold the shards back together. + # 3. Fold the shards back together, then make sure everything asked + # for actually came back -- see ``missing_results``. if not args.dry_run: print("merging shard results", flush=True) - for path in merge_results(shard_dirs, results_dir): + written = merge_results(shard_dirs, results_dir) + for path in written: print(f" wrote {path}", flush=True) + + absent = missing_results(written, selected, benchmarks) + if absent: + incomplete = True + print( + f"error: {len(absent)} of {len(selected)} selected instances " + "came back with no result:", + file=sys.stderr, + flush=True, + ) + for instance in absent: + print(f" {instance}", file=sys.stderr, flush=True) + else: + print(f" all {len(selected)} instances accounted for", flush=True) finally: if not args.dry_run: shutil.rmtree(scratch, ignore_errors=True) @@ -566,7 +650,10 @@ def main(argv=None): else: print(f" $ {shlex.join(serial)}", flush=True) - return 0 + # Non-zero rather than a warning buried in the log: a sharded run that lost + # a shard writes a well-formed results file with benchmarks missing from it, + # which is indistinguishable from a clean run unless the caller is told. + return 1 if incomplete else 0 if __name__ == "__main__": From 469854039c7a51c4cc50b2d53af2e9793137fda5 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:13:42 -0500 Subject: [PATCH 15/19] Skip warming the suite in interpreters that will not run it ASV drives one script in several modes, named by its first argument. `discover` imports every module in the suite to enumerate what is there and then calls none of it, but the modules warm their kernels at import so that forked benchmarks inherit them -- so discovery pays for warming it will never use. That is ~5s of the ~6s importing this suite costs, nearly all of it numba loading kernels, and it lands on the critical path of every `asv run`: once for the run, and once more for each concurrent shard. `setup_cache` and `check` are the same, though cheaper -- ASV imports one module for those rather than the tree. Measured warm, on a 12-core box: the discovery pass drops from 6.5s to 1.5s, and the fixed cost of an `asv run` from 13.1s to 7.8s. That second number is what caps sharding, since every shard pays it. Modes that do run something still warm, and so does anything argv cannot identify -- a REPL, a script -- so this only skips where nothing will be measured. Co-Authored-By: Claude Opus 5 --- benchmarks/helpers/_warmup.py | 45 ++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py index 21b68af85..866bf6bc4 100644 --- a/benchmarks/helpers/_warmup.py +++ b/benchmarks/helpers/_warmup.py @@ -9,11 +9,48 @@ launches the pool, and numba's OpenMP layer is not fork-safe. So this checks when a warmed kernel goes parallel, it becomes an import error rather than a benchmark that hangs on a cluster. + +Not every interpreter that imports the suite goes on to run something from it, +and the ones that do not should not pay to warm it -- see +:func:`will_run_benchmarks`. """ +import os +import sys + import numba -__all__ = ["warm_in_parent"] +__all__ = ["warm_in_parent", "will_run_benchmarks"] + +# ASV drives one script (``asv/benchmark.py``) in several modes, named by its +# first argument. These are the modes that import the suite without ever calling +# anything in it. +_NON_RUNNING_MODES = frozenset({"discover", "setup_cache", "check"}) + + +def will_run_benchmarks(): + """Whether this interpreter is going to execute a benchmark body. + + ``discover`` is the mode this exists for. It imports every module in the + suite to enumerate what is there, calls none of it, and sits on the critical + path of every ``asv run`` -- once for the run itself and once more for each + concurrent shard (:mod:`benchmarks.helpers._shard`). Warming it is pure + cost: measured at ~5s of the ~6s that importing this suite takes, nearly all + of it numba loading kernels that discovery will never call. + + ``setup_cache`` and ``check`` are the same story, though cheaper -- ASV + imports one module for those rather than the tree. + + The modes that do run something (``run``, ``run_server``, ``timing``) still + warm, and so does anything this cannot identify: an unrecognised argv means + warm, which is what happens when the suite is imported from a REPL or a + script. This only skips where it is certain nothing will be measured. + """ + return not ( + os.path.basename(sys.argv[0]) == "benchmark.py" + and len(sys.argv) > 1 + and sys.argv[1] in _NON_RUNNING_MODES + ) def _pool_running(): @@ -39,7 +76,13 @@ def warm_in_parent(warm, what): reports whatever it happened to warm as the cause of a pool that something else started earlier in the import, which sends you looking in the wrong module -- the two cases need different fixes, so they get different errors. + + A no-op in an interpreter that will not run a benchmark, per + :func:`will_run_benchmarks`. """ + if not will_run_benchmarks(): + return + already_running = _pool_running() warm() if not _pool_running(): From a48a45148571de27f0027ed298b355034120805a Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:13:42 -0500 Subject: [PATCH 16/19] Bring up the netCDF stack once, in the parent, not in every fork The first `xr.open_dataset` in a process spends ~164ms resolving xarray's backend entry points and loading the netCDF library, and only the first: reads after it cost single-digit milliseconds, 9ms even for the 18MB oQU120 grid. Once the fixture cache is warm nothing opens a netCDF file at import -- `preload_topologies` reads .npz through numpy -- so every forked benchmark that reached `cached_grid` or `cached_dataset` paid that start-up again on its own account. 134 of the ~235 benchmark processes in a run do, which is ~22s per commit spent bringing up the same library. `warm_netcdf` pays it once in the interpreter every benchmark is forked from, by opening the smallest grid in the repository -- so it works before the cache exists -- and closing it again. `preload_grids` and `preload_datasets` go further for the oQU meshes, ~22MB, which the forks then share copy-on-write instead of each holding a private read. The dyamond grids stay per-fork reads: that would be resident parent memory rather than transient. Also closes the netCDF handles in `_read_dataset` and `_build`. `load()` brings the arrays into memory but leaves the handle open in xarray's file cache, and on a cold cache `_build` runs at import -- so those handles were inherited by every forked benchmark, which HDF5 does not support. Serial track pass 57.7s -> 48.1s with the previous commit; four-way sharded 28.4s, six-way 26.0s. Values are unchanged: every track_nbytes_* and traced track_peakmem_* figure identical, the subprocess_peak_rss ones within 3%, which is their own run-to-run spread. Co-Authored-By: Claude Opus 5 --- benchmarks/face_bounds.py | 8 ++- benchmarks/helpers/_fixtures.py | 107 +++++++++++++++++++++++++++++--- benchmarks/mpas_ocean.py | 16 +++++ 3 files changed, 123 insertions(+), 8 deletions(-) diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py index f5d58ea47..0caf838df 100644 --- a/benchmarks/face_bounds.py +++ b/benchmarks/face_bounds.py @@ -1,8 +1,9 @@ import uxarray as ux -from .helpers._fixtures import GRIDS_BY_FORMAT, CachedFixtures +from .helpers._fixtures import GRIDS_BY_FORMAT, CachedFixtures, preload_grids from .helpers._memsize import grid_nbytes from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss +from .helpers._warmup import warm_in_parent grid_quad_hex = GRIDS_BY_FORMAT["ugrid-quad-hexagon"] grid_geoflow = GRIDS_BY_FORMAT["ugrid-geoflow"] @@ -87,3 +88,8 @@ def track_peakmem_open_and_bounds(self, grid_path): ) track_peakmem_open_and_bounds.unit = "bytes" + + +# Inherited by every forked benchmark rather than re-read in each; see the same +# note in ``mpas_ocean``. These four are the small ones, ~3.7MB all told. +warm_in_parent(lambda: preload_grids(GRIDS_BY_FORMAT.values()), "the source-format grids") diff --git a/benchmarks/helpers/_fixtures.py b/benchmarks/helpers/_fixtures.py index 2996be33a..99852457f 100644 --- a/benchmarks/helpers/_fixtures.py +++ b/benchmarks/helpers/_fixtures.py @@ -34,6 +34,8 @@ import uxarray as ux +from ._warmup import warm_in_parent + __all__ = [ "ALL_RESOLUTIONS", "DYAMOND_AVAILABLE", @@ -49,8 +51,11 @@ "cached_dataset", "cached_grid", "cached_topology", + "preload_datasets", + "preload_grids", "preload_topologies", "prime", + "warm_netcdf", ] BENCHMARK_DIR = Path(__file__).resolve().parents[1] @@ -159,8 +164,16 @@ def _write(dataset, artifact_path, writer): def _read_dataset(artifact_path): - """Reads back a cached ``xr.Dataset``.""" - return xr.open_dataset(artifact_path, mask_and_scale=False).load() + """Reads back a cached ``xr.Dataset``, holding no file handle afterwards. + + ``load()`` brings every array into memory, but on its own it leaves the + netCDF handle behind it open in xarray's file cache. Under + ``launch_method: forkserver`` a handle still open when the suite finishes + importing is inherited by every benchmark process forked from that + interpreter, and HDF5 does not support being used across a fork. + """ + with xr.open_dataset(artifact_path, mask_and_scale=False) as dataset: + return dataset.load() def _build(source): @@ -177,6 +190,13 @@ def _build(source): uxgrid._ds.load() # A ``UxDataset`` grid is cached separately. data_ds = uxds + uxds.close() + + # Everything is in memory by now, so the source files are done with. Closed + # rather than left to the garbage collector because on a cold cache this + # runs at import, in the interpreter every benchmark is forked from, and an + # inherited HDF5 handle is not something that library supports. + uxgrid._ds.close() _write( {name: getattr(uxgrid, name).data for name in ["node_lon", "node_lat", "face_node_connectivity"]}, @@ -225,13 +245,19 @@ def cached_grid(grid_path): 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") +def _cached_data_ds(grid_path, data_path): + """The cached dataset of a ``(grid, data)`` pair, held for this process.""" + artifact_path = _ensure((Path(grid_path), Path(data_path)), "data", ".nc") if artifact_path not in _loaded: _loaded[artifact_path] = _read_dataset(artifact_path) - return ux.UxDataset(_loaded[artifact_path].copy(), uxgrid=cached_grid(grid_path)) + return _loaded[artifact_path] + + +def cached_dataset(grid_path, data_path): + """A ``UxDataset`` over ``data_path``, on the cached grid.""" + return ux.UxDataset( + _cached_data_ds(grid_path, data_path).copy(), uxgrid=cached_grid(grid_path) + ) def prime(workers=1): @@ -301,6 +327,66 @@ def preload_topologies(grid_paths): return loaded +def preload_grids(grid_paths): + """Loads cached ``grid`` artifacts here so forked benchmarks inherit them. + + The counterpart of :func:`preload_topologies` for the flavor that carries + everything the reader produced, and it rests on the same thing: under + ``launch_method: forkserver`` a benchmark starts with whatever the + interpreter that imported the suite is holding, and reads its own copy of + anything else. + + Same caveat as well -- the arrays are shared rather than copied, so treat + them as read-only. ``cached_grid`` hands out a shallow copy of the dataset, + which is what keeps that safe. + + This costs the parent the size of every artifact named here for the length + of the run. Cheaper in total than the alternative, because the forks share + those pages copy-on-write instead of each holding a private read, but it is + resident rather than transient -- which is why the dyamond grids are left + out of the callers below and read per fork instead. + """ + loaded = 0 + for grid_path in grid_paths: + _cached_grid_ds(grid_path) # held by the process-level memo from here on + loaded += 1 + return loaded + + +def preload_datasets(sources): + """Loads cached ``(grid, data)`` artifacts, and their grids, for the same reason.""" + loaded = 0 + for grid_path, data_path in sources: + _cached_data_ds(grid_path, data_path) + _cached_grid_ds(grid_path) + loaded += 1 + return loaded + + +def warm_netcdf(): + """Pays netCDF4/HDF5's per-process start-up cost here, once, before any fork. + + The *first* ``xr.open_dataset`` in an interpreter spends ~164ms resolving + xarray's backend entry points and bringing up the netCDF library, and only + the first -- reads after it cost single-digit milliseconds even for the + 18MB oQU120 grid. Once the cache is warm nothing else in this module opens a + netCDF file at import (:func:`preload_topologies` goes through numpy), so + without this every forked benchmark that reaches ``cached_grid`` or + ``cached_dataset`` pays that start-up again on its own account. Roughly 134 + of the ~235 benchmark processes in a run do. + + Deliberately opens the smallest grid in the repository rather than a cached + artifact, so it still works before the cache exists, and closes it again so + nothing is inherited across the fork. + """ + try: + with xr.open_dataset(GRIDS_BY_FORMAT["ugrid-quad-hexagon"]) as dataset: + dataset.load() + except (KeyError, OSError, ValueError, RuntimeError): + # Nothing to warm is not a failure -- the forks just pay for themselves. + pass + + class CachedFixtures: """Mixin for benchmarks whose subject is not reading a file. @@ -326,6 +412,13 @@ def setup_cache(self): cached_dataset = staticmethod(cached_dataset) +# Run at import so the whole suite inherits an initialised netCDF stack. Through +# ``warm_in_parent`` for the two things it decides: skip this in the interpreters +# that only import the suite to enumerate it, and fail loudly if warming ever +# starts a numba thread pool a fork could not inherit. +warm_in_parent(warm_netcdf, "the netCDF backend") + + if __name__ == "__main__": # Fills the cache ahead of ``asv run``, so no benchmark -- and not even # ``setup_cache`` -- pays for reading a source grid. Worth a line in a batch diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 2136b2612..bc3719895 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -13,9 +13,12 @@ OQU_GRIDS, OQU_RESOLUTIONS, CachedFixtures, + preload_datasets, + preload_grids, ) from .helpers._memsize import grid_nbytes from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss +from .helpers._warmup import warm_in_parent data_var = 'bottomDepth' @@ -567,3 +570,16 @@ def track_peakmem_mean(self, resolution, chunking): return peak_allocated(lambda: self.nb.mean().compute()) track_peakmem_mean.unit = "bytes" + + +# Read at import rather than in ``setup``. ASV imports the suite once and forks +# every benchmark from that parent, so a fixture loaded here is inherited by all +# of them copy-on-write instead of being re-read by each -- and, more to the +# point, none of them has to bring up the netCDF stack for itself. Only the two +# oQU meshes, which come to ~22MB; the dyamond grids stay per-fork reads. +def _warm_parent(): + preload_grids(OQU_GRIDS.values()) + preload_datasets(OQU_DATASETS.values()) + + +warm_in_parent(_warm_parent, "the oQU fixtures") From e5676922fcddced2f10366410ad5f8390e21ca92 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:31:15 -0500 Subject: [PATCH 17/19] Prevent neighborhood filters from spawning threads early --- uxarray/grid/neighbors.py | 78 +++++++++++++++++---------------------- 1 file changed, 34 insertions(+), 44 deletions(-) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 7cf3fdf8b..1ae97f085 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -1285,44 +1285,34 @@ def _median(window, _): # ``Neighborhood`` method for it, so there is one place per reduction where its # kernel and parameter are chosen. # -# Each is built the first time it is named rather than at import. ``guvectorize`` -# with explicit signatures compiles at decoration time, and building a -# ``target="parallel"`` ufunc launches numba's OpenMP thread pool as part of -# that -- so assigning these at module scope made ``import uxarray`` start a -# thread pool before any user code ran. Besides charging every import for nine -# compiles it will usually not need, that leaves the process unsafe to fork: -# numba's OpenMP layer installs no at-fork handler, so a child inherits a pool -# it cannot use. Anything that forks workers after importing uxarray is affected -# -- multiprocessing's "fork" start method, and asv's ``forkserver`` launch -# method, which forks every benchmark from one interpreter that imported the -# suite. -_KERNEL_BUILDERS = { - "_MEAN_KERNEL": lambda: _make_kernel(lambda window, _: np.mean(window)), - "_SUM_KERNEL": lambda: _make_kernel(lambda window, _: np.sum(window)), - "_MIN_KERNEL": lambda: _make_kernel(lambda window, _: np.min(window)), - "_MAX_KERNEL": lambda: _make_kernel(lambda window, _: np.max(window)), - "_PTP_KERNEL": lambda: _make_kernel(lambda window, _: np.max(window) - np.min(window)), - "_MEDIAN_KERNEL": lambda: _make_kernel(_median), - "_VAR_KERNEL": lambda: _make_kernel(_variance), - "_STD_KERNEL": lambda: _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": lambda: _make_kernel(lambda window, q: np.quantile(window, q)), +# Built on first use rather than at import. ``guvectorize`` with explicit +# signatures compiles at decoration time, and building a ``target="parallel"`` +# ufunc launches numba's OpenMP thread pool as part of that. Assigning +# these at module scope made ``import uxarray`` start a thread pool before any +# user code ran. + +_REDUCERS = { + "mean": lambda window, _: np.mean(window), + "sum": lambda window, _: np.sum(window), + "min": lambda window, _: np.min(window), + "max": lambda window, _: np.max(window), + "ptp": lambda window, _: np.max(window) - np.min(window), + "median": _median, + "var": _variance, + "std": lambda window, ddof: np.sqrt(_variance(window, ddof)), + # ``percentile`` is ``quantile`` on a 0-100 scale, so both methods rescale + # onto this one kernel rather than compiling a near-duplicate. + "quantile": lambda window, q: np.quantile(window, q), } -_KERNEL_CACHE = {} +_KERNELS = {} -def __getattr__(name): - """Builds a reduction kernel the first time it is named (PEP 562).""" - builder = _KERNEL_BUILDERS.get(name) - if builder is None: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - if name not in _KERNEL_CACHE: - _KERNEL_CACHE[name] = builder() - return _KERNEL_CACHE[name] +def _kernel(reduction): + """The compiled kernel for ``reduction``, built once on first use.""" + if reduction not in _KERNELS: + _KERNELS[reduction] = _make_kernel(_REDUCERS[reduction]) + return _KERNELS[reduction] def _as_quantile(q, scale: float): @@ -1537,45 +1527,45 @@ def __repr__(self) -> str: def mean(self, uxda): """Mean of each neighborhood.""" - return self._apply_kernel(uxda, _MEAN_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("mean"), 0.0) def sum(self, uxda): """Sum of each neighborhood.""" - return self._apply_kernel(uxda, _SUM_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("sum"), 0.0) def min(self, uxda): """Smallest value in each neighborhood.""" - return self._apply_kernel(uxda, _MIN_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("min"), 0.0) def max(self, uxda): """Largest value in each neighborhood.""" - return self._apply_kernel(uxda, _MAX_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("max"), 0.0) def ptp(self, uxda): """Peak-to-peak spread (``max - min``) of each neighborhood.""" - return self._apply_kernel(uxda, _PTP_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("ptp"), 0.0) def median(self, uxda): """Median of each neighborhood.""" - return self._apply_kernel(uxda, _MEDIAN_KERNEL, 0.0) + return self._apply_kernel(uxda, _kernel("median"), 0.0) def var(self, uxda, ddof: int = 0): """Variance of each neighborhood, with ``ddof`` delta degrees of freedom.""" - return self._apply_kernel(uxda, _VAR_KERNEL, float(ddof)) + return self._apply_kernel(uxda, _kernel("var"), float(ddof)) def std(self, uxda, ddof: int = 0): """Standard deviation of each neighborhood, with ``ddof`` delta degrees of freedom.""" - return self._apply_kernel(uxda, _STD_KERNEL, float(ddof)) + return self._apply_kernel(uxda, _kernel("std"), float(ddof)) def quantile(self, uxda, q: float): """Quantile ``q`` (between 0 and 1) of each neighborhood.""" - return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 1.0)) + return self._apply_kernel(uxda, _kernel("quantile"), _as_quantile(q, 1.0)) def percentile(self, uxda, q: float): """Percentile ``q`` (between 0 and 100) of each neighborhood.""" - return self._apply_kernel(uxda, _QUANTILE_KERNEL, _as_quantile(q, 100.0)) + return self._apply_kernel(uxda, _kernel("quantile"), _as_quantile(q, 100.0)) def reduce(self, uxda, func: Callable): """Reduces each neighborhood with an arbitrary callable. From 4436f3ce40723ef3b122f51ddfe6c241328ec277 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:38:55 -0500 Subject: [PATCH 18/19] Report a pre-existing thread pool instead of refusing to benchmark Now that the neighborhood kernels are built lazily, the contender runs -- but the baseline is uxarray's main, which still builds them at import, so importing it starts an OpenMP pool. warm_in_parent raised on that, every module of the suite failed to import against that commit, every benchmark on it failed, and `asv continuous` compared the branch against nothing. Refusing does not help, which is the part I had wrong. The pool is started by importing uxarray, so it exists whether or not we warm; asv forks whether or not we raise; and the raise happens during an import asv runs with ignore_import_errors, so the parent keeps the pool and loses the module. The only thing it changed was that the commit could not be benchmarked at all. So a pool we did not start now reports itself once, on stderr, and the run continues. A pool that *warming* started is a different problem and still raises: that one is ours, it is preventable, and the fix is to move the work into setup. Warming still runs in the foreign-pool case, deliberately. Skipping it would leave forks on an unfixed commit compiling kernels that forks on a fixed commit inherit -- a difference between the two sides of a comparison that has nothing to do with the code being measured. Checked against uxarray at 644dd69d, which starts an omp pool on import: all nine benchmark modules import, against none before, and the notice is printed once rather than per module. Against this branch, where the import starts no pool, nothing is printed and warming is unchanged. Co-Authored-By: Claude Opus 5 --- benchmarks/helpers/_warmup.py | 51 ++++++++++++++++++++++++++--------- 1 file changed, 39 insertions(+), 12 deletions(-) diff --git a/benchmarks/helpers/_warmup.py b/benchmarks/helpers/_warmup.py index 866bf6bc4..96a63b3c2 100644 --- a/benchmarks/helpers/_warmup.py +++ b/benchmarks/helpers/_warmup.py @@ -67,15 +67,46 @@ def _pool_running(): return True +_REPORTED_FOREIGN_POOL = False + + +def _report_foreign_pool(what): + """Says once that this interpreter had a thread pool before we warmed anything.""" + global _REPORTED_FOREIGN_POOL + if _REPORTED_FOREIGN_POOL: + return + _REPORTED_FOREIGN_POOL = True + print( + f"asv: numba's thread pool was already running before {what} was warmed. " + "Something in the import of uxarray itself compiled or ran a parallel=True " + "kernel, so this interpreter was already unsafe to fork a numba benchmark " + "from -- warming is not the cause and cannot be the fix. Benchmarking " + "continues. To remove it, stop that kernel being built at import, or set " + "launch_method to spawn.", + file=sys.stderr, + ) + + def warm_in_parent(warm, what): - """Runs ``warm``, then fails if a numba thread pool is running afterwards. + """Runs ``warm``, then fails if *warming* started a numba thread pool. ``what`` names the thing being warmed, for the error message. - The pool is sampled before as well as after. Without that baseline this - reports whatever it happened to warm as the cause of a pool that something - else started earlier in the import, which sends you looking in the wrong - module -- the two cases need different fixes, so they get different errors. + The pool is sampled before as well as after, and the two cases are not the + same problem. A pool that warming started is ours and is preventable: move + that work into ``setup`` and the parent stays forkable. A pool that was + already running when we got here belongs to something in uxarray's own + import, and refusing to proceed does not help -- the pool exists either way, + asv forks either way, and raising only means the suite cannot benchmark that + commit at all. Which is what it did: every module failed to import against a + uxarray whose import builds a ``parallel=True`` kernel, so every benchmark + on that commit failed, and a run comparing against it compared against + nothing. So that case reports itself and carries on. + + Warming still happens in that case, deliberately. Skipping it would leave + the forks on such a commit compiling kernels that the forks on a fixed + commit inherit, which is a difference between the two sides of a comparison + that has nothing to do with the code being measured. A no-op in an interpreter that will not run a benchmark, per :func:`will_run_benchmarks`. @@ -89,13 +120,9 @@ def warm_in_parent(warm, what): return # nothing launched a pool, which is what makes this inheritable if already_running: - raise RuntimeError( - f"numba's thread pool was already running before {what} was warmed, " - "so this interpreter cannot safely fork a benchmark that uses numba. " - "Warming is not the cause: something earlier in the import of this " - "suite, or of uxarray itself, compiled or ran a parallel=True kernel. " - "Find that and move it, or drop launch_method: forkserver." - ) + _report_foreign_pool(what) + return + raise RuntimeError( f"warming {what} started numba's thread pool, which a forked benchmark " "cannot safely inherit -- warm it from setup() instead, now that these " From 1ec4c8e292116a8e433309bb88b9dd320df88418 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 24 Aug 2026 20:47:40 -0500 Subject: [PATCH 19/19] Spawn benchmarks rather than forking them, until main is fixed Forking is only safe while nothing has started numba's thread pool before the fork, and importing uxarray does exactly that on any commit whose neighbors.py builds its guvectorize kernels at module scope. This branch fixes that, but `asv continuous` installs the baseline commit too, and until the fix is in main every baseline -- and every historical commit before it -- imports with an OpenMP pool already running. The previous commit stops that being fatal, so those runs complete; it does not make the fork safe, it only declines to pretend otherwise. Spawning gives every benchmark a process that never forked, so the question does not arise on either side of a comparison. It costs what forking was buying: measured on bench_connectivity, 32 benchmark processes, 24s forked against 72s spawned -- roughly 1.5s per process, which over a full suite is minutes per commit. The config carries a TODO for putting it back, and how to know when: helpers/_warmup.py prints "asv: numba's thread pool was already running" from any process that starts with a pool it did not create, so a full run with no such line means forking is safe again. The parent-warming in the suite stays as it is. Under spawn it runs once per benchmark process instead of once per run, which is redundant with the setup methods rather than wrong -- _prepare and _warmup are both memoised per process, and asv does not time setup. Rewriting those comments now would only have to be undone when forkserver comes back. Co-Authored-By: Claude Opus 5 --- benchmarks/asv.conf.json | 32 +++++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/benchmarks/asv.conf.json b/benchmarks/asv.conf.json index 535ec6b2a..b664dc971 100644 --- a/benchmarks/asv.conf.json +++ b/benchmarks/asv.conf.json @@ -61,14 +61,32 @@ // 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 + // A fresh interpreter per benchmark, rather than forking each one from an + // interpreter that has already imported the suite. + // + // "forkserver" is what this suite is built for: it saves the ~0.9s uxarray + // import and the numba kernel load on every one of ~250 benchmark // processes, and lets a fixture loaded at import be inherited rather than - // re-read. 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", + // re-read. Measured on bench_connectivity alone, 32 benchmark processes: + // 24s forked against 72s spawned, so roughly 1.5s per process. + // + // It is "spawn" anyway because forking is only safe while nothing has + // started numba's thread pool before the fork, and importing uxarray does + // exactly that on any commit whose neighbors.py builds its guvectorize + // kernels at module scope -- guvectorize with explicit signatures compiles + // at decoration time, and building a target="parallel" ufunc launches the + // pool. This branch fixes that, but `asv continuous` installs the baseline + // commit too, and until the fix is in main every baseline -- and every + // historical commit before it -- imports with an OpenMP pool already + // running. Spawning gives every benchmark a process that never forked, so + // the question does not arise on either side of the comparison. + // + // TODO: put this back to "forkserver" once the lazy-kernel fix is in main + // and the commits being compared against carry it. The suite says when it + // is not safe yet: helpers/_warmup.py prints "asv: numba's thread pool was + // already running" from any process that starts with a pool it did not + // create. No such line over a full run means forking is safe again. + "launch_method": "spawn", // ``benchmark_timeout`` is not a key asv reads -- the one that sets the // default is ``default_benchmark_timeout`` (asv/config.py)