Cached benchmark file I/O

.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 }})"

.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

.gitignore

@@ -163,3 +163,4 @@ docs/user-guide/psi_healpix.nc
 benchmarks/env
 benchmarks/results
 benchmarks/html
+benchmarks/_io_cache

benchmarks/asv.conf.json

@@ -61,7 +61,18 @@
     // defaults to 10 min
     "install_timeout": 600,
 
-    "benchmark_timeout": 360,
+    // Fork each benchmark from one interpreter that has already imported the
+    // suite, instead of starting a fresh one per benchmark. Saves the ~0.9s
+    // uxarray import and the numba kernel load on every one of ~160 benchmark
+    // processes, and lets a fixture loaded at import be inherited rather than
+    // re-read. asv leaves this at "spawn" by default because fork and threads
+    // mix badly; nothing in this suite runs a parallel kernel at import time,
+    // which is what makes it safe here -- keep it that way.
+    "launch_method": "forkserver",
+
+    // ``benchmark_timeout`` is not a key asv reads -- the one that sets the
+    // default is ``default_benchmark_timeout`` (asv/config.py)
+    "default_benchmark_timeout": 360,
 
     // the base URL to show a commit for the project.
     "show_commit_url": "https://github.com/UXARRAY/uxarray/commit/",

benchmarks/bench_connectivity.py

@@ -1,53 +1,25 @@
-import os
-import urllib.request
-from pathlib import Path
-
 import uxarray as ux
 
-current_path = Path(os.path.dirname(os.path.realpath(__file__)))
-
-grid_filename_480 = "oQU480.grid.nc"
-grid_filename_120 = "oQU120.grid.nc"
-filenames = [grid_filename_480, grid_filename_120]
-
-for filename in filenames:
-    if not os.path.isfile(current_path / filename):
-        # downloads the files from Cookbook repo, if they haven't been downloaded locally yet
-        url = f"https://github.com/ProjectPythia/unstructured-grid-viz-cookbook/raw/main/meshfiles/{filename}"
-        _, headers = urllib.request.urlretrieve(url, filename=current_path / filename)
-
-oQU_path_dict = {"480km": current_path / grid_filename_480,
-                  "120km": current_path / grid_filename_120}
-
-# Paths to grid files on Glade
-dyamond_path_dict = {"30km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/30km/grid.nc",
-                  "15km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/15km/grid.nc",
-                  "7.5km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/7.5km/grid.nc",
-                  "3.75km": "/glade/campaign/cisl/vast/uxarray/data/dyamond/3.75km/grid.nc"}
-
-# Determines if all file paths exist and are accesible
-all_paths_exist = True
-for file_path in dyamond_path_dict.values():
-    all_paths_exist = all_paths_exist and os.path.exists(file_path)
-
-file_path_dict = oQU_path_dict
-if all_paths_exist:
-    file_path_dict = file_path_dict | dyamond_path_dict
+from .helpers._fixtures import (
+    ALL_RESOLUTIONS,
+    GRIDS_BY_RESOLUTION,
+    CachedFixtures,
+    cached_topology,
+    preload_topologies,
+)
+from .helpers._warmup import warm_in_parent
 
 
-class GridBenchmark:
+class GridBenchmark(CachedFixtures):
     """Class used as a template for benchmarks requiring a ``Grid`` in this
     module across both resolutions."""
-    param_names = ['resolution', ]
 
-    # Conditionally available; could get annoying if there are downstream tools relying on it.
-    if all_paths_exist:
-        params = [['480km', '120km', '30km', '15km', '7.5km', '3.75km'], ]
-    else:
-        params = [['480km', '120km'], ]
+    param_names = ['resolution', ]
+    params = [ALL_RESOLUTIONS, ]
+    timeout = 1200
 
     def setup(self, resolution, *args, **kwargs):
-        self.uxgrid = ux.open_grid(file_path_dict[resolution])
+        self.uxgrid = self.cached_grid(GRIDS_BY_RESOLUTION[resolution])
 
     def teardown(self, resolution, *args, **kwargs):
         del self.uxgrid
