diff --git a/.gitignore b/.gitignore index 254cba3..691657d 100644 --- a/.gitignore +++ b/.gitignore @@ -47,6 +47,7 @@ !/vast_api.py !/vast_config_node.py !/vast_executor.py +!/vast_image_build.py !/vast_leases.py !/vast_models.py !/vast_runtime.py diff --git a/README.md b/README.md index cfed6c3..83f5e8c 100644 --- a/README.md +++ b/README.md @@ -183,7 +183,7 @@ Vast startup uses every available lifecycle field instead of relying only on `ac Direct Vast SSH operations tolerate bounded connection-level closures, resets, and handshake interruptions with jittered exponential retries. 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. An owned lease may be adopted across controller-only fingerprint drift when its configured worker image is unchanged and its worker protocol is still compatible, avoiding a duplicate rental; changing the configured image still requires matching capacity. Stale registry entries are discarded whether Vast reports a missing contract through HTTP 404 or its v0 `instances: null` response. +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 `uv run python scripts/build_vast_worker_image.py --push` once for the current fingerprint, 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. Concurrent workflows share the completed publication. This automatic path requires `uv`, 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. Set credentials and the worker image in the ComfyUI process environment. Neither value is written to workflows, the lease registry, API responses, or logs: diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 2677d66..9961cc0 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -1,5 +1,12 @@ # Learnings +## 2026-08-28: A matching protocol number cannot make a stale Vast worker safe + +- A newly rented Vast instance can still report the same old runtime fingerprint as a destroyed instance because `COMFY_MODAL_VAST_IMAGE` is an immutable registry digest; destroying capacity does not rebuild that image. +- The worker staging CLI gained arguments and commands while the remote protocol remained at version 1. Allowing fingerprint drift whenever the protocol matched therefore admitted a worker that could pass readiness but fail model staging. +- Vast readiness now treats every source fingerprint mismatch as a typed drift error. Queue-time acquisition destroys the stale lease, coordinates a single local `uv run python scripts/build_vast_worker_image.py --push`, streams bounded subprocess output through the existing per-environment setup status, adopts the returned digest in the current process, and retries with fresh capacity. +- Automatic publication must remain fail-closed: missing `uv`, Docker, registry authentication, a failed build/push, a missing digest, or a second mismatched worker cancels preparation with the manual build command and configuration steps. + - 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. ## 2026-08-28 diff --git a/scripts/build_vast_worker_image.py b/scripts/build_vast_worker_image.py index d21137b..f7186a5 100644 --- a/scripts/build_vast_worker_image.py +++ b/scripts/build_vast_worker_image.py @@ -277,9 +277,12 @@ def build_image( _docker_build_command(image_tag, identity.fingerprint), input_payload=context, ) + logger.info("Finished building Vast worker image tag=%s.", image_tag) if not push: return image_tag + logger.info("Pushing Vast worker image tag=%s.", image_tag) _run(("docker", "push", image_tag)) + logger.info("Resolving immutable digest for Vast worker image tag=%s.", image_tag) repository_digests = _run( ( "docker", diff --git a/tests/conftest.py b/tests/conftest.py index 40e2768..b9a4302 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -273,6 +273,12 @@ def vast_runtime_module(extension_package: object) -> object: return importlib.import_module(f"{PACKAGE_NAME}.vast_runtime") +@pytest.fixture(scope="session") +def vast_image_build_module(extension_package: object) -> object: + """Return automatic Vast worker image publication helpers.""" + return importlib.import_module(f"{PACKAGE_NAME}.vast_image_build") + + @pytest.fixture(scope="session") def vast_executor_module(extension_package: object) -> object: """Return the direct Vast remote executor client.""" diff --git a/tests/test_vast_image_build.py b/tests/test_vast_image_build.py new file mode 100644 index 0000000..e26afca --- /dev/null +++ b/tests/test_vast_image_build.py @@ -0,0 +1,91 @@ +"""Tests for workflow-triggered Vast worker image publication.""" + +from __future__ import annotations + +import io +from pathlib import Path +from typing import Any + +import pytest + + +class FakeProcess: + """Expose deterministic merged subprocess output and an exit status.""" + + def __init__(self, output: str, return_code: int = 0) -> None: + """Initialize the readable output and return code.""" + self.stdout = io.StringIO(output) + self.return_code = return_code + + def wait(self) -> int: + """Return the configured process status.""" + return self.return_code + + +def test_builder_streams_progress_and_returns_digest( + vast_image_build_module: Any, + monkeypatch: Any, + tmp_path: Path, +) -> None: + """Build output should reach workflow status before the digest is returned.""" + digest = "ghcr.io/example/worker@sha256:" + "a" * 64 + observed: dict[str, Any] = {} + + def fake_popen(command: Any, **kwargs: Any) -> FakeProcess: + """Record safe process inputs and return a successful build.""" + observed["command"] = tuple(command) + observed.update(kwargs) + return FakeProcess( + "\x1b[32mBuilding layer 1/3\x1b[0m\n" + f"Pushing sha256:{'d' * 64}\n" + f"COMFY_MODAL_VAST_IMAGE={digest}\n" + ) + + monkeypatch.setattr(vast_image_build_module.subprocess, "Popen", fake_popen) + vast_image_build_module._BUILT_IMAGES_BY_FINGERPRINT.clear() + comfyui_root = tmp_path / "ComfyUI" + builder = vast_image_build_module.VastWorkerImageBuilder( + repo_root=tmp_path, + comfyui_root=comfyui_root, + environment={"PATH": "/usr/bin"}, + ) + statuses: list[str] = [] + + result = builder.build_and_push("f" * 64, status_callback=statuses.append) + + assert result == digest + assert observed["command"] == vast_image_build_module.VAST_IMAGE_BUILD_COMMAND + assert observed["cwd"] == tmp_path + assert observed["env"]["COMFYUI_ROOT"] == str(comfyui_root) + assert observed["env"]["BUILDKIT_PROGRESS"] == "plain" + assert statuses == [ + "Building the current Vast worker image", + "Vast image build: Building layer 1/3", + "Vast image build: Pushing sha256:[redacted]", + "Published the current Vast worker image", + ] + + +def test_builder_failure_requires_the_documented_manual_command( + vast_image_build_module: Any, + monkeypatch: Any, + tmp_path: Path, +) -> None: + """Missing local build prerequisites should cancel with recovery guidance.""" + monkeypatch.setattr( + vast_image_build_module.subprocess, + "Popen", + lambda *_args, **_kwargs: FakeProcess("docker: command not found\n", 127), + ) + vast_image_build_module._BUILT_IMAGES_BY_FINGERPRINT.clear() + builder = vast_image_build_module.VastWorkerImageBuilder( + repo_root=tmp_path, + comfyui_root=None, + environment={"PATH": "/usr/bin"}, + ) + + with pytest.raises( + vast_image_build_module.VastWorkerImageBuildError, + match=r"uv run python scripts/build_vast_worker_image\.py --push", + ): + builder.build_and_push("e" * 64) diff --git a/tests/test_vast_runtime.py b/tests/test_vast_runtime.py index d949b63..43a7c58 100644 --- a/tests/test_vast_runtime.py +++ b/tests/test_vast_runtime.py @@ -156,19 +156,26 @@ def test_runtime_manager_rejects_image_fingerprint_drift( ), ) - with pytest.raises(Exception, match="fingerprint mismatch"): + with pytest.raises( + vast_runtime_module.VastRuntimeFingerprintDriftError, + match="source fingerprint drift", + ) as raised: manager.ensure_worker() + assert raised.value.expected_fingerprint == "a" * 64 + assert raised.value.actual_fingerprint == "b" * 64 + assert raised.value.protocol_version == 0 -def test_runtime_manager_accepts_same_protocol_fingerprint_drift( + +def test_runtime_manager_rejects_same_protocol_fingerprint_drift( vast_runtime_module: Any, ) -> None: - """Controller-only source changes may reuse the same protocol-compatible image.""" + """A protocol match must not hide source-level worker incompatibility.""" manager = vast_runtime_module.VastRuntimeManager( runner=FakeRunner( [ { - "protocol_version": vast_runtime_module.REMOTE_PROTOCOL_VERSION, + "protocol_version": 1, "runtime_fingerprint": "b" * 64, "worker_socket_ready": True, } @@ -180,9 +187,12 @@ def test_runtime_manager_accepts_same_protocol_fingerprint_drift( ), ) - info = manager.ensure_worker() + with pytest.raises( + vast_runtime_module.VastRuntimeFingerprintDriftError + ) as raised: + manager.ensure_worker() - assert info["worker_socket_ready"] is True + assert raised.value.protocol_version == 1 def test_runtime_manager_logs_changed_ssh_readiness_diagnostics( diff --git a/tests/test_vast_service.py b/tests/test_vast_service.py index d9347d1..691b504 100644 --- a/tests/test_vast_service.py +++ b/tests/test_vast_service.py @@ -379,3 +379,104 @@ def test_acquire_replaces_instance_that_disappears_before_ssh( 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 + + def initialize_runtime(lease: Any) -> None: + """Reject only the worker baked into the first lease.""" + 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()), + 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" diff --git a/vast_image_build.py b/vast_image_build.py new file mode 100644 index 0000000..745a319 --- /dev/null +++ b/vast_image_build.py @@ -0,0 +1,179 @@ +"""Build and publish a current worker image after Vast source drift.""" + +from __future__ import annotations + +import logging +import os +import re +import subprocess +import threading +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import IO + +logger = logging.getLogger(__name__) + +VAST_IMAGE_BUILD_COMMAND = ( + "uv", + "run", + "python", + "scripts/build_vast_worker_image.py", + "--push", +) +_IMAGE_RESULT_PREFIX = "COMFY_MODAL_VAST_IMAGE=" +_DIGEST_REFERENCE_PATTERN = re.compile(r"^\S+@sha256:[0-9a-f]{64}$") +_DIGEST_IN_TEXT_PATTERN = re.compile(r"sha256:[0-9a-f]{64}") +_ANSI_ESCAPE_PATTERN = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]") +_MAXIMUM_STATUS_LINE_CHARACTERS = 300 +_BUILD_LOCK = threading.Lock() +_BUILT_IMAGES_BY_FINGERPRINT: dict[str, str] = {} + +VastImageBuildStatusCallback = Callable[[str], None] + + +class VastWorkerImageBuildError(RuntimeError): + """Report that the automatic Vast worker image build could not finish.""" + + +@dataclass(frozen=True) +class VastWorkerImageBuilder: + """Run the repository image builder and stream bounded progress messages.""" + + repo_root: Path + comfyui_root: Path | None + environment: Mapping[str, str] | None = None + + def build_and_push( + self, + expected_fingerprint: str, + *, + status_callback: VastImageBuildStatusCallback | None = None, + ) -> str: + """Build once per fingerprint and return the published digest reference.""" + with _BUILD_LOCK: + cached = _BUILT_IMAGES_BY_FINGERPRINT.get(expected_fingerprint) + if cached is not None: + self._emit( + status_callback, + "Using the worker image published by another workflow", + ) + return cached + image = self._run_build(status_callback=status_callback) + _BUILT_IMAGES_BY_FINGERPRINT[expected_fingerprint] = image + return image + + def _run_build( + self, + *, + status_callback: VastImageBuildStatusCallback | None, + ) -> str: + """Execute the image builder and extract its immutable result reference.""" + command_text = " ".join(VAST_IMAGE_BUILD_COMMAND) + self._emit(status_callback, "Building the current Vast worker image") + environment = dict(os.environ if self.environment is None else self.environment) + environment.setdefault("BUILDKIT_PROGRESS", "plain") + if self.comfyui_root is not None: + environment["COMFYUI_ROOT"] = str(self.comfyui_root) + try: + process = subprocess.Popen( + VAST_IMAGE_BUILD_COMMAND, + cwd=self.repo_root, + env=environment, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + ) + except OSError as exc: + raise VastWorkerImageBuildError( + self._manual_build_message( + f"Unable to start {VAST_IMAGE_BUILD_COMMAND[0]!r}: {exc}" + ) + ) from exc + image_reference, recent_output = self._consume_output( + process.stdout, + status_callback=status_callback, + ) + return_code = process.wait() + if return_code != 0: + diagnostic = recent_output[-1] if recent_output else "no build output" + raise VastWorkerImageBuildError( + self._manual_build_message( + f"{command_text} exited with status {return_code}: {diagnostic}" + ) + ) + if image_reference is None: + raise VastWorkerImageBuildError( + self._manual_build_message( + "The build completed without returning an immutable image digest." + ) + ) + self._emit(status_callback, "Published the current Vast worker image") + return image_reference + + def _consume_output( + self, + output: IO[str] | None, + *, + status_callback: VastImageBuildStatusCallback | None, + ) -> tuple[str | None, tuple[str, ...]]: + """Forward subprocess lines and return its digest plus recent diagnostics.""" + if output is None: + return None, () + image_reference: str | None = None + recent_output: list[str] = [] + for raw_line in output: + line = raw_line.strip() + if not line: + continue + if line.startswith(_IMAGE_RESULT_PREFIX): + candidate = line.removeprefix(_IMAGE_RESULT_PREFIX).strip() + if _DIGEST_REFERENCE_PATTERN.fullmatch(candidate): + image_reference = candidate + continue + bounded = self._safe_status_line(line) + recent_output.append(bounded) + del recent_output[:-10] + logger.info("Vast worker image build: %s", bounded) + self._emit(status_callback, f"Vast image build: {bounded}") + return image_reference, tuple(recent_output) + + @staticmethod + def _safe_status_line(line: str) -> str: + """Bound one progress line and remove terminal codes and exact digests.""" + without_ansi = _ANSI_ESCAPE_PATTERN.sub("", line) + redacted = _DIGEST_IN_TEXT_PATTERN.sub("sha256:[redacted]", without_ansi) + printable = "".join( + character for character in redacted if character.isprintable() + ) + return printable[:_MAXIMUM_STATUS_LINE_CHARACTERS] + + def _manual_build_message(self, cause: str) -> str: + """Return an actionable error without leaking subprocess environment data.""" + return ( + f"Automatic Vast worker image publication failed. {cause} Run `" + + " ".join(VAST_IMAGE_BUILD_COMMAND) + + "` from " + + str(self.repo_root) + + ", set COMFY_MODAL_VAST_IMAGE to the printed digest, restart ComfyUI, " + "and retry the workflow." + ) + + @staticmethod + def _emit( + status_callback: VastImageBuildStatusCallback | None, + message: str, + ) -> None: + """Send one build update when workflow progress is available.""" + if status_callback is not None: + status_callback(message) + + +__all__ = [ + "VAST_IMAGE_BUILD_COMMAND", + "VastWorkerImageBuildError", + "VastWorkerImageBuilder", +] diff --git a/vast_runtime.py b/vast_runtime.py index 5610339..af76115 100644 --- a/vast_runtime.py +++ b/vast_runtime.py @@ -13,13 +13,11 @@ from pathlib import PurePosixPath from typing import Any, Callable, Mapping if __package__: - from .remote_protocol import REMOTE_PROTOCOL_VERSION from .remote.vast_watchdog import VastWatchdogSnapshot from .vast_leases import VastLeaseRecord from .vast_models import VastInstanceLaunchSpec, VastResourceProfile from .vast_ssh import VastSshError, VastSshRunner, VastSshVolumeBackend else: # pragma: no cover - direct debugging imports. - from remote_protocol import REMOTE_PROTOCOL_VERSION from remote.vast_watchdog import VastWatchdogSnapshot from vast_leases import VastLeaseRecord from vast_models import VastInstanceLaunchSpec, VastResourceProfile @@ -36,6 +34,27 @@ DEFAULT_VAST_WATCHDOG_STATE_PATH = PurePosixPath( ) +class VastRuntimeFingerprintDriftError(VastSshError): + """Report that a published Vast worker does not match the local runtime.""" + + def __init__( + self, + *, + expected_fingerprint: str, + actual_fingerprint: str, + protocol_version: int, + ) -> None: + """Retain the compared identities for automated image replacement.""" + self.expected_fingerprint = expected_fingerprint + self.actual_fingerprint = actual_fingerprint + self.protocol_version = protocol_version + super().__init__( + "Vast worker image source fingerprint drift: expected " + f"{expected_fingerprint[:12]}, found " + f"{actual_fingerprint[:12] or 'empty'}, protocol {protocol_version}." + ) + + def _ssh_permission_repair_lines() -> tuple[str, ...]: """Return idempotent repairs for the SSH key path injected by Vast.""" return ( @@ -180,19 +199,10 @@ class VastRuntimeManager: actual_fingerprint = str(info.get("runtime_fingerprint") or "") if actual_fingerprint != self.configuration.runtime_fingerprint: actual_protocol = int(info.get("protocol_version") or 0) - if actual_protocol != REMOTE_PROTOCOL_VERSION: - raise VastSshError( - "Vast worker image fingerprint mismatch: expected " - f"{self.configuration.runtime_fingerprint[:12]}, found " - f"{actual_fingerprint[:12] or 'empty'}, with incompatible " - f"protocol {actual_protocol}." - ) - logger.warning( - "Reusing Vast worker with compatible protocol despite source " - "fingerprint drift expected=%s actual=%s protocol=%d.", - self.configuration.runtime_fingerprint[:12], - actual_fingerprint[:12] or "empty", - actual_protocol, + raise VastRuntimeFingerprintDriftError( + expected_fingerprint=self.configuration.runtime_fingerprint, + actual_fingerprint=actual_fingerprint, + protocol_version=actual_protocol, ) if bool(info.get("worker_socket_ready")): logger.info( @@ -289,5 +299,6 @@ __all__ = [ "DEFAULT_VAST_WATCHDOG_STATE_PATH", "VAST_IMAGE_ENV", "VastRuntimeConfiguration", + "VastRuntimeFingerprintDriftError", "VastRuntimeManager", ] diff --git a/vast_service.py b/vast_service.py index 548d49b..0439c9a 100644 --- a/vast_service.py +++ b/vast_service.py @@ -28,6 +28,11 @@ if __package__: 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, @@ -35,7 +40,12 @@ if __package__: vast_profile_fingerprint, ) from .vast_models import VastInstance, VastOffer, VastResourceProfile - from .vast_runtime import VastRuntimeConfiguration, VastRuntimeManager + from .vast_runtime import ( + VAST_IMAGE_ENV, + VastRuntimeConfiguration, + VastRuntimeFingerprintDriftError, + VastRuntimeManager, + ) from .vast_ssh import ( VastSshError, VastSshRunner, @@ -58,6 +68,11 @@ else: # pragma: no cover - direct debugging imports. 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, @@ -65,7 +80,12 @@ else: # pragma: no cover - direct debugging imports. vast_profile_fingerprint, ) from vast_models import VastInstance, VastOffer, VastResourceProfile - from vast_runtime import VastRuntimeConfiguration, VastRuntimeManager + from vast_runtime import ( + VAST_IMAGE_ENV, + VastRuntimeConfiguration, + VastRuntimeFingerprintDriftError, + VastRuntimeManager, + ) from vast_ssh import ( VastSshError, VastSshRunner, @@ -144,6 +164,8 @@ class VastService: 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 @@ -154,6 +176,11 @@ class VastService: 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), + ) + self.update_process_environment = update_process_environment self.huggingface_asset_registry = ( HuggingFaceAssetRegistry.for_user_directory(self.user_directory) ) @@ -162,10 +189,17 @@ class VastService: user_directory=self.user_directory, comfyui_root=getattr(settings, "comfyui_root", None), ) - self.lease_manager = VastLeaseManager( - api_client=api_client, - registry=registry, - owner_id=settings.app_name, + 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, @@ -223,6 +257,12 @@ class VastService: 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, + environment=(source if environment is not None else None), + ), + update_process_environment=environment is None, ) async def quote_best_profile( @@ -427,6 +467,7 @@ class VastService: 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() @@ -444,6 +485,48 @@ class VastService: status_callback("Initializing Vast.ai worker") try: await asyncio.to_thread(self._initialize_runtime, 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) @@ -509,6 +592,26 @@ class VastService: 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,