# -*- coding: utf-8 -*- """Checkpoints de treino — v5 (doc 16 §8). Substitui `dados/estados_hf.py` (REMOVIDO a pedido — estados/fase-som-v4 apagados do Hub e do disco). Regra v5 motivada pelo bloqueio do Hub a uploads parciais em sequência curta: • durante o treino, checkpoints ficam SEMPRE LOCAIS (safetensors + SHA-256, protocolo idêntico ao doc 10 §9) e o diretório mantém só o último (poupa disco, item c do escopo); • NENHUM upload parcial durante o treino; • o snapshot FINAL entra no commit ÚNICO da publicação (scripts/ 06_publicar_hf.py chama `preparar_para_publicacao`, que copia o par (modelo.safetensors, meta.json) para o stage) — um único commit organizado. Teorema 16.9 (integridade do snapshot): seja σ o par gravado atomicamente com hash h(σ)=SHA-256. Na leitura, se h(σ')=h(σ) então σ'=σ (colisão desprezível, 2^−256), logo retomadas reproduzem EXATAMENTE o estado salvo. """ from __future__ import annotations import hashlib import json import os import shutil import torch from safetensors.torch import load as _st_load from safetensors.torch import save as _st_save def _sha256_bytes(b: bytes) -> str: return hashlib.sha256(b).hexdigest() class GestorCheckpoints: """Checkpoints locais atômicos + snapshot final para publicação única. Args: dir_local: diretório de checkpoints (mantém apenas o último). pasta_publicacao: pasta dentro do repo HF onde o snapshot final entra NO commit único (ex.: 'estados/fase-v5'). """ def __init__(self, dir_local: str = "/home/z/my-project/aurora/estados_local", pasta_publicacao: str = "estados/fase-v5"): self.dir_local = dir_local self.pasta_publicacao = pasta_publicacao self.eventos: list[str] = [] os.makedirs(dir_local, exist_ok=True) # ---------------- protocolo (doc 10 §9 preservado) ---------------- def _caminho_tag(self, tag: str) -> str: return os.path.join(self.dir_local, tag) def salvar(self, modelo, tag: str, meta: dict) -> str: """Serializa safetensors + meta com hash; mantém apenas a última tag.""" sd = {k: v.detach().contiguous() for k, v in modelo.state_dict().items()} # pesos empatados (tied): uma única cópia (load com strict=False) vistos: dict[int, str] = {} sd_unico: dict[str, torch.Tensor] = {} for k, v in sd.items(): ptr = v.data_ptr() if ptr in vistos: continue vistos[ptr] = k sd_unico[k] = v dados = bytes(_st_save(sd_unico, metadata={"formato": "aurora-v5", "tag": tag})) meta = dict(meta or {}) meta["hash_sha256"] = _sha256_bytes(dados) meta["tag"] = tag destino = self._caminho_tag(tag) os.makedirs(destino, exist_ok=True) with open(os.path.join(destino, "modelo.safetensors"), "wb") as f: f.write(dados) with open(os.path.join(destino, "meta.json"), "w", encoding="utf-8") as f: json.dump(meta, f, ensure_ascii=False, indent=2) # higiene: mantém SOMENTE a última tag (poupa disco, item c) removidos = 0 for nome in os.listdir(self.dir_local): if nome != tag: shutil.rmtree(os.path.join(self.dir_local, nome), ignore_errors=True) removidos += 1 if removidos: self.eventos.append(f"{removidos} checkpoint(s) antigo(s) apagado(s)") self.eventos.append(f"checkpoint {tag}: local ({len(dados)/1e6:.1f} MB)") return os.path.join(destino, "modelo.safetensors") def carregar(self, modelo, tag: str | None = None) -> bool: """Carrega tolerante a forma (mesma política da v3: tensores por chave com mesma forma; divergências registradas).""" tag = tag or self.ultima_tag() if tag is None: return False caminho = os.path.join(self._caminho_tag(tag), "modelo.safetensors") if not os.path.exists(caminho): self.eventos.append(f"checkpoint {tag}: ausente") return False with open(caminho, "rb") as f: dados = f.read() meta_path = os.path.join(self._caminho_tag(tag), "meta.json") if os.path.exists(meta_path): with open(meta_path, encoding="utf-8") as f: esperado = json.load(f).get("hash_sha256") if esperado and _sha256_bytes(dados) != esperado: self.eventos.append(f"checkpoint {tag}: HASH divergente — rejeitado") return False try: sd = _st_load(dados) modelo_sd = modelo.state_dict() próprios = {k: v for k, v in sd.items() if k in modelo_sd and modelo_sd[k].shape == v.shape} pulados = len(sd) - len(próprios) faltantes = [k for k in modelo_sd if k not in próprios and k not in sd] modelo.load_state_dict(próprios, strict=False) self.eventos.append( f"checkpoint {tag}: {len(próprios)} tensores, {pulados} pulados, " f"{len(faltantes)} novos") return True except Exception as e: self.eventos.append(f"carregar {tag}: {type(e).__name__}: {e}") return False def listar(self) -> list: """Tags disponíveis localmente, ordenadas lexicograficamente (tags canônicas `epoca-*`/`fase-*` são cronológicas).""" if not os.path.isdir(self.dir_local): return [] return sorted(d for d in os.listdir(self.dir_local) if os.path.isdir(os.path.join(self.dir_local, d)) and os.path.exists(os.path.join(self.dir_local, d, "meta.json"))) def ultima_tag(self) -> str | None: if not os.path.isdir(self.dir_local): return None tags = [d for d in os.listdir(self.dir_local) if os.path.isdir(os.path.join(self.dir_local, d)) and os.path.exists(os.path.join(self.dir_local, d, "meta.json"))] if not tags: return None # mais avançado = maior passo_global no meta (ordem cronológica) melhor, melhor_passo = None, -1 for tag in tags: try: with open(os.path.join(self.dir_local, tag, "meta.json"), encoding="utf-8") as f: passo = int(json.load(f).get("passo_global", -1)) except Exception: continue if passo > melhor_passo: melhor, melhor_passo = tag, passo return melhor def meta_de(self, tag: str | None = None) -> dict: tag = tag or self.ultima_tag() if tag is None: return {} caminho = os.path.join(self._caminho_tag(tag), "meta.json") if not os.path.exists(caminho): return {} with open(caminho, encoding="utf-8") as f: return json.load(f) # ---------------- publicação ÚNICA (doc 16 §8) ---------------- def preparar_para_publicacao(self, stage_raiz: str) -> str | None: """Copia o snapshot final para o stage do repo (commit único do script 06). Devolve o caminho do peso no stage.""" tag = self.ultima_tag() if tag is None: self.eventos.append("publicação: nenhum checkpoint local") return None origem = self._caminho_tag(tag) destino = os.path.join(stage_raiz, self.pasta_publicacao) os.makedirs(destino, exist_ok=True) shutil.copy2(os.path.join(origem, "modelo.safetensors"), os.path.join(destino, "modelo.safetensors")) meta = self.meta_de(tag) meta["publicado_em_commit_unico"] = True with open(os.path.join(destino, "meta.json"), "w", encoding="utf-8") as f: json.dump(meta, f, ensure_ascii=False, indent=2) self.eventos.append(f"snapshot {tag} → stage ({self.pasta_publicacao})") return os.path.join(destino, "modelo.safetensors")