insitubatch 0.0.1__py3-none-any.whl

insitubatch/__init__.py

@@ -0,0 +1,44 @@
+"""insitubatch -- train in place on n-dimensional cloud tensors.
+
+The loader-orchestration layer that sits on top of *already-solved* async cloud
+IO (obstore / zarr v3 / icechunk): 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.
+
+See DESIGN.md for the full rationale.
+"""
+
+from __future__ import annotations
+
+from .buffer import BufferConfig, ShuffleBlockBuffer
+from .io import AsyncChunkReader, IOConfig
+from .plan import ReadPlan, build_read_plan, dedup_ratio
+from .shuffle import block_shuffled_order, chunk_permutation, shuffle_quality
+from .split import SplitManifest, split_by_chunk
+from .store import ensure_local_dir, open_geometries, store_from_url
+from .types import ArrayGeometry, Batch, ChunkRead, DecodedChunk, SplitName
+
+__version__ = "0.0.1"
+
+__all__ = [
+    "ArrayGeometry",
+    "AsyncChunkReader",
+    "Batch",
+    "BufferConfig",
+    "ChunkRead",
+    "DecodedChunk",
+    "IOConfig",
+    "ReadPlan",
+    "ShuffleBlockBuffer",
+    "SplitManifest",
+    "SplitName",
+    "block_shuffled_order",
+    "build_read_plan",
+    "chunk_permutation",
+    "dedup_ratio",
+    "ensure_local_dir",
+    "open_geometries",
+    "shuffle_quality",
+    "split_by_chunk",
+    "store_from_url",
+]

insitubatch/buffer.py

