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 datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Any | |
| 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) | |
| class TaskCheckpointStore: | |
| def __init__(self, task_dir: Path, request_hash: str, execution_epoch: int): | |
| self.task_dir = Path(task_dir) | |
| self.request_hash = str(request_hash) | |
| self.execution_epoch = int(execution_epoch) | |
| self.directory = self.task_dir / "checkpoints" | |
| self.directory.mkdir(parents=True, exist_ok=True) | |
| self.sequence = 0 | |
| def load_latest(self) -> dict[str, Any] | None: | |
| candidates: list[tuple[int, Path, dict[str, Any]]] = [] | |
| for path in self.directory.glob("checkpoint_*.json"): | |
| try: | |
| payload = json.loads(path.read_text(encoding="utf-8")) | |
| if not isinstance(payload, dict): | |
| continue | |
| if payload.get("request_hash") != self.request_hash: | |
| continue | |
| seq = int(payload.get("checkpoint_seq") or 0) | |
| if seq <= 0: | |
| continue | |
| candidates.append((seq, path, payload)) | |
| except Exception: | |
| continue | |
| if not candidates: | |
| return None | |
| seq, _, payload = max(candidates, key=lambda item: item[0]) | |
| self.sequence = seq | |
| return payload | |
| def save(self, stage: str, state: dict[str, Any]) -> tuple[int, str]: | |
| self.sequence += 1 | |
| stage_name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(stage or "checkpoint"))[:80] | |
| payload = dict(state or {}) | |
| payload.update( | |
| { | |
| "request_hash": self.request_hash, | |
| "checkpoint_seq": self.sequence, | |
| "checkpoint_stage": stage_name, | |
| "checkpoint_execution_epoch": self.execution_epoch, | |
| "checkpoint_saved_at": datetime.now(timezone.utc).isoformat(), | |
| } | |
| ) | |
| path = self.directory / f"checkpoint_{self.sequence:05d}_{stage_name}.json" | |
| _atomic_json(path, payload) | |
| _atomic_json(self.directory / "latest.json", payload) | |
| return self.sequence, stage_name | |