diff --git a/src/alignmentsieve.rs b/src/alignmentsieve.rs index 70e917077a..442dfb4ef0 100644 --- a/src/alignmentsieve.rs +++ b/src/alignmentsieve.rs @@ -4,6 +4,7 @@ use pyo3::prelude::*; use pyo3::types::PyList; use rust_htslib::bam::record::CigarString; use rust_htslib::bam::{self, Header, Read, Reader, Writer}; +use rust_htslib::tpool::ThreadPool; use std::collections::HashMap; use std::fs::File; use std::io::{BufWriter, Write}; @@ -108,7 +109,13 @@ pub fn r_alignmentsieve( None, ); - // Open output writers + // Open output writers. When both the main and filtered-out BAM writers are active, + // give them a single shared htslib thread pool instead of calling `set_threads` on + // each independently: `set_threads` spins up its own private pool per call, so two + // independent calls with `writerthreads` each would request `2 * writerthreads` BGZF + // compression threads total, oversubscribing the `writerthreads` budget we intended. + let writer_pool = ThreadPool::new(writerthreads as u32).ok(); + let mut obam = if !bed { Some( Writer::from_path(ofile, &header, bam::Format::Bam) @@ -118,7 +125,14 @@ pub fn r_alignmentsieve( None }; if let Some(ref mut w) = obam { - let _ = w.set_threads(writerthreads); + match &writer_pool { + Some(pool) => { + let _ = w.set_thread_pool(pool); + } + None => { + let _ = w.set_threads(writerthreads); + } + } } let mut obed = if bed { Some(BufWriter::new(File::create(ofile).unwrap_or_else(|e| { @@ -143,7 +157,14 @@ pub fn r_alignmentsieve( None }; if let Some(ref mut w) = ofilterbam { - let _ = w.set_threads(writerthreads); + match &writer_pool { + Some(pool) => { + let _ = w.set_thread_pool(pool); + } + None => { + let _ = w.set_threads(writerthreads); + } + } } let mut ofilterbed = if write_filters && bed { Some(BufWriter::new( diff --git a/src/bamcompare.rs b/src/bamcompare.rs index 85dd8b8fa8..b908a78b83 100644 --- a/src/bamcompare.rs +++ b/src/bamcompare.rs @@ -395,7 +395,7 @@ pub fn r_bamcompare( } else { Box::new(raw_lines) }; - write_covfile(lines, ofile, ofiletype, chromsizes); + write_covfile(lines, ofile, ofiletype, chromsizes, nproc); Ok(()) } diff --git a/src/bamcoverage.rs b/src/bamcoverage.rs index 02654de209..df0a6acfd2 100644 --- a/src/bamcoverage.rs +++ b/src/bamcoverage.rs @@ -289,6 +289,6 @@ pub fn r_bamcoverage( println!("Writing output to: {}", ofile); } - write_covfile(lines, ofile, ofiletype, chromsizes); + write_covfile(lines, ofile, ofiletype, chromsizes, nproc); Ok(()) } diff --git a/src/computematrix.rs b/src/computematrix.rs index ed00eeeb46..8594e89cfd 100644 --- a/src/computematrix.rs +++ b/src/computematrix.rs @@ -273,28 +273,45 @@ pub fn r_computematrix( // Discriminate between reference-point and scale-regions mode. - let matrix: Vec> = pool.install(|| { - bw_files + // bwintervals is single-threaded per bigwig file (BigWigRead::get_interval takes + // &mut self, so one reader can't be shared across threads). Parallelizing only + // across bw_files.par_iter() therefore leaves most of a reserved node idle whenever + // there are fewer tracks than threads, which is the common case. Add a second + // parallelism axis by also chunking the region list per file, mirroring the + // "open one reader per worker, reuse across a chunk of regions" pattern bam_pileup + // already uses for BAM reading. Chunk count in one dimension is intentionally + // rounded up so busy still finishes at nproc-ish task count even with 1 bigwig file. + let chunks_per_file = (nproc / bw_files.len().max(1)).max(1); + let region_chunks = chunk_ranges(regions.len(), chunks_per_file); + let tasks: Vec<(usize, usize, usize)> = (0..bw_files.len()) + .flat_map(|fi| region_chunks.iter().map(move |&(s, e)| (fi, s, e))) + .collect(); + + // rayon's par_iter().collect() on an indexed source (a Vec, here) is + // order-preserving: `results[k]` always corresponds to `tasks[k]`, regardless of + // which task finishes first. That lets the merge below stay a plain sequential + // walk instead of needing its own bookkeeping to figure out where each chunk goes. + let results: Vec>> = pool.install(|| { + tasks .par_iter() - .map(|i| { + .map(|&(fi, s, e)| { bwintervals( - &i, - ®ions, - &slopregions, + &bw_files[fi], + ®ions[s..e], + &slopregions[s..e], &scale_regions, blacklist_index.as_deref(), ) }) - .reduce( - || vec![vec![]; regions.len()], - |mut acc, vec_of_vecs| { - for (i, inner_vec) in vec_of_vecs.into_iter().enumerate() { - acc[i].extend(inner_vec); - } - acc - }, - ) + .collect() }); + + let mut matrix: Vec> = vec![Vec::new(); regions.len()]; + for (&(_fi, s, e), chunk_result) in tasks.iter().zip(results.into_iter()) { + for (local_ix, region_ix) in (s..e).enumerate() { + matrix[region_ix].extend(chunk_result[local_ix].iter().copied()); + } + } matrix_dump( sortregions, sortusing, @@ -312,6 +329,27 @@ pub fn r_computematrix( Ok(()) } +/// Split `[0, n)` into up to `k` contiguous, roughly-equal ranges. Returns a single +/// `(0, n)` range if `n == 0` or `k <= 1`, so callers never need to special-case +/// "no chunking" separately from "one chunk". +fn chunk_ranges(n: usize, k: usize) -> Vec<(usize, usize)> { + if n == 0 || k <= 1 { + return vec![(0, n)]; + } + let k = k.min(n); + let base = n / k; + let rem = n % k; + let mut ranges = Vec::with_capacity(k); + let mut start = 0; + for i in 0..k { + let len = base + if i < rem { 1 } else { 0 }; + let end = start + len; + ranges.push((start, end)); + start = end; + } + ranges +} + fn slop_region( region: &Region, scale_regions: &Scalingregions, diff --git a/src/filehandler.rs b/src/filehandler.rs index 703b263b2e..d20f5a820e 100644 --- a/src/filehandler.rs +++ b/src/filehandler.rs @@ -31,7 +31,13 @@ pub fn bam_ispaired(bam_ifile: &str) -> bool { return false; } -pub fn write_covfile
  • (lines: LI, ofile: &str, filetype: &str, chromsizes: HashMap) +pub fn write_covfile
  • ( + lines: LI, + ofile: &str, + filetype: &str, + chromsizes: HashMap, + nproc: usize, +) where LI: Iterator, { @@ -51,12 +57,26 @@ where } } else { let vals = BedParserStreamingIterator::wrap_infallible_iter(lines, false); - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(1) - .build() - .expect("Unable to create tokio runtime for bw writing."); - let writer = BigWigWrite::create_file(ofile, chromsizes) + let mut writer = BigWigWrite::create_file(ofile, chromsizes) .unwrap_or_else(|e| panic!("Failed to create output bigwig file '{}': {}", ofile, e)); + // bigtools spawns one async task per chromosome (zoom aggregation + section + // encoding) and funnels their output through a single sequential writer task + // that owns the actual file handle; the crate's own channels serialize that + // hand-off, so raising worker count only widens the per-chromosome encode + // stage, it does not add any concurrent writers to `ofile`. + // Mirrors bigtools' own bedgraphtobigwig CLI: skip the multi-thread runtime + // entirely at nproc == 1, since it would only ever have one consumer. + let runtime = if nproc <= 1 { + writer.options.channel_size = 0; + tokio::runtime::Builder::new_current_thread() + .build() + .expect("Unable to create tokio runtime for bw writing.") + } else { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(nproc) + .build() + .expect("Unable to create tokio runtime for bw writing.") + }; let _ = writer.write(vals, runtime); } } @@ -657,8 +677,8 @@ pub fn chrombounds_from_bam(bamfiles: Vec<&str>) -> HashMap { pub fn bwintervals( bwfile: &str, - regions: &Vec, - slopregions: &Vec>, + regions: &[Region], + slopregions: &[Vec], scale_regions: &Scalingregions, blacklist_index: Option<&crate::filtering::BlacklistIndex>, ) -> Vec> {