scripts/build_data.py

#!/usr/bin/env python3
"""Build the dashboard's public JSON payload from a W&B project."""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
from collections import defaultdict
from datetime import datetime, timezone
from itertools import islice

_SCORE_PREFIX = "miner/"
_SCORE_SUFFIX = "/score"
_LEN_SUFFIX = "/completion_len"
_CORRECTNESS_SUFFIX = "/correctness_score"
_ORIG_SCORE = "original/score"
_ORIG_CORRECTNESS = "original/correctness_score"
_ORIG_LEN = "original/completion_len"
_PROGRESS_STAGES = {"qualification", "full_evaluation"}
_PROGRESS_STATUSES = {"running", "completed", "failed"}
_STAGE_STATUSES = {
    "pending", "waiting", "preparing", "evaluating", "completed", "skipped", "failed"
}
_MINER_STATUSES = {
    "pending", "waiting", "evaluating", "finished", "failed", "skipped",
    "not_selected", "rejected",
}
_BASELINE_STATUSES = {
    "pending", "waiting", "evaluating", "finished", "failed", "skipped",
}


def _miner_hotkey_from_score_key(key: str) -> str | None:
    if key.startswith(_SCORE_PREFIX) and key.endswith(_SCORE_SUFFIX):
        hk = key[len(_SCORE_PREFIX) : -len(_SCORE_SUFFIX)]
        return hk or None
    return None


def _as_float(value) -> float | None:
    try:
        f = float(value)
    except (TypeError, ValueError):
        return None
    return f if math.isfinite(f) else None


def _score_to_signed_correctness(value) -> float | None:
    score = _as_float(value)
    if score is None:
        return None
    return (2.0 * score) - 1.0


def extract_rankings(summary: dict) -> list[dict]:
    """Return the current miner ranking from a run summary."""
    out = []
    for key, value in summary.items():
        hk = _miner_hotkey_from_score_key(key)
        if hk is None:
            continue
        score = _as_float(value)
        if score is None:
            continue
        out.append({"hotkey": hk, "score": score})
    out.sort(key=lambda r: r["score"], reverse=True)
    return out


def _summary_dict(run) -> dict:
    summary = getattr(run.summary, "_json_dict", run.summary)
    return dict(summary)


def extract_progress(summary: dict) -> dict | None:
    raw = summary.get("progress/snapshot")
    if isinstance(raw, str):
        try:
            raw = json.loads(raw)
        except json.JSONDecodeError:
            return None
    if not isinstance(raw, dict):
        return None

    epoch = _as_float(raw.get("epoch"))
    status = raw.get("status")
    if epoch is None or status not in _PROGRESS_STATUSES:
        return None

    stages = {}
    raw_stages = raw.get("stages")
    if isinstance(raw_stages, dict):
        for stage_name in _PROGRESS_STAGES:
            raw_stage = raw_stages.get(stage_name)
            if not isinstance(raw_stage, dict):
                continue
            stage_status = raw_stage.get("status")
            if stage_status not in _STAGE_STATUSES:
                stage_status = "pending"
            miners = {}
            raw_miners = raw_stage.get("miners")
            if isinstance(raw_miners, dict):
                for hotkey, miner_status in list(raw_miners.items())[:512]:
                    if (
                        isinstance(hotkey, str)
                        and hotkey
                        and miner_status in _MINER_STATUSES
                    ):
                        miners[hotkey] = miner_status
            stage = {"status": stage_status, "miners": miners}
            baseline_status = raw_stage.get("baseline")
            if baseline_status in _BASELINE_STATUSES:
                stage["baseline"] = baseline_status
            scores = {}
            raw_scores = raw_stage.get("scores")
            if isinstance(raw_scores, dict):
                for hotkey, score in list(raw_scores.items())[:512]:
                    if not isinstance(hotkey, str) or not hotkey:
                        continue
                    value = _as_float(score)
                    if value is not None:
                        scores[hotkey] = value
            if scores:
                stage["scores"] = scores
            stages[stage_name] = stage

    current_stage = raw.get("current_stage")
    if current_stage not in _PROGRESS_STAGES:
        current_stage = None
    updated_at = raw.get("updated_at")
    return {
        "schema_version": 1,
        "epoch": int(epoch) if epoch == int(epoch) else epoch,
        "status": status,
        "current_stage": current_stage,
        "updated_at": updated_at[:64] if isinstance(updated_at, str) else None,
        "stages": stages,
    }


def _parse_iso(value) -> datetime | None:
    if not isinstance(value, str):
        return None
    try:
        return datetime.fromisoformat(value)
    except ValueError:
        return None


