613 lines
20 KiB
Python
613 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Stage 08 joint genotyping from authoritative WGS Manifests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from datetime import datetime, timezone
|
|
import fcntl
|
|
import gzip
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from typing import Any
|
|
import uuid
|
|
|
|
from gfa_core.wgs_scientific_identity import (
|
|
build_wgs_analysis_id,
|
|
sha256_file,
|
|
)
|
|
from gfa_core.wgs_upstream_manifest import resolve_wgs_artifacts
|
|
|
|
|
|
STAGE = "wgs_joint_genotyping"
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Joint-genotype the authoritative Stage07 cohort GVCF"
|
|
)
|
|
parser.add_argument("--workspace", required=True)
|
|
parser.add_argument("--job-id", required=True)
|
|
parser.add_argument("--prefix", required=True)
|
|
parser.add_argument("--combine-manifest", required=True)
|
|
parser.add_argument("--index-manifest", required=True)
|
|
parser.add_argument("--gatk", default="gatk")
|
|
parser.add_argument("--max-alternate-alleles", type=int, default=6)
|
|
parser.add_argument("--calling-threshold", type=int, default=10)
|
|
parser.add_argument(
|
|
"--include-non-variant-sites",
|
|
type=lambda value: str(value).strip().lower() == "true",
|
|
default=False,
|
|
)
|
|
return parser
|
|
|
|
|
|
def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
|
|
temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
|
|
temporary.write_text(
|
|
json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def _safe_component(value: str, label: str) -> None:
|
|
if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", value or ""):
|
|
raise ValueError(f"{label} must be a non-empty safe path component")
|
|
|
|
|
|
def _authority_header(
|
|
path: str | Path,
|
|
*,
|
|
stage: str,
|
|
workspace: Path,
|
|
prefix: str,
|
|
job_id: str,
|
|
) -> dict[str, Any]:
|
|
try:
|
|
manifest = json.loads(Path(path).read_text(encoding="utf-8"))
|
|
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
|
|
raise ValueError(f"{stage} authority Manifest is unreadable") from exc
|
|
if (
|
|
not isinstance(manifest, dict)
|
|
or manifest.get("stage") != stage
|
|
or manifest.get("status") != "success"
|
|
):
|
|
raise ValueError(
|
|
f"{stage} authority Manifest is not a successful producer"
|
|
)
|
|
if (
|
|
manifest.get("prefix") != prefix
|
|
or manifest.get("job_id") != job_id
|
|
or manifest.get("workspace") != str(workspace)
|
|
):
|
|
raise ValueError(f"{stage} authority Manifest identity mismatch")
|
|
return manifest
|
|
|
|
|
|
def _invocation_path(command: str) -> Path:
|
|
resolved = shutil.which(command)
|
|
if not resolved:
|
|
raise RuntimeError(f"GATK executable is unavailable: {command}")
|
|
return Path(resolved).absolute()
|
|
|
|
|
|
def _executable_path(command: str) -> Path:
|
|
return _invocation_path(command).resolve()
|
|
|
|
|
|
def _checked(command: list[str], name: str) -> subprocess.CompletedProcess[str]:
|
|
environment = os.environ.copy()
|
|
executable_parent = str(_invocation_path(command[0]).parent)
|
|
environment["PATH"] = (
|
|
executable_parent + os.pathsep + environment.get("PATH", "")
|
|
)
|
|
result = subprocess.run(
|
|
command,
|
|
text=True,
|
|
capture_output=True,
|
|
check=False,
|
|
env=environment,
|
|
)
|
|
if result.returncode:
|
|
raise RuntimeError(
|
|
f"{name} failed with exit code {result.returncode}: "
|
|
f"{result.stderr.strip()}"
|
|
)
|
|
return result
|
|
|
|
|
|
def _tool_version(gatk: str) -> str:
|
|
result = _checked([gatk, "--version"], "GATK version check")
|
|
lines = (result.stdout or result.stderr).strip().splitlines()
|
|
if not lines:
|
|
raise RuntimeError("GATK version check produced no version")
|
|
return lines[0]
|
|
|
|
|
|
def _run_logged(command: list[str], log: Any, label: str) -> None:
|
|
log.write(f"{label} argv: {json.dumps(command)}\n")
|
|
log.flush()
|
|
result = subprocess.run(
|
|
command,
|
|
stdout=log,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
check=False,
|
|
env={
|
|
**os.environ,
|
|
"PATH": str(_invocation_path(command[0]).parent)
|
|
+ os.pathsep
|
|
+ os.environ.get("PATH", ""),
|
|
},
|
|
)
|
|
log.write(f"{label} exit={result.returncode}\n")
|
|
log.flush()
|
|
if result.returncode:
|
|
raise RuntimeError(f"{label} failed with exit code {result.returncode}")
|
|
|
|
|
|
def _vcf_header(
|
|
path: Path, *, require_non_ref: bool = False
|
|
) -> tuple[list[str], str]:
|
|
if not path.is_file() or path.stat().st_size <= 0:
|
|
raise RuntimeError("GVCF is missing or empty")
|
|
samples: list[str] | None = None
|
|
interval = ""
|
|
non_ref_declared = False
|
|
records = 0
|
|
try:
|
|
with gzip.open(path, "rt", encoding="utf-8") as handle:
|
|
for raw in handle:
|
|
line = raw.rstrip("\n")
|
|
if line.startswith("##ALT=<ID=NON_REF,"):
|
|
non_ref_declared = True
|
|
elif line.startswith("#CHROM\t"):
|
|
fields = line.split("\t")
|
|
if len(fields) < 10:
|
|
raise RuntimeError("GVCF has no sample columns")
|
|
samples = fields[9:]
|
|
elif line and not line.startswith("#"):
|
|
fields = line.split("\t")
|
|
if len(fields) < 10:
|
|
raise RuntimeError("VCF record is invalid")
|
|
if require_non_ref and "<NON_REF>" not in fields[4].split(","):
|
|
raise RuntimeError("combined GVCF record is invalid")
|
|
try:
|
|
position = int(fields[1])
|
|
except ValueError as exc:
|
|
raise RuntimeError(
|
|
"GVCF record position is invalid"
|
|
) from exc
|
|
if not fields[0] or position < 1:
|
|
raise RuntimeError("GVCF record position is invalid")
|
|
if not interval:
|
|
interval = f"{fields[0]}:{position}-{position}"
|
|
records += 1
|
|
except (OSError, UnicodeError) as exc:
|
|
raise RuntimeError("GVCF is unreadable") from exc
|
|
if (
|
|
samples is None
|
|
or not samples
|
|
or len(samples) != len(set(samples))
|
|
or (require_non_ref and not non_ref_declared)
|
|
or records < 1
|
|
or not interval
|
|
):
|
|
raise RuntimeError("VCF content is incomplete")
|
|
return samples, interval
|
|
|
|
|
|
def _validate_with_index(
|
|
gatk: str,
|
|
reference: Path,
|
|
gvcf: Path,
|
|
index: Path,
|
|
expected_samples: list[str],
|
|
log: Any,
|
|
label: str,
|
|
*,
|
|
require_non_ref: bool = False,
|
|
) -> None:
|
|
if index != Path(f"{gvcf}.tbi"):
|
|
raise RuntimeError(f"{label} index is not the declared VCF sidecar")
|
|
if not index.is_file() or index.stat().st_size <= 0:
|
|
raise RuntimeError(f"{label} index is missing or empty")
|
|
samples, interval = _vcf_header(
|
|
gvcf, require_non_ref=require_non_ref
|
|
)
|
|
if samples != expected_samples:
|
|
raise RuntimeError(f"{label} sample header mismatch")
|
|
_run_logged(
|
|
[
|
|
gatk,
|
|
"ValidateVariants",
|
|
"-R",
|
|
str(reference),
|
|
"-V",
|
|
str(gvcf),
|
|
"-L",
|
|
interval,
|
|
"--validation-type-to-exclude",
|
|
"ALLELES",
|
|
],
|
|
log,
|
|
label,
|
|
)
|
|
|
|
|
|
def _artifact(
|
|
workspace: Path,
|
|
source: Path,
|
|
destination: Path,
|
|
role: str,
|
|
) -> dict[str, Any]:
|
|
size = source.stat().st_size
|
|
return {
|
|
"artifact": role,
|
|
"role": role,
|
|
"path": str(destination.relative_to(workspace)),
|
|
"full_path": str(destination.resolve()),
|
|
"type": source.suffix.lstrip(".") or "file",
|
|
"size": size,
|
|
"size_bytes": size,
|
|
"sha256": sha256_file(source),
|
|
}
|
|
|
|
|
|
def _publish(
|
|
bundle: Path,
|
|
destination: Path,
|
|
manifest_work: Path,
|
|
authority: Path,
|
|
started_at: str,
|
|
) -> bool:
|
|
lock_path = authority.parent.parent / ".locks" / f"{authority.stem}.lock"
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with lock_path.open("a+", encoding="utf-8") as lock:
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
|
|
if authority.is_file():
|
|
try:
|
|
current = json.loads(authority.read_text(encoding="utf-8"))
|
|
if (
|
|
current.get("status") == "success"
|
|
and str(current.get("start_time") or "") >= started_at
|
|
):
|
|
return False
|
|
except Exception:
|
|
pass
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
os.replace(bundle, destination)
|
|
os.replace(manifest_work, authority)
|
|
return True
|
|
|
|
|
|
def _failure(
|
|
args: argparse.Namespace, run_id: str, started_at: str, error: str
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"manifest_version": "1.0",
|
|
"stage": STAGE,
|
|
"status": "failed",
|
|
"prefix": args.prefix,
|
|
"job_id": args.job_id,
|
|
"run_id": run_id,
|
|
"start_time": started_at,
|
|
"artifacts": {},
|
|
"errors": [error],
|
|
}
|
|
|
|
|
|
def _publish_failure(
|
|
authority: Path,
|
|
args: argparse.Namespace,
|
|
run_id: str,
|
|
started_at: str,
|
|
error: str,
|
|
) -> None:
|
|
lock_path = authority.parent.parent / ".locks" / f"{authority.stem}.lock"
|
|
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with lock_path.open("a+", encoding="utf-8") as lock:
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
|
|
failure_path = (
|
|
authority
|
|
if not authority.exists()
|
|
else authority.parent
|
|
/ f"{args.prefix}.{STAGE}.failed.{run_id}.manifest.json"
|
|
)
|
|
_atomic_json(
|
|
failure_path, _failure(args, run_id, started_at, error)
|
|
)
|
|
|
|
|
|
def run(args: argparse.Namespace) -> int:
|
|
workspace = Path(args.workspace).resolve()
|
|
run_id = uuid.uuid4().hex
|
|
started_at = datetime.now(timezone.utc).isoformat()
|
|
final = workspace / "final"
|
|
final.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
_safe_component(args.prefix, "prefix")
|
|
except ValueError as exc:
|
|
_atomic_json(
|
|
final / f"{STAGE}.failed.{run_id}.manifest.json",
|
|
_failure(args, run_id, started_at, str(exc)),
|
|
)
|
|
print(f"{STAGE} failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
authority = final / f"{args.prefix}.{STAGE}.manifest.json"
|
|
try:
|
|
combine_header = _authority_header(
|
|
args.combine_manifest,
|
|
stage="wgs_gvcf_combine",
|
|
workspace=workspace,
|
|
prefix=args.prefix,
|
|
job_id=args.job_id,
|
|
)
|
|
index_header = _authority_header(
|
|
args.index_manifest,
|
|
stage="wgs_index",
|
|
workspace=workspace,
|
|
prefix=args.prefix,
|
|
job_id=args.job_id,
|
|
)
|
|
sample_set_id = str(
|
|
combine_header.get("wgs_sample_set_id") or ""
|
|
)
|
|
reference_id = str(combine_header.get("wgs_reference_id") or "")
|
|
upstream_ids = {
|
|
"wgs_gvcf_combine": str(
|
|
combine_header.get("wgs_analysis_id") or ""
|
|
),
|
|
"wgs_index": str(index_header.get("wgs_analysis_id") or ""),
|
|
}
|
|
if (
|
|
not sample_set_id
|
|
or not reference_id
|
|
or not all(upstream_ids.values())
|
|
or index_header.get("wgs_reference_id") != reference_id
|
|
):
|
|
raise ValueError("two-way upstream scientific lineage mismatch")
|
|
producer_upstream = combine_header.get("upstream_wgs_analysis_ids")
|
|
if (
|
|
not isinstance(producer_upstream, dict)
|
|
or producer_upstream.get("wgs_index")
|
|
!= upstream_ids["wgs_index"]
|
|
):
|
|
raise ValueError(
|
|
"wgs_gvcf_combine index lineage does not match "
|
|
"the current wgs_index authority"
|
|
)
|
|
combine = resolve_wgs_artifacts(
|
|
args.combine_manifest,
|
|
"wgs_gvcf_combine",
|
|
("combined_gvcf", "combined_gvcf_index"),
|
|
sample_set_id,
|
|
reference_id,
|
|
upstream_ids["wgs_gvcf_combine"],
|
|
)
|
|
index = resolve_wgs_artifacts(
|
|
args.index_manifest,
|
|
"wgs_index",
|
|
("reference_fasta", "reference_fai", "reference_dict"),
|
|
expected_reference_id=reference_id,
|
|
expected_analysis_id=upstream_ids["wgs_index"],
|
|
)
|
|
if any(len(index[role]) != 1 for role in index):
|
|
raise RuntimeError(
|
|
"authoritative index must contain one FASTA, FAI and dictionary"
|
|
)
|
|
if any(len(combine[role]) != 1 for role in combine):
|
|
raise RuntimeError(
|
|
"authoritative combine must contain one GVCF and index"
|
|
)
|
|
metrics = combine_header.get("metrics") or {}
|
|
expected_samples = metrics.get("output_sample_ids")
|
|
if (
|
|
not isinstance(expected_samples, list)
|
|
or not expected_samples
|
|
or any(
|
|
not isinstance(sample, str) or not sample
|
|
for sample in expected_samples
|
|
)
|
|
or expected_samples != sorted(set(expected_samples))
|
|
or metrics.get("input_sample_ids") != expected_samples
|
|
):
|
|
raise RuntimeError("upstream combined GVCF sample set is incomplete")
|
|
if (
|
|
args.max_alternate_alleles != 6
|
|
or args.calling_threshold != 10
|
|
or args.include_non_variant_sites is not False
|
|
):
|
|
raise ValueError(
|
|
"joint genotyping parameters must be exactly 6, 10 and false"
|
|
)
|
|
version = _tool_version(args.gatk)
|
|
parameters = {
|
|
"gatk": args.gatk,
|
|
"operation": "GenotypeGVCFs",
|
|
"max_alternate_alleles": args.max_alternate_alleles,
|
|
"calling_threshold": args.calling_threshold,
|
|
"include_non_variant_sites": args.include_non_variant_sites,
|
|
}
|
|
consumed = [
|
|
item.sha256
|
|
for group in (combine, index)
|
|
for values in group.values()
|
|
for item in values
|
|
]
|
|
analysis_id = build_wgs_analysis_id(
|
|
STAGE,
|
|
consumed,
|
|
parameters,
|
|
{"gatk": version},
|
|
method="gatk_genotypegvcfs",
|
|
)
|
|
attempt = workspace / ".attempts" / run_id
|
|
bundle = attempt / "bundle"
|
|
bundle.mkdir(parents=True, exist_ok=False)
|
|
base = f"{args.prefix}.{STAGE}.cohort"
|
|
joint_vcf = bundle / f"{base}.joint.vcf.gz"
|
|
joint_vcf_index = Path(f"{joint_vcf}.tbi")
|
|
summary = bundle / f"{base}.summary.tsv"
|
|
log_path = bundle / f"{base}.joint.log"
|
|
reference = index["reference_fasta"][0].path
|
|
combined = combine["combined_gvcf"][0].path
|
|
combined_index = combine["combined_gvcf_index"][0].path
|
|
with log_path.open("w", encoding="utf-8") as log:
|
|
_validate_with_index(
|
|
args.gatk,
|
|
reference,
|
|
combined,
|
|
combined_index,
|
|
expected_samples,
|
|
log,
|
|
"ValidateVariants input",
|
|
require_non_ref=True,
|
|
)
|
|
genotype = [
|
|
args.gatk,
|
|
"GenotypeGVCFs",
|
|
"-R",
|
|
str(reference),
|
|
"-V",
|
|
str(combined),
|
|
"-O",
|
|
str(joint_vcf),
|
|
"--max-alternate-alleles",
|
|
str(args.max_alternate_alleles),
|
|
"--standard-min-confidence-threshold-for-calling",
|
|
str(args.calling_threshold),
|
|
"--include-non-variant-sites",
|
|
str(args.include_non_variant_sites).lower(),
|
|
]
|
|
_run_logged(genotype, log, "GenotypeGVCFs")
|
|
_run_logged(
|
|
[args.gatk, "IndexFeatureFile", "-I", str(joint_vcf)],
|
|
log,
|
|
"IndexFeatureFile",
|
|
)
|
|
_validate_with_index(
|
|
args.gatk,
|
|
reference,
|
|
joint_vcf,
|
|
joint_vcf_index,
|
|
expected_samples,
|
|
log,
|
|
"ValidateVariants",
|
|
)
|
|
_run_logged(
|
|
[
|
|
args.gatk,
|
|
"VariantsToTable",
|
|
"-V",
|
|
str(joint_vcf),
|
|
"-F",
|
|
"CHROM",
|
|
"-F",
|
|
"POS",
|
|
"-F",
|
|
"REF",
|
|
"-F",
|
|
"ALT",
|
|
"-F",
|
|
"QUAL",
|
|
"-O",
|
|
str(summary),
|
|
],
|
|
log,
|
|
"VariantsToTable",
|
|
)
|
|
try:
|
|
summary_lines = summary.read_text(encoding="utf-8").splitlines()
|
|
except (OSError, UnicodeError) as exc:
|
|
raise RuntimeError("joint summary is unreadable") from exc
|
|
if (
|
|
len(summary_lines) < 2
|
|
or summary_lines[0] != "CHROM\tPOS\tREF\tALT\tQUAL"
|
|
or any(len(line.split("\t")) != 5 for line in summary_lines[1:])
|
|
):
|
|
raise RuntimeError("joint summary content is invalid")
|
|
destination = (
|
|
final / ".wgs_joint_genotyping_runs" / args.prefix / run_id
|
|
)
|
|
artifacts = {
|
|
"joint_vcf": _artifact(
|
|
workspace,
|
|
joint_vcf,
|
|
destination / joint_vcf.name,
|
|
"joint_vcf",
|
|
),
|
|
"joint_vcf_index": _artifact(
|
|
workspace,
|
|
joint_vcf_index,
|
|
destination / joint_vcf_index.name,
|
|
"joint_vcf_index",
|
|
),
|
|
"joint_summary": _artifact(
|
|
workspace,
|
|
summary,
|
|
destination / summary.name,
|
|
"joint_summary",
|
|
),
|
|
"joint_log": _artifact(
|
|
workspace,
|
|
log_path,
|
|
destination / log_path.name,
|
|
"joint_log",
|
|
),
|
|
}
|
|
manifest = {
|
|
"manifest_version": "1.0",
|
|
"stage": STAGE,
|
|
"status": "success",
|
|
"prefix": args.prefix,
|
|
"job_id": args.job_id,
|
|
"run_id": run_id,
|
|
"start_time": started_at,
|
|
"workspace": str(workspace),
|
|
"wgs_sample_set_id": sample_set_id,
|
|
"wgs_reference_id": reference_id,
|
|
"upstream_wgs_analysis_ids": upstream_ids,
|
|
"wgs_analysis_id": analysis_id,
|
|
"parameters": parameters,
|
|
"software": {
|
|
"gatk": {
|
|
"path": str(_executable_path(args.gatk)),
|
|
"version": version,
|
|
}
|
|
},
|
|
"metrics": {
|
|
"input_sample_ids": expected_samples,
|
|
"output_sample_ids": expected_samples,
|
|
"total_samples": len(expected_samples),
|
|
},
|
|
"artifacts": artifacts,
|
|
}
|
|
manifest_work = attempt / "manifest.json"
|
|
_atomic_json(manifest_work, manifest)
|
|
if not _publish(
|
|
bundle, destination, manifest_work, authority, started_at
|
|
):
|
|
raise RuntimeError(
|
|
"stale combine attempt cannot replace a newer successful run"
|
|
)
|
|
return 0
|
|
except Exception as exc:
|
|
_publish_failure(authority, args, run_id, started_at, str(exc))
|
|
print(f"{STAGE} failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
return run(build_parser().parse_args(argv))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|