"""The adapter interface every forecasting model is reached through. Three ideas carry most of the weight here. **Capabilities are declared, not inferred.** An adapter says what kind of output it produces, what hardware it needs, and how much context it can take. The UI renders strictly from that declaration, so there is no place outside this package where a model is special-cased by name. Adding a family is a new module plus a registry entry, never an `if model == ...` in a renderer. **Seeding is part of the interface.** `predict` takes a seed and is required to be bit-reproducible under it. Some upstream models expose no seed argument at all; those are wrapped so that their sampling is seeded anyway. A forecast that cannot be reproduced cannot be audited, and an archive of unauditable forecasts is not a track record. **Contexts are validated, not trusted.** `check_context` runs on every call. It is the structural guarantee behind the no-lookahead rule: a context bar dated after the issue timestamp raises rather than quietly producing a forecast that would look brilliant. """ from __future__ import annotations import hashlib import logging from abc import ABC, abstractmethod from dataclasses import dataclass, field import numpy as np import pandas as pd from .. import config log = logging.getLogger("arena.adapters") OHLCV_COLUMNS = ("open", "high", "low", "close", "volume") # Output kinds. `ohlcv_paths` models emit whole sampled candlestick paths and # can drive the ghost-path and dispersion views; `quantile_line` models emit a # band on the close and nothing else. OUTPUT_QUANTILE_LINE = "quantile_line" OUTPUT_OHLCV_PATHS = "ohlcv_paths" OUTPUT_KINDS = (OUTPUT_QUANTILE_LINE, OUTPUT_OHLCV_PATHS) HARDWARE_CPU = "cpu" HARDWARE_GPU = "gpu" HARDWARE_TIERS = (HARDWARE_CPU, HARDWARE_GPU) class AdapterError(RuntimeError): """Anything an adapter refuses to do.""" class ModelNotAllowed(AdapterError): """The requested adapter family is not on the allow-list.""" class LookaheadError(AdapterError): """A context window reached past the moment the forecast is issued.""" class ContextError(AdapterError): """The context window is unusable -- gaps, NaNs, or too short.""" # -------------------------------------------------------------------------- # Capabilities # -------------------------------------------------------------------------- @dataclass(frozen=True) class Capabilities: """What a model can do, as the model itself declares it. `asset_generality` is a plain-language claim about training scope, not a quality score: "financial" means the model was pre-trained on market data, "general" means arbitrary time series. It drives one honest caveat in the UI and nothing else. """ output: str hardware: str max_context: int asset_generality: str = "general" seedable_natively: bool = True def __post_init__(self): if self.output not in OUTPUT_KINDS: raise AdapterError(f"unknown output kind {self.output!r}") if self.hardware not in HARDWARE_TIERS: raise AdapterError(f"unknown hardware tier {self.hardware!r}") if self.max_context < 1: raise AdapterError("max_context must be positive") @property def emits_paths(self) -> bool: return self.output == OUTPUT_OHLCV_PATHS def as_dict(self) -> dict: return { "output": self.output, "hardware": self.hardware, "max_context": self.max_context, "asset_generality": self.asset_generality, "seedable_natively": self.seedable_natively, } # -------------------------------------------------------------------------- # Result # -------------------------------------------------------------------------- @dataclass class ForecastResult: """One multi-step forecast. `quantiles` is (horizon, n_levels) over the *close* price, always present. `paths` is (n_samples, horizon, 5) of OHLCV and is present only for `ohlcv_paths` models -- the UI checks the capability, never this field, so that a model which merely happens to return paths cannot change how it is rendered. """ quantiles: np.ndarray levels: tuple[float, ...] horizon: int context_len: int inference_version: str seed: int n_samples: int paths: np.ndarray | None = None def __post_init__(self): self.quantiles = np.asarray(self.quantiles, dtype="float64") if self.quantiles.shape != (self.horizon, len(self.levels)): raise AdapterError( f"quantiles shape {self.quantiles.shape} != " f"{(self.horizon, len(self.levels))}" ) if not np.all(np.isfinite(self.quantiles)): raise AdapterError("forecast contains non-finite values") # Quantiles must not cross. Sorting along the level axis is the honest # repair for the tiny numerical inversions that sampling produces, and # it is what makes the coverage arithmetic downstream well-defined. self.quantiles = np.sort(self.quantiles, axis=1) if self.paths is not None: self.paths = np.asarray(self.paths, dtype="float64") if self.paths.ndim != 3 or self.paths.shape[1] != self.horizon: raise AdapterError(f"paths shape {self.paths.shape} is not (n, {self.horizon}, k)") def level_index(self, level: float) -> int: for i, lv in enumerate(self.levels): if abs(lv - level) < 1e-9: return i raise AdapterError(f"level {level} not in {self.levels}") def band(self, low: float = 0.1, high: float = 0.9): return (self.quantiles[:, self.level_index(low)], self.quantiles[:, self.level_index(high)]) def median(self) -> np.ndarray: return self.quantiles[:, self.level_index(0.5)] def dispersion(self) -> np.ndarray: """Per-step spread, as a fraction of the median. Sampled-path standard deviation where paths exist, band half-width otherwise. Both are a width; they are not the same statistic, and the UI labels which one it is showing. """ med = np.abs(self.median()) + 1e-12 if self.paths is not None: close = self.paths[:, :, OHLCV_COLUMNS.index("close")] return close.std(axis=0) / med lo, hi = self.band() return (hi - lo) / 2.0 / med # -------------------------------------------------------------------------- # Context validation # -------------------------------------------------------------------------- def check_context(context: pd.DataFrame, issued_ts: pd.Timestamp | None = None, min_len: int = 32) -> pd.DataFrame: """Validate a context window, or raise. This is the structural half of the no-lookahead guarantee. It is not a convention that callers are asked to honour: every adapter runs it on every call, so a forecast issued from data it should not have seen fails loudly at the point of use. """ if not isinstance(context, pd.DataFrame): raise ContextError("context must be a DataFrame") missing = [c for c in OHLCV_COLUMNS if c not in context.columns] if missing: raise ContextError(f"context is missing columns {missing}") if len(context) < min_len: raise ContextError(f"context has {len(context)} bars, need at least {min_len}") if "ts" not in context.columns: raise ContextError("context must carry a 'ts' column") ts = pd.to_datetime(context["ts"], utc=True) if ts.isna().any(): raise ContextError("context has unparseable timestamps") if not ts.is_monotonic_increasing: raise ContextError("context timestamps are not sorted ascending") if ts.duplicated().any(): raise ContextError("context has duplicate timestamps") prices = context[list(OHLCV_COLUMNS)] if not np.isfinite(prices.to_numpy(dtype="float64")).all(): raise ContextError("context contains NaN or infinite values") if (context[["open", "high", "low", "close"]].to_numpy(dtype="float64") <= 0).any(): raise ContextError("context contains non-positive prices") # A hole in the cache means the model sees two bars as adjacent when they # are weeks apart, and produces a confident wrong answer from it. # # But "irregular spacing" is not the same as "missing data". An hourly # equity series closes for ~17 hours every night and ~65 hours every # weekend; those gaps are the market, not the cache. Judging them against # the modal bar rejected every SPY and NVDA hourly context outright. # # So the threshold adapts to the series' own gap distribution: a session # boundary recurs and therefore sits inside the 99th percentile, while a # genuinely absent period stands outside it. if len(ts) > 20: deltas = ts.diff().dropna() modal = deltas.mode() if len(modal): step = modal.iloc[0] p99 = deltas.quantile(0.99) # Any *single* gap far beyond the series' own worst regular one is # a hole, however few there are. This is the check that survives a # series where absence has become the pattern -- an adaptive # threshold alone would quietly normalise that. hard = max(step * 24, p99 * 5) worst = deltas.max() if worst > hard: raise ContextError( f"context has a {worst} gap, far beyond its own session " f"pattern; the cache is incomplete for this window" ) # And an accumulation of smaller anomalies is a hole too. ceiling = max(step * 4, p99 * 1.5) bad = deltas[deltas > ceiling] if len(bad) > max(2, len(ts) // 50): raise ContextError( f"context has {len(bad)} gaps beyond its own session pattern " f"(> {ceiling}); the cache is incomplete for this window" ) if issued_ts is not None: issued = pd.Timestamp(issued_ts) if issued.tzinfo is None: issued = issued.tz_localize("UTC") last = ts.iloc[-1] if last > issued: raise LookaheadError( f"context ends at {last.isoformat()}, after issued_ts " f"{issued.isoformat()}: a forecast may only see data at or " f"before the moment it is issued" ) return context # -------------------------------------------------------------------------- # Determinism # -------------------------------------------------------------------------- def seed_everything(seed: int) -> None: """Pin every RNG an adapter might reach for. Seeding numpy alone is not enough: the sampling models draw through torch, and Kronos in particular calls `torch.multinomial` with no seed argument of its own. This is what makes "same inputs and seed produce bit-identical output" true rather than aspirational. """ seed = int(seed) % (2 ** 31 - 1) np.random.seed(seed) try: import random random.seed(seed) except Exception: # pragma: no cover pass try: import torch torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) except ImportError: pass # -------------------------------------------------------------------------- # The interface # -------------------------------------------------------------------------- class ForecastAdapter(ABC): """Uniform interface over multi-step probabilistic forecasters.""" family: str = "base" # Bumped when this adapter's own logic changes the numbers it produces. adapter_version: str = "1" def __init__(self, model_id: str, revision: str | None = None, device: str | None = None): self.model_id = model_id self.revision = revision self._device = device self._model = None self._resolved_revision: str | None = None @property def device(self) -> str: """Resolved on first use, never at construction. On ZeroGPU, `torch.cuda.is_available()` raises unless it is called inside a `@spaces.GPU` function -- so probing the device eagerly made merely *constructing* an adapter fatal, including for CPU-tier models that never wanted a GPU. Deferring it means construction is free and the probe happens inside the GPU call, where it is legal. """ if self._device is None: self._device = default_device() return self._device # -- interface -------------------------------------------------------- @abstractmethod def load(self, model_id: str | None = None, revision: str | None = None) -> "ForecastAdapter": """Materialise the model. Idempotent.""" @abstractmethod def predict(self, context_ohlcv: pd.DataFrame, horizon: int, n_samples: int = config.DEFAULT_N_SAMPLES, seed: int = 0, issued_ts: pd.Timestamp | None = None) -> ForecastResult: """Forecast `horizon` bars ahead from `context_ohlcv`. Must be bit-reproducible in `seed`. """ @abstractmethod def capabilities(self) -> Capabilities: """What this model can do. The UI renders from this and nothing else.""" # -- shared ----------------------------------------------------------- def component_versions(self) -> dict[str, str]: """Every external thing whose version changes the output. Overridden by families that load a companion repo (Kronos ships its tokenizer separately) or depend on an inference package whose version moves the numbers. """ return {"model": f"{self.model_id}@{self.resolved_revision}"} def inference_version(self) -> str: """Identity of everything that determines the output values. A short hash rather than the full component list, because it is written onto every archived row; the components themselves are recorded once in the registry so the hash can always be expanded. """ parts = [config.INFERENCE_VERSION, self.family, self.adapter_version] parts += [f"{k}={v}" for k, v in sorted(self.component_versions().items())] digest = hashlib.sha256("|".join(parts).encode()).hexdigest()[:12] return f"{config.INFERENCE_VERSION}+{self.family}.{digest}" @property def resolved_revision(self) -> str: return self._resolved_revision or self.revision or "unpinned" def resolve_revision(self) -> str: """Pin the model to an immutable commit sha before any inference runs. A floating `main` would mean two forecasts a week apart are not comparable and neither is reproducible, which defeats the archive. """ if self._resolved_revision: return self._resolved_revision from huggingface_hub import HfApi info = HfApi().model_info(self.model_id, revision=self.revision) self._resolved_revision = info.sha return self._resolved_revision # -- helpers for subclasses ------------------------------------------- def _trim(self, context: pd.DataFrame) -> pd.DataFrame: """Cut a context down to what the model can actually attend over.""" cap = self.capabilities().max_context return context.iloc[-cap:] if len(context) > cap else context @staticmethod def _quantiles_from_paths(paths_close: np.ndarray, levels=config.QUANTILE_LEVELS) -> np.ndarray: """(n_samples, horizon) -> (horizon, n_levels).""" return np.quantile(paths_close, list(levels), axis=0).T def default_device() -> str: """The device to run on, when the caller has not said. On ZeroGPU this never probes. `torch.cuda.is_available()` triggers a low-level CUDA init that ZeroGPU forbids outside a `@spaces.GPU` function, and it does not fail politely -- it takes the load down. Inside a GPU function the device is known to be cuda anyway, so `runtime` passes it explicitly and this is only the fallback for everywhere else. """ try: import spaces # noqa: F401 - presence is the signal return "cpu" except ImportError: pass try: import torch if torch.cuda.is_available(): return "cuda" except Exception: pass return "cpu"