From cb7d02face082e6509809515436cffbc7ac2a619 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Wed, 24 Jun 2026 10:13:59 -0700 Subject: [PATCH] fix: split mapped components with local reentry --- api_intercept.py | 25 +++++- docs/LEARNINGS.md | 1 + tests/test_modal_executor_node.py | 127 ++++++++++++++++++++++++++++++ 3 files changed, 151 insertions(+), 2 deletions(-) diff --git a/api_intercept.py b/api_intercept.py index 4fa3372..efd21aa 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -2672,8 +2672,13 @@ def _build_component_payload( return payload def build_phase_payloads_for_transportable_splits() -> list[dict[str, Any]] | None: - """Return ordered split-proxy phase payloads for a non-mapped coarse component.""" - if component.mapped_boundary_input_name is not None or len(component.execute_node_ids) <= 1: + """Return ordered split-proxy phase payloads for components with local feedback.""" + if len(component.execute_node_ids) <= 1: + return None + if component.mapped_boundary_input_name is not None and not _component_has_local_reentry_dependency( + prompt=signature_prompt, + component=component, + ): return None if component.local_tap_node_ids: if not _component_has_local_reentry_dependency( @@ -3320,6 +3325,19 @@ def _rewrite_component_into_proxy( is_output_node=contains_output_node(list(phase_payload["component_node_ids"])), ) + if component.mapped_boundary_source_node_id is not None: + mapped_node_id_set = set(component.mapped_node_ids) + for phase_payload in phase_payloads: + phase_node_ids = {str(node_id) for node_id in phase_payload.get("component_node_ids", [])} + if not (phase_node_ids & mapped_node_id_set): + continue + register_modal_map_input_warmup_context( + component.mapped_boundary_source_node_id, + phase_payload, + str(component.mapped_boundary_input_io_type or "*"), + ) + break + for node_id, prompt_node in list(rewritten_prompt.items()): if node_id in component_proxy_node_ids: continue @@ -3611,12 +3629,15 @@ def rewrite_prompt_for_modal( continue if isinstance(split_proxy_payloads, list): summary.remote_component_ids.extend(proxy_node_ids) + mapped_node_id_set = set(component.mapped_node_ids) for phase_payload in split_proxy_payloads: phase_proxy_node_id = str(phase_payload["component_id"]) phase_component_node_ids = [str(node_id) for node_id in phase_payload["component_node_ids"]] summary.component_node_ids_by_representative[phase_proxy_node_id] = phase_component_node_ids for node_id in phase_component_node_ids: summary.rewritten_node_id_map[node_id] = phase_proxy_node_id + if mapped_node_id_set and mapped_node_id_set.intersection(phase_component_node_ids): + mapped_proxy_component_ids.add(phase_proxy_node_id) continue summary.remote_component_ids.extend(proxy_node_ids) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 3147919..14daf2e 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -272,3 +272,4 @@ - Cancellation can also arrive before a Modal payload reaches the registered remote-call wrapper, especially during first-run auto-deploy or deployed-app readiness polling. The outer local waiter must write the shared cancel flag and enforce the same local release grace during deploy/provisioning, and the worker must refuse to start the remote payload if cancellation was already observed. - Users need immediate cancellation feedback even when Modal cleanup is still cooperative. The frontend should intercept ComfyUI `/interrupt` requests, mark the prompt as `cancelling`, and ignore nonterminal late progress until the native terminal event clears the UI. - Prompt-level Modal cancellation has to release both the remote work and the local prompt admission gate. If a user cancels during Modal deploy/readiness, ComfyUI can mark the prompt interrupted while the background Modal call still owns the local gate, leaving the next prompt stuck behind an already-cancelled prompt id. +- `ModalMapInput` does not remove the need for ordered split proxies. If a mapped remote component has a local re-entry dependency, such as an early sampler preview feeding a local text/image node that returns into a later mapped sampler, the static side must still split into ordered phases before the mapped phase or the rewritten prompt can create a false dependency cycle. diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 5d37c60..73cb433 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -155,6 +155,14 @@ class _FakeRewriteLocalSinkNode: OUTPUT_IS_LIST = (False,) +class _FakeRewriteLocalFeedbackNode: + """Fake local node that turns a remote latent into a transportable remote input.""" + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + OUTPUT_IS_LIST = (False,) + + def test_modal_proxy_waits_for_active_prompt_before_starting_next_prompt( modal_executor_module: Any, ) -> None: @@ -7555,6 +7563,125 @@ def test_split_hybrid_proxies_allow_local_downstream_work_before_mapped_completi ] +def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies( + api_intercept_module: Any, + settings_module: Any, + sync_engine_module: Any, + tmp_path: Path, +) -> None: + """A mapped input should not collapse local feedback back into the same 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": { + "RemoteModel": _FakeRewriteRemoteModelNode, + "RemoteSampler": _FakeRewriteRemoteSamplerNode, + "LatentSource": _FakeRewriteLatentSourceNode, + "ModalMapInput": _FakeRewriteModalMapInputNode, + "LocalFeedback": _FakeRewriteLocalFeedbackNode, + "LocalSink": _FakeRewriteLocalSinkNode, + }, + "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": False}}, + {"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": "First Sampler"}, + }, + "4": { + "class_type": "LocalFeedback", + "inputs": {"image": ["3", 0]}, + "_meta": {"title": "Local Feedback"}, + }, + "5": { + "class_type": "LatentSource", + "inputs": {}, + "_meta": {"title": "Batch Values"}, + }, + "6": { + "class_type": "ModalMapInput", + "inputs": {"value": ["5", 0]}, + "_meta": {"title": "Map Input"}, + }, + "7": { + "class_type": "RemoteSampler", + "inputs": {"model": ["1", 0], "latent": ["6", 0], "prompt": ["4", 0]}, + "_meta": {"title": "Mapped Sampler"}, + }, + "8": { + "class_type": "LocalSink", + "inputs": {"image": ["7", 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 api_intercept_module._find_prompt_dependency_cycles(rewritten_prompt) == [] + assert summary.mapped_component_ids == ["7"] + assert summary.component_execution_stages == [["3"], ["7"]] + assert rewritten_prompt["4"]["inputs"]["image"] == ["3", 0] + assert rewritten_prompt["7"]["inputs"]["remote_input_2"] == ["4", 0] + first_payload = rewritten_prompt["3"]["inputs"]["original_node_data"] + mapped_payload = rewritten_prompt["7"]["inputs"]["original_node_data"] + assert first_payload["component_node_ids"] == ["1", "3"] + assert mapped_payload["component_node_ids"] == ["7"] + assert { + boundary_input["proxy_input_name"] + for boundary_input in mapped_payload["boundary_inputs"] + } == {"phase_bridge_0", "remote_input_1", "remote_input_2"} + + def test_invoke_implicitly_mapped_subgraph_async_zips_batched_boundary_inputs( remote_modal_app_module: Any, serialization_module: Any,