server/health.py

"""Liveness + a hardware-adequacy signal for the UI banner and external monitors.

`health_status()` powers both `GET /health` and the on-page status banner. It
reports `degraded: true` when the *real* model would run on CPU-only hardware
(where extraction can be slow or time out). `FORCE_DEGRADED=1` forces it on so
the banner can be exercised without actually changing the Space's hardware.
"""
from __future__ import annotations

import os


def gpu_available() -> bool:
    """Best-effort GPU probe. This app serves via llama.cpp (no torch), so we
    can't rely on torch.cuda — check for an NVIDIA device or `nvidia-smi`."""
    import glob
    import shutil
    import subprocess

    if glob.glob("/dev/nvidia[0-9]*"):
        return True
    if shutil.which("nvidia-smi"):
        try:
            return subprocess.run(["nvidia-smi"], capture_output=True, timeout=5).returncode == 0
        except Exception:  # noqa: BLE001  any probe failure -> treat as no GPU
            return False
    return False


def health_status() -> dict:
    """Return liveness + the degraded/device/model signal (read at call time)."""
    if os.environ.get("FORCE_DEGRADED") == "1":
        return {
            "ok": True, "device": "cpu", "model": "real", "degraded": True,
            "reason": "Running the model on CPU-only hardware. Extraction may be slow "
                      "or time out. Upgrade to a GPU.",
        }
    # Stub mode is a deliberate, fast, free preview — not a degraded state.
    if os.environ.get("USE_STUB_EXTRACTOR") == "1":
        return {"ok": True, "device": "cpu", "model": "stub", "degraded": False, "reason": ""}

    # Real model. On the Space, INFERENCE_BASE_URL points at the *local* llama-server,
    # so local inference still depends on this host having a GPU. A non-localhost URL
    # means inference runs elsewhere, so this host's hardware is irrelevant.
    base = os.environ.get("INFERENCE_BASE_URL", "")
    local = (not base) or ("127.0.0.1" in base) or ("localhost" in base)
    gpu = gpu_available()
    degraded = local and not gpu
    device = "cuda" if (local and gpu) else ("cpu" if local else "remote")
    reason = (
        "Running the model on CPU-only hardware. Extraction may be slow or time out. "
        "Upgrade to a GPU." if degraded else ""
    )
    return {"ok": True, "device": device, "model": "real", "degraded": degraded, "reason": reason}