code/local_train/check.py

from __future__ import annotations

import hashlib
import json
import subprocess
import sys
from pathlib import Path

from .constants import GENESIS_DIR, METADATA_FILES


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


def check_metadata(model_dir: Path, genesis_dir: Path = GENESIS_DIR) -> dict:
    model_dir = Path(model_dir)
    genesis_dir = Path(genesis_dir)
    files = {}
    ok = True
    for name in METADATA_FILES:
        left = model_dir / name
        right = genesis_dir / name
        if not left.is_file():
            files[name] = {"ok": False, "reason": "missing in model"}
            ok = False
            continue
        if not right.is_file():
            files[name] = {"ok": False, "reason": "missing in genesis"}
            ok = False
            continue
        left_hash = sha256_file(left)
        right_hash = sha256_file(right)
        match = left_hash == right_hash
        files[name] = {
            "ok": match,
            "model": left_hash,
            "genesis": right_hash,
            "reason": "" if match else "sha256 mismatch (not byte-identical to genesis)",
        }
        ok = ok and match
    return {"ok": ok, "files": files}


def check_model(model_dir: Path, genesis_dir: Path = GENESIS_DIR) -> dict:
    """sha256 metadata vs genesis, then `albedo check-model` (no OpenSearch dedup)."""
    model_dir = Path(model_dir)
    meta = check_metadata(model_dir, genesis_dir)
    cli = _albedo_check(model_dir)
    report = {
        "path": str(model_dir),
        "metadata_hash": meta,
        "check_model": cli,
        "ok": bool(meta["ok"] and cli.get("ok")),
        "note": "check-model does not run OpenSearch 0.95 fingerprint dedup or live metadata pins.",
    }
    print(json.dumps(report, indent=2), flush=True)
    return report


def _albedo_check(model_dir: Path) -> dict:
    cmd = [sys.executable, "-m", "miner.cli", "check-model", "--path", str(model_dir)]
    try:
        proc = subprocess.run(
            cmd,
            cwd=str(Path(__file__).resolve().parents[1]),
            capture_output=True,
            text=True,
            check=False,
        )
    except OSError as exc:
        return {"ok": False, "reason": str(exc), "stdout": "", "stderr": ""}
    stdout = proc.stdout or ""
    return {
        "ok": proc.returncode == 0 and _valid_line(stdout),
        "returncode": proc.returncode,
        "stdout": stdout,
        "stderr": proc.stderr or "",
    }


def _valid_line(stdout: str) -> bool:
    lines = [line.strip() for line in stdout.splitlines() if line.strip()]
    return bool(lines) and lines[-1] == "VALID"