model_manager.py

"""Cache management and single-model runtimes for the ZeroGPU playground."""

from __future__ import annotations

import gc
import logging
import os
import re
import threading
from pathlib import Path
from typing import Any

import requests

# ZeroGPU must be imported before torch.  The lazy llama.cpp import below is
# also intentionally kept inside the GPU-side loader.
import spaces
import torch
from huggingface_hub import hf_hub_download, scan_cache_dir, snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

from backend_router import (
    BACKEND_AUTO,
    BACKEND_LLAMACPP,
    BACKEND_TRANSFORMERS,
    BackendRouter,
    BackendRouterError,
    ModelInspection,
    ResolvedBackend,
    inspection_from_files,
)


LOGGER = logging.getLogger(__name__)
MODEL_ID_PATTERN = re.compile(r"^[^/\\s]+/[^/\\s]+$")

# Some very new architectures ship their Transformers implementation inside
# the model repository. Never enable remote code globally in this playground:
# the model ID is user-controlled, so doing so would allow an arbitrary Hub
# repository to execute Python in the Space.
TRUST_REMOTE_CODE_ALLOWLIST = {
    "XHToken/Spark-X2.5-4B",
}


def _trust_remote_code(model_id: str) -> bool:
    return model_id in TRUST_REMOTE_CODE_ALLOWLIST


def validate_model_id(model_id: str) -> str:
    normalized = (model_id or "").strip()
    if not MODEL_ID_PATTERN.fullmatch(normalized):
        raise ValueError("Model ID must look like namespace/model-name.")
    return normalized


class UnsupportedModelError(RuntimeError):
    """Compatibility alias for callers that want a user-facing load error."""


class HubAccessError(RuntimeError):
    """Raised when a gated/private Hub artifact needs explicit Space access."""


def _hub_token() -> str:
    return (os.getenv("HF_TOKEN") or "").strip()


def _hub_access_diagnosis(model_id: str) -> str:
    """Explain a Hub auth failure without exposing the secret or account data."""

    token = _hub_token()
    if not token:
        return (
            f"`{model_id}` is gated, but this Space process cannot see an `HF_TOKEN` secret. "
            "Add it under Space Settings → Repository secrets and restart the Space."
        )

    try:
        response = requests.get(
            "https://huggingface.co/api/whoami-v2",
            headers={"Authorization": f"Bearer {token}"},
            timeout=15,
        )
        if response.status_code in {401, 403}:
            return (
                "The Space sees `HF_TOKEN`, but the token is invalid, expired, or revoked. "
                "Create a new read-scoped token and replace the Space secret."
            )
        if response.ok:
            return (
                f"The Space sees `HF_TOKEN`, but its token account has not been granted access to `{model_id}`. "
                "Accept the gated model access with that same Hugging Face account, "
                "then restart the Space."
            )
    except Exception:
        pass

    return (
        f"Hugging Face denied access to gated repo `{model_id}`. "
        "Check that the token is read-scoped, belongs to the account that accepted access, "
        "and restart the Space after changing the secret."
    )


