Leonid Pershin
5832c5cf75
- Updated the `provision_llm` function to utilize the `/api/tags` endpoint for verifying available models, improving accuracy in model management. - Introduced a new `already_have_ollama_tag` function to ensure exact tag matching, preventing mismatches during model checks. - Enhanced the `pull_stream` function to require a successful status from the API before proceeding, ensuring reliable model downloads. - Added logic to handle unwritten blob files, improving the robustness of the model pulling process. - Updated documentation and tests to reflect these changes, ensuring clarity and reliability in Ollama model operations.
262 lines
9.5 KiB
Python
262 lines
9.5 KiB
Python
"""Tests for remote tune_swarm_perf pip_ok / ExtraArgs gating."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
REMOTE = ROOT / "src" / "gpu_rent" / "remote" / "tune_swarm_perf.py"
|
|
|
|
|
|
def _load():
|
|
spec = importlib.util.spec_from_file_location("tune_swarm_perf", REMOTE)
|
|
assert spec and spec.loader
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
return mod
|
|
|
|
|
|
def test_pip_fail_skips_extra_args(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
data = tmp_path
|
|
backends = data / "Data" / "Backends.fds"
|
|
backends.parent.mkdir(parents=True)
|
|
backends.write_text("ExtraArgs: \n", encoding="utf-8")
|
|
gpu_json = data / ".gpu-rent-gpu.json"
|
|
gpu_json.write_text(
|
|
json.dumps(
|
|
{
|
|
"vram_mib": 24576,
|
|
"compute_cap": "8.9",
|
|
"uuid": "gpu-1",
|
|
"name": "RTX",
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
pip = data / "fake-pip"
|
|
pip.write_text("#!/bin/sh\n", encoding="utf-8")
|
|
|
|
monkeypatch.setattr(mod, "DATA", data)
|
|
monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
|
|
monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
monkeypatch.setattr(mod, "find_pip", lambda: pip)
|
|
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
|
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: False)
|
|
|
|
assert mod.main() == 0
|
|
marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
|
|
assert marker["pip_ok"] is False
|
|
assert marker["extra_args"] == ""
|
|
assert "--use-sage-attention" not in backends.read_text(encoding="utf-8")
|
|
|
|
|
|
def test_pip_ok_patches_extra_args(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
data = tmp_path
|
|
backends = data / "Data" / "Backends.fds"
|
|
backends.parent.mkdir(parents=True)
|
|
backends.write_text("ExtraArgs: \n", encoding="utf-8")
|
|
gpu_json = data / ".gpu-rent-gpu.json"
|
|
gpu_json.write_text(
|
|
json.dumps(
|
|
{
|
|
"vram_mib": 24576,
|
|
"compute_cap": "8.9",
|
|
"uuid": "gpu-1",
|
|
"name": "RTX",
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
pip = data / "fake-pip"
|
|
pip.write_text("#!/bin/sh\n", encoding="utf-8")
|
|
|
|
monkeypatch.setattr(mod, "DATA", data)
|
|
monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
|
|
monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
monkeypatch.setattr(mod, "find_pip", lambda: pip)
|
|
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
|
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
|
|
monkeypatch.setattr(mod, "triton_jit_ok", lambda _p: True)
|
|
|
|
assert mod.main() == 0
|
|
marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
|
|
assert marker["pip_ok"] is True
|
|
assert "--use-sage-attention" in marker["extra_args"]
|
|
assert "--use-sage-attention" in backends.read_text(encoding="utf-8")
|
|
|
|
|
|
def test_pip_fail_retries_next_run(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
data = tmp_path
|
|
backends = data / "Data" / "Backends.fds"
|
|
backends.parent.mkdir(parents=True)
|
|
backends.write_text("ExtraArgs: \n", encoding="utf-8")
|
|
gpu_json = data / ".gpu-rent-gpu.json"
|
|
gpu_json.write_text(
|
|
json.dumps(
|
|
{
|
|
"vram_mib": 24576,
|
|
"compute_cap": "8.9",
|
|
"uuid": "gpu-1",
|
|
"name": "RTX",
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
marker = data / ".gpu-rent-perf-tuned"
|
|
marker.write_text(
|
|
json.dumps(
|
|
{
|
|
"uuid": "gpu-1",
|
|
"extra_args": "",
|
|
"pip_ok": False,
|
|
"tier": "high",
|
|
"name": "RTX",
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
pip = data / "fake-pip"
|
|
pip.write_text("#!/bin/sh\n", encoding="utf-8")
|
|
|
|
monkeypatch.setattr(mod, "DATA", data)
|
|
monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
|
|
monkeypatch.setattr(mod, "MARKER", marker)
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
monkeypatch.setattr(mod, "find_pip", lambda: pip)
|
|
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
|
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
|
|
monkeypatch.setattr(mod, "triton_jit_ok", lambda _p: True)
|
|
|
|
assert mod.main() == 0
|
|
new_m = json.loads(marker.read_text(encoding="utf-8"))
|
|
assert new_m["pip_ok"] is True
|
|
assert "--use-sage-attention" in backends.read_text(encoding="utf-8")
|
|
|
|
|
|
def test_pip_install_uses_python_dash_m(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
py = tmp_path / "python"
|
|
py.write_text("", encoding="utf-8")
|
|
seen: list[list[str]] = []
|
|
|
|
def fake_call(cmd, env=None):
|
|
seen.append(list(cmd))
|
|
return 0
|
|
|
|
monkeypatch.setattr(mod, "sage_already_importable", lambda _p: False)
|
|
monkeypatch.setattr(mod.subprocess, "check_call", fake_call)
|
|
assert mod.pip_install_sage(py) is True
|
|
assert seen and seen[0][:4] == [str(py), "-m", "pip", "install"]
|
|
|
|
|
|
def test_patch_extra_args_fds_empty_x(tmp_path, monkeypatch):
|
|
"""FDS empty ExtraArgs is '\\x' — must replace, not append."""
|
|
mod = _load()
|
|
backends = tmp_path / "Backends.fds"
|
|
backends.write_text("0:\n\ttype: comfyui_selfstart\n\tExtraArgs: \\x\n", encoding="utf-8")
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
assert mod.patch_backends_extra_args("--use-sage-attention") is True
|
|
text = backends.read_text(encoding="utf-8")
|
|
assert "ExtraArgs: --use-sage-attention" in text
|
|
assert "\\x --use-sage" not in text
|
|
|
|
|
|
def test_sanitize_backends_fds_corruption(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
backends = tmp_path / "Backends.fds"
|
|
backends.write_text(
|
|
"ExtraArgs: \\x --use-sage-attention\n",
|
|
encoding="utf-8",
|
|
)
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
assert mod.sanitize_backends_fds() is True
|
|
assert backends.read_text(encoding="utf-8") == "ExtraArgs: --use-sage-attention\n"
|
|
|
|
|
|
def test_sanitize_does_not_eat_newline_on_bare_empty(tmp_path, monkeypatch):
|
|
"""Bare FDS empty ``\\x`` must not match — ``\\s+`` would merge next key."""
|
|
mod = _load()
|
|
backends = tmp_path / "Backends.fds"
|
|
original = "\tExtraArgs: \\x\n\tStartScript: /mnt/swarm_data/dlbackend/ComfyUI/main.py\n"
|
|
backends.write_text(original, encoding="utf-8")
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
assert mod.sanitize_backends_fds() is False
|
|
assert backends.read_text(encoding="utf-8") == original
|
|
|
|
|
|
def test_sanitize_drops_comment_garbage_extra_args(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
backends = tmp_path / "Backends.fds"
|
|
junk = (
|
|
'\tExtraArgs: "#If unchecked, the system will automatically add some '
|
|
'relevant arguments to the comfy launch."\n'
|
|
"\tStartScript: /mnt/x/main.py\n"
|
|
)
|
|
backends.write_text(junk, encoding="utf-8")
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
assert mod.sanitize_backends_fds() is True
|
|
text = backends.read_text(encoding="utf-8")
|
|
assert "unchecked" not in text
|
|
assert "StartScript: /mnt/x/main.py" in text
|
|
|
|
|
|
def test_ensure_absolute_start_script(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
data = tmp_path
|
|
main_py = data / "dlbackend" / "ComfyUI" / "main.py"
|
|
main_py.parent.mkdir(parents=True)
|
|
main_py.write_text("# comfy\n", encoding="utf-8")
|
|
backends = data / "Data" / "Backends.fds"
|
|
backends.parent.mkdir(parents=True)
|
|
backends.write_text(
|
|
"0:\n\ttype: comfyui_selfstart\n\tStartScript: dlbackend/ComfyUI/main.py\n",
|
|
encoding="utf-8",
|
|
)
|
|
monkeypatch.setattr(mod, "DATA", data)
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
assert mod.ensure_absolute_start_script() is True
|
|
text = backends.read_text(encoding="utf-8")
|
|
assert str(main_py.resolve()) in text
|
|
assert mod.ensure_absolute_start_script() is False
|
|
|
|
|
|
def test_jit_fail_strips_sage_extra_args(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
data = tmp_path
|
|
backends = data / "Data" / "Backends.fds"
|
|
backends.parent.mkdir(parents=True)
|
|
backends.write_text("ExtraArgs: --use-sage-attention\n", encoding="utf-8")
|
|
gpu_json = data / ".gpu-rent-gpu.json"
|
|
gpu_json.write_text(
|
|
json.dumps(
|
|
{
|
|
"vram_mib": 24576,
|
|
"compute_cap": "8.9",
|
|
"uuid": "gpu-1",
|
|
"name": "RTX",
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
pip = data / "fake-pip"
|
|
pip.write_text("#!/bin/sh\n", encoding="utf-8")
|
|
monkeypatch.setattr(mod, "DATA", data)
|
|
monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
|
|
monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
|
|
monkeypatch.setattr(mod, "BACKENDS", backends)
|
|
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
|
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
|
|
monkeypatch.setattr(mod, "triton_jit_ok", lambda _p: False)
|
|
assert mod.main() == 0
|
|
marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
|
|
assert marker["pip_ok"] is True
|
|
assert marker["jit_ok"] is False
|
|
assert "--use-sage-attention" not in backends.read_text(encoding="utf-8")
|