def _snapshot_for_epoch(raw, epoch: float) -> dict | None:
    if isinstance(raw, str):
        try:
            raw = json.loads(raw)
        except json.JSONDecodeError:
            return None
    if not isinstance(raw, dict):
        return None
    row_epoch = _as_float(raw.get("epoch"))
    if row_epoch is None or row_epoch != epoch:
        return None
    return raw


def extract_stage_timings(run, epoch) -> dict[str, dict]:
    """Replay a run's progress history to time how long each stage's baseline
    evaluation actually took for the given epoch.

    The live progress snapshot (summary) only ever holds the *latest* stage
    state, so a fast baseline/miner pass can flip from "evaluating" to
    "finished" between two dashboard polls without ever being observed live.
    Scanning history lets the dashboard show concrete proof ("base done in
    4.2s") instead of silently losing that transition.
    """
    epoch_value = _as_float(epoch)
    if epoch_value is None:
        return {}

    active_stage: dict[str, datetime] = {}
    active_baseline: dict[str, datetime] = {}
    timings: dict[str, dict] = {}

    try:
        rows = run.scan_history()
    except Exception:
        return {}

    for row in rows:
        snapshot = _snapshot_for_epoch(row.get("progress/snapshot"), epoch_value)
        if snapshot is None:
            continue
        ts = _parse_iso(snapshot.get("updated_at"))
        raw_stages = snapshot.get("stages")
        if not isinstance(raw_stages, dict):
            continue
        for stage_name in _PROGRESS_STAGES:
            stage = raw_stages.get(stage_name)
            if not isinstance(stage, dict):
                continue
            status = stage.get("status")
            if status == "evaluating" and stage_name not in active_stage:
                active_stage[stage_name] = ts
            elif (
                status in {"completed", "skipped", "failed"}
                and stage_name in active_stage
                and "duration_seconds" not in timings.get(stage_name, {})
            ):
                start = active_stage[stage_name]
                if start is not None and ts is not None:
                    timings.setdefault(stage_name, {})["duration_seconds"] = (
                        ts - start
                    ).total_seconds()

            baseline = stage.get("baseline")
            if baseline == "evaluating" and stage_name not in active_baseline:
                active_baseline[stage_name] = ts
            elif (
                baseline in {"finished", "skipped", "failed"}
                and stage_name in active_baseline
                and "baseline_seconds" not in timings.get(stage_name, {})
            ):
                start = active_baseline[stage_name]
                if start is not None and ts is not None:
                    timings.setdefault(stage_name, {})["baseline_seconds"] = (
                        ts - start
                    ).total_seconds()

    return timings


def progress_for_run(run, summary: dict) -> dict | None:
    progress = extract_progress(summary)
    if progress is None:
        return None
    if progress["status"] == "running" and getattr(run, "state", None) in {
        "crashed",
        "failed",
        "killed",
    }:
        progress["status"] = "failed"
    timings = extract_stage_timings(run, progress["epoch"])
    for stage_name, timing in timings.items():
        stage = progress["stages"].get(stage_name)
        if isinstance(stage, dict):
            stage.update(timing)
    return progress


