Spaces:
Running on Zero
Running on Zero
| """Persistence for the arena: battle logs + Elo state, stored in a | |
| Hugging Face Storage Bucket. | |
| Everything is written under `hf://buckets/<BUCKET_ID>/`: | |
| state/elo.json current ratings (rewritten per vote) | |
| generations/<arena>/<date>/<id>.json text, settings and both models | |
| votes/<arena>/<date>/<id>.json the vote for that battle + Elo deltas | |
| The bucket holds the logs and, under models/, the weights of every contender; every prompt and vote is readable by anyone who can read the bucket. | |
| Writes go through a single background worker so the UI never blocks on network | |
| I/O and so state/elo.json is never written by two threads at once. | |
| The HF token is read from the HF_TOKEN environment variable (a Space secret). | |
| It is only ever handed to huggingface_hub and is never logged or rendered. | |
| """ | |
| import json | |
| import os | |
| import queue | |
| import tempfile | |
| import threading | |
| import traceback | |
| from datetime import datetime, timezone | |
| BUCKET_ID = "Banaxi-Tech/tts-arena" | |
| STATE_PATH = "state/elo.json" | |
| _TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") | |
| BUCKET_URL = f"https://huggingface.co/buckets/{BUCKET_ID}" | |
| _LOCAL_MIRROR = "" # dev fallback: set to a folder to also write logs locally | |
| def utc_now() -> str: | |
| return datetime.now(timezone.utc).isoformat(timespec="milliseconds") | |
| def _api(): | |
| from huggingface_hub import HfApi | |
| return HfApi(token=_TOKEN) | |
| # ---------------------------------------------------------------- writing | |
| def _write_bytes(path: str, data: bytes) -> None: | |
| """Upload one object to the bucket. Raises on failure.""" | |
| api = _api() | |
| if hasattr(api, "batch_bucket_files"): | |
| api.batch_bucket_files(BUCKET_ID, add=[(data, path)]) | |
| return | |
| # Older huggingface_hub: fall back to the fsspec interface. | |
| from huggingface_hub import HfFileSystem | |
| fs = HfFileSystem(token=_TOKEN) | |
| with fs.open(f"hf://buckets/{BUCKET_ID}/{path}", "wb") as f: | |
| f.write(data) | |
| def _read_bytes(path: str) -> bytes: | |
| """Download one object from the bucket. Raises if it can't be read.""" | |
| from huggingface_hub import HfFileSystem | |
| fs = HfFileSystem(token=_TOKEN) | |
| uri = f"hf://buckets/{BUCKET_ID}/{path}" | |
| if not fs.exists(uri): | |
| raise FileNotFoundError(uri) | |
| with fs.open(uri, "rb") as f: | |
| return f.read() | |
| def list_paths(prefix: str) -> list[str]: | |
| """Every file under `prefix` in the bucket.""" | |
| api = _api() | |
| return [ | |
| item.path | |
| for item in api.list_bucket_tree(BUCKET_ID, prefix=prefix, recursive=True) | |
| if item.type == "file" | |
| ] | |
| def read_json_many(paths: list[str]) -> list[dict]: | |
| """Download and parse many JSON objects in one batched call.""" | |
| if not paths: | |
| return [] | |
| api = _api() | |
| out: list[dict] = [] | |
| with tempfile.TemporaryDirectory() as tmp: | |
| targets = [(p, os.path.join(tmp, p.replace("/", "__"))) for p in paths] | |
| api.download_bucket_files(BUCKET_ID, files=targets) | |
| for _, local in targets: | |
| try: | |
| with open(local, "rb") as f: | |
| out.append(json.loads(f.read())) | |
| except Exception: # one corrupt record must not sink the replay | |
| traceback.print_exc() | |
| return out | |
| def load_vote_log() -> list[dict]: | |
| """Every vote ever cast, oldest first. This is the authoritative history.""" | |
| records = read_json_many(list_paths("votes/")) | |
| records.sort(key=lambda r: r.get("created_at", "")) | |
| return records | |
| def _mirror(path: str, data: bytes) -> None: | |
| if not _LOCAL_MIRROR: | |
| return | |
| full = os.path.join(_LOCAL_MIRROR, path) | |
| os.makedirs(os.path.dirname(full), exist_ok=True) | |
| with open(full, "wb") as f: | |
| f.write(data) | |
| # ---------------------------------------------------------------- worker | |
| _q: "queue.Queue[tuple[str, bytes] | None]" = queue.Queue() | |
| _worker_started = False | |
| _worker_lock = threading.Lock() | |
| last_error: str | None = None | |
| def _worker() -> None: | |
| global last_error | |
| while True: | |
| item = _q.get() | |
| if item is None: | |
| return | |
| path, data = item | |
| try: | |
| _mirror(path, data) | |
| if _TOKEN: | |
| _write_bytes(path, data) | |
| else: | |
| last_error = "HF_TOKEN is not set — logs are not being uploaded." | |
| except Exception as exc: # keep the arena usable if the bucket is down | |
| last_error = f"{type(exc).__name__}: {exc}" | |
| traceback.print_exc() | |
| finally: | |
| _q.task_done() | |
| def _ensure_worker() -> None: | |
| global _worker_started | |
| with _worker_lock: | |
| if not _worker_started: | |
| threading.Thread(target=_worker, name="arena-uploader", daemon=True).start() | |
| _worker_started = True | |
| def enqueue_json(path: str, payload: dict) -> None: | |
| """Queue one JSON document for upload. Never raises.""" | |
| _ensure_worker() | |
| data = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8") | |
| _q.put((path, data)) | |
| def flush(timeout: float = 30.0) -> None: | |
| """Block until queued uploads are done (used by tests / shutdown).""" | |
| _ensure_worker() | |
| done = threading.Event() | |
| threading.Thread(target=lambda: (_q.join(), done.set()), daemon=True).start() | |
| done.wait(timeout) | |
| # ---------------------------------------------------------------- records | |
| def _day() -> str: | |
| return datetime.now(timezone.utc).strftime("%Y-%m-%d") | |
| def generation_path(arena: str, battle_id: str) -> str: | |
| return f"generations/{arena}/{_day()}/{battle_id}.json" | |
| def vote_path(arena: str, battle_id: str) -> str: | |
| return f"votes/{arena}/{_day()}/{battle_id}.json" | |
| def log_generation(record: dict) -> None: | |
| enqueue_json(generation_path(record["arena"], record["battle_id"]), record) | |
| def log_vote(record: dict) -> None: | |
| enqueue_json(vote_path(record["arena"], record["battle_id"]), record) | |
| def save_state(state: dict) -> None: | |
| state = dict(state, updated_at=utc_now()) | |
| enqueue_json(STATE_PATH, state) | |
| def load_state() -> dict | None: | |
| """The cached ratings blob. Returns None only if it genuinely isn't there.""" | |
| try: | |
| return json.loads(_read_bytes(STATE_PATH)) | |
| except FileNotFoundError: | |
| return None | |
| def token_configured() -> bool: | |
| return bool(_TOKEN) | |