@@ -0,0 +1,88 @@
+"""Bounded shuffle-block buffer: residency + batch assembly.
+
+Holds decoded chunks for a window of ``block_chunks`` chunks, draws shuffled
+batches across that window, and evicts chunks once fully drained. This is the
+memory-bounding component: peak residency is O(block_chunks), independent of the
+number of samples per epoch or the batch size.
+
+The batch assembly does a single coalesced gather per variable (one fancy-index
+copy), never a Python per-sample loop -- the constraint David's S3 benchmark
+imposed (Python per-chunk overhead bounds throughput).
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from .types import Batch, DecodedChunk
+
+
+@dataclass(slots=True)
+class BufferConfig:
+    block_chunks: int = 16
+    """Window size in chunks. Larger == better shuffle, more memory."""
+
+    batch_size: int = 32
+
+
+@dataclass(slots=True)
+class ShuffleBlockBuffer:
+    """Accumulates decoded chunks and emits shuffled, coalesced batches."""
+
+    config: BufferConfig
+    seed: int = 0
+    _chunks: dict[tuple[str, int], DecodedChunk] = field(default_factory=dict)
+    # Pending draws: rows of (array, chunk_index, within) flattened per variable.
+    _pending: deque[int] = field(default_factory=deque)
+
+    def add(self, chunk: DecodedChunk) -> None:
+        self._chunks[(chunk.read.array, chunk.read.chunk_index)] = chunk
+
+    def ready(self) -> bool:
+        """Enough buffered to safely emit a well-mixed batch."""
+        return len(self._chunks) >= self.config.block_chunks
+
+    def gather_batch(
+        self,
+        rows: np.ndarray,
+        variables: list[str],
+        sample_chunk_size: int,
+    ) -> Batch:
+        """Assemble one batch from ``rows`` of ``[chunk_id, within]`` draws.
+
+        ``rows`` are pre-shuffled draw coordinates (see shuffle.block_shuffled_order).
+        Draws are grouped by chunk so each resident array is touched once (one
+        coalesced fancy-index per chunk); ``data`` and ``sample_indices`` are
+        emitted in the *same* grouped order so row ``i`` of every variable and
+        ``sample_indices[i]`` refer to the same sample. Intra-batch order is thus
+        grouped-by-chunk -- irrelevant for training, and the cross-batch shuffle
+        is preserved.
+
+        ``sample_chunk_size`` is the array's true chunk length (from geometry),
+        used to recover global sample indices -- NOT inferred from a resident
+        chunk, which may be a short final chunk.
+        """
+        chunk_ids = rows[:, 0]
+        within = rows[:, 1]
+        uniq = np.unique(chunk_ids)
+
+        out: dict[str, list[np.ndarray]] = {v: [] for v in variables}
+        idx_pieces: list[np.ndarray] = []
+        for cid in uniq:
+            w = within[chunk_ids == cid]
+            idx_pieces.append(cid * sample_chunk_size + w)
+            for var in variables:
+                out[var].append(self._chunks[(var, int(cid))].data[w])
+
+        arrays = {var: np.concatenate(pieces, axis=0) for var, pieces in out.items()}
+        return Batch(arrays=arrays, sample_indices=np.concatenate(idx_pieces))
+
+    def evict_drained(self, still_needed: set[tuple[str, int]]) -> int:
+        """Drop chunks no longer referenced by any pending draw. Returns count."""
+        drop = [k for k in self._chunks if k not in still_needed]
+        for k in drop:
+            del self._chunks[k]
+        return len(drop)

insitubatch/io.py

@@ -0,0 +1,172 @@
+"""Async IO driver: the obstore win.
+
+This is where insitubatch *stands on* solved cloud IO rather than reinventing
+it. A single dedicated asyncio event loop (one OS thread) issues many concurrent
+chunk reads against a zarr v3 store -- ideally the obstore-backed store, whose
+``get_ranges_async`` coalesces concurrent range requests in a single coroutine
+and saturates the NIC without spawning Python threads per request.
+
+Design rules (DESIGN.md, "the inversion"):
+  * Parallelism lives *here*, in the event loop's concurrency, NOT in
+    torch.DataLoader worker processes. The torch surface runs with
+    ``num_workers=0``.
+  * A bounded in-flight window (semaphore of ``max_inflight`` chunks) caps memory
+    at O(in-flight chunks), independent of batch size.
+  * Decode releases the GIL (numcodecs C codecs do) so decode overlaps IO. The
+    Python hot path stays O(reads); never decode per-sample in Python.
+
+Status: SKELETON. The async store wiring is stubbed at the marked TODOs so the
+module imports and the control flow is reviewable without a live store.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+import numpy as np
+import zarr.api.asynchronous as za
+
+from .plan import ReadPlan
+from .store import store_from_url
+from .types import ArrayGeometry, ChunkRead, DecodedChunk
+
+
+@dataclass(slots=True)
+class IOConfig:
+    max_inflight: int = 16
+    """Upper bound on chunks in flight. Memory ~= max_inflight * chunk_nbytes."""
+
+    decode_threads: int = 4
+    """Worker threads for GIL-releasing decode, overlapped with IO."""
+
+
+class AsyncChunkReader:
+    """Owns one asyncio event loop on a background thread and fans out reads.
+
+    The public API is deliberately *synchronous and iterator-shaped* so the rest
+    of the engine (buffer, torch surface) needs no async knowledge -- the loop is
+    an implementation detail hidden behind a thread-safe queue.
+    """
+
+    def __init__(
+        self,
+        store_url: str,
+        geometries: dict[str, ArrayGeometry],
+        config: IOConfig | None = None,
+        **store_kwargs: object,
+    ) -> None:
+        self._url = store_url
+        self._store_kwargs = store_kwargs
+        self._geometries = geometries
+        self._config = config or IOConfig()
+        self._arrays: dict[str, za.AsyncArray] = {}  # opened lazily on the loop
+        self._loop = asyncio.new_event_loop()
+        self._sem: asyncio.Semaphore | None = None  # created on the loop bootstrap
+        self._open_lock: asyncio.Lock | None = None
+        self._ready = threading.Event()
+        self._thread = threading.Thread(target=self._run_loop, daemon=True, name="insitu-io")
+        self._thread.start()
+        self._ready.wait()  # don't return until the loop + primitives exist
+
+    # -- loop lifecycle -----------------------------------------------------
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._sem = asyncio.Semaphore(self._config.max_inflight)
+        self._open_lock = asyncio.Lock()
+        self._loop.call_soon(self._ready.set)
+        self._loop.run_forever()
+
+    def close(self) -> None:
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join(timeout=5)
+
+    def __enter__(self) -> AsyncChunkReader:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    # -- public, synchronous surface ---------------------------------------
+
+    _SENTINEL = object()
+
+    def read_plan(self, plan: ReadPlan) -> Iterator[DecodedChunk]:
+        """Fetch every read in ``plan`` concurrently; yield decoded chunks ASAP.
+
+        Yields in completion order (not plan order) so a slow chunk never stalls
+        the others -- the buffer downstream is responsible for reordering/gather.
+
+        The bridge uses a thread-safe ``queue.Queue`` and a done-callback so the
+        sentinel is *always* delivered -- even if the driver raises -- and any
+        exception is re-raised on the caller's thread rather than deadlocking it.
+        """
+        out_q: queue.Queue = queue.Queue()
+
+        def _on_done(fut: object) -> None:
+            try:
+                fut.result()  # type: ignore[attr-defined]  # surface driver errors
+            except Exception as exc:  # noqa: BLE001 - forwarded to the consumer
+                out_q.put(exc)
+            finally:
+                out_q.put(self._SENTINEL)
+
+        fut = asyncio.run_coroutine_threadsafe(self._drive(plan, out_q), self._loop)
+        fut.add_done_callback(_on_done)
+
+        while True:
+            item = out_q.get()
+            if item is self._SENTINEL:
+                break
+            if isinstance(item, Exception):
+                raise item
+            yield item
+
+    # -- async internals ----------------------------------------------------
+
+    async def _drive(self, plan: ReadPlan, out_q: queue.Queue) -> None:
+        assert self._sem is not None  # guaranteed by _ready.wait() in __init__
+        await self._ensure_arrays()
+
+        async def one(read: ChunkRead) -> None:
+            async with self._sem:  # bound in-flight chunks -> bounded memory
+                decoded = await self._fetch_and_decode(read)
+            out_q.put(decoded)  # stdlib queue: thread-safe, non-blocking
+
+        await asyncio.gather(*(one(r) for r in plan.reads))
+
+    async def _ensure_arrays(self) -> None:
+        """Open one AsyncArray per variable, once, sharing the store."""
+        if self._arrays:
+            return
+        assert self._open_lock is not None
+        async with self._open_lock:
+            if self._arrays:  # double-checked: another coroutine may have won
+                return
+            store = store_from_url(self._url, **self._store_kwargs)  # type: ignore[arg-type]
+            for name in self._geometries:
+                self._arrays[name] = await za.open_array(store=store, path=name, mode="r")
+
+    async def _fetch_and_decode(self, read: ChunkRead) -> DecodedChunk:
+        """Fetch + decode one chunk via the zarr v3 async codec pipeline.
+
+        The selection is exactly one chunk along the sample axis, full on the
+        inner dims (the v1 sample-geometry contract). zarr fans the underlying
+        byte-range reads out through obstore and runs the decode pipeline; for
+        single-chunk inner dims this touches exactly one stored chunk.
+
+        TODO(perf): decode currently runs inside zarr's pipeline on the loop. If
+        it shows up as a GIL bottleneck, move it to a thread pool
+        (numcodecs C codecs release the GIL) -- the bounded-fan-out structure
+        here already isolates that change to this method.
+        """
+        geom = self._geometries[read.array]
+        arr = self._arrays[read.array]
+        samples = geom.samples_in_chunk(read.chunk_index)
+        selection = (slice(samples.start, samples.stop), *(slice(None) for _ in geom.inner_shape))
+        block = await arr.getitem(selection)
+        return DecodedChunk(read=read, data=np.asarray(block), sample_offset=samples.start)

