dnascope-pangenome: KMC thread cap, NUMA alignment lanes, faster indel2cnv, vg haplotypes alongside bwa

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),
@@ -1274,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,
@@ -1291,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",
@@ -1303,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",
@@ -1390,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",
@@ -1411,6 +1454,7 @@ def cmd_minimap2_lift(
     )
 
     lift_cmd = Command(
+        *pin,
         "sentieon",
         "pgutil",
         "lift",
@@ -1425,6 +1469,7 @@ def cmd_minimap2_lift(
     )
 
     sort_cmd = Command(
+        *pin,
         "sentieon",
         "util",
         "sort",
@@ -1480,7 +1525,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),

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

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,13 +51,21 @@
     parse_rg_line,
     path_arg,
     require_versions,
+    split_alignment,
     total_memory,
 )
 from .shard import (
     determine_shards_from_fai,
     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"),
@@ -131,6 +140,10 @@ class SentieonPangenome(BasePangenome):
                 "help": argparse.SUPPRESS,
                 "action": "store_true",
             },
+            "no_split_alignment": {
+                "help": argparse.SUPPRESS,
+                "action": "store_true",
+            },
         }
     )
 
@@ -149,6 +162,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 +422,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 +480,24 @@ 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)
-            # Do not run vg-haplotypes with bwa in low-mem environments
-            if total_mem_gb < 70:
-                haplotype_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)
+            # 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
@@ -506,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
@@ -527,14 +558,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 +595,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 +605,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,10 +778,14 @@ 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(
-        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")
@@ -761,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
@@ -804,8 +849,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 +877,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