gpu-rent/src/gpu_rent/provision.py

"""After OS bootstrap: extensions, autocomplete, Civitai, then start SwarmUI."""

from __future__ import annotations

import json
import shlex
from collections.abc import Callable
from datetime import datetime, timezone
from importlib.resources import files
from pathlib import Path

import httpx

from gpu_rent.civitai import (
    fetch_model_version,
    file_size_bytes,
    pick_preview_image,
    pick_primary_file,
)
from gpu_rent.config import Config
from gpu_rent.errors import CloudError, GpuRentError
from gpu_rent.manifests import (
    MODEL_DIRS,
    extract_version_id,
    parse_extensions,
    parse_models,
    remote_root_for,
    repo_dirname,
    repo_matches_runtime,
)
from gpu_rent.ssh_ops import put_text, remote_exists, run_python, run_ssh
from gpu_rent.sync_files import pull_tree, push_tree

Log = Callable[[str], None]
DATA = "/mnt/swarm_data"

# Forwarded to remote install_*.sh (from .env / gpu-rent.vars → os.environ).
_OLLAMA_INSTALL_ENV = ("OLLAMA_VERSION", "OLLAMA_SHA256")


def _models_manifest_fp(cfg: Config) -> str:
    """Stable fingerprint of models.yaml so warm seed can skip API+fetch."""
    import hashlib

    path = Path(cfg.models_manifest)
    try:
        raw = path.read_bytes()
    except OSError:
        raw = b""
    return hashlib.sha256(raw).hexdigest()[:16]


def _remote_llm_env(cfg: Config, *keys: str) -> dict[str, str]:
    import os

    env: dict[str, str] = {"SWARM_USER": cfg.ssh_user}
    for key in keys:
        value = (os.environ.get(key) or "").strip()
        if value:
            env[key] = value
    return env


def _pkg_text(name: str) -> str:
    return files("gpu_rent.remote").joinpath(name).read_text(encoding="utf-8")


def probe_gpu(cfg: Config, host: str, log: Log) -> dict:
    """Write /mnt/swarm_data/.gpu-rent-gpu.json; return parsed dict."""
    out = run_python(
        cfg,
        host,
        _pkg_text("gpu_probe.py"),
        remote_path="/tmp/gpu-rent-gpu_probe.py",
        timeout=60,
        log=log,
    )
    # Last JSON line from script stdout
    data: dict = {}
    for line in reversed(out.splitlines()):
        line = line.strip()
        if line.startswith("{"):
            try:
                data = json.loads(line)
                break
            except json.JSONDecodeError:
                continue
    if data.get("ok"):
        log(
            f"GPU: {data.get('name')} "
            f"{data.get('vram_mib')} MiB cap={data.get('compute_cap')} "
            f"tier={data.get('tier')}"
        )
    else:
        log(f"GPU probe: {data.get('error') or 'нет данных'}")
    return data


def tune_swarm_perf(cfg: Config, host: str, log: Log) -> bool:
    """Install sage/triton into Comfy venv + ExtraArgs. Returns True if SwarmUI restart needed."""
    out = run_python(
        cfg,
        host,
        _pkg_text("tune_swarm_perf.py"),
        remote_path="/tmp/gpu-rent-tune_swarm_perf.py",
        timeout=900,
        log=log,
    )
    return "RESTART_SWARMUI=1" in out


def ensure_swarm_comfy_installed(cfg: Config, host: str, log: Log) -> None:
    """Headless Comfy install / recover errored backends before ready wait."""
    # Diag script next to install so recover-fail can subprocess it on the VM.
    put_text(cfg, host, "/tmp/gpu-rent-swarm_diag.py", _pkg_text("swarm_diag.py"))
    log("SwarmUI Comfy: install если empty, RestartBackends если errored…")
    run_python(
        cfg,
        host,
        _pkg_text("install_swarm_comfy.py"),
        remote_path="/tmp/gpu-rent-install_swarm_comfy.py",
        # Cold: git clone Comfy + pip torch — often 15–40 min.
        timeout=3900,
        log=log,
    )


