Add insitubatch streaming ERA5 data source and hindcast recipe

[emfdavid]

Earth2Studio Pull Request

Closes #961

Description

A recipe for streaming ERA5 into an Earth2Studio prognostic with
insitubatch — a streaming, read-planning cloud-zarr
loader — instead of the dense fetch_data grid, for IO-bound hindcast / scoring campaigns.

How it plugs in — around the model, not inside the xarray loop. insitubatch delivers
tensor batches and never builds xr.DataArray. The new
earth2studio.data.insitu.InSituForecastFeed reads the analysis store → DLPack →
(torch.Tensor, CoordSystem) and feeds prognostic.create_iterator(x, coords) directly, where
each forecast lead is a sample-axis shift view of one array (decode-once, no reshard). It
does not touch the DataSource → xr.DataArray → fetch_data path, so E2S's lexicon / coords /
regrid machinery is untouched — this is a complementary, opt-in feed, not a replacement.

Why it helps an IO-bound campaign. A scoring grid needs ERA5 at valid = init + lead for
every (init, lead); consecutive inits share valid times and a fat time-chunk holds several
steps, so the requested reads collapse onto far fewer stored chunks. insitubatch's read planning
de-duplicates those into one decode per stored chunk under a bounded max_inflight budget,
and streams lead-by-lead so scoring never materializes a dense verification tensor.

What's in the PR

• earth2studio/data/insitu.py — the InSituForecastFeed adapter (guarded optional import of
  insitubatch; yields (torch.Tensor, CoordSystem)).
• recipes/insitubatch_hindcast/ — three runnable benchmarks + a README:
  • bench_hindcast.py — verification-read de-duplication vs per-init fetch_data.
  • stream_score.py — streaming vs dense materialization (three modes, identical RMSE).
  • bench_cache.py — cross-run persistent cache (cold → warm).

Benchmark results

Preliminary, one n2-standard-8-class box (15 GB RAM), cold reads, obstore anon on both the
before and after side, so neither side carries a backend handicap. Public WeatherBench2 and ARCO
ERA5 stores. Run these with LOGURU_LEVEL=INFO — Earth2Studio's per-fetch debug logging sits
inside the timed region (14 760 lines on the e2s leg below) and is not free. Every number here is
measured with it off.

These numbers were re-measured on 2026-08-04 and supersede the ones this PR was opened with.
Earth2Studio's zarr data sources migrated to obstore in #955 while this PR was in review, which
silently turned the comparison into obstore (before) vs gcsfs (after). The feed now uses
insitubatch's obstore_store to match. The earlier figures (15.4× / ~1.9× / 39.6 s / ~2.2×) were
taken over gcsfs on both sides, before #955, and with the debug logging still on; they are
superseded rather than a controlled backend comparison, and shouldn't be quoted. The
de-duplication ratios did not change — identical chunk counts across the backend swap, which is
the substantive point: the win is read planning, not transport.

1. Verification-read de-duplication (bench_hindcast.py) — BEFORE = per-init fetch_data,
AFTER = the insitubatch feed:

Earth2Studio's data sources cache by default (cache=True), which materially changes the wall, so
the baseline runs in both configurations. LocalCachingStore holds compressed buffers keyed by
chunk path, below zarr's codec pipeline — so with the cache on a redundant read costs a local disk
hit instead of a round-trip, but the chunk is still decoded again. Decode counts are identical
either way.

store	chunks	requested → decodes	cache on (default)	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). The baseline's cache is written cold
on every repeat, so its device shows in the wall: ARCO's cache is 363 MB and moving it from the boot
disk to local NVMe took the cached leg 4.31 → 3.78 s, i.e. 1.4× → 1.2×. The NVMe figure is quoted —
the one favourable to the baseline. WB2's cache is ~10 MB and does not move. At 1.2× the ARCO
distributions overlap (feed 2.69–3.98 s against 3.71–4.02 s); that row is the boundary case, not a
win. Quote the cache-on column — it is how a stock run behaves. The
cache recovers only ~25% of the WB2 baseline wall, because WB2 chunks are ~116 KB and the network
was never the bottleneck there; the cost is 5760 decodes against 33, which no byte cache addresses.

