first commit

src/gpu_rent/inventory.py

@@ -0,0 +1,169 @@
+"""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")
+    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]