zerogpu scaffold: shared GPU entry point + cpu fallback

app.py

@@ -1,15 +1,15 @@
 """Oscar-1 demo Space: typed decisions (choice / score / noul) with calibrated probabilities.
 
 Layout and demo patterns adapted from convaiinnovations/laya-demo (Apache-2.0).
-This Space runs on CPU: the Oscar-1 checkpoints are small non-autoregressive encoders
-(17m / 32m), so a CPU tier is genuinely enough — RTX-5060-Ti p50 latency is 2.5 ms.
+ZeroGPU when available (premium quota), CPU fallback otherwise — the Oscar-1 checkpoints are
+small non-autoregressive encoders (17m / 32m): RTX-5060-Ti p50 latency is 2.5 ms.
 """
 import os
 
-DEVICE = os.environ.get("RL_AGENT_DEVICE", "cpu")  # "cpu" (default) or "zerogpu"
+DEVICE = os.environ.get("RL_AGENT_DEVICE", "zerogpu")   # "zerogpu" (default) or "cpu"
 ZERO = False
 if DEVICE == "zerogpu":
-    try:
+    try:  # ZeroGPU attaches a GPU only while a decorated function runs
         import spaces
         GPU = spaces.GPU(duration=30)
         os.environ["RL_AGENT_CUDA"] = "1"
@@ -23,17 +23,60 @@ if not ZERO:
         return fn
 
 import json      # noqa: E402  (everything below may import torch)
+import sys       # noqa: E402
 import time      # noqa: E402
 
+print("import order: spaces before torch =", "torch" not in sys.modules, flush=True)
+
 import gradio as gr   # noqa: E402
 
 import rl_agent_demo as D   # noqa: E402
 
 # Download and build the weights while the app boots, not on someone's first click.
-D.warmup()
+for _k in D.MODEL_REPOS:
+    try:
+        D.get_agent(_k)
+    except Exception as _e:
+        print("pre-load %s failed: %s %s" % (_k, type(_e).__name__, _e), flush=True)
+
+
+@GPU
+def _gpu_call(fn, *args):
+    """The one ZeroGPU entry point every tab shares — a single decorated function means the
+    start-up warm-up covers all of them, instead of each tab paying its own first-call attach."""
+    return fn(*args)
+
 
 _GPU_BROKEN = {"why": None}
 
+
+def gpu_call(fn, *args):
+    """Run `fn` on the GPU when one can be attached, otherwise on CPU.
+
+    ZeroGPU can fail to hand out a GPU — quota, capacity, or a torch/spaces version mismatch —
+    and it surfaces as `RuntimeError: No CUDA GPUs are available` raised inside worker_init.
+    The models run perfectly well on CPU, just slower, so fall back rather than fail,
+    and only say so once."""
+    t = time.perf_counter()
+    if ZERO and _GPU_BROKEN["why"] is None:
+        try:
+            out = _gpu_call(fn, *args)
+            print("%s: %.0f ms (gpu)" % (getattr(fn, "__name__", "call"),
+                                         (time.perf_counter() - t) * 1000), flush=True)
+            return out
+        except Exception as e:
+            _GPU_BROKEN["why"] = "%s: %s" % (type(e).__name__, str(e)[:160])
+            print("ZeroGPU unavailable (%s) - serving on CPU from here on" % _GPU_BROKEN["why"], flush=True)
+            os.environ.pop("RL_AGENT_CUDA", None)
+            t = time.perf_counter()
+    out = fn(*args)
+    print("%s: %.0f ms (cpu)" % (getattr(fn, "__name__", "call"),
+                                 (time.perf_counter() - t) * 1000), flush=True)
+    return out
+
+
+D.warmup()  # CPU warm-up: weights built at boot; ZeroGPU warm-up rides the first gpu_call
+
 MODEL_KEYS = list(D.MODEL_REPOS)
 ANSWER_COLS = ["question", "answer", "confidence"]
 
@@ -63,48 +106,48 @@ NOTE = """> **Preview checkpoints: 17M and 32M parameters**, trained with RLCD o
 
 
 def run_triage(model_key, message, tier, threshold):
-    rows, action, r = D.triage(model_key, message, tier, threshold)
+    rows, action, r = gpu_call(D.triage, model_key, message, tier, threshold)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), json.dumps(r, indent=2)
 
 
 def run_email(model_key, sender, subject, body):
-    rows, action, cleaned, r = D.email_triage(model_key, sender, subject, body)
+    rows, action, cleaned, r = gpu_call(D.email_triage, model_key, sender, subject, body)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), cleaned, json.dumps(r, indent=2)
 
 
 def run_intents(model_key, message):
-    rows, action, r = D.bank77(model_key, message)
+    rows, action, r = gpu_call(D.bank77, model_key, message)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), json.dumps(r, indent=2)
 
 
 def run_classify(model_key, text):
-    rows, action, r = D.classify(model_key, text)
+    rows, action, r = gpu_call(D.classify, model_key, text)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), json.dumps(r, indent=2)
 
 
 def run_guard(model_key, prompt, threshold):
-    rows, action, r = D.guardrail(model_key, prompt, threshold)
+    rows, action, r = gpu_call(D.guardrail, model_key, prompt, threshold)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), json.dumps(r, indent=2)
 
 
 def run_rag(model_key, query, passages, threshold):
-    table, summary = D.rag_filter(model_key, query, passages, threshold)
+    table, summary = gpu_call(D.rag_filter, model_key, query, passages, threshold)
     return table, "**%s**" % summary
 
 
 def run_mod(model_key, post, threshold):
-    rows, action, r = D.moderate(model_key, post, threshold)
+    rows, action, r = gpu_call(D.moderate, model_key, post, threshold)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), json.dumps(r, indent=2)
 
 
 def run_router(model_key, request, small, large):
-    rows, action, r = D.route_model(model_key, request, small, large)
+    rows, action, r = gpu_call(D.route_model, model_key, request, small, large)
     return rows, "**%s**  ·  %.0f ms" % (action, r["latency_ms"]), json.dumps(r, indent=2)
 
 
 def run_playground(model_key, state_text, questions_text):
     try:
-        rows, raw = D.playground(model_key, state_text, questions_text)
+        rows, raw = gpu_call(D.playground, model_key, state_text, questions_text)
         return rows, raw
     except Exception as e:
         return [], "%s: %s" % (type(e).__name__, e)

requirements.txt

@@ -4,3 +4,4 @@ transformers>=4.48,<5
 safetensors>=0.4
 huggingface_hub>=0.30
 numpy>=1.24
+spaces>=0.30