TreeWriter / TreeDataSource

TreeWriter: pytree-native writer for memory-mapped datasets.

class TreeWriter(directory: str | Path, expected_samples: int, *, on_overflow: Literal['error', 'trim', 'grow'] = 'grow', metadata: Dict[str, Any] | None = None, pbar=None, close_pbar: bool = False, exist_ok: bool = False, manifest_interval: float = 5.0)

Write any pytree to memory-mapped binary files.

Each leaf in the pytree becomes a separate .bin file. The tree structure is stored in manifest.json for reconstruction by TreeDataSource.

Accepts AudioTree objects, dicts of arrays, dicts of AudioTrees, or any combination. The schema is inferred from the first write() call.

Parameters:

  • directory – Directory where memmap files will be written

  • expected_samples – Number of samples to pre-allocate. This is an allocation hint, not a cap: what happens when more samples are offered is decided by on_overflow. Under-shooting it is always fine – close() truncates the files to what was written.

  • on_overflow –

    What to do with a batch that does not fit the current allocation.

    • "grow" (default) reallocates every leaf file to fit and carries on, so expected_samples really is only a hint.

    • "error" raises ValueError, naming how many samples were written, allocated and offered.

    • "trim" writes as much of the batch as fits and drops the rest, warning with the same counts. Once the allocation is full, every further write() warns and returns 0.

  • metadata – Optional dict of user metadata to store in manifest

  • pbar – Optional tqdm progress bar instance. Updated by batch_size after each write() call.

  • close_pbar – If True, close the progress bar when the writer closes. Default False.

  • manifest_interval – Seconds between refreshes of the manifest’s num_samples while writing. The manifest is what tells a reader how much of the dataset is real, so a hard kill must not leave it claiming zero; refreshing on a timer bounds how stale that count can be without rewriting the JSON on every batch. Pass 0 to refresh only on flush() and close().

Example

Pre-render a few batches of one-second mono AudioTree objects into a memory-mapped dataset, then read one back with TreeDataSource:

>>> import tempfile
>>> import jax.numpy as jnp
>>> from audiotree import AudioTree
>>> from audiotree.sources import TreeDataSource
>>> out_dir = tempfile.mkdtemp()
>>> batches = [AudioTree.create(jnp.zeros((8, 1, 16000)), 16000) for _ in range(3)]
>>> with TreeWriter(out_dir, expected_samples=8 * len(batches)) as w:
...     for batch in batches:
...         _ = w.write(batch)
>>> ds = TreeDataSource(out_dir)
>>> len(ds)
24
>>> ds[0].waveform.shape  # one sample, batch dim added back
(1, 1, 16000)
close()

Close all memmap and bagz files and write manifest.

If fewer samples were written than were allocated – because expected_samples overshot, or because on_overflow="grow" reallocated with slack – the memmap files are truncated to the actual sample count so no disk space is wasted and readers see the correct size.

Closing is terminal: the writer cannot be reopened, because its files have been truncated to what was already written and its memmaps released.

flush()

Flush all memmap files to disk and refresh the manifest.

get_stats() → Dict[str, Any]

Get statistics about written data.

open() → TreeWriter

Open the writer and create output directory.

Returns:

self for method chaining

Raises:

RuntimeError – If the writer is already open, or has been closed.

write(pytree) → int

Write a batch of samples to the memmap and bagz files.

The first call infers the schema from the pytree structure. Array leaves must have the same batch size (first dimension). String leaves (str or List[str]) are stored in bagz files.

Parameters:

pytree – Any JAX-compatible pytree (AudioTree, dict, nested). Array leaves must have a batch dimension. String leaves can be a single str (batch size 1) or List[str].

Returns:

Number of samples written. This is the batch size unless on_overflow="trim" dropped part of the batch, in which case it is smaller (possibly 0) and a warning is issued.

Raises:

  • RuntimeError – If writer is not open, has been closed, or could not grow its files for an earlier batch.

  • ValueError – If shapes don’t match, or if the batch does not fit the allocation and on_overflow="error".

class TreeDataSource(directory: str | Path, *, exclude_prefixes: List[str] | None = None, load_into_memory: bool = False)

Read pytrees from memory-mapped files created by TreeWriter.

Provides efficient random access to pre-rendered datasets without loading into RAM. Data is read from memory-mapped binary files and reconstructed into the original pytree structure (AudioTree, dict, etc.).

Each process holds one memmap per leaf, opened on first access and dropped on the way into a pickle, so the source stays safe to hand to grain’s multiprocessing DataLoader however that DataLoader starts its workers. Pass load_into_memory=True to instead load every leaf into RAM up front; a source in that mode opens no files at all after construction.

Parameters:

  • directory – Path to the directory containing manifest.json and memory-mapped data files.

  • exclude_prefixes – List of dot-separated leaf name prefixes to skip loading. For example, ["wet.waveform"] skips the wet.waveform memmap, and ["dry"] skips all leaves under the dry subtree. Excluded leaves are omitted from the reconstructed pytree (AudioTree fields default to None). Default: load all leaves.

  • load_into_memory – If True, load all non-excluded array leaves and string leaves into RAM at init time. Workers then read from pre-loaded numpy arrays instead of memmaps, eliminating disk I/O. With fork-based multiprocessing (default on Linux) the parent’s data is shared with workers via copy-on-write; with spawn it is pickled to them, so the RAM cost is per worker. Samples are copied out of the store on the way out, so a caller that writes into one cannot disturb the next reader. Default False.

Example

First, pre-render a small dataset with TreeWriter. Here each sample is a dict with dry and wet AudioTree branches:

>>> import tempfile
>>> import jax.numpy as jnp
>>> from audiotree import AudioTree
>>> from audiotree.tree_writer import TreeWriter
>>> dataset_dir = tempfile.mkdtemp()
>>> batch = {
...     "dry": AudioTree.create(jnp.zeros((4, 1, 16000)), 16000),
...     "wet": AudioTree.create(jnp.ones((4, 1, 16000)), 16000),
... }
>>> with TreeWriter(dataset_dir, expected_samples=4) as w:
...     _ = w.write(batch)

Read a sample back; the pytree structure is reconstructed:

>>> ds = TreeDataSource(dataset_dir)
>>> sample = ds[0]
>>> sorted(sample.keys())
['dry', 'wet']
>>> type(sample["dry"]).__name__
'AudioTree'

Skip loading some leaves with exclude_prefixes (they come back as None):

>>> ds = TreeDataSource(dataset_dir, exclude_prefixes=["wet.waveform"])
>>> ds[0]["wet"].waveform is None
True

Load everything into RAM up front for I/O-free random access:

>>> ds = TreeDataSource(dataset_dir, load_into_memory=True)
>>> ds[0]["dry"].waveform.shape  # reads from RAM, no disk I/O
(1, 1, 16000)
close() → None

Release this process’s memmaps and readers.

Holding the mappings open is what makes reads fast, but a mapped file cannot be deleted, moved, or replaced on Windows until it is unmapped – so a run that reads a dataset and then tries to clean it up fails with PermissionError: [WinError 32] while the source is alive. POSIX allows the unlink and hides the problem entirely.

Reading again reopens transparently, so this is a release rather than a teardown; TreeDataSource is also a context manager, which is the tidier way to scope the handles:

with TreeDataSource(directory) as source:
    tree = source[0]
get_metadata() → Dict

Get user metadata from the manifest.

Returns:

Dictionary containing user-provided metadata