From 4e1743b8bcfe923959e6ff536272674c668ee04d Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sat, 29 Aug 2026 08:35:15 -0700 Subject: [PATCH] Back off Vast SSH readiness retries --- README.md | 2 +- docs/LEARNINGS.md | 1 + tests/test_vast_runtime.py | 62 +++++++++++++++++++++++++++++++++++++- tests/test_vast_service.py | 24 +++++++++++++-- vast_runtime.py | 59 +++++++++++++++++++++++++++++++----- vast_service.py | 9 +++++- 6 files changed, 144 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index dd16f43..f5f612e 100644 --- a/README.md +++ b/README.md @@ -181,7 +181,7 @@ At queue time the planner raises **Any** or lower VRAM and RAM floors to whateve Vast startup uses every available lifecycle field instead of relying only on `actual_status`. Legitimate image-download and loading states retain the full configured readiness window, while an instance that continuously reports no lifecycle state, provider message, or SSH host fails after a two-minute grace period. When Vast publishes both a direct public mapping for container port 22 and an SSH proxy, the controller uses the same direct endpoint shown in the Vast console; this avoids stale proxy endpoints that close every handshake. Vast's SSH launch mode injects the account key before running the configured `onstart` script; that script normalizes `/root`, `.ssh`, and `authorized_keys` ownership and modes before starting the worker, preventing `sshd` StrictModes failures inherited from a custom image. Before every worker-readiness SSH probe, the controller rechecks that the exact Vast contract still exists. A contract that disappears during provider startup or worker initialization is removed from the local registry immediately, its marketplace offer is excluded, and the next compatible offer is cold-started as a replacement. A contract that remains non-SSH-ready for the configured startup window, or cannot initialize its worker afterward, is marked unusable and terminated through the Vast API before the controller cold-starts one replacement from a different offer. A reused contract consumes only the remainder of its original startup window, so restarting ComfyUI cannot restart the full wait. If the replacement also fails, acquisition stops with the failed instance IDs, cleanup outcome, and final provider/setup cause instead of entering an unbounded rental loop. Worker-readiness polling makes one connection per interval rather than nesting the normal command retry budget inside every probe. Commands on an established worker and interrupted asset transfers still retry transient connection closures. -Direct Vast SSH operations tolerate bounded connection-level closures, resets, and handshake interruptions with jittered exponential retries. Vast's direct public port mapping remains preferred, but the controller also retains the provider's `ssh*.vast.ai` proxy endpoint; if the direct readiness probe fails while the proxy succeeds, the proxy becomes the persisted endpoint for staging, execution, transfers, and watchdog updates. Authentication rejection and deterministic remote-command failures remain immediate. Streamed uploads reopen the local source from byte zero on every attempt, and the remote atomic writer verifies the expected byte count before publishing, so a killed connection cannot leave a truncated file at the content-addressed destination. +Direct Vast SSH operations tolerate bounded connection-level closures, resets, and handshake interruptions with jittered exponential retries. Worker-readiness SSH probes start one second apart and multiply the interval by 1.5 after every unsuccessful attempt, with each failure and next delay streamed to the workflow UI. Vast's direct public port mapping remains preferred, but the controller also retains the provider's `ssh*.vast.ai` proxy endpoint; if the direct readiness probe fails while the proxy succeeds, the proxy becomes the persisted endpoint for staging, execution, transfers, and watchdog updates. Authentication rejection and deterministic remote-command failures remain immediate. Streamed uploads reopen the local source from byte zero on every attempt, and the remote atomic writer verifies the expected byte count before publishing, so a killed connection cannot leave a truncated file at the content-addressed destination. Cancelling a workflow also cancels queue-time remote preparation. Intentional cancellation remains an interrupted workflow while Vast capacity is being requested or initialized; it is not recast as a provider-acquisition validation failure. File hashing checks the prompt-scoped cancellation signal between chunks, and active Vast SSH uploads or Hugging Face downloads terminate their local SSH process instead of continuing behind ComfyUI's cancelling status. Vast Hugging Face acquisition runs from a small content-addressed tool bundle uploaded by the controller, so an older worker image cannot force a multi-gigabyte fallback upload merely because it lacks the newest materializer module. The launch environment does not overwrite the runtime fingerprint baked into the published worker image. A fingerprint mismatch is never accepted solely because the protocol number matches: the controller terminates the stale lease, runs the image-build script once for the current fingerprint with the exact Python interpreter hosting ComfyUI, streams the build and push output through ComfyUI's remote-preparation progress, adopts the returned immutable digest for the current process, and rents fresh capacity. Invoking `sys.executable` directly avoids assumptions about virtualenv location or activation and prevents an obsolete copied `.venv` from hijacking the build. The build subprocess receives the resolved workflow GPU override and must report the same source fingerprint captured by the parent; if files changed afterward, publication stops before Docker work or another rental and asks for a restart. Concurrent workflows share the completed publication. This automatic path requires a running Docker service and registry push credentials on the ComfyUI host; if any prerequisite or publication step fails, the workflow stops with the exact manual recovery command instead of executing on stale source. Stale registry entries are discarded whether Vast reports a missing contract through HTTP 404 or its v0 `instances: null` response. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index e0b3910..a2a9081 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -10,6 +10,7 @@ - A custom-node deployment can contain a stale `.venv` whose interpreter symlink and `pyvenv.cfg` describe different Python installations. Automatic image publication must invoke the build script directly with `sys.executable` so neither project discovery nor a presumed virtualenv location can select that copied environment, and failure summaries must skip traceback footers such as `` in favor of the substantive error line. - A long-running ComfyUI process can retain fingerprint A while hot-deployed source on disk has fingerprint B. Pass A into the build subprocess, have it report and compare B before exporting a Docker context, and refuse to publish or rent until ComfyUI restarts; adopting B's digest while retaining A in runtime configuration guarantees a second drift failure. - A workflow-level Modal GPU override is part of the shared remote runtime identity even when the selected capacity is Vast. A fresh image-build process otherwise falls back to `COMFY_MODAL_GPU` from the host environment, so explicitly propagate the parent's resolved `settings.modal_gpu` into the child environment before comparing fingerprints or selecting accelerator packages. +- Vast may advertise a running contract before sshd accepts connections. Probe worker readiness after one second, back off each subsequent failure by 1.5, and stream the SSH diagnostic plus next delay through the environment setup callback so users can distinguish expected SSH warm-up from a wedged workflow. - The pinned vLLM wheel is large enough that pip's default five resumptions and 15-second socket timeout can be exhausted by an unstable connection during an SSH/Vast worker build. Give the accelerator-wheel layer an extended timeout and bounded resume budget while leaving earlier package layers unchanged so Docker can reuse their successful cache entries. diff --git a/tests/test_vast_runtime.py b/tests/test_vast_runtime.py index 4badb97..18b55cc 100644 --- a/tests/test_vast_runtime.py +++ b/tests/test_vast_runtime.py @@ -191,10 +191,12 @@ def test_runtime_manager_selects_proxy_when_direct_ssh_fails( ] ) selected: list[bool] = [] + statuses: list[str] = [] manager = vast_runtime_module.VastRuntimeManager( runner=RefusedRunner(), fallback_runner=proxy_runner, fallback_selected=lambda: selected.append(True), + status_callback=statuses.append, configuration=vast_runtime_module.VastRuntimeConfiguration( image="worker", runtime_fingerprint=fingerprint, @@ -207,6 +209,9 @@ def test_runtime_manager_selects_proxy_when_direct_ssh_fails( assert manager.runner is proxy_runner assert manager.fallback_runner is None assert selected == [True] + assert statuses == [ + "Vast direct SSH failed; continuing through the Vast SSH proxy" + ] def test_runtime_manager_rejects_same_protocol_fingerprint_drift( @@ -254,6 +259,7 @@ def test_runtime_manager_logs_changed_ssh_readiness_diagnostics( ) current_time = [0.0] + statuses: list[str] = [] def advance(seconds: float) -> None: """Advance the deterministic monotonic readiness clock.""" @@ -269,6 +275,7 @@ def test_runtime_manager_logs_changed_ssh_readiness_diagnostics( ), monotonic=lambda: current_time[0], sleep=advance, + status_callback=statuses.append, ) with caplog.at_level(logging.WARNING): @@ -281,9 +288,62 @@ def test_runtime_manager_logs_changed_ssh_readiness_diagnostics( if "Vast worker readiness probe" in record.message ] assert readiness_messages == [ - "Vast worker readiness probe attempt=1 failed: " + "Vast worker readiness probe attempt=1 failed; retrying in 1.00s: " "kex_exchange_identification: Connection closed by remote host" ] + assert statuses == [ + "Vast SSH attempt 1 failed: kex_exchange_identification: Connection " + "closed by remote host; retrying in 1.00s", + "Vast SSH attempt 2 failed: kex_exchange_identification: Connection " + "closed by remote host; retrying in 1.00s", + ] + + +def test_runtime_manager_uses_exponential_ssh_readiness_backoff( + vast_runtime_module: Any, +) -> None: + """Early SSH probes should start at one second and grow by 1.5 each time.""" + + class StartingRunner: + """Fail three SSH probes before returning ready runtime metadata.""" + + def __init__(self) -> None: + """Initialize the number of refused connections.""" + self.failures_remaining = 3 + + def run(self, argv: Any, **kwargs: Any) -> FakeResult: + """Refuse early probes and then report the worker ready.""" + del argv, kwargs + if self.failures_remaining > 0: + self.failures_remaining -= 1 + raise vast_runtime_module.VastSshError("Connection refused") + return FakeResult( + { + "runtime_fingerprint": "a" * 64, + "worker_socket_ready": True, + } + ) + + delays: list[float] = [] + statuses: list[str] = [] + manager = vast_runtime_module.VastRuntimeManager( + runner=StartingRunner(), + configuration=vast_runtime_module.VastRuntimeConfiguration( + image="worker", + runtime_fingerprint="a" * 64, + ), + sleep=delays.append, + status_callback=statuses.append, + ) + + manager.ensure_worker() + + assert delays == [1.0, 1.5, 2.25] + assert statuses == [ + "Vast SSH attempt 1 failed: Connection refused; retrying in 1.00s", + "Vast SSH attempt 2 failed: Connection refused; retrying in 1.50s", + "Vast SSH attempt 3 failed: Connection refused; retrying in 2.25s", + ] def test_runtime_manager_validates_instance_before_ssh_probe( diff --git a/tests/test_vast_service.py b/tests/test_vast_service.py index 2b16410..896e0e3 100644 --- a/tests/test_vast_service.py +++ b/tests/test_vast_service.py @@ -287,9 +287,16 @@ def test_acquire_replaces_lease_when_worker_initialization_fails( service.lease_manager = lease_manager service.registry = registry - def initialize_runtime(lease: Any) -> None: + 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" ) @@ -308,6 +315,7 @@ def test_acquire_replaces_lease_when_worker_initialization_fails( 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" @@ -364,8 +372,13 @@ def test_acquire_replaces_instance_that_disappears_before_ssh( service.lease_manager = lease_manager service.registry = registry - def initialize_runtime(lease: Any) -> None: + 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." @@ -460,8 +473,13 @@ def test_acquire_rebuilds_image_and_replaces_fingerprint_drift( adopted_images: list[str] = [] service._adopt_runtime_image = adopted_images.append - def initialize_runtime(lease: Any) -> None: + 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, diff --git a/vast_runtime.py b/vast_runtime.py index 5a694a8..eb587fc 100644 --- a/vast_runtime.py +++ b/vast_runtime.py @@ -81,7 +81,8 @@ class VastRuntimeConfiguration: remote_storage_root: PurePosixPath = DEFAULT_VAST_REMOTE_STORAGE_ROOT watchdog_state_path: PurePosixPath = DEFAULT_VAST_WATCHDOG_STATE_PATH startup_timeout_seconds: float = 900.0 - readiness_poll_seconds: float = 2.0 + readiness_poll_seconds: float = 1.0 + readiness_backoff_multiplier: float = 1.5 def __post_init__(self) -> None: """Validate launch and fingerprint values.""" @@ -99,8 +100,21 @@ class VastRuntimeConfiguration: ): if not path.is_absolute() or path == PurePosixPath("/"): raise ValueError(f"{field_name} must be an absolute non-root path.") - if self.startup_timeout_seconds <= 0 or self.readiness_poll_seconds <= 0: + if any( + not math.isfinite(value) or value <= 0 + for value in ( + self.startup_timeout_seconds, + self.readiness_poll_seconds, + ) + ): raise ValueError("Vast runtime readiness timing must be positive.") + if ( + not math.isfinite(self.readiness_backoff_multiplier) + or self.readiness_backoff_multiplier < 1.0 + ): + raise ValueError( + "Vast readiness backoff multiplier must be finite and at least 1." + ) @classmethod def from_environment( @@ -178,6 +192,7 @@ class VastRuntimeManager: fallback_runner: VastSshRunner | None = None fallback_selected: Callable[[], None] | None = None instance_validator: Callable[[], None] | None = None + status_callback: Callable[[str], None] | None = None clock: Callable[[], float] = time.time monotonic: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep @@ -189,13 +204,18 @@ class VastRuntimeManager: last_error: str | None = None last_logged_error: str | None = None attempt = 0 + retry_delay = self.configuration.readiness_poll_seconds while self.monotonic() < deadline: attempt += 1 + ssh_error: str | None = None if self.instance_validator is not None: self.instance_validator() try: info = self.runtime_info() - except (VastSshError, ValueError) as exc: + except VastSshError as exc: + last_error = str(exc) + ssh_error = last_error + except ValueError as exc: last_error = str(exc) else: actual_fingerprint = str(info.get("runtime_fingerprint") or "") @@ -213,16 +233,33 @@ class VastRuntimeManager: self.monotonic() - started_at, ) return info - self._start_supervisor() - last_error = "worker socket is not ready" + try: + self._start_supervisor() + except VastSshError as exc: + last_error = str(exc) + ssh_error = last_error + else: + last_error = "worker socket is not ready" + remaining_seconds = max(0.0, deadline - self.monotonic()) + if remaining_seconds <= 0: + break + actual_delay = min(retry_delay, remaining_seconds) if last_error != last_logged_error: logger.warning( - "Vast worker readiness probe attempt=%d failed: %s", + "Vast worker readiness probe attempt=%d failed; retrying in " + "%.2fs: %s", attempt, + actual_delay, last_error, ) last_logged_error = last_error - self.sleep(self.configuration.readiness_poll_seconds) + if ssh_error is not None: + self._emit_status( + f"Vast SSH attempt {attempt} failed: {ssh_error}; retrying in " + f"{actual_delay:.2f}s" + ) + self.sleep(actual_delay) + retry_delay *= self.configuration.readiness_backoff_multiplier raise TimeoutError( "Vast worker did not become ready within " f"{self.configuration.startup_timeout_seconds:.0f}s: " @@ -248,11 +285,19 @@ class VastRuntimeManager: self.fallback_runner = None if self.fallback_selected is not None: self.fallback_selected() + self._emit_status( + "Vast direct SSH failed; continuing through the Vast SSH proxy" + ) logger.warning( "Vast direct SSH failed; continuing through the Vast SSH proxy." ) return payload + def _emit_status(self, message: str) -> None: + """Forward one readiness update to the workflow UI when available.""" + if self.status_callback is not None: + self.status_callback(message) + @staticmethod def _runtime_info(runner: VastSshRunner) -> dict[str, Any]: """Return validated runtime metadata through one exact SSH endpoint.""" diff --git a/vast_service.py b/vast_service.py index c37522f..c510312 100644 --- a/vast_service.py +++ b/vast_service.py @@ -489,6 +489,7 @@ class VastService: initialized_lease = await asyncio.to_thread( self._initialize_runtime, lease, + status_callback=status_callback, ) if initialized_lease is not None: lease = initialized_lease @@ -833,7 +834,12 @@ class VastService: raise KeyError(f"Vast lease {instance_id} disappeared during setup.") return lease - def _initialize_runtime(self, lease: VastLeaseRecord) -> VastLeaseRecord: + 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.", @@ -846,6 +852,7 @@ class VastService: 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()