Introduce concurrency-bounded parallel processing for index cache layers
Implement an `obipipeline::Throttle` with an RAII guard to acquire and release concurrency slots. Expose new bounded parallel methods on `IndexCache` to process cached layers with a configurable cap. Refactor downstream aggregation logic to use single-pass parallel map-reduce instead of manual collect-map-reduce sequences, enforcing a maximum of 8 concurrent layer scans to bound memory usage.
This commit is contained in:
Eric Coissac
@@ -2,16 +2,17 @@
|
||||
//! (summing per-layer partials into one global partial) and the final
|
||||
//! distance calculation are both phylogenetic domain logic, so both live
|
||||
//! here, not in `obikidxcache` (which only ever hands out already-opened
|
||||
//! `KmerLayer`s via `IndexCache::iter()` — plain exposition, no
|
||||
//! computation of its own).
|
||||
//! `KmerLayer`s, via plain exposition — `IndexCache::iter()` — or bounded
|
||||
//! parallel map-reduce over them — `IndexCache::par_map_reduce()` — neither
|
||||
//! of which computes anything phylo-specific itself).
|
||||
//!
|
||||
//! [`AggregatedCount`]/[`AggregatedBit`] exist purely to "hang"
|
||||
//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a borrowed
|
||||
//! collection of `KmerLayer`s, so their already-written finalisation
|
||||
//! methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can be reused
|
||||
//! as-is instead of re-derived here — see `KmerLayer`'s own
|
||||
//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for
|
||||
//! the per-layer primitives these sum.
|
||||
//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a bounded
|
||||
//! parallel reduction over `IndexCache`'s layers, so their already-written
|
||||
//! finalisation methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can
|
||||
//! be reused as-is instead of re-derived here — see `KmerLayer`'s own
|
||||
//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for the
|
||||
//! per-layer primitives these sum.
|
||||
|
||||
use ndarray::{Array1, Array2};
|
||||
use obicompactvec::{BitPartials, ColumnWeights, CountPartials};
|
||||
@@ -64,77 +65,118 @@ impl DistanceMetric {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Aggregators — sum every cached layer's own partials into one global
|
||||
// partial, then hand the result off to `CountPartials`/`BitPartials`'s own
|
||||
// provided finalisation methods.
|
||||
// ── Aggregators — bounded parallel reduction of every cached layer's own
|
||||
// partial (of the wanted content kind) into one global partial, then handed
|
||||
// off to `CountPartials`/`BitPartials`'s own provided finalisation methods.
|
||||
// Layers of the *other* content kind contribute their zero element, same
|
||||
// as a layer that's simply absent — cheaper than pre-collecting a filtered
|
||||
// `Vec` first, since `IndexCache::par_map_reduce` never allocates one.
|
||||
|
||||
fn sum_array1(n: usize, items: impl Iterator<Item = Array1<u64>>) -> Array1<u64> {
|
||||
items.fold(Array1::zeros(n), |acc, x| acc + x)
|
||||
}
|
||||
fn sum_array2_u64(n: usize, items: impl Iterator<Item = Array2<u64>>) -> Array2<u64> {
|
||||
items.fold(Array2::zeros((n, n)), |acc, x| acc + x)
|
||||
}
|
||||
fn sum_array2_f64(n: usize, items: impl Iterator<Item = Array2<f64>>) -> Array2<f64> {
|
||||
items.fold(Array2::zeros((n, n)), |acc, x| acc + x)
|
||||
}
|
||||
/// Layers scanned at once per `par_map_reduce` call — bounds resident
|
||||
/// memory (each layer's own partial-computation working set), independent
|
||||
/// of how many `rayon` threads exist. See `IndexCache::par_map_reduce`'s
|
||||
/// own docs for why this is a separate knob from thread count.
|
||||
const MAX_CONCURRENT_LAYERS: usize = 8;
|
||||
|
||||
struct AggregatedCount<'a> {
|
||||
cache: &'a IndexCache,
|
||||
n_genomes: usize,
|
||||
layers: Vec<&'a KmerLayer>,
|
||||
}
|
||||
|
||||
impl AggregatedCount<'_> {
|
||||
fn reduce<T: Send>(
|
||||
&self,
|
||||
identity: impl Fn() -> T + Sync,
|
||||
per_layer: impl Fn(&KmerLayer) -> T + Sync,
|
||||
combine: impl Fn(T, T) -> T + Sync,
|
||||
) -> T {
|
||||
self.cache.par_map_reduce(
|
||||
MAX_CONCURRENT_LAYERS,
|
||||
&identity,
|
||||
|layer| if layer.content() == LayerContent::Count { per_layer(layer) } else { identity() },
|
||||
&combine,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ColumnWeights for AggregatedCount<'_> {
|
||||
fn col_weights(&self) -> Array1<u64> {
|
||||
sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights()))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array1::zeros(n), |l| l.col_weights(), |a, b| a + b)
|
||||
}
|
||||
}
|
||||
|
||||
impl CountPartials for AggregatedCount<'_> {
|
||||
fn partial_bray(&self) -> Array2<u64> {
|
||||
sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_bray()))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array2::zeros((n, n)), |l| l.partial_bray(), |a, b| a + b)
|
||||
}
|
||||
fn partial_euclidean(&self) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_euclidean()))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array2::zeros((n, n)), |l| l.partial_euclidean(), |a, b| a + b)
|
||||
}
|
||||
fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
|
||||
let n = self.n_genomes;
|
||||
self.layers.iter().map(|l| l.partial_threshold_jaccard(threshold)).fold(
|
||||
(Array2::zeros((n, n)), Array2::zeros((n, n))),
|
||||
|(inter, union), (i, u)| (inter + i, union + u),
|
||||
self.reduce(
|
||||
|| (Array2::zeros((n, n)), Array2::zeros((n, n))),
|
||||
|l| l.partial_threshold_jaccard(threshold),
|
||||
|(ia, ua), (ib, ub)| (ia + ib, ua + ub),
|
||||
)
|
||||
}
|
||||
fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_bray(global)))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array2::zeros((n, n)), |l| l.partial_relfreq_bray(global), |a, b| a + b)
|
||||
}
|
||||
fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_euclidean(global)))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array2::zeros((n, n)), |l| l.partial_relfreq_euclidean(global), |a, b| a + b)
|
||||
}
|
||||
fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_hellinger(global)))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array2::zeros((n, n)), |l| l.partial_hellinger(global), |a, b| a + b)
|
||||
}
|
||||
}
|
||||
|
||||
struct AggregatedBit<'a> {
|
||||
cache: &'a IndexCache,
|
||||
n_genomes: usize,
|
||||
layers: Vec<&'a KmerLayer>,
|
||||
}
|
||||
|
||||
impl AggregatedBit<'_> {
|
||||
fn reduce<T: Send>(
|
||||
&self,
|
||||
identity: impl Fn() -> T + Sync,
|
||||
per_layer: impl Fn(&KmerLayer) -> T + Sync,
|
||||
combine: impl Fn(T, T) -> T + Sync,
|
||||
) -> T {
|
||||
self.cache.par_map_reduce(
|
||||
MAX_CONCURRENT_LAYERS,
|
||||
&identity,
|
||||
|layer| if layer.content() == LayerContent::Presence { per_layer(layer) } else { identity() },
|
||||
&combine,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ColumnWeights for AggregatedBit<'_> {
|
||||
fn col_weights(&self) -> Array1<u64> {
|
||||
sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights()))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array1::zeros(n), |l| l.col_weights(), |a, b| a + b)
|
||||
}
|
||||
}
|
||||
|
||||
impl BitPartials for AggregatedBit<'_> {
|
||||
fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
|
||||
let n = self.n_genomes;
|
||||
self.layers.iter().map(|l| l.partial_jaccard()).fold(
|
||||
(Array2::zeros((n, n)), Array2::zeros((n, n))),
|
||||
|(inter, union), (i, u)| (inter + i, union + u),
|
||||
self.reduce(
|
||||
|| (Array2::zeros((n, n)), Array2::zeros((n, n))),
|
||||
|l| l.partial_jaccard(),
|
||||
|(ia, ua), (ib, ub)| (ia + ib, ua + ub),
|
||||
)
|
||||
}
|
||||
fn partial_hamming(&self) -> Array2<u64> {
|
||||
sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_hamming()))
|
||||
let n = self.n_genomes;
|
||||
self.reduce(|| Array2::zeros((n, n)), |l| l.partial_hamming(), |a, b| a + b)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -167,8 +209,7 @@ impl Metrics for IndexCache {
|
||||
let kmer_size = self.meta().config().kmer_size;
|
||||
|
||||
if use_counts {
|
||||
let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Count).collect();
|
||||
let agg = AggregatedCount { n_genomes, layers };
|
||||
let agg = AggregatedCount { cache: self, n_genomes };
|
||||
|
||||
let matrix = match metric {
|
||||
DistanceMetric::BrayCurtis => agg.bray_dist_matrix(),
|
||||
@@ -195,8 +236,7 @@ impl Metrics for IndexCache {
|
||||
|
||||
Ok(DistanceOutput { matrix, shared_kmers: shared })
|
||||
} else {
|
||||
let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Presence).collect();
|
||||
let agg = AggregatedBit { n_genomes, layers };
|
||||
let agg = AggregatedBit { cache: self, n_genomes };
|
||||
|
||||
let matrix = match metric {
|
||||
DistanceMetric::Jaccard => agg.jaccard_dist_matrix(),
|
||||
|
||||
Reference in New Issue
Block a user