| """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) |
|
|
