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.
2026-08-21 11:22:14 +03:00
parent d3d689052d
commit 241ada62cb
7 changed files with 300 additions and 9 deletions
@@ -460,6 +460,11 @@ def up(
        "--no-update",
        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(
        False,
        "--keep-on-fail",
@@ -627,7 +632,7 @@ def up(
            flavor=flavor,
            yes=yes,
            adopt=adopt,
-            update=False if no_update else None,
+            update=(False if no_update else True if force_update else None),
            confirm=confirm,
            ask=None if yes else ask,
            log=log,
@@ -138,6 +138,57 @@ WantedBy=timers.target
    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(
    cfg: Config,
    host: str,
@@ -147,6 +198,13 @@ def arm_idle_killer(
 ) -> None:
    if not 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)
    put_text(
        cfg,
@@ -38,6 +38,18 @@ DATA = "/mnt/swarm_data"
 _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]:
    import os

@@ -408,6 +420,41 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None:
        log("model-seed пропущен: манифест пуст — дефолт SwarmUI")
        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] = []
    for entry in entries:
        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):
        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))
    if cfg.civitai_api_token:
        put_text(
@@ -533,6 +605,12 @@ def seed_civitai(cfg: Config, host: str, log: Log) -> None:
        timeout=7200,
        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:
@@ -594,16 +672,39 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None:
        if not names:
            log("ollama-models.yaml пуст — pull skip")
        else:
-            put_text(cfg, host, "/tmp/gpu-rent-ollama-models.json", json.dumps(names, indent=2))
-            log(f"Ollama: pull {len(names)} из манифеста")
-            run_python(
+            # Fast path: all tags already present — skip upload/pull script.
+            listed = run_ssh(
                cfg,
                host,
-                _pkg_text("ollama_pull.py"),
-                remote_path="/tmp/gpu-rent-ollama_pull.py",
-                timeout=7200,
-                log=log,
+                "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(
+                    cfg,
+                    host,
+                    _pkg_text("ollama_pull.py"),
+                    remote_path="/tmp/gpu-rent-ollama_pull.py",
+                    timeout=7200,
+                    log=log,
+                )
    else:
        raise CloudError(f"неизвестный LLM_RUNTIME={runtime!r}")
    st = load_state()
@@ -317,6 +317,19 @@ def main() -> int:
        return 1
    MARKER.parent.mkdir(parents=True, exist_ok=True)
    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")
    return 0

@@ -143,7 +143,7 @@ while IFS= read -r line || [[ -n "$line" ]]; do
  ENV_LINES+="Environment=${line}"$'\n'
 done < "$OLLAMA_ENV_FILE"

-cat >/etc/systemd/system/${UNIT}.service <<EOF
+cat >/tmp/${UNIT}.service.new <<EOF
 [Unit]
 Description=gpu-rent Ollama (loopback, GPU-tuned)
 After=network-online.target local-fs.target
@@ -164,6 +164,15 @@ RestartSec=5
 WantedBy=multi-user.target
 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 enable "$UNIT"
 systemctl restart "$UNIT"
@@ -315,6 +315,15 @@ def cmd_up(
            state.server_name = getattr(existing, "name", None)
            if status == "ACTIVE":
                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 не создаём")
                    state.phase = "bootstrapping"
                    save_state(state)
@@ -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