diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index c208463..6adeb5c 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), @@ -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), 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 diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index ed445e8..8cf6c32 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 ( @@ -57,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"), @@ -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 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_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) 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) 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