insitubatch/plan.py

@@ -0,0 +1,112 @@
+"""Read planning: samples -> deduplicated chunk reads.
+
+This is the crux abstraction. Given the samples required for a window of the
+epoch and the array geometries, produce the *minimal* set of chunk reads plus a
+gather map describing where each sample lives once those chunks are decoded.
+
+Why this matters (DESIGN.md, "the spectrum"):
+  - Fat chunks: many samples share one chunk -> dedup collapses N samples to 1
+    read; the shared decoded chunk is gathered N times. This is the shared-cache
+    win that the classic per-worker DataLoader cannot get.
+  - GRIB-per-timestep: one sample per chunk -> no dedup possible, but the plan
+    still drives a single wide async fan-out (B samples == B concurrent reads),
+    which is exactly where obstore earns its keep.
+
+The Python hot path here is O(reads), never O(samples) once gathered, which is
+the constraint David's S3 benchmark imposed (Python per-chunk overhead bounds
+throughput; never loop per-sample in Python).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+
+from .types import ArrayGeometry, ChunkRead
+
+
+@dataclass(slots=True)
+class Gather:
+    """Where one sample lives within a decoded chunk.
+
+    ``read_index`` indexes into ``ReadPlan.reads``; ``within`` is the offset of
+    the sample inside that chunk's decoded array along the sample axis.
+    """
+
+    read_index: int
+    within: int
+
+
+@dataclass(slots=True)
+class ReadPlan:
+    """A deduplicated batch of chunk reads plus the gather map back to samples.
+
+    One ``ReadPlan`` typically covers enough samples to (a) saturate the async
+    fan-out and (b) fill the shuffle-block buffer. ``reads`` is what the IO
+    driver fetches; ``gathers[v]`` reconstructs the requested samples for
+    variable ``v`` from the decoded chunks.
+    """
+
+    reads: list[ChunkRead]
+    gathers: dict[str, list[Gather]]
+    sample_indices: np.ndarray  # global sample indices, in requested order
+
+    @property
+    def n_reads(self) -> int:
+        return len(self.reads)
+
+
+def build_read_plan(
+    sample_indices: Sequence[int],
+    geometries: dict[str, ArrayGeometry],
+) -> ReadPlan:
+    """Build a deduplicated read plan for ``sample_indices`` across all variables.
+
+    All variables are assumed aligned on the sample axis (same length, possibly
+    different chunking) -- the common case for co-registered NWP variables. A
+    sample at global index ``s`` requires chunk ``geom.chunk_of(s)`` from *each*
+    variable; identical chunks requested by multiple samples are read once.
+
+    Parameters
+    ----------
+    sample_indices:
+        Global sample-axis indices to fetch, in the order they should appear.
+    geometries:
+        Variable name -> :class:`ArrayGeometry`.
+
+    Returns
+    -------
+    ReadPlan
+    """
+    idx = np.asarray(sample_indices, dtype=np.int64)
+    reads: list[ChunkRead] = []
+    read_lookup: dict[ChunkRead, int] = {}
+    gathers: dict[str, list[Gather]] = {name: [] for name in geometries}
+
+    for name, geom in geometries.items():
+        # Vectorized chunk assignment for this variable across all samples.
+        chunk_ids = idx // geom.sample_chunk_size
+        within = idx - chunk_ids * geom.sample_chunk_size
+        for c, w in zip(chunk_ids.tolist(), within.tolist(), strict=True):
+            read = ChunkRead(array=name, chunk_index=int(c))
+            ri = read_lookup.get(read)
+            if ri is None:
+                ri = len(reads)
+                read_lookup[read] = ri
+                reads.append(read)
+            gathers[name].append(Gather(read_index=ri, within=int(w)))
+
+    return ReadPlan(reads=reads, gathers=gathers, sample_indices=idx)
+
+
+def dedup_ratio(plan: ReadPlan) -> float:
+    """Samples-per-read averaged over variables.
+
+    1.0 == degenerate (GRIB-per-timestep, no sharing); higher == fatter chunks
+    with more cache reuse. A quick lever for understanding which regime a dataset
+    + batch size lands in.
+    """
+    n_samples = len(plan.sample_indices) * max(len(plan.gathers), 1)
+    return n_samples / plan.n_reads if plan.n_reads else 0.0

insitubatch/shuffle.py

@@ -0,0 +1,74 @@
+"""Approximate-global shuffle for chunk-aligned data.
+
+True global shuffle is incompatible with chunk-aligned, low-copy reads: it would
+demand a random chunk per sample. The compromise (DESIGN.md, "shuffle"), adapted
+from MosaicML Streaming's shuffle-block algorithms (py1e / py1br), is two-level:
+
+  1. **Chunk permutation** -- shuffle the *order chunks are scheduled* each epoch.
+  2. **Shuffle-block buffer** -- hold samples from a window of B chunks and draw
+     batches across the whole window, so samples from different chunks interleave.
+
+Setting the block span B >= ~10x the samples-per-chunk yields shuffle quality
+close to global, at memory cost O(B chunks). B is the single quality<->memory
+knob. This module owns the *index math*; buffer.py owns the residency.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def chunk_permutation(chunk_ids: np.ndarray, *, seed: int, epoch: int) -> np.ndarray:
+    """Deterministically permute chunk ids for one epoch.
+
+    Determinism is keyed on (seed, epoch) only -- not on world size or worker
+    count -- so a run is reproducible and resumable across hardware (the
+    "canonical" property from MosaicML).
+    """
+    rng = np.random.default_rng((seed, epoch))
+    return rng.permutation(chunk_ids)
+
+
+def block_shuffled_order(
+    chunk_ids: np.ndarray,
+    samples_per_chunk: int,
+    *,
+    block_chunks: int,
+    seed: int,
+    epoch: int,
+) -> np.ndarray:
+    """Produce a globally-ordered list of (chunk_id, within) draws.
+
+    Emulates the shuffle-block draw order the live buffer will realise, useful
+    for the quality harness and for deterministic single-process iteration.
+    Returns an array of shape ``(n_samples, 2)`` of ``[chunk_id, within]`` rows.
+    """
+    perm = chunk_permutation(chunk_ids, seed=seed, epoch=epoch)
+    rng = np.random.default_rng((seed, epoch, 7919))
+
+    rows: list[np.ndarray] = []
+    for start in range(0, len(perm), block_chunks):
+        block = perm[start : start + block_chunks]
+        # Materialise every (chunk, within) pair in this block, then shuffle.
+        cc = np.repeat(block, samples_per_chunk)
+        ww = np.tile(np.arange(samples_per_chunk), len(block))
+        pairs = np.stack([cc, ww], axis=1)
+        rng.shuffle(pairs)  # in-place, along axis 0
+        rows.append(pairs)
+    return np.concatenate(rows, axis=0)
+
+
+def shuffle_quality(order: np.ndarray, samples_per_chunk: int) -> float:
+    """A 0..1 score for how well an emitted order mixes the source.
+
+    Heuristic: the mean absolute *source-rank* gap between consecutive emitted
+    samples, normalised by the gap a perfect global shuffle would give. 1.0 ~=
+    global; values near 0 mean adjacent samples still come out near each other
+    (poor mixing). Cheap to compute, good enough to tune ``block_chunks``.
+    """
+    source_rank = order[:, 0] * samples_per_chunk + order[:, 1]
+    gaps = np.abs(np.diff(source_rank.astype(np.int64)))
+    n = len(source_rank)
+    # Expected mean gap of a uniform random permutation of 0..n-1 is ~n/3.
+    expected = n / 3.0
+    return float(min(gaps.mean() / expected, 1.0)) if n > 1 and expected else 0.0