def seed_extensions(cfg: Config, host: str, log: Log, *, update: bool = True) -> bool:
    from gpu_rent.llm_runtime import normalize_runtime

    runtime = normalize_runtime(cfg.llm_runtime)
    all_repos = parse_extensions(cfg.extensions_manifest)
    repos = [r for r in all_repos if repo_matches_runtime(r, runtime)]
    skipped = [r for r in all_repos if r not in repos]
    for repo in skipped:
        log(
            f"extensions: пропуск {repo_dirname(repo)} "
            f"(requires={repo.requires}, LLM_RUNTIME={runtime})"
        )
    if not repos and not update:
        if all_repos:
            log("extensions: все строки отфильтрованы по requires — стоковый SwarmUI")
        else:
            log("extensions.yaml пуст — стоковый SwarmUI")
        return False
    if not repos and update:
        if not all_repos:
            log("extensions.yaml пуст — только update уже установленных на data (если есть)")
        else:
            log(
                "extensions: yaml отфильтрован по requires "
                f"(LLM_RUNTIME={runtime}) — только update установленных"
            )
    jobs = []
    for repo in repos:
        jobs.append(
            {
                "kind": repo.kind,
                "url": repo.url,
                "ref": repo.ref,
                "dest": remote_root_for(repo),
            }
        )
    put_text(cfg, host, "/tmp/gpu-rent-ext.json", json.dumps(jobs, indent=2))
    put_text(cfg, host, "/tmp/gpu-rent-update-git", "1\n" if update else "0\n")
    if cfg.git_token:
        put_text(cfg, host, "/tmp/gpu-rent-git.token", cfg.git_token + "\n", mode=0o600)
    if update:
        log(f"extensions: clone/update {len(jobs)} из yaml + установленные на data")
    else:
        log(f"extensions: только недостающие из yaml ({len(jobs)}), без git pull (--no-update)")
    out = run_python(
        cfg,
        host,
        _pkg_text("clone_ext.py"),
        remote_path="/tmp/gpu-rent-clone_ext.py",
        timeout=1800,
        log=log,
    )
    return "cloned " in out or "updated " in out


def _github_blob(cfg: Config) -> dict | None:
    url = (
        f"https://api.github.com/repos/{cfg.autocomplete_github_repo}/contents/"
        f"{cfg.autocomplete_github_path}?ref={cfg.autocomplete_github_ref}"
    )
    headers = {"Accept": "application/vnd.github+json", "User-Agent": "gpu-rent"}
    if cfg.git_token:
        headers["Authorization"] = f"Bearer {cfg.git_token}"
    try:
        with httpx.Client(timeout=20.0, follow_redirects=True) as client:
            response = client.get(url, headers=headers)
    except httpx.HTTPError as exc:
        log_skip = str(exc)
        return {"error": log_skip}
    if response.status_code == 403:
        return {"error": "GitHub rate limit — autocomplete не обновляю"}
    if response.status_code != 200:
        return {"error": f"GitHub HTTP {response.status_code}"}
    data = response.json()
    if not isinstance(data, dict):
        return {"error": "неожиданный JSON GitHub"}
    return data


