From d1f55b72202101f8329a6d3a36178699a43afc1c Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Fri, 7 Aug 2026 08:33:21 -0500 Subject: [PATCH] Add the DYAMOND resolution-ladder benchmark Runs one fixed request pipeline across 30/15/7.5/3.75 km MPAS output, a 64x range in cell count up to 41,943,042 faces, at a Globus Compute endpoint because the largest grid is 19.8 GB. Separates JIT compilation, per-grid bounds construction, and the per-request reduction, since collapsing them reports a 59.5 s first call as the cost of a 1.4 s operation. Uses both face-centred and node-centred fields so one request exercises two mesh locations. --- benchmarks/dyamond/README.md | 51 ++++++ benchmarks/dyamond/data_ladder_results.json | 190 ++++++++++++++++++++ benchmarks/dyamond/run_data_ladder.py | 167 +++++++++++++++++ 3 files changed, 408 insertions(+) create mode 100644 benchmarks/dyamond/README.md create mode 100644 benchmarks/dyamond/data_ladder_results.json create mode 100644 benchmarks/dyamond/run_data_ladder.py diff --git a/benchmarks/dyamond/README.md b/benchmarks/dyamond/README.md new file mode 100644 index 0000000..6685c42 --- /dev/null +++ b/benchmarks/dyamond/README.md @@ -0,0 +1,51 @@ +# DYAMOND resolution-ladder benchmark + +Measures one fixed request pipeline across the DYAMOND-1 MPAS resolution +ladder (30 / 15 / 7.5 / 3.75 km), a 64x range in cell count from 655,362 to +41,943,042 faces. + +This ladder is the same one UXarray's own ASV performance suite tracks, so +the axis is comparable with prior work rather than particular to this repo. +Unlike a bare-mesh benchmark it uses real model output: each diagnostic file +carries face-centred fields (for example `t2m`) and node-centred fields (for +example `vorticity_200hPa`), so one request exercises two mesh locations. + +## Running + +The data live on NSF NCAR GLADE and the largest grid is 19.8 GB, so this runs +at the endpoint rather than locally: + +```bash +uv run python benchmarks/dyamond/run_data_ladder.py # all four rungs +uv run python benchmarks/dyamond/run_data_ladder.py 30km 15km # a subset +``` + +Results append to `data_ladder_results.json` after every rung, so an +interrupted run keeps the rungs it already finished. + +## Reading the numbers + +Three costs are deliberately separated, because collapsing them is misleading: + +- **JIT compilation** — first call only, per worker deployment. UXarray's + spherical-geometry kernels are `numba`-jitted with `cache=True`, so the + first zonal reduction on a fresh worker can pay tens of seconds that no + later call pays. +- **Bounds construction** — once per grid, scales with cell count. The first + zonal reduction triggers it. +- **The reduction itself** — once per request. + +Our first 30 km run reported 59.5 s for the 1-degree zonal mean; the same call +afterwards took 1.4 s. Reporting the cold number as "the cost of a zonal mean" +would misinform every downstream scheduling decision. + +Timings are single-worker on one Casper node with numba resolving to two +threads, so they describe an unoptimized serial path and are upper bounds. + +## Reproducibility + +`data_ladder_results.json` is the run the paper table cites. A later repeat of +the 30 km rung on the same endpoint gave 7.98 s bounds and 1.33 s warm +1-degree zonal against 7.46 s and 1.40 s, so the run-to-run spread on a shared +node is a few percent — well inside the order-of-magnitude effects the table +is about. diff --git a/benchmarks/dyamond/data_ladder_results.json b/benchmarks/dyamond/data_ladder_results.json new file mode 100644 index 0000000..d1eba6b --- /dev/null +++ b/benchmarks/dyamond/data_ladder_results.json @@ -0,0 +1,190 @@ +[ + { + "resolution": "30km", + "grid_bytes": 403733477, + "data_bytes": 124750934, + "hostname": "casper04", + "python": "3.11.12", + "uxarray": "2026.7.0", + "ops": { + "open": 2.9407250322401524, + "load_face": 0.005672751925885677, + "mean_face": 0.00025865621864795685, + "bounds_cold": 7.460680991411209, + "zonal_1deg": 1.4338232101872563, + "zonal_1deg_warm": 1.398980274796486, + "zonal_5deg": 0.8103886498138309, + "zonal_5deg_warm": 0.8116860631853342, + "subset_bbox": 0.20600672252476215, + "zonal_on_subset": 0.03303864039480686, + "load_node": 0.015991318970918655 + }, + "n_face": 655362, + "n_node": 1310720, + "n_values_face": 655362, + "mean_t2m": 289.3864440917969, + "zonal_bins": 181, + "zonal_min": 218.15574645996094, + "zonal_max": 300.7034606933594, + "zonal5_bins": 37, + "subset_faces": 17667, + "n_values_node": 1310720, + "mean_vort": 1.8413867763911185e-08, + "_provenance": { + "tool": "remote_workflow", + "inputs": { + "args": [ + "30km" + ] + }, + "execution_venue": "hpc:ucar-uxarray-yac", + "timestamp_utc": "2026-08-07T11:02:57.637791+00:00", + "uxarray_version": "2026.7.0", + "python_version": "3.12.10", + "warnings": [], + "artifacts": [] + }, + "_wall_s": 22.495638847351074 + }, + { + "resolution": "15km", + "grid_bytes": 1614838757, + "data_bytes": 996507874, + "hostname": "casper04", + "python": "3.11.12", + "uxarray": "2026.7.0", + "ops": { + "open": 14.362089494243264, + "load_face": 0.23801299277693033, + "mean_face": 0.0009028185158967972, + "bounds_cold": 29.79861798323691, + "zonal_1deg": 4.87471757736057, + "zonal_1deg_warm": 4.763765844516456, + "zonal_5deg": 3.1069262735545635, + "zonal_5deg_warm": 3.1131048342213035, + "subset_bbox": 0.9503179062157869, + "zonal_on_subset": 0.09335884638130665, + "load_node": 0.5422558458521962 + }, + "n_face": 2621442, + "n_node": 5242880, + "n_values_face": 2621442, + "mean_t2m": 289.4427795410156, + "zonal_bins": 181, + "zonal_min": 218.10801696777344, + "zonal_max": 300.7080383300781, + "zonal5_bins": 37, + "subset_faces": 72052, + "n_values_node": 5242880, + "mean_vort": 7.081877306802653e-09, + "_provenance": { + "tool": "remote_workflow", + "inputs": { + "args": [ + "15km" + ] + }, + "execution_venue": "hpc:ucar-uxarray-yac", + "timestamp_utc": "2026-08-07T11:05:04.410536+00:00", + "uxarray_version": "2026.7.0", + "python_version": "3.12.10", + "warnings": [], + "artifacts": [] + }, + "_wall_s": 79.84460783004761 + }, + { + "resolution": "7.5km", + "grid_bytes": 4949311375, + "data_bytes": 1730121739, + "hostname": "casper04", + "python": "3.11.12", + "uxarray": "2026.7.0", + "ops": { + "open": 44.46716747339815, + "load_face": 0.19585166964679956, + "mean_face": 0.0035475092008709908, + "bounds_cold": 118.04297504108399, + "zonal_1deg": 15.274628319777548, + "zonal_1deg_warm": 12.797916687093675, + "zonal_5deg": 8.3560239225626, + "zonal_5deg_warm": 8.359011514112353, + "subset_bbox": 5.132943272590637, + "zonal_on_subset": 0.24271893221884966, + "load_node": 0.46354250609874725 + }, + "n_face": 10485762, + "n_node": 20971520, + "n_values_face": 10485762, + "mean_t2m": 289.4425964355469, + "zonal_bins": 181, + "zonal_min": 218.1250762939453, + "zonal_max": 300.7069091796875, + "zonal5_bins": 37, + "subset_faces": 289498, + "n_values_node": 20971520, + "mean_vort": 7.926760581256076e-09, + "_provenance": { + "tool": "remote_workflow", + "inputs": { + "args": [ + "7.5km" + ] + }, + "execution_venue": "hpc:ucar-uxarray-yac", + "timestamp_utc": "2026-08-07T11:09:25.192818+00:00", + "uxarray_version": "2026.7.0", + "python_version": "3.12.10", + "warnings": [], + "artifacts": [] + }, + "_wall_s": 260.78248381614685 + }, + { + "resolution": "3.75km", + "grid_bytes": 19797147535, + "data_bytes": 7721009652, + "hostname": "casper04", + "python": "3.11.12", + "uxarray": "2026.7.0", + "ops": { + "open": 174.8197883674875, + "load_face": 0.7597603499889374, + "mean_face": 0.014804086647927761, + "bounds_cold": 487.847321793437, + "zonal_1deg": 535.124601630494, + "zonal_1deg_warm": 528.600348033011, + "zonal_5deg": 486.91943696700037, + "zonal_5deg_warm": 489.1833068523556, + "subset_bbox": 25.208245809189975, + "zonal_on_subset": 1.0073513984680176, + "load_node": 1.270337413996458 + }, + "n_face": 41943042, + "n_node": 83886080, + "n_values_face": 41943042, + "mean_t2m": 289.44281005859375, + "zonal_bins": 181, + "zonal_min": 218.12278747558594, + "zonal_max": 300.70709228515625, + "zonal5_bins": 37, + "subset_faces": 1160500, + "n_values_node": 83886080, + "mean_vort": 8.223549841090971e-09, + "_provenance": { + "tool": "remote_workflow", + "inputs": { + "args": [ + "3.75km" + ] + }, + "execution_venue": "hpc:ucar-uxarray-yac", + "timestamp_utc": "2026-08-07T13:18:50.157103+00:00", + "uxarray_version": "2026.7.0", + "python_version": "3.12.10", + "warnings": [], + "artifacts": [] + }, + "_wall_s": 7756.7684071063995 + } +] diff --git a/benchmarks/dyamond/run_data_ladder.py b/benchmarks/dyamond/run_data_ladder.py new file mode 100644 index 0000000..a426ecc --- /dev/null +++ b/benchmarks/dyamond/run_data_ladder.py @@ -0,0 +1,167 @@ +"""Multi-operation DYAMOND workflow across the resolution ladder. + +Answers the question the paper actually needs: when an agent is asked +"subset this region and give me the zonal mean", what does each step cost, +and how does that cost scale with mesh resolution and with field location +(face-centred vs node-centred)? + +Runs on the NCAR Globus Compute endpoint because the 3.75 km grid alone is +19.8 GB. Data are DYAMOND-1 (Judt et al.), the same archive UXarray's ASV +suite uses, so the axis is comparable with prior work. +""" + +from __future__ import annotations + +import json +import sys +import time +import warnings + +warnings.filterwarnings("ignore") + +GRID = "/glade/campaign/cisl/vast/uxarray/data/dyamond/{res}/grid.nc" +DATA = ( + "/glade/campaign/mmm/wmr/fjudt/projects/dyamond_1/{res}/diag.2016-08-01_00.00.00.nc" +) +RES = ["30km", "15km", "7.5km", "3.75km"] +EP = "ucar-uxarray-yac" + + +def remote_workflow(res: str) -> dict: + """Self-contained: runs on the worker, no uxarray_mcp import.""" + import os + import platform + import socket + import time + + import numpy as np + import uxarray as ux + + grid = "/glade/campaign/cisl/vast/uxarray/data/dyamond/%s/grid.nc" % res + data = ( + "/glade/campaign/mmm/wmr/fjudt/projects/dyamond_1/%s/" + "diag.2016-08-01_00.00.00.nc" % res + ) + + out = { + "resolution": res, + "grid_bytes": os.path.getsize(grid), + "data_bytes": os.path.getsize(data), + "hostname": socket.gethostname(), + "python": platform.python_version(), + "uxarray": ux.__version__, + "ops": {}, + } + + def timed(key, fn): + t0 = time.perf_counter() + r = fn() + out["ops"][key] = time.perf_counter() - t0 + return r + + # 1. open grid + data together (what a user request actually triggers) + uxds = timed("open", lambda: ux.open_dataset(grid, data)) + g = uxds.uxgrid + out["n_face"] = int(g.n_face) + out["n_node"] = int(g.n_node) + + # 2. face-centred field: load values into memory + t2m = uxds["t2m"].isel(time=0) + timed("load_face", lambda: t2m.values) + out["n_values_face"] = int(t2m.size) + + # 3. global mean of the face field (cheap reduction over all values) + out["mean_t2m"] = float(timed("mean_face", lambda: t2m.values.mean())) + + # 4. one-time structural cost: spherical bounding box per face. + # The first zonal mean pays for this; every later one does not. + # Isolating it is the difference between a 60 s and a 0.8 s answer. + + timed("bounds_cold", lambda: g.bounds) + + # 5. zonal mean with bounds already built -- the true per-request cost. + zm = timed("zonal_1deg", lambda: t2m.zonal_mean(lat=(-90, 90, 1))) + zmv = np.asarray(zm.values) + out["zonal_bins"] = int(zmv.size) + out["zonal_min"] = float(np.nanmin(zmv)) + out["zonal_max"] = float(np.nanmax(zmv)) + reps = [] + for _ in range(3): + t0 = time.perf_counter() + t2m.zonal_mean(lat=(-90, 90, 1)) + reps.append(time.perf_counter() - t0) + out["ops"]["zonal_1deg_warm"] = float(np.median(reps)) + + zm5 = timed("zonal_5deg", lambda: t2m.zonal_mean(lat=(-90, 90, 5))) + out["zonal5_bins"] = int(np.asarray(zm5.values).size) + reps5 = [] + for _ in range(3): + t0 = time.perf_counter() + t2m.zonal_mean(lat=(-90, 90, 5)) + reps5.append(time.perf_counter() - t0) + out["ops"]["zonal_5deg_warm"] = float(np.median(reps5)) + + # 6. regional subset: a tropical-cyclone-sized box in the W Pacific + sub = timed("subset_bbox", lambda: t2m.subset.bounding_box((120, 160), (0, 30))) + out["subset_faces"] = int(sub.uxgrid.n_face) + + # 7. zonal mean on the subset -- cost after data reduction + timed("zonal_on_subset", lambda: sub.zonal_mean(lat=(0, 30, 1))) + + # 8. node-centred field: vorticity lives on vertices, not cells + vort = uxds["vorticity_200hPa"].isel(time=0) + timed("load_node", lambda: vort.values) + out["n_values_node"] = int(vort.size) + out["mean_vort"] = float(np.asarray(vort.values).mean()) + + uxds.close() + return out + + +def main() -> int: + which = sys.argv[1:] or RES + import asyncio + + from uxarray_mcp.remote.agent import UXarrayComputeAgent + from uxarray_mcp.remote.config import load_config + + cfg = load_config() + outfile = "benchmarks/dyamond/data_ladder_results.json" + try: + results = json.load(open(outfile)) + except Exception: + results = [] + results = [r for r in results if r.get("resolution") not in which] + + for res in which: + print("=== %s ===" % res, flush=True) + t0 = time.time() + try: + prof = cfg.for_endpoint(EP) + prof.timeout_seconds = 3000 # big rungs need it + agent = UXarrayComputeAgent(prof) + r = asyncio.run(agent._run_on_hpc(remote_workflow, res)) + r["_wall_s"] = time.time() - t0 + results.append(r) + print(json.dumps(r.get("ops", {}), indent=2), flush=True) + print("faces=%s wall=%.1fs" % (r.get("n_face"), r["_wall_s"]), flush=True) + except Exception as e: + print( + "%s FAILED after %.1fs: %s: %s" + % (res, time.time() - t0, type(e).__name__, str(e)[:400]), + flush=True, + ) + results.append( + { + "resolution": res, + "error": "%s: %s" % (type(e).__name__, str(e)[:400]), + } + ) + with open(outfile, "w") as f: + json.dump(results, f, indent=2, default=str) + print("-> wrote %s" % outfile, flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())