diff --git a/DevDocMD/implementation/partition_layer_cache.md b/DevDocMD/implementation/partition_layer_cache.md index 4f3a4c9e..0f37b661 100644 --- a/DevDocMD/implementation/partition_layer_cache.md +++ b/DevDocMD/implementation/partition_layer_cache.md @@ -62,7 +62,10 @@ layer_builder}`. (8) design agreed, item 1 done in (9) — IndexBuilder`, public, the four maintenance methods (`clear_output_for_create`/`create_skeleton`/`finalize_indexed`/`state`) shared with `merge`/`select`/`rebuild`/`reindex`. Item 2 from (8) -(`obikalgorithm::Algorithm`) still not started. Note: `Layer` renamed +(`obikalgorithm::Algorithm`) done in (12) — new crate, `type Output` + +`fn run(&mut self) -> SKResult`, `on_progress` moved off +`run()`'s signature entirely into a per-algorithm `.on_progress(...)` +setter. Note: `Layer` renamed `KmerLayer` (2026-08-21, outside this conversation). (5) itself still not implemented, still first on the "order of remaining work" list. Earlier mix-up, for @@ -1340,10 +1343,138 @@ end to end against the new `Arc`/on-disk-`IndexState` shape. Still not done: (5)'s `KmerPartition`/`Layer` self-naming redesign itself (only the `PartitionRouter`-`&mut`-removal piece landed, as a side -effect); (8)'s `obikalgorithm::Algorithm` trait (paused, not abandoned, -for this detour — points 2–4 from that discussion are still open); the +effect); the `distance.rs` → `obikphylo` relocation ((9), explicitly +deferred); the future cache-manager crate. (8)'s `obikalgorithm:: +Algorithm` trait, resumed and closed in (12) below. + +## (12) done (2026-08-21): `obikalgorithm::Algorithm` — the shared trait, resumed and closed in one session + +Resumed (8)'s point 2 through a point-by-point discussion of what's +actually common across the four pipeline algorithms, now that (11) made +`KmerIndex` itself immutable everywhere. Four sub-points, each closed +before moving to the next: + +**1. Receiver (`&self` vs `&mut self`)** — investigated whether (11)'s +removal of `&mut KmerIndex` also removed the need for `PartitionRouter:: +run` to take `&mut self`. It didn't: `PartitionRouter` holds real +per-run state of its own (`writers: Vec>`, open file +handles, purely in-process RAM — confirmed by checking where `writers` is +stored, nothing to do with `KmerIndex`/disk truth), unrelated to the +index. First proposal (wrap `writers` in `RefCell` so all four could +share a uniform `&self`) was retracted on pushback: manufacturing +interior mutability with runtime borrow checks to satisfy a cosmetic +uniformity that Rust doesn't even require is over-engineering — a trait +method's receiver must match exactly across implementors, but nothing +stops that shared receiver from being `&mut self` with three of the four +implementations simply not using the mutability. Settled: trait declares +`&mut self`; `Dereplicator`/`Counter`/`LayerBuilder` (previously `&self`) +now also take `&mut self`, unused. + +**2. `path_source` as a `PartitionRouter` setter, not a `run()` param** — +added a `files: Option + Send>>` field + +`.files(impl Iterator + Send + 'static) -> Self` setter +(boxed rather than a generic type parameter on `PartitionRouter<'a>`: +negligible cost — one `PathBuf` per input *file*, not per k-mer — for a +much more usable type when passing the builder around). `run` now does +`self.files.take().ok_or_else(...)`, erroring if `.files(...)` was never +called, instead of taking `path_source` as a parameter. + +**3. Unifying the three progress-callback bound shapes** — reopened, then +resolved differently than (8) originally framed it. First proposal +(force everything to `FnMut(Progress) + Send`) was rejected on the same +principle as point 1: `Dereplicator`/`Counter`'s `Fn(Progress) + Sync` +isn't arbitrary — their callback is invoked concurrently from multiple +rayon worker threads, and `FnMut` requires exclusive access, so forcing +it would mean wrapping the callback in a `Mutex` for zero benefit at the +one real call site (`pb.inc(1)`, already thread-safe). The actual +resolution: move `on_progress` off `run()`'s signature entirely, onto a +per-algorithm `.on_progress(...)` setter — same treatment as point 2's +`path_source` — so each algorithm keeps its own bound (`PartitionRouter`: +`FnMut(Progress) + 'a`, sequential, no `Send` needed; `LayerBuilder`: +`FnMut(Progress) + Send + 'a`, crosses into `PartitionRunner`'s +`thread::scope`-spawned controller thread once; `Dereplicator`/`Counter`: +`Fn(Progress) + Sync + 'a`, invoked concurrently from rayon workers). +This *also* dissolves the original problem `run()` had: once the +callback isn't part of `run`'s signature at all, there's nothing left to +unify there, and point 4 (below) becomes trivial. + +**4. `Output` as an associated type, `Error` fixed** — trivial once (3) +moved the callback out: `Error` was already uniform (all four return +`obiskio::SKResult` = `Result`), only `Output` varied +(`()`/`()`/`KmerSpectrum`/`usize`). First cut reused `obiskio::SKResult` +directly as the trait's return type — **caught and corrected the same +session**: `SKError` enumerates I/O-specific cases (`BadMagic`/ +`Truncated`/`Compression`/...), meaningless at the level of a generic +"algorithm" abstraction, and borrowing it made `obikalgorithm` — meant to +be minimal and neutral — depend on a low-level I/O crate purely to reuse +its error type. Textbook instance of the "petits pois" failure mode +(patch around a convenient existing type instead of asking what this +crate should actually own). Fixed to a genuinely generic, boxed error +type owned by `obikalgorithm` itself: + +```rust +// obikalgorithm — no dependency on obiskio or any other crate +pub type Error = Box; +pub type Result = std::result::Result; + +pub trait Algorithm { + type Output; + fn run(&mut self) -> Result; +} +``` + +Any concrete error (`SKError`, `std::io::Error`, ...) converts +automatically via `?`, through `std`'s own blanket `impl From for Box` — no custom `From` impl +needed, no dependency on the crate that defines the concrete error type. +The four algorithms' `run` bodies needed no change beyond the signature's +return type (every existing `?` on an `SKError`-returning subcall keeps +compiling, converting through the same blanket impl at the boundary). + +`PartitionRouter`/`Dereplicator`/`Counter`/`LayerBuilder` each `impl +Algorithm for X<'_> { type Output = ...; fn run(&mut self) -> obikalgorithm::Result<...> { ... } }` +— the old inherent `run` methods were removed outright (not kept as +duplicates), so callers now `use obikalgorithm::Algorithm;` to call +`.run()`. Every field-lifetime-bound boxed callback (`Box` etc.) is tied to the algorithm's own `'a` (the +`&'a KmerIndex` lifetime already on the struct), not `'static` — avoids +forcing callers' progress closures to `move`-capture (and therefore clone +or `Arc`-wrap) local state like `TracedBar`/EMA-rate accumulators that +they'd otherwise want to keep using by reference after `run()` returns. + +**Why a new crate, not a submodule of `obikindexer`**: `obikmer::cmd:: +index::mod` and `obikphylo`'s own test helpers both need to call `.run()` +on these algorithms, so the trait has to be reachable from outside +`obikindexer` — putting it in `obikindexer` itself would work file-wise +but conflates "the trait every algorithm implements" with "one crate's +particular four implementations of it", the same reasoning that already +separated `obikindex` (data model) from `obikindexer` (algorithms +operating on it). `obikalgorithm` has **no dependencies at all** (see +above); `obikindexer`, `obikmer`, and `obikphylo` (dev-dependency, for +its test helper) all depend on it. + +**Blast radius**: `obikindexer`'s four algorithm modules (struct field + +setter + trait impl each); `obikmer::cmd::index::mod` (three call sites: +`.on_progress(cb)` before `.run()`, unqualified now that the trait is in +scope); `obikindexer::algorithms::partitionner::tests` and `obikphylo:: +siblings::tests` (both had direct `.run(None::)`-shaped +calls needing the same treatment). New `obikalgorithm` crate registered +in the workspace `Cargo.toml`, depended on by `obikindexer`/`obikmer` +(regular) and `obikphylo` (dev). + +**Verification**: `cargo check --workspace --all-targets` and `cargo test +--workspace` both green (0 failures), `scripts/smoke_test_index.sh` green +(870 kmers, same as every prior round) — this round didn't repeat the +manual `merge`/`select`/`reindex` CLI exercise from (11), since nothing +in this pass touched those commands' code paths (only the four pipeline +algorithms and `cmd::index`, already covered by the smoke test). Reverified +after the `obiskio`-dependency fix above (same three checks, still green, +`obikalgorithm/Cargo.toml` now has zero `[dependencies]`). + +Still not done: (5)'s `KmerPartition`/`Layer` self-naming redesign; the `distance.rs` → `obikphylo` relocation ((9), explicitly deferred); the -future cache-manager crate. +future cache-manager crate (mentioned in (8) as a later, mirrored +extension-trait exercise, not started). ## The problem diff --git a/src/Cargo.lock b/src/Cargo.lock index 878a7162..0c7fa600 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1515,6 +1515,10 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "obikalgorithm" +version = "0.1.0" + [[package]] name = "obikentropy" version = "0.1.0" @@ -1564,6 +1568,7 @@ dependencies = [ "niffler", "obicompactvec", "obidebruinj", + "obikalgorithm", "obikindex", "obikrope", "obikseq", @@ -1591,6 +1596,7 @@ dependencies = [ "kodama", "obidebruinj", "obifastwrite", + "obikalgorithm", "obikindex", "obikindexer", "obikphylo", @@ -1619,6 +1625,7 @@ dependencies = [ "memmap2", "ndarray", "obicompactvec", + "obikalgorithm", "obikindex", "obikindexer", "obikseq", diff --git a/src/Cargo.toml b/src/Cargo.toml index 84df5883..296753ad 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -1,5 +1,5 @@ [workspace] resolver = "3" -members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obitaxonomy", "obikentropy", "obikphylo"] +members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm"] [profile.release] debug = 1 diff --git a/src/obikalgorithm/Cargo.toml b/src/obikalgorithm/Cargo.toml new file mode 100644 index 00000000..6b0e7f5b --- /dev/null +++ b/src/obikalgorithm/Cargo.toml @@ -0,0 +1,4 @@ +[package] +name = "obikalgorithm" +version = "0.1.0" +edition = "2024" diff --git a/src/obikalgorithm/src/lib.rs b/src/obikalgorithm/src/lib.rs new file mode 100644 index 00000000..3eb85278 --- /dev/null +++ b/src/obikalgorithm/src/lib.rs @@ -0,0 +1,36 @@ +//! Shared trait unifying the four indexing-pipeline algorithms in +//! `obikindexer::algorithms` (`PartitionRouter`/`Dereplicator`/`Counter`/ +//! `LayerBuilder`) — see `DevDocMD/implementation/partition_layer_cache.md`. + +/// Boxed, generic error for `Algorithm::run` — this crate defines the +/// *shape* algorithms conform to, not any specific failure domain, so it +/// must not depend on (and re-expose) any one algorithm's own I/O-flavored +/// error enum (e.g. `obiskio::SKError`). Any concrete error type +/// (`SKError`, `std::io::Error`, ...) converts into this automatically via +/// `?`, through `std`'s own blanket `From for +/// Box` — no dependency on the crate that defines +/// `E` needed here. +pub type Error = Box; + +pub type Result = std::result::Result; + +/// Common shape across the four pipeline algorithms: two-phase +/// construction (`new` + setters — not part of this trait, since `new`'s +/// argument and the setters differ per algorithm), then `run`. +/// +/// `run` takes `&mut self`, not `&self`: forced by `PartitionRouter`, +/// which holds real per-run state (open file writers) — the other three +/// don't need mutability but accept the same receiver rather than the +/// trait special-casing one implementor. +/// +/// Progress reporting is deliberately *not* part of this signature: the +/// four algorithms report progress under genuinely different concurrency +/// models (sequential, parallel-rayon, single-controller-thread), so each +/// keeps its own `.on_progress(...)` setter with its own callback bound +/// instead of a shared one here — forcing a single shape would mean +/// wrapping the parallel algorithms' callback in a `Mutex` for no +/// benefit. +pub trait Algorithm { + type Output; + fn run(&mut self) -> Result; +} diff --git a/src/obikindexer/Cargo.toml b/src/obikindexer/Cargo.toml index 2ad8111a..50f1bfdb 100644 --- a/src/obikindexer/Cargo.toml +++ b/src/obikindexer/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] obikindex = { path = "../obikindex" } +obikalgorithm = { path = "../obikalgorithm" } obikseq = { path = "../obikseq" } obidebruinj = { path = "../obidebruinj" } obiskio = { path = "../obiskio" } diff --git a/src/obikindexer/src/algorithms/counter/mod.rs b/src/obikindexer/src/algorithms/counter/mod.rs index eb3b737b..c96e10a1 100644 --- a/src/obikindexer/src/algorithms/counter/mod.rs +++ b/src/obikindexer/src/algorithms/counter/mod.rs @@ -13,6 +13,7 @@ use std::fs; use std::io; use std::sync::atomic::{AtomicU64, Ordering}; +use obikalgorithm::Algorithm; use obikindex::KmerIndex; use obiskio::SKResult; use obisys::Progress; @@ -38,12 +39,13 @@ pub struct KmerSpectrum { /// 4. Persists the MPHF to `mphf1.bin` for downstream use. /// /// Two-phase construction, same shape as `PartitionRouter`/`Dereplicator`: -/// `new` (no disk access), optional `keep_partial` setter, then `run`, -/// which returns the aggregated [`KmerSpectrum`]. +/// `new` (no disk access), optional `keep_partial`/`.on_progress(...)` +/// setters, then `run`, which returns the aggregated [`KmerSpectrum`]. pub struct Counter<'a> { index: &'a KmerIndex, n_partitions: usize, keep_partial: bool, + on_progress: Option>, } impl<'a> Counter<'a> { @@ -52,6 +54,7 @@ impl<'a> Counter<'a> { index, n_partitions: index.n_partitions(), keep_partial: false, + on_progress: None, } } @@ -62,17 +65,28 @@ impl<'a> Counter<'a> { self } - /// Count every partition in parallel, then aggregate their spectra. - /// - /// `on_progress`, when set, is called once per completed partition, from + /// Progress callback, called once per completed partition, from /// whichever rayon worker thread finished it — `Fn(...) + Sync`, not - /// `FnMut`, same reason as `Dereplicator::run`: the counting pass itself - /// is a parallel `par_iter`, so the callback must tolerate concurrent - /// calls. `total: Some(n_partitions)` — known up front. This algorithm - /// never renders anything itself. Writes `spectrums/{label}.json` and - /// marks the index as counted (`count.done`) once every partition - /// succeeds. - pub fn run(&self, on_progress: Option) -> SKResult { + /// `FnMut`, same reason as `Dereplicator::on_progress`: the counting + /// pass itself is a parallel `par_iter`, so the callback must tolerate + /// concurrent calls. `total: Some(n_partitions)` — known up front. This + /// algorithm never renders anything itself. + pub fn on_progress(mut self, cb: impl Fn(Progress) + Sync + 'a) -> Self { + self.on_progress = Some(Box::new(cb)); + self + } +} + +impl Algorithm for Counter<'_> { + type Output = KmerSpectrum; + + /// Count every partition in parallel, then aggregate their spectra. + /// Writes `spectrums/{label}.json` and marks the index as counted + /// (`count.done`) once every partition succeeds. + fn run(&mut self) -> obikalgorithm::Result { + let index = self.index; + let n_partitions = self.n_partitions; + let on_progress = &self.on_progress; let sys = System::new_all(); // available_memory() can return 0 on macOS when the compressor page count exceeds // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total. @@ -84,21 +98,21 @@ impl<'a> Counter<'a> { let chunk_kmers = chunk_size_from_ram(available / n_threads); let done = AtomicU64::new(0); - let results: Vec> = (0..self.n_partitions) + let results: Vec> = (0..n_partitions) .into_par_iter() .map(|i| { - let dir = self.index.layer_dir(i, 0); + let dir = index.layer_dir(i, 0); let dedup_path = obikindex::layer::dereplicated_superkmers_path(&dir); let result = if dedup_path.exists() { count_partition(&dir, &dedup_path, chunk_kmers) } else { Ok(()) }; - if let Some(cb) = &on_progress { + if let Some(cb) = on_progress { let pos = done.fetch_add(1, Ordering::Relaxed) + 1; cb(Progress { position: pos, - total: Some(self.n_partitions as u64), + total: Some(n_partitions as u64), }); } result @@ -114,8 +128,8 @@ impl<'a> Counter<'a> { let mut f0: u64 = 0; let mut f1: u64 = 0; - for i in 0..self.n_partitions { - let path = self.index.layer_dir(i, 0).join("kmer_spectrum_raw.json"); + for i in 0..n_partitions { + let path = index.layer_dir(i, 0).join("kmer_spectrum_raw.json"); if !path.exists() { continue; } @@ -135,10 +149,10 @@ impl<'a> Counter<'a> { } } - self.index + index .write_spectrum(f0, f1, &counts) .map_err(|e| io::Error::other(e.to_string()))?; - self.index + index .mark_counted() .map_err(|e| io::Error::other(e.to_string()))?; diff --git a/src/obikindexer/src/algorithms/dereplicator/mod.rs b/src/obikindexer/src/algorithms/dereplicator/mod.rs index ed90ff68..d24af805 100644 --- a/src/obikindexer/src/algorithms/dereplicator/mod.rs +++ b/src/obikindexer/src/algorithms/dereplicator/mod.rs @@ -15,6 +15,7 @@ use obisys::Progress; use rayon::prelude::*; use sysinfo::System; +use obikalgorithm::Algorithm; use obikindex::KmerIndex; use dereplicate::{dereplicate_partition, optimal_buckets}; @@ -24,13 +25,12 @@ use dereplicate::{dereplicate_partition, optimal_buckets}; /// merged and their counts summed. /// /// Two-phase construction, same shape as `PartitionRouter`: `new` (no disk -/// access), then `run`. Nothing to configure yet — no setters, unlike -/// `PartitionRouter`'s `level_max`/`theta`/etc. — added if a real need -/// shows up, not speculatively. +/// access), optional `.on_progress(...)` setter, then `run`. pub struct Dereplicator<'a> { index: &'a KmerIndex, n_partitions: usize, level: Level, + on_progress: Option>, } impl<'a> Dereplicator<'a> { @@ -39,9 +39,29 @@ impl<'a> Dereplicator<'a> { index, n_partitions: index.n_partitions(), level: Level::One, + on_progress: None, } } + /// Progress callback, called once per completed partition, from + /// whichever rayon worker thread finished it — `Fn(...) + Sync`, not + /// `FnMut`, unlike `PartitionRouter::on_progress`: that one is driven + /// from a single sequential loop, this one from a parallel `par_iter`, + /// so the callback itself must tolerate concurrent calls (same reason + /// `obisys::TracedBar`'s own methods take `&self`, not `&mut self`). + /// `total: Some(n_partitions)` — known up front here, unlike + /// `PartitionRouter`'s bases-processed count, so the caller can render + /// an actual progress bar rather than a spinner. This algorithm never + /// renders anything itself. + pub fn on_progress(mut self, cb: impl Fn(Progress) + Sync + 'a) -> Self { + self.on_progress = Some(Box::new(cb)); + self + } +} + +impl Algorithm for Dereplicator<'_> { + type Output = (); + /// Dereplicate every partition in parallel. /// /// Each partition file is processed in two phases to bound memory use: @@ -55,19 +75,11 @@ impl<'a> Dereplicator<'a> { /// /// If a merged count exceeds the 24-bit header limit, the sequence is /// emitted as multiple records whose counts sum to the true total. - /// - /// `on_progress`, when set, is called once per completed partition, from - /// whichever rayon worker thread finished it — `Fn(...) + Sync`, not - /// `FnMut`, unlike `PartitionRouter::run`'s callback: that one is driven - /// from a single sequential loop, this one from a parallel `par_iter`, - /// so the callback itself must tolerate concurrent calls (same reason - /// `obisys::TracedBar`'s own methods take `&self`, not `&mut self`). - /// `total: Some(n_partitions)` — known up front here, unlike - /// `PartitionRouter::run`'s bases-processed count, so the caller can - /// render an actual progress bar rather than a spinner. This algorithm - /// never renders anything itself. - pub fn run(&self, on_progress: Option) -> SKResult<()> { + fn run(&mut self) -> obikalgorithm::Result<()> { let level = self.level; + let index = self.index; + let n_partitions = self.n_partitions; + let on_progress = &self.on_progress; let sys = System::new_all(); // available_memory() can return 0 on macOS when the compressor page count exceeds // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total. @@ -79,10 +91,10 @@ impl<'a> Dereplicator<'a> { let available_per_thread = available / n_threads; let done = AtomicU64::new(0); - let results: Vec> = (0..self.n_partitions) + let results: Vec> = (0..n_partitions) .into_par_iter() .map(|i| { - let dir = self.index.layer_dir(i, 0); + let dir = index.layer_dir(i, 0); let result = if dir.exists() { let raw_path = obikindex::layer::raw_superkmers_path(&dir); let n_buckets = optimal_buckets(&raw_path, available_per_thread); @@ -90,9 +102,9 @@ impl<'a> Dereplicator<'a> { } else { Ok(()) }; - if let Some(cb) = &on_progress { + if let Some(cb) = on_progress { let pos = done.fetch_add(1, Ordering::Relaxed) + 1; - cb(Progress { position: pos, total: Some(self.n_partitions as u64) }); + cb(Progress { position: pos, total: Some(n_partitions as u64) }); } result }) diff --git a/src/obikindexer/src/algorithms/layer_builder/mod.rs b/src/obikindexer/src/algorithms/layer_builder/mod.rs index 9f745e5c..eb4225ef 100644 --- a/src/obikindexer/src/algorithms/layer_builder/mod.rs +++ b/src/obikindexer/src/algorithms/layer_builder/mod.rs @@ -4,8 +4,8 @@ //! `unitigs.bin` + matrix). See //! `DevDocMD/implementation/partition_layer_cache.md`. +use obikalgorithm::Algorithm; use obikindex::{KmerIndex, PartitionRunner}; -use obiskio::SKResult; use obisys::Progress; use crate::extensions::PrivateBuilder; @@ -14,9 +14,9 @@ use crate::extensions::PrivateBuilder; /// /// 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 +/// `max_abundance`/`keep_intermediate`/`.on_progress(...)`), 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 @@ -29,6 +29,7 @@ pub struct LayerBuilder<'a> { min_abundance: u32, max_abundance: Option, keep_intermediate: bool, + on_progress: Option>, } impl<'a> LayerBuilder<'a> { @@ -39,6 +40,7 @@ impl<'a> LayerBuilder<'a> { min_abundance: 1, max_abundance: None, keep_intermediate: false, + on_progress: None, } } @@ -64,36 +66,46 @@ impl<'a> LayerBuilder<'a> { self } + /// Progress callback, 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 on_progress(mut self, cb: impl FnMut(Progress) + Send + 'a) -> Self { + self.on_progress = Some(Box::new(cb)); + self + } +} + +impl Algorithm for LayerBuilder<'_> { + type Output = usize; + /// 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). Marks the index as fully built /// (`index.done`) once every partition succeeds. 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 { + fn run(&mut self) -> obikalgorithm::Result { 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 index = self.index; + let n_partitions = self.n_partitions; let mut total_kmers: usize = 0; let mut done: u64 = 0; + let on_progress = &mut self.on_progress; - let order: Vec = (0..self.n_partitions).collect(); + let order: Vec = (0..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) + index.build_index_layer(i, min_ab, max_ab, with_counts, &evidence, block_bits) }, |_i, n_kmers, _elapsed| { total_kmers += n_kmers; @@ -101,14 +113,14 @@ impl<'a> LayerBuilder<'a> { if let Some(cb) = on_progress.as_mut() { cb(Progress { position: done, - total: Some(self.n_partitions as u64), + total: Some(n_partitions as u64), }); } }, )?; if !self.keep_intermediate { - for i in 0..self.n_partitions { + for i in 0..n_partitions { self.index.remove_build_artifacts(i); } } diff --git a/src/obikindexer/src/algorithms/partitionner/router.rs b/src/obikindexer/src/algorithms/partitionner/router.rs index d79a3b55..3e54dd2d 100644 --- a/src/obikindexer/src/algorithms/partitionner/router.rs +++ b/src/obikindexer/src/algorithms/partitionner/router.rs @@ -19,6 +19,7 @@ use obipipeline::{ThrottleGuard, Throttled, throttle}; use obiread::NucPage; use crate::extensions::PrivateBuilder; +use obikalgorithm::Algorithm; // ── Pipeline plumbing, private to `run` ───────────────────────────────────── @@ -88,6 +89,8 @@ pub struct PartitionRouter<'a> { theta: f64, workers: usize, max_open: usize, + files: Option + Send>>, + on_progress: Option>, } impl<'a> PartitionRouter<'a> { @@ -108,6 +111,8 @@ impl<'a> PartitionRouter<'a> { theta: 0.7, workers, max_open: (workers / 4).max(1), + files: None, + on_progress: None, } } @@ -135,6 +140,25 @@ impl<'a> PartitionRouter<'a> { self } + /// Input files for `run`'s file-driven pipeline. Required before + /// calling `run` — `write`/`write_batch` callers (e.g. tests feeding + /// in-memory superkmers directly) don't need this at all. + pub fn files(mut self, path_source: impl Iterator + Send + 'static) -> Self { + self.files = Some(Box::new(path_source)); + self + } + + /// Progress callback for `run`'s file-driven pipeline — called + /// periodically (rate-limited, not once per batch) with cumulative + /// bases processed so far (`total: None`, unknown without pre-scanning + /// every input file). Optional — if never set, `run` reports nothing. + /// Sequential single-thread invocation, hence plain `FnMut`, no `Send` + /// bound needed (unlike `LayerBuilder::on_progress`). + pub fn on_progress(mut self, cb: impl FnMut(Progress) + 'a) -> Self { + self.on_progress = Some(Box::new(cb)); + self + } + /// Route and write one super-kmer to its partition's raw super-kmer file. pub fn write(&mut self, rsk: RoutableSuperKmer) -> SKResult<()> { self.check_not_closed()?; @@ -186,17 +210,58 @@ impl<'a> PartitionRouter<'a> { !self.closed } - /// Run the full scatter pipeline: normalise every file in - /// `path_source` -> build super-kmers -> route -> write, then close. - /// `on_progress`, when set, is called periodically (rate-limited, not - /// once per batch) with cumulative bases processed so far — + // ── private ─────────────────────────────────────────────────────────────── + + /// Directory of partition `i`'s layer 0 — every raw/dereplicated + /// superkmer file and provisional `mphf1.bin`/`counts1.bin` this router + /// produces lives here, alongside where `build_index_layer` + /// (`obikindex`) will later turn it into the real layer 0. + fn layer0_dir(&self, i: usize) -> PathBuf { + obikindex::layer::layer_dir(&self.index.index_dir(i), 0) + } + + fn check_not_closed(&self) -> SKResult<()> { + if self.closed { + Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into()) + } else { + Ok(()) + } + } + + fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> { + if self.writers[partition].is_none() { + let dir = self.layer0_dir(partition); + KmerLayer::create(&dir).map_err(|e| io::Error::other(e.to_string()))?; + let file_path = obikindex::layer::raw_superkmers_path(&dir); + let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?; + self.writers[partition] = Some(writer); + } + Ok(self.writers[partition].as_mut().unwrap()) + } +} + +impl Drop for PartitionRouter<'_> { + fn drop(&mut self) { + let _ = self.close(); + } +} + +impl Algorithm for PartitionRouter<'_> { + type Output = (); + + /// Run the full scatter pipeline: normalise every file from `.files(...)` + /// -> build super-kmers -> route -> write, then close. + /// `.on_progress(...)`, when set, is called periodically (rate-limited, + /// not once per batch) with cumulative bases processed so far — /// `total: None`, since the total isn't known without pre-scanning /// every input file. - pub fn run( - &mut self, - path_source: impl Iterator + Send + 'static, - mut on_progress: Option, - ) -> SKResult<()> { + /// + /// Errors if `.files(...)` was never called. + fn run(&mut self) -> obikalgorithm::Result<()> { + let path_source = self.files.take().ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "PartitionRouter::run: .files(...) was never called") + })?; + let k = self.index.kmer_size(); let level_max = self.level_max; let theta = self.theta; @@ -247,7 +312,7 @@ impl<'a> PartitionRouter<'a> { .iter() .map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap)) .sum::(); - if let Some(cb) = on_progress.as_mut() { + if let Some(cb) = self.on_progress.as_mut() { let now = Instant::now(); if now.duration_since(last_report).as_secs_f64() > REPORT_INTERVAL { last_report = now; @@ -259,47 +324,13 @@ impl<'a> PartitionRouter<'a> { } self.write_batch(batch)?; } - if let Some(cb) = on_progress.as_mut() { + if let Some(cb) = self.on_progress.as_mut() { cb(Progress { position: total_bases, total: None, }); } - self.close() - } - - // ── private ─────────────────────────────────────────────────────────────── - - /// Directory of partition `i`'s layer 0 — every raw/dereplicated - /// superkmer file and provisional `mphf1.bin`/`counts1.bin` this router - /// produces lives here, alongside where `build_index_layer` - /// (`obikindex`) will later turn it into the real layer 0. - fn layer0_dir(&self, i: usize) -> PathBuf { - obikindex::layer::layer_dir(&self.index.index_dir(i), 0) - } - - fn check_not_closed(&self) -> SKResult<()> { - if self.closed { - Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into()) - } else { - Ok(()) - } - } - - fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> { - if self.writers[partition].is_none() { - let dir = self.layer0_dir(partition); - KmerLayer::create(&dir).map_err(|e| io::Error::other(e.to_string()))?; - let file_path = obikindex::layer::raw_superkmers_path(&dir); - let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?; - self.writers[partition] = Some(writer); - } - Ok(self.writers[partition].as_mut().unwrap()) - } -} - -impl Drop for PartitionRouter<'_> { - fn drop(&mut self) { - let _ = self.close(); + self.close()?; + Ok(()) } } diff --git a/src/obikindexer/src/algorithms/partitionner/tests.rs b/src/obikindexer/src/algorithms/partitionner/tests.rs index 85713b0a..6905a285 100644 --- a/src/obikindexer/src/algorithms/partitionner/tests.rs +++ b/src/obikindexer/src/algorithms/partitionner/tests.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use crate::algorithms::counter::Counter; use crate::algorithms::dereplicator::Dereplicator; +use obikalgorithm::Algorithm; use obikindex::{IndexConfig, KmerIndex}; use obikindex::layer::IndexMode; use obikrope::Rope; @@ -60,19 +61,19 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) { let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0); let dir = tempfile::tempdir().unwrap(); - let mut index = test_index(&dir.path().join("idx")); - let mut kp = PartitionRouter::new(&mut index); + let index = test_index(&dir.path().join("idx")); + let mut kp = PartitionRouter::new(&index); kp.write_batch(superkmers).unwrap(); kp.close().unwrap(); drop(kp); // ends the borrow of `index` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope - Dereplicator::new(&index).run(None::).unwrap(); + Dereplicator::new(&index).run().unwrap(); let part_dir = index.layer_dir(0, 0); let dedup_path = obikindex::layer::dereplicated_superkmers_path(&part_dir); if !dedup_path.exists() { return (0, 0); } - let spectrum = Counter::new(&index).run(None::).unwrap(); + let spectrum = Counter::new(&index).run().unwrap(); (spectrum.f0, spectrum.f1) } diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 5b2cd922..3aaedb34 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -19,6 +19,7 @@ obisys = { path = "../obisys" } obiskio = { path = "../obiskio" } obikindex = { path = "../obikindex", default-features = false } obikindexer = { path = "../obikindexer" } +obikalgorithm = { path = "../obikalgorithm" } obikphylo = { path = "../obikphylo" } obitaxonomy = { path = "../obitaxonomy" } clap = { version = "4", features = ["derive"] } diff --git a/src/obikmer/src/cmd/index/mod.rs b/src/obikmer/src/cmd/index/mod.rs index 7d1f6fea..3159b250 100644 --- a/src/obikmer/src/cmd/index/mod.rs +++ b/src/obikmer/src/cmd/index/mod.rs @@ -8,6 +8,7 @@ use obikindexer::algorithms::layer_builder::LayerBuilder; use obikindexer::algorithms::partitionner::PartitionRouter; use obikindex::{validate_label, GenomeInfo, IndexBuilder, IndexConfig, IndexState, KmerIndex}; use obikindex::layer::IndexMode; +use obikalgorithm::Algorithm; fn current_state(idx: &KmerIndex) -> IndexState { idx.state().unwrap_or_else(|e| { @@ -236,7 +237,7 @@ pub fn run(args: IndexArgs) { } info }); - let mut idx = KmerIndex::create(&output, config, genome_info).unwrap_or_else(|e| { + let idx = KmerIndex::create(&output, config, genome_info).unwrap_or_else(|e| { eprintln!("error creating index: {e}"); std::process::exit(1); }); @@ -254,33 +255,32 @@ pub fn run(args: IndexArgs) { let mut last_bases: u64 = 0; const ALPHA: f64 = 0.15; - let mut router = PartitionRouter::new(&mut idx) + let mut router = PartitionRouter::new(&idx) .level_max(args.common.level_max) .theta(args.common.theta) .workers(n_workers) - .max_open(max_open); + .max_open(max_open) + .files(args.common.seqfile_paths()) + .on_progress(|p: Progress| { + let now = Instant::now(); + let dt = now.duration_since(last_t).as_secs_f64(); + if dt > 0.0 { + let instant = (p.position - last_bases) as f64 / dt; + ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate; + } + last_t = now; + last_bases = p.position; + let bp = p.position as f64; + let (count_str, rate_str) = if bp >= 1e9 { + (format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6)) + } else { + (format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6)) + }; + pb.set_message(format!("{count_str} {rate_str}")); + }); router - .run( - args.common.seqfile_paths(), - Some(|p: Progress| { - let now = Instant::now(); - let dt = now.duration_since(last_t).as_secs_f64(); - if dt > 0.0 { - let instant = (p.position - last_bases) as f64 / dt; - ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate; - } - last_t = now; - last_bases = p.position; - let bp = p.position as f64; - let (count_str, rate_str) = if bp >= 1e9 { - (format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6)) - } else { - (format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6)) - }; - pb.set_message(format!("{count_str} {rate_str}")); - }), - ) + .run() .unwrap_or_else(|e| { eprintln!("error: {e}"); std::process::exit(1); @@ -297,7 +297,8 @@ pub fn run(args: IndexArgs) { let t = Stage::start("dereplicate"); let pb = progress_bar("dereplication", idx.n_partitions() as u64, "partitions"); Dereplicator::new(&idx) - .run(Some(|_: Progress| pb.inc(1))) + .on_progress(|_: Progress| pb.inc(1)) + .run() .unwrap_or_else(|e| { eprintln!("error: {e}"); std::process::exit(1); @@ -311,7 +312,8 @@ pub fn run(args: IndexArgs) { // done (`count.done`) internally once every partition succeeds. Counter::new(&idx) .keep_partial(args.keep_intermediate) - .run(Some(|_: Progress| pb.inc(1))) + .on_progress(|_: Progress| pb.inc(1)) + .run() .unwrap_or_else(|e| { eprintln!("error: {e}"); std::process::exit(1); @@ -330,7 +332,8 @@ pub fn run(args: IndexArgs) { .min_abundance(args.min_abundance) .max_abundance(args.max_abundance) .keep_intermediate(args.keep_intermediate) - .run(Some(|_: Progress| pb.inc(1))) + .on_progress(|_: Progress| pb.inc(1)) + .run() .unwrap_or_else(|e| { eprintln!("error: {e}"); std::process::exit(1); diff --git a/src/obikphylo/Cargo.toml b/src/obikphylo/Cargo.toml index ef44f88b..df689c2c 100644 --- a/src/obikphylo/Cargo.toml +++ b/src/obikphylo/Cargo.toml @@ -18,7 +18,8 @@ rayon = "1" tracing = "0.1.44" [dev-dependencies] -obiread = { path = "../obiread" } -obikindexer = { path = "../obikindexer" } +obiread = { path = "../obiread" } +obikindexer = { path = "../obikindexer" } +obikalgorithm = { path = "../obikalgorithm" } tempfile = "3" tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index ad2972cd..85acd873 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -6,6 +6,7 @@ use obikindex::layer::MphfLayer; use obisys::Reporter; use tempfile::tempdir; +use obikalgorithm::Algorithm; use obikindexer::algorithms::counter::Counter; use obikindexer::algorithms::dereplicator::Dereplicator; use obikindexer::algorithms::layer_builder::LayerBuilder; @@ -68,11 +69,11 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { evidence: obikindex::layer::IndexMode::Exact, block_bits: 0, }; - let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label))) + let idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label))) .expect("create"); 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(&idx); for page in stream { let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0); router.write_batch(batch).expect("write_batch"); @@ -80,10 +81,10 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { router.close().expect("close partition writers"); // also marks scatter done drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope - Dereplicator::new(&idx).run(None::).expect("dereplicate"); - Counter::new(&idx).run(None::).expect("count_kmer"); // also writes the spectrum + marks counted + Dereplicator::new(&idx).run().expect("dereplicate"); + Counter::new(&idx).run().expect("count_kmer"); // also writes the spectrum + marks counted LayerBuilder::new(&idx) - .run(None::) + .run() .expect("build_layers"); // also marks indexed idx }