hubckpt: Content-Range + size guards on the bounded read-back, and prune_verified() (review E-042)

train/hubckpt.py

@@ -3,10 +3,13 @@
 
 # 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: the Hub's own file listing reports the expected path with the
-# expected size and the expected sha256 (the `lfs.oid` field, or the git sha for small files), and at
-# least one file is then actually downloaded back and re-hashed. A checkpoint that cannot be read back is
-# not a checkpoint, and the run must stop rather than delete the only local copy.
+# 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).
@@ -160,11 +163,19 @@ def verify_checkpoint(repo, dirpath, path_in_repo, token=None, repo_type="datase
         # 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] or 0
+        size = have[pick][0]
         url = f"{HUB}/{repo_type}s/{repo}/resolve/main/{prefix}{pick}"
         chunk = 4 << 20
-        spans = [(0, size - 1)] if size <= chunk else [(0, chunk - 1),
-                                                       (size // 2, min(size, size // 2 + chunk) - 1)]
+        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:
@@ -174,20 +185,27 @@ def verify_checkpoint(repo, dirpath, path_in_repo, token=None, repo_type="datase
                     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 = {"path": pick, "ranges": [[s, e] for s, e in spans],
-                  "bytes": sum(e - s + 1 for s, e in spans), "file_bytes": size}
+            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 = {"path": pick, "url": url,
-                  "error": f"{type(e).__name__}: {str(e)[:160]}"}
+            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
@@ -252,6 +270,24 @@ def prune_local(dirpath, res, min_free_gb=6.0):
     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
@@ -286,7 +322,7 @@ def push_and_prune(repo, dirpath, path_in_repo, api, repo_type="dataset", token=
                 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_local(dirpath, res, min_free_gb)
+        prune_verified(dirpath, res, min_free_gb)
     return res