OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

Refactor sibling annex to use mmap-backed concurrent storage

Browse Source 
Shift the sibling annex construction pipeline from an in-memory atomic mask to a memory-mapped file backend. This enables lock-free concurrent writes directly into the mapped region, streamlining the two-phase write process to accumulate bits atomically before finalization. Adjusted the cache lookup to return the specific matching layer index rather than a boolean flag, and added tests to verify correct layer tracking and cross-partition resolution in merged indexes.
This commit is contained in:
Eric Coissac
2026-08-16 21:07:39 +02:00
parent fecfe84ea6
commit 0ce934b111
4 changed files with 158 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikphylo/src/siblings/build.rs
+66 -35
View File
@@ -1,5 +1,5 @@
use std::path::Path;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::atomic::Ordering;
use std::sync::Arc;
use rayon::prelude::*;
@@ -102,7 +102,7 @@ impl SiblingAnnexBuildExt for KmerIndex {
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
let layer_dir = index_dir.join(format!("layer_{l}"));
part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, &cache)?;
part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, l, &cache)?;
}
total_slots += part_slots;
pb.inc(1);
@@ -122,6 +122,7 @@ fn build_layer_sibling_annex(
index: &KmerIndex,
layer_dir: &Path,
n_parts: usize,
l: usize,
cache: &Arc<PartitionCache>,
) -> OKIResult<u64> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
@@ -130,28 +131,51 @@ fn build_layer_sibling_annex(
let k = index.kmer_size();
let n = mphf.n();
// ── 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
// `docmd/architecture/siblings.md`: the MPHF is not invertible, and
// evidence answers membership, not identity). The annex is written
// under this same iteration order end to end, so a reader can later
// zip `iter_kmers()` with the annex file directly, with no
// MPHF/slot indirection at read time either.
// Whether this layer's family-member fields can carry a real layer
// number (fast path, exploited by a future scan-time reader) or must
// stay presence-only (today's behaviour) — a single fact about the
// whole index, decided once here, never stored: `n_layers` is
// guaranteed identical across every partition (a merge adds one layer
// to all of them at once), and `FamilyMask`'s 3-bit field only has room
// for layers `0..=6`.
let fast_mode = meta.n_layers <= 7;
if !fast_mode {
tracing::warn!(
"layer {l} ({layer_dir:?}): index has {} layers (>7) — sibling-annex fast layer \
lookup disabled for this build; consider compacting this index",
meta.n_layers
);
}
// ── 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 `docmd/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<Vec<AtomicU8>>`, not `Vec<AtomicU8>` — `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. `AtomicU8`, not `FamilyMask`,
// because the gather phase below parallelises across destination
// partitions (independent `query_partition_with` calls, safe to run
// concurrently) and their `Found` hits can land on arbitrary,
// possibly-shared source entries — a lock-free `fetch_or` avoids
// needing any synchronisation beyond that. ─────────────────────────
let mask: Arc<Vec<AtomicU8>> = Arc::new((0..n).map(|_| AtomicU8::new(0)).collect());
// `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
@@ -196,26 +220,28 @@ fn build_layer_sibling_annex(
// 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 into `mask` in this same
// pass, since this is exactly the iteration order `mask` is keyed
// on — no separate seeding pass needed.
let seed_mask = Arc::clone(&mask);
// 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)| {
seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
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 mask slots after seeding + cross-partition
// 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 = mask.iter().filter(|v| v.load(Ordering::Relaxed) == 0).count();
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:?})");
}
@@ -294,8 +320,9 @@ fn build_layer_sibling_annex(
.collect();
work.par_iter().for_each(|&(dest, chunk)| {
for &(variant, source_order, base) in chunk {
if cache.find(dest, variant) {
mask[source_order].fetch_or(1 << base, Ordering::Relaxed);
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);
}
}
});
@@ -315,10 +342,14 @@ fn build_layer_sibling_annex(
// `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).
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
let mut builder = SiblingAnnexBuilder::new(n, &annex_path)?;
// 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 = FamilyMask::from_bits(mask[order].load(Ordering::Relaxed));
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));
}