Enhance GPU probing and performance tuning in provisioning

- Introduced GPU probing functionality to gather and store GPU specifications in `/mnt/swarm_data/.gpu-rent-gpu.json`, aiding in performance tuning. - Updated `install_ollama.sh` and `install_llamacpp.sh` to utilize GPU information for configuring optimal runtime parameters. - Enhanced `provision.py` to include GPU probing and performance tuning logic, ensuring better resource allocation for LLM operations. - Improved documentation in `decisions.md`, `llm.md`, and `swarmui.md` to reflect changes in GPU handling and performance tuning processes. - Added new tests to validate the GPU probing and model resolution logic, ensuring robustness in handling various GPU configurations.
2026-08-21 06:10:24 +03:00
parent 603165a4ba
commit 2ccb03f7d2
16 changed files with 1270 additions and 138 deletions
@@ -0,0 +1,147 @@
+"""GPU performance tiers for SwarmUI + Ollama auto-tune."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+# Shared GPU with SwarmUI: reserve VRAM so image gen still fits (Krea2 Turbo fp8 ~10–14GB).
+TIER_LOW = "low"  # <16 GiB
+TIER_MID = "mid"  # 16–23
+TIER_HIGH = "high"  # 24–47
+TIER_ULTRA = "ultra"  # ≥48
+
+
+@dataclass(frozen=True)
+class GpuInfo:
+    name: str
+    vram_mib: int
+    compute_cap: str  # e.g. "8.9"
+    uuid: str
+    tier: str
+
+
+@dataclass(frozen=True)
+class OllamaTune:
+    flash_attention: bool
+    keep_alive: str
+    num_parallel: int
+    max_loaded_models: int
+    kv_cache_type: str | None
+    gpu_overhead_bytes: int
+    notes: str
+
+
+@dataclass(frozen=True)
+class SwarmTune:
+    use_sage_attention: bool
+    install_triton_sage: bool
+    comfy_extra_args: str
+    notes: str
+
+
+def tier_for_vram_mib(vram_mib: int) -> str:
+    gib = vram_mib / 1024.0
+    if gib < 16:
+        return TIER_LOW
+    if gib < 24:
+        return TIER_MID
+    if gib < 48:
+        return TIER_HIGH
+    return TIER_ULTRA
+
+
+def compute_cap_at_least(cap: str, major: int, minor: int = 0) -> bool:
+    """True if NVIDIA compute capability >= major.minor (Ampere=8.0)."""
+    try:
+        parts = str(cap).strip().split(".")
+        maj = int(parts[0])
+        mnr = int(parts[1]) if len(parts) > 1 else 0
+        return (maj, mnr) >= (major, minor)
+    except (TypeError, ValueError):
+        return False
+
+
+def ollama_tune_for(info: GpuInfo) -> OllamaTune:
+    """Tune Ollama for prompt-help beside SwarmUI (share one GPU)."""
+    ampere_plus = compute_cap_at_least(info.compute_cap, 8, 0)
+    flash = ampere_plus or "A100" in info.name.upper() or "H100" in info.name.upper()
+    # Reserve VRAM for Comfy/Krea so Ollama does not fill the card.
+    if info.tier == TIER_ULTRA:
+        overhead = 20 * 1024**3
+        keep = "30m"
+        kv = "q8_0"
+        note = "ultra: flash+q8 KV, 20GiB reserved for Swarm, keep 30m"
+    elif info.tier == TIER_HIGH:
+        overhead = 14 * 1024**3
+        keep = "15m"
+        kv = "q8_0"
+        note = "high: flash+q8 KV, 14GiB reserved for Swarm, keep 15m"
+    elif info.tier == TIER_MID:
+        overhead = 10 * 1024**3
+        keep = "5m"
+        kv = "q8_0"
+        note = "mid: flash+q8 KV, 10GiB reserved for Swarm, keep 5m"
+    else:
+        overhead = 6 * 1024**3
+        keep = "2m"
+        kv = "q4_0"
+        flash = False  # prefer stability on tiny cards
+        note = "low: conservative, 6GiB reserved, short keep-alive"
+
+    return OllamaTune(
+        flash_attention=flash,
+        keep_alive=keep,
+        num_parallel=1,
+        max_loaded_models=1,
+        kv_cache_type=kv,
+        gpu_overhead_bytes=overhead,
+        notes=note,
+    )
+
+
+def swarm_tune_for(info: GpuInfo) -> SwarmTune:
+    """Swarm/Comfy launch: speed without quality loss (sage on Ampere+)."""
+    ampere_plus = compute_cap_at_least(info.compute_cap, 8, 0)
+    # SageAttention: Linux + Triton; safe quality, faster attention (Swarm docs).
+    use_sage = ampere_plus and info.tier in {TIER_MID, TIER_HIGH, TIER_ULTRA}
+    if info.tier == TIER_LOW:
+        use_sage = False
+    args = "--use-sage-attention" if use_sage else ""
+    return SwarmTune(
+        use_sage_attention=use_sage,
+        install_triton_sage=use_sage,
+        comfy_extra_args=args,
+        notes=(
+            "sage-attention + triton (Ampere+)"
+            if use_sage
+            else "stock Comfy attention (low VRAM or pre-Ampere)"
+        ),
+    )
+
+
+def parse_probe_dict(data: dict[str, Any]) -> GpuInfo:
+    vram = int(data.get("vram_mib") or 0)
+    tier = str(data.get("tier") or tier_for_vram_mib(vram))
+    return GpuInfo(
+        name=str(data.get("name") or "unknown"),
+        vram_mib=vram,
+        compute_cap=str(data.get("compute_cap") or "0.0"),
+        uuid=str(data.get("uuid") or ""),
+        tier=tier,
+    )
+
+
+def ollama_env_lines(tune: OllamaTune) -> list[str]:
+    lines = [
+        f"Environment=OLLAMA_NUM_PARALLEL={tune.num_parallel}",
+        f"Environment=OLLAMA_MAX_LOADED_MODELS={tune.max_loaded_models}",
+        f"Environment=OLLAMA_KEEP_ALIVE={tune.keep_alive}",
+        f"Environment=OLLAMA_GPU_OVERHEAD={tune.gpu_overhead_bytes}",
+    ]
+    if tune.flash_attention:
+        lines.append("Environment=OLLAMA_FLASH_ATTENTION=1")
+    if tune.kv_cache_type:
+        lines.append(f"Environment=OLLAMA_KV_CACHE_TYPE={tune.kv_cache_type}")
+    return lines