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 trains on ERA5/weather over time, segments OME-NGFF microscopy volumes over Z (runnable example — raw image + label mask co-batched with no reshard), and maps cleanly onto radio-astronomy visibilities — one contract, any single sample axis, variables that chunk it differently.

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. See DESIGN.md.

Why

The classic PyTorch DataLoader spreads work across worker processes, each running a synchronous __getitem__. Against cloud Zarr that means 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. insitubatch inverts it: one async event loop streams stored chunks under a single concurrency budget and scatters them into a bounded pool that assembles batches — the pool doubles as the cache; torch runs num_workers=0.

The payoff is a two-regime story against the worker-process DataLoader. On a well-chunked store it matches a hand-tuned worker/xbatcher pool (swept to 32 workers) while running in one process at bounded memory, reaching first batch in ~ms rather than 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 pool: read planning decodes each shared chunk once where a per-sample __getitem__ re-reads it, so the win grows with samples-per-chunk (to ~25× at the fat end of the ERA5 sweep) and cross-epoch caching compounds it. The honest boundary: at the one-sample-per-chunk (GRIB) end there is nothing to amortize, so a tuned pool edges ahead on single-pass throughput, and against an unbounded concurrent gather on large fields bounded-inflight streaming trails per byte — the sweet spot is streaming with bounded memory, not a universal speed win. Full comparison: Benchmarks.

Status

🚧 alpha, but validated on real cloud IO. On an in-region S3 run (c6id.8xlarge, ERA5-shaped 721×1440 fields, sample_chunk=8), insitubatch delivers ~8× the throughput of a tuned xbatcher/worker DataLoader baseline (swept to 32 workers) and reaches its first batch ~10× sooner — the map-style baseline re-decodes a whole chunk per sample; insitubatch reads each chunk once. Full numbers + methodology: the benchmarks page.

The engine is the decoupled fetch scheduler: reads flatten to stored chunks under one max_inflight budget (no nested inner/outer concurrency caps), decoded tiles scatter into a ChunkPool that is the assembly buffer and the cache (byte budget + pin/LRU, heap or mmap-on-NVMe). Read concurrency and residency/shuffle span are independent dials — the decoupling reaches ~1 GB/s at flat, low memory (validated on S3; see below). Built: planner + chunk-aligned splits, async obstore reads, the scheduler + pool (with decode-once caching, cross-epoch and cross-run via persist=True), chunk/batch transforms (incl. a fitted StandardScaler), prefetch, the torch / JAX / TF surfaces, and runnable examples; validated free-threading-correct on 3.13t. Not yet built: Regrid + the GPU/device transform stage — see the roadmap in DESIGN.md.

📖 Docs: https://emfdavid.github.io/insitubatch/ (see Tuning for the chunks↔concurrency↔memory model).

Install

pip install insitubatch              # core engine (numpy Batch; no framework)
pip install "insitubatch[torch]"     # + torch DLPack adapter (insitubatch.frameworks)
pip install "insitubatch[jax]"       # + JAX adapter
pip install "insitubatch[tf]"        # + TensorFlow adapter

For development:

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

Tests

uv run pytest -q                              # the suite
uv run ruff check src tests bench             # lint
uv run mypy src                               # types

The torch-handoff tests skip unless torch is installed (uv sync --extra torch); the same is enforced in CI.

One framework per environment. torch, JAX and TensorFlow cannot coexist in one Python process — together they load duplicate OpenMP/XLA/protobuf runtimes and the process crashes (SIGSEGV / abort). Separate pytest processes in one env are not enough: TF (via its bundled Keras 3) transitively imports JAX whenever JAX is installed, so the two collide even if only the TF tests are selected. So install just one adapter at a time when running the framework tests:

uv sync --extra torch && uv run pytest -q   # torch adapter + core
uv sync --extra jax   && uv run pytest -q   # JAX adapter (others importorskip-skip)
uv sync --extra tf    && uv run pytest -q   # TF adapter

