engine.py

"""
llama.cpp inference engine for the Domain AI Space backend.

Key design:
  - /data/models persistent storage (HF Spaces mounted volume); MODELS_DIR env override.
  - Cache-first: skips download if the GGUF already exists on disk.
  - load_streaming(): async generator that yields SSE-ready progress dicts so the
    Android client can show real-time download % and a "loading" phase.
  - Descending GPU layer ladder — steps down on OOM, falls back to CPU.
  - asyncio.Lock serialises concurrent inference requests (single llama.cpp context).
"""

import asyncio
import json
import logging
import os
import traceback
from pathlib import Path
from typing import AsyncIterator, Optional

import requests
from huggingface_hub import hf_hub_url
from llama_cpp import Llama

from capabilities import recommended_batch_size, recommended_threads

log = logging.getLogger(__name__)

_MODELS_DIR = Path(os.environ.get("MODELS_DIR", "/data/models"))

# Use the full GPU ladder only when a CUDA device is actually visible.
# On CPU-only HF Spaces the ladder would retry 7 times before reaching 0,
# making a 0.5B load appear stuck for minutes.
_cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
_HAS_CUDA = (bool(_cuda_devices) and _cuda_devices != "-1") or os.path.exists("/dev/nvidia0")
_GPU_LADDER = [99, 32, 24, 16, 12, 8, 4, 0] if _HAS_CUDA else [0]
log.info("GPU available: %s — ladder: %s", _HAS_CUDA, _GPU_LADDER)


# ---------------------------------------------------------------------------
# Download helper (blocking — run in an executor)
# ---------------------------------------------------------------------------

def _download_to_file(
    url: str,
    target: Path,
    token: Optional[str],
    on_progress,  # callable(pct: int)
) -> None:
    """Stream-download url → target, calling on_progress(pct) per 4 MB chunk."""
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    tmp = target.with_suffix(target.suffix + ".tmp")
    with requests.get(url, headers=headers, stream=True, timeout=(20, None)) as resp:
        resp.raise_for_status()
        total = int(resp.headers.get("content-length", 0))
        received = 0
        with open(tmp, "wb") as f:
            for chunk in resp.iter_content(chunk_size=4 * 1024 * 1024):
                if chunk:
                    f.write(chunk)
                    received += len(chunk)
                    if total > 0:
                        on_progress(min(99, int(received * 100 / total)))
    tmp.rename(target)


# ---------------------------------------------------------------------------
# Engine
# ---------------------------------------------------------------------------

