Jungfraujoch/tools/battery/results/timing_platforms/timing_table.py

#!/usr/bin/env python3
"""Cross-platform report of the rugnux timing benchmark (battery.py timing).

    python3 -I timing_table.py platforms.json [-o out.md]

platforms.json: {"title", "intro", "timing_json", "stands_out": [str], "platforms": [entry]}.
"timing_json" is a copy of tools/battery/timing.json (path relative to platforms.json) and "benchmark"
(default inhouse) the benchmark in it the platforms ran, for the data
section. An entry is {"platform" (short name + CPU), "gpu", "role", "conditions", "os", "build",
"battery_cmd", "notes": [str], "peak_mem": {set: str}, "table2_note"} plus ONE source:
  "run":     a battery run directory (results.json + manifest.json; logs under logs/<set>/ or
             work/<arm>/<set>/),
  "archive": a .tar.gz/.tar plus "member", the run directory's path inside it; the archive is unpacked
             once into unpacked/<archive name>/ next to this script (members with absolute paths, '..'
             or links are refused),
  "manual":  typed in from a report with no run directory: {"sets": {set: {"verdict", "wall",
             "mr_decision", "steps": {step: s}}}, "repro": {field: text}}.
Whatever a source does not record is printed as "not recorded". Nothing here judges: the checks are
the benchmark's own (bench_checks in results.json).
"""
import json
import os
import re
import sys
import tarfile

SETS = ["myob_x10sa", "cytc_x10sa", "4nwv"]   # the benchmark's sets, set from timing.json in main()
HERE = os.path.dirname(os.path.abspath(__file__))
NR = "not recorded"


def unpack(archive, member):
    name = re.sub(r"\.tar(\.gz)?$", "", os.path.basename(archive))
    dest = os.path.join(HERE, "unpacked", name)
    if not os.path.isdir(dest):
        os.makedirs(dest)
        with tarfile.open(archive) as t:
            safe = [m for m in t.getmembers()
                    if not m.name.startswith("/") and ".." not in m.name.split("/") and (m.isfile() or m.isdir())]
            kw = {"filter": "data"} if hasattr(tarfile, "data_filter") else {}
            t.extractall(dest, members=safe, **kw)
    return os.path.join(dest, member)


def log_dir(run, row):
    for d in (os.path.join(run, "logs", row["set"]), os.path.join(run, "work", row.get("arm", ""), row["set"])):
        if os.path.isdir(d):
            return d
    return None


def from_log(d):
    """The binary's own header lines and where the work ran."""
    out = {}
    log = d and os.path.join(d, "p_log.txt")
    if not log or not os.path.exists(log):
        return out
    text = open(log, errors="replace").read()
    for key, pat in (("version", r"\] (rugnux Version [^\n]*)"), ("build_flags", r"\] Build flags: ([^\n]*)"),
                     ("gpu", r"\] GPU: ([^\n]*)"), ("threads", r"\] Threads: (\d+)")):
        m = re.search(pat, text)
        out[key] = m.group(1).strip() if m else None
    out["scaling"] = "GPU" if "RotationScaleMerge: GPU scaling" in text else "CPU"
    m = re.search(r"Model validation: structure factors[^\n]*?([\d.]+ GB) held of a ([\d.]+ GB) budget", text)
    out["mv"] = f"GPU ({m.group(1)} of a {m.group(2)} budget)" if m else "CPU"
    rep = os.path.join(d, "p_report.txt")
    if os.path.exists(rep):
        for line in open(rep, errors="replace"):
            if line.startswith("MR_SEARCH_DEVICE="):
                out["mr_dev"] = line.split("=", 1)[1].strip()
    return out


