commit ef75060ddb883f8ffca88bba48a48ce57421e335 Author: Ken Simpson Date: Fri Apr 17 10:18:17 2026 -0700 feat: add modal sync comfyui extension diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..de02703 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,32 @@ +## Dev environment tips +- Use uv for Python package management and manage this project as a uv project. +- There is a Python virtual environment in .venv; if for some reason you don't see it, run `scripts/setup_env.sh`. +- The `rg` command is installed; use it for quick searching. +- The git server we use is Github. +- Commit every change you make and ask the user to push changes when a significant batch of changes has been made. +- When you make a big discovery or change, note this in docs/LEARNINGS.md, even if user does not ask you to. +- Update README.md with each major commit. +- You can find the ComfyUI source code in ~/git/ComfyUI. + +## Dev process tips +- Use Python logging liberally to insert judicious info, warning, and debug messages in the code. +- Import logging into each module and set logger = logging.getLogger(__name__). Then use logger for logging in that module. +- In most cases where an error message is called for, you should raise an appropriate exception. We want to know. +- Let exceptions bubble up to the caller unless there is a logical way to handle the exception within the current scope +. +- Unless there is no more specific alternative, never catch the base `Exception` class. Always catch the most specific +exception you can. +- Commit every change you make and ask the user to push changes when a significant batch of changes has been made. +- Use Python type hints, data classes, Pydantic models, and other strong typing features everywhere. +- Never write a function that has no typing on parameters. +- Write a documentation string for every function and every class. +- Keep functions short and modular. If a function grows longer than 50 rows or so, break it into separate functions tha +t you can compose in a sequence. +- Take the time to keep your code clean, tidy, and readable. Comments should not be required if code is written in a mo +dular, readable manner. + +## Testing instructions +- Add or update tests for the code you change, even if nobody asked. +- Use pytest and create a harness so the user can just type "uv pytest". +- Always run the full test suite before every commit. Never commit if the test suite fails to pass. + diff --git a/README.md b/README.md new file mode 100644 index 0000000..df10741 --- /dev/null +++ b/README.md @@ -0,0 +1,59 @@ +# ComfyUI Modal-Sync + +`ComfyUI Modal-Sync` is a ComfyUI custom node extension that lets you mark individual nodes for remote execution. The frontend stores a per-node `is_modal_remote` flag in workflow metadata, the backend rewrites those nodes into a Modal proxy before queueing, and the proxy forwards execution payloads through a serialization layer that is safe for tensors and plain JSON values. + +## What is implemented + +- A frontend extension in [web/modal_toggle.js](/home/user/git/ComfyUI-Modal/web/modal_toggle.js) that injects a `Run on Modal` toggle into every non-internal node, draws a blue border around remote nodes, and submits queued prompts to `/modal/queue_prompt`. +- A backend route in [api_intercept.py](/home/user/git/ComfyUI-Modal/api_intercept.py) that reads the workflow snapshot from `extra_pnginfo.workflow`, identifies remote-marked nodes, syncs referenced assets, and rewrites each remote node into a signature-preserving Modal proxy node. +- A dynamic proxy registry in [modal_executor_node.py](/home/user/git/ComfyUI-Modal/modal_executor_node.py) that mirrors the original node’s output count and output types so ComfyUI validation still succeeds after rewrite. +- A content-addressable sync engine in [sync_engine.py](/home/user/git/ComfyUI-Modal/sync_engine.py) that mirrors model files and a zipped `custom_nodes/` bundle into storage keyed by SHA256. +- A remote runtime skeleton in [remote/modal_app.py](/home/user/git/ComfyUI-Modal/remote/modal_app.py) that supports local fallback execution for tests and can call into Modal when the SDK is available. + +## Repository layout + +```text +. +├── __init__.py +├── api_intercept.py +├── modal_executor_node.py +├── remote/ +│ ├── __init__.py +│ └── modal_app.py +├── serialization.py +├── settings.py +├── sync_engine.py +├── tests/ +└── web/ + └── modal_toggle.js +``` + +## Runtime behavior + +1. Queueing a prompt posts to `/modal/queue_prompt`. +2. The route loads the workflow metadata from `extra_data.extra_pnginfo.workflow`. +3. Nodes with `properties.is_modal_remote = true` are replaced with a generated `ModalUniversalExecutor_` node id. +4. File-like widget inputs ending in `.safetensors`, `.ckpt`, `.pt`, or `.vae` are resolved to local files and mirrored into storage. +5. The executor serializes tensors with `safetensors`, forwards the payload to the remote runtime, and deserializes the returned outputs back into ComfyUI values. + +## Configuration + +These environment variables are supported: + +- `COMFY_MODAL_ROUTE_PATH`: Override the queue endpoint. Default: `/modal/queue_prompt`. +- `COMFY_MODAL_MARKER_PROPERTY`: Override the workflow property used to mark nodes. Default: `is_modal_remote`. +- `COMFY_MODAL_LOCAL_STORAGE_ROOT`: Local mirror root used for sync tests and dry runs. Default: `/tmp/comfyui-modal-sync-storage`. +- `COMFY_MODAL_CUSTOM_NODES_DIR`: Override the `custom_nodes` directory to bundle and mirror. +- `COMFY_MODAL_EXECUTION_MODE`: Set to `local` for in-process fallback execution. Default: `local`. +- `COMFY_MODAL_APP_NAME` and `COMFY_MODAL_VOLUME_NAME`: Override Modal app and volume naming. + +## Development + +- Run tests with `uv run pytest`. +- The project includes a minimal [pyproject.toml](/home/user/git/ComfyUI-Modal/pyproject.toml) so pytest discovery is consistent under `uv`. + +## Current limitations + +- The graph rewrite is per-node, not whole-subgraph partitioning. Remote-to-remote chains will currently round-trip through the local executor boundary between nodes. +- The default sync backend is a local mirror used for development and tests. The Modal runtime entrypoint is present, but real cloud execution still depends on a working Modal SDK environment. +- Non-JSON, non-bytes, non-tensor values cannot be serialized across the remote boundary yet. Unsupported types raise immediately instead of silently degrading. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..026c9c4 --- /dev/null +++ b/__init__.py @@ -0,0 +1,50 @@ +"""ComfyUI Modal-Sync extension entrypoint.""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + +WEB_DIRECTORY = "./web" +try: + from comfy_api.latest import ComfyExtension, io + + from .api_intercept import setup_modal_queue_route + from .modal_executor_node import ModalUniversalExecutor +except ModuleNotFoundError: # pragma: no cover - used during local non-Comfy imports. + ComfyExtension = object # type: ignore[assignment] + io = None # type: ignore[assignment] + ModalUniversalExecutor = None # type: ignore[assignment] + + class ComfyModalSyncExtension: # type: ignore[no-redef] + """Fallback placeholder used when ComfyUI is not importable.""" + + async def comfy_entrypoint() -> object: # type: ignore[no-redef] + """Raise a helpful error when imported outside ComfyUI.""" + raise RuntimeError( + "ComfyUI Modal-Sync requires the ComfyUI runtime on sys.path to load." + ) + +else: + + class ComfyModalSyncExtension(ComfyExtension): + """Expose Modal-Sync nodes to the ComfyUI extension loader.""" + + async def get_node_list(self) -> list[type[io.ComfyNode]]: + """Return the node list after registering API routes.""" + setup_modal_queue_route() + return [ModalUniversalExecutor] + + async def comfy_entrypoint() -> ComfyExtension: + """Create the ComfyUI v3 extension entrypoint.""" + setup_modal_queue_route() + return ComfyModalSyncExtension() + + +__all__ = [ + "ComfyModalSyncExtension", + "ModalUniversalExecutor", + "WEB_DIRECTORY", + "comfy_entrypoint", +] diff --git a/api_intercept.py b/api_intercept.py new file mode 100644 index 0000000..4d31149 --- /dev/null +++ b/api_intercept.py @@ -0,0 +1,246 @@ +"""Prompt interception and graph rewriting for Modal-backed execution.""" + +from __future__ import annotations + +import copy +import logging +import time +import uuid +from dataclasses import dataclass, field +from typing import Any + +from aiohttp import web + +from .modal_executor_node import ensure_modal_proxy_node_registered +from .settings import ModalSyncSettings, get_settings +from .sync_engine import ModalAssetSyncEngine, SyncedAsset + +logger = logging.getLogger(__name__) + +_ROUTE_REGISTERED = False + + +@dataclass +class RewriteSummary: + """Summary of the prompt rewrite performed for a queue request.""" + + remote_node_ids: list[str] = field(default_factory=list) + synced_assets: list[SyncedAsset] = field(default_factory=list) + custom_nodes_bundle: SyncedAsset | None = None + + +def _get_nodes_module() -> Any: + """Import the ComfyUI nodes module lazily.""" + import nodes + + return nodes + + +def _get_server_module() -> Any: + """Import the ComfyUI server module lazily.""" + import server + + return server + + +def _get_execution_module() -> Any: + """Import the ComfyUI execution module lazily.""" + import execution + + return execution + + +def extract_remote_node_ids( + workflow: dict[str, Any] | None, + settings: ModalSyncSettings | None = None, +) -> set[str]: + """Return the node ids marked for remote execution in the workflow metadata.""" + if workflow is None: + return set() + + marker = (settings or get_settings()).marker_property + remote_node_ids: set[str] = set() + for node in workflow.get("nodes", []): + properties = node.get("properties") or {} + if properties.get(marker): + remote_node_ids.add(str(node.get("id"))) + return remote_node_ids + + +def rewrite_prompt_for_modal( + prompt: dict[str, Any], + workflow: dict[str, Any] | None, + sync_engine: ModalAssetSyncEngine | None = None, + settings: ModalSyncSettings | None = None, + nodes_module: Any | None = None, +) -> tuple[dict[str, Any], RewriteSummary]: + """Rewrite remote-marked nodes into signature-preserving Modal proxy nodes.""" + resolved_settings = settings or get_settings() + remote_node_ids = extract_remote_node_ids(workflow, resolved_settings) + summary = RewriteSummary(remote_node_ids=sorted(remote_node_ids)) + + if not remote_node_ids: + return copy.deepcopy(prompt), summary + + resolved_nodes_module = nodes_module or _get_nodes_module() + resolved_sync_engine = sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings) + rewritten_prompt = copy.deepcopy(prompt) + summary.custom_nodes_bundle = resolved_sync_engine.sync_custom_nodes_directory() + + for node_id in sorted(remote_node_ids): + prompt_node = rewritten_prompt.get(node_id) + if prompt_node is None: + logger.warning("Remote node id %s was not present in the prompt payload.", node_id) + continue + + original_class_type = str(prompt_node["class_type"]) + original_class = resolved_nodes_module.NODE_CLASS_MAPPINGS[original_class_type] + proxy_node_id = ensure_modal_proxy_node_registered( + original_class_type=original_class_type, + original_class=original_class, + nodes_module=resolved_nodes_module, + ) + + original_inputs = copy.deepcopy(prompt_node.get("inputs", {})) + remote_inputs, synced_assets = resolved_sync_engine.sync_prompt_inputs(original_inputs) + summary.synced_assets.extend(synced_assets) + + original_node_data = { + "node_id": node_id, + "class_type": original_class_type, + "inputs": remote_inputs, + "meta": copy.deepcopy(prompt_node.get("_meta", {})), + "custom_nodes_bundle": ( + summary.custom_nodes_bundle.remote_path + if summary.custom_nodes_bundle is not None + else None + ), + } + + prompt_node["class_type"] = proxy_node_id + prompt_node["inputs"] = remote_inputs + prompt_node["inputs"]["original_node_data"] = original_node_data + + logger.info( + "Rewrote node %s (%s) to Modal proxy %s.", + node_id, + original_class_type, + proxy_node_id, + ) + + return rewritten_prompt, summary + + +async def _queue_prompt_json(prompt_server: Any, json_data: dict[str, Any]) -> web.Response: + """Queue a possibly rewritten prompt using ComfyUI's native semantics.""" + execution = _get_execution_module() + json_data = prompt_server.trigger_on_prompt(json_data) + + if "number" in json_data: + number = float(json_data["number"]) + else: + number = prompt_server.number + if json_data.get("front"): + number = -number + prompt_server.number += 1 + + if "prompt" not in json_data: + return web.json_response( + { + "error": { + "type": "no_prompt", + "message": "No prompt provided", + "details": "No prompt provided", + "extra_info": {}, + } + }, + status=400, + ) + + prompt = json_data["prompt"] + prompt_id = str(json_data.get("prompt_id", uuid.uuid4())) + partial_execution_targets = json_data.get("partial_execution_targets") + valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets) + + extra_data = dict(json_data.get("extra_data", {})) + if "client_id" in json_data: + extra_data["client_id"] = json_data["client_id"] + + if not valid[0]: + logger.warning("invalid prompt: %s", valid[1]) + return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400) + + outputs_to_execute = valid[2] + sensitive: dict[str, Any] = {} + for sensitive_key in execution.SENSITIVE_EXTRA_DATA_KEYS: + if sensitive_key in extra_data: + sensitive[sensitive_key] = extra_data.pop(sensitive_key) + + extra_data["create_time"] = int(time.time() * 1000) + prompt_server.prompt_queue.put( + (number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive) + ) + return web.json_response( + {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]} + ) + + +def setup_modal_queue_route( + prompt_server: Any | None = None, + sync_engine: ModalAssetSyncEngine | None = None, + settings: ModalSyncSettings | None = None, +) -> None: + """Register the `/modal/queue_prompt` route once for the active PromptServer.""" + global _ROUTE_REGISTERED + if _ROUTE_REGISTERED: + return + + try: + resolved_server_module = _get_server_module() + except ModuleNotFoundError: + logger.debug("ComfyUI server module is not available; skipping route registration.") + return + + resolved_settings = settings or get_settings() + prompt_server = prompt_server or getattr(resolved_server_module.PromptServer, "instance", None) + if prompt_server is None: + logger.debug("PromptServer.instance is not available; skipping route registration.") + return + + resolved_sync_engine = sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings) + + @prompt_server.routes.post(resolved_settings.route_path) + async def modal_queue_prompt(request: web.Request) -> web.Response: + """Handle prompt queue requests that include Modal remote markers.""" + logger.info("Received Modal queue request.") + try: + json_data = await request.json() + extra_pnginfo = ((json_data.get("extra_data") or {}).get("extra_pnginfo") or {}) + workflow = extra_pnginfo.get("workflow") + if "prompt" in json_data: + rewritten_prompt, summary = rewrite_prompt_for_modal( + prompt=json_data["prompt"], + workflow=workflow, + sync_engine=resolved_sync_engine, + settings=resolved_settings, + ) + json_data["prompt"] = rewritten_prompt + json_data.setdefault("extra_data", {}).setdefault("modal", {}) + json_data["extra_data"]["modal"]["remote_node_ids"] = summary.remote_node_ids + json_data["extra_data"]["modal"]["synced_assets"] = [ + asset.remote_path for asset in summary.synced_assets + ] + if summary.custom_nodes_bundle is not None: + json_data["extra_data"]["modal"]["custom_nodes_bundle"] = ( + summary.custom_nodes_bundle.remote_path + ) + return await _queue_prompt_json(prompt_server, json_data) + except FileNotFoundError as exc: + logger.exception("Modal asset sync failed.") + return web.json_response({"error": str(exc), "node_errors": []}, status=400) + except Exception as exc: + logger.exception("Modal queue handler failed.") + return web.json_response({"error": str(exc), "node_errors": []}, status=500) + + _ROUTE_REGISTERED = True + logger.info("Registered Modal queue route at %s", resolved_settings.route_path) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md new file mode 100644 index 0000000..97d1138 --- /dev/null +++ b/docs/LEARNINGS.md @@ -0,0 +1,7 @@ +# Learnings + +## 2026-04-17 + +- The workflow metadata already rides along in `extra_data.extra_pnginfo.workflow`, so the backend can discover `node.properties.is_modal_remote` without any extra frontend metadata channel. +- A single wildcard-output proxy node is not enough for server-side graph rewriting. ComfyUI validates downstream link indices and output arity against the replacement class, so the backend has to generate proxy node classes that mirror the original node’s output signature before validation runs. +- Keeping the sync engine content-addressable makes prompt rewrites cheap after the first upload. The route can block on sync safely as long as repeated requests collapse to a hash marker check. diff --git a/modal_executor_node.py b/modal_executor_node.py new file mode 100644 index 0000000..a6f797d --- /dev/null +++ b/modal_executor_node.py @@ -0,0 +1,199 @@ +"""Dynamic Modal proxy nodes for ComfyUI execution.""" + +from __future__ import annotations + +import hashlib +import json +import logging +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Protocol + +from comfy_api.latest import _io as io + +from .serialization import deserialize_node_outputs, serialize_node_inputs + +logger = logging.getLogger(__name__) + + +class RemoteExecutorClient(Protocol): + """Execution client interface used by Modal proxy nodes.""" + + def execute_node(self, node_data: Mapping[str, Any], kwargs: Mapping[str, Any]) -> Sequence[Any]: + """Execute a node remotely and return its output tuple.""" + + +class ModalRemoteExecutorClient: + """Default execution client backed by the remote Modal app module.""" + + def execute_node(self, node_data: Mapping[str, Any], kwargs: Mapping[str, Any]) -> Sequence[Any]: + """Serialize inputs, invoke the remote engine, and deserialize outputs.""" + from .remote.modal_app import invoke_remote_engine + + response = invoke_remote_engine(dict(node_data), serialize_node_inputs(kwargs)) + return deserialize_node_outputs(response) + + +_REMOTE_EXECUTOR_CLIENT_FACTORY: Callable[[], RemoteExecutorClient] = ModalRemoteExecutorClient +_PROXY_NODE_CACHE: dict[str, type[io.ComfyNode]] = {} + + +def set_remote_executor_client_factory( + factory: Callable[[], RemoteExecutorClient] | None, +) -> None: + """Install a custom client factory, primarily for tests.""" + global _REMOTE_EXECUTOR_CLIENT_FACTORY + _REMOTE_EXECUTOR_CLIENT_FACTORY = factory or ModalRemoteExecutorClient + + +def get_remote_executor_client() -> RemoteExecutorClient: + """Instantiate the configured execution client.""" + return _REMOTE_EXECUTOR_CLIENT_FACTORY() + + +def _output_spec(io_type: str, name: str, is_list: bool) -> io.Output: + """Create a v3 output specification from a legacy ComfyUI return type.""" + comfy_type = io.AnyType if io_type == "*" else io.Custom(io_type) + return comfy_type.Output(display_name=name, is_output_list=is_list) + + +def _normalized_output_metadata( + original_class: type[Any], +) -> tuple[tuple[str, ...], tuple[str, ...], tuple[bool, ...]]: + """Normalize output metadata from a source node class.""" + if hasattr(original_class, "GET_SCHEMA"): + original_class.GET_SCHEMA() + + output_types = tuple(getattr(original_class, "RETURN_TYPES", ("*",))) or ("*",) + default_names = tuple(f"output_{index}" for index, _ in enumerate(output_types)) + output_names = tuple(getattr(original_class, "RETURN_NAMES", default_names)) + output_is_list = tuple(getattr(original_class, "OUTPUT_IS_LIST", (False,) * len(output_types))) + + if len(output_names) < len(output_types): + output_names = output_names + default_names[len(output_names) :] + if len(output_is_list) < len(output_types): + output_is_list = output_is_list + (False,) * (len(output_types) - len(output_is_list)) + + return output_types, output_names[: len(output_types)], output_is_list[: len(output_types)] + + +def _proxy_node_id(original_class_type: str, output_types: Sequence[str]) -> str: + """Build a stable proxy node identifier for an original node signature.""" + digest = hashlib.sha256( + json.dumps({"class_type": original_class_type, "outputs": list(output_types)}).encode( + "utf-8" + ) + ).hexdigest()[:12] + return f"ModalUniversalExecutor_{digest}" + + +def _build_proxy_node_class( + original_class_type: str, + output_types: tuple[str, ...], + output_names: tuple[str, ...], + output_is_list: tuple[bool, ...], +) -> type[io.ComfyNode]: + """Create a v3 proxy node that mirrors an original node output signature.""" + node_id = _proxy_node_id(original_class_type, output_types) + + class _DynamicModalExecutor(io.ComfyNode): + """Internal proxy node that forwards execution to Modal.""" + + @classmethod + def define_schema(cls) -> io.Schema: + """Return a schema that accepts any original node inputs.""" + outputs = [ + _output_spec(io_type, name, is_list) + for io_type, name, is_list in zip(output_types, output_names, output_is_list, strict=False) + ] + return io.Schema( + node_id=node_id, + display_name="Modal Universal Executor", + category="Modal", + description=( + "Internal proxy node that forwards the original node execution " + "to a Modal-backed runtime." + ), + inputs=[ + io.AnyType.Input("original_node_data"), + ], + outputs=outputs, + accept_all_inputs=True, + is_dev_only=True, + is_experimental=True, + ) + + @classmethod + def execute(cls, original_node_data: Any, **kwargs: Any) -> io.NodeOutput: + """Forward the original node payload to the configured remote executor.""" + if isinstance(original_node_data, str): + original_node_data = json.loads(original_node_data) + if not isinstance(original_node_data, Mapping): + raise TypeError("original_node_data must be a mapping or JSON object.") + + outputs = tuple(get_remote_executor_client().execute_node(original_node_data, kwargs)) + logger.debug( + "Remote execution completed for %s with %d outputs.", + original_node_data.get("class_type"), + len(outputs), + ) + return io.NodeOutput(*outputs) + + _DynamicModalExecutor.__name__ = f"DynamicModalExecutor_{node_id}" + return _DynamicModalExecutor + + +def ensure_modal_proxy_node_registered( + original_class_type: str, + original_class: type[Any], + nodes_module: Any, +) -> str: + """Register and return a proxy node id for the supplied original node class.""" + output_types, output_names, output_is_list = _normalized_output_metadata(original_class) + proxy_node_id = _proxy_node_id(original_class_type, output_types) + + if proxy_node_id in _PROXY_NODE_CACHE: + nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = _PROXY_NODE_CACHE[proxy_node_id] + nodes_module.NODE_DISPLAY_NAME_MAPPINGS[proxy_node_id] = "Modal Universal Executor" + return proxy_node_id + + proxy_class = _build_proxy_node_class( + original_class_type=original_class_type, + output_types=output_types, + output_names=output_names, + output_is_list=output_is_list, + ) + nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = proxy_class + nodes_module.NODE_DISPLAY_NAME_MAPPINGS[proxy_node_id] = "Modal Universal Executor" + _PROXY_NODE_CACHE[proxy_node_id] = proxy_class + logger.info("Registered Modal proxy node %s for %s", proxy_node_id, original_class_type) + return proxy_node_id + + +class ModalUniversalExecutor(io.ComfyNode): + """Base debug node for Modal execution routing.""" + + @classmethod + def define_schema(cls) -> io.Schema: + """Define a minimal dev-only schema for the base proxy node.""" + return io.Schema( + node_id="ModalUniversalExecutor", + display_name="Modal Universal Executor", + category="Modal", + description=( + "Debug entrypoint for Modal-backed execution. Production rewrites use " + "signature-preserving dynamic proxy variants." + ), + inputs=[io.AnyType.Input("original_node_data")], + outputs=[io.AnyType.Output(display_name="output")], + accept_all_inputs=True, + is_dev_only=True, + is_experimental=True, + ) + + @classmethod + def execute(cls, original_node_data: Any, **kwargs: Any) -> io.NodeOutput: + """Execute the base debug proxy node through the configured client.""" + if isinstance(original_node_data, str): + original_node_data = json.loads(original_node_data) + outputs = tuple(get_remote_executor_client().execute_node(original_node_data, kwargs)) + return io.NodeOutput(*outputs) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..eeb4633 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "comfyui-modal-sync" +version = "0.1.0" +description = "ComfyUI custom node extension for routing marked nodes to Modal-backed execution." +requires-python = ">=3.11" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.pytest.ini_options] +testpaths = ["tests"] + +[tool.uv] +package = false diff --git a/remote/__init__.py b/remote/__init__.py new file mode 100644 index 0000000..e0eaf4c --- /dev/null +++ b/remote/__init__.py @@ -0,0 +1,5 @@ +"""Remote execution helpers for Modal-Sync.""" + +from .modal_app import execute_node_locally, invoke_remote_engine + +__all__ = ["execute_node_locally", "invoke_remote_engine"] diff --git a/remote/modal_app.py b/remote/modal_app.py new file mode 100644 index 0000000..9c13ae5 --- /dev/null +++ b/remote/modal_app.py @@ -0,0 +1,149 @@ +"""Remote Modal runtime and local execution fallback.""" + +from __future__ import annotations + +import logging +import os +import sys +import tempfile +import zipfile +from pathlib import Path +from typing import Any + +from ..serialization import ( + deserialize_node_inputs, + serialize_node_outputs, +) +from ..settings import get_settings + +logger = logging.getLogger(__name__) + +try: + import modal # type: ignore +except ModuleNotFoundError: # pragma: no cover - exercised by local fallback tests. + modal = None + + +def _extract_custom_nodes_bundle(bundle_path: str | None) -> None: + """Extract a mirrored custom_nodes archive into a temporary import path.""" + if not bundle_path: + return + + settings = get_settings() + local_bundle = settings.local_storage_root / bundle_path.lstrip("/") + if not local_bundle.exists(): + logger.warning("Custom nodes bundle %s was not found in local storage.", local_bundle) + return + + extraction_root = Path(tempfile.gettempdir()) / "comfy-modal-sync-custom-nodes" + extraction_root.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(local_bundle, "r") as archive: + archive.extractall(extraction_root) + + if str(extraction_root) not in sys.path: + sys.path.insert(0, str(extraction_root)) + logger.info("Extracted remote custom_nodes bundle to %s", extraction_root) + + +def _load_nodes_module() -> Any: + """Import the ComfyUI nodes module lazily.""" + import nodes + + return nodes + + +def _invoke_original_node( + node_class: type[Any], + node_data: dict[str, Any], + kwargs: dict[str, Any], +) -> tuple[Any, ...]: + """Execute an original V1 or V3 node class and normalize its outputs.""" + class_type = node_data["class_type"] + logger.info("Executing remote node %s", class_type) + + if hasattr(node_class, "GET_SCHEMA"): + node_output = node_class.execute(**kwargs) + if hasattr(node_output, "result"): + result = node_output.result + return tuple(result) if result is not None else tuple() + return tuple(node_output) + + instance = node_class() + function_name = getattr(node_class, "FUNCTION", "execute") + function = getattr(instance, function_name) + result = function(**kwargs) + if result is None: + return tuple() + if isinstance(result, tuple): + return result + if isinstance(result, list): + return tuple(result) + return (result,) + + +def execute_node_locally( + node_data: dict[str, Any], + kwargs_payload: bytes | bytearray | str | dict[str, Any], + node_mapping: dict[str, type[Any]] | None = None, +) -> bytes: + """Execute the target node in-process and return serialized outputs.""" + _extract_custom_nodes_bundle(node_data.get("custom_nodes_bundle")) + kwargs = deserialize_node_inputs(kwargs_payload) + resolved_node_mapping = node_mapping or _load_nodes_module().NODE_CLASS_MAPPINGS + class_type = node_data["class_type"] + if class_type not in resolved_node_mapping: + raise KeyError(f"Remote node class {class_type!r} is not registered.") + + outputs = _invoke_original_node(resolved_node_mapping[class_type], node_data, kwargs) + return serialize_node_outputs(outputs) + + +def invoke_remote_engine(node_data: dict[str, Any], kwargs_payload: bytes) -> bytes: + """Invoke Modal when configured, or fall back to local in-process execution.""" + execution_mode = os.getenv("COMFY_MODAL_EXECUTION_MODE", "local") + if execution_mode == "local" or modal is None: + return execute_node_locally(node_data, kwargs_payload) + + engine = RemoteEngine() + remote_method = getattr(engine.execute_node, "remote", None) + if callable(remote_method): + return remote_method(node_data, kwargs_payload) + return engine.execute_node(node_data, kwargs_payload) + + +if modal is not None: # pragma: no branch - simple import-time configuration. + settings = get_settings() + app = modal.App(settings.app_name) + vol = modal.Volume.from_name(settings.volume_name, create_if_missing=True) + image = modal.Image.debian_slim().pip_install("torch", "safetensors", "pillow", "numpy") + + @app.cls( + gpu="A100", + volumes={settings.remote_storage_root: vol}, + scaledown_window=60, + image=image, + ) + class RemoteEngine: + """Modal runtime class that executes proxied ComfyUI nodes.""" + + @modal.enter() + def setup(self) -> None: + """Prepare the container process for headless node execution.""" + logger.info("RemoteEngine setup complete.") + + @modal.method() + def execute_node(self, node_data: dict[str, Any], kwargs_payload: bytes) -> bytes: + """Execute the proxied node inside the Modal container.""" + return execute_node_locally(node_data, kwargs_payload) + +else: + + class RemoteEngine: + """Local fallback runtime used when the Modal SDK is unavailable.""" + + def setup(self) -> None: + """No-op setup for local fallback execution.""" + + def execute_node(self, node_data: dict[str, Any], kwargs_payload: bytes) -> bytes: + """Execute the proxied node locally.""" + return execute_node_locally(node_data, kwargs_payload) diff --git a/serialization.py b/serialization.py new file mode 100644 index 0000000..aacfc27 --- /dev/null +++ b/serialization.py @@ -0,0 +1,152 @@ +"""Serialization helpers for Modal node execution payloads.""" + +from __future__ import annotations + +import base64 +import json +import logging +from collections.abc import Mapping, Sequence +from typing import Any + +logger = logging.getLogger(__name__) + +_KIND_KEY = "__comfy_modal_kind__" +_VALUE_KEY = "value" +_TENSOR_KIND = "tensor" +_BYTES_KIND = "bytes" +_TUPLE_KIND = "tuple" + + +def _is_scalar(value: Any) -> bool: + """Return whether a value is natively representable in JSON.""" + return value is None or isinstance(value, bool | int | float | str) + + +def _import_torch() -> Any: + """Import torch lazily so the module stays importable in light environments.""" + import torch + + return torch + + +def _serialize_tensor(value: Any) -> dict[str, str]: + """Serialize a torch.Tensor into a base64 safetensors payload.""" + from safetensors.torch import save + + torch = _import_torch() + if not isinstance(value, torch.Tensor): + raise TypeError("Expected a torch.Tensor payload.") + + tensor_bytes = save({_VALUE_KEY: value.detach().contiguous()}) + encoded = base64.b64encode(tensor_bytes).decode("ascii") + return {_KIND_KEY: _TENSOR_KIND, "payload": encoded} + + +def _deserialize_tensor(payload: Mapping[str, Any]) -> Any: + """Deserialize a base64 safetensors payload back into a tensor.""" + from safetensors.torch import load + + encoded = payload["payload"] + tensor_map = load(base64.b64decode(encoded.encode("ascii"))) + return tensor_map[_VALUE_KEY] + + +def serialize_value(value: Any) -> Any: + """Convert a Python value into a JSON-safe execution payload.""" + if _is_scalar(value): + return value + + try: + torch = _import_torch() + except ModuleNotFoundError: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + return _serialize_tensor(value) + + if isinstance(value, bytes): + return { + _KIND_KEY: _BYTES_KIND, + "payload": base64.b64encode(value).decode("ascii"), + } + + if isinstance(value, tuple): + return { + _KIND_KEY: _TUPLE_KIND, + "items": [serialize_value(item) for item in value], + } + + if isinstance(value, list): + return [serialize_value(item) for item in value] + + if isinstance(value, Mapping): + return {str(key): serialize_value(item) for key, item in value.items()} + + raise TypeError( + "ComfyUI-Modal can only serialize JSON-compatible values, bytes, " + "and torch tensors. Unsupported value type: " + f"{type(value)!r}" + ) + + +def deserialize_value(payload: Any) -> Any: + """Reconstruct a serialized execution payload back into Python values.""" + if _is_scalar(payload): + return payload + + if isinstance(payload, list): + return [deserialize_value(item) for item in payload] + + if not isinstance(payload, Mapping): + raise TypeError(f"Unsupported payload type: {type(payload)!r}") + + kind = payload.get(_KIND_KEY) + if kind == _TENSOR_KIND: + return _deserialize_tensor(payload) + if kind == _BYTES_KIND: + encoded = payload["payload"] + return base64.b64decode(encoded.encode("ascii")) + if kind == _TUPLE_KIND: + return tuple(deserialize_value(item) for item in payload["items"]) + + return {str(key): deserialize_value(value) for key, value in payload.items()} + + +def serialize_node_inputs(inputs: Mapping[str, Any]) -> bytes: + """Serialize node keyword arguments into transport bytes.""" + payload = {str(key): serialize_value(value) for key, value in inputs.items()} + return json.dumps(payload, sort_keys=True).encode("utf-8") + + +def deserialize_node_inputs(payload: bytes | bytearray | str | Mapping[str, Any]) -> dict[str, Any]: + """Deserialize node keyword arguments from transport bytes.""" + if isinstance(payload, Mapping): + raw_payload = dict(payload) + else: + if isinstance(payload, (bytes, bytearray)): + payload = payload.decode("utf-8") + raw_payload = json.loads(payload) + + if not isinstance(raw_payload, Mapping): + raise TypeError("Serialized node inputs must decode to a mapping.") + return {str(key): deserialize_value(value) for key, value in raw_payload.items()} + + +def serialize_node_outputs(outputs: Sequence[Any]) -> bytes: + """Serialize node outputs into transport bytes.""" + payload = [serialize_value(value) for value in outputs] + return json.dumps(payload).encode("utf-8") + + +def deserialize_node_outputs(payload: bytes | bytearray | str | Sequence[Any]) -> tuple[Any, ...]: + """Deserialize node outputs from transport bytes.""" + if isinstance(payload, Sequence) and not isinstance(payload, (bytes, bytearray, str)): + raw_payload = list(payload) + else: + if isinstance(payload, (bytes, bytearray)): + payload = payload.decode("utf-8") + raw_payload = json.loads(payload) + + if not isinstance(raw_payload, list): + raise TypeError("Serialized node outputs must decode to a list.") + return tuple(deserialize_value(value) for value in raw_payload) diff --git a/settings.py b/settings.py new file mode 100644 index 0000000..a3106e4 --- /dev/null +++ b/settings.py @@ -0,0 +1,98 @@ +"""Runtime configuration helpers for ComfyUI Modal-Sync.""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ModalSyncSettings: + """Resolved configuration for local and remote Modal-Sync operations.""" + + app_name: str + volume_name: str + route_path: str + marker_property: str + local_storage_root: Path + remote_storage_root: str + custom_nodes_archive_name: str + comfyui_root: Path | None + custom_nodes_dir: Path | None + + +def _read_path_env(name: str) -> Path | None: + """Resolve an environment variable into a path when present.""" + value = os.getenv(name) + if not value: + return None + return Path(value).expanduser().resolve() + + +def _discover_repo_root() -> Path: + """Return the repository root containing this module.""" + return Path(__file__).resolve().parent + + +def _discover_comfyui_root() -> Path | None: + """Locate the local ComfyUI checkout used for tests and path resolution.""" + env_root = _read_path_env("COMFYUI_ROOT") + if env_root is not None: + return env_root + + default_root = Path.home() / "git" / "ComfyUI" + if default_root.exists(): + return default_root.resolve() + + return None + + +def _discover_custom_nodes_dir(repo_root: Path, comfyui_root: Path | None) -> Path | None: + """Locate the custom_nodes directory that should be mirrored to Modal.""" + env_dir = _read_path_env("COMFY_MODAL_CUSTOM_NODES_DIR") + if env_dir is not None: + return env_dir + + if repo_root.parent.name == "custom_nodes": + return repo_root.parent.resolve() + + if comfyui_root is not None: + candidate = comfyui_root / "custom_nodes" + if candidate.exists(): + return candidate.resolve() + + return None + + +@lru_cache(maxsize=1) +def get_settings() -> ModalSyncSettings: + """Return cached extension settings derived from the environment.""" + repo_root = _discover_repo_root() + comfyui_root = _discover_comfyui_root() + custom_nodes_dir = _discover_custom_nodes_dir(repo_root, comfyui_root) + local_storage_root = ( + _read_path_env("COMFY_MODAL_LOCAL_STORAGE_ROOT") + or Path("/tmp/comfyui-modal-sync-storage") + ) + + settings = ModalSyncSettings( + app_name=os.getenv("COMFY_MODAL_APP_NAME", "comfy-modal-sync"), + volume_name=os.getenv("COMFY_MODAL_VOLUME_NAME", "comfy-universal-storage"), + route_path=os.getenv("COMFY_MODAL_ROUTE_PATH", "/modal/queue_prompt"), + marker_property=os.getenv("COMFY_MODAL_MARKER_PROPERTY", "is_modal_remote"), + local_storage_root=local_storage_root.resolve(), + remote_storage_root=os.getenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", "/storage"), + custom_nodes_archive_name=os.getenv( + "COMFY_MODAL_CUSTOM_NODES_ARCHIVE", + "custom_nodes_bundle.zip", + ), + comfyui_root=comfyui_root, + custom_nodes_dir=custom_nodes_dir, + ) + logger.debug("Resolved Modal-Sync settings: %s", settings) + return settings diff --git a/sync_engine.py b/sync_engine.py new file mode 100644 index 0000000..0806df7 --- /dev/null +++ b/sync_engine.py @@ -0,0 +1,236 @@ +"""Asset synchronization helpers for Modal-backed execution.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import tempfile +import zipfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Protocol + +from .settings import ModalSyncSettings, get_settings + +logger = logging.getLogger(__name__) + +_SYNC_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".vae"} +_SKIP_DIRS = {".git", "__pycache__", ".pytest_cache", ".mypy_cache"} + + +class VolumeBackend(Protocol): + """Minimal storage interface needed by the sync engine.""" + + def exists(self, remote_path: str) -> bool: + """Return whether the remote path already exists.""" + + def put_file(self, local_path: Path, remote_path: str) -> None: + """Upload a local file into the remote storage backend.""" + + def put_bytes(self, payload: bytes, remote_path: str) -> None: + """Upload raw bytes into the remote storage backend.""" + + +@dataclass(frozen=True) +class SyncedAsset: + """Description of a local asset mirrored into the remote storage root.""" + + local_path: Path + remote_path: str + sha256: str + + +class LocalMirrorVolume: + """Simple filesystem-backed volume used for tests and dry runs.""" + + def __init__(self, root: Path) -> None: + """Initialize the local mirror volume root.""" + self.root = root + self.root.mkdir(parents=True, exist_ok=True) + + def exists(self, remote_path: str) -> bool: + """Return whether a file already exists in the local mirror.""" + return self._resolve(remote_path).exists() + + def put_file(self, local_path: Path, remote_path: str) -> None: + """Copy a local file into the mirror volume.""" + target = self._resolve(remote_path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(local_path.read_bytes()) + + def put_bytes(self, payload: bytes, remote_path: str) -> None: + """Write bytes into the mirror volume.""" + target = self._resolve(remote_path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(payload) + + def _resolve(self, remote_path: str) -> Path: + """Resolve a remote storage path relative to the mirror root.""" + return self.root / remote_path.lstrip("/") + + +@dataclass +class ModalAssetSyncEngine: + """Content-addressable storage sync engine for files and custom nodes.""" + + volume: VolumeBackend + settings: ModalSyncSettings + + @classmethod + 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) + return cls(volume=backend, settings=resolved_settings) + + def sync_file(self, local_path: Path, remote_folder: str = "/assets") -> SyncedAsset: + """Sync a file into content-addressable remote storage.""" + resolved_path = local_path.expanduser().resolve() + if not resolved_path.is_file(): + raise FileNotFoundError(f"Asset not found: {resolved_path}") + + sha256 = self._hash_file(resolved_path) + marker_path = f"/hashes/{sha256}.done" + remote_path = f"{remote_folder.rstrip('/')}/{sha256}_{resolved_path.name}" + + if not self.volume.exists(marker_path): + logger.info("Syncing %s to %s", resolved_path, remote_path) + self.volume.put_file(resolved_path, remote_path) + self.volume.put_bytes( + json.dumps({"source": str(resolved_path), "remote_path": remote_path}).encode( + "utf-8" + ), + marker_path, + ) + else: + logger.debug("Asset already mirrored for sha=%s", sha256) + + return SyncedAsset(local_path=resolved_path, remote_path=remote_path, sha256=sha256) + + def sync_prompt_inputs(self, inputs: dict[str, Any]) -> tuple[dict[str, Any], list[SyncedAsset]]: + """Rewrite file-like prompt inputs to mirrored storage paths.""" + synced_assets: list[SyncedAsset] = [] + + def rewrite(value: Any) -> Any: + if isinstance(value, str): + maybe_path = self._resolve_model_path(value) + if maybe_path is not None: + synced_asset = self.sync_file(maybe_path) + synced_assets.append(synced_asset) + return synced_asset.remote_path + return value + if isinstance(value, list): + return [rewrite(item) for item in value] + if isinstance(value, dict): + return {str(key): rewrite(item) for key, item in value.items()} + return value + + return rewrite(inputs), synced_assets + + def sync_custom_nodes_directory(self) -> SyncedAsset | None: + """Zip and mirror the local custom_nodes directory when available.""" + custom_nodes_dir = self.settings.custom_nodes_dir + if custom_nodes_dir is None or not custom_nodes_dir.exists(): + logger.info("No custom_nodes directory detected for mirroring.") + return None + + directory_hash = self._hash_directory(custom_nodes_dir) + archive_name = self.settings.custom_nodes_archive_name + marker_path = f"/hashes/custom_nodes_{directory_hash}.done" + remote_path = f"/custom_nodes/{directory_hash}_{archive_name}" + + if self.volume.exists(marker_path): + return SyncedAsset( + local_path=custom_nodes_dir, + remote_path=remote_path, + sha256=directory_hash, + ) + + archive_path = self._create_archive(custom_nodes_dir) + try: + logger.info("Syncing custom_nodes bundle from %s to %s", custom_nodes_dir, remote_path) + self.volume.put_file(archive_path, remote_path) + self.volume.put_bytes( + json.dumps({"source": str(custom_nodes_dir), "remote_path": remote_path}).encode( + "utf-8" + ), + marker_path, + ) + finally: + archive_path.unlink(missing_ok=True) + + return SyncedAsset(local_path=custom_nodes_dir, remote_path=remote_path, sha256=directory_hash) + + def _resolve_model_path(self, value: str) -> Path | None: + """Resolve a prompt string into a local model file path when possible.""" + path = Path(value).expanduser() + if path.suffix.lower() not in _SYNC_EXTENSIONS: + return None + if path.is_file(): + return path.resolve() + + if os.path.isabs(value): + return None + + try: + import folder_paths + except ModuleNotFoundError: + folder_paths = None + + if folder_paths is not None: + for folder_name in folder_paths.folder_names_and_paths: + full_path = folder_paths.get_full_path(folder_name, value) + if full_path is not None: + return Path(full_path).resolve() + + if self.settings.comfyui_root is not None: + candidate = self.settings.comfyui_root / value + if candidate.is_file(): + return candidate.resolve() + + return None + + def _hash_file(self, path: Path) -> str: + """Compute the SHA256 digest for a file.""" + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def _hash_directory(self, path: Path) -> str: + """Compute a stable SHA256 digest for a directory tree.""" + digest = hashlib.sha256() + for child in sorted(self._iter_files(path), key=lambda item: item.relative_to(path).as_posix()): + relative_path = child.relative_to(path).as_posix() + digest.update(relative_path.encode("utf-8")) + digest.update(b"\0") + digest.update(self._hash_file(child).encode("ascii")) + digest.update(b"\0") + return digest.hexdigest() + + def _create_archive(self, path: Path) -> Path: + """Create a zip archive for the given directory tree.""" + with tempfile.NamedTemporaryFile( + prefix="comfy-modal-custom-nodes-", + suffix=".zip", + delete=False, + ) as handle: + archive_path = Path(handle.name) + + with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as archive: + for child in sorted(self._iter_files(path), key=lambda item: item.relative_to(path).as_posix()): + archive.write(child, arcname=child.relative_to(path)) + + return archive_path + + def _iter_files(self, path: Path) -> list[Path]: + """Yield files from a directory tree while skipping cache folders.""" + files: list[Path] = [] + for root, dirnames, filenames in os.walk(path): + dirnames[:] = [name for name in dirnames if name not in _SKIP_DIRS] + for filename in filenames: + files.append(Path(root) / filename) + return files diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..163dee0 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,113 @@ +"""Pytest fixtures for loading the ComfyUI Modal-Sync extension package.""" + +from __future__ import annotations + +import importlib +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +COMFYUI_ROOT = Path.home() / "git" / "ComfyUI" +PACKAGE_NAME = "comfyui_modal_sync_under_test" + + +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.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_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", (), {}) + + 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 + if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + if COMFYUI_ROOT.exists() and str(COMFYUI_ROOT) not in sys.path: + sys.path.insert(0, str(COMFYUI_ROOT)) + + +def _load_extension_package() -> object: + """Load the extension root as an importable package for test modules.""" + _ensure_import_paths() + if PACKAGE_NAME in sys.modules: + return sys.modules[PACKAGE_NAME] + + spec = importlib.util.spec_from_file_location( + PACKAGE_NAME, + REPO_ROOT / "__init__.py", + submodule_search_locations=[str(REPO_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to create an import spec for the extension package.") + + module = importlib.util.module_from_spec(spec) + sys.modules[PACKAGE_NAME] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="session") +def extension_package() -> object: + """Return the loaded extension package.""" + return _load_extension_package() + + +@pytest.fixture(scope="session") +def api_intercept_module(extension_package: object) -> object: + """Return the prompt interception module.""" + return importlib.import_module(f"{PACKAGE_NAME}.api_intercept") + + +@pytest.fixture(scope="session") +def modal_executor_module(extension_package: object) -> object: + """Return the dynamic proxy node module.""" + return importlib.import_module(f"{PACKAGE_NAME}.modal_executor_node") + + +@pytest.fixture(scope="session") +def remote_modal_app_module(extension_package: object) -> object: + """Return the remote execution module.""" + return importlib.import_module(f"{PACKAGE_NAME}.remote.modal_app") + + +@pytest.fixture(scope="session") +def serialization_module(extension_package: object) -> object: + """Return the serialization helpers module.""" + return importlib.import_module(f"{PACKAGE_NAME}.serialization") + + +@pytest.fixture(scope="session") +def settings_module(extension_package: object) -> object: + """Return the settings module.""" + return importlib.import_module(f"{PACKAGE_NAME}.settings") + + +@pytest.fixture(scope="session") +def sync_engine_module(extension_package: object) -> object: + """Return the sync engine module.""" + return importlib.import_module(f"{PACKAGE_NAME}.sync_engine") + + +@pytest.fixture(autouse=True) +def reset_modal_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Isolate Modal-Sync environment variables between tests.""" + monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "local") + monkeypatch.setenv("COMFY_MODAL_LOCAL_STORAGE_ROOT", str(tmp_path / "storage")) + monkeypatch.delenv("COMFY_MODAL_CUSTOM_NODES_DIR", raising=False) + monkeypatch.delenv("COMFYUI_ROOT", raising=False) diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py new file mode 100644 index 0000000..9bf8b2f --- /dev/null +++ b/tests/test_api_intercept.py @@ -0,0 +1,88 @@ +"""Tests for prompt rewriting and asset sync integration.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + + +class _FakeOriginalNode: + """Simple fake source node for rewrite tests.""" + + RETURN_TYPES = ("IMAGE", "INT") + RETURN_NAMES = ("image", "count") + OUTPUT_IS_LIST = (False, False) + + +def test_rewrite_prompt_for_remote_nodes( + api_intercept_module: Any, + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """Prompt rewrite should proxy remote-marked nodes and mirror their assets.""" + model_path = tmp_path / "weights.safetensors" + model_path.write_bytes(b"weights") + custom_nodes_dir = tmp_path / "custom_nodes" + custom_nodes_dir.mkdir() + (custom_nodes_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") + + settings = settings_module.ModalSyncSettings( + app_name="app", + 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=custom_nodes_dir, + ) + sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + + workflow = { + "nodes": [ + {"id": 1, "properties": {"is_modal_remote": True}}, + {"id": 2, "properties": {"is_modal_remote": False}}, + ] + } + prompt = { + "1": { + "class_type": "OriginalNode", + "inputs": { + "model_name": str(model_path), + "strength": 0.5, + }, + "_meta": {"title": "Remote Node"}, + }, + "2": { + "class_type": "OtherNode", + "inputs": {}, + "_meta": {"title": "Local Node"}, + }, + } + + rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal( + prompt=prompt, + workflow=workflow, + sync_engine=sync_engine, + settings=settings, + nodes_module=fake_nodes_module, + ) + + rewritten_node = rewritten_prompt["1"] + assert rewritten_node["class_type"].startswith("ModalUniversalExecutor_") + assert rewritten_node["inputs"]["model_name"].startswith("/assets/") + assert rewritten_node["inputs"]["original_node_data"]["class_type"] == "OriginalNode" + assert rewritten_prompt["2"]["class_type"] == "OtherNode" + assert summary.remote_node_ids == ["1"] + assert len(summary.synced_assets) == 1 + assert summary.custom_nodes_bundle is not None diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py new file mode 100644 index 0000000..6acb9b4 --- /dev/null +++ b/tests/test_modal_executor_node.py @@ -0,0 +1,95 @@ +"""Tests for dynamic Modal proxy nodes and local execution fallback.""" + +from __future__ import annotations + +from typing import Any + + +class _FakeOriginalNode: + """Simple fake legacy node for proxy signature mirroring.""" + + RETURN_TYPES = ("IMAGE", "INT") + RETURN_NAMES = ("image", "count") + OUTPUT_IS_LIST = (False, False) + FUNCTION = "run" + + def run(self, **kwargs: Any) -> tuple[Any, ...]: + """Return a tuple that exposes the inputs for verification.""" + return (kwargs["value"], 1) + + +def test_dynamic_proxy_node_preserves_output_signature( + modal_executor_module: Any, +) -> None: + """Dynamic Modal proxies should mirror the original output count and names.""" + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + + proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( + original_class_type="OriginalNode", + original_class=_FakeOriginalNode, + nodes_module=fake_nodes_module, + ) + + proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] + schema = proxy_class.GET_SCHEMA() + + assert schema.node_id == proxy_id + assert [output.display_name for output in schema.outputs] == ["image", "count"] + assert [output.io_type for output in schema.outputs] == ["IMAGE", "INT"] + + +def test_proxy_execution_uses_injected_remote_client( + modal_executor_module: Any, +) -> None: + """Proxy execution should delegate to the configured remote client.""" + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + + proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( + original_class_type="OriginalNode", + original_class=_FakeOriginalNode, + nodes_module=fake_nodes_module, + ) + proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] + + class FakeClient: + """Test client that returns deterministic outputs.""" + + def execute_node(self, node_data: dict[str, Any], kwargs: dict[str, Any]) -> tuple[str, int]: + """Return values derived from the proxied node payload.""" + return (f"{node_data['class_type']}::{kwargs['value']}", 3) + + modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) + try: + result = proxy_class.execute({"class_type": "OriginalNode"}, value="payload") + finally: + modal_executor_module.set_remote_executor_client_factory(None) + + assert result.result == ("OriginalNode::payload", 3) + + +def test_local_remote_app_executes_original_node( + remote_modal_app_module: Any, + serialization_module: Any, +) -> None: + """The local fallback remote app should execute a mapped legacy node.""" + payload = remote_modal_app_module.execute_node_locally( + node_data={"class_type": "OriginalNode"}, + kwargs_payload='{"value": "hello"}', + node_mapping={"OriginalNode": _FakeOriginalNode}, + ) + outputs = serialization_module.deserialize_node_outputs(payload) + assert outputs == ("hello", 1) diff --git a/tests/test_serialization.py b/tests/test_serialization.py new file mode 100644 index 0000000..884dd61 --- /dev/null +++ b/tests/test_serialization.py @@ -0,0 +1,37 @@ +"""Tests for Modal payload serialization helpers.""" + +from __future__ import annotations + +from typing import Any + +import pytest + + +def test_nested_payload_round_trip(serialization_module: Any) -> None: + """Nested JSON-safe values should survive a full round trip.""" + payload = { + "integer": 7, + "text": "hello", + "items": [1, True, None, {"nested": ("a", "b")}], + "raw_bytes": b"abc", + } + + encoded = serialization_module.serialize_node_inputs(payload) + decoded = serialization_module.deserialize_node_inputs(encoded) + + assert decoded["integer"] == 7 + assert decoded["text"] == "hello" + assert decoded["items"][3]["nested"] == ("a", "b") + assert decoded["raw_bytes"] == b"abc" + + +def test_tensor_round_trip(serialization_module: Any) -> None: + """Torch tensors should round-trip through the safetensors transport.""" + torch = pytest.importorskip("torch") + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + encoded = serialization_module.serialize_node_outputs((tensor,)) + decoded = serialization_module.deserialize_node_outputs(encoded) + + assert len(decoded) == 1 + assert torch.equal(decoded[0], tensor) diff --git a/tests/test_sync_engine.py b/tests/test_sync_engine.py new file mode 100644 index 0000000..cbdf749 --- /dev/null +++ b/tests/test_sync_engine.py @@ -0,0 +1,68 @@ +"""Tests for asset syncing and custom_nodes archiving.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + + +def test_sync_file_deduplicates_by_hash( + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """Repeated file syncs should reuse the same remote path and marker.""" + asset_path = tmp_path / "model.safetensors" + asset_path.write_bytes(b"model-bytes") + + settings = settings_module.ModalSyncSettings( + app_name="app", + 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) + first = engine.sync_file(asset_path) + second = engine.sync_file(asset_path) + + assert first.remote_path == second.remote_path + assert first.sha256 == second.sha256 + assert (settings.local_storage_root / first.remote_path.lstrip("/")).exists() + assert (settings.local_storage_root / f"hashes/{first.sha256}.done").exists() + + +def test_sync_custom_nodes_directory_creates_archive( + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """The sync engine should archive and mirror a custom_nodes directory.""" + custom_nodes_dir = tmp_path / "custom_nodes" + package_dir = custom_nodes_dir / "example" + package_dir.mkdir(parents=True) + (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8") + + settings = settings_module.ModalSyncSettings( + app_name="app", + 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=custom_nodes_dir, + ) + + engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + bundle = engine.sync_custom_nodes_directory() + + assert bundle is not None + assert bundle.remote_path.startswith("/custom_nodes/") + assert (settings.local_storage_root / bundle.remote_path.lstrip("/")).exists() diff --git a/web/modal_toggle.js b/web/modal_toggle.js new file mode 100644 index 0000000..0b0e6a9 --- /dev/null +++ b/web/modal_toggle.js @@ -0,0 +1,148 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const REMOTE_PROPERTY = "is_modal_remote"; +const MODAL_ROUTE = "/modal/queue_prompt"; +const REMOTE_BORDER_COLOR = "#1d9bf0"; +const INTERNAL_NODE_PREFIX = "ModalUniversalExecutor"; + +/** + * Return whether a node should show the Modal toggle. + * @param {LGraphNode} node + * @returns {boolean} + */ +function isEligibleNode(node) { + return Boolean(node?.comfyClass) && !String(node.comfyClass).startsWith(INTERNAL_NODE_PREFIX); +} + +/** + * Read the remote execution flag from node properties. + * @param {LGraphNode} node + * @returns {boolean} + */ +function isRemoteNode(node) { + return Boolean(node?.properties?.[REMOTE_PROPERTY]); +} + +/** + * Update the node state and redraw the canvas. + * @param {LGraphNode} node + * @param {boolean} value + */ +function setRemoteFlag(node, value) { + node.properties ||= {}; + node.properties[REMOTE_PROPERTY] = Boolean(value); + app.graph.setDirtyCanvas(true, true); +} + +/** + * Inject the toggle widget and remote border behavior into a node. + * @param {LGraphNode} node + */ +function decorateNode(node) { + if (!isEligibleNode(node) || node.__modalToggleInjected) { + return; + } + + node.properties ||= {}; + node.properties[REMOTE_PROPERTY] = Boolean(node.properties[REMOTE_PROPERTY]); + + const widget = node.addWidget( + "toggle", + "Run on Modal", + node.properties[REMOTE_PROPERTY], + (value) => setRemoteFlag(node, value), + { + on: "Enabled", + off: "Disabled", + serialize: false, + }, + ); + + widget.value = node.properties[REMOTE_PROPERTY]; + node.__modalToggleInjected = true; + node.__modalToggleWidget = widget; + + const originalDrawForeground = node.onDrawForeground; + node.onDrawForeground = function onDrawForeground(ctx) { + originalDrawForeground?.apply(this, arguments); + if (!isRemoteNode(this)) { + return; + } + + const titleHeight = this.constructor?.title_height ?? LiteGraph.NODE_TITLE_HEIGHT ?? 24; + const borderWidth = 3 / app.canvas.ds.scale; + ctx.save(); + ctx.strokeStyle = REMOTE_BORDER_COLOR; + ctx.lineWidth = borderWidth; + ctx.shadowColor = "rgba(29, 155, 240, 0.35)"; + ctx.shadowBlur = 8 / app.canvas.ds.scale; + ctx.strokeRect( + -borderWidth, + -titleHeight, + this.size[0] + borderWidth * 2, + this.size[1] + titleHeight + borderWidth, + ); + ctx.restore(); + }; +} + +/** + * Patch the queue API so prompt submission goes through the Modal route. + */ +function patchQueuePrompt() { + if (api.__modalQueuePromptPatched) { + return; + } + + const originalQueuePrompt = api.queuePrompt.bind(api); + api.queuePrompt = async function modalQueuePrompt(number, data, options) { + const { output: prompt, workflow } = data; + const body = { + client_id: this.clientId ?? "", + prompt, + ...(options?.partialExecutionTargets && { + partial_execution_targets: options.partialExecutionTargets, + }), + extra_data: { + auth_token_comfy_org: this.authToken, + api_key_comfy_org: this.apiKey, + extra_pnginfo: { workflow }, + }, + }; + + if (number === -1) { + body.front = true; + } else if (number !== 0) { + body.number = number; + } + + const response = await this.fetchApi(MODAL_ROUTE, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + }); + + if (response.status !== 200) { + return originalQueuePrompt(number, data, options); + } + + return await response.json(); + }; + + api.__modalQueuePromptPatched = true; +} + +app.registerExtension({ + name: "Comfy.ModalSync.Toggle", + + async init() { + patchQueuePrompt(); + }, + + async nodeCreated(node) { + decorateNode(node); + }, +});