538 lines
17 KiB
Python
538 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Stage 07 cohort GVCF combination 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_gvcf_combine"
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Combine every authoritative Stage06 GVCF into one cohort GVCF"
|
|
)
|
|
parser.add_argument("--workspace", required=True)
|
|
parser.add_argument("--job-id", required=True)
|
|
parser.add_argument("--prefix", required=True)
|
|
parser.add_argument("--haplotype-manifest", required=True)
|
|
parser.add_argument("--index-manifest", required=True)
|
|
parser.add_argument("--gatk", default="gatk")
|
|
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 _gvcf_header(path: Path) -> 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 or "<NON_REF>" not in fields[4].split(","):
|
|
raise RuntimeError("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 not non_ref_declared
|
|
or records < 1
|
|
or not interval
|
|
):
|
|
raise RuntimeError("GVCF 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,
|
|
) -> None:
|
|
if not index.is_file() or index.stat().st_size <= 0:
|
|
raise RuntimeError(f"{label} index is missing or empty")
|
|
samples, interval = _gvcf_header(gvcf)
|
|
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:
|
|
haplotype_header = _authority_header(
|
|
args.haplotype_manifest,
|
|
stage="wgs_haplotype_calling",
|
|
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(
|
|
haplotype_header.get("wgs_sample_set_id") or ""
|
|
)
|
|
reference_id = str(haplotype_header.get("wgs_reference_id") or "")
|
|
upstream_ids = {
|
|
"wgs_haplotype_calling": str(
|
|
haplotype_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 = haplotype_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_haplotype_calling index lineage does not match "
|
|
"the current wgs_index authority"
|
|
)
|
|
haplotype = resolve_wgs_artifacts(
|
|
args.haplotype_manifest,
|
|
"wgs_haplotype_calling",
|
|
("gvcf", "gvcf_index"),
|
|
sample_set_id,
|
|
reference_id,
|
|
upstream_ids["wgs_haplotype_calling"],
|
|
)
|
|
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"
|
|
)
|
|
gvcf_by_sample = {
|
|
item.sample_id: item for item in haplotype["gvcf"]
|
|
}
|
|
tbi_by_sample = {
|
|
item.sample_id: item for item in haplotype["gvcf_index"]
|
|
}
|
|
expected_samples = sorted(gvcf_by_sample)
|
|
metrics = haplotype_header.get("metrics") or {}
|
|
if (
|
|
not expected_samples
|
|
or len(gvcf_by_sample) != len(haplotype["gvcf"])
|
|
or len(tbi_by_sample) != len(haplotype["gvcf_index"])
|
|
or set(gvcf_by_sample) != set(tbi_by_sample)
|
|
or metrics.get("input_sample_ids") != expected_samples
|
|
or metrics.get("output_sample_ids") != expected_samples
|
|
):
|
|
raise RuntimeError("upstream GVCF sample set is incomplete")
|
|
version = _tool_version(args.gatk)
|
|
parameters = {
|
|
"gatk": args.gatk,
|
|
"operation": "CombineGVCFs",
|
|
}
|
|
consumed = [
|
|
item.sha256
|
|
for group in (haplotype, index)
|
|
for values in group.values()
|
|
for item in values
|
|
]
|
|
analysis_id = build_wgs_analysis_id(
|
|
STAGE,
|
|
consumed,
|
|
parameters,
|
|
{"gatk": version},
|
|
method="gatk_combinegvcfs",
|
|
)
|
|
attempt = workspace / ".attempts" / run_id
|
|
bundle = attempt / "bundle"
|
|
bundle.mkdir(parents=True, exist_ok=False)
|
|
base = f"{args.prefix}.{STAGE}.cohort"
|
|
combined = bundle / f"{base}.combined.g.vcf.gz"
|
|
combined_index = Path(f"{combined}.tbi")
|
|
log_path = bundle / f"{base}.combine.log"
|
|
reference = index["reference_fasta"][0].path
|
|
with log_path.open("w", encoding="utf-8") as log:
|
|
for sample in expected_samples:
|
|
_validate_with_index(
|
|
args.gatk,
|
|
reference,
|
|
gvcf_by_sample[sample].path,
|
|
tbi_by_sample[sample].path,
|
|
[sample],
|
|
log,
|
|
f"ValidateVariants input {sample}",
|
|
)
|
|
combine = [
|
|
args.gatk,
|
|
"CombineGVCFs",
|
|
"-R",
|
|
str(reference),
|
|
"-O",
|
|
str(combined),
|
|
]
|
|
for sample in expected_samples:
|
|
combine.extend(["-V", str(gvcf_by_sample[sample].path)])
|
|
_run_logged(combine, log, "CombineGVCFs")
|
|
_run_logged(
|
|
[args.gatk, "IndexFeatureFile", "-I", str(combined)],
|
|
log,
|
|
"IndexFeatureFile",
|
|
)
|
|
_validate_with_index(
|
|
args.gatk,
|
|
reference,
|
|
combined,
|
|
combined_index,
|
|
expected_samples,
|
|
log,
|
|
"ValidateVariants",
|
|
)
|
|
destination = (
|
|
final / ".wgs_gvcf_combine_runs" / args.prefix / run_id
|
|
)
|
|
artifacts = {
|
|
"combined_gvcf": _artifact(
|
|
workspace,
|
|
combined,
|
|
destination / combined.name,
|
|
"combined_gvcf",
|
|
),
|
|
"combined_gvcf_index": _artifact(
|
|
workspace,
|
|
combined_index,
|
|
destination / combined_index.name,
|
|
"combined_gvcf_index",
|
|
),
|
|
"combine_log": _artifact(
|
|
workspace,
|
|
log_path,
|
|
destination / log_path.name,
|
|
"combine_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())
|