gpu-rent/src/gpu_rent/inventory.py

"""Flavor preference matching and volume-type pick for the AZ."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass
class FlavorInfo:
    id: str
    name: str
    vcpus: int | None
    ram_mb: int | None
    disabled: bool
    extra: dict[str, Any]
    label: str | None = None


def _name(obj: Any) -> str:
    return (getattr(obj, "name", None) or "").strip()


def _id(obj: Any) -> str:
    return str(getattr(obj, "id", "") or "")


def is_disabled(flavor: Any) -> bool:
    if getattr(flavor, "is_disabled", False):
        return True
    extra = extra_specs(flavor)
    flag = extra.get("OS-FLV-DISABLED:disabled") or extra.get("disabled")
    if flag in (True, "True", "true", "1"):
        return True
    return False


def extra_specs(flavor: Any) -> dict[str, Any]:
    extra = getattr(flavor, "extra_specs", None)
    if isinstance(extra, dict):
        return extra
    blob = getattr(flavor, "get", None)
    if callable(blob):
        got = flavor.get("extra_specs")
        if isinstance(got, dict):
            return got
    return {}


def looks_like_gpu(flavor: Any) -> bool:
    name = _name(flavor).lower()
    extra = extra_specs(flavor)
    blob = " ".join(f"{k}={v}" for k, v in extra.items()).lower()
    hay = f"{name} {blob}"
    needles = (
        "gpu",
        "4090",
        "a5000",
        "a100",
        "a6000",
        "l40",
        "h100",
        "h200",
        "rtx",
        "tesla",
    )
    return any(n in hay for n in needles)


def match_label(label: str, flavor: Any) -> bool:
    """Match FLAVOR_PREFERENCE tokens to a live flavor name/extra specs."""
    name = _name(flavor).lower()
    extra = extra_specs(flavor)
    hay = name + " " + " ".join(str(v).lower() for v in extra.values())
    token = label.strip().lower()

    if token in {"4090-24", "4090_24", "rtx4090-24"}:
        return "4090" in hay and "48" not in hay
    if token in {"4090-48", "4090_48", "rtx4090-48"}:
        return "4090" in hay and "48" in hay
    if token in {"a5000", "rtx-a5000"}:
        return "a5000" in hay or "rtx a5000" in hay
    if token in {"a100-40", "a100_40"}:
        return "a100" in hay and "80" not in hay
    if token in {"a100-80", "a100_80"}:
        return "a100" in hay and "80" in hay
    return token.replace("_", "-") in hay or token.replace("-", " ") in hay


def flavor_info(flavor: Any, label: str | None = None) -> FlavorInfo:
    ram = getattr(flavor, "ram", None)
    vcpus = getattr(flavor, "vcpus", None)
    return FlavorInfo(
        id=_id(flavor),
        name=_name(flavor) or _id(flavor),
        vcpus=int(vcpus) if vcpus is not None else None,
        ram_mb=int(ram) if ram is not None else None,
        disabled=is_disabled(flavor),
        extra=extra_specs(flavor),
        label=label,
    )


def rank_flavors(flavors: list[Any], preference: tuple[str, ...]) -> list[FlavorInfo]:
    ranked: list[FlavorInfo] = []
    seen: set[str] = set()
    for label in preference:
        for flavor in flavors:
            fid = _id(flavor)
            if fid in seen or is_disabled(flavor):
                continue
            if match_label(label, flavor):
                ranked.append(flavor_info(flavor, label))
                seen.add(fid)
                break
    return ranked


def pick_volume_type(types: list[Any], az: str) -> str | None:
    az_l = az.lower()
    names = [_name(t) for t in types if _name(t)]
    for name in names:
        if az_l in name.lower() and "fast" in name.lower():
            return name
    for name in names:
        if az_l in name.lower():
            return name
    return names[0] if names else None


def gpu_quota_from_compute(quota: dict[str, Any]) -> int | None:
    """Return GPU limit if the quota dict exposes it; else None."""
    keys = []
    for key in quota:
        if "gpu" in str(key).lower():
            keys.append(key)
    if not keys:
        return None
    values = []
    for key in keys:
        raw = quota[key]
        if isinstance(raw, dict):
            raw = raw.get("limit", raw.get("in_use"))
        try:
            values.append(int(raw))
        except (TypeError, ValueError):
            continue
    if not values:
        return None
    return max(values)


def gpu_boot_image_score(name: str) -> int:
    """Higher is better. Canonical: Ubuntu 24.04 + driver 580, no Docker."""
    n = name.lower()
    if "gpu" not in n:
        return 0
    if "data science" in n or "analytics" in n:
        return 1
    score = 10
    if "docker" in n:
        score -= 30
    if "24.04" in n:
        score += 20
    elif "22.04" in n:
        score += 5
    if "580" in n:
        score += 15
    elif "535" in n:
        score += 4
    return score


def pick_boot_image(images: list[Any]) -> Any | None:
    ranked = [(gpu_boot_image_score(_name(img)), img) for img in images]
    ranked = [item for item in ranked if item[0] > 0]
    if not ranked:
        return None
    ranked.sort(key=lambda item: item[0], reverse=True)
    return ranked[0][1]


def resolve_flavor(
    flavors: list[Any],
    preference: tuple[str, ...],
    *,
    explicit: str | None = None,
    default_id: str | None = None,
    fallback: bool = True,
) -> FlavorInfo:
    if explicit:
        for flavor in flavors:
            if _id(flavor) == explicit or _name(flavor) == explicit:
                if is_disabled(flavor):
                    raise ValueError(f"flavor {explicit} disabled")
                return flavor_info(flavor, label="explicit")
        raise ValueError(f"flavor {explicit} не найден в регионе")
    if not fallback:
        if not default_id:
            raise ValueError("FLAVOR_FALLBACK=false требует DEFAULT_FLAVOR_ID или --flavor")
        return resolve_flavor(flavors, preference, explicit=default_id, fallback=True)
    gpu = [f for f in flavors if looks_like_gpu(f)]
    ranked = rank_flavors(gpu or flavors, preference)
    if ranked:
        return ranked[0]
    if default_id:
        return resolve_flavor(flavors, preference, explicit=default_id, fallback=True)
    raise ValueError("нет доступного GPU flavor из FLAVOR_PREFERENCE")