code/scripts/evaluate.py

#!/usr/bin/env python3
"""Pass@1 evaluation: immutable generations, 500 prompts, real 60s run windows."""
from __future__ import annotations

import argparse
from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED
import multiprocessing
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import statistics
import subprocess
import sys
import time

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from vibegame.project import ProjectError, canonical_json, parse_project
from vibegame.validate import sandbox_command, validate_project
from generate import add_backend_args, file_hash, make_generator


def atomic_json(path: Path, value, *, immutable: bool = False):
    if immutable:
        with path.open("x", encoding="utf-8") as handle:
            json.dump(value, handle, indent=2)
            handle.write("\n")
    else:
        temporary = path.with_suffix(path.suffix + ".tmp")
        temporary.write_text(json.dumps(value, indent=2) + "\n")
        temporary.replace(path)


def check_batch_outputs(generations, expected: int) -> None:
    if not isinstance(generations, list) or len(generations) != expected:
        raise RuntimeError("Generation batch result count mismatch")
    if any(not isinstance(row, dict) or not isinstance(row.get("output"), str) for row in generations):
        raise RuntimeError("Generation backend returned a malformed output")


def summarize(output: Path, total: int, run_seconds: float) -> dict:
    results = [json.loads(path.read_text()) for path in sorted((output / "items").glob("*/result.json"))]
    completed = len(results)
    compiled = sum(bool(item.get("compile_success")) for item in results)
    clean = sum(bool(item.get("clean_run_60s")) for item in results)
    infrastructure_failures = sum(bool(item.get("infrastructure_failure")) for item in results)
    generation_backend_failed = (output / "generation_infrastructure_error.json").exists()
    infrastructure_failures += int(generation_backend_failed)
    complete = completed == total and not infrastructure_failures
    latencies = [item["prompt_to_launch_seconds"] for item in results if item.get("compile_success") and item.get("prompt_to_launch_seconds") is not None]
    summary = {"status": "complete" if complete else "incomplete", "expected_prompts": total, "recorded_prompts": completed,
               "pending_prompts": total - completed, "compile_success_count": compiled,
               "generation_backend_failed": generation_backend_failed,
               "clean_run_60s_count": clean, "infrastructure_failures": infrastructure_failures,
               "compile_success_rate": compiled / total if complete and total else None,
               "clean_run_60s_rate": clean / total if complete and total and run_seconds >= 60 else None,
               "run_seconds": run_seconds, "pass_k": 1, "generation_retries": 0, "retries": 0,
               "interrupted_validation_attempts": sum(item.get("interrupted_validation_attempts", 0) for item in results), "sample_indices": [0, 200, 400],
               "prompt_to_launch_median_seconds": statistics.median(latencies) if latencies else None,
               "under_60s_launch_count": sum(value <= 60 for value in latencies),
               "limitations": "No input interaction is tested. A clean headless run does not prove gameplay correctness. Latency includes generation, project writing, import and startup; warm model excludes model loading. Infrastructure failures block final rates."}
    atomic_json(output / "summary.json", summary)
    return summary


def validate_saved_project(files: dict, item: Path, godot: str, seconds: float, queued_at: float | None = None) -> dict:
    """Resume only interrupted validation; never replace a completed failure."""
    attempts = sorted([path for path in item.glob("check*") if path.is_dir()])
    finished = [path for path in attempts if (path / "validation.json").exists()]
    if len(finished) > 1:
        raise RuntimeError("Multiple completed validation attempts found; evidence is ambiguous")
    if finished:
        result = json.loads((finished[0] / "validation.json").read_text())
        result["interrupted_validation_attempts"] = len(attempts) - 1
        result["validation_evidence"] = finished[0].name
        return result
    for path in attempts:
        note = path / "interruption.json"
        if not note.exists():
            atomic_json(note, {"status": "interrupted before final validation result", "resumed_utc": datetime.now(timezone.utc).isoformat(),
                              "policy": "Revalidate the same immutable generation; no resampling or source changes"}, immutable=True)
    destination = item / ("check" if not attempts else f"check-resume-{len(attempts):03d}")
    queue_seconds = max(0, time.monotonic() - queued_at) if queued_at is not None else None
    result = validate_project(files, destination, godot, seconds)
    result["validation_queue_seconds"] = queue_seconds
    result["interrupted_validation_attempts"] = len(attempts)
    result["validation_evidence"] = destination.name
    return result


