insitubatch

Train in place on n-dimensional cloud tensors.

insitubatch is the data-loader orchestration layer that sits on top of already-solved async cloud IO (obstore / zarr v3 / icechunk) for PyTorch, Jax and TensorFlow. It turns an existing Zarr archive into a shuffled, split-aware data source built to keep the GPU fed — with no reshard — and a Python hot path that scales with chunks, not samples.

It is domain-general: the sample axis is a role, not a fixed dimension. The same engine forecasts ERA5 weather over time, segments OME-NGFF microscopy volumes over Z, denoises Hubble telescope frames streamed straight out of FITS, and reconstructs SDSS galaxy spectra fiber by fiber — each a different geometry on a real public store, none of them resharded. See Examples and the use-case tables.

Quote

The IO race is over (obstore/icechunk saturate the NIC). The loader race is open. insitubatch builds the layer that projects like light-speed-io and hypergrib stopped one step short of.

Where it wins. On a well-chunked store it matches a hand-tuned worker DataLoader (swept to its best worker count) at a fraction of the memory — one process, bounded residency, ~ms to first batch instead of seconds of pool cold-start. When the chunk layout isn't sample-optimized — fat time-chunks, overlapping windows, verification grids — it pulls far ahead of even a tuned worker pool, because read planning decodes each shared chunk once where per-sample workers re-read it (the win grows with samples-per-chunk). It is not a universal speed win: at the one-sample-per-chunk (GRIB) end, or against an unbounded gather on large fields, a tuned pool can edge ahead per byte. Numbers: Benchmarks.

Cross-domain examples

Five runnable showcases, each a different geometry on a real public store — and each with an offline synthetic mode so it runs with no network or credentials. Details and commands: Examples.

	Domain & store	The geometry it exercises
advection	WeatherBench2 ERA5 (gs://, anonymous)	Input at t, target at t+24 h as offset views of one array — then the same CNN trained in torch, JAX and TF
microscopy	IDR OME-NGFF (T,C,Z,Y,X) on s3://idr	Samples a middle axis (Z); raw image and label mask chunked 1 vs 30 planes deep, co-batched
hubble	Hubble WFC3/IR frames on MAST's s3://stpubdata	FITS, not zarr — virtual byte-range references, streamed and trained in place
sdss	SDSS DR17 spectra, or the public reference stores	One archive, two chunk layouts by byte arithmetic alone: 64 spectra per chunk, or one per chunk at archive scale
wb2 pair	WeatherBench2 ERA5	The same task on insitu and on an xbatcher worker stack, for the cold-start trade-off

The problem, and the inversion

The classic PyTorch DataLoader puts parallelism in worker processes, each running a synchronous __getitem__. Against cloud Zarr that fights itself: no shared chunk cache (every worker re-reads the same chunk), no way to drive async obstore, and dask thread pools nested inside forked workers. The usual escape — resharding to one-sample-per-file — is a second copy of the dataset that throws away the chunk locality the store already has.

insitubatch keeps the data in place and inverts the loader:

Classic DataLoader: parallelism lives in num_workers OS processes, each running a synchronous __getitem__. insitubatch: parallelism lives in one async event loop; batch assembly is the consumer.

That single move unlocks async obstore, a shared chunk cache, bounded memory, and prefetch overlap with the training step; torch runs num_workers=0. Architecture has the full frictions breakdown, the loader/prefetch diagrams, and the read-plan abstraction; DESIGN.md has the why.

Shape of the API

The core InSituDataset is a framework-neutral iterable of numpy Batch objects; torch / JAX / TF handoff is a thin optional DLPack adapter, re-exported from the package root (defined in insitubatch.frameworks) — importing insitubatch pulls in no framework.

from insitubatch import (
    InSituDataset,
    as_tf_dataset,
    as_torch,
    obstore_store,
    open_geometries,
    split_by_chunk,
    to_jax,
)
from torch.utils.data import DataLoader

# The engine reads a zarr Store; obstore_store builds one for file://, s3://, gs://.
# (fsspec_store for GCS Rapid/requester-pays; arraylake_store for Icechunk sessions.)
store = obstore_store("file:///data/era5.zarr")  # or "s3://bucket/era5.zarr"
geoms = open_geometries(store)           # {var: ArrayGeometry} from zarr metadata
manifest = split_by_chunk(geoms["t2m"], fractions=(0.8, 0.1, 0.1))

ds = InSituDataset(store, manifest, batch_size=32, block_chunks=16)

for epoch in range(n_epochs):
    ds.set_epoch(epoch)
    for batch in ds.train:               # numpy Batch: {var: np.ndarray} + sample_indices
        ...
    for batch in ds.val:                 # deterministic; shares the pool with train
        ...

# Framework handoff (DLPack, zero-copy on CPU for torch/JAX; TF copies once):
loader = DataLoader(as_torch(ds.train), batch_size=None, num_workers=0)  # torch
jbatch = to_jax(next(iter(ds.train)))                                    # JAX:   {var: jax.Array}
tfds = as_tf_dataset(ds.val)                                             # TF:    tf.data.Dataset

See Examples for working CNN models built on InSituDataset — a three-framework ERA5 forecast, OME-NGFF microscopy segmentation over Z, Hubble frame denoising from FITS, and SDSS spectral reconstruction against a PCA baseline.

A runnable, network-free version of this — paralleling the Earthmover dataloader-demo, with a spatial subregion pulled out by a batch_transform — lives in examples/wb2_dataloader.py:

uv run python -m examples.wb2_dataloader            # tiny synthetic data, no network
uv run python -m examples.wb2_dataloader \
    --url s3://bucket/era5.zarr --var 2m_temperature --subregion 48,48 --request-payer

Install (dev)

uv sync                  # core engine + dev tools
uv sync --extra torch    # torch handoff (frameworks.as_torch)
uv sync --extra jax      # JAX handoff (frameworks.to_jax)
uv sync --extra tf       # TF handoff (frameworks.as_tf_dataset)
uv sync --extra bench    # benchmark suite (xbatcher baseline + plotly)
uv sync --extra gpu      # CUDA box only: cupy + kvikio zero-copy path

Status

Alpha — validated on real cloud IO. Built: planner + chunk-aligned splits, async obstore reads, the decoupled fetch Scheduler + ChunkPool (assembly buffer and cache — byte budget + pin/LRU, heap or mmap-on-NVMe, with cross-run persistence via persist=True), approximate (shuffle-block) shuffle, chunk/batch transforms (incl. a fitted StandardScaler), prefetch, windowed multi-offset sampling, an arbitrary sample axis with per-variable chunk sizes, and the torch / JAX / TF surfaces. Not yet built: a production Regrid + the GPU/device transform stage. (A window that crosses chunk boundaries as a single slab read is an explicit non-goal, not a gap — discrete offsets from an anchor already express it.)

DESIGN.md is the single source of truth for status, the roadmap, and the scope limits.

Contributions are welcome, and the project is being built to be maintained by more than one person: see Contributing for the scope limits and the dev setup, and GOVERNANCE.md for how decisions get made and how the core developer group grows.