train/hubckpt.py

# The §3.13 storage cycle as one callable: write -> upload -> verify FROM THE HUB -> delete local ->
# confirm free space. Every step is separate and none of them trusts the previous one.

# Why verification cannot be "the upload call returned": `upload_folder` returns a commit sha, which
# proves the Hub accepted a commit, not that the bytes are retrievable. For LFS-tracked files the commit
# is a pointer. So verification here means: every LFS file's stored `lfs.oid` -- which *is* the sha256 of
# the object the server holds -- equals the sha256 computed on the local bytes, every size agrees, and at
# least one file is then actually fetched back over the download path and compared byte for byte. A
# checkpoint that cannot be read back is not a checkpoint, and the run must stop rather than delete the
# only local copy. (Git-stored small files report a sha1 blob id, not a content sha256, so for those the
# size is all the listing proves; `verify_checkpoint` says so in `content_sha_comparable` rather than
# implying more than it has.)
#
# Used by the trainer (main run + Phase 3 preflight) and by the preflight report scripts. No credentials
# in this file: the token comes from the private store via ounce100m_credentials (D-006).

import hashlib
import json
import os
import shutil
import urllib.error
import urllib.request

HUB = "https://huggingface.co"


def sha256_file(path, block=1 << 20):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while True:
            b = f.read(block)
            if not b:
                break
            h.update(b)
    return h.hexdigest()


def local_manifest(dirpath):
    """{relpath: (bytes, sha256)} for everything under dirpath.

    A file that disappears between the walk and the hash is skipped with a warning rather than raising:
    it then shows up as a Hub-side `extra` in the comparison, which fails verification loudly and refuses
    the prune. Crashing in getsize instead (E-029) killed a completed training leg at its last line.
    """
    out = {}
    for root, _dirs, files in os.walk(dirpath):
        for fn in files:
            p = os.path.join(root, fn)
            rel = os.path.relpath(p, dirpath).replace(os.sep, "/")
            try:
                out[rel] = (os.path.getsize(p), sha256_file(p))
            except FileNotFoundError:
                print(f"   WARNING: {rel} vanished while hashing {dirpath} -- "
                      "something else is writing here", flush=True)
    return out


class RepoMissing(RuntimeError):
    """The Hub says this repository does not exist -- distinct from "the listing failed", which is what
    RuntimeError means here. Callers that treat a failure as an empty repo would restart training at step
    0 (E-018), so the two must stay distinguishable at the except clause."""


def hub_listing(repo, repo_type="dataset", revision="main", token=None, page=None):
    """{path_in_repo: (bytes, sha256_or_None)} straight from the Hub's tree endpoint, ALL pages.

    Recursive, because checkpoints are directories. The endpoint is paginated (`Link: rel="next"`), and a
    repo holding ten checkpoints is comfortably past one page: reading only the first page would make
    `verify_checkpoint` report phantom `missing` files, which blocks pruning and raises mid-run, and would
    make `download_checkpoint` silently under-download a checkpoint. Both are worse than an error.

    Raises rather than returning {} when the listing fails -- E-018: 'broken' and 'empty' are different
    answers and callers must see which one happened.
    """
    url = page or f"{HUB}/api/{repo_type}s/{repo}/tree/{revision}?recursive=true"
    req = urllib.request.Request(url)
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    try:
        with urllib.request.urlopen(req, timeout=120) as r:
            entries = json.loads(r.read().decode())
            link = r.headers.get("Link") or ""
    except urllib.error.HTTPError as e:
        # 404 here means the repository itself has never been created, which is a different fact from a
        # failed lookup -- and the first Phase 4 session depends on telling them apart.
        if e.code == 404:
            raise RepoMissing(f"hub repo {repo} does not exist (HTTP 404)") from e
        raise RuntimeError(f"hub listing failed for {repo}@{revision}: HTTPError {e.code}: "
                           f"{str(e)[:160]}") from e
    except Exception as e:
        raise RuntimeError(f"hub listing failed for {repo}@{revision}: {type(e).__name__}: "
                           f"{str(e)[:160]}") from e
    if not isinstance(entries, list):
        raise RuntimeError(f"hub listing for {repo}@{revision} was not a list: {str(entries)[:160]}")
    out = {}
    for e in entries:
        if e.get("type") != "file":
            continue
        sha = (e.get("lfs") or {}).get("oid") or e.get("sha")
        sz = e.get("size")
        # An absent size must not become -1: callers compare it against a real byte count and a -1 turns
        # into a bogus "hub lists -1 bytes, N arrived" abort.
        out[e["path"]] = (int(sz) if isinstance(sz, (int, float)) else None, sha)
    nxt = _next_link(link)
    if nxt:
        out.update(hub_listing(repo, repo_type, revision, token, page=nxt))
    return out


