Update ui/agents.py

ui/agents.py

@@ -9,6 +9,27 @@ equivalent, and it serves requests from a thread pool, so a naive "load if
 None" would let two simultaneous first-visitors each start a ~2 GB model load.
 Hence the explicit double-checked locking below: the lock is held across the
 load, and a second caller blocks and then sees the finished object.
+
+WHY LOADING NOW HAPPENS AT IMPORT, ON THE MAIN THREAD
+-----------------------------------------------------
+This used to load lazily: `load_agents()` from inside `demo.load`, and the
+chatbot models on a daemon thread via `warm_chatbot_models_async()`. That kept
+the boot splash responsive and got the Space green early — but it meant NOVA
+loaded *every* model on a thread that was not the main thread, which is exactly
+what ZeroGPU cannot tolerate. `spaces` fakes CUDA so nothing really initialises
+it before the GPU fork, but the TorchFunctionMode/TorchDispatchMode doing that
+work are thread-local to the thread that called `patch()` — the main thread, at
+import — while `torch.cuda.is_available()` is patched globally to return True
+for every thread. So off-main-thread model loading runs against a torch that
+claims a GPU exists with none of the interception that makes the claim safe.
+
+`preload_all()` below therefore runs at import, on the main thread, before
+Gradio ever spawns a worker. See the CUDA-VIRGINITY RULE in ui/gpu.py.
+
+The cost is honest and worth naming: the server now binds its port only after
+~2 GB of weights are resident, so the Space shows "Starting" for longer instead
+of showing NOVA's own splash. We trade a nicer boot animation for a GPU path
+that actually works.
 """
 
 import threading
@@ -19,6 +40,10 @@ _agents = None
 _chatbot_lock = threading.Lock()
 _chatbot_models = None
 
+# Set by preload_all() when a load fails, so nova_app.py can show a named error
+# on the splash instead of the Space dying at import with a raw traceback.
+BOOT_ERROR: str | None = None
+
 
 def load_agents():
     """Import the compiled LangGraph agents. Importing the search graph also
@@ -57,15 +82,37 @@ def warm_chatbot_models():
     return _chatbot_models
 
 
-def warm_chatbot_models_async():
-    """Kick the chatbot warm-up onto a daemon thread.
+def preload_all() -> "str | None":
+    """Load every model NOW, on the calling thread. Call this at import from
+    nova_app.py, which runs on the main thread.
+
+    Replaces the old pair of lazy paths (`load_agents()` inside `demo.load`, and
+    `warm_chatbot_models_async()`'s daemon thread). Both put multi-hundred-MB
+    transformer loads on non-main threads, which is the one thing a ZeroGPU
+    Space must not do — see the module docstring above and ui/gpu.py.
 
-    The Streamlit app warmed BOTH model sets behind one blocking splash, which
-    meant nobody saw a usable page until ~2 GB of weights had downloaded. Only
-    the agents are needed to act on the very first click, so we block on those
-    and let the chatbot models finish in the background — they have until the
-    user has framed an intent, run a search, and picked a paper, which is far
-    longer than the load takes. ensure_chat_ready() calls warm_chatbot_models()
-    anyway, so if the thread hasn't finished it simply blocks on the same lock.
+    Never raises. A failure here is almost always a missing API key, and the
+    old lazy path had the nice property of surfacing that as a readable message
+    on the splash rather than killing the Space at boot. That property is worth
+    keeping, so the error is stashed in BOOT_ERROR and rendered by do_boot().
     """
-    threading.Thread(target=warm_chatbot_models, daemon=True).start()
+    global BOOT_ERROR
+    current = threading.current_thread()
+    if current is not threading.main_thread():
+        # Not fatal — but it means the caller moved this off the main thread,
+        # which silently reintroduces the exact bug this function exists to
+        # prevent. Say so loudly rather than fail mysteriously an hour later.
+        print(
+            f"[preload_all] WARNING: running on thread {current.name!r}, not the main "
+            f"thread. ZeroGPU's torch patching is thread-local; models loaded here are "
+            f"not covered by it and the GPU fork may fail.",
+            flush=True,
+        )
+    try:
+        load_agents()
+        warm_chatbot_models()
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        BOOT_ERROR = f"{type(e).__name__}: {e}"
+    return BOOT_ERROR