Leonid Pershin
2ccb03f7d2
- Introduced GPU probing functionality to gather and store GPU specifications in `/mnt/swarm_data/.gpu-rent-gpu.json`, aiding in performance tuning. - Updated `install_ollama.sh` and `install_llamacpp.sh` to utilize GPU information for configuring optimal runtime parameters. - Enhanced `provision.py` to include GPU probing and performance tuning logic, ensuring better resource allocation for LLM operations. - Improved documentation in `decisions.md`, `llm.md`, and `swarmui.md` to reflect changes in GPU handling and performance tuning processes. - Added new tests to validate the GPU probing and model resolution logic, ensuring robustness in handling various GPU configurations.
595 lines
20 KiB
Python
595 lines
20 KiB
Python
"""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,
|
|
fetch_model_version_by_hash,
|
|
fetch_model_versions_by_hashes,
|
|
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 ResolveOutcome:
|
|
"""ok → item; unknown → not on Civitai; api_error → network/HTTP (retry later)."""
|
|
|
|
item: ModelCaptureItem | None = None
|
|
status: str = "unknown" # ok | unknown | api_error
|
|
detail: 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)
|
|
models_api_errors: list[str] = field(default_factory=list)
|
|
ext_new: list[ExtCaptureItem] = field(default_factory=list)
|
|
ext_updated: 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,
|
|
) -> ResolveOutcome:
|
|
kind = str(raw.get("kind") or "")
|
|
if kind not in MODEL_TYPES:
|
|
return ResolveOutcome(status="unknown", detail="bad kind")
|
|
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
|
|
|
|
def _ok(v: int, m: int, name: str = title) -> ResolveOutcome:
|
|
return ResolveOutcome(
|
|
item=ModelCaptureItem(
|
|
kind=kind,
|
|
version_id=v,
|
|
model_id=m,
|
|
url=civitai_model_url(m, v, link_host),
|
|
title=name,
|
|
rel=rel,
|
|
),
|
|
status="ok",
|
|
)
|
|
|
|
if vid is not None and mid is not None:
|
|
return _ok(vid, mid)
|
|
|
|
# Partial sidecar: have version id → GET /model-versions/{id} for modelId.
|
|
if vid is not None and mid is None:
|
|
try:
|
|
_h, version = fetch_model_version(token or "", api_host, vid)
|
|
except CloudError as exc:
|
|
msg = str(exc)
|
|
if "404" in msg or "пустой files" in msg:
|
|
return ResolveOutcome(
|
|
status="unknown",
|
|
detail=f"{rel} version_id={vid} ({msg})",
|
|
)
|
|
return ResolveOutcome(
|
|
status="api_error",
|
|
detail=f"{rel} version_id={vid} {msg}",
|
|
)
|
|
vid2, mid2 = version_ids_from_payload(version)
|
|
if vid2 is not None and mid2 is not None:
|
|
name = str(version.get("name") or title)
|
|
return _ok(vid2, mid2, name)
|
|
return ResolveOutcome(
|
|
status="unknown",
|
|
detail=f"{rel} version_id={vid} (нет modelId в ответе)",
|
|
)
|
|
|
|
sha = raw.get("sha256")
|
|
if not sha:
|
|
return ResolveOutcome(
|
|
status="unknown",
|
|
detail=f"{rel} (нет sidecar и нет sha256)",
|
|
)
|
|
try:
|
|
_host, version = fetch_model_version_by_hash(token or None, api_host, str(sha))
|
|
except CloudError as exc:
|
|
msg = str(exc)
|
|
# 404 = genuinely not on Civitai; other → api_error
|
|
if "HTTP 404" in msg or msg.rstrip().endswith("404"):
|
|
return ResolveOutcome(
|
|
status="unknown",
|
|
detail=f"{rel} sha={str(sha)[:12]}…",
|
|
)
|
|
return ResolveOutcome(
|
|
status="api_error",
|
|
detail=f"{rel} sha={str(sha)[:12]}… {msg}",
|
|
)
|
|
vid2, mid2 = version_ids_from_payload(version)
|
|
if vid2 is None or mid2 is None:
|
|
return ResolveOutcome(
|
|
status="unknown",
|
|
detail=f"{rel} sha={str(sha)[:12]}… (пустой payload)",
|
|
)
|
|
name = str(version.get("name") or title)
|
|
return _ok(vid2, mid2, name)
|
|
|
|
|
|
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 _keep_model_entry(it: dict) -> bool:
|
|
if not isinstance(it, dict):
|
|
return False
|
|
url = it.get("url")
|
|
if url:
|
|
return True
|
|
vid = it.get("version_id")
|
|
return vid not in (None, "", 0, "0")
|
|
|
|
|
|
def merge_models_yaml(
|
|
path: Path,
|
|
new_items: list[ModelCaptureItem],
|
|
*,
|
|
dry_run: bool,
|
|
) -> tuple[list[ModelCaptureItem], list[str]]:
|
|
"""Return (actually_new, skip_msgs). Dedupe by (kind, version_id)."""
|
|
existing = parse_models(path) if path.is_file() else []
|
|
have: set[tuple[str, int]] = set()
|
|
for e in existing:
|
|
vid = e.version_id
|
|
if vid is None and e.url:
|
|
vid = extract_version_id(e.url)
|
|
if vid is not None:
|
|
have.add((e.kind, vid))
|
|
|
|
added: list[ModelCaptureItem] = []
|
|
skipped: list[str] = []
|
|
seen_new: set[tuple[str, int]] = set()
|
|
for item in new_items:
|
|
key = (item.kind, item.version_id)
|
|
if key in have or key in seen_new:
|
|
skipped.append(f"{item.kind} {item.title} modelVersionId={item.version_id}")
|
|
continue
|
|
seen_new.add(key)
|
|
added.append(item)
|
|
|
|
if dry_run or not added:
|
|
return added, skipped
|
|
|
|
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 _keep_model_entry(it):
|
|
data[kind].append(dict(it))
|
|
|
|
for item in added:
|
|
data[item.kind].append({"url": item.url})
|
|
|
|
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[ExtCaptureItem], list[str]]:
|
|
"""Return (added, updated, skipped). Same dir + different URL → update url/ref."""
|
|
data: dict[str, list[dict[str, Any]]] = {"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))
|
|
|
|
def _dir_of(it: dict) -> str:
|
|
d = it.get("dir")
|
|
if d:
|
|
return str(d).lower()
|
|
url = strip_git_auth(str(it.get("url") or ""))
|
|
name = url.rstrip("/").rsplit("/", 1)[-1]
|
|
if name.endswith(".git"):
|
|
name = name[:-4]
|
|
return name.lower()
|
|
|
|
added: list[ExtCaptureItem] = []
|
|
updated: list[ExtCaptureItem] = []
|
|
skipped: list[str] = []
|
|
|
|
for item in new_items:
|
|
clean = strip_git_auth(item.url)
|
|
key_url = (item.kind, clean.rstrip("/").lower())
|
|
key_dir = (item.kind, item.directory.lower())
|
|
bucket = data.setdefault(item.kind, [])
|
|
|
|
matched_url = False
|
|
matched_dir_idx: int | None = None
|
|
for idx, it in enumerate(bucket):
|
|
it_url = strip_git_auth(str(it.get("url") or "")).rstrip("/").lower()
|
|
if (item.kind, it_url) == key_url:
|
|
matched_url = True
|
|
break
|
|
if (item.kind, _dir_of(it)) == key_dir:
|
|
matched_dir_idx = idx
|
|
|
|
if matched_url:
|
|
skipped.append(f"{item.kind} {item.directory} {clean}")
|
|
continue
|
|
|
|
if matched_dir_idx is not None:
|
|
old = bucket[matched_dir_idx]
|
|
old_url = strip_git_auth(str(old.get("url") or ""))
|
|
old_ref = str(old.get("ref") or "main")
|
|
if old_url.rstrip("/").lower() == clean.rstrip("/").lower() and old_ref == item.ref:
|
|
skipped.append(f"{item.kind} {item.directory} {clean}")
|
|
continue
|
|
old["url"] = clean
|
|
old["ref"] = item.ref
|
|
if not old.get("dir"):
|
|
old["dir"] = item.directory
|
|
updated.append(item)
|
|
continue
|
|
|
|
bucket.append({"url": clean, "ref": item.ref, "dir": item.directory})
|
|
added.append(item)
|
|
|
|
if dry_run or (not added and not updated):
|
|
return added, updated, skipped
|
|
|
|
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, updated, 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] = []
|
|
need_hash: list[dict] = []
|
|
link_host = cfg.civitai_api_host or "civitai.red"
|
|
token = cfg.civitai_api_token
|
|
api_host = cfg.civitai_api_host
|
|
|
|
for raw in raw_models:
|
|
if not isinstance(raw, dict):
|
|
continue
|
|
if kind_filter and str(raw.get("kind")) != kind_filter:
|
|
continue
|
|
vid = raw.get("version_id")
|
|
mid = raw.get("model_id")
|
|
try:
|
|
vid_i = int(vid) if vid is not None else None
|
|
except (TypeError, ValueError):
|
|
vid_i = None
|
|
try:
|
|
mid_i = int(mid) if mid is not None else None
|
|
except (TypeError, ValueError):
|
|
mid_i = None
|
|
sha = raw.get("sha256")
|
|
|
|
# Full sidecar / ids → no API. Partial vid → GET version. Else batch by-hash.
|
|
if vid_i is not None and mid_i is not None:
|
|
outcome = resolve_model_item(
|
|
raw, token=token, api_host=api_host, link_host=link_host
|
|
)
|
|
elif vid_i is not None:
|
|
outcome = resolve_model_item(
|
|
raw, token=token, api_host=api_host, link_host=link_host
|
|
)
|
|
elif sha:
|
|
need_hash.append(raw)
|
|
continue
|
|
else:
|
|
outcome = resolve_model_item(
|
|
raw, token=token, api_host=api_host, link_host=link_host
|
|
)
|
|
|
|
if outcome.status == "ok" and outcome.item is not None:
|
|
resolved.append(outcome.item)
|
|
elif outcome.status == "api_error":
|
|
report.models_api_errors.append(outcome.detail or "?")
|
|
else:
|
|
report.models_unknown.append(outcome.detail or str(raw.get("rel") or "?"))
|
|
|
|
if need_hash:
|
|
digests = [str(r["sha256"]).strip().lower() for r in need_hash if r.get("sha256")]
|
|
log(f"capture: by-hash batch {len(digests)} файл(ов)…")
|
|
try:
|
|
by_hash = fetch_model_versions_by_hashes(token or None, api_host, digests)
|
|
except CloudError as exc:
|
|
report.models_api_errors.append(f"by-hash batch: {exc}")
|
|
by_hash = {}
|
|
for raw in need_hash:
|
|
rel = str(raw.get("rel") or "?")
|
|
sha = str(raw.get("sha256") or "")[:12]
|
|
report.models_api_errors.append(f"{rel} sha={sha}… (batch failed)")
|
|
need_hash = []
|
|
|
|
for raw in need_hash:
|
|
sha = str(raw.get("sha256") or "").strip().lower()
|
|
rel = str(raw.get("rel") or raw.get("name") or "?")
|
|
version = by_hash.get(sha)
|
|
if not version:
|
|
report.models_unknown.append(f"{rel} sha={sha[:12]}…")
|
|
continue
|
|
vid2, mid2 = version_ids_from_payload(version)
|
|
if vid2 is None or mid2 is None:
|
|
report.models_unknown.append(f"{rel} sha={sha[:12]}… (пустой payload)")
|
|
continue
|
|
kind = str(raw.get("kind") or "")
|
|
if kind not in MODEL_TYPES:
|
|
continue
|
|
title = str(version.get("name") or Path(str(raw.get("name") or rel)).stem)
|
|
resolved.append(
|
|
ModelCaptureItem(
|
|
kind=kind,
|
|
version_id=vid2,
|
|
model_id=mid2,
|
|
url=civitai_model_url(mid2, vid2, link_host),
|
|
title=title,
|
|
rel=rel,
|
|
)
|
|
)
|
|
|
|
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, updated, skipped = merge_extensions_yaml(
|
|
cfg.extensions_manifest, resolved, dry_run=dry_run
|
|
)
|
|
report.ext_new = added
|
|
report.ext_updated = updated
|
|
report.ext_skip = skipped
|
|
report.wrote_extensions = (bool(added) or bool(updated)) 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,
|
|
models_api_errors=a.models_api_errors + b.models_api_errors,
|
|
ext_new=a.ext_new + b.ext_new,
|
|
ext_updated=a.ext_updated + b.ext_updated,
|
|
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,
|
|
show_models: bool = True,
|
|
show_extensions: bool = True,
|
|
) -> None:
|
|
prefix = "[dry-run] " if dry_run else ""
|
|
if show_models:
|
|
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), "
|
|
f"{len(report.models_api_errors)} api errors"
|
|
)
|
|
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}")
|
|
for line in report.models_api_errors:
|
|
log(f" ! {line}")
|
|
|
|
if show_extensions:
|
|
log(
|
|
f"{prefix}extensions: +{len(report.ext_new)} new, "
|
|
f"~{len(report.ext_updated)} updated, "
|
|
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 item in report.ext_updated:
|
|
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:
|
|
log("нечего добавлять")
|