source/study/analyze.py

#!/usr/bin/env python3
"""Post-hoc, task-weighted ProgramBench pilot curves. Never uses best-so-far scores."""

from __future__ import annotations

import argparse
from bisect import bisect_right
from contextlib import closing
import csv
from datetime import datetime, timezone
import json
import math
from pathlib import Path
import sqlite3


SCHEDULE = [0, 30, 60, 120, 300, 600, 1200, 2400, 4800, 7200]
POST_HOC_SCOPE = {
    "kind": "post_hoc_subset",
    "included_tasks": ["jqlang__jq.b33a763", "burntsushi__xsv.f430466",
                       "cmatsuoka__figlet.202a0a8", "sharkdp__hexyl.2e26437"],
    "excluded_tasks": ["burntsushi__ripgrep.3b7fd44"],
    "exclusion_reason": "slow_saved_artifact_grading",
}

CONDITIONS = ["single", "peers", "async"]
AXES = ["wall_seconds", "reference_seconds", "api_prefill_reference_seconds"]
CLOCK_EVENT_TYPES = {"episode_start", "agent_created", "message_sent", "message_delivered",
                     "tool_start", "tool_end", "snapshot", "model_request_end"}
CLOCK_EVENT_FIELDS = {"type", "elapsed_ms", "sequence", "utc", "agent_id", "parent_id",
                      "causal_event_id", "duration_ms", "path", "request_id", "id"}


def analysis_scope(source):
    """Validate the disclosed task-level exclusion; never change the original evidence."""
    if source is None:
        return None
    if not isinstance(source, dict) or set(source) != set(POST_HOC_SCOPE):
        raise ValueError("Unexpected analysis scope")
    for key, expected in POST_HOC_SCOPE.items():
        value = source[key]
        if isinstance(expected, list):
            if not isinstance(value, list) or len(value) != len(expected) or set(value) != set(expected):
                raise ValueError("Unexpected analysis task subset")
        elif value != expected:
            raise ValueError("Unexpected analysis scope category")
    return {key: list(value) if isinstance(value, list) else value for key, value in POST_HOC_SCOPE.items()}


def finite(value):
    return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value)


def read_json(path, default):
    return json.loads(path.read_text()) if path.exists() else default


def read_clock_events(path):
    """Read actual JSONL records, retaining only causal-clock metadata.

    str.splitlines() also splits Unicode separators inside valid JSON strings.
    File iteration respects JSONL newlines and bounds transcript memory to a single
    record instead of retaining every repeated prompt and tool response.
    """
    events, has_origin = [], False
    if not path.exists():
        return events
    with path.open(encoding="utf-8") as stream:
        for line in stream:
            if not line.strip():
                continue
            event = json.loads(line)
            is_origin = bool(event.get("utc")) and not has_origin
            if is_origin or event.get("type") in CLOCK_EVENT_TYPES:
                events.append({key: value for key, value in event.items() if key in CLOCK_EVENT_FIELDS})
            has_origin = has_origin or is_origin
    return events


def epoch(value):
    return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()


def clock_origin(events):
    return next((epoch(event["utc"]) - event.get("elapsed_ms", 0) / 1000
                 for event in events if event.get("utc")), None)


