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
@@ -0,0 +1,174 @@
|
||||
"""Selectel account balance watch: notify every N ₽ spent since up."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from gpu_rent.errors import CloudError
|
||||
from gpu_rent.paths import runtime_dir
|
||||
|
||||
Log = Callable[[str], None]
|
||||
|
||||
BALANCES_URL = "https://api.selectel.ru/v3/balances"
|
||||
DEFAULT_STEP_RUB = 200.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class BalanceWatchState:
|
||||
baseline_rub: float
|
||||
last_notified_step: int = 0
|
||||
low_notified: bool = False
|
||||
last_balance_rub: float | None = None
|
||||
updated_at: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> BalanceWatchState:
|
||||
return cls(
|
||||
baseline_rub=float(data.get("baseline_rub") or 0),
|
||||
last_notified_step=int(data.get("last_notified_step") or 0),
|
||||
low_notified=bool(data.get("low_notified")),
|
||||
last_balance_rub=(
|
||||
float(data["last_balance_rub"])
|
||||
if data.get("last_balance_rub") is not None
|
||||
else None
|
||||
),
|
||||
updated_at=data.get("updated_at"),
|
||||
)
|
||||
|
||||
|
||||
def balance_watch_path() -> Path:
|
||||
return runtime_dir() / "balance-watch.json"
|
||||
|
||||
|
||||
def utc_now_iso() -> str:
|
||||
return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
|
||||
|
||||
|
||||
def load_balance_state() -> BalanceWatchState | None:
|
||||
path = balance_watch_path()
|
||||
if not path.is_file():
|
||||
return None
|
||||
try:
|
||||
raw = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if not isinstance(raw, dict) or "baseline_rub" not in raw:
|
||||
return None
|
||||
return BalanceWatchState.from_dict(raw)
|
||||
|
||||
|
||||
def save_balance_state(state: BalanceWatchState) -> None:
|
||||
runtime_dir().mkdir(parents=True, exist_ok=True)
|
||||
state.updated_at = utc_now_iso()
|
||||
balance_watch_path().write_text(
|
||||
json.dumps(state.to_dict(), indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def clear_balance_state() -> None:
|
||||
path = balance_watch_path()
|
||||
if path.is_file():
|
||||
path.unlink()
|
||||
|
||||
|
||||
def balance_rub_from_payload(payload: dict[str, Any]) -> float:
|
||||
"""Selectel returns money as integer kopecks → rubles."""
|
||||
data = payload.get("data") if isinstance(payload.get("data"), dict) else payload
|
||||
billings = data.get("billings") if isinstance(data, dict) else None
|
||||
if not isinstance(billings, list) or not billings:
|
||||
raise CloudError("Selectel balances: нет data.billings")
|
||||
total_kop = 0
|
||||
for row in billings:
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
if row.get("balances_values_sum") is not None:
|
||||
total_kop += int(row["balances_values_sum"])
|
||||
elif row.get("final_sum") is not None:
|
||||
total_kop += int(row["final_sum"])
|
||||
return total_kop / 100.0
|
||||
|
||||
|
||||
def fetch_balance_rub(api_token: str, *, timeout: float = 20.0) -> float:
|
||||
token = (api_token or "").strip()
|
||||
if not token:
|
||||
raise CloudError("SELECTEL_API_TOKEN пуст — баланс недоступен")
|
||||
headers = {"X-Token": token, "Accept": "application/json"}
|
||||
try:
|
||||
with httpx.Client(timeout=timeout, follow_redirects=True) as client:
|
||||
resp = client.get(BALANCES_URL, headers=headers)
|
||||
except httpx.HTTPError as exc:
|
||||
raise CloudError(f"Selectel balances сеть: {exc}") from exc
|
||||
if resp.status_code in {401, 403}:
|
||||
raise CloudError(
|
||||
f"Selectel balances HTTP {resp.status_code} — нужен статический "
|
||||
"API-ключ аккаунта (X-Token), не OpenStack пароль"
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
raise CloudError(f"Selectel balances HTTP {resp.status_code}: {resp.text[:200]}")
|
||||
try:
|
||||
payload = resp.json()
|
||||
except json.JSONDecodeError as exc:
|
||||
raise CloudError("Selectel balances: не JSON") from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise CloudError("Selectel balances: неожиданный JSON")
|
||||
return balance_rub_from_payload(payload)
|
||||
|
||||
|
||||
def spend_steps(baseline_rub: float, current_rub: float, step_rub: float) -> int:
|
||||
"""How many full step intervals have been spent since baseline."""
|
||||
if step_rub <= 0:
|
||||
return 0
|
||||
spend = baseline_rub - current_rub
|
||||
if spend < 0:
|
||||
return 0
|
||||
return int(spend // step_rub)
|
||||
|
||||
|
||||
def plan_balance_notices(
|
||||
state: BalanceWatchState,
|
||||
current_rub: float,
|
||||
*,
|
||||
step_rub: float = DEFAULT_STEP_RUB,
|
||||
low_rub: float = 0.0,
|
||||
) -> tuple[BalanceWatchState, list[str]]:
|
||||
"""
|
||||
Update state for top-ups / spend steps / low watermark.
|
||||
Returns (new_state, human messages to show).
|
||||
"""
|
||||
messages: list[str] = []
|
||||
# Top-up: reset baseline so steps restart from new balance.
|
||||
if current_rub > state.baseline_rub + 0.01:
|
||||
state = BalanceWatchState(
|
||||
baseline_rub=current_rub,
|
||||
last_notified_step=0,
|
||||
low_notified=False,
|
||||
last_balance_rub=current_rub,
|
||||
)
|
||||
messages.append(f"Баланс пополнен: {current_rub:.0f} ₽ (новый baseline)")
|
||||
return state, messages
|
||||
|
||||
crossed = spend_steps(state.baseline_rub, current_rub, step_rub)
|
||||
while state.last_notified_step < crossed:
|
||||
state.last_notified_step += 1
|
||||
spent = state.last_notified_step * step_rub
|
||||
messages.append(
|
||||
f"Списано ~{spent:.0f} ₽ с момента up "
|
||||
f"(осталось {current_rub:.0f} ₽, шаг {step_rub:.0f} ₽)"
|
||||
)
|
||||
|
||||
if low_rub > 0 and current_rub < low_rub and not state.low_notified:
|
||||
state.low_notified = True
|
||||
messages.append(f"Баланс низкий: {current_rub:.0f} ₽ (< {low_rub:.0f} ₽)")
|
||||
|
||||
state.last_balance_rub = current_rub
|
||||
return state, messages
|
||||
Reference in New Issue
Block a user