diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml index 637c0937e..e2c92a755 100644 --- a/.github/workflows/check.yml +++ b/.github/workflows/check.yml @@ -61,6 +61,76 @@ jobs: - name: Clippy check run: cargo ci-lint + # Builds and tests the optional in-process bwa-mem3 aligner backend (feature + # `aligner-bwa-mem3`, off by default). This is the one feature that pulls a + # C++ toolchain into the build (the vendored bwa-mem3 kernels via + # `bwa-mem3-sys`'s `cc`-driven build.rs), so it gets its own job rather than + # folding into `test`/`lint`: those two stay C-toolchain-free, matching every + # other job's fast/portable assumption. + # + # Matrix covers both architectures the feature compiles on (the dependency is + # gated in Cargo.toml to `x86_64`/`aarch64` — see the `bwa-mem3-rs` entry). + # `ubuntu-24.04-arm` is a real (not emulated) arm64 GitHub-hosted runner, so + # this exercises the NEON build path, not just the x86_64 SIMD-tier ladder. + # + # Toolchain deps are deliberately minimal: `build-essential` (a C++17-capable + # g++) and `zlib1g-dev` (the crate links `z`) are everything `bwa-mem3-sys`'s + # build.rs actually needs for a normal build — mirrors fg-labs/bwa-mem3-rs's + # own `build`/`lint` jobs. Two toolchain pieces this job deliberately does + # NOT install: + # - `libclang-dev`: bindgen only runs under the sys crate's own + # `regenerate-bindings` feature (gated in its build.rs), which this job + # never enables — the committed bindings are used as-is. + # - clang-19 / libomp: bwa-mem3's Makefile compiler floor (clang >= 19 / + # gcc >= 15) applies to building the vendored bwa-mem3 CLI *from + # upstream* (fg-labs/bwa-mem3-rs's own `e2e` job does that, to diff + # against this crate's output byte-for-byte); `cc` enforces no such floor + # on the sys crate itself, and this job never builds the reference CLI + # binary. + aligner-ffi: + strategy: + fail-fast: false + matrix: + os: [ubuntu-24.04, ubuntu-24.04-arm] + runs-on: ${{ matrix.os }} + timeout-minutes: 30 + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Install build deps (C++17 toolchain for bwa-mem3-sys) + run: sudo apt-get update && sudo apt-get install -y build-essential zlib1g-dev + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + with: + components: clippy + - name: Set up compilation cache (sccache) + uses: mozilla-actions/sccache-action@fc920bf0ec8de6ee65d409111f7ec508035751ba # v0.0.11 + - name: Install nextest + uses: taiki-e/install-action@b20dedce73af6905cdc30d6611090c9b67557c8d # v2.85.12 + with: + tool: nextest + - name: Unit tests (aligner-bwa-mem3) + run: cargo nextest run --features aligner-bwa-mem3 -p fgumi --locked + - name: Clippy (aligner-bwa-mem3, pedantic) + # Lints compare,simulate,profile-adjacency,aligner-bwa-mem3 together: + # the first three are already part of every default dev/test build's + # assumed surface (`cargo ci-lint` compiles them), so linting them + # together with the aligner feature catches any interaction between the + # two rather than linting the aligner feature in isolation. + run: | + cargo clippy --workspace --all-targets \ + --features compare,simulate,profile-adjacency,aligner-bwa-mem3 --locked \ + -- -D warnings -W clippy::pedantic + - name: Rustdoc (aligner-bwa-mem3, -D warnings) + # `cargo ci-doc` (the `docs` job) does not enable `aligner-bwa-mem3`, so + # the feature-gated in-process backend's docs are rendered only here. + # Clippy does not report broken intra-doc links; rustdoc does. + env: + RUSTDOCFLAGS: "-D warnings" + run: cargo doc --no-deps -p fgumi --document-private-items --features aligner-bwa-mem3 --locked + # `merge_slots.rs` swaps `std::sync` for `loom::sync` under `--cfg loom`, and # `tests/loom_merge_slots.rs` is `#![cfg(loom)]`. Neither is built by any other # job, so without this one a broken model — or a `cfg(loom)` build that stopped diff --git a/CLAUDE.md b/CLAUDE.md index ddcab8e99..c8d755b97 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -39,6 +39,7 @@ cargo bench - `compare` - Enable compare subcommand (developer tools) - `simulate` - Enable simulate command for test data generation - `profile-adjacency` - Enable profiling output for adjacency UMI assigner +- `aligner-bwa-mem3` - Opt-in in-process bwa-mem3 aligner backend for `runall` (cohort math, engine abstraction; the backend itself is still being wired in); off by default, and compiling it requires a C++17 toolchain (see `Cargo.toml`'s `bwa-mem3-rs` dependency) Build with features: `cargo build --release --features compare,simulate` @@ -133,15 +134,35 @@ Uses `mimalloc` as global allocator for performance. `#[allow(unsafe_code)]` blocks are permitted only at the documented sites listed below; any new `unsafe` block requires updating this section with a written justification. +Enabling the optional `aligner-bwa-mem3` feature links `bwa-mem3-sys`/`bwa-mem3-rs` +(vendored C++ plus its Rust FFI bindings), which contain their own `unsafe` outside +this workspace's `deny(unsafe_code)`; fgumi's own crates carry no new `unsafe` for it +and remain `#![deny(unsafe_code)]` regardless of the feature. + ### Approved non-stdlib FFI exceptions The following external FFI calls are approved because they back core infrastructure: - **`libmimalloc_sys`** (`crates/fgumi-sort/src/memory_probe.rs`) — mimalloc is the configured - global allocator. The FFI wrappers (`force_mi_collect`, `process_rss_bytes`, and the - `memory-debug`-gated `print_mi_stats` calling `mi_stats_print_out`) are isolated to - a single `#[allow(unsafe_code)]` sub-module. mimalloc synchronizes these calls - internally, so they are safe to invoke concurrently. + global allocator. The FFI wrappers (`force_mi_collect`, `process_rss_bytes`, the + `memory-debug`-gated `print_mi_stats` calling `mi_stats_print_out`, and + `retain_freed_memory` / `mi_purge_delay_ms` calling `mi_option_set` / + `mi_option_get`) are isolated to a single `#[allow(unsafe_code)]` sub-module. + mimalloc synchronizes these calls internally, so they are safe to invoke + concurrently. `retain_freed_memory` sets mimalloc's purge delay to -1 (never + return freed pages to the OS); the align stage calls it at wiring (via + `retain_freed_memory_unless_user_set`) unless `MIMALLOC_PURGE_DELAY` is set, + because its per-batch buffer churn otherwise costs hundreds of thousands of page + faults. The setting is process-wide and never restored, so it also turns + fgumi-sort's `force_mi_collect()` into a no-op for later `runall` stages; measured + end to end (extract through consensus, 1M pairs, 32 threads) it is still ~4% + faster, sort spills included, for ~0.4 GB more peak RSS. + The subprocess route also passes `MIMALLOC_PURGE_DELAY=-1` to the aligner child + through its environment, which needs no FFI. A safe alternative does not exist: + mimalloc reads its environment only at process start, and `libmimalloc-sys` + exports no safe setter. + The option index (15) is not exported as a constant; a unit test pins it against + the linked mimalloc v3's 1000 ms default. - **`mach2`** (`crates/fgumi-sort/src/memory_probe.rs`, macOS only) — `task_info(TASK_VM_INFO)` is the only way to read `phys_footprint` (the RSS metric mimalloc reports accurately). Isolated to the same sub-module as the mimalloc FFI. diff --git a/Cargo.lock b/Cargo.lock index f2097c8d7..b4b20ddc8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -281,6 +281,27 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bwa-mem3-rs" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d605c934e6a7dbd00be78d3645335dbe664662f5234c3db3681db6b45167bce" +dependencies = [ + "bwa-mem3-sys", + "libc", + "thiserror 1.0.69", +] + +[[package]] +name = "bwa-mem3-sys" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "998de3690d2feda385b024c92a8f5f63ecce65fc2cb90f37cab87e5a5dea85ee" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "bytemuck" version = "1.25.2" @@ -712,6 +733,7 @@ dependencies = [ "approx", "bstr 1.13.1", "built", + "bwa-mem3-rs", "bytemuck", "bytes", "bytesize", diff --git a/Cargo.toml b/Cargo.toml index 5c34e3d36..0d5626049 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -181,6 +181,14 @@ mach2 = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] nix = { workspace = true } +# In-process bwa-mem3 aligner backend (feature `aligner-bwa-mem3`). Gated on the +# only two architectures whose vendored C++ SIMD tiers build (`x86_64`/`aarch64`); +# enabling the feature elsewhere hits a `compile_error!` in +# `pipeline::steps::align`. 0.3.0 is the first +# release with the three-phase (`seed_extend` / `infer_cohort` / `pair_emit`) API. +[target.'cfg(any(target_arch = "x86_64", target_arch = "aarch64"))'.dependencies] +bwa-mem3-rs = { version = "0.3.0", optional = true } + [features] default = ["consensus"] # Single umbrella feature for all consensus calling (simplex + duplex + codec). @@ -210,6 +218,19 @@ simulate = ["consensus"] profile-adjacency = [] # Enable slow stress tests for concurrency testing stress-tests = [] +# In-process bwa-mem3 aligner backend for `fgumi runall` (its steps and the +# `bwa-mem3-inproc` preset build on this). Off by default: it pulls the +# optional `bwa-mem3-rs` dependency, which compiles the vendored bwa-mem3 C++ (C++17 toolchain +# required) and links it statically. Default builds stay pure Rust with no C +# toolchain. The in-process backend code is `#[cfg(feature = "aligner-bwa-mem3")]` +# confined to `pipeline::steps::align::inproc`. +# +# `mimalloc/override` makes the global mimalloc also interpose the C `malloc`/ +# `free`, so the linked bwa-mem3 C++ uses mimalloc instead of the system +# allocator (whose per-thread arena contention regresses multithreaded +# alignment). Scoped to this feature because it is the only build that links the +# C++ aligner; a pure-Rust fgumi build gains nothing from the C interpose. +aligner-bwa-mem3 = ["dep:bwa-mem3-rs", "mimalloc/override"] # Test-only feature (not part of the supported public API). Its former gated # helper (the legacy multi-thread engine's `run_bam_pipeline_with_grouper`) was # removed in R6/C6; retained as the self dev-dependency's feature hook. Enabled diff --git a/crates/fgumi-sort/src/lib.rs b/crates/fgumi-sort/src/lib.rs index 7231606bd..d9f8f20a2 100644 --- a/crates/fgumi-sort/src/lib.rs +++ b/crates/fgumi-sort/src/lib.rs @@ -101,6 +101,10 @@ pub use memory_probe::print_mi_stats; /// pipeline thread-telemetry sampler, which needs an RSS reading on every tick /// regardless of whether that feature is enabled. pub use memory_probe::process_rss_bytes; +/// Tune mimalloc to keep freed pages instead of returning them to the OS, and +/// read the setting back. Re-exported for the align stage, which frees and +/// reallocates large per-batch buffers on every pool thread. +pub use memory_probe::{mi_purge_delay_ms, retain_freed_memory}; /// Background read-ahead record reader, re-exported for `fgumi compare bams`, /// which reads two inputs concurrently and needs each decode off the main thread. pub use read_ahead::RawReadAheadReader; diff --git a/crates/fgumi-sort/src/memory_probe.rs b/crates/fgumi-sort/src/memory_probe.rs index df9ca88e5..30952d8b7 100644 --- a/crates/fgumi-sort/src/memory_probe.rs +++ b/crates/fgumi-sort/src/memory_probe.rs @@ -45,7 +45,9 @@ use log::{Level, debug, log_enabled}; #[allow(unused_imports)] // consumed by main fgumi via the crate-root re-export pub use platform_ffi::print_mi_stats; -pub use platform_ffi::{force_mi_collect, process_rss_bytes}; +pub use platform_ffi::{ + force_mi_collect, mi_purge_delay_ms, process_rss_bytes, retain_freed_memory, +}; /// Log target for the memory probe. /// @@ -73,6 +75,33 @@ mod platform_ffi { } } + /// mimalloc's `mi_option_purge_delay`. `libmimalloc-sys` does not export it + /// as a constant, so this is its position in `mi_option_t`, which is 15 in + /// both the v2 and v3 `mimalloc.h` enums (deprecated slots keep their + /// places). `mi_purge_delay_ms_matches_mimalloc_default` pins it against the + /// linked mimalloc v3's default of 1000 ms. + const MI_OPTION_PURGE_DELAY: libmimalloc_sys::mi_option_t = 15; + + /// Stop mimalloc returning freed pages to the OS (`purge_delay = -1`), so a + /// page that is freed and soon reused is not decommitted and faulted back + /// in. Trades a higher peak RSS for fewer page faults and less system time. + pub fn retain_freed_memory() { + // SAFETY: mi_option_set only stores the option value in mimalloc's + // option table; mimalloc reads it on each purge, so setting it after + // start-up is supported, and the option index is fixed (see above). + unsafe { + libmimalloc_sys::mi_option_set(MI_OPTION_PURGE_DELAY, -1); + } + } + + /// mimalloc's current purge delay in milliseconds (`-1`: never purge). + #[must_use] + pub fn mi_purge_delay_ms() -> i64 { + // SAFETY: mi_option_get only reads mimalloc's option table. + #[allow(clippy::useless_conversion)] // `c_long` is `i32` on Windows + i64::from(unsafe { libmimalloc_sys::mi_option_get(MI_OPTION_PURGE_DELAY) }) + } + /// Print mimalloc allocator statistics to stderr via the default output handler. /// /// `mi_stats_print_out(None, null_mut())` uses mimalloc's internal synchronization, @@ -533,6 +562,23 @@ impl Default for MergeProbe { mod tests { use super::*; + /// Pins [`platform_ffi`]'s purge-delay option index: the linked mimalloc v3 + /// defaults `mi_option_purge_delay` to 1000 ms (v2 used 10 ms), and retaining + /// freed memory sets it to -1. Skipped when the environment overrides it. + #[test] + fn mi_purge_delay_ms_matches_mimalloc_default() { + // mimalloc matches its option names case-insensitively. + if std::env::vars_os().any(|(name, _)| { + name.eq_ignore_ascii_case("MIMALLOC_PURGE_DELAY") + || name.eq_ignore_ascii_case("MIMALLOC_RESET_DELAY") + }) { + return; + } + assert_eq!(mi_purge_delay_ms(), 1000); + retain_freed_memory(); + assert_eq!(mi_purge_delay_ms(), -1); + } + #[test] fn test_fmt_bytes_units() { assert_eq!(fmt_bytes(0), "0 B"); diff --git a/src/lib/aligner.rs b/src/lib/aligner.rs index 14c8d14c7..a66865544 100644 --- a/src/lib/aligner.rs +++ b/src/lib/aligner.rs @@ -74,6 +74,31 @@ pub struct AlignerProcess { stderr_thread: Option>>, } +/// The mimalloc option environment variables (`MIMALLOC_PURGE_DELAY` and its +/// legacy name) a user can set to choose the purge delay themselves. +const MIMALLOC_PURGE_ENV: [&str; 2] = ["MIMALLOC_PURGE_DELAY", "MIMALLOC_RESET_DELAY"]; + +/// Whether the user set mimalloc's purge delay in the environment. +pub(crate) fn user_set_mimalloc_purge() -> bool { + names_set_mimalloc_purge(std::env::vars_os().map(|(name, _)| name)) +} + +/// Whether `names` (environment variable names) include either +/// [`MIMALLOC_PURGE_ENV`] name. The comparison ignores ASCII case because +/// mimalloc's own lookup does, so `mimalloc_purge_delay=250` is the user's +/// choice too. Any value counts, including one mimalloc cannot parse: that one +/// is mimalloc's to reject, not ours to replace. +fn names_set_mimalloc_purge(names: I) -> bool +where + I: IntoIterator, + S: AsRef, +{ + names.into_iter().any(|name| { + let name = name.as_ref(); + MIMALLOC_PURGE_ENV.iter().any(|option| name.eq_ignore_ascii_case(option)) + }) +} + impl AlignerProcess { /// Spawn an aligner subprocess with the given shell command. /// @@ -97,13 +122,24 @@ impl AlignerProcess { /// configuring `Stdio::piped()`, which should never happen in /// practice. pub fn spawn(command: &str, ring_size: usize) -> Result { - let mut child = Command::new("/bin/bash") - .args(["-c", command]) + let mut cmd = Command::new("/bin/bash"); + cmd.args(["-c", command]) .stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .with_context(|| format!("failed to spawn aligner command: {command}"))?; + .stderr(Stdio::piped()); + // bwa-mem3 links mimalloc, whose default purge decommits freed pages + // after 1 s and costs the aligner page faults on every batch: never + // purging measured -0.5% wall on the bwa-mem3 CLI alone. A user-set + // value (either name, any case) is inherited unchanged; aligners + // without mimalloc ignore the variable. The variable reaches every + // process in the shell command, so with `--aligner::command` any + // mimalloc-linked tool in the user's pipeline also keeps its freed + // pages (higher RSS); set `MIMALLOC_PURGE_DELAY` to opt out. + if !user_set_mimalloc_purge() { + cmd.env("MIMALLOC_PURGE_DELAY", "-1"); + } + let mut child = + cmd.spawn().with_context(|| format!("failed to spawn aligner command: {command}"))?; let stdin = child.stdin.take(); let stdout = child.stdout.take(); @@ -441,9 +477,13 @@ impl AlignerPreset { /// Build the aligner shell command for this preset. /// - /// The command reads FASTQ from `/dev/stdin` and writes SAM to + /// The command reads FASTQ from `/dev/stdin` and writes its alignments to /// stdout, allowing it to be connected to fgumi's pipeline via - /// stdin/stdout pipes. + /// stdin/stdout pipes. `bwa-mem3` writes uncompressed BAM (`--bam=0`): + /// fgumi's single reader thread then takes records zero-copy instead of + /// parsing SAM text, which at 32 threads kept bwa-mem3 blocked on its + /// output for most of the run (a fused extract → correct → align ran 36% + /// faster). `bwa` has no BAM output and writes SAM. /// /// # Arguments /// @@ -472,7 +512,11 @@ impl AlignerPreset { None => self.binary_name().to_string(), }; let ref_str = reference.display(); - format!("{binary} mem -p -K {chunk_size} -t {threads} {ref_str} /dev/stdin") + let output = match self { + Self::BwaMem3 => " --bam=0", + Self::Bwa => "", + }; + format!("{binary} mem{output} -p -K {chunk_size} -t {threads} {ref_str} /dev/stdin") } /// Validate that the aligner binary and required index files are present. @@ -597,7 +641,7 @@ pub fn substitute_template(template: &str, reference: &Path, threads: usize) -> bail!( "--aligner::command template does not contain `{{ref}}`; \ the aligner needs a reference path. Example: \ - \"bwa-mem3 mem -p -K 150000000 -t {{threads}} {{ref}} /dev/stdin\"" + \"bwa-mem3 mem --bam=0 -p -K 150000000 -t {{threads}} {{ref}} /dev/stdin\"" ); } // Single left-to-right pass so injected text is never re-scanned. Two @@ -707,7 +751,17 @@ impl Default for AlignerOptions { pub(crate) enum ResolvedBackend { /// `--aligner::preset` or `--aligner::command`: the shell command to spawn /// via [`AlignerProcess::spawn`] (already substituted and validated). - Subprocess { command: String }, + Subprocess { + command: String, + /// Whether the aligner's output may carry bwa's mid-pair split (a + /// pair aligned as two unpaired reads because a `-K` chunk cut fell + /// between them). `true` only for the presets, which run `mem -p -K` + /// and so produce exactly that shape; a free-form `--aligner::command` + /// gets the loud "multiple primaries" error instead, because the same + /// shape from an aligner that never pairs would otherwise split every + /// pair silently. + accept_mid_pair_split: bool, + }, } /// Result of [`AlignerOptions::resolve`] — a ready-to-construct aligner @@ -751,6 +805,10 @@ pub(crate) enum ResolvedAlignerMode { Command, } +/// The largest `--aligner::chunk-size` a preset accepts: `i32::MAX`, since bwa +/// and bwa-mem3 parse `-K` with `atoi` into an `int`. +const MAX_PRESET_CHUNK_SIZE: u64 = 2_147_483_647; + impl AlignerOptions { /// Validate the option combination and produce a [`ResolvedAligner`] /// ready for the align stage's `backend_for` to construct. @@ -766,6 +824,7 @@ impl AlignerOptions { /// /// # Errors /// + /// - `--aligner::chunk-size` is zero, or exceeds `i32::MAX` with a preset. /// - Neither `--aligner::preset` nor `--aligner::command` was set. /// - Both were set (mutual exclusion violation). /// - Command mode + a preset-only flag (`--aligner-bin` or @@ -789,6 +848,17 @@ impl AlignerOptions { aligner's -K batch size and the pipeline's in-flight unmapped budget" ); } + // Every preset hands the chunk size to bwa's `-K`, which bwa and + // bwa-mem3 parse with `atoi` into an `int`: a larger value overflows + // there instead of reaching the aligner as given. Command mode sets its + // own `-K` (this flag only sizes its in-flight budget), so it is exempt. + if self.preset.is_some() && self.chunk_size > MAX_PRESET_CHUNK_SIZE { + bail!( + "--aligner::chunk-size must be at most {MAX_PRESET_CHUNK_SIZE} (got {}) with \ + --aligner::preset: bwa parses -K as a 32-bit int", + self.chunk_size + ); + } match (self.preset, self.command) { (None, None) => bail!( "--start-from align requires one of `--aligner::preset` \ @@ -811,7 +881,10 @@ impl AlignerOptions { let command = preset.build_command(reference, threads, self.chunk_size, aligner_bin); Ok(ResolvedAligner { - backend: ResolvedBackend::Subprocess { command }, + // Both presets run `mem -p -K`, whose smart pairing splits + // a pair across a chunk cut (bwa and bwa-mem3 share + // `bseq_read`'s even-count cut and `bseq_classify`). + backend: ResolvedBackend::Subprocess { command, accept_mid_pair_split: true }, chunk_size: self.chunk_size, threads: Some(threads), mode: ResolvedAlignerMode::Preset(preset), @@ -846,7 +919,7 @@ impl AlignerOptions { // hardcode any thread count). let command = substitute_template(&template, reference, top_threads)?; Ok(ResolvedAligner { - backend: ResolvedBackend::Subprocess { command }, + backend: ResolvedBackend::Subprocess { command, accept_mid_pair_split: false }, chunk_size: self.chunk_size, threads: None, mode: ResolvedAlignerMode::Command, @@ -923,6 +996,22 @@ mod tests { use super::*; + #[rstest] + #[case::current_name(&["MIMALLOC_PURGE_DELAY"], true)] + #[case::legacy_name(&["MIMALLOC_RESET_DELAY"], true)] + #[case::lowercase(&["mimalloc_purge_delay"], true)] + #[case::mixed_case_legacy(&["Mimalloc_Reset_Delay"], true)] + #[case::among_others(&["PATH", "HOME", "mimalloc_purge_delay"], true)] + #[case::unset(&["PATH", "HOME"], false)] + #[case::other_mimalloc_option(&["MIMALLOC_VERBOSE", "MIMALLOC_ARENA_EAGER_COMMIT"], false)] + #[case::prefix_only(&["MIMALLOC_PURGE_DELAY_MS", "X_MIMALLOC_PURGE_DELAY"], false)] + fn names_set_mimalloc_purge_matches_like_mimalloc( + #[case] names: &[&str], + #[case] expected: bool, + ) { + assert_eq!(names_set_mimalloc_purge(names), expected); + } + /// A non-UTF-8 stderr line must NOT halt the relay: draining has to /// continue to EOF, otherwise the aligner blocks once its stderr pipe fills /// (~64 KiB) and the whole align pipeline deadlocks. Regression test for the @@ -1054,7 +1143,7 @@ mod tests { AlignerPreset::BwaMem3, 4, 5_000_000, - "bwa-mem3 mem -p -K 5000000 -t 4 /ref/genome.fa /dev/stdin" + "bwa-mem3 mem --bam=0 -p -K 5000000 -t 4 /ref/genome.fa /dev/stdin" )] fn preset_build_command_default_binary( #[case] preset: AlignerPreset, @@ -1149,7 +1238,8 @@ mod tests { chunk_size: DEFAULT_ALIGNER_CHUNK_SIZE, }; let resolved = opts.resolve(&ref_path, 4, None).unwrap(); - let ResolvedBackend::Subprocess { command } = resolved.backend; + let ResolvedBackend::Subprocess { command, accept_mid_pair_split } = resolved.backend; + assert!(!accept_mid_pair_split, "command mode must not accept bwa's mid-pair split"); assert!(command.contains(&ref_path.display().to_string())); assert!(command.contains("-t 4")); assert!(matches!(resolved.mode, ResolvedAlignerMode::Command)); @@ -1220,6 +1310,26 @@ mod tests { assert!(msg.contains("bwa-mem3 index"), "expected fix-it hint, got: {msg}"); } + /// A preset that passes validation resolves to the subprocess backend with + /// bwa's mid-pair split accepted: both presets run `mem -p -K`. + #[rstest] + #[case::bwa_mem3(AlignerPreset::BwaMem3)] + #[case::bwa(AlignerPreset::Bwa)] + fn resolve_preset_accepts_mid_pair_split(#[case] preset: AlignerPreset) { + let tmp = tempfile::tempdir().unwrap(); + let ref_path = tmp.path().join("ref.fa"); + std::fs::write(&ref_path, b">chr1\nACGT\n").unwrap(); + for ext in preset.index_extensions() { + std::fs::write(append_extension(&ref_path, ext), b"x").unwrap(); + } + let bin = make_existing_binary(tmp.path()); + let opts = AlignerOptions { preset: Some(preset), ..AlignerOptions::default() }; + let resolved = opts.resolve(&ref_path, 4, Some(&bin)).unwrap(); + let ResolvedBackend::Subprocess { command, accept_mid_pair_split } = resolved.backend; + assert!(accept_mid_pair_split, "preset mode must accept bwa's mid-pair split"); + assert!(command.starts_with(&bin.display().to_string()), "got: {command}"); + } + /// `validate` errors out on a nonexistent `--aligner-bin` override path. #[test] fn test_validate_override_path_missing() { @@ -1327,6 +1437,40 @@ mod tests { assert_eq!(result, "bwa-mem3 mem -t 8 /data/{threads}/genome.fa /dev/stdin"); } + /// `resolve` caps a preset's `--aligner::chunk-size` at `i32::MAX`, since bwa + /// and bwa-mem3 parse `-K` with `atoi` into an `int`; the check runs before + /// preset validation, so the reference needs no index. + #[rstest] + #[case::bwa_mem3(AlignerPreset::BwaMem3, 2_147_483_648)] + #[case::bwa(AlignerPreset::Bwa, 4_294_967_296)] + fn resolve_rejects_preset_chunk_size_past_i32( + #[case] preset: AlignerPreset, + #[case] chunk_size: u64, + ) { + let tmp = tempfile::tempdir().unwrap(); + let ref_path = tmp.path().join("ref.fa"); + std::fs::write(&ref_path, b">chr1\nACGT\n").unwrap(); + let opts = AlignerOptions { preset: Some(preset), chunk_size, ..AlignerOptions::default() }; + let err = opts.resolve(&ref_path, 4, None).unwrap_err().to_string(); + let needle = format!("--aligner::chunk-size must be at most 2147483647 (got {chunk_size})"); + assert!(err.contains(&needle), "got: {err}"); + } + + /// Command mode sets its own `-K`, so a chunk size past `i32::MAX` (which + /// only sizes its in-flight budget) is accepted. + #[test] + fn resolve_accepts_command_mode_chunk_size_past_i32() { + let tmp = tempfile::tempdir().unwrap(); + let ref_path = tmp.path().join("ref.fa"); + std::fs::write(&ref_path, b">chr1\nACGT\n").unwrap(); + let opts = AlignerOptions { + command: Some("bwa-mem3 mem -t {threads} {ref} /dev/stdin".to_string()), + chunk_size: 2_147_483_648, + ..AlignerOptions::default() + }; + assert!(opts.resolve(&ref_path, 4, None).is_ok()); + } + /// `resolve` rejects `--aligner::chunk-size 0` (would reach the aligner as /// `-K 0` and zero the in-flight budget) with a flag-attributed error. #[test] diff --git a/src/lib/pipeline/steps/align/inproc/cohort.rs b/src/lib/pipeline/steps/align/inproc/cohort.rs new file mode 100644 index 000000000..9dfb51267 --- /dev/null +++ b/src/lib/pipeline/steps/align/inproc/cohort.rs @@ -0,0 +1,586 @@ +//! Cohort cutter, SE/PE layout, and per-sub-batch id-base arithmetic. +//! +//! These are the parity-critical *pure* functions of the in-process aligner: +//! they reproduce, exactly, three upstream bwa-mem3 rules so that the merged +//! output is byte-identical to the subprocess `bwa-mem3` preset: the `-K` cohort +//! cut, the SE/PE layout, and the per-sub-batch read-id bases. Each is pinned +//! by a proptest against a literal Rust port of the upstream rule in the test +//! module below. +//! +//! Nothing here calls FFI or wires a pipeline step. +//! +//! These pure primitives are the foundation the in-process backend is built on: +//! the pipeline steps that construct and thread the work items through the pool +//! (`AlignPrepareStep`, `AlignSeedExtendStep`, `CohortPeStatStep`, +//! `AlignPairEmitStep`) consume the types and functions here. + +use std::sync::Arc; + +use crate::pipeline::core::item::HeapSize; +use crate::template::Template; + +use super::engine::{AlignEngine, IdBases}; +use super::gate::CohortLease; + +/// Reproduces bwa-mem3's `-K` even-parity cohort cut +/// (`fast_reader_bseq.c:129-137`). +/// +/// bwa-mem3's fast reader appends reads to the current cohort one at a time, +/// accumulating `size += l_seq` after each, and closes the cohort as soon as +/// `size >= chunk_size && (n & 1) == 0` — i.e. the accumulated bases have +/// reached `-K` **and** an even number of reads has been buffered. The parity +/// test matters: with `-p` smart pairing an even read count keeps read pairs +/// together at a cohort boundary in the common all-paired case, but a cut can +/// still fall *between* the two reads of a pair when the running read count is +/// odd at the pair's start (a mixed SE/PE input) — bwa then classifies the two +/// reads as two separate singles in adjacent cohorts. +/// +/// The cut is therefore evaluated **per read**, not per template, which is why +/// the cutter's state is a running byte size and read count rather than a +/// per-template tally. +pub(crate) struct CohortCutter { + /// The `-K` chunk size in bases; a cohort closes once its accumulated read + /// bases reach this and the read count is even. + chunk_size: u64, + /// Accumulated `l_seq` of the reads buffered in the current (open) cohort. + size: u64, + /// Number of reads buffered in the current (open) cohort. + n_reads_in_cohort: u64, +} + +impl CohortCutter { + /// Create a cutter for the given `-K` chunk size (bases). Callers pass the + /// validated aligner chunk size (`>= 1`), matching the CLI. + pub(crate) fn new(chunk_size: u64) -> Self { + Self { chunk_size, size: 0, n_reads_in_cohort: 0 } + } + + /// Feed one template's surviving reads (1 for single-end, 2 for a pair) + /// through the per-read cut rule, in order. Returns `true` iff the cohort + /// boundary falls **exactly after this template** — i.e. the template's + /// last read triggered the even-parity cut, so the current cohort closes + /// here and the next template opens a new one. + /// + /// A cut that fires *between* a pair's two reads (possible only when the + /// running read count is odd at the pair's start, in a mixed SE/PE input) + /// leaves the boundary inside the template: the second read correctly opens + /// the next cohort in the cutter's state, but this method returns `false` + /// because the boundary is not after the template. Splitting such a + /// template into two singles across the two cohorts is the caller's concern + /// (`AlignPrepareStep` splits it); the cutter's own running state stays exact + /// regardless. + /// + /// Production drives the cutter one read at a time via [`Self::push_read`] (so + /// it can observe a mid-pair cut); this whole-template convenience is used + /// only by the parity proptest, hence `#[cfg(test)]`. + #[cfg(test)] + pub(crate) fn push_template(&mut self, l_seqs: impl Iterator) -> bool { + let mut cut_after_last = false; + for l_seq in l_seqs { + cut_after_last = self.push_read(l_seq); + } + cut_after_last + } + + /// Feed **one** read's `l_seq` through the per-read cut rule and return + /// `true` iff the cohort closes **after this read** — i.e. the accumulated + /// bases have reached `-K` and the running read count is now even + /// (`fast_reader_bseq.c:129-137`). On a cut the running accumulators reset, + /// so the next `push_read` opens a fresh cohort. + /// + /// This per-read granularity is what lets the caller + /// (`AlignPrepareStep`) observe a + /// **mid-pair cut**: feeding a pair's two reads individually, a `true` after + /// the *first* read means the cohort boundary falls between the pair's reads, + /// and bwa then classifies them as two separate singles in adjacent cohorts. + /// `Self::push_template` cannot express that (its single + /// `bool` reports only a cut after the *last* read), so the prepare step + /// drives the cutter one read at a time via this method. + pub(crate) fn push_read(&mut self, l_seq: u32) -> bool { + self.size += u64::from(l_seq); + self.n_reads_in_cohort += 1; + if self.size >= self.chunk_size && self.n_reads_in_cohort.is_multiple_of(2) { + // Close the cohort after this read and open a fresh one. + self.size = 0; + self.n_reads_in_cohort = 0; + true + } else { + false + } + } +} + +/// Where one template of a sub-batch sits in bwa-mem3's `-p` SE/PE +/// classification. `idx` is the template's ordinal **within its group** — its +/// index into bwa's `pairs[]` or `singles[]` array — which is what the id +/// formulas and the record-emission origin indices key off. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum TemplateSlot { + /// A paired template (two reads with equal names); the `idx`-th pair. + Pair { idx: u32 }, + /// A single-end template (one read); the `idx`-th single. + Single { idx: u32 }, +} + +/// The SE/PE classification of a sub-batch's templates, in input order. +/// +/// [`Layout::AllPairs`] is the production shape — every template is a full pair, +/// so a template's index is exactly its pair index and no per-template vector is +/// needed. [`Layout::Mixed`] carries one [`TemplateSlot`] per template for the +/// (rare) SE / mixed inputs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum Layout { + /// Every template contributes two reads: template `i` is pair `i`. + AllPairs, + /// One slot per template, in input order. + Mixed(Vec), +} + +impl Layout { + /// Classify a sub-batch's templates into bwa-mem3's `-p` SE/PE groups from + /// each template's surviving-read count (2 = pair, otherwise single), + /// reproducing `bseq_classify` (`bwa.cpp:258-274`). + /// + /// Upstream pairs two *consecutive equal-named reads*; because + /// `GroupByQueryname` never emits two consecutive same-name templates, a + /// template is a pair exactly when it contributes two reads and a single + /// otherwise. A non-empty sub-batch whose templates are + /// all pairs collapses to [`Layout::AllPairs`]; any single (or an empty + /// sub-batch) yields [`Layout::Mixed`]. + pub(crate) fn classify(read_counts: &[u8]) -> Self { + if !read_counts.is_empty() && read_counts.iter().all(|&c| c == 2) { + return Self::AllPairs; + } + let mut slots = Vec::with_capacity(read_counts.len()); + let mut pair_idx: u32 = 0; + let mut single_idx: u32 = 0; + for &count in read_counts { + if count == 2 { + slots.push(TemplateSlot::Pair { idx: pair_idx }); + pair_idx += 1; + } else { + slots.push(TemplateSlot::Single { idx: single_idx }); + single_idx += 1; + } + } + Self::Mixed(slots) + } +} + +/// Cohort-level facts a sub-batch carries so the barrier can compute its id +/// bases. Offsets are counts *within this cohort* accumulated +/// over the earlier sub-batches, in input order. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct CohortPos { + /// Global reads before this cohort — bwa's `n_processed` at cohort start. + pub(crate) cohort_read_base: u64, + /// Single-end reads in earlier sub-batches of this cohort. + pub(crate) se_offset: u64, + /// Paired templates in earlier sub-batches of this cohort. + pub(crate) pe_offset: u64, + /// Single-end reads in this sub-batch. + pub(crate) n_se: u32, + /// Paired templates in this sub-batch. + pub(crate) n_pe: u32, +} + +/// Identity of a sub-batch. `serial` is dense from 0 across the whole run and +/// becomes `BamTemplateBatch::batch_serial`; `(cohort, index_in_cohort)` is for +/// the cohort barrier. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct SubBatchId { + pub(crate) serial: u64, + pub(crate) cohort: u32, + pub(crate) index_in_cohort: u32, +} + +/// Attached to the **last** sub-batch of a cohort (only known at cut time): the +/// cohort's sub-batch count and total SE/PE tallies, which the barrier needs to +/// know a cohort is complete and to compute `pe_n_processed`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct CohortCloser { + pub(crate) n_sub_batches: u32, + pub(crate) cohort_n_se: u64, + pub(crate) cohort_n_pe: u64, +} + +/// Global-ordinal id bases for one sub-batch, reproducing bwa-mem3's `worker_sam` +/// ids over a `-p` cohort (`fastmap.cpp:908-958` + `bwamem.cpp:2795-2884`). +/// +/// bwa processes a cohort's SE group with `n_processed = cohort_read_base` and +/// its PE group with `n_processed = cohort_read_base + cohort_n_se` +/// (`fastmap.cpp:924-944`); `worker_sam` then assigns single `i` the id +/// `n_processed + i` and pair `i` the id `(n_processed >> 1) + i` +/// (`bwamem.cpp:2795-2884`). For a sub-batch that is preceded within its cohort +/// by `se_offset` singles and `pe_offset` pairs, the *first* single/pair ids are +/// therefore: +/// +/// - `first_single_id = cohort_read_base + se_offset` +/// - `first_pair_id = ((cohort_read_base + cohort_n_se) >> 1) + pe_offset` +/// +/// and the engine adds the per-read index `i` on top. The `>> 1` is bwa's own +/// truncating shift on the pair `n_processed`, reproduced verbatim. +pub(crate) fn id_bases(pos: &CohortPos, cohort_n_se: u64) -> IdBases { + IdBases { + first_single_id: pos.cohort_read_base + pos.se_offset, + first_pair_id: ((pos.cohort_read_base + cohort_n_se) >> 1) + pos.pe_offset, + } +} + +/// Heap footprint shared by the in-process work items: the summed +/// `Template::heap_size` of the moved-out unmapped templates. The C-side read +/// copies and alignment regions are not charged here: they live in the cohort's +/// resident state (owned by the [`CohortLease`], freed with its last clone), +/// not in any one sub-batch, and the lease's gate reservation is what bounds +/// how many cohorts hold them at once. +fn work_heap_size(unmapped: &[Template]) -> usize { + unmapped.iter().map(Template::heap_size).sum::() +} + +/// A sub-batch prepared for seeding: its identity, cohort position, SE/PE +/// layout, and the moved-out unmapped templates in input order, carrying the +/// cohort's [`CohortLease`]. Produced by `AlignPrepareStep`; consumed by +/// `AlignSeedExtendStep`. +pub(crate) struct AlignWork { + pub(crate) id: SubBatchId, + pub(crate) pos: CohortPos, + /// `Some` on the cohort's last sub-batch. + pub(crate) closer: Option, + pub(crate) layout: Layout, + pub(crate) unmapped: Vec