Robot_Arm/ml/common/ckpt.py

"""断点存档与"体面停止"（2026-10-18 第五轮补充：用户要求长训练意外中断后能接着跑，不从头来）。

- atomic_save：先写 <名>.tmp 再 os.replace —— 断电/被杀发生在写盘中途时，旧的 last.pt 仍完整可用。
- Saver.due()：按墙钟每 every_min 分钟存一次 last.pt（不再只在每 2000 步验证时存，断电最多丢 10 分钟）。
- Saver.stop()：收到 SIGTERM/SIGINT，或看到停止标记文件 ml/runs/STOP 时返回 True；
  训练循环据此先存 last.pt 再以 143 退出。Windows 上 taskkill /F 是硬杀、收不到信号，
  所以控制台「停止」先写 STOP 文件、等进程自己存档退出，超时才硬杀（见 monitor.Job.stop）。
"""
import os, random, signal, time
import torch

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
STOP = os.path.join(ROOT, 'ml', 'runs', 'STOP')


def atomic_save(obj, path):
    tmp = path + '.tmp'
    torch.save(obj, tmp)
    os.replace(tmp, path)


class Saver:
    def __init__(self, every_min=10.0):
        self.every = every_min * 60; self.last = time.time(); self._sig = False; self._chk = 0.0
        for s in ('SIGTERM', 'SIGINT', 'SIGBREAK', 'SIGHUP'):
            if hasattr(signal, s):
                try: signal.signal(getattr(signal, s), self._on)
                except (ValueError, OSError): pass   # 非主线程等情况：只靠 STOP 文件

    def _on(self, *_): self._sig = True

    def stop(self):
        if self._sig: return True
        now = time.time()
        if now - self._chk > 2:                      # 文件每 2 s 查一次，不拖慢训练步
            self._chk = now; self._sig = os.path.exists(STOP)
        return self._sig

    def due(self):
        if time.time() - self.last >= self.every:
            self.last = time.time(); return True
        return False

    def mark(self): self.last = time.time()


def rng_state():
    """2026-10-18 第七轮（遗留：续跑后数值轨迹与不断点不一致）：随机数状态一起存档。
    训练取 batch 用 torch.randint（全局生成器，CUDA 上取 CUDA 生成器）、噪声用 torch.randn、APG 选题用 random.choices，
    三者状态都恢复后，续跑的 batch 顺序与噪声与一口气跑完一致（同一设备、同一版本；cuDNN 非确定算子另计）。"""
    d = {'torch': torch.get_rng_state(), 'py': random.getstate()}
    if torch.cuda.is_available(): d['cuda'] = torch.cuda.get_rng_state_all()
    return d   # 不存 numpy 状态：训练不用 np.random，且 ndarray 过不了 torch.load 的 weights_only 默认


def set_rng_state(d):
    if not d: return False
    torch.set_rng_state(d['torch'].cpu() if hasattr(d['torch'], 'cpu') else d['torch']); random.setstate(d['py'])
    if 'cuda' in d and torch.cuda.is_available():
        try: torch.cuda.set_rng_state_all([t.cpu() for t in d['cuda']])
        except Exception: pass
    return True