fix(dataset): stop holding three copies of the corpus at load (#211)

[jayhesselberth]

Fixes #211.

LeechDataset.__init__ held the corpus three times over: the numpy members
load_chunks reads, one tensor per chunk built from them, and the contiguous
output torch.stack allocates while that list is still alive. The array
nulling at the end of __init__ is right but runs after all three exist, so it
never touched the peak. A 41 GB npz peaked at 116 GB and hit the 120 GB cap
before epoch 1.

What changed

Fill, don't stack. Each field fills a tensor allocated up front
(_TensorFill), writing runs of 256 rows with torch.stack(..., out=) — bulk
speed, no transient. Shape mismatches still degrade to the per-chunk list
__getitem__ already handles (now covered by a test; it wasn't before).

Stream the arrays. iter_npz_row_blocks walks signals_flat,
signal_residuals_flat and features_flat as sequential row blocks, and
load_chunks(..., defer=...) never reads them into the chunk dicts, so the
numpy source is never resident alongside the tensors built from it. Blocks are
sized by a byte budget rather than a row count, so a wide signal doesn't
silently allocate a huge buffer.

Read only what the run consumes. Residuals are skipped for --signal-mode signal, features for models without a feature branch, the base-to-signal maps
unless --seq-encoding signal_kmer asks. dwells_flat is never read at all:
_prepare_features only needed len(dwell), which the member header gives.

seq_to_sig_maps becomes CSR — seq_to_sig_values + seq_to_sig_offsets,
row i is values[offsets[i]:offsets[i+1]]. The pickled object member cost one
Python ndarray per chunk to unpickle and could not be read in row blocks.
data merge gathers those rows and rebuilds the offsets, normalizing legacy
inputs so a mixed-vintage merge writes one format.

Metadata goes in columns. The per-chunk dicts measured 780 bytes each —
5.2 GB for that corpus — to hold a handful of small integers and a few hundred
distinct strings. ChunkTable keeps the npz's own arrays as columns (text
packed to bytes, since numpy stores <U as UTF-32; integers narrowed to the
smallest dtype holding their range) and materialises a row view only when
something asks for a chunk: 112 B/chunk measured, with the conversion done
a block at a time so it has no transient either. A row is a read-only
Mapping, so chunk["label_int"], chunk.get("source_group") and
"feature_start" in chunk all work unchanged.

Also: load_chunks interns the repeated per-chunk strings, calibrate builds
its dataset from a path so it streams too, and the load_chunks docstring no
longer claims the data is memory-mapped — np.load never maps a zip member,
compressed or not, which is what made this path look lazy when it wasn't.

Measured

300k chunks, 1.8 GB npz, 1.6 GB of output tensors, on a rna compute node.
Tensors bit-identical across all three (checksummed):

	peak RSS	load
main (v0.6.7)	6.19 GB	21.5 s
this branch, pre-loaded chunks	4.22 GB	19.5 s
this branch, from chunk_path	2.20 GB	19.0 s

Per-chunk metadata, measured separately on the same corpus: 778 → 112 bytes.

62% off the peak, slightly faster. Scaled to the 450-context arm in the issue
that lands near 44 GB against the 116 GB that OOM'd — an extrapolation from a
synthetic corpus shaped like that one, since no real train.npz was available
here to measure.

Tests

tests/test_dataset_streaming.py (40 tests) and
tests/test_chunk_table.py (16). The gate is bit-identical parity
between the streaming and eager paths — every tensor, every assembled sample,
and the chunk metadata — across signal modes, asymmetric crops, feature and
non-feature models, dwell offsets, signal_kmer vs base_onehot, dwell
template channels, compressed and uncompressed files. Plus corpora with
label_int = -1 rows interleaved between labelled ones: row/chunk misalignment
there would pair signals with the wrong labels silently, so it gets its own
independent check by read_id.

Legacy corpora (object-array members, pickled seq_to_sig_maps) are covered
too — they take the eager path and produce identical tensors.

The columnar store is held to the same standard: every metadata field of every
chunk must read exactly as the dict load_chunks would have built, including
which keys are absent ("feature_start" in chunk decides a feature-window
convention in training.py) and which values come back as None.

Full suite: 1085 passed, 27 skipped. ruff, ty and the docs build clean.

Compat

Existing corpora read unchanged. The one edge is the other direction: a file
written by this version has no seq_to_sig_maps, so leech <= 0.6.7 reading it
falls back to base_onehot for a signal_kmer run — with the warning it
already emits, but a fallback nonetheless. Called out in the changelog.

Item 4, the memmap idea: ADR 0006

The CSR change means every member of an uncompressed corpus is now fixed-shape
and mappable in place, so the blocker is gone and the question is whether to
map. Measured on an rna node (dev-notes/adr/0006-memory-mapped-chunk-corpora.md
has the numbers and the caveats):

• sequential reads run at 724 MB/s against the 58 MB/s the loader
  consumes — 12× headroom;
• 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. Local
staging does not rescue it here either: these nodes' local disks are rotational
and /dev/shm is RAM with extra steps. So the ADR rejects mapping for the
shuffled read path and records what to do instead — sequential shard streaming
with a shuffle buffer, built on the row-block reader this PR adds, plus the
sampler constraints (--balance-groups, --oversample-minority and
--feature-noise-scale all assume global access) that design has to solve. It
also flags fp16 signals as the cheaper win to try first: signals are 28.8 GB of
the ~36 GB that remains resident.

Nothing here has to be undone to run the mapping experiment later if the
storage picture changes.