perf(bam): parallel BGZF and record encoding (draft: measured slower, investigating) - #164
perf(bam): parallel BGZF and record encoding (draft: measured slower, investigating)#164BenjaminDEMAILLE wants to merge 1 commit into
Conversation
BGZF deflate ran one block at a time on the thread doing the writing, while `noodles-bgzf` with the `libdeflate` feature and `libdeflater` were already dependencies and rayon's pool was already sized from `--runThreadN`. Record encoding was serial on that same thread. `make_bgzf_writer` now returns `bgzf::io::MultithreadedWriter`, and `encode_batch` splits a batch into sub-ranges encoded across the pool. Output is unchanged, and provably so rather than by inspection. The multithreaded writer stages into the same `MAX_BUF_SIZE` buffer, emits through the same `write_frame` and appends the same EOF block as the serial writer, and libdeflate is deterministic at a fixed level. `bam::io::Writer` keeps no state between records beyond a scratch buffer it clears per record, so a sub-range encoded against a fresh in-memory writer contributes exactly the bytes it would have contributed to the serial stream. One hazard, found by a test that hung: the multithreaded writer hands blocks to the pool and blocks on a channel bounded by the worker count, so writing to it from a rayon worker can deadlock at any pool size. rustar writes from the pipeline's dispatcher thread, never from a worker. A `debug_assert` now fails loudly instead of hanging if a future call site breaks that. The sorted writers encode in bounded slices rather than in one pass, so peak RSS does not gain a copy of the whole encoded stream.
|
Measured it. It is slower, not faster. Moving to draft. Machine at 87-92% CPU idle, 2M yeast reads (
All four compression-6 pairs go the same way: +1.8 s median on a 25.3 s run, about +7%. Why I think it regressedTo be checked rather than announced: at If that is right, the premise of this PR was wrong in the same way #153's was, just in the other direction: I assumed a serialised bottleneck parallelises, without checking there were free cores to absorb it. What is running nowA third binary, Three outcomes and what each means:
What stands regardlessThe output-neutrality work is independent of the performance question and holds either way: byte-identical BAM against So does the deadlock hazard documented on I will post the isolation numbers here and either narrow this PR to whatever actually wins or close it with the negative result recorded. It should not merge on the description it currently has. |
|
Isolation run done. The result picks outcome 1, and it also refutes the cause I proposed above. Three binaries, 2M yeast reads,
Two things fall out. The multithreaded BGZF writer is the whole regression. My "no spare cores" explanation was wrong. I said the pool is saturated at 16 threads so compression has nowhere to go. At 8 threads on 16 cores there are eight idle cores and the regression is larger, not smaller: +3.2 s, about +9%, with a 0.17 s spread across three rounds. Spare capacity does not rescue it, so capacity was not the constraint. What fits both thread counts is backpressure rather than capacity. RecommendationThis diff should not merge, in whole or in part. Both halves fail on their own terms: the writer half regresses at every thread count measured, and the encoding half is neutral, which under Three things from the work are worth keeping regardless of what happens to the diff, and none of them need this branch:
Happy to close this. If any of the three points above is worth landing on its own — the bench script with the |
BAM compression and record encoding move onto the rayon pool. Output-neutral, and verified byte-identical rather than assumed.
What changed
make_bgzf_writerreturnedbgzf::io::Writer, which deflates one block at a time on whichever thread is writing. It now returnsbgzf::io::MultithreadedWriter, which hands each block to the rayon pool and writes finished frames in order from a dedicated thread.run()already sizes the global pool from--runThreadN, andnoodles-bgzfwith thelibdeflatefeature pluslibdeflaterwere already dependencies, so no new dependency and no new configuration.encode_batchsplits a batch into sub-ranges encoded across the pool. Record encoding was serial on the writing thread;bam::io::Writer::write_alignment_recordruns the noodles encoder per record, and that is real CPU work in a BAM run.The sorted writers encode in bounded slices of 65 536 records rather than in one pass. Their whole record set is already resident as
RecordBufs, and encoding it in one go would have held the entire encoded stream alongside them.Why the output cannot move
Not an assertion about intent. The two writers are the same format machine:
MAX_BUF_SIZEand emit a block only when it fills,write_frame,BGZF_EOF,And
bam::io::Writerkeeps no state between records other than a scratch buffer it clears at the top of everywrite_alignment_record. A sub-range encoded against a fresh in-memory writer therefore contributes exactly the bytes it would have contributed to the serial stream, so concatenating sub-ranges in input order is the serial stream.par_chunksis an indexed parallel iterator, socollectrestores input order whatever the completion order.Measured on 200 000 real reads (yeast,
ERR12389696), old binary fromorigin/mainagainst this branch:BAM UnsortedAligned.out.bamBAM SortedByCoordinateAligned.sortedByCoord.out.bamWorth saying how that check went wrong first, because it is an easy trap: the files initially differed by one byte. The cause was the harness, not the code.
@PGrecordsCL:verbatim, so running./oldwith prefixout_old/against./newwith prefixout_new/makes the header differ and the compressed size with it. Both sides now run as./rustar-alignerwith prefix./from inside their own directory.One hazard, found by a test that hung
The multithreaded writer hands blocks to the pool via
rayon::spawnand blocks on a channel bounded by the worker count while a dedicated thread drains it. A caller that is itself occupying a rayon worker can fill that channel and then wait on a compression task that has no free worker to run on. That is a deadlock, not a slowdown, and it is reachable at any pool size, including one worker.I found it by writing a thread-count sweep that wrapped the whole writer in
ThreadPool::install. It hung; a stack sample showed the worker parked inLockLatch::wait_and_reset.rustar is safe as written:
run_batch_pipelinecallsconsumeon its own dispatcher thread, never on the pool, and the transcriptome writer is reached through a&mutthat cannot crosspar_iter. To keep it that way,write_batchcarriesso a future call site that writes from inside the pool fails loudly in tests rather than hanging in production. The constraint is also documented on
make_bgzf_writer, where someone changing the writer will read it.Verification
New unit tests in
src/io/bam.rs, each comparing against bytes produced by the pre-change path (a serialbgzf::io::Writerfed one record at a time), not merely against the other new half:multithreaded_bgzf_is_byte_identical_to_the_serial_writer— 4 000 records, several BGZF blocks, so block boundaries are actually crossedparallel_encoding_is_byte_identical_at_every_worker_count— pools of 1, 3, 4, 7 and 16 workers, including sizes that do not divide the batch evenlya_batch_below_the_fanout_threshold_encodes_the_same_bytes— the serial branch must agree with the parallel onechunked_encoding_matches_a_single_pass— the sorted writers' slice boundary must not show in the outputbam_file_is_byte_identical_to_the_serial_path— end to end throughBamWriterfinishing_twice_is_not_an_error—MultithreadedWriter::finishshuts the workers down and panics if called twice, andDropcalls it tooGate: 566 lib + 26 integration tests,
cargo clippy --all-targets -- -D warnings,cargo fmt --check, all green.Timings: measured, and they say this regresses
Draft, because the measurement contradicts the premise. At
--runThreadN 16on 16 cores, 2M yeast reads, machine at 87-92% CPU idle, interleaved with the order flipped on even pairs:origin/main)--outSAMtype NoneBAM Unsorted,--outBAMcompression 6All four compression-6 pairs go the same way: +1.8 s median on 25.3 s, about +7% slower.
Noneis neutral, which is what localises the difference to the BAM write path.Likely cause, being checked rather than announced: alignment already occupies every core at 16 threads,
MultithreadedWriterhands its blocks to that same rayon pool and blocks on a channel bounded by the worker count, so there is no spare capacity for compression to move to and the handoff is pure added cost.test/bench_bam_write.shreproduces the above. It times--outSAMtype Nonealongside the BAM modes deliberately: only the difference is the work this change touches, and a total that is mostly alignment hides it.See the comment thread for the isolation run in progress (parallel BGZF with serial encoding, plus a
--runThreadN 8leg where spare capacity exists). This PR should not merge on a performance claim until that resolves; the output-neutrality work above stands either way.