scripts/benchmark_workloads.py

"""Matched, repeated workload comparisons; preserve errors and raw predictions."""

import argparse
import hashlib
import json
import platform
import statistics
import subprocess
import time
from pathlib import Path

from gemma_rlcd.comparison import discrete_answers, generate_answers
from gemma_rlcd.core import DecisionEngine, Independent, State, parse_question
from gemma_rlcd.json_backend import JSONMLXBackend


def flatten(values, prefix=""):
    result = {}
    for name, value in values.items():
        key = f"{prefix}.{name}" if prefix else name
        if isinstance(value, dict):
            result.update(flatten(value, key))
        else:
            result[key] = value
    return result


def quality(actual, expected):
    actual = flatten(actual or {})
    expected = flatten(expected)
    checks = {
        key: key in actual
        and (actual[key] in value if isinstance(value, list) else actual[key] == value)
        for key, value in expected.items()
    }
    return {"correct": sum(checks.values()), "scored": len(checks), "checks": checks}


def agreement(left, right):
    if left is None or right is None:
        return None
    left, right = flatten(left), flatten(right)
    keys = left.keys() | right.keys()
    return {key: key in left and key in right and left[key] == right[key] for key in sorted(keys)}


def summarize(samples, methods):
    result = {}
    for method in methods:
        runs = [sample for sample in samples if sample["method"] == method]
        times = [sample["seconds"] for sample in runs]
        result[method] = {
            "median_seconds": statistics.median(times),
            "min_seconds": min(times),
            "max_seconds": max(times),
            "valid_runs": sum(sample["valid"] for sample in runs),
            "attempted_runs": len(runs),
            "correct": sum(sample["quality"]["correct"] for sample in runs),
            "scored": sum(sample["quality"]["scored"] for sample in runs),
            "stable_answers": all(sample["values"] == runs[0]["values"] for sample in runs),
        }
    valid = all(
        result[method]["valid_runs"] == result[method]["attempted_runs"] for method in methods
    )
    result["normal_over_batched"] = (
        result["normal"]["median_seconds"] / result["batched"]["median_seconds"] if valid else None
    )
    if "serial" in result:
        result["serial_over_batched"] = (
            result["serial"]["median_seconds"] / result["batched"]["median_seconds"]
            if valid
            else None
        )
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--cases", type=Path, default=Path("examples/demo-workloads.json"))
    parser.add_argument("--report", type=Path, required=True)
    parser.add_argument("--only", nargs="+")
    parser.add_argument("--repeats", type=int, default=4)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--max-input-tokens", type=int, default=16384)
    args = parser.parse_args()
    if args.repeats < 2:
        parser.error("Use at least two repetitions to vary execution order")
    cases = json.loads(args.cases.read_text())["cases"]
    if args.only:
        unknown = set(args.only) - {case["name"] for case in cases}
        if unknown:
            parser.error(f"Unknown cases: {sorted(unknown)}")
        cases = [case for case in cases if case["name"] in args.only]
    backend = JSONMLXBackend(
        args.model, branch_batch_size=args.batch_size, max_input_tokens=args.max_input_tokens
    )
    report = {
        "status": "development_workloads_not_a_held_out_benchmark",
        "platform": platform.platform(),
        "hardware": subprocess.check_output(
            ["sysctl", "-n", "machdep.cpu.brand_string"], text=True
        ).strip(),
        "host_swap_at_start": subprocess.check_output(
            ["sysctl", "vm.swapusage"], text=True
        ).strip(),
        "model_source": json.loads(Path("model-source.json").read_text()),
        "cases_sha256": hashlib.sha256(args.cases.read_bytes()).hexdigest(),
        "methodology": {
            "compute": "Same resident frozen 4-bit Gemma weights, float32 compute, all 35 layers.",
            "timing": "One excluded warmup per method per case, then repeated runs in rotating/reversed order. GPU synchronized at boundaries. Medians include prompt preparation and inference; exclude model loading, upload, and allocator reset.",
            "cache": "Fresh input/media KV per run. Every scorer run prefills all state and question definitions once. Serial control changes only branch batch size to one, retaining the same shared prefix.",
            "normal": "Greedy compact JSON with decisions only, no reasoning or requested probabilities; same complete input and schema as the scorer. No padded output requirement.",
            "quality": "Predeclared development expectations are scored where available. Invalid structured responses count as failed decisions. Unannotated policy questions are compared only for agreement. Agreement does not establish accuracy.",
            "input_limit": args.max_input_tokens,
            "input_limit_note": "The 255-choice example exceeds the default 8192-token UI limit. This standalone benchmark raises the limit without truncating the schema for either method.",
            "repeats": args.repeats,
            "batch_size": args.batch_size,
        },
        "cases": [],
    }

    def save():
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")

    for case in cases:
        state = State(text=case["text"])
        questions = {name: parse_question(value) for name, value in case["questions"].items()}
        methods = ["batched", "normal"] + (["serial"] if case.get("serial_control") else [])
        row = {
            "name": case["name"],
            "title": case["title"],
            "primitive_fields": sum(
                len(question.criteria) if isinstance(question, Independent) else 1
                for question in questions.values()
            ),
            "max_choices": max(
                2 if value["type"] in {"independent", "noul"} else len(value["criteria"])
                for value in case["questions"].values()
            ),
            "fixture": case,
            "warmups": [],
            "samples": [],
        }
        report["cases"].append(row)
        print(json.dumps({"start": case["name"], "methods": methods}), flush=True)

        def run(method):
            backend.branch_batch_size = 1 if method == "serial" else args.batch_size
            backend.mx.synchronize()
            backend.mx.clear_cache()
            backend.mx.reset_peak_memory()
            started = time.perf_counter()
            record = {"method": method, "valid": False, "values": None}
            try:
                if method == "normal":
                    output = generate_answers(backend, state, questions)
                    record.update(
                        valid=output["valid"], values=output["answers"] if output["valid"] else None
                    )
                else:
                    output = DecisionEngine(backend).system_one(state, questions)
                    output["execution"] = dict(backend.last_stats)
                    record.update(valid=True, values=discrete_answers(output["answers"]))
                record["output"] = output
            except Exception as exc:
                record["error"] = f"{type(exc).__name__}: {exc}"
            backend.mx.synchronize()
            record["seconds"] = time.perf_counter() - started
            record["peak_mlx_bytes"] = backend.mx.get_peak_memory()
            record["quality"] = quality(record["values"], case["expected"])
            return record

        for method in methods:
            row["warmups"].append(run(method))
            save()
        for repeat in range(args.repeats):
            order = methods[repeat % len(methods) :] + methods[: repeat % len(methods)]
            if len(methods) > 2 and repeat % 2:
                order.reverse()
            current = {}
            for method in order:
                sample = run(method)
                sample["repeat"] = repeat
                current[method] = sample["values"]
                row["samples"].append(sample)
                save()
            row.setdefault("agreements", []).append(
                {
                    "repeat": repeat,
                    "normal": agreement(current["batched"], current["normal"]),
                    "serial": agreement(current["batched"], current["serial"])
                    if "serial" in current
                    else None,
                }
            )
            print(
                json.dumps(
                    {
                        "case": case["name"],
                        "repeat": repeat,
                        "seconds": {
                            sample["method"]: round(sample["seconds"], 3)
                            for sample in row["samples"]
                            if sample["repeat"] == repeat
                        },
                    }
                ),
                flush=True,
            )
        row["summary"] = summarize(row["samples"], methods)
        save()
        print(json.dumps({"complete": case["name"], "summary": row["summary"]}), flush=True)


if __name__ == "__main__":
    main()