| """Best-effort JST daily counts. No disk/network I/O on the generation path.""" |
|
|
| import json |
| import logging |
| import os |
| import threading |
| from datetime import datetime, timedelta, timezone |
| from pathlib import Path |
|
|
| LOGGER = logging.getLogger(__name__) |
| JST = timezone(timedelta(hours=9), "Asia/Tokyo") |
|
|
|
|
| class GenerationStats: |
| """Single-process writer; deliberately not a distributed accounting DB.""" |
|
|
| def __init__(self, directory=None, interval=300): |
| self.path = Path(directory) / "qwen-moss" / "daily-counts.json" if directory else None |
| self.interval = interval |
| self._counts = {} |
| self._revision = self._saved_revision = 0 |
| self._loaded = False |
| self._lock = threading.Lock() |
| self._io_lock = threading.Lock() |
| self._stop = threading.Event() |
| self._thread = None |
|
|
| def record(self, family, model, count=1, completed_at=None): |
| if type(count) is not int or count <= 0: |
| return |
| if family not in ("qwen", "moss") or not isinstance(model, str) or not model or len(model) > 256: |
| raise ValueError("invalid model identifier") |
| day = (completed_at or datetime.now(JST)).astimezone(JST).date().isoformat() |
| with self._lock: |
| models = self._counts.setdefault(day, {}).setdefault(family, {}) |
| models[model] = models.get(model, 0) + count |
| self._revision += 1 |
|
|
| def snapshot(self): |
| with self._lock: |
| return self._snapshot_unlocked() |
|
|
| def _load(self): |
| if self._loaded or self.path is None: |
| return |
| try: |
| if self.path.stat().st_size > 10_000_000: |
| raise ValueError("statistics file too large") |
| raw = json.loads(self.path.read_text(encoding="utf-8")) |
| if raw.get("version") != 1 or not isinstance(raw.get("counts"), dict): |
| raise ValueError("invalid statistics schema") |
| if raw.get("timezone") != "Asia/Tokyo": |
| raise ValueError("invalid statistics timezone") |
| restored = {} |
| for day, families in raw["counts"].items(): |
| datetime.strptime(day, "%Y-%m-%d") |
| restored[day] = {} |
| for family, models in families.items(): |
| if family not in ("qwen", "moss") or not isinstance(models, dict): |
| raise ValueError("invalid model family") |
| restored[day][family] = {} |
| for model, count in models.items(): |
| if (not isinstance(model, str) or not model or len(model) > 256 |
| or type(count) is not int or not 0 <= count <= 10**12): |
| raise ValueError("invalid model count") |
| restored[day][family][model] = count |
| except FileNotFoundError: |
| restored = {} |
| with self._lock: |
| |
| for day, families in restored.items(): |
| for family, models in families.items(): |
| target = self._counts.setdefault(day, {}).setdefault(family, {}) |
| for model, count in models.items(): |
| target[model] = target.get(model, 0) + count |
| self._loaded = True |
|
|
| def flush(self): |
| if self.path is None: |
| return |
| with self._io_lock: |
| self._load() |
| with self._lock: |
| revision = self._revision |
| if revision == self._saved_revision: |
| return |
| data = self._snapshot_unlocked() |
| |
| self.path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = self.path.with_suffix(".json.tmp") |
| temporary.write_text(json.dumps(data, ensure_ascii=False, separators=(",", ":")), encoding="utf-8") |
| os.replace(temporary, self.path) |
| with self._lock: |
| self._saved_revision = revision |
|
|
| def _snapshot_unlocked(self): |
| return {"version": 1, "timezone": "Asia/Tokyo", |
| "counts": {day: {family: dict(models) for family, models in families.items()} |
| for day, families in self._counts.items()}} |
|
|
| def start(self): |
| if self._thread is not None: |
| return |
| if self.path is None: |
| LOGGER.warning("GENERATION_STATS_DIR is unset; counts are memory-only until a persistent directory is configured.") |
| return |
|
|
| def worker(): |
| try: |
| with self._io_lock: |
| self._load() |
| except Exception: |
| LOGGER.warning("Statistics restore failed; will retry without blocking generation.", exc_info=True) |
| while not self._stop.wait(self.interval): |
| try: |
| self.flush() |
| except Exception: |
| LOGGER.warning("Statistics save failed; will retry at the next interval.", exc_info=True) |
|
|
| self._thread = threading.Thread(target=worker, name="generation-statistics", daemon=True) |
| self._thread.start() |
|
|
| def stop(self): |
| self._stop.set() |
| if self._thread is not None: |
| self._thread.join(timeout=1) |
|
|
|
|
| stats = GenerationStats(os.environ.get("GENERATION_STATS_DIR")) |
|
|
