"""Remote component payload construction and proxy prompt rewriting.""" from __future__ import annotations import copy import logging import uuid from collections import deque from dataclasses import dataclass from typing import Any, Mapping if __package__: from .component_planning import ( _boundary_output_payload, _component_ancestors_of_local_source, _component_has_local_reentry_dependency, _component_has_parallel_local_remote_fanout, _component_upstream_closure, _filter_boundary_inputs_for_node_ids, _filter_boundary_outputs_for_node_ids, _order_execute_node_ids_for_transportable_splits, _proxy_boundary_output_is_list, _subgraph_topological_node_order, _subset_component_prompt, ) from .execution_environments import ExecutionProvider from .modal_executor_node import ( MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, MODAL_ARTIFACT_FINALIZER_NODE_ID, MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID, MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, ensure_modal_artifact_finalizer_registered, ensure_modal_component_proxy_node_registered, ensure_modal_local_bridge_materializer_registered, ensure_modal_parallel_local_passthrough_registered, register_cache_friendly_proxy_payload, register_modal_map_input_warmup_context, registered_proxy_execution_payload, update_registered_proxy_payload_fields, ) from .prompt_payload_metadata import ( _attach_snapshot_profile_key, _serialize_boundary_input_specs, ) from .remote_graph_analysis import ( _build_consumer_map, _downstream_node_ids_from_targets, _is_link, _iter_payload_input_strings, _node_output_refs, _normalize_output_metadata, _remote_output_io_type, _remote_output_is_list, ) from .remote_plan_types import ( BoundaryInputSpec, InputTarget, LinkedOutputRef, ModalPromptValidationError, ProducedPhaseOutputSpec, RemoteComponentPlan, ) from .session_state import RemoteSessionHandle from .settings import ModalSyncSettings from .sync_engine import AssetSyncRequestCache, ModalAssetSyncEngine, SyncedAsset else: # pragma: no cover - flat import inside the Modal container. from component_planning import ( _boundary_output_payload, _component_ancestors_of_local_source, _component_has_local_reentry_dependency, _component_has_parallel_local_remote_fanout, _component_upstream_closure, _filter_boundary_inputs_for_node_ids, _filter_boundary_outputs_for_node_ids, _order_execute_node_ids_for_transportable_splits, _proxy_boundary_output_is_list, _subgraph_topological_node_order, _subset_component_prompt, ) from execution_environments import ExecutionProvider from modal_executor_node import ( MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, MODAL_ARTIFACT_FINALIZER_NODE_ID, MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID, MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, ensure_modal_artifact_finalizer_registered, ensure_modal_component_proxy_node_registered, ensure_modal_local_bridge_materializer_registered, ensure_modal_parallel_local_passthrough_registered, register_cache_friendly_proxy_payload, register_modal_map_input_warmup_context, registered_proxy_execution_payload, update_registered_proxy_payload_fields, ) from prompt_payload_metadata import ( _attach_snapshot_profile_key, _serialize_boundary_input_specs, ) from remote_graph_analysis import ( _build_consumer_map, _downstream_node_ids_from_targets, _is_link, _iter_payload_input_strings, _node_output_refs, _normalize_output_metadata, _remote_output_io_type, _remote_output_is_list, ) from remote_plan_types import ( BoundaryInputSpec, InputTarget, LinkedOutputRef, ModalPromptValidationError, ProducedPhaseOutputSpec, RemoteComponentPlan, ) from session_state import RemoteSessionHandle from settings import ModalSyncSettings from sync_engine import AssetSyncRequestCache, ModalAssetSyncEngine, SyncedAsset logger = logging.getLogger(__name__) _SPECULATIVE_PREWARM_TARGET_KEY = "speculative_remote_prewarm_target" _SPECULATIVE_PREWARM_PAYLOAD_FIELDS = frozenset( { "component_id", "prompt_id", "modal_gpu", "execution_provider", "execution_environment_id", "remote_worker_affinity_group", "remote_local_gap_pool", "subgraph_prompt", "requires_volume_reload", "volume_reload_marker", "uploaded_volume_paths", "custom_nodes_bundle", "snapshot_profile_key", } ) def _sync_component_prompt_inputs( component: RemoteComponentPlan, rewritten_prompt: dict[str, Any], sync_engine: ModalAssetSyncEngine, request_cache: AssetSyncRequestCache, status_callback: Any | None = None, ) -> tuple[dict[str, Any], list[SyncedAsset]]: """Build a synced prompt payload for one remote component.""" component_prompt: dict[str, Any] = {} synced_assets: list[SyncedAsset] = [] logger.info( "Syncing prompt inputs for remote component %s with %d nodes.", component.representative_node_id, len(component.node_ids), ) for node_id in component.node_ids: prompt_node = rewritten_prompt[node_id] synced_inputs, node_assets = sync_engine.sync_prompt_inputs( copy.deepcopy(prompt_node.get("inputs", {})), status_callback=status_callback, request_cache=request_cache, ) synced_assets.extend(node_assets) component_prompt[node_id] = { "class_type": str(prompt_node["class_type"]), "inputs": synced_inputs, "_meta": copy.deepcopy(prompt_node.get("_meta", {})), } logger.info( "Prepared remote node %s (%s) with %d synced assets.", node_id, component_prompt[node_id]["class_type"], len(node_assets), ) logger.info( "Finished syncing remote component %s with %d total synced assets.", component.representative_node_id, len(synced_assets), ) return component_prompt, synced_assets def _deduplicate_synced_assets(synced_assets: list[SyncedAsset]) -> list[SyncedAsset]: """Return one request summary record per content-addressed remote path.""" unique_assets_by_remote_path: dict[str, SyncedAsset] = {} for synced_asset in synced_assets: unique_assets_by_remote_path.setdefault(synced_asset.remote_path, synced_asset) return list(unique_assets_by_remote_path.values()) @dataclass(frozen=True) class _PayloadBuildContext: """Hold immutable inputs shared by component payload builders.""" component_prompt: dict[str, Any] signature_prompt: dict[str, Any] prompt_id: Any extra_data: dict[str, Any] | None settings: ModalSyncSettings requires_volume_reload: bool volume_reload_marker: str | None custom_nodes_bundle_path: str | None uploaded_volume_paths: list[str] terminate_container_on_error: bool nodes_module: Any def _remote_worker_affinity_group( context: _PayloadBuildContext, component_node_ids: list[str], ) -> str: """Return the worker-pool group required by one remote component phase.""" class_types = { str((context.component_prompt.get(node_id) or {}).get("class_type") or "") for node_id in component_node_ids } return "llm" if "ModalLLM" in class_types else "comfy" def _build_subgraph_payload( context: _PayloadBuildContext, *, component_id: str, component_node_ids: list[str], boundary_inputs: list[BoundaryInputSpec], boundary_outputs: list[dict[str, Any]], execute_node_ids: list[str], remote_session: dict[str, Any] | None = None, clear_remote_session: bool = False, mapped_progress_display_node_id: str | None = None, ) -> dict[str, Any]: """Build one ordinary subgraph payload for a proxy node.""" payload = { "payload_kind": "subgraph", "component_id": component_id, "prompt_id": context.prompt_id, "modal_gpu": context.settings.modal_gpu, "remote_worker_affinity_group": _remote_worker_affinity_group( context, component_node_ids, ), "component_node_ids": list(component_node_ids), "subgraph_prompt": _subset_component_prompt( context.component_prompt, component_node_ids, ), "boundary_inputs": _serialize_boundary_input_specs( boundary_inputs, signature_prompt=context.signature_prompt, ), "boundary_outputs": copy.deepcopy(boundary_outputs), "execute_node_ids": list(execute_node_ids), "extra_data": copy.deepcopy(context.extra_data or {}), "requires_volume_reload": context.requires_volume_reload, "volume_reload_marker": context.volume_reload_marker, "uploaded_volume_paths": list(context.uploaded_volume_paths), "terminate_container_on_error": context.terminate_container_on_error, "custom_nodes_bundle": context.custom_nodes_bundle_path, } if remote_session is not None: payload["remote_session"] = copy.deepcopy(remote_session) if clear_remote_session: payload["clear_remote_session"] = True if mapped_progress_display_node_id is not None: payload["mapped_progress_display_node_id"] = mapped_progress_display_node_id return payload def _build_phase_payloads_for_transportable_splits( component: RemoteComponentPlan, context: _PayloadBuildContext, remote_session: dict[str, Any] | None, ) -> list[dict[str, Any]] | None: """Return ordered phases for local feedback or parallel local fanout.""" if len(component.execute_node_ids) <= 1: return None has_local_reentry_dependency = _component_has_local_reentry_dependency( prompt=context.signature_prompt, component=component, ) has_parallel_local_remote_fanout = _component_has_parallel_local_remote_fanout( component ) if not has_local_reentry_dependency and not has_parallel_local_remote_fanout: logger.info( "Keeping remote component %s as one proxy because execute targets %s " "have neither a local re-entry dependency nor parallel local/remote fanout.", component.representative_node_id, component.execute_node_ids, ) return None if component.local_tap_node_ids: logger.info( "Allowing remote component %s with local tap nodes %s to split because it has a local re-entry dependency.", component.representative_node_id, component.local_tap_node_ids, ) if has_parallel_local_remote_fanout: logger.info( "Forcing ordered phases for remote component %s because a remote output " "feeds both a non-returning local branch and later remote execution.", component.representative_node_id, ) component_node_id_set = set(component.node_ids) topological_node_ids = _subgraph_topological_node_order( context.component_prompt, component_node_id_set ) remaining_node_ids = set(component.node_ids) remaining_execute_node_ids = [ node_id for node_id in topological_node_ids if node_id in set(component.execute_node_ids) ] remaining_execute_node_ids = _order_execute_node_ids_for_transportable_splits( prompt=context.signature_prompt, component_prompt=context.component_prompt, component_node_ids=component_node_id_set, execute_node_ids=remaining_execute_node_ids, ) if len(remaining_execute_node_ids) <= 1: return None phase_payloads: list[dict[str, Any]] = [] produced_outputs_by_source: dict[LinkedOutputRef, ProducedPhaseOutputSpec] = {} local_boundary_outputs_by_source = { boundary_output.source: boundary_output for boundary_output in component.boundary_outputs } bridge_output_counter = 0 while remaining_execute_node_ids: target_node_id = remaining_execute_node_ids[0] phase_node_ids = sorted( _component_upstream_closure( prompt=context.component_prompt, seed_node_ids={target_node_id}, candidate_node_ids=remaining_node_ids, ) ) phase_node_id_set = set(phase_node_ids) if not phase_node_ids: raise ModalPromptValidationError( f"Unable to derive split phase nodes for remote component {component.representative_node_id}." ) phase_boundary_inputs = _filter_boundary_inputs_for_node_ids( component.boundary_inputs, phase_node_id_set, ) phase_boundary_inputs_by_name = { boundary_input.proxy_input_name: boundary_input for boundary_input in phase_boundary_inputs } for phase_node_id in phase_node_ids: prompt_node = context.component_prompt.get(phase_node_id) if prompt_node is None: continue for input_name, input_value in ( prompt_node.get("inputs") or {} ).items(): if not _is_link(input_value): continue source = LinkedOutputRef( node_id=str(input_value[0]), output_index=int(input_value[1]), ) if source.node_id in phase_node_id_set: continue produced_output = produced_outputs_by_source.get(source) if produced_output is None: continue boundary_input = phase_boundary_inputs_by_name.get( produced_output.proxy_output_name ) if boundary_input is None: boundary_input = BoundaryInputSpec( proxy_input_name=produced_output.proxy_output_name, source=source, io_type=produced_output.io_type, ) phase_boundary_inputs.append(boundary_input) phase_boundary_inputs_by_name[ boundary_input.proxy_input_name ] = boundary_input boundary_input.targets.append( InputTarget(node_id=phase_node_id, input_name=str(input_name)) ) phase_boundary_outputs: list[dict[str, Any]] = [] phase_output_names_by_source: dict[LinkedOutputRef, str] = {} for boundary_output in _filter_boundary_outputs_for_node_ids( component.boundary_outputs, phase_node_id_set, ): phase_boundary_outputs.append(_boundary_output_payload(boundary_output)) phase_output_names_by_source[ boundary_output.source ] = boundary_output.proxy_output_name pending_node_ids = remaining_node_ids - phase_node_id_set for pending_node_id in sorted(pending_node_ids): prompt_node = context.component_prompt.get(pending_node_id) if prompt_node is None: continue for input_value in (prompt_node.get("inputs") or {}).values(): if not _is_link(input_value): continue source = LinkedOutputRef( node_id=str(input_value[0]), output_index=int(input_value[1]), ) if ( source.node_id not in phase_node_id_set or source in produced_outputs_by_source ): continue local_boundary_output = local_boundary_outputs_by_source.get(source) io_type = ( local_boundary_output.io_type if local_boundary_output is not None else str( _remote_output_io_type( prompt=context.component_prompt, node_id=source.node_id, output_index=source.output_index, nodes_module=context.nodes_module, ) or "*" ) ) is_list = ( local_boundary_output.is_list if local_boundary_output is not None else _remote_output_is_list( prompt=context.component_prompt, node_id=source.node_id, output_index=source.output_index, nodes_module=context.nodes_module, ) ) proxy_output_name = phase_output_names_by_source.get(source) if proxy_output_name is None: proxy_output_name = f"phase_bridge_{bridge_output_counter}" bridge_output_counter += 1 phase_boundary_outputs.append( { "proxy_output_name": proxy_output_name, "node_id": source.node_id, "output_index": source.output_index, "io_type": io_type, "is_list": is_list, "preview_target_node_ids": [], "session_output": True, } ) phase_output_names_by_source[source] = proxy_output_name produced_outputs_by_source[source] = ProducedPhaseOutputSpec( proxy_output_name=proxy_output_name, source=source, io_type=io_type, is_list=is_list, session_output=True, ) phase_execute_node_ids = [ node_id for node_id in component.execute_node_ids if node_id in phase_node_id_set ] phase_payloads.append( _build_subgraph_payload( context, component_id=str(target_node_id), component_node_ids=phase_node_ids, boundary_inputs=phase_boundary_inputs, boundary_outputs=phase_boundary_outputs, execute_node_ids=phase_execute_node_ids, ) ) remaining_node_ids -= phase_node_id_set remaining_execute_node_ids = [ node_id for node_id in remaining_execute_node_ids if node_id not in phase_node_id_set ] has_session_bridges = any( bool(boundary_output.get("session_output")) for phase_payload in phase_payloads for boundary_output in phase_payload.get("boundary_outputs", []) ) if not phase_payloads or len(phase_payloads) <= 1: return None active_remote_session = remote_session if active_remote_session is None and has_session_bridges: active_remote_session = RemoteSessionHandle( session_id=uuid.uuid4().hex, prompt_id=( str(context.prompt_id) if context.prompt_id is not None else None ), owner_component_id=component.representative_node_id, ).to_payload() if active_remote_session is not None: for phase_index, phase_payload in enumerate(phase_payloads): phase_payload["remote_session"] = copy.deepcopy(active_remote_session) if phase_index == len(phase_payloads) - 1: phase_payload["clear_remote_session"] = True logger.info( "Split ordinary remote component %s into ordered phases: %s", component.representative_node_id, [ { "component_id": phase_payload["component_id"], "component_node_ids": phase_payload["component_node_ids"], "execute_node_ids": phase_payload["execute_node_ids"], } for phase_payload in phase_payloads ], ) return phase_payloads def _build_component_payload( component: RemoteComponentPlan, component_prompt: dict[str, Any], signature_prompt: dict[str, Any], extra_data: dict[str, Any] | None, settings: ModalSyncSettings, requires_volume_reload: bool, volume_reload_marker: str | None, custom_nodes_bundle: SyncedAsset | None, uploaded_volume_paths: list[str], terminate_container_on_error: bool, nodes_module: Any, remote_session: dict[str, Any] | None = None, ) -> dict[str, Any]: """Build the serialized execution payload for one remote component.""" prompt_id = (extra_data or {}).get("prompt_id") custom_nodes_bundle_path = ( custom_nodes_bundle.remote_path if custom_nodes_bundle is not None else None ) context = _PayloadBuildContext( component_prompt=component_prompt, signature_prompt=signature_prompt, prompt_id=prompt_id, extra_data=extra_data, settings=settings, requires_volume_reload=requires_volume_reload, volume_reload_marker=volume_reload_marker, custom_nodes_bundle_path=custom_nodes_bundle_path, uploaded_volume_paths=uploaded_volume_paths, terminate_container_on_error=terminate_container_on_error, nodes_module=nodes_module, ) split_phase_payloads = _build_phase_payloads_for_transportable_splits( component, context, remote_session, ) if split_phase_payloads is not None: return _attach_snapshot_profile_key( {"split_proxy_payloads": split_phase_payloads}, settings, ) payload = { "payload_kind": "mapped_subgraph" if component.mapped_boundary_input_name else "subgraph", "component_id": component.representative_node_id, "prompt_id": prompt_id, "modal_gpu": settings.modal_gpu, "remote_worker_affinity_group": _remote_worker_affinity_group( context, list(component.node_ids) ), "component_node_ids": list(component.node_ids), "subgraph_prompt": component_prompt, "boundary_inputs": _serialize_boundary_input_specs( component.boundary_inputs, signature_prompt=signature_prompt, ), "boundary_outputs": [ _boundary_output_payload( boundary_output, mapped_output=( boundary_output.source.node_id in set(component.mapped_node_ids) if component.mapped_boundary_input_name else None ), ) for boundary_output in component.boundary_outputs ], "execute_node_ids": list(component.execute_node_ids), "mapped_execute_node_ids": list(component.mapped_execute_node_ids), "static_execute_node_ids": list(component.static_execute_node_ids), "extra_data": copy.deepcopy(extra_data or {}), "requires_volume_reload": requires_volume_reload, "volume_reload_marker": volume_reload_marker, "uploaded_volume_paths": list(uploaded_volume_paths), "terminate_container_on_error": terminate_container_on_error, "custom_nodes_bundle": custom_nodes_bundle_path, "mapped_input": ( { "proxy_input_name": component.mapped_boundary_input_name, "io_type": str(component.mapped_boundary_input_io_type or "*"), } if component.mapped_boundary_input_name else None ), } if remote_session is not None: payload["remote_session"] = copy.deepcopy(remote_session) payload["clear_remote_session"] = True logger.info( "Built remote payload for component %s: boundary_inputs=%d boundary_outputs=%d execute_nodes=%s custom_nodes_bundle=%s", component.representative_node_id, len(payload["boundary_inputs"]), len(payload["boundary_outputs"]), payload["execute_node_ids"], payload["custom_nodes_bundle"], ) logger.info( "Remote payload for component %s requires_volume_reload=%s volume_reload_marker=%s", component.representative_node_id, requires_volume_reload, volume_reload_marker, ) if component.mapped_boundary_input_name: static_node_id_set = set(component.static_node_ids) mapped_node_id_set = set(component.mapped_node_ids) static_boundary_inputs = _filter_boundary_inputs_for_node_ids( component.boundary_inputs, static_node_id_set, ) mapped_boundary_inputs = _filter_boundary_inputs_for_node_ids( component.boundary_inputs, mapped_node_id_set, ) static_boundary_outputs = _filter_boundary_outputs_for_node_ids( component.boundary_outputs, static_node_id_set, ) mapped_boundary_outputs = _filter_boundary_outputs_for_node_ids( component.boundary_outputs, mapped_node_id_set, ) static_bridge_outputs = [ { "proxy_output_name": boundary_spec.proxy_name, "node_id": boundary_spec.source.node_id, "output_index": boundary_spec.source.output_index, "io_type": boundary_spec.io_type, "is_list": boundary_spec.is_list, "preview_target_node_ids": [], "session_output": True, } for boundary_spec in component.static_to_mapped_boundaries ] payload["static_to_mapped_boundaries"] = [ { "proxy_name": boundary_spec.proxy_name, "node_id": boundary_spec.source.node_id, "output_index": boundary_spec.source.output_index, "io_type": boundary_spec.io_type, "is_list": boundary_spec.is_list, "targets": [ {"node_id": target.node_id, "input_name": target.input_name} for target in boundary_spec.targets ], } for boundary_spec in component.static_to_mapped_boundaries ] payload["static_phase"] = { "component_node_ids": list(component.static_node_ids), "subgraph_prompt": _subset_component_prompt( component_prompt, component.static_node_ids ), "boundary_inputs": _serialize_boundary_input_specs( static_boundary_inputs, signature_prompt=signature_prompt, ), "boundary_outputs": [ _boundary_output_payload(boundary_output) for boundary_output in static_boundary_outputs ] + static_bridge_outputs, "execute_node_ids": list(component.static_execute_node_ids), } payload["mapped_phase"] = { "component_node_ids": list(component.mapped_node_ids), "subgraph_prompt": _subset_component_prompt( component_prompt, component.mapped_node_ids ), "boundary_inputs": _serialize_boundary_input_specs( mapped_boundary_inputs + [ BoundaryInputSpec( proxy_input_name=boundary_spec.proxy_name, source=boundary_spec.source, io_type=boundary_spec.io_type, targets=list(boundary_spec.targets), ) for boundary_spec in component.static_to_mapped_boundaries ], signature_prompt=signature_prompt, ), "boundary_outputs": [ _boundary_output_payload(boundary_output, mapped_output=True) for boundary_output in mapped_boundary_outputs ], "execute_node_ids": list(component.mapped_execute_node_ids), } if not component.static_node_ids: return _attach_snapshot_profile_key(payload, settings) if remote_session is None: remote_session = RemoteSessionHandle( session_id=uuid.uuid4().hex, prompt_id=(str(prompt_id) if prompt_id is not None else None), owner_component_id=component.representative_node_id, ).to_payload() logger.info( "Split hybrid component %s into static nodes=%s and mapped nodes=%s using remote_session session_id=%s with %d static bridge outputs.", component.representative_node_id, component.static_node_ids, component.mapped_node_ids, remote_session["session_id"], len(static_bridge_outputs), ) payload = { "split_proxy_payloads": { "static": _build_subgraph_payload( context, component_id=component.static_node_ids[0], component_node_ids=list(component.static_node_ids), boundary_inputs=static_boundary_inputs, boundary_outputs=[ _boundary_output_payload(boundary_output) for boundary_output in static_boundary_outputs ] + static_bridge_outputs, execute_node_ids=list(component.static_execute_node_ids), remote_session=remote_session, ), "mapped": _build_subgraph_payload( context, component_id=f"{component.representative_node_id}__mapped", component_node_ids=list(component.mapped_node_ids), boundary_inputs=mapped_boundary_inputs + [ BoundaryInputSpec( proxy_input_name=boundary_spec.proxy_name, source=boundary_spec.source, io_type=boundary_spec.io_type, targets=list(boundary_spec.targets), ) for boundary_spec in component.static_to_mapped_boundaries ], boundary_outputs=[ _boundary_output_payload(boundary_output) for boundary_output in mapped_boundary_outputs ], execute_node_ids=list(component.mapped_execute_node_ids), remote_session=remote_session, clear_remote_session=True, mapped_progress_display_node_id=component.static_node_ids[0], ), } } payload["split_proxy_payloads"]["mapped"]["static_to_mapped_boundaries"] = [ { "proxy_name": boundary_spec.proxy_name, "node_id": boundary_spec.source.node_id, "output_index": boundary_spec.source.output_index, "io_type": boundary_spec.io_type, "is_list": boundary_spec.is_list, "targets": [ {"node_id": target.node_id, "input_name": target.input_name} for target in boundary_spec.targets ], } for boundary_spec in component.static_to_mapped_boundaries ] payload["split_proxy_payloads"]["mapped"]["static_phase"] = { "component_node_ids": list(component.static_node_ids), "subgraph_prompt": _subset_component_prompt( component_prompt, component.static_node_ids ), "boundary_inputs": _serialize_boundary_input_specs( static_boundary_inputs, signature_prompt=signature_prompt, ), "boundary_outputs": copy.deepcopy(static_bridge_outputs), "execute_node_ids": list( dict.fromkeys( boundary_spec.source.node_id for boundary_spec in component.static_to_mapped_boundaries ) ), } return _attach_snapshot_profile_key(payload, settings) def _component_uploaded_volume_paths( *, component_prompt: dict[str, Any], synced_assets: list[SyncedAsset], custom_nodes_bundle: SyncedAsset | None, ) -> list[str]: """Return newly uploaded mounted-volume paths that this component can actually reference.""" referenced_paths: set[str] = set() for prompt_node in component_prompt.values(): if not isinstance(prompt_node, dict): continue inputs = prompt_node.get("inputs", {}) if not isinstance(inputs, dict): continue for input_value in inputs.values(): for candidate_path in _iter_payload_input_strings(input_value): if isinstance(candidate_path, str) and candidate_path.startswith("/"): referenced_paths.add(candidate_path) uploaded_paths = { asset.remote_path for asset in synced_assets if asset.uploaded and asset.remote_path in referenced_paths } if custom_nodes_bundle is not None and custom_nodes_bundle.uploaded: uploaded_paths.add(custom_nodes_bundle.remote_path) return sorted(uploaded_paths) def _contains_output_node( node_ids: list[str], *, rewritten_prompt: dict[str, Any], nodes_module: Any, ) -> bool: """Return whether one node subset contains an output node.""" for node_id in node_ids: prompt_node = rewritten_prompt.get(node_id) if prompt_node is None: continue node_class = nodes_module.NODE_CLASS_MAPPINGS.get( str(prompt_node["class_type"]) ) if node_class is not None and getattr(node_class, "OUTPUT_NODE", False): return True return False def _proxy_inputs_for_boundary_inputs( boundary_inputs: list[dict[str, Any]], *, rewritten_prompt: dict[str, Any], ) -> dict[str, Any]: """Resolve one proxy input mapping from the current prompt graph.""" proxy_inputs: dict[str, Any] = {} for boundary_input in boundary_inputs: current_input_value: Any = None for target in boundary_input.get("targets", []): target_prompt_node = rewritten_prompt.get(str(target["node_id"])) if target_prompt_node is None: continue target_input_value = (target_prompt_node.get("inputs") or {}).get( str(target["input_name"]) ) if _is_link(target_input_value): current_input_value = list(target_input_value) break if current_input_value is None: raise ModalPromptValidationError( "Unable to resolve proxy boundary input wiring while rewriting split Modal proxies." ) proxy_inputs[str(boundary_input["proxy_input_name"])] = current_input_value return proxy_inputs def _register_proxy_node( *, prompt_node_id: str, payload_mapping: dict[str, Any], proxy_inputs: dict[str, Any], meta: dict[str, Any], is_output_node: bool, rewritten_prompt: dict[str, Any], nodes_module: Any, ) -> None: """Insert one dynamic proxy node into the rewritten prompt.""" boundary_outputs = list(payload_mapping.get("boundary_outputs", [])) proxy_node_id = ensure_modal_component_proxy_node_registered( output_types=tuple(str(output["io_type"]) for output in boundary_outputs), output_names=tuple( str(output["proxy_output_name"]) for output in boundary_outputs ), output_is_list=tuple( _proxy_boundary_output_is_list(output) for output in boundary_outputs ), nodes_module=nodes_module, is_output_node=is_output_node, include_completion_output=True, ) proxy_inputs["original_node_data"] = register_cache_friendly_proxy_payload( prompt_node_id, payload_mapping, ) rewritten_prompt[prompt_node_id] = { "class_type": proxy_node_id, "inputs": proxy_inputs, "_meta": copy.deepcopy(meta), } def _rewrite_component_into_proxy( component: RemoteComponentPlan, rewritten_prompt: dict[str, Any], payload: dict[str, Any], nodes_module: Any, ) -> list[str]: """Replace a remote component with a single proxy node in the prompt.""" split_proxy_payloads = payload.get("split_proxy_payloads") component_node_id_set = set(component.node_ids) if isinstance(split_proxy_payloads, list): phase_payloads = [dict(phase_payload) for phase_payload in split_proxy_payloads] phase_proxy_node_ids: list[str] = [] produced_output_indices_by_name: dict[str, list[Any]] = {} replacement_output_indices: dict[LinkedOutputRef, list[Any]] = {} materializer_output_by_consumer_and_source: dict[ tuple[str, LinkedOutputRef], list[Any] ] = {} proxy_output_by_consumer_and_source: dict[ tuple[str, LinkedOutputRef], list[Any] ] = {} materializer_proxy_output_by_node_id: dict[str, list[Any]] = {} boundary_output_specs_by_source = { boundary_output.source: boundary_output for boundary_output in component.boundary_outputs } component_proxy_node_ids: set[str] = set() phase_proxy_inputs_by_node_id: dict[str, dict[str, Any]] = {} phase_proxy_meta_by_node_id: dict[str, dict[str, Any]] = {} if any( spec.local_materializer_node_id is not None for spec in component.boundary_outputs ): ensure_modal_local_bridge_materializer_registered(nodes_module) for phase_payload in phase_payloads: phase_proxy_node_id = str(phase_payload["component_id"]) while ( phase_proxy_node_id in rewritten_prompt and phase_proxy_node_id not in component_node_id_set ): phase_proxy_node_id = f"{phase_proxy_node_id}_proxy" phase_payload["component_id"] = phase_proxy_node_id phase_proxy_node_ids.append(phase_proxy_node_id) component_proxy_node_ids.add(phase_proxy_node_id) for phase_payload in phase_payloads: phase_proxy_node_id = str(phase_payload["component_id"]) phase_proxy_inputs = _proxy_inputs_for_boundary_inputs( list(phase_payload["boundary_inputs"]), rewritten_prompt=rewritten_prompt, ) for boundary_input in phase_payload.get("boundary_inputs", []): proxy_input_name = str(boundary_input["proxy_input_name"]) produced_output_index = produced_output_indices_by_name.get( proxy_input_name ) if produced_output_index is None: continue phase_proxy_inputs[proxy_input_name] = list(produced_output_index) first_phase_node_id = str(phase_payload["component_node_ids"][0]) phase_proxy_meta = copy.deepcopy( rewritten_prompt[first_phase_node_id].get("_meta", {}) ) phase_proxy_inputs_by_node_id[phase_proxy_node_id] = phase_proxy_inputs phase_proxy_meta_by_node_id[phase_proxy_node_id] = phase_proxy_meta for output_index, boundary_output in enumerate( phase_payload.get("boundary_outputs", []) ): proxy_output_name = str(boundary_output["proxy_output_name"]) proxy_output = [phase_proxy_node_id, output_index] produced_output_indices_by_name[proxy_output_name] = proxy_output source = LinkedOutputRef( node_id=str(boundary_output["node_id"]), output_index=int(boundary_output["output_index"]), ) boundary_output_spec = boundary_output_specs_by_source.get(source) if boundary_output_spec is not None: for ( consumer_node_id ) in boundary_output_spec.session_consumer_node_ids: proxy_output_by_consumer_and_source[ (consumer_node_id, source) ] = proxy_output materializer_node_id = ( boundary_output_spec.local_materializer_node_id ) if materializer_node_id is not None: materializer_proxy_output_by_node_id[ materializer_node_id ] = proxy_output for ( consumer_node_id ) in boundary_output_spec.local_materializer_consumer_node_ids: materializer_output_by_consumer_and_source[ (consumer_node_id, source) ] = [materializer_node_id, 0] if bool(boundary_output.get("session_output")): continue replacement_output_indices[source] = proxy_output for node_id in component.node_ids: rewritten_prompt.pop(node_id, None) for phase_payload in phase_payloads: phase_proxy_node_id = str(phase_payload["component_id"]) _register_proxy_node( prompt_node_id=phase_proxy_node_id, payload_mapping=phase_payload, proxy_inputs=phase_proxy_inputs_by_node_id[phase_proxy_node_id], meta=phase_proxy_meta_by_node_id[phase_proxy_node_id], is_output_node=_contains_output_node( list(phase_payload["component_node_ids"]), rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ), rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ) for ( materializer_node_id, proxy_output, ) in materializer_proxy_output_by_node_id.items(): rewritten_prompt[materializer_node_id] = { "class_type": MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID, "inputs": {"bridge_ref": list(proxy_output)}, "_meta": {"title": "Modal Local Bridge Materializer"}, } if component.mapped_boundary_source_node_id is not None: mapped_node_id_set = set(component.mapped_node_ids) for phase_payload in phase_payloads: phase_node_ids = { str(node_id) for node_id in phase_payload.get("component_node_ids", []) } if not (phase_node_ids & mapped_node_id_set): continue register_modal_map_input_warmup_context( component.mapped_boundary_source_node_id, phase_payload, str(component.mapped_boundary_input_io_type or "*"), ) break for node_id, prompt_node in list(rewritten_prompt.items()): if node_id in component_proxy_node_ids: continue for input_name, input_value in list( (prompt_node.get("inputs") or {}).items() ): if not _is_link(input_value): continue source = LinkedOutputRef( node_id=str(input_value[0]), output_index=int(input_value[1]) ) replacement_output = materializer_output_by_consumer_and_source.get( (node_id, source) ) if replacement_output is None: replacement_output = proxy_output_by_consumer_and_source.get( (node_id, source) ) if replacement_output is None: replacement_output = replacement_output_indices.get(source) if replacement_output is not None: prompt_node["inputs"][input_name] = list(replacement_output) logger.info( "Rewrote remote component %s into ordered proxies %s.", component.representative_node_id, phase_proxy_node_ids, ) return phase_proxy_node_ids if isinstance(split_proxy_payloads, dict): static_payload = dict(split_proxy_payloads["static"]) mapped_payload = dict(split_proxy_payloads["mapped"]) static_proxy_node_id = str(static_payload["component_id"]) mapped_proxy_node_id = str(mapped_payload["component_id"]) while ( mapped_proxy_node_id in rewritten_prompt and mapped_proxy_node_id not in component_node_id_set ): mapped_proxy_node_id = f"{mapped_proxy_node_id}_proxy" mapped_payload["component_id"] = mapped_proxy_node_id static_proxy_inputs = _proxy_inputs_for_boundary_inputs( list(static_payload["boundary_inputs"]), rewritten_prompt=rewritten_prompt, ) static_boundary_outputs = list(static_payload["boundary_outputs"]) static_proxy_meta = copy.deepcopy( rewritten_prompt[static_proxy_node_id].get("_meta", {}) ) mapped_proxy_meta = copy.deepcopy( rewritten_prompt[component.mapped_node_ids[0]].get("_meta", {}) ) bridge_output_indices = { str(boundary_output["proxy_output_name"]): output_index for output_index, boundary_output in enumerate(static_boundary_outputs) if bool(boundary_output.get("session_output")) } mapped_proxy_inputs = _proxy_inputs_for_boundary_inputs( list(mapped_payload["boundary_inputs"]), rewritten_prompt=rewritten_prompt, ) mapped_static_phase = mapped_payload.get("static_phase") if isinstance(mapped_static_phase, dict): for input_name, input_value in _proxy_inputs_for_boundary_inputs( list(mapped_static_phase.get("boundary_inputs", [])), rewritten_prompt=rewritten_prompt, ).items(): mapped_proxy_inputs.setdefault(input_name, input_value) for boundary_input in mapped_payload.get("boundary_inputs", []): proxy_input_name = str(boundary_input["proxy_input_name"]) if proxy_input_name not in bridge_output_indices: continue mapped_proxy_inputs[proxy_input_name] = [ static_proxy_node_id, bridge_output_indices[proxy_input_name], ] replacement_output_indices = { LinkedOutputRef( node_id=str(boundary_output["node_id"]), output_index=int(boundary_output["output_index"]), ): [static_proxy_node_id, output_index] for output_index, boundary_output in enumerate(static_boundary_outputs) if not bool(boundary_output.get("session_output")) } replacement_output_indices.update( { LinkedOutputRef( node_id=str(boundary_output["node_id"]), output_index=int(boundary_output["output_index"]), ): [mapped_proxy_node_id, output_index] for output_index, boundary_output in enumerate( mapped_payload.get("boundary_outputs", []) ) } ) for node_id in component.node_ids: rewritten_prompt.pop(node_id, None) _register_proxy_node( prompt_node_id=static_proxy_node_id, payload_mapping=static_payload, proxy_inputs=static_proxy_inputs, meta=static_proxy_meta, is_output_node=_contains_output_node( component.static_node_ids, rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ), rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ) _register_proxy_node( prompt_node_id=mapped_proxy_node_id, payload_mapping=mapped_payload, proxy_inputs=mapped_proxy_inputs, meta=mapped_proxy_meta, is_output_node=_contains_output_node( component.mapped_node_ids, rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ), rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ) if component.mapped_boundary_source_node_id is not None: register_modal_map_input_warmup_context( component.mapped_boundary_source_node_id, mapped_payload, str(component.mapped_boundary_input_io_type or "*"), ) for node_id, prompt_node in list(rewritten_prompt.items()): if node_id in {static_proxy_node_id, mapped_proxy_node_id}: continue for input_name, input_value in list( (prompt_node.get("inputs") or {}).items() ): if not _is_link(input_value): continue source = LinkedOutputRef( node_id=str(input_value[0]), output_index=int(input_value[1]) ) if source in replacement_output_indices: prompt_node["inputs"][input_name] = list( replacement_output_indices[source] ) logger.info( "Rewrote hybrid remote component %s into static proxy %s and mapped proxy %s.", component.representative_node_id, static_proxy_node_id, mapped_proxy_node_id, ) return [static_proxy_node_id, mapped_proxy_node_id] boundary_outputs = list(payload.get("boundary_outputs", [])) proxy_node_id = ensure_modal_component_proxy_node_registered( output_types=tuple(str(output["io_type"]) for output in boundary_outputs), output_names=tuple( str(output["proxy_output_name"]) for output in boundary_outputs ), output_is_list=tuple( _proxy_boundary_output_is_list(output) for output in boundary_outputs ), nodes_module=nodes_module, is_output_node=component.contains_output_node, include_completion_output=True, ) representative_node_id = component.representative_node_id proxy_inputs = _proxy_inputs_for_boundary_inputs( list(payload.get("boundary_inputs", [])), rewritten_prompt=rewritten_prompt, ) proxy_inputs["original_node_data"] = register_cache_friendly_proxy_payload( representative_node_id, payload, ) representative_meta = copy.deepcopy( rewritten_prompt[representative_node_id].get("_meta", {}) ) rewritten_prompt[representative_node_id] = { "class_type": proxy_node_id, "inputs": proxy_inputs, "_meta": representative_meta, } if component.mapped_boundary_source_node_id is not None: register_modal_map_input_warmup_context( component.mapped_boundary_source_node_id, payload, str(component.mapped_boundary_input_io_type or "*"), ) materializer_output_by_consumer_and_source: dict[ tuple[str, LinkedOutputRef], list[Any] ] = {} proxy_output_by_consumer_and_source: dict[ tuple[str, LinkedOutputRef], list[Any] ] = {} default_proxy_output_by_source: dict[LinkedOutputRef, list[Any]] = {} if any( spec.local_materializer_node_id is not None for spec in component.boundary_outputs ): ensure_modal_local_bridge_materializer_registered(nodes_module) for output_index, spec in enumerate(component.boundary_outputs): proxy_output = [representative_node_id, output_index] default_proxy_output_by_source.setdefault(spec.source, proxy_output) for consumer_node_id in spec.session_consumer_node_ids: proxy_output_by_consumer_and_source[ (consumer_node_id, spec.source) ] = proxy_output materializer_node_id = spec.local_materializer_node_id if materializer_node_id is None: continue rewritten_prompt[materializer_node_id] = { "class_type": MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID, "inputs": {"bridge_ref": proxy_output}, "_meta": {"title": "Modal Local Bridge Materializer"}, } for consumer_node_id in spec.local_materializer_consumer_node_ids: materializer_output_by_consumer_and_source[ (consumer_node_id, spec.source) ] = [materializer_node_id, 0] for node_id, prompt_node in list(rewritten_prompt.items()): if node_id in component_node_id_set and node_id != representative_node_id: del rewritten_prompt[node_id] continue if node_id == representative_node_id: continue for input_name, input_value in list((prompt_node.get("inputs") or {}).items()): if not _is_link(input_value): continue source = LinkedOutputRef( node_id=str(input_value[0]), output_index=int(input_value[1]) ) replacement_output = materializer_output_by_consumer_and_source.get( (node_id, source) ) if replacement_output is None: replacement_output = proxy_output_by_consumer_and_source.get( (node_id, source) ) if replacement_output is None: replacement_output = default_proxy_output_by_source.get(source) if replacement_output is not None: prompt_node["inputs"][input_name] = list(replacement_output) logger.info( "Rewrote remote component %s with %d nodes to Modal proxy %s.", representative_node_id, len(component.node_ids), proxy_node_id, ) return [representative_node_id] def _modal_component_completion_output_index( *, proxy_node_id: str, rewritten_prompt: dict[str, Any], nodes_module: Any, ) -> int: """Return the synthetic completion-token output index for one Modal proxy.""" prompt_node = rewritten_prompt.get(proxy_node_id) if prompt_node is None: raise ModalPromptValidationError( f"Modal proxy node {proxy_node_id!r} is missing before artifact finalization." ) class_type = str(prompt_node.get("class_type")) node_class = nodes_module.NODE_CLASS_MAPPINGS.get(class_type) if node_class is None: raise ModalPromptValidationError( f"Modal proxy class {class_type!r} is not registered before artifact finalization." ) _output_types, output_names, _output_is_list = _normalize_output_metadata( node_class ) try: return output_names.index(MODAL_COMPONENT_COMPLETION_OUTPUT_NAME) except ValueError as exc: raise ModalPromptValidationError( f"Modal proxy {proxy_node_id!r} does not expose its completion token." ) from exc def _attach_modal_artifact_finalizer( *, rewritten_prompt: dict[str, Any], remote_component_ids: list[str], nodes_module: Any, ) -> str: """Attach one internal output sink to every rewritten Modal component.""" if not remote_component_ids: raise ModalPromptValidationError( "Modal artifact finalization requires at least one remote component." ) if len(remote_component_ids) > MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS: raise ModalPromptValidationError( "Modal artifact finalization supports at most " f"{MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS} remote components per prompt." ) ensure_modal_artifact_finalizer_registered(nodes_module) finalizer_node_id = f"__{MODAL_ARTIFACT_FINALIZER_NODE_ID}__" while finalizer_node_id in rewritten_prompt: finalizer_node_id = f"{finalizer_node_id}_proxy" finalizer_inputs = { f"components.component_{component_index}": [ proxy_node_id, _modal_component_completion_output_index( proxy_node_id=proxy_node_id, rewritten_prompt=rewritten_prompt, nodes_module=nodes_module, ), ] for component_index, proxy_node_id in enumerate(remote_component_ids) } rewritten_prompt[finalizer_node_id] = { "class_type": MODAL_ARTIFACT_FINALIZER_NODE_ID, "inputs": finalizer_inputs, "_meta": {"title": "Modal Artifact Finalizer"}, } logger.info( "Attached Modal artifact finalizer %s to remote components %s.", finalizer_node_id, remote_component_ids, ) return finalizer_node_id def _nearest_downstream_remote_component_ids( *, rewritten_prompt: dict[str, Any], consumers: dict[LinkedOutputRef, list[InputTarget]], seed_targets: list[InputTarget], remote_component_id_set: set[str], nodes_module: Any, ) -> list[str]: """Return the first remote proxies reached along every downstream branch.""" nearest_remote_component_ids: set[str] = set() visited_node_ids: set[str] = set() pending_node_ids = deque(str(target.node_id) for target in seed_targets) while pending_node_ids: node_id = pending_node_ids.popleft() if node_id in visited_node_ids or node_id not in rewritten_prompt: continue visited_node_ids.add(node_id) if node_id in remote_component_id_set: nearest_remote_component_ids.add(node_id) continue for output_ref in _node_output_refs( rewritten_prompt, node_id, nodes_module, ): for downstream_target in consumers.get(output_ref, []): pending_node_ids.append(str(downstream_target.node_id)) return sorted(nearest_remote_component_ids) def _parallelize_non_returning_local_branches( *, rewritten_prompt: dict[str, Any], remote_component_ids: list[str], nodes_module: Any, ) -> list[str]: """Release local-only taps once the nearest remote continuations are in flight.""" remote_component_id_set = set(remote_component_ids) if len(remote_component_id_set) < 2: return [] consumers = _build_consumer_map(rewritten_prompt) ensure_modal_parallel_local_passthrough_registered(nodes_module) passthrough_node_ids: list[str] = [] dispatch_group_id = uuid.uuid4().hex for component_id in remote_component_ids: embedded_payload = ( rewritten_prompt.get(component_id, {}) .get("inputs", {}) .get("original_node_data") ) if not isinstance(embedded_payload, Mapping): continue payload_prompt_id = embedded_payload.get("prompt_id") if payload_prompt_id is not None and str(payload_prompt_id).strip(): dispatch_group_id = str(payload_prompt_id) break for source_component_id in remote_component_ids: for source in _node_output_refs( rewritten_prompt, source_component_id, nodes_module, ): output_consumers = list(consumers.get(source, [])) if not output_consumers: continue downstream_remote_component_ids = _nearest_downstream_remote_component_ids( rewritten_prompt=rewritten_prompt, consumers=consumers, seed_targets=output_consumers, remote_component_id_set=remote_component_id_set, nodes_module=nodes_module, ) downstream_remote_component_ids = [ component_id for component_id in downstream_remote_component_ids if component_id != source_component_id ] if not downstream_remote_component_ids: continue local_only_consumers: list[InputTarget] = [] for target in output_consumers: if target.node_id in remote_component_id_set: continue branch_node_ids = _downstream_node_ids_from_targets( prompt=rewritten_prompt, consumers=consumers, seed_targets=[target], nodes_module=nodes_module, ) if branch_node_ids & remote_component_id_set: continue local_only_consumers.append(target) if not local_only_consumers: continue passthrough_node_id = ( f"__{MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID}__" f"{len(passthrough_node_ids)}" ) while passthrough_node_id in rewritten_prompt: passthrough_node_id = f"{passthrough_node_id}_proxy" passthrough_inputs: dict[str, Any] = { "value": [source.node_id, source.output_index], "dispatch_context": { "dispatch_group_id": dispatch_group_id, "component_ids": downstream_remote_component_ids, }, } rewritten_prompt[passthrough_node_id] = { "class_type": MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, "inputs": passthrough_inputs, "_meta": {"title": "Modal Parallel Local Branch"}, } for target in local_only_consumers: rewritten_prompt[target.node_id]["inputs"][target.input_name] = [ passthrough_node_id, 0, ] for component_id in downstream_remote_component_ids: prompt_inputs = rewritten_prompt[component_id].get("inputs", {}) embedded_payload = prompt_inputs.get("original_node_data") if not isinstance(embedded_payload, Mapping): continue prompt_inputs[ "original_node_data" ] = update_registered_proxy_payload_fields( component_id, embedded_payload, { "parallel_local_dispatch_group_id": dispatch_group_id, "signal_parallel_local_dispatch": True, }, ) passthrough_node_ids.append(passthrough_node_id) logger.info( "Parallelized local-only consumers %s of remote output %s:%d after downstream Modal dispatches %s via passthrough=%s.", [target.node_id for target in local_only_consumers], source.node_id, source.output_index, downstream_remote_component_ids, passthrough_node_id, ) return passthrough_node_ids def _configure_local_gap_keepalive_payloads( *, rewritten_prompt: dict[str, Any], remote_component_ids: list[str], sandwiched_local_node_ids: set[str], ) -> None: """Mark only remote proxies that precede a local-to-remote execution gap.""" if not sandwiched_local_node_ids: return payloads_by_component_id: dict[str, Mapping[str, Any]] = {} for component_id in remote_component_ids: payload = ( rewritten_prompt.get(component_id, {}) .get("inputs", {}) .get("original_node_data") ) if not isinstance(payload, dict): continue payloads_by_component_id[component_id] = payload keepalive_component_ids: set[str] = set() continuation_component_ids: set[str] = set() for local_node_id in sandwiched_local_node_ids: upstream_component_ids: set[str] = set() downstream_component_ids: set[str] = set() for component_id in remote_component_ids: if _component_ancestors_of_local_source( prompt=rewritten_prompt, source_node_id=local_node_id, component_node_id_set={component_id}, ): upstream_component_ids.add(component_id) if _component_ancestors_of_local_source( prompt=rewritten_prompt, source_node_id=component_id, component_node_id_set={local_node_id}, ): downstream_component_ids.add(component_id) for upstream_component_id in upstream_component_ids: upstream_payload = payloads_by_component_id.get(upstream_component_id) if upstream_payload is None: continue upstream_environment_id = str( upstream_payload.get("execution_environment_id") or "" ) if upstream_payload.get("execution_provider") != ExecutionProvider.MODAL.value: continue for downstream_component_id in downstream_component_ids: downstream_payload = payloads_by_component_id.get( downstream_component_id ) if ( downstream_payload is None or downstream_payload.get("execution_provider") != ExecutionProvider.MODAL.value or str(downstream_payload.get("execution_environment_id") or "") != upstream_environment_id ): continue keepalive_component_ids.add(upstream_component_id) continuation_component_ids.add(downstream_component_id) for component_id, embedded_payload in payloads_by_component_id.items(): if component_id not in keepalive_component_ids | continuation_component_ids: continue payload_fields: dict[str, Any] = {"remote_local_gap_pool": True} if component_id in keepalive_component_ids: payload_fields["keepalive_after_remote_component"] = True if component_id in continuation_component_ids: payload_fields["stop_local_gap_keepalive_before_remote_component"] = True rewritten_prompt[component_id]["inputs"][ "original_node_data" ] = update_registered_proxy_payload_fields( component_id, embedded_payload, payload_fields, ) logger.info( "Configured local-gap Modal pool for components=%s keepalive_producers=%s continuations=%s.", sorted(keepalive_component_ids | continuation_component_ids), sorted(keepalive_component_ids), sorted(continuation_component_ids), ) def _remote_proxy_payload( rewritten_prompt: Mapping[str, Any], component_id: str ) -> Mapping[str, Any] | None: """Return the registered execution payload embedded in one remote proxy.""" prompt_node = rewritten_prompt.get(component_id) if not isinstance(prompt_node, Mapping): return None inputs = prompt_node.get("inputs") if not isinstance(inputs, Mapping): return None payload = inputs.get("original_node_data") if not isinstance(payload, Mapping): return None return registered_proxy_execution_payload(component_id, payload) def _remote_proxy_dependency_edges( rewritten_prompt: Mapping[str, Any], component_ids: set[str], ) -> dict[str, set[str]]: """Return nearest remote dependencies while traversing intervening local nodes.""" dependency_edges = {component_id: set() for component_id in component_ids} for downstream_component_id in component_ids: downstream_node = rewritten_prompt.get(downstream_component_id) if not isinstance(downstream_node, Mapping): continue pending_node_ids = [ str(input_value[0]) for input_value in (downstream_node.get("inputs") or {}).values() if _is_link(input_value) ] visited_node_ids: set[str] = set() while pending_node_ids: upstream_node_id = pending_node_ids.pop() if upstream_node_id in visited_node_ids: continue visited_node_ids.add(upstream_node_id) if upstream_node_id in component_ids: dependency_edges[upstream_node_id].add(downstream_component_id) continue upstream_node = rewritten_prompt.get(upstream_node_id) if not isinstance(upstream_node, Mapping): continue pending_node_ids.extend( str(input_value[0]) for input_value in (upstream_node.get("inputs") or {}).values() if _is_link(input_value) ) return dependency_edges def _component_descendant_distances( component_id: str, dependency_edges: Mapping[str, set[str]], ) -> dict[str, int]: """Return shortest downstream distances from one remote component.""" distances: dict[str, int] = {} pending = deque( (descendant_id, 1) for descendant_id in dependency_edges.get(component_id, set()) ) while pending: descendant_id, distance = pending.popleft() previous_distance = distances.get(descendant_id) if previous_distance is not None and previous_distance <= distance: continue distances[descendant_id] = distance pending.extend( (nested_descendant_id, distance + 1) for nested_descendant_id in dependency_edges.get(descendant_id, set()) ) return distances def _speculative_prewarm_target_payload( payload: Mapping[str, Any], ) -> dict[str, Any]: """Return the minimal side-effect-free payload needed to prepare one worker.""" return { field_name: copy.deepcopy(payload[field_name]) for field_name in _SPECULATIVE_PREWARM_PAYLOAD_FIELDS if field_name in payload } def _configure_speculative_affinity_prewarm_payloads( *, rewritten_prompt: dict[str, Any], execution_stages: list[list[str]], ) -> None: """Attach one next-affinity worker preparation target to each eligible proxy.""" stage_index_by_component_id = { component_id: stage_index for stage_index, stage in enumerate(execution_stages) for component_id in stage } payloads_by_component_id = { component_id: payload for component_id in stage_index_by_component_id if (payload := _remote_proxy_payload(rewritten_prompt, component_id)) is not None } dependency_edges = _remote_proxy_dependency_edges( rewritten_prompt, set(payloads_by_component_id), ) configured_targets: dict[str, str] = {} for component_id, payload in payloads_by_component_id.items(): if payload.get("execution_provider") != ExecutionProvider.MODAL.value: continue execution_environment_id = str( payload.get("execution_environment_id") or "" ) affinity_group = ( str(payload.get("remote_worker_affinity_group") or "comfy").strip().lower() ) descendant_distances = _component_descendant_distances( component_id, dependency_edges ) candidate_ids = sorted( ( descendant_id for descendant_id in descendant_distances if descendant_id in payloads_by_component_id and payloads_by_component_id[descendant_id].get( "execution_provider" ) == ExecutionProvider.MODAL.value and str( payloads_by_component_id[descendant_id].get( "execution_environment_id" ) or "" ) == execution_environment_id and str( payloads_by_component_id[descendant_id].get( "remote_worker_affinity_group" ) or "comfy" ) .strip() .lower() != affinity_group ), key=lambda descendant_id: ( descendant_distances[descendant_id], stage_index_by_component_id[descendant_id], descendant_id, ), ) if not candidate_ids: continue target_component_id = candidate_ids[0] target_payload = _speculative_prewarm_target_payload( payloads_by_component_id[target_component_id] ) rewritten_prompt[component_id]["inputs"][ "original_node_data" ] = update_registered_proxy_payload_fields( component_id, payload, {_SPECULATIVE_PREWARM_TARGET_KEY: target_payload}, ) configured_targets[component_id] = target_component_id if configured_targets: logger.info( "Configured one-step speculative Modal worker prewarm targets=%s.", configured_targets, )