La receta de entrenamiento

"""Tests for rukh.train.schedule: linear warmup then cosine decay to a floor."""

from __future__ import annotations

import pytest

from rukh.train import TrainConfig, lr_at

pytestmark = pytest.mark.unit

CFG = TrainConfig(lr=6e-4, min_lr_ratio=0.1, warmup=100, max_steps=1000)


def test_warmup_is_linear_and_reaches_the_peak() -> None:
    assert lr_at(0, CFG) == 0.0
    assert lr_at(50, CFG) == pytest.approx(CFG.lr / 2)
    assert lr_at(100, CFG) == pytest.approx(CFG.lr)
    rising = [lr_at(step, CFG) for step in range(101)]
    assert rising == sorted(rising)
    deltas = {round(b - a, 12) for a, b in zip(rising, rising[1:], strict=False)}
    assert len(deltas) == 1

def test_cosine_decays_to_the_floor_and_stays_there() -> None:
    floor = CFG.lr * CFG.min_lr_ratio
    assert lr_at(1000, CFG) == pytest.approx(floor)
    assert lr_at(5000, CFG) == pytest.approx(floor)
    falling = [lr_at(step, CFG) for step in range(100, 1001, 10)]
    assert falling == sorted(falling, reverse=True)
    assert all(value >= floor - 1e-12 for value in falling)
    halfway = lr_at(550, CFG)
    assert halfway == pytest.approx(floor + 0.5 * (CFG.lr - floor))


def test_without_warmup_the_first_step_is_the_peak() -> None:
    cfg = CFG.model_copy(update={"warmup": 0})
    assert lr_at(0, cfg) == pytest.approx(cfg.lr)


def test_degenerate_schedule_returns_the_floor() -> None:
    cfg = CFG.model_copy(update={"warmup": 1000, "max_steps": 1000})
    assert lr_at(1000, cfg) == pytest.approx(cfg.lr * cfg.min_lr_ratio)
    with pytest.raises(ValueError, match="non-negative"):
        lr_at(-1, CFG)


def test_config_rejects_unknown_keys() -> None:
    with pytest.raises(ValueError):
        TrainConfig(learning_rate=1e-3)  # type: ignore[call-arg]

from __future__ import annotations

import logging
import math
import os
import time
from collections.abc import Iterator
from contextlib import nullcontext
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Literal

import torch
from torch import nn
from torch.utils.data import DataLoader

from rukh import paths
from rukh.config import BaseConfig
from rukh.models import DecoderConfig, MoveDecoder, preset
from rukh.tokenize.loader import IGNORE_INDEX, PackedDataset, make_loader
from rukh.train.checkpoint import (
    BEST_NAME,
    load_checkpoint,
    read_manifest_sha,
    read_vocab_hash,
    restore,
    save_checkpoint,
    step_name,
)
from rukh.train.schedule import lr_at

log = logging.getLogger(__name__)

Batch = tuple[torch.Tensor, torch.Tensor]

def maybe_compile(
    model: MoveDecoder, enabled: bool, sample: torch.Tensor | None = None
) -> nn.Module:
    """``torch.compile`` the model when asked; a failure is a warning, never a stopped run.

    ``torch.compile`` is lazy: on Windows without MSVC it only fails when the first forward
    reaches Inductor. ``sample`` (one batch of the training shape) forces that compilation here,
    where it can still fall back to eager, and Dynamo's own error suppression covers any later
    recompilation for a different shape.
    """
    if not enabled:
        return model
    try:
        import torch._dynamo as dynamo

        dynamo.config.suppress_errors = True
        compiled = torch.compile(model)
        if sample is not None:
            _, loss = compiled(sample, sample)
            if loss is not None:
                loss.backward()
            model.zero_grad(set_to_none=True)
        return compiled
    except Exception as exc:  # noqa: BLE001 - compilation backends fail in many ways
        log.warning("torch.compile is unavailable, training eagerly: %s", exc)
        model.zero_grad(set_to_none=True)
        return model

def run_dir(cfg: TrainConfig, resume: Path | None = None) -> Path:
    """Where this run writes: the resumed run's folder, or ``out_dir/<run name>``.

    The name carries a timestamp unless ``unique_run_name`` is off, because two runs of the same
    config would otherwise write the same ``step-*.pt`` files and the second would quietly
    overwrite the checkpoint series of the first.
    """
    if resume is not None:
        return Path(resume).resolve().parent
    name = cfg.run_name or cfg.preset
    if cfg.unique_run_name:
        name = f"{name}-{datetime.now(UTC):%Y%m%d-%H%M%S}"
    return paths.resolve(cfg.out_dir) / name


