"""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")