def seed_autocomplete(cfg: Config, host: str, log: Log) -> bool:
    if not cfg.autocomplete_enabled:
        log("autocomplete выключен")
        return False
    dest_dir = f"{DATA}/Data/Autocompletions"
    dest = f"{dest_dir}/{cfg.autocomplete_filename}"
    meta_path = f"{dest}.gpu-rent-meta.json"
    blob = _github_blob(cfg)
    if blob is None:
        return False
    if blob.get("error"):
        log(str(blob["error"]))
        return False
    sha = str(blob.get("sha") or "")
    download_url = str(blob.get("download_url") or "")
    old_sha = ""
    if remote_exists(cfg, host, meta_path):
        raw = run_ssh(cfg, host, f"cat {meta_path}", check=False)
        try:
            old_sha = str(json.loads(raw).get("github_blob_sha") or "")
        except json.JSONDecodeError:
            old_sha = ""
    changed = sha != old_sha or not remote_exists(cfg, host, dest)
    if changed:
        if not download_url:
            log("GitHub не дал download_url")
            return False
        log(f"качаю {cfg.autocomplete_filename}")
        run_ssh(
            cfg,
            host,
            "mkdir -p {dir} && curl -fsSL -o {part} {url} && mv {part} {dest}".format(
                dir=shlex.quote(dest_dir),
                part=shlex.quote(dest + ".partial"),
                url=shlex.quote(download_url),
                dest=shlex.quote(dest),
            ),
            timeout=180,
        )
        meta = {
            "repo": cfg.autocomplete_github_repo,
            "path": cfg.autocomplete_github_path,
            "ref": cfg.autocomplete_github_ref,
            "github_blob_sha": sha,
            "filename": cfg.autocomplete_filename,
            "fetched_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
            "settings_applied": True,
        }
        put_text(cfg, host, meta_path, json.dumps(meta, indent=2) + "\n")
    settings = f"{DATA}/Data/Settings.fds"
    applied = False
    if remote_exists(cfg, host, meta_path):
        raw = run_ssh(cfg, host, f"cat {meta_path}", check=False)
        try:
            applied = bool(json.loads(raw).get("settings_applied"))
        except json.JSONDecodeError:
            applied = False
    if not applied:
        _merge_autocomplete_into_settings(cfg, host, settings, cfg.autocomplete_filename, log)
        if remote_exists(cfg, host, meta_path):
            try:
                meta_obj = json.loads(
                    run_ssh(cfg, host, f"cat {meta_path}", check=False)
                )
            except json.JSONDecodeError:
                meta_obj = {}
            if isinstance(meta_obj, dict):
                meta_obj["settings_applied"] = True
                put_text(cfg, host, meta_path, json.dumps(meta_obj, indent=2) + "\n")
    return changed


_AUTOCOMPLETE_MERGE_PY = r'''
#!/usr/bin/env python3
"""Merge AutoComplete.Source into Settings.fds without wiping other keys."""
from __future__ import annotations

import os
import re
import sys
import time
from pathlib import Path

p = Path(os.environ.get("GPU_RENT_SETTINGS_FDS") or "/mnt/swarm_data/Data/Settings.fds")
fname = (os.environ.get("GPU_RENT_AUTOCOMPLETE_FILE") or "").strip()
if not fname:
    print("no GPU_RENT_AUTOCOMPLETE_FILE", file=sys.stderr)
    raise SystemExit(1)

installed_block = (
    "IsInstalled: true\n"
    f"InstallDate: {time.strftime('%Y-%m-%d')}\n"
    "InstallVersion: gpu-rent\n"
)

block = (
    installed_block
    + "DefaultUser:\n"
    "    AutoComplete:\n"
    f"        Source: {fname}\n"
    "        EscapeParens: true\n"
)

if not p.is_file():
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(block, encoding="utf-8")
    print(f"created Settings.fds IsInstalled+AutoComplete.Source={fname}")
    raise SystemExit(0)

text = p.read_text(encoding="utf-8", errors="replace")
# Never leave a Settings.fds that sends the UI to /Install.
if not re.search(r"(?im)^\s*IsInstalled:\s*true\s*$", text):
    if re.search(r"(?im)^\s*IsInstalled:\s*", text):
        text = re.sub(r"(?im)^(\s*IsInstalled:\s*).*$", r"\1true", text, count=1)
    else:
        text = installed_block + text
    p.write_text(text, encoding="utf-8")
    print("patched IsInstalled: true (was missing/false)")
    text = p.read_text(encoding="utf-8", errors="replace")

if re.search(rf"^\s*Source:\s*{re.escape(fname)}\s*$", text, re.M):
    print(f"AutoComplete.Source already {fname}")
    raise SystemExit(0)

# Replace Source line if AutoComplete section exists
new, n = re.subn(
    r"(^[ \t]*Source:\s*).*$",
    rf"\1{fname}",
    text,
    count=1,
    flags=re.M,
)
if n and "AutoComplete" in text:
    p.write_text(new, encoding="utf-8")
    print(f"patched AutoComplete.Source={fname}")
    raise SystemExit(0)

if re.search(r"^DefaultUser:\s*$", text, re.M):
    new = re.sub(
        r"^(DefaultUser:\s*\n)",
        (
            r"\1    AutoComplete:\n"
            f"        Source: {fname}\n"
            "        EscapeParens: true\n"
        ),
        text,
        count=1,
        flags=re.M,
    )
    p.write_text(new, encoding="utf-8")
    print(f"inserted AutoComplete under DefaultUser Source={fname}")
else:
    p.write_text(text.rstrip() + "\n\n" + block, encoding="utf-8")
    print(f"appended DefaultUser.AutoComplete Source={fname}")
'''


