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.
This commit is contained in:
Eric Coissac
2026-08-21 10:20:03 +02:00
parent 346095b9eb
commit 4b7b3c3c1a
6 changed files with 245 additions and 70 deletions
@@ -53,8 +53,11 @@ detail, session ended (budget) before implementation — see "(5) design
agreed" below; **read it before touching `KmerPartition`/`Layer` agreed" below; **read it before touching `KmerPartition`/`Layer`
signatures**, the shape is fully specified. (6) done — `Counter`, a third signatures**, the shape is fully specified. (6) done — `Counter`, a third
algorithm, extracted from `PartitionRouter` the same way `Dereplicator` algorithm, extracted from `PartitionRouter` the same way `Dereplicator`
was in (4); (5) itself still not implemented, still first on the "order of was in (4). (7) done — `LayerBuilder`, the fourth and last pipeline
remaining work" list — see "(6) done" below. Earlier mix-up, for 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 context: an earlier
version of this doc used the name `KmerPartition` (singular) for what was version of this doc used the name `KmerPartition` (singular) for what was
actually the *collection* type (later renamed `KmerPartitions`, later 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` `&mut`→`&`), the future cache crate, `build_layers` (still a `KmerIndex`
inherent method, not an algorithm), and `obikalgorithm` itself. 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<u32>,
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<u32>) -> Self;
pub fn keep_intermediate(mut self, v: bool) -> Self;
pub fn run(&self, on_progress: Option<impl FnMut(Progress) + Send>) -> SKResult<usize>; // 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 ## The problem
Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps
+6 -58
View File
@@ -5,7 +5,6 @@ use std::path::{Path, PathBuf};
use crate::layer::meta::PartitionMeta; use crate::layer::meta::PartitionMeta;
use obisys::{Reporter, Stage, progress_bar}; use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*; use rayon::prelude::*;
use tracing::info;
use obikseq::{set_k, set_m}; use obikseq::{set_k, set_m};
@@ -230,6 +229,12 @@ impl KmerIndex {
Ok(()) 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 /// Write `spectrums/{label}.json` from an already-computed kmer
/// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather /// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather
/// than `obikindexer::algorithms::partitionner::KmerSpectrum` — `KmerIndex` /// than `obikindexer::algorithms::partitionner::KmerSpectrum` — `KmerIndex`
@@ -260,63 +265,6 @@ impl KmerIndex {
Ok(()) 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<u32>,
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<usize> = (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`. /// Path to the unitigs file for partition `part`, layer `layer`.
pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> PathBuf { pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> PathBuf {
self.layer_dir(part, layer).join("unitigs.bin") self.layer_dir(part, layer).join("unitigs.bin")
@@ -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<u32>,
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<u32>) -> 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<impl FnMut(Progress) + Send>) -> SKResult<usize> {
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<usize> = (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)
}
}
+4 -1
View File
@@ -4,8 +4,11 @@
//! algorithm is its own submodule: [`partitionner`] (routing raw super-kmers //! algorithm is its own submodule: [`partitionner`] (routing raw super-kmers
//! into partitions), [`dereplicator`] (deduplicating a partition's raw //! into partitions), [`dereplicator`] (deduplicating a partition's raw
//! super-kmers), [`counter`] (counting unique canonical kmers from a //! 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 counter;
pub mod dereplicator; pub mod dereplicator;
pub mod layer_builder;
pub mod partitionner; pub mod partitionner;
+17 -6
View File
@@ -4,6 +4,7 @@ use std::time::Instant;
use clap::Args; use clap::Args;
use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::counter::Counter;
use obikindexer::algorithms::dereplicator::Dereplicator; use obikindexer::algorithms::dereplicator::Dereplicator;
use obikindexer::algorithms::layer_builder::LayerBuilder;
use obikindexer::algorithms::partitionner::PartitionRouter; use obikindexer::algorithms::partitionner::PartitionRouter;
use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex}; use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
use obikindex::layer::IndexMode; use obikindex::layer::IndexMode;
@@ -328,15 +329,25 @@ pub fn run(args: IndexArgs) {
// ── Stage 3: build layered index ───────────────────────────────────────── // ── Stage 3: build layered index ─────────────────────────────────────────
if idx.state() < IndexState::Indexed { if idx.state() < IndexState::Indexed {
idx.build_layers( let t = Stage::start("index");
args.min_abundance, let pb = progress_bar("index", idx.n_partitions() as u64, "partitions");
args.max_abundance, let total_kmers = LayerBuilder::new(&idx)
args.keep_intermediate, .min_abundance(args.min_abundance)
&mut rep, .max_abundance(args.max_abundance)
).unwrap_or_else(|e| { .keep_intermediate(args.keep_intermediate)
.run(Some(|_: Progress| pb.inc(1)))
.unwrap_or_else(|e| {
eprintln!("error: {e}"); eprintln!("error: {e}");
std::process::exit(1); 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 { } else {
info!("index already built, skipping"); info!("index already built, skipping");
} }
+5 -2
View File
@@ -8,6 +8,7 @@ use tempfile::tempdir;
use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::counter::Counter;
use obikindexer::algorithms::dereplicator::Dereplicator; use obikindexer::algorithms::dereplicator::Dereplicator;
use obikindexer::algorithms::layer_builder::LayerBuilder;
use obikindexer::algorithms::partitionner::PartitionRouter; use obikindexer::algorithms::partitionner::PartitionRouter;
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode}; 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))) let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)))
.expect("create"); .expect("create");
let mut rep = Reporter::new();
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta"); let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
let mut router = PartitionRouter::new(&mut idx); let mut router = PartitionRouter::new(&mut idx);
for page in stream { 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.mark_scattered().expect("mark_scattered");
idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum"); idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum");
idx.mark_counted().expect("mark_counted"); idx.mark_counted().expect("mark_counted");
idx.build_layers(1, None, false, &mut rep).expect("build_layers"); LayerBuilder::new(&idx)
.run(None::<fn(obisys::Progress)>)
.expect("build_layers");
idx.mark_indexed().expect("mark_indexed");
idx idx
} }