diff --git a/README.md b/README.md index 6db8696..955a345 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,7 @@ Unix: `./gpu-rent.sh …` (один раз `chmod +x gpu-rent.sh`). | `tunnel` / `open` | Снова UI / браузер | | `hold` / `status` | Пауза killer / состояние | | `stop` / `destroy --i-understand-data-loss` | Стоп GPU / + диски | -| `push` / `pull-output` / `seed-*` | Файлы ↔ VM | +| `push` / `pull-output` / `seed-*` / `capture` | Файлы ↔ VM; capture = ссылки в yaml | Полный список: [docs/cli.md](docs/cli.md). diff --git a/docs/cli.md b/docs/cli.md index 1600290..7aea212 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -74,6 +74,10 @@ gpu-rent up --yes --ollama | `gpu-rent seed-extensions` | Доклонировать/обновить git-репы; restart swarmui | | `gpu-rent push` / `push-models` | Локальные деревья → VM | | `gpu-rent pull-output` | VM `Output/` → `./Output` | +| `gpu-rent capture` / `capture all` | Инвентарь VM → merge **ссылок** в `models.yaml` + `extensions.yaml` (веса не качать) | +| `gpu-rent capture models` / `--kind lora` | Только модели (фильтр по типу) | +| `gpu-rent capture extensions` | Только git Extensions / DLNodes | +| `gpu-rent capture --dry-run` | Отчёт без записи файлов | | `gpu-rent resize-data --gb 400` | Data volume **только вверх** | | `gpu-rent watchdog install` / `uninstall` / `status` / `tick` | Локальный safety-net (см. ниже) | diff --git a/docs/decisions.md b/docs/decisions.md index 0e796b5..75dc030 100644 --- a/docs/decisions.md +++ b/docs/decisions.md @@ -10,7 +10,7 @@ | Civitai хост | Дефолт API `civitai.red` (полный каталог). `.com` — SFW-витрина, NSFW с неё часто 404. Ссылки `.com`/`.red`/`.green` в манифесте принимаем. 404 → один retry на второй хост. Один токен на оба домена | | Пул GPU | Перед `up`/`flavors` сканируем `SCAN_POOLS` (дефолт `ru-6,ru-7`). `ru-6` — мультизональный: ходим на `https://ru-6.cloud.api.selcloud.ru/compute/` тем же токеном (SDK-каталог часто знает только RC-пул). Собираем типы GPU из extra_specs и совпадения с `FLAVOR_PREFERENCE`. Автоматом `.env` не пишем — печатаем рекомендацию `OS_REGION_NAME` / `GPU_RENT_AZ` | | Манифест моделей | `/models.yaml`, типы: checkpoint / lora / vae / embedding / controlnet / upscaler. В git только `models.example.yaml` | -| Расширения | `/extensions.yaml`: git-репы `swarmui` → `src/Extensions`, `comfy` → DLNodes. Клон на первом bootstrap до старта UI. Пустой файл — стоковый SwarmUI | +| Расширения | `/extensions.yaml`: git-репы `swarmui` → `src/Extensions`, `comfy` → DLNodes. Поле `requires: none\|ollama\|llamacpp\|any-llm` фильтрует по `LLM_RUNTIME`. Пример: swarm-assistent с `requires: ollama` | | Autocomplete | До первого старта: скачать word-list в `Data/Autocompletions`, прописать `DefaultUser.AutoComplete.Source`. На каждом `up` сверить GitHub blob sha и обновить файл, если изменился. Дефолт: `tags/danbooru.csv` из a1111-sd-webui-tagcomplete (как в доке SwarmUI) | | Доступ | Браузер на туннеле; MCP переключается на облако, пока оно живо; HTTP API SwarmUI через тот же туннель | | Preemptible | По умолчанию всегда. Обычный сервер — только `--no-spot` | diff --git a/docs/extensions.md b/docs/extensions.md index e417b18..1f12be6 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -8,6 +8,8 @@ copy extensions.example.yaml extensions.yaml .\gpu-rent.ps1 up --yes # доклонировать на уже живой VM: .\gpu-rent.ps1 seed-extensions +# записать установленные на VM репы обратно в extensions.yaml (без clone локально): +.\gpu-rent.ps1 capture extensions ``` Срабатывает на каждом `up` (и `seed-extensions`): diff --git a/docs/models.md b/docs/models.md index 5d9cf88..cce3119 100644 --- a/docs/models.md +++ b/docs/models.md @@ -16,6 +16,19 @@ copy models.example.yaml models.yaml Или положи веса в `Models/Stable-Diffusion/` (и т.п.) — уедут на `up` / `gpu-rent push-models`. +### Capture с VM (ссылки only) + +Если на облачном диске уже есть модели (seed, UI Civitai, ручная заливка) — записать их обратно в локальный манифест **без скачивания весов**: + +```powershell +.\gpu-rent.ps1 capture models # merge Civitai url → models.yaml +.\gpu-rent.ps1 capture models --kind lora +.\gpu-rent.ps1 capture --dry-run # models + extensions, без записи +.\gpu-rent.ps1 capture all +``` + +Берёт `{stem}.civitai.json` или SHA256 → Civitai `by-hash`. Дубли `modelVersionId` не дублируются. Перед записью — `models.yaml.bak`. Неизвестные файлы (нет в Civitai) — только в отчёте. + Локальный SwarmUI на `7801` не зеркалируем. Только дерево `./Models` приложения. ## Локальная папка → сервер diff --git a/src/gpu_rent/capture.py b/src/gpu_rent/capture.py new file mode 100644 index 0000000..a39539b --- /dev/null +++ b/src/gpu_rent/capture.py @@ -0,0 +1,429 @@ +"""Capture VM inventory into local models.yaml / extensions.yaml (links only).""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Callable +from dataclasses import dataclass, field +from importlib.resources import files +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit, urlunsplit + +import yaml + +from gpu_rent.civitai import ( + civitai_model_url, + fetch_model_version_by_hash, + version_ids_from_payload, +) +from gpu_rent.config import Config +from gpu_rent.errors import CloudError, GpuRentError +from gpu_rent.manifests import ( + MODEL_TYPES, + extract_version_id, + parse_extensions, + parse_models, +) +from gpu_rent.ssh_ops import run_python, run_ssh + +Log = Callable[[str], None] +INVENTORY_REMOTE = "/tmp/gpu-rent-inventory.json" + + +def _pkg_text(name: str) -> str: + return files("gpu_rent.remote").joinpath(name).read_text(encoding="utf-8") + + +def strip_git_auth(url: str) -> str: + parts = urlsplit(url.strip()) + host = parts.hostname or "" + if parts.port: + host = f"{host}:{parts.port}" + return urlunsplit((parts.scheme, host, parts.path, parts.query, parts.fragment)) + + +@dataclass +class ModelCaptureItem: + kind: str + version_id: int + model_id: int + url: str + title: str = "" + rel: str = "" + + +@dataclass +class ExtCaptureItem: + kind: str + url: str + ref: str + directory: str + + +@dataclass +class CaptureReport: + models_new: list[ModelCaptureItem] = field(default_factory=list) + models_skip: list[str] = field(default_factory=list) + models_unknown: list[str] = field(default_factory=list) + ext_new: list[ExtCaptureItem] = field(default_factory=list) + ext_skip: list[str] = field(default_factory=list) + ext_unknown: list[str] = field(default_factory=list) + models_path: Path | None = None + extensions_path: Path | None = None + wrote_models: bool = False + wrote_extensions: bool = False + + +def fetch_inventory(cfg: Config, host: str, log: Log) -> dict[str, Any]: + log("capture: сканирую Models / Extensions / DLNodes на VM…") + run_python( + cfg, + host, + _pkg_text("scan_inventory.py"), + remote_path="/tmp/gpu-rent-scan_inventory.py", + timeout=3600, + log=log, + ) + raw = run_ssh(cfg, host, f"cat {INVENTORY_REMOTE}", timeout=60) + try: + data = json.loads(raw) + except json.JSONDecodeError as exc: + raise CloudError(f"не разобрать inventory JSON: {exc}") from exc + if not isinstance(data, dict): + raise CloudError("inventory: ожидался object") + return data + + +def resolve_model_item( + raw: dict, + *, + token: str, + api_host: str, + link_host: str, +) -> ModelCaptureItem | None: + kind = str(raw.get("kind") or "") + if kind not in MODEL_TYPES: + return None + rel = str(raw.get("rel") or raw.get("name") or "?") + version_id = raw.get("version_id") + model_id = raw.get("model_id") + title = Path(str(raw.get("name") or rel)).stem + + try: + vid = int(version_id) if version_id is not None else None + except (TypeError, ValueError): + vid = None + try: + mid = int(model_id) if model_id is not None else None + except (TypeError, ValueError): + mid = None + + if vid is not None and mid is not None: + return ModelCaptureItem( + kind=kind, + version_id=vid, + model_id=mid, + url=civitai_model_url(mid, vid, link_host), + title=title, + rel=rel, + ) + + sha = raw.get("sha256") + if not sha: + return None + try: + _host, version = fetch_model_version_by_hash(token or None, api_host, str(sha)) + except CloudError: + return None + vid2, mid2 = version_ids_from_payload(version) + if vid2 is None or mid2 is None: + return None + name = version.get("name") or title + return ModelCaptureItem( + kind=kind, + version_id=vid2, + model_id=mid2, + url=civitai_model_url(mid2, vid2, link_host), + title=str(name), + rel=rel, + ) + + +def _backup(path: Path) -> None: + if path.is_file() and path.stat().st_size > 0: + shutil.copy2(path, path.with_suffix(path.suffix + ".bak")) + + +def merge_models_yaml( + path: Path, + new_items: list[ModelCaptureItem], + *, + dry_run: bool, +) -> tuple[list[ModelCaptureItem], list[str]]: + """Return (actually_new, skip_msgs). Writes path unless dry_run.""" + existing = parse_models(path) if path.is_file() else [] + have: set[int] = set() + for e in existing: + if e.version_id is not None: + have.add(e.version_id) + elif e.url: + vid = extract_version_id(e.url) + if vid is not None: + have.add(vid) + + added: list[ModelCaptureItem] = [] + skipped: list[str] = [] + seen_new: set[int] = set() + for item in new_items: + if item.version_id in have or item.version_id in seen_new: + skipped.append(f"{item.kind} {item.title} modelVersionId={item.version_id}") + continue + seen_new.add(item.version_id) + added.append(item) + + if dry_run or not added: + return added, skipped + + # Rebuild full mapping: keep existing entries, append new urls. + data: dict[str, list[dict[str, str]]] = {k: [] for k in MODEL_TYPES} + if path.is_file(): + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + if isinstance(raw, dict): + for kind in MODEL_TYPES: + items = raw.get(kind) or [] + if isinstance(items, list): + for it in items: + if isinstance(it, dict) and (it.get("url") or it.get("version_id") not in (None, 0, "0")): + data[kind].append(dict(it)) + + for item in added: + data[item.kind].append({"url": item.url}) + + # Drop empty kinds for cleaner file + out = {k: v for k, v in data.items() if v} + _backup(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump(out, allow_unicode=True, default_flow_style=False, sort_keys=False), + encoding="utf-8", + ) + return added, skipped + + +def merge_extensions_yaml( + path: Path, + new_items: list[ExtCaptureItem], + *, + dry_run: bool, +) -> tuple[list[ExtCaptureItem], list[str]]: + existing = parse_extensions(path) if path.is_file() else [] + have_urls: set[tuple[str, str]] = set() + have_dirs: set[tuple[str, str]] = set() + for e in existing: + have_urls.add((e.kind, strip_git_auth(e.url).rstrip("/").lower())) + dirname = e.directory or "" + if dirname: + have_dirs.add((e.kind, dirname.lower())) + + added: list[ExtCaptureItem] = [] + skipped: list[str] = [] + for item in new_items: + key_url = (item.kind, strip_git_auth(item.url).rstrip("/").lower()) + key_dir = (item.kind, item.directory.lower()) + if key_url in have_urls or key_dir in have_dirs: + skipped.append(f"{item.kind} {item.directory} {item.url}") + continue + have_urls.add(key_url) + have_dirs.add(key_dir) + added.append(item) + + if dry_run or not added: + return added, skipped + + data: dict[str, list[dict[str, str]]] = {"swarmui": [], "comfy": []} + if path.is_file(): + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + if isinstance(raw, dict): + for kind in ("swarmui", "comfy"): + items = raw.get(kind) or [] + if isinstance(items, list): + for it in items: + if isinstance(it, dict) and it.get("url"): + data[kind].append(dict(it)) + + for item in added: + entry: dict[str, str] = { + "url": strip_git_auth(item.url), + "ref": item.ref, + "dir": item.directory, + } + data[item.kind].append(entry) + + out = {k: v for k, v in data.items() if v} + _backup(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump(out, allow_unicode=True, default_flow_style=False, sort_keys=False), + encoding="utf-8", + ) + return added, skipped + + +def capture_models( + cfg: Config, + host: str, + inventory: dict[str, Any] | None, + *, + dry_run: bool, + kind_filter: str | None, + log: Log, +) -> CaptureReport: + report = CaptureReport(models_path=cfg.models_manifest) + inv = inventory if inventory is not None else fetch_inventory(cfg, host, log) + raw_models = inv.get("models") or [] + if not isinstance(raw_models, list): + raise CloudError("inventory.models: ожидался list") + + resolved: list[ModelCaptureItem] = [] + link_host = cfg.civitai_api_host or "civitai.red" + for raw in raw_models: + if not isinstance(raw, dict): + continue + if kind_filter and str(raw.get("kind")) != kind_filter: + continue + item = resolve_model_item( + raw, + token=cfg.civitai_api_token, + api_host=cfg.civitai_api_host, + link_host=link_host, + ) + if item is None: + rel = str(raw.get("rel") or raw.get("name") or "?") + sha = str(raw.get("sha256") or "")[:12] + report.models_unknown.append(f"{rel} sha={sha}…") + continue + resolved.append(item) + + added, skipped = merge_models_yaml( + cfg.models_manifest, resolved, dry_run=dry_run + ) + report.models_new = added + report.models_skip = skipped + report.wrote_models = bool(added) and not dry_run + return report + + +def capture_extensions( + cfg: Config, + host: str, + inventory: dict[str, Any] | None, + *, + dry_run: bool, + log: Log, +) -> CaptureReport: + report = CaptureReport(extensions_path=cfg.extensions_manifest) + inv = inventory if inventory is not None else fetch_inventory(cfg, host, log) + raw_ext = inv.get("extensions") or [] + if not isinstance(raw_ext, list): + raise CloudError("inventory.extensions: ожидался list") + + resolved: list[ExtCaptureItem] = [] + for raw in raw_ext: + if not isinstance(raw, dict): + continue + if raw.get("unknown") or not raw.get("url"): + name = str(raw.get("dir") or "?") + reason = str(raw.get("reason") or "unknown") + report.ext_unknown.append(f"{raw.get('kind')}/{name} ({reason})") + continue + kind = str(raw.get("kind") or "") + if kind not in ("swarmui", "comfy"): + continue + resolved.append( + ExtCaptureItem( + kind=kind, + url=strip_git_auth(str(raw["url"])), + ref=str(raw.get("ref") or "main"), + directory=str(raw.get("dir") or "extension"), + ) + ) + + added, skipped = merge_extensions_yaml( + cfg.extensions_manifest, resolved, dry_run=dry_run + ) + report.ext_new = added + report.ext_skip = skipped + report.wrote_extensions = bool(added) and not dry_run + return report + + +def merge_reports(a: CaptureReport, b: CaptureReport) -> CaptureReport: + return CaptureReport( + models_new=a.models_new + b.models_new, + models_skip=a.models_skip + b.models_skip, + models_unknown=a.models_unknown + b.models_unknown, + ext_new=a.ext_new + b.ext_new, + ext_skip=a.ext_skip + b.ext_skip, + ext_unknown=a.ext_unknown + b.ext_unknown, + models_path=a.models_path or b.models_path, + extensions_path=a.extensions_path or b.extensions_path, + wrote_models=a.wrote_models or b.wrote_models, + wrote_extensions=a.wrote_extensions or b.wrote_extensions, + ) + + +def capture_all( + cfg: Config, + host: str, + *, + dry_run: bool, + kind_filter: str | None = None, + log: Log, +) -> CaptureReport: + inv = fetch_inventory(cfg, host, log) + m = capture_models( + cfg, host, inv, dry_run=dry_run, kind_filter=kind_filter, log=log + ) + e = capture_extensions(cfg, host, inv, dry_run=dry_run, log=log) + return merge_reports(m, e) + + +def print_report(report: CaptureReport, log: Log, *, dry_run: bool) -> None: + prefix = "[dry-run] " if dry_run else "" + log( + f"{prefix}models: +{len(report.models_new)} new, " + f"{len(report.models_skip)} already in yaml, " + f"{len(report.models_unknown)} unknown (no Civitai)" + ) + for item in report.models_new: + log(f" + {item.kind} {item.title} modelVersionId={item.version_id}") + for line in report.models_unknown: + log(f" ? {line}") + + log( + f"{prefix}extensions: +{len(report.ext_new)} new, " + f"{len(report.ext_skip)} already, " + f"{len(report.ext_unknown)} unknown" + ) + for item in report.ext_new: + log(f" + {item.kind} {item.directory} {item.url} @{item.ref}") + for line in report.ext_unknown: + log(f" ? {line}") + + if dry_run: + log("dry-run — файлы не записаны") + return + if report.wrote_models and report.models_path: + log(f"wrote {report.models_path} (backup {report.models_path.name}.bak)") + if report.wrote_extensions and report.extensions_path: + log( + f"wrote {report.extensions_path} " + f"(backup {report.extensions_path.name}.bak)" + ) + if not report.wrote_models and not report.wrote_extensions: + if report.models_new or report.ext_new: + pass + else: + log("нечего добавлять") diff --git a/src/gpu_rent/civitai.py b/src/gpu_rent/civitai.py index f0e0dae..f530a83 100644 --- a/src/gpu_rent/civitai.py +++ b/src/gpu_rent/civitai.py @@ -92,7 +92,8 @@ def fetch_model_version(token: str, host: str, version_id: int, timeout: float = url = f"https://{candidate}/api/v1/model-versions/{version_id}" try: with httpx.Client(timeout=timeout, follow_redirects=True) as client: - response = client.get(url, headers={"Authorization": f"Bearer {token}"}) + headers = {"Authorization": f"Bearer {token}"} if token else {} + response = client.get(url, headers=headers) except httpx.HTTPError as exc: last_error = str(exc) continue @@ -106,3 +107,67 @@ def fetch_model_version(token: str, host: str, version_id: int, timeout: float = if response.status_code not in {404, 400}: break raise CloudError(f"Civitai version {version_id}: {last_error} (хосты {', '.join(seen)})") + + +def fetch_model_version_by_hash( + token: str | None, + host: str, + sha256: str, + timeout: float = 30.0, +) -> tuple[str, dict]: + """GET /api/v1/model-versions/by-hash/{sha}; public, token optional for NSFW/region.""" + digest = sha256.strip().lower() + if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise CloudError(f"Civitai by-hash: неверный SHA256 ({sha256[:16]}…)") + first = _normalize_host(host) + order = [first, other_host(first)] + last_error = "нет ответа" + seen: set[str] = set() + headers = {"Authorization": f"Bearer {token}"} if token else {} + for candidate in order: + if candidate in seen or candidate not in ALLOWED_HOSTS: + continue + seen.add(candidate) + url = f"https://{candidate}/api/v1/model-versions/by-hash/{digest}" + try: + with httpx.Client(timeout=timeout, follow_redirects=True) as client: + response = client.get(url, headers=headers) + except httpx.HTTPError as exc: + last_error = str(exc) + continue + if response.status_code == 200: + data = response.json() + if isinstance(data, dict) and data.get("id") is not None: + return candidate, data + last_error = "пустой ответ" + continue + last_error = f"HTTP {response.status_code}" + if response.status_code not in {404, 400}: + break + raise CloudError(f"Civitai by-hash {digest[:12]}…: {last_error} (хосты {', '.join(seen)})") + + +def civitai_model_url(model_id: int, version_id: int, host: str = "civitai.red") -> str: + """Canonical manifest URL (links only — no download).""" + h = _normalize_host(host) + if h not in ALLOWED_HOSTS: + h = "civitai.red" + if h == "civitai.green": + h = "civitai.com" + return f"https://{h}/models/{int(model_id)}?modelVersionId={int(version_id)}" + + +def version_ids_from_payload(version: dict) -> tuple[int | None, int | None]: + """Extract (version_id, model_id) from a Civitai version JSON object.""" + try: + vid = int(version["id"]) if version.get("id") is not None else None + except (TypeError, ValueError, KeyError): + vid = None + mid = version.get("modelId") + if mid is None and isinstance(version.get("model"), dict): + mid = version["model"].get("id") + try: + model_id = int(mid) if mid is not None else None + except (TypeError, ValueError): + model_id = None + return vid, model_id diff --git a/src/gpu_rent/cli.py b/src/gpu_rent/cli.py index e4c39e8..f24ab30 100644 --- a/src/gpu_rent/cli.py +++ b/src/gpu_rent/cli.py @@ -683,6 +683,120 @@ def seed_extensions_cmd() -> None: _die(exc) +capture_app = typer.Typer( + help=( + "Снять с VM инвентарь → локальные манифесты (только ссылки, без весов). " + "Merge в models.yaml / extensions.yaml." + ), + no_args_is_help=False, +) +app.add_typer(capture_app, name="capture") + + +@capture_app.callback(invoke_without_command=True) +def capture_root( + ctx: typer.Context, + dry_run: bool = typer.Option(False, "--dry-run", help="Не писать yaml, только отчёт"), + kind: Optional[str] = typer.Option( + None, "--kind", help="Только models: checkpoint|lora|vae|…" + ), +) -> None: + """Без подкоманды — capture all.""" + if ctx.invoked_subcommand is not None: + return + try: + from gpu_rent.capture import capture_all, print_report + from gpu_rent.manifests import MODEL_TYPES + + if kind and kind not in MODEL_TYPES: + raise GpuRentError(f"--kind: жду один из {', '.join(MODEL_TYPES)}") + cfg, host = _live() + report = capture_all( + cfg, + host, + dry_run=dry_run, + kind_filter=kind, + log=lambda m: console.print(m), + ) + print_report(report, lambda m: console.print(m), dry_run=dry_run) + except GpuRentError as exc: + _die(exc) + + +@capture_app.command("models") +def capture_models_cmd( + dry_run: bool = typer.Option(False, "--dry-run"), + kind: Optional[str] = typer.Option( + None, "--kind", help="checkpoint|lora|vae|embedding|controlnet|upscaler|clip" + ), +) -> None: + """Models на VM → merge Civitai url в models.yaml.""" + try: + from gpu_rent.capture import capture_models, print_report + from gpu_rent.manifests import MODEL_TYPES + + if kind and kind not in MODEL_TYPES: + raise GpuRentError(f"--kind: жду один из {', '.join(MODEL_TYPES)}") + cfg, host = _live() + report = capture_models( + cfg, + host, + None, + dry_run=dry_run, + kind_filter=kind, + log=lambda m: console.print(m), + ) + print_report(report, lambda m: console.print(m), dry_run=dry_run) + except GpuRentError as exc: + _die(exc) + + +@capture_app.command("extensions") +def capture_extensions_cmd( + dry_run: bool = typer.Option(False, "--dry-run"), +) -> None: + """Extensions/DLNodes на VM → merge git url в extensions.yaml.""" + try: + from gpu_rent.capture import capture_extensions, print_report + + cfg, host = _live() + report = capture_extensions( + cfg, + host, + None, + dry_run=dry_run, + log=lambda m: console.print(m), + ) + print_report(report, lambda m: console.print(m), dry_run=dry_run) + except GpuRentError as exc: + _die(exc) + + +@capture_app.command("all") +def capture_all_cmd( + dry_run: bool = typer.Option(False, "--dry-run"), + kind: Optional[str] = typer.Option(None, "--kind", help="Фильтр только для models"), +) -> None: + """models + extensions.""" + try: + from gpu_rent.capture import capture_all, print_report + from gpu_rent.manifests import MODEL_TYPES + + if kind and kind not in MODEL_TYPES: + raise GpuRentError(f"--kind: жду один из {', '.join(MODEL_TYPES)}") + cfg, host = _live() + report = capture_all( + cfg, + host, + dry_run=dry_run, + kind_filter=kind, + log=lambda m: console.print(m), + ) + print_report(report, lambda m: console.print(m), dry_run=dry_run) + except GpuRentError as exc: + _die(exc) + + @app.command("resize-data") def resize_data(gb: int = typer.Option(..., "--gb", help="Новый размер data volume, GB (только вверх)")) -> None: """Cinder extend data volume + resize2fs на VM.""" diff --git a/src/gpu_rent/remote/scan_inventory.py b/src/gpu_rent/remote/scan_inventory.py new file mode 100644 index 0000000..4811a96 --- /dev/null +++ b/src/gpu_rent/remote/scan_inventory.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Scan VM Models/ + git Extensions/DLNodes. Stdlib only. Writes JSON inventory.""" +from __future__ import annotations + +import hashlib +import json +import subprocess +import sys +from pathlib import Path +from urllib.parse import urlsplit, urlunsplit + +DATA = Path("/mnt/swarm_data") +MODELS = DATA / "Models" +OUT = Path("/tmp/gpu-rent-inventory.json") + +WEIGHT_SUFFIXES = {".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".sft", ".onnx"} + +# SwarmUI folder name → models.yaml kind +FOLDER_TO_KIND = { + "Stable-Diffusion": "checkpoint", + "Lora": "lora", + "VAE": "vae", + "Embeddings": "embedding", + "controlnet": "controlnet", + "upscale_models": "upscaler", + "clip": "clip", +} + +EXT_ROOTS = ( + ("swarmui", DATA / "Extensions"), + ("comfy", DATA / "DLNodes"), +) + + +def strip_auth(url: str) -> str: + parts = urlsplit(url) + host = parts.hostname or "" + if parts.port: + host = f"{host}:{parts.port}" + return urlunsplit((parts.scheme, host, parts.path, parts.query, parts.fragment)) + + +def sha256_file(path: Path, chunk: int = 1024 * 1024) -> str: + h = hashlib.sha256() + with path.open("rb") as fh: + while True: + block = fh.read(chunk) + if not block: + break + h.update(block) + return h.hexdigest() + + +def read_sidecar_ids(weight: Path) -> tuple[int | None, int | None]: + """Return (version_id, model_id) from {stem}.civitai.json if present.""" + sidecar = weight.parent / f"{weight.stem}.civitai.json" + if not sidecar.is_file(): + return None, None + try: + data = json.loads(sidecar.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None, None + if not isinstance(data, dict): + return None, None + vid = data.get("id") + mid = data.get("modelId") + try: + version_id = int(vid) if vid is not None else None + except (TypeError, ValueError): + version_id = None + try: + model_id = int(mid) if mid is not None else None + except (TypeError, ValueError): + model_id = None + return version_id, model_id + + +def scan_models() -> list[dict]: + items: list[dict] = [] + if not MODELS.is_dir(): + return items + for folder, kind in FOLDER_TO_KIND.items(): + root = MODELS / folder + if not root.is_dir(): + continue + for path in sorted(root.rglob("*")): + if not path.is_file(): + continue + if path.suffix.lower() not in WEIGHT_SUFFIXES: + continue + if path.name.startswith("."): + continue + rel = path.relative_to(MODELS).as_posix() + version_id, model_id = read_sidecar_ids(path) + try: + digest = sha256_file(path) + except OSError as exc: + items.append( + { + "kind": kind, + "rel": rel, + "name": path.name, + "sha256": None, + "version_id": version_id, + "model_id": model_id, + "error": str(exc), + } + ) + continue + items.append( + { + "kind": kind, + "rel": rel, + "name": path.name, + "sha256": digest, + "version_id": version_id, + "model_id": model_id, + } + ) + return items + + +def git_out(args: list[str], cwd: Path) -> str | None: + try: + return subprocess.check_output( + ["git", "-C", str(cwd), *args], + text=True, + stderr=subprocess.DEVNULL, + ).strip() + except (subprocess.CalledProcessError, FileNotFoundError, OSError): + return None + + +def scan_extensions() -> list[dict]: + items: list[dict] = [] + for kind, root in EXT_ROOTS: + if not root.is_dir(): + continue + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + if not (child / ".git").exists(): + items.append( + { + "kind": kind, + "dir": child.name, + "url": None, + "ref": None, + "unknown": True, + "reason": "no .git", + } + ) + continue + origin = git_out(["remote", "get-url", "origin"], child) + if not origin: + items.append( + { + "kind": kind, + "dir": child.name, + "url": None, + "ref": None, + "unknown": True, + "reason": "no origin", + } + ) + continue + url = strip_auth(origin) + branch = git_out(["rev-parse", "--abbrev-ref", "HEAD"], child) + if not branch or branch == "HEAD": + sha = git_out(["rev-parse", "--short", "HEAD"], child) + ref = sha or "main" + else: + ref = branch + items.append( + { + "kind": kind, + "dir": child.name, + "url": url, + "ref": ref, + "unknown": False, + } + ) + return items + + +def main() -> int: + payload = { + "models": scan_models(), + "extensions": scan_extensions(), + } + OUT.write_text(json.dumps(payload, indent=2), encoding="utf-8") + # One-line marker for local parsers; full JSON is in OUT. + print(f"inventory ok models={len(payload['models'])} extensions={len(payload['extensions'])}") + print(f"INVENTORY_PATH={OUT}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_capture_merge.py b/tests/test_capture_merge.py new file mode 100644 index 0000000..1d1be1a --- /dev/null +++ b/tests/test_capture_merge.py @@ -0,0 +1,111 @@ +"""Unit tests for capture merge / URL builders (no SSH).""" + +from pathlib import Path + +from gpu_rent.capture import ( + ExtCaptureItem, + ModelCaptureItem, + merge_extensions_yaml, + merge_models_yaml, + resolve_model_item, + strip_git_auth, +) +from gpu_rent.civitai import civitai_model_url, version_ids_from_payload +from gpu_rent.manifests import parse_extensions, parse_models + + +def test_civitai_model_url(): + assert ( + civitai_model_url(2786499, 3231611, "civitai.red") + == "https://civitai.red/models/2786499?modelVersionId=3231611" + ) + + +def test_version_ids_from_payload(): + vid, mid = version_ids_from_payload({"id": 9, "modelId": 1}) + assert vid == 9 and mid == 1 + vid2, mid2 = version_ids_from_payload({"id": 3, "model": {"id": 7}}) + assert vid2 == 3 and mid2 == 7 + + +def test_resolve_from_sidecar(): + item = resolve_model_item( + { + "kind": "lora", + "rel": "Lora/foo.safetensors", + "name": "foo.safetensors", + "version_id": 3107521, + "model_id": 2187487, + "sha256": "a" * 64, + }, + token="", + api_host="civitai.red", + link_host="civitai.red", + ) + assert item is not None + assert item.version_id == 3107521 + assert "modelVersionId=3107521" in item.url + + +def test_merge_models_dedupe(tmp_path: Path): + path = tmp_path / "models.yaml" + path.write_text( + "lora:\n - url: https://civitai.red/models/1?modelVersionId=100\n", + encoding="utf-8", + ) + items = [ + ModelCaptureItem("lora", 100, 1, "https://civitai.red/models/1?modelVersionId=100", "old"), + ModelCaptureItem("lora", 200, 2, "https://civitai.red/models/2?modelVersionId=200", "new"), + ModelCaptureItem("checkpoint", 300, 3, "https://civitai.red/models/3?modelVersionId=300", "ckpt"), + ] + added, skipped = merge_models_yaml(path, items, dry_run=False) + assert len(added) == 2 + assert any("100" in s for s in skipped) + entries = parse_models(path) + vids = {e.version_id for e in entries} + assert vids == {100, 200, 300} + assert (tmp_path / "models.yaml.bak").is_file() + + +def test_merge_models_dry_run(tmp_path: Path): + path = tmp_path / "models.yaml" + path.write_text("lora: []\n", encoding="utf-8") + before = path.read_text(encoding="utf-8") + items = [ + ModelCaptureItem("lora", 1, 1, "https://civitai.red/models/1?modelVersionId=1", "x"), + ] + added, _ = merge_models_yaml(path, items, dry_run=True) + assert len(added) == 1 + assert path.read_text(encoding="utf-8") == before + + +def test_merge_extensions_dedupe(tmp_path: Path): + path = tmp_path / "extensions.yaml" + path.write_text( + "swarmui:\n - url: https://github.com/org/A.git\n ref: main\n dir: A\n", + encoding="utf-8", + ) + items = [ + ExtCaptureItem("swarmui", "https://github.com/org/A.git", "main", "A"), + ExtCaptureItem("comfy", "https://github.com/org/B.git", "v1", "B"), + ExtCaptureItem( + "swarmui", + "https://x-access-token:SECRET@github.com/org/C.git", + "main", + "C", + ), + ] + added, skipped = merge_extensions_yaml(path, items, dry_run=False) + assert len(added) == 2 + assert len(skipped) == 1 + repos = parse_extensions(path) + urls = [r.url for r in repos] + assert "https://github.com/org/B.git" in urls + assert all("SECRET" not in u for u in urls) + + +def test_strip_git_auth(): + assert ( + strip_git_auth("https://x-access-token:tok@github.com/org/r.git") + == "https://github.com/org/r.git" + )