"""Mapped and implicit-batch remote execution orchestration.""" from __future__ import annotations import asyncio import copy from dataclasses import dataclass import logging import os import threading from typing import Any, Awaitable, Callable, Mapping from ..serialization import ( deserialize_node_inputs, deserialize_node_outputs, is_mapped_output_value, join_mapped_values_for_scheduler, serialize_node_inputs, serialize_node_outputs, split_mapped_value, ) from ..session_state import ( is_remote_session_bridge_ref_payload, is_remote_session_handle_payload, is_remote_session_value_ref_payload, ) from .host_session_bridge import _payload_remote_session_handle from .local_execution import ( RemoteSubgraphExecutionError, _execute_subgraph_prompt, _load_nodes_module, _node_input_type_map, _resolve_required_subgraph_nodes, ) from .local_ui_events import _emit_local_modal_progress from .modal_deployment import ( ModalRemoteInvocationError, _lookup_deployed_remote_engine, _mapped_lane_affinity_key, ) from .modal_warmup import ( _await_prompt_warmup_slots, _prompt_warmup_head_start_seconds, boost_mapped_component_warmup, ) from ..settings import get_settings logger = logging.getLogger(__name__) try: import modal # type: ignore except ModuleNotFoundError: # pragma: no cover - local fallback environments. modal = None _IMPLICIT_BATCH_PRESERVING_TARGETS = frozenset({("CreateVideo", "images")}) _MAPPED_PROGRESS_NODE_IDS_LOCK = threading.Lock() _MAPPED_PROGRESS_NODE_IDS: dict[tuple[str, str, str], str] = {} @dataclass(frozen=True) class MappedExecutionHooks: """Callbacks supplied by the host orchestrator for invocation and interrupts.""" invoke_remote_engine_async: Callable[..., Awaitable[bytes]] invoke_bound_remote_engine_async: Callable[[Any, dict[str, Any], bytes], Awaitable[bytes]] local_processing_interrupted: Callable[[], bool] raise_local_interrupt: Callable[[], None] exception_indicates_interruption: Callable[[BaseException], bool] _MAPPED_EXECUTION_HOOKS: MappedExecutionHooks | None = None def configure_mapped_execution_hooks(hooks: MappedExecutionHooks) -> None: """Install host callbacks without importing upward into the orchestrator.""" global _MAPPED_EXECUTION_HOOKS _MAPPED_EXECUTION_HOOKS = hooks def _mapped_execution_hooks() -> MappedExecutionHooks: """Return configured host callbacks or fail with a clear import-order error.""" if _MAPPED_EXECUTION_HOOKS is None: raise RuntimeError("Mapped execution hooks have not been configured.") return _MAPPED_EXECUTION_HOOKS async def invoke_remote_engine_async(*args: Any, **kwargs: Any) -> bytes: """Delegate remote invocation through the injected host callback.""" return await _mapped_execution_hooks().invoke_remote_engine_async(*args, **kwargs) async def _invoke_bound_remote_engine_async( remote_engine: Any, payload: dict[str, Any], kwargs_payload: bytes ) -> bytes: """Delegate affinity-bound invocation through the injected host callback.""" return await _mapped_execution_hooks().invoke_bound_remote_engine_async( remote_engine, payload, kwargs_payload ) def _local_processing_interrupted() -> bool: """Return local interrupt state through the injected host callback.""" return _mapped_execution_hooks().local_processing_interrupted() def _raise_local_interrupt() -> None: """Raise the local interrupt through the injected host callback.""" _mapped_execution_hooks().raise_local_interrupt() def _exception_indicates_interruption(exc: BaseException) -> bool: """Classify interruption through the injected host callback.""" return _mapped_execution_hooks().exception_indicates_interruption(exc) def _mapped_execution_parallelism(item_count: int) -> int: """Return the local worker width used to schedule mapped Modal item executions.""" settings = get_settings() configured_limit = settings.max_inflight_calls if settings.max_containers is not None: configured_limit = min(configured_limit, settings.max_containers) return max(1, min(item_count, configured_limit)) def _mapped_phase_definition( payload: dict[str, Any], phase_key: str ) -> dict[str, Any] | None: """Return one explicit mapped phase definition when queue-time planning provided it.""" phase_payload = payload.get(phase_key) if isinstance(phase_payload, dict): return phase_payload return None def _shared_subgraph_payload_fields(payload: dict[str, Any]) -> dict[str, Any]: """Return the payload fields shared by every explicit mapped phase.""" shared_fields = { "prompt_id": payload.get("prompt_id"), "extra_data": copy.deepcopy(payload.get("extra_data") or {}), "requires_volume_reload": bool(payload.get("requires_volume_reload", True)), "volume_reload_marker": payload.get("volume_reload_marker"), "uploaded_volume_paths": list(payload.get("uploaded_volume_paths", [])), "terminate_container_on_error": bool( payload.get("terminate_container_on_error", True) ), "custom_nodes_bundle": payload.get("custom_nodes_bundle"), } snapshot_profile_key = payload.get("snapshot_profile_key") if isinstance(snapshot_profile_key, str) and snapshot_profile_key.strip(): shared_fields["snapshot_profile_key"] = snapshot_profile_key.strip() remote_session = payload.get("remote_session") if remote_session is not None: shared_fields["remote_session"] = copy.deepcopy(remote_session) if bool(payload.get("clear_remote_session")): shared_fields["clear_remote_session"] = True return shared_fields def _build_phase_subgraph_payload( payload: dict[str, Any], phase_key: str, component_id: str, *, suppress_status_stream: bool = False, lane_id: str | None = None, item_index: int | None = None, ) -> dict[str, Any]: """Return one explicit static or mapped subgraph payload.""" phase_definition = _mapped_phase_definition(payload, phase_key) if phase_definition is None: raise KeyError(f"Mapped payload is missing phase definition {phase_key!r}.") phase_payload = { "payload_kind": "subgraph", "component_id": component_id, **_shared_subgraph_payload_fields(payload), "component_node_ids": [ str(node_id) for node_id in phase_definition.get("component_node_ids", []) if str(node_id) ], "subgraph_prompt": copy.deepcopy(phase_definition.get("subgraph_prompt", {})), "boundary_inputs": copy.deepcopy(phase_definition.get("boundary_inputs", [])), "boundary_outputs": copy.deepcopy(phase_definition.get("boundary_outputs", [])), "execute_node_ids": [ str(node_id) for node_id in phase_definition.get("execute_node_ids", []) if str(node_id) ], } if suppress_status_stream: phase_payload["suppress_status_stream"] = True if lane_id is not None: phase_payload["mapped_progress_lane_id"] = str(lane_id) phase_payload["mapped_progress_display_node_id"] = str( payload.get("component_id", "modal-subgraph") ) if item_index is not None: phase_payload["map_item_index"] = int(item_index) return phase_payload def _split_phase_outputs( phase_outputs: tuple[Any, ...], boundary_outputs: list[dict[str, Any]], internal_output_names: set[str], ) -> tuple[dict[str, Any], tuple[Any, ...]]: """Split one phase result tuple into bridge values and external outputs.""" internal_outputs: dict[str, Any] = {} external_outputs: list[Any] = [] for boundary_output, output_value in zip( boundary_outputs, phase_outputs, strict=True ): output_name = str(boundary_output.get("proxy_output_name") or "") if output_name in internal_output_names: internal_outputs[output_name] = output_value continue external_outputs.append(output_value) return internal_outputs, tuple(external_outputs) def _execute_mapped_subgraph_payload( payload: dict[str, Any], hydrated_inputs: dict[str, Any], node_mapping: dict[str, type[Any]] | None = None, ) -> tuple[Any, ...]: """Execute one mapped payload locally using explicit static and mapped phases.""" mapped_input = payload.get("mapped_input") or {} mapped_input_name = str(mapped_input.get("proxy_input_name") or "") if not mapped_input_name: raise ModalRemoteInvocationError( "Mapped remote payloads must define mapped_input.proxy_input_name." ) if mapped_input_name not in hydrated_inputs: raise KeyError( f"Mapped remote payload input {mapped_input_name!r} was not provided." ) mapped_items = split_mapped_value( hydrated_inputs[mapped_input_name], str(mapped_input.get("io_type", "*")), ) if not mapped_items: raise ValueError("Mapped remote execution requires at least one input item.") broadcast_inputs = dict(hydrated_inputs) broadcast_inputs.pop(mapped_input_name, None) static_to_mapped_boundaries = list(payload.get("static_to_mapped_boundaries", [])) bridge_output_names = { str(boundary_spec.get("proxy_name") or "") for boundary_spec in static_to_mapped_boundaries if str(boundary_spec.get("proxy_name") or "") } static_outputs: tuple[Any, ...] = () static_phase_payload: dict[str, Any] | None = None if payload.get("static_phase") is not None: static_phase_payload = _build_phase_subgraph_payload( payload, "static_phase", f"{payload.get('component_id', 'modal-subgraph')}::static", suppress_status_stream=True, ) static_phase_payload.pop("clear_remote_session", None) elif payload.get("static_execute_node_ids"): static_phase_payload = _build_static_mapped_payload(payload) if static_phase_payload is not None: if static_phase_payload.get("execute_node_ids"): logger.info( "Executing static mapped phase for component=%s with execute nodes=%s.", payload.get("component_id"), static_phase_payload.get("execute_node_ids", []), ) static_phase_outputs = _execute_subgraph_prompt( static_phase_payload, dict(broadcast_inputs), node_mapping, ) bridge_inputs, static_outputs = _split_phase_outputs( static_phase_outputs, list(static_phase_payload.get("boundary_outputs", [])), bridge_output_names, ) broadcast_inputs.update(bridge_inputs) total_items = len(mapped_items) _emit_local_mapped_progress(payload, 0, total_items) per_item_outputs: list[tuple[Any, ...]] = [] for item_index, item_value in enumerate(mapped_items): if _local_processing_interrupted(): _raise_local_interrupt() if payload.get("mapped_phase") is not None: item_payload = _build_phase_subgraph_payload( payload, "mapped_phase", f"{payload.get('component_id', 'modal-subgraph')}::item:{item_index}", suppress_status_stream=True, lane_id="0", item_index=item_index, ) else: item_payload = _build_mapped_item_payload(payload, item_index, 0) item_inputs = dict(broadcast_inputs) item_inputs[mapped_input_name] = item_value logger.info( "Executing mapped item %d/%d for component=%s with execute nodes=%s.", item_index + 1, total_items, payload.get("component_id"), item_payload.get("execute_node_ids", []), ) per_item_outputs.append( _execute_subgraph_prompt( item_payload, item_inputs, node_mapping, ) ) _emit_local_mapped_progress(payload, item_index + 1, total_items) if payload.get("mapped_phase") is not None: mapped_phase_payload = _build_phase_subgraph_payload( payload, "mapped_phase", f"{payload.get('component_id', 'modal-subgraph')}::mapped", ) else: mapped_phase_payload = { **payload, "boundary_outputs": [ boundary_output for boundary_output in payload.get("boundary_outputs", []) if _is_mapped_boundary_output(boundary_output, payload) ], } mapped_outputs = _aggregate_mapped_outputs( per_item_outputs, { **payload, "boundary_outputs": list(mapped_phase_payload.get("boundary_outputs", [])), }, ) return _merge_static_and_mapped_outputs( static_outputs=static_outputs, mapped_outputs=mapped_outputs, payload=payload, ) def _build_mapped_item_payload( payload: dict[str, Any], item_index: int, lane_index: int, ) -> dict[str, Any]: """Return one per-item subgraph payload derived from a mapped remote component payload.""" if _mapped_phase_definition(payload, "mapped_phase") is not None: return _build_phase_subgraph_payload( payload, "mapped_phase", f"{payload.get('component_id', 'modal-subgraph')}::item:{item_index}", suppress_status_stream=True, lane_id=str(lane_index), item_index=item_index, ) item_payload = copy.deepcopy(payload) item_payload["payload_kind"] = "subgraph" item_payload[ "component_id" ] = f"{payload.get('component_id', 'modal-subgraph')}::item:{item_index}" item_payload["mapped_input"] = None item_payload["suppress_status_stream"] = True item_payload["map_item_index"] = item_index item_payload["mapped_progress_lane_id"] = str(lane_index) item_payload["mapped_progress_display_node_id"] = str( payload.get("component_id", "modal-subgraph") ) item_payload.pop("clear_remote_session", None) item_payload["execute_node_ids"] = list( payload.get("mapped_execute_node_ids") or payload.get("execute_node_ids", []) ) item_payload["boundary_outputs"] = [ copy.deepcopy(boundary_output) for boundary_output in payload.get("boundary_outputs", []) if _is_mapped_boundary_output(boundary_output, payload) ] return item_payload def _aggregate_mapped_outputs( per_item_outputs: list[tuple[Any, ...]], payload: dict[str, Any], ) -> tuple[Any, ...]: """Reassemble ordered per-item outputs from mapped execution into one proxy result tuple.""" if not per_item_outputs: raise ValueError("Mapped execution produced no per-item outputs to aggregate.") output_count = len(per_item_outputs[0]) if any(len(item_outputs) != output_count for item_outputs in per_item_outputs): raise RemoteSubgraphExecutionError( "Mapped remote execution produced inconsistent output arity." ) boundary_outputs = list(payload.get("boundary_outputs", [])) aggregated_outputs: list[Any] = [] for output_index in range(output_count): boundary_output = ( boundary_outputs[output_index] if output_index < len(boundary_outputs) else {} ) aggregated_outputs.append( join_mapped_values_for_scheduler( [item_outputs[output_index] for item_outputs in per_item_outputs], io_type=str(boundary_output.get("io_type", "*")), is_list=bool(boundary_output.get("is_list", False)), scheduler_is_list=bool(boundary_output.get("scheduler_is_list", False)), ) ) return tuple(aggregated_outputs) def _build_remote_session_cleanup_payload( payload: dict[str, Any] ) -> dict[str, Any] | None: """Return a dedicated cleanup payload that clears one shared remote session once.""" remote_session = payload.get("remote_session") if not bool( payload.get("clear_remote_session") ) or not is_remote_session_handle_payload(remote_session): return None return { "payload_kind": "subgraph", "component_id": f"{payload.get('component_id', 'modal-subgraph')}::cleanup", **_shared_subgraph_payload_fields(payload), "component_node_ids": [], "subgraph_prompt": {}, "boundary_inputs": [], "boundary_outputs": [], "execute_node_ids": [], "remote_session": copy.deepcopy(remote_session), "clear_remote_session": True, "suppress_status_stream": True, "terminate_container_on_error": False, } def _is_mapped_boundary_output( boundary_output: dict[str, Any], payload: dict[str, Any] ) -> bool: """Return whether one boundary output belongs to the mapped per-item branch.""" mapped_output = boundary_output.get("mapped_output") if mapped_output is not None: return bool(mapped_output) return bool(payload.get("mapped_input")) and not bool( payload.get("static_execute_node_ids") ) def _build_static_mapped_payload(payload: dict[str, Any]) -> dict[str, Any]: """Return the one-time static subgraph payload for a hybrid mapped component.""" if _mapped_phase_definition(payload, "static_phase") is not None: return _build_phase_subgraph_payload( payload, "static_phase", f"{payload.get('component_id', 'modal-subgraph')}::static", suppress_status_stream=True, ) static_payload = copy.deepcopy(payload) static_payload["payload_kind"] = "subgraph" static_payload[ "component_id" ] = f"{payload.get('component_id', 'modal-subgraph')}::static" static_payload["mapped_input"] = None static_payload["suppress_status_stream"] = True static_payload.pop("clear_remote_session", None) static_payload["execute_node_ids"] = list( payload.get("static_execute_node_ids") or [] ) static_payload["boundary_outputs"] = [ copy.deepcopy(boundary_output) for boundary_output in payload.get("boundary_outputs", []) if not _is_mapped_boundary_output(boundary_output, payload) ] return static_payload def _merge_static_and_mapped_outputs( *, static_outputs: tuple[Any, ...], mapped_outputs: tuple[Any, ...], payload: dict[str, Any], ) -> tuple[Any, ...]: """Reassemble one hybrid mapped component's static and mapped outputs in original order.""" combined_outputs: list[Any] = [] static_output_index = 0 mapped_output_index = 0 for boundary_output in payload.get("boundary_outputs", []): if _is_mapped_boundary_output(boundary_output, payload): if mapped_output_index >= len(mapped_outputs): raise RemoteSubgraphExecutionError( "Mapped remote execution returned fewer mapped outputs than expected." ) combined_outputs.append(mapped_outputs[mapped_output_index]) mapped_output_index += 1 continue if static_output_index >= len(static_outputs): raise RemoteSubgraphExecutionError( "Mapped remote execution returned fewer static outputs than expected." ) combined_outputs.append(static_outputs[static_output_index]) static_output_index += 1 if static_output_index != len(static_outputs) or mapped_output_index != len( mapped_outputs ): raise RemoteSubgraphExecutionError( "Mapped remote execution produced extra outputs that did not match the declared boundary outputs." ) return tuple(combined_outputs) def _emit_local_mapped_progress( payload: dict[str, Any], completed_items: int, total_items: int, ) -> None: """Emit one aggregate mapped-execution progress update for the component representative node.""" prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) extra_data = payload.get("extra_data") or {} client_id = ( str(extra_data.get("client_id")) if extra_data.get("client_id") is not None else None ) display_node_id = str( payload.get("mapped_progress_display_node_id") or payload.get("component_id") or "" ) if not prompt_id or not client_id or not display_node_id: return _emit_local_modal_progress( prompt_id=prompt_id, client_id=client_id, node_id=display_node_id, value=float(completed_items), max_value=float(total_items), display_node_id=display_node_id, aggregate_only=True, ) def _mapped_progress_owner_component_id(payload: dict[str, Any]) -> str | None: """Return the stable component id used to track one mapped worker lane locally.""" owner_component_id = payload.get( "mapped_progress_display_node_id", payload.get("component_id") ) if owner_component_id is None: return None owner_component = str(owner_component_id) return owner_component or None def _remember_mapped_lane_node_id( payload: dict[str, Any], lane_id: str, node_id: str ) -> None: """Remember the last real node id that emitted progress for one mapped worker lane.""" prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) owner_component_id = _mapped_progress_owner_component_id(payload) if not prompt_id or not owner_component_id or not node_id: return with _MAPPED_PROGRESS_NODE_IDS_LOCK: _MAPPED_PROGRESS_NODE_IDS[(prompt_id, owner_component_id, lane_id)] = node_id def _pop_mapped_lane_node_id(payload: dict[str, Any], lane_id: str) -> str | None: """Forget and return the last real node id that emitted progress for one mapped worker lane.""" prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) owner_component_id = _mapped_progress_owner_component_id(payload) if not prompt_id or not owner_component_id: return None with _MAPPED_PROGRESS_NODE_IDS_LOCK: return _MAPPED_PROGRESS_NODE_IDS.pop( (prompt_id, owner_component_id, lane_id), None ) def _clear_local_mapped_lane_progress( payload: dict[str, Any], lane_index: int, item_index: int, ) -> None: """Remove one mapped worker lane from the local node overlay.""" prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) extra_data = payload.get("extra_data") or {} client_id = ( str(extra_data.get("client_id")) if extra_data.get("client_id") is not None else None ) lane_id = str(lane_index) display_node_id = _pop_mapped_lane_node_id(payload, lane_id) or str( payload.get("component_id") or "" ) if not prompt_id or not client_id or not display_node_id: return _emit_local_modal_progress( prompt_id=prompt_id, client_id=client_id, node_id=display_node_id, value=0.0, max_value=1.0, display_node_id=display_node_id, lane_id=lane_id, clear=True, item_index=item_index, ) def _emit_local_mapped_lane_progress_start( payload: dict[str, Any], lane_index: int, item_index: int | None = None, ) -> None: """Create or reset one mapped worker lane before remote progress begins arriving.""" prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) extra_data = payload.get("extra_data") or {} client_id = ( str(extra_data.get("client_id")) if extra_data.get("client_id") is not None else None ) display_node_id = str(payload.get("component_id") or "") if not prompt_id or not client_id or not display_node_id: return _emit_local_modal_progress( prompt_id=prompt_id, client_id=client_id, node_id=display_node_id, value=0.0, max_value=1.0, display_node_id=display_node_id, lane_id=str(lane_index), item_index=item_index, setup_only=True, ) def _implicit_batch_preserving_targets( payload: dict[str, Any], boundary_input: dict[str, Any], ) -> list[str]: """Return target sockets that must receive one complete tensor batch.""" prompt = payload.get("subgraph_prompt", {}) if not isinstance(prompt, dict): return [] preserving_targets: list[str] = [] for target in boundary_input.get("targets", []): node_id = str(target.get("node_id") or "") input_name = str(target.get("input_name") or "") prompt_node = prompt.get(node_id) if not node_id or not input_name or not isinstance(prompt_node, dict): continue class_type = str(prompt_node.get("class_type") or "") if (class_type, input_name) in _IMPLICIT_BATCH_PRESERVING_TARGETS: preserving_targets.append(f"{node_id}.{input_name}") return sorted(preserving_targets) def _split_batch_boundary_inputs( payload: dict[str, Any], hydrated_inputs: dict[str, Any], ) -> tuple[dict[str, list[Any]], int] | None: """Return zipped per-item boundary inputs when an ordinary subgraph receives batched values.""" implicitly_batchable_scalar_io_types = frozenset( {"BOOLEAN", "FLOAT", "INT", "STRING"} ) implicitly_batchable_transport_io_types = frozenset( {"IMAGE", "LATENT", "MASK", "NOISE", "SIGMAS"} ) split_inputs: dict[str, list[Any]] = {} def is_session_ref_list(value: Any) -> bool: """Return whether `value` is a non-empty list of remote session refs.""" return ( isinstance(value, list) and len(value) > 0 and all( is_remote_session_value_ref_payload(item) or is_remote_session_bridge_ref_payload(item) for item in value ) ) def unwrap_singleton_session_ref_list(value: Any) -> Any: """Unwrap Comfy's output-list wrapper around mapped remote session refs.""" if ( isinstance(value, list) and len(value) == 1 and is_session_ref_list(value[0]) ): return value[0] return value for boundary_input in payload.get("boundary_inputs", []): proxy_input_name = str(boundary_input.get("proxy_input_name") or "") if not proxy_input_name or proxy_input_name not in hydrated_inputs: continue input_value = unwrap_singleton_session_ref_list( hydrated_inputs[proxy_input_name] ) io_type = _implicit_batch_boundary_effective_io_type( payload=payload, boundary_input=boundary_input, input_value=input_value, ) preserving_targets = _implicit_batch_preserving_targets( payload, boundary_input, ) if preserving_targets: logger.info( "Skipping implicit batch split for boundary input %s io_type=%s " "because target sockets consume the complete tensor batch: %s.", proxy_input_name, io_type, preserving_targets, ) continue input_is_session_ref_list = is_session_ref_list(input_value) if ( isinstance(input_value, list) and not is_mapped_output_value(input_value) and not input_is_session_ref_list and io_type not in ( implicitly_batchable_scalar_io_types | implicitly_batchable_transport_io_types ) ): logger.info( "Skipping implicit batch split for boundary input %s io_type=%s because list-backed non-scalar values stay broadcast.", proxy_input_name, io_type, ) continue try: items = split_mapped_value( input_value, io_type, ) except (TypeError, ValueError): continue if len(items) <= 1: continue split_inputs[proxy_input_name] = items if not split_inputs: return None item_counts = {input_name: len(items) for input_name, items in split_inputs.items()} unique_counts = set(item_counts.values()) if len(unique_counts) != 1: raise ModalRemoteInvocationError( "Implicit Modal batch boundary inputs must all have the same item count. " f"Received counts: {item_counts!r}" ) return split_inputs, next(iter(unique_counts)) def _is_latent_like_mapping(value: Any) -> bool: """Return whether one runtime value looks like a ComfyUI LATENT mapping.""" return isinstance(value, Mapping) and "samples" in value def _list_is_latent_like_batch(value: Any) -> bool: """Return whether one runtime value is a list of LATENT-like mappings.""" return ( isinstance(value, list) and len(value) > 0 and all(_is_latent_like_mapping(item) for item in value) ) def _implicit_batch_boundary_target_input_types( payload: dict[str, Any], boundary_input: dict[str, Any], ) -> set[str]: """Return the declared input types of one boundary input's target sockets.""" prompt = payload.get("subgraph_prompt", {}) if not isinstance(prompt, dict): return set() try: resolved_node_mapping = _load_nodes_module().NODE_CLASS_MAPPINGS except ModuleNotFoundError: logger.debug( "Skipping target input type discovery for implicit batching because ComfyUI nodes are unavailable." ) return set() target_input_types: set[str] = set() for target in boundary_input.get("targets", []): node_id = str(target.get("node_id") or "") input_name = str(target.get("input_name") or "") if not node_id or not input_name: continue prompt_node = prompt.get(node_id) if not isinstance(prompt_node, dict): continue node_class = resolved_node_mapping.get(str(prompt_node.get("class_type"))) if node_class is None: continue prompt_inputs = prompt_node.get("inputs") or {} declared_type = _node_input_type_map(node_class, prompt_inputs).get(input_name) if isinstance(declared_type, str) and declared_type: target_input_types.add(declared_type) return target_input_types def _implicit_batch_boundary_effective_io_type( *, payload: dict[str, Any], boundary_input: dict[str, Any], input_value: Any, ) -> str: """Return the best effective io_type for implicit batching of one boundary input.""" declared_io_type = str(boundary_input.get("io_type", "*")) if declared_io_type != "*": return declared_io_type if _list_is_latent_like_batch(input_value): return "LATENT" target_input_types = _implicit_batch_boundary_target_input_types( payload, boundary_input ) if len(target_input_types) == 1: return next(iter(target_input_types)) return declared_io_type def _implicit_batch_input_is_list_target_node_ids( payload: dict[str, Any], split_inputs: dict[str, list[Any]], ) -> list[str]: """Return split-boundary target nodes that must consume the full list in one execution.""" prompt = payload.get("subgraph_prompt", {}) if not isinstance(prompt, dict): return [] try: resolved_node_mapping = _load_nodes_module().NODE_CLASS_MAPPINGS except ModuleNotFoundError: logger.debug( "Skipping INPUT_IS_LIST detection for implicit batching because ComfyUI nodes are unavailable." ) return [] target_node_ids: set[str] = set() for boundary_input in payload.get("boundary_inputs", []): proxy_input_name = str(boundary_input.get("proxy_input_name") or "") if proxy_input_name not in split_inputs: continue for target in boundary_input.get("targets", []): target_node_id = str(target.get("node_id") or "") if not target_node_id: continue prompt_node = prompt.get(target_node_id) if prompt_node is None: continue class_type = str(prompt_node.get("class_type")) node_class = resolved_node_mapping.get(class_type) if node_class is None: continue if bool(getattr(node_class, "INPUT_IS_LIST", False)): target_node_ids.add(target_node_id) return sorted(target_node_ids) def _partition_implicit_batched_execute_nodes( payload: dict[str, Any], split_inputs: dict[str, list[Any]], ) -> tuple[list[str], list[str]]: """Split one implicitly batched subgraph into static and per-item execute targets.""" prompt = payload.get("subgraph_prompt", {}) if not isinstance(prompt, dict): execute_node_ids = [ str(node_id) for node_id in payload.get("execute_node_ids", []) ] return [], execute_node_ids batched_target_node_ids: set[str] = set() for boundary_input in payload.get("boundary_inputs", []): proxy_input_name = str(boundary_input.get("proxy_input_name") or "") if proxy_input_name not in split_inputs: continue for target in boundary_input.get("targets", []): target_node_id = target.get("node_id") if target_node_id is not None: batched_target_node_ids.add(str(target_node_id)) execute_node_ids = [str(node_id) for node_id in payload.get("execute_node_ids", [])] static_execute_node_ids: list[str] = [] mapped_execute_node_ids: list[str] = [] for execute_node_id in execute_node_ids: required_node_ids = set( _resolve_required_subgraph_nodes( prompt=prompt, execute_node_ids=[execute_node_id], ) ) if required_node_ids & batched_target_node_ids: mapped_execute_node_ids.append(execute_node_id) continue static_execute_node_ids.append(execute_node_id) if not mapped_execute_node_ids and execute_node_ids: logger.warning( "Implicitly batched Modal component=%s had batched inputs %s but no execute target depended on them; " "falling back to per-item execution for all execute nodes.", payload.get("component_id"), sorted(split_inputs), ) return [], execute_node_ids logger.info( "Partitioned implicitly batched Modal component=%s into static execute nodes=%s and mapped execute nodes=%s.", payload.get("component_id"), static_execute_node_ids, mapped_execute_node_ids, ) return static_execute_node_ids, mapped_execute_node_ids def _annotate_implicit_batched_boundary_outputs( payload: dict[str, Any], mapped_execute_node_ids: list[str], ) -> list[dict[str, Any]]: """Mark which boundary outputs belong to the per-item branch of an implicitly batched subgraph.""" prompt = payload.get("subgraph_prompt", {}) if not isinstance(prompt, dict) or not mapped_execute_node_ids: return [ copy.deepcopy(boundary_output) for boundary_output in payload.get("boundary_outputs", []) ] mapped_required_node_ids = set( _resolve_required_subgraph_nodes( prompt=prompt, execute_node_ids=[str(node_id) for node_id in mapped_execute_node_ids], ) ) annotated_outputs: list[dict[str, Any]] = [] for boundary_output in payload.get("boundary_outputs", []): annotated_output = copy.deepcopy(boundary_output) annotated_output["mapped_output"] = ( str(boundary_output.get("node_id")) in mapped_required_node_ids ) annotated_outputs.append(annotated_output) return annotated_outputs def _log_detached_mapped_lane_result( task: asyncio.Task[None], *, component_id: str, lane_index: int, ) -> None: """Consume and log any late failure from one detached mapped worker lane.""" try: task.result() except asyncio.CancelledError: return except BaseException as exc: logger.warning( "Detached mapped worker lane=%d for component=%s finished with an unawaited failure: %s", lane_index, component_id, exc, ) async def _invoke_implicitly_mapped_subgraph_async( payload: dict[str, Any], kwargs_payload: bytes ) -> bytes: """Fan out one ordinary subgraph payload when batchable boundary inputs arrive zipped.""" hydrated_inputs = deserialize_node_inputs(kwargs_payload) split_batch_inputs = _split_batch_boundary_inputs(payload, hydrated_inputs) if split_batch_inputs is None: raise ModalRemoteInvocationError( "Implicit mapped subgraph execution requires at least one batched boundary input." ) split_inputs, total_items = split_batch_inputs input_is_list_target_node_ids = _implicit_batch_input_is_list_target_node_ids( payload, split_inputs, ) if input_is_list_target_node_ids: logger.info( "Executing implicitly batched Modal component=%s as one ordinary subgraph because split boundary inputs target INPUT_IS_LIST nodes=%s.", payload.get("component_id"), input_is_list_target_node_ids, ) return await invoke_remote_engine_async( payload, kwargs_payload, allow_implicit_mapping=False, ) parallelism, refined_prompt_warmup_target = boost_mapped_component_warmup( payload, total_items=total_items, reason="implicit_mapped_component_exact_parallelism", ) prompt_id = ( str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None ) if prompt_id is not None: await _await_prompt_warmup_slots( prompt_id, list(range(refined_prompt_warmup_target)), _prompt_warmup_head_start_seconds(), ) logger.info( "Scheduling implicitly mapped Modal component=%s for %d item(s) with local parallelism=%d prompt_warmup_target=%d across inputs=%s.", payload.get("component_id"), total_items, parallelism, refined_prompt_warmup_target, sorted(split_inputs), ) _emit_local_mapped_progress(payload, 0, total_items) broadcast_inputs = { input_name: value for input_name, value in hydrated_inputs.items() if input_name not in split_inputs } ( static_execute_node_ids, mapped_execute_node_ids, ) = _partition_implicit_batched_execute_nodes( payload, split_inputs, ) hybrid_payload = copy.deepcopy(payload) hybrid_payload["static_execute_node_ids"] = static_execute_node_ids hybrid_payload["mapped_execute_node_ids"] = mapped_execute_node_ids hybrid_payload["boundary_outputs"] = _annotate_implicit_batched_boundary_outputs( payload, mapped_execute_node_ids, ) cleanup_payload = _build_remote_session_cleanup_payload(hybrid_payload) execution_mode = os.getenv("COMFY_MODAL_EXECUTION_MODE", "local") use_seeded_remote_lanes = ( execution_mode != "local" and modal is not None and _payload_remote_session_handle(hybrid_payload) is not None and _mapped_phase_definition(hybrid_payload, "static_phase") is not None ) static_outputs: tuple[Any, ...] = () lane_remote_engines: list[Any] = [] seeded_lane_indices: set[int] = set() skip_cleanup_after_interrupt = False try: if use_seeded_remote_lanes: lane_remote_engines = [ _lookup_deployed_remote_engine( hybrid_payload, affinity_key_override=_mapped_lane_affinity_key( hybrid_payload, lane_index ), ) for lane_index in range(parallelism) ] logger.info( "Seeding %d mapped worker lane(s) for component=%s and allowing ready lanes to start item dispatch independently.", parallelism, payload.get("component_id"), ) async def seed_lane(lane_index: int) -> None: """Run the static bridge-producing phase once on one mapped worker lane.""" lane_seed_payload = _build_phase_subgraph_payload( hybrid_payload, "static_phase", f"{payload.get('component_id', 'modal-subgraph')}::seed:{lane_index}", suppress_status_stream=True, ) lane_seed_payload.pop("clear_remote_session", None) await _invoke_bound_remote_engine_async( lane_remote_engines[lane_index], lane_seed_payload, serialize_node_inputs(broadcast_inputs), ) seeded_lane_indices.add(lane_index) elif static_execute_node_ids: static_response = await invoke_remote_engine_async( _build_static_mapped_payload(hybrid_payload), serialize_node_inputs(broadcast_inputs), ) static_outputs = deserialize_node_outputs(static_response) else: lane_remote_engines = [] per_item_outputs: list[tuple[Any, ...] | None] = [None] * total_items completed_items = 0 all_items_completed = asyncio.Event() stop_dispatch_requested = asyncio.Event() worker_failure: asyncio.Future[ BaseException ] = asyncio.get_running_loop().create_future() item_queue: asyncio.Queue[int | None] = asyncio.Queue() for item_index in range(total_items): item_queue.put_nowait(item_index) for _ in range(parallelism): item_queue.put_nowait(None) def request_interrupt_stop() -> None: """Stop queued mapped item dispatch and suppress cleanup after interruption.""" nonlocal skip_cleanup_after_interrupt skip_cleanup_after_interrupt = True stop_dispatch_requested.set() def raise_if_local_interrupted() -> None: """Raise the native ComfyUI interrupt after marking mapped dispatch as stopped.""" if _local_processing_interrupted(): request_interrupt_stop() _raise_local_interrupt() async def run_worker(lane_index: int) -> None: """Execute queued implicit mapped items through one stable local worker lane.""" nonlocal completed_items try: raise_if_local_interrupted() if use_seeded_remote_lanes: _emit_local_mapped_lane_progress_start(payload, lane_index) await seed_lane(lane_index) while True: if stop_dispatch_requested.is_set(): return item_index = await item_queue.get() if item_index is None: return if stop_dispatch_requested.is_set(): return raise_if_local_interrupted() try: item_payload = _build_mapped_item_payload( hybrid_payload, item_index, lane_index ) item_inputs = dict(broadcast_inputs) for input_name, items in split_inputs.items(): item_inputs[input_name] = items[item_index] if use_seeded_remote_lanes: item_response = await _invoke_bound_remote_engine_async( lane_remote_engines[lane_index], item_payload, serialize_node_inputs(item_inputs), ) else: item_response = await invoke_remote_engine_async( item_payload, serialize_node_inputs(item_inputs), ) raise_if_local_interrupted() per_item_outputs[item_index] = deserialize_node_outputs( item_response ) completed_items += 1 _emit_local_mapped_progress( payload, completed_items, total_items ) if completed_items >= total_items: all_items_completed.set() finally: _clear_local_mapped_lane_progress( payload, lane_index, item_index ) except BaseException as exc: if all_items_completed.is_set(): logger.info( "Ignoring late mapped worker failure after component=%s already completed all items lane=%d: %s", payload.get("component_id"), lane_index, exc, ) return stop_dispatch_requested.set() if _exception_indicates_interruption(exc): request_interrupt_stop() if not worker_failure.done(): worker_failure.set_result(exc) raise tasks = [ asyncio.create_task(run_worker(lane_index)) for lane_index in range(parallelism) ] for lane_index, task in enumerate(tasks): task.add_done_callback( lambda completed_task, lane_index=lane_index: _log_detached_mapped_lane_result( completed_task, component_id=str(payload.get("component_id", "modal-subgraph")), lane_index=lane_index, ) ) completion_task = asyncio.create_task(all_items_completed.wait()) try: while not all_items_completed.is_set(): wait_set: set[asyncio.Task[Any] | asyncio.Future[BaseException]] = { completion_task, worker_failure, *tasks, } done, _ = await asyncio.wait( wait_set, return_when=asyncio.FIRST_COMPLETED ) if completion_task in done or all_items_completed.is_set(): break if worker_failure in done: stop_dispatch_requested.set() raise worker_failure.result() for task in done: if task in tasks: task_exc = task.exception() if task_exc is not None: stop_dispatch_requested.set() raise task_exc if not all(item_outputs is not None for item_outputs in per_item_outputs): await asyncio.gather(*tasks) except BaseException as exc: stop_dispatch_requested.set() if _exception_indicates_interruption(exc): skip_cleanup_after_interrupt = True for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) raise finally: completion_task.cancel() await asyncio.gather(completion_task, return_exceptions=True) return serialize_node_outputs( _merge_static_and_mapped_outputs( static_outputs=static_outputs, mapped_outputs=_aggregate_mapped_outputs( [ item_outputs for item_outputs in per_item_outputs if item_outputs is not None ], { **hybrid_payload, "boundary_outputs": [ boundary_output for boundary_output in hybrid_payload.get( "boundary_outputs", [] ) if _is_mapped_boundary_output( boundary_output, hybrid_payload ) ], }, ), payload=hybrid_payload, ) ) finally: if cleanup_payload is not None and not skip_cleanup_after_interrupt: if use_seeded_remote_lanes and lane_remote_engines: await asyncio.gather( *( _invoke_bound_remote_engine_async( lane_remote_engines[lane_index], { **cleanup_payload, "component_id": ( f"{payload.get('component_id', 'modal-subgraph')}::cleanup:{lane_index}" ), }, serialize_node_inputs({}), ) for lane_index in sorted(seeded_lane_indices) ) ) else: await invoke_remote_engine_async( cleanup_payload, serialize_node_inputs({}), ) elif cleanup_payload is not None: logger.info( "Skipping remote-session cleanup for implicitly mapped Modal component=%s because execution was interrupted.", payload.get("component_id"), ) async def _invoke_mapped_remote_engine_async( payload: dict[str, Any], kwargs_payload: bytes ) -> bytes: """Execute one mapped payload locally using explicit static and mapped phases.""" hydrated_inputs = deserialize_node_inputs(kwargs_payload) return await asyncio.to_thread( lambda: serialize_node_outputs( _execute_mapped_subgraph_payload( payload, hydrated_inputs, ) ) )