llm/zerogpu_backend.py

"""In-process transformers backend for HF Spaces ZeroGPU.

Loaded only when PHANTOM_GRID_LLM_PROVIDER=zerogpu_transformers. Follows the
canonical ZeroGPU pattern documented at
https://huggingface.co/docs/hub/en/spaces-zerogpu :

    "Models must be placed on `cuda` at the root module level. A PyTorch CUDA
    emulation mode is enabled outside @spaces.GPU functions, allowing CUDA
    operations without a real GPU. Inside @spaces.GPU, real CUDA is used."

So we load directly to `cuda` here; ZeroGPU's emulation handles it at import,
and the real device is attached only while the @spaces.GPU function runs.
"""
from __future__ import annotations

import os
from typing import Any

_MODEL_ID = os.getenv("PHANTOM_GRID_ZEROGPU_MODEL_ID", "openbmb/MiniCPM4.1-8B")
_DEFAULT_DURATION = int(os.getenv("PHANTOM_GRID_ZEROGPU_DURATION", "90"))

# ZeroGPU's `spaces` package patches torch's CUDA at import time. We must import
# it BEFORE torch so the emulation is in place when we then load to 'cuda'.
try:
    import spaces  # type: ignore
    _spaces_gpu = spaces.GPU
    _ON_ZEROGPU = True
except ImportError:  # local dev without the spaces package
    def _spaces_gpu(*_args, **_kwargs):
        def _wrap(fn):
            return fn
        return _wrap
    _ON_ZEROGPU = False

import threading  # noqa: E402
import torch  # noqa: E402  — must come after `import spaces`

# openbmb/MiniCPM4.1-8B's modeling_minicpm.py (loaded via trust_remote_code=True)
# imports `is_torch_fx_available` from transformers.utils.import_utils, which
# was removed in transformers >= 5.0. Inject a shim BEFORE we touch
# AutoModel/AutoTokenizer so the trust_remote_code import succeeds.
# torch.fx has been built-in since torch 2.1, so returning True is correct.
import transformers.utils.import_utils as _tx_import_utils  # noqa: E402
if not hasattr(_tx_import_utils, "is_torch_fx_available"):
    _tx_import_utils.is_torch_fx_available = lambda: True

from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402

_LOAD_ERROR: str | None = None
_model = None
_tokenizer = None
_load_done = threading.Event()


def _log(msg: str) -> None:
    """stdout flush so progress shows up in the HF container log immediately."""
    print(f"[zerogpu_backend] {msg}", flush=True)


