From 600043ecf17fc644748d172c7017da1e6672eb66 Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Mon, 24 Aug 2026 20:09:28 -0600 Subject: [PATCH 1/5] fix(dataset): stream chunk arrays instead of holding three copies LeechDataset.__init__ held the whole corpus three times over: the numpy members load_chunks reads, one tensor per chunk, and the contiguous output torch.stack allocates while that list is still alive. A 41 GB npz peaked at 116 GB and OOM'd before the first epoch (#211). - Fill a preallocated tensor per field (_TensorFill) instead of stacking a list, removing the stack transient. - Read signals, residuals and features from the npz in row blocks (iter_npz_row_blocks) and defer those members out of load_chunks, so the numpy source is never resident alongside the tensors built from it. Only members the run actually consumes are decompressed. - Store seq_to_sig_maps CSR-style (values + offsets) rather than as a pickled object array, so the member is streamable and costs no per-chunk Python object. Legacy files with the object member still load. - Intern the repeated per-chunk strings in load_chunks. - Correct the load_chunks docstring: np.load never memory-maps a zip member. --- src/leech/calibration.py | 11 +- src/leech/chunking/__init__.py | 22 +- src/leech/chunking/serialization.py | 358 ++++++++++++++++++-- src/leech/dataset.py | 447 ++++++++++++++++++++----- src/leech/splitting/splitter.py | 42 ++- tests/bench_prepare_backends.py | 11 +- tests/test_backend_parity.py | 4 +- tests/test_dataset_streaming.py | 498 ++++++++++++++++++++++++++++ 8 files changed, 1275 insertions(+), 118 deletions(-) create mode 100644 tests/test_dataset_streaming.py diff --git a/src/leech/calibration.py b/src/leech/calibration.py index ac2ef66..8a30335 100644 --- a/src/leech/calibration.py +++ b/src/leech/calibration.py @@ -23,7 +23,6 @@ from torch.optim import LBFGS from torch.utils.data import DataLoader -from leech.chunking import load_chunks from leech.dataset import LeechDataset, collate_fn, resolve_val_dataloader_workers from leech.models import get_model from leech.models.inference_wrapper import ModelInferenceWrapper @@ -178,10 +177,11 @@ def calibrate_model( # so the dataset builds a feature tensor with the same channel count the # model was trained on — otherwise the Conv1d in the feature branch # crashes with "expected N channels, but got M channels". - val_chunks = load_chunks(val_data_path) + # chunk_path (not pre-loaded chunks) so the dataset streams the arrays out + # of the npz instead of holding a second full copy of them (#211). dwell_template_table = config.get("dwell_template_table") or None val_dataset = LeechDataset( - chunks=val_chunks, + chunk_path=val_data_path, model_type=model_name, signal_len=signal_len, kmer_len=kmer_len, @@ -532,10 +532,11 @@ def calibrate_model_multiclass( # so the dataset builds a feature tensor with the same channel count the # model was trained on — otherwise the Conv1d in the feature branch # crashes with "expected N channels, but got M channels". - val_chunks = load_chunks(val_data_path) + # chunk_path (not pre-loaded chunks) so the dataset streams the arrays out + # of the npz instead of holding a second full copy of them (#211). dwell_template_table = config.get("dwell_template_table") or None val_dataset = LeechDataset( - chunks=val_chunks, + chunk_path=val_data_path, model_type=model_name, signal_len=signal_len, kmer_len=kmer_len, diff --git a/src/leech/chunking/__init__.py b/src/leech/chunking/__init__.py index bf6e544..0338e85 100644 --- a/src/leech/chunking/__init__.py +++ b/src/leech/chunking/__init__.py @@ -12,7 +12,19 @@ find_focus_bases, resolve_feature_window, ) -from leech.chunking.serialization import get_chunk_statistics, load_chunks, save_chunks +from leech.chunking.serialization import ( + DEFERRABLE_FIELDS, + csr_from_object_rows, + csr_gather_index, + csr_offsets_from_lens, + get_chunk_statistics, + iter_npz_row_blocks, + load_chunks, + load_seq_to_sig_csr, + npz_array_members, + npz_member_names, + save_chunks, +) __all__ = [ # Extraction @@ -25,4 +37,12 @@ "save_chunks", "load_chunks", "get_chunk_statistics", + "DEFERRABLE_FIELDS", + "iter_npz_row_blocks", + "load_seq_to_sig_csr", + "npz_array_members", + "npz_member_names", + "csr_from_object_rows", + "csr_gather_index", + "csr_offsets_from_lens", ] diff --git a/src/leech/chunking/serialization.py b/src/leech/chunking/serialization.py index 30a189e..182ce80 100644 --- a/src/leech/chunking/serialization.py +++ b/src/leech/chunking/serialization.py @@ -5,13 +5,234 @@ numpy format (.npz files). """ +import contextlib import logging +import zipfile +from collections.abc import Collection, Iterator from pathlib import Path import numpy as np logger = logging.getLogger("leech.chunking.serialization") +# Chunk fields `load_chunks(..., defer=...)` can skip. These are the per-chunk +# arrays; everything else in the file is a scalar or a string and stays cheap. +DEFERRABLE_FIELDS = frozenset( + { + "signal", + "signal_residual", + "dwell", + "features", + "seq_to_sig_map", + "sequence_with_kmer_context", + } +) + + +def _read_npy_header(fp) -> tuple[tuple[int, ...], bool, np.dtype]: + """Read the .npy magic + header from an open member stream. + + numpy 2.x moved the version-dispatching ``_read_array_header`` out of the + public namespace, so dispatch on the magic ourselves. + """ + major, _minor = np.lib.format.read_magic(fp) + if major == 1: + shape, fortran_order, dtype = np.lib.format.read_array_header_1_0(fp) + else: + shape, fortran_order, dtype = np.lib.format.read_array_header_2_0(fp) + return tuple(shape), fortran_order, dtype + + +# ZipExtFile has no native readinto: BufferedIOBase's default allocates a bytes +# object the size of the request. Fill the block in bounded slices so that +# transient stays small no matter how large a block is. +_READ_SLICE_BYTES = 1 << 20 + + +def _readinto_exact(fp, buf: memoryview) -> int: + """Fill ``buf`` from ``fp``, returning the bytes read (short only at EOF).""" + total = 0 + while total < len(buf): + end = min(total + _READ_SLICE_BYTES, len(buf)) + got = fp.readinto(buf[total:end]) + if not got: + break + total += got + return total + + +def csr_offsets_from_lens(lens: np.ndarray) -> np.ndarray: + """Row offsets for CSR storage: ``[0, lens[0], lens[0] + lens[1], ...]``.""" + offsets = np.zeros(len(lens) + 1, dtype=np.int64) + np.cumsum(lens, out=offsets[1:]) + return offsets + + +def csr_gather_index( + offsets: np.ndarray, rows: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Index arrays for reading a subset of CSR rows as one concatenated run. + + Args: + offsets: CSR row offsets. + rows: Row indices to gather, ascending. + + Returns: + ``(lens, col, src)``: per-row lengths, the within-row column of every + gathered element, and the index into the values array of every gathered + element. ``values[src]`` is the concatenation of the selected rows. + """ + lens = (offsets[rows + 1] - offsets[rows]).astype(np.int64) + total = int(lens.sum()) + col = np.arange(total, dtype=np.int64) - np.repeat(np.cumsum(lens) - lens, lens) + src = np.repeat(offsets[rows], lens) + col + return lens, col, src + + +def csr_from_object_rows(rows: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Convert a legacy object array of variable-length rows into CSR form.""" + lens = np.fromiter((len(r) for r in rows), dtype=np.int64, count=len(rows)) + offsets = csr_offsets_from_lens(lens) + values = ( + np.concatenate(list(rows)).astype(np.int32) + if offsets[-1] > 0 + else np.empty(0, dtype=np.int32) + ) + return values, offsets + + +def npz_member_names(input_path: Path) -> set[str]: + """Names of every member of an .npz, without reading any data. + + Unlike :func:`npz_array_members` this includes object (pickled) members, so + it answers "was this field written at all" for formats that + :func:`iter_npz_row_blocks` cannot stream. + """ + with zipfile.ZipFile(input_path) as zf: + return {name[:-4] for name in zf.namelist() if name.endswith(".npy")} + + +def npz_array_members(input_path: Path) -> dict[str, tuple[tuple[int, ...], np.dtype]]: + """Shapes and dtypes of the row-streamable members of an .npz, without reading data. + + Only the zip central directory and each member's .npy header are read, so + this is O(number of members) regardless of file size. Members that + :func:`iter_npz_row_blocks` cannot stream — object dtypes (pickled), + Fortran order, 0-d — are omitted, so ``name in npz_array_members(path)`` + is the test for "can I stream this". + + Args: + input_path: Path to .npz file + + Returns: + Mapping of member name (without the .npy suffix) to ``(shape, dtype)``. + + Examples: + >>> members = npz_array_members(Path("chunks.npz")) + >>> members["signals_flat"][0] # doctest: +SKIP + (6668328, 540) + """ + members: dict[str, tuple[tuple[int, ...], np.dtype]] = {} + with zipfile.ZipFile(input_path) as zf: + for info in zf.infolist(): + if not info.filename.endswith(".npy"): + continue + with zf.open(info) as fp: + try: + shape, fortran_order, dtype = _read_npy_header(fp) + except ValueError: # not a .npy stream we understand + continue + if fortran_order or dtype.hasobject or not shape: + continue + members[info.filename[:-4]] = (shape, dtype) + return members + + +def iter_npz_row_blocks( + input_path: Path, + names: Collection[str], + block_rows: int | None = None, + *, + block_bytes: int = 8 << 20, +) -> Iterator[tuple[int, dict[str, np.ndarray]]]: + """Yield ``(row_start, {name: block})`` over the rows of fixed-shape npz members. + + ``np.load`` reads a whole member at once — it never memory-maps a zip + member, compressed or not — so converting a large corpus means holding the + numpy source and the converted output at the same time (#211). This walks + the members as sequential row blocks instead, so only one block per member + is resident. + + **Blocks are reused between iterations.** The yielded arrays are views into + buffers that the next iteration overwrites; copy anything you need to keep. + + Args: + input_path: Path to .npz file + names: Member names to stream, without the .npy suffix. All must have + the same number of rows (see :func:`npz_array_members`). + block_rows: Rows per block. Default sizes the block from ``block_bytes`` + so a wide signal does not silently allocate a huge buffer, and + never exceeds the member's row count. + block_bytes: Target resident bytes across all streamed members, used + when ``block_rows`` is not given. + + Yields: + ``(row_start, blocks)`` where ``blocks[name]`` has ``min(block_rows, + n_rows - row_start)`` rows. + + Raises: + ValueError: If a member is not streamable, row counts disagree, or the + member is truncated. + """ + names = list(names) + if not names or (block_rows is not None and block_rows < 1): + return + + with contextlib.ExitStack() as stack: + opened: dict[str, tuple] = {} + n_rows: int | None = None + for name in names: + # One handle per member: each is then a single sequential read + # rather than interleaved seeks through one shared file object. + zf = stack.enter_context(zipfile.ZipFile(input_path)) + fp = stack.enter_context(zf.open(name + ".npy")) + shape, fortran_order, dtype = _read_npy_header(fp) + if fortran_order or dtype.hasobject or not shape: + raise ValueError(f"npz member '{name}' is not row-streamable") + if n_rows is None: + n_rows = shape[0] + elif shape[0] != n_rows: + raise ValueError(f"npz member '{name}' has {shape[0]} rows, expected {n_rows}") + row_shape = shape[1:] + row_bytes = int(np.prod(row_shape, dtype=np.int64)) * dtype.itemsize + opened[name] = (fp, row_shape, dtype, row_bytes) + + assert n_rows is not None + if block_rows is None: + per_row = sum(entry[3] for entry in opened.values()) + block_rows = max(1, block_bytes // max(per_row, 1)) + block_rows = min(block_rows, max(n_rows, 1)) + + streams: dict[str, tuple] = {} + for name, (fp, row_shape, dtype, row_bytes) in opened.items(): + block = np.empty((block_rows, *row_shape), dtype=dtype) + streams[name] = (fp, block, block.reshape(-1).view(np.uint8), row_bytes) + + start = 0 + while start < n_rows: + rows = min(block_rows, n_rows - start) + blocks = {} + for name, (fp, block, raw, row_bytes) in streams.items(): + want = rows * row_bytes + got = _readinto_exact(fp, memoryview(raw)[:want]) + if got != want: + raise ValueError( + f"npz member '{name}' truncated at row {start}: read {got} of {want} bytes" + ) + blocks[name] = block[:rows] + yield start, blocks + start += rows + def save_chunks(chunks: list[dict], output_path: Path, *, compressed: bool = True) -> None: """ @@ -36,6 +257,10 @@ def save_chunks(chunks: list[dict], output_path: Path, *, compressed: bool = Tru - labels_int: (N,) integer labels (0, 1, or -1 if unset) - read_ids: (N,) string array of read IDs - base_indices: (N,) base indices + - seq_to_sig_values / seq_to_sig_offsets: base-to-signal maps in CSR + form; row i is ``values[offsets[i]:offsets[i + 1]]``. Files written + before v0.6.8 carry a pickled object array named ``seq_to_sig_maps`` + instead, which :func:`load_chunks` still reads. Examples: >>> chunks = extract_training_chunks(read, motif="CCAGGC") @@ -106,8 +331,11 @@ def save_chunks(chunks: list[dict], output_path: Path, *, compressed: bool = Tru source_groups_arr = np.array(source_groups, dtype=str) reference_names_arr = np.array(reference_names, dtype=str) sequences_with_kmer_context_arr = np.array(sequences_with_kmer_context, dtype=str) - # seq_to_sig_maps are variable length (depend on read dwell times), keep as object - seq_to_sig_maps_arr = np.array(seq_to_sig_maps, dtype=object) + # seq_to_sig_maps are variable length (they depend on the read's dwell + # times), so store them CSR-style: one flat values array plus row offsets. + # An object array would be pickled, which costs a Python ndarray per chunk + # on load and makes the member unstreamable (#211). + seq_to_sig_values_arr, seq_to_sig_offsets_arr = csr_from_object_rows(seq_to_sig_maps) # Create parent directories if they don't exist output_path.parent.mkdir(parents=True, exist_ok=True) @@ -124,7 +352,8 @@ def save_chunks(chunks: list[dict], output_path: Path, *, compressed: bool = Tru "feature_ends": feature_ends_arr, "source_groups": source_groups_arr, "reference_names": reference_names_arr, - "seq_to_sig_maps": seq_to_sig_maps_arr, + "seq_to_sig_values": seq_to_sig_values_arr, + "seq_to_sig_offsets": seq_to_sig_offsets_arr, "sequences_with_kmer_context": sequences_with_kmer_context_arr, "cl_values": cl_values_arr, } @@ -170,21 +399,47 @@ def save_chunks(chunks: list[dict], output_path: Path, *, compressed: bool = Tru logger.info(f"Saved {len(chunks)} chunks to {output_path}") -def load_chunks(input_path: Path) -> list[dict]: +def load_seq_to_sig_csr(input_path: Path) -> tuple[np.ndarray, np.ndarray] | None: + """Load the base-to-signal maps in CSR form: ``(values, offsets)``. + + Row ``i`` is ``values[offsets[i]:offsets[i + 1]]``. Files written before + v0.6.8 store these as a pickled object array (``seq_to_sig_maps``); those + are converted here so callers see one representation. + + Args: + input_path: Path to .npz file + + Returns: + ``(values, offsets)``, or None if the file has no base-to-signal maps. + """ + with np.load(input_path, allow_pickle=True) as data: + if "seq_to_sig_values" in data: + return data["seq_to_sig_values"], data["seq_to_sig_offsets"] + if "seq_to_sig_maps" not in data: + return None + return csr_from_object_rows(data["seq_to_sig_maps"]) + + +def load_chunks(input_path: Path, *, defer: Collection[str] = ()) -> list[dict]: """ Load training chunks from compressed numpy format. Args: input_path: Path to .npz file + defer: Chunk array fields to leave unread (see :data:`DEFERRABLE_FIELDS`). + The key is still present on each chunk, set to None. Callers that + convert the arrays themselves — :class:`~leech.dataset.LeechDataset` + streams them with :func:`iter_npz_row_blocks` — use this to avoid + holding a second full copy. Returns: List of chunk dictionaries compatible with extract_training_chunks output Note: - This function loads all arrays into memory at once. The arrays are stored - as numpy object arrays (dtype=object) to handle variable-length signals. - The loaded data is kept in memory-mapped form when possible, but converting - to individual dictionaries will create copies in memory. + Every member requested is read in full: ``np.load`` does not memory-map + zip members, compressed or not, so there is no lazy path here. On a + large corpus this dominates peak memory, which is what ``defer`` and + :func:`iter_npz_row_blocks` exist to avoid (#211). Examples: >>> chunks = load_chunks(Path("output/chunks.npz")) @@ -192,35 +447,70 @@ def load_chunks(input_path: Path) -> list[dict]: >>> for chunk in chunks[:5]: ... print(f"{chunk['read_id']}: {chunk['label']}") """ - # Load all arrays at once (keeps data memory-mapped when possible) + deferred = frozenset(defer) + unknown = deferred - DEFERRABLE_FIELDS + if unknown: + raise ValueError( + f"Cannot defer unknown chunk field(s): {sorted(unknown)}. " + f"Deferrable: {sorted(DEFERRABLE_FIELDS)}" + ) + + # Repeated string fields: a few hundred distinct values across millions of + # chunks, and str(arr[i]) mints a fresh object every time. + _interned: dict[str, str] = {} + + def _intern(value) -> str: + text = str(value) + return _interned.setdefault(text, text) + with np.load(input_path, allow_pickle=True) as data: # Detect format: flat arrays (new, fast) vs object arrays (old, backward compat) has_flat_signals = "signals_flat" in data has_flat_dwells = "dwells_flat" in data has_flat_features = "features_flat" in data - signals = data["signals_flat"] if has_flat_signals else data["signals"] + signals = ( + None + if "signal" in deferred + else data["signals_flat" if has_flat_signals else "signals"] + ) sequences = data["sequences"] - dwells = data["dwells_flat"] if has_flat_dwells else data["dwells"] - features = data["features_flat"] if has_flat_features else data["features"] + dwells = ( + None if "dwell" in deferred else data["dwells_flat" if has_flat_dwells else "dwells"] + ) + features = ( + None + if "features" in deferred + else data["features_flat" if has_flat_features else "features"] + ) labels_arr = data["labels"] # String labels labels_int_arr = data["labels_int"] # Numeric labels read_ids = data["read_ids"] base_indices = data["base_indices"] - # New fields for signal_kmer encoding (backward compatible) - has_sig_kmer = "seq_to_sig_maps" in data + # Base-to-signal maps: CSR pair (new) or pickled object array (old) + has_sig_kmer = "seq_to_sig_values" in data or "seq_to_sig_maps" in data + seq_to_sig_values = seq_to_sig_offsets = seq_to_sig_maps = None + sequences_with_kmer_context = None if has_sig_kmer: - seq_to_sig_maps = data["seq_to_sig_maps"] - sequences_with_kmer_context = data["sequences_with_kmer_context"] + if "seq_to_sig_map" not in deferred: + if "seq_to_sig_values" in data: + seq_to_sig_values = data["seq_to_sig_values"] + seq_to_sig_offsets = data["seq_to_sig_offsets"] + else: + seq_to_sig_maps = data["seq_to_sig_maps"] + if "sequence_with_kmer_context" not in deferred: + sequences_with_kmer_context = data["sequences_with_kmer_context"] # Signal residual channel (backward compatible) has_signal_residual_data = "signal_residuals_flat" in data or "signal_residuals" in data - if has_signal_residual_data: + if has_signal_residual_data and "signal_residual" not in deferred: signal_residuals_loaded = ( data["signal_residuals_flat"] if "signal_residuals_flat" in data else data["signal_residuals"] ) + else: + signal_residuals_loaded = None # Feature window params (new format: feature_starts/feature_ends) has_feature_se = "feature_starts" in data @@ -259,22 +549,32 @@ def load_chunks(input_path: Path) -> list[dict]: # Create dictionaries with references to array elements for i in range(n_chunks): chunk = { - "signal": signals[i], - "sequence": str(sequences[i]), - "dwell": dwells[i], - "features": features[i], + "signal": None if signals is None else signals[i], + "sequence": _intern(sequences[i]), + "dwell": None if dwells is None else dwells[i], + "features": None if features is None else features[i], "read_id": str(read_ids[i]), "base_idx": int(base_indices[i]), - "label": str(labels_arr[i]) if labels_arr[i] != "" else None, + "label": _intern(labels_arr[i]) if labels_arr[i] != "" else None, "label_int": int(labels_int_arr[i]) if labels_int_arr[i] >= 0 else None, } if has_sig_kmer: - s2s = seq_to_sig_maps[i] - seq_ctx = str(sequences_with_kmer_context[i]) - chunk["seq_to_sig_map"] = s2s if len(s2s) > 0 else None - chunk["sequence_with_kmer_context"] = seq_ctx if seq_ctx else None + if seq_to_sig_offsets is not None: + s2s = seq_to_sig_values[seq_to_sig_offsets[i] : seq_to_sig_offsets[i + 1]] + elif seq_to_sig_maps is not None: + s2s = seq_to_sig_maps[i] + else: + s2s = None + chunk["seq_to_sig_map"] = s2s if s2s is not None and len(s2s) > 0 else None + if sequences_with_kmer_context is None: + chunk["sequence_with_kmer_context"] = None + else: + seq_ctx = str(sequences_with_kmer_context[i]) + chunk["sequence_with_kmer_context"] = seq_ctx if seq_ctx else None if has_signal_residual_data: - chunk["signal_residual"] = signal_residuals_loaded[i] + chunk["signal_residual"] = ( + None if signal_residuals_loaded is None else signal_residuals_loaded[i] + ) if has_feature_se: chunk["feature_start"] = int(feature_starts_loaded[i]) chunk["feature_end"] = int(feature_ends_loaded[i]) @@ -282,10 +582,10 @@ def load_chunks(input_path: Path) -> list[dict]: # Old format: convert dwell_margin_left to feature_left chunk["dwell_margin_left"] = int(dwell_margin_lefts[i]) if has_source_groups: - sg = str(source_groups[i]) + sg = _intern(source_groups[i]) chunk["source_group"] = sg if sg else None if has_reference_names: - rn = str(reference_names_loaded[i]) + rn = _intern(reference_names_loaded[i]) chunk["reference_name"] = rn if rn else "" if has_cl_values: cl_val = int(cl_values_loaded[i]) diff --git a/src/leech/dataset.py b/src/leech/dataset.py index 465bf6f..b4426c7 100644 --- a/src/leech/dataset.py +++ b/src/leech/dataset.py @@ -35,6 +35,7 @@ import logging import os +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -43,7 +44,13 @@ import torch from torch.utils.data import Dataset -from leech.chunking import load_chunks +from leech.chunking import ( + iter_npz_row_blocks, + load_chunks, + load_seq_to_sig_csr, + npz_array_members, + npz_member_names, +) from leech.constants import AUTO_DATALOADER_WORKERS from leech.features import encode_signal_kmer, sequence_to_int from leech.models.inference_wrapper import ModelInferenceWrapper @@ -179,6 +186,178 @@ def append_dwell_template_channels( return np.concatenate([features, template_channels], axis=0) +# Rows per block when expanding the CSR base-to-signal maps. Keeps the gather +# index arrays to tens of MB regardless of corpus size. +_S2S_BLOCK_ROWS = 65536 + + +class _TensorFill: + """Accumulate per-chunk tensors into one preallocated contiguous tensor. + + ``torch.stack`` allocates the whole output *while the input list is still + alive*, so stacking N chunk tensors peaks at twice the output — 33 GB of + transient on a large corpus (#211). The output shape is known before the + loop, so fill a buffer allocated up front instead and peak at once the + output. + + Chunk tensors whose shape disagrees with the first one fall back to the old + list, which ``__getitem__`` still handles. + """ + + def __init__(self, name: str, capacity: int): + self.name = name + self.capacity = capacity + self._tensor: torch.Tensor | None = None + self._items: list[torch.Tensor] = [] + self._n = 0 + + def append(self, tensor: torch.Tensor) -> None: + if self._items: # already degraded to list access + self._items.append(tensor) + return + if self._tensor is None: + self._tensor = torch.empty((self.capacity, *tensor.shape), dtype=tensor.dtype) + elif tuple(tensor.shape) != tuple(self._tensor.shape[1:]): + logger.warning( + "%s shapes differ (%s vs %s), falling back to list access", + self.name, + tuple(tensor.shape), + tuple(self._tensor.shape[1:]), + ) + self._items = [self._tensor[i].clone() for i in range(self._n)] + self._tensor = None + self._items.append(tensor) + return + self._tensor[self._n] = tensor + self._n += 1 + + def finish(self) -> tuple[torch.Tensor | None, list[torch.Tensor]]: + """Return ``(stacked_tensor, fallback_list)``; exactly one is populated.""" + if self._tensor is None: + # Nothing appended (the caller uses a different tensor) or shapes + # differed — both cases the old _try_stack signalled with None. + return None, self._items + return self._tensor[: self._n], [] + + +@dataclass +class _ArrayStream: + """Row-block source for the per-chunk arrays of an npz corpus. + + Reading the arrays this way — instead of through ``load_chunks``, which + materialises every member — keeps the numpy source out of the peak: only + one block per member is resident while the output tensors are filled. + + Attributes: + path: The npz being streamed. + members: Chunk field name -> npz member name, for the fields this run + actually consumes. Members nothing reads are never decompressed. + keep: One bool per npz row: False for rows dropped by label filtering, + so the stream stays aligned with ``LeechDataset.chunks``. + dwell_width: ``len(chunk["dwell"])``, constant in the flat format and + read from the member header — the only thing the tensorize loop + needs from the dwells, so that member is never read. + """ + + path: Path + members: dict[str, str] + dwell_width: int | None + keep: np.ndarray | None = None + + @classmethod + def build( + cls, + path: Path, + *, + needs_features: bool, + wants_residual: bool, + ) -> "_ArrayStream | None": + """Return a stream for ``path``, or None if it cannot be streamed. + + Corpora written before the flat format store the arrays as pickled + object members; those still take the eager path. + """ + members = npz_array_members(path) + if "signals_flat" not in members: + return None + wanted = {"signal": "signals_flat"} + if wants_residual: + if "signal_residuals_flat" in members: + wanted["signal_residual"] = "signal_residuals_flat" + elif "signal_residuals" in npz_member_names(path): + return None # residuals present but not streamable + if needs_features: + # dwells_flat carries no values the loop needs, only its width — + # but without it there is no way to know that width without reading + # the per-chunk dwells, so fall back. + if "features_flat" not in members or "dwells_flat" not in members: + return None + wanted["features"] = "features_flat" + dwell_width = members["dwells_flat"][0][1] if "dwells_flat" in members else None + return cls(path=path, members=wanted, dwell_width=dwell_width) + + def __iter__(self): + """Yield one dict of arrays per kept row, in ``LeechDataset.chunks`` order. + + The arrays are views into recycled block buffers — valid until the next + iteration, which is all the tensorize loop needs since it copies each + row into the output tensor. + """ + if self.keep is None: + raise RuntimeError("_ArrayStream.keep must be set before iterating") + for start, blocks in iter_npz_row_blocks(self.path, list(self.members.values())): + rows = len(next(iter(blocks.values()))) + for j in np.nonzero(self.keep[start : start + rows])[0]: + yield {field: blocks[member][j] for field, member in self.members.items()} + + +def _expand_seq_to_sig_csr( + values: np.ndarray, + offsets: np.ndarray, + rows: np.ndarray, + *, + signal_len: int, + crop_starts: np.ndarray | None, +) -> np.ndarray: + """Expand CSR base-to-signal maps into one padded ``(len(rows), max_len)`` array. + + Args: + values: Flat concatenated map values. + offsets: Row offsets, one per row plus a final total. + rows: npz row indices to expand, ascending. + signal_len: Padding value — the encoder's ``sig_start < signal_len`` + test fails at padded positions, so they contribute nothing. + crop_starts: Per-row signal offset to subtract (asymmetric crop), or + None to keep the stored coordinates. + + Returns: + int64 array of shape ``(len(rows), max_len)``. + """ + lens = (offsets[rows + 1] - offsets[rows]).astype(np.int64) + n = len(rows) + max_len = int(lens.max()) if n else 0 + padded = np.full((n, max_len), signal_len, dtype=np.int64) + starts = offsets[rows] + + # Blocked so the gather indices stay small on a multi-million-chunk corpus. + for block_start in range(0, n, _S2S_BLOCK_ROWS): + block_end = min(block_start + _S2S_BLOCK_ROWS, n) + block_lens = lens[block_start:block_end] + total = int(block_lens.sum()) + if total == 0: + continue + row_idx = np.repeat(np.arange(block_start, block_end), block_lens) + col_idx = np.arange(total) - np.repeat(np.cumsum(block_lens) - block_lens, block_lens) + gathered = values[np.repeat(starts[block_start:block_end], block_lens) + col_idx].astype( + np.int64 + ) + if crop_starts is not None: + gathered -= np.repeat(crop_starts[block_start:block_end], block_lens) + np.clip(gathered, 0, signal_len, out=gathered) + padded[row_idx, col_idx] = gathered + return padded + + class LeechDataset(Dataset): """ PyTorch Dataset for leech training chunks. @@ -269,27 +448,49 @@ def __init__( if dwell_template_table is not None: self._load_dwell_templates(Path(dwell_template_table)) - # Use pre-loaded chunks or load from file + self._needs_features = model_type in FEATURE_MODELS + + # Use pre-loaded chunks or load from file. Loading from a path streams + # the per-chunk arrays out of the npz a row block at a time instead of + # holding a full numpy copy alongside the tensors built from it (#211); + # pre-loaded chunks have already paid that cost. + self._array_stream: _ArrayStream | None = None + self._npz_members: set[str] = set() + self._s2s_csr: tuple[np.ndarray, np.ndarray] | None = None + self._s2s_rows: np.ndarray | None = None if chunks is not None: logger.info(f"Using {len(chunks)} pre-loaded chunks (skipping disk I/O)") self.chunks = chunks elif chunk_path is not None: - self.chunks = load_chunks(chunk_path) + self._load_from_path( + Path(chunk_path), signal_mode=signal_mode, seq_encoding=seq_encoding + ) else: raise ValueError("Either chunk_path or chunks must be provided") - # Filter chunks with valid numeric labels (label_int) + # Filter chunks with valid numeric labels (label_int). The streaming + # path has already applied this, and recorded the same mask so npz rows + # stay aligned with self.chunks — a mismatch here would pair signals + # with the wrong labels silently. self.chunks = [c for c in self.chunks if c["label_int"] is not None] if len(self.chunks) == 0: raise ValueError(f"No valid chunks found{f' in {chunk_path}' if chunk_path else ''}") - # Pre-tensorize: encode sequences, labels, signals, and features once + # Pre-tensorize: encode sequences, labels, signals, and features once. + # Each accumulator fills one preallocated contiguous tensor — see + # _TensorFill for why stacking a list instead doubles the peak. + n_chunks = len(self.chunks) + fill_encoded_seqs = _TensorFill("Encoded-sequence", n_chunks) + fill_labels = _TensorFill("Label", n_chunks) + fill_signals = _TensorFill("Signal", n_chunks) + fill_features = _TensorFill("Feature", n_chunks) + fill_confounds = _TensorFill("Confound", n_chunks) + fill_cl_targets = _TensorFill("CL target", n_chunks) self._encoded_seqs: list[torch.Tensor] = [] self._labels: list[torch.Tensor] = [] self._signals: list[torch.Tensor] = [] self._features: list[torch.Tensor] = [] - self._needs_features = model_type in FEATURE_MODELS self._confound_encoder = confound_encoder self._has_confound = confound_encoder is not None self._confound_labels: list[torch.Tensor] = [] @@ -300,18 +501,31 @@ def __init__( self._effective_seq_encoding = seq_encoding if seq_encoding == "signal_kmer": first = self.chunks[0] - if not first.get("seq_to_sig_map") is not None or not first.get( - "sequence_with_kmer_context" - ): + if self._s2s_csr is not None: + # Streaming path: the maps were deferred out of the chunk dicts, + # so ask the CSR arrays whether the first chunk has one. + _offsets = self._s2s_csr[1] + _row = self._s2s_rows[0] + has_seq_to_sig = bool(_offsets[_row + 1] > _offsets[_row]) + else: + has_seq_to_sig = first.get("seq_to_sig_map") is not None + if not has_seq_to_sig or not first.get("sequence_with_kmer_context"): logger.warning( "Chunks lack seq_to_sig_map/sequence_with_kmer_context; " "falling back to base_onehot encoding" ) self._effective_seq_encoding = "base_onehot" + self._s2s_csr = None # Detect signal_residual channel and apply signal_mode self._signal_mode = signal_mode - self._has_signal_residual = self.chunks[0].get("signal_residual") is not None + if self._array_stream is not None: + # Arrays were deferred, so presence is a property of the file. + self._has_signal_residual = bool( + {"signal_residuals_flat", "signal_residuals"} & self._npz_members + ) + else: + self._has_signal_residual = self.chunks[0].get("signal_residual") is not None if signal_mode == "both" and self._has_signal_residual: self.signal_channels = 2 else: @@ -330,38 +544,51 @@ def __init__( self._seq_ints: list[np.ndarray] = [] self._seq_to_sig: list[np.ndarray] = [] + # Arrays come either from the chunk dicts (pre-loaded / legacy corpus) + # or a row-block stream over the npz. Both yield the same field names, + # so the loop below does not care which. + array_iter = iter(self._array_stream) if self._array_stream is not None else None + stream_dwell_width = ( + self._array_stream.dwell_width if self._array_stream is not None else None + ) + for chunk in self.chunks: + arrays = chunk if array_iter is None else next(array_iter) + if self._effective_seq_encoding == "signal_kmer": seq_ctx = chunk["sequence_with_kmer_context"] seq_ints = sequence_to_int(seq_ctx).astype(np.int8) self._seq_ints.append(seq_ints) - s2s = chunk["seq_to_sig_map"].astype(np.int64, copy=True) - if self.left_context is not None and self.right_context is not None: - stored_focus = chunk.get("focus_signal_pos") - focus_pos = stored_focus if stored_focus is not None else int(s2s[-1]) // 2 - crop_start = focus_pos - self.left_context - s2s -= crop_start - np.clip(s2s, 0, signal_len, out=s2s) - self._seq_to_sig.append(s2s) + if self._s2s_csr is None: + # Non-streaming path: one map per chunk dict. The streaming + # path expands all of them at once, after this loop. + s2s = chunk["seq_to_sig_map"].astype(np.int64, copy=True) + if self.left_context is not None and self.right_context is not None: + stored_focus = chunk.get("focus_signal_pos") + focus_pos = stored_focus if stored_focus is not None else int(s2s[-1]) // 2 + crop_start = focus_pos - self.left_context + s2s -= crop_start + np.clip(s2s, 0, signal_len, out=s2s) + self._seq_to_sig.append(s2s) else: # Pre-encode sequence (vectorized, no Python loop) - self._encoded_seqs.append(self._encode_sequence(chunk["sequence"])) + fill_encoded_seqs.append(self._encode_sequence(chunk["sequence"])) # Pre-create label tensor: long for multi-class, float for binary if self._multiclass: - self._labels.append(torch.tensor(chunk["label_int"], dtype=torch.long)) + fill_labels.append(torch.tensor(chunk["label_int"], dtype=torch.long)) else: - self._labels.append(torch.tensor([chunk["label_int"]], dtype=torch.float32)) + fill_labels.append(torch.tensor([chunk["label_int"]], dtype=torch.float32)) # Pre-tensorize signal: pad/crop once instead of every __getitem__ call - signal = chunk["signal"] + signal = arrays["signal"] if signal.dtype != np.float32: signal = signal.astype(np.float32) - signal_residual = chunk.get("signal_residual") + signal_residual = arrays.get("signal_residual") if signal_residual is not None and signal_residual.dtype != np.float32: signal_residual = signal_residual.astype(np.float32) - self._signals.append( + fill_signals.append( self._prepare_signal( signal, signal_residual, focus_signal_pos=chunk.get("focus_signal_pos") ) @@ -369,52 +596,36 @@ def __init__( # Pre-tensorize features: apply dwell_offset slicing once if self._needs_features: - self._features.append(self._prepare_features(chunk)) - else: - self._features.append(torch.empty(0)) + dwell_width = ( + stream_dwell_width if stream_dwell_width is not None else len(chunk["dwell"]) + ) + fill_features.append(self._prepare_features(arrays["features"], dwell_width, chunk)) # Confound label for adversarial training. The encoder reads the # configured chunk field and maps it to a class int (-1 = ignore). if self._confound_encoder is not None: confound_class = self._confound_encoder.encode(chunk) - self._confound_labels.append(torch.tensor(confound_class, dtype=torch.long)) + fill_confounds.append(torch.tensor(confound_class, dtype=torch.long)) # CL regression target (cl_value / 255.0; sentinel -1.0 for missing) if self._cl_regression: cl_val = chunk.get("cl_value") if cl_val is not None and cl_val >= 0: - self._cl_targets.append(torch.tensor(cl_val / 255.0, dtype=torch.float32)) + fill_cl_targets.append(torch.tensor(cl_val / 255.0, dtype=torch.float32)) else: - self._cl_targets.append(torch.tensor(-1.0, dtype=torch.float32)) + fill_cl_targets.append(torch.tensor(-1.0, dtype=torch.float32)) - # Stack every per-chunk list into one contiguous tensor. This isn't just - # for cache friendliness — it's required for fork-safety. A DataLoader - # with num_workers > 0 forks worker processes that COW-inherit the - # parent's address space. CPython refcounts live inside each PyObject - # header, so a worker iterating a list of N tensors writes to N + # Every per-chunk tensor now lives in one contiguous buffer. That isn't + # just for cache friendliness — it's required for fork-safety. A + # DataLoader with num_workers > 0 forks worker processes that COW-inherit + # the parent's address space. CPython refcounts live inside each + # PyObject header, so a worker iterating a list of N tensors writes to N # separate page-resident headers and faults every page into a private # copy, multiplying peak RSS by (1 + num_workers). A single contiguous # tensor keeps its data buffer outside Python's GC, so the buffer # pages are genuinely shared across the fork. - def _try_stack(name: str, items: list[torch.Tensor]) -> torch.Tensor | None: - if not items: - # Expected when a feature mode isn't active (e.g. signal_kmer - # leaves _encoded_seqs empty); the consumer will use a different - # tensor instead. - return None - try: - return torch.stack(items) - except RuntimeError as e: - logger.warning("%s shapes differ, falling back to list access: %s", name, e) - return None - - self._signals_tensor = _try_stack("Signal", self._signals) - if self._signals_tensor is not None: - self._signals = [] - - self._encoded_seqs_tensor = _try_stack("Encoded-sequence", self._encoded_seqs) - if self._encoded_seqs_tensor is not None: - self._encoded_seqs = [] + self._signals_tensor, self._signals = fill_signals.finish() + self._encoded_seqs_tensor, self._encoded_seqs = fill_encoded_seqs.finish() # Compact signal_kmer inputs — only populated when encoding == signal_kmer. # Stacked into fork-safe int tensors. encode_signal_kmer is then called @@ -428,39 +639,49 @@ def _try_stack(name: str, items: list[torch.Tensor]) -> torch.Tensor | None: self._seq_to_sig_tensor: torch.Tensor | None = None if self._effective_seq_encoding == "signal_kmer" and self._seq_ints: max_seq_ints_len = max(s.shape[0] for s in self._seq_ints) - max_s2s_len = max(s.shape[0] for s in self._seq_to_sig) n = len(self._seq_ints) padded_seq_ints = np.full((n, max_seq_ints_len), -1, dtype=np.int8) - padded_s2s = np.full((n, max_s2s_len), signal_len, dtype=np.int64) - for i, (si, s2s) in enumerate(zip(self._seq_ints, self._seq_to_sig, strict=True)): + for i, si in enumerate(self._seq_ints): padded_seq_ints[i, : si.shape[0]] = si - padded_s2s[i, : s2s.shape[0]] = s2s self._seq_ints_tensor = torch.from_numpy(padded_seq_ints) - self._seq_to_sig_tensor = torch.from_numpy(padded_s2s) self._seq_ints = [] + + if self._s2s_csr is not None: + # Streaming path: expand every map at once from the CSR pair, + # instead of one astype/clip per chunk. + values, offsets = self._s2s_csr + crop_starts = None + if self.left_context is not None and self.right_context is not None: + crop_starts = self._crop_starts(values, offsets) - self.left_context + padded_s2s = _expand_seq_to_sig_csr( + values, + offsets, + self._s2s_rows, + signal_len=signal_len, + crop_starts=crop_starts, + ) + self._s2s_csr = None + else: + max_s2s_len = max(s.shape[0] for s in self._seq_to_sig) + padded_s2s = np.full((n, max_s2s_len), signal_len, dtype=np.int64) + for i, s2s in enumerate(self._seq_to_sig): + padded_s2s[i, : s2s.shape[0]] = s2s + self._seq_to_sig_tensor = torch.from_numpy(padded_s2s) self._seq_to_sig = [] - self._labels_tensor = _try_stack("Label", self._labels) - if self._labels_tensor is not None: - self._labels = [] + self._labels_tensor, self._labels = fill_labels.finish() self._features_tensor: torch.Tensor | None = None if self._needs_features: - self._features_tensor = _try_stack("Feature", self._features) - if self._features_tensor is not None: - self._features = [] + self._features_tensor, self._features = fill_features.finish() self._confound_labels_tensor: torch.Tensor | None = None if self._has_confound: - self._confound_labels_tensor = _try_stack("Confound", self._confound_labels) - if self._confound_labels_tensor is not None: - self._confound_labels = [] + self._confound_labels_tensor, self._confound_labels = fill_confounds.finish() self._cl_targets_tensor: torch.Tensor | None = None if self._cl_regression: - self._cl_targets_tensor = _try_stack("CL target", self._cl_targets) - if self._cl_targets_tensor is not None: - self._cl_targets = [] + self._cl_targets_tensor, self._cl_targets = fill_cl_targets.finish() # Drop the raw numpy arrays from self.chunks now that everything has # been pre-tensorized. External code (samplers, label tally, feature @@ -503,6 +724,70 @@ def _try_stack(name: str, items: list[torch.Tensor]) -> torch.Tensor | None: f"({_n_encoded} sequences encoded, encoding={self._effective_seq_encoding})" ) + def _load_from_path(self, chunk_path: Path, *, signal_mode: str, seq_encoding: str) -> None: + """Load chunk metadata from an npz, deferring the arrays a stream can supply. + + Sets ``self.chunks`` and, when the corpus is row-streamable, + ``self._array_stream`` (plus the CSR base-to-signal maps when the run + needs them). Falls back to loading everything eagerly for corpora + written before the flat array format. + """ + self._npz_members = npz_member_names(chunk_path) + stream = _ArrayStream.build( + chunk_path, + needs_features=self._needs_features, + wants_residual=signal_mode in ("both", "residual"), + ) + if stream is None: + logger.debug("%s is not row-streamable; loading arrays eagerly", chunk_path) + self.chunks = load_chunks(chunk_path) + return + + # Nothing downstream reads the raw arrays off the chunk dicts — the + # stream supplies signal/residual/features, the dwell width comes from + # the member header, and the base-to-signal maps are expanded from CSR + # only when signal_kmer needs them. So never read them into the dicts. + defer = {"signal", "signal_residual", "dwell", "features", "seq_to_sig_map"} + if seq_encoding != "signal_kmer": + defer.add("sequence_with_kmer_context") + raw = load_chunks(chunk_path, defer=defer) + + keep = np.fromiter((c["label_int"] is not None for c in raw), dtype=bool, count=len(raw)) + stream.keep = keep + self._array_stream = stream + self.chunks = [c for c, k in zip(raw, keep, strict=True) if k] + if seq_encoding == "signal_kmer": + self._s2s_csr = load_seq_to_sig_csr(chunk_path) + self._s2s_rows = np.nonzero(keep)[0] + + def _crop_starts(self, values: np.ndarray, offsets: np.ndarray) -> np.ndarray: + """Per-chunk focus signal position for the streamed base-to-signal maps. + + Mirrors the per-chunk rule: the stored ``focus_signal_pos`` when the + corpus has one, else half the map's last value (old symmetric data). + """ + rows = self._s2s_rows + assert rows is not None + focus = np.zeros(len(rows), dtype=np.int64) + missing: list[int] = [] + for i, chunk in enumerate(self.chunks): + stored = chunk.get("focus_signal_pos") + if stored is None: + missing.append(i) + else: + focus[i] = stored + if missing: + idx = np.asarray(missing) + starts = offsets[rows[idx]] + ends = offsets[rows[idx] + 1] + if np.any(ends <= starts): + raise ValueError( + "Chunk without focus_signal_pos has an empty seq_to_sig_map; " + "cannot place the asymmetric crop" + ) + focus[idx] = values[ends - 1].astype(np.int64) // 2 + return focus + def _prepare_signal( self, signal: np.ndarray, @@ -577,10 +862,18 @@ def _append_template_channels(self, features: np.ndarray, chunk: dict) -> np.nda template_min_pos=self._template_min_pos, ) - def _prepare_features(self, chunk: dict) -> torch.Tensor: - """Apply dwell_offset slicing and tensorize features. Called once during __init__.""" - dwell = chunk["dwell"] - features = chunk["features"] + def _prepare_features( + self, features: np.ndarray, dwell_width: int, chunk: dict + ) -> torch.Tensor: + """Apply dwell_offset slicing and tensorize features. Called once during __init__. + + Args: + features: The chunk's feature array, ``(num_features, feat_width)``. + dwell_width: ``len(chunk["dwell"])``. Only its width matters here, + so the streaming path takes it from the member header rather + than reading the dwells at all. + chunk: The chunk dict, for the feature-window metadata. + """ if features.dtype != np.float32: features = features.astype(np.float32) @@ -591,7 +884,7 @@ def _prepare_features(self, chunk: dict) -> torch.Tensor: if self.model_type in WIDE_FEATURE_MODELS: pass # full-width features - elif len(dwell) > self.kmer_len: + elif dwell_width > self.kmer_len: # Determine feature_start (signed offset from focus). # New chunks have it directly; old chunks need conversion. if "feature_start" in chunk: @@ -601,7 +894,7 @@ def _prepare_features(self, chunk: dict) -> torch.Tensor: elif "dwell_margin_left" in chunk: feat_start = -(kmer_context + int(chunk["dwell_margin_left"])) else: - feat_start = -(len(dwell) - 1) // 2 # symmetric fallback + feat_start = -(dwell_width - 1) // 2 # symmetric fallback # kmer-aligned start within the feature array # Feature array starts at focus + feat_start, kmer starts at focus - kmer_context kmer_start = (-kmer_context) - feat_start diff --git a/src/leech/splitting/splitter.py b/src/leech/splitting/splitter.py index a46b515..f3e10f8 100644 --- a/src/leech/splitting/splitter.py +++ b/src/leech/splitting/splitter.py @@ -13,10 +13,19 @@ import numpy as np -from leech.chunking import load_chunks +from leech.chunking import ( + csr_from_object_rows, + csr_gather_index, + csr_offsets_from_lens, + load_chunks, +) logger = logging.getLogger("leech.splitting.splitter") +# npz members holding the CSR base-to-signal maps; merged by row gather, not +# by boolean mask, so they are excluded from the generic member loop. +_S2S_MEMBERS = frozenset({"seq_to_sig_values", "seq_to_sig_offsets", "seq_to_sig_maps"}) + def _source_group_from_path(chunk_path: Path) -> str: """Extract source group name from a chunk file path. @@ -143,6 +152,11 @@ def _merge_arrays_by_split( # Accumulators: split_name -> array_key -> list of arrays accumulators: dict[str, dict[str, list[np.ndarray]]] = {s: {} for s in split_names} + # Base-to-signal maps are CSR (values + offsets), so they are selected by + # gathering rows rather than by masking, and their offsets are rebuilt from + # the accumulated row lengths at save time. + s2s_lens: dict[str, list[np.ndarray]] = {s: [] for s in split_names} + for chunk_path in input_paths: with np.load(chunk_path, allow_pickle=True) as data: # Get read_ids to build masks @@ -154,8 +168,18 @@ def _merge_arrays_by_split( for sname, rid_set in split_read_ids.items(): masks[sname] = np.array([r in rid_set for r in read_ids_str], dtype=bool) + # Base-to-signal maps: CSR pair, or a legacy object array normalized + # to CSR so a merge of mixed-vintage inputs writes one format. + if "seq_to_sig_values" in data: + s2s_values = data["seq_to_sig_values"] + s2s_offsets = data["seq_to_sig_offsets"] + elif "seq_to_sig_maps" in data: + s2s_values, s2s_offsets = csr_from_object_rows(data["seq_to_sig_maps"]) + else: + s2s_values = s2s_offsets = None + # Collect all array keys from the file - array_keys = list(data.keys()) + array_keys = [key for key in data.keys() if key not in _S2S_MEMBERS] for sname in split_names: mask = masks[sname] @@ -163,6 +187,15 @@ def _merge_arrays_by_split( continue count = int(mask.sum()) + if s2s_offsets is None: + # No maps in this file; keep the row count aligned so a + # merge with files that do have them stays row-indexable. + s2s_lens[sname].append(np.zeros(count, dtype=np.int64)) + else: + rows = np.nonzero(mask)[0] + lens, _col, src = csr_gather_index(s2s_offsets, rows) + s2s_lens[sname].append(lens) + accumulators[sname].setdefault("seq_to_sig_values", []).append(s2s_values[src]) for key in array_keys: arr = data[key] sliced = arr[mask] @@ -199,6 +232,11 @@ def _merge_arrays_by_split( for key, arr_list in acc.items(): save_kwargs[key] = np.concatenate(arr_list) + if s2s_lens[sname]: + lens = np.concatenate(s2s_lens[sname]) + save_kwargs["seq_to_sig_offsets"] = csr_offsets_from_lens(lens) + save_kwargs.setdefault("seq_to_sig_values", np.empty(0, dtype=np.int32)) + n_chunks = len(save_kwargs.get("read_ids", np.array([]))) counts[sname] = n_chunks diff --git a/tests/bench_prepare_backends.py b/tests/bench_prepare_backends.py index 06926d0..8fa8dc9 100644 --- a/tests/bench_prepare_backends.py +++ b/tests/bench_prepare_backends.py @@ -382,10 +382,15 @@ def _run_comparison( save_chunks(rs_matched, rs_npz, compressed=False) # Skip auxiliary keys that use different windowing approaches: - # - seq_to_sig_maps: Python=searchsorted-based, Rust=kmer-window-based - # - sequences_with_kmer_context: follows from seq_to_sig_maps windowing + # - seq_to_sig_*: Python=searchsorted-based, Rust=kmer-window-based + # - sequences_with_kmer_context: follows from the same windowing # These only matter for signal_kmer encoding (not base_onehot). - aux_keys = {"seq_to_sig_maps", "sequences_with_kmer_context"} + aux_keys = { + "seq_to_sig_maps", + "seq_to_sig_values", + "seq_to_sig_offsets", + "sequences_with_kmer_context", + } npz_match, npz_diffs = compare_npz(py_npz, rs_npz, atol=chunk_atol, skip_keys=aux_keys) for d in npz_diffs: print(d) diff --git a/tests/test_backend_parity.py b/tests/test_backend_parity.py index 8586d7c..370e4e0 100644 --- a/tests/test_backend_parity.py +++ b/tests/test_backend_parity.py @@ -71,6 +71,8 @@ "source_groups", "reference_names", "seq_to_sig_maps", + "seq_to_sig_values", + "seq_to_sig_offsets", "sequences_with_kmer_context", "cl_values", "focus_signal_pos", @@ -172,7 +174,7 @@ def _run_both_backends(config: PrepareConfig, tmp_path) -> tuple[dict, dict]: def _assert_field_equal(name: str, a: np.ndarray, b: np.ndarray) -> None: assert a.shape == b.shape, f"{name}: shape {a.shape} != {b.shape}" if a.dtype == object: - # Variable-length rows (seq_to_sig_maps). Compare row by row so a + # Variable-length rows (legacy object members). Compare row by row so a # length difference names the row rather than raising from numpy. for i, (ra, rb) in enumerate(zip(a, b, strict=True)): ra, rb = np.asarray(ra), np.asarray(rb) diff --git a/tests/test_dataset_streaming.py b/tests/test_dataset_streaming.py new file mode 100644 index 0000000..193de24 --- /dev/null +++ b/tests/test_dataset_streaming.py @@ -0,0 +1,498 @@ +"""Tests for streaming chunk arrays out of an npz instead of loading them whole. + +``LeechDataset`` used to hold three copies of the corpus at once: the numpy +members from ``load_chunks``, one tensor per chunk, and the contiguous output +``torch.stack`` built while that list was still alive (#211). It now fills a +preallocated tensor from row blocks read straight off the npz. + +The hard gate here is parity: a dataset built from a path (streaming) must be +bit-identical to one built from pre-loaded chunks (eager), field by field, over +the option matrix that changes how chunks are prepared. +""" + +import numpy as np +import pytest +import torch + +from leech.chunking import ( + csr_gather_index, + iter_npz_row_blocks, + load_chunks, + npz_array_members, + npz_member_names, + save_chunks, +) +from leech.dataset import LeechDataset + +STORED_SIGNAL_LEN = 64 +FEAT_WIDTH = 13 +KMER_LEN = 11 +NUM_FEATURES = 4 + + +def make_chunks( + n: int = 12, + *, + signal_len: int = STORED_SIGNAL_LEN, + with_residual: bool = True, + with_maps: bool = True, + unlabeled: tuple[int, ...] = (), + seed: int = 0, +) -> list[dict]: + """Build a small synthetic corpus. + + Focus positions vary per chunk so an asymmetric crop lands inside the + stored signal for some rows and overhangs it (the zero-pad branch) for + others — the per-row gather the streaming path has to reproduce. + """ + rng = np.random.default_rng(seed) + chunks = [] + for i in range(n): + chunk = { + "signal": rng.standard_normal(signal_len).astype(np.float32), + "sequence": "".join(rng.choice(list("ACGT"), KMER_LEN)), + "dwell": rng.integers(1, 9, FEAT_WIDTH).astype(np.float32), + "features": rng.standard_normal((NUM_FEATURES, FEAT_WIDTH)).astype(np.float32), + "label": "charged" if i % 2 else "uncharged", + "label_int": None if i in unlabeled else i % 2, + "read_id": f"read_{i:03d}", + "base_idx": 100 + i, + "source_group": "Ala" if i % 3 else "Gly", + "reference_name": "tRNA-Ala-AGC", + "feature_start": -(FEAT_WIDTH // 2), + "feature_end": FEAT_WIDTH // 2, + "cl_value": i % 5, + # 8, 22, 36, 50, 8, ... : the first and last overhang a 20/24 crop. + "focus_signal_pos": 8 + 14 * (i % 4), + } + if with_residual: + chunk["signal_residual"] = rng.standard_normal(signal_len).astype(np.float32) + if with_maps: + n_bases = 9 + (i % 4) + chunk["seq_to_sig_map"] = np.sort( + rng.choice(min(signal_len, STORED_SIGNAL_LEN), n_bases + 1, replace=False) + ).astype(np.int64) + chunk["sequence_with_kmer_context"] = "".join(rng.choice(list("ACGT"), n_bases)) + chunks.append(chunk) + return chunks + + +def build(path, **kwargs) -> tuple[LeechDataset, LeechDataset]: + """Return (streaming, eager) datasets built from the same file and options.""" + streamed = LeechDataset(chunk_path=path, **kwargs) + eager = LeechDataset(chunks=load_chunks(path), **kwargs) + return streamed, eager + + +def assert_datasets_equal(streamed: LeechDataset, eager: LeechDataset) -> None: + """Every tensor the two datasets expose must match bit for bit.""" + assert len(streamed) == len(eager) + assert streamed.signal_channels == eager.signal_channels + assert streamed._effective_seq_encoding == eager._effective_seq_encoding + assert streamed._has_signal_residual == eager._has_signal_residual + + for name in ( + "_signals_tensor", + "_features_tensor", + "_labels_tensor", + "_encoded_seqs_tensor", + "_seq_ints_tensor", + "_seq_to_sig_tensor", + "_confound_labels_tensor", + "_cl_targets_tensor", + ): + a, b = getattr(streamed, name), getattr(eager, name) + assert (a is None) == (b is None), f"{name}: one path produced a tensor, the other did not" + if a is not None: + assert a.shape == b.shape, f"{name}: {a.shape} != {b.shape}" + assert torch.equal(a, b), f"{name}: values differ" + + # And the assembled samples, which is what training actually consumes. + for idx in (0, len(streamed) // 2, len(streamed) - 1): + left, right = streamed[idx], eager[idx] + assert left.keys() == right.keys() + for key in left: + assert torch.equal(left[key], right[key]), f"item {idx} field {key}" + + # Metadata the samplers and training config read off the chunk dicts. + for a, b in zip(streamed.chunks, eager.chunks, strict=True): + for key in ("read_id", "label_int", "source_group", "base_idx", "feature_start"): + assert a.get(key) == b.get(key), f"chunk metadata {key}" + + +class TestNpzStreaming: + """The row-block reader underneath the dataset.""" + + @pytest.mark.parametrize("compressed", [True, False]) + @pytest.mark.parametrize("block_rows", [1, 5, 12, 100]) + def test_roundtrip_matches_np_load(self, tmp_path, compressed, block_rows): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path, compressed=compressed) + + names = ["signals_flat", "features_flat", "labels_int"] + blocks: dict[str, list[np.ndarray]] = {name: [] for name in names} + starts = [] + for start, block in iter_npz_row_blocks(path, names, block_rows): + starts.append(start) + for name in names: + blocks[name].append(block[name].copy()) + + assert starts == list(range(0, 12, block_rows)) + with np.load(path) as data: + for name in names: + np.testing.assert_array_equal(np.concatenate(blocks[name]), data[name]) + + def test_blocks_are_recycled(self, tmp_path): + """Documented contract: the yielded arrays are views into one buffer.""" + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(8), path) + seen = [ + block["signals_flat"] for _, block in iter_npz_row_blocks(path, ["signals_flat"], 4) + ] + assert seen[0].base is seen[1].base + + def test_pickled_member_raises(self, tmp_path): + path = tmp_path / "objects.npz" + np.savez(path, ragged=np.array([np.arange(3), np.arange(5)], dtype=object)) + with pytest.raises(ValueError, match="not row-streamable"): + list(iter_npz_row_blocks(path, ["ragged"], 2)) + + def test_block_size_defaults_to_a_byte_budget(self, tmp_path): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(64, signal_len=4096), path, compressed=False) + rows = [ + len(block["signals_flat"]) + for _, block in iter_npz_row_blocks(path, ["signals_flat"], block_bytes=1 << 18) + ] + # 4096 float32 = 16 KiB per row, so 16 rows fit the 256 KiB budget. + assert rows[0] == 16 + assert sum(rows) == 64 + + def test_members_exclude_pickled(self, tmp_path): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(4), path) + members = npz_array_members(path) + assert members["signals_flat"][0] == (4, STORED_SIGNAL_LEN) + assert members["features_flat"][0] == (4, NUM_FEATURES, FEAT_WIDTH) + # CSR replaced the pickled object member entirely. + assert "seq_to_sig_maps" not in npz_member_names(path) + assert {"seq_to_sig_values", "seq_to_sig_offsets"} <= set(members) + + def test_csr_gather_index(self, tmp_path): + path = tmp_path / "chunks.npz" + chunks = make_chunks(6) + save_chunks(chunks, path) + with np.load(path) as data: + values, offsets = data["seq_to_sig_values"], data["seq_to_sig_offsets"] + rows = np.array([1, 3, 5]) + lens, _col, src = csr_gather_index(offsets, rows) + np.testing.assert_array_equal(lens, [len(chunks[i]["seq_to_sig_map"]) for i in rows]) + np.testing.assert_array_equal( + values[src], np.concatenate([chunks[i]["seq_to_sig_map"] for i in rows]) + ) + + +class TestStreamingParity: + """Streaming and eager construction must agree exactly.""" + + @pytest.mark.parametrize("signal_mode", ["both", "signal", "residual"]) + def test_signal_modes(self, tmp_path, signal_mode): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path) + assert_datasets_equal( + *build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + signal_mode=signal_mode, + seq_encoding="base_onehot", + ) + ) + + def test_asymmetric_crop(self, tmp_path): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path) + assert_datasets_equal( + *build( + path, + signal_len=44, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + left_context=20, + right_context=24, + seq_encoding="base_onehot", + ) + ) + + @pytest.mark.parametrize("model_type", ["ConvLSTMDwell", "ConvLSTMBase"]) + def test_feature_and_non_feature_models(self, tmp_path, model_type): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path) + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type=model_type, + seq_encoding="base_onehot", + ) + assert_datasets_equal(streamed, eager) + if model_type == "ConvLSTMBase": + # Nothing reads the features, so that member is never decompressed. + assert "features" not in streamed._array_stream.members + + @pytest.mark.parametrize("dwell_offset", [0, 1]) + def test_dwell_offset(self, tmp_path, dwell_offset): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path) + assert_datasets_equal( + *build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + dwell_offset=dwell_offset, + seq_encoding="base_onehot", + ) + ) + + @pytest.mark.parametrize("compressed", [True, False]) + @pytest.mark.parametrize("left_right", [None, (20, 24)]) + def test_signal_kmer_encoding(self, tmp_path, compressed, left_right): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path, compressed=compressed) + kwargs = { + "signal_len": STORED_SIGNAL_LEN if left_right is None else 44, + "kmer_len": KMER_LEN, + "model_type": "ConvLSTMDwell", + "seq_encoding": "signal_kmer", + "signal_kmer_context": (2, 2), + } + if left_right is not None: + kwargs["left_context"], kwargs["right_context"] = left_right + streamed, eager = build(path, **kwargs) + assert streamed._effective_seq_encoding == "signal_kmer" + assert_datasets_equal(streamed, eager) + + def test_signal_kmer_falls_back_without_maps(self, tmp_path): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(8, with_maps=False), path) + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="signal_kmer", + ) + assert streamed._effective_seq_encoding == "base_onehot" + assert_datasets_equal(streamed, eager) + + def test_no_residual_channel(self, tmp_path): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(10, with_residual=False), path) + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + signal_mode="both", + seq_encoding="base_onehot", + ) + assert streamed.signal_channels == 1 + assert_datasets_equal(streamed, eager) + + def test_cl_regression_targets(self, tmp_path): + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(10), path) + assert_datasets_equal( + *build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + cl_regression=True, + ) + ) + + def test_dwell_template_channels(self, tmp_path): + table = tmp_path / "templates.tsv" + rows = ["aa\tposition\tdwell_mean"] + for aa in ("Ala", "Gly"): + for pos in range(-6, 7): + rows.append(f"{aa}\t{pos}\t{4.0 + pos * 0.1}") + table.write_text("\n".join(rows) + "\n") + + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(10), path) + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + dwell_template_table=table, + ) + assert streamed._features_tensor.shape[1] == NUM_FEATURES + 2 + assert_datasets_equal(streamed, eager) + + +class TestRowAlignment: + """Rows dropped by label filtering must not shift the stream.""" + + @pytest.mark.parametrize( + "unlabeled", + [(0,), (5,), (11,), (0, 1, 6, 11), tuple(range(1, 12, 2))], + ) + def test_unlabeled_rows_are_skipped_not_shifted(self, tmp_path, unlabeled): + path = tmp_path / "chunks.npz" + chunks = make_chunks(12, unlabeled=unlabeled) + save_chunks(chunks, path) + + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + ) + assert len(streamed) == 12 - len(unlabeled) + assert_datasets_equal(streamed, eager) + + # Independently of the eager path: chunk i must carry the signal of the + # source chunk with the same read_id, not of some earlier row. + by_read = {c["read_id"]: c for c in chunks} + for i, chunk in enumerate(streamed.chunks): + expected = by_read[chunk["read_id"]]["signal"] + np.testing.assert_array_equal(streamed._signals_tensor[i, 0].numpy(), expected) + + +class TestLegacyFormats: + """Old corpora keep working — they simply do not stream.""" + + def _write_object_format(self, path, chunks): + """Write the pre-flat-array npz layout, pickled object members and all.""" + np.savez_compressed( + path, + signals=np.array([c["signal"] for c in chunks], dtype=object), + sequences=np.array([c["sequence"] for c in chunks], dtype=str), + dwells=np.array([c["dwell"] for c in chunks], dtype=object), + features=np.array([c["features"] for c in chunks], dtype=object), + labels=np.array([c["label"] for c in chunks], dtype=str), + labels_int=np.array([c["label_int"] for c in chunks], dtype=np.int64), + read_ids=np.array([c["read_id"] for c in chunks], dtype=str), + base_indices=np.array([c["base_idx"] for c in chunks], dtype=np.int64), + feature_starts=np.array([c["feature_start"] for c in chunks], dtype=np.int64), + feature_ends=np.array([c["feature_end"] for c in chunks], dtype=np.int64), + source_groups=np.array([c["source_group"] for c in chunks], dtype=str), + reference_names=np.array([c["reference_name"] for c in chunks], dtype=str), + seq_to_sig_maps=np.array([c["seq_to_sig_map"] for c in chunks], dtype=object), + sequences_with_kmer_context=np.array( + [c["sequence_with_kmer_context"] for c in chunks], dtype=str + ), + cl_values=np.array([c["cl_value"] for c in chunks], dtype=np.int16), + focus_signal_pos=np.array([c["focus_signal_pos"] for c in chunks], dtype=np.int64), + ) + + def test_object_array_corpus_takes_eager_path(self, tmp_path): + path = tmp_path / "legacy.npz" + self._write_object_format(path, make_chunks(8)) + + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + ) + assert streamed._array_stream is None # not streamable, still correct + assert_datasets_equal(streamed, eager) + + def test_legacy_maps_match_csr_maps(self, tmp_path): + """A pickled seq_to_sig_maps member yields the same tensor as CSR.""" + chunks = make_chunks(8) + legacy, modern = tmp_path / "legacy.npz", tmp_path / "modern.npz" + self._write_object_format(legacy, chunks) + save_chunks(chunks, modern) + + kwargs = { + "signal_len": 44, + "kmer_len": KMER_LEN, + "model_type": "ConvLSTMDwell", + "seq_encoding": "signal_kmer", + "signal_kmer_context": (2, 2), + "left_context": 20, + "right_context": 24, + } + from_legacy = LeechDataset(chunk_path=legacy, **kwargs) + from_csr = LeechDataset(chunk_path=modern, **kwargs) + assert from_csr._effective_seq_encoding == "signal_kmer" + assert torch.equal(from_legacy._seq_to_sig_tensor, from_csr._seq_to_sig_tensor) + assert torch.equal(from_legacy._seq_ints_tensor, from_csr._seq_ints_tensor) + + +class TestPeakMemory: + """Guards on the two copies #211 removed.""" + + def test_init_does_not_stack(self, tmp_path, monkeypatch): + """Preallocate-and-fill, not torch.stack — the stack doubled the peak.""" + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path) + + def fail(*args, **kwargs): + raise AssertionError("torch.stack allocates a second copy of the corpus") + + monkeypatch.setattr(torch, "stack", fail) + dataset = LeechDataset( + chunk_path=path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + ) + assert dataset._signals_tensor.shape[0] == 12 + + def test_streaming_peak_does_not_scale_with_corpus(self, tmp_path): + """The point of streaming: the array copy is bounded by the block size. + + tracemalloc sees numpy and Python allocations but not torch, so this + measures exactly the thing #211 was about — whether the npz members + get materialised alongside the tensors built from them. Doubling the + corpus must not double what the load holds. + """ + import tracemalloc + + kwargs = { + "signal_len": 2048, + "kmer_len": KMER_LEN, + "model_type": "ConvLSTMDwell", + "seq_encoding": "base_onehot", + } + + def peaks(n_chunks): + path = tmp_path / f"chunks_{n_chunks}.npz" + # Wide signals so the members, not the chunk dicts, dominate. + save_chunks(make_chunks(n_chunks, signal_len=2048), path, compressed=False) + measured = {} + for name, build_one in ( + ("eager", lambda: LeechDataset(chunks=load_chunks(path), **kwargs)), + ("streamed", lambda: LeechDataset(chunk_path=path, **kwargs)), + ): + tracemalloc.start() + dataset = build_one() + _, measured[name] = tracemalloc.get_traced_memory() + tracemalloc.stop() + del dataset + return measured + + small, large = peaks(750), peaks(1500) + + eager_growth = large["eager"] - small["eager"] + streamed_growth = large["streamed"] - small["streamed"] + assert eager_growth > 0, "test corpus too small to measure growth" + assert streamed_growth < eager_growth * 0.25, ( + f"streaming peak grew {streamed_growth / 1e6:.1f} MB when the corpus " + f"doubled, against {eager_growth / 1e6:.1f} MB eager — the arrays are " + f"still being materialised" + ) + assert large["streamed"] < large["eager"] From fac64a87f3086792873fcf918c7f42fcbb63a20f Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Mon, 24 Aug 2026 20:17:32 -0600 Subject: [PATCH 2/5] perf(dataset): fill output tensors in batches; document #211 in the changelog torch.stack(..., out=) writes a run of rows at C speed with no transient, where per-row assignment paid a few microseconds of dispatch each -- seconds over a multi-million-chunk corpus. Streamed rows are copied out of the block buffer so a staged batch cannot outlive the block it came from. Measured on a 300k-chunk corpus (1.8 GB npz): peak RSS 6.19 -> 2.36 GB and load time 21.5 -> 19.7 s against main, with bit-identical tensors. --- CHANGELOG.md | 34 +++++++++++++++++ src/leech/dataset.py | 67 +++++++++++++++++++++++++-------- tests/test_dataset_streaming.py | 29 ++++++++++---- 3 files changed, 107 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c85d753..d323a27 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,40 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Fixed + +- **`LeechDataset` no longer holds three copies of the corpus while it loads** + (#211). `load_chunks` read every npz member, the tensorize loop built one + tensor per chunk from them, and `torch.stack` allocated the whole contiguous + output while that list was still alive — a 41 GB npz peaked at 116 GB and + hit the 120 GB cgroup limit before epoch 1. The fields now fill a + preallocated tensor in bounded batches (`torch.stack(..., out=)`), and the + arrays are read from the npz in row blocks rather than materialised, so the + numpy source is never resident alongside the tensors built from it. Measured + on a 300k-chunk corpus (1.8 GB npz, 1.6 GB of output tensors): peak RSS + 6.19 GB -> 2.36 GB, load time 21.5 s -> 19.7 s, tensors bit-identical. +- **Only the members a run consumes are decompressed.** `signal_residuals_flat` + is skipped for `--signal-mode signal`, `features_flat` for models without a + feature branch, and the base-to-signal maps unless `--seq-encoding + signal_kmer` asks for them — up to 20 GB of decompression that used to + happen on every load regardless. +- **`load_chunks`'s docstring no longer claims the data is memory-mapped.** + `np.load` never maps a zip member, compressed or not; it is always a full + read, which is what made this path look lazy when it was not. + +### Changed + +- **`seq_to_sig_maps` is stored as `seq_to_sig_values` + `seq_to_sig_offsets`** + (CSR: row `i` is `values[offsets[i]:offsets[i+1]]`) instead of a pickled + object array. The old member cost one Python ndarray per chunk to unpickle + and could not be read in row blocks. `load_chunks`, `data merge` and the + dataset still read the legacy member, so existing corpora stay valid — but a + file written by this version and read by leech <= 0.6.7 has no + `seq_to_sig_maps`, so a `signal_kmer` run on that older version falls back to + `base_onehot` (with the warning it already emits). + ## [0.6.7] - 2026-08-24 Promotes `0.6.7-rc.1` unchanged — no commits landed between the two tags. The diff --git a/src/leech/dataset.py b/src/leech/dataset.py index b4426c7..46a2115 100644 --- a/src/leech/dataset.py +++ b/src/leech/dataset.py @@ -204,11 +204,18 @@ class _TensorFill: list, which ``__getitem__`` still handles. """ + #: Rows staged before one bulk write. Assigning row by row costs a few + #: microseconds of dispatch each — seconds over a multi-million-chunk + #: corpus — while ``torch.stack(..., out=)`` writes a run at C speed with + #: no transient. Small enough that what it pins is irrelevant. + _BATCH_ROWS = 256 + def __init__(self, name: str, capacity: int): self.name = name self.capacity = capacity self._tensor: torch.Tensor | None = None self._items: list[torch.Tensor] = [] + self._pending: list[torch.Tensor] = [] self._n = 0 def append(self, tensor: torch.Tensor) -> None: @@ -217,19 +224,36 @@ def append(self, tensor: torch.Tensor) -> None: return if self._tensor is None: self._tensor = torch.empty((self.capacity, *tensor.shape), dtype=tensor.dtype) - elif tuple(tensor.shape) != tuple(self._tensor.shape[1:]): - logger.warning( - "%s shapes differ (%s vs %s), falling back to list access", - self.name, - tuple(tensor.shape), - tuple(self._tensor.shape[1:]), - ) - self._items = [self._tensor[i].clone() for i in range(self._n)] - self._tensor = None - self._items.append(tensor) + elif ( + tuple(tensor.shape) != tuple(self._tensor.shape[1:]) + or tensor.dtype != self._tensor.dtype + ): + self._degrade(tensor) + return + self._pending.append(tensor) + if len(self._pending) >= self._BATCH_ROWS: + self._flush() + + def _flush(self) -> None: + if not self._pending: return - self._tensor[self._n] = tensor - self._n += 1 + rows = len(self._pending) + torch.stack(self._pending, out=self._tensor[self._n : self._n + rows]) + self._n += rows + self._pending.clear() + + def _degrade(self, tensor: torch.Tensor) -> None: + """Fall back to a list of per-chunk tensors, keeping what was filled.""" + logger.warning( + "%s shapes differ (%s vs %s), falling back to list access", + self.name, + tuple(tensor.shape), + tuple(self._tensor.shape[1:]), + ) + self._flush() + self._items = [self._tensor[i].clone() for i in range(self._n)] + self._tensor = None + self._items.append(tensor) def finish(self) -> tuple[torch.Tensor | None, list[torch.Tensor]]: """Return ``(stacked_tensor, fallback_list)``; exactly one is populated.""" @@ -237,6 +261,7 @@ def finish(self) -> tuple[torch.Tensor | None, list[torch.Tensor]]: # Nothing appended (the caller uses a different tensor) or shapes # differed — both cases the old _try_stack signalled with None. return None, self._items + self._flush() return self._tensor[: self._n], [] @@ -299,16 +324,19 @@ def build( def __iter__(self): """Yield one dict of arrays per kept row, in ``LeechDataset.chunks`` order. - The arrays are views into recycled block buffers — valid until the next - iteration, which is all the tensorize loop needs since it copies each - row into the output tensor. + Rows are copied out of the block buffer, which the next block read + overwrites. Without the copy the caller would have to consume each row + before the next one arrives — and it does not: preparing a chunk can + return a tensor that aliases its input (a contiguous crop is a view), + and those are staged in batches. Copying a few KB per row is cheaper + than the alternatives and removes the lifetime question entirely. """ if self.keep is None: raise RuntimeError("_ArrayStream.keep must be set before iterating") for start, blocks in iter_npz_row_blocks(self.path, list(self.members.values())): rows = len(next(iter(blocks.values()))) for j in np.nonzero(self.keep[start : start + rows])[0]: - yield {field: blocks[member][j] for field, member in self.members.items()} + yield {field: blocks[member][j].copy() for field, member in self.members.items()} def _expand_seq_to_sig_csr( @@ -701,6 +729,13 @@ def __init__( if key in chunk: chunk[key] = None + # Same reasoning for the streaming bookkeeping: one row per chunk each, + # and a DataLoader that spawns workers pickles whatever is still here. + self._s2s_csr = None + self._s2s_rows = None + if self._array_stream is not None: + self._array_stream.keep = None + # Precompute per-channel feature stds for feature noise augmentation. # Reuse the already-stacked features tensor when available. self._feature_stds: torch.Tensor | None = None diff --git a/tests/test_dataset_streaming.py b/tests/test_dataset_streaming.py index 193de24..93666b5 100644 --- a/tests/test_dataset_streaming.py +++ b/tests/test_dataset_streaming.py @@ -434,15 +434,26 @@ def test_legacy_maps_match_csr_maps(self, tmp_path): class TestPeakMemory: """Guards on the two copies #211 removed.""" - def test_init_does_not_stack(self, tmp_path, monkeypatch): - """Preallocate-and-fill, not torch.stack — the stack doubled the peak.""" + def test_init_never_allocates_a_second_copy(self, tmp_path, monkeypatch): + """Every stack during init writes into the preallocated output. + + ``torch.stack(items)`` allocates the whole result while ``items`` is + still alive — that was the third copy in #211. Writing through ``out=`` + in bounded batches is what replaced it. + """ + from leech.dataset import _TensorFill + path = tmp_path / "chunks.npz" - save_chunks(make_chunks(12), path) + save_chunks(make_chunks(600), path) + + calls = [] + real_stack = torch.stack - def fail(*args, **kwargs): - raise AssertionError("torch.stack allocates a second copy of the corpus") + def record(tensors, *args, **kwargs): + calls.append((len(tensors), kwargs.get("out") is not None)) + return real_stack(tensors, *args, **kwargs) - monkeypatch.setattr(torch, "stack", fail) + monkeypatch.setattr(torch, "stack", record) dataset = LeechDataset( chunk_path=path, signal_len=STORED_SIGNAL_LEN, @@ -450,7 +461,11 @@ def fail(*args, **kwargs): model_type="ConvLSTMDwell", seq_encoding="base_onehot", ) - assert dataset._signals_tensor.shape[0] == 12 + assert dataset._signals_tensor.shape[0] == 600 + assert calls, "expected the fill to stack in batches" + allocating = [n for n, has_out in calls if not has_out] + assert not allocating, f"{len(allocating)} stacks allocated instead of writing to out=" + assert max(n for n, _ in calls) <= _TensorFill._BATCH_ROWS def test_streaming_peak_does_not_scale_with_corpus(self, tmp_path): """The point of streaming: the array copy is bounded by the block size. From d0d3ea6e1277893a1c79af5feeb3103905abbe75 Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Mon, 24 Aug 2026 20:20:48 -0600 Subject: [PATCH 3/5] docs: cover the npz row-block reader in the data-prep API page test: cover the shape-mismatch fallback in _TensorFill --- docs/api/data_prep.md | 21 +++++++++++++++++++++ tests/test_dataset_streaming.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/docs/api/data_prep.md b/docs/api/data_prep.md index 73ccc0c..5e0760e 100644 --- a/docs/api/data_prep.md +++ b/docs/api/data_prep.md @@ -101,6 +101,27 @@ Save and load training chunks. show_root_heading: true show_source: false +### Streaming Chunk Arrays + +Read the per-chunk arrays a row block at a time instead of materialising whole +members — how `LeechDataset` builds its tensors without holding a second copy +of the corpus. + +::: leech.chunking.serialization.npz_array_members + options: + show_root_heading: true + show_source: false + +::: leech.chunking.serialization.iter_npz_row_blocks + options: + show_root_heading: true + show_source: false + +::: leech.chunking.serialization.load_seq_to_sig_csr + options: + show_root_heading: true + show_source: false + ## Preparation Module (`leech.preparation`) ### Sequential Preparation diff --git a/tests/test_dataset_streaming.py b/tests/test_dataset_streaming.py index 93666b5..59105aa 100644 --- a/tests/test_dataset_streaming.py +++ b/tests/test_dataset_streaming.py @@ -367,6 +367,37 @@ def test_unlabeled_rows_are_skipped_not_shifted(self, tmp_path, unlabeled): np.testing.assert_array_equal(streamed._signals_tensor[i, 0].numpy(), expected) +class TestShapeMismatchFallback: + """A field whose per-chunk shapes disagree still works, via list access.""" + + def test_ragged_sequences_fall_back_to_a_list(self, tmp_path, caplog): + chunks = make_chunks(10) + for i, chunk in enumerate(chunks): # 11, 10, 11, 10, ... bases + chunk["sequence"] = chunk["sequence"][: KMER_LEN - (i % 2)] + + path = tmp_path / "chunks.npz" + save_chunks(chunks, path) + with caplog.at_level("WARNING", logger="leech.dataset"): + streamed, eager = build( + path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + ) + + assert "shapes differ" in caplog.text + assert streamed._encoded_seqs_tensor is None + assert len(streamed._encoded_seqs) == len(chunks) + # The rows filled before the mismatch survive it. + for i, chunk in enumerate(chunks): + assert streamed._encoded_seqs[i].shape == (4, len(chunk["sequence"])) + assert torch.equal(streamed._encoded_seqs[i], eager._encoded_seqs[i]) + # Everything else still fills a contiguous tensor. + assert streamed._signals_tensor is not None + assert torch.equal(streamed._signals_tensor, eager._signals_tensor) + + class TestLegacyFormats: """Old corpora keep working — they simply do not stream.""" From ebdc94887c70fda48ef1c430390067d2d8238e94 Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Mon, 24 Aug 2026 20:42:57 -0600 Subject: [PATCH 4/5] perf(dataset): hold chunk metadata in columns, not a dict per chunk The per-chunk dicts measured 780 bytes each -- 5.2 GB for the 6.7M-chunk corpus in #211 -- to hold a handful of small integers and a few hundred distinct strings. ChunkTable keeps the npz's own arrays as columns and materialises a row view only when something asks for a chunk: 112 B/chunk, measured, with the same mapping interface every consumer already uses. Text is packed to fixed-width bytes (numpy stores 2.20 GB, load 19.7 -> 19.0 s. --- CHANGELOG.md | 6 + docs/api/data_prep.md | 16 ++ src/leech/chunking/__init__.py | 3 + src/leech/chunking/table.py | 320 ++++++++++++++++++++++++++++++++ src/leech/dataset.py | 75 ++++---- tests/test_chunk_table.py | 243 ++++++++++++++++++++++++ tests/test_dataset_streaming.py | 39 +++- 7 files changed, 666 insertions(+), 36 deletions(-) create mode 100644 src/leech/chunking/table.py create mode 100644 tests/test_chunk_table.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d323a27..78fd7e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 feature branch, and the base-to-signal maps unless `--seq-encoding signal_kmer` asks for them — up to 20 GB of decompression that used to happen on every load regardless. +- **Chunk metadata is stored as columns, not a dict per chunk** (#211). The + dicts measured 780 bytes each — 5.2 GB for a 6.7M-chunk corpus — holding a + handful of small integers and a few hundred distinct strings. `ChunkTable` + keeps the npz's own arrays (text packed to bytes, integers narrowed) and + hands out a row view on demand: 112 B/chunk measured, with no conversion + transient, and `dataset.chunks` still reads as a sequence of mappings. - **`load_chunks`'s docstring no longer claims the data is memory-mapped.** `np.load` never maps a zip member, compressed or not; it is always a full read, which is what made this path look lazy when it was not. diff --git a/docs/api/data_prep.md b/docs/api/data_prep.md index 5e0760e..6fc9e73 100644 --- a/docs/api/data_prep.md +++ b/docs/api/data_prep.md @@ -122,6 +122,22 @@ of the corpus. show_root_heading: true show_source: false +### Chunk Metadata Table + +Per-chunk metadata as columns, read as a sequence of mappings — what +`LeechDataset.chunks` holds when the corpus is loaded from a path. + +::: leech.chunking.table.ChunkTable + options: + show_root_heading: true + show_source: false + members: + - from_npz + - select + - values + - require_values + - nbytes + ## Preparation Module (`leech.preparation`) ### Sequential Preparation diff --git a/src/leech/chunking/__init__.py b/src/leech/chunking/__init__.py index 0338e85..495d2b1 100644 --- a/src/leech/chunking/__init__.py +++ b/src/leech/chunking/__init__.py @@ -25,6 +25,7 @@ npz_member_names, save_chunks, ) +from leech.chunking.table import ChunkRow, ChunkTable __all__ = [ # Extraction @@ -42,6 +43,8 @@ "load_seq_to_sig_csr", "npz_array_members", "npz_member_names", + "ChunkTable", + "ChunkRow", "csr_from_object_rows", "csr_gather_index", "csr_offsets_from_lens", diff --git a/src/leech/chunking/table.py b/src/leech/chunking/table.py new file mode 100644 index 0000000..7b21c60 --- /dev/null +++ b/src/leech/chunking/table.py @@ -0,0 +1,320 @@ +"""Columnar chunk metadata. + +:func:`~leech.chunking.serialization.load_chunks` builds one dict per chunk. +Measured at 780 bytes each — with the repeated strings already interned — that +is 5.2 GB for a 6.7M-chunk corpus, spent on dicts whose values are a handful of +small integers and a few hundred distinct strings (#211). + +:class:`ChunkTable` keeps the npz's own arrays as columns and materialises a row +view only when something asks for a chunk, which costs roughly 130 bytes per +chunk for the same access patterns. Text is held as fixed-width bytes rather +than as `` (chunk field, is the negative value a missing marker) +_INT_FIELDS: tuple[tuple[str, str, bool], ...] = ( + ("base_indices", "base_idx", False), + ("labels_int", "label_int", True), + ("feature_starts", "feature_start", False), + ("feature_ends", "feature_end", False), + ("cl_values", "cl_value", True), + ("focus_signal_pos", "focus_signal_pos", False), +) + +#: npz member -> (chunk field, does the empty string mean None) +_TEXT_FIELDS: tuple[tuple[str, str, bool], ...] = ( + ("sequences", "sequence", False), + ("read_ids", "read_id", False), + ("labels", "label", True), + ("source_groups", "source_group", True), + # An absent reference name reads back as "" rather than None, matching + # load_chunks — callers concatenate it into report keys. + ("reference_names", "reference_name", False), + ("sequences_with_kmer_context", "sequence_with_kmer_context", True), +) + + +class _Column: + """One metadata field across every chunk.""" + + __slots__ = () + + def value(self, index: int): + raise NotImplementedError + + def take(self, rows: np.ndarray) -> "_Column": + raise NotImplementedError + + @property + def raw(self) -> np.ndarray | None: + """The underlying array, without the missing-value translation.""" + return None + + +class _IntColumn(_Column): + __slots__ = ("values", "none_if_negative") + + def __init__(self, values: np.ndarray, none_if_negative: bool): + self.values = values + self.none_if_negative = none_if_negative + + def value(self, index: int): + number = int(self.values[index]) + if self.none_if_negative and number < 0: + return None + return number + + def take(self, rows: np.ndarray) -> "_IntColumn": + return _IntColumn(self.values[rows], self.none_if_negative) + + @property + def raw(self) -> np.ndarray: + return self.values + + +class _TextColumn(_Column): + """Text as fixed-width bytes, decoded on access. + + Decoding allocates a str per read, but a transient one: the dicts this + replaces held every one of them for the life of the dataset. + """ + + __slots__ = ("values", "none_if_empty", "_encoded") + + def __init__(self, values: np.ndarray, none_if_empty: bool): + self.values = values + self.none_if_empty = none_if_empty + self._encoded = values.dtype.kind == "S" + + def value(self, index: int): + text = self.values[index] + text = text.decode() if self._encoded else str(text) + if self.none_if_empty and not text: + return None + return text + + def take(self, rows: np.ndarray) -> "_TextColumn": + return _TextColumn(self.values[rows], self.none_if_empty) + + @property + def raw(self) -> np.ndarray: + return self.values + + +class _ConstColumn(_Column): + """A field every chunk shares — currently only ``cl_value = None``.""" + + __slots__ = ("constant",) + + def __init__(self, constant): + self.constant = constant + + def value(self, index: int): + return self.constant + + def take(self, rows: np.ndarray) -> "_ConstColumn": + return self + + +def _narrow_ints(values: np.ndarray) -> np.ndarray: + """Cast to the smallest signed dtype that holds the column's range.""" + if values.size == 0: + return values.astype(np.int8) + low, high = int(values.min()), int(values.max()) + for dtype in (np.int8, np.int16, np.int32): + info = np.iinfo(dtype) + if info.min <= low and high <= info.max: + return values.astype(dtype) + return values.astype(np.int64) + + +def _read_text_member(input_path: Path, member: str, shape, dtype) -> np.ndarray: + """Read a text member as fixed-width bytes, four times smaller than `` bool: + # Mapping's default answers this by reading the value, which for a + # column means an index and a conversion. Presence is a property of the + # table, not of the row. + return key in self._table.columns + + def __iter__(self): + return iter(self._table.columns) + + def __len__(self) -> int: + return len(self._table.columns) + + def __repr__(self) -> str: + return f"ChunkRow({dict(self)!r})" + + +class ChunkTable(Sequence): + """Chunk metadata as columns, presented as a sequence of read-only mappings. + + Indexing yields a :class:`ChunkRow`; iterating yields one per chunk. Rows + are views, so they cost nothing to keep out of and nothing is shared with + the caller to mutate. + + Examples: + >>> table = ChunkTable.from_npz(Path("chunks.npz")) # doctest: +SKIP + >>> table[0]["label_int"] # doctest: +SKIP + 1 + >>> table.values("label_int") # raw column, for vectorized tallies + ... # doctest: +SKIP + array([1, 0, 1, ...], dtype=int8) + """ + + __slots__ = ("columns", "_n") + + def __init__(self, columns: dict[str, _Column], n_chunks: int): + self.columns = columns + self._n = n_chunks + + @classmethod + def from_npz( + cls, + input_path: Path, + *, + skip: Collection[str] = (), + ) -> "ChunkTable": + """Read a corpus's metadata members — never its per-chunk arrays. + + Args: + input_path: Path to .npz file. + skip: Chunk field names to leave out. Text the run will not read is + worth skipping: ``sequence_with_kmer_context`` is 56 bytes a + chunk that only ``signal_kmer`` encoding touches. + + Returns: + A table with one row per chunk in file order. + """ + skip = set(skip) + columns: dict[str, _Column] = {} + text_members = npz_array_members(input_path) + + # allow_pickle stays off: every member read here is a plain array, and + # the one pickled member (legacy seq_to_sig_maps) is never metadata. + with np.load(input_path, allow_pickle=False) as data: + n_chunks = len(data["labels_int"]) + has_feature_window = "feature_starts" in data + present = set(data.files) + + for member, field, none_if_negative in _INT_FIELDS: + if field in skip or member not in present: + continue + columns[field] = _IntColumn(_narrow_ints(data[member]), none_if_negative) + + # Old corpora carry dwell_margin_lefts instead of the signed window. + if not has_feature_window and "dwell_margin_lefts" in present: + columns["dwell_margin_left"] = _IntColumn( + _narrow_ints(data["dwell_margin_lefts"]), False + ) + + for member, field, none_if_empty in _TEXT_FIELDS: + if field in skip or member not in present: + continue + shape, dtype = text_members[member] + columns[field] = _TextColumn( + _read_text_member(input_path, member, shape, dtype), none_if_empty + ) + + # load_chunks reports cl_value as None when the corpus predates it, and + # callers read it unguarded. + if "cl_value" not in columns and "cl_value" not in skip: + columns["cl_value"] = _ConstColumn(None) + + return cls(columns, n_chunks) + + def select(self, mask: np.ndarray) -> "ChunkTable": + """Return a table holding only the rows where ``mask`` is True.""" + rows = np.nonzero(mask)[0] + return ChunkTable({name: col.take(rows) for name, col in self.columns.items()}, len(rows)) + + def values(self, field: str) -> np.ndarray | None: + """The raw column for ``field``, or None if the table lacks it. + + Raw means as stored: text comes back as bytes, and missing integers as + their negative sentinel rather than as None. Use it to tally a field + across a whole corpus without building a row per chunk. + """ + column = self.columns.get(field) + return None if column is None else column.raw + + def require_values(self, field: str) -> np.ndarray: + """The raw column for ``field``, raising if the corpus lacks it.""" + column = self.columns.get(field) + if column is None or column.raw is None: + raise KeyError(f"chunk metadata has no column '{field}'") + return column.raw + + def nbytes(self) -> int: + """Total bytes held by the columns.""" + return sum(col.raw.nbytes for col in self.columns.values() if col.raw is not None) + + def __len__(self) -> int: + return self._n + + def __getitem__(self, index: int) -> ChunkRow: + if isinstance(index, slice): + raise TypeError("ChunkTable indexes one chunk at a time; use select() for subsets") + if index < 0: + index += self._n + if not 0 <= index < self._n: + raise IndexError(f"chunk index out of range: {index}") + return ChunkRow(self, index) + + def __iter__(self): + for index in range(self._n): + yield ChunkRow(self, index) + + def __repr__(self) -> str: + return f"ChunkTable({self._n} chunks, fields={sorted(self.columns)})" diff --git a/src/leech/dataset.py b/src/leech/dataset.py index 46a2115..fbc0d5e 100644 --- a/src/leech/dataset.py +++ b/src/leech/dataset.py @@ -45,6 +45,7 @@ from torch.utils.data import Dataset from leech.chunking import ( + ChunkTable, iter_npz_row_blocks, load_chunks, load_seq_to_sig_csr, @@ -496,11 +497,12 @@ def __init__( else: raise ValueError("Either chunk_path or chunks must be provided") - # Filter chunks with valid numeric labels (label_int). The streaming - # path has already applied this, and recorded the same mask so npz rows - # stay aligned with self.chunks — a mismatch here would pair signals - # with the wrong labels silently. - self.chunks = [c for c in self.chunks if c["label_int"] is not None] + # Filter chunks with valid numeric labels (label_int). The columnar + # path applied this at load, recording the same mask so npz rows stay + # aligned with self.chunks — a mismatch here would pair signals with + # the wrong labels silently. + if not isinstance(self.chunks, ChunkTable): + self.chunks = [c for c in self.chunks if c["label_int"] is not None] if len(self.chunks) == 0: raise ValueError(f"No valid chunks found{f' in {chunk_path}' if chunk_path else ''}") @@ -576,11 +578,17 @@ def __init__( # or a row-block stream over the npz. Both yield the same field names, # so the loop below does not care which. array_iter = iter(self._array_stream) if self._array_stream is not None else None + # The one metadata field read for every chunk on the default encoding. + # Reading the column directly hands `_encode_sequence` the bytes it + # wants and skips a row view per chunk; dicts have nothing to hoist. + sequence_column = ( + self.chunks.values("sequence") if isinstance(self.chunks, ChunkTable) else None + ) stream_dwell_width = ( self._array_stream.dwell_width if self._array_stream is not None else None ) - for chunk in self.chunks: + for row, chunk in enumerate(self.chunks): arrays = chunk if array_iter is None else next(array_iter) if self._effective_seq_encoding == "signal_kmer": @@ -601,7 +609,8 @@ def __init__( self._seq_to_sig.append(s2s) else: # Pre-encode sequence (vectorized, no Python loop) - fill_encoded_seqs.append(self._encode_sequence(chunk["sequence"])) + sequence = chunk["sequence"] if sequence_column is None else sequence_column[row] + fill_encoded_seqs.append(self._encode_sequence(sequence)) # Pre-create label tensor: long for multi-class, float for binary if self._multiclass: @@ -717,17 +726,18 @@ def __init__( # so we keep self.chunks alive but null out the per-chunk arrays. # Without this, each chunk dict keeps a ~50 KB numpy view alive and # the same COW blowup hits during DataLoader fork. - for chunk in self.chunks: - for key in ( - "signal", - "signal_residual", - "dwell", - "features", - "seq_to_sig_map", - "sequence_with_kmer_context", - ): - if key in chunk: - chunk[key] = None + if not isinstance(self.chunks, ChunkTable): + for chunk in self.chunks: + for key in ( + "signal", + "signal_residual", + "dwell", + "features", + "seq_to_sig_map", + "sequence_with_kmer_context", + ): + if key in chunk: + chunk[key] = None # Same reasoning for the streaming bookkeeping: one row per chunk each, # and a DataLoader that spawns workers pickles whatever is still here. @@ -778,19 +788,17 @@ def _load_from_path(self, chunk_path: Path, *, signal_mode: str, seq_encoding: s self.chunks = load_chunks(chunk_path) return - # Nothing downstream reads the raw arrays off the chunk dicts — the - # stream supplies signal/residual/features, the dwell width comes from - # the member header, and the base-to-signal maps are expanded from CSR - # only when signal_kmer needs them. So never read them into the dicts. - defer = {"signal", "signal_residual", "dwell", "features", "seq_to_sig_map"} - if seq_encoding != "signal_kmer": - defer.add("sequence_with_kmer_context") - raw = load_chunks(chunk_path, defer=defer) + # Metadata goes into columns rather than a dict per chunk: the stream + # supplies signal/residual/features, the dwell width comes from the + # member header, and the base-to-signal maps are expanded from CSR only + # when signal_kmer needs them, so no per-chunk array is read at all. + skip = () if seq_encoding == "signal_kmer" else ("sequence_with_kmer_context",) + table = ChunkTable.from_npz(chunk_path, skip=skip) - keep = np.fromiter((c["label_int"] is not None for c in raw), dtype=bool, count=len(raw)) + keep = table.require_values("label_int") >= 0 stream.keep = keep self._array_stream = stream - self.chunks = [c for c, k in zip(raw, keep, strict=True) if k] + self.chunks = table if keep.all() else table.select(keep) if seq_encoding == "signal_kmer": self._s2s_csr = load_seq_to_sig_csr(chunk_path) self._s2s_rows = np.nonzero(keep)[0] @@ -1073,19 +1081,22 @@ def _apply_feature_noise(self, features: torch.Tensor) -> torch.Tensor: return features @staticmethod - def _encode_sequence(sequence: str) -> torch.Tensor: + def _encode_sequence(sequence: str | bytes) -> torch.Tensor: """Vectorized one-hot encoding of a DNA sequence. Uses a pre-built ASCII lookup table instead of a Python for-loop. Args: - sequence: DNA sequence string (A, C, G, T, N) + sequence: DNA sequence (A, C, G, T, N), str or ASCII bytes. The + columnar metadata store holds sequences as bytes, and the + lookup wants bytes, so accept them and skip the round trip. Returns: One-hot encoded tensor of shape (4, len(sequence)) """ - indices = _BASE_MAP[np.frombuffer(sequence.encode(), dtype=np.uint8)] - encoded = np.zeros((4, len(sequence)), dtype=np.float32) + raw = sequence if isinstance(sequence, bytes) else sequence.encode() + indices = _BASE_MAP[np.frombuffer(raw, dtype=np.uint8)] + encoded = np.zeros((4, len(raw)), dtype=np.float32) valid = indices < 4 encoded[indices[valid], np.where(valid)[0]] = 1.0 return torch.from_numpy(encoded) diff --git a/tests/test_chunk_table.py b/tests/test_chunk_table.py new file mode 100644 index 0000000..6f91d88 --- /dev/null +++ b/tests/test_chunk_table.py @@ -0,0 +1,243 @@ +"""Tests for the columnar chunk-metadata store. + +``ChunkTable`` replaces one dict per chunk — 780 bytes each, measured — with +columns plus a row view. Everything downstream reads chunks as mappings, so the +contract under test is that a row is indistinguishable from the dict +``load_chunks`` would have built, including which keys are absent and which +values come back as None. +""" + +import numpy as np +import pytest + +from leech.chunking import ChunkTable, load_chunks, save_chunks + +READ_ID_LEN = 36 # UUID-shaped, as dorado writes them + + +def make_chunks(n: int = 8, *, labelled: bool = True) -> list[dict]: + rng = np.random.default_rng(7) + chunks = [] + for i in range(n): + chunks.append( + { + "signal": rng.standard_normal(32).astype(np.float32), + "sequence": "".join(rng.choice(list("ACGT"), 11)), + "dwell": rng.integers(1, 9, 13).astype(np.float32), + "features": rng.standard_normal((4, 13)).astype(np.float32), + # Empty label and -1 label_int are how load_chunks spells None. + "label": ("charged" if i % 2 else "uncharged") if labelled else "", + "label_int": (i % 2) if labelled else None, + "read_id": f"{i:0{READ_ID_LEN}d}", + "base_idx": 1000 + i, + "source_group": "Ala" if i % 3 else "", + "reference_name": "tRNA-Ala-AGC-1-1" if i % 2 else "", + "feature_start": -6, + "feature_end": 6, + "cl_value": i if i % 4 else -1, + "focus_signal_pos": 8 + i, + "seq_to_sig_map": np.arange(12, dtype=np.int64), + "sequence_with_kmer_context": "ACGT" * 3, + } + ) + return chunks + + +@pytest.fixture +def corpus(tmp_path): + chunks = make_chunks(8) + path = tmp_path / "chunks.npz" + save_chunks(chunks, path) + return path, chunks + + +class TestRowsMatchLoadChunks: + """A row must read exactly like the dict load_chunks builds.""" + + def test_every_metadata_field_matches(self, corpus): + path, _ = corpus + table = ChunkTable.from_npz(path) + dicts = load_chunks(path) + + assert len(table) == len(dicts) + array_fields = { + "signal", + "signal_residual", + "dwell", + "features", + "seq_to_sig_map", + } + for row, chunk in zip(table, dicts, strict=True): + for key in set(chunk) - array_fields: + assert row.get(key) == chunk[key], key + + def test_missing_values_read_back_as_none(self, corpus): + path, chunks = corpus + table = ChunkTable.from_npz(path) + + # Sentinels: "" for text, -1 for ints — except reference_name, which + # load_chunks reports as "" rather than None. + assert table[0]["source_group"] is None # i % 3 == 0 wrote "" + assert table[1]["source_group"] == "Ala" + assert table[0]["reference_name"] == "" + assert table[1]["reference_name"] == "tRNA-Ala-AGC-1-1" + assert table[4]["cl_value"] is None # i % 4 == 0 wrote -1 + assert table[5]["cl_value"] == 5 + + def test_unlabelled_chunks_read_as_none(self, tmp_path): + path = tmp_path / "unlabelled.npz" + save_chunks(make_chunks(4, labelled=False), path) + table = ChunkTable.from_npz(path) + assert [row["label_int"] for row in table] == [None] * 4 + assert [row["label"] for row in table] == [None] * 4 + + +class TestMappingContract: + def test_absent_field_is_absent_not_none(self, corpus): + path, _ = corpus + table = ChunkTable.from_npz(path) + row = table[0] + + # `"feature_start" in chunk` is how training.py picks a feature-window + # convention, so absence has to mean absence. + assert "feature_start" in row + assert "dwell_margin_left" not in row + assert row.get("dwell_margin_left") is None + with pytest.raises(KeyError): + row["dwell_margin_left"] + + def test_row_is_a_mapping(self, corpus): + path, _ = corpus + table = ChunkTable.from_npz(path) + row = table[2] + assert dict(row) == {key: row[key] for key in row} + assert set(row.keys()) == set(row) + assert len(row) == len(list(row)) + assert row["read_id"] in repr(row) + + def test_indexing(self, corpus): + path, chunks = corpus + table = ChunkTable.from_npz(path) + assert table[-1]["read_id"] == chunks[-1]["read_id"] + with pytest.raises(IndexError): + table[len(chunks)] + with pytest.raises(TypeError, match="one chunk at a time"): + table[0:2] + + def test_skip_leaves_a_field_out(self, corpus): + path, _ = corpus + full = ChunkTable.from_npz(path) + trimmed = ChunkTable.from_npz(path, skip=("sequence_with_kmer_context",)) + assert "sequence_with_kmer_context" in full[0] + assert "sequence_with_kmer_context" not in trimmed[0] + assert trimmed[0]["sequence"] == full[0]["sequence"] + + +class TestSelect: + def test_select_keeps_the_masked_rows(self, corpus): + path, chunks = corpus + table = ChunkTable.from_npz(path) + mask = np.array([i % 3 == 0 for i in range(len(chunks))]) + + selected = table.select(mask) + expected = [c for c, keep in zip(chunks, mask, strict=True) if keep] + assert len(selected) == len(expected) + for row, chunk in zip(selected, expected, strict=True): + assert row["read_id"] == chunk["read_id"] + assert row["base_idx"] == chunk["base_idx"] + assert row["cl_value"] == (chunk["cl_value"] if chunk["cl_value"] >= 0 else None) + + def test_select_none(self, corpus): + path, chunks = corpus + table = ChunkTable.from_npz(path) + empty = table.select(np.zeros(len(chunks), dtype=bool)) + assert len(empty) == 0 + assert list(empty) == [] + + +class TestStorage: + """The two mechanisms that make it small, asserted directly.""" + + def test_integers_are_narrowed_without_changing_values(self, corpus): + path, chunks = corpus + table = ChunkTable.from_npz(path) + + assert table.values("label_int").dtype == np.int8 + assert table.values("feature_start").dtype == np.int8 + assert table.values("base_idx").dtype == np.int16 # 1000..1007 + # Negative sentinels survive narrowing. + assert table.values("cl_value").min() == -1 + assert [row["base_idx"] for row in table] == [c["base_idx"] for c in chunks] + + def test_text_is_stored_as_bytes(self, corpus): + path, _ = corpus + table = ChunkTable.from_npz(path) + assert table.values("read_id").dtype.kind == "S" + assert table.values("sequence").dtype.kind == "S" + assert isinstance(table[0]["read_id"], str) + + def test_non_ascii_text_still_loads(self, tmp_path): + chunks = make_chunks(4) + for chunk in chunks: + chunk["source_group"] = "Ångström" + path = tmp_path / "unicode.npz" + save_chunks(chunks, path) + + table = ChunkTable.from_npz(path) + assert table.values("source_group").dtype.kind == "U" # kept, not dropped + assert table[0]["source_group"] == "Ångström" + + def test_columns_are_far_smaller_than_dicts(self, tmp_path): + n = 2000 + path = tmp_path / "big.npz" + save_chunks(make_chunks(n), path) + + table = ChunkTable.from_npz(path) + per_chunk = table.nbytes() / n + # The dicts this replaces measured 780 B/chunk on a corpus with these + # same fields; the columns hold the same values in ~100. + assert per_chunk < 200, f"{per_chunk:.0f} B/chunk" + + def test_values_returns_none_for_an_absent_field(self, corpus): + path, _ = corpus + assert ChunkTable.from_npz(path).values("dwell_margin_left") is None + + +class TestLegacyCorpora: + def test_dwell_margin_left_replaces_the_signed_window(self, tmp_path): + """Pre-feature_start corpora carry dwell_margin_lefts instead.""" + chunks = make_chunks(4) + path = tmp_path / "legacy.npz" + np.savez( + path, + sequences=np.array([c["sequence"] for c in chunks], dtype=str), + labels=np.array([c["label"] for c in chunks], dtype=str), + labels_int=np.array([c["label_int"] for c in chunks], dtype=np.int64), + read_ids=np.array([c["read_id"] for c in chunks], dtype=str), + base_indices=np.array([c["base_idx"] for c in chunks], dtype=np.int64), + dwell_margin_lefts=np.full(len(chunks), 4, dtype=np.int64), + ) + + table = ChunkTable.from_npz(path) + row = table[0] + assert row["dwell_margin_left"] == 4 + assert "feature_start" not in row + # cl_value predates this format but callers read it unguarded. + assert row["cl_value"] is None + + +class TestPickling: + """A DataLoader that spawns workers pickles the dataset, table and all.""" + + def test_round_trips(self, corpus): + import pickle + + path, chunks = corpus + table = ChunkTable.from_npz(path) + restored = pickle.loads(pickle.dumps(table)) + + assert len(restored) == len(table) + for row, original in zip(restored, chunks, strict=True): + assert row["read_id"] == original["read_id"] + assert row["label_int"] == original["label_int"] + assert row["sequence"] == original["sequence"] diff --git a/tests/test_dataset_streaming.py b/tests/test_dataset_streaming.py index 59105aa..258eb47 100644 --- a/tests/test_dataset_streaming.py +++ b/tests/test_dataset_streaming.py @@ -15,6 +15,7 @@ import torch from leech.chunking import ( + ChunkTable, csr_gather_index, iter_npz_row_blocks, load_chunks, @@ -24,6 +25,19 @@ ) from leech.dataset import LeechDataset +#: Per-chunk arrays: streamed or deferred, so they are not metadata and the +#: columnar store does not carry them. +ARRAY_FIELDS = frozenset( + { + "signal", + "signal_residual", + "dwell", + "features", + "seq_to_sig_map", + "sequence_with_kmer_context", + } +) + STORED_SIGNAL_LEN = 64 FEAT_WIDTH = 13 KMER_LEN = 11 @@ -114,10 +128,12 @@ def assert_datasets_equal(streamed: LeechDataset, eager: LeechDataset) -> None: for key in left: assert torch.equal(left[key], right[key]), f"item {idx} field {key}" - # Metadata the samplers and training config read off the chunk dicts. - for a, b in zip(streamed.chunks, eager.chunks, strict=True): - for key in ("read_id", "label_int", "source_group", "base_idx", "feature_start"): - assert a.get(key) == b.get(key), f"chunk metadata {key}" + # Every metadata field, however it is stored: the columnar path must be + # indistinguishable from the chunk dicts to samplers, the label tally and + # the training-config introspection that read them. + for i, (a, b) in enumerate(zip(streamed.chunks, eager.chunks, strict=True)): + for key in (set(a) | set(b)) - ARRAY_FIELDS: + assert a.get(key) == b.get(key), f"chunk {i} metadata {key}" class TestNpzStreaming: @@ -210,6 +226,21 @@ def test_signal_modes(self, tmp_path, signal_mode): ) ) + def test_streaming_path_stores_metadata_columnar(self, tmp_path): + """The chunk dicts are the last per-chunk Python object at this scale.""" + path = tmp_path / "chunks.npz" + save_chunks(make_chunks(12), path) + streamed = LeechDataset( + chunk_path=path, + signal_len=STORED_SIGNAL_LEN, + kmer_len=KMER_LEN, + model_type="ConvLSTMDwell", + seq_encoding="base_onehot", + ) + assert isinstance(streamed.chunks, ChunkTable) + # Text the run never reads is not loaded at all. + assert "sequence_with_kmer_context" not in streamed.chunks[0] + def test_asymmetric_crop(self, tmp_path): path = tmp_path / "chunks.npz" save_chunks(make_chunks(12), path) From dddcca81ded83d9da4ee159e6789c8ee50b1d6e8 Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Mon, 24 Aug 2026 20:46:31 -0600 Subject: [PATCH 5/5] docs(adr): record the memory-mapped corpus decision (#211 item 4) PR #212's CSR change removed the blocker -- every member of an uncompressed corpus is now fixed-shape and mappable -- so the question is whether to map, not whether we can. Measured on an rna node: sequential reads run at 724 MB/s against the 58 MB/s the loader consumes, but a cold mmap page fault costs 0.60 ms, giving 1,668 chunks/s per thread where the loader wants 10,900. A mapped corpus runs at RAM speed while the page cache holds it and falls off a cliff when it does not, which is the state mapping exists to serve. Records sequential shard streaming with a shuffle buffer as the direction -- including the sampler constraints it has to solve -- and fp16 signals as the cheaper win to try first. --- .../adr/0006-memory-mapped-chunk-corpora.md | 176 ++++++++++++++++++ dev-notes/adr/README.md | 4 + 2 files changed, 180 insertions(+) create mode 100644 dev-notes/adr/0006-memory-mapped-chunk-corpora.md diff --git a/dev-notes/adr/0006-memory-mapped-chunk-corpora.md b/dev-notes/adr/0006-memory-mapped-chunk-corpora.md new file mode 100644 index 0000000..00f674d --- /dev/null +++ b/dev-notes/adr/0006-memory-mapped-chunk-corpora.md @@ -0,0 +1,176 @@ +# ADR 0006: Memory-Mapped Chunk Corpora + +**Status:** Accepted (mapping rejected for the training read path; sequential +shard streaming adopted as the direction) + +**Date:** 2026-08-24 + +## Context + +[#211](https://github.com/rnabioco/leech/issues/211) listed four fixes for +`LeechDataset` peaking at 116 GB on a 41 GB corpus. Three are done (PR #212): +the `torch.stack` transient is gone, the arrays are read in row blocks instead +of materialised, and the metadata is columnar. What remains is the output +tensors themselves — for that corpus, signals `(6668328, 2, 540)` fp32 = +28.8 GB and features `(6668328, 12, 21)` fp32 = 6.7 GB, since +`TCNDwellResidualLN` is a wide-feature model and keeps the full 21-base window. +Roughly 36 GB, resident for the whole run. + +Item 4 of the issue proposed mapping the file instead: under `--no-compress` +every member is written *stored*, so each could be `np.memmap`'d at its zip +offset and `__getitem__` could slice from disk. The pickled `seq_to_sig_maps` +member blocked that; PR #212 replaced it with a CSR pair, so **every member of a +freshly written uncompressed corpus is now fixed-shape, non-object, and +mappable in place**. The blocker is gone. The question left is whether mapping +is the right thing to do, which the issue explicitly flagged as a design call +rather than a patch. + +### What the training loop actually asks for + +Per chunk the loader touches three members at unrelated offsets — signal, +residual, features — about 5.3 KB in total. The reported production loader +wants ~10,900 chunks/s. + +| Access pattern | What that costs | +|---|---| +| Sequential | 10,900 × 5.3 KB = **58 MB/s** | +| Shuffled (mapped) | ≥3 page faults per chunk, more when a 2160-byte row straddles a page: **40–65k IOPS** | + +Those are wildly different asks of the same file, and shuffling is the only +difference. + +### Measured on this cluster + +An `rna` compute node, against a 1.83 GB corpus on BeeGFS. Client page cache +dropped with `posix_fadvise(POSIX_FADV_DONTNEED)` before each cold run: + +| | measured | +|---|---| +| sequential read, cold | **724 MB/s** (single stream) | +| random 4K `pread`, 1 thread | 4,712 IOPS | +| random 4K `pread`, 16 threads | 152,000 IOPS aggregate | +| **mmap page fault, 1 thread, cold** | **1,668 rows/s** (0.60 ms each) | +| mmap page fault, warm | 342,000 rows/s | + +**Treat the random numbers as upper bounds.** `POSIX_FADV_DONTNEED` drops the +*client's* cache; it cannot drop the BeeGFS servers', and the node was idle. +The production figure in #211 — ~113 random IOPS per thread against a loaded +filesystem — is three orders of magnitude below the idle-node number here, and +is the one to plan against. That spread *is* the finding: random-read +performance on this storage is a property of who else is using it. + +The gap between the two mmap rows is the whole story. A mapped corpus runs at +RAM speed while the page cache holds it, and at 1,668 chunks/s per thread when +it does not — and not holding it in RAM is the entire reason for mapping. + +Node-local staging does not rescue it here: the compute nodes' local disk is +rotational (`lsblk ROTA=1`), and `/dev/shm` is a 377 GB tmpfs, which is RAM +with extra steps — staging there to "save memory" spends exactly the memory it +claims to save. + +## Decision + +**Do not map the corpus for the shuffled training read path.** Keep the tensors +resident, and pursue the two directions below instead. + +### 1. Sequential shard streaming with a shuffle buffer (the real answer) + +The arithmetic above says the corpus streams at **12× the rate training +consumes it** (724 MB/s against 58 MB/s), and a full epoch of pure sequential +I/O over 41 GB is 57 seconds. Only the shuffle makes disk-resident training +hard, and a shuffle buffer is the standard trade: read shards sequentially, +shuffle within a window of B chunks, yield from the window. + +- RAM for the corpus drops from ~36 GB to `B × 5.3 KB` — 1.1 GB at B = 200,000 + chunks (3% of the corpus), plus prefetch. +- `iter_npz_row_blocks` from PR #212 is already the primitive: sequential row + blocks with a byte budget, one block resident. +- Randomness is approximate rather than exact. Interleaving several shards into + the window recovers most of it. + +**Constraints that must be designed around, not discovered later:** + +- `--balance-groups` and `--oversample-minority` build a `WeightedRandomSampler` + over the whole index, which assumes global random access. The counts they + need are now free (`ChunkTable.values("source_group")` is a column), but the + sampling itself has to move inside the buffer — per-window weighted draws or + rejection sampling against the global rate. +- `--feature-noise-scale` derives per-channel stds from the whole feature + tensor. Streaming needs a two-pass or running estimate. +- Workers must own disjoint shards, and the shuffle must be seedable, or runs + stop being reproducible. +- Validation and eval read in fixed order and need none of this. + +### 2. Halve the dominant term first (cheaper, orthogonal, no loader change) + +Signals are 28.8 GB of the 36 GB. They arrive as 16-bit ADC counts and are +stored normalised; **fp16 holds normalised signal to about 1e-3 at the +magnitudes involved**, which is far below the noise the model is being asked to +see through. Storing the signal tensor as fp16 and casting per batch takes the +resident corpus from ~36 GB to ~21 GB for a few lines and no change to how the +data is read. This should be tried before anything more ambitious, and needs an +accuracy check against a trained model, not just a memory measurement. + +### Where mapping *is* the right tool + +For the sequential passes — feature-std computation, label and group tallies, +anything that touches every chunk once in order — mapping is fine, and so is +the block reader, which is already there and does not depend on cache state. +Neither is worth adding for those. + +## Consequences + +**Positive** + +- No one spends a sprint building a mapped loader whose throughput is a + function of cluster load, and which degrades to 1,668 chunks/s exactly when + RAM is scarce enough to have wanted it. +- The direction that does work reuses the row-block reader already in the tree. +- fp16 gives most of the remaining win for a fraction of the effort. + +**Negative** + +- The 36 GB of tensors stays for now. Corpora much beyond the current one still + need a bigger memory request until the streaming loader exists. +- Shuffle-buffer training is a real project, with the sampler constraints + above, not a patch. + +**Neutral** + +- The mapping option stays open: PR #212's CSR change means an uncompressed + corpus is mappable today, so if the storage picture changes — node-local + NVMe, or a corpus small enough to stay in page cache — the experiment is a + short one. Nothing in this decision has to be undone to run it. + +## Alternatives Considered + +**Map the members and slice in `__getitem__` (issue item 4 as written).** +Rejected above: 1,668 chunks/s per thread cold on an idle node, ~113 IOPS per +thread under production load, against a loader wanting 10,900 chunks/s. + +**Stage to node-local disk, then map.** The usual fix for this problem, and it +is the reason the option is not dead in general — but these nodes have +rotational local disks. Revisit if the hardware changes. + +**`/dev/shm`.** tmpfs is RAM. Mapping from it does not reduce memory, it +relocates it out of the process's accounting where the scheduler can no longer +see it. + +**Keep compressed npz and map anyway.** Not possible: deflate members have no +byte-addressable layout. `--no-compress` is a precondition for any mapping +work, and costs disk (the reported corpus is 41 GB stored). + +**Convert to a format built for this (webdataset, FFCV, Arrow/Parquet).** +Would bring sharded sequential reads and a shuffle buffer with it rather than +hand-rolling them. Rejected for now only because it is a format migration on +top of a corpus that is already produced, validated and parity-tested; the +shuffle-buffer design above can be built against the existing npz and revisited +if it grows past what one module should own. + +## Notes + +Measurements: `srun -p rna -c 16`, 1.83 GB synthetic corpus with the shape of +the one in #211 (540-sample signals, residual channel, 12×21 features). +Reproduce with `os.posix_fadvise(fd, 0, 0, POSIX_FADV_DONTNEED)` before each +cold pass, and read the caveat above about server-side caching before quoting +any of the random-access numbers. diff --git a/dev-notes/adr/README.md b/dev-notes/adr/README.md index 44cb666..2675ed0 100644 --- a/dev-notes/adr/README.md +++ b/dev-notes/adr/README.md @@ -22,6 +22,10 @@ ADRs capture important architectural decisions along with their context and cons - [ADR 0003: Model Component Abstraction](0003-model-component-abstraction.md) - Reusable neural network branch components - [ADR 0004: Inference Wrapper Pattern](0004-inference-wrapper-pattern.md) - Unified forward pass interface +### Data Loading + +- [ADR 0006: Memory-Mapped Chunk Corpora](0006-memory-mapped-chunk-corpora.md) - Why the training read path stays resident, and what to do instead (#211) + ## ADR Format Each ADR follows this structure: