src/easytranslate/training/trainer.py

"""
训练器模块 — Person C 负责实现

这是训练的核心控制器，负责:
1. 训练循环 (training loop)
2. 验证循环 (validation loop)
3. 检查点保存/加载
4. 日志记录
5. 分布式训练支持
6. 混合精度训练
7. 梯度累积
8. 早停
"""

from __future__ import annotations

import json
import logging
import os
import shutil
import time
from pathlib import Path
from typing import Optional

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

from easytranslate.training.loss import LabelSmoothedCrossEntropyLoss
from easytranslate.training.optimizer import build_optimizer, build_scheduler

logger = logging.getLogger(__name__)


class Trainer:
    """
    翻译模型训练器。

    使用方法:
        trainer = Trainer(model, train_loader, val_loader, config)
        trainer.train()
    """

    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        config: dict,
        optimizer=None,
        scheduler=None,
        criterion=None,
        evaluator=None,
    ):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config

        train_cfg = config.get("training", {})

        self.device = self._resolve_device(train_cfg)
        self.model = self.model.to(self.device)

        self.fp16 = train_cfg.get("fp16", False)
        self.bf16 = train_cfg.get("bf16", False)
        self.gradient_accumulation_steps = int(train_cfg.get("gradient_accumulation_steps", 1))
        self.max_grad_norm = float(train_cfg.get("max_grad_norm", 1.0))
        self.gradient_checkpointing = train_cfg.get("gradient_checkpointing", False)

        if self.gradient_checkpointing and hasattr(self.model, "gradient_checkpointing_enable"):
            self.model.gradient_checkpointing_enable()

        self.num_epochs = int(train_cfg.get("epochs", 30))
        self.max_steps = int(train_cfg.get("max_steps", -1))

        self.scaler = None
        self.amp_dtype = None
        if self.fp16:
            self.scaler = torch.amp.GradScaler(self.device.type)
            self.amp_dtype = torch.float16
        elif self.bf16:
            self.amp_dtype = torch.bfloat16

        if optimizer is None:
            optimizer = build_optimizer(model, train_cfg)
        self.optimizer = optimizer

        total_steps = self.num_epochs * len(train_loader) // self.gradient_accumulation_steps
        if scheduler is None:
            scheduler = build_scheduler(optimizer, train_cfg, num_training_steps=total_steps)
        self.scheduler = scheduler

        reg_cfg = train_cfg.get("regularization", {})
        if criterion is None:
            criterion = LabelSmoothedCrossEntropyLoss(
                smoothing=float(reg_cfg.get("label_smoothing", 0.1)),
                pad_id=-100,  # must match TranslationCollator's label_pad_token_id
            )
        self.criterion = criterion

        self.evaluator = evaluator

        ckpt_cfg = train_cfg.get("checkpoint", {})
        self.checkpoint_dir = Path(ckpt_cfg.get("save_dir", "checkpoints/"))
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.save_every_n_steps = int(ckpt_cfg.get("save_every_n_steps", 5000))
        self.save_best = ckpt_cfg.get("save_best", True)
        self.metric_for_best = ckpt_cfg.get("metric_for_best", "bleu")
        self.max_checkpoints = int(ckpt_cfg.get("max_checkpoints", 5))

        es_cfg = train_cfg.get("early_stopping", {})
        self.early_stopping_enabled = es_cfg.get("enabled", True)
        self.patience = int(es_cfg.get("patience", 5))
        self.min_delta = float(es_cfg.get("min_delta", 0.1))

        log_cfg = config.get("logging", {})
        self.log_dir = Path(log_cfg.get("log_dir", "logs/"))
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.log_every_n_steps = int(log_cfg.get("log_every_n_steps", 100))
        self.log_backend = log_cfg.get("backend", "tensorboard")

        self._init_logger()

        self.global_step = 0
        self.current_epoch = 0
        self.best_metric_value = float("-inf")
        self.best_epoch = 0
        self.epochs_without_improvement = 0
        self.train_loss_history: list[float] = []
        self.val_metrics_history: list[dict] = []

        self._setup_distributed()

    def _resolve_device(self, train_cfg: dict) -> torch.device:
        device_str = train_cfg.get("device", "auto")
        if device_str == "auto":
            if torch.cuda.is_available():
                return torch.device("cuda")
            elif torch.backends.mps.is_available():
                return torch.device("mps")
            else:
                return torch.device("cpu")
        return torch.device(device_str)

    def _init_logger(self):
        self.writer = None
        if self.log_backend in ("tensorboard", "both"):
            try:
                from torch.utils.tensorboard import SummaryWriter
                self.writer = SummaryWriter(log_dir=str(self.log_dir))
            except ImportError:
                logger.warning("TensorBoard not available, skipping")

        self.wandb_run = None
        if self.log_backend in ("wandb", "both"):
            try:
                import wandb
                log_cfg = self.config.get("logging", {})
                self.wandb_run = wandb.init(
                    project=log_cfg.get("project_name", "EasyTranslate"),
                    config=self.config,
                    dir=str(self.log_dir),
                )
            except ImportError:
                logger.warning("WandB not available, skipping")

    def _setup_distributed(self):
        dist_cfg = self.config.get("training", {}).get("distributed", {})
        strategy = dist_cfg.get("strategy", "ddp")

        if strategy == "ddp" and torch.distributed.is_available() and torch.distributed.is_initialized():
            self.model = nn.parallel.DistributedDataParallel(self.model)
            self.is_distributed = True
        elif strategy == "deepspeed":
            try:
                import deepspeed
                ds_config = dist_cfg.get("deepspeed_config", "configs/deepspeed_config.json")
                self.model, self.optimizer, _, _ = deepspeed.initialize(
                    model=self.model,
                    optimizer=self.optimizer,
                    config_params=ds_config,
                )
                self.is_distributed = True
            except ImportError:
                logger.warning("DeepSpeed not available, falling back to single GPU")
                self.is_distributed = False
        else:
            self.is_distributed = False

    def train(self):
        logger.info("Starting training for %d epochs", self.num_epochs)
        logger.info("Device: %s, FP16: %s, BF16: %s", self.device, self.fp16, self.bf16)
        logger.info("Gradient accumulation steps: %d", self.gradient_accumulation_steps)

        for epoch in range(self.current_epoch, self.num_epochs):
            self.current_epoch = epoch
            logger.info("=" * 50)
            logger.info("Epoch %d/%d", epoch + 1, self.num_epochs)

            train_loss = self._train_one_epoch(epoch)
            self.train_loss_history.append(train_loss)

            val_metrics = self._validate(epoch)
            self.val_metrics_history.append(val_metrics)

            self._log_metrics(epoch, train_loss, val_metrics)
            self._save_checkpoint(epoch, val_metrics)

            if self._should_early_stop(val_metrics):
                logger.info("Early stopping triggered at epoch %d", epoch + 1)
                break

            if self.max_steps > 0 and self.global_step >= self.max_steps:
                logger.info("Reached max steps %d, stopping", self.max_steps)
                break

        self._log_final_results()
        self._cleanup()

    def _train_one_epoch(self, epoch: int) -> float:
        self.model.train()
        total_loss = 0.0
        num_batches = 0
        self.optimizer.zero_grad()

        pbar = tqdm(self.train_loader, desc=f"Train Epoch {epoch + 1}", leave=False)
        for batch_idx, batch in enumerate(pbar):
            batch = self._move_batch_to_device(batch)

            with torch.amp.autocast(self.device.type, enabled=self.amp_dtype is not None, dtype=self.amp_dtype):
                logits = self.model(
                    batch["src_ids"],
                    batch["tgt_input_ids"],
                    batch.get("src_padding_mask"),
                    batch.get("tgt_padding_mask"),
                )
                loss = self.criterion(logits, batch["labels"])

            loss = loss / self.gradient_accumulation_steps

            if self.scaler is not None:
                self.scaler.scale(loss).backward()
            else:
                loss.backward()

            total_loss += loss.item() * self.gradient_accumulation_steps
            num_batches += 1

            if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
                if self.scaler is not None:
                    self.scaler.unscale_(self.optimizer)
                    nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
                    self.optimizer.step()

                self.scheduler.step()
                self.optimizer.zero_grad()
                self.global_step += 1

                current_lr = self.scheduler.get_last_lr()[0]
                pbar.set_postfix({
                    "loss": f"{loss.item() * self.gradient_accumulation_steps:.4f}",
                    "lr": f"{current_lr:.2e}",
                    "step": self.global_step,
                })

                if self.global_step % self.log_every_n_steps == 0:
                    self._log_step_metrics(loss.item() * self.gradient_accumulation_steps, current_lr)

                if self.save_every_n_steps > 0 and self.global_step % self.save_every_n_steps == 0:
                    self._save_checkpoint(epoch, {"step": self.global_step}, prefix=f"step_{self.global_step}")

        avg_loss = total_loss / max(num_batches, 1)
        logger.info("Epoch %d - Train Loss: %.4f", epoch + 1, avg_loss)
        return avg_loss

    @torch.no_grad()
    def _validate(self, epoch: int) -> dict:
        self.model.eval()
        total_loss = 0.0
        num_batches = 0

        pbar = tqdm(self.val_loader, desc=f"Val Epoch {epoch + 1}", leave=False)
        for batch in pbar:
            batch = self._move_batch_to_device(batch)

            with torch.amp.autocast(self.device.type, enabled=self.amp_dtype is not None, dtype=self.amp_dtype):
                logits = self.model(
                    batch["src_ids"],
                    batch["tgt_input_ids"],
                    batch.get("src_padding_mask"),
                    batch.get("tgt_padding_mask"),
                )
                loss = self.criterion(logits, batch["labels"])
            total_loss += loss.item()
            num_batches += 1

        val_loss = total_loss / max(num_batches, 1)
        metrics = {"val_loss": round(val_loss, 4)}

        if self.evaluator is not None and self.config.get("evaluation", {}).get("eval_on_epoch_end", True):
            try:
                eval_results = self.evaluator.evaluate(self.val_loader)
                metrics.update(eval_results)
            except Exception as e:
                logger.warning("Evaluation failed during validation: %s", e)

        logger.info("Epoch %d - Val Loss: %.4f", epoch + 1, val_loss)
        for k, v in metrics.items():
            if k != "val_loss" and not isinstance(v, list):
                logger.info("  %s: %.4f", k, v)

        return metrics

    def _save_checkpoint(self, epoch: int, metrics: dict, prefix: str = ""):
        model_to_save = self.model.module if hasattr(self.model, "module") else self.model

        checkpoint = {
            "epoch": epoch,
            "step": self.global_step,
            "model_state_dict": model_to_save.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "metrics": metrics,
            "config": self.config,
            "train_loss_history": self.train_loss_history,
            "val_metrics_history": self.val_metrics_history,
        }

        if prefix:
            ckpt_path = self.checkpoint_dir / f"checkpoint_{prefix}.pt"
        else:
            ckpt_path = self.checkpoint_dir / f"checkpoint_epoch_{epoch + 1}.pt"

        torch.save(checkpoint, ckpt_path)
        logger.info("Checkpoint saved: %s", ckpt_path)

        current_metric = metrics.get(self.metric_for_best, metrics.get("val_loss", float("inf")))
        if self.metric_for_best == "val_loss":
            current_metric = -current_metric

        if self.save_best and current_metric > self.best_metric_value:
            self.best_metric_value = current_metric
            self.best_epoch = epoch
            best_path = self.checkpoint_dir / "best_model.pt"
            torch.save(checkpoint, best_path)
            logger.info("New best model saved: %s (metric: %.4f)", best_path, current_metric)

        self._cleanup_old_checkpoints()

    def _cleanup_old_checkpoints(self):
        ckpt_files = sorted(
            self.checkpoint_dir.glob("checkpoint_epoch_*.pt"),
            key=os.path.getmtime,
        )
        while len(ckpt_files) > self.max_checkpoints:
            oldest = ckpt_files.pop(0)
            oldest.unlink()
            logger.debug("Removed old checkpoint: %s", oldest)

    def _load_checkpoint(self, checkpoint_path: str):
        logger.info("Loading checkpoint from %s", checkpoint_path)
        checkpoint = torch.load(checkpoint_path, map_location=self.device)

        model_to_load = self.model.module if hasattr(self.model, "module") else self.model
        model_to_load.load_state_dict(checkpoint["model_state_dict"])

        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        self.current_epoch = checkpoint["epoch"] + 1
        self.global_step = checkpoint["step"]
        self.best_metric_value = checkpoint.get("metrics", {}).get(
            self.metric_for_best, float("-inf")
        )
        self.train_loss_history = checkpoint.get("train_loss_history", [])
        self.val_metrics_history = checkpoint.get("val_metrics_history", [])

        logger.info("Resumed from epoch %d, step %d", self.current_epoch, self.global_step)

    def _should_early_stop(self, metrics: dict) -> bool:
        if not self.early_stopping_enabled:
            return False

        current_metric = metrics.get(self.metric_for_best, metrics.get("val_loss", float("inf")))
        if self.metric_for_best == "val_loss":
            current_metric = -current_metric

        if current_metric > self.best_metric_value + self.min_delta:
            self.epochs_without_improvement = 0
            return False

        self.epochs_without_improvement += 1
        logger.info(
            "No improvement for %d epochs (best: %.4f, current: %.4f)",
            self.epochs_without_improvement, self.best_metric_value, current_metric,
        )
        return self.epochs_without_improvement >= self.patience

    def _log_metrics(self, epoch: int, train_loss: float, val_metrics: dict):
        if self.writer is not None:
            self.writer.add_scalar("Loss/train", train_loss, epoch)
            for k, v in val_metrics.items():
                if not isinstance(v, list):
                    self.writer.add_scalar(f"Metrics/{k}", v, epoch)
            self.writer.add_scalar("LR", self.scheduler.get_last_lr()[0], epoch)

        if self.wandb_run is not None:
            import wandb
            log_dict = {"epoch": epoch, "train_loss": train_loss}
            for k, v in val_metrics.items():
                if not isinstance(v, list):
                    log_dict[f"val/{k}"] = v
            log_dict["lr"] = self.scheduler.get_last_lr()[0]
            wandb.log(log_dict, step=self.global_step)

    def _log_step_metrics(self, loss: float, lr: float):
        if self.writer is not None:
            self.writer.add_scalar("Loss/train_step", loss, self.global_step)
            self.writer.add_scalar("LR/step", lr, self.global_step)

        if self.wandb_run is not None:
            import wandb
            wandb.log({"train/loss_step": loss, "lr": lr}, step=self.global_step)

    def _log_final_results(self):
        logger.info("=" * 50)
        logger.info("Training completed!")
        logger.info("Best epoch: %d", self.best_epoch + 1)
        logger.info("Best %s: %.4f", self.metric_for_best, self.best_metric_value)

        summary = {
            "best_epoch": self.best_epoch,
            "best_metric": self.best_metric_value,
            "metric_name": self.metric_for_best,
            "total_steps": self.global_step,
            "train_loss_history": self.train_loss_history,
            "val_metrics_history": self.val_metrics_history,
        }
        summary_path = self.checkpoint_dir / "training_summary.json"
        with open(summary_path, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
        logger.info("Training summary saved to %s", summary_path)

    def _cleanup(self):
        if self.writer is not None:
            self.writer.close()
        if self.wandb_run is not None:
            self.wandb_run.finish()

    def _move_batch_to_device(self, batch: dict) -> dict:
        return {
            k: v.to(self.device, non_blocking=True) if isinstance(v, torch.Tensor) else v
            for k, v in batch.items()
        }