From 1b1388086f2127c5cd7449a76cbc008af5b4a4f7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:52:39 +0000 Subject: [PATCH 1/4] Split pangenome alignment into one lane per NUMA node DNAscope already splits bwa into one job per NUMA node (halved until each piece has at most 64 CPUs) on large hosts. dnascope-pangenome ran bwa and minimap2 as single jobs across every node. With FASTQ input it now uses the same rule: - BwaExtractStage builds one bwa-extract job per lane: `fqidx extract -F j/N` feeds `bwa mem -p`, and bwa and `pgutil extract` are pinned to the lane's CPUs with taskset. Each lane writes sample-bwa-.bam and sample-bwa-.fq.gz. - One mm2-lift job per lane reads that lane's extracted reads, pinned to the same CPUs, and waits only for its own bwa lane. - Both dedup stages read every lane's output. Job names are unchanged, so job metrics still group the lanes as one step. - A hidden --no_split_alignment keeps the single jobs. BAM/CRAM input and hosts with 32 or fewer cores are unchanged. `fqidx extract` rejects -K below 10,000,000, while the pangenome bwa uses -K 2,000,000, so split runs use -K 10,000,000 for both. Unsplit runs keep 2,000,000. Co-Authored-By: Claude Opus 5.5 --- sentieon_cli/command_strings.py | 44 ++++++++- sentieon_cli/sentieon_pangenome.py | 75 ++++++++++----- sentieon_cli/stages/alignment.py | 76 ++++++++++----- tests/unit/test_pangenome_numa_lanes.py | 117 ++++++++++++++++++++++++ 4 files changed, 264 insertions(+), 48 deletions(-) create mode 100644 tests/unit/test_pangenome_numa_lanes.py diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index d1ad1f5..deb6517 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1378,8 +1378,15 @@ def cmd_bwa_extract( threads=1, unzip: str = "gzip", bwa_k: str = "2000000", + numa: Optional[str] = None, + split: Optional[str] = None, ) -> Pipeline: - """BWA alignment with pgutil extract""" + """BWA alignment with pgutil extract + + With ``split`` ("j/n") the job aligns one share of the read pairs, read + interleaved through `fqidx extract` with the same -K, so the batches and + the alignments match an unsplit run; ``numa`` pins it with `taskset`. + """ fq1_cmd = Pipeline( Command( unzip, @@ -1395,7 +1402,30 @@ def cmd_bwa_extract( ) ) + pin = ["taskset", "-c", numa] if numa else [] + if split: + reads = [ + InputProcSub( + Pipeline( + Command( + "sentieon", + "fqidx", + "extract", + "-F", + split, + "-K", + bwa_k, + InputProcSub(fq1_cmd), + InputProcSub(fq2_cmd), + ) + ) + ) + ] + else: + reads = [InputProcSub(fq1_cmd), InputProcSub(fq2_cmd)] + bwa_cmd = Command( + *pin, "sentieon", "bwa", "mem", @@ -1407,12 +1437,13 @@ def cmd_bwa_extract( str(bwa_model), "-K", bwa_k, + *(["-p"] if split else []), str(reference), - InputProcSub(fq1_cmd), - InputProcSub(fq2_cmd), + *reads, ) extract_cmd = Command( + *pin, "sentieon", "pgutil", "extract", @@ -1494,9 +1525,12 @@ def cmd_minimap2_lift( threads=1, mm2_xargs: List[str] = [], lift_prefix: str = "GRCh38#0#", + numa: Optional[str] = None, ) -> Pipeline: - """Minimap2 alignment with pgutil lift""" + """Minimap2 alignment with pgutil lift; ``numa`` pins it with `taskset`""" + pin = ["taskset", "-c", numa] if numa else [] mm2_cmd = Command( + *pin, "sentieon", "minimap2", "-t", @@ -1515,6 +1549,7 @@ def cmd_minimap2_lift( ) lift_cmd = Command( + *pin, "sentieon", "pgutil", "lift", @@ -1529,6 +1564,7 @@ def cmd_minimap2_lift( ) sort_cmd = Command( + *pin, "sentieon", "util", "sort", diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index dfef68b..83c48cb 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -26,6 +26,7 @@ from .logging import get_logger from .shell_pipeline import Command, Pipeline from .stages.alignment import ( + lane_path, BwaExtractStage, find_unzip, ) @@ -50,6 +51,7 @@ parse_rg_line, path_arg, require_versions, + split_alignment, total_memory, ) from .shard import ( @@ -131,6 +133,10 @@ class SentieonPangenome(BasePangenome): "help": argparse.SUPPRESS, "action": "store_true", }, + "no_split_alignment": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, } ) @@ -149,6 +155,7 @@ def __init__(self) -> None: self.skip_multiqc = False self.skip_pop_vcf_id_check: bool = False self.skip_model_apply = False + self.no_split_alignment = False self.skip_small_variants = False def _cnv_in_second_dag(self) -> bool: @@ -413,7 +420,10 @@ def collect_readgroups(self) -> None: def configure(self) -> None: """Configure pipeline parameters""" - pass + # bwa and minimap2 run as one lane per NUMA node (pieces of at most + # 64 CPUs, as in DNAscope) on large machines with fastq input + if self.r1_fastq and not self.no_split_alignment: + self.numa_nodes = split_alignment(self.cores) def build_dag(self) -> DAG: """Build the first DAG for the Sentieon pangenome pipeline""" @@ -468,17 +478,19 @@ def build_dag(self) -> DAG: dag.add_job(kmc_job) haplotype_dependencies.add(kmc_job) - # BWA alignment and extraction - bwa_result = self.bwa_extract_stage( - ctx, bwa_bam, ext_fastq - ).add_to(dag) - bwa_job = bwa_result.jobs[0] - mm2_dependencies.add(bwa_job) - bwa_lc_dependencies.add(bwa_job) + # BWA alignment and extraction, one lane per NUMA node + bwa_stage = self.bwa_extract_stage(ctx, bwa_bam, ext_fastq) + lanes = bwa_stage.lane_paths() + bwa_result = bwa_stage.add_to(dag) + bwa_jobs = list(bwa_result.jobs) + mm2_dependencies.update(bwa_jobs) + bwa_lc_dependencies.update(bwa_jobs) # Do not run vg-haplotypes with bwa in low-mem environments if total_mem_gb < 70: - haplotype_dependencies.add(bwa_job) + haplotype_dependencies.update(bwa_jobs) else: + lanes = [(bwa_bam, ext_fastq)] + bwa_jobs = [] dnascope_bams = copy.deepcopy(self.sample_input) # ReadWriter cannot write to /dev/stdout directly; pre-create a # symlink so the driver writes to a real path that resolves to @@ -532,14 +544,24 @@ def build_dag(self) -> DAG: # minimap2 alignment of the extracted fastq dnascope_dependencies = set() - mm2_job = self.build_minimap2_lift_job( - mm2_bam, - ext_fastq, - sample_fasta, - sample_gfa, - ) - dag.add_job(mm2_job, mm2_dependencies | {check_gfa_job, fasta_job}) - dnascope_dependencies.add(mm2_job) + n_lanes = len(lanes) + mm2_jobs: List[Job] = [] + mm2_bams: List[pathlib.Path] = [] + for j, (_, lane_fastq) in enumerate(lanes): + lane_bam = mm2_bam if n_lanes == 1 else lane_path(mm2_bam, j) + lane_job = self.build_minimap2_lift_job( + lane_bam, + lane_fastq, + sample_fasta, + sample_gfa, + lane=j if n_lanes > 1 else None, + n_lanes=n_lanes, + ) + lane_deps = {bwa_jobs[j]} if n_lanes > 1 else mm2_dependencies + dag.add_job(lane_job, lane_deps | {check_gfa_job, fasta_job}) + dnascope_dependencies.add(lane_job) + mm2_jobs.append(lane_job) + mm2_bams.append(lane_bam) # With fastq input, perform dedup and metrics sr_alignments: List[pathlib.Path] = [] @@ -559,7 +581,7 @@ def build_dag(self) -> DAG: bwa_dedup = DedupStage( ctx=ctx, tag="bwa", - inputs=[bwa_bam], + inputs=[bam for bam, _ in lanes], output=out_bwa_aln, score_file=self.tmp_dir.joinpath("sample-bwa-score.txt.gz"), dedup_metrics=dedup_metrics, @@ -569,14 +591,14 @@ def build_dag(self) -> DAG: mm2_dedup = DedupStage( ctx=ctx, tag="mm2", - inputs=[mm2_bam], + inputs=mm2_bams, output=out_mm2_aln, score_file=self.tmp_dir.joinpath("sample-mm2-score.txt.gz"), read_filters=[ "IndelLeftAlignReadTransform,rgid=" f"{self.fastq_readgroup['ID']}-mm2" ], - ).add_to(dag, {mm2_job}) + ).add_to(dag, set(mm2_jobs)) dnascope_dependencies.add(mm2_dedup.dedup_job) sr_alignments = [out_bwa_aln] @@ -743,6 +765,7 @@ def bwa_extract_stage( extract_model=bundle.joinpath(self.extract_model_name), bwa_model=bundle.joinpath("bwa.model"), unzip=find_unzip(self.logger), + numa_nodes=self.numa_nodes, ) def build_haplotypes_job( @@ -810,8 +833,12 @@ def build_minimap2_lift_job( ext_fastq: pathlib.Path, sample_fasta: pathlib.Path, sample_gfa: pathlib.Path, + lane: Optional[int] = None, + n_lanes: int = 1, ) -> Job: - """Build minimap2 alignment with pgutil lift job""" + """Build minimap2 alignment with pgutil lift job; with ``lane`` it is + that NUMA lane's share, pinned to its node""" + threads = max(1, int(self.cores / n_lanes)) bundle = self.required(self.model_bundle, "model_bundle") reference = self.required(self.reference, "reference") @@ -834,11 +861,13 @@ def build_minimap2_lift_job( reference, "@RG\\t" + "\\t".join([f"{x[0]}:{x[1]}" for x in rg2.items()]), mm2_model, - threads=self.cores, + threads=threads, lift_prefix=self.contig_prefix(), + numa=self.numa_nodes[lane] if lane is not None else None, ), "mm2-lift", - self.cores, + threads, + resources={f"node{lane}": 1} if lane is not None else None, task_name="pangenome-alignment", ) return mm2_job diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py index 32736df..3de0832 100644 --- a/sentieon_cli/stages/alignment.py +++ b/sentieon_cli/stages/alignment.py @@ -13,7 +13,7 @@ import logging import pathlib import shutil -from typing import Iterable, List, Optional, Union +from typing import Iterable, List, Optional, Tuple, Union import packaging.version @@ -377,6 +377,18 @@ def build(self) -> AlignResult: ) +# `sentieon fqidx extract` rejects -K below 10,000,000 (the pangenome bwa +# default is 2,000,000), so a split run aligns with this -K in both, keeping +# each lane's batches the ones an unsplit run at the same -K would form +SPLIT_BWA_K = "10000000" + + +def lane_path(path: pathlib.Path, lane: int) -> pathlib.Path: + """`sample-bwa.fq.gz` -> `sample-bwa-.fq.gz`""" + stem, _, ext = path.name.partition(".") + return path.with_name(f"{stem}-{lane}.{ext}" if ext else f"{stem}-{lane}") + + @dataclass(kw_only=True) class BwaExtractStage(AlignmentStage): """Align short-read fastq with bwa and extract the reads a pangenome @@ -396,29 +408,51 @@ class BwaExtractStage(AlignmentStage): unzip: str = "gzip" name: str = "bwa-extract" task_name: str = "alignment" + numa_nodes: List[str] = field(default_factory=list) + + def lane_paths(self) -> List[Tuple[pathlib.Path, pathlib.Path]]: + """(bam, fastq) per lane: one lane per NUMA node when there are + several, else the single output pair""" + if len(self.numa_nodes) <= 1: + return [(self.output_bam, self.output_fastq)] + return [ + (lane_path(self.output_bam, j), lane_path(self.output_fastq, j)) + for j in range(len(self.numa_nodes)) + ] def build(self) -> AlignResult: - job = Job( - cmds.cmd_bwa_extract( - self.output_bam, - self.output_fastq, - self.ctx.reference, - self.r1_fastq, - self.r2_fastq, - self.readgroup, - self.extract_model, - self.bwa_model, - self.ctx.cores, - unzip=self.unzip, - ), - self.name, - self.ctx.cores, - task_name=self.task_name, - ) + lanes = self.lane_paths() + n = len(lanes) + cores = max(1, int(self.ctx.cores / n)) + jobs: List[Job] = [] + for j, (bam, fastq) in enumerate(lanes): + jobs.append( + Job( + cmds.cmd_bwa_extract( + bam, + fastq, + self.ctx.reference, + self.r1_fastq, + self.r2_fastq, + self.readgroup, + self.extract_model, + self.bwa_model, + cores, + unzip=self.unzip, + **({"bwa_k": SPLIT_BWA_K} if n > 1 else {}), + numa=self.numa_nodes[j] if n > 1 else None, + split=f"{j}/{n}" if n > 1 else None, + ), + self.name, + cores, + resources={f"node{j}": 1} if n > 1 else None, + task_name=self.task_name, + ) + ) return AlignResult( - jobs=[job], - terminal={job}, - outputs=[self.output_bam, self.output_fastq], + jobs=jobs, + terminal=set(jobs), + outputs=[x for lane in lanes for x in lane], cleanup_paths=[], ) diff --git a/tests/unit/test_pangenome_numa_lanes.py b/tests/unit/test_pangenome_numa_lanes.py new file mode 100644 index 0000000..b88831e --- /dev/null +++ b/tests/unit/test_pangenome_numa_lanes.py @@ -0,0 +1,117 @@ +"""bwa-extract and mm2-lift run as one lane per NUMA node""" + +import os +import sys +from unittest.mock import patch + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +import test_sentieon_pangenome # noqa: E402 + +TWO_NODES = ["0-95", "96-191"] +LANES = ["0-47", "48-95", "96-143", "144-191"] + + +class TestNumaLanes: + # Use the pangenome tests' fixtures without inheriting, and so + # re-running, their tests + def setup_method(self): + self.base = test_sentieon_pangenome.TestSentieonPangenome() + self.base.setup_method() + + def lane_pipeline(self, cores=192, nodes=TWO_NODES, no_split=False): + pipeline = self.base.create_fastq_pipeline() + pipeline.cores = cores + pipeline.no_split_alignment = no_split + with patch("sentieon_cli.util.find_numa_nodes", return_value=nodes): + pipeline.configure() + return pipeline + + def jobs(self, dag, name): + everything = list(dag.waiting_jobs) + list(dag.ready_jobs) + return [j for j in everything if j.name == name] + + def test_configure_splits_like_dnascope(self): + assert self.lane_pipeline().numa_nodes == LANES + + def test_bwa_extract_lanes(self): + dag = self.lane_pipeline().build_dag() + bwa = self.jobs(dag, "bwa-extract") + assert len(bwa) == 4 + for j, cpus in enumerate(LANES): + job = next(b for b in bwa if f"taskset -c {cpus} " in str(b.shell)) + cmd = str(job.shell) + assert job.threads == 48 + assert job.resources == {f"node{j}": 1} + assert f"fqidx extract -F {j}/4 -K 10000000" in cmd + assert "bwa mem" in cmd and cmd.count("-K 10000000") == 2 + assert " -p " in cmd and "-t 48" in cmd + assert ( + f"sample-bwa-{j}.bam" in cmd and f"sample-bwa-{j}.fq.gz" in cmd + ) + # pgutil extract is pinned too + assert cmd.count(f"taskset -c {cpus}") == 2 + + def test_mm2_lanes_follow_their_bwa_lane(self): + pipeline = self.lane_pipeline() + dag = pipeline.build_dag() + mm2 = self.jobs(dag, "mm2-lift") + assert len(mm2) == 4 + for j, cpus in enumerate(LANES): + job = next( + m for m in mm2 if f"sample-bwa-{j}.fq.gz" in str(m.shell) + ) + cmd = str(job.shell) + assert job.threads == 48 and job.resources == {f"node{j}": 1} + assert cmd.count(f"taskset -c {cpus}") == 3 # minimap2, lift, sort + assert f"sample-mm2-{j}.bam" in cmd + deps = self.base._get_deps(dag, job) + lane_bwa = [d for d in deps if d.name == "bwa-extract"] + assert len(lane_bwa) == 1 + assert f"sample-bwa-{j}.fq.gz" in str(lane_bwa[0].shell) + assert {"check-sample-gfa", "vg-paths-fasta"} <= { + d.name for d in deps + } + + def test_dedups_read_every_lane(self): + dag = self.lane_pipeline().build_dag() + everything = list(dag.waiting_jobs) + list(dag.ready_jobs) + lc = [ + str(j.shell) + for j in everything + if "LocusCollector" in str(j.shell) + ] + bwa_lc = next(c for c in lc if "sample-bwa-0.bam" in c) + mm2_lc = next(c for c in lc if "sample-mm2-0.bam" in c) + for j in range(4): + assert f"sample-bwa-{j}.bam" in bwa_lc + assert f"sample-mm2-{j}.bam" in mm2_lc + + def test_no_split_alignment_keeps_one_job(self): + dag = self.lane_pipeline(no_split=True).build_dag() + bwa = self.jobs(dag, "bwa-extract") + mm2 = self.jobs(dag, "mm2-lift") + assert len(bwa) == 1 and len(mm2) == 1 + assert "taskset" not in str(bwa[0].shell) + str(mm2[0].shell) + assert "fqidx" not in str(bwa[0].shell) + assert "sample-bwa.bam" in str(bwa[0].shell) + assert "-K 2000000" in str(bwa[0].shell) # unsplit keeps the default + + def test_small_machine_keeps_one_job(self): + dag = self.lane_pipeline(cores=48, nodes=["0-47"]).build_dag() + assert len(self.jobs(dag, "bwa-extract")) == 1 + assert len(self.jobs(dag, "mm2-lift")) == 1 + + def test_bam_input_is_unchanged(self): + pipeline = self.base.create_pipeline() + pipeline.cores = 192 + with patch( + "sentieon_cli.util.find_numa_nodes", return_value=TWO_NODES + ): + pipeline.configure() + assert pipeline.numa_nodes == [] + dag = pipeline.build_dag() + mm2 = self.jobs(dag, "mm2-lift") + assert len(mm2) == 1 and "taskset" not in str(mm2[0].shell) From 0e50dbc209ef6dbf71e99bfa9224b0fe1efad2da Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:53:48 +0000 Subject: [PATCH 2/4] Hand indel2cnv's variants to workers one at a time indel2cnv sorts the candidate variants by length, longest first, so the expensive ones are dispatched early. It then grouped them into batches of 5 and called starmap with its default chunksize, which is about len(batches) / (4 * workers). The first chunk therefore held the longest variants of the whole run and went to a single worker, which finished long after the others went idle. Each task is now one variant, with chunksize=1. Results are collected in the same order as before, so the output does not change. On HG002 35x (HPRC v2.0), indel2cnv took 420 -> 68 s on 64 vCPUs and 156 -> 49 s on 192 vCPUs, and the CNV VCF body was identical in both cases. Co-Authored-By: Claude Opus 5.5 --- sentieon_cli/scripts/indel2cnv.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/sentieon_cli/scripts/indel2cnv.py b/sentieon_cli/scripts/indel2cnv.py index bc79ef2..e3801e1 100755 --- a/sentieon_cli/scripts/indel2cnv.py +++ b/sentieon_cli/scripts/indel2cnv.py @@ -1160,15 +1160,16 @@ def main(args): n_threads = args.threads or os.cpu_count() # Sort by expected difficulty (sequence length) descending so that # expensive variants are dispatched first and don't pile up at the end. - # Use small batches (5) for better load balancing across workers. + # Hand out one variant at a time: with larger batches, or starmap's + # default chunksize, the most expensive variants all land on one worker. variants.sort(key=lambda v: max(len(a) for a in [v.ref] + v.alt), reverse=True) - batch_size = 5 - args_in = [] - for i in range(0, len(variants), batch_size): - args_in.append((ref, variants[i:i+batch_size], args.thresh, args.extend_slack, args.match_thresh)) + args_in = [ + (ref, [v], args.thresh, args.extend_slack, args.match_thresh) + for v in variants + ] all_vars = [] with Pool(n_threads) as p: - for result in p.starmap(proc_batch, args_in): + for result in p.starmap(proc_batch, args_in, chunksize=1): for vs, t in result: all_vars += vs # Merge overlapping output calls (same array, different haplotypes) and write From 3812243f581e21d75c2be9b846a2eacc19eb9bd5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:54:34 +0000 Subject: [PATCH 3/4] Run vg haplotypes alongside bwa when both fit in memory vg-haplotypes only needs the k-mer counts, so it is meant to overlap bwa on hosts with enough memory. It did not: the job reserved every core, so the scheduler held it until bwa released its threads, whatever its dependencies said. When the host has room for vg haplotypes (100 GB) plus each bwa lane (32 GB), the job now reserves no threads, as kmc already does, and runs next to bwa. Otherwise it depends on every bwa lane, as before. This replaces the fixed 70 GB check. On HG002 35x with 4 lanes on a 192-vCPU, 384 GB host, the pipeline took 30:10 -> 26:24, with indel2cnv's change included. The output VCF body was unchanged. Co-Authored-By: Claude Opus 5.5 --- sentieon_cli/sentieon_pangenome.py | 32 +++++++++-- tests/unit/test_pangenome_vg_scheduling.py | 64 ++++++++++++++++++++++ 2 files changed, 91 insertions(+), 5 deletions(-) create mode 100644 tests/unit/test_pangenome_vg_scheduling.py diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 83c48cb..8ec80af 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -59,6 +59,13 @@ parse_fai, ) +# Peak memory, in GB, of vg haplotypes on the HPRC v2.0 and v2.1 graphs +# (about 60 with vg 1.78.0, up to 92 with a build that samples in parallel) +# and of one bwa-extract lane (about 20), measured on HG002 and HG005. +# vg-haplotypes runs alongside bwa only if the two fit together. +VG_HAPLOTYPES_MEM_GB = 100 +BWA_LANE_MEM_GB = 32 + SENT_PANGENOME_MIN_VERSIONS = { "kmc": None, "sentieon driver": packaging.version.Version("202503.02"), @@ -485,12 +492,17 @@ def build_dag(self) -> DAG: bwa_jobs = list(bwa_result.jobs) mm2_dependencies.update(bwa_jobs) bwa_lc_dependencies.update(bwa_jobs) - # Do not run vg-haplotypes with bwa in low-mem environments - if total_mem_gb < 70: + # Run vg-haplotypes alongside bwa when both fit in memory + vg_with_bwa = ( + total_mem_gb + >= VG_HAPLOTYPES_MEM_GB + BWA_LANE_MEM_GB * len(bwa_jobs) + ) + if not vg_with_bwa: haplotype_dependencies.update(bwa_jobs) else: lanes = [(bwa_bam, ext_fastq)] bwa_jobs = [] + vg_with_bwa = False dnascope_bams = copy.deepcopy(self.sample_input) # ReadWriter cannot write to /dev/stdout directly; pre-create a # symlink so the driver writes to a real path that resolves to @@ -523,7 +535,14 @@ def build_dag(self) -> DAG: haplotype_dependencies.add(extract_kmc_job) # vg haplotypes - create a sample-specific pangenome - haplotypes_job = self.build_haplotypes_job(sample_pangenome, kmer_file) + # Alongside bwa, vg-haplotypes reserves no threads, as kmc does. + # Otherwise the scheduler holds it until bwa releases its threads, + # whatever its dependencies say. + haplotypes_job = self.build_haplotypes_job( + sample_pangenome, + kmer_file, + job_threads=0 if vg_with_bwa else self.cores, + ) dag.add_job(haplotypes_job, haplotype_dependencies) # Confirm the sampled pangenome kept its reference paths before @@ -769,7 +788,10 @@ def bwa_extract_stage( ) def build_haplotypes_job( - self, output_gbz: pathlib.Path, kmer_file: pathlib.Path + self, + output_gbz: pathlib.Path, + kmer_file: pathlib.Path, + job_threads: Optional[int] = None, ) -> Job: """Build vg haplotypes job""" hapl_file = self.required(self.hapl, "hapl") @@ -790,7 +812,7 @@ def build_haplotypes_job( ], ), "vg-haplotypes", - self.cores, + self.cores if job_threads is None else job_threads, task_name="pangenome", ) return haplotypes_job diff --git a/tests/unit/test_pangenome_vg_scheduling.py b/tests/unit/test_pangenome_vg_scheduling.py new file mode 100644 index 0000000..9342129 --- /dev/null +++ b/tests/unit/test_pangenome_vg_scheduling.py @@ -0,0 +1,64 @@ +"""vg-haplotypes runs alongside bwa-extract when both fit in memory""" + +import os +import sys +from unittest.mock import patch + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +import test_sentieon_pangenome # noqa: E402 + +GB = 1024**3 +TWO_NODES = ["0-95", "96-191"] + + +class TestVgHaplotypesScheduling: + # Use the pangenome tests' fixtures without inheriting, and so + # re-running, their tests + def setup_method(self): + self.base = test_sentieon_pangenome.TestSentieonPangenome() + self.base.setup_method() + + def build(self, memory_gb, cores=64, nodes=None): + pipeline = self.base.create_fastq_pipeline() + pipeline.cores = cores + with patch( + "sentieon_cli.util.find_numa_nodes", return_value=nodes or [] + ): + pipeline.configure() + with patch( + "sentieon_cli.sentieon_pangenome.total_memory", + return_value=memory_gb * GB, + ): + dag = pipeline.build_dag() + everything = list(dag.waiting_jobs) + list(dag.ready_jobs) + vg = next(j for j in everything if j.name == "vg-haplotypes") + return pipeline, dag, vg + + def test_runs_alongside_bwa_with_enough_memory(self): + """No bwa dependency, and no threads reserved from the scheduler""" + _, dag, vg = self.build(384, cores=192, nodes=TWO_NODES) + deps = {d.name for d in self.base._get_deps(dag, vg)} + assert "bwa-extract" not in deps + assert "kmc" in deps + assert vg.threads == 0 + # vg itself still uses every core + assert "-t 192" in str(vg.shell) + + def test_waits_for_bwa_when_memory_is_short(self): + """A 128 GB host has no room for vg next to a bwa lane""" + pipeline, dag, vg = self.build(128) + deps = {d.name for d in self.base._get_deps(dag, vg)} + assert "bwa-extract" in deps + assert vg.threads == pipeline.cores + + def test_waits_for_every_bwa_lane(self): + """The memory check counts each lane, and vg waits for all of them""" + _, dag, vg = self.build(200, cores=192, nodes=TWO_NODES) + bwa = [ + d for d in self.base._get_deps(dag, vg) if d.name == "bwa-extract" + ] + assert len(bwa) == 4 + assert vg.threads == 192 From d64ae37414d5c7e8e302814d201df328e7644af9 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 18:28:01 +0000 Subject: [PATCH 4/4] Name each alignment lane's job, and use one -K for every pangenome run Review follow-up: - Job names are unique again: split runs name their lanes bwa-extract- and mm2-lift-, while unsplit runs keep bwa-extract and mm2-lift. - dnascope-pangenome aligns with -K 10,000,000 whether or not it splits, so its results no longer depend on how many NUMA lanes the host has. `fqidx extract` needs at least that, which is why split runs already used it. BwaExtractStage takes the -K as an option, so hybrid-pangenome keeps its current 2,000,000. Co-Authored-By: Claude Opus 5.5 --- sentieon_cli/sentieon_pangenome.py | 4 +++- sentieon_cli/stages/alignment.py | 16 ++++++++----- tests/unit/test_pangenome_numa_lanes.py | 28 ++++++++++++++++++---- tests/unit/test_pangenome_vg_scheduling.py | 10 +++++--- tests/unit/test_stage_alignment.py | 20 ++++++++++++++++ 5 files changed, 64 insertions(+), 14 deletions(-) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 8ec80af..b41c547 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -28,6 +28,7 @@ from .stages.alignment import ( lane_path, BwaExtractStage, + PANGENOME_BWA_K, find_unzip, ) from .stages.base import StageContext @@ -785,6 +786,7 @@ def bwa_extract_stage( bwa_model=bundle.joinpath("bwa.model"), unzip=find_unzip(self.logger), numa_nodes=self.numa_nodes, + bwa_k=PANGENOME_BWA_K, ) def build_haplotypes_job( @@ -887,7 +889,7 @@ def build_minimap2_lift_job( lift_prefix=self.contig_prefix(), numa=self.numa_nodes[lane] if lane is not None else None, ), - "mm2-lift", + "mm2-lift" if lane is None else f"mm2-lift-{lane}", threads, resources={f"node{lane}": 1} if lane is not None else None, task_name="pangenome-alignment", diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py index 3de0832..4656bbd 100644 --- a/sentieon_cli/stages/alignment.py +++ b/sentieon_cli/stages/alignment.py @@ -377,10 +377,11 @@ def build(self) -> AlignResult: ) -# `sentieon fqidx extract` rejects -K below 10,000,000 (the pangenome bwa -# default is 2,000,000), so a split run aligns with this -K in both, keeping -# each lane's batches the ones an unsplit run at the same -K would form -SPLIT_BWA_K = "10000000" +# `sentieon fqidx extract` rejects -K below 10,000,000, so a split run needs +# at least this. dnascope-pangenome uses it whether or not it splits: each +# lane then forms the batches an unsplit run would, and results match +# across machines +PANGENOME_BWA_K = "10000000" def lane_path(path: pathlib.Path, lane: int) -> pathlib.Path: @@ -409,6 +410,8 @@ class BwaExtractStage(AlignmentStage): name: str = "bwa-extract" task_name: str = "alignment" numa_nodes: List[str] = field(default_factory=list) + # None keeps cmd_bwa_extract's default, or PANGENOME_BWA_K when split + bwa_k: Optional[str] = None def lane_paths(self) -> List[Tuple[pathlib.Path, pathlib.Path]]: """(bam, fastq) per lane: one lane per NUMA node when there are @@ -424,6 +427,7 @@ def build(self) -> AlignResult: lanes = self.lane_paths() n = len(lanes) cores = max(1, int(self.ctx.cores / n)) + bwa_k = self.bwa_k or (PANGENOME_BWA_K if n > 1 else None) jobs: List[Job] = [] for j, (bam, fastq) in enumerate(lanes): jobs.append( @@ -439,11 +443,11 @@ def build(self) -> AlignResult: self.bwa_model, cores, unzip=self.unzip, - **({"bwa_k": SPLIT_BWA_K} if n > 1 else {}), + **({"bwa_k": bwa_k} if bwa_k else {}), numa=self.numa_nodes[j] if n > 1 else None, split=f"{j}/{n}" if n > 1 else None, ), - self.name, + f"{self.name}-{j}" if n > 1 else self.name, cores, resources={f"node{j}": 1} if n > 1 else None, task_name=self.task_name, diff --git a/tests/unit/test_pangenome_numa_lanes.py b/tests/unit/test_pangenome_numa_lanes.py index b88831e..4a3fb72 100644 --- a/tests/unit/test_pangenome_numa_lanes.py +++ b/tests/unit/test_pangenome_numa_lanes.py @@ -30,8 +30,13 @@ def lane_pipeline(self, cores=192, nodes=TWO_NODES, no_split=False): return pipeline def jobs(self, dag, name): + """The step's jobs: `name`, or `name-` when split""" everything = list(dag.waiting_jobs) + list(dag.ready_jobs) - return [j for j in everything if j.name == name] + return [ + j + for j in everything + if j.name == name or j.name.rpartition("-")[0] == name + ] def test_configure_splits_like_dnascope(self): assert self.lane_pipeline().numa_nodes == LANES @@ -43,6 +48,7 @@ def test_bwa_extract_lanes(self): for j, cpus in enumerate(LANES): job = next(b for b in bwa if f"taskset -c {cpus} " in str(b.shell)) cmd = str(job.shell) + assert job.name == f"bwa-extract-{j}" assert job.threads == 48 assert job.resources == {f"node{j}": 1} assert f"fqidx extract -F {j}/4 -K 10000000" in cmd @@ -64,12 +70,13 @@ def test_mm2_lanes_follow_their_bwa_lane(self): m for m in mm2 if f"sample-bwa-{j}.fq.gz" in str(m.shell) ) cmd = str(job.shell) + assert job.name == f"mm2-lift-{j}" assert job.threads == 48 and job.resources == {f"node{j}": 1} assert cmd.count(f"taskset -c {cpus}") == 3 # minimap2, lift, sort assert f"sample-mm2-{j}.bam" in cmd deps = self.base._get_deps(dag, job) - lane_bwa = [d for d in deps if d.name == "bwa-extract"] - assert len(lane_bwa) == 1 + lane_bwa = [d for d in deps if d.name.startswith("bwa-extract")] + assert [d.name for d in lane_bwa] == [f"bwa-extract-{j}"] assert f"sample-bwa-{j}.fq.gz" in str(lane_bwa[0].shell) assert {"check-sample-gfa", "vg-paths-fasta"} <= { d.name for d in deps @@ -97,13 +104,26 @@ def test_no_split_alignment_keeps_one_job(self): assert "taskset" not in str(bwa[0].shell) + str(mm2[0].shell) assert "fqidx" not in str(bwa[0].shell) assert "sample-bwa.bam" in str(bwa[0].shell) - assert "-K 2000000" in str(bwa[0].shell) # unsplit keeps the default + assert bwa[0].name == "bwa-extract" and mm2[0].name == "mm2-lift" + # the same -K as a split run, so the alignments match + assert "-K 10000000" in str(bwa[0].shell) + assert "-K 2000000" not in str(bwa[0].shell) def test_small_machine_keeps_one_job(self): dag = self.lane_pipeline(cores=48, nodes=["0-47"]).build_dag() assert len(self.jobs(dag, "bwa-extract")) == 1 assert len(self.jobs(dag, "mm2-lift")) == 1 + def test_job_names_are_unique(self): + dag = self.lane_pipeline().build_dag() + everything = list(dag.waiting_jobs) + list(dag.ready_jobs) + names = [ + j.name + for j in everything + if j.name.startswith(("bwa-extract", "mm2-lift")) + ] + assert len(names) == 8 and len(set(names)) == 8 + def test_bam_input_is_unchanged(self): pipeline = self.base.create_pipeline() pipeline.cores = 192 diff --git a/tests/unit/test_pangenome_vg_scheduling.py b/tests/unit/test_pangenome_vg_scheduling.py index 9342129..d09c2ce 100644 --- a/tests/unit/test_pangenome_vg_scheduling.py +++ b/tests/unit/test_pangenome_vg_scheduling.py @@ -41,7 +41,7 @@ def test_runs_alongside_bwa_with_enough_memory(self): """No bwa dependency, and no threads reserved from the scheduler""" _, dag, vg = self.build(384, cores=192, nodes=TWO_NODES) deps = {d.name for d in self.base._get_deps(dag, vg)} - assert "bwa-extract" not in deps + assert not any(d.startswith("bwa-extract") for d in deps) assert "kmc" in deps assert vg.threads == 0 # vg itself still uses every core @@ -58,7 +58,11 @@ def test_waits_for_every_bwa_lane(self): """The memory check counts each lane, and vg waits for all of them""" _, dag, vg = self.build(200, cores=192, nodes=TWO_NODES) bwa = [ - d for d in self.base._get_deps(dag, vg) if d.name == "bwa-extract" + d + for d in self.base._get_deps(dag, vg) + if d.name.startswith("bwa-extract") + ] + assert sorted(d.name for d in bwa) == [ + f"bwa-extract-{j}" for j in range(4) ] - assert len(bwa) == 4 assert vg.threads == 192 diff --git a/tests/unit/test_stage_alignment.py b/tests/unit/test_stage_alignment.py index 6e4829e..787c118 100644 --- a/tests/unit/test_stage_alignment.py +++ b/tests/unit/test_stage_alignment.py @@ -396,6 +396,26 @@ def test_unzip_reaches_the_command(self, tmp_path): assert "igzip -dc" in str(result.jobs[0].shell) + def test_bwa_k_default_and_override(self, tmp_path): + default = str(self.make(tmp_path).build().jobs[0].shell) + override = str( + self.make(tmp_path, bwa_k="10000000").build().jobs[0].shell + ) + + assert "-K 2000000" in default + assert "-K 10000000" in override and "-K 2000000" not in override + + def test_split_lanes_have_unique_names(self, tmp_path): + result = self.make(tmp_path, numa_nodes=["0-47", "48-95"]).build() + + assert [j.name for j in result.jobs] == [ + "bwa-extract-0", + "bwa-extract-1", + ] + # fqidx extract needs -K of at least 10,000,000 + for job in result.jobs: + assert str(job.shell).count("-K 10000000") == 2 + class TestReadWriterStage: """`sentieon driver --algo ReadWriter` merges and extracts"""