Eric Coissac
5c1584967f
The command now supports reopening existing indexes instead of failing when the output file exists. Control flow branches between opening an existing index and constructing a new one, moving configuration setup exclusively to the creation path. Directory existence is enforced upfront with proper I/O error propagation. The --force flag retains its original semantics by removing the target directory before proceeding with a fresh build.
372 lines
18 KiB
Rust
372 lines
18 KiB
Rust
use std::path::Path;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::Ordering;
|
|
|
|
use rayon::prelude::*;
|
|
|
|
use obikseq::CanonicalKmer;
|
|
use obikindex::layer::MphfLayer;
|
|
use obikindex::layer::meta::IndexMode;
|
|
use obipipeline::ThrottleGuard;
|
|
use obisys::progress_bar;
|
|
|
|
use obikindex::KmerIndex;
|
|
use obikindex::OKIResult;
|
|
|
|
use super::cache::PartitionCache;
|
|
use super::helpers::{central_base, is_minorant};
|
|
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder, olm_to_ok};
|
|
|
|
// ── obipipeline data types ─────────────────────────────────────────────────
|
|
|
|
/// A batch of this layer's distinct k-mers (iteration-order index + k-mer),
|
|
/// the pipeline's source item — batched, not one k-mer per item, so that
|
|
/// pipeline messages and their synchronisation cost stay amortised over
|
|
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on
|
|
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
|
|
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
|
|
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
|
|
///
|
|
/// The `usize` is this k-mer's position in `iter_kmers()`'s enumeration
|
|
/// order, **not** an MPHF slot — see `DevDocMD/architecture/siblings.md`. It
|
|
/// is the same index the annex file is written under.
|
|
struct SourceBatch {
|
|
items: Vec<(usize, CanonicalKmer)>,
|
|
_permit: ThrottleGuard,
|
|
}
|
|
|
|
/// One batch's worth of central-substitution variants (up to 3 per source
|
|
/// k-mer), each already routed to its destination partition and carrying
|
|
/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
|
|
/// hit. `(dest_partition, variant, source_order, base)` per entry, where
|
|
/// `source_order` is the source k-mer's iteration-order index (see
|
|
/// `SourceBatch`), not an MPHF slot.
|
|
struct VariantBatch {
|
|
items: Vec<(usize, CanonicalKmer, usize, u8)>,
|
|
_permit: ThrottleGuard,
|
|
}
|
|
|
|
enum SibData {
|
|
Batch(SourceBatch),
|
|
Variants(VariantBatch),
|
|
}
|
|
|
|
/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `KmerIndex`.
|
|
pub trait SiblingAnnexBuildExt {
|
|
/// Build the sibling-count/minorant annex for every layer of every
|
|
/// partition of this (already built) index, writing one annex file per
|
|
/// layer alongside its existing index files. Safe to call again later
|
|
/// (e.g. after a fresh `merge`) — each run simply overwrites the annex
|
|
/// files of the index it is called on.
|
|
///
|
|
/// Construction only — no statistics gathered here on purpose: this is
|
|
/// meant to run routinely (it is the artefact the SNP-family distances
|
|
/// will consume), while the sibling-count distribution
|
|
/// ([`sibling_annex_stats`](super::stats::SiblingStatsExt::sibling_annex_stats))
|
|
/// is a separate, occasional diagnostic pass over the result, not run
|
|
/// every time.
|
|
///
|
|
/// Cross-partition/cross-layer lookups are required (a k-mer's siblings
|
|
/// can live in any partition), but the layer loop itself — and thus the
|
|
/// annex file this produces — stays local to one layer at a time.
|
|
fn build_sibling_annex(&self) -> OKIResult<()>;
|
|
}
|
|
|
|
impl SiblingAnnexBuildExt for KmerIndex {
|
|
fn build_sibling_annex(&self) -> OKIResult<()> {
|
|
let n_parts = self.n_partitions();
|
|
|
|
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
|
|
let cache = Arc::new(PartitionCache::build(
|
|
self,
|
|
n_parts,
|
|
self.meta().config.with_counts,
|
|
)?);
|
|
|
|
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
|
|
let mut total_slots: u64 = 0;
|
|
for part in 0..n_parts {
|
|
let index_dir = self.index_dir(part);
|
|
if !index_dir.exists() {
|
|
pb.inc(1);
|
|
continue;
|
|
}
|
|
let meta = self.partition_meta(part)?;
|
|
|
|
let mut part_slots: u64 = 0;
|
|
for l in 0..meta.n_layers {
|
|
part_slots += build_layer_sibling_annex(
|
|
self,
|
|
&self.layer_dir(part, l),
|
|
&meta.mode,
|
|
n_parts,
|
|
l,
|
|
&cache,
|
|
)?;
|
|
}
|
|
total_slots += part_slots;
|
|
pb.inc(1);
|
|
pb.set_message(format!(
|
|
"partition {part}: {part_slots} kmers ({total_slots} total)"
|
|
));
|
|
}
|
|
pb.finish_and_clear();
|
|
tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions");
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Returns the number of distinct k-mers (annex entries) processed, for
|
|
/// progress reporting. A free function, not a `KmerIndex` method — called
|
|
/// only from `build_sibling_annex` above, in the same file.
|
|
fn build_layer_sibling_annex(
|
|
index: &KmerIndex,
|
|
layer_dir: &Path,
|
|
mode: &IndexMode,
|
|
n_parts: usize,
|
|
l: usize,
|
|
cache: &Arc<PartitionCache>,
|
|
) -> OKIResult<u64> {
|
|
let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_ok)?;
|
|
let k = index.kmer_size();
|
|
let n = mphf.n();
|
|
|
|
// Whether this layer's family-member fields can carry a real layer
|
|
// number (exploited by `family_scan::scan_layer_families`'s fast path)
|
|
// or must stay presence-only — a single fact about the whole index,
|
|
// decided once by `cache` (built from the same `PartitionMeta` this
|
|
// function would otherwise re-derive) so the writer and every reader
|
|
// agree; see `PartitionCache::fast_mode`'s docs.
|
|
let fast_mode = cache.fast_mode();
|
|
|
|
// ── The annex file itself is the reconciliation state, indexed by
|
|
// this layer's k-mer iteration order (the physical layout of
|
|
// `unitigs.bin`), never by MPHF slot — this layer's own k-mers are
|
|
// known members by construction, so no evidence check, no MPHF
|
|
// slot, and no slot -> k-mer reconstruction is needed or legitimate
|
|
// here (see `DevDocMD/architecture/siblings.md`: the MPHF is not
|
|
// invertible, and evidence answers membership, not identity). No
|
|
// separate in-memory accumulator: every concurrent write goes
|
|
// straight through `SiblingAnnexBuilder::atomic_slot` into the
|
|
// mmap — a layer can hold billions of k-mers, so a `Vec` shadowing
|
|
// the whole file in RAM just to copy it out again at the end (the
|
|
// previous design) doubles memory for no benefit once the builder
|
|
// itself can be written concurrently.
|
|
//
|
|
// `Arc<SiblingAnnexBuilder>`, not a bare `SiblingAnnexBuilder` —
|
|
// `Pipe::apply` requires its source iterator to be `Send + 'static`
|
|
// (its items are dispatched to worker threads that outlive this
|
|
// call), so the batch-generating closure below needs an owned
|
|
// handle it can move in, not a borrow of a local. `atomic_slot`,
|
|
// not `set`, for every concurrent write below: the gather phase
|
|
// parallelises across destination partitions (independent
|
|
// `cache.find` calls, safe to run concurrently) and their hits can
|
|
// land on arbitrary, possibly-shared source entries — a lock-free
|
|
// `fetch_or` avoids needing any synchronisation beyond that.
|
|
// Reclaimed as a plain owned value (`Arc::try_unwrap`) once every
|
|
// concurrent phase below is done, for the sequential finalisation
|
|
// pass. ─────────────────────────────────────────────────────────
|
|
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
|
let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?);
|
|
|
|
// ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
|
|
// the actual cross-partition lookup reuses
|
|
// `KmerPartition::query_partition_with` (the same batching mechanism
|
|
// `obikmer query` already uses: open a partition's files once,
|
|
// answer a whole batch of queries against it) instead of one lookup
|
|
// per pipeline item. A per-item lookup (tried first) reopened/
|
|
// re-mmap'd every target partition's files on every single variant
|
|
// — fine at toy scale, but ~90% system time against a real index,
|
|
// observed in practice. A *later* attempt still pushed one pipeline
|
|
// message per generated variant (a `Flat` stage, `SourceItem` =>
|
|
// `VariantQuery`, one k-mer in => up to 3 variants out as separate
|
|
// messages) — cheaper than reopening files, but sampling a real run
|
|
// showed most wall-clock time going into per-message channel
|
|
// send/notify syscalls instead of the lookup itself: the pipeline's
|
|
// whole point is amortising synchronisation over a batch, and a
|
|
// single k-mer's ≤3 variants is far too fine a granularity for
|
|
// that. Batching `BATCH_SIZE` source k-mers into one pipeline item
|
|
// — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of
|
|
// variants out, one message either way — keeps the per-message
|
|
// synchronisation cost amortised over thousands of lookups instead
|
|
// of one to three. `BATCH_SIZE` was originally tuned back when the
|
|
// per-k-mer cost inside this stage (minimiser via `RollingStat`)
|
|
// was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`,
|
|
// a direct O(k) bit-arithmetic replacement) — at the old per-k-mer
|
|
// cost, dispatch overhead (the scheduler's single-threaded `Select`
|
|
// loop: one channel round-trip per batch) was negligible next to
|
|
// the work each batch represented; now that the work itself is
|
|
// cheaper, that fixed per-batch overhead is proportionally larger,
|
|
// so a bigger batch amortises it over more k-mers again. ─────────
|
|
const BATCH_SIZE: usize = 32768;
|
|
let n_workers = obisys::effective_parallelism();
|
|
let capacity = 256;
|
|
|
|
// Throttling limits how many *batches* are in flight at once — the
|
|
// permit is acquired per batch (not per k-mer) in the source
|
|
// thread, and released once its `VariantBatch` has been read out of
|
|
// the pipeline by the accumulation loop below. See
|
|
// `obipipeline::throttle`'s docs for why this is required, not
|
|
// optional, once a `Flat`-style stage sits in the pipeline.
|
|
//
|
|
// Batches stream straight from `unitigs.bin` via
|
|
// `enumerate_kmers_batch` — never a full-layer `Vec` collect (a
|
|
// layer can hold billions of k-mers; see the "no full collect"
|
|
// rule). Each k-mer's own base is seeded straight into the annex in
|
|
// this same pass, since this is exactly the iteration order the
|
|
// annex is keyed on — no separate seeding pass needed.
|
|
let seed_builder = Arc::clone(&builder);
|
|
let batches = mphf
|
|
.enumerate_kmers_batch(BATCH_SIZE)
|
|
.map(move |(start, kmers)| {
|
|
kmers
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(i, kmer)| {
|
|
let base = central_base(kmer, k);
|
|
let bits = if fast_mode {
|
|
FamilyMask::layer_bits(base, l)
|
|
} else {
|
|
FamilyMask::presence_bits(base)
|
|
};
|
|
seed_builder
|
|
.atomic_slot(start + i)
|
|
.fetch_or(bits, Ordering::Relaxed);
|
|
(start + i, kmer)
|
|
})
|
|
.collect::<Vec<_>>()
|
|
});
|
|
|
|
// Diagnostic: count still-empty annex slots after seeding + cross-partition
|
|
// resolution. A non-zero count here means some k-mer's own base was never
|
|
// written (should be impossible after the batch_offset fix above).
|
|
let check_empty = |label: &str| {
|
|
let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count();
|
|
if empty > 0 {
|
|
tracing::warn!(
|
|
"{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})"
|
|
);
|
|
}
|
|
};
|
|
let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch {
|
|
items: t.item,
|
|
_permit: t.guard,
|
|
});
|
|
|
|
let pipe = obipipeline::make_pipe! {
|
|
SibData : SourceBatch => VariantBatch,
|
|
| {
|
|
move |batch: SourceBatch| -> VariantBatch {
|
|
let mut items = Vec::with_capacity(batch.items.len() * 3);
|
|
for (order, kmer) in batch.items {
|
|
for variant in kmer.central_canonical_neighbors() {
|
|
if variant == kmer {
|
|
continue;
|
|
}
|
|
items.push((
|
|
variant.partition(n_parts),
|
|
variant,
|
|
order,
|
|
central_base(variant, k),
|
|
));
|
|
}
|
|
}
|
|
VariantBatch { items, _permit: batch._permit }
|
|
}
|
|
} : Batch => Variants,
|
|
};
|
|
|
|
// ── Group generated variants by destination partition. `cache`
|
|
// holds every partition already mmap'd (no more `open()` cost), but
|
|
// `mmap` pages are still loaded on demand and can be evicted — a
|
|
// lookup is not free just because the file isn't reopened. Grouping
|
|
// keeps one partition's pages hot while its whole batch is resolved,
|
|
// instead of faulting pages in and out as lookups jump between
|
|
// partitions in whatever order the pipeline happens to produce
|
|
// them. Each batch's throttle permit drops here, once accumulated.
|
|
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> =
|
|
(0..n_parts).map(|_| Vec::new()).collect();
|
|
for vb in pipe.apply(throttled, n_workers, capacity) {
|
|
for (dest_partition, variant, source_order, base) in vb.items {
|
|
outgoing[dest_partition].push((variant, source_order, base));
|
|
}
|
|
}
|
|
|
|
check_empty("after_seeding");
|
|
|
|
// ── Resolve each partition's batch against the cache, parallelised
|
|
// over roughly equal-sized *chunks*, not over partitions — a plain
|
|
// `outgoing.par_iter()` over `n_parts` buckets gives one thread the
|
|
// whole of one partition's bucket, however large, and a lot of
|
|
// buckets are far from equal: a central-base substitution changes
|
|
// the minimiser (and thus routes to a different partition) only
|
|
// when the winning minimiser window overlaps the central position.
|
|
// For k=31/m=11 that is ~11 of the 21 possible windows — so *most*
|
|
// of the remaining ~10/21 send the variant right back to the
|
|
// partition already being built. That self-partition bucket ends up
|
|
// far larger than any other, so the naive per-partition split
|
|
// leaves one thread grinding through it alone long after every
|
|
// other partition's (tiny) bucket is done — confirmed by sampling a
|
|
// real run: one thread solely in `MphfLayer::find` while the rest
|
|
// of the pool sits idle. Splitting each non-empty bucket into
|
|
// `chunk_size`-sized pieces first keeps this pass's per-partition
|
|
// locality (each chunk is still contiguous within one partition,
|
|
// still resolved via `cache.find` grouped by that partition) while
|
|
// letting Rayon spread a single oversized bucket across several
|
|
// threads instead of pinning it to one. ──────────────────────────
|
|
let chunk_size = ((outgoing.iter().map(Vec::len).sum::<usize>() / n_workers).max(1)).min(4096);
|
|
let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(_, q)| !q.is_empty())
|
|
.flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c)))
|
|
.collect();
|
|
work.par_iter().for_each(|&(dest, chunk)| {
|
|
for &(variant, source_order, base) in chunk {
|
|
if let Some(layer) = cache.find(dest, variant) {
|
|
let bits = if fast_mode {
|
|
FamilyMask::layer_bits(base, layer)
|
|
} else {
|
|
FamilyMask::presence_bits(base)
|
|
};
|
|
builder
|
|
.atomic_slot(source_order)
|
|
.fetch_or(bits, Ordering::Relaxed);
|
|
}
|
|
}
|
|
});
|
|
|
|
check_empty("after_resolution");
|
|
|
|
// ── Write the layer's annex file, indexed by iteration order — a
|
|
// second streamed pass over `unitigs.bin` (via `enumerate_kmers`),
|
|
// now that every entry's mask is final; never a `Vec` hold of the
|
|
// whole layer. The minorant flag is computed here, not by every
|
|
// later reader: this is the one place the whole family's *final*
|
|
// mask and this entry's own k-mer (already in hand, no extra
|
|
// lookup) are both available together. Every consumer that only
|
|
// needs to know "is this k-mer the family's minorant" — the common
|
|
// case, since a family is tallied once, at its minorant — reads
|
|
// the bit straight back instead of re-deriving it (re-scanning
|
|
// `unitigs.bin` and re-hashing through the MPHF, the cost
|
|
// `is_minorant` was cheap to *compute* but expensive to *get the
|
|
// inputs for* every time).
|
|
// Every concurrent phase above is done and its `Arc` clone (the
|
|
// seeding pass's `seed_builder`) has already been dropped along
|
|
// with the now-fully-drained `batches` iterator — this is the only
|
|
// remaining strong reference, so reclaiming plain ownership for the
|
|
// sequential `set` calls below is guaranteed to succeed.
|
|
let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| {
|
|
panic!("sibling-annex builder still shared after every concurrent phase completed")
|
|
});
|
|
for (order, kmer) in mphf.enumerate_kmers() {
|
|
let final_mask = builder.get(order).expect("seeded by construction");
|
|
let minorant = is_minorant(kmer, final_mask, k);
|
|
builder.set(order, final_mask.with_minorant(minorant));
|
|
}
|
|
builder.close()?;
|
|
|
|
Ok(n as u64)
|
|
}
|