274 lines
9.8 KiB
Python
274 lines
9.8 KiB
Python
"""Tests for dynamic Modal proxy nodes and local execution fallback."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
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_payload(self, payload: dict[str, Any], kwargs: dict[str, Any]) -> tuple[str, int]:
|
|
"""Return values derived from the proxied node payload."""
|
|
return (f"{payload['class_type']}::{kwargs['value']}", 3)
|
|
|
|
modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
|
|
try:
|
|
result = proxy_class.execute(original_node_data={"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)
|
|
|
|
|
|
def test_stable_modal_cloud_entry_imports_without_modal_sdk(
|
|
modal_cloud_module: Any,
|
|
) -> None:
|
|
"""The stable Modal cloud module should stay importable when modal is unavailable."""
|
|
assert modal_cloud_module.__name__ == "comfyui_modal_sync_cloud"
|
|
assert hasattr(modal_cloud_module, "RemoteEngine")
|
|
|
|
|
|
def test_modal_cloud_ignores_heavy_comfyui_paths(
|
|
modal_cloud_module: Any,
|
|
) -> None:
|
|
"""The Modal cloud module should skip heavyweight ComfyUI runtime artifacts."""
|
|
from pathlib import Path
|
|
|
|
assert modal_cloud_module._should_ignore_comfyui_path(Path("models/checkpoint.safetensors"))
|
|
assert modal_cloud_module._should_ignore_comfyui_path(Path("output/run/output.png"))
|
|
assert modal_cloud_module._should_ignore_comfyui_path(Path("__pycache__/execution.pyc"))
|
|
assert not modal_cloud_module._should_ignore_comfyui_path(Path("execution.py"))
|
|
assert not modal_cloud_module._should_ignore_comfyui_path(Path("comfy/model_management.py"))
|
|
assert not modal_cloud_module._should_ignore_comfyui_path(Path("comfy/ldm/models/diffusion/ddpm.py"))
|
|
|
|
|
|
def test_modal_cloud_installs_comfyui_runtime_packages(
|
|
modal_cloud_module: Any,
|
|
) -> None:
|
|
"""The Modal cloud image should include the core packages ComfyUI imports at runtime."""
|
|
packages = set(modal_cloud_module._comfyui_runtime_packages())
|
|
|
|
assert "psutil" in packages
|
|
assert "torch" in packages
|
|
assert "torchvision" in packages
|
|
assert "torchsde" in packages
|
|
assert "transformers" in packages
|
|
assert "sentencepiece" in packages
|
|
assert "aiohttp" in packages
|
|
assert "opencv-python-headless" in packages
|
|
|
|
|
|
def test_modal_cloud_initializes_remote_comfy_runtime_once_per_custom_node_root(
|
|
modal_cloud_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Any,
|
|
) -> None:
|
|
"""The remote runtime should load built-in extras once and custom bundles per extracted root."""
|
|
init_calls: list[tuple[Any, ...]] = []
|
|
folder_path_calls: list[tuple[str, str, bool]] = []
|
|
|
|
fake_nodes_module = types.SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={},
|
|
NODE_DISPLAY_NAME_MAPPINGS={},
|
|
)
|
|
|
|
async def fake_init_extra_nodes(init_custom_nodes: bool = True, init_api_nodes: bool = True) -> None:
|
|
init_calls.append(("extra", init_custom_nodes, init_api_nodes))
|
|
|
|
async def fake_init_external_custom_nodes() -> None:
|
|
init_calls.append(("external",))
|
|
|
|
fake_nodes_module.init_extra_nodes = fake_init_extra_nodes
|
|
fake_nodes_module.init_external_custom_nodes = fake_init_external_custom_nodes
|
|
|
|
fake_folder_paths_module = types.SimpleNamespace(
|
|
add_model_folder_path=lambda folder_name, full_folder_path, is_default=False: folder_path_calls.append(
|
|
(folder_name, full_folder_path, is_default)
|
|
)
|
|
)
|
|
|
|
monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
|
|
monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths_module)
|
|
|
|
original_base_initialized = modal_cloud_module._COMFY_RUNTIME_BASE_INITIALIZED
|
|
original_custom_node_roots = set(modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS)
|
|
modal_cloud_module._COMFY_RUNTIME_BASE_INITIALIZED = False
|
|
modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.clear()
|
|
try:
|
|
custom_nodes_root = tmp_path / "custom_nodes"
|
|
custom_nodes_root.mkdir()
|
|
|
|
modal_cloud_module._ensure_comfy_runtime_initialized(None)
|
|
modal_cloud_module._ensure_comfy_runtime_initialized(custom_nodes_root)
|
|
modal_cloud_module._ensure_comfy_runtime_initialized(custom_nodes_root)
|
|
finally:
|
|
modal_cloud_module._COMFY_RUNTIME_BASE_INITIALIZED = original_base_initialized
|
|
modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.clear()
|
|
modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.update(original_custom_node_roots)
|
|
|
|
assert init_calls == [("extra", False, True), ("external",)]
|
|
assert folder_path_calls == [("custom_nodes", str(custom_nodes_root), True)]
|
|
|
|
|
|
class _BoundarySourceNode:
|
|
"""Simple source node used for subgraph execution tests."""
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("value",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "run"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
|
|
"""Return the minimal V1 input schema."""
|
|
return {"required": {"value": ("INT",)}}
|
|
|
|
def run(self, value: int) -> tuple[int]:
|
|
"""Increment the boundary input."""
|
|
return (value + 1,)
|
|
|
|
|
|
class _BoundarySinkNode:
|
|
"""Simple downstream node used for subgraph execution tests."""
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("value",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "run"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
|
|
"""Return the minimal V1 input schema."""
|
|
return {"required": {"value": ("INT",)}}
|
|
|
|
def run(self, value: int) -> tuple[int]:
|
|
"""Double the upstream value."""
|
|
return (value * 2,)
|
|
|
|
|
|
def test_local_remote_app_executes_subgraph_payload(
|
|
remote_modal_app_module: Any,
|
|
serialization_module: Any,
|
|
) -> None:
|
|
"""The local fallback remote app should execute rewritten subgraph payloads."""
|
|
payload = remote_modal_app_module.execute_subgraph_locally(
|
|
payload={
|
|
"payload_kind": "subgraph",
|
|
"component_id": "component-1",
|
|
"subgraph_prompt": {
|
|
"remote_1": {
|
|
"class_type": "BoundarySource",
|
|
"inputs": {"value": 0},
|
|
"_meta": {},
|
|
},
|
|
"remote_2": {
|
|
"class_type": "BoundarySink",
|
|
"inputs": {"value": ["remote_1", 0]},
|
|
"_meta": {},
|
|
},
|
|
},
|
|
"boundary_inputs": [
|
|
{
|
|
"proxy_input_name": "remote_input_0",
|
|
"targets": [{"node_id": "remote_1", "input_name": "value"}],
|
|
}
|
|
],
|
|
"boundary_outputs": [
|
|
{
|
|
"proxy_output_name": "remote_2_value",
|
|
"node_id": "remote_2",
|
|
"output_index": 0,
|
|
"io_type": "INT",
|
|
"is_list": False,
|
|
}
|
|
],
|
|
"execute_node_ids": ["remote_2"],
|
|
"extra_data": {},
|
|
"custom_nodes_bundle": None,
|
|
},
|
|
kwargs_payload='{"remote_input_0": 4}',
|
|
node_mapping={
|
|
"BoundarySource": _BoundarySourceNode,
|
|
"BoundarySink": _BoundarySinkNode,
|
|
},
|
|
)
|
|
outputs = serialization_module.deserialize_node_outputs(payload)
|
|
assert outputs == (10,)
|