431 lines
15 KiB
Python
431 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Stage 04 WGS duplicate removal from an authoritative alignment Manifest."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from datetime import datetime, timezone
|
|
import fcntl
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import uuid
|
|
from typing import Any
|
|
|
|
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_markdup"
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Remove duplicates from authoritative WGS alignments"
|
|
)
|
|
parser.add_argument("--workspace", required=True)
|
|
parser.add_argument("--job-id", required=True)
|
|
parser.add_argument("--prefix", required=True)
|
|
parser.add_argument("--align-manifest", required=True)
|
|
parser.add_argument("--gatk", required=True)
|
|
parser.add_argument("--samtools", required=True)
|
|
parser.add_argument("--max-jobs", required=True, type=int)
|
|
parser.add_argument("--remove-duplicates", required=True, choices=("true",))
|
|
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, *, 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("wgs_align authority Manifest is unreadable") from exc
|
|
if (
|
|
not isinstance(manifest, dict)
|
|
or manifest.get("stage") != "wgs_align"
|
|
or manifest.get("status") != "success"
|
|
):
|
|
raise ValueError("wgs_align 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("wgs_align authority Manifest identity mismatch")
|
|
return manifest
|
|
|
|
|
|
def _checked(command: list[str], name: str) -> subprocess.CompletedProcess[str]:
|
|
result = subprocess.run(command, text=True, capture_output=True, check=False)
|
|
if result.returncode:
|
|
raise RuntimeError(
|
|
f"{name} failed with exit code {result.returncode}: {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 _samtools_version(samtools: str) -> str:
|
|
result = _checked([samtools, "--version"], "samtools version check")
|
|
lines = (result.stdout or result.stderr).strip().splitlines()
|
|
if not lines:
|
|
raise RuntimeError("samtools version check produced no version")
|
|
return lines[0]
|
|
|
|
|
|
def _read_groups(header: str, sample_id: str) -> tuple[tuple[tuple[str, str], ...], ...]:
|
|
groups = []
|
|
for line in header.splitlines():
|
|
if not line.startswith("@RG\t"):
|
|
continue
|
|
fields = {}
|
|
for field in line.split("\t")[1:]:
|
|
if ":" in field:
|
|
key, value = field.split(":", 1)
|
|
fields[key] = value
|
|
groups.append(tuple(sorted(fields.items())))
|
|
if not groups or any(dict(group).get("SM") != sample_id for group in groups):
|
|
raise RuntimeError(f"sample {sample_id}: BAM read group sample identity mismatch")
|
|
return tuple(sorted(groups))
|
|
|
|
|
|
def _validate_metrics(path: Path, sample_id: str) -> None:
|
|
if not path.is_file() or path.stat().st_size <= 0:
|
|
raise RuntimeError(f"sample {sample_id}: markdup metrics are missing or empty")
|
|
lines = [
|
|
line
|
|
for line in path.read_text(encoding="utf-8", errors="strict").splitlines()
|
|
if line and not line.startswith("#")
|
|
]
|
|
required = {
|
|
"LIBRARY",
|
|
"UNPAIRED_READS_EXAMINED",
|
|
"READ_PAIRS_EXAMINED",
|
|
"PERCENT_DUPLICATION",
|
|
}
|
|
if len(lines) < 2:
|
|
raise RuntimeError(f"sample {sample_id}: markdup metrics content is invalid")
|
|
header = lines[0].split("\t")
|
|
values = lines[1].split("\t")
|
|
if not required.issubset(header) or len(values) != len(header):
|
|
raise RuntimeError(f"sample {sample_id}: markdup metrics content is invalid")
|
|
row = dict(zip(header, values))
|
|
try:
|
|
int(row["UNPAIRED_READS_EXAMINED"])
|
|
int(row["READ_PAIRS_EXAMINED"])
|
|
float(row["PERCENT_DUPLICATION"])
|
|
except (KeyError, ValueError) as exc:
|
|
raise RuntimeError(
|
|
f"sample {sample_id}: markdup metrics content is invalid"
|
|
) from exc
|
|
|
|
|
|
def _run_sample(
|
|
sample_id: str, input_bam: Path, bundle: Path, args: argparse.Namespace
|
|
) -> dict[str, Path]:
|
|
base = f"{args.prefix}.{STAGE}.{sample_id}"
|
|
bam = bundle / f"{base}.markdup.bam"
|
|
bai = bundle / f"{base}.markdup.bam.bai"
|
|
metrics = bundle / f"{base}.markdup.metrics.txt"
|
|
log_path = bundle / f"{base}.markdup.log"
|
|
input_header = _checked(
|
|
[args.samtools, "view", "-H", str(input_bam)],
|
|
f"sample {sample_id} input BAM header",
|
|
).stdout
|
|
input_groups = _read_groups(input_header, sample_id)
|
|
command = [
|
|
args.gatk,
|
|
"MarkDuplicates",
|
|
f"INPUT={input_bam}",
|
|
f"OUTPUT={bam}",
|
|
f"METRICS_FILE={metrics}",
|
|
"REMOVE_DUPLICATES=true",
|
|
"CREATE_INDEX=true",
|
|
"VALIDATION_STRINGENCY=LENIENT",
|
|
]
|
|
with log_path.open("w", encoding="utf-8") as log:
|
|
log.write("GATK argv: " + json.dumps(command) + "\n")
|
|
log.flush()
|
|
result = subprocess.run(
|
|
command,
|
|
stdout=log,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
log.write(f"GATK MarkDuplicates exit={result.returncode}\n")
|
|
if result.returncode:
|
|
raise RuntimeError(
|
|
f"sample {sample_id}: GATK MarkDuplicates failed with exit code "
|
|
f"{result.returncode}; see {log_path}"
|
|
)
|
|
alternate_bai = bam.with_suffix(".bai")
|
|
if alternate_bai != bai and alternate_bai.is_file() and not bai.exists():
|
|
os.replace(alternate_bai, bai)
|
|
if not bam.is_file() or bam.stat().st_size <= 0:
|
|
raise RuntimeError(f"sample {sample_id}: markdup BAM is missing or empty")
|
|
_checked([args.samtools, "quickcheck", str(bam)], f"sample {sample_id} quickcheck")
|
|
if not bai.is_file() or bai.stat().st_size <= 0:
|
|
raise RuntimeError(f"sample {sample_id}: markdup BAM index is missing or empty")
|
|
_checked([args.samtools, "idxstats", str(bam)], f"sample {sample_id} index validation")
|
|
output_header = _checked(
|
|
[args.samtools, "view", "-H", str(bam)], f"sample {sample_id} output BAM header"
|
|
).stdout
|
|
if _read_groups(output_header, sample_id) != input_groups:
|
|
raise RuntimeError(f"sample {sample_id}: BAM read group conservation failed")
|
|
_validate_metrics(metrics, sample_id)
|
|
return {
|
|
"markdup_bam": bam,
|
|
"markdup_bam_index": bai,
|
|
"markdup_metrics": metrics,
|
|
"markdup_log": log_path,
|
|
}
|
|
|
|
|
|
def _artifact(
|
|
workspace: Path,
|
|
source: Path,
|
|
destination: Path,
|
|
role: str,
|
|
sample_id: 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),
|
|
"sample_id": sample_id,
|
|
}
|
|
|
|
|
|
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 _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 _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 run(args: argparse.Namespace) -> int:
|
|
workspace = Path(args.workspace).resolve()
|
|
run_id = uuid.uuid4().hex
|
|
started_at = datetime.now(timezone.utc).isoformat()
|
|
try:
|
|
_safe_component(args.prefix, "prefix")
|
|
except ValueError as exc:
|
|
print(f"{STAGE} failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
final = workspace / "final"
|
|
final.mkdir(parents=True, exist_ok=True)
|
|
authority = final / f"{args.prefix}.{STAGE}.manifest.json"
|
|
try:
|
|
if args.max_jobs < 1:
|
|
raise ValueError("max-jobs must be positive")
|
|
header = _authority_header(
|
|
args.align_manifest,
|
|
workspace=workspace,
|
|
prefix=args.prefix,
|
|
job_id=args.job_id,
|
|
)
|
|
sample_set_id = str(header.get("wgs_sample_set_id") or "")
|
|
reference_id = str(header.get("wgs_reference_id") or "")
|
|
upstream_analysis_id = str(header.get("wgs_analysis_id") or "")
|
|
if not sample_set_id or not reference_id or not upstream_analysis_id:
|
|
raise ValueError("upstream scientific identities are missing")
|
|
resolved = resolve_wgs_artifacts(
|
|
args.align_manifest,
|
|
"wgs_align",
|
|
("sorted_bam", "sorted_bam_index"),
|
|
sample_set_id,
|
|
reference_id,
|
|
upstream_analysis_id,
|
|
)
|
|
bam_by_sample = {item.sample_id: item for item in resolved["sorted_bam"]}
|
|
index_by_sample = {
|
|
item.sample_id: item for item in resolved["sorted_bam_index"]
|
|
}
|
|
expected_samples = set(bam_by_sample)
|
|
if not expected_samples or expected_samples != set(index_by_sample):
|
|
raise RuntimeError("authoritative alignment input sample sets do not match")
|
|
version = _tool_version(args.gatk)
|
|
samtools_version = _samtools_version(args.samtools)
|
|
parameters = {
|
|
"gatk": args.gatk,
|
|
"samtools": args.samtools,
|
|
"max_jobs": args.max_jobs,
|
|
"remove_duplicates": True,
|
|
"create_index": True,
|
|
"validation_stringency": "LENIENT",
|
|
}
|
|
consumed = [
|
|
item.sha256 for group in resolved.values() for item in group
|
|
]
|
|
analysis_id = build_wgs_analysis_id(
|
|
STAGE,
|
|
consumed,
|
|
parameters,
|
|
{"gatk": version, "samtools": samtools_version},
|
|
method="gatk_markduplicates",
|
|
)
|
|
attempt = workspace / ".attempts" / run_id
|
|
bundle = attempt / "bundle"
|
|
bundle.mkdir(parents=True, exist_ok=False)
|
|
completed: dict[str, dict[str, Path]] = {}
|
|
failures: list[str] = []
|
|
with ThreadPoolExecutor(max_workers=args.max_jobs) as pool:
|
|
futures = {
|
|
pool.submit(
|
|
_run_sample, sample, bam_by_sample[sample].path, bundle, args
|
|
): sample
|
|
for sample in sorted(expected_samples)
|
|
}
|
|
for future in as_completed(futures):
|
|
sample = futures[future]
|
|
try:
|
|
completed[sample] = future.result()
|
|
except Exception as exc:
|
|
failures.append(str(exc))
|
|
if failures or set(completed) != expected_samples:
|
|
raise RuntimeError(
|
|
"; ".join(failures)
|
|
or "markdup input/output sample-set conservation failed"
|
|
)
|
|
destination = final / ".wgs_markdup_runs" / args.prefix / run_id
|
|
artifacts = {}
|
|
for sample in sorted(completed):
|
|
for role, source in completed[sample].items():
|
|
artifacts[f"{role}.{sample}"] = _artifact(
|
|
workspace,
|
|
source,
|
|
destination / source.name,
|
|
role,
|
|
sample,
|
|
)
|
|
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_id": upstream_analysis_id,
|
|
"wgs_analysis_id": analysis_id,
|
|
"parameters": parameters,
|
|
"software": {
|
|
"gatk": {"path": args.gatk, "version": version},
|
|
"samtools": {
|
|
"path": args.samtools,
|
|
"version": samtools_version,
|
|
},
|
|
},
|
|
"metrics": {
|
|
"input_sample_ids": sorted(expected_samples),
|
|
"output_sample_ids": sorted(completed),
|
|
"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 WGS markdup 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())
|