From 888ce6b0361734675c81968e846528d07f42fcc1 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 24 Jul 2026 09:30:35 -0700 Subject: [PATCH 01/58] Pass the `--pangenome_contig_prefix` to `PangenomeSV`. --- sentieon_cli/sentieon_pangenome.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 55616b8..6d73fcf 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -1226,6 +1226,7 @@ def build_dnascope_job( PangenomeSV( sv_vcf, gfa_file=gfa_file, + prefix=self.pangenome_contig_prefix, ) ) return Job( From 8d3d68e0b36f8a5ee1b01abad6e647cf45be48c6 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 12 Aug 2026 12:46:22 -0700 Subject: [PATCH 02/58] Fix shell-pipeline correctness bugs and size internal pipes natively Correctness fixes in the pipeline AST: - Command.__eq__ no longer mixes in self.proc (runtime state), so a command's identity is stable once its job runs. exec_kwargs (env, cwd) now participates in __eq__/__hash__ (canonicalised for order-independence) and renders into str() as a conventional bash prefix, so two jobs differing only in env/cwd no longer collide. - ProcSub.__eq__ requires the same subclass: <(x) != >(x). - Pipeline file/stream conflicts raise ValueError instead of calling sys.exit(), and the guards test `is not None` so fd 0 is not skipped. Fix a deadlock when a proc-sub FIFO is never opened: if the outer command failed to spawn or exited early, the background task's blocking open() never returned and Context.cleanup() hung the whole run, including the Ctrl-C path. cleanup() now unblocks still-unopened FIFOs (holding both ends non-blocking), lets the unblocked task return without spawning, and re-raises task exceptions only after all resources are released. Replace the injected `perl -MFcntl` pipe-buffer command with a native Pipeline(pipe_size=...) option. The injected command received the inherited stdout rather than the internal pipe it was meant to grow, so F_SETPIPE_SZ hit a non-pipe fd and the pipe kept its default 64 KiB size. The pipes are created with os.pipe() here, so they are now sized directly (best-effort; logged fallback on non-Linux or over fs.pipe-max-size) before the children inherit them. skip_pipe existed only for the injection and is removed; the bwa alignment pipelines in command_strings.py request the intended 256 MiB via pipe_size. Co-Authored-By: Claude Fable 5 --- sentieon_cli/command_strings.py | 34 ++-- sentieon_cli/shell_pipeline.py | 274 ++++++++++++++++++++--------- tests/unit/test_shell_pipeline.py | 275 ++++++++++++++++++++++++++++-- 3 files changed, 468 insertions(+), 115 deletions(-) diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 4bdebde..2595fe1 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -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, @@ -596,9 +601,6 @@ def cmd_samtools_fastq_bwa( if input_ref: ref_cmd = ["--reference", str(reference)] - pipebuf_cmd = Command( - "perl", "-MFcntl", "-e", "fcntl(STDOUT, 1031, 268435456)" - ) if collate: collate_cmd = Command( "samtools", @@ -667,20 +669,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, ) @@ -746,21 +746,10 @@ def cmd_fastq_bwa( 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 +769,6 @@ def cmd_fastq_bwa( ] if split: extract_cmd = Pipeline( - pipebuf_cmd, Command( "sentieon", "fqidx", @@ -792,7 +780,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 +803,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, ) diff --git a/sentieon_cli/shell_pipeline.py b/sentieon_cli/shell_pipeline.py index 02eb082..2469484 100644 --- a/sentieon_cli/shell_pipeline.py +++ b/sentieon_cli/shell_pipeline.py @@ -1,21 +1,77 @@ -"""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 dataclasses +import fcntl import io import os import pathlib 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, Union from .logging import get_logger logger = get_logger(__name__) +def _freeze(value: Any) -> Any: + """Return an order-independent, hashable view of ``value``. + + ``exec_kwargs`` (and a nested ``env`` dict) have no stable iteration + order, so hashing their raw form would let two *equal* commands hash + differently, breaking the ``a == b implies hash(a) == hash(b)`` + contract the DAG relies on. Canonicalise dicts/sets recursively. + """ + 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. ``F_SETPIPE_SZ`` is Linux-only and, for unprivileged + processes, capped by ``/proc/sys/fs/pipe-max-size`` (exceeding it + raises EPERM) -- both deployment concerns, so a failure is logged and + the pipe keeps its default size rather than aborting the run. A no-op + where the constant is unavailable (non-Linux), preserving POSIX + portability. + """ + 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.""" @@ -23,11 +79,15 @@ def __init__(self) -> 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] = [] + # 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 +98,48 @@ 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) 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 ShellNode(ABC): @@ -56,8 +153,8 @@ 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) @@ -68,9 +165,9 @@ def __init__( async def run( self, context: Context, - stdin: Union[IO, int, None] = None, - stdout: Union[IO, int, None] = None, - stderr: Union[IO, int, None] = None, + stdin: Union[IO[Any], int, None] = None, + stdout: Union[IO[Any], int, None] = None, + stderr: Union[IO[Any], int, None] = None, ) -> asyncio.subprocess.Process: # 1. Resolve Arguments (Handle Process Substitutions) final_args = [self.executable] @@ -96,27 +193,47 @@ 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: if not isinstance(other, Command): return NotImplemented + # Identity is the *definition* of the command: not self.proc, which + # is runtime state that would mutate identity once the job runs. return ( 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 +242,40 @@ 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, + stdin: Union[IO[Any], int, None] = None, + stdout: Union[IO[Any], int, None] = None, + stderr: Union[IO[Any], 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), + # We cannot have handles from both files and to run(). + # Test `is not None`, not truthiness: an int fd of 0 is a valid + # stream but falsy. + 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 +284,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 +307,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 +342,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 +351,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 +389,8 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: if not isinstance(other, ProcSub): return NotImplemented - return self.node == other.node + # An input and an output proc sub are not the same: <(x) != >(x). + return type(self) is type(other) and self.node == other.node def __hash__(self) -> int: return hash((self.__class__.__name__, self.node)) @@ -296,22 +404,26 @@ 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, @@ -321,10 +433,10 @@ async def run_inner(): 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 +450,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() 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/tests/unit/test_shell_pipeline.py b/tests/unit/test_shell_pipeline.py index 943d5b9..38ec2bd 100644 --- a/tests/unit/test_shell_pipeline.py +++ b/tests/unit/test_shell_pipeline.py @@ -3,6 +3,7 @@ """ import asyncio +import fcntl import os import pathlib import sys @@ -31,7 +32,6 @@ 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() @@ -73,7 +73,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 +81,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() - # Verify output file content with open(outfile_path, "r") as f: output = f.read().strip() @@ -132,18 +135,14 @@ 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) @@ -158,10 +157,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 +168,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.wait() + for sub in context.commands: + if sub.proc: + await sub.proc.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.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.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.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.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.wait() + + await asyncio.wait_for(context.cleanup(), timeout=10) + await asyncio.wait_for(context.cleanup(), timeout=10) + assert proc.returncode == 0 From a34554fec59930acd5ae29b7a6b12e03147030dd Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 12 Aug 2026 12:52:50 -0700 Subject: [PATCH 03/58] Rework the DAG/scheduler/executor engine layer Replace the undocumented send(None)/send(job) generator protocol between the DAG, scheduler, and executors with an explicit, typed interface: DAG.mark_finished(job) returns the jobs a completion unblocked, and BaseScheduler (new ABC) exposes start() / job_finished(job) with the contract documented and covered by a conformance-test harness. Make interrupt handling asyncio-correct and opt-in. The old handler ran time.sleep(10) and asyncio.run(context.cleanup()) inside a signal handler while the executor's event loop was running -- asyncio.run raises RuntimeError from within a running loop, and the blocking sleep froze it. Handlers (SIGINT/SIGTERM) are now installed on the running loop via loop.add_signal_handler only for the duration of a run (and only on the main thread), the previous handlers are restored afterwards, and shutdown is fully async: SIGTERM -> bounded grace period -> SIGKILL -> awaited cleanup. BasePipeline opts in, preserving Ctrl-C behavior. Executor correctness fixes: - A job that fails to launch is recorded in jobs_with_errors and stops new launches, instead of being dropped silently; the catch is narrowed to OSError so genuine bugs still propagate. - A later pipeline stage failing to spawn now terminates the stages that already started (e.g. `sleep 600 | no_such_cmd` no longer leaks sleep past the end of the run). - A proc-sub inner command that fails to spawn fails its own job instead of aborting the whole run mid-flight; teardown paths log (rather than re-raise) launch errors surfaced from cleanup so sibling contexts are still released. - A sub-command killed by SIGPIPE (downstream reader exited early) is no longer treated as a job failure, matching non-pipefail bash. Scheduler/DAG validation: - Jobs that can never be scheduled (threads over budget, or a managed resource request over capacity) raise DagExecutionError up front instead of stalling the run silently. - Adding a second job with an identical pipeline raises instead of silently collapsing onto the first. - DAG.add_job accepts any iterable of dependencies; a list or generator used to pass validation but was silently discarded, leaving the job immediately ready. Jobs may now declare inputs/outputs metadata: declared inputs are verified before a run and declared outputs after (StorageProvider seam, local presence checks), and a failed or interrupted job's declared outputs are deleted so a partial file cannot be mistaken for a good one. DAG.skip_satisfied plus the has_all_outputs predicate can prune already-satisfied jobs before a run, groundwork for a future resume option. The long-dead Job.run() method is removed. Co-Authored-By: Claude Fable 5 --- sentieon_cli/dag.py | 140 ++++-- sentieon_cli/executor.py | 475 +++++++++++++++---- sentieon_cli/job.py | 79 ++-- sentieon_cli/pipeline.py | 9 +- sentieon_cli/scheduler.py | 125 +++-- sentieon_cli/storage.py | 46 ++ tests/integration/test_executor.py | 636 ++++++++++++++++++++++++-- tests/integration/test_signals.py | 217 +++++++++ tests/unit/test_dag.py | 279 +++++++++++ tests/unit/test_dag_construction.py | 7 +- tests/unit/test_job.py | 52 +++ tests/unit/test_scheduler.py | 104 +++++ tests/unit/test_scheduler_contract.py | 109 +++++ tests/unit/test_storage.py | 49 ++ 14 files changed, 2070 insertions(+), 257 deletions(-) create mode 100644 sentieon_cli/storage.py create mode 100644 tests/integration/test_signals.py create mode 100644 tests/unit/test_dag.py create mode 100644 tests/unit/test_job.py create mode 100644 tests/unit/test_scheduler.py create mode 100644 tests/unit/test_scheduler_contract.py create mode 100644 tests/unit/test_storage.py diff --git a/sentieon_cli/dag.py b/sentieon_cli/dag.py index b3997b6..13c3343 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 Callable, Dict, Iterable, List, Optional, Set from .exceptions import DagExecutionError from .job import Job @@ -25,54 +25,100 @@ 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" - ) - - if isinstance(dependencies, set) and len(dependencies) > 0: - self.waiting_jobs[job] = dependencies.copy() - for dependency in dependencies: + def add_job( + self, job: Job, dependencies: Optional[Iterable[Job]] = None + ) -> None: + """Add a job to the DAG. + + ``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() - - 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) - - 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" - ) - - 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 + def mark_finished(self, job: Job) -> Dict[Job, None]: + """Record a finished job and return the jobs it unblocked. + + 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" + ) + + logger.debug("Marking finished: %s", job) + del self.ready_jobs[job] + self.finished_jobs.append(job) + + 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 + + def skip_satisfied(self, is_satisfied: Callable[[Job], bool]) -> List[Job]: + """Mark already-satisfied jobs finished before a run (resume). + + Walks the ready frontier: a job is skipped iff ``is_satisfied(job)`` + returns True and every one of its dependencies was itself skipped. + A job behind an unsatisfied ancestor is never offered to the + predicate -- it only becomes ready once all its dependencies are + finished, and during this walk "finished" means "skipped". + + Skipped jobs land in :attr:`finished_jobs`, so an executor never + sees them; unsatisfied jobs are left in place untouched. Call this + before handing the DAG to a scheduler. + + Returns the skipped jobs in topological order. + """ + skipped: List[Job] = [] + frontier: List[Job] = list(self.ready_jobs) + i = 0 + while i < len(frontier): + job = frontier[i] + i += 1 + if is_satisfied(job): + skipped.append(job) + frontier.extend(self.mark_finished(job)) + return skipped + + +def has_all_outputs(job: Job) -> bool: + """The default resume predicate for :meth:`DAG.skip_satisfied`. + + True iff ``job`` declares at least one output and every declared + output exists on the local filesystem. Jobs with no declared outputs + always return False -- they cannot be verified, so they must rerun. + Zero-byte files count as present; a dangling symlink does not + (``Path.exists()`` follows links) -- matching + :class:`~sentieon_cli.storage.LocalStorageProvider` semantics. + """ + return bool(job.outputs) and all(path.exists() for path in job.outputs) diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index 4d60f3e..ad0e39d 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -1,30 +1,71 @@ """Execute jobs""" import asyncio -import asyncio.subprocess +import contextlib 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 .scheduler import BaseScheduler from .shell_pipeline import Context +from .storage import LocalStorageProvider, StorageProvider logger = get_logger(__name__) -class BaseExecutor: +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 +76,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 +85,134 @@ 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). - def __init__(self, *args: Any, **kwargs: Any): - super().__init__(*args, **kwargs) + 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, + scheduler: BaseScheduler, + *, + install_signal_handlers: bool = False, + ) -> None: + super().__init__(scheduler) + self.install_signal_handlers = install_signal_handlers 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() + stop_event = asyncio.Event() + restore = self._install_signal_handlers(loop, stop_event) + try: + await self._drive(stop_event) + 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 async subprocesses.""" + + def __init__( + self, + scheduler: BaseScheduler, + *, + install_signal_handlers: bool = False, + shutdown_grace_period: float = 10.0, + storage: Optional[StorageProvider] = None, + ) -> None: + super().__init__( + scheduler, + install_signal_handlers=install_signal_handlers, + ) + self.shutdown_grace_period = shutdown_grace_period + self.storage: StorageProvider = ( + storage if storage is not None else LocalStorageProvider() + ) self.running: List[ Tuple[ Job, @@ -65,26 +221,6 @@ 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""" @@ -93,6 +229,8 @@ async def run_job(self, job: Job) -> None: context = Context() start_time = time.monotonic_ns() try: + for path in job.inputs: + await self.storage.stage_in(path) proc = await cmd.run( context, stderr=sys.stderr, @@ -105,85 +243,224 @@ async def run_job(self, job: Job) -> None: 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() + 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) + self._remove_outputs(job) - def execute(self) -> None: - """Execute jobs from the DAG""" - asyncio.run(self._execute()) + async def _terminate_context(self, context: Context) -> None: + """Tear down processes already spawned in an aborted context. - 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) + 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.wait()) for proc in live] + await asyncio.wait(waits, timeout=self.shutdown_grace_period) + _kill_survivors(context) + await asyncio.wait(waits) + + def _remove_outputs(self, job: Job) -> None: + """Delete a failed job's declared outputs. + + A partial or unverified output must not survive a failure: resume + (:meth:`~sentieon_cli.dag.DAG.skip_satisfied`) treats an existing + declared output as proof that its producer succeeded. Directories are + deliberately not removed (never recurse into user paths); a warning + is logged instead. + """ + for path in job.outputs: + try: + path.unlink() + except FileNotFoundError: + continue + except OSError as exc: + logger.warning( + "could not remove output %s of failed job %s: %s", + path, + job, + exc, + ) + else: + logger.info("Removed output of failed job %s: %s", job, path) + + 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}" - ) - cmd_failed = True - if ret == -9: + 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 + 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 + # 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( - "Sub-command received SIGKILL. " - "Possible out-of-memory?" + f"Return code {ret} for sub-command: " + f"{subcommand}" ) + cmd_failed = True + if ret == -9: + logger.error( + "Sub-command received SIGKILL. " + "Possible out-of-memory?" + ) + + if not cmd_failed: + for path in job.outputs: + try: + await self.storage.stage_out(path) + except OSError as exc: + logger.error( + "missing output %s for %s: %s", + path, + job.shell, + exc, + ) + cmd_failed = True + + # 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 cmd_failed: + logger.error("Command failure: %s", job.shell) + self.jobs_with_errors.append(job) + self.start_new_jobs = False + self._remove_outputs(job) + 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 - 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}" + 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) + + # Interrupted jobs never had their outputs verified; a partial file + # must not survive to fool a later resume. Jobs that completed in the + # same wake-up as the stop signal are conservatively included -- they + # were never verified either. + for job, _context, _task, _start in running: + self._remove_outputs(job) diff --git a/sentieon_cli/job.py b/sentieon_cli/job.py index dea8343..9d7d4d2 100644 --- a/sentieon_cli/job.py +++ b/sentieon_cli/job.py @@ -2,19 +2,29 @@ Job objects """ -import asyncio -import sys -import time -from typing import Dict, Optional +import pathlib +from typing import Dict, Iterable, List, Optional -from .logging import get_logger -from .shell_pipeline import Context, Pipeline - -logger = get_logger(__name__) +from .shell_pipeline import Pipeline 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). + * ``inputs`` / ``outputs`` -- local files the job reads / writes; when + declared, the executor verifies them around the run and removes the + outputs of a failed attempt. + * ``image`` -- container image (or environment) to run in on a backend. + + 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. + """ def __init__( self, @@ -22,53 +32,36 @@ def __init__( name: str, threads: int = 1, resources: Optional[Dict[str, int]] = None, - ): + *, + inputs: Optional[Iterable[pathlib.Path]] = None, + outputs: Optional[Iterable[pathlib.Path]] = None, + image: Optional[str] = None, + ) -> None: self.shell = pipeline self.name = name self.threads = threads self.resources = {} if resources is None else resources + self.inputs: List[pathlib.Path] = ( + [] if inputs is None else list(inputs) + ) + self.outputs: List[pathlib.Path] = ( + [] if outputs is None else list(outputs) + ) + self.image = image - def __hash__(self): + 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): + def __repr__(self) -> str: return f"Job({self.name})" - def __str__(self): + def __str__(self) -> str: 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}") diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 39672c1..ddb67e1 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -250,8 +250,13 @@ 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) self.logger.info("Starting execution") executor.execute() 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/storage.py b/sentieon_cli/storage.py new file mode 100644 index 0000000..5fb6263 --- /dev/null +++ b/sentieon_cli/storage.py @@ -0,0 +1,46 @@ +"""Storage staging for job inputs and outputs.""" + +import pathlib +from abc import ABC, abstractmethod + + +class StorageProvider(ABC): + """Stage a job's declared files in before it runs and out after. + + On a shared filesystem the files are already local, so staging + degenerates to verifying presence (see :class:`LocalStorageProvider`). + A cloud provider would instead pull inputs from and push outputs to + object storage, deriving each object's URI from a configured prefix and + the local path. The methods are async so they compose with the async + executors. + """ + + @abstractmethod + async def stage_in(self, local: pathlib.Path) -> None: + """Make ``local`` available before a job runs. + + Pull it from storage, or verify it is present. Raise if it cannot be + made available. + """ + + @abstractmethod + async def stage_out(self, local: pathlib.Path) -> None: + """Persist ``local`` after a job runs. + + Push it to storage, or verify it was produced. Raise if it is missing. + """ + + +class LocalStorageProvider(StorageProvider): + """Shared-filesystem provider: no data movement, just verify presence. + + A file counts as present when it exists, even if it is empty (zero bytes). + """ + + async def stage_in(self, local: pathlib.Path) -> None: + if not local.exists(): + raise FileNotFoundError(f"missing input: {local}") + + async def stage_out(self, local: pathlib.Path) -> None: + if not local.exists(): + raise FileNotFoundError(f"missing output: {local}") diff --git a/tests/integration/test_executor.py b/tests/integration/test_executor.py index 7199ecc..4f0a7f4 100644 --- a/tests/integration/test_executor.py +++ b/tests/integration/test_executor.py @@ -1,13 +1,13 @@ """ -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 pytest @@ -16,8 +16,13 @@ 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.dag import DAG, has_all_outputs # 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,64 +35,64 @@ 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", + ) 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) + pipeline = Pipeline( + Command("echo", "hello pipeline"), + Command("cat"), + file_output=cmd_out, + ) job = Job(pipeline, "pipeline-job") 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 + # This command will fail: the input file does not exist job = Job(Pipeline(Command("cat", str(cmd_in))), "failing-job") 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", @@ -100,18 +105,593 @@ def test_local_executor_proc_sub_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", + ) + 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") + b = Job(Pipeline(Command("echo", "b")), "b") + 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", + ) + 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")) + dag.add_job(Job(Pipeline(Command("echo", "2"), file_output=o2), "j2")) + + 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") + b = Job(Pipeline(Command("echo", "b"), file_output=second), "b") + 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") + 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}")) + 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)) + 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)) + 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")) + with pytest.raises(RuntimeError): + LocalExecutor(ThreadScheduler(dag, 1)).execute() + + +def test_missing_input_fails_before_run(tmp_path): + """A declared input that is missing fails the job before it runs.""" + missing = tmp_path / "in.txt" + marker = tmp_path / "ran.txt" + dag = DAG() + job = Job( + Pipeline(Command("touch", str(marker))), + "j", + inputs=[missing], + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + assert not marker.exists() # the command never ran + + +def test_present_inputs_and_outputs_allow_success(tmp_path): + """Empty (zero-byte) declared inputs/outputs still count as present.""" + inp = tmp_path / "in.txt" + inp.touch() # empty input file + out = tmp_path / "out.txt" + dag = DAG() + job = Job( + Pipeline(Command("touch", str(out))), # produces an empty output + "j", + inputs=[inp], + outputs=[out], + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert executor.jobs_with_errors == [] + assert out.exists() + + +def test_missing_output_fails_after_run(tmp_path): + """A declared output the command does not produce fails the job.""" + missing = tmp_path / "never.txt" + dag = DAG() + job = Job(Pipeline(Command("true")), "j", outputs=[missing]) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + + +def test_failed_job_outputs_are_removed(tmp_path): + """A failed job's declared outputs are deleted from disk.""" + out = tmp_path / "out.txt" + out.write_text("stale") + dag = DAG() + job = Job(Pipeline(Command("false")), "j", outputs=[out]) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + assert not out.exists() + + +def test_missing_output_failure_removes_produced_siblings(tmp_path): + """A missing declared output deletes the outputs that were produced.""" + produced = tmp_path / "out1.txt" + never = tmp_path / "out2.txt" + dag = DAG() + job = Job( + Pipeline(Command("touch", str(produced))), + "j", + outputs=[produced, never], + ) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + assert not produced.exists() + + +def test_launch_failure_removes_outputs(tmp_path): + """A job that fails to launch has its declared outputs deleted.""" + out = tmp_path / "out.txt" + out.write_text("stale") + dag = DAG() + job = Job(Pipeline(Command("no_such_cmd_zzz")), "j", outputs=[out]) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + assert not out.exists() + + +def test_missing_input_failure_removes_outputs(tmp_path): + """Any failed attempt clears the job's declared outputs -- even a + missing-input failure where the command never ran.""" + missing = tmp_path / "in.txt" + out = tmp_path / "out.txt" + out.write_text("stale") + dag = DAG() + job = Job(Pipeline(Command("true")), "j", inputs=[missing], outputs=[out]) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + assert not out.exists() + + +def test_successful_job_keeps_outputs(tmp_path): + """A successful job's outputs are never deleted.""" + out = tmp_path / "out.txt" + dag = DAG() + job = Job(Pipeline(Command("touch", str(out))), "j", outputs=[out]) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert executor.jobs_with_errors == [] + assert out.exists() + + +def test_failed_job_directory_output_is_left_in_place(tmp_path): + """Directories declared as outputs are never removed on failure.""" + out_dir = tmp_path / "out_dir" + out_dir.mkdir() + (out_dir / "keep.txt").write_text("keep") + dag = DAG() + job = Job(Pipeline(Command("false")), "j", outputs=[out_dir]) + dag.add_job(job) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert job in executor.jobs_with_errors + assert out_dir.is_dir() + assert (out_dir / "keep.txt").read_text() == "keep" + + +def test_local_executor_empty_dag_completes(): + """An empty DAG (e.g. everything skipped by resume) 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 test_resume_after_failure_reruns_only_failed_subtree(tmp_path): + """After a partial failure, skip_satisfied + has_all_outputs resumes + the run from the surviving intermediate files.""" + a_out = tmp_path / "a.txt" + b_out = tmp_path / "b.txt" + + def build_dag(b_cmd): + dag = DAG() + a = Job( + Pipeline(Command("echo", "a"), file_output=a_out), + "a", + outputs=[a_out], + ) + b = Job(Pipeline(b_cmd), "b", outputs=[b_out]) + dag.add_job(a) + dag.add_job(b, {a}) + return dag, a, b + + # First run: a succeeds, b fails (and b's partial output is deleted). + b_out.write_text("partial") + dag, _a, b = build_dag(Command("false")) + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert b in executor.jobs_with_errors + assert a_out.exists() + assert not b_out.exists() + + # Second run: rebuild the DAG with b fixed; resume skips a. + dag, a, b = build_dag(Command("touch", str(b_out))) + skipped = dag.skip_satisfied(has_all_outputs) + assert skipped == [a] + executor = LocalExecutor(ThreadScheduler(dag, 1)) + executor.execute() + assert executor.jobs_with_errors == [] + assert b_out.exists() + assert a_out.read_text().strip() == "a" # a was not rerun over + + +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", + ) + 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", + ) + 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", + ) + 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", + ) + good = Job(Pipeline(Command("true")), "sibling") + 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", + ) + 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_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", + ) + 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") + 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._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"), + bad, + asyncio.create_task(_done()), + 0, + ), + ( + Job(Pipeline(Command("false")), "j2"), + good, + asyncio.create_task(_done()), + 0, + ), + ] + await executor._shutdown() # must not raise + + assert bad.cleaned and good.cleaned diff --git a/tests/integration/test_signals.py b/tests/integration/test_signals.py new file mode 100644 index 0000000..4aeae88 --- /dev/null +++ b/tests/integration/test_signals.py @@ -0,0 +1,217 @@ +""" +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") + 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")) + 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_interrupted_job_outputs_are_removed(tmp_path): + """A job killed by an interrupt has its declared outputs deleted.""" + out = tmp_path / "out.txt" + out.write_text("partial") + dag = DAG() + job = Job(Pipeline(Command("sleep", "30")), "sleeper", outputs=[out]) + dag.add_job(job) + 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: + executor.execute() + finally: + timer.cancel() + + assert executor.running == [] + assert not out.exists() + + +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") + 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", + ) + 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/test_dag.py b/tests/unit/test_dag.py new file mode 100644 index 0000000..444d07e --- /dev/null +++ b/tests/unit/test_dag.py @@ -0,0 +1,279 @@ +""" +Unit tests for DAG traversal, duplicate detection, and the resume prune. +""" + +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, has_all_outputs # 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) + + +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") + j2 = Job(Pipeline(Command("echo", "x")), "name-b") + j3 = Job(Pipeline(Command("echo", "y")), "name-c") + 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")) + # A second job with an identical pipeline collides on identity. + with pytest.raises(ValueError): + dag.add_job(Job(Pipeline(Command("echo", "x")), "step-2")) + + 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") + j2 = Job(Pipeline(Command("run", exec_kwargs={"cwd": "/b"})), "b") + 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} + + +class TestSkipSatisfied: + """Test the pre-run resume prune""" + + def test_satisfied_root_is_marked_finished(self): + dag = DAG() + job = _job("a") + dag.add_job(job) + skipped = dag.skip_satisfied(lambda j: True) + assert skipped == [job] + assert job in dag.finished_jobs + assert len(dag.ready_jobs) == 0 + + def test_unsatisfied_root_stays_ready(self): + dag = DAG() + job = _job("a") + dag.add_job(job) + skipped = dag.skip_satisfied(lambda j: False) + assert skipped == [] + assert job in dag.ready_jobs + assert dag.finished_jobs == [] + + def test_cascades_down_a_chain_in_topological_order(self): + dag = DAG() + a, b, c = _job("a"), _job("b"), _job("c") + dag.add_job(a) + dag.add_job(b, {a}) + dag.add_job(c, {b}) + skipped = dag.skip_satisfied(lambda j: True) + assert skipped == [a, b, c] + assert dag.finished_jobs == [a, b, c] + assert len(dag.ready_jobs) == 0 + assert len(dag.waiting_jobs) == 0 + + def test_stops_at_the_first_unsatisfied_job(self): + dag = DAG() + a, b, c = _job("a"), _job("b"), _job("c") + dag.add_job(a) + dag.add_job(b, {a}) + dag.add_job(c, {b}) + skipped = dag.skip_satisfied(lambda j: j == a) + assert skipped == [a] + assert b in dag.ready_jobs + assert c in dag.waiting_jobs + + def test_blocked_jobs_are_not_offered_to_the_predicate(self): + dag = DAG() + a, b, c = _job("a"), _job("b"), _job("c") + dag.add_job(a) + dag.add_job(b, {a}) + dag.add_job(c, {b}) + seen = [] + + def predicate(job): + seen.append(job) + return job == a + + dag.skip_satisfied(predicate) + # b is offered (a was skipped, unblocking it); c never is -- the + # ancestor rule: a job behind an unsatisfied job cannot be skipped. + assert seen == [a, b] + + def test_diamond_with_partial_parents(self): + dag = DAG() + top = _job("top") + left, right = _job("left"), _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}) + skipped = dag.skip_satisfied(lambda j: j in (top, left)) + assert skipped == [top, left] + assert right in dag.ready_jobs + assert dag.waiting_jobs[bottom] == {right} + + def test_empty_dag_returns_empty(self): + assert DAG().skip_satisfied(lambda j: True) == [] + + +class TestHasAllOutputs: + """Test the default resume predicate""" + + def test_no_declared_outputs_is_never_satisfied(self): + assert has_all_outputs(_job("a")) is False + + def test_true_when_all_outputs_exist(self, tmp_path): + full = tmp_path / "full.txt" + full.write_text("data") + empty = tmp_path / "empty.txt" + empty.touch() # zero-byte files count as present + job = Job(Pipeline(Command("true")), "j", outputs=[full, empty]) + assert has_all_outputs(job) is True + + def test_false_when_any_output_is_missing(self, tmp_path): + present = tmp_path / "present.txt" + present.touch() + missing = tmp_path / "missing.txt" + job = Job(Pipeline(Command("true")), "j", outputs=[present, missing]) + assert has_all_outputs(job) is False + + def test_dangling_symlink_counts_as_missing(self, tmp_path): + target = tmp_path / "target.txt" + target.touch() + link = tmp_path / "link.txt" + link.symlink_to(target) + target.unlink() # now dangling + job = Job(Pipeline(Command("true")), "j", outputs=[link]) + assert has_all_outputs(job) is False diff --git a/tests/unit/test_dag_construction.py b/tests/unit/test_dag_construction.py index c42159f..e9bacdf 100644 --- a/tests/unit/test_dag_construction.py +++ b/tests/unit/test_dag_construction.py @@ -426,8 +426,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 +434,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_job.py b/tests/unit/test_job.py new file mode 100644 index 0000000..4f46942 --- /dev/null +++ b/tests/unit/test_job.py @@ -0,0 +1,52 @@ +""" +Unit tests for the Job model. +""" + +import os +import pathlib +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.job import Job # noqa: E402 +from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 + + +def _pipe(text): + return Pipeline(Command("echo", text)) + + +def test_new_fields_default_empty(): + job = Job(_pipe("a"), "a") + assert job.inputs == [] + assert job.outputs == [] + assert job.image is None + + +def test_new_fields_are_stored(): + ins = [pathlib.Path("in1"), pathlib.Path("in2")] + outs = [pathlib.Path("out")] + job = Job(_pipe("a"), "a", inputs=ins, outputs=outs, image="org/img:1") + assert job.inputs == ins + assert job.outputs == outs + assert job.image == "org/img:1" + + +def test_input_output_sequences_are_snapshotted(): + ins = [pathlib.Path("in1")] + job = Job(_pipe("a"), "a", inputs=ins) + ins.append(pathlib.Path("in2")) # mutating the caller's list... + assert job.inputs == [pathlib.Path("in1")] # ...doesn't change the job + + +def test_identity_ignores_metadata_fields(): + # Identity is keyed on the shell pipeline; the other fields don't + # affect it. + j1 = Job(_pipe("same"), "j1", inputs=[pathlib.Path("x")], image="a") + j2 = Job(_pipe("same"), "j2", outputs=[pathlib.Path("y")], image="b") + assert j1 == j2 + assert hash(j1) == hash(j2) diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py new file mode 100644 index 0000000..aaecf0e --- /dev/null +++ b/tests/unit/test_scheduler.py @@ -0,0 +1,104 @@ +""" +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) + + +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)) + 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}) + 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..0460403 --- /dev/null +++ b/tests/unit/test_scheduler_contract.py @@ -0,0 +1,109 @@ +""" +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) + + +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_storage.py b/tests/unit/test_storage.py new file mode 100644 index 0000000..800623f --- /dev/null +++ b/tests/unit/test_storage.py @@ -0,0 +1,49 @@ +""" +Unit tests for the storage providers. +""" + +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.storage import LocalStorageProvider # noqa: E402 + + +@pytest.mark.asyncio +async def test_stage_in_present_file_ok(tmp_path): + f = tmp_path / "in.txt" + f.write_text("data") + await LocalStorageProvider().stage_in(f) # no raise + + +@pytest.mark.asyncio +async def test_stage_in_empty_file_ok(tmp_path): + f = tmp_path / "in.txt" + f.touch() # a zero-byte file still counts as present + await LocalStorageProvider().stage_in(f) + + +@pytest.mark.asyncio +async def test_stage_in_missing_file_raises(tmp_path): + with pytest.raises(FileNotFoundError): + await LocalStorageProvider().stage_in(tmp_path / "nope.txt") + + +@pytest.mark.asyncio +async def test_stage_out_present_file_ok(tmp_path): + f = tmp_path / "out.txt" + f.touch() + await LocalStorageProvider().stage_out(f) + + +@pytest.mark.asyncio +async def test_stage_out_missing_file_raises(tmp_path): + with pytest.raises(FileNotFoundError): + await LocalStorageProvider().stage_out(tmp_path / "nope.txt") From faf6c9a2c2a341c5a70db138fe6fcb770dd46421 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 12 Aug 2026 12:55:24 -0700 Subject: [PATCH 04/58] Fix BasePipeline lifecycle issues - main() created the temp dir and only removed it after build_dag() and run() returned, so any raise in between (an infeasible DAG, a launch failure, a KeyboardInterrupt) leaked the directory. Wrap build/run/check in try/finally so the temp dir is always removed, and honor retain_tmpdir on failure too so a failed run can be inspected. - check_execution raised a bare ValueError("Execution failed"); raise DagExecutionError naming the failed jobs instead, to aid post-mortems. The unexecuted-jobs check uses DagExecutionError as well. - setup_logging set the level on the logger's parent, which resolved to the root logger -- so -v/--debug turned on DEBUG for every third-party library. Add logging.set_level(), which scopes the level to the sentieon_cli package logger; the module loggers are at NOTSET and inherit from it, and the root logger is left untouched. - add_arguments mutated the shared class-level params specs (deleting "flags" and injecting inferred "type" entries), so registering a pipeline class on a second parser lost its short flags (e.g. -r, -t). Build a fresh kwargs dict per argument instead. Co-Authored-By: Claude Fable 5 --- sentieon_cli/logging.py | 14 ++- sentieon_cli/pipeline.py | 47 ++++---- tests/unit/test_pipeline_lifecycle.py | 147 ++++++++++++++++++++++++++ 3 files changed, 187 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_pipeline_lifecycle.py diff --git a/sentieon_cli/logging.py b/sentieon_cli/logging.py index bd5d15a..128485d 100644 --- a/sentieon_cli/logging.py +++ b/sentieon_cli/logging.py @@ -4,9 +4,21 @@ handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s")) -def get_logger(name: str): +def get_logger(name: str) -> logging.Logger: """Return a logger with a StreamHandler.""" logger = logging.getLogger(name) logger.addHandler(handler) logger.propagate = False return logger + + +def set_level(level: int | str) -> None: + """Set the level of the package logger. + + The package's module loggers (``sentieon_cli.dag``, + ``sentieon_cli.executor``, ...) are all at NOTSET and inherit their + effective level from the ``sentieon_cli`` package logger, so setting it + here scopes verbosity to this package without touching the root logger + (and thus third-party libraries). + """ + logging.getLogger(__package__).setLevel(level) diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index ddb67e1..1433db2 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -15,9 +15,10 @@ 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_level from .scheduler import ThreadScheduler from .util import __version__, check_version, path_arg, tmp @@ -71,17 +72,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""" @@ -106,9 +107,10 @@ def handle_arguments(self, args: argparse.Namespace): 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) + # Scope verbosity to this package's loggers rather than the root + # logger, so -v/--debug does not turn on DEBUG for every + # third-party library. + set_level(args.loglevel) self.logger.info("Starting sentieon-cli version: %s", __version__) def __init__(self) -> None: @@ -130,13 +132,17 @@ def main(self, args: argparse.Namespace) -> None: 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) + # Always clean up the temp dir, even if build_dag()/run() raise (an + # infeasible DAG, a launch failure, or a KeyboardInterrupt) -- + # otherwise it would leak. retain_tmpdir is honored on failure too, + # so a failed run can still be inspected. + try: + dag = self.build_dag() + executor = self.run(dag) + self.check_execution(dag, executor) + finally: + if not self.retain_tmpdir: + shutil.rmtree(tmp_dir_str) def check_execution( self, @@ -145,10 +151,11 @@ 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) + raise DagExecutionError(f"Execution failed for jobs: {failed}") 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" diff --git a/tests/unit/test_pipeline_lifecycle.py b/tests/unit/test_pipeline_lifecycle.py new file mode 100644 index 0000000..b9659c1 --- /dev/null +++ b/tests/unit/test_pipeline_lifecycle.py @@ -0,0 +1,147 @@ +""" +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 + +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.pipeline import BasePipeline # 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 [] + + +def test_check_execution_names_the_failed_jobs(): + pipeline = _DummyPipeline() + pipeline.setup_logging(argparse.Namespace(loglevel="WARNING")) + job = Job(Pipeline(Command("false")), "broken-step") + + with pytest.raises(DagExecutionError, match="broken-step"): + pipeline.check_execution(DAG(), _StubExecutor([job])) + + +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")) + + with pytest.raises(DagExecutionError, match="unexecuted"): + pipeline.check_execution(dag, _StubExecutor()) From ffd4b603b325496306b9ac076bc4cc078124885b Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 13 Aug 2026 00:34:01 -0700 Subject: [PATCH 05/58] Trim unused engine functionality and condense comments Remove engine capabilities that nothing in the CLI uses: - DAG.skip_satisfied() and the has_all_outputs predicate; no pipeline offers a resume option, so the pre-run prune had no callers. - The storage module (StorageProvider/LocalStorageProvider) and the executor's input/output staging built on it, along with LocalExecutor._remove_outputs(); pipelines do not declare per-job files, so the staging and failed-output removal were dead paths. - Job.inputs/outputs/image; Job is back to (pipeline, name, threads, resources). Drop the tests that covered the removed functionality and condense some overly verbose code comments. Co-Authored-By: Claude Fable 5 --- sentieon_cli/dag.py | 41 +------ sentieon_cli/executor.py | 55 +-------- sentieon_cli/job.py | 18 +-- sentieon_cli/logging.py | 9 +- sentieon_cli/pipeline.py | 7 -- sentieon_cli/shell_pipeline.py | 20 +--- sentieon_cli/storage.py | 46 -------- tests/integration/test_executor.py | 174 +---------------------------- tests/integration/test_signals.py | 24 ---- tests/unit/test_dag.py | 114 +------------------ tests/unit/test_job.py | 52 --------- tests/unit/test_storage.py | 49 -------- 12 files changed, 10 insertions(+), 599 deletions(-) delete mode 100644 sentieon_cli/storage.py delete mode 100644 tests/unit/test_job.py delete mode 100644 tests/unit/test_storage.py diff --git a/sentieon_cli/dag.py b/sentieon_cli/dag.py index 13c3343..d40419d 100644 --- a/sentieon_cli/dag.py +++ b/sentieon_cli/dag.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import Callable, Dict, Iterable, List, Optional, Set +from typing import Dict, Iterable, List, Optional, Set from .exceptions import DagExecutionError from .job import Job @@ -83,42 +83,3 @@ def mark_finished(self, job: Job) -> Dict[Job, None]: new_ready_jobs[dependency] = None del self.waiting_jobs[dependency] return new_ready_jobs - - def skip_satisfied(self, is_satisfied: Callable[[Job], bool]) -> List[Job]: - """Mark already-satisfied jobs finished before a run (resume). - - Walks the ready frontier: a job is skipped iff ``is_satisfied(job)`` - returns True and every one of its dependencies was itself skipped. - A job behind an unsatisfied ancestor is never offered to the - predicate -- it only becomes ready once all its dependencies are - finished, and during this walk "finished" means "skipped". - - Skipped jobs land in :attr:`finished_jobs`, so an executor never - sees them; unsatisfied jobs are left in place untouched. Call this - before handing the DAG to a scheduler. - - Returns the skipped jobs in topological order. - """ - skipped: List[Job] = [] - frontier: List[Job] = list(self.ready_jobs) - i = 0 - while i < len(frontier): - job = frontier[i] - i += 1 - if is_satisfied(job): - skipped.append(job) - frontier.extend(self.mark_finished(job)) - return skipped - - -def has_all_outputs(job: Job) -> bool: - """The default resume predicate for :meth:`DAG.skip_satisfied`. - - True iff ``job`` declares at least one output and every declared - output exists on the local filesystem. Jobs with no declared outputs - always return False -- they cannot be verified, so they must rerun. - Zero-byte files count as present; a dangling symlink does not - (``Path.exists()`` follows links) -- matching - :class:`~sentieon_cli.storage.LocalStorageProvider` semantics. - """ - return bool(job.outputs) and all(path.exists() for path in job.outputs) diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index ad0e39d..66b1aaf 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -7,13 +7,12 @@ import threading import time from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Tuple from .job import Job from .logging import get_logger from .scheduler import BaseScheduler from .shell_pipeline import Context -from .storage import LocalStorageProvider, StorageProvider logger = get_logger(__name__) @@ -203,16 +202,12 @@ def __init__( *, install_signal_handlers: bool = False, shutdown_grace_period: float = 10.0, - storage: Optional[StorageProvider] = None, ) -> None: super().__init__( scheduler, install_signal_handlers=install_signal_handlers, ) self.shutdown_grace_period = shutdown_grace_period - self.storage: StorageProvider = ( - storage if storage is not None else LocalStorageProvider() - ) self.running: List[ Tuple[ Job, @@ -229,8 +224,6 @@ async def run_job(self, job: Job) -> None: context = Context() start_time = time.monotonic_ns() try: - for path in job.inputs: - await self.storage.stage_in(path) proc = await cmd.run( context, stderr=sys.stderr, @@ -258,7 +251,6 @@ async def run_job(self, job: Job) -> None: # race): the job is already recorded, so swallow it. await self._terminate_context(context) await _cleanup_quietly(context) - self._remove_outputs(job) async def _terminate_context(self, context: Context) -> None: """Tear down processes already spawned in an aborted context. @@ -282,30 +274,6 @@ async def _terminate_context(self, context: Context) -> None: _kill_survivors(context) await asyncio.wait(waits) - def _remove_outputs(self, job: Job) -> None: - """Delete a failed job's declared outputs. - - A partial or unverified output must not survive a failure: resume - (:meth:`~sentieon_cli.dag.DAG.skip_satisfied`) treats an existing - declared output as proof that its producer succeeded. Directories are - deliberately not removed (never recurse into user paths); a warning - is logged instead. - """ - for path in job.outputs: - try: - path.unlink() - except FileNotFoundError: - continue - except OSError as exc: - logger.warning( - "could not remove output %s of failed job %s: %s", - path, - job, - exc, - ) - else: - logger.info("Removed output of failed job %s: %s", job, path) - 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() @@ -375,19 +343,6 @@ async def _drive(self, stop_event: asyncio.Event) -> None: "Possible out-of-memory?" ) - if not cmd_failed: - for path in job.outputs: - try: - await self.storage.stage_out(path) - except OSError as exc: - logger.error( - "missing output %s for %s: %s", - path, - job.shell, - exc, - ) - cmd_failed = True - # 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 @@ -404,7 +359,6 @@ async def _drive(self, stop_event: asyncio.Event) -> None: logger.error("Command failure: %s", job.shell) self.jobs_with_errors.append(job) self.start_new_jobs = False - self._remove_outputs(job) else: logger.info( f"Finished command in " @@ -457,10 +411,3 @@ async def _shutdown(self) -> None: # tear down after an interrupt; log it rather than let it escape, # and keep releasing the remaining contexts. await _cleanup_quietly(context) - - # Interrupted jobs never had their outputs verified; a partial file - # must not survive to fool a later resume. Jobs that completed in the - # same wake-up as the stop signal are conservatively included -- they - # were never verified either. - for job, _context, _task, _start in running: - self._remove_outputs(job) diff --git a/sentieon_cli/job.py b/sentieon_cli/job.py index 9d7d4d2..3e36fa2 100644 --- a/sentieon_cli/job.py +++ b/sentieon_cli/job.py @@ -2,8 +2,7 @@ Job objects """ -import pathlib -from typing import Dict, Iterable, List, Optional +from typing import Dict, Optional from .shell_pipeline import Pipeline @@ -17,10 +16,6 @@ class Job: * ``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). - * ``inputs`` / ``outputs`` -- local files the job reads / writes; when - declared, the executor verifies them around the run and removes the - outputs of a failed attempt. - * ``image`` -- container image (or environment) to run in on a backend. 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. @@ -32,22 +27,11 @@ def __init__( name: str, threads: int = 1, resources: Optional[Dict[str, int]] = None, - *, - inputs: Optional[Iterable[pathlib.Path]] = None, - outputs: Optional[Iterable[pathlib.Path]] = None, - image: Optional[str] = None, ) -> None: self.shell = pipeline self.name = name self.threads = threads self.resources = {} if resources is None else resources - self.inputs: List[pathlib.Path] = ( - [] if inputs is None else list(inputs) - ) - self.outputs: List[pathlib.Path] = ( - [] if outputs is None else list(outputs) - ) - self.image = image def __hash__(self) -> int: return hash(self.shell) diff --git a/sentieon_cli/logging.py b/sentieon_cli/logging.py index 128485d..faca6d0 100644 --- a/sentieon_cli/logging.py +++ b/sentieon_cli/logging.py @@ -13,12 +13,5 @@ def get_logger(name: str) -> logging.Logger: def set_level(level: int | str) -> None: - """Set the level of the package logger. - - The package's module loggers (``sentieon_cli.dag``, - ``sentieon_cli.executor``, ...) are all at NOTSET and inherit their - effective level from the ``sentieon_cli`` package logger, so setting it - here scopes verbosity to this package without touching the root logger - (and thus third-party libraries). - """ + """Set the level of the package logger.""" logging.getLogger(__package__).setLevel(level) diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 1433db2..be82e0d 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -107,9 +107,6 @@ def handle_arguments(self, args: argparse.Namespace): def setup_logging(self, args: argparse.Namespace) -> None: self.logger = get_logger(__name__) - # Scope verbosity to this package's loggers rather than the root - # logger, so -v/--debug does not turn on DEBUG for every - # third-party library. set_level(args.loglevel) self.logger.info("Starting sentieon-cli version: %s", __version__) @@ -132,10 +129,6 @@ def main(self, args: argparse.Namespace) -> None: tmp_dir_str = tmp() self.tmp_dir = pathlib.Path(tmp_dir_str) - # Always clean up the temp dir, even if build_dag()/run() raise (an - # infeasible DAG, a launch failure, or a KeyboardInterrupt) -- - # otherwise it would leak. retain_tmpdir is honored on failure too, - # so a failed run can still be inspected. try: dag = self.build_dag() executor = self.run(dag) diff --git a/sentieon_cli/shell_pipeline.py b/sentieon_cli/shell_pipeline.py index 2469484..1d6d26d 100644 --- a/sentieon_cli/shell_pipeline.py +++ b/sentieon_cli/shell_pipeline.py @@ -22,13 +22,7 @@ def _freeze(value: Any) -> Any: - """Return an order-independent, hashable view of ``value``. - - ``exec_kwargs`` (and a nested ``env`` dict) have no stable iteration - order, so hashing their raw form would let two *equal* commands hash - differently, breaking the ``a == b implies hash(a) == hash(b)`` - contract the DAG relies on. Canonicalise dicts/sets recursively. - """ + """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)): @@ -42,12 +36,7 @@ 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. ``F_SETPIPE_SZ`` is Linux-only and, for unprivileged - processes, capped by ``/proc/sys/fs/pipe-max-size`` (exceeding it - raises EPERM) -- both deployment concerns, so a failure is logged and - the pipe keeps its default size rather than aborting the run. A no-op - where the constant is unavailable (non-Linux), preserving POSIX - portability. + stages. """ set_pipe_sz = getattr(fcntl, "F_SETPIPE_SZ", None) if set_pipe_sz is None: @@ -220,8 +209,6 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: if not isinstance(other, Command): return NotImplemented - # Identity is the *definition* of the command: not self.proc, which - # is runtime state that would mutate identity once the job runs. return ( self.executable == other.executable and self.args == other.args @@ -264,8 +251,6 @@ async def run( stderr: Union[IO[Any], int, None] = None, ) -> asyncio.subprocess.Process: # We cannot have handles from both files and to run(). - # Test `is not None`, not truthiness: an int fd of 0 is a valid - # stream but falsy. 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 " @@ -389,7 +374,6 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: if not isinstance(other, ProcSub): return NotImplemented - # An input and an output proc sub are not the same: <(x) != >(x). return type(self) is type(other) and self.node == other.node def __hash__(self) -> int: diff --git a/sentieon_cli/storage.py b/sentieon_cli/storage.py deleted file mode 100644 index 5fb6263..0000000 --- a/sentieon_cli/storage.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Storage staging for job inputs and outputs.""" - -import pathlib -from abc import ABC, abstractmethod - - -class StorageProvider(ABC): - """Stage a job's declared files in before it runs and out after. - - On a shared filesystem the files are already local, so staging - degenerates to verifying presence (see :class:`LocalStorageProvider`). - A cloud provider would instead pull inputs from and push outputs to - object storage, deriving each object's URI from a configured prefix and - the local path. The methods are async so they compose with the async - executors. - """ - - @abstractmethod - async def stage_in(self, local: pathlib.Path) -> None: - """Make ``local`` available before a job runs. - - Pull it from storage, or verify it is present. Raise if it cannot be - made available. - """ - - @abstractmethod - async def stage_out(self, local: pathlib.Path) -> None: - """Persist ``local`` after a job runs. - - Push it to storage, or verify it was produced. Raise if it is missing. - """ - - -class LocalStorageProvider(StorageProvider): - """Shared-filesystem provider: no data movement, just verify presence. - - A file counts as present when it exists, even if it is empty (zero bytes). - """ - - async def stage_in(self, local: pathlib.Path) -> None: - if not local.exists(): - raise FileNotFoundError(f"missing input: {local}") - - async def stage_out(self, local: pathlib.Path) -> None: - if not local.exists(): - raise FileNotFoundError(f"missing output: {local}") diff --git a/tests/integration/test_executor.py b/tests/integration/test_executor.py index 4f0a7f4..63c3779 100644 --- a/tests/integration/test_executor.py +++ b/tests/integration/test_executor.py @@ -16,7 +16,7 @@ 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) ) -from sentieon_cli.dag import DAG, has_all_outputs # noqa: E402 +from sentieon_cli.dag import DAG # noqa: E402 from sentieon_cli.exceptions import DagExecutionError # noqa: E402 from sentieon_cli.executor import ( # noqa: E402 DryRunExecutor, @@ -284,140 +284,8 @@ async def boom(self, *args, **kwargs): LocalExecutor(ThreadScheduler(dag, 1)).execute() -def test_missing_input_fails_before_run(tmp_path): - """A declared input that is missing fails the job before it runs.""" - missing = tmp_path / "in.txt" - marker = tmp_path / "ran.txt" - dag = DAG() - job = Job( - Pipeline(Command("touch", str(marker))), - "j", - inputs=[missing], - ) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - assert not marker.exists() # the command never ran - - -def test_present_inputs_and_outputs_allow_success(tmp_path): - """Empty (zero-byte) declared inputs/outputs still count as present.""" - inp = tmp_path / "in.txt" - inp.touch() # empty input file - out = tmp_path / "out.txt" - dag = DAG() - job = Job( - Pipeline(Command("touch", str(out))), # produces an empty output - "j", - inputs=[inp], - outputs=[out], - ) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert executor.jobs_with_errors == [] - assert out.exists() - - -def test_missing_output_fails_after_run(tmp_path): - """A declared output the command does not produce fails the job.""" - missing = tmp_path / "never.txt" - dag = DAG() - job = Job(Pipeline(Command("true")), "j", outputs=[missing]) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - - -def test_failed_job_outputs_are_removed(tmp_path): - """A failed job's declared outputs are deleted from disk.""" - out = tmp_path / "out.txt" - out.write_text("stale") - dag = DAG() - job = Job(Pipeline(Command("false")), "j", outputs=[out]) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - assert not out.exists() - - -def test_missing_output_failure_removes_produced_siblings(tmp_path): - """A missing declared output deletes the outputs that were produced.""" - produced = tmp_path / "out1.txt" - never = tmp_path / "out2.txt" - dag = DAG() - job = Job( - Pipeline(Command("touch", str(produced))), - "j", - outputs=[produced, never], - ) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - assert not produced.exists() - - -def test_launch_failure_removes_outputs(tmp_path): - """A job that fails to launch has its declared outputs deleted.""" - out = tmp_path / "out.txt" - out.write_text("stale") - dag = DAG() - job = Job(Pipeline(Command("no_such_cmd_zzz")), "j", outputs=[out]) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - assert not out.exists() - - -def test_missing_input_failure_removes_outputs(tmp_path): - """Any failed attempt clears the job's declared outputs -- even a - missing-input failure where the command never ran.""" - missing = tmp_path / "in.txt" - out = tmp_path / "out.txt" - out.write_text("stale") - dag = DAG() - job = Job(Pipeline(Command("true")), "j", inputs=[missing], outputs=[out]) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - assert not out.exists() - - -def test_successful_job_keeps_outputs(tmp_path): - """A successful job's outputs are never deleted.""" - out = tmp_path / "out.txt" - dag = DAG() - job = Job(Pipeline(Command("touch", str(out))), "j", outputs=[out]) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert executor.jobs_with_errors == [] - assert out.exists() - - -def test_failed_job_directory_output_is_left_in_place(tmp_path): - """Directories declared as outputs are never removed on failure.""" - out_dir = tmp_path / "out_dir" - out_dir.mkdir() - (out_dir / "keep.txt").write_text("keep") - dag = DAG() - job = Job(Pipeline(Command("false")), "j", outputs=[out_dir]) - dag.add_job(job) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert job in executor.jobs_with_errors - assert out_dir.is_dir() - assert (out_dir / "keep.txt").read_text() == "keep" - - def test_local_executor_empty_dag_completes(): - """An empty DAG (e.g. everything skipped by resume) runs cleanly.""" + """An empty DAG runs cleanly.""" executor = LocalExecutor(ThreadScheduler(DAG(), 1)) executor.execute() assert executor.jobs_with_errors == [] @@ -431,44 +299,6 @@ def test_dry_run_empty_dag_completes(capsys): assert executor.jobs_with_errors == [] -def test_resume_after_failure_reruns_only_failed_subtree(tmp_path): - """After a partial failure, skip_satisfied + has_all_outputs resumes - the run from the surviving intermediate files.""" - a_out = tmp_path / "a.txt" - b_out = tmp_path / "b.txt" - - def build_dag(b_cmd): - dag = DAG() - a = Job( - Pipeline(Command("echo", "a"), file_output=a_out), - "a", - outputs=[a_out], - ) - b = Job(Pipeline(b_cmd), "b", outputs=[b_out]) - dag.add_job(a) - dag.add_job(b, {a}) - return dag, a, b - - # First run: a succeeds, b fails (and b's partial output is deleted). - b_out.write_text("partial") - dag, _a, b = build_dag(Command("false")) - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert b in executor.jobs_with_errors - assert a_out.exists() - assert not b_out.exists() - - # Second run: rebuild the DAG with b fixed; resume skips a. - dag, a, b = build_dag(Command("touch", str(b_out))) - skipped = dag.skip_satisfied(has_all_outputs) - assert skipped == [a] - executor = LocalExecutor(ThreadScheduler(dag, 1)) - executor.execute() - assert executor.jobs_with_errors == [] - assert b_out.exists() - assert a_out.read_text().strip() == "a" # a was not rerun over - - def _execute_bounded(executor, timeout=15): """Run executor.execute() with a hang guard.""" holder = {} diff --git a/tests/integration/test_signals.py b/tests/integration/test_signals.py index 4aeae88..8aa2042 100644 --- a/tests/integration/test_signals.py +++ b/tests/integration/test_signals.py @@ -126,30 +126,6 @@ async def cleanup(self): assert not os.path.exists(temp_dir) # no leaked temp dir -def test_interrupted_job_outputs_are_removed(tmp_path): - """A job killed by an interrupt has its declared outputs deleted.""" - out = tmp_path / "out.txt" - out.write_text("partial") - dag = DAG() - job = Job(Pipeline(Command("sleep", "30")), "sleeper", outputs=[out]) - dag.add_job(job) - 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: - executor.execute() - finally: - timer.cancel() - - assert executor.running == [] - assert not out.exists() - - def test_grace_period_escalates_to_sigkill(): """A job that ignores SIGTERM is still killed after the grace period.""" script = ( diff --git a/tests/unit/test_dag.py b/tests/unit/test_dag.py index 444d07e..aa9d565 100644 --- a/tests/unit/test_dag.py +++ b/tests/unit/test_dag.py @@ -1,5 +1,5 @@ """ -Unit tests for DAG traversal, duplicate detection, and the resume prune. +Unit tests for DAG traversal and duplicate detection. """ import os @@ -13,7 +13,7 @@ os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), ) -from sentieon_cli.dag import DAG, has_all_outputs # 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.shell_pipeline import Command, Pipeline # noqa: E402 @@ -167,113 +167,3 @@ def test_stored_dependencies_are_independent_of_caller_set(self): deps.add(b) # mutating the caller's set must not affect the DAG assert dag.waiting_jobs[c] == {a} - - -class TestSkipSatisfied: - """Test the pre-run resume prune""" - - def test_satisfied_root_is_marked_finished(self): - dag = DAG() - job = _job("a") - dag.add_job(job) - skipped = dag.skip_satisfied(lambda j: True) - assert skipped == [job] - assert job in dag.finished_jobs - assert len(dag.ready_jobs) == 0 - - def test_unsatisfied_root_stays_ready(self): - dag = DAG() - job = _job("a") - dag.add_job(job) - skipped = dag.skip_satisfied(lambda j: False) - assert skipped == [] - assert job in dag.ready_jobs - assert dag.finished_jobs == [] - - def test_cascades_down_a_chain_in_topological_order(self): - dag = DAG() - a, b, c = _job("a"), _job("b"), _job("c") - dag.add_job(a) - dag.add_job(b, {a}) - dag.add_job(c, {b}) - skipped = dag.skip_satisfied(lambda j: True) - assert skipped == [a, b, c] - assert dag.finished_jobs == [a, b, c] - assert len(dag.ready_jobs) == 0 - assert len(dag.waiting_jobs) == 0 - - def test_stops_at_the_first_unsatisfied_job(self): - dag = DAG() - a, b, c = _job("a"), _job("b"), _job("c") - dag.add_job(a) - dag.add_job(b, {a}) - dag.add_job(c, {b}) - skipped = dag.skip_satisfied(lambda j: j == a) - assert skipped == [a] - assert b in dag.ready_jobs - assert c in dag.waiting_jobs - - def test_blocked_jobs_are_not_offered_to_the_predicate(self): - dag = DAG() - a, b, c = _job("a"), _job("b"), _job("c") - dag.add_job(a) - dag.add_job(b, {a}) - dag.add_job(c, {b}) - seen = [] - - def predicate(job): - seen.append(job) - return job == a - - dag.skip_satisfied(predicate) - # b is offered (a was skipped, unblocking it); c never is -- the - # ancestor rule: a job behind an unsatisfied job cannot be skipped. - assert seen == [a, b] - - def test_diamond_with_partial_parents(self): - dag = DAG() - top = _job("top") - left, right = _job("left"), _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}) - skipped = dag.skip_satisfied(lambda j: j in (top, left)) - assert skipped == [top, left] - assert right in dag.ready_jobs - assert dag.waiting_jobs[bottom] == {right} - - def test_empty_dag_returns_empty(self): - assert DAG().skip_satisfied(lambda j: True) == [] - - -class TestHasAllOutputs: - """Test the default resume predicate""" - - def test_no_declared_outputs_is_never_satisfied(self): - assert has_all_outputs(_job("a")) is False - - def test_true_when_all_outputs_exist(self, tmp_path): - full = tmp_path / "full.txt" - full.write_text("data") - empty = tmp_path / "empty.txt" - empty.touch() # zero-byte files count as present - job = Job(Pipeline(Command("true")), "j", outputs=[full, empty]) - assert has_all_outputs(job) is True - - def test_false_when_any_output_is_missing(self, tmp_path): - present = tmp_path / "present.txt" - present.touch() - missing = tmp_path / "missing.txt" - job = Job(Pipeline(Command("true")), "j", outputs=[present, missing]) - assert has_all_outputs(job) is False - - def test_dangling_symlink_counts_as_missing(self, tmp_path): - target = tmp_path / "target.txt" - target.touch() - link = tmp_path / "link.txt" - link.symlink_to(target) - target.unlink() # now dangling - job = Job(Pipeline(Command("true")), "j", outputs=[link]) - assert has_all_outputs(job) is False diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py deleted file mode 100644 index 4f46942..0000000 --- a/tests/unit/test_job.py +++ /dev/null @@ -1,52 +0,0 @@ -""" -Unit tests for the Job model. -""" - -import os -import pathlib -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.job import Job # noqa: E402 -from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 - - -def _pipe(text): - return Pipeline(Command("echo", text)) - - -def test_new_fields_default_empty(): - job = Job(_pipe("a"), "a") - assert job.inputs == [] - assert job.outputs == [] - assert job.image is None - - -def test_new_fields_are_stored(): - ins = [pathlib.Path("in1"), pathlib.Path("in2")] - outs = [pathlib.Path("out")] - job = Job(_pipe("a"), "a", inputs=ins, outputs=outs, image="org/img:1") - assert job.inputs == ins - assert job.outputs == outs - assert job.image == "org/img:1" - - -def test_input_output_sequences_are_snapshotted(): - ins = [pathlib.Path("in1")] - job = Job(_pipe("a"), "a", inputs=ins) - ins.append(pathlib.Path("in2")) # mutating the caller's list... - assert job.inputs == [pathlib.Path("in1")] # ...doesn't change the job - - -def test_identity_ignores_metadata_fields(): - # Identity is keyed on the shell pipeline; the other fields don't - # affect it. - j1 = Job(_pipe("same"), "j1", inputs=[pathlib.Path("x")], image="a") - j2 = Job(_pipe("same"), "j2", outputs=[pathlib.Path("y")], image="b") - assert j1 == j2 - assert hash(j1) == hash(j2) diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py deleted file mode 100644 index 800623f..0000000 --- a/tests/unit/test_storage.py +++ /dev/null @@ -1,49 +0,0 @@ -""" -Unit tests for the storage providers. -""" - -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.storage import LocalStorageProvider # noqa: E402 - - -@pytest.mark.asyncio -async def test_stage_in_present_file_ok(tmp_path): - f = tmp_path / "in.txt" - f.write_text("data") - await LocalStorageProvider().stage_in(f) # no raise - - -@pytest.mark.asyncio -async def test_stage_in_empty_file_ok(tmp_path): - f = tmp_path / "in.txt" - f.touch() # a zero-byte file still counts as present - await LocalStorageProvider().stage_in(f) - - -@pytest.mark.asyncio -async def test_stage_in_missing_file_raises(tmp_path): - with pytest.raises(FileNotFoundError): - await LocalStorageProvider().stage_in(tmp_path / "nope.txt") - - -@pytest.mark.asyncio -async def test_stage_out_present_file_ok(tmp_path): - f = tmp_path / "out.txt" - f.touch() - await LocalStorageProvider().stage_out(f) - - -@pytest.mark.asyncio -async def test_stage_out_missing_file_raises(tmp_path): - with pytest.raises(FileNotFoundError): - await LocalStorageProvider().stage_out(tmp_path / "nope.txt") From 56082bdd39b3f2a0d5c0e697c83273df3e40d0fe Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 13 Aug 2026 16:54:36 -0700 Subject: [PATCH 06/58] Add a per-run log directory with run.log and command.txt Move logging handlers to the package logger so verbosity is enforced per handler: the console follows -v/-d (default now INFO) while a new run.log always records DEBUG, both with timestamps. A new --log_dir argument (default: the output VCF with .vcf.gz replaced by _logs) holds run.log, a command.txt invocation record, and a task_logs/ directory for upcoming per-process tool logs. Runs now finish with a status/elapsed-time message, emitted from a finally so it also lands on failure paths. Co-Authored-By: Claude Fable 5 --- sentieon_cli/__init__.py | 2 +- sentieon_cli/logging.py | 62 ++++++-- sentieon_cli/pipeline.py | 88 ++++++++--- sentieon_cli/run_logs.py | 56 +++++++ sentieon_cli/sentieon_pangenome.py | 53 ++++--- tests/unit/conftest.py | 31 ++++ tests/unit/test_run_logs.py | 236 +++++++++++++++++++++++++++++ 7 files changed, 478 insertions(+), 50 deletions(-) create mode 100644 sentieon_cli/run_logs.py create mode 100644 tests/unit/conftest.py create mode 100644 tests/unit/test_run_logs.py diff --git a/sentieon_cli/__init__.py b/sentieon_cli/__init__.py index 3fcb34b..ceb76cb 100644 --- a/sentieon_cli/__init__.py +++ b/sentieon_cli/__init__.py @@ -16,7 +16,7 @@ def main(): action="store_const", dest="loglevel", const="INFO", - default="WARNING", + default="INFO", ) parser.add_argument( "-d", diff --git a/sentieon_cli/logging.py b/sentieon_cli/logging.py index faca6d0..df0a48a 100644 --- a/sentieon_cli/logging.py +++ b/sentieon_cli/logging.py @@ -1,17 +1,59 @@ +""" +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``/``-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) -handler = logging.StreamHandler() -handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s")) +# 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 logger with a StreamHandler.""" - logger = logging.getLogger(name) - logger.addHandler(handler) - logger.propagate = False - return 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) + + +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 set_level(level: int | str) -> None: - """Set the level of the package logger.""" - logging.getLogger(__package__).setLevel(level) +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/pipeline.py b/sentieon_cli/pipeline.py index be82e0d..f32e1b0 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -9,6 +9,7 @@ import pathlib import shutil import sys +import time from typing import Any, Dict, List, Optional import packaging.version @@ -18,7 +19,8 @@ from .exceptions import DagExecutionError from .executor import BaseExecutor, DryRunExecutor, LocalExecutor from .job import Job -from .logging import get_logger, set_level +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 @@ -53,6 +55,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, @@ -105,11 +114,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__) - set_level(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() @@ -117,25 +121,71 @@ 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 + + def setup_logging(self, args: argparse.Namespace) -> None: + """Configure the console handler and the run's log directory""" + self.logger = get_logger(__name__) + set_console_level(args.loglevel) + + # 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 not self.dry_run and ( + self.log_dir is not None or self.output_vcf is not None + ): + log_dir = self.log_dir + if log_dir is None: + # Check the suffix before deriving the directory name so an + # invalid output path cannot create a garbage-named log dir. + self.validate_output_suffix() + log_dir = pathlib.Path( + str(self.output_vcf).removesuffix(".vcf.gz") + "_logs" + ) + self.run_logs = RunLogs(log_dir) + self.run_logs.setup() + + # After the file handler is attached, so the banner reaches run.log + self.logger.info("Starting sentieon-cli version: %s", __version__) + if self.run_logs: + 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() + start_time = time.monotonic() + success = False + try: + self.validate() + self.configure() - tmp_dir_str = tmp() - self.tmp_dir = pathlib.Path(tmp_dir_str) + 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) + try: + dag = self.build_dag() + executor = self.run(dag) + self.check_execution(dag, executor) + finally: + if not self.retain_tmpdir: + shutil.rmtree(tmp_dir_str) + success = True finally: - if not self.retain_tmpdir: - shutil.rmtree(tmp_dir_str) + self.log_completion(success, start_time) def check_execution( self, @@ -158,10 +208,14 @@ def check_execution( 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(): diff --git a/sentieon_cli/run_logs.py b/sentieon_cli/run_logs.py new file mode 100644 index 0000000..2cfd5c8 --- /dev/null +++ b/sentieon_cli/run_logs.py @@ -0,0 +1,56 @@ +""" +The log directory for a single run +""" + +import datetime +import logging +import pathlib +import shlex +import shutil +import sys +from typing import Optional + +from .logging import add_file_handler, remove_file_handler +from .util import __version__ + + +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.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 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) + + 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/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 6d73fcf..456680e 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -8,6 +8,7 @@ import pathlib import shutil import sys +import time from typing import Dict, List, Optional, Set, Tuple, Union import packaging.version @@ -269,34 +270,42 @@ def main(self, args: argparse.Namespace) -> None: """Run the pipeline""" self.handle_arguments(args) self.setup_logging(args) - self.validate_ref() - - 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) - - self.validate() - self.shards = determine_shards_from_fai( - self.fai_data, 10 * 1000 * 1000 - ) + start_time = time.monotonic() + success = False + try: + self.validate_ref() - tmp_dir_str = tmp() - self.tmp_dir = pathlib.Path(tmp_dir_str) + 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) + + self.validate() + self.shards = determine_shards_from_fai( + self.fai_data, 10 * 1000 * 1000 + ) - dag = self.build_first_dag() - executor = self.run(dag) - self.check_execution(dag, executor) + tmp_dir_str = tmp() + self.tmp_dir = pathlib.Path(tmp_dir_str) - if self.expansion_catalog or self.segdup_caller is not None: - self.get_sex(self.ploidy_json) - dag = self.build_second_dag() + dag = self.build_first_dag() executor = self.run(dag) self.check_execution(dag, executor) - if not self.retain_tmpdir: - shutil.rmtree(tmp_dir_str) + 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) + success = True + finally: + self.log_completion(success, start_time) def validate(self) -> None: """Validate pipeline inputs""" diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..2c16132 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,31 @@ +""" +Shared unit-test fixtures +""" + +import logging + +import pytest + +from sentieon_cli import logging as cli_logging + + +@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_run_logs.py b/tests/unit/test_run_logs.py new file mode 100644 index 0000000..a16fb1b --- /dev/null +++ b/tests/unit/test_run_logs.py @@ -0,0 +1,236 @@ +""" +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.pipeline import BasePipeline # noqa: E402 +from sentieon_cli.run_logs import RunLogs # 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 _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 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" + pipeline.setup_logging(_args()) + + 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" + pipeline.setup_logging(_args()) + + 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" + pipeline.setup_logging(_args()) + + 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: + pipeline.setup_logging(_args()) + + assert excinfo.value.code == 2 + assert list(tmp_path.iterdir()) == [] + + +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" + pipeline.setup_logging(_args("INFO")) + 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" + pipeline.setup_logging(_args()) + 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 + pipeline.setup_logging(_args()) + 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" + pipeline.setup_logging(_args()) + + 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() + pipeline.setup_logging(_args()) + + assert pipeline.run_logs is None + assert list(tmp_path.iterdir()) == [] + + +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() From 0f7930acee3d24a5275b688719df701d1ac0684c Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 13 Aug 2026 22:40:52 -0700 Subject: [PATCH 07/58] Give jobs a task name and a unique, human-readable job id Job now requires a keyword-only task_name grouping jobs by pipeline stage, and assigns each job a job_id ({name}-{n}) from per-name counters that reset once per CLI invocation, so ids stay unique across multi-DAG runs. Jobs report themselves by job_id in scheduler and error messages. These fields feed the upcoming per-process task logs; neither is part of job identity, which stays keyed on the shell pipeline alone. Co-Authored-By: Claude Fable 5 --- sentieon_cli/__init__.py | 4 + sentieon_cli/base_pangenome.py | 2 + sentieon_cli/dnascope.py | 50 +++++++++++-- sentieon_cli/dnascope_hybrid.py | 61 +++++++++++++-- sentieon_cli/dnascope_longread.py | 50 +++++++++++-- sentieon_cli/job.py | 28 ++++++- sentieon_cli/pipeline.py | 1 + sentieon_cli/sentieon_pangenome.py | 28 ++++++- sentieon_cli/transfer.py | 3 + tests/integration/conftest.py | 23 ++++++ tests/integration/test_executor.py | 83 +++++++++++++++----- tests/integration/test_signals.py | 17 ++++- tests/unit/conftest.py | 12 +++ tests/unit/test_dag.py | 30 ++++++-- tests/unit/test_dag_construction.py | 45 +++++++++-- tests/unit/test_job.py | 104 ++++++++++++++++++++++++++ tests/unit/test_pipeline_lifecycle.py | 11 ++- tests/unit/test_scheduler.py | 16 +++- tests/unit/test_scheduler_contract.py | 4 +- 19 files changed, 504 insertions(+), 68 deletions(-) create mode 100644 tests/integration/conftest.py create mode 100644 tests/unit/test_job.py diff --git a/sentieon_cli/__init__.py b/sentieon_cli/__init__.py index ceb76cb..1de7890 100644 --- a/sentieon_cli/__init__.py +++ b/sentieon_cli/__init__.py @@ -2,6 +2,7 @@ from .dnascope import DNAscopePipeline from .dnascope_hybrid import DNAscopeHybridPipeline from .dnascope_longread import DNAscopeLRPipeline +from .job import Job from .sentieon_pangenome import SentieonPangenome from .util import __version__ @@ -59,6 +60,9 @@ def main(): dnascope_pangenome_subparser.set_defaults(pipeline=pipeline.main) args = parser.parse_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..7284de6 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -129,6 +129,7 @@ def build_kmc_job( ), "kmc", job_threads, + task_name="kmer-counting", ) return kmc_job @@ -150,6 +151,7 @@ def build_ploidy_job( ), "estimate-ploidy", 0, + task_name="ploidy", ) return ploidy_job diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index c4fa230..70a2e10 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -517,6 +517,7 @@ def sr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job], Job]: ), f"bam-align-{i}", self.cores, + task_name="alignment", ) res.append(out_aln) jobs.add(job) @@ -532,6 +533,7 @@ def sr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job], Job]: ), "rm-bam-aln", 0, + task_name="cleanup", ) return (res, jobs, rm_job) @@ -613,6 +615,7 @@ def sr_align_fastq( f"bam-align-{i}-{j}", split_cores, resources={f"node{j}": 1}, + task_name="alignment", ) res.append(out_aln) jobs.add(job) @@ -628,6 +631,7 @@ def sr_align_fastq( ), "rm-fq-aln", 0, + task_name="cleanup", ) return (res, jobs, rm_job) @@ -716,6 +720,7 @@ def dedup_and_metrics( Pipeline(Command(*driver.build_cmd())), "locuscollector", self.cores, + task_name="dedup", ) if self.sr_duplicate_marking == "none": @@ -743,7 +748,10 @@ def dedup_and_metrics( ) ) dedup_job = Job( - Pipeline(Command(*driver.build_cmd())), "dedup", self.cores + Pipeline(Command(*driver.build_cmd())), + "dedup", + self.cores, + task_name="dedup", ) if self.skip_metrics: @@ -767,7 +775,10 @@ def dedup_and_metrics( 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 + Pipeline(Command(*driver.build_cmd())), + "metrics", + 0, + task_name="metrics", ) # Run metrics in the background # Run WgsMetricsAlgo after duplicate marking to account for @@ -779,7 +790,10 @@ def dedup_and_metrics( ) driver.add_algo(CoverageMetrics(coverage_metrics)) metrics_job = Job( - Pipeline(Command(*driver.build_cmd())), "metrics", 0 + Pipeline(Command(*driver.build_cmd())), + "metrics", + 0, + task_name="metrics", ) # Run metrics in the background # Rehead WGS metrics so they are recognized by MultiQC @@ -801,6 +815,7 @@ def dedup_and_metrics( ), "Rehead metrics", 0, + task_name="metrics", ) return ([deduped], lc_job, dedup_job, metrics_job, rehead_job) @@ -874,6 +889,7 @@ def sr_call_variants( Pipeline(Command(*driver.build_cmd())), "variant-calling", self.cores, + task_name="variant-calling", ) # Genotyping and filtering with DNAModelApply @@ -889,12 +905,20 @@ def sr_call_variants( ) ) apply_job = Job( - Pipeline(Command(*driver.build_cmd())), "model-apply", self.cores + Pipeline(Command(*driver.build_cmd())), + "model-apply", + self.cores, + task_name="model-apply", ) # 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) + rm_job = Job( + Pipeline(Command(*rm_cmd, fail_ok=True)), + "rm-tmp-vcf", + 0, + task_name="cleanup", + ) # Genotype gVCFs gvcftyper_job = None @@ -911,7 +935,10 @@ def sr_call_variants( ) ) gvcftyper_job = Job( - Pipeline(Command(*driver.build_cmd())), "gvcftyper", self.cores + Pipeline(Command(*driver.build_cmd())), + "gvcftyper", + self.cores, + task_name="gvcftyper", ) # Call SVs @@ -930,7 +957,9 @@ def sr_call_variants( ) ) svsolver_job = Job( - Pipeline(Command(*driver.build_cmd())), "svsolver" + Pipeline(Command(*driver.build_cmd())), + "svsolver", + task_name="sv-calling", ) sv_rm_job = Job( Pipeline( @@ -943,6 +972,7 @@ def sr_call_variants( ), "rm-tmp-sv", 0, + task_name="cleanup", ) return ( @@ -991,7 +1021,10 @@ def call_cnvs( ) ) cnvscope_job = Job( - Pipeline(Command(*driver.build_cmd())), "CNVscope", cores + Pipeline(Command(*driver.build_cmd())), + "CNVscope", + cores, + task_name="cnv", ) driver = Driver( @@ -1009,6 +1042,7 @@ def call_cnvs( Pipeline(Command(*driver.build_cmd())), "CNVModelApply", cores, + task_name="cnv", ) return (cnvscope_job, cnvmodelapply_job) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index e1a9e28..6be5285 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -764,7 +764,10 @@ def call_variants( ) ) call_job = Job( - Pipeline(Command(*driver.build_cmd())), "calling-1", self.cores + Pipeline(Command(*driver.build_cmd())), + "calling-1", + self.cores, + task_name="variant-calling", ) # Region selection @@ -782,6 +785,7 @@ def call_variants( ), "hybrid-select", 0, + task_name="region-selection", ) mapq0_bed = self.tmp_dir.joinpath("hybrid_mapq0.bed") @@ -802,6 +806,7 @@ def call_variants( Pipeline(Command(*driver.build_cmd())), "mapq0-bed", self.cores, + task_name="region-selection", ) mapq0_slop_bed = self.tmp_dir.joinpath("hybrid_mapq0.ex1000.bed") @@ -814,6 +819,7 @@ def call_variants( ), "mapq0-bed-slop", 0, + task_name="region-selection", ) diff_bed = self.tmp_dir.joinpath("merged_diff.bed") @@ -825,9 +831,15 @@ def call_variants( ), "concat-merge-bed", 0, + task_name="region-selection", ) rm_cmd = ["rm", str(selected_bed), str(mapq0_slop_bed)] - rm_job1 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp1", 0) + rm_job1 = Job( + Pipeline(Command(*rm_cmd, fail_ok=True)), + "rm-tmp1", + 0, + task_name="cleanup", + ) stage1_ins_fa = self.tmp_dir.joinpath("stage1_ins.fa") stage1_ins_bed = self.tmp_dir.joinpath("stage1_ins.bed") @@ -881,6 +893,7 @@ def call_variants( ), "first-stage", self.cores, + task_name="hybrid-realignment", ) rm_cmd = [ "rm", @@ -888,7 +901,12 @@ def call_variants( str(stage1_ins_bed), str(stage1_hap_vcf), ] - rm_job2 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp2", 0) + rm_job2 = Job( + Pipeline(Command(*rm_cmd, fail_ok=True)), + "rm-tmp2", + 0, + task_name="cleanup", + ) stage2_bed = self.tmp_dir.joinpath("hybrid_stage2.bed") stage2_unmap_bam = self.tmp_dir.joinpath("hybrid_stage2_unmap.bam") @@ -911,10 +929,16 @@ def call_variants( Pipeline(Command(*driver.build_cmd())), "second-stage", self.cores, + task_name="hybrid-realignment", ) rm_cmd = ["rm", str(stage1_bam), str(stage1_hap_bam)] - rm_job3 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp3", 0) + rm_job3 = Job( + Pipeline(Command(*rm_cmd, fail_ok=True)), + "rm-tmp3", + 0, + task_name="cleanup", + ) suffix = "bam" if self.bam_format else "cram" stage3_aln = pathlib.Path( @@ -943,9 +967,15 @@ def call_variants( ), "third-stage", self.cores, + task_name="hybrid-realignment", ) rm_cmd = ["rm", str(stage2_unmap_bam), str(stage2_alt_bam)] - rm_job4 = Job(Pipeline(Command(*rm_cmd, fail_ok=True)), "rm-tmp4", 0) + rm_job4 = Job( + Pipeline(Command(*rm_cmd, fail_ok=True)), + "rm-tmp4", + 0, + task_name="cleanup", + ) # pass 2 of variant calling pass2_vcf = self.tmp_dir.joinpath("hybrid_pass2.vcf.gz") @@ -967,7 +997,10 @@ def call_variants( ) ) call2_job = Job( - Pipeline(Command(*driver.build_cmd())), "call2", self.cores + Pipeline(Command(*driver.build_cmd())), + "call2", + self.cores, + task_name="variant-calling", ) # Merge and normalize the VCFs @@ -981,6 +1014,7 @@ def call_variants( ), "subset-calls", 0, + task_name="vcf-merge", ) concat_job = Job( cmds.bcftools_concat( @@ -989,9 +1023,15 @@ def call_variants( ), "concat-calls", 0, + task_name="vcf-merge", ) 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) + rm_job5 = Job( + Pipeline(Command(*rm_cmd, fail_ok=True)), + "rm-tmp5", + 0, + task_name="cleanup", + ) # Annotate the output VCF hybrid_anno = pathlib.Path( @@ -1011,6 +1051,7 @@ def call_variants( ), "anno-calls", 0, + task_name="annotation", ) transfer_jobs: Optional[List[Job]] = None @@ -1077,7 +1118,10 @@ def call_variants( ) ) apply_job = Job( - Pipeline(Command(*driver.build_cmd())), "model-apply", self.cores + Pipeline(Command(*driver.build_cmd())), + "model-apply", + self.cores, + task_name="model-apply", ) # Final normalize @@ -1090,6 +1134,7 @@ def call_variants( ), "final-norm", 0, + task_name="vcf-norm", ) return ( call_job, diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index c256cd0..ca4aeb0 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -672,6 +672,7 @@ def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: ), f"bam-realign-{i}", self.cores, + task_name="alignment", ) ) res.append(out_aln) @@ -725,8 +726,9 @@ def lr_align_fastq(self) -> Tuple[List[pathlib.Path], Set[Job]]: self.minimap2_args, self.util_sort_args, ), - "align-{i}", + f"align-{i}", self.cores, + task_name="alignment", ) ) res.append(out_aln) @@ -764,6 +766,7 @@ def mosdepth(self, sample_input: List[pathlib.Path]) -> Set[Job]: ), f"mosdepth-{i}", 0, # Run in background + task_name="metrics", ) ) return mosdepth_jobs @@ -791,6 +794,7 @@ def merge_input_files( Pipeline(Command(*driver.build_cmd())), "merge-bam", 0, + task_name="alignment", ) return (merged_bam, merge_job) @@ -823,6 +827,7 @@ def pbsv( ), "pbsv-discover", 0, + task_name="sv-calling", ) # pbsv call @@ -840,6 +845,7 @@ def pbsv( Pipeline(Command(*call_cmd)), "pbsv-call", self.cores, + task_name="sv-calling", ) return (pbsv_discover, pbsv_call) @@ -875,7 +881,10 @@ 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 @@ -935,7 +944,10 @@ def lr_call_variants( ) ) first_calling_job = Job( - Pipeline(Command(*driver.build_cmd())), "first-pass", self.cores + Pipeline(Command(*driver.build_cmd())), + "first-pass", + self.cores, + task_name="variant-calling", ) # Transfer annotations to the tmp vcf @@ -970,6 +982,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "first-modelapply", self.cores, + task_name="model-apply", ) # Phasing and RepeatModel @@ -998,7 +1011,10 @@ def lr_call_variants( ) ) phaser_job = Job( - Pipeline(Command(*driver.build_cmd())), "variantphaser", self.cores + Pipeline(Command(*driver.build_cmd())), + "variantphaser", + self.cores, + task_name="phasing", ) bcftools_subset_phased_job = None @@ -1022,6 +1038,7 @@ def lr_call_variants( ), "bcftools-subset-phased", 0, + task_name="phasing", ) fai_to_bed_job = None @@ -1036,12 +1053,14 @@ def lr_call_variants( ), "fai-to-bed", 0, + task_name="phasing", ) bcftools_subtract_job = Job( cmds.cmd_bedtools_subtract(bed, phased_bed, unphased_bed), "bedtools-subtract", 0, + task_name="phasing", ) repeatmodel_job = None @@ -1068,6 +1087,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "repeatmodel", self.cores, + task_name="repeat-model", ) bcftools_subset_unphased_job = Job( @@ -1085,6 +1105,7 @@ def lr_call_variants( ), "bcftools-subset-unphased", 0, + task_name="phasing", ) # Second pass - phased variants @@ -1127,6 +1148,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "second-pass", self.cores, + task_name="variant-calling", ) ) @@ -1153,6 +1175,7 @@ def lr_call_variants( ), "patch", self.cores, + task_name="variant-patch", ) # Transfer annotations to the patched VCFs @@ -1198,6 +1221,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "second-modelapply", self.cores, + task_name="model-apply", ) ) @@ -1223,6 +1247,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "calling-unphased", self.cores, + task_name="variant-calling", ) # Patch DNA and DNAHP variants @@ -1236,7 +1261,9 @@ def lr_call_variants( self.cores, kwargs, ) - diploid_patch_job = Job(cmd, "diploid-patch", self.cores) + diploid_patch_job = Job( + cmd, "diploid-patch", self.cores, task_name="variant-patch" + ) # Transfer annotations to the diploid VCF unphased_transfer_jobs: List[Job] = [] @@ -1270,6 +1297,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "modelapply-unphased", self.cores, + task_name="model-apply", ) # merge calls to create the output @@ -1290,6 +1318,7 @@ def lr_call_variants( ), "merge", self.cores, + task_name="vcf-merge", ) gvcf_combine_job = None @@ -1304,6 +1333,7 @@ def lr_call_variants( ), "gvcf-combine", 0, + task_name="gvcf", ) haploid_calling_job = None @@ -1352,6 +1382,7 @@ def lr_call_variants( Pipeline(Command(*driver.build_cmd())), "haploid-calling", self.cores, + task_name="variant-calling", ) haploid_patch2_job = Job( @@ -1364,6 +1395,7 @@ def lr_call_variants( ), "haploid-patch2", self.cores, + task_name="variant-patch", ) haploid_concat_job = Job( cmds.bcftools_concat( @@ -1372,6 +1404,7 @@ def lr_call_variants( ), "haploid-diploid-concat", 0, + task_name="vcf-merge", ) if self.gvcf: @@ -1394,6 +1427,7 @@ def lr_call_variants( ), "haploid-gvcf-combine", 0, + task_name="gvcf", ) haploid_gvcf_concat_job = Job( cmds.bcftools_concat( @@ -1402,6 +1436,7 @@ def lr_call_variants( ), "haploid-gvcf-concat", 0, + task_name="gvcf", ) return LRCallVariantsResult( first_calling_job, @@ -1466,6 +1501,9 @@ def call_svs( ) ) longreadsv_job = Job( - Pipeline(Command(*driver.build_cmd())), "LongReadSV", self.cores + Pipeline(Command(*driver.build_cmd())), + "LongReadSV", + self.cores, + task_name="sv-calling", ) return longreadsv_job diff --git a/sentieon_cli/job.py b/sentieon_cli/job.py index 3e36fa2..bb893bb 100644 --- a/sentieon_cli/job.py +++ b/sentieon_cli/job.py @@ -16,22 +16,46 @@ class Job: * ``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, pipeline: Pipeline, 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 + count = Job._id_counters.get(name, 0) + 1 + Job._id_counters[name] = count + self.job_id = f"{name}-{count}" + + @classmethod + def reset_ids(cls) -> None: + """Restart id numbering. + + 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) @@ -45,7 +69,7 @@ def __ne__(self, other: object) -> bool: return not self == other def __repr__(self) -> str: - return f"Job({self.name})" + return f"Job({self.job_id})" def __str__(self) -> str: - return f"Job({self.name})" + return f"Job({self.job_id})" diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index f32e1b0..78726c9 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -288,6 +288,7 @@ def multiqc(self) -> Optional[Job]: ), "multiqc", 0, + task_name="multiqc", ) return multiqc_job diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 456680e..50f23af 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -705,6 +705,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) @@ -721,6 +722,7 @@ def build_first_dag(self) -> DAG: ), "extract-kmc", self.cores, + task_name="read-extraction", ) dag.add_job(extract_kmc_job, {ln_job}) haplotype_dependencies.add(extract_kmc_job) @@ -969,6 +971,7 @@ def build_alignment_job( ), "bwa-extract", self.cores, + task_name="alignment", ) return bwa_job @@ -999,6 +1002,7 @@ def build_haplotypes_job( ), "vg-haplotypes", self.cores, + task_name="pangenome", ) return haplotypes_job @@ -1015,6 +1019,7 @@ def build_gfa_job( ), "vg-convert-gfa", 0, + task_name="pangenome", ) return gfa_job @@ -1029,6 +1034,7 @@ def build_fasta_job( ), "vg-paths-fasta", 0, + task_name="pangenome", ) return fasta_job @@ -1073,6 +1079,7 @@ def build_minimap2_lift_job( ), "mm2-lift", self.cores, + task_name="pangenome-alignment", ) return mm2_job @@ -1107,6 +1114,7 @@ def build_dedup_job( Pipeline(Command(*driver.build_cmd())), f"locuscollector-{tag}", self.cores, + task_name="dedup", ) driver2 = Driver( @@ -1121,6 +1129,7 @@ def build_dedup_job( Pipeline(Command(*driver2.build_cmd())), f"dedup-{tag}", self.cores, + task_name="dedup", ) return lc_job, dedup_job @@ -1176,7 +1185,12 @@ def build_metrics_job( 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) + metrics_job = Job( + Pipeline(Command(*driver.build_cmd())), + "metrics", + 0, + task_name="metrics", + ) rehead_script = pathlib.Path( str( @@ -1194,6 +1208,7 @@ def build_metrics_job( ), "Rehead metrics", 0, + task_name="metrics", ) return (metrics_job, rehead_job) @@ -1242,6 +1257,7 @@ def build_dnascope_job( Pipeline(Command(*driver.build_cmd())), "dnascope-raw", self.cores, + task_name="variant-calling", ) def build_dnamodelapply_job( @@ -1268,6 +1284,7 @@ def build_dnamodelapply_job( Pipeline(Command(*driver.build_cmd())), "model-apply", self.cores, + task_name="model-apply", ) def build_gvcftyper_job( @@ -1295,6 +1312,7 @@ def build_gvcftyper_job( Pipeline(Command(*driver.build_cmd())), "gvcftyper", self.cores, + task_name="gvcftyper", ) def build_segdup_job( @@ -1333,6 +1351,7 @@ def build_segdup_job( ), "segdup-caller", self.cores, + task_name="segdup", ) def build_second_dag(self) -> DAG: @@ -1397,6 +1416,7 @@ def build_t1k_jobs( Pipeline(Command(*driver.build_cmd())), f"t1k-{tag}-extract", self.cores, + task_name="t1k", ) t1k_job = Job( @@ -1410,6 +1430,7 @@ def build_t1k_jobs( ), f"t1k-{tag}", self.cores, + task_name="t1k", ) return (extract_job, t1k_job) @@ -1435,6 +1456,7 @@ def build_expansion_job( ), "expansion-hunter", self.cores, + task_name="expansion-hunter", ) def _add_cnv_jobs( @@ -1506,6 +1528,7 @@ def _build_cnvscope_job( Pipeline(Command(*driver.build_cmd())), "cnvscope", self.cores, + task_name="cnv", ) def _build_cnv_model_apply_job( @@ -1530,6 +1553,7 @@ def _build_cnv_model_apply_job( Pipeline(Command(*driver.build_cmd())), "cnv-model-apply", self.cores, + task_name="cnv", ) def _build_indel2cnv_job( @@ -1556,6 +1580,7 @@ def _build_indel2cnv_job( ), "indel2cnv", 0, + task_name="cnv", ) def _build_combine_cnv_job( @@ -1583,4 +1608,5 @@ def _build_combine_cnv_job( ), "combine-cnv", 0, + task_name="cnv", ) diff --git a/sentieon_cli/transfer.py b/sentieon_cli/transfer.py index a5fed93..ed082fb 100644 --- a/sentieon_cli/transfer.py +++ b/sentieon_cli/transfer.py @@ -89,6 +89,7 @@ def build_transfer_jobs( ), "merge-trim-extra", 1, + task_name="annotation-transfer", ) sharded_merge_jobs.append(view_job) sharded_vcfs.append(subset_vcf) @@ -123,6 +124,7 @@ def build_transfer_jobs( ), f"merge-trim-{i}", 1, + task_name="annotation-transfer", ) sharded_merge_jobs.append(merge_job) sharded_vcfs.append(shard_vcf) @@ -136,5 +138,6 @@ def build_transfer_jobs( ), "merge-trim-concat", cores, + task_name="annotation-transfer", ) return (sharded_merge_jobs, concat_job) 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 63c3779..d0b1b6d 100644 --- a/tests/integration/test_executor.py +++ b/tests/integration/test_executor.py @@ -40,6 +40,7 @@ def test_local_executor_simple_job(): job = Job( Pipeline(Command("echo", "hello executor"), file_output=cmd_out), "echo-job", + task_name="test", ) dag.add_job(job) @@ -61,7 +62,7 @@ def test_local_executor_pipeline_job(): Command("cat"), file_output=cmd_out, ) - job = Job(pipeline, "pipeline-job") + job = Job(pipeline, "pipeline-job", task_name="test") dag.add_job(job) scheduler = ThreadScheduler(dag, 1) @@ -78,7 +79,11 @@ def test_local_executor_failing_job(): cmd_in = pathlib.Path(tmp_dir_str) / "test_in.txt" dag = DAG() # This command will fail: the input file does not exist - job = Job(Pipeline(Command("cat", str(cmd_in))), "failing-job") + job = Job( + Pipeline(Command("cat", str(cmd_in))), + "failing-job", + task_name="test", + ) dag.add_job(job) scheduler = ThreadScheduler(dag, 1) @@ -99,7 +104,11 @@ def test_local_executor_proc_sub_job(): 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 @@ -125,6 +134,7 @@ def test_dry_run_executor_prints_without_running(capsys): job = Job( Pipeline(Command("echo", "dryrun"), file_output=cmd_out), "dry-job", + task_name="test", ) dag.add_job(job) @@ -142,8 +152,8 @@ def test_dry_run_executor_prints_without_running(capsys): 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") - b = Job(Pipeline(Command("echo", "b")), "b") + 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}) @@ -172,6 +182,7 @@ def test_fail_ok_subcommand_does_not_fail_job(): file_output=out, ), "tolerant-job", + task_name="test", ) dag.add_job(job) @@ -189,8 +200,20 @@ def test_independent_jobs_all_complete(): o1 = tmp_dir / "o1.txt" o2 = tmp_dir / "o2.txt" dag = DAG() - dag.add_job(Job(Pipeline(Command("echo", "1"), file_output=o1), "j1")) - dag.add_job(Job(Pipeline(Command("echo", "2"), file_output=o2), "j2")) + 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) @@ -208,8 +231,16 @@ def test_dependent_job_runs_after_dependency(): first = tmp_dir / "first.txt" second = tmp_dir / "second.txt" dag = DAG() - a = Job(Pipeline(Command("echo", "a"), file_output=first), "a") - b = Job(Pipeline(Command("echo", "b"), file_output=second), "b") + 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}) @@ -225,7 +256,7 @@ def test_dependent_job_runs_after_dependency(): 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") + job = Job(Pipeline(Command("no_such_cmd_zzz")), "ghost", task_name="test") dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) executor.execute() @@ -239,7 +270,9 @@ def test_launch_loop_stops_after_a_failed_launch(): 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}")) + dag.add_job( + Job(Pipeline(Command("true", str(i))), f"j{i}", task_name="test") + ) executor = LocalExecutor(ThreadScheduler(dag, 5)) launched = [] @@ -259,14 +292,18 @@ async def fake_run_job(job): 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)) + 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)) + dag.add_job( + Job(Pipeline(Command("echo", "hi")), "big", 100, task_name="test") + ) with pytest.raises(DagExecutionError): DryRunExecutor(ThreadScheduler(dag, 2)).execute() @@ -279,7 +316,7 @@ async def boom(self, *args, **kwargs): monkeypatch.setattr("sentieon_cli.shell_pipeline.Pipeline.run", boom) dag = DAG() - dag.add_job(Job(Pipeline(Command("echo", "x")), "j")) + dag.add_job(Job(Pipeline(Command("echo", "x")), "j", task_name="test")) with pytest.raises(RuntimeError): LocalExecutor(ThreadScheduler(dag, 1)).execute() @@ -325,6 +362,7 @@ def test_procsub_job_with_early_exiting_outer_does_not_hang(): Command("false", InputProcSub(Pipeline(Command("echo", "x")))) ), "early-exit", + task_name="test", ) dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) @@ -344,6 +382,7 @@ def test_procsub_job_with_unspawnable_outer_does_not_hang(): ) ), "ghost-with-procsub", + task_name="test", ) dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) @@ -364,6 +403,7 @@ def test_procsub_inner_that_fails_to_launch_is_recorded(): ) ), "inner-fail", + task_name="test", ) dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) @@ -384,8 +424,9 @@ def test_procsub_inner_failure_does_not_skip_sibling_bookkeeping(): ) ), "inner-fail", + task_name="test", ) - good = Job(Pipeline(Command("true")), "sibling") + good = Job(Pipeline(Command("true")), "sibling", task_name="test") dag.add_job(bad) dag.add_job(good) executor = LocalExecutor(ThreadScheduler(dag, 2)) @@ -419,6 +460,7 @@ def __init__(self) -> None: Command("no_such_cmd_zzz"), ), "leak", + task_name="test", ) dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) @@ -446,6 +488,7 @@ def test_sigpipe_producer_does_not_fail_job(tmp_path): file_output=out, ), "sigpipe", + task_name="test", ) dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) @@ -458,7 +501,11 @@ 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") + job = Job( + Pipeline(Command("false"), Command("cat")), + "real-fail", + task_name="test", + ) dag.add_job(job) executor = LocalExecutor(ThreadScheduler(dag, 2)) _execute_bounded(executor) @@ -510,13 +557,13 @@ async def _done(): executor.running = [ ( - Job(Pipeline(Command("true")), "j1"), + Job(Pipeline(Command("true")), "j1", task_name="test"), bad, asyncio.create_task(_done()), 0, ), ( - Job(Pipeline(Command("false")), "j2"), + Job(Pipeline(Command("false")), "j2", task_name="test"), good, asyncio.create_task(_done()), 0, diff --git a/tests/integration/test_signals.py b/tests/integration/test_signals.py index 8aa2042..665a935 100644 --- a/tests/integration/test_signals.py +++ b/tests/integration/test_signals.py @@ -35,7 +35,11 @@ def _echo_dag(out): dag = DAG() - job = Job(Pipeline(Command("echo", "hi"), file_output=out), "echo") + job = Job( + Pipeline(Command("echo", "hi"), file_output=out), + "echo", + task_name="test", + ) dag.add_job(job) return dag @@ -103,7 +107,9 @@ async def cleanup(self): monkeypatch.setattr(executor_mod, "Context", SpyContext) dag = DAG() - dag.add_job(Job(Pipeline(Command("sleep", "30")), "sleeper")) + dag.add_job( + Job(Pipeline(Command("sleep", "30")), "sleeper", task_name="test") + ) executor = LocalExecutor( ThreadScheduler(dag, 1), install_signal_handlers=True, @@ -134,7 +140,11 @@ def test_grace_period_escalates_to_sigkill(): "time.sleep(30)" ) dag = DAG() - job = Job(Pipeline(Command(sys.executable, "-c", script)), "stubborn") + job = Job( + Pipeline(Command(sys.executable, "-c", script)), + "stubborn", + task_name="test", + ) dag.add_job(job) executor = LocalExecutor( ThreadScheduler(dag, 1), @@ -172,6 +182,7 @@ def test_sigint_with_blocked_procsub_cleans_up(): ) ), "blocked-procsub", + task_name="test", ) dag.add_job(job) executor = LocalExecutor( diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 2c16132..226413c 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -7,6 +7,18 @@ 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) diff --git a/tests/unit/test_dag.py b/tests/unit/test_dag.py index aa9d565..014f945 100644 --- a/tests/unit/test_dag.py +++ b/tests/unit/test_dag.py @@ -21,15 +21,17 @@ 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) + 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") - j2 = Job(Pipeline(Command("echo", "x")), "name-b") - j3 = Job(Pipeline(Command("echo", "y")), "name-c") + 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 @@ -101,17 +103,29 @@ def test_multilevel_chain(self): def test_add_duplicate_pipeline_raises(self): dag = DAG() - dag.add_job(Job(Pipeline(Command("echo", "x")), "step-1")) + 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")) + 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") - j2 = Job(Pipeline(Command("run", exec_kwargs={"cwd": "/b"})), "b") + 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 diff --git a/tests/unit/test_dag_construction.py b/tests/unit/test_dag_construction.py index e9bacdf..c0df2cf 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""" @@ -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) diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py new file mode 100644 index 0000000..4b37a58 --- /dev/null +++ b/tests/unit/test_job.py @@ -0,0 +1,104 @@ +""" +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 + + +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_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_pipeline_lifecycle.py b/tests/unit/test_pipeline_lifecycle.py index b9659c1..593a03f 100644 --- a/tests/unit/test_pipeline_lifecycle.py +++ b/tests/unit/test_pipeline_lifecycle.py @@ -131,9 +131,10 @@ def __init__(self, jobs_with_errors=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") + job = Job(Pipeline(Command("false")), "broken-step", task_name="test") - with pytest.raises(DagExecutionError, match="broken-step"): + # 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])) @@ -141,7 +142,9 @@ 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")) + dag.add_job( + Job(Pipeline(Command("echo", "x")), "unexecuted", task_name="test") + ) - with pytest.raises(DagExecutionError, match="unexecuted"): + with pytest.raises(DagExecutionError, match=r"Job\(unexecuted-1\)"): pipeline.check_execution(dag, _StubExecutor()) diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index aaecf0e..20327ae 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -21,7 +21,9 @@ def _job(name, threads=1, resources=None): - return Job(Pipeline(Command(name)), name, threads, resources) + return Job( + Pipeline(Command(name)), name, threads, resources, task_name="test" + ) def test_schedules_independent_jobs_together(): @@ -91,14 +93,22 @@ def test_resource_requirements_ignored_when_unmanaged(): def test_oversized_thread_request_is_rejected(): dag = DAG() - dag.add_job(Job(Pipeline(Command("echo", "hi")), "big", 100)) + 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}) + 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 index 0460403..4454575 100644 --- a/tests/unit/test_scheduler_contract.py +++ b/tests/unit/test_scheduler_contract.py @@ -31,7 +31,9 @@ def _job(name, threads=1, resources=None): - return Job(Pipeline(Command(name)), name, threads, resources) + return Job( + Pipeline(Command(name)), name, threads, resources, task_name="test" + ) def drive(scheduler): From 91c4d66551b116de855ca0d067b9fb4beb620639 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 14 Aug 2026 00:14:09 -0700 Subject: [PATCH 08/58] Capture each process's output to per-task log files Each job gets a JobLogSink on its execution Context, so every process it spawns -- including process substitutions, which previously bypassed stderr redirection entirely -- writes stderr to its own task_logs/{task}/{job_id}.{stage}.log file, beginning with a header recording the task, job, and command. Stdout is captured alongside (headerless, sharing the process's stage index when it would otherwise reach the terminal. Both streams of one process share a stage index, allocated in spawn order. On failure, the run log now reports the failing sub-command with its exit code, pid, log path, and the last 20 lines of its log, plus a pointer to the task-log directory; interrupted and launch-failed jobs list their logs. Successful jobs prune logs that hold nothing but a header. Co-Authored-By: Claude Fable 5 EOF ) --- sentieon_cli/executor.py | 110 +++++++- sentieon_cli/pipeline.py | 11 +- sentieon_cli/run_logs.py | 137 +++++++++- sentieon_cli/shell_pipeline.py | 27 +- tests/integration/test_executor.py | 1 + tests/integration/test_job_log_capture.py | 292 ++++++++++++++++++++++ tests/unit/test_job_logs.py | 203 +++++++++++++++ tests/unit/test_run_logs.py | 25 ++ 8 files changed, 790 insertions(+), 16 deletions(-) create mode 100644 tests/integration/test_job_log_capture.py create mode 100644 tests/unit/test_job_logs.py diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index 66b1aaf..3b9ad58 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -2,20 +2,46 @@ import asyncio import contextlib +import os +import pathlib import signal import sys import threading import time from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, List, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from .job import Job from .logging import get_logger +from .run_logs import JobLogSink, RunLogs from .scheduler import BaseScheduler -from .shell_pipeline import Context +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 + + +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 _signal_procs(context: Context, sig: int) -> None: """Send a signal to every live sub-process in a context.""" @@ -202,12 +228,14 @@ def __init__( *, install_signal_handlers: bool = False, shutdown_grace_period: float = 10.0, + run_logs: Optional[RunLogs] = None, ) -> None: super().__init__( scheduler, install_signal_handlers=install_signal_handlers, ) self.shutdown_grace_period = shutdown_grace_period + self.run_logs = run_logs self.running: List[ Tuple[ Job, @@ -221,13 +249,19 @@ 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, @@ -242,6 +276,8 @@ async def run_job(self, job: Job) -> None: # the caller's pipeline and is left to propagate. logger.error("failed to start command: %s", job.shell) logger.error("Error: %s", str(e)) + # 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 @@ -252,6 +288,47 @@ async def run_job(self, job: Job) -> None: await self._terminate_context(context) await _cleanup_quietly(context) + 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), + ) + + 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. @@ -310,6 +387,7 @@ async def _drive(self, stop_event: asyncio.Event) -> None: # Check if the command failed cmd_failed = False + failures: List[Tuple[Command, int]] = [] for subcommand in context.commands: if not subcommand.proc: logger.error( @@ -337,6 +415,7 @@ async def _drive(self, stop_event: asyncio.Event) -> None: f"{subcommand}" ) cmd_failed = True + failures.append((subcommand, ret)) if ret == -9: logger.error( "Sub-command received SIGKILL. " @@ -355,8 +434,18 @@ async def _drive(self, stop_event: asyncio.Event) -> None: logger.error("Error: %s", str(exc)) cmd_failed = True + sink = context.log_sink 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( + "Task logs for debugging: %s", + self.run_logs.task_logs, + ) self.jobs_with_errors.append(job) self.start_new_jobs = False else: @@ -364,6 +453,8 @@ async def _drive(self, stop_event: asyncio.Event) -> None: f"Finished command in " f"{total_seconds:.2f}: {job.shell}" ) + if sink is not None: + sink.finalize(success=not cmd_failed) if not self.start_new_jobs: # Don't start new jobs @@ -406,8 +497,11 @@ async def _shutdown(self) -> None: if tasks: await asyncio.wait(tasks) - for _job, context, _task, _start in running: + 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/pipeline.py b/sentieon_cli/pipeline.py index 78726c9..65231aa 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -195,7 +195,10 @@ def check_execution( """Check the DAG and executor after a run""" if executor.jobs_with_errors: failed = ", ".join(str(job) for job in executor.jobs_with_errors) - raise DagExecutionError(f"Execution failed for jobs: {failed}") + 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 DagExecutionError( @@ -311,7 +314,11 @@ def run(self, dag: DAG) -> BaseExecutor: 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) + 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 index 2cfd5c8..5a7f7c3 100644 --- a/sentieon_cli/run_logs.py +++ b/sentieon_cli/run_logs.py @@ -2,17 +2,148 @@ The log directory for a single run """ +from __future__ import annotations + +import dataclasses import datetime import logging import pathlib +import re import shlex import shutil import sys -from typing import Optional +from typing import IO, List, Optional, TYPE_CHECKING from .logging import add_file_handler, remove_file_handler from .util import __version__ +if TYPE_CHECKING: + from .job import Job + from .shell_pipeline import Command + +_UNSAFE = re.compile(r"[^A-Za-z0-9._-]") + +_STDERR = "stderr" +_STDOUT = "stdout" + + +def sanitize(component: str) -> str: + """Restrict a path component to filesystem-safe characters""" + return _UNSAFE.sub("-", component) + + +@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 _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""" @@ -39,6 +170,10 @@ def create_dirs(self) -> None: shutil.rmtree(self.task_logs) self.task_logs.mkdir(parents=True) + def job_sink(self, job: Job) -> JobLogSink: + """Create the log sink for a job""" + return JobLogSink(self.task_logs, job) + def write_command(self) -> None: """Record the invocation so the run can be reproduced""" timestamp = datetime.datetime.now().astimezone() diff --git a/sentieon_cli/shell_pipeline.py b/sentieon_cli/shell_pipeline.py index 1d6d26d..fa6aa7a 100644 --- a/sentieon_cli/shell_pipeline.py +++ b/sentieon_cli/shell_pipeline.py @@ -8,16 +8,18 @@ import asyncio import dataclasses import fcntl -import io import os import pathlib import tempfile import shlex from abc import ABC, abstractmethod -from typing import Any, Dict, IO, List, Optional, Union +from typing import Any, Dict, IO, List, Optional, TYPE_CHECKING, Union from .logging import get_logger +if TYPE_CHECKING: + from .run_logs import JobLogSink + logger = get_logger(__name__) @@ -64,7 +66,7 @@ class _ProcSubWait: 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) @@ -72,7 +74,9 @@ def __init__(self) -> None: # 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 @@ -169,7 +173,20 @@ 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, diff --git a/tests/integration/test_executor.py b/tests/integration/test_executor.py index d0b1b6d..41ce529 100644 --- a/tests/integration/test_executor.py +++ b/tests/integration/test_executor.py @@ -518,6 +518,7 @@ class _FakeContext: def __init__(self, cleanup_error=None): self.commands = [] self.cleaned = False + self.log_sink = None self._error = cleanup_error async def cleanup(self): 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/unit/test_job_logs.py b/tests/unit/test_job_logs.py new file mode 100644 index 0000000..e438533 --- /dev/null +++ b/tests/unit/test_job_logs.py @@ -0,0 +1,203 @@ +""" +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, sanitize # noqa: E402 +from sentieon_cli.shell_pipeline import ( # noqa: E402 + Command, + Context, + Pipeline, +) + + +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_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_run_logs.py b/tests/unit/test_run_logs.py index a16fb1b..02b8eed 100644 --- a/tests/unit/test_run_logs.py +++ b/tests/unit/test_run_logs.py @@ -19,8 +19,11 @@ 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" @@ -46,6 +49,13 @@ def build_dag(self) -> DAG: raise RuntimeError("boom") +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.""" @@ -210,6 +220,21 @@ def test_bare_pipeline_setup_logging_creates_nothing(tmp_path, monkeypatch): assert list(tmp_path.iterdir()) == [] +def test_check_execution_names_the_task_log_dir(tmp_path): + pipeline = _DummyPipeline() + pipeline.log_dir = tmp_path / "logs" + pipeline.setup_logging(_args()) + 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() From 87f54b2743efdd18c56027f9ddb877c09ee7775a Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 14 Aug 2026 12:54:53 -0700 Subject: [PATCH 09/58] Fix log-directory setup ordering, id sanitization, and verbosity flags Create the log directory only after validation: setup_logging is now console-only and a new start_run_logs(), called after validate(), derives and prepares the log directory. A typo'd output path again fails fast instead of being silently created, a rerun rejected by validation no longer destroys the previous run's logs, and a log path colliding with a regular file exits cleanly instead of raising. Key job-id counters on the sanitized name (sanitize moved to util.py) so log file paths stay unique even for names that differ only in unsafe characters. Make -v/-q/-d independent flags with explicit precedence: -d -v no longer downgrades DEBUG to INFO, and the new --quiet restores a warnings-only console. Document the flags and the per-run log directory in the README. Co-Authored-By: Claude Fable 5 --- README.md | 8 ++- sentieon_cli/__init__.py | 40 ++++++++++---- sentieon_cli/job.py | 8 ++- sentieon_cli/logging.py | 3 +- sentieon_cli/pipeline.py | 46 ++++++++++------ sentieon_cli/run_logs.py | 10 +--- sentieon_cli/sentieon_pangenome.py | 1 + sentieon_cli/util.py | 7 +++ tests/unit/test_cli_args.py | 40 ++++++++++++++ tests/unit/test_job.py | 12 +++++ tests/unit/test_job_logs.py | 3 +- tests/unit/test_run_logs.py | 87 ++++++++++++++++++++++++++---- 12 files changed, 213 insertions(+), 52 deletions(-) create mode 100644 tests/unit/test_cli_args.py diff --git a/README.md b/README.md index b51d89f..cccd5fc 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/`. 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/sentieon_cli/__init__.py b/sentieon_cli/__init__.py index 1de7890..59c9971 100644 --- a/sentieon_cli/__init__.py +++ b/sentieon_cli/__init__.py @@ -1,3 +1,5 @@ +import argparse + from . import argh_parser from .dnascope import DNAscopePipeline from .dnascope_hybrid import DNAscopeHybridPipeline @@ -7,26 +9,41 @@ 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="INFO", + 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", @@ -60,6 +77,7 @@ def main(): dnascope_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() diff --git a/sentieon_cli/job.py b/sentieon_cli/job.py index bb893bb..a6a15f7 100644 --- a/sentieon_cli/job.py +++ b/sentieon_cli/job.py @@ -5,6 +5,7 @@ from typing import Dict, Optional from .shell_pipeline import Pipeline +from .util import sanitize class Job: @@ -44,8 +45,11 @@ def __init__( self.threads = threads self.resources = {} if resources is None else resources self.task_name = task_name - count = Job._id_counters.get(name, 0) + 1 - Job._id_counters[name] = count + # 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 diff --git a/sentieon_cli/logging.py b/sentieon_cli/logging.py index df0a48a..fe12b53 100644 --- a/sentieon_cli/logging.py +++ b/sentieon_cli/logging.py @@ -4,7 +4,8 @@ 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``/``-d`` while the ``run.log`` file handler always records DEBUG. +follows ``-v``/``-q``/``-d`` while the ``run.log`` file handler always records +DEBUG. """ import logging diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 65231aa..55f267a 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -126,30 +126,43 @@ def __init__(self) -> None: self.run_logs: Optional[RunLogs] = None def setup_logging(self, args: argparse.Namespace) -> None: - """Configure the console handler and the run's log directory""" + """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 not self.dry_run and ( - self.log_dir is not None or self.output_vcf is not None - ): - log_dir = self.log_dir - if log_dir is None: - # Check the suffix before deriving the directory name so an - # invalid output path cannot create a garbage-named log dir. - self.validate_output_suffix() - log_dir = pathlib.Path( - str(self.output_vcf).removesuffix(".vcf.gz") + "_logs" - ) - self.run_logs = RunLogs(log_dir) - self.run_logs.setup() + 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__) - if self.run_logs: - self.logger.info("Writing logs to: %s", self.run_logs.log_dir) + 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""" @@ -171,6 +184,7 @@ def main(self, args: argparse.Namespace) -> None: success = False try: self.validate() + self.start_run_logs() self.configure() tmp_dir_str = tmp() diff --git a/sentieon_cli/run_logs.py b/sentieon_cli/run_logs.py index 5a7f7c3..adee9db 100644 --- a/sentieon_cli/run_logs.py +++ b/sentieon_cli/run_logs.py @@ -8,30 +8,22 @@ import datetime import logging import pathlib -import re import shlex import shutil import sys from typing import IO, List, Optional, TYPE_CHECKING from .logging import add_file_handler, remove_file_handler -from .util import __version__ +from .util import __version__, sanitize if TYPE_CHECKING: from .job import Job from .shell_pipeline import Command -_UNSAFE = re.compile(r"[^A-Za-z0-9._-]") - _STDERR = "stderr" _STDOUT = "stdout" -def sanitize(component: str) -> str: - """Restrict a path component to filesystem-safe characters""" - return _UNSAFE.sub("-", component) - - @dataclasses.dataclass class _ProcessLog: """A log file opened for one spawned process""" diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 50f23af..bb985db 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -284,6 +284,7 @@ def main(self, args: argparse.Namespace) -> None: self.logger.debug("VCF contigs are: %s", self.pop_vcf_contigs) self.validate() + self.start_run_logs() self.shards = determine_shards_from_fai( self.fai_data, 10 * 1000 * 1000 ) diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index 68021ea..fe9cc73 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -31,6 +31,13 @@ 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._-]") + + +def sanitize(component: str) -> str: + """Restrict a path component to filesystem-safe characters""" + return _UNSAFE.sub("-", component) + def tmp(): """Create a temporary directory for the current process.""" 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_job.py b/tests/unit/test_job.py index 4b37a58..f77205e 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -16,6 +16,7 @@ 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"): @@ -45,6 +46,17 @@ def test_each_name_has_its_own_counter(): 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" diff --git a/tests/unit/test_job_logs.py b/tests/unit/test_job_logs.py index e438533..29395bc 100644 --- a/tests/unit/test_job_logs.py +++ b/tests/unit/test_job_logs.py @@ -13,12 +13,13 @@ 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, sanitize # 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"): diff --git a/tests/unit/test_run_logs.py b/tests/unit/test_run_logs.py index 02b8eed..a9c84ff 100644 --- a/tests/unit/test_run_logs.py +++ b/tests/unit/test_run_logs.py @@ -49,6 +49,13 @@ 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.""" @@ -81,6 +88,12 @@ 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") @@ -92,7 +105,7 @@ def test_module_loggers_delegate_to_the_package_logger(): def test_default_log_dir_is_derived_from_the_output_vcf(tmp_path): pipeline = _DummyPipeline() pipeline.output_vcf = tmp_path / "sample.vcf.gz" - pipeline.setup_logging(_args()) + _start(pipeline) assert pipeline.run_logs is not None assert pipeline.run_logs.log_dir == tmp_path / "sample_logs" @@ -102,7 +115,7 @@ 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" - pipeline.setup_logging(_args()) + _start(pipeline) assert pipeline.run_logs.log_dir == tmp_path / "a.vcf.gz.rerun_logs" @@ -111,7 +124,7 @@ 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" - pipeline.setup_logging(_args()) + _start(pipeline) assert pipeline.run_logs.log_dir == tmp_path / "elsewhere" assert (tmp_path / "elsewhere" / "run.log").is_file() @@ -122,10 +135,64 @@ def test_invalid_output_suffix_exits_before_creating_a_log_dir(tmp_path): pipeline.output_vcf = tmp_path / "sample.bcf" with pytest.raises(SystemExit) as excinfo: - pipeline.setup_logging(_args()) + _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): @@ -160,7 +227,7 @@ def test_command_txt_records_the_invocation(tmp_path, monkeypatch): def test_run_log_captures_debug_while_the_console_stays_at_info(tmp_path): pipeline = _DummyPipeline() pipeline.output_vcf = tmp_path / "sample.vcf.gz" - pipeline.setup_logging(_args("INFO")) + _start(pipeline) logging.getLogger("sentieon_cli.example").debug("a debug record") run_log = pipeline.run_logs.run_log pipeline.run_logs.close() @@ -180,7 +247,7 @@ def test_repeated_setup_does_not_accumulate_handlers(tmp_path): for i in range(3): pipeline = _DummyPipeline() pipeline.output_vcf = tmp_path / f"sample{i}.vcf.gz" - pipeline.setup_logging(_args()) + _start(pipeline) assert len(package_logger.handlers) == before + 1 pipeline.run_logs.close() @@ -193,7 +260,7 @@ def test_rerun_truncates_the_previous_run_log(tmp_path): for _ in range(2): pipeline = _DummyPipeline() pipeline.log_dir = log_dir - pipeline.setup_logging(_args()) + _start(pipeline) run_log = pipeline.run_logs.run_log pipeline.run_logs.close() @@ -205,7 +272,7 @@ def test_dry_run_skips_file_logging(tmp_path): pipeline = _DummyPipeline() pipeline.dry_run = True pipeline.output_vcf = tmp_path / "sample.vcf.gz" - pipeline.setup_logging(_args()) + _start(pipeline) assert pipeline.run_logs is None assert list(tmp_path.iterdir()) == [] @@ -214,7 +281,7 @@ def test_dry_run_skips_file_logging(tmp_path): def test_bare_pipeline_setup_logging_creates_nothing(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) pipeline = _DummyPipeline() - pipeline.setup_logging(_args()) + _start(pipeline) assert pipeline.run_logs is None assert list(tmp_path.iterdir()) == [] @@ -223,7 +290,7 @@ def test_bare_pipeline_setup_logging_creates_nothing(tmp_path, monkeypatch): def test_check_execution_names_the_task_log_dir(tmp_path): pipeline = _DummyPipeline() pipeline.log_dir = tmp_path / "logs" - pipeline.setup_logging(_args()) + _start(pipeline) job = Job(Pipeline(Command("false")), "boom", task_name="failing") with pytest.raises(DagExecutionError) as excinfo: From 0d4150de2a42787b71a60265bfc89c5695278227 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 14 Aug 2026 14:45:56 -0700 Subject: [PATCH 10/58] Record per-process rusage by owning the child reap Spawn pipeline processes with RusagePopen, a subprocess.Popen subclass whose _try_wait() reaps with os.wait4() and records the child's rusage, replacing asyncio.create_subprocess_exec (which only provided spawn and reap here: every stdio stream is already a real FD). Waits move to async_wait() (Popen.wait in a worker thread), and the executor raises the loop's default thread pool to 512 lazily-created workers so blocked waits and proc-sub FIFO opens cannot starve each other. Context.cleanup() gains a non-blocking poll() sweep so children nobody awaited are still reaped, as asyncio's watcher used to do. Behavior parity otherwise: negative returncodes on signal death, SIGPIPE tolerance, and the signal/kill paths are unchanged. Metrics emission lands separately. Co-Authored-By: Claude Fable 5 --- sentieon_cli/executor.py | 16 +++++++-- sentieon_cli/shell_pipeline.py | 58 ++++++++++++++++++++++++++----- tests/unit/test_shell_pipeline.py | 24 ++++++------- 3 files changed, 75 insertions(+), 23 deletions(-) diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index 3b9ad58..a84388a 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -1,6 +1,7 @@ """Execute jobs""" import asyncio +import concurrent.futures import contextlib import os import pathlib @@ -201,6 +202,15 @@ async def _execute(self) -> None: """Execute jobs from the DAG""" self.jobs_with_errors = [] loop = asyncio.get_running_loop() + # Blocked wait() threads and the proc-sub FIFO-open threads share the + # default pool, so it must be large enough that neither starves the + # other. Threads are created lazily, so the 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=512, thread_name_prefix="sentieon-wait" + ) + ) stop_event = asyncio.Event() restore = self._install_signal_handlers(loop, stop_event) try: @@ -266,7 +276,7 @@ async def run_job(self, job: Job) -> None: ( job, context, - asyncio.create_task(proc.wait()), + asyncio.create_task(proc.async_wait()), start_time, ) ) @@ -346,7 +356,7 @@ async def _terminate_context(self, context: Context) -> None: if not live: return _signal_procs(context, signal.SIGTERM) - waits = [asyncio.create_task(proc.wait()) for proc in live] + 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) @@ -397,7 +407,7 @@ async def _drive(self, stop_event: asyncio.Event) -> None: cmd_failed = True continue ret = ( - await subcommand.proc.wait() + await subcommand.proc.async_wait() ) # Wait on all sub-commands # A subcommand killed by SIGPIPE is not a failure: # it was writing to a pipe whose reader exited early diff --git a/sentieon_cli/shell_pipeline.py b/sentieon_cli/shell_pipeline.py index fa6aa7a..2020954 100644 --- a/sentieon_cli/shell_pipeline.py +++ b/sentieon_cli/shell_pipeline.py @@ -10,10 +10,21 @@ import fcntl import os import pathlib +import resource +import subprocess import tempfile import shlex from abc import ABC, abstractmethod -from typing import Any, Dict, IO, List, Optional, TYPE_CHECKING, Union +from typing import ( + Any, + Dict, + IO, + List, + Optional, + Tuple, + TYPE_CHECKING, + Union, +) from .logging import get_logger @@ -124,6 +135,14 @@ async def cleanup(self) -> None: 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() @@ -135,6 +154,29 @@ async def cleanup(self) -> None: raise exc +class RusagePopen(subprocess.Popen[bytes]): + """A Popen that records the child's rusage when it is reaped.""" + + 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.""" + return await asyncio.to_thread(self.wait) + + class ShellNode(ABC): """Abstract base class for any node in the shell syntax tree.""" @@ -152,7 +194,7 @@ def __init__( 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( @@ -161,7 +203,7 @@ async def run( stdin: Union[IO[Any], int, None] = None, stdout: Union[IO[Any], int, None] = None, stderr: Union[IO[Any], int, None] = None, - ) -> asyncio.subprocess.Process: + ) -> RusagePopen: # 1. Resolve Arguments (Handle Process Substitutions) final_args = [self.executable] @@ -188,8 +230,8 @@ async def run( # 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, @@ -266,7 +308,7 @@ async def run( stdin: Union[IO[Any], int, None] = None, stdout: Union[IO[Any], int, None] = None, stderr: Union[IO[Any], int, None] = None, - ) -> asyncio.subprocess.Process: + ) -> 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( @@ -429,7 +471,7 @@ async def run_inner() -> None: context, stdout=write_handle, ) - await proc.wait() + await proc.async_wait() finally: write_handle.close() @@ -467,7 +509,7 @@ async def run_inner() -> None: 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() diff --git a/tests/unit/test_shell_pipeline.py b/tests/unit/test_shell_pipeline.py index 38ec2bd..c17b9b8 100644 --- a/tests/unit/test_shell_pipeline.py +++ b/tests/unit/test_shell_pipeline.py @@ -34,7 +34,7 @@ async def test_simple_command(): 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() @@ -56,7 +56,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() @@ -89,7 +89,7 @@ async def test_pipeline_with_file_io(): ) proc = await pipeline.run(context) - await proc.wait() + await proc.async_wait() with open(outfile_path, "r") as f: output = f.read().strip() @@ -114,7 +114,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() @@ -149,7 +149,7 @@ async def test_output_process_substitution(): 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) @@ -198,10 +198,10 @@ async def test_pipe_size_enlarges_internal_pipes(tmp_path): pipe_size=target, ) proc = await pipeline.run(context) - await proc.wait() + await proc.async_wait() for sub in context.commands: if sub.proc: - await sub.proc.wait() + await sub.proc.async_wait() await context.cleanup() assert int(report.read_text()) == target @@ -347,7 +347,7 @@ async def test_cleanup_after_outer_exits_without_opening(): context = Context() cmd = Command("false", InputProcSub(Pipeline(Command("echo", "x")))) proc = await cmd.run(context) - await proc.wait() + await proc.async_wait() await asyncio.wait_for(context.cleanup(), timeout=10) @@ -364,7 +364,7 @@ async def test_cleanup_unblocks_multiple_procsubs(): OutputProcSub(Pipeline(Command("cat"))), ) proc = await cmd.run(context) - await proc.wait() + await proc.async_wait() await asyncio.wait_for(context.cleanup(), timeout=10) @@ -380,7 +380,7 @@ async def test_cleanup_still_raises_inner_launch_failure(): "cat", InputProcSub(Pipeline(Command("no_such_cmd_zzz"))) ) proc = await cmd.run(context) - await proc.wait() + await proc.async_wait() with pytest.raises(FileNotFoundError): await asyncio.wait_for(context.cleanup(), timeout=10) @@ -408,7 +408,7 @@ async def test_cleanup_does_not_stall_running_inner_writer(): await asyncio.sleep(0.01) await asyncio.wait_for(context.cleanup(), timeout=10) - assert await proc.wait() == 0 + assert await proc.async_wait() == 0 @pytest.mark.asyncio @@ -417,7 +417,7 @@ async def test_cleanup_twice_is_safe(): context = Context() cmd = Command("cat", InputProcSub(Pipeline(Command("echo", "hello")))) proc = await cmd.run(context, stdout=asyncio.subprocess.DEVNULL) - await proc.wait() + await proc.async_wait() await asyncio.wait_for(context.cleanup(), timeout=10) await asyncio.wait_for(context.cleanup(), timeout=10) From be5b184ddad78e9704f9485dc3697a2d72475350 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 14 Aug 2026 15:08:13 -0700 Subject: [PATCH 11/58] Report per-process rusage in the run logs Each reaped process now logs a DEBUG line to run.log with its kernel resource usage (utime, stime, peak RSS, page faults, block I/O, context switches), emitted whatever the job's outcome -- signal-killed children included. Successful jobs extend the "Finished command" INFO line with the job's summed user/sys CPU and the per-process peak RSS (a sum would overstate concurrent stages that peak at different times). ru_maxrss is normalized across platforms (bytes on macOS, KiB elsewhere), and the DEBUG line renders it without an internal space (maxrss=33.6MiB) so the key=value line stays space-tokenizable for grep/awk. Co-Authored-By: Claude Fable 5 --- sentieon_cli/executor.py | 63 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index a84388a..0e8f64e 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -5,6 +5,7 @@ import contextlib import os import pathlib +import resource import signal import sys import threading @@ -44,6 +45,22 @@ def _log_tail(path: pathlib.Path, lines: int = TAIL_LINES) -> List[str]: 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: @@ -398,6 +415,14 @@ async def _drive(self, stop_event: asyncio.Event) -> None: # 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 for subcommand in context.commands: if not subcommand.proc: logger.error( @@ -409,6 +434,34 @@ async def _drive(self, stop_event: asyncio.Event) -> None: 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 + if rusage is not None: + rss = _maxrss_bytes(rusage) + 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, + ) # 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 @@ -458,6 +511,16 @@ async def _drive(self, stop_event: asyncio.Event) -> None: ) 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 " From fb9fa2c6c5dcebc382fe2d1d5fe7e3ba58f1faeb Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 14 Aug 2026 15:21:03 -0700 Subject: [PATCH 12/58] Test rusage capture and the per-process metrics report test_simple_command now asserts rusage is recorded at reap, and a new signal-death test pins the RusagePopen._try_wait override: if a future CPython changes that private method's shape, the override silently stops being called and rusage stays None, so the test asserts a SIGTERM'd child reports both its negative return code and its rusage. New tests/unit/test_rusage_metrics.py drives a job through LocalExecutor (the first unit-level file to do so) and checks the per-process "rusage for" DEBUG record, the aggregated "Finished command" INFO line, and _maxrss_bytes unit normalization on both platform branches. Co-Authored-By: Claude Fable 5 --- tests/unit/test_rusage_metrics.py | 99 +++++++++++++++++++++++++++++++ tests/unit/test_shell_pipeline.py | 24 ++++++++ 2 files changed, 123 insertions(+) create mode 100644 tests/unit/test_rusage_metrics.py 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_shell_pipeline.py b/tests/unit/test_shell_pipeline.py index c17b9b8..a6d704c 100644 --- a/tests/unit/test_shell_pipeline.py +++ b/tests/unit/test_shell_pipeline.py @@ -6,6 +6,7 @@ import fcntl import os import pathlib +import signal import sys import tempfile @@ -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 From 76394ba599ebea9d4b33052fab4f3fcff803c8d8 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 14 Aug 2026 16:42:43 -0700 Subject: [PATCH 13/58] Tear down children when the run loop fails; make the wait pool tunable Two fixes from a code review of the branch. First, an exception escaping _drive now tears the children down (_shutdown) before propagating: it previously unwound into asyncio.run's shutdown, which joins the pool threads blocked in wait4 on children nobody signalled, hanging the CLI until the children exited on their own. BaseException, so cancellation and KeyboardInterrupt are covered; a teardown failure is logged rather than allowed to mask the original error. Second, the wait/FIFO-open thread pool size is a keyword parameter instead of a hidden constant: None on AsyncExecutor (an executor that spawns nothing leaves the loop's default pool alone) and 512 on LocalExecutor, beside shutdown_grace_period. async_wait() returns immediately for an already-reaped child instead of burning a pool dispatch and a slot. Also document a known limitation honestly: a reap via poll() -- or send_signal/kill, which poll internally -- uses the base Popen machinery and records no rusage. The executor waits every process it reports metrics for, so those paths only collect processes whose rusage is never read. Co-Authored-By: Claude Fable 5 --- sentieon_cli/executor.py | 50 ++++++++++++++++++++++------ sentieon_cli/shell_pipeline.py | 13 +++++++- tests/integration/test_executor.py | 53 ++++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 11 deletions(-) diff --git a/sentieon_cli/executor.py b/sentieon_cli/executor.py index 0e8f64e..2fecf6f 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -144,6 +144,10 @@ class AsyncExecutor(BaseExecutor, ABC): 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. + + ``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__( @@ -151,9 +155,11 @@ def __init__( 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( @@ -219,19 +225,34 @@ async def _execute(self) -> None: """Execute jobs from the DAG""" self.jobs_with_errors = [] loop = asyncio.get_running_loop() - # Blocked wait() threads and the proc-sub FIFO-open threads share the - # default pool, so it must be large enough that neither starves the - # other. Threads are created lazily, so the 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=512, thread_name_prefix="sentieon-wait" + 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: - await self._drive(stop_event) + 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: @@ -247,7 +268,14 @@ async def _shutdown(self) -> None: class LocalExecutor(AsyncExecutor): - """Run jobs locally as async subprocesses.""" + """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, @@ -256,10 +284,12 @@ def __init__( 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 diff --git a/sentieon_cli/shell_pipeline.py b/sentieon_cli/shell_pipeline.py index 2020954..d2d1277 100644 --- a/sentieon_cli/shell_pipeline.py +++ b/sentieon_cli/shell_pipeline.py @@ -155,7 +155,15 @@ async def cleanup(self) -> None: class RusagePopen(subprocess.Popen[bytes]): - """A Popen that records the child's rusage when it is reaped.""" + """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 @@ -174,6 +182,9 @@ def _try_wait(self, wait_flags: int) -> Tuple[int, int]: 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) diff --git a/tests/integration/test_executor.py b/tests/integration/test_executor.py index 41ce529..d0ea7f5 100644 --- a/tests/integration/test_executor.py +++ b/tests/integration/test_executor.py @@ -8,6 +8,7 @@ import sys import tempfile import threading +import time import pytest @@ -475,6 +476,58 @@ def __init__(self) -> None: 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 From 58c9c4d6a6150c0ccf16d73b49232bea3a8b049c Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 15 Aug 2026 15:27:30 -0700 Subject: [PATCH 14/58] Add the hybrid-pangenome pipeline Combine the human pangenome with sample short-read fastq and aligned long-read data for variant calling: k-mer counting across both read sets, personalized diploid graph construction with vg, graph updates with PGHapUpdateAlgo using the long-read alignments and LongReadSV calls, re-alignment of extracted short reads to the personalized pangenome with liftover, short-read dedup and metrics, and PangenomeSV + DNAscope calling with LR readgroup attributes. The dedup and metrics job builders move from SentieonPangenome to BasePangenome so both pangenome pipelines share them; the emitted dnascope-pangenome commands are unchanged. Co-Authored-By: Claude Fable 5 --- sentieon_cli/__init__.py | 7 + sentieon_cli/base_pangenome.py | 145 ++++- sentieon_cli/command_strings.py | 116 ++++ sentieon_cli/driver.py | 20 + sentieon_cli/hybrid_pangenome.py | 911 ++++++++++++++++++++++++++++ sentieon_cli/sentieon_pangenome.py | 144 +---- tests/unit/test_hybrid_pangenome.py | 474 +++++++++++++++ 7 files changed, 1676 insertions(+), 141 deletions(-) create mode 100644 sentieon_cli/hybrid_pangenome.py create mode 100644 tests/unit/test_hybrid_pangenome.py diff --git a/sentieon_cli/__init__.py b/sentieon_cli/__init__.py index 59c9971..bffbb96 100644 --- a/sentieon_cli/__init__.py +++ b/sentieon_cli/__init__.py @@ -4,6 +4,7 @@ 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__ @@ -76,6 +77,12 @@ 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 diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index 7284de6..c126bbd 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -6,13 +6,28 @@ from enum import Enum import json import pathlib -from typing import List, Optional +import sys +from typing import List, Optional, Tuple from importlib.resources import files from . import command_strings as cmds +from .driver import ( + AlignmentStat, + BaseDistributionByCycle, + CoverageMetrics, + Dedup, + Driver, + GCBias, + InsertSizeMetricAlgo, + LocusCollector, + MeanQualityByCycle, + QualDistribution, + WgsMetricsAlgo, +) from .job import Job from .pipeline import BasePipeline +from .shell_pipeline import Command, Pipeline from .util import path_arg @@ -155,6 +170,134 @@ def build_ploidy_job( ) return ploidy_job + def build_dedup_job( + self, + output_bam, + input_bam: List[pathlib.Path], + tag: str, + left_align_rgid: Optional[str] = None, + 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_rgid: + read_filters.append( + f"IndelLeftAlignReadTransform,rgid={left_align_rgid}" + ) + + # 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, + task_name="dedup", + ) + + 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, + task_name="dedup", + ) + + 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, + task_name="metrics", + ) + + 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, + task_name="metrics", + ) + return (metrics_job, rehead_job) + def get_sex(self, ploidy_json: pathlib.Path) -> None: """Retrieve the sample sex""" if self.dry_run: diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 2595fe1..e09877a 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1415,6 +1415,122 @@ def cmd_minimap2_lift( return Pipeline(mm2_cmd, lift_cmd, sort_cmd) +def cmd_hybrid_kmc( + output_prefix: pathlib.Path, + fastq: List[pathlib.Path], + aln: List[pathlib.Path], + reference: 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. + """ + 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)) + for aln_file in aln: + cat_args.append( + InputProcSub( + Pipeline( + Command( + "samtools", + "fasta", + "--reference", + str(reference), + "-@", + str(threads), + str(aln_file), + ) + ) + ) + ) + cat_cmd = Command("cat", *cat_args) + kmc_cmd = Command( + "kmc", + f"-k{k}", + f"-m{memory}", + "-okff", + f"-t{threads}", + "-fa", + "/dev/stdin", + str(output_prefix), + str(tmp_dir), + ) + 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, +) -> 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("sort", "-k1,1", "-k2,2n") + 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/driver.py b/sentieon_cli/driver.py index c2bbb7f..e39fe80 100644 --- a/sentieon_cli/driver.py +++ b/sentieon_cli/driver.py @@ -542,6 +542,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""" diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py new file mode 100644 index 0000000..24ef897 --- /dev/null +++ b/sentieon_cli/hybrid_pangenome.py @@ -0,0 +1,911 @@ +""" +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 shutil +import sys +from typing import Dict, List, Optional, Set + +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 ( + DNAModelApply, + DNAscope, + Driver, + LongReadSV, + PangenomeSV, + PGHapUpdateAlgo, +) +from .job import Job +from .logging import get_logger +from .shard import ( + GRCH38_CONTIGS, + determine_shards_from_fai, + parse_fai, + vcf_contigs, +) +from .shell_pipeline import Command, Pipeline +from .transfer import build_transfer_jobs +from .util import ( + __version__, + check_kmc_patch, + check_version, + parse_rg_line, + path_arg, + total_memory, + vcf_id, +) + +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, +} + +# 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.", + }, + # Additional arguments + "bed": { + "flags": ["-b", "--bed"], + "help": ( + "Region BED file. Supplying this file will limit " + "small-variant calling to the intervals inside the BED " + "file." + ), + "type": path_arg(exists=True, is_file=True), + }, + "pangenome_ref_name": { + "default": "GRCh38", + "help": "Reference name in the pangenome (GRCh38).", + }, + "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", + }, + # Hidden arguments + "skip_contig_checks": { + "help": argparse.SUPPRESS, + "action": "store_true", + }, + "skip_model_apply": { + "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", + }, + "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.pangenome_ref_name = "GRCh38" + self.rgsm: Optional[str] = None + self.extract_model_name = "extract.model" + self.skip_metrics = False + self.skip_multiqc = False + self.skip_contig_checks = False + self.skip_model_apply = False + self.skip_pangenome_name_checks = False + self.skip_pop_vcf_id_check = False + self.skip_small_variants = False + self.skip_svs = False + + 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 + ) + 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) + + self.validate_bundle() + self.validate_output_vcf() + + if not self.r1_fastq or not self.readgroup: + self.logger.error( + "Please supply the short reads with the `--r1_fastq`, " + "`--r2_fastq`, and `--readgroup` arguments" + ) + 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) + if not self.lr_aln: + self.logger.error( + "Please supply the long-read alignments with the `--lr_aln` " + "argument" + ) + sys.exit(2) + + self.validate_bwa_index() + self.collect_readgroups() + self.validate_readgroups() + + if not self.skip_version_check: + for cmd, min_version in HYBRID_PANGENOME_MIN_VERSIONS.items(): + if not check_version(cmd, min_version): + sys.exit(2) + + 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( + "A BED file is recommended to avoid small-variant calling " + "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 are 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) + + 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 = { + "bwa.model", + "dnascope.model", + "longreadsv.model", + "minimap2.model", + self.extract_model_name, + } + 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) + + bundle_vcf_id = bundle_info.get("SentieonVcfID") + if ( + bundle_vcf_id + and not self.skip_pop_vcf_id_check + and not self.dry_run + ): + assert self.pop_vcf is not None + 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) + + def collect_readgroups(self) -> None: + """Collect readgroup tags from the inputs""" + assert self.readgroup is not None + 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: Dict[str, str] = parsed_rg + + self.lr_readgroups: List[List[Dict[str, str]]] = [] + for aln in self.lr_aln: + self.lr_readgroups.append([]) + for rg_line in cmds.get_rg_lines(aln, self.dry_run): + self.lr_readgroups[-1].append(parse_rg_line(rg_line)) + + def validate_readgroups(self) -> None: + """Confirm that all readgroups have a consistent SM tag""" + rg_sm = self.fastq_readgroup.get("SM") + for aln, aln_rgs in zip(self.lr_aln, self.lr_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 `--readgroup` argument.\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) + self.sample_sm: str = self.rgsm if self.rgsm else str(rg_sm) + + def configure(self) -> None: + """Configure pipeline parameters""" + pass + + def find_unzip(self) -> str: + """The decompression tool for fastq input""" + 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" + return unzip + + def build_dag(self) -> DAG: + """Build the DAG for the hybrid-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) + if not self.pop_vcf: + self.logger.error("pop_vcf is required") + sys.exit(2) + + self.logger.info("Building the hybrid-pangenome DAG") + dag = DAG() + + ref_fai = pathlib.Path(str(self.reference) + ".fai") + + # 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}") + ) + out_lift_aln = pathlib.Path( + str(self.output_vcf).replace(".vcf.gz", f"_lift_deduped.{suffix}") + ) + sv_vcf = pathlib.Path( + str(self.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") + + total_mem_gb = total_memory() / (1024.0**3) + + # 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: Set[Job] = {kmc_job} + + # BWA alignment and extraction + bwa_job = self.build_alignment_job(bwa_bam, ext_fastq) + dag.add_job(bwa_job) + # Do not run vg-haplotypes with bwa in low-mem environments + if total_mem_gb < 70: + haplotype_dependencies.add(bwa_job) + + # vg haplotypes - create a sample-specific pangenome + haplotypes_job = self.build_haplotypes_job(sample_pangenome, kmer_file) + dag.add_job(haplotypes_job, haplotype_dependencies) + + # convert the sample pangenome to GFA + gfa_job = self.build_gfa_job(hap_raw_gfa, sample_pangenome) + dag.add_job(gfa_job, {haplotypes_job}) + + # Graph update without the SV BED + update_raw_job = self.build_graph_update_job( + pangenome_raw_gfa, + hap_raw_gfa, + name="graph-update-raw", + ) + dag.add_job(update_raw_job, {gfa_job}) + + # Call SVs from the long reads and collect graph update regions + longreadsv_job = self.build_longreadsv_job(longread_sv_vcf) + dag.add_job(longreadsv_job) + sv_bed_job = Job( + cmds.cmd_longread_sv_bed(sv_bed, longread_sv_vcf), + "longread-sv-bed", + 0, + task_name="pangenome-update", + ) + dag.add_job(sv_bed_job, {longreadsv_job}) + + # Graph update with the SV BED + update_job = self.build_graph_update_job( + pangenome_gfa, + pangenome_raw_gfa, + 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 = Job( + cmds.cmd_pgutil_gfa2fa(pangenome_fasta, ref_fai, pangenome_gfa), + "gfa2fa", + 0, + task_name="pangenome", + ) + 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_bam, + ext_fastq, + pangenome_fasta, + pangenome_gfa, + ) + dag.add_job(mm2_job, {bwa_job, faidx_job, update_job}) + + # 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. + 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 + ) + lift_lc_job, lift_dedup_job = self.build_dedup_job( + out_lift_aln, + [lift_bam], + "lift", + ) + dag.add_job(bwa_lc_job, {bwa_job}) + dag.add_job(bwa_dedup_job, {bwa_lc_job}) + dag.add_job(lift_lc_job, {mm2_job}) + dag.add_job(lift_dedup_job, {lift_lc_job}) + + # Alignment metrics from the deduplicated short reads + if not self.skip_metrics: + metrics_job, rehead_job = self.build_metrics_job( + [out_bwa_aln, out_lift_aln], + ) + dag.add_job(metrics_job, {bwa_dedup_job, lift_dedup_job}) + dag.add_job(rehead_job, {metrics_job}) + if not self.skip_multiqc: + multiqc_job = self.multiqc() + if multiqc_job: + dag.add_job(multiqc_job, {rehead_job}) + + # Variant calling with the original, lifted, and long reads + calling_bams = [out_bwa_aln, out_lift_aln] + list(self.lr_aln) + replace_rg = self.build_replace_rg() + calling_dependencies = {bwa_dedup_job, lift_dedup_job} + + 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 + + dnascope_job = self.build_dnascope_job( + raw_vcf, + calling_bams, + replace_rg, + ) + dag.add_job(dnascope_job, calling_dependencies) + + # transfer annotations from the pop_vcf + transfer_target = ( + transfer_vcf if not self.skip_model_apply else self.output_vcf + ) + transfer_jobs, concat_job = build_transfer_jobs( + transfer_target, + 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)) + + if not self.skip_model_apply: + apply_job = self.build_dnamodelapply_job( + transfer_vcf, self.output_vcf + ) + dag.add_job(apply_job, {concat_job}) + + return dag + + def build_hybrid_kmc_job(self, kmer_prefix: pathlib.Path) -> Job: + """Count k-mers across the short-read fastq and long-read + alignments""" + assert self.reference is not None + kmc_job = Job( + cmds.cmd_hybrid_kmc( + kmer_prefix, + self.r1_fastq + self.r2_fastq, + self.lr_aln, + self.reference, + self.tmp_dir, + memory=self.kmer_memory, + threads=self.cores, + unzip=self.find_unzip(), + ), + "kmc", + 0, # run in the background + task_name="kmer-counting", + ) + return kmc_job + + def build_alignment_job( + self, + sample_bam: pathlib.Path, + sample_fastq: pathlib.Path, + ) -> Job: + """Build the bwa alignment and read extraction job""" + assert self.reference is not None + assert self.model_bundle is not None + + rg = copy.deepcopy(self.fastq_readgroup) + rg["SM"] = self.sample_sm + rg["LR"] = "0" + 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=self.find_unzip(), + ), + "bwa-extract", + self.cores, + task_name="alignment", + ) + return bwa_job + + 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.pangenome_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.pangenome_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, + 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(self.lr_aln), + ) + driver.add_algo( + PGHapUpdateAlgo(out_gfa, gfa_file=in_gfa, target_bed=bed) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + name, + self.cores, + task_name="pangenome-update", + ) + + def build_longreadsv_job(self, out_vcf: pathlib.Path) -> Job: + """Call SVs from the long reads for the graph update""" + assert self.model_bundle is not None + driver = Driver( + reference=self.reference, + thread_count=self.cores, + input=list(self.lr_aln), + ) + driver.add_algo( + LongReadSV( + out_vcf, + model=self.model_bundle.joinpath("longreadsv.model"), + min_sv_size=LONGREADSV_MIN_SV_SIZE, + ) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + "longreadsv", + self.cores, + task_name="pangenome-update", + ) + + 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 + + rg = copy.deepcopy(self.fastq_readgroup) + 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"], + ), + "mm2-lift", + self.cores, + task_name="pangenome-alignment", + ) + return mm2_job + + def build_replace_rg(self) -> List[List[str]]: + """`--replace_rg` arguments setting the LR readgroup attribute for + the variant calling stages. + + The bwa (LR:0) and lifted (LR:2) alignments are generated by the + pipeline with the LR attribute already in place; only the input + long-read readgroups are rewritten (LR:1), as they cannot be + assumed to carry the attribute. + """ + replace_rg: List[List[str]] = [[], []] # bwa and lifted alignments + for aln_rgs in self.lr_readgroups: + replace_rg.append([]) + for aln_rg in aln_rgs: + rg_id = aln_rg.get("ID") + sm = self.rgsm if self.rgsm else aln_rg.get("SM") + replace_rg[-1].append(f"{rg_id}=ID:{rg_id}\\tSM:{sm}\\tLR:1") + 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, + ) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + "pangenome-sv", + self.cores, + task_name="sv-calling", + ) + + def build_dnascope_job( + self, + out_vcf: pathlib.Path, + input_bams: List[pathlib.Path], + replace_rg: List[List[str]], + ) -> Job: + """Call small variants with the original, lifted, and long reads""" + assert self.model_bundle is not None + 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, + replace_rg=replace_rg, + ) + driver.add_algo( + DNAscope( + out_vcf, + model=self.model_bundle.joinpath("dnascope.model"), + pcr_indel_model=pcr_indel_model, + dbsnp=self.dbsnp, + ) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + "dnascope-raw", + self.cores, + task_name="variant-calling", + ) + + def build_dnamodelapply_job( + self, + in_vcf: pathlib.Path, + out_vcf: pathlib.Path, + ) -> Job: + """Apply the DNAscope model""" + assert self.model_bundle is not None + 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, + ) + ) + return Job( + Pipeline(Command(*driver.build_cmd())), + "model-apply", + self.cores, + task_name="model-apply", + ) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index bb985db..c6816ae 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -20,24 +20,14 @@ from .base_pangenome import BasePangenome, SampleSex from .dag import DAG from .driver import ( - AlignmentStat, - BaseDistributionByCycle, CNVModelApply, CNVscope, - CoverageMetrics, - Dedup, DNAModelApply, DNAscope, Driver, - GCBias, GVCFtyper, - InsertSizeMetricAlgo, - LocusCollector, - MeanQualityByCycle, PangenomeSV, - QualDistribution, ReadWriter, - WgsMetricsAlgo, ) from .job import Job from .logging import get_logger @@ -774,7 +764,10 @@ def build_first_dag(self) -> DAG: 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 + out_mm2_aln, + [mm2_bam], + "mm2", + left_align_rgid=f"{self.fastq_readgroup['ID']}-mm2", ) dag.add_job(bwa_lc_job, bwa_lc_dependencies) dag.add_job(bwa_dedup_job, {bwa_lc_job}) @@ -1084,135 +1077,6 @@ def build_minimap2_lift_job( ) 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, - task_name="dedup", - ) - - 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, - task_name="dedup", - ) - - 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, - task_name="metrics", - ) - - 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, - task_name="metrics", - ) - return (metrics_job, rehead_job) - def build_dnascope_job( self, out_vcf: pathlib.Path, diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py new file mode 100644 index 0000000..d2b99c7 --- /dev/null +++ b/tests/unit/test_hybrid_pangenome.py @@ -0,0 +1,474 @@ +""" +Unit tests for the HybridPangenome pipeline logic +""" + +import os +import pathlib +import sys +import tempfile +from unittest.mock import 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__), "..", "..")) +) + +from sentieon_cli.hybrid_pangenome import HybridPangenome +from sentieon_cli.command_strings import LONGREAD_SV_BED_AWK +from sentieon_cli.dag import DAG + + +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" + + 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, + ]: + 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.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 _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 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", + "pangenome-sv", + "dnascope-raw", + "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 + + 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 + + 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 + assert "sort -k1,1 -k2,2n" 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_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-raw"): + 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 + 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_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-raw").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-raw").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-raw").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-raw" 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-raw" 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-raw").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-raw").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 From 74f25ba364b864b82116967f2fc414ac1a540e33 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 15 Aug 2026 19:33:16 -0700 Subject: [PATCH 15/58] Support aligned short-read and unaligned long-read input in hybrid-pangenome Aligned short-read BAM/CRAM input via `--sr_aln`, mutually exclusive with fastq input. Aligned input is assumed deduplicated: no dedup or metrics jobs run, read extraction and k-mer counting share a single pass over the input (driver ReadWriter | pgutil extract | kmc, with the long-read FASTA streams concatenated), and the lifted alignment is written directly to the output. The calling stages rewrite the input readgroups with `--replace_rg` to set LR:0. Unaligned long-read (uBAM/uCRAM) input via `--lr_align_input` and `--lr_input_ref`. Each input is realigned with minimap2 using the new `minimap2_lr.model` bundle member; k-mer counting reads the original input files in parallel with the realignment. Also fix `cmd_samtools_fastq_minimap2` and `cmd_samtools_fastq_bwa` to decode the input file with `input_ref` rather than the alignment target reference, matching the flag's documented purpose. This changes behavior for dnascope-longread and dnascope-hybrid runs where the two references differ. Co-Authored-By: Claude Fable 5 --- sentieon_cli/command_strings.py | 174 ++++++-- sentieon_cli/hybrid_pangenome.py | 491 ++++++++++++++++----- tests/unit/test_command_strings_realign.py | 138 ++++++ tests/unit/test_hybrid_pangenome.py | 320 ++++++++++++++ 4 files changed, 964 insertions(+), 159 deletions(-) create mode 100644 tests/unit/test_command_strings_realign.py diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index e09877a..6a091ec 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 @@ -512,12 +512,15 @@ def cmd_samtools_fastq_minimap2( fastq_taglist: str = "*", minimap2_args: str = "-YL", util_sort_args: str = "--cram_write_options version=3.0,compressor=rans", + 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", @@ -530,6 +533,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", @@ -540,7 +546,7 @@ def cmd_samtools_fastq_minimap2( "-a", minimap2_args, "-x", - f"{model_bundle}/minimap2.model", + str(mm2_model), str(reference), "/dev/stdin", ) @@ -597,9 +603,11 @@ def cmd_samtools_fastq_bwa( util_sort_args: str = "--cram_write_options version=3.0,compressor=rans", ) -> 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)] if collate: collate_cmd = Command( @@ -879,19 +887,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, @@ -925,6 +936,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", @@ -1415,11 +1447,57 @@ 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[pathlib.Path], - reference: pathlib.Path, + aln: List[Tuple[pathlib.Path, pathlib.Path]], tmp_dir: pathlib.Path, k: int = 29, memory: int = 30, @@ -1431,7 +1509,8 @@ def cmd_hybrid_kmc( 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. + through stdin. Each alignment is supplied as an + `(alignment, decode reference)` pair. """ cat_args: List[Union[str, InputProcSub]] = [] if fastq: @@ -1444,34 +1523,47 @@ def cmd_hybrid_kmc( Command("awk", 'NR%4==1{print ">"substr($0,2)} NR%4==2{print}'), ) cat_args.append(InputProcSub(fq_fasta)) - for aln_file in aln: - cat_args.append( - InputProcSub( - Pipeline( - Command( - "samtools", - "fasta", - "--reference", - str(reference), - "-@", - str(threads), - str(aln_file), - ) - ) - ) - ) + cat_args.extend(_aln_fasta_procsubs(aln, threads)) cat_cmd = Command("cat", *cat_args) - kmc_cmd = Command( - "kmc", - f"-k{k}", - f"-m{memory}", - "-okff", - f"-t{threads}", - "-fa", - "/dev/stdin", - str(output_prefix), - str(tmp_dir), + 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) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 24ef897..a5a0d62 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -11,7 +11,7 @@ import pathlib import shutil import sys -from typing import Dict, List, Optional, Set +from typing import Dict, List, Optional, Set, Tuple import packaging.version @@ -90,7 +90,10 @@ class HybridPangenome(BasePangenome): "required": True, }, "readgroup": { - "help": "Readgroup information for the fastq files.", + "help": ( + "Readgroup information for the fastq files. Required " + "with fastq input; cannot be used with `--sr_aln`." + ), }, # Additional arguments "bed": { @@ -102,6 +105,21 @@ class HybridPangenome(BasePangenome): ), "type": path_arg(exists=True, is_file=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), + }, "pangenome_ref_name": { "default": "GRCh38", "help": "Reference name in the pangenome (GRCh38).", @@ -120,6 +138,15 @@ class HybridPangenome(BasePangenome): "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_contig_checks": { "help": argparse.SUPPRESS, @@ -156,11 +183,18 @@ def __init__(self) -> None: self.pop_vcf: Optional[pathlib.Path] = None self.readgroup: Optional[str] = None self.bed: Optional[pathlib.Path] = None + self.lr_align_input = False + self.lr_input_ref: Optional[pathlib.Path] = None self.pangenome_ref_name = "GRCh38" self.rgsm: Optional[str] = None self.extract_model_name = "extract.model" 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_contig_checks = False self.skip_model_apply = False self.skip_pangenome_name_checks = False @@ -183,18 +217,8 @@ def validate(self) -> None: self.validate_bundle() self.validate_output_vcf() - if not self.r1_fastq or not self.readgroup: - self.logger.error( - "Please supply the short reads with the `--r1_fastq`, " - "`--r2_fastq`, and `--readgroup` arguments" - ) - 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) + self.validate_sr_inputs() + if not self.lr_aln: self.logger.error( "Please supply the long-read alignments with the `--lr_aln` " @@ -202,7 +226,20 @@ def validate(self) -> None: ) sys.exit(2) - self.validate_bwa_index() + 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() @@ -272,6 +309,45 @@ def validate(self) -> None: ) sys.exit(2) + 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( @@ -310,12 +386,15 @@ def validate_bundle(self) -> None: self.extract_model_name = extract_candidate required_members = { - "bwa.model", "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") missing_members = required_members - bundle_members if missing_members: self.logger.error( @@ -340,29 +419,37 @@ def validate_bundle(self) -> None: def collect_readgroups(self) -> None: """Collect readgroup tags from the inputs""" - assert self.readgroup is not None - 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: Dict[str, str] = parsed_rg + 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 - self.lr_readgroups: List[List[Dict[str, str]]] = [] + # Read the readgroups from the ORIGINAL long-read input; minimap2 + # realignment preserves the readgroup IDs with `addreplacerg` + self.sr_readgroups = [] + for aln in self.sr_aln: + self.sr_readgroups.append([]) + for rg_line in cmds.get_rg_lines(aln, self.dry_run): + self.sr_readgroups[-1].append(parse_rg_line(rg_line)) + + self.lr_readgroups = [] for aln in self.lr_aln: self.lr_readgroups.append([]) for rg_line in cmds.get_rg_lines(aln, self.dry_run): @@ -370,8 +457,15 @@ def collect_readgroups(self) -> None: def validate_readgroups(self) -> None: """Confirm that all readgroups have a consistent SM tag""" - rg_sm = self.fastq_readgroup.get("SM") - for aln, aln_rgs in zip(self.lr_aln, self.lr_readgroups): + 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") + + aln_inputs = list(self.sr_aln) + list(self.lr_aln) + aln_readgroups = self.sr_readgroups + self.lr_readgroups + 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( @@ -393,7 +487,7 @@ def validate_readgroups(self) -> None: if sm != rg_sm: self.logger.error( "Input readgroup '%s' has a different RG-SM tag " - "from the `--readgroup` argument.\n" + "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), @@ -401,7 +495,7 @@ def validate_readgroups(self) -> None: rg_sm, ) sys.exit(2) - self.sample_sm: str = self.rgsm if self.rgsm else str(rg_sm) + self.sample_sm = self.rgsm if self.rgsm else str(rg_sm) def configure(self) -> None: """Configure pipeline parameters""" @@ -466,19 +560,49 @@ def build_dag(self) -> DAG: raw_vcf = self.tmp_dir.joinpath("sample-dnascope.vcf.gz") transfer_vcf = self.tmp_dir.joinpath("sample-dnascope_transfer.vcf.gz") - total_mem_gb = total_memory() / (1024.0**3) - - # 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: Set[Job] = {kmc_job} + # 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: + calling_lr, realign_jobs = self.lr_align_inputs() + for job in realign_jobs: + dag.add_job(job) + + 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_job = self.build_alignment_job(bwa_bam, ext_fastq) + dag.add_job(bwa_job) + 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) - # BWA alignment and extraction - bwa_job = self.build_alignment_job(bwa_bam, ext_fastq) - dag.add_job(bwa_job) - # Do not run vg-haplotypes with bwa in low-mem environments - if total_mem_gb < 70: - haplotype_dependencies.add(bwa_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) @@ -492,13 +616,14 @@ def build_dag(self) -> DAG: 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, {gfa_job}) + dag.add_job(update_raw_job, {gfa_job} | realign_jobs) # Call SVs from the long reads and collect graph update regions - longreadsv_job = self.build_longreadsv_job(longread_sv_vcf) - dag.add_job(longreadsv_job) + longreadsv_job = self.build_longreadsv_job(longread_sv_vcf, calling_lr) + dag.add_job(longreadsv_job, realign_jobs) sv_bed_job = Job( cmds.cmd_longread_sv_bed(sv_bed, longread_sv_vcf), "longread-sv-bed", @@ -511,6 +636,7 @@ def build_dag(self) -> DAG: update_job = self.build_graph_update_job( pangenome_gfa, pangenome_raw_gfa, + calling_lr, bed=sv_bed, name="graph-update", ) @@ -534,58 +660,67 @@ def build_dag(self) -> DAG: # minimap2 alignment of the extracted reads with liftover mm2_job = self.build_minimap2_lift_job( - lift_bam, + lift_aln, ext_fastq, pangenome_fasta, pangenome_gfa, ) - dag.add_job(mm2_job, {bwa_job, faidx_job, update_job}) + 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. + 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" + ) - # 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. - 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 ) - - bwa_lc_job, bwa_dedup_job = self.build_dedup_job( - out_bwa_aln, [bwa_bam], "bwa", metrics=dedup_metrics - ) - lift_lc_job, lift_dedup_job = self.build_dedup_job( - out_lift_aln, - [lift_bam], - "lift", - ) - dag.add_job(bwa_lc_job, {bwa_job}) - dag.add_job(bwa_dedup_job, {bwa_lc_job}) - dag.add_job(lift_lc_job, {mm2_job}) - dag.add_job(lift_dedup_job, {lift_lc_job}) - - # Alignment metrics from the deduplicated short reads - if not self.skip_metrics: - metrics_job, rehead_job = self.build_metrics_job( - [out_bwa_aln, out_lift_aln], + lift_lc_job, lift_dedup_job = self.build_dedup_job( + out_lift_aln, + [lift_bam], + "lift", ) - dag.add_job(metrics_job, {bwa_dedup_job, lift_dedup_job}) - dag.add_job(rehead_job, {metrics_job}) - if not self.skip_multiqc: - multiqc_job = self.multiqc() - if multiqc_job: - dag.add_job(multiqc_job, {rehead_job}) - - # Variant calling with the original, lifted, and long reads - calling_bams = [out_bwa_aln, out_lift_aln] + list(self.lr_aln) + dag.add_job(bwa_lc_job, {bwa_job}) + dag.add_job(bwa_dedup_job, {bwa_lc_job}) + dag.add_job(lift_lc_job, {mm2_job}) + dag.add_job(lift_dedup_job, {lift_lc_job}) + + # Alignment metrics from the deduplicated short reads + if not self.skip_metrics: + metrics_job, rehead_job = self.build_metrics_job( + [out_bwa_aln, out_lift_aln], + ) + dag.add_job(metrics_job, {bwa_dedup_job, lift_dedup_job}) + dag.add_job(rehead_job, {metrics_job}) + if not self.skip_multiqc: + multiqc_job = self.multiqc() + if multiqc_job: + dag.add_job(multiqc_job, {rehead_job}) + + calling_bams = [out_bwa_aln, out_lift_aln] + calling_lr + calling_dependencies = {bwa_dedup_job, lift_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} + calling_dependencies |= realign_jobs replace_rg = self.build_replace_rg() - calling_dependencies = {bwa_dedup_job, lift_dedup_job} if not self.skip_svs: pangenomesv_job = self.build_pangenomesv_job( @@ -633,16 +768,27 @@ def build_dag(self) -> DAG: 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""" - assert self.reference is not None kmc_job = Job( cmds.cmd_hybrid_kmc( kmer_prefix, self.r1_fastq + self.r2_fastq, - self.lr_aln, - self.reference, + self.lr_kmc_pairs(), self.tmp_dir, memory=self.kmer_memory, threads=self.cores, @@ -654,6 +800,85 @@ def build_hybrid_kmc_job(self, kmer_prefix: pathlib.Path) -> Job: ) 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_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: + """Align the long-read input files to the linear reference genome + with minimap2""" + assert self.output_vcf is not None + assert self.reference is not None + assert self.model_bundle is not None + + res: List[pathlib.Path] = [] + realign_jobs: Set[Job] = set() + suffix = "bam" if self.bam_format else "cram" + 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) + realign_jobs.add( + Job( + cmds.cmd_samtools_fastq_minimap2( + out_aln, + input_aln, + self.reference, + self.model_bundle, + self.cores, + rg_lines, + self.sample_sm, + self.lr_input_ref, + minimap2_model=self.model_bundle.joinpath( + "minimap2_lr.model" + ), + ), + f"bam-realign-{i}", + self.cores, + task_name="alignment", + ) + ) + res.append(out_aln) + return (res, realign_jobs) + def build_alignment_job( self, sample_bam: pathlib.Path, @@ -663,6 +888,8 @@ def build_alignment_job( assert self.reference is not None 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" @@ -733,6 +960,7 @@ 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: @@ -740,7 +968,7 @@ def build_graph_update_job( driver = Driver( reference=self.reference, thread_count=self.cores, - input=list(self.lr_aln), + input=list(lr_aln), ) driver.add_algo( PGHapUpdateAlgo(out_gfa, gfa_file=in_gfa, target_bed=bed) @@ -752,13 +980,17 @@ def build_graph_update_job( task_name="pangenome-update", ) - def build_longreadsv_job(self, out_vcf: pathlib.Path) -> Job: + def build_longreadsv_job( + self, + out_vcf: pathlib.Path, + lr_aln: List[pathlib.Path], + ) -> Job: """Call SVs from the long reads for the graph update""" assert self.model_bundle is not None driver = Driver( reference=self.reference, thread_count=self.cores, - input=list(self.lr_aln), + input=list(lr_aln), ) driver.add_algo( LongReadSV( @@ -786,7 +1018,14 @@ def build_minimap2_lift_job( assert self.reference is not None assert self.model_bundle is not None - rg = copy.deepcopy(self.fastq_readgroup) + # 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" @@ -808,22 +1047,38 @@ def build_minimap2_lift_job( ) 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 bwa (LR:0) and lifted (LR:2) alignments are generated by the - pipeline with the LR attribute already in place; only the input - long-read readgroups are rewritten (LR:1), as they cannot be - assumed to carry the attribute. + 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]] = [[], []] # bwa and lifted alignments + 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([]) - for aln_rg in aln_rgs: - rg_id = aln_rg.get("ID") - sm = self.rgsm if self.rgsm else aln_rg.get("SM") - replace_rg[-1].append(f"{rg_id}=ID:{rg_id}\\tSM:{sm}\\tLR:1") + replace_rg.append( + [self._replace_rg_arg(rg, "1") for rg in aln_rgs] + ) return replace_rg def build_pangenomesv_job( 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_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index d2b99c7..9953161 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -8,6 +8,8 @@ import tempfile from unittest.mock import MagicMock +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__), "..", "..")) @@ -37,6 +39,8 @@ def setup_method(self): 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, @@ -49,6 +53,8 @@ def setup_method(self): self.mock_bed, self.mock_r1, self.mock_r2, + self.mock_sr_bam, + self.mock_lr_ref, ]: file_path.touch() @@ -92,6 +98,28 @@ def create_pipeline(self): 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()) @@ -100,6 +128,11 @@ def _get_all_job_names(self, dag): 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() @@ -472,3 +505,290 @@ def test_rgsm_overrides_sm(self): 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", + "pangenome-sv", + "dnascope-raw", + "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-raw") == { + "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-raw").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-raw").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-raw", + ): + 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-raw"): + 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-raw").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-raw").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-raw") == { + "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 From e6d526fce1091b6095257f4cf849ad73a43ed456 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sun, 16 Aug 2026 00:23:54 -0700 Subject: [PATCH 16/58] Sort SV update regions in reference contig order Replace `sort -k1,1 -k2,2n` with `bedtools sort -faidx` in the LongReadSV BED generation so the regions sort correctly for references with an unusual contig order. Co-Authored-By: Claude Fable 5 --- sentieon_cli/command_strings.py | 3 ++- sentieon_cli/hybrid_pangenome.py | 2 +- tests/unit/test_hybrid_pangenome.py | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 6a091ec..e55ab93 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1581,11 +1581,12 @@ def cmd_hybrid_extract_kmc( 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("sort", "-k1,1", "-k2,2n") + sort_cmd = Command("bedtools", "sort", "-faidx", str(ref_fai), "-i", "-") merge_cmd = Command("bedtools", "merge") return Pipeline( zcat_cmd, diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index a5a0d62..d8d60be 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -625,7 +625,7 @@ def build_dag(self) -> DAG: longreadsv_job = self.build_longreadsv_job(longread_sv_vcf, calling_lr) dag.add_job(longreadsv_job, realign_jobs) sv_bed_job = Job( - cmds.cmd_longread_sv_bed(sv_bed, longread_sv_vcf), + cmds.cmd_longread_sv_bed(sv_bed, longread_sv_vcf, ref_fai), "longread-sv-bed", 0, task_name="pangenome-update", diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index 9953161..fc7b6f3 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -327,7 +327,8 @@ def test_longreadsv_and_bed(self): bed_cmd = str(self._get_job(all_jobs, "longread-sv-bed").shell) assert 'n=split($10,a,":")' in bed_cmd - assert "sort -k1,1 -k2,2n" 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 From 390c5cfe2d5444eaba6588954aaf18a6f03bd5ca Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sun, 16 Aug 2026 10:01:55 -0700 Subject: [PATCH 17/58] Thread the pangenome contig prefix through all graph consumers Add `--pangenome_contig_prefix` (default GRCh38#0#) and pass it to PGHapUpdateAlgo, pgutil lift, and PangenomeSV so a non-default pangenome reference name uses a consistent prefix throughout. Co-Authored-By: Claude Fable 5 --- sentieon_cli/hybrid_pangenome.py | 16 +++++++++++++++- tests/unit/test_hybrid_pangenome.py | 24 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index d8d60be..8207b68 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -120,6 +120,12 @@ class HybridPangenome(BasePangenome): ), "type": path_arg(exists=True, is_file=True), }, + "pangenome_contig_prefix": { + "default": "GRCh38#0#", + "help": ( + "Prefix to strip from pangenome contig names (GRCh38#0#)" + ), + }, "pangenome_ref_name": { "default": "GRCh38", "help": "Reference name in the pangenome (GRCh38).", @@ -185,6 +191,7 @@ def __init__(self) -> None: self.bed: Optional[pathlib.Path] = None self.lr_align_input = False self.lr_input_ref: Optional[pathlib.Path] = None + self.pangenome_contig_prefix = "GRCh38#0#" self.pangenome_ref_name = "GRCh38" self.rgsm: Optional[str] = None self.extract_model_name = "extract.model" @@ -971,7 +978,12 @@ def build_graph_update_job( input=list(lr_aln), ) driver.add_algo( - PGHapUpdateAlgo(out_gfa, gfa_file=in_gfa, target_bed=bed) + PGHapUpdateAlgo( + out_gfa, + gfa_file=in_gfa, + target_bed=bed, + prefix=self.pangenome_contig_prefix, + ) ) return Job( Pipeline(Command(*driver.build_cmd())), @@ -1040,6 +1052,7 @@ def build_minimap2_lift_job( self.model_bundle.joinpath("minimap2.model"), threads=self.cores, mm2_xargs=["--secondary=yes"], + lift_prefix=self.pangenome_contig_prefix, ), "mm2-lift", self.cores, @@ -1100,6 +1113,7 @@ def build_pangenomesv_job( out_vcf, gfa_file=pangenome_gfa, min_af=PANGENOME_SV_MIN_AF, + prefix=self.pangenome_contig_prefix, ) ) return Job( diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index fc7b6f3..7307c4e 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -304,6 +304,8 @@ def test_graph_update_commands(self): 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 @@ -311,6 +313,8 @@ def test_graph_update_commands(self): 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 @@ -385,11 +389,31 @@ def test_pangenomesv_command(self): 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_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() From f7b4a0f0e9762cf40e4738bf5904eb5c9f7ccdb3 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sun, 16 Aug 2026 15:05:08 -0700 Subject: [PATCH 18/58] Harden input readgroup validation and the KMC patch check Reject BAM/CRAM inputs with no @RG header lines (previously the LR attribute was silently never applied, or build_dag crashed), require readgroup IDs to be unique across all inputs, and report malformed @RG header lines with the offending file instead of a traceback. check_kmc_patch now parses the k-mer count from the KMC output and fails on a zero count, catching KMC builds that read nothing from piped stdin while still exiting successfully. Co-Authored-By: Claude Fable 5 --- sentieon_cli/hybrid_pangenome.py | 107 ++++++++++++++++++++++--- sentieon_cli/util.py | 20 ++++- tests/unit/test_hybrid_pangenome.py | 120 ++++++++++++++++++++++++++++ tests/unit/test_kmc_patch.py | 62 +++++++++++++- 4 files changed, 290 insertions(+), 19 deletions(-) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 8207b68..0f33b8e 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -450,28 +450,51 @@ def collect_readgroups(self) -> None: # Read the readgroups from the ORIGINAL long-read input; minimap2 # realignment preserves the readgroup IDs with `addreplacerg` - self.sr_readgroups = [] - for aln in self.sr_aln: - self.sr_readgroups.append([]) - for rg_line in cmds.get_rg_lines(aln, self.dry_run): - self.sr_readgroups[-1].append(parse_rg_line(rg_line)) - - self.lr_readgroups = [] - for aln in self.lr_aln: - self.lr_readgroups.append([]) - for rg_line in cmds.get_rg_lines(aln, self.dry_run): - self.lr_readgroups[-1].append(parse_rg_line(rg_line)) + 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") - aln_inputs = list(self.sr_aln) + list(self.lr_aln) - aln_readgroups = self.sr_readgroups + self.lr_readgroups for aln, aln_rgs in zip(aln_inputs, aln_readgroups): for aln_rg in aln_rgs: if not aln_rg.get("ID"): @@ -502,8 +525,66 @@ def validate_readgroups(self) -> None: 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 configure(self) -> None: """Configure pipeline parameters""" pass diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index fe9cc73..f03efaa 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -288,7 +288,8 @@ def check_kmc_patch(kmc_cmd: str = "kmc") -> bool: temp_path = pathlib.Path(temp_dir) output_prefix = temp_path / "kmc_test" - # Test input sequence + # Test input sequence. Both reads have an N-free stretch longer + # than the k-mer size, so a working KMC counts k-mers from them. test_input = ( ">206B4ABXX100825:7:1:1360:6029/1\n" "TGATTTTNNNNNNNNNNNTGAAGAACGCACCCATGTTAAAGAGCATGACAAANNNANNACAAGGCTAAGNGGCGNG\n" # noqa: E501 @@ -309,14 +310,25 @@ def check_kmc_patch(kmc_cmd: str = "kmc") -> bool: ] try: - sp.run( + res = sp.run( cmd, input=test_input, text=True, check=True, - stdout=sp.DEVNULL, + stdout=sp.PIPE, stderr=sp.DEVNULL, ) - return True except (sp.CalledProcessError, FileNotFoundError): return False + + # KMC prints a stats block after a successful run + match = re.search( + r"Total no\. of k-mers\s*:\s*(\d+)", res.stdout or "" + ) + if not match: + logger.debug( + "Could not find the k-mer count in the `%s` output", + kmc_cmd, + ) + return False + return int(match.group(1)) > 0 diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index 7307c4e..42c836c 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -15,6 +15,7 @@ 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) ) +from sentieon_cli import command_strings as cmds from sentieon_cli.hybrid_pangenome import HybridPangenome from sentieon_cli.command_strings import LONGREAD_SV_BED_AWK from sentieon_cli.dag import DAG @@ -817,3 +818,122 @@ def test_aligned_sr_with_lr_realign(self): 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" diff --git a/tests/unit/test_kmc_patch.py b/tests/unit/test_kmc_patch.py index 1589354..4c5f13a 100644 --- a/tests/unit/test_kmc_patch.py +++ b/tests/unit/test_kmc_patch.py @@ -3,20 +3,78 @@ """ import subprocess as sp -from unittest.mock import MagicMock, patch +from unittest.mock import patch from sentieon_cli.util import check_kmc_patch +def kmc_stdout(n_kmers: int) -> str: + """The stats block printed by KMC after a successful run""" + return ( + "1st stage: 0.068111s\n" + "2nd stage: 0.01257s\n" + "Total : 0.080681s\n" + "Tmp size : 0MB\n" + "\n" + "Stats:\n" + f" No. of k-mers below min. threshold : {n_kmers:>12}\n" + " No. of k-mers above max. threshold : 0\n" + f" No. of unique k-mers : {n_kmers:>12}\n" + " No. of unique counted k-mers : 0\n" + f" Total no. of k-mers : {n_kmers:>12}\n" + " Total no. of reads : 2\n" + " Total no. of super-k-mers : 3\n" + ) + + def test_check_kmc_patch_success(): - """Test check_kmc_patch returns True on success""" + """Test check_kmc_patch returns True when k-mers are counted""" with patch("subprocess.run") as mock_run: mock_run.return_value.returncode = 0 + mock_run.return_value.stdout = kmc_stdout(13) assert check_kmc_patch() is True mock_run.assert_called_once() args, kwargs = mock_run.call_args assert "kmc" == args[0][0] assert kwargs["input"] is not None + # The k-mer count is parsed from the captured stdout + assert kwargs["stdout"] == sp.PIPE + + +def test_check_kmc_patch_zero_kmers(): + """An unpatched KMC reads nothing from stdin but still exits 0""" + with patch("subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + mock_run.return_value.stdout = kmc_stdout(0) + assert check_kmc_patch() is False + + +def test_check_kmc_patch_missing_stats(): + """Test check_kmc_patch returns False without the k-mer count""" + with patch("subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + mock_run.return_value.stdout = "Stage 1: 100%\nStage 2: 100%\n" + assert check_kmc_patch() is False + + +def test_check_kmc_patch_no_stdout(): + """Test check_kmc_patch returns False without any output""" + with patch("subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + mock_run.return_value.stdout = None + assert check_kmc_patch() is False + + +def test_check_kmc_patch_custom_command(): + """Test check_kmc_patch runs the supplied kmc executable""" + with patch("subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + mock_run.return_value.stdout = kmc_stdout(13) + assert check_kmc_patch("/opt/kmc/kmc") is True + args, _kwargs = mock_run.call_args + assert args[0][0] == "/opt/kmc/kmc" + # KMC reads the test input from stdin + assert "/dev/stdin" in args[0] def test_check_kmc_patch_failure(): From feabf07f210ee1838107d2add3dbcbcb0084910b Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sun, 16 Aug 2026 22:14:19 -0700 Subject: [PATCH 19/58] Sort the HybridStage1 hap_bam output explicitly In Sentieon 202503.04 the HybridStage1 `--hap_bam` output is unsorted and `-` sends it to stdout. Send the hap BAM to stdout and pipe it to `sentieon util sort`, and move the algo's fastq output to a named fifo, `stage1_hap.fq`, read by bwa in the first-stage alignment. The haplotype job runs with zero scheduler threads so it always executes concurrently with the first-stage job reading the fifo. Requires sentieon driver 202503.04 or later. Co-Authored-By: Claude Fable 5 --- sentieon_cli/command_strings.py | 30 +++- sentieon_cli/dnascope_hybrid.py | 94 +++++++---- sentieon_cli/driver.py | 2 +- tests/unit/test_dnascope_hybrid_stage1.py | 187 ++++++++++++++++++++++ 4 files changed, 280 insertions(+), 33 deletions(-) create mode 100644 tests/unit/test_dnascope_hybrid_stage1.py diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index e55ab93..907c532 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -316,24 +316,46 @@ def cmd_pyexec_hybrid_anno( 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)), ) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 6be5285..012a5b0 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -51,7 +51,8 @@ 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"), @@ -640,6 +641,8 @@ def build_dag(self) -> DAG: mapq0_slop_job, cat_merge_job, rm_job1, + stage1_fifo_job, + stage1_hap_job, stage1_job, rm_job2, second_stage_job, @@ -661,8 +664,11 @@ def build_dag(self) -> DAG: 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}) + # The haplotype job writes the fifo read by the first-stage job + dag.add_job(stage1_fifo_job) + dag.add_job(stage1_hap_job, {stage1_fifo_job, cat_merge_job}) + dag.add_job(stage1_job, {stage1_fifo_job, cat_merge_job}) + dag.add_job(second_stage_job, {stage1_job, stage1_hap_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}) @@ -684,7 +690,7 @@ def build_dag(self) -> DAG: # 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_job2, {stage1_job, stage1_hap_job}) dag.add_job(rm_job3, {second_stage_job}) dag.add_job(rm_job4, {third_stage_job}) dag.add_job(rm_job5, {concat_job}) @@ -698,27 +704,29 @@ def call_variants( 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], + Job, # call_job + Job, # select_job + Job, # mapq0_job + Job, # mapq0_slop_job + Job, # cat_merge_job + Job, # rm_job1 + Job, # stage1_fifo_job + Job, # stage1_hap_job + Job, # stage1_job + Job, # rm_job2 + Job, # second_stage_job + Job, # rm_job3 + Job, # third_stage_job + Job, # rm_job4 + Job, # call2_job + Job, # subset_job + Job, # concat_job + Job, # rm_job5 + Job, # anno_job + Optional[List[Job]], # transfer_jobs + Optional[Job], # transfer_concat_job + Optional[Job], # apply_job + Optional[Job], # norm_job ]: """ Call SNVs and indels using the DNAscope hybrid pipeline @@ -862,6 +870,17 @@ 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", + ) + stage1_driver = Driver( reference=self.reference, thread_count=self.cores, @@ -872,13 +891,28 @@ def call_variants( ) stage1_driver.add_algo( HybridStage1( - "-", + stage1_fifo, model=self.model_bundle.joinpath("HybridStage1.model"), - hap_bam=stage1_hap_bam, + 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", + ) stage1_bam = self.tmp_dir.joinpath("hybrid_stage1.bam") stage1_job = Job( @@ -888,7 +922,7 @@ def call_variants( cores=self.cores, readgroup=f"@RG\\tID:hybrid-18893\\tSM:{self.hybrid_rg_sm}", ins_driver=ins_driver, - stage1_driver=stage1_driver, + hap_fastq_fifo=stage1_fifo, bwa_model=self.model_bundle.joinpath("HybridStage1_bwa.model"), ), "first-stage", @@ -1086,6 +1120,8 @@ def call_variants( mapq0_slop_job, cat_merge_job, rm_job1, + stage1_fifo_job, + stage1_hap_job, stage1_job, rm_job2, second_stage_job, @@ -1143,6 +1179,8 @@ def call_variants( mapq0_slop_job, cat_merge_job, rm_job1, + stage1_fifo_job, + stage1_hap_job, stage1_job, rm_job2, second_stage_job, diff --git a/sentieon_cli/driver.py b/sentieon_cli/driver.py index e39fe80..c6872c3 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, diff --git a/tests/unit/test_dnascope_hybrid_stage1.py b/tests/unit/test_dnascope_hybrid_stage1.py new file mode 100644 index 0000000..6b72aeb --- /dev/null +++ b/tests/unit/test_dnascope_hybrid_stage1.py @@ -0,0 +1,187 @@ +""" +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""" + with patch( + "sentieon_cli.dnascope_hybrid.check_version", return_value=True + ): + 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_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") From ecc766328f9320c2cb3e0e53e7c5e02000cb335d Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sun, 16 Aug 2026 23:03:17 -0700 Subject: [PATCH 20/58] Remove the temporary directory when a sentieon-pangenome run fails The two-DAG main() override performed its tmpdir cleanup inline after the last check_execution, so a failing run skipped it and leaked the directory. Wrap the body in try/finally, mirroring BasePipeline.main. Found by qualification (DEFECT_sentieon_pangenome_tmpdir_leak). Co-Authored-By: Claude Fable 5 --- sentieon_cli/sentieon_pangenome.py | 19 ++--- tests/unit/test_pipeline_lifecycle.py | 104 ++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 9 deletions(-) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index c6816ae..938faf4 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -282,18 +282,19 @@ def main(self, args: argparse.Namespace) -> None: 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() + try: + dag = self.build_first_dag() executor = self.run(dag) self.check_execution(dag, executor) - if not self.retain_tmpdir: - shutil.rmtree(tmp_dir_str) + 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) + finally: + if not self.retain_tmpdir: + shutil.rmtree(tmp_dir_str) success = True finally: self.log_completion(success, start_time) diff --git a/tests/unit/test_pipeline_lifecycle.py b/tests/unit/test_pipeline_lifecycle.py index 593a03f..aa701e8 100644 --- a/tests/unit/test_pipeline_lifecycle.py +++ b/tests/unit/test_pipeline_lifecycle.py @@ -16,10 +16,12 @@ os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), ) +from sentieon_cli import sentieon_pangenome # 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.sentieon_pangenome import SentieonPangenome # noqa: E402 from sentieon_cli.shell_pipeline import Command, Pipeline # noqa: E402 @@ -138,6 +140,108 @@ def test_check_execution_names_the_failed_jobs(): 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_ref(self) -> None: + pass + + def validate(self) -> None: + pass + + def configure(self) -> None: + pass + + def build_first_dag(self) -> DAG: + self.dags_built.append("first") + self.ploidy_json = self.tmp_dir.joinpath("ploidy.json") + return DAG() + + def build_second_dag(self) -> DAG: + 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)) + monkeypatch.setattr(sentieon_pangenome, "parse_fai", lambda fai: {}) + monkeypatch.setattr( + sentieon_pangenome, "determine_shards_from_fai", lambda fai, size: [] + ) + return _StubPangenome(**kwargs) + + +def test_pangenome_main_cleans_up_tmpdir_when_the_first_dag_fails( + tmp_path, monkeypatch +): + # SentieonPangenome overrides main() for its 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_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")) From 4f0750f2bdad6e518d3c3e042669c35a58069e49 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sun, 16 Aug 2026 23:17:44 -0700 Subject: [PATCH 21/58] Revert the k-mer count check in check_kmc_patch Restore the exit-status-only probe; the piped-KMC defect is being addressed in the KMC build itself. Co-Authored-By: Claude Fable 5 --- sentieon_cli/util.py | 20 +++--------- tests/unit/test_kmc_patch.py | 62 ++---------------------------------- 2 files changed, 6 insertions(+), 76 deletions(-) diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index f03efaa..fe9cc73 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -288,8 +288,7 @@ def check_kmc_patch(kmc_cmd: str = "kmc") -> bool: temp_path = pathlib.Path(temp_dir) output_prefix = temp_path / "kmc_test" - # Test input sequence. Both reads have an N-free stretch longer - # than the k-mer size, so a working KMC counts k-mers from them. + # Test input sequence test_input = ( ">206B4ABXX100825:7:1:1360:6029/1\n" "TGATTTTNNNNNNNNNNNTGAAGAACGCACCCATGTTAAAGAGCATGACAAANNNANNACAAGGCTAAGNGGCGNG\n" # noqa: E501 @@ -310,25 +309,14 @@ def check_kmc_patch(kmc_cmd: str = "kmc") -> bool: ] try: - res = sp.run( + sp.run( cmd, input=test_input, text=True, check=True, - stdout=sp.PIPE, + stdout=sp.DEVNULL, stderr=sp.DEVNULL, ) + return True except (sp.CalledProcessError, FileNotFoundError): return False - - # KMC prints a stats block after a successful run - match = re.search( - r"Total no\. of k-mers\s*:\s*(\d+)", res.stdout or "" - ) - if not match: - logger.debug( - "Could not find the k-mer count in the `%s` output", - kmc_cmd, - ) - return False - return int(match.group(1)) > 0 diff --git a/tests/unit/test_kmc_patch.py b/tests/unit/test_kmc_patch.py index 4c5f13a..1589354 100644 --- a/tests/unit/test_kmc_patch.py +++ b/tests/unit/test_kmc_patch.py @@ -3,78 +3,20 @@ """ import subprocess as sp -from unittest.mock import patch +from unittest.mock import MagicMock, patch from sentieon_cli.util import check_kmc_patch -def kmc_stdout(n_kmers: int) -> str: - """The stats block printed by KMC after a successful run""" - return ( - "1st stage: 0.068111s\n" - "2nd stage: 0.01257s\n" - "Total : 0.080681s\n" - "Tmp size : 0MB\n" - "\n" - "Stats:\n" - f" No. of k-mers below min. threshold : {n_kmers:>12}\n" - " No. of k-mers above max. threshold : 0\n" - f" No. of unique k-mers : {n_kmers:>12}\n" - " No. of unique counted k-mers : 0\n" - f" Total no. of k-mers : {n_kmers:>12}\n" - " Total no. of reads : 2\n" - " Total no. of super-k-mers : 3\n" - ) - - def test_check_kmc_patch_success(): - """Test check_kmc_patch returns True when k-mers are counted""" + """Test check_kmc_patch returns True on success""" with patch("subprocess.run") as mock_run: mock_run.return_value.returncode = 0 - mock_run.return_value.stdout = kmc_stdout(13) assert check_kmc_patch() is True mock_run.assert_called_once() args, kwargs = mock_run.call_args assert "kmc" == args[0][0] assert kwargs["input"] is not None - # The k-mer count is parsed from the captured stdout - assert kwargs["stdout"] == sp.PIPE - - -def test_check_kmc_patch_zero_kmers(): - """An unpatched KMC reads nothing from stdin but still exits 0""" - with patch("subprocess.run") as mock_run: - mock_run.return_value.returncode = 0 - mock_run.return_value.stdout = kmc_stdout(0) - assert check_kmc_patch() is False - - -def test_check_kmc_patch_missing_stats(): - """Test check_kmc_patch returns False without the k-mer count""" - with patch("subprocess.run") as mock_run: - mock_run.return_value.returncode = 0 - mock_run.return_value.stdout = "Stage 1: 100%\nStage 2: 100%\n" - assert check_kmc_patch() is False - - -def test_check_kmc_patch_no_stdout(): - """Test check_kmc_patch returns False without any output""" - with patch("subprocess.run") as mock_run: - mock_run.return_value.returncode = 0 - mock_run.return_value.stdout = None - assert check_kmc_patch() is False - - -def test_check_kmc_patch_custom_command(): - """Test check_kmc_patch runs the supplied kmc executable""" - with patch("subprocess.run") as mock_run: - mock_run.return_value.returncode = 0 - mock_run.return_value.stdout = kmc_stdout(13) - assert check_kmc_patch("/opt/kmc/kmc") is True - args, _kwargs = mock_run.call_args - assert args[0][0] == "/opt/kmc/kmc" - # KMC reads the test input from stdin - assert "/dev/stdin" in args[0] def test_check_kmc_patch_failure(): From 31ae7fd8aabf0e6b1a8402ae81175f45974159d9 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 18 Aug 2026 07:54:48 -0700 Subject: [PATCH 22/58] Increment version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8f866cb..d953ebd 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"}, From 5b60f8afb96cc4932afd19b01d12579f8157aca5 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 18 Aug 2026 08:19:10 -0700 Subject: [PATCH 23/58] Use the SENTIEON_VERSION repo variable in CI-smoke Replace the hardcoded sentieon-version matrix value with ${{ vars.SENTIEON_VERSION }} in all three smoke jobs, matching the pattern already used by the Docker workflow. Each "Install sentieon" step now fails with an explicit error when the variable is unset, instead of a confusing tar failure on a 404 from S3. Co-Authored-By: Claude Opus 5 --- .github/workflows/main.yml | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 0e5e624..1cee559 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) From be7b241867f42c960159f05358186bbc8a25c655 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 19 Aug 2026 14:16:43 -0700 Subject: [PATCH 24/58] Write per-process and per-job metrics tables to the log directory In addition to the rusage log lines, write two machine-readable TSV files to the run's log directory so metrics can be loaded into pandas/awk without parsing run.log: - process_metrics.txt: one row per reaped process, straight from rusage, with a job_id column linking each process to its job and a stage column matching the task_logs file names. Rows are written for every reaped process whatever the job's outcome; the exit code is the raw wait status. A process reaped without rusage keeps its pid/exit_code and leaves the usage columns empty. - job_metrics.txt: one row per job aggregated across its processes: status (ok/failed), wall clock, summed user/sys CPU, per-process peak RSS (matching the "Finished command" line's semantics), and process count. Jobs that never reach accounting (launch failure, interrupt) get no row. Headers are (re)written in RunLogs.create_dirs(), so a rerun into the same --log_dir replaces the previous run's rows; rows append as each job finishes accounting, which also keeps partial data on failure and lets the pangenome pipeline's two DAGs share one set of files. Whitespace in the command field is flattened so rows stay rectangular, and a metrics write failure degrades to a warning rather than failing the job. JobLogSink gains a non-allocating stage_index_for() lookup. Co-Authored-By: Claude Fable 5 --- README.md | 2 +- sentieon_cli/executor.py | 38 ++++++- sentieon_cli/run_logs.py | 166 ++++++++++++++++++++++++++++++- tests/unit/test_job_logs.py | 11 ++ tests/unit/test_metrics_files.py | 147 +++++++++++++++++++++++++++ 5 files changed, 357 insertions(+), 7 deletions(-) create mode 100644 tests/unit/test_metrics_files.py diff --git a/README.md b/README.md index cccd5fc..c5e56e8 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,7 @@ The `sentieon-cli` supports the following global arguments: - `--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/`. 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. +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/sentieon_cli/executor.py b/sentieon_cli/executor.py index 2fecf6f..9f3fec4 100644 --- a/sentieon_cli/executor.py +++ b/sentieon_cli/executor.py @@ -453,6 +453,8 @@ async def _drive(self, stop_event: asyncio.Event) -> None: 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( @@ -461,6 +463,7 @@ async def _drive(self, stop_event: asyncio.Event) -> None: ) cmd_failed = True continue + n_procs += 1 ret = ( await subcommand.proc.async_wait() ) # Wait on all sub-commands @@ -468,8 +471,12 @@ async def _drive(self, stop_event: asyncio.Event) -> None: # job's outcome; signal-killed children have rusage # too. rusage = subcommand.proc.rusage - if rusage is not None: - rss = _maxrss_bytes(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 @@ -492,6 +499,20 @@ async def _drive(self, stop_event: asyncio.Event) -> None: 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 @@ -527,7 +548,6 @@ async def _drive(self, stop_event: asyncio.Event) -> None: logger.error("Error: %s", str(exc)) cmd_failed = True - sink = context.log_sink if cmd_failed: logger.error("Command failure: %s", job.shell) # The sink's files are closed by cleanup(), so their @@ -556,6 +576,18 @@ async def _drive(self, stop_event: asyncio.Event) -> None: 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) diff --git a/sentieon_cli/run_logs.py b/sentieon_cli/run_logs.py index adee9db..6d605b1 100644 --- a/sentieon_cli/run_logs.py +++ b/sentieon_cli/run_logs.py @@ -11,18 +11,82 @@ import shlex import shutil import sys -from typing import IO, List, Optional, TYPE_CHECKING +from typing import IO, Iterable, List, Optional, TYPE_CHECKING -from .logging import add_file_handler, remove_file_handler +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: @@ -74,6 +138,17 @@ def _stage_for(self, command: Command) -> int: 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) @@ -145,6 +220,8 @@ def __init__(self, log_dir: pathlib.Path) -> None: 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: @@ -154,18 +231,101 @@ def setup(self) -> None: self.file_handler = add_file_handler(self.run_log) def create_dirs(self) -> None: - """Create the log directory, clearing logs from any previous run""" + """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() diff --git a/tests/unit/test_job_logs.py b/tests/unit/test_job_logs.py index 29395bc..841fa9e 100644 --- a/tests/unit/test_job_logs.py +++ b/tests/unit/test_job_logs.py @@ -127,6 +127,17 @@ def test_logs_are_matched_to_commands_by_identity(tmp_path): 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") 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] From dfb7f528c8c2abcf1d1f257af09972fb40c3314b Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 21 Aug 2026 20:48:13 -0700 Subject: [PATCH 25/58] Run pangenome metrics on the bwa alignment only The metrics job received both the bwa and mm2 deduped alignments, but the mm2 alignment is a re-alignment of a subset of the same reads, so the extracted reads were double-counted in the collected metrics. Co-Authored-By: Claude Fable 5 --- sentieon_cli/sentieon_pangenome.py | 4 ++-- tests/unit/test_sentieon_pangenome.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 938faf4..5fa6cbf 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -782,9 +782,9 @@ def build_first_dag(self) -> DAG: if not self.skip_metrics: metrics_job, rehead_job = self.build_metrics_job( - [out_bwa_aln, out_mm2_aln], + [out_bwa_aln], ) - dag.add_job(metrics_job, {bwa_dedup_job, mm2_dedup_job}) + dag.add_job(metrics_job, {bwa_dedup_job}) dag.add_job(rehead_job, {metrics_job}) if not self.skip_multiqc: multiqc_job = self.multiqc() diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index e6c1914..97319d9 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -133,6 +133,18 @@ 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 double-count the extracted reads it re-aligns.""" + pipeline = self.create_fastq_pipeline() + dag = pipeline.build_first_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() From 9c9471f0537ddaf567a451b2cf5c69bdef96c249 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Mon, 31 Aug 2026 12:04:20 -0700 Subject: [PATCH 26/58] Update the pangenome AD and DP counts from SAD or LAD Add scripts/sad_lad_update.py to the dnascope-pangenome pipeline as a final ad-update job. The joint pileup calls the bwa and mm2 alignments together; the script sets each sample's FORMAT/AD to either FORMAT/SAD or FORMAT/LAD -- whichever likelihood vector (SPL/LPL) better supports the called genotype -- and derives FORMAT/DP and INFO/DP from that choice. The previous values are kept as AD_ORIG/DP_ORIG and the choice is recorded in FORMAT/ADSRC and FORMAT/ADRULE. DNAModelApply (or the transfer concat under --skip_model_apply) now writes an intermediate in tmp_dir and the ad-update job writes the real output. With --gvcf the update runs on the gVCF ahead of GVCFtyper, so the genotyped VCF at --output_vcf carries the updated counts; gVCF hom-ref blocks pass through unchanged. No job is added under --skip_small_variants. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 1 + sentieon_cli/command_strings.py | 20 + sentieon_cli/scripts/sad_lad_update.py | 491 +++++++++++++++++++++++++ sentieon_cli/sentieon_pangenome.py | 44 ++- tests/unit/test_sentieon_pangenome.py | 108 +++++- 5 files changed, 649 insertions(+), 15 deletions(-) create mode 100755 sentieon_cli/scripts/sad_lad_update.py diff --git a/pyproject.toml b/pyproject.toml index d953ebd..29d7e12 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,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/command_strings.py b/sentieon_cli/command_strings.py index 907c532..f013faa 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -316,6 +316,26 @@ 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)) + + def hybrid_stage1_hap( out_hap_bam: pathlib.Path, stage1_driver: BaseDriver, 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 5fa6cbf..df4e3c0 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -878,6 +878,11 @@ def build_first_dag(self) -> 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 + 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 if self.pop_vcf: @@ -885,7 +890,7 @@ def build_first_dag(self) -> DAG: ( transfer_vcf if not self.skip_model_apply - else (small_variants_out) + else (snv_apply_vcf) ), self.pop_vcf, raw_vcf, @@ -905,13 +910,21 @@ def build_first_dag(self) -> DAG: if not self.skip_model_apply: # DNAModelApply apply_job = self.build_dnamodelapply_job( - transfer_vcf, small_variants_out + transfer_vcf, snv_apply_vcf ) dag.add_job(apply_job, apply_dependencies) small_variants_last_job = apply_job elif self.pop_vcf: small_variants_last_job = concat_job + # Update the overestimated AD/DP of the joint pileup + if small_variants_last_job is not None: + ad_update_job = self.build_count_ad_update_job( + small_variants_out, snv_apply_vcf + ) + dag.add_job(ad_update_job, {small_variants_last_job}) + small_variants_last_job = ad_update_job + # 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) @@ -1181,6 +1194,33 @@ def build_gvcftyper_job( task_name="gvcftyper", ) + def build_count_ad_update_job( + self, + out_vcf: pathlib.Path, + in_vcf: pathlib.Path, + ) -> Job: + """Update the AD/DP of the joint pileup. + + 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( + cmds.cmd_pyexec_sad_lad_update( + out_vcf, + in_vcf, + sad_lad_update, + self.cores, + ), + "sad-lad-update", + self.cores, + task_name="ad-update", + ) + def build_segdup_job( self, out_segdup: pathlib.Path, diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index 97319d9..0bf3605 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -135,7 +135,7 @@ def test_dedup_metrics_output(self): def test_metrics_input_bwa_only(self): """The metrics job reads only the bwa alignment; including the mm2 - alignment would double-count the extracted reads it re-aligns.""" + alignment would count the extracted reads it re-aligns twice.""" pipeline = self.create_fastq_pipeline() dag = pipeline.build_first_dag() @@ -156,7 +156,8 @@ def test_dedup_metrics_skipped(self): 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() @@ -167,6 +168,14 @@ 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() @@ -179,24 +188,37 @@ 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 + # 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(pipeline.output_vcf) in args_str + 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() @@ -204,27 +226,87 @@ def test_gvcf_adds_gvcftyper_and_routes_outputs(self): 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 + # DNAscope emits gVCF; model-apply writes gVCF for name in ("dnascope-raw", "model-apply"): job = next(j for j in all_jobs if j.name == name) cmd_str = str(job.shell) if name == "dnascope-raw": assert "--emit_mode gvcf" in cmd_str else: - assert expected_gvcf in cmd_str + assert apply_gvcf in cmd_str + assert expected_gvcf not in cmd_str - # GVCFtyper reads the gVCF and writes the final output VCF + # 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_first_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_first_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""" From 12c70a1693e737bac149704cd0112b0070701059 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 1 Sep 2026 12:12:55 -0700 Subject: [PATCH 27/58] Support the new CNVscope syntax from release 202503.04 CNVscope now takes --sex and, for male samples, a --par BED file. dnascope-pangenome and dnascope-hybrid always estimate the sample sex in their first DAG (writing _ploidy.json) and run CNV calling in a second DAG once the sex is known. The PAR BED file is selected from packaged data for recognized reference builds (hg38, hg19, b37, chm13) and can be overridden with --par_bed; --sample_sex overrides the estimated sex used for calling. CNV calling requires driver 202503.04 and a PAR BED file, checked during validation. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 4 + sentieon_cli/base_pangenome.py | 45 --- sentieon_cli/command_strings.py | 16 + sentieon_cli/data/par_b37.bed | 4 + sentieon_cli/data/par_chm13.bed | 4 + sentieon_cli/data/par_hg19.bed | 4 + sentieon_cli/data/par_hg38.bed | 4 + sentieon_cli/dnascope.py | 10 +- sentieon_cli/dnascope_hybrid.py | 117 +++++- sentieon_cli/driver.py | 5 + sentieon_cli/pipeline.py | 124 +++++- sentieon_cli/scripts/estimate_ploidy.py | 36 +- sentieon_cli/sentieon_pangenome.py | 111 ++++-- sentieon_cli/shard.py | 104 +++++ sentieon_cli/util.py | 44 ++- tests/integration/test_pipeline_execution.py | 6 +- tests/unit/test_cnv_sex.py | 375 +++++++++++++++++++ tests/unit/test_dnascope_hybrid_cnv.py | 226 +++++++++++ tests/unit/test_estimate_ploidy.py | 153 ++++++++ tests/unit/test_pipeline_lifecycle.py | 92 +++++ tests/unit/test_sentieon_pangenome.py | 178 ++++++++- tests/utils/test_helpers.py | 4 + 22 files changed, 1550 insertions(+), 116 deletions(-) create mode 100644 sentieon_cli/data/par_b37.bed create mode 100644 sentieon_cli/data/par_chm13.bed create mode 100644 sentieon_cli/data/par_hg19.bed create mode 100644 sentieon_cli/data/par_hg38.bed create mode 100644 tests/unit/test_cnv_sex.py create mode 100644 tests/unit/test_dnascope_hybrid_cnv.py create mode 100644 tests/unit/test_estimate_ploidy.py diff --git a/pyproject.toml b/pyproject.toml index 29d7e12..0b2a465 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index c126bbd..c617c48 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -3,8 +3,6 @@ """ import copy -from enum import Enum -import json import pathlib import sys from typing import List, Optional, Tuple @@ -31,12 +29,6 @@ from .util import path_arg -class SampleSex(Enum): - FEMALE = 1 - MALE = 2 - UNKNOWN = 3 - - class BasePangenome(BasePipeline): """A pipeline base class for short reads""" @@ -149,27 +141,6 @@ def build_kmc_job( 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, - task_name="ploidy", - ) - return ploidy_job - def build_dedup_job( self, output_bam, @@ -297,19 +268,3 @@ def build_metrics_job( task_name="metrics", ) return (metrics_job, rehead_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 diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index f013faa..4cd0ba1 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1223,6 +1223,10 @@ 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 = ( [ @@ -1236,6 +1240,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)) 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 70a2e10..2a3d837 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -42,7 +42,9 @@ from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline from .util import ( + SampleSex, check_version, + cnvscope_sex_args, library_preloaded, parse_rg_line, path_arg, @@ -64,7 +66,8 @@ } CNV_MIN_VERSIONS = { - "sentieon driver": packaging.version.Version("202308.03"), + # 202503.04 adds the CNVscope `--sex` and `--par` arguments + "sentieon driver": packaging.version.Version("202503.04"), } @@ -995,6 +998,8 @@ def call_cnvs( cores: int = mp.cpu_count(), skip_version_check: bool = False, replace_rg: Optional[List[List[str]]] = None, + sample_sex: Optional[SampleSex] = None, + par_bed: Optional[pathlib.Path] = None, **_kwargs: Any, ) -> Tuple[Job, Job]: """ @@ -1005,6 +1010,7 @@ def call_cnvs( if not check_version(cmd, min_version): sys.exit(2) + sex, par = cnvscope_sex_args(sample_sex, par_bed) cnvscope_vcf = tmp_dir.joinpath("cnvscope.vcf.gz") cnv_vcf = pathlib.Path(str(output_vcf).replace(".vcf.gz", ".cnv.vcf.gz")) driver = Driver( @@ -1018,6 +1024,8 @@ def call_cnvs( CNVscope( cnvscope_vcf, model_bundle.joinpath("cnv.model"), + sex=sex, + par=par, ) ) cnvscope_job = Job( diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 012a5b0..45a8cbc 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -26,7 +26,7 @@ HybridStage2, HybridStage3, ) -from .dnascope import call_cnvs, DNAscopePipeline +from .dnascope import call_cnvs, CNV_MIN_VERSIONS, DNAscopePipeline from .dnascope_longread import DNAscopeLRPipeline from .job import Job from .pipeline import BasePipeline @@ -37,12 +37,14 @@ library_preloaded, parse_rg_line, path_arg, + sample_sex_arg, split_alignment, vcf_id, ) from .shard import ( determine_shards_from_fai, parse_fai, + ploidy_contigs_for_build, vcf_contigs, ) from .transfer import build_transfer_jobs @@ -202,6 +204,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": ( @@ -209,6 +219,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 " @@ -302,6 +321,7 @@ def __init__(self) -> None: self.lr_aln: List[pathlib.Path] = [] 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 @@ -311,6 +331,10 @@ def __init__(self) -> None: self.assay = "WGS" self.skip_model_apply = False self.skip_pop_vcf_id_check = False + # 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")) @@ -325,6 +349,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: @@ -356,6 +382,24 @@ 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 + if not self.skip_version_check: + for cmd, min_version in CNV_MIN_VERSIONS.items(): + if not check_version(cmd, min_version): + sys.exit(2) + def validate_bundle(self) -> None: bundle_info_bytes = ar_load( str(self.model_bundle) + "/bundle_info.json" @@ -615,24 +659,27 @@ def build_dag(self) -> DAG: 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. + self.ploidy_json = pathlib.Path( + str(self.output_vcf).replace(".vcf.gz", "_ploidy.json") + ) + ploidy_job = self.build_ploidy_job( + self.ploidy_json, + sr_aln, + ploidy_contigs_for_build(self.reference_build), + ) + dag.add_job(ploidy_job, sr_preprocessing_jobs) + 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, @@ -697,6 +744,42 @@ def build_dag(self) -> DAG: return dag + def build_second_dag(self) -> Optional[DAG]: + """Build the second DAG, for sex-aware CNV calling""" + if self.skip_cnv: + return None + 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) + assert self.ploidy_json is not None + + self.get_sex(self.ploidy_json) + + self.logger.info("Building the DNAscope Hybrid CNV DAG") + dag = DAG() + cnvscope_job, cnvmodelapply_job = call_cnvs( + self.tmp_dir, + self.output_vcf, + self.reference, + self.cnv_sr_aln, + self.model_bundle, + self.bed, + self.cores, + self.skip_version_check, + replace_rg=self.cnv_replace_rg, + sample_sex=self.sample_sex, + par_bed=self.cnv_par_bed, + ) + dag.add_job(cnvscope_job) + dag.add_job(cnvmodelapply_job, {cnvscope_job}) + return dag + def call_variants( self, sr_aln: List[pathlib.Path], diff --git a/sentieon_cli/driver.py b/sentieon_cli/driver.py index c6872c3..4574df8 100644 --- a/sentieon_cli/driver.py +++ b/sentieon_cli/driver.py @@ -497,9 +497,14 @@ def __init__( self, output: Union[pathlib.Path, str], model: pathlib.Path, + *, + sex: Optional[str] = None, + par: Optional[pathlib.Path] = None, ): self.output = output self.model = model + self.sex = sex + self.par = par class CNVModelApply(BaseAlgo): diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 55f267a..f6d3cd6 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -4,16 +4,19 @@ from abc import ABC, abstractmethod import argparse +import json import multiprocessing as mp import os import pathlib import shutil import sys import time -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import packaging.version +from importlib.resources import files + from . import command_strings as cmds from .dag import DAG from .exceptions import DagExecutionError @@ -22,7 +25,19 @@ 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 ( + PloidyContigs, + detect_reference_build, + par_bed_for_build, +) +from .util import ( + SampleSex, + __version__, + check_version, + cnvscope_sex_args, + path_arg, + tmp, +) MULTIQC_MIN_VERSION = { "multiqc": packaging.version.Version("1.18"), @@ -124,6 +139,9 @@ def __init__(self) -> None: 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""" @@ -194,6 +212,13 @@ def main(self, args: argparse.Namespace) -> None: 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) @@ -278,6 +303,101 @@ 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 build_ploidy_job( + self, + ploidy_json: pathlib.Path, + deduped_bam: List[pathlib.Path], + ploidy_contigs: Optional[PloidyContigs] = None, + ) -> Job: + """Estimate sample ploidy and sex""" + estimate_ploidy = pathlib.Path( + str(files("sentieon_cli.scripts").joinpath("estimate_ploidy.py")) + ).resolve() + ploidy_contigs = ploidy_contigs or PloidyContigs() + ploidy_job = Job( + cmds.cmd_estimate_ploidy( + ploidy_json, + deduped_bam, + 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", + ) + return ploidy_job + + 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 cnv_sex_args(self) -> Tuple[Optional[str], Optional[pathlib.Path]]: + """The CNVscope `--sex` and `--par` arguments for this run""" + return cnvscope_sex_args(self.sample_sex, self.cnv_par_bed) + def multiqc(self) -> Optional[Job]: """Run MultiQC on the metrics files""" 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/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index df4e3c0..8099897 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -17,8 +17,9 @@ 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 .dnascope import CNV_MIN_VERSIONS from .driver import ( CNVModelApply, CNVscope, @@ -33,11 +34,13 @@ from .logging import get_logger from .shell_pipeline import Command, Pipeline from .util import ( + SampleSex, __version__, check_version, check_kmc_patch, parse_rg_line, path_arg, + sample_sex_arg, tmp, total_memory, vcf_id, @@ -46,6 +49,7 @@ GRCH38_CONTIGS, determine_shards_from_fai, parse_fai, + ploidy_contigs_for_build, vcf_contigs, ) from .transfer import build_transfer_jobs @@ -134,6 +138,23 @@ class SentieonPangenome(BasePangenome): "Prefix to strip from pangenome contig names (GRCh38#0#)" ), }, + "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, + }, "skip_metrics": { "help": "Skip metrics collection and multiQC", "action": "store_true", @@ -237,6 +258,8 @@ def __init__(self) -> None: self.bed: Optional[pathlib.Path] = None self.call_svs = False self.gvcf = False + self.par_bed: Optional[pathlib.Path] = None + self.cnv_input_bams: List[pathlib.Path] = [] self.pangenome_ref_name = "GRCh38" self.extract_model_name = "extract.model" self.pangenome_contig_prefix = "GRCh38#0#" @@ -256,6 +279,18 @@ def __init__(self) -> None: self.skip_model_apply = False self.skip_small_variants = False + 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 + + 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 main(self, args: argparse.Namespace) -> None: """Run the pipeline""" self.handle_arguments(args) @@ -287,7 +322,7 @@ def main(self, args: argparse.Namespace) -> None: executor = self.run(dag) self.check_execution(dag, executor) - if self.expansion_catalog or self.segdup_caller is not None: + if self._needs_second_dag(): self.get_sex(self.ploidy_json) dag = self.build_second_dag() executor = self.run(dag) @@ -327,6 +362,7 @@ def validate(self) -> None: self.validate_segdup() self.validate_expansion() self.validate_t1k() + self.validate_cnv() if not self.skip_version_check: for cmd, min_version in SENT_PANGENOME_MIN_VERSIONS.items(): @@ -442,6 +478,22 @@ def validate_expansion(self) -> None: if not check_version(cmd, min_version): sys.exit(2) + 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 + + # A PAR BED file is required for CNV calling, whatever the sex + self.validate_cnv_par(True) + + if not self.skip_version_check: + for cmd, min_version in CNV_MIN_VERSIONS.items(): + if not check_version(cmd, min_version): + sys.exit(2) + def validate_t1k(self) -> None: hla_requested = self.t1k_hla_seq is not None or ( self.t1k_hla_coord is not None @@ -796,13 +848,17 @@ def build_first_dag(self) -> DAG: # Stash the short-read alignment for the second DAG self.sr_alignment = cnv_input_bams[0] - - # 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) + self.cnv_input_bams = cnv_input_bams + + # 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_job = self.build_ploidy_job( + self.ploidy_json, + [cnv_input_bams[0]], + ploidy_contigs_for_build(self.reference_build), + ) + dag.add_job(ploidy_job, cnvscope_deps) # T1K HLA/KIR calling if self.t1k_hla_seq and self.t1k_hla_coord: @@ -855,14 +911,7 @@ def build_first_dag(self) -> DAG: gfa_file=sample_gfa, ) dag.add_job(dnascope_job, dnascope_dependencies) - if 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 apply_dependencies = set() @@ -930,14 +979,7 @@ def build_first_dag(self) -> DAG: gvcftyper_job = self.build_gvcftyper_job(self.output_vcf, out_gvcf) dag.add_job(gvcftyper_job, {small_variants_last_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 @@ -1265,6 +1307,13 @@ def build_second_dag(self) -> DAG: self.logger.info("Building the second pangenome DAG") dag = DAG() + # CNV calling with CNVscope, using the sample sex + if self._cnv_in_second_dag(): + sv_vcf = pathlib.Path( + str(self.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") + ) + self._add_cnv_jobs(dag, sv_vcf, self.cnv_input_bams) + if self.expansion_catalog and self.sr_alignment: out_expansions = pathlib.Path( str(self.output_vcf).replace(".vcf.gz", "_expansion") @@ -1370,8 +1419,6 @@ def _add_cnv_jobs( 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: @@ -1391,9 +1438,10 @@ def _add_cnv_jobs( str(self.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") ) - # CNVscope on BWA deduped BAM + # CNVscope on BWA deduped BAM. The alignment and the PangenomeSV + # output were both written by the first DAG. cnvscope_job = self._build_cnvscope_job(cnvscope_vcf, cnv_input_bams) - dag.add_job(cnvscope_job, cnvscope_deps) + dag.add_job(cnvscope_job) # CNVModelApply on CNVscope output cnv_model_apply_job = self._build_cnv_model_apply_job( @@ -1403,7 +1451,7 @@ def _add_cnv_jobs( # Convert PangenomeSV output to CNVs indel2cnv_job = self._build_indel2cnv_job(indel2cnv_vcf, sv_vcf) - dag.add_job(indel2cnv_job, {dnascope_job}) + dag.add_job(indel2cnv_job) # Combine CNVModelApply output with converted SVs combine_job = self._build_combine_cnv_job( @@ -1418,6 +1466,7 @@ def _build_cnvscope_job( ) -> Job: """Build a CNVscope job""" assert self.model_bundle is not None + sex, par = self.cnv_sex_args() driver = Driver( reference=self.reference, thread_count=self.cores, @@ -1428,6 +1477,8 @@ def _build_cnvscope_job( CNVscope( output=output_vcf, model=self.model_bundle.joinpath("cnv.model"), + sex=sex, + par=par, ) ) return Job( diff --git a/sentieon_cli/shard.py b/sentieon_cli/shard.py index 815f106..4506c40 100644 --- a/sentieon_cli/shard.py +++ b/sentieon_cli/shard.py @@ -7,6 +7,8 @@ import subprocess as sp from typing import Dict, List, NamedTuple, Optional +from importlib.resources import files + from .logging import get_logger logger = get_logger(__name__) @@ -41,6 +43,48 @@ } +# 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 start: int @@ -75,6 +119,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]: diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index fe9cc73..2499def 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -3,6 +3,7 @@ """ import argparse +from enum import Enum from importlib.metadata import PackageNotFoundError, version import multiprocessing as mp import os @@ -12,7 +13,7 @@ import shutil import subprocess as sp import tempfile -from typing import Callable, Dict, List, Optional +from typing import Callable, Dict, List, Optional, Tuple import packaging.version @@ -34,6 +35,47 @@ _UNSAFE = re.compile(r"[^A-Za-z0-9._-]") +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 sanitize(component: str) -> str: """Restrict a path component to filesystem-safe characters""" return _UNSAFE.sub("-", component) diff --git a/tests/integration/test_pipeline_execution.py b/tests/integration/test_pipeline_execution.py index 588ab9b..eba087e 100644 --- a/tests/integration/test_pipeline_execution.py +++ b/tests/integration/test_pipeline_execution.py @@ -155,8 +155,12 @@ def test_hybrid_dry_run_execution_rgsm(self): first_pass_job = [job for job in all_jobs if job.name == "calling-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 diff --git a/tests/unit/test_cnv_sex.py b/tests/unit/test_cnv_sex.py new file mode 100644 index 0000000..fa256b0 --- /dev/null +++ b/tests/unit/test_cnv_sex.py @@ -0,0 +1,375 @@ +""" +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.dnascope import call_cnvs +from sentieon_cli.driver import CNVscope +from sentieon_cli.pipeline import BasePipeline +from sentieon_cli.dag import DAG +from sentieon_cli.shard import ( + BUILD_SIGNATURES, + detect_reference_build, + par_bed_for_build, + ploidy_contigs_for_build, +) +from sentieon_cli.util import ( + SampleSex, + 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) + + +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 cmd[-1] == "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_cnv_sex_args_for_a_male_sample_with_a_par_bed(): + par_bed = pathlib.Path("/par.bed") + pipeline = _pipeline(sample_sex=SampleSex.MALE, cnv_par_bed=par_bed) + + assert pipeline.cnv_sex_args() == ("M", par_bed) + + +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) + + +# `call_cnvs` passes the sample sex through to CNVscope + + +def _cnvscope_cmd(**kwargs) -> str: + cnvscope_job, _ = call_cnvs( + pathlib.Path("/tmp/cnv"), + pathlib.Path("/out/sample.vcf.gz"), + pathlib.Path("/ref/reference.fa"), + [pathlib.Path("/in/sample.cram")], + pathlib.Path("/bundle/model.bundle"), + None, + mp.cpu_count(), + True, # skip_version_check + **kwargs, + ) + return str(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_dnascope_hybrid_cnv.py b/tests/unit/test_dnascope_hybrid_cnv.py new file mode 100644 index 0000000..20c4597 --- /dev/null +++ b/tests/unit/test_dnascope_hybrid_cnv.py @@ -0,0 +1,226 @@ +""" +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""" + with patch( + "sentieon_cli.dnascope_hybrid.check_version", return_value=True + ): + 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""" + with patch( + "sentieon_cli.dnascope_hybrid.check_version", return_value=True + ): + 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 "CNVModelApply" not in jobs + + second_dag, second_jobs = self.build_second_dag(pipeline) + assert "CNVscope" in second_jobs + assert "CNVModelApply" in second_jobs + assert second_dag.waiting_jobs[second_jobs["CNVModelApply"]] == { + 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_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_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_pipeline_lifecycle.py b/tests/unit/test_pipeline_lifecycle.py index aa701e8..9e84888 100644 --- a/tests/unit/test_pipeline_lifecycle.py +++ b/tests/unit/test_pipeline_lifecycle.py @@ -130,6 +130,84 @@ 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")) @@ -219,6 +297,20 @@ def test_pangenome_main_cleans_up_tmpdir_when_the_second_dag_fails( 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 ): diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index 0bf3605..f1d459e 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -14,7 +14,7 @@ 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 @@ -445,13 +445,78 @@ def _get_all_job_names(self, dag): ) 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_first_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_first_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 @@ -468,12 +533,68 @@ def test_cnv_jobs_not_present_without_call_svs(self): assert "indel2cnv" not in job_names assert "combine-cnv" not in job_names + # Nor to the second DAG + pipeline.sample_sex = SampleSex.FEMALE + second_job_names, _ = self._get_all_job_names( + pipeline.build_second_dag() + ) + assert "cnvscope" not in second_job_names + assert "combine-cnv" not in second_job_names + + 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_first_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_first_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_first_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 @@ -502,16 +623,23 @@ def test_call_svs_without_cnv_model(self): assert "indel2cnv" not in job_names assert "combine-cnv" not in job_names + # And they are not in the second DAG either + pipeline.sample_sex = SampleSex.FEMALE + second_job_names, _ = self._get_all_job_names( + pipeline.build_second_dag() + ) + assert "cnvscope" not in second_job_names + assert "combine-cnv" not in second_job_names + 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() + dag = self._build_cnv_dag(pipeline) 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 @@ -521,7 +649,7 @@ 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 @@ -533,14 +661,33 @@ 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 # 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 @@ -557,7 +704,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 @@ -574,7 +721,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 @@ -598,10 +745,15 @@ 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_first_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 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) From bd66813f978bdf4741e4830a00f6225f4ea4f700 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 00:28:54 -0700 Subject: [PATCH 28/58] Emit --interval_padding only when it is set `BaseDriver.interval_padding` now defaults to `None` instead of `0`, so driver commands carry `--interval_padding` only when a caller passes a value. The dnascope pipeline's attribute follows the argparse option, which has no default. The hybrid pipeline's DNAModelApply no longer passes `--interval`, matching every other model-apply invocation. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 2 +- sentieon_cli/dnascope_hybrid.py | 1 - sentieon_cli/driver.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 2a3d837..2885492 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -235,7 +235,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" diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 45a8cbc..dfb20e6 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -1227,7 +1227,6 @@ def call_variants( driver = Driver( reference=self.reference, thread_count=self.cores, - interval=self.bed, ) driver.add_algo( DNAModelApply( diff --git a/sentieon_cli/driver.py b/sentieon_cli/driver.py index 4574df8..e442a89 100644 --- a/sentieon_cli/driver.py +++ b/sentieon_cli/driver.py @@ -575,7 +575,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, From 6f7d77764653adc13a1cf326d788efd1b0becde5 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 00:35:47 -0700 Subject: [PATCH 29/58] Add the stage building blocks Introduce `sentieon_cli/stages/`, a layer of reusable pipeline stages that build their jobs and insert them into the DAG with the right edges. `stages/base.py` defines `StageContext` (run-wide settings), `StageResult` (the inserted jobs and the terminal jobs downstream stages depend on), the `Stage` base class with its `add_to` contract, and the `driver_job` / `rm_job` helpers for the two most common job shapes. `BasePipeline` gains `stage_context()` and `required()`, and `util` gains `require_versions()` / `versions_available()` for the version check idioms. Nothing uses them yet; the pipelines migrate stage by stage. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/pipeline.py | 46 +++- sentieon_cli/stages/__init__.py | 1 + sentieon_cli/stages/base.py | 127 +++++++++++ sentieon_cli/util.py | 42 +++- tests/unit/test_stage_base.py | 364 ++++++++++++++++++++++++++++++++ 5 files changed, 578 insertions(+), 2 deletions(-) create mode 100644 sentieon_cli/stages/__init__.py create mode 100644 sentieon_cli/stages/base.py create mode 100644 tests/unit/test_stage_base.py diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index f6d3cd6..a2732a4 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -11,7 +11,7 @@ import shutil import sys import time -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, TypeVar import packaging.version @@ -30,6 +30,7 @@ detect_reference_build, par_bed_for_build, ) +from .stages.base import StageContext from .util import ( SampleSex, __version__, @@ -45,6 +46,8 @@ BWA_INDEX_SUFFIXES = (".amb", ".ann", ".bwt", ".pac", ".sa") +T = TypeVar("T") + class BasePipeline(ABC): """A pipeline base class""" @@ -246,6 +249,47 @@ def check_execution( 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 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/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/util.py b/sentieon_cli/util.py index 2499def..67fa1c4 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -12,8 +12,9 @@ import shlex import shutil import subprocess as sp +import sys import tempfile -from typing import Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Mapping, Optional, Tuple import packaging.version @@ -137,6 +138,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 pipelines' 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, 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 From 54913fd174ab6d2fea09f617198446c3f1839b49 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 00:45:39 -0700 Subject: [PATCH 30/58] Move CNV calling and ploidy estimation into stages `CNVscopeStage` (CNVscope followed by CNVModelApply) replaces the module-level `call_cnvs` used by dnascope-hybrid and the two private builders in dnascope-pangenome; `PloidyStage` replaces `BasePipeline.build_ploidy_job` and derives the `_ploidy.json` path itself. `CNV_MIN_VERSIONS` moves to `stages/cnv.py`. The duplicate version check that ran inside `call_cnvs` is gone; `validate_cnv` already performs it before the first DAG runs. The hybrid pipeline's CNV jobs are now named `cnvscope` and `cnv-model-apply`, matching the pangenome pipeline. Commands and dependency edges are unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 80 +------ sentieon_cli/dnascope_hybrid.py | 46 ++-- sentieon_cli/pipeline.py | 30 --- sentieon_cli/sentieon_pangenome.py | 102 ++------- sentieon_cli/stages/cnv.py | 95 ++++++++ sentieon_cli/stages/ploidy.py | 67 ++++++ tests/integration/test_pipeline_execution.py | 2 +- tests/unit/test_cnv_sex.py | 36 +-- tests/unit/test_dnascope_hybrid_cnv.py | 18 +- tests/unit/test_stage_cnv.py | 217 +++++++++++++++++++ tests/unit/test_stage_ploidy.py | 130 +++++++++++ 11 files changed, 589 insertions(+), 234 deletions(-) create mode 100644 sentieon_cli/stages/cnv.py create mode 100644 sentieon_cli/stages/ploidy.py create mode 100644 tests/unit/test_stage_cnv.py create mode 100644 tests/unit/test_stage_ploidy.py diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 2885492..f006666 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -5,12 +5,11 @@ 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, Tuple import packaging.version @@ -21,8 +20,6 @@ from .driver import ( AlignmentStat, BaseDistributionByCycle, - CNVscope, - CNVModelApply, CoverageMetrics, Dedup, DNAModelApply, @@ -42,9 +39,7 @@ from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline from .util import ( - SampleSex, check_version, - cnvscope_sex_args, library_preloaded, parse_rg_line, path_arg, @@ -65,11 +60,6 @@ "sentieon driver": packaging.version.Version("202308"), } -CNV_MIN_VERSIONS = { - # 202503.04 adds the CNVscope `--sex` and `--par` arguments - "sentieon driver": packaging.version.Version("202503.04"), -} - class DNAscopePipeline(BasePipeline): """The DNAscope pipeline""" @@ -986,71 +976,3 @@ def sr_call_variants( 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, - sample_sex: Optional[SampleSex] = None, - par_bed: Optional[pathlib.Path] = 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) - - sex, par = cnvscope_sex_args(sample_sex, par_bed) - 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"), - sex=sex, - par=par, - ) - ) - cnvscope_job = Job( - Pipeline(Command(*driver.build_cmd())), - "CNVscope", - cores, - task_name="cnv", - ) - - 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, - task_name="cnv", - ) - - return (cnvscope_job, cnvmodelapply_job) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index dfb20e6..c29101e 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -26,11 +26,13 @@ HybridStage2, HybridStage3, ) -from .dnascope import call_cnvs, CNV_MIN_VERSIONS, DNAscopePipeline +from .dnascope import DNAscopePipeline from .dnascope_longread import DNAscopeLRPipeline from .job import Job from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline +from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage +from .stages.ploidy import PloidyStage from .util import ( __version__, check_version, @@ -44,7 +46,6 @@ from .shard import ( determine_shards_from_fai, parse_fai, - ploidy_contigs_for_build, vcf_contigs, ) from .transfer import build_transfer_jobs @@ -579,6 +580,7 @@ def build_dag(self) -> DAG: if not self.model_bundle: self.logger.error("model_bundle is required") sys.exit(2) + ctx = self.stage_context() rg_info = RgInfo( self.lr_aln_readgroups, @@ -662,15 +664,12 @@ def build_dag(self) -> DAG: # 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. - self.ploidy_json = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_ploidy.json") - ) - ploidy_job = self.build_ploidy_job( - self.ploidy_json, - sr_aln, - ploidy_contigs_for_build(self.reference_build), - ) - dag.add_job(ploidy_job, sr_preprocessing_jobs) + 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: # CNV calling is sex-aware and runs in the second DAG @@ -758,26 +757,25 @@ def build_second_dag(self) -> Optional[DAG]: self.logger.error("model_bundle is required") sys.exit(2) 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() - cnvscope_job, cnvmodelapply_job = call_cnvs( - self.tmp_dir, - self.output_vcf, - self.reference, - self.cnv_sr_aln, - self.model_bundle, - self.bed, - self.cores, - self.skip_version_check, - replace_rg=self.cnv_replace_rg, + CNVscopeStage( + ctx=ctx, + inputs=self.cnv_sr_aln, + model=self.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, - ) - dag.add_job(cnvscope_job) - dag.add_job(cnvmodelapply_job, {cnvscope_job}) + interval=self.bed, + replace_rg=self.cnv_replace_rg, + ).add_to(dag) return dag def call_variants( diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index a2732a4..3638390 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -15,8 +15,6 @@ import packaging.version -from importlib.resources import files - from . import command_strings as cmds from .dag import DAG from .exceptions import DagExecutionError @@ -26,7 +24,6 @@ from .run_logs import RunLogs from .scheduler import ThreadScheduler from .shard import ( - PloidyContigs, detect_reference_build, par_bed_for_build, ) @@ -356,33 +353,6 @@ def build_second_dag(self) -> Optional[DAG]: """ return None - def build_ploidy_job( - self, - ploidy_json: pathlib.Path, - deduped_bam: List[pathlib.Path], - ploidy_contigs: Optional[PloidyContigs] = None, - ) -> Job: - """Estimate sample ploidy and sex""" - estimate_ploidy = pathlib.Path( - str(files("sentieon_cli.scripts").joinpath("estimate_ploidy.py")) - ).resolve() - ploidy_contigs = ploidy_contigs or PloidyContigs() - ploidy_job = Job( - cmds.cmd_estimate_ploidy( - ploidy_json, - deduped_bam, - 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", - ) - return ploidy_job - def get_sex(self, ploidy_json: pathlib.Path) -> None: """Retrieve the sample sex""" if self.sample_sex is not None: diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 8099897..2002bc3 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -19,10 +19,7 @@ from .archive import ar_load from .base_pangenome import BasePangenome from .dag import DAG -from .dnascope import CNV_MIN_VERSIONS from .driver import ( - CNVModelApply, - CNVscope, DNAModelApply, DNAscope, Driver, @@ -33,6 +30,8 @@ from .job import Job from .logging import get_logger from .shell_pipeline import Command, Pipeline +from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage +from .stages.ploidy import PloidyStage from .util import ( SampleSex, __version__, @@ -49,7 +48,6 @@ GRCH38_CONTIGS, determine_shards_from_fai, parse_fai, - ploidy_contigs_for_build, vcf_contigs, ) from .transfer import build_transfer_jobs @@ -679,6 +677,8 @@ def build_first_dag(self) -> DAG: self.logger.error("output_vcf is required") sys.exit(2) + ctx = self.stage_context() + self.logger.info("Building the Sentieon pangenome DAG") dag = DAG() @@ -690,9 +690,6 @@ def build_first_dag(self) -> DAG: 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") - ) out_gvcf = pathlib.Path( str(self.output_vcf).replace(".vcf.gz", ".g.vcf.gz") ) @@ -853,12 +850,12 @@ def build_first_dag(self) -> DAG: # 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_job = self.build_ploidy_job( - self.ploidy_json, - [cnv_input_bams[0]], - ploidy_contigs_for_build(self.reference_build), - ) - dag.add_job(ploidy_job, cnvscope_deps) + ploidy_result = PloidyStage( + ctx=ctx, + inputs=[cnv_input_bams[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: @@ -1431,6 +1428,7 @@ def _add_cnv_jobs( self.logger.error("output_vcf is required") sys.exit(2) + ctx = self.stage_context() 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") @@ -1438,16 +1436,18 @@ def _add_cnv_jobs( str(self.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") ) - # CNVscope on BWA deduped BAM. The alignment and the PangenomeSV - # output were both written by the first DAG. - cnvscope_job = self._build_cnvscope_job(cnvscope_vcf, cnv_input_bams) - dag.add_job(cnvscope_job) - - # 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}) + # CNVscope on BWA deduped BAM, then CNVModelApply. The alignment + # and the PangenomeSV output were both written by the first DAG. + cnv_result = CNVscopeStage( + ctx=ctx, + inputs=cnv_input_bams, + model=self.model_bundle.joinpath("cnv.model"), + cnvscope_vcf=cnvscope_vcf, + cnv_vcf=cnv_apply_vcf, + sample_sex=self.sample_sex, + par_bed=self.cnv_par_bed, + interval=self.bed, + ).add_to(dag) # Convert PangenomeSV output to CNVs indel2cnv_job = self._build_indel2cnv_job(indel2cnv_vcf, sv_vcf) @@ -1457,61 +1457,7 @@ def _add_cnv_jobs( 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 - sex, par = self.cnv_sex_args() - 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"), - sex=sex, - par=par, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "cnvscope", - self.cores, - task_name="cnv", - ) - - 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, - task_name="cnv", - ) + dag.add_job(combine_job, {cnv_result.apply_job, indel2cnv_job}) def _build_indel2cnv_job( self, diff --git a/sentieon_cli/stages/cnv.py b/sentieon_cli/stages/cnv.py new file mode 100644 index 0000000..91904ad --- /dev/null +++ b/sentieon_cli/stages/cnv.py @@ -0,0 +1,95 @@ +""" +CNV calling with CNVscope +""" + +from dataclasses import dataclass +import pathlib +from typing import Iterable, List, Optional + +import packaging.version + +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. + """ + + 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 + 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, + ) + ], + 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, + ) 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/tests/integration/test_pipeline_execution.py b/tests/integration/test_pipeline_execution.py index eba087e..e7a25c1 100644 --- a/tests/integration/test_pipeline_execution.py +++ b/tests/integration/test_pipeline_execution.py @@ -160,7 +160,7 @@ def test_hybrid_dry_run_execution_rgsm(self): 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 second_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 diff --git a/tests/unit/test_cnv_sex.py b/tests/unit/test_cnv_sex.py index fa256b0..914d5a2 100644 --- a/tests/unit/test_cnv_sex.py +++ b/tests/unit/test_cnv_sex.py @@ -12,10 +12,11 @@ import pytest from sentieon_cli import shard -from sentieon_cli.dnascope import call_cnvs 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, @@ -335,22 +336,31 @@ def test_resolve_cnv_par_bed_skips_the_lookup_without_cnv_calling(messages): assert not any("PAR" in msg for msg in messages) -# `call_cnvs` passes the sample sex through to CNVscope +# `CNVscopeStage` passes the sample sex through to CNVscope def _cnvscope_cmd(**kwargs) -> str: - cnvscope_job, _ = call_cnvs( - pathlib.Path("/tmp/cnv"), - pathlib.Path("/out/sample.vcf.gz"), - pathlib.Path("/ref/reference.fa"), - [pathlib.Path("/in/sample.cram")], - pathlib.Path("/bundle/model.bundle"), - None, - mp.cpu_count(), - True, # skip_version_check - **kwargs, + 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, ) - return str(cnvscope_job.shell) + 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(): diff --git a/tests/unit/test_dnascope_hybrid_cnv.py b/tests/unit/test_dnascope_hybrid_cnv.py index 20c4597..924a712 100644 --- a/tests/unit/test_dnascope_hybrid_cnv.py +++ b/tests/unit/test_dnascope_hybrid_cnv.py @@ -105,14 +105,14 @@ def test_cnv_jobs_run_in_the_second_dag(self): pipeline.sample_sex = SampleSex.FEMALE _dag, jobs = self.build_dag(pipeline) - assert "CNVscope" not in jobs - assert "CNVModelApply" not in jobs + 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 "CNVModelApply" in second_jobs - assert second_dag.waiting_jobs[second_jobs["CNVModelApply"]] == { - second_jobs["CNVscope"] + 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): @@ -164,7 +164,7 @@ def test_a_female_sample_is_called_without_a_par_bed(self): self.build_dag(pipeline) _dag, jobs = self.build_second_dag(pipeline) - cmd_str = str(jobs["CNVscope"].shell) + cmd_str = str(jobs["cnvscope"].shell) assert "--sex F" in cmd_str assert "--par" not in cmd_str @@ -175,7 +175,7 @@ def test_a_male_sample_is_called_with_the_par_bed(self): self.build_dag(pipeline) _dag, jobs = self.build_second_dag(pipeline) - cmd_str = str(jobs["CNVscope"].shell) + cmd_str = str(jobs["cnvscope"].shell) assert "--sex M" in cmd_str assert f"--par {self.mock_par_bed}" in cmd_str @@ -187,7 +187,7 @@ def test_a_dry_run_assumes_a_male_sample(self): _dag, jobs = self.build_second_dag(pipeline) assert pipeline.sample_sex is SampleSex.MALE - cmd_str = str(jobs["CNVscope"].shell) + cmd_str = str(jobs["cnvscope"].shell) assert "--sex M" in cmd_str assert f"--par {self.mock_par_bed}" in cmd_str diff --git a/tests/unit/test_stage_cnv.py b/tests/unit/test_stage_cnv.py new file mode 100644 index 0000000..3fbf962 --- /dev/null +++ b/tests/unit/test_stage_cnv.py @@ -0,0 +1,217 @@ +""" +Unit tests for the CNVscope stage +""" + +import logging +import pathlib +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 +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_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 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} From c64e9419dc6de7f207812df472fc0324fe2fd9c1 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 00:58:16 -0700 Subject: [PATCH 31/58] Move duplicate marking and metrics into stages `DedupStage` (LocusCollector followed by Dedup) replaces the short-read `dedup_and_metrics` builder and the pangenome `build_dedup_job`; `MetricsStage` (one metrics pass plus the optional WGS rehead job) replaces the post-dedup metrics code in both places. `MetricsPaths` derives the metrics directory and every metrics file name once, instead of at five call sites. `DNAscopePipeline.add_sr_preprocessing` composes the two stages and inserts the jobs itself, returning the deduplicated alignment, the Dedup job, and the QC jobs MultiQC waits on; the hybrid pipeline inherits it. Renames: the rehead job is `rehead-metrics`, and with `--duplicate_marking none` the metrics-only pass is now the `metrics` job under the `metrics` task rather than `locuscollector` under `dedup`. Commands and dependency edges are unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/base_pangenome.py | 150 ++----------- sentieon_cli/dnascope.py | 324 +++++++++++----------------- sentieon_cli/dnascope_hybrid.py | 35 +-- sentieon_cli/hybrid_pangenome.py | 59 ++--- sentieon_cli/pipeline.py | 7 +- sentieon_cli/sentieon_pangenome.py | 64 +++--- sentieon_cli/stages/dedup.py | 94 ++++++++ sentieon_cli/stages/metrics.py | 181 ++++++++++++++++ tests/unit/test_hybrid_pangenome.py | 6 +- tests/unit/test_stage_dedup.py | 200 +++++++++++++++++ tests/unit/test_stage_metrics.py | 244 +++++++++++++++++++++ 11 files changed, 939 insertions(+), 425 deletions(-) create mode 100644 sentieon_cli/stages/dedup.py create mode 100644 sentieon_cli/stages/metrics.py create mode 100644 tests/unit/test_stage_dedup.py create mode 100644 tests/unit/test_stage_metrics.py diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index c617c48..f21e9a6 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -4,28 +4,23 @@ import copy import pathlib -import sys -from typing import List, Optional, Tuple - -from importlib.resources import files +from typing import List, Optional from . import command_strings as cmds from .driver import ( AlignmentStat, + BaseAlgo, BaseDistributionByCycle, CoverageMetrics, - Dedup, - Driver, GCBias, InsertSizeMetricAlgo, - LocusCollector, MeanQualityByCycle, QualDistribution, WgsMetricsAlgo, ) from .job import Job from .pipeline import BasePipeline -from .shell_pipeline import Command, Pipeline +from .stages.metrics import MetricsPaths from .util import path_arg @@ -141,130 +136,15 @@ def build_kmc_job( return kmc_job - def build_dedup_job( - self, - output_bam, - input_bam: List[pathlib.Path], - tag: str, - left_align_rgid: Optional[str] = None, - 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_rgid: - read_filters.append( - f"IndelLeftAlignReadTransform,rgid={left_align_rgid}" - ) - - # 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, - task_name="dedup", - ) - - 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, - task_name="dedup", - ) - - 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, - task_name="metrics", - ) - - 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, - task_name="metrics", - ) - return (metrics_job, rehead_job) + 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/dnascope.py b/sentieon_cli/dnascope.py index f006666..4ffb8fd 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -4,6 +4,7 @@ import argparse import copy +from dataclasses import dataclass, field import itertools import os import pathlib @@ -13,15 +14,13 @@ import packaging.version -from importlib.resources import files - from . import command_strings as cmds from .dag import DAG from .driver import ( AlignmentStat, + BaseAlgo, BaseDistributionByCycle, CoverageMetrics, - Dedup, DNAModelApply, DNAscope, Driver, @@ -30,7 +29,6 @@ HsMetricAlgo, InsertSizeMetricAlgo, MeanQualityByCycle, - LocusCollector, QualDistribution, SVSolver, WgsMetricsAlgo, @@ -38,6 +36,9 @@ from .job import Job from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline +from .stages.base import StageContext +from .stages.dedup import DedupStage +from .stages.metrics import MetricsPaths, MetricsStage from .util import ( check_version, library_preloaded, @@ -61,6 +62,21 @@ } +@dataclass +class SrPreprocessing: + """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] = None + qc_jobs: List[Job] = field(default_factory=list) + + class DNAscopePipeline(BasePipeline): """The DNAscope pipeline""" @@ -384,6 +400,7 @@ 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() @@ -399,25 +416,16 @@ def build_dag(self) -> DAG: sample_input += aligned_fastq # 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 + if dedup_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}) # Small variants if not self.skip_small_variants: @@ -447,16 +455,8 @@ def build_dag(self) -> DAG: # 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]: @@ -629,188 +629,116 @@ def sr_align_fastq( return (res, jobs, rm_job) - def dedup_and_metrics( + def sr_metrics_algos(self, paths: MetricsPaths) -> 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 self.assay == "WGS": + algos.append(GCBias(paths.gc_bias, summary=paths.gc_bias_summary)) + return algos + + def add_sr_preprocessing( self, - sample_input, - ) -> Tuple[ - List[pathlib.Path], - Optional[Job], - Optional[Job], - Optional[Job], - Optional[Job], - ]: - """Perform dedup and metrics collection""" + dag: DAG, + ctx: StageContext, + sample_input: List[pathlib.Path], + upstream: Set[Job], + ) -> SrPreprocessing: + """Mark duplicates and collect metrics from the short reads""" 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") + paths = MetricsPaths.from_output_vcf(self.output_vcf) + paths.ensure_dir(self.dry_run) - # 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, + if self.sr_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 SrPreprocessing( + deduped=sample_input, dedup_job=None, qc_jobs=[] ) + metrics_result = MetricsStage( + ctx=ctx, + inputs=sample_input, + algos=[ + InsertSizeMetricAlgo(paths.insert_size), + *self.sr_metrics_algos(paths), + ], + name="metrics", + task_name="metrics", + threads=ctx.cores, + ).add_to(dag, upstream) + return SrPreprocessing( + deduped=sample_input, + dedup_job=None, + qc_jobs=list(metrics_result.jobs), ) - # 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)) - + # Metrics that ride along on the LocusCollector pass + lc_extra_algos: List[BaseAlgo] = [] 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, - task_name="dedup", - ) + # 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(self.sr_metrics_algos(paths)) - 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, - task_name="dedup", - ) - - 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, + dedup_result = DedupStage( + ctx=ctx, + inputs=sample_input, + output=deduped, + score_file=paths.score, + consensus=self.consensus, + rmdup=(self.sr_duplicate_marking == "rmdup"), + cram_write_options="version=3.0,compressor=rans", + dedup_metrics=paths.dedup_metrics, + lc_extra_algos=lc_extra_algos, + ).add_to(dag, upstream) + 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}) + qc_jobs.extend(metrics_result.jobs) + + return SrPreprocessing( + deduped=[deduped], + dedup_job=dedup_result.dedup_job, + qc_jobs=qc_jobs, ) - 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, - task_name="metrics", - ) # 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, - task_name="metrics", - ) # 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, - task_name="metrics", - ) - return ([deduped], lc_job, dedup_job, metrics_job, rehead_job) def sr_call_variants( self, diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index c29101e..4036279 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -604,26 +604,13 @@ def build_dag(self) -> DAG: sr_aln.extend(aligned_fastq) # 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 + if dedup_job and fq_rm_job: + dag.add_job(fq_rm_job, {dedup_job}) # Long-read alignment realign_jobs: Set[Job] = set() @@ -640,16 +627,8 @@ def build_dag(self) -> DAG: 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( diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 0f33b8e..bf8f748 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -36,6 +36,8 @@ vcf_contigs, ) from .shell_pipeline import Command, Pipeline +from .stages.dedup import DedupStage +from .stages.metrics import MetricsPaths, MetricsStage from .transfer import build_transfer_jobs from .util import ( __version__, @@ -617,6 +619,7 @@ def build_dag(self) -> DAG: self.logger.info("Building the hybrid-pangenome DAG") dag = DAG() + ctx = self.stage_context() ref_fai = pathlib.Path(str(self.reference) + ".fai") @@ -762,42 +765,42 @@ def build_dag(self) -> DAG: # 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(self.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 - ) - lift_lc_job, lift_dedup_job = self.build_dedup_job( - out_lift_aln, - [lift_bam], - "lift", - ) - dag.add_job(bwa_lc_job, {bwa_job}) - dag.add_job(bwa_dedup_job, {bwa_lc_job}) - dag.add_job(lift_lc_job, {mm2_job}) - dag.add_job(lift_dedup_job, {lift_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_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_job, rehead_job = self.build_metrics_job( - [out_bwa_aln, out_lift_aln], - ) - dag.add_job(metrics_job, {bwa_dedup_job, lift_dedup_job}) - dag.add_job(rehead_job, {metrics_job}) + 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, {rehead_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} diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 3638390..428f3af 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -28,6 +28,7 @@ par_bed_for_build, ) from .stages.base import StageContext +from .stages.metrics import MetricsPaths from .util import ( SampleSex, __version__, @@ -428,9 +429,9 @@ def multiqc(self) -> Optional[Job]: ) 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, diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 2002bc3..f2f8d60 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -31,6 +31,8 @@ from .logging import get_logger from .shell_pipeline import Command, Pipeline from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage +from .stages.dedup import DedupStage +from .stages.metrics import MetricsPaths, MetricsStage from .stages.ploidy import PloidyStage from .util import ( SampleSex, @@ -798,47 +800,49 @@ 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(self.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_rgid=f"{self.fastq_readgroup['ID']}-mm2", - ) - 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] cnvscope_deps = {bwa_dedup_job} if not self.skip_metrics: - metrics_job, rehead_job = self.build_metrics_job( - [out_bwa_aln], - ) - dag.add_job(metrics_job, {bwa_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) 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/metrics.py b/sentieon_cli/stages/metrics.py new file mode 100644 index 0000000..ee9d82b --- /dev/null +++ b/sentieon_cli/stages/metrics.py @@ -0,0 +1,181 @@ +""" +Alignment metrics collection +""" + +from dataclasses import dataclass +from importlib.resources import files +import pathlib +import sys +from typing import Iterable, List, Optional + +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 + + +@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, + ) diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index 42c836c..9239e63 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -273,7 +273,7 @@ def test_metrics_job(self): ) assert "--algo GCBias" in cmd_str assert "--algo WgsMetricsAlgo" in cmd_str - assert "Rehead metrics" in job_names + assert "rehead-metrics" in job_names def test_skip_metrics(self): """skip_metrics removes the metrics jobs and the Dedup metrics""" @@ -283,7 +283,7 @@ def test_skip_metrics(self): job_names, all_jobs = self._get_all_job_names(dag) assert "metrics" not in job_names - assert "Rehead 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 @@ -561,7 +561,7 @@ def test_aligned_dag_jobs(self): "locuscollector-lift", "dedup-lift", "metrics", - "Rehead metrics", + "rehead-metrics", "multiqc", ): assert name not in job_names, f"unexpected job: {name}" 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_metrics.py b/tests/unit/test_stage_metrics.py new file mode 100644 index 0000000..1792560 --- /dev/null +++ b/tests/unit/test_stage_metrics.py @@ -0,0 +1,244 @@ +""" +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 + + +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 From b0b3a7334321daed74d324e3d0b4733c7ec1bb2b Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 01:08:24 -0700 Subject: [PATCH 32/58] Move small-variant calling into stages Add the small-variant stages: `DNAscopeStage` (any DNAscope-family pass), `ModelApplyStage`, `GVCFtyperStage`, `TransferStage` (the sharded annotation transfer, wrapping `build_transfer_jobs`, which gains a `tag` for pipelines that transfer more than once), and the `TransferApplyStage` composite for the common "transfer, then apply the model" tail, whose DNAModelApply waits on the raw-VCF producer as well as the transfer. The dnascope pipeline's `sr_call_variants` (a 6-tuple) becomes `add_small_variant_calling`, and the pangenome pipelines drop their `build_dnascope_job` / `build_dnamodelapply_job` / `build_gvcftyper_job` builders. The DNAscope job is now named `dnascope` everywhere (was `variant-calling` and `dnascope-raw`). Commands and dependency edges are unchanged, except that a dnascope-pangenome run without a population VCF now feeds DNAModelApply the raw VCF rather than the transfer output that no job would have written; validation requires `--pop_vcf`, so real runs are unaffected. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 163 +++----- sentieon_cli/hybrid_pangenome.py | 125 +++--- sentieon_cli/sentieon_pangenome.py | 246 +++++------- sentieon_cli/stages/small_variants.py | 235 +++++++++++ sentieon_cli/stages/transfer.py | 80 ++++ sentieon_cli/transfer.py | 17 +- tests/integration/test_pipeline_execution.py | 4 +- tests/unit/test_dag_construction.py | 2 +- tests/unit/test_hybrid_pangenome.py | 36 +- tests/unit/test_sentieon_pangenome.py | 24 +- tests/unit/test_stage_small_variants.py | 385 +++++++++++++++++++ tests/unit/test_stage_transfer.py | 133 +++++++ 12 files changed, 1057 insertions(+), 393 deletions(-) create mode 100644 sentieon_cli/stages/small_variants.py create mode 100644 sentieon_cli/stages/transfer.py create mode 100644 tests/unit/test_stage_small_variants.py create mode 100644 tests/unit/test_stage_transfer.py diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 4ffb8fd..2467fbe 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -21,11 +21,8 @@ BaseAlgo, BaseDistributionByCycle, CoverageMetrics, - DNAModelApply, DNAscope, - Driver, GCBias, - GVCFtyper, HsMetricAlgo, InsertSizeMetricAlgo, MeanQualityByCycle, @@ -36,9 +33,15 @@ from .job import Job from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline -from .stages.base import StageContext +from .stages.base import StageContext, driver_job, rm_job from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + GVCFtyperStage, + TransferApplyStage, +) from .util import ( check_version, library_preloaded, @@ -429,28 +432,15 @@ def build_dag(self) -> DAG: # 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: @@ -740,10 +730,13 @@ def add_sr_preprocessing( qc_jobs=qc_jobs, ) - def sr_call_variants( + 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") @@ -766,7 +759,6 @@ def sr_call_variants( out_svs = pathlib.Path( str(self.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") @@ -782,14 +774,8 @@ def sr_call_variants( tmp_vcf = pathlib.Path( str(self.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, @@ -797,110 +783,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, - task_name="variant-calling", - ) + 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, - task_name="model-apply", - ) + 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, - task_name="cleanup", + 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=self.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, - task_name="gvcftyper", - ) + ).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, - ) - driver.add_algo( - SVSolver( - output=out_svs, - vcf=out_svs_tmp, - ) - ) - svsolver_job = Job( - Pipeline(Command(*driver.build_cmd())), - "svsolver", + threads=1, + name="svsolver", task_name="sv-calling", ) - sv_rm_job = Job( - Pipeline( - Command( - "rm", - str(out_svs_tmp), - str(out_svs_tmp) + ".tbi", - fail_ok=True, - ) - ), - "rm-tmp-sv", - 0, - task_name="cleanup", + 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}, ) - - return ( - call_job, - apply_job, - rm_job, - gvcftyper_job, - svsolver_job, - sv_rm_job, - ) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index bf8f748..cc0e2f1 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -20,7 +20,6 @@ from .base_pangenome import BasePangenome from .dag import DAG from .driver import ( - DNAModelApply, DNAscope, Driver, LongReadSV, @@ -38,7 +37,8 @@ from .shell_pipeline import Command, Pipeline from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage -from .transfer import build_transfer_jobs +from .stages.small_variants import DNAscopeStage, ModelApplyStage +from .stages.transfer import TransferConfig, TransferStage from .util import ( __version__, check_kmc_patch, @@ -825,40 +825,57 @@ def build_dag(self) -> DAG: if self.skip_small_variants: return dag - dnascope_job = self.build_dnascope_job( - raw_vcf, - calling_bams, - replace_rg, - ) - dag.add_job(dnascope_job, calling_dependencies) + model = self.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 transfer_target = ( transfer_vcf if not self.skip_model_apply else self.output_vcf ) - transfer_jobs, concat_job = build_transfer_jobs( - transfer_target, - 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)) + transfer = TransferStage( + ctx=ctx, + config=self.transfer_config(), + raw_vcf=raw_vcf, + out_vcf=transfer_target, + ).add_to(dag, call.terminal) if not self.skip_model_apply: - apply_job = self.build_dnamodelapply_job( - transfer_vcf, self.output_vcf - ) - dag.add_job(apply_job, {concat_job}) + ModelApplyStage( + ctx=ctx, + model=model, + vcf=transfer_vcf, + output=self.output_vcf, + ).add_to(dag, transfer.terminal) return dag + def transfer_config(self) -> TransferConfig: + """The inputs the annotation-transfer stage needs. + + `build_dag` requires `--pop_vcf` before it calls this. + """ + assert self.pop_vcf is not None + return TransferConfig( + pop_vcf=self.pop_vcf, + shards=self.shards, + pop_vcf_contigs=self.pop_vcf_contigs, + fai_data=self.fai_data, + ) + def lr_kmc_pairs(self) -> List[Tuple[pathlib.Path, pathlib.Path]]: """`(alignment, decode reference)` pairs for the long-read k-mer counting. @@ -1206,59 +1223,3 @@ def build_pangenomesv_job( self.cores, task_name="sv-calling", ) - - def build_dnascope_job( - self, - out_vcf: pathlib.Path, - input_bams: List[pathlib.Path], - replace_rg: List[List[str]], - ) -> Job: - """Call small variants with the original, lifted, and long reads""" - assert self.model_bundle is not None - 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, - replace_rg=replace_rg, - ) - driver.add_algo( - DNAscope( - out_vcf, - model=self.model_bundle.joinpath("dnascope.model"), - pcr_indel_model=pcr_indel_model, - dbsnp=self.dbsnp, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "dnascope-raw", - self.cores, - task_name="variant-calling", - ) - - def build_dnamodelapply_job( - self, - in_vcf: pathlib.Path, - out_vcf: pathlib.Path, - ) -> Job: - """Apply the DNAscope model""" - assert self.model_bundle is not None - 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, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "model-apply", - self.cores, - task_name="model-apply", - ) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index f2f8d60..ccd4877 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -20,10 +20,9 @@ from .base_pangenome import BasePangenome from .dag import DAG from .driver import ( - DNAModelApply, + BaseAlgo, DNAscope, Driver, - GVCFtyper, PangenomeSV, ReadWriter, ) @@ -34,6 +33,14 @@ from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage from .stages.ploidy import PloidyStage +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + GVCFtyperStage, + TransferApplyStage, + TransferSpec, +) +from .stages.transfer import TransferConfig from .util import ( SampleSex, __version__, @@ -52,7 +59,6 @@ parse_fai, vcf_contigs, ) -from .transfer import build_transfer_jobs SENT_PANGENOME_MIN_VERSIONS = { "kmc": None, @@ -903,27 +909,48 @@ def build_first_dag(self) -> DAG: 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 = self.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", + ) + ) + if sv_vcf and gfa_file: + algos.append( + PangenomeSV( + sv_vcf, + gfa_file=gfa_file, + prefix=self.pangenome_contig_prefix, + ) ) - dag.add_job(dnascope_job, dnascope_dependencies) - # CNV calling is sex-aware and runs in the second DAG + 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. @@ -934,56 +961,60 @@ def build_first_dag(self) -> DAG: 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=self.transfer_config(), + out_vcf=( transfer_vcf if not self.skip_model_apply - else (snv_apply_vcf) + 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, snv_apply_vcf - ) - 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 - if small_variants_last_job is not None: - ad_update_job = self.build_count_ad_update_job( - small_variants_out, snv_apply_vcf - ) - dag.add_job(ad_update_job, {small_variants_last_job}) - small_variants_last_job = ad_update_job + 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=self.output_vcf, + interval=self.bed, + ).add_to(dag, {ad_update_job}) # CNV calling is sex-aware and runs in the second DAG return dag + def transfer_config(self) -> TransferConfig: + """The inputs the annotation-transfer stage needs. + + Only valid once `--pop_vcf` is known to be set; `validate_bundle` + requires it. + """ + assert self.pop_vcf is not None + return TransferConfig( + pop_vcf=self.pop_vcf, + shards=self.shards, + pop_vcf_contigs=self.pop_vcf_contigs, + fai_data=self.fai_data, + ) + def build_alignment_job( self, sample_bam: pathlib.Path, @@ -1134,109 +1165,6 @@ def build_minimap2_lift_job( ) return mm2_job - def build_dnascope_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, - prefix=self.pangenome_contig_prefix, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "dnascope-raw", - self.cores, - task_name="variant-calling", - ) - - 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) - - 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, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "model-apply", - self.cores, - task_name="model-apply", - ) - - 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, - task_name="gvcftyper", - ) - def build_count_ad_update_job( self, out_vcf: pathlib.Path, diff --git a/sentieon_cli/stages/small_variants.py b/sentieon_cli/stages/small_variants.py new file mode 100644 index 0000000..9757781 --- /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 *and* on ``upstream`` -- + the raw-VCF producer stays a direct dependency, matching the DAGs the + pipelines built by hand. + """ + + 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 = 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/transfer.py b/sentieon_cli/stages/transfer.py new file mode 100644 index 0000000..7ba294b --- /dev/null +++ b/sentieon_cli/stages/transfer.py @@ -0,0 +1,80 @@ +""" +Annotation transfer from a population VCF +""" + +from dataclasses import dataclass +import pathlib +from typing import Dict, Iterable, List, Optional + +from ..dag import DAG +from ..job import Job +from ..shard import Shard +from ..transfer import build_transfer_jobs +from .base import Stage, StageResult + + +@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]] + + +@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/transfer.py b/sentieon_cli/transfer.py index ed082fb..24f9c11 100644 --- a/sentieon_cli/transfer.py +++ b/sentieon_cli/transfer.py @@ -28,8 +28,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,7 +96,7 @@ def build_transfer_jobs( raw_vcf, regions_file=subset_bed, ), - "merge-trim-extra", + f"{name_base}-extra", 1, task_name="annotation-transfer", ) @@ -122,7 +131,7 @@ def build_transfer_jobs( ], view_xargs=["--no-version"], ), - f"merge-trim-{i}", + f"{name_base}-{i}", 1, task_name="annotation-transfer", ) @@ -136,7 +145,7 @@ def build_transfer_jobs( sharded_vcfs, xargs=["--no-version", "--threads", str(cores)], ), - "merge-trim-concat", + f"{name_base}-concat", cores, task_name="annotation-transfer", ) diff --git a/tests/integration/test_pipeline_execution.py b/tests/integration/test_pipeline_execution.py index e7a25c1..3f89de3 100644 --- a/tests/integration/test_pipeline_execution.py +++ b/tests/integration/test_pipeline_execution.py @@ -192,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 @@ -364,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/unit/test_dag_construction.py b/tests/unit/test_dag_construction.py index c0df2cf..1749f70 100644 --- a/tests/unit/test_dag_construction.py +++ b/tests/unit/test_dag_construction.py @@ -166,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) diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index 9239e63..f3f25e1 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -159,7 +159,7 @@ def test_dag_jobs_present(self): "dedup-lift", "metrics", "pangenome-sv", - "dnascope-raw", + "dnascope", "model-apply", ): assert name in job_names, f"missing job: {name}" @@ -366,7 +366,7 @@ def test_calling_inputs_and_replace_rg(self): 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-raw"): + 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 @@ -421,7 +421,7 @@ def test_dnascope_command(self): dag = pipeline.build_dag() _, all_jobs = self._get_all_job_names(dag) - cmd_str = str(self._get_job(all_jobs, "dnascope-raw").shell) + 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 @@ -435,7 +435,7 @@ def test_pcr_free(self): dag = pipeline.build_dag() _, all_jobs = self._get_all_job_names(dag) - cmd_str = str(self._get_job(all_jobs, "dnascope-raw").shell) + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) assert "--pcr_indel_model NONE" in cmd_str def test_bam_format(self): @@ -451,7 +451,7 @@ def test_bam_format(self): 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-raw").shell) + cmd_str = str(self._get_job(all_jobs, "dnascope").shell) assert bwa_bam in cmd_str assert lift_bam in cmd_str @@ -484,7 +484,7 @@ def test_skip_svs(self): job_names, _ = self._get_all_job_names(dag) assert "pangenome-sv" not in job_names - assert "dnascope-raw" in job_names + assert "dnascope" in job_names assert "graph-update" in job_names def test_skip_small_variants(self): @@ -494,7 +494,7 @@ def test_skip_small_variants(self): dag = pipeline.build_dag() job_names, _ = self._get_all_job_names(dag) - 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 assert "pangenome-sv" in job_names @@ -513,7 +513,7 @@ def test_multiple_lr_inputs(self): dag = pipeline.build_dag() _, all_jobs = self._get_all_job_names(dag) - cmd_str = str(self._get_job(all_jobs, "dnascope-raw").shell) + 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 @@ -526,7 +526,7 @@ def test_rgsm_overrides_sm(self): dag = pipeline.build_dag() _, all_jobs = self._get_all_job_names(dag) - cmd_str = str(self._get_job(all_jobs, "dnascope-raw").shell) + 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) @@ -548,7 +548,7 @@ def test_aligned_dag_jobs(self): "graph-update", "mm2-lift", "pangenome-sv", - "dnascope-raw", + "dnascope", "model-apply", ): assert name in job_names, f"missing job: {name}" @@ -619,7 +619,7 @@ def test_aligned_dag_dependencies(self): "extract-kmc" } assert "extract-kmc" in self._get_dep_names(dag, all_jobs, "mm2-lift") - assert self._get_dep_names(dag, all_jobs, "dnascope-raw") == { + assert self._get_dep_names(dag, all_jobs, "dnascope") == { "mm2-lift" } @@ -646,7 +646,7 @@ def test_aligned_replace_rg(self): 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-raw").shell) + 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)) @@ -669,7 +669,7 @@ def test_aligned_bam_format(self): 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-raw").shell) + assert lift_bam in str(self._get_job(all_jobs, "dnascope").shell) # Short-read input validation @@ -753,19 +753,19 @@ def test_lr_realign_downstream(self): "graph-update-raw", "graph-update", "pangenome-sv", - "dnascope-raw", + "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-raw"): + 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-raw").shell) + 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): @@ -803,13 +803,13 @@ def test_aligned_sr_with_lr_realign(self): 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-raw").shell) + 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-raw") == { + assert self._get_dep_names(dag, all_jobs, "dnascope") == { "mm2-lift", "bam-realign-0", } diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index f1d459e..5d17bb3 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -236,10 +236,10 @@ def test_gvcf_adds_gvcftyper_and_routes_outputs(self): assert apply_gvcf.endswith("sample-snv_apply.g.vcf.gz") # DNAscope emits gVCF; model-apply writes gVCF - for name in ("dnascope-raw", "model-apply"): + 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 apply_gvcf in cmd_str @@ -346,18 +346,18 @@ def test_call_svs(self): assert isinstance(dag, DAG) - # The dnascope-raw job should exist + # 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 @@ -382,7 +382,7 @@ def test_call_svs_disabled_by_default(self): ) 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 @@ -402,7 +402,7 @@ def test_skip_small_variants(self): 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 @@ -421,7 +421,7 @@ def test_skip_small_variants_with_call_svs(self): 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 @@ -430,7 +430,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 @@ -612,8 +612,8 @@ def test_call_svs_without_cnv_model(self): 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 diff --git a/tests/unit/test_stage_small_variants.py b/tests/unit/test_stage_small_variants.py new file mode 100644 index 0000000..2a561ec --- /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 both the transfer and the raw-VCF producer + assert dag.waiting_jobs[result.apply_job] == { + upstream, + 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_transfer.py b/tests/unit/test_stage_transfer.py new file mode 100644 index 0000000..ce3d807 --- /dev/null +++ b/tests/unit/test_stage_transfer.py @@ -0,0 +1,133 @@ +""" +Unit tests for the annotation-transfer stage +""" + +import pathlib + +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} From 11bf13c37b4864468ab2a6238d1bf5d795efe878 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 01:16:44 -0700 Subject: [PATCH 33/58] Insert the hybrid calling jobs directly into the DAG `DNAscopeHybridPipeline.call_variants` returned its 23 jobs as a tuple for `build_dag` to wire by hand. `add_variant_calling` now builds each job and inserts it with its dependencies as it goes, using `DNAscopeStage`, `ModelApplyStage`, `TransferStage` and the `driver_job` / `rm_job` helpers where they fit. Cleanup jobs are simply not created under `--retain_tmpdir`. The two DNAscope passes are now named `dnascope-1` and `dnascope-2` (were `calling-1` and `call2`). Commands and dependency edges are unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope_hybrid.py | 446 +++++++------------ tests/integration/test_pipeline_execution.py | 2 +- tests/unit/test_dnascope_hybrid_stage1.py | 22 + 3 files changed, 186 insertions(+), 284 deletions(-) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 4036279..4e35685 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,10 +18,8 @@ from . import command_strings as cmds from .dag import DAG from .driver import ( - BaseDriver, Driver, DNAscope, - DNAModelApply, HybridStage1, HybridStage2, HybridStage3, @@ -31,8 +29,11 @@ from .job import Job from .pipeline import BasePipeline from .shell_pipeline import Command, Pipeline +from .stages.base import StageContext, driver_job, rm_job from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage from .stages.ploidy import PloidyStage +from .stages.small_variants import DNAscopeStage, ModelApplyStage +from .stages.transfer import TransferConfig, TransferStage from .util import ( __version__, check_version, @@ -48,7 +49,6 @@ parse_fai, vcf_contigs, ) -from .transfer import build_transfer_jobs logger = get_logger(__name__) @@ -659,66 +659,14 @@ def build_dag(self) -> DAG: else None ) - ( - call_job, - select_job, - mapq0_job, - mapq0_slop_job, - cat_merge_job, - rm_job1, - stage1_fifo_job, - stage1_hap_job, - 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}) - # The haplotype job writes the fifo read by the first-stage job - dag.add_job(stage1_fifo_job) - dag.add_job(stage1_hap_job, {stage1_fifo_job, cat_merge_job}) - dag.add_job(stage1_job, {stage1_fifo_job, cat_merge_job}) - dag.add_job(second_stage_job, {stage1_job, stage1_hap_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, stage1_hap_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 @@ -757,39 +705,52 @@ def build_second_dag(self) -> Optional[DAG]: ).add_to(dag) return dag - def call_variants( + def transfer_config(self) -> TransferConfig: + """The inputs the annotation-transfer stage needs. + + Only valid once `--pop_vcf` is known to be set; the calling code + checks it first. + """ + assert self.pop_vcf is not None + return TransferConfig( + pop_vcf=self.pop_vcf, + shards=self.shards, + pop_vcf_contigs=self.pop_vcf_contigs, + fai_data=self.fai_data, + ) + + 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, # call_job - Job, # select_job - Job, # mapq0_job - Job, # mapq0_slop_job - Job, # cat_merge_job - Job, # rm_job1 - Job, # stage1_fifo_job - Job, # stage1_hap_job - Job, # stage1_job - Job, # rm_job2 - Job, # second_stage_job - Job, # rm_job3 - Job, # third_stage_job - Job, # rm_job4 - Job, # call2_job - Job, # subset_job - Job, # concat_job - Job, # rm_job5 - Job, # anno_job - Optional[List[Job]], # transfer_jobs - Optional[Job], # transfer_concat_job - Optional[Job], # apply_job - Optional[Job], # norm_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") @@ -808,35 +769,31 @@ def call_variants( if not check_version(cmd, min_version): sys.exit(2) + hybrid_model = self.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, - task_name="variant-calling", - ) + 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( @@ -855,27 +812,26 @@ def call_variants( 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, - read_filter=rg_info.lr_rg_read_filter + rg_info.sr_rg_read_filter, - ) - 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, + mapq0_job = driver_job( + ctx, + [ + HybridStage2( + model=self.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], ) + dag.add_job(mapq0_job, upstream) mapq0_slop_bed = self.tmp_dir.joinpath("hybrid_mapq0.ex1000.bed") mapq0_slop_job = Job( @@ -889,6 +845,7 @@ def call_variants( 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( @@ -901,12 +858,12 @@ def call_variants( 0, task_name="region-selection", ) - rm_cmd = ["rm", str(selected_bed), str(mapq0_slop_bed)] - rm_job1 = Job( - Pipeline(Command(*rm_cmd, fail_ok=True)), + dag.add_job(cat_merge_job, {mapq0_slop_job, select_job}) + self.add_cleanup( + dag, + [selected_bed, mapq0_slop_bed], "rm-tmp1", - 0, - task_name="cleanup", + {cat_merge_job}, ) stage1_ins_fa = self.tmp_dir.joinpath("stage1_ins.fa") @@ -940,6 +897,8 @@ def call_variants( 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, @@ -973,6 +932,7 @@ def call_variants( 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( @@ -989,56 +949,45 @@ def call_variants( self.cores, task_name="hybrid-realignment", ) - 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)), + 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", - 0, - task_name="cleanup", + {stage1_job, stage1_hap_job}, ) 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 = Job( - Pipeline(Command(*driver.build_cmd())), - "second-stage", - self.cores, + second_stage_job = driver_job( + ctx, + [ + 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, + ) + ], + name="second-stage", task_name="hybrid-realignment", + inputs=[stage1_bam, stage1_hap_bam], ) - - rm_cmd = ["rm", str(stage1_bam), str(stage1_hap_bam)] - rm_job3 = Job( - Pipeline(Command(*rm_cmd, fail_ok=True)), + dag.add_job(second_stage_job, {stage1_job, stage1_hap_job}) + self.add_cleanup( + dag, + [stage1_bam, stage1_hap_bam], "rm-tmp3", - 0, - task_name="cleanup", + {second_stage_job}, ) suffix = "bam" if self.bam_format else "cram" stage3_aln = pathlib.Path( str(self.output_vcf).replace(".vcf.gz", f"_sr_realigned.{suffix}") ) - driver = Driver( + stage3_driver = Driver( reference=self.reference, thread_count=self.cores, replace_rg=rg_info.replace_rg_args[0] + rg_info.replace_rg_args[1], @@ -1046,7 +995,7 @@ def call_variants( 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"), @@ -1056,46 +1005,40 @@ def call_variants( cmds.hybrid_stage3( stage3_aln, reference=self.reference, - driver=driver, + driver=stage3_driver, cores=self.cores, ), "third-stage", self.cores, task_name="hybrid-realignment", ) - rm_cmd = ["rm", str(stage2_unmap_bam), str(stage2_alt_bam)] - rm_job4 = Job( - Pipeline(Command(*rm_cmd, fail_ok=True)), + dag.add_job(third_stage_job, {second_stage_job}) + self.add_cleanup( + dag, + [stage2_unmap_bam, stage2_alt_bam], "rm-tmp4", - 0, - task_name="cleanup", + {third_stage_job}, ) # 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, - task_name="variant-calling", - ) + 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") @@ -1110,6 +1053,7 @@ def call_variants( 0, task_name="vcf-merge", ) + dag.add_job(subset_job, {second_stage_job}) concat_job = Job( cmds.bcftools_concat( combined_tmp_vcf, @@ -1119,12 +1063,12 @@ def call_variants( 0, task_name="vcf-merge", ) - rm_cmd = ["rm", str(combined_vcf), str(subset_vcf), str(pass2_vcf)] - rm_job5 = Job( - Pipeline(Command(*rm_cmd, fail_ok=True)), + dag.add_job(concat_job, {subset_job} | call2.terminal) + self.add_cleanup( + dag, + [combined_vcf, subset_vcf, pass2_vcf], "rm-tmp5", - 0, - task_name="cleanup", + {concat_job}, ) # Annotate the output VCF @@ -1147,11 +1091,10 @@ def call_variants( 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 - + apply_dependencies: Set[Job] = {anno_job} if self.pop_vcf: transfer_target = self.tmp_dir.joinpath( "combined_tmp_transfer.vcf.gz" @@ -1159,65 +1102,26 @@ def call_variants( 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 = TransferStage( + ctx=ctx, + config=self.transfer_config(), + raw_vcf=anno_target, + out_vcf=transfer_target, + ).add_to(dag, {anno_job}) input_to_apply = transfer_target + apply_dependencies = transfer.terminal if self.skip_model_apply: - return ( - call_job, - select_job, - mapq0_job, - mapq0_slop_job, - cat_merge_job, - rm_job1, - stage1_fifo_job, - stage1_hap_job, - 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, - ) + return # Model Apply apply_vcf = self.tmp_dir.joinpath("combined_apply.vcf.gz") - driver = Driver( - reference=self.reference, - thread_count=self.cores, - ) - 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, - task_name="model-apply", - ) + apply = ModelApplyStage( + ctx=ctx, + model=hybrid_model, + vcf=input_to_apply, + output=apply_vcf, + ).add_to(dag, apply_dependencies) # Final normalize norm_job = Job( @@ -1231,28 +1135,4 @@ def call_variants( 0, task_name="vcf-norm", ) - return ( - call_job, - select_job, - mapq0_job, - mapq0_slop_job, - cat_merge_job, - rm_job1, - stage1_fifo_job, - stage1_hap_job, - 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, apply.terminal) diff --git a/tests/integration/test_pipeline_execution.py b/tests/integration/test_pipeline_execution.py index 3f89de3..9ed9e6a 100644 --- a/tests/integration/test_pipeline_execution.py +++ b/tests/integration/test_pipeline_execution.py @@ -152,7 +152,7 @@ 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 diff --git a/tests/unit/test_dnascope_hybrid_stage1.py b/tests/unit/test_dnascope_hybrid_stage1.py index 6b72aeb..231fe8c 100644 --- a/tests/unit/test_dnascope_hybrid_stage1.py +++ b/tests/unit/test_dnascope_hybrid_stage1.py @@ -170,6 +170,28 @@ def test_stage1_dependencies(self): "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() From a82a97321718efd1c0f1d1a4d07f699506edb0d0 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 01:27:26 -0700 Subject: [PATCH 34/58] Insert the long-read calling jobs directly into the DAG `DNAscopeLRPipeline.lr_call_variants` returned a 27-field dataclass of jobs that `build_dag` wired by hand. `add_small_variant_calling` now composes the small-variant stages and inserts each job with its dependencies as it goes: `DNAscopeStage` for the five calling passes, `TransferApplyStage` for the diploid and unphased tails, and `TransferStage` plus `ModelApplyStage` for the two haplotypes, whose model-apply jobs wait on both haplotype transfers. The passes are now named after the stage: `dnascope-diploid`, `dnascope-hap1`/`hap2`, `dnascope-unphased`, `dnascope-haploid`, with matching `model-apply-*` jobs, and the four annotation transfers carry `diploid`/`hap1`/`hap2`/`unphased` tags in their job names. Commands and dependency edges are unchanged; a new unit test pins the non-obvious edges. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope_longread.py | 791 +++++++------------ tests/unit/test_dnascope_longread_dag.py | 218 +++++ tests/unit/test_dnascope_longread_pop_vcf.py | 10 +- 3 files changed, 522 insertions(+), 497 deletions(-) create mode 100644 tests/unit/test_dnascope_longread_dag.py diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index ca4aeb0..c0ddf13 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -4,7 +4,6 @@ import argparse import copy -from dataclasses import dataclass import json import pathlib import shutil @@ -19,10 +18,10 @@ from .archive import ar_load from .dag import DAG from .driver import ( + BaseAlgo, Driver, DNAscope, DNAscopeHP, - DNAModelApply, LongReadSV, ReadWriter, RepeatModel, @@ -36,7 +35,15 @@ vcf_contigs, ) from .shell_pipeline import Command, Pipeline -from .transfer import build_transfer_jobs +from .stages.base import StageContext, driver_job +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + ModelApplyStage, + TransferApplyStage, + TransferSpec, +) +from .stages.transfer import TransferConfig, TransferStage from .util import ( __version__, check_version, @@ -82,39 +89,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""" @@ -468,6 +442,7 @@ 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 realign_jobs: Set[Job] = set() @@ -500,122 +475,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}, + self.add_small_variant_calling( + dag, ctx, sample_input, realign_jobs.union(align_jobs) ) - 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, - }, - ) - - 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)) @@ -888,12 +751,34 @@ def hificnv( ) return hificnv_job - def lr_call_variants( + def transfer_config(self) -> TransferConfig: + """The inputs the annotation-transfer stage needs. + + Only valid once `--pop_vcf` is known to be set; the calling code + checks it first. + """ + assert self.pop_vcf is not None + return TransferConfig( + pop_vcf=self.pop_vcf, + shards=self.shards, + pop_vcf_contigs=self.pop_vcf_contigs, + fai_data=self.fai_data, + ) + + 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") @@ -909,26 +794,13 @@ def lr_call_variants( 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, - ) - # 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 = self.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, @@ -936,54 +808,40 @@ def lr_call_variants( model=self.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, - task_name="variant-calling", - ) - - # 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=self.transfer_config(), + 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, - task_name="model-apply", - ) + 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") @@ -996,28 +854,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, - 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, + 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, ) + 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( @@ -1040,8 +894,10 @@ def lr_call_variants( 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: @@ -1055,6 +911,8 @@ def lr_call_variants( 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), @@ -1062,33 +920,7 @@ def lr_call_variants( 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, - task_name="repeat-model", - ) + dag.add_job(bcftools_subtract_job, subtract_deps) bcftools_subset_unphased_job = Job( Pipeline( @@ -1107,50 +939,70 @@ def lr_call_variants( 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 = self.model_bundle.joinpath("haploid_model") + haploid_hp_model = self.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, - task_name="variant-calling", - ) - ) + 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( @@ -1177,128 +1029,108 @@ def lr_call_variants( 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 = [] + # Transfer annotations to the patched VCFs. Both model-apply jobs + # wait on both transfers, as they did when the DAG was wired by hand. + second_ma_deps: Set[Job] = {haploid_patch_job} 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 + for i, (patch_unanno_vcf, patch_vcf) in enumerate( + zip(patch_unanno_vcfs, patch_vcfs), start=1 ): - 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) + patch_transfer = TransferStage( + ctx=ctx, + config=self.transfer_config(), + raw_vcf=patch_unanno_vcf, + out_vcf=patch_vcf, + tag=f"hap{i}", + ).add_to(dag, {haploid_patch_job}) + second_ma_deps.update(patch_transfer.terminal) # apply trained model to the patched vcfs. 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, - task_name="model-apply", - ) - ) + merge_deps: Set[Job] = set() + for i, (patch_vcf, hap_vcf) in enumerate( + zip(patch_vcfs, hap_vcfs), start=1 + ): + hap_apply = ModelApplyStage( + ctx=ctx, + model=haploid_model, + vcf=patch_vcf, + output=hap_vcf, + name=f"model-apply-hap{i}", + ).add_to(dag, second_ma_deps) + 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=self.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, - task_name="variant-calling", - ) + 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", ) - diploid_patch_job = Job( - cmd, "diploid-patch", self.cores, task_name="variant-patch" + dag.add_job( + diploid_patch_job, + {bcftools_subset_unphased_job} | unphased_call.terminal, ) - # 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=self.transfer_config(), + 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, - task_name="model-apply", - ) + unphased_apply = TransferApplyStage( + ctx=ctx, + raw_vcf=diploid_unphased_patch, + transfer=unphased_transfer, + apply=ApplySpec( + model=self.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") @@ -1320,6 +1152,7 @@ def lr_call_variants( self.cores, task_name="vcf-merge", ) + dag.add_job(merge_job, merge_deps) gvcf_combine_job = None if self.gvcf: @@ -1335,137 +1168,111 @@ def lr_call_variants( 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, - dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("haploid_model"), - ) - ) - driver.add_algo( - DNAscopeHP( - haploid_hp_fn, + haploid_gvcf_fn, dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("haploid_hp_model"), - pcr_indel_model=self.repeat_model, + emit_mode="gvcf", + ploidy=1, + model=self.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, - task_name="variant-calling", - ) + 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( + self.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(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, - task_name="variant-patch", - ) - haploid_concat_job = Job( - cmds.bcftools_concat( - self.output_vcf, - [diploid_merged_vcf, haploid_out_fn], - ), - "haploid-diploid-concat", - 0, - task_name="vcf-merge", - ) + 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, - task_name="gvcf", - ) - haploid_gvcf_concat_job = Job( - cmds.bcftools_concat( - output_gvcf, - [diploid_gvcf, haploid_gvcf], - ), - "haploid-gvcf-concat", - 0, - task_name="gvcf", - ) - 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( diff --git a/tests/unit/test_dnascope_longread_dag.py b/tests/unit/test_dnascope_longread_dag.py new file mode 100644 index 0000000..c48bcca --- /dev/null +++ b/tests/unit/test_dnascope_longread_dag.py @@ -0,0 +1,218 @@ +""" +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_both_haplotype_applies_wait_on_both_transfers(self, tmp_path): + pop_vcf = tmp_path / "population.vcf.gz" + pop_vcf.touch() + dag = build_dag(tmp_path, pop_vcf=pop_vcf) + + # Each model-apply reads only its own haplotype's transfer, but + # both wait on both concat jobs, as the hand-wired DAG did + for hap in ("model-apply-hap1", "model-apply-hap2"): + assert deps_of(dag, hap) == [ + "merge-trim-hap1-concat", + "merge-trim-hap2-concat", + "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", + } + assert deps_of(dag, "model-apply-diploid") == [ + "dnascope-diploid", + "merge-trim-diploid-concat", + ] + assert deps_of(dag, "model-apply-unphased") == [ + "diploid-patch", + "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..0b65130 100644 --- a/tests/unit/test_dnascope_longread_pop_vcf.py +++ b/tests/unit/test_dnascope_longread_pop_vcf.py @@ -170,7 +170,7 @@ 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""" @@ -188,16 +188,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 From 6f0c225baaacf2437a6c4047ff78fc98fe56615b Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 01:42:14 -0700 Subject: [PATCH 35/58] Move alignment and read extraction into stages `stages/alignment.py` holds the bwa and minimap2 alignment stages (from BAM/CRAM input and from fastq), the pangenome bwa-plus-extract stage, and a ReadWriter stage, together with the shared `aln_suffix` and `find_unzip` helpers. Alignment jobs are always DAG roots, so these stages expose `build()` for callers that insert the jobs themselves as well as the usual `add_to()`. The hybrid-pangenome pipeline's copy of the minimap2 realignment and both pangenomes' bwa-extract builders are gone; the dnascope pipeline builds its alignment cleanup jobs from the stage's `cleanup_paths`. The long-read pipeline keeps its `lr_align_inputs`, `lr_align_fastq` and `merge_input_files` methods as thin wrappers. Commands, job names and dependency edges are unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 212 ++++--------- sentieon_cli/dnascope_hybrid.py | 16 +- sentieon_cli/dnascope_longread.py | 133 +++------ sentieon_cli/hybrid_pangenome.py | 121 +++----- sentieon_cli/sentieon_pangenome.py | 84 +++--- sentieon_cli/stages/alignment.py | 443 +++++++++++++++++++++++++++ tests/unit/test_stage_alignment.py | 464 +++++++++++++++++++++++++++++ 7 files changed, 1091 insertions(+), 382 deletions(-) create mode 100644 sentieon_cli/stages/alignment.py create mode 100644 tests/unit/test_stage_alignment.py diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 2467fbe..83397f7 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -10,11 +10,10 @@ import pathlib import shutil import sys -from typing import List, Optional, Set, Tuple +from typing import List, Optional, Set import packaging.version -from . import command_strings as cmds from .dag import DAG from .driver import ( AlignmentStat, @@ -32,7 +31,12 @@ ) from .job import Job from .pipeline import BasePipeline -from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + AlignResult, + BwaFastqStage, + BwaRealignStage, + find_unzip, +) from .stages.base import StageContext, driver_job, rm_job from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage @@ -410,13 +414,18 @@ def build_dag(self) -> DAG: sample_input = copy.deepcopy(self.sample_input) bam_rm_job = None 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_rm_job = rm_job(aln_result.cleanup_paths, "rm-bam-aln") + fq_result = self.add_sr_fastq_alignment(dag, ctx) + align_fastq_jobs = set(fq_result.jobs) + fq_rm_job = ( + rm_job(fq_result.cleanup_paths, "rm-fq-aln") + if fq_result.cleanup_paths + else None + ) + sample_input += fq_result.outputs # Dedup and metrics preprocessing = self.add_sr_preprocessing( @@ -449,11 +458,10 @@ def build_dag(self) -> DAG: 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) @@ -463,161 +471,51 @@ def sr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job], Job]: 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, - task_name="alignment", - ) - 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, - task_name="cleanup", - ) - - 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=self.model_bundle, + bam_format=self.bam_format, + duplicate_marking=self.sr_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) + return AlignResult( + jobs=[], terminal=set(), outputs=[], cleanup_paths=[] + ) 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) - 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 - ) - ): - 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}, - task_name="alignment", - ) - 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, - task_name="cleanup", - ) - - return (res, jobs, rm_job) + return BwaFastqStage( + ctx=ctx, + r1_fastq=self.sr_r1_fastq, + r2_fastq=self.sr_r2_fastq, + readgroups=self.sr_readgroups, + model_bundle=self.model_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 sr_metrics_algos(self, paths: MetricsPaths) -> List[BaseAlgo]: """The metrics that do not need duplicate-marked reads""" diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 4e35685..503f322 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -594,14 +594,14 @@ 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) + fq_rm_job = ( + rm_job(fq_result.cleanup_paths, "rm-fq-aln") + if fq_result.cleanup_paths + else None + ) + sr_aln.extend(fq_result.outputs) # Short-read dedup preprocessing = self.add_sr_preprocessing( diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index c0ddf13..e3f8da5 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -6,7 +6,6 @@ import copy import json import pathlib -import shutil import sys from typing import Dict, List, Optional, Set, Tuple @@ -23,7 +22,6 @@ DNAscope, DNAscopeHP, LongReadSV, - ReadWriter, RepeatModel, VariantPhaser, ) @@ -35,6 +33,12 @@ vcf_contigs, ) from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + Minimap2FastqStage, + Minimap2RealignStage, + ReadWriterStage, + find_unzip, +) from .stages.base import StageContext, driver_job from .stages.small_variants import ( ApplySpec, @@ -492,9 +496,6 @@ def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: 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) @@ -503,99 +504,45 @@ def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: 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, - ) - - 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, - task_name="alignment", - ) - ) - res.append(out_aln) - return (res, realign_jobs) + result = Minimap2RealignStage( + ctx=self.stage_context(), + inputs=self.lr_aln, + model_bundle=self.model_bundle, + sample_name=self.output_vcf.name.replace(".vcf.gz", ""), + bam_format=self.bam_format, + input_ref=self.lr_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] = [] 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) - 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, - ), - f"align-{i}", - self.cores, - task_name="alignment", - ) - ) - res.append(out_aln) - return (res, align_jobs) + result = Minimap2FastqStage( + ctx=self.stage_context(), + fastq=self.fastq, + readgroups=self.readgroups, + model_bundle=self.model_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""" @@ -647,19 +594,15 @@ def merge_input_files( # 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, + result = ReadWriterStage( + ctx=self.stage_context(), + inputs=sample_input, + output=merged_bam, + name="merge-bam", task_name="alignment", - ) - return (merged_bam, merge_job) + threads=0, # Run in the background + ).build() + return (merged_bam, result.jobs[0]) def pbsv( self, diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index cc0e2f1..a021dbf 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -8,8 +8,8 @@ import argparse import copy import json +import logging import pathlib -import shutil import sys from typing import Dict, List, Optional, Set, Tuple @@ -35,6 +35,12 @@ vcf_contigs, ) from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + BwaExtractStage, + Minimap2RealignStage, + find_unzip, +) +from .stages.base import StageContext from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage from .stages.small_variants import DNAscopeStage, ModelApplyStage @@ -591,17 +597,6 @@ def configure(self) -> None: """Configure pipeline parameters""" pass - def find_unzip(self) -> str: - """The decompression tool for fastq input""" - 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" - return unzip - def build_dag(self) -> DAG: """Build the DAG for the hybrid-pangenome pipeline""" if not self.reference: @@ -657,9 +652,9 @@ def build_dag(self) -> DAG: calling_lr = list(self.lr_aln) realign_jobs: Set[Job] = set() if self.lr_align_input: - calling_lr, realign_jobs = self.lr_align_inputs() - for job in realign_jobs: - dag.add_job(job) + 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() @@ -670,8 +665,10 @@ def build_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) # Do not run vg-haplotypes with bwa in low-mem environments total_mem_gb = total_memory() / (1024.0**3) @@ -900,7 +897,7 @@ def build_hybrid_kmc_job(self, kmer_prefix: pathlib.Path) -> Job: self.tmp_dir, memory=self.kmer_memory, threads=self.cores, - unzip=self.find_unzip(), + unzip=find_unzip(self.logger, logging.INFO), ), "kmc", 0, # run in the background @@ -947,78 +944,50 @@ def build_extract_kmc_jobs( ) return symlink_job, extract_kmc_job - def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: - """Align the long-read input files to the linear reference genome - with minimap2""" - assert self.output_vcf is not None - assert self.reference is not None + 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 - res: List[pathlib.Path] = [] - realign_jobs: Set[Job] = set() - suffix = "bam" if self.bam_format else "cram" - 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) - realign_jobs.add( - Job( - cmds.cmd_samtools_fastq_minimap2( - out_aln, - input_aln, - self.reference, - self.model_bundle, - self.cores, - rg_lines, - self.sample_sm, - self.lr_input_ref, - minimap2_model=self.model_bundle.joinpath( - "minimap2_lr.model" - ), - ), - f"bam-realign-{i}", - self.cores, - task_name="alignment", - ) - ) - res.append(out_aln) - return (res, realign_jobs) + 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 build_alignment_job( + def bwa_extract_stage( self, + ctx: StageContext, sample_bam: pathlib.Path, sample_fastq: pathlib.Path, - ) -> Job: - """Build the bwa alignment and read extraction job""" - assert self.reference is not None - assert self.model_bundle is not None + ) -> 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" - 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=self.find_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, - task_name="alignment", + extract_model=self.model_bundle.joinpath(self.extract_model_name), + bwa_model=self.model_bundle.joinpath("bwa.model"), + unzip=find_unzip(self.logger, logging.INFO), ) - return bwa_job def build_haplotypes_job( self, output_gbz: pathlib.Path, kmer_file: pathlib.Path diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index ccd4877..35c81ef 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -5,6 +5,7 @@ import argparse import copy import json +import logging import pathlib import shutil import sys @@ -22,13 +23,17 @@ from .driver import ( BaseAlgo, DNAscope, - Driver, PangenomeSV, - ReadWriter, ) from .job import Job from .logging import get_logger from .shell_pipeline import Command, Pipeline +from .stages.alignment import ( + BwaExtractStage, + ReadWriterStage, + find_unzip, +) +from .stages.base import StageContext from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage @@ -737,8 +742,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 @@ -1015,47 +1022,36 @@ def transfer_config(self) -> TransferConfig: fai_data=self.fai_data, ) - 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) + ) -> BwaExtractStage: + """The bwa alignment and read extraction stage. + + The bwa alignment reuses the input readgroup ID with a `-bwa` + suffix, so it stays distinct from the lifted alignment's. + """ if not self.model_bundle: self.logger.error("model_bundle is required") sys.exit(2) - 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" - 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, - task_name="alignment", + extract_model=self.model_bundle.joinpath(self.extract_model_name), + bwa_model=self.model_bundle.joinpath("bwa.model"), + unzip=find_unzip(self.logger, logging.INFO), ) - return bwa_job def build_haplotypes_job( self, output_gbz: pathlib.Path, kmer_file: pathlib.Path @@ -1289,19 +1285,15 @@ def build_t1k_jobs( # 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, + extract_result = ReadWriterStage( + ctx=self.stage_context(), + inputs=sr_alignments, + output=extracted_bam, + name=f"t1k-{tag}-extract", task_name="t1k", - ) + interval=locus, + ).build() + extract_job = extract_result.jobs[0] t1k_job = Job( cmds.cmd_t1k( diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py new file mode 100644 index 0000000..f474b42 --- /dev/null +++ b/sentieon_cli/stages/alignment.py @@ -0,0 +1,443 @@ +""" +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 + +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 + +# The `sentieon util sort` arguments every aligner defaults to +DEFAULT_UTIL_SORT_ARGS = "--cram_write_options version=3.0,compressor=rans" + + +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, level: int = logging.WARNING) -> str: + """The decompression tool to read gzipped fastq with. + + ``igzip`` is considerably faster but is not always installed, so fall + back to ``gzip``. Callers differ on how loudly that is worth + reporting, hence ``level``. + """ + unzip = "igzip" + if not shutil.which(unzip): + logger.log( + level, + "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 = DEFAULT_UTIL_SORT_ARGS + + 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) + cleanup_paths.append(out_aln) + cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) + if suffix == "cram": + cleanup_paths.append(pathlib.Path(str(out_aln) + ".crai")) + + 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 = DEFAULT_UTIL_SORT_ARGS + + 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) + cleanup_paths.append(out_aln) + cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) + if suffix == "cram": + cleanup_paths.append(pathlib.Path(str(out_aln) + ".crai")) + + 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 = DEFAULT_UTIL_SORT_ARGS + 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 = DEFAULT_UTIL_SORT_ARGS + + 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/tests/unit/test_stage_alignment.py b/tests/unit/test_stage_alignment.py new file mode 100644 index 0000000..3896437 --- /dev/null +++ b/tests/unit/test_stage_alignment.py @@ -0,0 +1,464 @@ +""" +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 + + def test_find_unzip_takes_the_log_level(self, caplog): + logger = logging.getLogger("test-unzip-level") + with patch( + "sentieon_cli.stages.alignment.shutil.which", return_value=None + ): + with caplog.at_level(logging.INFO, logger=logger.name): + assert find_unzip(logger, logging.INFO) == "gzip" + assert caplog.records[0].levelno == logging.INFO + + +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", + ] + assert "--bam_compression 1" not in str(result.jobs[0].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_bam_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_cram_cleanup_paths_add_the_crai(self, tmp_path, rg_lines): + result = self.make(tmp_path, duplicate_marking="none").build() + + out = str(tmp_path / "output_bwa_sorted_0.cram") + assert result.cleanup_paths[:3] == [ + pathlib.Path(out), + pathlib.Path(out + ".bai"), + pathlib.Path(out + ".crai"), + ] + + 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_outputs_and_cleanup_paths(self, tmp_path): + result = self.make(tmp_path, duplicate_marking="none").build() + + out = tmp_path / "output_bwa_sorted_fq_0_0.cram" + assert result.outputs == [out] + assert result.cleanup_paths == [ + out, + pathlib.Path(str(out) + ".bai"), + pathlib.Path(str(out) + ".crai"), + ] + + +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 == [] + + 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 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} From 4c5408858bc0318c558c1cab16f3725095d4dc2c Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 01:59:41 -0700 Subject: [PATCH 36/58] Use the shared version-check and required-argument helpers Replace the repeated version-check loops with `require_versions` / `versions_available`, and the repeated "argument is required" exits with `BasePipeline.required` or the values already carried by the stage context. Drop `BasePipeline.cnv_sex_args`, which lost its last caller when CNV calling moved into a stage. The long-read and hybrid tests no longer patch `check_version` in the pipeline modules: they skip version checks, so the patches were inert, and the modules no longer import the function. The pangenome KMC test, which does exercise the check, patches it in `util` instead. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 63 +++---- sentieon_cli/dnascope_hybrid.py | 82 +++------ sentieon_cli/dnascope_longread.py | 133 +++++--------- sentieon_cli/hybrid_pangenome.py | 55 +++--- sentieon_cli/pipeline.py | 29 ++- sentieon_cli/sentieon_pangenome.py | 167 ++++++------------ tests/unit/test_cnv_sex.py | 7 - tests/unit/test_dnascope_hybrid_cnv.py | 10 +- tests/unit/test_dnascope_hybrid_pop_vcf.py | 6 +- tests/unit/test_dnascope_hybrid_stage1.py | 5 +- tests/unit/test_dnascope_longread_pop_vcf.py | 6 +- ...test_sentieon_pangenome_kmc_integration.py | 4 +- 12 files changed, 190 insertions(+), 377 deletions(-) diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 83397f7..7ed4ba8 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -47,10 +47,10 @@ TransferApplyStage, ) from .util import ( - check_version, library_preloaded, parse_rg_line, path_arg, + require_versions, split_alignment, total_memory, ) @@ -272,9 +272,7 @@ def __init__(self) -> None: self.no_split_alignment = False def validate(self) -> None: - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) + self.required(self.output_vcf, "output_vcf") # uniquify pipeline attributes self.sr_r1_fastq = self.r1_fastq @@ -462,19 +460,14 @@ def add_sr_input_alignment( self, dag: DAG, ctx: StageContext ) -> AlignResult: """Align input BAM/CRAM/uBAM/uCRAM files with bwa""" - 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) + require_versions(ALN_MIN_VERSIONS, skip=self.skip_version_check) return BwaRealignStage( ctx=ctx, inputs=self.sample_input, - model_bundle=self.model_bundle, + model_bundle=bundle, bam_format=self.bam_format, duplicate_marking=self.sr_duplicate_marking, input_ref=self.input_ref, @@ -488,26 +481,21 @@ def add_sr_fastq_alignment( self, dag: DAG, ctx: StageContext ) -> AlignResult: """Align fastq files to the reference genome using bwa""" - 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.sr_r1_fastq and not self.sr_readgroups: return AlignResult( jobs=[], terminal=set(), outputs=[], cleanup_paths=[] ) - 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(FQ_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=self.model_bundle, + model_bundle=bundle, numa_nodes=self.numa_nodes, bam_format=self.bam_format, duplicate_marking=self.sr_duplicate_marking, @@ -537,12 +525,9 @@ def add_sr_preprocessing( upstream: Set[Job], ) -> SrPreprocessing: """Mark duplicates and collect metrics from the short reads""" - if not self.output_vcf: - self.logger.error("output_vcf is required") - sys.exit(2) suffix = "bam" if self.bam_format else "cram" - paths = MetricsPaths.from_output_vcf(self.output_vcf) + paths = MetricsPaths.from_output_vcf(ctx.output_vcf) paths.ensure_dir(self.dry_run) if self.sr_duplicate_marking == "none": @@ -578,7 +563,7 @@ def add_sr_preprocessing( lc_extra_algos.extend(self.sr_metrics_algos(paths)) deduped = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_deduped.{suffix}") + str(ctx.output_vcf).replace(".vcf.gz", f"_deduped.{suffix}") ) dedup_result = DedupStage( ctx=ctx, @@ -636,41 +621,33 @@ def add_small_variant_calling( 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") ) 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") ) algos: List[BaseAlgo] = [ @@ -716,7 +693,7 @@ def add_small_variant_calling( GVCFtyperStage( ctx=ctx, gvcf=out_gvcf, - output=self.output_vcf, + output=ctx.output_vcf, interval=self.bed, ).add_to(dag, transfer_apply.terminal) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 503f322..e31fd98 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -36,10 +36,10 @@ from .stages.transfer import TransferConfig, TransferStage from .util import ( __version__, - check_version, library_preloaded, parse_rg_line, path_arg, + require_versions, sample_sex_arg, split_alignment, vcf_id, @@ -396,10 +396,7 @@ def validate_cnv(self) -> None: # CNV calling now runs in the second DAG, so the version of the # driver is checked before the first DAG starts - if not self.skip_version_check: - for cmd, min_version in CNV_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(CNV_MIN_VERSIONS, skip=self.skip_version_check) def validate_bundle(self) -> None: bundle_info_bytes = ar_load( @@ -571,15 +568,7 @@ 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( @@ -674,15 +663,7 @@ def build_second_dag(self) -> Optional[DAG]: """Build the second DAG, for sex-aware CNV calling""" if self.skip_cnv: return None - 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") assert self.ploidy_json is not None ctx = self.stage_context() @@ -693,7 +674,7 @@ def build_second_dag(self) -> Optional[DAG]: CNVscopeStage( ctx=ctx, inputs=self.cnv_sr_aln, - model=self.model_bundle.joinpath("cnv.model"), + 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") @@ -752,24 +733,13 @@ def add_variant_calling( 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 = self.model_bundle.joinpath("hybrid.model") + hybrid_model = bundle.joinpath("hybrid.model") # First pass - combined variant calling vcf_suffix = ".g.vcf.gz" if self.gvcf else ".vcf.gz" @@ -819,9 +789,7 @@ def add_variant_calling( ctx, [ HybridStage2( - model=self.model_bundle.joinpath( - "HybridStage2_region.model" - ), + model=bundle.joinpath("HybridStage2_region.model"), all_bed=mapq0_bed, ) ], @@ -869,7 +837,7 @@ def add_variant_calling( 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, @@ -878,7 +846,7 @@ def add_variant_calling( 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, ) @@ -901,7 +869,7 @@ def add_variant_calling( 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, @@ -911,7 +879,7 @@ def add_variant_calling( stage1_driver.add_algo( HybridStage1( stage1_fifo, - model=self.model_bundle.joinpath("HybridStage1.model"), + model=bundle.joinpath("HybridStage1.model"), hap_bam="-", hap_bed=stage1_hap_bed, hap_vcf=stage1_hap_vcf, @@ -938,12 +906,12 @@ def add_variant_calling( 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, hap_fastq_fifo=stage1_fifo, - bwa_model=self.model_bundle.joinpath("HybridStage1_bwa.model"), + bwa_model=bundle.joinpath("HybridStage1_bwa.model"), ), "first-stage", self.cores, @@ -964,7 +932,7 @@ def add_variant_calling( ctx, [ HybridStage2( - model=self.model_bundle.joinpath("HybridStage2.model"), + model=bundle.joinpath("HybridStage2.model"), hap_bed=stage1_hap_bed, unmap_bam=stage2_unmap_bam, alt_bam=stage2_alt_bam, @@ -985,10 +953,10 @@ def add_variant_calling( 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}") ) stage3_driver = Driver( - reference=self.reference, + 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], @@ -998,13 +966,13 @@ def add_variant_calling( 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, + reference=ctx.reference, driver=stage3_driver, cores=self.cores, ), @@ -1077,7 +1045,7 @@ def add_variant_calling( ) 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( @@ -1100,7 +1068,7 @@ def add_variant_calling( "combined_tmp_transfer.vcf.gz" ) if self.skip_model_apply: - transfer_target = self.output_vcf + transfer_target = ctx.output_vcf transfer = TransferStage( ctx=ctx, @@ -1126,9 +1094,9 @@ def add_variant_calling( # 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", diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index e3f8da5..2ebdb32 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -50,11 +50,12 @@ from .stages.transfer import TransferConfig, TransferStage from .util import ( __version__, - check_version, + library_preloaded, parse_rg_line, path_arg, - library_preloaded, + require_versions, vcf_id, + versions_available, ) TOOL_MIN_VERSIONS = { @@ -493,22 +494,15 @@ 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.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) + output_vcf = self.required(self.output_vcf, "output_vcf") + bundle = self.required(self.model_bundle, "model_bundle") + require_versions(ALN_MIN_VERSIONS, skip=self.skip_version_check) result = Minimap2RealignStage( ctx=self.stage_context(), inputs=self.lr_aln, - model_bundle=self.model_bundle, - sample_name=self.output_vcf.name.replace(".vcf.gz", ""), + model_bundle=bundle, + sample_name=output_vcf.name.replace(".vcf.gz", ""), bam_format=self.bam_format, input_ref=self.lr_input_ref, fastq_taglist=self.fastq_taglist, @@ -521,22 +515,17 @@ def lr_align_fastq(self) -> Tuple[List[pathlib.Path], Set[Job]]: """ Align fastq to the reference genome using minimap2 """ - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") if self.fastq is None and self.readgroups is None: 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(FQ_MIN_VERSIONS, skip=self.skip_version_check) result = Minimap2FastqStage( ctx=self.stage_context(), fastq=self.fastq, readgroups=self.readgroups, - model_bundle=self.model_bundle, + model_bundle=bundle, bam_format=self.bam_format, unzip=find_unzip(self.logger), minimap2_args=self.minimap2_args, @@ -547,19 +536,14 @@ def lr_align_fastq(self) -> Tuple[List[pathlib.Path], Set[Job]]: 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() + 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 set() mosdepth_jobs = set() for i, input_file in enumerate(sample_input): @@ -587,10 +571,7 @@ 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") @@ -611,10 +592,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( @@ -662,15 +640,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 = [ @@ -723,24 +700,13 @@ def add_small_variant_calling( 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) + 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") - diploid_model = self.model_bundle.joinpath("diploid_model") + diploid_model = bundle.joinpath("diploid_model") diploid_algos: List[BaseAlgo] = [] if self.gvcf: diploid_algos.append( @@ -748,7 +714,7 @@ def add_small_variant_calling( diploid_gvcf_fn, dbsnp=self.dbsnp, emit_mode="gvcf", - model=self.model_bundle.joinpath("gvcf_model"), + model=bundle.joinpath("gvcf_model"), ) ) diploid_algos.append( @@ -847,7 +813,7 @@ def add_small_variant_calling( 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", @@ -909,8 +875,8 @@ def add_small_variant_calling( repeat_model_deps.add(repeatmodel_job) # Second pass - phased variants - haploid_model = self.model_bundle.joinpath("haploid_model") - haploid_hp_model = self.model_bundle.joinpath("haploid_hp_model") + 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" @@ -1022,7 +988,7 @@ def add_small_variant_calling( DNAscopeHP( diploid_unphased_hp, dbsnp=self.dbsnp, - model=self.model_bundle.joinpath("diploid_hp_model"), + model=bundle.joinpath("diploid_hp_model"), pcr_indel_model=self.repeat_model, ) ], @@ -1068,7 +1034,7 @@ def add_small_variant_calling( raw_vcf=diploid_unphased_patch, transfer=unphased_transfer, apply=ApplySpec( - model=self.model_bundle.joinpath("diploid_model_unphased"), + model=bundle.joinpath("diploid_model_unphased"), output=diploid_unphased, name="model-apply-unphased", ), @@ -1078,7 +1044,7 @@ def add_small_variant_calling( # 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( @@ -1101,7 +1067,7 @@ def add_small_variant_calling( 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, @@ -1142,7 +1108,7 @@ def add_small_variant_calling( dbsnp=self.dbsnp, emit_mode="gvcf", ploidy=1, - model=self.model_bundle.joinpath("gvcf_model"), + model=bundle.joinpath("gvcf_model"), ) ) haploid_call = DNAscopeStage( @@ -1169,7 +1135,7 @@ def add_small_variant_calling( haploid_concat_job = Job( cmds.bcftools_concat( - self.output_vcf, + ctx.output_vcf, [diploid_merged_vcf, haploid_out_fn], ), "haploid-diploid-concat", @@ -1182,7 +1148,7 @@ def add_small_variant_calling( return output_gvcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", ".g.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", ".g.vcf.gz") ) diploid_gvcf = pathlib.Path( str(diploid_merged_vcf).replace(".vcf.gz", ".g.vcf.gz") @@ -1192,7 +1158,7 @@ def add_small_variant_calling( ) haploid_gvcf_combine_job = Job( cmds.cmd_pyexec_gvcf_combine( - self.reference, + ctx.reference, str(haploid_gvcf_fn), str(haploid_out_fn), self.cores, @@ -1226,13 +1192,8 @@ def call_svs( """ 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) + bundle = self.required(self.model_bundle, "model_bundle") + require_versions(SV_MIN_VERSIONS, skip=self.skip_version_check) sv_vcf = pathlib.Path( str(self.output_vcf).replace(".vcf.gz", ".sv.vcf.gz") @@ -1247,7 +1208,7 @@ def call_svs( driver.add_algo( LongReadSV( sv_vcf, - self.model_bundle.joinpath("longreadsv.model"), + bundle.joinpath("longreadsv.model"), ) ) longreadsv_job = Job( diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index a021dbf..1245b3b 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -48,9 +48,9 @@ from .util import ( __version__, check_kmc_patch, - check_version, parse_rg_line, path_arg, + require_versions, total_memory, vcf_id, ) @@ -258,19 +258,18 @@ def validate(self) -> None: self.collect_readgroups() self.validate_readgroups() - if not self.skip_version_check: - for cmd, min_version in HYBRID_PANGENOME_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions( + HYBRID_PANGENOME_MIN_VERSIONS, skip=self.skip_version_check + ) - 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 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( @@ -599,35 +598,25 @@ def configure(self) -> None: def build_dag(self) -> DAG: """Build the DAG for the hybrid-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) - if not self.pop_vcf: - self.logger.error("pop_vcf is required") - sys.exit(2) + 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(self.reference) + ".fai") + ref_fai = pathlib.Path(str(ctx.reference) + ".fai") # 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_lift_aln = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", f"_lift_deduped.{suffix}") + str(ctx.output_vcf).replace(".vcf.gz", f"_lift_deduped.{suffix}") ) sv_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") + str(ctx.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") ) # Intermediate file paths @@ -762,7 +751,7 @@ def build_dag(self) -> DAG: # 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(self.output_vcf) + 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) @@ -822,7 +811,7 @@ def build_dag(self) -> DAG: if self.skip_small_variants: return dag - model = self.model_bundle.joinpath("dnascope.model") + model = bundle.joinpath("dnascope.model") pcr_indel_model = "NONE" if self.pcr_free else "CONSERVATIVE" call = DNAscopeStage( ctx=ctx, @@ -841,7 +830,7 @@ def build_dag(self) -> DAG: # transfer annotations from the pop_vcf transfer_target = ( - transfer_vcf if not self.skip_model_apply else self.output_vcf + transfer_vcf if not self.skip_model_apply else ctx.output_vcf ) transfer = TransferStage( ctx=ctx, @@ -855,7 +844,7 @@ def build_dag(self) -> DAG: ctx=ctx, model=model, vcf=transfer_vcf, - output=self.output_vcf, + output=ctx.output_vcf, ).add_to(dag, transfer.terminal) return dag diff --git a/sentieon_cli/pipeline.py b/sentieon_cli/pipeline.py index 428f3af..ed17af3 100644 --- a/sentieon_cli/pipeline.py +++ b/sentieon_cli/pipeline.py @@ -11,7 +11,7 @@ import shutil import sys import time -from typing import Any, Dict, List, Optional, Tuple, TypeVar +from typing import Any, Dict, List, Optional, TypeVar import packaging.version @@ -32,10 +32,9 @@ from .util import ( SampleSex, __version__, - check_version, - cnvscope_sex_args, path_arg, tmp, + versions_available, ) MULTIQC_MIN_VERSION = { @@ -409,25 +408,17 @@ def validate_cnv_par(self, cnv_will_run: bool) -> None: ) sys.exit(2) - def cnv_sex_args(self) -> Tuple[Optional[str], Optional[pathlib.Path]]: - """The CNVscope `--sex` and `--par` arguments for this run""" - return cnvscope_sex_args(self.sample_sex, self.cnv_par_bed) - 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 = MetricsPaths.from_output_vcf( pathlib.Path(str(self.output_vcf)) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 35c81ef..9807b07 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -49,10 +49,10 @@ from .util import ( SampleSex, __version__, - check_version, check_kmc_patch, parse_rg_line, path_arg, + require_versions, sample_sex_arg, tmp, total_memory, @@ -375,20 +375,19 @@ def validate(self) -> None: self.validate_t1k() self.validate_cnv() - 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) + require_versions( + SENT_PANGENOME_MIN_VERSIONS, skip=self.skip_version_check + ) - 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( @@ -461,10 +460,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_VERSION, skip=self.skip_version_check) def validate_expansion(self) -> None: if self.expansion_catalog is None: @@ -484,10 +480,7 @@ 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_VERSION, skip=self.skip_version_check) def validate_cnv(self) -> None: """Validate the arguments used for sex-aware CNV calling""" @@ -500,10 +493,7 @@ def validate_cnv(self) -> None: # A PAR BED file is required for CNV calling, whatever the sex self.validate_cnv_par(True) - if not self.skip_version_check: - for cmd, min_version in CNV_MIN_VERSIONS.items(): - if not check_version(cmd, min_version): - sys.exit(2) + require_versions(CNV_MIN_VERSIONS, skip=self.skip_version_check) def validate_t1k(self) -> None: hla_requested = self.t1k_hla_seq is not None or ( @@ -529,18 +519,11 @@ def validate_t1k(self) -> None: if not (hla_requested or kir_requested): 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) + require_versions(T1K_MIN_VERSION, 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) + pop_vcf = 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" ) @@ -573,7 +556,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.", @@ -613,7 +596,7 @@ def validate_bundle(self) -> None: ) if not self.skip_pop_vcf_id_check and not self.dry_run: - pop_vcf_id = vcf_id(self.pop_vcf) + pop_vcf_id = vcf_id(pop_vcf) if bundle_vcf_id != pop_vcf_id: self.logger.error( "The ID of the `--pop_vcf` does not match the model bundle" @@ -680,15 +663,7 @@ def build_dag(self) -> DAG: 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) + bundle = self.required(self.model_bundle, "model_bundle") ctx = self.stage_context() @@ -698,13 +673,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}") + 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 @@ -719,7 +694,7 @@ 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}" ) ) @@ -770,8 +745,8 @@ 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, @@ -813,7 +788,7 @@ 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(self.output_vcf) + 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) @@ -877,7 +852,7 @@ def build_first_dag(self) -> DAG: # 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") + str(ctx.output_vcf).replace(".vcf.gz", "_hla") ) extract_job, t1k_job = self.build_t1k_jobs( out_hla, @@ -892,7 +867,7 @@ def build_first_dag(self) -> DAG: 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") + str(ctx.output_vcf).replace(".vcf.gz", "_kir") ) extract_job, t1k_job = self.build_t1k_jobs( out_kir, @@ -910,7 +885,7 @@ def build_first_dag(self) -> DAG: 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: @@ -920,7 +895,7 @@ def build_first_dag(self) -> DAG: if self.tech.upper() == "ULTIMA": read_filters.append("UltimaReadFilter") pcr_indel_model = "NONE" if self.pcr_free else "CONSERVATIVE" - model = self.model_bundle.joinpath("dnascope.model") + model = bundle.joinpath("dnascope.model") gfa_file = sample_gfa if self.call_svs else None algos: List[BaseAlgo] = [] @@ -960,8 +935,8 @@ def build_first_dag(self) -> DAG: 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 @@ -1000,7 +975,7 @@ def build_first_dag(self) -> DAG: GVCFtyperStage( ctx=ctx, gvcf=out_gvcf, - output=self.output_vcf, + output=ctx.output_vcf, interval=self.bed, ).add_to(dag, {ad_update_job}) @@ -1033,9 +1008,7 @@ def bwa_extract_stage( The bwa alignment reuses the input readgroup ID with a `-bwa` suffix, so it stays distinct from the lifted alignment's. """ - if not self.model_bundle: - self.logger.error("model_bundle is required") - sys.exit(2) + bundle = self.required(self.model_bundle, "model_bundle") rg = copy.deepcopy(self.fastq_readgroup) rg["ID"] = rg["ID"] + "-bwa" @@ -1048,8 +1021,8 @@ def bwa_extract_stage( 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"), + extract_model=bundle.joinpath(self.extract_model_name), + bwa_model=bundle.joinpath("bwa.model"), unzip=find_unzip(self.logger, logging.INFO), ) @@ -1057,19 +1030,15 @@ 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", @@ -1124,12 +1093,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 @@ -1140,16 +1105,14 @@ 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, @@ -1196,12 +1159,8 @@ def build_segdup_job( 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) + reference = self.required(self.reference, "reference") + bundle = self.required(self.model_bundle, "model_bundle") sex = "male" if self.sample_sex == SampleSex.MALE else "female" @@ -1215,8 +1174,8 @@ def build_segdup_job( cmds.cmd_segdup_caller( out_segdup, sr_alignment, - reference=self.reference, - sr_bundle=self.model_bundle, + reference=reference, + sr_bundle=bundle, input_vcf=input_vcf, sex=sex, genes=genes, @@ -1279,9 +1238,7 @@ def build_t1k_jobs( 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) + self.required(self.reference, "reference") # Extract reads overlapping the locus to a BAM file with ReadWriter extracted_bam = self.tmp_dir.joinpath(f"sample_{tag}.bam") @@ -1317,15 +1274,13 @@ def build_expansion_job( expansion_catalog: pathlib.Path, ) -> Job: """Identify repeat expansions""" - if not self.reference: - self.logger.error("reference is required") - sys.exit(2) + reference = self.required(self.reference, "reference") return Job( cmds.cmd_expansion_hunter( out_expansions, sr_alignment, - reference=self.reference, + reference=reference, variant_catalog=expansion_catalog, sex="male" if self.sample_sex == SampleSex.MALE else "female", threads=self.cores, @@ -1342,22 +1297,14 @@ def _add_cnv_jobs( cnv_input_bams: List[pathlib.Path], ) -> 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) + bundle = self.required(self.model_bundle, "model_bundle") ctx = self.stage_context() 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") + str(ctx.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") ) # CNVscope on BWA deduped BAM, then CNVModelApply. The alignment @@ -1365,7 +1312,7 @@ def _add_cnv_jobs( cnv_result = CNVscopeStage( ctx=ctx, inputs=cnv_input_bams, - model=self.model_bundle.joinpath("cnv.model"), + model=bundle.joinpath("cnv.model"), cnvscope_vcf=cnvscope_vcf, cnv_vcf=cnv_apply_vcf, sample_sex=self.sample_sex, diff --git a/tests/unit/test_cnv_sex.py b/tests/unit/test_cnv_sex.py index 914d5a2..39efb2c 100644 --- a/tests/unit/test_cnv_sex.py +++ b/tests/unit/test_cnv_sex.py @@ -289,13 +289,6 @@ def test_validate_cnv_par_errors_in_a_dry_run_too(): assert excinfo.value.code == 2 -def test_cnv_sex_args_for_a_male_sample_with_a_par_bed(): - par_bed = pathlib.Path("/par.bed") - pipeline = _pipeline(sample_sex=SampleSex.MALE, cnv_par_bed=par_bed) - - assert pipeline.cnv_sex_args() == ("M", par_bed) - - 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") diff --git a/tests/unit/test_dnascope_hybrid_cnv.py b/tests/unit/test_dnascope_hybrid_cnv.py index 924a712..b77903d 100644 --- a/tests/unit/test_dnascope_hybrid_cnv.py +++ b/tests/unit/test_dnascope_hybrid_cnv.py @@ -82,19 +82,13 @@ def create_pipeline(self): def build_dag(self, pipeline): """Build the first DAG and index the jobs by name""" - with patch( - "sentieon_cli.dnascope_hybrid.check_version", return_value=True - ): - dag = pipeline.build_dag() + 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""" - with patch( - "sentieon_cli.dnascope_hybrid.check_version", return_value=True - ): - dag = pipeline.build_second_dag() + dag = pipeline.build_second_dag() if dag is None: return None, {} jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) 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 index 231fe8c..2c8c366 100644 --- a/tests/unit/test_dnascope_hybrid_stage1.py +++ b/tests/unit/test_dnascope_hybrid_stage1.py @@ -86,10 +86,7 @@ def create_pipeline(self): def build_dag(self, pipeline): """Build the DAG and index the jobs by name""" - with patch( - "sentieon_cli.dnascope_hybrid.check_version", return_value=True - ): - dag = pipeline.build_dag() + dag = pipeline.build_dag() jobs = list(dag.waiting_jobs.keys()) + list(dag.ready_jobs.keys()) return dag, {job.name: job for job in jobs} diff --git a/tests/unit/test_dnascope_longread_pop_vcf.py b/tests/unit/test_dnascope_longread_pop_vcf.py index 0b65130..8376b7e 100644 --- a/tests/unit/test_dnascope_longread_pop_vcf.py +++ b/tests/unit/test_dnascope_longread_pop_vcf.py @@ -154,8 +154,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())), \ @@ -177,8 +176,7 @@ def test_model_apply_dependency(self): 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())), \ diff --git a/tests/unit/test_sentieon_pangenome_kmc_integration.py b/tests/unit/test_sentieon_pangenome_kmc_integration.py index cfe4e61..9c0402f 100644 --- a/tests/unit/test_sentieon_pangenome_kmc_integration.py +++ b/tests/unit/test_sentieon_pangenome_kmc_integration.py @@ -33,7 +33,7 @@ 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.check_kmc_patch", return_value=True) as mock_check, \ patch("sys.exit") as mock_exit: @@ -59,7 +59,7 @@ def test_validate_kmc_patch_failure(self): 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.check_kmc_patch", return_value=False) as mock_check, \ patch("sys.exit") as mock_exit: From cfdc6a372a72153bf40795ffdd7e09b5fd859260 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Wed, 2 Sep 2026 09:57:00 -0700 Subject: [PATCH 37/58] Remove unused helpers Drop `cmd_bedtools_merge`, `cmd_vg_giraffe`, `cmd_vg_pack`, `cmd_sv_call`, `cmd_vg_surject` and `strip_ctg_prefix` from `command_strings`, `Shard.bcftools_str` from `shard`, and `get_read_length_aln` (with its regex and `shlex` import) from `util`. None had a caller in the package, the tests, or the docs. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/command_strings.py | 195 -------------------------------- sentieon_cli/shard.py | 3 - sentieon_cli/util.py | 44 ------- 3 files changed, 242 deletions(-) diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 4cd0ba1..4ed974f 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -48,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, @@ -1041,184 +1024,6 @@ 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], diff --git a/sentieon_cli/shard.py b/sentieon_cli/shard.py index 4506c40..9e509e5 100644 --- a/sentieon_cli/shard.py +++ b/sentieon_cli/shard.py @@ -93,9 +93,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""" diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index 67fa1c4..7475a65 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -9,7 +9,6 @@ import os import pathlib import re -import shlex import shutil import subprocess as sp import sys @@ -31,7 +30,6 @@ 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._-]") @@ -302,48 +300,6 @@ 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", - ) - - for line in res.stdout.split("\n"): - m = READ_LENGTH_PAT.match(line) - if m: - return int(m.groupdict()["length"]) - return 151 - - def vcf_id(in_vcf: pathlib.Path) -> Optional[str]: """Collect the SentieonVcfID header""" cmd = ["bcftools", "view", "-h", str(in_vcf)] From 98338cd1c74c8a0f04e9d9fa9e0622b883f42650 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 3 Sep 2026 20:48:58 -0700 Subject: [PATCH 38/58] Simplify the stages after review - `find_unzip` always warns; the per-caller log level is gone. - The bwa alignment stages report `cleanup_paths` only when the alignment is an intermediate that duplicate marking consumes, as the field's documentation already said; final outputs get none, and the pipelines build the `rm-*-aln` job only when there is something to remove. - `TransferApplyStage`'s DNAModelApply depends on the transfer's concat alone (the concat already depends on the raw-VCF producer), or on `upstream` when there is no transfer. The hybrid pipelines now use the composite too, and the long-read haplotype passes each run their own transfer-then-apply, so a haplotype's model-apply waits only on its own transfer. - `build_transfer_jobs` moves into `stages/transfer.py`, its only caller, and the top-level `transfer.py` is gone. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 23 ++-- sentieon_cli/dnascope_hybrid.py | 64 +++++----- sentieon_cli/dnascope_longread.py | 52 ++++---- sentieon_cli/hybrid_pangenome.py | 34 ++--- sentieon_cli/sentieon_pangenome.py | 3 +- sentieon_cli/stages/alignment.py | 22 ++-- sentieon_cli/stages/small_variants.py | 8 +- sentieon_cli/stages/transfer.py | 146 +++++++++++++++++++++- sentieon_cli/transfer.py | 152 ----------------------- tests/unit/test_dnascope_longread_dag.py | 26 ++-- tests/unit/test_stage_alignment.py | 35 +++--- tests/unit/test_stage_small_variants.py | 4 +- 12 files changed, 276 insertions(+), 293 deletions(-) delete mode 100644 sentieon_cli/transfer.py diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 7ed4ba8..431ea21 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -410,19 +410,14 @@ def build_dag(self) -> DAG: # 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: aln_result = self.add_sr_input_alignment(dag, ctx) sample_input = aln_result.outputs align_jobs = set(aln_result.jobs) - bam_rm_job = rm_job(aln_result.cleanup_paths, "rm-bam-aln") + bam_cleanup = aln_result.cleanup_paths fq_result = self.add_sr_fastq_alignment(dag, ctx) align_fastq_jobs = set(fq_result.jobs) - fq_rm_job = ( - rm_job(fq_result.cleanup_paths, "rm-fq-aln") - if fq_result.cleanup_paths - else None - ) sample_input += fq_result.outputs # Dedup and metrics @@ -431,11 +426,17 @@ def build_dag(self) -> DAG: ) 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_rm_job: - dag.add_job(bam_rm_job, {dedup_job}) - if fq_rm_job: - dag.add_job(fq_rm_job, {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: diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index e31fd98..fa767fe 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -32,8 +32,13 @@ from .stages.base import StageContext, driver_job, rm_job from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage from .stages.ploidy import PloidyStage -from .stages.small_variants import DNAscopeStage, ModelApplyStage -from .stages.transfer import TransferConfig, TransferStage +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + TransferApplyStage, + TransferSpec, +) +from .stages.transfer import TransferConfig from .util import ( __version__, library_preloaded, @@ -585,11 +590,6 @@ def build_dag(self) -> DAG: sr_aln = self.sr_aln fq_result = self.add_sr_fastq_alignment(dag, ctx) align_fastq_jobs = set(fq_result.jobs) - fq_rm_job = ( - rm_job(fq_result.cleanup_paths, "rm-fq-aln") - if fq_result.cleanup_paths - else None - ) sr_aln.extend(fq_result.outputs) # Short-read dedup @@ -598,8 +598,13 @@ def build_dag(self) -> DAG: ) sr_aln = preprocessing.deduped dedup_job = preprocessing.dedup_job - if dedup_job and fq_rm_job: - dag.add_job(fq_rm_job, {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() @@ -1061,36 +1066,37 @@ def add_variant_calling( ) dag.add_job(anno_job, {concat_job}) - input_to_apply = anno_target - apply_dependencies: Set[Job] = {anno_job} + 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 = ctx.output_vcf + transfer_spec = TransferSpec( + self.transfer_config(), transfer_target + ) - transfer = TransferStage( - ctx=ctx, - config=self.transfer_config(), - raw_vcf=anno_target, - out_vcf=transfer_target, - ).add_to(dag, {anno_job}) - input_to_apply = transfer_target - apply_dependencies = transfer.terminal + apply_spec: Optional[ApplySpec] = None + apply_vcf = self.tmp_dir.joinpath("combined_apply.vcf.gz") + 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 - # Model Apply - apply_vcf = self.tmp_dir.joinpath("combined_apply.vcf.gz") - apply = ModelApplyStage( - ctx=ctx, - model=hybrid_model, - vcf=input_to_apply, - output=apply_vcf, - ).add_to(dag, apply_dependencies) - # Final normalize norm_job = Job( cmds.filter_norm( @@ -1103,4 +1109,4 @@ def add_variant_calling( 0, task_name="vcf-norm", ) - dag.add_job(norm_job, apply.terminal) + dag.add_job(norm_job, transfer_apply.terminal) diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index 2ebdb32..0ee2864 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -43,11 +43,10 @@ from .stages.small_variants import ( ApplySpec, DNAscopeStage, - ModelApplyStage, TransferApplyStage, TransferSpec, ) -from .stages.transfer import TransferConfig, TransferStage +from .stages.transfer import TransferConfig from .util import ( __version__, library_preloaded, @@ -940,28 +939,9 @@ def add_small_variant_calling( ) dag.add_job(haploid_patch_job, haploid_patch_deps) - # Transfer annotations to the patched VCFs. Both model-apply jobs - # wait on both transfers, as they did when the DAG was wired by hand. - second_ma_deps: Set[Job] = {haploid_patch_job} - 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 i, (patch_unanno_vcf, patch_vcf) in enumerate( - zip(patch_unanno_vcfs, patch_vcfs), start=1 - ): - patch_transfer = TransferStage( - ctx=ctx, - config=self.transfer_config(), - raw_vcf=patch_unanno_vcf, - out_vcf=patch_vcf, - tag=f"hap{i}", - ).add_to(dag, {haploid_patch_job}) - second_ma_deps.update(patch_transfer.terminal) - - # 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) ] @@ -969,13 +949,25 @@ def add_small_variant_calling( for i, (patch_vcf, hap_vcf) in enumerate( zip(patch_vcfs, hap_vcfs), start=1 ): - hap_apply = ModelApplyStage( + patch_transfer: Optional[TransferSpec] = None + if self.pop_vcf: + patch_transfer = TransferSpec( + config=self.transfer_config(), + out_vcf=self.tmp_dir.joinpath( + f"out_hap{i}_patch_anno.vcf.gz" + ), + tag=f"hap{i}", + ) + hap_apply = TransferApplyStage( ctx=ctx, - model=haploid_model, - vcf=patch_vcf, - output=hap_vcf, - name=f"model-apply-hap{i}", - ).add_to(dag, second_ma_deps) + 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 diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 1245b3b..f6b09cb 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -8,7 +8,6 @@ import argparse import copy import json -import logging import pathlib import sys from typing import Dict, List, Optional, Set, Tuple @@ -43,8 +42,13 @@ from .stages.base import StageContext from .stages.dedup import DedupStage from .stages.metrics import MetricsPaths, MetricsStage -from .stages.small_variants import DNAscopeStage, ModelApplyStage -from .stages.transfer import TransferConfig, TransferStage +from .stages.small_variants import ( + ApplySpec, + DNAscopeStage, + TransferApplyStage, + TransferSpec, +) +from .stages.transfer import TransferConfig from .util import ( __version__, check_kmc_patch, @@ -828,25 +832,21 @@ def build_dag(self) -> DAG: replace_rg=replace_rg, ).add_to(dag, calling_dependencies) - # transfer annotations from the pop_vcf + # transfer annotations from the pop_vcf, then apply the model transfer_target = ( transfer_vcf if not self.skip_model_apply else ctx.output_vcf ) - transfer = TransferStage( + TransferApplyStage( ctx=ctx, - config=self.transfer_config(), raw_vcf=raw_vcf, - out_vcf=transfer_target, + transfer=TransferSpec(self.transfer_config(), transfer_target), + apply=( + ApplySpec(model, ctx.output_vcf) + if not self.skip_model_apply + else None + ), ).add_to(dag, call.terminal) - if not self.skip_model_apply: - ModelApplyStage( - ctx=ctx, - model=model, - vcf=transfer_vcf, - output=ctx.output_vcf, - ).add_to(dag, transfer.terminal) - return dag def transfer_config(self) -> TransferConfig: @@ -886,7 +886,7 @@ def build_hybrid_kmc_job(self, kmer_prefix: pathlib.Path) -> Job: self.tmp_dir, memory=self.kmer_memory, threads=self.cores, - unzip=find_unzip(self.logger, logging.INFO), + unzip=find_unzip(self.logger), ), "kmc", 0, # run in the background @@ -975,7 +975,7 @@ def bwa_extract_stage( ), extract_model=self.model_bundle.joinpath(self.extract_model_name), bwa_model=self.model_bundle.joinpath("bwa.model"), - unzip=find_unzip(self.logger, logging.INFO), + unzip=find_unzip(self.logger), ) def build_haplotypes_job( diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 9807b07..1420360 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -5,7 +5,6 @@ import argparse import copy import json -import logging import pathlib import shutil import sys @@ -1023,7 +1022,7 @@ def bwa_extract_stage( ), extract_model=bundle.joinpath(self.extract_model_name), bwa_model=bundle.joinpath("bwa.model"), - unzip=find_unzip(self.logger, logging.INFO), + unzip=find_unzip(self.logger), ) def build_haplotypes_job( diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py index f474b42..d12e6dc 100644 --- a/sentieon_cli/stages/alignment.py +++ b/sentieon_cli/stages/alignment.py @@ -30,17 +30,15 @@ def aln_suffix(bam_format: bool) -> str: return "bam" if bam_format else "cram" -def find_unzip(logger: logging.Logger, level: int = logging.WARNING) -> str: +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``. Callers differ on how loudly that is worth - reporting, hence ``level``. + back to ``gzip``. """ unzip = "igzip" if not shutil.which(unzip): - logger.log( - level, + logger.warning( "igzip is recommended for decompression, but is not " "available. Falling back to gzip.", ) @@ -151,10 +149,9 @@ def build(self) -> AlignResult: ) ) outputs.append(out_aln) - cleanup_paths.append(out_aln) - cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) - if suffix == "cram": - cleanup_paths.append(pathlib.Path(str(out_aln) + ".crai")) + if dedup_follows: + cleanup_paths.append(out_aln) + cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) return AlignResult( jobs=jobs, @@ -241,10 +238,9 @@ def build(self) -> AlignResult: ) ) outputs.append(out_aln) - cleanup_paths.append(out_aln) - cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) - if suffix == "cram": - cleanup_paths.append(pathlib.Path(str(out_aln) + ".crai")) + if dedup_follows: + cleanup_paths.append(out_aln) + cleanup_paths.append(pathlib.Path(str(out_aln) + ".bai")) return AlignResult( jobs=jobs, diff --git a/sentieon_cli/stages/small_variants.py b/sentieon_cli/stages/small_variants.py index 9757781..4880b55 100644 --- a/sentieon_cli/stages/small_variants.py +++ b/sentieon_cli/stages/small_variants.py @@ -174,9 +174,9 @@ 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 *and* on ``upstream`` -- - the raw-VCF producer stays a direct dependency, matching the DAGs the - pipelines built by hand. + 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 @@ -206,7 +206,7 @@ def add_to( ).add_to(dag, deps) jobs.extend(transfer_result.jobs) apply_in = self.transfer.out_vcf - apply_deps = deps | transfer_result.terminal + apply_deps = transfer_result.terminal apply_job: Optional[Job] = None if self.apply is not None: diff --git a/sentieon_cli/stages/transfer.py b/sentieon_cli/stages/transfer.py index 7ba294b..fe9a0cd 100644 --- a/sentieon_cli/stages/transfer.py +++ b/sentieon_cli/stages/transfer.py @@ -4,14 +4,156 @@ from dataclasses import dataclass import pathlib -from typing import Dict, Iterable, List, Optional +import re +import subprocess as sp +import tempfile +from typing import Dict, Iterable, List, Optional, Set, Tuple +from importlib.resources import files + +from .. import command_strings as cmds from ..dag import DAG from ..job import Job +from ..logging import get_logger from ..shard import Shard -from ..transfer import build_transfer_jobs from .base import Stage, StageResult +logger = get_logger(__name__) + + +def build_transfer_jobs( + out_vcf: pathlib.Path, + pop_vcf: pathlib.Path, + raw_vcf: pathlib.Path, + base_tmp_dir: pathlib.Path, + shards: List[Shard], + pop_vcf_contigs: Dict[str, Optional[int]], + 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. + + ``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) + tmp_dir = pathlib.Path(tmp_dir_str) + + # Generate merge rules from the population VCF + merge_rules = "AC_v20:sum,AF_v20:sum,AC_genomes:sum,AF_genomes:sum" + if not dry_run: + kvpat = re.compile(r'(.*?)=(".*?"|.*?)(?:,|$)') + cmd = ["bcftools", "view", "-h", str(pop_vcf)] + p = sp.run(cmd, capture_output=True, text=True) + id_fields: List[str] = [] + for line in p.stdout.split("\n"): + if not line.startswith("##INFO"): + continue + if ",Number=A" not in line: + continue + s = line.index("<") + e = line.index(">") + d = dict(kvpat.findall(line[s + 1 : e])) # noqa: E203 + id_fields.append(d["ID"]) + merge_rules = ",".join([x + ":sum" for x in id_fields]) + + # Merge VCFs by shards + sharded_vcfs: List[pathlib.Path] = [] + sharded_merge_jobs: List[Job] = [] + trim_script = pathlib.Path( + str(files("sentieon_cli.scripts").joinpath("trimalt.py")) + ).resolve() + seen_contigs: Set[str] = set() + for i, shard in enumerate(shards): + # Use a BED file for unusual contig names + subset_bed = tmp_dir.joinpath( + f"sample-dnascope_transfer-subset{i}.bed" + ) + + # Extract contigs not in the pop vcf as merge will fail + if shard.contig not in pop_vcf_contigs: + if shard.contig in seen_contigs: + continue + logger.info("Skipping transfer for contig: %s", shard.contig) + seen_contigs.add(shard.contig) + subset_vcf = tmp_dir.joinpath( + f"sample-dnascope_transfer-subset{i}.vcf.gz" + ) + + ctg_len = fai_data[shard.contig]["length"] + if not dry_run: + with open(subset_bed, "w") as fh: + print(f"{shard.contig}\t0\t{ctg_len}", file=fh) + + view_job = Job( + cmds.cmd_bcftools_view_regions( + subset_vcf, + raw_vcf, + regions_file=subset_bed, + ), + f"{name_base}-extra", + 1, + task_name="annotation-transfer", + ) + sharded_merge_jobs.append(view_job) + sharded_vcfs.append(subset_vcf) + else: + if not dry_run: + with open(subset_bed, "w") as fh: + print( + f"{shard.contig}\t{shard.start}\t{shard.stop}", + file=fh, + ) + + logger.debug("Transferring shard: %s", shard) + shard_vcf = tmp_dir.joinpath( + f"sample-dnascope_transfer-shard{i}.vcf.gz" + ) + merge_job = Job( + cmds.cmd_bcftools_merge_trim( + shard_vcf, + raw_vcf, + pop_vcf, + trim_script, + subset_bed, + merge_rules=merge_rules, + merge_xargs=[ + "--no-version", + "--regions-overlap", + "pos", + "-m", + "all", + ], + view_xargs=["--no-version"], + ), + f"{name_base}-{i}", + 1, + task_name="annotation-transfer", + ) + sharded_merge_jobs.append(merge_job) + sharded_vcfs.append(shard_vcf) + + # Concat all shards + concat_job = Job( + cmds.bcftools_concat( + out_vcf, + sharded_vcfs, + xargs=["--no-version", "--threads", str(cores)], + ), + f"{name_base}-concat", + cores, + task_name="annotation-transfer", + ) + return (sharded_merge_jobs, concat_job) + @dataclass(frozen=True) class TransferConfig: diff --git a/sentieon_cli/transfer.py b/sentieon_cli/transfer.py deleted file mode 100644 index 24f9c11..0000000 --- a/sentieon_cli/transfer.py +++ /dev/null @@ -1,152 +0,0 @@ -""" -Annotation transfer functionality -""" - -import pathlib -import re -import subprocess as sp -import tempfile -from typing import Dict, List, Optional, 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 - -logger = get_logger(__name__) - - -def build_transfer_jobs( - out_vcf: pathlib.Path, - pop_vcf: pathlib.Path, - raw_vcf: pathlib.Path, - base_tmp_dir: pathlib.Path, - shards: List[Shard], - pop_vcf_contigs: Dict[str, Optional[int]], - 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. - - ``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) - tmp_dir = pathlib.Path(tmp_dir_str) - - # Generate merge rules from the population VCF - merge_rules = "AC_v20:sum,AF_v20:sum,AC_genomes:sum,AF_genomes:sum" - if not dry_run: - kvpat = re.compile(r'(.*?)=(".*?"|.*?)(?:,|$)') - cmd = ["bcftools", "view", "-h", str(pop_vcf)] - p = sp.run(cmd, capture_output=True, text=True) - id_fields: List[str] = [] - for line in p.stdout.split("\n"): - if not line.startswith("##INFO"): - continue - if ",Number=A" not in line: - continue - s = line.index("<") - e = line.index(">") - d = dict(kvpat.findall(line[s + 1 : e])) # noqa: E203 - id_fields.append(d["ID"]) - merge_rules = ",".join([x + ":sum" for x in id_fields]) - - # Merge VCFs by shards - sharded_vcfs: List[pathlib.Path] = [] - sharded_merge_jobs: List[Job] = [] - trim_script = pathlib.Path( - str(files("sentieon_cli.scripts").joinpath("trimalt.py")) - ).resolve() - seen_contigs: Set[str] = set() - for i, shard in enumerate(shards): - # Use a BED file for unusual contig names - subset_bed = tmp_dir.joinpath( - f"sample-dnascope_transfer-subset{i}.bed" - ) - - # Extract contigs not in the pop vcf as merge will fail - if shard.contig not in pop_vcf_contigs: - if shard.contig in seen_contigs: - continue - logger.info("Skipping transfer for contig: %s", shard.contig) - seen_contigs.add(shard.contig) - subset_vcf = tmp_dir.joinpath( - f"sample-dnascope_transfer-subset{i}.vcf.gz" - ) - - ctg_len = fai_data[shard.contig]["length"] - if not dry_run: - with open(subset_bed, "w") as fh: - print(f"{shard.contig}\t0\t{ctg_len}", file=fh) - - view_job = Job( - cmds.cmd_bcftools_view_regions( - subset_vcf, - raw_vcf, - regions_file=subset_bed, - ), - f"{name_base}-extra", - 1, - task_name="annotation-transfer", - ) - sharded_merge_jobs.append(view_job) - sharded_vcfs.append(subset_vcf) - else: - if not dry_run: - with open(subset_bed, "w") as fh: - print( - f"{shard.contig}\t{shard.start}\t{shard.stop}", - file=fh, - ) - - logger.debug("Transferring shard: %s", shard) - shard_vcf = tmp_dir.joinpath( - f"sample-dnascope_transfer-shard{i}.vcf.gz" - ) - merge_job = Job( - cmds.cmd_bcftools_merge_trim( - shard_vcf, - raw_vcf, - pop_vcf, - trim_script, - subset_bed, - merge_rules=merge_rules, - merge_xargs=[ - "--no-version", - "--regions-overlap", - "pos", - "-m", - "all", - ], - view_xargs=["--no-version"], - ), - f"{name_base}-{i}", - 1, - task_name="annotation-transfer", - ) - sharded_merge_jobs.append(merge_job) - sharded_vcfs.append(shard_vcf) - - # Concat all shards - concat_job = Job( - cmds.bcftools_concat( - out_vcf, - sharded_vcfs, - xargs=["--no-version", "--threads", str(cores)], - ), - f"{name_base}-concat", - cores, - task_name="annotation-transfer", - ) - return (sharded_merge_jobs, concat_job) diff --git a/tests/unit/test_dnascope_longread_dag.py b/tests/unit/test_dnascope_longread_dag.py index c48bcca..e239146 100644 --- a/tests/unit/test_dnascope_longread_dag.py +++ b/tests/unit/test_dnascope_longread_dag.py @@ -129,20 +129,26 @@ def test_merge_waits_on_the_three_model_apply_jobs(self, tmp_path): "model-apply-unphased", ] - def test_both_haplotype_applies_wait_on_both_transfers(self, tmp_path): + 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) - # Each model-apply reads only its own haplotype's transfer, but - # both wait on both concat jobs, as the hand-wired DAG did - for hap in ("model-apply-hap1", "model-apply-hap2"): - assert deps_of(dag, hap) == [ - "merge-trim-hap1-concat", - "merge-trim-hap2-concat", - "patch", + # 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() @@ -157,12 +163,12 @@ def test_each_transfer_is_tagged(self, tmp_path): "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") == [ - "dnascope-diploid", "merge-trim-diploid-concat", ] assert deps_of(dag, "model-apply-unphased") == [ - "diploid-patch", "merge-trim-unphased-concat", ] diff --git a/tests/unit/test_stage_alignment.py b/tests/unit/test_stage_alignment.py index 3896437..e4b7df2 100644 --- a/tests/unit/test_stage_alignment.py +++ b/tests/unit/test_stage_alignment.py @@ -69,15 +69,6 @@ def test_find_unzip_falls_back_to_gzip(self, caplog): assert "igzip is recommended" in caplog.text assert caplog.records[0].levelno == logging.WARNING - def test_find_unzip_takes_the_log_level(self, caplog): - logger = logging.getLogger("test-unzip-level") - with patch( - "sentieon_cli.stages.alignment.shutil.which", return_value=None - ): - with caplog.at_level(logging.INFO, logger=logger.name): - assert find_unzip(logger, logging.INFO) == "gzip" - assert caplog.records[0].levelno == logging.INFO - class TestBwaRealignStage: """Re-aligning BAM/CRAM input with bwa""" @@ -143,7 +134,7 @@ def test_bam_format_without_duplicate_marking(self, tmp_path, rg_lines): tmp_path / "output_bwa_sorted_1.bam", ] - def test_bam_cleanup_paths(self, tmp_path, rg_lines): + def test_intermediate_cleanup_paths(self, tmp_path, rg_lines): result = self.make(tmp_path).build() assert result.cleanup_paths == [ @@ -153,15 +144,10 @@ def test_bam_cleanup_paths(self, tmp_path, rg_lines): pathlib.Path(str(tmp_path / "bwa_sorted_1.bam") + ".bai"), ] - def test_cram_cleanup_paths_add_the_crai(self, tmp_path, rg_lines): + def test_final_outputs_have_no_cleanup_paths(self, tmp_path, rg_lines): result = self.make(tmp_path, duplicate_marking="none").build() - out = str(tmp_path / "output_bwa_sorted_0.cram") - assert result.cleanup_paths[:3] == [ - pathlib.Path(out), - pathlib.Path(out + ".bai"), - pathlib.Path(out + ".crai"), - ] + assert result.cleanup_paths == [] def test_collate_and_input_ref(self, tmp_path, rg_lines): result = self.make( @@ -254,17 +240,24 @@ def test_unzip_reaches_the_command(self, tmp_path): assert "igzip -dc" in str(result.jobs[0].shell) - def test_outputs_and_cleanup_paths(self, tmp_path): - result = self.make(tmp_path, duplicate_marking="none").build() + def test_intermediate_outputs_and_cleanup_paths(self, tmp_path): + result = self.make(tmp_path).build() - out = tmp_path / "output_bwa_sorted_fq_0_0.cram" + out = tmp_path / "bwa_sorted_fq_0_0.bam" assert result.outputs == [out] assert result.cleanup_paths == [ out, pathlib.Path(str(out) + ".bai"), - pathlib.Path(str(out) + ".crai"), ] + 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 == [] + class TestMinimap2RealignStage: """Re-aligning long-read input with minimap2""" diff --git a/tests/unit/test_stage_small_variants.py b/tests/unit/test_stage_small_variants.py index 2a561ec..56ae89f 100644 --- a/tests/unit/test_stage_small_variants.py +++ b/tests/unit/test_stage_small_variants.py @@ -345,9 +345,9 @@ def test_both_halves(self, tmp_path): assert f"--vcf {tmp_path}/transfer.vcf.gz" in str( result.apply_job.shell ) - # ... and waits on both the transfer and the raw-VCF producer + # ... and waits on the transfer alone; the concat already depends + # on the raw-VCF producer assert dag.waiting_jobs[result.apply_job] == { - upstream, result.transfer.concat_job, } assert result.terminal == {result.apply_job} From 7ab0e4e1e9c71eb142dd76a078aa1deda677ed57 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 3 Sep 2026 22:13:55 -0700 Subject: [PATCH 39/58] Move short-read preprocessing, mosdepth and LongReadSV into stages `ShortReadPreprocessingStage` carries the duplicate-marking and metrics policy that lived in `DNAscopePipeline.add_sr_preprocessing`; `MosdepthStage` and `LongReadSVStage` replace the raw job construction in the long-read pipeline's `mosdepth` and `call_svs`. The alignment, mosdepth and LongReadSV version constants move next to the stages they guard. The pipeline methods remain as thin wrappers that hold the version-check policy, so the hybrid pipeline, which still inherits them, is unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 151 ++--------- sentieon_cli/dnascope_longread.py | 88 ++----- sentieon_cli/stages/alignment.py | 20 ++ sentieon_cli/stages/metrics.py | 52 ++++ sentieon_cli/stages/preprocessing.py | 179 +++++++++++++ sentieon_cli/stages/sv.py | 70 +++++ tests/unit/test_stage_metrics.py | 90 ++++++- tests/unit/test_stage_preprocessing.py | 342 +++++++++++++++++++++++++ tests/unit/test_stage_sv.py | 144 +++++++++++ 9 files changed, 940 insertions(+), 196 deletions(-) create mode 100644 sentieon_cli/stages/preprocessing.py create mode 100644 sentieon_cli/stages/sv.py create mode 100644 tests/unit/test_stage_preprocessing.py create mode 100644 tests/unit/test_stage_sv.py diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 431ea21..3d16a5d 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -4,7 +4,6 @@ import argparse import copy -from dataclasses import dataclass, field import itertools import os import pathlib @@ -16,30 +15,25 @@ from .dag import DAG from .driver import ( - AlignmentStat, BaseAlgo, - BaseDistributionByCycle, - CoverageMetrics, DNAscope, - GCBias, - HsMetricAlgo, - InsertSizeMetricAlgo, - MeanQualityByCycle, - QualDistribution, SVSolver, - WgsMetricsAlgo, ) from .job import Job from .pipeline import BasePipeline 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.dedup import DedupStage -from .stages.metrics import MetricsPaths, MetricsStage +from .stages.preprocessing import ( + ShortReadPreprocessingStage, + SrPreprocessingResult, +) from .stages.small_variants import ( ApplySpec, DNAscopeStage, @@ -55,35 +49,11 @@ 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"), } -@dataclass -class SrPreprocessing: - """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] = None - qc_jobs: List[Job] = field(default_factory=list) - - class DNAscopePipeline(BasePipeline): """The DNAscope pipeline""" @@ -463,7 +433,9 @@ def add_sr_input_alignment( """Align input BAM/CRAM/uBAM/uCRAM files with bwa""" bundle = self.required(self.model_bundle, "model_bundle") - require_versions(ALN_MIN_VERSIONS, skip=self.skip_version_check) + require_versions( + BWA_REALIGN_MIN_VERSIONS, skip=self.skip_version_check + ) return BwaRealignStage( ctx=ctx, @@ -489,7 +461,7 @@ def add_sr_fastq_alignment( jobs=[], terminal=set(), outputs=[], cleanup_paths=[] ) - require_versions(FQ_MIN_VERSIONS, skip=self.skip_version_check) + require_versions(BWA_FASTQ_MIN_VERSIONS, skip=self.skip_version_check) return BwaFastqStage( ctx=ctx, @@ -506,113 +478,24 @@ def add_sr_fastq_alignment( util_sort_args=self.util_sort_args, ).add_to(dag) - def sr_metrics_algos(self, paths: MetricsPaths) -> 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 self.assay == "WGS": - algos.append(GCBias(paths.gc_bias, summary=paths.gc_bias_summary)) - return algos - def add_sr_preprocessing( self, dag: DAG, ctx: StageContext, sample_input: List[pathlib.Path], upstream: Set[Job], - ) -> SrPreprocessing: + ) -> SrPreprocessingResult: """Mark duplicates and collect metrics from the short reads""" - suffix = "bam" if self.bam_format else "cram" - - paths = MetricsPaths.from_output_vcf(ctx.output_vcf) - paths.ensure_dir(self.dry_run) - - if self.sr_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 SrPreprocessing( - deduped=sample_input, dedup_job=None, qc_jobs=[] - ) - metrics_result = MetricsStage( - ctx=ctx, - inputs=sample_input, - algos=[ - InsertSizeMetricAlgo(paths.insert_size), - *self.sr_metrics_algos(paths), - ], - name="metrics", - task_name="metrics", - threads=ctx.cores, - ).add_to(dag, upstream) - return SrPreprocessing( - deduped=sample_input, - 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(self.sr_metrics_algos(paths)) - - deduped = pathlib.Path( - str(ctx.output_vcf).replace(".vcf.gz", f"_deduped.{suffix}") - ) - dedup_result = DedupStage( + return ShortReadPreprocessingStage( ctx=ctx, inputs=sample_input, - output=deduped, - score_file=paths.score, + duplicate_marking=self.sr_duplicate_marking, consensus=self.consensus, - rmdup=(self.sr_duplicate_marking == "rmdup"), - cram_write_options="version=3.0,compressor=rans", - dedup_metrics=paths.dedup_metrics, - lc_extra_algos=lc_extra_algos, + skip_metrics=self.skip_metrics, + assay=self.assay, + bed=self.bed, + bam_format=self.bam_format, ).add_to(dag, upstream) - 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}) - qc_jobs.extend(metrics_result.jobs) - - return SrPreprocessing( - deduped=[deduped], - dedup_job=dedup_result.dedup_job, - qc_jobs=qc_jobs, - ) def add_small_variant_calling( self, diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index 0ee2864..79e3a54 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -18,10 +18,8 @@ from .dag import DAG from .driver import ( BaseAlgo, - Driver, DNAscope, DNAscopeHP, - LongReadSV, RepeatModel, VariantPhaser, ) @@ -34,18 +32,22 @@ ) from .shell_pipeline import Command, Pipeline 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__, @@ -63,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, } @@ -495,7 +480,9 @@ def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: """ output_vcf = self.required(self.output_vcf, "output_vcf") bundle = self.required(self.model_bundle, "model_bundle") - require_versions(ALN_MIN_VERSIONS, skip=self.skip_version_check) + require_versions( + MINIMAP2_REALIGN_MIN_VERSIONS, skip=self.skip_version_check + ) result = Minimap2RealignStage( ctx=self.stage_context(), @@ -518,7 +505,9 @@ def lr_align_fastq(self) -> Tuple[List[pathlib.Path], Set[Job]]: if self.fastq is None and self.readgroups is None: return ([], set()) - require_versions(FQ_MIN_VERSIONS, skip=self.skip_version_check) + require_versions( + MINIMAP2_FASTQ_MIN_VERSIONS, skip=self.skip_version_check + ) result = Minimap2FastqStage( ctx=self.stage_context(), @@ -544,25 +533,14 @@ def mosdepth(self, sample_input: List[pathlib.Path]) -> Set[Job]: ) 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}") - ) - mosdepth_jobs.add( - Job( - cmds.cmd_mosdepth( - input_file, - mosdepth_dir, - fasta=self.reference, - threads=self.cores, - ), - f"mosdepth-{i}", - 0, # Run in background - task_name="metrics", - ) + 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] @@ -1185,28 +1163,16 @@ def call_svs( Call SVs using Sentieon LongReadSV """ bundle = self.required(self.model_bundle, "model_bundle") - require_versions(SV_MIN_VERSIONS, skip=self.skip_version_check) + require_versions(LONGREADSV_MIN_VERSIONS, skip=self.skip_version_check) - 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, - bundle.joinpath("longreadsv.model"), + return ( + LongReadSVStage( + ctx=self.stage_context(), + inputs=sample_input, + model=bundle.joinpath("longreadsv.model"), + interval=self.bed, + replace_rg=replace_rg, ) + .build() + .job ) - longreadsv_job = Job( - Pipeline(Command(*driver.build_cmd())), - "LongReadSV", - self.cores, - task_name="sv-calling", - ) - return longreadsv_job diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py index d12e6dc..ae6769a 100644 --- a/sentieon_cli/stages/alignment.py +++ b/sentieon_cli/stages/alignment.py @@ -15,6 +15,8 @@ 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 @@ -24,6 +26,24 @@ # The `sentieon util sort` arguments every aligner defaults to DEFAULT_UTIL_SORT_ARGS = "--cram_write_options version=3.0,compressor=rans" +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""" diff --git a/sentieon_cli/stages/metrics.py b/sentieon_cli/stages/metrics.py index ee9d82b..e03c9a3 100644 --- a/sentieon_cli/stages/metrics.py +++ b/sentieon_cli/stages/metrics.py @@ -8,12 +8,19 @@ 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: @@ -179,3 +186,48 @@ def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> MetricsResult: 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/preprocessing.py b/sentieon_cli/stages/preprocessing.py new file mode 100644 index 0000000..f2a29bd --- /dev/null +++ b/sentieon_cli/stages/preprocessing.py @@ -0,0 +1,179 @@ +""" +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"), + cram_write_options="version=3.0,compressor=rans", + 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/sv.py b/sentieon_cli/stages/sv.py new file mode 100644 index 0000000..03e0eef --- /dev/null +++ b/sentieon_cli/stages/sv.py @@ -0,0 +1,70 @@ +""" +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. + + The output VCF is derived from the run's output VCF, so the caller + only supplies the alignment and the model. + """ + + inputs: List[pathlib.Path] + model: pathlib.Path + 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 = pathlib.Path( + str(self.ctx.output_vcf).replace(".vcf.gz", ".sv.vcf.gz") + ) + job = driver_job( + self.ctx, + [LongReadSV(sv_vcf, self.model)], + 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/tests/unit/test_stage_metrics.py b/tests/unit/test_stage_metrics.py index 1792560..7170b1d 100644 --- a/tests/unit/test_stage_metrics.py +++ b/tests/unit/test_stage_metrics.py @@ -12,7 +12,11 @@ WgsMetricsAlgo, ) from sentieon_cli.stages.base import StageContext, rm_job -from sentieon_cli.stages.metrics import MetricsPaths, MetricsStage +from sentieon_cli.stages.metrics import ( + MetricsPaths, + MetricsStage, + MosdepthStage, +) def make_ctx(tmp_path: pathlib.Path, cores: int = 4) -> StageContext: @@ -242,3 +246,87 @@ def test_metrics_is_a_root_without_upstream(self, tmp_path): 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_preprocessing.py b/tests/unit/test_stage_preprocessing.py new file mode 100644 index 0000000..6ccf8b7 --- /dev/null +++ b/tests/unit/test_stage_preprocessing.py @@ -0,0 +1,342 @@ +""" +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_and_cram_options(self, tmp_path): + paths = MetricsPaths.from_output_vcf(tmp_path / "output.vcf.gz") + result = make_stage(tmp_path).add_to(DAG()) + + shell = str(result.dedup_job.shell) + assert f"--metrics {paths.dedup_metrics}" in shell + assert "--cram_write_options version=3.0,compressor=rans" 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_sv.py b/tests/unit/test_stage_sv.py new file mode 100644 index 0000000..c7f86aa --- /dev/null +++ b/tests/unit/test_stage_sv.py @@ -0,0 +1,144 @@ +""" +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 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) + + +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 From b3289b7bdeba4f6b4abdfc2e3fec851e8aabe478 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 3 Sep 2026 22:21:32 -0700 Subject: [PATCH 40/58] Make the hybrid pipeline a direct BasePipeline subclass `DNAscopeHybridPipeline` inherited from both `DNAscopePipeline` and `DNAscopeLRPipeline` to reuse six methods. It now subclasses `BasePipeline`, declares every attribute its arguments need, and composes the short-read preprocessing, minimap2 realignment, mosdepth and LongReadSV stages through its own small wrappers. The bwt_max_mem computation becomes `util.set_bwt_max_mem`, shared with the dnascope pipeline, and the long-read `call_svs` loses the `replace_rg` parameter only the hybrid used. Two latent defects go away with the inheritance: `--lr_fastq_taglist` now reaches the minimap2 command (the inherited method read the long-read pipeline's `fastq_taglist`), and a hybrid built without the argument parser gets `minimap2_args="-Y"`, its documented default, instead of the long-read pipeline's `-YL`. A new test asserts, for every pipeline, that each argument is declared in `__init__` with its argparse default. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 28 +-- sentieon_cli/dnascope_hybrid.py | 177 +++++++++++++++++-- sentieon_cli/dnascope_longread.py | 7 +- sentieon_cli/util.py | 30 ++++ tests/unit/test_dnascope_hybrid_longread.py | 186 ++++++++++++++++++++ tests/unit/test_pipeline_params.py | 48 +++++ tests/unit/test_pipeline_validation.py | 39 ++-- 7 files changed, 455 insertions(+), 60 deletions(-) create mode 100644 tests/unit/test_dnascope_hybrid_longread.py create mode 100644 tests/unit/test_pipeline_params.py diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 3d16a5d..66a489f 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -45,6 +45,7 @@ parse_rg_line, path_arg, require_versions, + set_bwt_max_mem, split_alignment, total_memory, ) @@ -319,8 +320,10 @@ 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: @@ -350,27 +353,6 @@ 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") diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index fa767fe..bc25793 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -24,20 +24,32 @@ HybridStage2, HybridStage3, ) -from .dnascope import 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__, @@ -46,8 +58,10 @@ path_arg, require_versions, sample_sex_arg, + set_bwt_max_mem, split_alignment, vcf_id, + versions_available, ) from .shard import ( determine_shards_from_fai, @@ -126,7 +140,7 @@ def __init__( ) -class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): +class DNAscopeHybridPipeline(BasePipeline): """The DNAscope Hybrid pipeline""" params = copy.deepcopy(BasePipeline.params) @@ -319,24 +333,43 @@ class DNAscopeHybridPipeline(DNAscopePipeline, DNAscopeLRPipeline): 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 = "-Y" + 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 = ( + "--cram_write_options version=3.0,compressor=rans" + ) # Stashed by `build_dag` for the second, sex-aware DAG self.ploidy_json: Optional[pathlib.Path] = None self.cnv_sr_aln: List[pathlib.Path] = [] @@ -549,7 +582,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] @@ -610,14 +643,12 @@ def build_dag(self) -> DAG: 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() @@ -625,10 +656,13 @@ def build_dag(self) -> DAG: 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) @@ -664,6 +698,111 @@ def build_dag(self) -> DAG: 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: diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index 79e3a54..0102a77 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -1154,11 +1154,7 @@ def add_small_variant_calling( {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 """ @@ -1171,7 +1167,6 @@ def call_svs( inputs=sample_input, model=bundle.joinpath("longreadsv.model"), interval=self.bed, - replace_rg=replace_rg, ) .build() .job diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index 7475a65..88e11f8 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -235,6 +235,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 = [] diff --git a/tests/unit/test_dnascope_hybrid_longread.py b/tests/unit/test_dnascope_hybrid_longread.py new file mode 100644 index 0000000..a1ff914 --- /dev/null +++ b/tests/unit/test_dnascope_hybrid_longread.py @@ -0,0 +1,186 @@ +""" +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 `-Y`, not the long-read `-YL`""" + 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 " -Y " in cmd_str + assert "-YL" not in cmd_str + + 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_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..64fce04 100644 --- a/tests/unit/test_pipeline_validation.py +++ b/tests/unit/test_pipeline_validation.py @@ -16,6 +16,7 @@ from sentieon_cli.dnascope import DNAscopePipeline from sentieon_cli.dnascope_longread import DNAscopeLRPipeline from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline +from sentieon_cli.util import set_bwt_max_mem from tests.utils.test_helpers import create_mock_args @@ -431,21 +432,35 @@ def test_total_input_size_calculation(self): 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()""" From 4ccade937306b79b303da0445ed4fac7bce4e141 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 3 Sep 2026 22:26:13 -0700 Subject: [PATCH 41/58] Drop the attribute renames from validate() The dnascope and long-read pipelines renamed their input attributes during validation (`r1_fastq` to `sr_r1_fastq`, `sample_input` to `lr_aln`, and so on) so that the hybrid pipeline could inherit from both without the names colliding. The hybrid no longer inherits from either, so each pipeline now reads its own argument names throughout. Two validation tests that set the renamed attributes before `validate()` ran, and so exercised the wrong branch, now assert the messages they were written for. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope.py | 38 +++++--------- sentieon_cli/dnascope_longread.py | 13 ++--- tests/unit/test_dag_construction.py | 8 +-- tests/unit/test_dnascope_longread_pop_vcf.py | 1 - tests/unit/test_pipeline_validation.py | 53 +++++++++++++------- 5 files changed, 54 insertions(+), 59 deletions(-) diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 66a489f..8f0729a 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -245,20 +245,8 @@ def __init__(self) -> None: def validate(self) -> None: self.required(self.output_vcf, "output_vcf") - # 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 - # 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" @@ -285,21 +273,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: @@ -329,9 +317,7 @@ def configure_alignment(self) -> None: 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 @@ -424,7 +410,7 @@ def add_sr_input_alignment( inputs=self.sample_input, model_bundle=bundle, bam_format=self.bam_format, - duplicate_marking=self.sr_duplicate_marking, + duplicate_marking=self.duplicate_marking, input_ref=self.input_ref, collate=self.collate_align, bwa_args=self.bwa_args, @@ -438,7 +424,7 @@ def add_sr_fastq_alignment( """Align fastq files to the reference genome using bwa""" bundle = self.required(self.model_bundle, "model_bundle") - if not self.sr_r1_fastq and not self.sr_readgroups: + if not self.r1_fastq and not self.readgroups: return AlignResult( jobs=[], terminal=set(), outputs=[], cleanup_paths=[] ) @@ -447,13 +433,13 @@ def add_sr_fastq_alignment( return BwaFastqStage( ctx=ctx, - r1_fastq=self.sr_r1_fastq, - r2_fastq=self.sr_r2_fastq, - readgroups=self.sr_readgroups, + 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.sr_duplicate_marking, + duplicate_marking=self.duplicate_marking, unzip=find_unzip(self.logger), bwa_args=self.bwa_args, bwa_k=str(self.bwa_k), @@ -471,7 +457,7 @@ def add_sr_preprocessing( return ShortReadPreprocessingStage( ctx=ctx, inputs=sample_input, - duplicate_marking=self.sr_duplicate_marking, + duplicate_marking=self.duplicate_marking, consensus=self.consensus, skip_metrics=self.skip_metrics, assay=self.assay, diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index 0102a77..f7b4c44 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -253,13 +253,8 @@ def __init__(self) -> None: 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" @@ -433,7 +428,7 @@ def build_dag(self) -> 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() @@ -486,11 +481,11 @@ def lr_align_inputs(self) -> Tuple[List[pathlib.Path], Set[Job]]: result = Minimap2RealignStage( ctx=self.stage_context(), - inputs=self.lr_aln, + inputs=self.sample_input, model_bundle=bundle, sample_name=output_vcf.name.replace(".vcf.gz", ""), bam_format=self.bam_format, - input_ref=self.lr_input_ref, + input_ref=self.input_ref, fastq_taglist=self.fastq_taglist, minimap2_args=self.minimap2_args, util_sort_args=self.util_sort_args, diff --git a/tests/unit/test_dag_construction.py b/tests/unit/test_dag_construction.py index 1749f70..b7faa77 100644 --- a/tests/unit/test_dag_construction.py +++ b/tests/unit/test_dag_construction.py @@ -66,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 diff --git a/tests/unit/test_dnascope_longread_pop_vcf.py b/tests/unit/test_dnascope_longread_pop_vcf.py index 8376b7e..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 diff --git a/tests/unit/test_pipeline_validation.py b/tests/unit/test_pipeline_validation.py index 64fce04..f628345 100644 --- a/tests/unit/test_pipeline_validation.py +++ b/tests/unit/test_pipeline_validation.py @@ -54,8 +54,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: @@ -84,8 +84,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() @@ -97,8 +97,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() @@ -109,11 +109,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""" @@ -121,8 +129,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 @@ -198,7 +206,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"] @@ -272,22 +280,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 @@ -426,8 +441,8 @@ 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 From 329891252ec2fa8380f48c79e485928e0eadc5af Mon Sep 17 00:00:00 2001 From: Don Freed Date: Thu, 3 Sep 2026 22:32:17 -0700 Subject: [PATCH 42/58] Update the require_versions docstring The version constants no longer all live in the pipeline modules. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index 88e11f8..e635991 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -146,7 +146,7 @@ def require_versions( `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 pipelines' module-level + 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. """ From 55ff1386b568e7b3b4e0f0cdd775bb3e9d9bcb4a Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 4 Sep 2026 09:57:54 -0700 Subject: [PATCH 43/58] Default the hybrid pipeline's minimap2 arguments to -YL The hybrid pipeline's long-read input may be ONT, whose very long reads can carry more than 65535 CIGAR operations, which the BAM format cannot store; `-L` moves such CIGAR strings into the `CG` tag. The hidden `--minimap2_args` default therefore becomes `-YL`, matching the long-read pipeline and the minimap2 stages, instead of `-Y`. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope_hybrid.py | 4 ++-- tests/unit/test_dnascope_hybrid_longread.py | 17 ++++++++++++++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index bc25793..ab4dfd9 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -306,7 +306,7 @@ class DNAscopeHybridPipeline(BasePipeline): "minimap2_args": { # help="Extra arguments for sentieon minimap2", "help": argparse.SUPPRESS, - "default": "-Y", + "default": "-YL", }, "no_split_alignment": { "help": argparse.SUPPRESS, @@ -362,7 +362,7 @@ def __init__(self) -> None: self.bwt_max_mem: Optional[str] = None self.lr_fastq_taglist = "*" self.lr_read_filter: Optional[str] = None - self.minimap2_args = "-Y" + self.minimap2_args = "-YL" self.no_split_alignment = False self.skip_model_apply = False self.skip_pop_vcf_id_check = False diff --git a/tests/unit/test_dnascope_hybrid_longread.py b/tests/unit/test_dnascope_hybrid_longread.py index a1ff914..7807977 100644 --- a/tests/unit/test_dnascope_hybrid_longread.py +++ b/tests/unit/test_dnascope_hybrid_longread.py @@ -99,14 +99,25 @@ def dep_names(self, dag, job): return {dep.name for dep in dag.waiting_jobs.get(job, set())} def test_minimap2_args_default(self): - """The hybrid pipeline aligns with `-Y`, not the long-read `-YL`""" + """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 " -Y " in cmd_str - assert "-YL" not in cmd_str + 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""" From c1066a741eff33ef68cbcbd85c8488d7ccaf3591 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 4 Sep 2026 10:38:43 -0700 Subject: [PATCH 44/58] Run the pangenome pipeline through BasePipeline.main `SentieonPangenome` overrode `main()` to parse the reference index and population VCF before validation, call `build_first_dag()` instead of `build_dag()`, and gate the second, sex-aware DAG itself. It now uses the base flow like every other pipeline: `validate()` parses the index, shards and VCF contigs (as the hybrid pangenome already did), the main DAG builder is `build_dag()`, and `build_second_dag()` returns `None` when no sex-aware caller is enabled and calls `get_sex` itself, as the hybrid DNAscope pipeline does. The reference index is validated once rather than twice. A new test asserts that no pipeline overrides `main()`. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/sentieon_pangenome.py | 75 ++++++------------- tests/unit/test_pipeline_lifecycle.py | 40 +++++++--- tests/unit/test_sentieon_pangenome.py | 67 +++++++---------- ...test_sentieon_pangenome_kmc_integration.py | 13 ++-- 4 files changed, 85 insertions(+), 110 deletions(-) diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 1420360..a41439e 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -6,9 +6,7 @@ import copy import json import pathlib -import shutil import sys -import time from typing import Dict, List, Optional, Set, Tuple, Union import packaging.version @@ -53,7 +51,6 @@ path_arg, require_versions, sample_sex_arg, - tmp, total_memory, vcf_id, ) @@ -269,7 +266,6 @@ def __init__(self) -> None: self.call_svs = False self.gvcf = False self.par_bed: Optional[pathlib.Path] = None - self.cnv_input_bams: List[pathlib.Path] = [] self.pangenome_ref_name = "GRCh38" self.extract_model_name = "extract.model" self.pangenome_contig_prefix = "GRCh38#0#" @@ -288,6 +284,10 @@ def __init__(self) -> None: self.skip_pop_vcf_id_check: bool = False self.skip_model_apply = False self.skip_small_variants = False + # Stashed by `build_dag` for the second, sex-aware DAG + self.ploidy_json: Optional[pathlib.Path] = None + self.sr_alignment: Optional[pathlib.Path] = None + self.cnv_input_bams: List[pathlib.Path] = [] def _cnv_in_second_dag(self) -> bool: """CNV calling runs in the second, sex-aware DAG""" @@ -301,55 +301,21 @@ def _needs_second_dag(self) -> bool: or self._cnv_in_second_dag() ) - def main(self, args: argparse.Namespace) -> None: - """Run the pipeline""" - self.handle_arguments(args) - self.setup_logging(args) - start_time = time.monotonic() - success = False - try: - self.validate_ref() - - 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) - - self.validate() - self.start_run_logs() - self.shards = determine_shards_from_fai( - self.fai_data, 10 * 1000 * 1000 - ) - - tmp_dir_str = tmp() - self.tmp_dir = pathlib.Path(tmp_dir_str) - - try: - dag = self.build_first_dag() - executor = self.run(dag) - self.check_execution(dag, executor) - - if self._needs_second_dag(): - self.get_sex(self.ploidy_json) - dag = self.build_second_dag() - executor = self.run(dag) - self.check_execution(dag, executor) - finally: - if not self.retain_tmpdir: - shutil.rmtree(tmp_dir_str) - success = True - finally: - self.log_completion(success, start_time) - 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 + ) + 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) + 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: @@ -658,10 +624,7 @@ def configure(self) -> None: pass def build_dag(self) -> DAG: - return DAG() - - def build_first_dag(self) -> DAG: - """Build the main DAG for the Sentieon pangenome pipeline""" + """Build the first DAG for the Sentieon pangenome pipeline""" bundle = self.required(self.model_bundle, "model_bundle") ctx = self.stage_context() @@ -1185,8 +1148,14 @@ def build_segdup_job( task_name="segdup", ) - 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() diff --git a/tests/unit/test_pipeline_lifecycle.py b/tests/unit/test_pipeline_lifecycle.py index 9e84888..adf67e9 100644 --- a/tests/unit/test_pipeline_lifecycle.py +++ b/tests/unit/test_pipeline_lifecycle.py @@ -7,6 +7,7 @@ import logging import os import sys +from typing import Optional import pytest @@ -16,9 +17,14 @@ os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")), ) -from sentieon_cli import sentieon_pangenome # noqa: E402 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 @@ -231,21 +237,20 @@ def __init__(self, fail_on=None): self.check_calls = 0 self.dags_built = [] - def validate_ref(self) -> None: - pass - def validate(self) -> None: pass def configure(self) -> None: pass - def build_first_dag(self) -> DAG: + 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) -> DAG: + def build_second_dag(self) -> Optional[DAG]: + if not self._needs_second_dag(): + return None self.dags_built.append("second") return DAG() @@ -264,18 +269,14 @@ def check_execution(self, dag, executor): def _stub_pangenome(monkeypatch, tmp_path, **kwargs): """A stubbed pangenome pipeline writing temp dirs into `tmp_path`""" monkeypatch.setenv("SENTIEON_TMPDIR", str(tmp_path)) - monkeypatch.setattr(sentieon_pangenome, "parse_fai", lambda fai: {}) - monkeypatch.setattr( - sentieon_pangenome, "determine_shards_from_fai", lambda fai, size: [] - ) return _StubPangenome(**kwargs) def test_pangenome_main_cleans_up_tmpdir_when_the_first_dag_fails( tmp_path, monkeypatch ): - # SentieonPangenome overrides main() for its two-DAG flow; the temp - # directory has to be removed on the failure path too. + # 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): @@ -344,3 +345,18 @@ def test_check_execution_flags_unexecuted_jobs(): 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_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index 5d17bb3..af96db3 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -96,7 +96,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 +117,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 @@ -137,7 +137,7 @@ 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_first_dag() + 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") @@ -149,7 +149,7 @@ 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") @@ -159,7 +159,7 @@ def test_model_apply_default(self): """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) @@ -180,7 +180,7 @@ 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) @@ -221,7 +221,7 @@ def test_gvcf_adds_gvcftyper_and_routes_outputs(self): 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] @@ -277,7 +277,7 @@ 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_first_dag() + dag = pipeline.build_dag() job_names, all_jobs = self._get_all_job_names(dag) assert job_names.count("sad-lad-update") == 1 @@ -303,7 +303,7 @@ 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_first_dag() + dag = pipeline.build_dag() job_names, _ = self._get_all_job_names(dag) assert "sad-lad-update" not in job_names @@ -311,7 +311,7 @@ def test_no_ad_update_with_skip_small_variants(self): 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] @@ -342,7 +342,7 @@ 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) @@ -375,7 +375,7 @@ def test_call_svs(self): 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() @@ -394,7 +394,7 @@ 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) @@ -411,7 +411,7 @@ 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) @@ -451,7 +451,7 @@ def _build_cnv_dag(self, pipeline, sample_sex=SampleSex.FEMALE): The first DAG has to be built first so that the CNV input alignments are stashed for the second DAG. """ - pipeline.build_first_dag() + pipeline.build_dag() pipeline.sample_sex = sample_sex return pipeline.build_second_dag() @@ -505,7 +505,7 @@ def test_cnv_jobs_with_call_svs(self): """CNV jobs run in the second DAG when --call_svs is enabled""" pipeline = self.create_pipeline() pipeline.call_svs = True - first_dag = pipeline.build_first_dag() + 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) @@ -525,7 +525,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 @@ -533,13 +533,9 @@ def test_cnv_jobs_not_present_without_call_svs(self): assert "indel2cnv" not in job_names assert "combine-cnv" not in job_names - # Nor to the second DAG + # And no second DAG at all, as nothing consumes the sample sex pipeline.sample_sex = SampleSex.FEMALE - second_job_names, _ = self._get_all_job_names( - pipeline.build_second_dag() - ) - assert "cnvscope" not in second_job_names - assert "combine-cnv" not in second_job_names + assert pipeline.build_second_dag() is None def test_second_dag_needed_for_cnv_only(self): """CNV calling alone triggers the second DAG""" @@ -562,7 +558,7 @@ def test_ploidy_estimation_runs_with_sample_sex(self): pipeline.call_svs = True pipeline.sample_sex = SampleSex.FEMALE - job_names, _ = self._get_all_job_names(pipeline.build_first_dag()) + 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): @@ -571,7 +567,7 @@ def test_ploidy_estimation_runs_without_sex_aware_callers(self): assert pipeline._needs_second_dag() is False job_names, all_jobs = self._get_all_job_names( - pipeline.build_first_dag() + pipeline.build_dag() ) assert "estimate-ploidy" in job_names @@ -586,7 +582,7 @@ def test_ploidy_estimation_runs_with_skip_small_variants(self): pipeline = self.create_pipeline() pipeline.skip_small_variants = True - job_names, _ = self._get_all_job_names(pipeline.build_first_dag()) + 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): @@ -607,7 +603,7 @@ 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) @@ -623,13 +619,9 @@ def test_call_svs_without_cnv_model(self): assert "indel2cnv" not in job_names assert "combine-cnv" not in job_names - # And they are not in the second DAG either + # And there is no second DAG either pipeline.sample_sex = SampleSex.FEMALE - second_job_names, _ = self._get_all_job_names( - pipeline.build_second_dag() - ) - assert "cnvscope" not in second_job_names - assert "combine-cnv" not in second_job_names + 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""" @@ -637,13 +629,8 @@ def test_call_svs_without_cnv_model_skip_small_variants(self): pipeline.skip_small_variants = True pipeline.call_svs = True pipeline.has_cnv_model = False - dag = self._build_cnv_dag(pipeline) - job_names, _ = self._get_all_job_names(dag) - 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""" @@ -746,7 +733,7 @@ def test_cnv_with_skip_model_apply(self): pipeline.call_svs = True pipeline.skip_model_apply = True first_job_names, _ = self._get_all_job_names( - pipeline.build_first_dag() + pipeline.build_dag() ) assert "model-apply" not in first_job_names diff --git a/tests/unit/test_sentieon_pangenome_kmc_integration.py b/tests/unit/test_sentieon_pangenome_kmc_integration.py index 9c0402f..eaad07e 100644 --- a/tests/unit/test_sentieon_pangenome_kmc_integration.py +++ b/tests/unit/test_sentieon_pangenome_kmc_integration.py @@ -21,10 +21,9 @@ 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 = {} - + # Mock dependencies + pipeline.logger = MagicMock() pipeline.handle_arguments = MagicMock() pipeline.setup_logging = MagicMock() pipeline.validate_bundle = MagicMock() @@ -34,6 +33,9 @@ def test_validate_kmc_patch_success(self): pipeline.collect_readgroups = MagicMock() 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.vcf_contigs", return_value={}), \ patch("sentieon_cli.sentieon_pangenome.check_kmc_patch", return_value=True) as mock_check, \ patch("sys.exit") as mock_exit: @@ -54,12 +56,13 @@ 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.logger = MagicMock() 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.vcf_contigs", return_value={}), \ patch("sentieon_cli.sentieon_pangenome.check_kmc_patch", return_value=False) as mock_check, \ patch("sys.exit") as mock_exit: From ac7c36ac37464ba27bd556c8156d46db595c6b61 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 4 Sep 2026 11:04:09 -0700 Subject: [PATCH 45/58] Collect the transfer inputs with TransferConfig.from_pipeline Four pipelines carried an identical `transfer_config()` helper building a `TransferConfig` from `pop_vcf`, `shards`, `pop_vcf_contigs` and `fai_data`. A `TransferInputs` protocol names those four attributes and `TransferConfig.from_pipeline` reads them structurally, raising a clear error when there is no population VCF. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope_hybrid.py | 16 +------------ sentieon_cli/dnascope_longread.py | 20 +++-------------- sentieon_cli/hybrid_pangenome.py | 17 +++----------- sentieon_cli/sentieon_pangenome.py | 16 +------------ sentieon_cli/stages/transfer.py | 34 +++++++++++++++++++++++++++- tests/unit/test_stage_transfer.py | 36 ++++++++++++++++++++++++++++++ 6 files changed, 77 insertions(+), 62 deletions(-) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index ab4dfd9..d69090e 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -830,20 +830,6 @@ def build_second_dag(self) -> Optional[DAG]: ).add_to(dag) return dag - def transfer_config(self) -> TransferConfig: - """The inputs the annotation-transfer stage needs. - - Only valid once `--pop_vcf` is known to be set; the calling code - checks it first. - """ - assert self.pop_vcf is not None - return TransferConfig( - pop_vcf=self.pop_vcf, - shards=self.shards, - pop_vcf_contigs=self.pop_vcf_contigs, - fai_data=self.fai_data, - ) - def add_cleanup( self, dag: DAG, @@ -1218,7 +1204,7 @@ def add_variant_calling( if self.skip_model_apply: transfer_target = ctx.output_vcf transfer_spec = TransferSpec( - self.transfer_config(), transfer_target + TransferConfig.from_pipeline(self), transfer_target ) apply_spec: Optional[ApplySpec] = None diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index f7b4c44..d6f081e 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -643,20 +643,6 @@ def hificnv( ) return hificnv_job - def transfer_config(self) -> TransferConfig: - """The inputs the annotation-transfer stage needs. - - Only valid once `--pop_vcf` is known to be set; the calling code - checks it first. - """ - assert self.pop_vcf is not None - return TransferConfig( - pop_vcf=self.pop_vcf, - shards=self.shards, - pop_vcf_contigs=self.pop_vcf_contigs, - fai_data=self.fai_data, - ) - def add_small_variant_calling( self, dag: DAG, @@ -709,7 +695,7 @@ def add_small_variant_calling( diploid_transfer: Optional[TransferSpec] = None if self.pop_vcf: diploid_transfer = TransferSpec( - config=self.transfer_config(), + config=TransferConfig.from_pipeline(self), out_vcf=self.tmp_dir.joinpath("out_diploid_tmp_anno.vcf.gz"), tag="diploid", ) @@ -925,7 +911,7 @@ def add_small_variant_calling( patch_transfer: Optional[TransferSpec] = None if self.pop_vcf: patch_transfer = TransferSpec( - config=self.transfer_config(), + config=TransferConfig.from_pipeline(self), out_vcf=self.tmp_dir.joinpath( f"out_hap{i}_patch_anno.vcf.gz" ), @@ -988,7 +974,7 @@ def add_small_variant_calling( unphased_transfer: Optional[TransferSpec] = None if self.pop_vcf: unphased_transfer = TransferSpec( - config=self.transfer_config(), + config=TransferConfig.from_pipeline(self), out_vcf=self.tmp_dir.joinpath( "out_diploid_unphased_patch_anno.vcf.gz" ), diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index f6b09cb..259a026 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -839,7 +839,9 @@ def build_dag(self) -> DAG: TransferApplyStage( ctx=ctx, raw_vcf=raw_vcf, - transfer=TransferSpec(self.transfer_config(), transfer_target), + transfer=TransferSpec( + TransferConfig.from_pipeline(self), transfer_target + ), apply=( ApplySpec(model, ctx.output_vcf) if not self.skip_model_apply @@ -849,19 +851,6 @@ def build_dag(self) -> DAG: return dag - def transfer_config(self) -> TransferConfig: - """The inputs the annotation-transfer stage needs. - - `build_dag` requires `--pop_vcf` before it calls this. - """ - assert self.pop_vcf is not None - return TransferConfig( - pop_vcf=self.pop_vcf, - shards=self.shards, - pop_vcf_contigs=self.pop_vcf_contigs, - fai_data=self.fai_data, - ) - def lr_kmc_pairs(self) -> List[Tuple[pathlib.Path, pathlib.Path]]: """`(alignment, decode reference)` pairs for the long-read k-mer counting. diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index a41439e..044247a 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -909,7 +909,7 @@ def build_dag(self) -> DAG: transfer: Optional[TransferSpec] = None if self.pop_vcf: transfer = TransferSpec( - config=self.transfer_config(), + config=TransferConfig.from_pipeline(self), out_vcf=( transfer_vcf if not self.skip_model_apply @@ -945,20 +945,6 @@ def build_dag(self) -> DAG: return dag - def transfer_config(self) -> TransferConfig: - """The inputs the annotation-transfer stage needs. - - Only valid once `--pop_vcf` is known to be set; `validate_bundle` - requires it. - """ - assert self.pop_vcf is not None - return TransferConfig( - pop_vcf=self.pop_vcf, - shards=self.shards, - pop_vcf_contigs=self.pop_vcf_contigs, - fai_data=self.fai_data, - ) - def bwa_extract_stage( self, ctx: StageContext, diff --git a/sentieon_cli/stages/transfer.py b/sentieon_cli/stages/transfer.py index fe9a0cd..b6b9fff 100644 --- a/sentieon_cli/stages/transfer.py +++ b/sentieon_cli/stages/transfer.py @@ -7,7 +7,7 @@ import re import subprocess as sp import tempfile -from typing import Dict, Iterable, List, Optional, Set, Tuple +from typing import Dict, Iterable, List, Optional, Protocol, Set, Tuple from importlib.resources import files @@ -155,6 +155,20 @@ def build_transfer_jobs( 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. @@ -169,6 +183,24 @@ class TransferConfig: 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): diff --git a/tests/unit/test_stage_transfer.py b/tests/unit/test_stage_transfer.py index ce3d807..d9c03f7 100644 --- a/tests/unit/test_stage_transfer.py +++ b/tests/unit/test_stage_transfer.py @@ -3,6 +3,9 @@ """ import pathlib +from types import SimpleNamespace + +import pytest from sentieon_cli.dag import DAG from sentieon_cli.shard import Shard @@ -131,3 +134,36 @@ def test_result_fields(self, tmp_path): 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) From f8f7280df3ae26febff3ad1fa3b8a0db0b9f32a8 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Fri, 4 Sep 2026 11:06:53 -0700 Subject: [PATCH 46/58] Let the hybrid pangenome pipeline use LongReadSVStage `LongReadSVStage` gains an explicit `output` (the derived `.sv.vcf.gz` path stays the default) and the LongReadSV algo options (`min_map_qual`, `min_sv_size`, `min_dp`, `min_af`), which is what the hybrid pangenome's graph-update SV call needed. Its hand-built `build_longreadsv_job` is gone; the job keeps its `longreadsv` name under the `pangenome-update` task and its command is unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/hybrid_pangenome.py | 41 +++++++------------------ sentieon_cli/stages/sv.py | 30 +++++++++++++++---- tests/unit/test_stage_sv.py | 51 ++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 36 deletions(-) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 259a026..6649837 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -21,7 +21,6 @@ from .driver import ( DNAscope, Driver, - LongReadSV, PangenomeSV, PGHapUpdateAlgo, ) @@ -48,6 +47,7 @@ TransferApplyStage, TransferSpec, ) +from .stages.sv import LongReadSVStage from .stages.transfer import TransferConfig from .util import ( __version__, @@ -703,15 +703,22 @@ def build_dag(self) -> DAG: dag.add_job(update_raw_job, {gfa_job} | realign_jobs) # Call SVs from the long reads and collect graph update regions - longreadsv_job = self.build_longreadsv_job(longread_sv_vcf, calling_lr) - dag.add_job(longreadsv_job, realign_jobs) + 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_job}) + dag.add_job(sv_bed_job, {longreadsv_result.job}) # Graph update with the SV BED update_job = self.build_graph_update_job( @@ -1040,32 +1047,6 @@ def build_graph_update_job( task_name="pangenome-update", ) - def build_longreadsv_job( - self, - out_vcf: pathlib.Path, - lr_aln: List[pathlib.Path], - ) -> Job: - """Call SVs from the long reads for the graph update""" - assert self.model_bundle is not None - driver = Driver( - reference=self.reference, - thread_count=self.cores, - input=list(lr_aln), - ) - driver.add_algo( - LongReadSV( - out_vcf, - model=self.model_bundle.joinpath("longreadsv.model"), - min_sv_size=LONGREADSV_MIN_SV_SIZE, - ) - ) - return Job( - Pipeline(Command(*driver.build_cmd())), - "longreadsv", - self.cores, - task_name="pangenome-update", - ) - def build_minimap2_lift_job( self, out_bam: pathlib.Path, diff --git a/sentieon_cli/stages/sv.py b/sentieon_cli/stages/sv.py index 03e0eef..c722172 100644 --- a/sentieon_cli/stages/sv.py +++ b/sentieon_cli/stages/sv.py @@ -30,12 +30,19 @@ class LongReadSVResult(StageResult): class LongReadSVStage(Stage): """Call SVs from long reads with Sentieon LongReadSV. - The output VCF is derived from the run's output VCF, so the caller - only supplies the alignment and the model. + ``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" @@ -43,12 +50,23 @@ class LongReadSVStage(Stage): def build(self) -> LongReadSVResult: """Construct this stage's job without touching a DAG""" - sv_vcf = pathlib.Path( - str(self.ctx.output_vcf).replace(".vcf.gz", ".sv.vcf.gz") - ) + 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, self.model)], + [ + 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, diff --git a/tests/unit/test_stage_sv.py b/tests/unit/test_stage_sv.py index c7f86aa..388ce81 100644 --- a/tests/unit/test_stage_sv.py +++ b/tests/unit/test_stage_sv.py @@ -87,6 +87,49 @@ def test_without_replace_rg(self, tmp_path): 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""" @@ -96,6 +139,14 @@ def test_sv_vcf(self, tmp_path): 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""" From a494662b84085c9f425ee80516048af0c8321958 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 5 Sep 2026 11:17:46 -0700 Subject: [PATCH 47/58] Skip CNV calling in the hybrid smoke tests The smoke reference is not a recognised genome build, so no PAR BED resolves and CNV calling now fails validation up front. The smoke runs exercise alignment, duplicate marking and variant calling, not CNVscope, so they skip it. Co-Authored-By: Claude Fable 5.1 --- .github/workflows/main.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1cee559..0a9ddac 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -308,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 @@ -326,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 @@ -343,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 From caa216fa7542129cea0bf5f0af396b5731296858 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 5 Sep 2026 13:36:22 -0700 Subject: [PATCH 48/58] Write CRAM output with the default util sort arguments Drop DEFAULT_UTIL_SORT_ARGS and the matching "--cram_write_options version=3.0,compressor=rans" defaults from the alignment stages, the command builders, and the hidden --util_sort_args argument of the dnascope, dnascope-longread and dnascope-hybrid pipelines, and stop passing the same option to Dedup. The default pipeline now writes CRAM with the tools' own defaults; the intermediate BAMs written before duplicate marking keep --bam_compression 1. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/command_strings.py | 8 ++++---- sentieon_cli/dnascope.py | 6 ++---- sentieon_cli/dnascope_hybrid.py | 6 ++---- sentieon_cli/dnascope_longread.py | 6 ++---- sentieon_cli/stages/alignment.py | 11 ++++------- sentieon_cli/stages/preprocessing.py | 1 - tests/unit/test_stage_alignment.py | 18 +++++++++++++++++- tests/unit/test_stage_preprocessing.py | 7 +++++-- 8 files changed, 36 insertions(+), 27 deletions(-) diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 4ed974f..2282594 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -536,7 +536,7 @@ 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""" @@ -625,7 +625,7 @@ 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 @@ -726,7 +726,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""" @@ -774,7 +774,7 @@ 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: diff --git a/sentieon_cli/dnascope.py b/sentieon_cli/dnascope.py index 8f0729a..c0c6006 100644 --- a/sentieon_cli/dnascope.py +++ b/sentieon_cli/dnascope.py @@ -205,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": "", }, } ) @@ -235,9 +235,7 @@ 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 diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index d69090e..5f907c3 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -326,7 +326,7 @@ class DNAscopeHybridPipeline(BasePipeline): "util_sort_args": { # help="Extra arguments for sentieon util sort", "help": argparse.SUPPRESS, - "default": "--cram_write_options version=3.0,compressor=rans", + "default": "", }, } ) @@ -367,9 +367,7 @@ def __init__(self) -> None: self.skip_model_apply = False self.skip_pop_vcf_id_check = False self.sr_read_filter: Optional[str] = None - self.util_sort_args = ( - "--cram_write_options version=3.0,compressor=rans" - ) + 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] = [] diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index d6f081e..0bd9eaf 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -218,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": "", }, } ) @@ -245,9 +245,7 @@ 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 diff --git a/sentieon_cli/stages/alignment.py b/sentieon_cli/stages/alignment.py index ae6769a..32736df 100644 --- a/sentieon_cli/stages/alignment.py +++ b/sentieon_cli/stages/alignment.py @@ -23,9 +23,6 @@ from ..job import Job from .base import Stage, StageResult, driver_job -# The `sentieon util sort` arguments every aligner defaults to -DEFAULT_UTIL_SORT_ARGS = "--cram_write_options version=3.0,compressor=rans" - BWA_REALIGN_MIN_VERSIONS = { "sentieon driver": packaging.version.Version("202308"), "samtools": packaging.version.Version("1.16"), @@ -122,7 +119,7 @@ class BwaRealignStage(AlignmentStage): collate: bool = False bwa_args: str = "" bwa_k: str = "20000000" - util_sort_args: str = DEFAULT_UTIL_SORT_ARGS + util_sort_args: str = "" def build(self) -> AlignResult: suffix = aln_suffix(self.bam_format) @@ -201,7 +198,7 @@ class BwaFastqStage(AlignmentStage): unzip: str = "gzip" bwa_args: str = "" bwa_k: str = "20000000" - util_sort_args: str = DEFAULT_UTIL_SORT_ARGS + util_sort_args: str = "" def build(self) -> AlignResult: n_alignment_jobs = max(1, len(self.numa_nodes)) @@ -285,7 +282,7 @@ class Minimap2RealignStage(AlignmentStage): input_ref: Optional[pathlib.Path] = None fastq_taglist: str = "*" minimap2_args: str = "-YL" - util_sort_args: str = DEFAULT_UTIL_SORT_ARGS + util_sort_args: str = "" minimap2_model: Optional[Union[pathlib.Path, str]] = None def build(self) -> AlignResult: @@ -340,7 +337,7 @@ class Minimap2FastqStage(AlignmentStage): bam_format: bool = False unzip: str = "gzip" minimap2_args: str = "-YL" - util_sort_args: str = DEFAULT_UTIL_SORT_ARGS + util_sort_args: str = "" def build(self) -> AlignResult: suffix = aln_suffix(self.bam_format) diff --git a/sentieon_cli/stages/preprocessing.py b/sentieon_cli/stages/preprocessing.py index f2a29bd..fac5d8b 100644 --- a/sentieon_cli/stages/preprocessing.py +++ b/sentieon_cli/stages/preprocessing.py @@ -133,7 +133,6 @@ def add_to( score_file=paths.score, consensus=self.consensus, rmdup=(self.duplicate_marking == "rmdup"), - cram_write_options="version=3.0,compressor=rans", dedup_metrics=paths.dedup_metrics, lc_extra_algos=lc_extra_algos, ).add_to(dag, upstream) diff --git a/tests/unit/test_stage_alignment.py b/tests/unit/test_stage_alignment.py index e4b7df2..6e4829e 100644 --- a/tests/unit/test_stage_alignment.py +++ b/tests/unit/test_stage_alignment.py @@ -122,7 +122,20 @@ def test_without_duplicate_marking_writes_cram_output( tmp_path / "output_bwa_sorted_0.cram", tmp_path / "output_bwa_sorted_1.cram", ] - assert "--bam_compression 1" not in str(result.jobs[0].shell) + # 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( @@ -257,6 +270,7 @@ def test_final_outputs_have_no_cleanup_paths(self, tmp_path): 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: @@ -280,6 +294,7 @@ def test_job_metadata_and_outputs(self, tmp_path, rg_lines): 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() @@ -323,6 +338,7 @@ def test_one_job_per_fastq(self, tmp_path): 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): diff --git a/tests/unit/test_stage_preprocessing.py b/tests/unit/test_stage_preprocessing.py index 6ccf8b7..efd24d7 100644 --- a/tests/unit/test_stage_preprocessing.py +++ b/tests/unit/test_stage_preprocessing.py @@ -182,13 +182,16 @@ def test_locuscollector_carries_the_pre_dedup_metrics(self, tmp_path): ["LocusCollector"] + PRE_DEDUP_ALGOS + ["GCBias"] ) - def test_dedup_writes_metrics_and_cram_options(self, tmp_path): + 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 version=3.0,compressor=rans" in shell + assert "--cram_write_options" not in shell assert "--rmdup" not in shell def test_post_dedup_metrics(self, tmp_path): From f1924dfcd486c30a006b5a8f6c7ebc083bf6a559 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 5 Sep 2026 13:36:22 -0700 Subject: [PATCH 49/58] Pass --dump_probes to CNVscope CNVscope can write its per-probe debugging output with --dump_probes. Add the argument to the CNVscope algo and CNVscopeStage, and have both sex-aware callers supply it: .cnv.probes in dnascope-hybrid and _cnv.probes in dnascope-pangenome, next to each CNV VCF. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/dnascope_hybrid.py | 3 +++ sentieon_cli/driver.py | 2 ++ sentieon_cli/sentieon_pangenome.py | 4 ++++ sentieon_cli/stages/cnv.py | 5 ++++- tests/unit/test_cnv_sex.py | 22 ++++++++++++++++++++++ tests/unit/test_dnascope_hybrid_cnv.py | 10 ++++++++++ tests/unit/test_sentieon_pangenome.py | 3 +++ tests/unit/test_stage_cnv.py | 7 +++++++ 8 files changed, 55 insertions(+), 1 deletion(-) diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index 5f907c3..d8e1f38 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -823,6 +823,9 @@ def build_second_dag(self) -> Optional[DAG]: ), 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) diff --git a/sentieon_cli/driver.py b/sentieon_cli/driver.py index e442a89..53b6131 100644 --- a/sentieon_cli/driver.py +++ b/sentieon_cli/driver.py @@ -500,11 +500,13 @@ def __init__( *, 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): diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 044247a..74d7f75 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -1260,6 +1260,9 @@ def _add_cnv_jobs( cnv_vcf = pathlib.Path( str(ctx.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") ) + cnv_probes = pathlib.Path( + str(ctx.output_vcf).replace(".vcf.gz", "_cnv.probes") + ) # CNVscope on BWA deduped BAM, then CNVModelApply. The alignment # and the PangenomeSV output were both written by the first DAG. @@ -1271,6 +1274,7 @@ def _add_cnv_jobs( cnv_vcf=cnv_apply_vcf, sample_sex=self.sample_sex, par_bed=self.cnv_par_bed, + dump_probes=cnv_probes, interval=self.bed, ).add_to(dag) diff --git a/sentieon_cli/stages/cnv.py b/sentieon_cli/stages/cnv.py index 91904ad..d8ff55b 100644 --- a/sentieon_cli/stages/cnv.py +++ b/sentieon_cli/stages/cnv.py @@ -35,7 +35,8 @@ class CNVscopeStage(Stage): 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. + pipelines, which run them during validation. ``dump_probes`` is the + path CNVscope writes its per-probe debugging output to. """ inputs: List[pathlib.Path] @@ -44,6 +45,7 @@ class CNVscopeStage(Stage): 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" @@ -62,6 +64,7 @@ def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> CNVResult: self.model, sex=sex, par=par, + dump_probes=self.dump_probes, ) ], inputs=self.inputs, diff --git a/tests/unit/test_cnv_sex.py b/tests/unit/test_cnv_sex.py index 39efb2c..27342ef 100644 --- a/tests/unit/test_cnv_sex.py +++ b/tests/unit/test_cnv_sex.py @@ -239,9 +239,31 @@ def test_cnvscope_algo_omits_the_sex_arguments_by_default(): 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 diff --git a/tests/unit/test_dnascope_hybrid_cnv.py b/tests/unit/test_dnascope_hybrid_cnv.py index b77903d..99287e4 100644 --- a/tests/unit/test_dnascope_hybrid_cnv.py +++ b/tests/unit/test_dnascope_hybrid_cnv.py @@ -173,6 +173,16 @@ def test_a_male_sample_is_called_with_the_par_bed(self): 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() diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index af96db3..a619f96 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -651,6 +651,9 @@ def test_cnvscope_command(self): # 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 diff --git a/tests/unit/test_stage_cnv.py b/tests/unit/test_stage_cnv.py index 3fbf962..5b97761 100644 --- a/tests/unit/test_stage_cnv.py +++ b/tests/unit/test_stage_cnv.py @@ -117,6 +117,13 @@ def test_replace_rg_is_passed_to_cnvscope_only(self, tmp_path): ) 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, From aa31de489cbb1105ac66940f0d2317a376ee2516 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 09:15:51 -0700 Subject: [PATCH 50/58] Extract the pangenome FASTA with vg paths on non-GRCh38 references `sentieon pgutil gfa2fa` in driver release 202503.04 handles only GRCh38 pangenomes; on a CHM13 graph it fails with "Failed to parse ref CHM13#0#chr1 in S line". Have hybrid-pangenome choose the tool for the gfa2fa job in validate(): when the reference is not GRCh38 (per detect_reference_build) and the driver is 202503.04 or older, extract the graph's reference paths with `vg paths -x -Q -F`, which is slower but works with every release; otherwise keep pgutil. An unknown driver version (--skip_version_check, or a failed probe) also takes the vg path for a non-GRCh38 reference, with the reason logged at INFO. The job keeps its name, threads, task and dependencies. Add util.executable_version, which returns a tool's parsed version, and build check_version on it; unparseable --version output now logs an error instead of raising. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/command_strings.py | 14 +++ sentieon_cli/hybrid_pangenome.py | 85 +++++++++++++++- sentieon_cli/util.py | 50 ++++++--- tests/unit/test_hybrid_pangenome.py | 98 ++++++++++++++++++ tests/unit/test_util_versions.py | 151 ++++++++++++++++++++++++++++ 5 files changed, 378 insertions(+), 20 deletions(-) create mode 100644 tests/unit/test_util_versions.py diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 2282594..0a0d132 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1247,6 +1247,20 @@ def cmd_vg_paths_fasta( return Pipeline(Command(*cmd), file_output=output_fasta) +def cmd_vg_paths_ref_fasta( + output_fasta: pathlib.Path, + gfa_file: pathlib.Path, + ref_name: str, +) -> Pipeline: + """Extract a graph's reference paths 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), "-Q", ref_name, "-F"] + return Pipeline(Command(*cmd), file_output=output_fasta) + + def cmd_minimap2_lift( out_bam: pathlib.Path, sample_fasta: pathlib.Path, diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 6649837..4f5f2f9 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -28,6 +28,7 @@ from .logging import get_logger from .shard import ( GRCH38_CONTIGS, + detect_reference_build, determine_shards_from_fai, parse_fai, vcf_contigs, @@ -52,6 +53,7 @@ from .util import ( __version__, check_kmc_patch, + executable_version, parse_rg_line, path_arg, require_versions, @@ -68,6 +70,11 @@ "bedtools": None, } +# The last driver release whose `pgutil gfa2fa` handles only GRCh38 +# pangenomes. Non-GRCh38 runs on it (and older) extract the graph's +# reference 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 @@ -207,6 +214,7 @@ def __init__(self) -> None: self.pangenome_ref_name = "GRCh38" 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] = [] @@ -327,6 +335,10 @@ def validate(self) -> None: ) sys.exit(2) + # After the version gate and the contig checks, so the tool choice + # is only logged for runs that proceed + self.resolve_gfa2fa_tool() + def validate_sr_inputs(self) -> None: """Validate the short-read input arguments. @@ -596,6 +608,49 @@ def validate_rg_ids( ) 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 @@ -731,11 +786,8 @@ def build_dag(self) -> DAG: dag.add_job(update_job, {update_raw_job, sv_bed_job}) # FASTA generation from the updated graph - gfa2fa_job = Job( - cmds.cmd_pgutil_gfa2fa(pangenome_fasta, ref_fai, pangenome_gfa), - "gfa2fa", - 0, - task_name="pangenome", + gfa2fa_job = self.build_gfa2fa_job( + pangenome_fasta, pangenome_gfa, ref_fai ) dag.add_job(gfa2fa_job, {update_job}) faidx_job = Job( @@ -1047,6 +1099,29 @@ def build_graph_update_job( 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_paths_ref_fasta( + pangenome_fasta, pangenome_gfa, self.pangenome_ref_name + ) + 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, diff --git a/sentieon_cli/util.py b/sentieon_cli/util.py index e635991..868dcd4 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -87,20 +87,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 = ( @@ -110,7 +103,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": @@ -122,7 +115,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 " diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index f3f25e1..be2dd3f 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -8,6 +8,7 @@ import tempfile from unittest.mock import MagicMock +import packaging.version import pytest # Add the parent directory to the path to import sentieon_cli @@ -16,6 +17,7 @@ ) 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 @@ -355,6 +357,102 @@ def test_gfa2fa_and_faidx(self): 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} -Q GRCh38 -F >'{fasta}'") + assert "pgutil gfa2fa" 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_gfa2fa_with_vg_paths_chm13(self): + """`vg paths` extracts the paths of the configured reference""" + pipeline = self.create_pipeline() + pipeline.gfa2fa_with_vg = True + pipeline.pangenome_ref_name = "CHM13" + dag = pipeline.build_dag() + _, all_jobs = self._get_all_job_names(dag) + + gfa2fa_cmd = str(self._get_job(all_jobs, "gfa2fa").shell) + assert "-Q CHM13" in gfa2fa_cmd + + 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""" 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 + ) From 819060c48c2edcda0d68af0ee508f631f467b967 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 10:28:31 -0700 Subject: [PATCH 51/58] Add segdup, ExpansionHunter, T1K and pangenome CNV stages Extract the optional pangenome callers into reusable stages so both pangenome pipelines can compose the same building blocks: `SegdupStage` (`stages/segdup.py`), `ExpansionHunterStage` (`stages/expansion.py`), `T1KStage` (`stages/t1k.py`, with the default HLA/KIR loci) and the composite `PangenomeCNVStage` in `stages/cnv.py` (CNVscope and CNVModelApply, indel2cnv.py on the PangenomeSV output, and combine_cnv.py over the two). `cmd_segdup_caller` gains optional `lr_alignments` / `lr_bundle` arguments for long-read calling; `cmd_pyexec_indel2cnv` and `cmd_pyexec_combine_cnv` move the two script invocations into `command_strings.py`. `util` gains `caller_sex_arg`, the shared `--sex male|female` mapping for ExpansionHunter and segdup-caller. No pipeline uses the new stages yet -- `sentieon_pangenome.py` keeps its inline builders, so the module-level version constants are duplicated for now -- and the DAG snapshots of all 88 configurations are identical. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/command_strings.py | 69 ++++++++++- sentieon_cli/stages/cnv.py | 108 ++++++++++++++++ sentieon_cli/stages/expansion.py | 71 +++++++++++ sentieon_cli/stages/segdup.py | 81 ++++++++++++ sentieon_cli/stages/t1k.py | 94 ++++++++++++++ sentieon_cli/util.py | 10 ++ tests/unit/test_cnv_sex.py | 14 +++ tests/unit/test_stage_cnv.py | 152 ++++++++++++++++++++++- tests/unit/test_stage_expansion.py | 130 +++++++++++++++++++ tests/unit/test_stage_segdup.py | 193 +++++++++++++++++++++++++++++ tests/unit/test_stage_t1k.py | 182 +++++++++++++++++++++++++++ 11 files changed, 1099 insertions(+), 5 deletions(-) create mode 100644 sentieon_cli/stages/expansion.py create mode 100644 sentieon_cli/stages/segdup.py create mode 100644 sentieon_cli/stages/t1k.py create mode 100644 tests/unit/test_stage_expansion.py create mode 100644 tests/unit/test_stage_segdup.py create mode 100644 tests/unit/test_stage_t1k.py diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 0a0d132..ae5a5de 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -319,6 +319,46 @@ def cmd_pyexec_sad_lad_update( return Pipeline(Command(*cmd)) +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, @@ -1122,16 +1162,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: diff --git a/sentieon_cli/stages/cnv.py b/sentieon_cli/stages/cnv.py index d8ff55b..d02e48e 100644 --- a/sentieon_cli/stages/cnv.py +++ b/sentieon_cli/stages/cnv.py @@ -3,11 +3,13 @@ """ 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 @@ -96,3 +98,109 @@ def add_to(self, dag: DAG, upstream: Iterable[Job] = ()) -> CNVResult: 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/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/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/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/util.py b/sentieon_cli/util.py index 868dcd4..9dd6794 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -75,6 +75,16 @@ def cnvscope_sex_args( 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) diff --git a/tests/unit/test_cnv_sex.py b/tests/unit/test_cnv_sex.py index 27342ef..3fc4f1d 100644 --- a/tests/unit/test_cnv_sex.py +++ b/tests/unit/test_cnv_sex.py @@ -25,6 +25,7 @@ ) from sentieon_cli.util import ( SampleSex, + caller_sex_arg, cnvscope_sex_args, sample_sex_arg, ) @@ -210,6 +211,19 @@ def test_cnvscope_sex_args_without_a_known_sex(sex, messages): 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"), diff --git a/tests/unit/test_stage_cnv.py b/tests/unit/test_stage_cnv.py index 5b97761..e0ee30c 100644 --- a/tests/unit/test_stage_cnv.py +++ b/tests/unit/test_stage_cnv.py @@ -2,8 +2,10 @@ Unit tests for the CNVscope stage """ +from importlib.resources import files import logging import pathlib +import sys from typing import List import packaging.version @@ -11,7 +13,11 @@ 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 +from sentieon_cli.stages.cnv import ( + CNV_MIN_VERSIONS, + CNVscopeStage, + PangenomeCNVStage, +) from sentieon_cli.util import SampleSex PACKAGE_LOGGER = "sentieon_cli" @@ -222,3 +228,147 @@ def test_result_fields(self, tmp_path): 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_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_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_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" From 6e3bc766221242ce6339d8d0756580f72374e5fe Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 10:36:56 -0700 Subject: [PATCH 52/58] Compose the dnascope-pangenome callers from the shared stages `SentieonPangenome` now builds its optional callers from the stages added in the previous commit, and everything both pangenome pipelines will need moved up to `BasePangenome`. `BasePangenome` gains the shared caller arguments -- `par_bed`, `sample_sex`, `segdup_caller`, `expansion_catalog` and the six `t1k_*` (the two loci defaulting to the constants in `stages/t1k.py`) -- with the matching `__init__` attributes, `has_cnv_model`, and the DAG-2 stashes `ploidy_json` / `sr_alignments` / `sr_replace_rg` (`sr_replace_rg` is unused until the hybrid pipeline needs it). `validate_t1k` moved up unchanged, and four thin helpers build the stages: `output_path(suffix)`, `add_t1k(dag, ctx, inputs, upstream)`, `add_expansion(dag, ctx, alignment)`, `add_segdup(dag, ctx, sr_alignment, *, lr_alignment, lr_bundle, overrides)` and `add_pangenome_cnv(dag, ctx, sv_vcf, inputs, *, replace_rg, interval)`. `--call_svs` stays on `SentieonPangenome`. `hybrid-pangenome` inherits the new flags, which are accepted but build no jobs yet. Deleted from `sentieon_pangenome.py`: the `SEGDUP_MIN_VERSION` / `EXPANSION_MIN_VERSION` / `T1K_MIN_VERSION` constants (now `*_MIN_VERSIONS` next to their stages) and the default T1K loci, the moved params and `__init__` lines, `validate_t1k`, `build_segdup_job`, `build_t1k_jobs`, `build_expansion_job`, `_add_cnv_jobs`, `_build_indel2cnv_job`, `_build_combine_cnv_job`, and the dead `self.rw_bam` that `build_dag` assigned but nothing read. The DAG snapshot over all 88 configurations is identical, including `compare --with-ids`. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/base_pangenome.py | 253 ++++++++++++++- sentieon_cli/sentieon_pangenome.py | 438 ++------------------------ tests/unit/test_sentieon_pangenome.py | 260 ++++++++++++--- 3 files changed, 496 insertions(+), 455 deletions(-) diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index f21e9a6..bf7f90f 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -4,9 +4,11 @@ import copy import pathlib -from typing import List, Optional +import sys +from typing import Iterable, List, Optional, Sequence from . import command_strings as cmds +from .dag import DAG from .driver import ( AlignmentStat, BaseAlgo, @@ -20,8 +22,19 @@ ) from .job import Job from .pipeline import BasePipeline +from .stages.base import StageContext +from .stages.cnv import PangenomeCNVResult, PangenomeCNVStage +from .stages.expansion import ExpansionHunterResult, ExpansionHunterStage from .stages.metrics import MetricsPaths -from .util import path_arg +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 path_arg, require_versions, sample_sex_arg class BasePangenome(BasePipeline): @@ -82,6 +95,84 @@ 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})." + ), + }, } ) @@ -98,6 +189,164 @@ 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 + # 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 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 diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 74d7f75..65c2375 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -7,7 +7,7 @@ import json import pathlib import sys -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Dict, List, Optional, Set, Union import packaging.version @@ -27,14 +27,15 @@ from .shell_pipeline import Command, Pipeline from .stages.alignment import ( BwaExtractStage, - ReadWriterStage, find_unzip, ) from .stages.base import StageContext -from .stages.cnv import CNV_MIN_VERSIONS, CNVscopeStage +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, @@ -44,13 +45,11 @@ ) from .stages.transfer import TransferConfig from .util import ( - SampleSex, __version__, check_kmc_patch, parse_rg_line, path_arg, require_versions, - sample_sex_arg, total_memory, vcf_id, ) @@ -69,21 +68,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__) @@ -145,23 +129,6 @@ class SentieonPangenome(BasePangenome): "Prefix to strip from pangenome contig names (GRCh38#0#)" ), }, - "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, - }, "skip_metrics": { "help": "Skip metrics collection and multiQC", "action": "store_true", @@ -170,67 +137,6 @@ 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, @@ -265,29 +171,16 @@ def __init__(self) -> None: self.bed: Optional[pathlib.Path] = None self.call_svs = False self.gvcf = False - self.par_bed: Optional[pathlib.Path] = None 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 - # Stashed by `build_dag` for the second, sex-aware DAG - self.ploidy_json: Optional[pathlib.Path] = None - self.sr_alignment: Optional[pathlib.Path] = None - self.cnv_input_bams: List[pathlib.Path] = [] def _cnv_in_second_dag(self) -> bool: """CNV calling runs in the second, sex-aware DAG""" @@ -425,7 +318,7 @@ def validate_segdup(self) -> None: ) sys.exit(2) - require_versions(SEGDUP_MIN_VERSION, skip=self.skip_version_check) + require_versions(SEGDUP_MIN_VERSIONS, skip=self.skip_version_check) def validate_expansion(self) -> None: if self.expansion_catalog is None: @@ -445,7 +338,7 @@ def validate_expansion(self) -> None: ) sys.exit(2) - require_versions(EXPANSION_MIN_VERSION, skip=self.skip_version_check) + require_versions(EXPANSION_MIN_VERSIONS, skip=self.skip_version_check) def validate_cnv(self) -> None: """Validate the arguments used for sex-aware CNV calling""" @@ -460,32 +353,6 @@ def validate_cnv(self) -> None: require_versions(CNV_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): - return - - require_versions(T1K_MIN_VERSION, skip=self.skip_version_check) - def validate_bundle(self) -> None: pop_vcf = self.required(self.pop_vcf, "pop_vcf") gbz = self.required(self.gbz, "gbz") @@ -662,8 +529,6 @@ def build_dag(self) -> DAG: ) 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() @@ -742,7 +607,7 @@ def build_dag(self) -> DAG: 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) @@ -779,7 +644,7 @@ def build_dag(self) -> DAG: ).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: @@ -795,53 +660,23 @@ def build_dag(self) -> DAG: 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] - self.cnv_input_bams = cnv_input_bams + # Stash the short-read alignments for the second DAG + self.sr_alignments = sr_alignments # 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=[cnv_input_bams[0]], + 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(ctx.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(ctx.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 @@ -1099,41 +934,6 @@ def build_count_ad_update_job( task_name="ad-update", ) - 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""" - reference = self.required(self.reference, "reference") - bundle = self.required(self.model_bundle, "model_bundle") - - 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=reference, - sr_bundle=bundle, - input_vcf=input_vcf, - sex=sex, - genes=genes, - overrides=overrides, - ), - "segdup-caller", - self.cores, - task_name="segdup", - ) - def build_second_dag(self) -> Optional[DAG]: """Build the second DAG for sex-aware downstream tools""" if not self._needs_second_dag(): @@ -1144,201 +944,33 @@ def build_second_dag(self) -> Optional[DAG]: self.logger.info("Building the second pangenome DAG") dag = DAG() + ctx = self.stage_context() # CNV calling with CNVscope, using the sample sex if self._cnv_in_second_dag(): - sv_vcf = pathlib.Path( - str(self.output_vcf).replace(".vcf.gz", "_sv.vcf.gz") - ) - self._add_cnv_jobs(dag, sv_vcf, self.cnv_input_bams) - - 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 + self.add_pangenome_cnv( + dag, + ctx, + self.output_path("_sv.vcf.gz"), + self.sr_alignments, + interval=self.bed, ) - 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, + if self.sr_alignments: + self.add_expansion(dag, ctx, self.sr_alignments[0]) + + # 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 () + ), ) - 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""" - self.required(self.reference, "reference") - - # Extract reads overlapping the locus to a BAM file with ReadWriter - extracted_bam = self.tmp_dir.joinpath(f"sample_{tag}.bam") - extract_result = ReadWriterStage( - ctx=self.stage_context(), - inputs=sr_alignments, - output=extracted_bam, - name=f"t1k-{tag}-extract", - task_name="t1k", - interval=locus, - ).build() - extract_job = extract_result.jobs[0] - - 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, - task_name="t1k", - ) - 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""" - reference = self.required(self.reference, "reference") - - return Job( - cmds.cmd_expansion_hunter( - out_expansions, - sr_alignment, - reference=reference, - variant_catalog=expansion_catalog, - sex="male" if self.sample_sex == SampleSex.MALE else "female", - threads=self.cores, - ), - "expansion-hunter", - self.cores, - task_name="expansion-hunter", - ) - - def _add_cnv_jobs( - self, - dag: DAG, - sv_vcf: pathlib.Path, - cnv_input_bams: List[pathlib.Path], - ) -> None: - """Add CNV calling jobs to the DAG""" - bundle = self.required(self.model_bundle, "model_bundle") - - ctx = self.stage_context() - 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(ctx.output_vcf).replace(".vcf.gz", "_cnv.vcf.gz") - ) - cnv_probes = pathlib.Path( - str(ctx.output_vcf).replace(".vcf.gz", "_cnv.probes") - ) - - # CNVscope on BWA deduped BAM, then CNVModelApply. The alignment - # and the PangenomeSV output were both written by the first DAG. - cnv_result = CNVscopeStage( - ctx=ctx, - inputs=cnv_input_bams, - model=bundle.joinpath("cnv.model"), - cnvscope_vcf=cnvscope_vcf, - cnv_vcf=cnv_apply_vcf, - sample_sex=self.sample_sex, - par_bed=self.cnv_par_bed, - dump_probes=cnv_probes, - interval=self.bed, - ).add_to(dag) - - # Convert PangenomeSV output to CNVs - indel2cnv_job = self._build_indel2cnv_job(indel2cnv_vcf, sv_vcf) - dag.add_job(indel2cnv_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_result.apply_job, indel2cnv_job}) - - 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, - task_name="cnv", - ) - - 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, - task_name="cnv", - ) diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index a619f96..ca6e9fa 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -11,7 +11,9 @@ 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.util import SampleSex @@ -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,14 @@ def create_pipeline(self): pipeline.skip_pangenome_name_checks = True pipeline.skip_pop_vcf_id_check = True pipeline.tmp_dir = self.mock_dir - + # 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} @@ -200,16 +202,16 @@ def test_skip_model_apply(self): assert concat_job is not None # 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]) + 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" - ) + 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 @@ -247,9 +249,7 @@ def test_gvcf_adds_gvcftyper_and_routes_outputs(self): # 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" - ) + 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 @@ -282,9 +282,7 @@ def test_ad_update_default(self): 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" - ) + 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 @@ -317,17 +315,24 @@ def test_no_gvcftyper_without_gvcf(self): 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): @@ -338,6 +343,180 @@ 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() @@ -347,9 +526,7 @@ def test_call_svs(self): assert isinstance(dag, DAG) # The dnascope job should exist - 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" in job_names @@ -367,9 +544,7 @@ 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): @@ -377,9 +552,7 @@ def test_call_svs_disabled_by_default(self): pipeline = self.create_pipeline() 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": @@ -398,9 +571,7 @@ def test_skip_small_variants(self): 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" not in job_names assert "model-apply" not in job_names @@ -415,9 +586,7 @@ def test_skip_small_variants_with_call_svs(self): 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 @@ -440,9 +609,7 @@ def test_skip_small_variants_with_call_svs(self): 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): @@ -566,9 +733,7 @@ def test_ploidy_estimation_runs_without_sex_aware_callers(self): pipeline = self.create_pipeline() assert pipeline._needs_second_dag() is False - job_names, all_jobs = self._get_all_job_names( - pipeline.build_dag() - ) + 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") @@ -725,9 +890,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): @@ -735,9 +898,7 @@ def test_cnv_with_skip_model_apply(self): pipeline = self.create_pipeline() pipeline.call_svs = True pipeline.skip_model_apply = True - first_job_names, _ = self._get_all_job_names( - pipeline.build_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 @@ -748,4 +909,3 @@ def test_cnv_with_skip_model_apply(self): assert "cnv-model-apply" in job_names assert "indel2cnv" in job_names assert "combine-cnv" in job_names - From 179ba3f7fe4171dd139c5edaad31866f55d03216 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 10:43:15 -0700 Subject: [PATCH 53/58] Validate the optional callers and estimate ploidy in hybrid-pangenome Give `hybrid-pangenome` the argument surface, validation, and first-DAG jobs of the optional callers that `dnascope-pangenome` already has. The new `--call_cnvs` flag opts into CNVscope calling, and the shared `--sample_sex`, `--par_bed`, `--segdup_caller`, `--expansion_catalog` and `--t1k_*` arguments come from `BasePangenome`. `validate_bundle` now records whether the bundle carries a 'cnv.model' file and requires `diploid_model` when `--segdup_caller` is set, since segdup-caller reads it from the bundle passed as its `--lr_model`. `validate()` runs `validate_segdup`, `validate_expansion`, `validate_t1k` and `validate_cnv` after the readgroup checks. segdup-caller and ExpansionHunter accept a single `--sr_aln` file (segdup-caller also a single `--lr_aln`), `--segdup_caller` cannot be combined with `--skip_small_variants`, and `--call_cnvs` requires SV calling and a bundle CNV model. CNV runs resolve the PAR BED file during validation and exit when none is available; a run without `--bed` warns that CNVscope will cover every contig. `--bed` now restricts CNV calling as well as small-variant calling. The first DAG always adds the `PloidyStage`, writing `_ploidy.json` whatever the enabled callers, and runs T1K HLA/KIR on the short reads. It also stashes the short-read alignments, their `--replace_rg` rows, and the single long-read alignment for the second DAG. `--call_cnvs`, `--segdup_caller` and `--expansion_catalog` are validated but build no second-DAG jobs yet; `build_second_dag` follows in the next commit. The DAG snapshots gain one `estimate-ploidy` row per hybrid-pangenome config and are otherwise identical. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/hybrid_pangenome.py | 147 +++++++++- tests/unit/test_hybrid_pangenome.py | 439 +++++++++++++++++++++++++++- 2 files changed, 580 insertions(+), 6 deletions(-) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 4f5f2f9..ce0a292 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -40,8 +40,12 @@ 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, @@ -119,11 +123,19 @@ class HybridPangenome(BasePangenome): "flags": ["-b", "--bed"], "help": ( "Region BED file. Supplying this file will limit " - "small-variant calling to the intervals inside the BED " - "file." + "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 " @@ -208,6 +220,7 @@ def __init__(self) -> None: 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.pangenome_contig_prefix = "GRCh38#0#" @@ -228,6 +241,21 @@ def __init__(self) -> None: 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""" @@ -270,6 +298,11 @@ def validate(self) -> None: 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 ) @@ -336,9 +369,86 @@ def validate(self) -> None: sys.exit(2) # After the version gate and the contig checks, so the tool choice - # is only logged for runs that proceed + # 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. @@ -425,6 +535,10 @@ def validate_bundle(self) -> None: 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( @@ -433,6 +547,8 @@ def validate_bundle(self) -> None: ) sys.exit(2) + self.has_cnv_model = "cnv.model" in bundle_members + bundle_vcf_id = bundle_info.get("SentieonVcfID") if ( bundle_vcf_id @@ -853,15 +969,40 @@ def build_dag(self) -> DAG: 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, diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index be2dd3f..ed7cd1b 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -160,6 +160,7 @@ def test_dag_jobs_present(self): "locuscollector-lift", "dedup-lift", "metrics", + "estimate-ploidy", "pangenome-sv", "dnascope", "model-apply", @@ -645,6 +646,7 @@ def test_aligned_dag_jobs(self): "vg-haplotypes", "graph-update", "mm2-lift", + "estimate-ploidy", "pangenome-sv", "dnascope", "model-apply", @@ -717,9 +719,7 @@ def test_aligned_dag_dependencies(self): "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" - } + 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""" @@ -1035,3 +1035,436 @@ def test_collect_readgroups_parses_input_headers(self, monkeypatch): 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 + + def test_build_second_dag_not_implemented_yet(self): + """The sex-aware callers do not build jobs yet""" + pipeline = self.create_pipeline() + pipeline.build_dag() + + assert pipeline.build_second_dag() is None + + # 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 From 3ff98a9c83fb393a3247b7ecc853794b0e0bd4ec Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 10:49:12 -0700 Subject: [PATCH 54/58] Call CNVs, repeat expansions and segdups in the hybrid-pangenome second DAG `HybridPangenome.build_second_dag` builds the sex-aware callers once ploidy estimation has run, using the short-read alignments, the sample sex and the long-read alignment stashed by `build_dag`. With `--call_cnvs` the DAG holds the `PangenomeCNVStage` jobs: CNVscope over the short reads only -- restricted by `--bed` and carrying the short-read `--replace_rg` rows for `--sr_aln` input -- CNVModelApply, `indel2cnv.py` over the PangenomeSV VCF, and `combine_cnv.py` writing `_cnv.vcf.gz`. `--expansion_catalog` adds ExpansionHunter over the short reads. `--segdup_caller` adds segdup-caller with both data types (`--short` and `--long`) and this pipeline's model bundle as both `--sr_model` and `--lr_model`, since it ships the `diploid_model`. The DAG-snapshot harness gains seven hybrid-pangenome configs for the new callers; every pre-existing config is unchanged. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/hybrid_pangenome.py | 39 ++++ tests/unit/test_hybrid_pangenome.py | 326 +++++++++++++++++++++++++++- 2 files changed, 361 insertions(+), 4 deletions(-) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index ce0a292..3c556ea 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -1051,6 +1051,45 @@ def build_dag(self) -> DAG: 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. diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index ed7cd1b..29d3ad8 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -2,10 +2,13 @@ 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 @@ -21,6 +24,32 @@ 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: @@ -1218,12 +1247,301 @@ def test_lr_alignment_stash_with_two_inputs(self): assert pipeline.lr_alignment is None - def test_build_second_dag_not_implemented_yet(self): - """The sex-aware callers do not build jobs yet""" - pipeline = self.create_pipeline() + # 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") - assert pipeline.build_second_dag() is None + 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 From 7121923ed7edca9d321d793b35b914f6415f8ed8 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 13:42:51 -0700 Subject: [PATCH 55/58] Drop the path selector from the vg paths gfa2fa fallback The `-Q ` (--paths-by) selector passed to `vg paths` selected incorrectly on the updated graph during testing. Run `vg paths -x -F` without a selector, which writes every path in the GFA file, and rename the builder to cmd_vg_gfa2fa, mirroring cmd_pgutil_gfa2fa, now that it takes only the FASTA and GFA paths. Co-Authored-By: Claude Fable 5.1 --- sentieon_cli/command_strings.py | 7 +++---- sentieon_cli/hybrid_pangenome.py | 6 ++---- tests/unit/test_hybrid_pangenome.py | 14 ++------------ 3 files changed, 7 insertions(+), 20 deletions(-) diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index ae5a5de..9b23910 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -1308,17 +1308,16 @@ def cmd_vg_paths_fasta( return Pipeline(Command(*cmd), file_output=output_fasta) -def cmd_vg_paths_ref_fasta( +def cmd_vg_gfa2fa( output_fasta: pathlib.Path, gfa_file: pathlib.Path, - ref_name: str, ) -> Pipeline: - """Extract a graph's reference paths as FASTA. + """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), "-Q", ref_name, "-F"] + cmd = ["vg", "paths", "-x", str(gfa_file), "-F"] return Pipeline(Command(*cmd), file_output=output_fasta) diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 3c556ea..78a1ab3 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -76,7 +76,7 @@ # The last driver release whose `pgutil gfa2fa` handles only GRCh38 # pangenomes. Non-GRCh38 runs on it (and older) extract the graph's -# reference paths with `vg paths` instead, which is much slower. +# 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 @@ -1287,9 +1287,7 @@ def build_gfa2fa_job( ) -> Job: """Generate the FASTA sequences of the updated graph""" shell = ( - cmds.cmd_vg_paths_ref_fasta( - pangenome_fasta, pangenome_gfa, self.pangenome_ref_name - ) + cmds.cmd_vg_gfa2fa(pangenome_fasta, pangenome_gfa) if self.gfa2fa_with_vg else cmds.cmd_pgutil_gfa2fa( pangenome_fasta, ref_fai, pangenome_gfa diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index 29d3ad8..4d35072 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -397,24 +397,14 @@ def test_gfa2fa_with_vg_paths(self): 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} -Q GRCh38 -F >'{fasta}'") + 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_gfa2fa_with_vg_paths_chm13(self): - """`vg paths` extracts the paths of the configured reference""" - pipeline = self.create_pipeline() - pipeline.gfa2fa_with_vg = True - pipeline.pangenome_ref_name = "CHM13" - dag = pipeline.build_dag() - _, all_jobs = self._get_all_job_names(dag) - - gfa2fa_cmd = str(self._get_job(all_jobs, "gfa2fa").shell) - assert "-Q CHM13" in gfa2fa_cmd - def test_resolve_gfa2fa_tool_hg38(self, monkeypatch): """A GRCh38 reference keeps `pgutil gfa2fa`, unprobed""" pipeline = self.create_pipeline() From e554e257ccc78aa9cb10b8e82568c0879a76bd93 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Tue, 8 Sep 2026 13:56:42 -0700 Subject: [PATCH 56/58] Import SEGDUP_MIN_VERSIONS from stages.segdup in the Dockerfile The image's segdup-caller version check imported `SEGDUP_MIN_VERSION` from `sentieon_cli.sentieon_pangenome`, which moved to `sentieon_cli.stages.segdup.SEGDUP_MIN_VERSIONS` when the caller stages were extracted, so the build failed with an ImportError. Co-Authored-By: Claude Fable 5.1 --- Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From 5d39e3a77ff00e7b75877546fdde593bb18b5cfd Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 12 Sep 2026 12:46:02 -0700 Subject: [PATCH 57/58] Detect the pangenome reference from the GBZ metadata A wrong `--pangenome_ref_name` was silent: `vg haplotypes --set-reference` only warns, the sampled graph loses its reference paths, `vg convert -Q` writes no rGFA tags and `pgutil lift` exits 0 with every read unmapped. Read the reference name and contig prefix from the GBWT metadata in the GBZ header instead (new `pangenome_meta` module, well under a second per graph). The backbone reference is the unfragmented reference sample. Make `--pangenome_ref_name` and `--pangenome_contig_prefix` hidden overrides that must agree with the graph, check the backbone contigs and build against the reference FASTA, replace the `grch38.gbz` filename-suffix check, apply the GRCh38 contig-length check only to GRCh38 backbones, and add `check-sample-gbz` / `check-sample-gfa` jobs to both DAGs so a sampled graph or GFA without the reference fails the run. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Lo4R3WSNwCYMcKzvGKCqSs --- sentieon_cli/base_pangenome.py | 342 +++++++- sentieon_cli/command_strings.py | 61 +- sentieon_cli/hybrid_pangenome.py | 98 +-- sentieon_cli/pangenome_meta.py | 782 ++++++++++++++++++ sentieon_cli/sentieon_pangenome.py | 104 +-- tests/data/pangenome/generate.sh | 40 + tests/data/pangenome/rgfa_no_header.gfa | 3 + tests/data/pangenome/rgfa_no_tags.gfa | 4 + tests/data/pangenome/rgfa_ok.gfa | 6 + tests/data/pangenome/rgfa_wrong_reference.gfa | 6 + tests/data/pangenome/tiny.gbz | Bin 0 -> 2184 bytes tests/data/pangenome/tiny.gfa | 28 + tests/data/pangenome/tiny.hapl | Bin 0 -> 760 bytes tests/data/pangenome/tiny_chm13.gbz | Bin 0 -> 2184 bytes tests/data/pangenome/tiny_chm13.gfa | 28 + tests/data/pangenome/tiny_no_backbone.gbz | Bin 0 -> 2240 bytes tests/data/pangenome/tiny_no_backbone.gfa | 30 + tests/data/pangenome/tiny_no_ref_tag.gbz | Bin 0 -> 2152 bytes tests/data/pangenome/tiny_no_ref_tag.gfa | 28 + tests/data/pangenome/tiny_sampled.gbz | Bin 0 -> 2184 bytes tests/unit/test_hybrid_pangenome.py | 70 ++ tests/unit/test_pangenome_meta.py | 394 +++++++++ tests/unit/test_pipeline_validation.py | 331 +++++++- tests/unit/test_sentieon_pangenome.py | 79 ++ ...test_sentieon_pangenome_kmc_integration.py | 7 + 25 files changed, 2279 insertions(+), 162 deletions(-) create mode 100644 sentieon_cli/pangenome_meta.py create mode 100755 tests/data/pangenome/generate.sh create mode 100644 tests/data/pangenome/rgfa_no_header.gfa create mode 100644 tests/data/pangenome/rgfa_no_tags.gfa create mode 100644 tests/data/pangenome/rgfa_ok.gfa create mode 100644 tests/data/pangenome/rgfa_wrong_reference.gfa create mode 100644 tests/data/pangenome/tiny.gbz create mode 100644 tests/data/pangenome/tiny.gfa create mode 100644 tests/data/pangenome/tiny.hapl create mode 100644 tests/data/pangenome/tiny_chm13.gbz create mode 100644 tests/data/pangenome/tiny_chm13.gfa create mode 100644 tests/data/pangenome/tiny_no_backbone.gbz create mode 100644 tests/data/pangenome/tiny_no_backbone.gfa create mode 100644 tests/data/pangenome/tiny_no_ref_tag.gbz create mode 100644 tests/data/pangenome/tiny_no_ref_tag.gfa create mode 100644 tests/data/pangenome/tiny_sampled.gbz create mode 100644 tests/unit/test_pangenome_meta.py diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index bf7f90f..9b6c82d 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -2,10 +2,11 @@ A base class for pangenome pipelines """ +import argparse import copy import pathlib import sys -from typing import Iterable, List, Optional, Sequence +from typing import Dict, Iterable, List, Optional, Sequence, Set from . import command_strings as cmds from .dag import DAG @@ -21,7 +22,14 @@ WgsMetricsAlgo, ) from .job import Job +from .pangenome_meta import ( + PangenomeMetadataError, + PangenomeReference, + detect_pangenome_reference, + read_hapl_header, +) from .pipeline import BasePipeline +from .shard import GRCH38_CONTIGS, detect_reference_build from .stages.base import StageContext from .stages.cnv import PangenomeCNVResult, PangenomeCNVStage from .stages.expansion import ExpansionHunterResult, ExpansionHunterStage @@ -36,6 +44,19 @@ ) from .util import path_arg, require_versions, sample_sex_arg +# 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#" + +# 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): """A pipeline base class for short reads""" @@ -173,6 +194,23 @@ class BasePangenome(BasePipeline): 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", + }, } ) @@ -199,6 +237,17 @@ def __init__(self) -> 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 + self.fai_data: Dict[str, Dict[str, int]] = {} + 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] = [] @@ -230,6 +279,297 @@ def validate_t1k(self) -> None: 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 validate_grch38_contigs(self) -> None: + """Check the reference and pop VCF contig lengths against GRCh38. + + The lengths only describe GRCh38; other pangenome references are + covered by the reference FASTA checks in + `resolve_pangenome_reference` instead. + """ + if self.skip_contig_checks: + return + if self.pangenome_ref_name != DEFAULT_PANGENOME_REF_NAME: + self.logger.info( + "The pangenome reference is '%s', so the GRCh38 " + "contig-length check is skipped", + self.pangenome_ref_name, + ) + return + + # Check the fai file contigs + mismatch_contigs: Set[str] = set() + for ctg, length in GRCH38_CONTIGS.items(): + fai_length = self.fai_data.get(ctg, {}).get("length", -1) + if length != fai_length: + mismatch_contigs.add(ctg) + if mismatch_contigs: + self.logger.error( + "Reference contigs with unexpected lengths: %s", + ", ".join(mismatch_contigs), + ) + sys.exit(2) + + # Check the pop VCF file contigs + if self.dry_run: + return + mismatch_contigs = set() + for ctg, length in GRCH38_CONTIGS.items(): + if length != self.pop_vcf_contigs.get(ctg, -1): + mismatch_contigs.add(ctg) + if mismatch_contigs: + self.logger.error( + "Pop VCF contigs with unexpected lengths: %s", + ", ".join(mismatch_contigs), + ) + sys.exit(2) + + 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") diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 9b23910..c208463 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -319,6 +319,63 @@ def cmd_pyexec_sad_lad_update( 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, @@ -1275,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 = [ diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index 78a1ab3..ea83b2b 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -27,7 +27,6 @@ from .job import Job from .logging import get_logger from .shard import ( - GRCH38_CONTIGS, detect_reference_build, determine_shards_from_fai, parse_fai, @@ -151,16 +150,6 @@ class HybridPangenome(BasePangenome): ), "type": path_arg(exists=True, is_file=True), }, - "pangenome_contig_prefix": { - "default": "GRCh38#0#", - "help": ( - "Prefix to strip from pangenome contig names (GRCh38#0#)" - ), - }, - "pangenome_ref_name": { - "default": "GRCh38", - "help": "Reference name in the pangenome (GRCh38).", - }, "rgsm": { "help": ( "Overwrite the SM tag of the input readgroups for " @@ -185,18 +174,10 @@ class HybridPangenome(BasePangenome): "type": path_arg(exists=True, is_file=True), }, # Hidden arguments - "skip_contig_checks": { - "help": argparse.SUPPRESS, - "action": "store_true", - }, "skip_model_apply": { "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", @@ -223,8 +204,6 @@ def __init__(self) -> None: self.call_cnvs = False self.lr_align_input = False self.lr_input_ref: Optional[pathlib.Path] = None - self.pangenome_contig_prefix = "GRCh38#0#" - self.pangenome_ref_name = "GRCh38" self.rgsm: Optional[str] = None self.extract_model_name = "extract.model" self.gfa2fa_with_vg = False @@ -235,9 +214,7 @@ def __init__(self) -> None: self.sr_readgroups: List[List[Dict[str, str]]] = [] self.lr_readgroups: List[List[Dict[str, str]]] = [] self.sample_sm = "" - self.skip_contig_checks = False self.skip_model_apply = False - self.skip_pangenome_name_checks = False self.skip_pop_vcf_id_check = False self.skip_small_variants = False self.skip_svs = False @@ -264,7 +241,10 @@ def validate(self) -> None: self.shards = determine_shards_from_fai( self.fai_data, 10 * 1000 * 1000 ) - self.pop_vcf_contigs: Dict[str, Optional[int]] = {} + # Before `validate_bundle`, which picks the bundle's + # `extract..model` member by the resolved name + self.resolve_pangenome_reference() + self.pop_vcf_contigs = {} 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) @@ -322,51 +302,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 are 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_grch38_contigs() # After the version gate and the contig checks, so the tool choice # is only logged for runs that proceed. The `reference_build` it @@ -860,9 +796,19 @@ def build_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 to GFA gfa_job = self.build_gfa_job(hap_raw_gfa, sample_pangenome) - dag.add_job(gfa_job, {haplotypes_job}) + 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( @@ -871,7 +817,7 @@ def build_dag(self) -> DAG: calling_lr, name="graph-update-raw", ) - dag.add_job(update_raw_job, {gfa_job} | realign_jobs) + 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( @@ -1224,7 +1170,7 @@ def build_haplotypes_job( "--include-reference", "--diploid-sampling", "--set-reference", - self.pangenome_ref_name, + self.ref_name(), ], ), "vg-haplotypes", @@ -1242,7 +1188,7 @@ 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, @@ -1269,7 +1215,7 @@ def build_graph_update_job( out_gfa, gfa_file=in_gfa, target_bed=bed, - prefix=self.pangenome_contig_prefix, + prefix=self.contig_prefix(), ) ) return Job( @@ -1334,7 +1280,7 @@ def build_minimap2_lift_job( self.model_bundle.joinpath("minimap2.model"), threads=self.cores, mm2_xargs=["--secondary=yes"], - lift_prefix=self.pangenome_contig_prefix, + lift_prefix=self.contig_prefix(), ), "mm2-lift", self.cores, @@ -1395,7 +1341,7 @@ def build_pangenomesv_job( out_vcf, gfa_file=pangenome_gfa, min_af=PANGENOME_SV_MIN_AF, - prefix=self.pangenome_contig_prefix, + prefix=self.contig_prefix(), ) ) return Job( 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/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 65c2375..27e258d 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -54,7 +54,6 @@ vcf_id, ) from .shard import ( - GRCH38_CONTIGS, determine_shards_from_fai, parse_fai, vcf_contigs, @@ -113,22 +112,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", @@ -138,14 +121,6 @@ class SentieonPangenome(BasePangenome): "action": "store_true", }, # 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", @@ -171,13 +146,9 @@ 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.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 @@ -201,7 +172,10 @@ def validate(self) -> None: self.shards = determine_shards_from_fai( self.fai_data, 10 * 1000 * 1000 ) - self.pop_vcf_contigs: Dict[str, Optional[int]] = {} + # Before `validate_bundle`, which picks the bundle's + # `extract..model` member by the resolved name + self.resolve_pangenome_reference() + self.pop_vcf_contigs = {} 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) @@ -253,51 +227,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_grch38_contigs() def validate_segdup(self) -> None: if self.segdup_caller is None: @@ -589,11 +519,21 @@ def build_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() @@ -603,7 +543,7 @@ def build_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 @@ -711,7 +651,7 @@ def build_dag(self) -> DAG: PangenomeSV( sv_vcf, gfa_file=gfa_file, - prefix=self.pangenome_contig_prefix, + prefix=self.contig_prefix(), ) ) call = DNAscopeStage( @@ -827,7 +767,7 @@ def build_haplotypes_job( "--include-reference", "--diploid-sampling", "--set-reference", - self.pangenome_ref_name, + self.ref_name(), ], ), "vg-haplotypes", @@ -845,7 +785,7 @@ 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, @@ -899,7 +839,7 @@ def build_minimap2_lift_job( "@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, 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 0000000000000000000000000000000000000000..d24bff4ddc989cf03dfaccdca0397c3aef83487b GIT binary patch literal 2184 zcmaJ?QAkr!7(VCTJGY$HGK*nR(}H}dIkh%}z}u=VP>HdD1!d`M)3v3pVH83yLQjGe zicyiE4+UWbSrR=3y;i=&^dWj!k@e6+_*B%ncmKEKxOFe?yZ?W_@B9CA&$-9Ik<%`P z{kuBoQ6tA9 zzML6mUBt5xE>j~mid){i{GK>A9@|)6G4XscFp<6U3MRUJBHDL=lWcG~w@ zY3sc+8P2}=wi?e}TN&1)_X4k;nk{-0-h|=3jr>A{m~R!uR>D1mn4kQ99?*CxfKDYs z1tIDMsnAZ!yWE~8pFa?64|R$6v+nh<2IArU{Gb?hvWF+{COf`HJej4LDl5;qJU?>_ z3n!jWWV7tlN5)2oznRcOh<uscvLi{!eHke1>YgkcJb#BM0=jB zRDvP;*h~GD5u$(0cV?6ZO2BqWUr90a?O2Esdejv?o0q))&9tFV2m2EH4t1877U%z0 z3OexRxc^LkQIUsNiuyu-wT}ih9-Lb29=WNnh$c?b_|>!Cgu%f+ljnu+#ZG(~&=>M+ z<7$yaedG}X@%>=j@R6U_#)i&_kwd)J&lp3Da|`p6*J&IBLR0L*FzU%M{P2E^V=(`n z$4$E}WC~py#prKh)}L{v%Lc{wWM1H59OonIebqlk1e;%p)2@Gbui7pk$Qs5 literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..7bf79bb9bda60fdc654433fedcf1845354a8811d GIT binary patch literal 760 zcmeZp4DexNU|?VZVh}*5A)*WnvQP#Gl!mDT=>-8s2*n@>rEyU(eK7U7_##jP3=EBo zO-#+qEi4n0l2cOC(sT0)KvM2dc~&T$3^fqu9$}}Ii;`7%!_P3j7P%fFy{;jc$>Q<~)Qu2D zVpJs5g+W+>l|@%Yx0N@syyzy2tcxzftD?TM^FC{QZk>aB&ilU4^S z9rURZt|OEPUlNuA=ztDEKlt8oY$w0l<85ec+OyZ^Z$8-*i}#*NB>VbL52Viw4v&mx z##lG;OoU6+h=uIt*U!JF4qi@Fl-s7iT;K6VjUQ-Lt1iCnOS>w=EvE8A)ipEz&vIMu zjj2fX>9^J7^7)l9Eq*Kb;!&p{QMvrb+Y@WZVq&Oje0ZBs_V?2*;m@n<%NaA zPp)ONtm`9Vx30oo!^~da%8bH{xtor8TQjtl3-N_-K}7yN~G zC3m7e{IG#IKZxsj#P32|*XcHVunY6)G1$1bFg|geE+N1*#U6xFPqZsYg!$30LF1e! zbCFFHhpvrm^p|nQx>`ofL;O7D9`F#y{fPRXNRM>}4>owX&xpPTpaWmDAwXVlLo7MW zo#+Sa<|HEWRLvMSkC?BNJs0<-^f8Nhu|MI*{f#=&d+4w@LRYF!hy7hnh&sC|P5r$4 zH|yLbvD;c87>NY5KqMH{A`va1aVTRKPE2|pn64@l+QY;oFS4F=sE7TQ?=xDwKX{=7 vU$nu)dg8zA2nA3N`=*ur`FzfLf2cEg7JB3t@4EqPU=$P(x_n-t4u<&`Z@z+2 literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..453f4602592b561b7770b6396737d8bfbacd850f GIT binary patch literal 2240 zcmaJ?O-NKx6u$SpH$VPNjXw*Ks0pJiFnwyy7zDni90QY(H(@O_bIh@E%rvo|O+-PP z1f!IqBB2&p&<_nI(JE*Y^s|b|MYOOeBU(hepnKn(&!%hBdko(_=X~Ee=iYbTo5x$)=OUb=M!aM^xchO51$X1 zKjvp=cRZR*r`ewOjExh&h_HeX{aC4J5^DpGVaQQE~du4{%?DEXP(j>6A{ zqJU}KFQ^5JbRQQqO=uXaq<--nu^$m~t~=+jB8LKN$%ej!Y^YK;R*F#3iH4jD^9I%F zlFSNqFsD_-2XF3S)!YCc0>!ccI`CC{!^Ih=1^6*frzi$VdMFL8%C$U+}VR8Ao3DVFS@O;xbS6I@?CQ?6}M$ZTMhk z=a(_q_}*dv>ia!`fGqrv3t`k#ZR9H;t^9muz7KK7gZYE5m3ZiH@@9CnY{Y%oO%Cu7 z$9EO=Kad{t03K}cF#kkf0?>i4+7KYGS=$pI7Ebho=U7HW_|^9?IzxmU09~#={k}Jj5_L9Gp2}49FV<1T*>x)%vF)%Gwj&YCwym%wz{5`N z8?U+Nx}cx2?)Qz0BFjmKdYCtyM>LwFK}82X^1y}~&lCU7BNRYA%uN&dq4$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 0000000000000000000000000000000000000000..a19ffcec2e298de409e0c7cef7226d418ca17358 GIT binary patch literal 2152 zcmah~OK1~O6rFDxn?$2l301JIg57AESWW6GU&k(t8!X8-KhJ3HShxx)#i_4>Pka$8I4pxS(Nd;47ci`1LH zcia5<;@evB^Oe;JEqC8|_0(=LR*ID@?X9GjB8&NMBHYATteJ=7ZzSwykOzGpBX0w} zgfrQEfvre+*jbOwb1m*-w8}o9DPLyqrRA`cBHqEvw=Y#^gFox z9}8>0WadL1k0hvqXg_SpWI>c<_RAhc#Um*KQin1(PZlt$#1b8MaB;@w~ze=8PVsEr2rVLkKNQC_67Z8z7MB};9|CM`2oVn@54cK$wyt* zvpr}V*aV!>O&=8M;N0z^IC#sSmY)Bs6!O65@!&~kQ{jhJnd+dw_W4moOwKN4CU0n~ ztcl|kzkDiIvN$+r{Jiiz_=wLN^aX!CZedQ;haVW$I#bp_ywQ7o%*wMce8Ba3Rty;T z66VM6(+mPu;RXX?)Z;Pqus+r`7uStsU#e)U{qH-zRn^Bs70)R9p6mgA^pE=w^;f7z z>@#@4;Nd(2Yd1#WChYKiIbr5uwvM=iEHJo*Q1s#Y)zkd0m`O_;Ej@ zPWll#Y;2L|)@LyA-cCfFy(H85Ec{dG-y;eSG~F;wUDHj&&`eX)HD;QGt7oRW9;-Ly zDeXyNniawQje0n5wLa60`-2yG;PV(f>?i(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 0000000000000000000000000000000000000000..c0e19be0c0eea49a89ebf40486ab527b1fac1709 GIT binary patch literal 2184 zcmaJ?O-NKx6u$SqH{;|_O=%%hm|Bqq=5Ni|%GaxA78)_U0SU^|F-IrOUyMSCKxidM zW*89(x(PxHlAu}zE!vc`$jViV3aqsd6m{Qw=d zQ}n76ZX)D_uL!dNbU?qN?^|{qyUAZwT~k|EfA~l+)NrA#J>GG#Gtt$3si*gH-@xF| z@Ca)qo`-OS8u5|c@cQ|WVcc@x_V2CqDMWsCwG<2guaTV zj(c~kkI%lZCB9r+9Wmnf!!MsWEe4apq@8;Q<%<$xzT3#oAuJ)p{PGq43B2Z@LENs)~@vC7!(rZ*o74_H%^OE(q@muE{}ypQeU z=1;eSQQ~hUtRX}{JUUk%+fL3503GI2uIPWwul>=a4|TkPc^G#dxuhE7LX*Gp(LH#C z&SW-C$QE?&lMUT3&XJ@GAr=C7I)WNyjOMKK>`IK7xUn%IoSMBbz z?xwIaB9_?}ybd>5=AW!pT-h zIiYJJ8~shqnZw0{-hTN#IRiYz)qCYL>9Nn?!A76B?}@$vpaWmEAwb@m+V;eNbfO>Z zTQL!lCotpQJZipd_Ix~;wCAkm#rcFE_c!Xq9-_m{2wk>5{RK;o5Ooewnu@1YzgbHq zXAcZBY+0sZTH&x^S%zsy$YIydk5?~vCbco+ap$-!vW0Z0hx4}HXDoAn@InW^YJ-RU r#DCdA3ZNd&O(Xf&^A%TwA}xt4&?CQk-}T@Cqo9J&t>;zhV48mcb%uUo literal 0 HcmV?d00001 diff --git a/tests/unit/test_hybrid_pangenome.py b/tests/unit/test_hybrid_pangenome.py index 4d35072..fbda6eb 100644 --- a/tests/unit/test_hybrid_pangenome.py +++ b/tests/unit/test_hybrid_pangenome.py @@ -118,6 +118,8 @@ def create_pipeline(self): 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" @@ -515,6 +517,74 @@ def test_pangenomesv_command(self): # 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() 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_validation.py b/tests/unit/test_pipeline_validation.py index f628345..c1bae30 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,9 +17,21 @@ 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.shard import GRCH38_CONTIGS from sentieon_cli.util import set_bwt_max_mem from tests.utils.test_helpers import create_mock_args +# A reference index that satisfies the GRCh38 contig-length check +GRCH38_CONTIGS_AS_FAI = { + ctg: {"length": length} for ctg, length in GRCH38_CONTIGS.items() +} + class TestDNAscopePipelineValidation: """Test DNAscope pipeline validation logic""" @@ -514,5 +527,321 @@ 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 TestGrch38ContigChecks: + """The GRCh38 contig-length check only runs for GRCh38 pangenomes""" + + def _pipeline(self, cls, **attributes): + pipeline = cls() + pipeline.logger = MagicMock() + pipeline.pangenome_ref_name = "GRCh38" + pipeline.fai_data = dict(GRCH38_CONTIGS_AS_FAI) + pipeline.pop_vcf_contigs = dict(GRCH38_CONTIGS) + for key, value in attributes.items(): + setattr(pipeline, key, value) + return pipeline + + def test_matching_contigs_pass(self, cls): + self._pipeline(cls).validate_grch38_contigs() + + def test_a_wrong_contig_length_exits(self, cls): + pipeline = self._pipeline(cls) + pipeline.fai_data["chr1"] = {"length": 1} + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_grch38_contigs() + assert excinfo.value.code == 2 + + def test_skip_contig_checks_bypasses_the_check(self, cls): + pipeline = self._pipeline(cls, skip_contig_checks=True) + pipeline.fai_data["chr1"] = {"length": 1} + pipeline.validate_grch38_contigs() + + def test_a_non_grch38_pangenome_skips_the_check(self, cls): + """A CHM13 reference has different contig lengths""" + pipeline = self._pipeline(cls, pangenome_ref_name="CHM13") + pipeline.fai_data = {"chr1": {"length": 1}} + pipeline.pop_vcf_contigs = {} + pipeline.validate_grch38_contigs() + assert "GRCh38" in pipeline.logger.info.call_args[0][0] + + def test_a_wrong_pop_vcf_contig_exits(self, cls): + pipeline = self._pipeline(cls) + pipeline.pop_vcf_contigs = {} + with pytest.raises(SystemExit) as excinfo: + pipeline.validate_grch38_contigs() + assert excinfo.value.code == 2 + + def test_a_dry_run_skips_the_pop_vcf_check(self, cls): + pipeline = self._pipeline(cls, dry_run=True) + pipeline.pop_vcf_contigs = {} + pipeline.validate_grch38_contigs() + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/test_sentieon_pangenome.py b/tests/unit/test_sentieon_pangenome.py index ca6e9fa..9d96fe2 100644 --- a/tests/unit/test_sentieon_pangenome.py +++ b/tests/unit/test_sentieon_pangenome.py @@ -75,6 +75,10 @@ def create_pipeline(self): 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()] @@ -607,6 +611,81 @@ 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()) diff --git a/tests/unit/test_sentieon_pangenome_kmc_integration.py b/tests/unit/test_sentieon_pangenome_kmc_integration.py index eaad07e..9b6d8b8 100644 --- a/tests/unit/test_sentieon_pangenome_kmc_integration.py +++ b/tests/unit/test_sentieon_pangenome_kmc_integration.py @@ -21,6 +21,10 @@ def test_validate_kmc_patch_success(self): pipeline.reference = "ref" pipeline.pop_vcf = "pop" pipeline.gbz = "ref.grch38.gbz" + pipeline.hapl = "ref.grch38.hapl" + # The mock graph files cannot be parsed; a dry run falls back to + # the default pangenome reference + pipeline.dry_run = True # Mock dependencies pipeline.logger = MagicMock() @@ -56,6 +60,9 @@ def test_validate_kmc_patch_failure(self): pipeline.validate_ref = MagicMock() pipeline.collect_readgroups = MagicMock() pipeline.bed = "bed" + pipeline.gbz = "ref.grch38.gbz" + pipeline.hapl = "ref.grch38.hapl" + pipeline.dry_run = True pipeline.logger = MagicMock() From fb193d34de276d1aee34e0451500e430a3357c05 Mon Sep 17 00:00:00 2001 From: Don Freed Date: Sat, 12 Sep 2026 14:10:21 -0700 Subject: [PATCH 58/58] Check pop VCF contig lengths against the reference for any backbone `validate_grch38_contigs` compared the reference index and the `--pop_vcf` with a hard-coded table of the 25 GRCh38 chromosome lengths, so it only ran for a GRCh38 backbone and never looked at the 170 unplaced contigs the HPRC graph also carries. Drop the table and compare the two files instead: every contig of the backbone reference detected in the GBZ must carry the same length in the `--pop_vcf` header as in the reference FASTA index. A CHM13 graph is now checked the same way a GRCh38 one is, a contig absent from the pop VCF or without a length is an error, and the message names the first ten offenders with both lengths. Only the hidden `--skip_contig_checks` flag bypasses the check; `--skip_pangenome_name_checks` does not. The pop VCF header is now parsed in pure Python. `vcf_contigs` and `vcf_id` shelled out to `bcftools view -h`, and `vcf_contigs` returned a fake three-contig dictionary in dry runs, which left the dry run checking nothing and the real run checking everything. `read_vcf_header` opens the file with `gzip` when its first two bytes say so -- the extension is not trusted -- and stops at the first line that is not a header line, so the body is never read: the header of the 10 GB population VCF parses in about 1 ms. Dry runs now parse it too, and fall back only when the file cannot be read at all, which is what the unit tests and the placeholder files behind them need. The `SentieonVcfID` comparison follows the same rule, so a dry run catches a pop VCF that does not belong with the model bundle. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01RWPD3suVQRpMZbJHijXhao --- sentieon_cli/base_pangenome.py | 173 ++++++++++--- sentieon_cli/dnascope_hybrid.py | 2 +- sentieon_cli/dnascope_longread.py | 2 +- sentieon_cli/hybrid_pangenome.py | 23 +- sentieon_cli/sentieon_pangenome.py | 20 +- sentieon_cli/shard.py | 74 ++---- sentieon_cli/util.py | 93 +++++-- tests/unit/test_pipeline_validation.py | 233 +++++++++++++++--- ...test_sentieon_pangenome_kmc_integration.py | 7 +- tests/unit/test_vcf_header.py | 175 +++++++++++++ 10 files changed, 628 insertions(+), 174 deletions(-) create mode 100644 tests/unit/test_vcf_header.py diff --git a/sentieon_cli/base_pangenome.py b/sentieon_cli/base_pangenome.py index 9b6c82d..5b2f328 100644 --- a/sentieon_cli/base_pangenome.py +++ b/sentieon_cli/base_pangenome.py @@ -6,7 +6,7 @@ import copy import pathlib import sys -from typing import Dict, Iterable, List, Optional, Sequence, Set +from typing import Dict, Iterable, List, Optional, Sequence from . import command_strings as cmds from .dag import DAG @@ -29,7 +29,7 @@ read_hapl_header, ) from .pipeline import BasePipeline -from .shard import GRCH38_CONTIGS, detect_reference_build +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 @@ -42,7 +42,14 @@ T1KResult, T1KStage, ) -from .util import path_arg, require_versions, sample_sex_arg +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 @@ -245,8 +252,11 @@ def __init__(self) -> 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 + # 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 @@ -490,50 +500,145 @@ def check_pangenome_hapl( detected.n_paths, ) - def validate_grch38_contigs(self) -> None: - """Check the reference and pop VCF contig lengths against GRCh38. + def load_pop_vcf_header(self, pop_vcf: Optional[pathlib.Path]) -> None: + """Parse the `--pop_vcf` header. - The lengths only describe GRCh38; other pangenome references are - covered by the reference FASTA checks in - `resolve_pangenome_reference` instead. + 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. """ - if self.skip_contig_checks: + self.pop_vcf_header = None + self.pop_vcf_contigs = {} + if pop_vcf is None: return - if self.pangenome_ref_name != DEFAULT_PANGENOME_REF_NAME: - self.logger.info( - "The pangenome reference is '%s', so the GRCh38 " - "contig-length check is skipped", - self.pangenome_ref_name, + 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, ) - return + 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) - # Check the fai file contigs - mismatch_contigs: Set[str] = set() - for ctg, length in GRCH38_CONTIGS.items(): - fai_length = self.fai_data.get(ctg, {}).get("length", -1) - if length != fai_length: - mismatch_contigs.add(ctg) - if mismatch_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( - "Reference contigs with unexpected lengths: %s", - ", ".join(mismatch_contigs), + "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) - # Check the pop VCF file contigs - if self.dry_run: + 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 - mismatch_contigs = set() - for ctg, length in GRCH38_CONTIGS.items(): - if length != self.pop_vcf_contigs.get(ctg, -1): - mismatch_contigs.add(ctg) - if mismatch_contigs: + + 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( - "Pop VCF contigs with unexpected lengths: %s", - ", ".join(mismatch_contigs), + "%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. diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index d8e1f38..8ad702b 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -377,7 +377,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 diff --git a/sentieon_cli/dnascope_longread.py b/sentieon_cli/dnascope_longread.py index 0bd9eaf..032f653 100644 --- a/sentieon_cli/dnascope_longread.py +++ b/sentieon_cli/dnascope_longread.py @@ -303,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 diff --git a/sentieon_cli/hybrid_pangenome.py b/sentieon_cli/hybrid_pangenome.py index ea83b2b..d3ae092 100644 --- a/sentieon_cli/hybrid_pangenome.py +++ b/sentieon_cli/hybrid_pangenome.py @@ -30,7 +30,6 @@ detect_reference_build, determine_shards_from_fai, parse_fai, - vcf_contigs, ) from .shell_pipeline import Command, Pipeline from .stages.alignment import ( @@ -61,7 +60,6 @@ path_arg, require_versions, total_memory, - vcf_id, ) HYBRID_PANGENOME_MIN_VERSIONS = { @@ -244,10 +242,7 @@ def validate(self) -> None: # Before `validate_bundle`, which picks the bundle's # `extract..model` member by the resolved name self.resolve_pangenome_reference() - self.pop_vcf_contigs = {} - 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) + self.load_pop_vcf_header(self.pop_vcf) self.validate_bundle() self.validate_output_vcf() @@ -302,7 +297,7 @@ def validate(self) -> None: "across decoy and unplaced contigs." ) - self.validate_grch38_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 @@ -486,18 +481,8 @@ def validate_bundle(self) -> None: 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 - and not self.dry_run - ): - assert self.pop_vcf is not None - 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 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""" diff --git a/sentieon_cli/sentieon_pangenome.py b/sentieon_cli/sentieon_pangenome.py index 27e258d..ed445e8 100644 --- a/sentieon_cli/sentieon_pangenome.py +++ b/sentieon_cli/sentieon_pangenome.py @@ -51,12 +51,10 @@ path_arg, require_versions, total_memory, - vcf_id, ) from .shard import ( determine_shards_from_fai, parse_fai, - vcf_contigs, ) SENT_PANGENOME_MIN_VERSIONS = { @@ -175,10 +173,7 @@ def validate(self) -> None: # Before `validate_bundle`, which picks the bundle's # `extract..model` member by the resolved name self.resolve_pangenome_reference() - self.pop_vcf_contigs = {} - 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) + self.load_pop_vcf_header(self.pop_vcf) self.validate_bundle() self.validate_fastq_rg() @@ -227,7 +222,7 @@ def validate(self) -> None: "across decoy and unplaced contigs." ) - self.validate_grch38_contigs() + self.validate_pangenome_contig_lengths() def validate_segdup(self) -> None: if self.segdup_caller is None: @@ -284,7 +279,7 @@ def validate_cnv(self) -> None: require_versions(CNV_MIN_VERSIONS, skip=self.skip_version_check) def validate_bundle(self) -> None: - pop_vcf = self.required(self.pop_vcf, "pop_vcf") + 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" @@ -357,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(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): diff --git a/sentieon_cli/shard.py b/sentieon_cli/shard.py index 9e509e5..8d04858 100644 --- a/sentieon_cli/shard.py +++ b/sentieon_cli/shard.py @@ -4,45 +4,16 @@ 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. @@ -192,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("<") @@ -224,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/util.py b/sentieon_cli/util.py index 9dd6794..7d79905 100644 --- a/sentieon_cli/util.py +++ b/sentieon_cli/util.py @@ -4,6 +4,7 @@ import argparse from enum import Enum +import gzip from importlib.metadata import PackageNotFoundError, version import multiprocessing as mp import os @@ -13,7 +14,16 @@ import subprocess as sp import sys import tempfile -from typing import Callable, Dict, List, Mapping, Optional, Tuple +from typing import ( + IO, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Tuple, +) import packaging.version @@ -33,6 +43,9 @@ _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 @@ -360,25 +373,77 @@ def parse_rg_line(rg_line: str) -> Dict[str, str]: return parsed -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"): +class VcfHeaderError(ValueError): + """A VCF header could not be read""" + + +def _open_vcf_text(in_vcf: pathlib.Path) -> IO[str]: + """Open a VCF for reading, transparently decompressing it. + + 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/unit/test_pipeline_validation.py b/tests/unit/test_pipeline_validation.py index c1bae30..ba5a546 100644 --- a/tests/unit/test_pipeline_validation.py +++ b/tests/unit/test_pipeline_validation.py @@ -23,15 +23,9 @@ PangenomeReference, ) from sentieon_cli.sentieon_pangenome import SentieonPangenome -from sentieon_cli.shard import GRCH38_CONTIGS -from sentieon_cli.util import set_bwt_max_mem +from sentieon_cli.util import VcfHeaderError, set_bwt_max_mem from tests.utils.test_helpers import create_mock_args -# A reference index that satisfies the GRCh38 contig-length check -GRCH38_CONTIGS_AS_FAI = { - ctg: {"length": length} for ctg, length in GRCH38_CONTIGS.items() -} - class TestDNAscopePipelineValidation: """Test DNAscope pipeline validation logic""" @@ -794,53 +788,222 @@ def test_both_overrides_survive_an_unreadable_graph(self, cls, tmp_path): @pytest.mark.parametrize( "cls", PANGENOME_PIPELINES, ids=lambda c: c.__name__ ) -class TestGrch38ContigChecks: - """The GRCh38 contig-length check only runs for GRCh38 pangenomes""" +class TestPangenomeContigLengthChecks: + """`BasePangenome.validate_pangenome_contig_lengths` in both pipelines. - def _pipeline(self, cls, **attributes): + 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_ref_name = "GRCh38" - pipeline.fai_data = dict(GRCH38_CONTIGS_AS_FAI) - pipeline.pop_vcf_contigs = dict(GRCH38_CONTIGS) + 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_contigs_pass(self, cls): - self._pipeline(cls).validate_grch38_contigs() + 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_wrong_contig_length_exits(self, cls): + def test_a_mismatching_length_exits(self, cls): pipeline = self._pipeline(cls) - pipeline.fai_data["chr1"] = {"length": 1} + pipeline.pop_vcf_contigs["chr2"] = 2001 with pytest.raises(SystemExit) as excinfo: - pipeline.validate_grch38_contigs() + 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.fai_data["chr1"] = {"length": 1} - pipeline.validate_grch38_contigs() - - def test_a_non_grch38_pangenome_skips_the_check(self, cls): - """A CHM13 reference has different contig lengths""" - pipeline = self._pipeline(cls, pangenome_ref_name="CHM13") - pipeline.fai_data = {"chr1": {"length": 1}} - pipeline.pop_vcf_contigs = {} - pipeline.validate_grch38_contigs() - assert "GRCh38" in pipeline.logger.info.call_args[0][0] - - def test_a_wrong_pop_vcf_contig_exits(self, cls): + 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.pop_vcf_contigs = {} + 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.validate_grch38_contigs() + pipeline.check_pop_vcf_id("population-other-20240101") assert excinfo.value.code == 2 - def test_a_dry_run_skips_the_pop_vcf_check(self, cls): + 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) - pipeline.pop_vcf_contigs = {} - pipeline.validate_grch38_contigs() + 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__": diff --git a/tests/unit/test_sentieon_pangenome_kmc_integration.py b/tests/unit/test_sentieon_pangenome_kmc_integration.py index 9b6d8b8..b4f5fe7 100644 --- a/tests/unit/test_sentieon_pangenome_kmc_integration.py +++ b/tests/unit/test_sentieon_pangenome_kmc_integration.py @@ -22,8 +22,9 @@ def test_validate_kmc_patch_success(self): pipeline.pop_vcf = "pop" pipeline.gbz = "ref.grch38.gbz" pipeline.hapl = "ref.grch38.hapl" - # The mock graph files cannot be parsed; a dry run falls back to - # the default pangenome reference + # 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 @@ -39,7 +40,6 @@ def test_validate_kmc_patch_success(self): 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.vcf_contigs", return_value={}), \ patch("sentieon_cli.sentieon_pangenome.check_kmc_patch", return_value=True) as mock_check, \ patch("sys.exit") as mock_exit: @@ -69,7 +69,6 @@ def test_validate_kmc_patch_failure(self): 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.vcf_contigs", 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_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__])