From 1a6293faa9b8fa2375a2e9fa7e99bdd4af7ba360 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Fri, 26 Jun 2026 10:12:29 -0700 Subject: [PATCH] fix: prevent implicit mapping recursion --- README.md | 1 + docs/LEARNINGS.md | 1 + remote/modal_app.py | 31 ++++++++--- tests/test_modal_executor_node.py | 88 ++++++++++++++++++++++++++++++- 4 files changed, 114 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 3047598..81579d6 100644 --- a/README.md +++ b/README.md @@ -331,6 +331,7 @@ Current mapped-execution rules: - if queue-time rewrite had to leave a mapped boundary as `io_type="*"`, the runtime still recovers obvious concrete cases such as `list` values of LATENT-like mappings or boundaries whose target sockets all declare the same concrete input type - ordinary remote components without `Modal Map Input` still preserve ComfyUI's usual zipped batch semantics at the remote boundary: if multiple boundary inputs arrive as compatible batches, Modal-Sync fans that component out per item instead of injecting the raw list into primitive widget sockets like `seed: INT` - that implicit batching path now partitions one-time and per-item execute nodes too, so if two remote samplers share the same upstream `MODEL` but only one sampler receives a multi-item batch input, the sibling sampler still runs once instead of repeating once per batch item +- if an implicit batch boundary feeds an `INPUT_IS_LIST` node, Modal-Sync runs the component once as an ordinary subgraph instead of recursively rediscovering the same implicit fan-out - non-mapped boundary inputs are broadcast unchanged to every per-item execution - mapped outputs are reassembled in item order, concatenating batchable tensors back together when possible and otherwise preserving the ordered per-item list when shapes differ - per-item remote node status updates are suppressed, and streamed UI events from mapped item runs are filtered to the nodes that actually belong to that per-item payload so static sibling branches do not repaint the UI on every item diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 6dd882f..d568d68 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -275,3 +275,4 @@ - `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. - Implicit mapped cancellation has to stop the local item scheduler, not only the remote payload that is currently running. If a sampler lane observes a prompt interrupt, queued sibling items and final cleanup payloads must be suppressed immediately or a quick cancel/restart can dispatch stale mapped sampler work after the UI already reports cancellation. - `ModalMapInput` should accept scalar primitive values as a one-item map. This keeps the marker useful as a no-op adapter when an upstream control value like `INT` is not batched, while preserving fan-out behavior for lists, tensor batches, and `LATENT` mappings. +- The implicit mapped dispatcher needs an explicit ordinary-subgraph bypass. When a batch-like boundary targets an `INPUT_IS_LIST` node, the helper correctly decides not to fan out, but re-entering the public dispatcher without disabling implicit mapping just rediscovers the same split and recurses forever. diff --git a/remote/modal_app.py b/remote/modal_app.py index f35c597..40517cb 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -4099,7 +4099,11 @@ async def _invoke_implicitly_mapped_subgraph_async(payload: dict[str, Any], kwar payload.get("component_id"), input_is_list_target_node_ids, ) - return await invoke_remote_engine_async(payload, kwargs_payload) + return await invoke_remote_engine_async( + payload, + kwargs_payload, + allow_implicit_mapping=False, + ) parallelism, refined_prompt_warmup_target = boost_mapped_component_warmup( payload, @@ -5543,7 +5547,12 @@ def _invoke_modal_payload_blocking( return result -def invoke_remote_engine(payload: dict[str, Any], kwargs_payload: bytes) -> bytes: +def invoke_remote_engine( + payload: dict[str, Any], + kwargs_payload: bytes, + *, + allow_implicit_mapping: bool = True, +) -> bytes: """Invoke Modal when configured, or fall back to local in-process execution.""" execution_mode = get_settings().execution_mode if payload.get("payload_kind") == "mapped_subgraph": @@ -5552,7 +5561,7 @@ def invoke_remote_engine(payload: dict[str, Any], kwargs_payload: bytes) -> byte return serialize_node_outputs( _execute_mapped_subgraph_payload(payload, hydrated_inputs) ) - if payload.get("payload_kind") == "subgraph": + if allow_implicit_mapping and payload.get("payload_kind") == "subgraph": hydrated_inputs = deserialize_node_inputs(kwargs_payload) if _split_batch_boundary_inputs(payload, hydrated_inputs) is not None: return asyncio.run(_invoke_implicitly_mapped_subgraph_async(payload, kwargs_payload)) @@ -5614,18 +5623,28 @@ def invoke_remote_engine(payload: dict[str, Any], kwargs_payload: bytes) -> byte return response -async def invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes: +async def invoke_remote_engine_async( + payload: dict[str, Any], + kwargs_payload: bytes, + *, + allow_implicit_mapping: bool = True, +) -> bytes: """Invoke Modal asynchronously so multiple proxy nodes can wait on remote work in parallel.""" execution_mode = get_settings().execution_mode if payload.get("payload_kind") == "mapped_subgraph": if execution_mode == "local" or modal is None: return await _invoke_mapped_remote_engine_async(payload, kwargs_payload) - if payload.get("payload_kind") == "subgraph": + if allow_implicit_mapping and payload.get("payload_kind") == "subgraph": hydrated_inputs = deserialize_node_inputs(kwargs_payload) if _split_batch_boundary_inputs(payload, hydrated_inputs) is not None: return await _invoke_implicitly_mapped_subgraph_async(payload, kwargs_payload) if execution_mode == "local" or modal is None: - return await asyncio.to_thread(invoke_remote_engine, payload, kwargs_payload) + return await asyncio.to_thread( + invoke_remote_engine, + payload, + kwargs_payload, + allow_implicit_mapping=allow_implicit_mapping, + ) logger.info( "Dispatching async Modal remote invocation for component=%s payload_kind=%s.", diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index cf5dede..b07c57f 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -8507,7 +8507,13 @@ def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets )() monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module) - async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes: + async def fake_invoke_remote_engine_async( + payload: dict[str, Any], + kwargs_payload: bytes, + *, + allow_implicit_mapping: bool = True, + ) -> bytes: + del allow_implicit_mapping hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload) execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", [])) observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs)) @@ -8568,6 +8574,86 @@ def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets ] +def test_invoke_remote_engine_async_bypasses_implicit_mapping_once_for_input_is_list_targets( + remote_modal_app_module: Any, + serialization_module: Any, + monkeypatch: Any, +) -> None: + """The async dispatcher should not recurse when implicit fan-out is suppressed.""" + observed_calls: list[tuple[bool, dict[str, Any]]] = [] + + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": { + "ImplicitBatchListSource": _ImplicitBatchListSourceNode, + "ImplicitBatchScalarConsumer": _ImplicitBatchScalarConsumerNode, + "ImplicitBatchListConsumer": _ImplicitBatchListConsumerNode, + } + }, + )() + monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module) + + def fake_invoke_remote_engine( + payload: dict[str, Any], + kwargs_payload: bytes, + *, + allow_implicit_mapping: bool = True, + ) -> bytes: + hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload) + observed_calls.append((allow_implicit_mapping, hydrated_inputs)) + return serialization_module.serialize_node_outputs(("ordinary-subgraph",)) + + monkeypatch.setattr( + remote_modal_app_module, + "invoke_remote_engine", + fake_invoke_remote_engine, + ) + + payload = { + "payload_kind": "subgraph", + "component_id": "17", + "prompt_id": "prompt-1", + "component_node_ids": ["4", "12", "17"], + "execute_node_ids": ["4", "12"], + "subgraph_prompt": { + "17": {"class_type": "ImplicitBatchListSource", "inputs": {"values": 0}}, + "4": { + "class_type": "ImplicitBatchScalarConsumer", + "inputs": {"value": ["17", 0]}, + }, + "12": { + "class_type": "ImplicitBatchListConsumer", + "inputs": {"values": ["17", 1]}, + }, + }, + "boundary_inputs": [ + { + "proxy_input_name": "remote_input_0", + "io_type": "INT", + "targets": [{"node_id": "17", "input_name": "values"}], + } + ], + "boundary_outputs": [ + {"node_id": "4", "io_type": "STRING", "is_list": False}, + ], + "extra_data": {"client_id": "client-1"}, + } + + response = asyncio.run( + remote_modal_app_module.invoke_remote_engine_async( + payload, + serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}), + ) + ) + + assert serialization_module.deserialize_node_outputs(response) == ("ordinary-subgraph",) + assert observed_calls == [ + (False, {"remote_input_0": [10, 11, 12]}), + ] + + def test_implicitly_mapped_subgraph_clears_remote_session_once_after_all_items_finish( remote_modal_app_module: Any, serialization_module: Any,