gpu-rent/src/gpu_rent/manifests.py

"""Parse models.yaml / extensions.yaml. version_id 0 is a placeholder."""

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlparse

import yaml

from gpu_rent.errors import ConfigError

MODEL_TYPES = (
    "checkpoint",
    "lora",
    "vae",
    "embedding",
    "controlnet",
    "upscaler",
    "clip",
)

MODEL_DIRS = {
    "checkpoint": "Stable-Diffusion",
    "lora": "Lora",
    "vae": "VAE",
    "embedding": "Embeddings",
    "controlnet": "controlnet",
    "upscaler": "upscale_models",
    "clip": "clip",
}

_VERSION_QS = re.compile(r"modelVersionId=(\d+)", re.I)
_VERSION_PATH = re.compile(r"/model-versions/(\d+)", re.I)
_DOWNLOAD_PATH = re.compile(r"/api/download/models/(\d+)", re.I)
_SHA_REF = re.compile(r"^[0-9a-fA-F]{40}$")


@dataclass
class ModelEntry:
    kind: str
    version_id: int | None
    url: str | None


VALID_REQUIRES = frozenset({"none", "ollama", "any-llm"})


@dataclass
class GitRepo:
    kind: str
    url: str
    ref: str
    directory: str | None
    requires: str = "none"


def _load_yaml(path: Path) -> Any:
    if not path.is_file():
        return None
    text = path.read_text(encoding="utf-8")
    if not text.strip():
        return {}
    try:
        return yaml.safe_load(text)
    except yaml.YAMLError as exc:
        raise ConfigError(f"Не разобрать YAML {path}: {exc}") from exc


def parse_models(path: Path) -> list[ModelEntry]:
    data = _load_yaml(path)
    if data is None:
        return []
    if data == {} or data is None:
        return []
    if not isinstance(data, dict):
        raise ConfigError(f"{path}: корень должен быть mapping типов моделей")
    entries: list[ModelEntry] = []
    for kind, items in data.items():
        if kind not in MODEL_TYPES:
            continue
        if not items:
            continue
        if not isinstance(items, list):
            raise ConfigError(f"{path}: {kind} должен быть списком")
        for item in items:
            if not isinstance(item, dict):
                raise ConfigError(f"{path}: элемент {kind} — объект с version_id или url")
            vid = item.get("version_id")
            url = item.get("url")
            if vid in (0, "0", None) and not url:
                continue
            version_id = int(vid) if vid not in (None, "", 0, "0") else None
            url_s = str(url) if url else None
            if version_id is None and url_s:
                version_id = extract_version_id(url_s)
            entries.append(ModelEntry(kind=kind, version_id=version_id, url=url_s))
    return entries


def normalize_requires(value: object | None) -> str:
    raw = str(value or "none").strip().lower().replace("_", "-")
    if raw in {"", "none", "always", "any"}:
        return "none"
    if raw in VALID_REQUIRES:
        return raw
    raise ConfigError(
        f"requires={value!r}: жду none|ollama|any-llm"
    )


def repo_matches_runtime(repo: GitRepo, llm_runtime: str) -> bool:
    """Whether this extension should be cloned for the active LLM_RUNTIME."""
    req = normalize_requires(repo.requires)
    runtime = (llm_runtime or "none").strip().lower()
    if req == "none":
        return True
    if req == "any-llm":
        return runtime == "ollama"
    return runtime == req


def parse_extensions(path: Path) -> list[GitRepo]:
    data = _load_yaml(path)
    if not data:
        return []
    if not isinstance(data, dict):
        raise ConfigError(f"{path}: корень swarmui: / comfy: / assistent:")
    repos: list[GitRepo] = []
    for kind in ("swarmui", "comfy", "assistent"):
        items = data.get(kind) or []
        if not items:
            continue
        if not isinstance(items, list):
            raise ConfigError(f"{path}: {kind} должен быть списком")
        for item in items:
            if not isinstance(item, dict) or not item.get("url"):
                raise ConfigError(f"{path}: у {kind} нужен url")
            try:
                requires = normalize_requires(item.get("requires"))
            except ConfigError as exc:
                raise ConfigError(f"{path}: {exc}") from exc
            repos.append(
                GitRepo(
                    kind=kind,
                    url=str(item["url"]).strip(),
                    ref=str(item.get("ref") or "main"),
                    directory=str(item["dir"]) if item.get("dir") else None,
                    requires=requires,
                )
            )
    return repos


def extract_version_id(url: str) -> int | None:
    text = url.strip()
    for rx in (_VERSION_QS, _VERSION_PATH, _DOWNLOAD_PATH):
        match = rx.search(text)
        if match:
            return int(match.group(1))
    parsed = urlparse(text)
    ids = parse_qs(parsed.query).get("modelVersionId") or parse_qs(parsed.query).get("modelversionid")
    if ids:
        try:
            return int(ids[0])
        except ValueError:
            return None
    return None


def repo_dirname(repo: GitRepo) -> str:
    if repo.directory:
        return repo.directory
    name = repo.url.rstrip("/").rsplit("/", 1)[-1]
    if name.endswith(".git"):
        name = name[:-4]
    return name or "extension"


def is_commit_sha(ref: str) -> bool:
    return bool(_SHA_REF.match(ref.strip()))


def remote_root_for(repo: GitRepo) -> str:
    if repo.kind == "swarmui":
        base = "/mnt/swarm_data/Extensions"
    elif repo.kind == "assistent":
        base = "/mnt/swarm_data/Assistent/extensions"
    else:
        base = "/mnt/swarm_data/DLNodes"
    return f"{base}/{repo_dirname(repo)}"