insitubatch/source.py

@@ -0,0 +1,118 @@
+"""Torch handoff surface.
+
+Ties the pieces together and exposes them to PyTorch *without* using the classic
+DataLoader worker model. Parallelism lives in :class:`AsyncChunkReader`'s event
+loop, so the recommended configuration is::
+
+    loader = DataLoader(InSituDataset(...), batch_size=None, num_workers=0)
+
+``batch_size=None`` because the dataset already yields assembled batches;
+``num_workers=0`` because forking workers would re-introduce exactly the
+redundant-read / nested-parallelism problems we set out to avoid.
+
+torch (and torchdata.nodes) are optional imports so the core engine stays
+framework-agnostic and importable on a box without torch installed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+import numpy as np
+
+from .buffer import BufferConfig, ShuffleBlockBuffer
+from .io import AsyncChunkReader, IOConfig
+from .plan import build_read_plan
+from .shuffle import block_shuffled_order
+from .split import SplitManifest
+from .store import open_geometries
+from .types import ArrayGeometry, Batch, SplitName
+
+try:  # optional torch surface
+    from torch.utils.data import IterableDataset
+
+    _HAS_TORCH = True
+except ImportError:  # pragma: no cover - exercised on torch-less installs
+    IterableDataset = object  # type: ignore[assignment,misc]
+    _HAS_TORCH = False
+
+
+class InSituDataset(IterableDataset):  # type: ignore[misc]
+    """An IterableDataset that streams shuffled batches from a Zarr archive.
+
+    One epoch = permute the split's chunks -> walk shuffle-blocks -> for each
+    block, async-fetch its chunks, fill the buffer, emit coalesced batches.
+    """
+
+    def __init__(
+        self,
+        store_url: str,
+        manifest: SplitManifest,
+        geometries: dict[str, ArrayGeometry] | None = None,
+        split: SplitName = SplitName.TRAIN,
+        *,
+        batch_size: int = 32,
+        block_chunks: int = 16,
+        max_inflight: int = 16,
+        seed: int = 0,
+        to_tensor: bool = True,
+        **store_kwargs: object,
+    ) -> None:
+        self.store_url = store_url
+        self.store_kwargs = store_kwargs
+        self.geometries = geometries if geometries is not None else open_geometries(store_url)
+        self.manifest = manifest
+        self.split = split
+        self.variables = list(self.geometries)
+        self.io_config = IOConfig(max_inflight=max_inflight)
+        self.buffer_config = BufferConfig(block_chunks=block_chunks, batch_size=batch_size)
+        self.seed = seed
+        self.to_tensor = to_tensor and _HAS_TORCH
+        self._epoch = 0
+
+    def set_epoch(self, epoch: int) -> None:
+        """Call from the training loop so each epoch reshuffles deterministically."""
+        self._epoch = epoch
+
+    def __iter__(self) -> Iterator[Batch | dict]:
+        geom = self.geometries[self.variables[0]]
+        chunk_ids = np.asarray(self.manifest.chunks[self.split.value], dtype=np.int64)
+        spc = geom.sample_chunk_size
+
+        order = block_shuffled_order(
+            chunk_ids,
+            spc,
+            block_chunks=self.buffer_config.block_chunks,
+            seed=self.seed,
+            epoch=self._epoch,
+        )
+
+        with AsyncChunkReader(
+            self.store_url, self.geometries, self.io_config, **self.store_kwargs
+        ) as reader:
+            buf = ShuffleBlockBuffer(self.buffer_config, seed=self.seed)
+            bs = self.buffer_config.batch_size
+
+            # Walk the emitted draw order in batch-sized windows. For each window
+            # we ensure its chunks are resident (async fan-out), then gather.
+            for start in range(0, len(order), bs):
+                rows = order[start : start + bs]
+                needed = {(v, int(c)) for v in self.variables for c in np.unique(rows[:, 0])}
+                missing_samples = [
+                    int(c) * spc for (v, c) in needed if (v, c) not in buf._chunks
+                ]
+                if missing_samples:
+                    plan = build_read_plan(sorted(set(missing_samples)), self.geometries)
+                    for decoded in reader.read_plan(plan):
+                        buf.add(decoded)
+
+                batch = buf.gather_batch(rows, self.variables, spc)
+                yield self._maybe_tensor(batch)
+                buf.evict_drained(needed)
+
+    def _maybe_tensor(self, batch: Batch) -> Batch | dict:
+        if not self.to_tensor:
+            return batch
+        import torch
+
+        return {k: torch.from_numpy(v) for k, v in batch.arrays.items()}

insitubatch/split.py

@@ -0,0 +1,97 @@
+"""Chunk-aligned train/val/test splits.
+
+Splits are done *ahead of time* and at *chunk granularity* along the sample
+axis. Two reasons (DESIGN.md, "splits"):
+
+  1. Leakage: splitting mid-chunk would scatter temporally adjacent, highly
+     autocorrelated samples across train and val. Chunk-aligned boundaries keep
+     a contiguous block of time in a single split.
+  2. Zero-copy: a split that respects chunk boundaries means every read serves
+     exactly one split, so the engine never decodes a chunk and throws half of
+     it away.
+
+The manifest is a plain, serializable record of which chunk indices belong to
+which split, so a run is reproducible and shareable.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import numpy as np
+
+from .types import ArrayGeometry, SplitName
+
+
+@dataclass(slots=True)
+class SplitManifest:
+    """Which sample-axis chunk indices belong to each split."""
+
+    n_chunks: int
+    sample_chunk_size: int
+    n_samples: int
+    chunks: dict[str, list[int]]  # SplitName.value -> sorted chunk indices
+    seed: int
+
+    def sample_indices(self, split: SplitName, geom: ArrayGeometry) -> np.ndarray:
+        """Expand a split's chunks into the global sample indices they contain."""
+        out: list[int] = []
+        for c in self.chunks[split.value]:
+            out.extend(geom.samples_in_chunk(c))
+        return np.asarray(out, dtype=np.int64)
+
+    def to_json(self, path: str | Path) -> None:
+        Path(path).write_text(json.dumps(asdict(self), indent=2))
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> SplitManifest:
+        return cls(**json.loads(Path(path).read_text()))
+
+
+def split_by_chunk(
+    geom: ArrayGeometry,
+    *,
+    fractions: tuple[float, float, float] = (0.8, 0.1, 0.1),
+    seed: int = 0,
+    contiguous: bool = True,
+) -> SplitManifest:
+    """Partition a variable's sample-axis chunks into train/val/test.
+
+    Parameters
+    ----------
+    fractions:
+        (train, val, test) fractions of *chunks* (not samples). Must sum to ~1.
+    contiguous:
+        If True (default), assign contiguous blocks of chunks to each split --
+        the safest choice for time series, where a randomly interleaved split
+        still risks leakage through autocorrelation across chunk boundaries. If
+        False, chunks are shuffled before partitioning (acceptable when samples
+        are exchangeable, e.g. independent scenes).
+    """
+    if abs(sum(fractions) - 1.0) > 1e-6:
+        raise ValueError(f"fractions must sum to 1.0, got {fractions} -> {sum(fractions)}")
+
+    n = geom.n_chunks
+    order = np.arange(n)
+    if not contiguous:
+        order = np.random.default_rng(seed).permutation(n)
+
+    n_train = int(round(fractions[0] * n))
+    n_val = int(round(fractions[1] * n))
+    train = sorted(order[:n_train].tolist())
+    val = sorted(order[n_train : n_train + n_val].tolist())
+    test = sorted(order[n_train + n_val :].tolist())
+
+    return SplitManifest(
+        n_chunks=n,
+        sample_chunk_size=geom.sample_chunk_size,
+        n_samples=geom.n_samples,
+        chunks={
+            SplitName.TRAIN.value: train,
+            SplitName.VAL.value: val,
+            SplitName.TEST.value: test,
+        },
+        seed=seed,
+    )

