Leonid Pershin
f437cd0373
- Clarified the behavior of `Ctrl+C` and `Ctrl+D` in the README and other documentation, specifying that `Ctrl+C` only stops the tunnel while keeping the GPU active, and `Ctrl+D` stops the GPU while preserving disk data. - Enhanced the CLI documentation to reflect these changes, ensuring users understand the implications of these commands during GPU operations. - Improved the handling of data bindings and remounting logic in the codebase to prevent issues with empty model tabs in the UI. - Added tests to validate the new command behaviors and ensure proper documentation alignment.
277 lines
9.5 KiB
Python
277 lines
9.5 KiB
Python
"""Warm re-up shortcuts: no git cascade, killer reuse, seed/ollama skips."""
|
|
|
|
from pathlib import Path
|
|
|
|
from gpu_rent.idle_killer import _remote_killer_reusable
|
|
from gpu_rent.provision import _models_manifest_fp, seed_civitai
|
|
|
|
|
|
def test_models_manifest_fp_stable(tmp_path):
|
|
man = tmp_path / "models.yaml"
|
|
man.write_text("models:\n - kind: checkpoint\n", encoding="utf-8")
|
|
|
|
class Cfg:
|
|
models_manifest = man
|
|
|
|
a = _models_manifest_fp(Cfg())
|
|
b = _models_manifest_fp(Cfg())
|
|
assert a == b
|
|
assert len(a) == 16
|
|
man.write_text("models:\n - kind: lora\n", encoding="utf-8")
|
|
assert _models_manifest_fp(Cfg()) != a
|
|
|
|
|
|
def test_seed_civitai_skips_when_cache_dests_present(tmp_path, monkeypatch):
|
|
man = tmp_path / "models.yaml"
|
|
man.write_text("models:\n - kind: checkpoint\n", encoding="utf-8")
|
|
from gpu_rent import provision
|
|
|
|
class Entry:
|
|
kind = "checkpoint"
|
|
url = "https://huggingface.co/org/model/resolve/main/a.safetensors"
|
|
version_id = None
|
|
|
|
monkeypatch.setattr(provision, "parse_models", lambda _p: [Entry()])
|
|
|
|
def fake_ssh(cfg, host, cmd, **kw):
|
|
if ".gpu-rent-models-jobs.json" in cmd and "cat" in cmd:
|
|
import json
|
|
|
|
return json.dumps(
|
|
{
|
|
"fp": provision._models_manifest_fp(cfg),
|
|
"dests": ["/mnt/swarm_data/Models/Stable-Diffusion/a.safetensors"],
|
|
}
|
|
)
|
|
if "ALL_OK" in cmd or "test -f" in cmd:
|
|
return "ALL_OK"
|
|
return ""
|
|
|
|
class Cfg:
|
|
models_manifest = man
|
|
civitai_api_token = "t"
|
|
civitai_api_host = "https://civitai.com"
|
|
hf_token = ""
|
|
|
|
logs = []
|
|
monkeypatch.setattr(provision, "run_ssh", fake_ssh)
|
|
seed_civitai(Cfg(), "1.2.3.4", logs.append)
|
|
assert any("skip" in x.lower() for x in logs)
|
|
|
|
|
|
def test_remote_killer_reusable_parses_ok(monkeypatch):
|
|
class Cfg:
|
|
pass
|
|
|
|
def fake_ssh(cfg, host, cmd, **kw):
|
|
return (
|
|
"sid=abc-123\n"
|
|
"want=abc-123\n"
|
|
"creds=yes\n"
|
|
"armed=yes\n"
|
|
"rootcreds=yes\n"
|
|
"timer=active\n"
|
|
)
|
|
|
|
from gpu_rent import idle_killer
|
|
|
|
monkeypatch.setattr(idle_killer, "run_ssh", fake_ssh)
|
|
assert _remote_killer_reusable(Cfg(), "h", "abc-123") is True
|
|
assert _remote_killer_reusable(Cfg(), "h", "other") is False
|
|
|
|
|
|
def test_install_ollama_skips_restart_when_unit_unchanged():
|
|
from importlib.resources import files
|
|
|
|
text = files("gpu_rent.remote").joinpath("install_ollama.sh").read_text(encoding="utf-8")
|
|
assert "cmp -s" in text
|
|
assert "skip restart" in text
|
|
|
|
|
|
def test_cli_has_update_flag():
|
|
from gpu_rent import cli
|
|
|
|
src = Path(cli.__file__).read_text(encoding="utf-8")
|
|
assert '"--update"' in src
|
|
assert "force_update" in src
|
|
|
|
|
|
def test_autocomplete_merge_sets_is_installed():
|
|
from gpu_rent import provision
|
|
from gpu_rent.provision import _AUTOCOMPLETE_MERGE_PY, _ENSURE_INSTALLED_PY
|
|
|
|
assert "IsInstalled: true" in _AUTOCOMPLETE_MERGE_PY
|
|
assert "IsInstalled: true" in _ENSURE_INSTALLED_PY
|
|
assert "comfy_on_disk" in _AUTOCOMPLETE_MERGE_PY
|
|
assert "comfy_on_disk" in _ENSURE_INSTALLED_PY
|
|
assert r"\1false" in _AUTOCOMPLETE_MERGE_PY
|
|
assert r"\1false" in _ENSURE_INSTALLED_PY
|
|
assert "GPU_RENT_COMFY_PRESENT" in _AUTOCOMPLETE_MERGE_PY
|
|
assert "GPU_RENT_COMFY_PRESENT" in _ENSURE_INSTALLED_PY
|
|
assert hasattr(provision, "ensure_settings_is_installed")
|
|
assert "set_autocomplete_source" in _AUTOCOMPLETE_MERGE_PY
|
|
assert "CHANGED inserted AutoComplete.Source" in _AUTOCOMPLETE_MERGE_PY
|
|
assert '"settings_applied": False' in Path(provision.__file__).read_text(
|
|
encoding="utf-8"
|
|
)
|
|
assert "if not applied:" not in Path(provision.__file__).read_text(encoding="utf-8")
|
|
|
|
|
|
def test_ensure_installed_clears_flag_without_comfy(tmp_path):
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
|
|
from gpu_rent.provision import _ENSURE_INSTALLED_PY
|
|
|
|
settings = tmp_path / "Settings.fds"
|
|
settings.write_text("IsInstalled: true\nTheme: modern_dark\n", encoding="utf-8")
|
|
script = tmp_path / "ensure.py"
|
|
script.write_text(_ENSURE_INSTALLED_PY, encoding="utf-8")
|
|
env = os.environ.copy()
|
|
env["GPU_RENT_SETTINGS_FDS"] = str(settings)
|
|
env["GPU_RENT_COMFY_PRESENT"] = "0"
|
|
out = subprocess.check_output([sys.executable, str(script)], env=env, text=True)
|
|
assert "IsInstalled false" in out or "cleared" in out
|
|
text = settings.read_text(encoding="utf-8")
|
|
assert re.search(r"(?im)^\s*IsInstalled:\s*false\s*$", text)
|
|
assert "Theme: modern_dark" in text
|
|
|
|
|
|
def test_ensure_installed_sets_true_when_comfy_present(tmp_path):
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
|
|
from gpu_rent.provision import _ENSURE_INSTALLED_PY
|
|
|
|
settings = tmp_path / "Settings.fds"
|
|
settings.write_text("IsInstalled: false\n", encoding="utf-8")
|
|
script = tmp_path / "ensure.py"
|
|
script.write_text(_ENSURE_INSTALLED_PY, encoding="utf-8")
|
|
env = os.environ.copy()
|
|
env["GPU_RENT_SETTINGS_FDS"] = str(settings)
|
|
env["GPU_RENT_COMFY_PRESENT"] = "1"
|
|
subprocess.check_call([sys.executable, str(script)], env=env)
|
|
text = settings.read_text(encoding="utf-8")
|
|
assert re.search(r"(?im)^\s*IsInstalled:\s*true\s*$", text)
|
|
|
|
|
|
def test_install_comfy_patches_is_installed_on_skip():
|
|
from importlib.resources import files
|
|
|
|
text = files("gpu_rent.remote").joinpath("install_swarm_comfy.py").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
assert "ensure_settings_installed_flag" in text
|
|
assert "leave /Install" in text
|
|
assert "prepare_first_install" in text
|
|
assert "already installed" in text
|
|
assert "IsInstalled=false" in text
|
|
assert "Open SwarmUI" not in text
|
|
|
|
|
|
def test_clear_settings_installed_flag(tmp_path, monkeypatch):
|
|
from gpu_rent.remote import install_swarm_comfy as inst
|
|
|
|
settings = tmp_path / "Settings.fds"
|
|
settings.write_text("IsInstalled: true\nTheme: x\n", encoding="utf-8")
|
|
monkeypatch.setattr(inst, "SETTINGS", settings)
|
|
assert inst.clear_settings_installed_flag() is True
|
|
text = settings.read_text(encoding="utf-8")
|
|
assert "IsInstalled: false" in text
|
|
assert "Theme: x" in text
|
|
assert inst.clear_settings_installed_flag() is False
|
|
|
|
|
|
def _run_autocomplete_merge(tmp_path, settings_text: str, fname: str = "danbooru.csv"):
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
from gpu_rent.provision import _AUTOCOMPLETE_MERGE_PY
|
|
|
|
settings = tmp_path / "Settings.fds"
|
|
settings.write_text(settings_text, encoding="utf-8")
|
|
script = tmp_path / "merge.py"
|
|
script.write_text(_AUTOCOMPLETE_MERGE_PY, encoding="utf-8")
|
|
env = os.environ.copy()
|
|
env["GPU_RENT_SETTINGS_FDS"] = str(settings)
|
|
env["GPU_RENT_AUTOCOMPLETE_FILE"] = fname
|
|
env["GPU_RENT_COMFY_PRESENT"] = "1"
|
|
out = subprocess.check_output([sys.executable, str(script)], env=env, text=True)
|
|
return out, settings.read_text(encoding="utf-8")
|
|
|
|
|
|
def test_autocomplete_merge_inserts_source_when_only_escapeparens(tmp_path):
|
|
out, text = _run_autocomplete_merge(
|
|
tmp_path,
|
|
"IsInstalled: true\nDefaultUser:\n AutoComplete:\n EscapeParens: true\n",
|
|
)
|
|
assert "CHANGED inserted AutoComplete.Source=danbooru.csv" in out
|
|
assert "Source: danbooru.csv" in text
|
|
assert "EscapeParens: true" in text
|
|
|
|
|
|
def test_autocomplete_merge_fills_empty_fds_source(tmp_path):
|
|
out, text = _run_autocomplete_merge(
|
|
tmp_path,
|
|
"DefaultUser:\n AutoComplete:\n Source: \\x\n EscapeParens: true\n",
|
|
)
|
|
assert "CHANGED patched AutoComplete.Source=danbooru.csv" in out
|
|
assert "Source: danbooru.csv" in text
|
|
|
|
|
|
def test_autocomplete_merge_keeps_user_source(tmp_path):
|
|
out, text = _run_autocomplete_merge(
|
|
tmp_path,
|
|
"DefaultUser:\n AutoComplete:\n Source: e621.csv\n",
|
|
)
|
|
assert "SKIP_USER" in out
|
|
assert "Source: e621.csv" in text
|
|
assert "danbooru.csv" not in text
|
|
|
|
|
|
def test_seed_autocomplete_merges_when_csv_already_present(monkeypatch):
|
|
import json
|
|
|
|
from gpu_rent import provision
|
|
|
|
class Cfg:
|
|
autocomplete_enabled = True
|
|
autocomplete_filename = "danbooru.csv"
|
|
autocomplete_github_repo = "org/repo"
|
|
autocomplete_github_path = "tags/danbooru.csv"
|
|
autocomplete_github_ref = "main"
|
|
|
|
monkeypatch.setattr(
|
|
provision,
|
|
"_github_blob",
|
|
lambda _cfg: {"sha": "abc", "download_url": "https://example.invalid/x"},
|
|
)
|
|
|
|
def fake_exists(_cfg, _host, path):
|
|
return path.endswith("danbooru.csv") or path.endswith(".json")
|
|
|
|
def fake_ssh(_cfg, _host, cmd, **_kw):
|
|
if "cat" in cmd:
|
|
return json.dumps({"github_blob_sha": "abc", "settings_applied": True})
|
|
return ""
|
|
|
|
merges: list[str] = []
|
|
|
|
def fake_merge(*_a, **_k):
|
|
merges.append("called")
|
|
return "changed"
|
|
|
|
monkeypatch.setattr(provision, "remote_exists", fake_exists)
|
|
monkeypatch.setattr(provision, "run_ssh", fake_ssh)
|
|
monkeypatch.setattr(provision, "_merge_autocomplete_into_settings", fake_merge)
|
|
monkeypatch.setattr(provision, "put_text", lambda *_a, **_k: None)
|
|
|
|
logs: list[str] = []
|
|
assert provision.seed_autocomplete(Cfg(), "1.2.3.4", logs.append) is True
|
|
assert merges == ["called"]
|