100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
|
"""Runtime registries for typed config variants.
|
||
|
|
|
||
|
|
Core Wall-X keeps only base config types in ``wall_x.config``. Optional
|
||
|
|
datasets, model families, and optimizer variants register their dataclasses
|
||
|
|
from their own packages, so trimmed distributions do not leave dead config
|
||
|
|
entries behind.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections.abc import Callable
|
||
|
|
from typing import Any, TypeVar
|
||
|
|
|
||
|
|
T = TypeVar("T")
|
||
|
|
|
||
|
|
_DATA_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||
|
|
_MODEL_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||
|
|
_OPTIMIZER_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||
|
|
_SCHEDULER_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||
|
|
|
||
|
|
|
||
|
|
def _register(
|
||
|
|
registry: dict[str, type[Any]], names: tuple[str, ...], cls: type[T]
|
||
|
|
) -> type[T]:
|
||
|
|
for name in names:
|
||
|
|
if not name:
|
||
|
|
raise ValueError("Config registry name must be non-empty")
|
||
|
|
existing = registry.get(name)
|
||
|
|
if existing is not None and existing is not cls:
|
||
|
|
raise ValueError(
|
||
|
|
f"Config name {name!r} is already registered by "
|
||
|
|
f"{existing.__module__}.{existing.__name__}"
|
||
|
|
)
|
||
|
|
registry[name] = cls
|
||
|
|
return cls
|
||
|
|
|
||
|
|
|
||
|
|
def register_data_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||
|
|
return lambda cls: _register(_DATA_CONFIG_REGISTRY, names, cls)
|
||
|
|
|
||
|
|
|
||
|
|
def register_model_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||
|
|
return lambda cls: _register(_MODEL_CONFIG_REGISTRY, names, cls)
|
||
|
|
|
||
|
|
|
||
|
|
def register_optimizer_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||
|
|
return lambda cls: _register(_OPTIMIZER_CONFIG_REGISTRY, names, cls)
|
||
|
|
|
||
|
|
|
||
|
|
def register_scheduler_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||
|
|
return lambda cls: _register(_SCHEDULER_CONFIG_REGISTRY, names, cls)
|
||
|
|
|
||
|
|
|
||
|
|
def get_data_config(name: str) -> type[Any] | None:
|
||
|
|
return _DATA_CONFIG_REGISTRY.get(name)
|
||
|
|
|
||
|
|
|
||
|
|
def get_model_config(name: str) -> type[Any] | None:
|
||
|
|
return _MODEL_CONFIG_REGISTRY.get(name)
|
||
|
|
|
||
|
|
|
||
|
|
def get_optimizer_config(name: str) -> type[Any] | None:
|
||
|
|
return _OPTIMIZER_CONFIG_REGISTRY.get(name)
|
||
|
|
|
||
|
|
|
||
|
|
def get_scheduler_config(name: str) -> type[Any] | None:
|
||
|
|
return _SCHEDULER_CONFIG_REGISTRY.get(name)
|
||
|
|
|
||
|
|
|
||
|
|
def registered_data_configs() -> list[str]:
|
||
|
|
return sorted(_DATA_CONFIG_REGISTRY)
|
||
|
|
|
||
|
|
|
||
|
|
def registered_model_configs() -> list[str]:
|
||
|
|
return sorted(_MODEL_CONFIG_REGISTRY)
|
||
|
|
|
||
|
|
|
||
|
|
def registered_optimizer_configs() -> list[str]:
|
||
|
|
return sorted(_OPTIMIZER_CONFIG_REGISTRY)
|
||
|
|
|
||
|
|
|
||
|
|
def registered_scheduler_configs() -> list[str]:
|
||
|
|
return sorted(_SCHEDULER_CONFIG_REGISTRY)
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"get_data_config",
|
||
|
|
"get_model_config",
|
||
|
|
"get_optimizer_config",
|
||
|
|
"get_scheduler_config",
|
||
|
|
"register_data_config",
|
||
|
|
"register_model_config",
|
||
|
|
"register_optimizer_config",
|
||
|
|
"register_scheduler_config",
|
||
|
|
"registered_data_configs",
|
||
|
|
"registered_model_configs",
|
||
|
|
"registered_optimizer_configs",
|
||
|
|
"registered_scheduler_configs",
|
||
|
|
]
|