src/lmr/checkpointing/checkpointing.py

import time
import re
import json
import shutil
from pathlib import Path
from glob import glob
from typing import Dict, Any, Tuple, List, Optional

import torch

from lmr.utils.logger import Logger
from lmr.utils.parsing import int_to_formatted_string
from lmr.ddp import unwrap_model


def _find_checkpoint_file_in_dir(path_dir: Path) -> Path:
    if not path_dir.is_dir():
        raise FileNotFoundError(f"{path_dir} is not a dir")

    idx = path_dir / "model.safetensors.index.json"
    if idx.exists():
        return idx

    safes = list(path_dir.glob("*.safetensors"))
    if safes:
        for cand in safes:
            if cand.name == "model.safetensors":
                return cand
        safes.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return safes[0]

    bins = list(path_dir.glob("pytorch_model.bin")) + list(path_dir.glob("*.pt"))
    if bins:
        for cand in bins:
            if cand.name == "pytorch_model.bin":
                return cand
        bins.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return bins[0]

    others = [p for p in path_dir.iterdir() if p.is_file()]
    if others:
        others.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return others[0]

    raise FileNotFoundError(f"No checkpoint file found in directory: {path_dir}")


def _load_safetensors_shards_from_index(index_path: Path, map_location="cpu") -> Dict[str, torch.Tensor]:
    from safetensors.torch import load_file

    base_dir = index_path.parent
    with open(index_path, "r") as f:
        index_data = json.load(f)
    weight_map = index_data.get("weight_map", {})
    shards = sorted(set(weight_map.values()))
    merged = {}
    for shard_name in shards:
        shard_path = base_dir / shard_name
        if not shard_path.exists():
            raise FileNotFoundError(f"Shard not found: {shard_path}")
        shard = load_file(str(shard_path), device=map_location)
        merged.update(shard)
    return merged


def _safe_torch_load(path: Path, map_location="cpu", allow_unsafe_fallback: bool = True) -> Any:
    try:
        return torch.load(str(path), map_location=map_location, weights_only=True)
    except TypeError:
        return torch.load(str(path), map_location=map_location)
    except Exception:
        if not allow_unsafe_fallback:
            raise
        return torch.load(str(path), map_location=map_location, weights_only=False)


