gpu-rent/src/gpu_rent/doctor.py

"""Preflight without creating a GPU. All checks print; exit 1 if session cannot start."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from gpu_rent.civitai import other_host, probe_me
from gpu_rent.config import Config, load_config
from gpu_rent.errors import CloudError, ConfigError
from gpu_rent.inventory import (
    gpu_quota_from_compute,
    looks_like_gpu,
    pick_boot_image,
    pick_volume_type,
    rank_flavors,
)
from gpu_rent.manifests import parse_extensions, parse_models
from gpu_rent.os_client import (
    compute_quotas,
    connect,
    find_snapshot_by_name,
    find_tagged_servers,
    iter_flavors,
    iter_images,
    iter_volume_types,
    volume_quotas,
)
from gpu_rent.payload import folder_bytes as _folder_bytes, has_payload as _has_payload
from gpu_rent.paths import env_path, runtime_dir
from gpu_rent.ssh_keys import key_ready
from gpu_rent.state import load_state


@dataclass
class Check:
    name: str
    ok: bool
    blocking: bool
    detail: str


def run_doctor() -> list[Check]:
    checks: list[Check] = []
    runtime_dir().mkdir(parents=True, exist_ok=True)

    env_file = env_path()
    if env_file.is_file():
        checks.append(Check("env file", True, True, str(env_file)))
    else:
        checks.append(
            Check(
                "env file",
                False,
                True,
                f"нет {env_file}. Скопируй env.example и заполни по docs/setup.md",
            )
        )
        return checks

    try:
        cfg = load_config(require_auth=True)
    except ConfigError as exc:
        checks.append(Check("OS_*", False, True, str(exc)))
        return checks
    checks.append(
        Check(
            "OS_*",
            True,
            True,
            f"project={cfg.os_project_id} region={cfg.os_region_name} az={cfg.gpu_rent_az}",
        )
    )

    if key_ready(cfg.ssh_private_key_path):
        checks.append(Check("SSH key", True, True, f"есть {cfg.ssh_private_key_path}"))
    else:
        checks.append(
            Check(
                "SSH key",
                True,
                True,
                f"ключа нет — CLI создаст {cfg.ssh_private_key_path} на первом up",
            )
        )

    conn = None
    try:
        conn = connect(cfg)
        checks.append(Check("Keystone", True, True, "IAM-токен выдан (TTL ~24 ч, sdk обновит)"))
    except CloudError as exc:
        checks.append(Check("Keystone", False, True, str(exc)))
        _local_manifests(cfg, checks)
        _local_folders(cfg, checks)
        return checks

    try:
        quota = compute_quotas(conn)
        gpu_limit = gpu_quota_from_compute(quota)
        if gpu_limit is None:
            checks.append(
                Check(
                    "квота GPU",
                    True,
                    False,
                    "в compute quota нет поля gpu — смотри панель IAM -> проект -> квоты. "
                    "Если там 0, напишите в поддержку Selectel (docs/setup.md).",
                )
            )
        elif gpu_limit <= 0:
            checks.append(
                Check(
                    "квота GPU",
                    False,
                    True,
                    "квота GPU = 0. Напиши в поддержку Selectel: 1× RTX 4090 24 GB "
                    f"в {cfg.os_region_name}/{cfg.gpu_rent_az}, проект {cfg.os_project_id}. "
                    "Текст тикета — docs/setup.md. CLI сервер не создаст, пока лимит 0.",
                )
            )
        else:
            checks.append(Check("квота GPU", True, True, f"limit={gpu_limit}"))
    except CloudError as exc:
        checks.append(Check("квота GPU", False, True, str(exc)))

    flavors = list(iter_flavors(conn))
    gpu_flavors = [f for f in flavors if looks_like_gpu(f)]
    ranked = rank_flavors(gpu_flavors or flavors, cfg.flavor_preference)
    if cfg.default_flavor_id:
        hit = next((f for f in flavors if getattr(f, "id", None) == cfg.default_flavor_id), None)
        if hit:
            checks.append(Check("flavor", True, True, f"DEFAULT_FLAVOR_ID={cfg.default_flavor_id}"))
        else:
            checks.append(
                Check(
                    "flavor",
                    False,
                    True,
                    f"DEFAULT_FLAVOR_ID={cfg.default_flavor_id} в регионе нет",
                )
            )
    elif not ranked:
        names = ", ".join(getattr(f, "name", "?") for f in gpu_flavors[:8]) or "нет GPU-flavors"
        checks.append(
            Check(
                "flavor",
                False,
                True,
                "ни один flavor из FLAVOR_PREFERENCE не найден в этом пуле. "
                f"Видно: {names}. Смени GPU_RENT_AZ / OS_REGION_NAME по матрице GPU.",
            )
        )
    else:
        first = ranked[0]
        rest = ", ".join(f"{x.label}:{x.name}" for x in ranked[1:3])
        extra = f"; дальше {rest}" if rest else ""
        checks.append(
            Check(
                "flavor",
                True,
                True,
                f"первый доступный {first.label} -> {first.name} ({first.id}){extra}",
            )
        )

    types = list(iter_volume_types(conn))
    vtype = pick_volume_type(types, cfg.gpu_rent_az)
    gigabytes = volume_quotas(conn).get("gigabytes")
    if vtype:
        disk_note = f"type={vtype}, data {cfg.data_volume_size_gb} GB"
        if isinstance(gigabytes, int) and gigabytes >= 0:
            disk_note += f", quota gigabytes={gigabytes}"
            if gigabytes < cfg.data_volume_size_gb:
                checks.append(
                    Check(
                        "диск",
                        False,
                        True,
                        f"{disk_note} — квота меньше {cfg.data_volume_size_gb} GB",
                    )
                )
            else:
                checks.append(Check("диск", True, True, disk_note))
        else:
            checks.append(Check("диск", True, False, disk_note + " (квоту дисков API не отдал)"))
    else:
        checks.append(
            Check(
                "диск",
                False,
                True,
                f"нет volume type для AZ {cfg.gpu_rent_az}. volume type list пуст?",
            )
        )

    chosen = pick_boot_image(list(iter_images(conn)))
    if chosen:
        name = getattr(chosen, "name", str(chosen))
        dockerish = "docker" in name.lower()
        checks.append(
            Check(
                "образ GPU",
                True,
                False,
                name
                + (
                    " (это Docker-образ: в пуле нет варианта без Docker)"
                    if dockerish
                    else ""
                ),
            )
        )
    else:
        checks.append(
            Check(
                "образ GPU",
                False,
                True,
                "не нашёл GPU-образ в Glance (ожидаем Ubuntu 24.04 Driver 580 без Docker)",
            )
        )

    servers = find_tagged_servers(conn)
    if servers:
        names = ", ".join(f"{s.name}:{s.status}" for s in servers)
        checks.append(Check("живой gpu-rent", True, False, names))
    else:
        checks.append(Check("живой gpu-rent", True, False, "серверов с тегом нет (это норма)"))

    snap = find_snapshot_by_name(conn, cfg.boot_snapshot_name)
    if snap:
        checks.append(Check("boot snapshot", True, False, cfg.boot_snapshot_name))
    else:
        checks.append(
            Check("boot snapshot", True, False, f"{cfg.boot_snapshot_name} ещё нет — будет после первого bootstrap")
        )

    _civitai(cfg, checks)
    _huggingface(cfg, checks)
    _local_manifests(cfg, checks)
    _local_folders(cfg, checks)
    return checks


def _civitai(cfg: Config, checks: list[Check]) -> None:
    if not cfg.civitai_api_token:
        checks.append(
            Check(
                "Civitai",
                True,
                False,
                "токена нет — на первом диске будет дефолт SwarmUI. Ключ: docs/setup.md §4",
            )
        )
        return
    probe = probe_me(cfg.civitai_api_token, cfg.civitai_api_host)
    if probe.ok:
        checks.append(Check("Civitai", True, True, f"{probe.host}: {probe.detail}"))
        return
    alt = probe_me(cfg.civitai_api_token, other_host(cfg.civitai_api_host))
    if alt.ok:
        checks.append(
            Check(
                "Civitai",
                True,
                False,
                f"{probe.host} не ответил ({probe.detail}); {alt.host} принял токен. "
                "Для NSFW оставь CIVITAI_API_HOST=civitai.red",
            )
        )
        return
    checks.append(
        Check(
            "Civitai",
            True,
            False,
            f"{probe.host}: {probe.detail}; fallback {alt.host}: {alt.detail}. "
            "seed-models недоступен, дефолт SwarmUI ок — up не блокируем",
        )
    )


def _huggingface(cfg: Config, checks: list[Check]) -> None:
    from gpu_rent.huggingface import is_huggingface_url, probe_whoami
    from gpu_rent.manifests import parse_models

    needs_hf = False
    try:
        for e in parse_models(cfg.models_manifest):
            if e.url and is_huggingface_url(e.url):
                needs_hf = True
                break
    except Exception:
        pass

    if not cfg.hf_token:
        checks.append(
            Check(
                "Hugging Face",
                True,
                False,
                (
                    "HF_TOKEN нет — gated HF URL в models.yaml дадут 401. "
                    "https://huggingface.co/settings/tokens"
                    if needs_hf
                    else "токена нет (опционально для HF URL / capture fallback)"
                ),
            )
        )
        return
    probe = probe_whoami(cfg.hf_token)
    if probe.ok:
        who = f" ({probe.name})" if probe.name else ""
        checks.append(Check("Hugging Face", True, True, f"токен ок{who}"))
    else:
        checks.append(
            Check(
                "Hugging Face",
                True,
                False,
                f"токен не принят: {probe.detail}",
            )
        )


def _local_manifests(cfg: Config, checks: list[Check]) -> None:
    try:
        models = parse_models(cfg.models_manifest)
        if cfg.models_manifest.is_file():
            checks.append(
                Check(
                    "models.yaml",
                    True,
                    True,
                    f"{cfg.models_manifest} — записей (без заглушек 0): {len(models)}",
                )
            )
        else:
            checks.append(
                Check(
                    "models.yaml",
                    True,
                    False,
                    f"нет {cfg.models_manifest} — seed Civitai пропустится",
                )
            )
    except ConfigError as exc:
        checks.append(Check("models.yaml", False, True, str(exc)))

    try:
        repos = parse_extensions(cfg.extensions_manifest)
        if not cfg.extensions_manifest.is_file():
            checks.append(Check("extensions.yaml", True, False, "файла нет — стоковый SwarmUI"))
        elif repos:
            checks.append(Check("extensions.yaml", True, True, f"{len(repos)} git-реп"))
        else:
            example = cfg.extensions_manifest.with_name("extensions.example.yaml")
            n_ex = 0
            if example.is_file():
                try:
                    n_ex = len(parse_extensions(example))
                except ConfigError:
                    n_ex = 0
            if n_ex:
                checks.append(
                    Check(
                        "extensions.yaml",
                        True,
                        False,
                        f"пустой; в example {n_ex} реп — скопируй и gpu-rent seed-extensions",
                    )
                )
            else:
                checks.append(
                    Check("extensions.yaml", True, False, "0 реп — стоковый SwarmUI")
                )
    except ConfigError as exc:
        checks.append(Check("extensions.yaml", False, True, str(exc)))


def _local_folders(cfg: Config, checks: list[Check]) -> None:
    size = _folder_bytes(cfg.local_models_dir)
    gi = size / (1024**3)
    payload = _has_payload(cfg.local_models_dir)
    if payload and gi > cfg.data_volume_size_gb * 0.8:
        checks.append(
            Check(
                "Models/",
                False,
                True,
                f"локально ~{gi:.1f} GB, диск {cfg.data_volume_size_gb} GB — не влезет. "
                "Урежь папку или gpu-rent resize-data после первого диска.",
            )
        )
    elif payload:
        checks.append(Check("Models/", True, False, f"есть файлы, ~{gi:.2f} GB — уедут на up"))
    else:
        checks.append(Check("Models/", True, False, "пусто — на up ничего не грузим"))

    for label, folder in (
        ("Wildcards/", cfg.local_wildcards_dir),
        ("CustomWorkflows/", cfg.local_workflows_dir),
    ):
        if _has_payload(folder):
            checks.append(Check(label, True, False, "не пусто — push на up"))
        else:
            checks.append(Check(label, True, False, "пусто — skip"))


def blocking_failed(checks: list[Check]) -> list[Check]:
    return [c for c in checks if c.blocking and not c.ok]


def dry_run_plan(checks: list[Check]) -> list[str]:
    cfg = load_config(require_auth=False)
    state = load_state()
    lines = [
        f"фаза state: {state.phase}",
        f"пул {cfg.os_region_name} / AZ {cfg.gpu_rent_az}",
        f"data volume: {cfg.data_volume_size_gb} GB (рост только вверх)",
        f"preemptible: {cfg.default_spot} (обычный сервер: gpu-rent up --no-spot)",
        f"idle-killer: {cfg.idle_minutes} мин пустой очереди, льгота {cfg.idle_grace_minutes} мин",
        "₽: в API нет — смотри панель; диск 24/7 даже после stop",
        f"туннель: localhost:{cfg.swarmui_local_port} -> VM :7801",
        "gpu-rent up --yes создаст GPU + SwarmUI и откроет туннель :17801 "
        "(Ctrl+C не гасит GPU, Ctrl+D гасит)",
        "только облако без туннеля: gpu-rent up --yes --no-tunnel",
    ]
    flavor = next((c.detail for c in checks if c.name == "flavor" and c.ok), None)
    if flavor:
        lines.insert(2, f"flavor: {flavor}")
    return lines