"""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"