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("нечего добавлять")
|
||||
+66
-1
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user