insitubatch/store.py

@@ -0,0 +1,75 @@
+"""Storage shim: one URL, any backend.
+
+The whole local-now / cloud-later story is a single function. ``obstore`` already
+dispatches on URL scheme (``file://``, ``s3://``, ``gs://``, ``az://``,
+``memory://``) and ``zarr.storage.ObjectStore`` wraps it for the async zarr path.
+So Phase 0 (local ``file://``) and Phase 1 (``s3://...``) differ only in the URL --
+no hot-path code change, and the read path stays pure Rust (no fsspec layer).
+
+We deliberately do *not* route through fsspec / universal_pathlib on the read
+hot path: the entire thesis is that obstore wins by bypassing the fsspec/s3fs
+Python layer. (obstore.fsspec exists if path-style ergonomics are ever wanted
+off the hot path -- but not here.)
+"""
+
+from __future__ import annotations
+
+import os
+from urllib.parse import urlparse
+
+import numpy as np
+import obstore
+import zarr
+import zarr.storage
+
+from .types import ArrayGeometry
+
+
+def store_from_url(
+    url: str, *, read_only: bool = True, **kwargs: object
+) -> zarr.storage.ObjectStore:
+    """Return a zarr ObjectStore for ``url`` (any obstore-supported scheme).
+
+    ``file:///abs/path.zarr`` for local; ``s3://bucket/path.zarr`` for cloud.
+    Extra ``kwargs`` pass through to ``obstore.store.from_url`` (region,
+    credentials, client options, ...).
+    """
+    obs = obstore.store.from_url(url, **kwargs)
+    return zarr.storage.ObjectStore(obs, read_only=read_only)
+
+
+def ensure_local_dir(url: str) -> str:
+    """For a ``file://`` URL, create the target directory so writes can land.
+
+    obstore's LocalStore will not create the prefix for you. No-op for non-file
+    schemes. Returns the URL unchanged for chaining.
+    """
+    parsed = urlparse(url)
+    if parsed.scheme in ("", "file"):
+        os.makedirs(parsed.path, exist_ok=True)
+    return url
+
+
+def open_geometries(
+    url: str,
+    variables: list[str] | None = None,
+    **kwargs: object,
+) -> dict[str, ArrayGeometry]:
+    """Introspect a zarr group at ``url`` into ``{name: ArrayGeometry}``.
+
+    Lets ``InSituDataset`` be built from a URL alone -- geometry (shape, chunks,
+    dtype) is read from the array metadata rather than hand-specified.
+    """
+    store = store_from_url(url, **kwargs)
+    group = zarr.open_group(store=store, mode="r")
+    names = variables if variables is not None else [k for k, _ in group.arrays()]
+    out: dict[str, ArrayGeometry] = {}
+    for name in names:
+        arr = group[name]
+        out[name] = ArrayGeometry(
+            name=name,
+            shape=tuple(arr.shape),
+            chunks=tuple(arr.chunks),
+            dtype=np.dtype(arr.dtype),
+        )
+    return out

