kernels/p7_probe_publish_source.py

"""Does `--ckpt-repo` actually choose the repo the publisher reads its weights from?

Hard-coding `Cion-lab/ounce100m-ckpt` as the only source was harmless while that repo existed and there was one
model. It is not harmless now: a second run publishing to its own repo would still resolve `latest.json` from
the first run's -- i.e. publish one model's weights under another model's card -- and since that repo has been
deleted the failure arrives as a confusing 404 instead of a wrong answer.

Why a 404 cannot be the test. A traceback prints the failing *source line*, which contains the literal
`{a.ckpt_repo}`, not the URL it resolved to. "It died" therefore says nothing about which repo it asked for,
and a probe asserting on that text would pass whether or not the fix works. So the decisive read is the
publisher's own `ckpt latest.json -> {...}` print, which names the repo it fetched -- that needs a repo which
answers. The kernel creates a throwaway **private** dataset repo holding nothing but
`latest.json = {"step": 0, "path_in_repo": "final"}`, reads that name back out of the publisher's stdout, and
deletes the repo in a `finally`.

Four cases, free CPU, nothing published:
  1  static  -- the fetch is built from the argument; no fetch is built from the constant.
  2  help    -- `--help` lists `--ckpt-repo`.
  3  positive-- the new file, told to read the temp repo, prints that repo's name and step 0, and stops before
                any upload (the temp repo has no `final/`, which is the expected place for it to stop).
  4  contrast-- the old file at REV `55f23ef5` rejects `--ckpt-repo` outright (argparse rc 2), so case 3 is
                evidence about this change and not about something that was always true.

Layout matters: the publisher resolves `ounce100m_credentials` from its own grandparent directory
(`sys.path.insert(0, dirname(dirname(__file__)))`), and a subprocess is a fresh interpreter that inherits
nothing this kernel puts on `sys.path` -- so both copies go exactly where Phase 5 put them,
`<work>/build/publish_model.py` beside `<work>/ounce100m_credentials.py`. (E-056 was this same bug.)
"""
import hashlib
import io
import json
import os
import subprocess
import sys
import urllib.request

CODE = "Cion-lab/ounce100m-code"
NEW_REV = "dd1b6b94cc90ae6deff94e0990d58da64c6637f1"
NEW_SHA = "d1f3c4eb2b80cd4bff35c1b1b90a67ec6204dd753e09f354673d0812cc231e01"
OLD_REV = "55f23ef5d0e8e6a60adad23bf8a2ffffc6bc7580"
CRED_SHA = "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"
TMP = "Cion-lab/ounce100m-p7-probe-tmp"
WORK = "/kaggle/working"


def fetch(rev, path):
    url = f"https://huggingface.co/{CODE}/resolve/{rev}/{path}"
    return urllib.request.urlopen(url, timeout=120).read()


def report(R, why=""):
    print("PROBE_JSON_BEGIN\n" + json.dumps(R, indent=1) + "\nPROBE_JSON_END", flush=True)
    oks = [v for k, v in R.items() if k.startswith("case")]
    ok = bool(oks) and all(oks)
    print(f"VERDICT PUBLISH_SOURCE {'PASS' if ok else 'FAIL'} {why}", flush=True)
    return ok


R = {}
try:
    for d in ("build", "build_old"):
        os.makedirs(os.path.join(WORK, d), exist_ok=True)
    cb = fetch(NEW_REV, "ounce100m_credentials.py")
    assert hashlib.sha256(cb).hexdigest() == CRED_SHA, "credentials module is not the Phase 5 bytes"
    open(os.path.join(WORK, "ounce100m_credentials.py"), "wb").write(cb)
    sys.path.insert(0, WORK)
    import ounce100m_credentials
    ounce100m_credentials.install(verify=True)
    from huggingface_hub import HfApi
    api = HfApi()

    nb = fetch(NEW_REV, "build/publish_model.py")
    R["pinned_bytes_ok"] = hashlib.sha256(nb).hexdigest() == NEW_SHA
    assert R["pinned_bytes_ok"], "the pinned publisher is not the bytes this probe was written against"
    newp = os.path.join(WORK, "build", "publish_model.py")
    oldp = os.path.join(WORK, "build_old", "publish_model.py")
    open(newp, "wb").write(nb)
    open(oldp, "wb").write(fetch(OLD_REV, "build/publish_model.py"))

    src = nb.decode("utf-8")
    R["case1_static_arg_wired"] = "datasets/{a.ckpt_repo}/resolve/main/latest.json" in src
    R["case1_static_no_constant_fetch"] = ("datasets/{CKPT}/resolve" not in src
                                           and "snapshot_download(CKPT" not in src)
    R["case2_help_lists_flag"] = "--ckpt-repo" in subprocess.run(
        [sys.executable, newp, "--help"], capture_output=True, text=True).stdout

    tok = os.environ["HF_TOKEN"]
    api.create_repo(repo_id=TMP, repo_type="dataset", private=True, exist_ok=True, token=tok)
    api.upload_file(path_or_fileobj=io.BytesIO(json.dumps({"step": 0, "path_in_repo": "final"}).encode()),
                    path_in_repo="latest.json", repo_id=TMP, repo_type="dataset", token=tok)
    made = True

    r = subprocess.run([sys.executable, newp, "--dry-run", "--ckpt-repo", TMP],
                       capture_output=True, text=True)
    out = r.stdout + r.stderr
    R["case3_names_the_repo_it_was_given"] = f"'repo': '{TMP}'" in out
    R["case3_read_the_pointer"] = "'step': 0" in out
    R["case3_stopped_before_upload"] = r.returncode != 0 and "uploading" not in out
    ro = subprocess.run([sys.executable, oldp, "--dry-run", "--ckpt-repo", TMP],
                        capture_output=True, text=True)
    R["case4_old_file_rejects_the_flag"] = ro.returncode == 2 and "unrecognized arguments" in ro.stderr
    R["observed_tail"] = out.strip().splitlines()[-1][:160] if out.strip() else ""
except Exception as e:
    R["aborted"] = f"{type(e).__name__}: {str(e)[:200]}"
finally:
    if locals().get("made"):
        try:
            api.delete_repo(repo_id=TMP, repo_type="dataset")
            R["temp_repo_deleted"] = True
        except Exception as e:
            R["temp_repo_deleted"] = f"FAILED {type(e).__name__} -- delete {TMP} by hand"

sys.exit(0 if report(R, R.get("aborted", "")) else 1)