_ENSURE_INSTALLED_PY = r'''
#!/usr/bin/env python3
"""Ensure Settings.fds has IsInstalled: true so UI skips /Install wizard."""
from __future__ import annotations

import os
import re
import sys
import time
from pathlib import Path

paths = []
env = (os.environ.get("GPU_RENT_SETTINGS_FDS") or "").strip()
if env:
    paths.append(Path(env))
paths.extend(
    [
        Path("/mnt/swarm_data/Data/Settings.fds"),
        Path("/opt/swarmui/Data/Settings.fds"),
    ]
)
# Unique while preserving order
seen = set()
uniq: list[Path] = []
for p in paths:
    key = str(p)
    if key in seen:
        continue
    seen.add(key)
    uniq.append(p)

changed = False
for p in uniq:
    try:
        same = False
        if p.is_file():
            for q in uniq:
                if q is not p and q.is_file():
                    try:
                        if p.resolve() == q.resolve():
                            same = True
                            break
                    except OSError:
                        pass
        if same and changed:
            continue
        if not p.is_file():
            # Only create on data volume path
            if "/mnt/swarm_data/" not in str(p):
                continue
            p.parent.mkdir(parents=True, exist_ok=True)
            p.write_text(
                "IsInstalled: true\n"
                f"InstallDate: {time.strftime('%Y-%m-%d')}\n"
                "InstallVersion: gpu-rent\n"
                "DefaultUser:\n"
                "    Theme: modern_dark\n",
                encoding="utf-8",
            )
            print(f"created {p} with IsInstalled")
            changed = True
            continue
        text = p.read_text(encoding="utf-8", errors="replace")
        if re.search(r"(?im)^\s*IsInstalled:\s*true\s*$", text):
            print(f"ok {p}: IsInstalled true")
            continue
        if re.search(r"(?im)^\s*IsInstalled:\s*", text):
            text = re.sub(r"(?im)^(\s*IsInstalled:\s*).*$", r"\1true", text, count=1)
        else:
            text = (
                "IsInstalled: true\n"
                f"InstallDate: {time.strftime('%Y-%m-%d')}\n"
                "InstallVersion: gpu-rent\n"
                + text
            )
        if not re.search(r"(?im)^\s*InstallDate:\s*", text):
            text = re.sub(
                r"(?im)^(\s*IsInstalled:\s*true\s*\n)",
                rf"\1InstallDate: {time.strftime('%Y-%m-%d')}\n",
                text,
                count=1,
            )
        p.write_text(text, encoding="utf-8")
        print(f"patched {p}: IsInstalled true")
        changed = True
    except OSError as exc:
        print(f"skip {p}: {exc}")

print("CHANGED" if changed else "NOOP")
'''


def _merge_autocomplete_into_settings(
    cfg: Config, host: str, settings_path: str, filename: str, log: Log
) -> None:
    """Patch AutoComplete.Source in Settings.fds without wiping the rest."""
    out = run_python(
        cfg,
        host,
        _AUTOCOMPLETE_MERGE_PY,
        remote_path="/tmp/gpu-rent-patch_autocomplete.py",
        timeout=60,
        log=None,
        env={
            "GPU_RENT_SETTINGS_FDS": settings_path,
            "GPU_RENT_AUTOCOMPLETE_FILE": filename,
        },
    )
    for line in (out or "").splitlines():
        if line.strip():
            log(line.strip())


def ensure_settings_is_installed(cfg: Config, host: str, log: Log) -> bool:
    """Make sure Settings.fds has IsInstalled:true (UI /Install wizard).

    Returns True if the file was created/patched (caller may need SwarmUI restart).
    """
    out = run_python(
        cfg,
        host,
        _ENSURE_INSTALLED_PY,
        remote_path="/tmp/gpu-rent-ensure_installed.py",
        timeout=60,
        log=None,
        env={"GPU_RENT_SETTINGS_FDS": f"{DATA}/Data/Settings.fds"},
    )
    changed = False
    for line in (out or "").splitlines():
        line = line.strip()
        if not line:
            continue
        log(f"Settings: {line}")
        if line == "CHANGED":
            changed = True
    return changed


