767 lines
31 KiB
Python
767 lines
31 KiB
Python
#!/usr/bin/env python3
|
|
"""Generic manifest, execution, recovery, and result CLI for SpectraREML."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import gzip
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Mapping, NoReturn, Sequence
|
|
|
|
for _name in (
|
|
"OMP_NUM_THREADS",
|
|
"OPENBLAS_NUM_THREADS",
|
|
"MKL_NUM_THREADS",
|
|
"NUMEXPR_NUM_THREADS",
|
|
):
|
|
os.environ.setdefault(_name, "1")
|
|
|
|
try:
|
|
import numpy as np
|
|
except ImportError as exc: # pragma: no cover
|
|
raise SystemExit("NumPy is required: {}".format(exc))
|
|
|
|
|
|
MANIFEST_FORMAT = "spectra-reml-manifest-v1"
|
|
BLOCK_FORMAT = "spectra-reml-block-v1"
|
|
RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v1"
|
|
FINALIZE_FORMAT = "spectra-reml-finalize-v1"
|
|
TASK_HEADER = (
|
|
"task_index",
|
|
"task_id",
|
|
"phenotype_row",
|
|
"n_extra_covariates",
|
|
)
|
|
SUMMARY_HEADER = (
|
|
"task_index",
|
|
"task_id",
|
|
"status",
|
|
"n_fixed",
|
|
"n_extra_covariates",
|
|
"beta_offset",
|
|
"cov_offset",
|
|
"sigma_g2",
|
|
"sigma_e2",
|
|
"h2",
|
|
"logL",
|
|
"iterations",
|
|
"line_search_steps",
|
|
"grad_inf",
|
|
"error",
|
|
)
|
|
SUCCESS_STATUSES = frozenset(("converged", "converged_boundary"))
|
|
BLOCK_PATTERNS = (
|
|
"block_*.summary.tsv",
|
|
"block_*.beta.f64.bin",
|
|
"block_*.cov.f64.bin",
|
|
"block_*.complete",
|
|
)
|
|
|
|
|
|
def fail(message: str) -> NoReturn:
|
|
raise RuntimeError(message)
|
|
|
|
|
|
def log(message: str) -> None:
|
|
print(time.strftime("%Y-%m-%d %H:%M:%S"), "|", message, flush=True)
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
|
|
|
|
def resolved(path: Path) -> Path:
|
|
return path.expanduser().resolve()
|
|
|
|
|
|
def file_identity(path: Path) -> dict[str, Any]:
|
|
path = resolved(path)
|
|
stat = path.stat()
|
|
return {
|
|
"path": str(path),
|
|
"size_bytes": int(stat.st_size),
|
|
"mtime_ns": int(stat.st_mtime_ns),
|
|
}
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def canonical_sha256(value: Any) -> str:
|
|
encoded = json.dumps(
|
|
value,
|
|
ensure_ascii=False,
|
|
allow_nan=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
return hashlib.sha256(encoded).hexdigest()
|
|
|
|
|
|
def atomic_json(path: Path, value: Mapping[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_name(path.name + ".tmp.{}".format(os.getpid()))
|
|
try:
|
|
with temporary.open("w", encoding="utf-8", newline="") as handle:
|
|
json.dump(value, handle, ensure_ascii=False, allow_nan=False, indent=2)
|
|
handle.write("\n")
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary, path)
|
|
finally:
|
|
try:
|
|
temporary.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def read_json(path: Path) -> dict[str, Any]:
|
|
try:
|
|
with path.open("r", encoding="utf-8-sig") as handle:
|
|
value = json.load(handle)
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
fail("Cannot read JSON {}: {}".format(path, exc))
|
|
if not isinstance(value, dict):
|
|
fail("JSON root is not an object: {}".format(path))
|
|
return value
|
|
|
|
|
|
def require_file(path: Path, label: str, allow_empty: bool = False) -> None:
|
|
if not path.is_file():
|
|
fail("{} is missing: {}".format(label, path))
|
|
if not allow_empty and path.stat().st_size == 0:
|
|
fail("{} is empty: {}".format(label, path))
|
|
|
|
|
|
def require_size(path: Path, expected: int, label: str) -> None:
|
|
require_file(path, label, allow_empty=(expected == 0))
|
|
observed = path.stat().st_size
|
|
if observed != expected:
|
|
fail("{} size mismatch: observed {}; expected {}: {}".format(
|
|
label, observed, expected, path
|
|
))
|
|
|
|
|
|
def manifest_paths(manifest: Mapping[str, Any]) -> dict[str, Path | None]:
|
|
raw = manifest.get("paths")
|
|
if not isinstance(raw, dict):
|
|
fail("Manifest has no paths object.")
|
|
required = (
|
|
"grm_bin",
|
|
"base_x_f64",
|
|
"phenotype_f64",
|
|
"tasks_tsv",
|
|
"extra_offsets_i64",
|
|
"extra_indices_i32",
|
|
"output_directory",
|
|
)
|
|
missing = [name for name in required if name not in raw]
|
|
if missing:
|
|
fail("Manifest paths are missing: {}".format(", ".join(missing)))
|
|
result: dict[str, Path | None] = {}
|
|
for name in required + ("grm_id", "extra_covariate_f32"):
|
|
value = raw.get(name)
|
|
result[name] = Path(str(value)) if value not in (None, "") else None
|
|
return result
|
|
|
|
|
|
def read_tasks(path: Path) -> list[dict[str, str]]:
|
|
with path.open("r", encoding="utf-8-sig", newline="") as handle:
|
|
reader = csv.DictReader(handle, delimiter="\t")
|
|
if tuple(reader.fieldnames or ()) != TASK_HEADER:
|
|
fail("Task header must be exactly: {}".format("\t".join(TASK_HEADER)))
|
|
return list(reader)
|
|
|
|
|
|
def validate_manifest(path: Path) -> dict[str, Any]:
|
|
path = resolved(path)
|
|
manifest = read_json(path)
|
|
if manifest.get("format") != MANIFEST_FORMAT:
|
|
fail("Unsupported manifest format: {!r}".format(manifest.get("format")))
|
|
dimensions = manifest.get("dimensions")
|
|
if not isinstance(dimensions, dict):
|
|
fail("Manifest has no dimensions object.")
|
|
names = (
|
|
"sample_count",
|
|
"base_covariate_count",
|
|
"phenotype_row_count",
|
|
"extra_covariate_row_count",
|
|
"task_count",
|
|
"extra_index_count",
|
|
)
|
|
parsed: dict[str, int] = {}
|
|
for name in names:
|
|
value = dimensions.get(name)
|
|
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
|
fail("Invalid dimension {}={!r}".format(name, value))
|
|
parsed[name] = value
|
|
if min(parsed["sample_count"], parsed["base_covariate_count"], parsed["task_count"]) <= 0:
|
|
fail("Sample, base-covariate, and task counts must be positive.")
|
|
|
|
paths = manifest_paths(manifest)
|
|
n = parsed["sample_count"]
|
|
require_size(paths["grm_bin"], n * (n + 1) // 2 * 4, "GRM") # type: ignore[arg-type]
|
|
require_size(paths["base_x_f64"], n * parsed["base_covariate_count"] * 8, "base design") # type: ignore[arg-type]
|
|
require_size(paths["phenotype_f64"], n * parsed["phenotype_row_count"] * 8, "phenotypes") # type: ignore[arg-type]
|
|
require_size(paths["extra_offsets_i64"], (parsed["task_count"] + 1) * 8, "extra offsets") # type: ignore[arg-type]
|
|
require_size(paths["extra_indices_i32"], parsed["extra_index_count"] * 4, "extra indices") # type: ignore[arg-type]
|
|
if parsed["extra_covariate_row_count"]:
|
|
if paths["extra_covariate_f32"] is None:
|
|
fail("extra_covariate_f32 is required when its row count is nonzero.")
|
|
require_size(paths["extra_covariate_f32"], n * parsed["extra_covariate_row_count"] * 4, "extra covariates")
|
|
if paths["grm_id"] is not None:
|
|
require_file(paths["grm_id"], "GRM IDs")
|
|
with paths["grm_id"].open("r", encoding="utf-8-sig") as handle:
|
|
if sum(bool(line.strip()) for line in handle) != n:
|
|
fail("GRM ID row count differs from sample_count.")
|
|
|
|
sources = manifest.get("source_files")
|
|
if not isinstance(sources, dict) or not sources:
|
|
fail("Manifest has no source identities.")
|
|
for label, expected in sources.items():
|
|
if not isinstance(expected, dict):
|
|
fail("Malformed source identity: {}".format(label))
|
|
source = Path(str(expected.get("path", "")))
|
|
if not source.is_file():
|
|
fail("Source is missing: {}: {}".format(label, source))
|
|
observed = source.stat()
|
|
if (
|
|
observed.st_size != int(expected.get("size_bytes", -1))
|
|
or observed.st_mtime_ns != int(expected.get("mtime_ns", -1))
|
|
):
|
|
fail("Source changed since manifest creation: {}: {}".format(label, source))
|
|
|
|
controls = {
|
|
"base_x_f64": paths["base_x_f64"],
|
|
"tasks_tsv": paths["tasks_tsv"],
|
|
"extra_offsets_i64": paths["extra_offsets_i64"],
|
|
"extra_indices_i32": paths["extra_indices_i32"],
|
|
}
|
|
hashes = manifest.get("control_sha256")
|
|
if not isinstance(hashes, dict) or set(hashes) != set(controls):
|
|
fail("control_sha256 does not cover all control files.")
|
|
for name, control in controls.items():
|
|
assert control is not None
|
|
if sha256_file(control) != hashes[name]:
|
|
fail("Control checksum mismatch: {}".format(control))
|
|
|
|
tasks = read_tasks(paths["tasks_tsv"]) # type: ignore[arg-type]
|
|
if len(tasks) != parsed["task_count"]:
|
|
fail("Task count differs from manifest.")
|
|
offsets = np.fromfile(paths["extra_offsets_i64"], dtype="<i8")
|
|
indices = np.fromfile(paths["extra_indices_i32"], dtype="<i4")
|
|
if offsets[0] != 0 or np.any(offsets[1:] < offsets[:-1]) or int(offsets[-1]) != indices.size:
|
|
fail("Extra-covariate CSR offsets are invalid.")
|
|
if indices.size and (
|
|
int(indices.min()) < 0
|
|
or int(indices.max()) >= parsed["extra_covariate_row_count"]
|
|
):
|
|
fail("An extra-covariate index is outside its matrix.")
|
|
seen_ids: set[str] = set()
|
|
for expected, row in enumerate(tasks):
|
|
try:
|
|
task_index = int(row["task_index"])
|
|
phenotype_row = int(row["phenotype_row"])
|
|
extra_count = int(row["n_extra_covariates"])
|
|
except ValueError as exc:
|
|
fail("Invalid task integer at row {}: {}".format(expected, exc))
|
|
if task_index != expected:
|
|
fail("Task indices must be consecutive and zero based.")
|
|
if not row["task_id"] or row["task_id"] in seen_ids:
|
|
fail("Task IDs must be nonempty and unique.")
|
|
seen_ids.add(row["task_id"])
|
|
if phenotype_row < 0 or phenotype_row >= parsed["phenotype_row_count"]:
|
|
fail("Task phenotype row is outside its matrix.")
|
|
if extra_count != int(offsets[expected + 1] - offsets[expected]):
|
|
fail("Task extra count disagrees with CSR offsets.")
|
|
return manifest
|
|
|
|
|
|
def make_manifest(args: argparse.Namespace) -> Path:
|
|
target = resolved(args.manifest)
|
|
paths: dict[str, Path | None] = {
|
|
"grm_bin": resolved(args.grm_bin),
|
|
"grm_id": resolved(args.grm_id) if args.grm_id else None,
|
|
"base_x_f64": resolved(args.base_x),
|
|
"phenotype_f64": resolved(args.phenotypes),
|
|
"extra_covariate_f32": resolved(args.extra_covariates) if args.extra_covariates else None,
|
|
"tasks_tsv": resolved(args.tasks),
|
|
"extra_offsets_i64": resolved(args.extra_offsets),
|
|
"extra_indices_i32": resolved(args.extra_indices),
|
|
"output_directory": resolved(args.output_dir),
|
|
}
|
|
tasks = read_tasks(paths["tasks_tsv"]) # type: ignore[arg-type]
|
|
extra_index_count = paths["extra_indices_i32"].stat().st_size // 4 # type: ignore[union-attr]
|
|
source_names = ("grm_bin", "phenotype_f64")
|
|
sources = {name: file_identity(paths[name]) for name in source_names} # type: ignore[arg-type]
|
|
for name in ("grm_id", "extra_covariate_f32"):
|
|
if paths[name] is not None:
|
|
sources[name] = file_identity(paths[name]) # type: ignore[arg-type]
|
|
controls = ("base_x_f64", "tasks_tsv", "extra_offsets_i64", "extra_indices_i32")
|
|
manifest = {
|
|
"format": MANIFEST_FORMAT,
|
|
"created_utc": utc_now(),
|
|
"dimensions": {
|
|
"sample_count": args.n_samples,
|
|
"base_covariate_count": args.n_base_covariates,
|
|
"phenotype_row_count": args.n_phenotype_rows,
|
|
"extra_covariate_row_count": args.n_extra_covariate_rows,
|
|
"task_count": len(tasks),
|
|
"extra_index_count": extra_index_count,
|
|
},
|
|
"paths": {name: str(value) if value is not None else None for name, value in paths.items()},
|
|
"source_files": sources,
|
|
"control_sha256": {name: sha256_file(paths[name]) for name in controls}, # type: ignore[arg-type]
|
|
}
|
|
atomic_json(target, manifest)
|
|
validate_manifest(target)
|
|
log("manifest created: {} tasks -> {}".format(len(tasks), target))
|
|
return target
|
|
|
|
|
|
def resolve_engine(value: str) -> Path:
|
|
candidate = Path(value).expanduser()
|
|
if candidate.is_file():
|
|
return candidate.resolve()
|
|
found = shutil.which(value)
|
|
if found:
|
|
return Path(found).resolve()
|
|
fail("SpectraREML engine was not found: {}".format(value))
|
|
|
|
|
|
def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argparse.Namespace) -> list[str]:
|
|
paths = manifest_paths(manifest)
|
|
dims = manifest["dimensions"]
|
|
command = [
|
|
str(engine),
|
|
"--grm-bin", str(paths["grm_bin"]),
|
|
"--base-x", str(paths["base_x_f64"]),
|
|
"--phenotypes", str(paths["phenotype_f64"]),
|
|
"--tasks", str(paths["tasks_tsv"]),
|
|
"--extra-offsets", str(paths["extra_offsets_i64"]),
|
|
"--extra-indices", str(paths["extra_indices_i32"]),
|
|
"--out-dir", str(paths["output_directory"]),
|
|
"--n-samples", str(dims["sample_count"]),
|
|
"--n-base-covariates", str(dims["base_covariate_count"]),
|
|
"--n-phenotype-rows", str(dims["phenotype_row_count"]),
|
|
"--n-extra-covariate-rows", str(dims["extra_covariate_row_count"]),
|
|
"--block-size", str(args.block_size),
|
|
"--threads", str(args.threads),
|
|
]
|
|
if paths["grm_id"] is not None:
|
|
command.extend(("--grm-id", str(paths["grm_id"])))
|
|
if int(dims["extra_index_count"]):
|
|
command.extend(("--extra-covariates", str(paths["extra_covariate_f32"])))
|
|
option_map = (
|
|
("max_iterations", "--max-iterations"),
|
|
("line_search_max_evals", "--line-search-max-evals"),
|
|
("line_search_max_zoom", "--line-search-max-zoom"),
|
|
("gradient_abs_tol", "--gradient-abs-tol"),
|
|
("gradient_rel_tol", "--gradient-rel-tol"),
|
|
("step_rel_tol", "--step-rel-tol"),
|
|
("likelihood_rel_tol", "--likelihood-rel-tol"),
|
|
("wolfe_c1", "--wolfe-c1"),
|
|
("wolfe_c2", "--wolfe-c2"),
|
|
("initial_step", "--initial-step"),
|
|
("maximum_step", "--maximum-step"),
|
|
("rank_tol", "--rank-tol"),
|
|
("covariance_floor", "--covariance-floor"),
|
|
("boundary_h2_trigger", "--boundary-h2-trigger"),
|
|
("boundary_score_tol", "--boundary-score-tol"),
|
|
("boundary_logl_rel_tol", "--boundary-logl-rel-tol"),
|
|
)
|
|
for attribute, option in option_map:
|
|
command.extend((option, repr(getattr(args, attribute))))
|
|
for attribute, option in (("initial_sigma_e", "--initial-sigma-e"), ("initial_sigma_g", "--initial-sigma-g")):
|
|
value = getattr(args, attribute)
|
|
if value is not None:
|
|
command.extend((option, repr(value)))
|
|
if args.resume:
|
|
command.append("--resume")
|
|
elif args.force:
|
|
command.append("--overwrite")
|
|
return command
|
|
|
|
|
|
def engine_identity(engine: Path) -> dict[str, Any]:
|
|
value = file_identity(engine)
|
|
value["sha256"] = sha256_file(engine)
|
|
return value
|
|
|
|
|
|
def signature_payload(manifest: Mapping[str, Any], engine: Path, args: argparse.Namespace) -> dict[str, Any]:
|
|
excluded = {"action", "manifest", "engine", "dry_run", "output"}
|
|
options = {
|
|
key: value for key, value in vars(args).items()
|
|
if key not in excluded and key not in {"resume", "force"}
|
|
}
|
|
return {
|
|
"manifest_sha256": canonical_sha256(manifest),
|
|
"engine": engine_identity(engine),
|
|
"options": options,
|
|
}
|
|
|
|
|
|
def make_signature(manifest: Mapping[str, Any], engine: Path, args: argparse.Namespace) -> dict[str, Any]:
|
|
payload = signature_payload(manifest, engine, args)
|
|
return {
|
|
"format": RUN_SIGNATURE_FORMAT,
|
|
"created_utc": utc_now(),
|
|
"signature_sha256": canonical_sha256(payload),
|
|
"payload": payload,
|
|
}
|
|
|
|
|
|
def read_signature(path: Path) -> dict[str, Any]:
|
|
value = read_json(path)
|
|
payload = value.get("payload")
|
|
if value.get("format") != RUN_SIGNATURE_FORMAT or not isinstance(payload, dict):
|
|
fail("Invalid run signature: {}".format(path))
|
|
if value.get("signature_sha256") != canonical_sha256(payload):
|
|
fail("Run signature integrity check failed: {}".format(path))
|
|
return value
|
|
|
|
|
|
def block_files(output: Path) -> list[Path]:
|
|
values: set[Path] = set()
|
|
for pattern in BLOCK_PATTERNS:
|
|
values.update(path for path in output.glob(pattern) if path.is_file())
|
|
return sorted(values)
|
|
|
|
|
|
def validate_resume_layout(output: Path, task_count: int, block_size: int) -> None:
|
|
for marker in output.glob("block_*.complete"):
|
|
digits = marker.name[len("block_"):-len(".complete")]
|
|
if len(digits) != 6 or not digits.isdigit():
|
|
fail("Invalid block marker: {}".format(marker))
|
|
index = int(digits)
|
|
expected = min(block_size, max(0, task_count - index * block_size))
|
|
values: dict[str, str] = {}
|
|
with marker.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
fields = line.rstrip("\r\n").split("\t", 1)
|
|
if len(fields) != 2:
|
|
fail("Malformed marker: {}".format(marker))
|
|
values[fields[0]] = fields[1]
|
|
if expected <= 0 or values.get("format") != BLOCK_FORMAT or int(values.get("tasks", -1)) != expected:
|
|
fail("Block marker is incompatible with this manifest/block size: {}".format(marker))
|
|
|
|
|
|
def read_marker(path: Path) -> dict[str, str]:
|
|
values: dict[str, str] = {}
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
fields = line.rstrip("\r\n").split("\t", 1)
|
|
if len(fields) != 2 or not fields[0] or fields[0] in values:
|
|
fail("Malformed completion marker: {}".format(path))
|
|
values[fields[0]] = fields[1]
|
|
return values
|
|
|
|
|
|
def run_engine(manifest: Mapping[str, Any], args: argparse.Namespace) -> None:
|
|
if args.dry_run:
|
|
print(json.dumps(build_engine_command(Path(args.engine), manifest, args), ensure_ascii=False))
|
|
return
|
|
engine = resolve_engine(args.engine)
|
|
expected_signature = make_signature(manifest, engine, args)
|
|
paths = manifest_paths(manifest)
|
|
output = paths["output_directory"]
|
|
assert output is not None
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
signature_path = output / "run.signature.json"
|
|
markers = list(output.glob("block_*.complete"))
|
|
if args.force:
|
|
try:
|
|
signature_path.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
for path in block_files(output):
|
|
path.unlink()
|
|
atomic_json(signature_path, expected_signature)
|
|
elif args.resume:
|
|
if markers:
|
|
if not signature_path.is_file():
|
|
fail("Completed blocks have no run signature; use --force.")
|
|
observed = read_signature(signature_path)
|
|
if observed["signature_sha256"] != expected_signature["signature_sha256"]:
|
|
fail("Run signature differs from engine, manifest, or options; use --force.")
|
|
validate_resume_layout(output, int(manifest["dimensions"]["task_count"]), args.block_size)
|
|
else:
|
|
atomic_json(signature_path, expected_signature)
|
|
else:
|
|
if block_files(output):
|
|
fail("Block outputs already exist; use --resume or --force.")
|
|
atomic_json(signature_path, expected_signature)
|
|
|
|
command = build_engine_command(engine, manifest, args)
|
|
environment = os.environ.copy()
|
|
environment.update({
|
|
"OMP_NUM_THREADS": str(args.threads),
|
|
"MKL_NUM_THREADS": str(args.blas_threads),
|
|
"OPENBLAS_NUM_THREADS": str(args.blas_threads),
|
|
"OMP_DYNAMIC": "FALSE",
|
|
})
|
|
log("starting SpectraREML engine")
|
|
with (output / "engine.log").open("a" if args.resume else "w", encoding="utf-8") as handle:
|
|
handle.write("# {}\n# command={}\n".format(utc_now(), json.dumps(command, ensure_ascii=False)))
|
|
handle.flush()
|
|
completed = subprocess.run(command, stdout=handle, stderr=subprocess.STDOUT, env=environment, check=False)
|
|
atomic_json(output / "run.json", {
|
|
"finished_utc": utc_now(),
|
|
"return_code": completed.returncode,
|
|
"command": command,
|
|
"run_signature_sha256": expected_signature["signature_sha256"],
|
|
})
|
|
if completed.returncode:
|
|
fail("SpectraREML failed with exit code {}; see engine.log".format(completed.returncode))
|
|
|
|
|
|
def finalizable_signature(manifest: Mapping[str, Any]) -> dict[str, Any]:
|
|
output = manifest_paths(manifest)["output_directory"]
|
|
assert output is not None
|
|
path = output / "run.signature.json"
|
|
if not path.is_file():
|
|
fail("Cannot finalize without a run signature.")
|
|
value = read_signature(path)
|
|
if value["payload"].get("manifest_sha256") != canonical_sha256(manifest):
|
|
fail("Run signature does not belong to the current manifest.")
|
|
return value
|
|
|
|
|
|
def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
|
|
output = manifest_paths(manifest)["output_directory"]
|
|
assert output is not None
|
|
task_count = int(manifest["dimensions"]["task_count"])
|
|
results: dict[int, dict[str, Any]] = {}
|
|
for marker in sorted(output.glob("block_*.complete")):
|
|
stem = marker.name[:-len(".complete")]
|
|
summary_path = output / (stem + ".summary.tsv")
|
|
beta_path = output / (stem + ".beta.f64.bin")
|
|
cov_path = output / (stem + ".cov.f64.bin")
|
|
marker_values = read_marker(marker)
|
|
try:
|
|
declared_tasks = int(marker_values["tasks"])
|
|
beta_elements = int(marker_values["beta_elements"])
|
|
cov_elements = int(marker_values["cov_elements"])
|
|
except (KeyError, ValueError) as exc:
|
|
fail("Malformed completion counts in {}: {}".format(marker, exc))
|
|
if marker_values.get("format") != BLOCK_FORMAT or min(
|
|
declared_tasks, beta_elements, cov_elements
|
|
) < 0:
|
|
fail("Invalid completion marker: {}".format(marker))
|
|
require_size(beta_path, beta_elements * 8, "block beta")
|
|
require_size(cov_path, cov_elements * 8, "block covariance")
|
|
with summary_path.open("r", encoding="utf-8-sig", newline="") as handle:
|
|
reader = csv.DictReader(handle, delimiter="\t")
|
|
if tuple(reader.fieldnames or ()) != SUMMARY_HEADER:
|
|
fail("Unexpected summary header: {}".format(summary_path))
|
|
rows = list(reader)
|
|
if len(rows) != declared_tasks:
|
|
fail("Summary row count disagrees with marker: {}".format(marker))
|
|
beta = np.fromfile(beta_path, dtype="<f8")
|
|
cov = np.fromfile(cov_path, dtype="<f8")
|
|
next_beta_offset = 0
|
|
next_cov_offset = 0
|
|
for row in rows:
|
|
index = int(row["task_index"])
|
|
if index in results or index < 0 or index >= task_count:
|
|
fail("Duplicate/out-of-range task index: {}".format(index))
|
|
p = int(row["n_fixed"])
|
|
beta_offset, cov_offset = int(row["beta_offset"]), int(row["cov_offset"])
|
|
if beta_offset < 0 or cov_offset < 0:
|
|
if beta_offset != cov_offset:
|
|
fail("Only one output offset is negative for task {}".format(index))
|
|
beta_values: list[float] = []
|
|
cov_values: list[float] = []
|
|
else:
|
|
packed = p * (p + 1) // 2
|
|
if beta_offset != next_beta_offset or cov_offset != next_cov_offset:
|
|
fail("Non-contiguous output offsets for task {}".format(index))
|
|
if beta_offset + p > beta.size or cov_offset + packed > cov.size:
|
|
fail("Output offsets exceed binary arrays for task {}".format(index))
|
|
beta_values = beta[beta_offset:beta_offset + p].tolist()
|
|
cov_values = cov[cov_offset:cov_offset + packed].tolist()
|
|
next_beta_offset += p
|
|
next_cov_offset += packed
|
|
result: dict[str, Any] = dict(row)
|
|
result["beta_json"] = json.dumps(beta_values, separators=(",", ":"))
|
|
result["covariance_packed_lower_json"] = json.dumps(cov_values, separators=(",", ":"))
|
|
results[index] = result
|
|
if next_beta_offset != beta_elements or next_cov_offset != cov_elements:
|
|
fail("Block binary arrays contain unused elements: {}".format(stem))
|
|
missing = sorted(set(range(task_count)).difference(results))
|
|
if missing:
|
|
fail("Completed blocks are missing {} tasks.".format(len(missing)))
|
|
return [results[index] for index in range(task_count)]
|
|
|
|
|
|
def finalize(manifest: Mapping[str, Any], output: Path) -> Path:
|
|
signature = finalizable_signature(manifest)
|
|
rows = read_completed(manifest)
|
|
output = resolved(output)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = output.with_name(output.name + ".tmp.{}".format(os.getpid()))
|
|
opener = gzip.open if output.suffix == ".gz" else open
|
|
fields = list(SUMMARY_HEADER) + ["beta_json", "covariance_packed_lower_json"]
|
|
try:
|
|
with opener(temporary, "wt", encoding="utf-8", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
os.replace(temporary, output)
|
|
finally:
|
|
try:
|
|
temporary.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
counts: dict[str, int] = {}
|
|
for row in rows:
|
|
counts[row["status"]] = counts.get(row["status"], 0) + 1
|
|
atomic_json(output.with_suffix(output.suffix + ".json"), {
|
|
"format": FINALIZE_FORMAT,
|
|
"finished_utc": utc_now(),
|
|
"manifest_sha256": canonical_sha256(manifest),
|
|
"run_signature_sha256": signature["signature_sha256"],
|
|
"tasks": len(rows),
|
|
"status_counts": counts,
|
|
})
|
|
log("finalized {} tasks -> {}".format(len(rows), output))
|
|
return output
|
|
|
|
|
|
def status(manifest: Mapping[str, Any]) -> None:
|
|
output = manifest_paths(manifest)["output_directory"]
|
|
assert output is not None
|
|
counts: dict[str, int] = {}
|
|
completed = 0
|
|
for marker in output.glob("block_*.complete"):
|
|
summary = output / (marker.name[:-len(".complete")] + ".summary.tsv")
|
|
with summary.open("r", encoding="utf-8-sig", newline="") as handle:
|
|
for row in csv.DictReader(handle, delimiter="\t"):
|
|
completed += 1
|
|
counts[row.get("status", "unknown")] = counts.get(row.get("status", "unknown"), 0) + 1
|
|
print(json.dumps({
|
|
"tasks": manifest["dimensions"]["task_count"],
|
|
"completed": completed,
|
|
"status_counts": counts,
|
|
"output_directory": str(output),
|
|
}, ensure_ascii=False, indent=2))
|
|
|
|
|
|
def positive_integer(value: str) -> int:
|
|
parsed = int(value)
|
|
if parsed <= 0:
|
|
raise argparse.ArgumentTypeError("must be positive")
|
|
return parsed
|
|
|
|
|
|
def add_run_options(parser: argparse.ArgumentParser) -> None:
|
|
parser.add_argument("--engine", default="spectra_reml")
|
|
parser.add_argument("--threads", type=positive_integer, default=1)
|
|
parser.add_argument("--blas-threads", type=positive_integer, default=1)
|
|
parser.add_argument("--block-size", type=positive_integer, default=256)
|
|
parser.add_argument("--max-iterations", type=positive_integer, default=100)
|
|
parser.add_argument("--line-search-max-evals", type=positive_integer, default=48)
|
|
parser.add_argument("--line-search-max-zoom", type=positive_integer, default=48)
|
|
parser.add_argument("--gradient-abs-tol", type=float, default=1e-7)
|
|
parser.add_argument("--gradient-rel-tol", type=float, default=1e-8)
|
|
parser.add_argument("--step-rel-tol", type=float, default=1e-9)
|
|
parser.add_argument("--likelihood-rel-tol", type=float, default=1e-11)
|
|
parser.add_argument("--wolfe-c1", type=float, default=1e-4)
|
|
parser.add_argument("--wolfe-c2", type=float, default=0.9)
|
|
parser.add_argument("--initial-step", type=float, default=1.0)
|
|
parser.add_argument("--maximum-step", type=float, default=64.0)
|
|
parser.add_argument("--rank-tol", type=float, default=1e-10)
|
|
parser.add_argument("--covariance-floor", type=float, default=1e-12)
|
|
parser.add_argument("--boundary-h2-trigger", type=float, default=1e-3)
|
|
parser.add_argument("--boundary-score-tol", type=float, default=1e-8)
|
|
parser.add_argument("--boundary-logl-rel-tol", type=float, default=1e-12)
|
|
parser.add_argument("--initial-sigma-e", type=float)
|
|
parser.add_argument("--initial-sigma-g", type=float)
|
|
mode = parser.add_mutually_exclusive_group()
|
|
mode.add_argument("--resume", action="store_true")
|
|
mode.add_argument("--force", action="store_true")
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
subparsers = parser.add_subparsers(dest="action", required=True)
|
|
make = subparsers.add_parser("make-manifest")
|
|
make.add_argument("--manifest", type=Path, required=True)
|
|
make.add_argument("--grm-bin", type=Path, required=True)
|
|
make.add_argument("--grm-id", type=Path)
|
|
make.add_argument("--base-x", type=Path, required=True)
|
|
make.add_argument("--phenotypes", type=Path, required=True)
|
|
make.add_argument("--extra-covariates", type=Path)
|
|
make.add_argument("--tasks", type=Path, required=True)
|
|
make.add_argument("--extra-offsets", type=Path, required=True)
|
|
make.add_argument("--extra-indices", type=Path, required=True)
|
|
make.add_argument("--output-dir", type=Path, required=True)
|
|
make.add_argument("--n-samples", type=positive_integer, required=True)
|
|
make.add_argument("--n-base-covariates", type=positive_integer, required=True)
|
|
make.add_argument("--n-phenotype-rows", type=positive_integer, required=True)
|
|
make.add_argument("--n-extra-covariate-rows", type=int, default=0)
|
|
for action in ("validate", "status"):
|
|
command = subparsers.add_parser(action)
|
|
command.add_argument("--manifest", type=Path, required=True)
|
|
finalize_parser = subparsers.add_parser("finalize")
|
|
finalize_parser.add_argument("--manifest", type=Path, required=True)
|
|
finalize_parser.add_argument("--output", type=Path, required=True)
|
|
for action in ("run", "all"):
|
|
command = subparsers.add_parser(action)
|
|
command.add_argument("--manifest", type=Path, required=True)
|
|
if action == "all":
|
|
command.add_argument("--output", type=Path, required=True)
|
|
add_run_options(command)
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
if args.action == "make-manifest":
|
|
make_manifest(args)
|
|
return 0
|
|
manifest = validate_manifest(args.manifest)
|
|
if args.action == "validate":
|
|
print(json.dumps({
|
|
"manifest": str(resolved(args.manifest)),
|
|
"manifest_sha256": canonical_sha256(manifest),
|
|
"dimensions": manifest["dimensions"],
|
|
}, ensure_ascii=False, indent=2))
|
|
elif args.action == "status":
|
|
status(manifest)
|
|
elif args.action == "finalize":
|
|
finalize(manifest, args.output)
|
|
elif args.action == "run":
|
|
run_engine(manifest, args)
|
|
elif args.action == "all":
|
|
run_engine(manifest, args)
|
|
if not args.dry_run:
|
|
finalize(manifest, args.output)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except KeyboardInterrupt:
|
|
raise SystemExit(130)
|
|
except Exception as exc:
|
|
print("ERROR: {}".format(exc), file=sys.stderr)
|
|
raise SystemExit(1)
|