Metadata-Version: 2.4
Name: insitubatch
Version: 0.1.0
Summary: Train in place on n-dimensional cloud tensors: the data-loader orchestration layer on top of solved async IO (obstore / zarr v3 / icechunk).
Keywords: zarr,xarray,pytorch,dataloader,obstore,machine-learning,cloud,async
Author: David Stuebe
Author-email: David Stuebe <stu3b3+emfdavid@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering
Requires-Dist: numpy>=1.26
Requires-Dist: zarr>=3.0
Requires-Dist: xarray>=2024.9
Requires-Dist: obstore>=0.3
Requires-Dist: arraylake>=0.14 ; extra == 'arraylake'
Requires-Dist: xbatcher>=0.3 ; extra == 'bench'
Requires-Dist: plotly>=5.20 ; extra == 'bench'
Requires-Dist: gcsfs>=2024.6 ; extra == 'bench'
Requires-Dist: py-spy>=0.4 ; extra == 'bench'
Requires-Dist: scikit-learn>=1.4 ; extra == 'bench'
Requires-Dist: psutil>=5.9 ; extra == 'bench'
Requires-Dist: cloudpickle>=3 ; extra == 'cache'
Requires-Dist: mkdocs-material>=9.5 ; extra == 'docs'
Requires-Dist: mkdocstrings[python]>=0.26 ; extra == 'docs'
Requires-Dist: gcsfs>=2024.6 ; extra == 'gcsfs'
Requires-Dist: cupy-cuda12x>=13.0 ; extra == 'gpu'
Requires-Dist: kvikio-cu12>=24.10 ; extra == 'gpu'
Requires-Dist: icechunk>=0.2 ; extra == 'icechunk'
Requires-Dist: jax>=0.4.30 ; extra == 'jax'
Requires-Dist: flax>=0.8 ; extra == 'jax'
Requires-Dist: optax>=0.2 ; extra == 'jax'
Requires-Dist: tensorflow>=2.16 ; extra == 'tf'
Requires-Dist: torch>=2.2 ; extra == 'torch'
Requires-Dist: torchdata>=0.10 ; extra == 'torch'
Requires-Dist: virtualizarr>=1.2 ; extra == 'virtual'
Requires-Python: >=3.12
Project-URL: Documentation, https://emfdavid.github.io/insitubatch/
Project-URL: Homepage, https://emfdavid.github.io/insitubatch/
Project-URL: Repository, https://github.com/emfdavid/insitubatch
Project-URL: Changelog, https://github.com/emfdavid/insitubatch/blob/main/CHANGELOG.md
Project-URL: Issues, https://github.com/emfdavid/insitubatch/issues
Provides-Extra: arraylake
Provides-Extra: bench
Provides-Extra: cache
Provides-Extra: docs
Provides-Extra: gcsfs
Provides-Extra: gpu
Provides-Extra: icechunk
Provides-Extra: jax
Provides-Extra: tf
Provides-Extra: torch
Provides-Extra: virtual
Description-Content-Type: text/markdown

# insitubatch

[![PyPI](https://img.shields.io/pypi/v/insitubatch.svg)](https://pypi.org/project/insitubatch/)
[![CI](https://github.com/emfdavid/insitubatch/actions/workflows/ci.yml/badge.svg)](https://github.com/emfdavid/insitubatch/actions/workflows/ci.yml)
[![docs](https://github.com/emfdavid/insitubatch/actions/workflows/docs.yml/badge.svg)](https://emfdavid.github.io/insitubatch/)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

**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](examples/microscopy/) — 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](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](https://emfdavid.github.io/insitubatch/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](https://emfdavid.github.io/insitubatch/benchmarks/).

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](examples/); validated free-threading-correct on 3.13t. Not yet built:
`Regrid` + the **GPU/device** transform stage — see the roadmap in [DESIGN.md](DESIGN.md).

📖 **Docs:** <https://emfdavid.github.io/insitubatch/>
(see [Tuning](https://emfdavid.github.io/insitubatch/tuning/) for the
chunks↔concurrency↔memory model).

## Install

```bash
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:

```bash
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

```bash
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:
> ```bash
> 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:

```bash
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](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.

```python
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`:

```python
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`](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:

```console
$ 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](LICENSE).
