Implement balance monitoring and notification for Selectel API integration
- Added support for balance tracking using `SELECTEL_API_TOKEN` in the configuration. - Introduced new balance notification logic in the local watchdog, alerting users on balance changes based on defined thresholds. - Updated documentation to include instructions for setting up balance notifications and the required environment variables. - Enhanced the `ready` and `session` modules to initialize balance state and handle notifications during GPU operations. - Refactored the CLI and related components to support the new balance monitoring features, ensuring a seamless user experience.
This commit is contained in:
Leonid Pershin
+295
-1
@@ -1,12 +1,21 @@
|
||||
"""Wait until SwarmUI HTTP is up and backend is Idle (on the VM)."""
|
||||
"""Wait until SwarmUI HTTP is up and backend is Idle (on the VM).
|
||||
|
||||
Also end-of-up stack verification: every enabled service must answer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import socket
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
from gpu_rent.config import Config
|
||||
from gpu_rent.errors import CloudError
|
||||
from gpu_rent.llm_runtime import normalize_runtime
|
||||
from gpu_rent.ssh_ops import run_ssh
|
||||
|
||||
Log = Callable[[str], None]
|
||||
@@ -57,6 +66,93 @@ while time.time() < deadline:
|
||||
print("WAIT timeout-slice")
|
||||
"""
|
||||
|
||||
# One-shot probe of configured stack endpoints on the VM (JSON line).
|
||||
_REMOTE_STACK_PROBE = r'''
|
||||
import json, urllib.error, urllib.request, subprocess
|
||||
|
||||
def http_ok(url, timeout=4.0):
|
||||
try:
|
||||
req = urllib.request.Request(url, method="GET")
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
code = getattr(resp, "status", 200) or 200
|
||||
body = resp.read(256)
|
||||
return True, f"HTTP {code} ({len(body)}b)"
|
||||
except Exception as exc:
|
||||
return False, str(exc)[:160]
|
||||
|
||||
def unit_active(name):
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
["systemctl", "is-active", name],
|
||||
text=True,
|
||||
stderr=subprocess.DEVNULL,
|
||||
).strip()
|
||||
return out
|
||||
except Exception:
|
||||
return "unknown"
|
||||
|
||||
checks = []
|
||||
want_swarm = WANT_SWARM
|
||||
want_ollama = WANT_OLLAMA
|
||||
want_llama = WANT_LLAMA
|
||||
|
||||
if want_swarm:
|
||||
ok, detail = http_ok("http://127.0.0.1:7801/")
|
||||
if not ok:
|
||||
# API may answer when / does not
|
||||
ok2, d2 = False, detail
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
"http://127.0.0.1:7801/API/GetNewSession",
|
||||
data=b"{}",
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
ok2 = True
|
||||
d2 = f"API session HTTP {getattr(resp, 'status', 200)}"
|
||||
except Exception as exc:
|
||||
d2 = str(exc)[:160]
|
||||
ok, detail = ok2, d2
|
||||
checks.append({
|
||||
"name": "swarmui",
|
||||
"ok": ok,
|
||||
"detail": detail,
|
||||
"unit": unit_active("swarmui"),
|
||||
})
|
||||
|
||||
if want_ollama:
|
||||
ok, detail = http_ok("http://127.0.0.1:11434/api/tags")
|
||||
checks.append({
|
||||
"name": "ollama",
|
||||
"ok": ok,
|
||||
"detail": detail,
|
||||
"unit": unit_active("gpu-rent-ollama"),
|
||||
})
|
||||
|
||||
if want_llama:
|
||||
ok, detail = http_ok("http://127.0.0.1:8080/health")
|
||||
if not ok:
|
||||
ok2, d2 = http_ok("http://127.0.0.1:8080/v1/models")
|
||||
ok, detail = ok2, d2
|
||||
checks.append({
|
||||
"name": "llamacpp",
|
||||
"ok": ok,
|
||||
"detail": detail,
|
||||
"unit": unit_active("gpu-rent-llamacpp"),
|
||||
})
|
||||
|
||||
print(json.dumps({"checks": checks}, ensure_ascii=False))
|
||||
'''
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ServiceCheck:
|
||||
name: str
|
||||
ok: bool
|
||||
detail: str
|
||||
where: str = "vm" # vm | local
|
||||
|
||||
|
||||
def wait_backend_idle(
|
||||
cfg: Config,
|
||||
@@ -93,3 +189,201 @@ def wait_backend_idle(
|
||||
f"backend не стал Idle за {int(timeout)} с. "
|
||||
"Проверь journalctl -u swarmui на VM; GPU всё ещё жив."
|
||||
)
|
||||
|
||||
|
||||
def _expected_services(cfg: Config) -> tuple[bool, bool, bool]:
|
||||
swarm = bool(getattr(cfg, "enable_swarmui", True))
|
||||
rt = normalize_runtime(getattr(cfg, "llm_runtime", "none"))
|
||||
return swarm, rt == "ollama", rt == "llamacpp"
|
||||
|
||||
|
||||
def _probe_vm_once(cfg: Config, host: str) -> list[ServiceCheck]:
|
||||
want_swarm, want_ollama, want_llama = _expected_services(cfg)
|
||||
script = (
|
||||
_REMOTE_STACK_PROBE.replace("WANT_SWARM", "True" if want_swarm else "False")
|
||||
.replace("WANT_OLLAMA", "True" if want_ollama else "False")
|
||||
.replace("WANT_LLAMA", "True" if want_llama else "False")
|
||||
)
|
||||
out = run_ssh(
|
||||
cfg,
|
||||
host,
|
||||
"python3 - <<'PY'\n" + script + "\nPY",
|
||||
check=False,
|
||||
timeout=60,
|
||||
).strip()
|
||||
line = ""
|
||||
for row in reversed(out.splitlines()):
|
||||
row = row.strip()
|
||||
if row.startswith("{"):
|
||||
line = row
|
||||
break
|
||||
if not line:
|
||||
return [
|
||||
ServiceCheck("stack", False, f"нет JSON от probe: {out[-200:]}", "vm")
|
||||
]
|
||||
try:
|
||||
data = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
return [ServiceCheck("stack", False, f"битый JSON: {line[:200]}", "vm")]
|
||||
checks: list[ServiceCheck] = []
|
||||
for item in data.get("checks") or []:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = str(item.get("name") or "?")
|
||||
detail = str(item.get("detail") or "")
|
||||
unit = str(item.get("unit") or "")
|
||||
if unit and unit != "unknown":
|
||||
detail = f"{detail}; unit={unit}"
|
||||
checks.append(
|
||||
ServiceCheck(name=name, ok=bool(item.get("ok")), detail=detail, where="vm")
|
||||
)
|
||||
return checks
|
||||
|
||||
|
||||
def verify_stack_on_vm(
|
||||
cfg: Config,
|
||||
host: str,
|
||||
log: Log,
|
||||
*,
|
||||
timeout: float = 300.0,
|
||||
poll_every: float = 8.0,
|
||||
raise_on_fail: bool = True,
|
||||
) -> list[ServiceCheck]:
|
||||
"""Poll until every enabled service answers on the VM loopback."""
|
||||
want_swarm, want_ollama, want_llama = _expected_services(cfg)
|
||||
if not (want_swarm or want_ollama or want_llama):
|
||||
log("проверка стека: нечего ждать (swarm off, LLM none)")
|
||||
return []
|
||||
|
||||
names = []
|
||||
if want_swarm:
|
||||
names.append("SwarmUI :7801")
|
||||
if want_ollama:
|
||||
names.append("Ollama :11434")
|
||||
if want_llama:
|
||||
names.append("llama.cpp :8080")
|
||||
log(f"проверка на VM: {', '.join(names)}")
|
||||
|
||||
deadline = time.time() + timeout
|
||||
last: list[ServiceCheck] = []
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
last = _probe_vm_once(cfg, host)
|
||||
except Exception as exc:
|
||||
last = [ServiceCheck("ssh", False, str(exc)[:200], "vm")]
|
||||
if last and all(c.ok for c in last):
|
||||
for c in last:
|
||||
log(f" [ok] {c.name}: {c.detail}")
|
||||
log("проверка VM: всё отвечает")
|
||||
return last
|
||||
bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
|
||||
log(f" … ещё нет: {bad}")
|
||||
time.sleep(poll_every)
|
||||
|
||||
for c in last:
|
||||
mark = "ok" if c.ok else "FAIL"
|
||||
log(f" [{mark}] {c.name}: {c.detail}")
|
||||
if raise_on_fail and any(not c.ok for c in last):
|
||||
failed = [c.name for c in last if not c.ok]
|
||||
raise CloudError(
|
||||
f"сервисы не ответили на VM за {int(timeout)} с: {', '.join(failed)}. "
|
||||
"GPU жив — смотри journalctl / gpu-rent logs"
|
||||
)
|
||||
return last
|
||||
|
||||
|
||||
def _tcp_ok(port: int, host: str = "127.0.0.1", timeout: float = 0.8) -> bool:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(timeout)
|
||||
try:
|
||||
return sock.connect_ex((host, port)) == 0
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def _http_local(url: str, timeout: float = 3.0) -> tuple[bool, str]:
|
||||
try:
|
||||
req = urllib.request.Request(url, method="GET")
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
code = getattr(resp, "status", 200) or 200
|
||||
return True, f"HTTP {code}"
|
||||
except Exception as exc:
|
||||
return False, str(exc)[:160]
|
||||
|
||||
|
||||
def verify_stack_local(
|
||||
cfg: Config,
|
||||
log: Log,
|
||||
*,
|
||||
timeout: float = 60.0,
|
||||
poll_every: float = 2.0,
|
||||
raise_on_fail: bool = True,
|
||||
) -> list[ServiceCheck]:
|
||||
"""After tunnel: local ports + light HTTP for enabled services."""
|
||||
want_swarm, want_ollama, want_llama = _expected_services(cfg)
|
||||
targets: list[tuple[str, int, str | None]] = []
|
||||
if want_swarm:
|
||||
targets.append(("swarmui", int(cfg.swarmui_local_port), None))
|
||||
if want_ollama:
|
||||
targets.append(
|
||||
("ollama", int(cfg.ollama_local_port), f"http://127.0.0.1:{cfg.ollama_local_port}/api/tags")
|
||||
)
|
||||
if want_llama:
|
||||
p = int(cfg.llamacpp_local_port)
|
||||
targets.append(("llamacpp", p, f"http://127.0.0.1:{p}/health"))
|
||||
|
||||
if not targets:
|
||||
return []
|
||||
|
||||
log("проверка туннеля (localhost): " + ", ".join(f"{n}:{port}" for n, port, _ in targets))
|
||||
deadline = time.time() + timeout
|
||||
last: list[ServiceCheck] = []
|
||||
while time.time() < deadline:
|
||||
last = []
|
||||
for name, port, url in targets:
|
||||
if not _tcp_ok(port):
|
||||
last.append(ServiceCheck(name, False, f"порт {port} закрыт", "local"))
|
||||
continue
|
||||
if url:
|
||||
ok, detail = _http_local(url)
|
||||
if not ok and name == "llamacpp":
|
||||
ok, detail = _http_local(
|
||||
f"http://127.0.0.1:{port}/v1/models"
|
||||
)
|
||||
last.append(ServiceCheck(name, ok, detail, "local"))
|
||||
else:
|
||||
# SwarmUI: TCP enough (API may need POST); try API too
|
||||
ok, detail = _http_local(f"http://127.0.0.1:{port}/")
|
||||
if not ok:
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{port}/API/GetNewSession",
|
||||
data=b"{}",
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=4) as resp:
|
||||
ok = True
|
||||
detail = f"API HTTP {getattr(resp, 'status', 200)}"
|
||||
except Exception as exc:
|
||||
detail = f"TCP ok; HTTP {exc}"[:160]
|
||||
# TCP open is enough for swarm local check
|
||||
ok = True
|
||||
detail = f"TCP :{port} open"
|
||||
last.append(ServiceCheck(name, ok, detail, "local"))
|
||||
if last and all(c.ok for c in last):
|
||||
for c in last:
|
||||
log(f" [ok] localhost {c.name}: {c.detail}")
|
||||
log("проверка туннеля: всё доступно")
|
||||
return last
|
||||
time.sleep(poll_every)
|
||||
|
||||
for c in last:
|
||||
mark = "ok" if c.ok else "FAIL"
|
||||
log(f" [{mark}] localhost {c.name}: {c.detail}")
|
||||
if raise_on_fail and any(not c.ok for c in last):
|
||||
failed = [c.name for c in last if not c.ok]
|
||||
raise CloudError(
|
||||
f"туннель поднят, но локально не отвечает: {', '.join(failed)}"
|
||||
)
|
||||
return last
|
||||
|
||||
Reference in New Issue
Block a user