def _load_model() -> None:
    """Heavy CPU work (download ~16 GB safetensors + load on CPU). Runs in a
    daemon thread so import returns immediately and the HTTP server can serve
    /api/setup/status without blocking.

    NOTE: We deliberately do NOT call `.to('cuda')` here. ZeroGPU's CUDA
    emulation only intercepts calls made from the main thread; a background
    thread that hits `to('cuda')` triggers a real `torch._C._cuda_init()`
    that fails because no GPU is attached yet. Instead, the @spaces.GPU
    function moves the model to cuda on each call, which is the safe pattern
    for off-main-thread initialization.
    """
    global _model, _tokenizer, _LOAD_ERROR
    try:
        _log(f"loading tokenizer for {_MODEL_ID}")
        _tokenizer = AutoTokenizer.from_pretrained(_MODEL_ID, trust_remote_code=True)
        _log("tokenizer ready; loading model weights (this downloads ~16 GB on first run)")
        model = AutoModelForCausalLM.from_pretrained(
            _MODEL_ID,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        model.eval()
        _model = model
        _log("READY (model on CPU; will move to cuda inside @spaces.GPU)")
    except Exception as exc:  # pragma: no cover
        _LOAD_ERROR = f"{exc.__class__.__name__}: {exc}"
        _log(f"LOAD FAILED: {_LOAD_ERROR}")
    finally:
        _load_done.set()


_loader_thread = threading.Thread(target=_load_model, name="zerogpu-model-loader", daemon=True)
_loader_thread.start()


_moved_to_cuda = False


@_spaces_gpu(duration=_DEFAULT_DURATION)
def _generate_on_gpu(
    messages: list[dict[str, Any]],
    max_new_tokens: int,
    temperature: float,
    json_mode: bool,
) -> str:
    global _moved_to_cuda
    if _model is None or _tokenizer is None:
        raise RuntimeError(_LOAD_ERROR or "zerogpu backend not initialized")

    # First call inside this ZeroGPU slot: move the CPU-resident model to the
    # real GPU. Cheap on subsequent calls (no-op if already on cuda).
    if _ON_ZEROGPU and not _moved_to_cuda:
        _model.to("cuda")
        _moved_to_cuda = True

    target_device = "cuda" if _ON_ZEROGPU else _model.device
    prompt = _tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = _tokenizer(prompt, return_tensors="pt").to(target_device)

    do_sample = temperature > 0
    gen_kwargs: dict[str, Any] = {
        "max_new_tokens": max_new_tokens,
        "do_sample": do_sample,
        "pad_token_id": _tokenizer.eos_token_id,
        # KV cache works with the transformers 4.55-4.57 pin in requirements.txt
        # — that's the range where the bundled modeling_minicpm.py's cache
        # ops match the installed transformers API. Decoding is O(n) per step
        # instead of O(n^2), so we can also raise the token cap.
        "use_cache": True,
    }
    if do_sample:
        gen_kwargs["temperature"] = temperature
        gen_kwargs["top_p"] = 0.9

    with torch.inference_mode():
        outputs = _model.generate(**inputs, **gen_kwargs)

    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
    text = _tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    text = _strip_think_tags(text)
    if json_mode:
        text = _extract_json(text)
    return text


_THINK_TAG_RE = None


def _strip_think_tags(text: str) -> str:
    """Remove <think>...</think> blocks (incl. empty ones) that MiniCPM4 and
    other hybrid-reasoning models emit even when /no_think is present in the
    system prompt. Also strips a stray opening/closing tag if unmatched."""
    import re
    global _THINK_TAG_RE
    if _THINK_TAG_RE is None:
        _THINK_TAG_RE = re.compile(r"<think>.*?</think>\s*", re.DOTALL | re.IGNORECASE)
    cleaned = _THINK_TAG_RE.sub("", text)
    cleaned = re.sub(r"</?think>\s*", "", cleaned, flags=re.IGNORECASE)
    return cleaned.strip()


def _extract_json(text: str) -> str:
    cleaned = text.strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.strip("`")
        if cleaned.lower().startswith("json"):
            cleaned = cleaned[4:].strip()
    start = cleaned.find("{")
    end = cleaned.rfind("}")
    if start >= 0 and end > start:
        return cleaned[start : end + 1]
    return text


_NO_THINK_HINTS = ("minicpm4", "minicpm-4", "minicpm_4", "qwen3", "qwen-3")


def _maybe_inject_no_think(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """MiniCPM4 / Qwen3 emit hidden <think> blocks that eat the token budget on
    small max_new_tokens. Append '/no_think' to the system turn so the model
    skips it. Mirrors llm/omni_client.py:_ensure_no_think."""
    if not any(hint in _MODEL_ID.lower() for hint in _NO_THINK_HINTS):
        return messages
    directive = "/no_think"
    for item in messages:
        if item.get("role") == "system" and isinstance(item.get("content"), str):
            if directive not in item["content"]:
                item["content"] = f"{item['content'].rstrip()} {directive}".strip()
            return messages
    return [{"role": "system", "content": directive}, *messages]


# Hard ceiling on tokens per generation. With use_cache=False (forced because
# the openbmb modeling code's cache path is broken on current transformers),
# attention is O(n^2) per step. Combined with the ZeroGPU 90-120 s per-call
# slot ceiling, anything above ~256 tokens risks getting aborted mid-stream.
# Witness chat and short story segments easily fit; longer chunks (full case
# briefings) will be truncated. Worth accepting for a working demo.
_ZEROGPU_MAX_TOKENS = int(os.getenv("PHANTOM_GRID_ZEROGPU_MAX_TOKENS", "384"))


def chat_completion(
    messages: list[dict[str, Any]],
    *,
    temperature: float = 0.4,
    max_tokens: int = 512,
    json_mode: bool = False,
) -> str:
    """Synchronous drop-in for an OpenAI chat-completions POST. Waits for the
    background model load to finish on the first call (no-op afterwards)."""
    _load_done.wait()
    if _LOAD_ERROR:
        raise RuntimeError(f"ZeroGPU model failed to load: {_LOAD_ERROR}")
    messages = _maybe_inject_no_think([dict(m) for m in messages])
    capped = min(max_tokens, _ZEROGPU_MAX_TOKENS)
    return _generate_on_gpu(messages, capped, temperature, json_mode)


def health() -> dict[str, Any]:
    """Non-blocking: reports whether the background load is done. NEVER waits."""
    loading = not _load_done.is_set()
    return {
        "reachable": True,
        "ready": _model is not None and _LOAD_ERROR is None and not loading,
        "detail": {
            "model_id": _MODEL_ID,
            "on_zerogpu": _ON_ZEROGPU,
            "loading": loading,
            "load_error": _LOAD_ERROR,
        },
    }