hooks.py

"""训练 Hook 系统：将 LR 调度、entropy 调度、计时、安全告警、早停等
从 train.py 主循环中剥离为可组合的 hook 类。"""
import math
import time

import numpy as np
import torch


# ── 基础：Hook 基类 ──────────────────────────────────────────────


class TrainingHook:
    """所有 hook 的基类。每个方法可选覆写，默认 no-op。"""

    def on_step_start(self, step: int, optimizer, config: dict) -> None:
        pass

    def on_after_rollout(self, episodes: list, step: int) -> list:
        """可返回过滤后的 episodes（否则用原列表）。"""
        return episodes

    def on_after_update(self, step: int, results: dict, episodes: list,
                        duration: float) -> None:
        pass

    def should_eval(self, step: int, config: dict) -> bool:
        return step % config["training"]["eval_interval"] == 0

    def on_after_eval(self, step: int, success_rate: float) -> bool:
        """返回 True 则触发早停。"""
        return False

    def should_save_ckpt(self, step: int, config: dict) -> bool:
        return step % config["training"]["ckpt_save_interval"] == 0

    def on_after_save(self, step: int, output_file: str) -> None:
        pass

    def on_step_end(self, step: int, logger) -> None:
        """每步最后调用，用于写 TensorBoard。"""
        pass

    def should_stop(self, step: int, config: dict) -> bool:
        max_steps = config["training"].get("max_steps")
        return max_steps is not None and step >= max_steps


# ── Hook 管理器 ──────────────────────────────────────────────────


class HookManager:
    """持有多个 hook，按注册顺序依次调用。"""

    def __init__(self, hooks: list[TrainingHook] | None = None):
        self._hooks = hooks or []

    def add(self, hook: TrainingHook) -> "HookManager":
        self._hooks.append(hook)
        return self

    def on_step_start(self, step, optimizer, config):
        for h in self._hooks:
            h.on_step_start(step, optimizer, config)

    def on_after_rollout(self, episodes, step):
        for h in self._hooks:
            episodes = h.on_after_rollout(episodes, step)
        return episodes

    def on_after_update(self, step, results, episodes, duration):
        for h in self._hooks:
            h.on_after_update(step, results, episodes, duration)

    def should_eval(self, step, config):
        return any(h.should_eval(step, config) for h in self._hooks)

    def on_after_eval(self, step, success_rate):
        return any(h.on_after_eval(step, success_rate) for h in self._hooks)

    def should_save_ckpt(self, step, config):
        return any(h.should_save_ckpt(step, config) for h in self._hooks)

    def on_after_save(self, step, output_file):
        for h in self._hooks:
            h.on_after_save(step, output_file)

    def on_step_end(self, step, logger):
        for h in self._hooks:
            h.on_step_end(step, logger)

    def should_stop(self, step, config):
        return any(h.should_stop(step, config) for h in self._hooks)

    def get_entropy_coeff(self, config: dict) -> float:
        """从 EntropySchedulerHook 取当前 ec，否则用 config 默认值。"""
        for h in self._hooks:
            if isinstance(h, EntropySchedulerHook):
                return h.current_ec
        return config.get("training", {}).get("entropy_coeff", 0.0)


# ── 内置 Hook 实现 ───────────────────────────────────────────────


class LRSchedulerHook(TrainingHook):
    """warmup_cosine 学习率调度。"""

    def on_step_start(self, step, optimizer, config):
        sched = config.get("lr_scheduler", {})
        max_lr = config["training"]["learning_rate"]
        warmup = sched.get("warmup_steps", 0)
        total = sched.get("total_steps", 600)
        min_lr = sched.get("min_lr", max_lr / 10)

        if step < warmup:
            lr = max_lr * (step + 1) / warmup
        elif step >= total:
            lr = min_lr
        else:
            progress = (step - warmup) / (total - warmup)
            lr = min_lr + (max_lr - min_lr) * (1 + math.cos(math.pi * progress)) / 2
        for group in optimizer.param_groups:
            group["lr"] = lr


class EntropySchedulerHook(TrainingHook):
    """piecewise 分段衰减 entropy coefficient。"""

    def __init__(self):
        self.current_ec = 0.0

    def on_step_start(self, step, optimizer, config):
        default = config["training"].get("entropy_coeff", 0.0)
        sched = config.get("entropy_schedule", {})
        milestones = sched.get("milestones", [])
        ec = default
        for m_step, m_val in milestones:
            if step >= m_step:
                ec = m_val
        self.current_ec = ec


class TimingHook(TrainingHook):
    """分段计时（rollout / update / eval / ckpt / mi / filter）。"""

    def __init__(self):
        self.timings: dict[str, float] = {}
        self._step_start = 0.0
        self._seg_start = 0.0

    def on_step_start(self, step, optimizer, config):
        self._step_start = time.time()
        self._seg_start = self._step_start
        self.timings = {}

    def time_segment(self, name: str):
        """记录一段分时（从上次 time_segment 或 step_start 开始）。"""
        now = time.time()
        self.timings[name] = now - self._seg_start
        self._seg_start = now

    def finish_step(self, step):
        """step 结束，同步 CUDA 并记录 total duration。"""
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        self.timings["duration"] = time.time() - self._step_start


