File size: 4,673 Bytes
2ea2835 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Deterministic identities for frozen files and JSON metadata."""
from __future__ import annotations
import hashlib
import json
import os
from typing import Callable, Optional
import torch.distributed as dist
def canonical_sha256(value) -> str:
"""Hash a JSON value with one canonical serialization."""
payload = json.dumps(
value,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
).encode("utf-8")
return hashlib.sha256(payload).hexdigest()
def canonical_identity_value(value):
"""Normalize tuples and lists through strict, finite canonical JSON."""
try:
payload = json.dumps(
value,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
allow_nan=False,
)
except (TypeError, ValueError) as exc:
raise ValueError(
"artifact identity must contain only finite JSON-serializable values"
) from exc
return json.loads(payload)
def file_sha256(path: str) -> str:
"""Hash a file without loading it wholly into memory."""
digest = hashlib.sha256()
with open(path, "rb") as handle:
while chunk := handle.read(8 * 1024 * 1024):
digest.update(chunk)
return digest.hexdigest()
def distributed_file_sha256(path: str) -> str:
"""Hash a shared file on rank zero and broadcast the result."""
if (
not dist.is_available()
or not dist.is_initialized()
or dist.get_world_size() == 1
):
return file_sha256(path)
payload = [None]
if dist.get_rank() == 0:
try:
payload[0] = {"ok": True, "sha256": file_sha256(path)}
except BaseException as exc:
payload[0] = {
"ok": False,
"error": f"{type(exc).__name__}: {exc}",
}
dist.broadcast_object_list(payload, src=0)
result = payload[0]
if not isinstance(result, dict) or not result.get("ok"):
error = (
result.get("error", "unknown error")
if isinstance(result, dict)
else "no rank-zero result"
)
raise RuntimeError(f"rank-zero file hashing failed: {error}")
return str(result["sha256"])
def runtime_file_identity(
path,
*,
file_hasher: Optional[Callable[[str], str]] = None,
) -> dict:
"""Return path, stable stat fields, and SHA256 for one file."""
resolved = os.path.realpath(os.path.expanduser(os.fspath(path)))
if not os.path.isfile(resolved):
raise FileNotFoundError(
f"runtime identity target is not a file: {resolved}")
before = os.stat(resolved)
digest = str((file_hasher or distributed_file_sha256)(resolved))
after = os.stat(resolved)
if (
int(before.st_size) != int(after.st_size)
or int(before.st_mtime_ns) != int(after.st_mtime_ns)
):
raise RuntimeError(
f"runtime identity target changed while hashing: {resolved}")
if (
len(digest) != 64
or any(character not in "0123456789abcdef" for character in digest)
):
raise ValueError(
f"runtime identity returned an invalid SHA256 for {resolved}")
return {
"path": resolved,
"size_bytes": int(after.st_size),
"mtime_ns": int(after.st_mtime_ns),
"sha256": digest,
}
def distributed_runtime_file_identity(path) -> dict:
"""Build the complete file identity on rank zero and broadcast it."""
if (
not dist.is_available()
or not dist.is_initialized()
or dist.get_world_size() == 1
):
return runtime_file_identity(path, file_hasher=file_sha256)
payload = [None]
if dist.get_rank() == 0:
try:
payload[0] = {
"ok": True,
"identity": runtime_file_identity(
path, file_hasher=file_sha256),
}
except BaseException as exc:
payload[0] = {
"ok": False,
"error": f"{type(exc).__name__}: {exc}",
}
dist.broadcast_object_list(payload, src=0)
result = payload[0]
if not isinstance(result, dict) or not result.get("ok"):
error = (
result.get("error", "unknown error")
if isinstance(result, dict)
else "no rank-zero result"
)
raise RuntimeError(
f"rank-zero runtime file identity failed: {error}")
identity = result.get("identity")
if not isinstance(identity, dict):
raise RuntimeError(
"rank-zero runtime file identity returned invalid data")
return dict(identity)
|