def load(p):
    """{set: row}, {repro field: text}"""
    if "manual" in p:
        return p["manual"]["sets"], dict(p["manual"].get("repro", {}))
    run = p["run"] if "run" in p else unpack(p["archive"], p["member"])
    man = json.load(open(os.path.join(run, "manifest.json")))
    sets = {}
    for r in json.load(open(os.path.join(run, "results.json"))):
        bad = [k for k, v in (r.get("bench_checks") or {}).items() if not v[0]]
        sets[r["set"]] = {"verdict": r.get("bench_verdict") or "-", "failed": bad, "wall": r.get("wall_s"),
                          "elapsed": r.get("elapsed_s"), "mr_decision": r.get("mr_decision"),
                          "steps": r.get("steps") or {}, "cmd": r.get("cmd"), **from_log(log_dir(run, r))}
    b, o, hw = man.get("binary", {}), man.get("options", {}), man.get("hardware") or {}
    first = sets.get("4nwv") or next(iter(sets.values()))
    gpus = ", ".join(f"{g.get('name')} (memory {rec(g.get('memory'))})" for g in hw.get("gpus", [])) or None
    repro = {
        "pack": (f"{os.path.basename(p['archive'])}, {p['member']}" if "archive" in p else "run directory on the machine itself"),
        "run directory": p.get("run_name") or os.path.basename(os.path.normpath(run)),
        "date/time": f"{man.get('start')} to {man.get('end')} (the machine's local time)",
        "host": man.get("host"),
        "rugnux (as the binary reports)": first.get("version"),
        "build flags (as the binary reports)": first.get("build_flags"),
        "build (battery manifest)": f"{b.get('flags')}; CMakeCache {b.get('cmake_cache')}; source HEAD {b.get('source_head')}"
                                    f"{' (tree dirty)' if b.get('source_dirty') else ''}; sha256 {b.get('sha256')}",
        "battery runner": f"tools/battery at {man.get('runner_git')}{' (dirty)' if man.get('runner_dirty') else ''}",
        "battery options (manifest)": f"label {man.get('label')}, threads {o.get('threads')} (0 = no -N), gpulock "
                                      f"{o.get('gpulock')}, extra args {o.get('extra_args')}, timeout {o.get('timeout')} s",
        "threads (rugnux reports)": "/".join(sorted({r.get("threads") or "?" for r in sets.values()})),
        "hardware (manifest)": (f"CPU {rec(hw.get('cpu_model'))}, {rec(hw.get('physical_cores'))} cores / {rec(hw.get('threads'))} threads, "
                                f"{hw.get('memory_gib'):.1f} GiB RAM, GPU {gpus}") if hw else None,
        "GPU (rugnux reports)": first.get("gpu"),
    }
    for s in SETS:
        if s in sets:
            repro[f"rugnux command, {s}"] = f"`{sets[s]['cmd']}`" if sets[s].get("cmd") else None
    return sets, repro


def rec(x):
    return NR if x in (None, "unknown", "[N/A]") else x


