diff --git a/README.md b/README.md index 1d8415d..f796993 100644 --- a/README.md +++ b/README.md @@ -127,7 +127,7 @@ When a prompt is queued: 6. Referenced model assets and, when enabled, `custom_nodes/` packages are mirrored into storage. 7. The rewritten prompt is submitted to ComfyUI's normal execution queue. 8. Local nodes execute normally until a proxy node is reached. -9. The proxy serializes boundary inputs, dispatches local or Modal execution, deserializes returned outputs, and exposes them as normal ComfyUI outputs. +9. The proxy dispatches local or Modal execution. Values needed by local nodes are materialized as normal ComfyUI outputs, while values consumed only by later remote components travel through the local graph as small Modal-backed references. Boundary-crossing values must be transportable. Supported evaluated values include: @@ -165,7 +165,9 @@ Warm containers can reuse loaded model state, `PromptExecutor` state, remote ses Independent Modal-backed components can overlap through ComfyUI's async proxy path and the local Modal call executor. Each Modal GPU container handles one active workflow execution at a time, so parallel ready components can scale out across containers instead of multiplexing several active executions onto one worker. `COMFY_MODAL_MAX_INFLIGHT_CALLS` bounds local dispatch independently from the local CPU count and the remote autoscaler. -Split proxies and mapped phases use prompt-scoped remote sessions for live non-transportable values. Durable bridge metadata is stored in a Modal `Dict` so later phases can rehydrate selected values after container churn. Oversized serialized bridge inputs and outputs are stored as integrity-checked, content-addressed objects on the shared Modal Volume instead of being embedded in `Dict` records. Sampler-producing bridges are not replayed as a fallback; losing those values is surfaced as a session-state error. +When a component output is consumed exclusively by other remote components, Modal-Sync keeps it remote. The producer returns a small durable bridge reference through the local ComfyUI proxy instead of returning the underlying `IMAGE`, `AUDIO`, `LATENT`, `MASK`, `SIGMAS`, `VIDEO`, or scalar value. A downstream remote component resolves that reference from the worker's warm bridge cache when possible, or restores it from shared Modal storage when Modal schedules the components on different containers. If any consumer is local, including an interim preview or final `SaveVideo`, the planner leaves that boundary materialized so the local node receives the ordinary ComfyUI value it expects. + +Split proxies, mapped phases, and remote-to-remote component edges use component-local remote sessions. Durable bridge metadata and small serialized values are stored in a Modal `Dict`. Serialized bridge values above `COMFY_MODAL_BRIDGE_INLINE_MAX_BYTES` are stored as integrity-checked, content-addressed objects on the shared Modal Volume, with only the object reference retained in the Dict record. The producer commits new objects before publishing a completed invocation; a consumer whose mounted Volume snapshot is older reads committed bytes through `Volume.read_file()` without reloading the mount or disturbing memory-mapped models. `NOISE` remains deliberately excluded because ComfyUI represents it as a deferred strategy object rather than a realized tensor. Each remote payload also carries a stable invocation id. The worker records its lifecycle in a shared Modal `Dict` and replays a completed result when the local client or Modal retries the same call. An overlapping delivery waits briefly for the active attempt to publish a terminal state instead of failing immediately. If a streamed call loses its consumer, the worker cancels unfinished compute and marks that attempt failed before the retry proceeds; uncommitted Volume writes remain unpublished. Large completed results use the same content-addressed Volume store, while failed attempts remain retryable. diff --git a/api_intercept.py b/api_intercept.py index 77445b6..229cf99 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -93,13 +93,14 @@ class BoundaryInputSpec: @dataclass class BoundaryOutputSpec: - """Describe one remote-to-local boundary value for a component.""" + """Describe one value exported across a remote component boundary.""" proxy_output_name: str source: LinkedOutputRef io_type: str is_list: bool preview_target_node_ids: list[str] = field(default_factory=list) + session_output: bool = False @dataclass @@ -2497,6 +2498,89 @@ def _build_component_plans( ] +def _mark_remote_to_remote_session_boundaries( + prompt: dict[str, Any], + components: list[RemoteComponentPlan], +) -> set[str]: + """Mark exclusively remote component edges for reference-based transport.""" + consumers = _build_consumer_map(prompt) + component_id_by_node_id = { + node_id: component.representative_node_id + for component in components + for node_id in component.node_ids + } + session_sources: set[LinkedOutputRef] = set() + + for component in components: + for boundary_output in component.boundary_outputs: + output_consumers = consumers.get(boundary_output.source, []) + if not output_consumers: + continue + if not all( + consumer.node_id in component_id_by_node_id + for consumer in output_consumers + ): + continue + boundary_output.session_output = True + session_sources.add(boundary_output.source) + logger.info( + "Keeping remote boundary output in Modal storage source=%s:%d io_type=%s producer_component=%s consumer_components=%s.", + boundary_output.source.node_id, + boundary_output.source.output_index, + boundary_output.io_type, + component.representative_node_id, + sorted( + { + component_id_by_node_id[consumer.node_id] + for consumer in output_consumers + } + ), + ) + + return _remote_session_component_ids( + components=components, + session_sources=session_sources, + ) + + +def _remote_session_component_ids( + *, + components: list[RemoteComponentPlan], + session_sources: set[LinkedOutputRef], +) -> set[str]: + """Return components that produce or consume Modal-backed boundary refs.""" + return { + component.representative_node_id + for component in components + if any(output.session_output for output in component.boundary_outputs) + or any( + boundary_input.source in session_sources + for boundary_input in component.boundary_inputs + ) + } + + +def _boundary_output_payload( + boundary_output: BoundaryOutputSpec, + *, + mapped_output: bool | None = None, +) -> dict[str, Any]: + """Serialize one component boundary output for a remote payload.""" + payload = { + "proxy_output_name": boundary_output.proxy_output_name, + "node_id": boundary_output.source.node_id, + "output_index": boundary_output.source.output_index, + "io_type": boundary_output.io_type, + "is_list": boundary_output.is_list, + "preview_target_node_ids": list(boundary_output.preview_target_node_ids), + } + if boundary_output.session_output: + payload["session_output"] = True + if mapped_output is not None: + payload["mapped_output"] = mapped_output + return payload + + def _mapped_boundary_origin_io_type( prompt: dict[str, Any], boundary_input: BoundaryInputSpec, @@ -2695,6 +2779,7 @@ def _build_component_payload( uploaded_volume_paths: list[str], terminate_container_on_error: bool, nodes_module: Any, + remote_session: dict[str, Any] | None = None, ) -> dict[str, Any]: """Build the serialized execution payload for one remote component.""" prompt_id = (extra_data or {}).get("prompt_id") @@ -2850,16 +2935,7 @@ def _build_component_payload( component.boundary_outputs, phase_node_id_set, ): - phase_boundary_outputs.append( - { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - } - ) + phase_boundary_outputs.append(_boundary_output_payload(boundary_output)) phase_output_names_by_source[boundary_output.source] = boundary_output.proxy_output_name pending_node_ids = remaining_node_ids - phase_node_id_set @@ -2912,7 +2988,7 @@ def _build_component_payload( "io_type": io_type, "is_list": is_list, "preview_target_node_ids": [], - "session_output": not _is_transportable_output_type(io_type), + "session_output": True, } ) phase_output_names_by_source[source] = proxy_output_name @@ -2921,7 +2997,7 @@ def _build_component_payload( source=source, io_type=io_type, is_list=is_list, - session_output=not _is_transportable_output_type(io_type), + session_output=True, ) phase_execute_node_ids = [ @@ -2950,14 +3026,16 @@ def _build_component_payload( ) if not phase_payloads or len(phase_payloads) <= 1: return None - if has_session_bridges: - remote_session = RemoteSessionHandle( + active_remote_session = remote_session + if active_remote_session is None and has_session_bridges: + active_remote_session = RemoteSessionHandle( session_id=uuid.uuid4().hex, prompt_id=(str(prompt_id) if prompt_id is not None else None), owner_component_id=component.representative_node_id, ).to_payload() + if active_remote_session is not None: for phase_index, phase_payload in enumerate(phase_payloads): - phase_payload["remote_session"] = copy.deepcopy(remote_session) + phase_payload["remote_session"] = copy.deepcopy(active_remote_session) if phase_index == len(phase_payloads) - 1: phase_payload["clear_remote_session"] = True logger.info( @@ -2992,25 +3070,13 @@ def _build_component_payload( signature_prompt=signature_prompt, ), "boundary_outputs": [ - ( - { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - "mapped_output": bool(boundary_output.source.node_id in set(component.mapped_node_ids)), - } - if component.mapped_boundary_input_name - else { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - } + _boundary_output_payload( + boundary_output, + mapped_output=( + boundary_output.source.node_id in set(component.mapped_node_ids) + if component.mapped_boundary_input_name + else None + ), ) for boundary_output in component.boundary_outputs ], @@ -3032,6 +3098,9 @@ def _build_component_payload( else None ), } + if remote_session is not None: + payload["remote_session"] = copy.deepcopy(remote_session) + payload["clear_remote_session"] = True logger.info( "Built remote payload for component %s: boundary_inputs=%d boundary_outputs=%d execute_nodes=%s custom_nodes_bundle=%s", component.representative_node_id, @@ -3099,14 +3168,7 @@ def _build_component_payload( signature_prompt=signature_prompt, ), "boundary_outputs": [ - { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - } + _boundary_output_payload(boundary_output) for boundary_output in static_boundary_outputs ] + static_bridge_outputs, @@ -3129,26 +3191,19 @@ def _build_component_payload( signature_prompt=signature_prompt, ), "boundary_outputs": [ - { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - "mapped_output": True, - } + _boundary_output_payload(boundary_output, mapped_output=True) for boundary_output in mapped_boundary_outputs ], "execute_node_ids": list(component.mapped_execute_node_ids), } if not component.static_node_ids: return _attach_snapshot_profile_key(payload, settings) - remote_session = RemoteSessionHandle( - session_id=uuid.uuid4().hex, - prompt_id=(str(prompt_id) if prompt_id is not None else None), - owner_component_id=component.representative_node_id, - ).to_payload() + if remote_session is None: + remote_session = RemoteSessionHandle( + session_id=uuid.uuid4().hex, + prompt_id=(str(prompt_id) if prompt_id is not None else None), + owner_component_id=component.representative_node_id, + ).to_payload() logger.info( "Split hybrid component %s into static nodes=%s and mapped nodes=%s using remote_session session_id=%s with %d static bridge outputs.", component.representative_node_id, @@ -3164,14 +3219,7 @@ def _build_component_payload( component_node_ids=list(component.static_node_ids), boundary_inputs=static_boundary_inputs, boundary_outputs=[ - { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - } + _boundary_output_payload(boundary_output) for boundary_output in static_boundary_outputs ] + static_bridge_outputs, @@ -3192,14 +3240,7 @@ def _build_component_payload( for boundary_spec in component.static_to_mapped_boundaries ], boundary_outputs=[ - { - "proxy_output_name": boundary_output.proxy_output_name, - "node_id": boundary_output.source.node_id, - "output_index": boundary_output.source.output_index, - "io_type": boundary_output.io_type, - "is_list": boundary_output.is_list, - "preview_target_node_ids": list(boundary_output.preview_target_node_ids), - } + _boundary_output_payload(boundary_output) for boundary_output in mapped_boundary_outputs ], execute_node_ids=list(component.mapped_execute_node_ids), @@ -3602,6 +3643,25 @@ def rewrite_prompt_for_modal( expanded_remote_node_ids, resolved_nodes_module, ) + session_component_ids = _mark_remote_to_remote_session_boundaries( + rewritten_prompt, + components, + ) + prompt_id = (extra_data or {}).get("prompt_id") + remote_sessions_by_component_id = { + component.representative_node_id: RemoteSessionHandle( + session_id=uuid.uuid4().hex, + prompt_id=(str(prompt_id) if prompt_id is not None else None), + owner_component_id=component.representative_node_id, + ).to_payload() + for component in components + if component.representative_node_id in session_component_ids + } + if remote_sessions_by_component_id: + logger.info( + "Enabled Modal-backed remote references for components=%s.", + sorted(remote_sessions_by_component_id), + ) validate_remote_component_transport_compatibility( prompt=rewritten_prompt, components=components, @@ -3682,6 +3742,9 @@ def rewrite_prompt_for_modal( uploaded_volume_paths=uploaded_volume_paths, terminate_container_on_error=resolved_settings.terminate_container_on_error, nodes_module=resolved_nodes_module, + remote_session=remote_sessions_by_component_id.get( + component.representative_node_id + ), ) proxy_node_ids = _rewrite_component_into_proxy( component=component, diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 7879dd2..2189b4d 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -126,7 +126,21 @@ _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", "LATENT"}) +_DURABLE_BRIDGE_SERIALIZATION_IO_TYPES = frozenset( + { + "AUDIO", + "BOOLEAN", + "CONDITIONING", + "FLOAT", + "IMAGE", + "INT", + "LATENT", + "MASK", + "SIGMAS", + "STRING", + "VIDEO", + } +) _DURABLE_BRIDGE_REHYDRATION_IO_TYPES = frozenset({"CLIP", "MODEL", "VAE"}) try: diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 2949ca0..9095a66 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -318,3 +318,5 @@ - Remote status maps use executable prompt paths such as `105:14`, while an inner node opened in the subgraph canvas still has the local id `14`. Node decoration must resolve `workflowNodePath(node)` before reading phase, progress, cache, component, or badge state; otherwise only the outer subgraph instance can reflect descendant execution. - An `IMAGE` tensor's leading dimension is not always an independently mappable batch. `CreateVideo.images` interprets it as one ordered frame sequence, so implicit boundary fan-out must preserve that complete tensor; splitting a 124-frame input runs `CreateVideo` 124 times and returns a list that local `SaveVideo` cannot consume. - Modal container task ids can be reused for many sequential payloads. Immediately stopping the local `modal container logs -f` watcher whenever one payload releases it makes the next watcher replay that container's complete history. Retain the watcher through a bounded idle grace period so nearby task reuse shares one continuous log cursor, then stop it if the task remains idle. +- A transport-safe type is not automatically a reason to materialize a component boundary locally. If every consumer of an `IMAGE`, `AUDIO`, `LATENT`, `MASK`, `SIGMAS`, `VIDEO`, or scalar output is another remote component, the planner should mark that output as a durable session bridge and let only its small reference pass through local ComfyUI. Any local consumer must keep the ordinary materialized boundary because local nodes do not understand bridge refs. +- Cross-container remote bridges need direct serialization for every supported transport type, not only `LATENT` and `CONDITIONING`. Otherwise the reference avoids local transfer but a fresh downstream worker can still replay the producer subgraph. Keep small metadata and payloads in Modal Dict, offload values above the configured inline threshold to content-addressed Volume objects, commit before publishing completion, and use the warm bridge cache or `Volume.read_file()` for restore. diff --git a/durable_state.py b/durable_state.py index 7bc7a8c..2be86bf 100644 --- a/durable_state.py +++ b/durable_state.py @@ -154,6 +154,13 @@ class FileDurableObjectStore: self._commit_callback() else: active_batch.wrote_object = True + logger.info( + "Resolved durable object namespace=%s path=%s size_bytes=%d created=%s.", + normalized_namespace, + relative_path.as_posix(), + len(payload), + wrote_object, + ) return DurableObjectRef( object_path=relative_path.as_posix(), sha256=sha256, diff --git a/pyproject.toml b/pyproject.toml index e0c190d..91266dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,8 @@ Repository = "https://github.com/ttulttul/ComfyUI-Modal" [dependency-groups] test = [ "aiohttp", + "comfy-aimdo==0.4.11", + "comfy-kitchen==0.2.26", "einops", "numpy", "pillow", @@ -29,6 +31,7 @@ test = [ "safetensors", "scipy", "torch", + "torchaudio>=2.11.0", "torchvision", "tqdm", "transformers", diff --git a/remote/modal_app.py b/remote/modal_app.py index 4703e8c..2e17956 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -112,7 +112,21 @@ _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", "LATENT"}) +_DURABLE_BRIDGE_SERIALIZATION_IO_TYPES = frozenset( + { + "AUDIO", + "BOOLEAN", + "CONDITIONING", + "FLOAT", + "IMAGE", + "INT", + "LATENT", + "MASK", + "SIGMAS", + "STRING", + "VIDEO", + } +) _DURABLE_BRIDGE_REHYDRATION_IO_TYPES = frozenset({"CLIP", "MODEL", "VAE"}) diff --git a/tests/conftest.py b/tests/conftest.py index bdc923e..3da6213 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -30,6 +30,7 @@ def _ensure_import_paths() -> None: if "av" not in sys.modules: av_module = types.ModuleType("av") av_module.__spec__ = importlib.machinery.ModuleSpec("av", loader=None) + av_module.__path__ = [] av_module.open = lambda *args, **kwargs: (_ for _ in ()).throw( RuntimeError("The av stub should not be used in tests.") ) @@ -51,11 +52,16 @@ def _ensure_import_paths() -> None: av_subtitles_module = types.ModuleType("av.subtitles") av_subtitles_stream_module = types.ModuleType("av.subtitles.stream") av_subtitles_stream_module.SubtitleStream = type("SubtitleStream", (), {}) + av_video_module = types.ModuleType("av.video") + av_video_reformatter_module = types.ModuleType("av.video.reformatter") + av_video_reformatter_module.ColorRange = type("ColorRange", (), {}) sys.modules["av"] = av_module sys.modules["av.container"] = av_container_module sys.modules["av.subtitles"] = av_subtitles_module sys.modules["av.subtitles.stream"] = av_subtitles_stream_module + sys.modules["av.video"] = av_video_module + sys.modules["av.video.reformatter"] = av_video_reformatter_module if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) comfyui_root = _comfyui_root() diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index ba1ed50..7f6e7ed 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -1073,11 +1073,12 @@ def test_rewrite_records_local_preview_targets_for_remote_boundary_images( def test_rewrite_splits_remote_chain_across_transportable_edges( api_intercept_module: Any, + modal_executor_module: Any, settings_module: Any, sync_engine_module: Any, tmp_path: Path, ) -> None: - """Transportable remote-to-remote edges should become ordered proxy boundaries.""" + """Transportable remote-to-remote edges should use Modal-backed references.""" settings = settings_module.ModalSyncSettings( app_name="app", auto_deploy=True, @@ -1160,6 +1161,7 @@ def test_rewrite_splits_remote_chain_across_transportable_edges( "io_type": "IMAGE", "is_list": False, "preview_target_node_ids": [], + "session_output": True, } ] @@ -1184,6 +1186,25 @@ def test_rewrite_splits_remote_chain_across_transportable_edges( assert rewritten_prompt["2"]["inputs"]["remote_input_0"] == ["1", 0] assert rewritten_prompt["3"]["inputs"]["image"] == ["2", 0] + first_execution_payload = modal_executor_module._rehydrate_proxy_payload( + first_payload, + unique_id="1", + ) + second_execution_payload = modal_executor_module._rehydrate_proxy_payload( + second_payload, + unique_id="2", + ) + assert "remote_session" not in first_payload + assert "remote_session" not in second_payload + assert first_execution_payload["clear_remote_session"] is True + assert second_execution_payload["clear_remote_session"] is True + assert first_execution_payload["remote_session"]["owner_component_id"] == "1" + assert second_execution_payload["remote_session"]["owner_component_id"] == "2" + assert ( + first_execution_payload["remote_session"]["session_id"] + != second_execution_payload["remote_session"]["session_id"] + ) + def test_rewrite_keeps_non_returning_local_preview_taps_local( api_intercept_module: Any, @@ -1278,6 +1299,18 @@ def test_rewrite_keeps_non_returning_local_preview_taps_local( assert all("9" not in payload["component_node_ids"] for payload in remote_payloads) assert all("9" not in payload["subgraph_prompt"] for payload in remote_payloads) assert all("9" not in payload["execute_node_ids"] for payload in remote_payloads) + assert rewritten_prompt["1"]["inputs"]["original_node_data"][ + "boundary_outputs" + ] == [ + { + "proxy_output_name": "1_image", + "node_id": "1", + "output_index": 0, + "io_type": "IMAGE", + "is_list": False, + "preview_target_node_ids": ["9"], + } + ] assert rewritten_prompt["9"]["inputs"]["images"] == ["1", 0] assert rewritten_prompt["3"]["inputs"]["image"] == ["2", 0] diff --git a/tests/test_remote_durability.py b/tests/test_remote_durability.py index 5e9eced..d555721 100644 --- a/tests/test_remote_durability.py +++ b/tests/test_remote_durability.py @@ -556,6 +556,57 @@ def test_cloud_large_bridge_inputs_and_output_use_object_store( assert torch.equal(restored_inputs["latent"]["samples"], latent["samples"]) +def test_cloud_large_image_bridge_output_uses_object_store( + modal_cloud_module: Any, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Remote IMAGE boundaries should restore from Volume bytes without replay.""" + import torch + + object_store = modal_cloud_module.FileDurableObjectStore(tmp_path) + monkeypatch.setattr(modal_cloud_module, "_DURABLE_OBJECT_STORE", object_store) + monkeypatch.setenv("COMFY_MODAL_BRIDGE_INLINE_MAX_BYTES", "1") + modal_cloud_module.get_settings.cache_clear() + image = torch.arange(48, dtype=torch.float32).reshape(1, 4, 4, 3) + target_handle = modal_cloud_module.RemoteSessionHandle( + session_id="session-image-object-backed", + prompt_id="prompt-1", + owner_component_id="component-2", + ) + + try: + record = modal_cloud_module._build_remote_session_bridge_record( + payload={"component_id": "component-1"}, + hydrated_inputs={}, + node_id="vae-decode-1", + output_index=0, + io_type="IMAGE", + output_value=image, + ) + monkeypatch.setattr( + modal_cloud_module, + "_execute_subgraph_prompt", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("durable IMAGE restore should skip producer replay") + ), + ) + restored_value = ( + modal_cloud_module._restore_serialized_remote_session_bridge_value( + record, + target_session_handle=target_handle, + ) + ) + finally: + modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle) + modal_cloud_module.get_settings.cache_clear() + + assert record.serialized_output is None + assert record.serialized_output_object is not None + assert record.serialized_output_io_type == "IMAGE" + assert torch.equal(restored_value, image) + + def test_local_client_attaches_stable_invocation_id( remote_modal_app_module: Any, monkeypatch: pytest.MonkeyPatch,