diff --git a/README.md b/README.md index 1db9aec..042b5de 100644 --- a/README.md +++ b/README.md @@ -321,7 +321,7 @@ For direct local `PreviewImage` consumers of a remote boundary `IMAGE` output, t Remote subgraph runs also persist transport-safe node outputs into a shared Modal `Dict` using ComfyUI's `CacheKeySetInputSignature` semantics. In practice that lets later workflow runs reuse small and medium cached outputs, including image tensors, even after Modal scales down and starts a fresh container, while still skipping non-serializable outputs and entries above the configured size cap. -Session-backed split proxies use a second Modal `Dict` for durable bridge metadata. That store does not try to serialize live remote-only objects like `MODEL` or `CONDITIONING`; instead it records how to replay the producing split phase, so cached bridge refs can lazily rebuild those values inside the current remote session if the original Modal container is gone. +Session-backed split proxies use a second Modal `Dict` for durable bridge metadata. That store still records the producer replay recipe for every bridge ref, and it now also persists directly restorable serialized payloads for selected remote-only types such as `CONDITIONING`. Fresh workers can rehydrate those bridge values into the current remote session without replay when a durable serialized form is available, while types like `MODEL` still fall back to producer replay. The extension also keeps a global activity badge visible during queue-time sync and remote execution, including the period before the prompt is formally queued. During setup it now surfaces more specific stages such as packaging the `custom_nodes` ZIP and uploading named assets, and after remote execution finishes it can briefly report that it is still receiving Modal outputs before the local proxy returns. @@ -405,6 +405,7 @@ COMFYUI_ROOT=/tmp/comfyui-modal-test/ComfyUI \ - Queue-time rewrite and implicit mapped execution now handle the list-shaped edge cases that showed up online: local `ModalMapInput` markers still produce mapped remote work, raw semantic lists like `CONDITIONING` stay broadcast, lists of remote session refs split per item, wildcard LATENT-list boundaries still itemize for concrete `LATENT` sockets like `KSampler.latent_image`, primitive list values from upstream nodes are allowed to follow ComfyUI's normal over-list behavior, singleton-list cleanup preserves semantic payloads, and `INPUT_IS_LIST` targets suppress outer Modal fan-out so the remote subgraph runs once unchanged. - Remote interrupt polling now shuts down cleanly when Modal closes the request-scoped client during container teardown. A late `modal.exception.ClientClosed` from the shared interrupt `Dict` is treated as normal cleanup rather than an unhandled post-success thread crash. - Split mapped proxies now carry a bridge-only `static_phase`, and the local async mapped scheduler seeds one bound remote lane per worker before dispatching per-item sampler calls. That keeps mapped fan-out from landing on a fresh worker and immediately replaying the full static producer branch just to reconstruct remote-only bridge inputs. +- Seeded mapped lanes now keep their prompt-scoped remote session alive until the explicit per-lane cleanup payload runs, so the first mapped item on that lane can reuse live session values instead of immediately dropping back to warm-bridge restore. Durable bridge records now also persist directly restorable `CONDITIONING` payloads, which lets fresh workers skip producer replay entirely for those bridge refs. - Frontend mapped-progress lanes are now scoped by component representative, and streamed runs still consume the proxy's final native completion event. That prevents parallel remote components from stealing each other's purple bars and stops single-node remote components from staying purple after completion. ## Current limitations diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 0e114e6..69e7296 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -91,6 +91,7 @@ _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LOCK = threading.Lock() _REMOTE_SESSION_BRIDGE_VALUE_CACHE: dict[str, Any] = {} _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER: list[str] = [] _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LIMIT = 32 +_DURABLE_BRIDGE_SERIALIZATION_IO_TYPES = frozenset({"CONDITIONING"}) try: import modal # type: ignore @@ -131,10 +132,13 @@ def _build_remote_session_bridge_record( hydrated_inputs: dict[str, Any], node_id: str, output_index: int, + io_type: str, + output_value: Any, ) -> RemoteSessionBridgeRecord: """Build one durable bridge record for a session-backed boundary output.""" producer_payload = _sanitize_payload_for_session_bridge_record(payload) producer_inputs = serialize_mapping(hydrated_inputs) + serialized_output = _serialize_durable_bridge_output(output_value, io_type) return RemoteSessionBridgeRecord( bridge_key=stable_session_bridge_key( producer_payload=producer_payload, @@ -146,6 +150,8 @@ def _build_remote_session_bridge_record( output_index=output_index, producer_payload=producer_payload, producer_inputs=producer_inputs, + serialized_output=serialized_output, + serialized_output_io_type=(str(io_type) if serialized_output is not None else None), ) @@ -220,6 +226,54 @@ def _get_remote_session_bridge_value(bridge_key: str) -> Any | None: return _clone_loader_cache_value(cached_value) +def _serialize_durable_bridge_output(output_value: Any, io_type: str) -> Any | None: + """Serialize one bridge output when its io_type supports durable direct restore.""" + normalized_io_type = str(io_type or "") + if normalized_io_type not in _DURABLE_BRIDGE_SERIALIZATION_IO_TYPES: + return None + try: + return serialize_value(output_value) + except TypeError: + logger.warning( + "Skipping durable bridge serialization for io_type=%s value_type=%s.", + normalized_io_type, + type(output_value).__name__, + ) + return None + + +def _restore_serialized_remote_session_bridge_value( + record: RemoteSessionBridgeRecord, + *, + target_session_handle: RemoteSessionHandle, + resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, +) -> Any | None: + """Restore one bridge value directly from a durable serialized payload.""" + if record.serialized_output is None: + return None + + restore_started_at = time.perf_counter() + restored_value = deserialize_value(record.serialized_output) + _REMOTE_SESSION_STORE.put_output( + target_session_handle, + node_id=record.node_id, + output_index=record.output_index, + value=restored_value, + ) + _store_remote_session_bridge_value(record.bridge_key, restored_value) + if resolution_stats is not None: + resolution_stats.durable_bridge_hits += 1 + resolution_stats.session_restore_writes += 1 + resolution_stats.direct_restore_seconds += time.perf_counter() - restore_started_at + logger.info( + "Restored remote session bridge bridge_key=%s from durable serialized %s payload into session_id=%s.", + record.bridge_key, + record.serialized_output_io_type or "bridge", + target_session_handle.session_id, + ) + return restored_value + + def _remote_session_bridge_replay_stack() -> set[str]: """Return the thread-local guard set for bridge replay recursion detection.""" replay_stack = getattr(_REMOTE_SESSION_BRIDGE_REPLAY_STATE, "bridge_keys", None) @@ -278,6 +332,13 @@ def _rehydrate_remote_session_bridge_value( resolution_stats.bridge_record_lookup_seconds += ( time.perf_counter() - record_lookup_started_at ) + restored_value = _restore_serialized_remote_session_bridge_value( + record, + target_session_handle=target_session_handle, + resolution_stats=resolution_stats, + ) + if restored_value is not None: + return restored_value replay_payload = copy.deepcopy(record.producer_payload) replay_payload["remote_session"] = target_session_handle.to_payload() replay_payload.pop("clear_remote_session", None) @@ -386,11 +447,12 @@ def _log_remote_session_resolution_summary( loader_hit_delta = loader_cache_after.get("hit", 0) - loader_cache_before.get("hit", 0) loader_miss_delta = loader_cache_after.get("miss", 0) - loader_cache_before.get("miss", 0) _emit_cloud_info( - "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d loader_cache_hits=%d loader_cache_misses=%d", + "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d durable_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d loader_cache_hits=%d loader_cache_misses=%d", component_id, resolution_stats.input_ref_count, resolution_stats.live_session_hits, resolution_stats.bridge_cache_hits, + resolution_stats.durable_bridge_hits, resolution_stats.bridge_record_lookups, resolution_stats.bridge_record_lookup_seconds, resolution_stats.replay_count, @@ -436,6 +498,7 @@ class _RemoteSessionBridgeResolutionStats: input_ref_count: int = 0 live_session_hits: int = 0 bridge_cache_hits: int = 0 + durable_bridge_hits: int = 0 bridge_record_lookups: int = 0 bridge_record_lookup_seconds: float = 0.0 replay_count: int = 0 @@ -3013,6 +3076,8 @@ def _execute_subgraph_prompt( hydrated_inputs=hydrated_inputs, node_id=node_id, output_index=output_index, + io_type=str(boundary_output.get("io_type") or "*"), + output_value=output_value, ) _store_remote_session_bridge_record(bridge_record) _store_remote_session_bridge_value(bridge_record.bridge_key, output_value) diff --git a/remote/modal_app.py b/remote/modal_app.py index b7b9fa9..4311307 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -30,6 +30,7 @@ from ..serialization import ( deserialize_node_inputs, deserialize_node_outputs, serialize_mapping, + serialize_value, split_mapped_value, serialize_node_outputs, serialize_node_inputs, @@ -73,6 +74,7 @@ _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LOCK = threading.Lock() _REMOTE_SESSION_BRIDGE_VALUE_CACHE: dict[str, Any] = {} _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER: list[str] = [] _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LIMIT = 32 +_DURABLE_BRIDGE_SERIALIZATION_IO_TYPES = frozenset({"CONDITIONING"}) @dataclass @@ -108,6 +110,7 @@ class _RemoteSessionBridgeResolutionStats: input_ref_count: int = 0 live_session_hits: int = 0 bridge_cache_hits: int = 0 + durable_bridge_hits: int = 0 bridge_record_lookups: int = 0 bridge_record_lookup_seconds: float = 0.0 replay_count: int = 0 @@ -475,10 +478,13 @@ def _build_remote_session_bridge_record( hydrated_inputs: dict[str, Any], node_id: str, output_index: int, + io_type: str, + output_value: Any, ) -> RemoteSessionBridgeRecord: """Build one durable bridge record for a session-backed boundary output.""" producer_payload = _sanitize_payload_for_session_bridge_record(payload) producer_inputs = serialize_mapping(hydrated_inputs) + serialized_output = _serialize_durable_bridge_output(output_value, io_type) return RemoteSessionBridgeRecord( bridge_key=stable_session_bridge_key( producer_payload=producer_payload, @@ -490,6 +496,8 @@ def _build_remote_session_bridge_record( output_index=output_index, producer_payload=producer_payload, producer_inputs=producer_inputs, + serialized_output=serialized_output, + serialized_output_io_type=(str(io_type) if serialized_output is not None else None), ) @@ -541,6 +549,54 @@ def _get_remote_session_bridge_value(bridge_key: str) -> Any | None: return _clone_cached_bridge_value(cached_value) +def _serialize_durable_bridge_output(output_value: Any, io_type: str) -> Any | None: + """Serialize one bridge output when its io_type supports durable direct restore.""" + normalized_io_type = str(io_type or "") + if normalized_io_type not in _DURABLE_BRIDGE_SERIALIZATION_IO_TYPES: + return None + try: + return serialize_value(output_value) + except TypeError: + logger.warning( + "Skipping durable bridge serialization for io_type=%s value_type=%s.", + normalized_io_type, + type(output_value).__name__, + ) + return None + + +def _restore_serialized_remote_session_bridge_value( + record: RemoteSessionBridgeRecord, + *, + target_session_handle: RemoteSessionHandle, + resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, +) -> Any | None: + """Restore one bridge value directly from a durable serialized payload.""" + if record.serialized_output is None: + return None + + restore_started_at = time.perf_counter() + restored_value = deserialize_value(record.serialized_output) + _REMOTE_SESSION_STORE.put_output( + target_session_handle, + node_id=record.node_id, + output_index=record.output_index, + value=restored_value, + ) + _store_remote_session_bridge_value(record.bridge_key, restored_value) + if resolution_stats is not None: + resolution_stats.durable_bridge_hits += 1 + resolution_stats.session_restore_writes += 1 + resolution_stats.direct_restore_seconds += time.perf_counter() - restore_started_at + logger.info( + "Resolved remote session bridge bridge_key=%s from durable serialized %s payload into session_id=%s.", + record.bridge_key, + record.serialized_output_io_type or "bridge", + target_session_handle.session_id, + ) + return restored_value + + def _remote_session_bridge_replay_stack() -> set[str]: """Return the thread-local guard set for bridge replay recursion detection.""" replay_stack = getattr(_REMOTE_SESSION_BRIDGE_REPLAY_STATE, "bridge_keys", None) @@ -596,6 +652,13 @@ def _rehydrate_remote_session_bridge_value( resolution_stats.bridge_record_lookup_seconds += ( time.perf_counter() - record_lookup_started_at ) + restored_value = _restore_serialized_remote_session_bridge_value( + record, + target_session_handle=target_session_handle, + resolution_stats=resolution_stats, + ) + if restored_value is not None: + return restored_value replay_payload = copy.deepcopy(record.producer_payload) replay_payload["remote_session"] = target_session_handle.to_payload() replay_payload.pop("clear_remote_session", None) @@ -685,11 +748,12 @@ def _log_remote_session_resolution_summary( if resolution_stats.input_ref_count <= 0: return logger.info( - "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d.", + "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d durable_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d.", component_id, resolution_stats.input_ref_count, resolution_stats.live_session_hits, resolution_stats.bridge_cache_hits, + resolution_stats.durable_bridge_hits, resolution_stats.bridge_record_lookups, resolution_stats.bridge_record_lookup_seconds, resolution_stats.replay_count, @@ -1377,6 +1441,8 @@ def _execute_subgraph_with_mapping( hydrated_inputs=hydrated_inputs, node_id=node_id, output_index=output_index, + io_type=str(boundary_output.get("io_type") or "*"), + output_value=output_value, ) _REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) _store_remote_session_bridge_value(bridge_record.bridge_key, output_value) @@ -1516,6 +1582,8 @@ def _execute_subgraph_prompt( hydrated_inputs=hydrated_inputs, node_id=node_id, output_index=output_index, + io_type=str(boundary_output.get("io_type") or "*"), + output_value=output_value, ) _REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) _store_remote_session_bridge_value(bridge_record.bridge_key, output_value) @@ -2516,6 +2584,7 @@ def _execute_mapped_subgraph_payload( f"{payload.get('component_id', 'modal-subgraph')}::static", suppress_status_stream=True, ) + static_phase_payload.pop("clear_remote_session", None) elif payload.get("static_execute_node_ids"): static_phase_payload = _build_static_mapped_payload(payload) @@ -3195,6 +3264,7 @@ async def _invoke_implicitly_mapped_subgraph_async(payload: dict[str, Any], kwar f"{payload.get('component_id', 'modal-subgraph')}::seed:{lane_index}", suppress_status_stream=True, ) + lane_seed_payload.pop("clear_remote_session", None) await _invoke_bound_remote_engine_async( lane_remote_engines[lane_index], lane_seed_payload, diff --git a/session_state.py b/session_state.py index 54e48ef..2ca2f18 100644 --- a/session_state.py +++ b/session_state.py @@ -141,6 +141,8 @@ class RemoteSessionBridgeRecord: output_index: int producer_payload: dict[str, Any] producer_inputs: dict[str, Any] + serialized_output: Any | None = None + serialized_output_io_type: str | None = None def to_payload(self) -> dict[str, Any]: """Serialize this bridge record into a JSON-safe payload mapping.""" @@ -151,6 +153,8 @@ class RemoteSessionBridgeRecord: "output_index": self.output_index, "producer_payload": self.producer_payload, "producer_inputs": self.producer_inputs, + "serialized_output": self.serialized_output, + "serialized_output_io_type": self.serialized_output_io_type, } @classmethod @@ -173,12 +177,19 @@ class RemoteSessionBridgeRecord: raise RemoteSessionStateError( "Remote session bridge records must define mapping producer_payload and producer_inputs." ) + serialized_output_io_type = payload.get("serialized_output_io_type") return cls( bridge_key=bridge_key, node_id=node_id, output_index=int(output_index), producer_payload=dict(producer_payload), producer_inputs=dict(producer_inputs), + serialized_output=payload.get("serialized_output"), + serialized_output_io_type=( + str(serialized_output_io_type) + if serialized_output_io_type is not None and str(serialized_output_io_type).strip() + else None + ), ) diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 9ca3014..8f1c0c7 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -1428,6 +1428,82 @@ def test_modal_cloud_rehydrates_bridge_refs_from_warm_value_cache_without_replay assert resolution_stats.session_restore_writes == 1 +def test_modal_cloud_rehydrates_conditioning_bridge_refs_from_durable_record_without_replay( + modal_cloud_module: Any, + monkeypatch: Any, +) -> None: + """Durably serialized CONDITIONING bridge values should restore without replay on fresh workers.""" + import torch + + target_handle = modal_cloud_module.RemoteSessionHandle( + session_id="session-target", + prompt_id="prompt-1", + owner_component_id="component-1", + ) + bridge_ref = modal_cloud_module.RemoteSessionBridgeRef( + bridge_key="RSB_conditioning_bridge", + node_id="node-9", + output_index=0, + session_id="session-source", + ) + conditioning = [ + [ + torch.arange(6, dtype=torch.float32).reshape(1, 2, 3), + {"pooled_output": torch.arange(4, dtype=torch.float32).reshape(1, 4)}, + ] + ] + monkeypatch.setattr( + modal_cloud_module, + "_load_remote_session_bridge_record", + lambda bridge_key: modal_cloud_module.RemoteSessionBridgeRecord( + bridge_key=bridge_key, + node_id="node-9", + output_index=0, + producer_payload={"component_id": "should-not-replay"}, + producer_inputs={}, + serialized_output=modal_cloud_module.serialize_value(conditioning), + serialized_output_io_type="CONDITIONING", + ), + ) + monkeypatch.setattr( + modal_cloud_module, + "_execute_subgraph_prompt", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("durable CONDITIONING restore should skip replay") + ), + ) + resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats() + + try: + restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value( + bridge_ref, + target_session_handle=target_handle, + custom_nodes_root=None, + cancellation_event=None, + interrupt_store=None, + interrupt_flag_key=None, + resolution_stats=resolution_stats, + ) + stored_value = modal_cloud_module._REMOTE_SESSION_STORE.get_output( + modal_cloud_module.RemoteSessionValueRef( + session_id=target_handle.session_id, + node_id="node-9", + output_index=0, + ) + ) + finally: + modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle) + + assert torch.equal(restored_value[0][0], conditioning[0][0]) + assert torch.equal(restored_value[0][1]["pooled_output"], conditioning[0][1]["pooled_output"]) + assert torch.equal(stored_value[0][0], conditioning[0][0]) + assert resolution_stats.bridge_cache_hits == 0 + assert resolution_stats.durable_bridge_hits == 1 + assert resolution_stats.bridge_record_lookups == 1 + assert resolution_stats.replay_count == 0 + assert resolution_stats.session_restore_writes == 1 + + def test_modal_cloud_logs_remote_session_resolution_summary( modal_cloud_module: Any, monkeypatch: Any, @@ -1446,6 +1522,7 @@ def test_modal_cloud_logs_remote_session_resolution_summary( input_ref_count=2, live_session_hits=1, bridge_cache_hits=1, + durable_bridge_hits=1, bridge_record_lookups=1, bridge_record_lookup_seconds=0.25, replay_count=1, @@ -1459,8 +1536,8 @@ def test_modal_cloud_logs_remote_session_resolution_summary( assert observed_logs == [ ( - "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d loader_cache_hits=%d loader_cache_misses=%d", - ("1::mapped", 2, 1, 1, 1, 0.25, 1, 1.5, 0.02, 1, 2, 1), + "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d durable_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d loader_cache_hits=%d loader_cache_misses=%d", + ("1::mapped", 2, 1, 1, 1, 1, 0.25, 1, 1.5, 0.02, 1, 2, 1), ) ] @@ -3927,6 +4004,79 @@ def test_local_remote_app_rehydrates_bridge_refs_from_warm_value_cache_without_r assert resolution_stats.session_restore_writes == 1 +def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_record_without_replay( + remote_modal_app_module: Any, + monkeypatch: Any, +) -> None: + """The local fallback should restore durably serialized CONDITIONING bridge values without replay.""" + import torch + + target_handle = remote_modal_app_module.RemoteSessionHandle( + session_id="session-target", + prompt_id="prompt-1", + owner_component_id="component-1", + ) + bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef( + bridge_key="RSB_local_conditioning_bridge", + node_id="node-9", + output_index=0, + session_id="session-source", + ) + conditioning = [ + [ + torch.arange(6, dtype=torch.float32).reshape(1, 2, 3), + {"pooled_output": torch.arange(4, dtype=torch.float32).reshape(1, 4)}, + ] + ] + monkeypatch.setattr( + remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, + "get_record", + lambda bridge_key: remote_modal_app_module.RemoteSessionBridgeRecord( + bridge_key=bridge_key, + node_id="node-9", + output_index=0, + producer_payload={"component_id": "should-not-replay"}, + producer_inputs={}, + serialized_output=remote_modal_app_module.serialize_value(conditioning), + serialized_output_io_type="CONDITIONING", + ), + ) + monkeypatch.setattr( + remote_modal_app_module, + "_execute_subgraph_prompt", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("durable CONDITIONING restore should skip local replay") + ), + ) + resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() + + try: + restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + bridge_ref, + target_session_handle=target_handle, + node_mapping=None, + resolution_stats=resolution_stats, + ) + stored_value = remote_modal_app_module._REMOTE_SESSION_STORE.get_output( + remote_modal_app_module.RemoteSessionValueRef( + session_id=target_handle.session_id, + node_id="node-9", + output_index=0, + ) + ) + finally: + remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + + assert torch.equal(restored_value[0][0], conditioning[0][0]) + assert torch.equal(restored_value[0][1]["pooled_output"], conditioning[0][1]["pooled_output"]) + assert torch.equal(stored_value[0][0], conditioning[0][0]) + assert resolution_stats.bridge_cache_hits == 0 + assert resolution_stats.durable_bridge_hits == 1 + assert resolution_stats.bridge_record_lookups == 1 + assert resolution_stats.replay_count == 0 + assert resolution_stats.session_restore_writes == 1 + + def test_local_phase_payload_builder_preserves_remote_session( remote_modal_app_module: Any, ) -> None: @@ -4354,7 +4504,7 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch( monkeypatch: Any, ) -> None: """Mapped remote scheduling should seed one bound worker lane before sending per-item calls there.""" - observed_calls: list[tuple[str, str, tuple[str, ...], dict[str, Any]]] = [] + observed_calls: list[tuple[str, str, bool, tuple[str, ...], dict[str, Any]]] = [] monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote") monkeypatch.setattr(remote_modal_app_module, "modal", object()) @@ -4375,7 +4525,15 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch( ) -> bytes: 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(remote_engine), str(payload["component_id"]), execute_node_ids, hydrated_inputs)) + observed_calls.append( + ( + str(remote_engine), + str(payload["component_id"]), + bool(payload.get("clear_remote_session")), + execute_node_ids, + hydrated_inputs, + ) + ) if str(payload["component_id"]).startswith("17::seed:"): return serialization_module.serialize_node_outputs(()) if str(payload["component_id"]).startswith("17::cleanup:"): @@ -4476,12 +4634,14 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch( ( "engine:session-1::lane:0", "17::seed:0", + False, ("4",), {"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True}}, ), ( "engine:session-1::lane:1", "17::seed:1", + False, ("4",), {"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True}}, ), @@ -4489,8 +4649,8 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch( assert {call[0] for call in item_calls} <= {"engine:session-1::lane:0", "engine:session-1::lane:1"} assert {call[0] for call in item_calls} <= {call[0] for call in seed_calls} assert sorted(cleanup_calls) == [ - ("engine:session-1::lane:0", "17::cleanup:0", (), {}), - ("engine:session-1::lane:1", "17::cleanup:1", (), {}), + ("engine:session-1::lane:0", "17::cleanup:0", True, (), {}), + ("engine:session-1::lane:1", "17::cleanup:1", True, (), {}), ] response_outputs = serialization_module.deserialize_node_outputs(response)[0] assert len(response_outputs) == 4