| """Versioned JSON plus non-executable safetensors checkpoint contract.""" |
|
|
| import json |
| import math |
| import os |
| import tempfile |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from typing import Never, cast |
|
|
| import torch |
| from safetensors import SafetensorError, safe_open |
| from safetensors.torch import save_file |
|
|
| from .features import ACTION_FEATURES, FEATURE_DIMENSION |
| from .model import DISPOSITION_COUNT, MAX_HIDDEN_DIMENSION, CandidateScorer |
|
|
| CHECKPOINT_FORMAT = "yaml-edit/ml-checkpoint-v1" |
| MAX_CONFIG_BYTES = 64 * 1024 |
| MAX_WEIGHTS_BYTES = 10 * 1024 * 1024 |
| _SHA256_LENGTH = 64 |
| _MAX_SEED = 2**31 - 1 |
| _MAX_EXAMPLES = 10_000 |
| _MAX_EPOCH = 100_000 |
| _CONFIG_KEYS = frozenset( |
| { |
| "format", |
| "feature_dimension", |
| "action_features", |
| "hidden_dimension", |
| "disposition_threshold", |
| "action_threshold", |
| "seed", |
| "dataset_sha256", |
| "train_examples", |
| "validation_examples", |
| "selected_epoch", |
| } |
| ) |
|
|
|
|
| class CheckpointError(ValueError): |
| """A checkpoint violates the bounded safe-artifact contract.""" |
|
|
|
|
| def _fail(message: str) -> Never: |
| raise CheckpointError(message) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class ArtifactConfig: |
| """All architecture and calibrated inference settings needed for loading.""" |
|
|
| feature_dimension: int |
| hidden_dimension: int |
| disposition_threshold: float |
| action_threshold: float |
| seed: int |
| dataset_sha256: str |
| train_examples: int |
| validation_examples: int |
| selected_epoch: int |
|
|
| def __post_init__(self) -> None: |
| """Reject incoherent in-memory configurations at their construction boundary.""" |
| self.validate() |
|
|
| def validate(self) -> None: |
| """Validate exact types, fixed architecture, bounds, and finite thresholds.""" |
| _bounded_int( |
| self.feature_dimension, "feature_dimension", FEATURE_DIMENSION, FEATURE_DIMENSION |
| ) |
| _bounded_int(self.hidden_dimension, "hidden_dimension", 1, MAX_HIDDEN_DIMENSION) |
| _threshold(self.disposition_threshold, "disposition_threshold") |
| _threshold(self.action_threshold, "action_threshold") |
| _bounded_int(self.seed, "seed", 0, _MAX_SEED) |
| _bounded_int(self.train_examples, "train_examples", 1, _MAX_EXAMPLES) |
| _bounded_int(self.validation_examples, "validation_examples", 1, _MAX_EXAMPLES) |
| _bounded_int(self.selected_epoch, "selected_epoch", 0, _MAX_EPOCH) |
| _digest(self.dataset_sha256) |
|
|
| def to_dict(self) -> dict[str, object]: |
| """Return the strict versioned JSON record.""" |
| self.validate() |
| result: dict[str, object] = dict(asdict(self)) |
| result.update({"format": CHECKPOINT_FORMAT, "action_features": ACTION_FEATURES}) |
| return result |
|
|
| @classmethod |
| def from_dict(cls, value: object) -> "ArtifactConfig": |
| """Validate untrusted decoded configuration before model allocation.""" |
| if not isinstance(value, dict) or set(value) != _CONFIG_KEYS: |
| _fail("Checkpoint config has missing or unknown fields.") |
| data = cast("dict[str, object]", value) |
| if data["format"] != CHECKPOINT_FORMAT: |
| _fail("Checkpoint format is unsupported.") |
| _bounded_int(data["action_features"], "action width", ACTION_FEATURES, ACTION_FEATURES) |
| return cls( |
| _bounded_int( |
| data["feature_dimension"], |
| "feature_dimension", |
| FEATURE_DIMENSION, |
| FEATURE_DIMENSION, |
| ), |
| _bounded_int(data["hidden_dimension"], "hidden_dimension", 1, MAX_HIDDEN_DIMENSION), |
| _threshold(data["disposition_threshold"], "disposition_threshold"), |
| _threshold(data["action_threshold"], "action_threshold"), |
| _bounded_int(data["seed"], "seed", 0, _MAX_SEED), |
| _digest(data["dataset_sha256"]), |
| _bounded_int(data["train_examples"], "train_examples", 1, _MAX_EXAMPLES), |
| _bounded_int(data["validation_examples"], "validation_examples", 1, _MAX_EXAMPLES), |
| _bounded_int(data["selected_epoch"], "selected_epoch", 0, _MAX_EPOCH), |
| ) |
|
|
|
|
| def _bounded_int(value: object, label: str, minimum: int, maximum: int) -> int: |
| if type(value) is not int or not minimum <= value <= maximum: |
| _fail(f"{label} is outside its integer bounds.") |
| return value |
|
|
|
|
| def _threshold(value: object, label: str) -> float: |
| if type(value) not in (int, float): |
| _fail(f"{label} must be finite and within [0, 1].") |
| numeric = cast("int | float", value) |
| if not math.isfinite(numeric) or not 0.0 <= numeric <= 1.0: |
| _fail(f"{label} must be finite and within [0, 1].") |
| return float(numeric) |
|
|
|
|
| def _digest(value: object) -> str: |
| if not isinstance(value, str) or len(value) != _SHA256_LENGTH: |
| _fail("dataset_sha256 must be a 64-character digest.") |
| if any(character not in "0123456789abcdef" for character in value): |
| _fail("dataset_sha256 must be lowercase hexadecimal.") |
| return value |
|
|
|
|
| def _expected_shapes(config: ArtifactConfig) -> dict[str, tuple[int, ...]]: |
| hidden = config.hidden_dimension |
| return { |
| "encoder.weight": (hidden, config.feature_dimension), |
| "encoder.bias": (hidden,), |
| "disposition.weight": (DISPOSITION_COUNT, hidden), |
| "disposition.bias": (DISPOSITION_COUNT,), |
| "action.weight": (hidden, ACTION_FEATURES), |
| "action.bias": (hidden,), |
| } |
|
|
|
|
| def _read_capped(path: Path, limit: int, label: str) -> bytes: |
| try: |
| with path.open("rb") as file: |
| raw = file.read(limit + 1) |
| except OSError as error: |
| raise CheckpointError(f"Checkpoint {label} cannot be read.") from error |
| if len(raw) > limit: |
| _fail(f"Checkpoint {label} size exceeds the hard limit.") |
| return raw |
|
|
|
|
| def _pairs(items: list[tuple[str, object]]) -> dict[str, object]: |
| result: dict[str, object] = {} |
| for key, value in items: |
| if key in result: |
| _fail(f"Duplicate JSON key: {key!r}.") |
| result[key] = value |
| return result |
|
|
|
|
| def _load_json(path: Path) -> ArtifactConfig: |
| raw = _read_capped(path, MAX_CONFIG_BYTES, "config") |
| try: |
| value: object = json.loads( |
| raw.decode("utf-8"), object_pairs_hook=_pairs, parse_constant=_nonfinite |
| ) |
| except (UnicodeError, json.JSONDecodeError, RecursionError) as error: |
| raise CheckpointError("Checkpoint config is not valid bounded JSON.") from error |
| return ArtifactConfig.from_dict(value) |
|
|
|
|
| def _nonfinite(value: str) -> Never: |
| _fail(f"Nonfinite JSON number is forbidden: {value}.") |
|
|
|
|
| def _validate_model(config: ArtifactConfig, model: CandidateScorer) -> None: |
| config.validate() |
| expected = _expected_shapes(config) |
| state = model.state_dict() |
| if set(state) != set(expected): |
| _fail("Checkpoint model keys do not match the model contract.") |
| for key, shape in expected.items(): |
| tensor = state[key] |
| if tuple(tensor.shape) != shape or tensor.dtype != torch.float32: |
| _fail("Checkpoint model is incompatible with its configuration.") |
| if not bool(torch.isfinite(tensor).all()): |
| _fail("Checkpoint tensors must contain only finite values.") |
|
|
|
|
| def _load_snapshot(raw: bytes, config: ArtifactConfig) -> CandidateScorer: |
| expected = _expected_shapes(config) |
| try: |
| with tempfile.TemporaryDirectory(prefix="yaml-edit-weights-") as directory: |
| snapshot = Path(directory) / "weights.safetensors" |
| snapshot.write_bytes(raw) |
| snapshot.chmod(0o600) |
| with safe_open( |
| snapshot, framework="pt", device="cpu" |
| ) as tensors: |
| if set(tensors.keys()) != set(expected): |
| _fail("Checkpoint tensor keys do not match the model contract.") |
| for key, shape in expected.items(): |
| tensor_slice = tensors.get_slice(key) |
| if tensor_slice.get_shape() != list(shape): |
| _fail(f"Checkpoint tensor {key!r} has the wrong shape.") |
| if tensor_slice.get_dtype() != "F32": |
| _fail(f"Checkpoint tensor {key!r} has the wrong dtype.") |
| model = CandidateScorer(config.feature_dimension, config.hidden_dimension) |
| state = {key: tensors.get_tensor(key) for key in expected} |
| except CheckpointError: |
| raise |
| except OSError as error: |
| raise CheckpointError("Checkpoint weight snapshot cannot be created.") from error |
| except SafetensorError as error: |
| raise CheckpointError("Checkpoint safetensors header or data is invalid.") from error |
| if any(not bool(torch.isfinite(tensor).all()) for tensor in state.values()): |
| _fail("Checkpoint tensors must contain only finite values.") |
| model.load_state_dict(state, strict=True) |
| model.eval() |
| return model |
|
|
|
|
| def load_checkpoint( |
| config_path: Path, weights_path: Path |
| ) -> tuple[ArtifactConfig, CandidateScorer]: |
| """Load one bounded immutable snapshot after strict pre-allocation validation.""" |
| config = _load_json(config_path) |
| weights = _read_capped(weights_path, MAX_WEIGHTS_BYTES, "weights") |
| return config, _load_snapshot(weights, config) |
|
|
|
|
| def save_checkpoint( |
| config_path: Path, |
| weights_path: Path, |
| config: ArtifactConfig, |
| model: CandidateScorer, |
| ) -> None: |
| """Privately publish a validated checkpoint pair without overwriting either path.""" |
| _validate_model(config, model) |
| if config_path.exists() or weights_path.exists(): |
| raise FileExistsError("Checkpoint output paths must not already exist.") |
| config_bytes = ( |
| json.dumps(config.to_dict(), indent=2, sort_keys=True, allow_nan=False) + "\n" |
| ).encode("utf-8") |
| if len(config_bytes) > MAX_CONFIG_BYTES: |
| _fail("Checkpoint config size exceeds the hard limit.") |
| state = {key: value.detach().cpu().contiguous() for key, value in model.state_dict().items()} |
| with tempfile.TemporaryDirectory(dir=config_path.parent) as directory: |
| temporary = Path(directory) |
| temporary_weights = temporary / "weights.safetensors" |
| temporary_config = temporary / "config.json" |
| save_file(state, temporary_weights) |
| if temporary_weights.stat().st_size > MAX_WEIGHTS_BYTES: |
| _fail("Checkpoint weights size exceeds the hard limit.") |
| temporary_config.write_bytes(config_bytes) |
| temporary_weights.chmod(0o600) |
| temporary_config.chmod(0o600) |
| os.link(temporary_weights, weights_path) |
| try: |
| os.link(temporary_config, config_path) |
| except OSError: |
| weights_path.unlink(missing_ok=True) |
| raise |
|
|