class Checkpointing:
    def __init__(self, model, checkpoint_dir, optimizer=None, scheduler=None, scaler=None, map_device="cpu"):
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.scaler = scaler
        self.map_device = map_device

        # state
        self.epoch = 0
        self.step = 0
        self.train_loss = float("inf")
        self.val_loss = float("inf")
        self.tokens_trained = 0

        # NEW: accuracies
        self.val_acc = None
        self.train_acc = None

        self.use_ddp = torch.distributed.is_available() and torch.distributed.is_initialized()
        self.is_main_process = (not self.use_ddp) or (torch.distributed.get_rank() == 0)

        self.checkpoint_dir = Path(checkpoint_dir)
        if not self.is_main_process:
            return

        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.log_path = self.checkpoint_dir / "_checkpoint_log.tsv"
        self._create_log()

        self.best_val_loss = self._get_best_val_loss()

    def _barrier(self):
        if self.use_ddp:
            torch.distributed.barrier()

    # ---------- best loss scan ----------
    def _get_best_val_loss(self):
        best_dirs = glob(str(self.checkpoint_dir / "best_epoch*_val=*"))
        best_val = float("inf")
        for dir_path in best_dirs:
            match = re.search(r"val=([0-9.]+)", dir_path)
            if match:
                try:
                    best_val = min(best_val, float(match.group(1)))
                except ValueError:
                    pass
        return best_val

    def _checkpoint_dirname(self, epoch, step=None, val_loss=None, tokens_trained=None, prefix=None):
        name = f"epoch_{epoch:03d}"
        if prefix is not None:
            name = f"{prefix}_{name}"
        if step is not None:
            name += f"_step_{step:09d}"
        if tokens_trained is not None:
            name += f"_tokens_{int_to_formatted_string(tokens_trained)}"
        if val_loss is not None:
            name += f"_val={val_loss:.4f}"
        return name

    # ---------- log ----------
    def _create_log(self):
        if self.log_path.exists():
            return
        header = "Time\tCheckpoint_Type\tEpoch\tStep\tTrain_Loss\tVal_Loss\tTrain_Acc\tVal_Acc\tTokens_Trained\tDirname\n"
        with open(self.log_path, "w", encoding="utf-8") as f:
            f.write(header)

    def _update_log(self, kind, dirname):
        ts = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        epoch_str = f"{self.epoch:03d}"
        step_str = f"{self.step:09d}"
        line = (
            f"{ts}\t{kind}\t{epoch_str}\t{step_str}\t"
            f"{'' if self.train_loss is None else self.train_loss}\t"
            f"{'' if self.val_loss is None else self.val_loss}\t"
            f"{'' if self.train_acc is None else self.train_acc}\t"
            f"{'' if self.val_acc is None else self.val_acc}\t"
            f"{self.tokens_trained}\t"
            f"{dirname}\n"
        )
        with open(self.log_path, "a", encoding="utf-8") as f:
            f.write(line)

    # ---------- remove old ----------
    def _remove_old(self, pattern: str):
        """
        Delete old checkpoint directories matching pattern, e.g. "recent_epoch*"
        Keeps only the most recent one if multiple exist.
        """
        if not self.is_main_process:
            return

        dirs = sorted(glob(str(self.checkpoint_dir / pattern)))
        if len(dirs) <= 1:
            return

        # keep newest by mtime
        dirs_paths = [Path(d) for d in dirs]
        dirs_paths.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        keep = dirs_paths[0]
        to_remove = dirs_paths[1:]

        # for p in to_remove:
        #     try:
        #         shutil.rmtree(p)
        #         Logger.log(f"🧹 Removed old checkpoint dir: {p.name}")
        #     except Exception as e:
        #         Logger.log(f"⚠️ Failed to remove old checkpoint dir {p}: {e}")

    # ---------- save ----------
    def _save_state(self, dirname, include_training_states=False):
        save_path = self.checkpoint_dir / dirname
        save_path.mkdir(parents=True, exist_ok=True)

        model = unwrap_model(self.model)

        if hasattr(model, "save_pretrained"):
            # Recommended: safetensors + shards
            model.save_pretrained(save_path, safe_serialization=True, max_shard_size="5GB")
        else:
            Logger.log("⚠️ Model does not have save_pretrained; saving state_dict.")
            from safetensors.torch import save_file
            state_dict = {k: v.cpu().contiguous() for k, v in model.state_dict().items()}
            save_file(state_dict, str(save_path / "model.safetensors"))

        # trainer metadata
        metadata = {
            "epoch": self.epoch,
            "step": self.step,
            "train_loss": self.train_loss,
            "val_loss": self.val_loss,
            "tokens_trained": self.tokens_trained,
            "train_acc": self.train_acc,
            "val_acc": self.val_acc,
        }
        with open(save_path / "trainer_state.json", "w") as f:
            json.dump(metadata, f, indent=4)

        if include_training_states:
            train_state = {}
            if self.optimizer is not None:
                train_state["optimizer"] = self.optimizer.state_dict()
            if self.scheduler is not None and hasattr(self.scheduler, "state_dict"):
                train_state["scheduler"] = self.scheduler.state_dict()
            if self.scaler is not None and hasattr(self.scaler, "state_dict"):
                train_state["scaler"] = self.scaler.state_dict()
            torch.save(train_state, save_path / "optimizer.pt")

        Logger.log(f"💾 Checkpoint saved to: {save_path}")

    def _update_state(
        self,
        epoch,
        step=0,
        train_loss=None,
        val_loss=None,
        tokens_trained=None,
        train_acc=None,
        val_acc=None,
    ):
        self.epoch = int(epoch)
        self.step = int(step) if step is not None else 0
        if train_loss is not None:
            self.train_loss = train_loss
        if val_loss is not None:
            self.val_loss = val_loss
        if tokens_trained is not None:
            self.tokens_trained = int(tokens_trained)

        # NEW
        if train_acc is not None:
            self.train_acc = float(train_acc)
        if val_acc is not None:
            self.val_acc = float(val_acc)

    def _save_best(self):
        self._remove_old("best_epoch*")
        dirname = self._checkpoint_dirname(self.epoch, self.step, self.val_loss, self.tokens_trained, prefix="best")
        self._save_state(dirname, include_training_states=False)
        self._update_log("best", dirname)

    def _save_recent(self):
        self._remove_old("recent_epoch*")
        dirname = self._checkpoint_dirname(self.epoch, self.step, self.val_loss, self.tokens_trained, prefix="recent")
        self._save_state(dirname, include_training_states=True)
        self._update_log("recent", dirname)

    def _save_epoch(self):
        dirname = self._checkpoint_dirname(self.epoch, None, self.val_loss, self.tokens_trained)
        self._save_state(dirname, include_training_states=False)
        self._update_log("epoch", dirname)

    # ---------- public save ----------
    def save_checkpoint(
        self,
        epoch,
        step=None,
        train_loss=None,
        val_loss=None,
        tokens_trained=None,
        val_acc=None,
        train_acc=None,
    ):
        if self.is_main_process:
            self._update_state(
                epoch,
                step=step,
                train_loss=train_loss,
                val_loss=val_loss,
                tokens_trained=tokens_trained,
                train_acc=train_acc,
                val_acc=val_acc,
            )

            self._save_recent()

            if step is None:
                self.step = 0
                self._save_epoch()

            if (self.val_loss is not None) and (self.val_loss < self.best_val_loss):
                self.best_val_loss = self.val_loss
                self._save_best()

        self._barrier()

    # ---------- find ckpt dir ----------
    def _get_checkpoint_path(self, checkpoint_type):
        if checkpoint_type == "best":
            pattern = "best_epoch*"
        elif checkpoint_type == "recent":
            pattern = "recent_epoch*"
        elif checkpoint_type.startswith("epoch_"):
            pattern = f"{checkpoint_type}*"
        else:
            path = self.checkpoint_dir / checkpoint_type
            if path.exists():
                return path
            return None

        candidates = sorted(glob(str(self.checkpoint_dir / pattern)))
        if not candidates:
            Logger.log(f"No checkpoint found for {checkpoint_type}")
            return None
        return Path(candidates[-1])

    # ---------- load model weights + metadata ----------
    def load_model_states(self, checkpoint_type="recent"):
        ckpt_dir = self._get_checkpoint_path(checkpoint_type)
        if not ckpt_dir:
            return None

        Logger.log(f"📂 Loading model from dir: {ckpt_dir}")

        meta_path = ckpt_dir / "trainer_state.json"
        if meta_path.exists():
            with open(meta_path, "r") as f:
                state = json.load(f)
            self.epoch = int(state.get("epoch", 0))
            self.step = int(state.get("step", 0))
            self.train_loss = state.get("train_loss", None)
            self.val_loss = state.get("val_loss", None)
            self.tokens_trained = int(state.get("tokens_trained", 0))
            self.train_acc = state.get("train_acc", None)
            self.val_acc = state.get("val_acc", None)

        try:
            self.load_only_model_weights(str(ckpt_dir), map_location=self.map_device, strict=False, verbose=True)
        except Exception as e:
            Logger.log(f"⚠️ Failed to load model weights from {ckpt_dir}: {e}")

        self._barrier()
        return ckpt_dir

    def load_only_model_weights(
        self,
        checkpoint_path: str,
        model: Optional[torch.nn.Module] = None,
        device: Optional[str] = None,
        map_location: Optional[str] = "cpu",
        save_filtered_to: Optional[str] = None,
        strict: bool = False,
        verbose: bool = True,
        allow_unsafe_fallback: bool = True,
        allow_continue_on_failure: bool = False,
        **kwargs
    ) -> Dict[str, Any]:
        load_model = model if model is not None else unwrap_model(self.model)
        p = Path(checkpoint_path)
        final_map = device or map_location

        # read raw ckpt
        try:
            if p.exists() and p.is_dir():
                candidate = _find_checkpoint_file_in_dir(p)
                if candidate.name == "model.safetensors.index.json":
                    if verbose:
                        Logger.log(f"[ckpt_loader] Detected safetensors sharded index at {candidate}; loading shards...")
                    raw_state = _load_safetensors_shards_from_index(candidate, map_location=final_map)
                else:
                    if verbose:
                        Logger.log(f"[ckpt_loader] Directory provided; selected checkpoint file: {candidate}")
                    if candidate.suffix == ".safetensors":
                        from safetensors.torch import load_file
                        raw_state = load_file(str(candidate), device=final_map)
                    else:
                        raw_state = _safe_torch_load(candidate, map_location=final_map, allow_unsafe_fallback=allow_unsafe_fallback)
            elif p.exists() and p.is_file():
                if p.suffix == ".safetensors":
                    from safetensors.torch import load_file
                    raw_state = load_file(str(p), device=final_map)
                else:
                    raw_state = _safe_torch_load(p, map_location=final_map, allow_unsafe_fallback=allow_unsafe_fallback)
            else:
                raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}")
        except Exception as e_load:
            msg = f"[ckpt_loader] Failed to load checkpoint '{checkpoint_path}': {e_load}"
            if not allow_continue_on_failure:
                raise RuntimeError(msg)
            Logger.log(msg + " -- continuing (allow_continue_on_failure=True).")
            return {"matched": 0, "total_target": len(load_model.state_dict()), "error": str(e_load)}

        # extract state_dict
        def _extract_state_dict(raw):
            if isinstance(raw, dict):
                sample_keys = list(raw.keys())[:10]
                if any(("weight" in k or "bias" in k or "embed" in k) for k in sample_keys):
                    return raw
                for candidate_key in ("model", "state_dict", "model_state_dict", "state"):
                    if candidate_key in raw and isinstance(raw[candidate_key], dict):
                        return raw[candidate_key]
            if isinstance(raw, dict):
                dicts = [v for v in raw.values() if isinstance(v, dict)]
                if dicts:
                    return max(dicts, key=lambda x: len(x))
            raise RuntimeError(f"Could not find a state_dict inside checkpoint (raw type={type(raw)})")

        state_dict = _extract_state_dict(raw_state)

        # normalize keys
        def _normalize_keys(state: Dict[str, Any], prefixes=("module.", "_orig_mod.", "model_state.")):
            normalized = {}
            for k, v in state.items():
                new_k = k
                for pfx in prefixes:
                    if new_k.startswith(pfx):
                        new_k = new_k[len(pfx):]
                normalized[new_k] = v
            return normalized

        normalized = _normalize_keys(state_dict)
        if verbose:
            print(f"[ckpt_loader] extracted {len(normalized)} params (sample: {list(normalized.keys())[:10]})")

        # match by name+shape
        target_state = load_model.state_dict()
        filtered = {}
        size_mismatch: List[Tuple[str, Any, Any]] = []
        missing: List[str] = []
        matched = 0

        for tk, tv in target_state.items():
            if tk in normalized:
                ck = normalized[tk]
                if getattr(ck, "shape", None) == getattr(tv, "shape", None):
                    try:
                        filtered[tk] = ck.to(tv.device) if hasattr(ck, "to") else ck
                    except Exception:
                        filtered[tk] = ck
                    matched += 1
                else:
                    size_mismatch.append((tk, getattr(ck, "shape", None), tv.shape))
            else:
                missing.append(tk)

        if verbose:
            print(f"[ckpt_loader] Matched: {matched}/{len(target_state)}; size_mismatch: {len(size_mismatch)}; missing: {len(missing)}")

        load_msg = load_model.load_state_dict(filtered, strict=False)

        import pdb
        # pdb.set_trace()
        if verbose:
            print(f"[ckpt_loader] load_state_dict: {load_msg}")

        if save_filtered_to:
            try:
                torch.save(filtered, save_filtered_to)
                if verbose:
                    print(f"[ckpt_loader] Saved filtered weights to {save_filtered_to}")
            except Exception as e_save:
                if verbose:
                    print(f"[ckpt_loader] Warning: failed to save filtered weights: {e_save}")

        return {
            "matched": matched,
            "total_target": len(target_state),
            "size_mismatch": size_mismatch,
            "missing": missing,
            "load_msg": load_msg,
        }

    # ---------- load optimizer/scheduler/scaler ----------
    def load_training_states(self, checkpoint_type="recent"):
        ckpt_dir = self._get_checkpoint_path(checkpoint_type)
        if not ckpt_dir:
            return

        opt_path = ckpt_dir / "optimizer.pt"
        if not opt_path.exists():
            Logger.log(f"⚠️ Optimizer state not found in {ckpt_dir}")
            return

        state = torch.load(opt_path, map_location=self.map_device)

        if self.optimizer is not None and "optimizer" in state:
            self.optimizer.load_state_dict(state["optimizer"])
        if self.scheduler is not None and "scheduler" in state:
            self.scheduler.load_state_dict(state["scheduler"])
        if self.scaler is not None and "scaler" in state:
            self.scaler.load_state_dict(state["scaler"])

        Logger.log(f"✅ Loaded training states from {opt_path}")
        self._barrier()