src/training/trainer_utils.py

"""Shared training utilities for pretraining and fine-tuning.

All training loops import from here. No duplicated optimizer construction,
scheduler logic, or checkpoint I/O across pretrain.py and finetune.py.

Functions:
  create_optimizer     — AdamW, weight decay on 2D+ params only
  create_scheduler     — Linear warmup + cosine decay to min_lr_fraction
  save_checkpoint      — Atomic write (.tmp then rename)
  load_checkpoint      — Restore with optional fingerprint verification
  setup_deterministic  — Seed all RNGs, deterministic algorithms
  nan_guard            — Per-step finiteness check with debug dump

Class:
  MetricsLogger        — Tensorboard + optional wandb
"""

from __future__ import annotations

import logging
import math
import os
import random
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR

log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Optimizer / Scheduler
# ---------------------------------------------------------------------------


def create_optimizer(
    model: nn.Module,
    lr: float = 3e-4,
    betas: tuple[float, float] = (0.9, 0.95),
    weight_decay: float = 0.1,
) -> AdamW:
    """AdamW with weight decay on matrices only (not norms, not biases)."""
    decay, no_decay = [], []
    for param in model.parameters():
        if not param.requires_grad:
            continue
        (decay if param.dim() >= 2 else no_decay).append(param)

    return AdamW(
        [{"params": decay, "weight_decay": weight_decay},
         {"params": no_decay, "weight_decay": 0.0}],
        lr=lr, betas=betas,
    )


def create_scheduler(
    optimizer: AdamW,
    warmup_steps: int,
    total_steps: int,
    min_lr_fraction: float = 0.1,
) -> LambdaLR:
    """Linear warmup then cosine decay to min_lr_fraction * peak LR."""

    def lr_lambda(step: int) -> float:
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return min_lr_fraction + (1 - min_lr_fraction) * 0.5 * (1 + math.cos(math.pi * progress))

    return LambdaLR(optimizer, lr_lambda)


# ---------------------------------------------------------------------------
# Checkpointing
# ---------------------------------------------------------------------------


def save_checkpoint(
    path: str | Path,
    model: nn.Module,
    optimizer: AdamW,
    scheduler: LambdaLR,
    step: int,
    config: dict[str, Any],
    fingerprint: str = "",
) -> None:
    """Atomic save: writes to .tmp then os.replace for crash safety."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")

    torch.save({
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict(),
        "step": step,
        "config": config,
        "fingerprint": fingerprint,
    }, tmp)
    os.replace(tmp, path)
    log.info("Checkpoint saved: %s (step %d)", path, step)


def load_checkpoint(
    path: str | Path,
    model: nn.Module,
    optimizer: AdamW | None = None,
    scheduler: LambdaLR | None = None,
    expected_fingerprint: str | None = None,
    strict: bool = True,
) -> dict[str, Any]:
    """Load checkpoint. Pass optimizer/scheduler=None to skip restoring them.

    strict=False allows loading pretrained weights into a model with extra
    parameters (e.g. a fraud head added for fine-tuning).
    Raises ValueError on fingerprint mismatch.
    """
    ckpt = torch.load(path, map_location="cpu", weights_only=False)

    if expected_fingerprint is not None:
        saved = ckpt.get("fingerprint", "")
        if saved != expected_fingerprint:
            raise ValueError(
                f"Fingerprint mismatch: checkpoint='{saved}', expected='{expected_fingerprint}'"
            )

    model.load_state_dict(ckpt["model_state_dict"], strict=strict)
    if optimizer is not None and "optimizer_state_dict" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    if scheduler is not None and "scheduler_state_dict" in ckpt:
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])

    log.info("Checkpoint loaded: %s (step %s)", path, ckpt.get("step", "?"))
    return ckpt


# ---------------------------------------------------------------------------
# Reproducibility
# ---------------------------------------------------------------------------


def setup_deterministic(
    seed: int = 42,
    warn_only: bool = True,
    cublas_workspace_config: str = ":4096:8",
) -> None:
    """Seed all RNGs and enable deterministic CUDA kernels.

    warn_only=True because depthwise Conv1d backward may lack a deterministic
    CUDA kernel — logs a warning instead of crashing (D19e).
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    torch.use_deterministic_algorithms(True, warn_only=warn_only)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = cublas_workspace_config

    log.info("Deterministic mode: seed=%d, warn_only=%s", seed, warn_only)


# ---------------------------------------------------------------------------
# NaN guard
# ---------------------------------------------------------------------------


class NaNError(RuntimeError):
    """Raised when loss contains NaN or Inf."""
    pass


def nan_guard(
    loss: torch.Tensor,
    step: int,
    output_dir: str | Path,
    model: nn.Module | None = None,
) -> None:
    """Check loss is finite. Dumps debug state and raises NaNError on failure."""
    if torch.isfinite(loss):
        return

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    debug_path = output_dir / f"nan_debug_step{step}.pt"

    state: dict[str, Any] = {"step": step, "loss": loss.detach().cpu()}
    if model is not None:
        state["model_state_dict"] = {
            k: v.detach().cpu() for k, v in model.state_dict().items()
        }

    torch.save(state, debug_path)
    log.error("NaN/Inf loss at step %d. Debug state: %s", step, debug_path)
    raise NaNError(f"Loss is {loss.item()} at step {step}. Debug: {debug_path}")


# ---------------------------------------------------------------------------
# Metrics logging
# ---------------------------------------------------------------------------


class MetricsLogger:
    """Tensorboard writer + optional wandb. Wandb activates on WANDB_API_KEY."""

    def __init__(
        self,
        log_dir: str | Path,
        experiment_name: str = "run",
        use_wandb: str = "auto",
        wandb_project: str = "lfm2-transactions",
    ) -> None:
        from torch.utils.tensorboard import SummaryWriter

        self.tb = SummaryWriter(log_dir=str(log_dir))
        self.wandb_run = None

        if use_wandb == "auto" and os.environ.get("WANDB_API_KEY"):
            try:
                import wandb

                self.wandb_run = wandb.init(
                    project=wandb_project, name=experiment_name, dir=str(log_dir),
                )
            except ImportError:
                log.info("wandb not installed, tensorboard only")

    def log_scalar(self, tag: str, value: float, step: int) -> None:
        self.tb.add_scalar(tag, value, step)
        if self.wandb_run is not None:
            import wandb

            wandb.log({tag: value}, step=step)

    def log_per_feature_losses(
        self,
        total_loss: float,
        per_feature_losses: dict[str, float],
        step: int,
    ) -> None:
        self.log_scalar("loss/total", total_loss, step)
        for name, val in per_feature_losses.items():
            self.log_scalar(f"loss/{name}", val, step)

    def log_config(self, config: dict[str, Any]) -> None:
        self.tb.add_text("config", str(config), 0)
        if self.wandb_run is not None:
            import wandb

            wandb.config.update(config)

    def close(self) -> None:
        self.tb.close()
        if self.wandb_run is not None:
            import wandb

            wandb.finish()