These ratios are against the live path, not against predownload.py. The baseline
re-requests every (init, lead) pair — what the pipelines do, not what predownload does, since it
de-duplicates valid times first. Against a valid-time-deduplicated baseline the advantage is
exactly the sample-axis steps per chunk: WB2 261 decodes → 33 = 7.9×; ARCO 162 → 162 =
none at all, because on a chunk-1 store chunk granularity is timestamp granularity. That is
arithmetic from the index ranges, not a measurement — the wall for that baseline is unmeasured.

These supersede the figures this PR previously carried (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. Both legs of every
comparison were measured in the same session, so the ratios hold regardless.

2. Streaming vs dense materialization (stream_score.py, N = 120 inits × 40 leads) — all
three modes produce identical RMSE (3.637 / 5.061 / 5.071 at 24 h / 120 h / 240 h):

mode	wall	peak RSS	field reads
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

The e2s mode computes that RMSE through an entirely independent path — live per-init
fetch_data via WB2ERA5_121x240, no insitubatch in the loop — so the agreement is a real
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 into a dense scoring buffer, and that buffer is
this harness's construction, not Earth2Studio's. compute_verification_times collapses the
(init, lead) grid onto unique valid times before any fetch — ~160 of them for this config, not
4920 pairs. The streaming-vs-dense memory result is unaffected (it compares two insitubatch
modes), but the e2s wall is not a fair status-quo baseline and is pending a rerun against a
valid-time-deduplicated fetch.

Streaming 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 does not — that bounded-memory property, not just throughput, is the point for a long
campaign. (Persistence is used as 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.)

Cross-run persistent cache

InSituForecastFeed(cache_dir=...) persists decoded chunks to local disk, covering the same
re-score need as the eval recipe's predownload.py without a separate phase: the first run decodes
each shared chunk once and caches it; a re-score of a different model against the same ERA5
reads from disk, fetching the cloud zero times. Only the chunks touched, no materialized copy of
the grid; a static reanalysis store never goes stale.

This is not an argument against predownload, which buys rank-parallel bulk fetch, resumability, and
a durable pre-regridded artifact. The trade is scope: the cache buys the same re-score property
with no ETL phase to schedule and no full copy to provision, which matters when there is no cluster
to run the phase on.

Measured cold→warm (cache on NVMe): WB2 33→0 cloud fetches, 1.16 s → 0.81 s (1.4×); ARCO
54→0, 0.81 s → 0.59 s (1.4×). Fetch elimination is the deterministic win. The wall speedup is
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. It 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 sits
slightly above the persist-off de-dup figures above.

Honest boundary

Not a universal speed win — but on a chunk-1 store with large fields the reason is not raw
throughput. Re-run ARCO with redundancy removed from both sides (one init at unit lead spacing, so
162 requested = 162 unique = 0.67 GB either way) and E2S takes 1.38 s (486 MB/s) against the
feed's 1.61 s (416 MB/s): ~17% slower per byte, roughly at parity, over three independent
8-repeat runs.

python bench_hindcast.py --store arco --vars t2m --lead-step-h 1 \
  --max-lead-h 162 --n-init 1 --repeats 8

The gap in the ARCO row above opens elsewhere: 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 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 (WB2) assembly is negligible and nearly the whole de-dup ratio converts.

A degenerate batch_size=N also throws away the memory advantage: dense above peaks at 7.63 GB,
worse than the dense-buffer baseline it replaces (3.10 GB). The sweet spot is streaming with
bounded memory, and the large wins land when the chunk layout maps many samples onto shared
chunks (overlapping windows, verification grids, fat chunks).

A further boundary is scope rather than misuse: insitubatch does not replace a rank-parallel bulk
ETL. recipes/eval's predownload.py scales a fetch across nodes via distribute_work, resumes
after a failure, and leaves a durable pre-regridded artifact; one async event loop does none of
those. What insitubatch replaces is the predownload-then-read cycle for streaming consumption on a
single box. Rank-parallel training and inference are unaffected — each DDP rank streams its own
shard — but the bulk ETL phase is not something this feed competes with.

Scope / caveats

• Preliminary, single environment — numbers to be cross-posted after NVIDIA-side runs on
  target infrastructure.
• 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 all-NaN fill in ~20 ms without a network request. --start now defaults per store
  (ARCO: 2020-01-01) and both benchmarks fail loudly if a window reads back entirely NaN. E2S'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 indexing is not yet wired in
  the adapter.
• insitubatch is an optional dependency (the module raises a clear install hint if absent);
  nothing in the core E2S path changes. It is deliberately not exported from
  earth2studio.data.__init__ — that module imports unconditionally, and insitubatch is a
  data-extra dependency gated to Python ≥ 3.12, so exporting it would break the core import for
  everyone else. Reachable as from earth2studio.data.insitu import InSituForecastFeed.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

RE: Docs — added a recipe README. Please advise on API docs: the adapter is intentionally not in
earth2studio.data's namespace (see above), so it has no docs/modules/ entry; happy to add one
if you'd prefer it listed with a note about the optional dependency.

Dependencies

• insitubatch>=0.1.0 (PyPI) — added to the data
  extra, gated python_version>='3.12'. Optional; the adapter guards its import.
  Licensed MIT — LICENSE,
  declared as license = "MIT" in pyproject.toml and reported as SPDX MIT by GitHub. Happy to
  revisit the license if that helps adoption.

[greptile-apps]

Greptile Summary

This PR adds an optional insitubatch-based ERA5 streaming path. The main changes are:

• New InSituForecastFeed adapter for tensor batches and forecast coordinates.
• Optional insitubatch dependency under the data extra.
• Offline tests for feed layout, shifts, cache behavior, and lead validation.
• Hindcast, cache, and streaming scoring benchmark recipes.

Confidence Score: 4/5

The shifted-lead feed path needs a bounds check before merging.

• The new adapter is isolated and optional.
• The dependency marker matches the optional integration boundary.
• A documented history or verification lead can read outside the requested time window at store boundaries.

earth2studio/data/insitu.py

Important Files Changed

Filename	Overview
earth2studio/data/insitu.py	Adds the streaming feed and batch conversion logic; shifted leads need bounds validation against the requested init window.
pyproject.toml	Adds the optional insitubatch data-extra dependency for supported Python versions.
test/data/test_insitu.py	Adds offline coverage for feed shape, values, deduplication, transpose behavior, persistent cache, and lead-step validation.
recipes/insitubatch_hindcast/stream_score.py	Adds a streaming versus dense scoring benchmark using Persistence.
recipes/insitubatch_hindcast/bench_hindcast.py	Adds a before/after hindcast verification-read benchmark.
recipes/insitubatch_hindcast/bench_cache.py	Adds a cold/warm persistent-cache benchmark.
recipes/insitubatch_hindcast/README.md	Documents setup, benchmark usage, results, and caveats for the new recipe.
CHANGELOG.md	Documents the new optional streaming feed.

_{Reviews (1): Last reviewed commit: "Sort recipe imports for the current ruff..." | Re-trigger Greptile}

[greptile-apps]

Shifted Leads Escape Sample Window

When sample_range starts at the first init and lead_times includes a history lead like -6h, this shift reads before index 0; when the range reaches the store tail, positive verification leads read past the last time step. The feed accepts those documented lead values, so iteration can fail or return boundary samples from the wrong stored time instead of rejecting the invalid init window.

[emfdavid]

Nick - I will update the branch anytime there are substantive conflicts but chasing the change log seems silly. I will fix that when we get close to merging.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[emfdavid]

@nsobhani — following up on your point in the call that Earth2Studio already de-duplicates. You were right, and pulling on it turned up three separate places where I'd measured Earth2Studio in a configuration it doesn't actually run — each one flattering to us. Benchmarking someone else's system fairly is harder than I'd given it credit for; I clearly should have read more of recipes/eval and the source defaults before publishing numbers against them. Corrected figures below, and they are all lower.

compute_verification_times (recipes/eval/src/predownload_utils.py) collapses the (init, lead) grid onto the unique valid times, and predownload.py then fetches those — partitioned across ranks by distribute_work, one timestamp at a time within a rank, each written and flushed with a resume marker before the next. This recipe claimed the predownload had "no read de-duplication" and materialized the grid "up front". Both wrong — and wrong since #835 in April, so not a description of older code. Fixed in ebdc695.

The consequence is that §2's e2s leg is mislabelled "dense predownload". That dense buffer is my harness's own construction, and its 14 760 field reads are the un-deduplicated product where a real predownload would collapse that config onto ~160 unique valid times. I've relabelled it and flagged it pending a rerun. The streaming-vs-dense memory result in that section is unaffected — it compares two insitubatch modes against each other.

§1 is unaffected — its baseline is the live path, fetch_data(time=[item.time], ...) per work item in pipelines/forecast.py, dlesym.py, assimilation.py, which has no memory of what a neighbouring init already read. That is a different thing from the ETL.

Reading the rest of predownload.py also corrected a bigger assumption of mine. It is a deliberate cluster-scale ETL, not a workaround for slow reads: rank-parallel via distribute_work, resumable through per-timestamp markers, and it leaves a durable store — pre-regridded onto the model grid in the StormScope case — that many checkpoints can be scored against. Separating the phases is the point on a GPU cluster, keeping bulk IO on cheap CPU nodes and off GPU time.

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 on a single box: no separate phase to schedule, no materialized copy of the verification set to provision. Rank-parallel training and inference are unaffected, since each DDP rank streams its own shard. I've rewritten the recipe's framing and §3 to state that scope plainly, and to present the cache as an alternative for the no-ETL-phase case rather than as a substitute for predownload.

So the honest claim is narrower than what I had written: Earth2Studio de-duplicates at timestamp granularity in an offline pass; insitubatch de-duplicates at chunk granularity in the live path. On WB2's (8, 240, 121) chunks that residual is still 8 stored steps served per decode.

Checking that against the code turned up a second thing I had wrong, so I re-measured §1. The baseline 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. Re-run in both configurations:

store	requested → decodes	cache on (default)	cache off
WB2	5760 → 33 (174×)	6.8×	9.4×
ARCO	576 → 162 (3.6×)	1.2×	1.7×

The cache recovers only ~25% of the WB2 baseline wall — WB2 chunks are ~116 KB, so the network was never the bottleneck; the cost is 5760 decodes against 33, and the cache sits below zarr's codec pipeline so decode counts are identical either way. The win survives on WB2, but as a decode-elimination result rather than a reads-eliminated one.

ARCO is now essentially parity. Its cached figure is also disk-sensitive — the 363 MB cache is written cold each repeat, and moving it from the boot disk to local NVMe took the baseline 4.31 → 3.78 s (1.4× → 1.2×). I've quoted the NVMe number, the one favourable to your side, and the two distributions overlap at that point (feed 2.69–3.98 s against 3.71–4.02 s). That row reads as the boundary, not a win.

These supersede the 12.8× / 1.5× this PR previously carried. 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, but I can't yet account for the shift. Both legs of every comparison were measured in the same session, so the ratios are unaffected.

Worth being explicit about what those ratios are measured against, given the above: the live per-work-item path, not predownload. Against a valid-time-deduplicated baseline the advantage is exactly the sample-axis steps per chunk — WB2 261 decodes → 33 = 7.9×, ARCO 162 → 162 = none at all, since on a chunk-1 store chunk granularity is timestamp granularity. That is arithmetic from the index ranges rather than a measurement; I haven't measured that baseline's wall.

Which raises the thing I'd most like a second pair of eyes on: is there something in the getitem path that would collapse those repeated decodes, and that my harness is defeating? My reading is that LocalCachingStore sits at Store.get holding compressed buffers, below zarr's codec pipeline, so a repeated zarr_array.getitem(time_index) on a (8, 240, 121) chunk decodes the whole chunk again every time — with no decoded-chunk cache anywhere in the chain to catch it. That is the entire basis for the WB2 result, so if it's wrong, or if there's a configuration or access pattern that avoids it, the numbers above change a lot and I'd rather hear it now than after they're quoted somewhere.

So: how much of the work — yours at NVIDIA, and what you see from external Earth2Studio users — actually runs through predownload, and how much is interactive or single-box scoring where standing up the ETL phase isn't worth it? That determines whether §2 and §3 are worth measuring properly or whether §1's live path is the only part that earns its place here. I'd rather build the benchmark you'd find persuasive than guess twice.