Leonid Pershin
f8f8dcc93e
- Revised README, architecture, CLI, and other documentation to specify that both `Ctrl+C` and `Ctrl+D` now stop the GPU while preserving disk data, correcting previous inaccuracies. - Enhanced access card and setup instructions to reflect the updated command behavior for better user understanding. - Updated tests to ensure the new command behaviors are validated and documented correctly, improving overall clarity in GPU management.
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/Ctrl+D ``stop`` 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:
|
|
"""Mark the lease detached so the local tick will not stop compute."""
|
|
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 / Ctrl+D на туннеле — stop. "
|
|
"При 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}")
|