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:
Leonid Pershin
2026-08-21 11:22:14 +03:00
parent d3d689052d
commit 241ada62cb
7 changed files with 300 additions and 9 deletions
+6 -1
View File
@@ -460,6 +460,11 @@ def up(
"--no-update", "--no-update",
help="Не делать git pull SwarmUI и установленных extensions", help="Не делать git pull SwarmUI и установленных extensions",
), ),
force_update: bool = typer.Option(
False,
"--update",
help="Принудительно git pull даже на warm ACTIVE (по умолчанию warm без pull)",
),
keep_on_fail: bool = typer.Option( keep_on_fail: bool = typer.Option(
False, False,
"--keep-on-fail", "--keep-on-fail",
@@ -627,7 +632,7 @@ def up(
flavor=flavor, flavor=flavor,
yes=yes, yes=yes,
adopt=adopt, adopt=adopt,
update=False if no_update else None, update=(False if no_update else True if force_update else None),
confirm=confirm, confirm=confirm,
ask=None if yes else ask, ask=None if yes else ask,
log=log, log=log,
+58
View File
@@ -138,6 +138,57 @@ WantedBy=timers.target
log(f"systemd timer {UNIT}.timer") log(f"systemd timer {UNIT}.timer")
def _remote_killer_reusable(cfg: Config, host: str, server_id: str) -> bool:
"""True if VM already has matching idle-killer creds + active timer."""
cmd = (
f"sid=$(cat {DATA}/.gpu-rent-server-id 2>/dev/null | tr -d '[:space:]'); "
f"echo sid=$sid; echo want={server_id}; "
f"test -f {DATA}/.gpu-rent-killer-creds-ok && echo creds=yes || echo creds=no; "
f"test -f {DATA}/.gpu-rent-killer-armed && echo armed=yes || echo armed=no; "
f"sudo -n test -f {CREDS_REMOTE} && echo rootcreds=yes || echo rootcreds=no; "
f"echo timer=$(systemctl is-active {UNIT}.timer 2>/dev/null || echo inactive)"
)
try:
out = run_ssh(cfg, host, cmd, check=False, timeout=25)
except Exception:
return False
parsed: dict[str, str] = {}
for line in out.splitlines():
if "=" in line:
k, v = line.strip().split("=", 1)
parsed[k] = v
return (
parsed.get("sid") == server_id
and parsed.get("want") == server_id
and parsed.get("creds") == "yes"
and parsed.get("armed") == "yes"
and parsed.get("rootcreds") == "yes"
and parsed.get("timer") == "active"
)
def _refresh_killer_grace(cfg: Config, host: str) -> None:
"""Bump grace_from / idle settings on existing remote creds; clear idle-since."""
run_ssh(
cfg,
host,
"sudo -n python3 -c \""
"import json,time; from pathlib import Path; "
f"p=Path('{CREDS_REMOTE}'); d=json.loads(p.read_text()); "
f"d['idle_minutes']={int(cfg.idle_minutes)}; "
f"d['grace_minutes']={int(cfg.idle_grace_minutes)}; "
"d['grace_from']=int(time.time()); "
"p.write_text(json.dumps(d,indent=2)+chr(10)); p.chmod(0o600)"
"\" && "
f"sudo -n rm -f {DATA}/.gpu-rent-idle-since && "
f"date -u +%Y-%m-%dT%H:%M:%SZ | sudo -n tee {DATA}/.gpu-rent-killer-armed >/dev/null && "
f"sudo -n chmod 644 {DATA}/.gpu-rent-killer-armed "
f"{DATA}/.gpu-rent-killer-creds-ok {DATA}/.gpu-rent-server-id 2>/dev/null || true",
check=False,
timeout=30,
)
def arm_idle_killer( def arm_idle_killer(
cfg: Config, cfg: Config,
host: str, host: str,
@@ -147,6 +198,13 @@ def arm_idle_killer(
) -> None: ) -> None:
if not server_id: if not server_id:
raise CloudError("idle-killer: нет server_id") raise CloudError("idle-killer: нет server_id")
if _remote_killer_reusable(cfg, host, server_id):
_refresh_killer_grace(cfg, host)
log(
f"idle-killer: reuse app cred (server {server_id[:12]}…) — "
f"льгота {cfg.idle_grace_minutes} мин, потом {cfg.idle_minutes} мин → delete"
)
return
creds = create_application_credential(conn, cfg, server_id, log) creds = create_application_credential(conn, cfg, server_id, log)
put_text( put_text(
cfg, cfg,
+103 -2
View File
@@ -38,6 +38,18 @@ DATA = "/mnt/swarm_data"
_OLLAMA_INSTALL_ENV = ("OLLAMA_VERSION", "OLLAMA_SHA256") _OLLAMA_INSTALL_ENV = ("OLLAMA_VERSION", "OLLAMA_SHA256")
def _models_manifest_fp(cfg: Config) -> str:
"""Stable fingerprint of models.yaml so warm seed can skip API+fetch."""
import hashlib
path = Path(cfg.models_manifest)
try:
raw = path.read_bytes()
except OSError:
raw = b""
return hashlib.sha256(raw).hexdigest()[:16]
def _remote_llm_env(cfg: Config, *keys: str) -> dict[str, str]: def _remote_llm_env(cfg: Config, *keys: str) -> dict[str, str]:
import os import os
@@ -408,6 +420,41 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None:
log("model-seed пропущен: манифест пуст — дефолт SwarmUI") log("model-seed пропущен: манифест пуст — дефолт SwarmUI")
return return
# Warm short-circuit: last successful jobs list still present on disk.
cache_remote = f"{DATA}/.gpu-rent-models-jobs.json"
try:
cached_raw = run_ssh(
cfg,
host,
f"test -f {cache_remote} && cat {cache_remote} || true",
check=False,
timeout=20,
).strip()
except Exception:
cached_raw = ""
if cached_raw.startswith("{"):
try:
cached = json.loads(cached_raw)
dests = [str(d) for d in (cached.get("dests") or []) if d]
fp = str(cached.get("fp") or "")
want_fp = _models_manifest_fp(cfg)
if dests and fp == want_fp:
check = " && ".join(f"test -f {shlex.quote(d)}" for d in dests)
ok = run_ssh(
cfg,
host,
f"if {check}; then echo ALL_OK; else echo MISSING; fi",
check=False,
timeout=30,
).strip()
if "ALL_OK" in ok:
log(
f"model-seed: skip — {len(dests)} файл(ов) уже на диске (cache)"
)
return
except (json.JSONDecodeError, TypeError, CloudError):
pass
jobs: list[dict] = [] jobs: list[dict] = []
for entry in entries: for entry in entries:
url = (entry.url or "").strip() url = (entry.url or "").strip()
@@ -514,6 +561,31 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None:
if not any(e.kind == "checkpoint" for e in entries): if not any(e.kind == "checkpoint" for e in entries):
log("в манифесте нет checkpoint — генерация может не стартовать") log("в манифесте нет checkpoint — генерация может не стартовать")
dests = [str(j.get("dest") or "") for j in jobs if j.get("dest")]
if dests:
check = " && ".join(f"test -f {shlex.quote(d)}" for d in dests)
present = run_ssh(
cfg,
host,
f"if {check}; then echo ALL_OK; else echo MISSING; fi",
check=False,
timeout=30,
).strip()
if "ALL_OK" in present:
put_text(
cfg,
host,
cache_remote,
json.dumps({"fp": _models_manifest_fp(cfg), "dests": dests}, indent=2)
+ "\n",
)
log(
f"model-seed: skip fetch — {len(dests)} файл(ов) уже на диске "
f"(civitai={civ_n}, huggingface={hf_n})"
)
return
put_text(cfg, host, "/tmp/gpu-rent-civitai-jobs.json", json.dumps(jobs, indent=2)) put_text(cfg, host, "/tmp/gpu-rent-civitai-jobs.json", json.dumps(jobs, indent=2))
if cfg.civitai_api_token: if cfg.civitai_api_token:
put_text( put_text(
@@ -533,6 +605,12 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None:
timeout=7200, timeout=7200,
log=log, log=log,
) )
put_text(
cfg,
host,
cache_remote,
json.dumps({"fp": _models_manifest_fp(cfg), "dests": dests}, indent=2) + "\n",
)
def ensure_swarmui_running(cfg: Config, host: str, log: Log, restart: bool) -> None: def ensure_swarmui_running(cfg: Config, host: str, log: Log, restart: bool) -> None:
@@ -594,8 +672,31 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None:
if not names: if not names:
log("ollama-models.yaml пуст — pull skip") log("ollama-models.yaml пуст — pull skip")
else: else:
put_text(cfg, host, "/tmp/gpu-rent-ollama-models.json", json.dumps(names, indent=2)) # Fast path: all tags already present — skip upload/pull script.
log(f"Ollama: pull {len(names)} из манифеста") listed = run_ssh(
cfg,
host,
"ollama list 2>/dev/null | awk 'NR>1 {print $1}' || true",
check=False,
timeout=30,
)
have = {ln.strip() for ln in listed.splitlines() if ln.strip()}
missing = []
for name in names:
if name in have or (
":" not in name and f"{name}:latest" in have
) or (
name.endswith(":latest") and name.rsplit(":", 1)[0] in have
):
continue
missing.append(name)
if not missing:
log(f"ollama pull: skip — все {len(names)} уже есть")
else:
put_text(
cfg, host, "/tmp/gpu-rent-ollama-models.json", json.dumps(missing, indent=2)
)
log(f"Ollama: pull {len(missing)} из манифеста (нет: {len(missing)})")
run_python( run_python(
cfg, cfg,
host, host,
+13
View File
@@ -317,6 +317,19 @@ def main() -> int:
return 1 return 1
MARKER.parent.mkdir(parents=True, exist_ok=True) MARKER.parent.mkdir(parents=True, exist_ok=True)
MARKER.write_text("ok\n", encoding="utf-8") MARKER.write_text("ok\n", encoding="utf-8")
# Dest list for warm seed short-circuit (local provision also writes fp).
try:
dests = []
for job in json.loads(JOBS_PATH.read_text(encoding="utf-8")):
if isinstance(job, dict) and job.get("dest"):
dests.append(str(job["dest"]))
if dests:
Path("/mnt/swarm_data/.gpu-rent-models-jobs.json").write_text(
json.dumps({"fp": "", "dests": dests}, indent=2) + "\n",
encoding="utf-8",
)
except Exception:
pass
print("civitai seed ok") print("civitai seed ok")
return 0 return 0
+10 -1
View File
@@ -143,7 +143,7 @@ while IFS= read -r line || [[ -n "$line" ]]; do
ENV_LINES+="Environment=${line}"$'\n' ENV_LINES+="Environment=${line}"$'\n'
done < "$OLLAMA_ENV_FILE" done < "$OLLAMA_ENV_FILE"
cat >/etc/systemd/system/${UNIT}.service <<EOF cat >/tmp/${UNIT}.service.new <<EOF
[Unit] [Unit]
Description=gpu-rent Ollama (loopback, GPU-tuned) Description=gpu-rent Ollama (loopback, GPU-tuned)
After=network-online.target local-fs.target After=network-online.target local-fs.target
@@ -164,6 +164,15 @@ RestartSec=5
WantedBy=multi-user.target WantedBy=multi-user.target
EOF EOF
UNIT_PATH=/etc/systemd/system/${UNIT}.service
if [[ -f "$UNIT_PATH" ]] && cmp -s "/tmp/${UNIT}.service.new" "$UNIT_PATH" \
&& systemctl is-active --quiet "$UNIT"; then
rm -f "/tmp/${UNIT}.service.new"
log "ok — unit без изменений и active, skip restart (OLLAMA_HOST=127.0.0.1:11434)"
exit 0
fi
mv "/tmp/${UNIT}.service.new" "$UNIT_PATH"
systemctl daemon-reload systemctl daemon-reload
systemctl enable "$UNIT" systemctl enable "$UNIT"
systemctl restart "$UNIT" systemctl restart "$UNIT"
+9
View File
@@ -315,6 +315,15 @@ def cmd_up(
state.server_name = getattr(existing, "name", None) state.server_name = getattr(existing, "name", None)
if status == "ACTIVE": if status == "ACTIVE":
if state.bootstrapped and state.floating_ip: if state.bootstrapped and state.floating_ip:
# Warm re-up: skip git pull/stop cascade unless --update.
if update is None:
do_update = False
log(
"warm ACTIVE — без git pull SwarmUI/extensions "
"(нужен свежий: --update)"
)
else:
do_update = update
log("сервер уже ACTIVE — второй GPU не создаём") log("сервер уже ACTIVE — второй GPU не создаём")
state.phase = "bootstrapping" state.phase = "bootstrapping"
save_state(state) save_state(state)
+96
View File
@@ -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