class ModelCache:
    """Keep standard snapshots and individually selected GGUF files in one cache."""

    def __init__(self, cache_dir: str | None = None) -> None:
        default_dir = Path.home() / ".cache" / "huggingface" / "llm-playground"
        self.root = Path(cache_dir or os.getenv("PLAYGROUND_CACHE_DIR", default_dir))
        self.root.mkdir(parents=True, exist_ok=True)
        self.router = BackendRouter()

    def inspect_remote(self, model_id: str) -> ModelInspection:
        return self.router.inspect_remote(validate_model_id(model_id), cache_dir=self.root)

    def download(self, model_id: str) -> Path:
        """Download a standard Transformers snapshot on CPU."""

        return Path(
            snapshot_download(
                repo_id=validate_model_id(model_id),
                repo_type="model",
                cache_dir=str(self.root),
                token=_hub_token() or None,
            )
        )

    def download_gguf(self, model_id: str, filename: str) -> Path:
        """Download exactly one GGUF file, never the whole repository."""

        model_id = validate_model_id(model_id)
        filename = (filename or "").strip()
        if not filename or not filename.lower().endswith(".gguf"):
            raise ValueError("Choose one `.gguf` file before downloading.")

        inspection = self.inspect_remote(model_id)
        if filename not in inspection.gguf_files:
            raise ValueError(f"`{filename}` is not a GGUF file in `{model_id}`.")

        try:
            return Path(
                hf_hub_download(
                    repo_id=model_id,
                    filename=filename,
                    repo_type="model",
                    cache_dir=str(self.root),
                    token=_hub_token() or None,
                )
            )
        except Exception as exc:
            error_text = f"{exc.__class__.__name__} {exc}".lower()
            if any(
                marker in error_text
                for marker in ("401", "403", "gatedrepoerror", "unauthorized", "forbidden")
            ):
                raise HubAccessError(_hub_access_diagnosis(model_id)) from exc
            raise

    def cached_snapshot(self, model_id: str) -> Path:
        model_id = validate_model_id(model_id)
        try:
            return Path(
                snapshot_download(
                    repo_id=model_id,
                    repo_type="model",
                    cache_dir=str(self.root),
                    local_files_only=True,
                    token=_hub_token() or None,
                )
            )
        except Exception as exc:
            raise FileNotFoundError(
                f"{model_id} is not downloaded yet. Click Download first."
            ) from exc

    def cached_gguf(self, model_id: str, filename: str) -> Path:
        model_id = validate_model_id(model_id)
        filename = (filename or "").strip()
        if not filename.lower().endswith(".gguf"):
            raise ValueError("Choose one `.gguf` file before loading.")
        try:
            return Path(
                hf_hub_download(
                    repo_id=model_id,
                    filename=filename,
                    repo_type="model",
                    cache_dir=str(self.root),
                    local_files_only=True,
                    token=_hub_token() or None,
                )
            )
        except Exception as exc:
            raise FileNotFoundError(
                f"`{filename}` is not downloaded yet. Click Download for this GGUF file first."
            ) from exc

    def cached_gguf_default(self, model_id: str) -> str | None:
        """Return the preferred already-downloaded GGUF, if one exists."""

        model_id = validate_model_id(model_id)
        info = scan_cache_dir(cache_dir=str(self.root))
        cached_names: set[str] = set()
        for repo in info.repos:
            if repo.repo_id != model_id:
                continue
            for revision in repo.revisions:
                cached_names.update(
                    file_info.file_name
                    for file_info in revision.files
                    if file_info.file_name.lower().endswith(".gguf")
                )
        if not cached_names:
            return None
        return inspection_from_files(model_id, cached_names, source="cache").default_gguf

    def resolve_cached(
        self,
        model_id: str,
        requested_backend: str | None,
        selected_file: str | None,
    ) -> ResolvedBackend:
        """Resolve a local artifact without doing network I/O on a GPU call."""

        model_id = validate_model_id(model_id)
        selected = (selected_file or "").strip()
        if selected:
            path = self.cached_gguf(model_id, selected)
            inspection = self.router.synthetic_gguf(model_id, selected)
            backend = self.router.resolve_backend(
                inspection,
                requested_backend=requested_backend,
                selected_file=selected,
            )
        else:
            try:
                path = self.cached_snapshot(model_id)
                inspection = self.router.inspect_snapshot(model_id, path)
                backend = self.router.resolve_backend(
                    inspection,
                    requested_backend=requested_backend,
                    selected_file=None,
                )
            except (FileNotFoundError, BackendRouterError):
                # API callers and a freshly opened browser may not yet have
                # the dynamic Dropdown value. If a GGUF was already selected
                # and downloaded, Auto can still pick that cached quant.
                if (requested_backend or BACKEND_AUTO).strip().lower() not in {
                    "auto",
                    "automatic",
                    "llama.cpp",
                    "llama-cpp",
                    "llamacpp",
                    "llama",
                }:
                    raise
                selected = self.cached_gguf_default(model_id) or ""
                if not selected:
                    raise
                path = self.cached_gguf(model_id, selected)
                inspection = self.router.synthetic_gguf(model_id, selected)
                backend = self.router.resolve_backend(
                    inspection,
                    requested_backend=requested_backend,
                    selected_file=selected,
                )
        if backend == BACKEND_LLAMACPP and not selected:
            raise BackendRouterError(
                "Choose the GGUF file you downloaded before loading it with llama.cpp."
            )
        return ResolvedBackend(
            backend=backend,
            path=path,
            inspection=inspection,
            selected_file=selected or None,
        )

    def describe(self, model_id: str, selected_file: str | None = None) -> str:
        model_id = (model_id or "").strip()
        if not model_id:
            return "Disk cache: no model selected."
        try:
            if selected_file:
                path = self.cached_gguf(model_id, selected_file)
                return f"Disk cache: GGUF ready (`{path.name}`)."
            path = self.cached_snapshot(model_id)
            return f"Disk cache: snapshot ready (`{path.name}`)."
        except (ValueError, FileNotFoundError):
            return f"Disk cache: `{model_id}` is not downloaded."

    def delete(self, model_id: str) -> bool:
        """Remove all cached revisions for one model, including selected GGUFs."""

        model_id = validate_model_id(model_id)
        cache_info = scan_cache_dir(cache_dir=str(self.root))
        revisions = []
        for repo in cache_info.repos:
            if repo.repo_id == model_id:
                revisions.extend(revision.commit_hash for revision in repo.revisions)
        if not revisions:
            return False
        cache_info.delete_revisions(*revisions).execute()
        return True


