988 lines
38 KiB
Python
988 lines
38 KiB
Python
"""Application service joining Vast marketplace, leases, SSH, sync, and execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass, replace
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Callable, Iterator, Mapping, Sequence
|
|
|
|
if __package__:
|
|
from .execution_environments import (
|
|
EnvironmentCapabilities,
|
|
EnvironmentHealth,
|
|
EnvironmentSchedulingState,
|
|
ExecutionProvider,
|
|
GpuCapability,
|
|
)
|
|
from .huggingface_assets import HuggingFaceAssetRegistry
|
|
from .huggingface_discovery import HuggingFaceAssetDiscovery
|
|
from .runtime_environment import build_remote_runtime_identity
|
|
from .r2_cache import R2CacheClient
|
|
from .settings import ModalSyncSettings, discover_comfyui_user_directory
|
|
from .sync_engine import ModalAssetSyncEngine
|
|
from .vast_api import VastApiClient, VastApiError, VastInstanceNotFoundError
|
|
from .vast_executor import VastExecutorClient
|
|
from .vast_image_build import (
|
|
VAST_IMAGE_BUILD_COMMAND,
|
|
VastWorkerImageBuildError,
|
|
VastWorkerImageBuilder,
|
|
)
|
|
from .vast_leases import (
|
|
VastLeaseManager,
|
|
VastLeaseRecord,
|
|
VastLeaseRegistry,
|
|
vast_profile_fingerprint,
|
|
)
|
|
from .vast_models import VastInstance, VastOffer, VastResourceProfile
|
|
from .vast_runtime import (
|
|
VAST_IMAGE_ENV,
|
|
VastRuntimeConfiguration,
|
|
VastRuntimeFingerprintDriftError,
|
|
VastRuntimeManager,
|
|
)
|
|
from .vast_ssh import (
|
|
VastSshError,
|
|
VastSshRunner,
|
|
VastSshVolumeBackend,
|
|
vast_connection_from_lease,
|
|
)
|
|
else: # pragma: no cover - direct debugging imports.
|
|
from execution_environments import (
|
|
EnvironmentCapabilities,
|
|
EnvironmentHealth,
|
|
EnvironmentSchedulingState,
|
|
ExecutionProvider,
|
|
GpuCapability,
|
|
)
|
|
from huggingface_assets import HuggingFaceAssetRegistry
|
|
from huggingface_discovery import HuggingFaceAssetDiscovery
|
|
from runtime_environment import build_remote_runtime_identity
|
|
from r2_cache import R2CacheClient
|
|
from settings import ModalSyncSettings, discover_comfyui_user_directory
|
|
from sync_engine import ModalAssetSyncEngine
|
|
from vast_api import VastApiClient, VastApiError, VastInstanceNotFoundError
|
|
from vast_executor import VastExecutorClient
|
|
from vast_image_build import (
|
|
VAST_IMAGE_BUILD_COMMAND,
|
|
VastWorkerImageBuildError,
|
|
VastWorkerImageBuilder,
|
|
)
|
|
from vast_leases import (
|
|
VastLeaseManager,
|
|
VastLeaseRecord,
|
|
VastLeaseRegistry,
|
|
vast_profile_fingerprint,
|
|
)
|
|
from vast_models import VastInstance, VastOffer, VastResourceProfile
|
|
from vast_runtime import (
|
|
VAST_IMAGE_ENV,
|
|
VastRuntimeConfiguration,
|
|
VastRuntimeFingerprintDriftError,
|
|
VastRuntimeManager,
|
|
)
|
|
from vast_ssh import (
|
|
VastSshError,
|
|
VastSshRunner,
|
|
VastSshVolumeBackend,
|
|
vast_connection_from_lease,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
VAST_API_KEY_ENV = "VAST_API_KEY"
|
|
VAST_API_BASE_URL_ENV = "COMFY_MODAL_VAST_API_BASE_URL"
|
|
VAST_SSH_IDENTITY_FILE_ENV = "COMFY_MODAL_VAST_SSH_IDENTITY_FILE"
|
|
VAST_OFFER_PREFETCH_CONCURRENCY = 8
|
|
VAST_INSTANCE_SETUP_REPLACEMENTS = 1
|
|
|
|
|
|
def _vast_startup_status_message(instance: VastInstance) -> str:
|
|
"""Translate one provider lifecycle record into stable user-facing progress."""
|
|
provider_message = (instance.status_message or "").casefold()
|
|
instance_label = f"Vast.ai instance {instance.instance_id}"
|
|
if "unauthorized" in provider_message or "denied" in provider_message:
|
|
return f"{instance_label} could not download the worker image"
|
|
if "download complete" in provider_message:
|
|
return f"{instance_label} is downloading the worker image (layer complete)"
|
|
if any(word in provider_message for word in ("download", "pull", "extract")):
|
|
return f"{instance_label} is downloading the worker image"
|
|
if instance.lifecycle_status.casefold() == "running":
|
|
return f"{instance_label} is waiting for SSH access"
|
|
return f"{instance_label} is starting ({instance.lifecycle_status})"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastProfileQuote:
|
|
"""Describe the least effective-hourly-price profile for one component."""
|
|
|
|
profile: VastResourceProfile
|
|
offer: VastOffer | None
|
|
existing_lease: VastLeaseRecord | None
|
|
predicted_incremental_cost_usd: float
|
|
retention_seconds_charged: float
|
|
predicted_execution_seconds: float
|
|
|
|
@property
|
|
def hourly_cost_usd(self) -> float:
|
|
"""Return the live lease or marketplace hourly price."""
|
|
if self.existing_lease is not None:
|
|
return self.existing_lease.hourly_cost_usd
|
|
if self.offer is None:
|
|
raise RuntimeError("Vast quote has neither an offer nor a lease.")
|
|
return self.offer.hourly_cost_usd
|
|
|
|
@property
|
|
def predicted_retention_cost_usd(self) -> float:
|
|
"""Return the incremental bill estimate including configured retention."""
|
|
return self.retention_seconds_charged / 3600.0 * self.hourly_cost_usd
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastSearchRequirements:
|
|
"""Describe component memory floors that affect a marketplace search."""
|
|
|
|
minimum_vram_bytes: int
|
|
minimum_ram_bytes: int
|
|
|
|
|
|
class VastService:
|
|
"""Provide one credential-safe controller for workflow Vast operations."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
settings: ModalSyncSettings,
|
|
repo_root: Path,
|
|
user_directory: Path,
|
|
api_client: VastApiClient,
|
|
runtime_configuration: VastRuntimeConfiguration,
|
|
registry: VastLeaseRegistry,
|
|
identity_file: Path | None = None,
|
|
r2_cache: R2CacheClient | None = None,
|
|
image_builder: VastWorkerImageBuilder | None = None,
|
|
update_process_environment: bool = False,
|
|
) -> None:
|
|
"""Initialize the controller and its persistent lease manager."""
|
|
self.settings = settings
|
|
self.repo_root = repo_root.resolve()
|
|
self.user_directory = user_directory.resolve()
|
|
self.api_client = api_client
|
|
self.runtime_configuration = runtime_configuration
|
|
self.registry = registry
|
|
self.identity_file = identity_file
|
|
self.r2_cache = r2_cache
|
|
self.image_builder = image_builder or VastWorkerImageBuilder(
|
|
repo_root=self.repo_root,
|
|
comfyui_root=getattr(settings, "comfyui_root", None),
|
|
modal_gpu=settings.modal_gpu,
|
|
)
|
|
self.update_process_environment = update_process_environment
|
|
self.huggingface_asset_registry = (
|
|
HuggingFaceAssetRegistry.for_user_directory(self.user_directory)
|
|
)
|
|
self.huggingface_asset_discovery = HuggingFaceAssetDiscovery(
|
|
registry=self.huggingface_asset_registry,
|
|
user_directory=self.user_directory,
|
|
comfyui_root=getattr(settings, "comfyui_root", None),
|
|
)
|
|
self.lease_manager = self._lease_manager_for(runtime_configuration)
|
|
|
|
def _lease_manager_for(
|
|
self,
|
|
runtime_configuration: VastRuntimeConfiguration,
|
|
) -> VastLeaseManager:
|
|
"""Return a lease manager bound to one immutable worker image."""
|
|
return VastLeaseManager(
|
|
api_client=self.api_client,
|
|
registry=self.registry,
|
|
owner_id=self.settings.app_name,
|
|
runtime_fingerprint=runtime_configuration.runtime_fingerprint,
|
|
launch_spec_factory=runtime_configuration.launch_spec,
|
|
startup_timeout_seconds=runtime_configuration.startup_timeout_seconds,
|
|
worker_image=runtime_configuration.image,
|
|
)
|
|
|
|
@classmethod
|
|
def from_environment(
|
|
cls,
|
|
settings: ModalSyncSettings,
|
|
*,
|
|
repo_root: Path,
|
|
environment: Mapping[str, str] | None = None,
|
|
) -> "VastService":
|
|
"""Resolve credentials, runtime identity, and local state from the environment."""
|
|
source = os.environ if environment is None else environment
|
|
api_key = str(source.get(VAST_API_KEY_ENV) or "").strip()
|
|
if not api_key:
|
|
raise RuntimeError(
|
|
f"Set {VAST_API_KEY_ENV} before selecting Vast.ai execution."
|
|
)
|
|
user_directory = discover_comfyui_user_directory(settings)
|
|
if user_directory is None:
|
|
raise RuntimeError(
|
|
"Vast.ai execution requires a persistent ComfyUI user directory."
|
|
)
|
|
identity = build_remote_runtime_identity(
|
|
repo_root=repo_root,
|
|
comfyui_root=settings.comfyui_root,
|
|
custom_nodes_dir=settings.custom_nodes_dir,
|
|
settings=settings,
|
|
)
|
|
runtime = VastRuntimeConfiguration.from_environment(
|
|
identity.fingerprint,
|
|
environment=source,
|
|
)
|
|
base_url = str(source.get(VAST_API_BASE_URL_ENV) or "").strip()
|
|
api_client = (
|
|
VastApiClient(api_key, base_url=base_url)
|
|
if base_url
|
|
else VastApiClient(api_key)
|
|
)
|
|
raw_identity_file = str(source.get(VAST_SSH_IDENTITY_FILE_ENV) or "").strip()
|
|
identity_file = (
|
|
Path(raw_identity_file).expanduser().resolve()
|
|
if raw_identity_file
|
|
else None
|
|
)
|
|
return cls(
|
|
settings=settings,
|
|
repo_root=repo_root,
|
|
user_directory=user_directory,
|
|
api_client=api_client,
|
|
runtime_configuration=runtime,
|
|
registry=VastLeaseRegistry.for_user_directory(user_directory),
|
|
identity_file=identity_file,
|
|
r2_cache=R2CacheClient.from_environment(source),
|
|
image_builder=VastWorkerImageBuilder(
|
|
repo_root=repo_root.resolve(),
|
|
comfyui_root=settings.comfyui_root,
|
|
modal_gpu=settings.modal_gpu,
|
|
environment=(source if environment is not None else None),
|
|
),
|
|
update_process_environment=environment is None,
|
|
)
|
|
|
|
async def quote_best_profile(
|
|
self,
|
|
profiles: Sequence[VastResourceProfile],
|
|
*,
|
|
minimum_vram_bytes: int,
|
|
minimum_ram_bytes: int,
|
|
predicted_execution_seconds: float,
|
|
) -> VastProfileQuote:
|
|
"""Return the cheapest compatible existing lease or current offer."""
|
|
if not profiles:
|
|
raise ValueError("At least one Vast resource profile is required.")
|
|
adjusted_profiles = tuple(
|
|
_profile_for_requirements(profile, minimum_vram_bytes, minimum_ram_bytes)
|
|
for profile in profiles
|
|
)
|
|
quotes = await asyncio.gather(
|
|
*(
|
|
self._quote_profile(profile, predicted_execution_seconds)
|
|
for profile in adjusted_profiles
|
|
),
|
|
return_exceptions=True,
|
|
)
|
|
candidates = [quote for quote in quotes if isinstance(quote, VastProfileQuote)]
|
|
if candidates:
|
|
return min(
|
|
candidates,
|
|
key=lambda quote: (
|
|
quote.hourly_cost_usd,
|
|
quote.profile.profile_id,
|
|
),
|
|
)
|
|
errors = [str(error) for error in quotes if isinstance(error, BaseException)]
|
|
raise VastApiError(
|
|
"No Vast.ai profile currently has a compatible offer: "
|
|
+ "; ".join(errors)
|
|
)
|
|
|
|
async def prefetch_offers(
|
|
self,
|
|
profiles: Sequence[VastResourceProfile],
|
|
requirements: Sequence[VastSearchRequirements],
|
|
) -> None:
|
|
"""Populate cached effective-profile searches with bounded parallelism."""
|
|
if not profiles or not requirements:
|
|
return
|
|
pending_profiles = self._distinct_marketplace_profiles(
|
|
profiles,
|
|
requirements,
|
|
)
|
|
if not pending_profiles:
|
|
return
|
|
logger.info(
|
|
"Prefetching %d unique Vast marketplace search(es) in parallel for "
|
|
"%d component requirement set(s).",
|
|
len(pending_profiles),
|
|
len(requirements),
|
|
)
|
|
results = await self._search_offer_profiles_in_parallel(pending_profiles)
|
|
cancelled = next(
|
|
(
|
|
result
|
|
for result in results
|
|
if isinstance(result, asyncio.CancelledError)
|
|
),
|
|
None,
|
|
)
|
|
if cancelled is not None:
|
|
raise cancelled
|
|
failure_count = sum(
|
|
isinstance(result, (OSError, RuntimeError, ValueError))
|
|
for result in results
|
|
)
|
|
logger.info(
|
|
"Completed Vast marketplace prefetch searches=%d successful=%d "
|
|
"failed=%d.",
|
|
len(results),
|
|
len(results) - failure_count,
|
|
failure_count,
|
|
)
|
|
|
|
def _distinct_marketplace_profiles(
|
|
self,
|
|
profiles: Sequence[VastResourceProfile],
|
|
requirements: Sequence[VastSearchRequirements],
|
|
) -> tuple[VastResourceProfile, ...]:
|
|
"""Return deduplicated effective profiles without reusable leases."""
|
|
unique_profiles: dict[str, VastResourceProfile] = {}
|
|
for requirement in requirements:
|
|
for profile in profiles:
|
|
adjusted_profile = _profile_for_requirements(
|
|
profile,
|
|
requirement.minimum_vram_bytes,
|
|
requirement.minimum_ram_bytes,
|
|
)
|
|
if self._existing_lease(adjusted_profile) is not None:
|
|
continue
|
|
query_key = json.dumps(
|
|
adjusted_profile.search_payload(),
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
)
|
|
unique_profiles.setdefault(query_key, adjusted_profile)
|
|
return tuple(unique_profiles.values())
|
|
|
|
async def _search_offer_profiles_in_parallel(
|
|
self,
|
|
profiles: Sequence[VastResourceProfile],
|
|
) -> tuple[object, ...]:
|
|
"""Search effective profiles in batches that respect the API bound."""
|
|
results: list[object] = []
|
|
for offset in range(0, len(profiles), VAST_OFFER_PREFETCH_CONCURRENCY):
|
|
batch = profiles[
|
|
offset : offset + VAST_OFFER_PREFETCH_CONCURRENCY
|
|
]
|
|
results.extend(
|
|
await asyncio.gather(
|
|
*(self.api_client.search_offers(profile) for profile in batch),
|
|
return_exceptions=True,
|
|
)
|
|
)
|
|
return tuple(results)
|
|
|
|
async def _quote_profile(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
predicted_execution_seconds: float,
|
|
) -> VastProfileQuote:
|
|
"""Quote compute at its hourly rate while tracking retention separately."""
|
|
existing = self._existing_lease(profile)
|
|
if existing is not None:
|
|
retention_cost = self.lease_manager.incremental_retention_cost_usd(
|
|
profile,
|
|
existing_lease=existing,
|
|
predicted_execution_seconds=predicted_execution_seconds,
|
|
)
|
|
charged_seconds = (
|
|
retention_cost / existing.hourly_cost_usd * 3600.0
|
|
if existing.hourly_cost_usd > 0
|
|
else 0.0
|
|
)
|
|
return VastProfileQuote(
|
|
profile=profile,
|
|
offer=None,
|
|
existing_lease=existing,
|
|
predicted_incremental_cost_usd=(
|
|
predicted_execution_seconds
|
|
/ 3600.0
|
|
* existing.hourly_cost_usd
|
|
),
|
|
retention_seconds_charged=charged_seconds,
|
|
predicted_execution_seconds=predicted_execution_seconds,
|
|
)
|
|
offers = await self.api_client.search_offers(profile)
|
|
if not offers:
|
|
raise VastApiError(
|
|
f"No offer satisfies Vast profile {profile.profile_name!r}."
|
|
)
|
|
offer = offers[0]
|
|
charged_seconds = predicted_execution_seconds + profile.idle_retention_seconds
|
|
return VastProfileQuote(
|
|
profile=profile,
|
|
offer=offer,
|
|
existing_lease=None,
|
|
predicted_incremental_cost_usd=(
|
|
predicted_execution_seconds / 3600.0 * offer.hourly_cost_usd
|
|
),
|
|
retention_seconds_charged=charged_seconds,
|
|
predicted_execution_seconds=predicted_execution_seconds,
|
|
)
|
|
|
|
async def acquire(
|
|
self,
|
|
quote: VastProfileQuote,
|
|
*,
|
|
slot: int = 0,
|
|
status_callback: Callable[[str], None] | None = None,
|
|
) -> VastLeaseRecord:
|
|
"""Reuse or rent the selected quoted profile capacity slot."""
|
|
if quote.existing_lease is None:
|
|
await self._preflight_published_image(status_callback=status_callback)
|
|
if status_callback is not None:
|
|
status_callback("Requesting Vast.ai capacity")
|
|
|
|
def emit_instance_status(instance: VastInstance) -> None:
|
|
"""Forward normalized provider progress when a caller is listening."""
|
|
if status_callback is not None:
|
|
status_callback(_vast_startup_status_message(instance))
|
|
|
|
return await self._acquire_with_replacement(
|
|
quote.profile,
|
|
slot=slot,
|
|
status_callback=status_callback,
|
|
instance_status_callback=(
|
|
emit_instance_status if status_callback is not None else None
|
|
),
|
|
)
|
|
|
|
async def _preflight_published_image(
|
|
self,
|
|
*,
|
|
status_callback: Callable[[str], None] | None,
|
|
) -> None:
|
|
"""Publish current runtime source before any billable Vast capacity request."""
|
|
configured_image = self.runtime_configuration.image
|
|
expected_fingerprint = self.runtime_configuration.runtime_fingerprint
|
|
image = await asyncio.to_thread(
|
|
self.image_builder.ensure_published_image,
|
|
configured_image,
|
|
expected_fingerprint,
|
|
status_callback=status_callback,
|
|
)
|
|
if image == configured_image:
|
|
return
|
|
self._adopt_runtime_image(image)
|
|
if status_callback is not None:
|
|
status_callback("Vast worker image updated before requesting capacity")
|
|
|
|
async def _acquire_with_replacement(
|
|
self,
|
|
profile: VastResourceProfile,
|
|
*,
|
|
slot: int,
|
|
status_callback: Callable[[str], None] | None,
|
|
instance_status_callback: Callable[[VastInstance], None] | None,
|
|
allow_image_rebuild: bool = True,
|
|
) -> VastLeaseRecord:
|
|
"""Initialize capacity and replace one disappeared or unusable contract."""
|
|
excluded_offer_ids: set[int] = set()
|
|
for replacement_attempt in range(
|
|
VAST_INSTANCE_SETUP_REPLACEMENTS + 1
|
|
):
|
|
lease = await self.lease_manager.ensure_lease(
|
|
profile,
|
|
slot=slot,
|
|
status_callback=instance_status_callback,
|
|
event_callback=status_callback,
|
|
excluded_offer_ids=frozenset(excluded_offer_ids),
|
|
)
|
|
if status_callback is not None:
|
|
status_callback("Initializing Vast.ai worker")
|
|
try:
|
|
initialized_lease = await asyncio.to_thread(
|
|
self._initialize_runtime,
|
|
lease,
|
|
status_callback=status_callback,
|
|
)
|
|
if initialized_lease is not None:
|
|
lease = initialized_lease
|
|
except VastRuntimeFingerprintDriftError as exc:
|
|
cleanup_completed = await self._discard_failed_runtime_lease(
|
|
lease,
|
|
str(exc),
|
|
)
|
|
if status_callback is not None:
|
|
cleanup_status = (
|
|
"terminated" if cleanup_completed else "marked unusable"
|
|
)
|
|
status_callback(
|
|
"Vast worker source drift detected; stale instance "
|
|
f"{lease.instance_id} was {cleanup_status}"
|
|
)
|
|
if not allow_image_rebuild:
|
|
raise RuntimeError(
|
|
self._manual_image_build_message(
|
|
"The automatically published worker image still reported "
|
|
f"fingerprint {exc.actual_fingerprint[:12]} instead of "
|
|
f"{exc.expected_fingerprint[:12]}."
|
|
)
|
|
) from exc
|
|
try:
|
|
image = self.image_builder.build_and_push(
|
|
exc.expected_fingerprint,
|
|
status_callback=status_callback,
|
|
)
|
|
except VastWorkerImageBuildError:
|
|
if status_callback is not None:
|
|
status_callback("Automatic Vast worker image build failed")
|
|
raise
|
|
self._adopt_runtime_image(image)
|
|
if status_callback is not None:
|
|
status_callback(
|
|
"Vast worker image updated; requesting fresh capacity"
|
|
)
|
|
return await self._acquire_with_replacement(
|
|
profile,
|
|
slot=slot,
|
|
status_callback=status_callback,
|
|
instance_status_callback=instance_status_callback,
|
|
allow_image_rebuild=False,
|
|
)
|
|
except VastInstanceNotFoundError as exc:
|
|
self.registry.remove(lease.instance_id)
|
|
excluded_offer_ids.add(lease.offer_id)
|
|
if replacement_attempt >= VAST_INSTANCE_SETUP_REPLACEMENTS:
|
|
if status_callback is not None:
|
|
status_callback("Vast.ai worker initialization failed")
|
|
raise RuntimeError(
|
|
"Vast.ai worker setup failed after the replacement "
|
|
f"instance also disappeared. Last instance: "
|
|
f"{lease.instance_id}."
|
|
) from exc
|
|
logger.warning(
|
|
"Vast instance disappeared before worker initialization "
|
|
"instance_id=%d offer=%d; cold-starting a replacement.",
|
|
lease.instance_id,
|
|
lease.offer_id,
|
|
)
|
|
if status_callback is not None:
|
|
status_callback(
|
|
"Vast.ai instance disappeared; requesting a replacement"
|
|
)
|
|
continue
|
|
except (TimeoutError, VastSshError, ValueError) as exc:
|
|
cleanup_completed = await self._discard_failed_runtime_lease(
|
|
lease,
|
|
str(exc),
|
|
)
|
|
excluded_offer_ids.add(lease.offer_id)
|
|
if replacement_attempt >= VAST_INSTANCE_SETUP_REPLACEMENTS:
|
|
if status_callback is not None:
|
|
status_callback("Vast.ai worker initialization failed")
|
|
cleanup_summary = (
|
|
"was destroyed"
|
|
if cleanup_completed
|
|
else "could not be confirmed destroyed"
|
|
)
|
|
raise RuntimeError(
|
|
"Vast.ai worker setup failed after a replacement attempt. "
|
|
f"Last instance {lease.instance_id} {cleanup_summary}. "
|
|
f"Last error: {exc}"
|
|
) from exc
|
|
logger.warning(
|
|
"Vast instance failed worker initialization instance_id=%d "
|
|
"offer=%d cleanup_completed=%s; cold-starting a replacement: %s",
|
|
lease.instance_id,
|
|
lease.offer_id,
|
|
cleanup_completed,
|
|
exc,
|
|
)
|
|
if status_callback is not None:
|
|
cleanup_status = (
|
|
"terminated"
|
|
if cleanup_completed
|
|
else "marked unusable"
|
|
)
|
|
status_callback(
|
|
f"Vast.ai instance {lease.instance_id} failed worker setup "
|
|
f"and was {cleanup_status}; requesting a replacement"
|
|
)
|
|
continue
|
|
if status_callback is not None:
|
|
status_callback("Vast.ai worker is ready")
|
|
return lease
|
|
raise RuntimeError("Vast replacement attempts ended without a lease.")
|
|
|
|
def _adopt_runtime_image(self, image: str) -> None:
|
|
"""Use one freshly published image for this and later process-local runs."""
|
|
configuration = replace(self.runtime_configuration, image=image)
|
|
self.runtime_configuration = configuration
|
|
self.lease_manager = self._lease_manager_for(configuration)
|
|
if self.update_process_environment:
|
|
os.environ[VAST_IMAGE_ENV] = image
|
|
logger.info(
|
|
"Adopted rebuilt Vast worker image fingerprint=%s.",
|
|
configuration.runtime_fingerprint[:12],
|
|
)
|
|
|
|
def _manual_image_build_message(self, cause: str) -> str:
|
|
"""Return the manual recovery required after an unusable rebuilt image."""
|
|
return (
|
|
f"{cause} Run `{' '.join(VAST_IMAGE_BUILD_COMMAND)}` from "
|
|
f"{self.repo_root}, set {VAST_IMAGE_ENV} to the printed digest, restart "
|
|
"ComfyUI, and retry the workflow."
|
|
)
|
|
|
|
async def _discard_failed_runtime_lease(
|
|
self,
|
|
lease: VastLeaseRecord,
|
|
error_message: str,
|
|
) -> bool:
|
|
"""Drain and destroy capacity that never produced a usable worker."""
|
|
self.registry.update(
|
|
lease.instance_id,
|
|
lambda current: replace(
|
|
current,
|
|
draining=True,
|
|
last_error=error_message,
|
|
),
|
|
)
|
|
try:
|
|
destroyed = await self.lease_manager.destroy_owned_lease(lease.instance_id)
|
|
except (RuntimeError, ValueError, VastApiError) as cleanup_error:
|
|
logger.warning(
|
|
"Unable to destroy unusable Vast worker instance_id=%d: %s",
|
|
lease.instance_id,
|
|
cleanup_error,
|
|
)
|
|
return False
|
|
if not destroyed:
|
|
logger.warning(
|
|
"Unusable Vast worker was not destroyed because ownership could not "
|
|
"be confirmed instance_id=%d.",
|
|
lease.instance_id,
|
|
)
|
|
return False
|
|
logger.info(
|
|
"Destroyed unusable Vast worker instance_id=%d after initialization failure.",
|
|
lease.instance_id,
|
|
)
|
|
return True
|
|
|
|
def quote_best_profile_sync(self, *args: object, **kwargs: object) -> VastProfileQuote:
|
|
"""Synchronously quote a profile from ComfyUI's queue worker thread."""
|
|
return asyncio.run(self.quote_best_profile(*args, **kwargs)) # type: ignore[arg-type]
|
|
|
|
def prefetch_offers_sync(
|
|
self,
|
|
profiles: Sequence[VastResourceProfile],
|
|
requirements: Sequence[VastSearchRequirements],
|
|
) -> None:
|
|
"""Synchronously prefetch searches from ComfyUI's queue worker thread."""
|
|
asyncio.run(self.prefetch_offers(profiles, requirements))
|
|
|
|
def acquire_sync(
|
|
self,
|
|
quote: VastProfileQuote,
|
|
*,
|
|
slot: int = 0,
|
|
status_callback: Callable[[str], None] | None = None,
|
|
) -> VastLeaseRecord:
|
|
"""Synchronously acquire a quote slot from ComfyUI's queue worker thread."""
|
|
return asyncio.run(
|
|
self.acquire(
|
|
quote,
|
|
slot=slot,
|
|
status_callback=status_callback,
|
|
)
|
|
)
|
|
|
|
def scheduling_state(self, quote: VastProfileQuote) -> EnvironmentSchedulingState:
|
|
"""Expose a quote to the provider-neutral cost scheduler."""
|
|
offer = quote.offer
|
|
lease = quote.existing_lease
|
|
gpu_name = lease.gpu_name if lease is not None else offer.gpu_name if offer else "Vast GPU"
|
|
gpu_count = lease.gpu_count if lease is not None else offer.num_gpus if offer else 1
|
|
gpu_ram_mb = lease.gpu_ram_mb if lease is not None else offer.gpu_ram_mb if offer else 0
|
|
cpu_ram_mb = lease.cpu_ram_mb if lease is not None else offer.cpu_ram_mb if offer else 0
|
|
environment_id = lease.environment_id if lease is not None else quote.profile.environment_id
|
|
return EnvironmentSchedulingState(
|
|
environment_id=environment_id,
|
|
provider=ExecutionProvider.VAST,
|
|
enabled=True,
|
|
health=EnvironmentHealth.READY,
|
|
cost_usd_per_second=quote.hourly_cost_usd / 3600.0,
|
|
capabilities=EnvironmentCapabilities(
|
|
architecture="x86_64",
|
|
operating_system="linux",
|
|
cpu_count=max(
|
|
1,
|
|
math.ceil(quote.profile.minimum_cpu_cores or 1.0),
|
|
),
|
|
total_ram_bytes=cpu_ram_mb * 1024**2,
|
|
available_ram_bytes=None,
|
|
available_disk_bytes=int(quote.profile.allocated_disk_gb * 1024**3),
|
|
docker_version="vast-container",
|
|
docker_rootless=False,
|
|
nvidia_container_runtime=True,
|
|
gpus=tuple(
|
|
GpuCapability(
|
|
uuid=f"vast-{environment_id}-{index}",
|
|
name=gpu_name,
|
|
total_vram_bytes=gpu_ram_mb * 1024**2,
|
|
)
|
|
for index in range(gpu_count)
|
|
),
|
|
),
|
|
maximum_workers=1,
|
|
cold_start_seconds=0.0,
|
|
)
|
|
|
|
def lease_for_environment_id(self, environment_id: str) -> VastLeaseRecord:
|
|
"""Resolve a concrete managed lease from an assigned environment ID."""
|
|
lease = next(
|
|
(
|
|
record
|
|
for record in self.registry.load().leases
|
|
if record.environment_id == environment_id
|
|
),
|
|
None,
|
|
)
|
|
if lease is None:
|
|
raise KeyError(f"Unknown managed Vast environment {environment_id!r}.")
|
|
return lease
|
|
|
|
def sync_engine(self, lease: VastLeaseRecord) -> ModalAssetSyncEngine:
|
|
"""Build a direct-filesystem content-addressed sync engine for one lease."""
|
|
runner = self._runner(lease)
|
|
runtime = VastRuntimeManager(
|
|
runner=runner,
|
|
configuration=self.runtime_configuration,
|
|
)
|
|
|
|
@contextmanager
|
|
def writeback_activity() -> Iterator[None]:
|
|
"""Keep the Vast lease alive around one idle cache transfer."""
|
|
active = self.lease_manager.begin_activity(lease.instance_id)
|
|
try:
|
|
runtime.update_watchdog(active)
|
|
yield
|
|
finally:
|
|
try:
|
|
finished = self.lease_manager.finish_activity(
|
|
lease.instance_id,
|
|
idle_retention_seconds=lease.idle_retention_seconds,
|
|
)
|
|
runtime.update_watchdog(finished)
|
|
except (KeyError, OSError, RuntimeError, ValueError) as exc:
|
|
logger.error(
|
|
"Unable to release Vast R2 write-back activity instance=%d: %s",
|
|
lease.instance_id,
|
|
exc,
|
|
)
|
|
|
|
vast_settings = replace(
|
|
self.settings,
|
|
execution_mode="vast",
|
|
local_storage_root=(
|
|
self.settings.local_storage_root / "vast" / str(lease.instance_id)
|
|
).resolve(),
|
|
remote_storage_root=str(self.runtime_configuration.remote_storage_root),
|
|
)
|
|
return ModalAssetSyncEngine(
|
|
volume=VastSshVolumeBackend(
|
|
runner,
|
|
storage_root=self.runtime_configuration.remote_storage_root,
|
|
),
|
|
settings=vast_settings,
|
|
huggingface_asset_registry=self.huggingface_asset_registry,
|
|
huggingface_asset_discovery=self.huggingface_asset_discovery,
|
|
r2_cache=self.r2_cache,
|
|
r2_writeback_activity=writeback_activity,
|
|
)
|
|
|
|
def executor(self) -> VastExecutorClient:
|
|
"""Return the direct worker executor sharing this controller's lease state."""
|
|
return VastExecutorClient(
|
|
registry=self.registry,
|
|
activity_manager=self.lease_manager,
|
|
runtime_configuration=self.runtime_configuration,
|
|
user_directory=self.user_directory,
|
|
settings=self.settings,
|
|
identity_file=self.identity_file,
|
|
)
|
|
|
|
def _runner(self, lease: VastLeaseRecord) -> VastSshRunner:
|
|
"""Return an SSH runner for the lease's selected endpoint."""
|
|
return VastSshRunner(
|
|
vast_connection_from_lease(
|
|
ssh_host=lease.ssh_host,
|
|
ssh_port=lease.ssh_port,
|
|
user_directory=self.user_directory,
|
|
identity_file=self.identity_file,
|
|
)
|
|
)
|
|
|
|
def _proxy_runner(self, lease: VastLeaseRecord) -> VastSshRunner | None:
|
|
"""Return the distinct Vast proxy endpoint available after direct failure."""
|
|
if (
|
|
lease.ssh_proxy_host is None
|
|
or lease.ssh_proxy_port is None
|
|
or (
|
|
lease.ssh_host == lease.ssh_proxy_host
|
|
and lease.ssh_port == lease.ssh_proxy_port
|
|
)
|
|
):
|
|
return None
|
|
return VastSshRunner(
|
|
vast_connection_from_lease(
|
|
ssh_host=lease.ssh_proxy_host,
|
|
ssh_port=lease.ssh_proxy_port,
|
|
user_directory=self.user_directory,
|
|
identity_file=self.identity_file,
|
|
)
|
|
)
|
|
|
|
def _promote_proxy_endpoint(self, lease: VastLeaseRecord) -> None:
|
|
"""Persist the proxy as primary after it proves reachable."""
|
|
if lease.ssh_proxy_host is None or lease.ssh_proxy_port is None:
|
|
raise ValueError("Vast lease does not expose a complete proxy endpoint.")
|
|
self.registry.update(
|
|
lease.instance_id,
|
|
lambda current: replace(
|
|
current,
|
|
ssh_host=lease.ssh_proxy_host,
|
|
ssh_port=lease.ssh_proxy_port,
|
|
),
|
|
)
|
|
logger.warning(
|
|
"Selected Vast SSH proxy after direct endpoint failure instance_id=%d.",
|
|
lease.instance_id,
|
|
)
|
|
|
|
def _current_lease(self, instance_id: int) -> VastLeaseRecord:
|
|
"""Return the latest persisted lease after endpoint selection."""
|
|
lease = next(
|
|
(
|
|
candidate
|
|
for candidate in self.registry.load().leases
|
|
if candidate.instance_id == instance_id
|
|
),
|
|
None,
|
|
)
|
|
if lease is None:
|
|
raise KeyError(f"Vast lease {instance_id} disappeared during setup.")
|
|
return lease
|
|
|
|
def _initialize_runtime(
|
|
self,
|
|
lease: VastLeaseRecord,
|
|
*,
|
|
status_callback: Callable[[str], None] | None = None,
|
|
) -> VastLeaseRecord:
|
|
"""Require the pinned worker and publish its initial idle fail-safe state."""
|
|
logger.info(
|
|
"Initializing Vast worker instance_id=%d environment=%s.",
|
|
lease.instance_id,
|
|
lease.environment_id,
|
|
)
|
|
runtime = VastRuntimeManager(
|
|
runner=self._runner(lease),
|
|
configuration=self.runtime_configuration,
|
|
fallback_runner=self._proxy_runner(lease),
|
|
fallback_selected=partial(self._promote_proxy_endpoint, lease),
|
|
instance_validator=partial(self._validate_live_instance, lease),
|
|
status_callback=status_callback,
|
|
)
|
|
try:
|
|
runtime.ensure_worker()
|
|
active_lease = self._current_lease(lease.instance_id)
|
|
runtime.update_watchdog(active_lease)
|
|
except (
|
|
TimeoutError,
|
|
VastInstanceNotFoundError,
|
|
VastSshError,
|
|
ValueError,
|
|
) as exc:
|
|
logger.error(
|
|
"Vast worker initialization failed instance_id=%d environment=%s: %s",
|
|
lease.instance_id,
|
|
lease.environment_id,
|
|
exc,
|
|
)
|
|
raise
|
|
logger.info(
|
|
"Vast worker initialization completed instance_id=%d environment=%s.",
|
|
lease.instance_id,
|
|
lease.environment_id,
|
|
)
|
|
return active_lease
|
|
|
|
def _validate_live_instance(self, lease: VastLeaseRecord) -> None:
|
|
"""Require the managed Vast contract to exist before an SSH probe."""
|
|
asyncio.run(self.api_client.show_instance(lease.instance_id))
|
|
|
|
def _existing_lease(self, profile: VastResourceProfile) -> VastLeaseRecord | None:
|
|
"""Return the oldest ready lease with an exact profile/runtime identity."""
|
|
fingerprint = vast_profile_fingerprint(profile)
|
|
matches = [
|
|
lease
|
|
for lease in self.registry.load().leases
|
|
if lease.owner_id == self.settings.app_name
|
|
and lease.profile_id == profile.profile_id
|
|
and lease.profile_fingerprint == fingerprint
|
|
and lease.worker_image in {None, self.runtime_configuration.image}
|
|
and lease.ready_for_work
|
|
]
|
|
return min(matches, key=lambda lease: lease.created_at_epoch, default=None)
|
|
|
|
|
|
def _profile_for_requirements(
|
|
profile: VastResourceProfile,
|
|
minimum_vram_bytes: int,
|
|
minimum_ram_bytes: int,
|
|
) -> VastResourceProfile:
|
|
"""Raise workflow profile memory floors to inferred component requirements."""
|
|
return replace(
|
|
profile,
|
|
minimum_gpu_ram_mb=_raised_memory_floor_mb(
|
|
profile.minimum_gpu_ram_mb,
|
|
minimum_vram_bytes,
|
|
),
|
|
minimum_cpu_ram_mb=_raised_memory_floor_mb(
|
|
profile.minimum_cpu_ram_mb,
|
|
minimum_ram_bytes,
|
|
),
|
|
)
|
|
|
|
|
|
def _raised_memory_floor_mb(
|
|
configured_floor_mb: int | None,
|
|
inferred_floor_bytes: int,
|
|
) -> int | None:
|
|
"""Apply an inferred requirement while preserving an unconstrained zero."""
|
|
inferred_floor_mb = math.ceil(inferred_floor_bytes / 1024**2)
|
|
if configured_floor_mb is None and inferred_floor_mb == 0:
|
|
return None
|
|
return max(configured_floor_mb or 0, inferred_floor_mb)
|
|
|
|
|
|
__all__ = [
|
|
"VAST_API_BASE_URL_ENV",
|
|
"VAST_API_KEY_ENV",
|
|
"VAST_SSH_IDENTITY_FILE_ENV",
|
|
"VastProfileQuote",
|
|
"VastSearchRequirements",
|
|
"VastService",
|
|
]
|