def clock_trace(events, requests, prefill_tps=10_000, decode_tps=265):
    """Two causal clocks from gateway operations and measured tool durations.

    Primary prefill counts estimated new context input; sensitivity prefill counts
    every billed prompt token. Gateway start/end times also capture compaction calls
    that bypass Pi's regular generation hooks. Only completed operations add time.
    """
    timeline, warnings = [], []
    origin = clock_origin(events)
    for index, event in enumerate(events):
        if finite(event.get("elapsed_ms")):
            timeline.append((event["elapsed_ms"] / 1000, index, event))
    ledger_ids = {r.get("client_request_id") for r in requests if r.get("client_request_id")}
    observed_ids = {e.get("request_id") for e in events if e.get("type") == "model_request_end"}
    if observed_ids - ledger_ids:
        warnings.append("model_events_missing_from_gateway_ledger")
    for index, row in enumerate(requests):
        if row.get("status") == "rejected":
            continue
        output, prompt, context = row.get("completion_tokens"), row.get("prompt_tokens"), row.get("context_charged")
        if not all(finite(v) and v >= 0 for v in (output, prompt, context)) or not row.get("ended_at"):
            warnings.append("request_has_unknown_usage_or_end_time")
            continue
        if origin is None:
            warnings.append("missing_episode_clock_origin")
            continue
        start, end = epoch(row["started_at"]) - origin, epoch(row["ended_at"]) - origin
        if end < start:
            warnings.append("invalid_gateway_timestamps")
            continue
        ident = row["id"]
        durations = (max(0, context - output) / prefill_tps + output / decode_tps,
                     prompt / prefill_tps + output / decode_tps)
        timeline.extend([(start, -2, {"type": "clock_model_start", "agent_id": row["agent_id"], "id": ident}),
                         (end, -1, {"type": "clock_model_end", "agent_id": row["agent_id"], "id": ident,
                                    "durations": durations})])
    agents, sent, starts = {}, {}, {}
    trace = [(0.0, 0.0, 0.0)]
    snapshot_clocks = {}
    for wall, _, event in sorted(timeline, key=lambda row: (row[0], row[1])):
        kind, agent = event.get("type"), event.get("agent_id")
        if agent:
            agents.setdefault(agent, (0.0, 0.0))
        if kind == "agent_created":
            agents[agent] = agents.get(event.get("parent_id"), (0.0, 0.0))
        elif kind == "message_sent":
            sent[event.get("sequence")] = agents.get(agent, (0.0, 0.0))
        elif kind == "message_delivered":
            inherited = sent.get(event.get("causal_event_id"))
            if inherited is None:
                warnings.append("missing_message_send_event")
            else:
                agents[agent] = tuple(max(a, b) for a, b in zip(agents[agent], inherited))
        elif kind in ("clock_model_start", "tool_start"):
            starts[(kind, event.get("id", event.get("sequence")))] = agents.get(agent, (0.0, 0.0))
        elif kind in ("clock_model_end", "tool_end"):
            if kind == "clock_model_end":
                key, durations = ("clock_model_start", event["id"]), event["durations"]
            else:
                duration = event.get("duration_ms")
                if not finite(duration) or duration < 0:
                    warnings.append("invalid_tool_duration")
                    continue
                key, durations = ("tool_start", event.get("causal_event_id")), (duration / 1000,) * 2
            anchor = starts.get(key, agents.get(agent, (0.0, 0.0)))
            agents[agent] = tuple(max(current, begin + duration)
                                  for current, begin, duration in zip(agents[agent], anchor, durations))
        current = tuple(max((clock[i] for clock in agents.values()), default=0.0) for i in (0, 1))
        trace.append((max(0, wall), *current))
        if kind == "snapshot" and event.get("path"):
            snapshot_clocks[str(event["path"])] = current
    if not events:
        warnings.append("missing_runner_events")
    return {"trace": trace, "snapshot_clocks": snapshot_clocks,
            "valid": not warnings, "warnings": sorted(set(warnings))}


def clocks_at(clocks, wall, snapshot_path=None):
    if snapshot_path in clocks["snapshot_clocks"]:
        return clocks["snapshot_clocks"][snapshot_path]
    rows = clocks["trace"]
    index = bisect_right([row[0] for row in rows], wall) - 1
    return rows[max(0, index)][1:]


def score_at(points, time_value, axis):
    eligible = [point for point in points if point[axis] <= time_value]
    if not eligible:
        return 0.0, False
    # Ordering includes original wall time for tied normalized clocks. Regressions stay.
    latest = max(eligible, key=lambda point: (point[axis], point["wall_seconds"], point["order"]))
    return latest["score"], True


