796 lines
29 KiB
Python
796 lines
29 KiB
Python
"""Tests for persistent Vast.ai lease acquisition and idle cleanup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Any, AsyncIterator
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
|
|
|
|
@dataclass
|
|
class FakeClock:
|
|
"""Controllable wall clock for lease-retention tests."""
|
|
|
|
value: float = 1_000_000.0
|
|
|
|
def __call__(self) -> float:
|
|
"""Return the current fake epoch."""
|
|
return self.value
|
|
|
|
def advance(self, seconds: float) -> None:
|
|
"""Move the fake epoch forward."""
|
|
self.value += seconds
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _running_simulator(app: web.Application) -> AsyncIterator[str]:
|
|
"""Serve one local simulator app on an ephemeral 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("Simulator 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 _runtime_fingerprint() -> str:
|
|
"""Return a deterministic runtime SHA-256 value."""
|
|
return "a" * 64
|
|
|
|
|
|
def _profile(vast_models_module: Any, **overrides: Any) -> Any:
|
|
"""Return one test profile with short retention."""
|
|
values = {
|
|
"profile_id": "node-17",
|
|
"profile_name": "default",
|
|
"idle_retention_seconds": 60.0,
|
|
}
|
|
values.update(overrides)
|
|
return vast_models_module.VastResourceProfile(**values)
|
|
|
|
|
|
def _launch_factory(vast_models_module: Any) -> Any:
|
|
"""Return a deterministic launch-spec factory."""
|
|
|
|
def factory(profile: Any, label: str) -> Any:
|
|
"""Build a launch spec matching one selected profile."""
|
|
return vast_models_module.VastInstanceLaunchSpec(
|
|
image="ghcr.io/example/comfy-vast@sha256:test",
|
|
disk_gb=profile.allocated_disk_gb,
|
|
label=label,
|
|
onstart="/opt/comfy/start-worker",
|
|
environment={"COMFY_MODAL_REMOTE_WORKER": "1"},
|
|
)
|
|
|
|
return factory
|
|
|
|
|
|
def test_manager_recovers_offer_race_then_reuses_persisted_lease(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""Rent the next candidate after a race and avoid duplicate instances."""
|
|
|
|
async def scenario() -> None:
|
|
state = vast_simulator_module.VastSimulatorState(
|
|
create_failures_remaining={1001: 1},
|
|
create_failure_status=400,
|
|
create_failure_message=(
|
|
"error 404/3603: no_such_ask Instance type by id 1001 is not available."
|
|
),
|
|
polls_until_running=1,
|
|
)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=vast_api_module.VastApiClient(
|
|
state.api_key,
|
|
base_url=base_url,
|
|
retry_attempts=1,
|
|
),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
profile = _profile(vast_models_module)
|
|
|
|
first = await manager.ensure_lease(profile)
|
|
second = await manager.ensure_lease(profile)
|
|
|
|
assert first.instance_id == second.instance_id
|
|
assert first.offer_id == 1002
|
|
assert len(state.instances) == 1
|
|
create_requests = [
|
|
request
|
|
for request in state.request_log
|
|
if request["path"].startswith("/api/v0/asks/")
|
|
]
|
|
assert [request["path"] for request in create_requests] == [
|
|
"/api/v0/asks/1001/",
|
|
"/api/v0/asks/1002/",
|
|
]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_manager_excludes_offer_from_disappeared_instance_replacement(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A cold replacement must not immediately rent the vanished offer again."""
|
|
|
|
async def scenario() -> None:
|
|
"""Exclude the cheapest offer and rent the next compatible candidate."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=vast_api_module.VastApiClient(
|
|
state.api_key,
|
|
base_url=base_url,
|
|
),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
|
|
lease = await manager.ensure_lease(
|
|
_profile(vast_models_module),
|
|
excluded_offer_ids=frozenset({1001}),
|
|
)
|
|
|
|
assert lease.offer_id == 1002
|
|
create_requests = [
|
|
request
|
|
for request in state.request_log
|
|
if request["path"].startswith("/api/v0/asks/")
|
|
]
|
|
assert [request["path"] for request in create_requests] == [
|
|
"/api/v0/asks/1002/"
|
|
]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_manager_replaces_contract_that_disappears_during_provider_startup(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A contract lost before SSH readiness should rent the next offer."""
|
|
|
|
class DisappearingClient(vast_api_module.VastApiClient):
|
|
"""Remove the first created contract during provider readiness."""
|
|
|
|
def __init__(self, *args: Any, state: Any, **kwargs: Any) -> None:
|
|
"""Retain simulator state and initialize the disappearance count."""
|
|
super().__init__(*args, **kwargs)
|
|
self.state = state
|
|
self.readiness_calls = 0
|
|
|
|
async def wait_until_ready(
|
|
self,
|
|
instance_id: int,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Lose the first contract and delegate later readiness checks."""
|
|
self.readiness_calls += 1
|
|
if self.readiness_calls == 1:
|
|
self.state.instances.pop(instance_id)
|
|
raise vast_api_module.VastInstanceNotFoundError(
|
|
f"Vast instance {instance_id} does not exist."
|
|
)
|
|
return await super().wait_until_ready(instance_id, **kwargs)
|
|
|
|
async def scenario() -> None:
|
|
"""Verify cold replacement and stale registry cleanup."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=DisappearingClient(
|
|
state.api_key,
|
|
base_url=base_url,
|
|
state=state,
|
|
),
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
|
|
lease = await manager.ensure_lease(_profile(vast_models_module))
|
|
|
|
assert lease.offer_id == 1002
|
|
assert [record.instance_id for record in registry.load().leases] == [
|
|
lease.instance_id
|
|
]
|
|
create_requests = [
|
|
request
|
|
for request in state.request_log
|
|
if request["path"].startswith("/api/v0/asks/")
|
|
]
|
|
assert [request["path"] for request in create_requests] == [
|
|
"/api/v0/asks/1001/",
|
|
"/api/v0/asks/1002/",
|
|
]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_manager_replaces_instance_that_times_out_during_provider_startup(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A loading instance that misses its deadline should be destroyed and replaced."""
|
|
|
|
class TimingOutClient(vast_api_module.VastApiClient):
|
|
"""Timeout the first rental and allow its replacement to become ready."""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
"""Initialize the deterministic readiness attempt counter."""
|
|
super().__init__(*args, **kwargs)
|
|
self.readiness_calls = 0
|
|
|
|
async def wait_until_ready(
|
|
self,
|
|
instance_id: int,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Fail only the initial instance's provider-readiness wait."""
|
|
self.readiness_calls += 1
|
|
if self.readiness_calls == 1:
|
|
raise TimeoutError(
|
|
f"Vast instance {instance_id} remained in 'loading'."
|
|
)
|
|
return await super().wait_until_ready(instance_id, **kwargs)
|
|
|
|
async def scenario() -> None:
|
|
"""Verify API destruction, offer exclusion, and replacement telemetry."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=TimingOutClient(state.api_key, base_url=base_url),
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
events: list[str] = []
|
|
|
|
lease = await manager.ensure_lease(
|
|
_profile(vast_models_module),
|
|
event_callback=events.append,
|
|
)
|
|
|
|
assert lease.offer_id == 1002
|
|
assert len(state.destroyed_instance_ids) == 1
|
|
assert state.destroyed_instance_ids[0] != lease.instance_id
|
|
assert [record.instance_id for record in registry.load().leases] == [
|
|
lease.instance_id
|
|
]
|
|
create_requests = [
|
|
request
|
|
for request in state.request_log
|
|
if request["path"].startswith("/api/v0/asks/")
|
|
]
|
|
assert [request["path"] for request in create_requests] == [
|
|
"/api/v0/asks/1001/",
|
|
"/api/v0/asks/1002/",
|
|
]
|
|
assert "failed setup and was terminated" in events[0]
|
|
assert "requesting a replacement" in events[0]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_manager_bounds_replacements_when_every_instance_times_out(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A bad marketplace pool must not create an unbounded rental loop."""
|
|
|
|
class AlwaysTimingOutClient(vast_api_module.VastApiClient):
|
|
"""Keep every newly created instance in a synthetic failed setup state."""
|
|
|
|
async def wait_until_ready(
|
|
self,
|
|
instance_id: int,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""Fail immediately while preserving the instance for API destruction."""
|
|
del kwargs
|
|
raise TimeoutError(f"Vast instance {instance_id} remained in 'loading'.")
|
|
|
|
async def scenario() -> None:
|
|
"""Verify exactly an initial attempt plus one replacement."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=AlwaysTimingOutClient(state.api_key, base_url=base_url),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
events: list[str] = []
|
|
|
|
with pytest.raises(
|
|
vast_leases_module.VastLeaseStartupExhaustedError,
|
|
match="after 2 attempts",
|
|
):
|
|
await manager.ensure_lease(
|
|
_profile(vast_models_module),
|
|
event_callback=events.append,
|
|
)
|
|
|
|
assert len(state.destroyed_instance_ids) == 2
|
|
assert "requesting a replacement" in events[0]
|
|
assert "no further automatic attempts remain" in events[-1]
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_manager_adopts_legacy_lease_for_unchanged_worker_image(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""Controller fingerprint drift should not rent again for the same worker image."""
|
|
|
|
async def scenario() -> None:
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
client = vast_api_module.VastApiClient(state.api_key, base_url=base_url)
|
|
profile = _profile(vast_models_module)
|
|
first_manager = vast_leases_module.VastLeaseManager(
|
|
api_client=client,
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint="a" * 64,
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
first = await first_manager.ensure_lease(profile)
|
|
second_manager = vast_leases_module.VastLeaseManager(
|
|
api_client=client,
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint="b" * 64,
|
|
worker_image="ghcr.io/example/worker@sha256:same",
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
|
|
second = await second_manager.ensure_lease(profile)
|
|
|
|
assert second.instance_id == first.instance_id
|
|
assert second.runtime_fingerprint == "b" * 64
|
|
assert second.worker_image == "ghcr.io/example/worker@sha256:same"
|
|
assert len(state.instances) == 1
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_active_work_suspends_idle_destruction_and_finish_resets_deadline(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""Retention begins at terminal activity and never destroys active work."""
|
|
|
|
async def scenario() -> None:
|
|
clock = FakeClock()
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=vast_api_module.VastApiClient(
|
|
state.api_key, base_url=base_url
|
|
),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
clock=clock,
|
|
)
|
|
profile = _profile(vast_models_module)
|
|
lease = await manager.ensure_lease(profile)
|
|
manager.begin_activity(lease.instance_id)
|
|
clock.advance(600)
|
|
|
|
assert await manager.destroy_expired() == ()
|
|
finished = manager.finish_activity(
|
|
lease.instance_id,
|
|
idle_retention_seconds=profile.idle_retention_seconds,
|
|
)
|
|
assert finished.idle_deadline_epoch == clock.value + 60
|
|
clock.advance(59)
|
|
assert await manager.destroy_expired() == ()
|
|
clock.advance(2)
|
|
|
|
assert await manager.destroy_expired() == (lease.instance_id,)
|
|
assert state.instances == {}
|
|
assert manager.registry.load().leases == ()
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_reconcile_after_restart_clears_stale_activity_count(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A dead local process must not leave a lease permanently busy."""
|
|
|
|
async def scenario() -> None:
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
client = vast_api_module.VastApiClient(state.api_key, base_url=base_url)
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=client,
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
lease = await manager.ensure_lease(_profile(vast_models_module))
|
|
manager.begin_activity(lease.instance_id)
|
|
|
|
restarted = vast_leases_module.VastLeaseManager(
|
|
api_client=client,
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
)
|
|
reconciled = await restarted.reconcile()
|
|
|
|
assert reconciled[0].active_invocations == 0
|
|
assert registry.load().leases[0].active_invocations == 0
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_inventory_refresh_preserves_live_activity_count(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A UI inventory refresh must not make an active lease appear idle."""
|
|
|
|
async def scenario() -> None:
|
|
"""Refresh provider state without applying restart reconciliation semantics."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=vast_api_module.VastApiClient(
|
|
state.api_key,
|
|
base_url=base_url,
|
|
),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
lease = await manager.ensure_lease(_profile(vast_models_module))
|
|
manager.begin_activity(lease.instance_id)
|
|
|
|
refreshed = await manager.refresh_owned_leases()
|
|
|
|
assert refreshed[0].active_invocations == 1
|
|
assert manager.registry.load().leases[0].active_invocations == 1
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_manual_destroy_requires_owned_matching_label(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""Never destroy a reused instance identity whose ownership label changed."""
|
|
|
|
async def scenario() -> None:
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=vast_api_module.VastApiClient(
|
|
state.api_key, base_url=base_url
|
|
),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(
|
|
tmp_path
|
|
),
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
lease = await manager.ensure_lease(_profile(vast_models_module))
|
|
state.instances[lease.instance_id]["label"] = "someone-else"
|
|
|
|
with pytest.raises(RuntimeError, match="label no longer matches"):
|
|
await manager.destroy_owned_lease(lease.instance_id)
|
|
|
|
assert lease.instance_id in state.instances
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_destroy_removes_registry_when_instance_vanishes_after_ownership_check(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""A destroy-time not-found response confirms cleanup of the owned lease."""
|
|
|
|
class VanishingOnDestroyClient(vast_api_module.VastApiClient):
|
|
"""Delete the contract immediately before its destroy API response."""
|
|
|
|
def __init__(self, *args: Any, state: Any, **kwargs: Any) -> None:
|
|
"""Retain simulator state for the synthetic provider race."""
|
|
super().__init__(*args, **kwargs)
|
|
self.state = state
|
|
|
|
async def destroy_instance(self, instance_id: int) -> None:
|
|
"""Simulate Vast removing the contract between show and destroy."""
|
|
self.state.instances.pop(instance_id, None)
|
|
raise vast_api_module.VastInstanceNotFoundError(
|
|
f"Vast instance {instance_id} does not exist."
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
"""Verify the provider race does not retain stale local inventory."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=VanishingOnDestroyClient(
|
|
state.api_key,
|
|
base_url=base_url,
|
|
state=state,
|
|
),
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
lease = await manager.ensure_lease(_profile(vast_models_module))
|
|
|
|
assert await manager.destroy_owned_lease(lease.instance_id)
|
|
assert registry.load().leases == ()
|
|
assert lease.instance_id not in state.instances
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_explicit_force_destroy_can_terminate_active_owned_lease(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
vast_simulator_module: Any,
|
|
) -> None:
|
|
"""The destructive UI path may kill active work after exact ownership checks."""
|
|
|
|
async def scenario() -> None:
|
|
"""Acquire, mark busy, and explicitly force-destroy one owned lease."""
|
|
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
|
|
async with _running_simulator(
|
|
vast_simulator_module.create_vast_simulator_app(state)
|
|
) as base_url:
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=vast_api_module.VastApiClient(
|
|
state.api_key,
|
|
base_url=base_url,
|
|
),
|
|
registry=registry,
|
|
owner_id="comfy-owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
startup_timeout_seconds=2.0,
|
|
)
|
|
lease = await manager.ensure_lease(_profile(vast_models_module))
|
|
manager.begin_activity(lease.instance_id)
|
|
|
|
with pytest.raises(RuntimeError, match="active work"):
|
|
await manager.destroy_owned_lease(lease.instance_id)
|
|
assert await manager.destroy_owned_lease(
|
|
lease.instance_id,
|
|
allow_active_work=True,
|
|
)
|
|
assert registry.load().leases == ()
|
|
assert lease.instance_id not in state.instances
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_registry_is_versioned_atomic_and_credential_free(
|
|
tmp_path: Any,
|
|
vast_leases_module: Any,
|
|
) -> None:
|
|
"""Durable lease state must contain no API or SSH credentials."""
|
|
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
|
|
label = vast_leases_module.vast_managed_label(
|
|
"owner", "profile", "b" * 64, "a" * 64
|
|
)
|
|
lease = vast_leases_module.VastLeaseRecord(
|
|
instance_id=9,
|
|
offer_id=7,
|
|
owner_id="owner",
|
|
profile_id="profile",
|
|
profile_name="Default",
|
|
profile_fingerprint="b" * 64,
|
|
runtime_fingerprint="a" * 64,
|
|
label=label,
|
|
actual_status="running",
|
|
ssh_host="ssh.example.invalid",
|
|
ssh_port=2222,
|
|
ssh_direct_host="192.0.2.4",
|
|
ssh_direct_port=40222,
|
|
ssh_proxy_host="ssh7.vast.ai",
|
|
ssh_proxy_port=2222,
|
|
gpu_name="GPU",
|
|
gpu_count=1,
|
|
gpu_ram_mb=24576,
|
|
cpu_ram_mb=65536,
|
|
hourly_cost_usd=0.5,
|
|
created_at_epoch=1.0,
|
|
last_activity_at_epoch=2.0,
|
|
idle_deadline_epoch=3.0,
|
|
)
|
|
|
|
registry.upsert(lease)
|
|
|
|
assert registry.load().leases == (lease,)
|
|
payload = registry.config_path.read_text(encoding="utf-8")
|
|
assert json.loads(payload)["version"] == 1
|
|
assert "api_key" not in payload.casefold()
|
|
assert "private_key" not in payload.casefold()
|
|
assert "instance_api_key" not in payload.casefold()
|
|
|
|
|
|
def test_retention_cost_accounts_for_existing_deadline(
|
|
tmp_path: Any,
|
|
vast_api_module: Any,
|
|
vast_leases_module: Any,
|
|
vast_models_module: Any,
|
|
) -> None:
|
|
"""Automatic placement should charge only the deadline extension on a warm lease."""
|
|
clock = FakeClock()
|
|
manager = vast_leases_module.VastLeaseManager(
|
|
api_client=object(),
|
|
registry=vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path),
|
|
owner_id="owner",
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
launch_spec_factory=_launch_factory(vast_models_module),
|
|
clock=clock,
|
|
)
|
|
profile = _profile(
|
|
vast_models_module,
|
|
maximum_hourly_cost_usd=1.0,
|
|
idle_retention_seconds=3600,
|
|
)
|
|
lease = vast_leases_module.VastLeaseRecord(
|
|
instance_id=9,
|
|
offer_id=7,
|
|
owner_id="owner",
|
|
profile_id="node-17",
|
|
profile_name="default",
|
|
profile_fingerprint=vast_leases_module.vast_profile_fingerprint(profile),
|
|
runtime_fingerprint=_runtime_fingerprint(),
|
|
label=vast_leases_module.vast_managed_label(
|
|
"owner",
|
|
"node-17",
|
|
vast_leases_module.vast_profile_fingerprint(profile),
|
|
_runtime_fingerprint(),
|
|
),
|
|
actual_status="running",
|
|
ssh_host="host",
|
|
ssh_port=22,
|
|
gpu_name="GPU",
|
|
gpu_count=1,
|
|
gpu_ram_mb=24576,
|
|
cpu_ram_mb=65536,
|
|
hourly_cost_usd=0.5,
|
|
created_at_epoch=clock.value,
|
|
last_activity_at_epoch=clock.value,
|
|
idle_deadline_epoch=clock.value + 3500,
|
|
)
|
|
|
|
new_cost = manager.incremental_retention_cost_usd(
|
|
profile,
|
|
existing_lease=None,
|
|
predicted_execution_seconds=60,
|
|
)
|
|
warm_cost = manager.incremental_retention_cost_usd(
|
|
profile,
|
|
existing_lease=lease,
|
|
predicted_execution_seconds=60,
|
|
)
|
|
|
|
assert new_cost == pytest.approx(3660 / 3600)
|
|
assert warm_cost == pytest.approx(160 / 3600 * 0.5)
|