def skip_batches(batches: Iterator[Batch], count: int) -> int:
    """Wind the batch stream forward ``count`` batches and return how many were skipped.

    A resumed run must not start again at the first window of the first epoch: it would train
    twice on the same games while the schedule believes it is halfway. Windows are memmap slices,
    so winding forward is cheap compared with a step, and it is logged because it is not free.
    """
    if count <= 0:
        return 0
    started = time.perf_counter()
    for _ in range(count):
        next(batches)
    log.info("skipped %d batches in %.1f s to resume", count, time.perf_counter() - started)
    return count

    use_bf16 = cfg.precision == "bf16" and where.type == "cuda"
    autocast = (
        torch.autocast(device_type="cuda", dtype=torch.bfloat16) if use_bf16 else nullcontext()
    )
    warmup = torch.ones((cfg.batch_size, cfg.block), dtype=torch.long, device=where)
    with autocast:  # compile under the same precision the loop will use
        runnable = maybe_compile(model, cfg.compile, warmup)
    runnable.train()

    out_dir = run_dir(cfg, resume)
    out_dir.mkdir(parents=True, exist_ok=True)
    vocab_hash = read_vocab_hash(tokens_dir / "train")
    manifest_sha = read_manifest_sha(paths.data_dir() / "raw" / "manifest.json")
    tokens_per_step = cfg.batch_size * cfg.grad_accum * cfg.block
    batches = forever(train_loader)
    skip_batches(batches, start_step * cfg.grad_accum)
    final = out_dir / step_name(cfg.max_steps)

    from rukh.tracking import git_sha, start_run

    params = {
        **cfg.model_dump(mode="json"),
        "vocab_hash": vocab_hash,
        "data_manifest_sha": manifest_sha,
        "device": str(where),
        "num_params": model.num_params(),
    }
    with start_run(out_dir.name, params, tags={"preset": cfg.preset}, run_id=run_id) as run:
        log.info("run %s in %s on %s", run.info.run_id, out_dir, where)
        this_run = str(run.info.run_id)

        def save(path: Path, step: int) -> Path:
            return save_checkpoint(
                path,
                step=step,
                model=model,
                optimizer=optimizer,
                cfg=cfg.model_dump(mode="json"),
                model_cfg=model_cfg.model_dump(mode="json"),
                vocab_hash=vocab_hash,
                data_manifest_sha=manifest_sha,
                git_sha=git_sha(),
                best_val=None if math.isinf(best_val) else best_val,
                run_id=this_run,
            )

def log_metrics(metrics: dict[str, float], step: int) -> None:
    """Send metrics to the active MLflow run; a tracking failure never stops training."""
    import mlflow

    clean = {key: value for key, value in metrics.items() if not math.isnan(value)}
    if not clean:
        return
    try:
        mlflow.log_metrics(clean, step=step)
    except Exception as exc:  # noqa: BLE001 - tracking is not worth a lost run
        log.warning("could not log metrics at step %d: %s", step, exc)

"""Checkpoints: what a run writes so it can be resumed, evaluated and published.

A payload holds only plain Python values and tensors, so it can be read back with
``torch.load(..., weights_only=True)``. Provenance (``vocab_hash``, ``data_manifest_sha``,
``git_sha``) is recorded on purpose: weights whose data or vocabulary cannot be identified are
not reproducible.
"""

from __future__ import annotations

import hashlib
import json
import random
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import nn

from rukh.models import DecoderConfig, MoveDecoder

BEST_NAME = "best.pt"
TIED_HEAD = "lm_head.weight"
"""Tied to ``tokens.weight``; a published state dict leaves it out and it is re-tied on load."""

def step_name(step: int) -> str:
    """File name of the checkpoint written after ``step`` optimizer steps."""
    return f"step-{step}.pt"


