Leonid Pershin
09b7c36f3b
- 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.
568 lines
16 KiB
Python
568 lines
16 KiB
Python
"""Local optional watchdog: stale tunnel lease → stop GPU.
|
|
|
|
VM idle-killer remains the default safety net. This module is opt-in via
|
|
`gpu-rent watchdog install`: while a tunnel is armed, a scheduled tick stops
|
|
compute if the laptop/process died without Ctrl+C detach or `gpu-rent stop`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
from gpu_rent.paths import (
|
|
app_root,
|
|
local_lease_path,
|
|
local_watchdog_marker_path,
|
|
runtime_dir,
|
|
)
|
|
|
|
Log = Callable[[str], None]
|
|
|
|
|
|
@dataclass
|
|
class LocalLease:
|
|
version: int = 1
|
|
armed: bool = False
|
|
detached: bool = False
|
|
pid: int | None = None
|
|
heartbeat_at: str | None = None
|
|
app_root: str | None = None
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> LocalLease:
|
|
known = {k: v for k, v in data.items() if k in cls.__dataclass_fields__}
|
|
return cls(**known)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TickDecision:
|
|
kind: str # noop | stop
|
|
detail: str
|
|
|
|
|
|
def utc_now() -> datetime:
|
|
return datetime.now(timezone.utc).replace(microsecond=0)
|
|
|
|
|
|
def utc_now_iso() -> str:
|
|
return utc_now().isoformat()
|
|
|
|
|
|
def load_lease() -> LocalLease | None:
|
|
path = local_lease_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):
|
|
return None
|
|
return LocalLease.from_dict(raw)
|
|
|
|
|
|
def save_lease(lease: LocalLease) -> None:
|
|
runtime_dir().mkdir(parents=True, exist_ok=True)
|
|
local_lease_path().write_text(
|
|
json.dumps(lease.to_dict(), indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
|
|
|
|
def clear_lease() -> None:
|
|
path = local_lease_path()
|
|
if path.is_file():
|
|
path.unlink(missing_ok=True)
|
|
|
|
|
|
def load_marker() -> dict[str, Any] | None:
|
|
path = local_watchdog_marker_path()
|
|
if not path.is_file():
|
|
return None
|
|
try:
|
|
raw = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
return raw if isinstance(raw, dict) else None
|
|
|
|
|
|
def watchdog_installed() -> bool:
|
|
return load_marker() is not None
|
|
|
|
|
|
def grace_seconds() -> int:
|
|
raw = (os.environ.get("LOCAL_WATCHDOG_GRACE_MINUTES") or "").strip()
|
|
try:
|
|
minutes = int(raw) if raw else 10
|
|
except ValueError:
|
|
minutes = 10
|
|
return max(1, minutes) * 60
|
|
|
|
|
|
def pid_alive(pid: int | None) -> bool:
|
|
if pid is None or pid <= 0:
|
|
return False
|
|
try:
|
|
os.kill(pid, 0)
|
|
except ProcessLookupError:
|
|
return False
|
|
except PermissionError:
|
|
return True
|
|
except OSError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _parse_iso(stamp: str | None) -> datetime | None:
|
|
if not stamp:
|
|
return None
|
|
try:
|
|
dt = datetime.fromisoformat(stamp)
|
|
except ValueError:
|
|
return None
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt
|
|
|
|
|
|
def decide_local_tick(
|
|
*,
|
|
installed: bool,
|
|
has_server: bool,
|
|
lease: LocalLease | None,
|
|
now: datetime,
|
|
process_alive: bool,
|
|
grace_sec: int,
|
|
) -> TickDecision:
|
|
if not installed:
|
|
return TickDecision("noop", "watchdog not installed")
|
|
if not has_server:
|
|
return TickDecision("noop", "no compute in state")
|
|
if lease is None:
|
|
return TickDecision("noop", "no local lease (tunnel never armed)")
|
|
if lease.detached:
|
|
return TickDecision("noop", "detached after Ctrl+C")
|
|
if not lease.armed:
|
|
return TickDecision("noop", "lease not armed")
|
|
if process_alive:
|
|
return TickDecision("noop", "lease pid alive")
|
|
hb = _parse_iso(lease.heartbeat_at)
|
|
if hb is None:
|
|
return TickDecision("stop", "armed lease without heartbeat")
|
|
age = (now - hb).total_seconds()
|
|
if age < grace_sec:
|
|
return TickDecision("noop", f"grace {int(age)}s/{grace_sec}s")
|
|
return TickDecision("stop", f"stale heartbeat {int(age)}s, pid dead")
|
|
|
|
|
|
def arm_lease_for_tunnel() -> LocalLease:
|
|
lease = LocalLease(
|
|
armed=True,
|
|
detached=False,
|
|
pid=os.getpid(),
|
|
heartbeat_at=utc_now_iso(),
|
|
app_root=str(app_root()),
|
|
)
|
|
save_lease(lease)
|
|
return lease
|
|
|
|
|
|
def touch_heartbeat() -> None:
|
|
lease = load_lease()
|
|
if lease is None or not lease.armed or lease.detached:
|
|
return
|
|
lease.heartbeat_at = utc_now_iso()
|
|
lease.pid = os.getpid()
|
|
save_lease(lease)
|
|
|
|
|
|
def detach_lease_keep_gpu() -> None:
|
|
"""Ctrl+C on tunnel: leave GPU running; local tick must not stop."""
|
|
lease = load_lease()
|
|
if lease is None:
|
|
lease = LocalLease()
|
|
lease.armed = False
|
|
lease.detached = True
|
|
lease.heartbeat_at = utc_now_iso()
|
|
lease.pid = None
|
|
lease.app_root = str(app_root())
|
|
save_lease(lease)
|
|
|
|
|
|
_heartbeat_stop: threading.Event | None = None
|
|
_heartbeat_thread: threading.Thread | None = None
|
|
|
|
|
|
def start_heartbeat_thread(*, interval_seconds: float = 30.0) -> None:
|
|
global _heartbeat_stop, _heartbeat_thread
|
|
stop_heartbeat_thread()
|
|
if not watchdog_installed():
|
|
return
|
|
arm_lease_for_tunnel()
|
|
stop = threading.Event()
|
|
_heartbeat_stop = stop
|
|
|
|
def _loop() -> None:
|
|
while not stop.wait(interval_seconds):
|
|
try:
|
|
touch_heartbeat()
|
|
except OSError:
|
|
pass
|
|
|
|
thread = threading.Thread(target=_loop, name="gpu-rent-lease-hb", daemon=True)
|
|
_heartbeat_thread = thread
|
|
thread.start()
|
|
|
|
|
|
def stop_heartbeat_thread() -> None:
|
|
global _heartbeat_stop, _heartbeat_thread
|
|
if _heartbeat_stop is not None:
|
|
_heartbeat_stop.set()
|
|
_heartbeat_stop = None
|
|
_heartbeat_thread = None
|
|
|
|
|
|
def _task_name(root: Path) -> str:
|
|
digest = hashlib.sha1(str(root.resolve()).encode("utf-8")).hexdigest()[:10]
|
|
return f"gpu-rent-local-watchdog-{digest}"
|
|
|
|
|
|
def install_watchdog(
|
|
*,
|
|
interval_minutes: int = 5,
|
|
log: Log = print,
|
|
) -> dict[str, Any]:
|
|
root = app_root().resolve()
|
|
runtime_dir().mkdir(parents=True, exist_ok=True)
|
|
interval = max(1, interval_minutes)
|
|
name = _task_name(root)
|
|
platform = sys.platform
|
|
if platform == "win32":
|
|
_install_windows(name, root, interval, log)
|
|
elif platform == "darwin":
|
|
_install_macos(name, root, interval, log)
|
|
else:
|
|
_install_linux(name, root, interval, log)
|
|
|
|
marker = {
|
|
"version": 1,
|
|
"installed_at": utc_now_iso(),
|
|
"platform": platform,
|
|
"task_name": name,
|
|
"app_root": str(root),
|
|
"python": sys.executable,
|
|
"interval_minutes": interval,
|
|
}
|
|
local_watchdog_marker_path().write_text(
|
|
json.dumps(marker, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
log(
|
|
f"local-watchdog установлен ({platform}): тик каждые {interval} мин. "
|
|
f"Grace {grace_seconds() // 60} мин после смерти процесса туннеля → stop. "
|
|
f"Ctrl+C на туннеле GPU не гасит. "
|
|
"При SELECTEL_API_TOKEN — toast каждые BALANCE_NOTIFY_STEP_RUB ₽ (дефолт 200)."
|
|
)
|
|
return marker
|
|
|
|
|
|
def uninstall_watchdog(*, log: Log = print) -> None:
|
|
marker = load_marker()
|
|
root = app_root().resolve()
|
|
name = (marker or {}).get("task_name") or _task_name(root)
|
|
platform = sys.platform
|
|
if platform == "win32":
|
|
_uninstall_windows(str(name), log)
|
|
elif platform == "darwin":
|
|
_uninstall_macos(str(name), log)
|
|
else:
|
|
_uninstall_linux(str(name), log)
|
|
path = local_watchdog_marker_path()
|
|
if path.is_file():
|
|
path.unlink(missing_ok=True)
|
|
clear_lease()
|
|
log("local-watchdog снят")
|
|
|
|
|
|
def watchdog_status_lines() -> list[str]:
|
|
marker = load_marker()
|
|
lease = load_lease()
|
|
lines: list[str] = []
|
|
if marker is None:
|
|
lines.append("не установлен (gpu-rent watchdog install)")
|
|
else:
|
|
lines.append(
|
|
f"установлен {marker.get('platform')} task={marker.get('task_name')} "
|
|
f"каждые {marker.get('interval_minutes')}м"
|
|
)
|
|
if lease is None:
|
|
lines.append("lease: нет")
|
|
else:
|
|
alive = pid_alive(lease.pid)
|
|
lines.append(
|
|
f"lease: armed={lease.armed} detached={lease.detached} "
|
|
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
|
|
|
|
state = load_state()
|
|
lease = load_lease()
|
|
decision = decide_local_tick(
|
|
installed=watchdog_installed(),
|
|
has_server=bool(state.server_id),
|
|
lease=lease,
|
|
now=utc_now(),
|
|
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
|
|
log(f"watchdog tick: STOP — {decision.detail}")
|
|
if dry_run:
|
|
return decision
|
|
from gpu_rent.config import load_config
|
|
|
|
cfg = load_config(require_auth=True)
|
|
cmd_stop(cfg, no_pull=True, log=log)
|
|
clear_lease()
|
|
return decision
|
|
|
|
|
|
def _install_windows(name: str, root: Path, interval: int, log: Log) -> None:
|
|
tr = (
|
|
f'"{sys.executable}" -m gpu_rent watchdog tick '
|
|
f'--project "{root}"'
|
|
)
|
|
cmd = [
|
|
"schtasks",
|
|
"/Create",
|
|
"/TN",
|
|
name,
|
|
"/SC",
|
|
"MINUTE",
|
|
"/MO",
|
|
str(interval),
|
|
"/TR",
|
|
tr,
|
|
"/F",
|
|
"/RL",
|
|
"LIMITED",
|
|
]
|
|
proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
|
if proc.returncode != 0:
|
|
err = (proc.stderr or proc.stdout or "").strip()
|
|
raise RuntimeError(f"schtasks failed: {err or proc.returncode}")
|
|
log(f"Task Scheduler: {name}")
|
|
|
|
|
|
def _uninstall_windows(name: str, log: Log) -> None:
|
|
subprocess.run(
|
|
["schtasks", "/Delete", "/TN", name, "/F"],
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
log(f"Task Scheduler удалён: {name}")
|
|
|
|
|
|
def _linux_unit_paths(name: str) -> tuple[Path, Path]:
|
|
base = Path.home() / ".config" / "systemd" / "user"
|
|
return base / f"{name}.service", base / f"{name}.timer"
|
|
|
|
|
|
def _install_linux(name: str, root: Path, interval: int, log: Log) -> None:
|
|
service_path, timer_path = _linux_unit_paths(name)
|
|
service_path.parent.mkdir(parents=True, exist_ok=True)
|
|
py = sys.executable
|
|
service_path.write_text(
|
|
f"""[Unit]
|
|
Description=gpu-rent local watchdog tick ({root})
|
|
|
|
[Service]
|
|
Type=oneshot
|
|
WorkingDirectory={root}
|
|
Environment=GPU_RENT_ROOT={root}
|
|
ExecStart={py} -m gpu_rent watchdog tick --project {root}
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
timer_path.write_text(
|
|
f"""[Unit]
|
|
Description=gpu-rent local watchdog every {interval} min
|
|
|
|
[Timer]
|
|
OnBootSec=2min
|
|
OnUnitActiveSec={interval}min
|
|
AccuracySec=1min
|
|
Persistent=true
|
|
|
|
[Install]
|
|
WantedBy=timers.target
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
subprocess.run(["systemctl", "--user", "daemon-reload"], check=False)
|
|
subprocess.run(
|
|
["systemctl", "--user", "enable", "--now", f"{name}.timer"],
|
|
check=False,
|
|
)
|
|
log(f"systemd user timer: {name}.timer")
|
|
|
|
|
|
def _uninstall_linux(name: str, log: Log) -> None:
|
|
subprocess.run(
|
|
["systemctl", "--user", "disable", "--now", f"{name}.timer"],
|
|
check=False,
|
|
capture_output=True,
|
|
)
|
|
service_path, timer_path = _linux_unit_paths(name)
|
|
service_path.unlink(missing_ok=True)
|
|
timer_path.unlink(missing_ok=True)
|
|
subprocess.run(["systemctl", "--user", "daemon-reload"], check=False)
|
|
log(f"systemd timer снят: {name}")
|
|
|
|
|
|
def _macos_plist_path(name: str) -> Path:
|
|
return Path.home() / "Library" / "LaunchAgents" / f"{name}.plist"
|
|
|
|
|
|
def _install_macos(name: str, root: Path, interval: int, log: Log) -> None:
|
|
plist = _macos_plist_path(name)
|
|
plist.parent.mkdir(parents=True, exist_ok=True)
|
|
seconds = interval * 60
|
|
py = sys.executable
|
|
body = f"""<?xml version="1.0" encoding="UTF-8"?>
|
|
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
|
|
<plist version="1.0">
|
|
<dict>
|
|
<key>Label</key>
|
|
<string>{name}</string>
|
|
<key>ProgramArguments</key>
|
|
<array>
|
|
<string>{py}</string>
|
|
<string>-m</string>
|
|
<string>gpu_rent</string>
|
|
<string>watchdog</string>
|
|
<string>tick</string>
|
|
<string>--project</string>
|
|
<string>{root}</string>
|
|
</array>
|
|
<key>WorkingDirectory</key>
|
|
<string>{root}</string>
|
|
<key>EnvironmentVariables</key>
|
|
<dict>
|
|
<key>GPU_RENT_ROOT</key>
|
|
<string>{root}</string>
|
|
</dict>
|
|
<key>StartInterval</key>
|
|
<integer>{seconds}</integer>
|
|
<key>RunAtLoad</key>
|
|
<true/>
|
|
</dict>
|
|
</plist>
|
|
"""
|
|
plist.write_text(body, encoding="utf-8")
|
|
subprocess.run(["launchctl", "unload", str(plist)], check=False, capture_output=True)
|
|
proc = subprocess.run(
|
|
["launchctl", "load", str(plist)],
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
if proc.returncode != 0:
|
|
err = (proc.stderr or proc.stdout or "").strip()
|
|
raise RuntimeError(f"launchctl load failed: {err or proc.returncode}")
|
|
log(f"launchd: {plist}")
|
|
|
|
|
|
def _uninstall_macos(name: str, log: Log) -> None:
|
|
plist = _macos_plist_path(name)
|
|
subprocess.run(["launchctl", "unload", str(plist)], check=False, capture_output=True)
|
|
plist.unlink(missing_ok=True)
|
|
log(f"launchd снят: {name}")
|