diff --git "a/fastplms/registry.py" "b/fastplms/registry.py" --- "a/fastplms/registry.py" +++ "b/fastplms/registry.py" @@ -1,1487 +1,1543 @@ -"""Typed access to the FastPLMs model and provenance manifest. - -The registry is intentionally independent of Torch and Transformers. Tooling can -therefore inspect supported checkpoints, licenses, and reference sources without -initializing a model runtime or downloading any files. -""" - -from __future__ import annotations - -import re -import tomllib -from collections.abc import Iterator, Mapping -from dataclasses import dataclass -from functools import lru_cache -from importlib import resources -from pathlib import Path, PurePosixPath, PureWindowsPath -from types import MappingProxyType -from typing import Any, Literal, cast -from urllib.parse import urlparse - - -_HEX_RE = re.compile(r"^[0-9a-f]+$") -_IDENTIFIER_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$") -_HUB_LICENSE_NAME_RE = re.compile(r"[^a-z0-9.]+") -_WINDOWS_INVALID_PATH_CHARACTERS = frozenset('<>:"|?*') -_WINDOWS_RESERVED_PATH_NAMES = frozenset( - {"AUX", "CON", "NUL", "PRN"} - | {f"COM{index}" for index in range(1, 10)} - | {f"LPT{index}" for index in range(1, 10)} -) -_REPOSITORY_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*$") -_REFERENCE_CONTAINER_RE = re.compile(r"^reference-[a-z0-9]+(?:-[a-z0-9]+)*$") -_REFERENCE_ADAPTER_RE = re.compile( - r"^tests\.parity\.support\.reference_adapters\.[a-z_][a-z0-9_]*$" -) -_DOCUMENTATION_FRAGMENT_RE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") -_ALLOWED_ATTENTION = frozenset( - {"eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"} -) -_ALLOWED_DTYPES = frozenset({"float32", "bfloat16"}) -_ALLOWED_PRECISIONS = frozenset({"default", "auto", "fp32", "bf16", "fp8"}) -_ALLOWED_BF16_EXECUTIONS = frozenset({"static_parameters", "fp32_parameters_autocast"}) -HUB_LICENSE_IDENTIFIERS = frozenset({"mit", "apache-2.0", "cc-by-nc-sa-4.0", "other"}) -_ALLOWED_TOKENIZER_MODES = frozenset({"tokenizer", "sequence", "structure"}) -_ALLOWED_SIZE_CATEGORIES = frozenset({"small", "medium", "large", "xlarge", "structure"}) -RuntimeExtra = Literal["core", "structure"] -TestTier = Literal["check", "compliance", "structure", "feature", "artifact", "benchmark"] -VramTier = Literal["sequence", "large-sequence", "structure", "structure-6b"] -GenerationContract = Literal["not_applicable", "required", "official_unavailable"] -RuntimeAssetTrustKind = Literal["hash_pinned_pickle"] -Bf16Execution = Literal["static_parameters", "fp32_parameters_autocast"] -DtypeName = Literal["float32", "bfloat16"] -_ALLOWED_EXTRAS = frozenset({"core", "structure"}) -_ALLOWED_TEST_TIERS = frozenset( - {"check", "compliance", "structure", "feature", "artifact", "benchmark"} -) -_ALLOWED_VRAM_TIERS = frozenset({"sequence", "large-sequence", "structure", "structure-6b"}) -_ALLOWED_GENERATION_CONTRACTS = frozenset({"not_applicable", "required", "official_unavailable"}) -_ALLOWED_RUNTIME_ASSET_TRUST_KINDS = frozenset({"hash_pinned_pickle"}) -_ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS = frozenset({"requires_cached_verified_file"}) -_ALLOWED_AUTO_CLASSES = frozenset( - { - "AutoConfig", - "AutoModel", - "AutoModelForMaskedLM", - "AutoModelForProteinFolding", - "AutoModelForSequenceClassification", - "AutoModelForSeq2SeqLM", - "AutoModelForTokenClassification", - } -) -_WEIGHT_SUFFIXES = (".bin", ".ckpt", ".pt", ".pth", ".safetensors") -_ALLOWED_ORACLE_ASSET_ROLES = frozenset({"weights", "contact_regression"}) -_FAIR_ESM_ASSET_HOST = "dl.fbaipublicfiles.com" -_ROOT_FIELDS = frozenset( - { - "schema_version", - "legal_files", - "attention_kernels", - "upstreams", - "families", - "models", - "runtime_assets", - } -) -_UPSTREAM_FIELDS = frozenset( - { - "id", - "path", - "url", - "revision", - "license", - "license_files", - "license_digests", - "distribution_files", - } -) -_FAMILY_FIELDS = frozenset( - { - "architecture", - "upstreams", - "tokenizer_mode", - "public_input", - "extra", - "reference_container", - "reference_adapter", - "attention", - "dtypes", - "bf16_execution", - "precisions", - "experimental_precisions", - "vram_tier", - "checkpoint_license", - "hub_license", - "hub_license_name", - "hub_license_link", - "state_transform", - "conversion_provenance", - "representative", - "documentation", - "test_tiers", - "runtime_paths", - "requires_complete_weight_publication", - "weights_publication_allowed", - "auto_map", - "tokenizer_class", - "backbone_model", - } -) -_MODEL_FIELDS = frozenset( - { - "id", - "family", - "size_category", - "generation_contract", - "fast_repo", - "fast_revision", - "fast_files", - "fast_unresolved_files", - "official_repo", - "official_revision", - "official_files", - "official_unresolved_files", - "oracle_assets", - "official_golden", - "artifact_source", - "canonical_state_sha256", - "tokenizer_source", - "auto_map", - "notes", - "msa_conditioning", - } -) -_RUNTIME_ASSET_FIELDS = frozenset( - { - "id", - "repository", - "revision", - "path", - "sha256", - "size", - "consumer_family", - "trust_kind", - "license", - "offline_behavior", - } -) - - -class RegistryError(ValueError): - """Raised when the model manifest is incomplete or internally inconsistent.""" - - -def _portable_relative_path(value: str, context: str) -> PurePosixPath: - """Return one normalized cross-platform relative path or fail closed.""" - - posix = PurePosixPath(value) - windows = PureWindowsPath(value) - unsafe_windows_part = any( - part.rstrip(" .") != part - or part.split(".", maxsplit=1)[0].upper() in _WINDOWS_RESERVED_PATH_NAMES - or any( - ord(character) < 32 or character in _WINDOWS_INVALID_PATH_CHARACTERS - for character in part - ) - for part in posix.parts - ) - if ( - not value - or not posix.parts - or posix == PurePosixPath(".") - or posix.is_absolute() - or windows.is_absolute() - or windows.drive - or "\\" in value - or "." in posix.parts - or ".." in posix.parts - or value != posix.as_posix() - or any( - part.lower() in {".git", ".cache", "__pycache__"} - for part in posix.parts - ) - or unsafe_windows_part - ): - raise RegistryError(f"{context} is not portable: {value!r}") - return posix - - -@dataclass(frozen=True, slots=True) -class FileDigest: - """Expected content identity for one pinned file.""" - - path: str - algorithm: str - digest: str - - @classmethod - def parse(cls, value: str) -> FileDigest: - try: - path, encoded_digest = value.split("=", maxsplit=1) - algorithm, digest = encoded_digest.split(":", maxsplit=1) - except ValueError as error: - raise RegistryError("File digests must use '=:'.") from error - - _portable_relative_path(path, "Checkpoint file path") - - expected_length = {"git-sha1": 40, "sha256": 64}.get(algorithm) - if expected_length is None: - raise RegistryError(f"Unsupported file digest algorithm: {algorithm!r}") - if len(digest) != expected_length or _HEX_RE.fullmatch(digest) is None: - raise RegistryError(f"Invalid {algorithm} digest for {path!r}: {digest!r}") - return cls(path=path, algorithm=algorithm, digest=digest) - - @property - def encoded(self) -> str: - return f"{self.algorithm}:{self.digest}" - - -@dataclass(frozen=True, slots=True) -class CheckpointSource: - """One immutable Hugging Face repository snapshot.""" - - repo_id: str - revision: str - files: tuple[FileDigest, ...] - unresolved_files: tuple[str, ...] = () - - @property - def file_map(self) -> Mapping[str, FileDigest]: - return MappingProxyType({item.path: item for item in self.files}) - - -@dataclass(frozen=True, slots=True) -class OracleAsset: - """Hash-pinned external file required by a native parity oracle.""" - - role: str - path: str - url: str - sha256: str - size: int - - -@dataclass(frozen=True, slots=True) -class RuntimeAsset: - """Immutable runtime data with an explicit deserialization trust boundary.""" - - id: str - repository: str - revision: str - path: str - sha256: str - size: int - consumer_family: str - trust_kind: RuntimeAssetTrustKind - license_expression: str - offline_behavior: str - - -@dataclass(frozen=True, slots=True) -class OfficialGolden: - """Hash-pinned official output bundle required by the check tier.""" - - metadata: FileDigest - tensors: FileDigest - - -@dataclass(frozen=True, slots=True) -class UpstreamSource: - """Pinned official implementation used as a parity oracle.""" - - id: str - path: str - url: str - revision: str - license_expression: str - license_files: tuple[str, ...] - license_digests: tuple[FileDigest, ...] = () - distribution_files: tuple[FileDigest, ...] = () - - -@dataclass(frozen=True, slots=True) -class AttentionKernelSpec: - """Immutable Hugging Face kernel used by one attention backend.""" - - implementation: str - repository: str - revision: str - version: int - expected_variant: str - dtypes: tuple[DtypeName, ...] - - -@dataclass(frozen=True, slots=True) -class ModelFamily: - """Shared runtime and compliance contract for one architecture family.""" - - id: str - architecture: str - upstreams: tuple[str, ...] - tokenizer_mode: str - public_input: str - extra: RuntimeExtra - reference_container: str - reference_adapter: str - attention: tuple[str, ...] - dtypes: tuple[DtypeName, ...] - bf16_execution: Bf16Execution - precisions: tuple[str, ...] - vram_tier: VramTier - checkpoint_license: str - hub_license: str - state_transform: str - representative: str - documentation: str - test_tiers: tuple[TestTier, ...] - runtime_paths: tuple[str, ...] - auto_map_items: tuple[tuple[str, str], ...] - requires_complete_weight_publication: bool = False - weights_publication_allowed: bool = False - experimental_precisions: tuple[str, ...] = () - tokenizer_class: str | None = None - hub_license_name: str | None = None - hub_license_link: str | None = None - conversion_provenance: str = "" - backbone_model: str | None = None - - @property - def auto_map(self) -> Mapping[str, str]: - return MappingProxyType(dict(self.auto_map_items)) - - @property - def hub_license_metadata(self) -> Mapping[str, str]: - """Return valid Hugging Face model-card license fields.""" - - metadata = {"license": self.hub_license} - if self.hub_license_name is not None: - # Hugging Face validates custom license names as lowercase slugs, - # while the manifest retains the reader-facing display name used - # in generated prose. - metadata["license_name"] = _HUB_LICENSE_NAME_RE.sub( - "-", - self.hub_license_name.lower(), - ).strip("-.") - if self.hub_license_link is not None: - metadata["license_link"] = self.hub_license_link - return MappingProxyType(metadata) - - @property - def stable_precisions(self) -> tuple[str, ...]: - """Return precision policies covered by the release contract.""" - - experimental = set(self.experimental_precisions) - return tuple(precision for precision in self.precisions if precision not in experimental) - - -@dataclass(frozen=True, slots=True) -class ModelSpec: - """Complete immutable source and runtime contract for one checkpoint.""" - - id: str - family: ModelFamily - fast: CheckpointSource - official: CheckpointSource - size_category: str - generation_contract: GenerationContract = "not_applicable" - oracle_assets: tuple[OracleAsset, ...] = () - official_golden: OfficialGolden | None = None - artifact_source: str = "fast" - canonical_state_sha256: str | None = None - tokenizer_source_id: str | None = None - auto_map_items: tuple[tuple[str, str], ...] = () - notes: str = "" - msa_conditioning: bool | None = None - - @property - def is_deep_reference(self) -> bool: - return self.id == self.family.representative - - @property - def auto_map(self) -> Mapping[str, str]: - if self.auto_map_items: - return MappingProxyType(dict(self.auto_map_items)) - return self.family.auto_map - - @property - def artifact_checkpoint(self) -> CheckpointSource: - """Return the checkpoint selected for local artifact construction.""" - - return self.fast if self.artifact_source == "fast" else self.official - - @property - def oracle_asset_map(self) -> Mapping[str, OracleAsset]: - """Return native oracle assets keyed by their declared role.""" - - return MappingProxyType({asset.role: asset for asset in self.oracle_assets}) - - -class ModelRegistry(Mapping[str, ModelSpec]): - """Validated mapping of model IDs to typed model specifications.""" - - def __init__( - self, - *, - schema_version: int, - upstreams: Mapping[str, UpstreamSource], - families: Mapping[str, ModelFamily], - models: Mapping[str, ModelSpec], - runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}), - attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}), - legal_files: tuple[FileDigest, ...] = (), - ) -> None: - self.schema_version = schema_version - self.upstreams = MappingProxyType(dict(upstreams)) - self.attention_kernels = MappingProxyType(dict(attention_kernels)) - self.families = MappingProxyType(dict(families)) - self._models = MappingProxyType(dict(models)) - self.runtime_assets = MappingProxyType(dict(runtime_assets)) - self.legal_files = legal_files - - def __getitem__(self, key: str) -> ModelSpec: - return self._models[key] - - def __iter__(self) -> Iterator[str]: - return iter(self._models) - - def __len__(self) -> int: - return len(self._models) - - def by_family(self, family_id: str) -> tuple[ModelSpec, ...]: - if family_id not in self.families: - raise KeyError(family_id) - return tuple(model for model in self._models.values() if model.family.id == family_id) - - def supported_attention_dtypes( - self, - family_id: str, - implementation: str, - ) -> tuple[DtypeName, ...]: - """Return manifest-supported dtypes for one family/backend pair.""" - - family = self.families[family_id] - if implementation not in family.attention: - raise KeyError( - f"Family {family_id!r} does not advertise attention backend " - f"{implementation!r}." - ) - kernel = self.attention_kernels.get(implementation) - if kernel is None: - return family.dtypes - return tuple(dtype for dtype in family.dtypes if dtype in kernel.dtypes) - +"""Typed access to the FastPLMs model and provenance manifest. + +The registry is intentionally independent of Torch and Transformers. Tooling can +therefore inspect supported checkpoints, licenses, and reference sources without +initializing a model runtime or downloading any files. +""" + +from __future__ import annotations + +import re +import tomllib +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from functools import lru_cache +from importlib import resources +from pathlib import Path, PurePosixPath, PureWindowsPath +from types import MappingProxyType +from typing import Any, Literal, cast +from urllib.parse import urlparse + + +_HEX_RE = re.compile(r"^[0-9a-f]+$") +_IDENTIFIER_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$") +_HUB_LICENSE_NAME_RE = re.compile(r"[^a-z0-9.]+") +_WINDOWS_INVALID_PATH_CHARACTERS = frozenset('<>:"|?*') +_WINDOWS_RESERVED_PATH_NAMES = frozenset( + {"AUX", "CON", "NUL", "PRN"} + | {f"COM{index}" for index in range(1, 10)} + | {f"LPT{index}" for index in range(1, 10)} +) +_REPOSITORY_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*$") +_REFERENCE_CONTAINER_RE = re.compile(r"^reference-[a-z0-9]+(?:-[a-z0-9]+)*$") +_REFERENCE_ADAPTER_RE = re.compile( + r"^tests\.parity\.support\.reference_adapters\.[a-z_][a-z0-9_]*$" +) +_DOCUMENTATION_FRAGMENT_RE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") +_ALLOWED_ATTENTION = frozenset( + {"eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"} +) +_ALLOWED_DTYPES = frozenset({"float32", "bfloat16"}) +_ALLOWED_PRECISIONS = frozenset({"default", "auto", "fp32", "bf16", "fp8"}) +_ALLOWED_BF16_EXECUTIONS = frozenset({"static_parameters", "fp32_parameters_autocast"}) +HUB_LICENSE_IDENTIFIERS = frozenset({"mit", "apache-2.0", "cc-by-nc-sa-4.0", "other"}) +_ALLOWED_TOKENIZER_MODES = frozenset({"tokenizer", "sequence", "structure"}) +_ALLOWED_SIZE_CATEGORIES = frozenset({"small", "medium", "large", "xlarge", "structure"}) +RuntimeExtra = Literal["core", "structure"] +TestTier = Literal["check", "compliance", "structure", "feature", "artifact", "benchmark"] +VramTier = Literal["sequence", "large-sequence", "structure", "structure-6b"] +GenerationContract = Literal["not_applicable", "required", "official_unavailable"] +RuntimeAssetTrustKind = Literal["hash_pinned_pickle"] +Bf16Execution = Literal["static_parameters", "fp32_parameters_autocast"] +DtypeName = Literal["float32", "bfloat16"] +_ALLOWED_EXTRAS = frozenset({"core", "structure"}) +_ALLOWED_TEST_TIERS = frozenset( + {"check", "compliance", "structure", "feature", "artifact", "benchmark"} +) +_ALLOWED_VRAM_TIERS = frozenset({"sequence", "large-sequence", "structure", "structure-6b"}) +_ALLOWED_GENERATION_CONTRACTS = frozenset({"not_applicable", "required", "official_unavailable"}) +_ALLOWED_RUNTIME_ASSET_TRUST_KINDS = frozenset({"hash_pinned_pickle"}) +_ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS = frozenset({"requires_cached_verified_file"}) +_ALLOWED_AUTO_CLASSES = frozenset( + { + "AutoConfig", + "AutoModel", + "AutoModelForMaskedLM", + "AutoModelForProteinFolding", + "AutoModelForSequenceClassification", + "AutoModelForSeq2SeqLM", + "AutoModelForTokenClassification", + } +) +_WEIGHT_SUFFIXES = (".bin", ".ckpt", ".pt", ".pth", ".safetensors") +_ALLOWED_ORACLE_ASSET_ROLES = frozenset({"weights", "contact_regression"}) +_FAIR_ESM_ASSET_HOST = "dl.fbaipublicfiles.com" +_ROOT_FIELDS = frozenset( + { + "schema_version", + "legal_files", + "attention_kernels", + "upstreams", + "families", + "models", + "runtime_assets", + } +) +_UPSTREAM_FIELDS = frozenset( + { + "id", + "path", + "url", + "revision", + "license", + "license_files", + "license_digests", + "distribution_files", + } +) +_FAMILY_FIELDS = frozenset( + { + "architecture", + "upstreams", + "tokenizer_mode", + "public_input", + "extra", + "reference_container", + "reference_adapter", + "attention", + "dtypes", + "bf16_execution", + "precisions", + "experimental_precisions", + "vram_tier", + "checkpoint_license", + "hub_license", + "hub_license_name", + "hub_license_link", + "state_transform", + "conversion_provenance", + "representative", + "documentation", + "test_tiers", + "runtime_paths", + "requires_complete_weight_publication", + "weights_publication_allowed", + "auto_map", + "tokenizer_class", + "backbone_model", + } +) +_MODEL_FIELDS = frozenset( + { + "id", + "family", + "size_category", + "generation_contract", + "fast_repo", + "fast_revision", + "fast_files", + "fast_unresolved_files", + "official_repo", + "official_revision", + "official_files", + "official_unresolved_files", + "oracle_assets", + "official_golden", + "artifact_source", + "canonical_state_sha256", + "tokenizer_source", + "auto_map", + "notes", + "msa_conditioning", + "backbone", + "backbone_model", + "publication_status", + } +) +_RUNTIME_ASSET_FIELDS = frozenset( + { + "id", + "repository", + "revision", + "path", + "sha256", + "size", + "consumer_family", + "trust_kind", + "license", + "offline_behavior", + } +) + + +class RegistryError(ValueError): + """Raised when the model manifest is incomplete or internally inconsistent.""" + + +def _portable_relative_path(value: str, context: str) -> PurePosixPath: + """Return one normalized cross-platform relative path or fail closed.""" + + posix = PurePosixPath(value) + windows = PureWindowsPath(value) + unsafe_windows_part = any( + part.rstrip(" .") != part + or part.split(".", maxsplit=1)[0].upper() in _WINDOWS_RESERVED_PATH_NAMES + or any( + ord(character) < 32 or character in _WINDOWS_INVALID_PATH_CHARACTERS + for character in part + ) + for part in posix.parts + ) + if ( + not value + or not posix.parts + or posix == PurePosixPath(".") + or posix.is_absolute() + or windows.is_absolute() + or windows.drive + or "\\" in value + or "." in posix.parts + or ".." in posix.parts + or value != posix.as_posix() + or any( + part.lower() in {".git", ".cache", "__pycache__"} + for part in posix.parts + ) + or unsafe_windows_part + ): + raise RegistryError(f"{context} is not portable: {value!r}") + return posix + + +@dataclass(frozen=True, slots=True) +class FileDigest: + """Expected content identity for one pinned file.""" + + path: str + algorithm: str + digest: str + + @classmethod + def parse(cls, value: str) -> FileDigest: + try: + path, encoded_digest = value.split("=", maxsplit=1) + algorithm, digest = encoded_digest.split(":", maxsplit=1) + except ValueError as error: + raise RegistryError("File digests must use '=:'.") from error + + _portable_relative_path(path, "Checkpoint file path") + + expected_length = {"git-sha1": 40, "sha256": 64}.get(algorithm) + if expected_length is None: + raise RegistryError(f"Unsupported file digest algorithm: {algorithm!r}") + if len(digest) != expected_length or _HEX_RE.fullmatch(digest) is None: + raise RegistryError(f"Invalid {algorithm} digest for {path!r}: {digest!r}") + return cls(path=path, algorithm=algorithm, digest=digest) + + @property + def encoded(self) -> str: + return f"{self.algorithm}:{self.digest}" + + +@dataclass(frozen=True, slots=True) +class CheckpointSource: + """One immutable Hugging Face repository snapshot.""" + + repo_id: str + revision: str + files: tuple[FileDigest, ...] + unresolved_files: tuple[str, ...] = () + + @property + def file_map(self) -> Mapping[str, FileDigest]: + return MappingProxyType({item.path: item for item in self.files}) + + +@dataclass(frozen=True, slots=True) +class OracleAsset: + """Hash-pinned external file required by a native parity oracle.""" + + role: str + path: str + url: str + sha256: str + size: int + + +@dataclass(frozen=True, slots=True) +class RuntimeAsset: + """Immutable runtime data with an explicit deserialization trust boundary.""" + + id: str + repository: str + revision: str + path: str + sha256: str + size: int + consumer_family: str + trust_kind: RuntimeAssetTrustKind + license_expression: str + offline_behavior: str + + +@dataclass(frozen=True, slots=True) +class OfficialGolden: + """Hash-pinned official output bundle required by the check tier.""" + + metadata: FileDigest + tensors: FileDigest + + +@dataclass(frozen=True, slots=True) +class UpstreamSource: + """Pinned official implementation used as a parity oracle.""" + + id: str + path: str + url: str + revision: str + license_expression: str + license_files: tuple[str, ...] + license_digests: tuple[FileDigest, ...] = () + distribution_files: tuple[FileDigest, ...] = () + + +@dataclass(frozen=True, slots=True) +class AttentionKernelSpec: + """Immutable Hugging Face kernel used by one attention backend.""" + + implementation: str + repository: str + revision: str + version: int + expected_variant: str + dtypes: tuple[DtypeName, ...] + + +@dataclass(frozen=True, slots=True) +class ModelFamily: + """Shared runtime and compliance contract for one architecture family.""" + + id: str + architecture: str + upstreams: tuple[str, ...] + tokenizer_mode: str + public_input: str + extra: RuntimeExtra + reference_container: str + reference_adapter: str + attention: tuple[str, ...] + dtypes: tuple[DtypeName, ...] + bf16_execution: Bf16Execution + precisions: tuple[str, ...] + vram_tier: VramTier + checkpoint_license: str + hub_license: str + state_transform: str + representative: str + documentation: str + test_tiers: tuple[TestTier, ...] + runtime_paths: tuple[str, ...] + auto_map_items: tuple[tuple[str, str], ...] + requires_complete_weight_publication: bool = False + weights_publication_allowed: bool = False + experimental_precisions: tuple[str, ...] = () + tokenizer_class: str | None = None + hub_license_name: str | None = None + hub_license_link: str | None = None + conversion_provenance: str = "" + backbone_model: str | None = None + + @property + def auto_map(self) -> Mapping[str, str]: + return MappingProxyType(dict(self.auto_map_items)) + + @property + def hub_license_metadata(self) -> Mapping[str, str]: + """Return valid Hugging Face model-card license fields.""" + + metadata = {"license": self.hub_license} + if self.hub_license_name is not None: + # Hugging Face validates custom license names as lowercase slugs, + # while the manifest retains the reader-facing display name used + # in generated prose. + metadata["license_name"] = _HUB_LICENSE_NAME_RE.sub( + "-", + self.hub_license_name.lower(), + ).strip("-.") + if self.hub_license_link is not None: + metadata["license_link"] = self.hub_license_link + return MappingProxyType(metadata) + + @property + def stable_precisions(self) -> tuple[str, ...]: + """Return precision policies covered by the release contract.""" + + experimental = set(self.experimental_precisions) + return tuple(precision for precision in self.precisions if precision not in experimental) + + +@dataclass(frozen=True, slots=True) +class ModelSpec: + """Complete immutable source and runtime contract for one checkpoint.""" + + id: str + family: ModelFamily + fast: CheckpointSource + official: CheckpointSource + size_category: str + generation_contract: GenerationContract = "not_applicable" + oracle_assets: tuple[OracleAsset, ...] = () + official_golden: OfficialGolden | None = None + artifact_source: str = "fast" + canonical_state_sha256: str | None = None + tokenizer_source_id: str | None = None + auto_map_items: tuple[tuple[str, str], ...] = () + notes: str = "" + msa_conditioning: bool | None = None + backbone: CheckpointSource | None = None + backbone_model: str | None = None + publication_status: str = "published" + + @property + def is_deep_reference(self) -> bool: + return self.id == self.family.representative + + @property + def auto_map(self) -> Mapping[str, str]: + if self.auto_map_items: + return MappingProxyType(dict(self.auto_map_items)) + return self.family.auto_map + + @property + def artifact_checkpoint(self) -> CheckpointSource: + """Return the checkpoint selected for local artifact construction.""" + + return self.fast if self.artifact_source == "fast" else self.official + + @property + def oracle_asset_map(self) -> Mapping[str, OracleAsset]: + """Return native oracle assets keyed by their declared role.""" + + return MappingProxyType({asset.role: asset for asset in self.oracle_assets}) + + +class ModelRegistry(Mapping[str, ModelSpec]): + """Validated mapping of model IDs to typed model specifications.""" + + def __init__( + self, + *, + schema_version: int, + upstreams: Mapping[str, UpstreamSource], + families: Mapping[str, ModelFamily], + models: Mapping[str, ModelSpec], + runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}), + attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}), + legal_files: tuple[FileDigest, ...] = (), + ) -> None: + self.schema_version = schema_version + self.upstreams = MappingProxyType(dict(upstreams)) + self.attention_kernels = MappingProxyType(dict(attention_kernels)) + self.families = MappingProxyType(dict(families)) + self._models = MappingProxyType(dict(models)) + self.runtime_assets = MappingProxyType(dict(runtime_assets)) + self.legal_files = legal_files + + def __getitem__(self, key: str) -> ModelSpec: + return self._models[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._models) + + def __len__(self) -> int: + return len(self._models) + + def by_family(self, family_id: str) -> tuple[ModelSpec, ...]: + if family_id not in self.families: + raise KeyError(family_id) + return tuple(model for model in self._models.values() if model.family.id == family_id) + + def supported_attention_dtypes( + self, + family_id: str, + implementation: str, + ) -> tuple[DtypeName, ...]: + """Return manifest-supported dtypes for one family/backend pair.""" + + family = self.families[family_id] + if implementation not in family.attention: + raise KeyError( + f"Family {family_id!r} does not advertise attention backend " + f"{implementation!r}." + ) + kernel = self.attention_kernels.get(implementation) + if kernel is None: + return family.dtypes + return tuple(dtype for dtype in family.dtypes if dtype in kernel.dtypes) + def require_resolved(self, model_id: str | None = None) -> None: - """Fail release validation when required file identities remain unresolved.""" - - selected = self._models.values() if model_id is None else (self._models[model_id],) + """Fail release validation when required file identities remain unresolved.""" + + selected = self._models.values() if model_id is None else (self._models[model_id],) unresolved: list[str] = [] for model in selected: + if model.publication_status == "pending": + unresolved.append(f"{model.id}.fast:unpublished") for label, checkpoint in (("fast", model.fast), ("official", model.official)): - for path in checkpoint.unresolved_files: - unresolved.append(f"{model.id}.{label}:{path}") - if unresolved: - detail = ", ".join(unresolved) - raise RegistryError(f"Release provenance is unresolved: {detail}") - - -def _reject_unknown_fields( - table: Mapping[str, Any], - allowed: frozenset[str], - context: str, -) -> None: - unknown = sorted(set(table).difference(allowed)) - if unknown: - raise RegistryError(f"{context} contains unknown fields: {unknown}.") - - -def _require_str(table: Mapping[str, Any], key: str, context: str) -> str: - value = table.get(key) - if not isinstance(value, str) or not value.strip(): - raise RegistryError(f"{context}.{key} must be a non-empty string.") - return value - - -def _require_enum( - table: Mapping[str, Any], - key: str, - context: str, - allowed: frozenset[str], -) -> str: - value = _require_str(table, key, context) - if value not in allowed: - raise RegistryError( - f"{context}.{key} must be one of {sorted(allowed)}; received {value!r}." - ) - return value - - -def _parse_reference_container(table: Mapping[str, Any], context: str) -> str: - value = _require_str(table, "reference_container", context) - if _REFERENCE_CONTAINER_RE.fullmatch(value) is None: - raise RegistryError( - f"{context}.reference_container must be a portable 'reference-' target." - ) - return value - - -def _parse_reference_adapter(table: Mapping[str, Any], context: str) -> str: - value = _require_str(table, "reference_adapter", context) - if _REFERENCE_ADAPTER_RE.fullmatch(value) is None: - raise RegistryError( - f"{context}.reference_adapter must name one module under " - "tests.parity.support.reference_adapters." - ) - return value - - -def _parse_documentation_path(table: Mapping[str, Any], context: str) -> str: - value = _require_str(table, "documentation", context) - if value.count("#") > 1 or "\\" in value: - raise RegistryError(f"{context}.documentation must be a portable documentation path.") - raw_path, separator, fragment = value.partition("#") - path = PurePosixPath(raw_path) - if ( - path.is_absolute() - or ".." in path.parts - or len(path.parts) < 2 - or path.parts[0] != "docs" - or path.suffix != ".md" - or path.as_posix() != raw_path - ): - raise RegistryError( - f"{context}.documentation must reference a normalized Markdown file under docs/." - ) - if separator and _DOCUMENTATION_FRAGMENT_RE.fullmatch(fragment) is None: - raise RegistryError(f"{context}.documentation has an invalid heading fragment.") - return value - - -def _require_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]: - value = table.get(key) - if not isinstance(value, list) or not value or any(not isinstance(item, str) for item in value): - raise RegistryError(f"{context}.{key} must be a non-empty string array.") - strings = tuple(value) - if len(set(strings)) != len(strings): - raise RegistryError(f"{context}.{key} contains duplicate values.") - return strings - - -def _optional_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]: - value = table.get(key, []) - if not isinstance(value, list) or any(not isinstance(item, str) for item in value): - raise RegistryError(f"{context}.{key} must be a string array.") - strings = tuple(value) - if len(set(strings)) != len(strings): - raise RegistryError(f"{context}.{key} contains duplicate values.") - return strings - - -def _optional_str(table: Mapping[str, Any], key: str, context: str) -> str | None: - value = table.get(key) - if value is None: - return None - if ( - not isinstance(value, str) - or not value.strip() - or value != value.strip() - or "\n" in value - or "\r" in value - ): - raise RegistryError(f"{context}.{key} must be a non-empty single-line string.") - return value - - -def _parse_hub_license( - table: Mapping[str, Any], - *, - checkpoint_license: str, - context: str, -) -> tuple[str, str | None, str | None]: - expected_fields = {"hub_license", "hub_license_name", "hub_license_link"} - unknown_fields = sorted( - key for key in table if key.startswith("hub_") and key not in expected_fields - ) - if unknown_fields: - raise RegistryError(f"{context} contains unsupported Hub license fields: {unknown_fields}.") - identifier = _require_str(table, "hub_license", context) - if identifier not in HUB_LICENSE_IDENTIFIERS: - raise RegistryError( - f"{context}.hub_license must be a supported Hugging Face license identifier." - ) - expected_identifier: str | None = None - for prefix, candidate in ( - ("MIT", "mit"), - ("Apache-2.0", "apache-2.0"), - ("CC-BY-NC-SA-4.0", "cc-by-nc-sa-4.0"), - ("Profluent-E1-Agreement", "other"), - ("Unresolved", "other"), - ): - if checkpoint_license.startswith(prefix): - expected_identifier = candidate - break - if expected_identifier is None: - raise RegistryError( - f"{context}.checkpoint_license has no declared Hugging Face identifier mapping." - ) - if identifier != expected_identifier: - raise RegistryError( - f"{context}.hub_license must be {expected_identifier!r} for " - f"checkpoint terms {checkpoint_license!r}." - ) - - name = _optional_str(table, "hub_license_name", context) - link = _optional_str(table, "hub_license_link", context) - if identifier != "other": - if name is not None or link is not None: - raise RegistryError( - f"{context} may define hub_license_name and hub_license_link only " - "when hub_license='other'." - ) - return identifier, None, None - if name is None or link is None: - raise RegistryError( - f"{context} must define hub_license_name and hub_license_link when hub_license='other'." - ) - parsed_link = urlparse(link) - if ( - parsed_link.scheme != "https" - or not parsed_link.netloc - or not parsed_link.path - or parsed_link.username is not None - or parsed_link.password is not None - ): - raise RegistryError(f"{context}.hub_license_link must be an absolute HTTPS URL.") - return identifier, name, link - - -def _require_digest_list( - table: Mapping[str, Any], key: str, context: str -) -> tuple[FileDigest, ...]: - encoded = _require_str_list(table, key, context) - digests = tuple(FileDigest.parse(value) for value in encoded) - paths = [item.path for item in digests] - if len(paths) != len(set(paths)): - raise RegistryError(f"{context}.{key} contains duplicate paths.") - return digests - - -def _validate_revision(revision: str, context: str) -> None: - if len(revision) != 40 or _HEX_RE.fullmatch(revision) is None: - raise RegistryError(f"{context} must be an immutable 40-character commit revision.") - - -def _parse_checkpoint(table: Mapping[str, Any], prefix: str, context: str) -> CheckpointSource: - repo_id = _require_str(table, f"{prefix}_repo", context) - if _REPOSITORY_ID_RE.fullmatch(repo_id) is None: - raise RegistryError(f"{context}.{prefix}_repo must be a Hugging Face repository ID.") - revision = _require_str(table, f"{prefix}_revision", context) - _validate_revision(revision, f"{context}.{prefix}_revision") - encoded_files = _require_str_list(table, f"{prefix}_files", context) - files = tuple(FileDigest.parse(value) for value in encoded_files) - paths = [item.path for item in files] - if len(paths) != len(set(paths)): - raise RegistryError(f"{context}.{prefix}_files contains duplicate paths.") - if not any(item.path.endswith(_WEIGHT_SUFFIXES) for item in files): - raise RegistryError(f"{context}.{prefix}_files does not identify a weight file.") - unresolved_files = _optional_str_list(table, f"{prefix}_unresolved_files", context) - for unresolved_path in unresolved_files: - _portable_relative_path(unresolved_path, "Unresolved checkpoint path") - if unresolved_path in paths: - raise RegistryError( - f"{context}.{prefix} marks {unresolved_path!r} both resolved and unresolved." - ) - return CheckpointSource( - repo_id=repo_id, - revision=revision, - files=files, - unresolved_files=unresolved_files, - ) - - -def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[OracleAsset, ...]: - raw = table.get("oracle_assets", []) - if not isinstance(raw, list): - raise RegistryError(f"{context}.oracle_assets must be an array of tables.") - assets: list[OracleAsset] = [] - for index, value in enumerate(raw): - asset_context = f"{context}.oracle_assets[{index}]" - if not isinstance(value, dict): - raise RegistryError(f"{asset_context} must be a table.") - expected_fields = {"role", "path", "url", "sha256", "size"} - if set(value) != expected_fields: - raise RegistryError(f"{asset_context} must contain exactly {sorted(expected_fields)}.") - role = _require_str(value, "role", asset_context) - if role not in _ALLOWED_ORACLE_ASSET_ROLES: - raise RegistryError(f"Unsupported oracle asset role: {role!r}.") - path = _require_str(value, "path", asset_context) - try: - normalized_path = _portable_relative_path(path, "Oracle asset path") - except RegistryError as error: - raise RegistryError(f"Invalid oracle asset path: {path!r}.") from error - if normalized_path.suffix != ".pt": - raise RegistryError(f"Invalid oracle asset path: {path!r}.") - url = _require_str(value, "url", asset_context) - parsed_url = urlparse(url) - if ( - parsed_url.scheme != "https" - or parsed_url.hostname != _FAIR_ESM_ASSET_HOST - or parsed_url.path != f"/fair-esm/{path}" - or parsed_url.params - or parsed_url.query - or parsed_url.fragment - ): - raise RegistryError(f"Invalid fair-esm oracle asset URL: {url!r}.") - sha256 = _require_str(value, "sha256", asset_context) - if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None: - raise RegistryError(f"Invalid oracle asset SHA-256 for {path!r}.") - size = value.get("size") - if isinstance(size, bool) or not isinstance(size, int) or size <= 0: - raise RegistryError(f"{asset_context}.size must be a positive byte count.") - assets.append( - OracleAsset( - role=role, - path=path, - url=url, - sha256=sha256, - size=size, - ) - ) - roles = [asset.role for asset in assets] - paths = [asset.path for asset in assets] - urls = [asset.url for asset in assets] - if ( - len(roles) != len(set(roles)) - or len(paths) != len(set(paths)) - or len(urls) != len(set(urls)) - ): - raise RegistryError(f"{context}.oracle_assets contains duplicate identities.") - return tuple(assets) - - -def _parse_official_golden( - table: Mapping[str, Any], - model_id: str, - context: str, -) -> OfficialGolden | None: - raw = table.get("official_golden") - if raw is None: - return None - if not isinstance(raw, dict) or set(raw) != {"metadata", "tensors"}: - raise RegistryError( - f"{context}.official_golden must contain exactly 'metadata' and 'tensors'." - ) - parsed: dict[str, FileDigest] = {} - for role in ("metadata", "tensors"): - value = raw[role] - if not isinstance(value, str): - raise RegistryError(f"{context}.official_golden.{role} must be a file digest.") - digest = FileDigest.parse(value) - if digest.algorithm != "sha256": - raise RegistryError( - f"{context}.official_golden.{role} must use an immutable SHA-256 digest." - ) - expected = f"tests/goldens/{model_id}.{'json' if role == 'metadata' else 'safetensors'}" - if digest.path != expected: - raise RegistryError(f"{context}.official_golden.{role} must use path {expected!r}.") - parsed[role] = digest - return OfficialGolden(metadata=parsed["metadata"], tensors=parsed["tensors"]) - - -def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]: - if not isinstance(raw, list) or not raw: - raise RegistryError("The manifest must contain [[attention_kernels]] entries.") - kernels: dict[str, AttentionKernelSpec] = {} - expected_variants = { - "flash_attention_2": "flash_attn2", - "flash_attention_3": "flash_attn3", - } - for index, value in enumerate(raw): - context = f"attention_kernels[{index}]" - if not isinstance(value, dict): - raise RegistryError(f"{context} must be a table.") - expected_fields = frozenset( - { - "implementation", - "repository", - "revision", - "version", - "expected_variant", - "dtypes", - } - ) - _reject_unknown_fields(value, expected_fields, context) - implementation = _require_str(value, "implementation", context) - if implementation not in expected_variants: - raise RegistryError(f"Unsupported attention kernel {implementation!r}.") - if implementation in kernels: - raise RegistryError(f"Duplicate attention kernel {implementation!r}.") - repository = _require_str(value, "repository", context) - if _REPOSITORY_ID_RE.fullmatch(repository) is None: - raise RegistryError(f"Invalid attention-kernel repository {repository!r}.") - revision = _require_str(value, "revision", context) - _validate_revision(revision, f"{context}.revision") - kernel_version = value.get("version") - if ( - isinstance(kernel_version, bool) - or not isinstance(kernel_version, int) - or kernel_version <= 0 - ): - raise RegistryError(f"{context}.version must be a positive integer.") - expected_variant = _require_str(value, "expected_variant", context) - if expected_variant != expected_variants[implementation]: - raise RegistryError( - f"{context}.expected_variant must be {expected_variants[implementation]!r}." - ) - dtypes = _require_str_list(value, "dtypes", context) - if not set(dtypes).issubset(_ALLOWED_DTYPES): - raise RegistryError(f"{context}.dtypes contains unsupported dtypes.") - kernels[implementation] = AttentionKernelSpec( - implementation=implementation, - repository=repository, - revision=revision, - version=kernel_version, - expected_variant=expected_variant, - dtypes=cast(tuple[DtypeName, ...], dtypes), - ) - if set(kernels) != set(expected_variants): - raise RegistryError("The manifest must pin both FlashAttention kernel versions.") - return kernels - - -def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]: - if not isinstance(raw, list) or not raw: - raise RegistryError("The manifest must contain at least one [[upstreams]] entry.") - upstreams: dict[str, UpstreamSource] = {} - paths: set[str] = set() - for index, value in enumerate(raw): - context = f"upstreams[{index}]" - if not isinstance(value, dict): - raise RegistryError(f"{context} must be a table.") - _reject_unknown_fields(value, _UPSTREAM_FIELDS, context) - source_id = _require_str(value, "id", context) - if _IDENTIFIER_RE.fullmatch(source_id) is None: - raise RegistryError(f"Invalid upstream ID: {source_id!r}") - if source_id in upstreams: - raise RegistryError(f"Duplicate upstream ID: {source_id!r}") - revision = _require_str(value, "revision", context) - _validate_revision(revision, f"{context}.revision") - path = _require_str(value, "path", context) - try: - normalized_path = _portable_relative_path(path, f"{context}.path") - except RegistryError as error: - raise RegistryError( - f"{context}.path must be a normalized directory directly under " - "'vendor/upstream/'." - ) from error - if ( - normalized_path.parts[:2] != ("vendor", "upstream") - or len(normalized_path.parts) != 3 - ): - raise RegistryError( - f"{context}.path must be a normalized directory directly under " - "'vendor/upstream/'." - ) - if path in paths: - raise RegistryError(f"Duplicate upstream path: {path!r}") - paths.add(path) - url = _require_str(value, "url", context) - if not url.startswith("https://github.com/") or not url.endswith(".git"): - raise RegistryError(f"{context}.url must be an HTTPS GitHub clone URL.") - license_files = _require_str_list(value, "license_files", context) - license_digests = _require_digest_list(value, "license_digests", context) - if tuple(item.path for item in license_digests) != license_files: - raise RegistryError( - f"{context}.license_digests must cover license_files in the same order." - ) - distribution_files = _require_digest_list(value, "distribution_files", context) - distribution_map = {item.path: item for item in distribution_files} - for canonical in license_digests: - distributed = distribution_map.get(canonical.path) - if distributed is None or distributed.encoded != canonical.encoded: - raise RegistryError( - f"{context}.distribution_files must include an exact copy of " - f"{canonical.path!r}." - ) - if source_id == "e1": - required_e1 = { - "LICENSE", - "ATTRIBUTION", - "NOTICE", - "Apache-2.0.txt", - "BSD-3-Clause.txt", - "MODIFICATIONS.md", - } - missing_e1 = sorted(required_e1.difference(distribution_map)) - if missing_e1: - raise RegistryError(f"{context} is missing E1 legal files: {missing_e1}") - upstreams[source_id] = UpstreamSource( - id=source_id, - path=path, - url=url, - revision=revision, - license_expression=_require_str(value, "license", context), - license_files=license_files, - license_digests=license_digests, - distribution_files=distribution_files, - ) - return upstreams - - -def _parse_families( - raw: object, - upstreams: Mapping[str, UpstreamSource], -) -> dict[str, ModelFamily]: - if not isinstance(raw, dict) or not raw: - raise RegistryError("The manifest must contain [families.] tables.") - families: dict[str, ModelFamily] = {} - for family_id, value in raw.items(): - context = f"families.{family_id}" - if _IDENTIFIER_RE.fullmatch(family_id) is None or not isinstance(value, dict): - raise RegistryError(f"Invalid family table: {family_id!r}") - checkpoint_license = _require_str(value, "checkpoint_license", context) - hub_license, hub_license_name, hub_license_link = _parse_hub_license( - value, - checkpoint_license=checkpoint_license, - context=context, - ) - _reject_unknown_fields(value, _FAMILY_FIELDS, context) - source_ids = _require_str_list(value, "upstreams", context) - unknown_sources = sorted(set(source_ids).difference(upstreams)) - if unknown_sources: - raise RegistryError(f"{context} references unknown upstreams: {unknown_sources}") - tokenizer_mode = _require_str(value, "tokenizer_mode", context) - if tokenizer_mode not in _ALLOWED_TOKENIZER_MODES: - raise RegistryError(f"Unsupported tokenizer mode in {context}: {tokenizer_mode!r}") - public_input = _require_str(value, "public_input", context) - attention = _require_str_list(value, "attention", context) - if not set(attention).issubset(_ALLOWED_ATTENTION): - raise RegistryError(f"Unsupported attention implementation in {context}.") - dtypes = _require_str_list(value, "dtypes", context) - if not set(dtypes).issubset(_ALLOWED_DTYPES): - raise RegistryError(f"Unsupported dtype in {context}.") - bf16_execution = cast( - Bf16Execution, - _require_enum( - value, - "bf16_execution", - context, - _ALLOWED_BF16_EXECUTIONS, - ), - ) - precisions = _require_str_list(value, "precisions", context) - if not set(precisions).issubset(_ALLOWED_PRECISIONS): - raise RegistryError(f"Unsupported precision policy in {context}.") - experimental_precisions = _optional_str_list( - value, - "experimental_precisions", - context, - ) - unknown_experimental_precisions = sorted( - set(experimental_precisions).difference(precisions) - ) - if unknown_experimental_precisions: - raise RegistryError( - f"{context}.experimental_precisions must be a subset of precisions; " - f"unknown values: {unknown_experimental_precisions}." - ) - extra = cast(RuntimeExtra, _require_enum(value, "extra", context, _ALLOWED_EXTRAS)) - vram_tier = cast( - VramTier, - _require_enum(value, "vram_tier", context, _ALLOWED_VRAM_TIERS), - ) - test_tiers_raw = _require_str_list(value, "test_tiers", context) - unknown_test_tiers = sorted(set(test_tiers_raw).difference(_ALLOWED_TEST_TIERS)) - if unknown_test_tiers: - raise RegistryError( - f"{context}.test_tiers contains unsupported tiers: {unknown_test_tiers}." - ) - test_tiers = cast(tuple[TestTier, ...], test_tiers_raw) - reference_container = _parse_reference_container(value, context) - reference_adapter = _parse_reference_adapter(value, context) - documentation = _parse_documentation_path(value, context) - runtime_paths = _require_str_list(value, "runtime_paths", context) - if len(runtime_paths) != len(set(runtime_paths)): - raise RegistryError(f"{context}.runtime_paths must not contain duplicates.") - for runtime_path in runtime_paths: - try: - _portable_relative_path(runtime_path, f"{context}.runtime_paths entry") - except RegistryError as error: - raise RegistryError( - f"Unsafe runtime path in {context}: {runtime_path!r}" - ) from error - if runtime_path.startswith("vendor/"): - raise RegistryError(f"Unsafe runtime path in {context}: {runtime_path!r}") - requires_complete_weight_publication = value.get( - "requires_complete_weight_publication", - False, - ) - if not isinstance(requires_complete_weight_publication, bool): - raise RegistryError( - f"{context}.requires_complete_weight_publication must be a boolean." - ) - if "weights_publication_allowed" not in value: - raise RegistryError( - f"{context}.weights_publication_allowed must be declared explicitly." - ) - weights_publication_allowed = value["weights_publication_allowed"] - if not isinstance(weights_publication_allowed, bool): - raise RegistryError(f"{context}.weights_publication_allowed must be a boolean.") - raw_auto_map = value.get("auto_map") - if not isinstance(raw_auto_map, dict) or not raw_auto_map: - raise RegistryError(f"{context}.auto_map must be a non-empty table.") - auto_map: list[tuple[str, str]] = [] - for auto_class, class_path in raw_auto_map.items(): - if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str): - raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}") - if not class_path.startswith("fastplms.") or class_path.count(".") < 2: - raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}") - auto_map.append((auto_class, class_path)) - tokenizer_class = value.get("tokenizer_class") - if tokenizer_class is not None: - if tokenizer_mode != "tokenizer": - raise RegistryError( - f"{context}.tokenizer_class requires tokenizer_mode='tokenizer'." - ) - if ( - not isinstance(tokenizer_class, str) - or not tokenizer_class.startswith("fastplms.") - or tokenizer_class.count(".") < 2 - ): - raise RegistryError( - f"Invalid tokenizer class path in {context}: {tokenizer_class!r}" - ) - backbone_model = value.get("backbone_model") - if backbone_model is not None and ( - not isinstance(backbone_model, str) - or _IDENTIFIER_RE.fullmatch(backbone_model) is None - ): - raise RegistryError( - f"{context}.backbone_model must be a valid manifest model ID." - ) - state_transform = _require_str(value, "state_transform", context) - conversion_provenance = _require_str(value, "conversion_provenance", context) - required_sections = ("Input:", "Transformation:", "Output:", "Validation:", "Limitation:") - missing_sections = [ - section for section in required_sections if section not in conversion_provenance - ] - if missing_sections or state_transform not in conversion_provenance: - raise RegistryError( - f"{context}.conversion_provenance must identify {state_transform!r} and " - f"contain mechanism-first sections; missing {missing_sections}." - ) - families[family_id] = ModelFamily( - id=family_id, - architecture=_require_str(value, "architecture", context), - upstreams=source_ids, - tokenizer_mode=tokenizer_mode, - public_input=public_input, - extra=extra, - reference_container=reference_container, - reference_adapter=reference_adapter, - attention=attention, - dtypes=cast(tuple[DtypeName, ...], dtypes), - bf16_execution=bf16_execution, - precisions=precisions, - experimental_precisions=experimental_precisions, - vram_tier=vram_tier, - checkpoint_license=checkpoint_license, - hub_license=hub_license, - state_transform=state_transform, - representative=_require_str(value, "representative", context), - documentation=documentation, - test_tiers=test_tiers, - runtime_paths=runtime_paths, - auto_map_items=tuple(auto_map), - requires_complete_weight_publication=requires_complete_weight_publication, - weights_publication_allowed=weights_publication_allowed, - tokenizer_class=tokenizer_class, - hub_license_name=hub_license_name, - hub_license_link=hub_license_link, - conversion_provenance=conversion_provenance, - backbone_model=backbone_model, - ) - return families - - -def _parse_runtime_assets( - raw: object, - families: Mapping[str, ModelFamily], -) -> dict[str, RuntimeAsset]: - if not isinstance(raw, list) or not raw: - raise RegistryError("The manifest must contain at least one [[runtime_assets]] entry.") - runtime_assets: dict[str, RuntimeAsset] = {} - identities: set[tuple[str, str, str]] = set() - for index, value in enumerate(raw): - context = f"runtime_assets[{index}]" - if not isinstance(value, dict): - raise RegistryError(f"{context} must be a table.") - _reject_unknown_fields(value, _RUNTIME_ASSET_FIELDS, context) - asset_id = _require_str(value, "id", context) - if _IDENTIFIER_RE.fullmatch(asset_id) is None: - raise RegistryError(f"Invalid runtime asset ID: {asset_id!r}") - if asset_id in runtime_assets: - raise RegistryError(f"Duplicate runtime asset ID: {asset_id!r}") - repository = _require_str(value, "repository", context) - if _REPOSITORY_ID_RE.fullmatch(repository) is None: - raise RegistryError(f"{context}.repository must be a Hugging Face repository ID.") - revision = _require_str(value, "revision", context) - _validate_revision(revision, f"{context}.revision") - path = _require_str(value, "path", context) - try: - normalized_path = _portable_relative_path(path, "Runtime asset path") - except RegistryError as error: - raise RegistryError(f"Runtime asset path is not portable: {path!r}") from error - sha256 = _require_str(value, "sha256", context) - if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None: - raise RegistryError(f"Invalid runtime asset SHA-256 for {path!r}.") - size = value.get("size") - if isinstance(size, bool) or not isinstance(size, int) or size <= 0: - raise RegistryError(f"{context}.size must be a positive byte count.") - consumer_family = _require_str(value, "consumer_family", context) - if consumer_family not in families: - raise RegistryError( - f"{context}.consumer_family references unknown family {consumer_family!r}." - ) - trust_kind = cast( - RuntimeAssetTrustKind, - _require_enum( - value, - "trust_kind", - context, - _ALLOWED_RUNTIME_ASSET_TRUST_KINDS, - ), - ) - license_expression = _require_str(value, "license", context) - offline_behavior = _require_str(value, "offline_behavior", context) - if offline_behavior not in _ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS: - raise RegistryError( - f"{context}.offline_behavior is unsupported: {offline_behavior!r}." - ) - if trust_kind == "hash_pinned_pickle" and normalized_path.suffix != ".pkl": - raise RegistryError( - f"{context}.path must end in '.pkl' for trust_kind='hash_pinned_pickle'." - ) - identity = (repository, revision, path) - if identity in identities: - raise RegistryError(f"Duplicate runtime asset identity: {identity!r}") - identities.add(identity) - runtime_assets[asset_id] = RuntimeAsset( - id=asset_id, - repository=repository, - revision=revision, - path=path, - sha256=sha256, - size=size, - consumer_family=consumer_family, - trust_kind=trust_kind, - license_expression=license_expression, - offline_behavior=offline_behavior, - ) - return runtime_assets - - -def _parse_models( - raw: object, - families: Mapping[str, ModelFamily], -) -> dict[str, ModelSpec]: - if not isinstance(raw, list) or not raw: - raise RegistryError("The manifest must contain at least one [[models]] entry.") - models: dict[str, ModelSpec] = {} - fast_repositories: set[str] = set() - for index, value in enumerate(raw): - context = f"models[{index}]" - if not isinstance(value, dict): - raise RegistryError(f"{context} must be a table.") - _reject_unknown_fields(value, _MODEL_FIELDS, context) - model_id = _require_str(value, "id", context) - if _IDENTIFIER_RE.fullmatch(model_id) is None: - raise RegistryError(f"Invalid model ID: {model_id!r}") - if model_id in models: - raise RegistryError(f"Duplicate model ID: {model_id!r}") - family_id = _require_str(value, "family", context) - if family_id not in families: - raise RegistryError(f"{context} references unknown family {family_id!r}.") - fast = _parse_checkpoint(value, "fast", context) - official = _parse_checkpoint(value, "official", context) - if fast.repo_id in fast_repositories: - raise RegistryError(f"Duplicate FastPLMs repository ID: {fast.repo_id!r}") - fast_repositories.add(fast.repo_id) - family = families[family_id] - oracle_assets = _parse_oracle_assets(value, context) - official_golden = _parse_official_golden(value, model_id, context) - size_category = _require_str(value, "size_category", context) - if size_category not in _ALLOWED_SIZE_CATEGORIES: - raise RegistryError(f"Unsupported size category in {context}: {size_category!r}") - generation_contract = cast( - GenerationContract, - _require_enum( - value, - "generation_contract", - context, - _ALLOWED_GENERATION_CONTRACTS, - ), - ) - if family.tokenizer_mode == "structure" and size_category != "structure": - raise RegistryError( - f"Structure checkpoint {model_id!r} must use size_category='structure'." - ) - artifact_source = value.get("artifact_source", "fast") - if artifact_source not in {"fast", "official"}: - raise RegistryError(f"{context}.artifact_source must be 'fast' or 'official'.") - canonical_state_sha256 = value.get("canonical_state_sha256") - if artifact_source == "official": - if ( - not isinstance(canonical_state_sha256, str) - or len(canonical_state_sha256) != 64 - or _HEX_RE.fullmatch(canonical_state_sha256) is None - ): - raise RegistryError( - f"{context}.canonical_state_sha256 must be a SHA-256 commitment " - "for an official-source artifact." - ) - elif canonical_state_sha256 is not None: - raise RegistryError( - f"{context}.canonical_state_sha256 is restricted to official-source artifacts." - ) - if family.tokenizer_mode == "tokenizer" and not any( - "tokenizer" in item.path or "vocab" in item.path for item in fast.files - ): - raise RegistryError(f"{context} does not pin a tokenizer asset.") - tokenizer_source_id = value.get("tokenizer_source") - if tokenizer_source_id is not None and ( - family.tokenizer_mode != "tokenizer" - or not isinstance(tokenizer_source_id, str) - or _IDENTIFIER_RE.fullmatch(tokenizer_source_id) is None - ): - raise RegistryError(f"{context}.tokenizer_source is invalid.") - notes = value.get("notes", "") - if not isinstance(notes, str): - raise RegistryError(f"{context}.notes must be a string.") - msa_conditioning = value.get("msa_conditioning") - if family_id == "esmfold2": - if not isinstance(msa_conditioning, bool): - raise RegistryError( - f"{context}.msa_conditioning must be an explicit boolean for " - "ESMFold2 checkpoints." - ) - elif "msa_conditioning" in value: - raise RegistryError( - f"{context}.msa_conditioning is only valid for ESMFold2 checkpoints." - ) - raw_auto_map = value.get("auto_map") - auto_map: list[tuple[str, str]] = [] - if raw_auto_map is not None: - if not isinstance(raw_auto_map, dict) or not raw_auto_map: - raise RegistryError(f"{context}.auto_map must be a non-empty table.") - for auto_class, class_path in raw_auto_map.items(): - if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str): - raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}") - if not class_path.startswith("fastplms.") or class_path.count(".") < 2: - raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}") - auto_map.append((auto_class, class_path)) - models[model_id] = ModelSpec( - id=model_id, - family=family, - fast=fast, - official=official, - size_category=size_category, - generation_contract=generation_contract, - oracle_assets=oracle_assets, - official_golden=official_golden, - artifact_source=artifact_source, - canonical_state_sha256=canonical_state_sha256, - tokenizer_source_id=tokenizer_source_id, - auto_map_items=tuple(auto_map), - notes=notes, - msa_conditioning=msa_conditioning, - ) - return models - - -def _validate_registry( - upstreams: Mapping[str, UpstreamSource], - attention_kernels: Mapping[str, AttentionKernelSpec], - families: Mapping[str, ModelFamily], - models: Mapping[str, ModelSpec], -) -> None: + for path in checkpoint.unresolved_files: + unresolved.append(f"{model.id}.{label}:{path}") + if unresolved: + detail = ", ".join(unresolved) + raise RegistryError(f"Release provenance is unresolved: {detail}") + + +def _reject_unknown_fields( + table: Mapping[str, Any], + allowed: frozenset[str], + context: str, +) -> None: + unknown = sorted(set(table).difference(allowed)) + if unknown: + raise RegistryError(f"{context} contains unknown fields: {unknown}.") + + +def _require_str(table: Mapping[str, Any], key: str, context: str) -> str: + value = table.get(key) + if not isinstance(value, str) or not value.strip(): + raise RegistryError(f"{context}.{key} must be a non-empty string.") + return value + + +def _require_enum( + table: Mapping[str, Any], + key: str, + context: str, + allowed: frozenset[str], +) -> str: + value = _require_str(table, key, context) + if value not in allowed: + raise RegistryError( + f"{context}.{key} must be one of {sorted(allowed)}; received {value!r}." + ) + return value + + +def _parse_reference_container(table: Mapping[str, Any], context: str) -> str: + value = _require_str(table, "reference_container", context) + if _REFERENCE_CONTAINER_RE.fullmatch(value) is None: + raise RegistryError( + f"{context}.reference_container must be a portable 'reference-' target." + ) + return value + + +def _parse_reference_adapter(table: Mapping[str, Any], context: str) -> str: + value = _require_str(table, "reference_adapter", context) + if _REFERENCE_ADAPTER_RE.fullmatch(value) is None: + raise RegistryError( + f"{context}.reference_adapter must name one module under " + "tests.parity.support.reference_adapters." + ) + return value + + +def _parse_documentation_path(table: Mapping[str, Any], context: str) -> str: + value = _require_str(table, "documentation", context) + if value.count("#") > 1 or "\\" in value: + raise RegistryError(f"{context}.documentation must be a portable documentation path.") + raw_path, separator, fragment = value.partition("#") + path = PurePosixPath(raw_path) + if ( + path.is_absolute() + or ".." in path.parts + or len(path.parts) < 2 + or path.parts[0] != "docs" + or path.suffix != ".md" + or path.as_posix() != raw_path + ): + raise RegistryError( + f"{context}.documentation must reference a normalized Markdown file under docs/." + ) + if separator and _DOCUMENTATION_FRAGMENT_RE.fullmatch(fragment) is None: + raise RegistryError(f"{context}.documentation has an invalid heading fragment.") + return value + + +def _require_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]: + value = table.get(key) + if not isinstance(value, list) or not value or any(not isinstance(item, str) for item in value): + raise RegistryError(f"{context}.{key} must be a non-empty string array.") + strings = tuple(value) + if len(set(strings)) != len(strings): + raise RegistryError(f"{context}.{key} contains duplicate values.") + return strings + + +def _optional_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]: + value = table.get(key, []) + if not isinstance(value, list) or any(not isinstance(item, str) for item in value): + raise RegistryError(f"{context}.{key} must be a string array.") + strings = tuple(value) + if len(set(strings)) != len(strings): + raise RegistryError(f"{context}.{key} contains duplicate values.") + return strings + + +def _optional_str(table: Mapping[str, Any], key: str, context: str) -> str | None: + value = table.get(key) + if value is None: + return None + if ( + not isinstance(value, str) + or not value.strip() + or value != value.strip() + or "\n" in value + or "\r" in value + ): + raise RegistryError(f"{context}.{key} must be a non-empty single-line string.") + return value + + +def _parse_hub_license( + table: Mapping[str, Any], + *, + checkpoint_license: str, + context: str, +) -> tuple[str, str | None, str | None]: + expected_fields = {"hub_license", "hub_license_name", "hub_license_link"} + unknown_fields = sorted( + key for key in table if key.startswith("hub_") and key not in expected_fields + ) + if unknown_fields: + raise RegistryError(f"{context} contains unsupported Hub license fields: {unknown_fields}.") + identifier = _require_str(table, "hub_license", context) + if identifier not in HUB_LICENSE_IDENTIFIERS: + raise RegistryError( + f"{context}.hub_license must be a supported Hugging Face license identifier." + ) + expected_identifier: str | None = None + for prefix, candidate in ( + ("MIT", "mit"), + ("Apache-2.0", "apache-2.0"), + ("CC-BY-NC-SA-4.0", "cc-by-nc-sa-4.0"), + ("Profluent-E1-Agreement", "other"), + ("Unresolved", "other"), + ): + if checkpoint_license.startswith(prefix): + expected_identifier = candidate + break + if expected_identifier is None: + raise RegistryError( + f"{context}.checkpoint_license has no declared Hugging Face identifier mapping." + ) + if identifier != expected_identifier: + raise RegistryError( + f"{context}.hub_license must be {expected_identifier!r} for " + f"checkpoint terms {checkpoint_license!r}." + ) + + name = _optional_str(table, "hub_license_name", context) + link = _optional_str(table, "hub_license_link", context) + if identifier != "other": + if name is not None or link is not None: + raise RegistryError( + f"{context} may define hub_license_name and hub_license_link only " + "when hub_license='other'." + ) + return identifier, None, None + if name is None or link is None: + raise RegistryError( + f"{context} must define hub_license_name and hub_license_link when hub_license='other'." + ) + parsed_link = urlparse(link) + if ( + parsed_link.scheme != "https" + or not parsed_link.netloc + or not parsed_link.path + or parsed_link.username is not None + or parsed_link.password is not None + ): + raise RegistryError(f"{context}.hub_license_link must be an absolute HTTPS URL.") + return identifier, name, link + + +def _require_digest_list( + table: Mapping[str, Any], key: str, context: str +) -> tuple[FileDigest, ...]: + encoded = _require_str_list(table, key, context) + digests = tuple(FileDigest.parse(value) for value in encoded) + paths = [item.path for item in digests] + if len(paths) != len(set(paths)): + raise RegistryError(f"{context}.{key} contains duplicate paths.") + return digests + + +def _validate_revision(revision: str, context: str) -> None: + if len(revision) != 40 or _HEX_RE.fullmatch(revision) is None: + raise RegistryError(f"{context} must be an immutable 40-character commit revision.") + + +def _parse_checkpoint(table: Mapping[str, Any], prefix: str, context: str) -> CheckpointSource: + repo_id = _require_str(table, f"{prefix}_repo", context) + if _REPOSITORY_ID_RE.fullmatch(repo_id) is None: + raise RegistryError(f"{context}.{prefix}_repo must be a Hugging Face repository ID.") + revision = _require_str(table, f"{prefix}_revision", context) + if not ( + prefix == "fast" + and table.get("publication_status") == "pending" + and revision == "unpublished" + ): + _validate_revision(revision, f"{context}.{prefix}_revision") + encoded_files = _require_str_list(table, f"{prefix}_files", context) + files = tuple(FileDigest.parse(value) for value in encoded_files) + paths = [item.path for item in files] + if len(paths) != len(set(paths)): + raise RegistryError(f"{context}.{prefix}_files contains duplicate paths.") + if not any(item.path.endswith(_WEIGHT_SUFFIXES) for item in files): + raise RegistryError(f"{context}.{prefix}_files does not identify a weight file.") + unresolved_files = _optional_str_list(table, f"{prefix}_unresolved_files", context) + for unresolved_path in unresolved_files: + _portable_relative_path(unresolved_path, "Unresolved checkpoint path") + if unresolved_path in paths: + raise RegistryError( + f"{context}.{prefix} marks {unresolved_path!r} both resolved and unresolved." + ) + return CheckpointSource( + repo_id=repo_id, + revision=revision, + files=files, + unresolved_files=unresolved_files, + ) + + +def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[OracleAsset, ...]: + raw = table.get("oracle_assets", []) + if not isinstance(raw, list): + raise RegistryError(f"{context}.oracle_assets must be an array of tables.") + assets: list[OracleAsset] = [] + for index, value in enumerate(raw): + asset_context = f"{context}.oracle_assets[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{asset_context} must be a table.") + expected_fields = {"role", "path", "url", "sha256", "size"} + if set(value) != expected_fields: + raise RegistryError(f"{asset_context} must contain exactly {sorted(expected_fields)}.") + role = _require_str(value, "role", asset_context) + if role not in _ALLOWED_ORACLE_ASSET_ROLES: + raise RegistryError(f"Unsupported oracle asset role: {role!r}.") + path = _require_str(value, "path", asset_context) + try: + normalized_path = _portable_relative_path(path, "Oracle asset path") + except RegistryError as error: + raise RegistryError(f"Invalid oracle asset path: {path!r}.") from error + if normalized_path.suffix != ".pt": + raise RegistryError(f"Invalid oracle asset path: {path!r}.") + url = _require_str(value, "url", asset_context) + parsed_url = urlparse(url) + if ( + parsed_url.scheme != "https" + or parsed_url.hostname != _FAIR_ESM_ASSET_HOST + or parsed_url.path != f"/fair-esm/{path}" + or parsed_url.params + or parsed_url.query + or parsed_url.fragment + ): + raise RegistryError(f"Invalid fair-esm oracle asset URL: {url!r}.") + sha256 = _require_str(value, "sha256", asset_context) + if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None: + raise RegistryError(f"Invalid oracle asset SHA-256 for {path!r}.") + size = value.get("size") + if isinstance(size, bool) or not isinstance(size, int) or size <= 0: + raise RegistryError(f"{asset_context}.size must be a positive byte count.") + assets.append( + OracleAsset( + role=role, + path=path, + url=url, + sha256=sha256, + size=size, + ) + ) + roles = [asset.role for asset in assets] + paths = [asset.path for asset in assets] + urls = [asset.url for asset in assets] + if ( + len(roles) != len(set(roles)) + or len(paths) != len(set(paths)) + or len(urls) != len(set(urls)) + ): + raise RegistryError(f"{context}.oracle_assets contains duplicate identities.") + return tuple(assets) + + +def _parse_official_golden( + table: Mapping[str, Any], + model_id: str, + context: str, +) -> OfficialGolden | None: + raw = table.get("official_golden") + if raw is None: + return None + if not isinstance(raw, dict) or set(raw) != {"metadata", "tensors"}: + raise RegistryError( + f"{context}.official_golden must contain exactly 'metadata' and 'tensors'." + ) + parsed: dict[str, FileDigest] = {} + for role in ("metadata", "tensors"): + value = raw[role] + if not isinstance(value, str): + raise RegistryError(f"{context}.official_golden.{role} must be a file digest.") + digest = FileDigest.parse(value) + if digest.algorithm != "sha256": + raise RegistryError( + f"{context}.official_golden.{role} must use an immutable SHA-256 digest." + ) + expected = f"tests/goldens/{model_id}.{'json' if role == 'metadata' else 'safetensors'}" + if digest.path != expected: + raise RegistryError(f"{context}.official_golden.{role} must use path {expected!r}.") + parsed[role] = digest + return OfficialGolden(metadata=parsed["metadata"], tensors=parsed["tensors"]) + + +def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain [[attention_kernels]] entries.") + kernels: dict[str, AttentionKernelSpec] = {} + expected_variants = { + "flash_attention_2": "flash_attn2", + "flash_attention_3": "flash_attn3", + } + for index, value in enumerate(raw): + context = f"attention_kernels[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + expected_fields = frozenset( + { + "implementation", + "repository", + "revision", + "version", + "expected_variant", + "dtypes", + } + ) + _reject_unknown_fields(value, expected_fields, context) + implementation = _require_str(value, "implementation", context) + if implementation not in expected_variants: + raise RegistryError(f"Unsupported attention kernel {implementation!r}.") + if implementation in kernels: + raise RegistryError(f"Duplicate attention kernel {implementation!r}.") + repository = _require_str(value, "repository", context) + if _REPOSITORY_ID_RE.fullmatch(repository) is None: + raise RegistryError(f"Invalid attention-kernel repository {repository!r}.") + revision = _require_str(value, "revision", context) + _validate_revision(revision, f"{context}.revision") + kernel_version = value.get("version") + if ( + isinstance(kernel_version, bool) + or not isinstance(kernel_version, int) + or kernel_version <= 0 + ): + raise RegistryError(f"{context}.version must be a positive integer.") + expected_variant = _require_str(value, "expected_variant", context) + if expected_variant != expected_variants[implementation]: + raise RegistryError( + f"{context}.expected_variant must be {expected_variants[implementation]!r}." + ) + dtypes = _require_str_list(value, "dtypes", context) + if not set(dtypes).issubset(_ALLOWED_DTYPES): + raise RegistryError(f"{context}.dtypes contains unsupported dtypes.") + kernels[implementation] = AttentionKernelSpec( + implementation=implementation, + repository=repository, + revision=revision, + version=kernel_version, + expected_variant=expected_variant, + dtypes=cast(tuple[DtypeName, ...], dtypes), + ) + if set(kernels) != set(expected_variants): + raise RegistryError("The manifest must pin both FlashAttention kernel versions.") + return kernels + + +def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain at least one [[upstreams]] entry.") + upstreams: dict[str, UpstreamSource] = {} + paths: set[str] = set() + for index, value in enumerate(raw): + context = f"upstreams[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + _reject_unknown_fields(value, _UPSTREAM_FIELDS, context) + source_id = _require_str(value, "id", context) + if _IDENTIFIER_RE.fullmatch(source_id) is None: + raise RegistryError(f"Invalid upstream ID: {source_id!r}") + if source_id in upstreams: + raise RegistryError(f"Duplicate upstream ID: {source_id!r}") + revision = _require_str(value, "revision", context) + _validate_revision(revision, f"{context}.revision") + path = _require_str(value, "path", context) + try: + normalized_path = _portable_relative_path(path, f"{context}.path") + except RegistryError as error: + raise RegistryError( + f"{context}.path must be a normalized directory directly under " + "'vendor/upstream/'." + ) from error + if ( + normalized_path.parts[:2] != ("vendor", "upstream") + or len(normalized_path.parts) != 3 + ): + raise RegistryError( + f"{context}.path must be a normalized directory directly under " + "'vendor/upstream/'." + ) + if path in paths: + raise RegistryError(f"Duplicate upstream path: {path!r}") + paths.add(path) + url = _require_str(value, "url", context) + if not url.startswith("https://github.com/") or not url.endswith(".git"): + raise RegistryError(f"{context}.url must be an HTTPS GitHub clone URL.") + license_files = _require_str_list(value, "license_files", context) + license_digests = _require_digest_list(value, "license_digests", context) + if tuple(item.path for item in license_digests) != license_files: + raise RegistryError( + f"{context}.license_digests must cover license_files in the same order." + ) + distribution_files = _require_digest_list(value, "distribution_files", context) + distribution_map = {item.path: item for item in distribution_files} + for canonical in license_digests: + distributed = distribution_map.get(canonical.path) + if distributed is None or distributed.encoded != canonical.encoded: + raise RegistryError( + f"{context}.distribution_files must include an exact copy of " + f"{canonical.path!r}." + ) + if source_id == "e1": + required_e1 = { + "LICENSE", + "ATTRIBUTION", + "NOTICE", + "Apache-2.0.txt", + "BSD-3-Clause.txt", + "MODIFICATIONS.md", + } + missing_e1 = sorted(required_e1.difference(distribution_map)) + if missing_e1: + raise RegistryError(f"{context} is missing E1 legal files: {missing_e1}") + upstreams[source_id] = UpstreamSource( + id=source_id, + path=path, + url=url, + revision=revision, + license_expression=_require_str(value, "license", context), + license_files=license_files, + license_digests=license_digests, + distribution_files=distribution_files, + ) + return upstreams + + +def _parse_families( + raw: object, + upstreams: Mapping[str, UpstreamSource], +) -> dict[str, ModelFamily]: + if not isinstance(raw, dict) or not raw: + raise RegistryError("The manifest must contain [families.] tables.") + families: dict[str, ModelFamily] = {} + for family_id, value in raw.items(): + context = f"families.{family_id}" + if _IDENTIFIER_RE.fullmatch(family_id) is None or not isinstance(value, dict): + raise RegistryError(f"Invalid family table: {family_id!r}") + checkpoint_license = _require_str(value, "checkpoint_license", context) + hub_license, hub_license_name, hub_license_link = _parse_hub_license( + value, + checkpoint_license=checkpoint_license, + context=context, + ) + _reject_unknown_fields(value, _FAMILY_FIELDS, context) + source_ids = _require_str_list(value, "upstreams", context) + unknown_sources = sorted(set(source_ids).difference(upstreams)) + if unknown_sources: + raise RegistryError(f"{context} references unknown upstreams: {unknown_sources}") + tokenizer_mode = _require_str(value, "tokenizer_mode", context) + if tokenizer_mode not in _ALLOWED_TOKENIZER_MODES: + raise RegistryError(f"Unsupported tokenizer mode in {context}: {tokenizer_mode!r}") + public_input = _require_str(value, "public_input", context) + attention = _require_str_list(value, "attention", context) + if not set(attention).issubset(_ALLOWED_ATTENTION): + raise RegistryError(f"Unsupported attention implementation in {context}.") + dtypes = _require_str_list(value, "dtypes", context) + if not set(dtypes).issubset(_ALLOWED_DTYPES): + raise RegistryError(f"Unsupported dtype in {context}.") + bf16_execution = cast( + Bf16Execution, + _require_enum( + value, + "bf16_execution", + context, + _ALLOWED_BF16_EXECUTIONS, + ), + ) + precisions = _require_str_list(value, "precisions", context) + if not set(precisions).issubset(_ALLOWED_PRECISIONS): + raise RegistryError(f"Unsupported precision policy in {context}.") + experimental_precisions = _optional_str_list( + value, + "experimental_precisions", + context, + ) + unknown_experimental_precisions = sorted( + set(experimental_precisions).difference(precisions) + ) + if unknown_experimental_precisions: + raise RegistryError( + f"{context}.experimental_precisions must be a subset of precisions; " + f"unknown values: {unknown_experimental_precisions}." + ) + extra = cast(RuntimeExtra, _require_enum(value, "extra", context, _ALLOWED_EXTRAS)) + vram_tier = cast( + VramTier, + _require_enum(value, "vram_tier", context, _ALLOWED_VRAM_TIERS), + ) + test_tiers_raw = _require_str_list(value, "test_tiers", context) + unknown_test_tiers = sorted(set(test_tiers_raw).difference(_ALLOWED_TEST_TIERS)) + if unknown_test_tiers: + raise RegistryError( + f"{context}.test_tiers contains unsupported tiers: {unknown_test_tiers}." + ) + test_tiers = cast(tuple[TestTier, ...], test_tiers_raw) + reference_container = _parse_reference_container(value, context) + reference_adapter = _parse_reference_adapter(value, context) + documentation = _parse_documentation_path(value, context) + runtime_paths = _require_str_list(value, "runtime_paths", context) + if len(runtime_paths) != len(set(runtime_paths)): + raise RegistryError(f"{context}.runtime_paths must not contain duplicates.") + for runtime_path in runtime_paths: + try: + _portable_relative_path(runtime_path, f"{context}.runtime_paths entry") + except RegistryError as error: + raise RegistryError( + f"Unsafe runtime path in {context}: {runtime_path!r}" + ) from error + if runtime_path.startswith("vendor/"): + raise RegistryError(f"Unsafe runtime path in {context}: {runtime_path!r}") + requires_complete_weight_publication = value.get( + "requires_complete_weight_publication", + False, + ) + if not isinstance(requires_complete_weight_publication, bool): + raise RegistryError( + f"{context}.requires_complete_weight_publication must be a boolean." + ) + if "weights_publication_allowed" not in value: + raise RegistryError( + f"{context}.weights_publication_allowed must be declared explicitly." + ) + weights_publication_allowed = value["weights_publication_allowed"] + if not isinstance(weights_publication_allowed, bool): + raise RegistryError(f"{context}.weights_publication_allowed must be a boolean.") + raw_auto_map = value.get("auto_map") + if not isinstance(raw_auto_map, dict) or not raw_auto_map: + raise RegistryError(f"{context}.auto_map must be a non-empty table.") + auto_map: list[tuple[str, str]] = [] + for auto_class, class_path in raw_auto_map.items(): + if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str): + raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}") + if not class_path.startswith("fastplms.") or class_path.count(".") < 2: + raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}") + auto_map.append((auto_class, class_path)) + tokenizer_class = value.get("tokenizer_class") + if tokenizer_class is not None: + if tokenizer_mode != "tokenizer": + raise RegistryError( + f"{context}.tokenizer_class requires tokenizer_mode='tokenizer'." + ) + if ( + not isinstance(tokenizer_class, str) + or not tokenizer_class.startswith("fastplms.") + or tokenizer_class.count(".") < 2 + ): + raise RegistryError( + f"Invalid tokenizer class path in {context}: {tokenizer_class!r}" + ) + backbone_model = value.get("backbone_model") + if backbone_model is not None and ( + not isinstance(backbone_model, str) + or _IDENTIFIER_RE.fullmatch(backbone_model) is None + ): + raise RegistryError( + f"{context}.backbone_model must be a valid manifest model ID." + ) + state_transform = _require_str(value, "state_transform", context) + conversion_provenance = _require_str(value, "conversion_provenance", context) + required_sections = ("Input:", "Transformation:", "Output:", "Validation:", "Limitation:") + missing_sections = [ + section for section in required_sections if section not in conversion_provenance + ] + if missing_sections or state_transform not in conversion_provenance: + raise RegistryError( + f"{context}.conversion_provenance must identify {state_transform!r} and " + f"contain mechanism-first sections; missing {missing_sections}." + ) + families[family_id] = ModelFamily( + id=family_id, + architecture=_require_str(value, "architecture", context), + upstreams=source_ids, + tokenizer_mode=tokenizer_mode, + public_input=public_input, + extra=extra, + reference_container=reference_container, + reference_adapter=reference_adapter, + attention=attention, + dtypes=cast(tuple[DtypeName, ...], dtypes), + bf16_execution=bf16_execution, + precisions=precisions, + experimental_precisions=experimental_precisions, + vram_tier=vram_tier, + checkpoint_license=checkpoint_license, + hub_license=hub_license, + state_transform=state_transform, + representative=_require_str(value, "representative", context), + documentation=documentation, + test_tiers=test_tiers, + runtime_paths=runtime_paths, + auto_map_items=tuple(auto_map), + requires_complete_weight_publication=requires_complete_weight_publication, + weights_publication_allowed=weights_publication_allowed, + tokenizer_class=tokenizer_class, + hub_license_name=hub_license_name, + hub_license_link=hub_license_link, + conversion_provenance=conversion_provenance, + backbone_model=backbone_model, + ) + return families + + +def _parse_runtime_assets( + raw: object, + families: Mapping[str, ModelFamily], +) -> dict[str, RuntimeAsset]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain at least one [[runtime_assets]] entry.") + runtime_assets: dict[str, RuntimeAsset] = {} + identities: set[tuple[str, str, str]] = set() + for index, value in enumerate(raw): + context = f"runtime_assets[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + _reject_unknown_fields(value, _RUNTIME_ASSET_FIELDS, context) + asset_id = _require_str(value, "id", context) + if _IDENTIFIER_RE.fullmatch(asset_id) is None: + raise RegistryError(f"Invalid runtime asset ID: {asset_id!r}") + if asset_id in runtime_assets: + raise RegistryError(f"Duplicate runtime asset ID: {asset_id!r}") + repository = _require_str(value, "repository", context) + if _REPOSITORY_ID_RE.fullmatch(repository) is None: + raise RegistryError(f"{context}.repository must be a Hugging Face repository ID.") + revision = _require_str(value, "revision", context) + _validate_revision(revision, f"{context}.revision") + path = _require_str(value, "path", context) + try: + normalized_path = _portable_relative_path(path, "Runtime asset path") + except RegistryError as error: + raise RegistryError(f"Runtime asset path is not portable: {path!r}") from error + sha256 = _require_str(value, "sha256", context) + if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None: + raise RegistryError(f"Invalid runtime asset SHA-256 for {path!r}.") + size = value.get("size") + if isinstance(size, bool) or not isinstance(size, int) or size <= 0: + raise RegistryError(f"{context}.size must be a positive byte count.") + consumer_family = _require_str(value, "consumer_family", context) + if consumer_family not in families: + raise RegistryError( + f"{context}.consumer_family references unknown family {consumer_family!r}." + ) + trust_kind = cast( + RuntimeAssetTrustKind, + _require_enum( + value, + "trust_kind", + context, + _ALLOWED_RUNTIME_ASSET_TRUST_KINDS, + ), + ) + license_expression = _require_str(value, "license", context) + offline_behavior = _require_str(value, "offline_behavior", context) + if offline_behavior not in _ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS: + raise RegistryError( + f"{context}.offline_behavior is unsupported: {offline_behavior!r}." + ) + if trust_kind == "hash_pinned_pickle" and normalized_path.suffix != ".pkl": + raise RegistryError( + f"{context}.path must end in '.pkl' for trust_kind='hash_pinned_pickle'." + ) + identity = (repository, revision, path) + if identity in identities: + raise RegistryError(f"Duplicate runtime asset identity: {identity!r}") + identities.add(identity) + runtime_assets[asset_id] = RuntimeAsset( + id=asset_id, + repository=repository, + revision=revision, + path=path, + sha256=sha256, + size=size, + consumer_family=consumer_family, + trust_kind=trust_kind, + license_expression=license_expression, + offline_behavior=offline_behavior, + ) + return runtime_assets + + +def _parse_models( + raw: object, + families: Mapping[str, ModelFamily], +) -> dict[str, ModelSpec]: + if not isinstance(raw, list) or not raw: + raise RegistryError("The manifest must contain at least one [[models]] entry.") + models: dict[str, ModelSpec] = {} + fast_repositories: set[str] = set() + for index, value in enumerate(raw): + context = f"models[{index}]" + if not isinstance(value, dict): + raise RegistryError(f"{context} must be a table.") + _reject_unknown_fields(value, _MODEL_FIELDS, context) + model_id = _require_str(value, "id", context) + if _IDENTIFIER_RE.fullmatch(model_id) is None: + raise RegistryError(f"Invalid model ID: {model_id!r}") + if model_id in models: + raise RegistryError(f"Duplicate model ID: {model_id!r}") + family_id = _require_str(value, "family", context) + if family_id not in families: + raise RegistryError(f"{context} references unknown family {family_id!r}.") + fast = _parse_checkpoint(value, "fast", context) + official = _parse_checkpoint(value, "official", context) + if fast.repo_id in fast_repositories: + raise RegistryError(f"Duplicate FastPLMs repository ID: {fast.repo_id!r}") + fast_repositories.add(fast.repo_id) + family = families[family_id] + publication_status = value.get("publication_status", "published") + if publication_status not in {"pending", "published"}: + raise RegistryError(f"{context}.publication_status must be pending or published.") + if publication_status == "pending" and fast.revision != "unpublished": + raise RegistryError(f"{context}: pending publication must use unpublished revision.") + backbone = None + if "backbone" in value: + raw_backbone = value["backbone"] + if family_id != "esmfold2" or not isinstance(raw_backbone, dict): + raise RegistryError(f"{context}.backbone is only valid for ESMFold2.") + _reject_unknown_fields(raw_backbone, {"repo", "revision", "files"}, context) + backbone = _parse_checkpoint( + {f"official_{key}": item for key, item in raw_backbone.items()}, + "official", + f"{context}.backbone", + ) + oracle_assets = _parse_oracle_assets(value, context) + official_golden = _parse_official_golden(value, model_id, context) + size_category = _require_str(value, "size_category", context) + if size_category not in _ALLOWED_SIZE_CATEGORIES: + raise RegistryError(f"Unsupported size category in {context}: {size_category!r}") + generation_contract = cast( + GenerationContract, + _require_enum( + value, + "generation_contract", + context, + _ALLOWED_GENERATION_CONTRACTS, + ), + ) + if family.tokenizer_mode == "structure" and size_category != "structure": + raise RegistryError( + f"Structure checkpoint {model_id!r} must use size_category='structure'." + ) + artifact_source = value.get("artifact_source", "fast") + if artifact_source not in {"fast", "official"}: + raise RegistryError(f"{context}.artifact_source must be 'fast' or 'official'.") + canonical_state_sha256 = value.get("canonical_state_sha256") + if artifact_source == "official": + if ( + not isinstance(canonical_state_sha256, str) + or len(canonical_state_sha256) != 64 + or _HEX_RE.fullmatch(canonical_state_sha256) is None + ): + raise RegistryError( + f"{context}.canonical_state_sha256 must be a SHA-256 commitment " + "for an official-source artifact." + ) + elif canonical_state_sha256 is not None: + raise RegistryError( + f"{context}.canonical_state_sha256 is restricted to official-source artifacts." + ) + if family.tokenizer_mode == "tokenizer" and not any( + "tokenizer" in item.path or "vocab" in item.path for item in fast.files + ): + raise RegistryError(f"{context} does not pin a tokenizer asset.") + tokenizer_source_id = value.get("tokenizer_source") + if tokenizer_source_id is not None and ( + family.tokenizer_mode != "tokenizer" + or not isinstance(tokenizer_source_id, str) + or _IDENTIFIER_RE.fullmatch(tokenizer_source_id) is None + ): + raise RegistryError(f"{context}.tokenizer_source is invalid.") + notes = value.get("notes", "") + if not isinstance(notes, str): + raise RegistryError(f"{context}.notes must be a string.") + msa_conditioning = value.get("msa_conditioning") + if family_id == "esmfold2": + if not isinstance(msa_conditioning, bool): + raise RegistryError( + f"{context}.msa_conditioning must be an explicit boolean for " + "ESMFold2 checkpoints." + ) + elif "msa_conditioning" in value: + raise RegistryError( + f"{context}.msa_conditioning is only valid for ESMFold2 checkpoints." + ) + raw_auto_map = value.get("auto_map") + auto_map: list[tuple[str, str]] = [] + if raw_auto_map is not None: + if not isinstance(raw_auto_map, dict) or not raw_auto_map: + raise RegistryError(f"{context}.auto_map must be a non-empty table.") + for auto_class, class_path in raw_auto_map.items(): + if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str): + raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}") + if not class_path.startswith("fastplms.") or class_path.count(".") < 2: + raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}") + auto_map.append((auto_class, class_path)) + models[model_id] = ModelSpec( + id=model_id, + family=family, + fast=fast, + official=official, + size_category=size_category, + generation_contract=generation_contract, + oracle_assets=oracle_assets, + official_golden=official_golden, + artifact_source=artifact_source, + canonical_state_sha256=canonical_state_sha256, + tokenizer_source_id=tokenizer_source_id, + auto_map_items=tuple(auto_map), + notes=notes, + backbone=backbone, + backbone_model=value.get("backbone_model"), + publication_status=publication_status, + msa_conditioning=msa_conditioning, + ) + return models + + +def _validate_registry( + upstreams: Mapping[str, UpstreamSource], + attention_kernels: Mapping[str, AttentionKernelSpec], + families: Mapping[str, ModelFamily], + models: Mapping[str, ModelSpec], +) -> None: + for spec in models.values(): + if spec.tokenizer_source_id is None: + continue + source = models.get(spec.tokenizer_source_id) + if source is None: + raise RegistryError( + f"Model {spec.id!r} references unknown tokenizer source " + f"{spec.tokenizer_source_id!r}." + ) + if not any( + PurePosixPath(item.path).name + in { + "added_tokens.json", + "merges.txt", + "sentencepiece.bpe.model", + "special_tokens_map.json", + "spiece.model", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "vocab.txt", + } + for item in source.official.files + ): + raise RegistryError( + f"Tokenizer source {source.id!r} has no official tokenizer assets." + ) for spec in models.values(): - if spec.tokenizer_source_id is None: - continue - source = models.get(spec.tokenizer_source_id) - if source is None: - raise RegistryError( - f"Model {spec.id!r} references unknown tokenizer source " - f"{spec.tokenizer_source_id!r}." - ) - if not any( - PurePosixPath(item.path).name - in { - "added_tokens.json", - "merges.txt", - "sentencepiece.bpe.model", - "special_tokens_map.json", - "spiece.model", - "tokenizer.json", - "tokenizer_config.json", - "vocab.json", - "vocab.txt", - } - for item in source.official.files - ): + if spec.backbone is not None and spec.backbone_model is None: raise RegistryError( - f"Tokenizer source {source.id!r} has no official tokenizer assets." - ) - expected_esmfold2 = { - "esmfold2": ("Synthyra/ESMFold2", "biohub/ESMFold2"), - "esmfold2_fast": ("Synthyra/ESMFold2-Fast", "biohub/ESMFold2-Fast"), - "esmfold2_experimental_cutoff2025": ( - "Synthyra/ESMFold2-Experimental-Cutoff2025", - "biohub/ESMFold2-Experimental-Cutoff2025", - ), - "esmfold2_experimental_fast_cutoff2025": ( - "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025", - "biohub/ESMFold2-Experimental-Fast-Cutoff2025", - ), - } - actual_esmfold2 = { - model.id: (model.fast.repo_id, model.official.repo_id) - for model in models.values() - if model.family.id == "esmfold2" - } - if actual_esmfold2 != expected_esmfold2: - raise RegistryError( - "ESMFold2 support must contain exactly the four approved model IDs and " - "official/Synthyra repositories." - ) - - golden_paths: list[str] = [] - for model in models.values(): - if model.official_golden is not None: - golden_paths.extend( - ( - model.official_golden.metadata.path, - model.official_golden.tensors.path, - ) + f"{spec.id}: backbone_model is required when a backbone source is declared." ) - if len(golden_paths) != len(set(golden_paths)): - raise RegistryError("Official golden paths must be unique across model declarations.") - unused_upstreams = sorted( - set(upstreams).difference( - source for family in families.values() for source in family.upstreams - ) - ) - if unused_upstreams: - raise RegistryError( - f"Upstream sources are not connected to a model family: {unused_upstreams}" - ) - advertised_flash = { - implementation - for family in families.values() - for implementation in family.attention - if implementation.startswith("flash_attention_") - } - missing_kernels = sorted(advertised_flash.difference(attention_kernels)) - if missing_kernels: - raise RegistryError( - f"Advertised FlashAttention backends lack kernel specs: {missing_kernels}." - ) - for family in families.values(): - for implementation in family.attention: - kernel = attention_kernels.get(implementation) - if kernel is not None and not set(family.dtypes).intersection(kernel.dtypes): + if spec.backbone_model is not None: + if ( + spec.family.id != "esmfold2" + or not isinstance(spec.backbone_model, str) + or spec.backbone_model not in models + ): + raise RegistryError(f"{spec.id}: unknown ESMFold2 backbone model.") + if models[spec.backbone_model].family.id != "esm_plusplus" or spec.backbone is None: raise RegistryError( - f"Family {family.id!r} and attention kernel {implementation!r} " - "have no supported dtype in common." + f"{spec.id}: backbone must pin an ESM++ dependency and its source." ) - family_models = [model for model in models.values() if model.family.id == family.id] - if not family_models: - raise RegistryError(f"Family {family.id!r} has no checkpoints.") - representative = models.get(family.representative) - if representative is None or representative.family.id != family.id: - raise RegistryError( - f"Family {family.id!r} has invalid representative {family.representative!r}." - ) - if family.backbone_model is not None and family.backbone_model not in models: - raise RegistryError( - f"Family {family.id!r} references unknown backbone model " - f"{family.backbone_model!r}." - ) - - -def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry: - try: - manifest = tomllib.loads(raw_bytes.decode("utf-8")) - except (UnicodeDecodeError, tomllib.TOMLDecodeError) as error: - raise RegistryError(f"Unable to parse model manifest: {error}") from error - _reject_unknown_fields(manifest, _ROOT_FIELDS, "manifest") - if manifest.get("schema_version") != 1: - raise RegistryError("Unsupported model manifest schema_version; expected 1.") - legal_files = _require_digest_list(manifest, "legal_files", "manifest") - required_legal_paths = {"LICENSE", "THIRD_PARTY_NOTICES.md"} - if {item.path for item in legal_files} != required_legal_paths: - raise RegistryError("manifest.legal_files must contain LICENSE and THIRD_PARTY_NOTICES.md.") - attention_kernels = _parse_attention_kernels(manifest.get("attention_kernels")) - upstreams = _parse_upstreams(manifest.get("upstreams")) - families = _parse_families(manifest.get("families"), upstreams) - runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families) - models = _parse_models(manifest.get("models"), families) - _validate_registry(upstreams, attention_kernels, families, models) - return ModelRegistry( - schema_version=1, - upstreams=upstreams, - attention_kernels=attention_kernels, - families=families, - models=models, - runtime_assets=runtime_assets, - legal_files=legal_files, - ) - - -def load_model_registry(path: str | Path | None = None) -> ModelRegistry: - """Load and validate a model manifest without importing model code.""" - - if path is None: - manifest = resources.files("fastplms").joinpath("models.toml") - return _load_manifest_bytes(manifest.read_bytes()) - return _load_manifest_bytes(Path(path).read_bytes()) - - -@lru_cache(maxsize=1) -def get_model_registry() -> ModelRegistry: - """Return the validated package registry, cached after its first read.""" - - return load_model_registry() - - -def get_model_spec(model_id: str) -> ModelSpec: - """Return one model specification by its stable manifest ID.""" - - try: - return get_model_registry()[model_id] - except KeyError as error: - supported = ", ".join(get_model_registry()) - raise KeyError( - f"Unknown FastPLMs model ID {model_id!r}. Supported IDs: {supported}" - ) from error - - -__all__ = [ - "HUB_LICENSE_IDENTIFIERS", - "CheckpointSource", - "FileDigest", - "GenerationContract", - "ModelFamily", - "ModelRegistry", - "ModelSpec", - "OracleAsset", - "RegistryError", - "RuntimeAsset", - "RuntimeAssetTrustKind", - "RuntimeExtra", - "TestTier", - "UpstreamSource", - "VramTier", - "get_model_registry", - "get_model_spec", - "load_model_registry", -] + expected_esmfold2 = { + "esmfold2_300": ( + "Synthyra/ESMFold2-300", + "biohub/ESMFold2-Experimental-Fast-base300M-step1500k", + ), + "esmfold2_600": ( + "Synthyra/ESMFold2-600", + "biohub/ESMFold2-Experimental-Fast-base600M-step1500k", + ), + "esmfold2": ("Synthyra/ESMFold2", "biohub/ESMFold2"), + "esmfold2_fast": ("Synthyra/ESMFold2-Fast", "biohub/ESMFold2-Fast"), + "esmfold2_experimental_cutoff2025": ( + "Synthyra/ESMFold2-Experimental-Cutoff2025", + "biohub/ESMFold2-Experimental-Cutoff2025", + ), + "esmfold2_experimental_fast_cutoff2025": ( + "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025", + "biohub/ESMFold2-Experimental-Fast-Cutoff2025", + ), + } + actual_esmfold2 = { + model.id: (model.fast.repo_id, model.official.repo_id) + for model in models.values() + if model.family.id == "esmfold2" + } + if actual_esmfold2 != expected_esmfold2: + raise RegistryError( + "ESMFold2 support must contain exactly the six approved model IDs and " + "official/Synthyra repositories." + ) + + golden_paths: list[str] = [] + for model in models.values(): + if model.official_golden is not None: + golden_paths.extend( + ( + model.official_golden.metadata.path, + model.official_golden.tensors.path, + ) + ) + if len(golden_paths) != len(set(golden_paths)): + raise RegistryError("Official golden paths must be unique across model declarations.") + unused_upstreams = sorted( + set(upstreams).difference( + source for family in families.values() for source in family.upstreams + ) + ) + if unused_upstreams: + raise RegistryError( + f"Upstream sources are not connected to a model family: {unused_upstreams}" + ) + advertised_flash = { + implementation + for family in families.values() + for implementation in family.attention + if implementation.startswith("flash_attention_") + } + missing_kernels = sorted(advertised_flash.difference(attention_kernels)) + if missing_kernels: + raise RegistryError( + f"Advertised FlashAttention backends lack kernel specs: {missing_kernels}." + ) + for family in families.values(): + for implementation in family.attention: + kernel = attention_kernels.get(implementation) + if kernel is not None and not set(family.dtypes).intersection(kernel.dtypes): + raise RegistryError( + f"Family {family.id!r} and attention kernel {implementation!r} " + "have no supported dtype in common." + ) + family_models = [model for model in models.values() if model.family.id == family.id] + if not family_models: + raise RegistryError(f"Family {family.id!r} has no checkpoints.") + representative = models.get(family.representative) + if representative is None or representative.family.id != family.id: + raise RegistryError( + f"Family {family.id!r} has invalid representative {family.representative!r}." + ) + if family.backbone_model is not None and family.backbone_model not in models: + raise RegistryError( + f"Family {family.id!r} references unknown backbone model " + f"{family.backbone_model!r}." + ) + + +def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry: + try: + manifest = tomllib.loads(raw_bytes.decode("utf-8")) + except (UnicodeDecodeError, tomllib.TOMLDecodeError) as error: + raise RegistryError(f"Unable to parse model manifest: {error}") from error + _reject_unknown_fields(manifest, _ROOT_FIELDS, "manifest") + if manifest.get("schema_version") != 1: + raise RegistryError("Unsupported model manifest schema_version; expected 1.") + legal_files = _require_digest_list(manifest, "legal_files", "manifest") + required_legal_paths = {"LICENSE", "THIRD_PARTY_NOTICES.md"} + if {item.path for item in legal_files} != required_legal_paths: + raise RegistryError("manifest.legal_files must contain LICENSE and THIRD_PARTY_NOTICES.md.") + attention_kernels = _parse_attention_kernels(manifest.get("attention_kernels")) + upstreams = _parse_upstreams(manifest.get("upstreams")) + families = _parse_families(manifest.get("families"), upstreams) + runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families) + models = _parse_models(manifest.get("models"), families) + _validate_registry(upstreams, attention_kernels, families, models) + return ModelRegistry( + schema_version=1, + upstreams=upstreams, + attention_kernels=attention_kernels, + families=families, + models=models, + runtime_assets=runtime_assets, + legal_files=legal_files, + ) + + +def load_model_registry(path: str | Path | None = None) -> ModelRegistry: + """Load and validate a model manifest without importing model code.""" + + if path is None: + manifest = resources.files("fastplms").joinpath("models.toml") + return _load_manifest_bytes(manifest.read_bytes()) + return _load_manifest_bytes(Path(path).read_bytes()) + + +@lru_cache(maxsize=1) +def get_model_registry() -> ModelRegistry: + """Return the validated package registry, cached after its first read.""" + + return load_model_registry() + + +def get_model_spec(model_id: str) -> ModelSpec: + """Return one model specification by its stable manifest ID.""" + + try: + return get_model_registry()[model_id] + except KeyError as error: + supported = ", ".join(get_model_registry()) + raise KeyError( + f"Unknown FastPLMs model ID {model_id!r}. Supported IDs: {supported}" + ) from error + + +__all__ = [ + "HUB_LICENSE_IDENTIFIERS", + "CheckpointSource", + "FileDigest", + "GenerationContract", + "ModelFamily", + "ModelRegistry", + "ModelSpec", + "OracleAsset", + "RegistryError", + "RuntimeAsset", + "RuntimeAssetTrustKind", + "RuntimeExtra", + "TestTier", + "UpstreamSource", + "VramTier", + "get_model_registry", + "get_model_spec", + "load_model_registry", +]