cloud/acceptance.py

"""Run an optional real model acceptance, isolated from production scores/state."""
from __future__ import annotations

from datetime import datetime, timezone
import json
import os
from pathlib import Path
import re
import signal
import sys
import tempfile
import time

ROOT = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(ROOT / "src"), str(ROOT)]
BRANCH = "acceptance"


def publish(api, repo, output, report):
    from huggingface_hub import CommitOperationAdd
    from tsfm_bench.remote_state import pack_files
    if not api.dataset_info(repo).private:
        raise ValueError("Acceptance artifacts require a private Dataset")
    api.create_branch(repo, branch=BRANCH, repo_type="dataset", exist_ok=True)
    parent = api.dataset_info(repo, revision=BRANCH).sha
    prefix = "acceptance/" + output.name
    with tempfile.TemporaryDirectory(prefix="livehouse-acceptance-") as temp:
        archive = Path(temp) / "artifacts.tar.gz"
        manifest = pack_files(ROOT, sorted(p for p in output.rglob("*") if p.is_file()), archive)
        return api.create_commit(repo, repo_type="dataset", revision=BRANCH, parent_commit=parent,
            commit_message=f"Cloud model acceptance {output.name}: {report['status']}", operations=[
                CommitOperationAdd(path_in_repo=prefix + "/artifacts.tar.gz", path_or_fileobj=archive),
                CommitOperationAdd(path_in_repo=prefix + "/manifest.json", path_or_fileobj=json.dumps(manifest).encode()),
                CommitOperationAdd(path_in_repo=prefix + "/report.json", path_or_fileobj=json.dumps(report, indent=2).encode()),
            ])


def main():
    from huggingface_hub import HfApi, hf_hub_download
    from scripts.run_online_daemon import run_managed_subprocess
    from tsfm_bench.remote_state import restore_archive
    identifier = os.environ["LIVEHOUSE_ACCEPTANCE_ID"]
    if not re.fullmatch(r"[a-z0-9][a-z0-9-]{1,60}", identifier):
        raise ValueError("Acceptance ID must contain only lowercase letters, digits and hyphens")
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
    repo = os.environ["HF_STATE_REPO"]
    token = os.environ["HF_TOKEN"]
    api = HfApi(token=token)
    if not api.dataset_info(repo).private:
        raise ValueError("Acceptance artifacts require a private Dataset")
    output = ROOT / "outputs/cloud-acceptance" / identifier
    output.mkdir(parents=True, exist_ok=True)
    prefix = "acceptance/" + identifier
    api.create_branch(repo, branch=BRANCH, repo_type="dataset", exist_ok=True)
    revision = api.dataset_info(repo, revision=BRANCH).sha
    def download(name):
        return Path(hf_hub_download(repo, prefix + "/" + name, repo_type="dataset",
                                   revision=revision, token=token))
    report = {"id":identifier, "isolated":True, "space":os.getenv("SPACE_ID"),
              "source_revision":os.environ["LIVEHOUSE_SOURCE_REVISION"],
              "started_at":datetime.now(timezone.utc).isoformat()}
    if api.file_exists(repo, prefix + "/report.json", repo_type="dataset", revision=revision):
        report = json.loads(download("report.json").read_text())
        if report["status"] in {"complete", "failed"}:
            print(f"Acceptance {identifier} already {report['status']}; no inference repeated", flush=True)
            if report["status"] == "failed":
                raise RuntimeError("This acceptance already failed; inspect its archived report")
            return
        manifest = json.loads(download("manifest.json").read_text())
        restore_archive(ROOT, download("artifacts.tar.gz"), manifest, allowed_roots=(output.relative_to(ROOT),))
    def run(phase):
        return run_managed_subprocess([sys.executable, "scripts/live_acceptance.py", phase,
            "--output-root", str(output)], timeout=1200, timeout_message="Cloud acceptance timed out")
    try:
        if report.get("status") != "issued":
            if run("issue") != 0:
                raise RuntimeError("At least one real model forecast failed; inspect archived acceptance-issue.json")
            report.update(status="issued", issue=json.loads((output / "acceptance-issue.json").read_text()))
            publish(api, repo, output, report)
        # The target bar starts at the forecast deadline and closes five minutes later.
        deadline = datetime.fromisoformat(report["issue"]["forecast_deadline"]).timestamp() + 310
        while time.time() < deadline:
            time.sleep(min(15, deadline - time.time()))
        for attempt in range(4):
            if run("resolve") == 0:
                break
            if attempt == 3:
                raise RuntimeError("Real future observations did not resolve every model")
            time.sleep(30)
        report.update(status="complete", resolve=json.loads((output / "acceptance-resolve.json").read_text()))
    except Exception as exc:
        report.update(status="failed", error_type=type(exc).__name__)
    report["finished_at"] = datetime.now(timezone.utc).isoformat()
    publish(api, repo, output, report)
    print(f"Cloud acceptance {identifier}: {report['status']}; artifacts on private Dataset branch {BRANCH}", flush=True)
    if report["status"] != "complete":
        raise RuntimeError("Cloud model acceptance failed; inspect its private artifacts")


if __name__ == "__main__":
    main()