diff --git a/remote/host_session_bridge.py b/remote/host_session_bridge.py new file mode 100644 index 0000000..954c56b --- /dev/null +++ b/remote/host_session_bridge.py @@ -0,0 +1,1073 @@ +"""Durable host-side session bridge offload, recovery, and rehydration.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +import hashlib +import importlib +import json +import logging +import os +from pathlib import Path +import threading +import time +from typing import Any, Callable, Iterable, Mapping + +from ..durable_state import ( + DurableObjectRef, + DurableStateError, + FileDurableObjectStore, + read_modal_volume_file, +) +from ..serialization import ( + deserialize_node_inputs, + deserialize_node_outputs, + deserialize_value, + serialize_mapping, + serialize_node_inputs, + serialize_node_outputs, + serialize_value, +) +from ..session_state import ( + InMemoryRemoteSessionBridgeStore, + InMemoryRemoteSessionStore, + RemoteSessionBridgeRecord, + RemoteSessionBridgeRecoveryKind, + RemoteSessionBridgeRef, + RemoteSessionHandle, + RemoteSessionStateError, + RemoteSessionValueRef, + is_remote_session_bridge_ref_payload, + is_remote_session_handle_payload, + is_remote_session_value_ref_payload, + stable_session_bridge_key, +) +from ..settings import get_settings + +logger = logging.getLogger(__name__) + +try: + import modal # type: ignore +except ModuleNotFoundError: # pragma: no cover - local fallback environments. + modal = None + +_REMOTE_SESSION_STORE = InMemoryRemoteSessionStore() +_REMOTE_SESSION_BRIDGE_STORE = InMemoryRemoteSessionBridgeStore() +_DURABLE_OBJECT_STORE_LOCK = threading.Lock() +_DURABLE_OBJECT_STORE: FileDurableObjectStore | None = None +_REMOTE_SESSION_BRIDGE_REPLAY_STATE = threading.local() +_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 +_MODAL_SESSION_BRIDGE_DICTS_LOCK = threading.Lock() +_MODAL_SESSION_BRIDGE_DICTS: dict[tuple[str, str | None], Any] = {} +_MODAL_DURABLE_VOLUMES_LOCK = threading.Lock() +_MODAL_DURABLE_VOLUMES: dict[tuple[str, str | None], Any] = {} +_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", "NOISE", "SAMPLER", "VAE"} +) + + +def _orchestrator_module() -> Any: + """Load the current host orchestrator for transitional execution adapters.""" + return importlib.import_module(".modal_app", __package__) + + +def _execute_node_locally_raw(*args: Any, **kwargs: Any) -> Any: + """Delegate replay of one node to the current local execution owner.""" + return _orchestrator_module()._execute_node_locally_raw(*args, **kwargs) + + +def _execute_subgraph_prompt(*args: Any, **kwargs: Any) -> Any: + """Delegate replay of one subgraph to the current local execution owner.""" + return _orchestrator_module()._execute_subgraph_prompt(*args, **kwargs) + + +def _is_link(value: Any) -> bool: + """Delegate prompt-link recognition to the current local execution owner.""" + return bool(_orchestrator_module()._is_link(value)) + + +def _normalize_prompt_input_value(value: Any) -> Any: + """Delegate prompt input normalization to the local execution owner.""" + return _orchestrator_module()._normalize_prompt_input_value(value) + + +def _resolve_required_subgraph_nodes(*args: Any, **kwargs: Any) -> Any: + """Delegate subgraph closure analysis to the local execution owner.""" + return _orchestrator_module()._resolve_required_subgraph_nodes(*args, **kwargs) + + +def _modal_environment_name() -> str | None: + """Return the active Modal environment name when explicitly configured.""" + environment_name = os.getenv("MODAL_ENVIRONMENT") + if environment_name is None: + return None + normalized = environment_name.strip() + return normalized or None + + +def _lookup_modal_session_bridge_store() -> Any: + """Return the shared Modal Dict that contains durable bridge records.""" + if modal is None or not hasattr(modal, "Dict"): + raise RuntimeError("Modal SDK is required to materialize a remote bridge locally.") + settings = get_settings() + cache_key = (settings.session_bridge_dict_name, _modal_environment_name()) + with _MODAL_SESSION_BRIDGE_DICTS_LOCK: + cached_store = _MODAL_SESSION_BRIDGE_DICTS.get(cache_key) + if cached_store is not None: + return cached_store + bridge_store = modal.Dict.from_name( + settings.session_bridge_dict_name, + environment_name=cache_key[1], + create_if_missing=True, + ) + with _MODAL_SESSION_BRIDGE_DICTS_LOCK: + _MODAL_SESSION_BRIDGE_DICTS[cache_key] = bridge_store + return bridge_store + + +def _lookup_modal_durable_volume() -> Any: + """Return the shared Modal Volume that contains durable bridge objects.""" + if modal is None or not hasattr(modal, "Volume"): + raise RuntimeError( + "Modal SDK is required to download a remote bridge object locally." + ) + settings = get_settings() + cache_key = (settings.volume_name, _modal_environment_name()) + with _MODAL_DURABLE_VOLUMES_LOCK: + cached_volume = _MODAL_DURABLE_VOLUMES.get(cache_key) + if cached_volume is not None: + return cached_volume + volume = modal.Volume.from_name( + settings.volume_name, + environment_name=cache_key[1], + create_if_missing=True, + ) + with _MODAL_DURABLE_VOLUMES_LOCK: + _MODAL_DURABLE_VOLUMES[cache_key] = volume + return volume + + +@dataclass +class _RemoteSessionBridgeResolutionStats: + """Track how one local fallback payload resolved session-backed inputs.""" + + 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 + replay_seconds: float = 0.0 + direct_restore_seconds: float = 0.0 + session_restore_writes: int = 0 +def _payload_remote_session_handle( + payload: dict[str, Any] +) -> RemoteSessionHandle | None: + """Return the decoded prompt-scoped remote session handle for one payload.""" + remote_session = payload.get("remote_session") + if not is_remote_session_handle_payload(remote_session): + return None + return RemoteSessionHandle.from_payload(remote_session) + + +def _sanitize_payload_for_session_bridge_record( + payload: dict[str, Any] +) -> dict[str, Any]: + """Strip run-scoped fields from one producer payload before persisting replay metadata.""" + sanitized_payload = copy.deepcopy(payload) + sanitized_payload.pop("prompt_id", None) + sanitized_payload.pop("remote_session", None) + sanitized_payload.pop("clear_remote_session", None) + sanitized_payload["extra_data"] = {} + return sanitized_payload + + +def _durable_object_store() -> FileDurableObjectStore: + """Return the process-local handle for durable binary objects.""" + global _DURABLE_OBJECT_STORE + + with _DURABLE_OBJECT_STORE_LOCK: + if _DURABLE_OBJECT_STORE is not None: + return _DURABLE_OBJECT_STORE + settings = get_settings() + commit_callback: Callable[[], Any] | None = None + committed_read_callback: Callable[[str], bytes] | None = None + if os.getenv("MODAL_IS_REMOTE") == "1" or os.getenv("MODAL_TASK_ID"): + object_root = Path(settings.remote_storage_root) / "durable_objects" + volume = globals().get("vol") + volume_commit = getattr(volume, "commit", None) + if callable(volume_commit): + commit_callback = volume_commit + if callable(getattr(volume, "read_file", None)): + + def read_committed_object(object_path: str) -> bytes: + """Read one durable object without reloading the mounted volume.""" + volume_path = (Path("durable_objects") / object_path).as_posix() + return read_modal_volume_file(volume, volume_path) + + committed_read_callback = read_committed_object + else: + object_root = settings.local_storage_root / "durable_objects" + _DURABLE_OBJECT_STORE = FileDurableObjectStore( + object_root, + commit_callback=commit_callback, + committed_read_callback=committed_read_callback, + ) + return _DURABLE_OBJECT_STORE + + +def _json_payload_size_bytes(payload: Any) -> int: + """Return the compact UTF-8 size of one JSON-safe payload.""" + return len( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + ) + + +def _select_remote_session_bridge_recovery_kind( + *, + serialized_output: Any | None, + rehydration_plan: Mapping[str, Any] | None, +) -> RemoteSessionBridgeRecoveryKind: + """Choose the least expensive complete recovery mechanism for one bridge.""" + if serialized_output is not None: + return RemoteSessionBridgeRecoveryKind.SERIALIZED_OUTPUT + if isinstance(rehydration_plan, Mapping): + plan_kind = str(rehydration_plan.get("kind") or "") + if plan_kind == "single_node_output": + return RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN + if plan_kind == "subgraph_output": + return RemoteSessionBridgeRecoveryKind.SUBGRAPH_PLAN + return RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY + + +def _remote_session_bridge_recovery_input_names( + *, + recovery_kind: RemoteSessionBridgeRecoveryKind, + rehydration_plan: Mapping[str, Any] | None, + hydrated_inputs: Mapping[str, Any], +) -> tuple[str, ...]: + """Return producer input names that the selected recovery mechanism needs.""" + if recovery_kind in { + RemoteSessionBridgeRecoveryKind.SERIALIZED_OUTPUT, + RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN, + }: + return () + if recovery_kind is RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY: + return tuple(str(input_name) for input_name in hydrated_inputs) + + plan_payload = ( + rehydration_plan.get("payload") + if isinstance(rehydration_plan, Mapping) + else None + ) + boundary_inputs = ( + plan_payload.get("boundary_inputs") + if isinstance(plan_payload, Mapping) + else None + ) + if not isinstance(boundary_inputs, list): + raise RemoteSessionStateError( + "Durable bridge subgraph recovery plans must define boundary_inputs." + ) + input_names = tuple( + dict.fromkeys( + str(boundary_input.get("proxy_input_name") or "").strip() + for boundary_input in boundary_inputs + if isinstance(boundary_input, Mapping) + and str(boundary_input.get("proxy_input_name") or "").strip() + ) + ) + missing_input_names = [ + input_name for input_name in input_names if input_name not in hydrated_inputs + ] + if missing_input_names: + raise RemoteSessionStateError( + "Durable bridge subgraph recovery inputs are missing hydrated values: " + f"{missing_input_names}." + ) + return input_names + + +def _offload_large_bridge_payloads( + *, + hydrated_inputs: Mapping[str, Any], + producer_inputs: dict[str, Any], + output_value: Any, + serialized_output: Any | None, +) -> tuple[ + dict[str, Any], + DurableObjectRef | None, + Any | None, + DurableObjectRef | None, +]: + """Move oversized bridge inputs and outputs into durable binary objects.""" + max_inline_bytes = get_settings().bridge_inline_max_bytes + producer_inputs_object: DurableObjectRef | None = None + serialized_output_object: DurableObjectRef | None = None + if producer_inputs and _json_payload_size_bytes(producer_inputs) > max_inline_bytes: + producer_inputs_object = _durable_object_store().put( + "bridge_objects", + serialize_node_inputs(hydrated_inputs), + ) + producer_inputs = {} + if ( + serialized_output is not None + and _json_payload_size_bytes(serialized_output) > max_inline_bytes + ): + serialized_output_object = _durable_object_store().put( + "bridge_objects", + serialize_node_outputs((output_value,)), + ) + serialized_output = None + return ( + producer_inputs, + producer_inputs_object, + serialized_output, + serialized_output_object, + ) + + +def _deserialize_remote_session_bridge_producer_inputs( + record: RemoteSessionBridgeRecord, +) -> dict[str, Any]: + """Restore bridge producer inputs from inline metadata or durable storage.""" + if record.producer_inputs_object is not None: + return deserialize_node_inputs( + _durable_object_store().get(record.producer_inputs_object) + ) + if record.producer_inputs_retained is False: + recovery_kind = ( + record.recovery_kind.value + if record.recovery_kind is not None + else "unspecified" + ) + raise RemoteSessionStateError( + "Remote session bridge producer inputs were intentionally omitted for " + f"bridge_key={record.bridge_key!r} recovery_kind={recovery_kind!r}." + ) + return deserialize_node_inputs(record.producer_inputs) + + +def _build_remote_session_bridge_record( + *, + payload: dict[str, Any], + 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_input_identity = serialize_mapping(hydrated_inputs) + serialized_output = _serialize_durable_bridge_output(output_value, io_type) + rehydration_plan = _build_durable_bridge_rehydration_plan( + payload=producer_payload, + node_id=node_id, + output_index=output_index, + io_type=io_type, + ) + recovery_kind = _select_remote_session_bridge_recovery_kind( + serialized_output=serialized_output, + rehydration_plan=rehydration_plan, + ) + recovery_input_names = _remote_session_bridge_recovery_input_names( + recovery_kind=recovery_kind, + rehydration_plan=rehydration_plan, + hydrated_inputs=hydrated_inputs, + ) + producer_inputs_retained = recovery_kind in { + RemoteSessionBridgeRecoveryKind.SUBGRAPH_PLAN, + RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY, + } + retained_hydrated_inputs = { + input_name: hydrated_inputs[input_name] for input_name in recovery_input_names + } + producer_inputs = { + input_name: producer_input_identity[input_name] + for input_name in recovery_input_names + } + bridge_key = stable_session_bridge_key( + producer_payload=producer_payload, + producer_inputs=producer_input_identity, + node_id=node_id, + output_index=output_index, + ) + ( + producer_inputs, + producer_inputs_object, + serialized_output, + serialized_output_object, + ) = _offload_large_bridge_payloads( + hydrated_inputs=retained_hydrated_inputs, + producer_inputs=producer_inputs, + output_value=output_value, + serialized_output=serialized_output, + ) + logger.info( + "Prepared remote session bridge bridge_key=%s recovery_kind=%s " + "producer_inputs_retained=%s retained_input_count=%d omitted_input_count=%d.", + bridge_key, + recovery_kind.value, + producer_inputs_retained, + len(recovery_input_names), + len(producer_input_identity) - len(recovery_input_names), + ) + return RemoteSessionBridgeRecord( + bridge_key=bridge_key, + node_id=node_id, + output_index=output_index, + producer_payload=producer_payload, + producer_inputs=producer_inputs, + recovery_kind=recovery_kind, + producer_inputs_retained=producer_inputs_retained, + producer_inputs_object=producer_inputs_object, + serialized_output=serialized_output, + serialized_output_object=serialized_output_object, + serialized_output_io_type=( + str(io_type) + if serialized_output is not None or serialized_output_object is not None + else None + ), + rehydration_plan=rehydration_plan, + rehydration_plan_io_type=( + str(io_type) if rehydration_plan is not None else None + ), + ) + + +def _record_remote_session_resolution_event( + resolution_stats: "_RemoteSessionBridgeResolutionStats | None", + event_name: str, + event_payload: Mapping[str, Any], +) -> None: + """Accumulate one local fallback remote-session resolution event.""" + del event_payload + if resolution_stats is None: + return + if event_name in {"session-value-hit", "bridge-target-hit", "bridge-source-hit"}: + resolution_stats.live_session_hits += 1 + + +def _clone_cached_bridge_value(value: Any) -> Any: + """Clone one cached bridge value when the runtime object exposes a safe clone path.""" + clone_method = getattr(value, "clone", None) + if callable(clone_method): + return clone_method() + return value + + +def _store_remote_session_bridge_value( + bridge_key: str, + value: Any, +) -> None: + """Retain one live bridge value in-process so later mapped phases can skip replay.""" + with _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LOCK: + _REMOTE_SESSION_BRIDGE_VALUE_CACHE[bridge_key] = _clone_cached_bridge_value( + value + ) + if bridge_key in _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER: + _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.remove(bridge_key) + _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.append(bridge_key) + while ( + len(_REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER) + > _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LIMIT + ): + evicted_key = _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.pop(0) + _REMOTE_SESSION_BRIDGE_VALUE_CACHE.pop(evicted_key, None) + + +def _get_remote_session_bridge_value(bridge_key: str) -> Any | None: + """Return one retained bridge value when the current process still has it.""" + with _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LOCK: + cached_value = _REMOTE_SESSION_BRIDGE_VALUE_CACHE.get(bridge_key) + if cached_value is None: + return None + if bridge_key in _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER: + _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.remove(bridge_key) + _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.append(bridge_key) + 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 and record.serialized_output_object is None: + return None + + restore_started_at = time.perf_counter() + if record.serialized_output_object is not None: + restored_outputs = deserialize_node_outputs( + _durable_object_store().get(record.serialized_output_object) + ) + if len(restored_outputs) != 1: + raise DurableStateError( + f"Bridge object {record.serialized_output_object.object_path!r} " + "must contain exactly one output." + ) + restored_value = restored_outputs[0] + storage_kind = "object-backed" + else: + restored_value = deserialize_value(record.serialized_output) + storage_kind = "inline" + _REMOTE_SESSION_STORE.put_bridge_output( + target_session_handle, + bridge_key=record.bridge_key, + 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 %s serialized %s payload into session_id=%s.", + record.bridge_key, + storage_kind, + record.serialized_output_io_type or "bridge", + target_session_handle.session_id, + ) + return restored_value + + +def materialize_remote_session_bridge_ref_locally( + ref_payload: Mapping[str, Any], +) -> Any: + """Download and deserialize one durable bridge value for local ComfyUI work.""" + if not is_remote_session_bridge_ref_payload(ref_payload): + raise TypeError( + "Local Modal bridge materialization requires a remote session bridge ref." + ) + ref = RemoteSessionBridgeRef.from_payload(ref_payload) + materialization_started_at = time.perf_counter() + logger.info( + "Starting local remote-bridge materialization bridge_key=%s node=%s output=%d.", + ref.bridge_key, + ref.node_id, + ref.output_index, + ) + record_lookup_started_at = time.perf_counter() + record = _load_local_materialization_bridge_record(ref.bridge_key) + logger.info( + "Loaded local materialization bridge record in %.3fs bridge_key=%s object_backed=%s.", + time.perf_counter() - record_lookup_started_at, + ref.bridge_key, + record.serialized_output_object is not None, + ) + if record.serialized_output_object is not None: + object_download_started_at = time.perf_counter() + restored_outputs = deserialize_node_outputs( + _read_local_materialization_bridge_object( + record.serialized_output_object + ) + ) + logger.info( + "Downloaded and deserialized local bridge object in %.3fs bridge_key=%s bytes=%d.", + time.perf_counter() - object_download_started_at, + ref.bridge_key, + record.serialized_output_object.size_bytes, + ) + if len(restored_outputs) != 1: + raise DurableStateError( + f"Bridge object {record.serialized_output_object.object_path!r} " + "must contain exactly one output." + ) + restored_value = restored_outputs[0] + storage_kind = "object-backed" + elif record.serialized_output is not None: + restored_value = deserialize_value(record.serialized_output) + storage_kind = "inline" + else: + raise RemoteSessionStateError( + "Remote bridge cannot be materialized locally without a durable " + f"serialized output bridge_key={ref.bridge_key!r}." + ) + logger.info( + "Materialized remote bridge locally in %.3fs bridge_key=%s node=%s output=%d storage=%s io_type=%s.", + time.perf_counter() - materialization_started_at, + ref.bridge_key, + ref.node_id, + ref.output_index, + storage_kind, + record.serialized_output_io_type or "bridge", + ) + return restored_value + + +def _load_local_materialization_bridge_record( + bridge_key: str, +) -> RemoteSessionBridgeRecord: + """Load a bridge record from local fallback state or Modal's shared Dict.""" + settings = get_settings() + if settings.execution_mode != "remote": + return _REMOTE_SESSION_BRIDGE_STORE.get_record(bridge_key) + bridge_store = _lookup_modal_session_bridge_store() + payload = bridge_store.get(bridge_key) + if not isinstance(payload, Mapping): + raise RemoteSessionStateError( + f"Remote session bridge record {bridge_key!r} was not found." + ) + return RemoteSessionBridgeRecord.from_payload(payload) + + +def _read_local_materialization_bridge_object( + object_ref: DurableObjectRef, +) -> bytes: + """Read and validate one bridge object locally from fallback or Modal storage.""" + settings = get_settings() + if settings.execution_mode != "remote": + return _durable_object_store().get(object_ref) + volume_path = (Path("durable_objects") / object_ref.object_path).as_posix() + payload = read_modal_volume_file( + _lookup_modal_durable_volume(), + volume_path, + ) + if len(payload) != object_ref.size_bytes: + raise DurableStateError( + f"Durable bridge object {object_ref.object_path!r} has an unexpected size." + ) + if hashlib.sha256(payload).hexdigest() != object_ref.sha256: + raise DurableStateError( + f"Durable bridge object {object_ref.object_path!r} failed its content-address check." + ) + return payload + + +def _build_durable_bridge_rehydration_plan( + *, + payload: dict[str, Any], + node_id: str, + output_index: int, + io_type: str, +) -> dict[str, Any] | None: + """Return a direct rehydration plan when one bridge output can be rebuilt without replay.""" + if str(io_type or "") not in _DURABLE_BRIDGE_REHYDRATION_IO_TYPES: + return None + prompt = payload.get("subgraph_prompt") + if not isinstance(prompt, dict): + return None + prompt_node = prompt.get(str(node_id)) + if not isinstance(prompt_node, dict): + return None + class_type = prompt_node.get("class_type") + inputs = prompt_node.get("inputs") + if ( + not isinstance(class_type, str) + or not class_type.strip() + or not isinstance(inputs, dict) + ): + return None + + normalized_inputs: dict[str, Any] = {} + has_linked_input = False + for input_name, input_value in inputs.items(): + normalized_value = _normalize_prompt_input_value(copy.deepcopy(input_value)) + if _is_link(normalized_value): + has_linked_input = True + continue + normalized_inputs[str(input_name)] = normalized_value + + if not has_linked_input: + node_data: dict[str, Any] = {"class_type": class_type} + custom_nodes_bundle = payload.get("custom_nodes_bundle") + if isinstance(custom_nodes_bundle, str) and custom_nodes_bundle.strip(): + node_data["custom_nodes_bundle"] = custom_nodes_bundle + return { + "kind": "single_node_output", + "node_data": node_data, + "node_inputs": normalized_inputs, + } + + required_node_ids = set( + _resolve_required_subgraph_nodes( + prompt=prompt, + execute_node_ids=[str(node_id)], + ) + ) + if str(node_id) not in required_node_ids: + return None + if _subgraph_contains_sampling_node(prompt, required_node_ids): + logger.info( + "Skipping durable %s bridge subgraph rehydration plan for node_id=%s class_type=%s because its dependency closure includes a sampler.", + io_type, + node_id, + class_type, + ) + return None + + rehydration_payload = copy.deepcopy(payload) + rehydration_payload[ + "component_id" + ] = f"{payload.get('component_id', 'component')}::rehydrate:{node_id}" + rehydration_payload["subgraph_prompt"] = { + str(current_node_id): copy.deepcopy(prompt[current_node_id]) + for current_node_id in prompt + if str(current_node_id) in required_node_ids + } + rehydration_payload["component_node_ids"] = sorted(required_node_ids) + rehydration_payload["execute_node_ids"] = [str(node_id)] + rehydration_payload["mapped_execute_node_ids"] = [] + rehydration_payload["static_execute_node_ids"] = [] + rehydration_payload["boundary_inputs"] = [ + { + **copy.deepcopy(boundary_input), + "targets": [ + copy.deepcopy(target) + for target in boundary_input.get("targets", []) + if str(target.get("node_id")) in required_node_ids + ], + } + for boundary_input in payload.get("boundary_inputs", []) + if any( + str(target.get("node_id")) in required_node_ids + for target in boundary_input.get("targets", []) + ) + ] + rehydration_payload["boundary_outputs"] = [ + { + "node_id": str(node_id), + "output_index": int(output_index), + "io_type": str(io_type), + "is_list": False, + "session_output": False, + "proxy_output_name": f"{node_id}_rehydrated_{output_index}", + } + ] + return { + "kind": "subgraph_output", + "payload": rehydration_payload, + } + + +def _restore_planned_remote_session_bridge_value( + record: RemoteSessionBridgeRecord, + *, + target_session_handle: RemoteSessionHandle, + node_mapping: dict[str, type[Any]] | None = None, + resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, +) -> Any | None: + """Restore one bridge value directly from a durable node rehydration plan.""" + if not isinstance(record.rehydration_plan, Mapping): + return None + + restore_started_at = time.perf_counter() + plan_kind = str(record.rehydration_plan.get("kind") or "") + if plan_kind == "single_node_output": + node_data = record.rehydration_plan.get("node_data") + node_inputs = record.rehydration_plan.get("node_inputs") + if not isinstance(node_data, Mapping) or not isinstance(node_inputs, Mapping): + return None + outputs = _execute_node_locally_raw( + dict(node_data), + dict(node_inputs), + node_mapping=node_mapping, + ) + output_index = record.output_index + elif plan_kind == "subgraph_output": + plan_payload = record.rehydration_plan.get("payload") + if not isinstance(plan_payload, Mapping): + return None + executable_payload = dict(plan_payload) + executable_payload["remote_session"] = target_session_handle.to_payload() + executable_payload.pop("clear_remote_session", None) + outputs = _execute_subgraph_prompt( + executable_payload, + _deserialize_remote_session_bridge_producer_inputs(record), + node_mapping, + ) + output_index = 0 + else: + return None + if output_index < 0 or output_index >= len(outputs): + raise RemoteSessionStateError( + f"Durable bridge rehydration plan for {record.bridge_key!r} did not produce output index {record.output_index}." + ) + restored_value = outputs[output_index] + _REMOTE_SESSION_STORE.put_bridge_output( + target_session_handle, + bridge_key=record.bridge_key, + 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 %s rehydration plan into session_id=%s.", + record.bridge_key, + record.rehydration_plan_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) + if replay_stack is None: + replay_stack = set() + _REMOTE_SESSION_BRIDGE_REPLAY_STATE.bridge_keys = replay_stack + return replay_stack + + +def _bridge_record_replays_sampling_node(record: RemoteSessionBridgeRecord) -> bool: + """Return whether replaying one bridge record would rerun sampler work.""" + execute_node_ids = [ + str(node_id) + for node_id in record.producer_payload.get("execute_node_ids", []) + if str(node_id) + ] + if not execute_node_ids: + return False + subgraph_prompt = record.producer_payload.get("subgraph_prompt") + if not isinstance(subgraph_prompt, dict): + return False + required_node_ids = _resolve_required_subgraph_nodes( + prompt=subgraph_prompt, + execute_node_ids=execute_node_ids, + ) + return _subgraph_contains_sampling_node(subgraph_prompt, required_node_ids) + + +def _subgraph_contains_sampling_node( + subgraph_prompt: Mapping[str, Any], + node_ids: Iterable[str], +) -> bool: + """Return whether any selected node looks like sampler work.""" + for node_id in node_ids: + prompt_node = subgraph_prompt.get(node_id) + if not isinstance(prompt_node, Mapping): + continue + class_type = str(prompt_node.get("class_type") or "") + if "sampler" in class_type.lower(): + return True + return False + + +def _rehydrate_remote_session_bridge_value( + ref: RemoteSessionBridgeRef, + *, + target_session_handle: RemoteSessionHandle | None, + node_mapping: dict[str, type[Any]] | None = None, + resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, +) -> Any: + """Replay one producer phase into the current session when the live value is gone.""" + if target_session_handle is None: + raise RemoteSessionStateError( + "Remote session bridge replay requires a target remote_session handle." + ) + + cached_value = _get_remote_session_bridge_value(ref.bridge_key) + if cached_value is not None: + restore_started_at = time.perf_counter() + _REMOTE_SESSION_STORE.put_bridge_output( + target_session_handle, + bridge_key=ref.bridge_key, + node_id=ref.node_id, + output_index=ref.output_index, + value=cached_value, + ) + if resolution_stats is not None: + resolution_stats.bridge_cache_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 directly from warm cache into session_id=%s.", + ref.bridge_key, + target_session_handle.session_id, + ) + return cached_value + + replay_stack = _remote_session_bridge_replay_stack() + if ref.bridge_key in replay_stack: + raise RemoteSessionStateError( + f"Detected recursive remote session bridge replay for {ref.bridge_key!r}." + ) + + record_lookup_started_at = time.perf_counter() + record = _REMOTE_SESSION_BRIDGE_STORE.get_record(ref.bridge_key) + if resolution_stats is not None: + resolution_stats.bridge_record_lookups += 1 + 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 + restored_value = _restore_planned_remote_session_bridge_value( + record, + target_session_handle=target_session_handle, + node_mapping=node_mapping, + resolution_stats=resolution_stats, + ) + if restored_value is not None: + return restored_value + if _bridge_record_replays_sampling_node(record): + message = ( + "Remote session bridge replay would rerun a sampler component " + f"for bridge_key={ref.bridge_key!r}; expected live or warm bridge reuse instead." + ) + logger.error(message) + raise RemoteSessionStateError(message) + replay_payload = copy.deepcopy(record.producer_payload) + replay_payload["remote_session"] = target_session_handle.to_payload() + replay_payload.pop("clear_remote_session", None) + if target_session_handle.prompt_id is not None: + replay_payload["prompt_id"] = target_session_handle.prompt_id + replay_inputs = _deserialize_remote_session_bridge_producer_inputs(record) + + logger.info( + "Replaying remote session bridge bridge_key=%s into session_id=%s via component=%s.", + ref.bridge_key, + target_session_handle.session_id, + replay_payload.get("component_id"), + ) + replay_stack.add(ref.bridge_key) + replay_started_at = time.perf_counter() + try: + _execute_subgraph_prompt(replay_payload, replay_inputs, node_mapping) + finally: + replay_stack.remove(ref.bridge_key) + if resolution_stats is not None: + resolution_stats.replay_count += 1 + resolution_stats.replay_seconds += time.perf_counter() - replay_started_at + + return _REMOTE_SESSION_STORE.get_output( + RemoteSessionValueRef( + session_id=target_session_handle.session_id, + node_id=ref.node_id, + output_index=ref.output_index, + ) + ) + + +def _resolve_remote_session_inputs( + hydrated_inputs: dict[str, Any], + *, + component_id: str | None = None, + target_session_handle: RemoteSessionHandle | None = None, + node_mapping: dict[str, type[Any]] | None = None, + resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, +) -> dict[str, Any]: + """Resolve any remote-session value refs embedded in boundary inputs.""" + + def count_refs(value: Any) -> int: + """Return the number of live or durable session refs nested in one value.""" + if is_remote_session_value_ref_payload( + value + ) or is_remote_session_bridge_ref_payload(value): + return 1 + if isinstance(value, list | tuple): + return sum(count_refs(item) for item in value) + if isinstance(value, Mapping): + return sum(count_refs(item) for item in value.values()) + return 0 + + ref_counts_by_input = { + input_name: count_refs(input_value) + for input_name, input_value in hydrated_inputs.items() + } + ref_input_names = [ + input_name + for input_name, ref_count in ref_counts_by_input.items() + if ref_count > 0 + ] + total_ref_count = sum(ref_counts_by_input.values()) + if ref_input_names: + logger.info( + "Resolving %d remote session input refs for component=%s inputs=%s.", + total_ref_count, + component_id or "", + sorted(ref_input_names), + ) + if resolution_stats is not None: + resolution_stats.input_ref_count += total_ref_count + return { + input_name: _REMOTE_SESSION_STORE.resolve_value_with_bridges( + input_value, + target_session_handle=target_session_handle, + resolution_callback=( + lambda event_name, event_payload: _record_remote_session_resolution_event( + resolution_stats, + event_name, + event_payload, + ) + ) + if resolution_stats is not None + else None, + bridge_resolver=lambda ref: _rehydrate_remote_session_bridge_value( + ref, + target_session_handle=target_session_handle, + node_mapping=node_mapping, + resolution_stats=resolution_stats, + ), + ) + for input_name, input_value in hydrated_inputs.items() + } + + +def _log_remote_session_resolution_summary( + *, + component_id: str, + resolution_stats: "_RemoteSessionBridgeResolutionStats", +) -> None: + """Emit one high-signal summary for local fallback bridge resolution work.""" + 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 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, + resolution_stats.replay_seconds, + resolution_stats.direct_restore_seconds, + resolution_stats.session_restore_writes, + ) diff --git a/remote/modal_app.py b/remote/modal_app.py index be0f82f..11b3d7a 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -126,6 +126,18 @@ from .modal_container_logs import ( list_active_modal_containers, stop_managed_modal_container, ) +from .host_session_bridge import ( + _RemoteSessionBridgeResolutionStats, + _build_remote_session_bridge_record, + _log_remote_session_resolution_summary, + _payload_remote_session_handle, + _rehydrate_remote_session_bridge_value, + _resolve_remote_session_inputs, + _restore_serialized_remote_session_bridge_value, + _store_remote_session_bridge_value, + materialize_remote_session_bridge_ref_locally, +) +from . import host_session_bridge as _host_session_bridge logger = logging.getLogger(__name__) @@ -145,11 +157,6 @@ _STAGED_LLM_PROFILE_RESULTS: dict[tuple[str, str], dict[str, Any]] = {} _MODAL_APP_STOP_TIMEOUT_SECONDS = 120.0 _MODAL_INTERRUPT_DICTS_LOCK = threading.Lock() _MODAL_INTERRUPT_DICTS: dict[tuple[str, str | None], Any] = {} -_MODAL_SESSION_BRIDGE_DICTS_LOCK = threading.Lock() -_MODAL_SESSION_BRIDGE_DICTS: dict[tuple[str, str | None], Any] = {} -_MODAL_DURABLE_VOLUMES_LOCK = threading.Lock() -_MODAL_DURABLE_VOLUMES: dict[tuple[str, str | None], Any] = {} - # Mapped execution, warmup, and keepalive state. _MAPPED_PROGRESS_NODE_IDS_LOCK = threading.Lock() _MAPPED_PROGRESS_NODE_IDS: dict[tuple[str, str, str], str] = {} @@ -180,33 +187,6 @@ _ACTIVE_REMOTE_INVOCATIONS_LOCK = threading.Lock() _ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT: dict[ str, dict[str, "_ActiveRemoteInvocation"] ] = {} -_REMOTE_SESSION_STORE = InMemoryRemoteSessionStore() -_REMOTE_SESSION_BRIDGE_STORE = InMemoryRemoteSessionBridgeStore() -_DURABLE_OBJECT_STORE_LOCK = threading.Lock() -_DURABLE_OBJECT_STORE: FileDurableObjectStore | None = None -_REMOTE_SESSION_BRIDGE_REPLAY_STATE = threading.local() -_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( - { - "AUDIO", - "BOOLEAN", - "CONDITIONING", - "FLOAT", - "IMAGE", - "INT", - "LATENT", - "MASK", - "SIGMAS", - "STRING", - "VIDEO", - } -) -_DURABLE_BRIDGE_REHYDRATION_IO_TYPES = frozenset( - {"CLIP", "MODEL", "NOISE", "SAMPLER", "VAE"} -) def _remote_modal_call_worker_count() -> int: """Return the number of local worker threads reserved for blocking Modal calls.""" @@ -308,20 +288,6 @@ class _ActiveRemoteInvocation: interrupt_remote_call: Callable[[], Any] | None -@dataclass -class _RemoteSessionBridgeResolutionStats: - """Track how one local fallback payload resolved session-backed inputs.""" - - 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 - replay_seconds: float = 0.0 - direct_restore_seconds: float = 0.0 - session_restore_writes: int = 0 try: @@ -416,909 +382,56 @@ def _close_remote_payload_stream(stream_events: Iterator[dict[str, Any]]) -> Non logger.debug("Ignoring async remote payload stream close failure: %s", exc) -def _payload_remote_session_handle( - payload: dict[str, Any] -) -> RemoteSessionHandle | None: - """Return the decoded prompt-scoped remote session handle for one payload.""" - remote_session = payload.get("remote_session") - if not is_remote_session_handle_payload(remote_session): - return None - return RemoteSessionHandle.from_payload(remote_session) - - -def _sanitize_payload_for_session_bridge_record( - payload: dict[str, Any] -) -> dict[str, Any]: - """Strip run-scoped fields from one producer payload before persisting replay metadata.""" - sanitized_payload = copy.deepcopy(payload) - sanitized_payload.pop("prompt_id", None) - sanitized_payload.pop("remote_session", None) - sanitized_payload.pop("clear_remote_session", None) - sanitized_payload["extra_data"] = {} - return sanitized_payload - - -def _durable_object_store() -> FileDurableObjectStore: - """Return the process-local handle for durable binary objects.""" - global _DURABLE_OBJECT_STORE - - with _DURABLE_OBJECT_STORE_LOCK: - if _DURABLE_OBJECT_STORE is not None: - return _DURABLE_OBJECT_STORE - settings = get_settings() - commit_callback: Callable[[], Any] | None = None - committed_read_callback: Callable[[str], bytes] | None = None - if os.getenv("MODAL_IS_REMOTE") == "1" or os.getenv("MODAL_TASK_ID"): - object_root = Path(settings.remote_storage_root) / "durable_objects" - volume = globals().get("vol") - volume_commit = getattr(volume, "commit", None) - if callable(volume_commit): - commit_callback = volume_commit - if callable(getattr(volume, "read_file", None)): - - def read_committed_object(object_path: str) -> bytes: - """Read one durable object without reloading the mounted volume.""" - volume_path = (Path("durable_objects") / object_path).as_posix() - return read_modal_volume_file(volume, volume_path) - - committed_read_callback = read_committed_object - else: - object_root = settings.local_storage_root / "durable_objects" - _DURABLE_OBJECT_STORE = FileDurableObjectStore( - object_root, - commit_callback=commit_callback, - committed_read_callback=committed_read_callback, - ) - return _DURABLE_OBJECT_STORE - - -def _json_payload_size_bytes(payload: Any) -> int: - """Return the compact UTF-8 size of one JSON-safe payload.""" - return len( - json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") - ) - - -def _select_remote_session_bridge_recovery_kind( - *, - serialized_output: Any | None, - rehydration_plan: Mapping[str, Any] | None, -) -> RemoteSessionBridgeRecoveryKind: - """Choose the least expensive complete recovery mechanism for one bridge.""" - if serialized_output is not None: - return RemoteSessionBridgeRecoveryKind.SERIALIZED_OUTPUT - if isinstance(rehydration_plan, Mapping): - plan_kind = str(rehydration_plan.get("kind") or "") - if plan_kind == "single_node_output": - return RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN - if plan_kind == "subgraph_output": - return RemoteSessionBridgeRecoveryKind.SUBGRAPH_PLAN - return RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY - - -def _remote_session_bridge_recovery_input_names( - *, - recovery_kind: RemoteSessionBridgeRecoveryKind, - rehydration_plan: Mapping[str, Any] | None, - hydrated_inputs: Mapping[str, Any], -) -> tuple[str, ...]: - """Return producer input names that the selected recovery mechanism needs.""" - if recovery_kind in { - RemoteSessionBridgeRecoveryKind.SERIALIZED_OUTPUT, - RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN, - }: - return () - if recovery_kind is RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY: - return tuple(str(input_name) for input_name in hydrated_inputs) - - plan_payload = ( - rehydration_plan.get("payload") - if isinstance(rehydration_plan, Mapping) - else None - ) - boundary_inputs = ( - plan_payload.get("boundary_inputs") - if isinstance(plan_payload, Mapping) - else None - ) - if not isinstance(boundary_inputs, list): - raise RemoteSessionStateError( - "Durable bridge subgraph recovery plans must define boundary_inputs." - ) - input_names = tuple( - dict.fromkeys( - str(boundary_input.get("proxy_input_name") or "").strip() - for boundary_input in boundary_inputs - if isinstance(boundary_input, Mapping) - and str(boundary_input.get("proxy_input_name") or "").strip() - ) - ) - missing_input_names = [ - input_name for input_name in input_names if input_name not in hydrated_inputs - ] - if missing_input_names: - raise RemoteSessionStateError( - "Durable bridge subgraph recovery inputs are missing hydrated values: " - f"{missing_input_names}." - ) - return input_names - - -def _offload_large_bridge_payloads( - *, - hydrated_inputs: Mapping[str, Any], - producer_inputs: dict[str, Any], - output_value: Any, - serialized_output: Any | None, -) -> tuple[ - dict[str, Any], - DurableObjectRef | None, - Any | None, - DurableObjectRef | None, -]: - """Move oversized bridge inputs and outputs into durable binary objects.""" - max_inline_bytes = get_settings().bridge_inline_max_bytes - producer_inputs_object: DurableObjectRef | None = None - serialized_output_object: DurableObjectRef | None = None - if producer_inputs and _json_payload_size_bytes(producer_inputs) > max_inline_bytes: - producer_inputs_object = _durable_object_store().put( - "bridge_objects", - serialize_node_inputs(hydrated_inputs), - ) - producer_inputs = {} - if ( - serialized_output is not None - and _json_payload_size_bytes(serialized_output) > max_inline_bytes - ): - serialized_output_object = _durable_object_store().put( - "bridge_objects", - serialize_node_outputs((output_value,)), - ) - serialized_output = None - return ( - producer_inputs, - producer_inputs_object, - serialized_output, - serialized_output_object, - ) - - -def _deserialize_remote_session_bridge_producer_inputs( - record: RemoteSessionBridgeRecord, -) -> dict[str, Any]: - """Restore bridge producer inputs from inline metadata or durable storage.""" - if record.producer_inputs_object is not None: - return deserialize_node_inputs( - _durable_object_store().get(record.producer_inputs_object) - ) - if record.producer_inputs_retained is False: - recovery_kind = ( - record.recovery_kind.value - if record.recovery_kind is not None - else "unspecified" - ) - raise RemoteSessionStateError( - "Remote session bridge producer inputs were intentionally omitted for " - f"bridge_key={record.bridge_key!r} recovery_kind={recovery_kind!r}." - ) - return deserialize_node_inputs(record.producer_inputs) - - -def _build_remote_session_bridge_record( - *, - payload: dict[str, Any], - 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_input_identity = serialize_mapping(hydrated_inputs) - serialized_output = _serialize_durable_bridge_output(output_value, io_type) - rehydration_plan = _build_durable_bridge_rehydration_plan( - payload=producer_payload, - node_id=node_id, - output_index=output_index, - io_type=io_type, - ) - recovery_kind = _select_remote_session_bridge_recovery_kind( - serialized_output=serialized_output, - rehydration_plan=rehydration_plan, - ) - recovery_input_names = _remote_session_bridge_recovery_input_names( - recovery_kind=recovery_kind, - rehydration_plan=rehydration_plan, - hydrated_inputs=hydrated_inputs, - ) - producer_inputs_retained = recovery_kind in { - RemoteSessionBridgeRecoveryKind.SUBGRAPH_PLAN, - RemoteSessionBridgeRecoveryKind.PRODUCER_REPLAY, - } - retained_hydrated_inputs = { - input_name: hydrated_inputs[input_name] for input_name in recovery_input_names - } - producer_inputs = { - input_name: producer_input_identity[input_name] - for input_name in recovery_input_names - } - bridge_key = stable_session_bridge_key( - producer_payload=producer_payload, - producer_inputs=producer_input_identity, - node_id=node_id, - output_index=output_index, - ) - ( - producer_inputs, - producer_inputs_object, - serialized_output, - serialized_output_object, - ) = _offload_large_bridge_payloads( - hydrated_inputs=retained_hydrated_inputs, - producer_inputs=producer_inputs, - output_value=output_value, - serialized_output=serialized_output, - ) - logger.info( - "Prepared remote session bridge bridge_key=%s recovery_kind=%s " - "producer_inputs_retained=%s retained_input_count=%d omitted_input_count=%d.", - bridge_key, - recovery_kind.value, - producer_inputs_retained, - len(recovery_input_names), - len(producer_input_identity) - len(recovery_input_names), - ) - return RemoteSessionBridgeRecord( - bridge_key=bridge_key, - node_id=node_id, - output_index=output_index, - producer_payload=producer_payload, - producer_inputs=producer_inputs, - recovery_kind=recovery_kind, - producer_inputs_retained=producer_inputs_retained, - producer_inputs_object=producer_inputs_object, - serialized_output=serialized_output, - serialized_output_object=serialized_output_object, - serialized_output_io_type=( - str(io_type) - if serialized_output is not None or serialized_output_object is not None - else None - ), - rehydration_plan=rehydration_plan, - rehydration_plan_io_type=( - str(io_type) if rehydration_plan is not None else None - ), - ) - - -def _record_remote_session_resolution_event( - resolution_stats: "_RemoteSessionBridgeResolutionStats | None", - event_name: str, - event_payload: Mapping[str, Any], -) -> None: - """Accumulate one local fallback remote-session resolution event.""" - del event_payload - if resolution_stats is None: - return - if event_name in {"session-value-hit", "bridge-target-hit", "bridge-source-hit"}: - resolution_stats.live_session_hits += 1 - - -def _clone_cached_bridge_value(value: Any) -> Any: - """Clone one cached bridge value when the runtime object exposes a safe clone path.""" - clone_method = getattr(value, "clone", None) - if callable(clone_method): - return clone_method() - return value - - -def _store_remote_session_bridge_value( - bridge_key: str, - value: Any, -) -> None: - """Retain one live bridge value in-process so later mapped phases can skip replay.""" - with _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LOCK: - _REMOTE_SESSION_BRIDGE_VALUE_CACHE[bridge_key] = _clone_cached_bridge_value( - value - ) - if bridge_key in _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER: - _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.remove(bridge_key) - _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.append(bridge_key) - while ( - len(_REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER) - > _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LIMIT - ): - evicted_key = _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.pop(0) - _REMOTE_SESSION_BRIDGE_VALUE_CACHE.pop(evicted_key, None) - - -def _get_remote_session_bridge_value(bridge_key: str) -> Any | None: - """Return one retained bridge value when the current process still has it.""" - with _REMOTE_SESSION_BRIDGE_VALUE_CACHE_LOCK: - cached_value = _REMOTE_SESSION_BRIDGE_VALUE_CACHE.get(bridge_key) - if cached_value is None: - return None - if bridge_key in _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER: - _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.remove(bridge_key) - _REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.append(bridge_key) - 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 and record.serialized_output_object is None: - return None - - restore_started_at = time.perf_counter() - if record.serialized_output_object is not None: - restored_outputs = deserialize_node_outputs( - _durable_object_store().get(record.serialized_output_object) - ) - if len(restored_outputs) != 1: - raise DurableStateError( - f"Bridge object {record.serialized_output_object.object_path!r} " - "must contain exactly one output." - ) - restored_value = restored_outputs[0] - storage_kind = "object-backed" - else: - restored_value = deserialize_value(record.serialized_output) - storage_kind = "inline" - _REMOTE_SESSION_STORE.put_bridge_output( - target_session_handle, - bridge_key=record.bridge_key, - 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 %s serialized %s payload into session_id=%s.", - record.bridge_key, - storage_kind, - record.serialized_output_io_type or "bridge", - target_session_handle.session_id, - ) - return restored_value - - -def materialize_remote_session_bridge_ref_locally( - ref_payload: Mapping[str, Any], -) -> Any: - """Download and deserialize one durable bridge value for local ComfyUI work.""" - if not is_remote_session_bridge_ref_payload(ref_payload): - raise TypeError( - "Local Modal bridge materialization requires a remote session bridge ref." - ) - ref = RemoteSessionBridgeRef.from_payload(ref_payload) - materialization_started_at = time.perf_counter() - logger.info( - "Starting local remote-bridge materialization bridge_key=%s node=%s output=%d.", - ref.bridge_key, - ref.node_id, - ref.output_index, - ) - record_lookup_started_at = time.perf_counter() - record = _load_local_materialization_bridge_record(ref.bridge_key) - logger.info( - "Loaded local materialization bridge record in %.3fs bridge_key=%s object_backed=%s.", - time.perf_counter() - record_lookup_started_at, - ref.bridge_key, - record.serialized_output_object is not None, - ) - if record.serialized_output_object is not None: - object_download_started_at = time.perf_counter() - restored_outputs = deserialize_node_outputs( - _read_local_materialization_bridge_object( - record.serialized_output_object - ) - ) - logger.info( - "Downloaded and deserialized local bridge object in %.3fs bridge_key=%s bytes=%d.", - time.perf_counter() - object_download_started_at, - ref.bridge_key, - record.serialized_output_object.size_bytes, - ) - if len(restored_outputs) != 1: - raise DurableStateError( - f"Bridge object {record.serialized_output_object.object_path!r} " - "must contain exactly one output." - ) - restored_value = restored_outputs[0] - storage_kind = "object-backed" - elif record.serialized_output is not None: - restored_value = deserialize_value(record.serialized_output) - storage_kind = "inline" - else: - raise RemoteSessionStateError( - "Remote bridge cannot be materialized locally without a durable " - f"serialized output bridge_key={ref.bridge_key!r}." - ) - logger.info( - "Materialized remote bridge locally in %.3fs bridge_key=%s node=%s output=%d storage=%s io_type=%s.", - time.perf_counter() - materialization_started_at, - ref.bridge_key, - ref.node_id, - ref.output_index, - storage_kind, - record.serialized_output_io_type or "bridge", - ) - return restored_value - - -def _load_local_materialization_bridge_record( - bridge_key: str, -) -> RemoteSessionBridgeRecord: - """Load a bridge record from local fallback state or Modal's shared Dict.""" - settings = get_settings() - if settings.execution_mode != "remote": - return _REMOTE_SESSION_BRIDGE_STORE.get_record(bridge_key) - bridge_store = _lookup_modal_session_bridge_store() - payload = bridge_store.get(bridge_key) - if not isinstance(payload, Mapping): - raise RemoteSessionStateError( - f"Remote session bridge record {bridge_key!r} was not found." - ) - return RemoteSessionBridgeRecord.from_payload(payload) - - -def _read_local_materialization_bridge_object( - object_ref: DurableObjectRef, -) -> bytes: - """Read and validate one bridge object locally from fallback or Modal storage.""" - settings = get_settings() - if settings.execution_mode != "remote": - return _durable_object_store().get(object_ref) - volume_path = (Path("durable_objects") / object_ref.object_path).as_posix() - payload = read_modal_volume_file( - _lookup_modal_durable_volume(), - volume_path, - ) - if len(payload) != object_ref.size_bytes: - raise DurableStateError( - f"Durable bridge object {object_ref.object_path!r} has an unexpected size." - ) - if hashlib.sha256(payload).hexdigest() != object_ref.sha256: - raise DurableStateError( - f"Durable bridge object {object_ref.object_path!r} failed its content-address check." - ) - return payload - - -def _build_durable_bridge_rehydration_plan( - *, - payload: dict[str, Any], - node_id: str, - output_index: int, - io_type: str, -) -> dict[str, Any] | None: - """Return a direct rehydration plan when one bridge output can be rebuilt without replay.""" - if str(io_type or "") not in _DURABLE_BRIDGE_REHYDRATION_IO_TYPES: - return None - prompt = payload.get("subgraph_prompt") - if not isinstance(prompt, dict): - return None - prompt_node = prompt.get(str(node_id)) - if not isinstance(prompt_node, dict): - return None - class_type = prompt_node.get("class_type") - inputs = prompt_node.get("inputs") - if ( - not isinstance(class_type, str) - or not class_type.strip() - or not isinstance(inputs, dict) - ): - return None - - normalized_inputs: dict[str, Any] = {} - has_linked_input = False - for input_name, input_value in inputs.items(): - normalized_value = _normalize_prompt_input_value(copy.deepcopy(input_value)) - if _is_link(normalized_value): - has_linked_input = True - continue - normalized_inputs[str(input_name)] = normalized_value - - if not has_linked_input: - node_data: dict[str, Any] = {"class_type": class_type} - custom_nodes_bundle = payload.get("custom_nodes_bundle") - if isinstance(custom_nodes_bundle, str) and custom_nodes_bundle.strip(): - node_data["custom_nodes_bundle"] = custom_nodes_bundle - return { - "kind": "single_node_output", - "node_data": node_data, - "node_inputs": normalized_inputs, - } - - required_node_ids = set( - _resolve_required_subgraph_nodes( - prompt=prompt, - execute_node_ids=[str(node_id)], - ) - ) - if str(node_id) not in required_node_ids: - return None - if _subgraph_contains_sampling_node(prompt, required_node_ids): - logger.info( - "Skipping durable %s bridge subgraph rehydration plan for node_id=%s class_type=%s because its dependency closure includes a sampler.", - io_type, - node_id, - class_type, - ) - return None - - rehydration_payload = copy.deepcopy(payload) - rehydration_payload[ - "component_id" - ] = f"{payload.get('component_id', 'component')}::rehydrate:{node_id}" - rehydration_payload["subgraph_prompt"] = { - str(current_node_id): copy.deepcopy(prompt[current_node_id]) - for current_node_id in prompt - if str(current_node_id) in required_node_ids - } - rehydration_payload["component_node_ids"] = sorted(required_node_ids) - rehydration_payload["execute_node_ids"] = [str(node_id)] - rehydration_payload["mapped_execute_node_ids"] = [] - rehydration_payload["static_execute_node_ids"] = [] - rehydration_payload["boundary_inputs"] = [ - { - **copy.deepcopy(boundary_input), - "targets": [ - copy.deepcopy(target) - for target in boundary_input.get("targets", []) - if str(target.get("node_id")) in required_node_ids - ], - } - for boundary_input in payload.get("boundary_inputs", []) - if any( - str(target.get("node_id")) in required_node_ids - for target in boundary_input.get("targets", []) - ) - ] - rehydration_payload["boundary_outputs"] = [ - { - "node_id": str(node_id), - "output_index": int(output_index), - "io_type": str(io_type), - "is_list": False, - "session_output": False, - "proxy_output_name": f"{node_id}_rehydrated_{output_index}", - } - ] - return { - "kind": "subgraph_output", - "payload": rehydration_payload, - } - - -def _restore_planned_remote_session_bridge_value( - record: RemoteSessionBridgeRecord, - *, - target_session_handle: RemoteSessionHandle, - node_mapping: dict[str, type[Any]] | None = None, - resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, -) -> Any | None: - """Restore one bridge value directly from a durable node rehydration plan.""" - if not isinstance(record.rehydration_plan, Mapping): - return None - - restore_started_at = time.perf_counter() - plan_kind = str(record.rehydration_plan.get("kind") or "") - if plan_kind == "single_node_output": - node_data = record.rehydration_plan.get("node_data") - node_inputs = record.rehydration_plan.get("node_inputs") - if not isinstance(node_data, Mapping) or not isinstance(node_inputs, Mapping): - return None - outputs = _execute_node_locally_raw( - dict(node_data), - dict(node_inputs), - node_mapping=node_mapping, - ) - output_index = record.output_index - elif plan_kind == "subgraph_output": - plan_payload = record.rehydration_plan.get("payload") - if not isinstance(plan_payload, Mapping): - return None - executable_payload = dict(plan_payload) - executable_payload["remote_session"] = target_session_handle.to_payload() - executable_payload.pop("clear_remote_session", None) - outputs = _execute_subgraph_prompt( - executable_payload, - _deserialize_remote_session_bridge_producer_inputs(record), - node_mapping, - ) - output_index = 0 - else: - return None - if output_index < 0 or output_index >= len(outputs): - raise RemoteSessionStateError( - f"Durable bridge rehydration plan for {record.bridge_key!r} did not produce output index {record.output_index}." - ) - restored_value = outputs[output_index] - _REMOTE_SESSION_STORE.put_bridge_output( - target_session_handle, - bridge_key=record.bridge_key, - 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 %s rehydration plan into session_id=%s.", - record.bridge_key, - record.rehydration_plan_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) - if replay_stack is None: - replay_stack = set() - _REMOTE_SESSION_BRIDGE_REPLAY_STATE.bridge_keys = replay_stack - return replay_stack - - -def _bridge_record_replays_sampling_node(record: RemoteSessionBridgeRecord) -> bool: - """Return whether replaying one bridge record would rerun sampler work.""" - execute_node_ids = [ - str(node_id) - for node_id in record.producer_payload.get("execute_node_ids", []) - if str(node_id) - ] - if not execute_node_ids: - return False - subgraph_prompt = record.producer_payload.get("subgraph_prompt") - if not isinstance(subgraph_prompt, dict): - return False - required_node_ids = _resolve_required_subgraph_nodes( - prompt=subgraph_prompt, - execute_node_ids=execute_node_ids, - ) - return _subgraph_contains_sampling_node(subgraph_prompt, required_node_ids) - - -def _subgraph_contains_sampling_node( - subgraph_prompt: Mapping[str, Any], - node_ids: Iterable[str], -) -> bool: - """Return whether any selected node looks like sampler work.""" - for node_id in node_ids: - prompt_node = subgraph_prompt.get(node_id) - if not isinstance(prompt_node, Mapping): - continue - class_type = str(prompt_node.get("class_type") or "") - if "sampler" in class_type.lower(): - return True - return False - - -def _rehydrate_remote_session_bridge_value( - ref: RemoteSessionBridgeRef, - *, - target_session_handle: RemoteSessionHandle | None, - node_mapping: dict[str, type[Any]] | None = None, - resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, -) -> Any: - """Replay one producer phase into the current session when the live value is gone.""" - if target_session_handle is None: - raise RemoteSessionStateError( - "Remote session bridge replay requires a target remote_session handle." - ) - - cached_value = _get_remote_session_bridge_value(ref.bridge_key) - if cached_value is not None: - restore_started_at = time.perf_counter() - _REMOTE_SESSION_STORE.put_bridge_output( - target_session_handle, - bridge_key=ref.bridge_key, - node_id=ref.node_id, - output_index=ref.output_index, - value=cached_value, - ) - if resolution_stats is not None: - resolution_stats.bridge_cache_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 directly from warm cache into session_id=%s.", - ref.bridge_key, - target_session_handle.session_id, - ) - return cached_value - - replay_stack = _remote_session_bridge_replay_stack() - if ref.bridge_key in replay_stack: - raise RemoteSessionStateError( - f"Detected recursive remote session bridge replay for {ref.bridge_key!r}." - ) - - record_lookup_started_at = time.perf_counter() - record = _REMOTE_SESSION_BRIDGE_STORE.get_record(ref.bridge_key) - if resolution_stats is not None: - resolution_stats.bridge_record_lookups += 1 - 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 - restored_value = _restore_planned_remote_session_bridge_value( - record, - target_session_handle=target_session_handle, - node_mapping=node_mapping, - resolution_stats=resolution_stats, - ) - if restored_value is not None: - return restored_value - if _bridge_record_replays_sampling_node(record): - message = ( - "Remote session bridge replay would rerun a sampler component " - f"for bridge_key={ref.bridge_key!r}; expected live or warm bridge reuse instead." - ) - logger.error(message) - raise RemoteSessionStateError(message) - replay_payload = copy.deepcopy(record.producer_payload) - replay_payload["remote_session"] = target_session_handle.to_payload() - replay_payload.pop("clear_remote_session", None) - if target_session_handle.prompt_id is not None: - replay_payload["prompt_id"] = target_session_handle.prompt_id - replay_inputs = _deserialize_remote_session_bridge_producer_inputs(record) - - logger.info( - "Replaying remote session bridge bridge_key=%s into session_id=%s via component=%s.", - ref.bridge_key, - target_session_handle.session_id, - replay_payload.get("component_id"), - ) - replay_stack.add(ref.bridge_key) - replay_started_at = time.perf_counter() - try: - _execute_subgraph_prompt(replay_payload, replay_inputs, node_mapping) - finally: - replay_stack.remove(ref.bridge_key) - if resolution_stats is not None: - resolution_stats.replay_count += 1 - resolution_stats.replay_seconds += time.perf_counter() - replay_started_at - - return _REMOTE_SESSION_STORE.get_output( - RemoteSessionValueRef( - session_id=target_session_handle.session_id, - node_id=ref.node_id, - output_index=ref.output_index, - ) - ) - - -def _resolve_remote_session_inputs( - hydrated_inputs: dict[str, Any], - *, - component_id: str | None = None, - target_session_handle: RemoteSessionHandle | None = None, - node_mapping: dict[str, type[Any]] | None = None, - resolution_stats: "_RemoteSessionBridgeResolutionStats | None" = None, -) -> dict[str, Any]: - """Resolve any remote-session value refs embedded in boundary inputs.""" - - def count_refs(value: Any) -> int: - """Return the number of live or durable session refs nested in one value.""" - if is_remote_session_value_ref_payload( - value - ) or is_remote_session_bridge_ref_payload(value): - return 1 - if isinstance(value, list | tuple): - return sum(count_refs(item) for item in value) - if isinstance(value, Mapping): - return sum(count_refs(item) for item in value.values()) - return 0 - - ref_counts_by_input = { - input_name: count_refs(input_value) - for input_name, input_value in hydrated_inputs.items() - } - ref_input_names = [ - input_name - for input_name, ref_count in ref_counts_by_input.items() - if ref_count > 0 - ] - total_ref_count = sum(ref_counts_by_input.values()) - if ref_input_names: - logger.info( - "Resolving %d remote session input refs for component=%s inputs=%s.", - total_ref_count, - component_id or "", - sorted(ref_input_names), - ) - if resolution_stats is not None: - resolution_stats.input_ref_count += total_ref_count - return { - input_name: _REMOTE_SESSION_STORE.resolve_value_with_bridges( - input_value, - target_session_handle=target_session_handle, - resolution_callback=( - lambda event_name, event_payload: _record_remote_session_resolution_event( - resolution_stats, - event_name, - event_payload, - ) - ) - if resolution_stats is not None - else None, - bridge_resolver=lambda ref: _rehydrate_remote_session_bridge_value( - ref, - target_session_handle=target_session_handle, - node_mapping=node_mapping, - resolution_stats=resolution_stats, - ), - ) - for input_name, input_value in hydrated_inputs.items() - } - - -def _log_remote_session_resolution_summary( - *, - component_id: str, - resolution_stats: "_RemoteSessionBridgeResolutionStats", -) -> None: - """Emit one high-signal summary for local fallback bridge resolution work.""" - 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 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, - resolution_stats.replay_seconds, - resolution_stats.direct_restore_seconds, - resolution_stats.session_restore_writes, - ) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + def _remote_worker_pool_affinity_key(slot_index: int) -> str: @@ -2338,9 +1451,9 @@ def _execute_subgraph_with_mapping( io_type=str(boundary_output.get("io_type") or "*"), output_value=output_value, ) - _REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) + _host_session_bridge._REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) _store_remote_session_bridge_value(bridge_record.bridge_key, output_value) - live_ref = _REMOTE_SESSION_STORE.put_bridge_output( + live_ref = _host_session_bridge._REMOTE_SESSION_STORE.put_bridge_output( session_handle, bridge_key=bridge_record.bridge_key, node_id=node_id, @@ -2495,11 +1608,11 @@ def _execute_subgraph_prompt( io_type=str(boundary_output.get("io_type") or "*"), output_value=output_value, ) - _REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) + _host_session_bridge._REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) _store_remote_session_bridge_value( bridge_record.bridge_key, output_value ) - live_ref = _REMOTE_SESSION_STORE.put_bridge_output( + live_ref = _host_session_bridge._REMOTE_SESSION_STORE.put_bridge_output( session_handle, bridge_key=bridge_record.bridge_key, node_id=node_id, @@ -2546,7 +1659,7 @@ def _short_circuit_restored_session_output_subgraph( node_id = str(boundary_output["node_id"]) output_index = int(boundary_output["output_index"]) try: - output_value = _REMOTE_SESSION_STORE.get_output( + output_value = _host_session_bridge._REMOTE_SESSION_STORE.get_output( RemoteSessionValueRef( session_id=session_handle.session_id, node_id=node_id, @@ -2564,9 +1677,9 @@ def _short_circuit_restored_session_output_subgraph( io_type=str(boundary_output.get("io_type") or "*"), output_value=output_value, ) - _REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) + _host_session_bridge._REMOTE_SESSION_BRIDGE_STORE.put_record(bridge_record) _store_remote_session_bridge_value(bridge_record.bridge_key, output_value) - live_ref = _REMOTE_SESSION_STORE.put_bridge_output( + live_ref = _host_session_bridge._REMOTE_SESSION_STORE.put_bridge_output( session_handle, bridge_key=bridge_record.bridge_key, node_id=node_id, @@ -2620,7 +1733,7 @@ def execute_subgraph_locally( payload.get("component_id"), session_handle.session_id, ) - _REMOTE_SESSION_STORE.clear_session(session_handle) + _host_session_bridge._REMOTE_SESSION_STORE.clear_session(session_handle) logger.info( "Local fallback subgraph %s completed with %d outputs.", payload.get("component_id"), @@ -5849,48 +4962,8 @@ def _modal_environment_name() -> str | None: return normalized or None -def _lookup_modal_session_bridge_store() -> Any: - """Return the shared Modal Dict that contains durable bridge records.""" - if modal is None or not hasattr(modal, "Dict"): - raise RuntimeError( - "Modal SDK is required to materialize a remote bridge locally." - ) - settings = get_settings() - cache_key = (settings.session_bridge_dict_name, _modal_environment_name()) - with _MODAL_SESSION_BRIDGE_DICTS_LOCK: - cached_store = _MODAL_SESSION_BRIDGE_DICTS.get(cache_key) - if cached_store is not None: - return cached_store - bridge_store = modal.Dict.from_name( - settings.session_bridge_dict_name, - environment_name=cache_key[1], - create_if_missing=True, - ) - with _MODAL_SESSION_BRIDGE_DICTS_LOCK: - _MODAL_SESSION_BRIDGE_DICTS[cache_key] = bridge_store - return bridge_store -def _lookup_modal_durable_volume() -> Any: - """Return the shared Modal Volume that contains durable bridge objects.""" - if modal is None or not hasattr(modal, "Volume"): - raise RuntimeError( - "Modal SDK is required to download a remote bridge object locally." - ) - settings = get_settings() - cache_key = (settings.volume_name, _modal_environment_name()) - with _MODAL_DURABLE_VOLUMES_LOCK: - cached_volume = _MODAL_DURABLE_VOLUMES.get(cache_key) - if cached_volume is not None: - return cached_volume - volume = modal.Volume.from_name( - settings.volume_name, - environment_name=cache_key[1], - create_if_missing=True, - ) - with _MODAL_DURABLE_VOLUMES_LOCK: - _MODAL_DURABLE_VOLUMES[cache_key] = volume - return volume def _modal_deploy_cache_key( diff --git a/tests/conftest.py b/tests/conftest.py index feae952..01b76f9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -455,6 +455,12 @@ def modal_container_logs_module(extension_package: object) -> object: return importlib.import_module(f"{PACKAGE_NAME}.remote.modal_container_logs") +@pytest.fixture(scope="session") +def host_session_bridge_module(extension_package: object) -> object: + """Return durable host session-bridge state and recovery helpers.""" + return importlib.import_module(f"{PACKAGE_NAME}.remote.host_session_bridge") + + @pytest.fixture(scope="session") def serialization_module(extension_package: object) -> object: """Return the serialization helpers module.""" diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index f55d550..c591970 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -1218,6 +1218,7 @@ def test_modal_local_bridge_materializer_downloads_without_blocking_event_loop( 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.""" @@ -1240,12 +1241,12 @@ def test_local_bridge_materialization_restores_durable_inline_output( serialized_output_io_type="IMAGE", ) monkeypatch.setattr( - remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, + host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record, ) - restored = remote_modal_app_module.materialize_remote_session_bridge_ref_locally( + restored = host_session_bridge_module.materialize_remote_session_bridge_ref_locally( bridge_ref.to_payload() ) @@ -1254,6 +1255,7 @@ def test_local_bridge_materialization_restores_durable_inline_output( 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.""" @@ -1303,7 +1305,7 @@ def test_local_bridge_materialization_downloads_object_backed_output( ) monkeypatch.setattr( - remote_modal_app_module, + host_session_bridge_module, "get_settings", lambda: types.SimpleNamespace( execution_mode="remote", @@ -1312,14 +1314,14 @@ def test_local_bridge_materialization_downloads_object_backed_output( ), ) monkeypatch.setattr( - remote_modal_app_module, + host_session_bridge_module, "modal", types.SimpleNamespace(Dict=FakeDict, Volume=FakeVolume), ) - remote_modal_app_module._MODAL_SESSION_BRIDGE_DICTS.clear() - remote_modal_app_module._MODAL_DURABLE_VOLUMES.clear() + host_session_bridge_module._MODAL_SESSION_BRIDGE_DICTS.clear() + host_session_bridge_module._MODAL_DURABLE_VOLUMES.clear() - restored = remote_modal_app_module.materialize_remote_session_bridge_ref_locally( + restored = host_session_bridge_module.materialize_remote_session_bridge_ref_locally( bridge_ref.to_payload() ) @@ -1336,8 +1338,8 @@ def test_local_bridge_materialization_downloads_object_backed_output( {"environment_name": None, "create_if_missing": True}, ) ] - remote_modal_app_module._MODAL_SESSION_BRIDGE_DICTS.clear() - remote_modal_app_module._MODAL_DURABLE_VOLUMES.clear() + host_session_bridge_module._MODAL_SESSION_BRIDGE_DICTS.clear() + host_session_bridge_module._MODAL_DURABLE_VOLUMES.clear() def test_parallel_local_dispatch_frontier_stops_at_nearest_remote_component( @@ -11016,6 +11018,7 @@ def test_modal_cloud_execute_mapped_subgraph_payload_preserves_assigned_lane_id( 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: @@ -11058,7 +11061,7 @@ def test_execute_subgraph_locally_round_trips_remote_session_bridge_refs( ) ) assert session_state_module.is_remote_session_bridge_ref_payload(static_outputs[0]) - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session( + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session( session_state_module.RemoteSessionHandle.from_payload(static_payload["remote_session"]) ) @@ -11131,6 +11134,7 @@ def test_execute_subgraph_locally_round_trips_remote_session_bridge_refs( def test_local_fallback_skips_seed_execution_when_session_outputs_are_already_restored( remote_modal_app_module: Any, + host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """Local fallback seed payloads should no-op when bridge outputs are already in session memory.""" @@ -11140,7 +11144,7 @@ def test_local_fallback_skips_seed_execution_when_session_outputs_are_already_re owner_component_id="component-1", ) restored_value = _FakeModelValue("restored-model") - remote_modal_app_module._REMOTE_SESSION_STORE.put_output( + host_session_bridge_module._REMOTE_SESSION_STORE.put_output( session_handle, node_id="5", output_index=0, @@ -11199,7 +11203,7 @@ def test_local_fallback_skips_seed_execution_when_session_outputs_are_already_re node_mapping={"ShouldNotRun": type("ShouldNotRun", (), {})}, ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(session_handle) + 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]) @@ -11207,6 +11211,7 @@ def test_local_fallback_skips_seed_execution_when_session_outputs_are_already_re def test_local_remote_app_rehydrates_bridge_refs_from_warm_value_cache_without_replay( remote_modal_app_module: Any, + host_session_bridge_module: Any, ) -> None: """The local fallback bridge fast path should restore retained values without replaying producers.""" bridge_key = "RSB_local_cached_bridge" @@ -11221,23 +11226,23 @@ def test_local_remote_app_rehydrates_bridge_refs_from_warm_value_cache_without_r output_index=0, session_id="session-source", ) - original_cache = dict(remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE) - original_order = list(remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER) + original_cache = dict(host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE) + original_order = list(host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER) try: seed_value = _CloneableCacheValue("warm-local-bridge-value") - remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear() - remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear() - remote_modal_app_module._store_remote_session_bridge_value(bridge_key, seed_value) + host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear() + host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear() + host_session_bridge_module._store_remote_session_bridge_value(bridge_key, seed_value) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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 = remote_modal_app_module._REMOTE_SESSION_STORE.get_output( + 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-7", @@ -11245,11 +11250,11 @@ def test_local_remote_app_rehydrates_bridge_refs_from_warm_value_cache_without_r ) ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) - remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear() - remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.update(original_cache) - remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear() - remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.extend(original_order) + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) + host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear() + host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.update(original_cache) + host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear() + host_session_bridge_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.extend(original_order) assert isinstance(restored_value, _CloneableCacheValue) assert restored_value.value == "warm-local-bridge-value" @@ -11263,6 +11268,7 @@ def test_local_remote_app_rehydrates_bridge_refs_from_warm_value_cache_without_r 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.""" @@ -11286,7 +11292,7 @@ def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_recor ] ] monkeypatch.setattr( - remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, + host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: remote_modal_app_module.RemoteSessionBridgeRecord( bridge_key=bridge_key, @@ -11308,13 +11314,13 @@ def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_recor resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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 = remote_modal_app_module._REMOTE_SESSION_STORE.get_output( + 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", @@ -11322,7 +11328,7 @@ def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_recor ) ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + 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"]) @@ -11336,6 +11342,7 @@ def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_recor 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( @@ -11343,14 +11350,14 @@ def test_remote_session_input_resolution_handles_nested_conditioning_bridge_refs prompt_id="prompt-1", owner_component_id="component-1", ) - first_ref = remote_modal_app_module._REMOTE_SESSION_STORE.put_bridge_output( + 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 = remote_modal_app_module._REMOTE_SESSION_STORE.put_bridge_output( + second_ref = host_session_bridge_module._REMOTE_SESSION_STORE.put_bridge_output( source_handle, bridge_key="RSB_b", node_id="508:item:1", @@ -11359,7 +11366,7 @@ def test_remote_session_input_resolution_handles_nested_conditioning_bridge_refs ) resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - resolved_inputs = remote_modal_app_module._resolve_remote_session_inputs( + resolved_inputs = host_session_bridge_module._resolve_remote_session_inputs( { "positive": [ remote_modal_app_module.RemoteSessionBridgeRef( @@ -11380,7 +11387,7 @@ def test_remote_session_input_resolution_handles_nested_conditioning_bridge_refs resolution_stats=resolution_stats, ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(source_handle) + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(source_handle) assert resolved_inputs["positive"] == [ [["conditioning-a", {"pooled_output": "pool-a"}]], @@ -11391,6 +11398,7 @@ def test_remote_session_input_resolution_handles_nested_conditioning_bridge_refs 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.""" @@ -11408,7 +11416,7 @@ def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_rec session_id="session-source", ) latent = {"samples": torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2)} - record = remote_modal_app_module._build_remote_session_bridge_record( + record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], @@ -11428,7 +11436,7 @@ def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_rec assert record.serialized_output is not None assert record.serialized_output_io_type == "LATENT" monkeypatch.setattr( - remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, + host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record, ) @@ -11442,13 +11450,13 @@ def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_rec resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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 = remote_modal_app_module._REMOTE_SESSION_STORE.get_output( + 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", @@ -11456,7 +11464,7 @@ def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_rec ) ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + 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"]) @@ -11476,6 +11484,7 @@ def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_rec ) def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_replay( remote_modal_app_module: Any, + host_session_bridge_module: Any, monkeypatch: Any, io_type: str, node_id: str, @@ -11494,7 +11503,7 @@ def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_r output_index=0, session_id="session-source", ) - record = remote_modal_app_module._build_remote_session_bridge_record( + record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "image-preview-component", "execute_node_ids": ["251"], @@ -11538,7 +11547,7 @@ def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_r } execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = [] monkeypatch.setattr( - remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, + host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record, ) @@ -11565,14 +11574,14 @@ def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_r resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + 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)] @@ -11582,6 +11591,7 @@ def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_r 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.""" @@ -11596,7 +11606,7 @@ def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without output_index=0, session_id="session-source", ) - record = remote_modal_app_module._build_remote_session_bridge_record( + record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "component-seed", "subgraph_prompt": { @@ -11612,7 +11622,7 @@ def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without io_type="MODEL", output_value=_FakeModelValue("seed-model"), ) - monkeypatch.setattr(remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record) + 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", @@ -11623,13 +11633,13 @@ def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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 = remote_modal_app_module._REMOTE_SESSION_STORE.get_output( + 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", @@ -11637,7 +11647,7 @@ def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without ) ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert isinstance(restored_value, _FakeModelValue) assert restored_value.value == "model::model.safetensors" @@ -11651,6 +11661,7 @@ def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without 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.""" @@ -11665,7 +11676,7 @@ def test_local_remote_app_rehydrates_clip_bridge_refs_from_durable_plan_without_ output_index=0, session_id="session-source", ) - record = remote_modal_app_module._build_remote_session_bridge_record( + record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], @@ -11692,7 +11703,7 @@ def test_local_remote_app_rehydrates_clip_bridge_refs_from_durable_plan_without_ ) assert record.rehydration_plan is not None assert record.rehydration_plan_io_type == "CLIP" - monkeypatch.setattr(remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record) + 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", @@ -11703,14 +11714,14 @@ def test_local_remote_app_rehydrates_clip_bridge_refs_from_durable_plan_without_ resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + 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 @@ -11719,6 +11730,7 @@ def test_local_remote_app_rehydrates_clip_bridge_refs_from_durable_plan_without_ def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgraph_plan( remote_modal_app_module: Any, + host_session_bridge_module: Any, monkeypatch: Any, ) -> None: """The local fallback should rehydrate linked MODEL bridges without sampler replay.""" @@ -11733,7 +11745,7 @@ def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgra output_index=0, session_id="session-source", ) - record = remote_modal_app_module._build_remote_session_bridge_record( + record = host_session_bridge_module._build_remote_session_bridge_record( payload={ "component_id": "sampler-component", "execute_node_ids": ["sampler-1"], @@ -11760,7 +11772,7 @@ def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgra ) assert record.rehydration_plan is not None assert record.rehydration_plan["kind"] == "subgraph_output" - monkeypatch.setattr(remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record) + 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( @@ -11781,14 +11793,14 @@ def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgra resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats() try: - restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value( + 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: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert isinstance(restored_value, _FakeModelValue) assert restored_value.value == "restored-lora-model" @@ -11799,6 +11811,7 @@ def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgra 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.""" @@ -11814,7 +11827,7 @@ def test_local_remote_app_refuses_sampler_ancestor_bridge_replay( session_id="session-source", ) monkeypatch.setattr( - remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, + host_session_bridge_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: remote_modal_app_module.RemoteSessionBridgeRecord( bridge_key=bridge_key, @@ -11849,14 +11862,14 @@ def test_local_remote_app_refuses_sampler_ancestor_bridge_replay( try: with pytest.raises(remote_modal_app_module.RemoteSessionStateError, match="rerun a sampler"): - remote_modal_app_module._rehydrate_remote_session_bridge_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: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) assert resolution_stats.bridge_record_lookups == 1 assert resolution_stats.replay_count == 0 diff --git a/tests/test_remote_durability.py b/tests/test_remote_durability.py index 93a8c8c..15f3d49 100644 --- a/tests/test_remote_durability.py +++ b/tests/test_remote_durability.py @@ -35,7 +35,7 @@ def test_remote_worker_affinity_separates_llm_and_comfy_pools( @pytest.mark.parametrize( "module_fixture_name", - ["modal_cloud_module", "remote_modal_app_module"], + ["modal_cloud_module", "host_session_bridge_module"], ) def test_modal_durable_store_reads_committed_object_without_volume_reload( module_fixture_name: str, @@ -793,6 +793,7 @@ def test_local_client_retries_exhausted_llm_memory_on_fresh_throughput_worker( def test_local_fallback_restores_object_backed_bridge_output( remote_modal_app_module: Any, + host_session_bridge_module: Any, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -800,7 +801,7 @@ def test_local_fallback_restores_object_backed_bridge_output( import torch object_store = remote_modal_app_module.FileDurableObjectStore(tmp_path) - monkeypatch.setattr(remote_modal_app_module, "_DURABLE_OBJECT_STORE", object_store) + monkeypatch.setattr(host_session_bridge_module, "_DURABLE_OBJECT_STORE", object_store) monkeypatch.setenv("COMFY_MODAL_BRIDGE_INLINE_MAX_BYTES", "1") remote_modal_app_module.get_settings.cache_clear() latent = {"samples": torch.arange(4, dtype=torch.float32).reshape(1, 1, 2, 2)} @@ -811,7 +812,7 @@ def test_local_fallback_restores_object_backed_bridge_output( ) try: - record = remote_modal_app_module._build_remote_session_bridge_record( + record = host_session_bridge_module._build_remote_session_bridge_record( payload={"component_id": "component-1"}, hydrated_inputs={"latent": latent}, node_id="node-1", @@ -820,13 +821,13 @@ def test_local_fallback_restores_object_backed_bridge_output( output_value=latent, ) restored_value = ( - remote_modal_app_module._restore_serialized_remote_session_bridge_value( + host_session_bridge_module._restore_serialized_remote_session_bridge_value( record, target_session_handle=target_handle, ) ) finally: - remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle) + host_session_bridge_module._REMOTE_SESSION_STORE.clear_session(target_handle) remote_modal_app_module.get_settings.cache_clear() assert record.producer_inputs_object is None @@ -840,7 +841,7 @@ def test_local_fallback_restores_object_backed_bridge_output( @pytest.mark.parametrize( "module_fixture_name", - ["modal_cloud_module", "remote_modal_app_module"], + ["modal_cloud_module", "host_session_bridge_module"], ) def test_direct_bridge_identity_includes_omitted_producer_inputs( module_fixture_name: str, @@ -875,7 +876,7 @@ def test_direct_bridge_identity_includes_omitted_producer_inputs( @pytest.mark.parametrize( "module_fixture_name", - ["modal_cloud_module", "remote_modal_app_module"], + ["modal_cloud_module", "host_session_bridge_module"], ) def test_literal_loader_plan_omits_producer_inputs( module_fixture_name: str, @@ -916,7 +917,7 @@ def test_literal_loader_plan_omits_producer_inputs( @pytest.mark.parametrize( "module_fixture_name", - ["modal_cloud_module", "remote_modal_app_module"], + ["modal_cloud_module", "host_session_bridge_module"], ) def test_linked_loader_plan_retains_only_required_boundary_inputs( module_fixture_name: str,