gpu-rent/src/gpu_rent/perf_tiers.py

"""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
    context_length: 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"
        ctx = 32768
        note = "ultra: flash+q8 KV, 20GiB reserved for Swarm, ctx 32k, keep 30m"
    elif info.tier == TIER_HIGH:
        overhead = 14 * 1024**3
        keep = "15m"
        kv = "q8_0"
        ctx = 16384
        note = "high: flash+q8 KV, 14GiB reserved for Swarm, ctx 16k, keep 15m"
    elif info.tier == TIER_MID:
        overhead = 10 * 1024**3
        keep = "5m"
        kv = "q8_0"
        ctx = 16384
        note = "mid: flash+q8 KV, 10GiB reserved for Swarm, ctx 16k, keep 5m"
    else:
        overhead = 6 * 1024**3
        keep = "2m"
        kv = "q4_0"
        ctx = 8192
        flash = False  # prefer stability on tiny cards
        note = "low: conservative, 6GiB reserved, ctx 8k, 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,
        context_length=ctx,
        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}",
        f"Environment=OLLAMA_CONTEXT_LENGTH={tune.context_length}",
    ]
    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