diff --git a/README.md b/README.md index fce8793..f5261d9 100644 --- a/README.md +++ b/README.md @@ -112,7 +112,7 @@ Mapped execution currently supports: One `Modal Map Input` boundary is supported per remote component. Non-mapped boundary inputs are broadcast unchanged to every item. Mapped outputs are reassembled in item order, concatenating batchable tensors when possible and otherwise preserving an ordered list. -Ordinary remote components without `Modal Map Input` still preserve ComfyUI's zipped batch behavior at the remote boundary. If a compatible batch reaches a primitive socket such as `seed: INT`, Modal-Sync itemizes it instead of injecting the whole list into the primitive widget input. If the target node declares `INPUT_IS_LIST`, Modal-Sync runs the component once as an ordinary subgraph. +Ordinary remote components without `Modal Map Input` still preserve ComfyUI's zipped batch behavior at the remote boundary. If a compatible batch reaches a primitive socket such as `seed: INT`, Modal-Sync itemizes it instead of injecting the whole list into the primitive widget input. If the target node declares `INPUT_IS_LIST`, Modal-Sync runs the component once as an ordinary subgraph. Scheduler-list output contracts remain stable for singleton runs: an ordinary one-item `IMAGE` result stays wrapped as one `[B,H,W,C]` value instead of losing its batch dimension when it returns to ComfyUI. Mapped components can contain both one-time execute targets and per-item execute targets. For example, two remote samplers may share one upstream model loader while only one sampler fans out over latents. Modal-Sync keeps the invariant upstream work separate from the per-item work so the sibling branch still runs once. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index da6a18f..c8a9187 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -1,5 +1,10 @@ # Learnings +## 2026-08-15 + +- ComfyUI `OUTPUT_IS_LIST` is a strict runtime container contract, not only planning metadata. If a proxy advertises a scheduler-list `IMAGE` but returns a raw `[1,H,W,C]` tensor, ComfyUI extends that tensor as the output list and passes `[H,W,C]` downstream; `PreviewImage` then treats height as batch size and saves `H` arrays shaped `[W,C]`. Normalize every scheduler-list proxy result to a Python list even when singleton execution skipped mapped aggregation. +- Multi-item mapped tests do not cover the singleton contract. Their aggregation path already wraps a joined tensor, so regression coverage also needs an ordinary remote result with one `IMAGE` and must verify that ComfyUI list merging preserves the complete `[1,H,W,C]` tensor as one scheduler item. + ## 2026-08-14 - A costly local re-entry is a graph property, not just a single-edge pattern. The planner can identify every affected local node by intersecting two local-only closures: nodes reachable downstream from any remote node and nodes that can reach any remote node upstream. This catches both one-node and multi-node remote-to-local-to-remote paths while excluding local preview branches that never re-enter remote execution. diff --git a/modal_executor_node.py b/modal_executor_node.py index 7b19400..992d986 100644 --- a/modal_executor_node.py +++ b/modal_executor_node.py @@ -303,6 +303,38 @@ def _normalize_proxy_payload(payload: Any) -> Mapping[str, Any]: return payload +def _normalize_scheduler_list_outputs( + payload: Mapping[str, Any], + outputs: Sequence[Any], +) -> tuple[Any, ...]: + """Make remote output containers match the proxy's scheduler-list declarations.""" + normalized_outputs = list(outputs) + boundary_outputs = payload.get("boundary_outputs") + if not isinstance(boundary_outputs, Sequence) or isinstance( + boundary_outputs, + str | bytes | bytearray, + ): + return tuple(normalized_outputs) + + for output_index, boundary_output in enumerate(boundary_outputs): + if output_index >= len(normalized_outputs): + break + if not isinstance(boundary_output, Mapping): + continue + if not bool(boundary_output.get("scheduler_is_list", False)): + continue + if isinstance(normalized_outputs[output_index], list): + continue + normalized_outputs[output_index] = [normalized_outputs[output_index]] + logger.debug( + "Wrapped singleton remote output %d for component=%s to satisfy its scheduler-list contract.", + output_index, + payload.get("component_id"), + ) + + return tuple(normalized_outputs) + + def _normalize_prompt_id(value: Any) -> str | None: """Return one non-empty prompt id string when available.""" if value is None: @@ -551,12 +583,13 @@ def _build_proxy_node_class( ) async with _modal_workflow_execution_slot(payload): - outputs = tuple( + outputs = _normalize_scheduler_list_outputs( + payload, await _execute_payload_async( get_remote_executor_client(), payload, _normalize_proxy_kwargs(kwargs), - ) + ), ) logger.debug( "Remote execution completed for payload kind=%s with %d outputs.", diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 70e0b00..fb10560 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -651,6 +651,74 @@ def test_mapped_boundary_payload_requests_scheduler_list_output( assert api_intercept_module._proxy_boundary_output_is_list(payload) is True +def test_scheduler_list_proxy_wraps_singleton_image_tensor( + modal_executor_module: Any, +) -> None: + """A singleton mapped IMAGE must retain its batch dimension in ComfyUI.""" + torch = pytest.importorskip("torch") + + fake_nodes_module = types.SimpleNamespace( + NODE_CLASS_MAPPINGS={}, + NODE_DISPLAY_NAME_MAPPINGS={}, + ) + proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered( + output_types=("IMAGE",), + output_names=("image",), + output_is_list=(True,), + nodes_module=fake_nodes_module, + is_output_node=False, + ) + proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] + image = torch.zeros((1, 8, 6, 3), dtype=torch.float32) + + class FakeClient: + """Return one ordinary IMAGE tensor from a singleton remote execution.""" + + async def execute_payload_async( + self, + payload: dict[str, Any], + kwargs: dict[str, Any], + ) -> tuple[Any, ...]: + """Return the unwrapped tensor produced by the ordinary subgraph path.""" + assert payload["boundary_outputs"][0]["scheduler_is_list"] is True + assert kwargs == {} + return (image,) + + modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) + try: + result = asyncio.run( + proxy_class.execute( + original_node_data={ + "payload_kind": "subgraph", + "prompt_id": "singleton-image-prompt", + "component_id": "singleton-image-component", + "boundary_outputs": [ + { + "proxy_output_name": "image", + "node_id": "11", + "output_index": 0, + "io_type": "IMAGE", + "is_list": False, + "scheduler_is_list": True, + } + ], + }, + unique_id="singleton-image-component", + ) + ) + finally: + modal_executor_module.set_remote_executor_client_factory(None) + + assert isinstance(result.result[0], list) + assert len(result.result[0]) == 1 + assert result.result[0][0] is image + scheduler_items: list[Any] = [] + scheduler_items.extend(result.result[0]) + assert len(scheduler_items) == 1 + assert scheduler_items[0] is image + assert tuple(scheduler_items[0].shape) == (1, 8, 6, 3) + + def test_component_proxy_emits_completion_token_after_remote_execution( modal_executor_module: Any, ) -> None: