Create llama_bootstrap.py

llama_bootstrap.py

@@ -0,0 +1,118 @@
+"""
+llama-cpp-python bootstrap — MUST import before gradio or anything else.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import sys
+
+CUDA_VARIANTS = (
+    os.environ.get("LLAMA_CPP_CUDA", "cu124"),
+    "cu124",
+    "cu121",
+)
+PIP_BASE = [sys.executable, "-m", "pip"]
+
+
+def has_nvidia_gpu() -> bool:
+    if os.environ.get("LLAMA_CPP_FORCE_CPU", "").strip().lower() in ("1", "true", "yes"):
+        print("[LLAMA] LLAMA_CPP_FORCE_CPU=1 — skipping GPU", flush=True)
+        return False
+    if "CUDA_VISIBLE_DEVICES" in os.environ and os.environ["CUDA_VISIBLE_DEVICES"].strip() == "":
+        print("[LLAMA] CUDA_VISIBLE_DEVICES='' — no GPU visible", flush=True)
+        return False
+    if shutil.which("nvidia-smi") is None:
+        return False
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "-L"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        ok = result.returncode == 0 and bool(result.stdout.strip())
+        if ok:
+            print(f"[LLAMA] GPU detected:\n{result.stdout.strip()}", flush=True)
+        return ok
+    except Exception as e:
+        print(f"[LLAMA] nvidia-smi check failed: {e}", flush=True)
+        return False
+
+
+def _pip(args: list) -> None:
+    cmd = [*PIP_BASE, *args]
+    print(f"[LLAMA] $ {' '.join(cmd)}", flush=True)
+    env = os.environ.copy()
+    env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
+    env["PIP_ROOT_USER_ACTION"] = "ignore"
+    subprocess.check_call(cmd, env=env)
+
+
+def pip_uninstall_llama() -> None:
+    try:
+        _pip(["uninstall", "-y", "llama-cpp-python"])
+    except Exception:
+        pass
+
+
+def pip_install_cpu() -> None:
+    print("[LLAMA] Installing CPU wheel (no CUDA)...", flush=True)
+    _pip(["install", "--no-cache-dir", "llama-cpp-python"])
+
+
+def pip_install_cuda(variant: str) -> None:
+    index = f"https://abetlen.github.io/llama-cpp-python/whl/{variant}"
+    print(f"[LLAMA] Installing CUDA wheel ({variant})...", flush=True)
+    _pip([
+        "install",
+        "--no-cache-dir",
+        "llama-cpp-python",
+        "--extra-index-url",
+        index,
+    ])
+
+
+def try_import_llama():
+    import llama_cpp  # noqa: F401
+    from llama_cpp import Llama  # noqa: F401
+    return llama_cpp
+
+
+def bootstrap_llama_cpp():
+    try:
+        mod = try_import_llama()
+        print(f"[LLAMA] Ready (gpu={has_nvidia_gpu()})", flush=True)
+        return mod
+    except ImportError:
+        print("[LLAMA] Not installed — bootstrapping...", flush=True)
+    except (RuntimeError, OSError) as exc:
+        print(f"[LLAMA] Broken install ({exc}) — reinstalling...", flush=True)
+        pip_uninstall_llama()
+
+    if has_nvidia_gpu():
+        last_err = None
+        for variant in CUDA_VARIANTS:
+            try:
+                pip_uninstall_llama()
+                pip_install_cuda(variant)
+                mod = try_import_llama()
+                os.environ.setdefault("DAVIDAU_N_GPU_LAYERS", "-1")
+                print(f"[LLAMA] CUDA ready ({variant}, n_gpu_layers=-1)", flush=True)
+                return mod
+            except Exception as exc:
+                last_err = exc
+                print(f"[LLAMA] CUDA {variant} failed: {exc}", flush=True)
+        print(f"[LLAMA] All CUDA variants failed ({last_err}) — CPU fallback", flush=True)
+        pip_uninstall_llama()
+
+    pip_install_cpu()
+    os.environ.setdefault("DAVIDAU_N_GPU_LAYERS", "0")
+    mod = try_import_llama()
+    print("[LLAMA] CPU ready (n_gpu_layers=0)", flush=True)
+    return mod
+
+
+bootstrap_llama_cpp()