def _next_link(header):
    """Follow the RFC 5988 Link header's rel="next" if the Hub sent one."""
    for part in header.split(","):
        if 'rel="next"' in part.replace(" ", "") or "rel=next" in part.replace(" ", ""):
            i = part.find("<")
            j = part.find(">", i)
            if i >= 0 and j > i:
                return part[i + 1:j]
    return None


def free_bytes(path="/kaggle/working"):
    return shutil.disk_usage(path).free


def verify_checkpoint(repo, dirpath, path_in_repo, token=None, repo_type="dataset",
                      readback_one=True):
    """Compare local bytes against the Hub's view, then prove one file really downloads.

    Returns a dict with ok/missing/mismatch/extra; never raises for a content problem, because the
    caller has to log the finding before deciding whether to prune.
    """
    prefix = path_in_repo.rstrip("/") + "/"
    want = local_manifest(dirpath)
    have = {k: v for k, v in hub_listing(repo, repo_type, token=token).items() if k.startswith(prefix)}
    have = {k[len(prefix):]: v for k, v in have.items()}
    missing = sorted(set(want) - set(have))
    extra = sorted(set(have) - set(want))
    mismatch = []
    for k in sorted(set(want) & set(have)):
        lsz, lsha = want[k]
        hsz, hsha = have[k]
        # The Hub reports LFS oids as sha256; small files come back as a git sha1, which is not
        # comparable, so a size match carries those and the read-back below carries the rest.
        comparable = hsha and len(hsha) == 64
        if (hsz is not None and hsz != lsz) or (comparable and hsha != lsha):
            mismatch.append({"path": k, "local": [lsz, lsha[:16]], "hub": [hsz, (hsha or "")[:16]]})
    rb = None
    rb_ok = False
    if readback_one and have:
        # Read back a bounded slice of the LARGEST file, from two offsets, and compare it byte for byte.
        # The whole-file content proof is already above: for every LFS object the Hub's `lfs.oid` *is* the
        # sha256 of the bytes it stores, and that is compared against the local sha256. What the listing
        # cannot establish is the retrieval path -- that `resolve/` really serves this object, because the
        # oid is what the commit *claims*. Picking the cheapest file (cursor.json, a few dozen git-stored
        # bytes) proved the git tree while leaving 1.7 GB of weights untested, and `ok=True` then
        # authorised deleting the only local copy (E-031).
        #
        # It used to download the entire largest file to prove that, and probe v5 measured what the tail
        # of that costs on this link: `CKPT 10 ... push+verify=797.7s` for an 850 MB optimizer.pt, against
        # 26 s for the identical cycle in the next leg when the CDN was warm. The run has eleven of those
        # cycles and the weekly quota has ~4 hours of slack, so a cold cache must not be able to eat two
        # of them. Two 4 MiB ranges -- head and mid-file, since a truncated object still has a good head --
        # answer the same question, and a response that is not 206 with exactly the bytes asked for is a
        # failure rather than a silent full download.
        pick = max(have, key=lambda k: (have[k][0] or 0))
        size = have[pick][0]
        url = f"{HUB}/{repo_type}s/{repo}/resolve/main/{prefix}{pick}"
        chunk = 4 << 20
        rb = {"path": pick, "url": url,
              "content_sha_comparable": bool(pick in want and have[pick][1] and
                                             len(have[pick][1]) == 64 and have[pick][1] == want[pick][1])}
        spans = []
        if not size or size <= 0:
            # `hub_listing` allows a None size, and a 0-byte span would "verify" by comparing nothing.
            rb["error"] = f"the Hub lists {pick} as {size!r} bytes -- nothing to read back"
        else:
            spans = [(0, min(chunk, size) - 1)] if size <= chunk else \
                    [(0, chunk - 1), (size // 2, size // 2 + chunk - 1)]
        try:
            with open(os.path.join(dirpath, pick), "rb") as loc:
                for start, end in spans:
                    rq = urllib.request.Request(url)
                    if token:
                        rq.add_header("Authorization", f"Bearer {token}")
                    rq.add_header("Range", f"bytes={start}-{end}")
                    with urllib.request.urlopen(rq, timeout=180) as fh:
                        status = getattr(fh, "status", 200)
                        crange = fh.headers.get("Content-Range")
                        data = fh.read((end - start + 1) + 1)
                    want_n = end - start + 1
                    if status != 206 or len(data) != want_n:
                        raise RuntimeError(f"range {start}-{end} returned HTTP {status} with "
                                           f"{len(data)} bytes; expected 206 with exactly {want_n}")
                    # The length alone is not the offset: a server that answers with the right number of
                    # the WRONG bytes passes a length check, and for a file that is mostly padding (the
                    # Phase 3 preflight pushed 4x425 MB of zeros) it passes the byte compare too.
                    if crange != f"bytes {start}-{end}/{size}":
                        raise RuntimeError(f"range {start}-{end} came back as Content-Range {crange!r}, "
                                           f"expected {f'bytes {start}-{end}/{size}'!r}")
                    loc.seek(start)
                    if loc.read(want_n) != data:
                        raise RuntimeError(f"bytes {start}-{end} of {pick} differ from the local file")
            rb_ok = True
            rb["ranges"] = [[s, e] for s, e in spans]
            rb["bytes"] = sum(e - s + 1 for s, e in spans)
            rb["file_bytes"] = size
        except Exception as e:
            rb["error"] = f"{type(e).__name__}: {str(e)[:200]}"
    # ok must include the read-back: a repo whose files are listed but not downloadable is not a
    # checkpoint, and pruning into that state loses the run. `extra` is gated too: a file on the Hub that
    # is not in this directory means the path being pushed to already held something else, which is a
    # naming collision, not a clean checkpoint (E-031 -- the docstring here used to claim otherwise).
    ok = bool(want) and not missing and not mismatch and not extra and (rb_ok or not readback_one)
    return {"ok": ok, "n_local": len(want), "n_hub": len(have), "missing": missing[:10],
            "mismatch": mismatch[:10], "extra": extra[:10], "readback": rb,
            "readback_ok": rb_ok}


def ensure_repo(repo, api, token=None, repo_type="dataset"):
    """Create the checkpoint repo and give it LFS patterns, in that order, before anything is pushed.

    A 1.6 GB checkpoint is far over the 10 MB non-LFS limit, and a fresh repo has no .gitattributes, so a
    first push fails with "should be tracked by LFS" -- ten minutes into the first checkpoint of a GPU
    session that is being billed while it waits. Found by preflight P1 on CPU, which is the whole point
    of running the cycle with mock bytes before the real run.
    """
    created = False
    try:
        api.create_repo(repo_id=repo, repo_type=repo_type, exist_ok=True, token=token)
        created = True
    except Exception as e:
        # exist_ok=True should absorb the common case; anything else is worth seeing, not swallowing.
        print(f"ensure_repo create: {type(e).__name__}: {str(e)[:160]}", flush=True)
    have = set()
    try:
        have = set(api.list_repo_files(repo_id=repo, repo_type=repo_type))
    except Exception as e:
        print(f"ensure_repo listing failed ({type(e).__name__}); writing attributes anyway", flush=True)
    if ".gitattributes" not in have:
        body = ("*.safetensors filter=lfs diff=lfs merge=lfs -text\n"
                "*.bin filter=lfs diff=lfs merge=lfs -text\n"
                "*.pt filter=lfs diff=lfs merge=lfs -text\n"
                "*.npy filter=lfs diff=lfs merge=lfs -text\n"
                "*.ckpt filter=lfs diff=lfs merge=lfs -text\n")
        api.upload_file(path_or_fileobj=body.encode(), path_in_repo=".gitattributes", repo_id=repo,
                        repo_type=repo_type, commit_message="track checkpoint tensors with LFS",
                        token=token)
    return {"repo": repo, "created": created, "had_gitattributes": ".gitattributes" in have,
            "files_before": len(have)}


def prune_local(dirpath, res, min_free_gb=6.0):
    """Delete a local checkpoint whose Hub copy has been verified, and record the space that freed.

    Separate from `push_and_prune` because the caller has to roll `latest.json` and read it back *between*
    the verification and the deletion: pruning first would destroy the only copy the pointer should be
    naming, and a pointer commit that failed silently would then be unrecoverable.
    """
    shutil.rmtree(dirpath, ignore_errors=True)
    res["pruned"] = True
    res["local_dir_exists_after"] = os.path.exists(dirpath)
    res["free_after_gb"] = round(free_bytes() / 1e9, 2)
    res["disk_ok"] = res["free_after_gb"] >= min_free_gb
    if not res["disk_ok"]:
        # Free space is a training-blocking metric (§3.13), not a footnote: say so loudly here rather
        # than letting the next checkpoint fail with an OSError.
        res["disk_warning"] = (f"only {res['free_after_gb']} GB free after pruning "
                               f"(floor {min_free_gb} GB)")
        print(f"   DISK WARNING: {res['disk_warning']}", flush=True)
    return res


def prune_verified(dirpath, res, min_free_gb=6.0):
    """`prune_local`, but only for a caller whose verification actually said yes.

    `push_and_prune(prune=True)` gates the deletion internally, and the mid-run save path deliberately does
    not use it: the pointer has to be rolled *between* the verification and the deletion (E-035/2), so the
    trainer calls `push_and_prune(..., prune=False)`, rolls the pointer, and then prunes itself -- which
    walked straight past the only guard standing between a failed verification and the loss of the last
    local copy. Review finding E-042. The check belongs next to the deleting, not only next to the upload.
    """
    ver = (res or {}).get("verify") or {}
    if ver.get("ok") is not True:
        raise RuntimeError(
            f"refusing to delete {dirpath}: the Hub copy is not verified "
            f"(missing={len(ver.get('missing') or [])} mismatched={len(ver.get('mismatch') or [])} "
            f"extra={len(ver.get('extra') or [])} readback_ok={ver.get('readback_ok')})")
    return prune_local(dirpath, res, min_free_gb)


def push_and_prune(repo, dirpath, path_in_repo, api, repo_type="dataset", token=None,
                   prune=True, min_free_gb=6.0, verify_tries=3):
    """One full §3.13 cycle. Raises if verification fails and pruning was requested -- the local copy is
    kept in exactly that case, because the Hub copy is not yet proven to exist."""
    files = sorted(os.listdir(dirpath))
    if not files:
        raise RuntimeError(f"refusing to upload empty checkpoint dir {dirpath}")
    free_before = free_bytes()
    if api is None:
        raise RuntimeError("push_and_prune needs an HfApi; refusing to skip the upload silently")
    up = api.upload_folder(repo_id=repo, folder_path=dirpath, path_in_repo=path_in_repo,
                           repo_type=repo_type, commit_message=f"checkpoint {path_in_repo}",
                           token=token)
    # Retry the verification, not the upload. `verify_checkpoint` reads the tree immediately after the
    # commit, and E-033 showed this Hub path answers transiently under load; treating a propagation race
    # as a content problem would kill a billed session over a checkpoint that is actually fine.
    ver = None
    for attempt in range(1, verify_tries + 1):
        ver = verify_checkpoint(repo, dirpath, path_in_repo, token=token, repo_type=repo_type)
        if ver["ok"]:
            break
        if attempt < verify_tries:
            print(f"   verify attempt {attempt}: {len(ver['missing'])} missing, "
                  f"{len(ver['mismatch'])} mismatched -- re-checking in {20 * attempt}s", flush=True)
            __import__("time").sleep(20 * attempt)
    res = {"uploaded_to": getattr(up, "commit_url", str(up)[:120]), "verify": ver,
           "free_before_gb": round(free_before / 1e9, 2), "pruned": False,
           "free_after_gb": round(free_bytes() / 1e9, 2)}
    if prune:
        if not ver["ok"]:
            res["prune_blocked_reason"] = (
                f"hub verification failed ({len(ver['missing'])} missing, "
                f"{len(ver['mismatch'])} mismatched) -- keeping the local copy in {dirpath}")
            raise RuntimeError(res["prune_blocked_reason"])
        prune_verified(dirpath, res, min_free_gb)
    return res


def latest_pointer(repo, repo_type="dataset", revision="main", token=None):
    """Which checkpoint the Hub believes is current. Read from the Hub, never from memory: an instance
    that died mid-cycle cannot be trusted to know what it finished uploading."""
    url = f"{HUB}/{repo_type}s/{repo}/resolve/{revision}/latest.json"
    req = urllib.request.Request(url)
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            ptr = json.loads(r.read().decode())
    except urllib.error.HTTPError as e:
        # 404 is a fact about the world ("this run has never checkpointed"), which legitimately means
        # start from step 0. Everything else -- 401, timeout, DNS -- is the lookup failing, and treating
        # that as "nothing yet" would silently retrain from the beginning on a mid-run cold restart.
        if e.code == 404:
            return {"not_found": True}
        return {"error": f"HTTPError {e.code}: {str(e)[:140]}"}
    except Exception as e:
        return {"error": f"{type(e).__name__}: {str(e)[:140]}"}
    if not isinstance(ptr, dict) or "path_in_repo" not in ptr or "step" not in ptr:
        return {"error": f"latest.json is not a usable pointer: {str(ptr)[:140]}"}
    return ptr


def download_checkpoint(repo, path_in_repo, dest, api, repo_type="dataset", token=None):
    """Pull one checkpoint down from the Hub into `dest` and prove the bytes are what the Hub lists.

    This is the cold-resume path: empty disk, nothing left behind, and the listing is checked against
    what lands. A partial download that looks like a checkpoint directory is worse than none.
    """
    import time as _t
    prefix = path_in_repo.rstrip("/") + "/"
    listing = hub_listing(repo, repo_type, token=token)
    if not any(k.startswith(prefix) for k in listing):
        raise RuntimeError(f"{repo} has no files under {prefix} -- refusing to resume from nothing")
    # Keys come back as full repo paths; strip the prefix once here so `rel` and `filename` agree.
    # (Passing prefix+full_path downstream asked the Hub for ckpt/x/ckpt/x/cursor.json -- a 404 that
    # looked like a missing checkpoint, found by preflight P1 on CPU rather than mid-main-run.)
    have = {k[len(prefix):]: v for k, v in listing.items() if k.startswith(prefix)}
    os.makedirs(dest, exist_ok=True)
    t0 = _t.time()
    tot = 0
    landed = []
    for rel in sorted(have):
        p = api.hf_hub_download(repo_id=repo, filename=prefix + rel, repo_type=repo_type,
                                local_dir=dest, token=token)
        sz = os.path.getsize(p)
        if have[rel][0] is not None and sz != have[rel][0]:
            raise RuntimeError(f"{rel}: hub lists {have[rel][0]} bytes, {sz} arrived")
        sha = have[rel][1]
        if sha and len(sha) == 64 and sha256_file(p) != sha:
            raise RuntimeError(f"{rel}: sha256 does not match the hub's lfs oid")
        tot += sz
        landed.append(p)
    # local_dir mirrors the repo layout, so a checkpoint at ckpt/checkpoint-7/ lands *inside* dest with
    # its subdirectories. The caller needs the directory that actually holds the weights, not the one it
    # asked for, or the resume silently points at an empty prefix and starts from step 0.
    dirs = {os.path.dirname(p) for p in landed}
    if len(dirs) != 1:
        raise RuntimeError(f"checkpoint files landed in {len(dirs)} directories: {sorted(dirs)[:4]}")
    return {"files": len(have), "bytes": tot, "seconds": round(_t.time() - t0, 1),
            "dir": dirs.pop(), "requested_dest": dest, "path_in_repo": path_in_repo}


def write_latest_pointer(repo, cur, api, repo_type="dataset", token=None, retries=3):
    """Roll `latest.json` and prove the Hub now serves it.

    The read-back is not ceremony: the caller prunes its local copy immediately afterwards, so a pointer
    commit that silently failed would leave a Hub-verified checkpoint that nothing points at, and the next
    session would resume from the previous one and re-read ~100 M tokens with every self-consistency guard
    passing (weights, cursor, fingerprint and seed all agree with each other -- only the pointer is wrong).
    """
    p = os.path.join("/tmp", "latest.json")
    with open(p, "w") as f:
        json.dump(cur, f, indent=1, sort_keys=True)
    want_step = cur.get("step")
    last = None
    for attempt in range(1, retries + 1):
        try:
            api.upload_file(path_or_fileobj=p, path_in_repo="latest.json", repo_id=repo,
                            repo_type=repo_type, commit_message=f"latest -> {want_step}", token=token)
        except Exception as e:
            last = f"upload {type(e).__name__}: {str(e)[:120]}"
        else:
            back = latest_pointer(repo, repo_type=repo_type, token=token)
            if back.get("step") == want_step and back.get("path_in_repo") == cur.get("path_in_repo"):
                return p
            last = f"hub reports {back.get('step')!r} at {back.get('path_in_repo')!r}, " \
                   f"we wrote {want_step!r} at {cur.get('path_in_repo')!r} ({str(back)[:120]})"
        print(f"   latest.json attempt {attempt} failed: {last}; retrying", flush=True)
        __import__("time").sleep(15 * attempt)
    raise RuntimeError(f"latest.json never landed at step {want_step} on {repo} after {retries} "
                       f"attempts: {last}. Refusing to prune the local copy that this pointer should "
                       "point at.")