Reject stale source during Vast image publication

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. 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. 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.
+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 must report the same source fingerprint captured when ComfyUI started; 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.
 
 Set credentials and the worker image in the ComfyUI process environment. Neither value is written to workflows, the lease registry, API responses, or logs:
 

docs/LEARNINGS.md

@@ -8,6 +8,7 @@
 - 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.
 - Vast can report a contract as successful and start the worker/watchdog internally while its advertised direct port mapping continues to refuse connections. Preserve both the direct and `ssh*.vast.ai` proxy endpoints in lease state, try the proxy during every failed direct readiness probe, and persist whichever endpoint first returns valid worker metadata so later operations do not fall back to the dead mapping.
 - 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 `<no Python frame>` 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.
 
 - 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.
 

scripts/build_vast_worker_image.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import argparse
 import logging
+import re
 import subprocess
 import sys
 import tomllib
@@ -26,6 +27,8 @@ logger = logging.getLogger(__name__)
 
 DEFAULT_TAG_TEMPLATE = "ghcr.io/{owner}/comfy-modal-worker:v{version}"
 VAST_WORKER_PLATFORM = "linux/amd64"
+SOURCE_FINGERPRINT_RESULT_PREFIX = "COMFY_MODAL_VAST_SOURCE_FINGERPRINT="
+_SHA256_FINGERPRINT_PATTERN = re.compile(r"^[0-9a-f]{64}$")
 
 
 def _parser() -> argparse.ArgumentParser:
@@ -66,9 +69,27 @@ def _parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Push the built tag and print its registry digest reference.",
     )
+    parser.add_argument(
+        "--expected-fingerprint",
+        type=_runtime_fingerprint,
+        help=(
+            "Abort before building if current source no longer matches this SHA-256 "
+            "runtime fingerprint. Used by automatic publication."
+        ),
+    )
     return parser
 
 
+def _runtime_fingerprint(value: str) -> str:
+    """Return one normalized SHA-256 runtime fingerprint for argparse."""
+    normalized = value.strip().casefold()
+    if not _SHA256_FINGERPRINT_PATTERN.fullmatch(normalized):
+        raise argparse.ArgumentTypeError(
+            "runtime fingerprint must be a 64-character lowercase SHA-256 value"
+        )
+    return normalized
+
+
 def _project_metadata(repo_root: Path = REPO_ROOT) -> Mapping[str, Any]:
     """Load and return the PEP 621 project metadata from pyproject.toml."""
     pyproject_path = repo_root / "pyproject.toml"
@@ -235,6 +256,7 @@ def build_image(
     *,
     push: bool,
     comfyui_root: Path | None = None,
+    expected_fingerprint: str | None = None,
 ) -> str:
     """Build the current runtime, optionally push it, and return its image reference."""
     image_tag = _validate_tag(tag)
@@ -262,6 +284,16 @@ def build_image(
         custom_nodes_dir=settings.custom_nodes_dir,
         settings=settings,
     )
+    print(f"{SOURCE_FINGERPRINT_RESULT_PREFIX}{identity.fingerprint}", flush=True)
+    if (
+        expected_fingerprint is not None
+        and identity.fingerprint != expected_fingerprint
+    ):
+        raise RuntimeError(
+            "Local runtime source changed after ComfyUI started: expected "
+            f"{expected_fingerprint[:12]}, found {identity.fingerprint[:12]}. "
+            "Restart ComfyUI before publishing a replacement worker image."
+        )
     context = export_worker_image_context(
         repo_root=repo_root,
         settings=settings,
@@ -315,6 +347,7 @@ def main(argv: Sequence[str] | None = None) -> int:
         image_tag,
         push=arguments.push,
         comfyui_root=arguments.comfyui_root,
+        expected_fingerprint=arguments.expected_fingerprint,
     )
     print(f"COMFY_MODAL_VAST_IMAGE={image_reference}")
     return 0

tests/test_build_vast_worker_image.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import importlib.util
 from pathlib import Path
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
@@ -104,6 +105,43 @@ def test_parser_allows_push_without_explicit_tag() -> None:
     assert arguments.tag is None
 
 
