Leonid Pershin
79cb0a7e25
- Updated the `render_access_panel` function to conditionally hide LLM errors when the Ollama model is operational, improving user experience by reducing unnecessary error visibility. - Introduced a new `_notes_from_disk` function to streamline the retrieval of notes from disk, enhancing state management during GPU environment checks. - Refactored error handling in the `_bind_access` function to ensure that stack and GPU environment errors are accurately recorded and managed, improving robustness in session state updates. - Added tests to validate the new behavior of error handling and state management, ensuring that LLM errors are appropriately suppressed when conditions are met.
280 lines
10 KiB
Python
280 lines
10 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")
|
|
|
|
|
|
def test_triton_jit_ok_runs_a_py_file(tmp_path, monkeypatch):
|
|
mod = _load()
|
|
py = tmp_path / "python"
|
|
py.write_text("", encoding="utf-8")
|
|
seen: list[list[str]] = []
|
|
|
|
def fake_output(cmd, **kw):
|
|
seen.append(list(cmd))
|
|
return "triton JIT ok\n"
|
|
|
|
monkeypatch.setattr(mod.subprocess, "check_output", fake_output)
|
|
assert mod.triton_jit_ok(py) is True
|
|
assert seen
|
|
assert seen[0][0] == str(py)
|
|
assert seen[0][1].endswith(".py")
|
|
assert "-c" not in seen[0]
|