def _download_url(host: str, version_id: int, file_info: dict) -> str:
    raw = str(file_info.get("downloadUrl") or "")
    if "civitai." in raw and "/api/download/" in raw:
        # NSFW downloadUrl often points at .com — use the host that had files.
        return f"https://{host}/api/download/models/{version_id}"
    if raw.startswith("https://"):
        return raw
    return f"https://{host}/api/download/models/{version_id}"


def seed_swarmui_api_keys(cfg: Config, host: str, log: Log) -> None:
    """Write CIVITAI_API_TOKEN (and HF if set) into SwarmUI user keys via SetAPIKey.

    Swarm stores them in Users.ldb GenericData — needed for Model Downloader in the UI.
    Call after SwarmUI HTTP is up (after wait_backend / verify).
    """
    keys: dict[str, str] = {}
    if cfg.civitai_api_token:
        keys["civitai_api"] = cfg.civitai_api_token
    hf = (cfg.hf_token or "").strip()
    if hf:
        keys["huggingface_api"] = hf
    if not keys:
        log("SwarmUI API keys: нет CIVITAI_API_TOKEN / HF_TOKEN — skip")
        return
    put_text(
        cfg,
        host,
        "/tmp/gpu-rent-swarm-api-keys.json",
        json.dumps(keys) + "\n",
        mode=0o600,
    )
    names = ", ".join(keys)
    log(f"SwarmUI: прокидываю API keys ({names})")
    try:
        run_python(
            cfg,
            host,
            _pkg_text("swarmui_set_api_keys.py"),
            remote_path="/tmp/gpu-rent-swarmui_set_api_keys.py",
            timeout=180,
            log=log,
        )
    except CloudError as exc:
        log(f"⚠ SwarmUI API keys: {exc}")
        run_ssh(cfg, host, "rm -f /tmp/gpu-rent-swarm-api-keys.json", check=False)


