API reference¶
The public surface is the top-level insitubatch package — everything in its
__all__. InSituDataset and the framework adapters are re-exported there, so
import them from the package root (from insitubatch import InSituDataset, to_torch),
not from submodules. The adapters are optional: they import torch / JAX / TF lazily,
only when called, so importing insitubatch never pulls a framework in.
insitubatch¶
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.
ChunkPool ¶
Byte-budgeted pool of outer-chunk slots, keyed (array, chunk_index).
The pool is the assembly buffer and the cache. A slot is pinned while the
current epoch needs it (in-flight or block-not-yet-drained) and unpinned once
its block is drained;
unpinned slots stay resident (retained for cross-epoch reuse) until budget
pressure evicts them in LRU order. budget_bytes is the single knob:
- small (~
2*block_chunksworth) -> read-once (unpinned evicted promptly); - large + persistent across epochs -> a decode-once cache (a still-resident prepped chunk is a hit, skipping fetch + decode + transform).
Eviction targets unpinned-LRU only; a slot is never unpinned before it is
ready+drained, so an in-flight or in-use chunk is never dropped. Backing is heap
or mmap (see backing_dir); chunk_transforms run once per outer chunk on
the assembled array, so a hit reflects decode + transform.
try_admit ¶
Reserve + allocate + reference one outer-chunk slot, evicting ready-LRU for room.
Admission takes one reference (incref) and claims the slot for this epoch (so
the consumer's :meth:wait_ready won't gather it until the driver has referenced
it -- see there), so the slot stays resident from its in-flight fetch through to
the consumer's release. The driver fetches each chunk once per epoch, so an
eviction before consume could not be re-fetched and would deadlock the waiter.
The consumer releases each chunk at its last use (:meth:unpin_keys); windowed
reads let one chunk be referenced by several blocks, hence reference counts, not a
boolean. Idempotent if already resident (incref only). Returns False only when
the budget is full of in-flight or referenced slots -- the caller awaits a release.
is_ready ¶
True if the chunk is resident, fully assembled, and not failed (a hit).
pin_if_ready ¶
Incref + return True iff the chunk is resident, ready, and not failed.
A cross-epoch (or, with persist, cross-run) cache hit the driver can skip
fetching -- but it must still be referenced so it stays resident through the
consumer's use (released at last use like an admitted chunk), else it could be
evicted before the waiter gathers it and, since the driver fetches each chunk
once, deadlock. One lock so the check and the incref cannot race an eviction in
between. A persisted-on-disk chunk is revived here on first touch (see
:meth:_revive), so a cross-run hit costs no fetch.
pin_keys ¶
Reference (incref) a set of (path, chunk_index) slots for a live block.
Windows make one chunk readable by several concurrent blocks, so pins are
reference-counted: each block that needs a slot increfs it on entry and
decrefs it on drain (:meth:unpin_keys). A slot with refcount > 0 is never
evicted. Pinning a not-yet-allocated key is fine -- the count is recorded and
the slot, once admitted, inherits it.
unpin_keys ¶
Release (decref) a block's (path, chunk_index) references.
A slot dropping to refcount 0 becomes LRU-evictable (retained for cross-epoch reuse until budget pressure drops it), not dropped now. Wakes any admit parked on a full budget.
unpin_all ¶
Epoch boundary reset: clear every pin and drop abandoned partials.
A pin is per-epoch working state, not cache membership -- ready chunks stay
resident (unpinned) for cross-epoch reuse. No pin may survive an epoch: an
aborted epoch (early break) leaves its read-ahead and un-drained current
block referenced, which would shrink the next epoch's budget until admission
can free no room and the driver deadlocks. A not-ready slot at this boundary
is an abandoned partial (its fetch was cancelled mid-flight) -- it can never be
a valid cache entry, so drop it; that also restores the in-epoch invariant
"not ready => in flight" that protects fetches from eviction. Called at the
next epoch's start, when the prior scheduler is fully closed (no race).
scatter ¶
Copy one decoded tile into its slot; complete the chunk if it was the last.
The memcpy happens before the lock (rule 1); the completion counter and
ready flip happen under the lock (rule 2). Completion (chunk
transforms on the assembled array) runs outside the lock -- no other thread
touches the slot once remaining hits 0 -- then ready is published.
fail ¶
Mark a slot failed so a waiting consumer re-raises instead of hanging.
Fail-fast: a fetch/decode error on any tile poisons its outer chunk; the
consumer's wait_ready surfaces it on the main thread.
set_error ¶
Poison the whole pool (the fetch driver died) so every waiter re-raises.
Unlike :meth:fail (one chunk), this unblocks consumers waiting on chunks
that may never be allocated -- the driver failed before reaching them. The
first error wins (later failures are usually cascade noise).
wait_ready ¶
Block until the outer chunk is assembled and claimed this epoch (or raise).
Waiting on claimed (set by the driver's admit / pin_if_ready) closes a
cross-epoch race: a chunk still resident-and-ready from the prior epoch would
otherwise be gathered before the driver references it, letting the consumer's
last-use release land before the driver's pin -- a lost release that leaks a
reference (and, worse, lets the driver evict a chunk mid-gather). Requiring the
claim orders pin-before-consume-before-release. Wakes on: ready+claimed, the
chunk failed (:meth:fail), or the pool was poisoned (:meth:set_error, covering
a driver death before this chunk was allocated).
gather ¶
Assemble one batch from [chunk_id, within] anchor draw rows.
Each row is one sample anchor t = chunk_id*ref_spc + within in the reference
(manifest) grid; each variable reads its array at t + offset (offset 0 is the
plain non-windowed case). Output is in anchor-row order: row i of every variable
is the same anchor, sample_indices[i] == t_i. Per variable the reads are grouped
by the variable's own (offset-shifted) chunk -- computed with that variable's
chunk size, so variables may chunk the sample axis differently -- one coalesced
fancy-index per chunk, never a Python per-sample loop. The caller must have waited
every referenced (path, offset-shifted chunk) ready.
close ¶
Free every remaining slot and release the log handle. Persist keeps ready cache files (each already recorded in the log at completion, so there is nothing to rewrite -- just flush + close the handle); heap/spill mmap files are unlinked. Idempotent.
Scheduler ¶
Owns one event loop + a decode pool; streams tiles into a caller-owned pool.
The :class:ChunkPool is passed in (dataset-owned, so it persists across epochs
as the cache). :meth:start streams the stored chunks of an ordered chunk list;
the consumer reads assembled chunks via :attr:pool and releases drained ones
via :meth:unpin. Per chunk the scheduler skips fetch if the pool already holds
it (a cross-epoch hit); misses are admitted against the pool's byte budget,
awaiting an unpin when the working set fills it.
close ¶
Cancel any in-flight driver, then stop the loop and decode pool.
Graceful: a consumer may close mid-epoch (early break) while _drive
is still streaming. We cancel outstanding tasks and let them unwind before
stopping the loop, so no coroutine is orphaned (which would surface as
GeneratorExit / "never awaited" warnings on GC).
start ¶
Begin streaming the stored chunks of chunk_ids (priority order).
chunk_ids are in the reference (manifest) grid; ref_spc is that grid's
sample-chunk size, used to map anchor chunks onto each variable's own chunks.
Returns the driver future; a failure there poisons the pool so consumers
re-raise. The consumer drives demand independently via :attr:pool.
unpin_block ¶
Release references on a set of drained (path, chunk_index) slots
(thread-safe): the slots that hit refcount 0 become LRU-evictable; wake any
admit parked on a full budget so it can evict them and proceed.
SchedulerConfig
dataclass
¶
max_inflight
class-attribute
instance-attribute
¶
Tiles in flight at once -- the single concurrency dial. Memory in flight ~= max_inflight * stored_chunk_nbytes (+ transform scratch). Residency is bounded separately by the pool's byte budget (admission evicts unpinned-LRU).
decode_threads
class-attribute
instance-attribute
¶
Size of the decode/scatter pool (GIL-releasing codec decode + the disjoint
scatter memcpy run here). 0 = auto = min(32, cpu+4).
on_bad_chunk
class-attribute
instance-attribute
¶
What to do when a stored chunk fails to fetch/decode (truncated/corrupt --
common in GRIB-under-zarr archives like HRRR). "raise" (default) fails fast;
"nan" fills that tile with NaN (float dtypes) or the fill value, so the chunk
assembles with a hole instead of poisoning the epoch -- the caller then handles
NaN with a chunk_transform (interpolate / drop). Bad reads are recorded in
Scheduler.bad_chunks.
InSituDataset ¶
A framework-neutral source of shuffled numpy batches from Zarr, split-aware.
The dataset is not itself iterated -- you iterate one of its split views:
:attr:train (shuffled), :attr:val, :attr:test, :attr:all (deterministic).
All four share one :class:ChunkPool, so a chunk that two splits both read --
e.g. a windowed read spilling across a split boundary -- is decoded once::
ds = InSituDataset(store, manifest, geometries=geoms, batch_size=32)
for batch in ds.train: ... # one epoch; ds.set_epoch(e) reshuffles
for batch in ds.val: ...
One epoch over a view = permute the split's chunks -> walk shuffle-blocks -> per
block, stream-fetch its stored chunks into the pool, gather coalesced batches, evict.
Batches are numpy :class:Batch; convert to a framework with
:mod:insitubatch.frameworks (as_torch / to_jax / as_tf_dataset). A
different per-split configuration (e.g. train-only augmentation) is a separate dataset.
Two preprocessing hooks, placed by cost (full model in the docs, "Transforms"):
chunk_transforms--(DecodedChunk) -> DecodedChunk, run per chunk before shuffle, seeing one variable. The cacheable home for elementwise, per-variable, deterministic work (scaling, unit conversion, dtype cast); amortized over every sample in the chunk and reused across epochs.batch_transforms--(Batch) -> Batch, run per assembled batch, seeing all variables aligned on the sample axis. For cross-variable derived fields and per-sample random augmentation; runs after the cache, so it is never cached.
Runnable side-by-side example: examples/transforms.py.
all
property
¶
Iterable over every split's chunks (deterministic) -- e.g. full-archive inference.
set_epoch ¶
Call from the training loop so each epoch reshuffles deterministically.
close ¶
Release the cache pool's backing (mmap handles, cached chunks) and any async store session.
The pool persists across epochs, so close it when done training -- not per
epoch. With persist=True the cache files + manifest are kept on disk for a
future run (only the in-memory handles are released); otherwise the mmap spill
files are unlinked. An fsspec/gcsfs store's aiohttp session is closed on its own
loop here (a no-op for obstore) so it does not leak or spew a teardown traceback
at GC; gcsfs recreates it lazily if the store is reused. Idempotent; also called
on GC.
SplitManifest
dataclass
¶
Which sample-axis chunk indices belong to each split.
sample_indices ¶
Expand a split's chunks into the global sample indices they contain.
BatchTransform ¶
Bases: Protocol
Per-batch transform applied after gather (not cached).
ChunkTransform ¶
Bases: Protocol
Per-chunk transform applied before shuffle/gather (cacheable).
StandardScaler
dataclass
¶
Global per-variable standardization with PRE-FIT, FIXED statistics.
mean/std are keyed by variable and shaped to broadcast over a chunk's
(n_samples, *inner) array WITHOUT the sample axis: a surface variable uses
shape (1, 1); per-level stats use (level, 1, 1). The same stats are
applied to every chunk of that variable -- never recomputed per chunk.
Pre-fit the stats however you like and pass them in. The recommended path is to
fit over the loader with sklearn's incremental StandardScaler.partial_fit
(which also warms the cache) and scale at the batch stage -- see
examples/fit_scaler.py; this class is the chunk-stage applier for when you
want the normalization cached with the decoded chunk.
ArrayGeometry
dataclass
¶
The minimal geometry the engine needs about one zarr array.
We only model the sample axis explicitly, because that is the axis we split,
shuffle, and batch along; the remaining dims are carried opaquely as
inner_shape and kept contiguous to preserve partial zero-copy. shape and
chunks are in physical (zarr) axis order -- they mirror the array's own
metadata -- and sample_axis names which physical axis is the sample axis
(0 by convention: time for ERA5/HRRR; e.g. 2 for the Z of an OME-NGFF
(T,C,Z,Y,X) microscopy stack sampled slice-by-slice). The engine works in a
logical view where the sample axis leads and the inner axes follow in physical
order; the one physical<->logical permutation is confined to the scheduler
(:meth:physical_chunk_coord for read addressing; a moveaxis on the decoded
tile). Everything downstream -- planning, pooling, gather -- is sample-first.
offset makes a variable a windowed view: it reads array[anchor + offset]
along the sample axis around a shared anchor. Two geometries with the same path
and different offset (e.g. g and g.shift(1)) are two views of one array --
they decode once and share slots. Offset 0 is not special; everything is relative to
the anchor.
sample_chunk_size
property
¶
How many samples live in one chunk along the sample axis.
inner_shape
property
¶
Shape of a single sample (every axis but the sample axis, physical order).
inner_chunks
property
¶
Stored-chunk shape on the inner (non-sample) axes (physical order).
shift ¶
A view of the same array read k samples later (composes: shift(1).shift(1)
is offset += 2). Declare a forecast target as g.shift(horizon).
physical_chunk_coord ¶
Full physical zarr stored-chunk coordinate for a logical read.
Reinsert the sample-axis chunk index chunk_index at sample_axis among the
inner-axis chunk coords (which are in physical inner order). With sample_axis == 0
this is exactly (chunk_index, *inner_coord) -- the identity the old code assumed.
samples_in_chunk ¶
The half-open range of global sample indices in chunk_index.
n_inner_chunks ¶
How many stored tiles compose one outer chunk (the inner-grid size).
Independent of chunk_index (a short final outer chunk is still one
axis-0 stored chunk), but kept index-keyed so the pool's completion count
reads naturally and the API survives a future per-axis sample chunking.
slot_shape ¶
Shape of the assembled outer chunk: (n_samples_in_chunk, *inner_shape).
Axis 0 uses the actual sample count so the final short chunk is sized exactly (no over-allocation, no out-of-range scatter).
tile_placement ¶
(dst, src) slices for scattering one decoded tile into its slot.
dst indexes the outer-chunk slot; src clips the full chunk-shaped
decoded tile to the (possibly partial) edge region -- both axis 0 (short
final outer chunk) and the inner edges. After the copy the tile is free.
Batch
dataclass
¶
A model-ready batch.
arrays maps variable label -> stacked array of shape (batch, *inner).
sample_indices is the per-row anchor sample index t (provenance for
determinism / resumption). offsets maps each label to its sample-axis read
offset, so label v of row i was read from global sample sample_indices[i]
+ offsets[v]. A plain (non-windowed) batch has every offset 0; a forecast batch
pairs e.g. an input at offset 0 with a target at offset horizon.
The batch stays a flat {label: array} dict -- there is no lead/role axis in the
engine. Use :meth:stack to assemble a multi-step window into one array and
:meth:read_indices for a label's true provenance.
read_indices ¶
Global sample index each row of label was read from: anchor + offset.
Provenance for a windowed view (e.g. to confirm a target leads its input by the intended horizon). Defaults the offset to 0 for a label without one recorded.
stack ¶
Stack several labels into one array along a new axis (default 1).
The obvious way to build a multi-step input window from a set of time-shifted
views, e.g. batch.stack(["t_m2", "t_m1", "t_0"]) -> (batch, 3, *inner).
Order follows labels; row i of every label shares anchor
sample_indices[i], so the stacked steps stay aligned. The caller chooses the
labels and their order -- the engine does not impose a window layout.
ChunkRead
dataclass
¶
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.
DecodedChunk
dataclass
¶
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.
StoredChunkRead
dataclass
¶
One stored chunk to fetch: a single tile of the chunk grid.
Reading a whole outer chunk per getitem lets zarr stitch the inner grid
under a second concurrency cap; fetching at stored-chunk granularity instead
-- (chunk_index, *inner_coord) -- lets a single max_inflight budget
span inner and outer reads, with no nested caps. chunk_index is the
sample-axis (outer) stored-chunk index; inner_coord is the stored-chunk
index on each inner axis (empty tuple when the inner dims are single-chunk --
the degenerate GRIB-per-timestep case).
Frozen + hashable so a plan can dedup tiles and key the in-flight set.
as_tf_dataset ¶
Wrap a split view (e.g. ds.val) as a tf.data.Dataset via from_generator.
output_signature is inferred from the view's geometries: each variable is
(None, *inner) (None = the variable last-batch size) with the variable's dtype.
Both from_generator here and :func:to_tf copy into the TF runtime -- TF has no
reliable zero-copy path from insitu's buffers (its experimental DLPack mishandles buffer
ownership; see :func:to_tf). Call :func:to_tf on the raw stream when you want plain
dict[str, tf.Tensor] batches instead of a tf.data.Dataset.
as_torch ¶
Wrap a split view (e.g. ds.train) as a torch IterableDataset for DataLoader.
Each yielded item is a dict[str, torch.Tensor] (via :func:to_torch). Use
DataLoader(as_torch(ds.train), batch_size=None, num_workers=0).
to_tf ¶
Convert a numpy Batch to a dict of tf.Tensor (one CPU copy per variable).
Unlike torch/JAX -- whose array-accepting from_dlpack manages the exported buffer's
lifetime correctly -- TensorFlow only exposes the experimental from_dlpack(capsule),
which mishandles ownership of the exported numpy buffer: under the concurrent allocation
of the prefetch decode threads it double-frees that buffer and aborts the process
(SIGABRT, no message). convert_to_tensor copies into a TF-owned tensor instead, so TF
never touches insitu-managed memory; the batch is already an owned array, so this is a
single CPU copy. (torch/JAX stay zero-copy; this is a TF-DLPack limitation, not ours.)
to_torch ¶
Convert a numpy Batch to a dict of torch tensors (DLPack; zero-copy on CPU).
build_stored_chunk_reads ¶
Expand outer chunk ids into deduped stored-chunk reads, in priority order.
There is no gather map: the scheduler scatters tiles into per-outer-chunk slots
in a :class:~insitubatch.pool.ChunkPool, and batches are gathered straight
from those assembled slots by (chunk_id, within) draw rows -- the same
coordinates the shuffle order already produces. So the result is just what to
fetch, in what order; the scheduler keeps max_inflight tiles in flight
across the list.
chunk_ids are outer (sample-axis) chunk indices in draw/priority order
(e.g. the next shuffle-block's chunks first), so the soonest-needed tiles go
first. Each outer chunk expands to its inner grid; every variable contributes
its own grid (variables may chunk the inner dims differently). Order is
chunk -> variable -> inner so a whole outer chunk's tiles are scheduled
together (it can be assembled and drained promptly). Reads are keyed by the
array path (not the dict label), so several windowed views of one array
(same path, different offset) collapse to a single fetch -- decode-once.
Dedup also makes the function safe to call with repeated ids.
chunk_ids are anchor chunks in the reference grid (ref_spc = the
manifest's sample-chunk size, which defines the shuffle/split anchor grid). A windowed
variable reads array[anchor + offset], and a variable that chunks the sample axis
differently from the reference maps those anchor samples onto its own chunks -- so one
anchor chunk expands to the (offset-shifted) chunks each variable needs. With every
offset == 0 and a uniform chunk size this is exactly anchor chunk -> itself.
block_shuffled_order ¶
Produce a shuffle-block-ordered list of [chunk_id, within] draws.
Chunks are permuted per epoch; within each window of block_chunks chunks
all samples are shuffled together. n_samples is the global sample-axis
length, used to size a short final chunk correctly. Returns an array of shape
(N, 2) where N is the number of samples covered by chunk_ids.
chunk_permutation ¶
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).
sequential_order ¶
In-order [chunk_id, within] draws (no permutation, no shuffle).
Used when shuffle=False (eval / inference / reconstruction): chunks in the
given order, samples in order within each. Honours a short final chunk.
shuffle_quality ¶
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.
split_by_chunk ¶
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).
sample_range:
Optional half-open (start, stop) window of sample (outer-axis) indices
to restrict the split to before partitioning -- e.g. train on one date
range of a long archive. The selection is chunk-aligned and contiguous:
every chunk that overlaps [start, stop) is kept whole, so a window
starting or ending mid-chunk pulls in that partial edge chunk (splits are
chunk-granular -- you subset whole chunks, never individual samples). Use it
for a single contiguous window; it is not a tool for scattered/boolean
selections (those would drag in straddling chunks and silently add samples).
valid_anchor_range ¶
Half-open [lo, hi) of anchor sample-indices whose every windowed read
anchor + offset stays in [0, n_samples) -- the anchors a windowed dataset
may draw, with array-edge anchors dropped.
Offsets {-1, 0, 1} over T samples -> anchors [1, T-1). Range is the
only validity the engine enforces; whether the user's offset choices define a
meaningful (non-leaky) task is theirs to decide (DESIGN, M-W). Empty/too-wide
windows return an empty range (lo, lo).
arraylake_store ¶
Open an Arraylake repo and return its read-only Icechunk session store.
Auth comes from a cached al auth login or ARRAYLAKE_TOKEN; the client
vends the bucket credentials for the repo. The returned object is a zarr-v3
Store bound to the branch snapshot -- exactly what the engine accepts.
Requires insitubatch[arraylake].
close_store ¶
Best-effort teardown for a store that holds an async fsspec session (gcsfs, s3fs).
Such a backend creates an aiohttp session on the first event loop that awaits it --
for a zarr store, that is zarr's loop, not fsspec's -- but gcsfs's finalizer captures
fs.loop (which is None here) and closes the session on the wrong loop at GC,
spewing a harmless-looking "Task was destroyed / attached to a different loop"
traceback and leaking the connection. Closing the session here on the loop it actually
lives on makes that finalizer a no-op.
A no-op for stores with no such session (obstore's ObjectStore has no .fs) and
for already-closed or not-running loops. gcsfs recreates the session lazily, so a
store closed here still works if reused -- but call this only when done with it.
ensure_local_dir ¶
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.
fsspec_store ¶
Return a zarr FsspecStore for url (any fsspec-supported backend).
Reaches stores via a backend fsspec filesystem -- notably GCS Rapid/zonal
buckets (gRPC) and GCS requester-pays, which obstore does not currently
support. **storage_options pass straight through to
FsspecStore.from_url (credentials, project, endpoint, Rapid config, ...).
Requires an fsspec backend for the URL scheme: insitubatch[gcsfs] for
gs://, or bring your own (s3fs, ...). A sync backend (e.g. local
file://) is auto-wrapped as async by zarr; gs:// via gcsfs is
natively async. See :func:obstore_store for the obstore-backed constructor.
obstore_store ¶
Return an obstore-backed zarr Store for url (any obstore 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, ...). The read path stays pure Rust -- no fsspec
Python layer.
open_geometries ¶
Introspect a zarr group Store into {name: ArrayGeometry}.
Lets InSituDataset be built from a store alone -- geometry (shape,
chunks, dtype) is read from the array metadata rather than hand-specified.
Build the store with :func:obstore_store / :func:fsspec_store /
:func:arraylake_store, or pass any prebuilt zarr Store.
sample_axis names which physical axis is the outer (sample) axis for
every returned variable -- 0 (default: time for ERA5/HRRR) or, e.g., the
Z of an OME-NGFF (T,C,Z,Y,X) stack sampled slice-by-slice (sample_axis=2).
Variables that need different sample axes are built individually (construct
:class:ArrayGeometry per array); the shape/chunks stay in physical order.
InSituDataset¶
insitubatch.InSituDataset ¶
A framework-neutral source of shuffled numpy batches from Zarr, split-aware.
The dataset is not itself iterated -- you iterate one of its split views:
:attr:train (shuffled), :attr:val, :attr:test, :attr:all (deterministic).
All four share one :class:ChunkPool, so a chunk that two splits both read --
e.g. a windowed read spilling across a split boundary -- is decoded once::
ds = InSituDataset(store, manifest, geometries=geoms, batch_size=32)
for batch in ds.train: ... # one epoch; ds.set_epoch(e) reshuffles
for batch in ds.val: ...
One epoch over a view = permute the split's chunks -> walk shuffle-blocks -> per
block, stream-fetch its stored chunks into the pool, gather coalesced batches, evict.
Batches are numpy :class:Batch; convert to a framework with
:mod:insitubatch.frameworks (as_torch / to_jax / as_tf_dataset). A
different per-split configuration (e.g. train-only augmentation) is a separate dataset.
Two preprocessing hooks, placed by cost (full model in the docs, "Transforms"):
chunk_transforms--(DecodedChunk) -> DecodedChunk, run per chunk before shuffle, seeing one variable. The cacheable home for elementwise, per-variable, deterministic work (scaling, unit conversion, dtype cast); amortized over every sample in the chunk and reused across epochs.batch_transforms--(Batch) -> Batch, run per assembled batch, seeing all variables aligned on the sample axis. For cross-variable derived fields and per-sample random augmentation; runs after the cache, so it is never cached.
Runnable side-by-side example: examples/transforms.py.
all
property
¶
Iterable over every split's chunks (deterministic) -- e.g. full-archive inference.
set_epoch ¶
Call from the training loop so each epoch reshuffles deterministically.
close ¶
Release the cache pool's backing (mmap handles, cached chunks) and any async store session.
The pool persists across epochs, so close it when done training -- not per
epoch. With persist=True the cache files + manifest are kept on disk for a
future run (only the in-memory handles are released); otherwise the mmap spill
files are unlinked. An fsspec/gcsfs store's aiohttp session is closed on its own
loop here (a no-op for obstore) so it does not leak or spew a teardown traceback
at GC; gcsfs recreates it lazily if the store is reused. Idempotent; also called
on GC.
Framework adapters¶
Thin, optional framework adapters: numpy Batch -> torch / JAX / TF via DLPack.
The core (:mod:insitubatch.source) yields numpy :class:Batch objects and imports
no framework. These adapters convert a batch's arrays to a framework's tensors with
DLPack (zero-copy on CPU where the framework supports it). The wrapping differs per
ecosystem -- there is no single cross-framework "dataset" base class:
-
torch has one.
DataLoaderrequires aDataset/IterableDatasetsubclass (itisinstance-checks), so :func:as_torchwraps the stream in one::DataLoader(as_torch(ds), batch_size=None, num_workers=0)batch_size=None(the stream already yields assembled batches) andnum_workers=0(parallelism is in our event loop; forking re-introduces the redundant-read problem). * JAX has none -- it is loader-agnostic. Iterate the dataset and call :func:to_jaxper batch. * TF adapts via a factory, not a base class: :func:as_tf_datasetwraps the stream intf.data.Dataset.from_generator.
Each framework is imported lazily inside its function, so importing this module costs
nothing and a missing framework raises a clear, actionable error. sample_indices
(provenance) stays on the numpy Batch; only the model-input arrays are converted.
to_torch ¶
Convert a numpy Batch to a dict of torch tensors (DLPack; zero-copy on CPU).
as_torch ¶
Wrap a split view (e.g. ds.train) as a torch IterableDataset for DataLoader.
Each yielded item is a dict[str, torch.Tensor] (via :func:to_torch). Use
DataLoader(as_torch(ds.train), batch_size=None, num_workers=0).
to_tf ¶
Convert a numpy Batch to a dict of tf.Tensor (one CPU copy per variable).
Unlike torch/JAX -- whose array-accepting from_dlpack manages the exported buffer's
lifetime correctly -- TensorFlow only exposes the experimental from_dlpack(capsule),
which mishandles ownership of the exported numpy buffer: under the concurrent allocation
of the prefetch decode threads it double-frees that buffer and aborts the process
(SIGABRT, no message). convert_to_tensor copies into a TF-owned tensor instead, so TF
never touches insitu-managed memory; the batch is already an owned array, so this is a
single CPU copy. (torch/JAX stay zero-copy; this is a TF-DLPack limitation, not ours.)
as_tf_dataset ¶
Wrap a split view (e.g. ds.val) as a tf.data.Dataset via from_generator.
output_signature is inferred from the view's geometries: each variable is
(None, *inner) (None = the variable last-batch size) with the variable's dtype.
Both from_generator here and :func:to_tf copy into the TF runtime -- TF has no
reliable zero-copy path from insitu's buffers (its experimental DLPack mishandles buffer
ownership; see :func:to_tf). Call :func:to_tf on the raw stream when you want plain
dict[str, tf.Tensor] batches instead of a tf.data.Dataset.