airfRANS-model-exploration/src/airfrans_frontier/remote/vast.py

from __future__ import annotations

import json
import math
import os
import urllib.parse
import urllib.request
from dataclasses import asdict, dataclass
from typing import Any, Mapping

from airfrans_frontier.remote.config import RemoteRunConfig, SelectionConfig


@dataclass(frozen=True)
class VastOffer:
    id: int
    gpu_name: str
    dph_total: float
    gpu_ram: float | None
    geolocation: str | None
    inet_down_cost_per_tb: float
    inet_up_cost_per_tb: float
    host_id: int | None
    verification: str | None
    reliability2: float | None
    cuda_max_good: float | None
    direct_port_count: int | None
    inet_down: float | None
    inet_up: float | None
    verified: bool | None

    @classmethod
    def from_mapping(cls, data: Mapping[str, Any]) -> VastOffer:
        return cls(
            id=_int(data, "id"),
            gpu_name=_string(data, "gpu_name"),
            dph_total=_float(data, "dph_total"),
            gpu_ram=_optional_float(data, "gpu_ram"),
            geolocation=_optional_string(data, "geolocation"),
            inet_down_cost_per_tb=_optional_float(data, "internet_down_cost_per_tb") or 0.0,
            inet_up_cost_per_tb=_optional_float(data, "internet_up_cost_per_tb") or 0.0,
            host_id=_optional_int(data, "host_id"),
            verification=_optional_string(data, "verification"),
            reliability2=_optional_float(data, "reliability2"),
            cuda_max_good=_optional_float(data, "cuda_max_good"),
            direct_port_count=_optional_int(data, "direct_port_count"),
            inet_down=_optional_float(data, "inet_down"),
            inet_up=_optional_float(data, "inet_up"),
            verified=_optional_bool(data, "verified"),
        )


@dataclass(frozen=True)
class SelectionResult:
    selected_offer: VastOffer
    candidate_count: int
    survivor_count: int
    effective_price: float
    query: dict[str, Any]
    policy: dict[str, Any]

    @property
    def selected_offer_id(self) -> int:
        return self.selected_offer.id

    def to_manifest(self) -> dict[str, Any]:
        offer = asdict(self.selected_offer)
        offer["effective_price"] = self.effective_price
        return {
            "selected_offer_id": self.selected_offer_id,
            "selected_offer": offer,
            "candidate_count": self.candidate_count,
            "survivor_count": self.survivor_count,
            "query": self.query,
            "policy": self.policy,
        }


def select_offer(config: RemoteRunConfig, *, api_key: str | None = None) -> SelectionResult:
    if config.provider.kind != "vastai":
        raise ValueError(f"Unsupported provider: {config.provider.kind}")
    resolved_key = api_key or os.environ.get("VAST_API_KEY")
    if not resolved_key:
        raise RuntimeError("VAST_API_KEY is required for Vast.ai offer selection")

    query = build_query(config)
    offers = search_offers(
        base_url=config.selection.base_url,
        api_key=resolved_key,
        query=query,
    )
    return choose_offer(offers, config, query=query)


def build_query(config: RemoteRunConfig) -> dict[str, Any]:
    selection = config.selection
    provider = config.provider
    query: dict[str, Any] = {
        "rentable": {"eq": True},
        "rented": {"eq": False},
        "reliability2": {"gte": selection.min_reliability},
        "cuda_max_good": {"gte": 12.6},
        "direct_port_count": {"gte": 1},
        "num_gpus": {"eq": provider.gpu.count},
        "inet_down": {"gte": selection.min_down_mbps},
        "limit": 5000,
    }
    if selection.min_up_mbps is not None:
        query["inet_up"] = {"gte": selection.min_up_mbps}
    if selection.require_verified:
        query["verified"] = {"eq": True}
    if provider.gpu.min_vram_gb is not None:
        query["gpu_ram"] = {"gte": provider.gpu.min_vram_gb * 1024}
    if provider.gpu.name:
        query["gpu_name"] = {"eq": provider.gpu.name}
    return query


def search_offers(*, base_url: str, api_key: str, query: Mapping[str, Any]) -> list[VastOffer]:
    encoded = urllib.parse.quote(json.dumps(query, separators=(",", ":")))
    url = f"{base_url.rstrip('/')}/api/v0/bundles/?q={encoded}"
    request = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"})
    try:
        with urllib.request.urlopen(request, timeout=45) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"Vast offer search HTTP {exc.code}: {body}") from exc
    except OSError as exc:
        raise RuntimeError(f"Vast offer search failed: {exc}") from exc

    raw_offers = payload.get("offers")
    if not isinstance(raw_offers, list):
        raise RuntimeError("Vast offer search response missing offers list")
    return [VastOffer.from_mapping(item) for item in raw_offers if isinstance(item, Mapping)]


