minifam/bootstrap.py

"""Stand up a local Ollama server inside a Hugging Face GPU Space — no Docker.

Why this exists
---------------
A *paid* GPU Space (e.g. Nvidia L4/A10G — **not** ZeroGPU) keeps a GPU attached
to the container for its whole lifetime. That means we can run an ordinary
`ollama serve` as a background process and talk to it over `localhost:11434`,
exactly like on the dev Mac. So none of the agent/config code changes — only
this file is Space-specific.

Why a runtime bootstrap instead of a Dockerfile: it keeps this a plain Gradio
SDK Space. On first boot we download the Ollama runtime and pull the model;
everything runs as a background thread so Gradio opens its port immediately and
the Space passes its health check while the (large) model is still downloading.

Turning it on (Space → Settings → Variables and secrets):
  MINIFAM_BOOTSTRAP_OLLAMA=1     # the on switch; leave UNSET locally (no-op)
  OLLAMA_MODELS=/data/ollama     # optional: persist the ~18 GB model across
  MINIFAM_OLLAMA_DIR=/data/runtime   # restarts (needs persistent storage at /data)
The model tag comes from MINIFAM_MODEL (config.MODEL, default qwen3:30b) and the
app already points MINIFAM_LLM_BASE_URL at localhost:11434 — leave both default.
"""

import os
import subprocess
import tarfile
import threading
import time
import urllib.request
from pathlib import Path

from . import config

# Official static build; bundles the CUDA runners (~1.5 GB download).
_OLLAMA_TGZ = "https://ollama.com/download/ollama-linux-amd64.tgz"
# Where the binary lives. Point MINIFAM_OLLAMA_DIR at persistent storage to keep
# it across restarts; otherwise it re-downloads on each cold start.
_PREFIX = Path(
    os.environ.get("MINIFAM_OLLAMA_DIR", str(Path.home() / ".cache" / "minifam" / "ollama"))
)
_BIN = _PREFIX / "bin" / "ollama"
_HOST = "127.0.0.1:11434"

# Bootstrap progress, read by warmup_message() so chat can say "still loading"
# instead of erroring. off → installing → starting → pulling → ready | error.
_status = "off"
_status_detail = ""
_lock = threading.Lock()


def enabled() -> bool:
    """True only when explicitly switched on in the Space (off locally)."""
    # The ZeroGPU backend runs the model in-process — there's no Ollama to start.
    if os.environ.get("MINIFAM_BACKEND", "").strip().lower() == "zerogpu":
        return False
    return os.environ.get("MINIFAM_BOOTSTRAP_OLLAMA", "").strip().lower() in {
        "1", "true", "yes", "on",
    }


def _set(status: str, detail: str = "") -> None:
    global _status, _status_detail
    _status, _status_detail = status, detail


def _env() -> dict:
    """Environment for the ollama subprocesses (inherits Space vars like
    OLLAMA_MODELS) with our host + the bundled CUDA libs on the search path."""
    env = dict(os.environ)
    env["OLLAMA_HOST"] = _HOST
    # Keep the 30B resident in VRAM between turns so each request is fast.
    env.setdefault("OLLAMA_KEEP_ALIVE", "-1")
    lib = str(_PREFIX / "lib" / "ollama")
    env["LD_LIBRARY_PATH"] = f"{lib}:{env.get('LD_LIBRARY_PATH', '')}".rstrip(":")
    return env


def _server_up(timeout: float = 2.0) -> bool:
    try:
        with urllib.request.urlopen(f"http://{_HOST}/api/tags", timeout=timeout):
            return True
    except Exception:
        return False


def _install_binary() -> None:
    if _BIN.exists():
        return
    _set("installing", "downloading the Ollama runtime (~1.5 GB)")
    _PREFIX.mkdir(parents=True, exist_ok=True)
    tgz = _PREFIX / "ollama.tgz"
    urllib.request.urlretrieve(_OLLAMA_TGZ, tgz)
    with tarfile.open(tgz) as tar:
        tar.extractall(_PREFIX, filter="fully_trusted")  # trusted: official build
    tgz.unlink(missing_ok=True)
    _BIN.chmod(0o755)


def _start_server() -> None:
    if _server_up():
        return
    _set("starting", "launching ollama serve")
    subprocess.Popen(
        [str(_BIN), "serve"],
        env=_env(),
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    for _ in range(60):  # serve is quick; give it up to a minute to bind
        if _server_up():
            return
        time.sleep(1)
    raise RuntimeError("ollama serve did not come up within 60s")


def _pull_model() -> None:
    model = config.MODEL
    # Already present (e.g. cached on persistent storage)? Skip the long pull.
    if not _model_present(model):
        _set("pulling", f"downloading {model} — first boot takes a few minutes")
        proc = subprocess.run([str(_BIN), "pull", model], env=_env())
        if proc.returncode != 0:
            raise RuntimeError(f"`ollama pull {model}` failed (exit {proc.returncode})")
    _set("ready", model)


def _model_present(model: str) -> bool:
    try:
        out = subprocess.run(
            [str(_BIN), "list"], env=_env(), capture_output=True, text=True, timeout=20
        ).stdout
    except Exception:
        return False
    # `ollama list` prints the tag (with or without an implicit :latest).
    base = model.split(":")[0]
    return any(base in line for line in out.splitlines()[1:])


def ensure_ollama() -> None:
    """Kick off the install → serve → pull sequence in the background.

    No-op unless MINIFAM_BOOTSTRAP_OLLAMA is set, so importing app.py on the dev
    machine does nothing and your own Ollama is used as before. Returns at once;
    progress is reported via status()/warmup_message().
    """
    if not enabled():
        return
    with _lock:
        if _status not in {"off", "error"}:
            return  # already running or finished
        _set("installing")

    def _run() -> None:
        try:
            _install_binary()
            _start_server()
            _pull_model()
        except Exception as e:  # surfaced to the user via warmup_message()
            _set("error", str(e))

    threading.Thread(target=_run, daemon=True, name="ollama-bootstrap").start()


def warmup_message() -> str | None:
    """A friendly 'still warming up' line for the chat box, or None.

    None means 'not bootstrapping, or already ready' — callers then fall back to
    their normal error handling. Locally this is always None.
    """
    if not enabled() or _status == "ready":
        return None
    if _status == "error":
        return (
            f"The model failed to start in this Space ({_status_detail}). "
            "Check the Space logs."
        )
    friendly = {
        "installing": "Setting up the model runtime in the Space…",
        "starting": "Starting the model server…",
        "pulling": (
            f"Downloading the model ({config.MODEL}, ~18 GB) — the first boot "
            "takes a few minutes. Try again shortly."
        ),
    }
    return friendly.get(_status, "Warming up — try again in a moment…")