115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Thread scaling of ``simurg genotype`` (plan §18 P9 gate: ≥ 20 samples/min/core on GSA-size arrays, linear
|
|
scaling to 32 threads).
|
|
|
|
For every worker count the same N samples are genotyped (``--no-cohort-qc``); the per-sample phase is timed from
|
|
the run log ("genotyping … sample(s)" → "done"), so manifest loading is reported separately and does not blur
|
|
the scaling. Writes a JSON summary.
|
|
|
|
python bench/scaling.py --bpm M.bpm --egt C.egt --idats DIR --out OUT --threads 1,2,4,8 [--samples N]
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import datetime as dt
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
def phase_seconds(log: Path) -> tuple[float, float]:
|
|
"""(manifest+setup seconds, per-sample phase seconds) from a genotype run log."""
|
|
ts: dict[str, dt.datetime] = {}
|
|
for line in log.read_text().splitlines():
|
|
rec = json.loads(line)
|
|
msg = rec.get("msg") or ""
|
|
t = dt.datetime.fromisoformat(rec["ts"].replace("Z", "+00:00"))
|
|
ts.setdefault("first", t)
|
|
if msg.startswith("genotyping "):
|
|
ts["start"] = t
|
|
elif msg.startswith("done:"):
|
|
ts["done"] = t
|
|
return (ts["start"] - ts["first"]).total_seconds(), (ts["done"] - ts["start"]).total_seconds()
|
|
|
|
|
|
def subset(idats: Path, n: int, dest: Path) -> Path:
|
|
grn = sorted(p for p in idats.rglob("*_Grn.idat*"))
|
|
if n <= 0 or n >= len(grn):
|
|
return idats
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
for g in grn[:n]:
|
|
d = dest / g.parent.name
|
|
d.mkdir(exist_ok=True)
|
|
for p in (g, g.with_name(g.name.replace("_Grn.", "_Red."))):
|
|
link = d / p.name
|
|
if not link.exists():
|
|
link.symlink_to(p.resolve())
|
|
return dest
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--bpm", required=True)
|
|
ap.add_argument("--egt", required=True)
|
|
ap.add_argument("--idats", required=True, type=Path)
|
|
ap.add_argument("--out", required=True, type=Path)
|
|
ap.add_argument("--threads", default="1,2,4")
|
|
ap.add_argument("--samples", type=int, default=0, help="use the first N samples (0 = all)")
|
|
ap.add_argument("--simurg", default=str(Path(sys.executable).with_name("simurg")))
|
|
ap.add_argument("--pure-python", action="store_true", help="run with SIMURG_PURE_PYTHON=1")
|
|
ap.add_argument("--work", type=Path, help="folder for the GTC outputs (e.g. a tmpfs); default: --out")
|
|
ap.add_argument("--clean", action="store_true", help="delete each run's GTCs after timing it")
|
|
a = ap.parse_args()
|
|
a.out.mkdir(parents=True, exist_ok=True)
|
|
idats = subset(a.idats, a.samples, a.out / "idats")
|
|
n = len(list(idats.rglob("*_Grn.idat*")))
|
|
env = {**os.environ, "SIMURG_PURE_PYTHON": "1" if a.pure_python else "0"}
|
|
runs = []
|
|
for t in (int(x) for x in a.threads.split(",")):
|
|
out = (a.work or a.out) / f"t{t}"
|
|
t0 = time.time()
|
|
cmd = [a.simurg, "genotype", "--bpm", a.bpm, "--egt", a.egt, "--idats", str(idats), "--out", str(out),
|
|
"--threads", str(t), "--no-cohort-qc", "--autocall-date", "1/1/2026 12:00:00 AM",
|
|
"--quiet", "--force"] # fmt: skip
|
|
rc = subprocess.run(cmd, env=env).returncode
|
|
wall = time.time() - t0
|
|
setup, phase = phase_seconds(out / "logs" / "genotype.log")
|
|
if a.clean:
|
|
shutil.rmtree(out / "samples", ignore_errors=True)
|
|
runs.append(
|
|
{
|
|
"threads": t,
|
|
"rc": rc,
|
|
"wall_s": round(wall, 1),
|
|
"setup_s": round(setup, 1),
|
|
"per_sample_phase_s": round(phase, 2),
|
|
"samples_per_min": round(60 * n / phase, 1) if phase else None,
|
|
"samples_per_min_per_core": round(60 * n / phase / t, 1) if phase else None,
|
|
}
|
|
)
|
|
print(json.dumps(runs[-1]), flush=True)
|
|
base = runs[0]["per_sample_phase_s"] * runs[0]["threads"]
|
|
for r in runs:
|
|
r["speedup_vs_first"] = round(base / r["per_sample_phase_s"], 2) if r["per_sample_phase_s"] else None
|
|
r["efficiency"] = round(r["speedup_vs_first"] / r["threads"], 3) if r["speedup_vs_first"] else None
|
|
from simurg._core import backend
|
|
|
|
summary = {
|
|
"samples": n,
|
|
"backend": "python" if a.pure_python else backend(),
|
|
"cpu_count": os.cpu_count(),
|
|
"machine": os.uname().machine,
|
|
"runs": runs,
|
|
}
|
|
(a.out / "scaling.json").write_text(json.dumps(summary, indent=1))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|