601 lines
25 KiB
Python
601 lines
25 KiB
Python
"""Strict artifact persistence for three-phase trigger-binding training.
|
|
|
|
This module is intentionally independent from trainer and model internals. It stores
|
|
one artifact class per safetensors file and treats every persisted manifest as an
|
|
untrusted input when loading.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import tempfile
|
|
from collections.abc import Mapping, Sequence
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Tuple, Union
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
ARTIFACT_SCHEMA = "ai-toolkit.trigger-binding-artifact"
|
|
ARTIFACT_SCHEMA_VERSION = 1
|
|
CHECKPOINT_SCHEMA = "ai-toolkit.trigger-binding-checkpoint"
|
|
CHECKPOINT_SCHEMA_VERSION = 1
|
|
ARTIFACT_TYPES = frozenset({"embedding", "te_adapter", "tap_adapter", "diffusion_lora"})
|
|
ARTIFACT_MANIFEST_SCHEMAS = {
|
|
artifact_type: f"{ARTIFACT_SCHEMA}.{artifact_type}.v{ARTIFACT_SCHEMA_VERSION}"
|
|
for artifact_type in ARTIFACT_TYPES
|
|
}
|
|
|
|
_METADATA_SCHEMA = "trigger_binding.schema"
|
|
_METADATA_VERSION = "trigger_binding.schema_version"
|
|
_METADATA_TYPE = "trigger_binding.artifact_type"
|
|
_METADATA_MANIFEST = "trigger_binding.manifest"
|
|
_METADATA_MANIFEST_SHA256 = "trigger_binding.manifest_sha256"
|
|
_REQUIRED_METADATA_KEYS = frozenset(
|
|
{
|
|
_METADATA_SCHEMA,
|
|
_METADATA_VERSION,
|
|
_METADATA_TYPE,
|
|
_METADATA_MANIFEST,
|
|
_METADATA_MANIFEST_SHA256,
|
|
}
|
|
)
|
|
|
|
PathLike = Union[str, os.PathLike]
|
|
|
|
|
|
class ArtifactValidationError(ValueError):
|
|
"""Raised when an artifact or checkpoint fails closed validation."""
|
|
|
|
|
|
def canonical_json_dumps(value: Any) -> str:
|
|
"""Serialize JSON deterministically for hashing and metadata."""
|
|
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)
|
|
|
|
|
|
def sha256_bytes(data: bytes) -> str:
|
|
return hashlib.sha256(data).hexdigest()
|
|
|
|
|
|
def sha256_file(path: PathLike, chunk_size: int = 1024 * 1024) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as handle:
|
|
while True:
|
|
chunk = handle.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def fingerprint(value: Any) -> str:
|
|
return sha256_bytes(canonical_json_dumps(to_json_compatible(value)).encode("utf-8"))
|
|
|
|
|
|
def phase_fingerprint(phase: Any) -> str:
|
|
return fingerprint(phase)
|
|
|
|
|
|
def source_fingerprint(source: Any) -> str:
|
|
return fingerprint(source)
|
|
|
|
|
|
def config_fingerprint(config: Any) -> str:
|
|
return fingerprint(config)
|
|
|
|
|
|
def _qualified_name(value: Any) -> str:
|
|
cls = value if isinstance(value, type) else type(value)
|
|
return f"{cls.__module__}.{cls.__qualname__}"
|
|
|
|
|
|
def _object_config(value: Any) -> Dict[str, Any]:
|
|
fields = {
|
|
key: field_value
|
|
for key, field_value in vars(value).items()
|
|
if not key.startswith("_") and not callable(field_value)
|
|
}
|
|
return {"__type__": "object", "class": _qualified_name(value), "fields": to_json_compatible(fields)}
|
|
|
|
|
|
def to_json_compatible(value: Any) -> Any:
|
|
"""Encode checkpoint state, including Python/NumPy/Torch RNG state, as JSON values."""
|
|
if value is None or isinstance(value, (str, bool, int)):
|
|
return value
|
|
if isinstance(value, float):
|
|
if not math.isfinite(value):
|
|
raise TypeError("Non-finite floats are not JSON compatible")
|
|
return value
|
|
if isinstance(value, bytes):
|
|
return {"__type__": "bytes", "base64": base64.b64encode(value).decode("ascii")}
|
|
if isinstance(value, bytearray):
|
|
return {"__type__": "bytearray", "base64": base64.b64encode(bytes(value)).decode("ascii")}
|
|
if isinstance(value, tuple):
|
|
return {"__type__": "tuple", "items": [to_json_compatible(item) for item in value]}
|
|
if isinstance(value, list):
|
|
return [to_json_compatible(item) for item in value]
|
|
if isinstance(value, Mapping):
|
|
if not all(isinstance(key, str) for key in value):
|
|
raise TypeError("JSON-compatible mappings require string keys")
|
|
return {key: to_json_compatible(item) for key, item in value.items()}
|
|
if isinstance(value, torch.dtype):
|
|
return {"__type__": "torch.dtype", "value": str(value).removeprefix("torch.")}
|
|
if isinstance(value, torch.Tensor):
|
|
tensor = value.detach().cpu().contiguous()
|
|
raw = _tensor_bytes(tensor)
|
|
return {
|
|
"__type__": "torch.Tensor",
|
|
"dtype": str(tensor.dtype).removeprefix("torch."),
|
|
"shape": list(tensor.shape),
|
|
"base64": base64.b64encode(raw).decode("ascii"),
|
|
}
|
|
module = type(value).__module__
|
|
if module.startswith("numpy"):
|
|
import numpy as np
|
|
|
|
if isinstance(value, np.ndarray):
|
|
array = np.ascontiguousarray(value)
|
|
return {
|
|
"__type__": "numpy.ndarray",
|
|
"dtype": array.dtype.str,
|
|
"shape": list(array.shape),
|
|
"base64": base64.b64encode(array.tobytes()).decode("ascii"),
|
|
}
|
|
if isinstance(value, np.generic):
|
|
return to_json_compatible(value.item())
|
|
if hasattr(value, "__dict__"):
|
|
return _object_config(value)
|
|
raise TypeError(f"Unsupported JSON-compatible value: {_qualified_name(value)}")
|
|
|
|
|
|
def from_json_compatible(value: Any) -> Any:
|
|
"""Decode values produced by :func:`to_json_compatible`."""
|
|
if isinstance(value, list):
|
|
return [from_json_compatible(item) for item in value]
|
|
if not isinstance(value, dict):
|
|
return value
|
|
tag = value.get("__type__")
|
|
if tag is None:
|
|
return {key: from_json_compatible(item) for key, item in value.items()}
|
|
if tag in {"bytes", "bytearray"}:
|
|
decoded = base64.b64decode(value["base64"], validate=True)
|
|
return decoded if tag == "bytes" else bytearray(decoded)
|
|
if tag == "tuple":
|
|
return tuple(from_json_compatible(item) for item in value["items"])
|
|
if tag == "torch.dtype":
|
|
dtype = getattr(torch, value["value"], None)
|
|
if not isinstance(dtype, torch.dtype):
|
|
raise ArtifactValidationError(f"Unknown torch dtype: {value['value']}")
|
|
return dtype
|
|
if tag == "torch.Tensor":
|
|
dtype = getattr(torch, value["dtype"], None)
|
|
if not isinstance(dtype, torch.dtype):
|
|
raise ArtifactValidationError(f"Unknown torch dtype: {value['dtype']}")
|
|
raw = base64.b64decode(value["base64"], validate=True)
|
|
tensor = torch.frombuffer(bytearray(raw), dtype=dtype).clone()
|
|
expected_numel = math.prod(value["shape"])
|
|
if tensor.numel() != expected_numel:
|
|
raise ArtifactValidationError("Encoded tensor byte count does not match its shape")
|
|
return tensor.reshape(value["shape"])
|
|
if tag == "numpy.ndarray":
|
|
import numpy as np
|
|
|
|
raw = base64.b64decode(value["base64"], validate=True)
|
|
array = np.frombuffer(raw, dtype=np.dtype(value["dtype"])).copy()
|
|
expected_size = math.prod(value["shape"])
|
|
if array.size != expected_size:
|
|
raise ArtifactValidationError("Encoded NumPy byte count does not match its shape")
|
|
return array.reshape(value["shape"])
|
|
if tag == "object":
|
|
return {"class": value["class"], "fields": from_json_compatible(value["fields"])}
|
|
raise ArtifactValidationError(f"Unknown JSON compatibility tag: {tag}")
|
|
|
|
|
|
def encode_rng_state(rng_state: Any) -> Any:
|
|
return to_json_compatible(rng_state)
|
|
|
|
|
|
def decode_rng_state(encoded: Any) -> Any:
|
|
return from_json_compatible(encoded)
|
|
|
|
|
|
def _tensor_bytes(tensor: torch.Tensor) -> bytes:
|
|
normalized = tensor.detach().cpu().contiguous()
|
|
if normalized.numel() == 0:
|
|
return b""
|
|
# PyTorch cannot reinterpret a zero-dimensional tensor as a dtype with a
|
|
# different element size. Flattening preserves the exact storage bytes and
|
|
# also handles BF16, scalar adapter scales, and non-contiguous inputs.
|
|
return normalized.reshape(-1).view(torch.uint8).numpy().tobytes()
|
|
|
|
|
|
def tensor_sha256(tensor: torch.Tensor) -> str:
|
|
return sha256_bytes(_tensor_bytes(tensor))
|
|
|
|
|
|
def _validate_tensor_mapping(tensors: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
|
if not isinstance(tensors, Mapping) or not tensors:
|
|
raise ArtifactValidationError("Artifact tensors must be a non-empty mapping")
|
|
normalized: Dict[str, torch.Tensor] = {}
|
|
for key, tensor in tensors.items():
|
|
if not isinstance(key, str) or not key:
|
|
raise ArtifactValidationError("Artifact tensor keys must be non-empty strings")
|
|
if key in normalized:
|
|
raise ArtifactValidationError(f"Duplicate tensor key: {key}")
|
|
if not isinstance(tensor, torch.Tensor):
|
|
raise ArtifactValidationError(f"Artifact value for {key!r} is not a tensor")
|
|
if tensor.layout != torch.strided:
|
|
raise ArtifactValidationError(f"Artifact tensor {key!r} must use strided layout")
|
|
normalized[key] = tensor.detach().cpu().contiguous()
|
|
return normalized
|
|
|
|
|
|
def build_artifact_manifest(
|
|
artifact_type: str,
|
|
tensors: Mapping[str, torch.Tensor],
|
|
*,
|
|
phase: Any,
|
|
source: Any,
|
|
config: Any,
|
|
extra: Optional[Mapping[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
if artifact_type not in ARTIFACT_TYPES:
|
|
raise ArtifactValidationError(f"Unsupported artifact type: {artifact_type!r}")
|
|
normalized = _validate_tensor_mapping(tensors)
|
|
specs = {
|
|
key: {
|
|
"shape": list(tensor.shape),
|
|
"dtype": str(tensor.dtype).removeprefix("torch."),
|
|
"sha256": tensor_sha256(tensor),
|
|
}
|
|
for key, tensor in sorted(normalized.items())
|
|
}
|
|
manifest = {
|
|
"schema": ARTIFACT_SCHEMA,
|
|
"schema_version": ARTIFACT_SCHEMA_VERSION,
|
|
"artifact_type": artifact_type,
|
|
"artifact_schema": ARTIFACT_MANIFEST_SCHEMAS[artifact_type],
|
|
"phase_fingerprint": phase_fingerprint(phase),
|
|
"source_fingerprint": source_fingerprint(source),
|
|
"config_fingerprint": config_fingerprint(config),
|
|
"tensors": specs,
|
|
}
|
|
if extra is not None:
|
|
manifest["extra"] = to_json_compatible(dict(extra))
|
|
return manifest
|
|
|
|
|
|
def _artifact_metadata(manifest: Mapping[str, Any]) -> Dict[str, str]:
|
|
manifest_json = canonical_json_dumps(manifest)
|
|
return {
|
|
_METADATA_SCHEMA: ARTIFACT_SCHEMA,
|
|
_METADATA_VERSION: str(ARTIFACT_SCHEMA_VERSION),
|
|
_METADATA_TYPE: str(manifest["artifact_type"]),
|
|
_METADATA_MANIFEST: manifest_json,
|
|
_METADATA_MANIFEST_SHA256: sha256_bytes(manifest_json.encode("utf-8")),
|
|
}
|
|
|
|
|
|
def _atomic_replace(temp_path: str, destination: Path) -> None:
|
|
try:
|
|
with open(temp_path, "r+b") as handle:
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
except OSError:
|
|
# Some Windows filesystems reject fsync on safetensors temp handles.
|
|
# os.replace remains atomic on the same volume.
|
|
pass
|
|
os.replace(temp_path, destination)
|
|
try:
|
|
directory_fd = os.open(str(destination.parent), os.O_RDONLY)
|
|
except (AttributeError, OSError):
|
|
return
|
|
try:
|
|
os.fsync(directory_fd)
|
|
finally:
|
|
os.close(directory_fd)
|
|
|
|
|
|
def save_artifact(
|
|
path: PathLike,
|
|
artifact_type: str,
|
|
tensors: Mapping[str, torch.Tensor],
|
|
*,
|
|
phase: Any,
|
|
source: Any,
|
|
config: Any,
|
|
extra: Optional[Mapping[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Atomically save a typed safetensors artifact and return its manifest."""
|
|
destination = Path(path)
|
|
if destination.suffix.lower() != ".safetensors":
|
|
raise ArtifactValidationError("Artifact path must end with .safetensors")
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
normalized = _validate_tensor_mapping(tensors)
|
|
manifest = build_artifact_manifest(
|
|
artifact_type,
|
|
normalized,
|
|
phase=phase,
|
|
source=source,
|
|
config=config,
|
|
extra=extra,
|
|
)
|
|
file_descriptor, temp_path = tempfile.mkstemp(
|
|
prefix=f".{destination.name}.", suffix=".tmp", dir=str(destination.parent)
|
|
)
|
|
os.close(file_descriptor)
|
|
try:
|
|
save_file(normalized, temp_path, metadata=_artifact_metadata(manifest))
|
|
_atomic_replace(temp_path, destination)
|
|
except BaseException:
|
|
try:
|
|
os.unlink(temp_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
return manifest
|
|
|
|
|
|
def _parse_artifact_manifest(metadata: Optional[Mapping[str, str]]) -> Dict[str, Any]:
|
|
if metadata is None or not _REQUIRED_METADATA_KEYS.issubset(metadata):
|
|
missing = sorted(_REQUIRED_METADATA_KEYS.difference(metadata or {}))
|
|
raise ArtifactValidationError(f"Missing required safetensors metadata: {missing}")
|
|
if metadata[_METADATA_SCHEMA] != ARTIFACT_SCHEMA:
|
|
raise ArtifactValidationError("Artifact metadata schema mismatch")
|
|
if metadata[_METADATA_VERSION] != str(ARTIFACT_SCHEMA_VERSION):
|
|
raise ArtifactValidationError("Unsupported artifact metadata schema version")
|
|
manifest_json = metadata[_METADATA_MANIFEST]
|
|
if sha256_bytes(manifest_json.encode("utf-8")) != metadata[_METADATA_MANIFEST_SHA256]:
|
|
raise ArtifactValidationError("Artifact manifest metadata hash mismatch")
|
|
try:
|
|
manifest = json.loads(manifest_json)
|
|
except (TypeError, json.JSONDecodeError) as exc:
|
|
raise ArtifactValidationError("Artifact manifest is not valid JSON") from exc
|
|
if canonical_json_dumps(manifest) != manifest_json:
|
|
raise ArtifactValidationError("Artifact manifest is not canonically encoded")
|
|
required = {
|
|
"schema",
|
|
"schema_version",
|
|
"artifact_type",
|
|
"artifact_schema",
|
|
"phase_fingerprint",
|
|
"source_fingerprint",
|
|
"config_fingerprint",
|
|
"tensors",
|
|
}
|
|
if set(manifest).difference(required | {"extra"}) or not required.issubset(manifest):
|
|
raise ArtifactValidationError("Artifact manifest has missing or unknown top-level keys")
|
|
if manifest["schema"] != ARTIFACT_SCHEMA or manifest["schema_version"] != ARTIFACT_SCHEMA_VERSION:
|
|
raise ArtifactValidationError("Artifact manifest schema mismatch")
|
|
if manifest["artifact_type"] not in ARTIFACT_TYPES:
|
|
raise ArtifactValidationError("Artifact manifest type is unsupported")
|
|
if manifest["artifact_schema"] != ARTIFACT_MANIFEST_SCHEMAS[manifest["artifact_type"]]:
|
|
raise ArtifactValidationError("Artifact-specific manifest schema mismatch")
|
|
if metadata[_METADATA_TYPE] != manifest["artifact_type"]:
|
|
raise ArtifactValidationError("Artifact type metadata disagrees with manifest")
|
|
for name in ("phase_fingerprint", "source_fingerprint", "config_fingerprint"):
|
|
value = manifest[name]
|
|
if not isinstance(value, str) or len(value) != 64:
|
|
raise ArtifactValidationError(f"Invalid {name}")
|
|
if not isinstance(manifest["tensors"], dict) or not manifest["tensors"]:
|
|
raise ArtifactValidationError("Artifact manifest tensor map is empty")
|
|
return manifest
|
|
|
|
|
|
def load_artifact(
|
|
path: PathLike,
|
|
*,
|
|
expected_type: Optional[str] = None,
|
|
expected_keys: Optional[Sequence[str]] = None,
|
|
expected_shapes: Optional[Mapping[str, Sequence[int]]] = None,
|
|
expected_phase_fingerprint: Optional[str] = None,
|
|
expected_source_fingerprint: Optional[str] = None,
|
|
expected_config_fingerprint: Optional[str] = None,
|
|
expected_file_sha256: Optional[str] = None,
|
|
device: Union[str, torch.device] = "cpu",
|
|
) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
|
|
"""Load an artifact only after exact key, shape, dtype and hash validation."""
|
|
artifact_path = Path(path)
|
|
if not artifact_path.is_file():
|
|
raise ArtifactValidationError(f"Artifact file does not exist: {artifact_path}")
|
|
if expected_file_sha256 is not None and sha256_file(artifact_path) != expected_file_sha256:
|
|
raise ArtifactValidationError("Artifact file SHA-256 mismatch")
|
|
try:
|
|
with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
|
|
manifest = _parse_artifact_manifest(handle.metadata())
|
|
file_keys = set(handle.keys())
|
|
manifest_keys = set(manifest["tensors"])
|
|
if file_keys != manifest_keys:
|
|
raise ArtifactValidationError(
|
|
f"Artifact tensor keys mismatch: file={sorted(file_keys)}, manifest={sorted(manifest_keys)}"
|
|
)
|
|
tensors = {key: handle.get_tensor(key) for key in sorted(file_keys)}
|
|
except ArtifactValidationError:
|
|
raise
|
|
except Exception as exc:
|
|
raise ArtifactValidationError(f"Unable to read safetensors artifact: {artifact_path}") from exc
|
|
|
|
if expected_type is not None and manifest["artifact_type"] != expected_type:
|
|
raise ArtifactValidationError("Artifact type does not match the expected type")
|
|
expected_fingerprints = {
|
|
"phase_fingerprint": expected_phase_fingerprint,
|
|
"source_fingerprint": expected_source_fingerprint,
|
|
"config_fingerprint": expected_config_fingerprint,
|
|
}
|
|
for name, expected in expected_fingerprints.items():
|
|
if expected is not None and manifest[name] != expected:
|
|
raise ArtifactValidationError(f"Artifact {name} mismatch")
|
|
if expected_keys is not None and set(expected_keys) != set(tensors):
|
|
raise ArtifactValidationError("Artifact keys do not match expected keys")
|
|
if expected_shapes is not None:
|
|
if set(expected_shapes) != set(tensors):
|
|
raise ArtifactValidationError("Expected shape keys do not match artifact keys")
|
|
for key, shape in expected_shapes.items():
|
|
if list(tensors[key].shape) != list(shape):
|
|
raise ArtifactValidationError(f"Artifact tensor {key!r} has unexpected shape")
|
|
|
|
for key, tensor in tensors.items():
|
|
spec = manifest["tensors"].get(key)
|
|
if not isinstance(spec, dict) or set(spec) != {"shape", "dtype", "sha256"}:
|
|
raise ArtifactValidationError(f"Invalid tensor manifest entry for {key!r}")
|
|
if list(tensor.shape) != spec["shape"]:
|
|
raise ArtifactValidationError(f"Artifact tensor {key!r} shape mismatch")
|
|
if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]:
|
|
raise ArtifactValidationError(f"Artifact tensor {key!r} dtype mismatch")
|
|
if tensor_sha256(tensor) != spec["sha256"]:
|
|
raise ArtifactValidationError(f"Artifact tensor {key!r} SHA-256 mismatch")
|
|
return {key: tensor.to(device) for key, tensor in tensors.items()}, manifest
|
|
|
|
|
|
def artifact_reference(path: PathLike, *, relative_to: Optional[PathLike] = None) -> Dict[str, Any]:
|
|
artifact_path = Path(path)
|
|
with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
|
|
manifest = _parse_artifact_manifest(handle.metadata())
|
|
reference_path = os.path.relpath(artifact_path, relative_to) if relative_to is not None else str(artifact_path)
|
|
return {
|
|
"path": reference_path.replace(os.sep, "/"),
|
|
"artifact_type": manifest["artifact_type"],
|
|
"sha256": sha256_file(artifact_path),
|
|
"manifest_sha256": sha256_bytes(canonical_json_dumps(manifest).encode("utf-8")),
|
|
}
|
|
|
|
|
|
def build_checkpoint_manifest(
|
|
*,
|
|
phase: Any,
|
|
source: Any,
|
|
config: Any,
|
|
step: int,
|
|
artifacts: Mapping[str, Mapping[str, Any]],
|
|
rng_state: Any = None,
|
|
extra: Optional[Mapping[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
if not isinstance(step, int) or isinstance(step, bool) or step < 0:
|
|
raise ArtifactValidationError("Checkpoint step must be a non-negative integer")
|
|
if not isinstance(artifacts, Mapping):
|
|
raise ArtifactValidationError("Checkpoint artifacts must be a mapping")
|
|
manifest = {
|
|
"schema": CHECKPOINT_SCHEMA,
|
|
"schema_version": CHECKPOINT_SCHEMA_VERSION,
|
|
"step": step,
|
|
"phase_fingerprint": phase_fingerprint(phase),
|
|
"source_fingerprint": source_fingerprint(source),
|
|
"config_fingerprint": config_fingerprint(config),
|
|
"artifacts": to_json_compatible(dict(artifacts)),
|
|
"rng_state": encode_rng_state(rng_state),
|
|
}
|
|
if extra is not None:
|
|
manifest["extra"] = to_json_compatible(dict(extra))
|
|
validate_checkpoint_manifest(manifest)
|
|
return manifest
|
|
|
|
|
|
def validate_checkpoint_manifest(manifest: Mapping[str, Any]) -> None:
|
|
required = {
|
|
"schema",
|
|
"schema_version",
|
|
"step",
|
|
"phase_fingerprint",
|
|
"source_fingerprint",
|
|
"config_fingerprint",
|
|
"artifacts",
|
|
"rng_state",
|
|
}
|
|
if not isinstance(manifest, Mapping):
|
|
raise ArtifactValidationError("Checkpoint manifest must be a mapping")
|
|
if set(manifest).difference(required | {"extra"}) or not required.issubset(manifest):
|
|
raise ArtifactValidationError("Checkpoint manifest has missing or unknown top-level keys")
|
|
if manifest["schema"] != CHECKPOINT_SCHEMA or manifest["schema_version"] != CHECKPOINT_SCHEMA_VERSION:
|
|
raise ArtifactValidationError("Checkpoint manifest schema mismatch")
|
|
if not isinstance(manifest["step"], int) or isinstance(manifest["step"], bool) or manifest["step"] < 0:
|
|
raise ArtifactValidationError("Checkpoint manifest step is invalid")
|
|
for name in ("phase_fingerprint", "source_fingerprint", "config_fingerprint"):
|
|
value = manifest[name]
|
|
if not isinstance(value, str) or len(value) != 64:
|
|
raise ArtifactValidationError(f"Checkpoint {name} is invalid")
|
|
if not isinstance(manifest["artifacts"], dict):
|
|
raise ArtifactValidationError("Checkpoint artifact references must be a mapping")
|
|
for name, reference in manifest["artifacts"].items():
|
|
if not isinstance(name, str) or not isinstance(reference, dict):
|
|
raise ArtifactValidationError("Checkpoint artifact reference is invalid")
|
|
if set(reference) != {"path", "artifact_type", "sha256", "manifest_sha256"}:
|
|
raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} has invalid keys")
|
|
if reference["artifact_type"] not in ARTIFACT_TYPES:
|
|
raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} has invalid type")
|
|
if not all(isinstance(reference[key], str) for key in reference):
|
|
raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} contains non-string values")
|
|
if len(reference["sha256"]) != 64 or len(reference["manifest_sha256"]) != 64:
|
|
raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} has invalid hash")
|
|
try:
|
|
decode_rng_state(manifest["rng_state"])
|
|
except Exception as exc:
|
|
raise ArtifactValidationError("Checkpoint RNG state encoding is invalid") from exc
|
|
|
|
|
|
def save_checkpoint_manifest(path: PathLike, manifest: Mapping[str, Any]) -> None:
|
|
validate_checkpoint_manifest(manifest)
|
|
destination = Path(path)
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = (canonical_json_dumps(manifest) + "\n").encode("utf-8")
|
|
file_descriptor, temp_path = tempfile.mkstemp(
|
|
prefix=f".{destination.name}.", suffix=".tmp", dir=str(destination.parent)
|
|
)
|
|
try:
|
|
with os.fdopen(file_descriptor, "wb") as handle:
|
|
handle.write(payload)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temp_path, destination)
|
|
except BaseException:
|
|
try:
|
|
os.unlink(temp_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def load_checkpoint_manifest(
|
|
path: PathLike,
|
|
*,
|
|
verify_artifacts: bool = True,
|
|
expected_phase_fingerprint: Optional[str] = None,
|
|
expected_source_fingerprint: Optional[str] = None,
|
|
expected_config_fingerprint: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
manifest_path = Path(path)
|
|
try:
|
|
with open(manifest_path, "r", encoding="utf-8") as handle:
|
|
manifest = json.load(handle)
|
|
except Exception as exc:
|
|
raise ArtifactValidationError(f"Unable to read checkpoint manifest: {manifest_path}") from exc
|
|
validate_checkpoint_manifest(manifest)
|
|
expected_fingerprints = {
|
|
"phase_fingerprint": expected_phase_fingerprint,
|
|
"source_fingerprint": expected_source_fingerprint,
|
|
"config_fingerprint": expected_config_fingerprint,
|
|
}
|
|
for name, expected in expected_fingerprints.items():
|
|
if expected is not None and manifest[name] != expected:
|
|
raise ArtifactValidationError(f"Checkpoint {name} mismatch")
|
|
if verify_artifacts:
|
|
for name, reference in manifest["artifacts"].items():
|
|
artifact_path = Path(reference["path"])
|
|
if not artifact_path.is_absolute():
|
|
artifact_path = manifest_path.parent / artifact_path
|
|
_, artifact_manifest = load_artifact(
|
|
artifact_path,
|
|
expected_type=reference["artifact_type"],
|
|
expected_file_sha256=reference["sha256"],
|
|
)
|
|
actual_manifest_hash = sha256_bytes(canonical_json_dumps(artifact_manifest).encode("utf-8"))
|
|
if actual_manifest_hash != reference["manifest_sha256"]:
|
|
raise ArtifactValidationError(f"Checkpoint artifact {name!r} manifest hash mismatch")
|
|
return deepcopy(manifest)
|