notebook-llm-cli 0.1.0__py3-none-any.whl

notebook_llm/estimate.py

@@ -0,0 +1,128 @@
+"""Will it fit? How fast will it run? (approximate, memory-bandwidth model)"""
+import re
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+from .estimate_data import BYTES_PER_PARAM, MOE
+
+GPU_EFF = 0.65        # fraction of peak bandwidth llama.cpp typically reaches
+MOE_EFF = 0.70        # extra penalty for expert routing
+CPU_BW_GBS = 15.0     # effective system-RAM bandwidth for offloaded layers
+OVERHEAD_PER_GPU = 0.7  # CUDA context + compute buffers (GB)
+KV_FACTORS = {"f16": 1.0, "q8_0": 0.53, "q4_0": 0.28}
+
+
+@dataclass
+class Estimate:
+    verdict: str            # fits | tight | offload | no | cpu
+    needed_gb: float
+    weights_gb: float
+    kv_gb: float
+    overhead_gb: float
+    gpu_fraction: float     # share of weights that sit in VRAM
+    tok_s: Optional[float]
+    context: int
+    note: str = ""
+
+    @property
+    def label(self) -> str:
+        return {"fits": "FITS", "tight": "TIGHT", "offload": "SLOW (CPU offload)",
+                "no": "TOO BIG", "cpu": "CPU ONLY"}[self.verdict]
+
+
+def parse_params(name: str, tag: str) -> Tuple[Optional[float], Optional[float]]:
+    """-> (total params in B, active params in B or None for dense)."""
+    key = f"{name}:{tag}"
+    base = f"{name}:{tag.split('-')[0]}"
+    for k in (key, base):
+        if k in MOE:
+            return MOE[k]
+    active = None
+    m = re.search(r"-a(\d+(?:\.\d+)?)b", tag)
+    if m:
+        active = float(m.group(1))
+    mx = re.match(r"(\d+)x(\d+(?:\.\d+)?)b", tag)
+    if mx:
+        return float(mx.group(1)) * float(mx.group(2)), active
+    m = re.match(r"e?(\d+(?:\.\d+)?)([bm])", tag)
+    if m:
+        v = float(m.group(1)) * (0.001 if m.group(2) == "m" else 1)
+        return v, active
+    return None, active
+
+
+def bytes_per_param(tag: str) -> float:
+    t = tag.lower()
+    for key, v in BYTES_PER_PARAM:
+        if key in t:
+            return v
+    return 0.61  # Ollama default tags are ~Q4_K_M
+
+
+def guess_size_gb(params_b: float, tag: str) -> float:
+    return round(params_b * bytes_per_param(tag), 2)
+
+
+def kv_cache_gb(params_b: float, ctx: int, kv_type: str = "f16") -> float:
+    per_token_mb = 0.03 * (max(params_b, 0.3) ** 0.85) * KV_FACTORS.get(kv_type, 1.0)
+    return per_token_mb * ctx / 1024
+
+
+def _overhead(hw) -> float:
+    return OVERHEAD_PER_GPU * max(1, hw.gpu.count)
+
+
+def auto_context(size_gb: float, params_b: float, hw, kv_type: str = "q8_0") -> int:
+    if not hw.gpu.available:
+        return 4096
+    room = hw.gpu.total_vram_gb * 0.92 - size_gb - _overhead(hw)
+    for ctx in (32768, 16384, 8192, 4096):
+        if kv_cache_gb(params_b, ctx, kv_type) <= 0.7 * room:
+            return ctx
+    return 4096
+
+
+def estimate(size_gb: float, params_b: Optional[float], active_b: Optional[float],
+             hw, ctx: int, kv_type: str = "q8_0") -> Estimate:
+    params = params_b or size_gb / 0.61
+    kv = kv_cache_gb(params, ctx, kv_type)
+    moe = bool(active_b and params and active_b < params * 0.8)
+    active_bytes = size_gb * (active_b / params) if moe else size_gb
+    eff_boost = MOE_EFF if moe else 1.0
+    gpu = hw.gpu
+
+    if not gpu.available:
+        needed = size_gb + kv + 0.5
+        t = active_bytes / (CPU_BW_GBS * eff_boost) + 0.003
+        verdict = "cpu" if needed <= hw.ram_gb * 0.8 else "no"
+        return Estimate(verdict, needed, size_gb, kv, 0.5, 0.0, 1 / t if verdict == "cpu" else None,
+                        ctx, "No GPU found: running from system RAM is slow.")
+
+    overhead = _overhead(hw)
+    needed = size_gb + kv + overhead
+    vram = gpu.total_vram_gb
+    note = ""
+    if needed <= vram * 0.92:
+        verdict = "fits"
+    elif needed <= vram:
+        verdict, note = "tight", "Barely fits: lower context or use a smaller quant if you hit OOM."
+    elif needed <= vram + hw.ram_gb * 0.7:
+        verdict, note = "offload", "Part of the model spills to system RAM, so speed drops a lot."
+    else:
+        verdict, note = "no", "Not enough VRAM + RAM."
+
+    weights_on_gpu = max(0.0, min(size_gb, vram * 0.95 - overhead - kv))
+    frac = weights_on_gpu / size_gb if size_gb else 1.0
+    if verdict == "no":
+        tok = None
+    else:
+        t = 0.0
+        for g in gpu.gpus:  # layers split by VRAM share; GPUs run in sequence
+            share = g.memory_total_mb / (vram * 1024)
+            t += active_bytes * frac * share / (g.bandwidth_gbs * GPU_EFF * eff_boost)
+        t += active_bytes * (1 - frac) / (CPU_BW_GBS * eff_boost)
+        t += 0.002 + 0.001 * (gpu.count - 1)
+        tok = 1 / t
+    if tok is not None and not all(g.bandwidth_known for g in gpu.gpus):
+        note = (note + " GPU model unknown: speed is a rough guess.").strip()
+    return Estimate(verdict, needed, size_gb, kv, overhead, frac, tok, ctx, note)

notebook_llm/estimate_data.py

@@ -0,0 +1,38 @@
+"""Static lookup tables (kept separate so gpu.py and estimate.py don't import each other)."""
+
+# (substring, memory bandwidth GB/s). First match wins -> keep specific names first.
+_GPU_BW = [
+    ("b200", 8000), ("h200", 4800), ("h100", 2400),
+    ("a100", 1555), ("l40", 864), ("a6000", 768), ("a5000", 768), ("a4000", 448),
+    ("a10g", 600), ("a10", 600), ("v100", 900), ("p100", 732), ("l4", 300), ("t4", 320),
+    ("rtx 4090", 1008), ("rtx 4080", 717), ("rtx 4070 ti", 504), ("rtx 4070", 504),
+    ("rtx 3090", 936), ("rtx 3080", 760), ("rtx 3070", 448), ("rtx 3060", 360),
+    ("rtx 2080 ti", 616), ("rtx 2080", 448), ("rtx 2070", 448), ("gtx 1080 ti", 484),
+    ("p4000", 243), ("p40", 346), ("k80", 240),
+]
+
+
+def gpu_bandwidth(name: str):
+    """Returns (GB/s, known?)."""
+    n = name.lower()
+    if "a100" in n and "80gb" in n:
+        return 2039, True
+    for key, bw in _GPU_BW:
+        if key in n:
+            return bw, True
+    return 300, False
+
+
+# MoE models: "name:tag" -> (total params B, active params B)
+MOE = {
+    "qwen3-coder:30b": (30.5, 3.3), "qwen3:30b": (30.5, 3.3), "qwen3:235b": (235, 22),
+    "qwen3-vl:30b": (31, 3.3), "gpt-oss:20b": (21, 3.6), "gpt-oss:120b": (117, 5.1),
+    "mixtral:8x7b": (46.7, 12.9), "mixtral:8x22b": (141, 39),
+    "deepseek-coder-v2:16b": (15.7, 2.4), "deepseek-coder-v2:236b": (236, 21),
+    "deepseek-v3:671b": (671, 37), "deepseek-r1:671b": (671, 37),
+    "llama4:scout": (109, 17), "llama4:maverick": (400, 17),
+}
+
+# Approx bytes per parameter for common GGUF quantizations.
+BYTES_PER_PARAM = [("fp16", 2.0), ("f16", 2.0), ("q8", 1.07), ("q6", 0.82), ("q5", 0.72),
+                   ("q4", 0.61), ("q3", 0.48), ("q2", 0.37)]

