1122 lines
43 KiB
Python
1122 lines
43 KiB
Python
"""Persistent Vast.ai lease acquisition, reuse, activity, and idle cleanup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from dataclasses import asdict, dataclass, field, replace
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Mapping
|
|
|
|
if __package__:
|
|
from .vast_api import (
|
|
VastApiClient,
|
|
VastApiError,
|
|
VastInstanceNotFoundError,
|
|
VastOfferUnavailableError,
|
|
)
|
|
from .vast_models import (
|
|
VastInstance,
|
|
VastInstanceLaunchSpec,
|
|
VastOffer,
|
|
VastResourceProfile,
|
|
)
|
|
else: # pragma: no cover - direct debugging imports.
|
|
from vast_api import (
|
|
VastApiClient,
|
|
VastApiError,
|
|
VastInstanceNotFoundError,
|
|
VastOfferUnavailableError,
|
|
)
|
|
from vast_models import (
|
|
VastInstance,
|
|
VastInstanceLaunchSpec,
|
|
VastOffer,
|
|
VastResourceProfile,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
VAST_LEASE_REGISTRY_VERSION = 1
|
|
VAST_LEASE_REGISTRY_FILENAME = "vast-leases.json"
|
|
VAST_MANAGED_LABEL_PREFIX = "comfy-modal-vast"
|
|
_SAFE_LABEL_PART_PATTERN = re.compile(r"[^a-zA-Z0-9._-]+")
|
|
VastInstanceStatusCallback = Callable[[VastInstance], None]
|
|
VastLeaseEventCallback = Callable[[str], None]
|
|
VAST_STARTUP_INSTANCE_REPLACEMENTS = 1
|
|
|
|
|
|
class VastLeaseStartupError(RuntimeError):
|
|
"""Describe one rented instance that failed before becoming usable."""
|
|
|
|
def __init__(
|
|
self,
|
|
instance_id: int,
|
|
offer_id: int,
|
|
message: str,
|
|
*,
|
|
cleanup_completed: bool,
|
|
) -> None:
|
|
"""Retain failed capacity identity and whether cleanup was confirmed."""
|
|
cleanup_status = (
|
|
"was destroyed"
|
|
if cleanup_completed
|
|
else "could not be confirmed destroyed"
|
|
)
|
|
super().__init__(
|
|
f"Vast instance {instance_id} from offer {offer_id} failed setup and "
|
|
f"{cleanup_status}: {message}"
|
|
)
|
|
self.instance_id = instance_id
|
|
self.offer_id = offer_id
|
|
self.cleanup_completed = cleanup_completed
|
|
self.failure_message = message
|
|
|
|
|
|
class VastLeaseStartupExhaustedError(RuntimeError):
|
|
"""Report that the initial Vast rental and its replacement both failed."""
|
|
|
|
|
|
def _startup_replacement_status(failure: VastLeaseStartupError) -> str:
|
|
"""Return a clear user-facing transition after one failed rental."""
|
|
cleanup = (
|
|
"terminated"
|
|
if failure.cleanup_completed
|
|
else "marked unusable, but automatic termination was not confirmed"
|
|
)
|
|
return (
|
|
f"Vast.ai instance {failure.instance_id} failed setup and was {cleanup}; "
|
|
"requesting a replacement"
|
|
)
|
|
|
|
|
|
def _startup_terminal_status(failure: VastLeaseStartupError) -> str:
|
|
"""Return a clear status when the bounded replacement also fails."""
|
|
cleanup = (
|
|
"was terminated"
|
|
if failure.cleanup_completed
|
|
else "could not be confirmed terminated"
|
|
)
|
|
return (
|
|
f"Replacement Vast.ai instance {failure.instance_id} also failed setup and "
|
|
f"{cleanup}; no further automatic attempts remain"
|
|
)
|
|
|
|
|
|
def _aggregate_startup_failure(
|
|
profile: VastResourceProfile,
|
|
failures: list[VastLeaseStartupError],
|
|
) -> VastLeaseStartupExhaustedError:
|
|
"""Summarize bounded initial plus replacement startup failures."""
|
|
last_failure = failures[-1]
|
|
instance_ids = ", ".join(str(failure.instance_id) for failure in failures)
|
|
cleanup_summary = (
|
|
"All failed instances were terminated."
|
|
if all(failure.cleanup_completed for failure in failures)
|
|
else "At least one failed instance could not be confirmed terminated."
|
|
)
|
|
return VastLeaseStartupExhaustedError(
|
|
f"Vast profile {profile.profile_name!r} could not obtain SSH-ready "
|
|
f"capacity after {len(failures)} attempts (instances {instance_ids}). "
|
|
f"{cleanup_summary} Last failure: {last_failure.failure_message}"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastLeaseRecord:
|
|
"""Persist one instance owned by a resource profile and ComfyUI installation."""
|
|
|
|
instance_id: int
|
|
offer_id: int
|
|
owner_id: str
|
|
profile_id: str
|
|
profile_name: str
|
|
profile_fingerprint: str
|
|
runtime_fingerprint: str
|
|
label: str
|
|
actual_status: str
|
|
ssh_host: str | None
|
|
ssh_port: int | None
|
|
gpu_name: str
|
|
gpu_count: int
|
|
gpu_ram_mb: int
|
|
cpu_ram_mb: int
|
|
hourly_cost_usd: float
|
|
created_at_epoch: float
|
|
last_activity_at_epoch: float
|
|
idle_deadline_epoch: float
|
|
idle_retention_seconds: float = 24 * 3600
|
|
active_invocations: int = 0
|
|
draining: bool = False
|
|
last_error: str | None = None
|
|
worker_image: str | None = None
|
|
ssh_direct_host: str | None = None
|
|
ssh_direct_port: int | None = None
|
|
ssh_proxy_host: str | None = None
|
|
ssh_proxy_port: int | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Validate ownership, lifecycle, and scheduling fields."""
|
|
for field_name, value in (
|
|
("instance_id", self.instance_id),
|
|
("offer_id", self.offer_id),
|
|
("gpu_count", self.gpu_count),
|
|
):
|
|
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
|
raise ValueError(f"{field_name} must be a positive integer.")
|
|
if self.active_invocations < 0:
|
|
raise ValueError("active_invocations must not be negative.")
|
|
if self.gpu_ram_mb < 0 or self.cpu_ram_mb < 0:
|
|
raise ValueError("Vast lease memory fields must not be negative.")
|
|
for field_name, value in (
|
|
("hourly_cost_usd", self.hourly_cost_usd),
|
|
("created_at_epoch", self.created_at_epoch),
|
|
("last_activity_at_epoch", self.last_activity_at_epoch),
|
|
("idle_deadline_epoch", self.idle_deadline_epoch),
|
|
("idle_retention_seconds", self.idle_retention_seconds),
|
|
):
|
|
if not math.isfinite(value) or value < 0:
|
|
raise ValueError(f"{field_name} must be finite and non-negative.")
|
|
if not self.owner_id.strip() or not self.profile_id.strip():
|
|
raise ValueError("Vast lease owner and profile identities must not be empty.")
|
|
if len(self.profile_fingerprint) != 64 or len(self.runtime_fingerprint) != 64:
|
|
raise ValueError("Vast lease fingerprints must be SHA-256 hex digests.")
|
|
expected_prefix = vast_managed_label_prefix(self.owner_id)
|
|
if not self.label.startswith(f"{expected_prefix}:"):
|
|
raise ValueError("Vast lease label does not match its owner identity.")
|
|
for field_name, port in (
|
|
("ssh_port", self.ssh_port),
|
|
("ssh_direct_port", self.ssh_direct_port),
|
|
("ssh_proxy_port", self.ssh_proxy_port),
|
|
):
|
|
if port is not None and port <= 0:
|
|
raise ValueError(f"{field_name} must be positive when present.")
|
|
for field_name, host, port in (
|
|
("direct", self.ssh_direct_host, self.ssh_direct_port),
|
|
("proxy", self.ssh_proxy_host, self.ssh_proxy_port),
|
|
):
|
|
if (host is None) != (port is None):
|
|
raise ValueError(
|
|
f"Vast SSH {field_name} host and port must be set together."
|
|
)
|
|
if self.worker_image is not None and not self.worker_image.strip():
|
|
raise ValueError("worker_image must be non-empty when present.")
|
|
|
|
@property
|
|
def environment_id(self) -> str:
|
|
"""Return the provider-neutral identity for this concrete lease."""
|
|
return f"vast:{self.profile_id}:{self.instance_id}"
|
|
|
|
@property
|
|
def ready_for_work(self) -> bool:
|
|
"""Return whether this lease is running, reachable, and accepting work."""
|
|
return (
|
|
not self.draining
|
|
and self.actual_status == "running"
|
|
and bool(self.ssh_host)
|
|
and self.ssh_port is not None
|
|
)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Return a credential-free JSON record."""
|
|
return asdict(self)
|
|
|
|
@classmethod
|
|
def from_dict(cls, payload: Mapping[str, Any]) -> "VastLeaseRecord":
|
|
"""Build one validated record from persisted state."""
|
|
return cls(
|
|
instance_id=int(payload["instance_id"]),
|
|
offer_id=int(payload["offer_id"]),
|
|
owner_id=str(payload["owner_id"]),
|
|
profile_id=str(payload["profile_id"]),
|
|
profile_name=str(payload["profile_name"]),
|
|
profile_fingerprint=str(payload["profile_fingerprint"]),
|
|
runtime_fingerprint=str(payload["runtime_fingerprint"]),
|
|
label=str(payload["label"]),
|
|
actual_status=str(payload.get("actual_status") or "unknown"),
|
|
ssh_host=(str(payload["ssh_host"]) if payload.get("ssh_host") else None),
|
|
ssh_port=(int(payload["ssh_port"]) if payload.get("ssh_port") else None),
|
|
gpu_name=str(payload.get("gpu_name") or "Unknown GPU"),
|
|
gpu_count=int(payload.get("gpu_count", 0)),
|
|
gpu_ram_mb=int(payload.get("gpu_ram_mb", 0)),
|
|
cpu_ram_mb=int(payload.get("cpu_ram_mb", 0)),
|
|
hourly_cost_usd=float(payload.get("hourly_cost_usd", 0.0)),
|
|
created_at_epoch=float(payload.get("created_at_epoch", 0.0)),
|
|
last_activity_at_epoch=float(
|
|
payload.get("last_activity_at_epoch", 0.0)
|
|
),
|
|
idle_deadline_epoch=float(payload.get("idle_deadline_epoch", 0.0)),
|
|
idle_retention_seconds=float(
|
|
payload.get("idle_retention_seconds", 24 * 3600)
|
|
),
|
|
active_invocations=int(payload.get("active_invocations", 0)),
|
|
draining=bool(payload.get("draining", False)),
|
|
last_error=(str(payload["last_error"]) if payload.get("last_error") else None),
|
|
worker_image=(
|
|
str(payload["worker_image"]) if payload.get("worker_image") else None
|
|
),
|
|
ssh_direct_host=(
|
|
str(payload["ssh_direct_host"])
|
|
if payload.get("ssh_direct_host")
|
|
else None
|
|
),
|
|
ssh_direct_port=(
|
|
int(payload["ssh_direct_port"])
|
|
if payload.get("ssh_direct_port")
|
|
else None
|
|
),
|
|
ssh_proxy_host=(
|
|
str(payload["ssh_proxy_host"])
|
|
if payload.get("ssh_proxy_host")
|
|
else None
|
|
),
|
|
ssh_proxy_port=(
|
|
int(payload["ssh_proxy_port"])
|
|
if payload.get("ssh_proxy_port")
|
|
else None
|
|
),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastLeaseRegistryState:
|
|
"""Describe the complete versioned Vast lease registry."""
|
|
|
|
leases: tuple[VastLeaseRecord, ...] = ()
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Return the persisted versioned document."""
|
|
return {
|
|
"version": VAST_LEASE_REGISTRY_VERSION,
|
|
"leases": [lease.to_dict() for lease in self.leases],
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, payload: Mapping[str, Any]) -> "VastLeaseRegistryState":
|
|
"""Build and validate one registry document."""
|
|
version = int(payload.get("version", 0))
|
|
if version != VAST_LEASE_REGISTRY_VERSION:
|
|
raise ValueError(f"Unsupported Vast lease registry version {version}.")
|
|
raw_leases = payload.get("leases")
|
|
lease_payloads = raw_leases if isinstance(raw_leases, list) else []
|
|
leases = tuple(
|
|
VastLeaseRecord.from_dict(lease)
|
|
for lease in lease_payloads
|
|
if isinstance(lease, Mapping)
|
|
)
|
|
instance_ids = [lease.instance_id for lease in leases]
|
|
if len(instance_ids) != len(set(instance_ids)):
|
|
raise ValueError("Vast lease registry contains duplicate instance IDs.")
|
|
return cls(leases=leases)
|
|
|
|
|
|
@dataclass
|
|
class VastLeaseRegistry:
|
|
"""Atomically persist managed Vast instances without credentials."""
|
|
|
|
config_path: Path
|
|
_lock: threading.RLock = field(default_factory=threading.RLock)
|
|
|
|
@classmethod
|
|
def for_user_directory(cls, user_directory: Path) -> "VastLeaseRegistry":
|
|
"""Create a registry below one ComfyUI user directory."""
|
|
return cls(
|
|
config_path=(
|
|
user_directory.expanduser().resolve()
|
|
/ "comfyui-modal"
|
|
/ VAST_LEASE_REGISTRY_FILENAME
|
|
)
|
|
)
|
|
|
|
def load(self) -> VastLeaseRegistryState:
|
|
"""Load the registry or return empty initial state."""
|
|
with self._lock:
|
|
if not self.config_path.exists():
|
|
return VastLeaseRegistryState()
|
|
try:
|
|
payload = json.loads(self.config_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise ValueError(
|
|
f"Vast lease registry {self.config_path} is unreadable."
|
|
) from exc
|
|
if not isinstance(payload, Mapping):
|
|
raise ValueError("Vast lease registry must be a JSON object.")
|
|
return VastLeaseRegistryState.from_dict(payload)
|
|
|
|
def save(self, state: VastLeaseRegistryState) -> None:
|
|
"""Validate and atomically persist complete registry state."""
|
|
validated = VastLeaseRegistryState.from_dict(state.to_dict())
|
|
serialized = json.dumps(validated.to_dict(), indent=2, sort_keys=True) + "\n"
|
|
with self._lock:
|
|
self.config_path.parent.mkdir(parents=True, exist_ok=True)
|
|
file_descriptor, temporary_name = tempfile.mkstemp(
|
|
prefix=f".{self.config_path.name}.",
|
|
suffix=".tmp",
|
|
dir=self.config_path.parent,
|
|
)
|
|
temporary_path = Path(temporary_name)
|
|
try:
|
|
with os.fdopen(file_descriptor, "w", encoding="utf-8") as output_file:
|
|
output_file.write(serialized)
|
|
output_file.flush()
|
|
os.fsync(output_file.fileno())
|
|
os.replace(temporary_path, self.config_path)
|
|
finally:
|
|
temporary_path.unlink(missing_ok=True)
|
|
|
|
def upsert(self, lease: VastLeaseRecord) -> VastLeaseRecord:
|
|
"""Insert or replace one record by instance identity."""
|
|
with self._lock:
|
|
state = self.load()
|
|
leases = [
|
|
existing
|
|
for existing in state.leases
|
|
if existing.instance_id != lease.instance_id
|
|
]
|
|
leases.append(lease)
|
|
self.save(
|
|
VastLeaseRegistryState(
|
|
leases=tuple(sorted(leases, key=lambda item: item.instance_id))
|
|
)
|
|
)
|
|
return lease
|
|
|
|
def remove(self, instance_id: int) -> VastLeaseRecord | None:
|
|
"""Remove and return one record when present."""
|
|
with self._lock:
|
|
state = self.load()
|
|
removed = next(
|
|
(lease for lease in state.leases if lease.instance_id == instance_id),
|
|
None,
|
|
)
|
|
if removed is None:
|
|
return None
|
|
self.save(
|
|
VastLeaseRegistryState(
|
|
leases=tuple(
|
|
lease
|
|
for lease in state.leases
|
|
if lease.instance_id != instance_id
|
|
)
|
|
)
|
|
)
|
|
return removed
|
|
|
|
def update(
|
|
self,
|
|
instance_id: int,
|
|
transform: Callable[[VastLeaseRecord], VastLeaseRecord],
|
|
) -> VastLeaseRecord:
|
|
"""Atomically replace one record through a deterministic transform."""
|
|
with self._lock:
|
|
state = self.load()
|
|
updated: VastLeaseRecord | None = None
|
|
leases: list[VastLeaseRecord] = []
|
|
for lease in state.leases:
|
|
if lease.instance_id != instance_id:
|
|
leases.append(lease)
|
|
continue
|
|
updated = transform(lease)
|
|
if updated.instance_id != instance_id:
|
|
raise ValueError("Vast lease transforms cannot change instance ID.")
|
|
leases.append(updated)
|
|
if updated is None:
|
|
raise KeyError(f"Unknown Vast lease {instance_id}.")
|
|
self.save(VastLeaseRegistryState(leases=tuple(leases)))
|
|
return updated
|
|
|
|
|
|
class VastLeaseManager:
|
|
"""Coordinate reusable leases and destroy them after profile idle retention."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
api_client: VastApiClient,
|
|
registry: VastLeaseRegistry,
|
|
owner_id: str,
|
|
runtime_fingerprint: str,
|
|
launch_spec_factory: Callable[
|
|
[VastResourceProfile, str], VastInstanceLaunchSpec
|
|
],
|
|
startup_timeout_seconds: float = 900.0,
|
|
clock: Callable[[], float] = time.time,
|
|
worker_image: str | None = None,
|
|
) -> None:
|
|
"""Configure acquisition, persistence, and lifecycle dependencies."""
|
|
if not owner_id.strip():
|
|
raise ValueError("Vast lease owner_id must not be empty.")
|
|
if len(runtime_fingerprint) != 64:
|
|
raise ValueError("runtime_fingerprint must be a SHA-256 hex digest.")
|
|
if startup_timeout_seconds <= 0:
|
|
raise ValueError("startup_timeout_seconds must be positive.")
|
|
self.api_client = api_client
|
|
self.registry = registry
|
|
self.owner_id = owner_id.strip()
|
|
self.runtime_fingerprint = runtime_fingerprint
|
|
self.launch_spec_factory = launch_spec_factory
|
|
self.startup_timeout_seconds = startup_timeout_seconds
|
|
self.clock = clock
|
|
self.worker_image = worker_image.strip() if worker_image is not None else None
|
|
self._profile_locks_guard = threading.Lock()
|
|
self._profile_locks: dict[str, asyncio.Lock] = {}
|
|
|
|
@classmethod
|
|
def for_inventory(
|
|
cls,
|
|
*,
|
|
api_client: VastApiClient,
|
|
registry: VastLeaseRegistry,
|
|
owner_id: str,
|
|
) -> "VastLeaseManager":
|
|
"""Build a lifecycle-only manager without requiring a worker image."""
|
|
|
|
def unavailable_launch_spec(
|
|
profile: VastResourceProfile,
|
|
label: str,
|
|
) -> VastInstanceLaunchSpec:
|
|
"""Reject acquisition through an inventory-only manager."""
|
|
del profile, label
|
|
raise RuntimeError("Inventory-only Vast manager cannot acquire capacity.")
|
|
|
|
return cls(
|
|
api_client=api_client,
|
|
registry=registry,
|
|
owner_id=owner_id,
|
|
runtime_fingerprint="0" * 64,
|
|
launch_spec_factory=unavailable_launch_spec,
|
|
)
|
|
|
|
async def ensure_lease(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
*,
|
|
slot: int = 0,
|
|
status_callback: VastInstanceStatusCallback | None = None,
|
|
event_callback: VastLeaseEventCallback | None = None,
|
|
excluded_offer_ids: frozenset[int] = frozenset(),
|
|
) -> VastLeaseRecord:
|
|
"""Reuse or rent one compatible SSH-ready lease for a profile slot."""
|
|
if slot < 0 or slot >= profile.maximum_instances:
|
|
raise ValueError(
|
|
f"Vast profile {profile.profile_name!r} supports slots 0 through "
|
|
f"{profile.maximum_instances - 1}."
|
|
)
|
|
slot_profile_id = _slot_profile_id(profile.profile_id, slot)
|
|
async with self._profile_lock(slot_profile_id):
|
|
unavailable_offer_ids = set(excluded_offer_ids)
|
|
existing = await self._reusable_lease(
|
|
profile,
|
|
slot_profile_id,
|
|
status_callback=status_callback,
|
|
event_callback=event_callback,
|
|
unavailable_offer_ids=unavailable_offer_ids,
|
|
)
|
|
if existing is not None:
|
|
return existing
|
|
return await self._rent_lease(
|
|
profile,
|
|
slot_profile_id,
|
|
status_callback=status_callback,
|
|
event_callback=event_callback,
|
|
excluded_offer_ids=frozenset(unavailable_offer_ids),
|
|
)
|
|
|
|
async def _reusable_lease(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
slot_profile_id: str,
|
|
*,
|
|
status_callback: VastInstanceStatusCallback | None,
|
|
event_callback: VastLeaseEventCallback | None,
|
|
unavailable_offer_ids: set[int],
|
|
) -> VastLeaseRecord | None:
|
|
"""Return one matching live lease after refreshing its API state."""
|
|
fingerprint = vast_profile_fingerprint(profile)
|
|
candidates = [
|
|
lease
|
|
for lease in self.registry.load().leases
|
|
if lease.owner_id == self.owner_id
|
|
and lease.profile_id == slot_profile_id
|
|
and lease.profile_fingerprint == fingerprint
|
|
and (
|
|
(
|
|
self.worker_image is None
|
|
and lease.runtime_fingerprint == self.runtime_fingerprint
|
|
)
|
|
or (
|
|
self.worker_image is not None
|
|
and lease.worker_image in {None, self.worker_image}
|
|
)
|
|
)
|
|
and not lease.draining
|
|
]
|
|
for lease in sorted(candidates, key=lambda item: item.created_at_epoch):
|
|
try:
|
|
instance = await self.api_client.show_instance(lease.instance_id)
|
|
except VastInstanceNotFoundError:
|
|
self.registry.remove(lease.instance_id)
|
|
unavailable_offer_ids.add(lease.offer_id)
|
|
continue
|
|
instance = await self._prepare_reusable_instance(
|
|
lease,
|
|
instance,
|
|
status_callback=status_callback,
|
|
event_callback=event_callback,
|
|
)
|
|
if instance is None:
|
|
unavailable_offer_ids.add(lease.offer_id)
|
|
continue
|
|
refreshed = self._refresh_record(lease, instance)
|
|
if refreshed.runtime_fingerprint != self.runtime_fingerprint:
|
|
logger.info(
|
|
"Adopting compatible Vast lease instance_id=%d from runtime=%s "
|
|
"to runtime=%s for unchanged worker image.",
|
|
refreshed.instance_id,
|
|
refreshed.runtime_fingerprint[:12],
|
|
self.runtime_fingerprint[:12],
|
|
)
|
|
refreshed = replace(
|
|
refreshed,
|
|
runtime_fingerprint=self.runtime_fingerprint,
|
|
worker_image=self.worker_image,
|
|
)
|
|
self.registry.upsert(refreshed)
|
|
if refreshed.ready_for_work:
|
|
logger.info(
|
|
"Reusing Vast lease profile=%s instance=%d gpu=%s hourly_cost=%.4f.",
|
|
profile.profile_name,
|
|
refreshed.instance_id,
|
|
refreshed.gpu_name,
|
|
refreshed.hourly_cost_usd,
|
|
)
|
|
return refreshed
|
|
return None
|
|
|
|
async def _prepare_reusable_instance(
|
|
self,
|
|
lease: VastLeaseRecord,
|
|
instance: VastInstance,
|
|
*,
|
|
status_callback: VastInstanceStatusCallback | None,
|
|
event_callback: VastLeaseEventCallback | None,
|
|
) -> VastInstance | None:
|
|
"""Ready an existing contract or destroy it when startup cannot recover."""
|
|
if instance.ready_for_ssh:
|
|
return instance
|
|
try:
|
|
timeout_seconds = self.startup_timeout_seconds
|
|
if instance.actual_status == "stopped":
|
|
await self.api_client.set_instance_state(lease.instance_id, "running")
|
|
else:
|
|
elapsed_seconds = max(0.0, self.clock() - lease.created_at_epoch)
|
|
timeout_seconds = self.startup_timeout_seconds - elapsed_seconds
|
|
if timeout_seconds <= 0:
|
|
raise TimeoutError(
|
|
f"Vast instance {lease.instance_id} exceeded its "
|
|
f"{self.startup_timeout_seconds:.0f}s SSH-readiness window."
|
|
)
|
|
return await self.api_client.wait_until_ready(
|
|
lease.instance_id,
|
|
timeout_seconds=timeout_seconds,
|
|
status_callback=status_callback,
|
|
)
|
|
except VastInstanceNotFoundError:
|
|
self.registry.remove(lease.instance_id)
|
|
return None
|
|
except (TimeoutError, VastApiError) as exc:
|
|
cleanup_completed = await self._discard_startup_lease(lease, str(exc))
|
|
failure = VastLeaseStartupError(
|
|
lease.instance_id,
|
|
lease.offer_id,
|
|
str(exc),
|
|
cleanup_completed=cleanup_completed,
|
|
)
|
|
logger.warning("%s", failure)
|
|
if event_callback is not None:
|
|
event_callback(_startup_replacement_status(failure))
|
|
return None
|
|
|
|
async def _discard_startup_lease(
|
|
self,
|
|
lease: VastLeaseRecord,
|
|
error_message: str,
|
|
) -> bool:
|
|
"""Mark and destroy one owned contract that never became SSH-ready."""
|
|
failed = replace(lease, draining=True, last_error=error_message)
|
|
self.registry.upsert(failed)
|
|
try:
|
|
await self.api_client.destroy_instance(lease.instance_id)
|
|
except VastInstanceNotFoundError:
|
|
self.registry.remove(lease.instance_id)
|
|
return True
|
|
except VastApiError as cleanup_error:
|
|
logger.warning(
|
|
"Unable to destroy failed Vast lease instance=%d: %s",
|
|
lease.instance_id,
|
|
cleanup_error,
|
|
)
|
|
return False
|
|
self.registry.remove(lease.instance_id)
|
|
return True
|
|
|
|
async def _rent_lease(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
slot_profile_id: str,
|
|
*,
|
|
status_callback: VastInstanceStatusCallback | None,
|
|
event_callback: VastLeaseEventCallback | None,
|
|
excluded_offer_ids: frozenset[int],
|
|
) -> VastLeaseRecord:
|
|
"""Search, rent, and persist the first still-available compatible offer."""
|
|
offers = await self.api_client.search_offers(profile)
|
|
if not offers:
|
|
raise VastApiError(
|
|
f"No Vast.ai offer satisfies profile {profile.profile_name!r}."
|
|
)
|
|
unavailable_offer_ids = set(excluded_offer_ids)
|
|
startup_failures: list[VastLeaseStartupError] = []
|
|
for search_attempt in range(2):
|
|
for offer in offers:
|
|
if offer.offer_id in unavailable_offer_ids:
|
|
continue
|
|
try:
|
|
return await self._create_from_offer(
|
|
profile,
|
|
slot_profile_id,
|
|
offer,
|
|
status_callback=status_callback,
|
|
)
|
|
except VastInstanceNotFoundError:
|
|
unavailable_offer_ids.add(offer.offer_id)
|
|
logger.warning(
|
|
"Vast instance disappeared during provider startup "
|
|
"profile=%s offer=%d; trying next candidate.",
|
|
profile.profile_name,
|
|
offer.offer_id,
|
|
)
|
|
except VastOfferUnavailableError:
|
|
unavailable_offer_ids.add(offer.offer_id)
|
|
logger.info(
|
|
"Vast offer disappeared before rental profile=%s offer=%d; trying next candidate.",
|
|
profile.profile_name,
|
|
offer.offer_id,
|
|
)
|
|
except VastLeaseStartupError as exc:
|
|
startup_failures.append(exc)
|
|
unavailable_offer_ids.add(offer.offer_id)
|
|
logger.warning(
|
|
"Vast instance failed during provider startup profile=%s "
|
|
"instance=%d offer=%d cleanup_completed=%s; trying a "
|
|
"replacement.",
|
|
profile.profile_name,
|
|
exc.instance_id,
|
|
exc.offer_id,
|
|
exc.cleanup_completed,
|
|
)
|
|
if len(startup_failures) > VAST_STARTUP_INSTANCE_REPLACEMENTS:
|
|
if event_callback is not None:
|
|
event_callback(_startup_terminal_status(exc))
|
|
raise _aggregate_startup_failure(
|
|
profile,
|
|
startup_failures,
|
|
) from exc
|
|
if event_callback is not None:
|
|
event_callback(_startup_replacement_status(exc))
|
|
if search_attempt == 0:
|
|
offers = await self.api_client.search_offers(
|
|
profile,
|
|
force_refresh=True,
|
|
)
|
|
if startup_failures:
|
|
raise _aggregate_startup_failure(profile, startup_failures)
|
|
raise VastOfferUnavailableError(
|
|
f"Compatible Vast.ai offers disappeared before profile "
|
|
f"{profile.profile_name!r} could be rented."
|
|
)
|
|
|
|
async def _create_from_offer(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
slot_profile_id: str,
|
|
offer: VastOffer,
|
|
*,
|
|
status_callback: VastInstanceStatusCallback | None,
|
|
) -> VastLeaseRecord:
|
|
"""Create, wait for, and persist one selected offer."""
|
|
label = vast_managed_label(
|
|
self.owner_id,
|
|
slot_profile_id,
|
|
vast_profile_fingerprint(profile),
|
|
self.runtime_fingerprint,
|
|
)
|
|
launch_spec = self.launch_spec_factory(profile, label)
|
|
created = await self.api_client.create_instance(offer.offer_id, launch_spec)
|
|
now = self.clock()
|
|
provisional = VastLeaseRecord(
|
|
instance_id=created.instance_id,
|
|
offer_id=offer.offer_id,
|
|
owner_id=self.owner_id,
|
|
profile_id=slot_profile_id,
|
|
profile_name=profile.profile_name,
|
|
profile_fingerprint=vast_profile_fingerprint(profile),
|
|
runtime_fingerprint=self.runtime_fingerprint,
|
|
worker_image=self.worker_image,
|
|
label=label,
|
|
actual_status="loading",
|
|
ssh_host=None,
|
|
ssh_port=None,
|
|
gpu_name=offer.gpu_name,
|
|
gpu_count=offer.num_gpus,
|
|
gpu_ram_mb=offer.gpu_ram_mb,
|
|
cpu_ram_mb=offer.cpu_ram_mb,
|
|
hourly_cost_usd=offer.hourly_cost_usd,
|
|
created_at_epoch=now,
|
|
last_activity_at_epoch=now,
|
|
idle_deadline_epoch=now + profile.idle_retention_seconds,
|
|
idle_retention_seconds=profile.idle_retention_seconds,
|
|
)
|
|
self.registry.upsert(provisional)
|
|
try:
|
|
instance = await self.api_client.wait_until_ready(
|
|
created.instance_id,
|
|
timeout_seconds=self.startup_timeout_seconds,
|
|
status_callback=status_callback,
|
|
)
|
|
except VastInstanceNotFoundError:
|
|
self.registry.remove(created.instance_id)
|
|
raise
|
|
except (TimeoutError, VastApiError) as exc:
|
|
cleanup_completed = await self._discard_startup_lease(
|
|
provisional,
|
|
str(exc),
|
|
)
|
|
raise VastLeaseStartupError(
|
|
created.instance_id,
|
|
offer.offer_id,
|
|
str(exc),
|
|
cleanup_completed=cleanup_completed,
|
|
) from exc
|
|
ready = self._refresh_record(provisional, instance)
|
|
self.registry.upsert(ready)
|
|
logger.info(
|
|
"Rented Vast lease profile=%s instance=%d offer=%d gpu=%s hourly_cost=%.4f idle_hours=%.1f.",
|
|
profile.profile_name,
|
|
ready.instance_id,
|
|
offer.offer_id,
|
|
ready.gpu_name,
|
|
ready.hourly_cost_usd,
|
|
profile.idle_retention_seconds / 3600.0,
|
|
)
|
|
return ready
|
|
|
|
def begin_activity(self, instance_id: int) -> VastLeaseRecord:
|
|
"""Mark one lease busy before staging or execution begins."""
|
|
now = self.clock()
|
|
return self.registry.update(
|
|
instance_id,
|
|
lambda lease: replace(
|
|
lease,
|
|
active_invocations=lease.active_invocations + 1,
|
|
last_activity_at_epoch=now,
|
|
last_error=None,
|
|
),
|
|
)
|
|
|
|
def finish_activity(
|
|
self,
|
|
instance_id: int,
|
|
*,
|
|
idle_retention_seconds: float,
|
|
error: str | None = None,
|
|
) -> VastLeaseRecord:
|
|
"""Release one activity and extend idle destruction from terminal time."""
|
|
if idle_retention_seconds < 0 or not math.isfinite(idle_retention_seconds):
|
|
raise ValueError("idle_retention_seconds must be finite and non-negative.")
|
|
now = self.clock()
|
|
|
|
def finish(lease: VastLeaseRecord) -> VastLeaseRecord:
|
|
"""Return terminal activity state without underflowing the count."""
|
|
if lease.active_invocations <= 0:
|
|
raise RuntimeError(
|
|
f"Vast lease {instance_id} has no active invocation to finish."
|
|
)
|
|
remaining = lease.active_invocations - 1
|
|
return replace(
|
|
lease,
|
|
active_invocations=remaining,
|
|
last_activity_at_epoch=now,
|
|
idle_deadline_epoch=(
|
|
now + idle_retention_seconds
|
|
if remaining == 0
|
|
else lease.idle_deadline_epoch
|
|
),
|
|
last_error=error,
|
|
)
|
|
|
|
return self.registry.update(instance_id, finish)
|
|
|
|
async def destroy_expired(self) -> tuple[int, ...]:
|
|
"""Destroy owned idle leases whose retention deadline has elapsed."""
|
|
now = self.clock()
|
|
expired = [
|
|
lease
|
|
for lease in self.registry.load().leases
|
|
if lease.owner_id == self.owner_id
|
|
and lease.active_invocations == 0
|
|
and lease.idle_deadline_epoch <= now
|
|
]
|
|
destroyed: list[int] = []
|
|
for lease in expired:
|
|
try:
|
|
await self.api_client.destroy_instance(lease.instance_id)
|
|
except VastInstanceNotFoundError:
|
|
pass
|
|
except VastApiError as exc:
|
|
error_message = str(exc)
|
|
self.registry.update(
|
|
lease.instance_id,
|
|
lambda record: replace(record, last_error=error_message),
|
|
)
|
|
logger.warning(
|
|
"Unable to destroy expired Vast lease instance=%d: %s",
|
|
lease.instance_id,
|
|
exc,
|
|
)
|
|
continue
|
|
self.registry.remove(lease.instance_id)
|
|
destroyed.append(lease.instance_id)
|
|
logger.info(
|
|
"Destroyed expired Vast lease instance=%d profile=%s.",
|
|
lease.instance_id,
|
|
lease.profile_name,
|
|
)
|
|
return tuple(destroyed)
|
|
|
|
async def destroy_owned_lease(
|
|
self,
|
|
instance_id: int,
|
|
*,
|
|
allow_active_work: bool = False,
|
|
) -> bool:
|
|
"""Destroy one exact registry-owned lease after checking its API label."""
|
|
lease = next(
|
|
(
|
|
candidate
|
|
for candidate in self.registry.load().leases
|
|
if candidate.instance_id == instance_id
|
|
and candidate.owner_id == self.owner_id
|
|
),
|
|
None,
|
|
)
|
|
if lease is None:
|
|
return False
|
|
if lease.active_invocations and not allow_active_work:
|
|
raise RuntimeError(
|
|
f"Vast lease {instance_id} has active work and cannot be destroyed."
|
|
)
|
|
try:
|
|
instance = await self.api_client.show_instance(instance_id)
|
|
except VastInstanceNotFoundError:
|
|
self.registry.remove(instance_id)
|
|
return True
|
|
if instance.label != lease.label:
|
|
raise RuntimeError(
|
|
f"Vast instance {instance_id} label no longer matches its managed lease."
|
|
)
|
|
try:
|
|
await self.api_client.destroy_instance(instance_id)
|
|
except VastInstanceNotFoundError:
|
|
logger.info(
|
|
"Vast lease disappeared between ownership verification and "
|
|
"destruction instance=%d.",
|
|
instance_id,
|
|
)
|
|
self.registry.remove(instance_id)
|
|
return True
|
|
|
|
async def refresh_owned_leases(
|
|
self,
|
|
*,
|
|
clear_stale_activity: bool = False,
|
|
) -> tuple[VastLeaseRecord, ...]:
|
|
"""Refresh managed inventory without normally changing live activity counts."""
|
|
visible_instances = {
|
|
instance.instance_id: instance
|
|
for instance in await self.api_client.list_instances()
|
|
}
|
|
reconciled: list[VastLeaseRecord] = []
|
|
for lease in self.registry.load().leases:
|
|
if lease.owner_id != self.owner_id:
|
|
reconciled.append(lease)
|
|
continue
|
|
instance = visible_instances.get(lease.instance_id)
|
|
if instance is None:
|
|
self.registry.remove(lease.instance_id)
|
|
continue
|
|
if instance.label != lease.label:
|
|
mismatched = replace(
|
|
lease,
|
|
draining=True,
|
|
last_error="Vast instance label no longer matches managed state.",
|
|
)
|
|
self.registry.upsert(mismatched)
|
|
reconciled.append(mismatched)
|
|
continue
|
|
refreshed = replace(
|
|
self._refresh_record(lease, instance),
|
|
active_invocations=(
|
|
0 if clear_stale_activity else lease.active_invocations
|
|
),
|
|
)
|
|
self.registry.upsert(refreshed)
|
|
reconciled.append(refreshed)
|
|
return tuple(reconciled)
|
|
|
|
async def reconcile(self) -> tuple[VastLeaseRecord, ...]:
|
|
"""Refresh leases after restart and clear activity owned by the dead process."""
|
|
return await self.refresh_owned_leases(clear_stale_activity=True)
|
|
|
|
def incremental_retention_cost_usd(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
*,
|
|
existing_lease: VastLeaseRecord | None,
|
|
predicted_execution_seconds: float,
|
|
) -> float:
|
|
"""Estimate compute cost added by a job and its resulting retention deadline."""
|
|
if predicted_execution_seconds < 0:
|
|
raise ValueError("predicted_execution_seconds must not be negative.")
|
|
hourly_cost = (
|
|
existing_lease.hourly_cost_usd
|
|
if existing_lease is not None
|
|
else profile.maximum_hourly_cost_usd
|
|
)
|
|
new_deadline = (
|
|
self.clock()
|
|
+ predicted_execution_seconds
|
|
+ profile.idle_retention_seconds
|
|
)
|
|
covered_until = (
|
|
existing_lease.idle_deadline_epoch
|
|
if existing_lease is not None
|
|
else self.clock()
|
|
)
|
|
incremental_seconds = max(0.0, new_deadline - covered_until)
|
|
return incremental_seconds / 3600.0 * hourly_cost
|
|
|
|
def _refresh_record(
|
|
self,
|
|
lease: VastLeaseRecord,
|
|
instance: VastInstance,
|
|
) -> VastLeaseRecord:
|
|
"""Apply current API status to one persisted record."""
|
|
proxy_selected = (
|
|
lease.ssh_proxy_host is not None
|
|
and lease.ssh_host == lease.ssh_proxy_host
|
|
and lease.ssh_port == lease.ssh_proxy_port
|
|
)
|
|
selected_host = (
|
|
instance.ssh_proxy_host if proxy_selected else instance.ssh_host
|
|
)
|
|
selected_port = (
|
|
instance.ssh_proxy_port if proxy_selected else instance.ssh_port
|
|
)
|
|
return replace(
|
|
lease,
|
|
actual_status=instance.actual_status,
|
|
ssh_host=selected_host or instance.ssh_host,
|
|
ssh_port=selected_port or instance.ssh_port,
|
|
ssh_direct_host=instance.ssh_direct_host,
|
|
ssh_direct_port=instance.ssh_direct_port,
|
|
ssh_proxy_host=instance.ssh_proxy_host,
|
|
ssh_proxy_port=instance.ssh_proxy_port,
|
|
gpu_name=instance.gpu_name or lease.gpu_name,
|
|
gpu_count=instance.num_gpus or lease.gpu_count,
|
|
gpu_ram_mb=instance.gpu_ram_mb or lease.gpu_ram_mb,
|
|
cpu_ram_mb=instance.cpu_ram_mb or lease.cpu_ram_mb,
|
|
hourly_cost_usd=(
|
|
instance.hourly_cost_usd
|
|
if instance.hourly_cost_usd is not None
|
|
else lease.hourly_cost_usd
|
|
),
|
|
last_error=None,
|
|
)
|
|
|
|
def _profile_lock(self, profile_id: str) -> asyncio.Lock:
|
|
"""Return the process-local acquisition lock for one profile slot."""
|
|
with self._profile_locks_guard:
|
|
lock = self._profile_locks.get(profile_id)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
self._profile_locks[profile_id] = lock
|
|
return lock
|
|
|
|
|
|
def vast_profile_fingerprint(profile: VastResourceProfile) -> str:
|
|
"""Return a deterministic fingerprint for all lease-shaping profile fields."""
|
|
payload = asdict(profile)
|
|
payload["rental_type"] = profile.rental_type.value
|
|
serialized = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode(
|
|
"utf-8"
|
|
)
|
|
return hashlib.sha256(serialized).hexdigest()
|
|
|
|
|
|
def vast_managed_label_prefix(owner_id: str) -> str:
|
|
"""Return the exact label prefix owned by one ComfyUI installation."""
|
|
owner_slug = _safe_label_part(owner_id, maximum_length=32)
|
|
return f"{VAST_MANAGED_LABEL_PREFIX}:{owner_slug}"
|
|
|
|
|
|
def vast_managed_label(
|
|
owner_id: str,
|
|
profile_id: str,
|
|
profile_fingerprint: str,
|
|
runtime_fingerprint: str,
|
|
) -> str:
|
|
"""Return a bounded ownership label for one profile/runtime combination."""
|
|
profile_slug = _safe_label_part(profile_id, maximum_length=24)
|
|
return (
|
|
f"{vast_managed_label_prefix(owner_id)}:{profile_slug}:"
|
|
f"{profile_fingerprint[:10]}:{runtime_fingerprint[:10]}"
|
|
)
|
|
|
|
|
|
def _safe_label_part(value: str, *, maximum_length: int) -> str:
|
|
"""Return one non-empty Vast label component."""
|
|
normalized = _SAFE_LABEL_PART_PATTERN.sub("-", value.strip()).strip("-._")
|
|
if not normalized:
|
|
normalized = hashlib.sha256(value.encode("utf-8")).hexdigest()[:12]
|
|
return normalized[:maximum_length]
|
|
|
|
|
|
def _slot_profile_id(profile_id: str, slot: int) -> str:
|
|
"""Return a deterministic profile identity for one capacity slot."""
|
|
return profile_id if slot == 0 else f"{profile_id}-slot-{slot}"
|
|
|
|
|
|
__all__ = [
|
|
"VAST_LEASE_REGISTRY_FILENAME",
|
|
"VAST_LEASE_REGISTRY_VERSION",
|
|
"VAST_MANAGED_LABEL_PREFIX",
|
|
"VastLeaseManager",
|
|
"VastLeaseRecord",
|
|
"VastLeaseRegistry",
|
|
"VastLeaseRegistryState",
|
|
"VastLeaseStartupError",
|
|
"VastLeaseStartupExhaustedError",
|
|
"vast_managed_label",
|
|
"vast_managed_label_prefix",
|
|
"vast_profile_fingerprint",
|
|
]
|