Leonid Pershin
241ada62cb
- Added a `--update` option to the `up` command in the CLI for forced updates during warm ACTIVE states. - Improved the idle-killer logic with new functions to check for reusable credentials and refresh settings on existing remote credentials. - Enhanced the `seed_civitai` function to implement caching for model jobs, optimizing the provisioning process. - Updated the installation script for Ollama to prevent unnecessary restarts when the service configuration has not changed. - Added diagnostics for model job caching and improved error handling in various functions.
97 lines
2.9 KiB
Python
97 lines
2.9 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
|