From 7b22d2336292ca9ba84a45b52f2f7f450a0afed8 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Sat, 4 Jul 2026 03:33:55 +0000 Subject: [PATCH 01/16] Add insitubatch-backed initial-condition feed New optional data module `earth2studio.data.insitu` that feeds initial conditions to `earth2studio.run` workflows without going through xarray. - `batch_to_xcoords`: pure converter from an insitubatch numpy `Batch` to the exact `fetch_data(legacy=True)` contract -- an `(time, lead_time, variable, lat, lon)` tensor plus the matching 5-key `CoordSystem` OrderedDict -- so it drops straight into `prognostic.create_iterator` after `map_coords`. - `InSituForecastFeed`: prefetched IC iterator over a contiguous window of an analysis store. Reads with insitubatch (bounded-fan-out async prefetch and a read plan that de-duplicates chunks shared across init times) instead of the per-`(time, variable)` `DataSource -> xr.DataArray -> fetch_data` path, whose gather is unbounded and re-reads overlapping chunks. Emits a single 0 h lead (an initial condition) today; multi-step history and verification-lead offsets are the next step. insitubatch is imported lazily and only required when this module is used. --- earth2studio/data/insitu.py | 153 ++++++++++++++++++++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 earth2studio/data/insitu.py diff --git a/earth2studio/data/insitu.py b/earth2studio/data/insitu.py new file mode 100644 index 000000000..512112f7f --- /dev/null +++ b/earth2studio/data/insitu.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""insitubatch-backed initial-condition / verification feed for Earth2Studio. + +Earth2Studio's standard path is ``DataSource -> xr.DataArray -> fetch_data -> +prep_data_array -> (torch.Tensor, coords)``; xarray is load-bearing down to +``prep_data_array``. For an IO-bound hindcast / scoring campaign over many init times, that +per-``(time, variable)`` fetch re-reads overlapping chunks and its ``gather`` is unbounded. + +This module skips xarray entirely: it reads the analysis store with **insitubatch** +(`InSituDataset` -- bounded-fan-out async prefetch, a read plan that de-duplicates chunks +shared across init times and lead offsets) and converts each numpy ``Batch`` to the exact +``(x, coords)`` tuple ``fetch_data(..., legacy=True)`` returns, so it is a drop-in for the +initial-condition feed of ``earth2studio.run`` workflows. See :func:`batch_to_xcoords` for +the contract and :class:`InSituForecastFeed` for the prefetched iterator. +""" + +from collections import OrderedDict + +import numpy as np +import torch +import zarr + +from earth2studio.utils.type import CoordSystem, TimeArray, VariableArray + +try: + from insitubatch import obstore_store, open_geometries, split_by_chunk + from insitubatch.frameworks import to_torch + from insitubatch.source import InSituDataset + from insitubatch.types import Batch +except ImportError as exc: # pragma: no cover - optional integration + raise ImportError( + "earth2studio.data.insitu needs insitubatch; install it with: pip install insitubatch" + ) from exc + + +def batch_to_xcoords( + batch: Batch, + *, + variables: VariableArray, + time: TimeArray, + lat: np.ndarray, + lon: np.ndarray, + device: torch.device | str = "cpu", + lead_time: np.ndarray | None = None, +) -> tuple[torch.Tensor, CoordSystem]: + """Convert one insitubatch ``Batch`` to the ``fetch_data(legacy=True)`` contract. + + The batch carries one ``(n_time, lat, lon)`` array per variable, keyed by the + Earth2Studio variable id used to build the dataset. The returned tensor has the model's + input layout ``(time, lead_time, variable, lat, lon)`` and ``coords`` is the matching + ``OrderedDict`` in that exact key order -- ``time`` (datetime64[ns]), ``lead_time`` + (timedelta64[ns]), ``variable`` (str), ``lat``/``lon`` (float32) -- so it drops straight + into ``prognostic.create_iterator`` after ``map_coords``. + + ``lead_time`` defaults to a single 0 h step (an initial condition); pass the model's + ``input_coords()["lead_time"]`` for a multi-step history window. + """ + variables = np.asarray(variables) + lead = np.array([np.timedelta64(0, "ns")]) if lead_time is None else np.asarray(lead_time) + tensors = to_torch(batch) # {var: (n_time, lat, lon)} via zero-copy DLPack + # (n_time, variable, lat, lon) -> insert the size-1 lead axis -> (n_time, lead, var, lat, lon) + x = torch.stack([tensors[str(v)] for v in variables], dim=1).unsqueeze(1).to(device) + coords: CoordSystem = OrderedDict( + [ + ("time", np.asarray(time, dtype="datetime64[ns]")), + ("lead_time", lead.astype("timedelta64[ns]")), + ("variable", variables), + ("lat", np.asarray(lat, dtype=np.float32)), + ("lon", np.asarray(lon, dtype=np.float32)), + ] + ) + return x, coords + + +class InSituForecastFeed: + """Iterate initial-condition ``(x, coords)`` batches for a hindcast window, prefetched. + + Reads the analysis ``store`` with insitubatch over a contiguous ``sample_range`` of the + time axis, batching ``batch_size`` consecutive init times per step with async read-ahead + bounded by ``max_inflight``. Each batch is converted to the ``fetch_data`` contract and + fed to ``prognostic.create_iterator``; while the model rolls out one batch, the loader + prefetches the next batch's ICs. ``variables`` are Earth2Studio ids; ``var_map`` maps them + to store array names when they differ (e.g. WB2 ``t2m -> 2m_temperature``). + """ + + def __init__( + self, + store_url: str, + variables: VariableArray, + *, + var_map: dict[str, str] | None = None, + time_name: str = "time", + lat_name: str = "latitude", + lon_name: str = "longitude", + sample_range: tuple[int, int] | None = None, + batch_size: int = 8, + max_inflight: int | None = None, + cache_dir: str | None = None, + device: torch.device | str = "cpu", + ) -> None: + self.variables = [str(v) for v in variables] + self.device = device + vmap = var_map or {v: v for v in self.variables} + store = obstore_store(store_url) + + group = zarr.open_group(store=store, mode="r") + self.time = np.asarray(group[time_name][:]).astype("datetime64[ns]") + self.lat = np.asarray(group[lat_name][:]).astype(np.float32) + self.lon = np.asarray(group[lon_name][:]).astype(np.float32) + + arrays = [vmap[v] for v in self.variables] + opened = open_geometries(store, variables=arrays) + geometries = {v: opened[vmap[v]] for v in self.variables} + manifest = split_by_chunk( + opened[arrays[0]], fractions=(1.0, 0.0, 0.0), sample_range=sample_range + ) + self.dataset = InSituDataset( + store, + manifest, + geometries=geometries, + batch_size=batch_size, + shuffle=False, + cache_dir=cache_dir, + max_inflight=max_inflight, + ) + + def __iter__(self): # type: ignore[no-untyped-def] + self.dataset.set_epoch(0) + for batch in self.dataset.train: + time = self.time[batch.sample_indices] + yield batch_to_xcoords( + batch, + variables=self.variables, + time=time, + lat=self.lat, + lon=self.lon, + device=self.device, + ) From 344e3a0f9821ee012b6c861cc425ad4c1ab9abc6 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Sat, 4 Jul 2026 16:15:58 +0000 Subject: [PATCH 02/16] Generalize the insitubatch feed to a lead axis with store injection Rework InSituForecastFeed from a single-lead initial-condition feed into a verification-capable feed: - Unify history (lead_time <= 0) and verification (lead_time > 0) into one lead axis. Each lead is a sample-axis `shift` view of one stored array, stacked into `(time, lead_time, variable, lat, lon)`, so the `(init, lead)` grid decodes each shared chunk exactly once. - Take an injected `store: Store` instead of building obstore from a URL, so a caller can pick any insitubatch backend (e.g. anonymous public buckets over gcsfs). - Decode the CF `time` coordinate via `cftime` (an existing dependency) rather than assuming a datetime64 encoding. - `transpose_inner` swaps a store's `(lon, lat)` field layout to the contract's `(lat, lon)`. `self.dataset` exposes the underlying InSituDataset cache counters. --- earth2studio/data/insitu.py | 157 ++++++++++++++++++++++++++---------- 1 file changed, 114 insertions(+), 43 deletions(-) diff --git a/earth2studio/data/insitu.py b/earth2studio/data/insitu.py index 512112f7f..433a556f7 100644 --- a/earth2studio/data/insitu.py +++ b/earth2studio/data/insitu.py @@ -18,27 +18,38 @@ Earth2Studio's standard path is ``DataSource -> xr.DataArray -> fetch_data -> prep_data_array -> (torch.Tensor, coords)``; xarray is load-bearing down to -``prep_data_array``. For an IO-bound hindcast / scoring campaign over many init times, that -per-``(time, variable)`` fetch re-reads overlapping chunks and its ``gather`` is unbounded. - -This module skips xarray entirely: it reads the analysis store with **insitubatch** -(`InSituDataset` -- bounded-fan-out async prefetch, a read plan that de-duplicates chunks -shared across init times and lead offsets) and converts each numpy ``Batch`` to the exact +``prep_data_array``, and ``fetch_data`` gathers one read per ``(time, lead_time)`` with +no de-duplication. For an IO-bound hindcast / scoring campaign the ``(init, lead)`` grid +maps many requested slices onto the **same** stored chunk (consecutive init times share +valid times; a fat time-chunk holds several steps), so that path re-reads and re-decodes +the same bytes over and over. + +This module skips xarray: it reads the analysis store with **insitubatch** +(:class:`~insitubatch.source.InSituDataset` -- bounded-fan-out async prefetch, a read plan +that decodes each shared chunk once) and converts each numpy ``Batch`` to the exact ``(x, coords)`` tuple ``fetch_data(..., legacy=True)`` returns, so it is a drop-in for the -initial-condition feed of ``earth2studio.run`` workflows. See :func:`batch_to_xcoords` for -the contract and :class:`InSituForecastFeed` for the prefetched iterator. +initial-condition feed of ``earth2studio.run`` workflows. + +The lead axis is unified: pass ``lead_times`` covering a model's input history (``<= 0``, +e.g. ``[-6h, 0]`` for a 2-step history model) and/or verification leads (``> 0``, for +scoring), each realised as a sample-axis ``shift`` view of one stored array -- no reshard. +See :func:`batch_to_xcoords` for the tensor contract and :class:`InSituForecastFeed` for +the prefetched iterator. """ from collections import OrderedDict +from collections.abc import Iterator +import cftime import numpy as np import torch import zarr +from zarr.abc.store import Store -from earth2studio.utils.type import CoordSystem, TimeArray, VariableArray +from earth2studio.utils.type import CoordSystem, VariableArray try: - from insitubatch import obstore_store, open_geometries, split_by_chunk + from insitubatch import open_geometries, split_by_chunk from insitubatch.frameworks import to_torch from insitubatch.source import InSituDataset from insitubatch.types import Batch @@ -48,38 +59,58 @@ ) from exc +def decode_cf_time(values: np.ndarray, units: str, calendar: str = "standard") -> np.ndarray: + """Decode a CF ``" since "`` integer time coordinate to datetime64[ns]. + + Handles the common reanalysis encoding (e.g. WB2/ARCO ERA5 store ``time`` as + ``"hours since 1959-01-01"``). Uses ``cftime`` (an Earth2Studio dependency) so odd + reference dates and units are handled the same way the rest of the ecosystem decodes + them. A non-standard calendar (``360_day`` / ``noleap``) yields ``cftime`` objects that + do not fit E2S's ``datetime64[ns]`` coordinate contract, so the cast raises here -- the + right place to fail for an out-of-contract analysis store. + """ + dates = cftime.num2date(values, units, calendar=calendar, only_use_cftime_datetimes=False) + return np.asarray(dates, dtype="datetime64[ns]") + + def batch_to_xcoords( batch: Batch, *, + labels: list[list[str]], variables: VariableArray, - time: TimeArray, + lead_time: np.ndarray, + time: np.ndarray, lat: np.ndarray, lon: np.ndarray, + transpose_inner: bool = False, device: torch.device | str = "cpu", - lead_time: np.ndarray | None = None, ) -> tuple[torch.Tensor, CoordSystem]: """Convert one insitubatch ``Batch`` to the ``fetch_data(legacy=True)`` contract. - The batch carries one ``(n_time, lat, lon)`` array per variable, keyed by the - Earth2Studio variable id used to build the dataset. The returned tensor has the model's - input layout ``(time, lead_time, variable, lat, lon)`` and ``coords`` is the matching - ``OrderedDict`` in that exact key order -- ``time`` (datetime64[ns]), ``lead_time`` - (timedelta64[ns]), ``variable`` (str), ``lat``/``lon`` (float32) -- so it drops straight - into ``prognostic.create_iterator`` after ``map_coords``. - - ``lead_time`` defaults to a single 0 h step (an initial condition); pass the model's - ``input_coords()["lead_time"]`` for a multi-step history window. + ``labels`` is a ``[lead][variable]`` grid of the batch keys (each a sample-axis + ``shift`` view of one stored array); the returned tensor has the model input layout + ``(time, lead_time, variable, lat, lon)`` and ``coords`` is the matching ``OrderedDict`` + in that exact key order -- ``time`` (datetime64[ns]), ``lead_time`` (timedelta64[ns]), + ``variable`` (str), ``lat``/``lon`` (float32) -- so it drops straight into + ``prognostic.create_iterator`` after ``map_coords``. ``transpose_inner`` swaps the two + field axes when the store lays fields out ``(lon, lat)`` but the contract wants + ``(lat, lon)``. """ - variables = np.asarray(variables) - lead = np.array([np.timedelta64(0, "ns")]) if lead_time is None else np.asarray(lead_time) - tensors = to_torch(batch) # {var: (n_time, lat, lon)} via zero-copy DLPack - # (n_time, variable, lat, lon) -> insert the size-1 lead axis -> (n_time, lead, var, lat, lon) - x = torch.stack([tensors[str(v)] for v in variables], dim=1).unsqueeze(1).to(device) + tensors = to_torch(batch) # {label: (n_time, *inner)} via zero-copy DLPack + # (n_time, var, *inner) per lead -> stack the lead axis -> (n_time, lead, var, *inner). + per_lead = [ + torch.stack([tensors[labels[li][vi]] for vi in range(len(variables))], dim=1) + for li in range(len(lead_time)) + ] + x = torch.stack(per_lead, dim=1) + if transpose_inner: + x = x.transpose(-1, -2) + x = x.contiguous().to(device) coords: CoordSystem = OrderedDict( [ ("time", np.asarray(time, dtype="datetime64[ns]")), - ("lead_time", lead.astype("timedelta64[ns]")), - ("variable", variables), + ("lead_time", np.asarray(lead_time, dtype="timedelta64[ns]")), + ("variable", np.asarray(variables)), ("lat", np.asarray(lat, dtype=np.float32)), ("lon", np.asarray(lon, dtype=np.float32)), ] @@ -88,22 +119,30 @@ def batch_to_xcoords( class InSituForecastFeed: - """Iterate initial-condition ``(x, coords)`` batches for a hindcast window, prefetched. + """Iterate ``(x, coords)`` batches over a hindcast window, prefetched and de-duplicated. Reads the analysis ``store`` with insitubatch over a contiguous ``sample_range`` of the - time axis, batching ``batch_size`` consecutive init times per step with async read-ahead - bounded by ``max_inflight``. Each batch is converted to the ``fetch_data`` contract and - fed to ``prognostic.create_iterator``; while the model rolls out one batch, the loader - prefetches the next batch's ICs. ``variables`` are Earth2Studio ids; ``var_map`` maps them - to store array names when they differ (e.g. WB2 ``t2m -> 2m_temperature``). + time axis, yielding ``batch_size`` consecutive init times per step with async read-ahead + bounded by ``max_inflight``. ``lead_times`` populates the ``lead_time`` axis: pass a + model's ``input_coords()["lead_time"]`` (values ``<= 0``) for a multi-step history + window, verification leads (``> 0``) for scoring, or their union. Each lead is a + sample-axis ``shift`` view of one stored array, so the ``(init, lead)`` grid decodes each + shared chunk exactly once (the win over per-``(time, lead)`` ``fetch_data``). + + ``variables`` are the ids to expose on the ``variable`` coordinate; ``var_map`` maps them + to store array names when they differ (e.g. ``t2m -> 2m_temperature``). Build ``store`` + with :func:`insitubatch.obstore_store` / :func:`insitubatch.fsspec_store` (e.g. anon + public buckets). ``self.dataset`` exposes the underlying :class:`InSituDataset` for its + ``cache_hits`` / ``cache_misses`` / ``resident_peak`` counters. """ def __init__( self, - store_url: str, + store: Store, variables: VariableArray, *, var_map: dict[str, str] | None = None, + lead_times: np.ndarray | None = None, time_name: str = "time", lat_name: str = "latitude", lon_name: str = "longitude", @@ -111,43 +150,75 @@ def __init__( batch_size: int = 8, max_inflight: int | None = None, cache_dir: str | None = None, + transpose_inner: bool = False, device: torch.device | str = "cpu", ) -> None: + self.store = store self.variables = [str(v) for v in variables] self.device = device + self.transpose_inner = transpose_inner vmap = var_map or {v: v for v in self.variables} - store = obstore_store(store_url) group = zarr.open_group(store=store, mode="r") - self.time = np.asarray(group[time_name][:]).astype("datetime64[ns]") + time_arr = np.asarray(group[time_name][:]) + attrs = dict(group[time_name].attrs) + units = attrs.get("units") + self.time = ( + decode_cf_time(time_arr, units, attrs.get("calendar", "standard")) + if units + else time_arr.astype("datetime64[ns]") + ) self.lat = np.asarray(group[lat_name][:]).astype(np.float32) self.lon = np.asarray(group[lon_name][:]).astype(np.float32) + # Sample-axis step of the store (dt); every lead must be an integer multiple of it. + dt = self.time[1] - self.time[0] + leads = np.array([np.timedelta64(0, "ns")]) if lead_times is None else np.asarray(lead_times) + self.lead_time = leads.astype("timedelta64[ns]") + steps = self.lead_time / dt + if not np.all(steps == np.round(steps)): + raise ValueError( + f"every lead_time must be an integer multiple of the store step {dt}; got {leads}" + ) + self.lead_steps = np.round(steps).astype(np.int64) + arrays = [vmap[v] for v in self.variables] opened = open_geometries(store, variables=arrays) - geometries = {v: opened[vmap[v]] for v in self.variables} + # One shifted geometry per (lead, variable); label grid indexes them for the stacker. + geometries: dict[str, object] = {} + self.labels: list[list[str]] = [] + for li, k in enumerate(self.lead_steps): + row = [] + for v in self.variables: + label = f"{v}#{li}" + geometries[label] = opened[vmap[v]].shift(int(k)) + row.append(label) + self.labels.append(row) + manifest = split_by_chunk( opened[arrays[0]], fractions=(1.0, 0.0, 0.0), sample_range=sample_range ) self.dataset = InSituDataset( store, manifest, - geometries=geometries, + geometries=geometries, # type: ignore[arg-type] batch_size=batch_size, shuffle=False, cache_dir=cache_dir, max_inflight=max_inflight, ) - def __iter__(self): # type: ignore[no-untyped-def] + def __iter__(self) -> Iterator[tuple[torch.Tensor, CoordSystem]]: self.dataset.set_epoch(0) - for batch in self.dataset.train: - time = self.time[batch.sample_indices] + for batch in self.dataset.all: yield batch_to_xcoords( batch, + labels=self.labels, variables=self.variables, - time=time, + lead_time=self.lead_time, + time=self.time[batch.sample_indices], lat=self.lat, lon=self.lon, + transpose_inner=self.transpose_inner, device=self.device, ) From deb2ced25cfa962a7481df5b3c6757e10c440c41 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Sat, 4 Jul 2026 16:15:58 +0000 Subject: [PATCH 03/16] Add insitubatch hindcast IO benchmark recipe Two runnable benchmarks that feed ERA5 into an Earth2Studio prognostic via the insitubatch feed instead of the dense `fetch_data` grid, with a README framing the results: - bench_hindcast.py: verification-read de-duplication over a hindcast grid (WB2 fat-chunk / ARCO chunk-1), before (E2S fetch) vs after (insitubatch feed). - stream_score.py: streaming vs dense verification materialization, scored against ERA5 with Persistence; reports wall, peak RSS, and matching RMSE. Both read anonymous public GCS over gcsfs on both sides, so the delta isolates the loader (not an obstore-vs-gcsfs artifact). Motivated by recipes/eval's predownload sentinel, which exists because live fetch_data is too slow. --- recipes/insitubatch_hindcast/README.md | 102 +++++++++++ .../insitubatch_hindcast/bench_hindcast.py | 160 ++++++++++++++++ recipes/insitubatch_hindcast/stream_score.py | 173 ++++++++++++++++++ 3 files changed, 435 insertions(+) create mode 100644 recipes/insitubatch_hindcast/README.md create mode 100644 recipes/insitubatch_hindcast/bench_hindcast.py create mode 100644 recipes/insitubatch_hindcast/stream_score.py diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md new file mode 100644 index 000000000..a1baf5e45 --- /dev/null +++ b/recipes/insitubatch_hindcast/README.md @@ -0,0 +1,102 @@ +# insitubatch × Earth2Studio: streaming hindcast IO + +Two runnable benchmarks that feed ERA5 into an Earth2Studio prognostic **without** the dense +`fetch_data` grid — reading the analysis store with [insitubatch](https://github.com/emfdavid/insitubatch) +(`earth2studio.data.insitu.InSituForecastFeed`) instead. They quantify what a streaming, +read-planning loader changes for an IO-bound hindcast / scoring campaign. + +The motivation is `recipes/eval`: its `predownload.py` sentinel exists because live `fetch_data` +is too slow for a scoring campaign. That predownload materializes the whole `(init, lead)` +verification grid up front. Both are consequences of a per-`(time, variable)` fetch with no read +de-duplication — exactly what insitubatch removes. + +## Setup + +```bash +pip install insitubatch # the loader (adds the earth2studio.data.insitu adapter's dep) +# earth2studio + its deps already present in this tree +``` + +Both stores are anonymous public GCS buckets (WeatherBench2 ERA5, ARCO ERA5); no credentials +needed. Every measurement below reads over **gcsfs anon on both the before and after side**, so the +delta isolates insitubatch's read-planning + streaming — it is *not* an obstore-vs-gcsfs artifact. + +## 1. `bench_hindcast.py` — verification-read de-duplication + +A scoring grid needs ERA5 at `valid = init + lead` for every `(init, lead)`. Consecutive init +times share valid times, and a fat time-chunk holds several steps, so the requested reads collapse +onto far fewer stored chunks. BEFORE = E2S's per-init `fetch_data`; AFTER = the insitubatch feed +(each lead a sample-axis `shift` view; each shared chunk decoded once). + +```bash +python bench_hindcast.py --store wb2 --vars t2m u10m v10m --n-init 48 --max-lead-h 240 --repeats 5 +python bench_hindcast.py --store arco --vars t2m --n-init 24 --lead-step-h 6 --max-lead-h 144 --repeats 3 +``` + +| store | layout | requested reads | unique decodes | **wall speedup** | +|-------|--------|-----------------|----------------|------------------| +| **WB2** 240×121 6-h | `chunks=(8,240,121)` (fat) | 5760 | **33** (174×) | **15.4×** (14.0 s → 0.91 s) | +| **ARCO** 721×1440 1-h | `chunks=(1,721,1440)` (chunk-1) | 576 | 162 (3.6×) | ~1.9× | + +WB2's fat time-chunk amortizes 8 steps per read, so the de-dup ratio is large and the fields are +small — insitubatch dominates. ARCO is the **honest** case (see caveats). + +## 2. `stream_score.py` — streaming vs dense materialization + +The model's `create_iterator` already streams the forecast lead-by-lead, and scoring is pointwise +per `(init, lead)` — so the verification never needs to be a dense tensor. Interleave instead: +roll out a window of inits, score each lead against a just-read verification slice, discard. Three +modes, all producing **identical RMSE** (a correctness check): + +```bash +for m in e2s dense stream; do python stream_score.py --mode $m --n-init 120 --n-leads 40; done +``` + +| mode | wall | **peak RSS** | field reads | +|------|------|--------------|-------------| +| `e2s` — dense predownload (redundant reads) | 39.6 s | 3.04 GB | 14 760 | +| `dense` — insitubatch, `batch_size=N` | 4.8 s | 7.53 GB | 60 | +| `stream` — insitubatch, `batch_size=W` | 3.3 s | **1.81 GB** | 60 | + +Streaming's peak memory is **flat at ~1.9 GB across N = 120 / 240 / 480**, while the dense grid is +7.53 GB at N = 120 and **OOMs a 15 GB box by ~N = 240**. Dense scales with campaign size; streaming +does not. That bounded-memory property — not just throughput — is the point for a long campaign. + +(Persistence is a checkpoint-free model that exercises the real `create_iterator` seam on CPU; a +real NVIDIA checkpoint — SFNO/FCN — is a drop-in with the same code on a GPU.) + +## How to read these numbers — framing insitubatch + +insitubatch is a **streaming batch loader** that trains/infers in place on cloud zarr: all +parallelism lives in one async event loop, the Python hot path is O(chunks) not O(samples), and +memory is bounded by a residency budget rather than the working set. The two benchmarks above +sharpen its positioning into three evidence-backed claims: + +1. **Competitive with an optimized parallel loader, at lower memory.** On a well-chunked store and + for streaming consumption it matches a hand-tuned concurrent fetch's throughput while holding + *bounded* memory (streaming: flat ~1.9 GB where dense predownload OOMs). Evidence: §2. +2. **Far ahead when the chunking strategy isn't sample-optimized.** When the access pattern maps + many samples onto shared chunks — overlapping windows, verification grids, fat chunks holding + several steps — its read planning de-duplicates and a per-sample parallel fetch re-reads. + Evidence: §1 WB2 (174× fewer decodes, 15× wall). +3. **Honest boundary — you can use it sub-optimally.** It is not a universal speed win. On a + chunk-1 store with large fields, against an *unbounded* concurrent gather, its bounded-inflight + scheduling trails per byte (ARCO ~2×; the reads are already minimal — verified — but the dense + output the model consumes must still be assembled, and E2S's flat gather saturates bandwidth on + 4 MB chunks). And a degenerate `batch_size=N` throws away the memory advantage. The tool is + **generally optimal for streaming with bounded memory** — that is the sweet spot. + +One line: *stream training/inference batches from cloud tensors in place, with bounded memory — +competitive with hand-tuned parallel loaders on optimized layouts, and far ahead when the chunking +causes duplicate reads.* + +## Caveats / methodology + +- **Single environment, preliminary.** One n2-standard-8-class box (15 GB RAM), cold reads, gcsfs + anon. Numbers to be **cross-posted** after NVIDIA-side runs on the target infrastructure. +- **gcsfs on both sides.** Isolates the loader's contribution from the store backend; obstore would + raise the AFTER throughput further but is not what these numbers measure. +- **Surface variables only** (`t2m`, `u10m`, `v10m`); pressure-level variables need level indexing, + not yet wired in the adapter. +- **The win is the IO-bound campaign** (many inits, verification-heavy — hindcast scoring, lagged + ensembles). A single-IC long rollout is compute-bound, where the loader is a rounding error. diff --git a/recipes/insitubatch_hindcast/bench_hindcast.py b/recipes/insitubatch_hindcast/bench_hindcast.py new file mode 100644 index 000000000..b24e39808 --- /dev/null +++ b/recipes/insitubatch_hindcast/bench_hindcast.py @@ -0,0 +1,160 @@ +"""Before/after hindcast verification-read benchmark on WB2 / ARCO ERA5. + +Scenario: score an ``N_init x len(leads)`` forecast grid against ERA5. Every ``(init, lead)`` +needs ERA5 at ``valid = init + lead``; consecutive init times share valid times, so the +requested reads collapse onto far fewer stored chunks. + +BEFORE = Earth2Studio's ERA5 source (``fetch`` gathers one read per (time, variable), no + dedup) -- the realistic per-init eval fetch. +AFTER = insitubatch InSituForecastFeed over the init window with the leads as shift views; + each shared chunk is decoded exactly once (``dataset.cache_misses``). + +Both read the SAME store over gcsfs anon, so the delta isolates insitubatch's +dedup + bounded prefetch (not obstore-vs-gcsfs). + +Two regimes: + wb2 = 240x121 6-hourly, chunks=(8,240,121): fat time-chunk -> high dedup ratio, tiny + fields -> wall gated by concurrency (wall speedup << read reduction). + arco = 721x1440 1-hourly, chunks=(1,721,1440): chunk-1 -> dedup = pure valid-time overlap, + 4MB fields -> genuinely IO-bound (wall speedup tracks read reduction). +""" + +import argparse +import time + +import numpy as np + +from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time +from insitubatch import fsspec_store + +# store id -> (url, before source class, inner (H,W), transpose store(lon,lat)->(lat,lon)) +STORES = { + "wb2": { + "url": "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr", + "before": "earth2studio.data.wb2:WB2ERA5_121x240", + "field_bytes": 240 * 121 * 4, + "chunk_steps": 8, + "transpose_inner": True, + }, + "arco": { + "url": "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", + "before": "earth2studio.data.arco:ARCO", + "field_bytes": 721 * 1440 * 4, + "chunk_steps": 1, + "transpose_inner": False, + }, +} +VAR_MAP = { + "t2m": "2m_temperature", + "u10m": "10m_u_component_of_wind", + "v10m": "10m_v_component_of_wind", +} + + +def anon_store(url): + return fsspec_store(url, token="anon", access="read_only") # noqa: S106 + + +def load_before_cls(spec): + mod, _, name = spec.partition(":") + import importlib + + return getattr(importlib.import_module(mod), name) + + +def run_after(cfg, variables, start, n_init, leads_h, batch_size, max_inflight): + leads = np.array([np.timedelta64(h, "h") for h in leads_h]) + feed = InSituForecastFeed( + anon_store(cfg["url"]), + variables=variables, + var_map={v: VAR_MAP[v] for v in variables}, + lead_times=leads, + sample_range=(start, start + n_init), + batch_size=batch_size, + max_inflight=max_inflight, + transpose_inner=cfg["transpose_inner"], + ) + t0 = time.perf_counter() + n_rows = 0 + for x, _coords in feed: + n_rows += x.shape[0] # gather returns eager numpy; no touch needed to force decode + wall = time.perf_counter() - t0 + feed.dataset.close() + return {"wall_s": wall, "init_rows": n_rows, + "chunk_decodes": feed.dataset.cache_misses, "resident_peak": feed.dataset.resident_peak} + + +def run_before(before_cls, variables, init_times64, leads_h): + src = before_cls(cache=False, verbose=False) + init_dt = init_times64.astype("datetime64[s]").astype("O") + leads_td = [np.timedelta64(h, "h") for h in leads_h] + t0 = time.perf_counter() + for it in init_dt: + valid = [(np.datetime64(it) + td).astype("datetime64[s]").astype("O") for td in leads_td] + src(valid, list(variables)) # realistic per-init verification fetch + return {"wall_s": time.perf_counter() - t0} + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--store", choices=list(STORES), default="wb2") + p.add_argument("--vars", nargs="+", default=["t2m"]) + p.add_argument("--start", type=int, default=1000) + p.add_argument("--n-init", type=int, default=24) + p.add_argument("--lead-step-h", type=int, default=6) + p.add_argument("--max-lead-h", type=int, default=120) + p.add_argument("--batch-size", type=int, default=8) + p.add_argument("--max-inflight", type=int, default=32) + p.add_argument("--repeats", type=int, default=1) + p.add_argument("--skip-before", action="store_true") + args = p.parse_args() + + cfg = STORES[args.store] + leads_h = list(range(args.lead_step_h, args.max_lead_h + 1, args.lead_step_h)) + requested = args.n_init * len(leads_h) * len(args.vars) + + import zarr + + g = zarr.open_group(store=anon_store(cfg["url"]), mode="r") + attrs = dict(g["time"].attrs) + times64 = decode_cf_time(np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard")) + init_times64 = times64[args.start : args.start + args.n_init] + + print(f"[{args.store}] grid: {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " + f"= {requested} requested field-reads ({requested*cfg['field_bytes']/1e9:.2f} GB naive)") + print(f"leads: {args.lead_step_h}h..{args.max_lead_h}h ; vars: {args.vars} ; repeats: {args.repeats}") + + def med3(w): + w = sorted(w) + return w[len(w) // 2], w[0], w[-1] + + before_cls = load_before_cls(cfg["before"]) + after_walls, before_walls = [], [] + decodes = resident = None + for r in range(args.repeats): + a = run_after(cfg, args.vars, args.start, args.n_init, leads_h, args.batch_size, args.max_inflight) + after_walls.append(a["wall_s"]) + decodes, resident = a["chunk_decodes"], a["resident_peak"] + if not args.skip_before: + before_walls.append(run_before(before_cls, args.vars, init_times64, leads_h)["wall_s"]) + print(f" repeat {r+1}/{args.repeats}: after={after_walls[-1]:.2f}s" + + (f" before={before_walls[-1]:.2f}s" if before_walls else "")) + + dedup = requested / decodes + a_med, a_lo, a_hi = med3(after_walls) + print("\n=== AFTER (insitubatch, gcsfs anon) ===") + print(f" wall (med/min/max): {a_med:.2f} / {a_lo:.2f} / {a_hi:.2f} s") + print(f" chunk decodes : {decodes} ({decodes*cfg['field_bytes']*cfg['chunk_steps']/1e9:.2f} GB)") + print(f" dedup ratio : {dedup:.1f}x ({requested} requested -> {decodes} decoded)") + print(f" resident peak : {resident} chunks") + if before_walls: + b_med, b_lo, b_hi = med3(before_walls) + print("\n=== BEFORE (E2S fetch, gcsfs anon, cache off) ===") + print(f" wall (med/min/max): {b_med:.2f} / {b_lo:.2f} / {b_hi:.2f} s") + print("\n=== HEADLINE (medians) ===") + print(f" speedup : {b_med/a_med:.1f}x wall ({b_med:.2f}s -> {a_med:.2f}s)") + print(f" read reduction : {dedup:.1f}x fewer chunk decodes") + + +if __name__ == "__main__": + main() diff --git a/recipes/insitubatch_hindcast/stream_score.py b/recipes/insitubatch_hindcast/stream_score.py new file mode 100644 index 000000000..d253c426e --- /dev/null +++ b/recipes/insitubatch_hindcast/stream_score.py @@ -0,0 +1,173 @@ +"""Streaming vs dense hindcast scoring: interleave verification with the model rollout. + +The Earth2Studio pattern (`fetch_data -> map_coords -> create_iterator`) materializes the whole +`(init, lead)` verification grid up front (its `recipes/eval` even has a predownload sentinel, +because live fetch is slow). That dense tensor is a *slow-IO shortcut*, not a requirement: the +model's `create_iterator` already STREAMS the forecast lead-by-lead, and scoring is pointwise per +`(init, lead)`. So the verification never needs to be dense -- read each lead's ground truth as the +rollout produces it, accumulate RMSE, discard. + +Three modes over one hindcast campaign (N inits x L leads x V vars), scored vs ERA5 with Persistence: + e2s = E2S WB2 fetch_data (redundant per-(time,var) reads) -> dense grid -> roll out + score. + dense = insitubatch, batch_size=N -> one dense materialization (the E2S shape, deduped reads). + stream = insitubatch, batch_size=W -> windowed stream; roll out + score each window inline, discard. + +Reports wall, peak RSS (getrusage), and RMSE-per-lead (must match across modes = correctness). +`dense` vs `stream` isolates the *materialization* axis (same reads, same backend): peak memory +N*L vs W*L. `e2s` adds the redundant-read wall of the status quo. +""" + +import argparse +import resource +import time +from collections import OrderedDict + +import numpy as np +import torch + +from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time +from earth2studio.models.px import Persistence +from earth2studio.utils.coords import map_coords +from insitubatch import fsspec_store + +URL = "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr" +VAR_MAP = {"t2m": "2m_temperature", "u10m": "10m_u_component_of_wind", "v10m": "10m_v_component_of_wind"} +DT = np.timedelta64(6, "h") + + +def anon_store(): + return fsspec_store(URL, token="anon", access="read_only") # noqa: S106 + + +def peak_rss_gb(): + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 # ru_maxrss is KB on Linux + + +class RmseAccumulator: + """Per-lead running MSE over (init, var, lat, lon); sqrt at the end.""" + + def __init__(self): + self.sse: dict[int, float] = {} + self.n: dict[int, int] = {} + + def update(self, lead_h: int, pred: torch.Tensor, truth: torch.Tensor): + e = (pred - truth).float() + self.sse[lead_h] = self.sse.get(lead_h, 0.0) + float((e * e).sum()) + self.n[lead_h] = self.n.get(lead_h, 0) + e.numel() + + def table(self): + return {h: (self.sse[h] / self.n[h]) ** 0.5 for h in sorted(self.sse)} + + +def build_feed(variables, start, n_init, leads_h, batch_size): + leads = np.array([np.timedelta64(h, "h") for h in leads_h]) # includes 0 (the IC) + return InSituForecastFeed( + anon_store(), variables=variables, var_map={v: VAR_MAP[v] for v in variables}, + lead_times=leads, sample_range=(start, start + n_init), + batch_size=batch_size, transpose_inner=True, + ) + + +def make_model(variables, feed): + domain = OrderedDict([("lat", feed.lat), ("lon", feed.lon)]) + return Persistence(variable=variables, domain_coords=domain, history=1, dt=DT) + + +def score_window(model, x_all, coords_all, leads_h, nsteps, acc): + """Roll Persistence out over one window and score each lead vs the pre-read verification slice. + + ``x_all`` is (W, L+1, V, lat, lon) with lead index 0 = IC and index k = verification at leads_h[k]. + """ + ic_x = x_all[:, 0:1] + ic_coords = OrderedDict( + [("time", coords_all["time"]), ("lead_time", coords_all["lead_time"][0:1]), + ("variable", coords_all["variable"]), ("lat", coords_all["lat"]), ("lon", coords_all["lon"])] + ) + ic_x, ic_coords = map_coords(ic_x, ic_coords, model.input_coords()) + for step, (fx, fcoords) in enumerate(model.create_iterator(ic_x, ic_coords)): + if step == 0: + continue # step 0 is the IC (lead 0); score forecast leads only + acc.update(leads_h[step], fx[:, -1], x_all[:, step]) + if step == nsteps: + break + + +def run_insitu(variables, start, n_init, leads_h, nsteps, batch_size): + feed = build_feed(variables, start, n_init, leads_h, batch_size) + model = make_model(variables, feed) + acc = RmseAccumulator() + for x_all, coords_all in feed: # one batch per window (stream) or the whole campaign (dense) + score_window(model, x_all, coords_all, leads_h, nsteps, acc) + decodes = feed.dataset.cache_misses + feed.dataset.close() + return acc, decodes + + +def run_e2s(variables, start, n_init, leads_h, nsteps): + from earth2studio.data.wb2 import WB2ERA5_121x240 + + src = WB2ERA5_121x240(cache=False, verbose=False) + g_store = anon_store() + import zarr + + g = zarr.open_group(store=g_store, mode="r") + attrs = dict(g["time"].attrs) + times = decode_cf_time(np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard")) + inits = times[start : start + n_init] + feed = build_feed(variables, start, n_init, leads_h, batch_size=n_init) # for coords/model only + model = make_model(variables, feed) + feed.dataset.close() + + # Dense predownload: fetch every (init, lead) valid time (redundant reads), materialize. + lat = np.asarray(g["latitude"][:]).astype(np.float32) + lon = np.asarray(g["longitude"][:]).astype(np.float32) + dense = np.empty((n_init, len(leads_h), len(variables), len(lat), len(lon)), dtype=np.float32) + for i, it in enumerate(inits): + # Realistic per-init verification fetch: all leads for this init in one call. + valids = [(np.datetime64(it) + np.timedelta64(h, "h")).astype("datetime64[s]").astype("O") + for h in leads_h] + da = src(valids, list(variables)) # (L+1, V, lat, lon) xr.DataArray + dense[i] = np.asarray(da.values) + x_all = torch.from_numpy(dense) + coords_all = OrderedDict( + [("time", inits), ("lead_time", np.array([np.timedelta64(h, "h") for h in leads_h], dtype="timedelta64[ns]")), + ("variable", np.asarray(variables)), ("lat", lat), ("lon", lon)] + ) + acc = RmseAccumulator() + score_window(model, x_all, coords_all, leads_h, nsteps, acc) + return acc, n_init * len(leads_h) * len(variables) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["stream", "dense", "e2s"], required=True) + p.add_argument("--vars", nargs="+", default=["t2m", "u10m", "v10m"]) + p.add_argument("--start", type=int, default=2000) + p.add_argument("--n-init", type=int, default=120) + p.add_argument("--n-leads", type=int, default=40) + p.add_argument("--window", type=int, default=8) + args = p.parse_args() + + leads_h = [6 * k for k in range(args.n_leads + 1)] # 0, 6, .., n_leads*6 (0 = IC) + nsteps = args.n_leads + + t0 = time.perf_counter() + if args.mode == "e2s": + acc, reads = run_e2s(args.vars, args.start, args.n_init, leads_h, nsteps) + else: + bs = args.window if args.mode == "stream" else args.n_init + acc, reads = run_insitu(args.vars, args.start, args.n_init, leads_h, nsteps, bs) + wall = time.perf_counter() - t0 + + rmse = acc.table() + print(f"mode={args.mode} grid={args.n_init}x{args.n_leads}x{len(args.vars)} window={args.window}") + print(f" wall : {wall:.2f} s") + print(f" peak RSS : {peak_rss_gb():.2f} GB") + print(f" reads/decodes: {reads}") + print(f" RMSE @ 24h/120h/240h: " + f"{rmse.get(24, float('nan')):.3f} / {rmse.get(120, float('nan')):.3f} / " + f"{rmse.get(min(240, nsteps*6), float('nan')):.3f}") + + +if __name__ == "__main__": + main() From 9da88a257a34f174074fc54608dc68c250304faa Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Mon, 6 Jul 2026 04:16:16 +0000 Subject: [PATCH 04/16] Pin released insitubatch 0.1.0 + use its public API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit insitubatch 0.1.0 is on PyPI and now exposes InSituDataset and the framework adapters at the package root, so: - earth2studio/data/insitu.py imports them from `insitubatch` (public surface) instead of reaching into `insitubatch.source` / `.frameworks` / `.types`. - pyproject: add `insitubatch>=0.1.0` to the `data` extra, gated `python_version>='3.12'` (insitubatch requires 3.12; E2S supports 3.11) — mirrors the existing intake-esgf marker. Resolves from PyPI (uv.lock updated). --- earth2studio/data/insitu.py | 11 +++++++---- pyproject.toml | 1 + uv.lock | 23 +++++++++++++++++++++-- 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/earth2studio/data/insitu.py b/earth2studio/data/insitu.py index 433a556f7..1e7d511b9 100644 --- a/earth2studio/data/insitu.py +++ b/earth2studio/data/insitu.py @@ -49,10 +49,13 @@ from earth2studio.utils.type import CoordSystem, VariableArray try: - from insitubatch import open_geometries, split_by_chunk - from insitubatch.frameworks import to_torch - from insitubatch.source import InSituDataset - from insitubatch.types import Batch + from insitubatch import ( + Batch, + InSituDataset, + open_geometries, + split_by_chunk, + to_torch, + ) except ImportError as exc: # pragma: no cover - optional integration raise ImportError( "earth2studio.data.insitu needs insitubatch; install it with: pip install insitubatch" diff --git a/pyproject.toml b/pyproject.toml index d8a4cd97c..60ec9e845 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,7 @@ data = [ "rioxarray>=0.15.5", "pyproj>=3.7.1", # TODO: pygrib imports this, so technically in core, to check "scipy>=1.15.2", + "insitubatch>=0.1.0; python_version>='3.12'", # streaming cloud-zarr loader — earth2studio.data.insitu "eumdac>=3.1.0", ] perturbation = [ diff --git a/uv.lock b/uv.lock index 76e4152eb..54f7b0ac1 100644 --- a/uv.lock +++ b/uv.lock @@ -2021,12 +2021,13 @@ all = [ { name = "fastapi" }, { name = "flax" }, { name = "graphcast" }, - { name = "gribberish", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens')" }, + { name = "gribberish", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "hiredis" }, { name = "httpx" }, { name = "hydra-core" }, { name = "icechunk", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "importlib-metadata" }, + { name = "insitubatch", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "intake-esgf", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "jax", extra = ["cuda13"] }, { name = "makani" }, @@ -2134,9 +2135,10 @@ data = [ { name = "eccodeslib" }, { name = "ecmwf-opendata" }, { name = "eumdac" }, - { name = "gribberish", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens')" }, + { name = "gribberish", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "httpx" }, { name = "icechunk", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, + { name = "insitubatch", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "intake-esgf", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "pcodec", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, { name = "planetary-computer" }, @@ -2449,6 +2451,8 @@ requires-dist = [ { name = "icechunk", marker = "python_full_version >= '3.12' and extra == 'data'", specifier = ">=2.0.0" }, { name = "importlib-metadata", marker = "extra == 'all'" }, { name = "importlib-metadata", marker = "extra == 'dlwp'" }, + { name = "insitubatch", marker = "python_full_version >= '3.12' and extra == 'all'", specifier = ">=0.1.0" }, + { name = "insitubatch", marker = "python_full_version >= '3.12' and extra == 'data'", specifier = ">=0.1.0" }, { name = "intake-esgf", marker = "python_full_version >= '3.12' and extra == 'all'", specifier = ">=2026.1.26" }, { name = "intake-esgf", marker = "python_full_version >= '3.12' and extra == 'data'", specifier = ">=2026.1.26" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'all'", specifier = ">=0.4.26" }, @@ -3997,6 +4001,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "insitubatch" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, + { name = "obstore", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, + { name = "xarray", marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, + { name = "zarr", version = "3.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-atlas') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-fcn3') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-perturbation') or (extra == 'extra-12-earth2studio-ace2' and extra == 'extra-12-earth2studio-sfno') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifs2ens') or (extra == 'extra-12-earth2studio-aifs2' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-aifs2ens' and extra == 'extra-12-earth2studio-aifsens') or (extra == 'extra-12-earth2studio-cosmo' and extra == 'extra-12-earth2studio-da-healda') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormcast-conus') or (extra == 'extra-12-earth2studio-da-healda' and extra == 'extra-12-earth2studio-stormscope')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/68/9e/9691819a9d8e2b86daf01a0852f478887fbdd26118ab0ac9f0b97793a09d/insitubatch-0.1.0.tar.gz", hash = "sha256:9d3b8d93a9249f7eedd7fc29daf28d676adc76f561f570f8e5960e8235df1743", size = 58708, upload-time = "2026-07-06T04:02:50.897Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/10/b9d2ae405202d34daa082b62cba983d611a193d3e1204fd38301bbfffc00/insitubatch-0.1.0-py3-none-any.whl", hash = "sha256:1022bc63086a843ad098ebc9eec05478ffb171e65573289bf957e63b1a6ecb6d", size = 65528, upload-time = "2026-07-06T04:02:49.641Z" }, +] + [[package]] name = "intake-esgf" version = "2026.6.4" From 18658b9facbc67b4a8e363bd684ddf0a9d3d95b1 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Mon, 6 Jul 2026 04:27:04 +0000 Subject: [PATCH 05/16] Update hindcast recipe README setup for the pinned data extra insitubatch is now a declared earth2studio `data`-extra dependency (insitubatch>=0.1.0, gated Python>=3.12), not a manual side install. Point the Setup at `earth2studio[data]` (or a direct pinned `insitubatch>=0.1.0`). --- recipes/insitubatch_hindcast/README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index a1baf5e45..bc7e562b6 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -13,8 +13,10 @@ de-duplication — exactly what insitubatch removes. ## Setup ```bash -pip install insitubatch # the loader (adds the earth2studio.data.insitu adapter's dep) -# earth2studio + its deps already present in this tree +# insitubatch is declared in earth2studio's `data` extra (needs Python >= 3.12): +pip install "earth2studio[data]" +# or, since earth2studio is already present in this tree, just the loader: +pip install "insitubatch>=0.1.0" ``` Both stores are anonymous public GCS buckets (WeatherBench2 ERA5, ARCO ERA5); no credentials From d15f11007306ef2f70261e135b7d6305b61d7c95 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Mon, 6 Jul 2026 04:28:47 +0000 Subject: [PATCH 06/16] Use uv sync --extra data in the recipe setup (house style) The earth2studio tree is uv-managed and the recipe runs in-tree, so match the repo convention (`uv sync --extra `) instead of pip; the data extra brings insitubatch. --- recipes/insitubatch_hindcast/README.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index bc7e562b6..53011b308 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -14,9 +14,7 @@ de-duplication — exactly what insitubatch removes. ```bash # insitubatch is declared in earth2studio's `data` extra (needs Python >= 3.12): -pip install "earth2studio[data]" -# or, since earth2studio is already present in this tree, just the loader: -pip install "insitubatch>=0.1.0" +uv sync --extra data ``` Both stores are anonymous public GCS buckets (WeatherBench2 ERA5, ARCO ERA5); no credentials From 4ee5d190971ea641bfcb5fede118d6779ec5c576 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Thu, 9 Jul 2026 03:40:04 +0000 Subject: [PATCH 07/16] Add offline tests + CHANGELOG for the insitubatch feed Cover InSituForecastFeed / batch_to_xcoords / decode_cf_time against a synthetic on-disk store (fat time-chunk, CF-encoded time coord): the (x, coords) contract, byte-correct sample-axis shift views, the decode-once dedup property (cache_misses), transpose_inner, and the integer-multiple-of-dt lead guard. No network / live bucket. Addresses the tests + CHANGELOG items of the PR checklist. --- CHANGELOG.md | 4 + test/data/test_insitu.py | 228 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 test/data/test_insitu.py diff --git a/CHANGELOG.md b/CHANGELOG.md index aec59bc96..343769f6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added IEM parsed ASOS/AWOS station observation data source (`IEM_ASOS`) - Added Zarr v3 sharding support to `AsyncZarrBackend` +- Added optional `InSituForecastFeed` (`earth2studio.data.insitu`), an insitubatch-backed + streaming initial-condition / verification feed that reads a cloud zarr analysis store + with a de-duplicating read plan and yields `(torch.Tensor, CoordSystem)` batches for + IO-bound hindcast / scoring campaigns. ### Changed diff --git a/test/data/test_insitu.py b/test/data/test_insitu.py new file mode 100644 index 000000000..8e21baeb6 --- /dev/null +++ b/test/data/test_insitu.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline tests for the insitubatch initial-condition / verification feed. + +The feed takes an injected zarr ``Store``, so the whole surface is exercised against a +small synthetic store on disk -- no network, no live bucket. The fixture mirrors the +public WB2 / ARCO layout the adapter targets: a fat time-chunk (several sample-axis steps +per stored chunk) plus 1-D ``time`` (CF-encoded) / ``latitude`` / ``longitude`` coordinate +arrays. +""" + +from importlib.metadata import version + +import numpy as np +import pytest + +pytest.importorskip("insitubatch", reason="insitubatch is an optional dependency") + +from insitubatch import ensure_local_dir, obstore_store # noqa: E402 + +from earth2studio.data.insitu import ( # noqa: E402 + InSituForecastFeed, + decode_cf_time, +) + +pytestmark = pytest.mark.skipif( + int(version("zarr").split(".")[0]) < 3, reason="Requires zarr v3" +) + +TIME_UNITS = "hours since 1959-01-01" +STEP_H = 6 # store sample-axis step (6-hourly, like WB2) + + +def write_store(tmp_path, *, n=48, spc=8, lat=4, lon=5, inner_latlon=True, seed=0): + """Write a synthetic analysis store; return ``(store, {array_name: source_ndarray})``. + + ``spc`` steps per stored chunk is the fat time-chunk that lets an overlapping + ``(init, lead)`` grid collapse onto few decodes. ``inner_latlon`` lays fields out + ``(lat, lon)`` (the contract order); set it ``False`` for the ``(lon, lat)`` layout that + ``transpose_inner=True`` is meant to fix. + """ + import zarr + + url = f"file://{tmp_path}/analysis.zarr" + ensure_local_dir(url) + store = obstore_store(url, read_only=False) + group = zarr.open_group(store=store, mode="w") + + # CF-encoded time coordinate: 6-hourly integers since a reanalysis epoch. + t = group.create_array("time", shape=(n,), chunks=(n,), dtype="i8") + t[:] = (np.arange(n) * STEP_H).astype("i8") + t.attrs["units"] = TIME_UNITS + + inner = (lat, lon) if inner_latlon else (lon, lat) + group.create_array("latitude", shape=(lat,), chunks=(lat,), dtype="f4")[:] = ( + np.linspace(90.0, -90.0, lat, dtype="f4") + ) + group.create_array("longitude", shape=(lon,), chunks=(lon,), dtype="f4")[:] = ( + np.linspace(0.0, 360.0, lon, endpoint=False, dtype="f4") + ) + + rng = np.random.default_rng(seed) + srcs: dict[str, np.ndarray] = {} + for name in ("2m_temperature", "10m_u_component_of_wind"): + arr = group.create_array( + name, shape=(n, *inner), chunks=(spc, *inner), dtype="f4" + ) + data = rng.standard_normal((n, *inner)).astype("f4") + arr[:] = data + srcs[name] = data + return store, srcs + + +def read_store_time(store): + import zarr + + g = zarr.open_group(store=store, mode="r") + attrs = dict(g["time"].attrs) + return decode_cf_time(np.asarray(g["time"][:]), attrs["units"]) + + +def test_decode_cf_time_matches_manual_offset(): + # "hours since 1959-01-01" -> datetime64[ns]; index 4 is +24h. + values = np.array([0, 6, 12, 18, 24], dtype="i8") + out = decode_cf_time(values, TIME_UNITS) + assert out.dtype == np.dtype("datetime64[ns]") + assert out[0] == np.datetime64("1959-01-01T00:00") + assert out[4] == np.datetime64("1959-01-02T00:00") + + +def test_feed_contract(tmp_path): + """One batch: tensor layout, coord keys/order/dtypes, and lead-axis values.""" + store, _ = write_store(tmp_path) + variables = ["t2m", "u10m"] + var_map = {"t2m": "2m_temperature", "u10m": "10m_u_component_of_wind"} + leads = np.array([np.timedelta64(h, "h") for h in (0, 6, 12)]) + + feed = InSituForecastFeed( + store, + variables=variables, + var_map=var_map, + lead_times=leads, + sample_range=(0, 8), + batch_size=8, + ) + batches = list(feed) + feed.dataset.close() + + assert len(batches) == 1 + x, coords = batches[0] + # (time, lead_time, variable, lat, lon) + assert x.shape == (8, len(leads), len(variables), 4, 5) + assert x.dtype.is_floating_point + + assert list(coords.keys()) == ["time", "lead_time", "variable", "lat", "lon"] + assert coords["time"].dtype == np.dtype("datetime64[ns]") + assert coords["lead_time"].dtype == np.dtype("timedelta64[ns]") + assert np.array_equal(coords["variable"], np.array(variables)) + assert coords["lat"].dtype == np.float32 and coords["lon"].dtype == np.float32 + assert np.array_equal(coords["lead_time"], leads.astype("timedelta64[ns]")) + # init times are the first 8 store steps + assert np.array_equal(coords["time"], read_store_time(store)[:8]) + + +def test_feed_values_are_shift_views(tmp_path): + """Each (lead, variable) cell is a sample-axis shift view of the stored array.""" + store, srcs = write_store(tmp_path) + variables = ["t2m", "u10m"] + var_map = {"t2m": "2m_temperature", "u10m": "10m_u_component_of_wind"} + lead_steps = (0, 1, 2) # in units of the 6-h store step + leads = np.array([np.timedelta64(k * STEP_H, "h") for k in lead_steps]) + + feed = InSituForecastFeed( + store, + variables=variables, + var_map=var_map, + lead_times=leads, + sample_range=(0, 8), + batch_size=8, + ) + ((x, _coords),) = list(feed) + feed.dataset.close() + x = x.numpy() + + for li, k in enumerate(lead_steps): + for vi, vid in enumerate(variables): + src = srcs[var_map[vid]] + for t in range(8): # init index t -> valid index t + k + np.testing.assert_array_equal(x[t, li, vi], src[t + k]) + + +def test_decode_once_dedup(tmp_path): + """The thesis: an overlapping (init, lead) grid decodes each shared chunk once. + + 48 requested field-reads (8 inits x 3 leads x 2 vars) touch sample indices 0..9, which + span two fat chunks (spc=8) per variable -> exactly 4 unique decodes. + """ + store, _ = write_store(tmp_path, n=48, spc=8) + variables = ["t2m", "u10m"] + var_map = {"t2m": "2m_temperature", "u10m": "10m_u_component_of_wind"} + leads = np.array([np.timedelta64(h, "h") for h in (0, 6, 12)]) + + feed = InSituForecastFeed( + store, + variables=variables, + var_map=var_map, + lead_times=leads, + sample_range=(0, 8), + batch_size=8, + ) + list(feed) + decodes = feed.dataset.cache_misses + feed.dataset.close() + + requested = 8 * len(leads) * len(variables) + # touched sample indices 0..9 -> chunks {0, 1} (spc=8), per each of 2 variables + touched = {i + k for i in range(8) for k in (0, 1, 2)} + unique_chunks = len({idx // 8 for idx in touched}) * len(variables) + assert decodes == unique_chunks == 4 + assert decodes < requested + + +def test_transpose_inner_swaps_field_axes(tmp_path): + """A store laid out (lon, lat) yields (lat, lon) fields under transpose_inner=True.""" + store, srcs = write_store(tmp_path, lat=4, lon=5, inner_latlon=False) + feed = InSituForecastFeed( + store, + variables=["t2m"], + var_map={"t2m": "2m_temperature"}, + sample_range=(0, 8), + batch_size=8, + transpose_inner=True, + ) + x, coords = list(feed)[0] + feed.dataset.close() + + assert x.shape[-2:] == (4, 5) # (lat, lon) + assert coords["lat"].shape[0] == 4 and coords["lon"].shape[0] == 5 + # value at (0,0) is the stored (lon, lat) field transposed + np.testing.assert_array_equal(x.numpy()[0, 0, 0], srcs["2m_temperature"][0].T) + + +def test_lead_not_multiple_of_store_step_raises(tmp_path): + store, _ = write_store(tmp_path) # 6-h store step + with pytest.raises(ValueError, match="integer multiple of the store step"): + InSituForecastFeed( + store, + variables=["t2m"], + var_map={"t2m": "2m_temperature"}, + lead_times=np.array( + [np.timedelta64(90, "m")] + ), # 1.5 h, not a multiple of 6 h + sample_range=(0, 4), + ) From 2517d229dd3788b8aa282ed17b390782d47c76d6 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Thu, 9 Jul 2026 04:16:15 +0000 Subject: [PATCH 08/16] Enable + document the cross-run persistent cache in the insitubatch feed cache_dir now sets persist=True, so decoded chunks survive across runs (the flag was passed through but never persisted). Add a cold-vs-warm bench (bench_cache.py), a cross-run persistence test, and a README section framing it as a predownload replacement. Measured cold->warm: WB2 33->0 cloud fetches (~1.4x wall), ARCO 54->0 (~2.2x wall). --- earth2studio/data/insitu.py | 23 +++- recipes/insitubatch_hindcast/README.md | 28 +++++ recipes/insitubatch_hindcast/bench_cache.py | 121 ++++++++++++++++++++ test/data/test_insitu.py | 27 +++++ 4 files changed, 196 insertions(+), 3 deletions(-) create mode 100644 recipes/insitubatch_hindcast/bench_cache.py diff --git a/earth2studio/data/insitu.py b/earth2studio/data/insitu.py index 1e7d511b9..48c055b7f 100644 --- a/earth2studio/data/insitu.py +++ b/earth2studio/data/insitu.py @@ -62,7 +62,9 @@ ) from exc -def decode_cf_time(values: np.ndarray, units: str, calendar: str = "standard") -> np.ndarray: +def decode_cf_time( + values: np.ndarray, units: str, calendar: str = "standard" +) -> np.ndarray: """Decode a CF ``" since "`` integer time coordinate to datetime64[ns]. Handles the common reanalysis encoding (e.g. WB2/ARCO ERA5 store ``time`` as @@ -72,7 +74,9 @@ def decode_cf_time(values: np.ndarray, units: str, calendar: str = "standard") - do not fit E2S's ``datetime64[ns]`` coordinate contract, so the cast raises here -- the right place to fail for an out-of-contract analysis store. """ - dates = cftime.num2date(values, units, calendar=calendar, only_use_cftime_datetimes=False) + dates = cftime.num2date( + values, units, calendar=calendar, only_use_cftime_datetimes=False + ) return np.asarray(dates, dtype="datetime64[ns]") @@ -137,6 +141,13 @@ class InSituForecastFeed: with :func:`insitubatch.obstore_store` / :func:`insitubatch.fsspec_store` (e.g. anon public buckets). ``self.dataset`` exposes the underlying :class:`InSituDataset` for its ``cache_hits`` / ``cache_misses`` / ``resident_peak`` counters. + + Setting ``cache_dir`` turns on a **cross-run persistent cache**: the decoded chunks a run + touches are written there (decode-once, no reshard) and a later run over the same store + reads them from local disk as ``cache_hits`` instead of re-fetching the cloud. Because a + reanalysis store is static, this is a drop-in replacement for a pre-download step when the + *same* ground truth is scored repeatedly (many models, one fixed verification set). The + path is the cache identity -- use a fresh ``cache_dir`` when the store or variables change. """ def __init__( @@ -176,7 +187,11 @@ def __init__( # Sample-axis step of the store (dt); every lead must be an integer multiple of it. dt = self.time[1] - self.time[0] - leads = np.array([np.timedelta64(0, "ns")]) if lead_times is None else np.asarray(lead_times) + leads = ( + np.array([np.timedelta64(0, "ns")]) + if lead_times is None + else np.asarray(lead_times) + ) self.lead_time = leads.astype("timedelta64[ns]") steps = self.lead_time / dt if not np.all(steps == np.round(steps)): @@ -208,6 +223,8 @@ def __init__( batch_size=batch_size, shuffle=False, cache_dir=cache_dir, + # cache_dir set => cross-run persistent cache (not just an in-run spill tier) + persist=cache_dir is not None, max_inflight=max_inflight, ) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index 53011b308..e0c60fa8c 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -65,6 +65,30 @@ does not. That bounded-memory property — not just throughput — is the point (Persistence is a checkpoint-free model that exercises the real `create_iterator` seam on CPU; a real NVIDIA checkpoint — SFNO/FCN — is a drop-in with the same code on a GPU.) +## 3. `bench_cache.py` — cross-run persistent cache + +The intro's `predownload.py` exists so a re-scored campaign doesn't re-fetch the same ground +truth. `InSituForecastFeed(cache_dir=...)` gives that for free: the first run decodes each shared +chunk once **and** persists it to local disk; a later run over the same store reads those chunks +back as cache hits, touching the cloud zero times. No predownload step, no reshard, only the chunks +actually touched — and because a reanalysis store is static, the cache never goes stale. This is the +common eval shape: many models (or checkpoints) scored against one fixed verification set. + +```bash +python bench_cache.py --store wb2 --vars t2m u10m v10m --n-init 48 --max-lead-h 240 --cache-dir /mnt/nvme/insitu_cache +python bench_cache.py --store arco --vars t2m --n-init 12 --max-lead-h 48 --cache-dir /mnt/nvme/insitu_cache +``` + +| store | field size | cold → warm wall | **cloud fetches (cold → warm)** | +|-------|------------|------------------|----------------------------------| +| **WB2** 240×121 | 116 KB | 1.41 s → 1.05 s (~1.4×) | **33 → 0** | +| **ARCO** 721×1440 | 4 MB | 1.22 s → 0.55 s (~2.2×) | **54 → 0** | + +The deterministic result is **zero cloud fetches on re-score** — the warm run serves every chunk +from local disk. The wall speedup is secondary and scales with how IO-bound the cold fetch is (tiny +WB2 fields ~1.4×; 4 MB ARCO fields ~2.2×); it is *understated* on this box's cheap same-region reads +and grows under metered egress, requester-pays, or cross-region access. + ## How to read these numbers — framing insitubatch insitubatch is a **streaming batch loader** that trains/infers in place on cloud zarr: all @@ -98,5 +122,9 @@ causes duplicate reads.* raise the AFTER throughput further but is not what these numbers measure. - **Surface variables only** (`t2m`, `u10m`, `v10m`); pressure-level variables need level indexing, not yet wired in the adapter. +- **Persistent cache footprint.** The cache stores *decoded* chunks, so per-chunk bytes exceed the + compressed store — but it is bounded to the unique chunks touched (decode-once), not the dense + grid a predownload materializes. The `cache_dir` path is the cache identity; use a fresh one when + the store or variable set changes. - **The win is the IO-bound campaign** (many inits, verification-heavy — hindcast scoring, lagged ensembles). A single-IC long rollout is compute-bound, where the loader is a rounding error. diff --git a/recipes/insitubatch_hindcast/bench_cache.py b/recipes/insitubatch_hindcast/bench_cache.py new file mode 100644 index 000000000..133037e9b --- /dev/null +++ b/recipes/insitubatch_hindcast/bench_cache.py @@ -0,0 +1,121 @@ +"""Cross-run persistent-cache benchmark for the insitubatch verification feed. + +Scoring a hindcast campaign is rarely a one-shot: the *same* ERA5 verification set is read +again every time another model (or another checkpoint / hyperparameter) is scored against +it. Earth2Studio's eval recipe handles this with a ``predownload.py`` sentinel -- a separate +step that materializes a dense local copy before the run. + +``InSituForecastFeed(cache_dir=...)`` replaces that: the first run decodes each shared chunk +once (the dedup win) AND persists it to local disk; a second run over the same store reads +those chunks back as ``cache_hits`` instead of re-fetching the cloud -- no predownload step, +no reshard, and only the chunks actually touched. Because reanalysis is static the cache +never goes stale. + +This measures the second-run win: same verification window, run COLD (empty cache) then WARM +(cache populated), over gcsfs anon. +""" + +import argparse +import shutil +import time + +import numpy as np +from insitubatch import fsspec_store + +from earth2studio.data.insitu import InSituForecastFeed + +STORES = { + "wb2": { + "url": "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr", + "transpose_inner": True, + }, + "arco": { + "url": "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", + "transpose_inner": False, + }, +} +VAR_MAP = { + "t2m": "2m_temperature", + "u10m": "10m_u_component_of_wind", + "v10m": "10m_v_component_of_wind", +} + + +def anon_store(url): + return fsspec_store(url, token="anon", access="read_only") # noqa: S106 + + +def run(cfg, variables, start, n_init, leads_h, batch_size, max_inflight, cache_dir): + leads = np.array([np.timedelta64(h, "h") for h in leads_h]) + feed = InSituForecastFeed( + anon_store(cfg["url"]), + variables=variables, + var_map={v: VAR_MAP[v] for v in variables}, + lead_times=leads, + sample_range=(start, start + n_init), + batch_size=batch_size, + max_inflight=max_inflight, + cache_dir=cache_dir, + transpose_inner=cfg["transpose_inner"], + ) + t0 = time.perf_counter() + for _x, _coords in feed: + pass + wall = time.perf_counter() - t0 + hits, misses = feed.dataset.cache_hits, feed.dataset.cache_misses + feed.dataset.close() + return {"wall_s": wall, "hits": hits, "misses": misses} + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--store", choices=list(STORES), default="wb2") + p.add_argument("--vars", nargs="+", default=["t2m", "u10m", "v10m"]) + p.add_argument("--start", type=int, default=1000) + p.add_argument("--n-init", type=int, default=48) + p.add_argument("--lead-step-h", type=int, default=6) + p.add_argument("--max-lead-h", type=int, default=240) + p.add_argument("--batch-size", type=int, default=8) + p.add_argument("--max-inflight", type=int, default=32) + p.add_argument("--cache-dir", default="/tmp/insitu_cache_bench") # noqa: S108 + args = p.parse_args() + + cfg = STORES[args.store] + leads_h = list(range(args.lead_step_h, args.max_lead_h + 1, args.lead_step_h)) + requested = args.n_init * len(leads_h) * len(args.vars) + shutil.rmtree(args.cache_dir, ignore_errors=True) # start cold + + print( + f"[{args.store}] {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " + f"= {requested} requested field-reads ; cache_dir={args.cache_dir}" + ) + + common = ( + cfg, + args.vars, + args.start, + args.n_init, + leads_h, + args.batch_size, + args.max_inflight, + ) + cold = run(*common, args.cache_dir) + warm = run(*common, args.cache_dir) + + print("\n=== COLD (empty cache: fetch + decode + persist) ===") + print( + f" wall: {cold['wall_s']:.2f} s ; misses={cold['misses']} hits={cold['hits']}" + ) + print("\n=== WARM (cache populated: local-disk hits) ===") + print( + f" wall: {warm['wall_s']:.2f} s ; misses={warm['misses']} hits={warm['hits']}" + ) + print("\n=== HEADLINE ===") + print( + f" cross-run speedup : {cold['wall_s'] / warm['wall_s']:.1f}x ({cold['wall_s']:.2f}s -> {warm['wall_s']:.2f}s)" + ) + print(f" cloud fetches : {cold['misses']} cold -> {warm['misses']} warm") + + +if __name__ == "__main__": + main() diff --git a/test/data/test_insitu.py b/test/data/test_insitu.py index 8e21baeb6..84a94efa1 100644 --- a/test/data/test_insitu.py +++ b/test/data/test_insitu.py @@ -214,6 +214,33 @@ def test_transpose_inner_swaps_field_axes(tmp_path): np.testing.assert_array_equal(x.numpy()[0, 0, 0], srcs["2m_temperature"][0].T) +def test_persistent_cache_across_runs(tmp_path): + """cache_dir persists decoded chunks: a second run over the same store hits the cache.""" + store, _ = write_store(tmp_path, n=48, spc=8) + kw = { + "variables": ["t2m"], + "var_map": {"t2m": "2m_temperature"}, + "lead_times": np.array([np.timedelta64(h, "h") for h in (0, 6, 12)]), + "sample_range": (0, 8), + "batch_size": 8, + "cache_dir": str(tmp_path / "cache"), + } + + cold = InSituForecastFeed(store, **kw) + list(cold) + cold_misses, cold_hits = cold.dataset.cache_misses, cold.dataset.cache_hits + cold.dataset.close() + + warm = InSituForecastFeed(store, **kw) + list(warm) + warm_misses, warm_hits = warm.dataset.cache_misses, warm.dataset.cache_hits + warm.dataset.close() + + assert cold_misses > 0 and cold_hits == 0 # cold run fetches + populates the cache + assert warm_hits == cold_misses # warm run serves every chunk from disk + assert warm_misses == 0 # ... and fetches nothing + + def test_lead_not_multiple_of_store_step_raises(tmp_path): store, _ = write_store(tmp_path) # 6-h store step with pytest.raises(ValueError, match="integer multiple of the store step"): From 0e52120e240fdfb6fbdd124b3129c1838b53cbc5 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Thu, 9 Jul 2026 04:24:19 +0000 Subject: [PATCH 09/16] Note the persist-write overhead in the cache cold wall Clarify that bench_cache cold includes the one-time persist write, so it runs slightly above the persist-off de-dup figure in section 1. --- recipes/insitubatch_hindcast/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index e0c60fa8c..6e631b1e0 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -87,7 +87,8 @@ python bench_cache.py --store arco --vars t2m --n-init 12 --max-lead-h 48 --cach The deterministic result is **zero cloud fetches on re-score** — the warm run serves every chunk from local disk. The wall speedup is secondary and scales with how IO-bound the cold fetch is (tiny WB2 fields ~1.4×; 4 MB ARCO fields ~2.2×); it is *understated* on this box's cheap same-region reads -and grows under metered egress, requester-pays, or cross-region access. +and grows under metered egress, requester-pays, or cross-region access. The cold wall includes the +one-time persist write, so it runs slightly above the persist-off de-dup figure in §1. ## How to read these numbers — framing insitubatch From d401977670b067c8ac6ed65707215819d9c5128f Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Thu, 9 Jul 2026 05:03:18 +0000 Subject: [PATCH 10/16] Sort recipe imports for the current ruff config Rebasing onto NVIDIA main picked up its import-grouping rules; re-sort bench_hindcast.py and stream_score.py (stdlib / third-party / first-party). --- .../insitubatch_hindcast/bench_hindcast.py | 65 +++++++++++---- recipes/insitubatch_hindcast/stream_score.py | 82 ++++++++++++++----- 2 files changed, 111 insertions(+), 36 deletions(-) diff --git a/recipes/insitubatch_hindcast/bench_hindcast.py b/recipes/insitubatch_hindcast/bench_hindcast.py index b24e39808..c4228b9ff 100644 --- a/recipes/insitubatch_hindcast/bench_hindcast.py +++ b/recipes/insitubatch_hindcast/bench_hindcast.py @@ -23,9 +23,9 @@ import time import numpy as np +from insitubatch import fsspec_store from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time -from insitubatch import fsspec_store # store id -> (url, before source class, inner (H,W), transpose store(lon,lat)->(lat,lon)) STORES = { @@ -77,11 +77,17 @@ def run_after(cfg, variables, start, n_init, leads_h, batch_size, max_inflight): t0 = time.perf_counter() n_rows = 0 for x, _coords in feed: - n_rows += x.shape[0] # gather returns eager numpy; no touch needed to force decode + n_rows += x.shape[ + 0 + ] # gather returns eager numpy; no touch needed to force decode wall = time.perf_counter() - t0 feed.dataset.close() - return {"wall_s": wall, "init_rows": n_rows, - "chunk_decodes": feed.dataset.cache_misses, "resident_peak": feed.dataset.resident_peak} + return { + "wall_s": wall, + "init_rows": n_rows, + "chunk_decodes": feed.dataset.cache_misses, + "resident_peak": feed.dataset.resident_peak, + } def run_before(before_cls, variables, init_times64, leads_h): @@ -90,7 +96,10 @@ def run_before(before_cls, variables, init_times64, leads_h): leads_td = [np.timedelta64(h, "h") for h in leads_h] t0 = time.perf_counter() for it in init_dt: - valid = [(np.datetime64(it) + td).astype("datetime64[s]").astype("O") for td in leads_td] + valid = [ + (np.datetime64(it) + td).astype("datetime64[s]").astype("O") + for td in leads_td + ] src(valid, list(variables)) # realistic per-init verification fetch return {"wall_s": time.perf_counter() - t0} @@ -117,12 +126,18 @@ def main(): g = zarr.open_group(store=anon_store(cfg["url"]), mode="r") attrs = dict(g["time"].attrs) - times64 = decode_cf_time(np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard")) + times64 = decode_cf_time( + np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard") + ) init_times64 = times64[args.start : args.start + args.n_init] - print(f"[{args.store}] grid: {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " - f"= {requested} requested field-reads ({requested*cfg['field_bytes']/1e9:.2f} GB naive)") - print(f"leads: {args.lead_step_h}h..{args.max_lead_h}h ; vars: {args.vars} ; repeats: {args.repeats}") + print( + f"[{args.store}] grid: {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " + f"= {requested} requested field-reads ({requested*cfg['field_bytes']/1e9:.2f} GB naive)" + ) + print( + f"leads: {args.lead_step_h}h..{args.max_lead_h}h ; vars: {args.vars} ; repeats: {args.repeats}" + ) def med3(w): w = sorted(w) @@ -132,27 +147,45 @@ def med3(w): after_walls, before_walls = [], [] decodes = resident = None for r in range(args.repeats): - a = run_after(cfg, args.vars, args.start, args.n_init, leads_h, args.batch_size, args.max_inflight) + a = run_after( + cfg, + args.vars, + args.start, + args.n_init, + leads_h, + args.batch_size, + args.max_inflight, + ) after_walls.append(a["wall_s"]) decodes, resident = a["chunk_decodes"], a["resident_peak"] if not args.skip_before: - before_walls.append(run_before(before_cls, args.vars, init_times64, leads_h)["wall_s"]) - print(f" repeat {r+1}/{args.repeats}: after={after_walls[-1]:.2f}s" - + (f" before={before_walls[-1]:.2f}s" if before_walls else "")) + before_walls.append( + run_before(before_cls, args.vars, init_times64, leads_h)["wall_s"] + ) + print( + f" repeat {r+1}/{args.repeats}: after={after_walls[-1]:.2f}s" + + (f" before={before_walls[-1]:.2f}s" if before_walls else "") + ) dedup = requested / decodes a_med, a_lo, a_hi = med3(after_walls) print("\n=== AFTER (insitubatch, gcsfs anon) ===") print(f" wall (med/min/max): {a_med:.2f} / {a_lo:.2f} / {a_hi:.2f} s") - print(f" chunk decodes : {decodes} ({decodes*cfg['field_bytes']*cfg['chunk_steps']/1e9:.2f} GB)") - print(f" dedup ratio : {dedup:.1f}x ({requested} requested -> {decodes} decoded)") + print( + f" chunk decodes : {decodes} ({decodes*cfg['field_bytes']*cfg['chunk_steps']/1e9:.2f} GB)" + ) + print( + f" dedup ratio : {dedup:.1f}x ({requested} requested -> {decodes} decoded)" + ) print(f" resident peak : {resident} chunks") if before_walls: b_med, b_lo, b_hi = med3(before_walls) print("\n=== BEFORE (E2S fetch, gcsfs anon, cache off) ===") print(f" wall (med/min/max): {b_med:.2f} / {b_lo:.2f} / {b_hi:.2f} s") print("\n=== HEADLINE (medians) ===") - print(f" speedup : {b_med/a_med:.1f}x wall ({b_med:.2f}s -> {a_med:.2f}s)") + print( + f" speedup : {b_med/a_med:.1f}x wall ({b_med:.2f}s -> {a_med:.2f}s)" + ) print(f" read reduction : {dedup:.1f}x fewer chunk decodes") diff --git a/recipes/insitubatch_hindcast/stream_score.py b/recipes/insitubatch_hindcast/stream_score.py index d253c426e..bac1a1074 100644 --- a/recipes/insitubatch_hindcast/stream_score.py +++ b/recipes/insitubatch_hindcast/stream_score.py @@ -24,14 +24,18 @@ import numpy as np import torch +from insitubatch import fsspec_store from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time from earth2studio.models.px import Persistence from earth2studio.utils.coords import map_coords -from insitubatch import fsspec_store URL = "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr" -VAR_MAP = {"t2m": "2m_temperature", "u10m": "10m_u_component_of_wind", "v10m": "10m_v_component_of_wind"} +VAR_MAP = { + "t2m": "2m_temperature", + "u10m": "10m_u_component_of_wind", + "v10m": "10m_v_component_of_wind", +} DT = np.timedelta64(6, "h") @@ -40,7 +44,9 @@ def anon_store(): def peak_rss_gb(): - return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 # ru_maxrss is KB on Linux + return ( + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 + ) # ru_maxrss is KB on Linux class RmseAccumulator: @@ -62,9 +68,13 @@ def table(self): def build_feed(variables, start, n_init, leads_h, batch_size): leads = np.array([np.timedelta64(h, "h") for h in leads_h]) # includes 0 (the IC) return InSituForecastFeed( - anon_store(), variables=variables, var_map={v: VAR_MAP[v] for v in variables}, - lead_times=leads, sample_range=(start, start + n_init), - batch_size=batch_size, transpose_inner=True, + anon_store(), + variables=variables, + var_map={v: VAR_MAP[v] for v in variables}, + lead_times=leads, + sample_range=(start, start + n_init), + batch_size=batch_size, + transpose_inner=True, ) @@ -80,8 +90,13 @@ def score_window(model, x_all, coords_all, leads_h, nsteps, acc): """ ic_x = x_all[:, 0:1] ic_coords = OrderedDict( - [("time", coords_all["time"]), ("lead_time", coords_all["lead_time"][0:1]), - ("variable", coords_all["variable"]), ("lat", coords_all["lat"]), ("lon", coords_all["lon"])] + [ + ("time", coords_all["time"]), + ("lead_time", coords_all["lead_time"][0:1]), + ("variable", coords_all["variable"]), + ("lat", coords_all["lat"]), + ("lon", coords_all["lon"]), + ] ) ic_x, ic_coords = map_coords(ic_x, ic_coords, model.input_coords()) for step, (fx, fcoords) in enumerate(model.create_iterator(ic_x, ic_coords)): @@ -96,7 +111,10 @@ def run_insitu(variables, start, n_init, leads_h, nsteps, batch_size): feed = build_feed(variables, start, n_init, leads_h, batch_size) model = make_model(variables, feed) acc = RmseAccumulator() - for x_all, coords_all in feed: # one batch per window (stream) or the whole campaign (dense) + for ( + x_all, + coords_all, + ) in feed: # one batch per window (stream) or the whole campaign (dense) score_window(model, x_all, coords_all, leads_h, nsteps, acc) decodes = feed.dataset.cache_misses feed.dataset.close() @@ -112,26 +130,46 @@ def run_e2s(variables, start, n_init, leads_h, nsteps): g = zarr.open_group(store=g_store, mode="r") attrs = dict(g["time"].attrs) - times = decode_cf_time(np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard")) + times = decode_cf_time( + np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard") + ) inits = times[start : start + n_init] - feed = build_feed(variables, start, n_init, leads_h, batch_size=n_init) # for coords/model only + feed = build_feed( + variables, start, n_init, leads_h, batch_size=n_init + ) # for coords/model only model = make_model(variables, feed) feed.dataset.close() # Dense predownload: fetch every (init, lead) valid time (redundant reads), materialize. lat = np.asarray(g["latitude"][:]).astype(np.float32) lon = np.asarray(g["longitude"][:]).astype(np.float32) - dense = np.empty((n_init, len(leads_h), len(variables), len(lat), len(lon)), dtype=np.float32) + dense = np.empty( + (n_init, len(leads_h), len(variables), len(lat), len(lon)), dtype=np.float32 + ) for i, it in enumerate(inits): # Realistic per-init verification fetch: all leads for this init in one call. - valids = [(np.datetime64(it) + np.timedelta64(h, "h")).astype("datetime64[s]").astype("O") - for h in leads_h] + valids = [ + (np.datetime64(it) + np.timedelta64(h, "h")) + .astype("datetime64[s]") + .astype("O") + for h in leads_h + ] da = src(valids, list(variables)) # (L+1, V, lat, lon) xr.DataArray dense[i] = np.asarray(da.values) x_all = torch.from_numpy(dense) coords_all = OrderedDict( - [("time", inits), ("lead_time", np.array([np.timedelta64(h, "h") for h in leads_h], dtype="timedelta64[ns]")), - ("variable", np.asarray(variables)), ("lat", lat), ("lon", lon)] + [ + ("time", inits), + ( + "lead_time", + np.array( + [np.timedelta64(h, "h") for h in leads_h], dtype="timedelta64[ns]" + ), + ), + ("variable", np.asarray(variables)), + ("lat", lat), + ("lon", lon), + ] ) acc = RmseAccumulator() score_window(model, x_all, coords_all, leads_h, nsteps, acc) @@ -160,13 +198,17 @@ def main(): wall = time.perf_counter() - t0 rmse = acc.table() - print(f"mode={args.mode} grid={args.n_init}x{args.n_leads}x{len(args.vars)} window={args.window}") + print( + f"mode={args.mode} grid={args.n_init}x{args.n_leads}x{len(args.vars)} window={args.window}" + ) print(f" wall : {wall:.2f} s") print(f" peak RSS : {peak_rss_gb():.2f} GB") print(f" reads/decodes: {reads}") - print(f" RMSE @ 24h/120h/240h: " - f"{rmse.get(24, float('nan')):.3f} / {rmse.get(120, float('nan')):.3f} / " - f"{rmse.get(min(240, nsteps*6), float('nan')):.3f}") + print( + f" RMSE @ 24h/120h/240h: " + f"{rmse.get(24, float('nan')):.3f} / {rmse.get(120, float('nan')):.3f} / " + f"{rmse.get(min(240, nsteps*6), float('nan')):.3f}" + ) if __name__ == "__main__": From 9f6c45562ae0998b89e262b94ee6473ca54077f9 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Thu, 9 Jul 2026 05:51:52 +0000 Subject: [PATCH 11/16] Reject leads that read past the store edge in the insitubatch feed Greptile flagged that a history (<0) or verification (>0) lead near a store boundary reads outside the requested init window: the engine silently drops those edge anchors, so a scoring campaign would cover a shorter window than asked without any signal. Validate the requested sample_range against valid_anchor_range(lead_steps, n_samples) and raise an actionable ValueError; when sample_range is unset, default to the in-bounds init window. Covered by boundary tests (past-end, before-start, none-defaults). --- earth2studio/data/insitu.py | 26 +++++++++++++++++++++ test/data/test_insitu.py | 46 +++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/earth2studio/data/insitu.py b/earth2studio/data/insitu.py index 48c055b7f..586d9c714 100644 --- a/earth2studio/data/insitu.py +++ b/earth2studio/data/insitu.py @@ -55,6 +55,7 @@ open_geometries, split_by_chunk, to_torch, + valid_anchor_range, ) except ImportError as exc: # pragma: no cover - optional integration raise ImportError( @@ -136,6 +137,10 @@ class InSituForecastFeed: sample-axis ``shift`` view of one stored array, so the ``(init, lead)`` grid decodes each shared chunk exactly once (the win over per-``(time, lead)`` ``fetch_data``). + Init times whose leads would read past either end of the store are rejected with a + ``ValueError`` (rather than silently dropped); with ``sample_range`` unset the feed spans + exactly the in-bounds init window for the requested leads. + ``variables`` are the ids to expose on the ``variable`` coordinate; ``var_map`` maps them to store array names when they differ (e.g. ``t2m -> 2m_temperature``). Build ``store`` with :func:`insitubatch.obstore_store` / :func:`insitubatch.fsspec_store` (e.g. anon @@ -202,6 +207,27 @@ def __init__( arrays = [vmap[v] for v in self.variables] opened = open_geometries(store, variables=arrays) + + # Each lead shifts the read to anchor + step, so init times near a store edge whose + # shifted read would leave [0, n_samples) are unusable. valid_anchor_range gives the + # in-bounds init window for these leads; the engine would silently drop out-of-range + # anchors, so validate here instead of scoring a shorter window than requested. + n_samples = opened[arrays[0]].n_samples + lo, hi = valid_anchor_range(self.lead_steps.tolist(), n_samples) + if lo >= hi: + raise ValueError( + f"lead steps {self.lead_steps.tolist()} span more than the store's " + f"{n_samples} samples; no init time can satisfy every lead" + ) + if sample_range is None: + sample_range = (lo, hi) # every init whose leads all fall within the store + elif sample_range[0] < lo or sample_range[1] > hi: + raise ValueError( + f"sample_range {sample_range} with lead steps {self.lead_steps.tolist()} " + f"reads outside the store [0, {n_samples}); the valid init range for these " + f"leads is [{lo}, {hi})" + ) + # One shifted geometry per (lead, variable); label grid indexes them for the stacker. geometries: dict[str, object] = {} self.labels: list[list[str]] = [] diff --git a/test/data/test_insitu.py b/test/data/test_insitu.py index 84a94efa1..6b46c9f75 100644 --- a/test/data/test_insitu.py +++ b/test/data/test_insitu.py @@ -253,3 +253,49 @@ def test_lead_not_multiple_of_store_step_raises(tmp_path): ), # 1.5 h, not a multiple of 6 h sample_range=(0, 4), ) + + +def test_verification_lead_past_store_end_raises(tmp_path): + """A positive lead whose read leaves the store end is rejected, not silently dropped.""" + store, _ = write_store(tmp_path, n=48) # 6-h step; 240 h = 40 steps + with pytest.raises(ValueError, match=r"outside the store|valid init range"): + InSituForecastFeed( + store, + variables=["t2m"], + var_map={"t2m": "2m_temperature"}, + lead_times=np.array([np.timedelta64(h, "h") for h in (0, 240)]), + sample_range=(0, 20), # init 20 + 40-step lead -> index 60 >> 48 + ) + + +def test_history_lead_before_store_start_raises(tmp_path): + """A negative (history) lead whose read precedes index 0 is rejected.""" + store, _ = write_store(tmp_path, n=48) + with pytest.raises(ValueError, match=r"outside the store|valid init range"): + InSituForecastFeed( + store, + variables=["t2m"], + var_map={"t2m": "2m_temperature"}, + lead_times=np.array( + [np.timedelta64(h, "h") for h in (-240, 0)] + ), # -40 steps + sample_range=(0, 44), # init 0 - 40-step history -> index -40 + ) + + +def test_sample_range_none_defaults_to_valid_window(tmp_path): + """With no sample_range, the feed covers exactly the inits whose leads all fit the store.""" + store, _ = write_store(tmp_path, n=48) # spc=8 + feed = InSituForecastFeed( + store, + variables=["t2m"], + var_map={"t2m": "2m_temperature"}, + lead_times=np.array( + [np.timedelta64(h, "h") for h in (0, 6, 12)] + ), # steps 0,1,2 + batch_size=8, + ) + n_inits = sum(x.shape[0] for x, _ in feed) + feed.dataset.close() + # valid_anchor_range([0,1,2], 48) = [0, 46): the last two inits would read past the end + assert n_inits == 46 From f78c36891c966245e17a23126fe3f801295e4437 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Tue, 4 Aug 2026 02:06:20 +0000 Subject: [PATCH 12/16] Meet the new GitHub CI gates in the hindcast recipe The CI transition (#1021) runs `make lint`, `make license` and `make interrogate` over all files, so the recipe scripts are now gated where they previously were not: - Annotate the three recipe scripts (mypy `-a` runs with `disallow_untyped_defs`; library objects stay `Any` since the hook resolves imports with `follow_imports=skip`). - Add the NVIDIA SPDX headers the license check requires. - Rewrap the README over markdownlint's line-length budget: the shell blocks use line continuations, and the WB2 table row is trimmed to the 100-col table limit. No behavior change; `--help` and the offline tests are unaffected. Co-Authored-By: Claude Opus 5 --- recipes/insitubatch_hindcast/README.md | 20 +++--- recipes/insitubatch_hindcast/bench_cache.py | 34 +++++++++-- .../insitubatch_hindcast/bench_hindcast.py | 47 +++++++++++--- recipes/insitubatch_hindcast/stream_score.py | 61 +++++++++++++++---- 4 files changed, 131 insertions(+), 31 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index 6e631b1e0..6e43ca9cc 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -29,14 +29,16 @@ onto far fewer stored chunks. BEFORE = E2S's per-init `fetch_data`; AFTER = the (each lead a sample-axis `shift` view; each shared chunk decoded once). ```bash -python bench_hindcast.py --store wb2 --vars t2m u10m v10m --n-init 48 --max-lead-h 240 --repeats 5 -python bench_hindcast.py --store arco --vars t2m --n-init 24 --lead-step-h 6 --max-lead-h 144 --repeats 3 +python bench_hindcast.py --store wb2 --vars t2m u10m v10m \ + --n-init 48 --max-lead-h 240 --repeats 5 +python bench_hindcast.py --store arco --vars t2m --lead-step-h 6 \ + --n-init 24 --max-lead-h 144 --repeats 3 ``` | store | layout | requested reads | unique decodes | **wall speedup** | |-------|--------|-----------------|----------------|------------------| -| **WB2** 240×121 6-h | `chunks=(8,240,121)` (fat) | 5760 | **33** (174×) | **15.4×** (14.0 s → 0.91 s) | -| **ARCO** 721×1440 1-h | `chunks=(1,721,1440)` (chunk-1) | 576 | 162 (3.6×) | ~1.9× | +| **WB2** 240×121 6-h | `chunks=(8,240,121)` fat | 5760 | **33** (174×) | **15.4×** (14.0→0.91 s) | +| **ARCO** 721×1440 1-h | `chunks=(1,721,1440)` chunk-1 | 576 | 162 (3.6×) | ~1.9× | WB2's fat time-chunk amortizes 8 steps per read, so the de-dup ratio is large and the fields are small — insitubatch dominates. ARCO is the **honest** case (see caveats). @@ -49,7 +51,9 @@ roll out a window of inits, score each lead against a just-read verification sli modes, all producing **identical RMSE** (a correctness check): ```bash -for m in e2s dense stream; do python stream_score.py --mode $m --n-init 120 --n-leads 40; done +for m in e2s dense stream; do + python stream_score.py --mode $m --n-init 120 --n-leads 40 +done ``` | mode | wall | **peak RSS** | field reads | @@ -75,8 +79,10 @@ actually touched — and because a reanalysis store is static, the cache never g common eval shape: many models (or checkpoints) scored against one fixed verification set. ```bash -python bench_cache.py --store wb2 --vars t2m u10m v10m --n-init 48 --max-lead-h 240 --cache-dir /mnt/nvme/insitu_cache -python bench_cache.py --store arco --vars t2m --n-init 12 --max-lead-h 48 --cache-dir /mnt/nvme/insitu_cache +python bench_cache.py --store wb2 --vars t2m u10m v10m \ + --n-init 48 --max-lead-h 240 --cache-dir /mnt/nvme/insitu_cache +python bench_cache.py --store arco --vars t2m \ + --n-init 12 --max-lead-h 48 --cache-dir /mnt/nvme/insitu_cache ``` | store | field size | cold → warm wall | **cloud fetches (cold → warm)** | diff --git a/recipes/insitubatch_hindcast/bench_cache.py b/recipes/insitubatch_hindcast/bench_cache.py index 133037e9b..cc4b585f7 100644 --- a/recipes/insitubatch_hindcast/bench_cache.py +++ b/recipes/insitubatch_hindcast/bench_cache.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Cross-run persistent-cache benchmark for the insitubatch verification feed. Scoring a hindcast campaign is rarely a one-shot: the *same* ERA5 verification set is read @@ -18,13 +34,14 @@ import argparse import shutil import time +from typing import Any import numpy as np from insitubatch import fsspec_store from earth2studio.data.insitu import InSituForecastFeed -STORES = { +STORES: dict[str, dict[str, Any]] = { "wb2": { "url": "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr", "transpose_inner": True, @@ -41,11 +58,20 @@ } -def anon_store(url): +def anon_store(url: str) -> Any: return fsspec_store(url, token="anon", access="read_only") # noqa: S106 -def run(cfg, variables, start, n_init, leads_h, batch_size, max_inflight, cache_dir): +def run( + cfg: dict[str, Any], + variables: list[str], + start: int, + n_init: int, + leads_h: list[int], + batch_size: int, + max_inflight: int, + cache_dir: str, +) -> dict[str, Any]: leads = np.array([np.timedelta64(h, "h") for h in leads_h]) feed = InSituForecastFeed( anon_store(cfg["url"]), @@ -67,7 +93,7 @@ def run(cfg, variables, start, n_init, leads_h, batch_size, max_inflight, cache_ return {"wall_s": wall, "hits": hits, "misses": misses} -def main(): +def main() -> None: p = argparse.ArgumentParser() p.add_argument("--store", choices=list(STORES), default="wb2") p.add_argument("--vars", nargs="+", default=["t2m", "u10m", "v10m"]) diff --git a/recipes/insitubatch_hindcast/bench_hindcast.py b/recipes/insitubatch_hindcast/bench_hindcast.py index c4228b9ff..a3123dd1a 100644 --- a/recipes/insitubatch_hindcast/bench_hindcast.py +++ b/recipes/insitubatch_hindcast/bench_hindcast.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Before/after hindcast verification-read benchmark on WB2 / ARCO ERA5. Scenario: score an ``N_init x len(leads)`` forecast grid against ERA5. Every ``(init, lead)`` @@ -21,6 +37,7 @@ import argparse import time +from typing import Any import numpy as np from insitubatch import fsspec_store @@ -28,7 +45,7 @@ from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time # store id -> (url, before source class, inner (H,W), transpose store(lon,lat)->(lat,lon)) -STORES = { +STORES: dict[str, dict[str, Any]] = { "wb2": { "url": "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr", "before": "earth2studio.data.wb2:WB2ERA5_121x240", @@ -51,18 +68,26 @@ } -def anon_store(url): +def anon_store(url: str) -> Any: return fsspec_store(url, token="anon", access="read_only") # noqa: S106 -def load_before_cls(spec): +def load_before_cls(spec: str) -> Any: mod, _, name = spec.partition(":") import importlib return getattr(importlib.import_module(mod), name) -def run_after(cfg, variables, start, n_init, leads_h, batch_size, max_inflight): +def run_after( + cfg: dict[str, Any], + variables: list[str], + start: int, + n_init: int, + leads_h: list[int], + batch_size: int, + max_inflight: int, +) -> dict[str, Any]: leads = np.array([np.timedelta64(h, "h") for h in leads_h]) feed = InSituForecastFeed( anon_store(cfg["url"]), @@ -90,7 +115,9 @@ def run_after(cfg, variables, start, n_init, leads_h, batch_size, max_inflight): } -def run_before(before_cls, variables, init_times64, leads_h): +def run_before( + before_cls: Any, variables: list[str], init_times64: Any, leads_h: list[int] +) -> dict[str, Any]: src = before_cls(cache=False, verbose=False) init_dt = init_times64.astype("datetime64[s]").astype("O") leads_td = [np.timedelta64(h, "h") for h in leads_h] @@ -104,7 +131,7 @@ def run_before(before_cls, variables, init_times64, leads_h): return {"wall_s": time.perf_counter() - t0} -def main(): +def main() -> None: p = argparse.ArgumentParser() p.add_argument("--store", choices=list(STORES), default="wb2") p.add_argument("--vars", nargs="+", default=["t2m"]) @@ -139,13 +166,15 @@ def main(): f"leads: {args.lead_step_h}h..{args.max_lead_h}h ; vars: {args.vars} ; repeats: {args.repeats}" ) - def med3(w): + def med3(w: list[float]) -> tuple[float, float, float]: w = sorted(w) return w[len(w) // 2], w[0], w[-1] before_cls = load_before_cls(cfg["before"]) - after_walls, before_walls = [], [] - decodes = resident = None + after_walls: list[float] = [] + before_walls: list[float] = [] + decodes: Any = None + resident: Any = None for r in range(args.repeats): a = run_after( cfg, diff --git a/recipes/insitubatch_hindcast/stream_score.py b/recipes/insitubatch_hindcast/stream_score.py index bac1a1074..cbda0c227 100644 --- a/recipes/insitubatch_hindcast/stream_score.py +++ b/recipes/insitubatch_hindcast/stream_score.py @@ -1,3 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Streaming vs dense hindcast scoring: interleave verification with the model rollout. The Earth2Studio pattern (`fetch_data -> map_coords -> create_iterator`) materializes the whole @@ -21,6 +37,7 @@ import resource import time from collections import OrderedDict +from typing import Any import numpy as np import torch @@ -39,11 +56,11 @@ DT = np.timedelta64(6, "h") -def anon_store(): +def anon_store() -> Any: return fsspec_store(URL, token="anon", access="read_only") # noqa: S106 -def peak_rss_gb(): +def peak_rss_gb() -> float: return ( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 ) # ru_maxrss is KB on Linux @@ -52,20 +69,26 @@ def peak_rss_gb(): class RmseAccumulator: """Per-lead running MSE over (init, var, lat, lon); sqrt at the end.""" - def __init__(self): + def __init__(self) -> None: self.sse: dict[int, float] = {} self.n: dict[int, int] = {} - def update(self, lead_h: int, pred: torch.Tensor, truth: torch.Tensor): + def update(self, lead_h: int, pred: torch.Tensor, truth: torch.Tensor) -> None: e = (pred - truth).float() self.sse[lead_h] = self.sse.get(lead_h, 0.0) + float((e * e).sum()) self.n[lead_h] = self.n.get(lead_h, 0) + e.numel() - def table(self): + def table(self) -> dict[int, float]: return {h: (self.sse[h] / self.n[h]) ** 0.5 for h in sorted(self.sse)} -def build_feed(variables, start, n_init, leads_h, batch_size): +def build_feed( + variables: list[str], + start: int, + n_init: int, + leads_h: list[int], + batch_size: int, +) -> InSituForecastFeed: leads = np.array([np.timedelta64(h, "h") for h in leads_h]) # includes 0 (the IC) return InSituForecastFeed( anon_store(), @@ -78,12 +101,19 @@ def build_feed(variables, start, n_init, leads_h, batch_size): ) -def make_model(variables, feed): +def make_model(variables: list[str], feed: InSituForecastFeed) -> Persistence: domain = OrderedDict([("lat", feed.lat), ("lon", feed.lon)]) return Persistence(variable=variables, domain_coords=domain, history=1, dt=DT) -def score_window(model, x_all, coords_all, leads_h, nsteps, acc): +def score_window( + model: Persistence, + x_all: torch.Tensor, + coords_all: OrderedDict[str, Any], + leads_h: list[int], + nsteps: int, + acc: "RmseAccumulator", +) -> None: """Roll Persistence out over one window and score each lead vs the pre-read verification slice. ``x_all`` is (W, L+1, V, lat, lon) with lead index 0 = IC and index k = verification at leads_h[k]. @@ -107,7 +137,14 @@ def score_window(model, x_all, coords_all, leads_h, nsteps, acc): break -def run_insitu(variables, start, n_init, leads_h, nsteps, batch_size): +def run_insitu( + variables: list[str], + start: int, + n_init: int, + leads_h: list[int], + nsteps: int, + batch_size: int, +) -> tuple["RmseAccumulator", int]: feed = build_feed(variables, start, n_init, leads_h, batch_size) model = make_model(variables, feed) acc = RmseAccumulator() @@ -121,7 +158,9 @@ def run_insitu(variables, start, n_init, leads_h, nsteps, batch_size): return acc, decodes -def run_e2s(variables, start, n_init, leads_h, nsteps): +def run_e2s( + variables: list[str], start: int, n_init: int, leads_h: list[int], nsteps: int +) -> tuple["RmseAccumulator", int]: from earth2studio.data.wb2 import WB2ERA5_121x240 src = WB2ERA5_121x240(cache=False, verbose=False) @@ -176,7 +215,7 @@ def run_e2s(variables, start, n_init, leads_h, nsteps): return acc, n_init * len(leads_h) * len(variables) -def main(): +def main() -> None: p = argparse.ArgumentParser() p.add_argument("--mode", choices=["stream", "dense", "e2s"], required=True) p.add_argument("--vars", nargs="+", default=["t2m", "u10m", "v10m"]) From 3374fcf97cd5c5462afcd0d04270b736b307f82b Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:41:53 +0000 Subject: [PATCH 13/16] Re-measure the hindcast recipe over obstore on both sides MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Earth2Studio's zarr data sources migrated to obstore in #955 while this PR was in review, so the comparison had silently become obstore (before) vs gcsfs (after). Move the feed to insitubatch's `obstore_store` to match, and refresh every number. - Both sides now read via obstore anon; the de-dup ratios are unchanged (33 decodes on WB2, 162 on ARCO), which is the point -- the win is read planning, not transport. - Fix ARCO's `--start`: its axis is `hours since 1900-01-01` spanning 1900-2050, but chunks outside ~1940-2023 were never written and read back as all-NaN fill without a network request. The default of 1000 was 1900, so the documented ARCO command in §1 failed outright and §3's row measured nothing. `--start` now defaults per store, and both benchmarks raise if a window reads back entirely NaN. - `bench_cache.py` confines its cold-start wipe to `/bench_cold_warm//` instead of `rmtree`-ing the user-supplied path. - Refresh §1/§2/§3 with 5-10 repeat medians, and record that the earlier figures were taken over gcsfs before #955 under different logging conditions -- superseded, not a controlled comparison. Co-Authored-By: Claude Opus 5 --- recipes/insitubatch_hindcast/README.md | 98 ++++++++++++++----- recipes/insitubatch_hindcast/bench_cache.py | 37 ++++--- .../insitubatch_hindcast/bench_hindcast.py | 37 ++++--- recipes/insitubatch_hindcast/stream_score.py | 4 +- 4 files changed, 123 insertions(+), 53 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index 6e43ca9cc..ec352f92c 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -12,14 +12,27 @@ de-duplication — exactly what insitubatch removes. ## Setup +insitubatch is declared in earth2studio's `data` extra (needs Python >= 3.12): + ```bash -# insitubatch is declared in earth2studio's `data` extra (needs Python >= 3.12): uv sync --extra data ``` Both stores are anonymous public GCS buckets (WeatherBench2 ERA5, ARCO ERA5); no credentials -needed. Every measurement below reads over **gcsfs anon on both the before and after side**, so the -delta isolates insitubatch's read-planning + streaming — it is *not* an obstore-vs-gcsfs artifact. +needed. Every measurement below reads over **obstore anon on both the before and after side**, so +the delta isolates insitubatch's read-planning + streaming rather than the storage backend. + +Run the benchmarks with Earth2Studio's per-fetch debug logging suppressed, so the timed region is +the same on both sides — it emits one line per `(time, variable)`, which is 5760 lines in §1 and +14 760 in §2, all on the baseline leg: + +```bash +export LOGURU_LEVEL=INFO +``` + +(`LOGURU_LEVEL` configures loguru's default handler. It gates these benchmarks, but several other +Earth2Studio data modules call `logger.remove()` at import and re-add a handler with no level, for +which the variable has no effect.) ## 1. `bench_hindcast.py` — verification-read de-duplication @@ -32,16 +45,18 @@ onto far fewer stored chunks. BEFORE = E2S's per-init `fetch_data`; AFTER = the python bench_hindcast.py --store wb2 --vars t2m u10m v10m \ --n-init 48 --max-lead-h 240 --repeats 5 python bench_hindcast.py --store arco --vars t2m --lead-step-h 6 \ - --n-init 24 --max-lead-h 144 --repeats 3 + --n-init 24 --max-lead-h 144 --repeats 10 ``` | store | layout | requested reads | unique decodes | **wall speedup** | |-------|--------|-----------------|----------------|------------------| -| **WB2** 240×121 6-h | `chunks=(8,240,121)` fat | 5760 | **33** (174×) | **15.4×** (14.0→0.91 s) | -| **ARCO** 721×1440 1-h | `chunks=(1,721,1440)` chunk-1 | 576 | 162 (3.6×) | ~1.9× | +| **WB2** 240×121 6-h | `chunks=(8,240,121)` fat | 5760 | **33** (174×) | **12.8×** (10.84→0.84 s) | +| **ARCO** 721×1440 1-h | `chunks=(1,721,1440)` chunk-1 | 576 | 162 (3.6×) | 1.5× (3.91→2.57 s) | WB2's fat time-chunk amortizes 8 steps per read, so the de-dup ratio is large and the fields are -small — insitubatch dominates. ARCO is the **honest** case (see caveats). +small — insitubatch dominates. ARCO is the **honest** case (see caveats). ARCO's per-repeat spread +is wide on this box (±15% on both legs); its row is the median of 10 repeats, and the BEFORE leg +converges downward with sample size as HTTP connections warm. ## 2. `stream_score.py` — streaming vs dense materialization @@ -58,12 +73,17 @@ done | mode | wall | **peak RSS** | field reads | |------|------|--------------|-------------| -| `e2s` — dense predownload (redundant reads) | 39.6 s | 3.04 GB | 14 760 | -| `dense` — insitubatch, `batch_size=N` | 4.8 s | 7.53 GB | 60 | -| `stream` — insitubatch, `batch_size=W` | 3.3 s | **1.81 GB** | 60 | +| `e2s` — dense predownload (redundant reads) | 29.0 s | 3.10 GB | 14 760 | +| `dense` — insitubatch, `batch_size=N` | 4.3 s | 7.63 GB | 60 | +| `stream` — insitubatch, `batch_size=W` | 2.9 s | **1.85 GB** | 60 | + +All three agree to three decimals on RMSE at every lead (3.637 / 5.061 / 5.071 at 24 h / 120 h / +240 h), and the `e2s` mode computes it through an entirely independent path — dense predownload via +`WB2ERA5_121x240`, no insitubatch in the loop. That agreement is the correctness check; throughput +alone would not catch a loader that silently aliased or double-lent a buffer. Streaming's peak memory is **flat at ~1.9 GB across N = 120 / 240 / 480**, while the dense grid is -7.53 GB at N = 120 and **OOMs a 15 GB box by ~N = 240**. Dense scales with campaign size; streaming +7.63 GB at N = 120 and **OOMs a 15 GB box by ~N = 240**. Dense scales with campaign size; streaming does not. That bounded-memory property — not just throughput — is the point for a long campaign. (Persistence is a checkpoint-free model that exercises the real `create_iterator` seam on CPU; a @@ -87,14 +107,19 @@ python bench_cache.py --store arco --vars t2m \ | store | field size | cold → warm wall | **cloud fetches (cold → warm)** | |-------|------------|------------------|----------------------------------| -| **WB2** 240×121 | 116 KB | 1.41 s → 1.05 s (~1.4×) | **33 → 0** | -| **ARCO** 721×1440 | 4 MB | 1.22 s → 0.55 s (~2.2×) | **54 → 0** | +| **WB2** 240×121 | 116 KB | 1.16 s → 0.81 s (1.4×) | **33 → 0** | +| **ARCO** 721×1440 | 4 MB | 0.81 s → 0.59 s (1.4×) | **54 → 0** | The deterministic result is **zero cloud fetches on re-score** — the warm run serves every chunk -from local disk. The wall speedup is secondary and scales with how IO-bound the cold fetch is (tiny -WB2 fields ~1.4×; 4 MB ARCO fields ~2.2×); it is *understated* on this box's cheap same-region reads -and grows under metered egress, requester-pays, or cross-region access. The cold wall includes the -one-time persist write, so it runs slightly above the persist-off de-dup figure in §1. +from local disk. The wall speedup is secondary and modest: 1.4× on both stores, despite a 35× +difference in field size, so it is *not* tracking how IO-bound the cold fetch is. On this box's +cheap same-region reads the cloud fetch simply isn't the bottleneck, so removing it entirely buys +little; the wall win grows under metered egress, requester-pays, or cross-region access, while the +fetch-elimination holds everywhere. The cold wall includes the one-time persist write, so it runs +slightly above the persist-off de-dup figure in §1. + +The benchmark wipes and rebuilds `/bench_cold_warm//` so the cold leg starts +empty; it never touches the rest of `--cache-dir`. ## How to read these numbers — framing insitubatch @@ -109,13 +134,14 @@ sharpen its positioning into three evidence-backed claims: 2. **Far ahead when the chunking strategy isn't sample-optimized.** When the access pattern maps many samples onto shared chunks — overlapping windows, verification grids, fat chunks holding several steps — its read planning de-duplicates and a per-sample parallel fetch re-reads. - Evidence: §1 WB2 (174× fewer decodes, 15× wall). + Evidence: §1 WB2 (174× fewer decodes, 12.8× wall). 3. **Honest boundary — you can use it sub-optimally.** It is not a universal speed win. On a chunk-1 store with large fields, against an *unbounded* concurrent gather, its bounded-inflight - scheduling trails per byte (ARCO ~2×; the reads are already minimal — verified — but the dense - output the model consumes must still be assembled, and E2S's flat gather saturates bandwidth on - 4 MB chunks). And a degenerate `batch_size=N` throws away the memory advantage. The tool is - **generally optimal for streaming with bounded memory** — that is the sweet spot. + scheduling trails per byte: on ARCO, E2S moves 2.39 GB in 3.91 s (~610 MB/s) while the feed moves + 0.67 GB in 2.57 s (~260 MB/s) — 2.3× slower per byte, ahead overall only because it reads 3.6× + fewer bytes, netting 1.5×. And a degenerate `batch_size=N` throws away the memory advantage: + §2 `dense` peaks at 7.63 GB, *worse* than the predownload it replaces. The tool is **generally + optimal for streaming with bounded memory** — that is the sweet spot. One line: *stream training/inference batches from cloud tensors in place, with bounded memory — competitive with hand-tuned parallel loaders on optimized layouts, and far ahead when the chunking @@ -123,10 +149,30 @@ causes duplicate reads.* ## Caveats / methodology -- **Single environment, preliminary.** One n2-standard-8-class box (15 GB RAM), cold reads, gcsfs - anon. Numbers to be **cross-posted** after NVIDIA-side runs on the target infrastructure. -- **gcsfs on both sides.** Isolates the loader's contribution from the store backend; obstore would - raise the AFTER throughput further but is not what these numbers measure. +- **Single environment, preliminary.** One n2-standard-8-class box (15 GB RAM), cold reads, + anonymous GCS. Numbers to be **cross-posted** after NVIDIA-side runs on the target infrastructure. +- **obstore on both sides.** Earth2Studio's zarr data sources migrated to obstore in + [#955](https://github.com/NVIDIA/earth2studio/pull/955); the feed uses insitubatch's + `obstore_store` to match, so neither side carries a backend handicap. The de-duplication ratios + are backend-independent — swapping both sides to `fsspec_store(url, token="anon")` changes the + wall clock but not the chunk counts. +- **These numbers supersede an earlier gcsfs measurement.** This recipe was first measured on + 2026-07-04, over gcsfs on both sides and before #955 landed, and reported 15.4× (§1 WB2), ~1.9× + (§1 ARCO), 39.6 s (§2 `e2s`) and ~2.2× (§3 ARCO). Every headline is lower now. Three things + differ between those runs and these — the storage backend on both sides (gcsfs → obstore), + Earth2Studio's per-fetch `logger.debug` output inside the timed region (on → suppressed), and a + month of drift on a shared box. We have **not** isolated their individual contributions, so the + earlier figures are superseded rather than a controlled comparison, and should not be quoted. + One difference *is* established: the §3 ARCO row was reading an unwritten region of the store + (below) — all-NaN fills that never touched the network — so it measured nothing, and its + "~2.2×, scales with field size" result was an artifact. +- **ARCO's time axis begins in 1900, its data in 1940.** The store declares + `hours since 1900-01-01` over 1 323 648 steps to 2050, but chunks outside ~1940–2023 were never + written and read back as NaN fill in ~20 ms without a network request. `--start` therefore + defaults per store (ARCO: `1051896` = 2020-01-01) and both benchmarks now fail loudly if a window + reads back entirely NaN. Earth2Studio's `ARCO` source validates this independently and refuses + pre-1940 requests; insitubatch does not, so a window outside the populated range returns fill + data rather than raising. - **Surface variables only** (`t2m`, `u10m`, `v10m`); pressure-level variables need level indexing, not yet wired in the adapter. - **Persistent cache footprint.** The cache stores *decoded* chunks, so per-chunk bytes exceed the diff --git a/recipes/insitubatch_hindcast/bench_cache.py b/recipes/insitubatch_hindcast/bench_cache.py index cc4b585f7..cdb606c74 100644 --- a/recipes/insitubatch_hindcast/bench_cache.py +++ b/recipes/insitubatch_hindcast/bench_cache.py @@ -28,16 +28,17 @@ never goes stale. This measures the second-run win: same verification window, run COLD (empty cache) then WARM -(cache populated), over gcsfs anon. +(cache populated), over obstore anon. """ import argparse +import os import shutil import time from typing import Any import numpy as np -from insitubatch import fsspec_store +from insitubatch import obstore_store from earth2studio.data.insitu import InSituForecastFeed @@ -45,10 +46,12 @@ "wb2": { "url": "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-240x121_equiangular_with_poles_conservative.zarr", "transpose_inner": True, + "start": 1000, # 1959-09-20; WB2 axis begins 1959 }, "arco": { "url": "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", "transpose_inner": False, + "start": 1051896, # 2020-01-01; ARCO axis begins 1900 but data only from 1940 }, } VAR_MAP = { @@ -59,7 +62,7 @@ def anon_store(url: str) -> Any: - return fsspec_store(url, token="anon", access="read_only") # noqa: S106 + return obstore_store(url, skip_signature=True) def run( @@ -85,11 +88,17 @@ def run( transpose_inner=cfg["transpose_inner"], ) t0 = time.perf_counter() - for _x, _coords in feed: - pass + saw_finite = False + for x, _coords in feed: + saw_finite = saw_finite or bool(x.isfinite().any()) wall = time.perf_counter() - t0 hits, misses = feed.dataset.cache_hits, feed.dataset.cache_misses feed.dataset.close() + if not saw_finite: + raise ValueError( + f"every field read from {cfg['url']} was fill/NaN: the window at " + f"start={start} is outside the store's populated range. Pass a valid --start." + ) return {"wall_s": wall, "hits": hits, "misses": misses} @@ -97,7 +106,7 @@ def main() -> None: p = argparse.ArgumentParser() p.add_argument("--store", choices=list(STORES), default="wb2") p.add_argument("--vars", nargs="+", default=["t2m", "u10m", "v10m"]) - p.add_argument("--start", type=int, default=1000) + p.add_argument("--start", type=int, default=None) # default: per-store (see STORES) p.add_argument("--n-init", type=int, default=48) p.add_argument("--lead-step-h", type=int, default=6) p.add_argument("--max-lead-h", type=int, default=240) @@ -107,26 +116,30 @@ def main() -> None: args = p.parse_args() cfg = STORES[args.store] + start = args.start if args.start is not None else cfg["start"] leads_h = list(range(args.lead_step_h, args.max_lead_h + 1, args.lead_step_h)) requested = args.n_init * len(leads_h) * len(args.vars) - shutil.rmtree(args.cache_dir, ignore_errors=True) # start cold + # Only ever wipe a subdirectory this benchmark owns -- --cache-dir may be a real + # cache root, and the cold leg requires starting empty. + cache_dir = os.path.join(args.cache_dir, "bench_cold_warm", args.store) + shutil.rmtree(cache_dir, ignore_errors=True) # start cold print( - f"[{args.store}] {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " - f"= {requested} requested field-reads ; cache_dir={args.cache_dir}" + f"[{args.store}] start={start} ; {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " + f"= {requested} requested field-reads ; cache_dir={cache_dir}" ) common = ( cfg, args.vars, - args.start, + start, args.n_init, leads_h, args.batch_size, args.max_inflight, ) - cold = run(*common, args.cache_dir) - warm = run(*common, args.cache_dir) + cold = run(*common, cache_dir) + warm = run(*common, cache_dir) print("\n=== COLD (empty cache: fetch + decode + persist) ===") print( diff --git a/recipes/insitubatch_hindcast/bench_hindcast.py b/recipes/insitubatch_hindcast/bench_hindcast.py index a3123dd1a..0fd538c2b 100644 --- a/recipes/insitubatch_hindcast/bench_hindcast.py +++ b/recipes/insitubatch_hindcast/bench_hindcast.py @@ -25,8 +25,11 @@ AFTER = insitubatch InSituForecastFeed over the init window with the leads as shift views; each shared chunk is decoded exactly once (``dataset.cache_misses``). -Both read the SAME store over gcsfs anon, so the delta isolates insitubatch's -dedup + bounded prefetch (not obstore-vs-gcsfs). +Both read the SAME store over obstore anon, so the delta isolates insitubatch's +dedup + bounded prefetch rather than the storage backend. Earth2Studio's zarr sources +moved to obstore in #955; the feed uses ``obstore_store`` to match. Swap both to +``fsspec_store(url, token="anon")`` to re-run the comparison over gcsfs -- the de-dup +ratio is backend-independent, only the wall moves. Two regimes: wb2 = 240x121 6-hourly, chunks=(8,240,121): fat time-chunk -> high dedup ratio, tiny @@ -40,7 +43,7 @@ from typing import Any import numpy as np -from insitubatch import fsspec_store +from insitubatch import obstore_store from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time @@ -52,6 +55,7 @@ "field_bytes": 240 * 121 * 4, "chunk_steps": 8, "transpose_inner": True, + "start": 1000, # 1959-09-20; WB2 axis begins 1959 }, "arco": { "url": "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", @@ -59,6 +63,7 @@ "field_bytes": 721 * 1440 * 4, "chunk_steps": 1, "transpose_inner": False, + "start": 1051896, # 2020-01-01; ARCO axis begins 1900 but data only from 1940 }, } VAR_MAP = { @@ -69,7 +74,7 @@ def anon_store(url: str) -> Any: - return fsspec_store(url, token="anon", access="read_only") # noqa: S106 + return obstore_store(url, skip_signature=True) def load_before_cls(spec: str) -> Any: @@ -101,12 +106,17 @@ def run_after( ) t0 = time.perf_counter() n_rows = 0 + saw_finite = False for x, _coords in feed: - n_rows += x.shape[ - 0 - ] # gather returns eager numpy; no touch needed to force decode + n_rows += x.shape[0] # decode is forced by the gather; no touch needed + saw_finite = saw_finite or bool(x.isfinite().any()) wall = time.perf_counter() - t0 feed.dataset.close() + if not saw_finite: + raise ValueError( + f"every field read from {cfg['url']} was fill/NaN: the window at " + f"start={start} is outside the store's populated range. Pass a valid --start." + ) return { "wall_s": wall, "init_rows": n_rows, @@ -135,7 +145,7 @@ def main() -> None: p = argparse.ArgumentParser() p.add_argument("--store", choices=list(STORES), default="wb2") p.add_argument("--vars", nargs="+", default=["t2m"]) - p.add_argument("--start", type=int, default=1000) + p.add_argument("--start", type=int, default=None) # default: per-store (see STORES) p.add_argument("--n-init", type=int, default=24) p.add_argument("--lead-step-h", type=int, default=6) p.add_argument("--max-lead-h", type=int, default=120) @@ -146,6 +156,7 @@ def main() -> None: args = p.parse_args() cfg = STORES[args.store] + start = args.start if args.start is not None else cfg["start"] leads_h = list(range(args.lead_step_h, args.max_lead_h + 1, args.lead_step_h)) requested = args.n_init * len(leads_h) * len(args.vars) @@ -156,10 +167,10 @@ def main() -> None: times64 = decode_cf_time( np.asarray(g["time"][:]), attrs["units"], attrs.get("calendar", "standard") ) - init_times64 = times64[args.start : args.start + args.n_init] + init_times64 = times64[start : start + args.n_init] print( - f"[{args.store}] grid: {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " + f"[{args.store}] start={start} ({init_times64[0]}) ; grid: {args.n_init} inits x {len(leads_h)} leads x {len(args.vars)} vars " f"= {requested} requested field-reads ({requested*cfg['field_bytes']/1e9:.2f} GB naive)" ) print( @@ -179,7 +190,7 @@ def med3(w: list[float]) -> tuple[float, float, float]: a = run_after( cfg, args.vars, - args.start, + start, args.n_init, leads_h, args.batch_size, @@ -198,7 +209,7 @@ def med3(w: list[float]) -> tuple[float, float, float]: dedup = requested / decodes a_med, a_lo, a_hi = med3(after_walls) - print("\n=== AFTER (insitubatch, gcsfs anon) ===") + print("\n=== AFTER (insitubatch, obstore anon) ===") print(f" wall (med/min/max): {a_med:.2f} / {a_lo:.2f} / {a_hi:.2f} s") print( f" chunk decodes : {decodes} ({decodes*cfg['field_bytes']*cfg['chunk_steps']/1e9:.2f} GB)" @@ -209,7 +220,7 @@ def med3(w: list[float]) -> tuple[float, float, float]: print(f" resident peak : {resident} chunks") if before_walls: b_med, b_lo, b_hi = med3(before_walls) - print("\n=== BEFORE (E2S fetch, gcsfs anon, cache off) ===") + print("\n=== BEFORE (E2S fetch, obstore anon, cache off) ===") print(f" wall (med/min/max): {b_med:.2f} / {b_lo:.2f} / {b_hi:.2f} s") print("\n=== HEADLINE (medians) ===") print( diff --git a/recipes/insitubatch_hindcast/stream_score.py b/recipes/insitubatch_hindcast/stream_score.py index cbda0c227..30e923adb 100644 --- a/recipes/insitubatch_hindcast/stream_score.py +++ b/recipes/insitubatch_hindcast/stream_score.py @@ -41,7 +41,7 @@ import numpy as np import torch -from insitubatch import fsspec_store +from insitubatch import obstore_store from earth2studio.data.insitu import InSituForecastFeed, decode_cf_time from earth2studio.models.px import Persistence @@ -57,7 +57,7 @@ def anon_store() -> Any: - return fsspec_store(URL, token="anon", access="read_only") # noqa: S106 + return obstore_store(URL, skip_signature=True) def peak_rss_gb() -> float: From d1df9df1bdcc6ee5dc7119048f88fec0c1036500 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:51:15 +0000 Subject: [PATCH 14/16] State the logging requirement plainly in the recipe README Earth2Studio's per-fetch debug logging sits inside the timed region and is not free -- say so, say the published numbers are all measured with it off, and drop the inconclusive breakdown of what changed since the first measurement. Co-Authored-By: Claude Opus 5 --- recipes/insitubatch_hindcast/README.md | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index ec352f92c..16fc89dae 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -22,9 +22,9 @@ Both stores are anonymous public GCS buckets (WeatherBench2 ERA5, ARCO ERA5); no needed. Every measurement below reads over **obstore anon on both the before and after side**, so the delta isolates insitubatch's read-planning + streaming rather than the storage backend. -Run the benchmarks with Earth2Studio's per-fetch debug logging suppressed, so the timed region is -the same on both sides — it emits one line per `(time, variable)`, which is 5760 lines in §1 and -14 760 in §2, all on the baseline leg: +Run the benchmarks with Earth2Studio's per-fetch debug logging suppressed. It emits one line per +`(time, variable)` — 5760 lines in §1, 14 760 in §2, essentially all on the baseline leg — from +inside the timed region, and it is not free. **Every number below is measured with it off:** ```bash export LOGURU_LEVEL=INFO @@ -158,14 +158,12 @@ causes duplicate reads.* wall clock but not the chunk counts. - **These numbers supersede an earlier gcsfs measurement.** This recipe was first measured on 2026-07-04, over gcsfs on both sides and before #955 landed, and reported 15.4× (§1 WB2), ~1.9× - (§1 ARCO), 39.6 s (§2 `e2s`) and ~2.2× (§3 ARCO). Every headline is lower now. Three things - differ between those runs and these — the storage backend on both sides (gcsfs → obstore), - Earth2Studio's per-fetch `logger.debug` output inside the timed region (on → suppressed), and a - month of drift on a shared box. We have **not** isolated their individual contributions, so the - earlier figures are superseded rather than a controlled comparison, and should not be quoted. - One difference *is* established: the §3 ARCO row was reading an unwritten region of the store - (below) — all-NaN fills that never touched the network — so it measured nothing, and its - "~2.2×, scales with field size" result was an artifact. + (§1 ARCO), 39.6 s (§2 `e2s`) and ~2.2× (§3 ARCO). Every headline is lower now. That run predates + #955, read over gcsfs on both sides, and was measured with the debug logging above still enabled + — so it is superseded rather than a controlled comparison, and should not be quoted. One + difference is established independently of all that: the §3 ARCO row was reading an unwritten + region of the store (below) — all-NaN fills that never touched the network — so it measured + nothing, and its "~2.2×, scales with field size" result was an artifact. - **ARCO's time axis begins in 1900, its data in 1940.** The store declares `hours since 1900-01-01` over 1 323 648 steps to 2050, but chunks outside ~1940–2023 were never written and read back as NaN fill in ~20 ms without a network request. `--start` therefore From 97cc393d00c02f54cdbd3a1ea2188ff512555572 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Tue, 4 Aug 2026 04:21:37 +0000 Subject: [PATCH 15/16] Attribute the ARCO gap to assembly, not throughput The "2.3x slower per byte" claim was an artifact: it divided Earth2Studio's *notional* byte count (576 requested reads x 4.15 MB) by wall time, crediting it with 3.6x redundancy it never uniquely moved. Measured with redundancy removed from both sides -- one init at unit lead spacing, so 162 requested = 162 unique = 0.67 GB either way -- E2S takes 1.38 s and the feed 1.61 s. The feed is ~17% slower per byte, roughly at parity, not 2.3x behind. The gap in the headline ARCO row comes from somewhere else: de-duplication removes the fetch and decode of a redundant sample but not its assembly, since the tensor the model consumes still has one slot per requested (init, lead). Subtracting the two configurations, each redundant sample costs E2S 6.1 ms and the feed 2.3 ms -- 2.6x cheaper, not 3.6x, which is why 3.6x fewer decodes nets only 1.5x wall. Document the zero-redundancy control so the split is reproducible. Co-Authored-By: Claude Opus 5 --- recipes/insitubatch_hindcast/README.md | 33 ++++++++++++++++++++------ 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index 16fc89dae..605661a78 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -58,6 +58,20 @@ small — insitubatch dominates. ARCO is the **honest** case (see caveats). ARCO is wide on this box (±15% on both legs); its row is the median of 10 repeats, and the BEFORE leg converges downward with sample size as HTTP connections warm. +**Separating read elimination from read throughput.** The speedups above conflate two things: +reading *less*, and reading *fast*. To isolate the second, run the same store with no redundancy at +all — one init at unit lead spacing, so every requested read is unique: + +```bash +python bench_hindcast.py --store arco --vars t2m --lead-step-h 1 \ + --max-lead-h 162 --n-init 1 --repeats 8 +``` + +162 requested = 162 unique = 0.67 GB moved on both sides. E2S: **1.38 s** (486 MB/s); the feed: +**1.61 s** (416 MB/s) — the feed is **~17% slower per byte**, roughly at parity (medians over three +independent 8-repeat runs). insitubatch's modest ARCO result is therefore *not* a throughput +deficit; see the honest boundary below for where the gap actually comes from. + ## 2. `stream_score.py` — streaming vs dense materialization The model's `create_iterator` already streams the forecast lead-by-lead, and scoring is pointwise @@ -135,13 +149,18 @@ sharpen its positioning into three evidence-backed claims: many samples onto shared chunks — overlapping windows, verification grids, fat chunks holding several steps — its read planning de-duplicates and a per-sample parallel fetch re-reads. Evidence: §1 WB2 (174× fewer decodes, 12.8× wall). -3. **Honest boundary — you can use it sub-optimally.** It is not a universal speed win. On a - chunk-1 store with large fields, against an *unbounded* concurrent gather, its bounded-inflight - scheduling trails per byte: on ARCO, E2S moves 2.39 GB in 3.91 s (~610 MB/s) while the feed moves - 0.67 GB in 2.57 s (~260 MB/s) — 2.3× slower per byte, ahead overall only because it reads 3.6× - fewer bytes, netting 1.5×. And a degenerate `batch_size=N` throws away the memory advantage: - §2 `dense` peaks at 7.63 GB, *worse* than the predownload it replaces. The tool is **generally - optimal for streaming with bounded memory** — that is the sweet spot. +3. **Honest boundary — you can use it sub-optimally.** It is not a universal speed win, and on a + chunk-1 store with large fields the reason is *not* raw throughput: with redundancy removed from + both sides the feed is only ~17% slower per byte than E2S's unbounded gather (§1). The gap opens + because **de-duplication removes the fetch and decode of a redundant sample, but not its + assembly** — the tensor the model consumes still has one slot per requested `(init, lead)`. + Subtracting the two §1 ARCO configurations, each of the 414 redundant samples costs E2S 6.1 ms + (it refetches) and the feed 2.3 ms (it re-assembles from an already-resident chunk): 2.6× cheaper, + not 3.6×, which is why 3.6× fewer decodes nets only 1.5× wall. Where fields are small (§1 WB2) + assembly is negligible and nearly the whole de-dup ratio converts. And a degenerate + `batch_size=N` throws away the memory advantage: §2 `dense` peaks at 7.63 GB, *worse* than the + predownload it replaces. The tool is **generally optimal for streaming with bounded memory** — + that is the sweet spot. One line: *stream training/inference batches from cloud tensors in place, with bounded memory — competitive with hand-tuned parallel loaders on optimized layouts, and far ahead when the chunking From ebdc69531b41b7eca3feddb69b6a34bfbf2b95f0 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Wed, 5 Aug 2026 01:38:45 +0000 Subject: [PATCH 16/16] =?UTF-8?q?Correct=20the=20predownload=20framing=20a?= =?UTF-8?q?nd=20re-measure=20=C2=A71=20against=20a=20cached=20baseline?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two corrections, both from reading `recipes/eval` and Earth2Studio's own defaults rather than assuming them. Predownload already de-duplicates: `compute_verification_times` collapses the (init, lead) grid onto unique valid times before any fetch. It is also a deliberate cluster-scale ETL -- rank-parallel via `distribute_work`, resumable, leaving a durable pre-regridded artifact -- not a workaround. State the scope honestly: insitubatch replaces the predownload-then-read cycle for streaming consumption on one box, not the bulk ETL phase, which a single async event loop cannot scale across nodes. The §1 baseline also ran with `cache=False`, which is not the source default. `LocalCachingStore` caches compressed buffers keyed by chunk path, so a stock run serves redundant reads from local disk. Add `--before-cache` to run the baseline as configured by default, wiping a bench-owned cache directory before each repeat so every repeat is a cold pass, and report both configurations. WB2 6.8x cached / 9.5x uncached; ARCO 1.4x / 1.7x. Decode counts are identical either way -- the cache is below zarr's codec pipeline. Also record what the de-dup ratios are measured against. 174x and 3.6x are against the live per-work-item path; against a valid-time-deduplicated baseline the advantage is exactly the sample-axis steps per chunk -- 7.9x on WB2, none at all on chunk-1 ARCO. Relabel the §2 `e2s` row, which measures this harness's own dense buffer rather than predownload, and flag it pending a rerun. Co-Authored-By: Claude Opus 5 --- recipes/insitubatch_hindcast/README.md | 169 ++++++++++++++---- .../insitubatch_hindcast/bench_hindcast.py | 73 +++++++- 2 files changed, 206 insertions(+), 36 deletions(-) diff --git a/recipes/insitubatch_hindcast/README.md b/recipes/insitubatch_hindcast/README.md index 605661a78..c6c95195f 100644 --- a/recipes/insitubatch_hindcast/README.md +++ b/recipes/insitubatch_hindcast/README.md @@ -5,10 +5,35 @@ Two runnable benchmarks that feed ERA5 into an Earth2Studio prognostic **without (`earth2studio.data.insitu.InSituForecastFeed`) instead. They quantify what a streaming, read-planning loader changes for an IO-bound hindcast / scoring campaign. -The motivation is `recipes/eval`: its `predownload.py` sentinel exists because live `fetch_data` -is too slow for a scoring campaign. That predownload materializes the whole `(init, lead)` -verification grid up front. Both are consequences of a per-`(time, variable)` fetch with no read -de-duplication — exactly what insitubatch removes. +The motivation is `recipes/eval`, which requires a `predownload.py` pass before `main.py`. + +That predownload is a **deliberate cluster-scale ETL, not a workaround**, and it already +de-duplicates. `compute_verification_times` (`src/predownload_utils.py`) collapses the +`(init, lead)` grid onto the set of unique valid times; `predownload.py` then fetches those +partitioned across ranks by `distribute_work` (documented invocation: +`torchrun --nproc_per_node=8`), one timestamp at a time within a rank, each written and flushed to +zarr with a resume marker before the next. Rank-parallel, resumable, and it leaves a durable store +many checkpoints can be scored against — pre-regridded onto the model grid in the StormScope case. +Separating the phases is itself the point on a GPU cluster: bulk IO runs on cheap CPU nodes and +expensive GPU time stays off the network. + +**insitubatch does not replace that, and structurally cannot** — its parallelism lives in one async +event loop rather than worker processes, so it does not scale a bulk fetch across nodes. What it +replaces is the predownload-then-read *cycle* for streaming consumption: one box, no separate phase, +no materialized copy of the verification set. Rank-parallel training and inference are unaffected +(each DDP rank streams its own shard); the bulk ETL phase is the part insitubatch has no answer for. + +Within that scope the gap is narrower than "no de-duplication", and lives in two places: + +1. **The live path has none.** Every pipeline call site fetches per work item + (`fetch_data(time=[item.time], ...)` in `src/pipelines/forecast.py`, `dlesym.py`, + `assimilation.py`) with no memory of what a neighbouring init already read. §1 measures this. +2. **Timestamp granularity is coarser than chunk granularity.** Even a perfect valid-time de-dup + issues one read per unique time; a fat time-chunk holding 8 steps serves 8 of them from one + decode. §1 WB2 measures that residual. + +§3 covers the remaining difference: predownload leaves a materialized copy of the verification set +on disk, where the persistent cache holds only the chunks actually touched. ## Setup @@ -41,22 +66,70 @@ times share valid times, and a fat time-chunk holds several steps, so the reques onto far fewer stored chunks. BEFORE = E2S's per-init `fetch_data`; AFTER = the insitubatch feed (each lead a sample-axis `shift` view; each shared chunk decoded once). +The BEFORE leg runs in two configurations, because Earth2Studio's data sources cache by default +(`cache=True`) and that materially changes the wall: + ```bash python bench_hindcast.py --store wb2 --vars t2m u10m v10m \ - --n-init 48 --max-lead-h 240 --repeats 5 + --n-init 48 --max-lead-h 240 --repeats 5 --before-cache python bench_hindcast.py --store arco --vars t2m --lead-step-h 6 \ - --n-init 24 --max-lead-h 144 --repeats 10 + --n-init 24 --max-lead-h 144 --repeats 10 --before-cache ``` -| store | layout | requested reads | unique decodes | **wall speedup** | -|-------|--------|-----------------|----------------|------------------| -| **WB2** 240×121 6-h | `chunks=(8,240,121)` fat | 5760 | **33** (174×) | **12.8×** (10.84→0.84 s) | -| **ARCO** 721×1440 1-h | `chunks=(1,721,1440)` chunk-1 | 576 | 162 (3.6×) | 1.5× (3.91→2.57 s) | - -WB2's fat time-chunk amortizes 8 steps per read, so the de-dup ratio is large and the fields are +Drop `--before-cache` for the uncached leg. `LocalCachingStore` sits at the `Store.get` level and +holds *compressed* buffers keyed by chunk path, so with the cache on a redundant read costs a local +disk hit instead of a network round-trip — but zarr still decodes the chunk again. Decode counts are +identical in both configurations; only the fetch component of the wall moves. + +| store | chunks | requested → decodes | cache **on** | cache off | +|-------|--------|---------------------|--------------|-----------| +| **WB2** 240×121 6-h | `(8,240,121)` fat | 5760 → **33** (174×) | **6.8×** | 9.4× | +| **ARCO** 721×1440 1-h | `(1,721,1440)` chunk-1 | 576 → 162 (3.6×) | **1.2×** | 1.7× | + +Medians: WB2 8.22→1.21 s cached (5 repeats), 10.95→1.17 s uncached (10 repeats); ARCO +3.78→3.16 s cached, 5.20→3.11 s uncached (10 repeats each). + +**Put the baseline's cache on your fastest local disk.** It is written cold on every repeat here, so +the device shows up in the wall. ARCO's cache is 363 MB (162 chunks) and moving it from the boot +disk to local NVMe took the cached baseline 4.31→3.78 s, i.e. 1.4×→**1.2×** — quoted above is the +NVMe figure, the one favourable to the baseline. WB2's cache is ~10 MB and does not move. Point +`TMPDIR` at the fast device: `TMPDIR=/mnt/nvme python bench_hindcast.py … --before-cache`. + +At 1.2× on ARCO the two distributions overlap (feed 2.69–3.98 s against a tight 3.71–4.02 s). Read +that row as the boundary case it is, not as a win. + +Quote the cache-on column: it is how a stock Earth2Studio run behaves. The cache recovers only ~25% +of the WB2 baseline wall (10.89→8.22 s) because WB2 chunks are ~116 KB — the network was never the +bottleneck there. The cost is 5760 decodes against 33, and no byte cache addresses decode. + +**These de-dup ratios are against the *live* path, not against `predownload.py`.** The BEFORE leg +re-requests every `(init, lead)` pair, which is what the pipelines do (`fetch_data(time=[item.time], +…)` per work item, no memory across items) but *not* what predownload does — it de-duplicates valid +times first. Against a valid-time-deduplicated baseline the advantage is smaller, and it is exactly +the number of sample-axis steps per chunk: + +| store | unique valid times × vars | insitubatch decodes | advantage | +|-------|--------------------------|---------------------|-----------| +| **WB2** | 87 × 3 = 261 | 33 | **7.9×** (= 8 steps/chunk) | +| **ARCO** | 162 × 1 = 162 | 162 | **1.0× — none** | + +That is arithmetic, not a measurement: WB2's 48 consecutive 6-h inits with leads +1…+40 span valid +indices 1001–1087 (87 times, chunks 125–135 = 11 × 3 vars = the 33 decoded); ARCO's 24 consecutive +1-h inits with leads +6…+144 cover every index in [6, 167] (162 times = the 162 decoded). On a +chunk-1 store, chunk granularity *is* timestamp granularity, so against a de-duplicated baseline +insitubatch decodes nothing fewer. The wall for that baseline is not measured here. + +WB2's fat time-chunk amortizes 8 steps per decode, so the de-dup ratio is large and the fields are small — insitubatch dominates. ARCO is the **honest** case (see caveats). ARCO's per-repeat spread -is wide on this box (±15% on both legs); its row is the median of 10 repeats, and the BEFORE leg -converges downward with sample size as HTTP connections warm. +is wide on this box (±15% on both legs); its row is the median of 10 repeats. + +These walls were measured on 2026-08-05 and supersede the figures this recipe carried previously +(12.8× WB2, 1.5× ARCO), which were cache-off only. One difference is **not** the cache: the +insitubatch leg no longer reproduces its earlier WB2 wall — 0.84 s then, 1.17 s now (median of 10 +repeats, range 1.02–1.24) against an unchanged 10.95 s baseline. The current figure reproduces +across independent runs and repeat counts and the earlier one does not, so it is the one quoted; +the cause of the shift is unexplained. Every comparison above has both legs measured in the same +session, so the ratios hold regardless. **Separating read elimination from read throughput.** The speedups above conflate two things: reading *less*, and reading *fast*. To isolate the second, run the same store with no redundancy at @@ -87,14 +160,22 @@ done | mode | wall | **peak RSS** | field reads | |------|------|--------------|-------------| -| `e2s` — dense predownload (redundant reads) | 29.0 s | 3.10 GB | 14 760 | +| `e2s` — live per-init `fetch_data`, dense buffer | 29.0 s | 3.10 GB | 14 760 | | `dense` — insitubatch, `batch_size=N` | 4.3 s | 7.63 GB | 60 | | `stream` — insitubatch, `batch_size=W` | 2.9 s | **1.85 GB** | 60 | All three agree to three decimals on RMSE at every lead (3.637 / 5.061 / 5.071 at 24 h / 120 h / -240 h), and the `e2s` mode computes it through an entirely independent path — dense predownload via -`WB2ERA5_121x240`, no insitubatch in the loop. That agreement is the correctness check; throughput -alone would not catch a loader that silently aliased or double-lent a buffer. +240 h), and the `e2s` mode computes it through an entirely independent path — live per-init +`fetch_data` via `WB2ERA5_121x240`, no insitubatch in the loop. That agreement is the correctness +check; throughput alone would not catch a loader that silently aliased or double-lent a buffer. + +> **The `e2s` leg is not `recipes/eval`'s predownload, and its 14 760 field reads overstate the +> status quo.** It fetches per init with all leads and accumulates into a dense scoring buffer — +> that dense buffer is this harness's construction, not Earth2Studio's. A real predownload would +> first collapse these 120 inits × 41 leads onto the ~160 unique valid times they span. The +> streaming-vs-dense memory result (the point of this section) is unaffected, since it compares the +> two insitubatch modes; the `e2s` wall is not a fair status-quo baseline and is pending a rerun +> against a valid-time-deduplicated fetch. Streaming's peak memory is **flat at ~1.9 GB across N = 120 / 240 / 480**, while the dense grid is 7.63 GB at N = 120 and **OOMs a 15 GB box by ~N = 240**. Dense scales with campaign size; streaming @@ -105,12 +186,18 @@ real NVIDIA checkpoint — SFNO/FCN — is a drop-in with the same code on a GPU ## 3. `bench_cache.py` — cross-run persistent cache -The intro's `predownload.py` exists so a re-scored campaign doesn't re-fetch the same ground -truth. `InSituForecastFeed(cache_dir=...)` gives that for free: the first run decodes each shared -chunk once **and** persists it to local disk; a later run over the same store reads those chunks -back as cache hits, touching the cloud zero times. No predownload step, no reshard, only the chunks -actually touched — and because a reanalysis store is static, the cache never goes stale. This is the -common eval shape: many models (or checkpoints) scored against one fixed verification set. +A re-scored campaign shouldn't re-fetch the same ground truth. On a cluster that is exactly what +`predownload.py` is for, and this is not an argument against it. `InSituForecastFeed(cache_dir=...)` +covers the same need without a separate phase: the first run decodes each shared chunk once **and** +persists it to local disk; a later run over the same store reads those chunks back as cache hits, +touching the cloud zero times. Only the chunks actually touched, no materialized copy of the grid — +and because a reanalysis store is static, the cache never goes stale. + +The trade is scope. Predownload buys rank-parallel bulk fetch, resumability, and a durable +pre-regridded artifact; the cache buys the same re-score property with no ETL phase to schedule and +no full copy to provision. Which one fits depends on whether you have a cluster to run the phase on +— for the common eval shape of many checkpoints against one fixed verification set on a single box, +the cache is the cheaper path. ```bash python bench_cache.py --store wb2 --vars t2m u10m v10m \ @@ -144,23 +231,36 @@ sharpen its positioning into three evidence-backed claims: 1. **Competitive with an optimized parallel loader, at lower memory.** On a well-chunked store and for streaming consumption it matches a hand-tuned concurrent fetch's throughput while holding - *bounded* memory (streaming: flat ~1.9 GB where dense predownload OOMs). Evidence: §2. + *bounded* memory (streaming: flat ~1.9 GB where a dense verification grid OOMs). Evidence: §2. 2. **Far ahead when the chunking strategy isn't sample-optimized.** When the access pattern maps many samples onto shared chunks — overlapping windows, verification grids, fat chunks holding several steps — its read planning de-duplicates and a per-sample parallel fetch re-reads. - Evidence: §1 WB2 (174× fewer decodes, 12.8× wall). + Evidence: §1 WB2 (174× fewer decodes, 6.8× wall against a default-configured source), measured + against the **live** `fetch_data` path, which de-duplicates nothing across work items. Against + `recipes/eval`'s offline valid-time de-dup the advantage narrows to exactly the steps per chunk + — 7.9× on WB2, and **nothing at all on a chunk-1 store like ARCO** — though it needs no offline + pass to get it. 3. **Honest boundary — you can use it sub-optimally.** It is not a universal speed win, and on a chunk-1 store with large fields the reason is *not* raw throughput: with redundancy removed from both sides the feed is only ~17% slower per byte than E2S's unbounded gather (§1). The gap opens because **de-duplication removes the fetch and decode of a redundant sample, but not its assembly** — the tensor the model consumes still has one slot per requested `(init, lead)`. - Subtracting the two §1 ARCO configurations, each of the 414 redundant samples costs E2S 6.1 ms - (it refetches) and the feed 2.3 ms (it re-assembles from an already-resident chunk): 2.6× cheaper, - not 3.6×, which is why 3.6× fewer decodes nets only 1.5× wall. Where fields are small (§1 WB2) - assembly is negligible and nearly the whole de-dup ratio converts. And a degenerate + Subtracting the two §1 ARCO configurations, each redundant sample costs E2S several ms (it + re-fetches, or re-reads from its local cache) against the feed's re-assembly from an + already-resident chunk — cheaper, but not by the full de-dup ratio, which is why 3.6× fewer + decodes nets only 1.2× wall against a default-configured baseline. (That per-sample split was + quantified against the cache-off + baseline and needs re-deriving for the cached one; the direction holds, the coefficients do + not.) Where fields are small (§1 WB2) assembly is negligible and much more of the de-dup ratio + converts. And a degenerate `batch_size=N` throws away the memory advantage: §2 `dense` peaks at 7.63 GB, *worse* than the - predownload it replaces. The tool is **generally optimal for streaming with bounded memory** — - that is the sweet spot. + dense-buffer baseline it replaces (3.10 GB). The tool is **generally optimal for streaming with + bounded memory** — that is the sweet spot. + +A fourth boundary is scope rather than misuse: **insitubatch does not replace a rank-parallel bulk +ETL.** `predownload.py` scales a fetch across nodes, resumes after a failure, and leaves a durable +pre-regridded artifact; one async event loop does none of those. Everything measured here compares +streaming consumption, not the ETL phase. One line: *stream training/inference batches from cloud tensors in place, with bounded memory — competitive with hand-tuned parallel loaders on optimized layouts, and far ahead when the chunking @@ -170,6 +270,11 @@ causes duplicate reads.* - **Single environment, preliminary.** One n2-standard-8-class box (15 GB RAM), cold reads, anonymous GCS. Numbers to be **cross-posted** after NVIDIA-side runs on the target infrastructure. +- **§2 and §3 were measured with the baseline's cache off.** Only §1 has been re-run in both + configurations. §2's `e2s` leg additionally lacks valid-time de-duplication (see the note there), + so its wall overstates the status quo on two counts, not one. Both are pending a re-run. +- **The zero-redundancy control below is unaffected by the cache setting** — with every requested + read unique there is nothing for a byte cache to serve. - **obstore on both sides.** Earth2Studio's zarr data sources migrated to obstore in [#955](https://github.com/NVIDIA/earth2studio/pull/955); the feed uses insitubatch's `obstore_store` to match, so neither side carries a backend handicap. The de-duplication ratios diff --git a/recipes/insitubatch_hindcast/bench_hindcast.py b/recipes/insitubatch_hindcast/bench_hindcast.py index 0fd538c2b..2ca00fa8c 100644 --- a/recipes/insitubatch_hindcast/bench_hindcast.py +++ b/recipes/insitubatch_hindcast/bench_hindcast.py @@ -25,6 +25,20 @@ AFTER = insitubatch InSituForecastFeed over the init window with the leads as shift views; each shared chunk is decoded exactly once (``dataset.cache_misses``). +The BEFORE leg has two configurations, and they measure different things: + + --before-cache OFF (default here) Earth2Studio's per-source cache disabled, so every + redundant read goes back to the cloud. Symmetric with the feed, which holds no local + byte cache either -- but NOT how a stock Earth2Studio run behaves. + --before-cache ON the source's own default (``cache=True``), which wraps + the store in ``LocalCachingStore``. Redundant reads then hit local disk instead of the + network. That cache sits at the ``Store.get`` level and holds *compressed* buffers, so + zarr still decodes a fat chunk once per requested step either way -- the de-dup ratio + is unchanged, only the fetch component of the wall moves. + +Report both. The cached run is the honest status quo for wall-clock; the uncached run +isolates read elimination from local-cache effects. + Both read the SAME store over obstore anon, so the delta isolates insitubatch's dedup + bounded prefetch rather than the storage backend. Earth2Studio's zarr sources moved to obstore in #955; the feed uses ``obstore_store`` to match. Swap both to @@ -39,6 +53,9 @@ """ import argparse +import os +import shutil +import tempfile import time from typing import Any @@ -125,10 +142,35 @@ def run_after( } +BEFORE_CACHE_DIRNAME = "insitubatch_bench_e2s_cache" + + +def before_cache_root() -> str: + """Bench-owned Earth2Studio cache root for the ``--before-cache`` leg. + + Deliberately not the user's real cache (``~/.cache/earth2studio`` or + ``$EARTH2STUDIO_CACHE``): this directory is wiped before every repeat so + each repeat measures a cold-cache campaign pass. + """ + return os.path.join(tempfile.gettempdir(), BEFORE_CACHE_DIRNAME) + + +def wipe_before_cache() -> None: + """Empty the bench-owned cache so the next repeat starts cold.""" + root = before_cache_root() + if os.path.basename(root) != BEFORE_CACHE_DIRNAME: + raise RuntimeError(f"refusing to wipe unexpected cache path: {root}") + shutil.rmtree(root, ignore_errors=True) + + def run_before( - before_cls: Any, variables: list[str], init_times64: Any, leads_h: list[int] + before_cls: Any, + variables: list[str], + init_times64: Any, + leads_h: list[int], + cache: bool, ) -> dict[str, Any]: - src = before_cls(cache=False, verbose=False) + src = before_cls(cache=cache, verbose=False) init_dt = init_times64.astype("datetime64[s]").astype("O") leads_td = [np.timedelta64(h, "h") for h in leads_h] t0 = time.perf_counter() @@ -153,8 +195,22 @@ def main() -> None: p.add_argument("--max-inflight", type=int, default=32) p.add_argument("--repeats", type=int, default=1) p.add_argument("--skip-before", action="store_true") + p.add_argument( + "--before-cache", + action="store_true", + help="run the BEFORE leg with Earth2Studio's per-source cache ON (its own " + "default), so redundant reads hit local disk instead of the cloud. Decode " + "still repeats per requested step either way. Uses a bench-owned cache " + "directory, wiped before each repeat so every repeat is a cold-cache pass.", + ) args = p.parse_args() + if args.before_cache: + # Redirect Earth2Studio's cache to a bench-owned directory before any source + # is constructed. DATA_CACHE takes precedence downstream, so set both. + os.environ["EARTH2STUDIO_CACHE"] = before_cache_root() + os.environ["EARTH2STUDIO_DATA_CACHE"] = before_cache_root() + cfg = STORES[args.store] start = args.start if args.start is not None else cfg["start"] leads_h = list(range(args.lead_step_h, args.max_lead_h + 1, args.lead_step_h)) @@ -199,8 +255,12 @@ def med3(w: list[float]) -> tuple[float, float, float]: after_walls.append(a["wall_s"]) decodes, resident = a["chunk_decodes"], a["resident_peak"] if not args.skip_before: + if args.before_cache: + wipe_before_cache() # every repeat is a cold-cache pass before_walls.append( - run_before(before_cls, args.vars, init_times64, leads_h)["wall_s"] + run_before( + before_cls, args.vars, init_times64, leads_h, args.before_cache + )["wall_s"] ) print( f" repeat {r+1}/{args.repeats}: after={after_walls[-1]:.2f}s" @@ -220,7 +280,12 @@ def med3(w: list[float]) -> tuple[float, float, float]: print(f" resident peak : {resident} chunks") if before_walls: b_med, b_lo, b_hi = med3(before_walls) - print("\n=== BEFORE (E2S fetch, obstore anon, cache off) ===") + cache_note = ( + "cache ON (E2S default; cold per repeat, redundant reads hit local disk)" + if args.before_cache + else "cache OFF (not the E2S default; every redundant read goes to the cloud)" + ) + print(f"\n=== BEFORE (E2S fetch, obstore anon, {cache_note}) ===") print(f" wall (med/min/max): {b_med:.2f} / {b_lo:.2f} / {b_hi:.2f} s") print("\n=== HEADLINE (medians) ===") print(