253 lines
13 KiB
Python
253 lines
13 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Run Stage 03 WGS paired-end alignment from authoritative Manifests."""
|
||
|
|
|
||
|
|
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 (
|
||
|
|
ResolvedWgsArtifact,
|
||
|
|
resolve_wgs_artifacts,
|
||
|
|
wgs_workspace_identity_matches,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
STAGE = "wgs_align"
|
||
|
|
|
||
|
|
|
||
|
|
def build_parser() -> argparse.ArgumentParser:
|
||
|
|
parser = argparse.ArgumentParser(description="Align authoritative WGS clean reads with BWA-MEM2")
|
||
|
|
parser.add_argument("--workspace", required=True)
|
||
|
|
parser.add_argument("--job-id", required=True)
|
||
|
|
parser.add_argument("--prefix", required=True)
|
||
|
|
parser.add_argument("--qc-manifest", required=True)
|
||
|
|
parser.add_argument("--index-manifest", required=True)
|
||
|
|
parser.add_argument("--bwa-mem2", default="bwa-mem2")
|
||
|
|
parser.add_argument("--samtools", default="samtools")
|
||
|
|
parser.add_argument("--threads", type=int, default=1)
|
||
|
|
parser.add_argument("--max-jobs", type=int, default=1)
|
||
|
|
parser.add_argument("--platform", default="illumina")
|
||
|
|
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 not wgs_workspace_identity_matches(manifest.get("workspace", ""), workspace)
|
||
|
|
):
|
||
|
|
raise ValueError(f"{stage} authority Manifest identity mismatch")
|
||
|
|
return manifest
|
||
|
|
|
||
|
|
|
||
|
|
def _tool_version(command: list[str], name: str) -> str:
|
||
|
|
result = subprocess.run(command, text=True, capture_output=True, check=False)
|
||
|
|
if result.returncode:
|
||
|
|
raise RuntimeError(f"{name} version check failed with exit code {result.returncode}: {result.stderr.strip()}")
|
||
|
|
lines = (result.stdout or result.stderr).strip().splitlines()
|
||
|
|
if not lines:
|
||
|
|
raise RuntimeError(f"{name} version check produced no version")
|
||
|
|
return lines[0]
|
||
|
|
|
||
|
|
|
||
|
|
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 _validate_header(header: str, sample_id: str, platform: str) -> None:
|
||
|
|
read_groups = []
|
||
|
|
for line in header.splitlines():
|
||
|
|
if line.startswith("@RG\t"):
|
||
|
|
fields = {}
|
||
|
|
for field in line.split("\t")[1:]:
|
||
|
|
if ":" in field:
|
||
|
|
key, value = field.split(":", 1)
|
||
|
|
fields[key] = value
|
||
|
|
read_groups.append(fields)
|
||
|
|
expected = {"ID": sample_id, "SM": sample_id, "PL": platform}
|
||
|
|
if len(read_groups) != 1 or any(read_groups[0].get(key) != value for key, value in expected.items()):
|
||
|
|
raise RuntimeError(f"sample {sample_id}: BAM read group identity mismatch")
|
||
|
|
|
||
|
|
|
||
|
|
def _run_sample(
|
||
|
|
sample_id: str, r1: Path, r2: Path, reference: Path, bundle: Path, args: argparse.Namespace,
|
||
|
|
) -> dict[str, Path]:
|
||
|
|
base = f"{args.prefix}.{STAGE}.{sample_id}"
|
||
|
|
bam = bundle / f"{base}.sorted.bam"
|
||
|
|
bai = bundle / f"{base}.sorted.bam.bai"
|
||
|
|
log_path = bundle / f"{base}.align.log"
|
||
|
|
rg = f"@RG\tID:{sample_id}\tSM:{sample_id}\tPL:{args.platform}"
|
||
|
|
bwa_command = [args.bwa_mem2, "mem", "-t", str(args.threads), "-R", rg, str(reference), str(r1), str(r2)]
|
||
|
|
sort_command = [args.samtools, "sort", "-@", str(args.threads), "-o", str(bam), "-"]
|
||
|
|
with log_path.open("w", encoding="utf-8") as log:
|
||
|
|
log.write("bwa-mem2 argv: " + json.dumps(bwa_command) + "\n")
|
||
|
|
log.write("samtools sort argv: " + json.dumps(sort_command) + "\n")
|
||
|
|
log.flush()
|
||
|
|
bwa = subprocess.Popen(bwa_command, stdout=subprocess.PIPE, stderr=log, text=True)
|
||
|
|
assert bwa.stdout is not None
|
||
|
|
sorter = subprocess.Popen(sort_command, stdin=bwa.stdout, stdout=subprocess.DEVNULL, stderr=log, text=True)
|
||
|
|
bwa.stdout.close()
|
||
|
|
sort_returncode = sorter.wait()
|
||
|
|
bwa_returncode = bwa.wait()
|
||
|
|
log.write(f"bwa-mem2 exit={bwa_returncode}; samtools sort exit={sort_returncode}\n")
|
||
|
|
if bwa_returncode:
|
||
|
|
raise RuntimeError(f"sample {sample_id}: bwa-mem2 failed with exit code {bwa_returncode}; see {log_path}")
|
||
|
|
if sort_returncode:
|
||
|
|
raise RuntimeError(f"sample {sample_id}: samtools sort failed with exit code {sort_returncode}; see {log_path}")
|
||
|
|
if not bam.is_file() or bam.stat().st_size <= 0:
|
||
|
|
raise RuntimeError(f"sample {sample_id}: sorted BAM is missing or empty")
|
||
|
|
_checked([args.samtools, "quickcheck", str(bam)], f"sample {sample_id} samtools quickcheck")
|
||
|
|
_checked([args.samtools, "index", str(bam), str(bai)], f"sample {sample_id} samtools index")
|
||
|
|
if not bai.is_file() or bai.stat().st_size <= 0:
|
||
|
|
raise RuntimeError(f"sample {sample_id}: BAM index is missing or empty")
|
||
|
|
_checked([args.samtools, "idxstats", str(bam)], f"sample {sample_id} indexed BAM validation")
|
||
|
|
header = _checked([args.samtools, "view", "-H", str(bam)], f"sample {sample_id} BAM header").stdout
|
||
|
|
_validate_header(header, sample_id, args.platform)
|
||
|
|
return {"sorted_bam": bam, "sorted_bam_index": bai, "align_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 _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"wgs_align 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.threads < 1 or args.max_jobs < 1:
|
||
|
|
raise ValueError("threads and max-jobs must be positive")
|
||
|
|
_safe_component(args.platform, "platform")
|
||
|
|
qc_header = _authority_header(args.qc_manifest, stage="wgs_qc", 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(qc_header.get("wgs_sample_set_id") or "")
|
||
|
|
reference_id = str(index_header.get("wgs_reference_id") or "")
|
||
|
|
if not sample_set_id or not reference_id:
|
||
|
|
raise ValueError("upstream scientific identities are missing")
|
||
|
|
qc = resolve_wgs_artifacts(args.qc_manifest, "wgs_qc", ("clean_reads_r1", "clean_reads_r2"), sample_set_id)
|
||
|
|
index = resolve_wgs_artifacts(
|
||
|
|
args.index_manifest, "wgs_index",
|
||
|
|
("reference_fasta", "reference_fai", "reference_dict", "bwa_mem2_index_component"),
|
||
|
|
expected_reference_id=reference_id,
|
||
|
|
)
|
||
|
|
r1_by_sample = {item.sample_id: item for item in qc["clean_reads_r1"]}
|
||
|
|
r2_by_sample = {item.sample_id: item for item in qc["clean_reads_r2"]}
|
||
|
|
expected_samples = set(r1_by_sample)
|
||
|
|
if expected_samples != set(r2_by_sample) or not expected_samples:
|
||
|
|
raise RuntimeError("authoritative QC input sample sets do not match")
|
||
|
|
reference = index["reference_fasta"][0].path
|
||
|
|
versions = {"bwa_mem2": _tool_version([args.bwa_mem2, "version"], "bwa-mem2"),
|
||
|
|
"samtools": _tool_version([args.samtools, "--version"], "samtools")}
|
||
|
|
parameters = {"bwa_mem2": args.bwa_mem2, "samtools": args.samtools, "threads": args.threads,
|
||
|
|
"max_jobs": args.max_jobs, "platform": args.platform}
|
||
|
|
consumed = [item.sha256 for group in (*qc.values(), *index.values()) for item in group]
|
||
|
|
analysis_id = build_wgs_analysis_id(STAGE, consumed, parameters, versions, method="bwa_mem2_samtools")
|
||
|
|
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, r1_by_sample[sample].path, r2_by_sample[sample].path,
|
||
|
|
reference, 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 "alignment input/output sample-set conservation failed")
|
||
|
|
destination = final / ".wgs_align_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, "wgs_analysis_id": analysis_id,
|
||
|
|
"parameters": parameters, "software": {name: {"path": getattr(args, name), "version": version} for name, version in versions.items()},
|
||
|
|
"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 alignment attempt cannot replace a newer successful run")
|
||
|
|
return 0
|
||
|
|
except Exception as exc:
|
||
|
|
failure_path = authority if not authority.exists() else final / f"{args.prefix}.{STAGE}.failed.{run_id}.manifest.json"
|
||
|
|
_atomic_json(failure_path, _failure(args, run_id, started_at, str(exc)))
|
||
|
|
print(f"wgs_align 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())
|