def attempt_data(directory, run, requests, prefill_tps, decode_tps):
    """Score one saved attempt using only requests belonging to its solver clock."""
    result = read_json(directory / "runner" / "result.json", {})
    controller_status = read_json(directory / "status.json", {})
    event_path = directory / "runner" / "events.jsonl"
    events = read_clock_events(event_path)
    elapsed = result.get("elapsed_seconds")
    origin = clock_origin(events)
    attempt_requests = requests
    if origin is not None:
        end = origin + elapsed if finite(elapsed) else math.inf
        # Event UTC timestamps have millisecond precision. Include that rounding
        # tolerance while excluding old attempts even if worker request IDs repeat.
        attempt_requests = [row for row in requests
                            if origin - 0.001 <= epoch(row["started_at"]) <= end + 0.001]
    clocks = clock_trace(events, attempt_requests, prefill_tps, decode_tps)
    grades = read_json(directory / "grades.json", [])
    if isinstance(grades, dict):
        grades = grades.get("grades", [])
    points, invalid_grades, discarded_grades, latest_grade = [], 0, 0, None
    for order, grade in enumerate(grades):
        score = grade.get("score")
        valid = grade.get("valid", True)
        if score is None and grade.get("sha256"):
            summary = read_json(directory / "evaluations" / grade["sha256"] / run["task"] / "summary.json", {})
            score = summary.get("analysis_score", summary.get("score"))
            valid = valid and summary.get("valid", summary.get("complete", False))
        wall = grade.get("elapsed_seconds")
        valid = bool(valid and finite(score) and 0 <= score <= 1 and finite(wall) and wall >= 0)
        if not valid:
            invalid_grades += 1
        # The original controller also repeated the final archive at scheduled
        # cutoffs beyond early termination. These are not additional solver states.
        if finite(elapsed) and finite(wall) and wall > elapsed + 1e-6:
            discarded_grades += 1
            continue
        if finite(wall) and wall >= 0:
            candidate = (wall, order, valid)
            if latest_grade is None or candidate[:2] >= latest_grade[:2]:
                latest_grade = candidate
        if not valid:
            continue
        primary, sensitivity = clocks_at(clocks, wall, grade.get("snapshot_path"))
        points.append({"order": order, "wall_seconds": wall, "reference_seconds": primary,
                       "api_prefill_reference_seconds": sensitivity, "score": score})
    points.sort(key=lambda point: (point["wall_seconds"], point["order"]))
    last_score = points[-1]["score"] if points else 0.0
    stop_reason = controller_status.get("stop_reason", result.get("stop_reason", "missing_result"))
    status = controller_status.get("status")
    if status not in {"completed", "failed", "interrupted", "budget_exhausted", "running", "pending"}:
        status = ("completed" if stop_reason in {"completed", "wall_time_limit", "wall_time_exhausted"}
                  else "budget_exhausted" if "budget" in stop_reason or stop_reason == "daily_quota_headroom"
                  else "interrupted" if stop_reason == "interrupted" else "failed")
    terminal_valid = (latest_grade is not None and latest_grade[2] and finite(elapsed)
                      and abs(latest_grade[0] - elapsed) <= 1e-6)
    return {"status": status, "stop_reason": stop_reason, "elapsed_seconds": elapsed,
            "final_score": last_score, "has_valid_grade": bool(points), "invalid_grade_count": invalid_grades,
            "has_valid_final_grade": bool(terminal_valid), "discarded_scheduled_grade_count": discarded_grades,
            "clock_valid": clocks["valid"], "clock_warnings": clocks["warnings"], "points": points,
            "clock_total": clocks_at(clocks, elapsed) if finite(elapsed) else clocks["trace"][-1][1:]}


def episode_data(root, run, requests, prefill_tps, decode_tps):
    history = root / "recovery-attempts" / run["run_id"]
    archives = sorted((path for path in history.glob("attempt-*") if path.is_dir()),
                      key=lambda path: int(path.name.removeprefix("attempt-")))
    directories = archives + [root / "episodes" / run["run_id"]]
    points, attempts, wall_offset, reference_offset, sensitivity_offset = [], [], 0.0, 0.0, 0.0
    warnings = set()
    for index, directory in enumerate(directories):
        attempt = attempt_data(directory, run, requests, prefill_tps, decode_tps)
        if index < len(archives) and not finite(attempt["elapsed_seconds"]):
            raise ValueError(f"Archived attempt has no solver duration: {run['run_id']} attempt {index + 1}")
        for point in attempt["points"]:
            points.append({**point, "order": len(points),
                           "wall_seconds": wall_offset + point["wall_seconds"],
                           "reference_seconds": reference_offset + point["reference_seconds"],
                           "api_prefill_reference_seconds": sensitivity_offset + point["api_prefill_reference_seconds"]})
        wall_offset += attempt["elapsed_seconds"] if finite(attempt["elapsed_seconds"]) else 0
        primary, sensitivity = attempt["clock_total"]
        reference_offset += primary
        sensitivity_offset += sensitivity
        warnings.update(attempt["clock_warnings"])
        attempts.append({key: value for key, value in attempt.items() if key not in {"points", "clock_total"}})
        attempts[-1]["attempt"] = index + 1
    active = attempts[-1]
    return {"run_id": run["run_id"], "task": run["task"], "condition": run["condition"],
            "status": active["status"], "stop_reason": active["stop_reason"],
            "elapsed_seconds": wall_offset if finite(active["elapsed_seconds"]) else None,
            "final_score": points[-1]["score"] if points else 0.0, "has_valid_grade": bool(points),
            "has_valid_final_grade": active["has_valid_final_grade"],
            "invalid_grade_count": sum(attempt["invalid_grade_count"] for attempt in attempts),
            "current_attempt_invalid_grade_count": active["invalid_grade_count"],
            "attempt_count": len(attempts), "attempts": attempts,
            "clock_valid": all(attempt["clock_valid"] for attempt in attempts),
            "clock_warnings": sorted(warnings), "points": points,
            "api_tokens_charged": sum(row.get("charged", 0) for row in requests)}


