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/ — two 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).

Benchmark results

Preliminary, one n2-standard-8-class box (15 GB RAM), cold reads, gcsfs anon on both the
before and after side — so the delta isolates insitubatch's read-planning + streaming, not an
obstore-vs-gcsfs artifact. Public WeatherBench2 and ARCO ERA5 stores.

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

store	layout	requested reads	unique decodes	wall speedup
WB2 240×121, 6-h	chunks=(8,240,121) (fat time-chunk)	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×

2. Streaming vs dense materialization (stream_score.py, N = 120 inits × 40 leads) — all
three modes produce identical RMSE:

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 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 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, replacing the eval
recipe's predownload.py sentinel: 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. No predownload step, no reshard, only the chunks touched; a static reanalysis store never
goes stale. Measured cold→warm (cache on NVMe): WB2 33→0 cloud fetches (~1.4× wall), ARCO
54→0 (~2.2× wall). Fetch elimination is the deterministic win; the wall speedup scales with how
IO-bound / metered the source is and is understated by this box's cheap same-region reads. 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. On a chunk-1 store with large fields, against an unbounded
concurrent gather, insitubatch's bounded-inflight scheduling trails per byte (ARCO ~2×; the reads
are already minimal, but the dense output the model consumes must still be assembled, and E2S's
flat gather saturates bandwidth on 4 MB chunks). 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).

Scope / caveats

• Preliminary, single environment — numbers to be cross-posted after NVIDIA-side runs on
  target infrastructure.
• 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.

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 recipe readme - please advise on API docs etc

Dependencies

• insitubatch>=0.1.0 (PyPI) — added to the data
  extra, gated python_version>='3.12'. Optional; the adapter guards its import.

[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.