Enhance CLI options and idle-killer functionality
- 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.
This commit is contained in:
@@ -0,0 +1,96 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user