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.
This commit is contained in:
Leonid Pershin
@@ -0,0 +1,165 @@
|
||||
#!/usr/bin/env python3
|
||||
"""First-boot Swarm/Comfy performance tune: sageattention ExtraArgs + pip libs.
|
||||
|
||||
Idempotent. Marker: /mnt/swarm_data/.gpu-rent-perf-tuned
|
||||
Re-runs if GPU uuid changed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
DATA = Path("/mnt/swarm_data")
|
||||
GPU_JSON = DATA / ".gpu-rent-gpu.json"
|
||||
MARKER = DATA / ".gpu-rent-perf-tuned"
|
||||
BACKENDS = DATA / "Data" / "Backends.fds"
|
||||
COMFY_VENV_CANDIDATES = [
|
||||
DATA / "dlbackend" / "comfy" / "venv" / "bin" / "pip",
|
||||
DATA / "dlbackend" / "comfy" / "ComfyUI" / "venv" / "bin" / "pip",
|
||||
Path("/opt/swarmui/dlbackend/comfy/venv/bin/pip"),
|
||||
Path("/opt/swarmui/dlbackend/comfy/ComfyUI/venv/bin/pip"),
|
||||
]
|
||||
|
||||
|
||||
def load_gpu() -> dict:
|
||||
if not GPU_JSON.is_file():
|
||||
return {}
|
||||
try:
|
||||
return json.loads(GPU_JSON.read_text(encoding="utf-8"))
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
|
||||
|
||||
def tier_notes(gpu: dict) -> dict:
|
||||
"""Mirror gpu_rent.perf_tiers.swarm_tune_for (stdlib-only on VM)."""
|
||||
vram = int(gpu.get("vram_mib") or 0)
|
||||
gib = vram / 1024.0
|
||||
if gib < 16:
|
||||
tier = "low"
|
||||
elif gib < 24:
|
||||
tier = "mid"
|
||||
elif gib < 48:
|
||||
tier = "high"
|
||||
else:
|
||||
tier = "ultra"
|
||||
cap = str(gpu.get("compute_cap") or "0.0")
|
||||
try:
|
||||
parts = cap.split(".")
|
||||
maj, mnr = int(parts[0]), int(parts[1]) if len(parts) > 1 else 0
|
||||
ampere = (maj, mnr) >= (8, 0)
|
||||
except ValueError:
|
||||
ampere = False
|
||||
use_sage = ampere and tier in {"mid", "high", "ultra"}
|
||||
return {
|
||||
"tier": tier,
|
||||
"use_sage": use_sage,
|
||||
"extra_args": "--use-sage-attention" if use_sage else "",
|
||||
"uuid": str(gpu.get("uuid") or ""),
|
||||
"name": str(gpu.get("name") or ""),
|
||||
}
|
||||
|
||||
|
||||
def find_pip() -> Path | None:
|
||||
for p in COMFY_VENV_CANDIDATES:
|
||||
if p.is_file():
|
||||
return p
|
||||
# glob
|
||||
for p in DATA.glob("dlbackend/**/venv/bin/pip"):
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def patch_backends_extra_args(extra: str) -> bool:
|
||||
"""Ensure ExtraArgs contains sage flag for Comfy self-start backends."""
|
||||
if not extra:
|
||||
return False
|
||||
if not BACKENDS.is_file():
|
||||
print(f"no {BACKENDS} yet — skip ExtraArgs (Comfy not registered)")
|
||||
return False
|
||||
text = BACKENDS.read_text(encoding="utf-8")
|
||||
if "--use-sage-attention" in text:
|
||||
print("Backends.fds already has --use-sage-attention")
|
||||
return False
|
||||
lines = text.splitlines()
|
||||
changed = False
|
||||
out = []
|
||||
for line in lines:
|
||||
if re.match(r"^(\s*)ExtraArgs:\s*$", line) or re.match(r"^(\s*)ExtraArgs:\s*\"\"\s*$", line):
|
||||
indent = re.match(r"^(\s*)", line).group(1)
|
||||
out.append(f"{indent}ExtraArgs: {extra}")
|
||||
changed = True
|
||||
elif re.match(r"^(\s*)ExtraArgs:\s+", line) and "--use-sage-attention" not in line:
|
||||
out.append(line.rstrip() + f" {extra}")
|
||||
changed = True
|
||||
else:
|
||||
out.append(line)
|
||||
if not changed:
|
||||
print("Backends.fds: no ExtraArgs field patched")
|
||||
return False
|
||||
BACKENDS.write_text("\n".join(out) + "\n", encoding="utf-8")
|
||||
print(f"patched {BACKENDS} ExtraArgs += {extra}")
|
||||
return True
|
||||
|
||||
|
||||
def pip_install_sage(pip: Path) -> None:
|
||||
print(f"pip install triton sageattention via {pip}")
|
||||
env = dict(os.environ)
|
||||
env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
|
||||
# Best-effort: do not fail whole tune if wheels missing for this torch.
|
||||
cmd = [str(pip), "install", "-U", "triton", "sageattention"]
|
||||
try:
|
||||
subprocess.check_call(cmd, env=env)
|
||||
print("triton + sageattention installed")
|
||||
except subprocess.CalledProcessError as exc:
|
||||
print(f"WARN: pip install failed ({exc}) — ExtraArgs may no-op until fixed")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
gpu = load_gpu()
|
||||
plan = tier_notes(gpu)
|
||||
prev = {}
|
||||
if MARKER.is_file():
|
||||
try:
|
||||
prev = json.loads(MARKER.read_text(encoding="utf-8"))
|
||||
except json.JSONDecodeError:
|
||||
prev = {}
|
||||
if prev.get("uuid") and prev.get("uuid") == plan["uuid"] and prev.get("extra_args") == plan["extra_args"]:
|
||||
if prev.get("pip_ok") or not plan["use_sage"]:
|
||||
print(f"perf tune already applied for {plan['name']} ({plan['tier']})")
|
||||
return 0
|
||||
|
||||
print(f"perf tune: {plan['name']} tier={plan['tier']} sage={plan['use_sage']}")
|
||||
pip_ok = False
|
||||
restarted_needed = False
|
||||
if plan["use_sage"]:
|
||||
pip = find_pip()
|
||||
if pip:
|
||||
pip_install_sage(pip)
|
||||
pip_ok = True
|
||||
else:
|
||||
print("Comfy venv pip not found yet — will retry next up")
|
||||
if patch_backends_extra_args(plan["extra_args"]):
|
||||
restarted_needed = True
|
||||
|
||||
marker = {
|
||||
"uuid": plan["uuid"],
|
||||
"name": plan["name"],
|
||||
"tier": plan["tier"],
|
||||
"extra_args": plan["extra_args"],
|
||||
"pip_ok": pip_ok or not plan["use_sage"],
|
||||
"restart_needed": restarted_needed,
|
||||
}
|
||||
MARKER.write_text(json.dumps(marker, indent=2) + "\n", encoding="utf-8")
|
||||
print("wrote", MARKER)
|
||||
if restarted_needed:
|
||||
print("RESTART_SWARMUI=1")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user