CI deploy local

src/adapters/timesfm.py

@@ -21,6 +21,8 @@ that -- which is why `_ensure_compiled` exists rather than a line in `load`.
 from __future__ import annotations
 
 import logging
+import threading
+from contextlib import contextmanager
 
 import numpy as np
 import pandas as pd
@@ -49,6 +51,43 @@ COMPILE_MAX_HORIZON = 256
 # The quantile grid TimesFM returns: column 0 is the mean, then 0.1 .. 0.9.
 QUANTILE_GRID = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
 
+# Serialises the CPU pin below, which mutates a torch global.
+_PIN_LOCK = threading.Lock()
+
+
+@contextmanager
+def _null_context():
+    yield
+
+
+@contextmanager
+def _pinned_to_cpu():
+    """Hide CUDA from TimesFM's loader for the duration of a CPU load.
+
+    `TimesFM_2p5_200M_torch.load_checkpoint` decides its own device:
+
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda:0")
+
+    There is no argument to override it. On ZeroGPU that probe runs outside a
+    `@spaces.GPU` function and trips a low-level CUDA init the emulation layer
+    cannot intercept, which fails the load outright -- so the probe has to not
+    happen. Patching the lookup TimesFM performs is narrower than making the
+    whole process CUDA-blind, and it is reverted immediately.
+    """
+    import torch
+
+    with _PIN_LOCK:
+        original_available = torch.cuda.is_available
+        original_count = torch.cuda.device_count
+        torch.cuda.is_available = lambda: False
+        torch.cuda.device_count = lambda: 1
+        try:
+            yield
+        finally:
+            torch.cuda.is_available = original_available
+            torch.cuda.device_count = original_count
+
 
 class TimesFMAdapter(ForecastAdapter):
     """`google/timesfm-2.5-*` quantile forecasters."""
@@ -116,14 +155,20 @@ class TimesFMAdapter(ForecastAdapter):
             )
 
         self.resolve_revision()
-        self._model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
-            self.model_id,
+        load_kwargs = dict(
             revision=self._resolved_revision,
             # torch.compile spends tens of seconds warming up and buys little
             # on a CPU box with batch size 1. It is the difference between
             # clearing the cold-start budget and blowing straight through it.
             torch_compile=(self.device == "cuda"),
         )
+        if self.device == "cuda":
+            self._model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+                self.model_id, **load_kwargs)
+        else:
+            with _pinned_to_cpu():
+                self._model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+                    self.model_id, **load_kwargs)
         return self
 
     def _ensure_compiled(self) -> None:
@@ -131,6 +176,11 @@ class TimesFMAdapter(ForecastAdapter):
             return
         import timesfm
 
+        with (_pinned_to_cpu() if self.device != "cuda"
+              else _null_context()):
+            self._compile(timesfm)
+
+    def _compile(self, timesfm) -> None:
         self._model.compile(timesfm.ForecastConfig(
             max_context=self._max_context,
             max_horizon=COMPILE_MAX_HORIZON,

src/runtime.py

@@ -235,20 +235,19 @@ def run_forecast(store: ArenaStore, model_slug: str, asset: str, timeframe: str,
     warm(entry["model_id"], entry.get("revision"))
 
     started = time.time()
-    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.
+    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"))
+                              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