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
+135
-10
@@ -1,6 +1,7 @@
|
||||
"""Unit tests for capture merge / URL builders (no SSH)."""
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from gpu_rent.capture import (
|
||||
ExtCaptureItem,
|
||||
@@ -11,6 +12,7 @@ from gpu_rent.capture import (
|
||||
strip_git_auth,
|
||||
)
|
||||
from gpu_rent.civitai import civitai_model_url, version_ids_from_payload
|
||||
from gpu_rent.errors import CloudError
|
||||
from gpu_rent.manifests import parse_extensions, parse_models
|
||||
|
||||
|
||||
@@ -29,22 +31,90 @@ def test_version_ids_from_payload():
|
||||
|
||||
|
||||
def test_resolve_from_sidecar():
|
||||
item = resolve_model_item(
|
||||
out = resolve_model_item(
|
||||
{
|
||||
"kind": "lora",
|
||||
"rel": "Lora/foo.safetensors",
|
||||
"name": "foo.safetensors",
|
||||
"version_id": 3107521,
|
||||
"model_id": 2187487,
|
||||
"sha256": "a" * 64,
|
||||
"sha256": None,
|
||||
},
|
||||
token="",
|
||||
api_host="civitai.red",
|
||||
link_host="civitai.red",
|
||||
)
|
||||
assert item is not None
|
||||
assert item.version_id == 3107521
|
||||
assert "modelVersionId=3107521" in item.url
|
||||
assert out.status == "ok"
|
||||
assert out.item is not None
|
||||
assert out.item.version_id == 3107521
|
||||
assert "modelVersionId=3107521" in out.item.url
|
||||
|
||||
|
||||
def test_resolve_vid_only_fetches_model_id():
|
||||
with patch(
|
||||
"gpu_rent.capture.fetch_model_version",
|
||||
return_value=("civitai.red", {"id": 10, "modelId": 20, "name": "X"}),
|
||||
) as mocked:
|
||||
out = resolve_model_item(
|
||||
{
|
||||
"kind": "lora",
|
||||
"rel": "Lora/x.safetensors",
|
||||
"name": "x.safetensors",
|
||||
"version_id": 10,
|
||||
"model_id": None,
|
||||
"sha256": None,
|
||||
},
|
||||
token="tok",
|
||||
api_host="civitai.red",
|
||||
link_host="civitai.red",
|
||||
)
|
||||
mocked.assert_called_once()
|
||||
assert out.status == "ok"
|
||||
assert out.item is not None
|
||||
assert out.item.model_id == 20
|
||||
assert "models/20?modelVersionId=10" in out.item.url
|
||||
|
||||
|
||||
def test_resolve_by_hash_404_is_unknown():
|
||||
with patch(
|
||||
"gpu_rent.capture.fetch_model_version_by_hash",
|
||||
side_effect=CloudError("Civitai by-hash abc: HTTP 404 (хосты civitai.red)"),
|
||||
):
|
||||
out = resolve_model_item(
|
||||
{
|
||||
"kind": "lora",
|
||||
"rel": "Lora/m.safetensors",
|
||||
"name": "m.safetensors",
|
||||
"sha256": "a" * 64,
|
||||
},
|
||||
token="",
|
||||
api_host="civitai.red",
|
||||
link_host="civitai.red",
|
||||
)
|
||||
assert out.status == "unknown"
|
||||
assert out.item is None
|
||||
|
||||
|
||||
def test_resolve_by_hash_network_is_api_error():
|
||||
with patch(
|
||||
"gpu_rent.capture.fetch_model_version_by_hash",
|
||||
side_effect=CloudError(
|
||||
"Civitai by-hash abc: Connection timeout (хосты civitai.red)"
|
||||
),
|
||||
):
|
||||
out = resolve_model_item(
|
||||
{
|
||||
"kind": "lora",
|
||||
"rel": "Lora/m.safetensors",
|
||||
"name": "m.safetensors",
|
||||
"sha256": "b" * 64,
|
||||
},
|
||||
token="",
|
||||
api_host="civitai.red",
|
||||
link_host="civitai.red",
|
||||
)
|
||||
assert out.status == "api_error"
|
||||
assert "timeout" in out.detail.lower() or "Connection" in out.detail
|
||||
|
||||
|
||||
def test_merge_models_dedupe(tmp_path: Path):
|
||||
@@ -54,9 +124,19 @@ def test_merge_models_dedupe(tmp_path: Path):
|
||||
encoding="utf-8",
|
||||
)
|
||||
items = [
|
||||
ModelCaptureItem("lora", 100, 1, "https://civitai.red/models/1?modelVersionId=100", "old"),
|
||||
ModelCaptureItem("lora", 200, 2, "https://civitai.red/models/2?modelVersionId=200", "new"),
|
||||
ModelCaptureItem("checkpoint", 300, 3, "https://civitai.red/models/3?modelVersionId=300", "ckpt"),
|
||||
ModelCaptureItem(
|
||||
"lora", 100, 1, "https://civitai.red/models/1?modelVersionId=100", "old"
|
||||
),
|
||||
ModelCaptureItem(
|
||||
"lora", 200, 2, "https://civitai.red/models/2?modelVersionId=200", "new"
|
||||
),
|
||||
ModelCaptureItem(
|
||||
"checkpoint",
|
||||
300,
|
||||
3,
|
||||
"https://civitai.red/models/3?modelVersionId=300",
|
||||
"ckpt",
|
||||
),
|
||||
]
|
||||
added, skipped = merge_models_yaml(path, items, dry_run=False)
|
||||
assert len(added) == 2
|
||||
@@ -72,13 +152,38 @@ def test_merge_models_dry_run(tmp_path: Path):
|
||||
path.write_text("lora: []\n", encoding="utf-8")
|
||||
before = path.read_text(encoding="utf-8")
|
||||
items = [
|
||||
ModelCaptureItem("lora", 1, 1, "https://civitai.red/models/1?modelVersionId=1", "x"),
|
||||
ModelCaptureItem(
|
||||
"lora", 1, 1, "https://civitai.red/models/1?modelVersionId=1", "x"
|
||||
),
|
||||
]
|
||||
added, _ = merge_models_yaml(path, items, dry_run=True)
|
||||
assert len(added) == 1
|
||||
assert path.read_text(encoding="utf-8") == before
|
||||
|
||||
|
||||
def test_merge_models_kind_scoped_dedupe(tmp_path: Path):
|
||||
"""Same version_id under different kinds can both be kept."""
|
||||
path = tmp_path / "models.yaml"
|
||||
path.write_text(
|
||||
"lora:\n - url: https://civitai.red/models/1?modelVersionId=100\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
items = [
|
||||
ModelCaptureItem(
|
||||
"checkpoint",
|
||||
100,
|
||||
1,
|
||||
"https://civitai.red/models/1?modelVersionId=100",
|
||||
"as-ckpt",
|
||||
),
|
||||
]
|
||||
added, skipped = merge_models_yaml(path, items, dry_run=False)
|
||||
assert len(added) == 1
|
||||
assert skipped == []
|
||||
entries = parse_models(path)
|
||||
assert {(e.kind, e.version_id) for e in entries} == {("lora", 100), ("checkpoint", 100)}
|
||||
|
||||
|
||||
def test_merge_extensions_dedupe(tmp_path: Path):
|
||||
path = tmp_path / "extensions.yaml"
|
||||
path.write_text(
|
||||
@@ -95,8 +200,9 @@ def test_merge_extensions_dedupe(tmp_path: Path):
|
||||
"C",
|
||||
),
|
||||
]
|
||||
added, skipped = merge_extensions_yaml(path, items, dry_run=False)
|
||||
added, updated, skipped = merge_extensions_yaml(path, items, dry_run=False)
|
||||
assert len(added) == 2
|
||||
assert updated == []
|
||||
assert len(skipped) == 1
|
||||
repos = parse_extensions(path)
|
||||
urls = [r.url for r in repos]
|
||||
@@ -104,6 +210,25 @@ def test_merge_extensions_dedupe(tmp_path: Path):
|
||||
assert all("SECRET" not in u for u in urls)
|
||||
|
||||
|
||||
def test_merge_extensions_updates_url_same_dir(tmp_path: Path):
|
||||
path = tmp_path / "extensions.yaml"
|
||||
path.write_text(
|
||||
"swarmui:\n - url: https://github.com/old/A.git\n ref: main\n dir: A\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
items = [
|
||||
ExtCaptureItem("swarmui", "https://github.com/new/A.git", "develop", "A"),
|
||||
]
|
||||
added, updated, skipped = merge_extensions_yaml(path, items, dry_run=False)
|
||||
assert added == []
|
||||
assert len(updated) == 1
|
||||
assert skipped == []
|
||||
repos = parse_extensions(path)
|
||||
assert len(repos) == 1
|
||||
assert repos[0].url == "https://github.com/new/A.git"
|
||||
assert repos[0].ref == "develop"
|
||||
|
||||
|
||||
def test_strip_git_auth():
|
||||
assert (
|
||||
strip_git_auth("https://x-access-token:tok@github.com/org/r.git")
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
from gpu_rent.perf_tiers import (
|
||||
GpuInfo,
|
||||
compute_cap_at_least,
|
||||
ollama_env_lines,
|
||||
ollama_tune_for,
|
||||
swarm_tune_for,
|
||||
tier_for_vram_mib,
|
||||
)
|
||||
|
||||
|
||||
def test_tier_boundaries():
|
||||
assert tier_for_vram_mib(8 * 1024) == "low"
|
||||
assert tier_for_vram_mib(16 * 1024) == "mid"
|
||||
assert tier_for_vram_mib(23 * 1024) == "mid"
|
||||
assert tier_for_vram_mib(24 * 1024) == "high"
|
||||
assert tier_for_vram_mib(40 * 1024) == "high"
|
||||
assert tier_for_vram_mib(48 * 1024) == "ultra"
|
||||
assert tier_for_vram_mib(80 * 1024) == "ultra"
|
||||
|
||||
|
||||
def test_compute_cap():
|
||||
assert compute_cap_at_least("8.0", 8, 0)
|
||||
assert compute_cap_at_least("8.9", 8, 0)
|
||||
assert not compute_cap_at_least("7.5", 8, 0)
|
||||
|
||||
|
||||
def test_ollama_high_reserves_vram_for_swarm():
|
||||
info = GpuInfo(
|
||||
name="NVIDIA A100-SXM4-40GB",
|
||||
vram_mib=40960,
|
||||
compute_cap="8.0",
|
||||
uuid="GPU-1",
|
||||
tier="high",
|
||||
)
|
||||
tune = ollama_tune_for(info)
|
||||
assert tune.flash_attention
|
||||
assert tune.num_parallel == 1
|
||||
assert tune.max_loaded_models == 1
|
||||
assert tune.kv_cache_type == "q8_0"
|
||||
assert tune.gpu_overhead_bytes == 14 * 1024**3
|
||||
env = "\n".join(ollama_env_lines(tune))
|
||||
assert "OLLAMA_FLASH_ATTENTION=1" in env
|
||||
assert "OLLAMA_GPU_OVERHEAD=" in env
|
||||
|
||||
|
||||
def test_swarm_sage_on_ampere_mid():
|
||||
info = GpuInfo(
|
||||
name="NVIDIA GeForce RTX 4090",
|
||||
vram_mib=24576,
|
||||
compute_cap="8.9",
|
||||
uuid="GPU-2",
|
||||
tier="high",
|
||||
)
|
||||
st = swarm_tune_for(info)
|
||||
assert st.use_sage_attention
|
||||
assert "--use-sage-attention" in st.comfy_extra_args
|
||||
assert st.install_triton_sage
|
||||
|
||||
|
||||
def test_swarm_no_sage_on_low():
|
||||
info = GpuInfo(
|
||||
name="NVIDIA T4",
|
||||
vram_mib=15360,
|
||||
compute_cap="7.5",
|
||||
uuid="GPU-3",
|
||||
tier="low",
|
||||
)
|
||||
st = swarm_tune_for(info)
|
||||
assert not st.use_sage_attention
|
||||
assert st.comfy_extra_args == ""
|
||||
Reference in New Issue
Block a user