@@ -65,40 +37,36 @@ def teardown(self, resolution, *args, **kwargs):
 
 _numba_warmed_up = False
 
-def _warmup(uxgrid):
+def _warmup():
     """Compiles the Numba kernels backing each connectivity variable.
 
-    ``_build_node_edge_connectivity`` is not disk-cached, so a fresh benchmark
-    process would otherwise charge ~240ms of JIT compilation to whichever sample
-    happened to touch it first.
+    ``_build_node_edge_connectivity`` is ``@njit`` without ``cache=True``, so a
+    fresh benchmark process would otherwise charge ~240ms of JIT compilation to
+    whichever sample happened to touch it first.
     """
     global _numba_warmed_up
     if _numba_warmed_up:
         return
+    uxgrid = ux.Grid.from_topology(*cached_topology(GRIDS_BY_RESOLUTION[ALL_RESOLUTIONS[0]]))
     for name in CONNECTIVITY_NAMES:
         getattr(uxgrid, name)
     _numba_warmed_up = True
 
 
 class Connectivity(GridBenchmark):
-    # Each connectivity variable is cached in ``Grid._ds`` once constructed, so a
-    # sample may only contain a single call; otherwise every call but the first
-    # would time a dictionary lookup.
+    # connectivity is cached in ``Grid._ds`` on construction, so only run them once
     number = 1
 
     def setup(self, resolution, *args, **kwargs):
         # The benchmark grids are MPAS meshes, which carry every connectivity
         # variable on disk. Reading one would time the MPAS parser rather than
-        # the construction routines, so reduce the grid down to the minimal
-        # UGRID topology and let each variable be built on demand.
-        source_grid = ux.open_grid(file_path_dict[resolution])
-        self.topology = (
-            source_grid.node_lon.data,
-            source_grid.node_lat.data,
-            source_grid.face_node_connectivity.data,
-        )
-
-        _warmup(self.minimal_grid())
+        # the construction routines, so this takes the minimal UGRID topology
+        # fixture and lets each variable be built on demand.
+        self.topology = self.cached_topology(GRIDS_BY_RESOLUTION[resolution])
+
+        # A no-op once the module-level warm below has run; kept so the class is
+        # still correct if that ever goes away.
+        _warmup()
         self.uxgrid = self.minimal_grid()
 
     def minimal_grid(self):
@@ -139,3 +107,15 @@ def time_edge_face(self, resolution):
 
     def time_node_face(self, resolution):
         _ = self.uxgrid.node_face_connectivity.compute()
+
+
+# Compiled at import rather than in ``setup``. ASV imports the suite once and forks
+# every benchmark from that parent, so kernels compiled here are inherited by all
+# of them. Only safe while the connectivity kernels are serial
+def _warm_parent():
+    _warmup()
+    # And the topologies themselves...
+    preload_topologies(GRIDS_BY_RESOLUTION[res] for res in ALL_RESOLUTIONS)
+
+
+warm_in_parent(_warm_parent, "the connectivity kernels")

benchmarks/face_bounds.py

@@ -1,18 +1,16 @@
-import os
-from pathlib import Path
-
 import uxarray as ux
-from .helpers._memsize import grid_nbytes
-from .helpers._peakmem import numba_threads, peak_allocated
 
-current_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[0]
+from .helpers._fixtures import GRIDS_BY_FORMAT, CachedFixtures, preload_grids
+from .helpers._memsize import grid_nbytes
+from .helpers._peakmem import numba_threads, peak_allocated, subprocess_peak_rss
+from .helpers._warmup import warm_in_parent
 
-grid_quad_hex = current_path / "test" / "meshfiles" / "ugrid" / "quad-hexagon" / "grid.nc"
-grid_geoflow = current_path / "test" / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc"
-grid_scrip = current_path / "test" / "meshfiles" / "scrip" / "outCSne8" / "outCSne8.nc"
-grid_mpas= current_path / "test" / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc"
+grid_quad_hex = GRIDS_BY_FORMAT["ugrid-quad-hexagon"]
+grid_geoflow = GRIDS_BY_FORMAT["ugrid-geoflow"]
+grid_scrip = GRIDS_BY_FORMAT["scrip-outCSne8"]
+grid_mpas = GRIDS_BY_FORMAT["mpas-oQU480"]
 
