src/runtime.py

"""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(),
    }