insitubatch/types.py

@@ -0,0 +1,108 @@
+"""Core data types shared across the insitubatch engine.
+
+The central design choice (see DESIGN.md): the unit of work is neither the
+*sample* nor the *chunk* in isolation, but a **read plan** that maps the samples
+required for a step to a *deduplicated* set of chunk reads. This lets the same
+scheduler serve the whole spectrum from fat chunks (many samples per chunk,
+shared-cache wins) to the degenerate GRIB-per-timestep case (one sample per
+chunk, async fan-out is everything).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+
+import numpy as np
+
+
+class SplitName(StrEnum):
+    TRAIN = "train"
+    VAL = "val"
+    TEST = "test"
+
+
+@dataclass(frozen=True, slots=True)
+class ArrayGeometry:
+    """The minimal geometry the engine needs about one zarr array.
+
+    We only model the *sample axis* (the outer dimension, axis 0 by convention:
+    time for ERA5/HRRR) explicitly, because that is the axis we split, shuffle,
+    and batch along. The trailing dims are carried opaquely as ``inner_shape``
+    and are kept contiguous to preserve partial zero-copy.
+    """
+
+    name: str
+    shape: tuple[int, ...]
+    chunks: tuple[int, ...]
+    dtype: np.dtype
+
+    @property
+    def n_samples(self) -> int:
+        """Length of the sample (outer) axis."""
+        return self.shape[0]
+
+    @property
+    def sample_chunk_size(self) -> int:
+        """How many samples live in one chunk along the sample axis."""
+        return self.chunks[0]
+
+    @property
+    def inner_shape(self) -> tuple[int, ...]:
+        """Shape of a single sample (everything past the sample axis)."""
+        return self.shape[1:]
+
+    @property
+    def n_chunks(self) -> int:
+        """Number of chunks along the sample axis."""
+        return -(-self.n_samples // self.sample_chunk_size)  # ceil div
+
+    def chunk_of(self, sample_index: int) -> int:
+        """Which sample-axis chunk a given sample index falls in."""
+        return sample_index // self.sample_chunk_size
+
+    def samples_in_chunk(self, chunk_index: int) -> range:
+        """The half-open range of global sample indices in ``chunk_index``."""
+        start = chunk_index * self.sample_chunk_size
+        stop = min(start + self.sample_chunk_size, self.n_samples)
+        return range(start, stop)
+
+
+@dataclass(frozen=True, slots=True)
+class ChunkRead:
+    """A single chunk to fetch, addressed along the sample axis.
+
+    ``array`` names which zarr array (variable) this read belongs to; a training
+    sample that concatenates several variables produces one ``ChunkRead`` per
+    variable that must be co-scheduled.
+    """
+
+    array: str
+    chunk_index: int
+
+
+@dataclass(slots=True)
+class DecodedChunk:
+    """A decoded, in-memory chunk, keyed by its read.
+
+    ``data`` has shape ``(n_samples_in_chunk, *inner_shape)``. The buffer holds a
+    bounded number of these; memory overhead is O(in-flight chunks), independent
+    of batch size.
+    """
+
+    read: ChunkRead
+    data: np.ndarray
+    sample_offset: int  # global sample index of data[0]
+
+
+@dataclass(slots=True)
+class Batch:
+    """A model-ready batch.
+
+    ``arrays`` maps variable name -> stacked array of shape ``(batch, *inner)``.
+    ``sample_indices`` records provenance (which global samples, in order) for
+    determinism checks and resumption.
+    """
+
+    arrays: dict[str, np.ndarray]
+    sample_indices: np.ndarray = field(default_factory=lambda: np.empty(0, dtype=np.int64))

insitubatch-0.0.1.dist-info/METADATA

@@ -0,0 +1,90 @@
+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).

insitubatch-0.0.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+insitubatch/__init__.py,sha256=9TTgS8anuSDAUm3QugS8EJgR6nIOCQpeEVFXUQNfZWo,1311
+insitubatch/buffer.py,sha256=JcjzcFjqIvaqWS0jp6_1bvPlLgXUXRUEOP9F9HLWAS4,3419
+insitubatch/io.py,sha256=kSXid_h1bu3tNoO3BbvlD1MDF5RlXh4eFtNI0fz0HyM,7090
+insitubatch/plan.py,sha256=cY3iYemKLmJvaHkTD1ahOU7WsfjJzG9z6wuxbxu9kOQ,4110
+insitubatch/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+insitubatch/shuffle.py,sha256=ARXoC1sAUU_xRa3YBN4adtKszSay_HqOTtBlvk1ik7Q,3136
+insitubatch/source.py,sha256=qhIcLj3HAtDePGiKHBWqyi0n9lgKV6-KlfyvKdrRAh4,4486
+insitubatch/split.py,sha256=xIQuzZmIuON4OANAEqVKHkdCc1BETYluv28HsztOTOg,3286
+insitubatch/store.py,sha256=Ny-DrEoQD-cv4y9NIkwIdzk4_QBrXduZyN7RgkNYIX8,2605
+insitubatch/types.py,sha256=6rquZxacOgekl-GKw37dfb3HOxD8muyHKRwaOqMBdQA,3455
+insitubatch-0.0.1.dist-info/licenses/LICENSE,sha256=6IPhQ_YPQ0SYyy6JY62NW7zhfJhoYj8mfYeb0LKMwJU,1069
+insitubatch-0.0.1.dist-info/WHEEL,sha256=8ZlpUMJ7mlDirmlHRhDirEx_nPnARrwDjeE92mlk68E,81
+insitubatch-0.0.1.dist-info/METADATA,sha256=TD_A5rHYJqcPWRezqOwbG9xmS8teVvYG9pBOwFliNYA,3365
+insitubatch-0.0.1.dist-info/RECORD,,

insitubatch-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.11.21
+Root-Is-Purelib: true
+Tag: py3-none-any

insitubatch-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 David Stuebe
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.