Blame for cref/source.py

from __future__ import annotations

import functools
import json
import logging
import math
import os
import sys
import threading
import time
import urllib.request
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import UTC, date, datetime, timedelta
from pathlib import Path
from typing import Any

from jsonschema import Draft202012Validator, FormatChecker

ROOT = Path(__file__).parent
SCHEMAS = ROOT / "schemas" / "v1"
LEADERBOARDS = {"release": "v1/leaderboard", "all": "v1/leaderboard/all_samples"}
TTL_SECONDS = 300
HISTORY_TTL_SECONDS = 3600
SOCKET_TIMEOUT_SECONDS = 10
READ_DEADLINE_SECONDS = 30
MAX_BYTES = 25_000_000
HISTORY_START = date(2026, 7, 1)

log = logging.getLogger(__name__)


class InvalidData(ValueError):
    pass


_FORMATS = FormatChecker(formats=())


@_FORMATS.checks("date-time", raises=ValueError)
def _instant(value: object) -> bool:
    return not isinstance(value, str) or 2000 <= datetime.fromisoformat(value).year <= 2200


@functools.cache
def _validator() -> Draft202012Validator:
    doc = json.loads((SCHEMAS / "leaderboard.schema.json").read_text(encoding="utf-8"))
    Draft202012Validator.check_schema(doc)
    return Draft202012Validator(doc, format_checker=_FORMATS)


def _finite(text: str) -> str:
    if not math.isfinite(float(text)):
        raise InvalidData(f"number out of range: {text[:40]}")
    return text


def overall_board(doc: Any) -> Any:
    if isinstance(doc, dict) and isinstance(doc.get("boards"), dict) and "overall" in doc["boards"]:
        return {**doc, "boards": {"overall": doc["boards"]["overall"]}}
    return doc


def parse(data: bytes, select: Callable[[Any], Any] | None = None) -> dict[str, Any]:
    if len(data) > MAX_BYTES:
        raise InvalidData(f"larger than {MAX_BYTES} bytes")
    try:
        doc = json.loads(
            data.decode("utf-8"),
            parse_constant=_finite,
            parse_float=lambda text: float(_finite(text)),
            parse_int=lambda text: int(_finite(text)),
        )
    except InvalidData:
        raise
    except (ValueError, RecursionError) as exc:
        raise InvalidData(str(exc)[:200]) from exc
    if select:
        doc = select(doc)
    error = next(_validator().iter_errors(doc), None)
    if error is not None:
        raise InvalidData(f"{error.json_path}: {error.message[:200]}")
    return doc


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args: Any, **kwargs: Any) -> None:
        return None


_OPENER = urllib.request.build_opener(_NoRedirect)


def read_url(url: str, max_bytes: int, deadline: float = READ_DEADLINE_SECONDS) -> bytes:
    request = urllib.request.Request(
        url, headers={"Accept": "application/json", "User-Agent": "cref-leaderboard"}
    )
    end = time.monotonic() + deadline
    chunks, size = [], 0
    with _OPENER.open(request, timeout=SOCKET_TIMEOUT_SECONDS) as response:
        while size <= max_bytes and (chunk := response.read1(65_536)):
            if time.monotonic() > end:
                raise TimeoutError(f"{url}: not read within {deadline} s")
            chunks.append(chunk)
            size += len(chunk)
    return b"".join(chunks)


def read(base: str, path: str, max_bytes: int = MAX_BYTES) -> bytes:
    if base.startswith("https://"):
        return read_url(f"{base.rstrip('/')}/{path}", max_bytes)
    with open(Path(base) / path, "rb") as file:
        return file.read(max_bytes + 1)


def check_base(base: str) -> str:
    if "://" in base and not base.startswith("https://"):
        raise ValueError(f"CREF_DATA_URL must start with https://, got {base[:80]!r}")
    return base


def base_from_env() -> str:
    return check_base(os.environ.get("CREF_DATA_URL", "").strip() or str(ROOT / "data"))


class Schedule:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._wake = threading.Event()
        self._jobs: dict[object, list[Any]] = {}
        self._thread: threading.Thread | None = None

    def add(self, seconds: float, update: Callable[[], None]) -> Callable[[], None]:
        key = object()
        with self._lock:
            self._jobs[key] = [time.monotonic() + seconds, seconds, update]
        self._wake.set()
        self.ensure_running()

        def remove() -> None:
            with self._lock:
                self._jobs.pop(key, None)

        return remove

    # Called on every page view: cheap while the thread runs, and restarts it if it has died.
    def ensure_running(self) -> None:
        if self._thread is not None and self._thread.is_alive():
            return
        with self._lock:
            if self._thread is None or not self._thread.is_alive():
                self._thread = threading.Thread(target=self._run, name="cref-reads", daemon=True)
                self._thread.start()

    def _run(self) -> None:
        while True:
            with self._lock:
                jobs = list(self._jobs.values())
            for job in jobs:
                if job[0] <= time.monotonic():
                    job[0] = time.monotonic() + job[1]
                    try:
                        job[2]()
                    except Exception:
                        log.exception("scheduled read failed")
            due = min((job[0] for job in jobs), default=None)
            self._wake.wait(None if due is None else max(0, due - time.monotonic()))
            self._wake.clear()


