feat: prefetching dataloader (buffered + double-buffered modes)

[d-laub]

Summary

Adds two new modes on Dataset.to_dataloader() that coarsen gvl fetches into large chunks and slice them into mini-batches on the consumer side. gvl's Dataset[r_idx, s_idx] throughput scales with fetch size (internal multithreading amortizes overhead), so one big call + slice outperforms many small per-batch calls.

• mode="buffered" — main process fetches one chunk per call (sized to buffer_bytes, default 2 GB), slices via slice_chunk.
• mode="double_buffered" — subprocess producer fills one of two POSIX shm slots while consumer drains the other; ping-pong via multiprocessing.Event. Refill latency is hidden behind the drain.

Footprint is computed exactly per (region, sample) via a new Dataset._output_bytes_per_instance method that walks haplotype_lengths, n_variants, and allele offset tables — no Zipf-style worst-case bounds, so slot capacity is tight.

API

loader = ds.to_dataloader(
    batch_size=32,
    mode="double_buffered",     # or "buffered", or None
    buffer_bytes=2 * 1024**3,
    copy=True,                  # default; zero-copy opt-out
    heartbeat_seconds=60.0,     # double_buffered: max wait per chunk before liveness check
)

Preconditions (raise ValueError at construction):

• with_seqs in {"haplotypes", "annotated"} requires deterministic=True.
• Spliced datasets are rejected.
• num_workers > 0 is rejected.
• A single mini-batch whose exact footprint exceeds slot capacity raises with remediation knobs (batch_size↓, buffer_bytes↑).
• mode="double_buffered" rejects datasets with non-default insertion-fill strategies (can't be serialized to the subprocess).

New files

• python/genvarloader/_chunked.py — ChunkPlanner, slice_chunk
• python/genvarloader/_buffered_loader.py — mode="buffered" factory
• python/genvarloader/_shm_layout.py — hand-rolled header pack/unpack for dense, Ragged, and RaggedVariants
• python/genvarloader/_producer.py — subprocess entrypoint with broad schema replay
• python/genvarloader/_double_buffered_loader.py — mode="double_buffered" factory

Modified

• _torch.py:get_dataloader — mode/buffer_bytes/copy/heartbeat_seconds dispatch
• _dataset/_impl.py — new _output_bytes_per_instance method (all four with_seqs modes + tracks + var_fields)
• _dataset/_haps.py — new _allele_bytes_sum helper (O(|V|) via RaggedAlleles.offsets diff)
• skills/genvarloader/SKILL.md — documents the new args

Test plan

• Unit tests for _allele_bytes_sum (cross-checked via independent _get_variants materialization).
• Unit tests for _output_bytes_per_instance across all (with_seqs × with_tracks × var_fields) combinations supported by the dummy fixture.
• Pure-logic tests for ChunkPlanner (slot respect, mini-batch boundary preservation, oversized-batch rejection, peak chunk size).
• Parity tests: slice_chunk output matches dataset[r, s] for every output mode.
• Shm layout round-trip: dense, Ragged, multi-Ragged (AnnotatedHaps-style), RaggedVariants (real dataset output), cross-process.
• Producer subprocess: signals correctly, surfaces exceptions via queue.
• mode="buffered": parity with direct indexing for all output modes, rejects num_workers, oversized batch, non-deterministic haplotypes.
• mode="double_buffered": parity with mode="buffered", parity under rc_neg=False, producer-crash heartbeat path, shm cleanup on close (Linux-only).
• Existing test suite: 248 passed, no regressions.

Spec / plan

• Spec: docs/superpowers/specs/2026-05-28-prefetching-dataloader-design.md
• Plan: docs/superpowers/plans/2026-05-28-prefetching-dataloader-implementation.md

Out of scope (intentional)

• DDP / multi-GPU (rank-aware sampler + per-rank producer).
• Non-deterministic / random-variant-selection datasets for haplotypes/annotated.
• Spliced datasets.
• Auto mode selection — user picks; the loader does not guess.

🤖 Generated with Claude Code