File size: 3,767 Bytes
c641d5f | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """Small atomic JSON cache used for answer resumability."""
from __future__ import annotations
import hashlib
import json
import os
import threading
from pathlib import Path
from typing import Any
class AnswerCache:
def __init__(self, path: Path, enabled: bool = True):
self.path = path
self.enabled = enabled
self._lock = threading.RLock()
@staticmethod
def key(task: dict[str, Any], file_path: Path | None = None) -> str:
digest = hashlib.sha256()
digest.update(str(task.get("task_id", "")).encode())
digest.update(b"\0")
digest.update(str(task.get("question", "")).encode())
if file_path and file_path.exists():
digest.update(b"\0")
with file_path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _read(self) -> dict[str, Any]:
if not self.enabled or not self.path.exists():
return {}
try:
data = json.loads(self.path.read_text(encoding="utf-8"))
return data if isinstance(data, dict) else {}
except (OSError, json.JSONDecodeError):
return {}
def _write(self, data: dict[str, Any]) -> None:
if not self.enabled:
return
self.path.parent.mkdir(parents=True, exist_ok=True)
temporary = self.path.with_suffix(self.path.suffix + ".tmp")
temporary.write_text(
json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True),
encoding="utf-8",
)
os.replace(temporary, self.path)
def get(self, key: str) -> dict[str, Any] | None:
with self._lock:
value = self._read().get(key)
return value if isinstance(value, dict) else None
def put(self, key: str, value: dict[str, Any]) -> None:
with self._lock:
data = self._read()
data[key] = value
self._write(data)
def save_run(self, results: list[dict[str, Any]]) -> None:
with self._lock:
self._write({"results": results})
def load_run(self) -> list[dict[str, Any]]:
with self._lock:
value = self._read().get("results", [])
return value if isinstance(value, list) else []
class ResultStore:
"""Atomic task-id-keyed checkpoints for evaluation and reruns."""
def __init__(self, path: Path):
self.path = path
self._cache = AnswerCache(path, enabled=True)
def load(self) -> dict[str, dict[str, Any]]:
with self._cache._lock:
data = self._cache._read()
results = data.get("results", {})
return results if isinstance(results, dict) else {}
def save_result(self, task_id: str, result: dict[str, Any]) -> None:
with self._cache._lock:
data = self._cache._read()
results = data.get("results", {})
if not isinstance(results, dict):
results = {}
results[str(task_id)] = result
data["results"] = results
self._cache._write(data)
def remove(self, task_id: str) -> None:
with self._cache._lock:
data = self._cache._read()
results = data.get("results", {})
if isinstance(results, dict):
results.pop(str(task_id), None)
data["results"] = results
self._cache._write(data)
def ordered(self, task_ids: list[str] | None = None) -> list[dict[str, Any]]:
results = self.load()
if task_ids is None:
return list(results.values())
return [results[task_id] for task_id in task_ids if task_id in results]
|