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.
2026-08-21 07:04:09 +03:00
parent d409e2e154
commit 1ec615c03e
7 changed files with 268 additions and 104 deletions
@@ -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)
-    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
-
-    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"
+        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)
+        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":