def seed_civitai(cfg: Config, host: str, log: Log) -> None:
    from gpu_rent.huggingface import is_huggingface_url

    entries = parse_models(cfg.models_manifest)
    if not entries:
        log("model-seed пропущен: манифест пуст — дефолт SwarmUI")
        return

    # Warm short-circuit: last successful jobs list still present on disk.
    cache_remote = f"{DATA}/.gpu-rent-models-jobs.json"
    try:
        cached_raw = run_ssh(
            cfg,
            host,
            f"test -f {cache_remote} && cat {cache_remote} || true",
            check=False,
            timeout=20,
        ).strip()
    except Exception:
        cached_raw = ""
    if cached_raw.startswith("{"):
        try:
            cached = json.loads(cached_raw)
            dests = [str(d) for d in (cached.get("dests") or []) if d]
            fp = str(cached.get("fp") or "")
            want_fp = _models_manifest_fp(cfg)
            if dests and fp == want_fp:
                check = " && ".join(f"test -f {shlex.quote(d)}" for d in dests)
                ok = run_ssh(
                    cfg,
                    host,
                    f"if {check}; then echo ALL_OK; else echo MISSING; fi",
                    check=False,
                    timeout=30,
                ).strip()
                if "ALL_OK" in ok:
                    log(
                        f"model-seed: skip — {len(dests)} файл(ов) уже на диске (cache)"
                    )
                    return
        except (json.JSONDecodeError, TypeError, CloudError):
            pass

    jobs: list[dict] = []
    for entry in entries:
        url = (entry.url or "").strip()
        if url and is_huggingface_url(url):
            name = url.rstrip("/").rsplit("/", 1)[-1].split("?", 1)[0] or "model.safetensors"
            folder = MODEL_DIRS.get(entry.kind, entry.kind)
            dest = f"{DATA}/Models/{folder}/{name}"
            stem = Path(name).stem
            swarm = {
                "name": stem,
                "title": stem,
                "description": f"Hugging Face: {url}",
                "trigger_phrase": "",
                "author": "",
                "tags": ["huggingface"],
            }
            jobs.append(
                {
                    "dest": dest,
                    "url": url,
                    "sha256": "",
                    "auth": "hf",
                    "sidecars": {
                        f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
                    },
                }
            )
            continue

        if not cfg.civitai_api_token:
            continue
        vid = entry.version_id or (extract_version_id(entry.url) if entry.url else None)
        if not vid:
            log(f"пропуск {entry.kind}: нет version_id / HF url")
            continue
        try:
            api_host, version = fetch_model_version(
                cfg.civitai_api_token, cfg.civitai_api_host, vid
            )
        except CloudError as exc:
            log(str(exc))
            continue
        info = pick_primary_file(version)
        if not info:
            log(f"version {vid}: нет files[]")
            continue
        name = str(info.get("name") or f"{vid}.safetensors")
        folder = MODEL_DIRS.get(entry.kind, entry.kind)
        dest = f"{DATA}/Models/{folder}/{name}"
        hashes = info.get("hashes") or {}
        sha = str((hashes.get("SHA256") or hashes.get("sha256") or "")).lower()
        stem = Path(name).stem
        civitai_json = json.dumps(version, ensure_ascii=False, indent=2)
        trained = version.get("trainedWords") or []
        phrase = trained[0] if isinstance(trained, list) and trained else ""
        swarm = {
            "name": stem,
            "title": version.get("name") or stem,
            "description": (version.get("description") or "")[:2000],
            "trigger_phrase": phrase,
            "author": (
                ((version.get("model") or {}) if isinstance(version.get("model"), dict) else {}).get(
                    "name"
                )
            ),
            "tags": version.get("tags") or [],
        }
        job: dict = {
            "dest": dest,
            "url": _download_url(api_host, vid, info),
            "sha256": sha,
            "auth": "civitai",
            "sidecars": {
                f"{stem}.civitai.json": civitai_json,
                f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
            },
        }
        size_b = file_size_bytes(info)
        if size_b:
            job["size"] = size_b
        preview = pick_preview_image(version)
        if preview:
            preview_url, preview_suffix = preview
            job["preview_url"] = preview_url
            job["preview_dest"] = f"{DATA}/Models/{folder}/{stem}{preview_suffix}"
        jobs.append(job)

    if not jobs:
        if not cfg.civitai_api_token:
            log("Civitai-seed пропущен: нет CIVITAI_API_TOKEN — дефолт SwarmUI")
        else:
            log("Civitai-seed: ни одной скачиваемой строки")
        return

    hf_n = sum(1 for j in jobs if j.get("auth") == "hf")
    civ_n = len(jobs) - hf_n
    if hf_n and not cfg.hf_token:
        log(
            "⚠ в манифесте есть Hugging Face URL, но нет HF_TOKEN — "
            "gated/abliterated файлы дадут 401. Токен: huggingface.co/settings/tokens"
        )
    if civ_n and not cfg.civitai_api_token:
        log("⚠ Civitai-строки пропущены: нет CIVITAI_API_TOKEN")

    if not any(e.kind == "checkpoint" for e in entries):
        log("в манифесте нет checkpoint — генерация может не стартовать")

    dests = [str(j.get("dest") or "") for j in jobs if j.get("dest")]
    if dests:
        check = " && ".join(f"test -f {shlex.quote(d)}" for d in dests)
        present = run_ssh(
            cfg,
            host,
            f"if {check}; then echo ALL_OK; else echo MISSING; fi",
            check=False,
            timeout=30,
        ).strip()
        if "ALL_OK" in present:
            put_text(
                cfg,
                host,
                cache_remote,
                json.dumps({"fp": _models_manifest_fp(cfg), "dests": dests}, indent=2)
                + "\n",
            )
            log(
                f"model-seed: skip fetch — {len(dests)} файл(ов) уже на диске "
                f"(civitai={civ_n}, huggingface={hf_n})"
            )
            return

    put_text(cfg, host, "/tmp/gpu-rent-civitai-jobs.json", json.dumps(jobs, indent=2))
    if cfg.civitai_api_token:
        put_text(
            cfg, host, "/tmp/gpu-rent-civitai.token", cfg.civitai_api_token + "\n", mode=0o600
        )
    if cfg.hf_token:
        put_text(cfg, host, "/tmp/gpu-rent-hf.token", cfg.hf_token + "\n", mode=0o600)
    log(
        f"model-seed: {len(jobs)} файл(ов) "
        f"(civitai={civ_n}, huggingface={hf_n}) — прогресс [N/{len(jobs)}]"
    )
    run_python(
        cfg,
        host,
        _pkg_text("civitai_fetch.py"),
        remote_path="/tmp/gpu-rent-civitai_fetch.py",
        timeout=7200,
        log=log,
    )
    put_text(
        cfg,
        host,
        cache_remote,
        json.dumps({"fp": _models_manifest_fp(cfg), "dests": dests}, indent=2) + "\n",
    )