def extract_history(run) -> dict:
    """Return aligned per-epoch miner and base-model series for one run."""
    rows_by_epoch: dict[float, dict] = {}
    stage_scores_by_epoch: dict[float, dict[str, dict[str, float]]] = {}
    miner_keys: set[str] = set()
    saw_orig_score = saw_orig_len = False

    for row in run.scan_history():
        progress = extract_progress(row)
        if progress is not None:
            progress_epoch = _as_float(progress.get("epoch"))
            if progress_epoch is not None:
                epoch_scores = stage_scores_by_epoch.setdefault(progress_epoch, {})
                for stage_name in _PROGRESS_STAGES:
                    scores = progress.get("stages", {}).get(stage_name, {}).get("scores")
                    if scores:
                        epoch_scores.setdefault(stage_name, {}).update(scores)

        epoch = _as_float(row.get("epoch"))
        if epoch is None:
            continue
        row_metrics = {}
        for key, value in row.items():
            hk = _miner_hotkey_from_score_key(key)
            if hk is not None:
                row_metrics[("miner", hk, "score")] = _as_float(value)
                miner_keys.add(hk)
            elif key.startswith(_SCORE_PREFIX) and key.endswith(_LEN_SUFFIX):
                hk2 = key[len(_SCORE_PREFIX) : -len(_LEN_SUFFIX)]
                if hk2:
                    row_metrics[("miner", hk2, "completion_len")] = _as_float(value)
                    miner_keys.add(hk2)
            elif key.startswith(_SCORE_PREFIX) and key.endswith(_CORRECTNESS_SUFFIX):
                hk2 = key[len(_SCORE_PREFIX) : -len(_CORRECTNESS_SUFFIX)]
                if hk2:
                    row_metrics[("miner", hk2, "correctness_score")] = _as_float(value)
                    miner_keys.add(hk2)
            elif key == _ORIG_SCORE:
                row_metrics[("orig", "score")] = _as_float(value)
                row_metrics.setdefault(
                    ("orig", "correctness_score"), _score_to_signed_correctness(value)
                )
                saw_orig_score = True
            elif key == _ORIG_CORRECTNESS:
                row_metrics[("orig", "correctness_score")] = _as_float(value)
                saw_orig_score = True
            elif key == _ORIG_LEN:
                row_metrics[("orig", "completion_len")] = _as_float(value)
                saw_orig_len = True
        if any(value is not None for value in row_metrics.values()):
            rows_by_epoch.setdefault(epoch, {}).update(row_metrics)

    epochs = sorted(rows_by_epoch)
    miners: dict[str, dict] = {}
    for hk in miner_keys:
        miners[hk] = {
            "score": [rows_by_epoch[e].get(("miner", hk, "score")) for e in epochs],
            "correctness_score": [
                rows_by_epoch[e].get(("miner", hk, "correctness_score")) for e in epochs
            ],
            "completion_len": [rows_by_epoch[e].get(("miner", hk, "completion_len")) for e in epochs],
        }
    original = {}
    if saw_orig_score:
        original["score"] = [rows_by_epoch[e].get(("orig", "score")) for e in epochs]
        original["correctness_score"] = [
            rows_by_epoch[e].get(("orig", "correctness_score")) for e in epochs
        ]
    if saw_orig_len:
        original["completion_len"] = [rows_by_epoch[e].get(("orig", "completion_len")) for e in epochs]

    stage_scores = [
        {
            "epoch": int(epoch) if epoch == int(epoch) else epoch,
            **stage_scores_by_epoch[epoch],
        }
        for epoch in sorted(stage_scores_by_epoch)
        if stage_scores_by_epoch[epoch]
    ]
    return {
        "epochs": [int(e) if e == int(e) else e for e in epochs],
        "miners": miners,
        "original": original,
        "stage_scores": stage_scores,
    }


def rankings_from_history(history: dict) -> list[dict]:
    epochs = history.get("epochs", [])
    if not epochs:
        return []
    latest_index = len(epochs) - 1
    rankings = []
    for hotkey, series in history.get("miners", {}).items():
        scores = series.get("score", [])
        score = scores[latest_index] if latest_index < len(scores) else None
        if score is not None:
            rankings.append({"hotkey": hotkey, "score": score})
    rankings.sort(key=lambda row: row["score"], reverse=True)
    return rankings


def validator_payload(run, hotkey: str) -> dict:
    history = extract_history(run)
    summary = _summary_dict(run)
    rankings = rankings_from_history(history) or (
        extract_rankings(summary) if not history.get("epochs") else []
    )
    payload = {"hotkey": hotkey, "rankings": rankings, "history": history}
    progress = progress_for_run(run, summary)
    if progress is not None:
        payload["progress"] = progress
    return payload


def build_aggregate(validators: list[dict]) -> dict:
    """Average miner scores across validators that ranked each miner."""
    score_sums: dict[str, float] = defaultdict(float)
    score_counts: dict[str, int] = defaultdict(int)
    for v in validators:
        for r in v["rankings"]:
            score_sums[r["hotkey"]] += r["score"]
            score_counts[r["hotkey"]] += 1

    rankings = [
        {"hotkey": hk, "score": score_sums[hk] / score_counts[hk], "validators": score_counts[hk]}
        for hk in score_sums
    ]
    rankings.sort(key=lambda r: r["score"], reverse=True)

    history = {"epochs": [], "miners": {}, "original": {}}
    if rankings:
        top = rankings[0]["hotkey"]
        history = _mean_history(validators, top)
    return {"rankings": rankings, "history": history}


