# Phase 6 launcher with a heartbeat. Kaggle kernel body: 2xT4, billed 1x wall clock.
#
# Why this exists: run_benchmarks.py is silent until its first push (the smoke pass writes smoke.json only
# after all eight smoke runs), and on this account `kernels.get` is denied, so a kernel that is queued,
# booting, installing or dead all look identical from outside. A heartbeat in the report repo makes those
# states distinguishable and costs a few kilobytes per push.
#
# Before launching: book the QUOTA.md row, then set REV/PUB/SHA pins to the files actually published at that
# revision, and read the pins back after fetch -- a kernel that runs code it did not verify is not reproducible.
import hashlib, os, runpy, subprocess, sys, threading, time

REV = "55f23ef5d0e8e6a60adad23bf8a2ffffc6bc7580"
RB_SHA = "8196d853f3e3d8f1c729ff0d76517ddc7814b9d460a09497e1f289b03ca8da69"
CRED_SHA = "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"
BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
REPORT_REPO = "Cion-lab/ounce100m-p5-report"          # private scratch; keeps the deliverable repo clean
LOG = "/kaggle/working/p6.log"
STAGE = {"name": "submit", "detail": ""}


def log(msg):
    line = "%s %s\n" % (time.strftime("%H:%M:%S", time.gmtime()), msg)
    with open(LOG, "a") as fh:
        fh.write(line)
    sys.__stdout__.write(line)
    sys.__stdout__.flush()


def push_heartbeat(api, tok):
    try:
        body = open(LOG, "rb").read()[-60000:]
        api.upload_file(path_or_fileobj=body, path_in_repo="p6_heartbeat.log",
                        repo_id=REPORT_REPO, repo_type="dataset",
                        commit_message="heartbeat %s" % STAGE["name"], token=tok)
    except Exception as e:                                # a lost beat must never kill the eval
        log("heartbeat push failed %r" % (e,))


def beat(api, tok):
    while True:
        time.sleep(120)
        push_heartbeat(api, tok)


os.chdir("/kaggle/working")
os.makedirs("/kaggle/working/lib/eval", exist_ok=True)
for rel, want in (("eval/run_benchmarks.py", RB_SHA), ("ounce100m_credentials.py", CRED_SHA)):
    dest = "/kaggle/working/lib/" + rel
    assert subprocess.run(["curl", "-sfL", BASE + "/" + rel, "-o", dest]).returncode == 0, rel
    got = hashlib.sha256(open(dest, "rb").read()).hexdigest()
    assert got == want, "%s: %s != %s" % (rel, got[:16], want[:16])
    log("PINNED %s %s" % (rel, got[:16]))

sys.path.insert(0, "/kaggle/working/lib")
import ounce100m_credentials

STAGE["name"] = "creds"
log("creds " + str(ounce100m_credentials.install(verify=True)))
tok = os.environ["HF_TOKEN"]

from huggingface_hub import HfApi

api = HfApi(token=tok)
threading.Thread(target=beat, args=(api, tok), daemon=True).start()
push_heartbeat(api, tok)                                   # first beat immediately: the kernel is alive

import site

up = site.getusersitepackages()                            # E-047: ~/.local is not on sys.path when created
if up and up not in sys.path:
    sys.path.append(up)                                    # mid-process, so add it explicitly
log("user_site " + str(up))

# tee: the runner streams its own progress, and that stream is the only status a denied-kernels.get account
# can still see once it is being shipped somewhere.
real = sys.stdout


class Tee:
    def write(self, s):
        real.write(s)
        real.flush()
        with open(LOG, "a") as fh:
            fh.write(s)
        return len(s)

    def flush(self):
        real.flush()

    def isatty(self):                                      # transformers' progress bars ask
        return False


STAGE["name"] = "run_benchmarks"
sys.argv = ["run_benchmarks.py", "--model", "Cion-lab/ounce100m-v1",
            "--repo", "Cion-lab/ounce100m-v1", "--full-gpu-hours", "0.3"]
log("ARGV " + " ".join(sys.argv))
sys.stdout = Tee()
err = None
try:
    runpy.run_path("/kaggle/working/lib/eval/run_benchmarks.py", run_name="__main__")
except SystemExit as e:
    err = None if e.code in (0, None) else "SystemExit(%r)" % (e.code,)
except BaseException as e:
    err = repr(e)
finally:
    sys.stdout = real
    STAGE["name"] = "done"
    log("P6_DONE err=%s" % (err,))
    push_heartbeat(api, tok)
print("VERDICT P6 err=%s" % (bool(err),), flush=True)