497 lines
16 KiB
Python
497 lines
16 KiB
Python
"""Tests for dynamic Modal proxy nodes and local execution fallback."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import copy
|
|
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
|
|
from dataclasses import replace
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from decimal import Decimal
|
|
|
|
import hashlib
|
|
|
|
import importlib
|
|
|
|
import importlib.util
|
|
|
|
import json
|
|
|
|
import pickle
|
|
|
|
import subprocess
|
|
|
|
import sys
|
|
|
|
import threading
|
|
|
|
import time
|
|
|
|
import types
|
|
|
|
from contextlib import nullcontext
|
|
|
|
from io import BytesIO, StringIO
|
|
|
|
from pathlib import Path
|
|
|
|
import logging
|
|
|
|
from typing import Any, Iterator
|
|
|
|
import pytest
|
|
|
|
DEFAULT_TEST_DEPLOYMENT_APP_NAME = "comfy-modal-sync-gpu-rtx-pro-6000"
|
|
|
|
def _cloud_session_bridge_owner() -> Any:
|
|
"""Return the module that owns cloud session-bridge state and replay helpers."""
|
|
return importlib.import_module("cloud_session_bridge")
|
|
|
|
def _cloud_remote_session_store() -> Any:
|
|
"""Return the process-local session store from its extracted owner."""
|
|
return _cloud_session_bridge_owner()._REMOTE_SESSION_STORE
|
|
|
|
def _cloud_bridge_value_cache() -> dict[str, Any]:
|
|
"""Return the warm bridge-value cache from its extracted owner."""
|
|
return _cloud_session_bridge_owner()._REMOTE_SESSION_BRIDGE_VALUE_CACHE
|
|
|
|
def _cloud_bridge_value_cache_order() -> list[str]:
|
|
"""Return the warm bridge-value eviction order from its extracted owner."""
|
|
return _cloud_session_bridge_owner()._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER
|
|
|
|
def _patch_cloud_session_bridge(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
name: str,
|
|
value: Any,
|
|
) -> None:
|
|
"""Patch a session-bridge helper at its extracted owner."""
|
|
monkeypatch.setattr(_cloud_session_bridge_owner(), name, value)
|
|
|
|
def _cloud_comfy_bootstrap_owner() -> Any:
|
|
"""Return the module that owns cloud ComfyUI bootstrap state and helpers."""
|
|
return importlib.import_module("cloud_comfy_bootstrap")
|
|
|
|
def _cloud_node_output_cache_owner() -> Any:
|
|
"""Return the module that owns persisted cloud node-output cache behavior."""
|
|
return importlib.import_module("cloud_node_output_cache")
|
|
|
|
def _cloud_prompt_execution_owner() -> Any:
|
|
"""Return the module that owns cloud prompt execution state and helpers."""
|
|
return importlib.import_module("cloud_prompt_execution")
|
|
|
|
def _cloud_mapped_execution_owner() -> Any:
|
|
"""Return the module that owns cloud mapped execution helpers."""
|
|
return importlib.import_module("cloud_mapped_execution")
|
|
|
|
def _cloud_execution_control_owner() -> Any:
|
|
"""Return the module that owns active cloud execution registration."""
|
|
return importlib.import_module("cloud_execution_control")
|
|
|
|
def _cloud_streaming_owner() -> Any:
|
|
"""Return the module that owns cloud stream buffering and payload dispatch."""
|
|
return importlib.import_module("cloud_streaming")
|
|
|
|
def _cloud_volume_reload_owner() -> Any:
|
|
"""Return the module that owns cloud volume reload state and read-through."""
|
|
return importlib.import_module("cloud_volume_reload")
|
|
|
|
def _cloud_prewarm_owner() -> Any:
|
|
"""Return the module that owns cloud warm-container preparation state."""
|
|
return importlib.import_module("cloud_prewarm")
|
|
|
|
def _patch_cloud_storage_root(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
modal_cloud_module: Any,
|
|
storage_root: Path,
|
|
) -> None:
|
|
"""Patch storage settings across the cloud entrypoint and bootstrap owner."""
|
|
|
|
def settings() -> Any:
|
|
"""Return storage settings for one isolated cloud test."""
|
|
return types.SimpleNamespace(
|
|
remote_storage_root=str(storage_root),
|
|
local_storage_root=None,
|
|
)
|
|
|
|
monkeypatch.setattr(modal_cloud_module, "get_settings", settings)
|
|
monkeypatch.setattr(_cloud_comfy_bootstrap_owner(), "get_settings", settings)
|
|
monkeypatch.setattr(_cloud_volume_reload_owner(), "get_settings", settings)
|
|
|
|
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)
|
|
|
|
class _CloneableCacheValue:
|
|
"""Simple cloneable object used for loader cache tests."""
|
|
|
|
def __init__(self, value: str) -> None:
|
|
"""Store an identifying value for later clone assertions."""
|
|
self.value = value
|
|
|
|
def clone(self) -> "_CloneableCacheValue":
|
|
"""Return a fresh object carrying the same value."""
|
|
return _CloneableCacheValue(self.value)
|
|
|
|
class _FakeModelValue:
|
|
"""Simple cloneable stand-in for a non-transportable MODEL output."""
|
|
|
|
def __init__(self, value: str) -> None:
|
|
"""Store an identifying value for later assertions."""
|
|
self.value = value
|
|
|
|
def clone(self) -> "_FakeModelValue":
|
|
"""Return a fresh model value carrying the same identifier."""
|
|
return _FakeModelValue(self.value)
|
|
|
|
class _FakeModelLoaderNode:
|
|
"""Fake self-contained loader node that returns one MODEL-like object."""
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "load_checkpoint"
|
|
|
|
def load_checkpoint(self, ckpt_name: str) -> tuple[_FakeModelValue]:
|
|
"""Return a deterministic model value derived from the checkpoint name."""
|
|
return (_FakeModelValue(f"model::{ckpt_name}"),)
|
|
|
|
class _PromptMetadataSerializationNode:
|
|
"""Output node double that JSON-serializes ComfyUI's hidden PROMPT input."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("metadata",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "run"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, Any]]:
|
|
"""Declare one hydrated tensor input and the hidden prompt metadata input."""
|
|
return {
|
|
"required": {"image": ("IMAGE",)},
|
|
"hidden": {"prompt": "PROMPT"},
|
|
}
|
|
|
|
def run(self, image: Any, prompt: dict[str, Any]) -> tuple[str]:
|
|
"""Serialize the prompt exactly as metadata-writing output nodes do."""
|
|
del image
|
|
return (json.dumps(prompt, sort_keys=True),)
|
|
|
|
class _FakeClipLoaderNode:
|
|
"""Fake self-contained loader node that returns one CLIP-like object."""
|
|
|
|
RETURN_TYPES = ("CLIP",)
|
|
RETURN_NAMES = ("clip",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "load_clip"
|
|
|
|
def load_clip(
|
|
self,
|
|
clip_name: str,
|
|
type: str = "stable_diffusion",
|
|
device: str = "default",
|
|
) -> tuple[str]:
|
|
"""Return a deterministic CLIP value derived from the loader inputs."""
|
|
return (f"clip::{clip_name}:{type}:{device}",)
|
|
|
|
class _FakeSessionValueNode:
|
|
"""Fake node that produces one remote-only STRING value."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self) -> tuple[str]:
|
|
"""Return a deterministic value that can be stored in session state."""
|
|
return ("shared-session-value",)
|
|
|
|
class _FakeSessionEchoNode:
|
|
"""Fake node that echoes a STRING input back to the caller."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, text: str) -> tuple[str]:
|
|
"""Return the supplied input unchanged for session-ref resolution tests."""
|
|
return (text,)
|
|
|
|
class _FakeRewriteRemoteModelNode:
|
|
"""Fake rewrite-time node that produces a non-transportable MODEL output."""
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteRemoteSamplerNode:
|
|
"""Fake rewrite-time node that produces a transportable LATENT output."""
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteLatentSourceNode:
|
|
"""Fake local source used to feed LATENT values into remote proxies."""
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteModalMapInputNode:
|
|
"""Fake rewrite-time Modal map marker node."""
|
|
|
|
RETURN_TYPES = ("*",)
|
|
RETURN_NAMES = ("value",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteLocalSinkNode:
|
|
"""Fake local sink used to model downstream local work in rewrite tests."""
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteLocalFeedbackNode:
|
|
"""Fake local node that turns a remote latent into a transportable remote input."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteRemoteImageNode:
|
|
"""Fake remote node that produces one locally previewable image."""
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewriteRemoteTextNode:
|
|
"""Fake remote node that consumes an image and produces text."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
|
|
class _FakeRewritePreviewImageNode:
|
|
"""Fake local PreviewImage output node."""
|
|
|
|
RETURN_TYPES: tuple[str, ...] = ()
|
|
RETURN_NAMES: tuple[str, ...] = ()
|
|
OUTPUT_IS_LIST: tuple[bool, ...] = ()
|
|
OUTPUT_NODE = True
|
|
|
|
def _current_remote_runtime_payload(remote_modal_app_module: Any) -> dict[str, Any]:
|
|
"""Return version metadata for a compatible remote-engine test double."""
|
|
return {
|
|
"protocol_version": remote_modal_app_module._REMOTE_APP_PROTOCOL_VERSION,
|
|
"runtime_fingerprint": remote_modal_app_module._expected_remote_runtime_fingerprint(),
|
|
}
|
|
|
|
class _FakeImplicitBatchKSamplerNode:
|
|
"""Fake sampler node exposing ComfyUI-style LATENT and primitive socket types."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
|
|
"""Return the minimal socket schema needed for implicit batch target inspection."""
|
|
return {
|
|
"required": {
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT",),
|
|
"positive": ("CONDITIONING",),
|
|
}
|
|
}
|
|
|
|
def _assert_node_module_identity(node_class: type[Any], expected_module: str) -> None:
|
|
"""Assert a registered V3 node exposes its loader-assigned module identity."""
|
|
assert node_class.RELATIVE_PYTHON_MODULE == expected_module
|
|
assert node_class.GET_NODE_INFO_V1()["python_module"] == expected_module
|
|
|
|
class _PersistentCacheNode:
|
|
"""Simple node used to verify persisted node-cache reuse across prompt runs."""
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("value",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "run"
|
|
invocation_count = 0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
|
|
"""Return the minimal V1 schema needed for cache-key generation."""
|
|
return {"required": {"value": ("INT",)}}
|
|
|
|
def run(self, value: int) -> tuple[int]:
|
|
"""Count real executions so persisted cache hits are visible to the test."""
|
|
type(self).invocation_count += 1
|
|
return (value + 1,)
|
|
|
|
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,)
|
|
|
|
class _PrimitiveEchoNode:
|
|
"""Simple node used to verify primitive widget coercion."""
|
|
|
|
RETURN_TYPES = ("INT", "FLOAT", "BOOLEAN", "STRING")
|
|
RETURN_NAMES = ("steps", "cfg", "enabled", "label")
|
|
OUTPUT_IS_LIST = (False, False, False, False)
|
|
FUNCTION = "run"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
|
|
"""Return one primitive input of each V1 widget type."""
|
|
return {
|
|
"required": {
|
|
"steps": ("INT",),
|
|
"cfg": ("FLOAT",),
|
|
"enabled": ("BOOLEAN",),
|
|
"label": ("STRING",),
|
|
}
|
|
}
|
|
|
|
def run(
|
|
self,
|
|
steps: int,
|
|
cfg: float,
|
|
enabled: bool,
|
|
label: str,
|
|
) -> tuple[int, float, bool, str]:
|
|
"""Echo primitive inputs after asserting their coerced Python types."""
|
|
assert isinstance(steps, int)
|
|
assert isinstance(cfg, float)
|
|
assert isinstance(enabled, bool)
|
|
assert isinstance(label, str)
|
|
return (steps, cfg, enabled, label)
|
|
|
|
def _v3_batch_images_node_class() -> type[Any]:
|
|
"""Return a dependency-light V3 Batch Images node with Autogrow inputs."""
|
|
class _V3BatchImagesNode:
|
|
"""Minimal V3 node matching ComfyUI's Batch Images input schema."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[Any, ...]]]:
|
|
"""Return the raw V3 Autogrow schema emitted by ComfyUI."""
|
|
return {
|
|
"required": {
|
|
"images": (
|
|
"COMFY_AUTOGROW_V3",
|
|
{
|
|
"template": {
|
|
"input": {"required": {"image": ("IMAGE", {})}},
|
|
"prefix": "image",
|
|
"min": 1,
|
|
"max": 50,
|
|
}
|
|
},
|
|
)
|
|
}
|
|
}
|
|
|
|
return _V3BatchImagesNode
|
|
|
|
def _install_fake_v3_input_finalizer(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Install a narrow ComfyUI V3 finalizer double for Autogrow unit tests."""
|
|
comfy_api_module = types.ModuleType("comfy_api")
|
|
comfy_api_module.__path__ = []
|
|
comfy_api_latest_module = types.ModuleType("comfy_api.latest")
|
|
|
|
def get_finalized_class_inputs(
|
|
raw_input_types: dict[str, Any],
|
|
live_inputs: dict[str, Any],
|
|
) -> tuple[dict[str, Any], None, dict[str, Any]]:
|
|
"""Expand the test Autogrow schema using ComfyUI's public prompt shape."""
|
|
_, config = raw_input_types["required"]["images"]
|
|
template = config["template"]
|
|
template_input = template["input"]["required"]["image"]
|
|
names = [f"{template['prefix']}{index}" for index in range(template["max"])]
|
|
finalized: dict[str, dict[str, Any]] = {"required": {}, "optional": {}}
|
|
dynamic_paths: dict[str, str] = {}
|
|
for index, name in enumerate(names):
|
|
expanded_name = f"images.{name}"
|
|
section_name = "required" if index < template["min"] else "optional"
|
|
finalized[section_name][expanded_name] = template_input
|
|
if expanded_name in live_inputs:
|
|
dynamic_paths[expanded_name] = expanded_name
|
|
return finalized, None, {"dynamic_paths": dynamic_paths}
|
|
|
|
comfy_api_latest_module._io = types.SimpleNamespace(
|
|
get_finalized_class_inputs=get_finalized_class_inputs
|
|
)
|
|
comfy_api_module.latest = comfy_api_latest_module
|
|
monkeypatch.setitem(sys.modules, "comfy_api", comfy_api_module)
|
|
monkeypatch.setitem(sys.modules, "comfy_api.latest", comfy_api_latest_module)
|
|
|
|
class _ImplicitBatchListSourceNode:
|
|
"""Fake node that consumes a whole list once and emits scalar and list outputs."""
|
|
|
|
RETURN_TYPES = ("INT", "INT")
|
|
RETURN_NAMES = ("first_value", "all_values")
|
|
OUTPUT_IS_LIST = (False, True)
|
|
INPUT_IS_LIST = True
|
|
FUNCTION = "run"
|
|
|
|
class _ImplicitBatchScalarConsumerNode:
|
|
"""Fake scalar consumer used for outer implicit batch regressions."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
OUTPUT_IS_LIST = (False,)
|
|
FUNCTION = "run"
|
|
|
|
class _ImplicitBatchListConsumerNode:
|
|
"""Fake list-aware consumer used for outer implicit batch regressions."""
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
INPUT_IS_LIST = True
|
|
FUNCTION = "run"
|
|
|
|
|
|
__all__ = tuple(name for name in globals() if not name.startswith("__"))
|