rnabioco · jayhesselberth · Aug 25, 2026 · Aug 25, 2026 · Aug 25, 2026 · Aug 25, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,46 @@ 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.
+- **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.
+
+### 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/dev-notes/adr/0006-memory-mapped-chunk-corpora.md 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
@@ -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:

diff --git a/docs/api/data_prep.md b/docs/api/data_prep.md
@@ -101,6 +101,43 @@ 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
+
+### 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/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
@@ -12,7 +12,20 @@
     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,
+)
+from leech.chunking.table import ChunkRow, ChunkTable
 
 __all__ = [
     # Extraction
@@ -25,4 +38,14 @@
     "save_chunks",
     "load_chunks",
     "get_chunk_statistics",
+    "DEFERRABLE_FIELDS",
+    "iter_npz_row_blocks",
+    "load_seq_to_sig_csr",
+    "npz_array_members",
+    "npz_member_names",
+    "ChunkTable",
+    "ChunkRow",
+    "csr_from_object_rows",
+    "csr_gather_index",
+    "csr_offsets_from_lens",
 ]