Implement idle-killer enhancements and swarm management improvements

- Marked critical bugs as resolved in the review documentation, including changes to the `arm_idle_killer` function to raise errors on credential creation failures and ensure proper file permissions for JSON credentials.
- Introduced a new `_try_arm_idle_killer` function in `provision.py` to manage idle-killer state more effectively, ensuring it arms correctly during provisioning.
- Updated the `swarm_busy` function in `remote/idle_killer.py` to allow idle state after a specified duration of Swarm unavailability, preventing unnecessary billing.
- Enhanced performance tuning logic in `tune_swarm_perf.py` to ensure proper handling of pip installation success before applying extra arguments.
- Added tests to validate the new idle-killer behavior and swarm management logic, ensuring robustness in handling idle states and error conditions.

docs/reviews/2026-08-21-review-3.md

@@ -81,10 +81,10 @@ N/A for CLI critical pass (or no critical issues).
 
 Critical only:
 
-- [ ] 1. [Bug] `arm_idle_killer`: on cred create failure **raise** or return False; never set `notes.idle_killer=armed`; surface same ⚠ as `"failed"` — `src/gpu_rent/idle_killer.py` line 151, `src/gpu_rent/provision.py` line 516
+- [x] 1. [Bug] `arm_idle_killer`: on cred create failure **raise** or return False; never set `notes.idle_killer=armed`; surface same ⚠ as `"failed"` — `src/gpu_rent/idle_killer.py` line 151, `src/gpu_rent/provision.py` line 516
-- [ ] 2. [Bug] Idle-killer: if Swarm unreachable longer than N minutes (e.g. 2× idle or fixed 60m), treat as idle/allow delete (llm-only already bypasses) — `src/gpu_rent/remote/idle_killer.py` line 75
+- [x] 2. [Bug] Idle-killer: if Swarm unreachable longer than N minutes (e.g. 2× idle or fixed 60m), treat as idle/allow delete (llm-only already bypasses) — `src/gpu_rent/remote/idle_killer.py` line 75
-- [ ] 3. [Logic] On mid-`up` failure after server create, arm killer anyway or fail loudly and refuse to leave session without killer / document mandatory `stop` — `src/gpu_rent/provision.py` / `session.py`
+- [x] 3. [Logic] On mid-`up` failure after server create, arm killer anyway or fail loudly and refuse to leave session without killer / document mandatory `stop` — `src/gpu_rent/provision.py` / `session.py`
-- [ ] 4. [Bug] Perf tune: set `pip_ok` only if pip succeeded; do **not** write `--use-sage-attention` ExtraArgs unless install OK (or allow retry when `pip_ok` false) — `src/gpu_rent/remote/tune_swarm_perf.py` line 109
+- [x] 4. [Bug] Perf tune: set `pip_ok` only if pip succeeded; do **not** write `--use-sage-attention` ExtraArgs unless install OK (or allow retry when `pip_ok` false) — `src/gpu_rent/remote/tune_swarm_perf.py` line 109
-- [ ] 5. [Security] Write idle-killer creds JSON with `mode=0o600` before `mv` (same as GIT_TOKEN) — `src/gpu_rent/idle_killer.py` line 156
+- [x] 5. [Security] Write idle-killer creds JSON with `mode=0o600` before `mv` (same as GIT_TOKEN) — `src/gpu_rent/idle_killer.py` line 156
 
 **Verified OK (critical):** `cmd_stop` delete path, tunnel Ctrl+C detach, local-watchdog stop on stale lease, `requires: ollama` filter, balance notify (toast only), Ollama unit bind to data dir, 127 tests green at review time.

src/gpu_rent/idle_killer.py

@@ -146,14 +146,15 @@ def arm_idle_killer(
     log: Log,
 ) -> None:
     if not server_id:
