Fix SSH worker startup with current ComfyUI

2026-10-29 11:58:16 -07:00
parent 39364494e7
commit 5b09bca6c1
10 changed files with 76 additions and 82 deletions
@@ -142,7 +142,7 @@ Each **Modal Configuration** selects a GPU type and maximum concurrent container

 Self-hosted execution runs the same workers on machines you control, reached over SSH and managed through Docker. Add one **SSH Configuration** per machine and connect it to the configurator. The extension deliberately stores no passwords or private-key paths — authentication, jump hosts, ports, and key selection stay in your normal SSH agent and `~/.ssh/config`.

-The installation-wide **Settings → Remote Execution: SSH environments** manager remains available for legacy workflows and host preflight. New configured workflows carry their own credential-free SSH destination and scheduling snapshot, including the probed GPU capabilities, so their execution does not depend on mutating that global registry. Older queued payloads without a capability snapshot are re-probed before launch. GPU workers are always started with the selected GPU device request; missing GPU/runtime capability is an explicit setup error instead of a silent CPU-only container. The worker image also pins and validates the PyAV API required by the copied ComfyUI source, so a stale video dependency fails during image construction instead of at the first remote node import.
+The installation-wide **Settings → Remote Execution: SSH environments** manager remains available for legacy workflows and host preflight. New configured workflows carry their own credential-free SSH destination and scheduling snapshot, including the probed GPU capabilities, so their execution does not depend on mutating that global registry. Older queued payloads without a capability snapshot are re-probed before launch. GPU workers are always started with the selected GPU device request; missing GPU/runtime capability is an explicit setup error instead of a silent CPU-only container. The worker image pins `comfy-aimdo==0.6.0` and `comfy-kitchen==0.2.40` and validates their storage and INT8 encoder APIs alongside PyAV during image construction.

 Worker images embed every file beneath ComfyUI's curated runtime package directories, including tokenizer vocabularies, model configuration JSON, SentencePiece data, and other non-Python resources. Those files participate in the runtime fingerprint, so adding or changing required package data automatically replaces a stale SSH worker and is reflected in the next published Vast worker image. The heavyweight apt, pinned PyTorch, CUDA, accelerator, and custom-requirement stack is retained as a separately fingerprinted dependency image: locally on self-hosted SSH machines and under a `deps-<fingerprint>` registry tag for published Vast images. A normal source edit therefore rebuilds only the roughly 52 MiB source overlay; the multi-gigabyte base is rebuilt and pushed only when its dependency inputs change. The pinned PyTorch stack is installed before general runtime and custom-node requirements so their dependency resolution cannot silently select a different Torch build. Before a self-hosted build, superseded local images tagged and labeled by this extension are removed when Docker no longer has a container reference to them, bounding the retained cache without pruning unrelated host images. The image also retains an explicit `gcc`/`g++` build toolchain and exports `CC`/`CXX`, because Triton may compile native launchers lazily when a ComfyUI node first executes; the dependency build validates that toolchain before publication.

@@ -152,7 +152,7 @@ Each host must provide:

 - Linux on `x86_64`
 - a non-interactive SSH login with an already trusted host key
- a working Docker CLI and daemon with the Buildx `default` builder available to that login
+- a working Docker CLI and daemon with the active context's built-in Buildx builder available to that login (including a `rootless` context)
 - NVIDIA drivers plus Docker's NVIDIA runtime or CDI configuration for GPU work
 - outbound package/image access during the first worker build

@@ -506,6 +506,8 @@ uv sync --group test
 uv run pytest
 ```

+The test group installs real PyAV and the same ComfyUI support-library pins as workers. Unit tests import the cloud entrypoint without constructing Modal resources, even when the remote extra is installed; live provider tests require explicit opt-in.
+
 Tests look for ComfyUI in `COMFYUI_ROOT` first, then `COMFY_MODAL_COMFYUI_ROOT`, then an installed parent checkout, then `~/git/ComfyUI`. To run against a temporary checkout:

 ```bash