notebook_llm/gpu.py

@@ -0,0 +1,127 @@
+"""Hardware detection: NVIDIA GPUs (nvidia-smi), system RAM, bandwidth lookup."""
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from .estimate_data import gpu_bandwidth
+
+_QUERY = "index,name,memory.total,memory.free,driver_version"
+
+
+@dataclass
+class GPU:
+    index: int
+    name: str
+    memory_total_mb: int
+    memory_free_mb: int
+
+    @property
+    def vram_gb(self) -> float:
+        return self.memory_total_mb / 1024
+
+    @property
+    def bandwidth_gbs(self) -> float:
+        return gpu_bandwidth(self.name)[0]
+
+    @property
+    def bandwidth_known(self) -> bool:
+        return gpu_bandwidth(self.name)[1]
+
+
+@dataclass
+class GPUInfo:
+    gpus: List[GPU] = field(default_factory=list)
+    driver: Optional[str] = None
+
+    @property
+    def count(self) -> int:
+        return len(self.gpus)
+
+    @property
+    def available(self) -> bool:
+        return self.count > 0
+
+    @property
+    def total_vram_gb(self) -> float:
+        return round(sum(g.memory_total_mb for g in self.gpus) / 1024, 1)
+
+    @property
+    def used_vram_gb(self) -> float:
+        return round(sum(g.memory_total_mb - g.memory_free_mb for g in self.gpus) / 1024, 1)
+
+    @property
+    def indices(self) -> List[int]:
+        return [g.index for g in self.gpus]
+
+    def summary(self) -> str:
+        if not self.available:
+            return "No NVIDIA GPU detected (CPU mode)"
+        names = {}
+        for g in self.gpus:
+            names[g.name] = names.get(g.name, 0) + 1
+        label = " + ".join(f"{n}x {k}" for k, n in names.items())
+        return f"{label} | {self.total_vram_gb} GB VRAM"
+
+
+@dataclass
+class Hardware:
+    gpu: GPUInfo
+    ram_gb: float
+    cpu_cores: int
+
+    def summary(self) -> str:
+        return f"{self.gpu.summary()} | {self.ram_gb:.0f} GB RAM | {self.cpu_cores} CPU cores"
+
+
+def parse_nvidia_smi(output: str) -> GPUInfo:
+    gpus, driver = [], None
+    for line in output.strip().splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) < 5:
+            continue
+        try:
+            gpus.append(GPU(int(parts[0]), parts[1], int(float(parts[2])), int(float(parts[3]))))
+            driver = parts[4]
+        except ValueError:
+            continue
+    return GPUInfo(gpus, driver)
+
+
+def detect_gpus() -> GPUInfo:
+    if not shutil.which("nvidia-smi"):
+        return GPUInfo()
+    try:
+        out = subprocess.run(
+            ["nvidia-smi", f"--query-gpu={_QUERY}", "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=20, check=True).stdout
+    except (subprocess.SubprocessError, OSError):
+        return GPUInfo()
+    return parse_nvidia_smi(out)
+
+
+def system_ram_gb() -> float:
+    total = 0.0
+    try:
+        with open("/proc/meminfo") as f:
+            for line in f:
+                if line.startswith("MemTotal"):
+                    total = int(line.split()[1]) / 1024 / 1024
+                    break
+    except OSError:
+        pass
+    for p in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"):
+        try:
+            v = open(p).read().strip()
+            if v.isdigit():
+                lim = int(v) / 1024 ** 3
+                if 0 < lim < (total or 1e9):
+                    total = lim
+        except OSError:
+            pass
+    return round(total, 1)
+
+
+def detect_hardware() -> Hardware:
+    return Hardware(detect_gpus(), system_ram_gb(), os.cpu_count() or 1)

notebook_llm/installer.py

@@ -0,0 +1,67 @@
+"""Install system deps, Ollama and cloudflared (idempotent)."""
+import os
+import platform
+import shutil
+import stat
+import subprocess
+import urllib.request
+
+from .env import log
+
+BIN_DIR = "/tmp/notebook_llm/bin"
+CLOUDFLARED_URL = "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-{arch}"
+
+
+def _sh(cmd: str, check: bool = True) -> int:
+    return subprocess.run(cmd, shell=True, check=check,
+                          stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT).returncode
+
+
+def ollama_installed() -> bool:
+    return bool(shutil.which("ollama"))
+
+
+def cloudflared_path():
+    p = shutil.which("cloudflared")
+    if p:
+        return p
+    local = os.path.join(BIN_DIR, "cloudflared")
+    return local if os.path.exists(local) else None
+
+
+def ensure_system_deps() -> None:
+    """zstd is needed by the Ollama installer; pciutils/lshw let it detect GPUs."""
+    missing = [p for p, b in (("zstd", "zstd"), ("pciutils", "lspci"), ("lshw", "lshw")) if not shutil.which(b)]
+    if not missing:
+        return
+    log(f"Installing system packages: {', '.join(missing)}")
+    _sh("apt-get update -qq", check=False)
+    _sh(f"DEBIAN_FRONTEND=noninteractive apt-get install -y -qq {' '.join(missing)}", check=False)
+
+
+def install_ollama() -> str:
+    path = shutil.which("ollama")
+    if path:
+        return path
+    ensure_system_deps()
+    log("Installing Ollama (about 30s)...")
+    _sh("curl -fsSL https://ollama.com/install.sh | sh")
+    path = shutil.which("ollama")
+    if not path:
+        raise RuntimeError("Ollama install failed. Is Internet enabled in the notebook settings?")
+    return path
+
+
+def install_cloudflared() -> str:
+    path = cloudflared_path()
+    if path:
+        return path
+    arch = {"x86_64": "amd64", "amd64": "amd64", "aarch64": "arm64", "arm64": "arm64"}.get(platform.machine().lower())
+    if not arch:
+        raise RuntimeError(f"Unsupported architecture: {platform.machine()}")
+    os.makedirs(BIN_DIR, exist_ok=True)
+    local = os.path.join(BIN_DIR, "cloudflared")
+    log("Downloading cloudflared...")
+    urllib.request.urlretrieve(CLOUDFLARED_URL.format(arch=arch), local)
+    os.chmod(local, os.stat(local).st_mode | stat.S_IEXEC)
+    return local

notebook_llm/library.py

@@ -0,0 +1,212 @@
+"""Search the Ollama library (ollama.com) and resolve exact model sizes via the registry."""
+import html
+import re
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import requests
+
+from .estimate import Estimate, auto_context, estimate, guess_size_gb, parse_params
+
+UA = {"User-Agent": "notebook_llm/0.1 (+ollama library search)"}
+_size_cache: Dict[str, Optional[float]] = {}
+
+
+class LibraryError(RuntimeError):
+    pass
+
+
+@dataclass
+class LibModel:
+    name: str
+    description: str = ""
+    sizes: List[str] = field(default_factory=list)
+    pulls: str = ""
+    capabilities: List[str] = field(default_factory=list)
+    updated: str = ""
+    cloud: bool = False          # also/only offered on Ollama Cloud
+
+    @property
+    def local(self) -> bool:
+        """Can it be pulled and run on our own GPU?"""
+        return bool(self.sizes) or not self.cloud
+
+
+@dataclass
+class Candidate:
+    ref: str                 # e.g. qwen3-coder:30b
+    size_gb: float
+    exact: bool              # size from registry (True) or estimated from params (False)
+    params_b: Optional[float]
+    active_b: Optional[float]
+    est: Estimate
+
+
+FALLBACK = [
+    LibModel("qwen3-coder", "Alibaba's coding model (MoE)", ["30b"]),
+    LibModel("qwen2.5-coder", "Code-specific Qwen models", ["0.5b", "1.5b", "3b", "7b", "14b", "32b"]),
+    LibModel("qwen3", "Qwen3 general models", ["0.6b", "1.7b", "4b", "8b", "14b", "30b", "32b"]),
+    LibModel("gpt-oss", "OpenAI open-weight models", ["20b", "120b"]),
+    LibModel("llama3.1", "Meta Llama 3.1", ["8b", "70b"]),
+    LibModel("llama3.2", "Meta Llama 3.2 small models", ["1b", "3b"]),
+    LibModel("gemma3", "Google Gemma 3", ["1b", "4b", "12b", "27b"]),
+    LibModel("deepseek-r1", "Reasoning models", ["1.5b", "7b", "8b", "14b", "32b", "70b"]),
+    LibModel("deepseek-coder-v2", "Coding MoE model", ["16b"]),
+    LibModel("phi4", "Microsoft Phi-4", ["14b"]),
+    LibModel("mistral", "Mistral 7B", ["7b"]),
+    LibModel("codellama", "Meta Code Llama", ["7b", "13b", "34b", "70b"]),
+]
+
+
+def _clean(s: str) -> str:
+    return html.unescape(re.sub(r"<[^>]+>", "", s)).strip()
+
+
+_SIZE_RE = re.compile(r"^(\d+(?:\.\d+)?[bm]|\d+x\d+(?:\.\d+)?b)$")  # lowercase only: "270m" is a size, "7M" is a pull count
+_PULL_RE = re.compile(r"^[\d.,]+[KMB]?$")
+_NOT_CAPS = {"pulls", "tags", "pull", "tag", "updated"}
+
+
+def _finish(name, desc, spans, updated="") -> LibModel:
+    sizes, caps, pulls, cloud = [], [], "", False
+    for t in spans:
+        t = _clean(t)
+        if not t:
+            continue
+        if t.lower() == "cloud":
+            cloud = True
+        elif _SIZE_RE.match(t):
+            sizes.append(t)
+        elif _PULL_RE.match(t) and not pulls and any(ch.isdigit() for ch in t):
+            pulls = t
+        elif re.match(r"^[a-z][a-z-]{1,14}$", t) and t not in _NOT_CAPS:
+            caps.append(t)
+    return LibModel(name, desc, sizes, pulls, caps, updated, cloud)
+
+
+def _parse_xtest(page: str) -> List[LibModel]:
+    out = []
+    for block in page.split("<li x-test-model")[1:]:
+        href = re.search(r'href="/([^"?#]+)"', block)
+        if not href:
+            continue
+        name = re.sub(r"^library/", "", href.group(1))
+        desc = re.search(r"<p[^>]*max-w-lg[^>]*>(.*?)</p>", block, re.S)
+        spans = (re.findall(r"x-test-size[^>]*>(.*?)</span>", block, re.S)
+                 + re.findall(r"x-test-capability[^>]*>(.*?)</span>", block, re.S)
+                 + re.findall(r"x-test-pull-count[^>]*>(.*?)</span>", block, re.S))
+        upd = re.search(r"x-test-updated[^>]*>(.*?)</span>", block, re.S)
+        out.append(_finish(name, _clean(desc.group(1)) if desc else "", spans, _clean(upd.group(1)) if upd else ""))
+    return out
+
+
+def _parse_generic(page: str) -> List[LibModel]:
+    """Layout-agnostic: find /library/<name> links, then read the <span>/<p> tags that follow."""
+    hits, seen = [], set()
+    for m in re.finditer(r'href="/library/([^"/?#]+)"', page):
+        if m.group(1) not in seen:
+            seen.add(m.group(1))
+            hits.append((m.group(1), m.start()))
+    out = []
+    for i, (name, pos) in enumerate(hits):
+        block = page[pos: hits[i + 1][1] if i + 1 < len(hits) else len(page)]
+        desc = re.search(r"<p[^>]*>(.*?)</p>", block, re.S)
+        spans = re.findall(r"<span[^>]*>\s*([^<]{1,30}?)\s*</span>", block)
+        upd = re.search(r"(\d+\s+(?:day|week|month|year)s?\s+ago)", block)
+        out.append(_finish(name, _clean(desc.group(1)) if desc else "", spans, upd.group(1) if upd else ""))
+    return out
+
+
+def parse_search_html(page: str) -> List[LibModel]:
+    return _parse_xtest(page) or _parse_generic(page)
+
+
+def search(query: str = "", page: int = 1, sort: Optional[str] = None, capability: Optional[str] = None,
+           timeout: int = 12) -> List[LibModel]:
+    """Search ollama.com/search (20 per page). Falls back to a built-in list if unreachable.
+    sort: 'popular' | 'newest'.  capability: tools | vision | thinking | embedding | cloud."""
+    params = {"q": query}
+    if page > 1:
+        params["page"] = page
+    if sort:
+        params["o"] = sort
+    if capability:
+        params["c"] = capability
+    try:
+        r = requests.get("https://ollama.com/search", params=params, headers=UA, timeout=timeout)
+        r.raise_for_status()
+        res = parse_search_html(r.text)
+        if capability:  # in case the site ignores the filter
+            res = [m for m in res if capability in m.capabilities or (capability == "cloud" and m.cloud)] or res
+        return res
+    except requests.RequestException as e:
+        if page > 1:
+            return []
+        q = query.lower()
+        res = [LibModel(m.name, m.description + " [offline list]", list(m.sizes)) for m in FALLBACK
+               if q in m.name.lower() or q in m.description.lower()]
+        if not res:
+            raise LibraryError(f"Could not reach ollama.com ({e}) and nothing matched offline.")
+        return res
+
+
+def registry_size_gb(name: str, tag: str, timeout: int = 8) -> Optional[float]:
+    """Exact download size = sum of manifest layer sizes."""
+    key = f"{name}:{tag}"
+    if key in _size_cache:
+        return _size_cache[key]
+    repo = name if "/" in name else f"library/{name}"
+    size = None
+    try:
+        r = requests.get(f"https://registry.ollama.ai/v2/{repo}/manifests/{tag}",
+                         headers={**UA, "Accept": "application/vnd.docker.distribution.manifest.v2+json"},
+                         timeout=timeout)
+        if r.ok:
+            m = r.json()
+            total = sum(l.get("size", 0) for l in m.get("layers", [])) + m.get("config", {}).get("size", 0)
+            size = round(total / 1024 ** 3, 2) if total else None
+    except (requests.RequestException, ValueError):
+        pass
+    _size_cache[key] = size
+    return size
+
+
+def active_hint(description: str, tag_params: Optional[float], single_size: bool) -> Optional[float]:
+    """Pull 'A3B' / '3B active' out of a model description (new MoE models aren't in our table)."""
+    d = description.lower()
+    m = re.search(r"(\d+(?:\.\d+)?)b[- ]?a(\d+(?:\.\d+)?)b", d)
+    if m and (tag_params is None or abs(float(m.group(1)) - tag_params) <= 0.15 * tag_params):
+        return float(m.group(2))
+    m = re.search(r"(\d+(?:\.\d+)?)b\s+active", d)
+    if m and single_size:
+        return float(m.group(1))
+    return None
+
+
+def make_candidate(ref: str, hw, ctx: Optional[int] = None, kv_type: str = "q8_0",
+                   online: bool = True, hint: str = "", single_size: bool = False) -> Candidate:
+    name, _, tag = ref.partition(":")
+    tag = tag or "latest"
+    params, active = parse_params(name, tag)
+    if active is None and hint:
+        active = active_hint(hint, params, single_size)
+    size = registry_size_gb(name, tag) if online else None
+    exact = size is not None
+    if size is None:
+        if params is None:
+            params = 7.0  # unknown: assume a 7B
+        size = guess_size_gb(params, tag)
+    if params is None:
+        params = size / 0.61
+    c = ctx or auto_context(size, params, hw, kv_type)
+    return Candidate(f"{name}:{tag}", size, exact, params, active, estimate(size, params, active, hw, c, kv_type))
+
+
+def candidates_for(model: LibModel, hw, ctx=None, kv_type="q8_0", online=True) -> List[Candidate]:
+    tags = [s for s in model.sizes if "cloud" not in s] or ["latest"]
+    refs = [f"{model.name}:{t}" for t in tags]
+    one = len(tags) == 1
+    with ThreadPoolExecutor(max_workers=8) as ex:
+        res = list(ex.map(lambda r: make_candidate(r, hw, ctx, kv_type, online, model.description, one), refs))
+    return sorted(res, key=lambda c: c.size_gb)