diff --git a/README.md b/README.md index 9be93e6..df6ed1f 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,7 @@ Important distinction: - the Modal image filter now preserves internal ComfyUI Python packages such as `comfy/ldm/models` while still excluding top-level runtime asset folders like `models/` and `output/` - the remote worker now initializes ComfyUI's built-in extra nodes and API nodes on first use, and loads extracted custom-node bundles through ComfyUI's normal `custom_nodes` registry path - mirrored asset references such as `/assets/_model.safetensors` are now materialized to container-local absolute paths before remote execution, and the worker patches ComfyUI's `folder_paths` lookups so normal loader nodes still accept them +- in real remote mode, queue-time asset sync now uploads into the named Modal volume instead of only mirroring to local disk, and the worker reloads that mounted volume before each execution so warm containers can see newly uploaded assets - the remote worker now force-loads ComfyUI's top-level `utils` package from the bundled ComfyUI source tree before API nodes initialize, which avoids third-party or stray `utils` modules shadowing `utils.install_util` If you want a stable reusable Modal deployment, that is still a separate step outside the current node-pack behavior. diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index e3792b3..e5666f2 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -547,6 +547,7 @@ if modal is not None: # pragma: no branch - remote entrypoint configuration. @modal.method() def execute_payload(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes: """Execute a proxied node or subgraph inside the Modal container.""" + vol.reload() if payload.get("payload_kind") == "subgraph": return execute_subgraph_locally(payload, kwargs_payload) return execute_node_locally(payload, kwargs_payload) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 18f608f..755f120 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -20,4 +20,5 @@ - If `/modal/queue_prompt` does expensive preparatory work before it returns, the frontend can keep ComfyUI's own queue/execution indicators alive by dispatching temporary synthetic `status`, `execution_start`, and `executing` events, then letting the real websocket events take over once the prompt is actually queued. - The synthetic frontend events have to match ComfyUI's websocket adapter shapes exactly, not just semantically. In practice that means `status` detail must be the inner status payload, `executing` detail should be the display node id string, and queue-list initialization state benefits from a matching `"Waiting for a machine"` notification payload keyed by prompt id. - Synced asset paths cannot stay as abstract `/assets/...` markers inside the Modal worker. Remote execution needs those markers materialized to real absolute files under the mounted storage root, and ComfyUI's `folder_paths` helpers need an absolute-path fast path so loader nodes can still resolve them. +- A local filesystem mirror is not enough in `COMFY_MODAL_EXECUTION_MODE=remote`. Queue-time sync must upload assets into the actual named Modal volume via the local SDK, and the remote worker must call `vol.reload()` before execution so warm containers can see newly uploaded files. - The Modal container can pick up a non-package `utils` module before ComfyUI initializes. Preloading ComfyUI's own top-level `utils` package from the bundled source tree prevents `utils.install_util` import failures during `init_extra_nodes(...)`. diff --git a/sync_engine.py b/sync_engine.py index fb015b6..55805f2 100644 --- a/sync_engine.py +++ b/sync_engine.py @@ -3,6 +3,7 @@ from __future__ import annotations import hashlib +import io import json import logging import os @@ -39,6 +40,11 @@ _SKIP_FILE_SUFFIXES = { ".tmp", } +try: + import modal # type: ignore +except ModuleNotFoundError: # pragma: no cover - exercised when Modal SDK is unavailable. + modal = None + class VolumeBackend(Protocol): """Minimal storage interface needed by the sync engine.""" @@ -91,6 +97,30 @@ class LocalMirrorVolume: return self.root / remote_path.lstrip("/") +class ModalVolumeBackend: + """Modal Volume-backed storage for real remote execution.""" + + def __init__(self, volume_name: str) -> None: + """Resolve a named Modal volume lazily from the local SDK client.""" + if modal is None: + raise RuntimeError("Modal SDK is required for ModalVolumeBackend.") + self._volume = modal.Volume.from_name(volume_name, create_if_missing=True) + + def exists(self, remote_path: str) -> bool: + """Return whether a file already exists in the Modal volume.""" + return len(self._volume.listdir(remote_path, recursive=False)) > 0 + + def put_file(self, local_path: Path, remote_path: str) -> None: + """Upload a local file into the Modal volume.""" + with self._volume.batch_upload() as batch: + batch.put_file(local_path, remote_path) + + def put_bytes(self, payload: bytes, remote_path: str) -> None: + """Upload bytes into the Modal volume.""" + with self._volume.batch_upload() as batch: + batch.put_file(io.BytesIO(payload), remote_path) + + @dataclass class ModalAssetSyncEngine: """Content-addressable storage sync engine for files and custom nodes.""" @@ -102,7 +132,19 @@ class ModalAssetSyncEngine: def from_environment(cls, settings: ModalSyncSettings | None = None) -> "ModalAssetSyncEngine": """Create a sync engine using the local mirror backend by default.""" resolved_settings = settings or get_settings() - backend = LocalMirrorVolume(resolved_settings.local_storage_root) + backend: VolumeBackend + if resolved_settings.execution_mode == "remote" and modal is not None: + logger.info( + "Using Modal volume backend %s for remote asset sync.", + resolved_settings.volume_name, + ) + backend = ModalVolumeBackend(resolved_settings.volume_name) + else: + if resolved_settings.execution_mode == "remote" and modal is None: + logger.warning( + "Modal SDK is unavailable in remote execution mode; falling back to local mirror storage." + ) + backend = LocalMirrorVolume(resolved_settings.local_storage_root) return cls(volume=backend, settings=resolved_settings) def sync_file(self, local_path: Path, remote_folder: str = "/assets") -> SyncedAsset: diff --git a/tests/test_sync_engine.py b/tests/test_sync_engine.py index d69282a..50c71c8 100644 --- a/tests/test_sync_engine.py +++ b/tests/test_sync_engine.py @@ -9,9 +9,11 @@ from typing import Any def test_sync_file_deduplicates_by_hash( settings_module: Any, sync_engine_module: Any, + monkeypatch: Any, tmp_path: Path, ) -> None: """Repeated file syncs should reuse the same remote path and marker.""" + monkeypatch.setattr(sync_engine_module, "modal", None) asset_path = tmp_path / "model.safetensors" asset_path.write_bytes(b"model-bytes") @@ -42,9 +44,11 @@ def test_sync_file_deduplicates_by_hash( def test_sync_custom_nodes_directory_creates_archive( settings_module: Any, sync_engine_module: Any, + monkeypatch: Any, tmp_path: Path, ) -> None: """The sync engine should archive and mirror a custom_nodes directory.""" + monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) @@ -75,9 +79,11 @@ def test_sync_custom_nodes_directory_creates_archive( def test_hash_directory_ignores_virtualenv_and_bytecode_artifacts( settings_module: Any, sync_engine_module: Any, + monkeypatch: Any, tmp_path: Path, ) -> None: """Ignored directories and file suffixes should not affect custom_nodes hashing.""" + monkeypatch.setattr(sync_engine_module, "modal", None) custom_nodes_dir = tmp_path / "custom_nodes" package_dir = custom_nodes_dir / "example" package_dir.mkdir(parents=True) @@ -112,3 +118,87 @@ def test_hash_directory_ignores_virtualenv_and_bytecode_artifacts( ignored_hash = engine._hash_directory(custom_nodes_dir) assert ignored_hash == baseline_hash + + +def test_remote_mode_uses_modal_volume_backend_when_sdk_is_available( + settings_module: Any, + sync_engine_module: Any, + monkeypatch: Any, + tmp_path: Path, +) -> None: + """Remote execution mode should upload into a real Modal volume when the SDK is available.""" + + class FakeBatch: + """Capture files uploaded through Modal batch_upload.""" + + def __init__(self, uploads: list[tuple[Any, str]]) -> None: + """Store uploaded file references.""" + self.uploads = uploads + + def __enter__(self) -> "FakeBatch": + """Return the active batch context.""" + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + """Close the fake batch context.""" + return None + + def put_file(self, source: Any, remote_path: str) -> None: + """Record the uploaded file or bytes payload.""" + self.uploads.append((source, remote_path)) + + class FakeVolume: + """Minimal Modal volume double.""" + + def __init__(self) -> None: + """Initialize the fake upload state.""" + self.paths: set[str] = set() + self.uploads: list[tuple[Any, str]] = [] + + def listdir(self, remote_path: str, recursive: bool = False) -> list[str]: + """Return a non-empty listing when the path has already been uploaded.""" + return [remote_path] if remote_path in self.paths else [] + + def batch_upload(self) -> FakeBatch: + """Return a fake batch uploader.""" + return FakeBatch(self.uploads) + + fake_volume = FakeVolume() + + class FakeModal: + """Minimal modal SDK double that returns a stable volume handle.""" + + class Volume: + """Namespace for volume lookups.""" + + @staticmethod + def from_name(name: str, create_if_missing: bool = False) -> FakeVolume: + """Return the fake volume for any lookup.""" + return fake_volume + + monkeypatch.setattr(sync_engine_module, "modal", FakeModal) + + settings = settings_module.ModalSyncSettings( + app_name="app", + execution_mode="remote", + sync_custom_nodes=True, + volume_name="volume", + route_path="/modal/queue_prompt", + marker_property="is_modal_remote", + local_storage_root=tmp_path / "storage", + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=None, + ) + + engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + assert type(engine.volume).__name__ == "ModalVolumeBackend" + + asset_path = tmp_path / "encoder.safetensors" + asset_path.write_bytes(b"weights") + synced = engine.sync_file(asset_path) + + uploaded_remote_paths = [remote_path for _, remote_path in fake_volume.uploads] + assert synced.remote_path in uploaded_remote_paths + assert f"/hashes/{synced.sha256}.done" in uploaded_remote_paths