def choose_offer(offers: list[VastOffer], config: RemoteRunConfig, *, query: Mapping[str, Any]) -> SelectionResult:
    survivors = reachable_offers(offers, config.selection)
    ranked = rank_survivors(survivors, config.selection)
    if config.provider.max_price_per_hour is not None:
        ranked = [offer for offer in ranked if effective_price(offer, config.selection) <= config.provider.max_price_per_hour]
    if not ranked:
        raise RuntimeError("No Vast offers survived quality filters and price cap")
    selected = ranked[0]
    return SelectionResult(
        selected_offer=selected,
        candidate_count=len(offers),
        survivor_count=len(ranked),
        effective_price=effective_price(selected, config.selection),
        query=dict(query),
        policy=selection_policy_manifest(config),
    )


def reachable_offers(offers: list[VastOffer], selection: SelectionConfig) -> list[VastOffer]:
    blacklist = set(selection.blacklist_hosts)
    blocked = tuple(item.upper() for item in selection.blocked_geos)
    result: list[VastOffer] = []
    for offer in offers:
        geo = (offer.geolocation or "").upper()
        if blocked and any(token and token in geo for token in blocked):
            continue
        if offer.host_id is not None and offer.host_id in blacklist:
            continue
        if offer.verification == "deverified":
            continue
        result.append(offer)
    return result


def rank_survivors(offers: list[VastOffer], selection: SelectionConfig) -> list[VastOffer]:
    by_model: dict[str, list[VastOffer]] = {}
    for offer in offers:
        by_model.setdefault(offer.gpu_name, []).append(offer)

    survivors: list[VastOffer] = []
    for group in by_model.values():
        group.sort(key=lambda offer: effective_price(offer, selection))
        drop = math.floor(selection.drop_cheap_frac * len(group))
        survivors.extend(group[drop:])
    survivors.sort(key=lambda offer: effective_price(offer, selection))
    return survivors


def effective_price(offer: VastOffer, selection: SelectionConfig) -> float:
    image_pull = 0.0
    if selection.image_size_gb is not None:
        image_pull = selection.image_size_gb * offer.inet_down_cost_per_tb / 1000.0
    return offer.dph_total + image_pull


def selection_policy_manifest(config: RemoteRunConfig) -> dict[str, Any]:
    return {
        "gpu_name": config.provider.gpu.name,
        "gpu_count": config.provider.gpu.count,
        "min_vram_gb": config.provider.gpu.min_vram_gb,
        "max_price_per_hour": config.provider.max_price_per_hour,
        "min_reliability": config.selection.min_reliability,
        "min_down_mbps": config.selection.min_down_mbps,
        "min_up_mbps": config.selection.min_up_mbps,
        "require_verified": config.selection.require_verified,
        "blocked_geos": list(config.selection.blocked_geos),
        "blacklist_hosts": list(config.selection.blacklist_hosts),
        "drop_cheap_frac": config.selection.drop_cheap_frac,
        "image_size_gb": config.selection.image_size_gb,
    }


def _string(data: Mapping[str, Any], key: str) -> str:
    value = data.get(key)
    if not isinstance(value, str):
        raise ValueError(f"Vast offer missing string field: {key}")
    return value


def _optional_string(data: Mapping[str, Any], key: str) -> str | None:
    value = data.get(key)
    return value if isinstance(value, str) else None


def _int(data: Mapping[str, Any], key: str) -> int:
    value = data.get(key)
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"Vast offer missing integer field: {key}")
    return value


def _optional_int(data: Mapping[str, Any], key: str) -> int | None:
    value = data.get(key)
    if isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, float) and value.is_integer():
        return int(value)
    return None


def _float(data: Mapping[str, Any], key: str) -> float:
    value = data.get(key)
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"Vast offer missing numeric field: {key}")
    return float(value)


def _optional_float(data: Mapping[str, Any], key: str) -> float | None:
    value = data.get(key)
    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return float(value)
    return None


def _optional_bool(data: Mapping[str, Any], key: str) -> bool | None:
    value = data.get(key)
    return value if isinstance(value, bool) else None