nodes/ttulttul-ComfyUI-Modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
9dcccd9e3e24bb89f576615f1d5d60862ac4bd38
ttulttul-ComfyUI-Modal/tests/test_cloud_cache_bootstrap.py
T

1996 lines
74 KiB
Python
Raw Blame History

"""Tests split from the Modal executor integration suite."""
from __future__ import annotations
from modal_executor_test_support import * # noqa: F401,F403
def test_dynamic_proxy_registration_sets_module_identity_on_new_and_cached_class(
modal_executor_module: Any,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Ordinary proxies should serialize a valid module after creation and cache reuse."""
expected_module = "custom_nodes.ComfyUI-Modal"
monkeypatch.setattr(
modal_executor_module.ModalUniversalExecutor,
"RELATIVE_PYTHON_MODULE",
expected_module,
)
fake_nodes_module = types.SimpleNamespace(
NODE_CLASS_MAPPINGS={"OriginalNodeForModuleIdentity": _FakeOriginalNode},
NODE_DISPLAY_NAME_MAPPINGS={},
)
proxy_node_id = modal_executor_module.ensure_modal_proxy_node_registered(
original_class_type="OriginalNodeForModuleIdentity",
original_class=_FakeOriginalNode,
nodes_module=fake_nodes_module,
)
proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
_assert_node_module_identity(proxy_class, expected_module)
proxy_class.RELATIVE_PYTHON_MODULE = None
fake_nodes_module.NODE_CLASS_MAPPINGS.clear()
fake_nodes_module.NODE_DISPLAY_NAME_MAPPINGS.clear()
cached_proxy_node_id = modal_executor_module.ensure_modal_proxy_node_registered(
original_class_type="OriginalNodeForModuleIdentity",
original_class=_FakeOriginalNode,
nodes_module=fake_nodes_module,
)
assert cached_proxy_node_id == proxy_node_id
assert fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] is proxy_class
_assert_node_module_identity(proxy_class, expected_module)
def test_component_proxy_registration_sets_module_identity_on_new_and_cached_class(
modal_executor_module: Any,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Component proxies should serialize a valid module after creation and cache reuse."""
expected_module = "custom_nodes.ComfyUI-Modal"
monkeypatch.setattr(
modal_executor_module.ModalUniversalExecutor,
"RELATIVE_PYTHON_MODULE",
expected_module,
)
fake_nodes_module = types.SimpleNamespace(
NODE_CLASS_MAPPINGS={},
NODE_DISPLAY_NAME_MAPPINGS={},
)
registration_kwargs = {
"output_types": ("MODULE_IDENTITY_COMPONENT",),
"output_names": ("value",),
"output_is_list": (False,),
"nodes_module": fake_nodes_module,
"is_output_node": False,
}
proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
**registration_kwargs
)
proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
_assert_node_module_identity(proxy_class, expected_module)
proxy_class.RELATIVE_PYTHON_MODULE = None
fake_nodes_module.NODE_CLASS_MAPPINGS.clear()
fake_nodes_module.NODE_DISPLAY_NAME_MAPPINGS.clear()
cached_proxy_node_id = (
modal_executor_module.ensure_modal_component_proxy_node_registered(
**registration_kwargs
)
)
assert cached_proxy_node_id == proxy_node_id
assert fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] is proxy_class
_assert_node_module_identity(proxy_class, expected_module)
def test_component_proxy_cache_distinguishes_scheduler_list_outputs(
modal_executor_module: Any,
) -> None:
"""Mapped list proxies must not reuse a cached scalar-output proxy class."""
fake_nodes_module = types.SimpleNamespace(
NODE_CLASS_MAPPINGS={},
NODE_DISPLAY_NAME_MAPPINGS={},
)
scalar_proxy_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
output_types=("LATENT",),
output_names=("samples",),
output_is_list=(False,),
nodes_module=fake_nodes_module,
is_output_node=False,
)
list_proxy_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
output_types=("LATENT",),
output_names=("samples",),
output_is_list=(True,),
nodes_module=fake_nodes_module,
is_output_node=False,
)
assert scalar_proxy_id != list_proxy_id
assert fake_nodes_module.NODE_CLASS_MAPPINGS[scalar_proxy_id].OUTPUT_IS_LIST == [False]
assert fake_nodes_module.NODE_CLASS_MAPPINGS[list_proxy_id].OUTPUT_IS_LIST == [True]
def test_cache_friendly_proxy_payload_ignores_volatile_queue_metadata(
modal_executor_module: Any,
proxy_payloads_module: Any,
) -> None:
"""Identical proxy work should sanitize to one local cache surface across prompt runs."""
first_payload = modal_executor_module.register_cache_friendly_proxy_payload(
"node-9",
{
"payload_kind": "subgraph",
"component_id": "component-9",
"prompt_id": "prompt-1",
"boundary_outputs": [],
"execute_node_ids": ["12"],
"extra_data": {
"client_id": "client-1",
"create_time": 1000,
"modal": {"remote_component_ids": ["12"]},
},
"requires_volume_reload": True,
"volume_reload_marker": "marker-1",
"uploaded_volume_paths": ["/storage/assets/a.bin"],
},
)
second_payload = modal_executor_module.register_cache_friendly_proxy_payload(
"node-9",
{
"payload_kind": "subgraph",
"component_id": "component-9",
"prompt_id": "prompt-2",
"boundary_outputs": [],
"execute_node_ids": ["12"],
"extra_data": {
"client_id": "client-1",
"create_time": 2000,
"modal": {"remote_component_ids": ["12", "39"]},
},
"requires_volume_reload": False,
"volume_reload_marker": None,
"uploaded_volume_paths": [],
},
)
assert first_payload == second_payload == {
"payload_kind": "subgraph",
"component_id": "component-9",
"boundary_outputs": [],
"execute_node_ids": ["12"],
proxy_payloads_module._PROXY_CACHE_CONTEXT_ID_KEY: "node-9",
}
assert modal_executor_module._rehydrate_proxy_payload(
first_payload,
unique_id="node-9",
prompt_id="prompt-1",
) == {
"payload_kind": "subgraph",
"component_id": "component-9",
"prompt_id": "prompt-1",
"boundary_outputs": [],
"execute_node_ids": ["12"],
"extra_data": {
"client_id": "client-1",
"create_time": 1000,
"modal": {"remote_component_ids": ["12"]},
},
"requires_volume_reload": True,
"volume_reload_marker": "marker-1",
"uploaded_volume_paths": ["/storage/assets/a.bin"],
}
assert modal_executor_module._rehydrate_proxy_payload(
second_payload,
unique_id="node-9",
prompt_id="prompt-2",
) == {
"payload_kind": "subgraph",
"component_id": "component-9",
"prompt_id": "prompt-2",
"boundary_outputs": [],
"execute_node_ids": ["12"],
"extra_data": {
"client_id": "client-1",
"create_time": 2000,
"modal": {"remote_component_ids": ["12", "39"]},
},
"requires_volume_reload": False,
"volume_reload_marker": None,
"uploaded_volume_paths": [],
}
def test_modal_cloud_reuses_extracted_custom_nodes_bundle(
modal_cloud_module: Any,
monkeypatch: Any,
tmp_path: Path,
) -> None:
"""The remote runtime should avoid re-extracting an unchanged custom_nodes bundle."""
storage_root = tmp_path / "storage"
bundle_path = storage_root / "custom_nodes" / "bundle.zip"
bundle_path.parent.mkdir(parents=True)
import zipfile
with zipfile.ZipFile(bundle_path, "w") as archive:
archive.writestr("example/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(storage_root))
modal_cloud_module.get_settings.cache_clear()
bootstrap_owner = _cloud_comfy_bootstrap_owner()
original_cache = dict(bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES)
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
try:
first_root = modal_cloud_module._extract_custom_nodes_bundle("/custom_nodes/bundle.zip")
monkeypatch.setattr(
bootstrap_owner.zipfile,
"ZipFile",
lambda *args, **kwargs: (_ for _ in ()).throw(
AssertionError("Expected cached extraction root to be reused.")
),
)
second_root = modal_cloud_module._extract_custom_nodes_bundle("/custom_nodes/bundle.zip")
finally:
modal_cloud_module.get_settings.cache_clear()
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.update(original_cache)
assert first_root is not None
assert second_root == first_root
def test_modal_cloud_extracts_custom_nodes_manifest_with_multiple_archives(
modal_cloud_module: Any,
monkeypatch: Any,
tmp_path: Path,
) -> None:
"""The remote runtime should resolve manifest-based custom_nodes bundles into several archives."""
storage_root = tmp_path / "storage"
entry_a_path = storage_root / "custom_nodes" / "entries" / "example_a" / "hash_a_bundle.zip"
entry_b_path = storage_root / "custom_nodes" / "entries" / "example_b" / "hash_b_bundle.zip"
manifest_path = storage_root / "custom_nodes" / "manifests" / "bundle_manifest.json"
entry_a_path.parent.mkdir(parents=True, exist_ok=True)
entry_b_path.parent.mkdir(parents=True, exist_ok=True)
manifest_path.parent.mkdir(parents=True, exist_ok=True)
import zipfile
with zipfile.ZipFile(entry_a_path, "w") as archive:
archive.writestr("example_a/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
with zipfile.ZipFile(entry_b_path, "w") as archive:
archive.writestr("example_b/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
manifest_path.write_text(
json.dumps(
{
"version": 1,
"bundle_sha256": "bundle-hash",
"entries": [
{
"entry_name": "example_a",
"display_name": "example_a",
"sha256": "hash-a",
"remote_path": "/custom_nodes/entries/example_a/hash_a_bundle.zip",
},
{
"entry_name": "example_b",
"display_name": "example_b",
"sha256": "hash-b",
"remote_path": "/custom_nodes/entries/example_b/hash_b_bundle.zip",
},
],
},
sort_keys=True,
),
encoding="utf-8",
)
monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(storage_root))
modal_cloud_module.get_settings.cache_clear()
bootstrap_owner = _cloud_comfy_bootstrap_owner()
original_cache = dict(bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES)
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
try:
extraction_root = modal_cloud_module._extract_custom_nodes_bundle(
"/custom_nodes/manifests/bundle_manifest.json"
)
finally:
modal_cloud_module.get_settings.cache_clear()
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.update(original_cache)
assert extraction_root is not None
assert (extraction_root / "example_a" / "__init__.py").exists()
assert (extraction_root / "example_b" / "__init__.py").exists()
def test_modal_cloud_links_version_two_custom_node_assets(
modal_cloud_module: Any,
monkeypatch: Any,
tmp_path: Path,
) -> None:
"""Version-two manifests should keep model assets on mounted storage."""
storage_root = tmp_path / "storage"
asset_sha256 = "a" * 64
archive_path = storage_root / "custom_nodes/entries/example/code_bundle.zip"
asset_path = storage_root / f"custom_nodes/assets/example/{asset_sha256}_model.pth"
manifest_path = storage_root / "custom_nodes/manifests/bundle_manifest_v2.json"
archive_path.parent.mkdir(parents=True)
asset_path.parent.mkdir(parents=True)
manifest_path.parent.mkdir(parents=True)
asset_path.write_bytes(b"mounted-model")
import zipfile
with zipfile.ZipFile(archive_path, "w") as archive:
archive.writestr("example/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
manifest_path.write_text(
json.dumps(
{
"version": 2,
"bundle_sha256": "bundle-hash",
"entries": [
{
"entry_name": "example",
"display_name": "example",
"sha256": "code-hash",
"remote_path": "/custom_nodes/entries/example/code_bundle.zip",
"assets": [
{
"relative_path": "example/checkpoints/model.pth",
"sha256": asset_sha256,
"size_bytes": len(b"mounted-model"),
"remote_path": (
f"/custom_nodes/assets/example/{asset_sha256}_model.pth"
),
}
],
}
],
},
sort_keys=True,
),
encoding="utf-8",
)
monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(storage_root))
modal_cloud_module.get_settings.cache_clear()
bootstrap_owner = _cloud_comfy_bootstrap_owner()
original_cache = dict(bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES)
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
try:
extraction_root = modal_cloud_module._extract_custom_nodes_bundle(
"/custom_nodes/manifests/bundle_manifest_v2.json"
)
finally:
modal_cloud_module.get_settings.cache_clear()
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
bootstrap_owner._EXTRACTED_CUSTOM_NODE_BUNDLES.update(original_cache)
assert extraction_root is not None
linked_asset = extraction_root / "example/checkpoints/model.pth"
assert linked_asset.is_symlink()
assert linked_asset.read_bytes() == b"mounted-model"
def test_local_fallback_links_version_two_custom_node_assets(
remote_modal_app_module: Any,
tmp_path: Path,
) -> None:
"""The non-Modal extractor should preserve version-two package asset links."""
storage_root = tmp_path / "storage"
extraction_root = tmp_path / "extracted"
asset_sha256 = "b" * 64
asset_path = storage_root / f"custom_nodes/assets/example/{asset_sha256}_model.pth"
manifest_path = storage_root / "custom_nodes/manifests/bundle_manifest_v2.json"
asset_path.parent.mkdir(parents=True)
manifest_path.parent.mkdir(parents=True)
extraction_root.mkdir()
asset_path.write_bytes(b"local-mounted-model")
manifest_path.write_text(
json.dumps(
{
"version": 2,
"entries": [
{
"entry_name": "example",
"assets": [
{
"relative_path": "example/checkpoints/model.pth",
"sha256": asset_sha256,
"size_bytes": len(b"local-mounted-model"),
"remote_path": (
f"/custom_nodes/assets/example/{asset_sha256}_model.pth"
),
}
],
}
],
},
sort_keys=True,
),
encoding="utf-8",
)
remote_modal_app_module._materialize_local_custom_node_assets(
manifest_path,
storage_root,
extraction_root,
)
linked_asset = extraction_root / "example/checkpoints/model.pth"
assert linked_asset.is_symlink()
assert linked_asset.read_bytes() == b"local-mounted-model"
def test_modal_cloud_readthrough_hydrates_custom_node_manifest_dependencies(
modal_cloud_module: Any,
monkeypatch: Any,
tmp_path: Path,
) -> None:
"""Read-through hydration should include archives and assets named by a manifest."""
manifest_payload = json.dumps(
{
"version": 2,
"entries": [
{
"remote_path": "/custom_nodes/packages/hash_package.zip",
"assets": [
{"remote_path": "/custom_nodes/assets/hash_model.safetensors"}
],
}
],
}
).encode()
committed_files = {
"custom_nodes/manifests/hash_manifest.json": manifest_payload,
"custom_nodes/packages/hash_package.zip": b"package",
"custom_nodes/assets/hash_model.safetensors": b"model",
}
class FakeVolume:
"""Modal Volume double backed by committed in-memory files."""
def read_file(self, path: str) -> Iterator[bytes]:
"""Yield the requested committed file."""
yield committed_files[path]
storage_root = tmp_path / "storage"
storage_root.mkdir()
readthrough_root = tmp_path / "readthrough"
_patch_cloud_storage_root(monkeypatch, modal_cloud_module, storage_root)
bootstrap_owner = _cloud_comfy_bootstrap_owner()
monkeypatch.setattr(
bootstrap_owner,
"_REMOTE_VOLUME_READTHROUGH_ROOT",
readthrough_root,
)
hydrated_paths = modal_cloud_module._hydrate_missing_payload_volume_paths(
FakeVolume(),
{"custom_nodes_bundle": "/custom_nodes/manifests/hash_manifest.json"},
)
assert hydrated_paths == [
readthrough_root / "custom_nodes" / "manifests" / "hash_manifest.json",
readthrough_root / "custom_nodes" / "assets" / "hash_model.safetensors",
readthrough_root / "custom_nodes" / "packages" / "hash_package.zip",
]
def test_modal_cloud_clears_each_owner_cache_before_volume_reload(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Volume preparation should clear every cache that can retain mounted files."""
cleared: list[str] = []
volume_reload_owner = _cloud_volume_reload_owner()
monkeypatch.setattr(
volume_reload_owner,
"clear_cloud_prompt_execution_warm_caches",
lambda: cleared.append("prompt-execution"),
)
monkeypatch.setattr(
volume_reload_owner,
"clear_cloud_session_bridge_warm_caches",
lambda: cleared.append("session-bridge"),
)
monkeypatch.setattr(
volume_reload_owner,
"clear_comfy_bootstrap_warm_caches",
lambda: cleared.append("comfy-bootstrap"),
)
modal_cloud_module._clear_warm_remote_caches()
assert cleared == ["prompt-execution", "session-bridge", "comfy-bootstrap"]
def test_modal_cloud_loader_cache_reuses_and_clones_outputs(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Cached loader wrappers should avoid repeated loads and clone cached outputs on hits."""
class FakeLoader:
"""Simple loader with one expensive method."""
def __init__(self) -> None:
"""Initialize call counter state."""
self.calls = 0
def load(self, model_name: str, device: str = "default") -> tuple[_CloneableCacheValue]:
"""Return a cloneable payload while counting underlying loads."""
self.calls += 1
return (_CloneableCacheValue(f"{model_name}:{device}:{self.calls}"),)
bootstrap_owner = _cloud_comfy_bootstrap_owner()
original_cache = dict(bootstrap_owner._LOADER_OUTPUT_CACHE)
original_wrapped = set(bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES)
original_metrics = dict(bootstrap_owner._LOADER_CACHE_METRICS)
bootstrap_owner._LOADER_OUTPUT_CACHE.clear()
bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES.clear()
bootstrap_owner._LOADER_CACHE_METRICS.clear()
bootstrap_owner._LOADER_CACHE_METRICS.update({"hit": 0, "miss": 0})
try:
modal_cloud_module._wrap_loader_method_with_cache(
"FakeLoader",
FakeLoader,
"load",
lambda kwargs: modal_cloud_module._serialize_loader_cache_key(kwargs),
)
loader = FakeLoader()
first = loader.load("model.safetensors", device="cpu")[0]
second = loader.load("model.safetensors", device="cpu")[0]
third = loader.load("other.safetensors", device="cpu")[0]
metrics_snapshot = modal_cloud_module._loader_cache_metric_snapshot()
finally:
bootstrap_owner._LOADER_OUTPUT_CACHE.clear()
bootstrap_owner._LOADER_OUTPUT_CACHE.update(original_cache)
bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES.clear()
bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES.update(original_wrapped)
bootstrap_owner._LOADER_CACHE_METRICS.clear()
bootstrap_owner._LOADER_CACHE_METRICS.update(original_metrics)
assert loader.calls == 2
assert first.value == "model.safetensors:cpu:1"
assert second.value == "model.safetensors:cpu:1"
assert third.value == "other.safetensors:cpu:2"
assert first is not second
assert metrics_snapshot == {"hit": 1, "miss": 2}
def test_modal_cloud_installs_loader_cache_wrappers_for_builtin_loaders(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""The runtime should patch the heavyweight built-in loaders once they are available."""
fake_nodes_module = types.SimpleNamespace(
NODE_CLASS_MAPPINGS={
"CheckpointLoader": type("CheckpointLoader", (), {"load_checkpoint": lambda self, config_name, ckpt_name: (config_name, ckpt_name)}),
"CheckpointLoaderSimple": type("CheckpointLoaderSimple", (), {"load_checkpoint": lambda self, ckpt_name: (ckpt_name,)}),
"UNETLoader": type("UNETLoader", (), {"load_unet": lambda self, unet_name, weight_dtype="default": (unet_name,)}),
"CLIPLoader": type("CLIPLoader", (), {"load_clip": lambda self, clip_name, type="stable_diffusion", device="default": (clip_name,)}),
"DualCLIPLoader": type("DualCLIPLoader", (), {"load_clip": lambda self, clip_name1, clip_name2, type, device="default": (clip_name1, clip_name2)}),
"VAELoader": type("VAELoader", (), {"load_vae": lambda self, vae_name: (vae_name,)}),
"unCLIPCheckpointLoader": type("unCLIPCheckpointLoader", (), {"load_checkpoint": lambda self, ckpt_name, output_vae=True, output_clip=True: (ckpt_name,)}),
"ImageOnlyCheckpointLoader": type("ImageOnlyCheckpointLoader", (), {"load_checkpoint": lambda self, ckpt_name, output_vae=True, output_clip=True: (ckpt_name,)}),
}
)
bootstrap_owner = _cloud_comfy_bootstrap_owner()
original_wrapped = set(bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES)
bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES.clear()
monkeypatch.setattr(bootstrap_owner, "_load_nodes_module", lambda: fake_nodes_module)
try:
modal_cloud_module._install_loader_cache_wrappers()
installed_wrappers = set(bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES)
finally:
bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES.clear()
bootstrap_owner._LOADER_CACHE_WRAPPED_CLASSES.update(original_wrapped)
assert {
"CheckpointLoader",
"CheckpointLoaderSimple",
"UNETLoader",
"CLIPLoader",
"DualCLIPLoader",
"VAELoader",
"unCLIPCheckpointLoader",
"ImageOnlyCheckpointLoader",
} <= installed_wrappers
def test_modal_cloud_node_cache_key_hashes_boundary_tensors(
modal_cloud_module: Any,
) -> None:
"""Boundary tensors inside ComfyUI cache signatures should produce stable cache keys."""
torch = pytest.importorskip("torch")
first_tensor = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3)
same_value_tensor = first_tensor.clone()
different_tensor = first_tensor + 1
signature = frozenset(
{
(
12,
frozenset(
{
(
4,
frozenset(
{
(0, "latent_image"),
(1, frozenset({("samples", first_tensor)})),
}
),
)
}
),
)
}
)
same_signature = frozenset(
{
(
12,
frozenset(
{
(
4,
frozenset(
{
(0, "latent_image"),
(1, frozenset({("samples", same_value_tensor)})),
}
),
)
}
),
)
}
)
different_signature = frozenset(
{
(
12,
frozenset(
{
(
4,
frozenset(
{
(0, "latent_image"),
(1, frozenset({("samples", different_tensor)})),
}
),
)
}
),
)
}
)
first_key = modal_cloud_module._node_output_cache_key(signature)
second_key = modal_cloud_module._node_output_cache_key(same_signature)
different_key = modal_cloud_module._node_output_cache_key(different_signature)
assert isinstance(first_key, str)
assert first_key.startswith("NC_")
assert second_key == first_key
assert different_key != first_key
def test_modal_cloud_node_cache_key_rebuilds_input_signature_before_unhashable_conversion(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Distributed node caching should bypass ComfyUI's precomputed `Unhashable` data key."""
torch = pytest.importorskip("torch")
class FakeDynPrompt:
"""Minimal dynamic-prompt stub for input-signature reconstruction."""
def __init__(self, node: dict[str, Any]) -> None:
"""Store one node payload under id `12`."""
self._node = node
def has_node(self, node_id: str) -> bool:
"""Return whether the requested node exists."""
return str(node_id) == "12"
def get_node(self, node_id: str) -> dict[str, Any]:
"""Return the stored node payload."""
if not self.has_node(node_id):
raise KeyError(node_id)
return self._node
class FakeUnhashable:
"""Stand-in for ComfyUI's `Unhashable` marker."""
tensor = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3)
dynprompt = FakeDynPrompt(
{
"class_type": "FakeSampler",
"inputs": {
"latent_image": {"samples": tensor},
"steps": 18,
},
}
)
cache_key_set = types.SimpleNamespace(
dynprompt=dynprompt,
is_changed_cache=types.SimpleNamespace(is_changed={"12": False}),
get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
include_node_id_in_input=lambda: False,
get_data_key=lambda node_id: frozenset({("latent_image", FakeUnhashable())}),
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_load_nodes_module",
lambda: types.SimpleNamespace(
NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})}
),
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_include_unique_id_in_input_signature",
lambda class_type: False,
)
rebuilt_key = modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "12")
direct_bad_key = modal_cloud_module._node_output_cache_key(cache_key_set.get_data_key("12"))
assert isinstance(rebuilt_key, str)
assert rebuilt_key.startswith("NC_")
assert direct_bad_key is None
def test_modal_cloud_node_cache_key_uses_boundary_source_signature_for_unhashable_inputs(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Boundary-fed runtime objects should hash by stable source provenance instead of object identity."""
class FakeDynPrompt:
"""Minimal dynamic prompt wrapper backed by one mutable prompt dict."""
def __init__(self, prompt: dict[str, Any]) -> None:
"""Store the prompt used by the cache-key rebuild."""
self._prompt = prompt
def has_node(self, node_id: str) -> bool:
"""Return whether the requested node exists in the prompt."""
return str(node_id) in self._prompt
def get_node(self, node_id: str) -> dict[str, Any]:
"""Return the stored node payload."""
return self._prompt[str(node_id)]
class FakeModelPatcher:
"""Stand-in for ComfyUI's unhashable `ModelPatcher` runtime object."""
prompt_one = {
"39": {
"class_type": "FakeSampler",
"inputs": {},
}
}
prompt_two = {
"39": {
"class_type": "FakeSampler",
"inputs": {},
}
}
prompt_three = {
"39": {
"class_type": "FakeSampler",
"inputs": {},
}
}
boundary_spec_one = [
{
"proxy_input_name": "remote_input_0",
"io_type": "MODEL",
"source_signature": "SRC_same_model",
"targets": [{"node_id": "39", "input_name": "model"}],
}
]
boundary_spec_three = [
{
"proxy_input_name": "remote_input_0",
"io_type": "MODEL",
"source_signature": "SRC_other_model",
"targets": [{"node_id": "39", "input_name": "model"}],
}
]
modal_cloud_module._apply_boundary_inputs(
prompt_one,
boundary_spec_one,
{"remote_input_0": FakeModelPatcher()},
)
modal_cloud_module._apply_boundary_inputs(
prompt_two,
boundary_spec_one,
{"remote_input_0": FakeModelPatcher()},
)
modal_cloud_module._apply_boundary_inputs(
prompt_three,
boundary_spec_three,
{"remote_input_0": FakeModelPatcher()},
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_load_nodes_module",
lambda: types.SimpleNamespace(
NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})}
),
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_include_unique_id_in_input_signature",
lambda class_type: False,
)
def cache_key_for(prompt: dict[str, Any]) -> str | None:
"""Build one distributed cache key for the prepared prompt."""
cache_key_set = types.SimpleNamespace(
dynprompt=FakeDynPrompt(prompt),
is_changed_cache=types.SimpleNamespace(is_changed={"39": False}),
get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
include_node_id_in_input=lambda: False,
get_data_key=lambda node_id: None,
)
return modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "39")
first_key = cache_key_for(prompt_one)
second_key = cache_key_for(prompt_two)
different_key = cache_key_for(prompt_three)
assert isinstance(first_key, str)
assert first_key.startswith("NC_")
assert second_key == first_key
assert different_key != first_key
def test_modal_cloud_node_cache_key_treats_nested_two_item_lists_as_data(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Nested conditioning-style two-item lists should not be mistaken for prompt links."""
class FakeDynPrompt:
"""Minimal dynamic prompt wrapper for one data-only node."""
def has_node(self, node_id: str) -> bool:
"""Return whether the requested node exists."""
return str(node_id) == "545"
def get_node(self, node_id: str) -> dict[str, Any]:
"""Return a node with nested list data shaped like Comfy conditioning."""
if not self.has_node(node_id):
raise KeyError(node_id)
return {
"class_type": "AnythingToMarkdown",
"inputs": {
"anything": [
[
{"pooled_output": "summary"},
[["not-a-node-id"], ["not-an-output-index"]],
]
]
},
}
cache_key_set = types.SimpleNamespace(
dynprompt=FakeDynPrompt(),
is_changed_cache=types.SimpleNamespace(is_changed={"545": False}),
get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
include_node_id_in_input=lambda: False,
get_data_key=lambda node_id: None,
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_load_nodes_module",
lambda: types.SimpleNamespace(
NODE_CLASS_MAPPINGS={"AnythingToMarkdown": type("AnythingToMarkdown", (), {})}
),
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_include_unique_id_in_input_signature",
lambda class_type: False,
)
cache_key = modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "545")
assert isinstance(cache_key, str)
assert cache_key.startswith("NC_")
def test_modal_cloud_collects_custom_node_runtime_packages(
modal_cloud_module: Any,
tmp_path: Path,
) -> None:
"""The Modal image should install requirements declared by bundled custom nodes."""
custom_nodes_dir = tmp_path / "custom_nodes"
first_package = custom_nodes_dir / "first"
second_package = custom_nodes_dir / "second"
first_package.mkdir(parents=True)
second_package.mkdir(parents=True)
(first_package / "requirements.txt").write_text(
"\n".join(
[
"# comments are ignored",
"omegaconf==2.3.0",
"hydra-core==1.3.2 # inline comment",
"--extra-index-url https://example.invalid/simple",
"-r extra-requirements.txt",
]
),
encoding="utf-8",
)
(first_package / "extra-requirements.txt").write_text(
"diffusers>=0.37.0\n",
encoding="utf-8",
)
(second_package / "requirements.txt").write_text(
"omegaconf==2.3.0\nsoundfile\n",
encoding="utf-8",
)
packages = modal_cloud_module._custom_node_runtime_packages(custom_nodes_dir)
assert packages == (
"omegaconf==2.3.0",
"hydra-core==1.3.2",
"diffusers>=0.37.0",
"soundfile",
)
def test_modal_cloud_returns_no_custom_node_packages_without_requirements(
modal_cloud_module: Any,
tmp_path: Path,
) -> None:
"""Custom-node package installation should be skipped when no requirements exist."""
custom_nodes_dir = tmp_path / "custom_nodes"
(custom_nodes_dir / "example").mkdir(parents=True)
assert modal_cloud_module._custom_node_runtime_packages(custom_nodes_dir) == ()
def test_modal_cloud_restores_runtime_pins_after_custom_node_packages(
modal_cloud_module: Any,
) -> None:
"""Custom requirements must not remain authoritative for shared runtime packages."""
class RecordingImage:
"""Record Modal image package layers."""
def __init__(self) -> None:
"""Initialize an empty call record."""
self.calls: list[tuple[str, ...]] = []
def pip_install(self, *packages: str, **_options: Any) -> "RecordingImage":
"""Record one package installation layer."""
self.calls.append(packages)
return self
image = RecordingImage()
result = modal_cloud_module._install_custom_node_packages(
image,
("numpy==1.26.4", "opencv-python-headless==4.11.0.86"),
)
assert result is image
assert image.calls[0] == (
"numpy==1.26.4",
"opencv-python-headless==4.11.0.86",
)
assert "numpy==2.3.5" in image.calls[1]
assert "opencv-python-headless==4.13.0.92" in image.calls[1]
def test_modal_cloud_missing_prompt_node_class_reports_unavailable_bundle_path(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Missing-class diagnostics should distinguish absent payload metadata from unavailable storage."""
fake_nodes_module = types.SimpleNamespace(NODE_CLASS_MAPPINGS={})
monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
with pytest.raises(modal_cloud_module.RemoteSubgraphExecutionError) as exc_info:
modal_cloud_module._ensure_prompt_node_classes_registered(
component_id="component-1",
prompt={"2": {"class_type": "KSamplerLoraSigmaInverse", "inputs": {}}},
custom_nodes_root=None,
custom_nodes_bundle_path="/custom_nodes/manifest.json",
)
message = str(exc_info.value)
assert "custom_nodes_bundle='/custom_nodes/manifest.json'" in message
assert "not available in Modal worker storage" in message
def test_modal_cloud_retries_custom_node_import_for_missing_prompt_class(
modal_cloud_module: Any,
monkeypatch: Any,
tmp_path: Path,
) -> None:
"""A missing prompt class should trigger one external custom-node import retry."""
import_calls: list[str] = []
custom_nodes_root = tmp_path / "custom_nodes_bundle"
custom_nodes_root.mkdir()
class RegisteredLaterNode:
"""Node type registered by the retry import."""
async def fake_init_external_custom_nodes() -> None:
"""Register the custom class when the retry import runs."""
import_calls.append("init")
fake_nodes_module.NODE_CLASS_MAPPINGS["KSamplerLoraSigmaInverse"] = RegisteredLaterNode
fake_nodes_module = types.SimpleNamespace(
NODE_CLASS_MAPPINGS={},
init_external_custom_nodes=fake_init_external_custom_nodes,
)
monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
bootstrap_owner = _cloud_comfy_bootstrap_owner()
monkeypatch.setattr(
bootstrap_owner,
"_register_custom_nodes_root",
lambda path: import_calls.append(str(path)),
)
monkeypatch.setattr(
bootstrap_owner,
"_install_loader_cache_wrappers",
lambda: import_calls.append("wrappers"),
)
resolved_mapping = modal_cloud_module._ensure_prompt_node_classes_registered(
component_id="component-1",
prompt={"2": {"class_type": "KSamplerLoraSigmaInverse", "inputs": {}}},
custom_nodes_root=custom_nodes_root,
)
assert resolved_mapping["KSamplerLoraSigmaInverse"] is RegisteredLaterNode
assert import_calls == [str(custom_nodes_root), "init", "wrappers"]
def test_modal_cloud_commits_after_each_genuine_triton_cache_miss(
modal_cloud_module: Any,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Repeated new specializations for one profile must each become durable."""
signal = {"size": 10}
runtime_module = types.ModuleType("modal_llm_runtime")
runtime_module.triton_compile_miss_signal_size = lambda: signal["size"]
runtime_module.triton_compile_listener_engine_pids = lambda: (1234,)
monkeypatch.setitem(sys.modules, "modal_llm_runtime", runtime_module)
payload = {
"payload_kind": "subgraph",
"execute_node_ids": ["llm"],
"subgraph_prompt": {
"llm": {
"class_type": "ModalLLM",
"inputs": {"model_profile": "same-profile"},
}
},
}
class FakeVolume:
"""Count explicit compile-cache commits."""
def __init__(self) -> None:
"""Initialize an empty commit count."""
self.commits = 0
def commit(self) -> None:
"""Record one explicit commit."""
self.commits += 1
volume = FakeVolume()
checkpoint = modal_cloud_module._llm_compile_miss_checkpoint(payload)
assert checkpoint is not None
assert (
modal_cloud_module._commit_actual_llm_compile_cache(checkpoint, volume)
is False
)
signal["size"] = 20
assert modal_cloud_module._commit_actual_llm_compile_cache(checkpoint, volume)
second_checkpoint = modal_cloud_module._llm_compile_miss_checkpoint(payload)
signal["size"] = 35
assert modal_cloud_module._commit_actual_llm_compile_cache(
second_checkpoint,
volume,
)
assert volume.commits == 2
def test_modal_cloud_does_not_claim_cache_hits_without_engine_listener(
modal_cloud_module: Any,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Missing EngineCore telemetry must skip commits without reporting a hit."""
runtime_module = types.ModuleType("modal_llm_runtime")
runtime_module.triton_compile_miss_signal_size = lambda: 0
runtime_module.triton_compile_listener_engine_pids = lambda: ()
monkeypatch.setitem(sys.modules, "modal_llm_runtime", runtime_module)
checkpoint = modal_cloud_module._LLMCompileMissCheckpoint(
profiles=("profile",),
signal_size=0,
listener_engine_pids=(),
)
class FakeVolume:
"""Reject any commit without listener evidence."""
def commit(self) -> None:
"""Fail if missing telemetry causes a commit."""
raise AssertionError("listener-less execution cannot commit")
with caplog.at_level(logging.WARNING):
committed = modal_cloud_module._commit_actual_llm_compile_cache(
checkpoint,
FakeVolume(),
)
assert committed is False
assert "no live vLLM EngineCore" in caplog.text
assert "every Triton lookup hit" not in caplog.text
def test_modal_cloud_reloads_compile_cache_before_restored_runtime(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Restored workers should refresh the cache Volume before importing runtimes."""
calls: list[str] = []
class FakeVolume:
"""Record one reload."""
def reload(self) -> None:
"""Record the refresh."""
calls.append("reload")
monkeypatch.setattr(
_cloud_prewarm_owner(),
"_ensure_comfy_runtime_initialized",
lambda custom_nodes_root: calls.append("runtime"),
)
monkeypatch.setattr(
_cloud_prewarm_owner(),
"_load_execution_module",
lambda: calls.append("execution"),
)
modal_cloud_module._prewarm_restored_runtime(FakeVolume())
assert calls == ["reload", "runtime", "execution"]
def test_get_hourly_modal_app_billing_filters_and_caches_gpu_app(
modal_billing_module: Any,
monkeypatch: Any,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Hourly billing should quietly select and cache one GPU app/environment."""
settings = modal_billing_module.get_settings()
selected_settings = modal_billing_module.settings_for_modal_gpu(
settings,
"B300",
)
app_name = modal_billing_module.modal_deployment_app_name(selected_settings)
interval_start = datetime(2026, 8, 19, 7, 0, tzinfo=timezone.utc)
report_calls: list[dict[str, Any]] = []
def workspace_billing_report(**kwargs: Any) -> list[dict[str, Any]]:
"""Return matching and unrelated hourly billing rows."""
report_calls.append(kwargs)
return [
{
"object_id": "ap-selected",
"description": app_name,
"environment_name": "main",
"interval_start": interval_start,
"cost": Decimal("0.11799566"),
},
{
"object_id": "ap-other-env",
"description": app_name,
"environment_name": "dev",
"interval_start": interval_start,
"cost": Decimal("9.99"),
},
{
"object_id": "ap-other-app",
"description": "unrelated",
"environment_name": "main",
"interval_start": interval_start,
"cost": Decimal("8.88"),
},
]
original_import_module = modal_billing_module.importlib.import_module
def fake_import_module(name: str) -> Any:
"""Supply the public Modal billing and environment SDK surfaces."""
if name == "modal._object":
return types.SimpleNamespace(_get_environment_name=lambda _environment: "main")
if name == "modal.environments":
return types.SimpleNamespace(ensure_env=lambda environment: environment or "main")
if name == "modal.billing":
return types.SimpleNamespace(
workspace_billing_report=workspace_billing_report
)
if name == "modal.exception":
return types.SimpleNamespace(Error=RuntimeError)
return original_import_module(name)
monkeypatch.setattr(modal_billing_module, "modal", object())
monkeypatch.setattr(
modal_billing_module.importlib,
"import_module",
fake_import_module,
)
modal_billing_module._MODAL_HOURLY_BILLING_CACHE.clear()
modal_billing_module._MODAL_HOURLY_BILLING_ERROR_CACHE.clear()
caplog.set_level(logging.INFO, logger=modal_billing_module.__name__)
now = datetime(2026, 8, 19, 8, 15, tzinfo=timezone.utc)
first_status = asyncio.run(
modal_billing_module.get_hourly_modal_app_billing(
"B300",
settings,
now=now,
)
)
cached_status = asyncio.run(
modal_billing_module.get_hourly_modal_app_billing(
"B300",
settings,
now=now + timedelta(minutes=30),
)
)
assert first_status is cached_status
assert first_status.app_id == "ap-selected"
assert first_status.app_name == app_name
assert first_status.modal_gpu == "B300"
assert first_status.environment_name == "main"
assert first_status.app_cost_usd_before_credits == Decimal("0.11799566")
assert first_status.has_usage is True
assert first_status.interval_start == interval_start
assert first_status.interval_end == datetime(
2026,
8,
19,
8,
tzinfo=timezone.utc,
)
assert first_status.next_refresh_at == datetime(
2026,
8,
19,
9,
10,
tzinfo=timezone.utc,
)
assert report_calls == [
{
"start": interval_start,
"end": datetime(2026, 8, 19, 8, 0, tzinfo=timezone.utc),
"resolution": "h",
}
]
assert "Fetched Modal hourly billing" not in caplog.text
def test_get_hourly_modal_app_billing_caches_report_failures(
modal_billing_module: Any,
monkeypatch: Any,
) -> None:
"""Fast status polling should attempt a failed billing interval only once."""
report_calls = 0
def workspace_billing_report(**_kwargs: Any) -> list[dict[str, Any]]:
"""Simulate a workspace that cannot expose billing reports."""
nonlocal report_calls
report_calls += 1
raise RuntimeError("billing access denied")
original_import_module = modal_billing_module.importlib.import_module
def fake_import_module(name: str) -> Any:
"""Supply a failing public Modal billing report."""
if name == "modal._object":
return types.SimpleNamespace(_get_environment_name=lambda _environment: "main")
if name == "modal.environments":
return types.SimpleNamespace(ensure_env=lambda environment: environment or "main")
if name == "modal.billing":
return types.SimpleNamespace(
workspace_billing_report=workspace_billing_report
)
if name == "modal.exception":
return types.SimpleNamespace(Error=RuntimeError)
return original_import_module(name)
monkeypatch.setattr(modal_billing_module, "modal", object())
monkeypatch.setattr(
modal_billing_module.importlib,
"import_module",
fake_import_module,
)
modal_billing_module._MODAL_HOURLY_BILLING_CACHE.clear()
modal_billing_module._MODAL_HOURLY_BILLING_ERROR_CACHE.clear()
now = datetime(2026, 8, 19, 8, 15, tzinfo=timezone.utc)
for _attempt in range(2):
with pytest.raises(
modal_billing_module.ModalBillingStatusError,
match="billing access denied",
):
asyncio.run(
modal_billing_module.get_hourly_modal_app_billing(
"L4",
now=now,
)
)
assert report_calls == 1
def test_modal_cloud_emits_restored_node_cache_events(
modal_cloud_module: Any,
) -> None:
"""Persisted node-cache restores should surface one cached-node marker per restored node."""
observed_updates: list[dict[str, Any]] = []
modal_cloud_module._emit_restored_node_cache_events(observed_updates.append, ["7", "9"])
assert observed_updates == [
{
"event_type": "node_cached",
"node_id": "7",
"display_node_id": "7",
"real_node_id": "7",
},
{
"event_type": "node_cached",
"node_id": "9",
"display_node_id": "9",
"real_node_id": "9",
},
]
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]] = []
original_nodes_module = sys.modules.get("nodes")
original_execution_module = sys.modules.get("execution")
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)
bootstrap_owner = _cloud_comfy_bootstrap_owner()
original_base_initialized = bootstrap_owner._COMFY_RUNTIME_BASE_INITIALIZED
original_custom_node_roots = set(bootstrap_owner._COMFY_RUNTIME_CUSTOM_NODE_ROOTS)
bootstrap_owner._COMFY_RUNTIME_BASE_INITIALIZED = False
bootstrap_owner._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:
bootstrap_owner._COMFY_RUNTIME_BASE_INITIALIZED = original_base_initialized
bootstrap_owner._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.clear()
bootstrap_owner._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.update(original_custom_node_roots)
loaded_execution_module = sys.modules.get("execution")
if original_execution_module is None:
sys.modules.pop("execution", None)
elif loaded_execution_module is not None and original_nodes_module is not None:
loaded_execution_module.nodes = original_nodes_module
assert init_calls == [("extra", False, True), ("external",)]
assert folder_path_calls == [("custom_nodes", str(custom_nodes_root), True)]
def test_modal_cloud_uses_comfy_prompt_executor_cache_defaults(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""The remote worker should mirror ComfyUI's prompt executor cache configuration."""
fake_args = types.SimpleNamespace(
cache_lru=0,
cache_ram=4.0,
cache_none=False,
)
fake_cli_args_module = types.SimpleNamespace(args=fake_args)
fake_execution_module = types.SimpleNamespace(
CacheType=types.SimpleNamespace(
CLASSIC="classic",
LRU="lru",
RAM_PRESSURE="ram-pressure",
NONE="none",
)
)
monkeypatch.setitem(sys.modules, "comfy.cli_args", fake_cli_args_module)
cache_type, cache_args = modal_cloud_module._prompt_executor_cache_config(fake_execution_module)
assert cache_type == "ram-pressure"
assert cache_args == {"lru": 0, "ram": 4.0}
def test_modal_cloud_uses_current_comfy_ram_pressure_cache_defaults(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Current list-valued cache arguments should use ComfyUI's calculated RAM thresholds."""
fake_args = types.SimpleNamespace(
cache_classic=False,
cache_lru=0,
cache_ram=[],
cache_none=False,
)
fake_cli_args_module = types.SimpleNamespace(args=fake_args)
fake_model_management_module = types.SimpleNamespace(total_ram=64 * 1024)
fake_execution_module = types.SimpleNamespace(
CacheType=types.SimpleNamespace(
CLASSIC="classic",
LRU="lru",
RAM_PRESSURE="ram-pressure",
NONE="none",
)
)
monkeypatch.setitem(sys.modules, "comfy.cli_args", fake_cli_args_module)
monkeypatch.setitem(
sys.modules,
"comfy.model_management",
fake_model_management_module,
)
cache_type, cache_args = modal_cloud_module._prompt_executor_cache_config(fake_execution_module)
assert cache_type == "ram-pressure"
assert cache_args == {"lru": 0, "ram": 6.4, "ram_inactive": 64.0}
def test_modal_cloud_sync_cache_reads_use_current_get_local_api(
modal_cloud_module: Any,
) -> None:
"""Synchronous progress paths should avoid creating current cache get coroutines."""
cache_entry = types.SimpleNamespace(outputs=[[512, 512]])
class FakeOutputsCache:
"""Current cache double with async get and synchronous get_local."""
async def get(self, node_id: str) -> Any:
"""Fail if the synchronous compatibility path calls async get."""
del node_id
raise AssertionError("async get should not be created")
def get_local(self, node_id: str) -> Any:
"""Return the already materialized local cache entry."""
assert node_id == "115"
return cache_entry
resolved_entry = modal_cloud_module._prompt_executor_cache_get_sync(
FakeOutputsCache(),
"115",
)
assert resolved_entry is cache_entry
def test_modal_cloud_reuses_prompt_executor_for_same_cache_scope(
modal_cloud_module: Any,
tmp_path: Path,
) -> None:
"""Warm-container subgraph runs should reuse one PromptExecutor per cache scope."""
class FakePromptExecutor:
"""Simple PromptExecutor double that records how many instances were created."""
instances_created = 0
def __init__(self, server: Any, cache_type: Any = False, cache_args: Any = None) -> None:
"""Capture initialization state for later assertions."""
type(self).instances_created += 1
self.server = server
self.cache_type = cache_type
self.cache_args = cache_args
self.status_messages = [("stale", {})]
self.success = False
self.history_result = {"stale": True}
fake_execution_module = types.SimpleNamespace(PromptExecutor=FakePromptExecutor)
first_server = types.SimpleNamespace(client_id="first", last_node_id="node-1")
second_server = types.SimpleNamespace(client_id="second", last_node_id="node-2")
prompt_execution_owner = _cloud_prompt_execution_owner()
original_states = dict(prompt_execution_owner._PROMPT_EXECUTOR_STATES)
prompt_execution_owner._PROMPT_EXECUTOR_STATES.clear()
try:
first_state = modal_cloud_module._get_or_create_prompt_executor_state(
execution=fake_execution_module,
prompt_server=first_server,
cache_type="classic",
cache_args={"lru": 0, "ram": 0.0},
custom_nodes_root=tmp_path / "bundle-a",
)
modal_cloud_module._reset_prompt_executor_request_state(first_state.executor, first_server)
second_state = modal_cloud_module._get_or_create_prompt_executor_state(
execution=fake_execution_module,
prompt_server=second_server,
cache_type="classic",
cache_args={"lru": 0, "ram": 0.0},
custom_nodes_root=tmp_path / "bundle-a",
)
modal_cloud_module._reset_prompt_executor_request_state(second_state.executor, second_server)
finally:
prompt_execution_owner._PROMPT_EXECUTOR_STATES.clear()
prompt_execution_owner._PROMPT_EXECUTOR_STATES.update(original_states)
assert FakePromptExecutor.instances_created == 1
assert first_state is second_state
assert second_state.executor.server is second_server
assert second_state.executor.status_messages == []
assert second_state.executor.success is True
assert second_state.executor.history_result == {}
assert second_server.client_id is None
assert second_server.last_node_id is None
def test_modal_cloud_separates_prompt_executor_cache_scopes_by_custom_nodes_root(
modal_cloud_module: Any,
tmp_path: Path,
) -> None:
"""Different custom-node bundle roots should not share a PromptExecutor cache scope."""
class FakePromptExecutor:
"""Simple PromptExecutor double used to count cache-scope creations."""
instances_created = 0
def __init__(self, server: Any, cache_type: Any = False, cache_args: Any = None) -> None:
"""Capture initialization state for later assertions."""
type(self).instances_created += 1
self.server = server
self.cache_type = cache_type
self.cache_args = cache_args
self.status_messages = []
self.success = True
self.history_result = {}
fake_execution_module = types.SimpleNamespace(PromptExecutor=FakePromptExecutor)
prompt_execution_owner = _cloud_prompt_execution_owner()
original_states = dict(prompt_execution_owner._PROMPT_EXECUTOR_STATES)
prompt_execution_owner._PROMPT_EXECUTOR_STATES.clear()
try:
first_state = modal_cloud_module._get_or_create_prompt_executor_state(
execution=fake_execution_module,
prompt_server=types.SimpleNamespace(client_id=None, last_node_id=None),
cache_type="classic",
cache_args={"lru": 0, "ram": 0.0},
custom_nodes_root=tmp_path / "bundle-a",
)
second_state = modal_cloud_module._get_or_create_prompt_executor_state(
execution=fake_execution_module,
prompt_server=types.SimpleNamespace(client_id=None, last_node_id=None),
cache_type="classic",
cache_args={"lru": 0, "ram": 0.0},
custom_nodes_root=tmp_path / "bundle-b",
)
finally:
prompt_execution_owner._PROMPT_EXECUTOR_STATES.clear()
prompt_execution_owner._PROMPT_EXECUTOR_STATES.update(original_states)
assert FakePromptExecutor.instances_created == 2
assert first_state is not second_state
def test_modal_cloud_awaits_async_node_output_cache_writes(
modal_cloud_module: Any,
) -> None:
"""Persisted cache writes should use Modal's non-blocking Dict interface."""
observed_writes: list[tuple[str, dict[str, Any]]] = []
class AsyncPut:
"""Expose the callable shape used by Modal's synchronized methods."""
async def aio(self, cache_key: str, record: dict[str, Any]) -> None:
"""Record one asynchronous cache write."""
observed_writes.append((cache_key, record))
class AsyncCacheStore:
"""Reject synchronous assignment while exposing an asynchronous put method."""
put = AsyncPut()
def __setitem__(self, cache_key: str, record: dict[str, Any]) -> None:
"""Fail if persistence regresses to Modal's blocking operator interface."""
del cache_key, record
raise AssertionError("cache write must use put.aio")
record = {"version": 1, "outputs_zlib": b"payload"}
asyncio.run(
modal_cloud_module._node_output_cache_store_put(
AsyncCacheStore(),
"NC_example",
record,
)
)
assert observed_writes == [("NC_example", record)]
def test_modal_cloud_restores_persisted_node_cache_across_prompt_executor_instances(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""A fresh PromptExecutor cache should round-trip persisted node outputs through Modal Dict state."""
monkeypatch.setitem(sys.modules, "torchsde", types.ModuleType("torchsde"))
modal_cloud_module._ensure_comfy_runtime_initialized(None)
import comfy_execution.caching as comfy_caching
execution = modal_cloud_module._load_execution_module()
nodes_module = modal_cloud_module._load_nodes_module()
cache_store: dict[str, Any] = {}
prompt = {
"node_1": {
"class_type": "PersistentCacheNode",
"inputs": {"value": 4},
"_meta": {},
}
}
_PersistentCacheNode.invocation_count = 0
monkeypatch.setitem(nodes_module.NODE_CLASS_MAPPINGS, "PersistentCacheNode", _PersistentCacheNode)
monkeypatch.setitem(
comfy_caching.nodes.NODE_CLASS_MAPPINGS,
"PersistentCacheNode",
_PersistentCacheNode,
)
monkeypatch.setitem(
nodes_module.NODE_DISPLAY_NAME_MAPPINGS,
"PersistentCacheNode",
"PersistentCacheNode",
)
cache_entry = execution.CacheEntry(ui={"output": {"value": [5]}}, outputs=[[5]])
first_executor = execution.PromptExecutor(
modal_cloud_module._NullPromptServer(),
cache_type=execution.CacheType.CLASSIC,
cache_args={"lru": 0, "ram": 0.0},
)
restored_first = asyncio.run(
modal_cloud_module._restore_persisted_node_output_cache_entries(
execution,
first_executor,
prompt_id="prompt-a",
prompt=copy.deepcopy(prompt),
cache_store=cache_store,
)
)
asyncio.run(
modal_cloud_module._await_maybe(
first_executor.caches.outputs.set("node_1", cache_entry)
)
)
persisted_nodes = asyncio.run(
modal_cloud_module._persist_node_output_cache_entries(
first_executor,
prompt=copy.deepcopy(prompt),
cache_store=cache_store,
)
)
second_executor = execution.PromptExecutor(
modal_cloud_module._NullPromptServer(),
cache_type=execution.CacheType.CLASSIC,
cache_args={"lru": 0, "ram": 0.0},
)
restored_cache_keys_by_node_id: dict[str, str] = {}
restored_second = asyncio.run(
modal_cloud_module._restore_persisted_node_output_cache_entries(
execution,
second_executor,
prompt_id="prompt-b",
prompt=copy.deepcopy(prompt),
cache_store=cache_store,
restored_cache_keys_by_node_id=restored_cache_keys_by_node_id,
)
)
restored_entry = asyncio.run(
modal_cloud_module._await_maybe(second_executor.caches.outputs.get("node_1"))
)
assert restored_first == []
assert persisted_nodes == ["node_1"]
assert restored_second == ["node_1"]
assert restored_cache_keys_by_node_id == {"node_1": next(iter(cache_store))}
assert list(cache_store) and all(key.startswith("NC_") for key in cache_store)
assert restored_entry == cache_entry
def test_modal_cloud_installs_persisted_cache_restore_after_live_set_prompt(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Persisted-cache restore should run after PromptExecutor prepares the active outputs cache."""
class FakeOutputsCache:
"""Minimal outputs-cache stub with a mutable cache-key-set marker."""
def __init__(self) -> None:
"""Initialize the fake cache-key-set marker."""
self.cache_key_set = None
async def set_prompt(self, dynprompt: Any, node_ids: Any, is_changed_cache: Any) -> None:
"""Simulate ComfyUI assigning the live cache-key set during prompt setup."""
del dynprompt, node_ids, is_changed_cache
self.cache_key_set = "live-cache-key-set"
outputs_cache = FakeOutputsCache()
executor = types.SimpleNamespace(caches=types.SimpleNamespace(outputs=outputs_cache))
observed_events: list[tuple[str, Any]] = []
async def fake_restore(
execution: Any,
prepared_outputs_cache: Any,
*,
prompt: dict[str, Any],
cache_store: Any,
required_materialized_node_ids: Any = None,
restored_cache_keys_by_node_id: dict[str, str] | None = None,
) -> list[str]:
"""Record the cache-key-set marker visible at restore time."""
del execution, required_materialized_node_ids
if restored_cache_keys_by_node_id is not None:
restored_cache_keys_by_node_id["12"] = "NC_example"
observed_events.append(
(
"restore",
prepared_outputs_cache.cache_key_set,
tuple(sorted(prompt)),
cache_store,
)
)
return ["12"]
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_restore_persisted_node_output_cache_entries_into_prepared_cache",
fake_restore,
)
restore_state = (
modal_cloud_module._install_prompt_executor_persisted_cache_restore(
object(),
executor,
component_id="component-1",
prompt={"12": {"class_type": "PersistentCacheNode", "inputs": {}}},
cache_store={"NC_example": {"version": 1}},
)
)
try:
asyncio.run(outputs_cache.set_prompt(object(), ["12"], object()))
finally:
restore_state.restore_original_method()
assert restore_state.restored_node_ids == ["12"]
assert restore_state.restored_cache_keys_by_node_id == {"12": "NC_example"}
assert observed_events == [
("restore", "live-cache-key-set", ("12",), {"NC_example": {"version": 1}})
]
assert outputs_cache.set_prompt.__func__ is FakeOutputsCache.set_prompt
def test_modal_cloud_skips_rewriting_restored_distributed_cache_entries(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Persist should skip distributed cache entries that were restored unchanged this run."""
monkeypatch.setitem(sys.modules, "torchsde", types.ModuleType("torchsde"))
modal_cloud_module._ensure_comfy_runtime_initialized(None)
execution = modal_cloud_module._load_execution_module()
cache_entry = execution.CacheEntry(ui={"output": {"value": [5]}}, outputs=[[5]])
cache_key = "NC_existing"
cache_store: dict[str, Any] = {cache_key: {"version": 1, "outputs_zlib": b"old"}}
observed_logs: list[tuple[Any, ...]] = []
class FakeOutputsCache:
"""Minimal outputs cache stub for persist-phase tests."""
def __init__(self) -> None:
"""Populate one persistent cache entry."""
self.cache_key_set = object()
async def get(self, node_id: str) -> Any:
"""Return the prepared cache entry for the target node only."""
if node_id == "node_1":
return cache_entry
return None
executor = types.SimpleNamespace(caches=types.SimpleNamespace(outputs=FakeOutputsCache()))
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_node_output_cache_key_from_key_set_sync",
lambda cache_key_set, node_id: cache_key if node_id == "node_1" else None,
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_emit_cloud_info",
lambda message, *args: observed_logs.append((message, *args)),
)
persisted_nodes = asyncio.run(
modal_cloud_module._persist_node_output_cache_entries(
executor,
prompt={"node_1": {"class_type": "PersistentCacheNode", "inputs": {"value": 4}}},
cache_store=cache_store,
restored_cache_keys_by_node_id={"node_1": cache_key},
)
)
assert persisted_nodes == []
assert cache_store == {cache_key: {"version": 1, "outputs_zlib": b"old"}}
assert observed_logs[-1] == (
"Node output cache write node=%s key_prefix=%s result=skip reason=restored-hit",
"node_1",
"NC_existing",
)
def test_modal_cloud_skips_restored_hit_before_cache_entry_serialization(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Persist should avoid cache-entry serialization entirely for restored-hit keys."""
cache_entry = object()
class FakeOutputsCache:
"""Minimal outputs cache stub for persist fast-path tests."""
def __init__(self) -> None:
"""Expose a cache-key-set marker for sync key generation."""
self.cache_key_set = object()
def get(self, node_id: str) -> Any:
"""Return a placeholder cache entry for the target node."""
if node_id == "node_1":
return cache_entry
return None
executor = types.SimpleNamespace(caches=types.SimpleNamespace(outputs=FakeOutputsCache()))
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_node_output_cache_key_from_key_set_sync",
lambda cache_key_set, node_id: "NC_existing" if node_id == "node_1" else None,
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_serialize_node_output_cache_entry",
lambda cache_entry, max_bytes: (_ for _ in ()).throw(
AssertionError("restored-hit should skip before serialization")
),
)
persisted_nodes = asyncio.run(
modal_cloud_module._persist_node_output_cache_entries(
executor,
prompt={"node_1": {"class_type": "PersistentCacheNode", "inputs": {"value": 4}}},
cache_store={},
restored_cache_keys_by_node_id={"node_1": "NC_existing"},
)
)
assert persisted_nodes == []
def test_modal_cloud_restores_persisted_node_cache_entries_in_parallel(
modal_cloud_module: Any,
monkeypatch: Any,
) -> None:
"""Prepared-cache restore should overlap independent distributed lookups."""
in_flight_gets = 0
max_in_flight_gets = 0
restored_values: dict[str, Any] = {}
class FakeOutputsCache:
"""Minimal outputs cache stub for restore concurrency tests."""
def __init__(self) -> None:
"""Expose the cache-key-set marker read by restore."""
self.cache_key_set = object()
async def get(self, node_id: str) -> Any:
"""Return any previously restored entry."""
return restored_values.get(node_id)
async def set(self, node_id: str, cache_entry: Any) -> None:
"""Record one restored cache entry."""
restored_values[node_id] = cache_entry
async def fake_key_from_key_set(cache_key_set: Any, node_id: str) -> str:
"""Yield briefly so multiple node lookups can queue together."""
del cache_key_set
await asyncio.sleep(0)
return f"NC_{node_id}"
async def fake_store_get(cache_store: Any, cache_key: str) -> Any:
"""Track how many distributed cache reads overlap."""
nonlocal in_flight_gets, max_in_flight_gets
del cache_store
in_flight_gets += 1
max_in_flight_gets = max(max_in_flight_gets, in_flight_gets)
await asyncio.sleep(0.01)
in_flight_gets -= 1
return {"cache_key": cache_key}
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_node_output_cache_key_from_key_set_async",
fake_key_from_key_set,
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_node_output_cache_store_get",
fake_store_get,
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(),
"_deserialize_node_output_cache_entry",
lambda execution, record: record,
)
monkeypatch.setattr(
_cloud_node_output_cache_owner(), "_emit_cloud_info", lambda *args: None
)
outputs_cache = FakeOutputsCache()
restored_node_ids = asyncio.run(
modal_cloud_module._restore_persisted_node_output_cache_entries_into_prepared_cache(
object(),
outputs_cache,
prompt={
"node_1": {"class_type": "PersistentCacheNode", "inputs": {"value": 1}},
"node_2": {"class_type": "PersistentCacheNode", "inputs": {"value": 2}},
"node_3": {"class_type": "PersistentCacheNode", "inputs": {"value": 3}},
},
cache_store={},
)
)
assert restored_node_ids == ["node_1", "node_2", "node_3"]
assert max_in_flight_gets >= 2
assert restored_values == {
"node_1": {"cache_key": "NC_node_1"},
"node_2": {"cache_key": "NC_node_2"},
"node_3": {"cache_key": "NC_node_3"},
}
def test_modal_cloud_creates_default_custom_nodes_dir_when_missing(
modal_cloud_module: Any,
monkeypatch: Any,
tmp_path: Path,
) -> None:
"""The remote runtime should create an empty default custom_nodes directory for ComfyUI."""
comfyui_root = tmp_path / "comfyui"
comfyui_root.mkdir()
bootstrap_owner = _cloud_comfy_bootstrap_owner()
monkeypatch.setattr(bootstrap_owner, "_REMOTE_COMFYUI_ROOT", comfyui_root)
monkeypatch.setattr(bootstrap_owner, "_LOCAL_COMFYUI_ROOT", tmp_path / "missing-local")
custom_nodes_dir = modal_cloud_module._ensure_default_custom_nodes_dir()
assert custom_nodes_dir == comfyui_root / "custom_nodes"
assert custom_nodes_dir is not None
assert custom_nodes_dir.exists()
assert custom_nodes_dir.is_dir()
Reference in New Issue View Git Blame Copy Permalink