From 04e0bed4955a4f5a1efa6d6d8c0cab17081d0e3a Mon Sep 17 00:00:00 2001 From: Joshua Allen Date: Mon, 2 Nov 2026 09:54:40 -0600 Subject: [PATCH 1/2] Shuffle gen-reads FASTQ output by default (shuffle_fastq, #870) gen-reads wrote reads in genome order, which no sequencer produces; on a 21x T2T-CHM13 simulation that made bwa-mem2 need 166 GB against 90 GB for the same reads shuffled. The combine step now does a two-pass external shuffle: records (pairs together) are scattered at random into bucket files, each bucket is permuted in memory, and the buckets are concatenated. Memory is bounded by bucket size and a 4 GiB pass-2 budget; temp disk stays near one copy because inputs are deleted as read. The shuffle draws from its own child of the run RNG, so reads, the truth VCF and the golden BAM are unchanged, and output is byte-identical across thread counts. shuffle_fastq: false restores genome order. Co-Authored-By: Claude --- docs-site/src/gen-reads/fastq-shuffling.md | 59 +- .../src/gen-reads/parallel-processing.md | 2 +- eidolon-core/src/file_tools.rs | 1 + eidolon-core/src/file_tools/block_gz.rs | 9 +- eidolon-core/src/file_tools/fastq_shuffle.rs | 549 ++++++++++++++++++ eidolon/src/gen_reads/utils/config.rs | 18 + eidolon/src/gen_reads/utils/runner.rs | 49 +- eidolon/tests/common/mod.rs | 10 +- eidolon/tests/determinism.rs | 25 +- eidolon/tests/fastq_shuffle_output.rs | 155 +++++ template_config/gen_reads_template.yml | 6 + 11 files changed, 858 insertions(+), 25 deletions(-) create mode 100644 eidolon-core/src/file_tools/fastq_shuffle.rs create mode 100644 eidolon/tests/fastq_shuffle_output.rs diff --git a/docs-site/src/gen-reads/fastq-shuffling.md b/docs-site/src/gen-reads/fastq-shuffling.md index 4043fa82..1ed4a0d3 100644 --- a/docs-site/src/gen-reads/fastq-shuffling.md +++ b/docs-site/src/gen-reads/fastq-shuffling.md @@ -1,27 +1,48 @@ # FASTQ Shuffling -`eidolon` writes reads in genome order: contig by contig, by position within each contig. A -real sequencer's FASTQ is in effectively random order. Variant calling does not depend on this, -since the BAM is coordinate-sorted before calling. Tools that read the FASTQ in batches, -aligners in particular, do see the difference: a batch can be filled entirely with reads from -one region, such as a long satellite array. Shuffle before aligning if that matters to you. -#870 tracks shuffling in `eidolon` itself. +`eidolon` writes its FASTQs in **random order** by default (`shuffle_fastq: true`), as a +sequencer does. Mates stay in register: line *n* of R1 and line *n* of R2 are the two reads of +one pair. -With `seqkit`, shuffle each mate separately **with the same seed**. seqkit has no paired mode, -but its permutation depends only on the seed and the number of records, and R1 and R2 have the -same count. So the two files come out in the same order and mates stay in register: +Reads are generated chunk by chunk along the genome, so without shuffling the FASTQ comes out in +genome order: contig by contig, by position within each contig. Variant calling does not depend +on this, because the BAM is coordinate-sorted before calling. Tools that read the FASTQ in +batches do see the difference, aligners in particular: a batch can be filled entirely with reads +from one region, such as a long satellite array. Measured on a 21x simulation of T2T-CHM13 +(#870), bwa-mem2 peaked at 166 GB on reads in genome order and at 90 GB on the same reads +shuffled, with the same alignments either way. -```bash -# single-ended -seqkit shuffle -s 11 sample.fastq.gz -o sample_shuffled.fastq.gz +```yaml +shuffle_fastq: true # default; false writes reads in genome order +``` + +**What it does and does not change.** +- Only the order of FASTQ records changes. The reads themselves, the truth VCF and the golden BAM + are identical with the shuffle on or off. The shuffle draws from its own branch of the random + number generator, so it does not disturb the streams that generate reads. +- Output is deterministic: the same seed gives byte-identical FASTQs at any `num_threads`. +- Turning the shuffle off restores genome order exactly as earlier versions wrote it. + +**Cost.** The shuffle is a two-pass external sort. Records are scattered at random into bucket +files in the run's temp directory, then each bucket is shuffled in memory and written out. Buckets +are shuffled in parallel only as far as they fit in about 4 GB, so memory stays bounded at any +genome size. Temp disk stays close to one copy of the output, because each input is deleted as it +is read. -# paired-ended: same -s on both mates +Measured on T2T-CHM13 chr21 at 20x (about 3M pairs, 8 threads), against `shuffle_fastq: false`: +- wall time 112 s against 99 s; +- peak memory 2.3 GB against 0.5 GB; +- **gzipped FASTQs about 60% larger.** Reads in genome order overlap their neighbors, which gzip + exploits; shuffled reads, like a real sequencer's, compress less well. + +**Shuffling output from older versions, or from `shuffle_fastq: false`.** With `seqkit`, shuffle +each mate separately **with the same seed**. seqkit has no paired mode, but its permutation +depends only on the seed and the number of records, and R1 and R2 have the same count: + +```bash seqkit shuffle -s 11 sample_r1.fastq.gz -o sample_r1_shuffled.fastq.gz seqkit shuffle -s 11 sample_r2.fastq.gz -o sample_r2_shuffled.fastq.gz ``` -**Memory:** `seqkit shuffle` reads every record into memory, so expect memory on the order of -the uncompressed FASTQ, per mate. Its low-memory `-2/--two-pass` mode supports FASTA only. -Do not pass both mates to one `seqkit shuffle` call: they would be read as one input and shuffled -together into a single output. - -(`seqkit` is an open-source FASTQ toolkit: https://github.com/shenwei356/seqkit.) +`seqkit shuffle` reads every record into memory; its low-memory `-2/--two-pass` mode supports +FASTA only. Do not pass both mates to one call: they would be shuffled together into a single +output. (`seqkit`: https://github.com/shenwei356/seqkit.) diff --git a/docs-site/src/gen-reads/parallel-processing.md b/docs-site/src/gen-reads/parallel-processing.md index 4d8c9106..23eb73fb 100644 --- a/docs-site/src/gen-reads/parallel-processing.md +++ b/docs-site/src/gen-reads/parallel-processing.md @@ -1,7 +1,7 @@ # Parallel Processing `eidolon gen-reads` processes contigs in parallel by default using rayon's work-stealing thread pool. Each contig is an independent unit of work — variant generation, fragment sampling, and FASTQ/BAM writing all happen concurrently across contigs — so references with many contigs scale well across cores. -Output is **byte-identical regardless of `num_threads`** (the same seed always produces the same reads in the same order), so you can change the thread count freely without affecting results. +Output is **byte-identical regardless of `num_threads`** (the same seed always produces the same reads in the same order, shuffled or not; see [FASTQ Shuffling](fastq-shuffling.md)), so you can change the thread count freely without affecting results. A note on scaling: read generation is largely **memory-bandwidth bound**, so on references dominated by one or a few large chromosomes the wall-time gain from extra cores is limited — the bottleneck is memory throughput, not CPU. An experimental sub-contig **chunk size** knob (below) can split a single chromosome across cores, but it is **disabled by default** because in our benchmarks it did not improve wall time and occasionally regressed it. diff --git a/eidolon-core/src/file_tools.rs b/eidolon-core/src/file_tools.rs index c9c76b9f..f30a9f3f 100644 --- a/eidolon-core/src/file_tools.rs +++ b/eidolon-core/src/file_tools.rs @@ -6,6 +6,7 @@ pub mod bam_writer; pub mod bed_reader; pub mod block_gz; pub mod fasta_stream; +pub mod fastq_shuffle; pub mod fastq_tools; pub mod file_io; pub mod folder_tools; diff --git a/eidolon-core/src/file_tools/block_gz.rs b/eidolon-core/src/file_tools/block_gz.rs index 09a89eb3..09ba83dd 100644 --- a/eidolon-core/src/file_tools/block_gz.rs +++ b/eidolon-core/src/file_tools/block_gz.rs @@ -70,11 +70,18 @@ impl BlockGzWriter { // CompressionLvl::default() is libdeflate level 6 — the same ratio // target as flate2's Compression::default(), so output size is // comparable and only the speed changes. + Self::with_level(inner, CompressionLvl::default()) + } + + /// As [`BlockGzWriter::new`] at a chosen level. Short-lived intermediates (the + /// FASTQ shuffle's bucket files) use the fastest level: they are read once and + /// deleted, so ratio buys nothing. + pub fn with_level(inner: W, level: CompressionLvl) -> Self { Self { inner, buf: Vec::with_capacity(BLOCK_SIZE), out: Vec::new(), - compressor: Compressor::new(CompressionLvl::default()), + compressor: Compressor::new(level), finished: false, wrote_any: false, } diff --git a/eidolon-core/src/file_tools/fastq_shuffle.rs b/eidolon-core/src/file_tools/fastq_shuffle.rs new file mode 100644 index 00000000..77cad0d3 --- /dev/null +++ b/eidolon-core/src/file_tools/fastq_shuffle.rs @@ -0,0 +1,549 @@ +//! Combine per-chunk temp FASTQs into the final output in random order (#870). +//! +//! **Why.** gen-reads generates reads chunk by chunk, so concatenating the temp files gives a +//! FASTQ in genome order, which no sequencer produces. Aligners read FASTQs in batches, and a +//! batch filled with one region's reads behaves differently from a mixed one: bwa-mem2 needed +//! 166 GB on a 21x T2T-CHM13 simulation as written and 90 GB on the same reads shuffled (#870). +//! The original in-memory shuffle was removed because it held every read at once. +//! +//! **How: a two-pass external shuffle with bounded memory.** +//! 1. Scatter. Read the temp files in order, and send each record (each pair, when paired) to one +//! of `K` bucket files chosen at random. Both mates of a pair go to the same bucket, in the +//! same position. Each input file is deleted once read, so temp disk stays near one copy. +//! 2. Shuffle. Load one bucket at a time, permute it in memory, and write it as gzip; then +//! concatenate the buckets (gzip members concatenate into one valid stream). +//! +//! Bucket assignment and then a uniform permutation within each bucket give a uniform +//! permutation of the whole file. `K` is chosen so a bucket fits the memory budget. +//! +//! **Determinism.** The scatter draws from one RNG, sequentially. Each bucket's permutation draws +//! from its own child RNG. Threads only change which bucket is processed when, so output is +//! byte-identical across thread counts. + +use std::fs::{self, File}; +use std::io::{BufRead, BufReader, BufWriter, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use flate2::read::MultiGzDecoder; +use libdeflater::CompressionLvl; + +use crate::file_tools::block_gz::BlockGzWriter; +use crate::file_tools::fastq_tools::FastqToolsError; +use crate::rng::NeatRng; + +/// Uncompressed bytes per bucket, both mates together. Pass 2 holds this much per thread, +/// roughly doubled by per-record allocation. +const BUCKET_TARGET_BYTES: u64 = 256 * 1024 * 1024; +/// Upper bound on buckets, so the scatter's open files (one per bucket per mate) stay well under +/// the common 1024 soft limit. Past this, buckets grow instead. +const MAX_BUCKETS: usize = 256; +/// Gzipped FASTQ expands about 3.5-4.5x; overestimating only adds buckets. +const EXPANSION_ESTIMATE: u64 = 5; +/// Pass-2 threads, capped because each holds a whole bucket in memory. +const MAX_SHUFFLE_THREADS: usize = 8; +/// Pass 2 runs at most as many buckets at once as fit in this much memory, judged by the +/// largest bucket. Small runs get full parallelism; at genome scale, where `MAX_BUCKETS` makes +/// buckets larger, concurrency drops so memory stays near this figure. +const SHUFFLE_MEMORY_BUDGET: u64 = 4 * 1024 * 1024 * 1024; + +fn read_err(path: &Path, e: impl std::fmt::Display) -> FastqToolsError { + FastqToolsError::FastqReadError(format!("{}: {e}", path.display())) +} + +fn write_err(path: &Path, e: impl std::fmt::Display) -> FastqToolsError { + FastqToolsError::FastqWriteError(format!("{}: {e}", path.display())) +} + +/// Reads one 4-line FASTQ record into `buf` (appending). `Ok(false)` at a clean end of input; +/// an error if the input ends partway through a record. +fn read_record( + reader: &mut impl BufRead, + buf: &mut Vec, + path: &Path, +) -> Result { + let start = buf.len(); + for line in 0..4 { + let n = reader + .read_until(b'\n', buf) + .map_err(|e| read_err(path, e))?; + if n == 0 { + if line == 0 { + return Ok(false); + } + return Err(read_err( + path, + format!("input ends {line} lines into a FASTQ record"), + )); + } + } + if buf.last() != Some(&b'\n') { + buf.push(b'\n'); + } + if buf[start] != b'@' { + return Err(read_err(path, "a FASTQ record does not start with '@'")); + } + Ok(true) +} + +fn open_gz(path: &Path) -> Result>, FastqToolsError> { + let file = File::open(path).map_err(|e| read_err(path, e))?; + Ok(BufReader::with_capacity(1 << 20, MultiGzDecoder::new(file))) +} + +/// Pass-2 concurrency: the requested threads, capped at `MAX_SHUFFLE_THREADS` and at however +/// many of the largest bucket fit in `SHUFFLE_MEMORY_BUDGET`. Always at least one. +fn shuffle_threads(requested: usize, largest_bucket_bytes: u64) -> usize { + let fit = (SHUFFLE_MEMORY_BUDGET / largest_bucket_bytes.max(1)) as usize; + requested.min(MAX_SHUFFLE_THREADS).min(fit).max(1) +} + +/// Number of buckets for inputs totalling `compressed_bytes` on disk. +fn bucket_count(compressed_bytes: u64) -> usize { + let estimate = compressed_bytes.saturating_mul(EXPANSION_ESTIMATE); + (estimate.div_ceil(BUCKET_TARGET_BYTES) as usize).clamp(1, MAX_BUCKETS) +} + +/// Writes `files_r1` (and `files_r2`, record for record, when paired) to `out_r1` (and +/// `out_r2`) in a uniformly random order, keeping mates in register. `rng` is consumed for the +/// scatter; each bucket's permutation uses `rng.derive_child(bucket)`. Temp files go in +/// `work_dir`, and every input file is deleted once read. +pub fn combine_temp_fastqs_shuffled( + files_r1: Vec, + files_r2: Vec, + out_r1: &Path, + out_r2: Option<&Path>, + rng: &mut NeatRng, + work_dir: &Path, + threads: usize, +) -> Result<(), FastqToolsError> { + let compressed: u64 = files_r1 + .iter() + .chain(files_r2.iter()) + .map(|p| fs::metadata(p).map(|m| m.len()).unwrap_or(0)) + .sum(); + shuffle_into_buckets( + files_r1, + files_r2, + out_r1, + out_r2, + rng, + work_dir, + threads, + bucket_count(compressed), + ) +} + +/// The shuffle with an explicit bucket count, so tests can exercise the multi-bucket path on +/// inputs small enough that `bucket_count` would choose one. +#[allow(clippy::too_many_arguments)] +fn shuffle_into_buckets( + files_r1: Vec, + files_r2: Vec, + out_r1: &Path, + out_r2: Option<&Path>, + rng: &mut NeatRng, + work_dir: &Path, + threads: usize, + k: usize, +) -> Result<(), FastqToolsError> { + let paired = out_r2.is_some(); + if paired && files_r1.len() != files_r2.len() { + return Err(FastqToolsError::FastqWriteError(format!( + "shuffle: {} read-1 temp files but {} read-2", + files_r1.len(), + files_r2.len() + ))); + } + let bucket_dir = work_dir.join("shuffle_buckets"); + fs::create_dir_all(&bucket_dir).map_err(|e| write_err(&bucket_dir, e))?; + let bucket_path = |b: usize, mate: u8, stage: &str| { + bucket_dir.join(format!("{stage}_{b:04}_r{mate}.fastq.gz")) + }; + + // ── Pass 1: scatter ─────────────────────────────────────────────────────────────────────── + let mates: Vec = if paired { vec![1, 2] } else { vec![1] }; + let mut writers: Vec>>> = Vec::with_capacity(k); + for b in 0..k { + let mut per_mate = Vec::with_capacity(mates.len()); + for &m in &mates { + let p = bucket_path(b, m, "bucket"); + let f = File::create(&p).map_err(|e| write_err(&p, e))?; + per_mate.push(BlockGzWriter::with_level( + BufWriter::new(f), + CompressionLvl::fastest(), + )); + } + writers.push(per_mate); + } + let mut counts = vec![0u64; k]; + let mut bytes = vec![0u64; k]; + let (mut rec1, mut rec2) = (Vec::with_capacity(1024), Vec::with_capacity(1024)); + for (i, p1) in files_r1.iter().enumerate() { + let mut in1 = open_gz(p1)?; + let mut in2 = if paired { + Some((open_gz(&files_r2[i])?, &files_r2[i])) + } else { + None + }; + loop { + rec1.clear(); + let more1 = read_record(&mut in1, &mut rec1, p1)?; + let more2 = match in2.as_mut() { + Some((r, p)) => { + rec2.clear(); + read_record(r, &mut rec2, p)? + } + None => more1, + }; + if more1 != more2 { + return Err(read_err( + p1, + "read-1 and read-2 temp files hold different numbers of records", + )); + } + if !more1 { + break; + } + let b = ((rng.random()? * k as f64) as usize).min(k - 1); + writers[b][0] + .write_all(&rec1) + .map_err(|e| write_err(&bucket_path(b, 1, "bucket"), e))?; + if paired { + writers[b][1] + .write_all(&rec2) + .map_err(|e| write_err(&bucket_path(b, 2, "bucket"), e))?; + } + counts[b] += 1; + bytes[b] += (rec1.len() + if paired { rec2.len() } else { 0 }) as u64; + } + drop(in1); + drop(in2); + let _ = fs::remove_file(p1); + if paired { + let _ = fs::remove_file(&files_r2[i]); + } + } + for (b, per_mate) in writers.into_iter().enumerate() { + for (j, w) in per_mate.into_iter().enumerate() { + w.finish() + .map_err(|e| write_err(&bucket_path(b, mates[j], "bucket"), e))?; + } + } + + // ── Pass 2: shuffle each bucket ─────────────────────────────────────────────────────────── + let next = AtomicUsize::new(0); + let first_error: Mutex> = Mutex::new(None); + let base = *rng; + let shuffle_bucket = |b: usize| -> Result<(), FastqToolsError> { + // One contiguous buffer per mate, with record start offsets: a Vec per record would + // cost an allocation and its slack for every read. + let mut records: Vec<(Vec, Vec)> = Vec::with_capacity(mates.len()); + for &m in &mates { + let p = bucket_path(b, m, "bucket"); + let mut r = open_gz(&p)?; + let mut data = Vec::new(); + let mut starts = Vec::with_capacity(counts[b] as usize + 1); + loop { + starts.push(data.len()); + if !read_record(&mut r, &mut data, &p)? { + break; + } + } + if starts.len() as u64 - 1 != counts[b] { + return Err(read_err( + &p, + format!( + "bucket holds {} records, {} were written", + starts.len() - 1, + counts[b] + ), + )); + } + drop(r); + let _ = fs::remove_file(&p); + records.push((data, starts)); + } + let mut order: Vec = (0..counts[b] as usize).collect(); + base.derive_child(b as u64).shuffle_in_place(&mut order)?; + for (j, &m) in mates.iter().enumerate() { + let p = bucket_path(b, m, "shuffled"); + let f = File::create(&p).map_err(|e| write_err(&p, e))?; + let mut w = BlockGzWriter::new(BufWriter::new(f)); + let (data, starts) = &records[j]; + for &i in &order { + w.write_all(&data[starts[i]..starts[i + 1]]) + .map_err(|e| write_err(&p, e))?; + } + w.finish().map_err(|e| write_err(&p, e))?; + } + Ok(()) + }; + std::thread::scope(|s| { + for _ in 0..shuffle_threads(threads, bytes.iter().copied().max().unwrap_or(0)) { + s.spawn(|| { + loop { + let b = next.fetch_add(1, Ordering::SeqCst); + if b >= k || first_error.lock().unwrap().is_some() { + break; + } + if let Err(e) = shuffle_bucket(b) { + first_error.lock().unwrap().get_or_insert(e); + break; + } + } + }); + } + }); + if let Some(e) = first_error.into_inner().unwrap() { + return Err(e); + } + + // ── Concatenate, in bucket order ────────────────────────────────────────────────────────── + for (j, &m) in mates.iter().enumerate() { + let out = if j == 0 { + out_r1 + } else { + out_r2.expect("paired") + }; + let mut w = BufWriter::new( + fs::OpenOptions::new() + .append(true) + .create(true) + .open(out) + .map_err(|e| write_err(out, e))?, + ); + for b in 0..k { + let p = bucket_path(b, m, "shuffled"); + let mut f = BufReader::new(File::open(&p).map_err(|e| read_err(&p, e))?); + std::io::copy(&mut f, &mut w).map_err(|e| write_err(out, e))?; + drop(f); + let _ = fs::remove_file(&p); + } + w.flush().map_err(|e| write_err(out, e))?; + } + let _ = fs::remove_dir(&bucket_dir); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Read; + + /// Writes `n` records named by index (and their mates) as `chunks` temp-file pairs, in order. + fn temp_inputs( + dir: &Path, + n: usize, + chunks: usize, + paired: bool, + ) -> (Vec, Vec) { + let (mut f1, mut f2) = (Vec::new(), Vec::new()); + let per = n.div_ceil(chunks); + for c in 0..chunks { + let p1 = dir.join(format!("t{c}_r1.fastq.gz")); + let p2 = dir.join(format!("t{c}_r2.fastq.gz")); + let mut w1 = BlockGzWriter::new(File::create(&p1).unwrap()); + let mut w2 = BlockGzWriter::new(File::create(&p2).unwrap()); + for i in c * per..((c + 1) * per).min(n) { + write!(w1, "@read_{i:08}/1\nACGT{i}\n+\nIIII{i}\n").unwrap(); + write!(w2, "@read_{i:08}/2\nTGCA{i}\n+\nJJJJ{i}\n").unwrap(); + } + w1.finish().unwrap(); + w2.finish().unwrap(); + f1.push(p1); + if paired { + f2.push(p2); + } + } + (f1, f2) + } + + fn names(path: &Path) -> Vec { + let mut s = String::new(); + MultiGzDecoder::new(File::open(path).unwrap()) + .read_to_string(&mut s) + .unwrap(); + s.lines().step_by(4).map(str::to_string).collect() + } + + fn index_of(name: &str) -> usize { + name[6..14].parse().unwrap() + } + + /// Shuffles `n` records from `chunks` temp-file pairs through `k` buckets. Tests pass k > 1 + /// explicitly: at these sizes `bucket_count` would always choose 1, leaving the scatter untested. + fn run_k( + n: usize, + chunks: usize, + seed: &str, + threads: usize, + k: usize, + ) -> (Vec, Vec) { + let dir = tempfile::tempdir().unwrap(); + let (f1, f2) = temp_inputs(dir.path(), n, chunks, true); + let (o1, o2) = ( + dir.path().join("o1.fastq.gz"), + dir.path().join("o2.fastq.gz"), + ); + let mut rng = NeatRng::new_from_seed(&vec![seed.to_string()]).unwrap(); + shuffle_into_buckets(f1, f2, &o1, Some(&o2), &mut rng, dir.path(), threads, k).unwrap(); + (names(&o1), names(&o2)) + } + + fn run(n: usize, chunks: usize, seed: &str, threads: usize) -> (Vec, Vec) { + run_k(n, chunks, seed, threads, 16) + } + + #[test] + fn output_is_a_permutation_with_mates_in_register() { + let n = 20_000; + let (r1, r2) = run(n, 7, "perm", 4); + assert_eq!(r1.len(), n); + let mut idx: Vec = r1.iter().map(|s| index_of(s)).collect(); + for (a, b) in r1.iter().zip(&r2) { + assert_eq!( + a.trim_end_matches("/1"), + b.trim_end_matches("/2"), + "mates out of register" + ); + } + idx.sort_unstable(); + assert_eq!(idx, (0..n).collect::>(), "every record exactly once"); + } + + #[test] + fn output_order_is_random_not_the_input_order() { + // Adjacent records: in input order every neighbor is the next index. Shuffled, a neighbor + // is within 100 indexes with probability ~ 2 * 100 / n = 0.01. + let n = 20_000; + let (r1, _) = run(n, 7, "order", 4); + let idx: Vec = r1.iter().map(|s| index_of(s) as i64).collect(); + let near = idx + .windows(2) + .filter(|w| (w[1] - w[0]).abs() <= 100) + .count() as f64 + / (n - 1) as f64; + assert!( + near < 0.03, + "adjacent records within 100 indexes: {near:.4}; the output is not shuffled" + ); + assert_ne!( + idx[..100], + (0..100).collect::>()[..], + "starts in input order" + ); + } + + #[test] + fn same_seed_same_output_across_thread_counts_and_different_seed_differs() { + let n = 5_000; + let a = run(n, 3, "det", 1); + assert_eq!(a, run(n, 3, "det", 8), "thread count changed the output"); + assert_ne!( + a.0, + run(n, 3, "det2", 1).0, + "a different seed gave the same order" + ); + } + + #[test] + fn single_ended_and_empty_inputs() { + let dir = tempfile::tempdir().unwrap(); + let (f1, _) = temp_inputs(dir.path(), 1_000, 2, false); + let o1 = dir.path().join("o1.fastq.gz"); + let mut rng = NeatRng::new_from_seed(&vec!["se".to_string()]).unwrap(); + combine_temp_fastqs_shuffled(f1, vec![], &o1, None, &mut rng, dir.path(), 2).unwrap(); + assert_eq!(names(&o1).len(), 1_000); + + let (e1, e2) = temp_inputs(dir.path(), 0, 1, true); + let (o1, o2) = ( + dir.path().join("e1.fastq.gz"), + dir.path().join("e2.fastq.gz"), + ); + combine_temp_fastqs_shuffled(e1, e2, &o1, Some(&o2), &mut rng, dir.path(), 2).unwrap(); + assert!(names(&o1).is_empty() && names(&o2).is_empty()); + } + + #[test] + fn unequal_mate_files_are_an_error() { + let dir = tempfile::tempdir().unwrap(); + let (f1, f2) = temp_inputs(dir.path(), 100, 1, true); + // Drop the last read-2 record. + let mut s = String::new(); + MultiGzDecoder::new(File::open(&f2[0]).unwrap()) + .read_to_string(&mut s) + .unwrap(); + let kept: Vec<&str> = s.lines().collect(); + let mut w = BlockGzWriter::new(File::create(&f2[0]).unwrap()); + for l in &kept[..kept.len() - 4] { + writeln!(w, "{l}").unwrap(); + } + w.finish().unwrap(); + let (o1, o2) = ( + dir.path().join("o1.fastq.gz"), + dir.path().join("o2.fastq.gz"), + ); + let mut rng = NeatRng::new_from_seed(&vec!["x".to_string()]).unwrap(); + let err = combine_temp_fastqs_shuffled(f1, f2, &o1, Some(&o2), &mut rng, dir.path(), 1) + .unwrap_err(); + assert!( + err.to_string().contains("different numbers of records"), + "{err}" + ); + } + + #[test] + fn one_bucket_and_many_buckets_both_give_a_shuffled_permutation() { + // k = 1 is all in-bucket shuffle; k = 64 is mostly scatter. Both must permute and mix. + for k in [1, 64] { + let n = 10_000; + let (r1, r2) = run_k(n, 5, "k", 3, k); + let mut idx: Vec = r1.iter().map(|s| index_of(s) as i64).collect(); + let near = idx + .windows(2) + .filter(|w| (w[1] - w[0]).abs() <= 100) + .count() as f64 + / (n - 1) as f64; + assert!(near < 0.05, "k={k}: adjacent within 100: {near:.4}"); + assert!( + r1.iter() + .zip(&r2) + .all(|(a, b)| a.trim_end_matches("/1") == b.trim_end_matches("/2")), + "k={k}" + ); + idx.sort_unstable(); + assert_eq!(idx, (0..n as i64).collect::>(), "k={k}"); + } + } + + #[test] + fn shuffle_threads_fit_the_memory_budget() { + let gib = 1024 * 1024 * 1024; + assert_eq!( + shuffle_threads(64, 256 * 1024 * 1024), + MAX_SHUFFLE_THREADS, + "small buckets: full cap" + ); + assert_eq!(shuffle_threads(64, gib), 4, "1 GiB buckets: 4 fit in 4 GiB"); + assert_eq!( + shuffle_threads(64, 10 * gib), + 1, + "a bucket over budget still runs, alone" + ); + assert_eq!(shuffle_threads(2, 1), 2, "never more than requested"); + assert_eq!(shuffle_threads(0, 0), 1, "at least one"); + } + + #[test] + fn bucket_count_scales_with_input_and_is_capped() { + assert_eq!(bucket_count(0), 1); + assert_eq!(bucket_count(1_000), 1); + assert_eq!( + bucket_count(BUCKET_TARGET_BYTES), + EXPANSION_ESTIMATE as usize + ); + assert_eq!(bucket_count(u64::MAX), MAX_BUCKETS); + } +} diff --git a/eidolon/src/gen_reads/utils/config.rs b/eidolon/src/gen_reads/utils/config.rs index f68c8b8a..ad08b9e9 100644 --- a/eidolon/src/gen_reads/utils/config.rs +++ b/eidolon/src/gen_reads/utils/config.rs @@ -151,6 +151,10 @@ pub struct RunConfiguration { // set it on its own to get the adapter-free short-insert arm. Default off → // short fragments discarded, output byte-identical to before. pub keep_short_fragments: bool, + // Write the final FASTQs in random order rather than genome order (#870). Default on: a real + // sequencer's FASTQ is unordered, and genome order made bwa-mem2 need 166 GB instead of 90 GB + // on a 21x CHM13 simulation. Off restores the byte-for-byte concatenation of chunk outputs. + pub shuffle_fastq: bool, // Scale applied to the SvModel's per-base SV rate at generation time. // 0.0 (default) disables de novo SV generation entirely — only SVs // from `input_vcf` flow through. Values > 0 enable sampling from @@ -230,6 +234,7 @@ impl Default for RunConfiguration { chunk_size: None, long_reads: false, keep_short_fragments: false, + shuffle_fastq: true, sv_rate_scale: 0.0, sv_max_length_fraction: 0.25, adapters: AdapterConfig::default(), @@ -273,6 +278,7 @@ pub const KNOWN_KEYS: &[&str] = &[ "reference", "rng_seed", "sequence_error_model", + "shuffle_fastq", "sv_max_length_fraction", "sv_rate_scale", "target_bed", @@ -720,6 +726,14 @@ impl RunConfiguration { ) })?; } + "shuffle_fastq" => { + config.shuffle_fastq = value.as_bool().ok_or_else(|| { + GenerateReadsError::ConfigReadError( + "shuffle_fastq".to_string(), + "boolean".to_string(), + ) + })?; + } "sv_rate_scale" => { config.sv_rate_scale = value.as_f64().ok_or_else(|| { GenerateReadsError::ConfigReadError( @@ -813,6 +827,9 @@ impl RunConfiguration { info!("\t>mutation rate: {}", rate) } info!("\t>paired ended: {}", config.paired_ended); + if config.produce_fastq { + info!("\t>shuffle FASTQ order: {}", config.shuffle_fastq); + } if config.long_reads { info!("\t>long reads mode: enabled (short fragments produce truncated reads)"); } @@ -1043,6 +1060,7 @@ mod tests { chunk_size: None, long_reads: false, keep_short_fragments: false, + shuffle_fastq: true, sv_rate_scale: 0.0, sv_max_length_fraction: 0.25, adapters: AdapterConfig::default(), diff --git a/eidolon/src/gen_reads/utils/runner.rs b/eidolon/src/gen_reads/utils/runner.rs index 17b7e769..f160f5ee 100644 --- a/eidolon/src/gen_reads/utils/runner.rs +++ b/eidolon/src/gen_reads/utils/runner.rs @@ -17,6 +17,7 @@ use crate::{ bam_writer::{BamBodyWriter, BamContext, BamRecordStager, concat_temp_bams}, bed_reader::read_bed, fasta_stream::{FastaStream, map_buffer, resolve_iupac_bases}, + fastq_shuffle::combine_temp_fastqs_shuffled, fastq_tools::{ HaplotypeContext, PlacedFragment, Strand, combine_temp_fastqs, generate_read, generate_reverse_mate, r2_quality_model, reverse_complement, write_block_fastq, @@ -84,6 +85,43 @@ struct ContigContext<'a> { contig_order: Arc>, } +/// Writes the final FASTQs from the per-chunk temp files: shuffled by default (#870), or +/// concatenated in contig order when `shuffle_fastq: false`. The shuffle draws from its own child +/// of the run's RNG, so it consumes nothing from the streams that generate reads; the truth VCF +/// and golden BAM are the same either way. +fn combine_fastqs( + config: &RunConfiguration, + ctx: &ContigContext<'_>, + all_r1: Vec, + all_r2: Vec, + filename1: &PathBuf, + filename2: Option<&PathBuf>, +) -> Result<(), GenerateReadsError> { + if !config.shuffle_fastq { + combine_temp_fastqs(all_r1, all_r2, filename1, filename2)?; + return Ok(()); + } + info!("Shuffling FASTQ order (shuffle_fastq: true)"); + let mut rng = ctx.base_rng.derive_child(SHUFFLE_RNG_CHILD); + let threads = config + .num_threads + .unwrap_or_else(|| std::thread::available_parallelism().map_or(1, |n| n.get())); + combine_temp_fastqs_shuffled( + all_r1, + all_r2, + filename1, + filename2.map(PathBuf::as_path), + &mut rng, + ctx.working_dir, + threads, + )?; + Ok(()) +} + +/// Child index of the run RNG reserved for the FASTQ shuffle. Contigs use their index and the +/// chimeric pass uses 999999; this is far from both. +const SHUFFLE_RNG_CHILD: u64 = 0x5348_5546_464c_4521; + pub fn run_neat( config: &RunConfiguration, rng: &mut NeatRng, @@ -545,7 +583,14 @@ pub fn run_neat( match &config.output_fastq_2 { Some(filename2) => { create_output_file(filename2, config.overwrite_output)?; - combine_temp_fastqs(all_r1, all_r2, filename1, Some(filename2))?; + combine_fastqs( + config, + &ctx, + all_r1, + all_r2, + filename1, + Some(filename2), + )?; } None => { return Err(GenerateReadsError::ConfigError(format!( @@ -554,7 +599,7 @@ pub fn run_neat( } } } else { - combine_temp_fastqs(all_r1, vec![], filename1, None)?; + combine_fastqs(config, &ctx, all_r1, vec![], filename1, None)?; } } None => { diff --git a/eidolon/tests/common/mod.rs b/eidolon/tests/common/mod.rs index 7f5a6fb2..c65b9433 100644 --- a/eidolon/tests/common/mod.rs +++ b/eidolon/tests/common/mod.rs @@ -96,6 +96,8 @@ pub struct GenReadsConfig { pub mutation_rate: Option, pub mutation_model: Option, pub sv_rate_scale: Option, + /// `None` leaves gen-reads' default, which shuffles the FASTQs (#870). + pub shuffle_fastq: Option, } impl GenReadsConfig { @@ -121,6 +123,7 @@ impl GenReadsConfig { mutation_rate: None, mutation_model: None, sv_rate_scale: None, + shuffle_fastq: None, } } @@ -169,6 +172,10 @@ impl GenReadsConfig { Some(s) => format!("sv_rate_scale: {s}\n"), None => String::new(), }; + let shuffle_section = match self.shuffle_fastq { + Some(b) => format!("shuffle_fastq: {b}\n"), + None => String::new(), + }; write!( f, "reference: {ref_}\n\ @@ -181,7 +188,7 @@ impl GenReadsConfig { output_filename: {name}\n\ overwrite_output: true\n\ rng_seed: {seed}\n\ - {pair}{model}{qmodel}{threads}{chunksz}{ivcf}{mrate}{mmodel}{svscale}", + {pair}{model}{qmodel}{threads}{chunksz}{ivcf}{mrate}{mmodel}{svscale}{shuffle}", ref_ = self.reference.display(), rl = self.read_len, cov = self.coverage, @@ -198,6 +205,7 @@ impl GenReadsConfig { chunksz = chunk_size_section, ivcf = input_vcf_section, mrate = mutation_rate_section, + shuffle = shuffle_section, mmodel = mutation_model_section, svscale = sv_rate_scale_section, ) diff --git a/eidolon/tests/determinism.rs b/eidolon/tests/determinism.rs index 2428545b..db0b0590 100644 --- a/eidolon/tests/determinism.rs +++ b/eidolon/tests/determinism.rs @@ -41,11 +41,23 @@ fn run_gen_reads_cfg( name: &str, threads: Option, sv_rate_scale: Option, +) { + run_gen_reads_full(seed, output_dir, name, threads, sv_rate_scale, None) +} + +fn run_gen_reads_full( + seed: &str, + output_dir: &std::path::Path, + name: &str, + threads: Option, + sv_rate_scale: Option, + shuffle_fastq: Option, ) { let mut config = GenReadsConfig::new(h1n1_reference(), output_dir.to_path_buf(), name); config.rng_seed = seed.to_string(); config.num_threads = threads; config.sv_rate_scale = sv_rate_scale; + config.shuffle_fastq = shuffle_fastq; let yaml = config.write_yaml(); eidolon() .args(["gen-reads", "-c"]) @@ -213,10 +225,21 @@ fn byte_identical_output_with_svs_across_thread_counts() { /// them and that section is legitimately not sorted by contig. Only the first section is /// asserted here; chimeric block ordering is left unpinned rather than pinned to something /// this test would have to guess at. +/// +/// Since #870 the default output is shuffled, so this pins the `shuffle_fastq: false` path. +/// That order still matters with the shuffle on: it is the input order the shuffle's +/// determinism rests on. #[test] fn ordinary_reads_are_emitted_in_reference_file_order() { let (_g, dir) = fresh_workdir(); - run_gen_reads_cfg("contig order seed", &dir, "order", Some(1), Some(50.0)); + run_gen_reads_full( + "contig order seed", + &dir, + "order", + Some(1), + Some(50.0), + Some(false), + ); let lines = read_gzip_fastq_lines(&dir.join("order_r1.fastq.gz")); let reference_order = [ diff --git a/eidolon/tests/fastq_shuffle_output.rs b/eidolon/tests/fastq_shuffle_output.rs new file mode 100644 index 00000000..8f0b4e48 --- /dev/null +++ b/eidolon/tests/fastq_shuffle_output.rs @@ -0,0 +1,155 @@ +//! gen-reads writes its FASTQs in random order by default (#870). +//! +//! Genome order made bwa-mem2 need 166 GB instead of 90 GB on a 21x T2T-CHM13 simulation, with +//! identical alignments either way. A real sequencer's FASTQ is unordered, so that is the default. +//! +//! **Known answer.** Read names carry each read's origin (`…___/1`). In genome +//! order nearly every read's neighbor starts within a few hundred bases of it; in random order on +//! H1N1 (13.5 kb), a neighbor on the same contig within 200 bp is rare. **Must not fire:** +//! `shuffle_fastq: false` keeps genome order, and the shuffle changes nothing but order. The same +//! records come out, mates stay in register, and the truth VCF and golden BAM are byte-identical, +//! because the shuffle draws from its own RNG child. + +mod common; + +use common::{eidolon, h1n1_reference}; +use noodles::bam; +use std::io::Read; +use std::path::Path; + +fn gz_text(path: &Path) -> String { + let mut s = String::new(); + flate2::read::MultiGzDecoder::new(std::fs::File::open(path).unwrap()) + .read_to_string(&mut s) + .unwrap(); + s +} + +fn run(dir: &Path, name: &str, shuffle: Option, threads: usize) { + let cfg = dir.join(format!("{name}.yml")); + let shuffle_line = shuffle + .map(|s| format!("shuffle_fastq: {s}\n")) + .unwrap_or_default(); + std::fs::write( + &cfg, + format!( + "reference: {ref}\nread_len: 101\ncoverage: 30\npaired_ended: true\n\ + fragment_mean: 300\nfragment_st_dev: 30\n\ + produce_fastq: true\nproduce_vcf: true\nproduce_bam: true\n\ + overwrite_output: true\noutput_dir: {out}\noutput_filename: {name}\n\ + rng_seed: fastq shuffle\nnum_threads: {threads}\n{shuffle_line}", + ref = h1n1_reference().display(), + out = dir.display(), + ), + ) + .unwrap(); + eidolon() + .args(["gen-reads", "-c"]) + .arg(&cfg) + .assert() + .success(); +} + +/// (contig, start) parsed from a read name, counting fields from the right because contig names +/// may contain underscores. +fn origin(name: &str) -> (String, i64) { + let core = name + .trim_start_matches('@') + .trim_start_matches("EIDOLON_generated_"); + let core = core.rsplit_once('/').map_or(core, |(c, _)| c); + let parts: Vec<&str> = core.rsplitn(4, '_').collect(); + (parts[3].to_string(), parts[2].parse().unwrap()) +} + +fn names(fastq: &str) -> Vec { + fastq.lines().step_by(4).map(str::to_string).collect() +} + +/// Fraction of adjacent reads on the same contig with starts within 200 bp. +fn near_neighbors(names: &[String]) -> f64 { + let o: Vec<(String, i64)> = names.iter().map(|n| origin(n)).collect(); + let near = o + .windows(2) + .filter(|w| w[0].0 == w[1].0 && (w[0].1 - w[1].1).abs() <= 200) + .count(); + near as f64 / (o.len() - 1) as f64 +} + +#[test] +fn default_output_is_shuffled_and_off_keeps_genome_order() { + let tmp = tempfile::tempdir().unwrap(); + let d = tmp.path(); + run(d, "on", None, 2); + run(d, "off", Some(false), 2); + + let on1 = names(&gz_text(&d.join("on_r1.fastq.gz"))); + let on2 = names(&gz_text(&d.join("on_r2.fastq.gz"))); + let off1 = names(&gz_text(&d.join("off_r1.fastq.gz"))); + // Random order on H1N1 puts a neighbor on the same contig ~15% of the time and within 200 bp + // of it ~23% of that (contigs ~1.7 kb): ~0.035 expected, sigma ~0.006 at 1000 reads. So 1000 + // reads keep the < 0.1 assertion below far from noise. + assert!( + on1.len() > 1000, + "only {} reads; the order statistics would be weak", + on1.len() + ); + + let (near_on, near_off) = (near_neighbors(&on1), near_neighbors(&off1)); + eprintln!("adjacent reads within 200 bp: shuffled {near_on:.4}, genome order {near_off:.4}"); + assert!( + near_off > 0.9, + "shuffle_fastq: false should keep genome order: {near_off:.4}" + ); + assert!( + near_on < 0.1, + "the default output is not shuffled: {near_on:.4}" + ); + + // Same records, mates in register. + for (a, b) in on1.iter().zip(&on2) { + assert_eq!( + a.trim_end_matches("/1"), + b.trim_end_matches("/2"), + "mates out of register" + ); + } + let (mut s_on, mut s_off) = (on1.clone(), off1.clone()); + s_on.sort(); + s_off.sort(); + assert_eq!( + s_on, s_off, + "the shuffle must change order only, not which reads exist" + ); + + // The shuffle draws from its own RNG child: the truth VCF and the golden BAM are unaffected. + assert_eq!( + gz_text(&d.join("on.vcf.gz")), + gz_text(&d.join("off.vcf.gz")), + "truth VCF changed" + ); + let bam_records = |p: &Path| -> Vec { + let mut r = bam::io::Reader::new(std::fs::File::open(p).unwrap()); + r.read_header().unwrap(); + r.records().map(|x| format!("{:?}", x.unwrap())).collect() + }; + assert_eq!( + bam_records(&d.join("on.bam")), + bam_records(&d.join("off.bam")), + "golden BAM changed" + ); +} + +#[test] +fn shuffled_output_is_identical_across_thread_counts() { + let tmp = tempfile::tempdir().unwrap(); + let d = tmp.path(); + run(d, "t1", None, 1); + run(d, "t4", None, 4); + for mate in ["r1", "r2"] { + assert_eq!( + gz_text(&d.join(format!("t1_{mate}.fastq.gz"))), + gz_text(&d.join(format!("t4_{mate}.fastq.gz"))), + "{mate} differs between 1 and 4 threads" + ); + } +} diff --git a/template_config/gen_reads_template.yml b/template_config/gen_reads_template.yml index e6225675..7aaba07e 100644 --- a/template_config/gen_reads_template.yml +++ b/template_config/gen_reads_template.yml @@ -55,6 +55,12 @@ fragment_st_dev: . # At least one of produce_fastq, produce_vcf, or produce_bam must be true. produce_fastq: true +# Write the FASTQ(s) in random order, as a sequencer does, rather than genome order. Mates stay +# in register. Only record order changes: the truth VCF and BAM are identical either way, and the +# same seed gives the same output at any thread count. Genome order made bwa-mem2 need ~1.8x the +# memory on a T2T-CHM13 simulation (#870). (default: true) +shuffle_fastq: true + # Produce a gzipped VCF file recording all inserted variants. (default: false) produce_vcf: false From 7b4e54c26fbb231a4bd4671a60191f2991a57597 Mon Sep 17 00:00:00 2001 From: Joshua Allen Date: Thu, 5 Nov 2026 08:38:22 -0600 Subject: [PATCH 2/2] shuffle: bound scatter file descriptors; compare gzip bytes across threads (#870) Review of #875. The scatter held one writer per bucket per mate, 2 * 256 descriptors at the cap, which fails under common 256/512 limits. Buckets now accumulate in memory and leave as gzip members appended to their file, which is closed again, so at most one bucket file is open at any bucket count. MAX_BUCKETS rises to 512 (the cap is now buffer memory: 256 MiB). A test re-runs itself at 512 paired buckets under ulimit -n 48; on the old code it failed with "Too many open files" at bucket 22. Thread-count determinism is now asserted on the raw gzip bytes as well as the decompressed records, in the unit and integration tests; an empty-member-on-multithread mutant passes the record check and fails the byte check. BlockGzWriter::with_level is dropped as unused. Co-Authored-By: Claude --- eidolon-core/src/file_tools/block_gz.rs | 9 +- eidolon-core/src/file_tools/fastq_shuffle.rs | 181 ++++++++++++++++--- eidolon/tests/fastq_shuffle_output.rs | 16 +- 3 files changed, 167 insertions(+), 39 deletions(-) diff --git a/eidolon-core/src/file_tools/block_gz.rs b/eidolon-core/src/file_tools/block_gz.rs index 09ba83dd..09a89eb3 100644 --- a/eidolon-core/src/file_tools/block_gz.rs +++ b/eidolon-core/src/file_tools/block_gz.rs @@ -70,18 +70,11 @@ impl BlockGzWriter { // CompressionLvl::default() is libdeflate level 6 — the same ratio // target as flate2's Compression::default(), so output size is // comparable and only the speed changes. - Self::with_level(inner, CompressionLvl::default()) - } - - /// As [`BlockGzWriter::new`] at a chosen level. Short-lived intermediates (the - /// FASTQ shuffle's bucket files) use the fastest level: they are read once and - /// deleted, so ratio buys nothing. - pub fn with_level(inner: W, level: CompressionLvl) -> Self { Self { inner, buf: Vec::with_capacity(BLOCK_SIZE), out: Vec::new(), - compressor: Compressor::new(level), + compressor: Compressor::new(CompressionLvl::default()), finished: false, wrote_any: false, } diff --git a/eidolon-core/src/file_tools/fastq_shuffle.rs b/eidolon-core/src/file_tools/fastq_shuffle.rs index 77cad0d3..709471f7 100644 --- a/eidolon-core/src/file_tools/fastq_shuffle.rs +++ b/eidolon-core/src/file_tools/fastq_shuffle.rs @@ -27,7 +27,7 @@ use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; use flate2::read::MultiGzDecoder; -use libdeflater::CompressionLvl; +use libdeflater::{CompressionLvl, Compressor}; use crate::file_tools::block_gz::BlockGzWriter; use crate::file_tools::fastq_tools::FastqToolsError; @@ -36,9 +36,14 @@ use crate::rng::NeatRng; /// Uncompressed bytes per bucket, both mates together. Pass 2 holds this much per thread, /// roughly doubled by per-record allocation. const BUCKET_TARGET_BYTES: u64 = 256 * 1024 * 1024; -/// Upper bound on buckets, so the scatter's open files (one per bucket per mate) stay well under -/// the common 1024 soft limit. Past this, buckets grow instead. -const MAX_BUCKETS: usize = 256; +/// Upper bound on buckets. The scatter holds one buffer per bucket per mate, so this bounds that +/// memory (512 * 2 * `BUCKET_BUFFER_BYTES` = 256 MiB). Past it, buckets grow instead. +const MAX_BUCKETS: usize = 512; +/// Scatter buffer per bucket per mate. A full buffer is compressed as one gzip member and +/// appended to its bucket file, which is then closed again, so the scatter holds at most one +/// bucket file open at any bucket count. One writer per bucket per mate would need +/// 2 * MAX_BUCKETS descriptors and failed under `ulimit -n` 256 or 512. +const BUCKET_BUFFER_BYTES: usize = 256 * 1024; /// Gzipped FASTQ expands about 3.5-4.5x; overestimating only adds buckets. const EXPANSION_ESTIMATE: u64 = 5; /// Pass-2 threads, capped because each holds a whole bucket in memory. @@ -92,6 +97,68 @@ fn open_gz(path: &Path) -> Result>, FastqToolsErr Ok(BufReader::with_capacity(1 << 20, MultiGzDecoder::new(file))) } +/// One bucket file being scattered into: records accumulate in `buf` and leave as a gzip member +/// appended to `path`, so the file is open only while a member is written. +struct BucketSink { + path: PathBuf, + buf: Vec, + wrote: bool, +} + +impl BucketSink { + fn new(path: PathBuf) -> Self { + Self { + path, + buf: Vec::new(), + wrote: false, + } + } + + fn push( + &mut self, + rec: &[u8], + c: &mut Compressor, + scratch: &mut Vec, + ) -> Result<(), FastqToolsError> { + self.buf.extend_from_slice(rec); + if self.buf.len() >= BUCKET_BUFFER_BYTES { + self.flush(c, scratch)?; + } + Ok(()) + } + + /// Appends the buffer as one gzip member. An empty member is written only for a bucket that + /// received nothing, so every bucket file is valid gzip. + fn flush(&mut self, c: &mut Compressor, scratch: &mut Vec) -> Result<(), FastqToolsError> { + if self.buf.is_empty() && self.wrote { + return Ok(()); + } + let bound = c.gzip_compress_bound(self.buf.len()); + if scratch.len() < bound { + scratch.resize(bound, 0); + } + let n = c + .gzip_compress(&self.buf, scratch) + .map_err(|e| write_err(&self.path, format!("gzip compression failed: {e:?}")))?; + let mut f = fs::OpenOptions::new() + .append(true) + .create(true) + .open(&self.path) + .map_err(|e| write_err(&self.path, e))?; + f.write_all(&scratch[..n]) + .map_err(|e| write_err(&self.path, e))?; + self.buf.clear(); + self.wrote = true; + Ok(()) + } + + fn finish(&mut self, c: &mut Compressor, scratch: &mut Vec) -> Result<(), FastqToolsError> { + self.flush(c, scratch)?; + self.buf = Vec::new(); + Ok(()) + } +} + /// Pass-2 concurrency: the requested threads, capped at `MAX_SHUFFLE_THREADS` and at however /// many of the largest bucket fit in `SHUFFLE_MEMORY_BUDGET`. Always at least one. fn shuffle_threads(requested: usize, largest_bucket_bytes: u64) -> usize { @@ -164,19 +231,16 @@ fn shuffle_into_buckets( // ── Pass 1: scatter ─────────────────────────────────────────────────────────────────────── let mates: Vec = if paired { vec![1, 2] } else { vec![1] }; - let mut writers: Vec>>> = Vec::with_capacity(k); - for b in 0..k { - let mut per_mate = Vec::with_capacity(mates.len()); - for &m in &mates { - let p = bucket_path(b, m, "bucket"); - let f = File::create(&p).map_err(|e| write_err(&p, e))?; - per_mate.push(BlockGzWriter::with_level( - BufWriter::new(f), - CompressionLvl::fastest(), - )); - } - writers.push(per_mate); - } + let mut sinks: Vec> = (0..k) + .map(|b| { + mates + .iter() + .map(|&m| BucketSink::new(bucket_path(b, m, "bucket"))) + .collect() + }) + .collect(); + let mut compressor = Compressor::new(CompressionLvl::fastest()); + let mut scratch = Vec::new(); let mut counts = vec![0u64; k]; let mut bytes = vec![0u64; k]; let (mut rec1, mut rec2) = (Vec::with_capacity(1024), Vec::with_capacity(1024)); @@ -207,13 +271,9 @@ fn shuffle_into_buckets( break; } let b = ((rng.random()? * k as f64) as usize).min(k - 1); - writers[b][0] - .write_all(&rec1) - .map_err(|e| write_err(&bucket_path(b, 1, "bucket"), e))?; + sinks[b][0].push(&rec1, &mut compressor, &mut scratch)?; if paired { - writers[b][1] - .write_all(&rec2) - .map_err(|e| write_err(&bucket_path(b, 2, "bucket"), e))?; + sinks[b][1].push(&rec2, &mut compressor, &mut scratch)?; } counts[b] += 1; bytes[b] += (rec1.len() + if paired { rec2.len() } else { 0 }) as u64; @@ -225,12 +285,10 @@ fn shuffle_into_buckets( let _ = fs::remove_file(&files_r2[i]); } } - for (b, per_mate) in writers.into_iter().enumerate() { - for (j, w) in per_mate.into_iter().enumerate() { - w.finish() - .map_err(|e| write_err(&bucket_path(b, mates[j], "bucket"), e))?; - } + for sink in sinks.iter_mut().flatten() { + sink.finish(&mut compressor, &mut scratch)?; } + drop(sinks); // ── Pass 2: shuffle each bucket ─────────────────────────────────────────────────────────── let next = AtomicUsize::new(0); @@ -392,6 +450,19 @@ mod tests { (names(&o1), names(&o2)) } + /// The output files' raw bytes, for comparing gzip streams rather than their contents. + fn raw_k(n: usize, chunks: usize, seed: &str, threads: usize, k: usize) -> (Vec, Vec) { + let dir = tempfile::tempdir().unwrap(); + let (f1, f2) = temp_inputs(dir.path(), n, chunks, true); + let (o1, o2) = ( + dir.path().join("o1.fastq.gz"), + dir.path().join("o2.fastq.gz"), + ); + let mut rng = NeatRng::new_from_seed(&vec![seed.to_string()]).unwrap(); + shuffle_into_buckets(f1, f2, &o1, Some(&o2), &mut rng, dir.path(), threads, k).unwrap(); + (fs::read(&o1).unwrap(), fs::read(&o2).unwrap()) + } + fn run(n: usize, chunks: usize, seed: &str, threads: usize) -> (Vec, Vec) { run_k(n, chunks, seed, threads, 16) } @@ -441,6 +512,12 @@ mod tests { let n = 5_000; let a = run(n, 3, "det", 1); assert_eq!(a, run(n, 3, "det", 8), "thread count changed the output"); + // The gzip bytes too: member boundaries or headers that varied with scheduling would + // pass the record comparison above. + assert!( + raw_k(n, 3, "det", 1, 16) == raw_k(n, 3, "det", 8, 16), + "thread count changed the gzip bytes" + ); assert_ne!( a.0, run(n, 3, "det2", 1).0, @@ -536,6 +613,54 @@ mod tests { assert_eq!(shuffle_threads(0, 0), 1, "at least one"); } + const FD_CHILD_ENV: &str = "EIDOLON_SHUFFLE_FD_CHILD"; + /// A file-descriptor limit well under what one writer per bucket per mate would need at the + /// cap (2 * MAX_BUCKETS), and well over what the shuffle needs. + const FD_LIMIT: usize = 48; + + /// Runs only as the child of `cap_path_runs_under_a_low_file_descriptor_limit`, which sets + /// the limit; a no-op otherwise. + #[test] + fn cap_path_child() { + if std::env::var(FD_CHILD_ENV).is_err() { + return; + } + let n = 3_000; + let (r1, r2) = run_k(n, 4, "fd", MAX_SHUFFLE_THREADS, MAX_BUCKETS); + assert_eq!(r1.len(), n); + assert!( + r1.iter() + .zip(&r2) + .all(|(a, b)| a.trim_end_matches("/1") == b.trim_end_matches("/2")) + ); + println!("CAP-PATH-OK {n} pairs through {MAX_BUCKETS} buckets"); + } + + #[cfg(target_os = "linux")] + #[test] + fn cap_path_runs_under_a_low_file_descriptor_limit() { + // Paired, at MAX_BUCKETS, under `ulimit -n 48`. The child prints a marker so a filter + // that matched nothing cannot pass for a run that succeeded. + let exe = std::env::current_exe().unwrap(); + let out = std::process::Command::new("bash") + .arg("-c") + .arg(format!( + "ulimit -n {FD_LIMIT} && exec \"$0\" --exact file_tools::fastq_shuffle::tests::cap_path_child --nocapture --test-threads=1" + )) + .arg(&exe) + .env(FD_CHILD_ENV, "1") + .output() + .unwrap(); + let (stdout, stderr) = ( + String::from_utf8_lossy(&out.stdout), + String::from_utf8_lossy(&out.stderr), + ); + assert!( + out.status.success() && stdout.contains("CAP-PATH-OK"), + "shuffle at {MAX_BUCKETS} buckets failed under ulimit -n {FD_LIMIT}:\n{stdout}\n{stderr}" + ); + } + #[test] fn bucket_count_scales_with_input_and_is_capped() { assert_eq!(bucket_count(0), 1); diff --git a/eidolon/tests/fastq_shuffle_output.rs b/eidolon/tests/fastq_shuffle_output.rs index 8f0b4e48..a5f47e73 100644 --- a/eidolon/tests/fastq_shuffle_output.rs +++ b/eidolon/tests/fastq_shuffle_output.rs @@ -146,10 +146,20 @@ fn shuffled_output_is_identical_across_thread_counts() { run(d, "t1", None, 1); run(d, "t4", None, 4); for mate in ["r1", "r2"] { + let (a, b) = ( + d.join(format!("t1_{mate}.fastq.gz")), + d.join(format!("t4_{mate}.fastq.gz")), + ); assert_eq!( - gz_text(&d.join(format!("t1_{mate}.fastq.gz"))), - gz_text(&d.join(format!("t4_{mate}.fastq.gz"))), - "{mate} differs between 1 and 4 threads" + gz_text(&a), + gz_text(&b), + "{mate} records differ between 1 and 4 threads" + ); + // The files themselves, not just their decompressed text: gzip member boundaries or + // headers that varied with thread count would pass the line above. + assert!( + std::fs::read(&a).unwrap() == std::fs::read(&b).unwrap(), + "{mate} gzip bytes differ between 1 and 4 threads" ); } }