From ee3f5e17c5a558402ccc4fabf30ab2dd719ce2f2 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Mon, 24 Aug 2026 13:37:34 -0700 Subject: [PATCH] Discard Vast leases that fail worker startup --- README.md | 2 +- docs/LEARNINGS.md | 1 + tests/test_vast_service.py | 69 ++++++++++++++++++++++++++++++++++++++ vast_service.py | 36 +++++++++++++++++++- 4 files changed, 106 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 63c60e2..100d4af 100644 --- a/README.md +++ b/README.md @@ -228,7 +228,7 @@ Subgraph containers inherit a derived summary decoration (marked with a `Σ` bad Active remote nodes show a small progress panel with per-bar `it/s` rates (LLM nodes show `tok/s`, token counts, and TTFT). When a prompt is scheduled across more than one remote environment, the panel also names where each component is running — the Modal container id or the SSH host name. Preview images and UI payloads from remote nodes stream back into the local session while the component is still running, and progress state is reconciled against queue history when the browser regains focus so cancelled prompts do not leave stale bars behind. -Remote planning and capacity acquisition count as active queue work even though they happen before the rewritten prompt enters ComfyUI's native execution queue. During Vast startup, the canvas receives readable updates for rental, worker-image download, SSH readiness, and runtime initialization. A provider-side terminal image-pull failure is reported immediately rather than appearing as an empty queue until the startup timeout expires. +Remote planning and capacity acquisition count as active queue work even though they happen before the rewritten prompt enters ComfyUI's native execution queue. During Vast startup, the canvas receives readable updates for rental, worker-image download, SSH readiness, and runtime initialization. A provider-side terminal image-pull failure is reported immediately rather than appearing as an empty queue until the startup timeout expires. A rental that never produces an SSH-usable worker is marked unusable and destroyed after the bounded initialization wait so a later workflow cannot silently reuse it. For workflows using a **Remote Execution Configurator**, that node is also the primary planning surface. As soon as scheduling finishes, the scheduler result is streamed to the panel before any Vast rental or SSH-readiness wait begins. The browser retains that plan across late node-panel mounts, tab refocus, and queue failure, so configuration names, selected targets, predicted completion time, and predicted cost remain inspectable even when capacity never becomes ready. Live preparation and execution status use the same durable event stream. Legacy workflows without a configurator continue to use the global status pill. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 355644f..69ded46 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -535,3 +535,4 @@ - A queue response cannot be the first carrier of a workflow plan when the same request performs capacity acquisition. If Vast rental or SSH initialization stalls, that response never arrives; publish the side-effect-free scheduler result before acquisition and retain it in prompt-scoped browser state. - ComfyUI DOM widgets may mount after their prompt events arrive or remain temporarily detached during canvas virtualization. Store plan and status independently of the element, replay the newest matching prompt after mount, and set a mounted guard only after `addDOMWidget` succeeds so a transient mount failure remains retryable. - A persisted Vast lease registry describes controller intent, not live SSH health. Worker readiness loops must log changed probe diagnostics as they occur and include instance/environment identity at the service boundary; otherwise a provider closing the connection during key exchange appears as a silent 15-minute initialization wait. +- Provider `running` plus an assigned SSH endpoint is not sufficient evidence that rented capacity is reusable. If bounded runtime initialization fails, mark the lease draining before cleanup, destroy the owned instance, and preserve the original failure even if destruction itself fails; otherwise every later plan can select the same poisoned registry entry. diff --git a/tests/test_vast_service.py b/tests/test_vast_service.py index c890155..521d018 100644 --- a/tests/test_vast_service.py +++ b/tests/test_vast_service.py @@ -8,6 +8,8 @@ from pathlib import Path from types import SimpleNamespace from typing import Any, AsyncIterator +import pytest + from aiohttp import web @@ -195,3 +197,70 @@ def test_prefetch_deduplicates_effective_profiles_and_searches_in_parallel( assert api_client.peak == 2 asyncio.run(scenario()) + + +def test_acquire_destroys_lease_when_worker_initialization_fails( + vast_service_module: Any, +) -> None: + """An SSH-unusable rental must be drained instead of reused next 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 and destroy one deterministic unusable rental.""" + + def __init__(self, lease: Any) -> None: + """Retain the fake lease and destruction log.""" + self.lease = lease + self.destroyed_instance_ids: list[int] = [] + + async def ensure_lease(self, profile: Any, **kwargs: Any) -> Any: + """Return the selected lease without provider calls.""" + del profile, kwargs + return self.lease + + 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 + + lease = SimpleNamespace(instance_id=42, environment_id="vast:test:42") + lease_manager = FakeLeaseManager(lease) + registry = FakeRegistry() + service = object.__new__(vast_service_module.VastService) + service.lease_manager = lease_manager + service.registry = registry + + def fail_runtime(_lease: Any) -> None: + """Simulate the observed provider SSH key-exchange failure.""" + raise vast_service_module.VastSshError( + "kex_exchange_identification: Connection closed by remote host" + ) + + service._initialize_runtime = fail_runtime + messages: list[str] = [] + + with pytest.raises( + vast_service_module.VastSshError, + match="Connection closed by remote host", + ): + asyncio.run( + service.acquire( + SimpleNamespace(profile=SimpleNamespace()), + status_callback=messages.append, + ) + ) + + assert registry.updated_instance_ids == [42] + assert lease_manager.destroyed_instance_ids == [42] + assert messages[-1] == "Vast.ai worker initialization failed" diff --git a/vast_service.py b/vast_service.py index 941a836..5a073fc 100644 --- a/vast_service.py +++ b/vast_service.py @@ -412,11 +412,45 @@ class VastService: ) if status_callback is not None: status_callback("Initializing Vast.ai worker") - await asyncio.to_thread(self._initialize_runtime, lease) + try: + await asyncio.to_thread(self._initialize_runtime, lease) + except (TimeoutError, VastSshError, ValueError) as exc: + if status_callback is not None: + status_callback("Vast.ai worker initialization failed") + await self._discard_failed_runtime_lease(lease, str(exc)) + raise if status_callback is not None: status_callback("Vast.ai worker is ready") return lease + async def _discard_failed_runtime_lease( + self, + lease: VastLeaseRecord, + error_message: str, + ) -> None: + """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: + 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 + logger.info( + "Destroyed unusable Vast worker instance_id=%d after initialization failure.", + lease.instance_id, + ) + 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]