def aggregate_at(episodes, tasks, condition, time_value, axis):
    selected = [row for row in episodes if row["condition"] == condition]
    clock_valid = axis == "wall_seconds" or all(row["clock_valid"] for row in selected)
    task_scores, observed = [], 0
    for task in tasks:
        repetitions = [row for row in selected if row["task"] == task]
        values = [score_at(row["points"], time_value, axis) for row in repetitions]
        task_scores.append(sum(value for value, _ in values) / len(values) if values else 0.0)
        observed += int(any(valid for _, valid in values))
    return {"axis": axis, "time_seconds": time_value, "condition": condition,
            "mean_hidden_test_fraction": sum(task_scores) / len(tasks) if clock_valid else None,
            "tasks_total": len(tasks), "tasks_with_valid_grade": observed,
            "episodes_configured": len(selected), "clock_valid": clock_valid}


def analyze(experiment_dir, config, prefill_tps=10_000, decode_tps=265):
    if prefill_tps <= 0 or decode_tps <= 0:
        raise ValueError("Reference rates must be positive")
    root = Path(experiment_dir)
    tasks = [task["instance_id"] for task in config["tasks"]]
    if not tasks or len(tasks) != len(set(tasks)):
        raise ValueError("Selected task IDs must be nonempty and unique")
    runs = [row for row in config["runs"] if row.get("category", "benchmark") == "benchmark"]
    if len({run["run_id"] for run in runs}) != len(runs):
        raise ValueError("Duplicate run ID")
    if any(run["task"] not in tasks or run["condition"] not in CONDITIONS for run in runs):
        raise ValueError("Unexpected task or condition")
    requests = []
    ledger = root / "budget.sqlite"
    if ledger.exists():
        with closing(sqlite3.connect(f"file:{ledger.resolve()}?mode=ro", uri=True)) as db:
            db.row_factory = sqlite3.Row
            requests = [dict(row) for row in db.execute("SELECT * FROM requests")]
    episodes = [episode_data(root, run, [row for row in requests if row["run_id"] == run["run_id"]],
                             prefill_tps, decode_tps) for run in runs]
    curves = []
    for axis in AXES:
        grid = set(SCHEDULE if axis == "wall_seconds" else [0.0])
        grid.update(point[axis] for row in episodes for point in row["points"])
        for instant in sorted(grid):
            curves.extend(aggregate_at(episodes, tasks, condition, instant, axis) for condition in CONDITIONS)
    final = {}
    for condition in CONDITIONS:
        scores = []
        selected = [row for row in episodes if row["condition"] == condition]
        for task in tasks:
            values = [row["final_score"] for row in selected if row["task"] == task]
            scores.append(sum(values) / len(values) if values else 0.0)
        final[condition] = {"mean_hidden_test_fraction": sum(scores) / len(tasks),
                            "tasks_total": len(tasks),
                            "tasks_with_valid_grade": len({row["task"] for row in selected if row["has_valid_grade"]}),
                            "episodes_configured": len(selected),
                            "episodes_completed": sum(row["status"] == "completed" for row in selected)}
    summary = {"generated_at_utc": datetime.now(timezone.utc).isoformat(), "tasks": tasks,
               "aggregation": "mean hidden-test fraction within each task, then equal-weight mean over every selected task",
               "failure_policy": "latest valid score carried forward; zero before any valid grade; no best-so-far filtering; retries accumulate prior solver and normalized time; scheduled grades after attempt end excluded",
               "reference_clock": {"decode_tokens_per_second": decode_tps, "prefill_tokens_per_second": prefill_tps,
                                   "primary_prefill": "estimated newly entering context tokens",
                                   "sensitivity_prefill": "all API prompt tokens",
                                   "label": "operational normalization; assumed reference rates; not Anthropic's unknown constants",
                                   "includes": "completed model operations, actual tool durations, parent/message causal handoffs",
                                   "limitation": "no partial in-flight decode/tool progress at snapshots"},
               "uncertainty": "exploratory five-task convenience sample; one repetition; no confidence intervals",
               "final": final, "episodes": episodes}
    root.mkdir(parents=True, exist_ok=True)
    (root / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    with (root / "curves.csv").open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(curves[0]))
        writer.writeheader()
        writer.writerows(curves)
    return summary, curves


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--experiment-dir", type=Path, required=True)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--reference-prefill-tps", type=float, default=10_000)
    parser.add_argument("--reference-decode-tps", type=float, default=265)
    args = parser.parse_args()
    summary, curves = analyze(args.experiment_dir, json.loads(args.config.read_text()),
                              args.reference_prefill_tps, args.reference_decode_tps)
    print(json.dumps({"summary": str(args.experiment_dir / "summary.json"),
                      "curves": str(args.experiment_dir / "curves.csv"), "final": summary["final"]}))


if __name__ == "__main__":
    main()