def ensure_swarmui_running(cfg: Config, host: str, log: Log, restart: bool) -> None:
    active = run_ssh(cfg, host, "systemctl is-active swarmui 2>/dev/null || true", check=False).strip()
    if restart and active == "active":
        log("systemctl restart swarmui (новые extensions/autocomplete)")
        run_ssh(cfg, host, "sudo -n systemctl restart swarmui", timeout=120)
        return
    if active != "active":
        log("systemctl start swarmui")
        run_ssh(cfg, host, "sudo -n systemctl start swarmui", timeout=120)
    run_ssh(cfg, host, "sudo -n systemctl enable swarmui", check=False)


def provision_llm(cfg: Config, host: str, log: Log) -> None:
    from gpu_rent.llm_runtime import normalize_runtime, parse_ollama_models
    from gpu_rent.ssh_ops import run_script_sudo
    from gpu_rent.state import load_state, save_state

    runtime = normalize_runtime(cfg.llm_runtime)

    def _stop_units(*names: str) -> None:
        for name in names:
            run_ssh(
                cfg,
                host,
                f"sudo -n systemctl stop {name} 2>/dev/null; "
                f"sudo -n systemctl disable {name} 2>/dev/null || true",
                check=False,
            )

    # Drop LLM units that should not hold VRAM for this runtime.
    if runtime == "none":
        log("LLM: none — останавливаю gpu-rent-ollama / gpu-rent-llamacpp если были")
        _stop_units("gpu-rent-ollama", "gpu-rent-llamacpp")
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["llm_runtime"] = "none"
        st.notes.pop("llm_error", None)
        save_state(st)
        return
    if runtime == "ollama":
        _stop_units("gpu-rent-llamacpp")
        log("LLM: ставим/запускаем Ollama")
        run_script_sudo(
            cfg,
            host,
            _pkg_text("install_ollama.sh"),
            remote_path="/tmp/gpu-rent-install_ollama.sh",
            timeout=900,
            env=_remote_llm_env(cfg, *_OLLAMA_INSTALL_ENV),
            log=log,
        )
        entries = parse_ollama_models(cfg.ollama_models_manifest)
        defaults = [e.name for e in entries if e.default]
        if defaults:
            log(f"Ollama preferred: {defaults[0]}")
        names = [e.name for e in entries]
        if not names:
            log("ollama-models.yaml пуст — pull skip")
        else:
            # Fast path: all tags already present — skip upload/pull script.
            listed = run_ssh(
                cfg,
                host,
                "ollama list 2>/dev/null | awk 'NR>1 {print $1}' || true",
                check=False,
                timeout=30,
            )
            have = {ln.strip() for ln in listed.splitlines() if ln.strip()}
            missing = []
            for name in names:
                if name in have or (
                    ":" not in name and f"{name}:latest" in have
                ) or (
                    name.endswith(":latest") and name.rsplit(":", 1)[0] in have
                ):
                    continue
                missing.append(name)
            if not missing:
                log(f"ollama pull: skip — все {len(names)} уже есть")
            else:
                put_text(
                    cfg, host, "/tmp/gpu-rent-ollama-models.json", json.dumps(missing, indent=2)
                )
                log(f"Ollama: pull {len(missing)} из манифеста (нет: {len(missing)})")
                run_python(
                    cfg,
                    host,
                    _pkg_text("ollama_pull.py"),
                    remote_path="/tmp/gpu-rent-ollama_pull.py",
                    timeout=7200,
                    log=log,
                )
    else:
        raise CloudError(f"неизвестный LLM_RUNTIME={runtime!r}")
    st = load_state()
    st.notes = dict(st.notes or {})
    st.notes["llm_runtime"] = runtime
    st.notes.pop("llm_error", None)
    save_state(st)