def finish_result(item: Path, result: dict, check: dict | None, output: Path, total: int, seconds: float):
    if check is not None:
        result.update(check)
        result["infrastructure_failure"] = bool(check.get("error") and not check["error"].startswith("ProjectError:")) or bool(check.get("import", {}).get("sandbox_failed"))
        if check.get("compile_success"):
            startup = check.get("runtime", {}).get("startup_seconds")
            if startup is not None and check.get("validation_queue_seconds") is not None and not check.get("interrupted_validation_attempts"):
                call_seconds = result.get("generation_call_seconds")
                result["generation_latency_source"] = "generation_call_seconds" if call_seconds is not None else "generation_seconds_legacy_fallback"
                generation_elapsed = call_seconds if call_seconds is not None else result["generation_seconds"]
                result["prompt_to_launch_seconds"] = generation_elapsed + check["validation_queue_seconds"] + check.get("preimport_seconds", 0) + check["import"]["wall_seconds"] + startup
    atomic_json(item / "result.json", result, immutable=True)
    summary = summarize(output, total, seconds)
    print(json.dumps({"index": result["index"], "compile_success": result["compile_success"], "clean_run_60s": result["clean_run_60s"], "recorded_prompts": summary["recorded_prompts"]}), flush=True)
    if result.get("infrastructure_failure"):
        raise RuntimeError("Environment validation failed; stopped without scoring unattempted prompts. Inspect retained logs.")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    add_backend_args(parser)
    parser.add_argument("--prompts", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--godot", required=True)
    parser.add_argument("--expected-count", type=int, default=500)
    parser.add_argument("--run-seconds", type=float, default=60)
    parser.add_argument("--workers", type=int, choices=[1, 2, 4], default=4, help="Concurrent CPU validation processes")
    parser.add_argument("--batch-size", type=int, choices=[1, 2, 4, 8, 16, 32], help="GPU generation batch, default 4 for Unsloth or 1 for Ollama")
    parser.add_argument("--smoke-test", action="store_true", help="Allow fewer than 500 prompts/short windows; not publication evaluation")
    args = parser.parse_args()
    args.batch_size = args.batch_size or (4 if args.backend == "unsloth" else 1)
    if args.backend == "ollama" and args.batch_size != 1:
        parser.error("Ollama evaluation uses batch size one")
    if not args.smoke_test and (args.expected_count != 500 or args.run_seconds != 60):
        parser.error("Publication evaluation requires exactly 500 prompts and 60 wall-clock seconds")
    rows = [json.loads(line) for line in args.prompts.read_text().splitlines() if line.strip()]
    if len(rows) != args.expected_count or not rows:
        parser.error(f"Expected {args.expected_count} prompts, got {len(rows)}")
    prompts = [row.get("instruction", row.get("prompt")) for row in rows]
    if any(not isinstance(prompt, str) or not prompt.strip() for prompt in prompts):
        parser.error("Each row requires a nonempty instruction string")
    if len(set(prompts)) != len(prompts):
        parser.error("Evaluation prompts must be unique")
    ids = [str(row.get("id", index)) for index, row in enumerate(rows)]
    if len(set(ids)) != len(ids):
        parser.error("Evaluation prompt ids must be unique")
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "items").mkdir(exist_ok=True)
    preflight = args.output / ".sandbox-preflight"
    preflight.mkdir(exist_ok=True)
    # Fail before expensive generation if the OS cannot isolate generated code.
    wrapped, _ = sandbox_command([str(Path(args.godot).resolve())], preflight)
    wrapped[-1] = "/usr/bin/true"
    probe = subprocess.run(wrapped, capture_output=True, timeout=15)
    if probe.returncode:
        raise RuntimeError("Sandbox preflight failed: " + probe.stderr.decode(errors="replace"))
    generator = make_generator(args)
    # Exercise the inference path before touching any blind prompt. This catches
    # backend errors and warms kernels without producing a scored generation.
    if args.backend == "unsloth":
        generator.generate_many(["Return a tiny Godot game."] * args.batch_size, 4)
    manifest = {"prompts_sha256": file_hash(args.prompts), "runtime_sha256": file_hash(Path(args.godot)),
                "expected_count": args.expected_count, "run_seconds": args.run_seconds, "workers": args.workers,
                "generation": generator.metadata, "max_new_tokens": args.max_new_tokens, "batch_size": args.batch_size,
                "latency_policy": "Whole generation call including chat formatting, tokenization, device transfers and decoding, plus validation queue, writing, import and startup. Interrupted/resumed validations excluded from latency median.",
                "latency_version": 2,
                "decoding": {"temperature": 0, "seed": 42}, "sample_indices": [0, 200, 400],
                "system": "no retries, no repairs, original order, warm model timing", "smoke_test": args.smoke_test}
    manifest_path = args.output / "manifest.json"
    if manifest_path.exists():
        if json.loads(manifest_path.read_text()) != manifest:
            raise ValueError("Resume manifest mismatch: model/runtime/data/settings changed; use a new evaluation directory")
    else:
        atomic_json(manifest_path, manifest, immutable=True)
    # spawn avoids inheriting a live CUDA context into CPU validation workers.
    with ProcessPoolExecutor(max_workers=args.workers, mp_context=multiprocessing.get_context("spawn")) as pool:
        pending = {}
        readiness_written = False

        def collect(block=False):
            if not pending:
                return
            ready, _ = wait(pending, timeout=None if block else 0, return_when=FIRST_COMPLETED)
            for future in ready:
                item, result = pending.pop(future)
                check = future.result()  # worker failure leaves raw/evidence untouched for resume
                finish_result(item, result, check, args.output, len(rows), args.run_seconds)

        for batch_start in range(0, len(prompts), args.batch_size):
            collect()
            candidates = []
            fresh = []
            for index in range(batch_start, min(batch_start + args.batch_size, len(prompts))):
                item = args.output / "items" / f"{index:04d}"
                item.mkdir(exist_ok=True)
                if (item / "result.json").exists():
                    continue
                generation_path = item / "generation.json"
                started = item / "attempt.json"
                if started.exists() and not generation_path.exists():
                    result = {"index": index, "prompt_id": ids[index], "compile_success": False, "clean_run_60s": False,
                              "error": "Generation interrupted before output was persisted; no retry under pass@1"}
                    finish_result(item, result, None, args.output, len(rows), args.run_seconds)
                    continue
                candidates.append((index, item, generation_path))
                if not generation_path.exists():
                    atomic_json(started, {"index": index, "prompt_id": ids[index], "instruction": prompts[index],
                                        "started_utc": datetime.now(timezone.utc).isoformat()}, immutable=True)
                    fresh.append((index, item, generation_path))
            ready_at = None
            if fresh:
                print(json.dumps({"event": "generation_batch_started", "indices": [index for index, _, _ in fresh],
                                  "batch_size": len(fresh), "max_new_tokens": args.max_new_tokens,
                                  "pending_validations": len(pending)}), flush=True)
                try:
                    call_started = time.monotonic()
                    if args.backend == "unsloth":
                        generations = generator.generate_many([prompts[index] for index, _, _ in fresh], args.max_new_tokens)
                    else:
                        generations = [generator.generate(prompts[fresh[0][0]], args.max_new_tokens)]
                    ready_at = time.monotonic()
                    call_elapsed = ready_at - call_started
                    check_batch_outputs(generations, len(fresh))
                except Exception as exc:
                    # Stop a broken backend instead of manufacturing 500 failures.
                    atomic_json(args.output / "generation_infrastructure_error.json", {"error": type(exc).__name__ + ": " + str(exc),
                                "indices": [index for index, _, _ in fresh]})
                    summarize(args.output, len(rows), args.run_seconds)
                    raise RuntimeError("Generation backend failed; retained attempt records and stopped before further prompts") from exc
                for (_, _, generation_path), generation in zip(fresh, generations):
                    # The outer call clock includes return/dispatch overhead;
                    # queue timing begins at exactly this call's end.
                    generation["generation_call_seconds"] = call_elapsed
                    generation["output_sha256"] = hashlib.sha256(generation["output"].encode()).hexdigest()
                    atomic_json(generation_path, generation, immutable=True)
                print(json.dumps({"event": "generation_batch_persisted", "indices": [index for index, _, _ in fresh],
                                  "batch_size": len(fresh), "whole_batch_seconds": max(row["generation_seconds"] for row in generations),
                                  "generation_call_seconds": call_elapsed,
                                  "per_prompt_tokens": [row.get("tokens") for row in generations],
                                  "per_prompt_done_reason": [row.get("done_reason") for row in generations]}), flush=True)
            fresh_indices = {index for index, _, _ in fresh}
            for index, item, generation_path in candidates:
                try:
                    generation = json.loads(generation_path.read_text())
                    if generation.get("output") is not None and generation.get("output_sha256") != hashlib.sha256(generation["output"].encode()).hexdigest():
                        raise ValueError("Persisted generation hash mismatch; refusing to validate modified output")
                except (json.JSONDecodeError, UnicodeError):
                    generation = {"output": None, "error": "Generation persistence interrupted; retained partial file and did not resample"}
                result = {"index": index, "prompt_id": ids[index], "compile_success": False, "clean_run_60s": False,
                          "generation_seconds": generation.get("generation_seconds"), "generation_batch_size": generation.get("generation_batch_size", 1),
                          "generation_call_seconds": generation.get("generation_call_seconds"),
                          "prompt_to_launch_seconds": None}
                if generation.get("error"):
                    result["error"] = generation["error"]
                    finish_result(item, result, None, args.output, len(rows), args.run_seconds)
                    continue
                try:
                    files = parse_project(generation["output"])
                except (ProjectError, ValueError) as exc:
                    result["error"] = f"{type(exc).__name__}: {exc}"
                    finish_result(item, result, None, args.output, len(rows), args.run_seconds)
                    continue
                future = pool.submit(validate_saved_project, files, item, args.godot, args.run_seconds,
                                     ready_at if index in fresh_indices else None)
                pending[future] = (item, result)
            if not readiness_written:
                # A fresh per-process marker lets the continuation runner wait
                # for actual model generation and worker submission on resume.
                # This runtime receipt is outside the frozen scoring manifest.
                atomic_json(args.output / f".runtime-ready-{os.getpid()}.json", {
                    "pid": os.getpid(), "ready_unix_seconds": time.time(),
                    "model_warmup_complete": True, "first_batch_processed": True,
                    "workers": args.workers, "pending_validations": len(pending)})
                readiness_written = True
            # Permit one batch of queued validators so CPU checks overlap the
            # next GPU batch; bound both memory and the latency queue explicitly.
            while len(pending) >= args.workers + args.batch_size:
                collect(block=True)
        while pending:
            collect(block=True)
    print(json.dumps(summarize(args.output, len(rows), args.run_seconds), indent=2))


if __name__ == "__main__":
    main()