CI does exactly this — one job per framework, each with a single adapter installed, plus a separate lint/types job that installs every extra but runs no pytest (mypy doesn't import the frameworks, so co-installation is harmless there). This is a framework-coexistence limitation, not an insitubatch one — the core engine and each adapter are independent.

Free-threaded (3.13t)

The engine is free-threading-correct by construction: the ChunkPool's scatter does its disjoint copy before the lock and publishes readiness under it, so the lock — not the GIL — is the happens-before edge to the consuming gather. The race probe is test_pool_concurrent_scatter_is_race_free (64 tiles, 32 threads).

Run the suite GIL-free on a free-threaded interpreter:

uv python install 3.13t
# Separate env so the default .venv stays put. numcodecs has no free-threaded
# wheel yet, so it compiles from sdist (needs a C/C++ compiler: Xcode CLT on
# macOS, gcc/gcc-c++ on Linux). torch/bench have no FT wheels -> core deps only.
UV_PROJECT_ENVIRONMENT=.venv-ft uv sync --python 3.13t

# numcodecs re-enables the GIL on import (not yet declared GIL-safe), so force it
# off and confirm it took before trusting the run:
PYTHON_GIL=0 UV_PROJECT_ENVIRONMENT=.venv-ft uv run --python 3.13t \
  python -c "import sys, zarr, numcodecs; assert not sys._is_gil_enabled(); print('GIL-free OK')"
PYTHON_GIL=0 UV_PROJECT_ENVIRONMENT=.venv-ft uv run --python 3.13t pytest -q

CI mirrors this: a {3.12, 3.13} matrix plus a 3.13t job that asserts the GIL is actually off before testing. Throughput is GIL-independent by design — fetch (obstore/Rust), decode (numcodecs zstd, C), and scatter/gather (vectorized numpy) all release the GIL — so 3.13t runs at the same speed as the GIL build, not faster. The free-threading work is correctness + future-proofing, not a speedup; not depending on the GIL is the point (see DESIGN.md).

Shape of the API

The core InSituDataset is a framework-neutral source of numpy Batch objects — it inherits nothing framework-specific. You iterate its split views (ds.train shuffled, ds.val / ds.test / ds.all deterministic), which all share one pool, so a chunk two splits both read decodes once. Handoff to torch / JAX / TF is a thin, optional DLPack adapter (re-exported from the package root; defined in insitubatch.frameworks) — the core imports no framework, and importing insitubatch pulls none in.

from insitubatch import InSituDataset, obstore_store, open_geometries, split_by_chunk

# The engine reads a zarr Store; build one per backend. obstore_store covers
# file://, s3://, gs://, az://. (fsspec_store reaches GCS Rapid/requester-pays;
# arraylake_store opens an Icechunk session — same InSituDataset below.)
store = obstore_store("file:///data/era5.zarr")  # or "s3://bucket/era5.zarr"
geoms = open_geometries(store)  # {var: ArrayGeometry} from zarr metadata
# contiguous chunk blocks by default (no time-series leakage);
# pass contiguous=False for exchangeable samples (independent scenes)
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
        ...

Hand off to a framework — zero-copy on CPU via DLPack for torch and JAX; TF takes one CPU copy (its experimental DLPack is unreliable — see frameworks.to_tf). The ecosystems differ — torch needs a Dataset subclass, JAX iterates directly, TF wraps via from_generator:

from insitubatch import as_tf_dataset, as_torch, to_jax
from torch.utils.data import DataLoader

# torch: parallelism is in our event loop, so num_workers=0, batch_size=None
loader = DataLoader(as_torch(ds.train), batch_size=None, num_workers=0)  # {var: torch.Tensor}

for batch in ds.train:      # JAX: iterate a view, convert each batch
    jbatch = to_jax(batch)  # {var: jax.Array}

tfds = as_tf_dataset(ds.val)  # a tf.data.Dataset

Transforms — and checking one before you train

Two hooks, placed by cost: a chunk_transform (DecodedChunk) -> DecodedChunk runs per decoded chunk (one variable), before the cache boundary, so its output is cached — the home for scaling, unit conversion, dtype cast, regrid; and a batch_transform (Batch) -> Batch runs per assembled batch (all variables aligned), uncached — for cross-variable derived fields and per-sample random augmentation. Both are pure numpy; see examples/transforms.py (K→C chunk stage + windspeed batch stage).

A chunk_transform must be vectorized numpy that releases the GIL (a per-element Python loop serializes the decode pool), and a reshaping one (regrid) must declare output_inner(geom) -> (inner_shape, dtype) so the cache can size its slot. Check both against one chunk of your real store before training:

$ insitubatch-check-transform \
    gs://weatherbench2/datasets/era5/1959-2022-6h-128x64_equiangular_with_poles_conservative.zarr \
    --var 2m_temperature --transform examples/transforms.py:kelvin_to_celsius --skip-signature

  sample axis : 92040 samples, 40/chunk, 2301 chunks
  chunk 0    : 40 samples -> source shape (40, 128, 64) = 1.3 MB decoded
transform output:
  (40, 128, 64) float32  ->  (40, 128, 64) float32   shape- and dtype-preserving
cacheability: shape/dtype-preserving, no output_inner needed -> cacheable as-is.
GIL-release probe (thread-scaling, 4 threads):
  speedup 3.50x (>= 2.40) -> releases the GIL (vectorized).
PASS: chunk_transform checks all passed.

The target is module:attr or path/to/file.py:attr (a transform class is instantiated). It reports the chunk geometry, validates a declared output_inner against the real output (catching the mismatch the cache would later reject), and gives a GIL-release verdict — a non-zero exit gates a pre-commit hook. Pass --no-gil-probe for a fast structural-only check; the GIL probe needs a realistically-sized chunk (a toy array is dominated by call overhead). For the reshaping path, try --transform examples/transforms.py:Coarsen — a chunk-local regrid that halves the grid and declares output_inner, so the report shows the validated shape change.

License

MIT — see LICENSE.