From 785a5b9fe33fba43623dcff25eb3cf74a3b40a47 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 27 Aug 2026 20:58:48 +0200 Subject: [PATCH 1/2] bench: add a divan harness and the first hot-path benchmarks Several open dependency questions (#162, #202, #205, #208) all start with "measure first", and each was about to invent its own measurement. This is the shared one. Divan rather than criterion: the questions that prompted a harness compare peak memory as much as speed, and divan reports allocation counts next to wall time without extra setup; a full run also takes seconds, so it stays usable in a pull request rather than only in a nightly job. Criterion is the better choice if the CI integration in #245 needs its report format, and swapping is a benches/ change rather than an API one. Three groups, chosen as the functions the open questions would replace: seed-extension scanning (`find_stop`, what a portable-SIMD crate would swap), the gene-overlap query and the annotation build (what an interval crate would swap). All are pure functions with no genome index to build, so a run takes seconds and a regression is attributable to one function. `find_stop` and the GTF record type are private to the crate; rather than widen the published API for a benchmark, both are exposed under a `bench` feature that only benches/ enables. Closes #204. --- Cargo.lock | 50 +++++++++++++ Cargo.toml | 15 ++++ benches/hot_paths.rs | 163 +++++++++++++++++++++++++++++++++++++++++++ src/align/mod.rs | 7 ++ src/junction/mod.rs | 6 ++ 5 files changed, 241 insertions(+) create mode 100644 benches/hot_paths.rs diff --git a/Cargo.lock b/Cargo.lock index 9f77c86d..729be5ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -237,6 +237,7 @@ dependencies = [ "anstyle", "clap_lex", "strsim", + "terminal_size", ] [[package]] @@ -263,6 +264,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "condtype" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf0a07a401f374238ab8e2f11a104d2851bf9ce711ec69804834de8af45c7af" + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -329,6 +336,31 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6184e33543162437515c2e2b48714794e37845ec9851711914eec9d308f6ebe8" +[[package]] +name = "divan" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a405457ec78b8fe08b0e32b4a3570ab5dff6dd16eb9e76a5ee0a9d9cbd898933" +dependencies = [ + "cfg-if", + "clap", + "condtype", + "divan-macros", + "libc", + "regex-lite", +] + +[[package]] +name = "divan-macros" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9556bc800956545d6420a640173e5ba7dfa82f38d3ea5a167eb555bc69ac3323" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "either" version = "1.15.0" @@ -959,6 +991,12 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "regex-lite" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" + [[package]] name = "regex-syntax" version = "0.8.9" @@ -978,6 +1016,7 @@ dependencies = [ "chrono", "clap", "dashmap", + "divan", "env_logger", "flate2", "libdeflater", @@ -989,6 +1028,7 @@ dependencies = [ "noodles-bgzf", "predicates", "rayon", + "rustar-aligner", "rustc-hash", "shlex", "tempfile", @@ -1134,6 +1174,16 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "terminal_size" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" +dependencies = [ + "rustix", + "windows-sys", +] + [[package]] name = "termtree" version = "0.5.1" diff --git a/Cargo.toml b/Cargo.toml index 8a4638f9..ae0ad47f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,9 +66,24 @@ libmimalloc-sys = { version = "0.1.49", features = ["extended"] } # mi_option_se libdeflater = "1.25.2" noodles-bgzf = { version = "0.49", features = ["libdeflate"] } +[features] +# Exposes a few internals to `benches/` without widening the published API. +# Enabled automatically for `cargo bench` through dev-dependencies below. +bench = [] + [dev-dependencies] assert_cmd = "2" predicates = "3" +# Benchmark harness. Divan over criterion for the questions that prompted it +# (#204): it reports allocation counts next to wall time, which is what the +# suffix-array and interval evaluations are actually comparing, and a run +# takes seconds rather than minutes. See DEPENDENCIES.md. +divan = "0.1" +rustar-aligner = { path = ".", features = ["bench"] } + +[[bench]] +name = "hot_paths" +harness = false [build-dependencies] chrono = { version = "0.4", default-features = false, features = ["clock"] } diff --git a/benches/hot_paths.rs b/benches/hot_paths.rs new file mode 100644 index 00000000..a353da9c --- /dev/null +++ b/benches/hot_paths.rs @@ -0,0 +1,163 @@ +//! Micro-benchmarks for the paths the open "measure first" questions land on. +//! +//! Several dependency decisions (#162, #202, #205, #208) all begin with a +//! measurement, and each was about to invent its own. These are the shared +//! ones, deliberately small: pure functions with no genome index to build, so +//! a run takes seconds and a regression is attributable to one function. +//! +//! ```text +//! cargo bench # everything +//! cargo bench -- seed_scan # one group +//! cargo bench -- --sample-count 200 # more samples +//! ``` +//! +//! Divan reports allocation counts alongside wall time, which matters here: +//! two of the open questions (`sufr`/`libsais` against `caps-sa`) are about +//! peak RSS as much as speed. + +use std::collections::HashMap; + +use rustar_aligner::align::simd_scan::find_stop; +use rustar_aligner::genome::Genome; +use rustar_aligner::junction::gtf::GtfRecord; +use rustar_aligner::quant::GeneAnnotation; + +fn main() { + divan::main(); +} + +/// Deterministic pseudo-random bases (0..=3), the same generator the tests use. +fn lcg_bases(seed: u32, length: usize) -> Vec { + let mut state = seed; + (0..length) + .map(|_| { + state = state.wrapping_mul(1_103_515_245).wrapping_add(12345); + ((state >> 16) & 3) as u8 + }) + .collect() +} + +// ── Seed extension ────────────────────────────────────────────────────────── +// +// `find_stop` is the innermost loop of seed extension: it walks read against +// genome and returns the first mismatch or padding byte. It runs once per +// candidate extension, which is millions of times per million reads, and it is +// the function a portable-SIMD crate (#205) would replace. + +#[divan::bench(args = [50, 100, 250])] +fn seed_scan_full_match(bencher: divan::Bencher, len: usize) { + let read = lcg_bases(1, len); + let genome = read.clone(); + bencher.bench(|| find_stop(divan::black_box(&read), divan::black_box(&genome))); +} + +#[divan::bench(args = [50, 100, 250])] +fn seed_scan_stop_at_half(bencher: divan::Bencher, len: usize) { + // The common case in practice: a long run of matches, then a mismatch. + let read = lcg_bases(1, len); + let mut genome = read.clone(); + genome[len / 2] = (genome[len / 2] + 1) % 4; + bencher.bench(|| find_stop(divan::black_box(&read), divan::black_box(&genome))); +} + +// ── Gene overlap ──────────────────────────────────────────────────────────── +// +// The segment-tree overlap query runs per read (twice for `GeneFull`) and was +// the top solo hotspot before it replaced a linear scan. It is what the +// interval-crate question (#208) would replace. + +fn synthetic_annotation(n_genes: usize) -> (GeneAnnotation, Genome) { + let genome_len = (n_genes as u64 + 2) * 1_000; + let genome = Genome { + transform_blocks: None, + sequence: vec![0u8; genome_len as usize].into(), + n_genome: genome_len, + n_genome_real: genome_len, + n_chr_real: 1, + chr_start: vec![0, genome_len], + chr_length: vec![genome_len], + chr_name: vec!["chr1".to_string()], + }; + + let exons: Vec = (0..n_genes) + .flat_map(|g| { + // Two exons per gene, genes 1 kb apart with a little overlap + // between neighbours so the query cannot stop at the first hit. + let base = g as u64 * 1_000 + 100; + [(base, base + 400), (base + 600, base + 1_100)] + .into_iter() + .map(move |(s, e)| { + let mut attributes = HashMap::new(); + attributes.insert("gene_id".to_string(), format!("G{g}")); + attributes.insert("transcript_id".to_string(), format!("G{g}_T1")); + GtfRecord { + seqname: "chr1".to_string(), + feature: "exon".to_string(), + start: s + 1, + end: e, + strand: '+', + attributes, + } + }) + }) + .collect(); + + (GeneAnnotation::from_gtf_exons(&exons, &genome), genome) +} + +#[divan::bench(args = [100, 2_000, 20_000])] +fn gene_overlap_query(bencher: divan::Bencher, n_genes: usize) { + use rustar_aligner::align::transcript::{Exon, Transcript}; + + let (ann, _genome) = synthetic_annotation(n_genes); + // A read landing in the middle of the annotation, spanning two exons. + let mid = (n_genes as u64 / 2) * 1_000 + 150; + let transcript = Transcript { + chr_idx: 0, + genome_start: mid, + genome_end: mid + 800, + is_reverse: false, + exons: vec![ + Exon { + genome_start: mid, + genome_end: mid + 200, + read_start: 0, + read_end: 200, + i_frag: 0, + }, + Exon { + genome_start: mid + 600, + genome_end: mid + 800, + read_start: 200, + read_end: 400, + i_frag: 0, + }, + ], + cigar: Vec::new(), + score: 0, + n_mismatch: 0, + n_gap: 0, + n_junction: 1, + junction_motifs: Vec::new(), + junction_annotated: Vec::new(), + }; + + let mut out = Vec::new(); + bencher.bench_local(|| { + ann.overlapping_genes_into(divan::black_box(&transcript), &mut out); + out.len() + }); +} + +// ── Annotation build ──────────────────────────────────────────────────────── +// +// Building the annotation is once per run, but it is O(exons log exons) and +// shows up on the solo startup path; it is also the allocation-heavy half of +// the interval question. + +#[divan::bench(args = [2_000, 20_000])] +fn gene_annotation_build(bencher: divan::Bencher, n_genes: usize) { + bencher + .with_inputs(|| n_genes) + .bench_values(|n| synthetic_annotation(n).0.n_genes()); +} diff --git a/src/align/mod.rs b/src/align/mod.rs index 1c0351ed..c4e15747 100644 --- a/src/align/mod.rs +++ b/src/align/mod.rs @@ -2,7 +2,14 @@ pub mod pe_overlap; pub mod read_align; pub mod score; pub mod seed; +// Private to the crate, except under the `bench` feature: `benches/hot_paths.rs` +// measures `find_stop` directly, because it is the function a portable-SIMD +// crate would replace (#205) and the one whose cost is worth watching. +#[cfg(not(feature = "bench"))] mod simd_scan; +#[cfg(feature = "bench")] +#[doc(hidden)] +pub mod simd_scan; pub mod stitch; pub mod transcript; diff --git a/src/junction/mod.rs b/src/junction/mod.rs index 13776ce4..36f483c2 100644 --- a/src/junction/mod.rs +++ b/src/junction/mod.rs @@ -6,7 +6,13 @@ /// - Junction lookup during alignment (annotated vs novel) /// - Junction statistics collection for SJ.out.tab output pub(crate) mod chr_start_end; +#[cfg(not(feature = "bench"))] pub(crate) mod gtf; +// See `align::simd_scan`: exposed only for `benches/`, which builds a gene +// annotation from in-memory GTF records rather than from a file on disk. +#[cfg(feature = "bench")] +#[doc(hidden)] +pub mod gtf; mod sj_output; pub mod sjdb_insert; From 6edbe741cdaa79bb611e6d07428ba5b32443efad Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 27 Aug 2026 21:05:49 +0200 Subject: [PATCH 2/2] ci: opt-in benchmark workflow, micro and end-to-end Two jobs, both opt-in the way the large-dataset job is (manual dispatch or a `benchmark` label), because a wall time from a busy shared runner is worse than no wall time. - micro: runs the divan benches, optionally against a baseline ref, and summarises them as a table with anything past 10% flagged. Handles a baseline older than the harness by saying so rather than failing. - end-to-end: times indexing and alignment separately for both aligners on the nf-core fixture, reporting wall time and peak RSS. Deliberately without thresholds: it is a data point, and only comparable within one machine. test/bench_report.py parses divan's tree output into the table. It was checked both ways: it recovers the group/argument names from real output, and with one baseline row perturbed it flags that row at +41% and leaves the others at 0%. Closes #245. --- .github/workflows/benchmark.yml | 111 ++++++++++++++++++++++++++++++++ CONTRIBUTING.md | 17 +++++ test/bench_report.py | 86 +++++++++++++++++++++++++ test/speed_bench.py | 107 ++++++++++++++++++++++++++++++ 4 files changed, 321 insertions(+) create mode 100644 .github/workflows/benchmark.yml create mode 100644 test/bench_report.py create mode 100644 test/speed_bench.py diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml new file mode 100644 index 00000000..d5454509 --- /dev/null +++ b/.github/workflows/benchmark.yml @@ -0,0 +1,111 @@ +name: Benchmarks + +# Benchmarks are opt-in for the same reason the large-dataset job is: they are +# long, and a number measured on a noisy shared runner is worse than no number. +on: + workflow_dispatch: + inputs: + runner: + description: "Runner label (e.g. ubuntu-latest, or a scverse AWS runner group)" + required: false + default: "ubuntu-latest" + baseline: + description: "Git ref to compare against (empty = no comparison)" + required: false + default: "main" + pull_request: + types: [labeled, synchronize] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + micro: + name: Micro-benchmarks (divan) + if: >- + github.event_name == 'workflow_dispatch' || + contains(github.event.pull_request.labels.*.name, 'benchmark') + runs-on: ${{ github.event.inputs.runner || 'ubuntu-latest' }} + timeout-minutes: 45 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # ratchet:actions/checkout@v7.0.1 + with: + fetch-depth: 0 + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@631a55b12751854ce901bb631d5902ceb48146f7 # ratchet:dtolnay/rust-toolchain@stable + + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # ratchet:Swatinem/rust-cache@v2.9.1 + + - name: Benchmark this ref + run: cargo bench --bench hot_paths | tee "${RUNNER_TEMP}/bench-head.txt" + + - name: Benchmark the baseline + if: ${{ github.event.inputs.baseline != '' }} + run: | + set -eu + base="${{ github.event.inputs.baseline || 'main' }}" + git worktree add "${RUNNER_TEMP}/baseline" "origin/${base}" + cd "${RUNNER_TEMP}/baseline" + # A baseline older than the harness has no benches to run; say so + # rather than failing the job. + if [ -f benches/hot_paths.rs ]; then + cargo bench --bench hot_paths | tee "${RUNNER_TEMP}/bench-base.txt" + else + echo "baseline ${base} predates benches/hot_paths.rs; nothing to compare" \ + | tee "${RUNNER_TEMP}/bench-base.txt" + fi + + - name: Summarise + if: always() + run: python3 test/bench_report.py "${RUNNER_TEMP}/bench-head.txt" "${RUNNER_TEMP}/bench-base.txt" >> "${GITHUB_STEP_SUMMARY}" + + - name: Upload raw output + if: always() + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # ratchet:actions/upload-artifact@v4.6.2 + with: + name: benchmark-output + path: ${{ runner.temp }}/bench-*.txt + if-no-files-found: warn + + end-to-end: + name: End-to-end wall time against STAR + if: >- + github.event_name == 'workflow_dispatch' || + contains(github.event.pull_request.labels.*.name, 'benchmark') + runs-on: ${{ github.event.inputs.runner || 'ubuntu-latest' }} + timeout-minutes: 60 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # ratchet:actions/checkout@v7.0.1 + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@631a55b12751854ce901bb631d5902ceb48146f7 # ratchet:dtolnay/rust-toolchain@stable + + - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # ratchet:Swatinem/rust-cache@v2.9.1 + + - run: cargo build --release + + - name: Install STAR + run: | + sudo apt-get update + sudo apt-get install -y rna-star + + - name: Time both aligners + run: >- + python3 test/speed_bench.py + --rustar ./target/release/rustar-aligner + --work "${RUNNER_TEMP}/speed" + --json "${RUNNER_TEMP}/speed.json" + >> "${GITHUB_STEP_SUMMARY}" + + - name: Upload timings + if: always() + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # ratchet:actions/upload-artifact@v4.6.2 + with: + name: speed-metrics + path: ${{ runner.temp }}/speed.json + if-no-files-found: warn diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 551505ff..83bb360f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -63,6 +63,23 @@ Adding a dependency — **especially a non-Rust one** (a C library via a `-sys` If you add a CLI flag that parses but is not yet implemented, mark it as such in the parameter-surface test and document it — do not silently accept a flag that does nothing. A user passing a flag should never be quietly ignored. +## Benchmarks + +`cargo bench` runs the divan micro-benchmarks in `benches/hot_paths.rs` +(seed-extension scanning, the gene-overlap query, the annotation build). A full +run takes seconds, and divan reports allocation counts next to wall time. + +```bash +cargo bench # everything +cargo bench -- seed_scan # one group +python3 test/speed_bench.py # end-to-end wall time and peak RSS vs STAR +``` + +In CI the `Benchmarks` workflow is opt-in: manual dispatch, or the `benchmark` +label on a pull request. Dispatch takes a runner label and a baseline ref; the +summary flags any micro-benchmark that moved by more than 10%. Treat a flag +from a shared runner as a prompt to re-run, not as a verdict. + ## Test data Integration tests in `tests/` use a bundled synthetic micro-genome and need no downloads. The differential benchmark below uses a small **public** yeast RNA-seq dataset that is not vendored; fetch it once and point `DATA` at wherever you keep it. diff --git a/test/bench_report.py b/test/bench_report.py new file mode 100644 index 00000000..f1f6a8b5 --- /dev/null +++ b/test/bench_report.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Turn divan's output into a Markdown summary, optionally against a baseline. + +Divan prints a tree; this pulls out the leaf rows (name, median) so a job +summary shows numbers rather than box drawing, and flags anything that moved +by more than 10% against the baseline run. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +# A leaf row: name, then four "value unit" columns separated by │. The third +# is the median, which is the one worth reporting. +ROW = re.compile( + r"^[\s│├╰─]*([A-Za-z0-9_.]+)\s+" + r"[\d.]+\s*(?:ns|µs|ms|s)\s*│\s*" + r"[\d.]+\s*(?:ns|µs|ms|s)\s*│\s*" + r"([\d.]+)\s*(ns|µs|ms|s)" +) +# A group header: a name with the columns empty. +GROUP = re.compile(r"^[\s│├╰─]*([A-Za-z0-9_]+)\s+│\s*│\s*│\s*│") +SCALE = {"ns": 1e-9, "µs": 1e-6, "ms": 1e-3, "s": 1.0} + + +def parse(path: Path) -> dict[str, float]: + out: dict[str, float] = {} + if not path.exists(): + return out + group = "" + for line in path.read_text(errors="replace").splitlines(): + g = GROUP.match(line) + if g: + group = g.group(1) + continue + m = ROW.match(line) + if not m: + continue + name, median, unit = m.groups() + key = f"{group}/{name}" if group and group != name else name + out[key] = float(median) * SCALE[unit] + return out + + +def fmt(seconds: float) -> str: + for unit, scale in (("s", 1.0), ("ms", 1e-3), ("µs", 1e-6), ("ns", 1e-9)): + if seconds >= scale: + return f"{seconds / scale:.2f} {unit}" + return f"{seconds * 1e9:.2f} ns" + + +def main() -> int: + head = parse(Path(sys.argv[1])) if len(sys.argv) > 1 else {} + base = parse(Path(sys.argv[2])) if len(sys.argv) > 2 else {} + + if not head: + print("### Benchmarks\n\nNo benchmark rows parsed; see the uploaded raw output.") + return 0 + + print("### Benchmarks (median)\n") + if base: + print("| benchmark | this ref | baseline | change |") + print("|---|---|---|---|") + for name in sorted(head): + h = head[name] + b = base.get(name) + if b is None: + print(f"| {name} | {fmt(h)} | — | new |") + continue + delta = (h - b) / b if b else 0.0 + flag = " ⚠️" if abs(delta) > 0.10 else "" + print(f"| {name} | {fmt(h)} | {fmt(b)} | {delta:+.1%}{flag} |") + print("\nA change above 10% is flagged; on a shared runner treat it as a") + print("prompt to re-run rather than as a verdict.") + else: + print("| benchmark | median |") + print("|---|---|") + for name in sorted(head): + print(f"| {name} | {fmt(head[name])} |") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/speed_bench.py b/test/speed_bench.py new file mode 100644 index 00000000..a2d91fe0 --- /dev/null +++ b/test/speed_bench.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Wall time and peak memory for rustar-aligner against STAR, end to end. + +Uses the same nf-core/rnaseq fixture as `nfcore_diff.py` (50 000 paired reads, +S. cerevisiae chrI plus the GFP transgene), timing indexing and alignment +separately because they are different questions: indexing is dominated by +suffix-array construction, alignment by the seed and stitch loops. + + python3 test/speed_bench.py --rustar ./target/release/rustar-aligner + +Prints a Markdown table, so the output can go straight into a job summary. No +thresholds and no non-zero exit: a wall time from a shared CI runner is a +data point, not a gate. Compare runs on the same machine. +""" + +from __future__ import annotations + +import argparse +import json +import resource +import subprocess +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from nfcore_diff import fetch # noqa: E402 - shared fixture download + + +def timed(cmd: list[str], log: Path) -> tuple[float, float]: + """Run `cmd`, returning (wall seconds, peak RSS in MB) for the child.""" + before = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss + start = time.monotonic() + with open(log, "w") as f: + proc = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT) + wall = time.monotonic() - start + after = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss + if proc.returncode != 0: + sys.exit(f"command failed ({proc.returncode}): {' '.join(cmd)}\nsee {log}") + # ru_maxrss is kilobytes on Linux, bytes on macOS. + peak = max(after, before) + peak_mb = peak / 1024 if sys.platform != "darwin" else peak / (1024 * 1024) + return wall, peak_mb + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--rustar", default="./target/release/rustar-aligner") + ap.add_argument("--star", default="STAR") + ap.add_argument("--work", default="/tmp/rustar-speed") + ap.add_argument("--threads", default="4") + ap.add_argument("--json") + args = ap.parse_args() + + work = Path(args.work) + fetch(work) + + results: dict[str, dict[str, float]] = {} + for tag, exe, run_mode in (("STAR", args.star, False), ("rustar", args.rustar, True)): + idx = work / f"{tag}_idx" + idx.mkdir(parents=True, exist_ok=True) + prefix = str(work / f"{tag}_") + + gen = [exe, "--runMode", "genomeGenerate", + "--genomeDir", str(idx), + "--genomeFastaFiles", str(work / "genome.fasta"), + "--genomeSAindexNbases", "9", + "--runThreadN", args.threads, + "--outFileNamePrefix", prefix + "idx_"] + index_wall, index_rss = timed(gen, work / f"{tag}_index.log") + + aln = [exe] + if run_mode: + aln += ["--runMode", "alignReads"] + aln += ["--genomeDir", str(idx), + "--readFilesIn", str(work / "reads_1.fastq"), str(work / "reads_2.fastq"), + "--outSAMtype", "SAM", + "--runThreadN", args.threads, + "--outFileNamePrefix", prefix] + align_wall, align_rss = timed(aln, work / f"{tag}_align.log") + + results[tag] = { + "index_seconds": index_wall, + "index_peak_mb": index_rss, + "align_seconds": align_wall, + "align_peak_mb": align_rss, + } + + print("### End-to-end timings\n") + print(f"Fixture: nf-core/rnaseq test data, 50 000 pairs, {args.threads} threads.\n") + print("| stage | STAR | rustar | ratio |") + print("|---|---|---|---|") + for stage, unit in (("index_seconds", "s"), ("align_seconds", "s")): + s, r = results["STAR"][stage], results["rustar"][stage] + ratio = r / s if s else float("inf") + print(f"| {stage.replace('_', ' ')} | {s:.1f}{unit} | {r:.1f}{unit} | {ratio:.2f}x |") + print() + print("Peak RSS is measured across the whole child process group, so it is") + print("only meaningful when the stages are run one at a time, as they are here.") + + if args.json: + Path(args.json).write_text(json.dumps(results, indent=2)) + return 0 + + +if __name__ == "__main__": + sys.exit(main())