Add package data for GPU rent and update CLI documentation

- Added package data configuration for the 'gpu_rent' package in pyproject.toml. - Updated README.md to include usage instructions for Windows and Unix launchers. - Enhanced CLI documentation in cli.md to reflect new commands and their functionalities. - Revised setup.md to clarify installation steps and environment setup. - Improved error handling and command descriptions in the CLI implementation. - Added new functions for model version handling and flavor resolution in the codebase. - Updated state management to include additional properties for better tracking.
2026-08-21 03:06:51 +03:00
parent 167d07a733
commit 615cf81493
34 changed files with 2733 additions and 117 deletions
@@ -0,0 +1,269 @@
+"""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, 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,
+)
+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"
+
+
+def _pkg_text(name: str) -> str:
+    return files("gpu_rent.remote").joinpath(name).read_text(encoding="utf-8")
+
+
+def seed_extensions(cfg: Config, host: str, log: Log) -> bool:
+    repos = parse_extensions(cfg.extensions_manifest)
+    if not repos:
+        log("extensions.yaml пуст — стоковый SwarmUI")
+        return False
+    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))
+    if cfg.git_token:
+        put_text(cfg, host, "/tmp/gpu-rent-git.token", cfg.git_token + "\n", mode=0o600)
+    log(f"clone {len(jobs)} git-реп на data volume")
+    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 remote_exists(cfg, host, settings) or not applied:
+        fds = (
+            "DefaultUser:\n"
+            "    AutoComplete:\n"
+            f"        Source: {cfg.autocomplete_filename}\n"
+            "        EscapeParens: true\n"
+        )
+        put_text(cfg, host, settings, fds)
+        log(f"Settings.fds AutoComplete.Source = {cfg.autocomplete_filename}")
+    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_civitai(cfg: Config, host: str, log: Log) -> None:
+    entries = parse_models(cfg.models_manifest)
+    if not cfg.civitai_api_token:
+        log("Civitai-seed пропущен: нет CIVITAI_API_TOKEN — дефолт SwarmUI")
+        return
+    if not entries:
+        log("Civitai-seed пропущен: манифест пуст — дефолт SwarmUI")
+        return
+    jobs = []
+    for entry in entries:
+        vid = entry.version_id or (extract_version_id(entry.url) if entry.url else None)
+        if not vid:
+            log(f"пропуск {entry.kind}: нет version_id")
+            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 [],
+        }
+        jobs.append(
+            {
+                "dest": dest,
+                "url": _download_url(api_host, vid, info),
+                "sha256": sha,
+                "sidecars": {
+                    f"{stem}.civitai.json": civitai_json,
+                    f"{stem}.swarm.json": json.dumps(swarm, ensure_ascii=False, indent=2),
+                },
+            }
+        )
+    if not jobs:
+        log("Civitai-seed: ни одной скачиваемой строки")
+        return
+    if not any(e.kind == "checkpoint" for e in entries):
+        log("в манифесте нет checkpoint — генерация может не стартовать")
+    put_text(cfg, host, "/tmp/gpu-rent-civitai-jobs.json", json.dumps(jobs, indent=2))
+    put_text(cfg, host, "/tmp/gpu-rent-civitai.token", cfg.civitai_api_token + "\n", mode=0o600)
+    log(f"Civitai: качаю {len(jobs)} файл(ов) на VM")
+    run_python(
+        cfg,
+        host,
+        _pkg_text("civitai_fetch.py"),
+        remote_path="/tmp/gpu-rent-civitai_fetch.py",
+        timeout=7200,
+        log=log,
+    )
+
+
+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_vm(cfg: Config, host: str, log: Log) -> None:
+    restart = False
+    try:
+        if seed_extensions(cfg, host, log):
+            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)
+    ensure_swarmui_running(cfg, host, log, restart=restart)
+    log("SwarmUI слушает 127.0.0.1:7801 — gpu-rent tunnel")