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