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

from __future__ import annotations

from datetime import UTC, datetime
import json
import math
import os
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import asdict, dataclass, field
from typing import Any, Iterable, 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
    disk_space: 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"),
            disk_space=_optional_float(data, "disk_space"),
            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]
    created_at: float = field(default_factory=time.time)

    @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
        now = time.time()
        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,
            "created_at": self.created_at,
            "created_at_iso": datetime.fromtimestamp(self.created_at, UTC).isoformat(),
            "age_seconds": max(0.0, now - self.created_at),
        }


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}
    query["disk_space"] = {"gte": provider.disk_gb}
    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 list_instances(*, base_url: str, api_key: str) -> list[dict[str, Any]]:
    payload = _vast_api_json_request(
        base_url=base_url,
        api_key=api_key,
        path="/api/v0/instances/",
        method="GET",
    )
    return _instances_from_payload(payload)


def destroy_instance(*, base_url: str, api_key: str, instance_id: int) -> Any:
    return _vast_api_json_request(
        base_url=base_url,
        api_key=api_key,
        path=f"/api/v0/instances/{int(instance_id)}/",
        method="DELETE",
    )


def summarize_instances(instances: list[Mapping[str, Any]]) -> list[dict[str, Any]]:
    fields = (
        "id",
        "instance_id",
        "machine_id",
        "host_id",
        "label",
        "status",
        "actual_status",
        "gpu_name",
        "num_gpus",
        "dph_total",
        "ssh_host",
        "ssh_port",
        "start_date",
    )
    summaries: list[dict[str, Any]] = []
    for instance in instances:
        summary = {field: instance[field] for field in fields if field in instance}
        summaries.append(summary)
    return summaries


def _vast_api_json_request(*, base_url: str, api_key: str, path: str, method: str) -> Any:
    url = f"{base_url.rstrip('/')}/{path.lstrip('/')}"
    request = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"}, method=method)
    try:
        with urllib.request.urlopen(request, timeout=45) as response:
            return 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 API {method} {path} HTTP {exc.code}: {body}") from exc
    except OSError as exc:
        raise RuntimeError(f"Vast API {method} {path} failed: {exc}") from exc


def _instances_from_payload(payload: Any) -> list[dict[str, Any]]:
    if isinstance(payload, list):
        raw_instances = payload
    elif isinstance(payload, Mapping):
        raw_instances = None
        for key in ("instances", "results", "items"):
            value = payload.get(key)
            if isinstance(value, list):
                raw_instances = value
                break
        if raw_instances is None:
            raise RuntimeError("Vast instances response missing instances list")
    else:
        raise RuntimeError("Vast instances response is not JSON object or list")
    return [dict(item) for item in raw_instances if isinstance(item, Mapping)]


def choose_offer(
    offers: list[VastOffer],
    config: RemoteRunConfig,
    *,
    query: Mapping[str, Any],
    reserved_host_ids: Iterable[int] = (),
    allow_reserved_hosts: bool = False,
) -> 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]
    reserved_hosts = set(reserved_host_ids)
    if reserved_hosts and not allow_reserved_hosts:
        ranked = [offer for offer in ranked if offer.host_id is None or offer.host_id not in reserved_hosts]
    if not ranked:
        raise RuntimeError("No Vast offers survived quality, price, and host anti-collision filters")
    selected = ranked[0]
    policy = selection_policy_manifest(config)
    policy["reserved_host_ids"] = sorted(reserved_hosts)
    policy["allow_reserved_hosts"] = bool(allow_reserved_hosts)
    return SelectionResult(
        selected_offer=selected,
        candidate_count=len(offers),
        survivor_count=len(ranked),
        effective_price=effective_price(selected, config.selection),
        query=dict(query),
        policy=policy,
    )


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