tests/test_runtime.py

"""Enrollment, the no-lookahead guarantee at the runtime layer, and tiering.

These cover the parts of the acceptance list that live between the adapters and
the UI: that enrollment is idempotent and refuses anything it cannot load
safely, that `load_context` cannot hand a model a bar from the future, and that
a CPU-tier model either meets its budget or is recorded as demoted.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from src import config, runtime
from src.adapters import LookaheadError
from tests.fixture import synth


@pytest.fixture
def seeded_store(store):
    """A store with prices and one enrolled baseline."""
    bars = synth(600)
    for year, chunk in bars.groupby(bars["ts"].dt.year):
        store.write_parquet(config.prices_path("BTC-USD", "1h", int(year)), chunk)
    outcome = runtime.enroll(store, "baseline", "baseline/random-walk",
                             enrolled_by="tests")
    assert outcome.ok, outcome.message
    return store


# --------------------------------------------------------------------------
# 5. Enrollment is idempotent, and refuses what it cannot vouch for
# --------------------------------------------------------------------------


def test_enrolling_the_same_revision_twice_is_a_no_op(store):
    first = runtime.enroll(store, "baseline", "baseline/drift", enrolled_by="a")
    registry = store.get_registry()
    second = runtime.enroll(store, "baseline", "baseline/drift", enrolled_by="b",
                            registry=registry)

    assert first.ok and second.ok
    assert second.already is True
    assert len(registry["models"]) == 1
    # The second attempt must not have rewritten who enrolled it.
    assert registry["models"]["drift"]["enrolled_by"] == "a"


def test_unsupported_family_is_refused_with_a_clear_message(store):
    outcome = runtime.enroll(store, "llama", "meta-llama/Llama-3-8B")
    assert not outcome.ok
    assert "not a supported adapter family" in outcome.message
    assert "baseline" in outcome.message      # names what *is* supported


def test_a_malformed_model_id_is_refused(store):
    for bad in ("not-a-model-id", "../../etc/passwd", "a/b/c", ""):
        outcome = runtime.enroll(store, "baseline", bad)
        assert not outcome.ok, f"{bad!r} was accepted"


def test_enrolling_under_the_wrong_family_is_refused(store):
    """A known Chronos model must not be loaded through the Kronos loader."""
    outcome = runtime.enroll(store, "kronos", "amazon/chronos-bolt-tiny")
    assert not outcome.ok
    assert "chronos" in outcome.message


def test_enrollment_records_what_it_verified(store):
    runtime.enroll(store, "baseline", "baseline/bootstrap", enrolled_by="tests")
    entry = store.get_registry()["models"]["bootstrap"]

    assert entry["revision"]                       # pinned
    assert entry["smoke_test"]["ok"] is True       # actually ran
    assert entry["capabilities"]["output"] == "ohlcv_paths"
    assert entry["enrolled_by"] == "tests"
    assert entry["inference_version"]


# --------------------------------------------------------------------------
# 2. No lookahead, at the runtime layer
# --------------------------------------------------------------------------


def test_load_context_never_returns_a_bar_after_the_cut(seeded_store):
    prices = seeded_store.get_prices("BTC-USD", "1h")
    cut = pd.to_datetime(prices["ts"], utc=True).iloc[300]

    context = runtime.load_context(seeded_store, "BTC-USD", "1h", as_of=cut)
    assert pd.to_datetime(context["ts"], utc=True).max() <= cut


def test_an_as_of_forecast_is_issued_from_the_cut_not_the_clock(seeded_store):
    prices = seeded_store.get_prices("BTC-USD", "1h")
    cut = pd.to_datetime(prices["ts"], utc=True).iloc[400]

    run = runtime.run_forecast(seeded_store, "random-walk", "BTC-USD", "1h",
                               horizon=6, seed=1, as_of=cut, archive=False)
    assert run.issued_ts <= cut
    assert run.target_ts[0] > cut


def test_a_context_reaching_past_the_issue_moment_raises(seeded_store):
    """The adapter re-checks, so a caller cannot smuggle one past it."""
    from src.adapters import get_adapter

    context = runtime.load_context(seeded_store, "BTC-USD", "1h")
    stale = pd.to_datetime(context["ts"], utc=True).iloc[-10]
    adapter = get_adapter("baseline", "baseline/random-walk").load()

    with pytest.raises(LookaheadError):
        adapter.predict(context, horizon=4, seed=0, issued_ts=stale)


def test_forecast_targets_start_after_the_context_ends(seeded_store):
    run = runtime.run_forecast(seeded_store, "random-walk", "BTC-USD", "1h",
                               horizon=8, seed=3, archive=False)
    assert run.target_ts[0] > run.issued_ts
    assert len(run.target_ts) == 8


# --------------------------------------------------------------------------
# 7. Tiering
# --------------------------------------------------------------------------


def test_a_cpu_tier_model_either_meets_the_budget_or_is_recorded_as_demoted(store):
    """The rule the bootstrap script applies, checked directly.

    The registry on the Hub is written by a measurement run; this asserts the
    rule that run enforces, so a change to the thresholds cannot silently let a
    model sit in the CPU tier while missing its budget.
    """
    from scripts.bootstrap_registry import measure

    context = synth(config.DEFAULT_CONTEXT_BARS)
    latency = measure("baseline", "baseline/random-walk", context)

    assert latency["cold_s"] <= config.CPU_BUDGET_COLD_S
    assert latency["warm_s"] <= config.CPU_BUDGET_WARM_S
    assert latency["hardware"] == "cpu"
    assert latency["demoted"] is False
    # The measurement is meaningless without the machine it was taken on.
    assert latency["machine"]


def test_the_demotion_rule_fires_when_a_budget_is_missed(monkeypatch):
    """A CPU-declared model that is too slow must come back as gpu."""
    from scripts import bootstrap_registry

    context = synth(config.DEFAULT_CONTEXT_BARS)
    monkeypatch.setattr(config, "CPU_BUDGET_WARM_S", 0.0)
    latency = bootstrap_registry.measure("baseline", "baseline/drift", context)

    assert latency["declared_hardware"] == "cpu"
    assert latency["hardware"] == "gpu"
    assert latency["demoted"] is True


# --------------------------------------------------------------------------
# Anonymous session cap
# --------------------------------------------------------------------------


def test_the_session_cap_is_a_real_number():
    assert config.ANON_SESSION_CAP > 0


# --------------------------------------------------------------------------
# ZeroGPU: constructing an adapter must not touch CUDA
# --------------------------------------------------------------------------


def test_constructing_an_adapter_never_probes_the_device(monkeypatch):
    """Regression: on ZeroGPU an out-of-context CUDA probe is fatal.

    `torch.cuda.is_available()` raises unless it is called inside a
    `@spaces.GPU` function. Probing eagerly in `__init__` made merely building
    an adapter fail -- including for CPU-tier models that never wanted a GPU,
    which took the whole forecast endpoint down on ZeroGPU.
    """
    from src.adapters import base, get_adapter

    def explode():
        raise AssertionError("device was probed at construction time")

    monkeypatch.setattr(base, "default_device", explode)

    for family, model_id in (("baseline", "baseline/random-walk"),
                             ("chronos", "amazon/chronos-bolt-tiny"),
                             ("kronos", "NeoQuasar/Kronos-mini"),
                             ("timesfm", "google/timesfm-2.5-200m-pytorch")):
        adapter = get_adapter(family, model_id)
        # Capabilities must be readable without a device, because that is how
        # `run_forecast` decides whether a GPU is even needed.
        caps = adapter.capabilities()
        assert caps.hardware in ("cpu", "gpu")
        assert caps.output in ("quantile_line", "ohlcv_paths")


def test_capabilities_are_readable_without_loading_weights():
    from src.adapters import get_adapter

    adapter = get_adapter("kronos", "NeoQuasar/Kronos-base")
    caps = adapter.capabilities()
    assert adapter._model is None, "capabilities() loaded the model"
    assert caps.output == "ohlcv_paths"
    assert caps.hardware == "gpu"