Enhance CLI and documentation for capturing VM inventory
- Introduced new `capture` commands in the CLI to allow users to merge VM inventory into local manifests without downloading weights. - Updated `README.md` and `cli.md` to include detailed instructions for the new capture functionality, including options for models and extensions. - Enhanced `decisions.md` to clarify the role of captured links in the manifest files. - Improved `extensions.md` to document the process of capturing installed extensions back to the local configuration. - Added new functions in `civitai.py` to support fetching model versions by hash and generating canonical URLs for models.
This commit is contained in:
Leonid Pershin
@@ -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("нечего добавлять")
|
||||
Reference in New Issue
Block a user