File size: 4,359 Bytes
3e363e2 | 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 | """Immutable model/dataset artifact metadata and promotion gates."""
from __future__ import annotations
import hashlib
import json
from dataclasses import asdict, dataclass
from enum import Enum, StrEnum
from pathlib import Path
class ArtifactType(StrEnum):
MODEL = "graph_reasoner_model"
TOKENIZER = "graph_reasoner_tokenizer"
CONFIG = "graph_reasoner_config"
CALIBRATION = "graph_reasoner_calibration"
EVALUATION = "graph_reasoner_evaluation"
DATASET = "graph_reasoner_dataset"
GENERATOR = "graph_reasoner_generator"
@dataclass(frozen=True)
class ArtifactMetadata:
artifact_type: ArtifactType
version: str
path: str
sha256: str
created_at: str
model_version: str | None = None
graph_version: str | None = None
parameter_count: int | None = None
evaluation_passed: bool = False
immutable: bool = True
architecture_version: str | None = None
dataset_version: str | None = None
generator_version: str | None = None
tokenizer_version: str | None = None
ontology_version: str | None = None
calibration_version: str | None = None
evaluation_version: str | None = None
config_hash: str | None = None
git_commit: str | None = None
pat_config_hash: str | None = None
class ArtifactRegistry:
def __init__(self, root: str | Path) -> None:
self.root = Path(root)
self.root.mkdir(parents=True, exist_ok=True)
self.index_path = self.root / "registry.json"
self.active_path = self.root / "active.json"
self._index: dict[str, ArtifactMetadata] = {}
if self.index_path.exists():
raw = json.loads(self.index_path.read_text(encoding="utf-8"))
self._index = {
key: ArtifactMetadata(
**{**value, "artifact_type": ArtifactType(value["artifact_type"])}
)
for key, value in raw.items()
}
def register(self, metadata: ArtifactMetadata) -> ArtifactMetadata:
key = f"{metadata.artifact_type.value}:{metadata.version}"
path = Path(metadata.path)
if not path.exists():
raise FileNotFoundError(path)
if sha256_file(path) != metadata.sha256:
raise ValueError(f"artifact digest does not match payload: {path}")
if key in self._index:
if self._index[key] != metadata:
raise ValueError(
f"immutable artifact metadata already registered: {key}"
)
return self._index[key]
self._index[key] = metadata
self._write()
return metadata
def activate(self, artifact_type: ArtifactType, version: str) -> ArtifactMetadata:
key = f"{artifact_type.value}:{version}"
if key not in self._index:
raise KeyError(key)
metadata = self._index[key]
if artifact_type == ArtifactType.MODEL and not metadata.evaluation_passed:
raise ValueError("model cannot be activated before evaluation gates pass")
# Registration hashed the payload; promotion must see the same bytes.
if sha256_file(metadata.path) != metadata.sha256:
raise ValueError(f"{key} changed after registration; refusing to activate")
active = {}
if self.active_path.exists():
active = json.loads(self.active_path.read_text(encoding="utf-8"))
active[artifact_type.value] = version
self.active_path.write_text(
json.dumps(active, indent=2, sort_keys=True), encoding="utf-8"
)
return metadata
def get(self, artifact_type: ArtifactType, version: str) -> ArtifactMetadata:
return self._index[f"{artifact_type.value}:{version}"]
def _write(self) -> None:
self.index_path.write_text(
json.dumps(
{key: asdict(value) for key, value in self._index.items()},
indent=2,
default=lambda value: value.value if isinstance(value, Enum) else value,
sort_keys=True,
),
encoding="utf-8",
)
def sha256_file(path: str | Path) -> str:
digest = hashlib.sha256()
with Path(path).open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
|