353 lines
13 KiB
Python
353 lines
13 KiB
Python
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
|