Metadata-Version: 2.4
Name: insitubatch
Version: 0.0.1
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: cupy-cuda12x>=13.0 ; extra == 'gpu'
Requires-Dist: kvikio-cu12>=24.10 ; extra == 'gpu'
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
Provides-Extra: gpu
Provides-Extra: torch
Provides-Extra: virtual
Description-Content-Type: text/markdown

# 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). It turns an
existing Zarr archive into a shuffled, split-aware, GPU-saturating PyTorch
source — **with no reshard** — and a Python hot path that scales with **chunks,
not samples**.

> 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 drives concurrent reads; a bounded
shuffle-block buffer assembles batches; torch runs `num_workers=0`.

## Status

🚧 **Pre-alpha skeleton.** Abstractions and control flow are in place; the live
store read in `io.py` and the GPU path are stubbed. Not yet usable for real
training — this is the design substrate.

## Install (dev)

```bash
uv sync                  # core engine + dev tools
uv sync --extra torch    # add the torch IterableDataset surface
uv sync --extra gpu      # CUDA box only: cupy + kvikio zero-copy path
```

## Shape of the API (target)

```python
from insitubatch import split_by_chunk, ArrayGeometry, SplitName
from insitubatch.source import InSituDataset
from torch.utils.data import DataLoader

geom = ArrayGeometry("t2m", shape=(8760, 721, 1440), chunks=(24, 721, 1440), dtype=...)
manifest = split_by_chunk(geom, fractions=(0.8, 0.1, 0.1))

ds = InSituDataset(store, {"t2m": geom}, manifest, SplitName.TRAIN,
                   batch_size=32, block_chunks=16)

# parallelism lives in insitubatch's event loop, not in workers:
loader = DataLoader(ds, batch_size=None, num_workers=0)
for epoch in range(n_epochs):
    ds.set_epoch(epoch)
    for batch in loader:
        ...
```

## License

MIT — see [LICENSE](LICENSE).
