space_runtime.py

"""Start the local MiniCPM llama.cpp backend when running in Hugging Face Spaces."""

from __future__ import annotations

import os
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path


ROOT = Path(__file__).parent
MODEL_REPO = os.getenv("MELTMIND_GGUF_REPO", "AyyYOO/MiniCPM4-8B-Q4_K_M-GGUF")
MODEL_FILE = os.getenv("MELTMIND_GGUF_FILE", "minicpm4-8b-q4_k_m.gguf")
MODEL_ALIAS = os.getenv("MELTMIND_LLM_MODEL", "openbmb/MiniCPM4-8B")
HOST = os.getenv("MELTMIND_LLM_HOST", "127.0.0.1")
PORT = int(os.getenv("MELTMIND_LLM_PORT", "8080"))
HEALTH_URL = f"http://{HOST}:{PORT}/v1/models"

_lock = threading.Lock()
_process: subprocess.Popen | None = None
_state = {
    "auto_start_enabled": False,
    "status": "not_started",
    "detail": "The local backend is managed externally.",
    "model_repo": MODEL_REPO,
    "model_file": MODEL_FILE,
}


def _auto_start_enabled() -> bool:
    explicit = os.getenv("MELTMIND_AUTO_START_LLM")
    if explicit is not None:
        return explicit.lower() in {"1", "true", "yes", "on"}
    return bool(os.getenv("SPACE_ID"))


def _healthy(timeout: float = 1.0) -> bool:
    try:
        with urllib.request.urlopen(HEALTH_URL, timeout=timeout) as response:
            return response.status < 400
    except (TimeoutError, urllib.error.URLError):
        return False


def _set_state(status: str, detail: str) -> None:
    with _lock:
        _state["status"] = status
        _state["detail"] = detail


def _run_backend() -> None:
    global _process
    try:
        if _healthy():
            _set_state("ready", "MiniCPM is already accepting requests.")
            return

        configured_model_path = os.getenv("MELTMIND_MODEL_PATH", "").strip()
        if configured_model_path and Path(configured_model_path).is_file():
            model_path = configured_model_path
        else:
            _set_state("downloading", "Downloading the quantized MiniCPM4-8B model from the Hub.")
            from huggingface_hub import hf_hub_download

            model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
        _set_state("loading", "Loading MiniCPM4-8B into the llama.cpp backend.")

        command = [
            sys.executable,
            "-m",
            "llama_cpp.server",
            "--model",
            model_path,
            "--model_alias",
            MODEL_ALIAS,
            "--host",
            HOST,
            "--port",
            str(PORT),
            "--n_ctx",
            os.getenv("MELTMIND_CONTEXT_SIZE", "8192"),
            "--n_threads",
            os.getenv("MELTMIND_THREADS", str(max(2, min(8, os.cpu_count() or 4)))),
            "--n_gpu_layers",
            os.getenv("MELTMIND_GPU_LAYERS", "0"),
            "--chat_format",
            "chatml",
            "--use_mlock",
            "false",
            "--cache",
            "false",
        ]
        _process = subprocess.Popen(command, cwd=ROOT)

        for _ in range(900):
            if _healthy():
                _set_state("ready", "MiniCPM4-8B is loaded and accepting requests.")
                return
            if _process.poll() is not None:
                raise RuntimeError(f"llama.cpp server exited with code {_process.returncode}")
            time.sleep(2)
        raise TimeoutError("MiniCPM did not become ready before the startup timeout.")
    except Exception as exc:  # The deterministic app must remain available.
        _set_state("fallback", f"MiniCPM backend unavailable: {exc}")


def start_space_backend() -> dict:
    enabled = _auto_start_enabled()
    with _lock:
        _state["auto_start_enabled"] = enabled
    if not enabled:
        return backend_status()
    if _healthy():
        _set_state("ready", "MiniCPM is already accepting requests.")
        return backend_status()
    with _lock:
        if _state["status"] in {"downloading", "loading"}:
            return dict(_state)
        _state["status"] = "starting"
        _state["detail"] = "Preparing the MiniCPM backend."
    threading.Thread(target=_run_backend, name="meltmind-space-backend", daemon=True).start()
    return backend_status()


def backend_status() -> dict:
    with _lock:
        status = dict(_state)
    status["healthy"] = _healthy()
    if status["healthy"]:
        status["status"] = "ready"
    return status