Files
ttulttul-ComfyUI-Modal/vast_models.py

665 lines
25 KiB
Python

"""Typed Vast.ai marketplace, configuration, and instance models."""
from __future__ import annotations
import math
import re
from dataclasses import dataclass
from enum import Enum
from typing import Any, Mapping, Sequence
VAST_API_BASE_URL = "https://console.vast.ai"
VAST_CONFIG_NODE_ID = "VastAILeaseConfiguration"
VAST_DEFAULT_IDLE_RETENTION_HOURS = 24.0
VAST_DEFAULT_MINIMUM_OFFER_DURATION_DAYS = 7.0
VAST_DEFAULT_EXCLUDED_COUNTRY_CODES = ("CN", "RU")
_PROFILE_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,62}$")
_COUNTRY_CODE_PATTERN = re.compile(r"^[A-Z]{2}$")
class VastRentalType(str, Enum):
"""Identify a supported Vast.ai marketplace rental type."""
ON_DEMAND = "on-demand"
INTERRUPTIBLE = "bid"
@dataclass(frozen=True)
class VastResourceProfile:
"""Describe one workflow-declared Vast.ai capacity pool."""
profile_id: str
profile_name: str
gpu_count: int | None = None
minimum_gpu_ram_mb: int | None = None
minimum_total_flops: float | None = None
minimum_cpu_ram_mb: int | None = None
minimum_cpu_cores: float | None = None
allocated_disk_gb: float = 200.0
maximum_hourly_cost_usd: float | None = None
idle_retention_seconds: float = VAST_DEFAULT_IDLE_RETENTION_HOURS * 3600
minimum_offer_duration_seconds: float | None = None
minimum_reliability: float | None = None
minimum_dlperf: float | None = None
minimum_download_mb_per_second: float | None = None
minimum_cuda_version: float = 13.0
verified_only: bool = False
allowed_geolocations: tuple[str, ...] = ()
excluded_country_codes: tuple[str, ...] = VAST_DEFAULT_EXCLUDED_COUNTRY_CODES
maximum_instances: int = 1
rental_type: VastRentalType = VastRentalType.ON_DEMAND
def __post_init__(self) -> None:
"""Reject malformed profile values before querying the marketplace."""
if not self.profile_id.strip():
raise ValueError("Vast profile_id must not be empty.")
if not _PROFILE_NAME_PATTERN.fullmatch(self.profile_name):
raise ValueError(
"Vast profile_name must contain letters, digits, dots, underscores, "
"or hyphens and be at most 63 characters."
)
if self.gpu_count is not None:
_require_positive_int(self.gpu_count, "gpu_count")
if self.minimum_gpu_ram_mb is not None:
_require_non_negative_int(
self.minimum_gpu_ram_mb,
"minimum_gpu_ram_mb",
)
if self.minimum_cpu_ram_mb is not None:
_require_non_negative_int(
self.minimum_cpu_ram_mb,
"minimum_cpu_ram_mb",
)
_require_positive_int(self.maximum_instances, "maximum_instances")
for field_name, value in (
("minimum_total_flops", self.minimum_total_flops),
("minimum_cpu_cores", self.minimum_cpu_cores),
("allocated_disk_gb", self.allocated_disk_gb),
("maximum_hourly_cost_usd", self.maximum_hourly_cost_usd),
("idle_retention_seconds", self.idle_retention_seconds),
("minimum_offer_duration_seconds", self.minimum_offer_duration_seconds),
("minimum_dlperf", self.minimum_dlperf),
(
"minimum_download_mb_per_second",
self.minimum_download_mb_per_second,
),
("minimum_cuda_version", self.minimum_cuda_version),
):
if value is not None:
_require_finite_non_negative(value, field_name)
if self.allocated_disk_gb <= 0:
raise ValueError("allocated_disk_gb must be positive.")
if (
self.maximum_hourly_cost_usd is not None
and self.maximum_hourly_cost_usd <= 0
):
raise ValueError("maximum_hourly_cost_usd must be positive.")
if self.minimum_reliability is not None and (
not math.isfinite(self.minimum_reliability)
or not 0.0 <= self.minimum_reliability <= 1.0
):
raise ValueError("minimum_reliability must be between 0 and 1.")
if any(
not location.strip()
or len(location.strip()) > 64
or any(character in location for character in ("\x00", "\n", "\r"))
for location in self.allowed_geolocations
):
raise ValueError("allowed_geolocations contains an invalid location.")
if any(
not _COUNTRY_CODE_PATTERN.fullmatch(country_code)
for country_code in self.excluded_country_codes
):
raise ValueError(
"excluded_country_codes must contain uppercase two-letter country codes."
)
if len(set(self.excluded_country_codes)) != len(self.excluded_country_codes):
raise ValueError("excluded_country_codes must not contain duplicates.")
allowed_country_codes = {
location.strip().upper()
for location in self.allowed_geolocations
if _COUNTRY_CODE_PATTERN.fullmatch(location.strip().upper())
}
overlap = allowed_country_codes.intersection(self.excluded_country_codes)
if overlap:
raise ValueError(
"Vast geolocation filters cannot both allow and exclude: "
+ ", ".join(sorted(overlap))
)
@property
def environment_id(self) -> str:
"""Return the provider-neutral environment identity for this profile."""
safe_profile_id = re.sub(r"[^A-Za-z0-9._-]+", "-", self.profile_id).strip(
"-"
)
return f"vast:{safe_profile_id or self.profile_name}"
def search_payload(self, *, limit: int = 25) -> dict[str, Any]:
"""Return the documented Vast marketplace search request."""
_require_positive_int(limit, "limit")
payload: dict[str, Any] = {
"limit": limit,
"type": self.rental_type.value,
"rentable": {"eq": True},
"rented": {"eq": False},
"gpu_arch": {"eq": "nvidia"},
"cpu_arch": {"eq": "amd64"},
"gpu_frac": {"eq": 1.0},
"gpu_display_active": {"eq": False},
"disk_space": {"gte": self.allocated_disk_gb},
"allocated_storage": self.allocated_disk_gb,
"cuda_max_good": {"gte": self.minimum_cuda_version},
"direct_port_count": {"gte": 1},
"order": [["dph_total", "asc"]],
}
if self.gpu_count is not None:
payload["num_gpus"] = {"eq": self.gpu_count}
if self.minimum_gpu_ram_mb is not None:
payload["gpu_ram"] = {"gte": self.minimum_gpu_ram_mb}
if self.minimum_cpu_ram_mb is not None:
payload["cpu_ram"] = {"gte": self.minimum_cpu_ram_mb}
if self.minimum_cpu_cores is not None:
payload["cpu_cores_effective"] = {"gte": self.minimum_cpu_cores}
if self.minimum_offer_duration_seconds is not None:
payload["duration"] = {"gte": self.minimum_offer_duration_seconds}
if self.minimum_reliability is not None:
payload["reliability"] = {"gte": self.minimum_reliability}
if self.minimum_download_mb_per_second is not None:
payload["inet_down"] = {
"gte": self.minimum_download_mb_per_second
}
if self.maximum_hourly_cost_usd is not None:
payload["dph_total"] = {"lte": self.maximum_hourly_cost_usd}
if self.minimum_total_flops is not None:
payload["total_flops"] = {"gte": self.minimum_total_flops}
if self.minimum_dlperf is not None:
payload["dlperf"] = {"gte": self.minimum_dlperf}
if self.verified_only:
payload["verified"] = {"eq": True}
geolocation_filter: dict[str, list[str]] = {}
if self.allowed_geolocations:
geolocation_filter["in"] = [
location.strip() for location in self.allowed_geolocations
]
if self.excluded_country_codes:
geolocation_filter["notin"] = list(self.excluded_country_codes)
if geolocation_filter:
payload["geolocation"] = geolocation_filter
return payload
@dataclass(frozen=True)
class VastOffer:
"""Describe one normalized Vast marketplace offer."""
offer_id: int
gpu_name: str
num_gpus: int
gpu_ram_mb: int
gpu_total_ram_mb: int
total_flops: float
cpu_ram_mb: int
cpu_cores_effective: float
disk_space_gb: float
duration_seconds: float
reliability: float
dlperf: float
download_mb_per_second: float
cuda_max_good: float
direct_port_count: int
hourly_cost_usd: float
storage_cost_usd_per_gb_month: float | None = None
geolocation: str | None = None
verification: str | None = None
@classmethod
def from_api(cls, payload: Mapping[str, Any]) -> "VastOffer":
"""Normalize one marketplace response record."""
return cls(
offer_id=_required_int(payload, "id"),
gpu_name=str(payload.get("gpu_name") or "Unknown GPU").strip(),
num_gpus=_non_negative_int(payload.get("num_gpus"), "num_gpus"),
gpu_ram_mb=_non_negative_int(payload.get("gpu_ram"), "gpu_ram"),
gpu_total_ram_mb=_non_negative_int(
payload.get("gpu_total_ram", payload.get("gpu_ram")),
"gpu_total_ram",
),
total_flops=_non_negative_float(
payload.get("total_flops"), "total_flops"
),
cpu_ram_mb=_non_negative_int(payload.get("cpu_ram"), "cpu_ram"),
cpu_cores_effective=_non_negative_float(
payload.get("cpu_cores_effective", payload.get("cpu_cores")),
"cpu_cores_effective",
),
disk_space_gb=_non_negative_float(
payload.get("disk_space"), "disk_space"
),
duration_seconds=_non_negative_float(
payload.get("duration"), "duration"
),
reliability=_non_negative_float(
payload.get("reliability"), "reliability"
),
dlperf=_non_negative_float(payload.get("dlperf"), "dlperf"),
download_mb_per_second=_non_negative_float(
payload.get("inet_down"), "inet_down"
),
cuda_max_good=_non_negative_float(
payload.get("cuda_max_good"), "cuda_max_good"
),
direct_port_count=_non_negative_int(
payload.get("direct_port_count"), "direct_port_count"
),
hourly_cost_usd=_non_negative_float(
payload.get("dph_total"), "dph_total"
),
storage_cost_usd_per_gb_month=_optional_non_negative_float(
payload.get("storage_cost"), "storage_cost"
),
geolocation=_optional_string(payload.get("geolocation")),
verification=_optional_string(payload.get("verification")),
)
def incompatibility_reason(self, profile: VastResourceProfile) -> str | None:
"""Return the first unmet hard constraint, if any."""
checks: Sequence[tuple[bool, str]] = (
(
profile.gpu_count is None or self.num_gpus == profile.gpu_count,
"GPU count does not match",
),
(
profile.minimum_gpu_ram_mb is None
or self.gpu_ram_mb >= profile.minimum_gpu_ram_mb,
"insufficient GPU RAM",
),
(
profile.minimum_total_flops is None
or self.total_flops >= profile.minimum_total_flops,
"insufficient total TFLOPS",
),
(
profile.minimum_cpu_ram_mb is None
or self.cpu_ram_mb >= profile.minimum_cpu_ram_mb,
"insufficient CPU RAM",
),
(
profile.minimum_cpu_cores is None
or self.cpu_cores_effective >= profile.minimum_cpu_cores,
"insufficient effective CPU cores",
),
(
self.disk_space_gb >= profile.allocated_disk_gb,
"insufficient disk space",
),
(
profile.minimum_offer_duration_seconds is None
or self.duration_seconds >= profile.minimum_offer_duration_seconds,
"insufficient offer duration",
),
(
profile.minimum_reliability is None
or self.reliability >= profile.minimum_reliability,
"insufficient reliability",
),
(
profile.minimum_dlperf is None
or self.dlperf >= profile.minimum_dlperf,
"insufficient DLPerf",
),
(
profile.minimum_download_mb_per_second is None
or self.download_mb_per_second
>= profile.minimum_download_mb_per_second,
"insufficient download bandwidth",
),
(
self.cuda_max_good >= profile.minimum_cuda_version,
"insufficient CUDA support",
),
(self.direct_port_count >= 1, "no direct SSH port is available"),
(
profile.maximum_hourly_cost_usd is None
or self.hourly_cost_usd <= profile.maximum_hourly_cost_usd,
"hourly price exceeds profile maximum",
),
)
for compatible, reason in checks:
if not compatible:
return reason
if profile.verified_only and self.verification not in {None, "verified"}:
return "offer is not verified"
if profile.allowed_geolocations and not _matches_geolocation(
self.geolocation, profile.allowed_geolocations
):
return "geolocation is not allowed"
if _matches_geolocation(self.geolocation, profile.excluded_country_codes):
return "geolocation is excluded"
return None
def ranking_key(self) -> tuple[float, float, float, float, int]:
"""Return deterministic price-first marketplace ordering."""
return (
self.hourly_cost_usd,
-self.reliability,
-self.dlperf,
-self.download_mb_per_second,
self.offer_id,
)
@dataclass(frozen=True)
class VastInstance:
"""Describe the lifecycle and connection fields used by the controller."""
instance_id: int
actual_status: str
intended_status: str | None
current_state: str | None
status_message: str | None
label: str | None
ssh_host: str | None
ssh_port: int | None
gpu_name: str | None
num_gpus: int
gpu_ram_mb: int
cpu_ram_mb: int
hourly_cost_usd: float | None
ssh_direct_host: str | None
ssh_direct_port: int | None
ssh_proxy_host: str | None
ssh_proxy_port: int | None
@classmethod
def from_api(cls, payload: Mapping[str, Any]) -> "VastInstance":
"""Normalize one Vast instance response record."""
(
ssh_host,
ssh_port,
ssh_direct_host,
ssh_direct_port,
ssh_proxy_host,
ssh_proxy_port,
) = _instance_ssh_endpoints(payload)
return cls(
instance_id=_required_int(payload, "id"),
actual_status=str(payload.get("actual_status") or "unknown").strip(),
intended_status=_optional_string(payload.get("intended_status")),
current_state=_optional_string(payload.get("cur_state")),
status_message=_bounded_status_string(payload.get("status_msg")),
label=_optional_string(payload.get("label")),
ssh_host=ssh_host,
ssh_port=ssh_port,
gpu_name=_optional_string(payload.get("gpu_name")),
num_gpus=_non_negative_int(payload.get("num_gpus"), "num_gpus"),
gpu_ram_mb=_non_negative_int(payload.get("gpu_ram"), "gpu_ram"),
cpu_ram_mb=_non_negative_int(payload.get("cpu_ram"), "cpu_ram"),
hourly_cost_usd=_optional_non_negative_float(
payload.get("dph_total"), "dph_total"
),
ssh_direct_host=ssh_direct_host,
ssh_direct_port=ssh_direct_port,
ssh_proxy_host=ssh_proxy_host,
ssh_proxy_port=ssh_proxy_port,
)
@property
def ready_for_ssh(self) -> bool:
"""Return whether the instance reports a usable SSH endpoint."""
return (
self.actual_status == "running"
and bool(self.ssh_host)
and self.ssh_port is not None
and self.ssh_port > 0
)
@property
def lifecycle_status(self) -> str:
"""Return the most informative provider lifecycle state available."""
for candidate in (
self.actual_status,
self.current_state,
self.intended_status,
):
normalized = str(candidate or "").strip()
if normalized and normalized.casefold() != "unknown":
return normalized
return "unknown"
@property
def terminal_failure(self) -> bool:
"""Return whether provider state proves startup cannot continue."""
terminal_states = {"destroyed", "error", "exited"}
if self.actual_status.casefold() in terminal_states:
return True
current_state = (self.current_state or "").casefold()
intended_status = (self.intended_status or "").casefold()
if current_state in terminal_states:
return True
return current_state == "stopped" and intended_status == "stopped"
@dataclass(frozen=True)
class VastInstanceLaunchSpec:
"""Describe the immutable launch settings for one managed Vast instance."""
image: str
disk_gb: float
label: str
onstart: str
environment: Mapping[str, str]
rental_type: VastRentalType = VastRentalType.ON_DEMAND
bid_price_usd: float | None = None
def __post_init__(self) -> None:
"""Validate fields before they can trigger a billable API call."""
if not self.image.strip():
raise ValueError("Vast launch image must not be empty.")
if not self.label.strip():
raise ValueError("Vast launch label must not be empty.")
_require_finite_non_negative(self.disk_gb, "disk_gb")
if self.disk_gb <= 0:
raise ValueError("disk_gb must be positive.")
if self.rental_type is VastRentalType.INTERRUPTIBLE:
if self.bid_price_usd is None or self.bid_price_usd <= 0:
raise ValueError("Interruptible Vast launches require a positive bid.")
elif self.bid_price_usd is not None:
raise ValueError("On-demand Vast launches must not include a bid price.")
for name, value in self.environment.items():
if not name.strip() or any(
character in name or character in value
for character in ("\x00", "\n", "\r")
):
raise ValueError("Vast launch environment contains an unsafe value.")
def to_api_payload(self) -> dict[str, Any]:
"""Return the documented create-instance request body."""
payload: dict[str, Any] = {
"image": self.image,
"disk": self.disk_gb,
"label": self.label,
"runtype": "ssh_direct",
"target_state": "running",
"cancel_unavail": True,
"python_utf8": True,
"lang_utf8": True,
"env": dict(self.environment),
"onstart": self.onstart,
}
if self.bid_price_usd is not None:
payload["price"] = self.bid_price_usd
return payload
def compatible_offers(
offers: Sequence[VastOffer], profile: VastResourceProfile
) -> tuple[VastOffer, ...]:
"""Return compatible offers in deterministic best-price order."""
return tuple(
sorted(
(
offer
for offer in offers
if offer.incompatibility_reason(profile) is None
),
key=VastOffer.ranking_key,
)
)
def _matches_geolocation(
value: str | None, allowed_geolocations: Sequence[str]
) -> bool:
"""Return whether an offer's location matches an allowed country or label."""
normalized_value = (value or "").strip().casefold()
return any(
normalized_value == allowed.strip().casefold()
or normalized_value.endswith(f", {allowed.strip().casefold()}")
for allowed in allowed_geolocations
)
def _instance_ssh_endpoints(
payload: Mapping[str, Any],
) -> tuple[
str | None,
int | None,
str | None,
int | None,
str | None,
int | None,
]:
"""Return the preferred endpoint while retaining direct and proxy choices."""
public_host = _optional_string(payload.get("public_ipaddr"))
direct_host: str | None = None
direct_port: int | None = None
raw_ports = payload.get("ports")
if public_host is not None and isinstance(raw_ports, Mapping):
raw_ssh_mappings = raw_ports.get("22/tcp")
ssh_mappings = (
raw_ssh_mappings
if isinstance(raw_ssh_mappings, list)
else [raw_ssh_mappings]
)
for raw_mapping in ssh_mappings:
if not isinstance(raw_mapping, Mapping):
continue
raw_host_port = raw_mapping.get("HostPort")
try:
host_port = int(raw_host_port)
except (TypeError, ValueError):
continue
if 1 <= host_port <= 65535:
direct_host = public_host
direct_port = host_port
break
raw_proxy_host = _optional_string(payload.get("ssh_host"))
raw_proxy_port = _optional_int(payload.get("ssh_port"), "ssh_port")
proxy_complete = raw_proxy_host is not None and raw_proxy_port is not None
proxy_host = raw_proxy_host if proxy_complete else None
proxy_port = raw_proxy_port if proxy_complete else None
primary_host = direct_host or raw_proxy_host
primary_port = direct_port if direct_host is not None else raw_proxy_port
return (
primary_host,
primary_port,
direct_host,
direct_port,
proxy_host,
proxy_port,
)
def _required_int(payload: Mapping[str, Any], field_name: str) -> int:
"""Return one required integer field."""
if field_name not in payload:
raise ValueError(f"Vast response omitted required field {field_name!r}.")
return _non_negative_int(payload[field_name], field_name)
def _non_negative_int(value: Any, field_name: str) -> int:
"""Return one non-negative integer field."""
if isinstance(value, bool):
raise ValueError(f"{field_name} must be a non-negative integer.")
try:
normalized = int(value or 0)
except (TypeError, ValueError) as exc:
raise ValueError(f"{field_name} must be a non-negative integer.") from exc
if normalized < 0:
raise ValueError(f"{field_name} must be a non-negative integer.")
return normalized
def _optional_int(value: Any, field_name: str) -> int | None:
"""Return one optional integer field."""
if value is None:
return None
return _non_negative_int(value, field_name)
def _non_negative_float(value: Any, field_name: str) -> float:
"""Return one finite non-negative float field."""
try:
normalized = float(value or 0.0)
except (TypeError, ValueError) as exc:
raise ValueError(f"{field_name} must be a non-negative number.") from exc
_require_finite_non_negative(normalized, field_name)
return normalized
def _optional_non_negative_float(value: Any, field_name: str) -> float | None:
"""Return one optional finite non-negative float field."""
if value is None:
return None
return _non_negative_float(value, field_name)
def _optional_string(value: Any) -> str | None:
"""Return a stripped optional string."""
if value is None:
return None
normalized = str(value).strip()
return normalized or None
def _bounded_status_string(value: Any) -> str | None:
"""Return one single-line provider status suitable for logs and errors."""
normalized = _optional_string(value)
if normalized is None:
return None
return " ".join(normalized.split())[:240]
def _require_positive_int(value: int, field_name: str) -> None:
"""Require a positive non-boolean integer."""
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ValueError(f"{field_name} must be a positive integer.")
def _require_non_negative_int(value: int, field_name: str) -> None:
"""Require a non-negative non-boolean integer."""
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{field_name} must be a non-negative integer.")
def _require_finite_non_negative(value: float, field_name: str) -> None:
"""Require a finite non-negative number."""
if not math.isfinite(value) or value < 0:
raise ValueError(f"{field_name} must be finite and non-negative.")
__all__ = [
"VAST_API_BASE_URL",
"VAST_CONFIG_NODE_ID",
"VAST_DEFAULT_EXCLUDED_COUNTRY_CODES",
"VAST_DEFAULT_IDLE_RETENTION_HOURS",
"VAST_DEFAULT_MINIMUM_OFFER_DURATION_DAYS",
"VastInstance",
"VastInstanceLaunchSpec",
"VastOffer",
"VastRentalType",
"VastResourceProfile",
"compatible_offers",
]