class LlamaEngine:
    def __init__(self) -> None:
        self._llama: Optional[Llama] = None
        self._model_label: Optional[str] = None
        self._lock = asyncio.Lock()       # serialises inference
        self._load_lock = asyncio.Lock()  # prevents concurrent load calls
        self._loading: bool = False
        self._load_error: Optional[str] = None

    @property
    def loaded(self) -> bool:
        return self._llama is not None

    @property
    def loading(self) -> bool:
        return self._loading

    @property
    def load_error(self) -> Optional[str]:
        return self._load_error

    @property
    def model_label(self) -> Optional[str]:
        return self._model_label

    # ------------------------------------------------------------------
    # load_streaming — yields SSE-ready dicts
    # ------------------------------------------------------------------

    async def load_streaming(
        self,
        repo_id: str,
        filename: str,
        hf_token: Optional[str] = None,
        n_gpu_layers: Optional[int] = None,
        n_ctx: int = 4096,
    ) -> AsyncIterator[dict]:
        async with self._load_lock:
            async for event in self._load_streaming_inner(
                repo_id, filename, hf_token, n_gpu_layers, n_ctx
            ):
                yield event

    async def _load_streaming_inner(
        self,
        repo_id: str,
        filename: str,
        hf_token: Optional[str],
        n_gpu_layers: Optional[int],
        n_ctx: int,
    ) -> AsyncIterator[dict]:
        self._loading = True
        self._load_error = None
        loop = asyncio.get_running_loop()
        target_dir = _MODELS_DIR / repo_id.replace("/", "--")
        target_dir.mkdir(parents=True, exist_ok=True)
        target_path = target_dir / filename

        try:
            # ── Download phase ──────────────────────────────────────────
            if target_path.exists():
                log.info("Cache hit: %s", target_path)
                yield {"status": "cached", "pct": 100}
            else:
                log.info("Downloading %s / %s …", repo_id, filename)
                queue: asyncio.Queue = asyncio.Queue()
                last_pct = -1

                def do_download() -> None:
                    try:
                        _download_to_file(
                            url=hf_hub_url(repo_id, filename),
                            target=target_path,
                            token=hf_token,
                            on_progress=lambda pct: loop.call_soon_threadsafe(
                                queue.put_nowait, {"status": "downloading", "pct": pct}
                            ),
                        )
                        loop.call_soon_threadsafe(queue.put_nowait, {"status": "downloaded"})
                    except Exception as exc:
                        log.error("Download failed: %s\n%s", exc, traceback.format_exc())
                        loop.call_soon_threadsafe(
                            queue.put_nowait, {"status": "error", "message": str(exc)}
                        )

                fut = loop.run_in_executor(None, do_download)
                while True:
                    item = await queue.get()
                    if item["status"] == "downloading":
                        pct = item["pct"]
                        if pct != last_pct:
                            last_pct = pct
                            yield item
                    elif item["status"] == "downloaded":
                        await fut
                        break
                    else:  # error
                        await fut
                        self._load_error = item.get("message", "Download failed")
                        yield item
                        return

            # ── Load phase ──────────────────────────────────────────────
            yield {"status": "loading"}
            log.info("Loading %s into memory (n_ctx=%d) …", filename, n_ctx)

            n_batch = recommended_batch_size()
            n_threads = recommended_threads()
            ladder = [n_gpu_layers] if n_gpu_layers is not None else _GPU_LADDER
            log.info("Using GPU ladder: %s, n_batch=%d, n_threads=%d", ladder, n_batch, n_threads)

            llama: Optional[Llama] = None
            for gpu_layers in ladder:
                try:
                    log.info("Attempting load with n_gpu_layers=%d …", gpu_layers)
                    llama = await loop.run_in_executor(
                        None,
                        lambda gl=gpu_layers: Llama(
                            model_path=str(target_path),
                            n_gpu_layers=gl,
                            n_ctx=n_ctx,
                            n_batch=n_batch,
                            n_ubatch=n_batch,
                            n_threads=n_threads,
                            n_threads_batch=n_threads,
                            use_mmap=True,
                            use_mlock=False,   # mlock is slow on HF Space storage
                            flash_attn=True,
                            verbose=True,      # visible in Space logs for debugging
                        ),
                    )
                    log.info("Model loaded successfully with n_gpu_layers=%d", gpu_layers)
                    break
                except Exception as exc:
                    log.warning(
                        "n_gpu_layers=%d failed: %s\n%s",
                        gpu_layers, exc, traceback.format_exc(),
                    )
                    if gpu_layers == ladder[-1]:
                        err = f"Model load failed: {exc}"
                        self._load_error = err
                        yield {"status": "error", "message": err}
                        return

            self._llama = llama
            self._model_label = f"{repo_id}/{filename}"
            log.info("Engine ready: %s", self._model_label)
            yield {"status": "ready", "model": self._model_label}

        finally:
            self._loading = False

    # ------------------------------------------------------------------
    # load — convenience wrapper for startup (logs but doesn't stream)
    # ------------------------------------------------------------------

    async def load(
        self,
        repo_id: str,
        filename: str,
        hf_token: Optional[str] = None,
        n_gpu_layers: Optional[int] = None,
        n_ctx: int = 4096,
    ) -> None:
        async for event in self.load_streaming(repo_id, filename, hf_token, n_gpu_layers, n_ctx):
            log.info("Startup load: %s", event)
            if event.get("status") == "error":
                raise RuntimeError(event.get("message", "Load failed"))

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    async def stream_chat(
        self,
        messages: list,
        max_tokens: int = 512,
        temperature: float = 0.7,
        top_p: float = 0.95,
        stop: Optional[list] = None,
    ) -> AsyncIterator[str]:
        if self._llama is None:
            raise RuntimeError("No model loaded")

        loop = asyncio.get_running_loop()
        queue: asyncio.Queue = asyncio.Queue()
        kwargs = dict(messages=messages, max_tokens=max_tokens,
                      temperature=temperature, top_p=top_p, stream=True)
        if stop:
            kwargs["stop"] = stop

        def _generate() -> None:
            try:
                for chunk in self._llama.create_chat_completion(**kwargs):  # type: ignore[union-attr]
                    loop.call_soon_threadsafe(queue.put_nowait, chunk)
            except Exception as exc:
                loop.call_soon_threadsafe(queue.put_nowait, {"__error__": str(exc)})
            finally:
                loop.call_soon_threadsafe(queue.put_nowait, None)

        async with self._lock:
            fut = loop.run_in_executor(None, _generate)
            while True:
                item = await queue.get()
                if item is None:
                    break
                if isinstance(item, dict) and "__error__" in item:
                    yield f"data: {json.dumps({'error': item['__error__']})}\n\n"
                    break
                yield f"data: {json.dumps(item)}\n\n"
            await fut
            yield "data: [DONE]\n\n"

    async def complete_chat(
        self,
        messages: list,
        max_tokens: int = 512,
        temperature: float = 0.7,
        top_p: float = 0.95,
        stop: Optional[list] = None,
    ) -> dict:
        if self._llama is None:
            raise RuntimeError("No model loaded")

        loop = asyncio.get_running_loop()
        kwargs: dict = dict(messages=messages, max_tokens=max_tokens,
                            temperature=temperature, top_p=top_p, stream=False)
        if stop:
            kwargs["stop"] = stop

        async with self._lock:
            result = await loop.run_in_executor(
                None, lambda: self._llama.create_chat_completion(**kwargs)  # type: ignore[union-attr]
            )
        return result  # type: ignore[return-value]


engine = LlamaEngine()