From 4b7b3c3c1a7fa2ea9403e59133b8d53c45859432 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 21 Aug 2026 09:18:09 +0200 Subject: [PATCH] Extract indexing stage into LayerBuilder with NUMA-aware scheduling Introduce a dedicated LayerBuilder struct to orchestrate parallel layer 0 construction across partitions. The fluent API supports configurable abundance thresholds and intermediate artifact retention. Parallel execution is delegated to a NUMA-aware PartitionRunner scheduler, while progress reporting and stage timing are shifted to the command layer. A new mark_indexed method cleanly separates state tracking from orchestration by generating a completion sentinel. --- .../implementation/partition_layer_cache.md | 105 ++++++++++++++++- src/obikindex/src/index/kmer_index.rs | 64 +--------- .../src/algorithms/layer_builder/mod.rs | 109 ++++++++++++++++++ src/obikindexer/src/algorithms/mod.rs | 5 +- src/obikmer/src/cmd/index/mod.rs | 25 ++-- src/obikphylo/src/siblings/tests.rs | 7 +- 6 files changed, 245 insertions(+), 70 deletions(-) create mode 100644 src/obikindexer/src/algorithms/layer_builder/mod.rs diff --git a/DevDocMD/implementation/partition_layer_cache.md b/DevDocMD/implementation/partition_layer_cache.md index 74a1286a..c4d5deee 100644 --- a/DevDocMD/implementation/partition_layer_cache.md +++ b/DevDocMD/implementation/partition_layer_cache.md @@ -53,8 +53,11 @@ detail, session ended (budget) before implementation — see "(5) design agreed" below; **read it before touching `KmerPartition`/`Layer` signatures**, the shape is fully specified. (6) done — `Counter`, a third algorithm, extracted from `PartitionRouter` the same way `Dereplicator` -was in (4); (5) itself still not implemented, still first on the "order of -remaining work" list — see "(6) done" below. Earlier mix-up, for +was in (4). (7) done — `LayerBuilder`, the fourth and last pipeline +algorithm; the indexing pipeline is now fully decomposed into +`obikindexer::algorithms::{partitionner, dereplicator, counter, +layer_builder}`. (5) itself still not implemented, still first on the +"order of remaining work" list — see "(7) done" below. Earlier mix-up, for context: an earlier version of this doc used the name `KmerPartition` (singular) for what was actually the *collection* type (later renamed `KmerPartitions`, later @@ -851,6 +854,104 @@ Still not done: (5) (`Layer`/`KmerPartition` redesign, `PartitionRouter`'s `&mut`→`&`), the future cache crate, `build_layers` (still a `KmerIndex` inherent method, not an algorithm), and `obikalgorithm` itself. +## (7) done (2026-08-21): `LayerBuilder` — fourth and last pipeline algorithm + +Closes out the indexing pipeline: `build_layers`/`build_index_layer` +(the last stage still living as `KmerIndex` inherent methods, flagged as +inconsistent since (6)) extracted into `obikindexer::algorithms:: +layer_builder::LayerBuilder`, same two-phase shape as the other three. + +```rust +pub struct LayerBuilder<'a> { + index: &'a KmerIndex, + n_partitions: usize, + min_abundance: u32, + max_abundance: Option, + keep_intermediate: bool, +} + +impl<'a> LayerBuilder<'a> { + pub fn new(index: &'a KmerIndex) -> Self; + pub fn min_abundance(mut self, v: u32) -> Self; + pub fn max_abundance(mut self, v: Option) -> Self; + pub fn keep_intermediate(mut self, v: bool) -> Self; + pub fn run(&self, on_progress: Option) -> SKResult; // returns total kmers built +} +``` + +**Different from all three prior extractions in one respect, deliberately +not "fixed" to match them**: the actual per-partition construction logic +(De Bruijn graph from dereplicated superkmers + provisional counts → +unitigs → MPHF → matrix) stayed put as `KmerIndex::build_index_layer`/ +`remove_build_artifacts` (both already `pub`) — not moved into +`obikindexer`. Checked first: unlike `dereplicate_partition`/ +`count_partition` (which only ever had one caller), `build_index_layer` +depends on several `obikindex`-internal helpers (`graph_pipeline:: +{write_graph_as_unitigs, materialize_layer}`, `common::olm_to_sk`) that +are `pub(crate)` and shared with `merge`/`select`/`rebuild`'s own +layer-construction paths — moving `build_index_layer` out would have +meant either exporting that internal surface just for this one algorithm +or duplicating it. Neither was needed: `build_index_layer`/ +`remove_build_artifacts` were *already* public `KmerIndex` methods, so +`LayerBuilder`'s job is purely the orchestration around them (scheduling, +config, progress) — the exact same "algorithm calls already-public +`KmerIndex` primitives" shape `PartitionRouter`/`Dereplicator`/`Counter` +already have, just at a coarser grain for this one stage. This is the +"is the producer's API actually deficient?" check from +[[feedback_no_spaghetti_petits_pois]] applied and answered "no" — not +skipped. + +**Two more real divergences, both forced by `PartitionRunner`, not +arbitrary:** +- Uses `obikindex::PartitionRunner` (NUMA-aware scheduler, already + `pub use`d from `obikindex`) instead of plain `rayon::into_par_iter` + like `Dereplicator`/`Counter` — matches what `build_layers` already used + before extraction; this stage is more CPU/memory-intensive per partition + (graph construction) than scatter/dereplicate/count. +- Callback bound is `FnMut(Progress) + Send` — a third variant, not + matching either prior shape. `PartitionRunner::run`'s `on_done` is + invoked from its own single controller thread (never concurrently, so + no `Sync` needed, unlike `Dereplicator`/`Counter`'s `Fn + Sync`), but + that controller thread is itself `std::thread::scope`-spawned, so the + closure still has to be `Send` to cross into it — caught immediately by + the compiler (`cannot be sent between threads safely`) when `Send` was + first omitted, not a design guess. `obikalgorithm`'s eventual shared + trait now has three real callback-bound data points to reconcile + (`FnMut` alone for `PartitionRouter::run`'s sequential loop, `FnMut + + Send` here, `Fn + Sync` for `Dereplicator`/`Counter`'s `rayon` + `par_iter`), not two. + +`KmerIndex::build_layers` deleted outright (`KmerIndex` stays a pure data +structure — no compute orchestration methods, consistent with `dereplicate`/ +`count_kmer`'s removal in (4)/(6)). New `KmerIndex::mark_indexed()` added, +symmetric to `mark_scattered`/`mark_counted`, replacing the inline +`touch(SENTINEL_INDEXED)` that used to live inside `build_layers`. +`Stage::start("index")`/`rep.push(...)` and the `progress_bar`/ +`"{n} total kmers indexed"` log line both moved to `cmd/index/mod.rs`, +same pattern as (3)/(4)/(6) — `LayerBuilder` renders nothing itself, just +reports `Progress`. + +All callers updated: `cmd/index/mod.rs` (Stage 3), `obikphylo`'s test +harness (also gained a `mark_indexed()` call it was missing — harmless +before since nothing checked `IndexState::Indexed` in that test, but now +correct). + +Full workspace suite green (`cargo check --workspace --all-targets` + +`cargo test --workspace`, exit code 0), plus the end-to-end CLI smoke test +(`scripts/smoke_test_index.sh`, built earlier specifically so this +verification step is a one-liner from now on) — 870 kmers indexed, query +round-trip confirmed, same numbers as (6). + +**The indexing pipeline is now fully decomposed**: `obikindexer:: +algorithms::{partitionner, dereplicator, counter, layer_builder}`, each a +`new`/(setters)/`run` algorithm operating on a `&KmerIndex` (or `&mut` for +`PartitionRouter`, not yet fixed — see (5)), `KmerIndex` itself holding no +pipeline-orchestration logic anymore. Still not done: (5), the future +cache crate, `obikalgorithm` (now unblocked — three real callback-bound +variants observed, worth revisiting whether a single trait can express +all three or whether that's itself the answer: it can't, and the trait +should not force it). + ## The problem Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps diff --git a/src/obikindex/src/index/kmer_index.rs b/src/obikindex/src/index/kmer_index.rs index 0d5fa63d..f6a4f747 100644 --- a/src/obikindex/src/index/kmer_index.rs +++ b/src/obikindex/src/index/kmer_index.rs @@ -5,7 +5,6 @@ use std::path::{Path, PathBuf}; use crate::layer::meta::PartitionMeta; use obisys::{Reporter, Stage, progress_bar}; use rayon::prelude::*; -use tracing::info; use obikseq::{set_k, set_m}; @@ -230,6 +229,12 @@ impl KmerIndex { Ok(()) } + /// Mark layer construction as complete and write `index.done`. + pub fn mark_indexed(&self) -> OKIResult<()> { + touch(&self.root_path.join(SENTINEL_INDEXED))?; + Ok(()) + } + /// Write `spectrums/{label}.json` from an already-computed kmer /// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather /// than `obikindexer::algorithms::partitionner::KmerSpectrum` — `KmerIndex` @@ -260,63 +265,6 @@ impl KmerIndex { Ok(()) } - /// Build the layered MPHF index for all partitions in parallel. - /// - /// Writes `index.done` upon completion. - pub fn build_layers( - &self, - min_ab: u32, - max_ab: Option, - keep_intermediate: bool, - rep: &mut Reporter, - ) -> OKIResult<()> { - let n = self.n_partitions(); - let t = Stage::start("index"); - let with_counts = self.meta.config.with_counts; - let evidence = self.meta.config.evidence.clone(); - let block_bits = self.meta.config.block_bits; - let mut total_kmers: usize = 0; - let pb = progress_bar("index", n as u64, "partitions"); - - let order: Vec = (0..n).collect(); - let runner = crate::index::numa::PartitionRunner::new(); - runner - .run( - &order, - |i| { - self.build_index_layer( - i, - min_ab, - max_ab, - with_counts, - &evidence, - block_bits, - ) - }, - |i, n_kmers, _| { - if n_kmers > 0 { - total_kmers += n_kmers; - pb.inc(1); - pb.set_message(format!("{i}: {n_kmers} kmers")); - } - }, - ) - .map_err(OKIError::Partition)?; - - pb.finish_and_clear(); - info!("done — {} total kmers indexed", total_kmers); - - if !keep_intermediate { - for i in 0..n { - self.remove_build_artifacts(i); - } - } - - touch(&self.root_path.join(SENTINEL_INDEXED))?; - rep.push(t.stop()); - Ok(()) - } - /// Path to the unitigs file for partition `part`, layer `layer`. pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> PathBuf { self.layer_dir(part, layer).join("unitigs.bin") diff --git a/src/obikindexer/src/algorithms/layer_builder/mod.rs b/src/obikindexer/src/algorithms/layer_builder/mod.rs new file mode 100644 index 00000000..4ffc3218 --- /dev/null +++ b/src/obikindexer/src/algorithms/layer_builder/mod.rs @@ -0,0 +1,109 @@ +//! Layer construction — the last stage of the indexing pipeline: turns +//! each partition's dereplicated superkmers + provisional counts (from +//! [`crate::algorithms::counter`]) into the real layer 0 (`mphf.bin` + +//! `unitigs.bin` + matrix). See +//! `DevDocMD/implementation/partition_layer_cache.md`. + +use obikindex::{KmerIndex, PartitionRunner}; +use obiskio::SKResult; +use obisys::Progress; + +/// Builds every partition's real layer 0 in parallel. +/// +/// Two-phase construction, same shape as the other pipeline algorithms: +/// `new` (no disk access), optional setters (`min_abundance`/ +/// `max_abundance`/`keep_intermediate`), then `run`. The actual +/// per-partition construction (De Bruijn graph, unitigs, MPHF, matrix) +/// stays a `KmerIndex` method (`build_index_layer`) — unlike +/// `Dereplicator`/`Counter`, it depends on several `obikindex`-internal +/// helpers (`graph_pipeline`, `common::olm_to_sk`) already shared with +/// `merge`/`select`/`rebuild`'s own layer-construction paths, so it +/// belongs there, not duplicated or newly exported just for this +/// algorithm. `LayerBuilder`'s own job is the orchestration around it: +/// scheduling, configuration, and progress reporting. +pub struct LayerBuilder<'a> { + index: &'a KmerIndex, + n_partitions: usize, + min_abundance: u32, + max_abundance: Option, + keep_intermediate: bool, +} + +impl<'a> LayerBuilder<'a> { + pub fn new(index: &'a KmerIndex) -> Self { + Self { + index, + n_partitions: index.n_partitions(), + min_abundance: 1, + max_abundance: None, + keep_intermediate: false, + } + } + + /// Minimum kmer abundance (inclusive) — kmers below this are filtered + /// out. Default 1 (no filtering). + pub fn min_abundance(mut self, v: u32) -> Self { + self.min_abundance = v; + self + } + + /// Maximum kmer abundance (inclusive) — kmers above this are filtered + /// out. Default `None` (no ceiling). + pub fn max_abundance(mut self, v: Option) -> Self { + self.max_abundance = v; + self + } + + /// Keep each partition's build artifacts (dereplicated superkmers, + /// `mphf1.bin`, `counts1.bin`) after their layer is built, instead of + /// deleting them (default: `false`). + pub fn keep_intermediate(mut self, v: bool) -> Self { + self.keep_intermediate = v; + self + } + + /// Build every partition's layer 0 in parallel via `PartitionRunner` + /// (NUMA-aware scheduling — this stage is more CPU/memory-intensive + /// per partition than scatter/dereplicate/count, unlike those it + /// doesn't use plain rayon). Returns the total number of kmers built + /// across all partitions. + /// + /// `on_progress`, when set, is called once per completed partition — + /// `FnMut + Send`, not `Fn + Sync`: `PartitionRunner::run` calls + /// `on_done` from its own single controller thread, not concurrently + /// from workers, so (unlike `Dereplicator`/`Counter`) there's no need + /// for a `Sync` bound — but that controller runs in a spawned scoped + /// thread, so the closure itself still has to cross a thread boundary + /// once, hence `Send`. `total: Some(n_partitions)` — known up front. + pub fn run(&self, mut on_progress: Option) -> SKResult { + let with_counts = self.index.with_counts(); + let evidence = self.index.evidence_mode().clone(); + let block_bits = self.index.block_bits(); + let min_ab = self.min_abundance; + let max_ab = self.max_abundance; + let mut total_kmers: usize = 0; + let mut done: u64 = 0; + + let order: Vec = (0..self.n_partitions).collect(); + let runner = PartitionRunner::new(); + runner.run( + &order, + |i| self.index.build_index_layer(i, min_ab, max_ab, with_counts, &evidence, block_bits), + |_i, n_kmers, _elapsed| { + total_kmers += n_kmers; + done += 1; + if let Some(cb) = on_progress.as_mut() { + cb(Progress { position: done, total: Some(self.n_partitions as u64) }); + } + }, + )?; + + if !self.keep_intermediate { + for i in 0..self.n_partitions { + self.index.remove_build_artifacts(i); + } + } + + Ok(total_kmers) + } +} diff --git a/src/obikindexer/src/algorithms/mod.rs b/src/obikindexer/src/algorithms/mod.rs index b20e5765..acf0fcf9 100644 --- a/src/obikindexer/src/algorithms/mod.rs +++ b/src/obikindexer/src/algorithms/mod.rs @@ -4,8 +4,11 @@ //! algorithm is its own submodule: [`partitionner`] (routing raw super-kmers //! into partitions), [`dereplicator`] (deduplicating a partition's raw //! super-kmers), [`counter`] (counting unique canonical kmers from a -//! partition's dereplicated super-kmers) — the three run in that order. +//! partition's dereplicated super-kmers), [`layer_builder`] (turning +//! dereplicated super-kmers + counts into the real layer 0) — the four run +//! in that order. pub mod counter; pub mod dereplicator; +pub mod layer_builder; pub mod partitionner; diff --git a/src/obikmer/src/cmd/index/mod.rs b/src/obikmer/src/cmd/index/mod.rs index 7e299b31..bf8ac1a2 100644 --- a/src/obikmer/src/cmd/index/mod.rs +++ b/src/obikmer/src/cmd/index/mod.rs @@ -4,6 +4,7 @@ use std::time::Instant; use clap::Args; use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::dereplicator::Dereplicator; +use obikindexer::algorithms::layer_builder::LayerBuilder; use obikindexer::algorithms::partitionner::PartitionRouter; use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex}; use obikindex::layer::IndexMode; @@ -328,13 +329,23 @@ pub fn run(args: IndexArgs) { // ── Stage 3: build layered index ───────────────────────────────────────── if idx.state() < IndexState::Indexed { - idx.build_layers( - args.min_abundance, - args.max_abundance, - args.keep_intermediate, - &mut rep, - ).unwrap_or_else(|e| { - eprintln!("error: {e}"); + let t = Stage::start("index"); + let pb = progress_bar("index", idx.n_partitions() as u64, "partitions"); + let total_kmers = LayerBuilder::new(&idx) + .min_abundance(args.min_abundance) + .max_abundance(args.max_abundance) + .keep_intermediate(args.keep_intermediate) + .run(Some(|_: Progress| pb.inc(1))) + .unwrap_or_else(|e| { + eprintln!("error: {e}"); + std::process::exit(1); + }); + pb.finish_and_clear(); + info!("done — {total_kmers} total kmers indexed"); + rep.push(t.stop()); + + idx.mark_indexed().unwrap_or_else(|e| { + eprintln!("error marking index done: {e}"); std::process::exit(1); }); } else { diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index cfe7d71d..cb3fae36 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -8,6 +8,7 @@ use tempfile::tempdir; use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::dereplicator::Dereplicator; +use obikindexer::algorithms::layer_builder::LayerBuilder; use obikindexer::algorithms::partitionner::PartitionRouter; use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode}; @@ -70,7 +71,6 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label))) .expect("create"); - let mut rep = Reporter::new(); let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta"); let mut router = PartitionRouter::new(&mut idx); for page in stream { @@ -86,7 +86,10 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { idx.mark_scattered().expect("mark_scattered"); idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum"); idx.mark_counted().expect("mark_counted"); - idx.build_layers(1, None, false, &mut rep).expect("build_layers"); + LayerBuilder::new(&idx) + .run(None::) + .expect("build_layers"); + idx.mark_indexed().expect("mark_indexed"); idx }