Blame for code/scripts/merge_eval_shards.py

#!/usr/bin/env python3
"""Merge disjoint agent-evaluation shards into canonical official UID order."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from collections import defaultdict
from pathlib import Path
from statistics import fmean
from typing import Any


IDENTITY_FIELDS = (
    "format",
    "classification",
    "benchmark_selected",
    "source_policy_gate_admitted",
    "schema_version",
    "benchmark",
    "model",
    "model_config_sha256",
    "model_manifest_sha256",
    "processor",
    "processor_sha256",
    "processor_contract",
    "controller_processor",
    "controller_processor_sha256",
    "answer_processor",
    "answer_processor_contract",
    "processor_geometry",
    "adapter",
    "evaluator",
    "evaluation_protocol",
    "evaluation_protocol_sha256",
    "protocol_validation",
    "policy_hash",
    "policy_head",
    "policy_head_sha256",
    "policy_report_sha256",
    "policy_calibration_sha256",
    "global_action_threshold",
    "policy_inputs",
    "budget",
    "policy_training_report",
    "one_shared_commit_policy",
    "benchmark_task_or_group_router",
    "question_keyword_router",
    "answer_space_router",
    "gold_answer_available_to_commit_policy",
    "commit_policy",
    "unparseable_or_error_policy",
    "maximum_unique_frames",
    "direct_prediction_cache",
    "direct_prediction_cache_sha256",
    "source_gate",
    "source_gate_sha256",
    "global_action_thresholds",
    "acquisition_identity",
    "action_abi_schema_sha256",
    "inference_policy_hash",
    "inference_policy",
    "policy_config",
    "comparable_under_locked_recent4_protocol",
    "generation_config",
    "execution_optimization",
    "system_prompt_sha256",
    "prompt_contract_sha256",
    "launch_marker",
    "launch_marker_sha256",
    "annotation",
    "annotation_sha256",
    "samples",
    "samples_sha256",
    "runner_sha256",
    "software",
    "policy",
    "identity_hash",
    "protocol_sha256",
    "evaluator_sha256",
    "generation",
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def load_expected_uids(path: Path) -> list[str]:
    uids: list[str] = []
    for line_number, line in enumerate(
        path.read_text(encoding="utf-8").splitlines(), start=1
    ):
        if not line.strip():
            continue
        uid = str(json.loads(line).get("uid") or "")
        if not uid:
            raise ValueError(f"{path}:{line_number} has no uid")
        uids.append(uid)
    if len(uids) != len(set(uids)):
        raise ValueError(f"{path} contains duplicate uids")
    return uids


def grouped_accuracy(rows: list[dict[str, Any]], field: str) -> dict[str, Any]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[str(row.get(field) or "unknown")].append(row)
    result: dict[str, Any] = {}
    for name in sorted(grouped):
        group_rows = grouped[name]
        correct = sum(int(row.get("correct") or 0) for row in group_rows)
        total = len(group_rows)
        result[name] = {
            "accuracy": correct / total if total else 0.0,
            "correct": correct,
            "total": total,
        }
    return result


def aggregate_score(
    score_rows: list[dict[str, Any]], benchmark: str
) -> dict[str, Any]:
    correct = sum(int(row.get("correct") or 0) for row in score_rows)
    total = len(score_rows)
    by_task = grouped_accuracy(score_rows, "task")
    by_group = grouped_accuracy(score_rows, "group")
    score: dict[str, Any] = {
        "accuracy": correct / total if total else 0.0,
        "by_group": by_group,
        "by_task": by_task,
        "num_queries": total,
        "rows": score_rows,
    }
    if benchmark == "ovo":
        tasks_by_group: dict[str, set[str]] = defaultdict(set)
        for row in score_rows:
            tasks_by_group[str(row.get("group") or "unknown")].add(
                str(row.get("task") or "unknown")
            )
        official_groups: dict[str, Any] = {}
        for group in sorted(tasks_by_group):
            task_names = sorted(tasks_by_group[group])
            official_groups[group] = {
                "accuracy": fmean(by_task[name]["accuracy"] for name in task_names),
                "num_tasks": len(task_names),
            }
        score["ovo_official"] = {
            "accuracy": fmean(
                value["accuracy"] for value in official_groups.values()
            ),
            "by_group": official_groups,
        }
    return score


def aggregate_efficiency(rows: list[dict[str, Any]]) -> dict[str, float]:
    def mean(field: str, *fallbacks: str) -> float:
        def value(row: dict[str, Any]) -> float:
            for name in (field, *fallbacks):
                if row.get(name) is not None:
                    return float(row[name])
            return 0.0

        return fmean(value(row) for row in rows) if rows else 0.0

    result = {
        "avg_frames_observed": mean(
            "frames_observed", "unique_frames_acquired", "unique_frames"
        ),
        "avg_model_processed_frames": mean(
            "model_processed_frames", "processed_frames"
        ),
        "avg_media_calls": mean("media_calls"),
        "avg_total_tokens": (
            mean("total_tokens")
            if any(row.get("total_tokens") is not None for row in rows)
            else mean("input_tokens") + mean("output_tokens")
        ),
        "avg_turns": mean("num_turns"),
    }
    if any(row.get("commit_latency_seconds") is not None for row in rows):
        result["avg_commit_latency_seconds"] = mean("commit_latency_seconds")
    if any(row.get("commit_visual_tokens") is not None for row in rows):
        result["avg_commit_visual_tokens"] = mean("commit_visual_tokens")
    if any(row.get("latency_seconds") is not None for row in rows):
        result["avg_latency_seconds"] = mean("latency_seconds")
    return result


def aggregate_agent_behavior(rows: list[dict[str, Any]]) -> dict[str, Any]:
    terminal_counts: dict[str, int] = defaultdict(int)
    for row in rows:
        terminal_counts[str(row.get("policy_terminal_event") or "unknown")] += 1
    total = len(rows)
    failures = sum(
        terminal_counts[event]
        for event in ("malformed_action", "illegal_observe", "missing_answer")
    )
    return {
        "terminal_counts": dict(sorted(terminal_counts.items())),
        "num_policy_failures": failures,
        "strict_terminal_answer_rate": (
            terminal_counts["answered"] / total if total else 0.0
        ),
        "malformed_action_rate": (
            terminal_counts["malformed_action"] / total if total else 0.0
        ),
        "illegal_observe_rate": (
            terminal_counts["illegal_observe"] / total if total else 0.0
        ),
        "missing_answer_rate": (
            terminal_counts["missing_answer"] / total if total else 0.0
        ),
    }


def aggregate_execution_statistics(
    shard_statistics: list[dict[str, Any]],
) -> dict[str, Any]:
    """Sum disjoint-shard cache counters and per-row stage timings."""

    counter_fields = (
        "decode_window_cache_hits",
        "decode_window_cache_misses",
        "decode_window_singleflight_hits",
        "temporal_budget_cache_hits",
        "temporal_budget_cache_misses",
        "video_reader_cache_hits",
        "video_reader_cache_misses",
        "stage_timing_samples",
    )
    stage_totals: dict[str, float] = defaultdict(float)
    result = {
        field: sum(int(stats.get(field) or 0) for stats in shard_statistics)
        for field in counter_fields
    }
    for stats in shard_statistics:
        for name, seconds in (stats.get("stage_seconds_total") or {}).items():
            stage_totals[str(name)] += float(seconds or 0.0)
    result["stage_seconds_total"] = {
        name: round(seconds, 6) for name, seconds in sorted(stage_totals.items())
    }
    return result


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--inputs", nargs="+", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--expected-count", type=int, default=-1)
    parser.add_argument(
        "--expected-uids",
        type=Path,
        help="Official JSONL export whose UID order becomes the canonical order.",
    )
    args = parser.parse_args()

    rows_by_uid: dict[str, dict[str, Any]] = {}
    score_by_uid: dict[str, dict[str, Any]] = {}
    seen_in: dict[str, str] = {}
    metadata: list[dict[str, Any]] = []
    execution_statistics: list[dict[str, Any]] = []
    row_repairs: list[dict[str, Any]] = []
    identity: dict[str, Any] | None = None
    benchmark: str | None = None

    for path in args.inputs:
        data = json.loads(path.read_text(encoding="utf-8"))
        # A sharded run intentionally records a different ``samples`` JSONL
        # path (and therefore a different samples hash) in every shard.  When
        # a frozen full-export is supplied via --expected-uids, that export is
        # the canonical merged sample manifest; all other identity fields
        # must still match exactly.
        comparable_identity_fields = tuple(
            field
            for field in IDENTITY_FIELDS
            if args.expected_uids is None
            or field not in {"samples", "samples_sha256"}
        )
        current_identity = {
            field: data.get(field) for field in comparable_identity_fields
        }
        if identity is None:
            identity = current_identity
            benchmark = str(data.get("benchmark") or "") or None
        elif current_identity != identity:
            changed = [
                field
                for field in comparable_identity_fields
                if current_identity.get(field) != identity.get(field)
            ]
            # Annotation files are frequently remounted from a volatile
            # location (for example /dev/shm) to persistent storage when an
            # evaluation is resumed.  The absolute path is not part of the
            # benchmark identity when both shards carry the same non-empty
            # content hash.  Keep the hash itself in the frozen identity so
            # different annotation contents remain a hard failure.
            if (
                "annotation" in changed
                and current_identity.get("annotation_sha256")
                and current_identity.get("annotation_sha256")
                == identity.get("annotation_sha256")
            ):
                changed.remove("annotation")
            if not changed:
                current_identity = identity
            else:
                raise ValueError(
                    f"{path} does not match the frozen shard identity: {changed}"
                )
        shard_rows = data.get("rows")
        if not isinstance(shard_rows, list):
            raise ValueError(f"{path} has no rows list")
        for row in shard_rows:
            uid = str(row.get("uid") or "")
            if not uid:
                raise ValueError(f"{path} contains a row without uid")
            if uid in rows_by_uid:
                raise ValueError(f"duplicate uid {uid!r} in {seen_in[uid]} and {path}")
            rows_by_uid[uid] = row
            seen_in[uid] = str(path)
            row_benchmark = str(row.get("benchmark") or "")
            if row_benchmark:
                if benchmark is None:
                    benchmark = row_benchmark
                elif row_benchmark != benchmark:
                    raise ValueError(
                        f"mixed benchmarks in shards: {benchmark!r} and {row_benchmark!r}"
                    )
        shard_score_rows = (data.get("score") or {}).get("rows") or []
        if not isinstance(shard_score_rows, list):
            raise ValueError(f"{path} score.rows is not a list")
        for score_row in shard_score_rows:
            uid = str(score_row.get("query_id") or "")
            if not uid:
                raise ValueError(f"{path} contains a score row without query_id")
            if uid in score_by_uid:
                raise ValueError(f"duplicate score query_id {uid!r}")
            score_by_uid[uid] = score_row
        metadata.append(
            {
                "path": str(path),
                "sha256": sha256(path),
                "num_samples": len(shard_rows),
                "sample_slice": data.get("sample_slice"),
                "samples": data.get("samples"),
                "samples_sha256": data.get("samples_sha256"),
                "annotation": data.get("annotation"),
                "annotation_sha256": data.get("annotation_sha256"),
            }
        )
        if isinstance(data.get("execution_statistics"), dict):
            execution_statistics.append(data["execution_statistics"])
        row_repairs.extend(data.get("row_repairs") or [])

    if args.expected_count >= 0 and len(rows_by_uid) != args.expected_count:
        raise ValueError(
            f"merged {len(rows_by_uid)} rows, expected {args.expected_count}"
        )

    expected_uids = (
        load_expected_uids(args.expected_uids)
        if args.expected_uids is not None
        else list(rows_by_uid)
    )
    expected_set = set(expected_uids)
    actual_set = set(rows_by_uid)
    if actual_set != expected_set:
        missing = [uid for uid in expected_uids if uid not in actual_set]
        extra = sorted(actual_set - expected_set)
        raise ValueError(
            f"official UID mismatch: missing={missing[:20]} extra={extra[:20]}"
        )
    if set(score_by_uid) != expected_set:
        missing = [uid for uid in expected_uids if uid not in score_by_uid]
        extra = sorted(set(score_by_uid) - expected_set)
        raise ValueError(
            f"score UID mismatch: missing={missing[:20]} extra={extra[:20]}"
        )

    rows = [rows_by_uid[uid] for uid in expected_uids]
    score_rows = [score_by_uid[uid] for uid in expected_uids]
    empty_predictions = [
        str(row.get("uid"))
        for row in rows
        if not str(row.get("prediction") or "").strip()
    ]
    errors = [
        row for row in rows if row.get("error") or row.get("commit_error")
    ]
    merged_identity = dict(identity or {})
    if args.expected_uids is not None:
        merged_identity["samples"] = str(args.expected_uids)
        merged_identity["samples_sha256"] = sha256(args.expected_uids)
    result = {
        "benchmark": benchmark,
        **merged_identity,
        "canonical_official_uid_order": args.expected_uids is not None,
        "efficiency": aggregate_efficiency(rows),
        "empty_prediction_uids": empty_predictions,
        "num_empty_predictions": len(empty_predictions),
        "num_errors": len(errors),
        "num_future_frame_violations": sum(
            bool(row.get("future_frame_violation")) for row in rows
        ),
        "num_samples": len(rows),
        "strict_terminal_answer_rate": (
            sum(bool(row.get("strict_terminal_answer")) for row in rows) / len(rows)
            if rows else 0.0
        ),
        "score": aggregate_score(score_rows, benchmark or ""),
        "row_repairs": row_repairs,
        "source_shards": metadata,
        "expected_uid_source": (
            {
                "path": str(args.expected_uids),
                "sha256": sha256(args.expected_uids),
            }
            if args.expected_uids is not None
            else None
        ),
        "rows": rows,
    }
    if execution_statistics:
        result["execution_statistics"] = aggregate_execution_statistics(
            execution_statistics
        )
    if any("policy_terminal_event" in row for row in rows):
        result["agent_behavior"] = aggregate_agent_behavior(rows)
        result["strict_terminal_answer_rate"] = result["agent_behavior"][
            "strict_terminal_answer_rate"
        ]
    if any("commit_action" in row for row in rows):
        result.update(
            {
                "num_commit_errors": sum(
                    bool(row.get("commit_error")) for row in rows
                ),
                "num_unparsed": sum(
                    row.get("commit_parsed") is False for row in rows
                ),
                "decision_counts": {
                    action: sum(row.get("commit_action") == action for row in rows)
                    for action in (
                        "not_applicable",
                        "unchanged",
                        "retain",
                        "commit",
                    )
                },
            }
        )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = args.output.with_suffix(args.output.suffix + f".tmp.{os.getpid()}")
    temporary.write_text(
        json.dumps(result, indent=2, sort_keys=True, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    temporary.replace(args.output)
    print(
        json.dumps(
            {
                "output": str(args.output),
                "num_samples": len(rows),
                "num_errors": len(errors),
                "num_empty_predictions": len(empty_predictions),
                "canonical_official_uid_order": args.expected_uids is not None,
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()