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.
This commit is contained in:
@@ -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
|
||||
- [ ] 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`
|
||||
- [ ] 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] 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] 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] 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] 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] 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.
|
||||
|
||||
@@ -146,14 +146,15 @@ def arm_idle_killer(
|
||||
log: Log,
|
||||
) -> None:
|
||||
if not server_id:
|
||||
log("idle-killer: нет server_id — пропуск")
|
||||
return
|
||||
try:
|
||||
raise CloudError("idle-killer: нет server_id")
|
||||
creds = create_application_credential(conn, cfg, server_id, log)
|
||||
except CloudError as exc:
|
||||
log(f"idle-killer слеп: {exc}")
|
||||
return
|
||||
put_text(cfg, host, "/tmp/gpu-rent-idle-killer.json", 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,
|
||||
|
||||
+44
-20
@@ -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,7 +485,11 @@ 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)
|
||||
try:
|
||||
if swarm:
|
||||
try:
|
||||
if seed_extensions(cfg, host, log, update=update):
|
||||
@@ -512,26 +551,11 @@ def provision_vm(
|
||||
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 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:
|
||||
log("SwarmUI слушает 127.0.0.1:7801 — gpu-rent tunnel")
|
||||
if rt == "ollama":
|
||||
|
||||
@@ -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 {}
|
||||
|
||||
@@ -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 = True
|
||||
pip_ok = pip_install_sage(pip)
|
||||
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"],
|
||||
"pip_ok": pip_ok or not plan["use_sage"],
|
||||
"extra_args": plan["extra_args"] if pip_ok else "",
|
||||
"pip_ok": pip_ok,
|
||||
"restart_needed": restarted_needed,
|
||||
}
|
||||
MARKER.write_text(json.dumps(marker, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user