def stage_table(L, cfg, data, SET, num, title):
        groups = [(g["box"], g["steps"]) for g in cfg["step_groups"]]
        mapped = {k for _, ks in groups for k in ks}
        for _, d, _ in data:   # a step no group names keeps a row of its own, so nothing is dropped
            for k in d.get(SET, {}).get("steps", {}):
                if k not in mapped:
                    groups.append((f"(not in the figure) {k}", [k]))
                    mapped.add(k)
        extra = cfg.get("step_extra_columns", [])   # [label, step]: one step shown again on its own
        L += ["", f"## Table {num}: {title}, time per pipeline stage (s)", "",
              "rugnux's `TIMING step` lines of " + SET + "'s p_log.txt, grouped into the boxes of the workflow figure "
              "(docs/images/rugnux_pipeline.svg); the mapping is " + ("below" if num == 2 else "under Table 2") + ". `-` = none of the stage's steps ran (from "
              "the log); `n/r` = a pack without logs reports none of them; `*` = it reports only some of them, so a sum "
              "marked `*` is partial. Wall - sum shows whether the steps cover rugnux's own "
              "clock (WALL_TIME). Elapsed (the battery's `elapsed_s`, from process launch to exit) minus wall is the time "
              "outside rugnux's clock: process start-up before step 1 and the exit after the last step.", ""]
        logged = {k for p, d, _ in data if "manual" not in p for k in d.get(SET, {}).get("steps", {})}
        hdr = (["Platform", "GPU", "run"] + [g for g, _ in groups] + [lab for lab, _ in extra]
               + ["**sum of steps**", "**wall (WALL_TIME)**", "wall - sum", "elapsed (elapsed_s)", "elapsed - wall"])
        L += ["| " + " | ".join(hdr) + " |", "|:--|:--|:--|" + "--:|" * (len(hdr) - 3)]
        for p, d, _ in data:
            r = d.get(SET, {})
            st = r.get("steps", {})
            s = sum(st.values()) if st else None
            w, e = r.get("wall"), r.get("elapsed")
            cells = [p["platform"], p["gpu"], p.get("role", "")]
            partial = False
            for _, ks in groups:
                have = [k for k in ks if k in st]
                if "manual" in p:   # a typed-in report: a step it does not list (but a log has) is unknown
                    ks = [k for k in ks if k in logged or k in st]
                    cells.append("n/r" if not have else f(sum(st[k] for k in have)) + ("*" if len(have) < len(ks) else ""))
                    partial = partial or len(have) < len(ks)
                else:               # from the log: a step that is not there did not run
                    cells.append(f(sum(st[k] for k in have)) if have else "-")
            cells += [f(st[k]) if k in st else "n/r" for _, k in extra]
            star = "*" if partial else ""
            cells += [f(s) + star, f(w), (f"{w - s:+.1f}" + star) if s is not None and w is not None else "-",
                      f(e) if e else NR, f"{e - w:+.1f}" if e and w else NR]
            L.append("| " + " | ".join(cells) + " |")
        if num != 2:
            L += ["", "Mapping of steps to figure boxes: as under Table 2."]
        else:
          L += ["", "Figure box <- rugnux steps (what each step covers, from rugnux/*.cpp):", ""]
          L += [f"- **{g['box']}** <- {' + '.join(g['steps'])}{': ' + g['why'] if g.get('why') else ''}" for g in cfg["step_groups"]]
          L += [f"- {n}" for n in cfg.get("step_notes", [])]
        notes = [f"- {p['platform']} ({p['gpu']}): {p['table_notes'][SET]}" for p, _, _ in data if p.get("table_notes", {}).get(SET)]
        if notes:
            L += [""] + notes


def f(x):
    return "-" if x is None else f"{x:.1f}"


