Spaces:
Sleeping
Sleeping
| # -*- 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) | |
| 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 | |