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
2026-08-21 06:52:51 +03:00
parent 7ed6a99df2
commit 09b7c36f3b
14 changed files with 759 additions and 12 deletions
+69 -1
View File
@@ -271,7 +271,8 @@ def install_watchdog(
log(
f"local-watchdog установлен ({platform}): тик каждые {interval} мин. "
f"Grace {grace_seconds() // 60} мин после смерти процесса туннеля → stop. "
f"Ctrl+C на туннеле GPU не гасит."
f"Ctrl+C на туннеле GPU не гасит. "
"При SELECTEL_API_TOKEN — toast каждые BALANCE_NOTIFY_STEP_RUB ₽ (дефолт 200)."
)
return marker
@@ -314,9 +315,64 @@ def watchdog_status_lines() -> list[str]:
f"pid={lease.pid} alive={alive} hb={lease.heartbeat_at}"
)
lines.append(f"grace: {grace_seconds() // 60} мин (LOCAL_WATCHDOG_GRACE_MINUTES)")
try:
from gpu_rent.balance import load_balance_state
bal = load_balance_state()
if bal:
lines.append(
f"balance: baseline {bal.baseline_rub:.0f}"
f"last={bal.last_balance_rub} step_notified={bal.last_notified_step}"
)
else:
lines.append("balance: нет baseline (появится на up при SELECTEL_API_TOKEN)")
except Exception:
pass
return lines
def check_balance_on_tick(cfg: Config, *, log: Log = print) -> list[str]:
"""Fetch balance; toast on each BALANCE_NOTIFY_STEP_RUB spent since baseline.
No-op without SELECTEL_API_TOKEN. Arms baseline on first successful fetch
if missing but a GPU session may be running.
"""
from gpu_rent.balance import (
BalanceWatchState,
fetch_balance_rub,
load_balance_state,
plan_balance_notices,
save_balance_state,
)
from gpu_rent.notify import notify_message
token = (cfg.selectel_api_token or "").strip()
if not token:
return []
step = float(cfg.balance_notify_step_rub or 200.0)
low = float(cfg.balance_notify_low_rub or 0.0)
try:
current = fetch_balance_rub(token)
except Exception as exc:
log(f"balance: {exc}")
return []
state = load_balance_state()
if state is None:
state = BalanceWatchState(baseline_rub=current, last_balance_rub=current)
save_balance_state(state)
log(f"balance: baseline {current:.0f} ₽ (шаг {step:.0f} ₽)")
return []
state, messages = plan_balance_notices(
state, current, step_rub=step, low_rub=low
)
save_balance_state(state)
for msg in messages:
notify_message("gpu-rent баланс", msg, log)
return messages
def run_tick(*, dry_run: bool = False, log: Log = print) -> TickDecision:
from gpu_rent.session import cmd_stop
from gpu_rent.state import load_state
@@ -331,6 +387,18 @@ def run_tick(*, dry_run: bool = False, log: Log = print) -> TickDecision:
process_alive=pid_alive(lease.pid if lease else None),
grace_sec=grace_seconds(),
)
# Balance steps run on every tick while watchdog is installed (same timer).
if watchdog_installed() and not dry_run:
try:
from gpu_rent.config import load_config
cfg = load_config(require_auth=False)
if state.server_id or (lease and lease.armed):
check_balance_on_tick(cfg, log=log)
except Exception as exc:
log(f"balance tick: {exc}")
if decision.kind == "noop":
log(f"watchdog tick: noop ({decision.detail})")
return decision