diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 0e5e624..0a9ddac 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -15,7 +15,7 @@ jobs: max-parallel: 1 matrix: python-version: ["3.11", "3.14"] - sentieon-version: ["202503.02"] + sentieon-version: ["${{ vars.SENTIEON_VERSION }}"] os: [ubuntu-22.04] #, macos-latest, windows-latest] runs-on: ${{ matrix.os }} steps: @@ -41,6 +41,10 @@ jobs: pip install multiqc - name: Install sentieon run: | + if [ -z "$SENTIEON_VERSION" ]; then + echo "::error::Repository variable SENTIEON_VERSION is not set." >&2 + exit 1 + fi url="https://s3.amazonaws.com/sentieon-release/software/sentieon-genomics-$SENTIEON_VERSION.tar.gz" bin_dir=$(pwd)/sentieon-genomics-$SENTIEON_VERSION/bin machine_arch=$(uname -m) @@ -98,7 +102,7 @@ jobs: max-parallel: 1 matrix: python-version: ["3.11", "3.14"] - sentieon-version: ["202503.02"] + sentieon-version: ["${{ vars.SENTIEON_VERSION }}"] os: [ubuntu-22.04] #, macos-latest, windows-latest] runs-on: ${{ matrix.os }} steps: @@ -136,6 +140,10 @@ jobs: sudo chmod ugo+x /usr/local/bin/mosdepth - name: Install sentieon run: | + if [ -z "$SENTIEON_VERSION" ]; then + echo "::error::Repository variable SENTIEON_VERSION is not set." >&2 + exit 1 + fi url="https://s3.amazonaws.com/sentieon-release/software/sentieon-genomics-$SENTIEON_VERSION.tar.gz" bin_dir=$(pwd)/sentieon-genomics-$SENTIEON_VERSION/bin machine_arch=$(uname -m) @@ -220,7 +228,7 @@ jobs: max-parallel: 1 matrix: python-version: ["3.11", "3.14"] - sentieon-version: ["202503.02"] + sentieon-version: ["${{ vars.SENTIEON_VERSION }}"] os: [ubuntu-22.04] #, macos-latest, windows-latest] runs-on: ${{ matrix.os }} steps: @@ -261,6 +269,10 @@ jobs: pip install multiqc - name: Install sentieon run: | + if [ -z "$SENTIEON_VERSION" ]; then + echo "::error::Repository variable SENTIEON_VERSION is not set." >&2 + exit 1 + fi url="https://s3.amazonaws.com/sentieon-release/software/sentieon-genomics-$SENTIEON_VERSION.tar.gz" bin_dir=$(pwd)/sentieon-genomics-$SENTIEON_VERSION/bin machine_arch=$(uname -m) @@ -296,6 +308,7 @@ jobs: --sr_aln "tests/smoke/illumina.cram" \ --lr_aln "tests/smoke/sam ple.cram" \ -m "DNAscope hybrid IlmnPB1.1.bundle" \ + --skip_cnv \ --sr_duplicate_marking none \ "output_hybrid1.vcf.gz" if [ ! -f "output_hybrid1.vcf.gz" ]; then @@ -314,6 +327,7 @@ jobs: --lr_align_input \ --lr_input_ref "tests/smoke/r ef.fa" \ -m "DNAscope hybrid IlmnPB1.1.bundle" \ + --skip_cnv \ "output_hybrid2.vcf.gz" if [ ! -f "output_hybrid2.vcf.gz" ]; then exit 1 @@ -331,6 +345,7 @@ jobs: --sr_aln "tests/smoke/illumina.cram" "tests/smoke/illumina_2.bam" \ --lr_aln "tests/smoke/sam ple.cram" \ -m "DNAscope hybrid IlmnPB1.1.bundle" \ + --skip_cnv \ --sr_duplicate_marking none \ "output_hybrid3.vcf.gz" if [ ! -f "output_hybrid3.vcf.gz" ]; then diff --git a/Dockerfile b/Dockerfile index c17c07d..4518bf7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -293,7 +293,7 @@ RUN sentieon driver --help && \ perl -c "$(command -v run-t1k)" && \ SENTIEON_LICENSE=placeholder segdup-caller --version && \ SENTIEON_LICENSE=placeholder python -c "import sys; from packaging.version import Version; \ -from sentieon_cli.sentieon_pangenome import SEGDUP_MIN_VERSION; \ +from sentieon_cli.stages.segdup import SEGDUP_MIN_VERSIONS; \ from sentieon_cli.util import check_version; \ -sys.exit(0 if all(check_version(c, v) for c, v in SEGDUP_MIN_VERSION.items()) else 1)" && \ +sys.exit(0 if all(check_version(c, v) for c, v in SEGDUP_MIN_VERSIONS.items()) else 1)" && \ sentieon-cli -h diff --git a/README.md b/README.md index b51d89f..c5e56e8 100644 --- a/README.md +++ b/README.md @@ -55,8 +55,12 @@ sentieon-cli ... ## Global arguments The `sentieon-cli` supports the following global arguments: -- `--verbose`: verbose logging. -- `--debug`: debugging mode for more verbose logging. +- `--verbose` (`-v`): verbose logging. This is the default. +- `--quiet` (`-q`): only log warnings and errors. +- `--debug` (`-d`): debugging mode for more verbose logging. Takes precedence over `--verbose` and `--quiet`. + +## Logging +Each run writes its log files to a directory next to the output VCF, named after the output file with the `.vcf.gz` suffix replaced by `_logs`, so `sample.vcf.gz` produces `sample_logs/`. The directory records the invocation in `command.txt`, the pipeline's own messages in `run.log`, and the output of each tool under `task_logs/`. Run metrics are also written as tab-separated tables for analysis: `process_metrics.txt` (one row per process, with its CPU time, peak memory and exit code) and `job_metrics.txt` (one row per pipeline job). Every pipeline accepts a `--log_dir` argument to write these files elsewhere. Rerunning a pipeline with the same output overwrites the logs of the previous run. ## Supported pipelines - [**DNAscope**](https://support.sentieon.com/docs/sentieon_cli/#dnascope) - DNAscope pipeline implementation for germline SNV and indel calling from short read data. diff --git a/pyproject.toml b/pyproject.toml index 8f866cb..0b2a465 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "sentieon_cli" -version = "1.7.0" +version = "2.0.0" description = "Pipeline implementations for the Sentieon software" authors = [ {name = "Don Freed", email = "don.freed@sentieon.com"}, @@ -27,6 +27,10 @@ dependencies = [ [project.scripts] sentieon-cli = "sentieon_cli:main" +[tool.poetry] +# Ship the packaged pseudo-autosomal region (PAR) BED files +include = [{ path = "sentieon_cli/data/*.bed", format = ["sdist", "wheel"] }] + [tool.mypy] ignore_missing_imports = true @@ -44,6 +48,7 @@ exclude = [ "sentieon_cli/scripts/hybrid_anno.py", "sentieon_cli/scripts/hybrid_select.py", "sentieon_cli/scripts/indel2cnv.py", + "sentieon_cli/scripts/sad_lad_update.py", "sentieon_cli/scripts/trimalt.py", "sentieon_cli/archive.py", "playground", diff --git a/sentieon_cli/__init__.py b/sentieon_cli/__init__.py index 3fcb34b..bffbb96 100644 --- a/sentieon_cli/__init__.py +++ b/sentieon_cli/__init__.py @@ -1,31 +1,50 @@ +import argparse + from . import argh_parser from .dnascope import DNAscopePipeline from .dnascope_hybrid import DNAscopeHybridPipeline from .dnascope_longread import DNAscopeLRPipeline +from .hybrid_pangenome import HybridPangenome +from .job import Job from .sentieon_pangenome import SentieonPangenome from .util import __version__ -def main(): - """main entry point for this project""" - parser = argh_parser.CustomArgparseParser() +def add_logging_args(parser: argparse.ArgumentParser) -> None: + """Add the console verbosity flags""" parser.add_argument( "-v", "--verbose", - help="Verbose logging", - action="store_const", - dest="loglevel", - const="INFO", - default="WARNING", + help="Verbose logging (the default)", + action="store_true", + ) + parser.add_argument( + "-q", + "--quiet", + help="Only log warnings and errors", + action="store_true", ) parser.add_argument( "-d", "--debug", help="Print debugging info", - action="store_const", - dest="loglevel", - const="DEBUG", + action="store_true", ) + + +def resolve_loglevel(args: argparse.Namespace) -> str: + """The console log level implied by the verbosity flags""" + if args.debug: + return "DEBUG" + if args.quiet and not args.verbose: + return "WARNING" + return "INFO" + + +def main(): + """main entry point for this project""" + parser = argh_parser.CustomArgparseParser() + add_logging_args(parser) parser.add_argument( "--version", action="version", @@ -58,7 +77,17 @@ def main(): pipeline.add_arguments(dnascope_pangenome_subparser) dnascope_pangenome_subparser.set_defaults(pipeline=pipeline.main) + # Hybrid pangenome + pipeline = HybridPangenome() + hybrid_pangenome_subparser = subparsers.add_parser("hybrid-pangenome") + pipeline.add_arguments(hybrid_pangenome_subparser) + hybrid_pangenome_subparser.set_defaults(pipeline=pipeline.main) + args = parser.parse_args() + args.loglevel = resolve_loglevel(args) + # Job ids must be unique for the whole run, which may execute more than + # one DAG, so numbering restarts here rather than per DAG. + Job.reset_ids() args.pipeline(args) diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index b15dfc3..5b2f328 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -2,24 +2,67 @@ A base class for pangenome pipelines """ +import argparse import copy -from enum import Enum -import json import pathlib -from typing import List, Optional - -from importlib.resources import files +import sys +from typing import Dict, Iterable, List, Optional, Sequence from . import command_strings as cmds +from .dag import DAG +from .driver import ( + AlignmentStat, + BaseAlgo, + BaseDistributionByCycle, + CoverageMetrics, + GCBias, + InsertSizeMetricAlgo, + MeanQualityByCycle, + QualDistribution, + WgsMetricsAlgo, +) from .job import Job +from .pangenome_meta import ( + PangenomeMetadataError, + PangenomeReference, + detect_pangenome_reference, + read_hapl_header, +) from .pipeline import BasePipeline -from .util import path_arg +from .shard import detect_reference_build, parse_vcf_contigs +from .stages.base import StageContext +from .stages.cnv import PangenomeCNVResult, PangenomeCNVStage +from .stages.expansion import ExpansionHunterResult, ExpansionHunterStage +from .stages.metrics import MetricsPaths +from .stages.segdup import SegdupResult, SegdupStage +from .stages.t1k import ( + DEFAULT_T1K_HLA_LOCUS, + DEFAULT_T1K_KIR_LOCUS, + T1K_MIN_VERSIONS, + T1KResult, + T1KStage, +) +from .util import ( + VcfHeaderError, + path_arg, + read_vcf_header, + require_versions, + sample_sex_arg, + vcf_id_from_header, +) +# The pangenome reference of the graphs the pipelines shipped with, used +# when a graph's own metadata cannot be read +DEFAULT_PANGENOME_REF_NAME = "GRCh38" +DEFAULT_PANGENOME_CONTIG_PREFIX = "GRCh38#0#" -class SampleSex(Enum): - FEMALE = 1 - MALE = 2 - UNKNOWN = 3 +# The linear reference build that each pangenome backbone reference +# implies. Contig names alone cannot tell the builds apart, as both name +# their chromosomes chr1..chrM. +PANGENOME_REFERENCE_BUILDS: Dict[str, str] = { + "GRCh38": "hg38", + "CHM13": "chm13", +} class BasePangenome(BasePipeline): @@ -80,6 +123,101 @@ class BasePangenome(BasePipeline): "help": "Use arguments for PCR-free data processing", "action": "store_true", }, + "par_bed": { + "help": ( + "A BED file of the pseudo-autosomal regions (PAR), used " + "for sex-aware CNV calling of male samples. Overrides the " + "PAR BED file selected for the reference genome." + ), + "type": path_arg(exists=True, is_file=True), + }, + "sample_sex": { + "help": ( + "The sample sex, used by the sex-aware callers. " + "Supplying this argument overrides the sex estimated " + "from read coverage." + ), + "metavar": "{male,female}", + "type": sample_sex_arg, + }, + "segdup_caller": { + "nargs": "*", + "help": ( + "Call variants in difficult segmental duplications with " + "segdup-caller. Supply the flag with no arguments to run " + "the caller's default gene set. Supply a comma-separated " + "list of gene names (e.g. 'CFH,CYP2D6,SMN1') to restrict " + "calling to those genes." + ), + }, + "expansion_catalog": { + "help": ( + "An ExpansionHunter variant catalog. Required for short " + "tandem repeat expansion calling." + ), + "type": path_arg(exists=True, is_file=True), + }, + "t1k_hla_seq": { + "help": ( + "The DNA HLA seq FASTA file for T1K. Required for HLA " + "calling." + ), + "type": path_arg(exists=True, is_file=True), + }, + "t1k_hla_coord": { + "help": ( + "The DNA HLA coord FASTA file for T1K. Required for HLA " + "calling." + ), + "type": path_arg(exists=True, is_file=True), + }, + "t1k_hla_locus": { + "default": DEFAULT_T1K_HLA_LOCUS, + "help": ( + "Reference interval covering the HLA locus. Reads " + "overlapping this region are extracted before being " + f"passed to T1K (default: {DEFAULT_T1K_HLA_LOCUS})." + ), + }, + "t1k_kir_seq": { + "help": ( + "The DNA KIR seq FASTA file for T1K. Required for KIR " + "calling." + ), + "type": path_arg(exists=True, is_file=True), + }, + "t1k_kir_coord": { + "help": ( + "The DNA KIR coord FASTA file for T1K. Required for KIR " + "calling." + ), + "type": path_arg(exists=True, is_file=True), + }, + "t1k_kir_locus": { + "default": DEFAULT_T1K_KIR_LOCUS, + "help": ( + "Reference interval covering the KIR locus. Reads " + "overlapping this region are extracted before being " + f"passed to T1K (default: {DEFAULT_T1K_KIR_LOCUS})." + ), + }, + # Hidden arguments. The pangenome reference name and contig + # prefix are read from the `--gbz` file; these override the + # detected values for a graph the detection cannot handle. + "pangenome_ref_name": { + "help": argparse.SUPPRESS, + }, + "pangenome_contig_prefix": { + "help": argparse.SUPPRESS, + }, + "skip_contig_checks": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, + "skip_pangenome_name_checks": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, } ) @@ -96,6 +234,564 @@ def __init__(self) -> None: self.dbsnp: Optional[pathlib.Path] = None self.kmer_memory = 30 self.pcr_free = False + self.par_bed: Optional[pathlib.Path] = None + self.segdup_caller: Optional[List[str]] = None + self.expansion_catalog: Optional[pathlib.Path] = None + self.t1k_hla_seq: Optional[pathlib.Path] = None + self.t1k_hla_coord: Optional[pathlib.Path] = None + self.t1k_hla_locus: str = DEFAULT_T1K_HLA_LOCUS + self.t1k_kir_seq: Optional[pathlib.Path] = None + self.t1k_kir_coord: Optional[pathlib.Path] = None + self.t1k_kir_locus: str = DEFAULT_T1K_KIR_LOCUS + self.has_cnv_model = False + # The pangenome reference, resolved from the graph's metadata by + # `resolve_pangenome_reference`. The two arguments are hidden + # overrides, so they stay `None` unless the user supplies them. + self.pangenome_ref_name: Optional[str] = None + self.pangenome_contig_prefix: Optional[str] = None + self.pangenome_reference: Optional[PangenomeReference] = None + self.skip_contig_checks = False + self.skip_pangenome_name_checks = False + # Parsed by `validate`; declared here for the shared checks. + # `pop_vcf_header` stays `None` when the `--pop_vcf` header + # could not be read, which only a dry run survives. + self.fai_data: Dict[str, Dict[str, int]] = {} + self.pop_vcf_header: Optional[List[str]] = None + self.pop_vcf_contigs: Dict[str, Optional[int]] = {} + # Stashed by `build_dag` for the second, sex-aware DAG + self.ploidy_json: Optional[pathlib.Path] = None + self.sr_alignments: List[pathlib.Path] = [] + self.sr_replace_rg: Optional[List[List[str]]] = None + + def validate_t1k(self) -> None: + hla_requested = self.t1k_hla_seq is not None or ( + self.t1k_hla_coord is not None + ) + if hla_requested and not (self.t1k_hla_seq and self.t1k_hla_coord): + self.logger.error( + "For HLA calling, both `--t1k_hla_seq` and `--t1k_hla_coord` " + "must be supplied." + ) + sys.exit(2) + + kir_requested = self.t1k_kir_seq is not None or ( + self.t1k_kir_coord is not None + ) + if kir_requested and not (self.t1k_kir_seq and self.t1k_kir_coord): + self.logger.error( + "For KIR calling, both `--t1k_kir_seq` and `--t1k_kir_coord` " + "must be supplied." + ) + sys.exit(2) + + if not (hla_requested or kir_requested): + return + + require_versions(T1K_MIN_VERSIONS, skip=self.skip_version_check) + + def resolve_pangenome_reference(self) -> None: + """Resolve the pangenome reference name and contig prefix. + + Both values are read from the `--gbz` file's own metadata. The + hidden `--pangenome_ref_name` and `--pangenome_contig_prefix` + arguments override the detected values, and disagreeing with the + graph is an error unless `--skip_pangenome_name_checks` is set. + + The check matters because a wrong reference name is silent: `vg + haplotypes --set-reference` only warns, writes a sampled graph + with no reference paths, `vg convert -Q` then writes a GFA with + no `SN:Z:` rGFA tags, and `pgutil lift` exits 0 with every read + unmapped. + + Called from `validate` once the reference index is parsed and + before `validate_bundle`, which selects the model bundle's + `extract..model` member by the resolved name. + """ + gbz = self.required(self.gbz, "gbz") + supplied_name = self.pangenome_ref_name + supplied_prefix = self.pangenome_contig_prefix + + detected: Optional[PangenomeReference] = None + try: + detected = detect_pangenome_reference(gbz) + except (PangenomeMetadataError, OSError) as err: + self._undetected_pangenome_reference(gbz, err) + + if detected is None: + self.pangenome_ref_name = ( + supplied_name or DEFAULT_PANGENOME_REF_NAME + ) + self.pangenome_contig_prefix = ( + supplied_prefix or DEFAULT_PANGENOME_CONTIG_PREFIX + ) + else: + self._check_pangenome_overrides(detected) + self.pangenome_reference = detected + self.pangenome_ref_name = supplied_name or detected.ref_name + self.pangenome_contig_prefix = ( + supplied_prefix or detected.contig_prefix + ) + + self.logger.info( + "Using pangenome reference '%s' (contig prefix '%s')", + self.pangenome_ref_name, + self.pangenome_contig_prefix, + ) + + if detected is not None: + self.check_pangenome_reference_fasta(detected) + self.check_pangenome_hapl(detected) + + def _undetected_pangenome_reference( + self, gbz: pathlib.Path, err: Exception + ) -> None: + """Handle a `--gbz` file whose reference could not be detected. + + Dry runs (and the unit tests behind them) parse placeholder files, + so they fall back to the historical defaults. A real run needs + both overrides to continue without the detected values. + """ + if self.dry_run: + self.logger.debug( + "Could not read the pangenome metadata of '%s': %s", gbz, err + ) + return + if self.pangenome_ref_name and self.pangenome_contig_prefix: + self.logger.warning( + "Could not read the pangenome metadata of '%s': %s. Using " + "the supplied `--pangenome_ref_name` and " + "`--pangenome_contig_prefix`.", + gbz, + err, + ) + return + self.logger.error( + "Could not read the pangenome metadata of '%s': %s. Supply " + "both `--pangenome_ref_name` and `--pangenome_contig_prefix` " + "to run without the detected values.", + gbz, + err, + ) + sys.exit(2) + + def _check_pangenome_overrides(self, detected: PangenomeReference) -> None: + """Compare the supplied overrides with the detected values""" + overrides = ( + ( + "--pangenome_ref_name", + self.pangenome_ref_name, + detected.ref_name, + ), + ( + "--pangenome_contig_prefix", + self.pangenome_contig_prefix, + detected.contig_prefix, + ), + ) + mismatched = False + for flag, supplied, found in overrides: + if supplied is None or supplied == found: + continue + mismatched = True + if self.skip_pangenome_name_checks: + self.logger.warning( + "The supplied `%s` value '%s' does not match the value " + "detected in the pangenome graph, '%s'. Using the " + "supplied value.", + flag, + supplied, + found, + ) + else: + self.logger.error( + "The supplied `%s` value '%s' does not match the value " + "detected in the pangenome graph, '%s'. Drop the " + "argument to use the detected value.", + flag, + supplied, + found, + ) + if mismatched and not self.skip_pangenome_name_checks: + sys.exit(2) + + def _pangenome_check_failed(self, msg: str, *args: object) -> None: + """Fail a pangenome consistency check. + + `--skip_pangenome_name_checks` downgrades every one of these + checks to a warning. + """ + if self.skip_pangenome_name_checks: + self.logger.warning(msg, *args) + return + self.logger.error(msg, *args) + sys.exit(2) + + def check_pangenome_reference_fasta( + self, detected: PangenomeReference + ) -> None: + """Check the reference FASTA against the pangenome backbone. + + Every contig of the backbone reference must be in the reference + index, and the build the backbone implies must match the build + detected from the index. Both checks are needed: contig names + alone cannot tell GRCh38 from CHM13. + """ + missing = [ctg for ctg in detected.contigs if ctg not in self.fai_data] + if missing: + self._pangenome_check_failed( + "%d contig(s) of the pangenome reference '%s' are missing " + "from the reference FASTA index: %s", + len(missing), + detected.ref_name, + ", ".join(missing[:10]), + ) + + expected_build = PANGENOME_REFERENCE_BUILDS.get(detected.ref_name) + fai_build = detect_reference_build(self.fai_data) + if expected_build is None or fai_build is None: + self.logger.info( + "The pangenome reference is '%s' and the reference FASTA " + "build is '%s'; the two were not compared", + detected.ref_name, + fai_build, + ) + return + if expected_build != fai_build: + self._pangenome_check_failed( + "The pangenome reference '%s' expects the '%s' reference " + "build, but the `--reference` FASTA is '%s'", + detected.ref_name, + expected_build, + fai_build, + ) + + def check_pangenome_hapl( + self, detected: Optional[PangenomeReference] + ) -> None: + """Check the `--hapl` file against the `--gbz` graph. + + A mismatched top-level chain count is only a warning: the count + equals the backbone's path count for every graph seen so far, but + it is a property of the snarl decomposition rather than a + guarantee. + """ + hapl = self.required(self.hapl, "hapl") + try: + header = read_hapl_header(hapl) + except (PangenomeMetadataError, OSError) as err: + if self.dry_run: + self.logger.debug( + "Could not read the haplotype file '%s': %s", hapl, err + ) + return + self._pangenome_check_failed("%s", err) + return + + if detected is None: + return + if header.top_level_chains != detected.n_paths: + self.logger.warning( + "The `--hapl` file '%s' has %d top-level chains, but the " + "pangenome reference '%s' has %d paths. The `--gbz` and " + "`--hapl` files may not be a matching pair.", + hapl, + header.top_level_chains, + detected.ref_name, + detected.n_paths, + ) + + def load_pop_vcf_header(self, pop_vcf: Optional[pathlib.Path]) -> None: + """Parse the `--pop_vcf` header. + + Sets `pop_vcf_header` and `pop_vcf_contigs`. The header is read + in dry runs too, so a dry run checks the pop VCF the way the real + run will. Dry runs (and the unit tests behind them) are given + placeholder files, so an unreadable header is only logged there + and leaves both attributes empty; a real run cannot continue + without the annotations the pop VCF carries. + """ + self.pop_vcf_header = None + self.pop_vcf_contigs = {} + if pop_vcf is None: + return + try: + header = read_vcf_header(pop_vcf) + except (VcfHeaderError, OSError) as err: + if self.dry_run: + self.logger.debug( + "Could not read the header of the `--pop_vcf` '%s': %s", + pop_vcf, + err, + ) + return + self.logger.error( + "Could not read the header of the `--pop_vcf` '%s': %s", + pop_vcf, + err, + ) + sys.exit(2) + self.pop_vcf_header = header + self.pop_vcf_contigs = parse_vcf_contigs(header) + self.logger.debug("VCF contigs are: %s", self.pop_vcf_contigs) + + def check_pop_vcf_id(self, bundle_vcf_id: str) -> None: + """Compare the `--pop_vcf` ID with the one the bundle expects. + + Called from `validate_bundle` in both pipelines, dry runs + included, so a dry run catches a pop VCF that does not belong + with the model bundle. Only a dry run reaches this without a + header: `load_pop_vcf_header` has already failed the real run. + """ + if self.pop_vcf_header is None: + self.logger.debug( + "The `--pop_vcf` header could not be read, so its " + "SentieonVcfID was not compared with the model bundle's" + ) + return + pop_vcf_id = vcf_id_from_header(self.pop_vcf_header) + if bundle_vcf_id != pop_vcf_id: + self.logger.error( + "The ID of the `--pop_vcf` does not match the model " + "bundle. The bundle expects '%s' and the `--pop_vcf` " + "carries '%s'.", + bundle_vcf_id, + pop_vcf_id, + ) + sys.exit(2) + + def validate_pangenome_contig_lengths(self) -> None: + """Check the `--pop_vcf` contig lengths against the reference. + + `DNAModelApply` transfers the pop VCF's annotations by position, + so a pop VCF built against another reference silently annotates + the wrong sites. Every contig of the pangenome backbone + reference must therefore carry the same length in the + `--pop_vcf` header as in the reference FASTA index, whatever the + build: the lengths come from the two files rather than from a + table of GRCh38 lengths. + + The backbone contigs come from the graph's own metadata. When + that could not be read -- a dry run with a placeholder graph, or + a real run driven by the hidden `--pangenome_ref_name` override + -- the check falls back to the contigs the index and the pop VCF + share. That fallback cannot report a contig missing from the pop + VCF, but it still catches a pop VCF from another build. + + `--skip_pangenome_name_checks` does not downgrade this check; + only the hidden `--skip_contig_checks` flag bypasses it. + """ + if self.skip_contig_checks: + return + + if not self.pop_vcf_contigs: + message = ( + "No `--pop_vcf` contigs are known, so their lengths were " + "not checked against the reference FASTA" + ) + if self.dry_run: + self.logger.debug(message) + else: + self.logger.warning(message) + return + + if self.pangenome_reference is not None: + contigs: List[str] = list(self.pangenome_reference.contigs) + else: + contigs = sorted(set(self.fai_data) & set(self.pop_vcf_contigs)) + + checked = 0 + mismatched: List[str] = [] + for ctg in contigs: + fai_length = self.fai_data.get(ctg, {}).get("length") + if fai_length is None: + # `check_pangenome_reference_fasta` has already reported + # a backbone contig that the reference index is missing + continue + checked += 1 + if ctg not in self.pop_vcf_contigs: + found = "absent" + else: + pop_length = self.pop_vcf_contigs[ctg] + if pop_length == fai_length: + continue + found = "no length" if pop_length is None else str(pop_length) + mismatched.append( + f"{ctg} (reference {fai_length}, pop VCF {found})" + ) + + if mismatched: + self.logger.error( + "%d of the %d contig(s) of the pangenome reference '%s' do " + "not have the reference FASTA's length in the `--pop_vcf`: " + "%s. The `--pop_vcf` must be built against the same " + "reference as the `--gbz` pangenome.", + len(mismatched), + checked, + self.pangenome_ref_name, + ", ".join(mismatched[:10]), + ) + sys.exit(2) + + self.logger.info( + "The `--pop_vcf` and the reference FASTA agree on the length of " + "all %d contig(s) of the pangenome reference '%s'", + checked, + self.pangenome_ref_name, + ) + + def ref_name(self) -> str: + """The resolved pangenome reference name. + + `resolve_pangenome_reference` always sets it; this narrows the + optional attribute for the consumers of the graph. + """ + return self.required(self.pangenome_ref_name, "pangenome_ref_name") + + def contig_prefix(self) -> str: + """The resolved pangenome contig prefix""" + return self.required( + self.pangenome_contig_prefix, "pangenome_contig_prefix" + ) + + def build_check_gbz_job(self, sample_gbz: pathlib.Path) -> Job: + """Check the sampled pangenome before anything consumes it""" + return Job( + cmds.cmd_check_pangenome_gbz( + sample_gbz, self.ref_name(), self.contig_prefix() + ), + "check-sample-gbz", + 1, + task_name="pangenome", + ) + + def build_check_gfa_job(self, sample_gfa: pathlib.Path) -> Job: + """Check the converted GFA before anything consumes it""" + return Job( + cmds.cmd_check_pangenome_gfa( + sample_gfa, self.ref_name(), self.contig_prefix() + ), + "check-sample-gfa", + 1, + task_name="pangenome", + ) + + def output_path(self, suffix: str) -> pathlib.Path: + """A path next to the output VCF, with `.vcf.gz` replaced""" + output_vcf = self.required(self.output_vcf, "output_vcf") + return pathlib.Path(str(output_vcf).replace(".vcf.gz", suffix)) + + def add_t1k( + self, + dag: DAG, + ctx: StageContext, + inputs: List[pathlib.Path], + upstream: Iterable[Job] = (), + ) -> List[T1KResult]: + """Genotype every enabled T1K gene group from `inputs`""" + deps = set(upstream) + gene_groups = ( + ( + "hla", + "hla-wgs", + self.t1k_hla_seq, + self.t1k_hla_coord, + self.t1k_hla_locus, + ), + ( + "kir", + "kir-wgs", + self.t1k_kir_seq, + self.t1k_kir_coord, + self.t1k_kir_locus, + ), + ) + + results: List[T1KResult] = [] + for tag, preset, gene_seq, gene_coord, locus in gene_groups: + if not (gene_seq and gene_coord): + continue + results.append( + T1KStage( + ctx=ctx, + inputs=inputs, + gene_seq=gene_seq, + gene_coord=gene_coord, + locus=locus, + preset=preset, + tag=tag, + output_dir=self.output_path(f"_{tag}"), + ).add_to(dag, deps) + ) + return results + + def add_expansion( + self, + dag: DAG, + ctx: StageContext, + alignment: pathlib.Path, + ) -> Optional[ExpansionHunterResult]: + """Genotype repeat expansions, when a catalog was supplied""" + if self.expansion_catalog is None: + return None + + return ExpansionHunterStage( + ctx=ctx, + alignment=alignment, + variant_catalog=self.expansion_catalog, + output_prefix=self.output_path("_expansion"), + sample_sex=self.sample_sex, + ).add_to(dag) + + def add_segdup( + self, + dag: DAG, + ctx: StageContext, + sr_alignment: pathlib.Path, + *, + lr_alignment: Optional[pathlib.Path] = None, + lr_bundle: Optional[pathlib.Path] = None, + overrides: Sequence[str] = (), + ) -> Optional[SegdupResult]: + """Call segmental duplications, when `--segdup_caller` was set""" + if self.segdup_caller is None: + return None + + return SegdupStage( + ctx=ctx, + sr_alignment=sr_alignment, + sr_bundle=self.required(self.model_bundle, "model_bundle"), + output_dir=self.output_path("_segdups"), + input_vcf=ctx.output_vcf, + sample_sex=self.sample_sex, + genes=",".join(self.segdup_caller) or None, + lr_alignment=lr_alignment, + lr_bundle=lr_bundle, + overrides=overrides, + ).add_to(dag) + + def add_pangenome_cnv( + self, + dag: DAG, + ctx: StageContext, + sv_vcf: pathlib.Path, + inputs: List[pathlib.Path], + *, + replace_rg: Optional[List[List[str]]] = None, + interval: Optional[pathlib.Path] = None, + ) -> PangenomeCNVResult: + """Call CNVs from the short reads and the PangenomeSV output""" + bundle = self.required(self.model_bundle, "model_bundle") + + return PangenomeCNVStage( + ctx=ctx, + inputs=inputs, + model=bundle.joinpath("cnv.model"), + sv_vcf=sv_vcf, + sample_sex=self.sample_sex, + par_bed=self.cnv_par_bed, + interval=interval, + replace_rg=replace_rg, + ).add_to(dag) def build_kmc_job( self, kmer_prefix: pathlib.Path, job_threads: int @@ -129,42 +825,20 @@ def build_kmc_job( ), "kmc", job_threads, + task_name="kmer-counting", ) return kmc_job - def build_ploidy_job( - self, - ploidy_json: pathlib.Path, - deduped_bam: List[pathlib.Path], - ) -> Job: - """Estimate sample ploidy and sex""" - estimate_ploidy = pathlib.Path( - str(files("sentieon_cli.scripts").joinpath("estimate_ploidy.py")) - ).resolve() - ploidy_job = Job( - cmds.cmd_estimate_ploidy( - ploidy_json, - deduped_bam, - estimate_ploidy, - ), - "estimate-ploidy", - 0, - ) - return ploidy_job - - def get_sex(self, ploidy_json: pathlib.Path) -> None: - """Retrieve the sample sex""" - if self.dry_run: - self.logger.info("Setting sample sex to MALE for dry-run") - self.sample_sex = SampleSex.MALE - return - with open(ploidy_json) as fh: - data = json.load(fh) - sex = data["sex"] - if sex == "female": - self.sample_sex = SampleSex.FEMALE - elif sex == "male": - self.sample_sex = SampleSex.MALE - else: - self.sample_sex = SampleSex.UNKNOWN + def pangenome_metrics_algos(self, paths: MetricsPaths) -> List[BaseAlgo]: + """The metrics collected from a deduplicated pangenome alignment""" + return [ + InsertSizeMetricAlgo(paths.insert_size), + MeanQualityByCycle(paths.mean_qual_by_cycle), + BaseDistributionByCycle(paths.base_distribution_by_cycle), + QualDistribution(paths.qual_distribution), + AlignmentStat(paths.alignment_stat), + GCBias(paths.gc_bias, summary=paths.gc_bias_summary), + WgsMetricsAlgo(paths.wgs, include_unpaired="true"), + CoverageMetrics(paths.coverage), + ] diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 4bdebde..c208463 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -10,7 +10,7 @@ import subprocess as sp import sys import typing -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Tuple, Union from .driver import BaseDriver, Driver, ReadWriter from .logging import get_logger @@ -18,6 +18,11 @@ logger = get_logger(__name__) +# Requested OS pipe-buffer size (bytes) for throughput-bound alignment +# pipelines. Applied best-effort: the kernel caps unprivileged requests at +# fs.pipe-max-size and non-Linux platforms keep their default size. +ALN_PIPE_SIZE = 268435456 + def cmd_fai_to_bed( fai: pathlib.Path, @@ -43,23 +48,6 @@ def cmd_bedtools_subtract( return Pipeline(Command(*cmd), file_output=unphased_bed) -def cmd_bedtools_merge( - in_bed: pathlib.Path, - out_bed: pathlib.Path, - distance: int = 0, -) -> Pipeline: - """Bedtools merge""" - cmd = [ - "bedtools", - "merge", - "-d", - str(distance), - "-i", - str(in_bed), - ] - return Pipeline(Command(*cmd), file_output=out_bed) - - def cmd_bedtools_slop( in_bed: pathlib.Path, out_bed: pathlib.Path, @@ -311,24 +299,163 @@ def cmd_pyexec_hybrid_anno( return Pipeline(Command(*cmd)) +def cmd_pyexec_sad_lad_update( + out_vcf: pathlib.Path, + in_vcf: pathlib.Path, + sad_lad_update: pathlib.Path, + threads: int, +) -> Pipeline: + """Update AD/DP by choosing between SAD and LAD""" + cmd = [ + sys.executable, + str(sad_lad_update), + "--input_vcf", + str(in_vcf), + "--output_vcf", + str(out_vcf), + "--threads", + str(threads), + ] + return Pipeline(Command(*cmd)) + + +PANGENOME_META_MODULE = "sentieon_cli.pangenome_meta" + +# `sentieon_cli/__init__.py` imports every pipeline, so `pangenome_meta` is +# already in `sys.modules` when `-m` re-executes it as `__main__`. The +# duplicate module is harmless -- the script holds no state -- but runpy +# warns about it on the stderr of every check job, so runpy's own warnings +# are silenced rather than left to clutter the job logs. +PANGENOME_META_WARNING_FILTER = "ignore::RuntimeWarning:runpy" + + +def cmd_pangenome_meta( + subcommand: str, + input_flag: str, + input_file: pathlib.Path, + reference_name: str, + contig_prefix: str, +) -> Pipeline: + """Run a `sentieon_cli.pangenome_meta` check on a pangenome file""" + cmd = [ + sys.executable, + "-W", + PANGENOME_META_WARNING_FILTER, + "-m", + PANGENOME_META_MODULE, + subcommand, + input_flag, + str(input_file), + "--reference_name", + reference_name, + "--contig_prefix", + contig_prefix, + ] + return Pipeline(Command(*cmd)) + + +def cmd_check_pangenome_gbz( + gbz: pathlib.Path, + reference_name: str, + contig_prefix: str, +) -> Pipeline: + """Check that a sampled pangenome kept its reference paths""" + return cmd_pangenome_meta( + "check-gbz", "--gbz", gbz, reference_name, contig_prefix + ) + + +def cmd_check_pangenome_gfa( + gfa: pathlib.Path, + reference_name: str, + contig_prefix: str, +) -> Pipeline: + """Check that a pangenome GFA carries the reference's rGFA tags""" + return cmd_pangenome_meta( + "check-gfa", "--gfa", gfa, reference_name, contig_prefix + ) + + +def cmd_pyexec_indel2cnv( + out_vcf: pathlib.Path, + in_vcf: pathlib.Path, + reference: pathlib.Path, + indel2cnv_script: pathlib.Path, + threads: int, +) -> Pipeline: + """Convert PangenomeSV INDELs to CNV calls""" + cmd = [ + sys.executable, + str(indel2cnv_script), + str(reference), + str(in_vcf), + str(out_vcf), + "-t", + str(threads), + ] + return Pipeline(Command(*cmd)) + + +def cmd_pyexec_combine_cnv( + out_vcf: pathlib.Path, + cnv_vcf: pathlib.Path, + converted_vcf: pathlib.Path, + combine_script: pathlib.Path, +) -> Pipeline: + """Combine CNVscope and converted SV calls""" + cmd = [ + sys.executable, + str(combine_script), + "--cnv", + str(cnv_vcf), + "--converted", + str(converted_vcf), + "-o", + str(out_vcf), + ] + return Pipeline(Command(*cmd)) + + +def hybrid_stage1_hap( + out_hap_bam: pathlib.Path, + stage1_driver: BaseDriver, + cores: int, +) -> Pipeline: + """Sort the haplotype alignments written to stdout by HybridStage1""" + sort_cmd = Command( + "sentieon", + "util", + "sort", + "-i", + "-", + "-t", + str(cores), + "-o", + str(out_hap_bam), + # No `--sam2bam`: the algo writes unsorted BAM, not SAM, to stdout + ) + return Pipeline(Command(*stage1_driver.build_cmd()), sort_cmd) + + def hybrid_stage1( out_aln: pathlib.Path, reference: pathlib.Path, cores: int, readgroup: str, ins_driver: BaseDriver, - stage1_driver: BaseDriver, + hap_fastq_fifo: pathlib.Path, bwa_model: pathlib.Path, ) -> Pipeline: bwa_env = dict(os.environ) _ = bwa_env.pop("bwt_max_mem", None) - # Send the input of both fq commands to bwa with cat - fq1_cmd = Command(*stage1_driver.build_cmd()) + # Send both sets of reads to bwa with cat. The HybridStage1 driver of + # the `hybrid_stage1_hap` job writes its fastq output to the fifo, so + # the fifo is read first to drain it while that job runs. fq2_cmd = Command(*ins_driver.build_cmd()) cat_cmd = Command( "cat", - InputProcSub(Pipeline(fq1_cmd)), + str(hap_fastq_fifo), InputProcSub(Pipeline(fq2_cmd)), ) @@ -506,13 +633,16 @@ def cmd_samtools_fastq_minimap2( input_ref: Optional[pathlib.Path] = None, fastq_taglist: str = "*", minimap2_args: str = "-YL", - util_sort_args: str = "--cram_write_options version=3.0,compressor=rans", + util_sort_args: str = "", + minimap2_model: Optional[Union[pathlib.Path, str]] = None, ) -> Pipeline: """Re-align an input BAM/CRAM/uBAM/uCRAM file with minimap2""" + # `input_ref` decodes the input file, which may use a different + # reference from the alignment target ref_cmd: List[str] = [] if input_ref: - ref_cmd = ["--reference", str(reference)] + ref_cmd = ["--reference", str(input_ref)] cmd_list = [ Command( "samtools", @@ -525,6 +655,9 @@ def cmd_samtools_fastq_minimap2( str(input_aln), ) ] + mm2_model = ( + minimap2_model if minimap2_model else f"{model_bundle}/minimap2.model" + ) cmd_list.append( Command( "sentieon", @@ -535,7 +668,7 @@ def cmd_samtools_fastq_minimap2( "-a", minimap2_args, "-x", - f"{model_bundle}/minimap2.model", + str(mm2_model), str(reference), "/dev/stdin", ) @@ -589,16 +722,15 @@ def cmd_samtools_fastq_bwa( bwa_args: str = "", bwa_k: str = "20000000", fastq_taglist: str = "RG", - util_sort_args: str = "--cram_write_options version=3.0,compressor=rans", + util_sort_args: str = "", ) -> Pipeline: """Re-align an input BAM/CRAM/uBAM/uCRAM file with bwa""" + # `input_ref` decodes the input file, which may use a different + # reference from the alignment target ref_cmd: List[str] = [] if input_ref: - ref_cmd = ["--reference", str(reference)] + ref_cmd = ["--reference", str(input_ref)] - pipebuf_cmd = Command( - "perl", "-MFcntl", "-e", "fcntl(STDOUT, 1031, 268435456)" - ) if collate: collate_cmd = Command( "samtools", @@ -667,20 +799,18 @@ def cmd_samtools_fastq_bwa( ) if collate_cmd: return Pipeline( - pipebuf_cmd, collate_cmd, fastq_cmd, bwa_cmd, sort_cmd, - skip_pipe=(0,), + pipe_size=ALN_PIPE_SIZE, ) else: return Pipeline( - pipebuf_cmd, fastq_cmd, bwa_cmd, sort_cmd, - skip_pipe=(0,), + pipe_size=ALN_PIPE_SIZE, ) @@ -693,7 +823,7 @@ def cmd_fastq_minimap2( cores: int, unzip: str = "gzip", minimap2_args: str = "-YL", - util_sort_args: str = "--cram_write_options version=3.0,compressor=rans", + util_sort_args: str = "", ) -> Pipeline: """Align an input fastq file with minimap2""" @@ -741,26 +871,15 @@ def cmd_fastq_bwa( unzip: str = "gzip", bwa_args: str = "", bwa_k: str = "20000000", - util_sort_args: str = "--cram_write_options version=3.0,compressor=rans", + util_sort_args: str = "", numa: Optional[str] = None, split: Optional[str] = None, ) -> Pipeline: """Align an input fastq file with bwa""" - pipebuf_cmd = Command( - "perl", "-MFcntl", "-e", "fcntl(STDOUT, 1031, 268435456)" - ) - r1_unzip = Pipeline( - pipebuf_cmd, - Command(str(unzip), "-dc", str(r1)), - skip_pipe=(0,), - ) + r1_unzip = Pipeline(Command(str(unzip), "-dc", str(r1))) r2_unzip = None if r2: - r2_unzip = Pipeline( - pipebuf_cmd, - Command(str(unzip), "-dc", str(r2)), - skip_pipe=(0,), - ) + r2_unzip = Pipeline(Command(str(unzip), "-dc", str(r2))) bwa_cmd_args: List[Union[str, InputProcSub]] = [ *(["taskset", "-c", numa] if numa else []), "sentieon", @@ -780,7 +899,6 @@ def cmd_fastq_bwa( ] if split: extract_cmd = Pipeline( - pipebuf_cmd, Command( "sentieon", "fqidx", @@ -792,7 +910,6 @@ def cmd_fastq_bwa( InputProcSub(r1_unzip), *([InputProcSub(r2_unzip)] if r2_unzip else []), ), - skip_pipe=(0,), ) bwa_cmd_args.append(InputProcSub(extract_cmd)) else: @@ -816,10 +933,9 @@ def cmd_fastq_bwa( *util_sort_args.split(), ) return Pipeline( - pipebuf_cmd, Command(str(bwa_cmd_args[0]), *bwa_cmd_args[1:]), sort_cmd, - skip_pipe=(0,), + pipe_size=ALN_PIPE_SIZE, ) @@ -893,19 +1009,22 @@ def cmd_kmc( return Pipeline(Command(*cmd)) -def cmd_extract_kmc( - output_prefix: pathlib.Path, - out_fastq: pathlib.Path, +def _readwriter_extract_cmds( input_aln: List[pathlib.Path], reference: pathlib.Path, extract_model: pathlib.Path, - tmp_dir: pathlib.Path, rw_bam: pathlib.Path, + out_fastq: pathlib.Path, threads: int = 1, -) -> Pipeline: - # `rw_bam` must already be a symlink to /dev/stdout; - # ReadWriter writes to it, which resolves to the driver's stdout (= the - # pipe to pgutil extract). +) -> Tuple[Command, Command]: + """Read an aligned input once, writing the extracted reads to a fastq + file and the full read stream to stdout. + + `rw_bam` must already be a symlink to /dev/stdout; ReadWriter writes to + it, which resolves to the driver's stdout (= the pipe to pgutil + extract). Duplicate, secondary, and supplementary reads are excluded + with `--output_flag_filter 0xf00:0`. + """ driver = Driver( reference=reference, thread_count=threads, @@ -939,6 +1058,27 @@ def cmd_extract_kmc( "-a", "-", ) + return driver_cmd, extract_cmd + + +def cmd_extract_kmc( + output_prefix: pathlib.Path, + out_fastq: pathlib.Path, + input_aln: List[pathlib.Path], + reference: pathlib.Path, + extract_model: pathlib.Path, + tmp_dir: pathlib.Path, + rw_bam: pathlib.Path, + threads: int = 1, +) -> Pipeline: + driver_cmd, extract_cmd = _readwriter_extract_cmds( + input_aln, + reference, + extract_model, + rw_bam, + out_fastq, + threads, + ) kmc_cmd = Command( "kmc", "-k29", @@ -981,188 +1121,14 @@ def cmd_vg_haplotypes( return Pipeline(Command(*cmd)) -def cmd_vg_giraffe( - sample_gam: pathlib.Path, - sample_pangenome: pathlib.Path, - fastq1: pathlib.Path, - fastq2: Optional[pathlib.Path], - readgroup="rg-1", - sample="sample", - threads=1, - max_fragment_length=3000, -) -> Pipeline: - vg_cmd = [ - "vg", - "giraffe", - "-t", - str(threads), - "-Z", - str(sample_pangenome), - "--read-group", - readgroup, - "--sample", - sample, - "-L", - str(max_fragment_length), - "--progress", - "-f", - str(fastq1), - ] - if fastq2: - vg_cmd.extend(["-f", str(fastq2)]) - - return Pipeline(Command(*vg_cmd), file_output=sample_gam) - - -def cmd_vg_pack( - sample_pack: pathlib.Path, - sample_gams: List[pathlib.Path], - gbz: pathlib.Path, - min_mapq=5, - threads=1, -) -> Pipeline: - cat_cmd = Command("cat", *[str(x) for x in sample_gams]) - vg_cmd = Command( - "vg", - "pack", - "-x", - str(gbz), - "-g", - "-", - "-o", - str(sample_pack), - "-Q", - str(min_mapq), - "--threads", - str(threads), - ) - - return Pipeline(cat_cmd, vg_cmd) - - -def cmd_sv_call( - out_svs: pathlib.Path, - sample_pack: pathlib.Path, - sv_header: pathlib.Path, - gbz: pathlib.Path, - snarls: pathlib.Path, - sample_name="", - ref_path="GRCh38", - min_indel_size=34, -) -> Pipeline: - # Call SVs with vg call - vg_call_cmd = Command( - "vg", - "call", - "-r", - str(snarls), - "-k", - str(sample_pack), - "-s", - sample_name, - "-z", - "--ref-sample", - ref_path, - "--progress", - str(gbz), - ) - - # Use bcftools to remove smaller variants - view_cmd = Command( - "bcftools", - "view", - "-i", - f"ABS(ILEN) > {min_indel_size}", - "-H", - "-", - ) - - # Rehead the vg call output - cat_cmd = Command( - "cat", - str(sv_header), - InputProcSub(Pipeline(vg_call_cmd, view_cmd)), - ) - - # Sort and compress the output - sort_cmd = Command( - "bcftools", - "sort", - "-O", - "v", - ) - convert_cmd = Command( - "sentieon", - "util", - "vcfconvert", - "-", - str(out_svs), - ) - return Pipeline(cat_cmd, sort_cmd, convert_cmd) - - -def cmd_vg_surject( - out_bam: pathlib.Path, - sample_gam: pathlib.Path, - xg: pathlib.Path, - surject_paths_dict: pathlib.Path, - threads=1, - surject_xargs: List[str] = [ - "--interleaved", - "--progress", - ], -) -> Pipeline: - surject_cmd_list = [ - "vg", - "surject", - "--sam-output", - "-t", - str(threads), - "-x", - str(xg), - "--into-paths", - shlex.quote(str(surject_paths_dict)), - ] - surject_cmd_list.extend(surject_xargs) - surject_cmd_list.append(str(sample_gam)) - - sort_cmd = Command( - "sentieon", - "util", - "sort", - "--sam2bam", - "-i", - "-", - "-o", - str(out_bam), - ) - - return Pipeline(Command(*surject_cmd_list), sort_cmd) - - -def strip_ctg_prefix( - output_header: pathlib.Path, - bam: pathlib.Path, - prefix: str, -) -> Pipeline: - # Also remove the M5 tag - samtools_cmd = Command("samtools", "view", "-H", str(bam)) - sed_cmd = Command( - "sed", - "-e", - f"s/\tSN:{prefix}/\tSN:/", - "-e", - "s/\tM5:[0-9a-zA-Z]*\t/\t/", - "-e", - "s/\tM5:[0-9a-zA-Z]*$//", - ) - return Pipeline(samtools_cmd, sed_cmd, file_output=output_header) - - def cmd_estimate_ploidy( output_json: pathlib.Path, aln_files: List[pathlib.Path], ploidy_script: pathlib.Path, + contigs: Optional[List[str]] = None, + autosomes: Optional[List[str]] = None, + x_contig: Optional[str] = None, + y_contig: Optional[str] = None, ) -> Pipeline: cmd = ( [ @@ -1176,6 +1142,18 @@ def cmd_estimate_ploidy( str(output_json), ] ) + # Contig names are only supplied when they differ from the defaults + # of the `estimate_ploidy.py` script + if contigs: + cmd.append("--contigs") + cmd.extend(contigs) + if autosomes: + cmd.append("--autosomes") + cmd.extend(autosomes) + if x_contig: + cmd.extend(["--x_contig", x_contig]) + if y_contig: + cmd.extend(["--y_contig", y_contig]) return Pipeline(Command(*cmd)) @@ -1241,16 +1219,37 @@ def cmd_segdup_caller( sex: Optional[str] = None, genes: Optional[str] = None, overrides: Optional[List[str]] = None, + lr_alignments: Optional[pathlib.Path] = None, + lr_bundle: Optional[pathlib.Path] = None, ) -> Pipeline: + """Call variants in difficult segmental duplications. + + Long reads are optional; segdup-caller needs both the long-read + alignment and the long-read model bundle to use them. + """ + if (lr_alignments is None) != (lr_bundle is None): + raise ValueError( + "segdup-caller needs both `lr_alignments` and `lr_bundle` to " + "call with long reads" + ) + cmd = [ "segdup-caller", "--short", str(sr_alignments), - "--reference", - str(reference), - "--sr_model", - str(sr_bundle), ] + if lr_alignments is not None: + cmd.extend(["--long", str(lr_alignments)]) + cmd.extend( + [ + "--reference", + str(reference), + "--sr_model", + str(sr_bundle), + ] + ) + if lr_bundle is not None: + cmd.extend(["--lr_model", str(lr_bundle)]) if input_vcf is not None: cmd.extend(["--input_vcf", str(input_vcf)]) if sex is not None: @@ -1333,8 +1332,8 @@ def cmd_bwa_extract( def cmd_vg_convert_gfa( output_gfa: pathlib.Path, input_gbz: pathlib.Path, - reference_name="GRCh38", - threads=1, + reference_name: str = "GRCh38", + threads: int = 1, ) -> Pipeline: """Convert GBZ to GFA""" cmd = [ @@ -1366,6 +1365,19 @@ def cmd_vg_paths_fasta( return Pipeline(Command(*cmd), file_output=output_fasta) +def cmd_vg_gfa2fa( + output_fasta: pathlib.Path, + gfa_file: pathlib.Path, +) -> Pipeline: + """Write every path of a GFA file as FASTA. + + A slower stand-in for `pgutil gfa2fa`, for driver releases whose + `gfa2fa` handles only GRCh38 pangenomes. + """ + cmd = ["vg", "paths", "-x", str(gfa_file), "-F"] + return Pipeline(Command(*cmd), file_output=output_fasta) + + def cmd_minimap2_lift( out_bam: pathlib.Path, sample_fasta: pathlib.Path, @@ -1429,6 +1441,183 @@ def cmd_minimap2_lift( return Pipeline(mm2_cmd, lift_cmd, sort_cmd) +def _aln_fasta_procsubs( + aln: List[Tuple[pathlib.Path, pathlib.Path]], + threads: int = 1, +) -> List[InputProcSub]: + """FASTA process substitutions for `(alignment, decode reference)` + pairs""" + procsubs: List[InputProcSub] = [] + for aln_file, decode_ref in aln: + procsubs.append( + InputProcSub( + Pipeline( + Command( + "samtools", + "fasta", + "--reference", + str(decode_ref), + "-@", + str(threads), + str(aln_file), + ) + ) + ) + ) + return procsubs + + +def _kmc_stdin_cmd( + output_prefix: pathlib.Path, + tmp_dir: pathlib.Path, + k: int = 29, + memory: int = 30, + threads: int = 1, +) -> Command: + """The patched KMC reading a FASTA stream from stdin""" + return Command( + "kmc", + f"-k{k}", + f"-m{memory}", + "-okff", + f"-t{threads}", + "-fa", + "/dev/stdin", + str(output_prefix), + str(tmp_dir), + ) + + +def cmd_hybrid_kmc( + output_prefix: pathlib.Path, + fastq: List[pathlib.Path], + aln: List[Tuple[pathlib.Path, pathlib.Path]], + tmp_dir: pathlib.Path, + k: int = 29, + memory: int = 30, + threads: int = 1, + unzip: str = "gzip", +) -> Pipeline: + """Count k-mers across fastq and aligned reads. + + KMC accepts a single input format per run, so the fastq files are + converted to FASTA and the aligned reads are extracted with + `samtools fasta`, and one FASTA stream is fed to the patched KMC + through stdin. Each alignment is supplied as an + `(alignment, decode reference)` pair. + """ + cat_args: List[Union[str, InputProcSub]] = [] + if fastq: + fq_fasta = Pipeline( + Command( + unzip, + "-dc", + *[str(x) for x in fastq], + ), + Command("awk", 'NR%4==1{print ">"substr($0,2)} NR%4==2{print}'), + ) + cat_args.append(InputProcSub(fq_fasta)) + cat_args.extend(_aln_fasta_procsubs(aln, threads)) + cat_cmd = Command("cat", *cat_args) + kmc_cmd = _kmc_stdin_cmd(output_prefix, tmp_dir, k, memory, threads) + return Pipeline(cat_cmd, kmc_cmd) + + +def cmd_hybrid_extract_kmc( + output_prefix: pathlib.Path, + out_fastq: pathlib.Path, + sr_aln: List[pathlib.Path], + lr_aln: List[Tuple[pathlib.Path, pathlib.Path]], + reference: pathlib.Path, + extract_model: pathlib.Path, + tmp_dir: pathlib.Path, + rw_bam: pathlib.Path, + k: int = 29, + memory: int = 30, + threads: int = 1, +) -> Pipeline: + """Extract reads from aligned short-read input and count k-mers across + the short and long reads in a single pass. + + The aligned short reads are read once: `pgutil extract` writes the + extracted reads to `out_fastq` and passes the full read stream on to + KMC. Each long-read alignment is supplied as an + `(alignment, decode reference)` pair. + """ + driver_cmd, extract_cmd = _readwriter_extract_cmds( + sr_aln, + reference, + extract_model, + rw_bam, + out_fastq, + threads, + ) + cat_args: List[Union[str, InputProcSub]] = [ + InputProcSub(Pipeline(driver_cmd, extract_cmd)) + ] + cat_args.extend(_aln_fasta_procsubs(lr_aln, threads)) + cat_cmd = Command("cat", *cat_args) + kmc_cmd = _kmc_stdin_cmd(output_prefix, tmp_dir, k, memory, threads) + return Pipeline(cat_cmd, kmc_cmd) + + +# Extract graph update regions from the phase sets of LongReadSV calls. +LONGREAD_SV_BED_AWK = """!/^#/ { + n=split($10,a,":"); + ps=a[n]; + if (ps != "." && ps ~ /^chr[^_]+_[0-9]+_[0-9]+$/) { + split(ps,b,"_"); + print b[1], b[2], b[3]; + } +}""" + + +def cmd_longread_sv_bed( + out_bed: pathlib.Path, + sv_vcf: pathlib.Path, + ref_fai: pathlib.Path, +) -> Pipeline: + """Generate a BED file of graph update regions from LongReadSV calls""" + zcat_cmd = Command("zcat", str(sv_vcf)) + awk_cmd = Command("awk", "-F\t", LONGREAD_SV_BED_AWK, "OFS=\t") + sort_cmd = Command("bedtools", "sort", "-faidx", str(ref_fai), "-i", "-") + merge_cmd = Command("bedtools", "merge") + return Pipeline( + zcat_cmd, + awk_cmd, + sort_cmd, + merge_cmd, + file_output=out_bed, + ) + + +def cmd_pgutil_gfa2fa( + out_fasta: pathlib.Path, + ref_fai: pathlib.Path, + gfa_file: pathlib.Path, +) -> Pipeline: + """Generate FASTA sequences from a pangenome graph""" + cmd = Command( + "sentieon", + "pgutil", + "gfa2fa", + "-F", + str(ref_fai), + "-g", + str(gfa_file), + "-o", + str(out_fasta), + ) + return Pipeline(cmd) + + +def cmd_samtools_faidx( + fasta: pathlib.Path, +) -> Pipeline: + """Index a FASTA file""" + return Pipeline(Command("samtools", "faidx", str(fasta))) + + def cmd_bcftools_merge_trim( shard_vcf: pathlib.Path, raw_vcf: pathlib.Path, diff --git a/sentieon_cli/dag.py b/sentieon_cli/dag.py index b3997b6..d40419d 100644 --- a/sentieon_cli/dag.py +++ b/sentieon_cli/dag.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import Dict, Generator, List, Optional, Set +from typing import Dict, Iterable, List, Optional, Set from .exceptions import DagExecutionError from .job import Job @@ -25,54 +25,61 @@ def __init__(self) -> None: self.ready_jobs: Dict[Job, None] = {} self.finished_jobs: List[Job] = [] - def add_job(self, job: Job, dependencies: Optional[Set[Job]] = None): - """Add a job to the DAG""" - if dependencies: - for dependency in dependencies: - if ( - dependency not in self.waiting_jobs - and dependency not in self.ready_jobs - ): - raise ValueError( - f"Dependency '{dependency}' is not in the DAG" - ) + def add_job( + self, job: Job, dependencies: Optional[Iterable[Job]] = None + ) -> None: + """Add a job to the DAG. - if isinstance(dependencies, set) and len(dependencies) > 0: - self.waiting_jobs[job] = dependencies.copy() - for dependency in dependencies: + ``dependencies`` may be any iterable of jobs already in the DAG; + it is consumed once and stored as a private set. + """ + deps: Set[Job] = ( + set(dependencies) if dependencies is not None else set() + ) + if job in self.waiting_jobs or job in self.ready_jobs: + raise ValueError( + f"{job} is already in the DAG; jobs are identified by their " + f"shell pipeline, so the same pipeline cannot be added twice." + ) + for dependency in deps: + if ( + dependency not in self.waiting_jobs + and dependency not in self.ready_jobs + ): + raise ValueError( + f"Dependency '{dependency}' is not in the DAG" + ) + + if deps: + self.waiting_jobs[job] = deps + for dependency in deps: downstream = self.planned_jobs.setdefault(dependency, []) downstream.append(job) else: self.ready_jobs[job] = None - def update_dag( - self, - ) -> Generator[Dict[Job, None], Optional[Job], None]: - """Update the DAG with a finished job""" - finished_job = yield self.ready_jobs.copy() + def mark_finished(self, job: Job) -> Dict[Job, None]: + """Record a finished job and return the jobs it unblocked. - while True: - logger.debug("Waiting jobs: %s", self.waiting_jobs) - logger.debug("Planned jobs: %s", self.planned_jobs) - logger.debug("Ready jobs: %s", self.ready_jobs) - logger.debug("Finished jobs: %s", self.finished_jobs) - logger.debug("Newly finished: %s", finished_job) + The returned mapping is used as an ordered set (values are ``None``). + Raises :class:`~sentieon_cli.exceptions.DagExecutionError` if ``job`` + was not currently ready. + """ + if job not in self.ready_jobs: + raise DagExecutionError( + f"Finished job '{job}' was not ready for execution" + ) - new_ready_jobs: Dict[Job, None] = {} - if isinstance(finished_job, Job): - if finished_job not in self.ready_jobs: - raise DagExecutionError( - f"Finished job '{finished_job}' was not ready for " - "execution" - ) + logger.debug("Marking finished: %s", job) + del self.ready_jobs[job] + self.finished_jobs.append(job) - del self.ready_jobs[finished_job] - self.finished_jobs.append(finished_job) - for dependency in self.planned_jobs.get(finished_job, []): - upstream = self.waiting_jobs[dependency] - upstream.remove(finished_job) - if len(upstream) < 1: - self.ready_jobs[dependency] = None - new_ready_jobs[dependency] = None - del self.waiting_jobs[dependency] - finished_job = yield new_ready_jobs + new_ready_jobs: Dict[Job, None] = {} + for dependency in self.planned_jobs.get(job, []): + upstream = self.waiting_jobs[dependency] + upstream.remove(job) + if len(upstream) < 1: + self.ready_jobs[dependency] = None + new_ready_jobs[dependency] = None + del self.waiting_jobs[dependency] + return new_ready_jobs diff --git a/sentieon_cli/data/par_b37.bed b/sentieon_cli/data/par_b37.bed new file mode 100644 index 0000000..e261c13 --- /dev/null +++ b/sentieon_cli/data/par_b37.bed @@ -0,0 +1,4 @@ +X 60000 2699520 PAR1 +Y 10000 2649520 PAR1 +X 154931043 155260560 PAR2 +Y 59034049 59363566 PAR2 diff --git a/sentieon_cli/data/par_chm13.bed b/sentieon_cli/data/par_chm13.bed new file mode 100644 index 0000000..fada96c --- /dev/null +++ b/sentieon_cli/data/par_chm13.bed @@ -0,0 +1,4 @@ +chrX 0 2394410 +chrX 153925834 154259566 +chrY 0 2458320 +chrY 62122809 62460029 diff --git a/sentieon_cli/data/par_hg19.bed b/sentieon_cli/data/par_hg19.bed new file mode 100644 index 0000000..01d0970 --- /dev/null +++ b/sentieon_cli/data/par_hg19.bed @@ -0,0 +1,4 @@ +chrX 60000 2699520 PAR1 +chrY 10000 2649520 PAR1 +chrX 154931043 155260560 PAR2 +chrY 59034049 59363566 PAR2 diff --git a/sentieon_cli/data/par_hg38.bed b/sentieon_cli/data/par_hg38.bed new file mode 100644 index 0000000..de11c6b --- /dev/null +++ b/sentieon_cli/data/par_hg38.bed @@ -0,0 +1,4 @@ +chrX 10000 2781479 PAR1 +chrY 10000 2781479 PAR1 +chrX 155701382 156030895 PAR2 +chrY 56887902 57217415 PAR2 diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index c4fa230..c0c6006 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -5,68 +5,55 @@ import argparse import copy import itertools -import multiprocessing as mp import os import pathlib import shutil import sys -from typing import Any, List, Optional, Set, Tuple +from typing import List, Optional, Set import packaging.version -from importlib.resources import files - -from . import command_strings as cmds from .dag import DAG from .driver import ( - AlignmentStat, - BaseDistributionByCycle, - CNVscope, - CNVModelApply, - CoverageMetrics, - Dedup, - DNAModelApply, + BaseAlgo, DNAscope, - Driver, - GCBias, - GVCFtyper, - HsMetricAlgo, - InsertSizeMetricAlgo, - MeanQualityByCycle, - LocusCollector, - QualDistribution, SVSolver, - WgsMetricsAlgo, ) from .job import Job from .pipeline import BasePipeline -from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + BWA_FASTQ_MIN_VERSIONS, + BWA_REALIGN_MIN_VERSIONS, + AlignResult, + BwaFastqStage, + BwaRealignStage, + find_unzip, +) +from .stages.base import StageContext, driver_job, rm_job +from .stages.preprocessing import ( + ShortReadPreprocessingStage, + SrPreprocessingResult, +) +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + GVCFtyperStage, + TransferApplyStage, +) from .util import ( - check_version, library_preloaded, parse_rg_line, path_arg, + require_versions, + set_bwt_max_mem, split_alignment, total_memory, ) -ALN_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308"), - "samtools": packaging.version.Version("1.16"), -} - -FQ_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308"), -} - VARIANTS_MIN_VERSIONS = { "sentieon driver": packaging.version.Version("202308"), } -CNV_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308.03"), -} - class DNAscopePipeline(BasePipeline): """The DNAscope pipeline""" @@ -218,7 +205,7 @@ class DNAscopePipeline(BasePipeline): "util_sort_args": { # "help": "Extra arguments for sentieon util sort", "help": argparse.SUPPRESS, - "default": "--cram_write_options version=3.0,compressor=rans", + "default": "", }, } ) @@ -232,7 +219,7 @@ def __init__(self) -> None: self.model_bundle: Optional[pathlib.Path] = None self.dbsnp: Optional[pathlib.Path] = None self.bed: Optional[pathlib.Path] = None - self.interval_padding = 0 + self.interval_padding: Optional[int] = None self.pcr_free = False self.gvcf = False self.duplicate_marking = "markdup" @@ -248,32 +235,16 @@ def __init__(self) -> None: self.bam_format = False self.bwa_args = "" self.bwa_k = 100000000 - self.util_sort_args = ( - "--cram_write_options version=3.0,compressor=rans" - ) + self.util_sort_args = "" self.bwt_max_mem: Optional[str] = None self.no_ramdisk = False self.no_split_alignment = False def validate(self) -> None: - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - - # uniquify pipeline attributes - self.sr_r1_fastq = self.r1_fastq - self.sr_r2_fastq = self.r2_fastq - self.sr_readgroups = self.readgroups - self.sr_duplicate_marking = self.duplicate_marking - del self.r1_fastq - del self.r2_fastq - del self.readgroups - del self.duplicate_marking + self.required(self.output_vcf, "output_vcf") # validate - if not self.sample_input and not ( - self.sr_r1_fastq and self.sr_readgroups - ): + if not self.sample_input and not (self.r1_fastq and self.readgroups): self.logger.error( "Please supply either the `--sample_input` or `--r1_fastq` " "and `--readgroups` arguments" @@ -300,21 +271,21 @@ def validate(self) -> None: "across decoy and unplaced contigs." ) - if len(self.sr_r1_fastq) != len(self.sr_readgroups): + if len(self.r1_fastq) != len(self.readgroups): self.logger.error( "The number of readgroups does not equal the number of fastq " "files" ) sys.exit(2) - for rg in self.sr_readgroups: + for rg in self.readgroups: try: parse_rg_line(rg.replace(r"\t", "\t")) except ValueError as e: self.logger.error("Invalid --readgroups value '%s': %s", rg, e) sys.exit(2) - if self.sr_r1_fastq or self.align or self.collate_align: + if self.r1_fastq or self.align or self.collate_align: self.validate_bwa_index() def configure(self) -> None: @@ -335,16 +306,16 @@ def configure_alignment(self) -> None: if shm_ok: os.environ["SENTIEON_TMPDIR"] = "/dev/shm" - self.set_bwt_max_mem( - total_input_size if shm_ok else 0, n_alignment_jobs + set_bwt_max_mem( + total_input_size if shm_ok else 0, + n_alignment_jobs, + override=self.bwt_max_mem, ) def total_input_size(self) -> int: """Find the total size of all inputs""" total_input_size = sum([x.stat().st_size for x in self.sample_input]) - for r1, r2 in itertools.zip_longest( - self.sr_r1_fastq, self.sr_r2_fastq - ): + for r1, r2 in itertools.zip_longest(self.r1_fastq, self.r2_fastq): for fq in (r1, r2): if isinstance(fq, pathlib.Path): total_input_size += fq.stat().st_size @@ -366,494 +337,170 @@ def check_shm(self, total_input_size: int, n_alignment_jobs) -> bool: return True return False - def set_bwt_max_mem( - self, - total_input_size: int, - n_alignment_jobs: int = 1, - ): - """Set the bwt_max_mem environment variable""" - if self.bwt_max_mem: - os.environ["bwt_max_mem"] = self.bwt_max_mem - return - - total_mem = total_memory() - total_mem_gb = total_mem / (1024.0**3) - align_mem_gb = ( - total_mem_gb - 4 - total_input_size / (1024.0**3) * 2.3 - ) # some memory for other system processes - bwa_mem_gb = max( - int((align_mem_gb / n_alignment_jobs) - 6), 0 - ) # some memory for other alignment processes - self.logger.debug("Setting bwt_max_mem to: %sG", bwa_mem_gb) - os.environ["bwt_max_mem"] = f"{bwa_mem_gb}G" - def build_dag(self) -> DAG: """Build the DAG for the pipeline""" self.logger.info("Building the DAG") dag = DAG() + ctx = self.stage_context() # Alignment align_jobs: Set[Job] = set() sample_input = copy.deepcopy(self.sample_input) - bam_rm_job = None + bam_cleanup: List[pathlib.Path] = [] if self.align or self.collate_align: - sample_input, align_jobs, bam_rm_job = self.sr_align_inputs() - for job in align_jobs: - dag.add_job(job) - aligned_fastq, align_fastq_jobs, fq_rm_job = self.sr_align_fastq() - for job in align_fastq_jobs: - dag.add_job(job) - sample_input += aligned_fastq + aln_result = self.add_sr_input_alignment(dag, ctx) + sample_input = aln_result.outputs + align_jobs = set(aln_result.jobs) + bam_cleanup = aln_result.cleanup_paths + fq_result = self.add_sr_fastq_alignment(dag, ctx) + align_fastq_jobs = set(fq_result.jobs) + sample_input += fq_result.outputs # Dedup and metrics - ( - deduped, - lc_job, - dedup_job, - metrics_job, - rehead_job, - ) = self.dedup_and_metrics(sample_input) - if lc_job: - dag.add_job(lc_job, align_jobs.union(align_fastq_jobs)) - if dedup_job: - dag.add_job(dedup_job, {lc_job}) - if metrics_job: - dag.add_job(metrics_job, {dedup_job}) - if rehead_job: - dag.add_job(rehead_job, {metrics_job}) - if bam_rm_job: - dag.add_job(bam_rm_job, {dedup_job}) - if fq_rm_job: - dag.add_job(fq_rm_job, {dedup_job}) + preprocessing = self.add_sr_preprocessing( + dag, ctx, sample_input, align_jobs.union(align_fastq_jobs) + ) + deduped = preprocessing.deduped + dedup_job = preprocessing.dedup_job + # An alignment has cleanup paths only when it is an intermediate, + # which is exactly when duplicate marking -- and so `dedup_job` -- + # follows it + if dedup_job: + if bam_cleanup: + dag.add_job(rm_job(bam_cleanup, "rm-bam-aln"), {dedup_job}) + if fq_result.cleanup_paths: + dag.add_job( + rm_job(fq_result.cleanup_paths, "rm-fq-aln"), + {dedup_job}, + ) # Small variants if not self.skip_small_variants: - ( - call_job, - apply_job, - rm_job, - gvcftyper_job, - svsolver_job, - sv_rm_job, - ) = self.sr_call_variants(deduped) call_dependencies: Set[Job] = set() if dedup_job: call_dependencies.add(dedup_job) else: call_dependencies.update(align_jobs) call_dependencies.update(align_fastq_jobs) - dag.add_job(call_job, call_dependencies) - dag.add_job(apply_job, {call_job}) - dag.add_job(rm_job, {apply_job}) - if gvcftyper_job: - dag.add_job(gvcftyper_job, {apply_job}) - if svsolver_job and sv_rm_job: - dag.add_job(svsolver_job, {call_job}) - dag.add_job(sv_rm_job, {svsolver_job}) + self.add_small_variant_calling( + dag, ctx, deduped, call_dependencies + ) # Multiqc if not self.skip_multiqc: multiqc_job = self.multiqc() - multiqc_dependencies: Set[Job] = set() - if lc_job: - multiqc_dependencies.add(lc_job) - if metrics_job: - multiqc_dependencies.add(metrics_job) - if rehead_job: - multiqc_dependencies.add(rehead_job) - if multiqc_job: - dag.add_job(multiqc_job, multiqc_dependencies) + dag.add_job(multiqc_job, set(preprocessing.qc_jobs)) return dag - def sr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job], Job]: + def add_sr_input_alignment( + self, dag: DAG, ctx: StageContext + ) -> AlignResult: """Align input BAM/CRAM/uBAM/uCRAM files with bwa""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") - if not self.skip_version_check: - for cmd, min_version in ALN_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) - - res: List[pathlib.Path] = [] - suffix = "bam" if self.bam_format else "cram" - util_sort_args = self.util_sort_args - if self.sr_duplicate_marking != "none": - suffix = "bam" - util_sort_args += " --bam_compression 1 " - jobs = set() - align_outputs: List[pathlib.Path] = [] - for i, input_aln in enumerate(self.sample_input): - out_aln = pathlib.Path( - str(self.output_vcf).replace( - ".vcf.gz", f"_bwa_sorted_{i}.{suffix}" - ) - ) - if self.sr_duplicate_marking != "none": - out_aln = self.tmp_dir.joinpath(f"bwa_sorted_{i}.{suffix}") - rg_lines = cmds.get_rg_lines(input_aln, self.dry_run) - rg_header = self.tmp_dir.joinpath(f"input_{i}.hdr") - with open(rg_header, "w", encoding="utf-8") as rg_fh: - for line in rg_lines: - print(line, file=rg_fh) - job = Job( - cmds.cmd_samtools_fastq_bwa( - out_aln, - input_aln, - self.reference, - self.model_bundle, - self.cores, - rg_header, - self.input_ref, - collate=self.collate_align, - bwa_args=self.bwa_args, - bwa_k=str(self.bwa_k), - util_sort_args=util_sort_args, - ), - f"bam-align-{i}", - self.cores, - ) - res.append(out_aln) - jobs.add(job) - align_outputs.append(out_aln) - align_outputs.append(pathlib.Path(str(out_aln) + ".bai")) - if suffix == "cram": - align_outputs.append(pathlib.Path(str(out_aln) + ".crai")) - - # Create an unscheduled job to remove the aligned inputs - rm_job = Job( - Pipeline( - Command("rm", *[str(x) for x in align_outputs], fail_ok=True) - ), - "rm-bam-aln", - 0, + require_versions( + BWA_REALIGN_MIN_VERSIONS, skip=self.skip_version_check ) - return (res, jobs, rm_job) - - def sr_align_fastq( - self, - ) -> Tuple[List[pathlib.Path], Set[Job], Optional[Job]]: + return BwaRealignStage( + ctx=ctx, + inputs=self.sample_input, + model_bundle=bundle, + bam_format=self.bam_format, + duplicate_marking=self.duplicate_marking, + input_ref=self.input_ref, + collate=self.collate_align, + bwa_args=self.bwa_args, + bwa_k=str(self.bwa_k), + util_sort_args=self.util_sort_args, + ).add_to(dag) + + def add_sr_fastq_alignment( + self, dag: DAG, ctx: StageContext + ) -> AlignResult: """Align fastq files to the reference genome using bwa""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - - res: List[pathlib.Path] = [] - jobs: Set[Job] = set() - if not self.sr_r1_fastq and not self.sr_readgroups: - return (res, jobs, None) - - if not self.skip_version_check: - for cmd, min_version in FQ_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") - unzip = "igzip" - if not shutil.which(unzip): - self.logger.warning( - "igzip is recommended for decompression, but is not " - "available. Falling back to gzip." - ) - unzip = "gzip" - - n_alignment_jobs = max(1, len(self.numa_nodes)) - suffix = "bam" if self.bam_format else "cram" - util_sort_args = self.util_sort_args - if self.sr_duplicate_marking != "none": - suffix = "bam" - util_sort_args += " --bam_compression 1 " - align_outputs: List[pathlib.Path] = [] - for i, (r1, r2, rg) in enumerate( - itertools.zip_longest( - self.sr_r1_fastq, self.sr_r2_fastq, self.sr_readgroups + if not self.r1_fastq and not self.readgroups: + return AlignResult( + jobs=[], terminal=set(), outputs=[], cleanup_paths=[] ) - ): - for j in range(n_alignment_jobs): - out_aln = pathlib.Path( - str(self.output_vcf).replace( - ".vcf.gz", f"_bwa_sorted_fq_{i}_{j}.{suffix}" - ) - ) - if self.sr_duplicate_marking != "none": - out_aln = self.tmp_dir.joinpath( - f"bwa_sorted_fq_{i}_{j}.{suffix}" - ) - numa = self.numa_nodes[j] if len(self.numa_nodes) > 1 else None - split = ( - f"{j}/{n_alignment_jobs}" - if len(self.numa_nodes) > 1 - else None - ) - split_cores = max(1, int(self.cores / n_alignment_jobs)) - job = Job( - cmds.cmd_fastq_bwa( - out_aln, - r1, - r2, - rg, - self.reference, - self.model_bundle, - split_cores, - unzip, - self.bwa_args, - str(self.bwa_k), - util_sort_args, - numa, - split, - ), - f"bam-align-{i}-{j}", - split_cores, - resources={f"node{j}": 1}, - ) - res.append(out_aln) - jobs.add(job) - align_outputs.append(out_aln) - align_outputs.append(pathlib.Path(str(out_aln) + ".bai")) - if suffix == "cram": - align_outputs.append(pathlib.Path(str(out_aln) + ".crai")) - - # Create an unscheduled job to remove the aligned inputs - rm_job = Job( - Pipeline( - Command("rm", *[str(x) for x in align_outputs], fail_ok=True) - ), - "rm-fq-aln", - 0, - ) - return (res, jobs, rm_job) - - def dedup_and_metrics( + require_versions(BWA_FASTQ_MIN_VERSIONS, skip=self.skip_version_check) + + return BwaFastqStage( + ctx=ctx, + r1_fastq=self.r1_fastq, + r2_fastq=self.r2_fastq, + readgroups=self.readgroups, + model_bundle=bundle, + numa_nodes=self.numa_nodes, + bam_format=self.bam_format, + duplicate_marking=self.duplicate_marking, + unzip=find_unzip(self.logger), + bwa_args=self.bwa_args, + bwa_k=str(self.bwa_k), + util_sort_args=self.util_sort_args, + ).add_to(dag) + + def add_sr_preprocessing( self, - sample_input, - ) -> Tuple[ - List[pathlib.Path], - Optional[Job], - Optional[Job], - Optional[Job], - Optional[Job], - ]: - """Perform dedup and metrics collection""" - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - suffix = "bam" if self.bam_format else "cram" - - # Create the metrics directory - sample_name = self.output_vcf.name.replace(".vcf.gz", "") - metric_base = sample_name + ".txt" - metrics_dir = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_metrics") - ) - if not self.dry_run: - metrics_dir.mkdir(exist_ok=True) - - # LocusCollector and Metrics - out_score = metrics_dir.joinpath(metric_base + ".score.txt.gz") - is_metrics = metrics_dir.joinpath(metric_base + ".insert_size.txt") - mqbc_metrics = metrics_dir.joinpath( - metric_base + ".mean_qual_by_cycle.txt" - ) - bdbc_metrics = metrics_dir.joinpath( - metric_base + ".base_distribution_by_cycle.txt" - ) - qualdist_metrics = metrics_dir.joinpath( - metric_base + ".qual_distribution.txt" - ) - as_metrics = metrics_dir.joinpath(metric_base + ".alignment_stat.txt") - coverage_metrics = metrics_dir.joinpath("coverage") - - # WES metrics - hs_metrics = metrics_dir.joinpath( - metric_base + ".hybrid-selection.txt" - ) - - # WGS metrics - wgs_metrics = metrics_dir.joinpath(metric_base + ".wgs.txt") - gc_metrics = metrics_dir.joinpath(metric_base + ".gc_bias.txt") - gc_summary = metrics_dir.joinpath(metric_base + ".gc_bias_summary.txt") - - lc_job = None - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - ) - if self.sr_duplicate_marking != "none": - driver.add_algo( - LocusCollector( - out_score, - consensus=self.consensus, - ) - ) - - # Prefer to run InsertSizeMetricAlgo after duplicate marking - if not self.skip_metrics and ( - (self.assay == "WES" and not self.bed) - or self.sr_duplicate_marking == "none" - ): - driver.add_algo(InsertSizeMetricAlgo(is_metrics)) - - if not self.skip_metrics: - driver.add_algo(MeanQualityByCycle(mqbc_metrics)) - driver.add_algo(BaseDistributionByCycle(bdbc_metrics)) - driver.add_algo(QualDistribution(qualdist_metrics)) - driver.add_algo(AlignmentStat(as_metrics)) - if self.assay == "WGS": - driver.add_algo(GCBias(gc_metrics, summary=gc_summary)) - - if not (self.sr_duplicate_marking == "none" and self.skip_metrics): - lc_job = Job( - Pipeline(Command(*driver.build_cmd())), - "locuscollector", - self.cores, - ) - - if self.sr_duplicate_marking == "none": - return (sample_input, lc_job, None, None, None) - - # Dedup - deduped = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_deduped.{suffix}") - ) - dedup_metrics = metrics_dir.joinpath( - metric_base + ".dedup_metrics.txt" - ) - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - ) - driver.add_algo( - Dedup( - deduped, - out_score, - cram_write_options="version=3.0,compressor=rans", - metrics=dedup_metrics, - rmdup=(self.sr_duplicate_marking == "rmdup"), - ) - ) - dedup_job = Job( - Pipeline(Command(*driver.build_cmd())), "dedup", self.cores - ) - - if self.skip_metrics: - return ([deduped], lc_job, dedup_job, None, None) - - # Run HsMetricAlgo after duplicate marking to account for - # duplicate reads - metrics_job = None - rehead_job = None - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=( - [deduped] - if self.sr_duplicate_marking != "none" - else sample_input - ), - interval=self.bed, - ) - if self.assay == "WES" and self.bed: - driver.add_algo(HsMetricAlgo(hs_metrics, self.bed, self.bed)) - driver.add_algo(InsertSizeMetricAlgo(is_metrics)) - metrics_job = Job( - Pipeline(Command(*driver.build_cmd())), "metrics", 0 - ) # Run metrics in the background - - # Run WgsMetricsAlgo after duplicate marking to account for - # duplicate reads - if self.assay == "WGS": - driver.add_algo(InsertSizeMetricAlgo(is_metrics)) - driver.add_algo( - WgsMetricsAlgo(wgs_metrics, include_unpaired="true") - ) - driver.add_algo(CoverageMetrics(coverage_metrics)) - metrics_job = Job( - Pipeline(Command(*driver.build_cmd())), "metrics", 0 - ) # Run metrics in the background - - # Rehead WGS metrics so they are recognized by MultiQC - rehead_script = pathlib.Path( - str( - files("sentieon_cli.scripts").joinpath( - "rehead_wgs_metrics.py" - ) - ) - ) - rehead_job = Job( - Pipeline( - Command( - sys.executable, - str(rehead_script), - "--metrics_file", - str(wgs_metrics), - ) - ), - "Rehead metrics", - 0, - ) - return ([deduped], lc_job, dedup_job, metrics_job, rehead_job) - - def sr_call_variants( + dag: DAG, + ctx: StageContext, + sample_input: List[pathlib.Path], + upstream: Set[Job], + ) -> SrPreprocessingResult: + """Mark duplicates and collect metrics from the short reads""" + return ShortReadPreprocessingStage( + ctx=ctx, + inputs=sample_input, + duplicate_marking=self.duplicate_marking, + consensus=self.consensus, + skip_metrics=self.skip_metrics, + assay=self.assay, + bed=self.bed, + bam_format=self.bam_format, + ).add_to(dag, upstream) + + def add_small_variant_calling( self, + dag: DAG, + ctx: StageContext, deduped: List[pathlib.Path], - ) -> Tuple[Job, Job, Job, Optional[Job], Optional[Job], Optional[Job]]: + upstream: Set[Job], + ) -> None: """Call SNVs, indels, and SVs using DNAscope""" - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") - if not self.skip_version_check: - for cmd, min_version in VARIANTS_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(VARIANTS_MIN_VERSIONS, skip=self.skip_version_check) out_gvcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", ".g.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", ".g.vcf.gz") ) out_svs_tmp = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_svs_tmp.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", "_svs_tmp.vcf.gz") ) out_svs = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_svs.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", "_svs.vcf.gz") ) - emit_mode = "gvcf" if self.gvcf else "variant" pcr_indel_model = "NONE" if self.pcr_free else "CONSERVATIVE" - model = self.model_bundle.joinpath("dnascope.model") + model = bundle.joinpath("dnascope.model") # Call variants with DNAscope emit_mode = "variant" - ds_out = self.output_vcf + ds_out = ctx.output_vcf tmp_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_tmp.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", "_tmp.vcf.gz") ) if self.gvcf: emit_mode = "gvcf" ds_out = out_gvcf tmp_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_tmp.g.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", "_tmp.g.vcf.gz") ) - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=deduped, - interval=self.bed, - interval_padding=self.interval_padding, - ) - driver.add_algo( + + algos: List[BaseAlgo] = [ DNAscope( tmp_vcf, dbsnp=self.dbsnp, @@ -861,154 +508,57 @@ def sr_call_variants( pcr_indel_model=pcr_indel_model, model=model, ) - ) + ] if not self.skip_svs: - driver.add_algo( + algos.append( DNAscope( out_svs_tmp, dbsnp=self.dbsnp, var_type="BND", ) ) - call_job = Job( - Pipeline(Command(*driver.build_cmd())), - "variant-calling", - self.cores, - ) + call = DNAscopeStage( + ctx=ctx, + algos=algos, + inputs=deduped, + interval=self.bed, + interval_padding=self.interval_padding, + ).add_to(dag, upstream) # Genotyping and filtering with DNAModelApply - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - driver.add_algo( - DNAModelApply( - model, - tmp_vcf, - ds_out, - ) - ) - apply_job = Job( - Pipeline(Command(*driver.build_cmd())), "model-apply", self.cores - ) + transfer_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=tmp_vcf, + apply=ApplySpec(model=model, output=ds_out), + ).add_to(dag, call.terminal) # Remove the tmp_vcf - rm_cmd = ["rm", str(tmp_vcf), str(tmp_vcf) + ".tbi"] - rm_job = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp-vcf", 0) + dag.add_job( + rm_job([tmp_vcf, str(tmp_vcf) + ".tbi"], "rm-tmp-vcf"), + transfer_apply.terminal, + ) # Genotype gVCFs - gvcftyper_job = None if self.gvcf: - driver = Driver( - reference=self.reference, - thread_count=self.cores, + GVCFtyperStage( + ctx=ctx, + gvcf=out_gvcf, + output=ctx.output_vcf, interval=self.bed, - ) - driver.add_algo( - GVCFtyper( - output=self.output_vcf, - vcf=out_gvcf, - ) - ) - gvcftyper_job = Job( - Pipeline(Command(*driver.build_cmd())), "gvcftyper", self.cores - ) + ).add_to(dag, transfer_apply.terminal) # Call SVs - svsolver_job = None - sv_rm_job = None if not self.skip_svs: - driver = Driver( - reference=self.reference, - thread_count=self.cores, + svsolver_job = driver_job( + ctx, + [SVSolver(output=out_svs, vcf=out_svs_tmp)], interval=self.bed, + threads=1, + name="svsolver", + task_name="sv-calling", ) - driver.add_algo( - SVSolver( - output=out_svs, - vcf=out_svs_tmp, - ) - ) - svsolver_job = Job( - Pipeline(Command(*driver.build_cmd())), "svsolver" + dag.add_job(svsolver_job, call.terminal) + dag.add_job( + rm_job([out_svs_tmp, str(out_svs_tmp) + ".tbi"], "rm-tmp-sv"), + {svsolver_job}, ) - sv_rm_job = Job( - Pipeline( - Command( - "rm", - str(out_svs_tmp), - str(out_svs_tmp) + ".tbi", - fail_ok=True, - ) - ), - "rm-tmp-sv", - 0, - ) - - return ( - call_job, - apply_job, - rm_job, - gvcftyper_job, - svsolver_job, - sv_rm_job, - ) - - -def call_cnvs( - tmp_dir: pathlib.Path, - output_vcf: pathlib.Path, - reference: pathlib.Path, - sample_input: List[pathlib.Path], - model_bundle: pathlib.Path, - bed: Optional[pathlib.Path] = None, - cores: int = mp.cpu_count(), - skip_version_check: bool = False, - replace_rg: Optional[List[List[str]]] = None, - **_kwargs: Any, -) -> Tuple[Job, Job]: - """ - Call CNVs using CNVscope - """ - if not skip_version_check: - for cmd, min_version in CNV_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) - - cnvscope_vcf = tmp_dir.joinpath("cnvscope.vcf.gz") - cnv_vcf = pathlib.Path(str(output_vcf).replace(".vcf.gz", ".cnv.vcf.gz")) - driver = Driver( - reference=reference, - thread_count=cores, - replace_rg=replace_rg, - input=sample_input, - interval=bed, - ) - driver.add_algo( - CNVscope( - cnvscope_vcf, - model_bundle.joinpath("cnv.model"), - ) - ) - cnvscope_job = Job( - Pipeline(Command(*driver.build_cmd())), "CNVscope", cores - ) - - driver = Driver( - reference=reference, - thread_count=cores, - ) - driver.add_algo( - CNVModelApply( - cnv_vcf, - model_bundle.joinpath("cnv.model"), - vcf=cnvscope_vcf, - ) - ) - cnvmodelapply_job = Job( - Pipeline(Command(*driver.build_cmd())), - "CNVModelApply", - cores, - ) - - return (cnvscope_job, cnvmodelapply_job) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index e1a9e28..8ad702b 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -7,7 +7,7 @@ import json import pathlib import sys -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Dict, List, Optional, Set, Tuple import packaging.version @@ -18,40 +18,63 @@ from . import command_strings as cmds from .dag import DAG from .driver import ( - BaseDriver, Driver, DNAscope, - DNAModelApply, HybridStage1, HybridStage2, HybridStage3, ) -from .dnascope import call_cnvs, DNAscopePipeline -from .dnascope_longread import DNAscopeLRPipeline from .job import Job from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + BWA_FASTQ_MIN_VERSIONS, + MINIMAP2_REALIGN_MIN_VERSIONS, + AlignResult, + BwaFastqStage, + Minimap2RealignStage, + find_unzip, +) +from .stages.base import StageContext, driver_job, rm_job +from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage +from .stages.metrics import MOSDEPTH_MIN_VERSIONS, MosdepthStage +from .stages.ploidy import PloidyStage +from .stages.preprocessing import ( + ShortReadPreprocessingStage, + SrPreprocessingResult, +) +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + TransferApplyStage, + TransferSpec, +) +from .stages.sv import LONGREADSV_MIN_VERSIONS, LongReadSVStage +from .stages.transfer import TransferConfig from .util import ( __version__, - check_version, library_preloaded, parse_rg_line, path_arg, + require_versions, + sample_sex_arg, + set_bwt_max_mem, split_alignment, vcf_id, + versions_available, ) from .shard import ( determine_shards_from_fai, parse_fai, vcf_contigs, ) -from .transfer import build_transfer_jobs logger = get_logger(__name__) CALLING_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202503.01"), + # 202503.04 writes unsorted BAM to `--hap_bam`, supporting stdout + "sentieon driver": packaging.version.Version("202503.04"), "bedtools": None, "bcftools": packaging.version.Version("1.22"), "samtools": packaging.version.Version("1.16"), @@ -117,7 +140,7 @@ def __init__( ) -class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): +class DNAscopeHybridPipeline(BasePipeline): """The DNAscope Hybrid pipeline""" params = copy.deepcopy(BasePipeline.params) @@ -201,6 +224,14 @@ class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): ), "type": path_arg(exists=True, is_file=True), }, + "par_bed": { + "help": ( + "A BED file of the pseudo-autosomal regions (PAR), used " + "for sex-aware CNV calling of male samples. Overrides the " + "PAR BED file selected for the reference genome." + ), + "type": path_arg(exists=True, is_file=True), + }, "pop_vcf": { "flags": ["--pop_vcf"], "help": ( @@ -208,6 +239,15 @@ class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): ), "type": path_arg(exists=True, is_file=True), }, + "sample_sex": { + "help": ( + "The sample sex, used for sex-aware CNV calling. " + "Supplying this argument overrides the sex estimated " + "from read coverage." + ), + "metavar": "{male,female}", + "type": sample_sex_arg, + }, "rgsm": { "help": ( "Overwrite the SM tag of the input readgroups for " @@ -266,7 +306,7 @@ class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): "minimap2_args": { # help="Extra arguments for sentieon minimap2", "help": argparse.SUPPRESS, - "default": "-Y", + "default": "-YL", }, "no_split_alignment": { "help": argparse.SUPPRESS, @@ -286,36 +326,58 @@ class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): "util_sort_args": { # help="Extra arguments for sentieon util sort", "help": argparse.SUPPRESS, - "default": "--cram_write_options version=3.0,compressor=rans", + "default": "", }, } ) def __init__(self) -> None: super().__init__() + # Required arguments + self.lr_aln: List[pathlib.Path] = [] + self.model_bundle: Optional[pathlib.Path] = None + self.sr_aln: List[pathlib.Path] = [] self.sr_r1_fastq: List[pathlib.Path] = [] self.sr_r2_fastq: List[pathlib.Path] = [] self.sr_readgroups: List[str] = [] - self.sr_aln: List[pathlib.Path] = [] - self.sr_duplicate_marking = "markdup" - self.lr_aln: List[pathlib.Path] = [] + # Additional arguments + self.bam_format = False + self.bed: Optional[pathlib.Path] = None + self.dbsnp: Optional[pathlib.Path] = None + self.gvcf = False self.lr_align_input = False self.lr_input_ref: Optional[pathlib.Path] = None + self.par_bed: Optional[pathlib.Path] = None self.pop_vcf: Optional[pathlib.Path] = None - self.bam_format = False self.rgsm: Optional[str] = None + self.skip_cnv = False + self.skip_metrics = False + self.skip_mosdepth = False + self.skip_multiqc = False + self.skip_svs = False + self.sr_duplicate_marking = "markdup" + # Hidden arguments + self.bwa_args = "" + self.bwa_k = 100000000 + self.bwt_max_mem: Optional[str] = None self.lr_fastq_taglist = "*" - self.sr_read_filter: Optional[str] = None self.lr_read_filter: Optional[str] = None - self.assay = "WGS" + self.minimap2_args = "-YL" + self.no_split_alignment = False self.skip_model_apply = False self.skip_pop_vcf_id_check = False + self.sr_read_filter: Optional[str] = None + self.util_sort_args = "" + # Stashed by `build_dag` for the second, sex-aware DAG + self.ploidy_json: Optional[pathlib.Path] = None + self.cnv_sr_aln: List[pathlib.Path] = [] + self.cnv_replace_rg: Optional[List[List[str]]] = None def validate(self) -> None: self.fai_data = parse_fai(pathlib.Path(str(self.reference) + ".fai")) self.pop_vcf_contigs: Dict[str, Optional[int]] = {} if self.pop_vcf: - self.pop_vcf_contigs = vcf_contigs(self.pop_vcf, self.dry_run) + self.pop_vcf_contigs = vcf_contigs(self.pop_vcf) self.logger.debug("VCF contigs are: %s", self.pop_vcf_contigs) self.shards = determine_shards_from_fai( self.fai_data, 10 * 1000 * 1000 @@ -324,6 +386,8 @@ def validate(self) -> None: self.validate_bundle() self.collect_readgroups() self.validate_readgroups() + # `validate_bundle` may have set `skip_cnv` + self.validate_cnv() self.validate_output_vcf() if not self.sr_aln and not self.sr_r1_fastq: @@ -355,6 +419,21 @@ def validate(self) -> None: if self.sr_r1_fastq: self.validate_bwa_index() + def validate_cnv(self) -> None: + """Validate the arguments used for sex-aware CNV calling""" + self.resolve_cnv_par_bed( + self.fai_data, self.par_bed, not self.skip_cnv + ) + if self.skip_cnv: + return + + # A PAR BED file is required for CNV calling, whatever the sex + self.validate_cnv_par(True) + + # CNV calling now runs in the second DAG, so the version of the + # driver is checked before the first DAG starts + require_versions(CNV_MIN_VERSIONS, skip=self.skip_version_check) + def validate_bundle(self) -> None: bundle_info_bytes = ar_load( str(self.model_bundle) + "/bundle_info.json" @@ -501,7 +580,7 @@ def configure(self) -> None: self.numa_nodes = split_alignment(self.cores) n_alignment_jobs = max(1, len(self.numa_nodes)) - self.set_bwt_max_mem(0, n_alignment_jobs) + set_bwt_max_mem(0, n_alignment_jobs, override=self.bwt_max_mem) def configure_readgroups(self) -> None: self.lr_aln_readgroups = self.all_readgroups[0] @@ -525,15 +604,8 @@ def build_dag(self) -> DAG: self.logger.info("Building the DAG") dag = DAG() - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + self.required(self.model_bundle, "model_bundle") + ctx = self.stage_context() rg_info = RgInfo( self.lr_aln_readgroups, @@ -547,225 +619,282 @@ def build_dag(self) -> DAG: # Short-read alignment sr_aln = self.sr_aln - ( - aligned_fastq, - align_fastq_jobs, - fq_rm_job, - ) = self.sr_align_fastq() - for job in align_fastq_jobs: - dag.add_job(job) - sr_aln.extend(aligned_fastq) + fq_result = self.add_sr_fastq_alignment(dag, ctx) + align_fastq_jobs = set(fq_result.jobs) + sr_aln.extend(fq_result.outputs) # Short-read dedup - ( - deduped, - lc_job, - dedup_job, - metrics_job, - rehead_job, - ) = self.dedup_and_metrics( - sample_input=sr_aln, + preprocessing = self.add_sr_preprocessing( + dag, ctx, sr_aln, set(align_fastq_jobs) ) - sr_aln = deduped - if lc_job: - dag.add_job(lc_job, align_fastq_jobs) - if dedup_job: - dag.add_job(dedup_job, {lc_job}) - if metrics_job and not self.skip_metrics: - dag.add_job(metrics_job, {dedup_job}) - if rehead_job: - dag.add_job(rehead_job, {metrics_job}) - if fq_rm_job: - dag.add_job(fq_rm_job, {dedup_job}) + sr_aln = preprocessing.deduped + dedup_job = preprocessing.dedup_job + # The fastq alignment has cleanup paths only when it is an + # intermediate, which is exactly when duplicate marking -- and so + # `dedup_job` -- follows it + if dedup_job and fq_result.cleanup_paths: + dag.add_job( + rm_job(fq_result.cleanup_paths, "rm-fq-aln"), {dedup_job} + ) # Long-read alignment realign_jobs: Set[Job] = set() lr_aln = self.lr_aln if self.lr_align_input: - lr_aln, realign_jobs = self.lr_align_inputs() - for job in realign_jobs: - dag.add_job(job) + realign_result = self.add_lr_realignment(dag, ctx) + lr_aln = realign_result.outputs + realign_jobs = set(realign_result.jobs) if not self.skip_mosdepth: - mosdepth_jobs = self.mosdepth(sample_input=lr_aln) - for job in mosdepth_jobs: - dag.add_job(job, realign_jobs) + self.add_mosdepth(dag, ctx, lr_aln, realign_jobs) if not self.skip_multiqc: multiqc_job = self.multiqc() - multiqc_dependencies: Set[Job] = set() - if lc_job: - multiqc_dependencies.add(lc_job) - if metrics_job: - multiqc_dependencies.add(metrics_job) - if rehead_job: - multiqc_dependencies.add(rehead_job) - if multiqc_job: - dag.add_job(multiqc_job, multiqc_dependencies) + dag.add_job(multiqc_job, set(preprocessing.qc_jobs)) if not self.skip_svs: - longreadsv_job = self.call_svs( - lr_aln, replace_rg=rg_info.replace_rg_args[0] + self.add_sv_calling( + dag, + ctx, + lr_aln, + rg_info.replace_rg_args[0], + realign_jobs, ) - dag.add_job(longreadsv_job, realign_jobs) sr_preprocessing_jobs: Set[Job] = set() sr_preprocessing_jobs.update(align_fastq_jobs) if dedup_job: sr_preprocessing_jobs.add(dedup_job) + # Estimate the sample ploidy and sex. The JSON output is always + # written; `--sample_sex` takes precedence for the sex used by + # sex-aware CNV calling. + ploidy_result = PloidyStage( + ctx=ctx, + inputs=sr_aln, + reference_build=self.reference_build, + ).add_to(dag, sr_preprocessing_jobs) + self.ploidy_json = ploidy_result.ploidy_json + if not self.skip_cnv: - cnvscope_job, cnvmodelapply_job = call_cnvs( - self.tmp_dir, - self.output_vcf, - self.reference, - sr_aln, - self.model_bundle, - self.bed, - self.cores, - self.skip_version_check, - replace_rg=( - rg_info.replace_rg_args[1] - if rg_info.replace_rg_args[1] - else None - ), + # CNV calling is sex-aware and runs in the second DAG + self.cnv_sr_aln = sr_aln + self.cnv_replace_rg = ( + rg_info.replace_rg_args[1] + if rg_info.replace_rg_args[1] + else None ) - dag.add_job(cnvscope_job, sr_preprocessing_jobs) - dag.add_job(cnvmodelapply_job, {cnvscope_job}) - - ( - call_job, - select_job, - mapq0_job, - mapq0_slop_job, - cat_merge_job, - rm_job1, - stage1_job, - rm_job2, - second_stage_job, - rm_job3, - third_stage_job, - rm_job4, - call2_job, - subset_job, - concat_job, - rm_job5, - anno_job, - transfer_jobs, - transfer_concat, - apply_job, - norm_job, - ) = self.call_variants(sr_aln, lr_aln, rg_info) - dag.add_job(call_job, realign_jobs | sr_preprocessing_jobs) - dag.add_job(select_job, {call_job}) - dag.add_job(mapq0_job, realign_jobs | sr_preprocessing_jobs) - dag.add_job(mapq0_slop_job, {mapq0_job}) - dag.add_job(cat_merge_job, {mapq0_slop_job, select_job}) - dag.add_job(stage1_job, {cat_merge_job}) - dag.add_job(second_stage_job, {stage1_job}) - dag.add_job(third_stage_job, {second_stage_job}) - dag.add_job(call2_job, {third_stage_job}) - dag.add_job(subset_job, {second_stage_job}) - dag.add_job(concat_job, {subset_job, call2_job}) - dag.add_job(anno_job, {concat_job}) - apply_dependencies = {anno_job} - if transfer_jobs and transfer_concat: - for job in transfer_jobs: - dag.add_job(job, {anno_job}) - dag.add_job(transfer_concat, set(transfer_jobs)) - apply_dependencies = {transfer_concat} - - if apply_job: - dag.add_job(apply_job, apply_dependencies) - if apply_job and norm_job: - dag.add_job(norm_job, {apply_job}) - - # Remove intermediate files during processing - if not self.retain_tmpdir: - dag.add_job(rm_job1, {cat_merge_job}) - dag.add_job(rm_job2, {stage1_job}) - dag.add_job(rm_job3, {second_stage_job}) - dag.add_job(rm_job4, {third_stage_job}) - dag.add_job(rm_job5, {concat_job}) + self.add_variant_calling( + dag, + ctx, + sr_aln, + lr_aln, + rg_info, + realign_jobs | sr_preprocessing_jobs, + ) + + return dag + + def add_sr_fastq_alignment( + self, dag: DAG, ctx: StageContext + ) -> AlignResult: + """Align the short-read fastq files with bwa""" + bundle = self.required(self.model_bundle, "model_bundle") + + if not self.sr_r1_fastq and not self.sr_readgroups: + return AlignResult( + jobs=[], terminal=set(), outputs=[], cleanup_paths=[] + ) + + require_versions(BWA_FASTQ_MIN_VERSIONS, skip=self.skip_version_check) + + return BwaFastqStage( + ctx=ctx, + r1_fastq=self.sr_r1_fastq, + r2_fastq=self.sr_r2_fastq, + readgroups=self.sr_readgroups, + model_bundle=bundle, + numa_nodes=self.numa_nodes, + bam_format=self.bam_format, + duplicate_marking=self.sr_duplicate_marking, + unzip=find_unzip(self.logger), + bwa_args=self.bwa_args, + bwa_k=str(self.bwa_k), + util_sort_args=self.util_sort_args, + ).add_to(dag) + + def add_sr_preprocessing( + self, + dag: DAG, + ctx: StageContext, + sr_aln: List[pathlib.Path], + upstream: Set[Job], + ) -> SrPreprocessingResult: + """Mark duplicates and collect metrics from the short reads""" + return ShortReadPreprocessingStage( + ctx=ctx, + inputs=sr_aln, + duplicate_marking=self.sr_duplicate_marking, + consensus=False, + skip_metrics=self.skip_metrics, + assay="WGS", + bed=self.bed, + bam_format=self.bam_format, + ).add_to(dag, upstream) + + def add_lr_realignment(self, dag: DAG, ctx: StageContext) -> AlignResult: + """Re-align the long-read input files with minimap2""" + bundle = self.required(self.model_bundle, "model_bundle") + require_versions( + MINIMAP2_REALIGN_MIN_VERSIONS, skip=self.skip_version_check + ) + + return Minimap2RealignStage( + ctx=ctx, + inputs=self.lr_aln, + model_bundle=bundle, + sample_name=ctx.output_vcf.name.replace(".vcf.gz", ""), + bam_format=self.bam_format, + input_ref=self.lr_input_ref, + fastq_taglist=self.lr_fastq_taglist, + minimap2_args=self.minimap2_args, + util_sort_args=self.util_sort_args, + ).add_to(dag) + + def add_mosdepth( + self, + dag: DAG, + ctx: StageContext, + inputs: List[pathlib.Path], + upstream: Set[Job], + ) -> None: + """Run mosdepth for QC, when it is available""" + if not versions_available( + MOSDEPTH_MIN_VERSIONS, skip=self.skip_version_check + ): + self.logger.warning( + "Skipping mosdepth. mosdepth version %s or later not found", + MOSDEPTH_MIN_VERSIONS["mosdepth"], + ) + return + + MosdepthStage(ctx=ctx, inputs=inputs).add_to(dag, upstream) + + def add_sv_calling( + self, + dag: DAG, + ctx: StageContext, + inputs: List[pathlib.Path], + replace_rg: Optional[List[List[str]]], + upstream: Set[Job], + ) -> None: + """Call SVs from the long reads with Sentieon LongReadSV""" + bundle = self.required(self.model_bundle, "model_bundle") + require_versions(LONGREADSV_MIN_VERSIONS, skip=self.skip_version_check) + + LongReadSVStage( + ctx=ctx, + inputs=inputs, + model=bundle.joinpath("longreadsv.model"), + interval=self.bed, + replace_rg=replace_rg, + ).add_to(dag, upstream) + def build_second_dag(self) -> Optional[DAG]: + """Build the second DAG, for sex-aware CNV calling""" + if self.skip_cnv: + return None + bundle = self.required(self.model_bundle, "model_bundle") + assert self.ploidy_json is not None + ctx = self.stage_context() + + self.get_sex(self.ploidy_json) + + self.logger.info("Building the DNAscope Hybrid CNV DAG") + dag = DAG() + CNVscopeStage( + ctx=ctx, + inputs=self.cnv_sr_aln, + model=bundle.joinpath("cnv.model"), + cnvscope_vcf=ctx.tmp_dir.joinpath("cnvscope.vcf.gz"), + cnv_vcf=pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", ".cnv.vcf.gz") + ), + sample_sex=self.sample_sex, + par_bed=self.cnv_par_bed, + dump_probes=pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", ".cnv.probes") + ), + interval=self.bed, + replace_rg=self.cnv_replace_rg, + ).add_to(dag) return dag - def call_variants( + def add_cleanup( + self, + dag: DAG, + paths: List[pathlib.Path], + name: str, + upstream: Set[Job], + ) -> None: + """Remove intermediate files once `upstream` has finished. + + A no-op under `--retain_tmpdir`, which keeps everything the run + wrote in the temporary directory. + """ + if self.retain_tmpdir: + return + dag.add_job(rm_job(paths, name), upstream) + + def add_variant_calling( self, + dag: DAG, + ctx: StageContext, sr_aln: List[pathlib.Path], lr_aln: List[pathlib.Path], rg_info: RgInfo, - **_kwargs: Any, - ) -> Tuple[ - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Job, - Optional[List[Job]], - Optional[Job], - Optional[Job], - Optional[Job], - ]: + upstream: Set[Job], + ) -> None: """ Call SNVs and indels using the DNAscope hybrid pipeline + + A first calling pass over the input alignments picks the regions + that need realignment; the three hybrid-realignment stages realign + the short reads there; a second pass recalls those regions, and the + two call sets are merged, annotated and filtered. """ - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") - ref_fai = pathlib.Path(str(self.reference) + ".fai") + ref_fai = pathlib.Path(str(ctx.reference) + ".fai") - if not self.skip_version_check: - for cmd, min_version in CALLING_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(CALLING_MIN_VERSIONS, skip=self.skip_version_check) + + hybrid_model = bundle.joinpath("hybrid.model") # First pass - combined variant calling vcf_suffix = ".g.vcf.gz" if self.gvcf else ".vcf.gz" combined_vcf = self.tmp_dir.joinpath("initial" + vcf_suffix) emit_mode = "gvcf" if self.gvcf else None - driver: BaseDriver = Driver( - reference=self.reference, - thread_count=self.cores, - replace_rg=rg_info.replace_rg_args[0] + rg_info.replace_rg_args[1], - input=lr_aln + sr_aln, + call = DNAscopeStage( + ctx=ctx, + algos=[ + DNAscope( + combined_vcf, + dbsnp=self.dbsnp, + model=hybrid_model, + pcr_indel_model="none", + emit_mode=emit_mode, + ) + ], + inputs=lr_aln + sr_aln, interval=self.bed, read_filter=rg_info.ultima_read_filter + rg_info.lr_rg_read_filter + rg_info.sr_rg_read_filter, - ) - driver.add_algo( - DNAscope( - combined_vcf, - dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("hybrid.model"), - pcr_indel_model="none", - emit_mode=emit_mode, - ) - ) - call_job = Job( - Pipeline(Command(*driver.build_cmd())), "calling-1", self.cores - ) + replace_rg=rg_info.replace_rg_args[0] + rg_info.replace_rg_args[1], + name="dnascope-1", + ).add_to(dag, upstream) # Region selection hybrid_select = pathlib.Path( @@ -782,27 +911,26 @@ def call_variants( ), "hybrid-select", 0, + task_name="region-selection", ) + dag.add_job(select_job, call.terminal) mapq0_bed = self.tmp_dir.joinpath("hybrid_mapq0.bed") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - replace_rg=rg_info.replace_rg_args[0] + rg_info.replace_rg_args[1], - input=lr_aln + sr_aln, + mapq0_job = driver_job( + ctx, + [ + HybridStage2( + model=bundle.joinpath("HybridStage2_region.model"), + all_bed=mapq0_bed, + ) + ], + name="mapq0-bed", + task_name="region-selection", + inputs=lr_aln + sr_aln, read_filter=rg_info.lr_rg_read_filter + rg_info.sr_rg_read_filter, + replace_rg=rg_info.replace_rg_args[0] + rg_info.replace_rg_args[1], ) - driver.add_algo( - HybridStage2( - model=self.model_bundle.joinpath("HybridStage2_region.model"), - all_bed=mapq0_bed, - ) - ) - mapq0_job = Job( - Pipeline(Command(*driver.build_cmd())), - "mapq0-bed", - self.cores, - ) + dag.add_job(mapq0_job, upstream) mapq0_slop_bed = self.tmp_dir.joinpath("hybrid_mapq0.ex1000.bed") mapq0_slop_job = Job( @@ -814,7 +942,9 @@ def call_variants( ), "mapq0-bed-slop", 0, + task_name="region-selection", ) + dag.add_job(mapq0_slop_job, {mapq0_job}) diff_bed = self.tmp_dir.joinpath("merged_diff.bed") cat_merge_job = Job( @@ -825,14 +955,20 @@ def call_variants( ), "concat-merge-bed", 0, + task_name="region-selection", + ) + dag.add_job(cat_merge_job, {mapq0_slop_job, select_job}) + self.add_cleanup( + dag, + [selected_bed, mapq0_slop_bed], + "rm-tmp1", + {cat_merge_job}, ) - rm_cmd = ["rm", str(selected_bed), str(mapq0_slop_bed)] - rm_job1 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp1", 0) stage1_ins_fa = self.tmp_dir.joinpath("stage1_ins.fa") stage1_ins_bed = self.tmp_dir.joinpath("stage1_ins.bed") ins_driver = Driver( - reference=self.reference, + reference=ctx.reference, thread_count=self.cores, replace_rg=rg_info.replace_rg_args[0], input=lr_aln, @@ -841,7 +977,7 @@ def call_variants( ins_driver.add_algo( HybridStage1( "-", - model=self.model_bundle.joinpath("HybridStage1_ins.model"), + model=bundle.joinpath("HybridStage1_ins.model"), fa_file=stage1_ins_fa, bed_file=stage1_ins_bed, ) @@ -850,8 +986,21 @@ def call_variants( stage1_hap_bam = self.tmp_dir.joinpath("stage1_hap.bam") stage1_hap_bed = self.tmp_dir.joinpath("stage1_hap.bed") stage1_hap_vcf = self.tmp_dir.joinpath("stage1_hap.vcf") + + # The algo writes its fastq output to a fifo read by the bwa job + # and its unsorted haplotype alignments to stdout + stage1_fifo = self.tmp_dir.joinpath("stage1_hap.fq") + stage1_fifo_job = Job( + Pipeline(Command("mkfifo", str(stage1_fifo))), + "stage1-fifo", + 1, + task_name="hybrid-realignment", + ) + # The haplotype job writes the fifo read by the first-stage job + dag.add_job(stage1_fifo_job) + stage1_driver = Driver( - reference=self.reference, + reference=ctx.reference, thread_count=self.cores, replace_rg=rg_info.replace_rg_args[0], input=lr_aln, @@ -860,115 +1009,135 @@ def call_variants( ) stage1_driver.add_algo( HybridStage1( - "-", - model=self.model_bundle.joinpath("HybridStage1.model"), - hap_bam=stage1_hap_bam, + stage1_fifo, + model=bundle.joinpath("HybridStage1.model"), + hap_bam="-", hap_bed=stage1_hap_bed, hap_vcf=stage1_hap_vcf, ) ) + stage1_hap_job = Job( + cmds.hybrid_stage1_hap( + stage1_hap_bam, + stage1_driver, + self.cores, + ), + "first-stage-hap", + # This job writes the fifo that the `first-stage` job reads, so + # the two need to run concurrently. Jobs requesting 0 threads + # start immediately instead of waiting for the thread budget; + # requesting `self.cores` here could let the scheduler serialize + # the two jobs and deadlock on the fifo. + 0, + task_name="hybrid-realignment", + ) + dag.add_job(stage1_hap_job, {stage1_fifo_job, cat_merge_job}) stage1_bam = self.tmp_dir.joinpath("hybrid_stage1.bam") stage1_job = Job( cmds.hybrid_stage1( out_aln=stage1_bam, - reference=self.reference, + reference=ctx.reference, cores=self.cores, readgroup=f"@RG\\tID:hybrid-18893\\tSM:{self.hybrid_rg_sm}", ins_driver=ins_driver, - stage1_driver=stage1_driver, - bwa_model=self.model_bundle.joinpath("HybridStage1_bwa.model"), + hap_fastq_fifo=stage1_fifo, + bwa_model=bundle.joinpath("HybridStage1_bwa.model"), ), "first-stage", self.cores, + task_name="hybrid-realignment", + ) + dag.add_job(stage1_job, {stage1_fifo_job, cat_merge_job}) + self.add_cleanup( + dag, + [stage1_ins_fa, stage1_ins_bed, stage1_hap_vcf], + "rm-tmp2", + {stage1_job, stage1_hap_job}, ) - rm_cmd = [ - "rm", - str(stage1_ins_fa), - str(stage1_ins_bed), - str(stage1_hap_vcf), - ] - rm_job2 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp2", 0) stage2_bed = self.tmp_dir.joinpath("hybrid_stage2.bed") stage2_unmap_bam = self.tmp_dir.joinpath("hybrid_stage2_unmap.bam") stage2_alt_bam = self.tmp_dir.joinpath("hybrid_stage2_alt.bam") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=[stage1_bam, stage1_hap_bam], - ) - driver.add_algo( - HybridStage2( - model=self.model_bundle.joinpath("HybridStage2.model"), - hap_bed=stage1_hap_bed, - unmap_bam=stage2_unmap_bam, - alt_bam=stage2_alt_bam, - all_bed=stage2_bed, - ) + second_stage_job = driver_job( + ctx, + [ + HybridStage2( + model=bundle.joinpath("HybridStage2.model"), + hap_bed=stage1_hap_bed, + unmap_bam=stage2_unmap_bam, + alt_bam=stage2_alt_bam, + all_bed=stage2_bed, + ) + ], + name="second-stage", + task_name="hybrid-realignment", + inputs=[stage1_bam, stage1_hap_bam], ) - second_stage_job = Job( - Pipeline(Command(*driver.build_cmd())), - "second-stage", - self.cores, + dag.add_job(second_stage_job, {stage1_job, stage1_hap_job}) + self.add_cleanup( + dag, + [stage1_bam, stage1_hap_bam], + "rm-tmp3", + {second_stage_job}, ) - rm_cmd = ["rm", str(stage1_bam), str(stage1_hap_bam)] - rm_job3 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp3", 0) - suffix = "bam" if self.bam_format else "cram" stage3_aln = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_sr_realigned.{suffix}") + str(ctx.output_vcf).replace(".vcf.gz", f"_sr_realigned.{suffix}") ) - driver = Driver( - reference=self.reference, + stage3_driver = Driver( + reference=ctx.reference, thread_count=self.cores, replace_rg=rg_info.replace_rg_args[0] + rg_info.replace_rg_args[1], input=lr_aln + sr_aln + [stage2_unmap_bam, stage2_alt_bam], interval=stage2_bed, read_filter=rg_info.lr_rg_read_filter + rg_info.sr_rg_read_filter, ) - driver.add_algo( + stage3_driver.add_algo( HybridStage3( "-", - model=self.model_bundle.joinpath("HybridStage3.model"), + model=bundle.joinpath("HybridStage3.model"), ) ) third_stage_job = Job( cmds.hybrid_stage3( stage3_aln, - reference=self.reference, - driver=driver, + reference=ctx.reference, + driver=stage3_driver, cores=self.cores, ), "third-stage", self.cores, + task_name="hybrid-realignment", + ) + dag.add_job(third_stage_job, {second_stage_job}) + self.add_cleanup( + dag, + [stage2_unmap_bam, stage2_alt_bam], + "rm-tmp4", + {third_stage_job}, ) - rm_cmd = ["rm", str(stage2_unmap_bam), str(stage2_alt_bam)] - rm_job4 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp4", 0) # pass 2 of variant calling pass2_vcf = self.tmp_dir.joinpath("hybrid_pass2.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - replace_rg=rg_info.replace_rg_args[0], - input=lr_aln + [stage3_aln], + call2 = DNAscopeStage( + ctx=ctx, + algos=[ + DNAscope( + pass2_vcf, + dbsnp=self.dbsnp, + model=hybrid_model, + pcr_indel_model="none", + emit_mode=emit_mode, + ) + ], + inputs=lr_aln + [stage3_aln], interval=stage2_bed, read_filter=rg_info.ultima_read_filter + rg_info.lr_rg_read_filter, - ) - driver.add_algo( - DNAscope( - pass2_vcf, - dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("hybrid.model"), - pcr_indel_model="none", - emit_mode=emit_mode, - ) - ) - call2_job = Job( - Pipeline(Command(*driver.build_cmd())), "call2", self.cores - ) + replace_rg=rg_info.replace_rg_args[0], + name="dnascope-2", + ).add_to(dag, {third_stage_job}) # Merge and normalize the VCFs subset_vcf = self.tmp_dir.joinpath("mix_subset.vcf.gz") @@ -981,7 +1150,9 @@ def call_variants( ), "subset-calls", 0, + task_name="vcf-merge", ) + dag.add_job(subset_job, {second_stage_job}) concat_job = Job( cmds.bcftools_concat( combined_tmp_vcf, @@ -989,9 +1160,15 @@ def call_variants( ), "concat-calls", 0, + task_name="vcf-merge", + ) + dag.add_job(concat_job, {subset_job} | call2.terminal) + self.add_cleanup( + dag, + [combined_vcf, subset_vcf, pass2_vcf], + "rm-tmp5", + {concat_job}, ) - rm_cmd = ["rm", str(combined_vcf), str(subset_vcf), str(pass2_vcf)] - rm_job5 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp5", 0) # Annotate the output VCF hybrid_anno = pathlib.Path( @@ -999,7 +1176,7 @@ def call_variants( ) anno_target = self.tmp_dir.joinpath("combined_tmp_anno.vcf.gz") if self.skip_model_apply and not self.pop_vcf: - anno_target = self.output_vcf + anno_target = ctx.output_vcf anno_job = Job( cmds.cmd_pyexec_hybrid_anno( @@ -1011,106 +1188,51 @@ def call_variants( ), "anno-calls", 0, + task_name="annotation", ) + dag.add_job(anno_job, {concat_job}) - transfer_jobs: Optional[List[Job]] = None - transfer_concat_job: Optional[Job] = None - input_to_apply = anno_target + if not self.pop_vcf and self.skip_model_apply: + # `anno_target` is already the output VCF + return + # Transfer annotations and apply the model + transfer_spec: Optional[TransferSpec] = None if self.pop_vcf: transfer_target = self.tmp_dir.joinpath( "combined_tmp_transfer.vcf.gz" ) if self.skip_model_apply: - transfer_target = self.output_vcf - - transfer_jobs, transfer_concat_job = build_transfer_jobs( - transfer_target, - self.pop_vcf, - anno_target, - self.tmp_dir, - self.shards, - self.pop_vcf_contigs, - self.fai_data, - self.dry_run, - self.cores, + transfer_target = ctx.output_vcf + transfer_spec = TransferSpec( + TransferConfig.from_pipeline(self), transfer_target ) - input_to_apply = transfer_target - if self.skip_model_apply: - return ( - call_job, - select_job, - mapq0_job, - mapq0_slop_job, - cat_merge_job, - rm_job1, - stage1_job, - rm_job2, - second_stage_job, - rm_job3, - third_stage_job, - rm_job4, - call2_job, - subset_job, - concat_job, - rm_job5, - anno_job, - transfer_jobs, - transfer_concat_job, - None, - None, - ) - - # Model Apply + apply_spec: Optional[ApplySpec] = None apply_vcf = self.tmp_dir.joinpath("combined_apply.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - interval=self.bed, - ) - driver.add_algo( - DNAModelApply( - model=self.model_bundle.joinpath("hybrid.model"), - vcf=input_to_apply, - output=apply_vcf, - ) - ) - apply_job = Job( - Pipeline(Command(*driver.build_cmd())), "model-apply", self.cores - ) + if not self.skip_model_apply: + apply_spec = ApplySpec(hybrid_model, apply_vcf) + + transfer_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=anno_target, + transfer=transfer_spec, + apply=apply_spec, + ).add_to(dag, {anno_job}) + + if self.skip_model_apply: + return # Final normalize norm_job = Job( cmds.filter_norm( - self.output_vcf, + ctx.output_vcf, apply_vcf, - self.reference, + ctx.reference, exclude_homref=not self.gvcf, ), "final-norm", 0, + task_name="vcf-norm", ) - return ( - call_job, - select_job, - mapq0_job, - mapq0_slop_job, - cat_merge_job, - rm_job1, - stage1_job, - rm_job2, - second_stage_job, - rm_job3, - third_stage_job, - rm_job4, - call2_job, - subset_job, - concat_job, - rm_job5, - anno_job, - transfer_jobs, - transfer_concat_job, - apply_job, - norm_job, - ) + dag.add_job(norm_job, transfer_apply.terminal) diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index c256cd0..032f653 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -4,10 +4,8 @@ import argparse import copy -from dataclasses import dataclass import json import pathlib -import shutil import sys from typing import Dict, List, Optional, Set, Tuple @@ -19,12 +17,9 @@ from .archive import ar_load from .dag import DAG from .driver import ( - Driver, + BaseAlgo, DNAscope, DNAscopeHP, - DNAModelApply, - LongReadSV, - ReadWriter, RepeatModel, VariantPhaser, ) @@ -36,14 +31,32 @@ vcf_contigs, ) from .shell_pipeline import Command, Pipeline -from .transfer import build_transfer_jobs +from .stages.alignment import ( + MINIMAP2_FASTQ_MIN_VERSIONS, + MINIMAP2_REALIGN_MIN_VERSIONS, + Minimap2FastqStage, + Minimap2RealignStage, + ReadWriterStage, + find_unzip, +) +from .stages.base import StageContext, driver_job +from .stages.metrics import MOSDEPTH_MIN_VERSIONS, MosdepthStage +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + TransferApplyStage, + TransferSpec, +) +from .stages.sv import LONGREADSV_MIN_VERSIONS, LongReadSVStage +from .stages.transfer import TransferConfig from .util import ( __version__, - check_version, + library_preloaded, parse_rg_line, path_arg, - library_preloaded, + require_versions, vcf_id, + versions_available, ) TOOL_MIN_VERSIONS = { @@ -52,23 +65,6 @@ "bedtools": None, } -SV_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308"), -} - -ALN_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308"), - "samtools": packaging.version.Version("1.16"), -} - -FQ_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308"), -} - -MOSDEPTH_MIN_VERSIONS = { - "mosdepth": packaging.version.Version("0.2.6"), -} - MERGE_MIN_VERSIONS = { "sentieon driver": None, } @@ -82,39 +78,6 @@ } -@dataclass -class LRCallVariantsResult: - """Results from lr_call_variants""" - - first_calling_job: Job - diploid_transfer_jobs: List[Job] - diploid_concat_job: Optional[Job] - first_modelapply_job: Job - phaser_job: Job - bcftools_subset_phased_job: Optional[Job] - fai_to_bed_job: Optional[Job] - bcftools_subtract_job: Job - repeatmodel_job: Optional[Job] - bcftools_subset_unphased_job: Job - second_calling_job: Set[Job] - haploid_patch_job: Job - patch_transfer_jobs: List[List[Job]] - patch_concat_jobs: List[Job] - second_modelapply_job: Set[Job] - calling_unphased_job: Job - diploid_patch_job: Job - unphased_transfer_jobs: List[Job] - unphased_concat_job: Optional[Job] - modelapply_unphased_job: Job - merge_job: Job - gvcf_combine_job: Optional[Job] - haploid_calling_job: Optional[Job] - haploid_patch2_job: Optional[Job] - haploid_concat_job: Optional[Job] - haploid_gvcf_combine_job: Optional[Job] - haploid_gvcf_concat_job: Optional[Job] - - class DNAscopeLRPipeline(BasePipeline): """The DNAscope LongRead pipeline""" @@ -255,7 +218,7 @@ class DNAscopeLRPipeline(BasePipeline): "util_sort_args": { # "help": "Extra arguments for sentieon util sort.", "help": argparse.SUPPRESS, - "default": "--cram_write_options version=3.0,compressor=rans", + "default": "", }, } ) @@ -282,21 +245,14 @@ def __init__(self) -> None: self.fastq_taglist = "*" self.bam_format = False self.minimap2_args = "-YL" - self.util_sort_args = ( - "--cram_write_options version=3.0,compressor=rans" - ) + self.util_sort_args = "" self.repeat_model: Optional[pathlib.Path] = None self.skip_pop_vcf_id_check = False self.use_pbsv = False def validate(self) -> None: - # uniquify pipeline attributes - self.lr_aln = self.sample_input - self.lr_input_ref = self.input_ref - del self.sample_input - # validate - if not (self.lr_aln or self.fastq): + if not (self.sample_input or self.fastq): self.logger.error( "Please supply either the `--sample_input` or `--fastq` and " "`--readgroups` arguments" @@ -347,7 +303,7 @@ def validate(self) -> None: self.fai_data = parse_fai(pathlib.Path(str(self.reference) + ".fai")) self.pop_vcf_contigs: Dict[str, Optional[int]] = {} if self.pop_vcf: - self.pop_vcf_contigs = vcf_contigs(self.pop_vcf, self.dry_run) + self.pop_vcf_contigs = vcf_contigs(self.pop_vcf) self.logger.debug("VCF contigs are: %s", self.pop_vcf_contigs) self.shards = determine_shards_from_fai( self.fai_data, 10 * 1000 * 1000 @@ -468,8 +424,9 @@ def configure(self) -> None: def build_dag(self) -> DAG: self.logger.info("Building the DAG") dag = DAG() + ctx = self.stage_context() - sample_input = self.lr_aln + sample_input = self.sample_input realign_jobs: Set[Job] = set() if self.align: sample_input, realign_jobs = self.lr_align_inputs() @@ -500,122 +457,10 @@ def build_dag(self) -> DAG: dag.add_job(hificnv_job, {merge_job}) if not self.skip_small_variants: - cv = self.lr_call_variants(sample_input) - dag.add_job(cv.first_calling_job, realign_jobs.union(align_jobs)) - - first_ma_deps = {cv.first_calling_job} - if cv.diploid_transfer_jobs and cv.diploid_concat_job: - for job in cv.diploid_transfer_jobs: - dag.add_job(job, {cv.first_calling_job}) - dag.add_job( - cv.diploid_concat_job, set(cv.diploid_transfer_jobs) - ) - first_ma_deps.add(cv.diploid_concat_job) - - dag.add_job(cv.first_modelapply_job, first_ma_deps) - dag.add_job(cv.phaser_job, {cv.first_modelapply_job}) - - haploid_patch_deps: Set[Job] = set() - if cv.bcftools_subset_phased_job: - dag.add_job(cv.bcftools_subset_phased_job, {cv.phaser_job}) - haploid_patch_deps.add(cv.bcftools_subset_phased_job) - - subtract_deps: Set[Job] = {cv.phaser_job} - if cv.fai_to_bed_job: - dag.add_job(cv.fai_to_bed_job) - subtract_deps.add(cv.fai_to_bed_job) - dag.add_job(cv.bcftools_subtract_job, subtract_deps) - dag.add_job( - cv.bcftools_subset_unphased_job, - {cv.bcftools_subtract_job}, - ) - - second_pass_deps: Set[Job] = {cv.phaser_job} - calling_unphased_deps: Set[Job] = {cv.bcftools_subtract_job} - haploid_calling_deps: Set[Job] = set() - if cv.repeatmodel_job: - dag.add_job(cv.repeatmodel_job, {cv.phaser_job}) - second_pass_deps.add(cv.repeatmodel_job) - calling_unphased_deps.add(cv.repeatmodel_job) - haploid_calling_deps.add(cv.repeatmodel_job) - - merge_deps: Set[Job] = set() - for job in cv.second_calling_job: - dag.add_job(job, second_pass_deps) - haploid_patch_deps.add(job) - dag.add_job(cv.haploid_patch_job, haploid_patch_deps) - - second_ma_deps = {cv.haploid_patch_job} - if cv.patch_transfer_jobs: # pop_vcf - for patch_transfer_hap, patch_concat in zip( - cv.patch_transfer_jobs, cv.patch_concat_jobs - ): - for job in patch_transfer_hap: - dag.add_job(job, {cv.haploid_patch_job}) - dag.add_job(patch_concat, set(patch_transfer_hap)) - second_ma_deps.add(patch_concat) - - for job in cv.second_modelapply_job: - dag.add_job(job, second_ma_deps) - merge_deps.add(job) - - dag.add_job(cv.calling_unphased_job, calling_unphased_deps) - dag.add_job( - cv.diploid_patch_job, - { - cv.bcftools_subset_unphased_job, - cv.calling_unphased_job, - }, + self.add_small_variant_calling( + dag, ctx, sample_input, realign_jobs.union(align_jobs) ) - unphased_ma_deps = {cv.diploid_patch_job} - if cv.unphased_transfer_jobs and cv.unphased_concat_job: # pop_vcf - for job in cv.unphased_transfer_jobs: - dag.add_job(job, {cv.diploid_patch_job}) - dag.add_job( - cv.unphased_concat_job, set(cv.unphased_transfer_jobs) - ) - unphased_ma_deps.add(cv.unphased_concat_job) - - dag.add_job(cv.modelapply_unphased_job, unphased_ma_deps) - merge_deps.add(cv.modelapply_unphased_job) - dag.add_job(cv.merge_job, merge_deps) - - if cv.gvcf_combine_job: - dag.add_job(cv.gvcf_combine_job, {cv.merge_job}) - - if ( - cv.haploid_calling_job - and cv.haploid_patch2_job - and cv.haploid_concat_job - ): - dag.add_job(cv.haploid_calling_job, haploid_calling_deps) - dag.add_job( - cv.haploid_patch2_job, - {cv.haploid_calling_job}, - ) - dag.add_job( - cv.haploid_concat_job, - {cv.haploid_patch2_job, cv.merge_job}, - ) - - if ( - cv.haploid_gvcf_combine_job - and cv.haploid_gvcf_concat_job - and cv.gvcf_combine_job - ): - dag.add_job( - cv.haploid_gvcf_combine_job, - {cv.haploid_patch2_job}, - ) - dag.add_job( - cv.haploid_gvcf_concat_job, - { - cv.haploid_gvcf_combine_job, - cv.gvcf_combine_job, - }, - ) - if not self.skip_svs: longreadsv_job = self.call_svs(sample_input) dag.add_job(longreadsv_job, realign_jobs.union(align_jobs)) @@ -626,147 +471,69 @@ def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: """ Align reads to the reference genome using minimap2 """ - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.skip_version_check: - for cmd, min_version in ALN_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) - - res: List[pathlib.Path] = [] - suffix = "bam" if self.bam_format else "cram" - sample_name = self.output_vcf.name.replace(".vcf.gz", "") - realign_jobs = set() - for i, input_aln in enumerate(self.lr_aln): - out_aln = pathlib.Path( - str(self.output_vcf).replace( - ".vcf.gz", f"_mm2_sorted_{i}.{suffix}" - ) - ) - rg_lines = cmds.get_rg_lines( - input_aln, - self.dry_run, - ) + output_vcf = self.required(self.output_vcf, "output_vcf") + bundle = self.required(self.model_bundle, "model_bundle") + require_versions( + MINIMAP2_REALIGN_MIN_VERSIONS, skip=self.skip_version_check + ) - realign_jobs.add( - Job( - cmds.cmd_samtools_fastq_minimap2( - out_aln, - input_aln, - self.reference, - self.model_bundle, - self.cores, - rg_lines, - sample_name, - self.lr_input_ref, - self.fastq_taglist, - self.minimap2_args, - self.util_sort_args, - ), - f"bam-realign-{i}", - self.cores, - ) - ) - res.append(out_aln) - return (res, realign_jobs) + result = Minimap2RealignStage( + ctx=self.stage_context(), + inputs=self.sample_input, + model_bundle=bundle, + sample_name=output_vcf.name.replace(".vcf.gz", ""), + bam_format=self.bam_format, + input_ref=self.input_ref, + fastq_taglist=self.fastq_taglist, + minimap2_args=self.minimap2_args, + util_sort_args=self.util_sort_args, + ).build() + return (result.outputs, set(result.jobs)) def lr_align_fastq(self) -> Tuple[List[pathlib.Path], Set[Job]]: """ Align fastq to the reference genome using minimap2 """ - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - res: List[pathlib.Path] = [] + bundle = self.required(self.model_bundle, "model_bundle") if self.fastq is None and self.readgroups is None: - return (res, set()) + return ([], set()) - if not self.skip_version_check: - for cmd, min_version in FQ_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions( + MINIMAP2_FASTQ_MIN_VERSIONS, skip=self.skip_version_check + ) - unzip = "igzip" - if not shutil.which(unzip): - self.logger.warning( - "igzip is recommended for decompression, but is not " - "available. Falling back to gzip." - ) - unzip = "gzip" - - suffix = "bam" if self.bam_format else "cram" - align_jobs = set() - for i, (fq, rg) in enumerate(zip(self.fastq, self.readgroups)): - out_aln = pathlib.Path( - str(self.output_vcf).replace( - ".vcf.gz", f"_mm2_sorted_fq_{i}.{suffix}" - ) - ) - align_jobs.add( - Job( - cmds.cmd_fastq_minimap2( - out_aln, - fq, - rg, - self.reference, - self.model_bundle, - self.cores, - unzip, - self.minimap2_args, - self.util_sort_args, - ), - "align-{i}", - self.cores, - ) - ) - res.append(out_aln) - return (res, align_jobs) + result = Minimap2FastqStage( + ctx=self.stage_context(), + fastq=self.fastq, + readgroups=self.readgroups, + model_bundle=bundle, + bam_format=self.bam_format, + unzip=find_unzip(self.logger), + minimap2_args=self.minimap2_args, + util_sort_args=self.util_sort_args, + ).build() + return (result.outputs, set(result.jobs)) def mosdepth(self, sample_input: List[pathlib.Path]) -> Set[Job]: """Run mosdepth for QC""" - if not self.skip_version_check: - if not all( - [ - check_version(cmd, min_version) - for (cmd, min_version) in MOSDEPTH_MIN_VERSIONS.items() - ] - ): - self.logger.warning( - "Skipping mosdepth. mosdepth version %s or later " - "not found", - MOSDEPTH_MIN_VERSIONS["mosdepth"], - ) - return set() - - mosdepth_jobs = set() - for i, input_file in enumerate(sample_input): - mosdepth_dir = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_mosdepth_{i}") + if not versions_available( + MOSDEPTH_MIN_VERSIONS, skip=self.skip_version_check + ): + self.logger.warning( + "Skipping mosdepth. mosdepth version %s or later not found", + MOSDEPTH_MIN_VERSIONS["mosdepth"], ) - mosdepth_jobs.add( - Job( - cmds.cmd_mosdepth( - input_file, - mosdepth_dir, - fasta=self.reference, - threads=self.cores, - ), - f"mosdepth-{i}", - 0, # Run in background - ) + return set() + + return set( + MosdepthStage( + ctx=self.stage_context(), + inputs=sample_input, ) - return mosdepth_jobs + .build() + .jobs + ) def merge_input_files( self, sample_input: List[pathlib.Path] @@ -774,25 +541,19 @@ def merge_input_files( """ Merge the aligned reads into a single BAM """ - if not self.skip_version_check: - for cmd, min_version in MERGE_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(MERGE_MIN_VERSIONS, skip=self.skip_version_check) # Merge the sample_input into a single BAM merged_bam = self.tmp_dir.joinpath("longread_merged.bam") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - ) - driver.add_algo(ReadWriter(merged_bam)) - merge_job = Job( - Pipeline(Command(*driver.build_cmd())), - "merge-bam", - 0, - ) - return (merged_bam, merge_job) + result = ReadWriterStage( + ctx=self.stage_context(), + inputs=sample_input, + output=merged_bam, + name="merge-bam", + task_name="alignment", + threads=0, # Run in the background + ).build() + return (merged_bam, result.jobs[0]) def pbsv( self, @@ -801,10 +562,7 @@ def pbsv( """ Call SVs with pbsv """ - if not self.skip_version_check: - for cmd, min_version in PBSV_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(PBSV_MIN_VERSIONS, skip=self.skip_version_check) # pbsv discover pbsv_discover_fn = str(self.output_vcf).replace( @@ -823,6 +581,7 @@ def pbsv( ), "pbsv-discover", 0, + task_name="sv-calling", ) # pbsv call @@ -840,6 +599,7 @@ def pbsv( Pipeline(Command(*call_cmd)), "pbsv-call", self.cores, + task_name="sv-calling", ) return (pbsv_discover, pbsv_call) @@ -850,15 +610,14 @@ def hificnv( """ Call CNVs with HiFiCNV """ - if not self.skip_version_check: - for cmd, min_version in HIFICNV_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - self.logger.warning( - "Skipping hificnv. hificnv version %s or later not " - "found", - HIFICNV_MIN_VERSIONS["hificnv"], - ) - return None + if not versions_available( + HIFICNV_MIN_VERSIONS, skip=self.skip_version_check + ): + self.logger.warning( + "Skipping hificnv. hificnv version %s or later not found", + HIFICNV_MIN_VERSIONS["hificnv"], + ) + return None hifi_cnv_fn = str(self.output_vcf).replace(".vcf.gz", ".hificnv") hificnv_cmd = [ @@ -875,102 +634,79 @@ def hificnv( if self.cnv_excluded_regions: hificnv_cmd.extend(["--exclude", str(self.cnv_excluded_regions)]) hificnv_job = Job( - Pipeline(Command(*hificnv_cmd)), "hificnv", self.cores + Pipeline(Command(*hificnv_cmd)), + "hificnv", + self.cores, + task_name="cnv", ) return hificnv_job - def lr_call_variants( + def add_small_variant_calling( self, + dag: DAG, + ctx: StageContext, sample_input: List[pathlib.Path], - ) -> LRCallVariantsResult: + upstream: Set[Job], + ) -> None: """ Call SNVs and indels using the DNAscope LongRead pipeline + + A diploid first pass phases the reads; the phased regions are then + called once per haplotype and the unphased regions are called with + the diploid HP model. The three call sets are patched, filtered and + merged into the output VCF, with optional haploid calling on top. """ - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - if not self.skip_version_check: - for check_cmd, min_version in TOOL_MIN_VERSIONS.items(): - if not check_version(check_cmd, min_version): - sys.exit(2) - - transfer_xargs = ( - self.tmp_dir, - self.shards, - self.pop_vcf_contigs, - self.fai_data, - self.dry_run, - self.cores, - ) + bundle = self.required(self.model_bundle, "model_bundle") + require_versions(TOOL_MIN_VERSIONS, skip=self.skip_version_check) # First pass - diploid calling diploid_gvcf_fn = self.tmp_dir.joinpath("out_diploid.g.vcf.gz") diploid_tmp_vcf = self.tmp_dir.joinpath("out_diploid_tmp.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - interval=self.bed, - ) + diploid_model = bundle.joinpath("diploid_model") + diploid_algos: List[BaseAlgo] = [] if self.gvcf: - driver.add_algo( + diploid_algos.append( DNAscope( diploid_gvcf_fn, dbsnp=self.dbsnp, emit_mode="gvcf", - model=self.model_bundle.joinpath("gvcf_model"), + model=bundle.joinpath("gvcf_model"), ) ) - driver.add_algo( + diploid_algos.append( DNAscope( diploid_tmp_vcf, dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("diploid_model"), + model=diploid_model, ) ) - first_calling_job = Job( - Pipeline(Command(*driver.build_cmd())), "first-pass", self.cores - ) - - # Transfer annotations to the tmp vcf - diploid_transfer_jobs: List[Job] = [] - diploid_concat_job: Optional[Job] = None - if self.pop_vcf: - diploid_unanno_tmp_vcf = diploid_tmp_vcf - diploid_tmp_vcf = self.tmp_dir.joinpath( - "out_diploid_tmp_anno.vcf.gz" - ) - diploid_transfer_jobs, diploid_concat_job = build_transfer_jobs( - diploid_tmp_vcf, - self.pop_vcf, - diploid_unanno_tmp_vcf, - *transfer_xargs, - ) + first_call = DNAscopeStage( + ctx=ctx, + algos=diploid_algos, + inputs=sample_input, + interval=self.bed, + name="dnascope-diploid", + ).add_to(dag, upstream) - # DNAModelApply on the diploid VCF + # Transfer annotations to the tmp vcf, then apply the diploid model diploid_vcf = self.tmp_dir.joinpath("out_diploid.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - driver.add_algo( - DNAModelApply( - self.model_bundle.joinpath("diploid_model"), - diploid_tmp_vcf, - diploid_vcf, + diploid_transfer: Optional[TransferSpec] = None + if self.pop_vcf: + diploid_transfer = TransferSpec( + config=TransferConfig.from_pipeline(self), + out_vcf=self.tmp_dir.joinpath("out_diploid_tmp_anno.vcf.gz"), + tag="diploid", ) - ) - first_modelapply_job = Job( - Pipeline(Command(*driver.build_cmd())), - "first-modelapply", - self.cores, - ) + first_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=diploid_tmp_vcf, + transfer=diploid_transfer, + apply=ApplySpec( + model=diploid_model, + output=diploid_vcf, + name="model-apply-diploid", + ), + ).add_to(dag, first_call.terminal) # Phasing and RepeatModel phased_bed = self.tmp_dir.joinpath("out_diploid_phased.bed") @@ -983,25 +719,24 @@ def lr_call_variants( phased_phased = self.tmp_dir.joinpath( "out_diploid_phased_phased.vcf.gz" ) - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, + phaser_job = driver_job( + ctx, + [ + VariantPhaser( + diploid_vcf, + phased_vcf, + out_bed=phased_bed, + out_ext=phased_ext, + ) + ], + name="variantphaser", + task_name="phasing", + inputs=sample_input, interval=self.bed, ) - driver.add_algo( - VariantPhaser( - diploid_vcf, - phased_vcf, - out_bed=phased_bed, - out_ext=phased_ext, - ) - ) - phaser_job = Job( - Pipeline(Command(*driver.build_cmd())), "variantphaser", self.cores - ) + dag.add_job(phaser_job, first_apply.terminal) - bcftools_subset_phased_job = None + haploid_patch_deps: Set[Job] = set() if self.tech.upper() == "ONT": bcftools_subset_phased_job = Job( Pipeline( @@ -1022,53 +757,35 @@ def lr_call_variants( ), "bcftools-subset-phased", 0, + task_name="phasing", ) + dag.add_job(bcftools_subset_phased_job, {phaser_job}) + haploid_patch_deps.add(bcftools_subset_phased_job) - fai_to_bed_job = None + subtract_deps: Set[Job] = {phaser_job} if self.bed: bed = self.bed else: bed = self.tmp_dir.joinpath("reference.bed") fai_to_bed_job = Job( cmds.cmd_fai_to_bed( - pathlib.Path(str(self.reference) + ".fai"), + pathlib.Path(str(ctx.reference) + ".fai"), bed, ), "fai-to-bed", 0, + task_name="phasing", ) + dag.add_job(fai_to_bed_job) + subtract_deps.add(fai_to_bed_job) bcftools_subtract_job = Job( cmds.cmd_bedtools_subtract(bed, phased_bed, unphased_bed), "bedtools-subtract", 0, + task_name="phasing", ) - - repeatmodel_job = None - if not self.repeat_model: - self.repeat_model = self.tmp_dir.joinpath("out_repeat.model") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - interval=phased_bed, - read_filter=[ - f"PhasedReadFilter,phased_vcf={phased_ext}," - "phase_select=tag" - ], # noqa: E501 - ) - driver.add_algo( - RepeatModel( - self.repeat_model, - phased=True, - read_flag_mask="drop=supplementary", - ) - ) - repeatmodel_job = Job( - Pipeline(Command(*driver.build_cmd())), - "repeatmodel", - self.cores, - ) + dag.add_job(bcftools_subtract_job, subtract_deps) bcftools_subset_unphased_job = Job( Pipeline( @@ -1085,50 +802,72 @@ def lr_call_variants( ), "bcftools-subset-unphased", 0, + task_name="phasing", ) + dag.add_job(bcftools_subset_unphased_job, {bcftools_subtract_job}) + + repeat_model_deps: Set[Job] = set() + if not self.repeat_model: + self.repeat_model = self.tmp_dir.joinpath("out_repeat.model") + repeatmodel_job = driver_job( + ctx, + [ + RepeatModel( + self.repeat_model, + phased=True, + read_flag_mask="drop=supplementary", + ) + ], + name="repeatmodel", + task_name="repeat-model", + inputs=sample_input, + interval=phased_bed, + read_filter=[ + f"PhasedReadFilter,phased_vcf={phased_ext}," + "phase_select=tag" + ], # noqa: E501 + ) + dag.add_job(repeatmodel_job, {phaser_job}) + repeat_model_deps.add(repeatmodel_job) # Second pass - phased variants - second_calling_job = set() + haploid_model = bundle.joinpath("haploid_model") + haploid_hp_model = bundle.joinpath("haploid_hp_model") for phase in (1, 2): hp_std_vcf = self.tmp_dir.joinpath( f"out_hap{phase}_nohp_tmp.vcf.gz" ) hp_vcf = self.tmp_dir.joinpath(f"out_hap{phase}_tmp.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - interval=phased_bed, - read_filter=[ - f"PhasedReadFilter,phased_vcf={phased_ext}" - f",phase_select={phase}" - ], - ) - + phase_algos: List[BaseAlgo] = [] if self.tech.upper() == "HIFI": # ONT doesn't do DNAscope in 2nd pass. - driver.add_algo( + phase_algos.append( DNAscope( hp_std_vcf, dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("haploid_model"), + model=haploid_model, ) ) - driver.add_algo( + phase_algos.append( DNAscopeHP( hp_vcf, dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("haploid_hp_model"), + model=haploid_hp_model, pcr_indel_model=self.repeat_model, ) ) - second_calling_job.add( - Job( - Pipeline(Command(*driver.build_cmd())), - "second-pass", - self.cores, - ) - ) + hap_call = DNAscopeStage( + ctx=ctx, + algos=phase_algos, + inputs=sample_input, + interval=phased_bed, + read_filter=[ + f"PhasedReadFilter,phased_vcf={phased_ext}" + f",phase_select={phase}" + ], + name=f"dnascope-hap{phase}", + ).add_to(dag, {phaser_job} | repeat_model_deps) + haploid_patch_deps.update(hap_call.terminal) kwargs: Dict[str, str] = dict() kwargs["gvcf_combine_py"] = str( @@ -1153,129 +892,108 @@ def lr_call_variants( ), "patch", self.cores, + task_name="variant-patch", ) + dag.add_job(haploid_patch_job, haploid_patch_deps) - # Transfer annotations to the patched VCFs - patch_transfer_jobs = [] - patch_concat_jobs = [] - if self.pop_vcf: - patch_unanno_vcfs = patch_vcfs.copy() - patch_vcfs = [ - self.tmp_dir.joinpath(f"out_hap{i}_patch_anno.vcf.gz") - for i in (1, 2) - ] - for patch_unanno_vcf, patch_vcf in zip( - patch_unanno_vcfs, patch_vcfs - ): - transfer_jobs, concat_job = build_transfer_jobs( - patch_vcf, - self.pop_vcf, - patch_unanno_vcf, - *transfer_xargs, - ) - patch_transfer_jobs.append(transfer_jobs) - patch_concat_jobs.append(concat_job) - - # apply trained model to the patched vcfs. + # Transfer annotations to each patched VCF, then apply the trained + # model. A haplotype's model-apply reads only that haplotype's + # VCF, so it waits on that haplotype's transfer alone. hap_vcfs = [ self.tmp_dir.joinpath(f"out_hap{i}.vcf.gz") for i in (1, 2) ] - second_modelapply_job = set() - for patch_vcf, hap_vcf in zip(patch_vcfs, hap_vcfs): - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - driver.add_algo( - DNAModelApply( - self.model_bundle.joinpath("haploid_model"), - patch_vcf, - hap_vcf, - ) - ) - second_modelapply_job.add( - Job( - Pipeline(Command(*driver.build_cmd())), - "second-modelapply", - self.cores, + merge_deps: Set[Job] = set() + for i, (patch_vcf, hap_vcf) in enumerate( + zip(patch_vcfs, hap_vcfs), start=1 + ): + patch_transfer: Optional[TransferSpec] = None + if self.pop_vcf: + patch_transfer = TransferSpec( + config=TransferConfig.from_pipeline(self), + out_vcf=self.tmp_dir.joinpath( + f"out_hap{i}_patch_anno.vcf.gz" + ), + tag=f"hap{i}", ) - ) + hap_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=patch_vcf, + transfer=patch_transfer, + apply=ApplySpec( + model=haploid_model, + output=hap_vcf, + name=f"model-apply-hap{i}", + ), + ).add_to(dag, {haploid_patch_job}) + merge_deps.update(hap_apply.terminal) # Second pass - unphased regions diploid_unphased_hp = self.tmp_dir.joinpath( "out_diploid_phased_unphased_hp.vcf.gz" ) - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, + unphased_call = DNAscopeStage( + ctx=ctx, + algos=[ + DNAscopeHP( + diploid_unphased_hp, + dbsnp=self.dbsnp, + model=bundle.joinpath("diploid_hp_model"), + pcr_indel_model=self.repeat_model, + ) + ], + inputs=sample_input, interval=unphased_bed, - ) - driver.add_algo( - DNAscopeHP( - diploid_unphased_hp, - dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("diploid_hp_model"), - pcr_indel_model=self.repeat_model, - ) - ) - calling_unphased_job = Job( - Pipeline(Command(*driver.build_cmd())), - "calling-unphased", - self.cores, - ) + name="dnascope-unphased", + ).add_to(dag, {bcftools_subtract_job} | repeat_model_deps) # Patch DNA and DNAHP variants diploid_unphased_patch = self.tmp_dir.joinpath( "out_diploid_unphased_patch.vcf.gz" ) - cmd = cmds.cmd_pyexec_vcf_mod_patch( - str(diploid_unphased_patch), - str(phased_unphased), - str(diploid_unphased_hp), + diploid_patch_job = Job( + cmds.cmd_pyexec_vcf_mod_patch( + str(diploid_unphased_patch), + str(phased_unphased), + str(diploid_unphased_hp), + self.cores, + kwargs, + ), + "diploid-patch", self.cores, - kwargs, + task_name="variant-patch", + ) + dag.add_job( + diploid_patch_job, + {bcftools_subset_unphased_job} | unphased_call.terminal, ) - diploid_patch_job = Job(cmd, "diploid-patch", self.cores) - - # Transfer annotations to the diploid VCF - unphased_transfer_jobs: List[Job] = [] - unphased_concat_job: Optional[Job] = None - if self.pop_vcf: - diploid_unphased_unanno = diploid_unphased_patch - diploid_unphased_patch = self.tmp_dir.joinpath( - "out_diploid_unphased_patch_anno.vcf.gz" - ) - unphased_transfer_jobs, unphased_concat_job = build_transfer_jobs( - diploid_unphased_patch, - self.pop_vcf, - diploid_unphased_unanno, - *transfer_xargs, - ) - # DNAModelApply on the unphased VCF + # Transfer annotations to the diploid VCF, then apply the model diploid_unphased = self.tmp_dir.joinpath("out_diploid_unphased.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - driver.add_algo( - DNAModelApply( - self.model_bundle.joinpath("diploid_model_unphased"), - diploid_unphased_patch, - diploid_unphased, + unphased_transfer: Optional[TransferSpec] = None + if self.pop_vcf: + unphased_transfer = TransferSpec( + config=TransferConfig.from_pipeline(self), + out_vcf=self.tmp_dir.joinpath( + "out_diploid_unphased_patch_anno.vcf.gz" + ), + tag="unphased", ) - ) - modelapply_unphased_job = Job( - Pipeline(Command(*driver.build_cmd())), - "modelapply-unphased", - self.cores, - ) + unphased_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=diploid_unphased_patch, + transfer=unphased_transfer, + apply=ApplySpec( + model=bundle.joinpath("diploid_model_unphased"), + output=diploid_unphased, + name="model-apply-unphased", + ), + ).add_to(dag, {diploid_patch_job}) + merge_deps.update(unphased_apply.terminal) # merge calls to create the output diploid_merged_vcf = self.tmp_dir.joinpath("out_diploid_merged.vcf.gz") merge_out_vcf = ( - diploid_merged_vcf if self.haploid_bed else self.output_vcf + diploid_merged_vcf if self.haploid_bed else ctx.output_vcf ) merge_job = Job( cmds.cmd_pyexec_vcf_mod_merge( @@ -1290,13 +1008,15 @@ def lr_call_variants( ), "merge", self.cores, + task_name="vcf-merge", ) + dag.add_job(merge_job, merge_deps) gvcf_combine_job = None if self.gvcf: gvcf_combine_job = Job( cmds.cmd_pyexec_gvcf_combine( - self.reference, + ctx.reference, str(diploid_gvcf_fn), str(merge_out_vcf), self.cores, @@ -1304,168 +1024,129 @@ def lr_call_variants( ), "gvcf-combine", 0, + task_name="gvcf", ) + dag.add_job(gvcf_combine_job, {merge_job}) - haploid_calling_job = None - haploid_patch2_job = None - haploid_concat_job = None - haploid_gvcf_combine_job = None - haploid_gvcf_concat_job = None - if self.haploid_bed: - # Haploid variant calling - haploid_fn = self.tmp_dir.joinpath("haploid.vcf.gz") - haploid_gvcf_fn = self.tmp_dir.joinpath("haploid.g.vcf.gz") - haploid_hp_fn = self.tmp_dir.joinpath("haploid_hp.vcf.gz") - haploid_out_fn = self.tmp_dir.joinpath("haploid_patched.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - interval=self.haploid_bed, - ) - driver.add_algo( + if not self.haploid_bed: + return + + # Haploid variant calling. This job depends only on the repeat + # model, not on the alignment jobs. + haploid_fn = self.tmp_dir.joinpath("haploid.vcf.gz") + haploid_gvcf_fn = self.tmp_dir.joinpath("haploid.g.vcf.gz") + haploid_hp_fn = self.tmp_dir.joinpath("haploid_hp.vcf.gz") + haploid_out_fn = self.tmp_dir.joinpath("haploid_patched.vcf.gz") + haploid_algos: List[BaseAlgo] = [ + DNAscope( + haploid_fn, + dbsnp=self.dbsnp, + model=haploid_model, + ), + DNAscopeHP( + haploid_hp_fn, + dbsnp=self.dbsnp, + model=haploid_hp_model, + pcr_indel_model=self.repeat_model, + ), + ] + if self.gvcf: + haploid_algos.append( DNAscope( - haploid_fn, + haploid_gvcf_fn, dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("haploid_model"), - ) - ) - driver.add_algo( - DNAscopeHP( - haploid_hp_fn, - dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("haploid_hp_model"), - pcr_indel_model=self.repeat_model, + emit_mode="gvcf", + ploidy=1, + model=bundle.joinpath("gvcf_model"), ) ) - if self.gvcf: - driver.add_algo( - DNAscope( - haploid_gvcf_fn, - dbsnp=self.dbsnp, - emit_mode="gvcf", - ploidy=1, - model=self.model_bundle.joinpath("gvcf_model"), - ) - ) - haploid_calling_job = Job( - Pipeline(Command(*driver.build_cmd())), - "haploid-calling", + haploid_call = DNAscopeStage( + ctx=ctx, + algos=haploid_algos, + inputs=sample_input, + interval=self.haploid_bed, + name="dnascope-haploid", + ).add_to(dag, repeat_model_deps) + + haploid_patch2_job = Job( + cmds.cmd_pyexec_vcf_mod_haploid_patch2( + str(haploid_out_fn), + str(haploid_fn), + str(haploid_hp_fn), self.cores, - ) + kwargs, + ), + "haploid-patch2", + self.cores, + task_name="variant-patch", + ) + dag.add_job(haploid_patch2_job, haploid_call.terminal) - haploid_patch2_job = Job( - cmds.cmd_pyexec_vcf_mod_haploid_patch2( - str(haploid_out_fn), - str(haploid_fn), - str(haploid_hp_fn), - self.cores, - kwargs, - ), - "haploid-patch2", + haploid_concat_job = Job( + cmds.bcftools_concat( + ctx.output_vcf, + [diploid_merged_vcf, haploid_out_fn], + ), + "haploid-diploid-concat", + 0, + task_name="vcf-merge", + ) + dag.add_job(haploid_concat_job, {haploid_patch2_job, merge_job}) + + if not self.gvcf or gvcf_combine_job is None: + return + + output_gvcf = pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", ".g.vcf.gz") + ) + diploid_gvcf = pathlib.Path( + str(diploid_merged_vcf).replace(".vcf.gz", ".g.vcf.gz") + ) + haploid_gvcf = pathlib.Path( + str(haploid_out_fn).replace(".vcf.gz", ".g.vcf.gz") + ) + haploid_gvcf_combine_job = Job( + cmds.cmd_pyexec_gvcf_combine( + ctx.reference, + str(haploid_gvcf_fn), + str(haploid_out_fn), self.cores, - ) - haploid_concat_job = Job( - cmds.bcftools_concat( - self.output_vcf, - [diploid_merged_vcf, haploid_out_fn], - ), - "haploid-diploid-concat", - 0, - ) + kwargs, + ), + "haploid-gvcf-combine", + 0, + task_name="gvcf", + ) + dag.add_job(haploid_gvcf_combine_job, {haploid_patch2_job}) - if self.gvcf: - output_gvcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", ".g.vcf.gz") - ) - diploid_gvcf = pathlib.Path( - str(diploid_merged_vcf).replace(".vcf.gz", ".g.vcf.gz") - ) - haploid_gvcf = pathlib.Path( - str(haploid_out_fn).replace(".vcf.gz", ".g.vcf.gz") - ) - haploid_gvcf_combine_job = Job( - cmds.cmd_pyexec_gvcf_combine( - self.reference, - str(haploid_gvcf_fn), - str(haploid_out_fn), - self.cores, - kwargs, - ), - "haploid-gvcf-combine", - 0, - ) - haploid_gvcf_concat_job = Job( - cmds.bcftools_concat( - output_gvcf, - [diploid_gvcf, haploid_gvcf], - ), - "haploid-gvcf-concat", - 0, - ) - return LRCallVariantsResult( - first_calling_job, - diploid_transfer_jobs, - diploid_concat_job, - first_modelapply_job, - phaser_job, - bcftools_subset_phased_job, - fai_to_bed_job, - bcftools_subtract_job, - repeatmodel_job, - bcftools_subset_unphased_job, - second_calling_job, - haploid_patch_job, - patch_transfer_jobs, - patch_concat_jobs, - second_modelapply_job, - calling_unphased_job, - diploid_patch_job, - unphased_transfer_jobs, - unphased_concat_job, - modelapply_unphased_job, - merge_job, - gvcf_combine_job, - haploid_calling_job, - haploid_patch2_job, - haploid_concat_job, - haploid_gvcf_combine_job, + haploid_gvcf_concat_job = Job( + cmds.bcftools_concat( + output_gvcf, + [diploid_gvcf, haploid_gvcf], + ), + "haploid-gvcf-concat", + 0, + task_name="gvcf", + ) + dag.add_job( haploid_gvcf_concat_job, + {haploid_gvcf_combine_job, gvcf_combine_job}, ) - def call_svs( - self, - sample_input: List[pathlib.Path], - replace_rg: Optional[List[List[str]]] = None, - ) -> Job: + def call_svs(self, sample_input: List[pathlib.Path]) -> Job: """ Call SVs using Sentieon LongReadSV """ - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.skip_version_check: - for cmd, min_version in SV_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) - - sv_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", ".sv.vcf.gz") - ) - driver = Driver( - reference=self.reference, - thread_count=self.cores, - replace_rg=replace_rg, - input=sample_input, - interval=self.bed, - ) - driver.add_algo( - LongReadSV( - sv_vcf, - self.model_bundle.joinpath("longreadsv.model"), + bundle = self.required(self.model_bundle, "model_bundle") + require_versions(LONGREADSV_MIN_VERSIONS, skip=self.skip_version_check) + + return ( + LongReadSVStage( + ctx=self.stage_context(), + inputs=sample_input, + model=bundle.joinpath("longreadsv.model"), + interval=self.bed, ) + .build() + .job ) - longreadsv_job = Job( - Pipeline(Command(*driver.build_cmd())), "LongReadSV", self.cores - ) - return longreadsv_job diff --git a/sentieon_cli/driver.py b/sentieon_cli/driver.py index c2bbb7f..53b6131 100644 --- a/sentieon_cli/driver.py +++ b/sentieon_cli/driver.py @@ -410,7 +410,7 @@ def __init__( fa_file: Optional[pathlib.Path] = None, bed_file: Optional[pathlib.Path] = None, cut_indel: Optional[int] = None, - hap_bam: Optional[pathlib.Path] = None, + hap_bam: Optional[Union[pathlib.Path, str]] = None, hap_bed: Optional[pathlib.Path] = None, cut_len: Optional[int] = None, split_size: Optional[int] = None, @@ -497,9 +497,16 @@ def __init__( self, output: Union[pathlib.Path, str], model: pathlib.Path, + *, + sex: Optional[str] = None, + par: Optional[pathlib.Path] = None, + dump_probes: Optional[pathlib.Path] = None, ): self.output = output self.model = model + self.sex = sex + self.par = par + self.dump_probes = dump_probes class CNVModelApply(BaseAlgo): @@ -542,6 +549,26 @@ def __init__( self.prefix = prefix +class PGHapUpdateAlgo(BaseAlgo): + """algo PGHapUpdateAlgo""" + + name = "PGHapUpdateAlgo" + + def __init__( + self, + output: pathlib.Path, + gfa_file: pathlib.Path, + target_bed: Optional[pathlib.Path] = None, + min_map_qual: Optional[int] = None, + prefix: Optional[str] = None, + ): + self.output = output + self.gfa_file = gfa_file + self.target_bed = target_bed + self.min_map_qual = min_map_qual + self.prefix = prefix + + class BaseDriver: """A base class for the Sentieon driver""" @@ -550,7 +577,7 @@ def __init__( reference: Optional[pathlib.Path] = None, thread_count: Optional[int] = None, interval: Optional[Union[pathlib.Path, str]] = None, - interval_padding: int = 0, + interval_padding: Optional[int] = None, read_filter: Optional[List[str]] = None, replace_rg: Optional[List[List[str]]] = None, input: Optional[List[pathlib.Path]] = None, diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index 4d60f3e..9f3fec4 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -1,30 +1,114 @@ """Execute jobs""" import asyncio -import asyncio.subprocess +import concurrent.futures +import contextlib +import os +import pathlib +import resource import signal import sys +import threading import time -from typing import Any, List, Tuple +from abc import ABC, abstractmethod +from typing import Any, Callable, Dict, List, Optional, Tuple from .job import Job from .logging import get_logger -from .scheduler import ThreadScheduler -from .shell_pipeline import Context +from .run_logs import JobLogSink, RunLogs +from .scheduler import BaseScheduler +from .shell_pipeline import Command, Context logger = get_logger(__name__) +# Lines of a failing process's log to quote in the failure report +TAIL_LINES = 20 +# Only the tail of a log is read; tool logs can be arbitrarily large +TAIL_BYTES = 64 * 1024 -class BaseExecutor: + +def _log_tail(path: pathlib.Path, lines: int = TAIL_LINES) -> List[str]: + """Return the last ``lines`` lines of a log file. + + Reads at most ``TAIL_BYTES`` from the end, so quoting the tail of a + multi-gigabyte log is cheap. A line split by the seek is dropped with the + rest of the truncated prefix. + """ + try: + with open(path, "rb") as handle: + size = handle.seek(0, os.SEEK_END) + handle.seek(max(0, size - TAIL_BYTES)) + data = handle.read() + except OSError as exc: + logger.debug("could not read %s: %s", path, exc) + return [] + return data.decode("utf-8", errors="replace").splitlines()[-lines:] + + +def _maxrss_bytes(ru: resource.struct_rusage) -> int: + """Return a child's peak resident set size in bytes. + + ``ru_maxrss`` is already bytes on macOS; every other platform we run on + reports kibibytes. + """ + if sys.platform == "darwin": + return int(ru.ru_maxrss) + return int(ru.ru_maxrss) * 1024 + + +def _mib(nbytes: int) -> str: + """Render a byte count as MiB for the human-readable log lines.""" + return f"{nbytes / (1024 * 1024):.1f} MiB" + + +def _signal_procs(context: Context, sig: int) -> None: + """Send a signal to every live sub-process in a context.""" + for subcommand in context.commands: + proc = subcommand.proc + if proc is not None and proc.returncode is None: + try: + proc.send_signal(sig) + except ProcessLookupError: + pass + + +def _kill_survivors(context: Context) -> None: + """SIGKILL any sub-process that is still running.""" + for subcommand in context.commands: + proc = subcommand.proc + if proc is not None and proc.returncode is None: + try: + proc.kill() + except ProcessLookupError: + pass + + +async def _cleanup_quietly(context: Context) -> None: + """Release a context, logging -- not raising -- a surfaced launch failure. + + A proc-sub inner command that fails to spawn surfaces its OSError from + ``context.cleanup()``'s gather. On the abort and shutdown paths the owning + job is already recorded (or the whole run is being torn down), so a stored + OSError must not escape ``execute()``; log it and release the rest. A + non-OSError stays a bug and propagates, matching ``_drive``'s per-job + handling. + """ + try: + await context.cleanup() + except OSError as exc: + logger.error("error releasing resources during teardown: %s", exc) + + +class BaseExecutor(ABC): """Execute jobs""" - def __init__(self, scheduler: ThreadScheduler): + def __init__(self, scheduler: BaseScheduler): self.scheduler = scheduler self.jobs_with_errors: List[Job] = [] + @abstractmethod def execute(self) -> None: """Execute jobs from the DAG""" - raise NotImplementedError class DryRunExecutor(BaseExecutor): @@ -35,8 +119,7 @@ def run_job(self, job: Job) -> None: print(job.shell) def execute(self) -> None: - scheduler_gen = self.scheduler.schedule() - ready_jobs = scheduler_gen.send(None) + ready_jobs = self.scheduler.start() for job in ready_jobs: self.run_job(job) @@ -45,18 +128,171 @@ def execute(self) -> None: ready_jobs = { new_job for completed_job in finished_jobs - for new_job in scheduler_gen.send(completed_job) + for new_job in self.scheduler.job_finished(completed_job) } for job in ready_jobs: self.run_job(job) -class LocalExecutor(BaseExecutor): - """Run jobs locally""" +class AsyncExecutor(BaseExecutor, ABC): + """Base for executors that run an asyncio event loop. + + Provides the ``execute()`` entry point and opt-in SIGINT/SIGTERM handling + with graceful shutdown; subclasses implement :meth:`_drive` (the run loop) + and :meth:`_shutdown` (teardown after an interrupt). + + Set ``install_signal_handlers=True`` to install handlers (on the main + thread only) for the duration of a run; the previous handlers are + restored afterwards, so a caller's signal state is left untouched. - def __init__(self, *args: Any, **kwargs: Any): - super().__init__(*args, **kwargs) + ``thread_pool_size`` sizes the loop's default thread pool for the run; + ``None`` leaves the loop's own default alone, which is what an executor + that spawns no processes wants. + """ + + def __init__( + self, + scheduler: BaseScheduler, + *, + install_signal_handlers: bool = False, + thread_pool_size: Optional[int] = None, + ) -> None: + super().__init__(scheduler) + self.install_signal_handlers = install_signal_handlers + self.thread_pool_size = thread_pool_size self.start_new_jobs = True + + def _install_signal_handlers( + self, + loop: asyncio.AbstractEventLoop, + stop_event: asyncio.Event, + ) -> Callable[[], None]: + """Install SIGINT/SIGTERM handlers on the running loop. + + Returns a callable that removes the handlers and restores the + previous ones. A no-op is returned when handlers are disabled, or + when we are not on the main thread (Python can only manage signals + there, so we skip rather than crash). + """ + if not self.install_signal_handlers: + return lambda: None + if threading.current_thread() is not threading.main_thread(): + logger.debug("Signal handlers require the main thread; skipping.") + return lambda: None + + def on_signal(signum: int) -> None: + logger.error( + "Received %s; terminating jobs.", + signal.Signals(signum).name, + ) + self.start_new_jobs = False + stop_event.set() + + signums = (signal.SIGINT, signal.SIGTERM) + previous: Dict[int, Any] = { + signum: signal.getsignal(signum) for signum in signums + } + installed: List[int] = [] + for signum in signums: + try: + loop.add_signal_handler(signum, on_signal, signum) + installed.append(signum) + except (NotImplementedError, RuntimeError, ValueError) as exc: + logger.debug( + "Could not install handler for %s: %s", signum, exc + ) + + def restore() -> None: + for signum in installed: + try: + loop.remove_signal_handler(signum) + except (RuntimeError, ValueError): + pass + handler = previous[signum] + if handler is not None: + try: + signal.signal(signum, handler) + except (OSError, ValueError, TypeError): + pass + + return restore + + def execute(self) -> None: + """Execute jobs from the DAG""" + asyncio.run(self._execute()) + + async def _execute(self) -> None: + """Execute jobs from the DAG""" + self.jobs_with_errors = [] + loop = asyncio.get_running_loop() + if self.thread_pool_size is not None: + # Threads are created lazily, so a generous cap costs nothing + # unless it is used, and asyncio.run() shuts the executor down on + # exit. + loop.set_default_executor( + concurrent.futures.ThreadPoolExecutor( + max_workers=self.thread_pool_size, + thread_name_prefix="sentieon-wait", + ) + ) + stop_event = asyncio.Event() + restore = self._install_signal_handlers(loop, stop_event) + try: + try: + await self._drive(stop_event) + except BaseException: + # Tear the children down before the exception unwinds into + # asyncio.run's Runner.close, which joins the wait threads -- + # and those threads are blocked in wait4 on live children. + # BaseException so cancellation and KeyboardInterrupt are + # covered too. + self.start_new_jobs = False + try: + await self._shutdown() + except Exception as exc: + # A failed teardown must not mask the original error. + logger.error("teardown after failure also failed: %s", exc) + raise + if stop_event.is_set(): + await self._shutdown() + finally: + restore() + + @abstractmethod + async def _drive(self, stop_event: asyncio.Event) -> None: + """Submit/run jobs until the DAG drains or a stop is requested.""" + + @abstractmethod + async def _shutdown(self) -> None: + """Tear down whatever is still running after a stop request.""" + + +class LocalExecutor(AsyncExecutor): + """Run jobs locally as sub-processes, driven by an asyncio loop. + + ``thread_pool_size`` defaults high because the loop's default pool serves + two kinds of blocking work at once: one thread per in-flight ``wait()`` + and the proc-sub FIFO opens. A FIFO open queued behind saturated wait + threads would deadlock the run, so the pool must comfortably exceed the + number of processes a run keeps in flight. + """ + + def __init__( + self, + scheduler: BaseScheduler, + *, + install_signal_handlers: bool = False, + shutdown_grace_period: float = 10.0, + run_logs: Optional[RunLogs] = None, + thread_pool_size: int = 512, + ) -> None: + super().__init__( + scheduler, + install_signal_handlers=install_signal_handlers, + thread_pool_size=thread_pool_size, + ) + self.shutdown_grace_period = shutdown_grace_period + self.run_logs = run_logs self.running: List[ Tuple[ Job, @@ -65,125 +301,342 @@ def __init__(self, *args: Any, **kwargs: Any): int, ] ] = [] - signal.signal(signal.SIGINT, self.exit_gracefully) - signal.signal(signal.SIGTERM, self.exit_gracefully) - - def exit_gracefully(self, signum, frame) -> None: - logger.error("Termination signal detected. Terminating jobs.") - self.start_new_jobs = False - for running_job in self.running: - for subcommand in running_job[1].commands: - if not subcommand.proc: - continue - subcommand.proc.send_signal(signal=signal.SIGTERM) - time.sleep(10) - for running_job in self.running: - context = running_job[1] - for subcommand in context.commands: - if not subcommand.proc: - continue - if subcommand.proc.returncode is None: - subcommand.proc.kill() - asyncio.run(context.cleanup()) async def run_job(self, job: Job) -> None: """Run a job""" cmd = job.shell logger.info("Running: %s", cmd) - context = Context() + sink = self.run_logs.job_sink(job) if self.run_logs else None + context = Context() if sink is None else Context(log_sink=sink) start_time = time.monotonic_ns() try: - proc = await cmd.run( - context, - stderr=sys.stderr, - ) + if sink is None: + # No log directory: the processes inherit our stderr + proc = await cmd.run( + context, + stderr=sys.stderr, + ) + else: + # Unset streams are resolved against the sink + proc = await cmd.run(context) self.running.append( ( job, context, - asyncio.create_task(proc.wait()), + asyncio.create_task(proc.async_wait()), start_time, ) ) - except Exception as e: + except OSError as e: + # OSError covers the expected launch failures (command not found, + # permission denied, pipe/FIFO errors). Anything else is a bug in + # the caller's pipeline and is left to propagate. logger.error("failed to start command: %s", job.shell) logger.error("Error: %s", str(e)) - await context.cleanup() + # Stages that did spawn have logs; keep them and point at them. + self._report_logs(job, sink) + self.jobs_with_errors.append(job) + self.start_new_jobs = False + # A later pipeline stage can fail to spawn after earlier stages + # are already running; terminate them so they do not outlive the + # workflow, then release the remaining resources. A proc-sub inner + # launch failure can also surface from cleanup() here (narrow + # race): the job is already recorded, so swallow it. + await self._terminate_context(context) + await _cleanup_quietly(context) - def execute(self) -> None: - """Execute jobs from the DAG""" - asyncio.run(self._execute()) + def _report_logs(self, job: Job, sink: Optional[JobLogSink]) -> None: + """Point at the logs of a job that was aborted rather than reaped.""" + if sink is None: + return + paths = sink.log_paths() + if paths: + logger.error( + "Logs from %s: %s", + job, + ", ".join(str(path) for path in paths), + ) - async def _execute(self) -> None: - """Execute jobs from the DAG""" - self.jobs_with_errors: List[Job] = [] - scheduler_gen = self.scheduler.schedule() - ready_jobs = scheduler_gen.send(None) + def _report_failure( + self, + job: Job, + subcommand: Command, + ret: int, + sink: Optional[JobLogSink], + ) -> None: + """Report a failed sub-command with its pid, log path and log tail. + + Without a log -- no log directory, or a process that never spawned -- + the caller's messages are the whole report. + """ + log_path = sink.stderr_log_for(subcommand) if sink else None + if log_path is None: + return + pid = subcommand.proc.pid if subcommand.proc else None + report = [ + f"Failed sub-command of {job}: {subcommand}", + f" exit code: {ret}, pid: {pid}", + f" log: {log_path}", + ] + tail = _log_tail(log_path) + if tail: + report.append(f" last {len(tail)} line(s) of the log:") + report.extend(f" {line}" for line in tail) + else: + report.append(" the log is empty") + logger.error("\n".join(report)) + + async def _terminate_context(self, context: Context) -> None: + """Tear down processes already spawned in an aborted context. + + Used when a job aborts mid-launch (a later pipeline stage failed to + spawn): earlier stages are already running and would otherwise be + leaked. SIGTERM them, allow a grace period, then SIGKILL any + survivors, waiting for each to be reaped. + """ + live = [ + subcommand.proc + for subcommand in context.commands + if subcommand.proc is not None + and subcommand.proc.returncode is None + ] + if not live: + return + _signal_procs(context, signal.SIGTERM) + waits = [asyncio.create_task(proc.async_wait()) for proc in live] + await asyncio.wait(waits, timeout=self.shutdown_grace_period) + _kill_survivors(context) + await asyncio.wait(waits) + + async def _drive(self, stop_event: asyncio.Event) -> None: + """Schedule and run jobs until the DAG drains or we are stopped.""" + ready_jobs = self.scheduler.start() for job in ready_jobs: + if not self.start_new_jobs: + # A job in this batch failed to launch; drain, don't start + # the rest -- matching the policy the main loop applies. + break await self.run_job(job) - while self.running: - done, _running = await asyncio.wait( - [job[2] for job in self.running], - return_when=asyncio.FIRST_COMPLETED, - ) + stop_task = asyncio.ensure_future(stop_event.wait()) + try: + while self.running and not stop_event.is_set(): + wait_on: List[Any] = [job[2] for job in self.running] + wait_on.append(stop_task) + done, _pending = await asyncio.wait( + wait_on, + return_when=asyncio.FIRST_COMPLETED, + ) - finished_jobs = [ - self.running.pop(i) - for i in reversed(range(len(self.running))) - if self.running[i][2] in done - ] + if stop_event.is_set(): + break - # Check job execution - for job, context, _task, start_time in finished_jobs: - end_time = time.monotonic_ns() - total_seconds = (end_time - start_time) / 1e9 - - # Check if the command failed - cmd_failed = False - for subcommand in context.commands: - if not subcommand.proc: - logger.error( - "Sub-command has no process: %s", - subcommand, - ) - cmd_failed = True - continue - ret = ( - await subcommand.proc.wait() - ) # Wait on all sub-commands - if ret != 0 and not subcommand.fail_ok: - logger.error( - f"Return code {ret} for sub-command: {subcommand}" + finished_jobs = [ + self.running.pop(i) + for i in reversed(range(len(self.running))) + if self.running[i][2] in done + ] + + # Check job execution + for job, context, _task, start_time in finished_jobs: + end_time = time.monotonic_ns() + total_seconds = (end_time - start_time) / 1e9 + + # Check if the command failed + cmd_failed = False + failures: List[Tuple[Command, int]] = [] + # Kernel resource usage, aggregated over the job's + # processes. RSS is the per-process peak rather than a + # sum: concurrent pipeline stages peak at different + # times, so a sum would overstate the job. + have_rusage = False + job_utime = 0.0 + job_stime = 0.0 + job_maxrss = 0 + n_procs = 0 + sink = context.log_sink + for subcommand in context.commands: + if not subcommand.proc: + logger.error( + "Sub-command has no process: %s", + subcommand, + ) + cmd_failed = True + continue + n_procs += 1 + ret = ( + await subcommand.proc.async_wait() + ) # Wait on all sub-commands + # Reported for every reaped process, whatever the + # job's outcome; signal-killed children have rusage + # too. + rusage = subcommand.proc.rusage + rss = ( + _maxrss_bytes(rusage) + if rusage is not None + else None ) + if rusage is not None and rss is not None: + have_rusage = True + job_utime += rusage.ru_utime + job_stime += rusage.ru_stime + job_maxrss = max(job_maxrss, rss) + logger.debug( + "rusage for %s [pid %d, %s]: utime=%.2fs " + "stime=%.2fs maxrss=%.1fMiB minflt=%d " + "majflt=%d inblock=%d oublock=%d " + "nvcsw=%d nivcsw=%d", + job, + subcommand.proc.pid, + subcommand, + rusage.ru_utime, + rusage.ru_stime, + rss / (1024 * 1024), + rusage.ru_minflt, + rusage.ru_majflt, + rusage.ru_inblock, + rusage.ru_oublock, + rusage.ru_nvcsw, + rusage.ru_nivcsw, + ) + if self.run_logs is not None: + self.run_logs.record_process_metrics( + job=job, + command=subcommand, + pid=subcommand.proc.pid, + exit_code=ret, + stage=( + sink.stage_index_for(subcommand) + if sink is not None + else None + ), + rusage=rusage, + maxrss_bytes=rss, + ) + # A subcommand killed by SIGPIPE is not a failure: + # it was writing to a pipe whose reader exited early + # (e.g. `... | head`), which is normal in default + # (non-pipefail) bash. Real crashes surface as other + # non-zero codes and still fail the job. + if ret == -signal.SIGPIPE: + logger.debug( + "Sub-command exited on SIGPIPE (ok): %s", + subcommand, + ) + elif ret != 0 and not subcommand.fail_ok: + logger.error( + f"Return code {ret} for sub-command: " + f"{subcommand}" + ) + cmd_failed = True + failures.append((subcommand, ret)) + if ret == -9: + logger.error( + "Sub-command received SIGKILL. " + "Possible out-of-memory?" + ) + + # Release resources before recording the verdict. A + # proc-sub inner command that failed to launch surfaces + # its OSError from the background task here; treat it as + # a launch failure for this job rather than letting it + # abort the run while siblings are still mid-flight. + try: + await context.cleanup() + except OSError as exc: + logger.error("failed to start command: %s", job.shell) + logger.error("Error: %s", str(exc)) cmd_failed = True - if ret == -9: + + if cmd_failed: + logger.error("Command failure: %s", job.shell) + # The sink's files are closed by cleanup(), so their + # tails can be read now. + for failed, code in failures: + self._report_failure(job, failed, code, sink) + if self.run_logs is not None: logger.error( - "Sub-command received SIGKILL. " - "Possible out-of-memory?" + "Task logs for debugging: %s", + self.run_logs.task_logs, ) + self.jobs_with_errors.append(job) + self.start_new_jobs = False + elif have_rusage: + logger.info( + "Finished command in %.2fs (user %.1fs, " + "sys %.1fs, max proc RSS %s): %s", + total_seconds, + job_utime, + job_stime, + _mib(job_maxrss), + job.shell, + ) + else: + logger.info( + f"Finished command in " + f"{total_seconds:.2f}: {job.shell}" + ) + if self.run_logs is not None: + self.run_logs.record_job_metrics( + job=job, + failed=cmd_failed, + wall_s=total_seconds, + user_s=job_utime if have_rusage else None, + sys_s=job_stime if have_rusage else None, + max_proc_rss_bytes=( + job_maxrss if have_rusage else None + ), + processes=n_procs, + ) + if sink is not None: + sink.finalize(success=not cmd_failed) - if cmd_failed: - logger.error("Command failure: %s", job.shell) - self.jobs_with_errors.append(job) - self.start_new_jobs = False - else: - logger.info( - f"Finished command in " - f"{total_seconds:.2f}: {job.shell}" + if not self.start_new_jobs: + # Don't start new jobs + continue + + ready_jobs = { + new_job + for completed_job in finished_jobs + for new_job in self.scheduler.job_finished( + completed_job[0] ) - await context.cleanup() + } - if not self.start_new_jobs: - # Don't start new jobs - continue + # Run the ready jobs + for job in ready_jobs: + if not self.start_new_jobs: + # A job in this batch failed to launch; drain the + # rest of the run without starting more. + break + await self.run_job(job) + finally: + stop_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stop_task - ready_jobs = { - new_job - for completed_job in finished_jobs - for new_job in scheduler_gen.send(completed_job[0]) - } + async def _shutdown(self) -> None: + """Terminate running jobs, escalate to SIGKILL, then clean up.""" + running = self.running + self.running = [] - # Run the ready jobs - for job in ready_jobs: - await self.run_job(job) + for _job, context, _task, _start in running: + _signal_procs(context, signal.SIGTERM) + + tasks = [tup[2] for tup in running] + if tasks: + await asyncio.wait(tasks, timeout=self.shutdown_grace_period) + + for _job, context, _task, _start in running: + _kill_survivors(context) + if tasks: + await asyncio.wait(tasks) + + for job, context, _task, _start in running: + # A proc-sub inner launch failure can surface from cleanup() as we + # tear down after an interrupt; log it rather than let it escape, + # and keep releasing the remaining contexts. + await _cleanup_quietly(context) + # An interrupted job keeps all of its logs; they explain how far + # it got. + self._report_logs(job, context.log_sink) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py new file mode 100644 index 0000000..d3ae092 --- /dev/null +++ b/sentieon_cli/hybrid_pangenome.py @@ -0,0 +1,1337 @@ +""" +The Sentieon hybrid-pangenome pipeline + +Combines the human pangenome with sample short-read and long-read data +for highly accurate variant calling. +""" + +import argparse +import copy +import json +import pathlib +import sys +from typing import Dict, List, Optional, Set, Tuple + +import packaging.version + +from . import command_strings as cmds +from .archive import ar_load +from .base_pangenome import BasePangenome +from .dag import DAG +from .driver import ( + DNAscope, + Driver, + PangenomeSV, + PGHapUpdateAlgo, +) +from .job import Job +from .logging import get_logger +from .shard import ( + detect_reference_build, + determine_shards_from_fai, + parse_fai, +) +from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + BwaExtractStage, + Minimap2RealignStage, + find_unzip, +) +from .stages.base import StageContext +from .stages.cnv import CNV_MIN_VERSIONS +from .stages.dedup import DedupStage +from .stages.expansion import EXPANSION_MIN_VERSIONS +from .stages.metrics import MetricsPaths, MetricsStage +from .stages.ploidy import PloidyStage +from .stages.segdup import SEGDUP_MIN_VERSIONS +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + TransferApplyStage, + TransferSpec, +) +from .stages.sv import LongReadSVStage +from .stages.transfer import TransferConfig +from .util import ( + __version__, + check_kmc_patch, + executable_version, + parse_rg_line, + path_arg, + require_versions, + total_memory, +) + +HYBRID_PANGENOME_MIN_VERSIONS = { + "kmc": None, + "sentieon driver": packaging.version.Version("202503.04"), + "vg": None, + "bcftools": packaging.version.Version("1.22"), + "samtools": packaging.version.Version("1.16"), + "bedtools": None, +} + +# The last driver release whose `pgutil gfa2fa` handles only GRCh38 +# pangenomes. Non-GRCh38 runs on it (and older) extract the graph's +# paths with `vg paths` instead, which is much slower. +PGUTIL_GFA2FA_GRCH38_ONLY_VERSION = packaging.version.Version("202503.04") + +# LongReadSV settings for finding graph update regions +LONGREADSV_MIN_SV_SIZE = 20 + +# PangenomeSV settings +PANGENOME_SV_MIN_AF = 0.1 + +logger = get_logger(__name__) + + +class HybridPangenome(BasePangenome): + """The Sentieon hybrid-pangenome pipeline""" + + params = copy.deepcopy(BasePangenome.params) + params.update( + { + # Required arguments + "lr_aln": { + "nargs": "*", + "help": ( + "Long-read BAM or CRAM files aligned to the linear " + "reference genome." + ), + "required": True, + "type": path_arg(exists=True, is_file=True), + }, + "pop_vcf": { + "flags": ["--pop_vcf"], + "help": ( + "A VCF containing annotations for use with DNAModelApply." + ), + "type": path_arg(exists=True, is_file=True), + "required": True, + }, + "readgroup": { + "help": ( + "Readgroup information for the fastq files. Required " + "with fastq input; cannot be used with `--sr_aln`." + ), + }, + # Additional arguments + "bed": { + "flags": ["-b", "--bed"], + "help": ( + "Region BED file. Supplying this file will limit " + "small-variant and CNV calling to the intervals inside " + "the BED file." + ), + "type": path_arg(exists=True, is_file=True), + }, + "call_cnvs": { + "help": ( + "Call copy number variants with CNVscope and the " + "PangenomeSV output. Requires a model bundle with a " + "'cnv.model' file." + ), + "action": "store_true", + }, + "lr_align_input": { + "help": ( + "Align the `--lr_aln` input files to the linear " + "reference genome with minimap2. Use with unaligned " + "(uBAM or uCRAM) long-read input." + ), + "action": "store_true", + }, + "lr_input_ref": { + "help": ( + "A reference FASTA used to decode the `--lr_aln` input " + "files. Only used with `--lr_align_input`." + ), + "type": path_arg(exists=True, is_file=True), + }, + "rgsm": { + "help": ( + "Overwrite the SM tag of the input readgroups for " + "compatibility" + ), + }, + "skip_metrics": { + "help": "Skip metrics collection and multiQC", + "action": "store_true", + }, + "skip_multiqc": { + "help": "Skip multiQC report generation", + "action": "store_true", + }, + "sr_aln": { + "nargs": "*", + "help": ( + "Aligned short-read BAM or CRAM files. Assumed " + "deduplicated or duplicate-marked. Cannot be used with " + "fastq input." + ), + "type": path_arg(exists=True, is_file=True), + }, + # Hidden arguments + "skip_model_apply": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, + "skip_pop_vcf_id_check": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, + "skip_small_variants": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, + "skip_svs": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, + } + ) + + positionals = BasePangenome.positionals + + def __init__(self) -> None: + super().__init__() + self.lr_aln: List[pathlib.Path] = [] + self.pop_vcf: Optional[pathlib.Path] = None + self.readgroup: Optional[str] = None + self.bed: Optional[pathlib.Path] = None + self.call_cnvs = False + self.lr_align_input = False + self.lr_input_ref: Optional[pathlib.Path] = None + self.rgsm: Optional[str] = None + self.extract_model_name = "extract.model" + self.gfa2fa_with_vg = False + self.skip_metrics = False + self.skip_multiqc = False + self.sr_aln: List[pathlib.Path] = [] + self.fastq_readgroup: Optional[Dict[str, str]] = None + self.sr_readgroups: List[List[Dict[str, str]]] = [] + self.lr_readgroups: List[List[Dict[str, str]]] = [] + self.sample_sm = "" + self.skip_model_apply = False + self.skip_pop_vcf_id_check = False + self.skip_small_variants = False + self.skip_svs = False + # The single long-read alignment, stashed by `build_dag` for the + # second, sex-aware DAG + self.lr_alignment: Optional[pathlib.Path] = None + + def _cnv_in_second_dag(self) -> bool: + """CNV calling runs in the second, sex-aware DAG""" + return self.call_cnvs and self.has_cnv_model + + def _needs_second_dag(self) -> bool: + """The run has jobs that depend on the estimated sample sex""" + return bool( + self.expansion_catalog + or self.segdup_caller is not None + or self._cnv_in_second_dag() + ) + + def validate(self) -> None: + """Validate pipeline inputs""" + self.validate_ref() + self.fai_data = parse_fai(pathlib.Path(str(self.reference) + ".fai")) + self.shards = determine_shards_from_fai( + self.fai_data, 10 * 1000 * 1000 + ) + # Before `validate_bundle`, which picks the bundle's + # `extract..model` member by the resolved name + self.resolve_pangenome_reference() + self.load_pop_vcf_header(self.pop_vcf) + + self.validate_bundle() + self.validate_output_vcf() + + self.validate_sr_inputs() + + if not self.lr_aln: + self.logger.error( + "Please supply the long-read alignments with the `--lr_aln` " + "argument" + ) + sys.exit(2) + + if self.lr_input_ref and not self.lr_align_input: + self.logger.warning( + "The `--lr_input_ref` argument is only used with " + "`--lr_align_input` and will be ignored" + ) + if self.lr_align_input and not self.lr_input_ref: + if any(str(aln).endswith(".cram") for aln in self.lr_aln): + self.logger.warning( + "CRAM input may fail to decode without the " + "`--lr_input_ref` argument" + ) + + if self.r1_fastq: + self.validate_bwa_index() + self.collect_readgroups() + self.validate_readgroups() + + self.validate_segdup() + self.validate_expansion() + self.validate_t1k() + self.validate_cnv() + + require_versions( + HYBRID_PANGENOME_MIN_VERSIONS, skip=self.skip_version_check + ) + + if not self.skip_version_check and not check_kmc_patch("kmc"): + self.logger.error( + "Error: The 'kmc' executable in the PATH does not " + "support reading from stdin. Please ensure " + "you are using the patched version of KMC from " + "https://github.com/Sentieon/KMC/releases." + ) + sys.exit(2) + + if self.bed is None: + self.logger.info( + "A BED file is recommended to avoid small-variant calling " + "across decoy and unplaced contigs." + ) + + self.validate_pangenome_contig_lengths() + + # After the version gate and the contig checks, so the tool choice + # is only logged for runs that proceed. The `reference_build` it + # sets is the value `validate_cnv` already resolved. + self.resolve_gfa2fa_tool() + + def validate_segdup(self) -> None: + """Validate the arguments used by segdup-caller""" + if self.segdup_caller is None: + return + + if len(self.sr_aln) > 1: + self.logger.error( + "`--segdup_caller` accepts only a single `--sr_aln` " + "BAM/CRAM file." + ) + sys.exit(2) + + if len(self.lr_aln) > 1: + self.logger.error( + "`--segdup_caller` accepts only a single `--lr_aln` " + "BAM/CRAM file." + ) + sys.exit(2) + + if self.skip_small_variants: + self.logger.error( + "`--segdup_caller` requires the small-variant VCF as " + "`--input_vcf` and cannot be combined with " + "`--skip_small_variants`." + ) + sys.exit(2) + + require_versions(SEGDUP_MIN_VERSIONS, skip=self.skip_version_check) + + def validate_expansion(self) -> None: + """Validate the arguments used by ExpansionHunter""" + if self.expansion_catalog is None: + return + + if len(self.sr_aln) > 1: + self.logger.error( + "`--expansion_catalog` accepts only a single `--sr_aln` " + "BAM/CRAM file." + ) + sys.exit(2) + + require_versions(EXPANSION_MIN_VERSIONS, skip=self.skip_version_check) + + def validate_cnv(self) -> None: + """Validate the arguments used for sex-aware CNV calling""" + if self.call_cnvs and self.skip_svs: + self.logger.error( + "`--call_cnvs` requires SV calling and cannot be combined " + "with `--skip_svs`." + ) + sys.exit(2) + + # `validate_bundle` has already set `self.has_cnv_model` + if self.call_cnvs and not self.has_cnv_model: + self.logger.error( + "`--call_cnvs` requires a model bundle with a 'cnv.model' " + "file." + ) + sys.exit(2) + + if self.call_cnvs and self.bed is None: + self.logger.warning( + "No `--bed` supplied; CNVscope will run on every contig in " + "the reference, including decoy and unplaced contigs" + ) + + cnv_will_run = self._cnv_in_second_dag() + self.resolve_cnv_par_bed(self.fai_data, self.par_bed, cnv_will_run) + if not cnv_will_run: + return + + # A PAR BED file is required for CNV calling, whatever the sex + self.validate_cnv_par(True) + + require_versions(CNV_MIN_VERSIONS, skip=self.skip_version_check) + + def validate_sr_inputs(self) -> None: + """Validate the short-read input arguments. + + Short reads are supplied either as fastq (with `--readgroup`) or as + aligned BAM/CRAM files (`--sr_aln`), but not both. + """ + if self.r1_fastq and self.sr_aln: + self.logger.error( + "Supplying both fastq (`--r1_fastq`) and aligned " + "(`--sr_aln`) short reads is not supported" + ) + sys.exit(2) + if not self.r1_fastq and not self.sr_aln: + self.logger.error( + "Please supply the short reads with the `--sr_aln` argument " + "or with the `--r1_fastq`, `--r2_fastq`, and `--readgroup` " + "arguments" + ) + sys.exit(2) + + if self.r1_fastq: + if not self.readgroup: + self.logger.error( + "The `--readgroup` argument is required with fastq input" + ) + sys.exit(2) + if len(self.r1_fastq) != len(self.r2_fastq): + self.logger.error( + "The number of input `--r1_fastq` files does not equal " + "the number of `--r2_fastq` files" + ) + sys.exit(2) + elif self.readgroup: + self.logger.error( + "The `--readgroup` argument cannot be used with aligned " + "short-read input (`--sr_aln`)" + ) + sys.exit(2) + + def validate_bundle(self) -> None: + """Validate the model bundle""" + bundle_info_bytes = ar_load( + str(self.model_bundle) + "/bundle_info.json" + ) + if isinstance(bundle_info_bytes, list): + bundle_info_bytes = b"{}" + bundle_info = json.loads(bundle_info_bytes.decode()) + + req_version_s = bundle_info.get("minScriptVersion") + if req_version_s: + req_version = packaging.version.Version(req_version_s) + if req_version > packaging.version.Version(__version__): + self.logger.error( + "The model bundle requires version %s or later of the " + "sentieon-cli.", + req_version, + ) + sys.exit(2) + + bundle_pipeline = bundle_info.get("pipeline") + if bundle_pipeline and bundle_pipeline != "Hybrid pangenome": + self.logger.error("The model bundle is for a different pipeline.") + sys.exit(2) + + bundle_members = set(ar_load(str(self.model_bundle))) + + # Prefer a reference-specific extract model. Fall back to the generic + # 'extract.model' only for the default 'GRCh38' reference. + extract_candidate = f"extract.{self.pangenome_ref_name}.model" + if extract_candidate in bundle_members: + self.extract_model_name = extract_candidate + elif self.pangenome_ref_name == "GRCh38": + self.extract_model_name = "extract.model" + else: + self.extract_model_name = extract_candidate + + required_members = { + "dnascope.model", + "longreadsv.model", + "minimap2.model", + self.extract_model_name, + } + if self.r1_fastq: + required_members.add("bwa.model") + if self.lr_align_input: + required_members.add("minimap2_lr.model") + if self.segdup_caller is not None: + # segdup-caller reads `/diploid_model` and the + # pipeline passes this bundle as its `--lr_model` + required_members.add("diploid_model") + missing_members = required_members - bundle_members + if missing_members: + self.logger.error( + "Expected model files not found in the model bundle file: %s", + ", ".join(sorted(missing_members)), + ) + sys.exit(2) + + self.has_cnv_model = "cnv.model" in bundle_members + + bundle_vcf_id = bundle_info.get("SentieonVcfID") + if bundle_vcf_id and not self.skip_pop_vcf_id_check: + self.check_pop_vcf_id(bundle_vcf_id) + + def collect_readgroups(self) -> None: + """Collect readgroup tags from the inputs""" + if self.readgroup: + try: + parsed_rg = parse_rg_line(self.readgroup.replace(r"\t", "\t")) + except ValueError as e: + self.logger.error( + "Invalid --readgroup value '%s': %s", self.readgroup, e + ) + sys.exit(2) + if not parsed_rg.get("ID"): + self.logger.error( + "Readgroup '%s' does not have a RGID tag", + self.readgroup, + ) + sys.exit(2) + if not parsed_rg.get("SM"): + self.logger.error( + "Readgroup '%s' does not have a RGSM tag", + self.readgroup, + ) + sys.exit(2) + self.fastq_readgroup = parsed_rg + + # Read the readgroups from the ORIGINAL long-read input; minimap2 + # realignment preserves the readgroup IDs with `addreplacerg` + self.sr_readgroups = [ + self.parse_aln_readgroups(aln) for aln in self.sr_aln + ] + self.lr_readgroups = [ + self.parse_aln_readgroups(aln) for aln in self.lr_aln + ] + + def parse_aln_readgroups(self, aln: pathlib.Path) -> List[Dict[str, str]]: + """Parse the @RG lines from the header of an input alignment""" + parsed: List[Dict[str, str]] = [] + for rg_line in cmds.get_rg_lines(aln, self.dry_run): + try: + parsed.append(parse_rg_line(rg_line)) + except ValueError as e: + self.logger.error( + "Invalid readgroup line in '%s': '%s': %s", + aln, + rg_line, + e, + ) + sys.exit(2) + return parsed + + def validate_readgroups(self) -> None: + """Confirm that all readgroups have a consistent SM tag""" + aln_inputs = list(self.sr_aln) + list(self.lr_aln) + aln_readgroups = self.sr_readgroups + self.lr_readgroups + + # Every input needs readgroups; the variant calling stages set the + # LR attribute of the input readgroups with `--replace_rg` + for aln, aln_rgs in zip(aln_inputs, aln_readgroups): + if not aln_rgs: + self.logger.error( + "No @RG lines found in the header of '%s'. Please add " + "readgroups to the input file", + aln, + ) + sys.exit(2) + + rg_sm: Optional[str] = None + if self.fastq_readgroup: + rg_sm = self.fastq_readgroup.get("SM") + elif self.sr_readgroups and self.sr_readgroups[0]: + rg_sm = self.sr_readgroups[0][0].get("SM") + + for aln, aln_rgs in zip(aln_inputs, aln_readgroups): + for aln_rg in aln_rgs: + if not aln_rg.get("ID"): + self.logger.error( + "Found a readgroup without an ID tag in '%s': %s", + aln, + str(aln_rg), + ) + sys.exit(2) + sm = aln_rg.get("SM") + if not sm: + self.logger.error( + "Found a readgroup without a SM tag in '%s': %s", + aln, + str(aln_rg), + ) + sys.exit(2) + if self.dry_run or self.rgsm: + continue + if sm != rg_sm: + self.logger.error( + "Input readgroup '%s' has a different RG-SM tag " + "from the other sample inputs.\n" + "found='%s' expected='%s'. Please set the `--rgsm` " + "argument to override the SM tag in the input files", + str(aln_rg), + sm, + rg_sm, + ) + sys.exit(2) + + # The dry-run readgroups are synthetic and identical for every + # input, so they would always collide + if not self.dry_run: + self.validate_rg_ids(aln_inputs, aln_readgroups) + + if not self.rgsm and not rg_sm: + self.logger.error( + "Could not determine the sample name from the inputs. " + "Please set the `--rgsm` argument" + ) + sys.exit(2) + self.sample_sm = self.rgsm if self.rgsm else str(rg_sm) + + def validate_rg_ids( + self, + aln_inputs: List[pathlib.Path], + aln_readgroups: List[List[Dict[str, str]]], + ) -> None: + """Confirm that every input readgroup ID is unique. + + The inputs are merged by readgroup ID during variant calling, so a + collision would give one readgroup two conflicting `--replace_rg` + rewrites of the LR attribute. + """ + rg_sources: Dict[str, str] = {} + if self.fastq_readgroup: + rg_sources[self.fastq_readgroup["ID"]] = ( + "the `--readgroup` argument" + ) + for aln, aln_rgs in zip(aln_inputs, aln_readgroups): + for aln_rg in aln_rgs: + rg_id = aln_rg["ID"] + if rg_id in rg_sources: + self.logger.error( + "Duplicate readgroup ID '%s' found in %s and '%s'. " + "Readgroup IDs need to be unique across all inputs", + rg_id, + rg_sources[rg_id], + aln, + ) + sys.exit(2) + rg_sources[rg_id] = f"'{aln}'" + + # The lifted alignment reuses the first short-read readgroup ID + # with a '-pg' suffix + lift_rg = self.fastq_readgroup + if not lift_rg and self.sr_readgroups and self.sr_readgroups[0]: + lift_rg = self.sr_readgroups[0][0] + if lift_rg: + lift_id = lift_rg["ID"] + "-pg" + if lift_id in rg_sources: + self.logger.error( + "The readgroup ID '%s' in %s is reserved for the lifted " + "alignment. Please rename the input readgroup", + lift_id, + rg_sources[lift_id], + ) + sys.exit(2) + + def resolve_gfa2fa_tool(self) -> None: + """Choose the tool that converts the updated graph to FASTA. + + `pgutil gfa2fa` handles only GRCh38 pangenomes up to driver + release `PGUTIL_GFA2FA_GRCH38_ONLY_VERSION`; other references on + those releases fall back to `vg paths`. When the driver version + is unknown (`--skip_version_check`, or the probe fails) a + non-GRCh38 reference also uses `vg paths`, the tool that works + with every release. + """ + self.reference_build = detect_reference_build(self.fai_data) + self.gfa2fa_with_vg = False + if self.reference_build == "hg38": + return + + driver_version = ( + None + if self.skip_version_check + else executable_version("sentieon driver") + ) + if ( + driver_version is not None + and driver_version > PGUTIL_GFA2FA_GRCH38_ONLY_VERSION + ): + return + + self.gfa2fa_with_vg = True + if driver_version is None: + self.logger.info( + "The reference is not GRCh38 and the `sentieon driver` " + "version is not known. The pangenome FASTA will be " + "generated with `vg paths`, which works with every " + "driver release but is slower than `pgutil gfa2fa`." + ) + else: + self.logger.info( + "The reference is not GRCh38 and `sentieon driver` " + "version '%s' supports only GRCh38 pangenomes in " + "`pgutil gfa2fa`. The pangenome FASTA will be generated " + "with `vg paths`, which is slower.", + driver_version, + ) + + def configure(self) -> None: + """Configure pipeline parameters""" + pass + + def build_dag(self) -> DAG: + """Build the DAG for the hybrid-pangenome pipeline""" + bundle = self.required(self.model_bundle, "model_bundle") + self.required(self.pop_vcf, "pop_vcf") + + self.logger.info("Building the hybrid-pangenome DAG") + dag = DAG() + ctx = self.stage_context() + + ref_fai = pathlib.Path(str(ctx.reference) + ".fai") + + # Output files + suffix = "bam" if self.bam_format else "cram" + out_bwa_aln = pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", f"_bwa_deduped.{suffix}") + ) + out_lift_aln = pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", f"_lift_deduped.{suffix}") + ) + sv_vcf = pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") + ) + + # Intermediate file paths + bwa_bam = self.tmp_dir.joinpath("sample-bwa.bam") + lift_bam = self.tmp_dir.joinpath("sample-lift.bam") + ext_fastq = self.tmp_dir.joinpath("sample-extract.fq.gz") + kmer_prefix = self.tmp_dir.joinpath("sample") + kmer_file = pathlib.Path(str(kmer_prefix) + ".kff") + sample_pangenome = self.tmp_dir.joinpath("sample_pangenome.gbz") + hap_raw_gfa = self.tmp_dir.joinpath("sample-hap.raw.gfa") + pangenome_raw_gfa = self.tmp_dir.joinpath("sample-pangenome-raw.gfa") + pangenome_gfa = self.tmp_dir.joinpath("sample-pangenome.gfa") + pangenome_fasta = self.tmp_dir.joinpath("sample-pangenome.fa") + longread_sv_vcf = self.tmp_dir.joinpath("sample-longread-sv.vcf.gz") + sv_bed = self.tmp_dir.joinpath("sample-sv.bed") + raw_vcf = self.tmp_dir.joinpath("sample-dnascope.vcf.gz") + transfer_vcf = self.tmp_dir.joinpath("sample-dnascope_transfer.vcf.gz") + + # Realign unaligned (uBAM or uCRAM) long-read input. K-mer counting + # reads the original input files, so it runs without waiting for + # the realignment. + calling_lr = list(self.lr_aln) + realign_jobs: Set[Job] = set() + if self.lr_align_input: + lr_result = self.lr_realign_stage(ctx).add_to(dag) + calling_lr = lr_result.outputs + realign_jobs = set(lr_result.jobs) + + haplotype_dependencies: Set[Job] = set() + mm2_dependencies: Set[Job] = set() + if self.r1_fastq: + # KMC k-mer counting across the short and long reads + kmc_job = self.build_hybrid_kmc_job(kmer_prefix) + 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) + # Do not run vg-haplotypes with bwa in low-mem environments + total_mem_gb = total_memory() / (1024.0**3) + if total_mem_gb < 70: + haplotype_dependencies.add(bwa_job) + + # The lifted alignment is deduplicated before variant calling + lift_aln = lift_bam + else: + # Read extraction and k-mer counting in a single pass over the + # aligned short-read input + symlink_job, extract_kmc_job = self.build_extract_kmc_jobs( + kmer_prefix, ext_fastq + ) + dag.add_job(symlink_job) + dag.add_job(extract_kmc_job, {symlink_job}) + haplotype_dependencies.add(extract_kmc_job) + mm2_dependencies.add(extract_kmc_job) + + # Aligned input is not deduplicated, so the lifted alignment is + # written directly to the output file + lift_aln = out_lift_aln + + # vg haplotypes - create a sample-specific pangenome + haplotypes_job = self.build_haplotypes_job(sample_pangenome, kmer_file) + dag.add_job(haplotypes_job, haplotype_dependencies) + + # Confirm the sampled pangenome kept its reference paths before + # anything consumes it + check_gbz_job = self.build_check_gbz_job(sample_pangenome) + dag.add_job(check_gbz_job, {haplotypes_job}) + + # convert the sample pangenome to GFA + gfa_job = self.build_gfa_job(hap_raw_gfa, sample_pangenome) + dag.add_job(gfa_job, {check_gbz_job}) + + # Confirm the GFA carries the reference's rGFA tags, without which + # `pgutil lift` leaves every read unmapped + check_gfa_job = self.build_check_gfa_job(hap_raw_gfa) + dag.add_job(check_gfa_job, {gfa_job}) + + # Graph update without the SV BED + update_raw_job = self.build_graph_update_job( + pangenome_raw_gfa, + hap_raw_gfa, + calling_lr, + name="graph-update-raw", + ) + dag.add_job(update_raw_job, {check_gfa_job} | realign_jobs) + + # Call SVs from the long reads and collect graph update regions + longreadsv_result = LongReadSVStage( + ctx=ctx, + inputs=list(calling_lr), + model=bundle.joinpath("longreadsv.model"), + output=longread_sv_vcf, + min_sv_size=LONGREADSV_MIN_SV_SIZE, + name="longreadsv", + task_name="pangenome-update", + ).add_to(dag, realign_jobs) + sv_bed_job = Job( + cmds.cmd_longread_sv_bed(sv_bed, longread_sv_vcf, ref_fai), + "longread-sv-bed", + 0, + task_name="pangenome-update", + ) + dag.add_job(sv_bed_job, {longreadsv_result.job}) + + # Graph update with the SV BED + update_job = self.build_graph_update_job( + pangenome_gfa, + pangenome_raw_gfa, + calling_lr, + bed=sv_bed, + name="graph-update", + ) + dag.add_job(update_job, {update_raw_job, sv_bed_job}) + + # FASTA generation from the updated graph + gfa2fa_job = self.build_gfa2fa_job( + pangenome_fasta, pangenome_gfa, ref_fai + ) + dag.add_job(gfa2fa_job, {update_job}) + faidx_job = Job( + cmds.cmd_samtools_faidx(pangenome_fasta), + "faidx", + 0, + task_name="pangenome", + ) + dag.add_job(faidx_job, {gfa2fa_job}) + + # minimap2 alignment of the extracted reads with liftover + mm2_job = self.build_minimap2_lift_job( + lift_aln, + ext_fastq, + pangenome_fasta, + pangenome_gfa, + ) + dag.add_job(mm2_job, mm2_dependencies | {faidx_job, update_job}) + + # Variant calling with the short reads, the lifted reads, and the + # long reads + if self.r1_fastq: + # Deduplicate the short-read alignments. The `--lr_aln` input + # is assumed to be deduplicated already. + # Emit Dedup metrics for the primary (bwa) short-read alignment + # so they land in the metrics directory scanned by MultiQC. + paths = MetricsPaths.from_output_vcf(ctx.output_vcf) + dedup_metrics: Optional[pathlib.Path] = None + if not self.skip_metrics: + paths.ensure_dir(self.dry_run) + dedup_metrics = paths.dedup_metrics + + bwa_dedup = DedupStage( + ctx=ctx, + tag="bwa", + inputs=[bwa_bam], + output=out_bwa_aln, + score_file=self.tmp_dir.joinpath("sample-bwa-score.txt.gz"), + dedup_metrics=dedup_metrics, + ).add_to(dag, {bwa_job}) + bwa_dedup_job = bwa_dedup.dedup_job + lift_dedup = DedupStage( + ctx=ctx, + tag="lift", + inputs=[lift_bam], + output=out_lift_aln, + score_file=self.tmp_dir.joinpath("sample-lift-score.txt.gz"), + ).add_to(dag, {mm2_job}) + lift_dedup_job = lift_dedup.dedup_job + + # Alignment metrics from the deduplicated short reads + if not self.skip_metrics: + metrics_result = MetricsStage( + ctx=ctx, + inputs=[out_bwa_aln, out_lift_aln], + algos=self.pangenome_metrics_algos(paths), + rehead_metrics=paths.wgs, + ).add_to(dag, {bwa_dedup_job, lift_dedup_job}) + if not self.skip_multiqc: + multiqc_job = self.multiqc() + if multiqc_job: + dag.add_job(multiqc_job, metrics_result.terminal) + + calling_bams = [out_bwa_aln, out_lift_aln] + calling_lr + calling_dependencies = {bwa_dedup_job, lift_dedup_job} + sr_alignments = [out_bwa_aln] + sr_deps: Set[Job] = {bwa_dedup_job} + else: + # Aligned short reads are used as-is; duplicate, secondary, and + # supplementary reads are excluded from the lifted alignment + # during read extraction + calling_bams = list(self.sr_aln) + [out_lift_aln] + calling_lr + calling_dependencies = {mm2_job} + sr_alignments = list(self.sr_aln) + sr_deps = set() + calling_dependencies |= realign_jobs + replace_rg = self.build_replace_rg() + + # Stash the short-read inputs for the second, sex-aware DAG. The + # `replace_rg` rows follow the calling input order, so the + # short-read rows are the leading ones; all-empty rows are dropped. + sr_replace_rg = replace_rg[: len(sr_alignments)] + self.sr_alignments = sr_alignments + self.sr_replace_rg = sr_replace_rg if any(sr_replace_rg) else None + self.lr_alignment = calling_lr[0] if len(calling_lr) == 1 else None + + # Estimate the sample ploidy and sex. The JSON output is always + # written; `--sample_sex` takes precedence for the sex used by the + # sex-aware callers. + ploidy_result = PloidyStage( + ctx=ctx, + inputs=[sr_alignments[0]], + reference_build=self.reference_build, + ).add_to(dag, sr_deps) + self.ploidy_json = ploidy_result.ploidy_json + + # T1K HLA/KIR calling from the short reads + self.add_t1k(dag, ctx, sr_alignments, sr_deps) + + if not self.skip_svs: + pangenomesv_job = self.build_pangenomesv_job( + sv_vcf, + calling_bams, + pangenome_gfa, + replace_rg, + ) + dag.add_job(pangenomesv_job, calling_dependencies | {update_job}) + + if self.skip_small_variants: + return dag + + model = bundle.joinpath("dnascope.model") + pcr_indel_model = "NONE" if self.pcr_free else "CONSERVATIVE" + call = DNAscopeStage( + ctx=ctx, + algos=[ + DNAscope( + raw_vcf, + model=model, + pcr_indel_model=pcr_indel_model, + dbsnp=self.dbsnp, + ) + ], + inputs=calling_bams, + interval=self.bed, + replace_rg=replace_rg, + ).add_to(dag, calling_dependencies) + + # transfer annotations from the pop_vcf, then apply the model + transfer_target = ( + transfer_vcf if not self.skip_model_apply else ctx.output_vcf + ) + TransferApplyStage( + ctx=ctx, + raw_vcf=raw_vcf, + transfer=TransferSpec( + TransferConfig.from_pipeline(self), transfer_target + ), + apply=( + ApplySpec(model, ctx.output_vcf) + if not self.skip_model_apply + else None + ), + ).add_to(dag, call.terminal) + + return dag + + def build_second_dag(self) -> Optional[DAG]: + """Build the second DAG for sex-aware downstream tools""" + if not self._needs_second_dag(): + return None + + assert self.ploidy_json is not None + self.get_sex(self.ploidy_json) + + self.logger.info("Building the second hybrid-pangenome DAG") + dag = DAG() + ctx = self.stage_context() + + # CNV calling with CNVscope on the short reads, using the sample sex + if self._cnv_in_second_dag(): + self.add_pangenome_cnv( + dag, + ctx, + self.output_path("_sv.vcf.gz"), + self.sr_alignments, + replace_rg=self.sr_replace_rg, + interval=self.bed, + ) + + if self.sr_alignments: + self.add_expansion(dag, ctx, self.sr_alignments[0]) + + # segdup-caller consumes the small-variant VCF, the sample sex, + # and the long reads. The bundle serves as both the short- and + # long-read model, as it ships the `diploid_model`. + self.add_segdup( + dag, + ctx, + self.sr_alignments[0], + lr_alignment=self.lr_alignment, + lr_bundle=self.model_bundle if self.lr_alignment else None, + ) + + return dag + + def lr_kmc_pairs(self) -> List[Tuple[pathlib.Path, pathlib.Path]]: + """`(alignment, decode reference)` pairs for the long-read k-mer + counting. + + K-mer counting reads the ORIGINAL `--lr_aln` input, which may use a + different reference from the alignment target. + """ + assert self.reference is not None + decode_ref = self.reference + if self.lr_align_input and self.lr_input_ref: + decode_ref = self.lr_input_ref + return [(aln, decode_ref) for aln in self.lr_aln] + + def build_hybrid_kmc_job(self, kmer_prefix: pathlib.Path) -> Job: + """Count k-mers across the short-read fastq and long-read + alignments""" + kmc_job = Job( + cmds.cmd_hybrid_kmc( + kmer_prefix, + self.r1_fastq + self.r2_fastq, + self.lr_kmc_pairs(), + self.tmp_dir, + memory=self.kmer_memory, + threads=self.cores, + unzip=find_unzip(self.logger), + ), + "kmc", + 0, # run in the background + task_name="kmer-counting", + ) + return kmc_job + + def build_extract_kmc_jobs( + self, + kmer_prefix: pathlib.Path, + ext_fastq: pathlib.Path, + ) -> Tuple[Job, Job]: + """Extract reads from the aligned short-read input and count k-mers + across the short and long reads in a single pass""" + assert self.reference is not None + assert self.model_bundle is not None + + # ReadWriter cannot write to /dev/stdout directly; pre-create a + # symlink so the driver writes to a real path that resolves to its + # stdout (the pipe to pgutil extract). + rw_bam = self.tmp_dir.joinpath("extract-kmc-rw.bam") + symlink_job = Job( + Pipeline(Command("ln", "-sf", "/dev/stdout", str(rw_bam))), + "extract-kmc-symlink", + 1, + task_name="read-extraction", + ) + extract_kmc_job = Job( + cmds.cmd_hybrid_extract_kmc( + kmer_prefix, + ext_fastq, + self.sr_aln, + self.lr_kmc_pairs(), + self.reference, + self.model_bundle.joinpath(self.extract_model_name), + self.tmp_dir, + rw_bam, + memory=self.kmer_memory, + threads=self.cores, + ), + "extract-kmc", + self.cores, + task_name="read-extraction", + ) + return symlink_job, extract_kmc_job + + def lr_realign_stage(self, ctx: StageContext) -> Minimap2RealignStage: + """Re-align the long-read input to the linear reference genome""" + assert self.model_bundle is not None + + return Minimap2RealignStage( + ctx=ctx, + inputs=self.lr_aln, + model_bundle=self.model_bundle, + sample_name=self.sample_sm, + bam_format=self.bam_format, + input_ref=self.lr_input_ref, + minimap2_model=self.model_bundle.joinpath("minimap2_lr.model"), + ) + + def bwa_extract_stage( + self, + ctx: StageContext, + sample_bam: pathlib.Path, + sample_fastq: pathlib.Path, + ) -> BwaExtractStage: + """The bwa alignment and read extraction stage. + + The short reads are given the sample name and marked as short + reads (`LR:0`) so the hybrid caller can tell them apart. + """ + assert self.model_bundle is not None + assert self.fastq_readgroup is not None + + rg = copy.deepcopy(self.fastq_readgroup) + rg["SM"] = self.sample_sm + rg["LR"] = "0" + return BwaExtractStage( + ctx=ctx, + output_bam=sample_bam, + output_fastq=sample_fastq, + r1_fastq=self.r1_fastq, + r2_fastq=self.r2_fastq, + readgroup=( + "@RG\\t" + "\\t".join([f"{x[0]}:{x[1]}" for x in rg.items()]) + ), + extract_model=self.model_bundle.joinpath(self.extract_model_name), + bwa_model=self.model_bundle.joinpath("bwa.model"), + unzip=find_unzip(self.logger), + ) + + def build_haplotypes_job( + self, output_gbz: pathlib.Path, kmer_file: pathlib.Path + ) -> Job: + """Build vg haplotypes job""" + assert self.hapl is not None + assert self.gbz is not None + + haplotypes_job = Job( + cmds.cmd_vg_haplotypes( + output_gbz, + kmer_file, + self.hapl, + self.gbz, + threads=self.cores, + xargs=[ + "--include-reference", + "--diploid-sampling", + "--set-reference", + self.ref_name(), + ], + ), + "vg-haplotypes", + self.cores, + task_name="pangenome", + ) + return haplotypes_job + + def build_gfa_job( + self, output_gfa: pathlib.Path, input_gbz: pathlib.Path + ) -> Job: + """Build vg convert to GFA job""" + gfa_job = Job( + cmds.cmd_vg_convert_gfa( + output_gfa, + input_gbz, + threads=self.cores, + reference_name=self.ref_name(), + ), + "vg-convert-gfa", + 0, + task_name="pangenome", + ) + return gfa_job + + def build_graph_update_job( + self, + out_gfa: pathlib.Path, + in_gfa: pathlib.Path, + lr_aln: List[pathlib.Path], + bed: Optional[pathlib.Path] = None, + name: str = "graph-update", + ) -> Job: + """Update the personalized graph using the long-read alignments""" + driver = Driver( + reference=self.reference, + thread_count=self.cores, + input=list(lr_aln), + ) + driver.add_algo( + PGHapUpdateAlgo( + out_gfa, + gfa_file=in_gfa, + target_bed=bed, + prefix=self.contig_prefix(), + ) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + name, + self.cores, + task_name="pangenome-update", + ) + + def build_gfa2fa_job( + self, + pangenome_fasta: pathlib.Path, + pangenome_gfa: pathlib.Path, + ref_fai: pathlib.Path, + ) -> Job: + """Generate the FASTA sequences of the updated graph""" + shell = ( + cmds.cmd_vg_gfa2fa(pangenome_fasta, pangenome_gfa) + if self.gfa2fa_with_vg + else cmds.cmd_pgutil_gfa2fa( + pangenome_fasta, ref_fai, pangenome_gfa + ) + ) + return Job( + shell, + "gfa2fa", + 0, + task_name="pangenome", + ) + + def build_minimap2_lift_job( + self, + out_bam: pathlib.Path, + ext_fastq: pathlib.Path, + pangenome_fasta: pathlib.Path, + pangenome_gfa: pathlib.Path, + ) -> Job: + """Align the extracted reads to the personalized pangenome and lift + the alignments back to the linear reference""" + assert self.reference is not None + assert self.model_bundle is not None + + # With aligned input, the first readgroup of the first alignment + # file seeds the readgroup of the lifted alignment + rg_source = ( + self.fastq_readgroup + if self.fastq_readgroup + else self.sr_readgroups[0][0] + ) + rg = copy.deepcopy(rg_source) + rg["ID"] = rg["ID"] + "-pg" + rg["SM"] = self.sample_sm + rg["LR"] = "2" + mm2_job = Job( + cmds.cmd_minimap2_lift( + out_bam, + pangenome_fasta, + ext_fastq, + pangenome_gfa, + self.reference, + "@RG\\t" + "\\t".join([f"{x[0]}:{x[1]}" for x in rg.items()]), + self.model_bundle.joinpath("minimap2.model"), + threads=self.cores, + mm2_xargs=["--secondary=yes"], + lift_prefix=self.contig_prefix(), + ), + "mm2-lift", + self.cores, + task_name="pangenome-alignment", + ) + return mm2_job + + def _replace_rg_arg(self, aln_rg: Dict[str, str], lr: str) -> str: + """A `--replace_rg` argument setting the LR attribute of an input + readgroup""" + rg_id = aln_rg.get("ID") + sm = self.rgsm if self.rgsm else aln_rg.get("SM") + return f"{rg_id}=ID:{rg_id}\\tSM:{sm}\\tLR:{lr}" + + def build_replace_rg(self) -> List[List[str]]: + """`--replace_rg` arguments setting the LR readgroup attribute for + the variant calling stages. + + The rows match the variant-calling input order: the short-read + alignments, the lifted alignment, then the long-read alignments. + Alignments generated by the pipeline (the bwa alignment with LR:0 + and the lifted alignment with LR:2) already carry the LR attribute + and take an empty row. User-supplied alignments are rewritten, as + they cannot be assumed to carry the attribute: `--sr_aln` input + with LR:0 and `--lr_aln` input with LR:1. + """ + replace_rg: List[List[str]] = [] + if self.r1_fastq: + replace_rg.append([]) # the bwa alignment + else: + for aln_rgs in self.sr_readgroups: + replace_rg.append( + [self._replace_rg_arg(rg, "0") for rg in aln_rgs] + ) + replace_rg.append([]) # the lifted alignment + for aln_rgs in self.lr_readgroups: + replace_rg.append( + [self._replace_rg_arg(rg, "1") for rg in aln_rgs] + ) + return replace_rg + + def build_pangenomesv_job( + self, + out_vcf: pathlib.Path, + input_bams: List[pathlib.Path], + pangenome_gfa: pathlib.Path, + replace_rg: List[List[str]], + ) -> Job: + """Call SVs with the original, lifted, and long reads""" + driver = Driver( + reference=self.reference, + thread_count=self.cores, + input=input_bams, + replace_rg=replace_rg, + ) + driver.add_algo( + PangenomeSV( + out_vcf, + gfa_file=pangenome_gfa, + min_af=PANGENOME_SV_MIN_AF, + prefix=self.contig_prefix(), + ) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + "pangenome-sv", + self.cores, + task_name="sv-calling", + ) diff --git a/sentieon_cli/job.py b/sentieon_cli/job.py index dea8343..a6a15f7 100644 --- a/sentieon_cli/job.py +++ b/sentieon_cli/job.py @@ -2,19 +2,34 @@ Job objects """ -import asyncio -import sys -import time from typing import Dict, Optional -from .logging import get_logger -from .shell_pipeline import Context, Pipeline - -logger = get_logger(__name__) +from .shell_pipeline import Pipeline +from .util import sanitize class Job: - """A job for execution""" + """A unit of work: a shell pipeline plus its execution metadata. + + Fields: + + * ``shell`` -- the pipeline to run. + * ``name`` -- a human-readable label (not part of identity). + * ``threads`` -- CPU threads the job needs (a local scheduling budget). + * ``resources`` -- named resource counts (e.g. NUMA-node tokens). + * ``task_name`` -- the pipeline stage this job belongs to; it groups the + job's log files, so every shard of an operation shares one value. + * ``job_id`` -- ``{name}-{n}``, unique across a run (not part of + identity). + + A job's identity is keyed only on ``shell``: two jobs with the same + pipeline are equal (and collide in a DAG) regardless of the other fields. + So two Job objects built from identical pipelines are "the same job" to + the DAG even though each carries its own ``job_id``. + """ + + # Per-name counters backing ``job_id``; see ``reset_ids``. + _id_counters: Dict[str, int] = {} def __init__( self, @@ -22,53 +37,43 @@ def __init__( name: str, threads: int = 1, resources: Optional[Dict[str, int]] = None, - ): + *, + task_name: str, + ) -> None: self.shell = pipeline self.name = name self.threads = threads self.resources = {} if resources is None else resources + self.task_name = task_name + # Log file names are sanitized, so ids must stay unique after + # sanitization; the id itself keeps the readable name. + key = sanitize(name) + count = Job._id_counters.get(key, 0) + 1 + Job._id_counters[key] = count + self.job_id = f"{name}-{count}" + + @classmethod + def reset_ids(cls) -> None: + """Restart id numbering. - def __hash__(self): + Ids must stay unique for a whole run, which can execute more than one + DAG, so this is called once per CLI invocation -- never between DAGs. + """ + cls._id_counters.clear() + + def __hash__(self) -> int: return hash(self.shell) - def __eq__(self, other: object): + def __eq__(self, other: object) -> bool: if isinstance(other, Job): return self.shell == other.shell return False - def __ne__(self, other: object): + def __ne__(self, other: object) -> bool: return not self == other - def __repr__(self): - return f"Job({self.name})" - - def __str__(self): - return f"Job({self.name})" - - def run(self, dry_run: bool = False): - """Run a command""" - if dry_run: - print(self.shell) - return - - logger.info("running command: %s", self.shell) - t0 = time.time() - context = Context() - try: - asyncio.run( - self.shell.run(context, stdout=sys.stdout, stderr=sys.stderr) - ) - for subcommand in context.commands: - if not subcommand.proc: - logger.error("subcommand has no process: %s", subcommand) - continue - ret = asyncio.run(subcommand.proc.wait()) - if ret != 0 and not subcommand.fail_ok: - logger.error( - f"subcommand failed with code {ret}: {subcommand}" - ) - except Exception as e: - logger.error(f"Failed to run command: {e}") - finally: - asyncio.run(context.cleanup()) - logger.info("finished in: %s seconds", f"{time.time() - t0:.1f}") + def __repr__(self) -> str: + return f"Job({self.job_id})" + + def __str__(self) -> str: + return f"Job({self.job_id})" diff --git a/sentieon_cli/logging.py b/sentieon_cli/logging.py index bd5d15a..fe12b53 100644 --- a/sentieon_cli/logging.py +++ b/sentieon_cli/logging.py @@ -1,12 +1,60 @@ +""" +Logging configuration + +Handlers live only on the package logger; the loggers returned by +``get_logger`` are plain children that propagate to it. The package logger is +always at DEBUG so verbosity can be enforced per handler: the console handler +follows ``-v``/``-q``/``-d`` while the ``run.log`` file handler always records +DEBUG. +""" + import logging +import os +from typing import Optional + +LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s" + +_package_logger = logging.getLogger(__package__) +_package_logger.setLevel(logging.DEBUG) +_package_logger.propagate = False # avoid duplicates through the root logger + +_console_handler = logging.StreamHandler() +_console_handler.setFormatter(logging.Formatter(LOG_FORMAT)) +_console_handler.setLevel(logging.INFO) +_package_logger.addHandler(_console_handler) + +# The file handler installed by `add_file_handler`, kept so repeated setup in +# a single process replaces it rather than accumulating handlers. +_file_handler: Optional[logging.FileHandler] = None + + +def get_logger(name: str) -> logging.Logger: + """Return a module logger propagating to the package logger.""" + return logging.getLogger(name) + + +def set_console_level(level: int | str) -> None: + """Set the verbosity of the console handler.""" + _console_handler.setLevel(level) + -handler = logging.StreamHandler() -handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s")) +def add_file_handler(path: "os.PathLike[str] | str") -> logging.FileHandler: + """Attach a DEBUG file handler, replacing any previously added one.""" + global _file_handler + if _file_handler is not None: + remove_file_handler(_file_handler) + handler = logging.FileHandler(path, mode="w") + handler.setFormatter(logging.Formatter(LOG_FORMAT)) + handler.setLevel(logging.DEBUG) + _package_logger.addHandler(handler) + _file_handler = handler + return handler -def get_logger(name: str): - """Return a logger with a StreamHandler.""" - logger = logging.getLogger(name) - logger.addHandler(handler) - logger.propagate = False - return logger +def remove_file_handler(handler: logging.FileHandler) -> None: + """Detach and close a handler added by ``add_file_handler``.""" + global _file_handler + _package_logger.removeHandler(handler) + handler.close() + if _file_handler is handler: + _file_handler = None diff --git a/sentieon_cli/pangenome_meta.py b/sentieon_cli/pangenome_meta.py new file mode 100644 index 0000000..000eb42 --- /dev/null +++ b/sentieon_cli/pangenome_meta.py @@ -0,0 +1,782 @@ +""" +Read pangenome graph metadata without loading the graph. + +A GBZ file is a `simple-sds` serialization of:: + + GBZ header (16 B) | GBZ tags | GBWT | GBWTGraph sequences ... + +and the GBWT inside it is:: + + GBWT header (48 B) | GBWT tags | BWT | DA samples | Option + +Every component carries its own length, so the multi-GB BWT body is +skipped with a seek and only the few MB holding the path names are read. +Parsing a whole-genome HPRC graph this way takes well under a second. + +The metadata answers the two questions the pangenome pipelines ask before +they run: + +* `pangenome_ref_name`, the sample name that `vg haplotypes + --set-reference` and `vg convert -Q` expect. A wrong name is not an + error for `vg`: it warns, writes a sampled graph with an empty + `reference_samples` tag and no reference paths, and then converts a GFA + with no `SN:Z:` rGFA tags. `pgutil lift` exits 0 with every read + unmapped, so the run completes with silently ruined accuracy. +* `pangenome_contig_prefix`, the `##` prefix of the + reference path names, stripped by `pgutil lift` and by the `PangenomeSV` + and `PGHapUpdate` algos. + +The backbone reference is the sample in the `reference_samples` tag whose +paths are unfragmented: one path per contig, all with fragment 0. Any +other listed sample is chopped into fragments by the graph construction. + +The module also runs as a script, for the run-time checks the pipelines +wire into their DAGs:: + + python -m sentieon_cli.pangenome_meta detect --gbz graph.gbz + python -m sentieon_cli.pangenome_meta check-gbz --gbz sample.gbz + --reference_name GRCh38 --contig_prefix 'GRCh38#0#' + python -m sentieon_cli.pangenome_meta check-gfa --gfa sample.gfa + --reference_name GRCh38 --contig_prefix 'GRCh38#0#' +""" + +from __future__ import annotations + +import argparse +import pathlib +import struct +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from typing import BinaryIO, Dict, List, Optional, Sequence, Tuple, Union + +from .logging import get_logger + +# Spelled out rather than taken from `__name__`: run as a script with +# `python -m sentieon_cli.pangenome_meta`, `__name__` is "__main__" and the +# logger would no longer be a child of the package logger, so its messages +# would miss the console handler and the run log. +logger = get_logger("sentieon_cli.pangenome_meta") + +# Bytes in a `simple-sds` word +WORD = 8 + +# Magic numbers of the serialized structures +GBZ_TAG = 0x205A4247 +GBWT_TAG = 0x6B376B37 +METADATA_TAG = 0x6B375E7A +HAPL_MAGIC = 0x4C504148 + +# How much of a GFA `check_gfa` reads before giving up on the `SN:Z:` tags +DEFAULT_MAX_GFA_LINES = 100000 + +# A GBWT path name record: (sample, contig, phase, fragment) +PathName = Tuple[int, int, int, int] + +PathLike = Union[str, pathlib.Path] + + +class PangenomeMetadataError(ValueError): + """A pangenome file could not be parsed, or does not describe one + unambiguous reference sample.""" + + +def unpack_bits(words: bytes, count: int, width: int) -> List[int]: + """Unpack `count` little-endian `width`-bit integers from 64-bit words. + + `simple-sds` packs an integer vector into a stream of 64-bit words with + no padding, so an element may straddle a word boundary. The running + buffer keeps the decode linear in `count`. + """ + mask = (1 << width) - 1 + out: List[int] = [] + buf = 0 + nbuf = 0 + offset = 0 + while len(out) < count: + while nbuf < width: + end = offset + WORD + buf |= int.from_bytes(words[offset:end], "little") << nbuf + nbuf += 64 + offset = end + out.append(buf & mask) + buf >>= width + nbuf -= width + return out + + +class _Reader: + """A cursor over the `simple-sds` serialization of a GBZ or GBWT. + + Every method reads (or skips) exactly one serialized structure and + leaves the cursor on the next one. + """ + + def __init__(self, fh: BinaryIO, path: PathLike) -> None: + self.fh = fh + self.path = str(path) + self.size = fh.seek(0, 2) + fh.seek(0) + + def _fail(self, msg: str) -> PangenomeMetadataError: + return PangenomeMetadataError(f"{self.path}: {msg}") + + def _read(self, nbytes: int) -> bytes: + data = self.fh.read(nbytes) + if len(data) != nbytes: + raise self._fail("the file ends before its contents do") + return data + + def _checked(self, nbytes: int) -> int: + """Reject a length that cannot fit in the file""" + if nbytes < 0 or nbytes > self.size: + raise self._fail( + f"implausible length {nbytes} in the file structure; the " + "file may be truncated or may not be a pangenome file" + ) + return nbytes + + # --- simple-sds primitives ------------------------------------------- + def u64(self) -> int: + return int(struct.unpack(" Tuple[int, int]: + low, high = struct.unpack(" None: + self.fh.seek(self._checked(nbytes), 1) + + def vec_u8(self) -> bytes: + count = self._checked(self.u64()) + data = self._read(count) + self._read((-count) % WORD) # padding to a whole word + return data + + def skip_vec_u8(self) -> None: + count = self.u64() + self.skip(count + (-count) % WORD) + + def vec_u64_words(self) -> bytes: + return self._read(self._checked(self.u64() * WORD)) + + def skip_vec_u64(self) -> None: + self.skip(self.u64() * WORD) + + def skip_option(self) -> None: + """Skip an `Option<...>` stored as a word count plus its words""" + self.skip(self.u64() * WORD) + + def raw_vector(self) -> Tuple[int, bytes]: + nbits = self.u64() + return nbits, self.vec_u64_words() + + def skip_raw_vector(self) -> None: + self.u64() + self.skip_vec_u64() + + def int_vector(self) -> List[int]: + count = self._checked(self.u64()) + width = self.u64() + if not 0 < width <= 64: + raise self._fail(f"unsupported integer width {width}") + nbits, words = self.raw_vector() + if nbits != count * width: + raise self._fail("inconsistent integer vector") + return unpack_bits(words, count, width) + + def skip_int_vector(self) -> None: + self.u64() + self.u64() + self.skip_raw_vector() + + def bit_vector(self) -> Tuple[int, bytes]: + self.u64() # the number of set bits + nbits, words = self.raw_vector() + self.skip_option() # rank support + self.skip_option() # select support + self.skip_option() # select-zero support + return nbits, words + + def skip_bit_vector(self) -> None: + self.u64() + self.skip_raw_vector() + self.skip_option() + self.skip_option() + self.skip_option() + + def sparse_vector_ones(self) -> List[int]: + """The positions of the set bits of an Elias-Fano sparse vector""" + self.u64() # universe size + _nbits, words = self.bit_vector() + count = self._checked(self.u64()) + width = self.u64() + lbits, lwords = self.raw_vector() + if lbits != count * width: + raise self._fail("inconsistent sparse vector") + low = unpack_bits(lwords, count, width) + + # The high part is unary-coded: the i-th set bit at position `p` + # contributes the bucket `p - i`. + result: List[int] = [] + index = 0 + for offset in range(0, len(words), WORD): + end = offset + WORD + word = int.from_bytes(words[offset:end], "little") + base = offset * 8 + while word: + lsb = word & -word + position = base + lsb.bit_length() - 1 + word ^= lsb + bucket = position - index + result.append((bucket << width) | low[index]) + index += 1 + if index == count: + return result + return result + + def skip_sparse_vector(self) -> None: + self.u64() + self.skip_bit_vector() + self.skip_int_vector() + + # --- gbwt-rs structures ---------------------------------------------- + def string_array(self) -> List[str]: + """A `StringArray`: offsets, a byte alphabet and packed indices""" + starts = self.sparse_vector_ones() + alphabet = self.vec_u8() + packed = self.int_vector() + try: + data = bytes(alphabet[index] for index in packed) + except IndexError as err: + raise self._fail( + "a string array indexes its alphabet out of range" + ) from err + starts.append(len(data)) + return [ + data[begin:end].decode("utf-8", "replace") + for begin, end in zip(starts, starts[1:]) + ] + + def tags(self) -> Dict[str, str]: + """A `Tags`: a string array of alternating keys and values""" + array = self.string_array() + return {array[2 * i]: array[2 * i + 1] for i in range(len(array) // 2)} + + def dictionary(self) -> List[str]: + """A `Dictionary`: a string array plus a sorted-id index""" + strings = self.string_array() + self.skip_int_vector() # sorted_ids + return strings + + +@dataclass +class GBZMetadata: + """The headers, tags and path names of a GBZ file""" + + gbz_version: int + gbz_tags: Dict[str, str] + gbwt_version: int + gbwt_tags: Dict[str, str] + sample_names: List[str] + contig_names: List[str] + path_names: List[PathName] + sample_count: int = 0 + haplotype_count: int = 0 + contig_count: int = 0 + sequence_count: int = 0 + bytes_read: int = 0 + + @property + def reference_samples(self) -> List[str]: + """The sample names listed by the `reference_samples` GBWT tag""" + return self.reference_samples_tag.split() + + @property + def reference_samples_tag(self) -> str: + """The raw `reference_samples` GBWT tag, `''` when it is absent""" + return self.gbwt_tags.get("reference_samples", "") + + def sample_of(self, path: PathName) -> str: + """The sample name of a path record""" + return self.sample_names[path[0]] + + def contig_of(self, path: PathName) -> str: + """The contig name of a path record""" + return self.contig_names[path[1]] + + def full_path_name(self, path: PathName) -> str: + """The `##` name of a path record""" + return f"{self.sample_of(path)}#{path[2]}#{self.contig_of(path)}" + + +@dataclass +class HaplHeader: + """The header of a `vg haplotypes` (.hapl) file""" + + version: int + top_level_chains: int + construction_jobs: int + total_subchains: int + total_kmers: int + k: int + tags: Dict[str, str] = field(default_factory=dict) + + +@dataclass +class PangenomeReference: + """The backbone reference sample of a pangenome graph""" + + ref_name: str + contig_prefix: str + contigs: List[str] + reference_samples: List[str] + n_paths: int + + +@dataclass +class _SampleSummary: + """What the path names say about one sample""" + + n_paths: int + phases: List[int] + contigs: List[str] + fragmented: bool + + +def read_gbz_metadata(path: PathLike) -> GBZMetadata: + """Read the tags and path names of a GBZ file. + + Only the file's headers, tags and metadata block are read; the BWT and + the graph sequences are skipped with seeks. + + Raises: + PangenomeMetadataError: the file is not a GBZ, is truncated, or + carries no GBWT metadata. + """ + try: + with open(path, "rb") as fh: + return _read_gbz_metadata(fh, path) + except PangenomeMetadataError: + raise + except (struct.error, IndexError, ValueError, EOFError) as err: + raise PangenomeMetadataError( + f"{path}: the file could not be parsed as a GBZ pangenome " + f"({err})" + ) from err + + +def _read_gbz_metadata(fh: BinaryIO, path: PathLike) -> GBZMetadata: + reader = _Reader(fh, path) + + tag, gbz_version = reader.u32x2() + if tag != GBZ_TAG: + raise PangenomeMetadataError(f"{path}: not a GBZ pangenome file") + reader.u64() # GBZ flags + gbz_tags = reader.tags() + + tag, gbwt_version = reader.u32x2() + if tag != GBWT_TAG: + raise PangenomeMetadataError( + f"{path}: the GBWT header was not found where the GBZ format " + "puts it" + ) + sequence_count = reader.u64() + reader.u64() # size + reader.u64() # offset + reader.u64() # alphabet size + reader.u64() # GBWT flags + gbwt_tags = reader.tags() + + # The BWT body: a sparse-vector index and a (possibly compressed) + # Vec. Both are skipped without reading their contents. + reader.skip_sparse_vector() + reader.skip_vec_u8() + reader.skip_vec_u64() # the document array samples + + # Option + if reader.u64() == 0: + raise PangenomeMetadataError( + f"{path}: the GBWT carries no metadata, so its path names are " + "unknown" + ) + tag, _meta_version = reader.u32x2() + if tag != METADATA_TAG: + raise PangenomeMetadataError( + f"{path}: the GBWT metadata header was not found" + ) + sample_count = reader.u64() + haplotype_count = reader.u64() + contig_count = reader.u64() + reader.u64() # metadata flags + n_paths = reader._checked(reader.u64()) + raw = reader._read(n_paths * 16) + path_names: List[PathName] = [ + (rec[0], rec[1], rec[2], rec[3]) + for rec in struct.iter_unpack(" HaplHeader: + """Read the header of a `vg haplotypes` (.hapl) file. + + Raises: + PangenomeMetadataError: the file is not a `vg haplotypes` file or + is too short to hold a header. + """ + try: + with open(path, "rb") as fh: + reader = _Reader(fh, path) + magic, version = reader.u32x2() + if magic != HAPL_MAGIC: + raise PangenomeMetadataError( + f"{path}: not a vg haplotypes file" + ) + return HaplHeader( + version=version, + top_level_chains=reader.u64(), + construction_jobs=reader.u64(), + total_subchains=reader.u64(), + total_kmers=reader.u64(), + k=reader.u64(), + # Tags were added to the header in version 6 + tags=reader.tags() if version >= 6 else {}, + ) + except PangenomeMetadataError: + raise + except (struct.error, IndexError, ValueError, EOFError) as err: + raise PangenomeMetadataError( + f"{path}: the file could not be parsed as a vg haplotypes " + f"file ({err})" + ) from err + + +def _summarize_samples(meta: GBZMetadata) -> Dict[str, _SampleSummary]: + """Group the path names of a GBZ by sample. + + A sample is "fragmented" when any of its paths carries a non-zero + fragment, or when one (contig, phase) pair has more than one path. + """ + by_sample: Dict[int, List[PathName]] = defaultdict(list) + for path in meta.path_names: + by_sample[path[0]].append(path) + + summaries: Dict[str, _SampleSummary] = {} + for sample_id, paths in by_sample.items(): + if sample_id >= len(meta.sample_names): + continue + per_contig: Dict[Tuple[int, int], int] = defaultdict(int) + for path in paths: + per_contig[(path[1], path[2])] += 1 + summaries[meta.sample_names[sample_id]] = _SampleSummary( + n_paths=len(paths), + phases=sorted({path[2] for path in paths}), + contigs=sorted({meta.contig_of(path) for path in paths}), + fragmented=( + any(path[3] != 0 for path in paths) + or any(count > 1 for count in per_contig.values()) + ), + ) + return summaries + + +def detect_pangenome_reference(gbz_path: PathLike) -> PangenomeReference: + """Detect the backbone reference sample of a pangenome graph. + + The GBWT `reference_samples` tag lists the graph's reference samples, + for example `"GRCh38 CHM13"`. Exactly one of them is the backbone: its + paths are unfragmented, one per contig. The others were chopped into + fragments when the graph was built. + + Raises: + PangenomeMetadataError: the file is not a GBZ, has no metadata, + names no reference sample, or names no single unfragmented + one. + """ + meta = read_gbz_metadata(gbz_path) + ref_samples = meta.reference_samples + if not ref_samples: + raise PangenomeMetadataError( + f"{gbz_path}: the GBWT 'reference_samples' tag is missing or " + "empty, so the graph names no reference sample" + ) + + summaries = _summarize_samples(meta) + present = [name for name in ref_samples if name in summaries] + if not present: + raise PangenomeMetadataError( + f"{gbz_path}: no sample named by the 'reference_samples' tag " + f"({' '.join(ref_samples)}) has any path in the graph" + ) + + backbone = [name for name in present if not summaries[name].fragmented] + if not backbone: + raise PangenomeMetadataError( + f"{gbz_path}: every reference sample ({', '.join(present)}) is " + "split into fragments, so the backbone reference cannot be " + "identified" + ) + if len(backbone) > 1: + raise PangenomeMetadataError( + f"{gbz_path}: more than one reference sample is unfragmented " + f"({', '.join(backbone)}), so the backbone reference is " + "ambiguous" + ) + + name = backbone[0] + summary = summaries[name] + if len(summary.phases) != 1: + phases = ", ".join(str(phase) for phase in summary.phases) + raise PangenomeMetadataError( + f"{gbz_path}: the reference sample '{name}' has paths on more " + f"than one haplotype ({phases}), so its contig prefix is " + "ambiguous" + ) + + reference = PangenomeReference( + ref_name=name, + contig_prefix=f"{name}#{summary.phases[0]}#", + contigs=summary.contigs, + reference_samples=ref_samples, + n_paths=summary.n_paths, + ) + logger.debug( + "Detected pangenome reference '%s' (contig prefix '%s', %d paths) " + "in '%s'", + reference.ref_name, + reference.contig_prefix, + reference.n_paths, + gbz_path, + ) + return reference + + +def check_sample_gbz( + gbz_path: PathLike, + ref_name: str, + contig_prefix: str, +) -> None: + """Check that a sampled pangenome kept its reference paths. + + `vg haplotypes --set-reference NAME` only warns when `NAME` is not a + reference sample of the input graph. The sampled graph it writes then + has an empty `reference_samples` tag and no reference paths at all, + which every downstream step turns into empty output. + + Raises: + PangenomeMetadataError: the file could not be read, its + `reference_samples` tag is not exactly `ref_name`, or the + reference sample has no path under `contig_prefix`. + """ + meta = read_gbz_metadata(gbz_path) + observed = meta.reference_samples + if observed != [ref_name]: + raise PangenomeMetadataError( + f"{gbz_path}: the sampled pangenome has " + f"reference_samples='{meta.reference_samples_tag}', expected " + f"'{ref_name}'. The pangenome reference name may be wrong: " + "`vg haplotypes --set-reference` drops every reference path " + "when the name it is given is not a reference sample of the " + "input graph." + ) + + reference_paths = [ + path for path in meta.path_names if meta.sample_of(path) == ref_name + ] + if not any( + meta.full_path_name(path).startswith(contig_prefix) + for path in reference_paths + ): + raise PangenomeMetadataError( + f"{gbz_path}: none of the {len(reference_paths)} paths of the " + f"reference sample '{ref_name}' starts with the contig prefix " + f"'{contig_prefix}'. The pangenome reference name or contig " + "prefix may be wrong." + ) + logger.info( + "The sampled pangenome '%s' has %d reference paths for '%s'", + gbz_path, + len(reference_paths), + ref_name, + ) + + +def _gfa_header_reference_samples(header: str) -> List[str]: + """The samples named by the `RS:Z:` tag of a GFA header line""" + for gfa_field in header.rstrip("\n").split("\t")[1:]: + if gfa_field.startswith("RS:Z:"): + return gfa_field.removeprefix("RS:Z:").split() + return [] + + +def check_gfa( + gfa_path: PathLike, + ref_name: str, + contig_prefix: str, + max_lines: int = DEFAULT_MAX_GFA_LINES, +) -> None: + """Check that a GFA carries the rGFA tags of the reference. + + `vg convert -f -Q NAME` writes an rGFA whose `S` lines carry + `SN:Z:##`, and lists the graph's remaining + reference samples in the header's `RS:Z:` tag. When `NAME` does not + match a reference sample, `vg` writes no `SN:Z:` tags at all and + leaves `NAME` in `RS:Z:`; `pgutil lift` then leaves every read + unmapped. + + Only the first `max_lines` lines are read: a whole-genome GFA runs to + many GB, and its first segment lines already carry the tags. + + Raises: + PangenomeMetadataError: the header lists `ref_name` among its + reference samples, or no segment line in the head of the file + carries an `SN:Z:` tag. + """ + sn_tag = f"SN:Z:{contig_prefix}" + header_checked = False + lines_read = 0 + with open(gfa_path, "r") as fh: + for line in fh: + if lines_read >= max_lines: + break + lines_read += 1 + if line.startswith("H\t") and not header_checked: + header_checked = True + rs_samples = _gfa_header_reference_samples(line) + if ref_name in rs_samples: + raise PangenomeMetadataError( + f"{gfa_path}: the GFA header still lists " + f"'{ref_name}' among its reference samples " + f"(RS:Z:{' '.join(rs_samples)}), so the rGFA tags " + "were written for a different reference. The " + "pangenome reference name may be wrong." + ) + elif line.startswith("S\t") and sn_tag in line: + logger.info( + "The GFA '%s' carries '%s' rGFA tags", gfa_path, sn_tag + ) + return + + raise PangenomeMetadataError( + f"{gfa_path}: no segment line in the first {lines_read} lines " + f"carries an '{sn_tag}' tag. `vg convert -Q` writes no rGFA tags " + "when its reference name does not match a reference sample of the " + "graph, and `pgutil lift` then leaves every read unmapped." + ) + + +def _build_parser() -> argparse.ArgumentParser: + """The argument parser of the `pangenome_meta` script""" + parser = argparse.ArgumentParser( + prog="python -m sentieon_cli.pangenome_meta", + description="Inspect and check pangenome graph metadata", + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + detect = subparsers.add_parser( + "detect", + help="Print the reference sample detected in a pangenome GBZ", + ) + detect.add_argument( + "--gbz", required=True, type=pathlib.Path, help="The GBZ file" + ) + + check_gbz = subparsers.add_parser( + "check-gbz", + help="Check that a sampled GBZ kept its reference paths", + ) + check_gbz.add_argument( + "--gbz", + required=True, + type=pathlib.Path, + help="The sampled GBZ file", + ) + check_gbz.add_argument( + "--reference_name", + required=True, + help="The expected pangenome reference sample name", + ) + check_gbz.add_argument( + "--contig_prefix", + required=True, + help="The expected pangenome contig prefix", + ) + + check_gfa_parser = subparsers.add_parser( + "check-gfa", + help="Check that a GFA carries the reference's rGFA tags", + ) + check_gfa_parser.add_argument( + "--gfa", required=True, type=pathlib.Path, help="The GFA file" + ) + check_gfa_parser.add_argument( + "--reference_name", + required=True, + help="The expected pangenome reference sample name", + ) + check_gfa_parser.add_argument( + "--contig_prefix", + required=True, + help="The expected pangenome contig prefix", + ) + check_gfa_parser.add_argument( + "--max_lines", + type=int, + default=DEFAULT_MAX_GFA_LINES, + help=( + "Stop looking for the rGFA tags after this many lines " + f"(default: {DEFAULT_MAX_GFA_LINES})" + ), + ) + return parser + + +def main(argv: Optional[Sequence[str]] = None) -> int: + """Run a pangenome metadata check; return 0 on success and 1 on + failure""" + args = _build_parser().parse_args(argv) + try: + if args.command == "detect": + reference = detect_pangenome_reference(args.gbz) + print(f"pangenome_ref_name\t{reference.ref_name}") + print(f"pangenome_contig_prefix\t{reference.contig_prefix}") + print( + "reference_samples\t" + " ".join(reference.reference_samples) + ) + print(f"n_paths\t{reference.n_paths}") + print(f"n_contigs\t{len(reference.contigs)}") + print("contigs\t" + ",".join(reference.contigs)) + elif args.command == "check-gbz": + check_sample_gbz(args.gbz, args.reference_name, args.contig_prefix) + else: + check_gfa( + args.gfa, + args.reference_name, + args.contig_prefix, + max_lines=args.max_lines, + ) + except (PangenomeMetadataError, OSError) as err: + logger.error("%s", err) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 39672c1..ed17af3 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -4,22 +4,38 @@ from abc import ABC, abstractmethod import argparse +import json import multiprocessing as mp import os import pathlib import shutil import sys -from typing import Any, Dict, List, Optional +import time +from typing import Any, Dict, List, Optional, TypeVar import packaging.version from . import command_strings as cmds from .dag import DAG +from .exceptions import DagExecutionError from .executor import BaseExecutor, DryRunExecutor, LocalExecutor from .job import Job -from .logging import get_logger +from .logging import get_logger, set_console_level +from .run_logs import RunLogs from .scheduler import ThreadScheduler -from .util import __version__, check_version, path_arg, tmp +from .shard import ( + detect_reference_build, + par_bed_for_build, +) +from .stages.base import StageContext +from .stages.metrics import MetricsPaths +from .util import ( + SampleSex, + __version__, + path_arg, + tmp, + versions_available, +) MULTIQC_MIN_VERSION = { "multiqc": packaging.version.Version("1.18"), @@ -27,6 +43,8 @@ BWA_INDEX_SUFFIXES = (".amb", ".ann", ".bwt", ".pac", ".sa") +T = TypeVar("T") + class BasePipeline(ABC): """A pipeline base class""" @@ -52,6 +70,13 @@ class BasePipeline(ABC): "help": "Print the commands without running them.", "action": "store_true", }, + "log_dir": { + "help": ( + "Directory for the run's log files. Defaults to the output " + "VCF with the '.vcf.gz' suffix replaced by '_logs'." + ), + "type": path_arg(), + }, # Hidden arguments "retain_tmpdir": { "help": argparse.SUPPRESS, @@ -71,17 +96,17 @@ class BasePipeline(ABC): @classmethod def add_arguments(cls, parser: argparse.ArgumentParser): - for k, kwargs in cls.params.items(): - flags = ["--" + k] + # Build a fresh kwargs dict per argument so repeated calls do not + # mutate the shared class-level params/positionals specs. + for k, spec in cls.params.items(): + kwargs = dict(spec) + flags = kwargs.pop("flags", ["--" + k]) if "default" in kwargs and "type" not in kwargs: kwargs["type"] = type(kwargs["default"]) - if "flags" in kwargs: - flags = kwargs["flags"] - del kwargs["flags"] parser.add_argument(*flags, **kwargs) - for k, kwargs in cls.positionals.items(): - parser.add_argument(k, **kwargs) + for k, spec in cls.positionals.items(): + parser.add_argument(k, **dict(spec)) def handle_arguments(self, args: argparse.Namespace): """Update self using the argparse object""" @@ -104,13 +129,6 @@ def handle_arguments(self, args: argparse.Namespace): if k in args.__dict__: setattr(self, k, getattr(args, k)) - def setup_logging(self, args: argparse.Namespace) -> None: - self.logger = get_logger(__name__) - if not self.logger.parent: - raise RuntimeError("Logger has no parent logger") - self.logger.parent.setLevel(args.loglevel) - self.logger.info("Starting sentieon-cli version: %s", __version__) - def __init__(self) -> None: self.reference: Optional[pathlib.Path] = None self.cores = mp.cpu_count() @@ -118,25 +136,95 @@ def __init__(self) -> None: self.dry_run = False self.retain_tmpdir = False self.skip_version_check = False + self.log_dir: Optional[pathlib.Path] = None self.output_vcf: Optional[pathlib.Path] = None + self.run_logs: Optional[RunLogs] = None + self.sample_sex: Optional[SampleSex] = None + self.cnv_par_bed: Optional[pathlib.Path] = None + self.reference_build: Optional[str] = None + + def setup_logging(self, args: argparse.Namespace) -> None: + """Configure console logging""" + self.logger = get_logger(__name__) + set_console_level(args.loglevel) + + def start_run_logs(self) -> None: + """Create the run's log directory and start writing `run.log`. + + Called after `validate`, so a run rejected for an invalid output path + neither creates directories nor clobbers a previous run's logs. + """ + # File logging is skipped for dry-runs and when there is nothing to + # derive a log directory from (a bare pipeline, as used by tests). + if self.dry_run or (self.log_dir is None and self.output_vcf is None): + self.logger.info("Starting sentieon-cli version: %s", __version__) + return + + log_dir = self.log_dir + if log_dir is None: + # Defensive: `validate` has already checked the output path, but + # not every pipeline's validation covers the suffix. + self.validate_output_suffix() + log_dir = pathlib.Path( + str(self.output_vcf).removesuffix(".vcf.gz") + "_logs" + ) + run_logs = RunLogs(log_dir) + try: + run_logs.setup() + except OSError as exc: + self.logger.error( + "Could not prepare the log directory %s: %s", log_dir, exc + ) + sys.exit(2) + self.run_logs = run_logs + + # After the file handler is attached, so the banner reaches run.log + self.logger.info("Starting sentieon-cli version: %s", __version__) + self.logger.info("Writing logs to: %s", self.run_logs.log_dir) + + def log_completion(self, success: bool, start_time: float) -> None: + """Report the outcome and duration of the run""" + self.logger.info( + "Finished sentieon-cli (status: %s, elapsed: %.1fs)", + "succeeded" if success else "failed", + time.monotonic() - start_time, + ) + if not success and self.run_logs: + self.logger.info( + "Logs from this run are in: %s", self.run_logs.log_dir + ) def main(self, args: argparse.Namespace) -> None: """Run the DNAscope pipeline""" self.handle_arguments(args) self.setup_logging(args) - self.validate() - self.configure() - - tmp_dir_str = tmp() - self.tmp_dir = pathlib.Path(tmp_dir_str) - - dag = self.build_dag() - executor = self.run(dag) - - if not self.retain_tmpdir: - shutil.rmtree(tmp_dir_str) - - self.check_execution(dag, executor) + start_time = time.monotonic() + success = False + try: + self.validate() + self.start_run_logs() + self.configure() + + tmp_dir_str = tmp() + self.tmp_dir = pathlib.Path(tmp_dir_str) + + try: + dag = self.build_dag() + executor = self.run(dag) + self.check_execution(dag, executor) + + # Jobs that depend on the results of the first DAG, such + # as sex-aware calling after ploidy estimation + second_dag = self.build_second_dag() + if second_dag is not None: + executor = self.run(second_dag) + self.check_execution(second_dag, executor) + finally: + if not self.retain_tmpdir: + shutil.rmtree(tmp_dir_str) + success = True + finally: + self.log_completion(success, start_time) def check_execution( self, @@ -145,23 +233,72 @@ def check_execution( ): """Check the DAG and executor after a run""" if executor.jobs_with_errors: - raise ValueError("Execution failed") + failed = ", ".join(str(job) for job in executor.jobs_with_errors) + message = f"Execution failed for jobs: {failed}" + if self.run_logs: + message += f"\nTask logs are in: {self.run_logs.task_logs}" + raise DagExecutionError(message) if len(dag.waiting_jobs) > 0 or len(dag.ready_jobs) > 0: - raise ValueError( + raise DagExecutionError( "The DAG has some unexecuted jobs\n" f"Waiting jobs: {dag.waiting_jobs}\n" f"Ready jobs: {dag.ready_jobs}\n" ) + def required(self, value: Optional[T], name: str) -> T: + """Return `value`, ending the run when it was not supplied. + + Narrows an optional pipeline attribute to its value, so callers do + not have to repeat the "missing argument" check. + """ + if value is None: + self.logger.error("%s is required", name) + sys.exit(2) + return value + + def stage_context(self) -> StageContext: + """The run-wide settings this pipeline's stages need. + + Built fresh on every call, so it always reflects the pipeline's + current attributes. Only available once `main` has created the + run's temporary directory. + """ + if self.reference is None: + self.logger.error("reference is required") + sys.exit(2) + if self.output_vcf is None: + self.logger.error("output_vcf is required") + sys.exit(2) + tmp_dir: Optional[pathlib.Path] = getattr(self, "tmp_dir", None) + if tmp_dir is None: + self.logger.error( + "The temporary directory has not been created yet; " + "`stage_context` is only available after `main` sets " + "`tmp_dir`" + ) + sys.exit(2) + return StageContext( + reference=self.reference, + output_vcf=self.output_vcf, + tmp_dir=tmp_dir, + cores=self.cores, + dry_run=self.dry_run, + skip_version_check=self.skip_version_check, + ) + @abstractmethod def validate(self) -> None: pass - def validate_output_vcf(self) -> None: + def validate_output_suffix(self) -> None: + """Confirm the output VCF file name ends in '.vcf.gz'""" if not str(self.output_vcf).endswith(".vcf.gz"): self.logger.error("The output file should end with '.vcf.gz'") sys.exit(2) + + def validate_output_vcf(self) -> None: + self.validate_output_suffix() assert self.output_vcf is not None parent = self.output_vcf.resolve().parent if not parent.is_dir(): @@ -207,25 +344,85 @@ def configure(self) -> None: def build_dag(self) -> DAG: pass + def build_second_dag(self) -> Optional[DAG]: + """Build a second DAG, run after the first one has finished. + + Pipelines with jobs that depend on the results of the first DAG, + such as sex-aware calling after ploidy estimation, override this + hook. Returning `None` runs a single DAG. + """ + return None + + def get_sex(self, ploidy_json: pathlib.Path) -> None: + """Retrieve the sample sex""" + if self.sample_sex is not None: + # Supplied through `--sample_sex` + return + if self.dry_run: + self.logger.info("Setting sample sex to MALE for dry-run") + self.sample_sex = SampleSex.MALE + return + with open(ploidy_json) as fh: + data = json.load(fh) + sex = data["sex"] + if sex == "female": + self.sample_sex = SampleSex.FEMALE + elif sex == "male": + self.sample_sex = SampleSex.MALE + else: + self.sample_sex = SampleSex.UNKNOWN + + def resolve_cnv_par_bed( + self, + fai_data: Dict[str, Dict[str, int]], + par_bed: Optional[pathlib.Path] = None, + cnv_will_run: bool = True, + ) -> None: + """Identify the reference build and select the PAR BED file. + + `self.reference_build` is always identified, as the ploidy + estimation contigs follow it. The PAR BED file is only used by + CNV calling, so it is looked up only when CNVs will be called. + """ + self.reference_build = detect_reference_build(fai_data) + if not cnv_will_run: + return + if par_bed is not None: + self.cnv_par_bed = par_bed + return + self.cnv_par_bed = par_bed_for_build(self.reference_build) + + def validate_cnv_par(self, cnv_will_run: bool) -> None: + """Confirm a PAR BED file is available for CNV calling. + + The check runs during validation, before any job starts, and does + not depend on the sample sex. + """ + if not cnv_will_run or self.cnv_par_bed is not None: + return + + self.logger.error( + "CNV calling uses a BED file of the pseudo-autosomal regions " + "(PAR), and no PAR BED file is available for this reference " + "genome. Please supply the `--par_bed` argument." + ) + sys.exit(2) + def multiqc(self) -> Optional[Job]: """Run MultiQC on the metrics files""" - if not self.skip_version_check: - if not all( - [ - check_version(cmd, min_version) - for (cmd, min_version) in MULTIQC_MIN_VERSION.items() - ] - ): - self.logger.warning( - "Skipping MultiQC. MultiQC version %s or later not found", - MULTIQC_MIN_VERSION["multiqc"], - ) - return None + if not versions_available( + MULTIQC_MIN_VERSION, skip=self.skip_version_check + ): + self.logger.warning( + "Skipping MultiQC. MultiQC version %s or later not found", + MULTIQC_MIN_VERSION["multiqc"], + ) + return None - metrics_dir = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_metrics") - ) + metrics_dir = MetricsPaths.from_output_vcf( + pathlib.Path(str(self.output_vcf)) + ).metrics_dir multiqc_job = Job( cmds.cmd_multiqc( metrics_dir, @@ -234,6 +431,7 @@ def multiqc(self) -> Optional[Job]: ), "multiqc", 0, + task_name="multiqc", ) return multiqc_job @@ -250,8 +448,17 @@ def run(self, dag: DAG) -> BaseExecutor: ) self.logger.debug("Creating the executor") - Executor = DryRunExecutor if self.dry_run else LocalExecutor - executor = Executor(scheduler) + executor: BaseExecutor + if self.dry_run: + executor = DryRunExecutor(scheduler) + else: + # Handle Ctrl-C/SIGTERM by terminating running jobs gracefully; + # the handlers are installed only for the duration of the run. + executor = LocalExecutor( + scheduler, + install_signal_handlers=True, + run_logs=self.run_logs, + ) self.logger.info("Starting execution") executor.execute() diff --git a/sentieon_cli/run_logs.py b/sentieon_cli/run_logs.py new file mode 100644 index 0000000..6d605b1 --- /dev/null +++ b/sentieon_cli/run_logs.py @@ -0,0 +1,343 @@ +""" +The log directory for a single run +""" + +from __future__ import annotations + +import dataclasses +import datetime +import logging +import pathlib +import shlex +import shutil +import sys +from typing import IO, Iterable, List, Optional, TYPE_CHECKING + +from .logging import add_file_handler, get_logger, remove_file_handler +from .util import __version__, sanitize + +if TYPE_CHECKING: + import resource + + from .job import Job + from .shell_pipeline import Command + +logger = get_logger(__name__) + +_STDERR = "stderr" +_STDOUT = "stdout" + +# The rusage-derived columns of a process row, empty as a group when a +# process was reaped without rusage. +_RUSAGE_COLUMNS = [ + "user_s", + "sys_s", + "maxrss_mib", + "minflt", + "majflt", + "inblock", + "oublock", + "nvcsw", + "nivcsw", +] + +PROCESS_METRICS_COLUMNS = [ + "job_id", + "task_name", + "pid", + "stage", + "exit_code", + *_RUSAGE_COLUMNS, + "command", +] + +JOB_METRICS_COLUMNS = [ + "job_id", + "task_name", + "threads", + "status", + "wall_s", + "user_s", + "sys_s", + "max_proc_rss_mib", + "processes", +] + +_MIB = 1024 * 1024 + + +def _tsv_field(value: object) -> str: + """Render a value as a TSV field that cannot break the row layout""" + text = "" if value is None else str(value) + for char in ("\t", "\n", "\r"): + text = text.replace(char, " ") + return text + + +def _write_row(path: pathlib.Path, fields: Iterable[object]) -> None: + """Append one row to a metrics file. + + A metrics file is a by-product of the run: a write failure (full disk, + read-only log directory) is reported once and otherwise ignored. + """ + row = "\t".join(_tsv_field(field) for field in fields) + try: + with open(path, "a") as handle: + handle.write(row + "\n") + except OSError as exc: + logger.warning("could not write %s: %s", path, exc) + + +@dataclasses.dataclass +class _ProcessLog: + """A log file opened for one spawned process""" + + command: Command + stage_index: int + path: pathlib.Path + kind: str + header_bytes: int + + +class JobLogSink: + """The per-process log files of a single job. + + Attached to the job's ``Context``, so every process spawned for the job -- + including process substitutions -- writes its stderr (and its stdout, when + that would otherwise be inherited) to its own file. + """ + + def __init__(self, task_logs: pathlib.Path, job: Job) -> None: + self.job_id = job.job_id + self.name = job.name + self.task_name = job.task_name + self.log_dir = task_logs / sanitize(job.task_name) + self._prefix = sanitize(job.job_id) + self._stage_index = 0 # Next unallocated stage index + self._logs: List[_ProcessLog] = [] + + def open_stderr(self, command: Command) -> IO[bytes]: + """Open the stderr log of the next process to spawn""" + return self._open(command, _STDERR) + + def open_stdout(self, command: Command) -> IO[bytes]: + """Open the stdout log of the next process to spawn""" + return self._open(command, _STDOUT) + + def _stage_for(self, command: Command) -> int: + """The stage index of a command, allocated on first use. + + Both streams of one process share its index, so an already-seen + command reuses it. Matched by identity: equal-but-distinct commands + are distinct processes. + """ + for log in self._logs: + if log.command is command: + return log.stage_index + stage_index = self._stage_index + self._stage_index += 1 + return stage_index + + def stage_index_for(self, command: Command) -> Optional[int]: + """The stage index already allocated to a command, if any. + + Unlike ``_stage_for`` this never allocates: a command whose streams + were all explicit (so no log was opened) has no stage index. + """ + for log in self._logs: + if log.command is command: + return log.stage_index + return None + + def _open(self, command: Command, kind: str) -> IO[bytes]: + self.log_dir.mkdir(parents=True, exist_ok=True) + stage_index = self._stage_for(command) + suffix = ".log" if kind == _STDERR else ".stdout.log" + path = self.log_dir / f"{self._prefix}.{stage_index}{suffix}" + handle = open(path, "wb") + header_bytes = 0 + if kind == _STDERR: + # The child inherits the file offset, so the header must be on + # disk before the FD is handed over or it would interleave with + # (or land after) the process's own output. Stdout logs get no + # header: they hold exactly what the tool would have printed. + header = self._header(command, stage_index) + handle.write(header) + handle.flush() + header_bytes = len(header) + self._logs.append( + _ProcessLog(command, stage_index, path, kind, header_bytes) + ) + return handle + + def _header(self, command: Command, stage_index: int) -> bytes: + timestamp = datetime.datetime.now().astimezone() + return ( + f"# timestamp: {timestamp.isoformat(timespec='seconds')}\n" + f"# task: {self.task_name}\n" + f"# job: {self.name} ({self.job_id})\n" + f"# stage: {stage_index}\n" + f"# command: {command}\n" + ).encode() + + def stderr_log_for(self, command: Command) -> Optional[pathlib.Path]: + """The stderr log of a command, matched by identity. + + Equal-but-distinct ``Command`` objects can share a context, so the + lookup must not fall back to ``__eq__``. + """ + for log in self._logs: + if log.kind == _STDERR and log.command is command: + return log.path + return None + + def log_paths(self) -> List[pathlib.Path]: + """Every log file opened for this job""" + return [log.path for log in self._logs] + + def finalize(self, success: bool) -> None: + """Drop the logs of a successful job that hold nothing but a header. + + Empty logs are common (cleanup jobs, unused stdout) and only add + clutter. After a failure everything is kept -- an empty log is itself + informative. + """ + if not success: + return + for log in self._logs: + try: + if log.path.stat().st_size <= log.header_bytes: + log.path.unlink() + except OSError: + pass + + +class RunLogs: + """The log directory of a single sentieon-cli invocation""" + + def __init__(self, log_dir: pathlib.Path) -> None: + self.log_dir = log_dir + self.run_log = log_dir / "run.log" + self.command_txt = log_dir / "command.txt" + self.task_logs = log_dir / "task_logs" + self.process_metrics = log_dir / "process_metrics.txt" + self.job_metrics = log_dir / "job_metrics.txt" + self.file_handler: Optional[logging.FileHandler] = None + + def setup(self) -> None: + """Prepare the log directory and start writing `run.log`""" + self.create_dirs() + self.write_command() + self.file_handler = add_file_handler(self.run_log) + + def create_dirs(self) -> None: + """Create the log directory, clearing logs from any previous run. + + The metrics files are (re-)written with their header row here, so a + rerun into the same directory replaces the previous run's rows and a + run with no finished jobs still leaves a readable, empty table. + """ + # The log directory itself is never removed - the user may point + # `--log_dir` at an existing directory holding unrelated files. + self.log_dir.mkdir(parents=True, exist_ok=True) + if self.task_logs.exists(): + shutil.rmtree(self.task_logs) + self.task_logs.mkdir(parents=True) + self.process_metrics.write_text( + "\t".join(PROCESS_METRICS_COLUMNS) + "\n" + ) + self.job_metrics.write_text("\t".join(JOB_METRICS_COLUMNS) + "\n") + + def job_sink(self, job: Job) -> JobLogSink: + """Create the log sink for a job""" + return JobLogSink(self.task_logs, job) + + def record_process_metrics( + self, + *, + job: Job, + command: Command, + pid: int, + exit_code: int, + stage: Optional[int], + rusage: Optional["resource.struct_rusage"], + maxrss_bytes: Optional[int], + ) -> None: + """Append one process to `process_metrics.txt`. + + Written for every reaped process, whatever the job's outcome; the + rusage columns are empty for a process reaped without one. + """ + if rusage is None: + usage: List[object] = [None] * len(_RUSAGE_COLUMNS) + else: + rss = "" if maxrss_bytes is None else f"{maxrss_bytes / _MIB:.1f}" + usage = [ + f"{rusage.ru_utime:.2f}", + f"{rusage.ru_stime:.2f}", + rss, + rusage.ru_minflt, + rusage.ru_majflt, + rusage.ru_inblock, + rusage.ru_oublock, + rusage.ru_nvcsw, + rusage.ru_nivcsw, + ] + _write_row( + self.process_metrics, + [job.job_id, job.task_name, pid, stage, exit_code] + + usage + + [command], + ) + + def record_job_metrics( + self, + *, + job: Job, + failed: bool, + wall_s: float, + user_s: Optional[float], + sys_s: Optional[float], + max_proc_rss_bytes: Optional[int], + processes: int, + ) -> None: + """Append one job to `job_metrics.txt`. + + Only jobs that reach accounting get a row: a job that failed to + launch, or one cut short by an interrupt, is absent. + """ + rss = ( + "" + if max_proc_rss_bytes is None + else f"{max_proc_rss_bytes / _MIB:.1f}" + ) + _write_row( + self.job_metrics, + [ + job.job_id, + job.task_name, + job.threads, + "failed" if failed else "ok", + f"{wall_s:.2f}", + "" if user_s is None else f"{user_s:.2f}", + "" if sys_s is None else f"{sys_s:.2f}", + rss, + processes, + ], + ) + + def write_command(self) -> None: + """Record the invocation so the run can be reproduced""" + timestamp = datetime.datetime.now().astimezone() + self.command_txt.write_text( + f"command: {shlex.join(sys.argv)}\n" + f"version: {__version__}\n" + f"directory: {pathlib.Path.cwd()}\n" + f"timestamp: {timestamp.isoformat(timespec='seconds')}\n" + ) + + def close(self) -> None: + """Stop writing `run.log`""" + if self.file_handler is not None: + remove_file_handler(self.file_handler) + self.file_handler = None diff --git a/sentieon_cli/scheduler.py b/sentieon_cli/scheduler.py index ffc0bb8..4ed5329 100644 --- a/sentieon_cli/scheduler.py +++ b/sentieon_cli/scheduler.py @@ -1,15 +1,48 @@ """Schedule jobs""" -from typing import Dict, Generator, Optional, Set +from abc import ABC, abstractmethod +from typing import Dict, Optional, Set from .dag import DAG +from .exceptions import DagExecutionError from .job import Job from .logging import get_logger logger = get_logger(__name__) -class ThreadScheduler: +class BaseScheduler(ABC): + """Selects which ready jobs an executor may run, subject to a policy. + + An executor drives a scheduler with two calls: + + 1. :meth:`start` once, to get the first batch of jobs to run. + 2. :meth:`job_finished` for each job as it completes, to get the jobs + that its completion unblocked (possibly an empty set). + + Contract / invariants: + + * :meth:`start` is called exactly once, before any :meth:`job_finished`. + * Every job passed to :meth:`job_finished` was previously returned by + :meth:`start` or :meth:`job_finished`, and is reported at most once. + * A scheduler instance drives a single run; do not reuse it. + * :meth:`start` may raise + :class:`~sentieon_cli.exceptions.DagExecutionError` if the DAG + contains a job that can never be scheduled. + * The executor stops when it has no running jobs and the scheduler + returns no new ones. + """ + + @abstractmethod + def start(self) -> Set[Job]: + """Return the first batch of jobs to run.""" + + @abstractmethod + def job_finished(self, job: Job) -> Set[Job]: + """Record a finished job and return the jobs it unblocked.""" + + +class ThreadScheduler(BaseScheduler): """Schedule jobs as threads are available""" def __init__( @@ -17,41 +50,65 @@ def __init__( dag: DAG, threads: int = 1, resources: Optional[Dict[str, int]] = None, - ): + ) -> None: self.dag = dag self.threads = threads self.available_threads = threads self.resources = {} if resources is None else resources + # Jobs that are ready to run but have not yet been scheduled. + self._ready: Dict[Job, None] = {} + + def _check_feasible(self) -> None: + """Ensure every job can eventually be scheduled. + + A job that requests more threads than the total budget, or more of a + managed resource than exists, can never be released and would + otherwise stall the run silently. + """ + jobs = list(self.dag.waiting_jobs) + list(self.dag.ready_jobs) + for job in jobs: + if job.threads > self.threads: + raise DagExecutionError( + f"{job} requests {job.threads} threads but the scheduler " + f"budget is only {self.threads}" + ) + for resource, needed in job.resources.items(): + capacity = self.resources.get(resource) + if capacity is not None and needed > capacity: + raise DagExecutionError( + f"{job} requests {needed} of resource '{resource}' " + f"but the scheduler capacity is only {capacity}" + ) + + def _fit(self) -> Set[Job]: + """Reserve budget for as many ready jobs as fit; return that set.""" + logger.debug("Ready pool: %s", self._ready) - def schedule(self) -> Generator[Set[Job], Optional[Job], None]: - """Schedule a job for execution""" - dag_gen = self.dag.update_dag() - ready_jobs = dag_gen.send(None) - - while True: - logger.debug("Ready jobs: %s", ready_jobs) - - scheduled_jobs: Set[Job] = set() - for ready_job in ready_jobs: - if self.available_threads - ready_job.threads >= 0 and all( - [ - available - ready_job.resources.get(resource, 0) >= 0 - for resource, available in self.resources.items() - ] - ): - self.available_threads -= ready_job.threads - for resource, used in ready_job.resources.items(): - if resource in self.resources: - self.resources[resource] -= used - scheduled_jobs.add(ready_job) - - for job in scheduled_jobs: - del ready_jobs[job] - - finished_job = yield scheduled_jobs - if isinstance(finished_job, Job): - self.available_threads += finished_job.threads - for resource, used in finished_job.resources.items(): + scheduled_jobs: Set[Job] = set() + for ready_job in self._ready: + if self.available_threads - ready_job.threads >= 0 and all( + available - ready_job.resources.get(resource, 0) >= 0 + for resource, available in self.resources.items() + ): + self.available_threads -= ready_job.threads + for resource, used in ready_job.resources.items(): if resource in self.resources: - self.resources[resource] += used - ready_jobs.update(dag_gen.send(finished_job)) + self.resources[resource] -= used + scheduled_jobs.add(ready_job) + + for job in scheduled_jobs: + del self._ready[job] + return scheduled_jobs + + def start(self) -> Set[Job]: + self._check_feasible() + self._ready = dict(self.dag.ready_jobs) + return self._fit() + + def job_finished(self, job: Job) -> Set[Job]: + self.available_threads += job.threads + for resource, used in job.resources.items(): + if resource in self.resources: + self.resources[resource] += used + self._ready.update(self.dag.mark_finished(job)) + return self._fit() diff --git a/sentieon_cli/scripts/estimate_ploidy.py b/sentieon_cli/scripts/estimate_ploidy.py index a5f3a5d..c42e0a3 100644 --- a/sentieon_cli/scripts/estimate_ploidy.py +++ b/sentieon_cli/scripts/estimate_ploidy.py @@ -20,14 +20,22 @@ ] DEFAULT_AUTOSOMES = DEFAULT_CONTIGS[:-2] -DEFAULT_SEXES = { - "male": {"chrX": 1, "chrY": 1}, - "female": {"chrX": 2, "chrY": 0}, -} +DEFAULT_X_CONTIG = "chrX" +DEFAULT_Y_CONTIG = "chrY" IDXSTATS_CMD = ["samtools", "idxstats"] +def sex_requirements( + x_contig: str, y_contig: str +) -> Dict[str, Dict[str, int]]: + """The sex chromosome ploidy expected for each sex""" + return { + "male": {x_contig: 1, y_contig: 1}, + "female": {x_contig: 2, y_contig: 0}, + } + + def process_args(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( @@ -49,6 +57,16 @@ def process_args(): default=DEFAULT_AUTOSOMES, help="Autosome contigs", ) + parser.add_argument( + "--x_contig", + default=DEFAULT_X_CONTIG, + help="The name of the X chromosome", + ) + parser.add_argument( + "--y_contig", + default=DEFAULT_Y_CONTIG, + help="The name of the Y chromosome", + ) parser.add_argument( "--outfile", default=sys.stdout, @@ -85,6 +103,10 @@ def main(args: argparse.Namespace): total_reads = sum( [x[1] for contig, x in idxstats_results.items() if contig in autosomes] ) + if total_bases == 0: + # No recognized autosomes, so ploidy cannot be estimated + json.dump({"contigs": {}, "sex": "Unknown"}, args.outfile, indent=2) + return sys.exit(os.EX_OK) average_reads_per_base = total_reads / total_bases # Calculate read coverage by contig @@ -111,14 +133,16 @@ def main(args: argparse.Namespace): # Determine sex sex = "Unknown" - for possible_sex, required_chroms in DEFAULT_SEXES.items(): + sexes = sex_requirements(args.x_contig, args.y_contig) + for possible_sex, required_chroms in sexes.items(): match = True for contig, req_ploidy in required_chroms.items(): try: found_ploidy = int(results["contigs"][contig]["ploidy"]) if found_ploidy != req_ploidy: match = False - except ValueError: + except (KeyError, ValueError): + # An absent or indeterminate sex chromosome match = False if match: diff --git a/sentieon_cli/scripts/sad_lad_update.py b/sentieon_cli/scripts/sad_lad_update.py new file mode 100755 index 0000000..6a3e3e8 --- /dev/null +++ b/sentieon_cli/scripts/sad_lad_update.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python3 +"""sad_lad_update.py -- alignment-free pangenome AD/DP update. + +The dnascope-pangenome joint pileup calls a bwa CRAM (all reads) together with +a minimap2 pangenome-realignment CRAM (a subset of the same reads, tagged +LR:1). Reads that appear in both CRAMs are counted twice, so FORMAT/AD = +SAD + LAD, FORMAT/DP and INFO/DP are inflated. Plain VCF and gVCF input are +both accepted. + +This script picks, per record and per sample, EITHER the short-read depth +vector (FORMAT/SAD) OR the pangenome one (FORMAT/LAD) as the updated +FORMAT/AD -- whichever likelihood vector (SPL / LPL) better supports the +called genotype -- and derives FORMAT/DP and INFO/DP from that choice. No +alignment file is read; only the fields already present in the VCF are used. +gVCF hom-ref blocks are the one exception: nothing there is chosen, read or +written (see "gVCF input" below). + +Decision rule, per record and sample. Let g be the PL index of the called +genotype (VCF ordering) and, for a likelihood vector X, + + conf(X) = min(X[j] for j != g) - X[g], capped at --conf_cap + +so conf is GQ-like: large positive means X strongly favours the called GT, a +negative value means X prefers a different genotype. conf is None +("uninformative") when X is missing, has fewer than two entries, or has no +value at index g. + + branch 1 GT missing (./., .) or SAD or LAD missing + -> sample left untouched, ADSRC=NONE, ADRULE=0 + branch 2 sum(LAD) == 0 -> SAD (nothing to gain; + sum(SAD) == 0 and sum(LAD) > 0 -> LAD this also covers + every LPL=. record) + branch 3 exactly one conf is None -> the informative one + both None -> SAD if sum(SAD) >= sum(LAD) else LAD + branch 4 conf_S - conf_L > --similar_threshold -> SAD + conf_L - conf_S > --similar_threshold -> LAD + branch 5 "similar": LAD, UNLESS sum(SAD) > (1 + --lad_ratio) * sum(LAD), + in which case SAD. + +ADRULE is written as 0 for branch 1 (untouched) as declared in the header, and +as 2..5 for the branches that actually choose a source; the value 1 is never +emitted. + +Output. FORMAT/AD becomes the chosen vector, FORMAT/DP its sum, INFO/DP the +sum of the new FORMAT/DP over the samples of the record. The originals are +kept as FORMAT/AD_ORIG, FORMAT/DP_ORIG and INFO/DP_ORIG, and the choice is +recorded as FORMAT/ADSRC (SAD/LAD/NONE) and FORMAT/ADRULE. For an untouched +sample AD and DP keep their original values and AD_ORIG / DP_ORIG are written +as '.' -- the '.' is the signal that nothing was updated there. + +gVCF input. A record whose ALT is exactly (or <*>) is a +homozygous-reference block and is copied to the output byte for byte: its +original line is emitted, nothing is parsed back out, and it is counted only +as records_ref_block -- never in samples_seen, the by-filter tables or +records_updated / records_untouched. No update is possible there +(a block carries no SAD/LAD, so no source can be chosen) and nothing may be +added either: INFO/DP or extra FORMAT keys on a block change what GVCFtyper +reads back, so the safe action and the right one are the same. gVCF +*variant* records are handled exactly like plain-VCF ones -- their trailing + allele is present in AD/SAD/LAD (Number=R) and in PL/SPL +(Number=G) alike, so the vectors stay parallel and gt_index() / conf_of() / +select() need no gVCF special case. Output naming needs none either: +'out.g.vcf.gz' satisfies the --output_vcf '.vcf.gz' suffix check. + +Implementation. Built on Sentieon's vcflib: records are read as parsed +Variant objects, v.info / v.samples are mutated in place, v.line is cleared +and vcflib's VCF.format() re-emits the record. Re-formatting is therefore +expected and intended -- format() sorts the INFO keys, writes FORMAT as GT +followed by the remaining keys in sorted order, prints QUAL as '%4.2f' and +re-encodes floats through str() (QUAL 0 -> 0.00, FS=0 -> FS=0.0). Values are +unchanged; only their spelling is. The header declarations are added through +vcflib's copy_header(update=...), so the new lines land at the end of their +##INFO / ##FORMAT group. vcflib also writes the output .tbi during the write +pass; it answers region queries like a tabix index but carries no htslib count +pseudo-bin, so `bcftools index -n` reports 0 on it. + +Usage: + sad_lad_update.py --input_vcf in.vcf.gz --output_vcf out.vcf.gz \\ + [--stats stats.tsv] [--similar_threshold 10] \\ + [--lad_ratio 0.10] [--conf_cap 99] [--threads N] + +With --threads > 1 vcflib buffers each shard in a temporary file under +$SENTIEON_TMPDIR (or the system temp dir); point it at a writable filesystem +with room for the uncompressed output (~4 GB for a WGS VCF). +""" + +import argparse +import json +import resource +import sys +import time + +try: + import vcflib +except ImportError: + sys.stderr.write( + "sad_lad_update: cannot import vcflib.\n" + " Install it with `pip install vcflib`, or point PYTHONPATH at a\n" + " vcflib checkout, e.g.\n" + " PYTHONPATH=/path/to/vcflib python3 sad_lad_update.py ...\n") + sys.exit(1) + +INFO_DP_ORIG_HDR = ('##INFO=') +FMT_AD_ORIG_HDR = ('##FORMAT=') +FMT_DP_ORIG_HDR = ('##FORMAT=') +FMT_ADSRC_HDR = ('##FORMAT=') +FMT_ADRULE_HDR = ('##FORMAT=') + +SRC_NONE, SRC_SAD, SRC_LAD = 0, 1, 2 +SRC_NAME = ('NONE', 'SAD', 'LAD') +FILTER_CLASSES = ('PASS', 'MLrejected', 'LowQual', 'other') + +COUNTERS = ( + 'records', + 'records_ref_block', + 'records_updated', + 'records_untouched', + 'records_no_samples', + 'samples_seen', + 'samples_updated', + 'samples_untouched', + 'src_SAD', 'src_LAD', 'src_NONE', + 'rule_0', 'rule_1', 'rule_2', 'rule_3', 'rule_4', 'rule_5', + 'info_DP_added', + 'fmt_AD_added', + 'fmt_DP_added', + 'AD_changed', + 'FMT_DP_changed', + 'INFO_DP_changed', +) +BYFILT_KEYS = ('records', 'src_SAD', 'src_LAD', 'src_NONE', + 'rule_0', 'rule_1', 'rule_2', 'rule_3', 'rule_4', 'rule_5') + + +# ------------------------------------------------------------ decision rule -- +def gt_index(gt): + """PL index of the called genotype, or None if GT is missing. + + VCF ordering: for alleles sorted ascending a_0..a_{P-1}, + index = sum_k C(a_{k-1} + k - 1, k). Diploid a<=b gives b(b+1)/2 + a and + haploid a gives a. + """ + if not gt or gt == '.' or gt == './.' or gt == '.|.': + return None + parts = gt.split('|') if '|' in gt else gt.split('/') + try: + alleles = sorted(int(p) for p in parts) + except ValueError: # any missing allele -> treat GT as missing + return None + if not alleles: + return None + if len(alleles) == 1: + return alleles[0] + if len(alleles) == 2: + a, b = alleles + return (b * (b + 1)) // 2 + a + idx = 0 # general polyploid fallback + for k, a in enumerate(alleles, start=1): + n, num, den = a + k - 1, 1, 1 + for i in range(k): + num *= (n - i) + den *= (i + 1) + idx += num // den + return idx + + +def conf_of(pl, g, cap): + """min(pl[j], j != g) - pl[g], capped; None when uninformative.""" + if pl is None: + return None + n = len(pl) + if n < 2 or g >= n: + return None + x = pl[g] + if x is None: + return None + best = None + for j in range(n): + if j == g: + continue + y = pl[j] + if y is not None and (best is None or y < best): + best = y + if best is None: + return None + c = best - x + return c if c < cap else cap + + +def sum_ad(v): + s = 0 + for x in v: + if x is not None: + s += x + return s + + +def select(g, sad, lad, spl, lpl, cap, thr, ratio): + """-> (source, branch). See the module docstring.""" + if g is None or sad is None or lad is None: + return SRC_NONE, 0 + ssum = sum_ad(sad) + lsum = sum_ad(lad) + if lsum == 0: + return SRC_SAD, 2 + if ssum == 0: + return SRC_LAD, 2 + cs = conf_of(spl, g, cap) + cl = conf_of(lpl, g, cap) + if cs is None or cl is None: + if cs is None and cl is None: + return (SRC_SAD if ssum >= lsum else SRC_LAD), 3 + return (SRC_LAD if cs is None else SRC_SAD), 3 + if cs - cl > thr: + return SRC_SAD, 4 + if cl - cs > thr: + return SRC_LAD, 4 + if ssum > (1.0 + ratio) * lsum: + return SRC_SAD, 5 + return SRC_LAD, 5 + + +def filt_class(filters): + """v.filter is a list; [] means '.'""" + if not filters or filters == ['PASS']: + return 'PASS' + if filters == ['MLrejected']: + return 'MLrejected' + if filters == ['LowQual']: + return 'LowQual' + return 'other' + + +# ----------------------------------------------------------------- the map -- +def fix_shard(vcfi, vcfo, cap=99.0, thr=10.0, ratio=0.10): + """Rewrite one shard (or, at --threads 1, the whole file). + + vcfi/vcfo are VCF or VCFReader/VCFWriter. Returns (counters, by_filter) + which the parent reduces in shard order. + """ + cnt = dict.fromkeys(COUNTERS, 0) + byfilt = dict((c, dict.fromkeys(BYFILT_KEYS, 0)) for c in FILTER_CLASSES) + # VCFWriter.emit drops v.pos >= end but NOT v.pos < start, so a record + # that straddles the shard start belongs to the previous shard and would + # otherwise be written twice (and make tabix.add assert). + lo = getattr(vcfo, 'start', 0) + + for v in vcfi: + if v.pos < lo: + continue + cnt['records'] += 1 + if v.alt == [''] or v.alt == ['<*>']: + # gVCF hom-ref block: no SAD/LAD to choose between, and adding + # INFO/DP or FORMAT keys would change what GVCFtyper reads back. + cnt['records_ref_block'] += 1 + vcfo.emit(v) # v.line is intact + continue + if not v.samples: + cnt['records_no_samples'] += 1 + vcfo.emit(v) # v.line is intact + continue + + bf = byfilt[filt_class(v.filter)] + bf['records'] += 1 + + info = v.info + has_info_dp = 'DP' in info + old_info_dp = info.get('DP') + total = 0 + any_upd = False + + for s in v.samples: + cnt['samples_seen'] += 1 + sad = s.get('SAD') + lad = s.get('LAD') + src, rule = select(gt_index(s.get('GT')), sad, lad, + s.get('SPL'), s.get('LPL'), cap, thr, ratio) + has_ad = 'AD' in s + has_dp = 'DP' in s + old_ad = s.get('AD') + old_dp = s.get('DP') + if not has_ad: + cnt['fmt_AD_added'] += 1 + if not has_dp: + cnt['fmt_DP_added'] += 1 + + if src == SRC_NONE: + cnt['samples_untouched'] += 1 + cnt['src_NONE'] += 1 + bf['src_NONE'] += 1 + s['AD_ORIG'] = None + s['DP_ORIG'] = None + if old_dp is not None: + total += old_dp + else: + cnt['samples_updated'] += 1 + any_upd = True + new_ad = list(sad if src == SRC_SAD else lad) + new_dp = sum_ad(new_ad) + if src == SRC_SAD: + cnt['src_SAD'] += 1 + bf['src_SAD'] += 1 + else: + cnt['src_LAD'] += 1 + bf['src_LAD'] += 1 + s['AD_ORIG'] = old_ad if has_ad else None + s['DP_ORIG'] = old_dp if has_dp else None + s['AD'] = new_ad + s['DP'] = new_dp + total += new_dp + if new_ad != old_ad: + cnt['AD_changed'] += 1 + if new_dp != old_dp: + cnt['FMT_DP_changed'] += 1 + s['ADSRC'] = SRC_NAME[src] + s['ADRULE'] = rule + cnt['rule_%d' % rule] += 1 # rule_0 .. rule_5 + bf['rule_%d' % rule] += 1 + + if any_upd: + cnt['records_updated'] += 1 + else: + cnt['records_untouched'] += 1 + + if has_info_dp: + info['DP_ORIG'] = old_info_dp + else: + cnt['info_DP_added'] += 1 + if total != old_info_dp: + cnt['INFO_DP_changed'] += 1 + info['DP'] = total + + v.line = None # force VCF.format() to re-emit + vcfo.emit(v) + + return cnt, byfilt + + +# -------------------------------------------------------------- the driver -- +def sharded_run(nthr, contigs, func, step, *args, **kwargs): + """vcflib.Sharder driver, after sentieon_cli.vcf_mod.sharded_run().""" + sharder = vcflib.Sharder(nthr) + intervals = ((c, 0, int(t['length'])) for c, t in contigs.items()) + shards = sharder.cut(intervals, step) + return sharder.run(shards, func, [], *args, **kwargs) + + +def open_vcfs(input_vcf, output_vcf): + """after sentieon_cli.vcf_mod.open_vcfs(); the header is copied from the + input with the five new declarations added through vcflib.""" + try: + vcfi = vcflib.VCF(input_vcf, 'r') + except EnvironmentError as err: + sys.exit('sad_lad_update: cannot read %s or its index: %s' % + (input_vcf, err)) + if not vcfi.contigs: + sys.exit('sad_lad_update: the input VCF header has no ##contig lines; ' + 'vcflib cannot index the output without them') + for k in ('SAD', 'LAD'): + if k not in vcfi.formats: + sys.stderr.write('sad_lad_update: WARNING the input does not declare ' + 'FORMAT/%s; every sample will be left ' + 'untouched\n' % k) + + if ('DP_ORIG' in vcfi.infos or 'AD_ORIG' in vcfi.formats or + 'DP_ORIG' in vcfi.formats): + sys.stderr.write( + 'sad_lad_update: WARNING the input already declares DP_ORIG and/or ' + 'AD_ORIG (has this VCF already been updated?); the *current* ' + 'DP/AD are what will be stored as _ORIG\n') + update = [INFO_DP_ORIG_HDR, FMT_AD_ORIG_HDR, FMT_DP_ORIG_HDR, + FMT_ADSRC_HDR, FMT_ADRULE_HDR] + + try: + vcfo = vcflib.VCF(output_vcf, 'w') + vcfo.copy_header(vcfi, update=update) + vcfo.emit_header() + except EnvironmentError as err: + sys.exit('sad_lad_update: cannot write %s: %s' % (output_vcf, err)) + return vcfi, vcfo + + +def write_stats(path, cnt, byfilt, args, elapsed, rss_mb, nshards): + meta = { + 'input_vcf': args.input_vcf, + 'output_vcf': args.output_vcf, + 'similar_threshold': args.similar_threshold, + 'lad_ratio': args.lad_ratio, + 'conf_cap': args.conf_cap, + 'threads': args.threads, + 'shards': nshards, + 'elapsed_s': round(elapsed, 1), + 'peak_rss_mb': round(rss_mb, 1), + } + if path.endswith('.json'): + with open(path, 'w') as fh: + json.dump({'params': meta, 'counters': cnt, 'by_filter': byfilt}, + fh, indent=1, sort_keys=True) + fh.write('\n') + return + with open(path, 'w') as fh: + fh.write('#class\tkey\tvalue\n') + for k in sorted(meta): + fh.write('param\t%s\t%s\n' % (k, meta[k])) + for k in sorted(cnt): + fh.write('total\t%s\t%d\n' % (k, cnt[k])) + for c in FILTER_CLASSES: + for k in sorted(byfilt[c]): + fh.write('filter:%s\t%s\t%d\n' % (c, k, byfilt[c][k])) + + +def main(args): + if not args.output_vcf.endswith('.vcf.gz'): + sys.exit('sad_lad_update: --output_vcf must end in .vcf.gz') + t0 = time.time() + vcfi, vcfo = open_vcfs(args.input_vcf, args.output_vcf) + + kw = dict(cap=args.conf_cap, thr=args.similar_threshold, + ratio=args.lad_ratio) + if args.threads > 1: + results = sharded_run(args.threads, vcfi.contigs, fix_shard, + args.step_size, vcfi, vcfo, **kw) + else: + results = [fix_shard(vcfi, vcfo, **kw)] + + total = dict.fromkeys(COUNTERS, 0) + byfilt = dict((c, dict.fromkeys(BYFILT_KEYS, 0)) for c in FILTER_CLASSES) + for cnt, bf in results: + for k, n in cnt.items(): + total[k] += n + for c in FILTER_CLASSES: + for k, n in bf[c].items(): + byfilt[c][k] += n + + for f in (vcfi, vcfo): + f.close() + + elapsed = time.time() - t0 + rss = (resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss) / 1024.0 + for k in COUNTERS: + sys.stderr.write('sad_lad_update: %-22s= %d\n' % (k, total[k])) + sys.stderr.write('sad_lad_update: %-22s= %d\n' % ('shards', len(results))) + sys.stderr.write('sad_lad_update: %-22s= %.1f\n' % ('elapsed_s', elapsed)) + sys.stderr.write('sad_lad_update: %-22s= %.1f\n' % ('peak_rss_mb', rss)) + if args.stats: + write_stats(args.stats, total, byfilt, args, elapsed, rss, + len(results)) + return 0 + + +def parse_args(argv=None): + p = argparse.ArgumentParser( + description=__doc__.split('\n\n')[1], + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p.add_argument('--input_vcf', required=True, + help='input VCF or gVCF (bgzf + index); gVCF hom-ref ' + 'blocks are passed through unchanged') + p.add_argument('--output_vcf', required=True, + help='output VCF; must end in .vcf.gz (so .g.vcf.gz is ' + 'accepted), .tbi is written too') + p.add_argument('--stats', default=None, + help='write the counters here (.json -> JSON, else TSV)') + p.add_argument('--similar_threshold', type=float, default=10.0, + help='|conf_S - conf_L| at or below which the two sources ' + 'are called "similar" (branch 4)') + p.add_argument('--lad_ratio', type=float, default=0.10, + help='branch 5 uses LAD unless sum(SAD) exceeds sum(LAD) ' + 'by more than this fraction') + p.add_argument('--conf_cap', type=float, default=99.0, + help='cap applied to conf(SPL) and conf(LPL)') + p.add_argument('--threads', type=int, default=1, + help='worker processes; 1 runs in-process, >1 shards') + p.add_argument('--step_size', type=int, default=25 * 1000 * 1000, + help='shard size in bp (only used when --threads > 1)') + return p.parse_args(argv) + + +if __name__ == '__main__': + sys.exit(main(parse_args())) + +# vim: ts=4 sw=4 expandtab diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 55616b8..ed445e8 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -6,9 +6,8 @@ import copy import json import pathlib -import shutil import sys -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Dict, List, Optional, Set, Union import packaging.version @@ -16,48 +15,47 @@ from . import command_strings as cmds from .archive import ar_load -from .base_pangenome import BasePangenome, SampleSex +from .base_pangenome import BasePangenome from .dag import DAG from .driver import ( - AlignmentStat, - BaseDistributionByCycle, - CNVModelApply, - CNVscope, - CoverageMetrics, - Dedup, - DNAModelApply, + BaseAlgo, DNAscope, - Driver, - GCBias, - GVCFtyper, - InsertSizeMetricAlgo, - LocusCollector, - MeanQualityByCycle, PangenomeSV, - QualDistribution, - ReadWriter, - WgsMetricsAlgo, ) from .job import Job from .logging import get_logger from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + BwaExtractStage, + find_unzip, +) +from .stages.base import StageContext +from .stages.cnv import CNV_MIN_VERSIONS +from .stages.dedup import DedupStage +from .stages.expansion import EXPANSION_MIN_VERSIONS +from .stages.metrics import MetricsPaths, MetricsStage +from .stages.ploidy import PloidyStage +from .stages.segdup import SEGDUP_MIN_VERSIONS +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + GVCFtyperStage, + TransferApplyStage, + TransferSpec, +) +from .stages.transfer import TransferConfig from .util import ( __version__, - check_version, check_kmc_patch, parse_rg_line, path_arg, - tmp, + require_versions, total_memory, - vcf_id, ) from .shard import ( - GRCH38_CONTIGS, determine_shards_from_fai, parse_fai, - vcf_contigs, ) -from .transfer import build_transfer_jobs SENT_PANGENOME_MIN_VERSIONS = { "kmc": None, @@ -67,21 +65,6 @@ "samtools": packaging.version.Version("1.16"), } -SEGDUP_MIN_VERSION = { - "segdup-caller": packaging.version.Version("0.7.0"), -} - -EXPANSION_MIN_VERSION = { - "ExpansionHunter": None, -} - -T1K_MIN_VERSION = { - "run-t1k": None, -} - -DEFAULT_T1K_HLA_LOCUS = "chr6:28510020-33480577" -DEFAULT_T1K_KIR_LOCUS = "chr19:53100000-55800000" - logger = get_logger(__name__) @@ -127,22 +110,6 @@ class SentieonPangenome(BasePangenome): "help": "Generate a gVCF output file.", "action": "store_true", }, - "pangenome_ref_name": { - "default": "GRCh38", - "help": ( - "Reference name in the pangenome (GRCh38). The " - "'extract..model' member of the model " - "bundle is preferred; if it is absent and the reference " - "name is 'GRCh38', the pipeline falls back to " - "'extract.model'." - ), - }, - "pangenome_contig_prefix": { - "default": "GRCh38#0#", - "help": ( - "Prefix to strip from pangenome contig names (GRCh38#0#)" - ), - }, "skip_metrics": { "help": "Skip metrics collection and multiQC", "action": "store_true", @@ -151,76 +118,7 @@ class SentieonPangenome(BasePangenome): "help": "Skip multiQC report generation", "action": "store_true", }, - "segdup_caller": { - "nargs": "*", - "help": ( - "Call variants in difficult segmental duplications with " - "segdup-caller. Supply the flag with no arguments to run " - "the caller's default gene set. Supply a comma-separated " - "list of gene names (e.g. 'CFH,CYP2D6,SMN1') to restrict " - "calling to those genes." - ), - }, - "expansion_catalog": { - "help": ( - "An ExpansionHunter variant catalog. Required for short " - "tandem repeat expansion calling." - ), - "type": path_arg(exists=True, is_file=True), - }, - "t1k_hla_seq": { - "help": ( - "The DNA HLA seq FASTA file for T1K. Required for HLA " - "calling." - ), - "type": path_arg(exists=True, is_file=True), - }, - "t1k_hla_coord": { - "help": ( - "The DNA HLA coord FASTA file for T1K. Required for HLA " - "calling." - ), - "type": path_arg(exists=True, is_file=True), - }, - "t1k_hla_locus": { - "default": DEFAULT_T1K_HLA_LOCUS, - "help": ( - "Reference interval covering the HLA locus. Reads " - "overlapping this region are extracted before being " - f"passed to T1K (default: {DEFAULT_T1K_HLA_LOCUS})." - ), - }, - "t1k_kir_seq": { - "help": ( - "The DNA KIR seq FASTA file for T1K. Required for KIR " - "calling." - ), - "type": path_arg(exists=True, is_file=True), - }, - "t1k_kir_coord": { - "help": ( - "The DNA KIR coord FASTA file for T1K. Required for KIR " - "calling." - ), - "type": path_arg(exists=True, is_file=True), - }, - "t1k_kir_locus": { - "default": DEFAULT_T1K_KIR_LOCUS, - "help": ( - "Reference interval covering the KIR locus. Reads " - "overlapping this region are extracted before being " - f"passed to T1K (default: {DEFAULT_T1K_KIR_LOCUS})." - ), - }, # Hidden arguments - "skip_contig_checks": { - "help": argparse.SUPPRESS, - "action": "store_true", - }, - "skip_pangenome_name_checks": { - "help": argparse.SUPPRESS, - "action": "store_true", - }, "skip_pop_vcf_id_check": { "help": argparse.SUPPRESS, "action": "store_true", @@ -246,64 +144,40 @@ def __init__(self) -> None: self.bed: Optional[pathlib.Path] = None self.call_svs = False self.gvcf = False - self.pangenome_ref_name = "GRCh38" self.extract_model_name = "extract.model" - self.pangenome_contig_prefix = "GRCh38#0#" self.skip_metrics = False self.skip_multiqc = False - self.segdup_caller: Optional[List[str]] = None - self.expansion_catalog: Optional[pathlib.Path] = None - self.t1k_hla_seq: Optional[pathlib.Path] = None - self.t1k_hla_coord: Optional[pathlib.Path] = None - self.t1k_hla_locus: str = DEFAULT_T1K_HLA_LOCUS - self.t1k_kir_seq: Optional[pathlib.Path] = None - self.t1k_kir_coord: Optional[pathlib.Path] = None - self.t1k_kir_locus: str = DEFAULT_T1K_KIR_LOCUS - self.skip_contig_checks: bool = False - self.skip_pangenome_name_checks: bool = False self.skip_pop_vcf_id_check: bool = False self.skip_model_apply = False self.skip_small_variants = False - def main(self, args: argparse.Namespace) -> None: - """Run the pipeline""" - self.handle_arguments(args) - self.setup_logging(args) - self.validate_ref() + def _cnv_in_second_dag(self) -> bool: + """CNV calling runs in the second, sex-aware DAG""" + return self.call_svs and self.has_cnv_model - self.fai_data = parse_fai(pathlib.Path(str(self.reference) + ".fai")) - self.pop_vcf_contigs: Dict[str, Optional[int]] = {} - if self.pop_vcf: - self.pop_vcf_contigs = vcf_contigs(self.pop_vcf, self.dry_run) - self.logger.debug("VCF contigs are: %s", self.pop_vcf_contigs) + def _needs_second_dag(self) -> bool: + """The run has jobs that depend on the estimated sample sex""" + return bool( + self.expansion_catalog + or self.segdup_caller is not None + or self._cnv_in_second_dag() + ) - self.validate() + def validate(self) -> None: + """Validate pipeline inputs""" + self.validate_ref() + self.fai_data = parse_fai(pathlib.Path(str(self.reference) + ".fai")) self.shards = determine_shards_from_fai( self.fai_data, 10 * 1000 * 1000 ) + # Before `validate_bundle`, which picks the bundle's + # `extract..model` member by the resolved name + self.resolve_pangenome_reference() + self.load_pop_vcf_header(self.pop_vcf) - tmp_dir_str = tmp() - self.tmp_dir = pathlib.Path(tmp_dir_str) - - dag = self.build_first_dag() - executor = self.run(dag) - self.check_execution(dag, executor) - - if self.expansion_catalog or self.segdup_caller is not None: - self.get_sex(self.ploidy_json) - dag = self.build_second_dag() - executor = self.run(dag) - self.check_execution(dag, executor) - - if not self.retain_tmpdir: - shutil.rmtree(tmp_dir_str) - - def validate(self) -> None: - """Validate pipeline inputs""" self.validate_bundle() self.validate_fastq_rg() self.validate_output_vcf() - self.validate_ref() self.collect_readgroups() if not self.sample_input and not self.r1_fastq: @@ -326,21 +200,21 @@ def validate(self) -> None: self.validate_segdup() self.validate_expansion() self.validate_t1k() + self.validate_cnv() + + require_versions( + SENT_PANGENOME_MIN_VERSIONS, skip=self.skip_version_check + ) - if not self.skip_version_check: - for cmd, min_version in SENT_PANGENOME_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) - - if self.sample_input: - if not check_kmc_patch("kmc"): - self.logger.error( - "Error: The 'kmc' executable in the PATH does not " - "support reading from stdin. Please ensure " - "you are using the patched version of KMC from " - "https://github.com/Sentieon/KMC/releases." - ) - sys.exit(2) + if not self.skip_version_check and self.sample_input: + if not check_kmc_patch("kmc"): + self.logger.error( + "Error: The 'kmc' executable in the PATH does not " + "support reading from stdin. Please ensure " + "you are using the patched version of KMC from " + "https://github.com/Sentieon/KMC/releases." + ) + sys.exit(2) if self.bed is None: self.logger.info( @@ -348,51 +222,7 @@ def validate(self) -> None: "across decoy and unplaced contigs." ) - if not self.skip_pangenome_name_checks: - if not str(self.gbz).endswith("grch38.gbz"): - self.logger.error( - "The `--gbz` file does not have the expected suffix. " - "Check that you are using a GRCh38 pangenome." - ) - sys.exit(2) - - if not str(self.hapl).endswith("grch38.hapl"): - self.logger.error( - "The `--hapl` file does not have the expected suffix. " - "Check that you re using a GRCh38 pangenome." - ) - sys.exit(2) - - if not self.skip_contig_checks: - # Check the fai file contigs - mismatch_contigs: Set[str] = set() - for ctg, length in GRCH38_CONTIGS.items(): - d = self.fai_data.get(ctg, {}) - fai_length = d.get("length", -1) - if length != fai_length: - mismatch_contigs.add(ctg) - if mismatch_contigs: - mismatch_contigs_s = ", ".join(mismatch_contigs) - self.logger.error( - "Reference contigs with unexpected lengths: %s", - mismatch_contigs_s, - ) - sys.exit(2) - - # Check the pop VCF file contigs - if not self.dry_run: - mismatch_contigs = set() - for ctg, length in GRCH38_CONTIGS.items(): - vcf_length = self.pop_vcf_contigs.get(ctg, -1) - if length != vcf_length: - mismatch_contigs.add(ctg) - if mismatch_contigs: - mismatch_contigs_s = ", ".join(mismatch_contigs) - self.logger.error( - "Pop VCF contigs with unexpected lengths: %s", - mismatch_contigs_s, - ) - sys.exit(2) + self.validate_pangenome_contig_lengths() def validate_segdup(self) -> None: if self.segdup_caller is None: @@ -413,10 +243,7 @@ def validate_segdup(self) -> None: ) sys.exit(2) - if not self.skip_version_check: - for cmd, min_version in SEGDUP_MIN_VERSION.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(SEGDUP_MIN_VERSIONS, skip=self.skip_version_check) def validate_expansion(self) -> None: if self.expansion_catalog is None: @@ -436,47 +263,24 @@ def validate_expansion(self) -> None: ) sys.exit(2) - if not self.skip_version_check: - for cmd, min_version in EXPANSION_MIN_VERSION.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(EXPANSION_MIN_VERSIONS, skip=self.skip_version_check) - def validate_t1k(self) -> None: - hla_requested = self.t1k_hla_seq is not None or ( - self.t1k_hla_coord is not None - ) - if hla_requested and not (self.t1k_hla_seq and self.t1k_hla_coord): - self.logger.error( - "For HLA calling, both `--t1k_hla_seq` and `--t1k_hla_coord` " - "must be supplied." - ) - sys.exit(2) - - kir_requested = self.t1k_kir_seq is not None or ( - self.t1k_kir_coord is not None - ) - if kir_requested and not (self.t1k_kir_seq and self.t1k_kir_coord): - self.logger.error( - "For KIR calling, both `--t1k_kir_seq` and `--t1k_kir_coord` " - "must be supplied." - ) - sys.exit(2) - - if not (hla_requested or kir_requested): + def validate_cnv(self) -> None: + """Validate the arguments used for sex-aware CNV calling""" + # `validate_bundle` has already set `self.has_cnv_model` + cnv_will_run = self._cnv_in_second_dag() + self.resolve_cnv_par_bed(self.fai_data, self.par_bed, cnv_will_run) + if not cnv_will_run: return - if not self.skip_version_check: - for cmd, min_version in T1K_MIN_VERSION.items(): - if not check_version(cmd, min_version): - sys.exit(2) + # A PAR BED file is required for CNV calling, whatever the sex + self.validate_cnv_par(True) + + require_versions(CNV_MIN_VERSIONS, skip=self.skip_version_check) def validate_bundle(self) -> None: - if not self.pop_vcf: - self.logger.error("pop_vcf is required") - sys.exit(2) - if not self.gbz: - self.logger.error("gbz is required") - sys.exit(2) + self.required(self.pop_vcf, "pop_vcf") + gbz = self.required(self.gbz, "gbz") bundle_info_bytes = ar_load( str(self.model_bundle) + "/bundle_info.json" ) @@ -509,7 +313,7 @@ def validate_bundle(self) -> None: if bundle_pipeline != "Sentieon pangenome": self.logger.error("The model bundle is for a different pipeline.") sys.exit(2) - if self.gbz.name != bundle_pangenome: + if gbz.name != bundle_pangenome: self.logger.warning( "The `--gbz` file name is not '%s'. " "This model is optimized for the %s pangenome.", @@ -548,13 +352,8 @@ def validate_bundle(self) -> None: "PangenomeSV will still run." ) - if not self.skip_pop_vcf_id_check and not self.dry_run: - pop_vcf_id = vcf_id(self.pop_vcf) - if bundle_vcf_id != pop_vcf_id: - self.logger.error( - "The ID of the `--pop_vcf` does not match the model bundle" - ) - sys.exit(2) + if not self.skip_pop_vcf_id_check: + self.check_pop_vcf_id(bundle_vcf_id) def validate_fastq_rg(self) -> None: if len(self.r1_fastq) != len(self.r2_fastq): @@ -612,19 +411,10 @@ def configure(self) -> None: pass def build_dag(self) -> DAG: - return DAG() + """Build the first DAG for the Sentieon pangenome pipeline""" + bundle = self.required(self.model_bundle, "model_bundle") - def build_first_dag(self) -> DAG: - """Build the main DAG for the Sentieon pangenome pipeline""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) + ctx = self.stage_context() self.logger.info("Building the Sentieon pangenome DAG") dag = DAG() @@ -632,16 +422,13 @@ def build_first_dag(self) -> DAG: # Output files suffix = "bam" if self.bam_format else "cram" out_bwa_aln = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_bwa_deduped.{suffix}") + str(ctx.output_vcf).replace(".vcf.gz", f"_bwa_deduped.{suffix}") ) out_mm2_aln = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_mm2_deduped.{suffix}") - ) - self.ploidy_json = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_ploidy.json") + str(ctx.output_vcf).replace(".vcf.gz", f"_mm2_deduped.{suffix}") ) out_gvcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", ".g.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", ".g.vcf.gz") ) # Intermediate file paths @@ -656,14 +443,12 @@ def build_first_dag(self) -> DAG: if not self.r1_fastq: # with bam/cram input, output the realigned bam/cram mm2_bam = pathlib.Path( - str(self.output_vcf).replace( + str(ctx.output_vcf).replace( ".vcf.gz", f"_mm2_deduped.{suffix}" ) ) raw_vcf = self.tmp_dir.joinpath("sample-dnascope.vcf.gz") transfer_vcf = self.tmp_dir.joinpath("sample-dnascope_transfer.vcf.gz") - # - Ensure we have a bam file for t1k - self.rw_bam = self.tmp_dir.joinpath("sample_deduped.bam") bwa_lc_dependencies: Set[Job] = set() haplotype_dependencies: Set[Job] = set() @@ -679,8 +464,10 @@ def build_first_dag(self) -> DAG: haplotype_dependencies.add(kmc_job) # BWA alignment and extraction - bwa_job = self.build_alignment_job(bwa_bam, ext_fastq) - dag.add_job(bwa_job) + 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 @@ -696,6 +483,7 @@ def build_first_dag(self) -> DAG: Pipeline(Command("ln", "-sf", "/dev/stdout", str(rw_bam))), "extract-kmc-symlink", 1, + task_name="read-extraction", ) dag.add_job(ln_job) @@ -704,14 +492,15 @@ def build_first_dag(self) -> DAG: kmer_prefix, ext_fastq, self.sample_input, - self.reference, - self.model_bundle.joinpath(self.extract_model_name), + ctx.reference, + bundle.joinpath(self.extract_model_name), self.tmp_dir, rw_bam, threads=self.cores, ), "extract-kmc", self.cores, + task_name="read-extraction", ) dag.add_job(extract_kmc_job, {ln_job}) haplotype_dependencies.add(extract_kmc_job) @@ -720,11 +509,21 @@ def build_first_dag(self) -> DAG: haplotypes_job = self.build_haplotypes_job(sample_pangenome, kmer_file) dag.add_job(haplotypes_job, haplotype_dependencies) + # Confirm the sampled pangenome kept its reference paths before + # anything consumes it + check_gbz_job = self.build_check_gbz_job(sample_pangenome) + dag.add_job(check_gbz_job, {haplotypes_job}) + # convert the sample pangenome gfa_job = self.build_gfa_job(sample_gfa, sample_pangenome) fasta_job = self.build_fasta_job(sample_fasta, sample_pangenome) - dag.add_job(gfa_job, {haplotypes_job}) - dag.add_job(fasta_job, {haplotypes_job}) + dag.add_job(gfa_job, {check_gbz_job}) + dag.add_job(fasta_job, {check_gbz_job}) + + # Confirm the GFA carries the reference's rGFA tags, without which + # `pgutil lift` leaves every read unmapped + check_gfa_job = self.build_check_gfa_job(sample_gfa) + dag.add_job(check_gfa_job, {gfa_job}) # minimap2 alignment of the extracted fastq dnascope_dependencies = set() @@ -734,11 +533,11 @@ def build_first_dag(self) -> DAG: sample_fasta, sample_gfa, ) - dag.add_job(mm2_job, mm2_dependencies | {gfa_job, fasta_job}) + dag.add_job(mm2_job, mm2_dependencies | {check_gfa_job, fasta_job}) dnascope_dependencies.add(mm2_job) # With fastq input, perform dedup and metrics - cnv_input_bams: List[pathlib.Path] = [] + sr_alignments: List[pathlib.Path] = [] cnvscope_deps: Set[Job] = set() if self.r1_fastq: dnascope_bams.append(out_bwa_aln) @@ -746,250 +545,224 @@ def build_first_dag(self) -> DAG: # Emit Dedup metrics for the primary (bwa) short-read alignment so # they land in the metrics directory scanned by MultiQC. + paths = MetricsPaths.from_output_vcf(ctx.output_vcf) dedup_metrics: Optional[pathlib.Path] = None if not self.skip_metrics: - metrics_dir = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_metrics") - ) - if not self.dry_run: - metrics_dir.mkdir(exist_ok=True) - sample_name = self.output_vcf.name.replace(".vcf.gz", "") - dedup_metrics = metrics_dir.joinpath( - sample_name + ".txt.dedup_metrics.txt" - ) - - bwa_lc_job, bwa_dedup_job = self.build_dedup_job( - out_bwa_aln, [bwa_bam], "bwa", metrics=dedup_metrics - ) - mm2_lc_job, mm2_dedup_job = self.build_dedup_job( - out_mm2_aln, [mm2_bam], "mm2", left_align=True - ) - dag.add_job(bwa_lc_job, bwa_lc_dependencies) - dag.add_job(bwa_dedup_job, {bwa_lc_job}) + paths.ensure_dir(self.dry_run) + dedup_metrics = paths.dedup_metrics + + bwa_dedup = DedupStage( + ctx=ctx, + tag="bwa", + inputs=[bwa_bam], + output=out_bwa_aln, + score_file=self.tmp_dir.joinpath("sample-bwa-score.txt.gz"), + dedup_metrics=dedup_metrics, + ).add_to(dag, bwa_lc_dependencies) + bwa_dedup_job = bwa_dedup.dedup_job dnascope_dependencies.add(bwa_dedup_job) - dag.add_job(mm2_lc_job, {mm2_job}) - dag.add_job(mm2_dedup_job, {mm2_lc_job}) - dnascope_dependencies.add(mm2_dedup_job) + mm2_dedup = DedupStage( + ctx=ctx, + tag="mm2", + inputs=[mm2_bam], + 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}) + dnascope_dependencies.add(mm2_dedup.dedup_job) - cnv_input_bams = [out_bwa_aln] + sr_alignments = [out_bwa_aln] cnvscope_deps = {bwa_dedup_job} if not self.skip_metrics: - metrics_job, rehead_job = self.build_metrics_job( - [out_bwa_aln, out_mm2_aln], - ) - dag.add_job(metrics_job, {bwa_dedup_job, mm2_dedup_job}) - dag.add_job(rehead_job, {metrics_job}) + metrics_result = MetricsStage( + ctx=ctx, + inputs=[out_bwa_aln], + algos=self.pangenome_metrics_algos(paths), + rehead_metrics=paths.wgs, + ).add_to(dag, {bwa_dedup_job}) if not self.skip_multiqc: multiqc_job = self.multiqc() if multiqc_job: - dag.add_job(multiqc_job, {rehead_job}) + dag.add_job(multiqc_job, metrics_result.terminal) else: dnascope_bams.append(mm2_bam) - cnv_input_bams = list(self.sample_input) + sr_alignments = list(self.sample_input) - # Stash the short-read alignment for the second DAG - self.sr_alignment = cnv_input_bams[0] + # Stash the short-read alignments for the second DAG + self.sr_alignments = sr_alignments - # Estimate ploidy when needed for sex-aware downstream tools - if self.expansion_catalog or self.segdup_caller is not None: - ploidy_job = self.build_ploidy_job( - self.ploidy_json, [cnv_input_bams[0]] - ) - dag.add_job(ploidy_job, cnvscope_deps) + # Estimate the sample ploidy and sex. The JSON output is always + # written; `--sample_sex` takes precedence for the sex used by + # the sex-aware callers. + ploidy_result = PloidyStage( + ctx=ctx, + inputs=[self.sr_alignments[0]], + reference_build=self.reference_build, + ).add_to(dag, cnvscope_deps) + self.ploidy_json = ploidy_result.ploidy_json # T1K HLA/KIR calling - if self.t1k_hla_seq and self.t1k_hla_coord: - out_hla = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_hla") - ) - extract_job, t1k_job = self.build_t1k_jobs( - out_hla, - cnv_input_bams, - self.t1k_hla_seq, - self.t1k_hla_coord, - self.t1k_hla_locus, - "hla-wgs", - tag="hla", - ) - dag.add_job(extract_job, cnvscope_deps) - dag.add_job(t1k_job, {extract_job}) - if self.t1k_kir_seq and self.t1k_kir_coord: - out_kir = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_kir") - ) - extract_job, t1k_job = self.build_t1k_jobs( - out_kir, - cnv_input_bams, - self.t1k_kir_seq, - self.t1k_kir_coord, - self.t1k_kir_locus, - "kir-wgs", - tag="kir", - ) - dag.add_job(extract_job, cnvscope_deps) - dag.add_job(t1k_job, {extract_job}) + self.add_t1k(dag, ctx, self.sr_alignments, cnvscope_deps) # DNAscope calling with bwa and mm2 input sv_vcf = None if self.call_svs: sv_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") ) if self.skip_small_variants and not self.call_svs: return dag - if self.skip_small_variants and self.call_svs and sv_vcf: - # SV calling only - dnascope_job = self.build_dnascope_job( - raw_vcf, - dnascope_bams, - sv_vcf=sv_vcf, - gfa_file=sample_gfa, + read_filters: List[str] = [] + if self.tech.upper() == "ULTIMA": + read_filters.append("UltimaReadFilter") + pcr_indel_model = "NONE" if self.pcr_free else "CONSERVATIVE" + model = bundle.joinpath("dnascope.model") + gfa_file = sample_gfa if self.call_svs else None + + algos: List[BaseAlgo] = [] + if not self.skip_small_variants: + algos.append( + DNAscope( + raw_vcf, + model=model, + pcr_indel_model=pcr_indel_model, + dbsnp=self.dbsnp, + emit_mode="gvcf" if self.gvcf else "variant", + ) ) - dag.add_job(dnascope_job, dnascope_dependencies) - if self.has_cnv_model: - self._add_cnv_jobs( - dag, + if sv_vcf and gfa_file: + algos.append( + PangenomeSV( sv_vcf, - cnv_input_bams, - cnvscope_deps, - dnascope_job, + gfa_file=gfa_file, + prefix=self.contig_prefix(), ) + ) + call = DNAscopeStage( + ctx=ctx, + algos=algos, + inputs=dnascope_bams, + interval=self.bed, + read_filter=read_filters, + ).add_to(dag, dnascope_dependencies) + + if self.skip_small_variants: + # SV calling only. CNV calling is sex-aware and runs in the + # second DAG return dag - apply_dependencies = set() - dnascope_job = self.build_dnascope_job( - raw_vcf, - dnascope_bams, - sv_vcf=sv_vcf, - gfa_file=sample_gfa if self.call_svs else None, - ) - dag.add_job(dnascope_job, dnascope_dependencies) - apply_dependencies.add(dnascope_job) + if self.skip_model_apply and not self.pop_vcf: + # Nothing post-processes the raw VCF + return dag # When --gvcf is set, the model-apply / transfer outputs are - # gVCFs; GVCFtyper produces the final VCF at self.output_vcf. - small_variants_out = out_gvcf if self.gvcf else self.output_vcf + # gVCFs; GVCFtyper produces the final VCF at ctx.output_vcf. + small_variants_out = out_gvcf if self.gvcf else ctx.output_vcf + snv_apply_vcf = self.tmp_dir.joinpath( + "sample-snv_apply.g.vcf.gz" + if self.gvcf + else "sample-snv_apply.vcf.gz" + ) - # transfer annotations from the pop_vcf + # Transfer annotations from the pop_vcf, then apply the model + transfer: Optional[TransferSpec] = None if self.pop_vcf: - transfer_jobs, concat_job = build_transfer_jobs( - ( + transfer = TransferSpec( + config=TransferConfig.from_pipeline(self), + out_vcf=( transfer_vcf if not self.skip_model_apply - else (small_variants_out) + else snv_apply_vcf ), - self.pop_vcf, - raw_vcf, - self.tmp_dir, - self.shards, - self.pop_vcf_contigs, - self.fai_data, - self.dry_run, - self.cores, ) - for job in transfer_jobs: - dag.add_job(job, {dnascope_job}) - dag.add_job(concat_job, set(transfer_jobs)) - apply_dependencies.add(concat_job) - - small_variants_last_job: Optional[Job] = None + apply_spec: Optional[ApplySpec] = None if not self.skip_model_apply: - # DNAModelApply - apply_job = self.build_dnamodelapply_job( - transfer_vcf, small_variants_out - ) - dag.add_job(apply_job, apply_dependencies) - small_variants_last_job = apply_job - elif self.pop_vcf: - small_variants_last_job = concat_job + apply_spec = ApplySpec(model=model, output=snv_apply_vcf) + transfer_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=raw_vcf, + transfer=transfer, + apply=apply_spec, + ).add_to(dag, call.terminal) + + # Update the overestimated AD/DP of the joint pileup + ad_update_job = self.build_count_ad_update_job( + small_variants_out, snv_apply_vcf + ) + dag.add_job(ad_update_job, transfer_apply.terminal) # Genotype the gVCF to also produce a regular VCF at output_vcf - if self.gvcf and small_variants_last_job is not None: - gvcftyper_job = self.build_gvcftyper_job(self.output_vcf, out_gvcf) - dag.add_job(gvcftyper_job, {small_variants_last_job}) + if self.gvcf: + GVCFtyperStage( + ctx=ctx, + gvcf=out_gvcf, + output=ctx.output_vcf, + interval=self.bed, + ).add_to(dag, {ad_update_job}) - if self.call_svs and sv_vcf and self.has_cnv_model: - self._add_cnv_jobs( - dag, - sv_vcf, - cnv_input_bams, - cnvscope_deps, - dnascope_job, - ) + # CNV calling is sex-aware and runs in the second DAG return dag - def build_alignment_job( + def bwa_extract_stage( self, + ctx: StageContext, sample_bam: pathlib.Path, sample_fastq: pathlib.Path, - ) -> Job: - """Build the alignment and extract jobs""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + ) -> BwaExtractStage: + """The bwa alignment and read extraction stage. - unzip = "igzip" - if not shutil.which(unzip): - self.logger.info( - "igzip is recommended for decompression, but is not " - "available. Falling back to gzip." - ) - unzip = "gzip" + The bwa alignment reuses the input readgroup ID with a `-bwa` + suffix, so it stays distinct from the lifted alignment's. + """ + bundle = self.required(self.model_bundle, "model_bundle") rg = copy.deepcopy(self.fastq_readgroup) rg["ID"] = rg["ID"] + "-bwa" - bwa_job = Job( - cmds.cmd_bwa_extract( - sample_bam, - sample_fastq, - self.reference, - self.r1_fastq, - self.r2_fastq, - "@RG\\t" + "\\t".join([f"{x[0]}:{x[1]}" for x in rg.items()]), - self.model_bundle.joinpath(self.extract_model_name), - self.model_bundle.joinpath("bwa.model"), - self.cores, - unzip=unzip, + return BwaExtractStage( + ctx=ctx, + output_bam=sample_bam, + output_fastq=sample_fastq, + r1_fastq=self.r1_fastq, + r2_fastq=self.r2_fastq, + readgroup=( + "@RG\\t" + "\\t".join([f"{x[0]}:{x[1]}" for x in rg.items()]) ), - "bwa-extract", - self.cores, + extract_model=bundle.joinpath(self.extract_model_name), + bwa_model=bundle.joinpath("bwa.model"), + unzip=find_unzip(self.logger), ) - return bwa_job def build_haplotypes_job( self, output_gbz: pathlib.Path, kmer_file: pathlib.Path ) -> Job: """Build vg haplotypes job""" - if not self.hapl: - self.logger.error("hapl is required") - sys.exit(2) - if not self.gbz: - self.logger.error("gbz is required") - sys.exit(2) + hapl_file = self.required(self.hapl, "hapl") + gbz_file = self.required(self.gbz, "gbz") haplotypes_job = Job( cmds.cmd_vg_haplotypes( output_gbz, kmer_file, - self.hapl, - self.gbz, + hapl_file, + gbz_file, threads=self.cores, xargs=[ "--include-reference", "--diploid-sampling", "--set-reference", - self.pangenome_ref_name, + self.ref_name(), ], ), "vg-haplotypes", self.cores, + task_name="pangenome", ) return haplotypes_job @@ -1002,10 +775,11 @@ def build_gfa_job( output_gfa, input_gbz, threads=self.cores, - reference_name=self.pangenome_ref_name, + reference_name=self.ref_name(), ), "vg-convert-gfa", 0, + task_name="pangenome", ) return gfa_job @@ -1020,6 +794,7 @@ def build_fasta_job( ), "vg-paths-fasta", 0, + task_name="pangenome", ) return fasta_job @@ -1031,12 +806,8 @@ def build_minimap2_lift_job( sample_gfa: pathlib.Path, ) -> Job: """Build minimap2 alignment with pgutil lift job""" - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") + reference = self.required(self.reference, "reference") rg = ( self.fastq_readgroup @@ -1047,530 +818,89 @@ def build_minimap2_lift_job( rg2["ID"] = rg2["ID"] + "-mm2" rg2["LR"] = "1" - mm2_model: Union[str, pathlib.Path] = self.model_bundle.joinpath( - "minimap2.model" - ) + mm2_model: Union[str, pathlib.Path] = bundle.joinpath("minimap2.model") mm2_job = Job( cmds.cmd_minimap2_lift( mm2_bam, sample_fasta, ext_fastq, sample_gfa, - self.reference, + reference, "@RG\\t" + "\\t".join([f"{x[0]}:{x[1]}" for x in rg2.items()]), mm2_model, threads=self.cores, - lift_prefix=self.pangenome_contig_prefix, + lift_prefix=self.contig_prefix(), ), "mm2-lift", self.cores, + task_name="pangenome-alignment", ) return mm2_job - def build_dedup_job( - self, - output_bam, - input_bam: List[pathlib.Path], - tag: str, - left_align=False, - metrics: Optional[pathlib.Path] = None, - ) -> Tuple[Job, Job]: - """Build deduplication job""" - score_file = self.tmp_dir.joinpath(f"sample-{tag}-score.txt.gz") - - read_filters = [] - if left_align: - read_filters.append( - "IndelLeftAlignReadTransform," - f"rgid={self.fastq_readgroup['ID']}-mm2" - ) - - # LocusCollector + Dedup - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=input_bam, - read_filter=read_filters, - ) - driver.add_algo(LocusCollector(score_file)) - - lc_job = Job( - Pipeline(Command(*driver.build_cmd())), - f"locuscollector-{tag}", - self.cores, - ) - - driver2 = Driver( - reference=self.reference, - thread_count=self.cores, - input=input_bam, - read_filter=read_filters, - ) - driver2.add_algo(Dedup(output_bam, score_file, metrics=metrics)) - - dedup_job = Job( - Pipeline(Command(*driver2.build_cmd())), - f"dedup-{tag}", - self.cores, - ) - - return lc_job, dedup_job - - def build_metrics_job( - self, - sample_input: List[pathlib.Path], - ) -> Tuple[Job, Job]: - """Build a metrics job""" - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - - # Create the metrics directory - sample_name = self.output_vcf.name.replace(".vcf.gz", "") - metric_base = sample_name + ".txt" - metrics_dir = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_metrics") - ) - if not self.dry_run: - metrics_dir.mkdir(exist_ok=True) - - is_metrics = metrics_dir.joinpath(metric_base + ".insert_size.txt") - mqbc_metrics = metrics_dir.joinpath( - metric_base + ".mean_qual_by_cycle.txt" - ) - bdbc_metrics = metrics_dir.joinpath( - metric_base + ".base_distribution_by_cycle.txt" - ) - qualdist_metrics = metrics_dir.joinpath( - metric_base + ".qual_distribution.txt" - ) - as_metrics = metrics_dir.joinpath(metric_base + ".alignment_stat.txt") - coverage_metrics = metrics_dir.joinpath("coverage") - - # WGS metrics - wgs_metrics = metrics_dir.joinpath(metric_base + ".wgs.txt") - gc_metrics = metrics_dir.joinpath(metric_base + ".gc_bias.txt") - gc_summary = metrics_dir.joinpath(metric_base + ".gc_bias_summary.txt") - - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sample_input, - ) - - driver.add_algo(InsertSizeMetricAlgo(is_metrics)) - driver.add_algo(MeanQualityByCycle(mqbc_metrics)) - driver.add_algo(BaseDistributionByCycle(bdbc_metrics)) - driver.add_algo(QualDistribution(qualdist_metrics)) - driver.add_algo(AlignmentStat(as_metrics)) - driver.add_algo(GCBias(gc_metrics, summary=gc_summary)) - driver.add_algo(WgsMetricsAlgo(wgs_metrics, include_unpaired="true")) - driver.add_algo(CoverageMetrics(coverage_metrics)) - - metrics_job = Job(Pipeline(Command(*driver.build_cmd())), "metrics", 0) - - rehead_script = pathlib.Path( - str( - files("sentieon_cli.scripts").joinpath("rehead_wgs_metrics.py") - ) - ) - rehead_job = Job( - Pipeline( - Command( - sys.executable, - str(rehead_script), - "--metrics_file", - str(wgs_metrics), - ) - ), - "Rehead metrics", - 0, - ) - return (metrics_job, rehead_job) - - def build_dnascope_job( + def build_count_ad_update_job( self, out_vcf: pathlib.Path, - input_bams: List[pathlib.Path], - sv_vcf: Optional[pathlib.Path] = None, - gfa_file: Optional[pathlib.Path] = None, - ) -> Job: - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - - read_filters = [] - if self.tech.upper() == "ULTIMA": - read_filters.append("UltimaReadFilter") - - pcr_indel_model = "NONE" if self.pcr_free else "CONSERVATIVE" - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=input_bams, - interval=self.bed, - read_filter=read_filters, - ) - if not self.skip_small_variants: - driver.add_algo( - DNAscope( - out_vcf, - model=self.model_bundle.joinpath("dnascope.model"), - pcr_indel_model=pcr_indel_model, - dbsnp=self.dbsnp, - emit_mode="gvcf" if self.gvcf else "variant", - ) - ) - if sv_vcf and gfa_file: - driver.add_algo( - PangenomeSV( - sv_vcf, - gfa_file=gfa_file, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "dnascope-raw", - self.cores, - ) - - def build_dnamodelapply_job( - self, in_vcf: pathlib.Path, - out_vcf: pathlib.Path, ) -> Job: - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + """Update the AD/DP of the joint pileup. - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - driver.add_algo( - DNAModelApply( - model=self.model_bundle.joinpath("dnascope.model"), - vcf=in_vcf, - output=out_vcf, - ) + Reads that are present in both the bwa and the mm2 alignment are + counted twice, so FORMAT/AD, FORMAT/DP and INFO/DP are inflated. + The script picks FORMAT/SAD or FORMAT/LAD per sample and derives + the updated depths from that choice. + """ + sad_lad_update = pathlib.Path( + str(files("sentieon_cli.scripts").joinpath("sad_lad_update.py")) ) return Job( - Pipeline(Command(*driver.build_cmd())), - "model-apply", - self.cores, - ) - - def build_gvcftyper_job( - self, - out_vcf: pathlib.Path, - in_gvcf: pathlib.Path, - ) -> Job: - """Genotype a gVCF to produce a VCF""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - - driver = Driver( - reference=self.reference, - thread_count=self.cores, - interval=self.bed, - ) - driver.add_algo( - GVCFtyper( - output=out_vcf, - vcf=in_gvcf, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "gvcftyper", - self.cores, - ) - - def build_segdup_job( - self, - out_segdup: pathlib.Path, - sr_alignment: pathlib.Path, - input_vcf: pathlib.Path, - genes: Optional[str], - ) -> Job: - """Call variants in difficult SegDups""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - - sex = "male" if self.sample_sex == SampleSex.MALE else "female" - - # segdup-caller's default `main.min_map_qual` of 45 is too strict - # for Ultima alignments. - overrides = [] - if self.tech.upper() == "ULTIMA": - overrides.append("main.min_map_qual=30") - - return Job( - cmds.cmd_segdup_caller( - out_segdup, - sr_alignment, - reference=self.reference, - sr_bundle=self.model_bundle, - input_vcf=input_vcf, - sex=sex, - genes=genes, - overrides=overrides, + cmds.cmd_pyexec_sad_lad_update( + out_vcf, + in_vcf, + sad_lad_update, + self.cores, ), - "segdup-caller", + "sad-lad-update", self.cores, + task_name="ad-update", ) - def build_second_dag(self) -> DAG: + def build_second_dag(self) -> Optional[DAG]: """Build the second DAG for sex-aware downstream tools""" + if not self._needs_second_dag(): + return None + + assert self.ploidy_json is not None + self.get_sex(self.ploidy_json) + self.logger.info("Building the second pangenome DAG") dag = DAG() + ctx = self.stage_context() - if self.expansion_catalog and self.sr_alignment: - out_expansions = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_expansion") - ) - expansion_job = self.build_expansion_job( - out_expansions, self.sr_alignment, self.expansion_catalog - ) - dag.add_job(expansion_job) - - # SegDup calling consumes the small-variant VCF and the inferred sex - if ( - self.segdup_caller is not None - and self.sr_alignment - and self.output_vcf - ): - out_segdup = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_segdups") - ) - genes = ",".join(self.segdup_caller) or None - segdup_job = self.build_segdup_job( - out_segdup, - self.sr_alignment, - self.output_vcf, - genes, + # CNV calling with CNVscope, using the sample sex + if self._cnv_in_second_dag(): + self.add_pangenome_cnv( + dag, + ctx, + self.output_path("_sv.vcf.gz"), + self.sr_alignments, + interval=self.bed, ) - dag.add_job(segdup_job) - - return dag - - def build_t1k_jobs( - self, - out_basename: pathlib.Path, - sr_alignments: List[pathlib.Path], - gene_seq: pathlib.Path, - gene_coord: pathlib.Path, - locus: str, - preset: str, - tag: str, - ) -> Tuple[Job, Job]: - """Extract reads at the T1K locus and genotype them with T1K""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - - # Extract reads overlapping the locus to a BAM file with ReadWriter - extracted_bam = self.tmp_dir.joinpath(f"sample_{tag}.bam") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=sr_alignments, - interval=locus, - ) - driver.add_algo(ReadWriter(extracted_bam)) - extract_job = Job( - Pipeline(Command(*driver.build_cmd())), - f"t1k-{tag}-extract", - self.cores, - ) - - t1k_job = Job( - cmds.cmd_t1k( - out_basename, - extracted_bam, - gene_seq=gene_seq, - gene_coord=gene_coord, - preset=preset, - threads=self.cores, - ), - f"t1k-{tag}", - self.cores, - ) - return (extract_job, t1k_job) - def build_expansion_job( - self, - out_expansions: pathlib.Path, - sr_alignment: pathlib.Path, - expansion_catalog: pathlib.Path, - ) -> Job: - """Identify repeat expansions""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) + if self.sr_alignments: + self.add_expansion(dag, ctx, self.sr_alignments[0]) - return Job( - cmds.cmd_expansion_hunter( - out_expansions, - sr_alignment, - reference=self.reference, - variant_catalog=expansion_catalog, - sex="male" if self.sample_sex == SampleSex.MALE else "female", - threads=self.cores, - ), - "expansion-hunter", - self.cores, - ) - - def _add_cnv_jobs( - self, - dag: DAG, - sv_vcf: pathlib.Path, - cnv_input_bams: List[pathlib.Path], - cnvscope_deps: Set[Job], - dnascope_job: Job, - ) -> None: - """Add CNV calling jobs to the DAG""" - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) - - cnvscope_vcf = self.tmp_dir.joinpath("sample-cnvscope.vcf.gz") - cnv_apply_vcf = self.tmp_dir.joinpath("sample-cnv_model_apply.vcf.gz") - indel2cnv_vcf = self.tmp_dir.joinpath("sample-sv_cnv.vcf.gz") - cnv_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") - ) - - # CNVscope on BWA deduped BAM - cnvscope_job = self._build_cnvscope_job(cnvscope_vcf, cnv_input_bams) - dag.add_job(cnvscope_job, cnvscope_deps) - - # CNVModelApply on CNVscope output - cnv_model_apply_job = self._build_cnv_model_apply_job( - cnv_apply_vcf, cnvscope_vcf - ) - dag.add_job(cnv_model_apply_job, {cnvscope_job}) - - # Convert PangenomeSV output to CNVs - indel2cnv_job = self._build_indel2cnv_job(indel2cnv_vcf, sv_vcf) - dag.add_job(indel2cnv_job, {dnascope_job}) - - # Combine CNVModelApply output with converted SVs - combine_job = self._build_combine_cnv_job( - cnv_vcf, cnv_apply_vcf, indel2cnv_vcf - ) - dag.add_job(combine_job, {cnv_model_apply_job, indel2cnv_job}) - - def _build_cnvscope_job( - self, - output_vcf: pathlib.Path, - input_bams: List[pathlib.Path], - ) -> Job: - """Build a CNVscope job""" - assert self.model_bundle is not None - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=input_bams, - interval=self.bed, - ) - driver.add_algo( - CNVscope( - output=output_vcf, - model=self.model_bundle.joinpath("cnv.model"), + # SegDup calling consumes the small-variant VCF and the + # inferred sex. segdup-caller's default `main.min_map_qual` + # of 45 is too strict for Ultima alignments. + self.add_segdup( + dag, + ctx, + self.sr_alignments[0], + overrides=( + ["main.min_map_qual=30"] + if self.tech.upper() == "ULTIMA" + else () + ), ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "cnvscope", - self.cores, - ) - def _build_cnv_model_apply_job( - self, - output_vcf: pathlib.Path, - input_vcf: pathlib.Path, - ) -> Job: - """Build a CNVModelApply job""" - assert self.model_bundle is not None - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - driver.add_algo( - CNVModelApply( - output=output_vcf, - model=self.model_bundle.joinpath("cnv.model"), - vcf=input_vcf, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "cnv-model-apply", - self.cores, - ) - - def _build_indel2cnv_job( - self, - output_vcf: pathlib.Path, - input_vcf: pathlib.Path, - ) -> Job: - """Convert PangenomeSV INDELs to CNV calls""" - indel2cnv_script = pathlib.Path( - str(files("sentieon_cli.scripts").joinpath("indel2cnv.py")) - ) - # Run in background - return Job( - Pipeline( - Command( - sys.executable, - str(indel2cnv_script), - str(self.reference), - str(input_vcf), - str(output_vcf), - "-t", - str(self.cores), - ) - ), - "indel2cnv", - 0, - ) - - def _build_combine_cnv_job( - self, - output_vcf: pathlib.Path, - cnv_vcf: pathlib.Path, - converted_vcf: pathlib.Path, - ) -> Job: - """Combine CNVscope and converted SV calls""" - combine_script = pathlib.Path( - str(files("sentieon_cli.scripts").joinpath("combine_cnv.py")) - ) - return Job( - Pipeline( - Command( - sys.executable, - str(combine_script), - "--cnv", - str(cnv_vcf), - "--converted", - str(converted_vcf), - "-o", - str(output_vcf), - ) - ), - "combine-cnv", - 0, - ) + return dag diff --git a/sentieon_cli/shard.py b/sentieon_cli/shard.py index 815f106..8d04858 100644 --- a/sentieon_cli/shard.py +++ b/sentieon_cli/shard.py @@ -4,42 +4,57 @@ import pathlib import re -import subprocess as sp -from typing import Dict, List, NamedTuple, Optional +from typing import Dict, Iterable, List, NamedTuple, Optional + +from importlib.resources import files from .logging import get_logger +from .util import VcfHeaderError, read_vcf_header logger = get_logger(__name__) -GRCH38_CONTIGS: Dict[str, int] = { - "chr1": 248956422, - "chr2": 242193529, - "chr3": 198295559, - "chr4": 190214555, - "chr5": 181538259, - "chr6": 170805979, - "chr7": 159345973, - "chr8": 145138636, - "chr9": 138394717, - "chr10": 133797422, - "chr11": 135086622, - "chr12": 133275309, - "chr13": 114364328, - "chr14": 107043718, - "chr15": 101991189, - "chr16": 90338345, - "chr17": 83257441, - "chr18": 80373285, - "chr19": 58617616, - "chr20": 64444167, - "chr21": 46709983, - "chr22": 50818468, - "chrX": 156040895, - "chrY": 57227415, - "chrM": 16569, +# Signature contigs used to recognize the common human reference builds. +# A build is only recognized when every signature contig is present with +# the exact expected length. +BUILD_SIGNATURES: Dict[str, Dict[str, int]] = { + "hg38": { + "chr1": 248956422, + "chr2": 242193529, + "chrX": 156040895, + }, + "hg19": { + "chr1": 249250621, + "chr2": 243199373, + "chrX": 155270560, + }, + "b37": { + "1": 249250621, + "2": 243199373, + "X": 155270560, + }, + "chm13": { + "chr1": 248387328, + "chr2": 242696752, + "chrX": 154259566, + }, } +DEFAULT_PLOIDY_CONTIGS = [f"chr{i}" for i in range(1, 23)] + ["chrX", "chrY"] +B37_PLOIDY_CONTIGS = [str(i) for i in range(1, 23)] + ["X", "Y"] + + +class PloidyContigs(NamedTuple): + """Contig names used for ploidy estimation. + + Fields are ``None`` when the defaults of `estimate_ploidy.py` apply. + """ + + contigs: Optional[List[str]] = None + autosomes: Optional[List[str]] = None + x_contig: Optional[str] = None + y_contig: Optional[str] = None + class Shard(NamedTuple): contig: str @@ -49,9 +64,6 @@ class Shard(NamedTuple): def __str__(self) -> str: return f"{self.contig}:{self.start}-{self.stop}" - def bcftools_str(self) -> str: - return f"{{{self.contig}}}:{self.start}-{self.stop}" - def parse_fai(ref_fai: pathlib.Path) -> Dict[str, Dict[str, int]]: """Parse a faidx index""" @@ -75,6 +87,66 @@ def parse_fai(ref_fai: pathlib.Path) -> Dict[str, Dict[str, int]]: return contigs +def detect_reference_build( + fai_data: Dict[str, Dict[str, int]], +) -> Optional[str]: + """Identify the reference build from the fasta index. + + A build is reported only when all of its signature contigs are present + with the expected lengths. Unrecognized references return `None`; the + `--par_bed` argument covers those references. + """ + for build, signature in BUILD_SIGNATURES.items(): + if all( + fai_data.get(ctg, {}).get("length") == length + for ctg, length in signature.items() + ): + logger.debug("Identified the reference build as '%s'", build) + return build + logger.debug("Could not identify the reference build") + return None + + +def par_bed_for_build(build: Optional[str]) -> Optional[pathlib.Path]: + """The packaged pseudo-autosomal region (PAR) BED for a build. + + Returns `None` when the build is unknown or the packaged BED file is + absent. A zero-length data file is treated the same as a missing one, + so the wiring stays safe until the BED file contents are shipped. + """ + if build is None: + return None + + try: + par_bed = pathlib.Path( + str(files("sentieon_cli.data").joinpath(f"par_{build}.bed")) + ) + except ModuleNotFoundError: + par_bed = None # type: ignore[assignment] + + if par_bed is None or not par_bed.is_file() or par_bed.stat().st_size == 0: + logger.warning( + "No pseudo-autosomal region (PAR) BED file is packaged for the " + "'%s' reference build.", + build, + ) + return None + return par_bed + + +def ploidy_contigs_for_build(build: Optional[str]) -> PloidyContigs: + """Contig names used for ploidy estimation with a reference build""" + if build == "b37": + return PloidyContigs( + contigs=list(B37_PLOIDY_CONTIGS), + autosomes=list(B37_PLOIDY_CONTIGS[:-2]), + x_contig="X", + y_contig="Y", + ) + # The `estimate_ploidy.py` defaults are the chr-prefixed contig names + return PloidyContigs() + + def determine_shards_from_fai( fai_data: Dict[str, Dict[str, int]], step: int ) -> List[Shard]: @@ -91,29 +163,16 @@ def determine_shards_from_fai( return shards -def vcf_contigs( - in_vcf: pathlib.Path, dry_run=False +def parse_vcf_contigs( + header: Iterable[str], ) -> Dict[str, Optional[int]]: - """Report the contigs in the input VCF""" - if dry_run: - return { - "chr1": 100, - "chr2": 200, - "chr3": 300, - } + """The ID and length of every `##contig` line of a VCF header. + + A `##contig` line without a `length` maps to `None`. + """ kvpat = re.compile(r'(.*?)=(".*?"|.*?)(?:,|$)') - cmd = ["bcftools", "view", "-h", str(in_vcf)] - p = sp.run(cmd, capture_output=True, text=True) - if p.returncode != 0: - logger.error( - "`%s` failed with return code %d: %s", - " ".join(cmd), - p.returncode, - p.stderr.strip(), - ) - return {} contigs: Dict[str, Optional[int]] = {} - for line in p.stdout.split("\n"): + for line in header: if not line.startswith("##contig"): continue s = line.index("<") @@ -123,3 +182,17 @@ def vcf_contigs( length: Optional[str] = d.get("length", None) contigs[ctg] = int(length) if length else None return contigs + + +def vcf_contigs(in_vcf: pathlib.Path) -> Dict[str, Optional[int]]: + """Report the contigs in the input VCF. + + Returns an empty dictionary, and logs why, when the header could not + be read. + """ + try: + header = read_vcf_header(in_vcf) + except (VcfHeaderError, OSError) as err: + logger.error("%s", err) + return {} + return parse_vcf_contigs(header) diff --git a/sentieon_cli/shell_pipeline.py b/sentieon_cli/shell_pipeline.py index 02eb082..d2d1277 100644 --- a/sentieon_cli/shell_pipeline.py +++ b/sentieon_cli/shell_pipeline.py @@ -1,33 +1,97 @@ -"""Complex shell pipelines with process substitution""" +"""Complex shell pipelines with process substitution. + +POSIX only: pipelines are wired together with OS pipes and FIFOs, and process +substitution relies on named pipes, so this module does not work on Windows. +""" from __future__ import annotations import asyncio -import io +import dataclasses +import fcntl import os import pathlib +import resource +import subprocess import tempfile import shlex -import sys from abc import ABC, abstractmethod -from typing import Any, Dict, IO, Iterable, List, Optional, Union +from typing import ( + Any, + Dict, + IO, + List, + Optional, + Tuple, + TYPE_CHECKING, + Union, +) from .logging import get_logger +if TYPE_CHECKING: + from .run_logs import JobLogSink + logger = get_logger(__name__) +def _freeze(value: Any) -> Any: + """Return an order-independent, hashable view of ``value``.""" + if isinstance(value, dict): + return frozenset((k, _freeze(v)) for k, v in value.items()) + if isinstance(value, (list, tuple)): + return tuple(_freeze(v) for v in value) + if isinstance(value, (set, frozenset)): + return frozenset(_freeze(v) for v in value) + return value + + +def _set_pipe_size(fd: int, size: int) -> None: + """Best-effort enlarge the OS pipe buffer backing ``fd``. + + Lets a throughput-bound pipeline hold more in-flight data between + stages. + """ + set_pipe_sz = getattr(fcntl, "F_SETPIPE_SZ", None) + if set_pipe_sz is None: + return + try: + fcntl.fcntl(fd, set_pipe_sz, size) + except OSError as exc: + logger.warning( + "could not set pipe size to %d (%s); using the default. " + "Raise fs.pipe-max-size to allow larger buffers.", + size, + exc, + ) + + +@dataclasses.dataclass +class _ProcSubWait: + """A proc-sub FIFO whose blocking open may need unblocking.""" + + fifo_path: str + opened: bool = False # run_inner's open() returned (set pre-await) + unblocked: bool = False # cleanup already opened both ends + + class Context: """Holds shared resources.""" - def __init__(self) -> None: + def __init__(self, log_sink: Optional[JobLogSink] = None) -> None: # This directory is already unique per pipeline run self.temp_dir = tempfile.TemporaryDirectory() # Background tasks to wait for (e.g., proc sub processes) - self.tasks: List[asyncio.Task] = [] + self.tasks: List[asyncio.Task[Any]] = [] # Hold the commands run in this context self.commands: List[Command] = [] self._counter = 0 # Monotonic counter - self.file_handles: List[io.IOBase] = [] + self.file_handles: List[IO[Any]] = [] + # Per-process log files; without a sink, stdout/stderr are inherited + self.log_sink = log_sink + # Set by cleanup(); a proc sub whose FIFO open was unblocked by + # cleanup must not spawn its inner command. + self.closing: bool = False + self._procsub_waits: List[_ProcSubWait] = [] def get_new_fifo(self) -> str: """Creates a new unique FIFO and returns its path.""" @@ -38,11 +102,90 @@ def get_new_fifo(self) -> str: return path async def cleanup(self) -> None: - if self.tasks: - await asyncio.gather(*self.tasks) + """Wait for background tasks, then release all resources. + + A proc sub's blocking FIFO open() never returns if the outer + command exited without opening the other end, so cleanup first + unblocks those opens, then waits for the tasks. Task exceptions + are re-raised only after every resource has been released. + """ + self.closing = True + unblock_fds: List[int] = [] + try: + while True: + for wait in self._procsub_waits: + if wait.opened or wait.unblocked: + continue + wait.unblocked = True + # Hold both ends open: this unblocks a blocked open() + # of either end, and read-then-write cannot raise + # ENXIO. + unblock_fds.append( + os.open(wait.fifo_path, os.O_RDONLY | os.O_NONBLOCK) + ) + unblock_fds.append( + os.open(wait.fifo_path, os.O_WRONLY | os.O_NONBLOCK) + ) + # Tasks can register new proc subs (nested pipelines), so + # re-snapshot until everything is done. + pending = [t for t in self.tasks if not t.done()] + if not pending: + break + await asyncio.gather(*pending, return_exceptions=True) + finally: + for fd in unblock_fds: + os.close(fd) + # asyncio's child watcher reaped children nobody awaited; Popen does + # not, so sweep up already-exited processes here to avoid zombies for + # callers that only await the final pipeline stage. poll() is + # non-blocking: a process still running is left to its owner. + for command in self.commands: + proc = command.proc + if proc is not None and proc.returncode is None: + proc.poll() for fh in self.file_handles: fh.close() self.temp_dir.cleanup() + for task in self.tasks: + if task.cancelled(): + continue + exc = task.exception() + if exc is not None and not isinstance(exc, asyncio.CancelledError): + raise exc + + +class RusagePopen(subprocess.Popen[bytes]): + """A Popen that records the child's rusage when it is reaped. + + Only the ``wait``/``async_wait`` path reaps through ``_try_wait`` and + sets ``rusage``; a reap via ``poll`` -- or ``send_signal``/``kill``, + which poll internally -- goes through the base class's machinery and + leaves ``rusage`` None. The executor waits every process it reports + metrics for, so those paths only collect processes whose rusage is + never read. + """ + + rusage: Optional[resource.struct_rusage] = None + + def _try_wait(self, wait_flags: int) -> Tuple[int, int]: + # Mirrors the base class, substituting wait4 for waitpid so the + # child's rusage is captured at reap time. Callers hold + # self._waitpid_lock (base-class contract). + try: + pid, sts, res = os.wait4(self.pid, wait_flags) + except ChildProcessError: + pid, sts = self.pid, 0 + else: + if pid == self.pid: + self.rusage = res + return (pid, sts) + + async def async_wait(self) -> int: + """Await the child's exit without blocking the event loop.""" + if self.returncode is not None: + # Already reaped; no need for a pool thread. + return self.returncode + return await asyncio.to_thread(self.wait) class ShellNode(ABC): @@ -56,22 +199,22 @@ def __init__( self, executable: str, *args: str | ProcSub, - fail_ok=False, - exec_kwargs: Optional[Dict] = None, + fail_ok: bool = False, + exec_kwargs: Optional[Dict[str, Any]] = None, ) -> None: self.executable = executable self.args = list(args) self.fail_ok = fail_ok - self.proc: Union[asyncio.subprocess.Process, None] = None + self.proc: Optional[RusagePopen] = None self.exec_kwargs: Dict[str, Any] = exec_kwargs if exec_kwargs else {} async def run( self, context: Context, - stdin: Union[IO, int, None] = None, - stdout: Union[IO, int, None] = None, - stderr: Union[IO, int, None] = None, - ) -> asyncio.subprocess.Process: + stdin: Union[IO[Any], int, None] = None, + stdout: Union[IO[Any], int, None] = None, + stderr: Union[IO[Any], int, None] = None, + ) -> RusagePopen: # 1. Resolve Arguments (Handle Process Substitutions) final_args = [self.executable] @@ -83,10 +226,23 @@ async def run( else: final_args.append(str(arg)) - # 2. Run the process + # 2. Capture the streams the caller left unset + # An explicit stream (a pipe between stages, a file_output, a proc-sub + # FIFO) always wins; only an inherited stream goes to the sink. + if context.log_sink is not None: + if stderr is None: + stderr_log = context.log_sink.open_stderr(self) + context.file_handles.append(stderr_log) + stderr = stderr_log + if stdout is None: + stdout_log = context.log_sink.open_stdout(self) + context.file_handles.append(stdout_log) + stdout = stdout_log + + # 3. Run the process # We allow the caller to wait for this process - self.proc = await asyncio.create_subprocess_exec( - *final_args, + self.proc = RusagePopen( + final_args, stdin=stdin, stdout=stdout, stderr=stderr, @@ -96,13 +252,28 @@ async def run( return self.proc def __str__(self) -> str: - return shlex.join([self.executable] + [str(x) for x in self.args]) + rendered = shlex.join([self.executable] + [str(x) for x in self.args]) + # Render env/cwd so the command is self-contained bash (backends + # submit str(job.shell)). Other exec_kwargs (close_fds, ...) have no + # shell rendering and are omitted. env is a conventional K=V prefix + # that merges into the inherited environment. + env = self.exec_kwargs.get("env") + if env: + assignments = " ".join( + f"{key}={shlex.quote(str(val))}" for key, val in env.items() + ) + rendered = f"{assignments} {rendered}" + cwd = self.exec_kwargs.get("cwd") + if cwd: + rendered = f"(cd {shlex.quote(str(cwd))} && {rendered})" + return rendered def __repr__(self) -> str: return ( f"{self.__class__.__name__}({self.executable}, " + ", ".join([repr(x) for x in self.args]) - + f", fail_ok={self.fail_ok})" + + f", fail_ok={self.fail_ok}, " + + f"exec_kwargs={self.exec_kwargs!r})" ) def __eq__(self, other: object) -> bool: @@ -112,11 +283,14 @@ def __eq__(self, other: object) -> bool: self.executable == other.executable and self.args == other.args and self.fail_ok == other.fail_ok - and self.proc == other.proc + and self.exec_kwargs == other.exec_kwargs ) def __hash__(self) -> int: - return hash(tuple([self.executable, self.fail_ok] + self.args)) + return hash( + tuple([self.executable, self.fail_ok] + self.args) + + (_freeze(self.exec_kwargs),) + ) class Pipeline(ShellNode): @@ -125,39 +299,38 @@ class Pipeline(ShellNode): def __init__( self, *nodes: Command, - skip_pipe: Optional[Iterable[int]] = None, + pipe_size: Optional[int] = None, file_input: Optional[pathlib.Path] = None, file_output: Optional[pathlib.Path] = None, ): self.nodes = list(nodes) if not self.nodes: raise ValueError("Pipeline nodes cannot be empty") - self.skip_pipe = set(skip_pipe) if skip_pipe else set() + # Best-effort OS pipe-buffer size (bytes) for every internal pipe; + # None keeps the kernel default. A local-execution tuning hint, so + # it is not part of pipeline identity (like Job.threads). + self.pipe_size = pipe_size self.file_input = file_input self.file_output = file_output async def run( self, context: Context, - stdin: Union[IO, int, None] = None, - stdout: Union[IO, int, None] = None, - stderr: Union[IO, int, None] = None, - ) -> asyncio.subprocess.Process: - # We cannot have handles from both files and to run() - if self.file_input and stdin: - logger.error( - "Pipeline %s with both file and stdin input. %s", - self, - str(stdin), + stdin: Union[IO[Any], int, None] = None, + stdout: Union[IO[Any], int, None] = None, + stderr: Union[IO[Any], int, None] = None, + ) -> RusagePopen: + # We cannot have handles from both files and to run(). + if self.file_input is not None and stdin is not None: + raise ValueError( + f"Pipeline {self} was given both a file_input and a stdin " + f"stream ({stdin!r}); provide only one." ) - sys.exit(1) - if self.file_output and stdout: - logger.error( - "Pipeline %s with both file and stdout output. %s", - self, - str(stdout), + if self.file_output is not None and stdout is not None: + raise ValueError( + f"Pipeline {self} was given both a file_output and a stdout " + f"stream ({stdout!r}); provide only one." ) - sys.exit(1) processes = [] # We need to close these FDs in the parent after passing them to @@ -166,19 +339,18 @@ async def run( # pass the file handle as input to the pipeline current_stdin = stdin - if self.file_input: + if self.file_input is not None: current_stdin = open(self.file_input, "rb") context.file_handles.append(current_stdin) try: # Chain all commands except the last one for i in range(len(self.nodes) - 1): - read_fd = None - write_fd = None - if i not in self.skip_pipe: - # Create OS pipe - read_fd, write_fd = os.pipe() - fds_to_close.extend([read_fd, write_fd]) + # Create the OS pipe linking this command to the next. + read_fd, write_fd = os.pipe() + fds_to_close.extend([read_fd, write_fd]) + if self.pipe_size is not None: + _set_pipe_size(write_fd, self.pipe_size) # Run the node (Command or sub-Pipeline) # stdout goes to the write-end of the pipe @@ -190,16 +362,15 @@ async def run( ) processes.append(proc) - if write_fd: - # Close write_fd in parent (child has it now) - os.close(write_fd) - fds_to_close.remove(write_fd) + # Close write_fd in parent (child has it now) + os.close(write_fd) + fds_to_close.remove(write_fd) # Next command reads from the read-end current_stdin = read_fd # Send the pipeline output to a file - if self.file_output: + if self.file_output is not None: stdout = open(self.file_output, "wb") context.file_handles.append(stdout) @@ -226,13 +397,7 @@ def __str__(self) -> str: res = [] if self.file_input: res.append(f"<'{self.file_input}' ") - res.append(str(self.nodes[0])) - for i, node in enumerate(self.nodes[1:]): - if i in self.skip_pipe: - res.append("; ") - else: - res.append(" | ") - res.append(str(node)) + res.append(" | ".join(str(node) for node in self.nodes)) if self.file_output: res.append(f" >'{self.file_output}'") return "".join(res) @@ -241,26 +406,23 @@ def __repr__(self) -> str: return ( f"{self.__class__.__name__}(" + ", ".join(repr(x) for x in self.nodes) - + f"skip_pipe={self.skip_pipe},file_input={self.file_input}," + + f", pipe_size={self.pipe_size}, " + + f"file_input={self.file_input}, " + f"file_output={self.file_output})" ) def __eq__(self, other: object) -> bool: if not isinstance(other, Pipeline): return NotImplemented + # pipe_size is a tuning hint, not part of identity (see __init__). return ( self.nodes == other.nodes - and self.skip_pipe == other.skip_pipe and self.file_input == other.file_input and self.file_output == other.file_output ) def __hash__(self) -> int: attrs: List[Any] = self.nodes + [self.file_input, self.file_output] - if self.skip_pipe is None: - attrs.append(None) - else: - attrs.append(tuple(self.skip_pipe)) return hash(tuple(attrs)) @@ -282,7 +444,7 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: if not isinstance(other, ProcSub): return NotImplemented - return self.node == other.node + return type(self) is type(other) and self.node == other.node def __hash__(self) -> int: return hash((self.__class__.__name__, self.node)) @@ -296,35 +458,39 @@ class InputProcSub(ProcSub): """ async def setup(self, context: Context) -> str: - self.fifo_path = context.get_new_fifo() + fifo_path = context.get_new_fifo() + self.fifo_path = fifo_path + wait = _ProcSubWait(fifo_path) # We must open the FIFO for writing. # WARNING: opening a FIFO blocks until the other end is opened. # We run the open() in a thread to prevent blocking the async event # loop. - def open_fifo_write(): - if not self.fifo_path: - raise RuntimeError("FIFO path not set") - return open(self.fifo_path, "w") + def open_fifo_write() -> IO[Any]: + return open(fifo_path, "w") # Start the task that waits for the open, then runs the process - async def run_inner(): - # This awaits until the outer command opens the file for reading! + async def run_inner() -> None: + # This awaits until the outer command opens the file for + # reading (or until cleanup() unblocks the open)! write_handle = await asyncio.to_thread(open_fifo_write) + wait.opened = True try: + if context.closing: + return # cleanup unblocked us; do not spawn proc = await self.node.run( context, stdout=write_handle, ) - await proc.wait() + await proc.async_wait() finally: write_handle.close() # Schedule this task in the background - task = asyncio.create_task(run_inner()) - context.tasks.append(task) + context.tasks.append(asyncio.create_task(run_inner())) + context._procsub_waits.append(wait) - return self.fifo_path + return fifo_path def __str__(self) -> str: return "<(" + str(self.node) + ")" @@ -338,26 +504,30 @@ class OutputProcSub(ProcSub): """ async def setup(self, context: Context) -> str: - self.fifo_path = context.get_new_fifo() + fifo_path = context.get_new_fifo() + self.fifo_path = fifo_path + wait = _ProcSubWait(fifo_path) - def open_fifo_read(): - if not self.fifo_path: - raise RuntimeError("FIFO path not set") - return open(self.fifo_path, "r") + def open_fifo_read() -> IO[Any]: + return open(fifo_path, "r") - async def run_inner(): - # This awaits until the outer command opens the file for writing! + async def run_inner() -> None: + # This awaits until the outer command opens the file for + # writing (or until cleanup() unblocks the open)! read_handle = await asyncio.to_thread(open_fifo_read) + wait.opened = True try: + if context.closing: + return # cleanup unblocked us; do not spawn proc = await self.node.run(context, stdin=read_handle) - await proc.wait() + await proc.async_wait() finally: read_handle.close() - task = asyncio.create_task(run_inner()) - context.tasks.append(task) + context.tasks.append(asyncio.create_task(run_inner())) + context._procsub_waits.append(wait) - return self.fifo_path + return fifo_path def __str__(self) -> str: return ">(" + str(self.node) + ")" diff --git a/sentieon_cli/stages/__init__.py b/sentieon_cli/stages/__init__.py new file mode 100644 index 0000000..e74d5da --- /dev/null +++ b/sentieon_cli/stages/__init__.py @@ -0,0 +1 @@ +"""Reusable pipeline stages""" diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py new file mode 100644 index 0000000..32736df --- /dev/null +++ b/sentieon_cli/stages/alignment.py @@ -0,0 +1,456 @@ +""" +Read alignment and read extraction + +Alignment jobs are the roots of every pipeline DAG: they read the run's +input files and depend on nothing. Because a few pipelines need the jobs +before they insert them, these stages split the work in two -- ``build`` +constructs the jobs, ``add_to`` inserts what ``build`` returned. +""" + +from abc import abstractmethod +from dataclasses import dataclass, field +import itertools +import logging +import pathlib +import shutil +from typing import Iterable, List, Optional, Union + +import packaging.version + +from .. import command_strings as cmds +from ..dag import DAG +from ..driver import ReadWriter +from ..job import Job +from .base import Stage, StageResult, driver_job + +BWA_REALIGN_MIN_VERSIONS = { + "sentieon driver": packaging.version.Version("202308"), + "samtools": packaging.version.Version("1.16"), +} + +BWA_FASTQ_MIN_VERSIONS = { + "sentieon driver": packaging.version.Version("202308"), +} + +MINIMAP2_REALIGN_MIN_VERSIONS = { + "sentieon driver": packaging.version.Version("202308"), + "samtools": packaging.version.Version("1.16"), +} + +MINIMAP2_FASTQ_MIN_VERSIONS = { + "sentieon driver": packaging.version.Version("202308"), +} + + +def aln_suffix(bam_format: bool) -> str: + """The file extension for an aligned output file""" + return "bam" if bam_format else "cram" + + +def find_unzip(logger: logging.Logger) -> str: + """The decompression tool to read gzipped fastq with. + + ``igzip`` is considerably faster but is not always installed, so fall + back to ``gzip``. + """ + unzip = "igzip" + if not shutil.which(unzip): + logger.warning( + "igzip is recommended for decompression, but is not " + "available. Falling back to gzip.", + ) + unzip = "gzip" + return unzip + + +@dataclass(kw_only=True) +class AlignResult(StageResult): + """The jobs and outputs of an `AlignmentStage`. + + * ``outputs`` -- the aligned files, in input order. + * ``cleanup_paths`` -- files a later ``rm`` job should delete once the + alignment has been consumed; empty when the stage writes only + final outputs. + """ + + outputs: List[pathlib.Path] + cleanup_paths: List[pathlib.Path] + + +@dataclass(kw_only=True) +class AlignmentStage(Stage): + """A group of alignment jobs, all of them DAG roots. + + Subclasses implement :meth:`build`, which constructs the jobs and + inserts nothing. :meth:`add_to` then adds each of them with the same + dependencies, since no alignment job depends on another. + """ + + @abstractmethod + def build(self) -> AlignResult: + """Construct this stage's jobs without touching a DAG""" + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> AlignResult: + deps = set(upstream) + result = self.build() + for job in result.jobs: + dag.add_job(job, deps) + return result + + +@dataclass(kw_only=True) +class BwaRealignStage(AlignmentStage): + """Re-align BAM/CRAM/uBAM/uCRAM input to the reference with bwa. + + The @RG lines of each input are read at build time and written to a + header file in the run's temporary directory, which the alignment + command passes to `samtools fastq`. + + When duplicate marking follows, the alignment is an intermediate: it + is written to the temporary directory as a lightly compressed BAM + instead of to the output directory. + """ + + inputs: List[pathlib.Path] + model_bundle: pathlib.Path + bam_format: bool = False + duplicate_marking: str = "markdup" + input_ref: Optional[pathlib.Path] = None + collate: bool = False + bwa_args: str = "" + bwa_k: str = "20000000" + util_sort_args: str = "" + + def build(self) -> AlignResult: + suffix = aln_suffix(self.bam_format) + util_sort_args = self.util_sort_args + dedup_follows = self.duplicate_marking != "none" + if dedup_follows: + suffix = "bam" + util_sort_args += " --bam_compression 1 " + + jobs: List[Job] = [] + outputs: List[pathlib.Path] = [] + cleanup_paths: List[pathlib.Path] = [] + for i, input_aln in enumerate(self.inputs): + out_aln = pathlib.Path( + str(self.ctx.output_vcf).replace( + ".vcf.gz", f"_bwa_sorted_{i}.{suffix}" + ) + ) + if dedup_follows: + out_aln = self.ctx.tmp_dir.joinpath(f"bwa_sorted_{i}.{suffix}") + rg_lines = cmds.get_rg_lines(input_aln, self.ctx.dry_run) + rg_header = self.ctx.tmp_dir.joinpath(f"input_{i}.hdr") + with open(rg_header, "w", encoding="utf-8") as rg_fh: + for line in rg_lines: + print(line, file=rg_fh) + jobs.append( + Job( + cmds.cmd_samtools_fastq_bwa( + out_aln, + input_aln, + self.ctx.reference, + self.model_bundle, + self.ctx.cores, + rg_header, + self.input_ref, + collate=self.collate, + bwa_args=self.bwa_args, + bwa_k=self.bwa_k, + util_sort_args=util_sort_args, + ), + f"bam-align-{i}", + self.ctx.cores, + task_name="alignment", + ) + ) + outputs.append(out_aln) + if dedup_follows: + cleanup_paths.append(out_aln) + cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) + + return AlignResult( + jobs=jobs, + terminal=set(jobs), + outputs=outputs, + cleanup_paths=cleanup_paths, + ) + + +@dataclass(kw_only=True) +class BwaFastqStage(AlignmentStage): + """Align short-read fastq pairs to the reference with bwa. + + With more than one NUMA node the reads of each pair are split across + the nodes: one job per (fastq pair, node), each pinned with `taskset` + and taking a share of the cores and a node token, so the scheduler + runs exactly one alignment per node. + """ + + r1_fastq: List[pathlib.Path] + r2_fastq: List[pathlib.Path] + readgroups: List[str] + model_bundle: pathlib.Path + numa_nodes: List[str] = field(default_factory=list) + bam_format: bool = False + duplicate_marking: str = "markdup" + unzip: str = "gzip" + bwa_args: str = "" + bwa_k: str = "20000000" + util_sort_args: str = "" + + def build(self) -> AlignResult: + n_alignment_jobs = max(1, len(self.numa_nodes)) + split_alignment = len(self.numa_nodes) > 1 + suffix = aln_suffix(self.bam_format) + util_sort_args = self.util_sort_args + dedup_follows = self.duplicate_marking != "none" + if dedup_follows: + suffix = "bam" + util_sort_args += " --bam_compression 1 " + + jobs: List[Job] = [] + outputs: List[pathlib.Path] = [] + cleanup_paths: List[pathlib.Path] = [] + for i, (r1, r2, rg) in enumerate( + itertools.zip_longest( + self.r1_fastq, self.r2_fastq, self.readgroups + ) + ): + for j in range(n_alignment_jobs): + out_aln = pathlib.Path( + str(self.ctx.output_vcf).replace( + ".vcf.gz", f"_bwa_sorted_fq_{i}_{j}.{suffix}" + ) + ) + if dedup_follows: + out_aln = self.ctx.tmp_dir.joinpath( + f"bwa_sorted_fq_{i}_{j}.{suffix}" + ) + numa = self.numa_nodes[j] if split_alignment else None + split = f"{j}/{n_alignment_jobs}" if split_alignment else None + split_cores = max(1, int(self.ctx.cores / n_alignment_jobs)) + jobs.append( + Job( + cmds.cmd_fastq_bwa( + out_aln, + r1, + r2, + rg, + self.ctx.reference, + self.model_bundle, + split_cores, + self.unzip, + self.bwa_args, + self.bwa_k, + util_sort_args, + numa, + split, + ), + f"bam-align-{i}-{j}", + split_cores, + resources={f"node{j}": 1}, + task_name="alignment", + ) + ) + outputs.append(out_aln) + if dedup_follows: + cleanup_paths.append(out_aln) + cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) + + return AlignResult( + jobs=jobs, + terminal=set(jobs), + outputs=outputs, + cleanup_paths=cleanup_paths, + ) + + +@dataclass(kw_only=True) +class Minimap2RealignStage(AlignmentStage): + """Re-align long-read BAM/CRAM/uBAM/uCRAM input with minimap2. + + ``minimap2_model`` overrides the model the aligner picks up from the + bundle; the pangenome pipelines use it to select the long-read model. + """ + + inputs: List[pathlib.Path] + model_bundle: pathlib.Path + sample_name: str + bam_format: bool = False + input_ref: Optional[pathlib.Path] = None + fastq_taglist: str = "*" + minimap2_args: str = "-YL" + util_sort_args: str = "" + minimap2_model: Optional[Union[pathlib.Path, str]] = None + + def build(self) -> AlignResult: + suffix = aln_suffix(self.bam_format) + jobs: List[Job] = [] + outputs: List[pathlib.Path] = [] + for i, input_aln in enumerate(self.inputs): + out_aln = pathlib.Path( + str(self.ctx.output_vcf).replace( + ".vcf.gz", f"_mm2_sorted_{i}.{suffix}" + ) + ) + rg_lines = cmds.get_rg_lines(input_aln, self.ctx.dry_run) + jobs.append( + Job( + cmds.cmd_samtools_fastq_minimap2( + out_aln, + input_aln, + self.ctx.reference, + self.model_bundle, + self.ctx.cores, + rg_lines, + self.sample_name, + self.input_ref, + self.fastq_taglist, + self.minimap2_args, + self.util_sort_args, + self.minimap2_model, + ), + f"bam-realign-{i}", + self.ctx.cores, + task_name="alignment", + ) + ) + outputs.append(out_aln) + + return AlignResult( + jobs=jobs, + terminal=set(jobs), + outputs=outputs, + cleanup_paths=[], + ) + + +@dataclass(kw_only=True) +class Minimap2FastqStage(AlignmentStage): + """Align long-read fastq to the reference with minimap2""" + + fastq: List[pathlib.Path] + readgroups: List[str] + model_bundle: pathlib.Path + bam_format: bool = False + unzip: str = "gzip" + minimap2_args: str = "-YL" + util_sort_args: str = "" + + def build(self) -> AlignResult: + suffix = aln_suffix(self.bam_format) + jobs: List[Job] = [] + outputs: List[pathlib.Path] = [] + for i, (fq, rg) in enumerate(zip(self.fastq, self.readgroups)): + out_aln = pathlib.Path( + str(self.ctx.output_vcf).replace( + ".vcf.gz", f"_mm2_sorted_fq_{i}.{suffix}" + ) + ) + jobs.append( + Job( + cmds.cmd_fastq_minimap2( + out_aln, + fq, + rg, + self.ctx.reference, + self.model_bundle, + self.ctx.cores, + self.unzip, + self.minimap2_args, + self.util_sort_args, + ), + f"align-{i}", + self.ctx.cores, + task_name="alignment", + ) + ) + outputs.append(out_aln) + + return AlignResult( + jobs=jobs, + terminal=set(jobs), + outputs=outputs, + cleanup_paths=[], + ) + + +@dataclass(kw_only=True) +class BwaExtractStage(AlignmentStage): + """Align short-read fastq with bwa and extract the reads a pangenome + pipeline needs, in a single pass. + + ``readgroup`` is the rendered `@RG` string; the pangenome pipelines + derive it differently from the input readgroup. + """ + + output_bam: pathlib.Path + output_fastq: pathlib.Path + r1_fastq: List[pathlib.Path] + r2_fastq: List[pathlib.Path] + readgroup: str + extract_model: pathlib.Path + bwa_model: pathlib.Path + unzip: str = "gzip" + name: str = "bwa-extract" + task_name: str = "alignment" + + 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, + ) + return AlignResult( + jobs=[job], + terminal={job}, + outputs=[self.output_bam, self.output_fastq], + cleanup_paths=[], + ) + + +@dataclass(kw_only=True) +class ReadWriterStage(AlignmentStage): + """Rewrite alignments with `sentieon driver --algo ReadWriter`. + + With several inputs the driver merges them; with an ``interval`` it + extracts the reads overlapping that region. + """ + + inputs: List[pathlib.Path] + output: pathlib.Path + name: str + task_name: str + threads: Optional[int] = None + interval: Optional[Union[pathlib.Path, str]] = None + + def build(self) -> AlignResult: + job = driver_job( + self.ctx, + [ReadWriter(self.output)], + inputs=self.inputs, + interval=self.interval, + name=self.name, + task_name=self.task_name, + threads=self.threads, + ) + return AlignResult( + jobs=[job], + terminal={job}, + outputs=[self.output], + cleanup_paths=[], + ) diff --git a/sentieon_cli/stages/base.py b/sentieon_cli/stages/base.py new file mode 100644 index 0000000..ecee86f --- /dev/null +++ b/sentieon_cli/stages/base.py @@ -0,0 +1,127 @@ +""" +Building blocks for pipeline stages + +A *stage* is a reusable piece of a pipeline: it knows how to build the jobs +for one operation (dedup, metrics, small-variant calling, ...) and how to +wire them into a :class:`~sentieon_cli.dag.DAG`. Stages are configured with +plain dataclass fields plus a :class:`StageContext` of run-wide settings. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +import pathlib +from typing import Iterable, List, Optional, Sequence, Set, Union + +from ..dag import DAG +from ..driver import BaseAlgo, Driver +from ..job import Job +from ..shell_pipeline import Command, Pipeline + + +@dataclass(frozen=True) +class StageContext: + """Run-wide settings every stage needs.""" + + reference: pathlib.Path + output_vcf: pathlib.Path + tmp_dir: pathlib.Path + cores: int + dry_run: bool + skip_version_check: bool + + +@dataclass(kw_only=True) +class StageResult: + """The jobs a stage added to the DAG. + + * ``jobs`` -- every job the stage inserted, in insertion order. It + exists for the occasional edge that has to attach to a job in the + middle of a stage rather than to its output. + * ``terminal`` -- the stage's sink job(s), the ones that carry its + declared outputs. Downstream stages depend on these. + + Subclasses add named fields for the files the stage produces. + """ + + jobs: List[Job] + terminal: Set[Job] + + +@dataclass(kw_only=True) +class Stage(ABC): + """A reusable group of jobs that inserts itself into a DAG. + + Subclasses are dataclasses (``@dataclass(kw_only=True)``) that add + their own configuration fields and return a :class:`StageResult` + subclass naming the outputs they produce. + + The :meth:`add_to` contract: + + * ``upstream`` is any iterable of jobs already in the DAG; it is + consumed once into a set. + * Every *entry* job -- an in-stage job with no in-stage dependency -- + is inserted with that set as its dependencies. Edges between the + stage's own jobs are added on top of that. + * Insertion is topological (a job is added only after the jobs it + depends on) and deterministic: the same stage and DAG produce the + same jobs, in the same order, every time. + * The jobs are created inside ``add_to``; constructing a stage has no + side effects and assigns no job ids. + * ``add_to`` may be called only once per stage instance. Job identity + is the shell pipeline, so a second call raises in + :meth:`~sentieon_cli.dag.DAG.add_job`. + """ + + ctx: StageContext + + @abstractmethod + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> StageResult: + """Add this stage's jobs to ``dag``, downstream of ``upstream``""" + + +def driver_job( + ctx: StageContext, + algos: Sequence[BaseAlgo], + *, + name: str, + task_name: str, + threads: Optional[int] = None, + inputs: Optional[List[pathlib.Path]] = None, + interval: Optional[Union[pathlib.Path, str]] = None, + interval_padding: Optional[int] = None, + read_filter: Optional[List[str]] = None, + replace_rg: Optional[List[List[str]]] = None, +) -> Job: + """A job running `sentieon driver` with ``algos``. + + ``threads`` defaults to the run's core count; the driver arguments are + passed through untouched so the command matches what the pipelines + build by hand. + """ + driver = Driver( + reference=ctx.reference, + thread_count=ctx.cores, + interval=interval, + interval_padding=interval_padding, + read_filter=read_filter, + replace_rg=replace_rg, + input=inputs, + ) + for algo in algos: + driver.add_algo(algo) + return Job( + Pipeline(Command(*driver.build_cmd())), + name, + ctx.cores if threads is None else threads, + task_name=task_name, + ) + + +def rm_job(paths: Iterable[Union[pathlib.Path, str]], name: str) -> Job: + """A job removing intermediate files, tolerating missing ones""" + return Job( + Pipeline(Command("rm", *[str(p) for p in paths], fail_ok=True)), + name, + 0, + task_name="cleanup", + ) diff --git a/sentieon_cli/stages/cnv.py b/sentieon_cli/stages/cnv.py new file mode 100644 index 0000000..d02e48e --- /dev/null +++ b/sentieon_cli/stages/cnv.py @@ -0,0 +1,206 @@ +""" +CNV calling with CNVscope +""" + +from dataclasses import dataclass +from importlib.resources import files +import pathlib +from typing import Iterable, List, Optional + +import packaging.version + +from .. import command_strings as cmds +from ..dag import DAG +from ..driver import CNVModelApply, CNVscope +from ..job import Job +from ..util import SampleSex, cnvscope_sex_args +from .base import Stage, StageResult, driver_job + +CNV_MIN_VERSIONS = { + # 202503.04 adds the CNVscope `--sex` and `--par` arguments + "sentieon driver": packaging.version.Version("202503.04"), +} + + +@dataclass(kw_only=True) +class CNVResult(StageResult): + """The jobs and output of `CNVscopeStage`""" + + cnvscope_job: Job + apply_job: Job + cnv_vcf: pathlib.Path + + +@dataclass(kw_only=True) +class CNVscopeStage(Stage): + """Call CNVs with CNVscope, then filter them with CNVModelApply. + + The sample sex has to be known before the jobs are built, so callers + run this stage after ploidy estimation. Version checks stay with the + pipelines, which run them during validation. ``dump_probes`` is the + path CNVscope writes its per-probe debugging output to. + """ + + inputs: List[pathlib.Path] + model: pathlib.Path + cnvscope_vcf: pathlib.Path + cnv_vcf: pathlib.Path + sample_sex: Optional[SampleSex] = None + par_bed: Optional[pathlib.Path] = None + dump_probes: Optional[pathlib.Path] = None + interval: Optional[pathlib.Path] = None + replace_rg: Optional[List[List[str]]] = None + name: str = "cnvscope" + apply_name: str = "cnv-model-apply" + task_name: str = "cnv" + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> CNVResult: + deps = set(upstream) + sex, par = cnvscope_sex_args(self.sample_sex, self.par_bed) + + cnvscope = driver_job( + self.ctx, + [ + CNVscope( + self.cnvscope_vcf, + self.model, + sex=sex, + par=par, + dump_probes=self.dump_probes, + ) + ], + inputs=self.inputs, + interval=self.interval, + replace_rg=self.replace_rg, + name=self.name, + task_name=self.task_name, + ) + apply_job = driver_job( + self.ctx, + [ + CNVModelApply( + self.cnv_vcf, + self.model, + vcf=self.cnvscope_vcf, + ) + ], + name=self.apply_name, + task_name=self.task_name, + ) + + dag.add_job(cnvscope, deps) + dag.add_job(apply_job, {cnvscope}) + + return CNVResult( + jobs=[cnvscope, apply_job], + terminal={apply_job}, + cnvscope_job=cnvscope, + apply_job=apply_job, + cnv_vcf=self.cnv_vcf, + ) + + +@dataclass(kw_only=True) +class PangenomeCNVResult(StageResult): + """The jobs and output of `PangenomeCNVStage`""" + + cnvscope_job: Job + apply_job: Job + indel2cnv_job: Job + combine_job: Job + cnv_vcf: pathlib.Path + + +@dataclass(kw_only=True) +class PangenomeCNVStage(Stage): + """CNV calling for the pangenome pipelines. + + CNVscope and CNVModelApply run on the short reads, the PangenomeSV + INDELs are converted to CNVs with `indel2cnv.py`, and the two call + sets are combined by `combine_cnv.py` into `_cnv.vcf.gz`. + Both inputs -- the alignments and the SV VCF -- come from an earlier + DAG, and the sample sex has to be known before the jobs are built. + """ + + inputs: List[pathlib.Path] + model: pathlib.Path + sv_vcf: pathlib.Path + sample_sex: Optional[SampleSex] = None + par_bed: Optional[pathlib.Path] = None + interval: Optional[pathlib.Path] = None + replace_rg: Optional[List[List[str]]] = None + task_name: str = "cnv" + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> PangenomeCNVResult: + deps = set(upstream) + + cnvscope_vcf = self.ctx.tmp_dir.joinpath("sample-cnvscope.vcf.gz") + apply_vcf = self.ctx.tmp_dir.joinpath("sample-cnv_model_apply.vcf.gz") + converted_vcf = self.ctx.tmp_dir.joinpath("sample-sv_cnv.vcf.gz") + cnv_vcf = pathlib.Path( + str(self.ctx.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") + ) + cnv_probes = pathlib.Path( + str(self.ctx.output_vcf).replace(".vcf.gz", "_cnv.probes") + ) + + cnv_result = CNVscopeStage( + ctx=self.ctx, + inputs=self.inputs, + model=self.model, + cnvscope_vcf=cnvscope_vcf, + cnv_vcf=apply_vcf, + sample_sex=self.sample_sex, + par_bed=self.par_bed, + dump_probes=cnv_probes, + interval=self.interval, + replace_rg=self.replace_rg, + task_name=self.task_name, + ).add_to(dag, deps) + + # Convert the PangenomeSV output to CNVs + indel2cnv_script = pathlib.Path( + str(files("sentieon_cli.scripts").joinpath("indel2cnv.py")) + ) + indel2cnv_job = Job( + cmds.cmd_pyexec_indel2cnv( + converted_vcf, + self.sv_vcf, + self.ctx.reference, + indel2cnv_script, + self.ctx.cores, + ), + "indel2cnv", + 0, + task_name=self.task_name, + ) + dag.add_job(indel2cnv_job, deps) + + # Combine the CNVModelApply output with the converted SVs + combine_script = pathlib.Path( + str(files("sentieon_cli.scripts").joinpath("combine_cnv.py")) + ) + combine_job = Job( + cmds.cmd_pyexec_combine_cnv( + cnv_vcf, + apply_vcf, + converted_vcf, + combine_script, + ), + "combine-cnv", + 0, + task_name=self.task_name, + ) + dag.add_job(combine_job, {cnv_result.apply_job, indel2cnv_job}) + + return PangenomeCNVResult( + jobs=[*cnv_result.jobs, indel2cnv_job, combine_job], + terminal={combine_job}, + cnvscope_job=cnv_result.cnvscope_job, + apply_job=cnv_result.apply_job, + indel2cnv_job=indel2cnv_job, + combine_job=combine_job, + cnv_vcf=cnv_vcf, + ) diff --git a/sentieon_cli/stages/dedup.py b/sentieon_cli/stages/dedup.py new file mode 100644 index 0000000..534dd19 --- /dev/null +++ b/sentieon_cli/stages/dedup.py @@ -0,0 +1,94 @@ +""" +Duplicate marking with LocusCollector and Dedup +""" + +from dataclasses import dataclass, field +import pathlib +from typing import Iterable, List, Optional + +from ..dag import DAG +from ..driver import BaseAlgo, Dedup, LocusCollector +from ..job import Job +from .base import Stage, StageResult, driver_job + + +@dataclass(kw_only=True) +class DedupResult(StageResult): + """The jobs and output of `DedupStage`""" + + lc_job: Job + dedup_job: Job + output: pathlib.Path + + +@dataclass(kw_only=True) +class DedupStage(Stage): + """Mark duplicates: LocusCollector collects the scores, Dedup applies + them. + + Both passes read the same input files with the same read filters, so + the two jobs only differ in their algo. Metrics that are cheap to + collect before duplicate marking can ride along on the LocusCollector + pass through ``lc_extra_algos``. + """ + + inputs: List[pathlib.Path] + output: pathlib.Path + score_file: pathlib.Path + tag: Optional[str] = None + consensus: bool = False + rmdup: bool = False + cram_write_options: Optional[str] = None + dedup_metrics: Optional[pathlib.Path] = None + read_filters: List[str] = field(default_factory=list) + lc_extra_algos: List[BaseAlgo] = field(default_factory=list) + threads: Optional[int] = None + task_name: str = "dedup" + + def _name(self, base: str) -> str: + return base if self.tag is None else f"{base}-{self.tag}" + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> DedupResult: + deps = set(upstream) + read_filter = self.read_filters or None + + lc_job = driver_job( + self.ctx, + [ + LocusCollector(self.score_file, consensus=self.consensus), + *self.lc_extra_algos, + ], + inputs=self.inputs, + read_filter=read_filter, + name=self._name("locuscollector"), + task_name=self.task_name, + threads=self.threads, + ) + dedup_job = driver_job( + self.ctx, + [ + Dedup( + self.output, + self.score_file, + cram_write_options=self.cram_write_options, + metrics=self.dedup_metrics, + rmdup=self.rmdup, + ) + ], + inputs=self.inputs, + read_filter=read_filter, + name=self._name("dedup"), + task_name=self.task_name, + threads=self.threads, + ) + + dag.add_job(lc_job, deps) + dag.add_job(dedup_job, {lc_job}) + + return DedupResult( + jobs=[lc_job, dedup_job], + terminal={dedup_job}, + lc_job=lc_job, + dedup_job=dedup_job, + output=self.output, + ) diff --git a/sentieon_cli/stages/expansion.py b/sentieon_cli/stages/expansion.py new file mode 100644 index 0000000..e75155c --- /dev/null +++ b/sentieon_cli/stages/expansion.py @@ -0,0 +1,71 @@ +""" +Repeat expansion genotyping with ExpansionHunter +""" + +from dataclasses import dataclass +import pathlib +from typing import Dict, Iterable, Optional + +import packaging.version + +from .. import command_strings as cmds +from ..dag import DAG +from ..job import Job +from ..util import SampleSex, caller_sex_arg +from .base import Stage, StageResult + +EXPANSION_MIN_VERSIONS: Dict[str, Optional[packaging.version.Version]] = { + "ExpansionHunter": None, +} + + +@dataclass(kw_only=True) +class ExpansionHunterResult(StageResult): + """The job and output of `ExpansionHunterStage`""" + + job: Job + output_prefix: pathlib.Path + + +@dataclass(kw_only=True) +class ExpansionHunterStage(Stage): + """Genotype repeat expansions with ExpansionHunter. + + ExpansionHunter is sex-aware, so callers run this stage after ploidy + estimation. ``output_prefix`` is the `--output-prefix` the tool + writes its files under. + """ + + alignment: pathlib.Path + variant_catalog: pathlib.Path + output_prefix: pathlib.Path + sample_sex: Optional[SampleSex] = None + name: str = "expansion-hunter" + task_name: str = "expansion-hunter" + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> ExpansionHunterResult: + deps = set(upstream) + + job = Job( + cmds.cmd_expansion_hunter( + self.output_prefix, + self.alignment, + reference=self.ctx.reference, + variant_catalog=self.variant_catalog, + sex=caller_sex_arg(self.sample_sex), + threads=self.ctx.cores, + ), + self.name, + self.ctx.cores, + task_name=self.task_name, + ) + dag.add_job(job, deps) + + return ExpansionHunterResult( + jobs=[job], + terminal={job}, + job=job, + output_prefix=self.output_prefix, + ) diff --git a/sentieon_cli/stages/metrics.py b/sentieon_cli/stages/metrics.py new file mode 100644 index 0000000..e03c9a3 --- /dev/null +++ b/sentieon_cli/stages/metrics.py @@ -0,0 +1,233 @@ +""" +Alignment metrics collection +""" + +from dataclasses import dataclass +from importlib.resources import files +import pathlib +import sys +from typing import Iterable, List, Optional + +import packaging.version + +from .. import command_strings as cmds +from ..dag import DAG +from ..driver import BaseAlgo +from ..job import Job +from ..shell_pipeline import Command, Pipeline +from .base import Stage, StageResult, driver_job + +MOSDEPTH_MIN_VERSIONS = { + "mosdepth": packaging.version.Version("0.2.6"), +} + + +@dataclass(frozen=True) +class MetricsPaths: + """The metrics directory and the files the pipelines write into it. + + Every metrics file is named after the output VCF, so the whole set is + derived once, from :meth:`from_output_vcf`, and passed around instead + of being rebuilt at each call site. + """ + + metrics_dir: pathlib.Path + metric_base: str + + @classmethod + def from_output_vcf(cls, output_vcf: pathlib.Path) -> "MetricsPaths": + """The metrics paths belonging to an output VCF""" + return cls( + metrics_dir=pathlib.Path( + str(output_vcf).replace(".vcf.gz", "_metrics") + ), + metric_base=output_vcf.name.replace(".vcf.gz", "") + ".txt", + ) + + def ensure_dir(self, dry_run: bool) -> None: + """Create the metrics directory, unless this is a dry run""" + if not dry_run: + self.metrics_dir.mkdir(exist_ok=True) + + def _metric(self, suffix: str) -> pathlib.Path: + return self.metrics_dir.joinpath(self.metric_base + suffix) + + @property + def score(self) -> pathlib.Path: + """The LocusCollector score file""" + return self._metric(".score.txt.gz") + + @property + def insert_size(self) -> pathlib.Path: + """InsertSizeMetricAlgo output""" + return self._metric(".insert_size.txt") + + @property + def mean_qual_by_cycle(self) -> pathlib.Path: + """MeanQualityByCycle output""" + return self._metric(".mean_qual_by_cycle.txt") + + @property + def base_distribution_by_cycle(self) -> pathlib.Path: + """BaseDistributionByCycle output""" + return self._metric(".base_distribution_by_cycle.txt") + + @property + def qual_distribution(self) -> pathlib.Path: + """QualDistribution output""" + return self._metric(".qual_distribution.txt") + + @property + def alignment_stat(self) -> pathlib.Path: + """AlignmentStat output""" + return self._metric(".alignment_stat.txt") + + @property + def coverage(self) -> pathlib.Path: + """CoverageMetrics output prefix (not named after the VCF)""" + return self.metrics_dir.joinpath("coverage") + + @property + def hybrid_selection(self) -> pathlib.Path: + """HsMetricAlgo output, for WES runs""" + return self._metric(".hybrid-selection.txt") + + @property + def wgs(self) -> pathlib.Path: + """WgsMetricsAlgo output, for WGS runs""" + return self._metric(".wgs.txt") + + @property + def gc_bias(self) -> pathlib.Path: + """GCBias output""" + return self._metric(".gc_bias.txt") + + @property + def gc_bias_summary(self) -> pathlib.Path: + """The GCBias summary output""" + return self._metric(".gc_bias_summary.txt") + + @property + def dedup_metrics(self) -> pathlib.Path: + """The Dedup metrics output""" + return self._metric(".dedup_metrics.txt") + + +@dataclass(kw_only=True) +class MetricsResult(StageResult): + """The jobs of `MetricsStage`""" + + metrics_job: Job + rehead_job: Optional[Job] + + +@dataclass(kw_only=True) +class MetricsStage(Stage): + """Collect alignment metrics with a single `sentieon driver` pass. + + The caller supplies the ordered algo list, since which metrics are + worth collecting depends on the assay and on whether the alignment + has been deduplicated yet. When ``rehead_metrics`` names a WGS + metrics file, a second job rewrites its header so MultiQC recognizes + it. + """ + + inputs: List[pathlib.Path] + algos: List[BaseAlgo] + interval: Optional[pathlib.Path] = None + rehead_metrics: Optional[pathlib.Path] = None + name: str = "metrics" + task_name: str = "metrics" + threads: int = 0 + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> MetricsResult: + deps = set(upstream) + + metrics = driver_job( + self.ctx, + self.algos, + inputs=self.inputs, + interval=self.interval, + name=self.name, + task_name=self.task_name, + threads=self.threads, + ) + dag.add_job(metrics, deps) + jobs = [metrics] + + rehead: Optional[Job] = None + if self.rehead_metrics is not None: + rehead_script = pathlib.Path( + str( + files("sentieon_cli.scripts").joinpath( + "rehead_wgs_metrics.py" + ) + ) + ) + rehead = Job( + Pipeline( + Command( + sys.executable, + str(rehead_script), + "--metrics_file", + str(self.rehead_metrics), + ) + ), + "rehead-metrics", + 0, + task_name="metrics", + ) + dag.add_job(rehead, {metrics}) + jobs.append(rehead) + + return MetricsResult( + jobs=jobs, + terminal={rehead} if rehead else {metrics}, + metrics_job=metrics, + rehead_job=rehead, + ) + + +@dataclass(kw_only=True) +class MosdepthResult(StageResult): + """The jobs of `MosdepthStage`""" + + +@dataclass(kw_only=True) +class MosdepthStage(Stage): + """Coverage QC with mosdepth, one job per input file. + + The jobs are independent of each other, so ``build`` constructs them + all and ``add_to`` inserts each with the same dependencies. + """ + + inputs: List[pathlib.Path] + + def build(self) -> MosdepthResult: + """Construct this stage's jobs without touching a DAG""" + jobs: List[Job] = [] + for i, input_file in enumerate(self.inputs): + mosdepth_dir = pathlib.Path( + str(self.ctx.output_vcf).replace(".vcf.gz", f"_mosdepth_{i}") + ) + jobs.append( + Job( + cmds.cmd_mosdepth( + input_file, + mosdepth_dir, + fasta=self.ctx.reference, + threads=self.ctx.cores, + ), + f"mosdepth-{i}", + 0, # Run in background + task_name="metrics", + ) + ) + return MosdepthResult(jobs=jobs, terminal=set(jobs)) + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> MosdepthResult: + deps = set(upstream) + result = self.build() + for job in result.jobs: + dag.add_job(job, deps) + return result diff --git a/sentieon_cli/stages/ploidy.py b/sentieon_cli/stages/ploidy.py new file mode 100644 index 0000000..49e87da --- /dev/null +++ b/sentieon_cli/stages/ploidy.py @@ -0,0 +1,67 @@ +""" +Sample ploidy and sex estimation +""" + +from dataclasses import dataclass +from importlib.resources import files +import pathlib +from typing import Iterable, List, Optional + +from .. import command_strings as cmds +from ..dag import DAG +from ..job import Job +from ..shard import ploidy_contigs_for_build +from .base import Stage, StageResult + + +@dataclass(kw_only=True) +class PloidyResult(StageResult): + """The job and output of `PloidyStage`""" + + job: Job + ploidy_json: pathlib.Path + + +@dataclass(kw_only=True) +class PloidyStage(Stage): + """Estimate the sample ploidy and sex from an alignment. + + The JSON report is written next to the output VCF and is read after + the DAG has run, by the pipelines that call variants with a + sex-aware caller. + """ + + inputs: List[pathlib.Path] + reference_build: Optional[str] = None + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> PloidyResult: + deps = set(upstream) + ploidy_json = pathlib.Path( + str(self.ctx.output_vcf).replace(".vcf.gz", "_ploidy.json") + ) + estimate_ploidy = pathlib.Path( + str(files("sentieon_cli.scripts").joinpath("estimate_ploidy.py")) + ).resolve() + ploidy_contigs = ploidy_contigs_for_build(self.reference_build) + job = Job( + cmds.cmd_estimate_ploidy( + ploidy_json, + self.inputs, + estimate_ploidy, + contigs=ploidy_contigs.contigs, + autosomes=ploidy_contigs.autosomes, + x_contig=ploidy_contigs.x_contig, + y_contig=ploidy_contigs.y_contig, + ), + "estimate-ploidy", + 0, + task_name="ploidy", + ) + dag.add_job(job, deps) + + return PloidyResult( + jobs=[job], + terminal={job}, + job=job, + ploidy_json=ploidy_json, + ) diff --git a/sentieon_cli/stages/preprocessing.py b/sentieon_cli/stages/preprocessing.py new file mode 100644 index 0000000..fac5d8b --- /dev/null +++ b/sentieon_cli/stages/preprocessing.py @@ -0,0 +1,178 @@ +""" +Short-read preprocessing: duplicate marking and metrics collection +""" + +from dataclasses import dataclass +import pathlib +from typing import Iterable, List, Optional + +from ..dag import DAG +from ..driver import ( + AlignmentStat, + BaseAlgo, + BaseDistributionByCycle, + CoverageMetrics, + GCBias, + HsMetricAlgo, + InsertSizeMetricAlgo, + MeanQualityByCycle, + QualDistribution, + WgsMetricsAlgo, +) +from ..job import Job +from .base import Stage, StageResult +from .dedup import DedupStage +from .metrics import MetricsPaths, MetricsStage + + +def pre_dedup_metrics_algos(paths: MetricsPaths, assay: str) -> List[BaseAlgo]: + """The metrics that do not need duplicate-marked reads""" + algos: List[BaseAlgo] = [ + MeanQualityByCycle(paths.mean_qual_by_cycle), + BaseDistributionByCycle(paths.base_distribution_by_cycle), + QualDistribution(paths.qual_distribution), + AlignmentStat(paths.alignment_stat), + ] + if assay == "WGS": + algos.append(GCBias(paths.gc_bias, summary=paths.gc_bias_summary)) + return algos + + +@dataclass(kw_only=True) +class SrPreprocessingResult(StageResult): + """The short-read dedup and metrics jobs added to a DAG. + + * ``deduped`` -- the alignment to call variants from; the input files + unchanged when duplicate marking is off. + * ``dedup_job`` -- the Dedup job, absent when duplicate marking is off. + * ``qc_jobs`` -- every metrics-producing job, for MultiQC to wait on. + """ + + deduped: List[pathlib.Path] + dedup_job: Optional[Job] + qc_jobs: List[Job] + + +@dataclass(kw_only=True) +class ShortReadPreprocessingStage(Stage): + """Mark duplicates and collect metrics from the short reads. + + Which metrics are collected, and whether they ride along on the + LocusCollector pass or run after duplicate marking, depends on the + assay and on whether a BED file restricts the run. + + ``terminal`` is the Dedup job, so downstream stages read the + deduplicated alignment. With ``duplicate_marking="none"`` there is no + Dedup job and the stage adds only a leaf metrics job that nothing + should wait on, so ``terminal`` is empty and callers depend on the + jobs that produced ``inputs`` instead. + """ + + inputs: List[pathlib.Path] + duplicate_marking: str = "markdup" + consensus: bool = False + skip_metrics: bool = False + assay: str = "WGS" + bed: Optional[pathlib.Path] = None + bam_format: bool = False + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> SrPreprocessingResult: + ctx = self.ctx + suffix = "bam" if self.bam_format else "cram" + + paths = MetricsPaths.from_output_vcf(ctx.output_vcf) + paths.ensure_dir(ctx.dry_run) + + if self.duplicate_marking == "none": + # Without duplicate marking there is nothing to collect after + # the fact, so every metric comes from one pass over the input + if self.skip_metrics: + return SrPreprocessingResult( + jobs=[], + terminal=set(), + deduped=self.inputs, + dedup_job=None, + qc_jobs=[], + ) + metrics_result = MetricsStage( + ctx=ctx, + inputs=self.inputs, + algos=[ + InsertSizeMetricAlgo(paths.insert_size), + *pre_dedup_metrics_algos(paths, self.assay), + ], + name="metrics", + task_name="metrics", + threads=ctx.cores, + ).add_to(dag, upstream) + return SrPreprocessingResult( + jobs=list(metrics_result.jobs), + terminal=set(), + deduped=self.inputs, + dedup_job=None, + qc_jobs=list(metrics_result.jobs), + ) + + # Metrics that ride along on the LocusCollector pass + lc_extra_algos: List[BaseAlgo] = [] + if not self.skip_metrics: + # Prefer to run InsertSizeMetricAlgo after duplicate marking + if self.assay == "WES" and not self.bed: + lc_extra_algos.append(InsertSizeMetricAlgo(paths.insert_size)) + lc_extra_algos.extend(pre_dedup_metrics_algos(paths, self.assay)) + + deduped = pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", f"_deduped.{suffix}") + ) + dedup_result = DedupStage( + ctx=ctx, + inputs=self.inputs, + output=deduped, + score_file=paths.score, + consensus=self.consensus, + rmdup=(self.duplicate_marking == "rmdup"), + dedup_metrics=paths.dedup_metrics, + lc_extra_algos=lc_extra_algos, + ).add_to(dag, upstream) + jobs: List[Job] = list(dedup_result.jobs) + qc_jobs: List[Job] = [dedup_result.lc_job] + + # HsMetricAlgo and WgsMetricsAlgo run after duplicate marking to + # account for duplicate reads + post_algos: List[BaseAlgo] = [] + rehead_metrics: Optional[pathlib.Path] = None + if not self.skip_metrics: + if self.assay == "WES" and self.bed: + post_algos = [ + HsMetricAlgo(paths.hybrid_selection, self.bed, self.bed), + InsertSizeMetricAlgo(paths.insert_size), + ] + elif self.assay == "WGS": + post_algos = [ + InsertSizeMetricAlgo(paths.insert_size), + WgsMetricsAlgo(paths.wgs, include_unpaired="true"), + CoverageMetrics(paths.coverage), + ] + # Rehead WGS metrics so they are recognized by MultiQC + rehead_metrics = paths.wgs + if post_algos: + metrics_result = MetricsStage( + ctx=ctx, + inputs=[deduped], + algos=post_algos, + interval=self.bed, + rehead_metrics=rehead_metrics, + threads=0, # Run metrics in the background + ).add_to(dag, {dedup_result.dedup_job}) + jobs.extend(metrics_result.jobs) + qc_jobs.extend(metrics_result.jobs) + + return SrPreprocessingResult( + jobs=jobs, + terminal={dedup_result.dedup_job}, + deduped=[deduped], + dedup_job=dedup_result.dedup_job, + qc_jobs=qc_jobs, + ) diff --git a/sentieon_cli/stages/segdup.py b/sentieon_cli/stages/segdup.py new file mode 100644 index 0000000..5ecabb2 --- /dev/null +++ b/sentieon_cli/stages/segdup.py @@ -0,0 +1,81 @@ +""" +Variant calling in segmental duplications with segdup-caller +""" + +from dataclasses import dataclass +import pathlib +from typing import Dict, Iterable, Optional, Sequence + +import packaging.version + +from .. import command_strings as cmds +from ..dag import DAG +from ..job import Job +from ..util import SampleSex, caller_sex_arg +from .base import Stage, StageResult + +SEGDUP_MIN_VERSIONS: Dict[str, Optional[packaging.version.Version]] = { + # 0.7.0 adds the `--set` argument + "segdup-caller": packaging.version.Version("0.7.0"), +} + + +@dataclass(kw_only=True) +class SegdupResult(StageResult): + """The job and output of `SegdupStage`""" + + job: Job + output_dir: pathlib.Path + + +@dataclass(kw_only=True) +class SegdupStage(Stage): + """Call variants in difficult segmental duplications. + + segdup-caller is sex-aware, so callers run this stage after ploidy + estimation. Long reads are optional and need both ``lr_alignment`` + and ``lr_bundle``. ``genes`` is a comma-separated gene list; without + one the caller's default gene set is used. ``overrides`` become + ``--set KEY=VALUE`` arguments. + """ + + sr_alignment: pathlib.Path + sr_bundle: pathlib.Path + output_dir: pathlib.Path + input_vcf: Optional[pathlib.Path] = None + sample_sex: Optional[SampleSex] = None + genes: Optional[str] = None + lr_alignment: Optional[pathlib.Path] = None + lr_bundle: Optional[pathlib.Path] = None + overrides: Sequence[str] = () + name: str = "segdup-caller" + task_name: str = "segdup" + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> SegdupResult: + deps = set(upstream) + + job = Job( + cmds.cmd_segdup_caller( + self.output_dir, + self.sr_alignment, + reference=self.ctx.reference, + sr_bundle=self.sr_bundle, + input_vcf=self.input_vcf, + sex=caller_sex_arg(self.sample_sex), + genes=self.genes, + overrides=list(self.overrides), + lr_alignments=self.lr_alignment, + lr_bundle=self.lr_bundle, + ), + self.name, + self.ctx.cores, + task_name=self.task_name, + ) + dag.add_job(job, deps) + + return SegdupResult( + jobs=[job], + terminal={job}, + job=job, + output_dir=self.output_dir, + ) diff --git a/sentieon_cli/stages/small_variants.py b/sentieon_cli/stages/small_variants.py new file mode 100644 index 0000000..4880b55 --- /dev/null +++ b/sentieon_cli/stages/small_variants.py @@ -0,0 +1,235 @@ +""" +Small-variant calling, annotation transfer, and genotyping +""" + +from dataclasses import dataclass +import pathlib +from typing import Iterable, List, Optional, Set, Union + +from ..dag import DAG +from ..driver import BaseAlgo, DNAModelApply, GVCFtyper +from ..job import Job +from .base import Stage, StageResult, driver_job +from .transfer import TransferConfig, TransferResult, TransferStage + + +@dataclass(kw_only=True) +class DriverResult(StageResult): + """The job of a single-job `sentieon driver` stage""" + + job: Job + + +@dataclass(kw_only=True) +class DriverStage(Stage): + """A single `sentieon driver` invocation, wired into the DAG. + + The driver arguments are passed through untouched; a stage that needs + more structure than "run these algos" subclasses this one and fixes + the defaults it cares about. + """ + + algos: List[BaseAlgo] + name: str + task_name: str + threads: Optional[int] = None + inputs: Optional[List[pathlib.Path]] = None + interval: Optional[Union[pathlib.Path, str]] = None + interval_padding: Optional[int] = None + read_filter: Optional[List[str]] = None + replace_rg: Optional[List[List[str]]] = None + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> DriverResult: + job = driver_job( + self.ctx, + self.algos, + name=self.name, + task_name=self.task_name, + threads=self.threads, + inputs=self.inputs, + interval=self.interval, + interval_padding=self.interval_padding, + read_filter=self.read_filter, + replace_rg=self.replace_rg, + ) + dag.add_job(job, set(upstream)) + return DriverResult(jobs=[job], terminal={job}, job=job) + + +@dataclass(kw_only=True) +class DNAscopeStage(DriverStage): + """The DNAscope-family calling pass. + + The caller supplies the algos in the order the driver should see + them -- SNV/indel calling first, then a BND or PangenomeSV pass in + the same pileup. + """ + + name: str = "dnascope" + task_name: str = "variant-calling" + + +@dataclass(kw_only=True) +class ModelApplyResult(StageResult): + """The job and output of `ModelApplyStage`""" + + job: Job + output: pathlib.Path + + +@dataclass(kw_only=True) +class ModelApplyStage(Stage): + """Genotype and filter a raw VCF with DNAModelApply""" + + model: pathlib.Path + vcf: pathlib.Path + output: pathlib.Path + name: str = "model-apply" + task_name: str = "model-apply" + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> ModelApplyResult: + job = driver_job( + self.ctx, + [DNAModelApply(self.model, self.vcf, self.output)], + name=self.name, + task_name=self.task_name, + ) + dag.add_job(job, set(upstream)) + return ModelApplyResult( + jobs=[job], + terminal={job}, + job=job, + output=self.output, + ) + + +@dataclass(kw_only=True) +class GVCFtyperResult(StageResult): + """The job and output of `GVCFtyperStage`""" + + job: Job + output: pathlib.Path + + +@dataclass(kw_only=True) +class GVCFtyperStage(Stage): + """Genotype a gVCF to produce a VCF""" + + gvcf: pathlib.Path + output: pathlib.Path + interval: Optional[pathlib.Path] = None + name: str = "gvcftyper" + task_name: str = "gvcftyper" + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> GVCFtyperResult: + job = driver_job( + self.ctx, + [GVCFtyper(output=self.output, vcf=self.gvcf)], + interval=self.interval, + name=self.name, + task_name=self.task_name, + ) + dag.add_job(job, set(upstream)) + return GVCFtyperResult( + jobs=[job], + terminal={job}, + job=job, + output=self.output, + ) + + +@dataclass(frozen=True) +class TransferSpec: + """Ask `TransferApplyStage` for an annotation transfer""" + + config: TransferConfig + out_vcf: pathlib.Path + tag: Optional[str] = None + + +@dataclass(frozen=True) +class ApplySpec: + """Ask `TransferApplyStage` for a DNAModelApply pass""" + + model: pathlib.Path + output: pathlib.Path + name: str = "model-apply" + + +@dataclass(kw_only=True) +class TransferApplyResult(StageResult): + """The jobs and output of `TransferApplyStage`""" + + transfer: Optional[TransferResult] + apply_job: Optional[Job] + output: pathlib.Path + + +@dataclass(kw_only=True) +class TransferApplyStage(Stage): + """Post-process a raw VCF: transfer annotations, then apply the model. + + Either half may be omitted. When both run, DNAModelApply reads the + transfer output and waits on the transfer alone -- the transfer's + concat job already depends on the raw-VCF producer. Without a + transfer, DNAModelApply waits on ``upstream``. + """ + + raw_vcf: pathlib.Path + transfer: Optional[TransferSpec] = None + apply: Optional[ApplySpec] = None + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> TransferApplyResult: + deps = set(upstream) + if self.transfer is None and self.apply is None: + raise ValueError( + "TransferApplyStage needs a transfer, an apply, or both" + ) + + jobs: List[Job] = [] + transfer_result: Optional[TransferResult] = None + apply_in = self.raw_vcf + apply_deps = deps + if self.transfer is not None: + transfer_result = TransferStage( + ctx=self.ctx, + config=self.transfer.config, + raw_vcf=self.raw_vcf, + out_vcf=self.transfer.out_vcf, + tag=self.transfer.tag, + ).add_to(dag, deps) + jobs.extend(transfer_result.jobs) + apply_in = self.transfer.out_vcf + apply_deps = transfer_result.terminal + + apply_job: Optional[Job] = None + if self.apply is not None: + apply_result = ModelApplyStage( + ctx=self.ctx, + model=self.apply.model, + vcf=apply_in, + output=self.apply.output, + name=self.apply.name, + ).add_to(dag, apply_deps) + apply_job = apply_result.job + jobs.append(apply_job) + terminal: Set[Job] = {apply_job} + output = self.apply.output + else: + assert transfer_result is not None + terminal = transfer_result.terminal + output = transfer_result.out_vcf + + return TransferApplyResult( + jobs=jobs, + terminal=terminal, + transfer=transfer_result, + apply_job=apply_job, + output=output, + ) diff --git a/sentieon_cli/stages/sv.py b/sentieon_cli/stages/sv.py new file mode 100644 index 0000000..c722172 --- /dev/null +++ b/sentieon_cli/stages/sv.py @@ -0,0 +1,88 @@ +""" +Structural variant calling with LongReadSV +""" + +from dataclasses import dataclass +import pathlib +from typing import Iterable, List, Optional + +import packaging.version + +from ..dag import DAG +from ..driver import LongReadSV +from ..job import Job +from .base import Stage, StageResult, driver_job + +LONGREADSV_MIN_VERSIONS = { + "sentieon driver": packaging.version.Version("202308"), +} + + +@dataclass(kw_only=True) +class LongReadSVResult(StageResult): + """The job and output of `LongReadSVStage`""" + + job: Job + sv_vcf: pathlib.Path + + +@dataclass(kw_only=True) +class LongReadSVStage(Stage): + """Call SVs from long reads with Sentieon LongReadSV. + + ``output`` names the SV VCF; when it is None the path is derived from + the run's output VCF, so the caller only supplies the alignment and + the model. The remaining fields are LongReadSV algo options, emitted + only when set. + """ + + inputs: List[pathlib.Path] + model: pathlib.Path + output: Optional[pathlib.Path] = None + min_map_qual: Optional[int] = None + min_sv_size: Optional[int] = None + min_dp: Optional[int] = None + min_af: Optional[float] = None + interval: Optional[pathlib.Path] = None + replace_rg: Optional[List[List[str]]] = None + name: str = "LongReadSV" + task_name: str = "sv-calling" + + def build(self) -> LongReadSVResult: + """Construct this stage's job without touching a DAG""" + sv_vcf = self.output + if sv_vcf is None: + sv_vcf = pathlib.Path( + str(self.ctx.output_vcf).replace(".vcf.gz", ".sv.vcf.gz") + ) + job = driver_job( + self.ctx, + [ + LongReadSV( + sv_vcf, + model=self.model, + min_map_qual=self.min_map_qual, + min_sv_size=self.min_sv_size, + min_dp=self.min_dp, + min_af=self.min_af, + ) + ], + inputs=self.inputs, + interval=self.interval, + replace_rg=self.replace_rg, + name=self.name, + task_name=self.task_name, + ) + return LongReadSVResult( + jobs=[job], + terminal={job}, + job=job, + sv_vcf=sv_vcf, + ) + + def add_to( + self, dag: DAG, upstream: Iterable[Job] = () + ) -> LongReadSVResult: + result = self.build() + dag.add_job(result.job, set(upstream)) + return result diff --git a/sentieon_cli/stages/t1k.py b/sentieon_cli/stages/t1k.py new file mode 100644 index 0000000..6e2b264 --- /dev/null +++ b/sentieon_cli/stages/t1k.py @@ -0,0 +1,94 @@ +""" +HLA and KIR genotyping with T1K +""" + +from dataclasses import dataclass +import pathlib +from typing import Dict, Iterable, List, Optional + +import packaging.version + +from .. import command_strings as cmds +from ..dag import DAG +from ..job import Job +from .alignment import ReadWriterStage +from .base import Stage, StageResult + +T1K_MIN_VERSIONS: Dict[str, Optional[packaging.version.Version]] = { + "run-t1k": None, +} + +DEFAULT_T1K_HLA_LOCUS = "chr6:28510020-33480577" +DEFAULT_T1K_KIR_LOCUS = "chr19:53100000-55800000" + + +@dataclass(kw_only=True) +class T1KResult(StageResult): + """The jobs and outputs of `T1KStage`""" + + extract_job: Job + t1k_job: Job + extracted_bam: pathlib.Path + output_dir: pathlib.Path + + +@dataclass(kw_only=True) +class T1KStage(Stage): + """Extract the reads at a gene locus and genotype them with T1K. + + ``tag`` names the gene group ("hla" or "kir"); it labels the jobs and + the extracted BAM. ``preset`` is T1K's own `--preset` ("hla-wgs" or + "kir-wgs") and ``locus`` the region the reads are extracted from. + """ + + inputs: List[pathlib.Path] + gene_seq: pathlib.Path + gene_coord: pathlib.Path + locus: str + preset: str + tag: str + output_dir: pathlib.Path + task_name: str = "t1k" + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> T1KResult: + deps = set(upstream) + + # Extract reads overlapping the locus to a BAM file with ReadWriter + extracted_bam = self.ctx.tmp_dir.joinpath(f"sample_{self.tag}.bam") + extract_job = ( + ReadWriterStage( + ctx=self.ctx, + inputs=self.inputs, + output=extracted_bam, + name=f"t1k-{self.tag}-extract", + task_name=self.task_name, + interval=self.locus, + ) + .build() + .jobs[0] + ) + dag.add_job(extract_job, deps) + + t1k_job = Job( + cmds.cmd_t1k( + self.output_dir, + extracted_bam, + gene_seq=self.gene_seq, + gene_coord=self.gene_coord, + preset=self.preset, + threads=self.ctx.cores, + ), + f"t1k-{self.tag}", + self.ctx.cores, + task_name=self.task_name, + ) + dag.add_job(t1k_job, {extract_job}) + + return T1KResult( + jobs=[extract_job, t1k_job], + terminal={t1k_job}, + extract_job=extract_job, + t1k_job=t1k_job, + extracted_bam=extracted_bam, + output_dir=self.output_dir, + ) diff --git a/sentieon_cli/transfer.py b/sentieon_cli/stages/transfer.py similarity index 50% rename from sentieon_cli/transfer.py rename to sentieon_cli/stages/transfer.py index a5fed93..b6b9fff 100644 --- a/sentieon_cli/transfer.py +++ b/sentieon_cli/stages/transfer.py @@ -1,19 +1,22 @@ """ -Annotation transfer functionality +Annotation transfer from a population VCF """ +from dataclasses import dataclass import pathlib import re import subprocess as sp import tempfile -from typing import Dict, List, Optional, Set, Tuple +from typing import Dict, Iterable, List, Optional, Protocol, Set, Tuple from importlib.resources import files -from . import command_strings as cmds -from .job import Job -from .logging import get_logger -from .shard import Shard +from .. import command_strings as cmds +from ..dag import DAG +from ..job import Job +from ..logging import get_logger +from ..shard import Shard +from .base import Stage, StageResult logger = get_logger(__name__) @@ -28,8 +31,17 @@ def build_transfer_jobs( fai_data: Dict[str, Dict[str, int]], dry_run: bool = False, cores: int = 1, + *, + tag: Optional[str] = None, ) -> Tuple[List[Job], Job]: - """Transfer annotations from the pop_vcf to the raw_vcf""" + """Transfer annotations from the pop_vcf to the raw_vcf. + + ``tag`` is inserted into the job names (``merge-trim-{tag}-...``) so a + pipeline that transfers annotations more than once in a run can tell + the two fan-outs apart. + """ + + name_base = "merge-trim" if tag is None else f"merge-trim-{tag}" # Get a unique tmpdir tmp_dir_str = tempfile.mkdtemp(dir=base_tmp_dir) @@ -87,8 +99,9 @@ def build_transfer_jobs( raw_vcf, regions_file=subset_bed, ), - "merge-trim-extra", + f"{name_base}-extra", 1, + task_name="annotation-transfer", ) sharded_merge_jobs.append(view_job) sharded_vcfs.append(subset_vcf) @@ -121,8 +134,9 @@ def build_transfer_jobs( ], view_xargs=["--no-version"], ), - f"merge-trim-{i}", + f"{name_base}-{i}", 1, + task_name="annotation-transfer", ) sharded_merge_jobs.append(merge_job) sharded_vcfs.append(shard_vcf) @@ -134,7 +148,107 @@ def build_transfer_jobs( sharded_vcfs, xargs=["--no-version", "--threads", str(cores)], ), - "merge-trim-concat", + f"{name_base}-concat", cores, + task_name="annotation-transfer", ) return (sharded_merge_jobs, concat_job) + + +class TransferInputs(Protocol): + """A pipeline that carries the annotation transfer's inputs. + + Every pipeline that transfers annotations holds these four values under + these names, so `TransferConfig.from_pipeline` can collect them + structurally instead of each pipeline repeating the same helper. + """ + + pop_vcf: Optional[pathlib.Path] + shards: List[Shard] + pop_vcf_contigs: Dict[str, Optional[int]] + fai_data: Dict[str, Dict[str, int]] + + +@dataclass(frozen=True) +class TransferConfig: + """The run-wide inputs the annotation transfer needs. + + Every pipeline that transfers annotations builds the same fan-out from + the same four values, so they travel together instead of being threaded + through each call site. + """ + + pop_vcf: pathlib.Path + shards: List[Shard] + pop_vcf_contigs: Dict[str, Optional[int]] + fai_data: Dict[str, Dict[str, int]] + + @classmethod + def from_pipeline(cls, pipeline: TransferInputs) -> "TransferConfig": + """Collect the transfer's inputs from a pipeline. + + The four values are set together while the pipeline validates its + arguments, so they are read back together here. Callers check + `pop_vcf` before asking for a transfer at all. + """ + pop_vcf = pipeline.pop_vcf + if pop_vcf is None: + raise ValueError("The annotation transfer needs a population VCF") + return cls( + pop_vcf=pop_vcf, + shards=pipeline.shards, + pop_vcf_contigs=pipeline.pop_vcf_contigs, + fai_data=pipeline.fai_data, + ) + + +@dataclass(kw_only=True) +class TransferResult(StageResult): + """The jobs and output of `TransferStage`""" + + shard_jobs: List[Job] + concat_job: Job + out_vcf: pathlib.Path + + +@dataclass(kw_only=True) +class TransferStage(Stage): + """Transfer annotations from the population VCF onto a raw VCF. + + One job per genomic shard merges the shard's annotations, then a + single concat job assembles ``out_vcf``. ``tag`` disambiguates the + job names when a pipeline transfers more than once in a run. + """ + + config: TransferConfig + raw_vcf: pathlib.Path + out_vcf: pathlib.Path + tag: Optional[str] = None + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> TransferResult: + deps = set(upstream) + + shard_jobs, concat_job = build_transfer_jobs( + self.out_vcf, + self.config.pop_vcf, + self.raw_vcf, + self.ctx.tmp_dir, + self.config.shards, + self.config.pop_vcf_contigs, + self.config.fai_data, + self.ctx.dry_run, + self.ctx.cores, + tag=self.tag, + ) + + for job in shard_jobs: + dag.add_job(job, deps) + dag.add_job(concat_job, set(shard_jobs)) + + return TransferResult( + jobs=[*shard_jobs, concat_job], + terminal={concat_job}, + shard_jobs=shard_jobs, + concat_job=concat_job, + out_vcf=self.out_vcf, + ) diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index 68021ea..7d79905 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -3,16 +3,27 @@ """ import argparse +from enum import Enum +import gzip from importlib.metadata import PackageNotFoundError, version import multiprocessing as mp import os import pathlib import re -import shlex import shutil import subprocess as sp +import sys import tempfile -from typing import Callable, Dict, List, Optional +from typing import ( + IO, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Tuple, +) import packaging.version @@ -29,7 +40,67 @@ PRELOAD_SEP_PAT = re.compile(PRELOAD_SEP) NUMA_NODE_PAT = re.compile(r"^NUMA node. CPU\(s\):\s+(?P.*)$") -READ_LENGTH_PAT = re.compile(r"SN\taverage length:\t(?P\d*)$") + +_UNSAFE = re.compile(r"[^A-Za-z0-9._-]") + +# The first two bytes of a gzip (and so of a BGZF) file +GZIP_MAGIC = b"\x1f\x8b" + + +class SampleSex(Enum): + FEMALE = 1 + MALE = 2 + UNKNOWN = 3 + + +def sample_sex_arg(value: str) -> SampleSex: + """Parse the `--sample_sex` argument""" + sex = value.strip().lower() + if sex == "male": + return SampleSex.MALE + if sex == "female": + return SampleSex.FEMALE + raise argparse.ArgumentTypeError( + f"invalid sample sex '{value}'. Please supply 'male' or 'female'" + ) + + +def cnvscope_sex_args( + sample_sex: Optional[SampleSex], + par_bed: Optional[pathlib.Path], +) -> Tuple[Optional[str], Optional[pathlib.Path]]: + """The CNVscope `--sex` and `--par` arguments for a sample. + + Male samples are called with the pseudo-autosomal regions (PAR) BED + file. Female samples do not need a PAR BED file. When the sample sex + is not known, both arguments are omitted and CNVscope assumes a + diploid genome, matching the behavior of previous releases. + """ + if sample_sex is SampleSex.MALE: + return ("M", par_bed) + if sample_sex is SampleSex.FEMALE: + return ("F", None) + logger.warning( + "The sample sex is not known. CNVscope will assume a diploid " + "genome, matching the behavior of previous releases. Supply " + "`--sample_sex` for sex-aware CNV calling." + ) + return (None, None) + + +def caller_sex_arg(sample_sex: Optional[SampleSex]) -> str: + """The `--sex male|female` value for ExpansionHunter and segdup-caller. + + Both tools accept only the two values, so an unknown sample sex + collapses to "female", unchanged from the previous + dnascope-pangenome behavior. + """ + return "male" if sample_sex is SampleSex.MALE else "female" + + +def sanitize(component: str) -> str: + """Restrict a path component to filesystem-safe characters""" + return _UNSAFE.sub("-", component) def tmp(): @@ -39,20 +110,13 @@ def tmp(): return tmp_dir -def check_version( - cmd: str, - version: Optional[packaging.version.Version], -) -> bool: - """Check the version of an executable""" - cmd_list: List[str] = cmd.split() - exec_file = shutil.which(cmd_list[0]) - if not exec_file: - logger.error("Error: no '%s' found in the PATH", cmd) - return False - - if version is None: - return True +def executable_version(cmd: str) -> Optional[packaging.version.Version]: + """The version reported by ` --version`. + Returns `None`, after logging the reason, when the command cannot be + run or its output is not a version. + """ + cmd_list: List[str] = cmd.split() cmd_list.append("--version") try: cmd_version_str = ( @@ -62,7 +126,7 @@ def check_version( logger.error( "Error: could not determine the version of '%s': %s", cmd, e ) - return False + return None if cmd_list[0] == "sentieon": cmd_version_str = cmd_version_str.split("-")[-1] elif cmd_list[0] == "pbsv": @@ -74,7 +138,34 @@ def check_version( cmd_version_str = ( cmd_version_str.split("\n")[0].split()[-1].split("-")[0] ) - cmd_version = packaging.version.Version(cmd_version_str) + try: + return packaging.version.Version(cmd_version_str) + except packaging.version.InvalidVersion: + logger.error( + "Error: could not parse the version of '%s': '%s'", + cmd, + cmd_version_str, + ) + return None + + +def check_version( + cmd: str, + version: Optional[packaging.version.Version], +) -> bool: + """Check the version of an executable""" + cmd_list: List[str] = cmd.split() + exec_file = shutil.which(cmd_list[0]) + if not exec_file: + logger.error("Error: no '%s' found in the PATH", cmd) + return False + + if version is None: + return True + + cmd_version = executable_version(cmd) + if cmd_version is None: + return False if cmd_version < version: logger.error( "Error: the pipeline requires %s version '%s' or later " @@ -88,6 +179,45 @@ def check_version( return True +def require_versions( + min_versions: Mapping[str, Optional[packaging.version.Version]], + *, + skip: bool = False, +) -> None: + """Exit unless every executable meets its minimum version. + + `check_version` has already logged the reason, so the exit is silent. + Pass `skip=True` (the pipelines' `--skip_version_check`) to do nothing. + + A `Mapping` rather than a `Dict`, so the module-level + `*_MIN_VERSIONS` constants -- some of which mypy infers as + `Dict[str, Version]`, with no `None` entry -- are accepted. + """ + if skip: + return + for cmd, min_version in min_versions.items(): + if not check_version(cmd, min_version): + sys.exit(2) + + +def versions_available( + min_versions: Mapping[str, Optional[packaging.version.Version]], + *, + skip: bool = False, +) -> bool: + """Whether every executable meets its minimum version. + + For optional tools, where a missing or outdated executable skips a step + rather than ending the run. `skip=True` reports them as available. + """ + if skip: + return True + return all( + check_version(cmd, min_version) + for cmd, min_version in min_versions.items() + ) + + def path_arg( exists: Optional[bool] = None, is_dir: Optional[bool] = None, @@ -148,6 +278,36 @@ def total_memory() -> int: return total_mem +def set_bwt_max_mem( + total_input_size: int, + n_alignment_jobs: int = 1, + override: Optional[str] = None, +) -> str: + """Set the `bwt_max_mem` environment variable and return its value. + + `override` short-circuits the calculation, for the hidden + `--bwt_max_mem` argument. Otherwise the value is derived from the + memory available to the run, the size of the inputs staged in + memory, and the number of concurrent alignment jobs. + """ + if override: + os.environ["bwt_max_mem"] = override + return override + + total_mem = total_memory() + total_mem_gb = total_mem / (1024.0**3) + align_mem_gb = ( + total_mem_gb - 4 - total_input_size / (1024.0**3) * 2.3 + ) # some memory for other system processes + bwa_mem_gb = max( + int((align_mem_gb / n_alignment_jobs) - 6), 0 + ) # some memory for other alignment processes + logger.debug("Setting bwt_max_mem to: %sG", bwa_mem_gb) + bwt_max_mem = f"{bwa_mem_gb}G" + os.environ["bwt_max_mem"] = bwt_max_mem + return bwt_max_mem + + def find_numa_nodes() -> List[str]: """Find NUMA nodes on the system, if available""" numa_nodes = [] @@ -213,67 +373,77 @@ def parse_rg_line(rg_line: str) -> Dict[str, str]: return parsed -def get_read_length_aln( - aln: pathlib.Path, - reference: pathlib.Path, - n_reads: int = 100000, -) -> int: - """Get the average read length for an alignment file""" - cmds = [] - cmds.append( - [ - "samtools", - "view", - "-h", - "--reference", - shlex.quote(str(reference)), - shlex.quote(str(aln)), - ] - ) - cmds.append( - [ - "head", - "-n", - str(n_reads), - ] - ) - cmds.append(["samtools", "stats", "-"]) - all_cmds = [shlex.join(x) for x in cmds] - cmd = " | ".join(all_cmds) - res = sp.run( - cmd, - shell=True, - capture_output=True, - text=True, - executable="/bin/bash", - ) +class VcfHeaderError(ValueError): + """A VCF header could not be read""" - for line in res.stdout.split("\n"): - m = READ_LENGTH_PAT.match(line) - if m: - return int(m.groupdict()["length"]) - return 151 +def _open_vcf_text(in_vcf: pathlib.Path) -> IO[str]: + """Open a VCF for reading, transparently decompressing it. -def vcf_id(in_vcf: pathlib.Path) -> Optional[str]: - """Collect the SentieonVcfID header""" - cmd = ["bcftools", "view", "-h", str(in_vcf)] - p = sp.run(cmd, capture_output=True, text=True) - if p.returncode != 0: - logger.error( - "`%s` failed with return code %d: %s", - " ".join(cmd), - p.returncode, - p.stderr.strip(), - ) - return None - for line in p.stdout.split("\n"): + The first two bytes decide: a BGZF file is a gzip file, and a VCF is + routinely named without regard to how it is compressed. + """ + with open(in_vcf, "rb") as raw: + compressed = raw.read(len(GZIP_MAGIC)) == GZIP_MAGIC + if compressed: + return gzip.open(in_vcf, "rt") + return open(in_vcf, "rt") + + +def read_vcf_header(in_vcf: pathlib.Path) -> List[str]: + """Read the header lines of a VCF file. + + Lines are read until the first one that is not a header line, so the + body of the file is never touched: the header of a whole-genome + population VCF parses in milliseconds, with no `bcftools` on the + PATH and no subprocess. + + Raises: + VcfHeaderError: the file holds no VCF header, or its header + could not be decompressed or decoded. + OSError: the file does not exist or could not be read. + """ + header: List[str] = [] + try: + with _open_vcf_text(in_vcf) as fh: + for line in fh: + if not line.startswith("#"): + break + header.append(line.rstrip("\n")) + if line.startswith("#CHROM"): + break + except (gzip.BadGzipFile, EOFError, UnicodeDecodeError) as err: + raise VcfHeaderError( + f"{in_vcf}: the VCF header could not be read ({err})" + ) from err + if not header: + raise VcfHeaderError(f"{in_vcf}: the file has no VCF header") + return header + + +def vcf_id_from_header(header: Iterable[str]) -> Optional[str]: + """The `##SentieonVcfID` of VCF header lines, `None` when absent""" + for line in header: if line.startswith("##SentieonVcfID="): i = line.index("=") return line[i + 1 :] # noqa: E203 return None +def vcf_id(in_vcf: pathlib.Path) -> Optional[str]: + """Collect the SentieonVcfID header. + + Returns `None` when the header has no ID and when it could not be + read at all; an unreadable header is logged. + """ + try: + header = read_vcf_header(in_vcf) + except (VcfHeaderError, OSError) as err: + logger.error("%s", err) + return None + return vcf_id_from_header(header) + + def check_kmc_patch(kmc_cmd: str = "kmc") -> bool: """Check if the KMC version supports the required patch""" # Create a temporary directory for the test diff --git a/tests/data/pangenome/generate.sh b/tests/data/pangenome/generate.sh new file mode 100755 index 0000000..1d51556 --- /dev/null +++ b/tests/data/pangenome/generate.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# Regenerate the binary pangenome fixtures in this directory. +# +# The `.gfa` files are the checked-in sources; `vg` (tested with v1.72.0) +# turns them into the GBZ graphs and the `vg haplotypes` index that +# `tests/unit/test_pangenome_meta.py` parses. Everything here is a toy +# graph of 8 segments and two 30 bp contigs, so the outputs stay a few KB. +# +# The graphs differ only in their walks and in the `RS:Z:` header tag, +# which `vg` turns into the GBWT `reference_samples` tag: +# +# tiny.gfa GRCh38 is the unfragmented backbone, CHM13 is +# split into two fragments on chr1 +# tiny_chm13.gfa the same graph with the two roles swapped +# tiny_no_backbone.gfa both reference samples are fragmented +# tiny_no_ref_tag.gfa no `RS:Z:` tag, so no reference sample at all +# +# `tiny_sampled.gbz` stands in for the graph `vg haplotypes +# --set-reference GRCh38` writes: one reference sample, as +# `check_sample_gbz` requires. +# +# The `rgfa_*.gfa` files are hand-written heads of a `vg convert -f -Q` +# rGFA and need no regeneration. +set -euo pipefail +cd "$(dirname "$0")" + +for name in tiny tiny_chm13 tiny_no_backbone tiny_no_ref_tag; do + vg gbwt -G "${name}.gfa" --gbz-format -g "${name}.gbz" +done + +vg gbwt -Z tiny.gbz --set-reference GRCh38 -g tiny_sampled.gbz + +vg index -j tiny.dist tiny.gbz +vg gbwt -Z tiny.gbz -r tiny.ri +vg haplotypes -d tiny.dist -r tiny.ri -H tiny.hapl tiny.gbz +rm -f tiny.dist tiny.ri + +# What the fixtures should contain +vg gbwt -Z tiny.gbz --tags +vg paths -x tiny.gbz -L diff --git a/tests/data/pangenome/rgfa_no_header.gfa b/tests/data/pangenome/rgfa_no_header.gfa new file mode 100644 index 0000000..afa1dce --- /dev/null +++ b/tests/data/pangenome/rgfa_no_header.gfa @@ -0,0 +1,3 @@ +S 1 ACGTACGTAC SN:Z:GRCh38#0#chr1 SO:i:0 SR:i:0 +S 2 GGGGGGGGGG SN:Z:GRCh38#0#chr1 SO:i:10 SR:i:0 +L 1 + 2 + 0M diff --git a/tests/data/pangenome/rgfa_no_tags.gfa b/tests/data/pangenome/rgfa_no_tags.gfa new file mode 100644 index 0000000..c4dcc80 --- /dev/null +++ b/tests/data/pangenome/rgfa_no_tags.gfa @@ -0,0 +1,4 @@ +H VN:Z:1.1 +S 1 ACGTACGTAC +S 2 GGGGGGGGGG +L 1 + 2 + 0M diff --git a/tests/data/pangenome/rgfa_ok.gfa b/tests/data/pangenome/rgfa_ok.gfa new file mode 100644 index 0000000..c6a8256 --- /dev/null +++ b/tests/data/pangenome/rgfa_ok.gfa @@ -0,0 +1,6 @@ +H VN:Z:1.1 RS:Z:CHM13 +S 1 ACGTACGTAC SN:Z:GRCh38#0#chr1 SO:i:0 SR:i:0 +S 2 GGGGGGGGGG SN:Z:GRCh38#0#chr1 SO:i:10 SR:i:0 +S 3 TTTTTTTTTT SR:i:1 +L 1 + 2 + 0M +L 1 + 3 + 0M diff --git a/tests/data/pangenome/rgfa_wrong_reference.gfa b/tests/data/pangenome/rgfa_wrong_reference.gfa new file mode 100644 index 0000000..46a1708 --- /dev/null +++ b/tests/data/pangenome/rgfa_wrong_reference.gfa @@ -0,0 +1,6 @@ +H VN:Z:1.1 RS:Z:CHM13 GRCh38 +S 1 ACGTACGTAC +S 2 GGGGGGGGGG +S 3 TTTTTTTTTT +L 1 + 2 + 0M +L 1 + 3 + 0M diff --git a/tests/data/pangenome/tiny.gbz b/tests/data/pangenome/tiny.gbz new file mode 100644 index 0000000..d24bff4 Binary files /dev/null and b/tests/data/pangenome/tiny.gbz differ diff --git a/tests/data/pangenome/tiny.gfa b/tests/data/pangenome/tiny.gfa new file mode 100644 index 0000000..519c6e5 --- /dev/null +++ b/tests/data/pangenome/tiny.gfa @@ -0,0 +1,28 @@ +H VN:Z:1.1 RS:Z:GRCh38 CHM13 +S s1 ACGTACGTAC +S s2 GGGGGGGGGG +S s3 TTTTTTTTTT +S s4 CCCCCCCCCC +S s5 AAAAAAAAAA +S s6 ACACACACAC +S s7 TGTGTGTGTG +S s8 GATTACAGAT +L s1 + s2 + 0M +L s1 + s3 + 0M +L s2 + s4 + 0M +L s3 + s4 + 0M +L s5 + s6 + 0M +L s5 + s7 + 0M +L s6 + s8 + 0M +L s7 + s8 + 0M +W GRCh38 0 chr1 0 30 >s1>s2>s4 +W GRCh38 0 chr2 0 30 >s5>s6>s8 +W CHM13 0 chr1 0 20 >s1>s3 +W CHM13 0 chr1 20 30 >s4 +W CHM13 0 chr2 0 30 >s5>s7>s8 +W HG01 1 chr1 0 30 >s1>s3>s4 +W HG01 2 chr1 0 30 >s1>s2>s4 +W HG01 1 chr2 0 30 >s5>s6>s8 +W HG01 2 chr2 0 30 >s5>s7>s8 +W HG02 1 chr1 0 30 >s1>s2>s4 +W HG02 2 chr2 0 30 >s5>s6>s8 diff --git a/tests/data/pangenome/tiny.hapl b/tests/data/pangenome/tiny.hapl new file mode 100644 index 0000000..7bf79bb Binary files /dev/null and b/tests/data/pangenome/tiny.hapl differ diff --git a/tests/data/pangenome/tiny_chm13.gbz b/tests/data/pangenome/tiny_chm13.gbz new file mode 100644 index 0000000..f7a304c Binary files /dev/null and b/tests/data/pangenome/tiny_chm13.gbz differ diff --git a/tests/data/pangenome/tiny_chm13.gfa b/tests/data/pangenome/tiny_chm13.gfa new file mode 100644 index 0000000..40451f7 --- /dev/null +++ b/tests/data/pangenome/tiny_chm13.gfa @@ -0,0 +1,28 @@ +H VN:Z:1.1 RS:Z:CHM13 GRCh38 +S s1 ACGTACGTAC +S s2 GGGGGGGGGG +S s3 TTTTTTTTTT +S s4 CCCCCCCCCC +S s5 AAAAAAAAAA +S s6 ACACACACAC +S s7 TGTGTGTGTG +S s8 GATTACAGAT +L s1 + s2 + 0M +L s1 + s3 + 0M +L s2 + s4 + 0M +L s3 + s4 + 0M +L s5 + s6 + 0M +L s5 + s7 + 0M +L s6 + s8 + 0M +L s7 + s8 + 0M +W CHM13 0 chr1 0 30 >s1>s2>s4 +W CHM13 0 chr2 0 30 >s5>s6>s8 +W GRCh38 0 chr1 0 20 >s1>s3 +W GRCh38 0 chr1 20 30 >s4 +W GRCh38 0 chr2 0 30 >s5>s7>s8 +W HG01 1 chr1 0 30 >s1>s3>s4 +W HG01 2 chr1 0 30 >s1>s2>s4 +W HG01 1 chr2 0 30 >s5>s6>s8 +W HG01 2 chr2 0 30 >s5>s7>s8 +W HG02 1 chr1 0 30 >s1>s2>s4 +W HG02 2 chr2 0 30 >s5>s6>s8 diff --git a/tests/data/pangenome/tiny_no_backbone.gbz b/tests/data/pangenome/tiny_no_backbone.gbz new file mode 100644 index 0000000..453f460 Binary files /dev/null and b/tests/data/pangenome/tiny_no_backbone.gbz differ diff --git a/tests/data/pangenome/tiny_no_backbone.gfa b/tests/data/pangenome/tiny_no_backbone.gfa new file mode 100644 index 0000000..753d81d --- /dev/null +++ b/tests/data/pangenome/tiny_no_backbone.gfa @@ -0,0 +1,30 @@ +H VN:Z:1.1 RS:Z:GRCh38 CHM13 +S s1 ACGTACGTAC +S s2 GGGGGGGGGG +S s3 TTTTTTTTTT +S s4 CCCCCCCCCC +S s5 AAAAAAAAAA +S s6 ACACACACAC +S s7 TGTGTGTGTG +S s8 GATTACAGAT +L s1 + s2 + 0M +L s1 + s3 + 0M +L s2 + s4 + 0M +L s3 + s4 + 0M +L s5 + s6 + 0M +L s5 + s7 + 0M +L s6 + s8 + 0M +L s7 + s8 + 0M +W GRCh38 0 chr1 0 20 >s1>s2 +W GRCh38 0 chr1 20 30 >s4 +W GRCh38 0 chr2 0 20 >s5>s6 +W GRCh38 0 chr2 20 30 >s8 +W CHM13 0 chr1 0 20 >s1>s3 +W CHM13 0 chr1 20 30 >s4 +W CHM13 0 chr2 0 30 >s5>s7>s8 +W HG01 1 chr1 0 30 >s1>s3>s4 +W HG01 2 chr1 0 30 >s1>s2>s4 +W HG01 1 chr2 0 30 >s5>s6>s8 +W HG01 2 chr2 0 30 >s5>s7>s8 +W HG02 1 chr1 0 30 >s1>s2>s4 +W HG02 2 chr2 0 30 >s5>s6>s8 diff --git a/tests/data/pangenome/tiny_no_ref_tag.gbz b/tests/data/pangenome/tiny_no_ref_tag.gbz new file mode 100644 index 0000000..a19ffce Binary files /dev/null and b/tests/data/pangenome/tiny_no_ref_tag.gbz differ diff --git a/tests/data/pangenome/tiny_no_ref_tag.gfa b/tests/data/pangenome/tiny_no_ref_tag.gfa new file mode 100644 index 0000000..ceacf56 --- /dev/null +++ b/tests/data/pangenome/tiny_no_ref_tag.gfa @@ -0,0 +1,28 @@ +H VN:Z:1.1 +S s1 ACGTACGTAC +S s2 GGGGGGGGGG +S s3 TTTTTTTTTT +S s4 CCCCCCCCCC +S s5 AAAAAAAAAA +S s6 ACACACACAC +S s7 TGTGTGTGTG +S s8 GATTACAGAT +L s1 + s2 + 0M +L s1 + s3 + 0M +L s2 + s4 + 0M +L s3 + s4 + 0M +L s5 + s6 + 0M +L s5 + s7 + 0M +L s6 + s8 + 0M +L s7 + s8 + 0M +W GRCh38 0 chr1 0 30 >s1>s2>s4 +W GRCh38 0 chr2 0 30 >s5>s6>s8 +W CHM13 0 chr1 0 20 >s1>s3 +W CHM13 0 chr1 20 30 >s4 +W CHM13 0 chr2 0 30 >s5>s7>s8 +W HG01 1 chr1 0 30 >s1>s3>s4 +W HG01 2 chr1 0 30 >s1>s2>s4 +W HG01 1 chr2 0 30 >s5>s6>s8 +W HG01 2 chr2 0 30 >s5>s7>s8 +W HG02 1 chr1 0 30 >s1>s2>s4 +W HG02 2 chr2 0 30 >s5>s6>s8 diff --git a/tests/data/pangenome/tiny_sampled.gbz b/tests/data/pangenome/tiny_sampled.gbz new file mode 100644 index 0000000..c0e19be Binary files /dev/null and b/tests/data/pangenome/tiny_sampled.gbz differ diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 0000000..0436aa7 --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,23 @@ +""" +Shared integration-test fixtures +""" + +import pytest + +from sentieon_cli.job import Job + +# `pytest --doctest-modules` collects .py files as modules, and two +# conftest.py files outside a package collide on the module name. This one +# holds no doctests, so skip collecting it. +collect_ignore = ["conftest.py"] + + +@pytest.fixture(autouse=True) +def reset_job_ids(): + """Restart job id numbering for every test. + + ``job_id`` comes from a class-level per-name counter that the CLI resets + once per invocation, so without this a test's ids would depend on which + tests ran before it. + """ + Job.reset_ids() diff --git a/tests/integration/test_executor.py b/tests/integration/test_executor.py index 7199ecc..d0ea7f5 100644 --- a/tests/integration/test_executor.py +++ b/tests/integration/test_executor.py @@ -1,13 +1,14 @@ """ -Integration tests for the Executor +Integration tests for the executors. """ -import io +import asyncio import os import pathlib import sys import tempfile -from unittest.mock import patch +import threading +import time import pytest @@ -17,7 +18,12 @@ ) from sentieon_cli.dag import DAG # noqa: E402 -from sentieon_cli.executor import LocalExecutor # noqa: E402 +from sentieon_cli.exceptions import DagExecutionError # noqa: E402 +from sentieon_cli.executor import ( # noqa: E402 + DryRunExecutor, + LocalExecutor, + _cleanup_quietly, +) from sentieon_cli.job import Job # noqa: E402 from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 from sentieon_cli.shell_pipeline import ( # noqa: E402 @@ -30,88 +36,593 @@ def test_local_executor_simple_job(): """Test LocalExecutor with a single, simple job""" with tempfile.TemporaryDirectory() as tmp_dir_str: - tmp_dir = pathlib.Path(tmp_dir_str) - cmd_out = tmp_dir / "test_out.txt" + cmd_out = pathlib.Path(tmp_dir_str) / "test_out.txt" dag = DAG() - job = Job(Pipeline(Command("echo", "hello executor"), file_output=cmd_out), "echo-job") + job = Job( + Pipeline(Command("echo", "hello executor"), file_output=cmd_out), + "echo-job", + task_name="test", + ) dag.add_job(job) scheduler = ThreadScheduler(dag, 1) executor = LocalExecutor(scheduler) - executor.execute() - cmd_out_str = open(cmd_out).read() - assert "hello executor" in cmd_out_str + assert "hello executor" in cmd_out.read_text() assert len(executor.jobs_with_errors) == 0 def test_local_executor_pipeline_job(): """Test LocalExecutor with a job containing a pipeline""" with tempfile.TemporaryDirectory() as tmp_dir_str: - tmp_dir = pathlib.Path(tmp_dir_str) - cmd_out = tmp_dir / "test_out.txt" + cmd_out = pathlib.Path(tmp_dir_str) / "test_out.txt" dag = DAG() - pipeline = Pipeline(Command("echo", "hello pipeline"), Command("cat"), file_output=cmd_out) - job = Job(pipeline, "pipeline-job") + pipeline = Pipeline( + Command("echo", "hello pipeline"), + Command("cat"), + file_output=cmd_out, + ) + job = Job(pipeline, "pipeline-job", task_name="test") dag.add_job(job) scheduler = ThreadScheduler(dag, 1) executor = LocalExecutor(scheduler) - executor.execute() - cmd_out_str = open(cmd_out).read() - assert "hello pipeline" in cmd_out_str + assert "hello pipeline" in cmd_out.read_text() assert len(executor.jobs_with_errors) == 0 + def test_local_executor_failing_job(): """Test LocalExecutor with a job that fails""" with tempfile.TemporaryDirectory() as tmp_dir_str: - tmp_dir = pathlib.Path(tmp_dir_str) - cmd_in = tmp_dir / "test_in.txt" + cmd_in = pathlib.Path(tmp_dir_str) / "test_in.txt" dag = DAG() - # This command will fail - job = Job(Pipeline(Command("cat", str(cmd_in))), "failing-job") + # This command will fail: the input file does not exist + job = Job( + Pipeline(Command("cat", str(cmd_in))), + "failing-job", + task_name="test", + ) dag.add_job(job) scheduler = ThreadScheduler(dag, 1) executor = LocalExecutor(scheduler) - executor.execute() assert len(executor.jobs_with_errors) == 1 assert executor.jobs_with_errors[0] == job + def test_local_executor_proc_sub_job(): """Test LocalExecutor with a job using process substitution""" with tempfile.TemporaryDirectory() as tmp_dir_str: - tmp_dir = pathlib.Path(tmp_dir_str) - cmd_out = tmp_dir / "test_out.txt" + cmd_out = pathlib.Path(tmp_dir_str) / "test_out.txt" dag = DAG() command = Command( "diff", InputProcSub(Pipeline(Command("echo", "a"))), InputProcSub(Pipeline(Command("echo", "b"))), ) - job = Job(Pipeline(command, file_output=cmd_out), "proc-sub-job") + job = Job( + Pipeline(command, file_output=cmd_out), + "proc-sub-job", + task_name="test", + ) dag.add_job(job) # Need at least 2 threads for the two proc subs scheduler = ThreadScheduler(dag, 2) executor = LocalExecutor(scheduler) - executor.execute() # diff will write to stdout and fail - output = open(cmd_out).read() + output = cmd_out.read_text() assert "1c1" in output assert "< a" in output assert "> b" in output - # The job should be marked as an error because the diff command fails + # The job is marked as an error because the diff command fails assert len(executor.jobs_with_errors) == 1 -if __name__ == "__main__": - pytest.main([__file__]) \ No newline at end of file +def test_dry_run_executor_prints_without_running(capsys): + """DryRunExecutor prints the shell but does not execute it""" + with tempfile.TemporaryDirectory() as tmp_dir_str: + cmd_out = pathlib.Path(tmp_dir_str) / "out.txt" + dag = DAG() + job = Job( + Pipeline(Command("echo", "dryrun"), file_output=cmd_out), + "dry-job", + task_name="test", + ) + dag.add_job(job) + + scheduler = ThreadScheduler(dag, 1) + executor = DryRunExecutor(scheduler) + executor.execute() + + printed = capsys.readouterr().out + assert "echo" in printed + assert "dryrun" in printed + assert not cmd_out.exists() + assert len(executor.jobs_with_errors) == 0 + + +def test_dry_run_executor_drains_dependencies(capsys): + """DryRunExecutor walks the whole DAG, including dependent jobs""" + dag = DAG() + a = Job(Pipeline(Command("echo", "a")), "a", task_name="test") + b = Job(Pipeline(Command("echo", "b")), "b", task_name="test") + dag.add_job(a) + dag.add_job(b, {a}) + + scheduler = ThreadScheduler(dag, 4) + executor = DryRunExecutor(scheduler) + executor.execute() + + printed = capsys.readouterr().out + assert "echo a" in printed + assert "echo b" in printed + + +def test_fail_ok_subcommand_does_not_fail_job(): + """A failing but fail_ok sub-command does not mark the job as errored""" + with tempfile.TemporaryDirectory() as tmp_dir_str: + tmp_dir = pathlib.Path(tmp_dir_str) + missing = tmp_dir / "missing.txt" + out = tmp_dir / "out.txt" + + # `cat missing | wc -c`; cat fails (no file) but is tolerated. + dag = DAG() + job = Job( + Pipeline( + Command("cat", str(missing), fail_ok=True), + Command("wc", "-c"), + file_output=out, + ), + "tolerant-job", + task_name="test", + ) + dag.add_job(job) + + scheduler = ThreadScheduler(dag, 1) + executor = LocalExecutor(scheduler) + executor.execute() + + assert len(executor.jobs_with_errors) == 0 + + +def test_independent_jobs_all_complete(): + """Two independent jobs both run and produce output""" + with tempfile.TemporaryDirectory() as tmp_dir_str: + tmp_dir = pathlib.Path(tmp_dir_str) + o1 = tmp_dir / "o1.txt" + o2 = tmp_dir / "o2.txt" + dag = DAG() + dag.add_job( + Job( + Pipeline(Command("echo", "1"), file_output=o1), + "j1", + task_name="test", + ) + ) + dag.add_job( + Job( + Pipeline(Command("echo", "2"), file_output=o2), + "j2", + task_name="test", + ) + ) + + scheduler = ThreadScheduler(dag, 2) + executor = LocalExecutor(scheduler) + executor.execute() + + assert o1.read_text().strip() == "1" + assert o2.read_text().strip() == "2" + assert len(executor.jobs_with_errors) == 0 + + +def test_dependent_job_runs_after_dependency(): + """A dependent job runs and both jobs complete successfully""" + with tempfile.TemporaryDirectory() as tmp_dir_str: + tmp_dir = pathlib.Path(tmp_dir_str) + first = tmp_dir / "first.txt" + second = tmp_dir / "second.txt" + dag = DAG() + a = Job( + Pipeline(Command("echo", "a"), file_output=first), + "a", + task_name="test", + ) + b = Job( + Pipeline(Command("echo", "b"), file_output=second), + "b", + task_name="test", + ) + dag.add_job(a) + dag.add_job(b, {a}) + + scheduler = ThreadScheduler(dag, 1) + executor = LocalExecutor(scheduler) + executor.execute() + + assert first.read_text().strip() == "a" + assert second.read_text().strip() == "b" + assert len(executor.jobs_with_errors) == 0 + + +def test_job_that_fails_to_start_is_recorded(): + """A job whose command cannot be spawned is recorded as an error.""" + dag = DAG() + job = Job(Pipeline(Command("no_such_cmd_zzz")), "ghost", task_name="test") + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + executor.execute() + assert job in executor.jobs_with_errors + + +def test_launch_loop_stops_after_a_failed_launch(): + """Once a job in a ready batch fails to launch, the loop stops starting + the rest of the batch instead of launching all of it.""" + dag = DAG() + for i in range(5): + # Distinct args -> distinct pipeline identities (else the DAG + # rejects them as duplicates). + dag.add_job( + Job(Pipeline(Command("true", str(i))), f"j{i}", task_name="test") + ) + executor = LocalExecutor(ThreadScheduler(dag, 5)) + + launched = [] + + async def fake_run_job(job): + launched.append(job) + executor.start_new_jobs = False # simulate a launch failure + + executor.run_job = fake_run_job + executor.execute() + + # The guard stopped the loop after the first job rather than starting + # all five. + assert len(launched) == 1 + + +def test_infeasible_job_raises_instead_of_deadlocking(): + """A job larger than the budget raises rather than stalling silently.""" + dag = DAG() + dag.add_job( + Job(Pipeline(Command("echo", "hi")), "big", 100, task_name="test") + ) + with pytest.raises(DagExecutionError): + LocalExecutor(ThreadScheduler(dag, 2)).execute() + + +def test_dry_run_also_rejects_infeasible_job(): + dag = DAG() + dag.add_job( + Job(Pipeline(Command("echo", "hi")), "big", 100, task_name="test") + ) + with pytest.raises(DagExecutionError): + DryRunExecutor(ThreadScheduler(dag, 2)).execute() + + +def test_non_oserror_on_start_propagates(monkeypatch): + """A non-OSError while starting a job is a bug and must propagate.""" + + async def boom(self, *args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr("sentieon_cli.shell_pipeline.Pipeline.run", boom) + dag = DAG() + dag.add_job(Job(Pipeline(Command("echo", "x")), "j", task_name="test")) + with pytest.raises(RuntimeError): + LocalExecutor(ThreadScheduler(dag, 1)).execute() + + +def test_local_executor_empty_dag_completes(): + """An empty DAG runs cleanly.""" + executor = LocalExecutor(ThreadScheduler(DAG(), 1)) + executor.execute() + assert executor.jobs_with_errors == [] + + +def test_dry_run_empty_dag_completes(capsys): + """DryRunExecutor handles an empty DAG without printing anything.""" + executor = DryRunExecutor(ThreadScheduler(DAG(), 1)) + executor.execute() + assert capsys.readouterr().out == "" + assert executor.jobs_with_errors == [] + + +def _execute_bounded(executor, timeout=15): + """Run executor.execute() with a hang guard.""" + holder = {} + + def run(): + try: + executor.execute() + except BaseException as exc: # pragma: no cover - regression guard + holder["exc"] = exc + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=timeout) + assert not thread.is_alive(), "executor.execute() hung" + assert "exc" not in holder, holder.get("exc") + + +def test_procsub_job_with_early_exiting_outer_does_not_hang(): + """An outer command that exits without opening its proc-sub FIFO must + not hang the executor (regression: cleanup deadlock).""" + dag = DAG() + job = Job( + Pipeline( + Command("false", InputProcSub(Pipeline(Command("echo", "x")))) + ), + "early-exit", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + assert job in executor.jobs_with_errors + + +def test_procsub_job_with_unspawnable_outer_does_not_hang(): + """A missing outer executable with a proc-sub argument must not hang + the executor's error path (regression: cleanup deadlock).""" + dag = DAG() + job = Job( + Pipeline( + Command( + "no_such_cmd_zzz", + InputProcSub(Pipeline(Command("echo", "x"))), + ) + ), + "ghost-with-procsub", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + assert job in executor.jobs_with_errors + + +def test_procsub_inner_that_fails_to_launch_is_recorded(): + """An inner proc-sub command that cannot be spawned fails its job + instead of crashing the executor (regression: a raw FileNotFoundError + from the background task propagated out of execute()).""" + dag = DAG() + job = Job( + Pipeline( + Command( + "cat", + InputProcSub(Pipeline(Command("no_such_cmd_zzz"))), + ) + ), + "inner-fail", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + assert job in executor.jobs_with_errors + + +def test_procsub_inner_failure_does_not_skip_sibling_bookkeeping(): + """A job whose inner proc-sub fails to launch must not abort the run + for independent sibling jobs (regression: the exception escaped mid- + flight and skipped other jobs' bookkeeping).""" + dag = DAG() + bad = Job( + Pipeline( + Command( + "cat", + InputProcSub(Pipeline(Command("no_such_cmd_zzz"))), + ) + ), + "inner-fail", + task_name="test", + ) + good = Job(Pipeline(Command("true")), "sibling", task_name="test") + dag.add_job(bad) + dag.add_job(good) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + assert bad in executor.jobs_with_errors + # The sibling ran to completion and was recorded as a success, not + # silently dropped when the bad job's failure surfaced. + assert good not in executor.jobs_with_errors + + +def test_mid_pipeline_spawn_failure_terminates_running_stages(monkeypatch): + """When a later pipeline stage fails to spawn, earlier stages that are + already running must be terminated, not leaked (regression: a + Pipeline(sleep 600, no_such_cmd) left sleep alive past the workflow).""" + import sentieon_cli.executor as executor_mod + + contexts = [] + real_context = executor_mod.Context + + class RecordingContext(real_context): # type: ignore[valid-type,misc] + def __init__(self) -> None: + super().__init__() + contexts.append(self) + + monkeypatch.setattr(executor_mod, "Context", RecordingContext) + + dag = DAG() + job = Job( + Pipeline( + Command("sleep", "600"), + Command("no_such_cmd_zzz"), + ), + "leak", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + + assert job in executor.jobs_with_errors + # The first stage spawned; every spawned proc must be reaped, not left + # running to outlive the workflow. + spawned = [sub.proc for c in contexts for sub in c.commands if sub.proc] + assert spawned, "expected the first stage to have spawned" + for proc in spawned: + assert proc.returncode is not None + + +def test_error_in_the_run_loop_terminates_running_jobs(monkeypatch): + """An exception escaping the run loop must tear the children down first. + + Otherwise it unwinds into asyncio.run's shutdown, which joins the wait + threads -- and those are blocked in wait4 on children nobody signalled, + so the run hangs until they exit on their own (regression: a `sleep 30` + held the CLI for 30s and was never terminated).""" + import sentieon_cli.executor as executor_mod + + contexts = [] + real_context = executor_mod.Context + + class RecordingContext(real_context): # type: ignore[valid-type,misc] + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + contexts.append(self) + + monkeypatch.setattr(executor_mod, "Context", RecordingContext) + + dag = DAG() + dag.add_job(Job(Pipeline(Command("sleep", "30")), "slow", task_name="t")) + dag.add_job(Job(Pipeline(Command("true")), "quick", task_name="t")) + + scheduler = ThreadScheduler(dag, 2) + + def boom(job): + # Fires when the quick job finishes, with the slow one still running + raise RuntimeError("scheduler exploded") + + monkeypatch.setattr(scheduler, "job_finished", boom) + + executor = LocalExecutor(scheduler, shutdown_grace_period=0.5) + start = time.monotonic() + with pytest.raises(RuntimeError): + executor.execute() + elapsed = time.monotonic() - start + + assert elapsed < 15, "the run waited for the sleep instead of killing it" + spawned = [sub.proc for c in contexts for sub in c.commands if sub.proc] + assert spawned, "expected both jobs to have spawned" + for proc in spawned: + assert proc.returncode is not None + + +def test_thread_pool_size_is_configurable(): + """The wait/FIFO-open pool cap is a parameter, not a hidden constant.""" + scheduler = ThreadScheduler(DAG(), 1) + + assert LocalExecutor(scheduler).thread_pool_size == 512 + assert LocalExecutor(scheduler, thread_pool_size=8).thread_pool_size == 8 + + +def test_sigpipe_producer_does_not_fail_job(tmp_path): + """A producer killed by SIGPIPE when a consumer exits early (e.g. + `seq | head -1`) produces the right output and must not fail the job + (regression: the producer's rc -13 was treated as a failure).""" + out = tmp_path / "out.txt" + dag = DAG() + job = Job( + Pipeline( + Command("seq", "1000000"), + Command("head", "-1"), + file_output=out, + ), + "sigpipe", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + assert job not in executor.jobs_with_errors + assert out.read_text().strip() == "1" + + +def test_nonzero_producer_still_fails_job(): + """SIGPIPE forgiveness must not mask a producer that exits non-zero + for another reason: `false | cat` still fails the job.""" + dag = DAG() + job = Job( + Pipeline(Command("false"), Command("cat")), + "real-fail", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 2)) + _execute_bounded(executor) + assert job in executor.jobs_with_errors + + +class _FakeContext: + """A stand-in Context whose cleanup() can be told to raise.""" + + def __init__(self, cleanup_error=None): + self.commands = [] + self.cleaned = False + self.log_sink = None + self._error = cleanup_error + + async def cleanup(self): + self.cleaned = True + if self._error is not None: + raise self._error + + +@pytest.mark.asyncio +async def test_cleanup_quietly_swallows_oserror(): + """A launch failure surfacing from cleanup() during teardown is logged, + not raised -- so it cannot escape execute() on the abort/interrupt + paths (symmetry with _drive's per-job handling).""" + context = _FakeContext(FileNotFoundError("no_such_inner")) + await _cleanup_quietly(context) # must not raise + assert context.cleaned + + +@pytest.mark.asyncio +async def test_cleanup_quietly_propagates_non_oserror(): + """A non-OSError from cleanup() is a bug and must still propagate.""" + with pytest.raises(RuntimeError): + await _cleanup_quietly(_FakeContext(RuntimeError("bug"))) + + +@pytest.mark.asyncio +async def test_shutdown_releases_all_contexts_despite_cleanup_error(): + """A cleanup() OSError during interrupt shutdown is swallowed and the + remaining contexts are still released, instead of the error escaping + _shutdown() and skipping their teardown.""" + executor = LocalExecutor(ThreadScheduler(DAG(), 1)) + bad = _FakeContext(FileNotFoundError("no_such_inner")) + good = _FakeContext() + + async def _done(): + return 0 + + executor.running = [ + ( + Job(Pipeline(Command("true")), "j1", task_name="test"), + bad, + asyncio.create_task(_done()), + 0, + ), + ( + Job(Pipeline(Command("false")), "j2", task_name="test"), + good, + asyncio.create_task(_done()), + 0, + ), + ] + await executor._shutdown() # must not raise + + assert bad.cleaned and good.cleaned diff --git a/tests/integration/test_job_log_capture.py b/tests/integration/test_job_log_capture.py new file mode 100644 index 0000000..0940cec --- /dev/null +++ b/tests/integration/test_job_log_capture.py @@ -0,0 +1,292 @@ +""" +Integration tests for capturing job output to per-process log files. +""" + +import asyncio +import logging +import os +import sys +from typing import List + +import pytest + +# Add the parent directory to the path +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.executor import LocalExecutor # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.run_logs import RunLogs # noqa: E402 +from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 +from sentieon_cli.shell_pipeline import ( # noqa: E402 + Command, + Context, + InputProcSub, + OutputProcSub, + Pipeline, +) + + +class _RecordingHandler(logging.Handler): + """Collects formatted messages from the package logger.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.messages: List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def messages(): + """Capture the package logger, which does not propagate to the root.""" + handler = _RecordingHandler() + package_logger = logging.getLogger("sentieon_cli") + package_logger.addHandler(handler) + yield handler.messages + package_logger.removeHandler(handler) + + +def _run(dag: DAG, log_dir, cores: int = 2): + """Execute a DAG with its output captured under ``log_dir``""" + run_logs = RunLogs(log_dir) + run_logs.create_dirs() + executor = LocalExecutor(ThreadScheduler(dag, cores), run_logs=run_logs) + executor.execute() + return executor, run_logs + + +def _task_dir(run_logs: RunLogs, task_name: str): + return run_logs.task_logs / task_name + + +def test_each_pipeline_stage_gets_its_own_stderr_log(tmp_path): + """Every stage's stderr is captured; only uncaptured stdout is.""" + dag = DAG() + job = Job( + Pipeline( + Command("sh", "-c", "echo payload; echo first >&2"), + Command("sh", "-c", "cat; echo second >&2"), + ), + "staged", + task_name="two-stage", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert executor.jobs_with_errors == [] + task_dir = _task_dir(run_logs, "two-stage") + assert sorted(p.name for p in task_dir.iterdir()) == [ + "staged-1.0.log", + "staged-1.1.log", + "staged-1.1.stdout.log", + ] + assert "first" in (task_dir / "staged-1.0.log").read_text() + assert "second" in (task_dir / "staged-1.1.log").read_text() + # The last stage's stdout was inherited, so it is captured verbatim + # (sharing the stage index of that stage's stderr log); the stdout + # piped between the stages is not. + assert (task_dir / "staged-1.1.stdout.log").read_text() == "payload\n" + + +def test_a_file_output_wins_over_the_sink(tmp_path): + """An explicit destination is never redirected into the log dir.""" + out = tmp_path / "out.txt" + dag = DAG() + job = Job( + Pipeline(Command("echo", "hello"), file_output=out), + "redirected", + task_name="file-out", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert executor.jobs_with_errors == [] + assert out.read_text() == "hello\n" + # The command was silent, so its header-only stderr log was pruned. + assert list(_task_dir(run_logs, "file-out").iterdir()) == [] + + +def test_input_proc_sub_stderr_is_captured(tmp_path): + """An inner <(...) command writes to the job's sink, not the terminal.""" + out = tmp_path / "out.txt" + dag = DAG() + job = Job( + Pipeline( + Command( + "cat", + InputProcSub( + Pipeline( + Command("sh", "-c", "echo inner >&2; echo payload") + ) + ), + ), + file_output=out, + ), + "reader", + task_name="proc-sub", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert executor.jobs_with_errors == [] + assert out.read_text() == "payload\n" + logs = list(_task_dir(run_logs, "proc-sub").iterdir()) + assert [p.name for p in logs] == ["reader-1.1.log"] + assert "inner" in logs[0].read_text() + + +def test_output_proc_sub_stderr_is_captured(tmp_path): + """An inner >(...) command writes to the job's sink too.""" + dag = DAG() + job = Job( + Pipeline( + Command( + "sh", + "-c", + 'echo payload > "$1"', + "sh", + OutputProcSub( + Pipeline(Command("sh", "-c", "cat >&2; echo done >&2")) + ), + ) + ), + "writer", + task_name="proc-sub", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert executor.jobs_with_errors == [] + contents = [ + path.read_text() for path in _task_dir(run_logs, "proc-sub").iterdir() + ] + assert any("payload" in text and "done" in text for text in contents) + + +def test_a_successful_job_prunes_its_empty_logs(tmp_path): + dag = DAG() + dag.add_job(Job(Pipeline(Command("true")), "quiet", task_name="cleanup")) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert executor.jobs_with_errors == [] + assert list(_task_dir(run_logs, "cleanup").iterdir()) == [] + + +def test_a_failing_job_is_reported_with_its_log_and_tail(tmp_path, messages): + dag = DAG() + job = Job( + Pipeline(Command("sh", "-c", "echo boom >&2; exit 3")), + "boomer", + task_name="failing", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert job in executor.jobs_with_errors + log_path = _task_dir(run_logs, "failing") / "boomer-1.0.log" + assert log_path.is_file() # nothing is pruned after a failure + assert "boom" in log_path.read_text() + + report = "\n".join(messages) + pid = job.shell.nodes[0].proc.pid + assert f"Failed sub-command of {job}" in report + assert f"exit code: 3, pid: {pid}" in report + assert str(log_path) in report + assert " boom" in report + assert f"Task logs for debugging: {run_logs.task_logs}" in report + + +def test_a_failing_stage_reports_its_own_log(tmp_path, messages): + """The tail comes from the failing process, not a sibling stage.""" + dag = DAG() + job = Job( + Pipeline( + Command("sh", "-c", "echo upstream-noise >&2"), + Command("sh", "-c", "cat; echo downstream-boom >&2; exit 4"), + ), + "mixed", + task_name="failing", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert job in executor.jobs_with_errors + report = "\n".join(messages) + # Quoted log lines are indented; the sibling's output is not quoted. + assert " downstream-boom" in report + assert " upstream-noise" not in report + assert str(_task_dir(run_logs, "failing") / "mixed-1.1.log") in report + + +def test_a_job_that_cannot_start_points_at_its_partial_logs( + tmp_path, messages +): + dag = DAG() + job = Job( + Pipeline( + Command("sh", "-c", "echo started >&2; sleep 5"), + Command("no_such_cmd_zzz"), + ), + "unstartable", + task_name="launch-failure", + ) + dag.add_job(job) + executor, run_logs = _run(dag, tmp_path / "logs") + + assert job in executor.jobs_with_errors + report = "\n".join(messages) + # The stage that did spawn keeps its log, and the report names it. + log_path = _task_dir(run_logs, "launch-failure") / "unstartable-1.0.log" + assert log_path.is_file() + assert f"Logs from {job}" in report + assert str(log_path) in report + + +@pytest.mark.asyncio +async def test_an_interrupted_job_keeps_and_reports_its_logs( + tmp_path, messages +): + """Shutdown after an interrupt prunes nothing and names the logs.""" + run_logs = RunLogs(tmp_path / "logs") + run_logs.create_dirs() + job = Job( + Pipeline(Command("sleep", "30")), + "sleeper", + task_name="interrupted", + ) + sink = run_logs.job_sink(job) + context = Context(log_sink=sink) + context.file_handles.append(sink.open_stderr(Command("sleep", "30"))) + + async def _done() -> int: + return 0 + + executor = LocalExecutor(ThreadScheduler(DAG(), 1), run_logs=run_logs) + executor.running = [(job, context, asyncio.create_task(_done()), 0)] + await executor._shutdown() + + log_path = sink.log_paths()[0] + assert log_path.is_file() + assert str(log_path) in "\n".join(messages) + + +def test_without_run_logs_no_files_are_written(tmp_path): + """An executor with no log directory behaves exactly as before.""" + out = tmp_path / "out.txt" + dag = DAG() + dag.add_job( + Job( + Pipeline(Command("echo", "hi"), file_output=out), + "plain", + task_name="no-logs", + ) + ) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + + assert executor.jobs_with_errors == [] + assert [p.name for p in tmp_path.iterdir()] == ["out.txt"] diff --git a/tests/integration/test_pipeline_execution.py b/tests/integration/test_pipeline_execution.py index 588ab9b..9ed9e6a 100644 --- a/tests/integration/test_pipeline_execution.py +++ b/tests/integration/test_pipeline_execution.py @@ -152,11 +152,15 @@ def test_hybrid_dry_run_execution_rgsm(self): all_jobs = [job for job in list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys())] # with --rgsm should use replace_rg during first pass - first_pass_job = [job for job in all_jobs if job.name == "calling-1"][0] + first_pass_job = [job for job in all_jobs if job.name == "dnascope-1"][0] assert "--replace_rg" in first_pass_job.shell.nodes[0].args + # CNV calling is sex-aware and runs in the second DAG + second_dag = pipeline.build_second_dag() + second_jobs = list(second_dag.waiting_jobs.keys()) + list(second_dag.ready_jobs.keys()) + # should use replace_rg in CNV calling - cnv_job = [job for job in all_jobs if job.name == "CNVscope"][0] + cnv_job = [job for job in second_jobs if job.name == "cnvscope"][0] assert "--replace_rg" in cnv_job.shell.nodes[0].args # should use replace_rg in SV calling @@ -188,7 +192,7 @@ def test_dnascope_command_parameters(self): dag = pipeline.build_dag() # Find variant calling job - variant_job = DAGAnalyzer.get_job_by_name(dag, "variant-calling") + variant_job = DAGAnalyzer.get_job_by_name(dag, "dnascope") assert variant_job is not None, "Should have variant calling job" # Check that command includes expected parameters @@ -360,7 +364,7 @@ def test_variant_calling_dependencies(self): dag = pipeline.build_dag() # Variant calling should depend on either dedup or input processing - variant_jobs = DAGAnalyzer.get_jobs_by_type(dag, "variant") + variant_jobs = DAGAnalyzer.get_jobs_by_type(dag, "dnascope") assert len(variant_jobs) > 0, "Should have variant calling jobs" diff --git a/tests/integration/test_signals.py b/tests/integration/test_signals.py new file mode 100644 index 0000000..665a935 --- /dev/null +++ b/tests/integration/test_signals.py @@ -0,0 +1,204 @@ +""" +Signal-handling tests for LocalExecutor. + +These deliver real signals to the current process. That is safe here because +the executor installs an asyncio loop signal handler that suppresses the +default disposition (e.g. SIGINT -> KeyboardInterrupt) for the duration of the +run, and the timer that raises the signal is always armed *after* the handler +is installed. If these ever prove flaky in a constrained CI, the same behavior +can be exercised by running the workflow in a subprocess and signalling the +child. +""" + +import os +import signal +import sys +import threading +import time + +# Add the parent directory to the path +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +import sentieon_cli.executor as executor_mod # noqa: E402 +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.executor import LocalExecutor # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 +from sentieon_cli.shell_pipeline import ( # noqa: E402 + Command, + InputProcSub, + Pipeline, +) + + +def _echo_dag(out): + dag = DAG() + job = Job( + Pipeline(Command("echo", "hi"), file_output=out), + "echo", + task_name="test", + ) + dag.add_job(job) + return dag + + +def test_off_main_thread_does_not_crash(tmp_path): + """LocalExecutor with signal handlers runs fine off the main thread.""" + out = tmp_path / "o.txt" + holder = {} + + def run(): + try: + LocalExecutor( + ThreadScheduler(_echo_dag(out), 1), + install_signal_handlers=True, + ).execute() + except Exception as exc: # pragma: no cover - regression guard + holder["exc"] = exc + + thread = threading.Thread(target=run) + thread.start() + thread.join(timeout=10) + + assert not thread.is_alive() + assert "exc" not in holder, holder.get("exc") + assert out.read_text().strip() == "hi" + + +def test_default_does_not_touch_signal_handlers(tmp_path): + """The default executor installs no handlers and leaves them intact.""" + before_int = signal.getsignal(signal.SIGINT) + before_term = signal.getsignal(signal.SIGTERM) + + out = tmp_path / "o.txt" + LocalExecutor(ThreadScheduler(_echo_dag(out), 1)).execute() + + assert signal.getsignal(signal.SIGINT) is before_int + assert signal.getsignal(signal.SIGTERM) is before_term + + +def test_opt_in_restores_signal_handlers(tmp_path): + """Opt-in handlers are restored to their prior values after a run.""" + before_int = signal.getsignal(signal.SIGINT) + before_term = signal.getsignal(signal.SIGTERM) + + out = tmp_path / "o.txt" + LocalExecutor( + ThreadScheduler(_echo_dag(out), 1), + install_signal_handlers=True, + ).execute() + + assert signal.getsignal(signal.SIGINT) is before_int + assert signal.getsignal(signal.SIGTERM) is before_term + + +def test_sigint_triggers_graceful_shutdown(monkeypatch): + """A delivered SIGINT terminates running jobs and cleans up.""" + temp_dirs = [] + real_context = executor_mod.Context + + class SpyContext(real_context): + async def cleanup(self): + temp_dirs.append(self.temp_dir.name) + await super().cleanup() + + monkeypatch.setattr(executor_mod, "Context", SpyContext) + + dag = DAG() + dag.add_job( + Job(Pipeline(Command("sleep", "30")), "sleeper", task_name="test") + ) + executor = LocalExecutor( + ThreadScheduler(dag, 1), + install_signal_handlers=True, + shutdown_grace_period=1.0, + ) + + timer = threading.Timer(0.5, lambda: os.kill(os.getpid(), signal.SIGINT)) + timer.start() + try: + start = time.monotonic() + executor.execute() + elapsed = time.monotonic() - start + finally: + timer.cancel() + + assert elapsed < 10 # returned well before `sleep 30` would finish + assert executor.running == [] + assert temp_dirs # cleanup ran for the interrupted job + for temp_dir in temp_dirs: + assert not os.path.exists(temp_dir) # no leaked temp dir + + +def test_grace_period_escalates_to_sigkill(): + """A job that ignores SIGTERM is still killed after the grace period.""" + script = ( + "import signal, time; " + "signal.signal(signal.SIGTERM, signal.SIG_IGN); " + "time.sleep(30)" + ) + dag = DAG() + job = Job( + Pipeline(Command(sys.executable, "-c", script)), + "stubborn", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor( + ThreadScheduler(dag, 1), + install_signal_handlers=True, + shutdown_grace_period=0.5, + ) + + timer = threading.Timer(0.5, lambda: os.kill(os.getpid(), signal.SIGINT)) + timer.start() + try: + start = time.monotonic() + executor.execute() + elapsed = time.monotonic() - start + finally: + timer.cancel() + + assert elapsed < 10 # SIGKILL fired after grace, not after time.sleep(30) + assert executor.running == [] + + +def test_sigint_with_blocked_procsub_cleans_up(): + """SIGINT while a proc-sub FIFO open is still blocked must not hang + the shutdown path (regression: cleanup deadlock).""" + dag = DAG() + # sh never opens the FIFO in "$1", so the proc sub's open() stays + # blocked until cleanup unblocks it. + job = Job( + Pipeline( + Command( + "sh", + "-c", + "sleep 30", + "sh", + InputProcSub(Pipeline(Command("echo", "x"))), + ) + ), + "blocked-procsub", + task_name="test", + ) + dag.add_job(job) + executor = LocalExecutor( + ThreadScheduler(dag, 2), + install_signal_handlers=True, + shutdown_grace_period=1.0, + ) + + timer = threading.Timer(0.5, lambda: os.kill(os.getpid(), signal.SIGINT)) + timer.start() + try: + start = time.monotonic() + executor.execute() + elapsed = time.monotonic() - start + finally: + timer.cancel() + + assert elapsed < 10 + assert executor.running == [] diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..226413c --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,43 @@ +""" +Shared unit-test fixtures +""" + +import logging + +import pytest + +from sentieon_cli import logging as cli_logging +from sentieon_cli.job import Job + + +@pytest.fixture(autouse=True) +def reset_job_ids(): + """Restart job id numbering for every test. + + ``job_id`` comes from a class-level per-name counter that the CLI resets + once per invocation, so without this a test's ids would depend on which + tests ran before it. + """ + Job.reset_ids() + + +@pytest.fixture(autouse=True) +def no_handler_leaks(): + """Drop any handler a test leaves on the package logger. + + `setup_logging` attaches a `run.log` FileHandler to the package logger, + so any test that reaches it with an output VCF (or log_dir) set would + otherwise leak that handler into later tests. + """ + package_logger = logging.getLogger("sentieon_cli") + before = list(package_logger.handlers) + console_level = cli_logging._console_handler.level + yield + for handler in list(package_logger.handlers): + if handler in before: + continue + if isinstance(handler, logging.FileHandler): + cli_logging.remove_file_handler(handler) + else: + package_logger.removeHandler(handler) + cli_logging.set_console_level(console_level) diff --git a/tests/unit/test_cli_args.py b/tests/unit/test_cli_args.py new file mode 100644 index 0000000..28909ed --- /dev/null +++ b/tests/unit/test_cli_args.py @@ -0,0 +1,40 @@ +""" +Unit tests for the top-level CLI arguments +""" + +import argparse +import os +import sys + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli import add_logging_args, resolve_loglevel # noqa: E402 + + +def _loglevel(argv): + parser = argparse.ArgumentParser() + add_logging_args(parser) + return resolve_loglevel(parser.parse_args(argv)) + + +@pytest.mark.parametrize( + "argv,expected", + [ + ([], "INFO"), + (["-v"], "INFO"), + (["-q"], "WARNING"), + (["-d"], "DEBUG"), + # `-d` wins wherever it appears; it used to be undone by a later + # `-v` because both flags wrote to one destination. + (["-d", "-v"], "DEBUG"), + (["-q", "-v"], "INFO"), + ], +) +def test_the_verbosity_flags_resolve_to_a_console_level(argv, expected): + assert _loglevel(argv) == expected diff --git a/tests/unit/test_cnv_sex.py b/tests/unit/test_cnv_sex.py new file mode 100644 index 0000000..3fc4f1d --- /dev/null +++ b/tests/unit/test_cnv_sex.py @@ -0,0 +1,414 @@ +""" +Unit tests for sex-aware CNV calling: reference build detection, PAR BED +selection, the CNVscope algo arguments, and the `--sample_sex` argument. +""" + +import argparse +import logging +import multiprocessing as mp +import pathlib +from typing import List + +import pytest + +from sentieon_cli import shard +from sentieon_cli.driver import CNVscope +from sentieon_cli.pipeline import BasePipeline +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext +from sentieon_cli.stages.cnv import CNVscopeStage +from sentieon_cli.shard import ( + BUILD_SIGNATURES, + detect_reference_build, + par_bed_for_build, + ploidy_contigs_for_build, +) +from sentieon_cli.util import ( + SampleSex, + caller_sex_arg, + cnvscope_sex_args, + sample_sex_arg, +) + +PACKAGE_LOGGER = "sentieon_cli" + + +class _RecordingHandler(logging.Handler): + """Collects formatted messages from the package logger.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.messages: List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def messages(): + """Capture the package logger, which does not propagate to the root.""" + handler = _RecordingHandler() + package_logger = logging.getLogger(PACKAGE_LOGGER) + package_logger.addHandler(handler) + yield handler.messages + package_logger.removeHandler(handler) + + +def _fai(contigs) -> dict: + """A parsed fasta index with the supplied contig lengths""" + return { + ctg: {"length": length, "offset": 0, "linebases": 60, "linewidth": 61} + for ctg, length in contigs.items() + } + + +class _MinimalPipeline(BasePipeline): + """A concrete pipeline for exercising the shared sex machinery""" + + def validate(self) -> None: + pass + + def configure(self) -> None: + pass + + def build_dag(self) -> DAG: + return DAG() + + +def _pipeline(**kwargs) -> _MinimalPipeline: + pipeline = _MinimalPipeline() + pipeline.setup_logging(argparse.Namespace(loglevel="DEBUG")) + for key, value in kwargs.items(): + setattr(pipeline, key, value) + return pipeline + + +# Reference build detection + + +@pytest.mark.parametrize("build", sorted(BUILD_SIGNATURES)) +def test_detect_reference_build_identifies_each_build(build): + assert detect_reference_build(_fai(BUILD_SIGNATURES[build])) == build + + +def test_detect_reference_build_distinguishes_hg19_from_b37(): + # hg19 and b37 share contig lengths and differ only by the 'chr' prefix + hg19 = detect_reference_build(_fai(BUILD_SIGNATURES["hg19"])) + b37 = detect_reference_build(_fai(BUILD_SIGNATURES["b37"])) + assert hg19 == "hg19" + assert b37 == "b37" + + +def test_detect_reference_build_requires_all_signature_contigs(): + signature = dict(BUILD_SIGNATURES["hg38"]) + del signature["chrX"] + assert detect_reference_build(_fai(signature)) is None + + +def test_detect_reference_build_requires_exact_lengths(): + signature = dict(BUILD_SIGNATURES["hg38"]) + signature["chr2"] = signature["chr2"] - 1 + assert detect_reference_build(_fai(signature)) is None + + +def test_detect_reference_build_with_an_empty_index(): + assert detect_reference_build({}) is None + + +def test_detect_reference_build_with_an_unknown_reference(): + assert detect_reference_build(_fai({"contig1": 1000})) is None + + +# PAR BED selection + + +def test_par_bed_for_build_without_a_build(): + assert par_bed_for_build(None) is None + + +def test_par_bed_for_build_with_an_absent_file(messages): + assert par_bed_for_build("not_a_build") is None + assert any("PAR" in msg for msg in messages) + + +def test_par_bed_for_build_treats_an_empty_file_as_absent( + tmp_path, monkeypatch, messages +): + # The packaged BED files ship empty until their contents are added + empty_bed = tmp_path.joinpath("par_hg38.bed") + empty_bed.touch() + monkeypatch.setattr(shard, "files", lambda pkg: tmp_path) + + assert par_bed_for_build("hg38") is None + assert any("PAR" in msg for msg in messages) + + +def test_par_bed_for_build_returns_a_packaged_file(tmp_path, monkeypatch): + par_bed = tmp_path.joinpath("par_hg38.bed") + par_bed.write_text("chrX\t10000\t2781479\n") + monkeypatch.setattr(shard, "files", lambda pkg: tmp_path) + + assert par_bed_for_build("hg38") == par_bed + + +# Ploidy estimation contigs + + +def test_ploidy_contigs_for_b37_use_bare_contig_names(): + contigs = ploidy_contigs_for_build("b37") + assert contigs.contigs is not None + assert contigs.contigs[0] == "1" + assert contigs.autosomes == [str(i) for i in range(1, 23)] + assert contigs.x_contig == "X" + assert contigs.y_contig == "Y" + + +@pytest.mark.parametrize("build", ["hg38", "hg19", "chm13", None]) +def test_ploidy_contigs_default_to_the_script_defaults(build): + contigs = ploidy_contigs_for_build(build) + assert contigs.contigs is None + assert contigs.autosomes is None + assert contigs.x_contig is None + assert contigs.y_contig is None + + +# The `--sample_sex` argument + + +@pytest.mark.parametrize("value", ["male", "MALE", " Male "]) +def test_sample_sex_arg_parses_male(value): + assert sample_sex_arg(value) is SampleSex.MALE + + +@pytest.mark.parametrize("value", ["female", "Female", "FEMALE"]) +def test_sample_sex_arg_parses_female(value): + assert sample_sex_arg(value) is SampleSex.FEMALE + + +@pytest.mark.parametrize("value", ["m", "f", "unknown", ""]) +def test_sample_sex_arg_rejects_other_values(value): + with pytest.raises(argparse.ArgumentTypeError): + sample_sex_arg(value) + + +# The CNVscope `--sex` and `--par` arguments + + +def test_cnvscope_sex_args_for_a_male_sample(): + par_bed = pathlib.Path("/par.bed") + assert cnvscope_sex_args(SampleSex.MALE, par_bed) == ("M", par_bed) + + +def test_cnvscope_sex_args_for_a_female_sample(): + # A female sample does not need the PAR regions + par_bed = pathlib.Path("/par.bed") + assert cnvscope_sex_args(SampleSex.FEMALE, par_bed) == ("F", None) + + +@pytest.mark.parametrize("sex", [None, SampleSex.UNKNOWN]) +def test_cnvscope_sex_args_without_a_known_sex(sex, messages): + assert cnvscope_sex_args(sex, None) == (None, None) + assert any("diploid" in msg for msg in messages) + + +# The ExpansionHunter and segdup-caller `--sex` argument + + +def test_caller_sex_arg_for_a_male_sample(): + assert caller_sex_arg(SampleSex.MALE) == "male" + + +@pytest.mark.parametrize("sex", [SampleSex.FEMALE, SampleSex.UNKNOWN, None]) +def test_caller_sex_arg_for_every_other_sample(sex): + # The two callers take only "male" or "female" + assert caller_sex_arg(sex) == "female" + + +def test_cnvscope_algo_emits_the_sex_arguments(): + cmd = CNVscope( + pathlib.Path("out.vcf.gz"), + pathlib.Path("bundle/cnv.model"), + sex="M", + par=pathlib.Path("par.bed"), + ).build_cmd() + + assert cmd == [ + "--algo", + "CNVscope", + "--model", + "bundle/cnv.model", + "--sex", + "M", + "--par", + "par.bed", + "out.vcf.gz", + ] + + +def test_cnvscope_algo_omits_the_sex_arguments_by_default(): + cmd = CNVscope( + pathlib.Path("out.vcf.gz"), + pathlib.Path("bundle/cnv.model"), + ).build_cmd() + + assert "--sex" not in cmd + assert "--par" not in cmd + assert "--dump_probes" not in cmd + assert cmd[-1] == "out.vcf.gz" + + +def test_cnvscope_algo_emits_dump_probes(): + cmd = CNVscope( + pathlib.Path("out.vcf.gz"), + pathlib.Path("bundle/cnv.model"), + sex="F", + dump_probes=pathlib.Path("out.cnv.probes"), + ).build_cmd() + + assert cmd == [ + "--algo", + "CNVscope", + "--model", + "bundle/cnv.model", + "--sex", + "F", + "--dump_probes", + "out.cnv.probes", + "out.vcf.gz", + ] + + +# The shared fail-fast helpers + + +@pytest.mark.parametrize( + "sample_sex", [None, SampleSex.MALE, SampleSex.FEMALE, SampleSex.UNKNOWN] +) +def test_validate_cnv_par_requires_a_par_bed_for_every_sex( + sample_sex, messages +): + # CNV calling needs the PAR BED file whatever the sample sex is, and + # the run stops during validation without one + pipeline = _pipeline(sample_sex=sample_sex, cnv_par_bed=None) + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_cnv_par(True) + assert excinfo.value.code == 2 + assert any("--par_bed" in msg for msg in messages) + + +@pytest.mark.parametrize( + "sample_sex", [None, SampleSex.MALE, SampleSex.FEMALE, SampleSex.UNKNOWN] +) +def test_validate_cnv_par_accepts_a_par_bed(sample_sex): + pipeline = _pipeline( + sample_sex=sample_sex, cnv_par_bed=pathlib.Path("/par.bed") + ) + + pipeline.validate_cnv_par(True) # no SystemExit + + +def test_validate_cnv_par_is_skipped_when_cnvs_are_not_called(): + # A PAR BED file is only needed by CNV calling + pipeline = _pipeline(sample_sex=SampleSex.MALE, cnv_par_bed=None) + + pipeline.validate_cnv_par(False) # no SystemExit + + +def test_validate_cnv_par_errors_in_a_dry_run_too(): + # There is no dry-run exception: a dry-run of a real command needs + # the same arguments the real run would use + pipeline = _pipeline(cnv_par_bed=None, dry_run=True) + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_cnv_par(True) + assert excinfo.value.code == 2 + + +def test_resolve_cnv_par_bed_prefers_the_supplied_file(tmp_path): + par_bed = tmp_path.joinpath("custom_par.bed") + par_bed.write_text("chrX\t10000\t2781479\n") + pipeline = _pipeline() + + pipeline.resolve_cnv_par_bed(_fai(BUILD_SIGNATURES["hg38"]), par_bed) + + assert pipeline.cnv_par_bed == par_bed + # The build is still identified, for the ploidy estimation contigs + assert pipeline.reference_build == "hg38" + + +def test_resolve_cnv_par_bed_falls_back_to_the_packaged_file( + tmp_path, monkeypatch +): + packaged = tmp_path.joinpath("par_b37.bed") + packaged.write_text("X\t60001\t2699520\n") + monkeypatch.setattr(shard, "files", lambda pkg: tmp_path) + pipeline = _pipeline() + + pipeline.resolve_cnv_par_bed(_fai(BUILD_SIGNATURES["b37"])) + + assert pipeline.reference_build == "b37" + assert pipeline.cnv_par_bed == packaged + + +def test_resolve_cnv_par_bed_skips_the_lookup_without_cnv_calling(messages): + # The build is still identified, for the ploidy estimation contigs, + # but a run without CNV calling needs no PAR BED file + pipeline = _pipeline() + + pipeline.resolve_cnv_par_bed( + _fai(BUILD_SIGNATURES["hg38"]), None, cnv_will_run=False + ) + + assert pipeline.reference_build == "hg38" + assert pipeline.cnv_par_bed is None + assert not any("PAR" in msg for msg in messages) + + +# `CNVscopeStage` passes the sample sex through to CNVscope + + +def _cnvscope_cmd(**kwargs) -> str: + tmp_dir = pathlib.Path("/tmp/cnv") + output_vcf = pathlib.Path("/out/sample.vcf.gz") + ctx = StageContext( + reference=pathlib.Path("/ref/reference.fa"), + output_vcf=output_vcf, + tmp_dir=tmp_dir, + cores=mp.cpu_count(), + dry_run=True, + skip_version_check=True, + ) + result = CNVscopeStage( + ctx=ctx, + inputs=[pathlib.Path("/in/sample.cram")], + model=pathlib.Path("/bundle/model.bundle/cnv.model"), + cnvscope_vcf=tmp_dir.joinpath("cnvscope.vcf.gz"), + cnv_vcf=pathlib.Path( + str(output_vcf).replace(".vcf.gz", ".cnv.vcf.gz") + ), + **kwargs, + ).add_to(DAG()) + return str(result.cnvscope_job.shell) + + +def test_call_cnvs_passes_the_sex_of_a_male_sample(): + cmd = _cnvscope_cmd( + sample_sex=SampleSex.MALE, par_bed=pathlib.Path("/par.bed") + ) + assert "--sex M" in cmd + assert "--par /par.bed" in cmd + + +def test_call_cnvs_passes_the_sex_of_a_female_sample(): + cmd = _cnvscope_cmd( + sample_sex=SampleSex.FEMALE, par_bed=pathlib.Path("/par.bed") + ) + assert "--sex F" in cmd + assert "--par" not in cmd + + +def test_call_cnvs_without_a_sample_sex(): + cmd = _cnvscope_cmd() + assert "--sex" not in cmd + assert "--par" not in cmd diff --git a/tests/unit/test_command_strings_realign.py b/tests/unit/test_command_strings_realign.py new file mode 100644 index 0000000..abfb56d --- /dev/null +++ b/tests/unit/test_command_strings_realign.py @@ -0,0 +1,138 @@ +""" +Unit tests for the BAM/CRAM re-alignment command builders +""" + +import os +import pathlib +import sys + +# Add the parent directory to the path to import sentieon_cli +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +from sentieon_cli.command_strings import ( + cmd_samtools_fastq_bwa, + cmd_samtools_fastq_minimap2, +) + + +REFERENCE = pathlib.Path("/ref/hg38.fa") +INPUT_REF = pathlib.Path("/ref/hs37d5.fa") +MODEL_BUNDLE = pathlib.Path("/bundle/sample.bundle") +INPUT_ALN = pathlib.Path("/data/sample.cram") +OUT_ALN = pathlib.Path("/out/sample_mm2_sorted_0.cram") +RG_HEADER = pathlib.Path("/tmp/sample.hdr") + + +class TestSamtoolsFastqMinimap2: + """Test the minimap2 re-alignment command""" + + def build(self, **kwargs) -> str: + args = { + "out_aln": OUT_ALN, + "input_aln": INPUT_ALN, + "reference": REFERENCE, + "model_bundle": MODEL_BUNDLE, + "cores": 4, + "rg_lines": ["@RG\tID:rg1\tSM:sample1"], + "sample_name": "sample1", + } + args.update(kwargs) + return str(cmd_samtools_fastq_minimap2(**args)) + + def test_input_ref_decodes_the_input(self): + """`input_ref` decodes the input file, not the target reference""" + cmd_str = self.build(input_ref=INPUT_REF) + assert f"samtools fastq --reference {INPUT_REF}" in cmd_str + assert f"samtools fastq --reference {REFERENCE}" not in cmd_str + # The target reference is still used for alignment and sorting + assert f"--reference {REFERENCE}" in cmd_str + + def test_no_input_ref(self): + """No `--reference` is passed to samtools without `input_ref`""" + cmd_str = self.build() + assert "--reference" not in cmd_str.split("sentieon minimap2")[0] + assert "samtools fastq -@ 4" in cmd_str + + def test_default_minimap2_model(self): + """The minimap2 model defaults to the bundle's minimap2.model""" + cmd_str = self.build() + assert f"-x {MODEL_BUNDLE}/minimap2.model" in cmd_str + + def test_explicit_minimap2_model(self): + """An explicit minimap2 model overrides the default""" + model = MODEL_BUNDLE.joinpath("minimap2_lr.model") + cmd_str = self.build(minimap2_model=model) + assert f"-x {model}" in cmd_str + assert f"-x {MODEL_BUNDLE}/minimap2.model" not in cmd_str + + def test_sm_backfill(self): + """A missing SM tag is backfilled with the sample name""" + cmd_str = self.build(rg_lines=["@RG\tID:rg1"], sample_name="sample1") + assert "addreplacerg" in cmd_str + assert "SM:sample1" in cmd_str + + def test_sm_not_overwritten(self): + """An existing SM tag is left alone""" + cmd_str = self.build( + rg_lines=["@RG\tID:rg1\tSM:from_header"], + sample_name="sample1", + ) + assert "SM:from_header" in cmd_str + assert "SM:sample1" not in cmd_str + + def test_one_addreplacerg_per_readgroup(self): + """Each readgroup gets its own addreplacerg command""" + cmd_str = self.build( + rg_lines=["@RG\tID:rg1\tSM:s1", "@RG\tID:rg2\tSM:s1"], + ) + assert cmd_str.count("samtools addreplacerg") == 2 + + def test_output_and_sort(self): + """The realigned reads are sorted into the output file""" + cmd_str = self.build() + assert f"-o {OUT_ALN}" in cmd_str + assert "sentieon util sort" in cmd_str + assert "--sam2bam" in cmd_str + + +class TestSamtoolsFastqBwa: + """Test the bwa re-alignment command""" + + def build(self, **kwargs) -> str: + args = { + "out_aln": OUT_ALN, + "input_aln": INPUT_ALN, + "reference": REFERENCE, + "model_bundle": MODEL_BUNDLE, + "cores": 4, + "rg_header": RG_HEADER, + } + args.update(kwargs) + return str(cmd_samtools_fastq_bwa(**args)) + + def test_input_ref_without_collate(self): + """`input_ref` decodes the input file read by samtools fastq""" + cmd_str = self.build(input_ref=INPUT_REF) + assert f"samtools fastq --reference {INPUT_REF}" in cmd_str + assert "samtools collate" not in cmd_str + + def test_input_ref_with_collate(self): + """`input_ref` decodes the input file read by samtools collate""" + cmd_str = self.build(input_ref=INPUT_REF, collate=True) + assert f"samtools collate --reference {INPUT_REF}" in cmd_str + # The fastq stage reads the collated stream, not the input file + assert f"samtools fastq --reference {INPUT_REF}" not in cmd_str + + def test_no_input_ref(self): + """No `--reference` is passed to samtools without `input_ref`""" + cmd_str = self.build() + assert "--reference" not in cmd_str.split("sentieon bwa")[0] + assert "samtools fastq -@ 4" in cmd_str + + def test_bwa_model_and_header(self): + """bwa uses the bundle model and the readgroup header file""" + cmd_str = self.build() + assert f"-x {MODEL_BUNDLE}/bwa.model" in cmd_str + assert f"-H {RG_HEADER}" in cmd_str diff --git a/tests/unit/test_dag.py b/tests/unit/test_dag.py new file mode 100644 index 0000000..014f945 --- /dev/null +++ b/tests/unit/test_dag.py @@ -0,0 +1,183 @@ +""" +Unit tests for DAG traversal and duplicate detection. +""" + +import os +import sys + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.exceptions import DagExecutionError # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + + +def _job(name, threads=1, resources=None): + """A trivial job whose identity is its (unique) command name.""" + return Job( + Pipeline(Command(name)), name, threads, resources, task_name="test" + ) + + +def test_job_identity_is_the_pipeline(): + # Two jobs with the same pipeline are equal even if named + # differently; a DAG deduplicates on this identity. + j1 = Job(Pipeline(Command("echo", "x")), "name-a", task_name="test") + j2 = Job(Pipeline(Command("echo", "x")), "name-b", task_name="test") + j3 = Job(Pipeline(Command("echo", "y")), "name-c", task_name="test") + assert j1 == j2 + assert hash(j1) == hash(j2) + assert j1 != j3 + + +class TestDAGTraversal: + """Test DAG bookkeeping through mark_finished""" + + def test_add_job_with_unknown_dependency_raises(self): + dag = DAG() + known = _job("known") + dag.add_job(known) + + unknown = _job("unknown") + with pytest.raises(ValueError): + dag.add_job(_job("downstream"), {unknown}) + + def test_finishing_an_unready_job_raises(self): + dag = DAG() + a = _job("a") + b = _job("b") + dag.add_job(a) + dag.add_job(b, {a}) + + with pytest.raises(DagExecutionError): + dag.mark_finished(b) # b is not ready + + def test_diamond_dependencies(self): + dag = DAG() + top = _job("top") + left = _job("left") + right = _job("right") + bottom = _job("bottom") + + dag.add_job(top) + dag.add_job(left, {top}) + dag.add_job(right, {top}) + dag.add_job(bottom, {left, right}) + + assert set(dag.ready_jobs) == {top} + + assert set(dag.mark_finished(top)) == {left, right} + + # bottom needs BOTH left and right + assert set(dag.mark_finished(left)) == set() + assert set(dag.mark_finished(right)) == {bottom} + + def test_multilevel_chain(self): + dag = DAG() + a = _job("a") + b = _job("b") + c = _job("c") + + dag.add_job(a) + dag.add_job(b, {a}) + dag.add_job(c, {b}) + + assert set(dag.ready_jobs) == {a} + assert set(dag.mark_finished(a)) == {b} + assert set(dag.mark_finished(b)) == {c} + + # Everything has been consumed and marked finished. + dag.mark_finished(c) + assert a in dag.finished_jobs + assert b in dag.finished_jobs + assert c in dag.finished_jobs + assert len(dag.waiting_jobs) == 0 + assert len(dag.ready_jobs) == 0 + + def test_add_duplicate_pipeline_raises(self): + dag = DAG() + dag.add_job( + Job(Pipeline(Command("echo", "x")), "step-1", task_name="test") + ) + # A second job with an identical pipeline collides on identity. + with pytest.raises(ValueError): + dag.add_job( + Job(Pipeline(Command("echo", "x")), "step-2", task_name="test") + ) + + def test_jobs_differing_only_in_exec_kwargs_are_not_duplicates(self): + # A command's env/cwd is part of its identity, so two jobs that + # differ only there are distinct, not rejected as duplicates. + dag = DAG() + j1 = Job( + Pipeline(Command("run", exec_kwargs={"cwd": "/a"})), + "a", + task_name="test", + ) + j2 = Job( + Pipeline(Command("run", exec_kwargs={"cwd": "/b"})), + "b", + task_name="test", + ) + dag.add_job(j1) + dag.add_job(j2) # must not raise "already in the DAG" + assert j1 in dag.ready_jobs + assert j2 in dag.ready_jobs + + def test_add_job_accepts_list_dependencies(self): + # Regression: non-set dependencies used to be silently dropped, + # leaving the job immediately ready. + dag = DAG() + a = _job("a") + b = _job("b") + dag.add_job(a) + dag.add_job(b, [a]) + + assert b not in dag.ready_jobs + assert dag.waiting_jobs[b] == {a} + assert set(dag.mark_finished(a)) == {b} + + def test_add_job_accepts_any_iterable_dependencies(self): + # Tuples, frozensets, and single-use generators all work; the + # generator case also guards against consuming the iterable twice. + for make in (tuple, frozenset, iter): + dag = DAG() + a = _job("a") + b = _job("b") + dag.add_job(a) + dag.add_job(b, make([a])) + assert dag.waiting_jobs[b] == {a}, make + + def test_add_job_rejects_unknown_dependency_in_list(self): + dag = DAG() + dag.add_job(_job("known")) + with pytest.raises(ValueError): + dag.add_job(_job("downstream"), [_job("unknown")]) + + def test_add_job_with_empty_iterable_is_ready(self): + for deps in ([], iter([])): + dag = DAG() + job = _job("solo") + dag.add_job(job, deps) + assert job in dag.ready_jobs + + def test_stored_dependencies_are_independent_of_caller_set(self): + dag = DAG() + a = _job("a") + b = _job("b") + c = _job("c") + dag.add_job(a) + dag.add_job(b) + + deps = {a} + dag.add_job(c, deps) + deps.add(b) # mutating the caller's set must not affect the DAG + + assert dag.waiting_jobs[c] == {a} diff --git a/tests/unit/test_dag_construction.py b/tests/unit/test_dag_construction.py index c42159f..b7faa77 100644 --- a/tests/unit/test_dag_construction.py +++ b/tests/unit/test_dag_construction.py @@ -5,6 +5,7 @@ import json import pathlib import pytest +import re import tempfile import shlex import sys @@ -20,6 +21,17 @@ from sentieon_cli.job import Job from sentieon_cli.shell_pipeline import Pipeline, Command +# task_name becomes a log directory name, so it must be kebab-case. +TASK_NAME_RE = re.compile(r"[a-z0-9]+(-[a-z0-9]+)*") + + +def assert_task_names(dag): + """Every job in a built DAG carries a kebab-case task_name""" + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) + assert all_jobs + for job in all_jobs: + assert TASK_NAME_RE.fullmatch(job.task_name), job.name + class TestDAGConstruction: """Test DAG construction for pipelines""" @@ -54,10 +66,10 @@ def create_basic_dnascope_pipeline(self): pipeline.reference = self.mock_ref pipeline.model_bundle = self.mock_bundle pipeline.sample_input = [self.mock_bam] - pipeline.sr_r1_fastq = [] - pipeline.sr_r2_fastq = [] - pipeline.sr_readgroups = [] - pipeline.sr_duplicate_marking = "markdup" + pipeline.r1_fastq = [] + pipeline.r2_fastq = [] + pipeline.readgroups = [] + pipeline.duplicate_marking = "markdup" pipeline.cores = 2 pipeline.dry_run = True # Important: use dry run to avoid actual execution pipeline.skip_version_check = True @@ -154,7 +166,7 @@ def test_skip_small_variants_dag(self, mock_lib_preloaded): job_names = [job.name for job in all_jobs] # Should not have variant calling jobs - variant_job_keywords = ["variant-calling", "model-apply", "gvcftyper"] + variant_job_keywords = ["dnascope", "model-apply", "gvcftyper"] for keyword in variant_job_keywords: assert not any(keyword in name for name in job_names) @@ -187,6 +199,18 @@ def test_gvcf_mode_dag_differences(self, mock_lib_preloaded): assert not any("gvcftyper" in name for name in job_names1) assert any("gvcftyper" in name for name in job_names2) + @patch('sentieon_cli.util.library_preloaded') + def test_jobs_have_task_names(self, mock_lib_preloaded): + """Every job built by the pipeline groups under a task_name""" + mock_lib_preloaded.return_value = True + + pipeline = self.create_basic_dnascope_pipeline() + pipeline.gvcf = True + pipeline.validate() + pipeline.configure() + + assert_task_names(pipeline.build_dag()) + class TestDAGJobProperties: """Test properties of individual jobs in the DAG""" @@ -198,6 +222,7 @@ def test_job_thread_allocation(self): Pipeline(Command(*shlex.split("sentieon driver --algo DNAscope"))), "variant-calling", 8, + task_name="variant-calling", ) assert job.threads == 8 @@ -206,6 +231,7 @@ def test_job_thread_allocation(self): Pipeline(Command("rm", "temp_file.vcf")), "cleanup", 0, + task_name="cleanup", ) assert job.threads == 0 @@ -217,6 +243,7 @@ def test_job_resource_requirements(self): "alignment", 4, resources={"node0": 1}, + task_name="alignment", ) assert job.resources == {"node0": 1} @@ -225,6 +252,7 @@ def test_job_resource_requirements(self): Pipeline(Command("simple-command")), "simple", 2, + task_name="test", ) assert job.resources == {} @@ -235,6 +263,7 @@ def test_job_failure_tolerance(self): Pipeline(Command("rm", "temp_files", fail_ok=True)), "cleanup", 0, + task_name="cleanup", ) assert job.shell.nodes[0].fail_ok is True @@ -243,6 +272,7 @@ def test_job_failure_tolerance(self): Pipeline(Command("sentieon", "driver")), "variant-calling", 4, + task_name="variant-calling", ) assert job.shell.nodes[0].fail_ok is False @@ -328,6 +358,9 @@ def test_longread_dag_complexity(self, mock_lib_preloaded): # Should have multiple phases of variant calling assert total_jobs > 10 # Rough estimate for complex phased calling + # Every job groups under a task_name, phasing jobs included + assert_task_names(dag) + @patch('sentieon_cli.util.library_preloaded') def test_ont_vs_hifi_dag_differences(self, mock_lib_preloaded): """Test DAG differences between ONT and HiFi technologies""" @@ -393,9 +426,9 @@ def test_dag_has_no_cycles(self): dag = DAG() # Create test jobs - job1 = Job(Pipeline(Command("command1")), "job1") - job2 = Job(Pipeline(Command("command2")), "job2") - job3 = Job(Pipeline(Command("command3")), "job3") + job1 = Job(Pipeline(Command("command1")), "job1", task_name="test") + job2 = Job(Pipeline(Command("command2")), "job2", task_name="test") + job3 = Job(Pipeline(Command("command3")), "job3", task_name="test") # Add jobs with dependencies: job1 -> job2 -> job3 dag.add_job(job1) @@ -416,9 +449,9 @@ def test_dag_execution_simulation(self): dag = DAG() # Create test jobs - job1 = Job(Pipeline(Command("command1")), "job1") - job2 = Job(Pipeline(Command("command2")), "job2") - job3 = Job(Pipeline(Command("command3")), "job3") + job1 = Job(Pipeline(Command("command1")), "job1", task_name="test") + job2 = Job(Pipeline(Command("command2")), "job2", task_name="test") + job3 = Job(Pipeline(Command("command3")), "job3", task_name="test") # Add jobs: job1 and job2 can run in parallel, then job3 dag.add_job(job1) @@ -426,8 +459,7 @@ def test_dag_execution_simulation(self): dag.add_job(job3, {job1, job2}) # Simulate execution - dag_gen = dag.update_dag() - ready_jobs = dag_gen.send(None) + ready_jobs = dag.ready_jobs # Initially job1 and job2 should be ready assert job1 in ready_jobs @@ -435,11 +467,11 @@ def test_dag_execution_simulation(self): assert job3 not in ready_jobs # Complete job1 - new_ready = dag_gen.send(job1) + new_ready = dag.mark_finished(job1) assert len(new_ready) == 0 # job3 still waiting for job2 # Complete job2 - new_ready = dag_gen.send(job2) + new_ready = dag.mark_finished(job2) assert job3 in new_ready # Now job3 is ready diff --git a/tests/unit/test_dnascope_hybrid_cnv.py b/tests/unit/test_dnascope_hybrid_cnv.py new file mode 100644 index 0000000..99287e4 --- /dev/null +++ b/tests/unit/test_dnascope_hybrid_cnv.py @@ -0,0 +1,230 @@ +""" +Unit tests for sex-aware CNV calling in the DNAscope hybrid pipeline +""" + +import os +import pathlib +import sys +import tempfile +from unittest.mock import MagicMock, patch + +import pytest + +# Add the parent directory to the path to import sentieon_cli +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline +from sentieon_cli.util import SampleSex + + +class TestDNAscopeHybridCnv: + """CNV calling runs in a second DAG, after the sample sex is known""" + + def setup_method(self): + self.temp_dir = tempfile.mkdtemp() + self.mock_dir = pathlib.Path(self.temp_dir) + + self.mock_vcf = self.mock_dir / "output.vcf.gz" + self.mock_ref = self.mock_dir / "reference.fa" + self.mock_sr_aln = [self.mock_dir / "short.bam"] + self.mock_lr_aln = [self.mock_dir / "long.bam"] + self.mock_bundle = self.mock_dir / "model.bundle" + self.mock_bed = self.mock_dir / "interval.bed" + self.mock_par_bed = self.mock_dir / "par.bed" + + for file_path in [ + self.mock_ref, + self.mock_sr_aln[0], + self.mock_lr_aln[0], + self.mock_bundle, + self.mock_bed, + self.mock_par_bed, + ]: + file_path.touch() + + def create_pipeline(self): + """Create a DNAscopeHybridPipeline for testing""" + with patch("sys.exit"): + pipeline = DNAscopeHybridPipeline() + + pipeline.logger = MagicMock() + + pipeline.output_vcf = self.mock_vcf + pipeline.reference = self.mock_ref + pipeline.model_bundle = self.mock_bundle + pipeline.bed = self.mock_bed + pipeline.cores = 2 + pipeline.dry_run = True + pipeline.skip_version_check = True + pipeline.tmp_dir = self.mock_dir + pipeline.pop_vcf = None + + # State normally set by validate() + pipeline.fai_data = {"chr1": {"length": 1000}} + pipeline.shards = [MagicMock()] + pipeline.shards[0].contig = "chr1" + pipeline.shards[0].start = 1 + pipeline.shards[0].stop = 1000 + pipeline.pop_vcf_contigs = {"chr1": 1000} + pipeline.lr_aln_readgroups = [[{"ID": "lr_rg1", "SM": "sample1"}]] + pipeline.sr_aln_readgroups = [[{"ID": "sr_rg1", "SM": "sample1"}]] + pipeline.hybrid_rg_sm = "sample1" + pipeline.hybrid_set_rg = False + pipeline.shortread_tech = "Illumina" + pipeline.longread_tech = "ONT" + pipeline.sr_aln = self.mock_sr_aln + pipeline.lr_aln = self.mock_lr_aln + pipeline.skip_cnv = False + + return pipeline + + def build_dag(self, pipeline): + """Build the first DAG and index the jobs by name""" + dag = pipeline.build_dag() + jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) + return dag, {job.name: job for job in jobs} + + def build_second_dag(self, pipeline): + """Build the second DAG and index the jobs by name""" + dag = pipeline.build_second_dag() + if dag is None: + return None, {} + jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) + return dag, {job.name: job for job in jobs} + + def test_cnv_jobs_run_in_the_second_dag(self): + pipeline = self.create_pipeline() + pipeline.sample_sex = SampleSex.FEMALE + _dag, jobs = self.build_dag(pipeline) + + assert "cnvscope" not in jobs + assert "cnv-model-apply" not in jobs + + second_dag, second_jobs = self.build_second_dag(pipeline) + assert "cnvscope" in second_jobs + assert "cnv-model-apply" in second_jobs + assert second_dag.waiting_jobs[second_jobs["cnv-model-apply"]] == { + second_jobs["cnvscope"] + } + + def test_the_first_dag_estimates_the_sample_ploidy(self): + pipeline = self.create_pipeline() + _dag, jobs = self.build_dag(pipeline) + + assert "estimate-ploidy" in jobs + cmd_str = str(jobs["estimate-ploidy"].shell) + assert "estimate_ploidy.py" in cmd_str + # Ploidy is estimated from the deduped short-read alignment + assert "output_deduped.cram" in cmd_str + assert str(pipeline.ploidy_json).endswith("output_ploidy.json") + # The chr-prefixed script defaults apply to this reference + assert "--contigs" not in cmd_str + assert "--x_contig" not in cmd_str + + def test_the_ploidy_job_runs_with_sample_sex(self): + # `--sample_sex` overrides the estimate, which still runs + pipeline = self.create_pipeline() + pipeline.sample_sex = SampleSex.FEMALE + _dag, jobs = self.build_dag(pipeline) + + assert "estimate-ploidy" in jobs + + def test_ploidy_contigs_follow_the_reference_build(self): + pipeline = self.create_pipeline() + pipeline.reference_build = "b37" + _dag, jobs = self.build_dag(pipeline) + + cmd_str = str(jobs["estimate-ploidy"].shell) + assert "--contigs 1 2 " in cmd_str + assert "--autosomes 1 2 " in cmd_str + assert "--x_contig X" in cmd_str + assert "--y_contig Y" in cmd_str + + def test_the_ploidy_job_runs_when_cnvs_are_skipped(self): + # The ploidy JSON is written by every run, but there is no + # second DAG without CNV calling + pipeline = self.create_pipeline() + pipeline.skip_cnv = True + _dag, jobs = self.build_dag(pipeline) + + assert "estimate-ploidy" in jobs + assert self.build_second_dag(pipeline)[0] is None + + def test_a_female_sample_is_called_without_a_par_bed(self): + pipeline = self.create_pipeline() + pipeline.sample_sex = SampleSex.FEMALE + self.build_dag(pipeline) + _dag, jobs = self.build_second_dag(pipeline) + + cmd_str = str(jobs["cnvscope"].shell) + assert "--sex F" in cmd_str + assert "--par" not in cmd_str + + def test_a_male_sample_is_called_with_the_par_bed(self): + pipeline = self.create_pipeline() + pipeline.sample_sex = SampleSex.MALE + pipeline.cnv_par_bed = self.mock_par_bed + self.build_dag(pipeline) + _dag, jobs = self.build_second_dag(pipeline) + + cmd_str = str(jobs["cnvscope"].shell) + assert "--sex M" in cmd_str + assert f"--par {self.mock_par_bed}" in cmd_str + + def test_cnvscope_dumps_the_probes_next_to_the_cnv_vcf(self): + pipeline = self.create_pipeline() + pipeline.sample_sex = SampleSex.FEMALE + self.build_dag(pipeline) + _dag, jobs = self.build_second_dag(pipeline) + + probes = self.mock_dir / "output.cnv.probes" + assert f"--dump_probes {probes}" in str(jobs["cnvscope"].shell) + assert "--dump_probes" not in str(jobs["cnv-model-apply"].shell) + + def test_a_dry_run_assumes_a_male_sample(self): + # Validation has already confirmed a PAR BED file is available + pipeline = self.create_pipeline() + pipeline.cnv_par_bed = self.mock_par_bed + self.build_dag(pipeline) + _dag, jobs = self.build_second_dag(pipeline) + + assert pipeline.sample_sex is SampleSex.MALE + cmd_str = str(jobs["cnvscope"].shell) + assert "--sex M" in cmd_str + assert f"--par {self.mock_par_bed}" in cmd_str + + def _validate_cnv(self, pipeline): + """Run the CNV validation with the reference index parsed""" + pipeline.fai_data = {"chr1": {"length": 1000}} + pipeline.validate_cnv() + + @pytest.mark.parametrize( + "sample_sex", [None, SampleSex.MALE, SampleSex.FEMALE] + ) + def test_validation_requires_a_par_bed_for_cnv_calling(self, sample_sex): + # The reference build is not recognized, so no packaged PAR BED + # file can be selected and the run stops during validation + pipeline = self.create_pipeline() + pipeline.sample_sex = sample_sex + + with pytest.raises(SystemExit) as excinfo: + self._validate_cnv(pipeline) + assert excinfo.value.code == 2 + + def test_validation_accepts_the_par_bed_argument(self): + pipeline = self.create_pipeline() + pipeline.par_bed = self.mock_par_bed + + self._validate_cnv(pipeline) # no SystemExit + + assert pipeline.cnv_par_bed == self.mock_par_bed + + def test_validation_needs_no_par_bed_when_cnvs_are_skipped(self): + pipeline = self.create_pipeline() + pipeline.skip_cnv = True + + self._validate_cnv(pipeline) # no SystemExit + + assert pipeline.cnv_par_bed is None diff --git a/tests/unit/test_dnascope_hybrid_longread.py b/tests/unit/test_dnascope_hybrid_longread.py new file mode 100644 index 0000000..7807977 --- /dev/null +++ b/tests/unit/test_dnascope_hybrid_longread.py @@ -0,0 +1,197 @@ +""" +Unit tests for the DNAscope hybrid pipeline's long-read jobs + +The hybrid pipeline used to inherit its long-read realignment, mosdepth and +SV calling from `DNAscopeLRPipeline`; it now composes the same stages +itself, with its own argument defaults. +""" + +import pathlib +import tempfile +from abc import ABC +from unittest.mock import MagicMock, patch + +import pytest + +from sentieon_cli.dnascope import DNAscopePipeline +from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline +from sentieon_cli.dnascope_longread import DNAscopeLRPipeline +from sentieon_cli.pipeline import BasePipeline + +RG_LINES = ["@RG\tID:lr_rg1\tSM:sample1"] + + +class TestDNAscopeHybridLongRead: + """The long-read half of the hybrid DAG""" + + def setup_method(self): + self.temp_dir = tempfile.mkdtemp() + self.mock_dir = pathlib.Path(self.temp_dir) + + self.mock_vcf = self.mock_dir / "output.vcf.gz" + self.mock_ref = self.mock_dir / "reference.fa" + self.mock_input_ref = self.mock_dir / "input_reference.fa" + self.mock_sr_aln = [self.mock_dir / "short.bam"] + self.mock_lr_aln = [self.mock_dir / "long.bam"] + self.mock_bundle = self.mock_dir / "model.bundle" + self.mock_bed = self.mock_dir / "interval.bed" + + for file_path in [ + self.mock_ref, + self.mock_input_ref, + self.mock_sr_aln[0], + self.mock_lr_aln[0], + self.mock_bundle, + self.mock_bed, + ]: + file_path.touch() + + def create_pipeline(self): + """Create a DNAscopeHybridPipeline that realigns its long reads""" + with patch("sys.exit"): + pipeline = DNAscopeHybridPipeline() + + pipeline.logger = MagicMock() + + pipeline.output_vcf = self.mock_vcf + pipeline.reference = self.mock_ref + pipeline.model_bundle = self.mock_bundle + pipeline.bed = self.mock_bed + pipeline.cores = 2 + pipeline.dry_run = True + pipeline.skip_version_check = True + pipeline.tmp_dir = self.mock_dir + pipeline.pop_vcf = None + + # State normally set by validate() + pipeline.fai_data = {"chr1": {"length": 1000}} + pipeline.shards = [MagicMock()] + pipeline.shards[0].contig = "chr1" + pipeline.shards[0].start = 1 + pipeline.shards[0].stop = 1000 + pipeline.pop_vcf_contigs = {"chr1": 1000} + pipeline.lr_aln_readgroups = [[{"ID": "lr_rg1", "SM": "sample1"}]] + pipeline.sr_aln_readgroups = [[{"ID": "sr_rg1", "SM": "sample1"}]] + pipeline.hybrid_rg_sm = "sample1" + pipeline.hybrid_set_rg = False + pipeline.shortread_tech = "Illumina" + pipeline.longread_tech = "ONT" + pipeline.sr_aln = self.mock_sr_aln + pipeline.lr_aln = self.mock_lr_aln + + # Realign the long reads + pipeline.lr_align_input = True + pipeline.lr_input_ref = self.mock_input_ref + + return pipeline + + def build_dag(self, pipeline): + """Build the DAG and index the jobs by name""" + with patch( + "sentieon_cli.command_strings.get_rg_lines", return_value=RG_LINES + ): + dag = pipeline.build_dag() + jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) + return dag, {job.name: job for job in jobs} + + def dep_names(self, dag, job): + """The names of a job's dependencies""" + return {dep.name for dep in dag.waiting_jobs.get(job, set())} + + def test_minimap2_args_default(self): + """The hybrid pipeline aligns with `-YL`, like the long-read one + + `-L` moves CIGAR strings longer than 65535 operations - which ONT + reads can produce and BAM cannot store - into the `CG` tag, so + every minimap2 call that may see ONT input needs it. + """ + pipeline = self.create_pipeline() + _dag, jobs = self.build_dag(pipeline) + + assert "bam-realign-0" in jobs + cmd_str = str(jobs["bam-realign-0"].shell) + assert " -YL " in cmd_str + + def test_minimap2_args_match_the_long_read_pipeline(self): + """The hybrid and long-read defaults agree""" + assert ( + DNAscopeHybridPipeline().minimap2_args + == DNAscopeLRPipeline().minimap2_args + ) + + def test_lr_fastq_taglist_is_used(self): + """`--lr_fastq_taglist` reaches the minimap2 command""" + pipeline = self.create_pipeline() + pipeline.lr_fastq_taglist = "RG,MM" + _dag, jobs = self.build_dag(pipeline) + + assert "-T RG,MM" in str(jobs["bam-realign-0"].shell) + + def test_long_read_jobs_wait_for_realignment(self): + """mosdepth and LongReadSV read the realigned long reads""" + pipeline = self.create_pipeline() + dag, jobs = self.build_dag(pipeline) + + assert self.dep_names(dag, jobs["mosdepth-0"]) == {"bam-realign-0"} + assert self.dep_names(dag, jobs["LongReadSV"]) == {"bam-realign-0"} + + def test_realigned_long_reads_are_called(self): + """The realigned alignment, not the input, is passed downstream""" + pipeline = self.create_pipeline() + _dag, jobs = self.build_dag(pipeline) + + realigned = str(self.mock_vcf).replace(".vcf.gz", "_mm2_sorted_0.cram") + for name in ("mosdepth-0", "LongReadSV", "dnascope-1"): + assert realigned in str(jobs[name].shell) + + +class TestDNAscopeHybridBases: + """The hybrid pipeline no longer inherits from the other pipelines""" + + def test_mro(self): + assert DNAscopeHybridPipeline.__mro__ == ( + DNAscopeHybridPipeline, + BasePipeline, + ABC, + object, + ) + + def test_not_a_dnascope_or_longread_pipeline(self): + with patch("sys.exit"): + pipeline = DNAscopeHybridPipeline() + + assert not isinstance(pipeline, DNAscopePipeline) + assert not isinstance(pipeline, DNAscopeLRPipeline) + + @pytest.mark.parametrize( + "attribute", + [ + "sample_input", + "r1_fastq", + "r2_fastq", + "readgroups", + "duplicate_marking", + "align", + "collate_align", + "consensus", + "fastq", + "fastq_taglist", + "haploid_bed", + "input_ref", + "interval_padding", + "no_ramdisk", + "pcr_free", + "repeat_model", + "skip_small_variants", + "tech", + "use_pbsv", + "cnv_excluded_regions", + "assay", + ], + ) + def test_parent_only_attributes_are_gone(self, attribute): + """The attributes that came from the parent `__init__`s""" + with patch("sys.exit"): + pipeline = DNAscopeHybridPipeline() + + assert not hasattr(pipeline, attribute) diff --git a/tests/unit/test_dnascope_hybrid_pop_vcf.py b/tests/unit/test_dnascope_hybrid_pop_vcf.py index 7add9f1..b7be312 100644 --- a/tests/unit/test_dnascope_hybrid_pop_vcf.py +++ b/tests/unit/test_dnascope_hybrid_pop_vcf.py @@ -135,8 +135,7 @@ def test_transfer_jobs_creation(self): pipeline.pop_vcf = self.mock_pop_vcf # Build DAG - with patch("sentieon_cli.dnascope_hybrid.check_version", return_value=True): - dag = pipeline.build_dag() + dag = pipeline.build_dag() # Verify transfer jobs exist job_names = [job.name for job in dag.waiting_jobs] @@ -150,8 +149,7 @@ def test_model_apply_input_with_pop_vcf(self): pipeline = self.create_pipeline() pipeline.pop_vcf = self.mock_pop_vcf - with patch("sentieon_cli.dnascope_hybrid.check_version", return_value=True): - dag = pipeline.build_dag() + dag = pipeline.build_dag() # Find model-apply job apply_job = None diff --git a/tests/unit/test_dnascope_hybrid_stage1.py b/tests/unit/test_dnascope_hybrid_stage1.py new file mode 100644 index 0000000..2c8c366 --- /dev/null +++ b/tests/unit/test_dnascope_hybrid_stage1.py @@ -0,0 +1,206 @@ +""" +Unit tests for the DNAscope hybrid first-stage realignment jobs +""" + +import os +import pathlib +import sys +import tempfile +from unittest.mock import MagicMock, patch + +import packaging.version + +# Add the parent directory to the path to import sentieon_cli +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +from sentieon_cli.dnascope_hybrid import ( + CALLING_MIN_VERSIONS, + DNAscopeHybridPipeline, +) + + +class TestDNAscopeHybridStage1: + """Test the first-stage haplotype alignment jobs""" + + def setup_method(self): + """Setup test fixtures""" + self.temp_dir = tempfile.mkdtemp() + self.mock_dir = pathlib.Path(self.temp_dir) + + self.mock_vcf = self.mock_dir / "output.vcf.gz" + self.mock_ref = self.mock_dir / "reference.fa" + self.mock_sr_aln = [self.mock_dir / "short.bam"] + self.mock_lr_aln = [self.mock_dir / "long.bam"] + self.mock_bundle = self.mock_dir / "model.bundle" + self.mock_bed = self.mock_dir / "interval.bed" + + for file_path in [ + self.mock_ref, + self.mock_sr_aln[0], + self.mock_lr_aln[0], + self.mock_bundle, + self.mock_bed, + ]: + file_path.touch() + + self.stage1_fifo = self.mock_dir / "stage1_hap.fq" + self.stage1_hap_bam = self.mock_dir / "stage1_hap.bam" + + def create_pipeline(self): + """Create a DNAscopeHybridPipeline for testing""" + with patch("sys.exit"): + pipeline = DNAscopeHybridPipeline() + + pipeline.logger = MagicMock() + + # Configure arguments + pipeline.output_vcf = self.mock_vcf + pipeline.reference = self.mock_ref + pipeline.model_bundle = self.mock_bundle + pipeline.bed = self.mock_bed + pipeline.cores = 2 + pipeline.dry_run = True + pipeline.skip_version_check = True + pipeline.tmp_dir = self.mock_dir + pipeline.pop_vcf = None + + # State normally set by validate() + pipeline.fai_data = {"chr1": {"length": 1000}} + pipeline.shards = [MagicMock()] + pipeline.shards[0].contig = "chr1" + pipeline.shards[0].start = 1 + pipeline.shards[0].stop = 1000 + pipeline.pop_vcf_contigs = {"chr1": 1000} + pipeline.lr_aln_readgroups = [[{"ID": "lr_rg1", "SM": "sample1"}]] + pipeline.sr_aln_readgroups = [[{"ID": "sr_rg1", "SM": "sample1"}]] + pipeline.hybrid_rg_sm = "sample1" + pipeline.hybrid_set_rg = False + pipeline.shortread_tech = "Illumina" + pipeline.longread_tech = "ONT" + pipeline.sr_aln = self.mock_sr_aln + pipeline.lr_aln = self.mock_lr_aln + + return pipeline + + def build_dag(self, pipeline): + """Build the DAG and index the jobs by name""" + dag = pipeline.build_dag() + jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) + return dag, {job.name: job for job in jobs} + + def dep_names(self, dag, job): + """The names of a job's dependencies""" + return {dep.name for dep in dag.waiting_jobs.get(job, set())} + + def test_fifo_job(self): + """The haplotype fastq fifo is created before the stage1 jobs""" + pipeline = self.create_pipeline() + dag, jobs = self.build_dag(pipeline) + + assert "stage1-fifo" in jobs + cmd_str = str(jobs["stage1-fifo"].shell) + assert cmd_str == f"mkfifo {self.stage1_fifo}" + assert self.dep_names(dag, jobs["stage1-fifo"]) == set() + + def test_stage1_hap_command(self): + """The haplotype alignments are sorted into stage1_hap.bam""" + pipeline = self.create_pipeline() + _dag, jobs = self.build_dag(pipeline) + + assert "first-stage-hap" in jobs + hap_job = jobs["first-stage-hap"] + cmd_str = str(hap_job.shell) + assert "--algo HybridStage1" in cmd_str + assert "HybridStage1.model" in cmd_str + # The unsorted haplotype BAM is written to stdout and sorted + assert "--hap_bam -" in cmd_str + assert "sentieon util sort" in cmd_str + assert f"-o {self.stage1_hap_bam}" in cmd_str + # The driver writes BAM, not SAM, so no conversion is needed + assert "--sam2bam" not in cmd_str + # The fastq output goes to the fifo + assert str(self.stage1_fifo) in cmd_str + + # The job must run concurrently with the fifo reader + assert hap_job.threads == 0 + + def test_stage1_reads_the_fifo(self): + """The bwa job reads the haplotype fastq from the fifo""" + pipeline = self.create_pipeline() + _dag, jobs = self.build_dag(pipeline) + + cmd_str = str(jobs["first-stage"].shell) + # The fifo is a plain `cat` argument, read before the proc sub + assert cmd_str.startswith(f"cat {self.stage1_fifo} ") + assert f"<({self.stage1_fifo}" not in cmd_str + # The haplotype driver moved to its own job; only the insertion + # driver remains in the bwa pipeline + assert "--hap_bam" not in cmd_str + assert "HybridStage1.model" not in cmd_str + assert "HybridStage1_ins.model" in cmd_str + assert "sentieon bwa mem" in cmd_str + assert "--sam2bam" in cmd_str + + def test_stage1_dependencies(self): + """The stage1 jobs wait for the fifo and the merged BED""" + pipeline = self.create_pipeline() + dag, jobs = self.build_dag(pipeline) + + assert self.dep_names(dag, jobs["first-stage-hap"]) == { + "stage1-fifo", + "concat-merge-bed", + } + assert self.dep_names(dag, jobs["first-stage"]) == { + "stage1-fifo", + "concat-merge-bed", + } + # The second stage reads both stage1 outputs + assert self.dep_names(dag, jobs["second-stage"]) == { + "first-stage", + "first-stage-hap", + } + # The haplotype VCF removed here is written by the haplotype job + assert self.dep_names(dag, jobs["rm-tmp2"]) == { + "first-stage", + "first-stage-hap", + } + + def test_cleanup_jobs(self): + """Each cleanup job waits for the jobs writing what it removes""" + pipeline = self.create_pipeline() + dag, jobs = self.build_dag(pipeline) + + assert self.dep_names(dag, jobs["rm-tmp1"]) == {"concat-merge-bed"} + assert self.dep_names(dag, jobs["rm-tmp2"]) == { + "first-stage", + "first-stage-hap", + } + assert self.dep_names(dag, jobs["rm-tmp3"]) == {"second-stage"} + assert self.dep_names(dag, jobs["rm-tmp4"]) == {"third-stage"} + assert self.dep_names(dag, jobs["rm-tmp5"]) == {"concat-calls"} + + def test_retain_tmpdir_skips_cleanup(self): + """`--retain_tmpdir` keeps the intermediate files""" + pipeline = self.create_pipeline() + pipeline.retain_tmpdir = True + _dag, jobs = self.build_dag(pipeline) + + assert [name for name in jobs if name.startswith("rm-tmp")] == [] + + def test_second_stage_inputs(self): + """The second stage consumes the sorted haplotype BAM""" + pipeline = self.create_pipeline() + _dag, jobs = self.build_dag(pipeline) + + cmd_str = str(jobs["second-stage"].shell) + assert "--algo HybridStage2" in cmd_str + assert f"--input {self.stage1_hap_bam}" in cmd_str + assert f"--input {self.mock_dir / 'hybrid_stage1.bam'}" in cmd_str + + def test_driver_min_version(self): + """The unsorted `--hap_bam` output requires 202503.04""" + assert CALLING_MIN_VERSIONS[ + "sentieon driver" + ] >= packaging.version.Version("202503.04") diff --git a/tests/unit/test_dnascope_longread_dag.py b/tests/unit/test_dnascope_longread_dag.py new file mode 100644 index 0000000..e239146 --- /dev/null +++ b/tests/unit/test_dnascope_longread_dag.py @@ -0,0 +1,224 @@ +""" +The non-obvious edges of the DNAscope LongRead small-variant DAG +""" + +import json +import pathlib +from typing import Dict, List, Set +from unittest.mock import patch + +from sentieon_cli.dag import DAG +from sentieon_cli.dnascope_longread import DNAscopeLRPipeline +from sentieon_cli.job import Job + +from tests.utils.test_helpers import create_mock_args + +BUNDLE_MEMBERS = [ + "diploid_hp_model", + "diploid_model", + "diploid_model_unphased", + "gvcf_model", + "haploid_hp_model", + "haploid_model", + "longreadsv.model", + "minimap2.model", +] + +BUNDLE_INFO = { + "platform": "HiFi", + "minScriptVersion": "1.5.2", + "pipeline": "DNAscope LongRead", +} + + +def build_dag(tmp_path: pathlib.Path, **kwargs) -> DAG: + """Build the first DAG of a HiFi long-read run""" + reference = tmp_path / "reference.fa" + reference.touch() + (tmp_path / "reference.fa.fai").write_text( + "chr1\t1000\t10\t60\t61\nchr2\t2000\t1030\t60\t61\n" + ) + bam = tmp_path / "longread.bam" + bam.touch() + bundle = tmp_path / "model.bundle" + bundle.touch() + bed = tmp_path / "regions.bed" + bed.touch() + + pipeline = DNAscopeLRPipeline() + pipeline.setup_logging(create_mock_args()) + pipeline.output_vcf = tmp_path / "output.vcf.gz" + pipeline.reference = reference + pipeline.model_bundle = bundle + pipeline.sample_input = [bam] + pipeline.fastq = [] + pipeline.readgroups = [] + pipeline.bed = bed + pipeline.cores = 2 + pipeline.dry_run = True + pipeline.skip_version_check = True + pipeline.tmp_dir = tmp_path + # Only the small-variant jobs are under test + pipeline.skip_mosdepth = True + pipeline.skip_cnv = True + pipeline.skip_svs = True + for key, value in kwargs.items(): + setattr(pipeline, key, value) + + bundle_info = dict(BUNDLE_INFO) + if pipeline.pop_vcf: + bundle_info["SentieonVcfID"] = "test-pop-vcf" + pipeline.skip_pop_vcf_id_check = True + + with patch("sentieon_cli.dnascope_longread.ar_load") as mock_ar_load: + mock_ar_load.side_effect = [ + BUNDLE_MEMBERS, + json.dumps(bundle_info).encode(), + ] + pipeline.validate() + pipeline.configure() + return pipeline.build_dag() + + +def deps_of(dag: DAG, name: str) -> List[str]: + """The sorted dependency names of the single job called `name`""" + return sorted(job.name for job in dag.waiting_jobs[job_named(dag, name)]) + + +def job_named(dag: DAG, name: str) -> Job: + """The one job called `name`""" + matches = [job for job in all_jobs(dag) if job.name == name] + assert len(matches) == 1, f"{name}: {len(matches)} jobs" + return matches[0] + + +def all_jobs(dag: DAG) -> List[Job]: + return list(dag.waiting_jobs) + list(dag.ready_jobs) + + +def roots(dag: DAG) -> Set[str]: + """The names of the jobs with no dependencies""" + names = {job.name for job in dag.ready_jobs} + names.update( + job.name for job, deps in dag.waiting_jobs.items() if not deps + ) + return names + + +class TestSmallVariantEdges: + """The edges `add_small_variant_calling` wires by hand""" + + def test_the_calling_passes_follow_the_phaser(self, tmp_path): + dag = build_dag(tmp_path) + + assert deps_of(dag, "variantphaser") == ["model-apply-diploid"] + for hap in ("dnascope-hap1", "dnascope-hap2"): + assert deps_of(dag, hap) == ["repeatmodel", "variantphaser"] + assert deps_of(dag, "dnascope-unphased") == [ + "bedtools-subtract", + "repeatmodel", + ] + assert deps_of(dag, "patch") == ["dnascope-hap1", "dnascope-hap2"] + + def test_merge_waits_on_the_three_model_apply_jobs(self, tmp_path): + dag = build_dag(tmp_path) + + assert deps_of(dag, "merge") == [ + "model-apply-hap1", + "model-apply-hap2", + "model-apply-unphased", + ] + + def test_each_haplotype_apply_waits_on_its_own_transfer(self, tmp_path): + pop_vcf = tmp_path / "population.vcf.gz" + pop_vcf.touch() + dag = build_dag(tmp_path, pop_vcf=pop_vcf) + + # A model-apply reads only its own haplotype's transfer, so it + # waits on that haplotype's concat alone + for i in (1, 2): + assert deps_of(dag, f"model-apply-hap{i}") == [ + f"merge-trim-hap{i}-concat", + ] + + def test_the_haplotype_applies_wait_on_the_patch_without_a_pop_vcf( + self, tmp_path + ): + dag = build_dag(tmp_path) + + for hap in ("model-apply-hap1", "model-apply-hap2"): + assert deps_of(dag, hap) == ["patch"] + + def test_each_transfer_is_tagged(self, tmp_path): + pop_vcf = tmp_path / "population.vcf.gz" + pop_vcf.touch() + dag = build_dag(tmp_path, pop_vcf=pop_vcf) + + concats = { + job.name for job in all_jobs(dag) if job.name.endswith("-concat") + } + assert concats == { + "merge-trim-diploid-concat", + "merge-trim-hap1-concat", + "merge-trim-hap2-concat", + "merge-trim-unphased-concat", + } + # With a transfer, the apply waits on the concat alone -- the + # concat already depends on the raw-VCF producer + assert deps_of(dag, "model-apply-diploid") == [ + "merge-trim-diploid-concat", + ] + assert deps_of(dag, "model-apply-unphased") == [ + "merge-trim-unphased-concat", + ] + + def test_fai_to_bed_is_a_root_without_a_bed(self, tmp_path): + dag = build_dag(tmp_path, bed=None) + + assert "fai-to-bed" in roots(dag) + assert deps_of(dag, "bedtools-subtract") == [ + "fai-to-bed", + "variantphaser", + ] + + def test_haploid_calling_only_follows_the_repeat_model(self, tmp_path): + haploid_bed = tmp_path / "haploid.bed" + haploid_bed.touch() + dag = build_dag(tmp_path, haploid_bed=haploid_bed) + + # It does not wait on the alignment jobs, only on the repeat model + assert deps_of(dag, "dnascope-haploid") == ["repeatmodel"] + assert deps_of(dag, "haploid-patch2") == ["dnascope-haploid"] + assert deps_of(dag, "haploid-diploid-concat") == [ + "haploid-patch2", + "merge", + ] + + def test_haploid_calling_is_a_root_with_a_repeat_model(self, tmp_path): + haploid_bed = tmp_path / "haploid.bed" + haploid_bed.touch() + repeat_model = tmp_path / "given_repeat.model" + repeat_model.touch() + dag = build_dag( + tmp_path, haploid_bed=haploid_bed, repeat_model=repeat_model + ) + + assert "repeatmodel" not in {job.name for job in all_jobs(dag)} + assert "dnascope-haploid" in roots(dag) + + +class TestJobIdentity: + """Job names and task names of the migrated stages""" + + def test_task_names(self, tmp_path): + dag = build_dag(tmp_path) + + tasks: Dict[str, str] = { + job.name: job.task_name for job in all_jobs(dag) + } + assert tasks["dnascope-diploid"] == "variant-calling" + assert tasks["dnascope-hap1"] == "variant-calling" + assert tasks["dnascope-unphased"] == "variant-calling" + assert tasks["model-apply-diploid"] == "model-apply" + assert tasks["model-apply-hap2"] == "model-apply" + assert tasks["repeatmodel"] == "repeat-model" diff --git a/tests/unit/test_dnascope_longread_pop_vcf.py b/tests/unit/test_dnascope_longread_pop_vcf.py index bdcfd6c..b96e613 100644 --- a/tests/unit/test_dnascope_longread_pop_vcf.py +++ b/tests/unit/test_dnascope_longread_pop_vcf.py @@ -75,7 +75,6 @@ def create_pipeline(self): # Mock sample inputs pipeline.sample_input = self.mock_aln - pipeline.lr_aln = self.mock_aln return pipeline @@ -154,8 +153,7 @@ def test_transfer_jobs_creation(self): pipeline.pop_vcf = self.mock_pop_vcf # Mock methods to avoid real execution/checking - with patch("sentieon_cli.dnascope_longread.check_version", return_value=True), \ - patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.lr_align_inputs", return_value=([], set())), \ + with patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.lr_align_inputs", return_value=([], set())), \ patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.lr_align_fastq", return_value=([], set())), \ patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.mosdepth", return_value=set()), \ patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.merge_input_files", return_value=(pathlib.Path("merged.bam"), MagicMock())), \ @@ -170,15 +168,14 @@ def test_transfer_jobs_creation(self): # Expect merge-trim jobs assert any("merge-trim" in name for name in job_names) - assert "merge-trim-concat" in job_names + assert "merge-trim-diploid-concat" in job_names def test_model_apply_dependency(self): """Test that model apply uses the transfer output when pop_vcf is present""" pipeline = self.create_pipeline() pipeline.pop_vcf = self.mock_pop_vcf - with patch("sentieon_cli.dnascope_longread.check_version", return_value=True), \ - patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.lr_align_inputs", return_value=([], set())), \ + with patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.lr_align_inputs", return_value=([], set())), \ patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.lr_align_fastq", return_value=([], set())), \ patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.mosdepth", return_value=set()), \ patch("sentieon_cli.dnascope_longread.DNAscopeLRPipeline.merge_input_files", return_value=(pathlib.Path("merged.bam"), MagicMock())), \ @@ -188,16 +185,16 @@ def test_model_apply_dependency(self): dag = pipeline.build_dag() - # Find first-modelapply job + # Find the diploid model-apply job apply_job = None for job in dag.waiting_jobs: - if job.name == "first-modelapply": + if job.name == "model-apply-diploid": apply_job = job break assert apply_job is not None - # Check dependency: apply_job should depend on merge-trim-concat (transfer_vcf_job) + # Check dependency: apply_job should depend on the transfer concat job deps = dag.waiting_jobs[apply_job] dep_names = [j.name for j in deps] - assert "merge-trim-concat" in dep_names + assert "merge-trim-diploid-concat" in dep_names diff --git a/tests/unit/test_estimate_ploidy.py b/tests/unit/test_estimate_ploidy.py new file mode 100644 index 0000000..07116d4 --- /dev/null +++ b/tests/unit/test_estimate_ploidy.py @@ -0,0 +1,153 @@ +""" +Unit tests for the `estimate_ploidy.py` helper script. + +The script is not importable as part of the package (the `scripts` +directory has no `__init__.py`), so it is loaded from its file path. +""" + +import importlib.util +import io +import json +import pathlib +import sys +from types import SimpleNamespace + +import pytest + +SCRIPT = ( + pathlib.Path(__file__).resolve().parents[2] + / "sentieon_cli" + / "scripts" + / "estimate_ploidy.py" +) + + +def _load_script(): + spec = importlib.util.spec_from_file_location("estimate_ploidy", SCRIPT) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +estimate_ploidy = _load_script() + + +def _idxstats(contigs): + """Format idxstats output for the supplied {contig: (length, reads)}""" + lines = [ + f"{ctg}\t{length}\t{reads}\t0" for ctg, (length, reads) in contigs + ] + lines.append("*\t0\t0\t0") + return "\n".join(lines) + "\n" + + +def _chr_contigs(x_reads, y_reads, prefix="chr"): + """Autosomes at ploidy two plus the requested sex chromosome reads""" + contigs = [(f"{prefix}{i}", (1000, 100)) for i in range(1, 23)] + contigs.append((f"{prefix}X", (1000, x_reads))) + if y_reads is not None: + contigs.append((f"{prefix}Y", (1000, y_reads))) + return contigs + + +def _run(monkeypatch, idxstats_contigs, **kwargs): + """Run the script's `main` with a mocked `samtools idxstats`""" + monkeypatch.setattr( + estimate_ploidy.subprocess, + "run", + lambda *a, **kw: SimpleNamespace(stdout=_idxstats(idxstats_contigs)), + ) + outfile = io.StringIO() + args = SimpleNamespace( + input_bam=["sample.bam"], + contigs=kwargs.get("contigs", estimate_ploidy.DEFAULT_CONTIGS), + autosomes=kwargs.get("autosomes", estimate_ploidy.DEFAULT_AUTOSOMES), + x_contig=kwargs.get("x_contig", estimate_ploidy.DEFAULT_X_CONTIG), + y_contig=kwargs.get("y_contig", estimate_ploidy.DEFAULT_Y_CONTIG), + outfile=outfile, + ) + with pytest.raises(SystemExit) as excinfo: + estimate_ploidy.main(args) + assert excinfo.value.code == 0 + return json.loads(outfile.getvalue()) + + +def test_female_sample_with_chr_prefixed_contigs(monkeypatch): + results = _run(monkeypatch, _chr_contigs(x_reads=100, y_reads=0)) + + assert results["sex"] == "female" + assert results["contigs"]["chrX"]["ploidy"] == "2" + assert results["contigs"]["chrY"]["ploidy"] == "0" + + +def test_male_sample_with_chr_prefixed_contigs(monkeypatch): + results = _run(monkeypatch, _chr_contigs(x_reads=50, y_reads=50)) + + assert results["sex"] == "male" + assert results["contigs"]["chrX"]["ploidy"] == "1" + assert results["contigs"]["chrY"]["ploidy"] == "1" + + +def test_male_sample_with_bare_contig_names(monkeypatch): + # A b37-style reference, using the contig name arguments + results = _run( + monkeypatch, + _chr_contigs(x_reads=50, y_reads=50, prefix=""), + contigs=[str(i) for i in range(1, 23)] + ["X", "Y"], + autosomes=[str(i) for i in range(1, 23)], + x_contig="X", + y_contig="Y", + ) + + assert results["sex"] == "male" + assert results["contigs"]["X"]["ploidy"] == "1" + + +def test_a_missing_y_contig_reports_an_unknown_sex(monkeypatch): + # An absent sex chromosome used to raise an uncaught KeyError + results = _run(monkeypatch, _chr_contigs(x_reads=100, y_reads=None)) + + assert results["sex"] == "Unknown" + assert "chrY" not in results["contigs"] + + +def test_no_matching_contigs_reports_an_unknown_sex(monkeypatch): + # A reference whose contig names do not match the supplied ones + results = _run(monkeypatch, [("scaffold1", (1000, 100))]) + + assert results == {"contigs": {}, "sex": "Unknown"} + + +def test_the_contig_arguments_are_parsed(monkeypatch): + argv = [ + "estimate_ploidy.py", + "-i", + "sample.bam", + "--contigs", + "1", + "2", + "--autosomes", + "1", + "2", + "--x_contig", + "X", + "--y_contig", + "Y", + ] + monkeypatch.setattr(sys, "argv", argv) + args = estimate_ploidy.process_args() + + assert args.contigs == ["1", "2"] + assert args.autosomes == ["1", "2"] + assert args.x_contig == "X" + assert args.y_contig == "Y" + + +def test_the_contig_arguments_default_to_chr_prefixed_names(monkeypatch): + monkeypatch.setattr(sys, "argv", ["estimate_ploidy.py", "-i", "s.bam"]) + args = estimate_ploidy.process_args() + + assert args.contigs == estimate_ploidy.DEFAULT_CONTIGS + assert args.x_contig == "chrX" + assert args.y_contig == "chrY" diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py new file mode 100644 index 0000000..fbda6eb --- /dev/null +++ b/tests/unit/test_hybrid_pangenome.py @@ -0,0 +1,1848 @@ +""" +Unit tests for the HybridPangenome pipeline logic +""" + +from importlib.resources import files +import logging +import os +import pathlib +import sys +import tempfile +from typing import List +from unittest.mock import MagicMock + +import packaging.version +import pytest + +# Add the parent directory to the path to import sentieon_cli +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) + +from sentieon_cli import command_strings as cmds +from sentieon_cli import hybrid_pangenome +from sentieon_cli.hybrid_pangenome import HybridPangenome +from sentieon_cli.command_strings import LONGREAD_SV_BED_AWK +from sentieon_cli.dag import DAG +from sentieon_cli.util import SampleSex + +PACKAGE_LOGGER = "sentieon_cli" + + +class _RecordingHandler(logging.Handler): + """Collects formatted messages from the package logger.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.messages: List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def messages(): + """`caplog` does not see the package logger, which does not propagate""" + logger = logging.getLogger(PACKAGE_LOGGER) + handler = _RecordingHandler() + logger.addHandler(handler) + try: + yield handler.messages + finally: + logger.removeHandler(handler) + + +class TestHybridPangenome: + """Test the hybrid-pangenome pipeline logic""" + + def setup_method(self): + """Setup test fixtures""" + self.temp_dir = tempfile.mkdtemp() + self.mock_dir = pathlib.Path(self.temp_dir) + + self.mock_vcf = self.mock_dir / "output.vcf.gz" + self.mock_ref = self.mock_dir / "reference.fa" + self.mock_lr_bam = self.mock_dir / "longreads.bam" + self.mock_bundle = self.mock_dir / "model.bundle" + self.mock_gbz = self.mock_dir / "pangenome.grch38.gbz" + self.mock_hapl = self.mock_dir / "pangenome.grch38.hapl" + self.mock_pop_vcf = self.mock_dir / "population.vcf.gz" + self.mock_dbsnp = self.mock_dir / "dbsnp.vcf.gz" + self.mock_bed = self.mock_dir / "autosomes.bed" + self.mock_r1 = self.mock_dir / "sample_R1.fastq.gz" + self.mock_r2 = self.mock_dir / "sample_R2.fastq.gz" + self.mock_sr_bam = self.mock_dir / "shortreads.bam" + self.mock_lr_ref = self.mock_dir / "lr_reference.fa" + + for file_path in [ + self.mock_ref, + self.mock_lr_bam, + self.mock_bundle, + self.mock_gbz, + self.mock_hapl, + self.mock_pop_vcf, + self.mock_dbsnp, + self.mock_bed, + self.mock_r1, + self.mock_r2, + self.mock_sr_bam, + self.mock_lr_ref, + ]: + file_path.touch() + + with open(str(self.mock_ref) + ".fai", "w") as f: + f.write("chr1\t1000\t0\t80\t81\n") + + def create_pipeline(self): + """Create a HybridPangenome pipeline for testing""" + pipeline = HybridPangenome() + + pipeline.logger = MagicMock() + + # Configure arguments + pipeline.output_vcf = self.mock_vcf + pipeline.reference = self.mock_ref + pipeline.model_bundle = self.mock_bundle + pipeline.r1_fastq = [self.mock_r1] + pipeline.r2_fastq = [self.mock_r2] + pipeline.lr_aln = [self.mock_lr_bam] + pipeline.gbz = self.mock_gbz + pipeline.hapl = self.mock_hapl + pipeline.pop_vcf = self.mock_pop_vcf + pipeline.dbsnp = self.mock_dbsnp + pipeline.bed = self.mock_bed + pipeline.cores = 2 + pipeline.dry_run = True + pipeline.skip_version_check = True + pipeline.skip_multiqc = True + pipeline.tmp_dir = self.mock_dir + + # State normally set by validate() + pipeline.pangenome_ref_name = "GRCh38" + pipeline.pangenome_contig_prefix = "GRCh38#0#" + pipeline.fai_data = {"chr1": {"length": 1000}} + pipeline.shards = [MagicMock()] + pipeline.shards[0].contig = "chr1" + pipeline.shards[0].start = 1 + pipeline.shards[0].stop = 1000 + pipeline.pop_vcf_contigs = {"chr1": 1000} + pipeline.fastq_readgroup = {"ID": "rg1", "SM": "sample1"} + pipeline.lr_readgroups = [[{"ID": "lr-rg1", "SM": "sample1"}]] + pipeline.sample_sm = "sample1" + + return pipeline + + def create_aligned_pipeline(self): + """Create a pipeline with aligned short-read input""" + pipeline = self.create_pipeline() + pipeline.r1_fastq = [] + pipeline.r2_fastq = [] + pipeline.readgroup = None + pipeline.fastq_readgroup = None + pipeline.sr_aln = [self.mock_sr_bam] + pipeline.sr_readgroups = [[{"ID": "sr-rg1", "SM": "sample1"}]] + return pipeline + + def create_lr_realign_pipeline(self, aligned_sr=False): + """Create a pipeline that realigns the long-read input""" + pipeline = ( + self.create_aligned_pipeline() + if aligned_sr + else self.create_pipeline() + ) + pipeline.lr_align_input = True + pipeline.lr_input_ref = self.mock_lr_ref + return pipeline + + def _get_all_job_names(self, dag): + """Helper to get all job names from a DAG""" + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) + return [job.name for job in all_jobs], all_jobs + + def _get_job(self, all_jobs, name): + return next(j for j in all_jobs if j.name == name) + + def _get_dep_names(self, dag, all_jobs, name): + """Helper to get the dependency names of a job""" + job = self._get_job(all_jobs, name) + return {dep.name for dep in dag.waiting_jobs.get(job, set())} + + def test_dag_jobs_present(self): + """All expected jobs are in the DAG""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + assert isinstance(dag, DAG) + + job_names, _ = self._get_all_job_names(dag) + for name in ( + "kmc", + "bwa-extract", + "vg-haplotypes", + "vg-convert-gfa", + "graph-update-raw", + "longreadsv", + "longread-sv-bed", + "graph-update", + "gfa2fa", + "faidx", + "mm2-lift", + "locuscollector-bwa", + "dedup-bwa", + "locuscollector-lift", + "dedup-lift", + "metrics", + "estimate-ploidy", + "pangenome-sv", + "dnascope", + "model-apply", + ): + assert name in job_names, f"missing job: {name}" + + def test_kmc_command(self): + """KMC counts the short-read fastq and long-read fasta together""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "kmc").shell) + assert "-k29" in cmd_str + assert "-m30" in cmd_str + assert "-okff" in cmd_str + assert "-fa /dev/stdin" in cmd_str + assert "samtools fasta" in cmd_str + assert str(self.mock_lr_bam) in cmd_str + assert str(self.mock_r1) in cmd_str + assert str(self.mock_r2) in cmd_str + + def test_bwa_readgroup_lr0(self): + """The bwa alignment carries the LR:0 readgroup attribute""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "bwa-extract").shell) + # The user's RGID is used unchanged + assert r"ID:rg1\t" in cmd_str + assert "ID:rg1-bwa" not in cmd_str + assert "LR:0" in cmd_str + assert "pgutil extract" in cmd_str + assert "bwa.model" in cmd_str + # The alignment is written to the temporary directory for dedup + bwa_bam = self.mock_dir / "sample-bwa.bam" + assert f"-b {bwa_bam}" in cmd_str + + def test_mm2_lift_readgroup_lr2(self): + """The lifted alignment carries the LR:2 readgroup attribute""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "mm2-lift").shell) + assert "ID:rg1-pg" in cmd_str + assert "LR:2" in cmd_str + assert "--secondary=yes" in cmd_str + assert "pgutil lift" in cmd_str + assert "--prefix" in cmd_str + assert "GRCh38#0#" in cmd_str + assert "minimap2.model" in cmd_str + # `util sort` writes the sorted alignment to the temporary directory + assert "util sort" in cmd_str + lift_bam = self.mock_dir / "sample-lift.bam" + assert f"-o {lift_bam}" in cmd_str + + def test_dedup_commands(self): + """The short-read alignments are deduplicated, with Dedup metrics + from the bwa alignment only""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + bwa_bam = self.mock_dir / "sample-bwa.bam" + lift_bam = self.mock_dir / "sample-lift.bam" + out_bwa = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + out_lift = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.cram") + dedup_metrics = str( + self.mock_dir / "output_metrics" / "output.txt.dedup_metrics.txt" + ) + + lc_cmd = str(self._get_job(all_jobs, "locuscollector-bwa").shell) + assert "--algo LocusCollector" in lc_cmd + assert str(bwa_bam) in lc_cmd + + bwa_cmd = str(self._get_job(all_jobs, "dedup-bwa").shell) + assert "--algo Dedup" in bwa_cmd + assert str(bwa_bam) in bwa_cmd + assert out_bwa in bwa_cmd + assert f"--metrics {dedup_metrics}" in bwa_cmd + assert "IndelLeftAlignReadTransform" not in bwa_cmd + + lift_cmd = str(self._get_job(all_jobs, "dedup-lift").shell) + assert "--algo Dedup" in lift_cmd + assert str(lift_bam) in lift_cmd + assert out_lift in lift_cmd + assert "IndelLeftAlignReadTransform" not in lift_cmd + # Dedup metrics are only collected from the bwa alignment + assert "--metrics" not in lift_cmd + + # The long-read input is assumed to be deduplicated already + for job_name in ("locuscollector-bwa", "dedup-bwa", "dedup-lift"): + assert str(self.mock_lr_bam) not in str( + self._get_job(all_jobs, job_name).shell + ) + + def test_metrics_job(self): + """Metrics are collected from the deduplicated alignments""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "metrics").shell) + assert ( + str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + in cmd_str + ) + assert ( + str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.cram") + in cmd_str + ) + assert "--algo GCBias" in cmd_str + assert "--algo WgsMetricsAlgo" in cmd_str + assert "rehead-metrics" in job_names + + def test_skip_metrics(self): + """skip_metrics removes the metrics jobs and the Dedup metrics""" + pipeline = self.create_pipeline() + pipeline.skip_metrics = True + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + assert "metrics" not in job_names + assert "rehead-metrics" not in job_names + assert "multiqc" not in job_names + # Deduplication still runs + assert "dedup-bwa" in job_names + assert "dedup-lift" in job_names + assert "--metrics" not in str( + self._get_job(all_jobs, "dedup-bwa").shell + ) + + def test_graph_update_commands(self): + """PGHapUpdateAlgo runs without and then with the SV BED""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + raw_cmd = str(self._get_job(all_jobs, "graph-update-raw").shell) + assert "--algo PGHapUpdateAlgo" in raw_cmd + assert "--gfa_file " in raw_cmd + assert "sample-hap.raw.gfa" in raw_cmd + assert "sample-pangenome-raw.gfa" in raw_cmd + assert "--target_bed" not in raw_cmd + assert str(self.mock_lr_bam) in raw_cmd + assert "--prefix" in raw_cmd + assert "GRCh38#0#" in raw_cmd + + update_cmd = str(self._get_job(all_jobs, "graph-update").shell) + assert "--algo PGHapUpdateAlgo" in update_cmd + assert "--gfa_file " in update_cmd + assert "sample-pangenome-raw.gfa" in update_cmd + assert "--target_bed " in update_cmd + assert "sample-sv.bed" in update_cmd + assert "--prefix" in update_cmd + assert "GRCh38#0#" in update_cmd + + def test_longreadsv_and_bed(self): + """LongReadSV runs on the long reads and the awk BED script is + retained verbatim""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + sv_cmd = str(self._get_job(all_jobs, "longreadsv").shell) + assert "--algo LongReadSV" in sv_cmd + assert "longreadsv.model" in sv_cmd + assert "--min_sv_size 20" in sv_cmd + assert str(self.mock_lr_bam) in sv_cmd + + bed_cmd = str(self._get_job(all_jobs, "longread-sv-bed").shell) + assert 'n=split($10,a,":")' in bed_cmd + # Sort in reference contig order for bedtools merge + assert f"bedtools sort -faidx {self.mock_ref}.fai -i -" in bed_cmd + assert "bedtools merge" in bed_cmd + assert "sample-sv.bed" in bed_cmd + # The awk script matches the validated implementation + assert "chr[^_]+_[0-9]+_[0-9]+" in LONGREAD_SV_BED_AWK + + def test_gfa2fa_and_faidx(self): + """The updated graph is converted to an indexed FASTA""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + gfa2fa_cmd = str(self._get_job(all_jobs, "gfa2fa").shell) + assert "pgutil gfa2fa" in gfa2fa_cmd + assert str(self.mock_ref) + ".fai" in gfa2fa_cmd + assert "sample-pangenome.gfa" in gfa2fa_cmd + assert "sample-pangenome.fa" in gfa2fa_cmd + + faidx_cmd = str(self._get_job(all_jobs, "faidx").shell) + assert "samtools faidx" in faidx_cmd + assert "sample-pangenome.fa" in faidx_cmd + + def test_gfa2fa_with_vg_paths(self): + """`vg paths` replaces `pgutil gfa2fa` when selected""" + pipeline = self.create_pipeline() + pipeline.gfa2fa_with_vg = True + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + gfa = self.mock_dir / "sample-pangenome.gfa" + fasta = self.mock_dir / "sample-pangenome.fa" + gfa2fa_cmd = str(self._get_job(all_jobs, "gfa2fa").shell) + assert gfa2fa_cmd == (f"vg paths -x {gfa} -F >'{fasta}'") + assert "pgutil gfa2fa" not in gfa2fa_cmd + assert "-Q" not in gfa2fa_cmd + + # The DAG shape is unchanged + assert self._get_dep_names(dag, all_jobs, "faidx") == {"gfa2fa"} + assert "faidx" in self._get_dep_names(dag, all_jobs, "mm2-lift") + + def test_resolve_gfa2fa_tool_hg38(self, monkeypatch): + """A GRCh38 reference keeps `pgutil gfa2fa`, unprobed""" + pipeline = self.create_pipeline() + pipeline.skip_version_check = False + pipeline.fai_data = { + "chr1": {"length": 248956422}, + "chr2": {"length": 242193529}, + "chrX": {"length": 156040895}, + } + probe = MagicMock() + monkeypatch.setattr(hybrid_pangenome, "executable_version", probe) + + pipeline.resolve_gfa2fa_tool() + assert pipeline.reference_build == "hg38" + assert pipeline.gfa2fa_with_vg is False + probe.assert_not_called() + + def test_resolve_gfa2fa_tool_old_driver(self, monkeypatch): + """A non-GRCh38 reference on the GRCh38-only driver uses vg""" + pipeline = self.create_pipeline() + pipeline.skip_version_check = False + monkeypatch.setattr( + hybrid_pangenome, + "executable_version", + lambda cmd: packaging.version.Version("202503.04"), + ) + + pipeline.resolve_gfa2fa_tool() + assert pipeline.reference_build is None + assert pipeline.gfa2fa_with_vg is True + + def test_resolve_gfa2fa_tool_new_driver(self, monkeypatch): + """A newer driver keeps `pgutil gfa2fa` for any reference""" + pipeline = self.create_pipeline() + pipeline.skip_version_check = False + monkeypatch.setattr( + hybrid_pangenome, + "executable_version", + lambda cmd: packaging.version.Version("202503.05"), + ) + + pipeline.resolve_gfa2fa_tool() + assert pipeline.gfa2fa_with_vg is False + + def test_resolve_gfa2fa_tool_skip_version_check(self, monkeypatch): + """`--skip_version_check` does not probe and prefers vg""" + pipeline = self.create_pipeline() + pipeline.skip_version_check = True + probe = MagicMock() + monkeypatch.setattr(hybrid_pangenome, "executable_version", probe) + + pipeline.resolve_gfa2fa_tool() + assert pipeline.gfa2fa_with_vg is True + probe.assert_not_called() + + def test_resolve_gfa2fa_tool_probe_failure(self, monkeypatch): + """An unknown driver version prefers vg""" + pipeline = self.create_pipeline() + pipeline.skip_version_check = False + monkeypatch.setattr( + hybrid_pangenome, + "executable_version", + lambda cmd: None, + ) + + pipeline.resolve_gfa2fa_tool() + assert pipeline.gfa2fa_with_vg is True + + def test_calling_inputs_and_replace_rg(self): + """PangenomeSV and DNAscope use the bwa, lifted, and long reads, + rewriting only the long-read readgroups with LR:1""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + bwa_aln = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + lift_aln = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.cram") + replace_arg = r"lr-rg1=ID:lr-rg1\tSM:sample1\tLR:1" + + for job_name in ("pangenome-sv", "dnascope"): + cmd_str = str(self._get_job(all_jobs, job_name).shell) + assert bwa_aln in cmd_str + assert lift_aln in cmd_str + assert str(self.mock_lr_bam) in cmd_str + assert replace_arg in cmd_str + # The long-read input is preceded by its --replace_rg argument + assert cmd_str.index(replace_arg) < cmd_str.index( + str(self.mock_lr_bam) + ) + # The pipeline-generated alignments are not rewritten + assert cmd_str.count("--replace_rg") == 1 + + def test_pangenomesv_command(self): + """PangenomeSV uses the updated graph and min_af""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "pangenome-sv").shell) + assert "--algo PangenomeSV" in cmd_str + assert "--gfa_file" in cmd_str + assert "sample-pangenome.gfa" in cmd_str + assert "--min_af 0.1" in cmd_str + assert "--prefix" in cmd_str + assert "GRCh38#0#" in cmd_str + sv_vcf = str(self.mock_vcf).replace(".vcf.gz", "_sv.vcf.gz") + assert sv_vcf in cmd_str + # SV calling is not restricted to the small-variant BED + assert f"--interval {self.mock_bed}" not in cmd_str + + def test_pangenome_check_jobs(self): + """The graph checks sit between vg and everything downstream. + + A wrong `--pangenome_ref_name` is silent in `vg`: the sampled + graph loses its reference paths and the GFA loses its rGFA tags, + and `pgutil lift` then leaves every read unmapped. + """ + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + assert "check-sample-gbz" in job_names + assert "check-sample-gfa" in job_names + + assert self._get_dep_names(dag, all_jobs, "check-sample-gbz") == { + "vg-haplotypes" + } + assert self._get_dep_names(dag, all_jobs, "vg-convert-gfa") == { + "check-sample-gbz" + } + assert self._get_dep_names(dag, all_jobs, "check-sample-gfa") == { + "vg-convert-gfa" + } + # The raw graph update is the only consumer of the converted GFA + assert self._get_dep_names(dag, all_jobs, "graph-update-raw") == { + "check-sample-gfa" + } + + # The checks share the log group of the jobs they guard + haplotypes = self._get_job(all_jobs, "vg-haplotypes") + check_gbz = self._get_job(all_jobs, "check-sample-gbz") + check_gfa = self._get_job(all_jobs, "check-sample-gfa") + assert check_gbz.task_name == haplotypes.task_name + assert check_gfa.task_name == haplotypes.task_name + + def test_pangenome_check_commands(self): + """The checks run the resolved name and prefix""" + pipeline = self.create_pipeline() + pipeline.pangenome_ref_name = "CHM13" + pipeline.pangenome_contig_prefix = "CHM13#0#" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + gbz_cmd = str(self._get_job(all_jobs, "check-sample-gbz").shell) + assert "sentieon_cli.pangenome_meta check-gbz" in gbz_cmd + assert f"--gbz {self.mock_dir}/sample_pangenome.gbz" in gbz_cmd + assert "--reference_name CHM13" in gbz_cmd + # The `#` characters are shell-quoted in the rendered command + assert "--contig_prefix 'CHM13#0#'" in gbz_cmd + + gfa_cmd = str(self._get_job(all_jobs, "check-sample-gfa").shell) + assert "sentieon_cli.pangenome_meta check-gfa" in gfa_cmd + assert f"--gfa {self.mock_dir}/sample-hap.raw.gfa" in gfa_cmd + assert "--reference_name CHM13" in gfa_cmd + assert "--contig_prefix 'CHM13#0#'" in gfa_cmd + + def test_the_resolved_reference_name_reaches_vg(self): + """`vg haplotypes` and `vg convert` use the resolved name""" + pipeline = self.create_pipeline() + pipeline.pangenome_ref_name = "CHM13" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + haplotypes = str(self._get_job(all_jobs, "vg-haplotypes").shell) + assert "--set-reference CHM13" in haplotypes + convert = str(self._get_job(all_jobs, "vg-convert-gfa").shell) + assert "-Q CHM13" in convert + + def test_pangenome_contig_prefix(self): + """`--pangenome_contig_prefix` reaches every graph consumer""" + pipeline = self.create_pipeline() + pipeline.pangenome_contig_prefix = "CHM13#0#" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + for job_name in ( + "mm2-lift", + "graph-update-raw", + "graph-update", + "pangenome-sv", + ): + cmd_str = str(self._get_job(all_jobs, job_name).shell) + assert "--prefix" in cmd_str, job_name + assert "CHM13#0#" in cmd_str, job_name + assert "GRCh38#0#" not in cmd_str, job_name + + def test_dnascope_command(self): + """DNAscope runs with the model, interval, and pcr_indel_model""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert "--algo DNAscope" in cmd_str + assert "dnascope.model" in cmd_str + assert "--pcr_indel_model CONSERVATIVE" in cmd_str + assert f"--interval {self.mock_bed}" in cmd_str + assert f"--dbsnp {self.mock_dbsnp}" in cmd_str + + def test_pcr_free(self): + """--pcr_free calls DNAscope with `--pcr_indel_model NONE`""" + pipeline = self.create_pipeline() + pipeline.pcr_free = True + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert "--pcr_indel_model NONE" in cmd_str + + def test_bam_format(self): + """--bam_format switches the deduplicated outputs to BAM""" + pipeline = self.create_pipeline() + pipeline.bam_format = True + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + bwa_bam = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.bam") + lift_bam = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.bam") + + assert bwa_bam in str(self._get_job(all_jobs, "dedup-bwa").shell) + assert lift_bam in str(self._get_job(all_jobs, "dedup-lift").shell) + + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert bwa_bam in cmd_str + assert lift_bam in cmd_str + + def test_model_apply_writes_output_vcf(self): + """DNAModelApply produces the final output VCF""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "model-apply").shell) + assert "--algo DNAModelApply" in cmd_str + assert str(self.mock_vcf) in cmd_str + + def test_skip_model_apply(self): + """With skip_model_apply the transfer writes the final VCF""" + pipeline = self.create_pipeline() + pipeline.skip_model_apply = True + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + assert "model-apply" not in job_names + concat_job = self._get_job(all_jobs, "merge-trim-concat") + assert str(pipeline.output_vcf) in str(concat_job.shell) + + def test_skip_svs(self): + """skip_svs removes PangenomeSV but not the graph jobs""" + pipeline = self.create_pipeline() + pipeline.skip_svs = True + dag = pipeline.build_dag() + job_names, _ = self._get_all_job_names(dag) + + assert "pangenome-sv" not in job_names + assert "dnascope" in job_names + assert "graph-update" in job_names + + def test_skip_small_variants(self): + """skip_small_variants removes DNAscope, transfer, and model-apply""" + pipeline = self.create_pipeline() + pipeline.skip_small_variants = True + dag = pipeline.build_dag() + job_names, _ = self._get_all_job_names(dag) + + assert "dnascope" not in job_names + assert "model-apply" not in job_names + assert "merge-trim-concat" not in job_names + assert "pangenome-sv" in job_names + + def test_multiple_lr_inputs(self): + """Each long-read input gets its own --replace_rg arguments""" + lr_bam2 = self.mock_dir / "longreads2.bam" + lr_bam2.touch() + + pipeline = self.create_pipeline() + pipeline.lr_aln = [self.mock_lr_bam, lr_bam2] + pipeline.lr_readgroups = [ + [{"ID": "lr-rg1", "SM": "sample1"}], + [{"ID": "lr-rg2", "SM": "sample1"}], + ] + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert r"lr-rg1=ID:lr-rg1\tSM:sample1\tLR:1" in cmd_str + assert r"lr-rg2=ID:lr-rg2\tSM:sample1\tLR:1" in cmd_str + assert str(lr_bam2) in cmd_str + + def test_rgsm_overrides_sm(self): + """--rgsm overrides the SM tag in the rewritten readgroups""" + pipeline = self.create_pipeline() + pipeline.rgsm = "override_sm" + pipeline.sample_sm = "override_sm" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert r"lr-rg1=ID:lr-rg1\tSM:override_sm\tLR:1" in cmd_str + + bwa_cmd = str(self._get_job(all_jobs, "bwa-extract").shell) + assert "SM:override_sm" in bwa_cmd + + # Aligned short-read input + + def test_aligned_dag_jobs(self): + """Aligned short-read input skips alignment, dedup, and metrics""" + pipeline = self.create_aligned_pipeline() + assert not pipeline.skip_metrics + dag = pipeline.build_dag() + job_names, _ = self._get_all_job_names(dag) + + for name in ( + "extract-kmc-symlink", + "extract-kmc", + "vg-haplotypes", + "graph-update", + "mm2-lift", + "estimate-ploidy", + "pangenome-sv", + "dnascope", + "model-apply", + ): + assert name in job_names, f"missing job: {name}" + + for name in ( + "kmc", + "bwa-extract", + "locuscollector-bwa", + "dedup-bwa", + "locuscollector-lift", + "dedup-lift", + "metrics", + "rehead-metrics", + "multiqc", + ): + assert name not in job_names, f"unexpected job: {name}" + + def test_aligned_extract_kmc_command(self): + """Read extraction and k-mer counting run in a single pass""" + pipeline = self.create_aligned_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + symlink_cmd = str(self._get_job(all_jobs, "extract-kmc-symlink").shell) + rw_bam = self.mock_dir / "extract-kmc-rw.bam" + assert f"ln -sf /dev/stdout {rw_bam}" in symlink_cmd + + cmd_str = str(self._get_job(all_jobs, "extract-kmc").shell) + # One pass over the aligned short reads, without duplicate, + # secondary, or supplementary reads + assert "--algo ReadWriter" in cmd_str + assert "--output_flag_filter 0xf00:0" in cmd_str + assert str(self.mock_sr_bam) in cmd_str + assert str(rw_bam) in cmd_str + # writing the extracted reads to the fastq and the reads to kmc + assert "pgutil extract" in cmd_str + ext_fastq = self.mock_dir / "sample-extract.fq.gz" + assert f"-o {ext_fastq}" in cmd_str + assert "-a -" in cmd_str + # concatenated with the long reads + assert "cat " in cmd_str + assert "samtools fasta" in cmd_str + assert str(self.mock_lr_bam) in cmd_str + assert "-fa /dev/stdin" in cmd_str + assert "-k29" in cmd_str + assert "-m30" in cmd_str + + def test_aligned_extract_kmc_memory(self): + """`--kmer_memory` is passed to the single-pass KMC""" + pipeline = self.create_aligned_pipeline() + pipeline.kmer_memory = 64 + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "extract-kmc").shell) + assert "-m64" in cmd_str + + def test_aligned_dag_dependencies(self): + """The pangenome is built from the extracted k-mers""" + pipeline = self.create_aligned_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + assert self._get_dep_names(dag, all_jobs, "extract-kmc") == { + "extract-kmc-symlink" + } + assert self._get_dep_names(dag, all_jobs, "vg-haplotypes") == { + "extract-kmc" + } + assert "extract-kmc" in self._get_dep_names(dag, all_jobs, "mm2-lift") + assert self._get_dep_names(dag, all_jobs, "dnascope") == {"mm2-lift"} + + def test_aligned_lift_output(self): + """The lifted alignment is the final short-read output""" + pipeline = self.create_aligned_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + lift_cram = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.cram") + cmd_str = str(self._get_job(all_jobs, "mm2-lift").shell) + assert f"-o {lift_cram}" in cmd_str + # The readgroup is seeded from the first input readgroup + assert "ID:sr-rg1-pg" in cmd_str + assert "LR:2" in cmd_str + + def test_aligned_replace_rg(self): + """The aligned short reads are rewritten with LR:0""" + pipeline = self.create_aligned_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + lift_cram = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.cram") + sr_arg = r"sr-rg1=ID:sr-rg1\tSM:sample1\tLR:0" + lr_arg = r"lr-rg1=ID:lr-rg1\tSM:sample1\tLR:1" + + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert cmd_str.count("--replace_rg") == 2 + # Each row precedes the input file it applies to + assert cmd_str.index(sr_arg) < cmd_str.index(str(self.mock_sr_bam)) + assert cmd_str.index(lr_arg) < cmd_str.index(str(self.mock_lr_bam)) + # The calling inputs are ordered short reads, lifted, long reads + assert ( + cmd_str.index(str(self.mock_sr_bam)) + < cmd_str.index(lift_cram) + < cmd_str.index(str(self.mock_lr_bam)) + ) + # The lifted alignment carries its LR tag already + assert "LR:2" not in cmd_str + + def test_aligned_bam_format(self): + """--bam_format switches the lifted output to BAM""" + pipeline = self.create_aligned_pipeline() + pipeline.bam_format = True + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + lift_bam = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.bam") + assert lift_bam in str(self._get_job(all_jobs, "mm2-lift").shell) + assert lift_bam in str(self._get_job(all_jobs, "dnascope").shell) + + # Short-read input validation + + def test_validate_sr_inputs_fastq(self): + """fastq input with a readgroup is accepted""" + pipeline = self.create_pipeline() + pipeline.readgroup = r"@RG\tID:rg1\tSM:sample1" + pipeline.validate_sr_inputs() + + def test_validate_sr_inputs_aligned(self): + """Aligned input without a readgroup is accepted""" + pipeline = self.create_aligned_pipeline() + pipeline.validate_sr_inputs() + + def test_validate_sr_inputs_both(self): + """fastq and aligned short reads cannot be combined""" + pipeline = self.create_pipeline() + pipeline.readgroup = r"@RG\tID:rg1\tSM:sample1" + pipeline.sr_aln = [self.mock_sr_bam] + with pytest.raises(SystemExit): + pipeline.validate_sr_inputs() + + def test_validate_sr_inputs_neither(self): + """Short reads are required""" + pipeline = self.create_pipeline() + pipeline.r1_fastq = [] + pipeline.r2_fastq = [] + with pytest.raises(SystemExit): + pipeline.validate_sr_inputs() + + def test_validate_sr_inputs_fastq_without_readgroup(self): + """fastq input requires a readgroup""" + pipeline = self.create_pipeline() + pipeline.readgroup = None + with pytest.raises(SystemExit): + pipeline.validate_sr_inputs() + + def test_validate_sr_inputs_fastq_length_mismatch(self): + """The r1 and r2 fastq lists must have the same length""" + pipeline = self.create_pipeline() + pipeline.readgroup = r"@RG\tID:rg1\tSM:sample1" + pipeline.r2_fastq = [] + with pytest.raises(SystemExit): + pipeline.validate_sr_inputs() + + def test_validate_sr_inputs_aligned_with_readgroup(self): + """`--readgroup` cannot be used with aligned input""" + pipeline = self.create_aligned_pipeline() + pipeline.readgroup = r"@RG\tID:rg1\tSM:sample1" + with pytest.raises(SystemExit): + pipeline.validate_sr_inputs() + + # Unaligned (uBAM/uCRAM) long-read input + + def test_lr_realign_job(self): + """The long-read input is realigned with minimap2""" + pipeline = self.create_lr_realign_pipeline() + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + assert "bam-realign-0" in job_names + cmd_str = str(self._get_job(all_jobs, "bam-realign-0").shell) + # The input reference decodes the input file + assert f"samtools fastq --reference {self.mock_lr_ref}" in cmd_str + # The long-read minimap2 model aligns to the linear reference + assert "minimap2_lr.model" in cmd_str + assert str(self.mock_ref) in cmd_str + realigned = str(self.mock_vcf).replace(".vcf.gz", "_mm2_sorted_0.cram") + assert f"-o {realigned}" in cmd_str + assert self._get_dep_names(dag, all_jobs, "bam-realign-0") == set() + + def test_lr_realign_downstream(self): + """Downstream jobs consume the realigned long reads""" + pipeline = self.create_lr_realign_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + realigned = str(self.mock_vcf).replace(".vcf.gz", "_mm2_sorted_0.cram") + for job_name in ( + "longreadsv", + "graph-update-raw", + "graph-update", + "pangenome-sv", + "dnascope", + ): + cmd_str = str(self._get_job(all_jobs, job_name).shell) + assert realigned in cmd_str, job_name + assert str(self.mock_lr_bam) not in cmd_str, job_name + + for job_name in ("longreadsv", "graph-update-raw", "dnascope"): + assert "bam-realign-0" in self._get_dep_names( + dag, all_jobs, job_name + ), job_name + + # The readgroups of the realigned input are unchanged + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert r"lr-rg1=ID:lr-rg1\tSM:sample1\tLR:1" in cmd_str + + def test_lr_realign_kmc_reads_original_input(self): + """K-mer counting reads the original long-read input""" + pipeline = self.create_lr_realign_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "kmc").shell) + assert f"samtools fasta --reference {self.mock_lr_ref}" in cmd_str + assert str(self.mock_lr_bam) in cmd_str + assert "_mm2_sorted_0" not in cmd_str + assert self._get_dep_names(dag, all_jobs, "kmc") == set() + + def test_lr_realign_without_input_ref(self): + """The target reference decodes the input without `lr_input_ref`""" + pipeline = self.create_lr_realign_pipeline() + pipeline.lr_input_ref = None + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + cmd_str = str(self._get_job(all_jobs, "kmc").shell) + assert f"samtools fasta --reference {self.mock_ref}" in cmd_str + + def test_aligned_sr_with_lr_realign(self): + """Aligned short reads combine with realigned long reads""" + pipeline = self.create_lr_realign_pipeline(aligned_sr=True) + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + assert "bam-realign-0" in job_names + assert "extract-kmc" in job_names + assert "kmc" not in job_names + assert "dedup-bwa" not in job_names + + realigned = str(self.mock_vcf).replace(".vcf.gz", "_mm2_sorted_0.cram") + lift_cram = str(self.mock_vcf).replace(".vcf.gz", "_lift_deduped.cram") + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) + assert ( + cmd_str.index(str(self.mock_sr_bam)) + < cmd_str.index(lift_cram) + < cmd_str.index(realigned) + ) + assert self._get_dep_names(dag, all_jobs, "dnascope") == { + "mm2-lift", + "bam-realign-0", + } + + # k-mer counting still reads the original long-read input + extract_cmd = str(self._get_job(all_jobs, "extract-kmc").shell) + assert str(self.mock_lr_bam) in extract_cmd + assert f"--reference {self.mock_lr_ref}" in extract_cmd + + # Readgroup validation + # + # The readgroups are read from real input headers, so these tests set + # the parsed readgroups directly (or patch `get_rg_lines`) rather than + # relying on the synthetic readgroup of a dry run. + + def create_rg_pipeline(self): + """A pipeline with aligned inputs and readgroup checks enabled""" + pipeline = self.create_aligned_pipeline() + pipeline.dry_run = False + pipeline.sr_readgroups = [[{"ID": "sr-rg1", "SM": "sample1"}]] + pipeline.lr_readgroups = [[{"ID": "lr-rg1", "SM": "sample1"}]] + return pipeline + + def test_validate_readgroups_ok(self): + """Unique readgroups with a shared SM tag are accepted""" + pipeline = self.create_rg_pipeline() + pipeline.validate_readgroups() + assert pipeline.sample_sm == "sample1" + + def test_validate_readgroups_sr_without_rg_lines(self): + """An aligned short-read input without @RG lines is rejected""" + pipeline = self.create_rg_pipeline() + pipeline.sr_readgroups = [[]] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_lr_without_rg_lines(self): + """A long-read input without @RG lines is rejected""" + pipeline = self.create_rg_pipeline() + pipeline.lr_readgroups = [[]] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_no_rg_lines_with_rgsm(self): + """`--rgsm` does not bypass the missing readgroup check""" + pipeline = self.create_rg_pipeline() + pipeline.rgsm = "override_sm" + pipeline.sr_readgroups = [[]] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_duplicate_ids_across_inputs(self): + """A readgroup ID shared by the short and long reads is rejected""" + pipeline = self.create_rg_pipeline() + pipeline.sr_readgroups = [[{"ID": "rg1", "SM": "sample1"}]] + pipeline.lr_readgroups = [[{"ID": "rg1", "SM": "sample1"}]] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_duplicate_ids_within_input(self): + """A readgroup ID repeated inside one input is rejected""" + pipeline = self.create_rg_pipeline() + pipeline.sr_readgroups = [ + [{"ID": "rg1", "SM": "sample1"}, {"ID": "rg1", "SM": "sample1"}] + ] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_duplicate_with_fastq_readgroup(self): + """A long-read ID matching the `--readgroup` ID is rejected""" + pipeline = self.create_pipeline() + pipeline.dry_run = False + pipeline.fastq_readgroup = {"ID": "rg1", "SM": "sample1"} + pipeline.lr_readgroups = [[{"ID": "rg1", "SM": "sample1"}]] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_lift_id_collision(self): + """An input ID reserved for the lifted alignment is rejected""" + pipeline = self.create_pipeline() + pipeline.dry_run = False + pipeline.fastq_readgroup = {"ID": "rg1", "SM": "sample1"} + pipeline.lr_readgroups = [[{"ID": "rg1-pg", "SM": "sample1"}]] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_validate_readgroups_without_sample_name(self): + """A sample name is required for the output readgroups""" + pipeline = self.create_rg_pipeline() + pipeline.fastq_readgroup = None + pipeline.sr_aln = [] + pipeline.lr_aln = [] + pipeline.sr_readgroups = [] + pipeline.lr_readgroups = [] + with pytest.raises(SystemExit): + pipeline.validate_readgroups() + + def test_collect_readgroups_malformed_rg_line(self, monkeypatch): + """A malformed @RG line in an input header is rejected""" + pipeline = self.create_aligned_pipeline() + pipeline.dry_run = False + monkeypatch.setattr( + cmds, + "get_rg_lines", + lambda aln, dry_run: ["@RG\tID:sr-rg1\t"], + ) + with pytest.raises(SystemExit): + pipeline.collect_readgroups() + + def test_collect_readgroups_parses_input_headers(self, monkeypatch): + """The @RG lines of every input are parsed""" + pipeline = self.create_aligned_pipeline() + pipeline.dry_run = False + headers = { + str(self.mock_sr_bam): ["@RG\tID:sr-rg1\tSM:sample1"], + str(self.mock_lr_bam): ["@RG\tID:lr-rg1\tSM:sample1"], + } + monkeypatch.setattr( + cmds, + "get_rg_lines", + lambda aln, dry_run: headers[str(aln)], + ) + pipeline.collect_readgroups() + assert pipeline.sr_readgroups == [[{"ID": "sr-rg1", "SM": "sample1"}]] + assert pipeline.lr_readgroups == [[{"ID": "lr-rg1", "SM": "sample1"}]] + pipeline.validate_readgroups() + assert pipeline.sample_sm == "sample1" + + # Ploidy estimation + + def test_estimate_ploidy_job_with_fastq_input(self): + """Ploidy is estimated from the deduplicated bwa alignment""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + bwa_aln = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + job = self._get_job(all_jobs, "estimate-ploidy") + assert "estimate_ploidy.py" in str(job.shell) + assert f"-i {bwa_aln}" in str(job.shell) + assert job.task_name == "ploidy" + assert job.threads == 0 + assert self._get_dep_names(dag, all_jobs, "estimate-ploidy") == { + "dedup-bwa" + } + + def test_estimate_ploidy_job_with_aligned_input(self): + """Aligned short reads are used as-is, without dependencies""" + pipeline = self.create_aligned_pipeline() + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + job = self._get_job(all_jobs, "estimate-ploidy") + assert f"-i {self.mock_sr_bam}" in str(job.shell) + assert self._get_dep_names(dag, all_jobs, "estimate-ploidy") == set() + + def test_ploidy_json_stashed(self): + """The ploidy JSON path is stashed for the second DAG""" + pipeline = self.create_pipeline() + pipeline.build_dag() + + assert pipeline.ploidy_json == pathlib.Path( + str(self.mock_vcf).replace(".vcf.gz", "_ploidy.json") + ) + + def test_estimate_ploidy_with_skip_small_variants(self): + """Ploidy estimation runs before the small-variant early return""" + pipeline = self.create_pipeline() + pipeline.skip_small_variants = True + job_names, _ = self._get_all_job_names(pipeline.build_dag()) + + assert "estimate-ploidy" in job_names + + # T1K HLA/KIR genotyping + + def enable_t1k(self, pipeline): + """Supply the T1K reference files""" + for attr, name in ( + ("t1k_hla_seq", "hla_seq.fa"), + ("t1k_hla_coord", "hla_coord.fa"), + ("t1k_kir_seq", "kir_seq.fa"), + ("t1k_kir_coord", "kir_coord.fa"), + ): + path = self.mock_dir / name + path.touch() + setattr(pipeline, attr, path) + return pipeline + + def test_no_t1k_jobs_by_default(self): + """T1K runs only when its reference files are supplied""" + pipeline = self.create_pipeline() + job_names, _ = self._get_all_job_names(pipeline.build_dag()) + + assert not [name for name in job_names if name.startswith("t1k")] + + def test_t1k_jobs_with_fastq_input(self): + """T1K genotypes the deduplicated short reads""" + pipeline = self.enable_t1k(self.create_pipeline()) + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + for name in ( + "t1k-hla-extract", + "t1k-hla", + "t1k-kir-extract", + "t1k-kir", + ): + assert name in job_names, f"missing job: {name}" + + bwa_aln = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + extract = self._get_job(all_jobs, "t1k-hla-extract") + assert str(extract.shell) == ( + f"sentieon driver --input {bwa_aln} " + f"--reference {self.mock_ref} --thread_count {pipeline.cores} " + "--interval chr6:28510020-33480577 " + f"--algo ReadWriter {self.mock_dir}/sample_hla.bam" + ) + assert str(self._get_job(all_jobs, "t1k-hla").shell) == ( + f"run-t1k --abnormalUnmapFlag -t {pipeline.cores} " + "--preset hla-wgs " + f"-f {self.mock_dir}/hla_seq.fa " + f"-c {self.mock_dir}/hla_coord.fa " + f"--od {self.mock_dir}/output_hla " + f"-b {self.mock_dir}/sample_hla.bam" + ) + assert str(self._get_job(all_jobs, "t1k-kir").shell) == ( + f"run-t1k --abnormalUnmapFlag -t {pipeline.cores} " + "--preset kir-wgs " + f"-f {self.mock_dir}/kir_seq.fa " + f"-c {self.mock_dir}/kir_coord.fa " + f"--od {self.mock_dir}/output_kir " + f"-b {self.mock_dir}/sample_kir.bam" + ) + + assert self._get_dep_names(dag, all_jobs, "t1k-hla") == { + "t1k-hla-extract" + } + assert self._get_dep_names(dag, all_jobs, "t1k-hla-extract") == { + "dedup-bwa" + } + + def test_t1k_jobs_with_aligned_input(self): + """T1K extracts from the `--sr_aln` input without dependencies""" + pipeline = self.enable_t1k(self.create_aligned_pipeline()) + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + extract = self._get_job(all_jobs, "t1k-kir-extract") + assert str(extract.shell) == ( + f"sentieon driver --input {self.mock_sr_bam} " + f"--reference {self.mock_ref} --thread_count {pipeline.cores} " + "--interval chr19:53100000-55800000 " + f"--algo ReadWriter {self.mock_dir}/sample_kir.bam" + ) + assert self._get_dep_names(dag, all_jobs, "t1k-kir-extract") == set() + assert self._get_dep_names(dag, all_jobs, "t1k-kir") == { + "t1k-kir-extract" + } + + # The second-DAG stashes + + def test_stashes_with_fastq_input(self): + """The bwa alignment is stashed without a `--replace_rg` row""" + pipeline = self.create_pipeline() + pipeline.build_dag() + + bwa_aln = pathlib.Path( + str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + ) + assert pipeline.sr_alignments == [bwa_aln] + assert pipeline.sr_replace_rg is None + assert pipeline.lr_alignment == self.mock_lr_bam + + def test_stashes_with_aligned_input(self): + """The `--sr_aln` input is stashed with its `--replace_rg` row""" + pipeline = self.create_aligned_pipeline() + pipeline.build_dag() + + assert pipeline.sr_alignments == [self.mock_sr_bam] + assert pipeline.sr_replace_rg == [ + [pipeline._replace_rg_arg({"ID": "sr-rg1", "SM": "sample1"}, "0")] + ] + assert pipeline.sr_replace_rg == [ + [r"sr-rg1=ID:sr-rg1\tSM:sample1\tLR:0"] + ] + + def test_lr_alignment_stash_with_realignment(self): + """The realigned long-read alignment is stashed""" + pipeline = self.create_lr_realign_pipeline() + pipeline.build_dag() + + assert pipeline.lr_alignment == pathlib.Path( + str(self.mock_vcf).replace(".vcf.gz", "_mm2_sorted_0.cram") + ) + + def test_lr_alignment_stash_with_two_inputs(self): + """segdup-caller takes a single long-read alignment""" + lr_bam2 = self.mock_dir / "longreads2.bam" + lr_bam2.touch() + + pipeline = self.create_pipeline() + pipeline.lr_aln = [self.mock_lr_bam, lr_bam2] + pipeline.lr_readgroups = [ + [{"ID": "lr-rg1", "SM": "sample1"}], + [{"ID": "lr-rg2", "SM": "sample1"}], + ] + pipeline.build_dag() + + assert pipeline.lr_alignment is None + + # The second, sex-aware DAG + + def _build_second_dag(self, pipeline, sample_sex=SampleSex.FEMALE): + """Build both DAGs and return the second. + + The first DAG has to be built first so that the short-read + alignments and the ploidy JSON are stashed for the second. + """ + pipeline.build_dag() + pipeline.sample_sex = sample_sex + return pipeline.build_second_dag() + + def _second_dag_job(self, pipeline, name, sample_sex=SampleSex.FEMALE): + """Build both DAGs and return the named job of the second""" + dag = self._build_second_dag(pipeline, sample_sex) + _, all_jobs = self._get_all_job_names(dag) + return self._get_job(all_jobs, name) + + def enable_cnv(self, pipeline): + """Call CNVs with a bundle CNV model and a PAR BED file""" + par_bed = self.mock_dir / "par.bed" + par_bed.write_text("chrX\t10000\t2781479\n") + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + pipeline.cnv_par_bed = par_bed + return pipeline + + def script_path(self, name): + """A path to a packaged helper script, as the stages build it""" + return pathlib.Path(str(files("sentieon_cli.scripts").joinpath(name))) + + CNV_JOBS = ("cnvscope", "cnv-model-apply", "indel2cnv", "combine-cnv") + + def test_no_second_dag_without_callers(self): + """No second DAG when nothing consumes the sample sex""" + pipeline = self.create_pipeline() + + assert self._build_second_dag(pipeline) is None + + def test_no_second_dag_without_a_cnv_model(self): + """`--call_cnvs` without a bundle CNV model is rejected earlier""" + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = False + + assert self._build_second_dag(pipeline) is None + + def test_second_dag_for_each_caller(self): + """Every sex-aware caller on its own builds the second DAG""" + pipeline = self.create_pipeline() + pipeline.expansion_catalog = self.mock_bed + job_names, _ = self._get_all_job_names( + self._build_second_dag(pipeline) + ) + assert job_names == ["expansion-hunter"] + + pipeline = self.create_pipeline() + pipeline.segdup_caller = [] + job_names, _ = self._get_all_job_names( + self._build_second_dag(pipeline) + ) + assert job_names == ["segdup-caller"] + + pipeline = self.enable_cnv(self.create_pipeline()) + job_names, _ = self._get_all_job_names( + self._build_second_dag(pipeline) + ) + assert sorted(job_names) == sorted(self.CNV_JOBS) + + def test_cnv_jobs_in_the_second_dag(self): + """The CNV jobs and their dependencies""" + pipeline = self.enable_cnv(self.create_pipeline()) + first_dag = pipeline.build_dag() + + # CNV calling is sex-aware, so it is not in the first DAG + first_job_names, _ = self._get_all_job_names(first_dag) + for name in self.CNV_JOBS: + assert name not in first_job_names + + pipeline.sample_sex = SampleSex.FEMALE + dag = pipeline.build_second_dag() + job_names, all_jobs = self._get_all_job_names(dag) + assert sorted(job_names) == sorted(self.CNV_JOBS) + + assert self._get_dep_names(dag, all_jobs, "cnvscope") == set() + assert self._get_dep_names(dag, all_jobs, "indel2cnv") == set() + assert self._get_dep_names(dag, all_jobs, "cnv-model-apply") == { + "cnvscope" + } + assert self._get_dep_names(dag, all_jobs, "combine-cnv") == { + "cnv-model-apply", + "indel2cnv", + } + + for name in self.CNV_JOBS: + job = self._get_job(all_jobs, name) + assert job.task_name == "cnv", name + for name in ("cnvscope", "cnv-model-apply"): + assert self._get_job(all_jobs, name).threads == pipeline.cores + for name in ("indel2cnv", "combine-cnv"): + assert self._get_job(all_jobs, name).threads == 0 + + def test_cnvscope_command_with_fastq_input(self): + """CNVscope runs on the deduplicated short reads""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job(pipeline, "cnvscope") + + bwa_aln = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + probes = str(self.mock_vcf).replace(".vcf.gz", "_cnv.probes") + assert str(job.shell) == ( + f"sentieon driver --input {bwa_aln} " + f"--reference {self.mock_ref} --thread_count {pipeline.cores} " + f"--interval {self.mock_bed} " + f"--algo CNVscope --model {self.mock_bundle}/cnv.model " + "--sex F " + f"--dump_probes {probes} " + f"{self.mock_dir}/sample-cnvscope.vcf.gz" + ) + # The pipeline-generated alignment carries the LR attribute already + assert "--replace_rg" not in str(job.shell) + + def test_cnvscope_command_with_aligned_input(self): + """The `--sr_aln` input is used with its `--replace_rg` row""" + pipeline = self.enable_cnv(self.create_aligned_pipeline()) + job = self._second_dag_job(pipeline, "cnvscope") + + probes = str(self.mock_vcf).replace(".vcf.gz", "_cnv.probes") + assert str(job.shell) == ( + "sentieon driver " + r"--replace_rg 'sr-rg1=ID:sr-rg1\tSM:sample1\tLR:0' " + f"--input {self.mock_sr_bam} " + f"--reference {self.mock_ref} --thread_count {pipeline.cores} " + f"--interval {self.mock_bed} " + f"--algo CNVscope --model {self.mock_bundle}/cnv.model " + "--sex F " + f"--dump_probes {probes} " + f"{self.mock_dir}/sample-cnvscope.vcf.gz" + ) + + def test_cnvscope_without_a_bed(self): + """Without `--bed`, CNVscope runs across the whole reference""" + pipeline = self.enable_cnv(self.create_pipeline()) + pipeline.bed = None + job = self._second_dag_job(pipeline, "cnvscope") + + assert "--interval" not in str(job.shell) + + def test_cnv_model_apply_command(self): + """CNVModelApply filters the CNVscope output""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job(pipeline, "cnv-model-apply") + + assert str(job.shell) == ( + f"sentieon driver --reference {self.mock_ref} " + f"--thread_count {pipeline.cores} " + f"--algo CNVModelApply --model {self.mock_bundle}/cnv.model " + f"--vcf {self.mock_dir}/sample-cnvscope.vcf.gz " + f"{self.mock_dir}/sample-cnv_model_apply.vcf.gz" + ) + # The BED restricts CNVscope, not the model apply + assert "--interval" not in str(job.shell) + + def test_cnvscope_male_sample(self): + """A male sample is called with the PAR BED file""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job( + pipeline, "cnvscope", sample_sex=SampleSex.MALE + ) + + assert "--sex M" in str(job.shell) + assert f"--par {pipeline.cnv_par_bed}" in str(job.shell) + + def test_cnvscope_female_sample(self): + """A female sample does not need the PAR regions""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job( + pipeline, "cnvscope", sample_sex=SampleSex.FEMALE + ) + + assert "--sex F" in str(job.shell) + assert "--par" not in str(job.shell) + + def test_cnvscope_unknown_sex(self, messages): + """An unknown sex calls a diploid genome, with a warning""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job( + pipeline, "cnvscope", sample_sex=SampleSex.UNKNOWN + ) + + assert "--sex" not in str(job.shell) + assert "--par" not in str(job.shell) + assert any("diploid" in msg for msg in messages) + + def test_indel2cnv_command(self): + """The PangenomeSV output is converted to CNVs""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job(pipeline, "indel2cnv") + + sv_vcf = str(self.mock_vcf).replace(".vcf.gz", "_sv.vcf.gz") + assert str(job.shell) == ( + f"{sys.executable} {self.script_path('indel2cnv.py')} " + f"{self.mock_ref} {sv_vcf} " + f"{self.mock_dir}/sample-sv_cnv.vcf.gz -t {pipeline.cores}" + ) + + def test_combine_cnv_command(self): + """The CNV and converted SV calls are combined""" + pipeline = self.enable_cnv(self.create_pipeline()) + job = self._second_dag_job(pipeline, "combine-cnv") + + cnv_vcf = str(self.mock_vcf).replace(".vcf.gz", "_cnv.vcf.gz") + assert str(job.shell) == ( + f"{sys.executable} {self.script_path('combine_cnv.py')} " + f"--cnv {self.mock_dir}/sample-cnv_model_apply.vcf.gz " + f"--converted {self.mock_dir}/sample-sv_cnv.vcf.gz " + f"-o {cnv_vcf}" + ) + + def test_expansion_hunter_command(self): + """ExpansionHunter genotypes the short reads""" + catalog = self.mock_dir / "catalog.json" + catalog.touch() + + pipeline = self.create_pipeline() + pipeline.expansion_catalog = catalog + dag = self._build_second_dag(pipeline) + _, all_jobs = self._get_all_job_names(dag) + job = self._get_job(all_jobs, "expansion-hunter") + + bwa_aln = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + assert str(job.shell) == ( + f"ExpansionHunter --reads {bwa_aln} " + f"--reference {self.mock_ref} " + f"--variant-catalog {catalog} " + "--sex female " + f"--threads {pipeline.cores} " + f"--output-prefix {self.mock_dir}/output_expansion" + ) + assert job.task_name == "expansion-hunter" + assert job.threads == pipeline.cores + assert self._get_dep_names(dag, all_jobs, "expansion-hunter") == set() + + def test_segdup_command_with_fastq_input(self): + """segdup-caller uses the short and long reads, and one bundle""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = ["CFH", "CYP2D6", "SMN1"] + first_job_names, _ = self._get_all_job_names(pipeline.build_dag()) + assert "segdup-caller" not in first_job_names + + pipeline.sample_sex = SampleSex.MALE + dag = pipeline.build_second_dag() + _, all_jobs = self._get_all_job_names(dag) + job = self._get_job(all_jobs, "segdup-caller") + + bwa_aln = str(self.mock_vcf).replace(".vcf.gz", "_bwa_deduped.cram") + assert str(job.shell) == ( + f"segdup-caller --short {bwa_aln} " + f"--long {self.mock_lr_bam} " + f"--reference {self.mock_ref} " + f"--sr_model {self.mock_bundle} " + f"--lr_model {self.mock_bundle} " + f"--input_vcf {self.mock_vcf} " + "--sex male " + "--genes CFH,CYP2D6,SMN1 " + f"--outdir {self.mock_dir}/output_segdups" + ) + assert job.task_name == "segdup" + assert job.threads == pipeline.cores + assert self._get_dep_names(dag, all_jobs, "segdup-caller") == set() + + def test_segdup_without_genes(self): + """An empty gene list runs segdup-caller's own default set""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = [] + job = self._second_dag_job(pipeline, "segdup-caller") + + assert "--genes" not in str(job.shell) + assert "--sex female" in str(job.shell) + + def test_segdup_command_with_aligned_input(self): + """The `--sr_aln` input is passed to segdup-caller as-is""" + pipeline = self.create_aligned_pipeline() + pipeline.segdup_caller = [] + job = self._second_dag_job(pipeline, "segdup-caller") + + assert f"--short {self.mock_sr_bam} " in str(job.shell) + + def test_segdup_command_with_realigned_long_reads(self): + """The realigned long-read alignment reaches segdup-caller""" + pipeline = self.create_lr_realign_pipeline() + pipeline.segdup_caller = [] + job = self._second_dag_job(pipeline, "segdup-caller") + + lr_aln = str(self.mock_vcf).replace(".vcf.gz", "_mm2_sorted_0.cram") + assert f"--long {lr_aln} " in str(job.shell) + + # Second-DAG gates + + @pytest.mark.parametrize( + "call_cnvs,has_cnv_model,expected", + [ + (False, False, False), + (False, True, False), + (True, False, False), + (True, True, True), + ], + ) + def test_cnv_in_second_dag(self, call_cnvs, has_cnv_model, expected): + """CNVs are called with `--call_cnvs` and a bundle CNV model""" + pipeline = self.create_pipeline() + pipeline.call_cnvs = call_cnvs + pipeline.has_cnv_model = has_cnv_model + + assert pipeline._cnv_in_second_dag() is expected + + def test_needs_second_dag(self): + """Each sex-aware caller requires the second DAG""" + pipeline = self.create_pipeline() + assert pipeline._needs_second_dag() is False + + pipeline = self.create_pipeline() + pipeline.expansion_catalog = self.mock_bed + assert pipeline._needs_second_dag() is True + + pipeline = self.create_pipeline() + pipeline.segdup_caller = [] + assert pipeline._needs_second_dag() is True + + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + assert pipeline._needs_second_dag() is True + + # `--call_cnvs` without a CNV model is rejected during validation + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + assert pipeline._needs_second_dag() is False + + # Optional caller validation + + def test_validate_segdup_single_input(self): + """One short-read and one long-read alignment are accepted""" + pipeline = self.create_aligned_pipeline() + pipeline.segdup_caller = [] + + pipeline.validate_segdup() # no SystemExit + + def test_validate_segdup_two_sr_inputs(self): + """segdup-caller takes a single `--sr_aln` file""" + pipeline = self.create_aligned_pipeline() + pipeline.segdup_caller = [] + pipeline.sr_aln = [self.mock_sr_bam, self.mock_sr_bam] + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_segdup() + assert excinfo.value.code == 2 + + def test_validate_segdup_two_lr_inputs(self): + """segdup-caller takes a single `--lr_aln` file""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = ["CFH"] + pipeline.lr_aln = [self.mock_lr_bam, self.mock_lr_bam] + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_segdup() + assert excinfo.value.code == 2 + + def test_validate_segdup_skip_small_variants(self): + """segdup-caller reads the small-variant VCF""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = [] + pipeline.skip_small_variants = True + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_segdup() + assert excinfo.value.code == 2 + + def test_validate_segdup_without_the_argument(self): + """Nothing is validated without `--segdup_caller`""" + pipeline = self.create_pipeline() + pipeline.skip_small_variants = True + + pipeline.validate_segdup() # no SystemExit + + def test_validate_expansion_two_sr_inputs(self): + """ExpansionHunter takes a single `--sr_aln` file""" + pipeline = self.create_aligned_pipeline() + pipeline.expansion_catalog = self.mock_bed + pipeline.sr_aln = [self.mock_sr_bam, self.mock_sr_bam] + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_expansion() + assert excinfo.value.code == 2 + + def test_validate_expansion_single_input(self): + """A single short-read alignment is accepted""" + pipeline = self.create_aligned_pipeline() + pipeline.expansion_catalog = self.mock_bed + + pipeline.validate_expansion() # no SystemExit + + def test_validate_cnv_requires_sv_calling(self): + """indel2cnv reads the PangenomeSV output""" + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + pipeline.skip_svs = True + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_cnv() + assert excinfo.value.code == 2 + + def test_validate_cnv_requires_a_cnv_model(self): + """`--call_cnvs` needs a bundle with a 'cnv.model' file""" + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = False + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_cnv() + assert excinfo.value.code == 2 + + def test_validate_cnv_requires_a_par_bed(self): + # The mock reference is not a recognized build, so no packaged + # PAR BED file can be selected and validation stops the run + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_cnv() + assert excinfo.value.code == 2 + + def test_validate_cnv_accepts_the_par_bed_argument(self): + par_bed = self.mock_dir / "par.bed" + par_bed.write_text("chrX\t10000\t2781479\n") + + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + pipeline.par_bed = par_bed + + pipeline.validate_cnv() # no SystemExit + + assert pipeline.cnv_par_bed == par_bed + + def test_validate_cnv_needs_no_par_bed_without_call_cnvs(self): + pipeline = self.create_pipeline() + pipeline.has_cnv_model = True + + pipeline.validate_cnv() # no SystemExit + + assert pipeline.cnv_par_bed is None + + def test_validate_cnv_warns_without_a_bed(self): + """CNVscope runs across every contig without `--bed`""" + par_bed = self.mock_dir / "par.bed" + par_bed.write_text("chrX\t10000\t2781479\n") + + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + pipeline.par_bed = par_bed + pipeline.bed = None + + pipeline.validate_cnv() + + assert any( + "No `--bed` supplied" in str(call.args[0]) + for call in pipeline.logger.warning.call_args_list + ) + + def test_validate_cnv_does_not_warn_with_a_bed(self): + par_bed = self.mock_dir / "par.bed" + par_bed.write_text("chrX\t10000\t2781479\n") + + pipeline = self.create_pipeline() + pipeline.call_cnvs = True + pipeline.has_cnv_model = True + pipeline.par_bed = par_bed + + pipeline.validate_cnv() + + pipeline.logger.warning.assert_not_called() + + # Model bundle validation + + def bundle_pipeline(self, monkeypatch, members): + """A pipeline whose bundle holds `members`""" + pipeline = self.create_pipeline() + pipeline.skip_pop_vcf_id_check = True + + def fake_ar_load(path): + if str(path).endswith("bundle_info.json"): + return b'{"pipeline": "Hybrid pangenome"}' + return list(members) + + monkeypatch.setattr(hybrid_pangenome, "ar_load", fake_ar_load) + return pipeline + + BUNDLE_MEMBERS = [ + "dnascope.model", + "longreadsv.model", + "minimap2.model", + "extract.model", + "bwa.model", + ] + + def test_validate_bundle_sets_has_cnv_model(self, monkeypatch): + pipeline = self.bundle_pipeline( + monkeypatch, self.BUNDLE_MEMBERS + ["cnv.model"] + ) + pipeline.validate_bundle() + assert pipeline.has_cnv_model is True + + pipeline = self.bundle_pipeline(monkeypatch, self.BUNDLE_MEMBERS) + pipeline.validate_bundle() + assert pipeline.has_cnv_model is False + + def test_validate_bundle_requires_the_diploid_model(self, monkeypatch): + """segdup-caller reads the bundle's `diploid_model`""" + pipeline = self.bundle_pipeline(monkeypatch, self.BUNDLE_MEMBERS) + pipeline.segdup_caller = [] + + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_bundle() + assert excinfo.value.code == 2 + + pipeline = self.bundle_pipeline( + monkeypatch, self.BUNDLE_MEMBERS + ["diploid_model"] + ) + pipeline.segdup_caller = [] + pipeline.validate_bundle() # no SystemExit + + def test_validate_bundle_diploid_model_not_required_by_default( + self, monkeypatch + ): + pipeline = self.bundle_pipeline(monkeypatch, self.BUNDLE_MEMBERS) + pipeline.validate_bundle() # no SystemExit diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py new file mode 100644 index 0000000..f77205e --- /dev/null +++ b/tests/unit/test_job.py @@ -0,0 +1,116 @@ +""" +Unit tests for Job identity: task_name, job_id, and id numbering. +""" + +import os +import sys + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 +from sentieon_cli.util import sanitize # noqa: E402 + + +def _job(name, arg, task_name="test"): + """A job whose identity is its (unique) command argument.""" + return Job(Pipeline(Command("echo", arg)), name, task_name=task_name) + + +def test_task_name_is_required(): + with pytest.raises(TypeError): + Job( # type: ignore[call-arg] + Pipeline(Command("echo", "x")), "no-task" + ) + + +def test_ids_are_numbered_per_name(): + jobs = [_job("shard", str(i)) for i in range(3)] + assert [job.job_id for job in jobs] == ["shard-1", "shard-2", "shard-3"] + + +def test_each_name_has_its_own_counter(): + first = _job("dedup", "a") + other = _job("metrics", "b") + second = _job("dedup", "c") + + assert first.job_id == "dedup-1" + assert other.job_id == "metrics-1" + assert second.job_id == "dedup-2" + + +def test_names_differing_only_in_unsafe_characters_share_a_counter(): + # Log file names are sanitized, so ids that only differ in unsafe + # characters would name the same file and truncate each other. + first = _job("a b", "x") + second = _job("a-b", "y") + + assert (first.job_id, second.job_id) == ("a b-1", "a-b-2") + assert sanitize(first.job_id) == "a-b-1" + assert sanitize(second.job_id) == "a-b-2" + + +def test_reset_ids_restarts_the_sequence(): + assert _job("multiqc", "a").job_id == "multiqc-1" + assert _job("multiqc", "b").job_id == "multiqc-2" + + Job.reset_ids() + + assert _job("multiqc", "c").job_id == "multiqc-1" + + +def test_ids_keep_counting_across_two_dags(): + # The pangenome pipeline executes two DAGs in one process; ids must stay + # unique for the whole run, so they are never reset between DAGs. + def build_dag(tag): + dag = DAG() + for i in range(2): + dag.add_job(_job("calling", f"{tag}-{i}")) + return dag + + first = build_dag("a") + second = build_dag("b") + + ids = sorted( + job.job_id + for dag in (first, second) + for job in list(dag.ready_jobs) + list(dag.waiting_jobs) + ) + assert ids == ["calling-1", "calling-2", "calling-3", "calling-4"] + + +def test_task_name_and_job_id_are_not_part_of_identity(): + j1 = Job(Pipeline(Command("echo", "x")), "a", task_name="alignment") + j2 = Job(Pipeline(Command("echo", "x")), "b", task_name="cleanup") + + assert j1.job_id != j2.job_id + assert j1.task_name != j2.task_name + assert j1 == j2 + assert hash(j1) == hash(j2) + + +def test_jobs_with_different_task_names_still_collide_in_a_dag(): + dag = DAG() + dag.add_job(Job(Pipeline(Command("echo", "x")), "a", task_name="dedup")) + with pytest.raises(ValueError): + dag.add_job( + Job(Pipeline(Command("echo", "x")), "b", task_name="metrics") + ) + + +def test_repr_and_str_report_the_job_id(): + job = _job("locuscollector", "x") + assert repr(job) == "Job(locuscollector-1)" + assert str(job) == "Job(locuscollector-1)" + + +def test_task_name_is_stored(): + job = _job("dnascope", "x", task_name="variant-calling") + assert job.task_name == "variant-calling" diff --git a/tests/unit/test_job_logs.py b/tests/unit/test_job_logs.py new file mode 100644 index 0000000..841fa9e --- /dev/null +++ b/tests/unit/test_job_logs.py @@ -0,0 +1,215 @@ +""" +Unit tests for the per-process job log files +""" + +import os +import sys + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.executor import TAIL_LINES, _log_tail # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.run_logs import RunLogs # noqa: E402 +from sentieon_cli.shell_pipeline import ( # noqa: E402 + Command, + Context, + Pipeline, +) +from sentieon_cli.util import sanitize # noqa: E402 + + +def _sink(tmp_path, name: str = "bwa", task_name: str = "alignment"): + job = Job(Pipeline(Command("true")), name, task_name=task_name) + return RunLogs(tmp_path / "logs").job_sink(job) + + +def test_sanitize_keeps_safe_characters(): + assert sanitize("bwa-1.0_x") == "bwa-1.0_x" + assert sanitize("call vars/chr1:1-2") == "call-vars-chr1-1-2" + + +def test_stage_indices_follow_spawn_order_across_both_streams(tmp_path): + sink = _sink(tmp_path) + for handle in ( + sink.open_stderr(Command("a")), + sink.open_stdout(Command("b")), + sink.open_stderr(Command("c")), + ): + handle.close() + + assert [path.name for path in sink.log_paths()] == [ + "bwa-1.0.log", + "bwa-1.1.stdout.log", + "bwa-1.2.log", + ] + + +def test_both_streams_of_one_process_share_a_stage_index(tmp_path): + sink = _sink(tmp_path) + first = Command("a") + second = Command("b") + for handle in ( + sink.open_stderr(first), + sink.open_stdout(first), + sink.open_stderr(second), + sink.open_stdout(second), + ): + handle.close() + + assert [path.name for path in sink.log_paths()] == [ + "bwa-1.0.log", + "bwa-1.0.stdout.log", + "bwa-1.1.log", + "bwa-1.1.stdout.log", + ] + + +def test_the_task_directory_is_created_lazily(tmp_path): + sink = _sink(tmp_path) + assert not sink.log_dir.exists() + + sink.open_stderr(Command("true")).close() + + assert sink.log_dir == tmp_path / "logs" / "task_logs" / "alignment" + assert sink.log_dir.is_dir() + + +def test_path_components_are_sanitized(tmp_path): + sink = _sink(tmp_path, name="call vars/2", task_name="variant calling") + sink.open_stderr(Command("true")).close() + + path = sink.log_paths()[0] + assert path.parent.name == "variant-calling" + assert path.name == "call-vars-2-1.0.log" + + +def test_the_stderr_log_starts_with_a_header(tmp_path): + sink = _sink(tmp_path) + handle = sink.open_stderr(Command("echo", "hello world")) + handle.write(b"from the child\n") + handle.close() + + contents = sink.log_paths()[0].read_text() + assert "# timestamp: " in contents + assert "# task: alignment" in contents + assert "# job: bwa (bwa-1)" in contents + assert "# stage: 0" in contents + assert "# command: echo 'hello world'" in contents + # The header is flushed before the child inherits the FD, so child + # output always lands after it. + assert contents.splitlines()[-1] == "from the child" + + +def test_the_stdout_log_has_no_header(tmp_path): + sink = _sink(tmp_path) + handle = sink.open_stdout(Command("echo", "hello")) + handle.write(b"payload\n") + handle.close() + + assert sink.log_paths()[0].read_bytes() == b"payload\n" + + +def test_logs_are_matched_to_commands_by_identity(tmp_path): + sink = _sink(tmp_path) + first = Command("echo", "same") + second = Command("echo", "same") + assert first == second # equal, but two distinct processes + + sink.open_stderr(first).close() + sink.open_stderr(second).close() + + assert sink.stderr_log_for(first).name == "bwa-1.0.log" + assert sink.stderr_log_for(second).name == "bwa-1.1.log" + assert sink.stderr_log_for(Command("elsewhere")) is None + + +def test_stage_index_for_does_not_allocate(tmp_path): + sink = _sink(tmp_path) + unseen = Command("never-logged") + + assert sink.stage_index_for(unseen) is None + + logged = Command("logged") + sink.open_stderr(logged).close() + assert sink.stage_index_for(logged) == 0 + + +def test_stdout_logs_are_not_returned_as_stderr_logs(tmp_path): + sink = _sink(tmp_path) + command = Command("echo", "hi") + sink.open_stdout(command).close() + + assert sink.stderr_log_for(command) is None + + +def test_finalize_prunes_uninformative_logs_after_a_success(tmp_path): + sink = _sink(tmp_path) + header_only = Command("quiet") + empty_stdout = Command("quiet") + noisy = Command("loud") + sink.open_stderr(header_only).close() + sink.open_stdout(empty_stdout).close() + handle = sink.open_stderr(noisy) + handle.write(b"something happened\n") + handle.close() + + sink.finalize(success=True) + + assert not sink.stderr_log_for(header_only).exists() + assert not sink.log_paths()[1].exists() + assert sink.stderr_log_for(noisy).exists() + + +def test_finalize_keeps_every_log_after_a_failure(tmp_path): + sink = _sink(tmp_path) + sink.open_stderr(Command("quiet")).close() + sink.open_stdout(Command("quiet")).close() + + sink.finalize(success=False) + + assert all(path.exists() for path in sink.log_paths()) + + +def test_a_context_without_a_sink_captures_nothing(): + assert Context().log_sink is None + + +def test_log_tail_returns_the_last_lines(tmp_path): + log = tmp_path / "big.log" + log.write_text("".join(f"line {i}\n" for i in range(100))) + + tail = _log_tail(log) + + assert len(tail) == TAIL_LINES + assert tail[0] == "line 80" + assert tail[-1] == "line 99" + + +def test_log_tail_returns_a_short_file_whole(tmp_path): + log = tmp_path / "small.log" + log.write_text("one\ntwo\nthree\n") + + assert _log_tail(log) == ["one", "two", "three"] + + +def test_log_tail_only_reads_the_end_of_a_large_file(tmp_path): + log = tmp_path / "huge.log" + padding = "x" * 200 + with open(log, "w") as handle: + for i in range(2000): + handle.write(f"line {i} {padding}\n") + assert log.stat().st_size > 64 * 1024 + + tail = _log_tail(log) + + assert len(tail) == TAIL_LINES + assert tail[-1].startswith("line 1999 ") + assert not any(line.startswith("line 0 ") for line in tail) + + +def test_log_tail_of_a_missing_file_is_empty(tmp_path): + assert _log_tail(tmp_path / "gone.log") == [] diff --git a/tests/unit/test_metrics_files.py b/tests/unit/test_metrics_files.py new file mode 100644 index 0000000..ed5264f --- /dev/null +++ b/tests/unit/test_metrics_files.py @@ -0,0 +1,147 @@ +""" +Unit tests for the machine-readable metrics files +""" + +import os +import pathlib +import sys +from typing import List + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.executor import LocalExecutor # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.run_logs import ( # noqa: E402 + JOB_METRICS_COLUMNS, + PROCESS_METRICS_COLUMNS, + RunLogs, +) +from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + + +def _run(log_dir: pathlib.Path, pipeline: Pipeline) -> RunLogs: + """Run a single job with its metrics collected under ``log_dir``""" + dag = DAG() + dag.add_job(Job(pipeline, "metrics", task_name="metrics")) + run_logs = RunLogs(log_dir) + run_logs.create_dirs() + LocalExecutor(ThreadScheduler(dag, 1), run_logs=run_logs).execute() + return run_logs + + +def _rows(path: pathlib.Path) -> List[List[str]]: + """The rows of a metrics file, header included, split on tabs""" + lines = path.read_text().splitlines() + return [line.split("\t") for line in lines] + + +def _row_dict(path: pathlib.Path, index: int = 1) -> dict: + rows = _rows(path) + return dict(zip(rows[0], rows[index])) + + +def test_process_metrics_has_a_row_per_process(tmp_path): + run_logs = _run(tmp_path / "logs", Pipeline(Command("true"))) + + rows = _rows(run_logs.process_metrics) + assert rows[0] == PROCESS_METRICS_COLUMNS + assert len(rows) == 2 + assert len(rows[1]) == len(PROCESS_METRICS_COLUMNS) + row = dict(zip(rows[0], rows[1])) + assert row["job_id"] == "metrics-1" + assert row["task_name"] == "metrics" + assert row["exit_code"] == "0" + assert int(row["pid"]) > 0 + assert float(row["user_s"]) >= 0 + assert float(row["maxrss_mib"]) > 0 + assert row["command"] == "true" + + +def test_job_metrics_has_a_row_per_job(tmp_path): + # `wall_s` is rendered with two decimals, so the job has to take long + # enough to be visible there. + run_logs = _run(tmp_path / "logs", Pipeline(Command("sleep", "0.05"))) + + rows = _rows(run_logs.job_metrics) + assert rows[0] == JOB_METRICS_COLUMNS + assert len(rows) == 2 + row = dict(zip(rows[0], rows[1])) + assert row["job_id"] == "metrics-1" + assert row["status"] == "ok" + assert float(row["wall_s"]) > 0 + assert row["processes"] == "1" + assert row["threads"] == "1" + + +def test_process_rows_link_to_job_rows(tmp_path): + run_logs = _run( + tmp_path / "logs", + Pipeline(Command("printf", "x"), Command("cat")), + ) + + process_rows = _rows(run_logs.process_metrics)[1:] + assert len(process_rows) == 2 + header = PROCESS_METRICS_COLUMNS + pids = {row[header.index("pid")] for row in process_rows} + assert len(pids) == 2 + job_row = _row_dict(run_logs.job_metrics) + assert {row[header.index("job_id")] for row in process_rows} == { + job_row["job_id"] + } + assert job_row["processes"] == "2" + + +def test_stage_links_rows_to_task_log_files(tmp_path): + run_logs = _run(tmp_path / "logs", Pipeline(Command("true"))) + + # The log file itself may be gone: a successful job prunes logs that + # hold nothing but their header. + assert _row_dict(run_logs.process_metrics)["stage"] == "0" + + +def test_failed_job_still_records_metrics(tmp_path): + run_logs = _run(tmp_path / "logs", Pipeline(Command("false"))) + + process_row = _row_dict(run_logs.process_metrics) + assert process_row["exit_code"] == "1" + assert process_row["user_s"] != "" + assert process_row["maxrss_mib"] != "" + assert _row_dict(run_logs.job_metrics)["status"] == "failed" + + +def test_command_field_keeps_the_row_rectangular(tmp_path): + run_logs = _run(tmp_path / "logs", Pipeline(Command("printf", "a\tb\nc"))) + + rows = _rows(run_logs.process_metrics) + assert len(rows) == 2 + assert len(rows[1]) == len(PROCESS_METRICS_COLUMNS) + command = rows[1][PROCESS_METRICS_COLUMNS.index("command")] + assert "\t" not in command + assert "\n" not in command + + +def test_rerun_truncates_metrics_files(tmp_path): + log_dir = tmp_path / "logs" + _run(log_dir, Pipeline(Command("false"))) + + Job.reset_ids() + run_logs = _run(log_dir, Pipeline(Command("true"))) + + assert len(_rows(run_logs.process_metrics)) == 2 + assert _row_dict(run_logs.process_metrics)["exit_code"] == "0" + assert len(_rows(run_logs.job_metrics)) == 2 + assert _row_dict(run_logs.job_metrics)["status"] == "ok" + + +def test_metrics_files_start_as_headers_only(tmp_path): + run_logs = RunLogs(tmp_path / "logs") + run_logs.create_dirs() + + assert _rows(run_logs.process_metrics) == [PROCESS_METRICS_COLUMNS] + assert _rows(run_logs.job_metrics) == [JOB_METRICS_COLUMNS] diff --git a/tests/unit/test_pangenome_meta.py b/tests/unit/test_pangenome_meta.py new file mode 100644 index 0000000..aa37e42 --- /dev/null +++ b/tests/unit/test_pangenome_meta.py @@ -0,0 +1,394 @@ +""" +Unit tests for the pangenome metadata reader and its run-time checks. + +The binary fixtures in `tests/data/pangenome` are built by the +`generate.sh` script next to them; see that script for the `vg` commands +and for what each graph contains. +""" + +import pathlib +import struct +import subprocess +import sys + +import pytest + +from sentieon_cli.pangenome_meta import ( + HAPL_MAGIC, + PangenomeMetadataError, + check_gfa, + check_sample_gbz, + detect_pangenome_reference, + main, + read_gbz_metadata, + read_hapl_header, + unpack_bits, +) + +DATA = pathlib.Path(__file__).resolve().parents[1] / "data" / "pangenome" + +# The tiny graphs hold two reference samples over two 30 bp contigs +TINY_CONTIGS = ["chr1", "chr2"] + + +class TestUnpackBits: + """The packed-integer decoder behind every `simple-sds` vector""" + + def test_byte_wide_values(self): + words = struct.pack("` writes""" + with pytest.raises(PangenomeMetadataError) as excinfo: + check_sample_gbz( + DATA / "tiny_no_ref_tag.gbz", "GRCh38", "GRCh38#0#" + ) + message = str(excinfo.value) + assert "reference_samples=''" in message + assert "reference name may be wrong" in message + + def test_a_reference_sample_without_the_prefix_fails(self): + with pytest.raises( + PangenomeMetadataError, match="contig prefix 'GRCh38#1#'" + ): + check_sample_gbz( + DATA / "tiny_sampled.gbz", "GRCh38", "GRCh38#1#" + ) + + +class TestCheckGfa: + """The run-time check on the GFA `vg convert -Q` writes""" + + def test_an_rgfa_with_the_reference_tags_passes(self): + check_gfa(DATA / "rgfa_ok.gfa", "GRCh38", "GRCh38#0#") + + def test_a_header_that_still_lists_the_reference_fails(self): + with pytest.raises(PangenomeMetadataError) as excinfo: + check_gfa( + DATA / "rgfa_wrong_reference.gfa", "GRCh38", "GRCh38#0#" + ) + message = str(excinfo.value) + assert "RS:Z:CHM13 GRCh38" in message + assert "reference name may be wrong" in message + + def test_a_file_without_sn_tags_fails(self): + with pytest.raises(PangenomeMetadataError, match="SN:Z:GRCh38#0#"): + check_gfa(DATA / "rgfa_no_tags.gfa", "GRCh38", "GRCh38#0#") + + def test_a_file_without_a_header_line_is_tolerated(self): + check_gfa(DATA / "rgfa_no_header.gfa", "GRCh38", "GRCh38#0#") + + def test_the_wrong_contig_prefix_fails(self): + """The header passes, but no segment carries the prefix""" + with pytest.raises(PangenomeMetadataError, match="SN:Z:GRCh38#1#"): + check_gfa(DATA / "rgfa_ok.gfa", "GRCh38", "GRCh38#1#") + + def test_a_reference_name_in_the_header_fails_first(self): + """`RS:Z:` naming the reference is checked before the tags""" + with pytest.raises(PangenomeMetadataError, match="RS:Z:CHM13"): + check_gfa(DATA / "rgfa_ok.gfa", "CHM13", "CHM13#0#") + + def test_only_the_head_of_the_file_is_read(self): + """The tags are past `max_lines`, so the check gives up""" + with pytest.raises( + PangenomeMetadataError, match="first 1 lines" + ): + check_gfa( + DATA / "rgfa_ok.gfa", "GRCh38", "GRCh38#0#", max_lines=1 + ) + + +class TestCli: + """The `python -m sentieon_cli.pangenome_meta` entry point""" + + def test_detect_prints_the_reference(self, capsys): + assert main(["detect", "--gbz", str(DATA / "tiny.gbz")]) == 0 + out = capsys.readouterr().out + assert "pangenome_ref_name\tGRCh38" in out + assert "pangenome_contig_prefix\tGRCh38#0#" in out + assert "n_contigs\t2" in out + + def test_detect_fails_on_a_graph_without_a_backbone(self): + argv = ["detect", "--gbz", str(DATA / "tiny_no_backbone.gbz")] + assert main(argv) == 1 + + def test_check_gbz_passes(self): + argv = [ + "check-gbz", + "--gbz", + str(DATA / "tiny_sampled.gbz"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ] + assert main(argv) == 0 + + def test_check_gbz_fails(self): + argv = [ + "check-gbz", + "--gbz", + str(DATA / "tiny.gbz"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ] + assert main(argv) == 1 + + def test_check_gbz_fails_on_a_missing_file(self): + argv = [ + "check-gbz", + "--gbz", + str(DATA / "does-not-exist.gbz"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ] + assert main(argv) == 1 + + def test_check_gfa_passes(self): + argv = [ + "check-gfa", + "--gfa", + str(DATA / "rgfa_ok.gfa"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ] + assert main(argv) == 0 + + def test_check_gfa_fails(self): + argv = [ + "check-gfa", + "--gfa", + str(DATA / "rgfa_wrong_reference.gfa"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ] + assert main(argv) == 1 + + def test_check_gfa_honors_max_lines(self): + argv = [ + "check-gfa", + "--gfa", + str(DATA / "rgfa_ok.gfa"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + "--max_lines", + "1", + ] + assert main(argv) == 1 + + def test_the_module_runs_as_a_script(self): + """The DAG jobs run the module with `python -m`""" + failing = subprocess.run( + [ + sys.executable, + "-m", + "sentieon_cli.pangenome_meta", + "check-gbz", + "--gbz", + str(DATA / "tiny.gbz"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ], + capture_output=True, + text=True, + ) + assert failing.returncode == 1 + assert "reference name may be wrong" in failing.stderr + + passing = subprocess.run( + [ + sys.executable, + "-m", + "sentieon_cli.pangenome_meta", + "check-gbz", + "--gbz", + str(DATA / "tiny_sampled.gbz"), + "--reference_name", + "GRCh38", + "--contig_prefix", + "GRCh38#0#", + ], + capture_output=True, + text=True, + ) + assert passing.returncode == 0 diff --git a/tests/unit/test_pipeline_lifecycle.py b/tests/unit/test_pipeline_lifecycle.py new file mode 100644 index 0000000..adf67e9 --- /dev/null +++ b/tests/unit/test_pipeline_lifecycle.py @@ -0,0 +1,362 @@ +""" +Unit tests for the BasePipeline lifecycle (argument registration, logging +scope, temp-dir cleanup, and post-run checks). +""" + +import argparse +import logging +import os +import sys +from typing import Optional + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.dnascope import DNAscopePipeline # noqa: E402 +from sentieon_cli.dnascope_hybrid import ( # noqa: E402 + DNAscopeHybridPipeline, +) +from sentieon_cli.dnascope_longread import DNAscopeLRPipeline # noqa: E402 +from sentieon_cli.exceptions import DagExecutionError # noqa: E402 +from sentieon_cli.hybrid_pangenome import HybridPangenome # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.pipeline import BasePipeline # noqa: E402 +from sentieon_cli.sentieon_pangenome import SentieonPangenome # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + + +class _DummyPipeline(BasePipeline): + """A minimal concrete pipeline for exercising BasePipeline plumbing.""" + + def validate(self) -> None: + pass + + def configure(self) -> None: + pass + + def build_dag(self) -> DAG: + return DAG() + + +class _FailingPipeline(_DummyPipeline): + """A pipeline whose DAG construction raises.""" + + def build_dag(self) -> DAG: + raise RuntimeError("boom") + + +def _option_strings(parser): + return [opt for action in parser._actions for opt in action.option_strings] + + +def test_add_arguments_does_not_mutate_the_class_specs(): + before_reference = dict(BasePipeline.params["reference"]) + before_cores = dict(BasePipeline.params["cores"]) + + parser = argparse.ArgumentParser() + BasePipeline.add_arguments(parser) + + assert BasePipeline.params["reference"] == before_reference + assert BasePipeline.params["cores"] == before_cores + # In particular the specs keep their flags and gain no inferred type. + assert "flags" in BasePipeline.params["reference"] + assert "type" not in BasePipeline.params["cores"] + + +def test_add_arguments_keeps_short_flags_across_parsers(): + # Registering the same pipeline class on a second parser must produce + # the same options; the short flags used to be lost because the first + # call popped them from the shared class-level spec. + p1 = argparse.ArgumentParser() + p2 = argparse.ArgumentParser() + BasePipeline.add_arguments(p1) + BasePipeline.add_arguments(p2) + + for parser in (p1, p2): + opts = _option_strings(parser) + assert "-r" in opts + assert "--reference" in opts + assert "-t" in opts + assert "--cores" in opts + + +def test_setup_logging_does_not_touch_the_root_logger(): + root_level_before = logging.getLogger().level + package_logger = logging.getLogger("sentieon_cli") + package_level_before = package_logger.level + try: + pipeline = _DummyPipeline() + pipeline.setup_logging(argparse.Namespace(loglevel=logging.DEBUG)) + + assert logging.getLogger().level == root_level_before + assert package_logger.level == logging.DEBUG + finally: + package_logger.setLevel(package_level_before) + + +def test_main_cleans_up_tmpdir_when_build_dag_raises(tmp_path, monkeypatch): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _FailingPipeline() + + with pytest.raises(RuntimeError, match="boom"): + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert list(tmp_path.iterdir()) == [] # no leaked temp dir + + +def test_main_retains_tmpdir_on_failure_when_requested(tmp_path, monkeypatch): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _FailingPipeline() + pipeline.retain_tmpdir = True + + with pytest.raises(RuntimeError, match="boom"): + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert len(list(tmp_path.iterdir())) == 1 # kept for inspection + + +def test_main_cleans_up_tmpdir_on_success(tmp_path, monkeypatch): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _DummyPipeline() + pipeline.dry_run = True + + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert list(tmp_path.iterdir()) == [] + + +class _StubExecutor: + def __init__(self, jobs_with_errors=None): + self.jobs_with_errors = jobs_with_errors or [] + + +class _TwoDagPipeline(_DummyPipeline): + """A pipeline that counts the DAGs run through the `main` hook. + + `second_dag` is the DAG returned by the hook; `fail_on` is the 1-based + index of the `check_execution` call that raises. + """ + + def __init__(self, second_dag=None, fail_on=None): + super().__init__() + self.second_dag = second_dag + self.fail_on = fail_on + self.dags_run = 0 + self.check_calls = 0 + + def build_second_dag(self): + return self.second_dag + + def run(self, dag): + self.dags_run += 1 + return _StubExecutor() + + def check_execution(self, dag, executor): + self.check_calls += 1 + if self.fail_on == self.check_calls: + raise DagExecutionError("a job failed") + + +def test_main_runs_one_dag_when_the_second_dag_hook_returns_none( + tmp_path, monkeypatch +): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _TwoDagPipeline() + + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert pipeline.dags_run == 1 + assert pipeline.check_calls == 1 + assert list(tmp_path.iterdir()) == [] + + +def test_main_runs_the_dag_returned_by_the_second_dag_hook( + tmp_path, monkeypatch +): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _TwoDagPipeline(second_dag=DAG()) + + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert pipeline.dags_run == 2 + assert pipeline.check_calls == 2 + assert list(tmp_path.iterdir()) == [] + + +def test_main_cleans_up_tmpdir_when_the_second_dag_fails( + tmp_path, monkeypatch +): + # The second DAG runs inside the same try/finally, so the temp + # directory is removed on its failure path too + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _TwoDagPipeline(second_dag=DAG(), fail_on=2) + + with pytest.raises(DagExecutionError): + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert pipeline.dags_run == 2 + assert list(tmp_path.iterdir()) == [] + + +def test_hybrid_second_dag_hook_is_skipped_without_cnv_calling(): + from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline + + pipeline = DNAscopeHybridPipeline() + pipeline.setup_logging(argparse.Namespace(loglevel="WARNING")) + pipeline.skip_cnv = True + + assert pipeline.build_second_dag() is None + + +def test_check_execution_names_the_failed_jobs(): + pipeline = _DummyPipeline() + pipeline.setup_logging(argparse.Namespace(loglevel="WARNING")) + job = Job(Pipeline(Command("false")), "broken-step", task_name="test") + + # Jobs report themselves by job_id, not by name. + with pytest.raises(DagExecutionError, match=r"Job\(broken-step-1\)"): + pipeline.check_execution(DAG(), _StubExecutor([job])) + + +class _StubPangenome(SentieonPangenome): + """A SentieonPangenome with its run stages stubbed out. + + `fail_on` is the 1-based index of the `check_execution` call that + raises, mimicking a job failure in the first or the second DAG. + """ + + def __init__(self, fail_on=None): + super().__init__() + self.fail_on = fail_on + self.check_calls = 0 + self.dags_built = [] + + def validate(self) -> None: + pass + + def configure(self) -> None: + pass + + def build_dag(self) -> DAG: + self.dags_built.append("first") + self.ploidy_json = self.tmp_dir.joinpath("ploidy.json") + return DAG() + + def build_second_dag(self) -> Optional[DAG]: + if not self._needs_second_dag(): + return None + self.dags_built.append("second") + return DAG() + + def get_sex(self, ploidy_json) -> None: + pass + + def run(self, dag): + return _StubExecutor() + + def check_execution(self, dag, executor): + self.check_calls += 1 + if self.fail_on == self.check_calls: + raise DagExecutionError("a job failed") + + +def _stub_pangenome(monkeypatch, tmp_path, **kwargs): + """A stubbed pangenome pipeline writing temp dirs into `tmp_path`""" + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + return _StubPangenome(**kwargs) + + +def test_pangenome_main_cleans_up_tmpdir_when_the_first_dag_fails( + tmp_path, monkeypatch +): + # SentieonPangenome runs the base two-DAG flow; the temp directory + # has to be removed on the failure path too. + pipeline = _stub_pangenome(monkeypatch, tmp_path, fail_on=1) + + with pytest.raises(DagExecutionError): + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert list(tmp_path.iterdir()) == [] # no leaked temp dir + + +def test_pangenome_main_cleans_up_tmpdir_when_the_second_dag_fails( + tmp_path, monkeypatch +): + pipeline = _stub_pangenome(monkeypatch, tmp_path, fail_on=2) + pipeline.segdup_caller = [] # request the second DAG + + with pytest.raises(DagExecutionError): + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert pipeline.dags_built == ["first", "second"] + assert list(tmp_path.iterdir()) == [] + + +def test_pangenome_main_runs_the_second_dag_for_cnv_calling( + tmp_path, monkeypatch +): + # CNV calling is sex-aware, so it alone triggers the second DAG + pipeline = _stub_pangenome(monkeypatch, tmp_path) + pipeline.call_svs = True + pipeline.has_cnv_model = True + + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert pipeline.dags_built == ["first", "second"] + assert pipeline.check_calls == 2 + + +def test_pangenome_main_retains_tmpdir_on_failure_when_requested( + tmp_path, monkeypatch +): + pipeline = _stub_pangenome(monkeypatch, tmp_path, fail_on=1) + pipeline.retain_tmpdir = True + + with pytest.raises(DagExecutionError): + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert len(list(tmp_path.iterdir())) == 1 # kept for inspection + + +def test_pangenome_main_cleans_up_tmpdir_on_success(tmp_path, monkeypatch): + pipeline = _stub_pangenome(monkeypatch, tmp_path) + pipeline.segdup_caller = [] # request the second DAG + + pipeline.main(argparse.Namespace(loglevel="WARNING")) + + assert pipeline.dags_built == ["first", "second"] + assert pipeline.check_calls == 2 + assert list(tmp_path.iterdir()) == [] + + +def test_check_execution_flags_unexecuted_jobs(): + pipeline = _DummyPipeline() + pipeline.setup_logging(argparse.Namespace(loglevel="WARNING")) + dag = DAG() + dag.add_job( + Job(Pipeline(Command("echo", "x")), "unexecuted", task_name="test") + ) + + with pytest.raises(DagExecutionError, match=r"Job\(unexecuted-1\)"): + pipeline.check_execution(dag, _StubExecutor()) + + +@pytest.mark.parametrize( + "pipeline_cls", + [ + DNAscopePipeline, + DNAscopeLRPipeline, + DNAscopeHybridPipeline, + SentieonPangenome, + HybridPangenome, + ], +) +def test_no_pipeline_overrides_main(pipeline_cls): + """Every pipeline runs through `BasePipeline.main`""" + assert "main" not in pipeline_cls.__dict__ diff --git a/tests/unit/test_pipeline_params.py b/tests/unit/test_pipeline_params.py new file mode 100644 index 0000000..d0fb54e --- /dev/null +++ b/tests/unit/test_pipeline_params.py @@ -0,0 +1,48 @@ +""" +Every pipeline argument must be an attribute of a freshly built pipeline. + +`BasePipeline.handle_arguments` raises when a `params` key is missing from +the instance, and argparse leaves an unset argument as `None`, which it +skips -- so the `__init__` value is the effective default. These tests keep +the two in step. +""" + +from typing import Type + +import pytest + +from sentieon_cli.dnascope import DNAscopePipeline +from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline +from sentieon_cli.dnascope_longread import DNAscopeLRPipeline +from sentieon_cli.hybrid_pangenome import HybridPangenome +from sentieon_cli.pipeline import BasePipeline +from sentieon_cli.sentieon_pangenome import SentieonPangenome + +PIPELINES = [ + DNAscopePipeline, + DNAscopeLRPipeline, + DNAscopeHybridPipeline, + SentieonPangenome, + HybridPangenome, +] + + +@pytest.mark.parametrize("cls", PIPELINES, ids=lambda c: c.__name__) +def test_every_argument_is_an_attribute(cls: Type[BasePipeline]): + """`handle_arguments` needs an attribute per param and positional""" + attributes = set(vars(cls())) + arguments = set(cls.params) | set(cls.positionals) + assert arguments <= attributes, sorted(arguments - attributes) + + +@pytest.mark.parametrize("cls", PIPELINES, ids=lambda c: c.__name__) +def test_argparse_defaults_match_the_attributes(cls: Type[BasePipeline]): + """An argparse default must match the value `__init__` assigns""" + pipeline = cls() + specs = list(cls.params.items()) + list(cls.positionals.items()) + mismatched = { + key: (getattr(pipeline, key), spec["default"]) + for key, spec in specs + if "default" in spec and getattr(pipeline, key) != spec["default"] + } + assert not mismatched diff --git a/tests/unit/test_pipeline_validation.py b/tests/unit/test_pipeline_validation.py index b3ae579..ba5a546 100644 --- a/tests/unit/test_pipeline_validation.py +++ b/tests/unit/test_pipeline_validation.py @@ -8,7 +8,8 @@ import sys import os import json -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, patch # Add the parent directory to the path so we can import sentieon_cli sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) @@ -16,6 +17,13 @@ from sentieon_cli.dnascope import DNAscopePipeline from sentieon_cli.dnascope_longread import DNAscopeLRPipeline from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline +from sentieon_cli.hybrid_pangenome import HybridPangenome +from sentieon_cli.pangenome_meta import ( + PangenomeMetadataError, + PangenomeReference, +) +from sentieon_cli.sentieon_pangenome import SentieonPangenome +from sentieon_cli.util import VcfHeaderError, set_bwt_max_mem from tests.utils.test_helpers import create_mock_args @@ -53,8 +61,8 @@ def test_valid_configuration_with_bam_input(self): self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle self.pipeline.sample_input = [self.mock_bam] - self.pipeline.sr_r1_fastq = [] - self.pipeline.sr_readgroups = [] + self.pipeline.r1_fastq = [] + self.pipeline.readgroups = [] # Should not raise any exceptions try: @@ -83,8 +91,8 @@ def test_missing_inputs_raises_error(self): self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle self.pipeline.sample_input = [] - self.pipeline.sr_r1_fastq = [] - self.pipeline.sr_readgroups = [] + self.pipeline.r1_fastq = [] + self.pipeline.readgroups = [] with pytest.raises(SystemExit): self.pipeline.validate() @@ -96,8 +104,8 @@ def test_invalid_output_extension_raises_error(self): self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle self.pipeline.sample_input = [self.mock_bam] - self.pipeline.sr_r1_fastq = [] - self.pipeline.sr_readgroups = [] + self.pipeline.r1_fastq = [] + self.pipeline.readgroups = [] with pytest.raises(SystemExit): self.pipeline.validate() @@ -108,11 +116,19 @@ def test_mismatched_fastq_readgroups_raises_error(self): self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle self.pipeline.sample_input = [] - self.pipeline.sr_r1_fastq = [self.mock_fastq] - self.pipeline.sr_readgroups = [] # Empty readgroups with non-empty fastq + # Two fastq files, but only one readgroup. Both must be non-empty to + # get past the "supply --sample_input or --r1_fastq" guard. + self.pipeline.r1_fastq = [self.mock_fastq, self.mock_fastq] + self.pipeline.readgroups = ["@RG\\tID:a\\tSM:s"] - with pytest.raises(SystemExit): - self.pipeline.validate() + with patch.object(self.pipeline.logger, "error") as mock_error: + with pytest.raises(SystemExit) as excinfo: + self.pipeline.validate() + + assert excinfo.value.code == 2 + mock_error.assert_any_call( + "The number of readgroups does not equal the number of fastq files" + ) def test_skip_multiqc_when_skip_metrics(self): """Test that skip_multiqc is set when skip_metrics is True""" @@ -120,8 +136,8 @@ def test_skip_multiqc_when_skip_metrics(self): self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle self.pipeline.sample_input = [self.mock_bam] - self.pipeline.sr_r1_fastq = [] - self.pipeline.sr_readgroups = [] + self.pipeline.r1_fastq = [] + self.pipeline.readgroups = [] self.pipeline.skip_metrics = True self.pipeline.skip_multiqc = False @@ -197,7 +213,7 @@ def test_valid_configuration_with_fastq_input(self): self.pipeline.output_vcf = self.mock_vcf self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle - self.pipeline.lr_aln = [] + self.pipeline.sample_input = [] self.pipeline.fastq = [self.mock_fastq] self.pipeline.readgroups = ["@RG\\tID:test\\tSM:sample"] @@ -271,22 +287,29 @@ def test_haploid_bed_requires_diploid_bed(self): self.pipeline.output_vcf = self.mock_vcf self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle - self.pipeline.lr_aln = [self.mock_bam] + self.pipeline.sample_input = [self.mock_bam] self.pipeline.fastq = [] self.pipeline.readgroups = [] self.pipeline.haploid_bed = self.mock_bed self.pipeline.bed = None # No diploid bed self.pipeline.skip_small_variants = False - with pytest.raises(SystemExit): - self.pipeline.validate() + with patch.object(self.pipeline.logger, "error") as mock_error: + with pytest.raises(SystemExit) as excinfo: + self.pipeline.validate() + + assert excinfo.value.code == 2 + mock_error.assert_any_call( + "Please supply a BED file of diploid regions to distinguish " + "haploid and diploid regions of the genome." + ) def test_mismatched_fastq_readgroups_raises_error(self): """Test that mismatched FASTQ and readgroup counts raise an error""" self.pipeline.output_vcf = self.mock_vcf self.pipeline.reference = self.mock_ref self.pipeline.model_bundle = self.mock_bundle - self.pipeline.lr_aln = [] + self.pipeline.sample_input = [] self.pipeline.fastq = [self.mock_fastq] self.pipeline.readgroups = [] # Empty readgroups with non-empty fastq @@ -425,27 +448,41 @@ def test_total_input_size_calculation(self): fastq_file.write_bytes(b"y" * 500) # 0.5KB pipeline.sample_input = [bam_file] - pipeline.sr_r1_fastq = [fastq_file] - pipeline.sr_r2_fastq = [] + pipeline.r1_fastq = [fastq_file] + pipeline.r2_fastq = [] total_size = pipeline.total_input_size() assert total_size == 1500 # 1KB + 0.5KB - def test_memory_calculation_logic(self): - """Test bwt_max_mem calculation logic""" - pipeline = DNAscopePipeline() + def test_bwt_max_mem_override(self, monkeypatch): + """`--bwt_max_mem` short-circuits the calculation""" + monkeypatch.delenv("bwt_max_mem", raising=False) - # Setup logging first - args = create_mock_args() - pipeline.setup_logging(args) - - # Test with specific memory values - pipeline.bwt_max_mem = "10G" - pipeline.set_bwt_max_mem(0, 1) - - import os + assert set_bwt_max_mem(0, 1, override="10G") == "10G" assert os.environ.get("bwt_max_mem") == "10G" + @pytest.mark.parametrize( + "total_input_size,n_alignment_jobs,expected", + [ + (0, 1, "54G"), # 64 - 4 - 0 = 60; 60 / 1 - 6 = 54 + (0, 2, "24G"), # 60 / 2 - 6 = 24 + (30 * 1024**3, 1, "0G"), # 60 - 30 * 2.3 < 0, floored at 0 + ], + ) + def test_bwt_max_mem_calculation( + self, monkeypatch, total_input_size, n_alignment_jobs, expected + ): + """bwt_max_mem is derived from the memory and the input size""" + monkeypatch.setenv("bwt_max_mem", "unset") + + with patch( + "sentieon_cli.util.total_memory", return_value=64 * 1024**3 + ): + result = set_bwt_max_mem(total_input_size, n_alignment_jobs) + + assert result == expected + assert os.environ.get("bwt_max_mem") == expected + class TestValidateBwaIndex: """Test BasePipeline.validate_bwa_index()""" @@ -484,5 +521,490 @@ def test_partial_index_files_missing_raises(self): assert excinfo.value.code == 2 +PANGENOME_PIPELINES = [SentieonPangenome, HybridPangenome] + +GRCH38_REFERENCE = PangenomeReference( + ref_name="GRCh38", + contig_prefix="GRCh38#0#", + contigs=["chr1", "chr2"], + reference_samples=["GRCh38", "CHM13"], + n_paths=2, +) +CHM13_REFERENCE = PangenomeReference( + ref_name="CHM13", + contig_prefix="CHM13#0#", + contigs=["chr1", "chr2"], + reference_samples=["CHM13", "GRCh38"], + n_paths=2, +) + + +@pytest.mark.parametrize( + "cls", PANGENOME_PIPELINES, ids=lambda c: c.__name__ +) +class TestResolvePangenomeReference: + """`BasePangenome.resolve_pangenome_reference` in both pipelines. + + The reference name and contig prefix come from the graph's own + metadata, so a graph built against the wrong reference cannot reach + `vg haplotypes --set-reference`, which would silently drop every + reference path. + """ + + def _pipeline(self, cls, tmp_path, **attributes): + pipeline = cls() + pipeline.logger = MagicMock() + pipeline.gbz = tmp_path / "graph.gbz" + pipeline.hapl = tmp_path / "graph.hapl" + pipeline.gbz.touch() + pipeline.hapl.touch() + pipeline.fai_data = {"chr1": {"length": 10}, "chr2": {"length": 10}} + for key, value in attributes.items(): + setattr(pipeline, key, value) + return pipeline + + def _detect(self, reference): + return patch( + "sentieon_cli.base_pangenome.detect_pangenome_reference", + return_value=reference, + ) + + def _build(self, build): + return patch( + "sentieon_cli.base_pangenome.detect_reference_build", + return_value=build, + ) + + def _hapl(self, chains=2): + return patch( + "sentieon_cli.base_pangenome.read_hapl_header", + return_value=SimpleNamespace(top_level_chains=chains), + ) + + def test_the_detected_reference_populates_both_attributes( + self, cls, tmp_path + ): + pipeline = self._pipeline(cls, tmp_path) + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "GRCh38" + assert pipeline.pangenome_contig_prefix == "GRCh38#0#" + assert pipeline.pangenome_reference is GRCH38_REFERENCE + + def test_a_chm13_graph_resolves_to_the_chm13_prefix( + self, cls, tmp_path + ): + pipeline = self._pipeline(cls, tmp_path) + with self._detect(CHM13_REFERENCE), self._build("chm13"), ( + self._hapl() + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "CHM13" + assert pipeline.pangenome_contig_prefix == "CHM13#0#" + + def test_a_matching_override_is_accepted(self, cls, tmp_path): + pipeline = self._pipeline( + cls, + tmp_path, + pangenome_ref_name="GRCh38", + pangenome_contig_prefix="GRCh38#0#", + ) + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "GRCh38" + + def test_a_mismatched_name_override_exits(self, cls, tmp_path): + pipeline = self._pipeline(cls, tmp_path, pangenome_ref_name="hg38") + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + with pytest.raises(SystemExit) as excinfo: + pipeline.resolve_pangenome_reference() + assert excinfo.value.code == 2 + pipeline.logger.error.assert_called() + + def test_a_mismatched_prefix_override_exits(self, cls, tmp_path): + pipeline = self._pipeline( + cls, tmp_path, pangenome_contig_prefix="GRCh38#1#" + ) + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + with pytest.raises(SystemExit) as excinfo: + pipeline.resolve_pangenome_reference() + assert excinfo.value.code == 2 + + def test_the_skip_flag_honors_a_mismatched_override( + self, cls, tmp_path + ): + pipeline = self._pipeline( + cls, + tmp_path, + pangenome_ref_name="hg38", + skip_pangenome_name_checks=True, + ) + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "hg38" + assert pipeline.pangenome_contig_prefix == "GRCh38#0#" + pipeline.logger.warning.assert_called() + pipeline.logger.error.assert_not_called() + + def test_a_contig_missing_from_the_fai_exits(self, cls, tmp_path): + pipeline = self._pipeline(cls, tmp_path) + pipeline.fai_data = {"chr1": {"length": 10}} + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + with pytest.raises(SystemExit) as excinfo: + pipeline.resolve_pangenome_reference() + assert excinfo.value.code == 2 + message = pipeline.logger.error.call_args[0] + assert "missing from the reference FASTA index" in message[0] + assert "chr2" in message + + def test_the_skip_flag_downgrades_a_missing_contig(self, cls, tmp_path): + pipeline = self._pipeline( + cls, tmp_path, skip_pangenome_name_checks=True + ) + pipeline.fai_data = {"chr1": {"length": 10}} + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + pipeline.resolve_pangenome_reference() + pipeline.logger.warning.assert_called() + + def test_a_reference_build_mismatch_exits(self, cls, tmp_path): + """A CHM13 graph with an hg38 reference FASTA. + + Contig names alone cannot catch this: both builds name their + chromosomes chr1..chrM. + """ + pipeline = self._pipeline(cls, tmp_path) + with self._detect(CHM13_REFERENCE), self._build("hg38"), ( + self._hapl() + ): + with pytest.raises(SystemExit) as excinfo: + pipeline.resolve_pangenome_reference() + assert excinfo.value.code == 2 + assert "reference build" in pipeline.logger.error.call_args[0][0] + + def test_an_unknown_reference_build_is_only_logged(self, cls, tmp_path): + pipeline = self._pipeline(cls, tmp_path) + with self._detect(GRCH38_REFERENCE), self._build(None), ( + self._hapl() + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "GRCh38" + pipeline.logger.error.assert_not_called() + + def test_a_hapl_chain_count_mismatch_only_warns(self, cls, tmp_path): + pipeline = self._pipeline(cls, tmp_path) + with self._detect(GRCH38_REFERENCE), self._build("hg38"), ( + self._hapl(chains=195) + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "GRCh38" + pipeline.logger.warning.assert_called() + + def test_an_unreadable_hapl_exits(self, cls, tmp_path): + pipeline = self._pipeline(cls, tmp_path) + hapl_error = PangenomeMetadataError("not a vg haplotypes file") + with self._detect(GRCH38_REFERENCE), self._build("hg38"), patch( + "sentieon_cli.base_pangenome.read_hapl_header", + side_effect=hapl_error, + ): + with pytest.raises(SystemExit) as excinfo: + pipeline.resolve_pangenome_reference() + assert excinfo.value.code == 2 + + def test_a_dry_run_falls_back_to_the_defaults(self, cls, tmp_path): + """Mock graph files cannot be parsed, but dry runs still work""" + pipeline = self._pipeline(cls, tmp_path, dry_run=True) + error = PangenomeMetadataError("not a GBZ pangenome file") + with patch( + "sentieon_cli.base_pangenome.detect_pangenome_reference", + side_effect=error, + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "GRCh38" + assert pipeline.pangenome_contig_prefix == "GRCh38#0#" + assert pipeline.pangenome_reference is None + + def test_a_dry_run_keeps_the_overrides(self, cls, tmp_path): + pipeline = self._pipeline( + cls, + tmp_path, + dry_run=True, + pangenome_ref_name="CHM13", + pangenome_contig_prefix="CHM13#0#", + ) + error = PangenomeMetadataError("not a GBZ pangenome file") + with patch( + "sentieon_cli.base_pangenome.detect_pangenome_reference", + side_effect=error, + ): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "CHM13" + assert pipeline.pangenome_contig_prefix == "CHM13#0#" + + def test_an_unreadable_graph_exits_outside_a_dry_run( + self, cls, tmp_path + ): + pipeline = self._pipeline(cls, tmp_path) + error = PangenomeMetadataError("not a GBZ pangenome file") + with patch( + "sentieon_cli.base_pangenome.detect_pangenome_reference", + side_effect=error, + ): + with pytest.raises(SystemExit) as excinfo: + pipeline.resolve_pangenome_reference() + assert excinfo.value.code == 2 + + def test_both_overrides_survive_an_unreadable_graph(self, cls, tmp_path): + pipeline = self._pipeline( + cls, + tmp_path, + pangenome_ref_name="CHM13", + pangenome_contig_prefix="CHM13#0#", + ) + error = PangenomeMetadataError("not a GBZ pangenome file") + with patch( + "sentieon_cli.base_pangenome.detect_pangenome_reference", + side_effect=error, + ), self._hapl(): + pipeline.resolve_pangenome_reference() + assert pipeline.pangenome_ref_name == "CHM13" + assert pipeline.pangenome_contig_prefix == "CHM13#0#" + pipeline.logger.warning.assert_called() + + +@pytest.mark.parametrize( + "cls", PANGENOME_PIPELINES, ids=lambda c: c.__name__ +) +class TestPangenomeContigLengthChecks: + """`BasePangenome.validate_pangenome_contig_lengths` in both pipelines. + + The pop VCF is compared with the reference FASTA rather than with a + table of GRCh38 lengths, so the check covers every backbone. + """ + + def _pipeline(self, cls, reference=GRCH38_REFERENCE, **attributes): + pipeline = cls() + pipeline.logger = MagicMock() + pipeline.pangenome_reference = reference + pipeline.pangenome_ref_name = ( + None if reference is None else reference.ref_name + ) + pipeline.fai_data = { + "chr1": {"length": 1000}, + "chr2": {"length": 2000}, + } + pipeline.pop_vcf_contigs = {"chr1": 1000, "chr2": 2000} + for key, value in attributes.items(): + setattr(pipeline, key, value) + return pipeline + + def test_matching_lengths_pass(self, cls): + pipeline = self._pipeline(cls) + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.error.assert_not_called() + # The INFO line reports how many backbone contigs were compared + assert pipeline.logger.info.call_args[0][1] == 2 + + def test_a_mismatching_length_exits(self, cls): + pipeline = self._pipeline(cls) + pipeline.pop_vcf_contigs["chr2"] = 2001 + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_pangenome_contig_lengths() + assert excinfo.value.code == 2 + reported = pipeline.logger.error.call_args[0][-1] + assert "chr2" in reported + assert "2000" in reported and "2001" in reported + + def test_a_contig_missing_from_the_pop_vcf_exits(self, cls): + pipeline = self._pipeline(cls) + del pipeline.pop_vcf_contigs["chr2"] + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_pangenome_contig_lengths() + assert excinfo.value.code == 2 + assert "absent" in pipeline.logger.error.call_args[0][-1] + + def test_a_contig_without_a_length_exits(self, cls): + pipeline = self._pipeline(cls) + pipeline.pop_vcf_contigs["chr2"] = None + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_pangenome_contig_lengths() + assert excinfo.value.code == 2 + assert "no length" in pipeline.logger.error.call_args[0][-1] + + def test_skip_contig_checks_bypasses_the_check(self, cls): + pipeline = self._pipeline(cls, skip_contig_checks=True) + pipeline.pop_vcf_contigs["chr1"] = 1 + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.error.assert_not_called() + + def test_the_name_check_flag_does_not_bypass_the_check(self, cls): + """Only `--skip_contig_checks` turns this check off""" + pipeline = self._pipeline(cls, skip_pangenome_name_checks=True) + pipeline.pop_vcf_contigs["chr1"] = 1 + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_pangenome_contig_lengths() + assert excinfo.value.code == 2 + + def test_a_chm13_backbone_is_checked(self, cls): + """The old GRCh38-only check skipped every other backbone""" + pipeline = self._pipeline(cls, reference=CHM13_REFERENCE) + pipeline.pop_vcf_contigs["chr1"] = 248956422 + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_pangenome_contig_lengths() + assert excinfo.value.code == 2 + assert "CHM13" in pipeline.logger.error.call_args[0] + + def test_a_matching_chm13_backbone_passes(self, cls): + pipeline = self._pipeline(cls, reference=CHM13_REFERENCE) + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.error.assert_not_called() + + def test_empty_pop_vcf_contigs_warn_and_skip(self, cls): + pipeline = self._pipeline(cls, pop_vcf_contigs={}) + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.warning.assert_called() + pipeline.logger.error.assert_not_called() + + def test_empty_pop_vcf_contigs_only_debug_in_a_dry_run(self, cls): + pipeline = self._pipeline(cls, pop_vcf_contigs={}, dry_run=True) + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.warning.assert_not_called() + pipeline.logger.debug.assert_called() + + def test_an_undetected_backbone_checks_the_shared_contigs(self, cls): + """Without graph metadata, the fai/pop VCF intersection is used""" + pipeline = self._pipeline(cls, reference=None) + pipeline.pangenome_ref_name = "GRCh38" + pipeline.pop_vcf_contigs = {"chr1": 1001, "chr3": 30} + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_pangenome_contig_lengths() + assert excinfo.value.code == 2 + reported = pipeline.logger.error.call_args[0][-1] + assert "chr1" in reported and "chr2" not in reported + + def test_a_contig_missing_from_the_fai_is_not_reported_twice(self, cls): + """`check_pangenome_reference_fasta` already reported it""" + pipeline = self._pipeline(cls) + del pipeline.fai_data["chr2"] + del pipeline.pop_vcf_contigs["chr2"] + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.error.assert_not_called() + assert pipeline.logger.info.call_args[0][1] == 1 + + +@pytest.mark.parametrize( + "cls", PANGENOME_PIPELINES, ids=lambda c: c.__name__ +) +class TestPopVcfHeader: + """`load_pop_vcf_header` and `check_pop_vcf_id` in both pipelines""" + + HEADER = ( + "##fileformat=VCFv4.2\n" + "##SentieonVcfID=population-test-20260101\n" + "##contig=\n" + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" + ) + + def _pipeline(self, cls, **attributes): + pipeline = cls() + pipeline.logger = MagicMock() + for key, value in attributes.items(): + setattr(pipeline, key, value) + return pipeline + + def _pop_vcf(self, tmp_path, text=None): + pop_vcf = tmp_path / "pop.vcf.gz" + pop_vcf.write_text(self.HEADER if text is None else text) + return pop_vcf + + def test_the_header_and_the_contigs_are_parsed(self, cls, tmp_path): + pipeline = self._pipeline(cls) + pipeline.load_pop_vcf_header(self._pop_vcf(tmp_path)) + assert pipeline.pop_vcf_contigs == {"chr1": 1000} + assert pipeline.pop_vcf_header is not None + + def test_no_pop_vcf_leaves_both_empty(self, cls): + pipeline = self._pipeline(cls) + pipeline.load_pop_vcf_header(None) + assert pipeline.pop_vcf_contigs == {} + assert pipeline.pop_vcf_header is None + + def test_a_dry_run_falls_back_on_an_unparsable_pop_vcf( + self, cls, tmp_path + ): + """Unit tests and dry runs supply placeholder files""" + pipeline = self._pipeline(cls, dry_run=True) + placeholder = tmp_path / "placeholder.vcf.gz" + placeholder.touch() + pipeline.load_pop_vcf_header(placeholder) + assert pipeline.pop_vcf_header is None + assert pipeline.pop_vcf_contigs == {} + pipeline.logger.debug.assert_called() + pipeline.logger.error.assert_not_called() + + def test_an_unparsable_pop_vcf_exits_outside_a_dry_run( + self, cls, tmp_path + ): + pipeline = self._pipeline(cls) + placeholder = tmp_path / "placeholder.vcf.gz" + placeholder.touch() + with pytest.raises(SystemExit) as excinfo: + pipeline.load_pop_vcf_header(placeholder) + assert excinfo.value.code == 2 + + def test_a_missing_pop_vcf_exits_outside_a_dry_run(self, cls, tmp_path): + pipeline = self._pipeline(cls) + with pytest.raises(SystemExit) as excinfo: + pipeline.load_pop_vcf_header(tmp_path / "absent.vcf.gz") + assert excinfo.value.code == 2 + + def test_a_matching_vcf_id_passes(self, cls, tmp_path): + pipeline = self._pipeline(cls) + pipeline.load_pop_vcf_header(self._pop_vcf(tmp_path)) + pipeline.check_pop_vcf_id("population-test-20260101") + pipeline.logger.error.assert_not_called() + + def test_a_mismatching_vcf_id_exits_in_a_dry_run(self, cls, tmp_path): + """The ID check no longer waits for a real run""" + pipeline = self._pipeline(cls, dry_run=True) + pipeline.load_pop_vcf_header(self._pop_vcf(tmp_path)) + with pytest.raises(SystemExit) as excinfo: + pipeline.check_pop_vcf_id("population-other-20240101") + assert excinfo.value.code == 2 + + def test_an_unread_header_skips_the_vcf_id_check(self, cls): + pipeline = self._pipeline(cls, dry_run=True) + pipeline.check_pop_vcf_id("population-test-20260101") + pipeline.logger.error.assert_not_called() + pipeline.logger.debug.assert_called() + + def test_an_unparsable_header_does_not_raise_in_a_dry_run( + self, cls, tmp_path + ): + """The whole fall-back path, as a `--dry_run` walks it""" + pipeline = self._pipeline(cls, dry_run=True) + with patch( + "sentieon_cli.base_pangenome.read_vcf_header", + side_effect=VcfHeaderError("no VCF header"), + ): + pipeline.load_pop_vcf_header(tmp_path / "pop.vcf.gz") + pipeline.check_pop_vcf_id("population-test-20260101") + pipeline.validate_pangenome_contig_lengths() + pipeline.logger.error.assert_not_called() + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/test_run_logs.py b/tests/unit/test_run_logs.py new file mode 100644 index 0000000..a9c84ff --- /dev/null +++ b/tests/unit/test_run_logs.py @@ -0,0 +1,328 @@ +""" +Unit tests for the run log directory and the logging plumbing +""" + +import argparse +import logging +import os +import pathlib +import sys +from typing import List + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli import logging as cli_logging # noqa: E402 +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.exceptions import DagExecutionError # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.pipeline import BasePipeline # noqa: E402 +from sentieon_cli.run_logs import RunLogs # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 +from sentieon_cli.util import __version__ # noqa: E402 + +PACKAGE_LOGGER = "sentieon_cli" + + +class _DummyPipeline(BasePipeline): + """A minimal concrete pipeline for exercising the logging plumbing.""" + + def validate(self) -> None: + pass + + def configure(self) -> None: + pass + + def build_dag(self) -> DAG: + return DAG() + + +class _FailingPipeline(_DummyPipeline): + """A pipeline whose DAG construction raises.""" + + def build_dag(self) -> DAG: + raise RuntimeError("boom") + + +class _ValidatingPipeline(_DummyPipeline): + """A pipeline that checks its output path, as the real ones do.""" + + def validate(self) -> None: + self.validate_output_vcf() + + +class _ErroredExecutor: + """A stand-in executor that finished with failed jobs.""" + + def __init__(self, jobs: List[Job]) -> None: + self.jobs_with_errors = jobs + + +class _RecordingHandler(logging.Handler): + """Collects formatted messages from the package logger.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.messages: List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def messages(): + """Capture the package logger, which does not propagate to the root.""" + handler = _RecordingHandler() + package_logger = logging.getLogger(PACKAGE_LOGGER) + package_logger.addHandler(handler) + yield handler.messages + package_logger.removeHandler(handler) + + +def _args(loglevel: str = "INFO") -> argparse.Namespace: + return argparse.Namespace(loglevel=loglevel) + + +def _start(pipeline: BasePipeline, loglevel: str = "INFO") -> None: + """Set up logging the way `main` does, minus the run itself.""" + pipeline.setup_logging(_args(loglevel)) + pipeline.start_run_logs() + + +def test_module_loggers_delegate_to_the_package_logger(): + module_logger = cli_logging.get_logger("sentieon_cli.example") + + assert module_logger.handlers == [] + assert module_logger.propagate + assert not logging.getLogger(PACKAGE_LOGGER).propagate + + +def test_default_log_dir_is_derived_from_the_output_vcf(tmp_path): + pipeline = _DummyPipeline() + pipeline.output_vcf = tmp_path / "sample.vcf.gz" + _start(pipeline) + + assert pipeline.run_logs is not None + assert pipeline.run_logs.log_dir == tmp_path / "sample_logs" + + +def test_default_log_dir_only_strips_the_trailing_suffix(tmp_path): + # `str.replace` would mangle a name with '.vcf.gz' in the middle. + pipeline = _DummyPipeline() + pipeline.output_vcf = tmp_path / "a.vcf.gz.rerun.vcf.gz" + _start(pipeline) + + assert pipeline.run_logs.log_dir == tmp_path / "a.vcf.gz.rerun_logs" + + +def test_explicit_log_dir_skips_the_output_suffix_check(tmp_path): + pipeline = _DummyPipeline() + pipeline.output_vcf = tmp_path / "sample.bcf" + pipeline.log_dir = tmp_path / "elsewhere" + _start(pipeline) + + assert pipeline.run_logs.log_dir == tmp_path / "elsewhere" + assert (tmp_path / "elsewhere" / "run.log").is_file() + + +def test_invalid_output_suffix_exits_before_creating_a_log_dir(tmp_path): + pipeline = _DummyPipeline() + pipeline.output_vcf = tmp_path / "sample.bcf" + + with pytest.raises(SystemExit) as excinfo: + _start(pipeline) + + assert excinfo.value.code == 2 + assert list(tmp_path.iterdir()) == [] + + +def test_an_invalid_output_path_creates_no_log_dir(tmp_path, messages): + # `validate` runs first, so a typo in the output path is rejected before + # anything is written -- including the log directory next to it. + pipeline = _ValidatingPipeline() + pipeline.output_vcf = tmp_path / "typo" / "sample.vcf.gz" + + with pytest.raises(SystemExit) as excinfo: + pipeline.main(_args()) + + assert excinfo.value.code == 2 + assert pipeline.run_logs is None + assert list(tmp_path.iterdir()) == [] + assert any("status: failed" in msg for msg in messages) + + +def test_a_rejected_rerun_keeps_the_previous_runs_logs(tmp_path): + log_dir = tmp_path / "logs" + first = _ValidatingPipeline() + first.output_vcf = tmp_path / "sample.vcf.gz" + first.log_dir = log_dir + _start(first) + stale = first.run_logs.task_logs / "alignment" + stale.mkdir(parents=True) + (stale / "bwa-1.0.log").write_text("from the previous run") + first.run_logs.close() + run_log = (log_dir / "run.log").read_text() + + second = _ValidatingPipeline() + second.output_vcf = tmp_path / "typo" / "sample.vcf.gz" + second.log_dir = log_dir + with pytest.raises(SystemExit): + second.main(_args()) + + assert (stale / "bwa-1.0.log").read_text() == "from the previous run" + assert (log_dir / "run.log").read_text() == run_log + + +def test_a_log_dir_colliding_with_a_file_exits_cleanly(tmp_path, messages): + collision = tmp_path / "logs" + collision.write_text("not a directory") + pipeline = _DummyPipeline() + pipeline.log_dir = collision + + with pytest.raises(SystemExit) as excinfo: + _start(pipeline) + + assert excinfo.value.code == 2 + assert pipeline.run_logs is None + assert collision.read_text() == "not a directory" + assert any( + "Could not prepare the log directory" in msg for msg in messages + ) + + +def test_setup_wipes_stale_task_logs_but_keeps_the_log_dir(tmp_path): + log_dir = tmp_path / "logs" + stale = log_dir / "task_logs" / "alignment" + stale.mkdir(parents=True) + (stale / "bwa-1.0.log").write_text("from the previous run") + (log_dir / "unrelated.txt").write_text("keep me") + + RunLogs(log_dir).create_dirs() + + assert (log_dir / "unrelated.txt").read_text() == "keep me" + assert (log_dir / "task_logs").is_dir() + assert list((log_dir / "task_logs").iterdir()) == [] + + +def test_command_txt_records_the_invocation(tmp_path, monkeypatch): + monkeypatch.setattr( + sys, "argv", ["sentieon-cli", "dnascope", "a b.vcf.gz"] + ) + run_logs = RunLogs(tmp_path / "logs") + run_logs.create_dirs() + run_logs.write_command() + + contents = run_logs.command_txt.read_text() + assert "command: sentieon-cli dnascope 'a b.vcf.gz'" in contents + assert f"version: {__version__}" in contents + assert f"directory: {pathlib.Path.cwd()}" in contents + assert "timestamp: " in contents + + +def test_run_log_captures_debug_while_the_console_stays_at_info(tmp_path): + pipeline = _DummyPipeline() + pipeline.output_vcf = tmp_path / "sample.vcf.gz" + _start(pipeline) + logging.getLogger("sentieon_cli.example").debug("a debug record") + run_log = pipeline.run_logs.run_log + pipeline.run_logs.close() + + contents = run_log.read_text() + assert "a debug record" in contents + assert "Starting sentieon-cli version" in contents + assert f"Writing logs to: {pipeline.run_logs.log_dir}" in contents + assert cli_logging._console_handler.level == logging.INFO + + +def test_repeated_setup_does_not_accumulate_handlers(tmp_path): + package_logger = logging.getLogger(PACKAGE_LOGGER) + before = len(package_logger.handlers) + + pipeline = None + for i in range(3): + pipeline = _DummyPipeline() + pipeline.output_vcf = tmp_path / f"sample{i}.vcf.gz" + _start(pipeline) + assert len(package_logger.handlers) == before + 1 + + pipeline.run_logs.close() + assert len(package_logger.handlers) == before + + +def test_rerun_truncates_the_previous_run_log(tmp_path): + log_dir = tmp_path / "logs" + pipeline = None + for _ in range(2): + pipeline = _DummyPipeline() + pipeline.log_dir = log_dir + _start(pipeline) + run_log = pipeline.run_logs.run_log + pipeline.run_logs.close() + + banner = "Starting sentieon-cli version" + assert run_log.read_text().count(banner) == 1 + + +def test_dry_run_skips_file_logging(tmp_path): + pipeline = _DummyPipeline() + pipeline.dry_run = True + pipeline.output_vcf = tmp_path / "sample.vcf.gz" + _start(pipeline) + + assert pipeline.run_logs is None + assert list(tmp_path.iterdir()) == [] + + +def test_bare_pipeline_setup_logging_creates_nothing(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + pipeline = _DummyPipeline() + _start(pipeline) + + assert pipeline.run_logs is None + assert list(tmp_path.iterdir()) == [] + + +def test_check_execution_names_the_task_log_dir(tmp_path): + pipeline = _DummyPipeline() + pipeline.log_dir = tmp_path / "logs" + _start(pipeline) + job = Job(Pipeline(Command("false")), "boom", task_name="failing") + + with pytest.raises(DagExecutionError) as excinfo: + pipeline.check_execution(DAG(), _ErroredExecutor([job])) + pipeline.run_logs.close() + + message = str(excinfo.value) + assert "Job(boom-1)" in message + assert str(tmp_path / "logs" / "task_logs") in message + + +def test_end_of_run_message_on_success(tmp_path, monkeypatch, messages): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _DummyPipeline() + pipeline.dry_run = True + + pipeline.main(_args()) + + assert any("status: succeeded" in msg for msg in messages) + + +def test_end_of_run_message_on_failure(tmp_path, monkeypatch, messages): + monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) + pipeline = _FailingPipeline() + pipeline.log_dir = tmp_path / "logs" + + with pytest.raises(RuntimeError, match="boom"): + pipeline.main(_args()) + + run_log = pipeline.run_logs.run_log + pipeline.run_logs.close() + + assert any("status: failed" in msg for msg in messages) + assert any(str(tmp_path / "logs") in msg for msg in messages) + assert "status: failed" in run_log.read_text() diff --git a/tests/unit/test_rusage_metrics.py b/tests/unit/test_rusage_metrics.py new file mode 100644 index 0000000..a7c0b53 --- /dev/null +++ b/tests/unit/test_rusage_metrics.py @@ -0,0 +1,99 @@ +""" +Unit tests for the per-process resource-usage metrics +""" + +import logging +import os +import pathlib +import resource +import sys +from typing import List + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.executor import ( # noqa: E402 + LocalExecutor, + _maxrss_bytes, +) +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.run_logs import RunLogs # noqa: E402 +from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + + +class _RecordingHandler(logging.Handler): + """Collects formatted messages from the package logger.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.messages: List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def messages(): + """Capture the package logger, which does not propagate to the root.""" + handler = _RecordingHandler() + package_logger = logging.getLogger("sentieon_cli") + package_logger.addHandler(handler) + yield handler.messages + package_logger.removeHandler(handler) + + +def _run_one_job(log_dir: pathlib.Path) -> None: + """Run a single trivial job with its output captured under ``log_dir``""" + dag = DAG() + dag.add_job(Job(Pipeline(Command("true")), "metrics", task_name="metrics")) + run_logs = RunLogs(log_dir) + run_logs.create_dirs() + LocalExecutor(ThreadScheduler(dag, 1), run_logs=run_logs).execute() + + +def test_each_process_of_a_job_reports_its_rusage(tmp_path, messages): + _run_one_job(tmp_path / "logs") + + rusage_lines = [m for m in messages if m.startswith("rusage for ")] + assert len(rusage_lines) == 1 + line = rusage_lines[0] + assert "Job(metrics-1)" in line + for field in ("utime=", "stime=", "maxrss=", "minflt=", "nivcsw="): + assert field in line + + +def test_a_finished_job_reports_its_aggregated_metrics(tmp_path, messages): + _run_one_job(tmp_path / "logs") + + finished = [m for m in messages if m.startswith("Finished command in")] + assert len(finished) == 1 + assert "user " in finished[0] + assert "sys " in finished[0] + assert "max proc RSS" in finished[0] + + +def test_maxrss_is_reported_in_bytes(): + ru = resource.getrusage(resource.RUSAGE_SELF) + + maxrss = _maxrss_bytes(ru) + + assert isinstance(maxrss, int) + assert maxrss > 0 + + +def test_maxrss_units_follow_the_platform(monkeypatch): + """macOS reports bytes; everywhere else reports kibibytes.""" + ru = resource.getrusage(resource.RUSAGE_SELF) + + monkeypatch.setattr(sys, "platform", "darwin") + assert _maxrss_bytes(ru) == ru.ru_maxrss + + monkeypatch.setattr(sys, "platform", "linux") + assert _maxrss_bytes(ru) == ru.ru_maxrss * 1024 diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py new file mode 100644 index 0000000..20327ae --- /dev/null +++ b/tests/unit/test_scheduler.py @@ -0,0 +1,114 @@ +""" +Unit tests for the ThreadScheduler. +""" + +import os +import sys + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.exceptions import DagExecutionError # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + + +def _job(name, threads=1, resources=None): + return Job( + Pipeline(Command(name)), name, threads, resources, task_name="test" + ) + + +def test_schedules_independent_jobs_together(): + dag = DAG() + j1 = _job("j1", threads=1) + j2 = _job("j2", threads=1) + dag.add_job(j1) + dag.add_job(j2) + + scheduler = ThreadScheduler(dag, threads=4) + assert scheduler.start() == {j1, j2} + + +def test_thread_budget_gates_scheduling(): + dag = DAG() + j1 = _job("j1", threads=2) + j2 = _job("j2", threads=2) + dag.add_job(j1) + dag.add_job(j2) + + # A 2-thread budget only fits one 2-thread job at a time. + scheduler = ThreadScheduler(dag, threads=2) + + first = scheduler.start() + assert len(first) == 1 + running = next(iter(first)) + + # The second job starts only once the first frees its threads. + second = scheduler.job_finished(running) + assert len(second) == 1 + remaining = next(iter(second)) + + assert {running, remaining} == {j1, j2} + + +def test_named_resource_gates_scheduling(): + dag = DAG() + j1 = _job("j1", threads=1, resources={"node0": 1}) + j2 = _job("j2", threads=1, resources={"node0": 1}) + dag.add_job(j1) + dag.add_job(j2) + + # Plenty of threads, but only a single node0 token. + scheduler = ThreadScheduler(dag, threads=8, resources={"node0": 1}) + + first = scheduler.start() + assert len(first) == 1 + running = next(iter(first)) + + second = scheduler.job_finished(running) + assert len(second) == 1 + assert {running, next(iter(second))} == {j1, j2} + + +def test_resource_requirements_ignored_when_unmanaged(): + # Jobs may request a resource the scheduler does not track; those + # requirements are simply ignored. + dag = DAG() + j1 = _job("j1", threads=1, resources={"gpu": 1}) + j2 = _job("j2", threads=1, resources={"gpu": 1}) + dag.add_job(j1) + dag.add_job(j2) + + scheduler = ThreadScheduler(dag, threads=8) + assert scheduler.start() == {j1, j2} + + +def test_oversized_thread_request_is_rejected(): + dag = DAG() + dag.add_job( + Job(Pipeline(Command("echo", "hi")), "big", 100, task_name="test") + ) + with pytest.raises(DagExecutionError): + ThreadScheduler(dag, threads=2).start() + + +def test_oversized_resource_request_is_rejected(): + dag = DAG() + job = Job( + Pipeline(Command("echo", "hi")), + "greedy", + 1, + {"node0": 3}, + task_name="test", + ) + dag.add_job(job) + with pytest.raises(DagExecutionError): + ThreadScheduler(dag, threads=8, resources={"node0": 1}).start() diff --git a/tests/unit/test_scheduler_contract.py b/tests/unit/test_scheduler_contract.py new file mode 100644 index 0000000..4454575 --- /dev/null +++ b/tests/unit/test_scheduler_contract.py @@ -0,0 +1,111 @@ +""" +Conformance tests for the BaseScheduler contract. + +``drive`` mimics what an executor does: ask the scheduler for the initial +batch, then report each job finished (FIFO) and collect whatever that +unblocks, until the DAG drains. Any BaseScheduler implementation can be +dropped into ``drive`` to check that it honors the contract; the +dependency-order tests run against every scheduler in ``SCHEDULERS``. +""" + +import os +import sys + +import pytest + +# Add the parent directory to the path so we can import sentieon_cli +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), +) + +from sentieon_cli.dag import DAG # noqa: E402 +from sentieon_cli.exceptions import DagExecutionError # noqa: E402 +from sentieon_cli.job import Job # noqa: E402 +from sentieon_cli.scheduler import ThreadScheduler # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + +SCHEDULERS = [ + pytest.param(lambda dag: ThreadScheduler(dag, threads=8), id="thread"), +] + + +def _job(name, threads=1, resources=None): + return Job( + Pipeline(Command(name)), name, threads, resources, task_name="test" + ) + + +def drive(scheduler): + """Drive a scheduler to completion; return the job finish order.""" + order = [] + queue = list(scheduler.start()) + while queue: + job = queue.pop(0) + order.append(job) + queue.extend(scheduler.job_finished(job)) + return order + + +@pytest.mark.parametrize("make_scheduler", SCHEDULERS) +def test_drives_a_diamond_to_completion(make_scheduler): + dag = DAG() + top = _job("top") + left = _job("left") + right = _job("right") + bottom = _job("bottom") + dag.add_job(top) + dag.add_job(left, {top}) + dag.add_job(right, {top}) + dag.add_job(bottom, {left, right}) + + order = drive(make_scheduler(dag)) + + # Every job runs exactly once. + assert len(order) == len(set(order)) + assert sorted(j.name for j in order) == [ + "bottom", + "left", + "right", + "top", + ] + # Dependencies precede their dependents. + pos = {job: i for i, job in enumerate(order)} + assert pos[top] < pos[left] < pos[bottom] + assert pos[top] < pos[right] < pos[bottom] + + +@pytest.mark.parametrize("make_scheduler", SCHEDULERS) +def test_drives_a_chain_to_completion(make_scheduler): + dag = DAG() + a = _job("a") + b = _job("b") + c = _job("c") + dag.add_job(a) + dag.add_job(b, {a}) + dag.add_job(c, {b}) + + order = drive(make_scheduler(dag)) + assert [j.name for j in order] == ["a", "b", "c"] + + +@pytest.mark.parametrize("make_scheduler", SCHEDULERS) +def test_job_finished_on_unknown_job_raises(make_scheduler): + dag = DAG() + dag.add_job(_job("a")) + scheduler = make_scheduler(dag) + scheduler.start() + + with pytest.raises(DagExecutionError): + scheduler.job_finished(_job("never-scheduled")) + + +def test_budget_constrained_run_still_completes(): + # ThreadScheduler-specific: a 1-thread budget serializes two independent + # jobs, but driving to completion still runs both exactly once. + dag = DAG() + dag.add_job(_job("a", threads=1)) + dag.add_job(_job("b", threads=1)) + + order = drive(ThreadScheduler(dag, threads=1)) + assert sorted(j.name for j in order) == ["a", "b"] diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index e6c1914..9d96fe2 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -11,10 +11,12 @@ from unittest.mock import patch, MagicMock # Add the parent directory to the path to import sentieon_cli -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) from sentieon_cli.sentieon_pangenome import SentieonPangenome -from sentieon_cli.base_pangenome import SampleSex +from sentieon_cli.util import SampleSex from sentieon_cli.dag import DAG @@ -53,10 +55,10 @@ def setup_method(self): def create_pipeline(self): """Create a SentieonPangenome pipeline for testing""" pipeline = SentieonPangenome() - + # Setup mocks pipeline.logger = MagicMock() - + # Configure arguments pipeline.output_vcf = self.mock_vcf pipeline.reference = self.mock_ref @@ -72,14 +74,18 @@ def create_pipeline(self): pipeline.skip_pangenome_name_checks = True pipeline.skip_pop_vcf_id_check = True pipeline.tmp_dir = self.mock_dir - + + # The pangenome reference, normally resolved by validate() + pipeline.pangenome_ref_name = "GRCh38" + pipeline.pangenome_contig_prefix = "GRCh38#0#" + # Mock parsing fai pipeline.fai_data = {"chr1": {"length": 1000}} pipeline.shards = [MagicMock()] pipeline.shards[0].contig = "chr1" pipeline.shards[0].start = 1 pipeline.shards[0].stop = 1000 - + # Mock pop vcf contigs pipeline.pop_vcf_contigs = {"chr1": 1000} @@ -96,7 +102,7 @@ def create_fastq_pipeline(self): """Create a fastq-input SentieonPangenome pipeline for testing. The default fixture uses BAM/CRAM input, which does not exercise the - dedup/metrics branch of ``build_first_dag``. This variant provides + dedup/metrics branch of ``build_dag``. This variant provides FASTQ input so the bwa/mm2 dedup jobs are created. """ mock_r1 = self.mock_dir / "sample_R1.fastq.gz" @@ -117,7 +123,7 @@ def test_dedup_metrics_output(self): """Dedup on the primary (bwa) alignment emits a --metrics file that lands in the metrics directory scanned by MultiQC.""" pipeline = self.create_fastq_pipeline() - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() job_names, all_jobs = self._get_all_job_names(dag) assert "dedup-bwa" in job_names @@ -133,20 +139,33 @@ def test_dedup_metrics_output(self): mm2_dedup = next(j for j in all_jobs if j.name == "dedup-mm2") assert "--metrics" not in str(mm2_dedup.shell) + def test_metrics_input_bwa_only(self): + """The metrics job reads only the bwa alignment; including the mm2 + alignment would count the extracted reads it re-aligns twice.""" + pipeline = self.create_fastq_pipeline() + dag = pipeline.build_dag() + + _, all_jobs = self._get_all_job_names(dag) + metrics_job = next(j for j in all_jobs if j.name == "metrics") + metrics_cmd = str(metrics_job.shell) + assert "output_bwa_deduped.cram" in metrics_cmd + assert "output_mm2_deduped.cram" not in metrics_cmd + def test_dedup_metrics_skipped(self): """No Dedup --metrics output when metrics collection is skipped.""" pipeline = self.create_fastq_pipeline() pipeline.skip_metrics = True - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() _, all_jobs = self._get_all_job_names(dag) bwa_dedup = next(j for j in all_jobs if j.name == "dedup-bwa") assert "--metrics" not in str(bwa_dedup.shell) def test_model_apply_default(self): - """Test that model apply job is created by default""" + """Model apply runs by default and writes the + intermediate, not the final output VCF.""" pipeline = self.create_pipeline() - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() # Verify DAG is created assert isinstance(dag, DAG) @@ -155,11 +174,19 @@ def test_model_apply_default(self): job_names = [job.name for job in dag.waiting_jobs.keys()] assert "model-apply" in job_names + apply_job = next( + j for j in dag.waiting_jobs if j.name == "model-apply" + ) + cmd_str = str(apply_job.shell) + assert str(self.snv_apply_vcf(pipeline)) in cmd_str + # The final output is written by the AD-update job + assert str(pipeline.output_vcf) not in cmd_str + def test_skip_model_apply(self): """Test that model apply job is skipped when requested""" pipeline = self.create_pipeline() pipeline.skip_model_apply = True - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() # Verify DAG is created assert isinstance(dag, DAG) @@ -167,73 +194,149 @@ def test_skip_model_apply(self): # Check that model-apply job is NOT present job_names = [job.name for job in dag.waiting_jobs] assert "model-apply" not in job_names - - # Verify that the concat job writes to the final output + + # Verify that the concat job writes the intermediate, + # which the AD-update job then rewrites to the final output # Find the merge-trim-concat job concat_job = None for job in dag.waiting_jobs: if job.name == "merge-trim-concat": concat_job = job break - + assert concat_job is not None - # Check that the first argument (output file) is set to the final output VCF - # Check that the output file is in the arguments - args_str = " ".join([str(arg) for arg in concat_job.shell.nodes[0].args]) - assert str(pipeline.output_vcf) in args_str + # Check that the first argument (output file) is set to the intermediate + args_str = " ".join( + [str(arg) for arg in concat_job.shell.nodes[0].args] + ) + assert str(self.snv_apply_vcf(pipeline)) in args_str + assert str(pipeline.output_vcf) not in args_str + + # The AD-update job writes the final output and depends + # on the concat job + _, all_jobs = self._get_all_job_names(dag) + ad_update_job = next(j for j in all_jobs if j.name == "sad-lad-update") + update_cmd = str(ad_update_job.shell) + assert f"--input_vcf {self.snv_apply_vcf(pipeline)}" in update_cmd + assert f"--output_vcf {pipeline.output_vcf}" in update_cmd + assert dag.waiting_jobs[ad_update_job] == {concat_job} def test_gvcf_adds_gvcftyper_and_routes_outputs(self): - """--gvcf produces a .g.vcf.gz from model-apply and runs - GVCFtyper to produce the final VCF at output_vcf.""" + """--gvcf produces a .g.vcf.gz from model-apply, + updates it into the final .g.vcf.gz, then runs GVCFtyper on the + updated gVCF to produce the final VCF at output_vcf.""" pipeline = self.create_pipeline() pipeline.gvcf = True - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) job_names = [job.name for job in all_jobs] assert "model-apply" in job_names + assert "sad-lad-update" in job_names assert "gvcftyper" in job_names expected_gvcf = str(pipeline.output_vcf).replace( ".vcf.gz", ".g.vcf.gz" ) + apply_gvcf = str(self.snv_apply_vcf(pipeline)) + assert apply_gvcf.endswith("sample-snv_apply.g.vcf.gz") - # DNAscope emits gVCF and model-apply writes to .g.vcf.gz - for name in ("dnascope-raw", "model-apply"): + # DNAscope emits gVCF; model-apply writes gVCF + for name in ("dnascope", "model-apply"): job = next(j for j in all_jobs if j.name == name) cmd_str = str(job.shell) - if name == "dnascope-raw": + if name == "dnascope": assert "--emit_mode gvcf" in cmd_str else: - assert expected_gvcf in cmd_str - - # GVCFtyper reads the gVCF and writes the final output VCF + assert apply_gvcf in cmd_str + assert expected_gvcf not in cmd_str + + # The update job rewrites it to the final .g.vcf.gz + apply_job = next(j for j in all_jobs if j.name == "model-apply") + ad_update_job = next(j for j in all_jobs if j.name == "sad-lad-update") + update_cmd = str(ad_update_job.shell) + assert f"--input_vcf {apply_gvcf}" in update_cmd + assert f"--output_vcf {expected_gvcf}" in update_cmd + assert dag.waiting_jobs[ad_update_job] == {apply_job} + + # GVCFtyper reads the updated gVCF and writes the final VCF gvcftyper_job = next(j for j in all_jobs if j.name == "gvcftyper") cmd_str = str(gvcftyper_job.shell) assert "--algo GVCFtyper" in cmd_str assert expected_gvcf in cmd_str + assert apply_gvcf not in cmd_str assert str(pipeline.output_vcf) in cmd_str + assert dag.waiting_jobs[gvcftyper_job] == {ad_update_job} + + def snv_apply_vcf(self, pipeline): + """The small-variant model-apply intermediate for a pipeline""" + name = ( + "sample-snv_apply.g.vcf.gz" + if getattr(pipeline, "gvcf", False) + else "sample-snv_apply.vcf.gz" + ) + return pipeline.tmp_dir.joinpath(name) + + def test_ad_update_default(self): + """A single AD-update job reads the model-apply + intermediate and writes the final output VCF.""" + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + + job_names, all_jobs = self._get_all_job_names(dag) + assert job_names.count("sad-lad-update") == 1 + + ad_update_job = next(j for j in all_jobs if j.name == "sad-lad-update") + cmd_str = str(ad_update_job.shell) + assert "sad_lad_update.py" in cmd_str + assert f"--input_vcf {self.snv_apply_vcf(pipeline)}" in cmd_str + assert f"--output_vcf {pipeline.output_vcf}" in cmd_str + assert f"--threads {pipeline.cores}" in cmd_str + assert ad_update_job.task_name == "ad-update" + + # Model apply writes the intermediate, not the final output + apply_job = next(j for j in all_jobs if j.name == "model-apply") + assert str(self.snv_apply_vcf(pipeline)) in str(apply_job.shell) + + # The AD-update job depends on model-apply + assert dag.waiting_jobs[ad_update_job] == {apply_job} + + def test_no_ad_update_with_skip_small_variants(self): + """No AD-update job when small variants are skipped""" + pipeline = self.create_pipeline() + pipeline.skip_small_variants = True + dag = pipeline.build_dag() + + job_names, _ = self._get_all_job_names(dag) + assert "sad-lad-update" not in job_names def test_no_gvcftyper_without_gvcf(self): """Without --gvcf, no GVCFtyper job is added""" pipeline = self.create_pipeline() - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) job_names = [job.name for job in all_jobs] assert "gvcftyper" not in job_names + def _second_dag_job(self, pipeline, name, sample_sex=SampleSex.FEMALE): + """Build both DAGs and return the named job of the second. + + The first DAG has to be built first so that the short-read + alignments and the ploidy JSON are stashed for the second. + """ + pipeline.build_dag() + pipeline.sample_sex = sample_sex + dag = pipeline.build_second_dag() + _, all_jobs = self._get_all_job_names(dag) + return next(j for j in all_jobs if j.name == name) + def build_segdup_cmd(self, tech): """Build the segdup-caller command string for a given platform""" pipeline = self.create_pipeline() pipeline.tech = tech - pipeline.sample_sex = SampleSex.FEMALE - job = pipeline.build_segdup_job( - self.mock_dir / "output_segdups", - self.mock_bam, - self.mock_vcf, - None, - ) + pipeline.segdup_caller = [] + job = self._second_dag_job(pipeline, "segdup-caller") return str(job.shell) def test_segdup_ultima_lowers_min_map_qual(self): @@ -244,26 +347,198 @@ def test_segdup_no_override_for_short_reads(self): """Non-Ultima input leaves the segdup-caller defaults alone""" assert "--set" not in self.build_segdup_cmd("Illumina") + def test_segdup_command(self): + """The short-read segdup-caller command built by the second DAG""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = ["CFH", "CYP2D6", "SMN1"] + job = self._second_dag_job(pipeline, "segdup-caller") + + assert str(job.shell) == ( + f"segdup-caller --short {self.mock_bam} " + f"--reference {self.mock_ref} " + f"--sr_model {self.mock_bundle} " + f"--input_vcf {self.mock_vcf} " + "--sex female " + "--genes CFH,CYP2D6,SMN1 " + f"--outdir {self.mock_dir}/output_segdups" + ) + assert job.name == "segdup-caller" + assert job.task_name == "segdup" + assert job.threads == pipeline.cores + + def test_segdup_without_genes(self): + """An empty gene list runs segdup-caller's own default set""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = [] + job = self._second_dag_job(pipeline, "segdup-caller") + + assert "--genes" not in str(job.shell) + + def test_segdup_male_sample(self): + """The estimated sex reaches segdup-caller""" + pipeline = self.create_pipeline() + pipeline.segdup_caller = [] + job = self._second_dag_job( + pipeline, "segdup-caller", sample_sex=SampleSex.MALE + ) + + assert "--sex male" in str(job.shell) + + def test_no_segdup_job_by_default(self): + """Without `--segdup_caller` there is no segdup job, or DAG 2""" + pipeline = self.create_pipeline() + pipeline.build_dag() + pipeline.sample_sex = SampleSex.FEMALE + + assert pipeline.build_second_dag() is None + + def test_expansion_hunter_command(self): + """The ExpansionHunter command built by the second DAG""" + catalog = self.mock_dir / "catalog.json" + catalog.touch() + + pipeline = self.create_pipeline() + pipeline.expansion_catalog = catalog + job = self._second_dag_job(pipeline, "expansion-hunter") + + assert str(job.shell) == ( + f"ExpansionHunter --reads {self.mock_bam} " + f"--reference {self.mock_ref} " + f"--variant-catalog {catalog} " + "--sex female " + f"--threads {pipeline.cores} " + f"--output-prefix {self.mock_dir}/output_expansion" + ) + assert job.name == "expansion-hunter" + assert job.task_name == "expansion-hunter" + assert job.threads == pipeline.cores + + def enable_t1k(self, pipeline): + """Supply the T1K HLA and KIR reference files""" + for name in ("hla_seq", "hla_coord", "kir_seq", "kir_coord"): + (self.mock_dir / f"{name}.fa").touch() + pipeline.t1k_hla_seq = self.mock_dir / "hla_seq.fa" + pipeline.t1k_hla_coord = self.mock_dir / "hla_coord.fa" + pipeline.t1k_kir_seq = self.mock_dir / "kir_seq.fa" + pipeline.t1k_kir_coord = self.mock_dir / "kir_coord.fa" + return pipeline + + def _get_deps(self, dag, job): + """The dependencies of a job in a DAG""" + return dag.waiting_jobs.get(job, set()) + + def _get_dep_names(self, dag, job): + """The names of a job's dependencies""" + return {dep.name for dep in self._get_deps(dag, job)} + + def test_t1k_jobs_in_the_first_dag(self): + """Both T1K gene groups add an extract and a genotyping job""" + pipeline = self.enable_t1k(self.create_pipeline()) + dag = pipeline.build_dag() + + job_names, all_jobs = self._get_all_job_names(dag) + for name in ( + "t1k-hla-extract", + "t1k-hla", + "t1k-kir-extract", + "t1k-kir", + ): + assert name in job_names + + for tag in ("hla", "kir"): + extract = next( + j for j in all_jobs if j.name == f"t1k-{tag}-extract" + ) + t1k = next(j for j in all_jobs if j.name == f"t1k-{tag}") + assert self._get_deps(dag, t1k) == {extract} + # With BAM/CRAM input the extract reads the user's alignment + assert self._get_deps(dag, extract) == set() + assert extract.task_name == "t1k" + assert t1k.task_name == "t1k" + + def test_t1k_extract_command(self): + """The HLA extract runs ReadWriter over the T1K locus""" + pipeline = self.enable_t1k(self.create_pipeline()) + dag = pipeline.build_dag() + + _, all_jobs = self._get_all_job_names(dag) + extract = next(j for j in all_jobs if j.name == "t1k-hla-extract") + assert str(extract.shell) == ( + f"sentieon driver --input {self.mock_bam} " + f"--reference {self.mock_ref} --thread_count {pipeline.cores} " + "--interval chr6:28510020-33480577 " + f"--algo ReadWriter {self.mock_dir}/sample_hla.bam" + ) + + def test_t1k_command(self): + """run-t1k reads the extracted BAM and writes next to the VCF""" + pipeline = self.enable_t1k(self.create_pipeline()) + dag = pipeline.build_dag() + + _, all_jobs = self._get_all_job_names(dag) + t1k = next(j for j in all_jobs if j.name == "t1k-kir") + assert str(t1k.shell) == ( + f"run-t1k --abnormalUnmapFlag -t {pipeline.cores} " + "--preset kir-wgs " + f"-f {self.mock_dir}/kir_seq.fa " + f"-c {self.mock_dir}/kir_coord.fa " + f"--od {self.mock_dir}/output_kir " + f"-b {self.mock_dir}/sample_kir.bam" + ) + + def test_t1k_extract_depends_on_dedup_with_fastq_input(self): + """With fastq input the extract waits for the bwa alignment""" + pipeline = self.enable_t1k(self.create_fastq_pipeline()) + dag = pipeline.build_dag() + + _, all_jobs = self._get_all_job_names(dag) + extract = next(j for j in all_jobs if j.name == "t1k-hla-extract") + assert self._get_dep_names(dag, extract) == {"dedup-bwa"} + assert str(self.mock_dir / "output_bwa_deduped.cram") in str( + extract.shell + ) + + def test_no_t1k_jobs_by_default(self): + """T1K runs only when its reference files are supplied""" + pipeline = self.create_pipeline() + job_names, _ = self._get_all_job_names(pipeline.build_dag()) + + assert not [name for name in job_names if name.startswith("t1k")] + + def test_sr_alignments_stashed_with_bam_input(self): + """The second DAG reads the short-read alignments off the pipeline""" + pipeline = self.create_pipeline() + pipeline.build_dag() + + assert pipeline.sr_alignments == [self.mock_bam] + + def test_sr_alignments_stashed_with_fastq_input(self): + """With fastq input it is the deduplicated bwa alignment""" + pipeline = self.create_fastq_pipeline() + pipeline.build_dag() + + assert pipeline.sr_alignments == [ + self.mock_dir / "output_bwa_deduped.cram" + ] + def test_call_svs(self): """Test that PangenomeSV is added when --call_svs is enabled""" pipeline = self.create_pipeline() pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() assert isinstance(dag, DAG) - # The dnascope-raw job should exist - all_jobs = list(dag.waiting_jobs.keys()) + list( - dag.ready_jobs.keys() - ) + # The dnascope job should exist + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) job_names = [job.name for job in all_jobs] - assert "dnascope-raw" in job_names + assert "dnascope" in job_names - # Find the dnascope-raw job and check its command includes + # Find the dnascope job and check its command includes # both DNAscope and PangenomeSV dnascope_job = None for job in all_jobs: - if job.name == "dnascope-raw": + if job.name == "dnascope": dnascope_job = job break assert dnascope_job is not None @@ -273,22 +548,18 @@ def test_call_svs(self): assert "--gfa_file" in cmd_str # SV output should use _sv.vcf.gz suffix - sv_vcf = str(pipeline.output_vcf).replace( - ".vcf.gz", "_sv.vcf.gz" - ) + sv_vcf = str(pipeline.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") assert sv_vcf in cmd_str def test_call_svs_disabled_by_default(self): """Test that PangenomeSV is not added by default""" pipeline = self.create_pipeline() - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() - all_jobs = list(dag.waiting_jobs.keys()) + list( - dag.ready_jobs.keys() - ) + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) dnascope_job = None for job in all_jobs: - if job.name == "dnascope-raw": + if job.name == "dnascope": dnascope_job = job break assert dnascope_job is not None @@ -300,15 +571,13 @@ def test_skip_small_variants(self): """Test that DNAscope, transfer, and model-apply are skipped""" pipeline = self.create_pipeline() pipeline.skip_small_variants = True - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() assert isinstance(dag, DAG) - all_jobs = list(dag.waiting_jobs.keys()) + list( - dag.ready_jobs.keys() - ) + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) job_names = [job.name for job in all_jobs] - assert "dnascope-raw" not in job_names + assert "dnascope" not in job_names assert "model-apply" not in job_names assert "merge-trim-concat" not in job_names @@ -317,17 +586,15 @@ def test_skip_small_variants_with_call_svs(self): pipeline = self.create_pipeline() pipeline.skip_small_variants = True pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() assert isinstance(dag, DAG) - all_jobs = list(dag.waiting_jobs.keys()) + list( - dag.ready_jobs.keys() - ) + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) job_names = [job.name for job in all_jobs] # The driver job should still run for SV calling - assert "dnascope-raw" in job_names + assert "dnascope" in job_names # Transfer and model-apply should be skipped assert "model-apply" not in job_names @@ -336,7 +603,7 @@ def test_skip_small_variants_with_call_svs(self): # The driver command should have PangenomeSV but NOT DNAscope dnascope_job = None for job in all_jobs: - if job.name == "dnascope-raw": + if job.name == "dnascope": dnascope_job = job break assert dnascope_job is not None @@ -344,20 +611,158 @@ def test_skip_small_variants_with_call_svs(self): assert "--algo PangenomeSV" in cmd_str assert "--algo DNAscope" not in cmd_str + def test_pangenome_check_jobs(self): + """The graph checks sit between vg and everything downstream. + + A wrong `--pangenome_ref_name` is silent in `vg`: the sampled + graph loses its reference paths and the GFA loses its rGFA tags, + and `pgutil lift` then leaves every read unmapped. + """ + pipeline = self.create_pipeline() + dag = pipeline.build_dag() + job_names, all_jobs = self._get_all_job_names(dag) + + assert "check-sample-gbz" in job_names + assert "check-sample-gfa" in job_names + + check_gbz = next(j for j in all_jobs if j.name == "check-sample-gbz") + check_gfa = next(j for j in all_jobs if j.name == "check-sample-gfa") + haplotypes = next(j for j in all_jobs if j.name == "vg-haplotypes") + gfa = next(j for j in all_jobs if j.name == "vg-convert-gfa") + fasta = next(j for j in all_jobs if j.name == "vg-paths-fasta") + mm2 = next(j for j in all_jobs if j.name == "mm2-lift") + + assert self._get_deps(dag, check_gbz) == {haplotypes} + assert self._get_deps(dag, gfa) == {check_gbz} + assert self._get_deps(dag, fasta) == {check_gbz} + assert self._get_deps(dag, check_gfa) == {gfa} + assert check_gfa in self._get_deps(dag, mm2) + assert "vg-convert-gfa" not in self._get_dep_names(dag, mm2) + + # The checks share the log group of the jobs they guard + assert check_gbz.task_name == haplotypes.task_name + assert check_gfa.task_name == gfa.task_name + + def test_pangenome_check_commands(self): + """The checks run the resolved name and prefix""" + pipeline = self.create_pipeline() + pipeline.pangenome_ref_name = "CHM13" + pipeline.pangenome_contig_prefix = "CHM13#0#" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + gbz_cmd = str( + next(j for j in all_jobs if j.name == "check-sample-gbz").shell + ) + assert "sentieon_cli.pangenome_meta check-gbz" in gbz_cmd + # runpy's duplicate-module warning is kept out of the job logs + assert "-W ignore::RuntimeWarning:runpy" in gbz_cmd + assert f"--gbz {self.mock_dir}/sample_pangenome.gbz" in gbz_cmd + assert "--reference_name CHM13" in gbz_cmd + # The `#` characters are shell-quoted in the rendered command + assert "--contig_prefix 'CHM13#0#'" in gbz_cmd + + gfa_cmd = str( + next(j for j in all_jobs if j.name == "check-sample-gfa").shell + ) + assert "sentieon_cli.pangenome_meta check-gfa" in gfa_cmd + assert f"--gfa {self.mock_dir}/sample-hap.gfa" in gfa_cmd + assert "--reference_name CHM13" in gfa_cmd + assert "--contig_prefix 'CHM13#0#'" in gfa_cmd + + def test_the_resolved_reference_name_reaches_vg(self): + """`vg haplotypes` and `vg convert` use the resolved name""" + pipeline = self.create_pipeline() + pipeline.pangenome_ref_name = "CHM13" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + haplotypes = str( + next(j for j in all_jobs if j.name == "vg-haplotypes").shell + ) + assert "--set-reference CHM13" in haplotypes + convert = str( + next(j for j in all_jobs if j.name == "vg-convert-gfa").shell + ) + assert "-Q CHM13" in convert + def _get_all_job_names(self, dag): """Helper to get all job names from a DAG""" - all_jobs = list(dag.waiting_jobs.keys()) + list( - dag.ready_jobs.keys() - ) + all_jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) return [job.name for job in all_jobs], all_jobs + def _build_cnv_dag(self, pipeline, sample_sex=SampleSex.FEMALE): + """Build the second DAG, which holds the sex-aware CNV jobs. + + The first DAG has to be built first so that the CNV input + alignments are stashed for the second DAG. + """ + pipeline.build_dag() + pipeline.sample_sex = sample_sex + return pipeline.build_second_dag() + + def validate_cnv(self, pipeline): + """Run the CNV validation with the reference index parsed""" + pipeline.validate_cnv() + + @pytest.mark.parametrize( + "sample_sex", [None, SampleSex.MALE, SampleSex.FEMALE] + ) + def test_validate_cnv_requires_a_par_bed(self, sample_sex): + # The mock reference is not a recognized build, so no packaged + # PAR BED file can be selected and validation stops the run + pipeline = self.create_pipeline() + pipeline.call_svs = True + pipeline.sample_sex = sample_sex + + with pytest.raises(SystemExit) as excinfo: + self.validate_cnv(pipeline) + assert excinfo.value.code == 2 + + def test_validate_cnv_accepts_the_par_bed_argument(self): + par_bed = self.mock_dir / "par.bed" + par_bed.write_text("chrX\t10000\t2781479\n") + + pipeline = self.create_pipeline() + pipeline.call_svs = True + pipeline.par_bed = par_bed + + self.validate_cnv(pipeline) # no SystemExit + + assert pipeline.cnv_par_bed == par_bed + + def test_validate_cnv_needs_no_par_bed_without_call_svs(self): + pipeline = self.create_pipeline() + + self.validate_cnv(pipeline) # no SystemExit + + assert pipeline.cnv_par_bed is None + + def test_validate_cnv_needs_no_par_bed_without_a_cnv_model(self): + pipeline = self.create_pipeline() + pipeline.call_svs = True + pipeline.has_cnv_model = False + + self.validate_cnv(pipeline) # no SystemExit + + assert pipeline.cnv_par_bed is None + def test_cnv_jobs_with_call_svs(self): - """Test that CNV jobs are added when --call_svs is enabled""" + """CNV jobs run in the second DAG when --call_svs is enabled""" pipeline = self.create_pipeline() pipeline.call_svs = True - dag = pipeline.build_first_dag() + first_dag = pipeline.build_dag() - job_names, all_jobs = self._get_all_job_names(dag) + # CNV calling is sex-aware, so it is not in the first DAG + first_job_names, _ = self._get_all_job_names(first_dag) + assert "cnvscope" not in first_job_names + assert "cnv-model-apply" not in first_job_names + assert "indel2cnv" not in first_job_names + assert "combine-cnv" not in first_job_names + + pipeline.sample_sex = SampleSex.FEMALE + dag = pipeline.build_second_dag() + job_names, _ = self._get_all_job_names(dag) assert "cnvscope" in job_names assert "cnv-model-apply" in job_names assert "indel2cnv" in job_names @@ -366,7 +771,7 @@ def test_cnv_jobs_with_call_svs(self): def test_cnv_jobs_not_present_without_call_svs(self): """Test that CNV jobs are not added by default""" pipeline = self.create_pipeline() - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() job_names, _ = self._get_all_job_names(dag) assert "cnvscope" not in job_names @@ -374,12 +779,62 @@ def test_cnv_jobs_not_present_without_call_svs(self): assert "indel2cnv" not in job_names assert "combine-cnv" not in job_names + # And no second DAG at all, as nothing consumes the sample sex + pipeline.sample_sex = SampleSex.FEMALE + assert pipeline.build_second_dag() is None + + def test_second_dag_needed_for_cnv_only(self): + """CNV calling alone triggers the second DAG""" + pipeline = self.create_pipeline() + pipeline.call_svs = True + + assert pipeline._cnv_in_second_dag() is True + assert pipeline._needs_second_dag() is True + + def test_no_second_dag_without_a_sex_aware_caller(self): + """No second DAG when nothing consumes the sample sex""" + pipeline = self.create_pipeline() + + assert pipeline._cnv_in_second_dag() is False + assert pipeline._needs_second_dag() is False + + def test_ploidy_estimation_runs_with_sample_sex(self): + """--sample_sex overrides the estimate, which still runs""" + pipeline = self.create_pipeline() + pipeline.call_svs = True + pipeline.sample_sex = SampleSex.FEMALE + + job_names, _ = self._get_all_job_names(pipeline.build_dag()) + assert "estimate-ploidy" in job_names + + def test_ploidy_estimation_runs_without_sex_aware_callers(self): + """The ploidy JSON is written by every run""" + pipeline = self.create_pipeline() + + assert pipeline._needs_second_dag() is False + job_names, all_jobs = self._get_all_job_names(pipeline.build_dag()) + assert "estimate-ploidy" in job_names + + ploidy_job = next(j for j in all_jobs if j.name == "estimate-ploidy") + cmd_str = str(ploidy_job.shell) + assert "estimate_ploidy.py" in cmd_str + assert str(pipeline.ploidy_json) in cmd_str + assert str(pipeline.ploidy_json).endswith("output_ploidy.json") + + def test_ploidy_estimation_runs_with_skip_small_variants(self): + """The ploidy job precedes the SV-only early return""" + pipeline = self.create_pipeline() + pipeline.skip_small_variants = True + + job_names, _ = self._get_all_job_names(pipeline.build_dag()) + assert "estimate-ploidy" in job_names + def test_cnv_jobs_with_skip_small_variants(self): """Test CNV jobs are added in SV-only mode""" pipeline = self.create_pipeline() pipeline.skip_small_variants = True pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = self._build_cnv_dag(pipeline) job_names, _ = self._get_all_job_names(dag) assert "cnvscope" in job_names @@ -392,13 +847,13 @@ def test_call_svs_without_cnv_model(self): pipeline = self.create_pipeline() pipeline.call_svs = True pipeline.has_cnv_model = False - dag = pipeline.build_first_dag() + dag = pipeline.build_dag() job_names, all_jobs = self._get_all_job_names(dag) # PangenomeSV still runs - assert "dnascope-raw" in job_names - dnascope_job = next(j for j in all_jobs if j.name == "dnascope-raw") + assert "dnascope" in job_names + dnascope_job = next(j for j in all_jobs if j.name == "dnascope") cmd_str = str(dnascope_job.shell) assert "--algo PangenomeSV" in cmd_str @@ -408,26 +863,24 @@ def test_call_svs_without_cnv_model(self): assert "indel2cnv" not in job_names assert "combine-cnv" not in job_names + # And there is no second DAG either + pipeline.sample_sex = SampleSex.FEMALE + assert pipeline.build_second_dag() is None + def test_call_svs_without_cnv_model_skip_small_variants(self): """SV-only mode: SVs run but CNV jobs are skipped without cnv.model""" pipeline = self.create_pipeline() pipeline.skip_small_variants = True pipeline.call_svs = True pipeline.has_cnv_model = False - dag = pipeline.build_first_dag() - job_names, _ = self._get_all_job_names(dag) - assert "dnascope-raw" in job_names - assert "cnvscope" not in job_names - assert "cnv-model-apply" not in job_names - assert "indel2cnv" not in job_names - assert "combine-cnv" not in job_names + assert self._build_cnv_dag(pipeline) is None def test_cnvscope_command(self): """Test CNVscope driver command has correct algo and model""" pipeline = self.create_pipeline() pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = self._build_cnv_dag(pipeline) _, all_jobs = self._get_all_job_names(dag) cnvscope_job = None @@ -439,14 +892,36 @@ def test_cnvscope_command(self): cmd_str = str(cnvscope_job.shell) assert "--algo CNVscope" in cmd_str assert "cnv.model" in cmd_str + # Sex-aware calling of a female sample needs no PAR BED file + assert "--sex F" in cmd_str + assert "--par" not in cmd_str + # Per-probe debugging output lands next to the CNV VCF + probes = str(self.mock_vcf).replace(".vcf.gz", "_cnv.probes") + assert f"--dump_probes {probes}" in cmd_str # Input should be the sample BAM (BAM input mode) assert str(pipeline.sample_input[0]) in cmd_str + def test_cnvscope_command_male_sample(self): + """A male sample is called with the PAR BED file""" + par_bed = self.mock_dir / "par.bed" + par_bed.touch() + + pipeline = self.create_pipeline() + pipeline.call_svs = True + pipeline.cnv_par_bed = par_bed + dag = self._build_cnv_dag(pipeline, sample_sex=SampleSex.MALE) + + _, all_jobs = self._get_all_job_names(dag) + cnvscope_job = next(j for j in all_jobs if j.name == "cnvscope") + cmd_str = str(cnvscope_job.shell) + assert "--sex M" in cmd_str + assert f"--par {par_bed}" in cmd_str + def test_cnv_model_apply_command(self): """Test CNVModelApply driver command""" pipeline = self.create_pipeline() pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = self._build_cnv_dag(pipeline) _, all_jobs = self._get_all_job_names(dag) job = None @@ -463,7 +938,7 @@ def test_indel2cnv_command(self): """Test indel2cnv script command""" pipeline = self.create_pipeline() pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = self._build_cnv_dag(pipeline) _, all_jobs = self._get_all_job_names(dag) job = None @@ -480,7 +955,7 @@ def test_combine_cnv_command(self): """Test combine_cnv script command""" pipeline = self.create_pipeline() pipeline.call_svs = True - dag = pipeline.build_first_dag() + dag = self._build_cnv_dag(pipeline) _, all_jobs = self._get_all_job_names(dag) job = None @@ -494,9 +969,7 @@ def test_combine_cnv_command(self): assert "--cnv" in cmd_str assert "--converted" in cmd_str # Output should use _cnv.vcf.gz suffix - cnv_vcf = str(pipeline.output_vcf).replace( - ".vcf.gz", "_cnv.vcf.gz" - ) + cnv_vcf = str(pipeline.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") assert cnv_vcf in cmd_str def test_cnv_with_skip_model_apply(self): @@ -504,12 +977,14 @@ def test_cnv_with_skip_model_apply(self): pipeline = self.create_pipeline() pipeline.call_svs = True pipeline.skip_model_apply = True - dag = pipeline.build_first_dag() + first_job_names, _ = self._get_all_job_names(pipeline.build_dag()) + assert "model-apply" not in first_job_names + + pipeline.sample_sex = SampleSex.FEMALE + dag = pipeline.build_second_dag() job_names, _ = self._get_all_job_names(dag) - assert "model-apply" not in job_names assert "cnvscope" in job_names assert "cnv-model-apply" in job_names assert "indel2cnv" in job_names assert "combine-cnv" in job_names - diff --git a/tests/unit/test_sentieon_pangenome_kmc_integration.py b/tests/unit/test_sentieon_pangenome_kmc_integration.py index cfe4e61..b4f5fe7 100644 --- a/tests/unit/test_sentieon_pangenome_kmc_integration.py +++ b/tests/unit/test_sentieon_pangenome_kmc_integration.py @@ -21,10 +21,14 @@ def test_validate_kmc_patch_success(self): pipeline.reference = "ref" pipeline.pop_vcf = "pop" pipeline.gbz = "ref.grch38.gbz" - pipeline.fai_data = {} - pipeline.pop_vcf_contigs = {} - + pipeline.hapl = "ref.grch38.hapl" + # The mock graph and pop VCF files cannot be parsed; a dry run + # falls back to the default pangenome reference and to an + # empty set of pop VCF contigs + pipeline.dry_run = True + # Mock dependencies + pipeline.logger = MagicMock() pipeline.handle_arguments = MagicMock() pipeline.setup_logging = MagicMock() pipeline.validate_bundle = MagicMock() @@ -33,7 +37,9 @@ def test_validate_kmc_patch_success(self): pipeline.validate_ref = MagicMock() pipeline.collect_readgroups = MagicMock() - with patch("sentieon_cli.sentieon_pangenome.check_version", return_value=True), \ + with patch("sentieon_cli.util.check_version", return_value=True), \ + patch("sentieon_cli.sentieon_pangenome.parse_fai", return_value={}), \ + patch("sentieon_cli.sentieon_pangenome.determine_shards_from_fai", return_value=[]), \ patch("sentieon_cli.sentieon_pangenome.check_kmc_patch", return_value=True) as mock_check, \ patch("sys.exit") as mock_exit: @@ -54,12 +60,15 @@ def test_validate_kmc_patch_failure(self): pipeline.validate_ref = MagicMock() pipeline.collect_readgroups = MagicMock() pipeline.bed = "bed" - pipeline.fai_data = {} - pipeline.pop_vcf_contigs = {} + pipeline.gbz = "ref.grch38.gbz" + pipeline.hapl = "ref.grch38.hapl" + pipeline.dry_run = True pipeline.logger = MagicMock() - with patch("sentieon_cli.sentieon_pangenome.check_version", return_value=True), \ + with patch("sentieon_cli.util.check_version", return_value=True), \ + patch("sentieon_cli.sentieon_pangenome.parse_fai", return_value={}), \ + patch("sentieon_cli.sentieon_pangenome.determine_shards_from_fai", return_value=[]), \ patch("sentieon_cli.sentieon_pangenome.check_kmc_patch", return_value=False) as mock_check, \ patch("sys.exit") as mock_exit: diff --git a/tests/unit/test_shell_pipeline.py b/tests/unit/test_shell_pipeline.py index 943d5b9..a6d704c 100644 --- a/tests/unit/test_shell_pipeline.py +++ b/tests/unit/test_shell_pipeline.py @@ -3,8 +3,10 @@ """ import asyncio +import fcntl import os import pathlib +import signal import sys import tempfile @@ -31,10 +33,9 @@ async def test_simple_command(): cmd = Command("echo", "hello world") context = Context() - # Use a file for stdout with tempfile.NamedTemporaryFile(mode="w+", delete=False) as stdout_file: proc = await cmd.run(context, stdout=stdout_file) - await proc.wait() + await proc.async_wait() stdout_file.seek(0) output = stdout_file.read().strip() @@ -44,6 +45,29 @@ async def test_simple_command(): assert output == "hello world" assert proc.returncode == 0 + # wait4 captured the child's resource usage as it was reaped + assert proc.rusage is not None + assert proc.rusage.ru_maxrss > 0 + + +@pytest.mark.asyncio +async def test_a_signalled_command_reports_its_signal_and_its_rusage(): + """A child killed by a signal is still reaped through ``wait4``. + + This pins ``RusagePopen._try_wait``: if a future CPython changes that + private method's shape, the override stops being called and ``rusage`` + silently stays None while the negative return code still works. + """ + cmd = Command("sleep", "5") + context = Context() + + proc = await cmd.run(context) + proc.send_signal(signal.SIGTERM) + + assert await proc.async_wait() == -signal.SIGTERM + assert proc.rusage is not None + + await context.cleanup() @pytest.mark.asyncio @@ -56,7 +80,7 @@ async def test_simple_pipeline(): with tempfile.NamedTemporaryFile(mode="w+", delete=False) as stdout_file: proc = await pipeline.run(context, stdout=stdout_file) - await proc.wait() + await proc.async_wait() stdout_file.seek(0) output = stdout_file.read().strip() @@ -73,7 +97,6 @@ async def test_pipeline_with_file_io(): """Test a pipeline with file input and output""" context = Context() - # Create input and output files with tempfile.NamedTemporaryFile(mode="w", delete=False) as infile: infile.write("hello from file") infile_path = pathlib.Path(infile.name) @@ -82,12 +105,16 @@ async def test_pipeline_with_file_io(): cmd1 = Command("cat") cmd2 = Command("wc", "-c") - pipeline = Pipeline(cmd1, cmd2, file_input=infile_path, file_output=outfile_path) + pipeline = Pipeline( + cmd1, + cmd2, + file_input=infile_path, + file_output=outfile_path, + ) proc = await pipeline.run(context) - await proc.wait() + await proc.async_wait() - # Verify output file content with open(outfile_path, "r") as f: output = f.read().strip() @@ -111,7 +138,7 @@ async def test_input_process_substitution(): with tempfile.NamedTemporaryFile(mode="w+", delete=False) as stdout_file: proc = await cmd.run(context, stdout=stdout_file) - await proc.wait() + await proc.async_wait() stdout_file.seek(0) output = stdout_file.read().strip() @@ -132,25 +159,21 @@ async def test_output_process_substitution(): """Test output process substitution: tee >(wc -c)""" context = Context() - # We will write the output of wc -c to a file to check it outfile_path = pathlib.Path(context.temp_dir.name) / "wc_output.txt" - # The pipeline will be: echo "hello" | tee >(wc -c > outfile.txt) + # echo "hello" | tee >(wc -c > outfile.txt) # The main output of the pipeline should be "hello" - # The output of the process substitution should be the character count "6" - - # Inner pipeline for process substitution + # The output of the process substitution should be the char count. wc_cmd = Command("wc", "-c") wc_pipeline = Pipeline(wc_cmd, file_output=outfile_path) - # Main pipeline echo_cmd = Command("echo", "hello") tee_cmd = Command("tee", OutputProcSub(wc_pipeline)) main_pipeline = Pipeline(echo_cmd, tee_cmd) with tempfile.NamedTemporaryFile(mode="w+", delete=False) as stdout_file: proc = await main_pipeline.run(context, stdout=stdout_file) - await proc.wait() + await proc.async_wait() # Wait for background tasks from process substitution to finish await asyncio.gather(*context.tasks) @@ -158,10 +181,8 @@ async def test_output_process_substitution(): stdout_file.seek(0) main_output = stdout_file.read().strip() - # Verify main output assert main_output == "hello" - # Verify process substitution output with open(outfile_path, "r") as f: wc_output = f.read().strip() @@ -171,3 +192,257 @@ async def test_output_process_substitution(): await context.cleanup() os.remove(stdout_file.name) assert proc.returncode == 0 + + +@pytest.mark.asyncio +async def test_pipe_size_enlarges_internal_pipes(tmp_path): + """pipe_size raises F_SETPIPE_SZ on the pipeline's internal pipes, + which the downstream child inherits.""" + if not hasattr(fcntl, "F_SETPIPE_SZ"): + pytest.skip("F_SETPIPE_SZ not available (non-Linux)") + try: + pipe_max = int(open("/proc/sys/fs/pipe-max-size").read()) + except OSError: + pytest.skip("cannot read fs.pipe-max-size") + target = 128 * 1024 # a power of two, so the kernel sets it exactly + if pipe_max < target: + pytest.skip("fs.pipe-max-size too small for this test") + + report = tmp_path / "size.txt" + # The first stage's stdout is the internal pipe to `cat`; report that + # pipe's buffer size (1032 == F_GETPIPE_SZ) rather than writing to it. + reporter = ( + "import fcntl, sys; " + "open(sys.argv[1], 'w').write(str(fcntl.fcntl(1, 1032)))" + ) + context = Context() + pipeline = Pipeline( + Command(sys.executable, "-c", reporter, str(report)), + Command("cat"), + pipe_size=target, + ) + proc = await pipeline.run(context) + await proc.async_wait() + for sub in context.commands: + if sub.proc: + await sub.proc.async_wait() + await context.cleanup() + + assert int(report.read_text()) == target + + +def test_command_and_pipeline_hash_identity(): + """Equal commands/pipelines hash equally (DAG relies on this).""" + c1 = Command("echo", "x") + c2 = Command("echo", "x") + c3 = Command("echo", "y") + assert c1 == c2 + assert hash(c1) == hash(c2) + assert c1 != c3 + + p1 = Pipeline(Command("echo", "x"), Command("cat")) + p2 = Pipeline(Command("echo", "x"), Command("cat")) + p3 = Pipeline(Command("echo", "x")) + assert p1 == p2 + assert hash(p1) == hash(p2) + assert p1 != p3 + + +def test_procsub_input_and_output_are_distinct(): + """<(x) and >(x) wrap the same pipeline but are not the same node: they + must be unequal and hash differently.""" + p = Pipeline(Command("echo", "x")) + assert InputProcSub(p) != OutputProcSub(p) + assert hash(InputProcSub(p)) != hash(OutputProcSub(p)) + # Each still equals another of its own kind wrapping an equal pipeline. + other = Pipeline(Command("echo", "x")) + assert InputProcSub(p) == InputProcSub(other) + assert hash(InputProcSub(p)) == hash(InputProcSub(other)) + + +def test_command_identity_is_stable_across_running(): + """A Command's identity must not change once it has started: __eq__ + used to mix in self.proc, so running a job mutated its identity.""" + c1 = Command("echo", "x") + c2 = Command("echo", "x") + assert c1 == c2 and hash(c1) == hash(c2) + c1.proc = object() # stand in for a started asyncio process + assert c1 == c2 + assert hash(c1) == hash(c2) + + +def test_command_identity_includes_exec_kwargs(): + """exec_kwargs participates in identity: commands differing only in + env/cwd are distinct, and equal env hashes the same regardless of the + dict's insertion order.""" + assert Command("run", exec_kwargs={"cwd": "/a"}) != Command( + "run", exec_kwargs={"cwd": "/b"} + ) + assert Command("run", exec_kwargs={"cwd": "/a"}) != Command("run") + e1 = Command("run", exec_kwargs={"env": {"A": "1", "B": "2"}}) + e2 = Command("run", exec_kwargs={"env": {"B": "2", "A": "1"}}) + assert e1 == e2 + assert hash(e1) == hash(e2) + + +def test_str_renders_env_and_cwd(): + """str() renders exec_kwargs env/cwd as a conventional bash prefix, + quoting values that need it.""" + env_cmd = Command("samtools", "sort", exec_kwargs={"env": {"REF": "hg"}}) + assert str(env_cmd) == "REF=hg samtools sort" + cwd_cmd = Command("samtools", "sort", exec_kwargs={"cwd": "/data"}) + assert str(cwd_cmd) == "(cd /data && samtools sort)" + both = Command( + "run", exec_kwargs={"cwd": "/data dir", "env": {"K": "v w"}} + ) + assert str(both) == "(cd '/data dir' && K='v w' run)" + + +def test_empty_pipeline_raises(): + with pytest.raises(ValueError): + Pipeline() + + +@pytest.mark.asyncio +async def test_file_output_and_stdout_conflict_raises(): + """A file/stdout conflict raises ValueError (not sys.exit).""" + context = Context() + out = pathlib.Path(context.temp_dir.name) / "out.txt" + pipeline = Pipeline(Command("echo", "x"), file_output=out) + with pytest.raises(ValueError): + await pipeline.run(context, stdout=1) + await context.cleanup() + + +@pytest.mark.asyncio +async def test_file_input_and_stdin_fd0_conflict_raises(tmp_path): + """file_input plus a stdin of fd 0 conflicts: fd 0 is a valid stream but + falsy, so the guard must test `is not None`, not truthiness.""" + inp = tmp_path / "in.txt" + inp.write_text("data") + context = Context() + pipeline = Pipeline(Command("cat"), file_input=inp) + with pytest.raises(ValueError): + await pipeline.run(context, stdin=0) + await context.cleanup() + + +@pytest.mark.asyncio +async def test_file_input_and_stdin_conflict_raises(): + """A file/stdin conflict raises ValueError (not sys.exit).""" + context = Context() + infile = pathlib.Path(context.temp_dir.name) / "in.txt" + pipeline = Pipeline(Command("cat"), file_input=infile) + with pytest.raises(ValueError): + await pipeline.run(context, stdin=1) + await context.cleanup() + + +@pytest.mark.asyncio +async def test_cleanup_unblocks_unopened_input_procsub(): + """cleanup() must not hang when an <(...) FIFO was never opened.""" + context = Context() + proc_sub = InputProcSub(Pipeline(Command("echo", "x"))) + fifo_path = await proc_sub.setup(context) + assert os.path.exists(fifo_path) + + await asyncio.wait_for(context.cleanup(), timeout=10) + + # The inner command never spawned: nothing opened the FIFO. + assert context.commands == [] + + +@pytest.mark.asyncio +async def test_cleanup_unblocks_unopened_output_procsub(): + """cleanup() must not hang when a >(...) FIFO was never opened.""" + context = Context() + proc_sub = OutputProcSub(Pipeline(Command("cat"))) + await proc_sub.setup(context) + + await asyncio.wait_for(context.cleanup(), timeout=10) + + assert context.commands == [] + + +@pytest.mark.asyncio +async def test_cleanup_after_outer_exits_without_opening(): + """An outer command that exits without opening its proc-sub FIFO + (here: false ignores its arguments) must not hang cleanup().""" + context = Context() + cmd = Command("false", InputProcSub(Pipeline(Command("echo", "x")))) + proc = await cmd.run(context) + await proc.async_wait() + + await asyncio.wait_for(context.cleanup(), timeout=10) + + assert [c.executable for c in context.commands] == ["false"] + + +@pytest.mark.asyncio +async def test_cleanup_unblocks_multiple_procsubs(): + """Several never-opened proc subs are all unblocked by cleanup().""" + context = Context() + cmd = Command( + "false", + InputProcSub(Pipeline(Command("echo", "x"))), + OutputProcSub(Pipeline(Command("cat"))), + ) + proc = await cmd.run(context) + await proc.async_wait() + + await asyncio.wait_for(context.cleanup(), timeout=10) + + assert all(task.done() for task in context.tasks) + + +@pytest.mark.asyncio +async def test_cleanup_still_raises_inner_launch_failure(): + """An inner command that cannot be spawned still propagates from + cleanup(), but only after all resources have been released.""" + context = Context() + cmd = Command( + "cat", InputProcSub(Pipeline(Command("no_such_cmd_zzz"))) + ) + proc = await cmd.run(context) + await proc.async_wait() + + with pytest.raises(FileNotFoundError): + await asyncio.wait_for(context.cleanup(), timeout=10) + + assert not os.path.exists(context.temp_dir.name) + + +@pytest.mark.asyncio +async def test_cleanup_does_not_stall_running_inner_writer(): + """cleanup() must not touch the FIFO of a proc sub whose inner command + is already running; holding a read end there would prevent the EPIPE + that lets the inner writer exit.""" + context = Context() + cmd = Command( + "sh", + "-c", + 'sleep 0.5 <"$1"', + "sh", + InputProcSub(Pipeline(Command("yes"))), + ) + proc = await cmd.run(context) + # Wait until the inner command has spawned (sh + yes). + async with asyncio.timeout(5): + while len(context.commands) < 2: + await asyncio.sleep(0.01) + + await asyncio.wait_for(context.cleanup(), timeout=10) + assert await proc.async_wait() == 0 + + +@pytest.mark.asyncio +async def test_cleanup_twice_is_safe(): + """A second cleanup() call is a harmless no-op.""" + context = Context() + cmd = Command("cat", InputProcSub(Pipeline(Command("echo", "hello")))) + proc = await cmd.run(context, stdout=asyncio.subprocess.DEVNULL) + await proc.async_wait() + + await asyncio.wait_for(context.cleanup(), timeout=10) + await asyncio.wait_for(context.cleanup(), timeout=10) + assert proc.returncode == 0 diff --git a/tests/unit/test_stage_alignment.py b/tests/unit/test_stage_alignment.py new file mode 100644 index 0000000..6e4829e --- /dev/null +++ b/tests/unit/test_stage_alignment.py @@ -0,0 +1,473 @@ +""" +Unit tests for the alignment and read-extraction stages +""" + +import logging +import pathlib +from unittest.mock import patch + +import pytest + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.alignment import ( + BwaExtractStage, + BwaFastqStage, + BwaRealignStage, + Minimap2FastqStage, + Minimap2RealignStage, + ReadWriterStage, + aln_suffix, + find_unzip, +) +from sentieon_cli.stages.base import StageContext, rm_job + +RG_LINES = ["@RG\tID:rg1\tSM:sample"] + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +@pytest.fixture +def rg_lines(): + """Stub out the `samtools view -H` call that reads the @RG lines""" + with patch( + "sentieon_cli.command_strings.get_rg_lines", return_value=RG_LINES + ) as mock: + yield mock + + +class TestHelpers: + """The shared helpers the stages and pipelines use""" + + def test_aln_suffix(self): + assert aln_suffix(True) == "bam" + assert aln_suffix(False) == "cram" + + def test_find_unzip_prefers_igzip(self): + with patch( + "sentieon_cli.stages.alignment.shutil.which", + return_value="/usr/bin/igzip", + ): + assert find_unzip(logging.getLogger("test-unzip")) == "igzip" + + def test_find_unzip_falls_back_to_gzip(self, caplog): + logger = logging.getLogger("test-unzip-fallback") + with patch( + "sentieon_cli.stages.alignment.shutil.which", return_value=None + ): + with caplog.at_level(logging.WARNING, logger=logger.name): + assert find_unzip(logger) == "gzip" + assert "igzip is recommended" in caplog.text + assert caplog.records[0].levelno == logging.WARNING + + +class TestBwaRealignStage: + """Re-aligning BAM/CRAM input with bwa""" + + def make(self, tmp_path, **kwargs) -> BwaRealignStage: + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "in0.bam", tmp_path / "in1.bam"], + model_bundle=tmp_path / "model.bundle", + ) + defaults.update(kwargs) + return BwaRealignStage(**defaults) # type: ignore[arg-type] + + def test_one_job_per_input(self, tmp_path, rg_lines): + result = self.make(tmp_path).build() + + assert [job.name for job in result.jobs] == [ + "bam-align-0", + "bam-align-1", + ] + for job in result.jobs: + assert job.threads == 4 + assert job.task_name == "alignment" + assert result.terminal == set(result.jobs) + + def test_header_files_are_written(self, tmp_path, rg_lines): + self.make(tmp_path).build() + + for i in range(2): + header = tmp_path / f"input_{i}.hdr" + assert header.read_text() == RG_LINES[0] + "\n" + assert rg_lines.call_count == 2 + + def test_duplicate_marking_writes_bam_to_the_tmp_dir( + self, tmp_path, rg_lines + ): + result = self.make(tmp_path).build() + + assert result.outputs == [ + tmp_path / "bwa_sorted_0.bam", + tmp_path / "bwa_sorted_1.bam", + ] + assert "--bam_compression 1" in str(result.jobs[0].shell) + + def test_without_duplicate_marking_writes_cram_output( + self, tmp_path, rg_lines + ): + result = self.make(tmp_path, duplicate_marking="none").build() + + assert result.outputs == [ + tmp_path / "output_bwa_sorted_0.cram", + tmp_path / "output_bwa_sorted_1.cram", + ] + # CRAM is written with the `util sort` defaults + shell = str(result.jobs[0].shell) + assert "--bam_compression 1" not in shell + assert "--cram_write_options" not in shell + + def test_util_sort_args_reach_the_command(self, tmp_path, rg_lines): + result = self.make( + tmp_path, + duplicate_marking="none", + util_sort_args="--cram_write_options version=3.1", + ).build() + + shell = str(result.jobs[0].shell) + assert "--sam2bam --cram_write_options version=3.1" in shell + + def test_bam_format_without_duplicate_marking(self, tmp_path, rg_lines): + result = self.make( + tmp_path, duplicate_marking="none", bam_format=True + ).build() + + assert result.outputs == [ + tmp_path / "output_bwa_sorted_0.bam", + tmp_path / "output_bwa_sorted_1.bam", + ] + + def test_intermediate_cleanup_paths(self, tmp_path, rg_lines): + result = self.make(tmp_path).build() + + assert result.cleanup_paths == [ + tmp_path / "bwa_sorted_0.bam", + pathlib.Path(str(tmp_path / "bwa_sorted_0.bam") + ".bai"), + tmp_path / "bwa_sorted_1.bam", + pathlib.Path(str(tmp_path / "bwa_sorted_1.bam") + ".bai"), + ] + + def test_final_outputs_have_no_cleanup_paths(self, tmp_path, rg_lines): + result = self.make(tmp_path, duplicate_marking="none").build() + + assert result.cleanup_paths == [] + + def test_collate_and_input_ref(self, tmp_path, rg_lines): + result = self.make( + tmp_path, + inputs=[tmp_path / "in0.cram"], + collate=True, + input_ref=tmp_path / "decode.fa", + ).build() + + shell = str(result.jobs[0].shell) + assert "samtools collate" in shell + assert f"--reference {tmp_path}/decode.fa" in shell + + def test_build_inserts_nothing(self, tmp_path, rg_lines): + dag = DAG() + self.make(tmp_path).build() + + assert not dag.ready_jobs and not dag.waiting_jobs + + +class TestBwaFastqStage: + """Aligning short-read fastq with bwa""" + + def make(self, tmp_path, **kwargs) -> BwaFastqStage: + defaults = dict( + ctx=make_ctx(tmp_path), + r1_fastq=[tmp_path / "r1.fq.gz"], + r2_fastq=[tmp_path / "r2.fq.gz"], + readgroups=[r"@RG\tID:rg1\tSM:sample"], + model_bundle=tmp_path / "model.bundle", + ) + defaults.update(kwargs) + return BwaFastqStage(**defaults) # type: ignore[arg-type] + + def test_without_numa_one_job_takes_every_core(self, tmp_path): + result = self.make(tmp_path).build() + + assert [job.name for job in result.jobs] == ["bam-align-0-0"] + job = result.jobs[0] + assert job.threads == 4 + assert job.resources == {"node0": 1} + assert job.task_name == "alignment" + assert "taskset" not in str(job.shell) + assert "--split" not in str(job.shell) + + def test_numa_splits_the_pair_across_the_nodes(self, tmp_path): + result = self.make( + tmp_path, ctx=make_ctx(tmp_path, cores=8), numa_nodes=["0-3", "4-7"] + ).build() + + assert [job.name for job in result.jobs] == [ + "bam-align-0-0", + "bam-align-0-1", + ] + assert [job.threads for job in result.jobs] == [4, 4] + assert [job.resources for job in result.jobs] == [ + {"node0": 1}, + {"node1": 1}, + ] + first, second = (str(job.shell) for job in result.jobs) + assert "taskset -c 0-3" in first and "0/2" in first + assert "taskset -c 4-7" in second and "1/2" in second + + def test_one_job_per_pair_and_node(self, tmp_path): + result = self.make( + tmp_path, + r1_fastq=[tmp_path / "a1.fq.gz", tmp_path / "b1.fq.gz"], + r2_fastq=[tmp_path / "a2.fq.gz", tmp_path / "b2.fq.gz"], + readgroups=[r"@RG\tID:a", r"@RG\tID:b"], + numa_nodes=["0-3", "4-7"], + ).build() + + assert [job.name for job in result.jobs] == [ + "bam-align-0-0", + "bam-align-0-1", + "bam-align-1-0", + "bam-align-1-1", + ] + + def test_no_fastq_builds_nothing(self, tmp_path): + result = self.make(tmp_path, r1_fastq=[], r2_fastq=[], readgroups=[]) + + built = result.build() + assert built.jobs == [] + assert built.outputs == [] + assert built.cleanup_paths == [] + + def test_unzip_reaches_the_command(self, tmp_path): + result = self.make(tmp_path, unzip="igzip").build() + + assert "igzip -dc" in str(result.jobs[0].shell) + + def test_intermediate_outputs_and_cleanup_paths(self, tmp_path): + result = self.make(tmp_path).build() + + out = tmp_path / "bwa_sorted_fq_0_0.bam" + assert result.outputs == [out] + assert result.cleanup_paths == [ + out, + pathlib.Path(str(out) + ".bai"), + ] + + def test_final_outputs_have_no_cleanup_paths(self, tmp_path): + result = self.make(tmp_path, duplicate_marking="none").build() + + assert result.outputs == [ + tmp_path / "output_bwa_sorted_fq_0_0.cram" + ] + assert result.cleanup_paths == [] + assert "--cram_write_options" not in str(result.jobs[0].shell) + + +class TestMinimap2RealignStage: + """Re-aligning long-read input with minimap2""" + + def make(self, tmp_path, **kwargs) -> Minimap2RealignStage: + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "lr.bam"], + model_bundle=tmp_path / "model.bundle", + sample_name="sample", + ) + defaults.update(kwargs) + return Minimap2RealignStage(**defaults) # type: ignore[arg-type] + + def test_job_metadata_and_outputs(self, tmp_path, rg_lines): + result = self.make(tmp_path).build() + + assert [job.name for job in result.jobs] == ["bam-realign-0"] + assert result.jobs[0].threads == 4 + assert result.jobs[0].task_name == "alignment" + assert result.outputs == [tmp_path / "output_mm2_sorted_0.cram"] + assert result.cleanup_paths == [] + assert "--cram_write_options" not in str(result.jobs[0].shell) + + def test_default_model_comes_from_the_bundle(self, tmp_path, rg_lines): + result = self.make(tmp_path).build() + + assert f"-x {tmp_path}/model.bundle/minimap2.model" in str( + result.jobs[0].shell + ) + + def test_minimap2_model_overrides_the_bundle_default( + self, tmp_path, rg_lines + ): + model = tmp_path / "model.bundle" / "minimap2_lr.model" + result = self.make(tmp_path, minimap2_model=model).build() + + assert f"-x {model}" in str(result.jobs[0].shell) + + def test_bam_format_and_input_ref(self, tmp_path, rg_lines): + result = self.make( + tmp_path, bam_format=True, input_ref=tmp_path / "lr_ref.fa" + ).build() + + assert result.outputs == [tmp_path / "output_mm2_sorted_0.bam"] + assert f"--reference {tmp_path}/lr_ref.fa" in str(result.jobs[0].shell) + + +class TestMinimap2FastqStage: + """Aligning long-read fastq with minimap2""" + + def test_one_job_per_fastq(self, tmp_path): + result = Minimap2FastqStage( + ctx=make_ctx(tmp_path), + fastq=[tmp_path / "a.fq.gz", tmp_path / "b.fq.gz"], + readgroups=[r"@RG\tID:a", r"@RG\tID:b"], + model_bundle=tmp_path / "model.bundle", + unzip="igzip", + ).build() + + assert [job.name for job in result.jobs] == ["align-0", "align-1"] + assert result.outputs == [ + tmp_path / "output_mm2_sorted_fq_0.cram", + tmp_path / "output_mm2_sorted_fq_1.cram", + ] + assert "igzip -dc" in str(result.jobs[0].shell) + assert "--cram_write_options" not in str(result.jobs[0].shell) + assert result.jobs[0].task_name == "alignment" + + def test_no_fastq_builds_nothing(self, tmp_path): + result = Minimap2FastqStage( + ctx=make_ctx(tmp_path), + fastq=[], + readgroups=[], + model_bundle=tmp_path / "model.bundle", + ).build() + + assert result.jobs == [] + assert result.outputs == [] + + +class TestBwaExtractStage: + """The pangenome bwa alignment and read-extraction pass""" + + def make(self, tmp_path, **kwargs) -> BwaExtractStage: + defaults = dict( + ctx=make_ctx(tmp_path), + output_bam=tmp_path / "sample-bwa.bam", + output_fastq=tmp_path / "sample-extract.fq.gz", + r1_fastq=[tmp_path / "r1.fq.gz"], + r2_fastq=[tmp_path / "r2.fq.gz"], + readgroup=r"@RG\tID:rg1-bwa\tSM:sample", + extract_model=tmp_path / "model.bundle" / "extract.model", + bwa_model=tmp_path / "model.bundle" / "bwa.model", + ) + defaults.update(kwargs) + return BwaExtractStage(**defaults) # type: ignore[arg-type] + + def test_single_job(self, tmp_path): + result = self.make(tmp_path).build() + + assert len(result.jobs) == 1 + job = result.jobs[0] + assert job.name == "bwa-extract" + assert job.threads == 4 + assert job.task_name == "alignment" + assert result.outputs == [ + tmp_path / "sample-bwa.bam", + tmp_path / "sample-extract.fq.gz", + ] + assert result.cleanup_paths == [] + + def test_readgroup_is_used_verbatim(self, tmp_path): + result = self.make( + tmp_path, readgroup=r"@RG\tID:rg1\tSM:sample\tLR:0" + ).build() + + assert r"@RG\tID:rg1\tSM:sample\tLR:0" in str(result.jobs[0].shell) + + def test_unzip_reaches_the_command(self, tmp_path): + result = self.make(tmp_path, unzip="igzip").build() + + assert "igzip -dc" in str(result.jobs[0].shell) + + +class TestReadWriterStage: + """`sentieon driver --algo ReadWriter` merges and extracts""" + + def test_merge(self, tmp_path): + result = ReadWriterStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + output=tmp_path / "merged.bam", + name="merge-bam", + task_name="alignment", + threads=0, + ).build() + + job = result.jobs[0] + assert job.name == "merge-bam" + assert job.threads == 0 + assert job.task_name == "alignment" + assert str(job.shell) == ( + f"sentieon driver --input {tmp_path}/one.bam " + f"--input {tmp_path}/two.bam --reference {tmp_path}/ref.fa " + f"--thread_count 4 --algo ReadWriter {tmp_path}/merged.bam" + ) + assert result.outputs == [tmp_path / "merged.bam"] + + def test_string_interval_extracts_a_locus(self, tmp_path): + result = ReadWriterStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + output=tmp_path / "sample_hla.bam", + name="t1k-hla-extract", + task_name="t1k", + interval="chr6:28510020-33480577", + ).build() + + job = result.jobs[0] + assert job.name == "t1k-hla-extract" + assert job.threads == 4 # defaults to the run's core count + assert job.task_name == "t1k" + assert "--interval chr6:28510020-33480577" in str(job.shell) + + +class TestAddTo: + """`add_to` inserts what `build` returned""" + + def test_jobs_are_roots_without_upstream(self, tmp_path): + dag = DAG() + result = BwaFastqStage( + ctx=make_ctx(tmp_path, cores=8), + r1_fastq=[tmp_path / "r1.fq.gz"], + r2_fastq=[tmp_path / "r2.fq.gz"], + readgroups=[r"@RG\tID:rg1"], + model_bundle=tmp_path / "model.bundle", + numa_nodes=["0-3", "4-7"], + ).add_to(dag) + + assert len(result.jobs) == 2 + for job in result.jobs: + assert job in dag.ready_jobs + assert not dag.waiting_jobs + + def test_every_job_waits_on_the_upstream_jobs(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = Minimap2FastqStage( + ctx=make_ctx(tmp_path), + fastq=[tmp_path / "a.fq.gz", tmp_path / "b.fq.gz"], + readgroups=[r"@RG\tID:a", r"@RG\tID:b"], + model_bundle=tmp_path / "model.bundle", + ).add_to(dag, [upstream]) + + for job in result.jobs: + assert dag.waiting_jobs[job] == {upstream} diff --git a/tests/unit/test_stage_base.py b/tests/unit/test_stage_base.py new file mode 100644 index 0000000..46468ed --- /dev/null +++ b/tests/unit/test_stage_base.py @@ -0,0 +1,364 @@ +""" +Unit tests for the pipeline stage building blocks +""" + +import argparse +import dataclasses +import pathlib +from typing import Iterable +from unittest.mock import patch + +import packaging.version +import pytest + +from sentieon_cli import util +from sentieon_cli.dag import DAG +from sentieon_cli.dnascope import DNAscopePipeline +from sentieon_cli.driver import LocusCollector +from sentieon_cli.job import Job +from sentieon_cli.stages.base import ( + Stage, + StageContext, + StageResult, + driver_job, + rm_job, +) +from sentieon_cli.util import require_versions, versions_available + + +def create_mock_args(loglevel: str = "WARNING") -> argparse.Namespace: + """The argparse namespace `setup_logging` needs""" + return argparse.Namespace(loglevel=loglevel) + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +@dataclasses.dataclass(kw_only=True) +class ToyStage(Stage): + """Two jobs: `b` depends on `a`, `a` is the only entry job""" + + tag: str = "toy" + + def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> StageResult: + deps = set(upstream) + job_a = rm_job( + [self.ctx.tmp_dir / f"{self.tag}-a"], f"{self.tag}-first" + ) + job_b = rm_job( + [self.ctx.tmp_dir / f"{self.tag}-b"], f"{self.tag}-second" + ) + dag.add_job(job_a, deps) + dag.add_job(job_b, {job_a}) + return StageResult(jobs=[job_a, job_b], terminal={job_b}) + + +class TestStageContext: + """The run-wide settings object""" + + def test_fields(self, tmp_path): + ctx = make_ctx(tmp_path) + assert ctx.reference == tmp_path / "ref.fa" + assert ctx.output_vcf == tmp_path / "output.vcf.gz" + assert ctx.tmp_dir == tmp_path + assert ctx.cores == 4 + assert ctx.dry_run is True + assert ctx.skip_version_check is True + + def test_frozen(self, tmp_path): + ctx = make_ctx(tmp_path) + with pytest.raises(dataclasses.FrozenInstanceError): + ctx.cores = 8 # type: ignore[misc] + + +class TestStage: + """The `add_to` contract""" + + def test_upstream_lands_on_entry_jobs_only(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = ToyStage(ctx=make_ctx(tmp_path)).add_to(dag, [upstream]) + + job_a, job_b = result.jobs + assert result.jobs == [job_a, job_b] + assert result.terminal == {job_b} + # The entry job gets the upstream dependencies... + assert dag.waiting_jobs[job_a] == {upstream} + # ...and the internal edge is added on top of them. + assert dag.waiting_jobs[job_b] == {job_a} + + def test_no_upstream_makes_the_entry_job_ready(self, tmp_path): + dag = DAG() + result = ToyStage(ctx=make_ctx(tmp_path)).add_to(dag) + + job_a, job_b = result.jobs + assert job_a in dag.ready_jobs + assert dag.waiting_jobs[job_b] == {job_a} + + def test_upstream_iterable_consumed_once(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = ToyStage(ctx=make_ctx(tmp_path)).add_to(dag, iter([upstream])) + assert dag.waiting_jobs[result.jobs[0]] == {upstream} + + def test_add_to_twice_raises(self, tmp_path): + dag = DAG() + stage = ToyStage(ctx=make_ctx(tmp_path)) + stage.add_to(dag) + # Job identity is the shell pipeline, so the same jobs collide + with pytest.raises(ValueError): + stage.add_to(dag) + + def test_constructing_a_stage_adds_nothing(self, tmp_path): + dag = DAG() + ToyStage(ctx=make_ctx(tmp_path)) + assert not dag.waiting_jobs + assert not dag.ready_jobs + + +class TestDriverJob: + """The `sentieon driver` job helper""" + + def test_basic_command(self, tmp_path): + ctx = make_ctx(tmp_path, cores=8) + job = driver_job( + ctx, + [LocusCollector(tmp_path / "score.txt.gz")], + name="locuscollector", + task_name="dedup", + inputs=[tmp_path / "sample.bam"], + ) + expected = ( + f"sentieon driver --input {tmp_path}/sample.bam " + f"--reference {tmp_path}/ref.fa --thread_count 8 " + f"--algo LocusCollector {tmp_path}/score.txt.gz" + ) + assert str(job.shell) == expected + assert job.name == "locuscollector" + assert job.task_name == "dedup" + assert job.threads == 8 + + def test_no_interval_padding_by_default(self, tmp_path): + job = driver_job( + make_ctx(tmp_path), + [LocusCollector(tmp_path / "score.txt.gz")], + name="locuscollector", + task_name="dedup", + ) + assert "--interval_padding" not in str(job.shell) + assert "--interval" not in str(job.shell) + assert "--read_filter" not in str(job.shell) + assert "--input" not in str(job.shell) + + def test_interval_and_padding(self, tmp_path): + bed = tmp_path / "regions.bed" + job = driver_job( + make_ctx(tmp_path), + [LocusCollector(tmp_path / "score.txt.gz")], + name="locuscollector", + task_name="dedup", + interval=bed, + interval_padding=5, + ) + assert f"--interval {bed} --interval_padding 5" in str(job.shell) + + def test_interval_as_a_string(self, tmp_path): + job = driver_job( + make_ctx(tmp_path), + [LocusCollector(tmp_path / "score.txt.gz")], + name="readwriter", + task_name="extract", + interval="chr6:28510020-33480577", + ) + assert "--interval chr6:28510020-33480577" in str(job.shell) + + def test_threads_zero(self, tmp_path): + job = driver_job( + make_ctx(tmp_path), + [LocusCollector(tmp_path / "score.txt.gz")], + name="locuscollector", + task_name="dedup", + threads=0, + ) + # `threads` is the scheduling budget only; the driver still gets + # the run's core count. + assert job.threads == 0 + assert "--thread_count 4" in str(job.shell) + + def test_read_filter_repeats_the_flag(self, tmp_path): + job = driver_job( + make_ctx(tmp_path), + [LocusCollector(tmp_path / "score.txt.gz")], + name="locuscollector", + task_name="dedup", + read_filter=["QualCalFilter,table=t", "IndelLeftAlign,rgid=rg"], + ) + assert ( + "--read_filter QualCalFilter,table=t " + "--read_filter IndelLeftAlign,rgid=rg" in str(job.shell) + ) + + def test_replace_rg_interleaves_with_input(self, tmp_path): + bam1 = tmp_path / "one.bam" + bam2 = tmp_path / "two.bam" + job = driver_job( + make_ctx(tmp_path), + [LocusCollector(tmp_path / "score.txt.gz")], + name="locuscollector", + task_name="dedup", + inputs=[bam1, bam2], + replace_rg=[["rg1=@RG\\tID:rg1"], ["rg2=@RG\\tID:rg2"]], + ) + shell = str(job.shell) + assert ( + f"--replace_rg 'rg1=@RG\\tID:rg1' --input {bam1} " + f"--replace_rg 'rg2=@RG\\tID:rg2' --input {bam2}" in shell + ) + + def test_algos_are_added_in_order(self, tmp_path): + job = driver_job( + make_ctx(tmp_path), + [ + LocusCollector(tmp_path / "first.txt.gz"), + LocusCollector(tmp_path / "second.txt.gz"), + ], + name="locuscollector", + task_name="dedup", + ) + shell = str(job.shell) + assert shell.index("first.txt.gz") < shell.index("second.txt.gz") + + +class TestRmJob: + """The cleanup job helper""" + + def test_command_and_metadata(self, tmp_path): + vcf = tmp_path / "tmp.vcf.gz" + job = rm_job([vcf, str(vcf) + ".tbi"], "rm-tmp-vcf") + assert str(job.shell) == f"rm {vcf} {vcf}.tbi" + assert job.name == "rm-tmp-vcf" + assert job.task_name == "cleanup" + assert job.threads == 0 + assert job.shell.nodes[0].fail_ok is True + + +class TestRequireVersions: + """`require_versions` and `versions_available`""" + + min_versions = {"sentieon driver": packaging.version.Version("202308")} + + def test_require_skips(self): + with patch.object(util, "check_version") as check: + require_versions(self.min_versions, skip=True) + check.assert_not_called() + + def test_require_passes(self): + with patch.object(util, "check_version", return_value=True) as check: + require_versions(self.min_versions) + check.assert_called_once_with( + "sentieon driver", packaging.version.Version("202308") + ) + + def test_require_exits_on_failure(self): + with patch.object(util, "check_version", return_value=False): + with pytest.raises(SystemExit) as excinfo: + require_versions(self.min_versions) + assert excinfo.value.code == 2 + + def test_require_exits_on_the_first_failure(self): + versions = { + "samtools": packaging.version.Version("1.16"), + "bcftools": packaging.version.Version("1.10"), + } + with patch.object(util, "check_version", return_value=False) as check: + with pytest.raises(SystemExit): + require_versions(versions) + assert check.call_count == 1 + + def test_available_skips(self): + with patch.object(util, "check_version") as check: + assert versions_available(self.min_versions, skip=True) is True + check.assert_not_called() + + def test_available_true(self): + with patch.object(util, "check_version", return_value=True): + assert versions_available(self.min_versions) is True + + def test_available_false(self): + with patch.object(util, "check_version", return_value=False): + assert versions_available(self.min_versions) is False + + +class TestPipelineStageContext: + """`BasePipeline.stage_context` and `BasePipeline.required`""" + + def make_pipeline(self, tmp_path) -> DNAscopePipeline: + pipeline = DNAscopePipeline() + pipeline.setup_logging(create_mock_args()) + pipeline.reference = tmp_path / "ref.fa" + pipeline.output_vcf = tmp_path / "output.vcf.gz" + pipeline.cores = 3 + pipeline.dry_run = True + pipeline.skip_version_check = True + pipeline.tmp_dir = tmp_path + return pipeline + + def test_stage_context(self, tmp_path): + ctx = self.make_pipeline(tmp_path).stage_context() + assert ctx == StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=3, + dry_run=True, + skip_version_check=True, + ) + + def test_missing_reference(self, tmp_path): + pipeline = self.make_pipeline(tmp_path) + pipeline.reference = None + with pytest.raises(SystemExit) as excinfo: + pipeline.stage_context() + assert excinfo.value.code == 2 + + def test_missing_output_vcf(self, tmp_path): + pipeline = self.make_pipeline(tmp_path) + pipeline.output_vcf = None + with pytest.raises(SystemExit) as excinfo: + pipeline.stage_context() + assert excinfo.value.code == 2 + + def test_missing_tmp_dir(self, tmp_path): + pipeline = DNAscopePipeline() + pipeline.setup_logging(create_mock_args()) + pipeline.reference = tmp_path / "ref.fa" + pipeline.output_vcf = tmp_path / "output.vcf.gz" + assert not hasattr(pipeline, "tmp_dir") + with pytest.raises(SystemExit) as excinfo: + pipeline.stage_context() + assert excinfo.value.code == 2 + + def test_required_returns_the_value(self, tmp_path): + pipeline = self.make_pipeline(tmp_path) + assert pipeline.required(pipeline.output_vcf, "output_vcf") == ( + tmp_path / "output.vcf.gz" + ) + + def test_required_exits_when_missing(self, tmp_path): + pipeline = self.make_pipeline(tmp_path) + with pytest.raises(SystemExit) as excinfo: + pipeline.required(None, "model_bundle") + assert excinfo.value.code == 2 diff --git a/tests/unit/test_stage_cnv.py b/tests/unit/test_stage_cnv.py new file mode 100644 index 0000000..e0ee30c --- /dev/null +++ b/tests/unit/test_stage_cnv.py @@ -0,0 +1,374 @@ +""" +Unit tests for the CNVscope stage +""" + +from importlib.resources import files +import logging +import pathlib +import sys +from typing import List + +import packaging.version +import pytest + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.cnv import ( + CNV_MIN_VERSIONS, + CNVscopeStage, + PangenomeCNVStage, +) +from sentieon_cli.util import SampleSex + +PACKAGE_LOGGER = "sentieon_cli" + + +class _RecordingHandler(logging.Handler): + """Collects formatted messages from the package logger.""" + + def __init__(self) -> None: + super().__init__(logging.DEBUG) + self.messages: List[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def messages(): + """`caplog` does not see the package logger, which does not propagate""" + logger = logging.getLogger(PACKAGE_LOGGER) + handler = _RecordingHandler() + logger.addHandler(handler) + try: + yield handler.messages + finally: + logger.removeHandler(handler) + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> CNVscopeStage: + """A CNVscopeStage with the paths a caller supplies""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.cram"], + model=tmp_path / "bundle" / "cnv.model", + cnvscope_vcf=tmp_path / "cnvscope.vcf.gz", + cnv_vcf=tmp_path / "output.cnv.vcf.gz", + ) + defaults.update(kwargs) + return CNVscopeStage(**defaults) # type: ignore[arg-type] + + +class TestMinVersions: + """The version constant moved here with the stage""" + + def test_cnvscope_sex_arguments_need_202503_04(self): + assert CNV_MIN_VERSIONS["sentieon driver"] == ( + packaging.version.Version("202503.04") + ) + + +class TestCommands: + """The `sentieon driver` commands the stage builds""" + + def test_cnvscope_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert str(result.cnvscope_job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.cram " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + f"--algo CNVscope --model {tmp_path}/bundle/cnv.model " + f"{tmp_path}/cnvscope.vcf.gz" + ) + + def test_cnv_model_apply_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + # CNVModelApply reads the CNVscope VCF; it takes no alignment + assert str(result.apply_job.shell) == ( + f"sentieon driver --reference {tmp_path}/ref.fa " + f"--thread_count 4 --algo CNVModelApply " + f"--model {tmp_path}/bundle/cnv.model " + f"--vcf {tmp_path}/cnvscope.vcf.gz " + f"{tmp_path}/output.cnv.vcf.gz" + ) + + def test_interval_is_passed_to_cnvscope_only(self, tmp_path): + bed = tmp_path / "regions.bed" + result = make_stage(tmp_path, interval=bed).add_to(DAG()) + + assert f"--interval {bed}" in str(result.cnvscope_job.shell) + assert "--interval" not in str(result.apply_job.shell) + + def test_replace_rg_is_passed_to_cnvscope_only(self, tmp_path): + result = make_stage( + tmp_path, + replace_rg=[["rg1=ID:rg1\\tSM:sample"]], + ).add_to(DAG()) + + assert "--replace_rg 'rg1=ID:rg1\\tSM:sample'" in str( + result.cnvscope_job.shell + ) + assert "--replace_rg" not in str(result.apply_job.shell) + + def test_dump_probes_is_passed_to_cnvscope_only(self, tmp_path): + probes = tmp_path / "output.cnv.probes" + result = make_stage(tmp_path, dump_probes=probes).add_to(DAG()) + + assert f"--dump_probes {probes}" in str(result.cnvscope_job.shell) + assert "--dump_probes" not in str(result.apply_job.shell) + + def test_multiple_inputs(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + ).add_to(DAG()) + + shell = str(result.cnvscope_job.shell) + assert f"--input {tmp_path}/one.bam" in shell + assert f"--input {tmp_path}/two.bam" in shell + + +class TestSexArguments: + """`--sex` and `--par` follow the sample sex""" + + def test_male_sample(self, tmp_path): + par_bed = tmp_path / "par.bed" + result = make_stage( + tmp_path, sample_sex=SampleSex.MALE, par_bed=par_bed + ).add_to(DAG()) + + shell = str(result.cnvscope_job.shell) + assert "--sex M" in shell + assert f"--par {par_bed}" in shell + + def test_female_sample(self, tmp_path): + # A female sample does not need the PAR regions + result = make_stage( + tmp_path, + sample_sex=SampleSex.FEMALE, + par_bed=tmp_path / "par.bed", + ).add_to(DAG()) + + shell = str(result.cnvscope_job.shell) + assert "--sex F" in shell + assert "--par" not in shell + + @pytest.mark.parametrize("sample_sex", [None, SampleSex.UNKNOWN]) + def test_without_a_known_sex(self, tmp_path, sample_sex, messages): + result = make_stage(tmp_path, sample_sex=sample_sex).add_to(DAG()) + + shell = str(result.cnvscope_job.shell) + assert "--sex" not in shell + assert "--par" not in shell + assert any("diploid" in msg for msg in messages) + + +class TestJobMetadata: + """Job names, task name and thread counts""" + + def test_default_names(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.cnvscope_job.name == "cnvscope" + assert result.apply_job.name == "cnv-model-apply" + + def test_overridden_names(self, tmp_path): + result = make_stage( + tmp_path, name="CNVscope", apply_name="CNVModelApply" + ).add_to(DAG()) + + assert result.cnvscope_job.name == "CNVscope" + assert result.apply_job.name == "CNVModelApply" + + def test_task_name_and_threads(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + for job in result.jobs: + assert job.task_name == "cnv" + assert job.threads == 4 + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_lands_on_cnvscope_only(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.cnvscope_job] == {upstream} + assert dag.waiting_jobs[result.apply_job] == {result.cnvscope_job} + + def test_cnvscope_is_a_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.cnvscope_job in dag.ready_jobs + assert dag.waiting_jobs[result.apply_job] == {result.cnvscope_job} + + def test_result_fields(self, tmp_path): + cnv_vcf = tmp_path / "output.cnv.vcf.gz" + result = make_stage(tmp_path, cnv_vcf=cnv_vcf).add_to(DAG()) + + assert result.jobs == [result.cnvscope_job, result.apply_job] + assert result.terminal == {result.apply_job} + assert result.cnv_vcf == cnv_vcf + + +def make_pangenome_stage( + tmp_path: pathlib.Path, **kwargs +) -> PangenomeCNVStage: + """A PangenomeCNVStage over the short-read alignment and the SV VCF""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.cram"], + model=tmp_path / "bundle" / "cnv.model", + sv_vcf=tmp_path / "output_sv.vcf.gz", + ) + defaults.update(kwargs) + return PangenomeCNVStage(**defaults) # type: ignore[arg-type] + + +def script_path(name: str) -> str: + """The packaged helper script, unresolved as the stage builds it""" + return str(files("sentieon_cli.scripts").joinpath(name)) + + +class TestPangenomeCNVStage: + """CNVscope, CNVModelApply, indel2cnv and combine_cnv together""" + + def test_job_names_threads_and_task(self, tmp_path): + result = make_pangenome_stage(tmp_path).add_to(DAG()) + + assert [job.name for job in result.jobs] == [ + "cnvscope", + "cnv-model-apply", + "indel2cnv", + "combine-cnv", + ] + assert [job.threads for job in result.jobs] == [4, 4, 0, 0] + for job in result.jobs: + assert job.task_name == "cnv" + + def test_derived_paths(self, tmp_path): + result = make_pangenome_stage(tmp_path).add_to(DAG()) + + # The intermediates live in the temporary directory; the CNV VCF + # and the probes file sit next to the output VCF + cnvscope_shell = str(result.cnvscope_job.shell) + assert f"{tmp_path}/sample-cnvscope.vcf.gz" in cnvscope_shell + assert f"--dump_probes {tmp_path}/output_cnv.probes" in cnvscope_shell + + apply_shell = str(result.apply_job.shell) + assert f"--vcf {tmp_path}/sample-cnvscope.vcf.gz" in apply_shell + assert f"{tmp_path}/sample-cnv_model_apply.vcf.gz" in apply_shell + + assert result.cnv_vcf == tmp_path / "output_cnv.vcf.gz" + + def test_indel2cnv_command(self, tmp_path): + result = make_pangenome_stage(tmp_path).add_to(DAG()) + + assert str(result.indel2cnv_job.shell) == ( + f"{sys.executable} {script_path('indel2cnv.py')} " + f"{tmp_path}/ref.fa {tmp_path}/output_sv.vcf.gz " + f"{tmp_path}/sample-sv_cnv.vcf.gz -t 4" + ) + + def test_combine_cnv_command(self, tmp_path): + result = make_pangenome_stage(tmp_path).add_to(DAG()) + + assert str(result.combine_job.shell) == ( + f"{sys.executable} {script_path('combine_cnv.py')} " + f"--cnv {tmp_path}/sample-cnv_model_apply.vcf.gz " + f"--converted {tmp_path}/sample-sv_cnv.vcf.gz " + f"-o {tmp_path}/output_cnv.vcf.gz" + ) + + def test_upstream_lands_on_the_entry_jobs(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_pangenome_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.cnvscope_job] == {upstream} + assert dag.waiting_jobs[result.indel2cnv_job] == {upstream} + assert dag.waiting_jobs[result.apply_job] == {result.cnvscope_job} + assert dag.waiting_jobs[result.combine_job] == { + result.apply_job, + result.indel2cnv_job, + } + + def test_entry_jobs_are_roots_without_upstream(self, tmp_path): + dag = DAG() + result = make_pangenome_stage(tmp_path).add_to(dag) + + assert result.cnvscope_job in dag.ready_jobs + assert result.indel2cnv_job in dag.ready_jobs + + def test_interval_and_replace_rg_reach_cnvscope_only(self, tmp_path): + bed = tmp_path / "regions.bed" + result = make_pangenome_stage( + tmp_path, + interval=bed, + replace_rg=[["rg1=ID:rg1\\tSM:sample"]], + ).add_to(DAG()) + + cnvscope_shell = str(result.cnvscope_job.shell) + assert f"--interval {bed}" in cnvscope_shell + assert "--replace_rg 'rg1=ID:rg1\\tSM:sample'" in cnvscope_shell + for job in ( + result.apply_job, + result.indel2cnv_job, + result.combine_job, + ): + assert "--interval" not in str(job.shell) + assert "--replace_rg" not in str(job.shell) + + def test_sex_arguments_reach_cnvscope(self, tmp_path): + par_bed = tmp_path / "par.bed" + result = make_pangenome_stage( + tmp_path, sample_sex=SampleSex.MALE, par_bed=par_bed + ).add_to(DAG()) + + shell = str(result.cnvscope_job.shell) + assert "--sex M" in shell + assert f"--par {par_bed}" in shell + + def test_female_sample_needs_no_par(self, tmp_path): + result = make_pangenome_stage( + tmp_path, + sample_sex=SampleSex.FEMALE, + par_bed=tmp_path / "par.bed", + ).add_to(DAG()) + + shell = str(result.cnvscope_job.shell) + assert "--sex F" in shell + assert "--par" not in shell + + def test_result_fields(self, tmp_path): + result = make_pangenome_stage(tmp_path).add_to(DAG()) + + assert result.jobs == [ + result.cnvscope_job, + result.apply_job, + result.indel2cnv_job, + result.combine_job, + ] + assert result.terminal == {result.combine_job} + assert result.cnv_vcf == tmp_path / "output_cnv.vcf.gz" diff --git a/tests/unit/test_stage_dedup.py b/tests/unit/test_stage_dedup.py new file mode 100644 index 0000000..b48081e --- /dev/null +++ b/tests/unit/test_stage_dedup.py @@ -0,0 +1,200 @@ +""" +Unit tests for the duplicate-marking stage +""" + +import pathlib + +from sentieon_cli.dag import DAG +from sentieon_cli.driver import AlignmentStat, GCBias +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.dedup import DedupStage + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> DedupStage: + """A DedupStage with the paths a caller supplies""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + output=tmp_path / "deduped.cram", + score_file=tmp_path / "score.txt.gz", + ) + defaults.update(kwargs) + return DedupStage(**defaults) # type: ignore[arg-type] + + +class TestCommands: + """The `sentieon driver` commands the stage builds""" + + def test_locuscollector_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert str(result.lc_job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.bam " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + f"--algo LocusCollector {tmp_path}/score.txt.gz" + ) + + def test_dedup_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + # The defaults produce the shortest possible Dedup command + assert str(result.dedup_job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.bam " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + f"--algo Dedup --score_info {tmp_path}/score.txt.gz " + f"{tmp_path}/deduped.cram" + ) + + def test_consensus_reads(self, tmp_path): + result = make_stage(tmp_path, consensus=True).add_to(DAG()) + + assert "--algo LocusCollector --consensus" in str(result.lc_job.shell) + assert "--consensus" not in str(result.dedup_job.shell) + + def test_rmdup_and_cram_write_options(self, tmp_path): + result = make_stage( + tmp_path, + rmdup=True, + cram_write_options="version=3.0,compressor=rans", + ).add_to(DAG()) + + shell = str(result.dedup_job.shell) + assert "--cram_write_options version=3.0,compressor=rans" in shell + assert "--rmdup" in shell + # The options belong to Dedup alone + assert "--rmdup" not in str(result.lc_job.shell) + + def test_dedup_metrics(self, tmp_path): + metrics = tmp_path / "sample.dedup_metrics.txt" + result = make_stage(tmp_path, dedup_metrics=metrics).add_to(DAG()) + + assert f"--metrics {metrics}" in str(result.dedup_job.shell) + + def test_without_optional_dedup_arguments(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + shell = str(result.dedup_job.shell) + for flag in ("--cram_write_options", "--metrics", "--rmdup"): + assert flag not in shell + + def test_read_filters_reach_both_passes(self, tmp_path): + read_filter = "IndelLeftAlignReadTransform,rgid=rg1-mm2" + result = make_stage(tmp_path, read_filters=[read_filter]).add_to(DAG()) + + for job in (result.lc_job, result.dedup_job): + assert f"--read_filter {read_filter}" in str(job.shell) + + def test_no_read_filters_emits_nothing(self, tmp_path): + result = make_stage(tmp_path, read_filters=[]).add_to(DAG()) + + for job in (result.lc_job, result.dedup_job): + assert "--read_filter" not in str(job.shell) + + def test_multiple_inputs(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + ).add_to(DAG()) + + for job in (result.lc_job, result.dedup_job): + shell = str(job.shell) + assert f"--input {tmp_path}/one.bam" in shell + assert f"--input {tmp_path}/two.bam" in shell + + +class TestExtraAlgos: + """Metrics that ride along on the LocusCollector pass""" + + def test_extra_algos_follow_locuscollector(self, tmp_path): + result = make_stage( + tmp_path, + lc_extra_algos=[ + AlignmentStat(tmp_path / "as.txt"), + GCBias(tmp_path / "gc.txt", summary=tmp_path / "sum.txt"), + ], + ).add_to(DAG()) + + shell = str(result.lc_job.shell) + assert shell.index("--algo LocusCollector") < shell.index( + "--algo AlignmentStat" + ) + assert shell.index("--algo AlignmentStat") < shell.index( + "--algo GCBias" + ) + # They stay off the Dedup pass + assert "--algo AlignmentStat" not in str(result.dedup_job.shell) + + +class TestJobMetadata: + """Job names, task name and thread counts""" + + def test_default_names(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.lc_job.name == "locuscollector" + assert result.dedup_job.name == "dedup" + + def test_tagged_names(self, tmp_path): + result = make_stage(tmp_path, tag="bwa").add_to(DAG()) + + assert result.lc_job.name == "locuscollector-bwa" + assert result.dedup_job.name == "dedup-bwa" + + def test_threads_default_to_the_core_count(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + for job in result.jobs: + assert job.threads == 4 + assert job.task_name == "dedup" + + def test_overridden_threads_and_task_name(self, tmp_path): + result = make_stage( + tmp_path, threads=1, task_name="preprocess" + ).add_to(DAG()) + + for job in result.jobs: + assert job.threads == 1 + assert job.task_name == "preprocess" + # The driver still uses every core + assert "--thread_count 4" in str(result.lc_job.shell) + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_lands_on_locuscollector_only(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.lc_job] == {upstream} + assert dag.waiting_jobs[result.dedup_job] == {result.lc_job} + + def test_locuscollector_is_a_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.lc_job in dag.ready_jobs + assert dag.waiting_jobs[result.dedup_job] == {result.lc_job} + + def test_result_fields(self, tmp_path): + output = tmp_path / "deduped.cram" + result = make_stage(tmp_path, output=output).add_to(DAG()) + + assert result.jobs == [result.lc_job, result.dedup_job] + assert result.terminal == {result.dedup_job} + assert result.output == output diff --git a/tests/unit/test_stage_expansion.py b/tests/unit/test_stage_expansion.py new file mode 100644 index 0000000..78026b3 --- /dev/null +++ b/tests/unit/test_stage_expansion.py @@ -0,0 +1,130 @@ +""" +Unit tests for the ExpansionHunter stage +""" + +import pathlib + +import pytest + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.expansion import ( + EXPANSION_MIN_VERSIONS, + ExpansionHunterStage, +) +from sentieon_cli.util import SampleSex + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> ExpansionHunterStage: + """An ExpansionHunterStage with the paths a caller supplies""" + defaults = dict( + ctx=make_ctx(tmp_path), + alignment=tmp_path / "sample.cram", + variant_catalog=tmp_path / "catalog.json", + output_prefix=tmp_path / "output_expansion", + ) + defaults.update(kwargs) + return ExpansionHunterStage(**defaults) # type: ignore[arg-type] + + +class TestMinVersions: + """ExpansionHunter has no minimum version""" + + def test_no_minimum_version(self): + assert EXPANSION_MIN_VERSIONS == {"ExpansionHunter": None} + + +class TestCommands: + """The `ExpansionHunter` command the stage builds""" + + def test_expansion_hunter_command(self, tmp_path): + result = make_stage(tmp_path, sample_sex=SampleSex.MALE).add_to(DAG()) + + assert str(result.job.shell) == ( + f"ExpansionHunter --reads {tmp_path}/sample.cram " + f"--reference {tmp_path}/ref.fa " + f"--variant-catalog {tmp_path}/catalog.json " + "--sex male --threads 4 " + f"--output-prefix {tmp_path}/output_expansion" + ) + + @pytest.mark.parametrize( + "sample_sex,expected", + [ + (SampleSex.MALE, "male"), + (SampleSex.FEMALE, "female"), + (SampleSex.UNKNOWN, "female"), + (None, "female"), + ], + ) + def test_sex_argument(self, tmp_path, sample_sex, expected): + result = make_stage(tmp_path, sample_sex=sample_sex).add_to(DAG()) + + assert f"--sex {expected} " in str(result.job.shell) + + def test_threads_follow_the_run_cores(self, tmp_path): + result = make_stage(tmp_path, ctx=make_ctx(tmp_path, cores=16)).add_to( + DAG() + ) + + assert "--threads 16" in str(result.job.shell) + + +class TestJobMetadata: + """Job name, task name and thread count""" + + def test_default_names_and_threads(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.job.name == "expansion-hunter" + assert result.job.task_name == "expansion-hunter" + assert result.job.threads == 4 + + def test_overridden_names(self, tmp_path): + result = make_stage( + tmp_path, name="expansions", task_name="expansions" + ).add_to(DAG()) + + assert result.job.name == "expansions" + assert result.job.task_name == "expansions" + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_becomes_the_dependencies(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.job] == {upstream} + + def test_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.job in dag.ready_jobs + + def test_result_fields(self, tmp_path): + output_prefix = tmp_path / "output_expansion" + result = make_stage(tmp_path, output_prefix=output_prefix).add_to( + DAG() + ) + + assert result.jobs == [result.job] + assert result.terminal == {result.job} + assert result.output_prefix == output_prefix diff --git a/tests/unit/test_stage_metrics.py b/tests/unit/test_stage_metrics.py new file mode 100644 index 0000000..7170b1d --- /dev/null +++ b/tests/unit/test_stage_metrics.py @@ -0,0 +1,332 @@ +""" +Unit tests for the metrics stage and the metrics paths +""" + +import pathlib +import sys + +from sentieon_cli.dag import DAG +from sentieon_cli.driver import ( + AlignmentStat, + InsertSizeMetricAlgo, + WgsMetricsAlgo, +) +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.metrics import ( + MetricsPaths, + MetricsStage, + MosdepthStage, +) + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> MetricsStage: + """A MetricsStage with the paths a caller supplies""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "deduped.cram"], + algos=[AlignmentStat(tmp_path / "as.txt")], + ) + defaults.update(kwargs) + return MetricsStage(**defaults) # type: ignore[arg-type] + + +class TestMetricsPaths: + """Every metrics file is derived from the output VCF""" + + def test_directory_and_base(self): + paths = MetricsPaths.from_output_vcf( + pathlib.Path("/data/run/sample.vcf.gz") + ) + + assert paths.metrics_dir == pathlib.Path("/data/run/sample_metrics") + assert paths.metric_base == "sample.txt" + + def test_metric_files(self): + paths = MetricsPaths.from_output_vcf(pathlib.Path("/data/out.vcf.gz")) + base = pathlib.Path("/data/out_metrics") + + assert paths.score == base / "out.txt.score.txt.gz" + assert paths.insert_size == base / "out.txt.insert_size.txt" + assert ( + paths.mean_qual_by_cycle == base / "out.txt.mean_qual_by_cycle.txt" + ) + assert paths.base_distribution_by_cycle == ( + base / "out.txt.base_distribution_by_cycle.txt" + ) + assert paths.qual_distribution == ( + base / "out.txt.qual_distribution.txt" + ) + assert paths.alignment_stat == base / "out.txt.alignment_stat.txt" + assert paths.hybrid_selection == base / "out.txt.hybrid-selection.txt" + assert paths.wgs == base / "out.txt.wgs.txt" + assert paths.gc_bias == base / "out.txt.gc_bias.txt" + assert paths.gc_bias_summary == base / "out.txt.gc_bias_summary.txt" + assert paths.dedup_metrics == base / "out.txt.dedup_metrics.txt" + + def test_coverage_is_not_named_after_the_vcf(self): + paths = MetricsPaths.from_output_vcf(pathlib.Path("/data/out.vcf.gz")) + + assert paths.coverage == pathlib.Path("/data/out_metrics/coverage") + + def test_ensure_dir_creates_the_directory(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + paths.ensure_dir(dry_run=False) + + assert paths.metrics_dir.is_dir() + + def test_ensure_dir_creates_nothing_on_a_dry_run(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + paths.ensure_dir(dry_run=True) + + assert not paths.metrics_dir.exists() + + def test_ensure_dir_tolerates_an_existing_directory(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + paths.metrics_dir.mkdir() + paths.ensure_dir(dry_run=False) + + assert paths.metrics_dir.is_dir() + + +class TestCommands: + """The commands the stage builds""" + + def test_metrics_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert str(result.metrics_job.shell) == ( + f"sentieon driver --input {tmp_path}/deduped.cram " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + f"--algo AlignmentStat --adapter_seq '' {tmp_path}/as.txt" + ) + + def test_algo_order_is_the_callers(self, tmp_path): + result = make_stage( + tmp_path, + algos=[ + InsertSizeMetricAlgo(tmp_path / "is.txt"), + WgsMetricsAlgo(tmp_path / "wgs.txt", include_unpaired="true"), + AlignmentStat(tmp_path / "as.txt"), + ], + ).add_to(DAG()) + + shell = str(result.metrics_job.shell) + assert shell.index("--algo InsertSizeMetricAlgo") < shell.index( + "--algo WgsMetricsAlgo" + ) + assert shell.index("--algo WgsMetricsAlgo") < shell.index( + "--algo AlignmentStat" + ) + + def test_interval(self, tmp_path): + bed = tmp_path / "regions.bed" + result = make_stage(tmp_path, interval=bed).add_to(DAG()) + + assert f"--interval {bed}" in str(result.metrics_job.shell) + + def test_without_an_interval(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert "--interval" not in str(result.metrics_job.shell) + + def test_multiple_inputs(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.cram", tmp_path / "two.cram"], + ).add_to(DAG()) + + shell = str(result.metrics_job.shell) + assert f"--input {tmp_path}/one.cram" in shell + assert f"--input {tmp_path}/two.cram" in shell + + +class TestRehead: + """The optional WGS-metrics rehead job""" + + def test_rehead_command(self, tmp_path): + wgs = tmp_path / "output.txt.wgs.txt" + result = make_stage(tmp_path, rehead_metrics=wgs).add_to(DAG()) + + assert result.rehead_job is not None + shell = str(result.rehead_job.shell) + assert shell.startswith(sys.executable) + assert "rehead_wgs_metrics.py" in shell + assert f"--metrics_file {wgs}" in shell + + def test_rehead_job_metadata(self, tmp_path): + result = make_stage( + tmp_path, rehead_metrics=tmp_path / "wgs.txt" + ).add_to(DAG()) + + assert result.rehead_job is not None + assert result.rehead_job.name == "rehead-metrics" + assert result.rehead_job.threads == 0 + assert result.rehead_job.task_name == "metrics" + + def test_no_rehead_job_by_default(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.rehead_job is None + assert result.jobs == [result.metrics_job] + assert result.terminal == {result.metrics_job} + + def test_terminal_is_the_rehead_job(self, tmp_path): + result = make_stage( + tmp_path, rehead_metrics=tmp_path / "wgs.txt" + ).add_to(DAG()) + + assert result.jobs == [result.metrics_job, result.rehead_job] + assert result.terminal == {result.rehead_job} + + +class TestJobMetadata: + """Job name, task name and thread count""" + + def test_defaults(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.metrics_job.name == "metrics" + assert result.metrics_job.task_name == "metrics" + # Metrics run in the background by default + assert result.metrics_job.threads == 0 + + def test_custom_name_task_and_threads(self, tmp_path): + result = make_stage( + tmp_path, + name="sr-metrics", + task_name="qc", + threads=4, + ).add_to(DAG()) + + assert result.metrics_job.name == "sr-metrics" + assert result.metrics_job.task_name == "qc" + assert result.metrics_job.threads == 4 + + def test_rehead_keeps_the_metrics_task(self, tmp_path): + result = make_stage( + tmp_path, + name="sr-metrics", + task_name="qc", + rehead_metrics=tmp_path / "wgs.txt", + ).add_to(DAG()) + + assert result.rehead_job is not None + assert result.rehead_job.task_name == "metrics" + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_lands_on_the_metrics_job(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage( + tmp_path, rehead_metrics=tmp_path / "wgs.txt" + ).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.metrics_job] == {upstream} + assert dag.waiting_jobs[result.rehead_job] == {result.metrics_job} + + def test_metrics_is_a_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.metrics_job in dag.ready_jobs + + +class TestMosdepth: + """Coverage QC, one job per input file""" + + def test_mosdepth_command(self, tmp_path): + result = MosdepthStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + ).add_to(DAG()) + + (job,) = result.jobs + assert str(job.shell) == ( + f"mosdepth --fasta {tmp_path}/ref.fa --threads 4 " + "--by 500 --no-per-base --use-median " + "-T 1,3,5,10,15,20,30,40,50 " + f"{tmp_path}/output_mosdepth_0 {tmp_path}/sample.bam" + ) + + def test_one_job_per_input(self, tmp_path): + result = MosdepthStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + ).add_to(DAG()) + + assert [job.name for job in result.jobs] == [ + "mosdepth-0", + "mosdepth-1", + ] + assert f"{tmp_path}/output_mosdepth_0" in str(result.jobs[0].shell) + assert f"{tmp_path}/output_mosdepth_1" in str(result.jobs[1].shell) + + def test_job_metadata(self, tmp_path): + result = MosdepthStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + ).add_to(DAG()) + + (job,) = result.jobs + assert job.task_name == "metrics" + # mosdepth runs in the background + assert job.threads == 0 + + def test_no_inputs_adds_nothing(self, tmp_path): + dag = DAG() + result = MosdepthStage(ctx=make_ctx(tmp_path), inputs=[]).add_to(dag) + + assert result.jobs == [] + assert result.terminal == set() + assert not dag.ready_jobs + assert not dag.waiting_jobs + + def test_build_inserts_nothing(self, tmp_path): + dag = DAG() + result = MosdepthStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + ).build() + + assert result.terminal == set(result.jobs) + assert not dag.ready_jobs + assert not dag.waiting_jobs + + def test_upstream_lands_on_every_job(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = MosdepthStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + ).add_to(dag, [upstream]) + + for job in result.jobs: + assert dag.waiting_jobs[job] == {upstream} + + def test_roots_without_upstream(self, tmp_path): + dag = DAG() + result = MosdepthStage( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + ).add_to(dag) + + assert result.jobs[0] in dag.ready_jobs diff --git a/tests/unit/test_stage_ploidy.py b/tests/unit/test_stage_ploidy.py new file mode 100644 index 0000000..d696903 --- /dev/null +++ b/tests/unit/test_stage_ploidy.py @@ -0,0 +1,130 @@ +""" +Unit tests for the ploidy estimation stage +""" + +import pathlib + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.ploidy import PloidyStage + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> PloidyStage: + """A PloidyStage over one alignment""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.cram"], + ) + defaults.update(kwargs) + return PloidyStage(**defaults) # type: ignore[arg-type] + + +class TestOutput: + """The ploidy JSON sits next to the output VCF""" + + def test_ploidy_json_derivation(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.ploidy_json == tmp_path / "output_ploidy.json" + assert str(result.ploidy_json) in str(result.job.shell) + + +class TestCommand: + """The `estimate_ploidy.py` command line""" + + def test_script_and_inputs(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + ).add_to(DAG()) + + cmd = str(result.job.shell) + assert cmd.startswith("python3 ") + assert "estimate_ploidy.py" in cmd + assert f"-i {tmp_path}/one.bam {tmp_path}/two.bam" in cmd + assert f"--outfile {tmp_path}/output_ploidy.json" in cmd + + def test_the_script_defaults_apply_without_a_build(self, tmp_path): + # The `estimate_ploidy.py` defaults are the chr-prefixed names + cmd = str(make_stage(tmp_path).add_to(DAG()).job.shell) + + assert "--contigs" not in cmd + assert "--autosomes" not in cmd + assert "--x_contig" not in cmd + assert "--y_contig" not in cmd + + def test_an_unrecognized_build_uses_the_defaults(self, tmp_path): + cmd = str( + make_stage(tmp_path, reference_build="hg38") + .add_to(DAG()) + .job.shell + ) + + assert "--contigs" not in cmd + assert "--x_contig" not in cmd + + def test_b37_contig_names(self, tmp_path): + cmd = str( + make_stage(tmp_path, reference_build="b37").add_to(DAG()).job.shell + ) + + assert "--contigs 1 2 " in cmd + assert "--autosomes 1 2 " in cmd + # The sex chromosomes are in `--contigs` but not in `--autosomes` + assert ( + "--contigs " + + " ".join([str(i) for i in range(1, 23)] + ["X", "Y"]) + in cmd + ) + assert "--autosomes " + " ".join([str(i) for i in range(1, 23)]) in cmd + assert "--x_contig X" in cmd + assert "--y_contig Y" in cmd + + +class TestJobMetadata: + """Job name, task name and thread count""" + + def test_metadata(self, tmp_path): + job = make_stage(tmp_path).add_to(DAG()).job + + assert job.name == "estimate-ploidy" + assert job.task_name == "ploidy" + # The script is single-threaded and does not reserve a core + assert job.threads == 0 + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_edge(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.job] == {upstream} + + def test_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.job in dag.ready_jobs + + def test_result_fields(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.jobs == [result.job] + assert result.terminal == {result.job} diff --git a/tests/unit/test_stage_preprocessing.py b/tests/unit/test_stage_preprocessing.py new file mode 100644 index 0000000..efd24d7 --- /dev/null +++ b/tests/unit/test_stage_preprocessing.py @@ -0,0 +1,345 @@ +""" +Unit tests for the short-read preprocessing stage +""" + +import pathlib + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.metrics import MetricsPaths +from sentieon_cli.stages.preprocessing import ( + ShortReadPreprocessingStage, + pre_dedup_metrics_algos, +) + +# The pre-dedup metrics, in the order the stage collects them +PRE_DEDUP_ALGOS = [ + "MeanQualityByCycle", + "BaseDistributionByCycle", + "QualDistribution", + "AlignmentStat", +] + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage( + tmp_path: pathlib.Path, **kwargs +) -> ShortReadPreprocessingStage: + """A preprocessing stage over one input alignment""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + ) + defaults.update(kwargs) + return ShortReadPreprocessingStage(**defaults) # type: ignore[arg-type] + + +def algo_order(shell: str) -> list: + """The algos a driver command runs, in command-line order""" + tokens = shell.split() + return [ + tokens[i + 1] + for i, token in enumerate(tokens[:-1]) + if token == "--algo" + ] + + +class TestPreDedupAlgos: + """The metrics that do not need duplicate-marked reads""" + + def test_wgs_adds_gc_bias(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + algos = pre_dedup_metrics_algos(paths, "WGS") + + assert [algo.name for algo in algos] == PRE_DEDUP_ALGOS + ["GCBias"] + + def test_wes_omits_gc_bias(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + algos = pre_dedup_metrics_algos(paths, "WES") + + assert [algo.name for algo in algos] == PRE_DEDUP_ALGOS + + def test_the_outputs_come_from_the_metrics_paths(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + algos = pre_dedup_metrics_algos(paths, "WGS") + + assert algos[0].output == paths.mean_qual_by_cycle + assert algos[1].output == paths.base_distribution_by_cycle + assert algos[2].output == paths.qual_distribution + assert algos[3].output == paths.alignment_stat + assert algos[4].output == paths.gc_bias + assert algos[4].summary == paths.gc_bias_summary + + +class TestNoDuplicateMarking: + """`duplicate_marking="none"` collects everything in one pass""" + + def test_skip_metrics_adds_no_jobs(self, tmp_path): + dag = DAG() + inputs = [tmp_path / "sample.bam"] + result = make_stage( + tmp_path, + inputs=inputs, + duplicate_marking="none", + skip_metrics=True, + ).add_to(dag) + + assert result.jobs == [] + assert result.terminal == set() + assert result.qc_jobs == [] + assert result.dedup_job is None + assert result.deduped == inputs + assert not dag.ready_jobs + assert not dag.waiting_jobs + + def test_one_metrics_job(self, tmp_path): + inputs = [tmp_path / "sample.bam"] + result = make_stage( + tmp_path, inputs=inputs, duplicate_marking="none" + ).add_to(DAG()) + + (metrics,) = result.jobs + assert metrics.name == "metrics" + assert metrics.task_name == "metrics" + # This pass reads the whole alignment, so it takes every core + assert metrics.threads == 4 + assert result.dedup_job is None + assert result.deduped == inputs + assert result.qc_jobs == [metrics] + + def test_wgs_algo_order(self, tmp_path): + result = make_stage(tmp_path, duplicate_marking="none").add_to(DAG()) + + assert algo_order(str(result.jobs[0].shell)) == ( + ["InsertSizeMetricAlgo"] + PRE_DEDUP_ALGOS + ["GCBias"] + ) + + def test_wes_omits_gc_bias(self, tmp_path): + result = make_stage( + tmp_path, duplicate_marking="none", assay="WES" + ).add_to(DAG()) + + assert algo_order(str(result.jobs[0].shell)) == ( + ["InsertSizeMetricAlgo"] + PRE_DEDUP_ALGOS + ) + + def test_the_metrics_job_reads_the_inputs(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + duplicate_marking="none", + ).add_to(DAG()) + + shell = str(result.jobs[0].shell) + assert f"--input {tmp_path}/one.bam" in shell + assert f"--input {tmp_path}/two.bam" in shell + + def test_upstream_lands_on_the_metrics_job(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path, duplicate_marking="none").add_to( + dag, [upstream] + ) + + assert dag.waiting_jobs[result.jobs[0]] == {upstream} + # Nothing should wait on this leaf + assert result.terminal == set() + + +class TestMarkdupWgs: + """The default: duplicate marking plus WGS metrics""" + + def test_result_fields(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + lc, dedup, metrics, rehead = result.jobs + assert lc.name == "locuscollector" + assert dedup.name == "dedup" + assert metrics.name == "metrics" + assert rehead.name == "rehead-metrics" + assert result.dedup_job is dedup + assert result.terminal == {dedup} + assert result.qc_jobs == [lc, metrics, rehead] + assert result.deduped == [tmp_path / "output_deduped.cram"] + + def test_locuscollector_carries_the_pre_dedup_metrics(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert algo_order(str(result.jobs[0].shell)) == ( + ["LocusCollector"] + PRE_DEDUP_ALGOS + ["GCBias"] + ) + + def test_dedup_writes_metrics_with_default_cram_options( + self, tmp_path + ): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + result = make_stage(tmp_path).add_to(DAG()) + + # The default pipeline writes CRAM with the driver defaults + shell = str(result.dedup_job.shell) + assert f"--metrics {paths.dedup_metrics}" in shell + assert "--cram_write_options" not in shell + assert "--rmdup" not in shell + + def test_post_dedup_metrics(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + metrics = result.jobs[2] + assert algo_order(str(metrics.shell)) == [ + "InsertSizeMetricAlgo", + "WgsMetricsAlgo", + "CoverageMetrics", + ] + assert "--include_unpaired true" in str(metrics.shell) + # Metrics run in the background + assert metrics.threads == 0 + assert f"--input {tmp_path}/output_deduped.cram" in str(metrics.shell) + + def test_the_post_job_waits_on_dedup(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert dag.waiting_jobs[result.jobs[2]] == {result.dedup_job} + assert dag.waiting_jobs[result.jobs[3]] == {result.jobs[2]} + + def test_upstream_lands_on_locuscollector(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.jobs[0]] == {upstream} + + def test_insert_size_is_not_on_the_locuscollector_pass(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert "InsertSizeMetricAlgo" not in algo_order( + str(result.jobs[0].shell) + ) + + +class TestWes: + """WES runs collect hybrid-selection metrics instead""" + + def test_with_a_bed(self, tmp_path): + bed = tmp_path / "regions.bed" + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + result = make_stage(tmp_path, assay="WES", bed=bed).add_to(DAG()) + + lc, dedup, metrics = result.jobs + assert algo_order(str(metrics.shell)) == [ + "HsMetricAlgo", + "InsertSizeMetricAlgo", + ] + assert f"--interval {bed}" in str(metrics.shell) + assert f"--targets_list {bed}" in str(metrics.shell) + assert str(paths.hybrid_selection) in str(metrics.shell) + # No WGS metrics means no rehead job + assert result.qc_jobs == [lc, metrics] + assert result.terminal == {dedup} + + def test_insert_size_is_not_on_the_locuscollector_pass_with_a_bed( + self, tmp_path + ): + result = make_stage( + tmp_path, assay="WES", bed=tmp_path / "regions.bed" + ).add_to(DAG()) + + assert algo_order(str(result.jobs[0].shell)) == ( + ["LocusCollector"] + PRE_DEDUP_ALGOS + ) + + def test_without_a_bed(self, tmp_path): + result = make_stage(tmp_path, assay="WES").add_to(DAG()) + + lc, dedup = result.jobs + # InsertSize moves onto the LocusCollector pass + assert algo_order(str(lc.shell)) == ( + ["LocusCollector", "InsertSizeMetricAlgo"] + PRE_DEDUP_ALGOS + ) + assert result.qc_jobs == [lc] + assert result.terminal == {dedup} + + +class TestOptions: + """The flags a caller passes through""" + + def test_rmdup(self, tmp_path): + result = make_stage(tmp_path, duplicate_marking="rmdup").add_to(DAG()) + + assert "--rmdup" in str(result.dedup_job.shell) + + def test_consensus(self, tmp_path): + result = make_stage(tmp_path, consensus=True).add_to(DAG()) + + assert "--algo LocusCollector --consensus" in str(result.jobs[0].shell) + + def test_bam_format(self, tmp_path): + result = make_stage(tmp_path, bam_format=True).add_to(DAG()) + + assert result.deduped == [tmp_path / "output_deduped.bam"] + + def test_skip_metrics_keeps_duplicate_marking(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + result = make_stage(tmp_path, skip_metrics=True).add_to(DAG()) + + lc, dedup = result.jobs + assert algo_order(str(lc.shell)) == ["LocusCollector"] + # Dedup still writes its own metrics + assert f"--metrics {paths.dedup_metrics}" in str(dedup.shell) + assert result.qc_jobs == [lc] + assert result.terminal == {dedup} + + +class TestMetricsDirectory: + """The metrics directory is created before the jobs are built""" + + def test_created_when_not_a_dry_run(self, tmp_path): + ctx = StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=4, + dry_run=False, + skip_version_check=True, + ) + make_stage(tmp_path, ctx=ctx).add_to(DAG()) + + assert (tmp_path / "output_metrics").is_dir() + + def test_not_created_on_a_dry_run(self, tmp_path): + make_stage(tmp_path).add_to(DAG()) + + assert not (tmp_path / "output_metrics").exists() + + def test_created_even_when_no_jobs_are_added(self, tmp_path): + ctx = StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=4, + dry_run=False, + skip_version_check=True, + ) + make_stage( + tmp_path, + ctx=ctx, + duplicate_marking="none", + skip_metrics=True, + ).add_to(DAG()) + + assert (tmp_path / "output_metrics").is_dir() diff --git a/tests/unit/test_stage_segdup.py b/tests/unit/test_stage_segdup.py new file mode 100644 index 0000000..35c7106 --- /dev/null +++ b/tests/unit/test_stage_segdup.py @@ -0,0 +1,193 @@ +""" +Unit tests for the segdup-caller stage +""" + +import pathlib + +import packaging.version +import pytest + +from sentieon_cli import command_strings as cmds +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.segdup import ( + SEGDUP_MIN_VERSIONS, + SegdupStage, +) +from sentieon_cli.util import SampleSex + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> SegdupStage: + """A SegdupStage with the paths a caller supplies""" + defaults = dict( + ctx=make_ctx(tmp_path), + sr_alignment=tmp_path / "sample.cram", + sr_bundle=tmp_path / "bundle", + output_dir=tmp_path / "output_segdups", + input_vcf=tmp_path / "output.vcf.gz", + sample_sex=SampleSex.MALE, + ) + defaults.update(kwargs) + return SegdupStage(**defaults) # type: ignore[arg-type] + + +class TestMinVersions: + """The version constant lives next to the stage""" + + def test_set_argument_needs_0_7_0(self): + assert SEGDUP_MIN_VERSIONS["segdup-caller"] == ( + packaging.version.Version("0.7.0") + ) + + +class TestCommands: + """The `segdup-caller` commands the stage builds""" + + def test_short_read_command(self, tmp_path): + # The command dnascope-pangenome builds today + result = make_stage(tmp_path, genes="CFH,CYP2D6,SMN1").add_to(DAG()) + + assert str(result.job.shell) == ( + f"segdup-caller --short {tmp_path}/sample.cram " + f"--reference {tmp_path}/ref.fa " + f"--sr_model {tmp_path}/bundle " + f"--input_vcf {tmp_path}/output.vcf.gz " + "--sex male --genes CFH,CYP2D6,SMN1 " + f"--outdir {tmp_path}/output_segdups" + ) + + def test_long_read_command(self, tmp_path): + result = make_stage( + tmp_path, + lr_alignment=tmp_path / "long.cram", + lr_bundle=tmp_path / "bundle", + ).add_to(DAG()) + + assert str(result.job.shell) == ( + f"segdup-caller --short {tmp_path}/sample.cram " + f"--long {tmp_path}/long.cram " + f"--reference {tmp_path}/ref.fa " + f"--sr_model {tmp_path}/bundle " + f"--lr_model {tmp_path}/bundle " + f"--input_vcf {tmp_path}/output.vcf.gz " + "--sex male " + f"--outdir {tmp_path}/output_segdups" + ) + + def test_overrides_become_set_arguments(self, tmp_path): + result = make_stage( + tmp_path, + overrides=["main.min_map_qual=30", "main.other=1"], + ).add_to(DAG()) + + assert "--set main.min_map_qual=30 --set main.other=1 --outdir" in str( + result.job.shell + ) + + def test_without_an_input_vcf_or_genes(self, tmp_path): + result = make_stage(tmp_path, input_vcf=None).add_to(DAG()) + + shell = str(result.job.shell) + assert "--input_vcf" not in shell + assert "--genes" not in shell + assert "--set" not in shell + + @pytest.mark.parametrize( + "sample_sex,expected", + [ + (SampleSex.MALE, "male"), + (SampleSex.FEMALE, "female"), + (SampleSex.UNKNOWN, "female"), + (None, "female"), + ], + ) + def test_sex_argument(self, tmp_path, sample_sex, expected): + result = make_stage(tmp_path, sample_sex=sample_sex).add_to(DAG()) + + assert f"--sex {expected} " in str(result.job.shell) + + +class TestLongReadValidation: + """Long-read calling needs both the alignment and the model bundle""" + + def test_alignment_without_a_bundle(self, tmp_path): + with pytest.raises(ValueError): + cmds.cmd_segdup_caller( + tmp_path / "out", + tmp_path / "sample.cram", + reference=tmp_path / "ref.fa", + sr_bundle=tmp_path / "bundle", + lr_alignments=tmp_path / "long.cram", + ) + + def test_bundle_without_an_alignment(self, tmp_path): + with pytest.raises(ValueError): + cmds.cmd_segdup_caller( + tmp_path / "out", + tmp_path / "sample.cram", + reference=tmp_path / "ref.fa", + sr_bundle=tmp_path / "bundle", + lr_bundle=tmp_path / "bundle", + ) + + def test_the_stage_propagates_the_error(self, tmp_path): + with pytest.raises(ValueError): + make_stage(tmp_path, lr_bundle=tmp_path / "bundle").add_to(DAG()) + + +class TestJobMetadata: + """Job name, task name and thread count""" + + def test_default_names_and_threads(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.job.name == "segdup-caller" + assert result.job.task_name == "segdup" + assert result.job.threads == 4 + + def test_overridden_names(self, tmp_path): + result = make_stage( + tmp_path, name="segdups", task_name="segdups" + ).add_to(DAG()) + + assert result.job.name == "segdups" + assert result.job.task_name == "segdups" + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_becomes_the_dependencies(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.job] == {upstream} + + def test_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.job in dag.ready_jobs + + def test_result_fields(self, tmp_path): + output_dir = tmp_path / "output_segdups" + result = make_stage(tmp_path, output_dir=output_dir).add_to(DAG()) + + assert result.jobs == [result.job] + assert result.terminal == {result.job} + assert result.output_dir == output_dir diff --git a/tests/unit/test_stage_small_variants.py b/tests/unit/test_stage_small_variants.py new file mode 100644 index 0000000..56ae89f --- /dev/null +++ b/tests/unit/test_stage_small_variants.py @@ -0,0 +1,385 @@ +""" +Unit tests for the small-variant stages +""" + +import pathlib + +import pytest + +from sentieon_cli.dag import DAG +from sentieon_cli.driver import DNAscope, SVSolver +from sentieon_cli.shard import Shard +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.small_variants import ( + ApplySpec, + DNAscopeStage, + DriverStage, + GVCFtyperStage, + ModelApplyStage, + TransferApplyStage, + TransferSpec, +) +from sentieon_cli.stages.transfer import TransferConfig + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + # `build_transfer_jobs` shells out to bcftools unless this is set + dry_run=True, + skip_version_check=True, + ) + + +def make_call_stage(tmp_path: pathlib.Path, **kwargs) -> DNAscopeStage: + """A DNAscopeStage calling SNVs and indels""" + defaults = dict( + ctx=make_ctx(tmp_path), + algos=[DNAscope(tmp_path / "raw.vcf.gz")], + inputs=[tmp_path / "sample.cram"], + ) + defaults.update(kwargs) + return DNAscopeStage(**defaults) # type: ignore[arg-type] + + +def make_transfer_config(tmp_path: pathlib.Path) -> TransferConfig: + """One shard the pop VCF has, one it lacks""" + return TransferConfig( + pop_vcf=tmp_path / "pop.vcf.gz", + shards=[Shard("chr1", 0, 1000), Shard("chrUn_extra", 0, 500)], + pop_vcf_contigs={"chr1": 1000}, + fai_data={ + "chr1": {"length": 1000}, + "chrUn_extra": {"length": 500}, + }, + ) + + +class TestDriverStage: + """The generic single-job driver stage""" + + def test_minimal_command(self, tmp_path): + result = DriverStage( + ctx=make_ctx(tmp_path), + algos=[ + SVSolver( + output=tmp_path / "sv.vcf.gz", + vcf=tmp_path / "sv_tmp.vcf.gz", + ) + ], + name="svsolver", + task_name="sv-calling", + ).add_to(DAG()) + + assert str(result.job.shell) == ( + f"sentieon driver --reference {tmp_path}/ref.fa " + f"--thread_count 4 --algo SVSolver " + f"--vcf {tmp_path}/sv_tmp.vcf.gz {tmp_path}/sv.vcf.gz" + ) + + def test_name_task_name_and_default_threads(self, tmp_path): + result = DriverStage( + ctx=make_ctx(tmp_path), + algos=[DNAscope(tmp_path / "raw.vcf.gz")], + name="a-name", + task_name="a-task", + ).add_to(DAG()) + + assert result.job.name == "a-name" + assert result.job.task_name == "a-task" + assert result.job.threads == 4 + + def test_threads_override_leaves_the_driver_alone(self, tmp_path): + result = DriverStage( + ctx=make_ctx(tmp_path), + algos=[DNAscope(tmp_path / "raw.vcf.gz")], + name="a-name", + task_name="a-task", + threads=1, + ).add_to(DAG()) + + assert result.job.threads == 1 + assert "--thread_count 4" in str(result.job.shell) + + def test_result_fields(self, tmp_path): + result = make_call_stage(tmp_path).add_to(DAG()) + + assert result.jobs == [result.job] + assert result.terminal == {result.job} + + +class TestDNAscopeStage: + """The DNAscope calling pass""" + + def test_default_name_and_task_name(self, tmp_path): + result = make_call_stage(tmp_path).add_to(DAG()) + + assert result.job.name == "dnascope" + assert result.job.task_name == "variant-calling" + + def test_inputs_and_interval(self, tmp_path): + bed = tmp_path / "regions.bed" + result = make_call_stage( + tmp_path, + inputs=[tmp_path / "one.cram", tmp_path / "two.cram"], + interval=bed, + ).add_to(DAG()) + + shell = str(result.job.shell) + assert f"--input {tmp_path}/one.cram" in shell + assert f"--input {tmp_path}/two.cram" in shell + assert f"--interval {bed}" in shell + + def test_interval_padding_is_absent_unless_set(self, tmp_path): + result = make_call_stage(tmp_path).add_to(DAG()) + + assert "--interval_padding" not in str(result.job.shell) + + def test_interval_padding(self, tmp_path): + result = make_call_stage(tmp_path, interval_padding=100).add_to(DAG()) + + assert "--interval_padding 100" in str(result.job.shell) + + def test_read_filter(self, tmp_path): + result = make_call_stage( + tmp_path, read_filter=["UltimaReadFilter"] + ).add_to(DAG()) + + assert "--read_filter UltimaReadFilter" in str(result.job.shell) + + def test_empty_read_filter_emits_nothing(self, tmp_path): + result = make_call_stage(tmp_path, read_filter=[]).add_to(DAG()) + + assert "--read_filter" not in str(result.job.shell) + + def test_replace_rg_precedes_its_input(self, tmp_path): + result = make_call_stage( + tmp_path, + inputs=[tmp_path / "one.cram", tmp_path / "two.cram"], + replace_rg=[["rg1=rg1-lr"], ["rg2=rg2-lr"]], + ).add_to(DAG()) + + assert str(result.job.shell).startswith( + f"sentieon driver --replace_rg rg1=rg1-lr " + f"--input {tmp_path}/one.cram " + f"--replace_rg rg2=rg2-lr --input {tmp_path}/two.cram" + ) + + def test_algos_keep_their_order(self, tmp_path): + result = make_call_stage( + tmp_path, + algos=[ + DNAscope(tmp_path / "raw.vcf.gz"), + DNAscope(tmp_path / "sv_tmp.vcf.gz", var_type="BND"), + ], + ).add_to(DAG()) + + shell = str(result.job.shell) + assert shell.index("raw.vcf.gz") < shell.index("--var_type BND") + + def test_upstream_becomes_the_dependency(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_call_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.job] == {upstream} + + +class TestModelApplyStage: + """DNAModelApply""" + + def test_command(self, tmp_path): + result = ModelApplyStage( + ctx=make_ctx(tmp_path), + model=tmp_path / "bundle" / "dnascope.model", + vcf=tmp_path / "raw.vcf.gz", + output=tmp_path / "applied.vcf.gz", + ).add_to(DAG()) + + assert str(result.job.shell) == ( + f"sentieon driver --reference {tmp_path}/ref.fa " + f"--thread_count 4 --algo DNAModelApply " + f"--model {tmp_path}/bundle/dnascope.model " + f"--vcf {tmp_path}/raw.vcf.gz " + f"{tmp_path}/applied.vcf.gz" + ) + + def test_defaults_and_result(self, tmp_path): + output = tmp_path / "applied.vcf.gz" + result = ModelApplyStage( + ctx=make_ctx(tmp_path), + model=tmp_path / "dnascope.model", + vcf=tmp_path / "raw.vcf.gz", + output=output, + ).add_to(DAG()) + + assert result.job.name == "model-apply" + assert result.job.task_name == "model-apply" + assert result.job.threads == 4 + assert result.output == output + assert result.terminal == {result.job} + + def test_renamed_job(self, tmp_path): + result = ModelApplyStage( + ctx=make_ctx(tmp_path), + model=tmp_path / "dnascope.model", + vcf=tmp_path / "raw.vcf.gz", + output=tmp_path / "applied.vcf.gz", + name="snv-apply", + ).add_to(DAG()) + + assert result.job.name == "snv-apply" + + +class TestGVCFtyperStage: + """Genotyping a gVCF""" + + def test_command_without_an_interval(self, tmp_path): + result = GVCFtyperStage( + ctx=make_ctx(tmp_path), + gvcf=tmp_path / "sample.g.vcf.gz", + output=tmp_path / "sample.vcf.gz", + ).add_to(DAG()) + + assert str(result.job.shell) == ( + f"sentieon driver --reference {tmp_path}/ref.fa " + f"--thread_count 4 --algo GVCFtyper " + f"--vcf {tmp_path}/sample.g.vcf.gz " + f"{tmp_path}/sample.vcf.gz" + ) + assert result.job.name == "gvcftyper" + assert result.job.task_name == "gvcftyper" + + def test_interval(self, tmp_path): + bed = tmp_path / "regions.bed" + result = GVCFtyperStage( + ctx=make_ctx(tmp_path), + gvcf=tmp_path / "sample.g.vcf.gz", + output=tmp_path / "sample.vcf.gz", + interval=bed, + ).add_to(DAG()) + + assert f"--interval {bed}" in str(result.job.shell) + assert result.output == tmp_path / "sample.vcf.gz" + + +class TestTransferApplyStage: + """Transfer, apply, or both""" + + def test_neither_half_is_an_error(self, tmp_path): + with pytest.raises(ValueError): + TransferApplyStage( + ctx=make_ctx(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + ).add_to(DAG()) + + def test_apply_only_reads_the_raw_vcf(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = TransferApplyStage( + ctx=make_ctx(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + apply=ApplySpec( + model=tmp_path / "dnascope.model", + output=tmp_path / "applied.vcf.gz", + ), + ).add_to(dag, [upstream]) + + assert result.transfer is None + assert result.apply_job is not None + assert f"--vcf {tmp_path}/raw.vcf.gz" in str(result.apply_job.shell) + assert dag.waiting_jobs[result.apply_job] == {upstream} + assert result.terminal == {result.apply_job} + assert result.output == tmp_path / "applied.vcf.gz" + + def test_transfer_only(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = TransferApplyStage( + ctx=make_ctx(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + transfer=TransferSpec( + config=make_transfer_config(tmp_path), + out_vcf=tmp_path / "transfer.vcf.gz", + ), + ).add_to(dag, [upstream]) + + assert result.apply_job is None + assert result.transfer is not None + assert result.terminal == {result.transfer.concat_job} + assert result.output == tmp_path / "transfer.vcf.gz" + for job in result.transfer.shard_jobs: + assert dag.waiting_jobs[job] == {upstream} + + def test_both_halves(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = TransferApplyStage( + ctx=make_ctx(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + transfer=TransferSpec( + config=make_transfer_config(tmp_path), + out_vcf=tmp_path / "transfer.vcf.gz", + ), + apply=ApplySpec( + model=tmp_path / "dnascope.model", + output=tmp_path / "applied.vcf.gz", + ), + ).add_to(dag, [upstream]) + + assert result.transfer is not None + assert result.apply_job is not None + # DNAModelApply reads the transfer output + assert f"--vcf {tmp_path}/transfer.vcf.gz" in str( + result.apply_job.shell + ) + # ... and waits on the transfer alone; the concat already depends + # on the raw-VCF producer + assert dag.waiting_jobs[result.apply_job] == { + result.transfer.concat_job, + } + assert result.terminal == {result.apply_job} + assert result.jobs == [ + *result.transfer.jobs, + result.apply_job, + ] + + def test_transfer_tag_reaches_the_job_names(self, tmp_path): + result = TransferApplyStage( + ctx=make_ctx(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + transfer=TransferSpec( + config=make_transfer_config(tmp_path), + out_vcf=tmp_path / "transfer.vcf.gz", + tag="hp", + ), + ).add_to(DAG()) + + assert result.transfer is not None + assert result.transfer.concat_job.name == "merge-trim-hp-concat" + + def test_apply_name_override(self, tmp_path): + result = TransferApplyStage( + ctx=make_ctx(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + apply=ApplySpec( + model=tmp_path / "dnascope.model", + output=tmp_path / "applied.vcf.gz", + name="snv-apply", + ), + ).add_to(DAG()) + + assert result.apply_job is not None + assert result.apply_job.name == "snv-apply" diff --git a/tests/unit/test_stage_sv.py b/tests/unit/test_stage_sv.py new file mode 100644 index 0000000..388ce81 --- /dev/null +++ b/tests/unit/test_stage_sv.py @@ -0,0 +1,195 @@ +""" +Unit tests for the long-read SV calling stage +""" + +import pathlib + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.sv import LongReadSVStage + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> LongReadSVStage: + """A LongReadSVStage over one alignment""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.bam"], + model=tmp_path / "bundle" / "longreadsv.model", + ) + defaults.update(kwargs) + return LongReadSVStage(**defaults) # type: ignore[arg-type] + + +class TestCommand: + """The `sentieon driver` command the stage builds""" + + def test_longreadsv_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert str(result.job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.bam " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + f"--algo LongReadSV " + f"--model {tmp_path}/bundle/longreadsv.model " + f"{tmp_path}/output.sv.vcf.gz" + ) + + def test_interval(self, tmp_path): + bed = tmp_path / "regions.bed" + result = make_stage(tmp_path, interval=bed).add_to(DAG()) + + assert f"--interval {bed}" in str(result.job.shell) + + def test_without_an_interval(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert "--interval" not in str(result.job.shell) + + def test_multiple_inputs(self, tmp_path): + result = make_stage( + tmp_path, inputs=[tmp_path / "one.bam", tmp_path / "two.bam"] + ).add_to(DAG()) + + shell = str(result.job.shell) + assert f"--input {tmp_path}/one.bam" in shell + assert f"--input {tmp_path}/two.bam" in shell + + def test_replace_rg_interleaves_with_the_inputs(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.bam", tmp_path / "two.bam"], + replace_rg=[["rg1=@RG\\tID:rg1"], ["rg2=@RG\\tID:rg2"]], + ).add_to(DAG()) + + shell = str(result.job.shell) + assert ( + "--replace_rg 'rg1=@RG\\tID:rg1' " + f"--input {tmp_path}/one.bam " + "--replace_rg 'rg2=@RG\\tID:rg2' " + f"--input {tmp_path}/two.bam" + ) in shell + + def test_without_replace_rg(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert "--replace_rg" not in str(result.job.shell) + + +class TestAlgoOptions: + """The optional LongReadSV algo arguments""" + + def test_min_sv_size(self, tmp_path): + result = make_stage(tmp_path, min_sv_size=20).add_to(DAG()) + + assert str(result.job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.bam " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + f"--algo LongReadSV " + f"--model {tmp_path}/bundle/longreadsv.model " + f"--min_sv_size 20 " + f"{tmp_path}/output.sv.vcf.gz" + ) + + def test_min_map_qual(self, tmp_path): + result = make_stage(tmp_path, min_map_qual=30).add_to(DAG()) + + assert ( + f"--model {tmp_path}/bundle/longreadsv.model --min_map_qual 30 " + f"{tmp_path}/output.sv.vcf.gz" + ) in str(result.job.shell) + + def test_all_options(self, tmp_path): + result = make_stage( + tmp_path, + min_map_qual=30, + min_sv_size=20, + min_dp=5, + min_af=0.1, + ).add_to(DAG()) + + assert ( + "--min_map_qual 30 --min_sv_size 20 --min_dp 5 --min_af 0.1" + ) in str(result.job.shell) + + def test_options_default_to_absent(self, tmp_path): + shell = str(make_stage(tmp_path).add_to(DAG()).job.shell) + + for flag in ("min_map_qual", "min_sv_size", "min_dp", "min_af"): + assert f"--{flag}" not in shell + + +class TestOutput: + """The SV VCF is derived from the run's output VCF""" + + def test_sv_vcf(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.sv_vcf == tmp_path / "output.sv.vcf.gz" + assert str(result.sv_vcf) in str(result.job.shell) + + def test_an_explicit_output(self, tmp_path): + out_vcf = tmp_path / "tmp" / "sample-longread-sv.vcf.gz" + result = make_stage(tmp_path, output=out_vcf).add_to(DAG()) + + assert result.sv_vcf == out_vcf + assert str(result.job.shell).endswith(f" {out_vcf}") + assert "output.sv.vcf.gz" not in str(result.job.shell) + + +class TestJobMetadata: + """Job name, task name and thread count""" + + def test_defaults(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.job.name == "LongReadSV" + assert result.job.task_name == "sv-calling" + assert result.job.threads == 4 + + def test_custom_name_and_task(self, tmp_path): + result = make_stage( + tmp_path, name="longreadsv-lr", task_name="sv" + ).add_to(DAG()) + + assert result.job.name == "longreadsv-lr" + assert result.job.task_name == "sv" + + +class TestDagWiring: + """`build` touches no DAG; `add_to` inserts the job""" + + def test_build_inserts_nothing(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).build() + + assert result.jobs == [result.job] + assert result.terminal == {result.job} + assert not dag.ready_jobs + assert not dag.waiting_jobs + + def test_upstream_lands_on_the_job(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.job] == {upstream} + + def test_a_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.job in dag.ready_jobs diff --git a/tests/unit/test_stage_t1k.py b/tests/unit/test_stage_t1k.py new file mode 100644 index 0000000..b08b265 --- /dev/null +++ b/tests/unit/test_stage_t1k.py @@ -0,0 +1,182 @@ +""" +Unit tests for the T1K HLA/KIR genotyping stage +""" + +import pathlib + +from sentieon_cli.dag import DAG +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.t1k import ( + DEFAULT_T1K_HLA_LOCUS, + DEFAULT_T1K_KIR_LOCUS, + T1K_MIN_VERSIONS, + T1KStage, +) + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + dry_run=True, + skip_version_check=True, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> T1KStage: + """A T1KStage genotyping HLA, unless the caller says otherwise""" + defaults = dict( + ctx=make_ctx(tmp_path), + inputs=[tmp_path / "sample.cram"], + gene_seq=tmp_path / "hla_seq.fa", + gene_coord=tmp_path / "hla_coord.fa", + locus=DEFAULT_T1K_HLA_LOCUS, + preset="hla-wgs", + tag="hla", + output_dir=tmp_path / "output_hla", + ) + defaults.update(kwargs) + return T1KStage(**defaults) # type: ignore[arg-type] + + +def make_kir_stage(tmp_path: pathlib.Path, **kwargs) -> T1KStage: + """A T1KStage genotyping KIR""" + return make_stage( + tmp_path, + gene_seq=tmp_path / "kir_seq.fa", + gene_coord=tmp_path / "kir_coord.fa", + locus=DEFAULT_T1K_KIR_LOCUS, + preset="kir-wgs", + tag="kir", + output_dir=tmp_path / "output_kir", + **kwargs, + ) + + +class TestConstants: + """The version constant and the default loci live with the stage""" + + def test_no_minimum_version(self): + assert T1K_MIN_VERSIONS == {"run-t1k": None} + + def test_default_loci(self): + assert DEFAULT_T1K_HLA_LOCUS == "chr6:28510020-33480577" + assert DEFAULT_T1K_KIR_LOCUS == "chr19:53100000-55800000" + + +class TestCommands: + """The extraction and `run-t1k` commands the stage builds""" + + def test_hla_extract_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert str(result.extract_job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.cram " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + "--interval chr6:28510020-33480577 " + f"--algo ReadWriter {tmp_path}/sample_hla.bam" + ) + + def test_hla_t1k_command(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert str(result.t1k_job.shell) == ( + "run-t1k --abnormalUnmapFlag -t 4 --preset hla-wgs " + f"-f {tmp_path}/hla_seq.fa -c {tmp_path}/hla_coord.fa " + f"--od {tmp_path}/output_hla -b {tmp_path}/sample_hla.bam" + ) + + def test_kir_extract_command(self, tmp_path): + result = make_kir_stage(tmp_path).add_to(DAG()) + + assert str(result.extract_job.shell) == ( + f"sentieon driver --input {tmp_path}/sample.cram " + f"--reference {tmp_path}/ref.fa --thread_count 4 " + "--interval chr19:53100000-55800000 " + f"--algo ReadWriter {tmp_path}/sample_kir.bam" + ) + + def test_kir_t1k_command(self, tmp_path): + result = make_kir_stage(tmp_path).add_to(DAG()) + + assert str(result.t1k_job.shell) == ( + "run-t1k --abnormalUnmapFlag -t 4 --preset kir-wgs " + f"-f {tmp_path}/kir_seq.fa -c {tmp_path}/kir_coord.fa " + f"--od {tmp_path}/output_kir -b {tmp_path}/sample_kir.bam" + ) + + def test_multiple_inputs_are_merged_by_the_extract(self, tmp_path): + result = make_stage( + tmp_path, + inputs=[tmp_path / "one.cram", tmp_path / "two.cram"], + ).add_to(DAG()) + + shell = str(result.extract_job.shell) + assert f"--input {tmp_path}/one.cram" in shell + assert f"--input {tmp_path}/two.cram" in shell + + def test_a_custom_locus(self, tmp_path): + result = make_stage(tmp_path, locus="chr6:1-100").add_to(DAG()) + + assert "--interval chr6:1-100 " in str(result.extract_job.shell) + + +class TestJobMetadata: + """Job names, task name and thread counts""" + + def test_hla_job_names(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.extract_job.name == "t1k-hla-extract" + assert result.t1k_job.name == "t1k-hla" + + def test_kir_job_names(self, tmp_path): + result = make_kir_stage(tmp_path).add_to(DAG()) + + assert result.extract_job.name == "t1k-kir-extract" + assert result.t1k_job.name == "t1k-kir" + + def test_task_name_and_threads(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + for job in result.jobs: + assert job.task_name == "t1k" + assert job.threads == 4 + + def test_overridden_task_name(self, tmp_path): + result = make_stage(tmp_path, task_name="hla").add_to(DAG()) + + for job in result.jobs: + assert job.task_name == "hla" + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_upstream_lands_on_the_extract_only(self, tmp_path): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + assert dag.waiting_jobs[result.extract_job] == {upstream} + assert dag.waiting_jobs[result.t1k_job] == {result.extract_job} + + def test_the_extract_is_a_root_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + assert result.extract_job in dag.ready_jobs + assert dag.waiting_jobs[result.t1k_job] == {result.extract_job} + + def test_result_fields(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.jobs == [result.extract_job, result.t1k_job] + assert result.terminal == {result.t1k_job} + assert result.extracted_bam == tmp_path / "sample_hla.bam" + assert result.output_dir == tmp_path / "output_hla" diff --git a/tests/unit/test_stage_transfer.py b/tests/unit/test_stage_transfer.py new file mode 100644 index 0000000..d9c03f7 --- /dev/null +++ b/tests/unit/test_stage_transfer.py @@ -0,0 +1,169 @@ +""" +Unit tests for the annotation-transfer stage +""" + +import pathlib +from types import SimpleNamespace + +import pytest + +from sentieon_cli.dag import DAG +from sentieon_cli.shard import Shard +from sentieon_cli.stages.base import StageContext, rm_job +from sentieon_cli.stages.transfer import TransferConfig, TransferStage + + +def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: + """A StageContext over a temporary directory""" + return StageContext( + reference=tmp_path / "ref.fa", + output_vcf=tmp_path / "output.vcf.gz", + tmp_dir=tmp_path, + cores=cores, + # `build_transfer_jobs` shells out to bcftools unless this is set + dry_run=True, + skip_version_check=True, + ) + + +def make_config(tmp_path: pathlib.Path) -> TransferConfig: + """Two shards on a contig the pop VCF has, one on a contig it lacks""" + return TransferConfig( + pop_vcf=tmp_path / "pop.vcf.gz", + shards=[ + Shard("chr1", 0, 1000), + Shard("chr1", 1000, 2000), + Shard("chrUn_extra", 0, 500), + ], + pop_vcf_contigs={"chr1": 2000}, + fai_data={ + "chr1": {"length": 2000}, + "chrUn_extra": {"length": 500}, + }, + ) + + +def make_stage(tmp_path: pathlib.Path, **kwargs) -> TransferStage: + """A TransferStage over `make_config`""" + defaults = dict( + ctx=make_ctx(tmp_path), + config=make_config(tmp_path), + raw_vcf=tmp_path / "raw.vcf.gz", + out_vcf=tmp_path / "transfer.vcf.gz", + ) + defaults.update(kwargs) + return TransferStage(**defaults) # type: ignore[arg-type] + + +class TestJobs: + """The jobs the fan-out produces""" + + def test_one_job_per_shard_plus_the_extra_contig(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert [job.name for job in result.shard_jobs] == [ + "merge-trim-0", + "merge-trim-1", + "merge-trim-extra", + ] + assert result.concat_job.name == "merge-trim-concat" + + def test_task_name_and_threads(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + for job in result.jobs: + assert job.task_name == "annotation-transfer" + for job in result.shard_jobs: + assert job.threads == 1 + # The concat runs with every core + assert result.concat_job.threads == 4 + + def test_shard_jobs_read_the_raw_vcf(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + for job in result.shard_jobs: + assert str(tmp_path / "raw.vcf.gz") in str(job.shell) + + def test_concat_writes_the_output(self, tmp_path): + out_vcf = tmp_path / "annotated.vcf.gz" + result = make_stage(tmp_path, out_vcf=out_vcf).add_to(DAG()) + + assert str(out_vcf) in str(result.concat_job.shell) + assert result.out_vcf == out_vcf + + +class TestTag: + """`tag` disambiguates two fan-outs in one run""" + + def test_tagged_names(self, tmp_path): + result = make_stage(tmp_path, tag="hp").add_to(DAG()) + + assert [job.name for job in result.shard_jobs] == [ + "merge-trim-hp-0", + "merge-trim-hp-1", + "merge-trim-hp-extra", + ] + assert result.concat_job.name == "merge-trim-hp-concat" + + +class TestDagWiring: + """The edges the stage inserts""" + + def test_shards_depend_on_upstream_and_concat_on_the_shards( + self, tmp_path + ): + dag = DAG() + upstream = rm_job([tmp_path / "upstream"], "upstream") + dag.add_job(upstream) + + result = make_stage(tmp_path).add_to(dag, [upstream]) + + for job in result.shard_jobs: + assert dag.waiting_jobs[job] == {upstream} + assert dag.waiting_jobs[result.concat_job] == set(result.shard_jobs) + + def test_shards_are_roots_without_upstream(self, tmp_path): + dag = DAG() + result = make_stage(tmp_path).add_to(dag) + + for job in result.shard_jobs: + assert job in dag.ready_jobs + + def test_result_fields(self, tmp_path): + result = make_stage(tmp_path).add_to(DAG()) + + assert result.jobs == [*result.shard_jobs, result.concat_job] + assert result.terminal == {result.concat_job} + + +class TestFromPipeline: + """`TransferConfig.from_pipeline` collecting a pipeline's attributes""" + + @staticmethod + def make_pipeline(tmp_path: pathlib.Path, **kwargs) -> SimpleNamespace: + """A stand-in carrying the four `TransferInputs` attributes""" + config = make_config(tmp_path) + defaults = dict( + pop_vcf=config.pop_vcf, + shards=config.shards, + pop_vcf_contigs=config.pop_vcf_contigs, + fai_data=config.fai_data, + ) + defaults.update(kwargs) + return SimpleNamespace(**defaults) + + def test_collects_the_four_values(self, tmp_path): + pipeline = self.make_pipeline(tmp_path) + + config = TransferConfig.from_pipeline(pipeline) + + assert config.pop_vcf == pipeline.pop_vcf + assert config.shards == pipeline.shards + assert config.pop_vcf_contigs == pipeline.pop_vcf_contigs + assert config.fai_data == pipeline.fai_data + + def test_without_a_pop_vcf(self, tmp_path): + pipeline = self.make_pipeline(tmp_path, pop_vcf=None) + + with pytest.raises(ValueError, match="population VCF"): + TransferConfig.from_pipeline(pipeline) diff --git a/tests/unit/test_util_versions.py b/tests/unit/test_util_versions.py new file mode 100644 index 0000000..610bd7d --- /dev/null +++ b/tests/unit/test_util_versions.py @@ -0,0 +1,151 @@ +""" +Unit tests for the executable version helpers in `sentieon_cli.util` +""" + +import subprocess as sp +from unittest.mock import MagicMock + +import packaging.version + +from sentieon_cli import util + + +class TestExecutableVersion: + """Test `util.executable_version`""" + + def test_sentieon_driver_version(self, monkeypatch): + """The sentieon release is taken from the package name""" + check_output = MagicMock(return_value=b"sentieon-genomics-202503.04\n") + monkeypatch.setattr(util.sp, "check_output", check_output) + + assert util.executable_version( + "sentieon driver" + ) == packaging.version.Version("202503.04") + assert check_output.call_args[0][0] == [ + "sentieon", + "driver", + "--version", + ] + + def test_multiline_version(self, monkeypatch): + """Only the first line of a multi-line `--version` is parsed""" + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock( + return_value=( + b"bcftools 1.22\n" + b"Using htslib 1.22\n" + b"Copyright (C) 2024 Genome Research Ltd.\n" + ) + ), + ) + + assert util.executable_version( + "bcftools" + ) == packaging.version.Version("1.22") + + def test_called_process_error(self, monkeypatch): + """A failing `--version` call reports no version""" + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock(side_effect=sp.CalledProcessError(1, "bcftools")), + ) + + assert util.executable_version("bcftools") is None + + def test_os_error(self, monkeypatch): + """An un-runnable command reports no version""" + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock(side_effect=OSError("no such file")), + ) + + assert util.executable_version("bcftools") is None + + def test_unparseable_version(self, monkeypatch): + """Output that is not a version reports no version""" + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock(return_value=b"bcftools unknown-build\n"), + ) + + assert util.executable_version("bcftools") is None + + +class TestCheckVersion: + """Test `util.check_version`""" + + def test_missing_executable(self, monkeypatch): + """A command that is not in the PATH fails the check""" + monkeypatch.setattr(util.shutil, "which", MagicMock(return_value=None)) + check_output = MagicMock() + monkeypatch.setattr(util.sp, "check_output", check_output) + + assert ( + util.check_version("bcftools", packaging.version.Version("1.22")) + is False + ) + check_output.assert_not_called() + + def test_no_minimum_version_skips_probe(self, monkeypatch): + """Tools listed without a minimum version are not probed""" + monkeypatch.setattr( + util.shutil, "which", MagicMock(return_value="/usr/bin/vg") + ) + check_output = MagicMock() + monkeypatch.setattr(util.sp, "check_output", check_output) + + assert util.check_version("vg", None) is True + check_output.assert_not_called() + + def test_old_version(self, monkeypatch): + """An outdated executable fails the check""" + monkeypatch.setattr( + util.shutil, "which", MagicMock(return_value="/usr/bin/bcftools") + ) + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock(return_value=b"bcftools 1.16\n"), + ) + + assert ( + util.check_version("bcftools", packaging.version.Version("1.22")) + is False + ) + + def test_current_version(self, monkeypatch): + """An up-to-date executable passes the check""" + monkeypatch.setattr( + util.shutil, "which", MagicMock(return_value="/usr/bin/bcftools") + ) + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock(return_value=b"bcftools 1.22\n"), + ) + + assert ( + util.check_version("bcftools", packaging.version.Version("1.22")) + is True + ) + + def test_unparseable_version_fails(self, monkeypatch): + """Unparseable `--version` output fails the check""" + monkeypatch.setattr( + util.shutil, "which", MagicMock(return_value="/usr/bin/bcftools") + ) + monkeypatch.setattr( + util.sp, + "check_output", + MagicMock(return_value=b"bcftools unknown-build\n"), + ) + + assert ( + util.check_version("bcftools", packaging.version.Version("1.22")) + is False + ) diff --git a/tests/unit/test_vcf_header.py b/tests/unit/test_vcf_header.py new file mode 100644 index 0000000..859455a --- /dev/null +++ b/tests/unit/test_vcf_header.py @@ -0,0 +1,175 @@ +""" +Unit tests for the pure-Python VCF header reader +""" + +import gzip +import pathlib +import struct +import zlib + +import pytest + +from sentieon_cli.shard import parse_vcf_contigs, vcf_contigs +from sentieon_cli.util import ( + VcfHeaderError, + read_vcf_header, + vcf_id, + vcf_id_from_header, +) + +HEADER = """##fileformat=VCFv4.2 +##SentieonVcfID=population-test-20260101 +##contig= +##contig= +##contig= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO +""" + +BODY = "chr1\t1\t.\tA\tC\t50\tPASS\t.\n#not-a-header-line\n" + + +def bgzf(data: bytes) -> bytes: + """A single BGZF block, the way a `bgzip`-compressed VCF starts. + + BGZF is gzip with an extra `BC` subfield, so the magic bytes the + reader sniffs are the same and `gzip` decompresses it unchanged. + """ + deflater = zlib.compressobj(6, zlib.DEFLATED, -zlib.MAX_WBITS) + deflated = deflater.compress(data) + deflater.flush() + bsize = 18 + len(deflated) + 8 - 1 + header = struct.pack( + " pathlib.Path: + """A plain-text VCF, named `.vcf.gz` to defeat the extension""" + path = tmp_path / "plain.vcf.gz" + path.write_text(HEADER + BODY) + return path + + +@pytest.fixture +def gzip_vcf(tmp_path: pathlib.Path) -> pathlib.Path: + """A gzip-compressed VCF, named `.vcf` to defeat the extension""" + path = tmp_path / "compressed.vcf" + with gzip.open(path, "wt") as fh: + fh.write(HEADER + BODY) + return path + + +@pytest.fixture +def bgzf_vcf(tmp_path: pathlib.Path) -> pathlib.Path: + """A BGZF-compressed VCF, as `bgzip` would write it""" + path = tmp_path / "bgzf.vcf.gz" + path.write_bytes(bgzf((HEADER + BODY).encode())) + return path + + +class TestReadVcfHeader: + """`read_vcf_header` decompresses by magic bytes, not by name""" + + def test_a_plain_text_vcf_is_read(self, plain_vcf): + assert read_vcf_header(plain_vcf) == HEADER.splitlines() + + def test_a_gzip_vcf_is_read(self, gzip_vcf): + assert read_vcf_header(gzip_vcf) == HEADER.splitlines() + + def test_a_bgzf_vcf_is_read(self, bgzf_vcf): + assert read_vcf_header(bgzf_vcf) == HEADER.splitlines() + + def test_the_body_is_never_read(self, plain_vcf): + """Reading stops at `#CHROM`, so a `#` line below it is not a + header line""" + header = read_vcf_header(plain_vcf) + assert header[-1].startswith("#CHROM") + assert "#not-a-header-line" not in header + + def test_a_header_without_chrom_stops_at_the_body(self, tmp_path): + path = tmp_path / "no_chrom.vcf" + path.write_text("##fileformat=VCFv4.2\nchr1\t1\t.\tA\tC\n") + assert read_vcf_header(path) == ["##fileformat=VCFv4.2"] + + def test_a_file_with_no_header_raises(self, tmp_path): + path = tmp_path / "headerless.vcf" + path.write_text("chr1\t1\t.\tA\tC\t50\tPASS\t.\n") + with pytest.raises(VcfHeaderError): + read_vcf_header(path) + + def test_an_empty_file_raises(self, tmp_path): + path = tmp_path / "empty.vcf.gz" + path.touch() + with pytest.raises(VcfHeaderError): + read_vcf_header(path) + + def test_a_missing_file_raises_oserror(self, tmp_path): + with pytest.raises(OSError): + read_vcf_header(tmp_path / "absent.vcf.gz") + + def test_truncated_gzip_raises(self, tmp_path, gzip_vcf): + path = tmp_path / "truncated.vcf.gz" + path.write_bytes(gzip_vcf.read_bytes()[:20]) + with pytest.raises((VcfHeaderError, OSError)): + read_vcf_header(path) + + +class TestVcfContigs: + """`vcf_contigs` keeps its return type and its `None` lengths""" + + def test_contigs_and_lengths_are_parsed(self, bgzf_vcf): + assert vcf_contigs(bgzf_vcf) == { + "chr1": 248956422, + "chr2": 242193529, + "chrUn_KI270302v1": None, + } + + def test_a_contig_without_a_length_is_none(self): + contigs = parse_vcf_contigs(["##contig="]) + assert contigs == {"chr1": None} + + def test_a_quoted_assembly_does_not_confuse_the_parser(self): + contigs = parse_vcf_contigs( + ['##contig='] + ) + assert contigs == {"chr1": 10} + + def test_a_missing_file_returns_an_empty_dict(self, tmp_path): + assert vcf_contigs(tmp_path / "absent.vcf.gz") == {} + + def test_a_headerless_file_returns_an_empty_dict(self, tmp_path): + path = tmp_path / "headerless.vcf" + path.write_text("chr1\t1\t.\tA\tC\n") + assert vcf_contigs(path) == {} + + +class TestVcfId: + """`vcf_id` reports `None` the same way `vcf_contigs` reports `{}`""" + + def test_the_id_is_read(self, gzip_vcf): + assert vcf_id(gzip_vcf) == "population-test-20260101" + + def test_a_header_without_an_id(self, tmp_path): + path = tmp_path / "no_id.vcf" + path.write_text("##fileformat=VCFv4.2\n") + assert vcf_id(path) is None + + def test_a_missing_file_returns_none(self, tmp_path): + assert vcf_id(tmp_path / "absent.vcf.gz") is None + + def test_an_empty_file_returns_none(self, tmp_path): + path = tmp_path / "empty.vcf.gz" + path.touch() + assert vcf_id(path) is None + + def test_vcf_id_from_header(self): + assert vcf_id_from_header(HEADER.splitlines()) == ( + "population-test-20260101" + ) + assert vcf_id_from_header([]) is None + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py index cfaf676..af59da6 100644 --- a/tests/utils/test_helpers.py +++ b/tests/utils/test_helpers.py @@ -202,6 +202,10 @@ def create_hybrid_pipeline(self, **kwargs) -> DNAscopeHybridPipeline: _fai = self.fs.create_file("reference.fa.fai") pipeline.model_bundle = self.fs.create_model_bundle("model.bundle", "DNAscope Hybrid") + # CNV calling requires a PAR BED file, and the mock reference index + # is empty, so no packaged PAR BED file can be selected for it + pipeline.par_bed = self.fs.create_file("par.bed", b"chrX\t10000\t2781479\n") + # Create BWA index files (required when short-read fastq is supplied) for suf in (".amb", ".ann", ".bwt", ".pac", ".sa"): self.fs.create_file("reference.fa" + suf)