| """Validation and loading for the experiment coverage registry.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Literal |
|
|
| import yaml |
| from pydantic import BaseModel, ConfigDict, Field, model_validator |
|
|
|
|
| MANY_STEPS = [1, 2, 4, 8, 16, 32, 1024] |
| FEW_STEPS = [1, 2, 4, 8, 16, 32] |
| FIXED_1024_STEPS = [1024] |
| RDLM_OFFICIAL_STEPS = [1000, 1024] |
|
|
| MODEL_IDENTIFIERS = { |
| "flm": ("many", "dlb-flm", "flm", "flm"), |
| "fmlm": ("few", "dlb-flm", "flm", "flm"), |
| "langflow": ("many", "dlb-langflow", "langflow", "langflow"), |
| "duo": ("many", "dlb-duo", "duo", "duo"), |
| "duo_dcd": ("few", "dlb-duo", "duo", "duo"), |
| "mdlm": ("many", "dlb-mdlm", "mdlm", "mdlm"), |
| "candi": ("many", "dlb-candi", "candi", "candi"), |
| "rdlm": ("many", "dlb-rdlm", "rdlm", "rdlm"), |
| "mdlm_sdtt": ("few", "dlb-sdtt", "sdtt", "sdtt"), |
| "duo_di4c": ("few", "dlb-di4c", "di4c", "di4c"), |
| "mdlm_di4c": ("few", "dlb-di4c", "di4c", "di4c"), |
| } |
|
|
| Category = Literal["many", "few", "fixed_1024"] |
| SupportStatus = Literal["supported", "unsupported"] |
| Provenance = Literal["official", "reference_reproduction", "self_trained"] |
|
|
|
|
| class RegistryModel(BaseModel): |
| model_config = ConfigDict(extra="forbid") |
|
|
|
|
| class DatasetSupport(RegistryModel): |
| status: SupportStatus |
| provenance: Provenance | None = None |
| reason: str | None = None |
| train_recipe: str | None = Field(default=None, pattern=r"^[a-z][a-z0-9_]*$") |
|
|
| @model_validator(mode="after") |
| def validate_support_details(self) -> "DatasetSupport": |
| if self.status == "supported": |
| if self.provenance is None or self.reason is not None: |
| raise ValueError("supported datasets require provenance and no reason") |
| elif ( |
| self.reason is None |
| or self.reason.strip() == "" |
| or self.provenance is not None |
| or self.train_recipe is not None |
| ): |
| raise ValueError( |
| "unsupported datasets require a reason and no provenance or train recipe" |
| ) |
| return self |
|
|
|
|
| class ModelRegistryEntry(RegistryModel): |
| category: Category |
| step_override: list[int] | None = None |
| environment: str = Field(pattern=r"^[a-z][a-z0-9-]*$") |
| adapter: str = Field(pattern=r"^[a-z][a-z0-9-]*$") |
| source: str = Field(pattern=r"^[a-z][a-z0-9-]*$") |
| datasets: dict[Literal["lm1b", "owt"], DatasetSupport] |
|
|
| @model_validator(mode="after") |
| def require_all_datasets(self) -> "ModelRegistryEntry": |
| if set(self.datasets) != {"lm1b", "owt"}: |
| raise ValueError("each model must declare lm1b and owt support") |
| if self.step_override is not None: |
| if self.step_override != sorted(set(self.step_override)): |
| raise ValueError("step overrides must be sorted and unique") |
| if any(step <= 0 for step in self.step_override): |
| raise ValueError("step overrides must be positive") |
| return self |
|
|
|
|
| class ExperimentRegistry(RegistryModel): |
| step_grids: dict[Category, list[int]] |
| models: dict[str, ModelRegistryEntry] |
|
|
| @model_validator(mode="after") |
| def validate_coverage(self) -> "ExperimentRegistry": |
| if self.step_grids != { |
| "many": MANY_STEPS, |
| "few": FEW_STEPS, |
| "fixed_1024": FIXED_1024_STEPS, |
| }: |
| raise ValueError("step grids must match the prescribed schedules") |
| if set(self.models) != set(MODEL_IDENTIFIERS): |
| raise ValueError("registry must contain the complete baseline model scope") |
| for model_id, model in self.models.items(): |
| if (model.category, model.environment, model.adapter, model.source) != ( |
| MODEL_IDENTIFIERS[model_id] |
| ): |
| raise ValueError(f"registry identifiers do not match {model_id}") |
| if model_id == "rdlm": |
| if model.step_override != RDLM_OFFICIAL_STEPS: |
| raise ValueError("RDLM must use the official/default step override") |
| elif model.step_override is not None: |
| raise ValueError(f"{model_id} must use its category step grid") |
| return self |
|
|
|
|
| def load_registry(path: Path) -> ExperimentRegistry: |
| """Load and validate a registry YAML document from *path*.""" |
|
|
| with path.open(encoding="utf-8") as registry_file: |
| document = yaml.safe_load(registry_file) |
| if not isinstance(document, dict): |
| raise ValueError("registry document must be a mapping") |
| return ExperimentRegistry.model_validate(document) |
|
|
|
|
| def step_grid_for_model(registry: ExperimentRegistry, model_id: str) -> list[int]: |
| """Return the canonical step grid for one logical baseline model.""" |
|
|
| model = registry.models[model_id] |
| return list(model.step_override or registry.step_grids[model.category]) |
|
|