-class FaceBounds:
+class FaceBounds(CachedFixtures):
 
     params = [grid_quad_hex, grid_geoflow, grid_scrip, grid_mpas]
 
@@ -24,8 +22,8 @@ def setup(self, grid_path):
         # compiled before anything is measured. ``track_peakmem_*`` would
         # otherwise charge the first sample for loading them off numba's disk
         # cache, which inflates the reported peak by ~3%.
-        ux.open_grid(grid_quad_hex).bounds
-        self.uxgrid = ux.open_grid(grid_path)
+        self.cached_grid(grid_quad_hex).bounds
+        self.uxgrid = self.cached_grid(grid_path)
 
     def teardown(self, n):
         del self.uxgrid
@@ -63,21 +61,35 @@ def track_peakmem_face_bounds(self, grid_path):
 class FaceBoundsColdStartRss:
     """Peak memory of a cold start: import uxarray, open a grid, get its bounds.
 
-    Whole-process ``ru_maxrss``, not tracemalloc -- the ~250MB uxarray import is
-    part of the number by design, because the cold start is the subject. For the
-    cost of ``bounds`` alone see ``FaceBounds.track_peakmem_face_bounds``, which
-    runs one to three orders of magnitude lower.
+    Whole-process peak resident memory, not tracemalloc -- the ~250MB uxarray
+    import is part of the number by design, because the cold start is the
+    subject. For the cost of ``bounds`` alone see
+    ``FaceBounds.track_peakmem_face_bounds``, which runs one to three orders of
+    magnitude lower.
+
+    Measured in a subprocess of its own rather than through asv's ``peakmem_*``,
+    which reports ``ru_maxrss`` for the benchmark process.
     """
 
     params = FaceBounds.params
     param_names = ["grid_path"]
 
     def setup_cache(self):
-        """Compile the njit kernels before anything is measured."""
+        """Compile the njit kernels before anything is measured.
+        """
         for grid_path in self.params:
             ux.open_grid(grid_path).bounds
 
     setup_cache.timeout = 1800
 
-    def peakmem_open_and_bounds(self, grid_path):
-        ux.open_grid(grid_path).bounds
+    def track_peakmem_open_and_bounds(self, grid_path):
+        return subprocess_peak_rss(
+            f"import uxarray as ux; ux.open_grid({str(grid_path)!r}).bounds"
+        )
+
+    track_peakmem_open_and_bounds.unit = "bytes"
+
+
+# Inherited by every forked benchmark rather than re-read in each; see the same
+# note in ``mpas_ocean``. These four are the small ones, ~3.7MB all told.
+warm_in_parent(lambda: preload_grids(GRIDS_BY_FORMAT.values()), "the source-format grids")

benchmarks/geometry_kernels.py

@@ -15,6 +15,8 @@
 
 import numpy as np
 
+from .helpers._warmup import warm_in_parent
+
 
 def _unit(v):
     return v / np.linalg.norm(v)
@@ -205,3 +207,25 @@ def time_try_gca_const_lat_intersection(self):
     def time_gca_const_lat_intersection(self):
         """Layer 3: dispatcher (full public API)."""
         self.gca_const_lat_intersection(self.gca_cart, _CONST_Z)
+
+
+def _warm_classes():
+    """Compiles what each class's ``setup`` compiles, once per process.
+
+    Runs the setups themselves rather than a copy of their warm calls, so this
+    cannot drift out of step with them.
+    """
+    for cls in (
+        EFTPrimitives,
+        AccucrossKernels,
+        OrientPredicates,
+        GCAGCAIntersection,
+        GCAConstLatIntersection,
+    ):
+        try:
+            cls().setup()
+        except Exception:
+            pass
+
+
+warm_in_parent(_warm_classes, "the geometry kernels")

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()