Introduce obikalgorithm crate to unify pipeline algorithms
Define a shared Algorithm trait with an associated Output type and a parameterless run(&mut self) method. Refactor PartitionRouter, Dereplicator, Counter, and LayerBuilder to implement the trait, standardizing receivers to &mut self and moving configuration and progress callbacks to dedicated builder setters. Decouple error handling using a generic boxed error type and update workspace dependencies accordingly.
This commit is contained in:
Eric Coissac
Generated
+7
@@ -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",
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
[package]
|
||||
name = "obikalgorithm"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
@@ -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<E: Error + Send + Sync> for
|
||||
/// Box<dyn Error + Send + Sync>` — no dependency on the crate that defines
|
||||
/// `E` needed here.
|
||||
pub type Error = Box<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
|
||||
/// 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<Self::Output>;
|
||||
}
|
||||
@@ -5,6 +5,7 @@ edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikindex = { path = "../obikindex" }
|
||||
obikalgorithm = { path = "../obikalgorithm" }
|
||||
obikseq = { path = "../obikseq" }
|
||||
obidebruinj = { path = "../obidebruinj" }
|
||||
obiskio = { path = "../obiskio" }
|
||||
|
||||
@@ -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<Box<dyn Fn(Progress) + Sync + 'a>>,
|
||||
}
|
||||
|
||||
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<impl Fn(Progress) + Sync>) -> SKResult<KmerSpectrum> {
|
||||
/// `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<KmerSpectrum> {
|
||||
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<SKResult<()>> = (0..self.n_partitions)
|
||||
let results: Vec<SKResult<()>> = (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()))?;
|
||||
|
||||
|
||||
@@ -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<Box<dyn Fn(Progress) + Sync + 'a>>,
|
||||
}
|
||||
|
||||
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<impl Fn(Progress) + Sync>) -> 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<SKResult<()>> = (0..self.n_partitions)
|
||||
let results: Vec<SKResult<()>> = (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
|
||||
})
|
||||
|
||||
@@ -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<u32>,
|
||||
keep_intermediate: bool,
|
||||
on_progress: Option<Box<dyn FnMut(Progress) + Send + 'a>>,
|
||||
}
|
||||
|
||||
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<impl FnMut(Progress) + Send>) -> SKResult<usize> {
|
||||
fn run(&mut self) -> obikalgorithm::Result<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 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<usize> = (0..self.n_partitions).collect();
|
||||
let order: Vec<usize> = (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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Box<dyn Iterator<Item = PathBuf> + Send>>,
|
||||
on_progress: Option<Box<dyn FnMut(Progress) + 'a>>,
|
||||
}
|
||||
|
||||
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<Item = PathBuf> + 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<Item = PathBuf> + Send + 'static,
|
||||
mut on_progress: Option<impl FnMut(Progress)>,
|
||||
) -> 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::<u64>();
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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::<fn(obisys::Progress)>).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::<fn(obisys::Progress)>).unwrap();
|
||||
let spectrum = Counter::new(&index).run().unwrap();
|
||||
(spectrum.f0, spectrum.f1)
|
||||
}
|
||||
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
@@ -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::<fn(obisys::Progress)>).expect("dereplicate");
|
||||
Counter::new(&idx).run(None::<fn(obisys::Progress)>).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::<fn(obisys::Progress)>)
|
||||
.run()
|
||||
.expect("build_layers"); // also marks indexed
|
||||
idx
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user