@@ -736,3 +736,11 @@ Upload, Hugging Face discovery, Hugging Face download, and R2 download progress
 ## Compatibility imports can hide an unfinished extraction

 Moving behavior and mutable state out of a stable entrypoint does not make the entrypoint thin if it still eagerly imports every extracted private helper. Import only the hooks and operations used by production glue, then resolve legacy read-only private access through the focused owner modules. This preserves a migration surface without creating hundreds of stale bindings or obscuring which module actually owns execution.
+
+## SSH builds must use the active Docker context
+
+Docker's built-in Buildx builders are named after their contexts. Forcing `--builder default` fails when the active context is `rootless`; switching only the build to another context would also put the image in a different daemon from subsequent worker commands. Use `docker context show` to select the builder for both dependency and source-overlay images.
+
+## Keep worker and test dependencies aligned with ComfyUI
+
+ComfyUI 0.43 requires `comfy-aimdo==0.6.0` for `comfy_aimdo.storage` and `comfy-kitchen==0.2.40` for `comfy_kitchen.int8_decode`. Validate these APIs at image-build time. Tests need real PyAV rather than a partial module stub that breaks whenever ComfyUI imports another video API. Import the cloud entrypoint with Modal unavailable in unit tests so installing the optional SDK cannot trigger authenticated resource lookup. Registry-initialization tests that stub ComfyUI dependencies must also isolate the server import.
@@ -33,8 +33,9 @@ Repository = "https://github.com/ttulttul/ComfyUI-Modal"
 [dependency-groups]
 test = [
    "aiohttp",
-    "comfy-aimdo==0.4.13",
-    "comfy-kitchen==0.2.31",
+    "av==18.1.0",
+    "comfy-aimdo==0.6.0",
+    "comfy-kitchen==0.2.40",
    "einops",
    "numpy",
    "pillow",
@@ -53,8 +53,8 @@ _REMOTE_RUNTIME_PACKAGES = (
    REMOTE_AV_SPEC,
    "blake3==1.0.9",
    "comfy-angle==0.1.0",
-    "comfy-aimdo==0.4.13",
-    "comfy-kitchen==0.2.31",
+    "comfy-aimdo==0.6.0",
+    "comfy-kitchen==0.2.40",
    "einops==0.8.2",
    "kornia==0.8.2",
    REMOTE_NUMPY_SPEC,
@@ -253,9 +253,11 @@ def remote_runtime_packages() -> tuple[str, ...]:


 def remote_runtime_validation_command() -> str:
-    """Return a build-time check for ComfyUI's required PyAV API."""
+    """Return a build-time check for ComfyUI's required video and storage APIs."""
    validation_script = (
        "from importlib import metadata; "
+        "import comfy_aimdo.storage; "
+        "from comfy_kitchen.int8_decode import Int8DecodeCache; "
        "from av.video.reformatter import ColorPrimaries; "
        f"expected_av={REMOTE_AV_SPEC.split('==', maxsplit=1)[1]!r}; "
        "actual_av=metadata.version('av'); "
@@ -277,6 +277,10 @@ class SshRuntimeManager:
        status_callback: Callable[[str], None] | None = None,
    ) -> None:
        """Build a stable dependency image and apply the current source overlay."""
+        builder_context = self.controller.docker(("context", "show")).stdout_text.strip()
+        if not builder_context:
+            raise SshDockerError("Remote Docker did not report its active context.")
+        logger.info("Building SSH worker with Docker context %s.", builder_context)
        self._remove_stale_runtime_images(
            spec,
            status_callback=status_callback,
@@ -300,6 +304,7 @@ class SshRuntimeManager:
            )
            self.controller.docker(
                self._buildx_arguments(
+                    builder_context=builder_context,
                    image_tag=dependency_image_tag,
                    labels={_DEPENDENCY_LABEL: dependency_fingerprint},
                    pull=True,
@@ -328,6 +333,7 @@ class SshRuntimeManager:
        )
        self.controller.docker(
            self._buildx_arguments(
+                builder_context=builder_context,
                image_tag=spec.image_tag,
                labels={
                    _RUNTIME_LABEL: spec.identity.fingerprint,
@@ -453,12 +459,13 @@ class SshRuntimeManager:
    def _buildx_arguments(
        self,
        *,
+        builder_context: str,
        image_tag: str,
        labels: dict[str, str],
        pull: bool,
    ) -> tuple[str, ...]:
-        """Return a build command pinned to Docker's persistent local builder."""
-        arguments = ["buildx", "build", "--builder", "default"]
+        """Use the active context's local builder for the worker's Docker daemon."""
+        arguments = ["buildx", "build", "--builder", builder_context]
        if pull:
            arguments.append("--pull")
        arguments.append("--load")
@@ -3,11 +3,9 @@
 from __future__ import annotations

 import importlib
-import importlib.machinery
 import importlib.util
 import os
 import sys
-import types
 from pathlib import Path

 import pytest
@@ -30,58 +28,6 @@ def _comfyui_root() -> Path:

 def _ensure_import_paths() -> None:
    """Add the repository and local ComfyUI checkout to sys.path when present."""
-    if "av" not in sys.modules:
-        av_module = types.ModuleType("av")
-        av_module.__spec__ = importlib.machinery.ModuleSpec("av", loader=None)
-        av_module.__path__ = []
-        av_module.open = lambda *args, **kwargs: (_ for _ in ()).throw(
-            RuntimeError("The av stub should not be used in tests.")
-        )
-        av_module.time_base = 1
-        av_module.VideoStream = type("VideoStream", (), {})
-        av_module.AVError = RuntimeError
-        av_module.FFmpegError = RuntimeError
-        av_module.logging = types.SimpleNamespace(
-            ERROR="ERROR",
-            set_level=lambda *args, **kwargs: None,
-        )
-        av_module.video = types.SimpleNamespace(
-            frame=types.SimpleNamespace(
-                VideoFrame=type("VideoFrame", (), {"pict_type": None})
-            )
-        )
-
-        av_container_module = types.ModuleType("av.container")
-        av_container_module.InputContainer = type("InputContainer", (), {})
-
-        av_subtitles_module = types.ModuleType("av.subtitles")
-        av_subtitles_stream_module = types.ModuleType("av.subtitles.stream")
-        av_subtitles_stream_module.SubtitleStream = type("SubtitleStream", (), {})
-        av_video_module = types.ModuleType("av.video")
-        av_video_reformatter_module = types.ModuleType("av.video.reformatter")
-        av_video_reformatter_module.ColorPrimaries = type(
-            "ColorPrimaries", (), {"BT2020": 9, "BT709": 1}
-        )
-        av_video_reformatter_module.ColorRange = type(
-            "ColorRange", (), {"JPEG": 2, "MPEG": 1}
-        )
-        av_video_reformatter_module.ColorTrc = type(
-            "ColorTrc",
-            (),
-            {
-                "ARIB_STD_B67": 18,
-                "BT709": 1,
-                "IEC61966_2_1": 13,
-                "SMPTE2084": 16,
-            },
-        )
-
-        sys.modules["av"] = av_module
-        sys.modules["av.container"] = av_container_module
-        sys.modules["av.subtitles"] = av_subtitles_module
-        sys.modules["av.subtitles.stream"] = av_subtitles_stream_module
-        sys.modules["av.video"] = av_video_module
-        sys.modules["av.video.reformatter"] = av_video_reformatter_module
    if str(REPO_ROOT) not in sys.path:
        sys.path.insert(0, str(REPO_ROOT))
    comfyui_root = _comfyui_root()
@@ -665,7 +611,9 @@ def sync_r2_transfer_module(extension_package: object) -> object:
 def modal_cloud_module() -> object:
    """Return the stable Modal cloud entry module."""
    _ensure_import_paths()
-    return importlib.import_module("comfyui_modal_sync_cloud")
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setitem(sys.modules, "modal", None)
+        return importlib.import_module("comfyui_modal_sync_cloud")


@pytest.fixture(scope="session")
@@ -1378,6 +1378,7 @@ def test_modal_cloud_initializes_remote_comfy_runtime_once_per_custom_node_root(
        )
    )

+    monkeypatch.setitem(sys.modules, "server", types.SimpleNamespace())
    monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
    monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths_module)

@@ -3,9 +3,12 @@
 from __future__ import annotations

 import importlib
+import sys
 from pathlib import Path
 from types import ModuleType

+import pytest
+
 from conftest import PACKAGE_NAME, REPO_ROOT

 _CLOUD_EXCEPTION_NAMES = (
@@ -27,9 +30,11 @@ def _root_module_names(repo_root: Path) -> tuple[str, ...]:

 def test_root_modules_support_package_and_flat_imports(
    extension_package: object,
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
    """Every shipped root module must load locally and in the flat cloud layout."""
    del extension_package
+    monkeypatch.setitem(sys.modules, "modal", None)
    module_names = _root_module_names(REPO_ROOT)

    for module_name in module_names:
@@ -39,11 +44,10 @@ def test_root_modules_support_package_and_flat_imports(


 def test_stable_cloud_entrypoint_exports_host_runtime_surface(
-    extension_package: object,
+    modal_cloud_module: ModuleType,
 ) -> None:
    """The stable flat cloud module must retain every name read by the host."""
-    del extension_package
-    cloud_module = importlib.import_module("comfyui_modal_sync_cloud")
+    cloud_module = modal_cloud_module

    assert isinstance(cloud_module, ModuleType)
    assert hasattr(cloud_module, "app")
@@ -54,11 +58,10 @@ def test_stable_cloud_entrypoint_exports_host_runtime_surface(


 def test_cloud_entrypoint_eagerly_imports_volume_hydration_helper(
-    extension_package: object,
+    modal_cloud_module: ModuleType,
 ) -> None:
    """Method globals must not depend on the module compatibility fallback."""
-    del extension_package
-    cloud_module = importlib.import_module("comfyui_modal_sync_cloud")
+    cloud_module = modal_cloud_module

    assert "_hydrate_missing_payload_volume_paths" in vars(cloud_module)
    assert callable(vars(cloud_module)["_hydrate_missing_payload_volume_paths"])
@@ -3,6 +3,7 @@
 from __future__ import annotations

 import shlex
+import tomllib
 from pathlib import Path
 from types import SimpleNamespace
 from typing import Any
@@ -87,8 +88,8 @@ def test_remote_environment_is_fully_pinned(runtime_environment_module: Any) ->
    assert len(runtime_packages) == len(set(runtime_packages))
    assert "blake3==1.0.9" in runtime_packages
    assert "comfy-angle==0.1.0" in runtime_packages
-    assert "comfy-aimdo==0.4.13" in runtime_packages
-    assert "comfy-kitchen==0.2.31" in runtime_packages
+    assert "comfy-aimdo==0.6.0" in runtime_packages
+    assert "comfy-kitchen==0.2.40" in runtime_packages
    assert "pyopengl==3.1.10" in runtime_packages
    assert "pypdf==6.16.1" in runtime_packages
    assert "numpy==2.3.5" in runtime_packages
@@ -100,6 +101,15 @@ def test_remote_environment_is_fully_pinned(runtime_environment_module: Any) ->
    assert "simpleeval==1.0.7" in runtime_packages


+def test_test_group_matches_worker_support_libraries(runtime_environment_module: Any) -> None:
+    """Local tests must exercise the same ComfyUI support APIs as worker images."""
+    project = tomllib.loads((Path(__file__).resolve().parents[1] / "pyproject.toml").read_text())
+    test_packages = project["dependency-groups"]["test"]
+    for package in runtime_environment_module.remote_runtime_packages():
+        if package.startswith(("av==", "comfy-aimdo==", "comfy-kitchen==")):
+            assert package in test_packages
+
+
 def test_remote_video_stack_is_pinned_and_validated(
    runtime_environment_module: Any,
 ) -> None:
@@ -109,6 +119,8 @@ def test_remote_video_stack_is_pinned_and_validated(
    assert runtime_environment_module.REMOTE_AV_SPEC == "av==18.1.0"
    assert command.startswith("python -c ")
    assert "ColorPrimaries" in command
+    assert "import comfy_aimdo.storage" in command
+    assert "from comfy_kitchen.int8_decode import Int8DecodeCache" in command
    assert "Expected PyAV" in command
    assert "18.1.0" in command

@@ -7,6 +7,8 @@ from pathlib import Path
 from types import SimpleNamespace
 from typing import Any

+import pytest
+

 def test_worker_lifecycle_reports_build_and_readiness_status(
    ssh_runtime_module: Any,
@@ -95,9 +97,11 @@ def test_image_only_preparation_builds_before_storage_helpers_run(
    ]


+@pytest.mark.parametrize("docker_context", ["default", "rootless"])
 def test_worker_build_loads_image_into_the_remote_daemon(
    ssh_runtime_module: Any,
    monkeypatch: Any,
+    docker_context: str,
 ) -> None:
    """BuildKit docker-container drivers must export the image for docker run."""
    calls: list[tuple[tuple[str, ...], dict[str, Any]]] = []
@@ -105,6 +109,8 @@ def test_worker_build_loads_image_into_the_remote_daemon(
    def docker(arguments: tuple[str, ...], **kwargs: Any) -> Any:
        """Record one remote Docker invocation."""
        calls.append((arguments, kwargs))
+        if arguments == ("context", "show"):
+            return SimpleNamespace(returncode=0, stdout_text=f"{docker_context}\n")
        if arguments[:2] == ("image", "inspect"):
            return SimpleNamespace(returncode=1, stdout_text="")
        return None
@@ -144,23 +150,24 @@ def test_worker_build_loads_image_into_the_remote_daemon(

    manager._build_image(spec)

-    assert calls[0][0][:2] == ("image", "inspect")
-    dependency_arguments, dependency_kwargs = calls[1]
+    assert calls[0][0] == ("context", "show")
+    assert calls[1][0][:2] == ("image", "inspect")
+    dependency_arguments, dependency_kwargs = calls[2]
    assert dependency_arguments[:6] == (
        "buildx",
        "build",
        "--builder",
-        "default",
+        docker_context,
        "--pull",
        "--load",
    )
    assert dependency_kwargs["input_payload"] == b"dependencies"
-    overlay_arguments, overlay_kwargs = calls[2]
+    overlay_arguments, overlay_kwargs = calls[3]
    assert overlay_arguments[:5] == (
        "buildx",
        "build",
        "--builder",
-        "default",
+        docker_context,
        "--load",
    )
    assert "--pull" not in overlay_arguments
@@ -177,6 +184,8 @@ def test_worker_build_reuses_retained_dependency_image(
    def docker(arguments: tuple[str, ...], **kwargs: Any) -> Any:
        """Report a current dependency image and record the overlay build."""
        calls.append((arguments, kwargs))
+        if arguments == ("context", "show"):
+            return SimpleNamespace(returncode=0, stdout_text="rootless\n")
        if arguments[:2] == ("image", "inspect"):
            return SimpleNamespace(
                returncode=0,
@@ -220,16 +229,17 @@ def test_worker_build_reuses_retained_dependency_image(

    manager._build_image(spec)

-    assert len(calls) == 2
-    assert calls[0][0][:2] == ("image", "inspect")
-    assert calls[1][0][:5] == (
+    assert len(calls) == 3
+    assert calls[0][0] == ("context", "show")
+    assert calls[1][0][:2] == ("image", "inspect")
+    assert calls[2][0][:5] == (
        "buildx",
        "build",
        "--builder",
-        "default",
+        "rootless",
        "--load",
    )
-    assert calls[1][1]["input_payload"] == b"52MB-source-overlay"
+    assert calls[2][1]["input_payload"] == b"52MB-source-overlay"


 def test_worker_build_reclaims_only_superseded_managed_images(