Leonid Pershin
ccba40a228
- Updated the `_print_checks` function to replace console prints with logging functions for better traceability. - Introduced timing functionality in the `doctor`, `dry_run`, and `up` functions to log the duration of preflight checks. - Modified the `wait_ssh` function to accept a logging callback, improving SSH wait feedback. - Enhanced the `mark` method in `PhaseTimes` to log phase durations, aiding in performance analysis. - Updated various remote scripts to ensure error messages are printed to stderr for better error handling.
578 lines
21 KiB
Python
578 lines
21 KiB
Python
"""up / stop / destroy: GPU lifetime; SwarmUI bootstrap after SSH."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Callable
|
||
|
||
from gpu_rent.cloud import (
|
||
create_gpu_server,
|
||
delete_floating_ip,
|
||
delete_server,
|
||
ensure_boot_volume,
|
||
ensure_data_volume,
|
||
ensure_floating_ip,
|
||
ensure_keypair,
|
||
ensure_network,
|
||
ensure_security_group,
|
||
guess_operator_cidr,
|
||
pick_existing_server,
|
||
server_status,
|
||
unshelve,
|
||
wait_volume,
|
||
)
|
||
from gpu_rent.bootstrap import run_bootstrap
|
||
from gpu_rent.provision import provision_vm, seed_swarmui_api_keys, tune_swarm_perf
|
||
from gpu_rent.ready import verify_gpu_env, verify_stack_on_vm, wait_backend_idle
|
||
from gpu_rent.snapshot import ensure_boot_snapshot
|
||
from gpu_rent.notify import notify_ready
|
||
from gpu_rent.config import Config
|
||
from gpu_rent.errors import CloudError, GpuRentError
|
||
from gpu_rent.inventory import (
|
||
gpu_quota_from_compute,
|
||
pick_boot_image,
|
||
pick_volume_type,
|
||
resolve_flavor,
|
||
)
|
||
from gpu_rent.ux import print_up_preview, prompt_server_plan
|
||
from gpu_rent.pools import best_offer, format_pool_scan, scan_pools
|
||
from gpu_rent.lock import SessionLock
|
||
from gpu_rent.os_client import (
|
||
KEYPAIR_NAME,
|
||
SG_NAME,
|
||
compute_quotas,
|
||
connect,
|
||
iter_flavors,
|
||
iter_images,
|
||
iter_volume_types,
|
||
)
|
||
from gpu_rent.ssh_keys import ensure_ed25519
|
||
from gpu_rent.ssh_ops import probe_ssh, run_ssh, wait_ssh
|
||
from gpu_rent.state import SessionState, load_state, save_state, utc_now
|
||
from gpu_rent.timing import PhaseTimes
|
||
|
||
Log = Callable[[str], None]
|
||
|
||
|
||
def _log_default(msg: str) -> None:
|
||
if msg.startswith("\r"):
|
||
print("\r" + msg[1:], end="", flush=True)
|
||
return
|
||
print(msg)
|
||
|
||
|
||
def _require_gpu_quota(conn) -> None:
|
||
quota = compute_quotas(conn)
|
||
limit = gpu_quota_from_compute(quota)
|
||
if limit is not None and limit <= 0:
|
||
raise CloudError(
|
||
"квота GPU = 0. Напиши в поддержку Selectel (текст в docs/setup.md). "
|
||
"CLI не создаст сервер."
|
||
)
|
||
|
||
|
||
def _bind_access(
|
||
conn,
|
||
server,
|
||
state: SessionState,
|
||
cfg: Config,
|
||
log: Log,
|
||
*,
|
||
update: bool = True,
|
||
phases: PhaseTimes | None = None,
|
||
) -> SessionState:
|
||
clock = phases or PhaseTimes()
|
||
ip, fip_id = ensure_floating_ip(
|
||
conn, server, state.floating_ip_id, state.floating_ip, log
|
||
)
|
||
state.floating_ip = ip
|
||
if fip_id:
|
||
state.floating_ip_id = fip_id
|
||
save_state(state)
|
||
wait_ssh(cfg, ip, log=log)
|
||
clock.mark("SSH", log)
|
||
log(f"SSH {cfg.ssh_user}@{ip}")
|
||
state.phase = "bootstrapping"
|
||
save_state(state)
|
||
swarm = bool(getattr(cfg, "enable_swarmui", True))
|
||
if swarm and update:
|
||
active = run_ssh(cfg, ip, "systemctl is-active swarmui 2>/dev/null || true", check=False).strip()
|
||
if active == "active":
|
||
log("systemctl stop swarmui перед git update")
|
||
run_ssh(cfg, ip, "sudo -n systemctl stop swarmui", timeout=120, check=False)
|
||
|
||
# Detect existing markers so light/full choice is explicit (swarm ↔ llm-only).
|
||
probe = run_ssh(
|
||
cfg,
|
||
ip,
|
||
"echo swarm=$(test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no); "
|
||
"echo data=$(test -f /mnt/swarm_data/.gpu-rent-ready && echo yes || echo no); "
|
||
"echo llm=$(test -f /mnt/swarm_data/.gpu-rent-llm-only && echo yes || echo no)",
|
||
check=False,
|
||
)
|
||
flags: dict[str, str] = {}
|
||
for line in probe.splitlines():
|
||
if "=" in line:
|
||
k, v = line.strip().split("=", 1)
|
||
flags[k] = v
|
||
has_swarm = flags.get("swarm") == "yes"
|
||
has_data = flags.get("data") == "yes"
|
||
was_llm_only = flags.get("llm") == "yes"
|
||
|
||
if swarm:
|
||
light = has_swarm
|
||
if was_llm_only and not has_swarm:
|
||
log("bootstrap: был llm-only → полный проход (ставим SwarmUI)")
|
||
elif light:
|
||
why = (
|
||
"локальный bootstrapped был сброшен, маркер на VM есть"
|
||
if not state.bootstrapped
|
||
else "маркер /opt/swarmui/.gpu-rent-bootstrapped"
|
||
)
|
||
log(f"bootstrap: LIGHT (без apt) — {why}")
|
||
else:
|
||
log("bootstrap: FULL (apt + SwarmUI) — маркера bootstrapped нет")
|
||
else:
|
||
light = has_data
|
||
if has_swarm and not was_llm_only:
|
||
log("bootstrap: llm-only на диске со SwarmUI — LIGHT data, Swarm unit stop")
|
||
if light:
|
||
log("bootstrap: LIGHT llm-only (без apt) — есть .gpu-rent-ready")
|
||
else:
|
||
log("bootstrap: FULL llm-only — маркера data ready нет")
|
||
|
||
run_bootstrap(cfg, ip, log, update=update and swarm, light=light)
|
||
clock.mark("bootstrap", log)
|
||
provision_vm(
|
||
cfg,
|
||
ip,
|
||
log,
|
||
conn=conn,
|
||
server_id=getattr(server, "id", None) or state.server_id,
|
||
update=update,
|
||
)
|
||
clock.mark("provision", log)
|
||
if swarm:
|
||
try:
|
||
wait_backend_idle(cfg, ip, log)
|
||
except CloudError as exc:
|
||
log(f"ready: {exc}")
|
||
clock.mark("Idle", log)
|
||
try:
|
||
if tune_swarm_perf(cfg, ip, log):
|
||
log("systemctl restart swarmui (perf ExtraArgs)")
|
||
run_ssh(cfg, ip, "sudo -n systemctl restart swarmui", timeout=120)
|
||
try:
|
||
wait_backend_idle(cfg, ip, log)
|
||
except CloudError as exc:
|
||
log(f"ready after perf: {exc}")
|
||
else:
|
||
log("perf tune: restart не нужен")
|
||
except Exception as exc:
|
||
log(f"perf tune: {exc}")
|
||
clock.mark("perf", log)
|
||
try:
|
||
seed_swarmui_api_keys(cfg, ip, log)
|
||
except Exception as exc:
|
||
log(f"SwarmUI API keys: {exc}")
|
||
clock.mark("api-keys", log)
|
||
else:
|
||
log("ready: llm-only (без ожидания SwarmUI Idle)")
|
||
|
||
try:
|
||
checks = verify_stack_on_vm(cfg, ip, log, timeout=300.0)
|
||
state.notes = dict(state.notes or {})
|
||
state.notes["stack_vm"] = [
|
||
{"name": c.name, "ok": c.ok, "detail": c.detail} for c in checks
|
||
]
|
||
state.notes.pop("stack_vm_error", None)
|
||
except CloudError as exc:
|
||
state.notes = dict(state.notes or {})
|
||
state.notes["stack_vm_error"] = str(exc)[:500]
|
||
save_state(state)
|
||
raise
|
||
clock.mark("verify", log)
|
||
|
||
try:
|
||
gpu_checks = verify_gpu_env(cfg, ip, log, timeout=600.0)
|
||
state.notes = dict(state.notes or {})
|
||
state.notes["gpu_env"] = [
|
||
{"name": c.name, "ok": c.ok, "detail": c.detail} for c in gpu_checks
|
||
]
|
||
state.notes.pop("gpu_env_error", None)
|
||
except CloudError as exc:
|
||
state.notes = dict(state.notes or {})
|
||
state.notes["gpu_env_error"] = str(exc)[:500]
|
||
save_state(state)
|
||
raise
|
||
clock.mark("gpu-env", log)
|
||
|
||
try:
|
||
ensure_boot_snapshot(
|
||
conn,
|
||
boot_volume_id=state.boot_volume_id,
|
||
cfg=cfg,
|
||
log=log,
|
||
)
|
||
except CloudError as exc:
|
||
log(f"snapshot: {exc}")
|
||
notify_ready(cfg, log)
|
||
try:
|
||
from gpu_rent.balance import (
|
||
BalanceWatchState,
|
||
fetch_balance_rub,
|
||
save_balance_state,
|
||
)
|
||
|
||
if cfg.selectel_api_token:
|
||
rub = fetch_balance_rub(cfg.selectel_api_token)
|
||
save_balance_state(
|
||
BalanceWatchState(baseline_rub=rub, last_balance_rub=rub)
|
||
)
|
||
log(
|
||
f"balance: baseline {rub:.0f} ₽ "
|
||
f"(уведомление каждые {cfg.balance_notify_step_rub:.0f} ₽, "
|
||
"нужен local-watchdog)"
|
||
)
|
||
else:
|
||
log("balance: SELECTEL_API_TOKEN не задан — шаги по ₽ skip")
|
||
except Exception as exc:
|
||
log(f"balance baseline: {exc}")
|
||
state.bootstrapped = True
|
||
state.phase = "ready_cloud"
|
||
state.notes = dict(state.notes or {})
|
||
state.notes["enable_swarmui"] = swarm
|
||
state.notes["up_timing"] = clock.summary_line()
|
||
save_state(state)
|
||
for line in clock.summary_lines():
|
||
log(line)
|
||
return state
|
||
|
||
|
||
def adopt_server(cfg: Config, log: Log = _log_default, *, update: bool = True) -> SessionState:
|
||
conn = connect(cfg)
|
||
server = pick_existing_server(conn)
|
||
if not server:
|
||
raise CloudError("нечего adopt: нет сервера с тегом/именем gpu-rent")
|
||
state = load_state()
|
||
state.server_id = server.id
|
||
state.server_name = getattr(server, "name", None)
|
||
state.flavor_id = getattr(server, "flavor_id", None) or (
|
||
(server.flavor or {}).get("id") if isinstance(getattr(server, "flavor", None), dict) else None
|
||
)
|
||
state.phase = "ready_cloud"
|
||
state.availability_zone = cfg.gpu_rent_az
|
||
save_state(state)
|
||
log(f"подхватили {server.id} статус {server_status(server)}")
|
||
if server_status(server) == "ACTIVE":
|
||
state.phase = "bootstrapping"
|
||
save_state(state)
|
||
_bind_access(conn, server, state, cfg, log, update=update)
|
||
return state
|
||
|
||
|
||
def cmd_up(
|
||
cfg: Config,
|
||
*,
|
||
no_spot: bool = False,
|
||
flavor: str | None = None,
|
||
yes: bool = False,
|
||
adopt: bool = False,
|
||
update: bool | None = None,
|
||
confirm: Callable[[str], bool] | None = None,
|
||
ask: Callable[[str, str], str] | None = None,
|
||
log: Log = _log_default,
|
||
) -> SessionState:
|
||
do_update = cfg.update_git if update is None else update
|
||
with SessionLock():
|
||
if adopt:
|
||
return adopt_server(cfg, log=log, update=do_update)
|
||
conn = connect(cfg)
|
||
_require_gpu_quota(conn)
|
||
state = load_state()
|
||
existing = pick_existing_server(conn)
|
||
if existing:
|
||
status = server_status(existing)
|
||
state.server_id = existing.id
|
||
state.server_name = getattr(existing, "name", None)
|
||
if status == "ACTIVE":
|
||
if state.bootstrapped and state.floating_ip:
|
||
log("сервер уже ACTIVE — второй GPU не создаём")
|
||
state.phase = "bootstrapping"
|
||
save_state(state)
|
||
_bind_access(conn, existing, state, cfg, log, update=do_update)
|
||
return state
|
||
|
||
# Bootstrap не завершён: почти всегда VM без authorized_keys.
|
||
fip = state.floating_ip
|
||
if not fip:
|
||
try:
|
||
fip, fip_id = ensure_floating_ip(
|
||
conn, existing, state.floating_ip_id, state.floating_ip, log
|
||
)
|
||
state.floating_ip = fip
|
||
if fip_id:
|
||
state.floating_ip_id = fip_id
|
||
save_state(state)
|
||
except Exception as exc:
|
||
log(f"FIP: {exc}")
|
||
|
||
outcome = probe_ssh(cfg, fip, attempts=2) if fip else "down"
|
||
if outcome == "ok":
|
||
log("сервер ACTIVE, SSH ок — продолжаем bootstrap")
|
||
state.phase = "bootstrapping"
|
||
save_state(state)
|
||
_bind_access(conn, existing, state, cfg, log, update=do_update)
|
||
return state
|
||
|
||
log(
|
||
f"ACTIVE без bootstrap, SSH={outcome} — "
|
||
"удаляем compute (диски оставляем), create с user_data Base64"
|
||
)
|
||
delete_server(conn, existing, log)
|
||
state.server_id = None
|
||
state.server_name = None
|
||
state.bootstrapped = False
|
||
state.phase = "idle"
|
||
save_state(state)
|
||
for vid in (state.boot_volume_id, state.data_volume_id):
|
||
if not vid:
|
||
continue
|
||
try:
|
||
wait_volume(conn, conn.block_storage.get_volume(vid), "available")
|
||
except Exception as vol_exc:
|
||
log(f"wait volume {vid}: {vol_exc}")
|
||
existing = None
|
||
if existing is not None and status in {"EXPIRED", "SHELVED", "SHELVED_OFFLOADED"}:
|
||
existing = unshelve(conn, existing, log)
|
||
state.server_id = existing.id
|
||
state.phase = "bootstrapping"
|
||
state.unshelved_at = utc_now()
|
||
save_state(state)
|
||
_bind_access(conn, existing, state, cfg, log, update=do_update)
|
||
return state
|
||
if existing is not None:
|
||
raise CloudError(f"сервер gpu-rent в статусе {status} — разбери в панели")
|
||
|
||
if state.server_id:
|
||
log(f"в state был server {state.server_id}, в облаке нет — создаём заново")
|
||
state.server_id = None
|
||
state.server_name = None
|
||
state.floating_ip = None
|
||
state.floating_ip_id = None
|
||
state.bootstrapped = False
|
||
state.phase = "idle"
|
||
save_state(state)
|
||
|
||
flavors = list(iter_flavors(conn))
|
||
try:
|
||
picked = resolve_flavor(
|
||
flavors,
|
||
cfg.flavor_preference,
|
||
explicit=flavor,
|
||
default_id=cfg.default_flavor_id or None,
|
||
fallback=cfg.flavor_fallback,
|
||
)
|
||
except ValueError as exc:
|
||
raise CloudError(str(exc)) from exc
|
||
|
||
image = pick_boot_image(list(iter_images(conn)))
|
||
if not image:
|
||
raise CloudError("нет GPU-образа Ubuntu Driver 580/535 в Glance")
|
||
vtype = pick_volume_type(list(iter_volume_types(conn)), cfg.gpu_rent_az)
|
||
spot = cfg.default_spot and not no_spot
|
||
|
||
try:
|
||
offers = scan_pools(cfg)
|
||
for line in format_pool_scan(offers, cfg.flavor_preference):
|
||
log(line)
|
||
best = best_offer(offers)
|
||
if best and best.region != cfg.os_region_name:
|
||
log(
|
||
f"! сейчас OS_REGION_NAME={cfg.os_region_name}, "
|
||
f"а предпочтение лучше закрывается в {best.region} — "
|
||
f"поставь OS_REGION_NAME={best.region} и GPU_RENT_AZ={best.region}a "
|
||
f"в .env (диски ещё не созданы) и повтори up"
|
||
)
|
||
except Exception as exc:
|
||
log(f"скан пулов: {exc}")
|
||
|
||
print_up_preview(cfg, flavors, picked=picked, spot=spot, log=log)
|
||
|
||
if not yes and ask is not None and flavor is None:
|
||
try:
|
||
plan = prompt_server_plan(
|
||
cfg,
|
||
flavors,
|
||
picked=picked,
|
||
spot=spot,
|
||
ask=ask,
|
||
confirm=confirm,
|
||
)
|
||
except ValueError as exc:
|
||
raise GpuRentError(str(exc)) from exc
|
||
picked = plan.flavor
|
||
spot = plan.spot
|
||
if plan.data_gb != cfg.data_volume_size_gb:
|
||
from dataclasses import replace
|
||
|
||
cfg = replace(cfg, data_volume_size_gb=plan.data_gb)
|
||
log(
|
||
f"выбрано: {'preemptible ' if spot else ''}{picked.name}, "
|
||
f"data {cfg.data_volume_size_gb} GB"
|
||
)
|
||
|
||
prompt = (
|
||
f"Создать {'preemptible ' if spot else ''}GPU {picked.name} "
|
||
f"в {cfg.gpu_rent_az}, образ {getattr(image, 'name', image.id)}, "
|
||
f"data {cfg.data_volume_size_gb} GB.\n"
|
||
f"Диск тарифицируется всегда; idle-killer через {cfg.idle_minutes} мин "
|
||
f"простоя (льгота {cfg.idle_grace_minutes} мин). Продолжить?"
|
||
)
|
||
if not yes:
|
||
ok = confirm(prompt) if confirm else False
|
||
if not ok:
|
||
raise GpuRentError("отменено")
|
||
|
||
private_path, pub = ensure_ed25519(cfg.ssh_private_key_path)
|
||
del private_path
|
||
public_key = pub.read_text(encoding="utf-8")
|
||
ensure_keypair(conn, public_key, log)
|
||
|
||
net, _subnet = ensure_network(conn, log)
|
||
cidr = guess_operator_cidr()
|
||
if cidr == "0.0.0.0/0":
|
||
import os
|
||
|
||
if (os.environ.get("GPU_RENT_SSH_CIDR") or "").strip():
|
||
log("SG SSH: GPU_RENT_SSH_CIDR=0.0.0.0/0")
|
||
else:
|
||
log("не удалось узнать твой IP — SG откроет SSH с 0.0.0.0/0")
|
||
sg = ensure_security_group(conn, cidr, log)
|
||
|
||
boot = ensure_boot_volume(
|
||
conn,
|
||
az=cfg.gpu_rent_az,
|
||
volume_type=vtype,
|
||
image_id=image.id,
|
||
snapshot_name=cfg.boot_snapshot_name,
|
||
existing_id=state.boot_volume_id or cfg.boot_volume_id,
|
||
log=log,
|
||
)
|
||
data = ensure_data_volume(
|
||
conn,
|
||
az=cfg.gpu_rent_az,
|
||
volume_type=vtype,
|
||
size_gb=cfg.data_volume_size_gb,
|
||
existing_id=state.data_volume_id or cfg.data_volume_id,
|
||
log=log,
|
||
)
|
||
|
||
state.phase = "provisioning"
|
||
state.flavor_id = picked.id
|
||
state.flavor_name = picked.name
|
||
state.boot_volume_id = boot.id
|
||
state.data_volume_id = data.id
|
||
state.image_id = image.id
|
||
state.network_id = net.id
|
||
state.security_group_id = sg.id
|
||
state.availability_zone = cfg.gpu_rent_az
|
||
state.spot = spot
|
||
state.keypair_name = KEYPAIR_NAME
|
||
save_state(state)
|
||
|
||
server = create_gpu_server(
|
||
conn,
|
||
flavor_id=picked.id,
|
||
net_id=net.id,
|
||
sg_name=SG_NAME,
|
||
boot_volume_id=boot.id,
|
||
data_volume_id=data.id,
|
||
az=cfg.gpu_rent_az,
|
||
spot=spot,
|
||
public_key=public_key,
|
||
log=log,
|
||
)
|
||
state.server_id = server.id
|
||
state.server_name = getattr(server, "name", None)
|
||
state.created_at = utc_now()
|
||
state.unshelved_at = None
|
||
state.phase = "bootstrapping"
|
||
save_state(state)
|
||
_bind_access(conn, server, state, cfg, log, update=do_update)
|
||
return state
|
||
|
||
|
||
def cmd_stop(
|
||
cfg: Config,
|
||
*,
|
||
destroy_disks: bool = False,
|
||
no_pull: bool = False,
|
||
log: Log = _log_default,
|
||
) -> SessionState:
|
||
with SessionLock():
|
||
conn = connect(cfg)
|
||
state = load_state()
|
||
if cfg.pull_output and not no_pull and state.floating_ip:
|
||
try:
|
||
from gpu_rent.sync_files import pull_tree
|
||
|
||
pull_tree(cfg, state.floating_ip, "/mnt/swarm_data/Output", cfg.local_output_dir, log)
|
||
except Exception as exc:
|
||
log(f"pull Output не удался: {exc}")
|
||
server = None
|
||
if state.server_id:
|
||
try:
|
||
server = conn.compute.get_server(state.server_id)
|
||
except Exception:
|
||
server = None
|
||
if server is None:
|
||
server = pick_existing_server(conn)
|
||
if server:
|
||
delete_server(conn, server, log)
|
||
else:
|
||
log("compute уже нет")
|
||
|
||
try:
|
||
from gpu_rent.idle_killer import revoke_old_credentials
|
||
|
||
revoke_old_credentials(conn, log)
|
||
except Exception as exc:
|
||
log(f"revoke idle-killer app cred: {exc}")
|
||
|
||
if not cfg.keep_floating_ip:
|
||
delete_floating_ip(conn, state.floating_ip_id, state.floating_ip, log)
|
||
state.floating_ip = None
|
||
state.floating_ip_id = None
|
||
|
||
if destroy_disks:
|
||
for vid in (state.data_volume_id, state.boot_volume_id):
|
||
if not vid:
|
||
continue
|
||
try:
|
||
conn.block_storage.delete_volume(vid, ignore_missing=True)
|
||
log(f"удалён том {vid}")
|
||
except Exception as exc:
|
||
log(f"том {vid} не удалился: {exc}")
|
||
state.boot_volume_id = None
|
||
state.data_volume_id = None
|
||
|
||
state.server_id = None
|
||
state.server_name = None
|
||
state.bootstrapped = False
|
||
state.phase = "idle"
|
||
save_state(state)
|
||
try:
|
||
from gpu_rent.local_watchdog import clear_lease
|
||
|
||
clear_lease()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
from gpu_rent.balance import clear_balance_state
|
||
|
||
clear_balance_state()
|
||
except Exception:
|
||
pass
|
||
log("фаза idle" + ("" if destroy_disks else " (диски на месте)"))
|
||
return state
|