class ModelRuntime:
    """Route one active model to Transformers or llama.cpp."""

    def __init__(self, cache: ModelCache) -> None:
        self.cache = cache
        self._model: Any | None = None
        self._tokenizer: Any | None = None
        self._llama: Any | None = None
        self._model_id: str | None = None
        self._backend: str | None = None
        self._selected_file: str | None = None
        self._inspection: ModelInspection | None = None
        self._lock = threading.RLock()

    @property
    def active_model_id(self) -> str | None:
        return self._model_id

    @property
    def active_backend(self) -> str | None:
        return self._backend

    @property
    def active_selected_file(self) -> str | None:
        return self._selected_file

    @property
    def active_inspection(self) -> ModelInspection | None:
        return self._inspection

    def active_label(self) -> str:
        if not self._model_id:
            return "No model loaded"
        suffix = f" · {self._backend}"
        if self._selected_file:
            suffix += f" · {self._selected_file}"
        return f"{self._model_id}{suffix}"

    def unload(self) -> None:
        """Release both possible runtimes and clear CUDA allocator state."""

        with self._lock:
            old_model = self._model
            old_llama = self._llama
            self._model = None
            self._tokenizer = None
            self._llama = None
            self._model_id = None
            self._backend = None
            self._selected_file = None
            self._inspection = None

            if old_llama is not None:
                close = getattr(old_llama, "close", None)
                if callable(close):
                    try:
                        close()
                    except Exception:
                        LOGGER.debug("llama.cpp close failed during cleanup", exc_info=True)
                del old_llama
            if old_model is not None:
                del old_model

            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def ensure_loaded(
        self,
        model_id: str,
        requested_backend: str | None = BACKEND_AUTO,
        selected_file: str | None = None,
    ) -> ResolvedBackend:
        """Load one local artifact on GPU, replacing any active model."""

        model_id = validate_model_id(model_id)
        target = self.cache.resolve_cached(model_id, requested_backend, selected_file)
        identity = (model_id, target.backend, target.selected_file)

        with self._lock:
            current = (self._model_id, self._backend, self._selected_file)
            if current == identity and (self._model is not None or self._llama is not None):
                return target

            # The one-model invariant is enforced before constructing either
            # a new Transformers model or a new llama.cpp context.
            self.unload()
            if target.backend == BACKEND_LLAMACPP:
                self._load_llama(target)
            else:
                self._load_transformers(target)

            self._model_id = model_id
            self._backend = target.backend
            self._selected_file = target.selected_file
            self._inspection = target.inspection
            return target

    @staticmethod
    def _preferred_dtype(inspection: ModelInspection) -> Any:
        value = str(
            inspection.config.get("torch_dtype") or inspection.config.get("dtype") or ""
        ).lower()
        if "float16" in value or value in {"fp16", "half"}:
            return torch.float16
        if "float32" in value or value == "fp32":
            return torch.float32
        if "float8" in value or "fp8" in value:
            return "auto"
        return torch.bfloat16

    @staticmethod
    def _bnb_config(inspection: ModelInspection) -> Any | None:
        """Create a BitsAndBytesConfig only for filename-only quant repos."""

        quant_config = inspection.config.get("quantization_config")
        if isinstance(quant_config, dict) and any(
            key in quant_config
            for key in ("load_in_4bit", "load_in_8bit", "_load_in_4bit", "_load_in_8bit")
        ):
            return None  # Transformers will consume the repository config itself.

        if inspection.quantization_kind != "bitsandbytes":
            return None
        try:
            from transformers import BitsAndBytesConfig

            text = inspection.format_label.lower()
            load_in_8bit = "8-bit" in text
            return BitsAndBytesConfig(
                load_in_4bit=not load_in_8bit,
                load_in_8bit=load_in_8bit,
                bnb_4bit_compute_dtype=torch.bfloat16,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
            )
        except Exception:
            LOGGER.debug("Could not construct a BitsAndBytesConfig", exc_info=True)
            return None

    def _load_transformers(self, target: ResolvedBackend) -> None:
        inspection = target.inspection
        trust_remote_code = _trust_remote_code(inspection.model_id)
        tokenizer = AutoTokenizer.from_pretrained(
            str(target.path),
            local_files_only=True,
            use_fast=True,
            trust_remote_code=trust_remote_code,
        )

        quantized = inspection.is_transformers_quantized
        load_kwargs: dict[str, Any] = {
            "local_files_only": True,
            "low_cpu_mem_usage": True,
            "trust_remote_code": trust_remote_code,
            "dtype": "auto" if quantized else self._preferred_dtype(inspection),
        }
        if quantized:
            # Quantized modules must be placed by Accelerate/Transformers and
            # must not receive a later blanket `.to("cuda")` call.
            load_kwargs["device_map"] = "cuda"
            bnb_config = self._bnb_config(inspection)
            if bnb_config is not None:
                load_kwargs["quantization_config"] = bnb_config

        try:
            model = AutoModelForCausalLM.from_pretrained(str(target.path), **load_kwargs)

            # Spark-X2.5 ships generation_config.top_k = -1 to mean
            # "disable top-k". Recent Transformers validates GenerationConfig
            # before generation and rejects negative top_k values, even when
            # generate() receives a positive top_k override.
            if inspection.model_id == "XHToken/Spark-X2.5-4B":
                if getattr(model, "generation_config", None) is not None:
                    model.generation_config.top_k = 20

            if not quantized:
                model = model.to("cuda")
            model = model.eval()
        except Exception as exc:
            label = inspection.format_label
            raise RuntimeError(
                f"Could not load `{label}` with Transformers. "
                "The repository's quantization runtime may need a compatible loader package. "
                f"Details: {str(exc).splitlines()[0][:260]}"
            ) from exc

        if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
            tokenizer.pad_token = tokenizer.eos_token
        if getattr(model.config, "pad_token_id", None) is None:
            model.config.pad_token_id = tokenizer.pad_token_id

        self._tokenizer = tokenizer
        self._model = model

    def _load_llama(self, target: ResolvedBackend) -> None:
        try:
            from llama_cpp import Llama
        except Exception as exc:
            raise RuntimeError(
                "llama.cpp is not available. The Space needs the CUDA-enabled llama-cpp-python wheel."
            ) from exc

        kwargs: dict[str, Any] = {
            "model_path": str(target.path),
            "n_gpu_layers": -1,
            "n_ctx": 8192,
            "n_batch": 512,
            "n_threads": max(2, min(8, os.cpu_count() or 4)),
            "verbose": False,
        }
        try:
            llama = Llama(**kwargs, flash_attn=True)
        except TypeError:
            # Keep compatibility with older wheels that predate flash_attn in
            # the Python constructor; the GPU offload remains explicit.
            llama = Llama(**kwargs)
        except Exception as exc:
            raise RuntimeError(
                f"Could not load `{target.selected_file}` with llama.cpp GPU offload. "
                f"Details: {str(exc).splitlines()[0][:260]}"
            ) from exc
        self._llama = llama

    @staticmethod
    def _history_to_messages(history: list[Any] | None) -> list[dict[str, str]]:
        messages: list[dict[str, str]] = []
        for item in history or []:
            if isinstance(item, dict):
                role = str(item.get("role", ""))
                content = item.get("content", "")
                if role in {"user", "assistant"} and isinstance(content, str):
                    messages.append({"role": role, "content": content})
            elif isinstance(item, (list, tuple)) and len(item) == 2:
                user_text, assistant_text = item
                if isinstance(user_text, str) and user_text:
                    messages.append({"role": "user", "content": user_text})
                if isinstance(assistant_text, str) and assistant_text:
                    messages.append({"role": "assistant", "content": assistant_text})
        return messages

    @staticmethod
    def _plain_prompt(messages: list[dict[str, str]]) -> str:
        lines = [f"{m['role'].capitalize()}: {m['content']}" for m in messages]
        return "\n".join(lines) + "\nAssistant:"

    def _tokenize(self, messages: list[dict[str, str]]) -> Any:
        assert self._tokenizer is not None
        tokenizer = self._tokenizer
        if hasattr(tokenizer, "apply_chat_template"):
            try:
                return tokenizer.apply_chat_template(
                    messages,
                    add_generation_prompt=True,
                    tokenize=True,
                    return_tensors="pt",
                    return_dict=True,
                )
            except TypeError:
                try:
                    return tokenizer.apply_chat_template(
                        messages,
                        add_generation_prompt=True,
                        tokenize=True,
                        return_tensors="pt",
                    )
                except Exception:
                    LOGGER.debug("Chat template without return_dict failed", exc_info=True)
            except Exception:
                LOGGER.debug("Chat template failed; using plain prompt", exc_info=True)
        return tokenizer(self._plain_prompt(messages), return_tensors="pt")

    def generate(
        self,
        model_id: str,
        requested_backend: str | None,
        selected_file: str | None,
        message: str,
        history: list[Any] | None,
        system_prompt: str,
        max_new_tokens: int,
        temperature: float,
        top_p: float,
    ) -> str:
        target = self.ensure_loaded(model_id, requested_backend, selected_file)
        messages: list[dict[str, str]] = []
        if (system_prompt or "").strip():
            messages.append({"role": "system", "content": system_prompt.strip()})
        messages.extend(self._history_to_messages(history))
        messages.append({"role": "user", "content": (message or "").strip()})

        with self._lock:
            if target.backend == BACKEND_LLAMACPP:
                if self._llama is None:
                    raise RuntimeError("llama.cpp runtime is not loaded.")
                result = self._llama.create_chat_completion(
                    messages=messages,
                    max_tokens=max_new_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
                answer = ""
                if isinstance(result, dict) and result.get("choices"):
                    answer = str(
                        result["choices"][0].get("message", {}).get("content", "")
                    )
                return answer.strip() or "The model returned an empty response."

            if self._model is None or self._tokenizer is None:
                raise RuntimeError("Transformers runtime is not loaded.")
            encoded = self._tokenize(messages)
            encoded = {
                key: value.to("cuda")
                for key, value in encoded.items()
                if torch.is_tensor(value)
            }
            input_length = int(encoded["input_ids"].shape[-1])
            generation_kwargs: dict[str, Any] = {
                "max_new_tokens": max_new_tokens,
                "do_sample": temperature > 0,
                "top_k": 20,
            }

            if temperature > 0:
                generation_kwargs.update({
                    "temperature": temperature,
                    "top_p": top_p,
                    "top_k": 20,
                })
            with torch.inference_mode():
                generated = self._model.generate(**encoded, **generation_kwargs)
            new_tokens = generated[0, input_length:]
            answer = self._tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
            return answer or "The model returned an empty response."