Spaces:
Sleeping
Sleeping
File size: 2,344 Bytes
5db30ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | from __future__ import annotations
import asyncio
import json
from pathlib import Path
from typing import Any
class CheckpointStore:
def __init__(self, path: Path, *, flush_interval: int = 1) -> None:
self.path = path
self.path.parent.mkdir(parents=True, exist_ok=True)
self._lock = asyncio.Lock()
self._flush_interval = max(1, int(flush_interval))
self._dirty_operations = 0
self.completed_ids: set[str] = set()
self.failed_ids: set[str] = set()
self.retry_counts: dict[str, int] = {}
def load(self) -> None:
if not self.path.exists():
return
data = json.loads(self.path.read_text(encoding="utf-8"))
self.completed_ids = set(data.get("completed_ids", []))
self.failed_ids = set(data.get("failed_ids", []))
self.retry_counts = {str(k): int(v) for k, v in data.get("retry_counts", {}).items()}
async def mark_completed(self, sample_key: str) -> None:
self.completed_ids.add(sample_key)
self.failed_ids.discard(sample_key)
await self._save_if_needed()
async def mark_failed(self, sample_key: str) -> None:
self.failed_ids.add(sample_key)
await self._save_if_needed()
async def set_retry_count(self, sample_key: str, retry_count: int) -> None:
self.retry_counts[sample_key] = retry_count
await self._save_if_needed()
def is_completed(self, sample_key: str) -> bool:
return sample_key in self.completed_ids
async def flush(self) -> None:
await self._save_if_needed(force=True)
async def save(self) -> None:
async with self._lock:
payload: dict[str, Any] = {
"completed_ids": sorted(self.completed_ids),
"failed_ids": sorted(self.failed_ids),
"retry_counts": self.retry_counts,
}
temp_path = self.path.with_suffix(self.path.suffix + ".tmp")
temp_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
temp_path.replace(self.path)
self._dirty_operations = 0
async def _save_if_needed(self, *, force: bool = False) -> None:
self._dirty_operations += 1
if force or self._dirty_operations >= self._flush_interval:
await self.save()
|