Extract phylogenetic sibling logic into new obikphylo crate
Relocate the `siblings` and `cardcomp` modules from `obikindex` to a dedicated `obikphylo` workspace member. Convert inherent methods on `KmerIndex` into extension traits, update import paths across `obikmer`, and add supporting accessor methods to `obikseq` and `obilayeredmap`. This restructuring reduces the public API surface of `obikindex` while organizing phylogenetic iteration, caching, and distance calculation logic under a dedicated crate.
This commit is contained in:
Eric Coissac
@@ -0,0 +1,301 @@
|
||||
//! 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 obicompactvec::{FamilyMask, PersistentBitMatrix, PersistentCompactIntMatrix, SiblingAnnex};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::Layer;
|
||||
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, ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||||
|
||||
/// 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;
|
||||
|
||||
/// 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().part_dir(part).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
continue;
|
||||
}
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||||
for l in 0..meta.n_layers {
|
||||
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||||
let annex_path = 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(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 `docmd/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)`.
|
||||
outgoing: Vec<Vec<(CanonicalKmer, usize, 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 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>,
|
||||
mut on_family: impl FnMut(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 use_counts = with_counts && layer_dir.join("counts").exists();
|
||||
let mat = if use_counts {
|
||||
Mat::Count(Layer::<PersistentCompactIntMatrix>::open(layer_dir, &meta.mode).map_err(olm_to_ok)?)
|
||||
} else {
|
||||
Mat::Presence(Layer::<PersistentBitMatrix>::open(layer_dir, &meta.mode).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 `docmd/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,
|
||||
});
|
||||
|
||||
let worker_ctx = Arc::clone(&ctx);
|
||||
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)>> = (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.
|
||||
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);
|
||||
for other in kmer.central_canonical_neighbors() {
|
||||
if other == kmer {
|
||||
continue; // local — resolved below straight from `mat`, no lookup
|
||||
}
|
||||
let b = central_base(other, ctx.k);
|
||||
if !mask.has(b) {
|
||||
continue;
|
||||
}
|
||||
let dest = other.partition(ctx.n_parts);
|
||||
outgoing[dest].push((other, i, b));
|
||||
}
|
||||
}
|
||||
|
||||
// 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();
|
||||
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
||||
cache.find_presence_batch(dest, queries, n_genomes, |i, base, g| {
|
||||
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
|
||||
});
|
||||
});
|
||||
|
||||
for i in 0..n {
|
||||
for (g, dst) in scratch.iter_mut().enumerate() {
|
||||
*dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed);
|
||||
}
|
||||
on_family(batch.masks[i], &scratch);
|
||||
}
|
||||
next_expected += n;
|
||||
}
|
||||
}
|
||||
debug_assert!(pending.is_empty(), "every generated batch must have been replayed");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user