def _try_arm_idle_killer(
    cfg: Config,
    host: str,
    log: Log,
    *,
    conn,
    server_id: str | None,
) -> bool:
    """Arm idle-killer; update state notes. Returns True if armed."""
    from gpu_rent.idle_killer import arm_idle_killer
    from gpu_rent.state import load_state, save_state

    if conn is None or not server_id:
        return False
    try:
        arm_idle_killer(cfg, host, conn, server_id, log)
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["idle_killer"] = "armed"
        st.notes.pop("idle_killer_error", None)
        save_state(st)
        return True
    except GpuRentError as exc:
        log(f"idle-killer: {exc}")
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["idle_killer"] = "failed"
        st.notes["idle_killer_error"] = str(exc)[:500]
        save_state(st)
        log(
            "⚠ idle-killer НЕ вооружён — GPU может тарифицироваться без авто-stop. "
            "Сделай gpu-rent stop или почини identity/application_credential_create."
        )
        return False


def provision_vm(
    cfg: Config,
    host: str,
    log: Log,
    *,
    conn=None,
    server_id: str | None = None,
    update: bool = True,
) -> None:
    from gpu_rent.llm_runtime import normalize_runtime
    from gpu_rent.state import load_state, save_state

    swarm = bool(getattr(cfg, "enable_swarmui", True))
    rt = normalize_runtime(cfg.llm_runtime)
    if not swarm and rt == "none":
        raise CloudError(
            "llm-only: нужен LLM_RUNTIME=ollama (или --ollama / --llm ollama)"
        )

    # Arm ASAP so mid-provision failures still leave auto-stop on the VM.
    armed = _try_arm_idle_killer(cfg, host, log, conn=conn, server_id=server_id)

    restart = bool(update)
    try:
        if swarm:
            try:
                if seed_extensions(cfg, host, log, update=update):
                    restart = True
            except GpuRentError as exc:
                log(f"extensions: {exc}")
                raise
            try:
                if seed_autocomplete(cfg, host, log):
                    restart = True
            except GpuRentError as exc:
                log(f"autocomplete: {exc}")
            seed_civitai(cfg, host, log)
            push_tree(cfg, host, cfg.local_models_dir, f"{DATA}/Models", log, models=True)
            push_tree(
                cfg, host, cfg.local_wildcards_dir, f"{DATA}/Data/Wildcards", log, models=False
            )
            push_tree(
                cfg,
                host,
                cfg.local_workflows_dir,
                f"{DATA}/CustomWorkflows",
                log,
                models=False,
            )
            if cfg.pull_output:
                pull_tree(cfg, host, f"{DATA}/Output", cfg.local_output_dir, log)
            # Avoid /Install wizard when backends already exist but Settings lack flag.
            if ensure_settings_is_installed(cfg, host, log):
                restart = True
        else:
            log("SwarmUI: skip (llm-only)")
            run_ssh(
                cfg,
                host,
                "sudo -n systemctl stop swarmui 2>/dev/null; "
                "sudo -n systemctl disable swarmui 2>/dev/null || true; "
                "echo llm-only | sudo -n tee /mnt/swarm_data/.gpu-rent-llm-only >/dev/null",
                check=False,
            )

        try:
            probe_gpu(cfg, host, log)
        except Exception as exc:
            log(f"GPU probe: {exc}")

        if swarm:
            ensure_swarmui_running(cfg, host, log, restart=restart)
            run_ssh(
                cfg,
                host,
                "sudo -n rm -f /mnt/swarm_data/.gpu-rent-llm-only",
                check=False,
            )

        try:
            provision_llm(cfg, host, log)
        except Exception as exc:
            st = load_state()
            st.notes = dict(st.notes or {})
            st.notes["llm_error"] = str(exc)[:500]
            st.notes["llm_runtime"] = "none"
            save_state(st)
            raise CloudError(f"LLM runtime: {exc}") from exc
    finally:
        # If first arm failed (SSH race), retry once after seeds.
        if not armed:
            _try_arm_idle_killer(cfg, host, log, conn=conn, server_id=server_id)

    log("Hold killer: gpu-rent hold   |  Стоп GPU: gpu-rent stop")