CI deploy local

src/runtime.py

@@ -22,7 +22,7 @@ from __future__ import annotations
 
 import logging
 import time
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 
 import pandas as pd
 
@@ -52,17 +52,31 @@ _ADAPTERS: dict[tuple, ForecastAdapter] = {}
 
 
 def _device(tier: str) -> str | None:
-    """The device to hand an adapter, or None to let it decide.
-
-    On ZeroGPU the answer is never "work it out yourself": a CUDA probe outside
-    a `@spaces.GPU` function is fatal. Everywhere else -- local, CI, cpu-basic
-    -- the adapter's own detection is right.
-    """
+    """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."""
@@ -99,11 +113,11 @@ def clear_adapter_cache() -> None:
 @gpu_dispatch.gpu()
 def _predict(family: str, model_id: str, revision: str | None,
              context_ohlcv, horizon: int, n_samples: int, seed: int,
-             issued_ts) -> ForecastResult:
+             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("gpu"))
+                          device=_device(tier))
     return adapter.predict(context_ohlcv, horizon=horizon, n_samples=n_samples,
                            seed=int(seed), issued_ts=issued_ts)
 
@@ -193,6 +207,12 @@ def run_forecast(store: ArenaStore, model_slug: str, asset: str, timeframe: str,
     # 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.
@@ -211,18 +231,24 @@ def run_forecast(store: ArenaStore, model_slug: str, asset: str, timeframe: str,
             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 a CPU-tier model
-        # through it would spend someone's ZeroGPU quota to run a forecast that
-        # finishes in 20ms on the CPU that is already sitting there.
+    if gpu_dispatch.HAS_SPACES:
+        # On ZeroGPU *everything* goes through the GPU function, CPU-tier
+        # models included. Outside one, ZeroGPU runs a CUDA emulation that does
+        # not intercept every operation -- TimesFM's loader trips straight
+        # through it into a low-level init and takes the request down. Inside,
+        # every torch operation is legal, and these calls are seconds long.
+        #
+        # The tier still decides who may run what; it no longer decides where.
         result = _predict(entry["family"], entry["model_id"],
                           entry.get("revision"), context, horizon, n_samples,
-                          seed, issued_ts)
+                          seed, issued_ts, caps.hardware)
     else:
         adapter = adapter_for(entry["family"], entry["model_id"],
-                              revision=entry.get("revision"),
-                              device=_device("cpu"))
+                              revision=entry.get("revision"))
         result = adapter.predict(context, horizon=horizon, n_samples=n_samples,
                                  seed=int(seed), issued_ts=issued_ts)
     elapsed = time.time() - started