Rust backend: parallelism fixes for bigWig writer, alignmentSieve, computeMatrix - #1445
Rust backend: parallelism fixes for bigWig writer, alignmentSieve, computeMatrix#1445adRn-s wants to merge 5 commits into
Conversation
This comment has been minimized.
This comment has been minimized.
|
Four commits, three "confirmed against pinned source" citations, a spec doc with a section literally titled "Memory tradeoff of the combined change set," and not one single em dash anywhere in sight. Either this was written by someone who's spent way too long reading bigtools internals for fun, or it's the most thoroughly laundered AI slop you've ever reviewed. Benchmarks on raven will settle which... 👍🏽 maybe it's both |
|
Pushed one more commit (349c594): gates the parallel indexed-read path (item 1b) behind a size threshold on the merged bedgraph. This was flagged in the spec as the expected fallback if benchmarks disagreed on item 1, and raven's 24-thread run confirmed it does: bamCoverage/bamCompare on WGS-scale samples regressed to 0.78-0.85x wall time under the always-on parallel path, while chip/rna-scale samples sped up up to 3.8x, with memory flat in both cases. Full numbers and rationale are in the commit message. Set up a proper A/B benchmark Snakefile (4.0.0 pre-branch vs this branch, side by side, two pixi environments so both binaries coexist) on https://github.com/deeptools/deeptools25/tree/dt4-ab-benchmark, |
349c594 to
117f41e
Compare
|
Status update, replacing the outdated spec above (minimized, kept for history). Branch is now rebased onto latest
Worth calling out: the rebase surfaced that upstream's gz-support PR had already restructured 232 tests passing (up from 153 pre-rebase, thanks to #1443's exhaustive suite). Re-running the raven A/B benchmark (4.0.0 pre-branch vs this branch) now that the rebase is in, still draft until those numbers are in. |
31a9852 to
052bdf5
Compare
|
Re-ran the A/B benchmark after rebasing on Reverted
Back at parity with baseline everywhere. The two outliers above 1.5x run identical code to Still working through computeMatrix/alignmentSieve/multibamSummary numbers before deciding if this is ready to leave draft. |
write_covfile hardcoded worker_threads(1), starving bigtools' own per-chromosome zoom/encode parallelism regardless of -p. Pass nproc through from bamCoverage/bamCompare instead. At nproc == 1, mirror bigtools' own bedgraphtobigwig CLI and use a current_thread runtime (channel_size = 0) rather than paying for a multi-thread runtime that will only ever have one consumer. Verified byte-identical bigWig output across worker counts (1/4/12) on single- and multi-chromosome fixtures; full pytest suite passes.
set_threads() spins up a private htslib thread pool per call. With --filteredOutReads set, both obam and ofilterbam independently requested writerthreads threads, so total BGZF compression threads could reach 2 * writerthreads against a writerthreads (or nproc) budget. Build one ThreadPool and share it via set_thread_pool() instead, per rust-htslib's own documented pattern for multiple concurrent readers/writers (bam/mod.rs tests, v1.0.1). Verified byte-identical main and filtered-out BAM output across -p 1/4/12 with both writers active; full pytest suite passes.
Matrix building only parallelized across bw_files.par_iter(); each task then processed every region for that file in a strictly sequential loop (bwintervals reuses one BigWigRead reader, since get_interval takes &mut self and can't be shared across threads). So core utilization was bounded by min(nproc, bw_files.len()), leaving most of a reserved node idle on typical single- or few-sample runs. Widen bwintervals to take region/slopregion slices instead of &Vec, and parallelize over (bigwig file, region chunk) pairs instead of just bigwig file, opening one reader per task and reusing it across that task's chunk of regions, mirroring the pattern bam_pileup already uses for BAM reading. Chunk count targets roughly nproc total tasks (nproc / bw_files.len(), rounded up to 1) so a single-track run still fans out across the whole reservation. rayon's par_iter().collect() into a Vec is order-preserving for an indexed source, so results[k] always corresponds to tasks[k] regardless of completion order; the merge step is a plain sequential walk rather than needing its own reordering logic. Verified the raw --outFileNameMatrix output is identical across -p 1/4/12 with two bigwig tracks, and matches the pre-change baseline exactly. Full pytest suite passes.
wrap_infallible_iter already drives BigWigWrite directly off the in-memory iterator with no disk round-trip; that stays the path at nproc <= 1; skip building a multi-thread runtime there too, since it would only ever have one consumer (mirrors bigtools' own bedgraphtobigwig CLI at nthreads == 1). At nproc > 1, spool to one merged bedgraph temp file and hand it to bigtools' own indexed parallel reader (BedParserParallelStreamingIterator), the same path bigtools' own bedgraphtobigwig CLI uses by default. That reads and encodes multiple chromosomes concurrently straight off disk, instead of the single sequential stream wrap_infallible_iter produces. This unavoidably costs one extra write pass versus the nproc <= 1 path, since index_chroms() needs a real sorted file on disk to build its byte-offset index. Gate that parallel read behind a size threshold on the merged file, inverted from bigtools' own ~200MB auto-enable-above heuristic. Raven 24-thread benchmarks (this branch vs 4.0.0 pre-branch, 10 reps) showed it as a real win on chip/rna-scale outputs (up to ~3.8x faster) but a regression on whole-genome-scale ones (bamCoverage/bamCompare on WGS samples dropped to 0.78-0.85x), with memory flat in both cases. Likely cause: index_chroms() does one full linear scan of the merged bedgraph before BedParserParallelStreamingIterator does its own per-chromosome reads; cheap next to the parse time it saves on small/medium files, but dominant on multi-hundred-MB WGS ones. Data-driven, not a hunch: this exact fallback was flagged in the parallelism spec's "alternative if benchmarks disagree" section for this item before benchmarking. Full pytest suite passes.
…put size" This reverts commit 117f41e. The size-threshold gate this commit added doesn't fix the WGS regression it was meant to fix, and it regresses chip/rna-scale outputs it claimed to speed up. Root cause: at nproc>1, the reverted code always spools `lines` to a merged temp bedgraph file before deciding anything - the gate only chooses between an indexed parallel *read* of that file (under threshold) or a plain sequential *read* of it (over threshold). Either way the disk round-trip (format to text, write, reopen, reparse) happens unconditionally, on top of whatever the read path costs. The pre-regression baseline never touched disk here at all: it streamed `lines` straight into BigWigWrite in memory. Fresh raven 24-thread A/B benchmark (10 reps, this branch rebased onto 4.0.0 117f41e's parent vs pre-branch 4.0.0) confirms this end to end - bamCoverage/bamCompare regressed across every chip/wgs sample tested, not just WGS as the original commit's benchmark suggested: bcov_human_wgs 0.68x bcom_human_chip 0.74x bcov_triticum_wgs 0.58x bcom_triticum_chip 0.73x bcov_human_chip 0.67x bcov_triticum_chip 0.95x Deciding the fast path correctly would require knowing the output size before consuming `lines` (a single-pass streaming iterator), which isn't available cheaply for bamCoverage/bamCompare's use case. Given real deepTools output sizes appear to exceed the 200MB threshold routinely (bedgraph text is far less compact than the bigwig it produces), the disk-based parallel-read path doesn't pay for itself in practice. Reverting restores the proven-correct in-memory streaming write for nproc>1, matching 4.0.0 baseline exactly.
052bdf5 to
5291287
Compare
|
Quick disclaimer up front: computeMatrix numbers below are stale and a rerun is in progress, see the note at the end before reading too much into them. Full A/B table after the bigwig-writer revert (10 reps, raven 24-thread profile), for the record:
alignmentSieve, bamCoverage/bamCompare and multibamSummary are all clean, matching the expected effect of each commit. multibamSummary is untouched by any of this branch's commits and sits flat at ~1.00x, which is a good sanity check that the rest isn't noise. The computeMatrix row is where it gets complicated. Our Rebased the branch onto current |
|
Dug into why bamCoverage's triticum_chip (1.53x) and triticum_rna (3.59x) stand out from the rest of the table, which is otherwise close to parity. Short version: it's real, not noise, and it comes from genome size.
Checked the two reference genomes used in the benchmark:
Triticum is about 4.7x bigger genome-wide, in 22 uniformly large chromosomes. Human has more contigs on paper, but 169 of them are unplaced/alt scaffolds that finish almost instantly. The real work is in ~25 chromosomes, the largest of which is a third the size of triticum's largest. So parallelizing the same write stage has far more raw work to divide across threads for triticum. That also explains why WGS shows no gain in either species (0.98x both) while chip/rna do. WGS read depth is high enough that the upstream BAM coverage computation dominates total wall time, so a faster write stage barely moves the total. Chip/rna have much lower depth, so the write stage (which scales with genome length, not depth) is a bigger share of runtime in both species, and that's where triticum's heavier per-chromosome write cost, now finally parallelized instead of serial, shows up. computeMatrix numbers are still pending the rerun against the corrected baseline mentioned above, will follow up separately once that's done. |
|
Updated the PR description to match. Corrected computeMatrix numbers, rerun against the fixed baseline (current
For reference, the old (stale-baseline) numbers were human_15 1.01x/0.82, human_3 2.23x/1.20, triticum_15 1.03x/0.85, triticum_3 1.54x/1.03. The The One thing to call out: So the chunking commit still earns its complexity on the low-bigwig-count case, but it's a real speed/memory tradeoff rather than a free win. The PR description should say so instead of leading with the old 2.23x number. |
1 participant
Welcome to deepTools GitHub repository! Please check the following regarding
your pull request :
Summary
Three bounded parallelism fixes in the Rust backend, found while checking why bamCoverage/bamCompare's bigWig writer, alignmentSieve, and computeMatrix don't scale as well as expected with
-p. Each commit is independently buildable and testable. Rebased onto current4.0.0(0f07abf3, which independently fixed computeMatrix's hashmap-based memory footprint) to keep the comparison fair.4f483b94bigwig writer: usenprocfor the tokio runtime, skip the multi-thread runtime entirely at-p 1(mirrors bigtools' ownbedgraphtobigwigCLI behavior).03a525b5alignmentSieve: share one htslib thread pool between the main and--filteredOutReadsBAM writers, instead of two independentset_threads()calls that were oversubscribing cores (up to2 * (nproc - 2)threads competing fornprocreserved cores).ee7d9e8ccomputeMatrix: add a region-chunk axis to matrix-building parallelism, so core usage isn't bounded bymin(nproc, bw_files.len())on single- or few-sample runs.A fourth commit (
07aca1fa, reading the bigwig writer's input from a merged, indexed bedgraph instead of a sequential one) is also in the branch history but reverted (5291287b): benchmarking showed it wrote a temp file to disk at-p> 1 that the pre-change baseline never touched, which made it strictly worse than baseline above roughly 200MB of output, i.e. most chip/wgs samples. Kept thebwintervalsslice-signature widening it introduced, since commit 3 needs it.Root causes for items 1-3 were confirmed against the actual pinned library source, not just documentation:
rust-htslibv1.0.1 (cloned at that tag) for item 2, andbigtools0.5.8 (cloned at HEAD, one commit past our pinned tag) for item 1.Testing
4.0.0).-p 1/4/12for the affected tools.-p 24on the file axis alone (15 files here). With few files (3 here), the region-chunking commit still gives a real ~1.7x speedup, down from an earlier (invalid, stale-baseline) reading of up to 2.23x now that both sides share4.0.0's hashmap-to-vector fix. That speedup comes with a real memory cost on the largest genome (triticum, 1.19x baseline RSS), not the near-free win it first looked like.cc colleagues for input before this merges.