class SafetyWarningHook(TrainingHook):
    """安全告警：entropy 过早坍缩 / grad_norm 过高 / format_reward 过低。"""

    def on_after_update(self, step, results, episodes, duration):
        if "entropy" not in results:
            return
        entropy = results["entropy"]
        grad_norm = results["grad_norm"]
        reward = [ep.reward_info.get("format_reward", 0) for ep in episodes]
        format_reward = np.mean(reward)
        if step < 100 and entropy < 0.15:
            print(f"\r⚠️ Early entropy collapse at step {step}: entropy={entropy:.3f}",
                  flush=True)
        if grad_norm > 3.0:
            print(f"\r⚠️ High grad_norm at step {step}: {grad_norm:.2f}", flush=True)
        if step > 100 and format_reward < 0.8:
            print(f"\r⚠️ Low format_reward at step {step}: {format_reward:.3f}",
                  flush=True)


class EarlyStopHook(TrainingHook):
    """Val 不再提升时触发早停。"""

    def __init__(self, patience: int = 50, eval_interval: int = 20):
        self.patience = patience
        self.eval_interval = eval_interval
        self.best = -1.0
        self.steps_without_improvement = 0

    def on_after_eval(self, step, success_rate):
        if success_rate > self.best:
            self.best = success_rate
            self.steps_without_improvement = 0
        else:
            self.steps_without_improvement += self.eval_interval
        if self.steps_without_improvement >= self.patience:
            print(f"\rEval not improved for {self.patience} steps "
                  f"(best={self.best:.2f}), stopping.")
            return True
        return False


class MIDetectionHook(TrainingHook):
    """MI collapse 检测（每 compute_freq 步运行一次）。

    需要在 build_hooks 之后调用 setup() 注入 model / device / dtype。
    """

    def __init__(self, compute_freq: int = 5, num_samples: int = 64,
                 pad_token_id: int = 0):
        self.compute_freq = compute_freq
        self.num_samples = num_samples
        self.pad_token_id = pad_token_id
        self.last_metrics = {}
        self._model = None
        self._device = torch.device("cuda")
        self._dtype = torch.bfloat16

    def setup(self, model, device, dtype):
        """由 train_v2 在模型加载后调用，注入运行时对象。"""
        self._model = model
        self._device = device
        self._dtype = dtype

    def on_after_rollout(self, episodes, step):
        if step % self.compute_freq != 0:
            return episodes
        if self._model is None:
            print("\r⚠️ MIDetectionHook: model not injected, skipping", flush=True)
            return episodes
        try:
            from mi_metrics import compute_mi_metrics
            self.last_metrics = compute_mi_metrics(
                episodes=episodes,
                model=self._model,
                device=self._device,
                dtype=self._dtype,
                pad_token_id=self.pad_token_id,
                num_samples=self.num_samples,
            )
        except Exception as e:
            print(f"\r⚠️ MI detection failed at step {step}: {e}", flush=True)
            self.last_metrics = {}
        return episodes


class RolloutFilterHook(TrainingHook):
    """SNR-Adaptive rollout 过滤。"""

    def __init__(self, top_p: float = 0.9, include_zero: bool = False,
                 selection_eps: float = 0.01):
        self.top_p = top_p
        self.include_zero = include_zero
        self.selection_eps = selection_eps

    def on_after_rollout(self, episodes, step):
        from rollout_filter import filter_episodes_by_reward_variance
        n_before = len(episodes)
        episodes = filter_episodes_by_reward_variance(
            episodes, top_p=self.top_p, include_zero=self.include_zero,
            selection_eps=self.selection_eps,
        )
        if n_before != len(episodes):
            print(f"\r* Rollout filter: kept {len(episodes)}/{n_before} episodes",
                  end=" ")
        return episodes


# ── 便捷工厂：从 config 构建 hook 集合 ─────────────────────────────


def build_hooks(config: dict) -> HookManager:
    """读取 config，自动组装对应的 hook。"""
    training = config.get("training", {})
    hooks: list[TrainingHook] = []

    # Timing（始终开启）
    hooks.append(TimingHook())

    # LR scheduler
    if config.get("lr_scheduler"):
        hooks.append(LRSchedulerHook())

    # Entropy scheduler
    if config.get("entropy_schedule"):
        hooks.append(EntropySchedulerHook())

    # Safety warnings
    if training.get("safety_warnings", True):
        hooks.append(SafetyWarningHook())

    # Early stop
    if training.get("early_stop_enabled", True):
        hooks.append(EarlyStopHook(
            patience=training.get("eval_patience", 50),
            eval_interval=training.get("eval_interval", 20),
        ))

    # MI detection
    mi = config.get("mi_detection", {})
    if mi.get("enabled", False):
        hooks.append(MIDetectionHook(
            compute_freq=mi.get("compute_freq", 5),
            num_samples=mi.get("num_samples", 64),
            pad_token_id=mi.get("pad_token_id", 0),
        ))

    # Rollout filter
    rf = config.get("rollout_filter", {})
    if rf.get("enabled", False):
        hooks.append(RolloutFilterHook(
            top_p=rf.get("top_p", 0.9),
            include_zero=rf.get("include_zero", False),
            selection_eps=rf.get("selection_eps", 0.01),
        ))

    return HookManager(hooks)