"""Read and write access to the signal store. The store is a Hugging Face dataset repo laid out as parquet under stable paths. `local_root` is a working mirror: reads fall back to the Hub when a file is absent locally, writes stage into the mirror, and `flush()` pushes them as one commit. The rule that shapes this whole module is **frozen at issue**. An archived forecast row is written once and never rewritten. Resolution appends to a different tree; it does not revise the archive. So `append_forecast` is first-write-wins on its dedup key rather than last-write-wins, which is the opposite of what a cache would do and is the entire point: a track record whose history can be edited after the fact is not evidence of anything. """ from __future__ import annotations import hashlib import json import logging import os import threading from pathlib import Path import numpy as np import pandas as pd from . import config log = logging.getLogger("arena.store") class StoreError(RuntimeError): pass # -------------------------------------------------------------------------- # Schemas # -------------------------------------------------------------------------- # One row per (forecast, step). Flat rather than array-valued because every # consumer -- the resolver, the standings, the chart -- wants steps as rows. FORECAST_COLUMNS = [ "forecast_id", "model_slug", "asset", "timeframe", "issued_ts", "target_ts", "step", "horizon_bars", "seed", "n_samples", "context_len", "inference_version", "dispersion", "path_ref", "backfilled", ] + [f"q{int(round(q * 100)):02d}" for q in config.QUANTILE_LEVELS] QUANTILE_COLUMNS = [f"q{int(round(q * 100)):02d}" for q in config.QUANTILE_LEVELS] # Resolved outcomes. Append-only, never revised. TRACKRECORD_COLUMNS = [ "forecast_id", "model_slug", "asset", "timeframe", "issued_ts", "target_ts", "step", "horizon_bars", "realized_close", "q10", "q50", "q90", "q20", "q80", "inside_80", "inside_60", "abs_error", "pct_error", "resolved_ts", "backfilled", ] STANDINGS_COLUMNS = [ "model_slug", "asset_class", "resolved_count", "forecast_count", "coverage_80", "coverage_60", "median_abs_pct_error", "calibration_gap", "grade", "updated_ts", ] def forecast_id(model_slug: str, asset: str, timeframe: str, issued_ts, seed: int) -> str: """The dedup key, as a short stable hash. Deliberately *excludes* `inference_version`. Re-running the same request after a model upgrade must not append a second row: the first answer is the one that was issued, and the archive records what was issued. The version that produced it is kept as a column so the row is still traceable. """ payload = "|".join([ str(model_slug), str(asset), str(timeframe), _iso(issued_ts), str(int(seed)), ]) return hashlib.sha256(payload.encode()).hexdigest()[:16] def _utc(ts) -> pd.Timestamp: t = pd.Timestamp(ts) return t.tz_localize("UTC") if t.tzinfo is None else t.tz_convert("UTC") def _iso(ts) -> str: return _utc(ts).isoformat() def now_utc() -> pd.Timestamp: """A tz-aware "now". `pd.Timestamp.utcnow()` returns tz-aware on current pandas and tz-naive on older ones, so localising it unconditionally raises. Every "now" in the app comes through here so that difference is handled once. """ return _utc(pd.Timestamp.utcnow()) # -------------------------------------------------------------------------- # Store # -------------------------------------------------------------------------- class ArenaStore: """Parquet access to the Arena's trees in the signal store. Set `offline=True` (or leave `repo_id=None`) for a purely local store -- what the tests and the backfill script use before they have a token. """ def __init__(self, repo_id: str | None = config.STORE_REPO, local_root: str | os.PathLike | None = None, token: str | None = None, offline: bool = False, revision: str = "main") -> None: self.repo_id = repo_id self.revision = revision self.token = token or os.environ.get("HF_WRITE_TOKEN") or os.environ.get("HF_TOKEN") self.offline = offline or repo_id is None self.local_root = Path(local_root or os.environ.get("BIT_STORE_CACHE", ".cache/store")) self.local_root.mkdir(parents=True, exist_ok=True) self._pending: set[str] = set() self._price_cache: dict[tuple[str, str], pd.DataFrame] = {} self._flush_timer: threading.Timer | None = None self._lock = threading.RLock() # -- low-level -------------------------------------------------------- def _local(self, repo_path: str) -> Path: return self.local_root / repo_path def _fetch(self, repo_path: str) -> Path | None: local = self._local(repo_path) if local.exists(): return local if self.offline: return None try: from huggingface_hub import hf_hub_download return Path(hf_hub_download( repo_id=self.repo_id, repo_type=config.STORE_REPO_TYPE, filename=repo_path, revision=self.revision, token=self.token, )) except Exception: # Absent from the Hub is a normal "no data" answer, not a fault. return None def read_parquet(self, repo_path: str) -> pd.DataFrame | None: p = self._fetch(repo_path) if p is None: return None try: return pd.read_parquet(p) except Exception as e: # pragma: no cover log.warning("unreadable parquet %s: %s", repo_path, e) return None def read_json(self, repo_path: str) -> dict | None: p = self._fetch(repo_path) if p is None: return None try: return json.loads(Path(p).read_text()) except Exception as e: # pragma: no cover log.warning("unreadable json %s: %s", repo_path, e) return None def _stage(self, repo_path: str, write) -> None: local = self._local(repo_path) local.parent.mkdir(parents=True, exist_ok=True) write(local) with self._lock: self._pending.add(repo_path) def write_parquet(self, repo_path: str, df: pd.DataFrame) -> None: self._stage(repo_path, lambda p: df.to_parquet(p, index=False)) if repo_path.startswith("prices/"): self.invalidate_prices() def write_json(self, repo_path: str, payload: dict) -> None: self._stage(repo_path, lambda p: Path(p).write_text(json.dumps(payload, indent=2, sort_keys=True))) @property def pending(self) -> list[str]: with self._lock: return sorted(self._pending) def schedule_flush(self, message: str = "arena: archive forecasts", every: float = 30.0) -> None: """Commit staged files in the background, batched. Flushing inline put an HTTP commit on the end of every forecast -- seconds of latency on a request whose actual work took milliseconds, and one commit per click in the store's history. This coalesces: the first call starts a timer, later calls within the window join the same commit. Deliberately not `huggingface_hub.CommitScheduler`, which commits an entire folder. This mirror also holds price files pulled *down* from the Hub, and a whole-folder scheduler would keep re-committing those. Only what was explicitly staged is ever pushed. """ if self.offline: return with self._lock: if not self._pending or self._flush_timer is not None: return timer = threading.Timer(every, self._flush_now, args=(message,)) timer.daemon = True self._flush_timer = timer timer.start() def _flush_now(self, message: str) -> None: with self._lock: self._flush_timer = None try: pushed = self.flush(message) if pushed: log.info("archived %d file(s) to the store", pushed) except Exception as e: # pragma: no cover - network # The rows are already in the local mirror; the next flush retries. log.warning("could not push the archive: %s", e) def flush(self, message: str = "arena: update") -> int: """Push everything staged as one commit. Returns the file count.""" with self._lock: paths = sorted(self._pending) if not paths or self.offline: self._pending.clear() return 0 from huggingface_hub import CommitOperationAdd, HfApi ops = [CommitOperationAdd(path_in_repo=p, path_or_fileobj=str(self._local(p))) for p in paths] HfApi(token=self.token).create_commit( repo_id=self.repo_id, repo_type=config.STORE_REPO_TYPE, operations=ops, commit_message=message, revision=self.revision, ) with self._lock: self._pending.clear() return len(ops) # -- prices ----------------------------------------------------------- def get_prices(self, asset: str, timeframe: str, start=None, end=None) -> pd.DataFrame: """OHLCV slice with a `ts` column, sorted ascending. Reads the existing price cache written by the shared pipeline; the Arena never fetches a price in the request path. The full series is memoised per (asset, timeframe). Without it the backfill re-read and re-concatenated ~9,600 rows of parquet for every one of its thousands of as-of forecasts, which dominated the run; `invalidate_prices` drops the memo after a refresh writes new bars. """ key = (asset, timeframe) cached = self._price_cache.get(key) if cached is not None: return self._slice(cached, start, end) full = self._read_prices(asset, timeframe) self._price_cache[key] = full return self._slice(full, start, end) @staticmethod def _slice(frame: pd.DataFrame, start, end) -> pd.DataFrame: if start is None and end is None: return frame out = frame if start is not None: out = out[out["ts"] >= _utc(start)] if end is not None: out = out[out["ts"] <= _utc(end)] return out.reset_index(drop=True) def invalidate_prices(self, asset: str | None = None, timeframe: str | None = None) -> None: if asset is None: self._price_cache.clear() return self._price_cache.pop((asset, timeframe), None) def _read_prices(self, asset: str, timeframe: str) -> pd.DataFrame: start = end = None years = _years_between(start, end) frames = [] for y in years: df = self.read_parquet(config.prices_path(asset, timeframe, y)) if df is not None and len(df): frames.append(df) if not frames: return pd.DataFrame(columns=["ts", "open", "high", "low", "close", "volume"]) out = pd.concat(frames, ignore_index=True) out["ts"] = pd.to_datetime(out["ts"], utc=True) return out.drop_duplicates(subset=["ts"]).sort_values("ts").reset_index(drop=True) # -- forecasts (the archive) ------------------------------------------ def get_forecasts(self, model_slug: str, asset: str, timeframe: str, years: list[int] | None = None) -> pd.DataFrame: years = years or _years_between(None, None) frames = [] for y in years: df = self.read_parquet(config.signals_path(model_slug, asset, timeframe, y)) if df is not None and len(df): frames.append(df) if not frames: return pd.DataFrame(columns=FORECAST_COLUMNS) out = pd.concat(frames, ignore_index=True) out["issued_ts"] = pd.to_datetime(out["issued_ts"], utc=True) out["target_ts"] = pd.to_datetime(out["target_ts"], utc=True) return out def append_forecast(self, rows: pd.DataFrame) -> int: """Archive one forecast. Returns rows actually written (0 if a dup). First write wins. A forecast already present under this id is left exactly as it was -- see the module docstring. """ if rows is None or not len(rows): return 0 missing = [c for c in FORECAST_COLUMNS if c not in rows.columns] if missing: raise StoreError(f"forecast rows missing columns {missing}") fid = str(rows["forecast_id"].iloc[0]) if rows["forecast_id"].nunique() != 1: raise StoreError("append_forecast takes one forecast at a time") model_slug = str(rows["model_slug"].iloc[0]) asset = str(rows["asset"].iloc[0]) tf = str(rows["timeframe"].iloc[0]) year = int(_utc(rows["issued_ts"].iloc[0]).year) path = config.signals_path(model_slug, asset, tf, year) existing = self.read_parquet(path) if existing is not None and len(existing): if fid in set(existing["forecast_id"].astype(str)): return 0 merged = pd.concat([existing, rows[FORECAST_COLUMNS]], ignore_index=True) else: merged = rows[FORECAST_COLUMNS].copy() merged = merged.sort_values(["issued_ts", "step"]).reset_index(drop=True) self.write_parquet(path, merged) return len(rows) # -- track record ----------------------------------------------------- def get_trackrecord(self, model_slug: str, asset: str, timeframe: str) -> pd.DataFrame: df = self.read_parquet(config.trackrecord_path(model_slug, asset, timeframe)) if df is None or not len(df): return pd.DataFrame(columns=TRACKRECORD_COLUMNS) df["issued_ts"] = pd.to_datetime(df["issued_ts"], utc=True) df["target_ts"] = pd.to_datetime(df["target_ts"], utc=True) return df def append_trackrecord(self, model_slug: str, asset: str, timeframe: str, rows: pd.DataFrame) -> int: """Append resolved outcomes. Idempotent on (forecast_id, step).""" if rows is None or not len(rows): return 0 path = config.trackrecord_path(model_slug, asset, timeframe) existing = self.get_trackrecord(model_slug, asset, timeframe) if len(existing): seen = set(zip(existing["forecast_id"].astype(str), existing["step"].astype(int))) mask = [(str(f), int(s)) not in seen for f, s in zip(rows["forecast_id"], rows["step"])] rows = rows[np.asarray(mask)] if not len(rows): return 0 merged = pd.concat([existing, rows[TRACKRECORD_COLUMNS]], ignore_index=True) else: merged = rows[TRACKRECORD_COLUMNS].copy() merged = merged.sort_values(["issued_ts", "step"]).reset_index(drop=True) self.write_parquet(path, merged) return len(rows) def all_trackrecords(self) -> pd.DataFrame: """Every resolved row in the store, for the standings rebuild. **Enumerates the repo, not the local mirror.** Walking the mirror alone was wrong in exactly the environment that matters: a fresh CI runner has downloaded only the handful of files it happened to touch, so the rebuild saw a fraction of the track record and published standings that dropped from 16 rows to 10 -- silently, because a partial answer looks just like a complete one. Offline, the mirror is all there is, which is correct for tests. """ root = None if not self.offline: try: from huggingface_hub import snapshot_download # One parallel snapshot, not ninety sequential downloads: the # file-at-a-time version took over ten minutes on a cold # mirror and timed out. root = Path(snapshot_download( self.repo_id, repo_type=config.STORE_REPO_TYPE, revision=self.revision, token=self.token, allow_patterns=["arena/trackrecord/**"], )) / "arena" / "trackrecord" except Exception as e: # pragma: no cover - network log.warning("could not sync the track record: %s", e) root = None frames = [] if root is not None and root.exists(): for p in sorted(root.rglob("*.parquet")): try: frames.append(pd.read_parquet(p)) except Exception as e: # pragma: no cover log.warning("skipping unreadable trackrecord %s: %s", p, e) else: root = self.local_root / "arena" / "trackrecord" if not root.exists(): return pd.DataFrame(columns=TRACKRECORD_COLUMNS) for p in sorted(root.rglob("*.parquet")): try: frames.append(pd.read_parquet(p)) except Exception as e: # pragma: no cover log.warning("skipping unreadable trackrecord %s: %s", p, e) if not frames: return pd.DataFrame(columns=TRACKRECORD_COLUMNS) out = pd.concat(frames, ignore_index=True) out["issued_ts"] = pd.to_datetime(out["issued_ts"], utc=True) out["target_ts"] = pd.to_datetime(out["target_ts"], utc=True) return out # -- standings and registry ------------------------------------------- def get_standings(self) -> pd.DataFrame: df = self.read_parquet(config.STANDINGS_PATH) if df is None: return pd.DataFrame(columns=STANDINGS_COLUMNS) return df def put_standings(self, df: pd.DataFrame, force: bool = False) -> None: """Write the standings, refusing a suspicious shrink. Standings are derived, so a rebuild from an incomplete read produces a smaller table that is indistinguishable from a correct one. This is the backstop that turns that into a loud refusal instead of a quiet regression. It has already caught one. """ if not force: existing = self.get_standings() if len(existing) and len(df) < len(existing): raise StoreError( f"refusing to shrink the standings from {len(existing)} to " f"{len(df)} rows -- the track record was probably read " f"incompletely. Pass force=True if the shrink is real." ) self.write_parquet(config.STANDINGS_PATH, df[STANDINGS_COLUMNS]) def get_panels(self) -> dict: """Precomputed Track Record panels. See `trackrecord.build_panels`.""" doc = self.read_json(config.PANELS_PATH) if not doc: return {"version": 1, "panels": {}, "totals": {}} doc.setdefault("panels", {}) doc.setdefault("totals", {}) return doc def put_panels(self, doc: dict, force: bool = False) -> None: """Write the panel summaries, refusing a suspicious shrink.""" if not force: existing = self.get_panels().get("panels", {}) new = doc.get("panels", {}) if existing and len(new) < len(existing): raise StoreError( f"refusing to shrink the panel summaries from " f"{len(existing)} to {len(new)} -- the track record was " f"probably read incompletely. Pass force=True if real." ) self.write_json(config.PANELS_PATH, doc) def get_registry(self) -> dict: reg = self.read_json(config.REGISTRY_PATH) if not reg: return {"version": 1, "models": {}} reg.setdefault("models", {}) return reg def put_registry(self, registry: dict) -> None: self.write_json(config.REGISTRY_PATH, registry) def _years_between(start, end) -> list[int]: """Year partitions to look in. Defaults to a window around now.""" now = now_utc() s = _utc(start).year if start is not None else now.year - 3 e = _utc(end).year if end is not None else now.year + 1 return list(range(int(s), int(e) + 1))