"""Tests split from the Modal executor integration suite.""" from __future__ import annotations from modal_executor_test_support import * # noqa: F401,F403 def test_modal_local_bridge_materializer_downloads_without_blocking_event_loop( modal_executor_module: Any, remote_modal_app_module: Any, session_state_module: Any, monkeypatch: Any, ) -> None: """The local bridge node should perform its durable download in a worker thread.""" materialization_started = threading.Event() release_materialization = threading.Event() def fake_materialize(ref_payload: dict[str, Any]) -> str: """Block the worker thread until the async test permits completion.""" assert session_state_module.is_remote_session_bridge_ref_payload(ref_payload) materialization_started.set() assert release_materialization.wait(timeout=1.0) return "local-image" monkeypatch.setattr( remote_modal_app_module, "materialize_remote_session_bridge_ref_locally", fake_materialize, ) bridge_ref = session_state_module.RemoteSessionBridgeRef( bridge_key="RSB_local_materializer", node_id="251", output_index=0, session_id="session-source", ).to_payload() async def run_scenario() -> tuple[Any, ...]: """Verify another coroutine runs while durable materialization is blocked.""" materializer_task = asyncio.create_task( modal_executor_module.ModalLocalBridgeMaterializer.execute(bridge_ref) ) await asyncio.to_thread(materialization_started.wait, 1.0) assert materializer_task.done() is False await asyncio.sleep(0) release_materialization.set() result = await asyncio.wait_for(materializer_task, timeout=1.0) return result.result assert asyncio.run(run_scenario()) == ("local-image",) def test_local_bridge_materialization_restores_durable_inline_output( remote_modal_app_module: Any, host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """A local branch should restore the producer's durable serialized output.""" import torch bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef( bridge_key="RSB_local_inline_bridge", node_id="251", output_index=0, session_id="session-source", ) image = torch.arange(12, dtype=torch.float32).reshape(1, 2, 2, 3) record = remote_modal_app_module.RemoteSessionBridgeRecord( bridge_key=bridge_ref.bridge_key, node_id=bridge_ref.node_id, output_index=bridge_ref.output_index, producer_payload={"component_id": "251"}, producer_inputs={}, serialized_output=remote_modal_app_module.serialize_value(image), serialized_output_io_type="IMAGE", ) monkeypatch.setattr( host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record, ) restored = host_session_bridge_module.materialize_remote_session_bridge_ref_locally( bridge_ref.to_payload() ) assert torch.equal(restored, image) def test_local_bridge_materialization_downloads_object_backed_output( remote_modal_app_module: Any, host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """A large bridge should download its content-addressed output object locally.""" monkeypatch.delenv("MODAL_ENVIRONMENT", raising=False) bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef( bridge_key="RSB_local_object_bridge", node_id="251", output_index=0, session_id="session-source", ) stored_payload = remote_modal_app_module.serialize_node_outputs(("large-image",)) object_ref = remote_modal_app_module.DurableObjectRef( object_path="bridge-outputs/sha256/value.bin", sha256=hashlib.sha256(stored_payload).hexdigest(), size_bytes=len(stored_payload), ) record = remote_modal_app_module.RemoteSessionBridgeRecord( bridge_key=bridge_ref.bridge_key, node_id=bridge_ref.node_id, output_index=bridge_ref.output_index, producer_payload={"component_id": "251"}, producer_inputs={}, serialized_output_object=object_ref, serialized_output_io_type="IMAGE", ) dict_calls: list[tuple[Any, ...]] = [] volume_calls: list[tuple[Any, ...]] = [] class FakeDict: """Expose the durable bridge record through Modal's shared Dict API.""" @staticmethod def from_name(*args: Any, **kwargs: Any) -> Any: """Return a Dict handle that records bridge lookups.""" dict_calls.append((*args, kwargs)) return types.SimpleNamespace(get=lambda bridge_key: record.to_payload()) class FakeVolume: """Expose the durable object through Modal's direct Volume API.""" @staticmethod def from_name(*args: Any, **kwargs: Any) -> Any: """Return a Volume handle that records object downloads.""" volume_calls.append((*args, kwargs)) return types.SimpleNamespace( read_file=lambda volume_path: iter((stored_payload[:7], stored_payload[7:])) ) monkeypatch.setattr( host_session_bridge_module, "get_settings", lambda: types.SimpleNamespace( execution_mode="remote", session_bridge_dict_name="bridge-dict", volume_name="bridge-volume", ), ) monkeypatch.setattr( host_session_bridge_module, "modal", types.SimpleNamespace(Dict=FakeDict, Volume=FakeVolume), ) host_session_bridge_module._MODAL_SESSION_BRIDGE_DICTS.clear() host_session_bridge_module._MODAL_DURABLE_VOLUMES.clear() restored = host_session_bridge_module.materialize_remote_session_bridge_ref_locally( bridge_ref.to_payload() ) assert restored == "large-image" assert dict_calls == [ ( "bridge-dict", {"environment_name": None, "create_if_missing": True}, ) ] assert volume_calls == [ ( "bridge-volume", {"environment_name": None, "create_if_missing": True}, ) ] host_session_bridge_module._MODAL_SESSION_BRIDGE_DICTS.clear() host_session_bridge_module._MODAL_DURABLE_VOLUMES.clear() @pytest.mark.parametrize("hidden_delivery", ["v3_class_clone", "legacy_kwargs"]) def test_cache_friendly_proxy_payload_rehydrates_prompt_id_at_execution( modal_executor_module: Any, hidden_delivery: str, ) -> None: """V3 and legacy hidden inputs should select the correct overlapping prompt.""" fake_nodes_module = type( "FakeNodesModule", (), { "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode}, "NODE_DISPLAY_NAME_MAPPINGS": {}, }, )() proxy_id = modal_executor_module.ensure_modal_proxy_node_registered( original_class_type="OriginalNode", original_class=_FakeOriginalNode, nodes_module=fake_nodes_module, ) proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id] payload = modal_executor_module.register_cache_friendly_proxy_payload( "node-1", { "payload_kind": "subgraph", "component_id": "component-1", "prompt_id": "prompt-1", "boundary_outputs": [], "execute_node_ids": [], }, ) next_payload = modal_executor_module.register_cache_friendly_proxy_payload( "node-1", { "payload_kind": "subgraph", "component_id": "component-1", "prompt_id": "prompt-2", "boundary_outputs": [], "execute_node_ids": [], }, ) class FakeClient: """Test client that captures the rehydrated payload.""" async def execute_payload_async( self, payload: dict[str, Any], kwargs: dict[str, Any], ) -> tuple[str, int]: """Return values derived from the restored prompt id.""" return (str(payload.get("prompt_id")), len(kwargs)) modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient()) try: hidden_metadata = { modal_executor_module.MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY: "prompt-1" } execution_proxy_class = proxy_class hidden_kwargs: dict[str, Any] = {} if hidden_delivery == "v3_class_clone": execution_proxy_class = proxy_class.PREPARE_CLASS_CLONE( { "hidden_inputs": { modal_executor_module.io.Hidden.unique_id: "node-1", modal_executor_module.io.Hidden.extra_pnginfo: hidden_metadata, } } ) else: hidden_kwargs = { "unique_id": ["node-1"], "extra_pnginfo": [hidden_metadata], } result = asyncio.run( execution_proxy_class.execute( original_node_data=[payload], value=["payload"], **hidden_kwargs, ) ) finally: modal_executor_module.set_remote_executor_client_factory(None) assert "prompt_id" not in payload assert payload == next_payload assert result.result == ("prompt-1", 1) def test_register_cache_friendly_proxy_payload_strips_session_fields_and_rehydrates_them( modal_executor_module: Any, ) -> None: """Session-backed proxy payloads should strip run-scoped fields from the local cache surface.""" payload = modal_executor_module.register_cache_friendly_proxy_payload( "node-1", { "payload_kind": "subgraph", "component_id": "component-1", "prompt_id": "prompt-1", "remote_session": { "session_id": "session-1", "prompt_id": "prompt-1", "owner_component_id": "component-1", }, "boundary_outputs": [], "execute_node_ids": [], "clear_remote_session": True, "extra_data": { "prompt_id": "prompt-1", "create_time": 1234567890, "modal": {"remote_node_ids": ["12"], "estimated_max_parallel_requests": 1}, }, "requires_volume_reload": True, "volume_reload_marker": "marker-1", "uploaded_volume_paths": ["/storage/assets/example.safetensors"], }, ) assert "prompt_id" not in payload assert "remote_session" not in payload assert "clear_remote_session" not in payload assert "extra_data" not in payload assert "requires_volume_reload" not in payload assert "volume_reload_marker" not in payload assert "uploaded_volume_paths" not in payload assert modal_executor_module._rehydrate_proxy_payload(payload, unique_id="node-1") == { "payload_kind": "subgraph", "component_id": "component-1", "prompt_id": "prompt-1", "remote_session": { "session_id": "session-1", "prompt_id": "prompt-1", "owner_component_id": "component-1", }, "boundary_outputs": [], "execute_node_ids": [], "clear_remote_session": True, "extra_data": { "prompt_id": "prompt-1", "create_time": 1234567890, "modal": {"remote_node_ids": ["12"], "estimated_max_parallel_requests": 1}, }, "requires_volume_reload": True, "volume_reload_marker": "marker-1", "uploaded_volume_paths": ["/storage/assets/example.safetensors"], } def test_cache_friendly_proxy_payload_rehydrates_without_hidden_unique_id( modal_executor_module: Any, ) -> None: """Cache-friendly proxy payloads should still rehydrate when ComfyUI omits hidden unique_id.""" payload = modal_executor_module.register_cache_friendly_proxy_payload( "node-7", { "payload_kind": "subgraph", "component_id": "component-7", "prompt_id": "prompt-7", "remote_session": { "session_id": "session-7", "prompt_id": "prompt-7", "owner_component_id": "component-7", }, "boundary_outputs": [], "execute_node_ids": [], }, ) assert modal_executor_module._rehydrate_proxy_payload(payload, unique_id=None) == { "payload_kind": "subgraph", "component_id": "component-7", "prompt_id": "prompt-7", "remote_session": { "session_id": "session-7", "prompt_id": "prompt-7", "owner_component_id": "component-7", }, "boundary_outputs": [], "execute_node_ids": [], } 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)}, ] ] _patch_cloud_session_bridge( monkeypatch, "_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", ), ) _patch_cloud_session_bridge( monkeypatch, "_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 = _cloud_remote_session_store().get_output( modal_cloud_module.RemoteSessionValueRef( session_id=target_handle.session_id, node_id="node-9", output_index=0, ) ) finally: _cloud_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 @pytest.mark.parametrize( ("io_type", "node_id", "class_type", "node_inputs"), [ ("NOISE", "15", "RandomNoise", {"noise_seed": 42}), ("SAMPLER", "17", "KSamplerSelect", {"sampler_name": "euler"}), ], ) def test_modal_cloud_rehydrates_literal_sampling_strategy_bridges_without_replay( modal_cloud_module: Any, monkeypatch: Any, io_type: str, node_id: str, class_type: str, node_inputs: dict[str, Any], ) -> None: """Literal NOISE and SAMPLER strategy nodes should rebuild without producer replay.""" 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=f"RSB_{io_type.lower()}_bridge", node_id=node_id, output_index=0, session_id="session-source", ) record = modal_cloud_module._build_remote_session_bridge_record( payload={ "component_id": "image-preview-component", "execute_node_ids": ["251"], "subgraph_prompt": { "14": { "class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["15", 0], "sampler": ["17", 0]}, }, "15": { "class_type": "RandomNoise", "inputs": {"noise_seed": 42}, }, "17": { "class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}, }, "250": { "class_type": "VAEDecode", "inputs": {"samples": ["14", 0]}, }, "251": { "class_type": "ImageFromBatch", "inputs": {"image": ["250", 0]}, }, }, }, hydrated_inputs={}, node_id=node_id, output_index=0, io_type=io_type, output_value=object(), ) assert record.recovery_kind is modal_cloud_module.RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN assert record.rehydration_plan == { "kind": "single_node_output", "node_data": {"class_type": class_type}, "node_inputs": node_inputs, } execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = [] _patch_cloud_session_bridge( monkeypatch, "_load_remote_session_bridge_record", lambda bridge_key: record, ) _patch_cloud_session_bridge( monkeypatch, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("literal strategy bridge should not replay its producer component") ), ) _patch_cloud_session_bridge( monkeypatch, "_execute_node_locally_raw", lambda node_data, kwargs_payload, **kwargs: ( execute_calls.append((dict(node_data), dict(kwargs_payload))), (f"restored-{io_type.lower()}",), )[1], ) _patch_cloud_session_bridge( monkeypatch, "_store_remote_session_bridge_value", lambda bridge_key, value: None, ) 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, ) finally: _cloud_remote_session_store().clear_session(target_handle) assert restored_value == f"restored-{io_type.lower()}" assert execute_calls == [({"class_type": class_type}, node_inputs)] assert resolution_stats.durable_bridge_hits == 1 assert resolution_stats.replay_count == 0 def test_modal_cloud_rehydrates_sampler_latent_bridge_refs_from_durable_record_without_replay( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Durably serialized LATENT bridge values should restore without replaying sampler producers.""" 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_sampler_latent_bridge", node_id="sampler-1", output_index=0, session_id="session-source", ) latent = {"samples": torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2)} record = modal_cloud_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], "subgraph_prompt": { "sampler-1": { "class_type": "KSampler", "inputs": {}, } }, }, hydrated_inputs={}, node_id="sampler-1", output_index=0, io_type="LATENT", output_value=latent, ) assert record.serialized_output is not None assert record.serialized_output_io_type == "LATENT" _patch_cloud_session_bridge( monkeypatch, "_load_remote_session_bridge_record", lambda bridge_key: record, ) _patch_cloud_session_bridge( monkeypatch, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("durable LATENT restore should skip sampler 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 = _cloud_remote_session_store().get_output( modal_cloud_module.RemoteSessionValueRef( session_id=target_handle.session_id, node_id="sampler-1", output_index=0, ) ) finally: _cloud_remote_session_store().clear_session(target_handle) assert torch.equal(restored_value["samples"], latent["samples"]) assert torch.equal(stored_value["samples"], latent["samples"]) 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_rehydrates_model_bridge_refs_from_durable_plan_without_replay( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Durable MODEL bridge plans should rebuild one self-contained loader output without replay.""" 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_model_bridge", node_id="node-5", output_index=0, session_id="session-source", ) record = modal_cloud_module._build_remote_session_bridge_record( payload={ "component_id": "component-seed", "subgraph_prompt": { "node-5": { "class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}, } }, }, hydrated_inputs={}, node_id="node-5", output_index=0, io_type="MODEL", output_value=_FakeModelValue("seed-model"), ) execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = [] _patch_cloud_session_bridge(monkeypatch, "_load_remote_session_bridge_record", lambda bridge_key: record) _patch_cloud_session_bridge( monkeypatch, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("durable MODEL rehydration should skip replay") ), ) _patch_cloud_session_bridge( monkeypatch, "_execute_node_locally_raw", lambda node_data, kwargs_payload, **kwargs: ( execute_calls.append((dict(node_data), dict(kwargs_payload))), (_FakeModelValue("restored-model"),), )[1], ) 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 = _cloud_remote_session_store().get_output( modal_cloud_module.RemoteSessionValueRef( session_id=target_handle.session_id, node_id="node-5", output_index=0, ) ) finally: _cloud_remote_session_store().clear_session(target_handle) assert isinstance(restored_value, _FakeModelValue) assert restored_value.value == "restored-model" assert stored_value is restored_value assert execute_calls == [ ( {"class_type": "CheckpointLoaderSimple"}, {"ckpt_name": "model.safetensors"}, ) ] 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_rehydrates_clip_bridge_refs_from_durable_plan_without_replay( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Durable CLIP bridge plans should rebuild one self-contained loader output without replay.""" 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_clip_bridge", node_id="clip-1", output_index=0, session_id="session-source", ) record = modal_cloud_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], "subgraph_prompt": { "clip-1": { "class_type": "CLIPLoader", "inputs": { "clip_name": "text.safetensors", "type": "flux", "device": "default", }, }, "sampler-1": { "class_type": "KSampler", "inputs": {"clip": ["clip-1", 0]}, }, }, }, hydrated_inputs={}, node_id="clip-1", output_index=0, io_type="CLIP", output_value="seed-clip", ) assert record.rehydration_plan is not None assert record.rehydration_plan_io_type == "CLIP" _patch_cloud_session_bridge(monkeypatch, "_load_remote_session_bridge_record", lambda bridge_key: record) _patch_cloud_session_bridge( monkeypatch, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("durable CLIP rehydration should skip replay") ), ) _patch_cloud_session_bridge( monkeypatch, "_execute_node_locally_raw", lambda node_data, kwargs_payload, **kwargs: ( "clip::" f"{kwargs_payload['clip_name']}:" f"{kwargs_payload['type']}:" f"{kwargs_payload['device']}", ), ) 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, ) finally: _cloud_remote_session_store().clear_session(target_handle) assert restored_value == "clip::text.safetensors:flux:default" assert resolution_stats.durable_bridge_hits == 1 assert resolution_stats.replay_count == 0 def test_modal_cloud_rehydrates_linked_model_bridge_with_non_sampler_subgraph_plan( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Linked MODEL bridge plans should rerun only the non-sampler model dependency closure.""" 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_linked_model_bridge", node_id="lora-1", output_index=0, session_id="session-source", ) record = modal_cloud_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], "subgraph_prompt": { "loader-1": { "class_type": "UNETLoader", "inputs": {"unet_name": "base.safetensors"}, }, "lora-1": { "class_type": "LoraLoaderModelOnly", "inputs": {"model": ["loader-1", 0], "lora_name": "style.safetensors"}, }, "sampler-1": { "class_type": "KSampler", "inputs": {"model": ["lora-1", 0]}, }, }, }, hydrated_inputs={}, node_id="lora-1", output_index=0, io_type="MODEL", output_value=_FakeModelValue("lora-model"), ) assert record.rehydration_plan is not None assert record.rehydration_plan["kind"] == "subgraph_output" _patch_cloud_session_bridge(monkeypatch, "_load_remote_session_bridge_record", lambda bridge_key: record) observed_payloads: list[dict[str, Any]] = [] def fake_execute_subgraph_prompt( payload: dict[str, Any], hydrated_inputs: dict[str, Any], *_args: Any, ) -> tuple[_FakeModelValue]: """Assert the durable plan avoids rerunning the sampler component.""" del hydrated_inputs observed_payloads.append(copy.deepcopy(payload)) assert payload["execute_node_ids"] == ["lora-1"] assert payload["remote_session"]["session_id"] == target_handle.session_id assert sorted(payload["subgraph_prompt"]) == ["loader-1", "lora-1"] assert "sampler-1" not in payload["subgraph_prompt"] return (_FakeModelValue("restored-lora-model"),) _patch_cloud_session_bridge(monkeypatch, "_execute_subgraph_prompt", fake_execute_subgraph_prompt) 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, ) finally: _cloud_remote_session_store().clear_session(target_handle) assert isinstance(restored_value, _FakeModelValue) assert restored_value.value == "restored-lora-model" assert len(observed_payloads) == 1 assert resolution_stats.durable_bridge_hits == 1 assert resolution_stats.replay_count == 0 def test_modal_cloud_refuses_sampler_ancestor_bridge_replay( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Producer replay should fail when a terminal node's dependency closure samples.""" 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_sampler_bridge", node_id="unsupported-bridge", output_index=0, session_id="session-source", ) _patch_cloud_session_bridge( monkeypatch, "_load_remote_session_bridge_record", lambda bridge_key: modal_cloud_module.RemoteSessionBridgeRecord( bridge_key=bridge_key, node_id="unsupported-bridge", output_index=0, producer_payload={ "component_id": "sampler-component", "execute_node_ids": ["251"], "subgraph_prompt": { "14": {"class_type": "KSampler", "inputs": {}}, "250": { "class_type": "VAEDecode", "inputs": {"samples": ["14", 0]}, }, "251": { "class_type": "ImageFromBatch", "inputs": {"image": ["250", 0]}, }, }, }, producer_inputs={}, ), ) _patch_cloud_session_bridge( monkeypatch, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("sampler bridge replay should be blocked") ), ) resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats() with pytest.raises(modal_cloud_module.RemoteSessionStateError, match="rerun a sampler"): 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, ) assert resolution_stats.bridge_record_lookups == 1 assert resolution_stats.replay_count == 0 def test_modal_cloud_skips_seed_execution_when_session_outputs_are_already_restored( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Cloud seed payloads should no-op when bridge outputs are already in session memory.""" session_handle = modal_cloud_module.RemoteSessionHandle( session_id="session-restored", prompt_id="prompt-1", owner_component_id="component-1", ) restored_value = _FakeModelValue("restored-model") _cloud_remote_session_store().put_output( session_handle, node_id="5", output_index=0, value=restored_value, ) payload = { "payload_kind": "subgraph", "component_id": "1__mapped::seed:0", "prompt_id": "prompt-1", "component_node_ids": ["5"], "subgraph_prompt": { "5": { "class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}, } }, "boundary_inputs": [ { "proxy_input_name": "static_input_0", "io_type": "MODEL", "targets": [{"node_id": "5", "input_name": "model"}], } ], "boundary_outputs": [ { "proxy_output_name": "static_input_0", "node_id": "5", "output_index": 0, "io_type": "MODEL", "is_list": False, "session_output": True, } ], "execute_node_ids": ["5"], "remote_session": session_handle.to_payload(), } hydrated_inputs = { "static_input_0": modal_cloud_module.RemoteSessionValueRef( session_id=session_handle.session_id, node_id="5", output_index=0, ).to_payload() } monkeypatch.setattr( _cloud_prompt_execution_owner(), "_load_execution_module", lambda: (_ for _ in ()).throw(AssertionError("seed short-circuit should skip PromptExecutor")), ) try: outputs = modal_cloud_module._execute_subgraph_prompt( payload, hydrated_inputs, None, ) finally: _cloud_remote_session_store().clear_session(session_handle) assert len(outputs) == 1 assert modal_cloud_module.is_remote_session_bridge_ref_payload(outputs[0]) def test_modal_cloud_logs_remote_session_resolution_summary( modal_cloud_module: Any, monkeypatch: Any, ) -> None: """Mapped-phase bridge resolution should emit one summary block with replay and loader deltas.""" observed_logs: list[tuple[str, tuple[Any, ...]]] = [] _patch_cloud_session_bridge( monkeypatch, "_emit_cloud_info", lambda message, *args: observed_logs.append((message, args)), ) modal_cloud_module._log_remote_session_resolution_summary( component_id="1::mapped", resolution_stats=modal_cloud_module._RemoteSessionBridgeResolutionStats( 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, replay_seconds=1.5, direct_restore_seconds=0.02, session_restore_writes=1, ), loader_cache_before={"hit": 3, "miss": 4}, loader_cache_after={"hit": 5, "miss": 5}, ) assert observed_logs == [ ( "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), ) ] 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( _cloud_node_output_cache_owner(), "_load_nodes_module", lambda: types.SimpleNamespace( NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})} ), ) monkeypatch.setattr( _cloud_node_output_cache_owner(), "_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_execute_subgraph_locally_round_trips_remote_session_bridge_refs( remote_modal_app_module: Any, host_session_bridge_module: Any, serialization_module: Any, session_state_module: Any, ) -> None: """Split proxy subgraphs should return durable bridge refs that survive lost sessions.""" static_payload = { "payload_kind": "subgraph", "component_id": "1", "prompt_id": "prompt-1", "component_node_ids": ["1"], "subgraph_prompt": { "1": { "class_type": "SessionValueNode", "inputs": {}, } }, "boundary_inputs": [], "boundary_outputs": [ { "proxy_output_name": "static_input_0", "node_id": "1", "output_index": 0, "io_type": "MODEL", "is_list": False, "session_output": True, } ], "execute_node_ids": ["1"], "remote_session": session_state_module.RemoteSessionHandle( session_id="session-runtime-1", prompt_id="prompt-1", owner_component_id="1", ).to_payload(), } static_outputs = serialization_module.deserialize_node_outputs( remote_modal_app_module.execute_subgraph_locally( static_payload, serialization_module.serialize_node_inputs({}), node_mapping={"SessionValueNode": _FakeSessionValueNode}, ) ) assert session_state_module.is_remote_session_bridge_ref_payload(static_outputs[0]) host_session_bridge_module._REMOTE_SESSION_STORE.clear_session( session_state_module.RemoteSessionHandle.from_payload(static_payload["remote_session"]) ) mapped_payload = { "payload_kind": "subgraph", "component_id": "1__mapped", "prompt_id": "prompt-1", "component_node_ids": ["7"], "subgraph_prompt": { "7": { "class_type": "SessionEchoNode", "inputs": {"text": ["static_input_0", 0]}, } }, "boundary_inputs": [ { "proxy_input_name": "static_input_0", "io_type": "MODEL", "targets": [{"node_id": "7", "input_name": "text"}], } ], "boundary_outputs": [ { "proxy_output_name": "7_text", "node_id": "7", "output_index": 0, "io_type": "STRING", "is_list": False, } ], "execute_node_ids": ["7"], "remote_session": session_state_module.RemoteSessionHandle( session_id="session-runtime-2", prompt_id="prompt-1", owner_component_id="1", ).to_payload(), "clear_remote_session": True, } mapped_outputs = serialization_module.deserialize_node_outputs( remote_modal_app_module.execute_subgraph_locally( mapped_payload, serialization_module.serialize_node_inputs({"static_input_0": static_outputs[0]}), node_mapping={ "SessionEchoNode": _FakeSessionEchoNode, "SessionValueNode": _FakeSessionValueNode, }, ) ) assert mapped_outputs == ("shared-session-value",) replay_outputs = remote_modal_app_module._execute_subgraph_prompt( { **mapped_payload, "remote_session": session_state_module.RemoteSessionHandle( session_id="session-runtime-3", prompt_id="prompt-1", owner_component_id="1", ).to_payload(), "clear_remote_session": False, }, {"static_input_0": static_outputs[0]}, { "SessionEchoNode": _FakeSessionEchoNode, "SessionValueNode": _FakeSessionValueNode, }, ) assert replay_outputs == ("shared-session-value",) def test_local_fallback_skips_seed_execution_when_session_outputs_are_already_restored( remote_modal_app_module: Any, host_session_bridge_module: Any, local_execution_module: Any, monkeypatch: Any, ) -> None: """Local fallback seed payloads should no-op when bridge outputs are already in session memory.""" session_handle = remote_modal_app_module.RemoteSessionHandle( session_id="session-restored", prompt_id="prompt-1", owner_component_id="component-1", ) restored_value = _FakeModelValue("restored-model") host_session_bridge_module._REMOTE_SESSION_STORE.put_output( session_handle, node_id="5", output_index=0, value=restored_value, ) payload = { "payload_kind": "subgraph", "component_id": "1__mapped::seed:0", "prompt_id": "prompt-1", "component_node_ids": ["5"], "subgraph_prompt": { "5": { "class_type": "ShouldNotRun", "inputs": {"model": ["static_input_0", 0]}, } }, "boundary_inputs": [ { "proxy_input_name": "static_input_0", "io_type": "MODEL", "targets": [{"node_id": "5", "input_name": "model"}], } ], "boundary_outputs": [ { "proxy_output_name": "static_input_0", "node_id": "5", "output_index": 0, "io_type": "MODEL", "is_list": False, "session_output": True, } ], "execute_node_ids": ["5"], "remote_session": session_handle.to_payload(), } hydrated_inputs = { "static_input_0": remote_modal_app_module.RemoteSessionValueRef( session_id=session_handle.session_id, node_id="5", output_index=0, ).to_payload() } monkeypatch.setattr( local_execution_module, "_invoke_original_node", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("seed short-circuit should skip local node execution") ), ) try: outputs = remote_modal_app_module._execute_subgraph_prompt( payload, hydrated_inputs, node_mapping={"ShouldNotRun": type("ShouldNotRun", (), {})}, ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(session_handle) assert len(outputs) == 1 assert remote_modal_app_module.is_remote_session_bridge_ref_payload(outputs[0]) def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_record_without_replay( remote_modal_app_module: Any, host_session_bridge_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( host_session_bridge_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 = host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping=None, resolution_stats=resolution_stats, ) stored_value = host_session_bridge_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: host_session_bridge_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_remote_session_input_resolution_handles_nested_conditioning_bridge_refs( remote_modal_app_module: Any, host_session_bridge_module: Any, ) -> None: """Boundary input resolution should resolve bridge refs nested inside list outputs.""" source_handle = remote_modal_app_module.RemoteSessionHandle( session_id="session-source", prompt_id="prompt-1", owner_component_id="component-1", ) first_ref = host_session_bridge_module._REMOTE_SESSION_STORE.put_bridge_output( source_handle, bridge_key="RSB_a", node_id="508:item:0", output_index=0, value=[["conditioning-a", {"pooled_output": "pool-a"}]], ) second_ref = host_session_bridge_module._REMOTE_SESSION_STORE.put_bridge_output( source_handle, bridge_key="RSB_b", node_id="508:item:1", output_index=0, value=[["conditioning-b", {"pooled_output": "pool-b"}]], ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: resolved_inputs = host_session_bridge_module._resolve_remote_session_inputs( { "positive": [ remote_modal_app_module.RemoteSessionBridgeRef( bridge_key="RSB_a", node_id=first_ref.node_id, output_index=first_ref.output_index, session_id=first_ref.session_id, ).to_payload(), remote_modal_app_module.RemoteSessionBridgeRef( bridge_key="RSB_b", node_id=second_ref.node_id, output_index=second_ref.output_index, session_id=second_ref.session_id, ).to_payload(), ] }, component_id="sampler-component", resolution_stats=resolution_stats, ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(source_handle) assert resolved_inputs["positive"] == [ [["conditioning-a", {"pooled_output": "pool-a"}]], [["conditioning-b", {"pooled_output": "pool-b"}]], ] assert resolution_stats.input_ref_count == 2 def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_record_without_replay( remote_modal_app_module: Any, host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """The local fallback should restore serialized LATENT bridge values without sampler 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_sampler_latent_bridge", node_id="sampler-1", output_index=0, session_id="session-source", ) latent = {"samples": torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2)} record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], "subgraph_prompt": { "sampler-1": { "class_type": "KSampler", "inputs": {}, } }, }, hydrated_inputs={}, node_id="sampler-1", output_index=0, io_type="LATENT", output_value=latent, ) assert record.serialized_output is not None assert record.serialized_output_io_type == "LATENT" monkeypatch.setattr( host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record, ) monkeypatch.setattr( remote_modal_app_module, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("durable LATENT restore should skip local sampler replay") ), ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: restored_value = host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping=None, resolution_stats=resolution_stats, ) stored_value = host_session_bridge_module._REMOTE_SESSION_STORE.get_output( remote_modal_app_module.RemoteSessionValueRef( session_id=target_handle.session_id, node_id="sampler-1", output_index=0, ) ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert torch.equal(restored_value["samples"], latent["samples"]) assert torch.equal(stored_value["samples"], latent["samples"]) 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 @pytest.mark.parametrize( ("io_type", "node_id", "class_type", "node_inputs"), [ ("NOISE", "15", "RandomNoise", {"noise_seed": 42}), ("SAMPLER", "17", "KSamplerSelect", {"sampler_name": "euler"}), ], ) def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_replay( remote_modal_app_module: Any, host_session_bridge_module: Any, local_execution_module: Any, monkeypatch: Any, io_type: str, node_id: str, class_type: str, node_inputs: dict[str, Any], ) -> None: """The local fallback should rebuild literal sampling strategies without replay.""" 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=f"RSB_local_{io_type.lower()}_bridge", node_id=node_id, output_index=0, session_id="session-source", ) record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "image-preview-component", "execute_node_ids": ["251"], "subgraph_prompt": { "14": { "class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["15", 0], "sampler": ["17", 0]}, }, "15": { "class_type": "RandomNoise", "inputs": {"noise_seed": 42}, }, "17": { "class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}, }, "250": { "class_type": "VAEDecode", "inputs": {"samples": ["14", 0]}, }, "251": { "class_type": "ImageFromBatch", "inputs": {"image": ["250", 0]}, }, }, }, hydrated_inputs={}, node_id=node_id, output_index=0, io_type=io_type, output_value=object(), ) assert ( record.recovery_kind is remote_modal_app_module.RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN ) assert record.rehydration_plan == { "kind": "single_node_output", "node_data": {"class_type": class_type}, "node_inputs": node_inputs, } execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = [] monkeypatch.setattr( host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record, ) monkeypatch.setattr( local_execution_module, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("literal strategy bridge should not replay its producer component") ), ) monkeypatch.setattr( local_execution_module, "_execute_node_locally_raw", lambda node_data, kwargs_payload, **kwargs: ( execute_calls.append((dict(node_data), dict(kwargs_payload))), (f"restored-{io_type.lower()}",), )[1], ) monkeypatch.setattr( remote_modal_app_module, "_store_remote_session_bridge_value", lambda bridge_key, value: None, ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: restored_value = host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping=None, resolution_stats=resolution_stats, ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert restored_value == f"restored-{io_type.lower()}" assert execute_calls == [({"class_type": class_type}, node_inputs)] assert resolution_stats.durable_bridge_hits == 1 assert resolution_stats.replay_count == 0 def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without_replay( remote_modal_app_module: Any, host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """The local fallback should rebuild one MODEL bridge output from a durable plan without replay.""" 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_model_bridge", node_id="node-5", output_index=0, session_id="session-source", ) record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "component-seed", "subgraph_prompt": { "node-5": { "class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}, } }, }, hydrated_inputs={}, node_id="node-5", output_index=0, io_type="MODEL", output_value=_FakeModelValue("seed-model"), ) monkeypatch.setattr(host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record) monkeypatch.setattr( remote_modal_app_module, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("durable MODEL rehydration should skip local replay") ), ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: restored_value = host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping={"CheckpointLoaderSimple": _FakeModelLoaderNode}, resolution_stats=resolution_stats, ) stored_value = host_session_bridge_module._REMOTE_SESSION_STORE.get_output( remote_modal_app_module.RemoteSessionValueRef( session_id=target_handle.session_id, node_id="node-5", output_index=0, ) ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert isinstance(restored_value, _FakeModelValue) assert restored_value.value == "model::model.safetensors" assert stored_value is restored_value 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_remote_app_rehydrates_clip_bridge_refs_from_durable_plan_without_replay( remote_modal_app_module: Any, host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """The local fallback should rebuild one CLIP bridge output from a durable plan.""" 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_clip_bridge", node_id="clip-1", output_index=0, session_id="session-source", ) record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], "subgraph_prompt": { "clip-1": { "class_type": "CLIPLoader", "inputs": { "clip_name": "text.safetensors", "type": "flux", "device": "default", }, }, "sampler-1": { "class_type": "KSampler", "inputs": {"clip": ["clip-1", 0]}, }, }, }, hydrated_inputs={}, node_id="clip-1", output_index=0, io_type="CLIP", output_value="seed-clip", ) assert record.rehydration_plan is not None assert record.rehydration_plan_io_type == "CLIP" monkeypatch.setattr(host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record) monkeypatch.setattr( remote_modal_app_module, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("durable CLIP rehydration should skip local replay") ), ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: restored_value = host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping={"CLIPLoader": _FakeClipLoaderNode}, resolution_stats=resolution_stats, ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert restored_value == "clip::text.safetensors:flux:default" assert resolution_stats.durable_bridge_hits == 1 assert resolution_stats.replay_count == 0 def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgraph_plan( remote_modal_app_module: Any, host_session_bridge_module: Any, local_execution_module: Any, monkeypatch: Any, ) -> None: """The local fallback should rehydrate linked MODEL bridges without sampler replay.""" 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_linked_model_bridge", node_id="lora-1", output_index=0, session_id="session-source", ) record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], "subgraph_prompt": { "loader-1": { "class_type": "UNETLoader", "inputs": {"unet_name": "base.safetensors"}, }, "lora-1": { "class_type": "LoraLoaderModelOnly", "inputs": {"model": ["loader-1", 0], "lora_name": "style.safetensors"}, }, "sampler-1": { "class_type": "KSampler", "inputs": {"model": ["lora-1", 0]}, }, }, }, hydrated_inputs={}, node_id="lora-1", output_index=0, io_type="MODEL", output_value=_FakeModelValue("lora-model"), ) assert record.rehydration_plan is not None assert record.rehydration_plan["kind"] == "subgraph_output" monkeypatch.setattr(host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record) observed_payloads: list[dict[str, Any]] = [] def fake_execute_subgraph_prompt( payload: dict[str, Any], hydrated_inputs: dict[str, Any], _node_mapping: dict[str, type[Any]] | None = None, ) -> tuple[_FakeModelValue]: """Assert the durable plan avoids rerunning the sampler component.""" del hydrated_inputs, _node_mapping observed_payloads.append(copy.deepcopy(payload)) assert payload["execute_node_ids"] == ["lora-1"] assert payload["remote_session"]["session_id"] == target_handle.session_id assert sorted(payload["subgraph_prompt"]) == ["loader-1", "lora-1"] assert "sampler-1" not in payload["subgraph_prompt"] return (_FakeModelValue("restored-lora-model"),) monkeypatch.setattr(local_execution_module, "_execute_subgraph_prompt", fake_execute_subgraph_prompt) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: restored_value = host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping=None, resolution_stats=resolution_stats, ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert isinstance(restored_value, _FakeModelValue) assert restored_value.value == "restored-lora-model" assert len(observed_payloads) == 1 assert resolution_stats.durable_bridge_hits == 1 assert resolution_stats.replay_count == 0 def test_local_remote_app_refuses_sampler_ancestor_bridge_replay( remote_modal_app_module: Any, host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """The local fallback should reject replay when a terminal depends on sampling.""" 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_sampler_bridge", node_id="unsupported-bridge", output_index=0, session_id="session-source", ) monkeypatch.setattr( host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: remote_modal_app_module.RemoteSessionBridgeRecord( bridge_key=bridge_key, node_id="unsupported-bridge", output_index=0, producer_payload={ "component_id": "sampler-component", "execute_node_ids": ["251"], "subgraph_prompt": { "14": {"class_type": "KSampler", "inputs": {}}, "250": { "class_type": "VAEDecode", "inputs": {"samples": ["14", 0]}, }, "251": { "class_type": "ImageFromBatch", "inputs": {"image": ["250", 0]}, }, }, }, producer_inputs={}, ), ) monkeypatch.setattr( remote_modal_app_module, "_execute_subgraph_prompt", lambda *args, **kwargs: (_ for _ in ()).throw( AssertionError("sampler bridge replay should be blocked") ), ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: with pytest.raises(remote_modal_app_module.RemoteSessionStateError, match="rerun a sampler"): host_session_bridge_module._rehydrate_remote_session_bridge_value( bridge_ref, target_session_handle=target_handle, node_mapping=None, resolution_stats=resolution_stats, ) finally: host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert resolution_stats.bridge_record_lookups == 1 assert resolution_stats.replay_count == 0