"""Tests for prompt rewriting and asset sync integration.""" from __future__ import annotations import asyncio import json from pathlib import Path import threading from types import SimpleNamespace from typing import Any import pytest class _FakeRemoteModelNode: """Fake node that produces a non-transportable MODEL output.""" RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) OUTPUT_IS_LIST = (False,) class _FakeCheckpointLoaderSimpleNode: """Fake root loader node that produces a non-transportable MODEL output.""" RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) OUTPUT_IS_LIST = (False,) class _FakeRemoteSamplerNode: """Fake node that consumes a model and produces a transportable latent.""" RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("latent",) OUTPUT_IS_LIST = (False,) class _FakeRemoteNoiseNode: """Fake node that produces a non-transportable NOISE strategy object.""" RETURN_TYPES = ("NOISE",) RETURN_NAMES = ("noise",) OUTPUT_IS_LIST = (False,) class _FakeVAELoaderNode: """Fake VAE loader that produces a non-transportable VAE output.""" RETURN_TYPES = ("VAE",) RETURN_NAMES = ("vae",) OUTPUT_IS_LIST = (False,) class _FakeVAEDecodeNode: """Fake VAE decoder that consumes latent and VAE inputs and produces an image.""" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) OUTPUT_IS_LIST = (False,) class _FakeVAEEncodeNode: """Fake VAE encoder that consumes image and VAE inputs and produces a latent.""" RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("latent",) OUTPUT_IS_LIST = (False,) class _FakeLatentSourceNode: """Fake node that produces a transportable LATENT output.""" RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("latent",) OUTPUT_IS_LIST = (False,) class _FakeRemoteClipNode: """Fake node that produces a non-transportable CLIP output.""" RETURN_TYPES = ("CLIP",) RETURN_NAMES = ("clip",) OUTPUT_IS_LIST = (False,) class _FakeLocalSinkNode: """Fake local node used to verify downstream rewiring.""" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) OUTPUT_IS_LIST = (False,) class _FakeRemoteConditioningNode: """Fake node that produces a non-transportable CONDITIONING output.""" RETURN_TYPES = ("CONDITIONING",) RETURN_NAMES = ("conditioning",) OUTPUT_IS_LIST = (False,) class _FakeRemoteImageNode: """Fake remote node that produces a transportable IMAGE output.""" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) OUTPUT_IS_LIST = (False,) class _FakeRemoteVideoNode: """Fake remote node that produces a transportable VIDEO output.""" RETURN_TYPES = ("VIDEO",) RETURN_NAMES = ("video",) OUTPUT_IS_LIST = (False,) class _FakeRemoteAudioNode: """Fake remote node that produces a transportable AUDIO output.""" RETURN_TYPES = ("AUDIO",) RETURN_NAMES = ("audio",) OUTPUT_IS_LIST = (False,) class _FakeRemoteImageConsumerNode: """Fake remote node that consumes IMAGE and produces IMAGE.""" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) OUTPUT_IS_LIST = (False,) class _FakePreviewImageNode: """Fake preview node that behaves like a terminal UI output.""" RETURN_TYPES = () RETURN_NAMES = () OUTPUT_IS_LIST = () OUTPUT_NODE = True class _FakeRemoteArtifactWriterNode: """Fake terminal remote node that saves artifacts without ComfyUI outputs.""" RETURN_TYPES = () RETURN_NAMES = () OUTPUT_IS_LIST = () OUTPUT_NODE = False def _artifact_finalizer_node_id(summary: Any) -> str: """Return the finalizer id after asserting that prompt rewrite attached it.""" finalizer_node_id = summary.artifact_finalizer_node_id assert isinstance(finalizer_node_id, str) assert finalizer_node_id return finalizer_node_id def test_rewrite_remote_mode_rejects_local_sync_backend( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Remote execution must not queue payloads whose synced assets only exist in local mirror storage.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=True, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=tmp_path / "storage", remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=None, ) sync_engine = sync_engine_module.ModalAssetSyncEngine( volume=sync_engine_module.LocalMirrorVolume(settings.local_storage_root), settings=settings, ) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": {"RemoteImage": _FakeRemoteImageNode}, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() with pytest.raises(api_intercept_module.ModalPromptValidationError) as exc_info: api_intercept_module.rewrite_prompt_for_modal( prompt={"1": {"class_type": "RemoteImage", "inputs": {}}}, workflow={"nodes": [{"id": 1, "properties": {"is_modal_remote": True}}]}, sync_engine=sync_engine, settings=settings, nodes_module=fake_nodes_module, ) assert "requires asset sync to use the Modal volume backend" in str(exc_info.value) class _FakeRemoteModelAndImageNode: """Fake remote node that produces both MODEL and IMAGE outputs.""" RETURN_TYPES = ("MODEL", "IMAGE") RETURN_NAMES = ("model", "image") OUTPUT_IS_LIST = (False, False) class _FakeRemoteModelAndImageConsumerNode: """Fake remote node that consumes MODEL and IMAGE and produces IMAGE.""" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) OUTPUT_IS_LIST = (False,) class _FakePromptListNode: """Fake upstream node that represents a prompt-list producer.""" RETURN_TYPES = ("STRING",) RETURN_NAMES = ("text",) OUTPUT_IS_LIST = (False,) class _FakeModalMapInputNode: """Fake Modal map marker node that passes a wildcard value through.""" RETURN_TYPES = ("*",) RETURN_NAMES = ("value",) OUTPUT_IS_LIST = (False,) class _FakeRemoteStringEchoNode: """Fake remote node that echoes a string output.""" RETURN_TYPES = ("STRING",) RETURN_NAMES = ("text",) OUTPUT_IS_LIST = (False,) class _FakeLocalStringSinkNode: """Fake local node used to consume remote STRING outputs.""" RETURN_TYPES = ("STRING",) RETURN_NAMES = ("text",) OUTPUT_IS_LIST = (False,) def test_queue_prompt_json_includes_resolved_modal_metadata( api_intercept_module: Any, monkeypatch: Any, ) -> None: """Successful queue responses should include resolved remote nodes and component membership.""" class FakePromptQueue: """Minimal prompt queue sink.""" def __init__(self) -> None: """Initialize captured queue items.""" self.items: list[tuple[Any, ...]] = [] def put(self, item: tuple[Any, ...]) -> None: """Record one queued prompt item.""" self.items.append(item) class FakePromptServer: """Minimal PromptServer double for queue-response tests.""" def __init__(self) -> None: """Initialize queue state.""" self.number = 0 self.prompt_queue = FakePromptQueue() def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]: """Return the prompt unchanged.""" return json_data class FakeExecutionModule: """Minimal execution module exposing prompt validation.""" SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = () @staticmethod async def validate_prompt( prompt_id: str, prompt: dict[str, Any], partial_execution_targets: Any, ) -> tuple[bool, None, list[str], list[Any]]: """Accept the supplied prompt with one fake execution target.""" return True, None, ["1"], [] monkeypatch.setattr(api_intercept_module, "_get_execution_module", lambda: FakeExecutionModule) prompt_server = FakePromptServer() response = asyncio.run( api_intercept_module._queue_prompt_json( prompt_server, { "prompt_id": "prompt-1", "prompt": {"1": {"class_type": "Anything", "inputs": {}}}, "extra_data": {}, }, modal_response_payload={ "modal_remote_node_ids": ["1", "2"], "modal_components": [ { "representative_node_id": "1", "node_ids": ["1", "2"], } ], }, ) ) response_payload = json.loads(response.text) assert response_payload["prompt_id"] == "prompt-1" assert response_payload["modal_remote_node_ids"] == ["1", "2"] assert response_payload["modal_components"] == [ { "representative_node_id": "1", "node_ids": ["1", "2"], } ] def test_rewritten_prompt_diagnostics_reports_dependency_cycles( api_intercept_module: Any, monkeypatch: pytest.MonkeyPatch, ) -> None: """Rewritten prompt diagnostics should name local dependency cycles before Comfy executes.""" prompt = { "1": { "class_type": "ModalUniversalExecutor_a", "inputs": {"remote_input_0": ["2", 0]}, }, "2": { "class_type": "ModalUniversalExecutor_b", "inputs": {"remote_input_0": ["1", 0]}, }, } diagnostics = api_intercept_module._modal_rewritten_prompt_diagnostics(prompt) assert diagnostics["cycles"] == [["1", "2", "1"]] warning_messages: list[str] = [] log_messages: list[str] = [] def record_warning(message: str, *args: Any, **_kwargs: Any) -> None: """Record one warning log message.""" warning_messages.append(message % args) def record_log(_level: int, message: str, *args: Any, **_kwargs: Any) -> None: """Record one generic log message.""" log_messages.append(message % args) monkeypatch.setattr(api_intercept_module.logger, "warning", record_warning) monkeypatch.setattr(api_intercept_module.logger, "log", record_log) api_intercept_module._log_modal_rewritten_prompt_diagnostics( prompt_id="prompt-cycle", prompt=prompt, reason="test", ) assert any("Modal rewritten prompt contains dependency cycle(s)" in item for item in warning_messages) assert any("prompt-cycle" in item for item in warning_messages) assert any("Modal rewritten prompt diagnostics" in item for item in log_messages) def test_queue_prompt_json_logs_rewritten_modal_diagnostics_on_validation_failure( api_intercept_module: Any, monkeypatch: pytest.MonkeyPatch, ) -> None: """Validation failures for Modal prompts should log the rewritten dependency graph.""" class FakePromptQueue: """Minimal prompt queue sink.""" def put(self, _item: tuple[Any, ...]) -> None: """Fail if an invalid prompt reaches the queue.""" raise AssertionError("invalid prompt must not be queued") class FakePromptServer: """Minimal PromptServer double for validation-failure tests.""" number = 0 prompt_queue = FakePromptQueue() def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]: """Return the prompt unchanged.""" return json_data class FakeExecutionModule: """Minimal execution module that rejects the prompt.""" SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = () @staticmethod async def validate_prompt( prompt_id: str, prompt: dict[str, Any], partial_execution_targets: Any, ) -> tuple[bool, dict[str, Any], list[str], dict[str, Any]]: """Reject the supplied prompt with a dependency-cycle shaped error.""" del prompt_id, prompt, partial_execution_targets return ( False, { "type": "execution_error", "message": "Dependency cycle detected", "details": "", "extra_info": {}, }, [], {}, ) prompt = { "1": { "class_type": "ModalUniversalExecutor_a", "inputs": {"remote_input_0": ["2", 0]}, }, "2": { "class_type": "ModalUniversalExecutor_b", "inputs": {"remote_input_0": ["1", 0]}, }, } diagnostic_calls: list[dict[str, Any]] = [] def record_diagnostics(**kwargs: Any) -> None: """Record one rewritten-prompt diagnostics request.""" diagnostic_calls.append(dict(kwargs)) monkeypatch.setattr(api_intercept_module, "_get_execution_module", lambda: FakeExecutionModule) monkeypatch.setattr( api_intercept_module, "_log_modal_rewritten_prompt_diagnostics", record_diagnostics, ) response = asyncio.run( api_intercept_module._queue_prompt_json( FakePromptServer(), { "prompt_id": "prompt-cycle", "prompt": prompt, "extra_data": { "modal": { "remote_component_ids": ["1", "2"], } }, }, ) ) assert response.status == 400 assert diagnostic_calls == [ { "prompt_id": "prompt-cycle", "prompt": prompt, "reason": "comfy_validation_failure", "level": api_intercept_module.logging.WARNING, } ] def test_split_phase_order_accounts_for_local_feedback_dependencies( api_intercept_module: Any, ) -> None: """Split phase ordering should treat local re-entry paths as real dependencies.""" prompt = { "3": { "class_type": "RemoteSampler", "inputs": {"model": ["14", 0]}, }, "11": { "class_type": "VAEDecode", "inputs": {"samples": ["3", 0]}, }, "14": { "class_type": "RemoteModel", "inputs": {}, }, "191": { "class_type": "RemoteSampler", "inputs": {"model": ["14", 0], "conditioning": ["358", 0]}, }, "357": { "class_type": "BetterGrok", "inputs": {"prompt_images": ["11", 0]}, }, "358": { "class_type": "RemoteTextEncode", "inputs": {"text": ["357", 1]}, }, } component_prompt = { "3": prompt["3"], "14": prompt["14"], "191": prompt["191"], "358": prompt["358"], } ordered_execute_node_ids = ( api_intercept_module._order_execute_node_ids_for_transportable_splits( prompt=prompt, component_prompt=component_prompt, component_node_ids={"3", "14", "191", "358"}, execute_node_ids=["191", "3"], ) ) assert ordered_execute_node_ids == ["3", "191"] def test_queue_prompt_route_does_not_warm_modal_at_queue_time( api_intercept_module: Any, remote_modal_app_module: Any, settings_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Accepting a queued prompt should not launch Modal warmup containers.""" class FakeRoutes: """Capture aiohttp route registrations.""" def __init__(self) -> None: """Initialize the route handler map.""" self.handlers: dict[str, Any] = {} def post(self, path: str) -> Any: """Return a decorator that records one POST handler.""" def register(handler: Any) -> Any: """Store the decorated handler unchanged.""" self.handlers[path] = handler return handler return register class FakePromptQueue: """Minimal prompt queue sink.""" def __init__(self) -> None: """Initialize captured queue items.""" self.items: list[tuple[Any, ...]] = [] def put(self, item: tuple[Any, ...]) -> None: """Record one queued prompt item.""" self.items.append(item) class FakePromptServer: """Minimal PromptServer double with route registration.""" def __init__(self) -> None: """Initialize routing and queue state.""" self.number = 0 self.routes = FakeRoutes() self.prompt_queue = FakePromptQueue() def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]: """Return the prompt unchanged.""" return json_data class FakeRequest: """Minimal aiohttp request double.""" async def json(self) -> dict[str, Any]: """Return one Modal-marked prompt request.""" return { "prompt_id": "prompt-queue-warmup", "prompt": {"1": {"class_type": "RemoteImage", "inputs": {}}}, "extra_data": { "extra_pnginfo": { "workflow": { "extra": {"comfy_modal": {"gpu": "B300"}}, "nodes": [ { "id": 1, "type": "RemoteImage", "properties": {"is_modal_remote": True}, } ] } } }, } class FakeExecutionModule: """Minimal execution module exposing prompt validation.""" SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = () @staticmethod async def validate_prompt( prompt_id: str, prompt: dict[str, Any], partial_execution_targets: Any, ) -> tuple[bool, None, list[str], list[Any]]: """Accept the supplied prompt with one fake execution target.""" return True, None, ["1"], [] def fail_queue_time_warmup(*_args: Any, **_kwargs: Any) -> int: """Fail the test if queue handling tries to launch proactive warmup.""" raise AssertionError("queue route must not schedule Modal warmup") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=False, 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, ) prompt_server = FakePromptServer() summary = api_intercept_module.RewriteSummary( remote_node_ids=["1"], remote_component_ids=["1"], component_node_ids_by_representative={"1": ["1"]}, sandwiched_local_node_ids=["4"], component_execution_stages=[["1"]], estimated_max_parallel_requests=1, max_parallel_requests_upper_bound=1, ) observed_rewrite_settings: list[Any] = [] def capture_rewrite_settings(**kwargs: Any) -> tuple[dict[str, Any], Any]: """Capture the workflow-derived settings passed into prompt rewriting.""" observed_rewrite_settings.append(kwargs["settings"]) return kwargs["prompt"], summary monkeypatch.setattr(api_intercept_module, "_ROUTE_REGISTERED", False) monkeypatch.setattr( api_intercept_module, "_get_server_module", lambda: SimpleNamespace(PromptServer=SimpleNamespace(instance=prompt_server)), ) monkeypatch.setattr(api_intercept_module, "_get_execution_module", lambda: FakeExecutionModule) monkeypatch.setattr(api_intercept_module, "_emit_modal_status", lambda **_kwargs: None) monkeypatch.setattr( api_intercept_module, "rewrite_prompt_for_modal", capture_rewrite_settings, ) monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", fail_queue_time_warmup) api_intercept_module.setup_modal_queue_route( prompt_server=prompt_server, sync_engine=object(), settings=settings, ) response = asyncio.run(prompt_server.routes.handlers["/modal/queue_prompt"](FakeRequest())) response_payload = json.loads(response.text) assert response_payload["prompt_id"] == "prompt-queue-warmup" assert response_payload["modal_gpu"] == "B300" assert response_payload["modal_remote_node_ids"] == ["1"] assert response_payload["modal_sandwiched_local_node_ids"] == ["4"] assert observed_rewrite_settings[0].modal_gpu == "B300" queued_extra_data = prompt_server.prompt_queue.items[0][3] assert queued_extra_data["modal"]["gpu"] == "B300" assert len(prompt_server.prompt_queue.items) == 1 def test_modal_prompt_rewrite_keeps_event_loop_responsive( api_intercept_module: Any, monkeypatch: Any, ) -> None: """Hashing and upload preparation should execute outside the ComfyUI event loop.""" rewrite_started = threading.Event() release_rewrite = threading.Event() def blocking_rewrite(**kwargs: Any) -> tuple[dict[str, Any], Any]: """Hold one fake rewrite until the async test proves the loop is responsive.""" rewrite_started.set() assert release_rewrite.wait(timeout=1.0) return kwargs["prompt"], api_intercept_module.RewriteSummary() monkeypatch.setattr(api_intercept_module, "rewrite_prompt_for_modal", blocking_rewrite) async def run_test() -> None: """Run the blocking rewrite and an independent event-loop callback together.""" rewrite_task = asyncio.create_task( api_intercept_module.rewrite_prompt_for_modal_async( prompt={"1": {"class_type": "RemoteImage", "inputs": {}}}, workflow=None, ) ) deadline = asyncio.get_running_loop().time() + 1.0 while not rewrite_started.is_set(): if asyncio.get_running_loop().time() >= deadline: raise AssertionError("background rewrite did not start") await asyncio.sleep(0) loop_progress: list[str] = [] asyncio.get_running_loop().call_soon(loop_progress.append, "responsive") await asyncio.sleep(0) assert loop_progress == ["responsive"] assert not rewrite_task.done() release_rewrite.set() await asyncio.wait_for(rewrite_task, timeout=1.0) asyncio.run(run_test()) def test_queue_prompt_route_without_remote_nodes_skips_modal_status_and_rewrite( api_intercept_module: Any, settings_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Prompts with no Modal-enabled workflow nodes should queue without Modal UI setup.""" class FakeRoutes: """Capture aiohttp route registrations.""" def __init__(self) -> None: """Initialize the route handler map.""" self.handlers: dict[str, Any] = {} def post(self, path: str) -> Any: """Return a decorator that records one POST handler.""" def register(handler: Any) -> Any: """Store the decorated handler unchanged.""" self.handlers[path] = handler return handler return register class FakePromptQueue: """Minimal prompt queue sink.""" def __init__(self) -> None: """Initialize captured queue items.""" self.items: list[tuple[Any, ...]] = [] def put(self, item: tuple[Any, ...]) -> None: """Record one queued prompt item.""" self.items.append(item) class FakePromptServer: """Minimal PromptServer double with route registration.""" def __init__(self) -> None: """Initialize routing and queue state.""" self.number = 0 self.routes = FakeRoutes() self.prompt_queue = FakePromptQueue() def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]: """Return the prompt unchanged.""" return json_data class FakeRequest: """Minimal aiohttp request double.""" async def json(self) -> dict[str, Any]: """Return one ordinary prompt request.""" return { "prompt_id": "prompt-no-modal", "prompt": {"1": {"class_type": "LocalImage", "inputs": {}}}, "extra_data": {"extra_pnginfo": {"workflow": {"nodes": []}}}, } class FakeExecutionModule: """Minimal execution module exposing prompt validation.""" SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = () @staticmethod async def validate_prompt( prompt_id: str, prompt: dict[str, Any], partial_execution_targets: Any, ) -> tuple[bool, None, list[str], list[Any]]: """Accept the supplied prompt with one fake execution target.""" return True, None, ["1"], [] def fail_modal_status(*_args: Any, **_kwargs: Any) -> None: """Fail if the no-remote fast path emits Modal UI state.""" raise AssertionError("no-remote prompts must not emit Modal status") def fail_rewrite(*_args: Any, **_kwargs: Any) -> tuple[dict[str, Any], Any]: """Fail if the no-remote fast path enters Modal prompt rewriting.""" raise AssertionError("no-remote prompts must not be rewritten for Modal") settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=False, 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, ) prompt_server = FakePromptServer() monkeypatch.setattr(api_intercept_module, "_ROUTE_REGISTERED", False) monkeypatch.setattr( api_intercept_module, "_get_server_module", lambda: SimpleNamespace(PromptServer=SimpleNamespace(instance=prompt_server)), ) monkeypatch.setattr(api_intercept_module, "_get_execution_module", lambda: FakeExecutionModule) monkeypatch.setattr(api_intercept_module, "_emit_modal_status", fail_modal_status) monkeypatch.setattr(api_intercept_module, "rewrite_prompt_for_modal", fail_rewrite) api_intercept_module.setup_modal_queue_route( prompt_server=prompt_server, sync_engine=object(), settings=settings, ) response = asyncio.run(prompt_server.routes.handlers["/modal/queue_prompt"](FakeRequest())) response_payload = json.loads(response.text) assert response_payload["prompt_id"] == "prompt-no-modal" assert "modal_remote_node_ids" not in response_payload assert len(prompt_server.prompt_queue.items) == 1 def test_rewrite_groups_connected_remote_nodes_into_single_proxy( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Connected remote nodes should collapse into one proxy-backed component.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteModel": _FakeRemoteModelNode, "RemoteSampler": _FakeRemoteSamplerNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteModel", "inputs": {"model_name": str(model_path)}, "_meta": {"title": "Model"}, }, "2": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}, "_meta": {"title": "Sampler"}, }, "3": { "class_type": "LocalSink", "inputs": {"latent": ["2", 0]}, "_meta": {"title": "Sink"}, }, } 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, extra_data={"extra_pnginfo": {"workflow": workflow}}, ) assert set(rewritten_prompt) == {"1", "3", _artifact_finalizer_node_id(summary)} rewritten_node = rewritten_prompt["1"] payload = rewritten_node["inputs"]["original_node_data"] assert rewritten_node["class_type"].startswith("ModalUniversalExecutor_") assert payload["payload_kind"] == "subgraph" assert "prompt_id" not in payload assert payload["component_node_ids"] == ["1", "2"] assert payload["subgraph_prompt"]["1"]["inputs"]["model_name"].startswith("/assets/") assert payload["execute_node_ids"] == ["2"] assert "requires_volume_reload" not in payload assert "volume_reload_marker" not in payload assert "uploaded_volume_paths" not in payload assert payload["terminate_container_on_error"] is True assert payload["boundary_inputs"] == [] assert payload["boundary_outputs"] == [ { "proxy_output_name": "2_latent", "node_id": "2", "output_index": 0, "io_type": "LATENT", "is_list": False, "preview_target_node_ids": [], } ] assert rewritten_prompt["3"]["inputs"]["latent"] == ["1", 0] assert summary.remote_node_ids == ["1", "2"] assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} assert len(summary.synced_assets) == 1 assert summary.synced_assets[0].uploaded is True def test_rewrite_anchors_terminal_artifact_only_remote_node( api_intercept_module: Any, modal_executor_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """A terminal remote side-effect node should remain executable through the finalizer.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteArtifactWriter": _FakeRemoteArtifactWriterNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = {"nodes": [{"id": 1, "properties": {"is_modal_remote": True}}]} prompt = { "1": { "class_type": "RemoteArtifactWriter", "inputs": {}, "_meta": {"title": "Remote Artifact Writer"}, } } 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, ) finalizer_node_id = _artifact_finalizer_node_id(summary) assert set(rewritten_prompt) == {"1", finalizer_node_id} payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["execute_node_ids"] == ["1"] assert payload["boundary_outputs"] == [] proxy_class_type = rewritten_prompt["1"]["class_type"] proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_class_type] proxy_schema = proxy_class.GET_SCHEMA() assert [output.io_type for output in proxy_schema.outputs] == ["BOOLEAN"] assert [output.display_name for output in proxy_schema.outputs] == [ modal_executor_module.MODAL_COMPONENT_COMPLETION_OUTPUT_NAME ] assert rewritten_prompt[finalizer_node_id] == { "class_type": modal_executor_module.MODAL_ARTIFACT_FINALIZER_NODE_ID, "inputs": {"components.component_0": ["1", 0]}, "_meta": {"title": "Modal Artifact Finalizer"}, } finalizer_class = fake_nodes_module.NODE_CLASS_MAPPINGS[ modal_executor_module.MODAL_ARTIFACT_FINALIZER_NODE_ID ] assert finalizer_class.GET_SCHEMA().is_output_node is True assert finalizer_class.OUTPUT_NODE is True finalized_inputs, _hidden_inputs, _v3_data = ( modal_executor_module.io.get_finalized_class_inputs( finalizer_class.INPUT_TYPES(), rewritten_prompt[finalizer_node_id]["inputs"], ) ) assert "components.component_0" in finalized_inputs["required"] def test_rewrite_strips_prompt_id_from_cache_safe_proxy_payload( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Cache-safe remote proxies should not bake prompt_id into original_node_data inputs.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteModel": _FakeRemoteModelNode, "RemoteSampler": _FakeRemoteSamplerNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteModel", "inputs": {}, "_meta": {"title": "Model"}, }, "2": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}, "_meta": {"title": "Sampler"}, }, "3": { "class_type": "LocalSink", "inputs": {"latent": ["2", 0]}, "_meta": {"title": "Sink"}, }, } 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, extra_data={"prompt_id": "prompt-1", "client_id": "client-1"}, ) payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert "prompt_id" not in payload def test_rewrite_records_local_preview_targets_for_remote_boundary_images( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Boundary IMAGE outputs should remember direct local PreviewImage consumers.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteImage": _FakeRemoteImageNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 9, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote Image"}, }, "9": { "class_type": "PreviewImage", "inputs": {"images": ["1", 0]}, "_meta": {"title": "Preview"}, }, } 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, ) payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["boundary_outputs"] == [ { "proxy_output_name": "1_image", "node_id": "1", "output_index": 0, "io_type": "IMAGE", "is_list": False, "preview_target_node_ids": ["9"], } ] def test_rewrite_colocates_remote_chain_across_large_transportable_edges( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Large transportable remote-to-remote values should remain inside one component.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteImage": _FakeRemoteImageNode, "RemoteImageConsumer": _FakeRemoteImageConsumerNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote Image"}, }, "2": { "class_type": "RemoteImageConsumer", "inputs": {"image": ["1", 0]}, "_meta": {"title": "Remote Image Consumer"}, }, "3": { "class_type": "LocalSink", "inputs": {"image": ["2", 0]}, "_meta": {"title": "Local Sink"}, }, } 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, ) assert set(rewritten_prompt) == {"1", "3", _artifact_finalizer_node_id(summary)} assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["component_node_ids"] == ["1", "2"] assert payload["boundary_inputs"] == [] assert payload["boundary_outputs"] == [ { "proxy_output_name": "2_image", "node_id": "2", "output_index": 0, "io_type": "IMAGE", "is_list": False, "preview_target_node_ids": [], } ] assert payload["execute_node_ids"] == ["2"] assert rewritten_prompt["3"]["inputs"]["image"] == ["1", 0] assert api_intercept_module._is_inexpensive_remote_boundary_type("IMAGE") is False assert api_intercept_module._is_inexpensive_remote_boundary_type("STRING") is True def test_rewrite_colocates_sampler_decoders_and_video_assembly( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Sampler media branches should stay together through final video assembly.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteModel": _FakeRemoteModelNode, "RemoteSampler": _FakeRemoteSamplerNode, "VAELoader": _FakeVAELoaderNode, "VAEDecode": _FakeVAEDecodeNode, "VAEDecodeAudio": _FakeRemoteAudioNode, "CreateVideo": _FakeRemoteVideoNode, "SaveVideo": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": node_id, "properties": {"is_modal_remote": node_id != 8}} for node_id in range(1, 9) ] } prompt = { "1": {"class_type": "RemoteModel", "inputs": {}}, "2": {"class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}}, "3": {"class_type": "VAELoader", "inputs": {}}, "4": { "class_type": "VAEDecode", "inputs": {"samples": ["2", 0], "vae": ["3", 0]}, }, "5": {"class_type": "VAELoader", "inputs": {}}, "6": { "class_type": "VAEDecodeAudio", "inputs": {"samples": ["2", 0], "vae": ["5", 0]}, }, "7": { "class_type": "CreateVideo", "inputs": {"images": ["4", 0], "audio": ["6", 0]}, }, "8": {"class_type": "SaveVideo", "inputs": {"video": ["7", 0]}}, } 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, ) assert set(rewritten_prompt) == {"1", "8", _artifact_finalizer_node_id(summary)} assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == { "1": ["1", "2", "3", "4", "5", "6", "7"], } assert summary.rewritten_node_id_map == { str(node_id): "1" for node_id in range(1, 8) } payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["component_node_ids"] == ["1", "2", "3", "4", "5", "6", "7"] assert payload["execute_node_ids"] == ["7"] assert payload["boundary_inputs"] == [] assert payload["boundary_outputs"] == [ { "proxy_output_name": "7_video", "node_id": "7", "output_index": 0, "io_type": "VIDEO", "is_list": False, "preview_target_node_ids": [], } ] assert rewritten_prompt["8"]["inputs"]["video"] == ["1", 0] def test_rewrite_keeps_non_returning_local_preview_taps_local( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Local preview branches should stay local even when a remote chain continues.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteImage": _FakeRemoteImageNode, "RemoteImageConsumer": _FakeRemoteImageConsumerNode, "LocalSink": _FakeLocalSinkNode, "PreviewImage": _FakePreviewImageNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, {"id": 9, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote Image"}, }, "2": { "class_type": "RemoteImageConsumer", "inputs": {"image": ["1", 0]}, "_meta": {"title": "Remote Image Consumer"}, }, "3": { "class_type": "LocalSink", "inputs": {"image": ["2", 0]}, "_meta": {"title": "Local Sink"}, }, "9": { "class_type": "PreviewImage", "inputs": {"images": ["1", 0]}, "_meta": {"title": "Interim Preview"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "3", "9", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} remote_payloads = [ rewritten_node["inputs"]["original_node_data"] for rewritten_node in rewritten_prompt.values() if isinstance(rewritten_node.get("inputs"), dict) and "original_node_data" in rewritten_node["inputs"] ] assert remote_payloads assert all("9" not in payload["component_node_ids"] for payload in remote_payloads) assert all("9" not in payload["subgraph_prompt"] for payload in remote_payloads) assert all("9" not in payload["execute_node_ids"] for payload in remote_payloads) payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["boundary_outputs"] == [ { "proxy_output_name": "1_image", "node_id": "1", "output_index": 0, "io_type": "IMAGE", "is_list": False, "preview_target_node_ids": ["9"], }, { "proxy_output_name": "2_image", "node_id": "2", "output_index": 0, "io_type": "IMAGE", "is_list": False, "preview_target_node_ids": [], }, ] assert rewritten_prompt["9"]["inputs"]["images"] == ["1", 0] assert rewritten_prompt["3"]["inputs"]["image"] == ["1", 1] def test_rewrite_keeps_unmarked_preview_subgraph_nodes_local( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Unmarked preview producer nodes must not execute remotely.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteSampler": _FakeRemoteSamplerNode, "LocalSink": _FakeLocalSinkNode, "PreviewImage": _FakePreviewImageNode, "VAEDecode": _FakeVAEDecodeNode, "VAEEncode": _FakeVAEEncodeNode, "VAELoader": _FakeVAELoaderNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, {"id": 7, "properties": {"is_modal_remote": False}}, {"id": 8, "properties": {"is_modal_remote": False}}, {"id": 9, "properties": {"is_modal_remote": False}}, {"id": 11, "properties": {"is_modal_remote": False}}, {"id": 90, "properties": {"is_modal_remote": False}}, {"id": 192, "properties": {"is_modal_remote": False}}, ] } prompt = { "9": { "class_type": "VAELoader", "inputs": {"vae_name": "vae.safetensors"}, "_meta": {"title": "VAE Loader"}, }, "1": { "class_type": "RemoteSampler", "inputs": {}, "_meta": {"title": "Remote Sampler 1"}, }, "192": { "class_type": "VAEDecode", "inputs": {"samples": ["1", 0], "vae": ["9", 0]}, "_meta": {"title": "VAE Decode Preview"}, }, "8": { "class_type": "VAEEncode", "inputs": {"pixels": ["192", 0], "vae": ["9", 0]}, "_meta": {"title": "Local VAE Encode"}, }, "7": { "class_type": "LocalSink", "inputs": {"image": ["8", 0]}, "_meta": {"title": "Local Encoded Sink"}, }, "11": { "class_type": "VAEDecode", "inputs": {"samples": ["1", 0], "vae": ["9", 0]}, "_meta": {"title": "Local VAE Decode"}, }, "90": { "class_type": "PreviewImage", "inputs": {"images": ["192", 0]}, "_meta": {"title": "Preview"}, }, "2": { "class_type": "RemoteSampler", "inputs": {"latent": ["1", 0]}, "_meta": {"title": "Remote Sampler 2"}, }, "3": { "class_type": "LocalSink", "inputs": {"image": ["2", 0]}, "_meta": {"title": "Local Sink"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "3", "7", "8", "9", "11", "90", "192", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1"} remote_payloads = [ rewritten_node["inputs"]["original_node_data"] for rewritten_node in rewritten_prompt.values() if isinstance(rewritten_node.get("inputs"), dict) and "original_node_data" in rewritten_node["inputs"] ] local_node_ids = {"7", "8", "9", "11", "90", "192"} assert len(remote_payloads) == 1 for payload in remote_payloads: assert not (local_node_ids & set(payload["component_node_ids"])) assert not (local_node_ids & set(payload["subgraph_prompt"])) assert not (local_node_ids & set(payload["execute_node_ids"])) assert rewritten_prompt["90"]["inputs"]["images"] == ["192", 0] assert rewritten_prompt["3"]["inputs"]["image"] == ["1", 1] def test_rewrite_keeps_local_branches_that_feed_remote_as_boundaries( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Local branches that later feed remote work are dependencies, not preview taps.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteImage": _FakeRemoteImageNode, "RemoteImageConsumer": _FakeRemoteImageConsumerNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote Image"}, }, "4": { "class_type": "LocalSink", "inputs": {"image": ["1", 0]}, "_meta": {"title": "Local Transform"}, }, "2": { "class_type": "RemoteImageConsumer", "inputs": {"image": ["4", 0]}, "_meta": {"title": "Remote Image Consumer"}, }, "3": { "class_type": "LocalSink", "inputs": {"image": ["2", 0]}, "_meta": {"title": "Local Sink"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "2", "3", "4", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["1", "2"] assert summary.component_node_ids_by_representative == {"1": ["1"], "2": ["2"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "2"} assert summary.sandwiched_local_node_ids == ["4"] assert rewritten_prompt["4"]["inputs"]["image"] == ["1", 0] assert rewritten_prompt["2"]["inputs"]["remote_input_0"] == ["4", 0] assert rewritten_prompt["3"]["inputs"]["image"] == ["2", 0] second_payload = rewritten_prompt["2"]["inputs"]["original_node_data"] assert second_payload["component_node_ids"] == ["2"] assert second_payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_0", "io_type": "IMAGE", "targets": [{"node_id": "2", "input_name": "image"}], } ] def test_component_local_reentry_dependency_detection( api_intercept_module: Any, ) -> None: """Boundary inputs that trace back to the same component require a split-capable proxy.""" prompt = { "1": { "class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote Image"}, }, "4": { "class_type": "LocalSink", "inputs": {"image": ["1", 0]}, "_meta": {"title": "Local Transform"}, }, "2": { "class_type": "RemoteImageConsumer", "inputs": {"image": ["4", 0]}, "_meta": {"title": "Remote Consumer"}, }, } component = api_intercept_module.RemoteComponentPlan( node_ids=["1", "2"], representative_node_id="1", boundary_inputs=[ api_intercept_module.BoundaryInputSpec( proxy_input_name="remote_input_0", source=api_intercept_module.LinkedOutputRef(node_id="4", output_index=0), io_type="IMAGE", targets=[ api_intercept_module.InputTarget( node_id="2", input_name="image", ) ], ) ], boundary_outputs=[], execute_node_ids=["1", "2"], contains_output_node=False, local_tap_node_ids=["9"], ) assert api_intercept_module._component_has_local_reentry_dependency( prompt=prompt, component=component, ) def test_sandwiched_local_nodes_include_only_remote_reentry_paths( api_intercept_module: Any, ) -> None: """Planner warnings should cover local chains that leave and re-enter remote work.""" prompt = { "1": {"class_type": "RemoteSource", "inputs": {}}, "2": {"class_type": "LocalTransform", "inputs": {"value": ["1", 0]}}, "3": { "class_type": "LocalTransform", "inputs": {"value": ["2", 0], "local_only": ["7", 0]}, }, "4": {"class_type": "RemoteSink", "inputs": {"value": ["3", 0]}}, "5": {"class_type": "LocalPreview", "inputs": {"value": ["1", 0]}}, "6": {"class_type": "RemoteSink", "inputs": {"value": ["8", 0]}}, "7": {"class_type": "LocalSource", "inputs": {}}, "8": {"class_type": "LocalSource", "inputs": {}}, "9": {"class_type": "RemoteSink", "inputs": {"value": ["1", 1]}}, } assert api_intercept_module._sandwiched_local_node_ids( prompt, {"1", "4", "6", "9"}, ) == {"2", "3"} def test_rewrite_reports_parallel_component_stages( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Prompt rewrites should report best-effort concurrent stages for independent remote components.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteImage": _FakeRemoteImageNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": {"class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote A"}}, "2": {"class_type": "RemoteImage", "inputs": {}, "_meta": {"title": "Remote B"}}, "3": {"class_type": "LocalSink", "inputs": {"image": ["1", 0]}, "_meta": {"title": "Sink A"}}, "4": {"class_type": "LocalSink", "inputs": {"image": ["2", 0]}, "_meta": {"title": "Sink B"}}, } _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, ) assert summary.component_execution_stages == [["1", "2"]] assert summary.component_dependency_ids_by_representative == {"1": [], "2": []} assert summary.mapped_component_ids == [] assert summary.estimated_max_parallel_requests == 2 assert summary.max_parallel_requests_upper_bound == 2 def test_rewrite_reports_mapped_parallelism_upper_bound( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Mapped components should warm only the single container needed for one in-process mapped run.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", max_containers=5, ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "PromptList": _FakePromptListNode, "ModalMapInput": _FakeModalMapInputNode, "RemoteStringEcho": _FakeRemoteStringEchoNode, "LocalStringSink": _FakeLocalStringSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": {"class_type": "PromptList", "inputs": {}, "_meta": {"title": "Prompt List"}}, "2": {"class_type": "ModalMapInput", "inputs": {"value": ["1", 0]}, "_meta": {"title": "Map"}}, "3": {"class_type": "RemoteStringEcho", "inputs": {"text": ["2", 0]}, "_meta": {"title": "Echo"}}, "4": {"class_type": "LocalStringSink", "inputs": {"text": ["3", 0]}, "_meta": {"title": "Sink"}}, } _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, ) assert summary.component_execution_stages == [["2"]] assert summary.mapped_component_ids == ["2"] assert summary.estimated_max_parallel_requests == 1 assert summary.max_parallel_requests_upper_bound == 1 def test_rewrite_uses_one_request_wide_volume_reload_marker_across_components( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """All components in one rewritten prompt should share one reload marker and decision.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteString": _FakeRemoteStringEchoNode, "RemoteStringConsumer": _FakeRemoteStringEchoNode, "LocalSink": _FakeLocalStringSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteString", "inputs": {}, "_meta": {"title": "Remote String"}, }, "2": { "class_type": "RemoteStringConsumer", "inputs": {"text": ["1", 0]}, "_meta": {"title": "Remote String Consumer"}, }, "3": { "class_type": "LocalSink", "inputs": {"text": ["2", 0]}, "_meta": {"title": "Local Sink"}, }, } uploaded_asset = sync_engine_module.SyncedAsset( local_path=tmp_path / "uploaded.bin", remote_path="/assets/uploaded.bin", sha256="uploaded", uploaded=True, ) def fake_sync_component_prompt_inputs( *, component: Any, rewritten_prompt: dict[str, Any], sync_engine: Any, request_cache: Any, status_callback: Any = None, ) -> tuple[dict[str, Any], list[Any]]: del sync_engine, request_cache, status_callback if component.representative_node_id == "1": return {"1": rewritten_prompt["1"]}, [] return { "2": { "class_type": rewritten_prompt["2"]["class_type"], "inputs": {"text_path": uploaded_asset.remote_path}, "_meta": rewritten_prompt["2"]["_meta"], } }, [uploaded_asset] monkeypatch.setattr( api_intercept_module, "_sync_component_prompt_inputs", fake_sync_component_prompt_inputs, ) 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, ) first_payload = rewritten_prompt["1"]["inputs"]["original_node_data"] second_payload = rewritten_prompt["2"]["inputs"]["original_node_data"] assert summary.remote_component_ids == ["1", "2"] assert summary.synced_assets == [uploaded_asset] assert "requires_volume_reload" not in first_payload assert "requires_volume_reload" not in second_payload assert "volume_reload_marker" not in first_payload assert "volume_reload_marker" not in second_payload assert "uploaded_volume_paths" not in first_payload assert "uploaded_volume_paths" not in second_payload assert summary.requires_volume_reload is True assert isinstance(summary.volume_reload_marker, str) assert summary.volume_reload_marker assert summary.uploaded_volume_paths == [uploaded_asset.remote_path] def test_rewrite_merges_cyclic_coarse_components_back_into_single_proxy( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """A cyclic quotient between coarse groups should collapse back into one remote proxy.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteModelAndImage": _FakeRemoteModelAndImageNode, "RemoteImageConsumer": _FakeRemoteImageConsumerNode, "RemoteModelAndImageConsumer": _FakeRemoteModelAndImageConsumerNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteModelAndImage", "inputs": {}, "_meta": {"title": "Remote Model And Image"}, }, "2": { "class_type": "RemoteImageConsumer", "inputs": {"image": ["1", 1]}, "_meta": {"title": "Remote Image Consumer"}, }, "3": { "class_type": "RemoteModelAndImageConsumer", "inputs": {"model": ["1", 0], "image": ["2", 0]}, "_meta": {"title": "Remote Model And Image Consumer"}, }, "4": { "class_type": "PreviewImage", "inputs": {"images": ["3", 0]}, "_meta": {"title": "Preview"}, }, } 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, ) assert set(rewritten_prompt) == {"1", "4", _artifact_finalizer_node_id(summary)} assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2", "3"]} payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["component_node_ids"] == ["1", "2", "3"] assert payload["boundary_inputs"] == [] assert payload["boundary_outputs"] == [ { "proxy_output_name": "3_image", "node_id": "3", "output_index": 0, "io_type": "IMAGE", "is_list": False, "preview_target_node_ids": ["4"], } ] assert rewritten_prompt["4"]["inputs"]["images"] == ["1", 0] def test_rewrite_marks_modal_map_boundary_as_mapped_subgraph( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """A remote component fed through ModalMapInput should rewrite to a mapped payload.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "PromptList": _FakePromptListNode, "ModalMapInput": _FakeModalMapInputNode, "RemoteStringEcho": _FakeRemoteStringEchoNode, "LocalStringSink": _FakeLocalStringSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 5, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "PromptList", "inputs": {}, "_meta": {"title": "Prompt List"}, }, "2": { "class_type": "ModalMapInput", "inputs": {"value": ["1", 0]}, "_meta": {"title": "Map Input"}, }, "3": { "class_type": "RemoteStringEcho", "inputs": {"text": ["2", 0]}, "_meta": {"title": "Remote Echo"}, }, "5": { "class_type": "RemoteStringEcho", "inputs": {"text": ["3", 0]}, "_meta": {"title": "Remote Echo 2"}, }, "4": { "class_type": "LocalStringSink", "inputs": {"text": ["5", 0]}, "_meta": {"title": "Local Sink"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "2", "4", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["2"] payload = rewritten_prompt["2"]["inputs"]["original_node_data"] assert payload["payload_kind"] == "mapped_subgraph" assert payload["component_node_ids"] == ["2", "3", "5"] assert payload["mapped_input"] == { "proxy_input_name": "remote_input_0", "io_type": "STRING", } assert payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_0", "io_type": "STRING", "targets": [{"node_id": "2", "input_name": "value"}], } ] assert payload["static_to_mapped_boundaries"] == [] assert payload["static_phase"] == { "component_node_ids": [], "subgraph_prompt": {}, "boundary_inputs": [], "boundary_outputs": [], "execute_node_ids": [], } assert payload["mapped_phase"] == { "component_node_ids": ["2", "3", "5"], "subgraph_prompt": { "2": prompt["2"], "3": prompt["3"], "5": prompt["5"], }, "boundary_inputs": [ { "proxy_input_name": "remote_input_0", "io_type": "STRING", "targets": [{"node_id": "2", "input_name": "value"}], } ], "boundary_outputs": [ { "proxy_output_name": "5_text", "node_id": "5", "output_index": 0, "io_type": "STRING", "is_list": False, "preview_target_node_ids": [], "mapped_output": True, "scheduler_is_list": True, } ], "execute_node_ids": ["5"], } assert rewritten_prompt["4"]["inputs"]["text"] == ["2", 0] def test_rewrite_marks_local_modal_map_source_as_mapped_subgraph( api_intercept_module: Any, modal_executor_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """A local ModalMapInput feeding a remote node should still rewrite to mapped remote execution.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) settings.custom_nodes_dir.mkdir() sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "PromptList": _FakePromptListNode, "ModalMapInput": _FakeModalMapInputNode, "RemoteStringEcho": _FakeRemoteStringEchoNode, "LocalStringSink": _FakeLocalStringSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "PromptList", "inputs": {}, "_meta": {"title": "Prompt List"}, }, "2": { "class_type": "ModalMapInput", "inputs": {"value": ["1", 0]}, "_meta": {"title": "Map Input"}, }, "3": { "class_type": "RemoteStringEcho", "inputs": {"text": ["2", 0]}, "_meta": {"title": "Remote Echo"}, }, "4": { "class_type": "LocalStringSink", "inputs": {"text": ["3", 0]}, "_meta": {"title": "Local Sink"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "2", "3", "4", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["3"] payload = rewritten_prompt["3"]["inputs"]["original_node_data"] assert payload["payload_kind"] == "mapped_subgraph" assert payload["component_node_ids"] == ["3"] assert payload["mapped_input"] == { "proxy_input_name": "remote_input_0", "io_type": "STRING", } assert payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_0", "io_type": "*", "targets": [{"node_id": "3", "input_name": "text"}], } ] assert payload["mapped_phase"] == { "component_node_ids": ["3"], "subgraph_prompt": { "3": prompt["3"], }, "boundary_inputs": [ { "proxy_input_name": "remote_input_0", "io_type": "*", "targets": [{"node_id": "3", "input_name": "text"}], } ], "boundary_outputs": [ { "proxy_output_name": "3_text", "node_id": "3", "output_index": 0, "io_type": "STRING", "is_list": False, "preview_target_node_ids": [], "mapped_output": True, "scheduler_is_list": True, } ], "execute_node_ids": ["3"], } with modal_executor_module._MODAL_MAP_WARMUP_CONTEXTS_LOCK: warmup_context = modal_executor_module._MODAL_MAP_WARMUP_CONTEXTS["2"] assert warmup_context.mapped_io_type == "STRING" assert warmup_context.execution_payload["component_id"] == "3" assert rewritten_prompt["4"]["inputs"]["text"] == ["3", 0] def test_rewrite_supports_mapped_branch_that_shares_non_transportable_upstream_with_unmapped_sibling( api_intercept_module: Any, modal_executor_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Mapped execution should separate static and per-item execute targets within one coarse component.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteModel": _FakeRemoteModelNode, "RemoteSampler": _FakeRemoteSamplerNode, "LatentSource": _FakeLatentSourceNode, "ModalMapInput": _FakeModalMapInputNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, {"id": 5, "properties": {"is_modal_remote": False}}, {"id": 6, "properties": {"is_modal_remote": True}}, {"id": 7, "properties": {"is_modal_remote": True}}, {"id": 8, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteModel", "inputs": {}, "_meta": {"title": "Shared Model"}, }, "2": { "class_type": "LatentSource", "inputs": {}, "_meta": {"title": "Single Latent"}, }, "3": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0], "latent": ["2", 0]}, "_meta": {"title": "Unmapped Sampler"}, }, "4": { "class_type": "LocalSink", "inputs": {"image": ["3", 0]}, "_meta": {"title": "Local Sink 1"}, }, "5": { "class_type": "LatentSource", "inputs": {}, "_meta": {"title": "Batch Latent Source"}, }, "6": { "class_type": "ModalMapInput", "inputs": {"value": ["5", 0]}, "_meta": {"title": "Map Input"}, }, "7": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0], "latent": ["6", 0]}, "_meta": {"title": "Mapped Sampler"}, }, "8": { "class_type": "LocalSink", "inputs": {"image": ["7", 0]}, "_meta": {"title": "Local Sink 2"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "1__mapped", "2", "4", "5", "8", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["1", "1__mapped"] assert summary.component_node_ids_by_representative == { "1": ["1", "3"], "1__mapped": ["6", "7"], } assert summary.component_dependency_ids_by_representative == { "1": [], "1__mapped": ["1"], } assert summary.component_execution_stages == [["1"], ["1__mapped"]] assert summary.mapped_component_ids == ["1__mapped"] static_payload = rewritten_prompt["1"]["inputs"]["original_node_data"] mapped_payload = rewritten_prompt["1__mapped"]["inputs"]["original_node_data"] static_execution_payload = modal_executor_module._rehydrate_proxy_payload( static_payload, unique_id="1", ) mapped_execution_payload = modal_executor_module._rehydrate_proxy_payload( mapped_payload, unique_id="1__mapped", ) assert static_payload["payload_kind"] == "subgraph" assert static_payload["component_node_ids"] == ["1", "3"] assert static_payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_0", "io_type": "LATENT", "targets": [{"node_id": "3", "input_name": "latent"}], } ] assert static_payload["boundary_outputs"] == [ { "proxy_output_name": "3_latent", "node_id": "3", "output_index": 0, "io_type": "LATENT", "is_list": False, "preview_target_node_ids": [], }, { "proxy_output_name": "static_input_0", "node_id": "1", "output_index": 0, "io_type": "MODEL", "is_list": False, "preview_target_node_ids": [], "session_output": True, }, ] assert static_payload["execute_node_ids"] == ["1", "3"] assert "remote_session" not in static_payload assert static_execution_payload["remote_session"]["owner_component_id"] == "1" assert mapped_payload["payload_kind"] == "subgraph" assert mapped_payload["component_node_ids"] == ["6", "7"] assert mapped_payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_1", "io_type": "LATENT", "targets": [{"node_id": "6", "input_name": "value"}], }, { "proxy_input_name": "static_input_0", "io_type": "MODEL", "targets": [{"node_id": "7", "input_name": "model"}], "source_signature": api_intercept_module._boundary_source_signature( prompt, api_intercept_module.LinkedOutputRef(node_id="1", output_index=0), ), }, ] assert mapped_payload["boundary_outputs"] == [ { "proxy_output_name": "7_latent", "node_id": "7", "output_index": 0, "io_type": "LATENT", "is_list": False, "preview_target_node_ids": [], } ] assert mapped_payload["execute_node_ids"] == ["7"] assert "clear_remote_session" not in mapped_payload assert mapped_payload["mapped_progress_display_node_id"] == "1" assert mapped_execution_payload["clear_remote_session"] is True assert ( mapped_execution_payload["remote_session"]["session_id"] == static_execution_payload["remote_session"]["session_id"] ) assert rewritten_prompt["1__mapped"]["inputs"]["remote_input_0"] == ["2", 0] assert rewritten_prompt["1__mapped"]["inputs"]["static_input_0"] == ["1", 1] assert rewritten_prompt["4"]["inputs"]["image"] == ["1", 0] assert rewritten_prompt["8"]["inputs"]["image"] == ["1__mapped", 0] def test_rewrite_stamps_snapshot_profile_on_split_static_and_mapped_payloads( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Split static and mapped payloads should inherit the same loader snapshot profile.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=True, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", modal_gpu="L40S", ) sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "CheckpointLoaderSimple": _FakeCheckpointLoaderSimpleNode, "RemoteSampler": _FakeRemoteSamplerNode, "LatentSource": _FakeLatentSourceNode, "ModalMapInput": _FakeModalMapInputNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, {"id": 5, "properties": {"is_modal_remote": False}}, {"id": 6, "properties": {"is_modal_remote": True}}, {"id": 7, "properties": {"is_modal_remote": True}}, {"id": 8, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "base.safetensors"}, "_meta": {"title": "Checkpoint"}, }, "2": { "class_type": "LatentSource", "inputs": {}, "_meta": {"title": "Single Latent"}, }, "3": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0], "latent": ["2", 0]}, "_meta": {"title": "Unmapped Sampler"}, }, "4": { "class_type": "LocalSink", "inputs": {"image": ["3", 0]}, "_meta": {"title": "Local Sink 1"}, }, "5": { "class_type": "LatentSource", "inputs": {}, "_meta": {"title": "Batch Latent Source"}, }, "6": { "class_type": "ModalMapInput", "inputs": {"value": ["5", 0]}, "_meta": {"title": "Map Input"}, }, "7": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0], "latent": ["6", 0]}, "_meta": {"title": "Mapped Sampler"}, }, "8": { "class_type": "LocalSink", "inputs": {"image": ["7", 0]}, "_meta": {"title": "Local Sink 2"}, }, } rewritten_prompt, _ = api_intercept_module.rewrite_prompt_for_modal( prompt=prompt, workflow=workflow, sync_engine=sync_engine, settings=settings, nodes_module=fake_nodes_module, ) static_payload = rewritten_prompt["1"]["inputs"]["original_node_data"] mapped_payload = rewritten_prompt["1__mapped"]["inputs"]["original_node_data"] assert static_payload["snapshot_profile_key"].startswith("loader-profile:") assert mapped_payload["snapshot_profile_key"] == static_payload["snapshot_profile_key"] assert static_payload["modal_gpu"] == "L40S" assert mapped_payload["modal_gpu"] == "L40S" def test_rewrite_keeps_unmapped_remote_siblings_without_local_reentry_together( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Ordinary remote execute siblings should remain one proxy without local re-entry.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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=tmp_path / "custom_nodes", ) sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteModel": _FakeRemoteModelNode, "RemoteSampler": _FakeRemoteSamplerNode, "LatentSource": _FakeLatentSourceNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, {"id": 5, "properties": {"is_modal_remote": False}}, {"id": 6, "properties": {"is_modal_remote": True}}, {"id": 7, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteModel", "inputs": {}, "_meta": {"title": "Shared Model"}, }, "2": { "class_type": "LatentSource", "inputs": {}, "_meta": {"title": "Latent A"}, }, "3": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0], "latent": ["2", 0]}, "_meta": {"title": "Sampler A"}, }, "4": { "class_type": "LocalSink", "inputs": {"image": ["3", 0]}, "_meta": {"title": "Sink A"}, }, "5": { "class_type": "LatentSource", "inputs": {}, "_meta": {"title": "Latent B"}, }, "6": { "class_type": "RemoteSampler", "inputs": {"model": ["1", 0], "latent": ["5", 0]}, "_meta": {"title": "Sampler B"}, }, "7": { "class_type": "LocalSink", "inputs": {"image": ["6", 0]}, "_meta": {"title": "Sink B"}, }, } 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, ) assert set(rewritten_prompt) == { "1", "2", "4", "5", "7", _artifact_finalizer_node_id(summary), } assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == { "1": ["1", "3", "6"], } assert summary.component_dependency_ids_by_representative == { "1": [], } assert summary.component_execution_stages == [["1"]] assert summary.rewritten_node_id_map == {"1": "1", "3": "1", "6": "1"} payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["payload_kind"] == "subgraph" assert payload["component_node_ids"] == ["1", "3", "6"] assert payload["boundary_inputs"] == [ { "proxy_input_name": "remote_input_0", "io_type": "LATENT", "targets": [{"node_id": "3", "input_name": "latent"}], }, { "proxy_input_name": "remote_input_1", "io_type": "LATENT", "targets": [{"node_id": "6", "input_name": "latent"}], }, ] assert payload["boundary_outputs"] == [ { "proxy_output_name": "3_latent", "node_id": "3", "output_index": 0, "io_type": "LATENT", "is_list": False, "preview_target_node_ids": [], }, { "proxy_output_name": "6_latent", "node_id": "6", "output_index": 0, "io_type": "LATENT", "is_list": False, "preview_target_node_ids": [], }, ] assert payload["execute_node_ids"] == ["3", "6"] assert "remote_session" not in payload assert rewritten_prompt["4"]["inputs"]["image"] == ["1", 0] assert rewritten_prompt["7"]["inputs"]["image"] == ["1", 1] def test_boundary_source_signature_changes_with_upstream_prompt_structure( api_intercept_module: Any, ) -> None: """Non-transportable boundary provenance should change when the upstream prompt changes.""" source = api_intercept_module.LinkedOutputRef(node_id="2", output_index=0) base_prompt = { "1": { "class_type": "CheckpointLoader", "inputs": {"ckpt_name": "base.safetensors"}, }, "2": { "class_type": "LoraLoader", "inputs": {"model": ["1", 0], "strength_model": 0.8}, }, } changed_prompt = { "1": { "class_type": "CheckpointLoader", "inputs": {"ckpt_name": "base.safetensors"}, }, "2": { "class_type": "LoraLoader", "inputs": {"model": ["1", 0], "strength_model": 0.5}, }, } first_signature = api_intercept_module._boundary_source_signature(base_prompt, source) second_signature = api_intercept_module._boundary_source_signature(base_prompt, source) changed_signature = api_intercept_module._boundary_source_signature(changed_prompt, source) assert first_signature == second_signature assert changed_signature != first_signature def test_extract_remote_node_ids_recurses_into_nested_subgraph_workflows( api_intercept_module: Any, settings_module: Any, ) -> None: """Modal marker extraction should find nodes nested inside saved subgraph metadata.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) workflow = { "nodes": [ { "id": 100, "properties": {"is_modal_remote": False}, "subgraph": { "nodes": [ {"id": 11, "properties": {"is_modal_remote": True}}, {"id": 12, "properties": {"is_modal_remote": False}}, ] }, } ] } assert api_intercept_module.extract_remote_node_ids(workflow, settings) == {"11"} assert api_intercept_module.extract_remote_node_ids( workflow, settings, prompt_node_ids={"100"}, ) == {"100"} assert api_intercept_module.extract_remote_node_ids( workflow, settings, prompt_node_ids={"100:11"}, ) == {"100:11"} def test_extract_remote_node_ids_maps_subgraph_container_to_descendant_prompt_nodes( api_intercept_module: Any, settings_module: Any, ) -> None: """A marked subgraph container should remote its expanded descendant prompt nodes.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) workflow = { "nodes": [ { "id": 24, "properties": {"is_modal_remote": True}, "subgraph": { "nodes": [ {"id": 23, "properties": {"is_modal_remote": False}}, {"id": 25, "properties": {"is_modal_remote": False}}, ] }, } ] } assert api_intercept_module.extract_remote_node_ids( workflow, settings, prompt_node_ids={"24:23", "24:25", "99"}, ) == {"24:23", "24:25"} def test_extract_remote_node_ids_maps_defined_subgraph_nodes_through_instances( api_intercept_module: Any, settings_module: Any, ) -> None: """Markers in reusable subgraph definitions should map to every executable instance path.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) subgraph_type = "4c314f31-ecda-4b08-ae98-faaba1bf613f" workflow = { "nodes": [ {"id": 105, "type": subgraph_type, "properties": {"is_modal_remote": False}}, {"id": 205, "type": subgraph_type, "properties": {"is_modal_remote": False}}, ], "definitions": { "subgraphs": [ { "id": subgraph_type, "nodes": [ {"id": 11, "type": "VAELoader", "properties": {"is_modal_remote": True}}, { "id": 14, "type": "SamplerCustomAdvanced", "properties": {"is_modal_remote": True}, }, {"id": 107, "type": "ComfyMathExpression", "properties": {}}, ], } ] }, } prompt_node_ids = { "105:11", "105:14", "105:107", "205:11", "205:14", "205:107", "300", } assert api_intercept_module.extract_remote_node_ids( workflow, settings, prompt_node_ids=prompt_node_ids, ) == {"105:11", "105:14", "205:11", "205:14"} assert api_intercept_module._extract_marked_workflow_node_paths( workflow, settings, ) == {"105:11", "105:14", "205:11", "205:14"} def test_extract_remote_node_ids_maps_nested_defined_subgraph_instances( api_intercept_module: Any, settings_module: Any, ) -> None: """Nested reusable definitions should retain every instance ancestor in prompt ids.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) outer_type = "outer-subgraph" inner_type = "inner-subgraph" workflow = { "nodes": [ {"id": 105, "type": outer_type, "properties": {"is_modal_remote": False}}, ], "definitions": { "subgraphs": [ { "id": outer_type, "nodes": [ {"id": 7, "type": inner_type, "properties": {"is_modal_remote": False}}, ], }, { "id": inner_type, "nodes": [ {"id": 11, "type": "VAELoader", "properties": {"is_modal_remote": True}}, ], }, ] }, } assert api_intercept_module.extract_remote_node_ids( workflow, settings, prompt_node_ids={"105:7:11", "300"}, ) == {"105:7:11"} def test_rewrite_rejects_non_transportable_remote_inputs( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Remote nodes should absorb a single non-transportable upstream dependency automatically.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteConsumer": _FakeRemoteSamplerNode, "ModelSource": _FakeRemoteModelNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": True}}, ] } prompt = { "1": { "class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model Source"}, }, "2": { "class_type": "RemoteConsumer", "inputs": {"model": ["1", 0]}, "_meta": {"title": "Remote Consumer"}, }, } 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, ) assert list(rewritten_prompt) == ["1", _artifact_finalizer_node_id(summary)] assert summary.remote_node_ids == ["1", "2"] assert summary.remote_component_ids == ["1"] def test_rewrite_detects_remote_marker_inside_nested_subgraph_workflow( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Prompt rewrite should honor Modal markers found inside nested subgraph metadata.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteConsumer": _FakeRemoteSamplerNode, "LocalConsumer": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ { "id": 99, "properties": {"is_modal_remote": False}, "subgraph": { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": True}}, ] }, }, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "99": { "class_type": "RemoteConsumer", "inputs": {}, "_meta": {"title": "Subgraph Container"}, }, "4": { "class_type": "LocalConsumer", "inputs": {"latent": ["99", 0]}, "_meta": {"title": "Local Consumer"}, }, } 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, ) assert list(rewritten_prompt) == ["99", "4", _artifact_finalizer_node_id(summary)] assert rewritten_prompt["4"]["inputs"]["latent"] == ["99", 0] assert summary.remote_node_ids == ["99"] assert summary.remote_component_ids == ["99"] def test_rewrite_detects_marked_inner_subgraph_prompt_node_ids( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """A marked nested workflow node should resolve to its composed prompt id.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteClip": _FakeRemoteClipNode, "RemoteConsumer": _FakeRemoteSamplerNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ { "id": 24, "properties": {"is_modal_remote": False}, "subgraph": { "nodes": [ {"id": 23, "properties": {"is_modal_remote": True}}, ] }, }, {"id": 30, "properties": {"is_modal_remote": True}}, ] } prompt = { "30": { "class_type": "RemoteClip", "inputs": {}, "_meta": {"title": "Remote VAE Source"}, }, "24:23": { "class_type": "RemoteConsumer", "inputs": {"clip": ["30", 0]}, "_meta": {"title": "Nested Remote Consumer"}, }, } 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, ) assert list(rewritten_prompt) == ["24:23", _artifact_finalizer_node_id(summary)] assert summary.remote_node_ids == ["24:23", "30"] assert summary.remote_component_ids == ["24:23"] def test_rewrite_auto_expands_upstream_non_transportable_dependencies( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Marked remote nodes should absorb upstream non-transportable producers automatically.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "ModelSource": _FakeRemoteModelNode, "ConditioningSource": _FakeRemoteConditioningNode, "RemoteConsumer": _FakeRemoteSamplerNode, "LocalConsumer": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": True}}, {"id": 4, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model Source"}, }, "2": { "class_type": "ConditioningSource", "inputs": {}, "_meta": {"title": "Conditioning Source"}, }, "3": { "class_type": "RemoteConsumer", "inputs": { "model": ["1", 0], "conditioning": ["2", 0], }, "_meta": {"title": "Remote Consumer"}, }, "4": { "class_type": "LocalConsumer", "inputs": {"latent": ["3", 0]}, "_meta": {"title": "Local Consumer"}, }, } 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, ) assert set(rewritten_prompt) == {"1", "4", _artifact_finalizer_node_id(summary)} assert summary.remote_node_ids == ["1", "2", "3"] assert summary.remote_component_ids == ["1"] assert summary.component_node_ids_by_representative == {"1": ["1", "2", "3"]} assert summary.rewritten_node_id_map == {"1": "1", "2": "1", "3": "1"} payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["boundary_inputs"] == [] assert payload["execute_node_ids"] == ["3"] assert rewritten_prompt["4"]["inputs"]["latent"] == ["1", 0] def test_analyze_remote_node_selection_returns_nodes_to_mark_and_reasons( api_intercept_module: Any, settings_module: Any, ) -> None: """Dry-run analysis should surface the clicked node plus required upstream nodes.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "ModelSource": _FakeRemoteModelNode, "ConditioningSource": _FakeRemoteConditioningNode, "RemoteConsumer": _FakeRemoteSamplerNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": False}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": {"class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model"}}, "2": { "class_type": "ConditioningSource", "inputs": {}, "_meta": {"title": "Conditioning"}, }, "3": { "class_type": "RemoteConsumer", "inputs": {"model": ["1", 0], "conditioning": ["2", 0]}, "_meta": {"title": "Remote Consumer"}, }, } analysis = api_intercept_module.analyze_remote_node_selection( prompt=prompt, workflow=workflow, seed_workflow_node_paths=["3"], settings=settings, nodes_module=fake_nodes_module, ) assert analysis.requested_node_ids == ["3"] assert analysis.requested_workflow_node_paths == ["3"] assert analysis.current_remote_node_ids == [] assert analysis.current_remote_workflow_node_paths == [] assert analysis.resolved_remote_node_ids == ["1", "2", "3"] assert analysis.resolved_workflow_node_paths == ["1", "2", "3"] assert analysis.added_node_ids == ["1", "2", "3"] assert analysis.added_workflow_node_paths == ["1", "2", "3"] assert [(reason.node_id, reason.required_by_node_id) for reason in analysis.reasons] == [ ("1", "3"), ("2", "3"), ] def test_analyze_remote_node_selection_reports_local_reentry( api_intercept_module: Any, settings_module: Any, ) -> None: """Dry-run analysis should report local nodes between existing remote regions.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "RemoteString": _FakeRemoteStringEchoNode, "LocalString": _FakeLocalStringSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": False}}, {"id": 3, "properties": {"is_modal_remote": True}}, ] } prompt = { "1": {"class_type": "RemoteString", "inputs": {}}, "2": {"class_type": "LocalString", "inputs": {"text": ["1", 0]}}, "3": {"class_type": "RemoteString", "inputs": {"text": ["2", 0]}}, } analysis = api_intercept_module.analyze_remote_node_selection( prompt=prompt, workflow=workflow, seed_workflow_node_paths=[], settings=settings, nodes_module=fake_nodes_module, ) assert analysis.resolved_remote_node_ids == ["1", "3"] assert analysis.sandwiched_local_node_ids == ["2"] def test_analyze_remote_node_selection_prefers_nested_workflow_paths( api_intercept_module: Any, settings_module: Any, ) -> None: """Nested prompt ids should map back to the specific inner workflow node path.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": { "ModelSource": _FakeRemoteModelNode, "RemoteConsumer": _FakeRemoteSamplerNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ { "id": 24, "properties": {"is_modal_remote": False}, "subgraph": { "nodes": [ {"id": 23, "properties": {"is_modal_remote": True}}, ] }, }, {"id": 30, "properties": {"is_modal_remote": False}}, ] } prompt = { "30": {"class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model"}}, "24:23": { "class_type": "RemoteConsumer", "inputs": {"model": ["30", 0]}, "_meta": {"title": "Nested Consumer"}, }, } analysis = api_intercept_module.analyze_remote_node_selection( prompt=prompt, workflow=workflow, seed_workflow_node_paths=["24:23"], settings=settings, nodes_module=fake_nodes_module, ) assert analysis.requested_node_ids == ["24:23"] assert analysis.current_remote_node_ids == ["24:23"] assert analysis.current_remote_workflow_node_paths == ["24:23"] assert analysis.resolved_remote_node_ids == ["24:23", "30"] assert analysis.resolved_workflow_node_paths == ["24:23", "30"] assert analysis.added_node_ids == ["30"] assert analysis.added_workflow_node_paths == ["30"] assert [(reason.node_id, reason.required_by_node_id) for reason in analysis.reasons] == [ ("30", "24:23"), ] def test_rewrite_rejects_non_transportable_remote_outputs( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Remote component boundaries should reject non-transportable local downstream edges.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteClip": _FakeRemoteClipNode, "LocalConsumer": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "RemoteClip", "inputs": {}, "_meta": {"title": "Remote Clip"}, }, "2": { "class_type": "LocalConsumer", "inputs": {"clip": ["1", 0]}, "_meta": {"title": "Local Consumer"}, }, } try: api_intercept_module.rewrite_prompt_for_modal( prompt=prompt, workflow=workflow, sync_engine=sync_engine, settings=settings, nodes_module=fake_nodes_module, ) except api_intercept_module.ModalPromptValidationError as exc: message = str(exc) else: raise AssertionError("Expected ModalPromptValidationError to be raised.") assert "exports node 1 (RemoteClip) output index 0 of type 'CLIP'" in message assert "cannot cross the current local/remote boundary" in message def test_rewrite_allows_video_and_audio_across_remote_boundaries( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Current ComfyUI VIDEO and AUDIO values should pass boundary validation.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "CreateVideo": _FakeRemoteVideoNode, "SaveVideo": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": { "class_type": "CreateVideo", "inputs": {}, "_meta": {"title": "Create Video"}, }, "2": { "class_type": "SaveVideo", "inputs": {"video": ["1", 0]}, "_meta": {"title": "Save Video"}, }, } 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, ) payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["boundary_outputs"][0]["io_type"] == "VIDEO" assert rewritten_prompt["2"]["inputs"]["video"] == ["1", 0] assert summary.remote_node_ids == ["1"] assert api_intercept_module._is_transportable_output_type("AUDIO") is True def test_rewrite_keeps_remote_noise_producer_with_remote_sampler( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """NOISE strategy objects should remain inside one remote component.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RandomNoise": _FakeRemoteNoiseNode, "RemoteSampler": _FakeRemoteSamplerNode, "LocalSink": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 1, "properties": {"is_modal_remote": True}}, {"id": 2, "properties": {"is_modal_remote": True}}, {"id": 3, "properties": {"is_modal_remote": False}}, ] } prompt = { "1": {"class_type": "RandomNoise", "inputs": {"noise_seed": 42}}, "2": { "class_type": "RemoteSampler", "inputs": {"noise": ["1", 0]}, }, "3": { "class_type": "LocalSink", "inputs": {"latent": ["2", 0]}, }, } 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, ) payload = rewritten_prompt["1"]["inputs"]["original_node_data"] assert payload["component_node_ids"] == ["1", "2"] assert set(payload["subgraph_prompt"]) == {"1", "2"} assert payload["boundary_outputs"][0]["io_type"] == "LATENT" assert rewritten_prompt["3"]["inputs"]["latent"] == ["1", 0] assert summary.remote_component_ids == ["1"] assert api_intercept_module._is_transportable_output_type("NOISE") is False def test_extract_remote_node_ids_prefers_nested_prompt_id_over_colliding_root_id( api_intercept_module: Any, settings_module: Any, ) -> None: """Nested Modal markers should resolve to their composed prompt ids when root ids collide.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, volume_name="volume", route_path="/modal/queue_prompt", marker_property="is_modal_remote", local_storage_root=Path("/tmp/storage"), remote_storage_root="/storage", custom_nodes_archive_name="custom_nodes_bundle.zip", comfyui_root=None, custom_nodes_dir=Path("/tmp/custom_nodes"), ) workflow = { "nodes": [ {"id": 27, "properties": {"is_modal_remote": False}}, { "id": 195, "properties": {"is_modal_remote": False}, "subgraph": { "nodes": [ {"id": 27, "properties": {"is_modal_remote": True}}, ] }, }, ] } assert api_intercept_module.extract_remote_node_ids( workflow, settings, prompt_node_ids={"27", "195:27", "222", "223"}, ) == {"195:27"} def test_rewrite_keeps_nested_remote_nodes_remote_when_root_ids_collide( api_intercept_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: """Nested remote markers should survive prompt-id collisions with root workflow nodes.""" 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", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="local", sync_custom_nodes=False, 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": { "RemoteClip": _FakeRemoteClipNode, "RemoteConsumer": _FakeRemoteSamplerNode, "LocalConsumer": _FakeLocalSinkNode, }, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() workflow = { "nodes": [ {"id": 27, "properties": {"is_modal_remote": False}}, {"id": 222, "properties": {"is_modal_remote": True}}, { "id": 195, "properties": {"is_modal_remote": False}, "subgraph": { "nodes": [ {"id": 27, "properties": {"is_modal_remote": True}}, ] }, }, {"id": 223, "properties": {"is_modal_remote": False}}, ] } prompt = { "27": { "class_type": "LocalConsumer", "inputs": {}, "_meta": {"title": "Root Local Consumer"}, }, "222": { "class_type": "RemoteClip", "inputs": {}, "_meta": {"title": "Remote Clip Source"}, }, "195:27": { "class_type": "RemoteConsumer", "inputs": {"clip": ["222", 0]}, "_meta": {"title": "Nested Remote Consumer"}, }, "223": { "class_type": "LocalConsumer", "inputs": {"latent": ["195:27", 0]}, "_meta": {"title": "Local Sink"}, }, } 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, ) assert set(rewritten_prompt) == { "27", "195:27", "223", _artifact_finalizer_node_id(summary), } assert summary.remote_node_ids == ["195:27", "222"] assert summary.remote_component_ids == ["195:27"] def test_emit_modal_status_targets_prompt_client( api_intercept_module: Any, ) -> None: """Modal status events should preserve prompt and component metadata for the UI.""" api_intercept_module._MODAL_UI_EVENTS_BY_CLIENT.clear() class FakePromptServer: """Capture websocket events emitted by the queue route.""" def __init__(self) -> None: """Initialize the event sink.""" self.messages: list[tuple[str, dict[str, Any], str | None]] = [] def send_sync(self, event: str, data: dict[str, Any], sid: str | None) -> None: """Record an emitted websocket message.""" self.messages.append((event, data, sid)) prompt_server = FakePromptServer() api_intercept_module._emit_modal_status( prompt_server=prompt_server, phase="executing", client_id="client-1", prompt_id="prompt-1", node_ids=["4", "5"], modal_gpu="B300", component_node_ids_by_representative={"4": ["4", "5"]}, active_node_id="5", active_node_class_type="KSampler", active_node_role="sampling", ) assert prompt_server.messages == [ ( "modal_status", { "phase": "executing", "prompt_id": "prompt-1", "node_ids": ["4", "5"], "modal_gpu": "B300", "active_node_id": "5", "active_node_class_type": "KSampler", "active_node_role": "sampling", "components": [ { "representative_node_id": "4", "node_ids": ["4", "5"], } ], }, "client-1", ) ] replay_events = api_intercept_module.modal_ui_events_for_client("client-1") assert replay_events == [ { "event": "modal_status", "payload": { "phase": "executing", "prompt_id": "prompt-1", "node_ids": ["4", "5"], "modal_gpu": "B300", "active_node_id": "5", "active_node_class_type": "KSampler", "active_node_role": "sampling", "components": [ { "representative_node_id": "4", "node_ids": ["4", "5"], } ], }, "updated_at": replay_events[0]["updated_at"], } ] def test_modal_ui_event_replay_is_client_scoped(api_intercept_module: Any) -> None: """Refocus replay should only return events for the requesting ComfyUI client.""" api_intercept_module._MODAL_UI_EVENTS_BY_CLIENT.clear() api_intercept_module.record_modal_ui_event( "modal_progress", {"prompt_id": "prompt-1", "node_id": "4", "value": 2.0, "max": 10.0}, "client-1", ) api_intercept_module.record_modal_ui_event( "modal_status", {"prompt_id": "prompt-2", "phase": "executing", "node_ids": ["9"]}, "client-2", ) replay_events = api_intercept_module.modal_ui_events_for_client("client-1") assert len(replay_events) == 1 assert replay_events[0]["event"] == "modal_progress" assert replay_events[0]["payload"] == { "prompt_id": "prompt-1", "node_id": "4", "value": 2.0, "max": 10.0, } assert api_intercept_module.modal_ui_events_for_client(None) == [] def test_progress_state_route_is_queue_route_sibling(api_intercept_module: Any) -> None: """The frontend should have a stable sibling route for Modal UI replay.""" assert api_intercept_module._progress_state_route_path("/modal/queue_prompt") == ( "/modal/progress_state" ) assert api_intercept_module._progress_state_route_path("/custom/modal") == ( "/custom/modal/progress_state" ) def test_container_status_route_is_queue_route_sibling(api_intercept_module: Any) -> None: """The frontend should have a stable sibling route for active Modal containers.""" assert api_intercept_module._container_status_route_path("/modal/queue_prompt") == ( "/modal/container_status" ) assert api_intercept_module._container_status_route_path("/custom/modal") == ( "/custom/modal/container_status" ) def test_modal_reset_route_paths_are_queue_route_siblings(api_intercept_module: Any) -> None: """The frontend should have stable sibling routes for Modal maintenance actions.""" assert api_intercept_module._delete_modal_caches_route_path("/modal/queue_prompt") == ( "/modal/delete_caches" ) assert api_intercept_module._delete_modal_volume_route_path("/modal/queue_prompt") == ( "/modal/delete_volume" ) assert api_intercept_module._delete_modal_caches_route_path("/custom/modal") == ( "/custom/modal/delete_caches" ) assert api_intercept_module._delete_modal_volume_route_path("/custom/modal") == ( "/custom/modal/delete_volume" ) def test_delete_modal_cache_dicts_deletes_configured_dicts( api_intercept_module: Any, settings_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Deleting Modal caches should clear and delete every configured cache Dict.""" class FakeNotFoundError(Exception): """Stand-in for Modal object misses.""" class FakeDictObject: """Minimal Modal Dict object used only for existence checks.""" def __init__(self, name: str) -> None: """Store the configured Dict name.""" self.name = name def delete(self, name: str) -> None: """Fail if the deprecated instance deletion path is used.""" raise AssertionError(f"deprecated Dict.delete path used for {name}") class FakeDictObjects: """Minimal Modal Dict manager namespace.""" @staticmethod def delete(name: str, allow_missing: bool = False) -> None: """Record a manager delete call.""" assert allow_missing is True deleted.append(name) class FakeDict: """Minimal Modal Dict namespace.""" objects = FakeDictObjects() @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeDictObject: """Return fake Dict objects, except for one missing cache.""" assert create_if_missing is False if name == "app-interrupts": raise FakeNotFoundError(name) return FakeDictObject(name) class FakeModal: """Minimal Modal SDK double.""" exception = SimpleNamespace(NotFoundError=FakeNotFoundError) Dict = FakeDict settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=False, 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, interrupt_dict_name="app-interrupts", node_output_cache_dict_name="app-node-cache", session_bridge_dict_name="app-session-bridges", sync_index_dict_name="app-sync-index", snapshot_profile_dict_name="app-snapshot-profiles", ) deleted: list[str] = [] monkeypatch.setattr(api_intercept_module, "modal", FakeModal) result = asyncio.run(api_intercept_module.delete_modal_cache_dicts(settings)) assert result == { "deleted": [ "app-node-cache", "app-session-bridges", "app-sync-index", "app-snapshot-profiles", ], "skipped": ["app-interrupts"], } assert deleted == result["deleted"] def test_delete_modal_volume_deletes_configured_volume( api_intercept_module: Any, settings_module: Any, monkeypatch: Any, tmp_path: Path, ) -> None: """Deleting the Modal volume should target only the configured volume name.""" class FakeVolumeObject: """Minimal Modal Volume object used only for existence checks.""" def __init__(self, name: str) -> None: """Store the configured Volume name.""" self.name = name def delete(self, name: str) -> None: """Fail if the deprecated instance deletion path is used.""" raise AssertionError(f"deprecated Volume.delete path used for {name}") class FakeVolumeObjects: """Minimal Modal Volume manager namespace.""" @staticmethod def delete(name: str, allow_missing: bool = False) -> None: """Record a manager delete call.""" assert allow_missing is True deleted.append(name) class FakeVolume: """Minimal Modal Volume namespace.""" objects = FakeVolumeObjects() @staticmethod def from_name(name: str, create_if_missing: bool = False) -> FakeVolumeObject: """Return a fake Volume object.""" assert create_if_missing is False return FakeVolumeObject(name) class FakeModal: """Minimal Modal SDK double.""" exception = SimpleNamespace() Volume = FakeVolume settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, allow_ephemeral_fallback=False, enable_memory_snapshot=True, enable_gpu_memory_snapshot=False, execution_mode="remote", sync_custom_nodes=False, volume_name="configured-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, ) deleted: list[str] = [] monkeypatch.setattr(api_intercept_module, "modal", FakeModal) result = asyncio.run(api_intercept_module.delete_modal_volume(settings)) assert result == {"deleted": ["configured-volume"], "skipped": []} assert deleted == ["configured-volume"] def test_modal_interrupt_queue_bridge_exposes_active_remote_prompts( api_intercept_module: Any, remote_modal_app_module: Any, ) -> None: """Targeted ComfyUI interrupts should see prompts currently blocked on Modal work.""" class FakePromptQueue: """Minimal ComfyUI prompt queue with no native running prompts.""" def get_current_queue(self) -> tuple[list[Any], list[Any]]: """Return an empty running queue and one pending item.""" return [], ["queued"] prompt_queue = FakePromptQueue() prompt_server = SimpleNamespace(prompt_queue=prompt_queue) cancellation_event = remote_modal_app_module.threading.Event() api_intercept_module._install_modal_interrupt_queue_bridge(prompt_server) with remote_modal_app_module._registered_active_remote_invocation( {"prompt_id": "prompt-1", "component_id": "component-1"}, cancellation_event, None, ): running, queued = prompt_queue.get_current_queue() assert queued == ["queued"] assert [item[1] for item in running] == ["prompt-1"]