ui/gpu.py

"""
ui/gpu.py
---------
ZeroGPU wiring — the only place in NOVA that touches a GPU.

HF's ZeroGPU hardware hands a Space a GPU *only* for the duration of a call to an
@spaces.GPU-decorated function, and refuses to boot at all if it can't find one
at import time ("No @spaces.GPU function detected during startup"). That single
constraint drives the whole design here.

NOVA is CPU-first and stays that way. Exactly one operation runs on the GPU: the
bulk embedding of a paper's chunks during vectorizing, which is by far the
slowest thing in the app (tens of seconds on CPU, a few on GPU). Everything else
— SPECTER reranking, query embedding, the cross-encoder — is pinned to CPU *on
purpose*, because those run outside any GPU window and a cuda-resident model
there would fail on first use. That's why the three backend call sites now take
an explicit `device` instead of auto-detecting.

Off ZeroGPU (local, or Spaces "CPU basic") @spaces.GPU is a transparent
passthrough and ON_ZEROGPU is False, so this module quietly degrades to plain
CPU work and nothing else in the app changes.
"""

import os

import spaces

# Set by the ZeroGPU runtime; `spaces.config` reads the same variable.
ON_ZEROGPU = os.getenv("SPACES_ZERO_GPU", "").lower() in ("1", "t", "true")

# The device to use *inside* a GPU window. Outside one, always "cpu".
GPU_DEVICE = "cuda" if ON_ZEROGPU else "cpu"

# Generous but bounded. The window has to cover PDF text extraction and chunking
# (CPU work that unavoidably happens inside build_vectorstore) plus the encode
# itself. A long paper on a cold cache is the worst case.
_VECTORIZE_SECONDS = 120


@spaces.GPU(duration=_VECTORIZE_SECONDS)
def vectorize_on_gpu(pdf_path: str) -> None:
    """Build and persist this paper's vectorstore with the embedder on GPU.

    Returns None deliberately. ZeroGPU runs this in its own GPU worker, so a
    Chroma handle created here would carry a cuda-resident embedding model back
    to a caller that no longer holds the GPU — useless at best, a crash at worst.
    What crosses the boundary is the *persisted vectorstore on disk*, which is
    device-independent.

    The caller then re-opens it on CPU, which costs nothing: build_vectorstore
    short-circuits to a plain load as soon as the persist dir exists.
    """
    from vectorizeer import build_vectorstore
    build_vectorstore(pdf_path, device=GPU_DEVICE)