+def test_build_stops_before_docker_when_expected_source_changed(
+    monkeypatch: Any,
+    tmp_path: Path,
+    capsys: Any,
+) -> None:
+    """Automatic publication must compare startup identity before exporting context."""
+    module = _module()
+    expected = "a" * 64
+    actual = "b" * 64
+    settings = SimpleNamespace(
+        comfyui_root=tmp_path,
+        custom_nodes_dir=tmp_path,
+    )
+    monkeypatch.setattr(module, "get_settings", lambda: settings)
+    monkeypatch.setattr(
+        module,
+        "build_remote_runtime_identity",
+        lambda **_kwargs: SimpleNamespace(fingerprint=actual),
+    )
+    monkeypatch.setattr(
+        module,
+        "export_worker_image_context",
+        lambda **_kwargs: pytest.fail("build context was exported after source drift"),
+    )
+
+    with pytest.raises(RuntimeError, match="Restart ComfyUI"):
+        module.build_image(
+            "ghcr.io/example/worker:v1",
+            push=True,
+            expected_fingerprint=expected,
+        )
+
+    assert capsys.readouterr().out.strip() == (
+        f"{module.SOURCE_FINGERPRINT_RESULT_PREFIX}{actual}"
+    )
+
+
 def test_docker_build_targets_vast_x86_64_platform() -> None:
     """Apple Silicon builders must produce the architecture used by Vast workers."""
     module = _module()

tests/test_vast_image_build.py

@@ -30,6 +30,7 @@ def test_builder_streams_progress_and_returns_digest(
 ) -> None:
     """Build output should reach workflow status before the digest is returned."""
     digest = "ghcr.io/example/worker@sha256:" + "a" * 64
+    fingerprint = "f" * 64
     observed: dict[str, Any] = {}
 
     def fake_popen(command: Any, **kwargs: Any) -> FakeProcess:
@@ -39,6 +40,7 @@ def test_builder_streams_progress_and_returns_digest(
         return FakeProcess(
             "\x1b[32mBuilding layer 1/3\x1b[0m\n"
             f"Pushing sha256:{'d' * 64}\n"
+            f"COMFY_MODAL_VAST_SOURCE_FINGERPRINT={fingerprint}\n"
             f"COMFY_MODAL_VAST_IMAGE={digest}\n"
         )
 
@@ -52,13 +54,15 @@ def test_builder_streams_progress_and_returns_digest(
     )
     statuses: list[str] = []
 
-    result = builder.build_and_push("f" * 64, status_callback=statuses.append)
+    result = builder.build_and_push(fingerprint, status_callback=statuses.append)
 
     assert result == digest
     assert observed["command"] == (
         sys.executable,
         "scripts/build_vast_worker_image.py",
         "--push",
+        "--expected-fingerprint",
+        fingerprint,
     )
     assert observed["cwd"] == tmp_path
     assert observed["env"]["COMFYUI_ROOT"] == str(comfyui_root)
@@ -103,3 +107,37 @@ def test_builder_failure_requires_the_documented_manual_command(
 
     assert "ModuleNotFoundError: No module named 'encodings'" in str(raised.value)
     assert "<no Python frame> Run" not in str(raised.value)
+
+
+def test_builder_rejects_source_changed_after_comfyui_started(
+    vast_image_build_module: Any,
+    monkeypatch: Any,
+    tmp_path: Path,
+) -> None:
+    """A newly built source identity must not be adopted under an old fingerprint."""
+    expected = "a" * 64
+    actual = "b" * 64
+    monkeypatch.setattr(
+        vast_image_build_module.subprocess,
+        "Popen",
+        lambda *_args, **_kwargs: FakeProcess(
+            f"COMFY_MODAL_VAST_SOURCE_FINGERPRINT={actual}\n"
+            "RuntimeError: Local runtime source changed after ComfyUI started\n",
+            1,
+        ),
+    )
+    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="no manual image build is needed",
+    ) as raised:
+        builder.build_and_push(expected)
+
+    assert expected[:12] in str(raised.value)
+    assert actual[:12] in str(raised.value)

vast_image_build.py

@@ -24,7 +24,9 @@ VAST_IMAGE_BUILD_COMMAND = (
     "--push",
 )
 _IMAGE_RESULT_PREFIX = "COMFY_MODAL_VAST_IMAGE="
+_SOURCE_FINGERPRINT_RESULT_PREFIX = "COMFY_MODAL_VAST_SOURCE_FINGERPRINT="
 _DIGEST_REFERENCE_PATTERN = re.compile(r"^\S+@sha256:[0-9a-f]{64}$")
+_FINGERPRINT_PATTERN = re.compile(r"^[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
@@ -61,17 +63,21 @@ class VastWorkerImageBuilder:
                     "Using the worker image published by another workflow",
                 )
                 return cached
-            image = self._run_build(status_callback=status_callback)
+            image = self._run_build(
+                expected_fingerprint=expected_fingerprint,
+                status_callback=status_callback,
+            )
             _BUILT_IMAGES_BY_FINGERPRINT[expected_fingerprint] = image
             return image
 
     def _run_build(
         self,
         *,
+        expected_fingerprint: str,
         status_callback: VastImageBuildStatusCallback | None,
     ) -> str:
         """Execute the image builder and extract its immutable result reference."""
-        command = self._automatic_build_command()
+        command = self._automatic_build_command(expected_fingerprint)
         command_text = shlex.join(command)
         self._emit(status_callback, "Building the current Vast worker image")
         environment = dict(os.environ if self.environment is None else self.environment)
@@ -96,11 +102,22 @@ class VastWorkerImageBuilder:
                     f"Unable to start {command[0]!r}: {exc}"
                 )
             ) from exc
