"""User-funded coverage extension and add-model, on ZeroGPU.
Inference runs inside a `@spaces.GPU` function so it draws on the *signed-in
user's* ZeroGPU quota, not the Space owner's. Anonymous visitors keep full read
and backtest access; only extension is gated.
Safety properties enforced here:
* **Dedup** — a range the manifest already covers is never recomputed.
* **Caps** — per-request range limits keep one user from monopolising the queue.
* **Single writer** — a process-wide lock serialises commits, so concurrent
extensions cannot interleave and corrupt the manifest.
* **Allow-list** — only vetted adapter families load, and model ids are
validated before they reach the Hub.
"""
from __future__ import annotations
import logging
import os
import threading
from dataclasses import dataclass
import gradio as gr
import pandas as pd
from . import catalog, comparisons, config, runtime
from .adapters import AdapterError, ModelNotAllowed, build_windows, get_adapter, validate_model_id
from .store import _utc
log = logging.getLogger("bit.extension")
# One writer at a time. Commits to the store are serialised process-wide so a
# second extension cannot land between another's data write and its manifest
# update.
_WRITE_LOCK = threading.Lock()
try:
import spaces # provided by the ZeroGPU runtime
HAS_SPACES = True
except Exception: # running locally or on CPU-only hardware
spaces = None
HAS_SPACES = False
def _gpu(duration=120):
"""Apply @spaces.GPU when the runtime offers it, otherwise run on CPU."""
def deco(fn):
if HAS_SPACES:
return spaces.GPU(duration=duration)(fn)
return fn
return deco
class ExtensionError(RuntimeError):
pass
class QuotaExhausted(ExtensionError):
pass
# --------------------------------------------------------------------------
# Estimation & guardrails
# --------------------------------------------------------------------------
@dataclass
class Estimate:
model_slug: str
asset: str
timeframe: str
start: pd.Timestamp
end: pd.Timestamp
steps: int
already_covered: bool
capped: bool
cap_days: int
note: str = ""
def estimate(model_slug: str, asset: str, timeframe: str, start, end) -> Estimate:
if model_slug not in config.SEED_MODELS:
raise ExtensionError(f"unknown model {model_slug!r}")
if asset not in config.ASSETS:
raise ExtensionError(f"unknown asset {asset!r}")
if timeframe not in config.TIMEFRAMES:
raise ExtensionError(f"unknown timeframe {timeframe!r}")
try:
s, e = _utc(start), _utc(end)
except Exception as exc:
raise ExtensionError(f"could not parse the date range: {exc}") from exc
if s >= e:
raise ExtensionError("start must be before end")
cap_days = config.CAPS.max_days.get(timeframe, 365)
capped = (e - s).days > cap_days
if capped:
s = e - pd.Timedelta(days=cap_days)
store = runtime.get_store()
prices = store.get_prices(asset, timeframe, s, e)
if prices.empty:
raise ExtensionError(
f"No cached prices for {asset} {timeframe} in that range. "
"Price coverage has to exist before signals can be generated."
)
spec = config.SEED_MODELS[model_slug]
ctx = min(spec.context_len, max(64, len(prices) // 3))
steps = max(0, len(prices) - ctx)
if steps > config.CAPS.max_steps_per_run:
steps = config.CAPS.max_steps_per_run
# Dedup must compare against the range this request would actually
# *produce*, not the range the user typed. A forecast needs a full trailing
# context window, so the first producible timestamp sits `ctx` bars after
# the start of the price slice. Comparing the typed range instead would
# report an already-covered slice as uncovered and pay for it twice.
stamps, _ = build_windows(prices["close"], ctx)
if len(stamps) == 0:
produced_start, produced_end = s, e
steps = 0
else:
produced_start = _utc(stamps[0])
produced_end = _utc(stamps[min(steps, len(stamps)) - 1]) if steps else produced_start
# Revision is resolved lazily to avoid a Hub round-trip on every keystroke,
# so this matches any revision of the model.
covered = any(
ent.model_slug == model_slug and ent.asset == asset
and ent.timeframe == timeframe
and _utc(ent.start_ts) <= produced_start and _utc(ent.end_ts) >= produced_end
for ent in store.load_manifest().signals.values()
)
note = ""
if capped:
note = (f"Range trimmed to the {cap_days}-day cap for {timeframe} bars.")
return Estimate(model_slug, asset, timeframe, s, e, steps, covered,
capped, cap_days, note)
# --------------------------------------------------------------------------
# GPU inference
# --------------------------------------------------------------------------
@_gpu(duration=180)
def run_inference(model_id: str, family: str, values: list, ctx_len: int) -> dict:
"""Inference on the caller's ZeroGPU allocation.
Kept deliberately small and picklable: it takes plain values and returns
plain lists, so nothing in the store or the app leaks into the GPU worker.
"""
import numpy as np
adapter = get_adapter(family, model_id, context_len=ctx_len)
adapter.load()
windows = np.asarray(values, dtype="float32")
forecast = adapter.predict(windows)
return {
"q10": forecast.q10.tolist(),
"q50": forecast.q50.tolist(),
"q90": forecast.q90.tolist(),
"context_len": int(forecast.context_len),
"revision": adapter.resolved_revision,
"inference_version": adapter.inference_version(),
}
def _is_quota_error(exc: Exception) -> bool:
text = f"{type(exc).__name__} {exc}".lower()
return any(k in text for k in ("quota", "gpu task aborted", "exceeded",
"no gpu available", "zerogpu"))
QUOTA_FALLBACK = (
''
"Your ZeroGPU quota is exhausted, so this extension could not run. "
"The quota refills over time. If you need to run a large batch now, "
f'
duplicate this Space and run it on your '
"own hardware — the signal store is public, so a duplicate reads the same data."
"
That range is already covered — nothing was '
"recomputed. Coverage is deduplicated against the manifest.
")
if est.steps <= 0:
return ('Not enough cached price '
"history in that range to build a single context window.
")
spec = config.SEED_MODELS[model_slug]
store = runtime.get_store()
prices = store.get_prices(asset, timeframe, est.start, est.end)
close = prices["close"]
ctx = min(spec.context_len, max(64, len(prices) // 3))
if progress:
progress(0.15, desc=f"Preparing {est.steps} context windows")
stamps, windows = build_windows(close, ctx)
if len(stamps) == 0:
return ('Not enough bars for a '
"context window.
")
stamps, windows = stamps[:est.steps], windows[:est.steps]
if progress:
progress(0.35, desc=f"Running {spec.display} on your GPU quota")
try:
out = run_inference(spec.model_id, spec.family, windows.tolist(), ctx)
except Exception as exc:
if _is_quota_error(exc):
log.warning("ZeroGPU quota exhausted for %s: %s", username, exc)
return QUOTA_FALLBACK
log.exception("extension inference failed")
return (f'Inference failed: '
f"{type(exc).__name__}: {exc}
")
frame = pd.DataFrame({
"ts": stamps, "q10": out["q10"], "q50": out["q50"], "q90": out["q90"],
"context_len": out["context_len"], "inference_version": out["inference_version"],
})
if progress:
progress(0.75, desc="Committing to the signal store")
with _WRITE_LOCK:
entry = store.write_signals(
model_slug, spec.model_id, out["revision"], asset, timeframe, frame,
inference_version=out["inference_version"],
contributed_by=username,
)
if progress:
progress(0.9, desc="Regenerating comparison tables")
# Coverage just changed, so every derived view is now stale. Both are
# rebuilt inside the same write lock, before the commit, so readers
# never see new signals alongside an old leaderboard.
try:
comparisons.regenerate(store, assets=[asset])
except Exception:
log.exception("comparison regeneration failed (coverage still written)")
try:
catalog.build(store)
except Exception:
log.exception("catalog rebuild failed (coverage still written)")
oid = store.flush(f"Extend {model_slug}/{asset}/{timeframe} by @{username}")
runtime.cache_clear()
return (f'Coverage extended by @{username} — '
f"{len(frame):,} new steps for {model_slug} on {asset} {timeframe} "
f"({entry.start_ts[:10]} → {entry.end_ts[:10]})."
+ (f" Commit {str(oid)[:8]}." if oid else "")
+ "
")
def add_model(family: str, model_id: str, username: str = "anonymous") -> str:
"""Smoke-test a user-supplied model, then register it if it passes."""
try:
model_id = validate_model_id(model_id)
except AdapterError as e:
return f'{e}
'
if family not in config.ALLOWED_ADAPTER_FAMILIES:
return (f'Adapter family {family!r} '
"is not on the allow-list.
")
store = runtime.get_store()
asset, timeframe = "BTC-USD", "1d"
prices = store.get_prices(asset, timeframe)
if prices.empty:
return ('No cached prices to smoke '
"test against.
")
n = config.CAPS.smoke_test_steps
close = prices["close"]
ctx = min(512, max(64, len(close) // 3))
stamps, windows = build_windows(close, ctx)
if len(stamps) < n:
return ('Not enough history for a '
f"{n}-step smoke test.
")
stamps, windows = stamps[-n:], windows[-n:]
try:
out = run_inference(model_id, family, windows.tolist(), ctx)
except ModelNotAllowed as e:
return f'{e}
'
except Exception as exc:
if _is_quota_error(exc):
return QUOTA_FALLBACK
return (f'Smoke test failed: '
f"{type(exc).__name__}: {exc}
")
slug = model_id.split("/")[-1].lower()
frame = pd.DataFrame({
"ts": stamps, "q10": out["q10"], "q50": out["q50"], "q90": out["q90"],
"context_len": out["context_len"], "inference_version": out["inference_version"],
})
with _WRITE_LOCK:
store.write_signals(slug, model_id, out["revision"], asset, timeframe, frame,
inference_version=out["inference_version"],
contributed_by=username)
store.flush(f"Add model {model_id} (smoke test) by @{username}")
runtime.cache_clear()
return (f'Smoke test passed: {model_id} '
f"(revision {str(out['revision'])[:8]}) produced {n} "
f"schema-valid steps and now appears in the coverage map as "
f"{slug}.
")
# --------------------------------------------------------------------------
# Gradio bindings
# --------------------------------------------------------------------------
def status_html() -> str:
if HAS_SPACES:
return ('ZEROGPU AVAILABLE · EXTENSION RUNS ON '
"YOUR OWN QUOTA WHEN SIGNED IN
")
return ('This Space is running on CPU, so coverage '
"extension is disabled. Reading and backtesting the existing store "
"works normally.
")
def _username(profile) -> str | None:
if profile is None:
return None
return getattr(profile, "username", None) or getattr(profile, "name", None)
def estimate_ui(model_slug, asset, timeframe, start, end):
if not (model_slug and asset and timeframe and start and end):
return 'PICK A MODEL, ASSET, TIMEFRAME AND RANGE
'
try:
est = estimate(model_slug, asset, timeframe, start, end)
except ExtensionError as e:
return f'{e}
'
if est.already_covered:
return ('Already covered — running this would '
"recompute nothing. Pick a wider range.
")
return (f'About {est.steps:,} inference steps for '
f"{est.start.date()} → {est.end.date()}. "
+ (est.note + " " if est.note else "")
+ "This runs on your ZeroGPU quota once you sign in.
")
def extend_ui(model_slug, asset, timeframe, start, end,
profile: gr.OAuthProfile | None = None, progress=None):
user = _username(profile)
if user is None:
return ('Sign in with Hugging Face to '
"extend coverage. Reading and backtesting stay open to everyone; "
"extension spends your own GPU quota, so it needs an account.
",
runtime.coverage_frame())
if not HAS_SPACES:
return (status_html(), runtime.coverage_frame())
try:
html = extend_coverage(model_slug, asset, timeframe, start, end,
username=user, progress=progress)
except ExtensionError as e:
html = f'{e}
'
except Exception as e:
log.exception("extend failed")
html = f'{type(e).__name__}: {e}
'
return html, runtime.coverage_frame()
def add_model_ui(family, model_id, profile: gr.OAuthProfile | None = None):
user = _username(profile)
if user is None:
return ('Sign in with Hugging Face to '
"add a model — the smoke test runs on your GPU quota.
",
runtime.coverage_frame())
if not HAS_SPACES:
return (status_html(), runtime.coverage_frame())
try:
html = add_model(family, model_id, username=user)
except Exception as e:
log.exception("add model failed")
html = f'{type(e).__name__}: {e}
'
return html, runtime.coverage_frame()