Eric Coissac
76cbd3a886
Replaced scattered direct metadata loading with centralized instance methods on `KmerPartition` to guarantee consistent error mapping and legacy recovery. Introduced `StorageKind`, `LayerContent`, and `EvidenceKind` enums alongside lightweight disk-probe methods that inspect file presence without opening heavy data structures. Updated callers across the index, partitioner, and phylo modules to use the new partition API, and added unit tests validating the introspection behavior.
352 lines
17 KiB
Rust
352 lines
17 KiB
Rust
//! Shared per-layer family traversal, used by every sibling-annex consumer
|
|
//! (`snp_pseudo_alignment`, `cardinality_tally`, `scan_family_pairs`,
|
|
//! `sibling_annex_stats`) — resolves each minorant family's per-genome
|
|
//! base-presence, with the same locality discipline as
|
|
//! `build_sibling_annex` (see `cache.rs`/`build.rs`'s docs): cross-partition
|
|
//! lookups are grouped by destination partition and resolved in one
|
|
//! contiguous sweep per partition, instead of one lookup jumping between
|
|
//! partitions in family order.
|
|
//!
|
|
//! Two concerns, kept on two different mechanisms because they need
|
|
//! opposite things:
|
|
//!
|
|
//! - **Generating** each batch's cross-partition queries is pure CPU (bit
|
|
//! tests, one hash per variant) — cheap, and safe to run for several
|
|
//! batches at once. It runs on an `obipipeline::throttle` + `make_pipe!`
|
|
//! pipeline, the same mechanism `build_layer_sibling_annex` uses, for the
|
|
//! same reason: it overlaps this CPU work with the *previous* batch's
|
|
//! resolution instead of leaving cores idle between batches.
|
|
//! - **Resolving** those queries against `PartitionCache` is I/O (mmap page
|
|
//! faults on a real index far larger than RAM) — one batch at a time,
|
|
//! `rayon`-parallel *across partitions* (like `build_sibling_annex`'s own
|
|
//! `outgoing.par_iter()`), never several batches concurrently. An earlier
|
|
//! version resolved each batch fully inside its own pipeline worker, so
|
|
//! `n_workers` batches were resolved concurrently — each one spreading its
|
|
//! queries thin across every partition of the layer at once. Sampling a
|
|
//! real run showed the fix hadn't helped at all: 16 threads "busy" in
|
|
//! `find_presence`, but mostly blocked on page faults (~1.7 GB/s pagein),
|
|
//! because 16 concurrent sweeps across the same partition space is exactly
|
|
//! the scattering the grouping was meant to prevent — just spread over
|
|
//! threads instead of over layers this time. Resolving one batch's worth
|
|
//! at a time, with each of `rayon`'s threads owning one partition
|
|
//! contiguously until that batch is done, keeps only one partition set
|
|
//! "hot" at once, restoring the locality the batching is for.
|
|
//!
|
|
//! `obipipeline` does not guarantee output order, so generated batches carry
|
|
//! their own starting family index and are replayed through a small reorder
|
|
//! buffer before resolution — bounded by the throttle's own concurrency
|
|
//! (`n_workers` batches), not by the layer's size.
|
|
//!
|
|
//! [`FAMILY_BATCH`] bounds a single batch's memory to `FAMILY_BATCH *
|
|
//! genome_count` bytes, and is chosen large enough that a batch still gives
|
|
//! each partition a decent number of queries to resolve in one go — too
|
|
//! small a batch starves that density regardless of how the resolution step
|
|
//! is parallelised. An earlier version materialised the whole layer's
|
|
//! families up front before resolving anything: O(layer's family count ×
|
|
//! genome count) memory, measured pushing a real run into heavy VM
|
|
//! compression/swap.
|
|
|
|
use std::collections::HashMap;
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicU8, Ordering};
|
|
|
|
use rayon::prelude::*;
|
|
|
|
use obikseq::CanonicalKmer;
|
|
use obilayeredmap::meta::PartitionMeta;
|
|
use obipipeline::{ThrottleGuard, throttle};
|
|
|
|
use obikindex::{OKIError, OKIResult};
|
|
use obikindex::KmerIndex;
|
|
|
|
use super::cache::{Mat, PartitionCache};
|
|
use super::helpers::central_base;
|
|
use super::iter::SiblingEntry;
|
|
use super::{olm_to_ok, FamilyMask, SiblingAnnex, ANNEX_FILE_NAME};
|
|
|
|
/// Families per batch — see the module docs for the memory-vs-per-partition-
|
|
/// density trade-off this picks a point on. At ~90 genomes and a few
|
|
/// hundred partitions, this keeps a batch's resolution memory in the low
|
|
/// tens of MB while still giving each partition on the order of a thousand
|
|
/// queries per batch to amortise against.
|
|
const FAMILY_BATCH: usize = 65536;
|
|
|
|
/// Restricts [`scan_layer_families`] to a subset of a layer's families,
|
|
/// identified by their iteration-order index — see
|
|
/// `DevDocMD/architecture/siblings.md`, "`--subsample`/`--shannon`", for why:
|
|
/// bounding which families actually pay the expensive cross-partition
|
|
/// resolution cost is what makes a sampled run cheap on an index far larger
|
|
/// than the sample itself. `All` (the default for every pre-existing
|
|
/// caller) keeps today's behaviour exactly.
|
|
pub(super) enum Selection<'a> {
|
|
All,
|
|
Some(&'a std::collections::HashSet<usize>),
|
|
}
|
|
|
|
/// Shared by the pipeline worker (via the `Arc`-cloned `Option<HashSet>`,
|
|
/// see below — a `Selection` itself doesn't cross the worker thread
|
|
/// boundary) and the final replay loop, so both sides agree on membership
|
|
/// without duplicating the match.
|
|
fn is_selected(selected: &Option<std::collections::HashSet<usize>>, family_idx: usize) -> bool {
|
|
match selected {
|
|
None => true,
|
|
Some(set) => set.contains(&family_idx),
|
|
}
|
|
}
|
|
|
|
/// Every (partition, layer) directory carrying a sibling annex, checked
|
|
/// up front so a missing one is reported before any real work starts.
|
|
/// Shared by every sibling-annex consumer's extension-trait impl (`alignment`,
|
|
/// `build`, `cardinality`, `distance`, `stats`) — a free function, not a
|
|
/// `KmerIndex` method, since it is crate-internal only and `KmerIndex` lives
|
|
/// in `obikindex`, a foreign crate from here (orphan rule).
|
|
pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult<Vec<PathBuf>> {
|
|
let n_parts = index.n_partitions();
|
|
let mut layer_dirs = Vec::new();
|
|
for part in 0..n_parts {
|
|
let index_dir = index.partition().index_dir(part);
|
|
if !index_dir.exists() {
|
|
continue;
|
|
}
|
|
let n_layers = index.partition().n_layers(part)?;
|
|
for l in 0..n_layers {
|
|
let this_layer_dir = index.partition().layer_dir(part, l);
|
|
let annex_path = this_layer_dir.join(ANNEX_FILE_NAME);
|
|
if !annex_path.exists() {
|
|
return Err(OKIError::InvalidInput(format!(
|
|
"no sibling annex at {} — run build_sibling_annex first",
|
|
annex_path.display()
|
|
)));
|
|
}
|
|
layer_dirs.push(this_layer_dir);
|
|
}
|
|
}
|
|
Ok(layer_dirs)
|
|
}
|
|
|
|
/// Read-only state shared (via `Arc`) across every pipeline worker
|
|
/// generating this layer's batches — opened once, not per batch. No `cache`
|
|
/// here: generation never touches the cross-partition cache, only this
|
|
/// layer's own already-open matrix. No separate MPHF/`slot_kmer` either —
|
|
/// `mat` (a `Layer<D>`) already bundles the MPHF, and each `SiblingEntry`
|
|
/// arrives with its kmer and mask already in hand from `iter_minorants_batch`.
|
|
struct LayerCtx {
|
|
mat: Mat,
|
|
n_parts: usize,
|
|
n_genomes: usize,
|
|
n_cols: usize,
|
|
k: usize,
|
|
}
|
|
|
|
struct SourceBatch {
|
|
start_family_idx: usize,
|
|
/// One entry per minorant family in this batch, straight from
|
|
/// `iter_minorants_batch` — `order` (iteration-order index, not an MPHF
|
|
/// slot; see `DevDocMD/architecture/siblings.md`), `kmer`, and `mask`
|
|
/// already carried through, no second annex read.
|
|
entries: Vec<SiblingEntry>,
|
|
_permit: ThrottleGuard,
|
|
}
|
|
|
|
/// One batch's generated work: local presence already resolved (straight
|
|
/// from this layer's own matrix, no lookup needed), cross-partition queries
|
|
/// collected but not yet resolved against the cache.
|
|
struct GeneratedBatch {
|
|
start_family_idx: usize,
|
|
masks: Vec<FamilyMask>,
|
|
/// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`.
|
|
genome_mask: Vec<u8>,
|
|
/// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base,
|
|
/// layer)` — `layer` is the annex-recorded destination layer
|
|
/// (`FamilyMask::family_members`'s `layer_value`, `0` if absent),
|
|
/// trustworthy only when `PartitionCache::fast_mode()` is true.
|
|
outgoing: Vec<Vec<(CanonicalKmer, usize, u8, u8)>>,
|
|
_permit: ThrottleGuard,
|
|
}
|
|
|
|
enum FamData {
|
|
Batch(SourceBatch),
|
|
Generated(GeneratedBatch),
|
|
}
|
|
|
|
/// Visits every minorant family of one layer, in iteration order, batched
|
|
/// [`FAMILY_BATCH`] at a time — see the module docs for why generation and
|
|
/// resolution use different concurrency. `on_family` is called once per
|
|
/// family, in iteration order, with that family's iteration-order index
|
|
/// (the same numbering [`Selection`] indices are drawn from), its own annex
|
|
/// mask, and its per-genome base-presence (`genome_mask[g]`: bit `b` set
|
|
/// iff genome `g` carries the member whose own canonical central base is
|
|
/// `b`), backed by a scratch buffer reused across every call — callers that
|
|
/// need to keep data past the call must copy it themselves.
|
|
pub(super) fn scan_layer_families(
|
|
layer_dir: &Path,
|
|
n_parts: usize,
|
|
n_genomes: usize,
|
|
with_counts: bool,
|
|
k: usize,
|
|
cache: &Arc<PartitionCache>,
|
|
selection: &Selection,
|
|
mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
|
|
) -> OKIResult<()> {
|
|
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
|
|
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
|
|
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?);
|
|
|
|
let mat = Mat::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?;
|
|
let n_cols = mat.n_cols().min(n_genomes);
|
|
|
|
let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k });
|
|
|
|
// Streamed straight from `iter_minorants_batch` (zips this layer's own
|
|
// `iter_kmers()` with the annex, both in iteration order — never an
|
|
// MPHF slot; see `DevDocMD/architecture/siblings.md`) — never collected
|
|
// into a `Vec` first: a layer can hold billions of k-mers, so
|
|
// materialising every minorant family up front is exactly the memory
|
|
// blowup an earlier version of this traversal was rewritten to avoid
|
|
// (see the module docs). `.scan()` computes each batch's starting
|
|
// family index lazily, mirroring what the eager `chunks()`+running
|
|
// `offset` used to do.
|
|
let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| {
|
|
let start = *offset;
|
|
*offset += batch.len();
|
|
Some((start, batch))
|
|
});
|
|
|
|
let n_workers = obisys::effective_parallelism();
|
|
let capacity = 4;
|
|
let throttled = throttle(batches, n_workers).map(|t| SourceBatch {
|
|
start_family_idx: t.item.0,
|
|
entries: t.item.1,
|
|
_permit: t.guard,
|
|
});
|
|
|
|
// `Selection::Some` borrows a `HashSet` whose lifetime doesn't span the
|
|
// pipeline's worker threads — cloned once into an `Arc` so every worker
|
|
// can share it cheaply instead of requiring `selection` itself to be
|
|
// `'static`. `None` means [`Selection::All`], checked with a plain
|
|
// `map_or` at each use instead of allocating an always-true set.
|
|
let selected: Arc<Option<std::collections::HashSet<usize>>> = Arc::new(match selection {
|
|
Selection::All => None,
|
|
Selection::Some(set) => Some((*set).clone()),
|
|
});
|
|
|
|
let worker_ctx = Arc::clone(&ctx);
|
|
let worker_selected = Arc::clone(&selected);
|
|
let pipe = obipipeline::make_pipe! {
|
|
FamData : SourceBatch => GeneratedBatch,
|
|
| {
|
|
move |batch: SourceBatch| -> GeneratedBatch {
|
|
let ctx = &worker_ctx;
|
|
let n = batch.entries.len();
|
|
let mut masks = Vec::with_capacity(n);
|
|
let mut bases = Vec::with_capacity(n);
|
|
let mut genome_mask = vec![0u8; n * ctx.n_genomes];
|
|
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8, u8)>> = (0..ctx.n_parts).map(|_| Vec::new()).collect();
|
|
|
|
// Pass 1: cheap, no matrix access — own base and this
|
|
// batch's cross-partition queries. Kmer and mask already in
|
|
// hand from `iter_minorants_batch`, no second annex read,
|
|
// no slot lookup needed for this pass. `family_members`
|
|
// (not a hand-rolled `central_canonical_neighbors` +
|
|
// `mask.has` loop) already filters to present members and
|
|
// hands back each one's annex-recorded layer alongside.
|
|
// Families outside `selected` (sampling only) still get a
|
|
// mask/base entry — pass 2 indexes uniformly by `i` — but
|
|
// never an outgoing query: that's the expensive part a
|
|
// sampled run exists to avoid paying for every family.
|
|
for (i, entry) in batch.entries.iter().enumerate() {
|
|
let (kmer, mask) = (entry.kmer, entry.mask);
|
|
masks.push(mask);
|
|
let base = central_base(kmer, ctx.k);
|
|
bases.push(base);
|
|
let family_idx = batch.start_family_idx + i;
|
|
if !is_selected(&worker_selected, family_idx) {
|
|
continue;
|
|
}
|
|
for (member, layer) in mask.family_members(kmer, ctx.k) {
|
|
if member == kmer {
|
|
continue; // local — resolved below straight from `mat`, no lookup
|
|
}
|
|
let b = central_base(member, ctx.k);
|
|
let dest = member.partition(ctx.n_parts);
|
|
outgoing[dest].push((member, i, b, layer.unwrap_or(0)));
|
|
}
|
|
}
|
|
|
|
// Pass 2: genome-major, not family-major — `mat` is stored
|
|
// one contiguous block per genome (column), slot as the
|
|
// offset within it (see `PartitionCache::find_presence_batch`'s
|
|
// docs for the full rationale). `fill_sub_matrix_carries`
|
|
// sorts the slots internally for a sequential mmap sweep per
|
|
// column, then restores this batch's order — no hand-rolled
|
|
// sort/genome-major loop needed here. The presence/count
|
|
// matrix is still MPHF-slot-indexed (unlike the annex), so
|
|
// each entry's kmer is mapped to its slot via `index_batch`
|
|
// — a pure MPHF lookup, no evidence check, since these are
|
|
// this layer's own kmers, known members by construction.
|
|
let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect();
|
|
let slots = ctx.mat.index_batch(&kmers);
|
|
let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect();
|
|
ctx.mat.fill_sub_matrix_carries(&slots, &mut carries);
|
|
for (g, col) in carries.iter().enumerate() {
|
|
for (i, &carries_it) in col.iter().enumerate() {
|
|
if carries_it {
|
|
genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
GeneratedBatch { start_family_idx: batch.start_family_idx, masks, genome_mask, outgoing, _permit: batch._permit }
|
|
}
|
|
} : Batch => Generated,
|
|
};
|
|
|
|
// Batches finish generation in whatever order their worker completes
|
|
// them, not submission order — replay in order through a buffer bounded
|
|
// by the throttle's own concurrency (`n_workers` batches can be in
|
|
// flight at once, so at most that many can be waiting here).
|
|
let mut pending: HashMap<usize, GeneratedBatch> = HashMap::new();
|
|
let mut next_expected = 0usize;
|
|
let mut scratch = vec![0u8; n_genomes];
|
|
for generated in pipe.apply(throttled, n_workers, capacity) {
|
|
pending.insert(generated.start_family_idx, generated);
|
|
while let Some(batch) = pending.remove(&next_expected) {
|
|
let n = batch.masks.len();
|
|
|
|
// ── Resolve this one batch's cross-partition queries — the
|
|
// only place that touches `cache`. Parallel across partitions
|
|
// (one batch at a time, never several concurrently — see the
|
|
// module docs), each thread owning one partition's queries
|
|
// contiguously until this batch is done.
|
|
let genome_mask: Vec<AtomicU8> = batch.genome_mask.into_iter().map(AtomicU8::new).collect();
|
|
let fast_mode = cache.fast_mode();
|
|
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
|
let on_hit = |i: usize, base: u8, g: usize| {
|
|
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
|
};
|
|
if fast_mode {
|
|
cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit);
|
|
} else {
|
|
cache.find_presence_batch(dest, queries, n_genomes, on_hit);
|
|
}
|
|
});
|
|
|
|
for i in 0..n {
|
|
let family_idx = batch.start_family_idx + i;
|
|
if !is_selected(&selected, family_idx) {
|
|
continue; // not in the sample — never resolved above, nothing to report
|
|
}
|
|
for (g, dst) in scratch.iter_mut().enumerate() {
|
|
*dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed);
|
|
}
|
|
on_family(family_idx, batch.masks[i], &scratch);
|
|
}
|
|
next_expected += n;
|
|
}
|
|
}
|
|
debug_assert!(pending.is_empty(), "every generated batch must have been replayed");
|
|
|
|
Ok(())
|
|
}
|