def main():
    cfg_path = sys.argv[1]
    cfg = json.load(open(cfg_path))
    out_path = sys.argv[sys.argv.index("-o") + 1] if "-o" in sys.argv else None
    tj = json.load(open(os.path.join(os.path.dirname(os.path.abspath(cfg_path)), cfg["timing_json"])))
    benchmark = tj["benchmarks"][cfg.get("benchmark", "inhouse")]
    SETS[:] = [t["set"] for t in benchmark]
    short = {s: s.split("_")[0] for s in SETS}
    data = [(p, *load(p)) for p in cfg["platforms"]]
    L = [f"# {cfg['title']}", "", cfg.get("intro", ""), ""]

    L += ["## Table 1: wall time (s), timing benchmark", "",
          "Wall = rugnux's own WALL_TIME. Checks = the benchmark's six checks on each set (space_group, placed, "
          "mr_group, r_free, confirmed, placement_cc); a failing check would be named.", "",
          "| Platform | GPU | run | " + " | ".join(SETS) + " | checks | MR decision ("
          + " / ".join(short[s] for s in SETS) + ") |",
          "|:--|:--|:--|" + "--:|" * len(SETS) + ":--|:--|"]
    for p, d, _ in data:
        ver = [d[s]["verdict"] + (" (" + ", ".join(d[s]["failed"]) + ")" if d[s].get("failed") else "")
               for s in SETS if s in d]
        checks = "all pass" if len(ver) == len(SETS) and all(v == "pass" for v in ver) else " / ".join(ver)
        L.append(f"| {p['platform']} | {p['gpu']} | {p.get('role', '')} | "
                 + " | ".join(f(d.get(s, {}).get("wall")) for s in SETS)
                 + f" | {checks} | " + " / ".join(d.get(s, {}).get("mr_decision") or NR for s in SETS) + " |")

    titles = {"cytc_x10sa": "cytochrome C (cytc_x10sa)", "myob_x10sa": "myoglobin (myob_x10sa)"}
    # The large set first (Table 2 carries the mapping), then the others from the last to the first.
    for num, s_ in enumerate(["4nwv"] + [s for s in reversed(SETS) if s != "4nwv"], start=2):
        stage_table(L, cfg, data, s_, num, titles.get(s_, s_))
    L += ["", "Where 4nwv's work ran and peak memory:", "",
          "| Platform | GPU | run | scaling | model validation | MR search | peak memory ("
          + " / ".join(short[s] for s in SETS) + ") |",
          "|:--|:--|:--|:--|:--|:--|:--|"]
    for p, d, _ in data:
        r = d.get("4nwv", {})
        dev = [r.get("scaling"), r.get("mv"), r.get("mr_dev")] if r.get("scaling") else p.get("devices", [NR] * 3)
        pm = p.get("peak_mem")
        L.append(f"| {p['platform']} | {p['gpu']} | {p.get('role', '')} | " + " | ".join(x or NR for x in dev)
                 + " | " + (" / ".join(pm.get(s, NR) for s in SETS) if pm else NR) + " |")

    L += ["", "## Data", "",
          f"The {len(SETS)} sets of the `{cfg.get('benchmark', 'inhouse')}` benchmark in tools/battery/timing.json (copy: `{cfg['timing_json']}` beside this report). Each "
          "is started from a model it has to place by molecular replacement:", "",
          "| set | arm | start model (--model) | expected: MR group (sgno), R_free max, placement CC min |",
          "|:--|:--|:--|:--|"]
    for t in benchmark:
        st = t["start"]
        if "alphafold" in st:
            sm = f"`tools/battery/{st['alphafold']}`: {st['why']}"
        else:
            mv = st["moved"]
            sm = (f"deposited {mv['pdb'].upper()} from the site's pdb_cache, strict NCS expanded, turned {mv['deg']} deg "
                  f"about the unit axis {tuple(mv['axis'])} through the centroid c of the expanded asymmetric unit and "
                  f"shifted t = {tuple(mv['shift'])} A (x' = R (x - c) + c + t), written as 4nwv_moved.cif.gz")
        e = t["expect"]
        L.append(f"| {t['set']} | {t['arm']} | {sm} | {e['mr_sgno']}, {e['rfree_max']}, {e['cc_min']} |")
    L += ["", "Each platform's input path for each set is in its rugnux commands below."]

    L += ["", "## Reproducibility, per platform and run", ""]
    order = ["pack", "run directory", "date/time", "host", "OS", "rugnux (as the binary reports)",
             "build flags (as the binary reports)", "build (battery manifest)", "compiler / CUDA / arch",
             "battery command", "battery runner", "battery options (manifest)", "threads (rugnux reports)",
             "hardware (manifest)", "GPU (rugnux reports)", "conditions"] + [f"rugnux command, {s}" for s in SETS]
    for p, d, repro in data:
        for k, field in (("OS", "os"), ("compiler / CUDA / arch", "build"), ("battery command", "battery_cmd"),
                         ("conditions", "conditions")):
            repro.setdefault(k, p.get(field))
        L += [f"### {p['platform']} | {p['gpu']} | {p.get('role', '')}", ""]
        L += [f"- **{k}**: {repro.get(k) or NR}" for k in order if k in repro]
        L += [f"- **{k}**: {v or NR}" for k, v in repro.items() if k not in order]
        L += [f"- {n}" for n in p.get("notes", [])] + [""]

    if cfg.get("stands_out"):
        L += ["## What stands out", ""] + [f"- {x}" for x in cfg["stands_out"]]
    text = "\n".join(L) + "\n"
    if out_path:
        open(out_path, "w").write(text)
    print(text)


if __name__ == "__main__":
    main()