From fce996475efe6c62ce2f0350ddf7f8a8b7701a34 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Fri, 26 Jun 2026 12:31:15 -0700 Subject: [PATCH] fix: include bridge refs in remote cache keys --- README.md | 1 + comfyui_modal_sync_cloud.py | 84 +++++++++++++++++++++++++++-- docs/LEARNINGS.md | 1 + remote/modal_app.py | 2 +- tests/test_modal_executor_node.py | 88 +++++++++++++++++++++++++++++++ 5 files changed, 170 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 1825485..56079c1 100644 --- a/README.md +++ b/README.md @@ -236,6 +236,7 @@ Remote runtime behavior: - The distributed cache hot path is also lighter now: `restored-hit` nodes short-circuit before any output serialization work in the persist phase, and restore fan-outs independent Modal Dict lookups so mapped item runs spend less time waiting on node-cache bookkeeping. - Distributed cache restore now preserves required boundary producers. If a downstream cache hit would skip an upstream node whose output must be collected as a boundary or session bridge, Modal-Sync ignores that downstream hit for the current subgraph so ComfyUI materializes the required producer instead of failing with a missing local cache entry. - Distributed cache signature rebuilds now distinguish real prompt links from nested runtime data. This keeps first-run cache restore from misreading list-shaped values such as `CONDITIONING` payloads as graph edges. +- Distributed cache signatures for resolved remote-session bridge inputs now include the durable `bridge_key`, so mapped downstream sampler items fed by different `CONDITIONING` bridge refs cannot reuse the first item's cached sampler output. - Remote session bridge reuse is keyed by durable bridge id as well as node output coordinates, so mapped outputs from the same node slot, such as per-prompt `CONDITIONING` values, do not collapse to the first resolved item. - Implicit mapped execution recognizes Comfy's singleton output wrapper around remote session ref lists. A `CONDITIONING` boundary shaped like `[[bridge_ref_a, bridge_ref_b]]` is itemized as two mapped inputs, while ordinary list-shaped conditioning still stays broadcast. - The remote headless `PromptServer.instance` now implements ComfyUI's progress-text method, so API nodes such as `GetImageSize` can report node text without crashing in Modal. diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 5444e0c..733c9d2 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -32,7 +32,7 @@ _REPO_ROOT = Path(__file__).resolve().parent _REMOTE_REPO_ROOT = Path("/root/comfyui_modal_sync_repo") _LOCAL_COMFYUI_ROOT = (Path.home() / "git" / "ComfyUI").resolve() _REMOTE_COMFYUI_ROOT = Path("/root/comfyui_src") -_REMOTE_APP_PROTOCOL_VERSION = 3 +_REMOTE_APP_PROTOCOL_VERSION = 4 _PYTORCH_CUDA_INDEX_URL = "https://download.pytorch.org/whl/cu128" _COMFYUI_TORCH_VERSION = "2.10.0" _COMFYUI_TORCHVISION_VERSION = "0.25.0" @@ -2462,7 +2462,7 @@ def _build_node_output_cache_immediate_signature( boundary_signature = None if isinstance(boundary_input_signatures, dict): candidate_signature = boundary_input_signatures.get(str(key)) - if isinstance(candidate_signature, str) and candidate_signature: + if candidate_signature is not None: boundary_signature = candidate_signature if boundary_signature is not None: signature.append((key, ("BOUNDARY_SOURCE", boundary_signature))) @@ -3230,10 +3230,75 @@ def _execute_node_locally_raw( return _invoke_original_node(resolved_node_mapping[class_type], node_data, kwargs) +def _remote_session_ref_cache_signature(value: Any) -> Any | None: + """Return cache-key metadata for any remote-session refs nested in `value`.""" + if is_remote_session_bridge_ref_payload(value): + return { + "kind": "remote_session_bridge_ref", + "bridge_key": value.get("bridge_key"), + "node_id": value.get("node_id"), + "output_index": value.get("output_index"), + } + if is_remote_session_value_ref_payload(value): + return { + "kind": "remote_session_value_ref", + "session_id": value.get("session_id"), + "node_id": value.get("node_id"), + "output_index": value.get("output_index"), + } + if isinstance(value, list): + items = [ + _remote_session_ref_cache_signature(item) + for item in value + ] + if any(item is not None for item in items): + return {"kind": "list", "items": items} + return None + if isinstance(value, tuple): + items = [ + _remote_session_ref_cache_signature(item) + for item in value + ] + if any(item is not None for item in items): + return {"kind": "tuple", "items": items} + return None + if isinstance(value, Mapping): + items = { + str(key): _remote_session_ref_cache_signature(item) + for key, item in value.items() + } + filtered_items = { + key: item + for key, item in items.items() + if item is not None + } + if filtered_items: + return {"kind": "mapping", "items": filtered_items} + return None + + +def _boundary_input_cache_signature( + *, + source_signature: Any, + cache_value: Any, +) -> Any | None: + """Return a cache signature for one boundary input and its unresolved source value.""" + ref_signature = _remote_session_ref_cache_signature(cache_value) + if ref_signature is None: + if isinstance(source_signature, str) and source_signature: + return source_signature + return None + return { + "source_signature": source_signature if isinstance(source_signature, str) else None, + "remote_session_refs": ref_signature, + } + + def _apply_boundary_inputs( prompt: dict[str, Any], boundary_input_specs: list[dict[str, Any]], hydrated_inputs: dict[str, Any], + cache_signature_inputs: dict[str, Any] | None = None, ) -> None: """Inject hydrated local boundary inputs into a remote subgraph prompt.""" for boundary_input in boundary_input_specs: @@ -3254,11 +3319,19 @@ def _apply_boundary_inputs( value, io_type=io_type, ) - source_signature = boundary_input.get("source_signature") - if isinstance(source_signature, str) and source_signature: + cache_signature = _boundary_input_cache_signature( + source_signature=boundary_input.get("source_signature"), + cache_value=( + cache_signature_inputs.get(proxy_input_name) + if cache_signature_inputs is not None + and proxy_input_name in cache_signature_inputs + else value + ), + ) + if cache_signature is not None: boundary_signatures = prompt_node.setdefault(_BOUNDARY_INPUT_SIGNATURES_KEY, {}) if isinstance(boundary_signatures, dict): - boundary_signatures[input_name] = source_signature + boundary_signatures[input_name] = cache_signature def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any: @@ -3788,6 +3861,7 @@ def _execute_subgraph_prompt( prompt=prompt, boundary_input_specs=list(normalized_payload.get("boundary_inputs", [])), hydrated_inputs=resolved_inputs, + cache_signature_inputs=hydrated_inputs, ) with _timed_phase("load_execution_module", component=component_id): execution = _load_execution_module() diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index c5bae28..0688c9d 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -281,3 +281,4 @@ - Rebuilt remote node-cache signatures must use stricter prompt-link detection than general prompt rewriting. Runtime values such as `CONDITIONING` can contain nested two-item lists, and treating those data lists as `[node_id, output_index]` edges makes first-run distributed cache restore fail before the target node executes. - Remote session bridge refs need bridge-keyed live cache entries, not only `(node_id, output_index)` entries. A mapped CLIP output can produce multiple bridge refs for the same node output slot, and target-session lookup by node/output alone makes downstream mapped sampler items reuse the first conditioning vector. - Comfy can wrap a proxy output value in an outer list before handing it to an `INPUT_IS_LIST`-aware local/debug node or another proxy. For mapped remote session refs this creates `[[bridge_ref_a, bridge_ref_b]]`; implicit mapped splitting must unwrap exactly that singleton ref-list wrapper before deciding whether nontransportable `CONDITIONING` should broadcast. +- Distributed remote node-cache signatures also need bridge identity, not only boundary source provenance. By the time `PromptExecutor` prepares its cache keys, bridge refs have already been resolved to live runtime objects, so the remote worker must carry the unresolved `bridge_key` into boundary cache metadata or mapped sampler items with different conditioning refs can collide and restore the first item's cached output. diff --git a/remote/modal_app.py b/remote/modal_app.py index 3e320c8..c0d4e31 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -57,7 +57,7 @@ from ..settings import get_settings logger = logging.getLogger(__name__) _MODAL_CLOUD_MODULE_NAME = "comfyui_modal_sync_cloud" -_REMOTE_APP_PROTOCOL_VERSION = 3 +_REMOTE_APP_PROTOCOL_VERSION = 4 _MODAL_AUTO_DEPLOY_LOCK = threading.Lock() _MODAL_AUTO_DEPLOY_STATES: dict[tuple[str, str | None], "_ModalAutoDeployState"] = {} _MODAL_REMOTE_APP_VERSION_OK: set[tuple[str, str | None]] = set() diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 84794ed..952a4ce 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -3051,6 +3051,94 @@ def test_modal_cloud_node_cache_key_uses_boundary_source_signature_for_unhashabl assert different_key != first_key +def test_modal_cloud_node_cache_key_includes_remote_session_bridge_key( + modal_cloud_module: Any, + monkeypatch: Any, +) -> None: + """Mapped remote-session bridge inputs should not collide in the distributed cache.""" + + class FakeDynPrompt: + """Minimal dynamic prompt wrapper backed by one mutable prompt dict.""" + + def __init__(self, prompt: dict[str, Any]) -> None: + """Store the prompt used by the cache-key rebuild.""" + self._prompt = prompt + + def has_node(self, node_id: str) -> bool: + """Return whether the requested node exists in the prompt.""" + return str(node_id) in self._prompt + + def get_node(self, node_id: str) -> dict[str, Any]: + """Return the stored node payload.""" + return self._prompt[str(node_id)] + + class FakeConditioning: + """Stand-in for a resolved CONDITIONING object.""" + + def build_prompt(bridge_key: str) -> dict[str, Any]: + """Build one sampler prompt with a resolved conditioning and original bridge ref metadata.""" + prompt = { + "507": { + "class_type": "FakeSampler", + "inputs": {}, + } + } + boundary_spec = [ + { + "proxy_input_name": "phase_bridge_3", + "io_type": "CONDITIONING", + "source_signature": "SRC_node_508_output_0", + "targets": [{"node_id": "507", "input_name": "positive"}], + } + ] + bridge_ref = modal_cloud_module.RemoteSessionBridgeRef( + bridge_key=bridge_key, + node_id="508", + output_index=0, + session_id="session_1", + ).to_payload() + modal_cloud_module._apply_boundary_inputs( + prompt, + boundary_spec, + {"phase_bridge_3": FakeConditioning()}, + cache_signature_inputs={"phase_bridge_3": bridge_ref}, + ) + return prompt + + monkeypatch.setattr( + modal_cloud_module, + "_load_nodes_module", + lambda: types.SimpleNamespace( + NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})} + ), + ) + monkeypatch.setattr( + modal_cloud_module, + "_include_unique_id_in_input_signature", + lambda class_type: False, + ) + + def cache_key_for(prompt: dict[str, Any]) -> str | None: + """Build one distributed cache key for the prepared prompt.""" + cache_key_set = types.SimpleNamespace( + dynprompt=FakeDynPrompt(prompt), + is_changed_cache=types.SimpleNamespace(is_changed={"507": False}), + get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}), + include_node_id_in_input=lambda: False, + get_data_key=lambda node_id: None, + ) + return modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "507") + + first_key = cache_key_for(build_prompt("RSB_first")) + same_key = cache_key_for(build_prompt("RSB_first")) + second_key = cache_key_for(build_prompt("RSB_second")) + + assert isinstance(first_key, str) + assert first_key.startswith("NC_") + assert same_key == first_key + assert second_key != first_key + + def test_modal_cloud_node_cache_key_treats_nested_two_item_lists_as_data( modal_cloud_module: Any, monkeypatch: Any,