enlace/enlace/config/load.py

"""Composición y carga de configs.

Un YAML raíz declara qué capa usar de cada familia; se fusionan con OmegaConf y
el resultado se valida con pydantic. Los overrides de línea de comandos existen
para probar variantes sin editar archivos, no para esconder configuración.

    python -m enlace.train.train configs/runs/smoke.yaml hardware.compile=false
"""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Any

from omegaconf import DictConfig, OmegaConf

from enlace.config.schema import AgentConfig, Config, DistillConfig

# Familias de capas que puede declarar un YAML raíz, en el orden en que se
# fusionan. El orden solo importa para los mensajes de error.
LAYERS = ("hardware", "model", "train", "data")


class ConfigError(RuntimeError):
    """Error de configuración legible, sin traceback de pydantic encima."""


def _resolve(path: str | Path, root: Path) -> Path:
    p = Path(path)
    full = p if p.is_absolute() else root / p
    if not full.is_file():
        raise ConfigError(f"no existe el archivo de config: {full}")
    return full


def _load_yaml(path: Path) -> DictConfig:
    cfg = OmegaConf.load(path)
    if not isinstance(cfg, DictConfig):
        raise ConfigError(f"{path}: el YAML raíz debe ser un mapeo, no una lista.")
    return cfg


def compose(path: str | Path, overrides: list[str] | None = None) -> dict[str, Any]:
    """Compone un YAML raíz en un dict plano, sin validar todavía.

    El YAML raíz tiene la forma:

        include:
          hardware: hardware/turing-2060.yaml
          model: model/tiny-50m.yaml
          train: train/pretrain.yaml
          data: data/corpus.yaml

        # opcional: ajustes puntuales encima de las capas incluidas
        train:
          run_name: mi-corrida
    """
    path = Path(path).resolve()
    if not path.is_file():
        raise ConfigError(f"no existe el archivo de config: {path}")

    raw = _load_yaml(path)
    # Las rutas de `include` se resuelven contra el directorio configs/, que es
    # el padre del directorio del YAML raíz (configs/runs/foo.yaml -> configs/).
    configs_root = path.parent.parent if path.parent.name == "runs" else path.parent

    includes = raw.pop("include", None)
    merged = OmegaConf.create({})
    if includes is not None:
        for layer in LAYERS:
            ref = includes.get(layer)
            if ref is None:
                continue
            merged[layer] = _load_yaml(_resolve(ref, configs_root))
        unknown = set(includes.keys()) - set(LAYERS)
        if unknown:
            raise ConfigError(
                f"{path}: capas desconocidas en `include`: {sorted(unknown)}. "
                f"Válidas: {list(LAYERS)}."
            )

    # Lo que quede en el YAML raíz pisa a las capas incluidas.
    merged = OmegaConf.merge(merged, raw)

    if overrides:
        merged = OmegaConf.merge(merged, OmegaConf.from_dotlist(list(overrides)))

    resolved = OmegaConf.to_container(merged, resolve=True)
    assert isinstance(resolved, dict)
    return resolved


def load_config(path: str | Path, overrides: list[str] | None = None) -> Config:
    """Compone, valida y devuelve la config. Falla temprano y con claridad."""
    data = compose(path, overrides)
    try:
        return Config.model_validate(data)
    except Exception as exc:  # pydantic.ValidationError y los ValueError propios
        raise ConfigError(f"config inválida ({path}):\n{exc}") from None


CONFIGS_ROOT = Path(__file__).resolve().parents[2] / "configs"


def _load_single(path: Path, overrides: list[str] | None, model_cls, etiqueta: str):
    """Carga un YAML suelto y lo valida. No compone capas.

    Lo usan los árboles de config que no son de entrenamiento (agente,
    destilación): no necesitan saber nada del optimizador ni del schedule.
    Los valores que salen de `.env` se resuelven con `${oc.env:...}`, así que
    los secretos nunca entran en un archivo commiteado.
    """
    if not path.is_file():
        raise ConfigError(f"no existe el archivo de config: {path}")

    merged = _load_yaml(path)
    if overrides:
        merged = OmegaConf.merge(merged, OmegaConf.from_dotlist(list(overrides)))
    data = OmegaConf.to_container(merged, resolve=True)
    assert isinstance(data, dict)

    # OmegaConf devuelve cadena vacía cuando la variable de entorno no está
    # definida; para pydantic eso tiene que ser ausencia, no un valor vacío.
    _vaciar_a_none(data)

    try:
        return model_cls.model_validate(data)
    except Exception as exc:
        raise ConfigError(f"config de {etiqueta} inválida ({path}):\n{exc}") from None


def _vaciar_a_none(data: dict) -> None:
    """Convierte las cadenas vacías de `${oc.env:VAR,""}` en ausencia, recursivo."""
    for clave, valor in data.items():
        if isinstance(valor, dict):
            _vaciar_a_none(valor)
        elif valor == "":
            data[clave] = None


def load_agent_config(
    path: str | Path | None = None, overrides: list[str] | None = None
) -> AgentConfig:
    """Carga la config del runtime del agente."""
    path = Path(path) if path is not None else CONFIGS_ROOT / "agent" / "tools.yaml"
    return _load_single(path, overrides, AgentConfig, "agente")


def load_distill_config(
    path: str | Path | None = None, overrides: list[str] | None = None
) -> DistillConfig:
    """Carga la config de generación de datos sintéticos por lotes."""
    path = Path(path) if path is not None else CONFIGS_ROOT / "data" / "distill.yaml"
    return _load_single(path, overrides, DistillConfig, "destilación")


def load_config_from_argv(argv: list[str] | None = None) -> Config:
    """`prog config.yaml [clave.sub=valor ...]`, para los entrypoints."""
    args = list(sys.argv[1:] if argv is None else argv)
    if not args:
        raise ConfigError("uso: <programa> <config.yaml> [clave.sub=valor ...]")
    return load_config(args[0], args[1:])


def to_yaml(config: Config) -> str:
    """Serializa la config resuelta, para guardarla junto al checkpoint.

    Un checkpoint sin su config no es reproducible: esto es lo que se escribe
    en el snapshot.
    """
    return OmegaConf.to_yaml(OmegaConf.create(config.model_dump(mode="json")))