702 lines
24 KiB
Python
702 lines
24 KiB
Python
"""Tests for Vast marketplace quotation and scheduler integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path, PurePosixPath
|
|
from types import SimpleNamespace
|
|
from typing import Any, AsyncIterator
|
|
|
|
import pytest
|
|
|
|
from aiohttp import web
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _running_app(app: web.Application) -> AsyncIterator[str]:
|
|
"""Serve one aiohttp application on an ephemeral loopback port."""
|
|
runner = web.AppRunner(app)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, "127.0.0.1", 0)
|
|
await site.start()
|
|
server = site._server
|
|
if server is None or not server.sockets:
|
|
await runner.cleanup()
|
|
raise RuntimeError("Test server did not expose a listening socket.")
|
|
port = int(server.sockets[0].getsockname()[1])
|
|
try:
|
|
yield f"http://127.0.0.1:{port}"
|
|
finally:
|
|
await runner.cleanup()
|
|
|
|
|
|
def test_quote_raises_memory_floors_and_compares_effective_hourly_price(
|
|
vast_api_module: Any,
|
|
vast_models_module: Any,
|
|
vast_runtime_module: Any,
|
|
vast_service_module: Any,
|
|
vast_simulator_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Scheduling should exclude retention cost while still exposing its estimate."""
|
|
|
|
async def scenario() -> None:
|
|
"""Exercise the production client and service against the simulator."""
|
|
simulator_state = vast_simulator_module.VastSimulatorState()
|
|
simulator = vast_simulator_module.VastApiSimulator(simulator_state)
|
|
async with _running_app(simulator.app) as base_url:
|
|
settings = SimpleNamespace(
|
|
app_name="test-owner",
|
|
modal_gpu="RTX-PRO-6000",
|
|
)
|
|
registry = vast_service_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
)
|
|
service = vast_service_module.VastService(
|
|
settings=settings,
|
|
repo_root=tmp_path,
|
|
user_directory=tmp_path,
|
|
api_client=vast_api_module.VastApiClient(
|
|
simulator_state.api_key,
|
|
base_url=base_url,
|
|
),
|
|
runtime_configuration=vast_runtime_module.VastRuntimeConfiguration(
|
|
image="example.invalid/comfy-worker:test",
|
|
runtime_fingerprint="a" * 64,
|
|
),
|
|
registry=registry,
|
|
)
|
|
profile = vast_models_module.VastResourceProfile(
|
|
profile_id="77",
|
|
profile_name="workflow-default",
|
|
maximum_hourly_cost_usd=2.0,
|
|
idle_retention_seconds=24 * 3600,
|
|
)
|
|
|
|
quote = await service.quote_best_profile(
|
|
(profile,),
|
|
minimum_vram_bytes=40 * 1024**3,
|
|
minimum_ram_bytes=96 * 1024**3,
|
|
predicted_execution_seconds=120.0,
|
|
)
|
|
scheduling_state = service.scheduling_state(quote)
|
|
|
|
assert quote.offer.offer_id == 1002
|
|
assert quote.profile.minimum_gpu_ram_mb == 40 * 1024
|
|
assert quote.profile.minimum_cpu_ram_mb == 96 * 1024
|
|
assert quote.predicted_incremental_cost_usd == 120 / 3600 * 0.74
|
|
assert quote.predicted_retention_cost_usd == (
|
|
(24 * 3600 + 120) / 3600 * 0.74
|
|
)
|
|
assert scheduling_state.provider.value == "vast"
|
|
assert scheduling_state.cost_usd_per_second == 0.74 / 3600
|
|
assert scheduling_state.capabilities.maximum_vram_bytes == 48 * 1024**3
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_from_environment_requires_credential_before_other_setup(
|
|
vast_service_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Vast-only selection should fail clearly without touching marketplace state."""
|
|
settings = SimpleNamespace()
|
|
|
|
try:
|
|
vast_service_module.VastService.from_environment(
|
|
settings,
|
|
repo_root=tmp_path,
|
|
environment={},
|
|
)
|
|
except RuntimeError as exc:
|
|
assert "VAST_API_KEY" in str(exc)
|
|
else:
|
|
raise AssertionError("Missing Vast credential was accepted.")
|
|
|
|
|
|
def test_sync_engine_keeps_vast_lease_active_during_r2_writeback(
|
|
settings_module: Any,
|
|
vast_service_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Idle R2 transfers should hold lease activity and refresh both watchdog states."""
|
|
events: list[tuple[str, Any]] = []
|
|
|
|
class FakeLeaseManager:
|
|
"""Record the activity lifecycle selected for background write-back."""
|
|
|
|
def begin_activity(self, instance_id: int) -> Any:
|
|
"""Return one active lease snapshot."""
|
|
events.append(("begin", instance_id))
|
|
return SimpleNamespace(state="active")
|
|
|
|
def finish_activity(
|
|
self,
|
|
instance_id: int,
|
|
*,
|
|
idle_retention_seconds: float,
|
|
) -> Any:
|
|
"""Return one idle lease snapshot."""
|
|
events.append(("finish", (instance_id, idle_retention_seconds)))
|
|
return SimpleNamespace(state="idle")
|
|
|
|
class FakeRuntime:
|
|
"""Record watchdog publications without opening SSH."""
|
|
|
|
def __init__(self, **_kwargs: Any) -> None:
|
|
"""Accept the production runtime constructor arguments."""
|
|
|
|
def update_watchdog(self, lease: Any) -> None:
|
|
"""Record one active or idle watchdog snapshot."""
|
|
events.append(("watchdog", lease.state))
|
|
|
|
class FakeVolume:
|
|
"""Accept the production Vast storage constructor arguments."""
|
|
|
|
def __init__(self, *_args: Any, **_kwargs: Any) -> None:
|
|
"""Initialize the no-op test backend."""
|
|
|
|
monkeypatch.setattr(vast_service_module, "VastRuntimeManager", FakeRuntime)
|
|
monkeypatch.setattr(vast_service_module, "VastSshVolumeBackend", FakeVolume)
|
|
settings = settings_module.ModalSyncSettings(
|
|
app_name="app",
|
|
auto_deploy=True,
|
|
allow_ephemeral_fallback=False,
|
|
enable_memory_snapshot=True,
|
|
enable_gpu_memory_snapshot=False,
|
|
execution_mode="vast",
|
|
sync_custom_nodes=False,
|
|
volume_name="volume",
|
|
route_path="/modal/queue_prompt",
|
|
marker_property="is_modal_remote",
|
|
local_storage_root=tmp_path / "storage",
|
|
remote_storage_root="/storage",
|
|
custom_nodes_archive_name="custom_nodes_bundle.zip",
|
|
comfyui_root=None,
|
|
custom_nodes_dir=None,
|
|
)
|
|
service = SimpleNamespace(
|
|
settings=settings,
|
|
runtime_configuration=SimpleNamespace(
|
|
remote_storage_root=PurePosixPath("/storage")
|
|
),
|
|
lease_manager=FakeLeaseManager(),
|
|
huggingface_asset_registry=None,
|
|
huggingface_asset_discovery=None,
|
|
r2_cache=None,
|
|
_runner=lambda _lease: object(),
|
|
)
|
|
lease = SimpleNamespace(instance_id=42, idle_retention_seconds=90.0)
|
|
|
|
engine = vast_service_module.VastService.sync_engine(service, lease)
|
|
assert engine.r2_writeback_activity is not None
|
|
with engine.r2_writeback_activity():
|
|
events.append(("transfer", 42))
|
|
|
|
assert events == [
|
|
("begin", 42),
|
|
("watchdog", "active"),
|
|
("transfer", 42),
|
|
("finish", (42, 90.0)),
|
|
("watchdog", "idle"),
|
|
]
|
|
|
|
|
|
def test_startup_status_translates_vast_image_progress(
|
|
vast_models_module: Any,
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""Provider layer messages should become useful workflow status without hashes."""
|
|
instance = vast_models_module.VastInstance.from_api(
|
|
{
|
|
"id": 48534708,
|
|
"actual_status": "loading",
|
|
"intended_status": "running",
|
|
"cur_state": "running",
|
|
"status_msg": "4c7d45f7c63c: Download complete",
|
|
"num_gpus": 1,
|
|
"gpu_ram": 48 * 1024,
|
|
"cpu_ram": 96 * 1024,
|
|
}
|
|
)
|
|
|
|
message = vast_service_module._vast_startup_status_message(instance)
|
|
|
|
assert message == (
|
|
"Vast.ai instance 48534708 is downloading the worker image (layer complete)"
|
|
)
|
|
assert "4c7d45f7c63c" not in message
|
|
|
|
|
|
def test_startup_status_uses_current_state_when_actual_status_is_unknown(
|
|
vast_models_module: Any,
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""Missing actual status should not render a misleading unknown state."""
|
|
instance = vast_models_module.VastInstance.from_api(
|
|
{
|
|
"id": 48597015,
|
|
"actual_status": None,
|
|
"intended_status": "running",
|
|
"cur_state": "loading",
|
|
"num_gpus": 1,
|
|
"gpu_ram": 96 * 1024,
|
|
"cpu_ram": 128 * 1024,
|
|
}
|
|
)
|
|
|
|
assert vast_service_module._vast_startup_status_message(instance) == (
|
|
"Vast.ai instance 48597015 is starting (loading)"
|
|
)
|
|
|
|
|
|
def test_prefetch_deduplicates_effective_profiles_and_searches_in_parallel(
|
|
vast_models_module: Any,
|
|
vast_runtime_module: Any,
|
|
vast_service_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Planning should concurrently warm only distinct marketplace queries."""
|
|
|
|
class DelayedApiClient:
|
|
"""Record overlapping searches without making network requests."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize concurrency counters and searched VRAM floors."""
|
|
self.active = 0
|
|
self.peak = 0
|
|
self.minimum_gpu_ram_mb: list[int] = []
|
|
|
|
async def search_offers(self, profile: Any) -> tuple[Any, ...]:
|
|
"""Hold the event loop briefly so independent searches overlap."""
|
|
self.active += 1
|
|
self.peak = max(self.peak, self.active)
|
|
self.minimum_gpu_ram_mb.append(profile.minimum_gpu_ram_mb)
|
|
await asyncio.sleep(0.02)
|
|
self.active -= 1
|
|
return ()
|
|
|
|
async def scenario() -> None:
|
|
"""Prefetch repeated low floors and one raised floor."""
|
|
api_client = DelayedApiClient()
|
|
service = vast_service_module.VastService(
|
|
settings=SimpleNamespace(
|
|
app_name="test-owner",
|
|
modal_gpu="RTX-PRO-6000",
|
|
),
|
|
repo_root=tmp_path,
|
|
user_directory=tmp_path,
|
|
api_client=api_client,
|
|
runtime_configuration=vast_runtime_module.VastRuntimeConfiguration(
|
|
image="example.invalid/comfy-worker:test",
|
|
runtime_fingerprint="b" * 64,
|
|
),
|
|
registry=vast_service_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
)
|
|
profile = vast_models_module.VastResourceProfile(
|
|
profile_id="parallel",
|
|
profile_name="parallel",
|
|
minimum_gpu_ram_mb=48 * 1024,
|
|
)
|
|
requirements = (
|
|
vast_service_module.VastSearchRequirements(0, 0),
|
|
vast_service_module.VastSearchRequirements(16 * 1024**3, 0),
|
|
vast_service_module.VastSearchRequirements(80 * 1024**3, 0),
|
|
)
|
|
|
|
await service.prefetch_offers((profile,), requirements)
|
|
|
|
assert sorted(api_client.minimum_gpu_ram_mb) == [48 * 1024, 80 * 1024]
|
|
assert api_client.peak == 2
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_acquire_replaces_lease_when_worker_initialization_fails(
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""An SSH-unusable rental must be destroyed and replaced in the same run."""
|
|
|
|
class FakeRegistry:
|
|
"""Record that the unusable lease was marked before destruction."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the updated instance list."""
|
|
self.updated_instance_ids: list[int] = []
|
|
|
|
def update(self, instance_id: int, updater: Any) -> None:
|
|
"""Record the update without requiring a complete lease fixture."""
|
|
del updater
|
|
self.updated_instance_ids.append(instance_id)
|
|
|
|
class FakeLeaseManager:
|
|
"""Return an unusable rental followed by a healthy replacement."""
|
|
|
|
def __init__(self, leases: list[Any]) -> None:
|
|
"""Retain the fake leases, exclusions, and destruction log."""
|
|
self.leases = leases
|
|
self.destroyed_instance_ids: list[int] = []
|
|
self.excluded_offer_ids: list[frozenset[int]] = []
|
|
|
|
async def ensure_lease(
|
|
self,
|
|
profile: Any,
|
|
*,
|
|
excluded_offer_ids: frozenset[int],
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Return each selected lease while recording rejected offers."""
|
|
del profile, kwargs
|
|
self.excluded_offer_ids.append(excluded_offer_ids)
|
|
return self.leases.pop(0)
|
|
|
|
async def destroy_owned_lease(self, instance_id: int) -> bool:
|
|
"""Record cleanup of the unusable capacity."""
|
|
self.destroyed_instance_ids.append(instance_id)
|
|
return True
|
|
|
|
failed = SimpleNamespace(
|
|
instance_id=42,
|
|
offer_id=1001,
|
|
environment_id="vast:test:42",
|
|
)
|
|
replacement = SimpleNamespace(
|
|
instance_id=43,
|
|
offer_id=1002,
|
|
environment_id="vast:test:43",
|
|
)
|
|
lease_manager = FakeLeaseManager([failed, replacement])
|
|
registry = FakeRegistry()
|
|
service = object.__new__(vast_service_module.VastService)
|
|
service.lease_manager = lease_manager
|
|
service.registry = registry
|
|
|
|
async def skip_preflight(*, status_callback: Any) -> None:
|
|
"""Keep this test focused on post-rental SSH replacement."""
|
|
del status_callback
|
|
|
|
service._preflight_published_image = skip_preflight
|
|
|
|
def initialize_runtime(
|
|
lease: Any,
|
|
*,
|
|
status_callback: Any,
|
|
) -> None:
|
|
"""Simulate the observed SSH failure only on the first rental."""
|
|
if lease.instance_id == failed.instance_id:
|
|
status_callback(
|
|
"Vast SSH attempt 1 failed: Connection refused; retrying in 1.00s"
|
|
)
|
|
raise vast_service_module.VastSshError(
|
|
"kex_exchange_identification: Connection closed by remote host"
|
|
)
|
|
|
|
service._initialize_runtime = initialize_runtime
|
|
messages: list[str] = []
|
|
|
|
acquired = asyncio.run(
|
|
service.acquire(
|
|
SimpleNamespace(profile=SimpleNamespace(), existing_lease=None),
|
|
status_callback=messages.append,
|
|
)
|
|
)
|
|
|
|
assert acquired is replacement
|
|
assert registry.updated_instance_ids == [42]
|
|
assert lease_manager.destroyed_instance_ids == [42]
|
|
assert lease_manager.excluded_offer_ids == [frozenset(), frozenset({1001})]
|
|
assert any("Vast SSH attempt 1 failed" in message for message in messages)
|
|
assert any("failed worker setup" in message for message in messages)
|
|
assert messages[-1] == "Vast.ai worker is ready"
|
|
|
|
|
|
def test_acquire_replaces_instance_that_disappears_before_ssh(
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""A missing live contract should cold-start the next marketplace offer."""
|
|
|
|
class FakeRegistry:
|
|
"""Record stale contract removal before replacement."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the removed instance log."""
|
|
self.removed_instance_ids: list[int] = []
|
|
|
|
def remove(self, instance_id: int) -> None:
|
|
"""Record removal of a provider-missing lease."""
|
|
self.removed_instance_ids.append(instance_id)
|
|
|
|
class FakeLeaseManager:
|
|
"""Return a vanished lease followed by a fresh replacement."""
|
|
|
|
def __init__(self, leases: list[Any]) -> None:
|
|
"""Retain sequenced leases and offer exclusions."""
|
|
self.leases = leases
|
|
self.excluded_offer_ids: list[frozenset[int]] = []
|
|
|
|
async def ensure_lease(
|
|
self,
|
|
profile: Any,
|
|
*,
|
|
excluded_offer_ids: frozenset[int],
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Return the next lease while recording rejected offers."""
|
|
del profile, kwargs
|
|
self.excluded_offer_ids.append(excluded_offer_ids)
|
|
return self.leases.pop(0)
|
|
|
|
vanished = SimpleNamespace(
|
|
instance_id=42,
|
|
offer_id=1001,
|
|
environment_id="vast:test:42",
|
|
)
|
|
replacement = SimpleNamespace(
|
|
instance_id=43,
|
|
offer_id=1002,
|
|
environment_id="vast:test:43",
|
|
)
|
|
lease_manager = FakeLeaseManager([vanished, replacement])
|
|
registry = FakeRegistry()
|
|
service = object.__new__(vast_service_module.VastService)
|
|
service.lease_manager = lease_manager
|
|
service.registry = registry
|
|
|
|
async def skip_preflight(*, status_callback: Any) -> None:
|
|
"""Keep this test focused on a disappeared provider contract."""
|
|
del status_callback
|
|
|
|
service._preflight_published_image = skip_preflight
|
|
|
|
def initialize_runtime(
|
|
lease: Any,
|
|
*,
|
|
status_callback: Any,
|
|
) -> None:
|
|
"""Report only the first provider contract as missing."""
|
|
del status_callback
|
|
if lease.instance_id == vanished.instance_id:
|
|
raise vast_service_module.VastInstanceNotFoundError(
|
|
"Vast instance 42 does not exist."
|
|
)
|
|
|
|
service._initialize_runtime = initialize_runtime
|
|
messages: list[str] = []
|
|
|
|
acquired = asyncio.run(
|
|
service.acquire(
|
|
SimpleNamespace(profile=SimpleNamespace(), existing_lease=None),
|
|
status_callback=messages.append,
|
|
)
|
|
)
|
|
|
|
assert acquired is replacement
|
|
assert registry.removed_instance_ids == [42]
|
|
assert lease_manager.excluded_offer_ids == [frozenset(), frozenset({1001})]
|
|
assert "Vast.ai instance disappeared; requesting a replacement" in messages
|
|
assert messages[-1] == "Vast.ai worker is ready"
|
|
|
|
|
|
def test_acquire_rebuilds_image_and_replaces_fingerprint_drift(
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""A stale published worker should be rebuilt once before capacity is retried."""
|
|
|
|
class FakeRegistry:
|
|
"""Record stale worker draining before its destruction."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the updated instance log."""
|
|
self.updated_instance_ids: list[int] = []
|
|
|
|
def update(self, instance_id: int, updater: Any) -> None:
|
|
"""Record one stale lease update."""
|
|
del updater
|
|
self.updated_instance_ids.append(instance_id)
|
|
|
|
class FakeLeaseManager:
|
|
"""Return one stale lease followed by a fresh-image lease."""
|
|
|
|
def __init__(self, leases: list[Any]) -> None:
|
|
"""Retain sequenced leases and destruction calls."""
|
|
self.leases = leases
|
|
self.destroyed_instance_ids: list[int] = []
|
|
|
|
async def ensure_lease(self, profile: Any, **kwargs: Any) -> Any:
|
|
"""Return the next lease without contacting Vast."""
|
|
del profile, kwargs
|
|
return self.leases.pop(0)
|
|
|
|
async def destroy_owned_lease(self, instance_id: int) -> bool:
|
|
"""Record stale capacity cleanup."""
|
|
self.destroyed_instance_ids.append(instance_id)
|
|
return True
|
|
|
|
class FakeImageBuilder:
|
|
"""Publish a deterministic replacement image."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize requested fingerprint tracking."""
|
|
self.fingerprints: list[str] = []
|
|
|
|
def build_and_push(
|
|
self,
|
|
expected_fingerprint: str,
|
|
*,
|
|
status_callback: Any,
|
|
) -> str:
|
|
"""Emit progress and return a digest-pinned image."""
|
|
self.fingerprints.append(expected_fingerprint)
|
|
status_callback("Building replacement worker")
|
|
return "ghcr.io/example/worker@sha256:" + "c" * 64
|
|
|
|
stale = SimpleNamespace(
|
|
instance_id=42,
|
|
offer_id=1001,
|
|
environment_id="vast:test:42",
|
|
)
|
|
replacement = SimpleNamespace(
|
|
instance_id=43,
|
|
offer_id=1001,
|
|
environment_id="vast:test:43",
|
|
)
|
|
lease_manager = FakeLeaseManager([stale, replacement])
|
|
image_builder = FakeImageBuilder()
|
|
service = object.__new__(vast_service_module.VastService)
|
|
service.lease_manager = lease_manager
|
|
service.registry = FakeRegistry()
|
|
service.image_builder = image_builder
|
|
adopted_images: list[str] = []
|
|
service._adopt_runtime_image = adopted_images.append
|
|
|
|
async def skip_preflight(*, status_callback: Any) -> None:
|
|
"""Keep runtime drift as a defense-in-depth behavior in this test."""
|
|
del status_callback
|
|
|
|
service._preflight_published_image = skip_preflight
|
|
|
|
def initialize_runtime(
|
|
lease: Any,
|
|
*,
|
|
status_callback: Any,
|
|
) -> None:
|
|
"""Reject only the worker baked into the first lease."""
|
|
del status_callback
|
|
if lease.instance_id == stale.instance_id:
|
|
raise vast_service_module.VastRuntimeFingerprintDriftError(
|
|
expected_fingerprint="a" * 64,
|
|
actual_fingerprint="b" * 64,
|
|
protocol_version=1,
|
|
)
|
|
|
|
service._initialize_runtime = initialize_runtime
|
|
messages: list[str] = []
|
|
|
|
acquired = asyncio.run(
|
|
service.acquire(
|
|
SimpleNamespace(profile=SimpleNamespace(), existing_lease=None),
|
|
status_callback=messages.append,
|
|
)
|
|
)
|
|
|
|
assert acquired is replacement
|
|
assert service.registry.updated_instance_ids == [42]
|
|
assert lease_manager.destroyed_instance_ids == [42]
|
|
assert image_builder.fingerprints == ["a" * 64]
|
|
assert adopted_images == ["ghcr.io/example/worker@sha256:" + "c" * 64]
|
|
assert "Building replacement worker" in messages
|
|
assert "Vast worker image updated; requesting fresh capacity" in messages
|
|
assert messages[-1] == "Vast.ai worker is ready"
|
|
|
|
|
|
def test_preflight_rebuilds_stale_image_before_capacity_request(
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""Registry drift should be repaired before lease acquisition can begin."""
|
|
|
|
class FakeImageBuilder:
|
|
"""Return a replacement digest and record preflight inputs."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the call log."""
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
def ensure_published_image(
|
|
self,
|
|
image: str,
|
|
expected_fingerprint: str,
|
|
*,
|
|
status_callback: Any,
|
|
) -> str:
|
|
"""Simulate a stale registry image replacement."""
|
|
self.calls.append((image, expected_fingerprint))
|
|
status_callback("Published Vast worker image is stale")
|
|
return "ghcr.io/example/worker@sha256:" + "d" * 64
|
|
|
|
builder = FakeImageBuilder()
|
|
service = object.__new__(vast_service_module.VastService)
|
|
service.runtime_configuration = SimpleNamespace(
|
|
image="ghcr.io/example/worker:v1",
|
|
runtime_fingerprint="a" * 64,
|
|
)
|
|
service.image_builder = builder
|
|
adopted: list[str] = []
|
|
service._adopt_runtime_image = adopted.append
|
|
messages: list[str] = []
|
|
|
|
asyncio.run(service._preflight_published_image(status_callback=messages.append))
|
|
|
|
assert builder.calls == [("ghcr.io/example/worker:v1", "a" * 64)]
|
|
assert adopted == ["ghcr.io/example/worker@sha256:" + "d" * 64]
|
|
assert messages[-1] == "Vast worker image updated before requesting capacity"
|
|
|
|
|
|
def test_proxy_endpoint_promotion_persists_the_working_route(
|
|
vast_service_module: Any,
|
|
) -> None:
|
|
"""Successful proxy readiness must replace the dead direct primary endpoint."""
|
|
|
|
@dataclass(frozen=True)
|
|
class Lease:
|
|
"""Represent the endpoint fields changed by proxy promotion."""
|
|
|
|
instance_id: int
|
|
ssh_host: str
|
|
ssh_port: int
|
|
ssh_proxy_host: str
|
|
ssh_proxy_port: int
|
|
|
|
class UpdatingRegistry:
|
|
"""Apply one lease update in memory."""
|
|
|
|
def __init__(self, lease: Any) -> None:
|
|
"""Retain the current lease."""
|
|
self.lease = lease
|
|
|
|
def update(self, instance_id: int, updater: Any) -> Any:
|
|
"""Apply and retain the requested endpoint change."""
|
|
assert instance_id == self.lease.instance_id
|
|
self.lease = updater(self.lease)
|
|
return self.lease
|
|
|
|
lease = Lease(
|
|
instance_id=42,
|
|
ssh_host="192.0.2.4",
|
|
ssh_port=40112,
|
|
ssh_proxy_host="ssh7.vast.ai",
|
|
ssh_proxy_port=22017,
|
|
)
|
|
registry = UpdatingRegistry(lease)
|
|
service = object.__new__(vast_service_module.VastService)
|
|
service.registry = registry
|
|
|
|
service._promote_proxy_endpoint(lease)
|
|
|
|
assert registry.lease.ssh_host == "ssh7.vast.ai"
|
|
assert registry.lease.ssh_port == 22017
|