def read_vocab_hash(tokens_dir: Path) -> str | None:
    """``vocab_hash`` from a pack ``meta.json``, or None when it cannot be read."""
    try:
        meta = json.loads((Path(tokens_dir) / "meta.json").read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    value = meta.get("vocab_hash")
    return str(value) if isinstance(value, str) else None


def read_manifest_sha(manifest: Path) -> str | None:
    """SHA-256 of a data manifest, or None when the file is absent or unreadable.

    Training must never depend on ``data/`` being present: a checkpoint trained from a pack
    copied elsewhere simply records ``None``.
    """
    try:
        payload = Path(manifest).read_bytes()
    except OSError:
        return None
    return hashlib.sha256(payload).hexdigest()

def rng_state() -> dict[str, Any]:
    """Snapshot of the Python, NumPy and torch CPU generators."""
    name, keys, pos, has_gauss, cached = np.random.get_state(legacy=True)
    return {
        "python": json.dumps(random.getstate()),
        "numpy": json.dumps(
            [name, np.asarray(keys).tolist(), int(pos), int(has_gauss), float(cached)]
        ),
        "torch": torch.get_rng_state(),
    }


def set_rng_state(state: Mapping[str, Any]) -> None:
    """Restore a snapshot taken by ``rng_state`` (missing entries are skipped)."""
    python = state.get("python")
    if isinstance(python, str):
        version, keys, gauss = json.loads(python)
        random.setstate((version, tuple(keys), gauss))
    numpy = state.get("numpy")
    if isinstance(numpy, str):
        name, keys, pos, has_gauss, cached = json.loads(numpy)
        np.random.set_state((name, np.array(keys, dtype=np.uint32), pos, has_gauss, cached))
    torch_state = state.get("torch")
    if isinstance(torch_state, torch.Tensor):
        torch.set_rng_state(torch_state.to(torch.uint8).cpu())

def save_checkpoint(
    path: Path,
    *,
    step: int,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None,
    cfg: Mapping[str, Any],
    model_cfg: Mapping[str, Any],
    vocab_hash: str | None = None,
    data_manifest_sha: str | None = None,
    git_sha: str | None = None,
    best_val: float | None = None,
    run_id: str | None = None,
) -> Path:
    """Write one checkpoint atomically (temporary file plus replace) and return its path.

    ``run_id`` is the MLflow run that wrote it, so ``--resume`` can carry on logging into the
    same run and the loss curve stays one line instead of two.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    payload: dict[str, Any] = {
        "step": int(step),
        "model_state": {k: v.detach().cpu() for k, v in model.state_dict().items()},
        "opt_state": optimizer.state_dict() if optimizer is not None else None,
        "cfg": dict(cfg),
        "model_cfg": dict(model_cfg),
        "vocab_hash": vocab_hash,
        "data_manifest_sha": data_manifest_sha,
        "git_sha": git_sha,
        "best_val": best_val,
        "run_id": run_id,
        "rng": rng_state(),
    }
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, tmp)
    tmp.replace(path)
    return path

def load_checkpoint(path: Path, map_location: str | torch.device = "cpu") -> dict[str, Any]:
    """Read a checkpoint written by ``save_checkpoint``."""
    payload = torch.load(Path(path), map_location=map_location, weights_only=True)
    if not isinstance(payload, dict) or "model_state" not in payload:
        raise ValueError(f"{path} is not a rukh checkpoint")
    return payload


def load_state(model: nn.Module, state: Mapping[str, Any]) -> None:
    """Load weights, accepting a state dict whose tied head was left out when it was saved.

    ``MoveDecoder`` ties ``lm_head.weight`` to ``tokens.weight`` in its constructor, so the tied
    head is already correct once the embedding is loaded: a state dict without it (the one
    ``rukh publish`` writes, because two names for one tensor is what makes ``safetensors``
    refuse the file) loads cleanly and nothing else may be missing.
    """
    missing, unexpected = model.load_state_dict(dict(state), strict=False)
    absent = [name for name in missing if name != TIED_HEAD]
    if absent or unexpected:
        raise ValueError(
            f"the weights do not match the model: missing {absent}, unexpected {list(unexpected)}"
        )

def load_model(
    path: Path, map_location: str | torch.device = "cpu"
) -> tuple[MoveDecoder, dict[str, Any]]:
    """Rebuild the decoder a checkpoint describes, in eval mode, plus the whole payload."""
    payload = load_checkpoint(path, map_location=map_location)
    model = MoveDecoder(DecoderConfig.model_validate(payload["model_cfg"]))
    load_state(model, payload["model_state"])
    return model.eval(), payload


def restore(
    payload: Mapping[str, Any],
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
    with_rng: bool = True,
) -> int:
    """Load weights (and optionally the optimizer and RNG) and return the step reached."""
    load_state(model, payload["model_state"])
    opt_state = payload.get("opt_state")
    if optimizer is not None and opt_state is not None:
        optimizer.load_state_dict(opt_state)
    rng = payload.get("rng")
    if with_rng and isinstance(rng, Mapping):
        set_rng_state(rng)
    return int(payload["step"])

"""Training: the loop, the learning-rate schedule and checkpoint handling."""

from rukh.train.checkpoint import (
    BEST_NAME,
    TIED_HEAD,
    load_checkpoint,
    load_model,
    load_state,
    read_manifest_sha,
    read_vocab_hash,
    restore,
    save_checkpoint,
    step_name,
)
from rukh.train.loop import (
    TrainConfig,
    evaluate,
    param_groups,
    pick_device,
    run_dir,
    skip_batches,
    train,
)
from rukh.train.schedule import lr_at

__all__ = [
    "BEST_NAME",
    "TIED_HEAD",
    "TrainConfig",
    "evaluate",
    "load_checkpoint",
    "load_model",
    "load_state",
    "lr_at",
    "param_groups",
    "pick_device",
    "read_manifest_sha",
    "read_vocab_hash",
    "restore",
    "run_dir",
    "save_checkpoint",
    "skip_batches",
    "step_name",
    "train",
]

"""Tests for rukh.train.checkpoint: round-trip, provenance and RNG restoration."""

from __future__ import annotations

import json
import random
from pathlib import Path

import numpy as np
import pytest
import torch

from rukh.models import DecoderConfig, MoveDecoder
from rukh.train import (
    TrainConfig,
    load_checkpoint,
    read_manifest_sha,
    read_vocab_hash,
    restore,
    save_checkpoint,
    step_name,
)

pytestmark = pytest.mark.unit

TOY = DecoderConfig(vocab_size=32, n_layer=2, n_head=2, d_model=32, block=16)


def make_pair() -> tuple[MoveDecoder, torch.optim.Optimizer]:
    torch.manual_seed(0)
    model = MoveDecoder(TOY)
    return model, torch.optim.AdamW(model.parameters(), lr=1e-3)

def test_save_and_load_reproduce_weights_and_step(tmp_path: Path) -> None:
    model, optimizer = make_pair()
    idx = torch.randint(1, TOY.vocab_size, (2, 8))
    _, loss = model(idx, idx)
    assert loss is not None
    loss.backward()
    optimizer.step()

    path = save_checkpoint(
        tmp_path / step_name(7),
        step=7,
        model=model,
        optimizer=optimizer,
        cfg=TrainConfig(max_steps=7).model_dump(mode="json"),
        model_cfg=TOY.model_dump(mode="json"),
        vocab_hash="abc",
        data_manifest_sha=None,
        git_sha="deadbeef",
        best_val=1.25,
    )
    assert path.name == "step-7.pt"
    assert not list(tmp_path.glob("*.tmp"))

    payload = load_checkpoint(path)
    assert payload["step"] == 7
    assert payload["vocab_hash"] == "abc"
    assert payload["data_manifest_sha"] is None
    assert payload["git_sha"] == "deadbeef"
    assert payload["best_val"] == pytest.approx(1.25)
    assert payload["cfg"]["max_steps"] == 7
    assert payload["model_cfg"]["d_model"] == 32

    fresh = MoveDecoder(TOY)
    fresh_opt = torch.optim.AdamW(fresh.parameters(), lr=1e-3)
    assert restore(payload, fresh, fresh_opt) == 7
    for a, b in zip(model.state_dict().values(), fresh.state_dict().values(), strict=True):
        assert torch.equal(a, b)
    with torch.no_grad():
        assert torch.allclose(model.eval()(idx)[0], fresh.eval()(idx)[0], atol=1e-6)

def test_restoring_the_rng_repeats_the_same_draws(tmp_path: Path) -> None:
    model, optimizer = make_pair()
    random.seed(11)
    np.random.seed(11)
    torch.manual_seed(11)
    path = save_checkpoint(
        tmp_path / "rng.pt",
        step=0,
        model=model,
        optimizer=optimizer,
        cfg={},
        model_cfg=TOY.model_dump(mode="json"),
    )
    expected = (random.random(), float(np.random.random()), torch.rand(3).tolist())

    random.seed(99)
    np.random.seed(99)
    torch.manual_seed(99)
    restore(load_checkpoint(path), MoveDecoder(TOY))
    assert (random.random(), float(np.random.random()), torch.rand(3).tolist()) == expected

def test_vocab_hash_and_manifest_sha_are_defensive(tmp_path: Path) -> None:
    assert read_vocab_hash(tmp_path) is None
    (tmp_path / "meta.json").write_text(json.dumps({"n_games": 1}), encoding="utf-8")
    assert read_vocab_hash(tmp_path) is None
    (tmp_path / "meta.json").write_text(json.dumps({"vocab_hash": "cafe"}), encoding="utf-8")
    assert read_vocab_hash(tmp_path) == "cafe"

    assert read_manifest_sha(tmp_path / "does-not-exist" / "manifest.json") is None
    manifest = tmp_path / "manifest.json"
    manifest.write_bytes(b"{}")
    sha = read_manifest_sha(manifest)
    assert sha is not None and len(sha) == 64


def test_a_foreign_file_is_not_a_checkpoint(tmp_path: Path) -> None:
    path = tmp_path / "other.pt"
    torch.save({"weights": torch.zeros(2)}, path)
    with pytest.raises(ValueError, match="not a rukh checkpoint"):
        load_checkpoint(path)

def start_run(
    name: str,
    config: Mapping[str, Any],
    tags: Mapping[str, str] | None = None,
    run_id: str | None = None,
) -> Iterator[Any]:
    """Start an MLflow run in the local store with ``config`` logged as flattened params.

    Tags always include ``rukh_version`` and, when the project lives in a git repo, ``git_sha``.
    With ``run_id`` the existing run is reopened instead of a new one being created, so a
    training run that is resumed keeps one curve rather than starting a second one; a run id
    that cannot be reopened (deleted, or from another store) falls back to a fresh run.
    Yields the active ``mlflow.ActiveRun``.
    """

    started = None
    if run_id:
        try:
            started = mlflow.start_run(run_id=run_id, tags=run_tags)
        except Exception as exc:  # noqa: BLE001 - a lost run must never stop a training run
            log.warning("could not reopen the MLflow run %s (%s); starting a new one", run_id, exc)
    with started or mlflow.start_run(run_name=name, tags=run_tags) as run:
        params = flatten(config)
        if params:
            # A resumed run already carries these; re-logging a changed value is an error there.
            try:
                mlflow.log_params(params)
            except Exception as exc:  # noqa: BLE001 - the params are already recorded
                log.warning("could not log the run parameters: %s", exc)
        yield run

@app.command("train")
def train_cmd(
    config: Annotated[
        Path,
        typer.Option(
            "--config", exists=True, dir_okay=False, readable=True, help="Training YAML config."
        ),
    ],
    model_preset: Annotated[
        str | None, typer.Option("--preset", help="Override the preset: tiny, small or medium.")
    ] = None,
    resume: Annotated[
        Path | None,
        typer.Option(
            "--resume", exists=True, dir_okay=False, readable=True, help="Checkpoint to continue."
        ),
    ] = None,
    max_steps: Annotated[
        int | None, typer.Option("--max-steps", help="Override max_steps from the config.")
    ] = None,
) -> None:
    """Train a MoveDecoder from a packed token stream, logging the run to MLflow."""
    from rukh.config import load_yaml
    from rukh.models import PRESETS
    from rukh.train import TrainConfig, train

    cfg = load_yaml(config, TrainConfig)
    if model_preset is not None:
        if model_preset not in PRESETS:
            typer.echo(f"error: --preset must be one of {', '.join(PRESETS)}", err=True)
            raise typer.Exit(code=2)
        cfg = cfg.model_copy(update={"preset": model_preset, "model": None})
    if max_steps is not None:
        cfg = cfg.model_copy(update={"max_steps": max_steps})
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    try:
        checkpoint = train(cfg, resume=resume)
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    typer.echo(f"preset:     {cfg.preset}")
    typer.echo(f"steps:      {cfg.max_steps}")
    typer.echo(f"checkpoint: {checkpoint}")