SCHEDULE = Schedule()


@dataclass(frozen=True)
class Snapshot:
    board: dict[str, Any] | None
    stale: bool


class Source:
    def __init__(
        self, base: str, prefix: str = LEADERBOARDS["release"], ttl: float = TTL_SECONDS
    ) -> None:
        self.base = check_base(base)
        self._path = f"{prefix}/latest.json"
        self._ttl = ttl
        self._lock = threading.Lock()
        self._snapshot = Snapshot(None, stale=False)
        self._remove: Callable[[], None] | None = None

    @classmethod
    def from_env(cls, variant: str = "release") -> Source:
        return cls(base_from_env(), LEADERBOARDS[variant])

    def get(self) -> Snapshot:
        # Only the first read waits for a page view; later ones run on the schedule, never on demand.
        if self._remove is None:
            with self._lock:
                if self._remove is None:
                    self._refresh()
                    self._remove = SCHEDULE.add(self._ttl, self._refresh)
        SCHEDULE.ensure_running()
        return self._snapshot

    def close(self) -> None:
        if self._remove is not None:
            self._remove()

    def _refresh(self) -> None:
        try:
            self._snapshot = Snapshot(parse(read(self.base, self._path)), stale=False)
        except Exception as exc:
            # Logged once per outage, not on every scheduled read.
            report = log.debug if self._snapshot.stale else log.warning
            report("leaderboard read failed: %s: %s", type(exc).__name__, exc)
            self._snapshot = Snapshot(self._snapshot.board, stale=True)


class History:
    def __init__(
        self,
        base: str,
        start: date,
        prefix: str = LEADERBOARDS["release"],
        ttl: float = HISTORY_TTL_SECONDS,
        today: Callable[[], date] = lambda: datetime.now(UTC).date(),
    ) -> None:
        self.base = check_base(base)
        self._start = start
        self._prefix = prefix
        self._ttl = ttl
        self._today = today
        self._lock = threading.Lock()
        self._days: dict[str, list[dict[str, Any]]] = {}
        self._remove: Callable[[], None] | None = None

    @classmethod
    def from_env(cls, variant: str = "release") -> History:
        return cls(base_from_env(), HISTORY_START, LEADERBOARDS[variant])

    def _read(self, day: str) -> list[dict[str, Any]] | None:
        try:
            data = read(self.base, f"{self._prefix}/history/{day}.json")
            board = parse(data, select=overall_board)
        except InvalidData as exc:
            log.warning("history %s rejected: %s", day, exc)
            return None
        except Exception as exc:
            log.debug("history %s not available: %s", day, type(exc).__name__)
            return None
        overall = board["boards"].get("overall", {"sample_count": 0, "rankings": []})
        samples = overall["sample_count"]
        return [
            {
                "date": day,
                "model_id": row["model_id"],
                **{
                    f"{m}_{k}": row["metrics"][m][k]
                    for m in ("mase", "sql")
                    for k in ("skill_score", "win_rate")
                },
                "coverage_samples": 1 - row["missing_samples"] / samples if samples else None,
                "coverage_weight": None
                if row["missing_weight"] is None
                else 1 - row["missing_weight"],
            }
            for row in overall["rankings"]
        ]

    def rows(self) -> list[dict[str, Any]]:
        if self._remove is None:
            with self._lock:
                if self._remove is None:
                    self._update()
                    self._remove = SCHEDULE.add(self._ttl, self._update)
        SCHEDULE.ensure_running()
        days = self._days
        return [row for day in sorted(days) for row in days[day]]

    def close(self) -> None:
        if self._remove is not None:
            self._remove()

    def _update(self) -> None:
        days = range((self._today() - self._start).days + 1)
        due = [
            d
            for d in ((self._start + timedelta(n)).isoformat() for n in days)
            if d not in self._days
        ]
        if not due:
            return
        with ThreadPoolExecutor(16) as pool:
            read = zip(due, pool.map(self._read, due), strict=True)
            found = {day: rows for day, rows in read if rows is not None}
        self._days = {**self._days, **found}


def main(paths: list[str]) -> int:
    status = 0
    for path in paths:
        try:
            board = parse(Path(path).read_bytes())
        except (OSError, InvalidData) as exc:
            print(f"{path}: INVALID: {exc}")
            status = 1
        else:
            print(f"{path}: OK, {board['timestamp']}, {len(board['boards'])} boards")
    return status


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))