From 2de907c8fa4fe96709faa4c8013e0dd05ab98b28 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Mon, 20 Apr 2026 13:21:24 -0700 Subject: [PATCH] fix: preserve singleton conditioning payloads --- README.md | 1 + comfyui_modal_sync_cloud.py | 37 +++++++++++++++++++++++++++---- docs/LEARNINGS.md | 1 + docs/UPNEXT.md | 1 + remote/modal_app.py | 37 +++++++++++++++++++++++++++---- tests/test_modal_executor_node.py | 37 +++++++++++++++++++++++++++++++ 6 files changed, 106 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 30ba416..a4580f1 100644 --- a/README.md +++ b/README.md @@ -392,6 +392,7 @@ COMFYUI_ROOT=/tmp/comfyui-modal-test/ComfyUI \ - [`modal_test_workflow.json`](modal_test_workflow.json) is a checked-in smoke artifact from a successful Modal-path run, not a pristine authoring workflow. - Implicit mapped execution only treats explicitly batchable transport types and list-wrapped primitive sockets as per-item inputs. List-backed semantic types like `CONDITIONING` stay broadcast across mapped item runs so the remote worker does not accidentally slice them into invalid fragments. - Remote primitive-socket validation only rejects raw widget list literals now. If a primitive socket like `INT` is hydrated from a boundary input or upstream output and carries a list value, the runtime leaves it alone so ComfyUI can apply its normal over-list execution semantics for nodes like `NextSeeds`. +- Boundary input normalization only unwraps singleton wrappers for scalar-like values and true prompt links. Single-entry semantic lists like `CONDITIONING` must stay wrapped as lists, or ComfyUI will misread `[[tensor, meta]]` as `[tensor, meta]` and crash during conditioning conversion. ## Current limitations diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 8e53d9f..4a5d3b6 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -1993,10 +1993,18 @@ def _apply_boundary_inputs( if proxy_input_name not in hydrated_inputs: raise KeyError(f"Missing hydrated boundary input {proxy_input_name!r}.") value = hydrated_inputs[proxy_input_name] + io_type = ( + str(boundary_input["io_type"]) + if boundary_input.get("io_type") is not None + else None + ) for target in boundary_input.get("targets", []): node_id = str(target["node_id"]) input_name = str(target["input_name"]) - prompt[node_id]["inputs"][input_name] = _normalize_prompt_input_value(value) + prompt[node_id]["inputs"][input_name] = _normalize_prompt_input_value( + value, + io_type=io_type, + ) def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any: @@ -2286,9 +2294,30 @@ def _normalize_link_output_index(value: Any) -> Any: return value -def _normalize_prompt_input_value(value: Any) -> Any: - """Unwrap transport-added singleton lists around scalar prompt input values.""" - while isinstance(value, list) and len(value) == 1: +def _unwrap_wrapped_prompt_link(value: Any) -> Any: + """Collapse nested singleton wrappers around one serialized prompt link when present.""" + candidate = value + while isinstance(candidate, list) and len(candidate) == 1: + candidate = candidate[0] + if _is_link(candidate): + return [candidate[0], _normalize_link_output_index(candidate[1])] + return value + + +def _normalize_prompt_input_value(value: Any, io_type: str | None = None) -> Any: + """Unwrap transport-added singleton wrappers only for scalar-like prompt input values.""" + wrapped_link = _unwrap_wrapped_prompt_link(value) + if wrapped_link is not value: + return wrapped_link + while ( + isinstance(value, list) + and len(value) == 1 + and ( + io_type in _PRIMITIVE_WIDGET_INPUT_TYPES + or value[0] is None + or isinstance(value[0], bool | int | float | str) + ) + ): value = value[0] if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str): return [value[0], _normalize_link_output_index(value[1])] diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 5a29836..0f10e15 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -154,3 +154,4 @@ - Distributed node-cache restore also benefits from fan-out/fan-in structure. The old per-node `await key -> await Dict get -> deserialize` loop made total restore time roughly the sum of all independent Dict reads, while batching those lookups concurrently preserves prompt-order hydration but trims the network-bound portion substantially. - Implicit mapped fan-out cannot treat every Python `list` as a batch axis. ComfyUI semantic payloads like `CONDITIONING` are also list-backed, and splitting them per item corrupts the conditioning structure into fragments like `["cond", {...}] -> "cond"` that later explode in `sampler_helpers.convert_cond`. Only explicitly batchable transport types and list-wrapped primitive sockets should participate in that automatic zipping path. - Primitive widget sockets are not always scalar once a graph is running. Custom nodes can legally emit `list(INT)` or similar values into an `INT` input and rely on ComfyUI's normal over-list execution behavior, so remote validation should only reject list values that are still raw widget literals, not boundary-fed or upstream-fed runtime values. +- Singleton-list cleanup on the remote boundary needs the same nuance. It is correct to unwrap transport-added wrappers like `[4] -> 4` and `[[["node", [0]]]] -> ["node", 0]`, but wrong to flatten single-entry semantic payloads like `[[tensor, meta]]` for `CONDITIONING`; that shape change is enough to turn `convert_cond()` into a tensor indexing crash. diff --git a/docs/UPNEXT.md b/docs/UPNEXT.md index 9a416cd..4f391cb 100644 --- a/docs/UPNEXT.md +++ b/docs/UPNEXT.md @@ -43,5 +43,6 @@ Partial-completion semantics for one proxy would require ComfyUI scheduler behav - Online testing confirmed that `RemoteEngine(session_affinity_key)` kept the split static and mapped proxies on the same deployed Modal worker for the tested workflow. - Online testing also exposed one implicit-mapped runtime edge: list-backed `CONDITIONING` inputs must stay broadcast across mapped item runs. That regression is now covered so future batching changes do not reintroduce the `convert_cond` failure shape. - Online testing also exposed a primitive-socket validation edge: upstream or boundary-fed `list(INT)` values are legitimate for nodes that rely on ComfyUI's over-list execution semantics, so only raw widget list literals should fail early. +- Online testing also exposed a boundary-normalization edge: singleton wrapper cleanup must preserve single-entry semantic lists like `CONDITIONING` while still collapsing true scalar wrappers and wrapped prompt links. - Run with `COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS=true` during broader production workflows and watch the mirrored session lifecycle logs for any real session misses under scale-out. - If production logs ever show split-proxy session misses under real scale-out, add a stronger affinity fallback or explicit leased-worker/session routing instead of relying on best-effort class-instance affinity. diff --git a/remote/modal_app.py b/remote/modal_app.py index 9e25380..6162193 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -603,6 +603,11 @@ def _apply_boundary_inputs( if proxy_input_name not in hydrated_inputs: raise KeyError(f"Missing hydrated boundary input {proxy_input_name!r}.") value = hydrated_inputs[proxy_input_name] + io_type = ( + str(boundary_input["io_type"]) + if boundary_input.get("io_type") is not None + else None + ) logger.info( "Applying boundary input %s to %d targets.", proxy_input_name, @@ -611,7 +616,10 @@ def _apply_boundary_inputs( for target in boundary_input.get("targets", []): node_id = str(target["node_id"]) input_name = str(target["input_name"]) - prompt[node_id]["inputs"][input_name] = _normalize_prompt_input_value(value) + prompt[node_id]["inputs"][input_name] = _normalize_prompt_input_value( + value, + io_type=io_type, + ) def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any: @@ -933,9 +941,30 @@ def _normalize_link_output_index(value: Any) -> Any: return value -def _normalize_prompt_input_value(value: Any) -> Any: - """Unwrap transport-added singleton lists around scalar prompt input values.""" - while isinstance(value, list) and len(value) == 1: +def _unwrap_wrapped_prompt_link(value: Any) -> Any: + """Collapse nested singleton wrappers around one serialized prompt link when present.""" + candidate = value + while isinstance(candidate, list) and len(candidate) == 1: + candidate = candidate[0] + if _is_link(candidate): + return [candidate[0], _normalize_link_output_index(candidate[1])] + return value + + +def _normalize_prompt_input_value(value: Any, io_type: str | None = None) -> Any: + """Unwrap transport-added singleton wrappers only for scalar-like prompt input values.""" + wrapped_link = _unwrap_wrapped_prompt_link(value) + if wrapped_link is not value: + return wrapped_link + while ( + isinstance(value, list) + and len(value) == 1 + and ( + io_type in _PRIMITIVE_WIDGET_INPUT_TYPES + or value[0] is None + or isinstance(value[0], bool | int | float | str) + ) + ): value = value[0] if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str): return [value[0], _normalize_link_output_index(value[1])] diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index b586e42..68ffc15 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -5082,6 +5082,43 @@ def test_apply_boundary_inputs_normalizes_wrapped_scalar_values( assert prompt["remote_1"]["inputs"]["value"] == 4 +@pytest.mark.parametrize( + ("module_fixture_name",), + [ + ("remote_modal_app_module",), + ("modal_cloud_module",), + ], +) +def test_apply_boundary_inputs_preserves_singleton_conditioning_lists( + request: Any, + module_fixture_name: str, +) -> None: + """Boundary input hydration must not flatten singleton CONDITIONING payloads.""" + target_module = request.getfixturevalue(module_fixture_name) + conditioning = [["cond", {"pooled_output": None}]] + prompt = { + "remote_1": { + "class_type": "BoundarySource", + "inputs": {"value": 0}, + "_meta": {}, + } + } + + target_module._apply_boundary_inputs( + prompt=prompt, + boundary_input_specs=[ + { + "proxy_input_name": "remote_input_0", + "io_type": "CONDITIONING", + "targets": [{"node_id": "remote_1", "input_name": "value"}], + } + ], + hydrated_inputs={"remote_input_0": conditioning}, + ) + + assert prompt["remote_1"]["inputs"]["value"] == conditioning + + @pytest.mark.parametrize( ("module_fixture_name",), [