# -*- coding: utf-8 -*- """Lightweight task-scoped checkpoint persistence for SmartPaper executors. This module intentionally depends only on the Python standard library. It can be used immediately after a Worker claim without importing DeepSeek, RAG, statistics, document, or Streamlit modules. """ from __future__ import annotations import json import os import re from contextlib import contextmanager from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterator try: # Hugging Face production runs on Linux. import fcntl # type: ignore except ImportError: # pragma: no cover - Windows-only local development fallback. fcntl = None import msvcrt # type: ignore CLAIM_HANDSHAKE_FILENAME = "executor_claim.json" AUTHORITY_LOCK_FILENAME = ".checkpoint_authority.lock" def _atomic_json(path: Path, value: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(json.dumps(value, ensure_ascii=False), encoding="utf-8") os.replace(temporary, path) @contextmanager def checkpoint_authority_lock(task_dir: Path) -> Iterator[None]: """Serialize local claim-marker replacement and checkpoint writes. Worker/D1 remains the global authority. This small process lock only closes the local race where a superseded executor could otherwise validate an old marker immediately before a newer epoch replaces it, then write afterwards. """ lock_path = Path(task_dir) / AUTHORITY_LOCK_FILENAME lock_path.parent.mkdir(parents=True, exist_ok=True) handle = lock_path.open("a+b") try: if fcntl is not None: fcntl.flock(handle.fileno(), fcntl.LOCK_EX) else: # pragma: no cover - Windows-only local development fallback. handle.seek(0, os.SEEK_END) if handle.tell() == 0: handle.write(b"0") handle.flush() handle.seek(0) msvcrt.locking(handle.fileno(), msvcrt.LK_LOCK, 1) yield finally: try: if fcntl is not None: fcntl.flock(handle.fileno(), fcntl.LOCK_UN) else: # pragma: no cover - Windows-only local development fallback. handle.seek(0) msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) finally: handle.close() class TaskCheckpointStore: def __init__( self, task_dir: Path, request_hash: str, execution_epoch: int, *, task_id: str = "", executor_instance_id: str = "", confirmed_checkpoint_seq: int | None = None, confirmed_checkpoint_stage: str = "", ): self.task_dir = Path(task_dir) self.request_hash = str(request_hash) self.execution_epoch = int(execution_epoch) self.task_id = str(task_id or "") self.executor_instance_id = str(executor_instance_id or "") if bool(self.task_id) != bool(self.executor_instance_id): raise ValueError("checkpoint_authority_identity_incomplete") self.confirmed_checkpoint_seq = ( None if confirmed_checkpoint_seq is None else max(0, int(confirmed_checkpoint_seq)) ) self.confirmed_checkpoint_stage = str(confirmed_checkpoint_stage or "")[:120] self.directory = self.task_dir / "checkpoints" self.directory.mkdir(parents=True, exist_ok=True) self.sequence = int(self.confirmed_checkpoint_seq or 0) def _read_payload(self, path: Path) -> dict[str, Any] | None: try: payload = json.loads(path.read_text(encoding="utf-8")) except Exception: return None if not isinstance(payload, dict): return None if payload.get("request_hash") != self.request_hash: return None try: seq = int(payload.get("checkpoint_seq") or 0) writer_epoch = int(payload.get("checkpoint_execution_epoch") or 0) except (TypeError, ValueError): return None if seq <= 0 or writer_epoch <= 0 or writer_epoch > self.execution_epoch: return None return payload def load_confirmed( self, checkpoint_seq: int | None = None, checkpoint_stage: str | None = None, ) -> dict[str, Any] | None: """Load only the exact checkpoint sequence acknowledged by Worker/D1. Unacknowledged local snapshots from an older network-partitioned epoch are intentionally ignored even when they have a larger sequence number. """ seq = self.confirmed_checkpoint_seq if checkpoint_seq is None else max(0, int(checkpoint_seq)) stage = self.confirmed_checkpoint_stage if checkpoint_stage is None else str(checkpoint_stage or "")[:120] if seq is None: raise ValueError("confirmed_checkpoint_seq_required") self.sequence = int(seq) if seq <= 0: return None candidates: list[dict[str, Any]] = [] for path in self.directory.glob(f"checkpoint_{seq:05d}_*.json"): payload = self._read_payload(path) if payload is None: continue if int(payload.get("checkpoint_seq") or 0) != seq: continue if stage and str(payload.get("checkpoint_stage") or "") != stage: continue candidates.append(payload) if len(candidates) != 1: return None return candidates[0] def load_latest(self) -> dict[str, Any] | None: # Authoritative production callers always pin recovery to the exact # Worker-confirmed checkpoint. Keep the legacy scan only for isolated # compatibility/tests that do not yet supply confirmed metadata. if self.confirmed_checkpoint_seq is not None: return self.load_confirmed() candidates: list[tuple[int, Path, dict[str, Any]]] = [] for path in self.directory.glob("checkpoint_*.json"): payload = self._read_payload(path) if payload is None: continue seq = int(payload.get("checkpoint_seq") or 0) candidates.append((seq, path, payload)) if not candidates: return None seq, _, payload = max(candidates, key=lambda item: item[0]) self.sequence = seq return payload def _assert_local_authority_locked(self) -> None: if not self.task_id: return marker_path = self.task_dir / CLAIM_HANDSHAKE_FILENAME try: marker = json.loads(marker_path.read_text(encoding="utf-8")) except Exception as exc: raise RuntimeError("checkpoint_authority_marker_unavailable") from exc if not isinstance(marker, dict): raise RuntimeError("checkpoint_authority_marker_invalid") expected = { "task_id": self.task_id, "request_hash": self.request_hash, "execution_epoch": self.execution_epoch, "executor_instance_id": self.executor_instance_id, } for key, value in expected.items(): actual = marker.get(key) if key == "execution_epoch": try: actual = int(actual or 0) except (TypeError, ValueError): actual = 0 else: actual = str(actual or "") if actual != value: raise RuntimeError("checkpoint_authority_lost") if self.confirmed_checkpoint_seq is not None: try: marker_base_seq = int(marker.get("worker_checkpoint_seq") or 0) except (TypeError, ValueError): marker_base_seq = -1 marker_base_stage = str(marker.get("worker_checkpoint_stage") or "")[:120] if marker_base_seq != int(self.confirmed_checkpoint_seq): raise RuntimeError("checkpoint_authority_base_mismatch") if self.confirmed_checkpoint_stage and marker_base_stage != self.confirmed_checkpoint_stage: raise RuntimeError("checkpoint_authority_base_mismatch") def save(self, stage: str, state: dict[str, Any]) -> tuple[int, str]: stage_name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(stage or "checkpoint"))[:80] with checkpoint_authority_lock(self.task_dir): self._assert_local_authority_locked() next_sequence = self.sequence + 1 payload = dict(state or {}) payload.update( { "request_hash": self.request_hash, "checkpoint_seq": next_sequence, "checkpoint_stage": stage_name, "checkpoint_execution_epoch": self.execution_epoch, "checkpoint_executor_instance_id": self.executor_instance_id, "checkpoint_saved_at": datetime.now(timezone.utc).isoformat(), } ) path = self.directory / f"checkpoint_{next_sequence:05d}_{stage_name}.json" _atomic_json(path, payload) _atomic_json(self.directory / "latest.json", payload) self.sequence = next_sequence return self.sequence, stage_name