58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
from threading import Lock
|
||
|
|
from typing import Dict, Optional
|
||
|
|
|
||
|
|
|
||
|
|
class StatusStore:
|
||
|
|
"""Persistent status tracker for resumable batch inference."""
|
||
|
|
|
||
|
|
def __init__(self, status_file: Path):
|
||
|
|
self.status_file = status_file
|
||
|
|
self.lock = Lock()
|
||
|
|
self._status: Dict[str, Dict[str, str]] = {}
|
||
|
|
self.status_file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
if self.status_file.exists():
|
||
|
|
try:
|
||
|
|
self._status = json.loads(self.status_file.read_text())
|
||
|
|
except Exception:
|
||
|
|
self._status = {}
|
||
|
|
|
||
|
|
def is_done(self, task_id: str) -> bool:
|
||
|
|
return self._status.get(task_id, {}).get("state") == "done"
|
||
|
|
|
||
|
|
def get(self, task_id: str) -> Optional[Dict[str, str]]:
|
||
|
|
return self._status.get(task_id)
|
||
|
|
|
||
|
|
def mark_running(self, task_id: str, detail: str) -> None:
|
||
|
|
with self.lock:
|
||
|
|
self._status[task_id] = {"state": "running", "detail": detail}
|
||
|
|
self._flush()
|
||
|
|
|
||
|
|
def mark_done(self, task_id: str, detail: str) -> None:
|
||
|
|
with self.lock:
|
||
|
|
self._status[task_id] = {"state": "done", "detail": detail}
|
||
|
|
self._flush()
|
||
|
|
|
||
|
|
def mark_failed(self, task_id: str, detail: str) -> None:
|
||
|
|
with self.lock:
|
||
|
|
self._status[task_id] = {"state": "failed", "detail": detail}
|
||
|
|
self._flush()
|
||
|
|
|
||
|
|
def summary(self) -> Dict[str, int]:
|
||
|
|
done = 0
|
||
|
|
running = 0
|
||
|
|
failed = 0
|
||
|
|
for item in self._status.values():
|
||
|
|
state = item.get("state")
|
||
|
|
if state == "done":
|
||
|
|
done += 1
|
||
|
|
elif state == "running":
|
||
|
|
running += 1
|
||
|
|
elif state == "failed":
|
||
|
|
failed += 1
|
||
|
|
return {"done": done, "running": running, "failed": failed, "total": len(self._status)}
|
||
|
|
|
||
|
|
def _flush(self) -> None:
|
||
|
|
self.status_file.write_text(json.dumps(self._status, indent=2, ensure_ascii=False))
|