"""Running a forecast, and enrolling a model. This is the layer the UI calls. It owns three things the UI must not: **Adapter caching.** Loading weights costs seconds; doing it per request would put every model outside its latency budget. Adapters are cached per (family, model id, revision) and reused. **The issue moment.** `issued_ts` is the timestamp of the last bar in the validated price cache, never the wall clock. That is what makes a live forecast and a backfilled one structurally the same operation -- the backfill just supplies an earlier cut. There is no code path where a forecast can see a bar it should not have. **Enrollment.** Adding a model is: validate the id, check the family is on the allow-list, pin the revision to an immutable sha, run a smoke test, write the registry entry. No user-supplied code is ever imported or executed, and an unsupported family is rejected with a message that says so. """ from __future__ import annotations import logging import time from dataclasses import dataclass, field, replace import pandas as pd from . import config, gpu as gpu_dispatch, trackrecord from .adapters import (ALLOWED_ADAPTER_FAMILIES, AdapterError, ForecastAdapter, ForecastResult, ModelNotAllowed, family_for, get_adapter, validate_model_id) from .adapters import model_slug as slug_for from .store import ArenaStore, now_utc log = logging.getLogger("arena.runtime") class ForecastUnavailable(RuntimeError): """A forecast could not be produced. Carries a UI-renderable reason.""" def __init__(self, message: str, kind: str = "load_failure"): super().__init__(message) self.kind = kind # -------------------------------------------------------------------------- # Adapter cache # -------------------------------------------------------------------------- _ADAPTERS: dict[tuple, ForecastAdapter] = {} def _device(tier: str) -> str | None: """The device to hand an adapter, or None to let it decide.""" if not gpu_dispatch.HAS_SPACES: return None return "cuda" if tier == "gpu" else "cpu" def warm(model_id: str, revision: str | None) -> None: """Pull a model's weights to local disk, outside any GPU call. Downloading is network and disk, not CUDA, so it is legal anywhere -- and doing it here means the GPU call spends its 60-second budget on compute instead of on a 400 MB download it might not finish. """ if model_id.startswith("baseline/"): return try: from huggingface_hub import snapshot_download snapshot_download(model_id, revision=revision, allow_patterns=["*.json", "*.safetensors", "*.ckpt"]) except Exception as e: # pragma: no cover # A failed prefetch is not fatal: the loader will fetch what it needs. log.info("could not prefetch %s: %s", model_id, e) def adapter_for(family: str, model_id: str, revision: str | None = None, **kwargs) -> ForecastAdapter: """A loaded adapter, cached. Raises `ForecastUnavailable` on load failure.""" key = (family, model_id, revision or "pinned", tuple(sorted(kwargs.items()))) cached = _ADAPTERS.get(key) if cached is not None: return cached try: adapter = get_adapter(family, model_id, revision=revision, **kwargs) adapter.load() except ModelNotAllowed: raise except Exception as e: # Weights that will not download or will not fit are the single most # likely runtime failure, and the UI has a designed state for it. raise ForecastUnavailable( f"{model_id} could not be loaded: {e}", kind="load_failure") from e _ADAPTERS[key] = adapter return adapter def clear_adapter_cache() -> None: _ADAPTERS.clear() # The GPU entry point. # # Defined at module scope and decorated unconditionally, because ZeroGPU scans # for `@spaces.GPU` at startup and refuses to boot a Space that declares none. # It takes plain arguments and resolves the adapter through the cache rather # than receiving one, so nothing that holds CUDA state crosses the boundary. # # Off ZeroGPU the decorator is the identity and this is an ordinary call. @gpu_dispatch.gpu() def _predict(family: str, model_id: str, revision: str | None, context_ohlcv, horizon: int, n_samples: int, seed: int, issued_ts, tier: str = "gpu") -> ForecastResult: # Inside the GPU function the device is known, so it is stated rather than # probed -- see `base.default_device` for why probing is not an option. adapter = adapter_for(family, model_id, revision=revision, device=_device(tier)) return adapter.predict(context_ohlcv, horizon=horizon, n_samples=n_samples, seed=int(seed), issued_ts=issued_ts) # -------------------------------------------------------------------------- # Forecasting # -------------------------------------------------------------------------- @dataclass class ForecastRun: """Everything the UI needs about one issued forecast.""" model_slug: str model_id: str family: str asset: str timeframe: str horizon: int issued_ts: pd.Timestamp target_ts: pd.DatetimeIndex result: ForecastResult context: pd.DataFrame forecast_id: str = "" archived_rows: int = 0 elapsed_s: float = 0.0 capabilities: dict = field(default_factory=dict) # True when this was rebuilt from the archive rather than just computed. # The UI says so: "issued 3h ago" is a different claim from "just run". from_cache: bool = False backfilled: bool = False # Bars that printed *after* the forecast was issued, up to the end of its # horizon. Only ever populated for a cached run: a forecast issued a moment # ago has nothing to show yet. This is what turns "here is an old forecast" # into "here is how that forecast is doing". realized: pd.DataFrame | None = None def future_timestamps(context: pd.DataFrame, horizon: int) -> pd.DatetimeIndex: """Continue the context's own cadence forward by `horizon` bars.""" ts = pd.to_datetime(context["ts"], utc=True) deltas = ts.diff().dropna() if not len(deltas): raise AdapterError("cannot infer cadence from a single bar") modal = deltas.mode() step = modal.iloc[0] if len(modal) else deltas.median() last = ts.iloc[-1] return pd.DatetimeIndex([last + step * (i + 1) for i in range(horizon)]) def load_context(store: ArenaStore, asset: str, timeframe: str, as_of=None, context_len: int = 512) -> pd.DataFrame: """The most recent validated bars at or before `as_of`. Slicing on `as_of` here rather than in the caller is what makes the no-lookahead guarantee structural: every context the app ever builds goes through this function, and it cannot return a bar past the cut. """ prices = store.get_prices(asset, timeframe) if not len(prices): raise ForecastUnavailable( f"no price history cached for {asset} {timeframe}", kind="no_data") if as_of is not None: cut = pd.Timestamp(as_of) cut = cut.tz_localize("UTC") if cut.tzinfo is None else cut.tz_convert("UTC") prices = prices[prices["ts"] <= cut] if len(prices) < 64: raise ForecastUnavailable( f"only {len(prices)} bars available for {asset} {timeframe}", kind="no_data") return prices.iloc[-context_len:].reset_index(drop=True) def run_forecast(store: ArenaStore, model_slug: str, asset: str, timeframe: str, horizon: int | None = None, n_samples: int = config.DEFAULT_N_SAMPLES, seed: int = 0, as_of=None, registry: dict | None = None, archive: bool = True, backfilled: bool = False) -> ForecastRun: """Issue one forecast, archive it, and return everything needed to draw it.""" registry = registry if registry is not None else store.get_registry() entry = registry.get("models", {}).get(model_slug) if entry is None: raise ForecastUnavailable( f"{model_slug} is not enrolled in the Arena", kind="not_enrolled") horizon = int(horizon or config.DEFAULT_HORIZON.get(timeframe, 24)) max_h = config.MAX_HORIZON.get(timeframe, 168) if not 1 <= horizon <= max_h: raise ForecastUnavailable( f"horizon must be between 1 and {max_h} for {timeframe}", kind="bad_request") # Capabilities without loading. Constructing an adapter touches no # weights and -- since `device` is lazy -- no CUDA, so this is safe to do # outside the GPU call. Loading here instead would initialise CUDA in a # context ZeroGPU forbids, which failed even for CPU-tier models. caps = get_adapter(entry["family"], entry["model_id"], revision=entry.get("revision")).capabilities() # The registry's recorded hardware wins over the adapter's declared # default: it is the one that was measured on real hardware, and a # demotion recorded there must actually govern who can run the model. recorded = (entry.get("capabilities") or {}).get("hardware") if recorded in ("cpu", "gpu") and recorded != caps.hardware: caps = replace(caps, hardware=recorded) # The smaller of what the model can take and what the Arena spends. See # config.DEFAULT_CONTEXT_BARS for why the ceiling is not the model's own. context = load_context(store, asset, timeframe, as_of=as_of, context_len=min(caps.max_context, config.DEFAULT_CONTEXT_BARS)) issued_ts = pd.to_datetime(context["ts"], utc=True).iloc[-1] targets = future_timestamps(context, horizon) # A GPU-tier model on hardware that has no GPU cannot be made to work by # trying: it renders its designed unavailable state instead of holding a # spinner for minutes. if caps.hardware == "gpu" and not gpu_dispatch.available(): raise ForecastUnavailable( f"{model_slug} needs GPU hardware, which this Space does not " f"currently have. CPU-tier models are unaffected.", kind="no_gpu") # Weights land on disk before the GPU clock starts. warm(entry["model_id"], entry.get("revision")) started = time.time() if caps.hardware == "gpu": # Only GPU-tier models take the GPU path. Routing CPU-tier models # through it was tried and is wrong: ZeroGPU's anonymous run limit is # exhausted in a couple of calls, so a visitor who clicked Forecast # twice on a model that runs in 20ms on CPU got locked out of the GPU # models they actually needed it for. result = _predict(entry["family"], entry["model_id"], entry.get("revision"), context, horizon, n_samples, seed, issued_ts, caps.hardware) else: adapter = adapter_for(entry["family"], entry["model_id"], revision=entry.get("revision"), device=_device("cpu")) result = adapter.predict(context, horizon=horizon, n_samples=n_samples, seed=int(seed), issued_ts=issued_ts) elapsed = time.time() - started run = ForecastRun( model_slug=model_slug, model_id=entry["model_id"], family=entry["family"], asset=asset, timeframe=timeframe, horizon=horizon, issued_ts=issued_ts, target_ts=targets, result=result, context=context, elapsed_s=elapsed, capabilities=caps.as_dict(), ) if archive: fid, written = trackrecord.archive( store, result, model_slug, asset, timeframe, issued_ts, targets, backfilled=backfilled) run.forecast_id, run.archived_rows = fid, written return run def cached_run(store: ArenaStore, model_slug: str, asset: str, timeframe: str, registry: dict | None = None) -> ForecastRun | None: """Rebuild the most recently archived forecast for a series, or None. This is what a visitor sees before pressing anything. It reads the small latest-forecast cache and the price history around the issue time -- no model is loaded, so it costs a parquet read rather than an inference. The forecast returned is the one that was issued: same quantiles, same seed, same issue timestamp. Nothing is recomputed, because recomputing it would be a *different* forecast wearing the old one's timestamp. """ latest = store.get_latest() if not len(latest): return None rows = latest[(latest["model_slug"] == model_slug) & (latest["asset"] == asset) & (latest["timeframe"] == timeframe)] if not len(rows): return None rows = rows.sort_values("step") registry = registry if registry is not None else store.get_registry() entry = registry.get("models", {}).get(model_slug, {}) caps = (entry.get("capabilities") or {}).copy() issued_ts = pd.Timestamp(rows["issued_ts"].iloc[0]) horizon = int(rows["horizon_bars"].iloc[0]) levels = tuple(config.QUANTILE_LEVELS) quantiles = rows[[f"q{int(round(q * 100)):02d}" for q in levels]] \ .to_numpy(dtype="float64") paths = None if caps.get("output") == "ohlcv_paths": stored = store.get_latest_paths(model_slug, asset, timeframe) if stored is not None and stored.shape[1] == horizon: paths = stored result = ForecastResult( quantiles=quantiles, levels=levels, horizon=horizon, context_len=int(rows["context_len"].iloc[0]), inference_version=str(rows["inference_version"].iloc[0]), seed=int(rows["seed"].iloc[0]), n_samples=int(rows["n_samples"].iloc[0]), paths=paths, ) try: context = load_context(store, asset, timeframe, as_of=issued_ts, context_len=config.DEFAULT_CONTEXT_BARS) except ForecastUnavailable: return None targets = pd.DatetimeIndex(rows["target_ts"]) # What actually happened since. The context is deliberately frozen at the # issue moment -- the model must be shown what it saw -- but the chart is # far more useful with the realised path drawn through the frozen fan, and # a landing page whose candles stopped three days ago just looks broken. realized = None try: bars = store.get_prices(asset, timeframe, start=issued_ts, end=targets.max()) bars = bars[pd.to_datetime(bars["ts"], utc=True) > issued_ts] if len(bars): realized = bars.reset_index(drop=True) except Exception as e: # pragma: no cover log.info("no realised bars for %s %s: %s", asset, timeframe, e) return ForecastRun( model_slug=model_slug, model_id=entry.get("model_id", model_slug), family=entry.get("family", ""), asset=asset, timeframe=timeframe, horizon=horizon, issued_ts=issued_ts, target_ts=targets, result=result, context=context, forecast_id=str(rows["forecast_id"].iloc[0]), archived_rows=0, elapsed_s=0.0, capabilities=caps, from_cache=True, backfilled=bool(rows["backfilled"].any()), realized=realized, ) # -------------------------------------------------------------------------- # Enrollment # -------------------------------------------------------------------------- @dataclass class Enrollment: ok: bool model_slug: str = "" message: str = "" entry: dict | None = None already: bool = False def enroll(store: ArenaStore, family: str, model_id: str, enrolled_by: str = "anonymous", registry: dict | None = None, smoke_steps: int = config.CAPS.smoke_test_steps) -> Enrollment: """Validate, pin, smoke-test and register a model. Enrolling an (id, revision) that is already registered is a no-op rather than an error: the flow is idempotent so that a double-submitted form or a retried job cannot fork the registry. """ try: model_id = validate_model_id(model_id) except AdapterError as e: return Enrollment(ok=False, message=str(e)) fam = (family or "").strip().lower() if fam not in ALLOWED_ADAPTER_FAMILIES: return Enrollment(ok=False, message=( f"'{family}' is not a supported adapter family. The Arena runs " f"models only through vetted loaders, so a model outside " f"{', '.join(ALLOWED_ADAPTER_FAMILIES)} cannot be enrolled.")) known = family_for(model_id) if known is not None and known != fam: return Enrollment(ok=False, message=( f"{model_id} is a {known} model; enrolling it as {fam} would load " f"it with the wrong loader.")) registry = registry if registry is not None else store.get_registry() slug = slug_for(model_id) try: adapter = get_adapter(fam, model_id) revision = adapter.resolve_revision() except ModelNotAllowed as e: return Enrollment(ok=False, message=str(e)) except Exception as e: return Enrollment(ok=False, message=( f"could not reach {model_id} on the Hub: {e}")) existing = registry.get("models", {}).get(slug) if existing and existing.get("revision") == revision: return Enrollment(ok=True, model_slug=slug, already=True, entry=existing, message=f"{model_id} is already enrolled at this revision.") try: smoke = _smoke_test(adapter, steps=smoke_steps) except Exception as e: return Enrollment(ok=False, message=f"{model_id} failed its smoke test: {e}") caps = adapter.capabilities() entry = { "model_slug": slug, "model_id": model_id, "family": fam, "revision": revision, "display": model_id.split("/", 1)[1], "capabilities": caps.as_dict(), "components": adapter.component_versions(), "inference_version": adapter.inference_version(), "enrolled_by": str(enrolled_by or "anonymous")[:64], "enrolled_ts": now_utc().isoformat(), "smoke_test": smoke, } registry.setdefault("models", {})[slug] = entry store.put_registry(registry) return Enrollment(ok=True, model_slug=slug, entry=entry, message=f"{model_id} enrolled as '{slug}'.") def _smoke_test(adapter: ForecastAdapter, steps: int) -> dict: """Forecast a synthetic series and check the output is usable. Synthetic rather than real prices so enrollment works before the cache has any coverage for a new asset, and so the test is identical every time. """ import numpy as np n = max(128, steps) rng = np.random.default_rng(0) close = 100.0 * np.exp(np.cumsum(rng.normal(0, 0.01, n))) frame = pd.DataFrame({ "ts": pd.date_range("2025-01-01", periods=n, freq="1h", tz="UTC"), "open": close, "high": close * 1.001, "low": close * 0.999, "close": close, "volume": 1000.0, }) horizon = min(16, max(1, steps // 8)) started = time.time() result = adapter.predict(frame, horizon=horizon, n_samples=8, seed=0) elapsed = time.time() - started lo, hi = result.band() if not (lo <= hi).all(): raise AdapterError("smoke test produced a crossed band") return { "ok": True, "steps": int(n), "horizon": int(horizon), "elapsed_s": round(elapsed, 3), "emits_paths": bool(result.paths is not None), "ran_ts": now_utc().isoformat(), }