-        image_reference, recent_output = self._consume_output(
+        image_reference, source_fingerprint, recent_output = self._consume_output(
             process.stdout,
             status_callback=status_callback,
         )
         return_code = process.wait()
+        if (
+            source_fingerprint is not None
+            and source_fingerprint != expected_fingerprint
+        ):
+            raise VastWorkerImageBuildError(
+                "Automatic Vast worker image publication stopped because local "
+                "runtime source changed after ComfyUI started: expected "
+                f"{expected_fingerprint[:12]}, found "
+                f"{source_fingerprint[:12]}. Restart ComfyUI and retry the "
+                "workflow; no manual image build is needed."
+            )
         if return_code != 0:
             diagnostic = self._failure_diagnostic(recent_output)
             raise VastWorkerImageBuildError(
@@ -114,16 +131,24 @@ class VastWorkerImageBuilder:
                     "The build completed without returning an immutable image digest."
                 )
             )
+        if source_fingerprint is None:
+            raise VastWorkerImageBuildError(
+                self._manual_build_message(
+                    "The build did not report its baked runtime fingerprint."
+                )
+            )
         self._emit(status_callback, "Published the current Vast worker image")
         return image_reference
 
     @staticmethod
-    def _automatic_build_command() -> tuple[str, ...]:
+    def _automatic_build_command(expected_fingerprint: str) -> tuple[str, ...]:
         """Run with the exact interpreter hosting the active ComfyUI process."""
         return (
             sys.executable,
             "scripts/build_vast_worker_image.py",
             "--push",
+            "--expected-fingerprint",
+            expected_fingerprint,
         )
 
     @staticmethod
@@ -140,11 +165,12 @@ class VastWorkerImageBuilder:
         output: IO[str] | None,
         *,
         status_callback: VastImageBuildStatusCallback | None,
-    ) -> tuple[str | None, tuple[str, ...]]:
-        """Forward subprocess lines and return its digest plus recent diagnostics."""
+    ) -> tuple[str | None, str | None, tuple[str, ...]]:
+        """Forward output and return its digest, source identity, and diagnostics."""
         if output is None:
-            return None, ()
+            return None, None, ()
         image_reference: str | None = None
+        source_fingerprint: str | None = None
         recent_output: list[str] = []
         for raw_line in output:
             line = raw_line.strip()
@@ -155,12 +181,19 @@ class VastWorkerImageBuilder:
                 if _DIGEST_REFERENCE_PATTERN.fullmatch(candidate):
                     image_reference = candidate
                 continue
+            if line.startswith(_SOURCE_FINGERPRINT_RESULT_PREFIX):
+                candidate = line.removeprefix(
+                    _SOURCE_FINGERPRINT_RESULT_PREFIX
+                ).strip()
+                if _FINGERPRINT_PATTERN.fullmatch(candidate):
+                    source_fingerprint = 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)
+        return image_reference, source_fingerprint, tuple(recent_output)
 
     @staticmethod
     def _safe_status_line(line: str) -> str:
@@ -174,7 +207,7 @@ class VastWorkerImageBuilder:
 
     def _manual_build_message(self, cause: str) -> str:
         """Return an actionable error without leaking subprocess environment data."""
-        recovery_command = shlex.join(self._automatic_build_command())
+        recovery_command = shlex.join(VAST_IMAGE_BUILD_COMMAND)
         return (
             f"Automatic Vast worker image publication failed. {cause} Run `"
             + recovery_command