-        log("idle-killer: нет server_id — пропуск")
+        raise CloudError("idle-killer: нет server_id")
-        return
+    creds = create_application_credential(conn, cfg, server_id, log)
-    try:
+    put_text(
-        creds = create_application_credential(conn, cfg, server_id, log)
+        cfg,
-    except CloudError as exc:
+        host,
-        log(f"idle-killer слеп: {exc}")
+        "/tmp/gpu-rent-idle-killer.json",
-        return
+        json.dumps(creds, indent=2) + "\n",
-    put_text(cfg, host, "/tmp/gpu-rent-idle-killer.json", json.dumps(creds, indent=2) + "\n")
+        mode=0o600,
+    )
     run_ssh(
         cfg,
         host,

src/gpu_rent/provision.py

@@ -25,7 +25,6 @@ from gpu_rent.manifests import (
 )
 from gpu_rent.ssh_ops import put_text, remote_exists, run_python, run_ssh
 from gpu_rent.sync_files import pull_tree, push_tree
-from gpu_rent.idle_killer import arm_idle_killer
 
 Log = Callable[[str], None]
 DATA = "/mnt/swarm_data"
@@ -431,6 +430,42 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None:
     save_state(st)
 
 
+def _try_arm_idle_killer(
+    cfg: Config,
+    host: str,
+    log: Log,
+    *,
+    conn,
+    server_id: str | None,
+) -> bool:
+    """Arm idle-killer; update state notes. Returns True if armed."""
+    from gpu_rent.idle_killer import arm_idle_killer
+    from gpu_rent.state import load_state, save_state
+
+    if conn is None or not server_id:
+        return False
+    try:
+        arm_idle_killer(cfg, host, conn, server_id, log)
+        st = load_state()
+        st.notes = dict(st.notes or {})
+        st.notes["idle_killer"] = "armed"
+        st.notes.pop("idle_killer_error", None)
+        save_state(st)
+        return True
+    except GpuRentError as exc:
+        log(f"idle-killer: {exc}")
+        st = load_state()
+        st.notes = dict(st.notes or {})
+        st.notes["idle_killer"] = "failed"
+        st.notes["idle_killer_error"] = str(exc)[:500]
+        save_state(st)
+        log(
+            "⚠ idle-killer НЕ вооружён — GPU может тарифицироваться без авто-stop. "
+            "Сделай gpu-rent stop или почини identity/application_credential_create."
+        )
+        return False
+
+
 def provision_vm(
     cfg: Config,
     host: str,
@@ -450,88 +485,77 @@ def provision_vm(
             "llm-only: нужен LLM_RUNTIME=ollama|llamacpp (или --ollama / --llamacpp)"
         )
 
+    # Arm ASAP so mid-provision failures still leave auto-stop on the VM.
+    armed = _try_arm_idle_killer(cfg, host, log, conn=conn, server_id=server_id)
+
     restart = bool(update)
-    if swarm:
-        try:
-            if seed_extensions(cfg, host, log, update=update):
-                restart = True
-        except GpuRentError as exc:
-            log(f"extensions: {exc}")
-            raise
-        try:
-            if seed_autocomplete(cfg, host, log):
-                restart = True
-        except GpuRentError as exc:
-            log(f"autocomplete: {exc}")
-        seed_civitai(cfg, host, log)
-        push_tree(cfg, host, cfg.local_models_dir, f"{DATA}/Models", log, models=True)
-        push_tree(
-            cfg, host, cfg.local_wildcards_dir, f"{DATA}/Data/Wildcards", log, models=False
-        )
-        push_tree(
-            cfg,
-            host,
-            cfg.local_workflows_dir,
-            f"{DATA}/CustomWorkflows",
-            log,
-            models=False,
-        )
-        if cfg.pull_output:
-            pull_tree(cfg, host, f"{DATA}/Output", cfg.local_output_dir, log)
-    else:
-        log("SwarmUI: skip (llm-only)")
-        run_ssh(
-            cfg,
-            host,
-            "sudo -n systemctl stop swarmui 2>/dev/null; "
-            "sudo -n systemctl disable swarmui 2>/dev/null || true; "
-            "echo llm-only | sudo -n tee /mnt/swarm_data/.gpu-rent-llm-only >/dev/null",
-            check=False,
-        )
-
     try:
-        probe_gpu(cfg, host, log)
+        if swarm:
-    except Exception as exc:
+            try:
-        log(f"GPU probe: {exc}")
+                if seed_extensions(cfg, host, log, update=update):
-
+                    restart = True
-    if swarm:
+            except GpuRentError as exc:
-        ensure_swarmui_running(cfg, host, log, restart=restart)
+                log(f"extensions: {exc}")
-        run_ssh(
+                raise
-            cfg,
+            try:
-            host,
+                if seed_autocomplete(cfg, host, log):
-            "sudo -n rm -f /mnt/swarm_data/.gpu-rent-llm-only",
+                    restart = True
-            check=False,
+            except GpuRentError as exc:
-        )
+                log(f"autocomplete: {exc}")
-
+            seed_civitai(cfg, host, log)
-    try:
+            push_tree(cfg, host, cfg.local_models_dir, f"{DATA}/Models", log, models=True)
-        provision_llm(cfg, host, log)
+            push_tree(
-    except Exception as exc:
+                cfg, host, cfg.local_wildcards_dir, f"{DATA}/Data/Wildcards", log, models=False
-        st = load_state()
-        st.notes = dict(st.notes or {})
-        st.notes["llm_error"] = str(exc)[:500]
-        st.notes["llm_runtime"] = "none"
-        save_state(st)
-        raise CloudError(f"LLM runtime: {exc}") from exc
-
-    if conn is not None and server_id:
-        try:
-            arm_idle_killer(cfg, host, conn, server_id, log)
-            st = load_state()
-            st.notes = dict(st.notes or {})
-            st.notes["idle_killer"] = "armed"
-            st.notes.pop("idle_killer_error", None)
-            save_state(st)
-        except GpuRentError as exc:
-            log(f"idle-killer: {exc}")
-            st = load_state()
-            st.notes = dict(st.notes or {})
-            st.notes["idle_killer"] = "failed"
-            st.notes["idle_killer_error"] = str(exc)[:500]
-            save_state(st)
-            log(
-                "⚠ idle-killer НЕ вооружён — GPU может тарифицироваться без авто-stop. "
-                "См. status / docs/setup.md"
             )
+            push_tree(
+                cfg,
+                host,
+                cfg.local_workflows_dir,
+                f"{DATA}/CustomWorkflows",
+                log,
+                models=False,
+            )
+            if cfg.pull_output:
+                pull_tree(cfg, host, f"{DATA}/Output", cfg.local_output_dir, log)
+        else:
+            log("SwarmUI: skip (llm-only)")
+            run_ssh(
+                cfg,
+                host,
+                "sudo -n systemctl stop swarmui 2>/dev/null; "
+                "sudo -n systemctl disable swarmui 2>/dev/null || true; "
+                "echo llm-only | sudo -n tee /mnt/swarm_data/.gpu-rent-llm-only >/dev/null",
+                check=False,
+            )
+
+        try:
+            probe_gpu(cfg, host, log)
+        except Exception as exc:
+            log(f"GPU probe: {exc}")
+
+        if swarm:
+            ensure_swarmui_running(cfg, host, log, restart=restart)
+            run_ssh(
+                cfg,
+                host,
+                "sudo -n rm -f /mnt/swarm_data/.gpu-rent-llm-only",
+                check=False,
+            )
+
+        try:
+            provision_llm(cfg, host, log)
+        except Exception as exc:
+            st = load_state()
+            st.notes = dict(st.notes or {})
+            st.notes["llm_error"] = str(exc)[:500]
+            st.notes["llm_runtime"] = "none"
+            save_state(st)
+            raise CloudError(f"LLM runtime: {exc}") from exc
+    finally:
+        # If first arm failed (SSH race), retry once after seeds.
+        if not armed:
+            _try_arm_idle_killer(cfg, host, log, conn=conn, server_id=server_id)
+
     if swarm:
         log("SwarmUI слушает 127.0.0.1:7801 — gpu-rent tunnel")
     if rt == "ollama":

src/gpu_rent/remote/idle_killer.py

@@ -13,9 +13,12 @@ DATA = Path("/mnt/swarm_data")
 CREDS = Path("/root/.gpu-rent/idle-killer.json")
 HOLD = DATA / ".gpu-rent-hold-until"
 IDLE_SINCE = DATA / ".gpu-rent-idle-since"
+SWARM_DOWN_SINCE = DATA / ".gpu-rent-swarm-down-since"
 ARMED = DATA / ".gpu-rent-killer-armed"
 LOG = DATA / ".gpu-rent-killer.log"
 SERVER_ID_FILE = DATA / ".gpu-rent-server-id"
+# After this many seconds of continuous Swarm unreachable → treat as idle (don't bill forever).
+SWARM_UNREACHABLE_IDLE_SEC = 60 * 60
 
 
 def log(msg: str) -> None:
@@ -49,7 +52,11 @@ def write_ts(path: Path, value: float | None) -> None:
 
 
 def swarm_busy(swarm_url: str, timeout: float = 8.0) -> tuple[bool, str]:
-    """Return (busy, detail). Treat unreachable UI as busy (don't kill mid-boot)."""
+    """Return (busy, detail).
+
+    Unreachable UI is busy during boot, but after SWARM_UNREACHABLE_IDLE_SEC of
+    continuous failure we stop counting it as busy so idle clock can run / delete.
+    """
     ctx = ssl.create_default_context()
     try:
         req = urllib.request.Request(
@@ -73,7 +80,18 @@ def swarm_busy(swarm_url: str, timeout: float = 8.0) -> tuple[bool, str]:
         with urllib.request.urlopen(req2, timeout=timeout, context=ctx) as resp:
             data = json.loads(resp.read().decode("utf-8"))
     except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, json.JSONDecodeError, OSError) as exc:
-        return True, f"swarm unreachable: {exc}"
+        now = time.time()
+        since = read_ts(SWARM_DOWN_SINCE)
+        if since is None:
+            write_ts(SWARM_DOWN_SINCE, now)
+            return True, f"swarm unreachable (clock start): {exc}"
+        down_for = now - since
+        if down_for >= SWARM_UNREACHABLE_IDLE_SEC:
+            return False, f"swarm unreachable {int(down_for)}s >= {SWARM_UNREACHABLE_IDLE_SEC}s — allow idle"
+        return True, f"swarm unreachable {int(down_for)}s / {SWARM_UNREACHABLE_IDLE_SEC}s: {exc}"
+
+    # Reachable again — clear down clock.
+    write_ts(SWARM_DOWN_SINCE, None)
 
     status = data.get("status") or {}
     backend = data.get("backend_status") or {}

src/gpu_rent/remote/tune_swarm_perf.py

@@ -106,17 +106,18 @@ def patch_backends_extra_args(extra: str) -> bool:
     return True
 
 
-def pip_install_sage(pip: Path) -> None:
+def pip_install_sage(pip: Path) -> bool:
     print(f"pip install triton sageattention via {pip}")
     env = dict(os.environ)
     env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
-    # Best-effort: do not fail whole tune if wheels missing for this torch.
     cmd = [str(pip), "install", "-U", "triton", "sageattention"]
     try:
         subprocess.check_call(cmd, env=env)
         print("triton + sageattention installed")
+        return True
     except subprocess.CalledProcessError as exc:
-        print(f"WARN: pip install failed ({exc}) — ExtraArgs may no-op until fixed")
+        print(f"WARN: pip install failed ({exc}) — ExtraArgs NOT applied")
+        return False
 
 
 def main() -> int:
@@ -134,24 +135,25 @@ def main() -> int:
             return 0
 
     print(f"perf tune: {plan['name']} tier={plan['tier']} sage={plan['use_sage']}")
-    pip_ok = False
+    pip_ok = not plan["use_sage"]
     restarted_needed = False
     if plan["use_sage"]:
         pip = find_pip()
         if pip:
-            pip_install_sage(pip)
+            pip_ok = pip_install_sage(pip)
-            pip_ok = True
         else:
             print("Comfy venv pip not found yet — will retry next up")
-        if patch_backends_extra_args(plan["extra_args"]):
+            pip_ok = False
+        # Only patch ExtraArgs when wheels installed — otherwise Comfy may break.
+        if pip_ok and patch_backends_extra_args(plan["extra_args"]):
             restarted_needed = True
 
     marker = {
         "uuid": plan["uuid"],
         "name": plan["name"],
         "tier": plan["tier"],
-        "extra_args": plan["extra_args"],
+        "extra_args": plan["extra_args"] if pip_ok else "",
-        "pip_ok": pip_ok or not plan["use_sage"],
+        "pip_ok": pip_ok,
         "restart_needed": restarted_needed,
     }
     MARKER.write_text(json.dumps(marker, indent=2) + "\n", encoding="utf-8")

tests/test_hold_killer.py

@@ -109,3 +109,34 @@ def test_classify_busy_queue(monkeypatch):
     busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
     assert busy is True
     assert "waiting=2" in detail
+
+
+def test_swarm_unreachable_starts_busy_then_allows_idle(tmp_path, monkeypatch):
+    mod = _load_remote()
+    monkeypatch.setattr(mod, "DATA", tmp_path)
+    monkeypatch.setattr(mod, "SWARM_DOWN_SINCE", tmp_path / ".gpu-rent-swarm-down-since")
+    monkeypatch.setattr(mod, "SWARM_UNREACHABLE_IDLE_SEC", 100)
+
+    def boom(*_a, **_k):
+        raise mod.urllib.error.URLError("down")
+
+    monkeypatch.setattr(mod.urllib.request, "urlopen", boom)
+
+    busy1, d1 = mod.swarm_busy("http://127.0.0.1:7801")
+    assert busy1 is True
+    assert "clock start" in d1
+    assert mod.SWARM_DOWN_SINCE.is_file()
+
+    # Still within window
+    since = mod.read_ts(mod.SWARM_DOWN_SINCE)
+    assert since is not None
+    mod.write_ts(mod.SWARM_DOWN_SINCE, since - 50)
+    busy2, d2 = mod.swarm_busy("http://127.0.0.1:7801")
+    assert busy2 is True
+    assert "50s" in d2 or "/ 100s" in d2
+
+    # Past window → not busy so idle clock can run
+    mod.write_ts(mod.SWARM_DOWN_SINCE, since - 120)
+    busy3, d3 = mod.swarm_busy("http://127.0.0.1:7801")
+    assert busy3 is False
+    assert "allow idle" in d3

tests/test_tune_swarm_perf.py

@@ -0,0 +1,88 @@
+"""Tests for remote tune_swarm_perf pip_ok / ExtraArgs gating."""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+REMOTE = ROOT / "src" / "gpu_rent" / "remote" / "tune_swarm_perf.py"
+
+
+def _load():
+    spec = importlib.util.spec_from_file_location("tune_swarm_perf", REMOTE)
+    assert spec and spec.loader
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_pip_fail_skips_extra_args(tmp_path, monkeypatch):
+    mod = _load()
+    data = tmp_path
+    backends = data / "Data" / "Backends.fds"
+    backends.parent.mkdir(parents=True)
+    backends.write_text("ExtraArgs: \n", encoding="utf-8")
+    gpu_json = data / ".gpu-rent-gpu.json"
+    gpu_json.write_text(
+        json.dumps(
+            {
+                "vram_mib": 24576,
+                "compute_cap": "8.9",
+                "uuid": "gpu-1",
+                "name": "RTX",
+            }
+        ),
+        encoding="utf-8",
+    )
+    pip = data / "fake-pip"
+    pip.write_text("#!/bin/sh\n", encoding="utf-8")
+
+    monkeypatch.setattr(mod, "DATA", data)
+    monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
+    monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
+    monkeypatch.setattr(mod, "BACKENDS", backends)
+    monkeypatch.setattr(mod, "find_pip", lambda: pip)
+    monkeypatch.setattr(mod, "pip_install_sage", lambda _p: False)
+
+    assert mod.main() == 0
+    marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
+    assert marker["pip_ok"] is False
+    assert marker["extra_args"] == ""
+    assert "--use-sage-attention" not in backends.read_text(encoding="utf-8")
+
+
+def test_pip_ok_patches_extra_args(tmp_path, monkeypatch):
+    mod = _load()
+    data = tmp_path
+    backends = data / "Data" / "Backends.fds"
+    backends.parent.mkdir(parents=True)
+    backends.write_text("ExtraArgs: \n", encoding="utf-8")
+    gpu_json = data / ".gpu-rent-gpu.json"
+    gpu_json.write_text(
+        json.dumps(
+            {
+                "vram_mib": 24576,
+                "compute_cap": "8.9",
+                "uuid": "gpu-1",
+                "name": "RTX",
+            }
+        ),
+        encoding="utf-8",
+    )
+    pip = data / "fake-pip"
+    pip.write_text("#!/bin/sh\n", encoding="utf-8")
+
+    monkeypatch.setattr(mod, "DATA", data)
+    monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
+    monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
+    monkeypatch.setattr(mod, "BACKENDS", backends)
+    monkeypatch.setattr(mod, "find_pip", lambda: pip)
+    monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
+
+    assert mod.main() == 0
+    marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
+    assert marker["pip_ok"] is True
+    assert "--use-sage-attention" in marker["extra_args"]
+    assert "--use-sage-attention" in backends.read_text(encoding="utf-8")