training/checkpoint.py

"""
Checkpoint management for training.

Single Responsibility: Only handles saving and loading checkpoints.
"""

import os
import torch
import threading
from typing import Dict, Any, Optional, List
from pathlib import Path
from dataclasses import dataclass
from .utils import log


@dataclass
class TrainingState:
    """Immutable training state for checkpointing."""
    step: int
    epoch: int
    loss: float
    text_ratio: float
    best_loss: float = float("inf")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "step": self.step,
            "epoch": self.epoch,
            "loss": self.loss,
            "text_ratio": self.text_ratio,
            "best_loss": self.best_loss,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> 'TrainingState':
        return cls(
            step=d.get("step", 0),
            epoch=d.get("epoch", 0),
            loss=d.get("loss", 0.0),
            text_ratio=d.get("text_ratio", 0.9),
            best_loss=d.get("best_loss", float("inf")),
        )


class CheckpointManager:
    """
    Manages checkpoint saving and loading with async support.

    Single Responsibility: Only handles checkpoint I/O.
    Open/Closed: Can extend with new checkpoint formats without modification.
    """

    def __init__(
        self,
        output_dir: str,
        prefix: str = "checkpoint",
        verbose: bool = True,
        max_checkpoints: Optional[int] = None
    ):
        """
        Initialize checkpoint manager.

        Args:
            output_dir: Directory for saving checkpoints
            prefix: Prefix for checkpoint filenames
            verbose: Whether to log operations
            max_checkpoints: Maximum checkpoints to keep (None = keep all)
        """
        self.output_dir = Path(output_dir)
        self.prefix = prefix
        self.verbose = verbose
        self.max_checkpoints = max_checkpoints
        self._save_threads: List[threading.Thread] = []

        # Create output directory
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def save(
        self,
        state_dict: Dict[str, Any],
        filename: str,
        async_save: bool = True
    ) -> str:
        """
        Save checkpoint.

        Args:
            state_dict: State dictionary to save
            filename: Checkpoint filename
            async_save: Whether to save asynchronously

        Returns:
            Path to saved checkpoint
        """
        path = self.output_dir / filename

        if async_save:
            self._save_async(state_dict, path)
        else:
            self._save_sync(state_dict, path)

        return str(path)

    def save_step(
        self,
        adapter_state: Dict[str, Any],
        optimizer_state: Dict[str, Any],
        training_state: TrainingState,
        async_save: bool = True
    ) -> str:
        """Save step checkpoint."""
        state_dict = {
            "adapter": adapter_state,
            "optimizer": optimizer_state,
            **training_state.to_dict()
        }
        filename = f"{self.prefix}_step{training_state.step}.pt"
        return self.save(state_dict, filename, async_save)

    def save_epoch(
        self,
        adapter_state: Dict[str, Any],
        optimizer_state: Dict[str, Any],
        training_state: TrainingState,
        async_save: bool = True
    ) -> str:
        """Save epoch checkpoint."""
        state_dict = {
            "adapter": adapter_state,
            "optimizer": optimizer_state,
            **training_state.to_dict()
        }
        filename = f"{self.prefix}_epoch{training_state.epoch}.pt"
        return self.save(state_dict, filename, async_save)

    def save_best(
        self,
        adapter_state: Dict[str, Any],
        training_state: TrainingState,
        lora_state: Optional[Dict[str, Any]] = None,
        async_save: bool = True
    ) -> str:
        """Save best model checkpoint."""
        state_dict = {
            "adapter": adapter_state,
            **training_state.to_dict()
        }
        if lora_state is not None:
            state_dict["lora"] = lora_state

        filename = f"{self.prefix}_best.pt"
        return self.save(state_dict, filename, async_save)

    def load(self, path: str) -> Dict[str, Any]:
        """
        Load checkpoint.

        Args:
            path: Path to checkpoint

        Returns:
            Loaded state dictionary
        """
        if self.verbose:
            log(f"Loading checkpoint: {path}")
        return torch.load(path, map_location="cpu", weights_only=False)

    def load_latest(self) -> Optional[Dict[str, Any]]:
        """Load the most recent checkpoint."""
        checkpoints = self._get_checkpoints()
        if not checkpoints:
            return None
        return self.load(str(checkpoints[-1]))

    def wait_for_saves(self):
        """Wait for all async saves to complete."""
        for t in self._save_threads:
            t.join()
        self._save_threads = []

    def _save_sync(self, state_dict: Dict[str, Any], path: Path):
        """Synchronous save."""
        # Copy tensors to CPU
        state_copy = self._copy_to_cpu(state_dict)
        torch.save(state_copy, path)
        if self.verbose:
            log(f"[Checkpoint] Saved: {path.name}")

    def _save_async(self, state_dict: Dict[str, Any], path: Path):
        """Asynchronous save."""
        # Clean up completed threads
        self._save_threads = [t for t in self._save_threads if t.is_alive()]

        # Copy tensors to CPU
        state_copy = self._copy_to_cpu(state_dict)

        def _save():
            try:
                torch.save(state_copy, path)
                if self.verbose:
                    log(f"[Checkpoint] Saved: {path.name}")
            except Exception as e:
                if self.verbose:
                    log(f"[Checkpoint] Error saving {path.name}: {e}")

        thread = threading.Thread(target=_save, daemon=True)
        thread.start()
        self._save_threads.append(thread)

        # Cleanup old checkpoints if needed
        if self.max_checkpoints:
            self._cleanup_old_checkpoints()

    def _copy_to_cpu(self, obj: Any) -> Any:
        """Recursively copy tensors to CPU."""
        if isinstance(obj, torch.Tensor):
            return obj.detach().cpu().clone()
        elif isinstance(obj, dict):
            return {k: self._copy_to_cpu(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self._copy_to_cpu(v) for v in obj]
        return obj

    def _get_checkpoints(self) -> List[Path]:
        """Get sorted list of checkpoint files."""
        pattern = f"{self.prefix}_step*.pt"
        checkpoints = sorted(
            self.output_dir.glob(pattern),
            key=lambda p: int(p.stem.split("step")[-1])
        )
        return checkpoints

    def _cleanup_old_checkpoints(self):
        """Remove old checkpoints beyond max_checkpoints."""
        if not self.max_checkpoints:
            return

        checkpoints = self._get_checkpoints()
        while len(checkpoints) > self.max_checkpoints:
            oldest = checkpoints.pop(0)
            try:
                oldest.unlink()
                if self.verbose:
                    log(f"[Checkpoint] Removed old: {oldest.name}")
            except Exception:
                pass


class Stage1CheckpointManager(CheckpointManager):
    """Checkpoint manager for Stage 1 training."""

    def __init__(self, output_dir: str, **kwargs):
        super().__init__(output_dir, prefix="stage1", **kwargs)


class Stage2CheckpointManager(CheckpointManager):
    """Checkpoint manager for Stage 2 training."""

    def __init__(self, output_dir: str, **kwargs):
        super().__init__(output_dir, prefix="stage2", **kwargs)

    def save_best(
        self,
        adapter_state: Dict[str, Any],
        training_state: TrainingState,
        lora_state: Dict[str, Any],
        async_save: bool = True
    ) -> str:
        """Save best model with LoRA weights."""
        return super().save_best(
            adapter_state, training_state,
            lora_state=lora_state, async_save=async_save
        )