414 lines
14 KiB
Python
414 lines
14 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from threading import Barrier, Thread
|
||
|
|
import zipfile
|
||
|
|
|
||
|
|
import pysam
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
|
||
|
|
SAMPLES = "wgs-samples:" + "a" * 64
|
||
|
|
REFERENCE = "wgs-reference:" + "b" * 64
|
||
|
|
FILTER_ANALYSIS = "wgs-analysis:" + "c" * 64
|
||
|
|
SITE_KEYS = [
|
||
|
|
("chr1", 10, "A", "C"),
|
||
|
|
("chr1", 20, "G", "T"),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _record(workspace: Path, path: Path, role: str) -> dict:
|
||
|
|
return {
|
||
|
|
"artifact": role,
|
||
|
|
"role": role,
|
||
|
|
"path": str(path.relative_to(workspace)),
|
||
|
|
"full_path": str(path.resolve()),
|
||
|
|
"size": path.stat().st_size,
|
||
|
|
"size_bytes": path.stat().st_size,
|
||
|
|
"sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _write_bgzip_vcf(
|
||
|
|
path: Path,
|
||
|
|
*,
|
||
|
|
samples: tuple[str, ...] = ("102", "108", "112"),
|
||
|
|
site_keys: list[tuple[str, int, str, str]] | None = None,
|
||
|
|
) -> None:
|
||
|
|
plain = path.with_suffix("").with_suffix("")
|
||
|
|
with plain.open("w", encoding="utf-8") as handle:
|
||
|
|
handle.write("##fileformat=VCFv4.2\n")
|
||
|
|
handle.write("##contig=<ID=chr1,length=1000>\n")
|
||
|
|
handle.write(
|
||
|
|
'##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">\n'
|
||
|
|
)
|
||
|
|
handle.write(
|
||
|
|
"#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t"
|
||
|
|
+ "\t".join(samples)
|
||
|
|
+ "\n"
|
||
|
|
)
|
||
|
|
for chrom, pos, ref, alt in site_keys or SITE_KEYS:
|
||
|
|
handle.write(
|
||
|
|
f"{chrom}\t{pos}\t.\t{ref}\t{alt}\t50\tPASS\t.\tGT\t"
|
||
|
|
+ "\t".join("./." if index == 1 else "0/1" for index, _ in enumerate(samples))
|
||
|
|
+ "\n"
|
||
|
|
)
|
||
|
|
pysam.tabix_compress(str(plain), str(path), force=True)
|
||
|
|
pysam.tabix_index(str(path), preset="vcf", force=True)
|
||
|
|
plain.unlink()
|
||
|
|
|
||
|
|
|
||
|
|
def _filter_authority(
|
||
|
|
workspace: Path,
|
||
|
|
*,
|
||
|
|
method: str = "gatk",
|
||
|
|
samples: tuple[str, ...] = ("102", "108", "112"),
|
||
|
|
) -> Path:
|
||
|
|
bundle = (
|
||
|
|
workspace
|
||
|
|
/ "final"
|
||
|
|
/ ".wgs_variant_filter_runs"
|
||
|
|
/ "Cohort"
|
||
|
|
/ method
|
||
|
|
/ ("8" * 32)
|
||
|
|
)
|
||
|
|
bundle.mkdir(parents=True)
|
||
|
|
vcf = bundle / "Cohort.wgs_variant_filter.cohort.pass.snp.vcf.gz"
|
||
|
|
_write_bgzip_vcf(vcf, samples=samples)
|
||
|
|
index = Path(f"{vcf}.tbi")
|
||
|
|
report = bundle / "Cohort.wgs_variant_filter.cohort.report.md"
|
||
|
|
report.write_text(f"method: {method}\n", encoding="utf-8")
|
||
|
|
authority = (
|
||
|
|
workspace
|
||
|
|
/ "final"
|
||
|
|
/ f"Cohort.wgs_variant_filter.{method}.manifest.json"
|
||
|
|
)
|
||
|
|
authority.write_text(
|
||
|
|
json.dumps(
|
||
|
|
{
|
||
|
|
"manifest_version": "1.0",
|
||
|
|
"stage": "wgs_variant_filter",
|
||
|
|
"status": "success",
|
||
|
|
"prefix": "Cohort",
|
||
|
|
"job_id": "job-1",
|
||
|
|
"run_id": "8" * 32,
|
||
|
|
"workspace": str(workspace.resolve()),
|
||
|
|
"wgs_sample_set_id": SAMPLES,
|
||
|
|
"wgs_reference_id": REFERENCE,
|
||
|
|
"wgs_analysis_id": FILTER_ANALYSIS,
|
||
|
|
"method": method,
|
||
|
|
"parameters": {"method": method},
|
||
|
|
"metrics": {
|
||
|
|
"input_sample_ids": list(samples),
|
||
|
|
"output_sample_ids": list(samples),
|
||
|
|
},
|
||
|
|
"artifacts": {
|
||
|
|
"pass_snp_vcf": _record(
|
||
|
|
workspace, vcf, "pass_snp_vcf"
|
||
|
|
),
|
||
|
|
"pass_snp_vcf_index": _record(
|
||
|
|
workspace, index, "pass_snp_vcf_index"
|
||
|
|
),
|
||
|
|
"variant_filter_report": _record(
|
||
|
|
workspace, report, "variant_filter_report"
|
||
|
|
),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
),
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
return authority
|
||
|
|
|
||
|
|
|
||
|
|
def _java(tmp_path: Path, version: str = "23-test") -> Path:
|
||
|
|
java = tmp_path / "java"
|
||
|
|
java.write_text(
|
||
|
|
"""#!/home/xia-lab/miniconda3/envs/test4py311/bin/python
|
||
|
|
import gzip, json, os, pathlib, sys, tempfile
|
||
|
|
import pysam
|
||
|
|
args = sys.argv[1:]
|
||
|
|
if args == ['-version']:
|
||
|
|
print('openjdk version "JAVA_VERSION"', file=sys.stderr)
|
||
|
|
raise SystemExit(0)
|
||
|
|
calls = os.environ.get('WGS_IMPUTE_CALLS')
|
||
|
|
if calls:
|
||
|
|
with open(calls, 'a', encoding='utf-8') as handle:
|
||
|
|
handle.write(json.dumps(args) + '\\n')
|
||
|
|
if os.environ.get('FAIL_POINT') == 'beagle':
|
||
|
|
raise SystemExit(23)
|
||
|
|
def value(key):
|
||
|
|
return next(item.split('=', 1)[1] for item in args if item.startswith(key + '='))
|
||
|
|
source = pathlib.Path(value('gt'))
|
||
|
|
out = pathlib.Path(value('out'))
|
||
|
|
with gzip.open(source, 'rt', encoding='utf-8') as handle:
|
||
|
|
lines = list(handle)
|
||
|
|
mode = os.environ.get('BEAGLE_OUTPUT', '')
|
||
|
|
if mode == 'empty':
|
||
|
|
lines = []
|
||
|
|
elif mode == 'invalid':
|
||
|
|
lines = ['not-a-vcf\\n']
|
||
|
|
elif mode == 'sample-mismatch':
|
||
|
|
lines = [line.replace('\\t112\\n', '\\t999\\n') if line.startswith('#CHROM') else line for line in lines]
|
||
|
|
elif mode == 'site-drop':
|
||
|
|
records = [line for line in lines if not line.startswith('#')]
|
||
|
|
lines = [line for line in lines if line.startswith('#')] + records[:1]
|
||
|
|
elif mode == 'site-reorder':
|
||
|
|
records = [line for line in lines if not line.startswith('#')]
|
||
|
|
lines = [line for line in lines if line.startswith('#')] + list(reversed(records))
|
||
|
|
plain = pathlib.Path(str(out) + '.vcf')
|
||
|
|
plain.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with plain.open('w', encoding='utf-8') as handle:
|
||
|
|
for line in lines:
|
||
|
|
if not line.startswith('#') and mode not in ('empty', 'invalid'):
|
||
|
|
fields = line.rstrip().split('\\t')
|
||
|
|
fields[9:] = [gt.replace('/', '|').replace('.', '0') for gt in fields[9:]]
|
||
|
|
line = '\\t'.join(fields) + '\\n'
|
||
|
|
handle.write(line)
|
||
|
|
pysam.tabix_compress(str(plain), str(out) + '.vcf.gz', force=True)
|
||
|
|
plain.unlink()
|
||
|
|
pathlib.Path(str(out) + '.log').write_text('beagle test log\\n', encoding='utf-8')
|
||
|
|
""".replace("JAVA_VERSION", version),
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
java.chmod(0o755)
|
||
|
|
return java
|
||
|
|
|
||
|
|
|
||
|
|
def _class_jar(tmp_path: Path, class_major: int) -> Path:
|
||
|
|
jar = tmp_path / f"beagle-java-{class_major - 44}.jar"
|
||
|
|
class_header = (
|
||
|
|
b"\xca\xfe\xba\xbe"
|
||
|
|
+ (0).to_bytes(2, "big")
|
||
|
|
+ class_major.to_bytes(2, "big")
|
||
|
|
)
|
||
|
|
with zipfile.ZipFile(jar, "w") as archive:
|
||
|
|
archive.writestr("main/Main.class", class_header)
|
||
|
|
return jar
|
||
|
|
|
||
|
|
|
||
|
|
def _argv(workspace: Path, authority: Path, java: Path) -> list[str]:
|
||
|
|
return [
|
||
|
|
"--workspace",
|
||
|
|
str(workspace),
|
||
|
|
"--job-id",
|
||
|
|
"job-1",
|
||
|
|
"--prefix",
|
||
|
|
"Cohort",
|
||
|
|
"--filter-manifest",
|
||
|
|
str(authority),
|
||
|
|
"--java",
|
||
|
|
str(java),
|
||
|
|
"--beagle-jar",
|
||
|
|
str(Path("/tmp/beagle.29Oct24.c8e.jar")),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _authority_path(workspace: Path) -> Path:
|
||
|
|
return workspace / "final" / "Cohort.wgs_impute.manifest.json"
|
||
|
|
|
||
|
|
|
||
|
|
def _failed(workspace: Path) -> dict:
|
||
|
|
paths = sorted((workspace / "final").glob("*wgs_impute*failed*.json"))
|
||
|
|
if paths:
|
||
|
|
path = paths[-1]
|
||
|
|
else:
|
||
|
|
path = _authority_path(workspace)
|
||
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
||
|
|
assert payload["status"] == "failed"
|
||
|
|
return payload
|
||
|
|
|
||
|
|
|
||
|
|
def test_cli_success_locks_parameters_lineage_roles_samples_sites_and_index(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.cli.run_wgs_impute import main
|
||
|
|
|
||
|
|
workspace = tmp_path / "workspace"
|
||
|
|
authority = _filter_authority(workspace, method="bcftools")
|
||
|
|
java = _java(tmp_path)
|
||
|
|
calls = tmp_path / "calls.jsonl"
|
||
|
|
monkeypatch.setenv("WGS_IMPUTE_CALLS", str(calls))
|
||
|
|
assert main(_argv(workspace, authority, java)) == 0
|
||
|
|
manifest = json.loads(_authority_path(workspace).read_text())
|
||
|
|
assert set(manifest["artifacts"]) == {
|
||
|
|
"imputed_snp_vcf",
|
||
|
|
"imputed_snp_vcf_index",
|
||
|
|
"impute_report",
|
||
|
|
"impute_log",
|
||
|
|
}
|
||
|
|
assert manifest["upstream_wgs_analysis_ids"] == {
|
||
|
|
"wgs_variant_filter": FILTER_ANALYSIS
|
||
|
|
}
|
||
|
|
assert manifest["upstream_method"] == "bcftools"
|
||
|
|
assert manifest["wgs_sample_set_id"] == SAMPLES
|
||
|
|
assert manifest["wgs_reference_id"] == REFERENCE
|
||
|
|
assert manifest["metrics"]["input_sample_ids"] == ["102", "108", "112"]
|
||
|
|
assert manifest["metrics"]["output_sample_ids"] == ["102", "108", "112"]
|
||
|
|
assert manifest["metrics"]["input_site_keys"] == [
|
||
|
|
["chr1", 10, "A", "C"],
|
||
|
|
["chr1", 20, "G", "T"],
|
||
|
|
]
|
||
|
|
assert manifest["metrics"]["output_site_keys"] == manifest["metrics"]["input_site_keys"]
|
||
|
|
assert manifest["parameters"] == {
|
||
|
|
"method": "beagle",
|
||
|
|
"java": str(java),
|
||
|
|
"beagle_jar": "/tmp/beagle.29Oct24.c8e.jar",
|
||
|
|
"memory": "8g",
|
||
|
|
"window": 40.0,
|
||
|
|
"overlap": 2.0,
|
||
|
|
"ne": 100000,
|
||
|
|
"iterations": 12,
|
||
|
|
"phase_states": 280,
|
||
|
|
"imp_states": 1600,
|
||
|
|
"cluster": 0.005,
|
||
|
|
}
|
||
|
|
command = json.loads(calls.read_text().splitlines()[0])
|
||
|
|
assert command == [
|
||
|
|
"-Xmx8g",
|
||
|
|
"-jar",
|
||
|
|
"/tmp/beagle.29Oct24.c8e.jar",
|
||
|
|
f"gt={json.loads(authority.read_text())['artifacts']['pass_snp_vcf']['full_path']}",
|
||
|
|
"out=" + str(workspace / ".attempts" / manifest["run_id"] / "beagle"),
|
||
|
|
"window=40.0",
|
||
|
|
"overlap=2.0",
|
||
|
|
"ne=100000",
|
||
|
|
"iterations=12",
|
||
|
|
"phase-states=280",
|
||
|
|
"imp-states=1600",
|
||
|
|
"cluster=0.005",
|
||
|
|
]
|
||
|
|
output = Path(manifest["artifacts"]["imputed_snp_vcf"]["full_path"])
|
||
|
|
with pysam.VariantFile(str(output)) as handle:
|
||
|
|
assert [record.pos for record in handle.fetch("chr1", 0, 1000)] == [10, 20]
|
||
|
|
|
||
|
|
|
||
|
|
def test_cli_rejects_java_older_than_beagle_class_version(
|
||
|
|
tmp_path: Path
|
||
|
|
) -> None:
|
||
|
|
from apps.cli.run_wgs_impute import main
|
||
|
|
|
||
|
|
workspace = tmp_path / "workspace"
|
||
|
|
java = _java(tmp_path, version="11.0.13")
|
||
|
|
argv = _argv(workspace, _filter_authority(workspace), java)
|
||
|
|
argv[-1] = str(_class_jar(tmp_path, class_major=65))
|
||
|
|
|
||
|
|
assert main(argv) == 1
|
||
|
|
error = _failed(workspace)["errors"][0]
|
||
|
|
assert "requires Java 21" in error
|
||
|
|
assert "configured Java is 11" in error
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("method", ["gatk", "bcftools"])
|
||
|
|
def test_cli_accepts_either_legal_stage10_method(
|
||
|
|
tmp_path: Path, method: str
|
||
|
|
) -> None:
|
||
|
|
from apps.cli.run_wgs_impute import main
|
||
|
|
|
||
|
|
workspace = tmp_path / method
|
||
|
|
assert main(
|
||
|
|
_argv(workspace, _filter_authority(workspace, method=method), _java(tmp_path))
|
||
|
|
) == 0
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
("failure", "message"),
|
||
|
|
[
|
||
|
|
("beagle", "Beagle failed with exit code 23"),
|
||
|
|
("sample-mismatch", "sample header mismatch"),
|
||
|
|
("site-drop", "site keys are not conserved"),
|
||
|
|
("site-reorder", "building of index"),
|
||
|
|
("empty", "invalid"),
|
||
|
|
("invalid", "invalid"),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_cli_rejects_tool_and_postflight_failures_with_empty_artifacts(
|
||
|
|
tmp_path: Path, monkeypatch, failure: str, message: str
|
||
|
|
) -> None:
|
||
|
|
from apps.cli.run_wgs_impute import main
|
||
|
|
|
||
|
|
workspace = tmp_path / "workspace"
|
||
|
|
upstream = _filter_authority(workspace)
|
||
|
|
if failure == "beagle":
|
||
|
|
monkeypatch.setenv("FAIL_POINT", failure)
|
||
|
|
else:
|
||
|
|
monkeypatch.setenv("BEAGLE_OUTPUT", failure)
|
||
|
|
assert main(_argv(workspace, upstream, _java(tmp_path))) == 1
|
||
|
|
payload = _failed(workspace)
|
||
|
|
assert payload["artifacts"] == {}
|
||
|
|
assert message.lower() in payload["errors"][0].lower()
|
||
|
|
|
||
|
|
|
||
|
|
def test_cli_index_failure_is_not_overwritten_by_later_validation(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.cli import run_wgs_impute
|
||
|
|
|
||
|
|
workspace = tmp_path / "workspace"
|
||
|
|
upstream = _filter_authority(workspace)
|
||
|
|
|
||
|
|
def fail_index(*_args, **_kwargs):
|
||
|
|
raise OSError("controlled index failure")
|
||
|
|
|
||
|
|
monkeypatch.setattr(run_wgs_impute, "_tabix_index", fail_index)
|
||
|
|
assert run_wgs_impute.main(_argv(workspace, upstream, _java(tmp_path))) == 1
|
||
|
|
payload = _failed(workspace)
|
||
|
|
assert payload["artifacts"] == {}
|
||
|
|
assert "controlled index failure" in payload["errors"][0]
|
||
|
|
|
||
|
|
|
||
|
|
def test_failed_rerun_preserves_old_authority_and_bundle(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.cli.run_wgs_impute import main
|
||
|
|
|
||
|
|
workspace = tmp_path / "workspace"
|
||
|
|
upstream = _filter_authority(workspace)
|
||
|
|
argv = _argv(workspace, upstream, _java(tmp_path))
|
||
|
|
assert main(argv) == 0
|
||
|
|
authority = _authority_path(workspace)
|
||
|
|
before = authority.read_bytes()
|
||
|
|
manifest = json.loads(before)
|
||
|
|
bundle_hashes = {
|
||
|
|
role: hashlib.sha256(Path(item["full_path"]).read_bytes()).hexdigest()
|
||
|
|
for role, item in manifest["artifacts"].items()
|
||
|
|
}
|
||
|
|
monkeypatch.setenv("FAIL_POINT", "beagle")
|
||
|
|
assert main(argv) == 1
|
||
|
|
assert authority.read_bytes() == before
|
||
|
|
assert {
|
||
|
|
role: hashlib.sha256(Path(item["full_path"]).read_bytes()).hexdigest()
|
||
|
|
for role, item in manifest["artifacts"].items()
|
||
|
|
} == bundle_hashes
|
||
|
|
|
||
|
|
|
||
|
|
def test_concurrent_success_publishes_one_complete_authority(
|
||
|
|
tmp_path: Path, monkeypatch
|
||
|
|
) -> None:
|
||
|
|
from apps.cli.run_wgs_impute import main
|
||
|
|
|
||
|
|
workspace = tmp_path / "workspace"
|
||
|
|
argv = _argv(workspace, _filter_authority(workspace), _java(tmp_path))
|
||
|
|
barrier = Barrier(2)
|
||
|
|
original = __import__(
|
||
|
|
"apps.cli.run_wgs_impute", fromlist=["_run_beagle"]
|
||
|
|
)._run_beagle
|
||
|
|
|
||
|
|
def synchronized(*args, **kwargs):
|
||
|
|
barrier.wait()
|
||
|
|
return original(*args, **kwargs)
|
||
|
|
|
||
|
|
monkeypatch.setattr("apps.cli.run_wgs_impute._run_beagle", synchronized)
|
||
|
|
results: list[int] = []
|
||
|
|
threads = [Thread(target=lambda: results.append(main(argv))) for _ in range(2)]
|
||
|
|
for thread in threads:
|
||
|
|
thread.start()
|
||
|
|
for thread in threads:
|
||
|
|
thread.join(timeout=20)
|
||
|
|
assert sorted(results) in ([0, 0], [0, 1])
|
||
|
|
manifest = json.loads(_authority_path(workspace).read_text())
|
||
|
|
assert manifest["status"] == "success"
|
||
|
|
assert all(Path(item["full_path"]).is_file() for item in manifest["artifacts"].values())
|