def _mean_history(validators: list[dict], top_hotkey: str) -> dict:
    """Average per-epoch history for the leading miner and base model."""
    epochs = sorted({e for v in validators for e in v["history"].get("epochs", [])})

    def mean_series(getter):
        out = []
        for e in epochs:
            vals = []
            for v in validators:
                hist = v["history"]
                if e not in hist.get("epochs", []):
                    continue
                pos = hist["epochs"].index(e)
                val = getter(hist, pos)
                if val is not None and val == val:
                    vals.append(val)
            out.append(sum(vals) / len(vals) if vals else None)
        return out

    miner_score = mean_series(lambda h, p: (h["miners"].get(top_hotkey, {}).get("score") or [None] * (p + 1))[p] if h["miners"].get(top_hotkey) else None)
    miner_correctness = mean_series(lambda h, p: (h["miners"].get(top_hotkey, {}).get("correctness_score") or [None] * (p + 1))[p] if h["miners"].get(top_hotkey) else None)
    miner_len = mean_series(lambda h, p: (h["miners"].get(top_hotkey, {}).get("completion_len") or [None] * (p + 1))[p] if h["miners"].get(top_hotkey) else None)
    orig_score = mean_series(lambda h, p: (h["original"].get("score") or [None] * (p + 1))[p] if h["original"].get("score") else None)
    orig_correctness = mean_series(lambda h, p: (h["original"].get("correctness_score") or [None] * (p + 1))[p] if h["original"].get("correctness_score") else None)
    orig_len = mean_series(lambda h, p: (h["original"].get("completion_len") or [None] * (p + 1))[p] if h["original"].get("completion_len") else None)

    original = {}
    if any(x is not None for x in orig_score):
        original["score"] = orig_score
    if any(x is not None for x in orig_correctness):
        original["correctness_score"] = orig_correctness
    if any(x is not None for x in orig_len):
        original["completion_len"] = orig_len
    return {
        "epochs": epochs,
        "miners": {
            top_hotkey: {
                "score": miner_score,
                "correctness_score": miner_correctness,
                "completion_len": miner_len,
            }
        },
        "original": original,
    }


def wandb_path(project: str, entity: str | None) -> str:
    project = project.strip()
    entity = entity.strip() if entity else None
    if "/" in project:
        project_entity, project_name = (part.strip() for part in project.rsplit("/", 1))
        if not project_entity or not project_name:
            raise ValueError("W&B project must be 'project' or 'entity/project'")
        project = project_name
        entity = entity or project_entity
    if not project:
        raise ValueError("W&B project cannot be empty")
    return f"{entity}/{project}" if entity else project


def _is_running_run(run) -> bool:
    return getattr(run, "state", None) == "running"


def build(project: str, entity: str | None, api=None) -> dict:
    if api is None:
        import wandb

        api = wandb.Api()
    path = wandb_path(project, entity)

    order: list[str] = []
    fallbacks: dict[str, dict] = {}
    selected: dict[str, dict] = {}
    for run in api.runs(path, order="-created_at"):
        if not _is_running_run(run):
            continue
        hotkey = dict(run.config).get("validator_hotkey")
        if not isinstance(hotkey, str) or not hotkey or hotkey in selected:
            continue
        payload = validator_payload(run, hotkey)
        if hotkey not in fallbacks:
            order.append(hotkey)
            fallbacks[hotkey] = payload
        if payload["rankings"]:
            selected[hotkey] = payload

    validators = []
    for hotkey in order:
        latest = fallbacks[hotkey]
        chosen = dict(selected.get(hotkey, latest))
        if "progress" in latest:
            chosen["progress"] = latest["progress"]
        else:
            chosen.pop("progress", None)
        validators.append(chosen)

    return {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "validators": validators,
        "aggregate": build_aggregate(validators),
    }


def build_progress(
    project: str, entity: str | None, api=None, run_limit: int = 500
) -> dict[str, dict]:
    if api is None:
        import wandb

        api = wandb.Api()
    path = wandb_path(project, entity)
    progress_by_validator = {}
    seen = set()
    runs = api.runs(path, order="-created_at")
    for run in islice(runs, max(1, run_limit)):
        if not _is_running_run(run):
            continue
        hotkey = dict(run.config).get("validator_hotkey")
        if not isinstance(hotkey, str) or not hotkey or hotkey in seen:
            continue
        seen.add(hotkey)
        progress = progress_for_run(run, _summary_dict(run))
        if progress is not None:
            progress_by_validator[hotkey] = progress
    return progress_by_validator


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate the static dashboard's data.json from wandb.")
    parser.add_argument("--wandb-project", required=True, help="Shared wandb project every validator logs into")
    parser.add_argument("--wandb-entity", default=None, help="Wandb entity/team (default: the API key's default entity)")
    parser.add_argument("--out", default="data.json", help="Output path for the generated JSON")
    args = parser.parse_args(argv)

    key = os.environ.get("WANDB_KEY") or os.environ.get("WANDB_API_KEY")
    if not key:
        print("error: set WANDB_KEY (or WANDB_API_KEY) in the environment", file=sys.stderr)
        return 2
    os.environ["WANDB_API_KEY"] = key

    data = build(args.wandb_project, args.wandb_entity)
    with open(args.out, "w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2)
        fh.write("\n")
    print(f"wrote {args.out}: {len(data['validators'])} validators, "
          f"{len(data['aggregate']['rankings'])} ranked miners")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())