From b2b57c25c4238cde1731d6e5c6f365c6978a490f Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:51:32 +0000 Subject: [PATCH 1/4] Cap KMC at 128 threads KMC rejects more than 128 threads ("number of threads for 2nd stage must be in range <1,128>"), so dnascope-pangenome failed at k-mer counting on hosts with more than 128 logical CPUs, such as 192-vCPU instances. The cap is applied in the three KMC command builders, which covers the pangenome and hybrid-pangenome pipelines for FASTQ and aligned input. The other commands in those pipelines keep every core. Co-Authored-By: Claude Opus 5.5 --- sentieon_cli/command_strings.py | 15 ++++++-- tests/unit/test_kmc_threads.py | 61 +++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_kmc_threads.py diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index c208463..25eb454 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -23,6 +23,15 @@ # fs.pipe-max-size and non-Linux platforms keep their default size. ALN_PIPE_SIZE = 268435456 +# KMC rejects more threads than this: "number of threads for 2nd stage +# must be in range <1,128>" +KMC_MAX_THREADS = 128 + + +def kmc_threads(threads: int) -> int: + """The thread count to pass to KMC""" + return min(threads, KMC_MAX_THREADS) + def cmd_fai_to_bed( fai: pathlib.Path, @@ -1000,7 +1009,7 @@ def cmd_kmc( f"-k{k}", f"-m{memory}", "-okff", - f"-t{threads}", + f"-t{kmc_threads(threads)}", "-hp", f"@{file_list}", str(output_prefix), @@ -1084,7 +1093,7 @@ def cmd_extract_kmc( "-k29", "-m32", "-okff", - f"-t{threads}", + f"-t{kmc_threads(threads)}", "-fa", "/dev/stdin", str(output_prefix), @@ -1480,7 +1489,7 @@ def _kmc_stdin_cmd( f"-k{k}", f"-m{memory}", "-okff", - f"-t{threads}", + f"-t{kmc_threads(threads)}", "-fa", "/dev/stdin", str(output_prefix), diff --git a/tests/unit/test_kmc_threads.py b/tests/unit/test_kmc_threads.py new file mode 100644 index 0000000..d6ea1b7 --- /dev/null +++ b/tests/unit/test_kmc_threads.py @@ -0,0 +1,61 @@ +""" +Unit tests for the KMC thread cap +""" + +import pathlib +import re + +from sentieon_cli.command_strings import ( + KMC_MAX_THREADS, + cmd_extract_kmc, + cmd_hybrid_kmc, + cmd_kmc, +) + +TMP = pathlib.Path("/tmp/kmc") +PREFIX = pathlib.Path("/tmp/sample") + + +def kmc_t(cmd_str: str) -> str: + """The -t value of the kmc command in a pipeline""" + match = re.search(r"\bkmc -k\S+ -m\S+ -okff -t(\d+)", cmd_str) + assert match, cmd_str + return match.group(1) + + +def test_kmc_threads_capped(): + """More cores than KMC accepts are capped at its limit""" + cmd = cmd_kmc(PREFIX, pathlib.Path("/tmp/fq.list"), TMP, threads=192) + assert kmc_t(str(cmd)) == str(KMC_MAX_THREADS) + + +def test_kmc_threads_below_cap_unchanged(): + """A thread count KMC accepts is passed through""" + cmd = cmd_kmc(PREFIX, pathlib.Path("/tmp/fq.list"), TMP, threads=64) + assert kmc_t(str(cmd)) == "64" + + +def test_extract_kmc_caps_only_kmc(): + """The driver and pgutil keep every core; only KMC is capped""" + cmd_str = str( + cmd_extract_kmc( + PREFIX, + pathlib.Path("/tmp/ext.fq.gz"), + [pathlib.Path("/data/sample.bam")], + pathlib.Path("/ref/hg38.fa"), + pathlib.Path("/bundle/extract.model"), + TMP, + pathlib.Path("/tmp/rw.bam"), + threads=192, + ) + ) + assert kmc_t(cmd_str) == str(KMC_MAX_THREADS) + assert "-t 192" in cmd_str + + +def test_hybrid_kmc_threads_capped(): + """The stdin KMC used by the hybrid pipelines is capped too""" + cmd = cmd_hybrid_kmc( + PREFIX, [pathlib.Path("/data/r1.fq.gz")], [], TMP, threads=192 + ) + assert kmc_t(str(cmd)) == str(KMC_MAX_THREADS) From 704711ae3257c6563c1f718fffee43ab300d79d0 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:52:39 +0000 Subject: [PATCH 2/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 | 111 ++++++++++++++++++++++++ 4 files changed, 258 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 25eb454..6adeb5c 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1283,8 +1283,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, @@ -1300,7 +1307,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", @@ -1312,12 +1342,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", @@ -1399,9 +1430,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", @@ -1420,6 +1454,7 @@ def cmd_minimap2_lift( ) lift_cmd = Command( + *pin, "sentieon", "pgutil", "lift", @@ -1434,6 +1469,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 ed445e8..78fb8cd 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: @@ -408,7 +415,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""" @@ -463,17 +473,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 @@ -527,14 +539,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] = [] @@ -554,7 +576,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, @@ -564,14 +586,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] @@ -737,6 +759,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( @@ -804,8 +827,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") @@ -828,11 +855,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..61cdcea --- /dev/null +++ b/tests/unit/test_pangenome_numa_lanes.py @@ -0,0 +1,111 @@ +"""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__), "..", "..")) +) + +from test_sentieon_pangenome import TestSentieonPangenome # noqa: E402 + +TWO_NODES = ["0-95", "96-191"] +LANES = ["0-47", "48-95", "96-143", "144-191"] + + +class TestNumaLanes(TestSentieonPangenome): + def lane_pipeline(self, cores=192, nodes=TWO_NODES, no_split=False): + pipeline = self.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._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.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 320f33744b5750b274a66f3f094f954714a6865a Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:53:48 +0000 Subject: [PATCH 3/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 | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/sentieon_cli/scripts/indel2cnv.py b/sentieon_cli/scripts/indel2cnv.py index 68713f9..6f112b6 100755 --- a/sentieon_cli/scripts/indel2cnv.py +++ b/sentieon_cli/scripts/indel2cnv.py @@ -709,15 +709,13 @@ 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_in = [(ref, [v], args.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 bf2dfbc96a4330f123d9a504c1c269f8505c359f Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 19:54:34 +0000 Subject: [PATCH 4/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 | 56 ++++++++++++++++++++++ 2 files changed, 83 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 78fb8cd..8cf6c32 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"), @@ -480,12 +487,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 @@ -518,7 +530,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 @@ -763,7 +782,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") @@ -784,7 +806,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..558e468 --- /dev/null +++ b/tests/unit/test_pangenome_vg_scheduling.py @@ -0,0 +1,56 @@ +"""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__), "..", "..")) +) + +from test_sentieon_pangenome import TestSentieonPangenome # noqa: E402 + +GB = 1024**3 +TWO_NODES = ["0-95", "96-191"] + + +class TestVgHaplotypesScheduling(TestSentieonPangenome): + def build(self, memory_gb, cores=64, nodes=None): + pipeline = self.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._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._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._get_deps(dag, vg) if d.name == "bwa-extract"] + assert len(bwa) == 4 + assert vg.threads == 192