Files
gpu-rent/src/gpu_rent/session.py
T
Leonid Pershin 57d38bd9f6 Enhance backend loading diagnostics and remount logic
- Introduced a new `loading_fail_sec` parameter in the `wait_backend_idle` function to handle prolonged loading states, improving error handling for backend readiness.
- Updated the `ensure_dlbackend_bind` function to stop SwarmUI before remounting, preventing target busy errors and ensuring consistent data mounts.
- Enhanced the `recover_errored_backends` function to account for the new remount logic, improving backend recovery processes.
- Refactored tests to validate the new loading failure conditions and ensure proper handling of backend states during diagnostics.
2026-08-21 12:24:31 +03:00

657 lines
25 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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 (
ensure_swarm_comfy_installed,
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:
ensure_swarm_comfy_installed(cfg, ip, log)
except Exception as exc:
log(f"SwarmUI Comfy install: {exc}")
try:
from gpu_rent.ready import collect_swarm_diagnostics
collect_swarm_diagnostics(cfg, ip, log)
except Exception as diag_exc:
log(f"diag: {diag_exc}")
raise
clock.mark("comfy-install", log)
try:
wait_backend_idle(cfg, ip, log)
except CloudError as exc:
# wait_backend_idle already collected diag on errored/timeout
log(f"ready: {exc}")
raise
clock.mark("Idle", log)
try:
if tune_swarm_perf(cfg, ip, log):
# Stop before remount — otherwise umount fails with target busy.
run_ssh(
cfg,
ip,
"sudo -n systemctl stop swarmui",
check=False,
timeout=120,
)
run_ssh(
cfg,
ip,
"sudo -n bash -c '"
"for p in dlbackend Data Models Output; do "
"src=/mnt/swarm_data/$p; dst=/opt/swarmui/$p; "
"mkdir -p \"$src\" \"$dst\"; "
"cur=$(findmnt -n -o SOURCE --target \"$dst\" 2>/dev/null || true); "
"if [[ \"$cur\" != \"$src\" ]]; then "
"umount -l \"$dst\" 2>/dev/null || umount \"$dst\" 2>/dev/null || true; "
"mount --bind \"$src\" \"$dst\" || true; "
"echo remounted $dst; "
"fi; "
"done'",
check=False,
timeout=60,
)
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} — recover backends")
try:
ensure_swarm_comfy_installed(cfg, ip, log)
wait_backend_idle(cfg, ip, log)
except Exception as exc2:
log(f"recover after perf: {exc2}")
raise exc from exc2
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":
# Prefer FIP from state; may still need ensure below.
fip = state.floating_ip
# Warm if local state says so, OR VM still has bootstrap marker
# (local bootstrapped often cleared on stop while disk/VM remain).
vm_warm = False
if fip and not (state.bootstrapped and state.floating_ip):
try:
if probe_ssh(cfg, fip, attempts=1) == "ok":
probe = run_ssh(
cfg,
fip,
"test -f /opt/swarmui/.gpu-rent-bootstrapped && echo yes || echo no",
check=False,
timeout=20,
).strip()
vm_warm = probe.endswith("yes") or probe == "yes"
except Exception:
vm_warm = False
if (state.bootstrapped and state.floating_ip) or vm_warm:
if update is None:
do_update = False
log(
"warm ACTIVE — без git pull SwarmUI/extensions "
"(нужен свежий: --update)"
)
else:
do_update = update
if vm_warm and not state.bootstrapped:
log("warm: маркер на VM — восстанавливаю local bootstrapped")
state.bootstrapped = True
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.
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