1434 lines
55 KiB
Python
1434 lines
55 KiB
Python
"""Remote component construction, replication, and boundary validation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import hashlib
|
|
import logging
|
|
from collections import defaultdict, deque
|
|
from typing import Any, Iterable, Mapping
|
|
|
|
if __package__:
|
|
from .execution_environments import ExecutionAssignment, ExecutionProvider
|
|
from .modal_executor_node import MODAL_MAP_INPUT_NODE_ID
|
|
from .remote_graph_analysis import (
|
|
_build_consumer_map,
|
|
_build_remote_components,
|
|
_expand_component_for_non_transportable_local_outputs,
|
|
_is_link,
|
|
_is_transportable_output_type,
|
|
_non_returning_local_output_consumers,
|
|
_non_returning_local_tap_node_ids,
|
|
_normalize_output_metadata,
|
|
_remote_component_partition_groups,
|
|
_remote_output_io_type,
|
|
_remote_output_is_list,
|
|
)
|
|
from .remote_plan_types import (
|
|
BoundaryInputSpec,
|
|
BoundaryOutputSpec,
|
|
InputTarget,
|
|
LinkedOutputRef,
|
|
ModalPromptValidationError,
|
|
RemoteComponentPlan,
|
|
StaticToMappedBoundarySpec,
|
|
)
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from execution_environments import ExecutionAssignment, ExecutionProvider
|
|
from modal_executor_node import MODAL_MAP_INPUT_NODE_ID
|
|
from remote_graph_analysis import (
|
|
_build_consumer_map,
|
|
_build_remote_components,
|
|
_expand_component_for_non_transportable_local_outputs,
|
|
_is_link,
|
|
_is_transportable_output_type,
|
|
_non_returning_local_output_consumers,
|
|
_non_returning_local_tap_node_ids,
|
|
_normalize_output_metadata,
|
|
_remote_component_partition_groups,
|
|
_remote_output_io_type,
|
|
_remote_output_is_list,
|
|
)
|
|
from remote_plan_types import (
|
|
BoundaryInputSpec,
|
|
BoundaryOutputSpec,
|
|
InputTarget,
|
|
LinkedOutputRef,
|
|
ModalPromptValidationError,
|
|
RemoteComponentPlan,
|
|
StaticToMappedBoundarySpec,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_REMOTE_REPLICA_NODE_PREFIX = "__ComfyModalReplica__"
|
|
_REPLICABLE_REMOTE_OBJECT_TYPES = frozenset(
|
|
{"CLIP", "MODEL", "NOISE", "SAMPLER", "VAE"}
|
|
)
|
|
|
|
|
|
def _build_component_plan(
|
|
component_node_ids: list[str],
|
|
prompt: dict[str, Any],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
remote_node_ids: set[str],
|
|
nodes_module: Any,
|
|
) -> RemoteComponentPlan:
|
|
"""Build rewrite metadata for a connected remote component."""
|
|
original_component_node_id_set = set(component_node_ids)
|
|
representative_node_id = component_node_ids[0]
|
|
local_tap_node_ids, local_tap_terminal_node_ids = _non_returning_local_tap_node_ids(
|
|
prompt=prompt,
|
|
component_node_ids=original_component_node_id_set,
|
|
remote_node_ids=remote_node_ids,
|
|
consumers=consumers,
|
|
nodes_module=nodes_module,
|
|
)
|
|
if local_tap_node_ids:
|
|
logger.info(
|
|
"Absorbing non-returning local tap nodes into remote component %s: %s",
|
|
representative_node_id,
|
|
sorted(local_tap_node_ids),
|
|
)
|
|
if local_tap_terminal_node_ids:
|
|
logger.info(
|
|
"Keeping non-returning local tap terminal nodes outside remote component %s: %s",
|
|
representative_node_id,
|
|
sorted(local_tap_terminal_node_ids),
|
|
)
|
|
component_node_id_set = original_component_node_id_set | local_tap_node_ids
|
|
if local_tap_node_ids:
|
|
(
|
|
expanded_component_node_id_set,
|
|
tap_dependency_node_ids,
|
|
) = _expand_component_for_non_transportable_local_outputs(
|
|
prompt=prompt,
|
|
component_node_ids=component_node_id_set,
|
|
remote_node_ids=remote_node_ids,
|
|
consumers=consumers,
|
|
nodes_module=nodes_module,
|
|
)
|
|
if tap_dependency_node_ids:
|
|
logger.info(
|
|
"Absorbing non-transportable boundary dependencies for remote component %s preview taps: %s",
|
|
representative_node_id,
|
|
sorted(tap_dependency_node_ids),
|
|
)
|
|
local_tap_node_ids.update(tap_dependency_node_ids)
|
|
component_node_id_set = expanded_component_node_id_set
|
|
component_node_ids = sorted(component_node_id_set)
|
|
boundary_inputs_by_source: dict[LinkedOutputRef, BoundaryInputSpec] = {}
|
|
boundary_outputs_by_source: dict[LinkedOutputRef, BoundaryOutputSpec] = {}
|
|
output_execution_targets: set[str] = set()
|
|
contains_output_node = False
|
|
|
|
for node_id in component_node_ids:
|
|
prompt_node = prompt[node_id]
|
|
class_type = str(prompt_node["class_type"])
|
|
node_class = nodes_module.NODE_CLASS_MAPPINGS[class_type]
|
|
output_types, output_names, output_is_list = _normalize_output_metadata(
|
|
node_class
|
|
)
|
|
|
|
if getattr(node_class, "OUTPUT_NODE", False):
|
|
contains_output_node = True
|
|
output_execution_targets.add(node_id)
|
|
|
|
for input_name, input_value in (prompt_node.get("inputs") or {}).items():
|
|
if not _is_link(input_value):
|
|
continue
|
|
upstream_node_id = str(input_value[0])
|
|
if upstream_node_id in component_node_id_set:
|
|
continue
|
|
source = LinkedOutputRef(
|
|
node_id=upstream_node_id, output_index=int(input_value[1])
|
|
)
|
|
spec = boundary_inputs_by_source.get(source)
|
|
if spec is None:
|
|
source_io_type = _remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=source.node_id,
|
|
output_index=source.output_index,
|
|
nodes_module=nodes_module,
|
|
)
|
|
spec = BoundaryInputSpec(
|
|
proxy_input_name=f"remote_input_{len(boundary_inputs_by_source)}",
|
|
source=source,
|
|
io_type=source_io_type,
|
|
)
|
|
boundary_inputs_by_source[source] = spec
|
|
spec.targets.append(
|
|
InputTarget(node_id=node_id, input_name=str(input_name))
|
|
)
|
|
|
|
has_downstream_consumer = False
|
|
for output_index, io_type in enumerate(output_types):
|
|
source = LinkedOutputRef(node_id=node_id, output_index=output_index)
|
|
output_consumers = consumers.get(source, [])
|
|
if output_consumers:
|
|
has_downstream_consumer = True
|
|
local_consumers = [
|
|
consumer
|
|
for consumer in output_consumers
|
|
if consumer.node_id not in component_node_id_set
|
|
]
|
|
if not local_consumers:
|
|
continue
|
|
output_execution_targets.add(node_id)
|
|
if source in boundary_outputs_by_source:
|
|
continue
|
|
output_name = output_names[output_index]
|
|
boundary_outputs_by_source[source] = BoundaryOutputSpec(
|
|
proxy_output_name=f"{node_id}_{output_name}",
|
|
source=source,
|
|
io_type=str(io_type),
|
|
is_list=bool(output_is_list[output_index]),
|
|
preview_target_node_ids=_preview_target_node_ids(
|
|
prompt=prompt,
|
|
local_consumers=local_consumers,
|
|
),
|
|
)
|
|
|
|
if not has_downstream_consumer:
|
|
output_execution_targets.add(node_id)
|
|
|
|
mapped_boundary_spec: BoundaryInputSpec | None = None
|
|
mapped_boundary_input_io_type: str | None = None
|
|
for boundary_input in boundary_inputs_by_source.values():
|
|
source_prompt_node = prompt.get(boundary_input.source.node_id)
|
|
source_class_type = (
|
|
str(source_prompt_node.get("class_type"))
|
|
if source_prompt_node is not None
|
|
else None
|
|
)
|
|
mapped_targets = [
|
|
target
|
|
for target in boundary_input.targets
|
|
if str(prompt[target.node_id]["class_type"]) == MODAL_MAP_INPUT_NODE_ID
|
|
]
|
|
source_is_modal_map_input = source_class_type == MODAL_MAP_INPUT_NODE_ID
|
|
if not mapped_targets and not source_is_modal_map_input:
|
|
continue
|
|
if mapped_targets and len(mapped_targets) != len(boundary_input.targets):
|
|
raise ModalPromptValidationError(
|
|
"Mapped remote execution requires the mapped boundary input to feed only ModalMapInput nodes."
|
|
)
|
|
if mapped_boundary_spec is not None:
|
|
raise ModalPromptValidationError(
|
|
"Remote components currently support only one mapped ModalMapInput boundary."
|
|
)
|
|
mapped_boundary_spec = boundary_input
|
|
mapped_boundary_input_io_type = _mapped_boundary_origin_io_type(
|
|
prompt,
|
|
boundary_input,
|
|
nodes_module,
|
|
)
|
|
|
|
mapped_node_ids: list[str] = []
|
|
mapped_execute_node_ids: list[str] = []
|
|
static_execute_node_ids: list[str] = []
|
|
static_node_ids: list[str] = []
|
|
static_to_mapped_boundaries: list[StaticToMappedBoundarySpec] = []
|
|
if mapped_boundary_spec is not None:
|
|
mapped_reachable_node_ids = _component_downstream_closure(
|
|
seed_node_ids={target.node_id for target in mapped_boundary_spec.targets},
|
|
component_node_id_set=component_node_id_set,
|
|
consumers=consumers,
|
|
)
|
|
mapped_node_ids = sorted(mapped_reachable_node_ids)
|
|
mapped_node_id_set = set(mapped_node_ids)
|
|
static_node_ids = sorted(component_node_id_set - mapped_node_id_set)
|
|
static_to_mapped_boundaries_by_source: dict[
|
|
LinkedOutputRef, StaticToMappedBoundarySpec
|
|
] = {}
|
|
for mapped_node_id in mapped_node_ids:
|
|
prompt_node = prompt[mapped_node_id]
|
|
for input_name, input_value in (prompt_node.get("inputs") or {}).items():
|
|
if not _is_link(input_value):
|
|
continue
|
|
upstream_node_id = str(input_value[0])
|
|
if (
|
|
upstream_node_id not in component_node_id_set
|
|
or upstream_node_id in mapped_node_id_set
|
|
):
|
|
continue
|
|
source = LinkedOutputRef(
|
|
node_id=upstream_node_id,
|
|
output_index=int(input_value[1]),
|
|
)
|
|
boundary_spec = static_to_mapped_boundaries_by_source.get(source)
|
|
if boundary_spec is None:
|
|
boundary_spec = StaticToMappedBoundarySpec(
|
|
proxy_name=f"static_input_{len(static_to_mapped_boundaries_by_source)}",
|
|
source=source,
|
|
io_type=str(
|
|
_remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=source.node_id,
|
|
output_index=source.output_index,
|
|
nodes_module=nodes_module,
|
|
)
|
|
or "*"
|
|
),
|
|
is_list=_remote_output_is_list(
|
|
prompt=prompt,
|
|
node_id=source.node_id,
|
|
output_index=source.output_index,
|
|
nodes_module=nodes_module,
|
|
),
|
|
)
|
|
static_to_mapped_boundaries_by_source[source] = boundary_spec
|
|
boundary_spec.targets.append(
|
|
InputTarget(node_id=mapped_node_id, input_name=str(input_name))
|
|
)
|
|
static_to_mapped_boundaries = sorted(
|
|
static_to_mapped_boundaries_by_source.values(),
|
|
key=lambda spec: (spec.source.node_id, spec.source.output_index),
|
|
)
|
|
mapped_execute_node_ids = sorted(output_execution_targets & mapped_node_id_set)
|
|
static_execute_node_ids = sorted(
|
|
(output_execution_targets - mapped_node_id_set)
|
|
| {
|
|
boundary_spec.source.node_id
|
|
for boundary_spec in static_to_mapped_boundaries
|
|
}
|
|
)
|
|
|
|
component = RemoteComponentPlan(
|
|
node_ids=component_node_ids,
|
|
representative_node_id=representative_node_id,
|
|
boundary_inputs=sorted(
|
|
boundary_inputs_by_source.values(),
|
|
key=lambda spec: (spec.source.node_id, spec.source.output_index),
|
|
),
|
|
boundary_outputs=sorted(
|
|
boundary_outputs_by_source.values(),
|
|
key=lambda spec: (spec.source.node_id, spec.source.output_index),
|
|
),
|
|
execute_node_ids=sorted(output_execution_targets),
|
|
contains_output_node=contains_output_node,
|
|
mapped_boundary_input_name=(
|
|
mapped_boundary_spec.proxy_input_name
|
|
if mapped_boundary_spec is not None
|
|
else None
|
|
),
|
|
mapped_boundary_input_io_type=mapped_boundary_input_io_type,
|
|
mapped_boundary_source_node_id=(
|
|
mapped_boundary_spec.source.node_id
|
|
if mapped_boundary_spec is not None
|
|
and prompt.get(mapped_boundary_spec.source.node_id, {}).get("class_type")
|
|
== MODAL_MAP_INPUT_NODE_ID
|
|
else None
|
|
),
|
|
static_node_ids=static_node_ids,
|
|
mapped_node_ids=mapped_node_ids,
|
|
mapped_execute_node_ids=mapped_execute_node_ids,
|
|
static_execute_node_ids=static_execute_node_ids,
|
|
static_to_mapped_boundaries=static_to_mapped_boundaries,
|
|
local_tap_node_ids=sorted(local_tap_node_ids),
|
|
local_tap_terminal_node_ids=sorted(local_tap_terminal_node_ids),
|
|
)
|
|
logger.info(
|
|
"Planned remote component %s: nodes=%s boundary_inputs=%d boundary_outputs=%d execute_nodes=%s output_node=%s mapped_input=%s static_nodes=%s mapped_nodes=%s mapped_execute_nodes=%s static_execute_nodes=%s local_tap_nodes=%s local_tap_terminal_nodes=%s static_to_mapped_boundaries=%s",
|
|
component.representative_node_id,
|
|
component.node_ids,
|
|
len(component.boundary_inputs),
|
|
len(component.boundary_outputs),
|
|
component.execute_node_ids,
|
|
component.contains_output_node,
|
|
component.mapped_boundary_input_name,
|
|
component.static_node_ids,
|
|
component.mapped_node_ids,
|
|
component.mapped_execute_node_ids,
|
|
component.static_execute_node_ids,
|
|
component.local_tap_node_ids,
|
|
component.local_tap_terminal_node_ids,
|
|
[
|
|
{
|
|
"proxy_name": boundary_spec.proxy_name,
|
|
"source": (
|
|
boundary_spec.source.node_id,
|
|
boundary_spec.source.output_index,
|
|
),
|
|
"targets": [
|
|
(target.node_id, target.input_name)
|
|
for target in boundary_spec.targets
|
|
],
|
|
}
|
|
for boundary_spec in component.static_to_mapped_boundaries
|
|
],
|
|
)
|
|
return component
|
|
|
|
|
|
def _filter_boundary_inputs_for_node_ids(
|
|
boundary_inputs: list[BoundaryInputSpec],
|
|
allowed_node_ids: set[str],
|
|
) -> list[BoundaryInputSpec]:
|
|
"""Return boundary inputs whose targets belong to one node-id subset."""
|
|
filtered_boundary_inputs: list[BoundaryInputSpec] = []
|
|
for boundary_input in boundary_inputs:
|
|
filtered_targets = [
|
|
target
|
|
for target in boundary_input.targets
|
|
if target.node_id in allowed_node_ids
|
|
]
|
|
if not filtered_targets:
|
|
continue
|
|
filtered_boundary_inputs.append(
|
|
BoundaryInputSpec(
|
|
proxy_input_name=boundary_input.proxy_input_name,
|
|
source=boundary_input.source,
|
|
io_type=boundary_input.io_type,
|
|
targets=filtered_targets,
|
|
)
|
|
)
|
|
return filtered_boundary_inputs
|
|
|
|
|
|
def _filter_boundary_outputs_for_node_ids(
|
|
boundary_outputs: list[BoundaryOutputSpec],
|
|
allowed_node_ids: set[str],
|
|
) -> list[BoundaryOutputSpec]:
|
|
"""Return boundary outputs exported by one node-id subset."""
|
|
return [
|
|
boundary_output
|
|
for boundary_output in boundary_outputs
|
|
if boundary_output.source.node_id in allowed_node_ids
|
|
]
|
|
|
|
|
|
def _subset_component_prompt(
|
|
component_prompt: dict[str, Any],
|
|
node_ids: list[str],
|
|
) -> dict[str, Any]:
|
|
"""Return one deep-copied prompt subset for a phase-local node set."""
|
|
return {
|
|
node_id: copy.deepcopy(component_prompt[node_id])
|
|
for node_id in node_ids
|
|
if node_id in component_prompt
|
|
}
|
|
|
|
|
|
def _preview_target_node_ids(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
local_consumers: list[InputTarget],
|
|
) -> list[str]:
|
|
"""Return direct local PreviewImage consumers for one remote boundary output."""
|
|
preview_target_node_ids: set[str] = set()
|
|
for local_consumer in local_consumers:
|
|
consumer_prompt_node = prompt.get(local_consumer.node_id)
|
|
if consumer_prompt_node is None:
|
|
continue
|
|
if str(consumer_prompt_node.get("class_type")) != "PreviewImage":
|
|
continue
|
|
preview_target_node_ids.add(str(local_consumer.node_id))
|
|
return sorted(preview_target_node_ids)
|
|
|
|
|
|
def _component_downstream_closure(
|
|
*,
|
|
seed_node_ids: set[str],
|
|
component_node_id_set: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
) -> set[str]:
|
|
"""Return component-local nodes reachable downstream from one seed set."""
|
|
reachable_node_ids: set[str] = set()
|
|
pending_node_ids = list(sorted(seed_node_ids))
|
|
while pending_node_ids:
|
|
current_node_id = pending_node_ids.pop()
|
|
if (
|
|
current_node_id in reachable_node_ids
|
|
or current_node_id not in component_node_id_set
|
|
):
|
|
continue
|
|
reachable_node_ids.add(current_node_id)
|
|
for consumer_source, consumer_targets in consumers.items():
|
|
if consumer_source.node_id != current_node_id:
|
|
continue
|
|
for consumer_target in consumer_targets:
|
|
if consumer_target.node_id in component_node_id_set:
|
|
pending_node_ids.append(consumer_target.node_id)
|
|
return reachable_node_ids
|
|
|
|
|
|
def _component_upstream_closure(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
seed_node_ids: set[str],
|
|
candidate_node_ids: set[str],
|
|
) -> set[str]:
|
|
"""Return one candidate-local upstream closure for the supplied seed nodes."""
|
|
reachable_node_ids: set[str] = set()
|
|
pending_node_ids = list(sorted(seed_node_ids))
|
|
|
|
while pending_node_ids:
|
|
current_node_id = pending_node_ids.pop()
|
|
if (
|
|
current_node_id in reachable_node_ids
|
|
or current_node_id not in candidate_node_ids
|
|
):
|
|
continue
|
|
reachable_node_ids.add(current_node_id)
|
|
prompt_node = prompt.get(current_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
|
|
upstream_node_id = str(input_value[0])
|
|
if upstream_node_id in candidate_node_ids:
|
|
pending_node_ids.append(upstream_node_id)
|
|
return reachable_node_ids
|
|
|
|
|
|
def _component_ancestors_of_local_source(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
source_node_id: str,
|
|
component_node_id_set: set[str],
|
|
) -> set[str]:
|
|
"""Return component nodes that feed one local boundary-input source."""
|
|
ancestor_node_ids: set[str] = set()
|
|
visited_node_ids: set[str] = set()
|
|
pending_node_ids = [source_node_id]
|
|
|
|
while pending_node_ids:
|
|
current_node_id = str(pending_node_ids.pop())
|
|
if current_node_id in visited_node_ids:
|
|
continue
|
|
visited_node_ids.add(current_node_id)
|
|
if current_node_id in component_node_id_set:
|
|
ancestor_node_ids.add(current_node_id)
|
|
continue
|
|
prompt_node = prompt.get(current_node_id)
|
|
if prompt_node is None:
|
|
continue
|
|
for input_value in (prompt_node.get("inputs") or {}).values():
|
|
if _is_link(input_value):
|
|
pending_node_ids.append(str(input_value[0]))
|
|
return ancestor_node_ids
|
|
|
|
|
|
def _component_has_local_reentry_dependency(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
component: RemoteComponentPlan,
|
|
) -> bool:
|
|
"""Return whether a component boundary input depends on that same component's output."""
|
|
component_node_id_set = set(component.node_ids)
|
|
for boundary_input in component.boundary_inputs:
|
|
if boundary_input.source.node_id in component_node_id_set:
|
|
continue
|
|
if _component_ancestors_of_local_source(
|
|
prompt=prompt,
|
|
source_node_id=boundary_input.source.node_id,
|
|
component_node_id_set=component_node_id_set,
|
|
):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _component_has_parallel_local_remote_fanout(
|
|
component: RemoteComponentPlan,
|
|
) -> bool:
|
|
"""Return whether one remote output feeds local work and a remote continuation."""
|
|
return any(
|
|
boundary_output.local_materializer_node_id is not None
|
|
and bool(boundary_output.session_consumer_node_ids)
|
|
for boundary_output in component.boundary_outputs
|
|
)
|
|
|
|
|
|
def _order_execute_node_ids_for_transportable_splits(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
component_prompt: dict[str, Any],
|
|
component_node_ids: set[str],
|
|
execute_node_ids: list[str],
|
|
) -> list[str]:
|
|
"""Return split execute targets ordered by component and local feedback dependencies."""
|
|
base_order = list(execute_node_ids)
|
|
if len(base_order) <= 1:
|
|
return base_order
|
|
|
|
closure_by_execute_node_id = {
|
|
execute_node_id: _component_upstream_closure(
|
|
prompt=component_prompt,
|
|
seed_node_ids={execute_node_id},
|
|
candidate_node_ids=component_node_ids,
|
|
)
|
|
for execute_node_id in base_order
|
|
}
|
|
producer_execute_node_ids_by_component_node_id: dict[str, list[str]] = defaultdict(
|
|
list
|
|
)
|
|
for execute_node_id in base_order:
|
|
for component_node_id in closure_by_execute_node_id[execute_node_id]:
|
|
producer_execute_node_ids_by_component_node_id[component_node_id].append(
|
|
execute_node_id
|
|
)
|
|
|
|
base_index_by_execute_node_id = {
|
|
execute_node_id: index for index, execute_node_id in enumerate(base_order)
|
|
}
|
|
dependency_edges: dict[str, set[str]] = {
|
|
execute_node_id: set() for execute_node_id in base_order
|
|
}
|
|
indegree_by_execute_node_id: dict[str, int] = {
|
|
execute_node_id: 0 for execute_node_id in base_order
|
|
}
|
|
|
|
def producer_for_component_node(component_node_id: str) -> str | None:
|
|
"""Return the earliest execute target that produces a component node."""
|
|
producer_execute_node_ids = producer_execute_node_ids_by_component_node_id.get(
|
|
component_node_id,
|
|
[],
|
|
)
|
|
if not producer_execute_node_ids:
|
|
return None
|
|
exact_producers = [
|
|
execute_node_id
|
|
for execute_node_id in producer_execute_node_ids
|
|
if execute_node_id == component_node_id
|
|
]
|
|
if exact_producers:
|
|
return exact_producers[0]
|
|
return min(
|
|
producer_execute_node_ids,
|
|
key=lambda execute_node_id: base_index_by_execute_node_id[execute_node_id],
|
|
)
|
|
|
|
for target_execute_node_id in base_order:
|
|
target_closure = closure_by_execute_node_id[target_execute_node_id]
|
|
for phase_node_id in sorted(target_closure):
|
|
prompt_node = component_prompt.get(phase_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_node_id = str(input_value[0])
|
|
if source_node_id in component_node_ids:
|
|
continue
|
|
source_ancestor_node_ids = _component_ancestors_of_local_source(
|
|
prompt=prompt,
|
|
source_node_id=source_node_id,
|
|
component_node_id_set=component_node_ids,
|
|
)
|
|
for source_ancestor_node_id in sorted(source_ancestor_node_ids):
|
|
if source_ancestor_node_id in target_closure:
|
|
continue
|
|
producer_execute_node_id = producer_for_component_node(
|
|
source_ancestor_node_id
|
|
)
|
|
if (
|
|
producer_execute_node_id is None
|
|
or producer_execute_node_id == target_execute_node_id
|
|
or target_execute_node_id
|
|
in dependency_edges[producer_execute_node_id]
|
|
):
|
|
continue
|
|
dependency_edges[producer_execute_node_id].add(
|
|
target_execute_node_id
|
|
)
|
|
indegree_by_execute_node_id[target_execute_node_id] += 1
|
|
|
|
ready_execute_node_ids = [
|
|
execute_node_id
|
|
for execute_node_id in base_order
|
|
if indegree_by_execute_node_id[execute_node_id] == 0
|
|
]
|
|
ordered_execute_node_ids: list[str] = []
|
|
while ready_execute_node_ids:
|
|
ready_execute_node_ids.sort(
|
|
key=lambda node_id: base_index_by_execute_node_id[node_id]
|
|
)
|
|
execute_node_id = ready_execute_node_ids.pop(0)
|
|
ordered_execute_node_ids.append(execute_node_id)
|
|
for downstream_execute_node_id in sorted(
|
|
dependency_edges[execute_node_id],
|
|
key=lambda node_id: base_index_by_execute_node_id[node_id],
|
|
):
|
|
indegree_by_execute_node_id[downstream_execute_node_id] -= 1
|
|
if indegree_by_execute_node_id[downstream_execute_node_id] == 0:
|
|
ready_execute_node_ids.append(downstream_execute_node_id)
|
|
|
|
if len(ordered_execute_node_ids) != len(base_order):
|
|
logger.warning(
|
|
"Split phase local-feedback dependency ordering encountered a cycle; keeping base execute order %s.",
|
|
base_order,
|
|
)
|
|
return base_order
|
|
|
|
if ordered_execute_node_ids != base_order:
|
|
logger.info(
|
|
"Reordered split execute targets from %s to %s using local-feedback dependencies %s.",
|
|
base_order,
|
|
ordered_execute_node_ids,
|
|
{
|
|
execute_node_id: sorted(downstream_execute_node_ids)
|
|
for execute_node_id, downstream_execute_node_ids in dependency_edges.items()
|
|
if downstream_execute_node_ids
|
|
},
|
|
)
|
|
return ordered_execute_node_ids
|
|
|
|
|
|
def _subgraph_topological_node_order(
|
|
prompt: dict[str, Any],
|
|
node_ids: set[str],
|
|
) -> list[str]:
|
|
"""Return one best-effort topological order for a prompt node subset."""
|
|
indegree_by_node_id: dict[str, int] = {node_id: 0 for node_id in node_ids}
|
|
downstream_node_ids_by_node_id: dict[str, set[str]] = {
|
|
node_id: set() for node_id in node_ids
|
|
}
|
|
|
|
for node_id in sorted(node_ids):
|
|
prompt_node = prompt.get(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
|
|
upstream_node_id = str(input_value[0])
|
|
if upstream_node_id not in node_ids:
|
|
continue
|
|
if node_id in downstream_node_ids_by_node_id[upstream_node_id]:
|
|
continue
|
|
downstream_node_ids_by_node_id[upstream_node_id].add(node_id)
|
|
indegree_by_node_id[node_id] += 1
|
|
|
|
ready_node_ids = deque(
|
|
sorted(
|
|
node_id
|
|
for node_id, indegree in indegree_by_node_id.items()
|
|
if indegree == 0
|
|
)
|
|
)
|
|
ordered_node_ids: list[str] = []
|
|
|
|
while ready_node_ids:
|
|
node_id = ready_node_ids.popleft()
|
|
ordered_node_ids.append(node_id)
|
|
for downstream_node_id in sorted(downstream_node_ids_by_node_id[node_id]):
|
|
indegree_by_node_id[downstream_node_id] -= 1
|
|
if indegree_by_node_id[downstream_node_id] == 0:
|
|
ready_node_ids.append(downstream_node_id)
|
|
|
|
if len(ordered_node_ids) == len(node_ids):
|
|
return ordered_node_ids
|
|
|
|
logger.warning(
|
|
"Component phase planning encountered a node-level cycle; falling back to stable node order for %s.",
|
|
sorted(node_ids),
|
|
)
|
|
return sorted(node_ids)
|
|
|
|
|
|
def _remote_dependency_closure(
|
|
prompt: Mapping[str, Any],
|
|
source_node_id: str,
|
|
) -> set[str]:
|
|
"""Return the complete upstream dependency closure for one remote producer."""
|
|
closure: set[str] = set()
|
|
pending = [source_node_id]
|
|
while pending:
|
|
node_id = pending.pop()
|
|
if node_id in closure:
|
|
continue
|
|
prompt_node = prompt.get(node_id)
|
|
if not isinstance(prompt_node, Mapping):
|
|
continue
|
|
closure.add(node_id)
|
|
pending.extend(
|
|
str(input_value[0])
|
|
for input_value in (prompt_node.get("inputs") or {}).values()
|
|
if _is_link(input_value)
|
|
)
|
|
return closure
|
|
|
|
|
|
def _can_replicate_remote_dependency_closure(
|
|
*,
|
|
prompt: Mapping[str, Any],
|
|
source_node_id: str,
|
|
dependency_node_ids: set[str],
|
|
remote_node_ids: set[str],
|
|
) -> bool:
|
|
"""Return whether a producer closure can be safely rebuilt on another worker."""
|
|
if not dependency_node_ids or not dependency_node_ids.issubset(remote_node_ids):
|
|
return False
|
|
source_node = prompt.get(source_node_id)
|
|
if not isinstance(source_node, Mapping):
|
|
return False
|
|
source_inputs = source_node.get("inputs")
|
|
if not isinstance(source_inputs, Mapping):
|
|
return False
|
|
if not any(_is_link(value) for value in source_inputs.values()):
|
|
return True
|
|
return all(
|
|
"sampler"
|
|
not in str((prompt.get(node_id) or {}).get("class_type") or "").casefold()
|
|
for node_id in dependency_node_ids
|
|
)
|
|
|
|
|
|
def _replica_node_id_mapping(
|
|
*,
|
|
prompt: Mapping[str, Any],
|
|
dependency_node_ids: set[str],
|
|
source_node_id: str,
|
|
target_group_id: str,
|
|
) -> dict[str, str]:
|
|
"""Return collision-free deterministic ids for one replicated dependency closure."""
|
|
replica_digest = hashlib.sha256(
|
|
f"{source_node_id}\0{target_group_id}".encode("utf-8")
|
|
).hexdigest()[:12]
|
|
reserved_node_ids = set(prompt)
|
|
mapping: dict[str, str] = {}
|
|
for node_id in sorted(dependency_node_ids):
|
|
base_id = f"{_REMOTE_REPLICA_NODE_PREFIX}{replica_digest}__{node_id}"
|
|
replica_node_id = base_id
|
|
suffix = 1
|
|
while replica_node_id in reserved_node_ids:
|
|
replica_node_id = f"{base_id}_{suffix}"
|
|
suffix += 1
|
|
mapping[node_id] = replica_node_id
|
|
reserved_node_ids.add(replica_node_id)
|
|
return mapping
|
|
|
|
|
|
def _install_remote_dependency_replica(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
dependency_node_ids: set[str],
|
|
source_node_id: str,
|
|
target_group_id: str,
|
|
) -> str:
|
|
"""Clone one safe producer closure and return its replicated source node id."""
|
|
replica_ids = _replica_node_id_mapping(
|
|
prompt=prompt,
|
|
dependency_node_ids=dependency_node_ids,
|
|
source_node_id=source_node_id,
|
|
target_group_id=target_group_id,
|
|
)
|
|
for node_id in sorted(dependency_node_ids):
|
|
replica_node = copy.deepcopy(prompt[node_id])
|
|
for input_name, input_value in list(
|
|
(replica_node.get("inputs") or {}).items()
|
|
):
|
|
if _is_link(input_value) and str(input_value[0]) in replica_ids:
|
|
replica_node["inputs"][input_name] = [
|
|
replica_ids[str(input_value[0])],
|
|
int(input_value[1]),
|
|
]
|
|
prompt[replica_ids[node_id]] = replica_node
|
|
remote_node_ids.update(replica_ids.values())
|
|
return replica_ids[source_node_id]
|
|
|
|
|
|
def _cross_group_replicable_boundaries(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
component_groups: Mapping[str, set[str]],
|
|
nodes_module: Any,
|
|
) -> list[tuple[str, int, str, str, str]]:
|
|
"""Return replicable non-transportable edges crossing tentative components."""
|
|
group_by_node_id = {
|
|
node_id: group_id
|
|
for group_id, node_ids in component_groups.items()
|
|
for node_id in node_ids
|
|
}
|
|
boundaries: list[tuple[str, int, str, str, str]] = []
|
|
for target_node_id in sorted(remote_node_ids):
|
|
target_node = prompt.get(target_node_id)
|
|
if not isinstance(target_node, Mapping):
|
|
continue
|
|
for input_name, input_value in (target_node.get("inputs") or {}).items():
|
|
if not _is_link(input_value):
|
|
continue
|
|
source_node_id = str(input_value[0])
|
|
if (
|
|
source_node_id not in remote_node_ids
|
|
or group_by_node_id.get(source_node_id)
|
|
== group_by_node_id.get(target_node_id)
|
|
):
|
|
continue
|
|
output_index = int(input_value[1])
|
|
io_type = _remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=source_node_id,
|
|
output_index=output_index,
|
|
nodes_module=nodes_module,
|
|
)
|
|
if io_type not in _REPLICABLE_REMOTE_OBJECT_TYPES:
|
|
continue
|
|
boundaries.append(
|
|
(
|
|
source_node_id,
|
|
output_index,
|
|
target_node_id,
|
|
str(input_name),
|
|
str(group_by_node_id[target_node_id]),
|
|
)
|
|
)
|
|
return boundaries
|
|
|
|
|
|
def _replicate_safe_nontransportable_provider_boundaries(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
nodes_module: Any,
|
|
) -> None:
|
|
"""Replicate safe object producers instead of forcing provider cycles together."""
|
|
for _iteration in range(len(prompt) + 1):
|
|
consumers = _build_consumer_map(prompt)
|
|
tentative_groups = _remote_component_partition_groups(
|
|
prompt,
|
|
remote_node_ids,
|
|
consumers,
|
|
nodes_module,
|
|
preserve_nontransportable_affinity=False,
|
|
)
|
|
boundaries = _cross_group_replicable_boundaries(
|
|
prompt=prompt,
|
|
remote_node_ids=remote_node_ids,
|
|
component_groups=tentative_groups,
|
|
nodes_module=nodes_module,
|
|
)
|
|
replica_source_by_group: dict[tuple[str, str], str] = {}
|
|
changed = False
|
|
for (
|
|
source_node_id,
|
|
output_index,
|
|
target_node_id,
|
|
input_name,
|
|
group_id,
|
|
) in boundaries:
|
|
dependency_node_ids = _remote_dependency_closure(prompt, source_node_id)
|
|
if not _can_replicate_remote_dependency_closure(
|
|
prompt=prompt,
|
|
source_node_id=source_node_id,
|
|
dependency_node_ids=dependency_node_ids,
|
|
remote_node_ids=remote_node_ids,
|
|
):
|
|
continue
|
|
replica_key = (source_node_id, group_id)
|
|
replica_source_node_id = replica_source_by_group.get(replica_key)
|
|
if replica_source_node_id is None:
|
|
replica_source_node_id = _install_remote_dependency_replica(
|
|
prompt=prompt,
|
|
remote_node_ids=remote_node_ids,
|
|
dependency_node_ids=dependency_node_ids,
|
|
source_node_id=source_node_id,
|
|
target_group_id=group_id,
|
|
)
|
|
replica_source_by_group[replica_key] = replica_source_node_id
|
|
logger.info(
|
|
"Replicated non-transportable producer %s and dependencies %s "
|
|
"for remote component group %s.",
|
|
source_node_id,
|
|
sorted(dependency_node_ids),
|
|
group_id,
|
|
)
|
|
prompt[target_node_id]["inputs"][input_name] = [
|
|
replica_source_node_id,
|
|
output_index,
|
|
]
|
|
changed = True
|
|
if not changed:
|
|
return
|
|
raise ModalPromptValidationError(
|
|
"Unable to stabilize provider-aware non-transportable dependency replicas."
|
|
)
|
|
|
|
|
|
def _workflow_visible_remote_node_ids(node_ids: Iterable[str]) -> list[str]:
|
|
"""Exclude internal dependency replicas from workflow-facing rewrite metadata."""
|
|
return [
|
|
str(node_id)
|
|
for node_id in node_ids
|
|
if not str(node_id).startswith(_REMOTE_REPLICA_NODE_PREFIX)
|
|
]
|
|
|
|
|
|
def _build_component_plans(
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
nodes_module: Any,
|
|
) -> list[RemoteComponentPlan]:
|
|
"""Build plans for every connected remote component."""
|
|
_replicate_safe_nontransportable_provider_boundaries(
|
|
prompt=prompt,
|
|
remote_node_ids=remote_node_ids,
|
|
nodes_module=nodes_module,
|
|
)
|
|
consumers = _build_consumer_map(prompt)
|
|
components = _build_remote_components(
|
|
prompt, remote_node_ids, consumers, nodes_module
|
|
)
|
|
return [
|
|
_build_component_plan(
|
|
component, prompt, consumers, remote_node_ids, nodes_module
|
|
)
|
|
for component in components
|
|
]
|
|
|
|
|
|
def _mark_remote_to_remote_session_boundaries(
|
|
prompt: dict[str, Any],
|
|
components: list[RemoteComponentPlan],
|
|
nodes_module: Any,
|
|
assignments_by_component_id: Mapping[str, ExecutionAssignment],
|
|
) -> set[str]:
|
|
"""Keep only same-environment remote edges in provider-local session storage."""
|
|
consumers = _build_consumer_map(prompt)
|
|
component_id_by_node_id = {
|
|
node_id: component.representative_node_id
|
|
for component in components
|
|
for node_id in component.node_ids
|
|
}
|
|
session_sources: set[LinkedOutputRef] = set()
|
|
|
|
for component in components:
|
|
for boundary_output in component.boundary_outputs:
|
|
output_consumers = consumers.get(boundary_output.source, [])
|
|
if not output_consumers:
|
|
continue
|
|
remote_consumers = [
|
|
consumer
|
|
for consumer in output_consumers
|
|
if consumer.node_id in component_id_by_node_id
|
|
]
|
|
if not remote_consumers:
|
|
continue
|
|
producer_assignment = assignments_by_component_id[
|
|
component.representative_node_id
|
|
]
|
|
consumer_component_ids = {
|
|
component_id_by_node_id[consumer.node_id]
|
|
for consumer in remote_consumers
|
|
}
|
|
cross_environment_consumers = sorted(
|
|
consumer_component_id
|
|
for consumer_component_id in consumer_component_ids
|
|
if assignments_by_component_id[consumer_component_id].environment_id
|
|
!= producer_assignment.environment_id
|
|
)
|
|
if cross_environment_consumers:
|
|
logger.info(
|
|
"Materializing remote boundary output through ComfyUI across "
|
|
"execution environments source=%s:%d producer=%s environment=%s "
|
|
"consumer_components=%s.",
|
|
boundary_output.source.node_id,
|
|
boundary_output.source.output_index,
|
|
component.representative_node_id,
|
|
producer_assignment.environment_id,
|
|
cross_environment_consumers,
|
|
)
|
|
continue
|
|
(
|
|
non_returning_local_consumers,
|
|
returning_local_consumers,
|
|
) = _non_returning_local_output_consumers(
|
|
prompt=prompt,
|
|
source=boundary_output.source,
|
|
remote_node_ids=set(component_id_by_node_id),
|
|
consumers=consumers,
|
|
nodes_module=nodes_module,
|
|
)
|
|
if returning_local_consumers:
|
|
logger.info(
|
|
"Materializing remote boundary output normally because local consumers %s later return to remote execution source=%s:%d.",
|
|
[consumer.node_id for consumer in returning_local_consumers],
|
|
boundary_output.source.node_id,
|
|
boundary_output.source.output_index,
|
|
)
|
|
continue
|
|
if (
|
|
non_returning_local_consumers
|
|
and producer_assignment.provider is not ExecutionProvider.MODAL
|
|
):
|
|
logger.info(
|
|
"Materializing remote boundary output through ComfyUI because "
|
|
"provider-local bridge storage is not accessible to local consumers "
|
|
"source=%s:%d provider=%s producer_component=%s local_consumers=%s.",
|
|
boundary_output.source.node_id,
|
|
boundary_output.source.output_index,
|
|
producer_assignment.provider.value,
|
|
component.representative_node_id,
|
|
[
|
|
consumer.node_id
|
|
for consumer in non_returning_local_consumers
|
|
],
|
|
)
|
|
continue
|
|
if boundary_output.is_list and _is_transportable_output_type(
|
|
boundary_output.io_type
|
|
):
|
|
logger.info(
|
|
"Materializing transportable list boundary output through ComfyUI "
|
|
"to preserve scheduler item semantics source=%s:%d io_type=%s "
|
|
"producer_component=%s consumer_components=%s.",
|
|
boundary_output.source.node_id,
|
|
boundary_output.source.output_index,
|
|
boundary_output.io_type,
|
|
component.representative_node_id,
|
|
sorted(consumer_component_ids),
|
|
)
|
|
continue
|
|
boundary_output.session_output = True
|
|
boundary_output.session_consumer_node_ids = sorted(
|
|
{consumer.node_id for consumer in remote_consumers}
|
|
)
|
|
if non_returning_local_consumers:
|
|
materializer_node_id = (
|
|
f"__ModalLocalBridgeMaterializer__"
|
|
f"{boundary_output.source.node_id}_"
|
|
f"{boundary_output.source.output_index}"
|
|
)
|
|
reserved_node_ids = {
|
|
output.local_materializer_node_id
|
|
for planned_component in components
|
|
for output in planned_component.boundary_outputs
|
|
if output.local_materializer_node_id is not None
|
|
}
|
|
while (
|
|
materializer_node_id in prompt
|
|
or materializer_node_id in reserved_node_ids
|
|
):
|
|
materializer_node_id = f"{materializer_node_id}_proxy"
|
|
boundary_output.local_materializer_node_id = materializer_node_id
|
|
boundary_output.local_materializer_consumer_node_ids = sorted(
|
|
{consumer.node_id for consumer in non_returning_local_consumers}
|
|
)
|
|
boundary_output.preview_target_node_ids = []
|
|
session_sources.add(boundary_output.source)
|
|
logger.info(
|
|
"Keeping remote boundary output in provider-local storage source=%s:%d io_type=%s provider=%s producer_component=%s consumer_components=%s local_materializer=%s local_consumers=%s.",
|
|
boundary_output.source.node_id,
|
|
boundary_output.source.output_index,
|
|
boundary_output.io_type,
|
|
producer_assignment.provider.value,
|
|
component.representative_node_id,
|
|
sorted(
|
|
{
|
|
component_id_by_node_id[consumer.node_id]
|
|
for consumer in output_consumers
|
|
if consumer.node_id in component_id_by_node_id
|
|
}
|
|
),
|
|
boundary_output.local_materializer_node_id,
|
|
boundary_output.local_materializer_consumer_node_ids,
|
|
)
|
|
|
|
return _remote_session_component_ids(
|
|
components=components,
|
|
session_sources=session_sources,
|
|
)
|
|
|
|
|
|
def _remote_session_component_ids(
|
|
*,
|
|
components: list[RemoteComponentPlan],
|
|
session_sources: set[LinkedOutputRef],
|
|
) -> set[str]:
|
|
"""Return components that produce or consume provider-local boundary refs."""
|
|
return {
|
|
component.representative_node_id
|
|
for component in components
|
|
if any(output.session_output for output in component.boundary_outputs)
|
|
or any(
|
|
boundary_input.source in session_sources
|
|
for boundary_input in component.boundary_inputs
|
|
)
|
|
}
|
|
|
|
|
|
def _boundary_output_payload(
|
|
boundary_output: BoundaryOutputSpec,
|
|
*,
|
|
mapped_output: bool | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Serialize one component boundary output for a remote payload."""
|
|
payload = {
|
|
"proxy_output_name": boundary_output.proxy_output_name,
|
|
"node_id": boundary_output.source.node_id,
|
|
"output_index": boundary_output.source.output_index,
|
|
"io_type": boundary_output.io_type,
|
|
"is_list": boundary_output.is_list,
|
|
"preview_target_node_ids": list(boundary_output.preview_target_node_ids),
|
|
}
|
|
if boundary_output.session_output:
|
|
payload["session_output"] = True
|
|
if mapped_output is not None:
|
|
payload["mapped_output"] = mapped_output
|
|
if mapped_output:
|
|
payload["scheduler_is_list"] = True
|
|
return payload
|
|
|
|
|
|
def _proxy_boundary_output_is_list(boundary_output: Mapping[str, Any]) -> bool:
|
|
"""Return whether ComfyUI should expose one proxy boundary as a list output."""
|
|
return bool(
|
|
boundary_output.get("is_list", False)
|
|
or boundary_output.get("scheduler_is_list", False)
|
|
)
|
|
|
|
|
|
def _implicitly_mapped_boundary_output_sources(
|
|
*,
|
|
component: RemoteComponentPlan,
|
|
original_prompt: dict[str, Any],
|
|
rewritten_prompt: dict[str, Any],
|
|
nodes_module: Any,
|
|
) -> set[LinkedOutputRef]:
|
|
"""Return component outputs transitively driven by scheduler-list proxy inputs."""
|
|
mapped_target_node_ids: set[str] = set()
|
|
list_boundary_input_names: set[str] = set()
|
|
|
|
for boundary_input in component.boundary_inputs:
|
|
input_is_scheduler_list = False
|
|
for target in boundary_input.targets:
|
|
target_prompt_node = rewritten_prompt.get(target.node_id)
|
|
if target_prompt_node is None:
|
|
continue
|
|
current_input_value = (target_prompt_node.get("inputs") or {}).get(
|
|
target.input_name
|
|
)
|
|
if not _is_link(current_input_value):
|
|
continue
|
|
if not _remote_output_is_list(
|
|
prompt=rewritten_prompt,
|
|
node_id=str(current_input_value[0]),
|
|
output_index=int(current_input_value[1]),
|
|
nodes_module=nodes_module,
|
|
):
|
|
continue
|
|
input_is_scheduler_list = True
|
|
mapped_target_node_ids.add(target.node_id)
|
|
|
|
if input_is_scheduler_list:
|
|
list_boundary_input_names.add(boundary_input.proxy_input_name)
|
|
|
|
if not mapped_target_node_ids:
|
|
return set()
|
|
|
|
consumers = _build_consumer_map(original_prompt)
|
|
mapped_node_ids = _component_downstream_closure(
|
|
seed_node_ids=mapped_target_node_ids,
|
|
component_node_id_set=set(component.node_ids),
|
|
consumers=consumers,
|
|
)
|
|
mapped_output_sources = {
|
|
boundary_output.source
|
|
for boundary_output in component.boundary_outputs
|
|
if boundary_output.source.node_id in mapped_node_ids
|
|
}
|
|
logger.info(
|
|
"Propagating scheduler-list metadata through remote component %s from "
|
|
"boundary inputs=%s to outputs=%s.",
|
|
component.representative_node_id,
|
|
sorted(list_boundary_input_names),
|
|
sorted(
|
|
f"{source.node_id}:{source.output_index}"
|
|
for source in mapped_output_sources
|
|
),
|
|
)
|
|
return mapped_output_sources
|
|
|
|
|
|
def _mark_payload_scheduler_list_outputs(
|
|
payload: dict[str, Any],
|
|
scheduler_list_output_sources: set[LinkedOutputRef],
|
|
) -> None:
|
|
"""Mark matching outputs across a component payload and any split phase payloads."""
|
|
if not scheduler_list_output_sources:
|
|
return
|
|
|
|
boundary_outputs = payload.get("boundary_outputs")
|
|
if isinstance(boundary_outputs, list):
|
|
for boundary_output in boundary_outputs:
|
|
if not isinstance(boundary_output, dict):
|
|
continue
|
|
source = LinkedOutputRef(
|
|
node_id=str(boundary_output.get("node_id")),
|
|
output_index=int(boundary_output.get("output_index", -1)),
|
|
)
|
|
if source in scheduler_list_output_sources:
|
|
boundary_output["scheduler_is_list"] = True
|
|
|
|
for phase_name in ("static_phase", "mapped_phase"):
|
|
phase_payload = payload.get(phase_name)
|
|
if isinstance(phase_payload, dict):
|
|
_mark_payload_scheduler_list_outputs(
|
|
phase_payload,
|
|
scheduler_list_output_sources,
|
|
)
|
|
|
|
split_proxy_payloads = payload.get("split_proxy_payloads")
|
|
if isinstance(split_proxy_payloads, dict):
|
|
for split_payload in split_proxy_payloads.values():
|
|
if isinstance(split_payload, dict):
|
|
_mark_payload_scheduler_list_outputs(
|
|
split_payload,
|
|
scheduler_list_output_sources,
|
|
)
|
|
elif isinstance(split_proxy_payloads, list):
|
|
for split_payload in split_proxy_payloads:
|
|
if isinstance(split_payload, dict):
|
|
_mark_payload_scheduler_list_outputs(
|
|
split_payload,
|
|
scheduler_list_output_sources,
|
|
)
|
|
|
|
|
|
def _mapped_boundary_origin_io_type(
|
|
prompt: dict[str, Any],
|
|
boundary_input: BoundaryInputSpec,
|
|
nodes_module: Any,
|
|
) -> str | None:
|
|
"""Return the effective io_type for one mapped boundary, unwrapping local ModalMapInput markers."""
|
|
source_prompt_node = prompt.get(boundary_input.source.node_id)
|
|
if source_prompt_node is None:
|
|
return boundary_input.io_type
|
|
|
|
if str(source_prompt_node.get("class_type")) != MODAL_MAP_INPUT_NODE_ID:
|
|
return _remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=boundary_input.source.node_id,
|
|
output_index=boundary_input.source.output_index,
|
|
nodes_module=nodes_module,
|
|
)
|
|
|
|
mapped_value = (source_prompt_node.get("inputs") or {}).get("value")
|
|
if not _is_link(mapped_value):
|
|
return boundary_input.io_type
|
|
|
|
return _remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=str(mapped_value[0]),
|
|
output_index=int(mapped_value[1]),
|
|
nodes_module=nodes_module,
|
|
)
|
|
|
|
|
|
def _describe_output_boundary_error(
|
|
component: RemoteComponentPlan,
|
|
source: LinkedOutputRef,
|
|
source_class_type: str,
|
|
io_type: str,
|
|
local_consumer: InputTarget,
|
|
local_consumer_class_type: str,
|
|
) -> str:
|
|
"""Format a human-readable component-output transport validation error."""
|
|
return (
|
|
"Remote component rooted at node "
|
|
f"{component.representative_node_id} exports node {source.node_id} "
|
|
f"({source_class_type}) output index {source.output_index} of type '{io_type}' "
|
|
f"to node {local_consumer.node_id} ({local_consumer_class_type}) outside that "
|
|
f"component at input '{local_consumer.input_name}', which cannot cross the "
|
|
"current component boundary. "
|
|
"Current ComfyUI-Modal transport only supports JSON-compatible values, bytes, "
|
|
"and media or tensor-like outputs such as VIDEO, AUDIO, IMAGE, MASK, LATENT, "
|
|
"SIGMAS, INT, FLOAT, BOOLEAN, and STRING."
|
|
)
|
|
|
|
|
|
def _describe_input_boundary_error(
|
|
component: RemoteComponentPlan,
|
|
target: InputTarget,
|
|
target_class_type: str,
|
|
source: LinkedOutputRef,
|
|
source_class_type: str,
|
|
io_type: str,
|
|
) -> str:
|
|
"""Format a human-readable component-input transport validation error."""
|
|
return (
|
|
"Remote node "
|
|
f"{target.node_id} ({target_class_type}) input '{target.input_name}' "
|
|
f"depends on node {source.node_id} ({source_class_type}) outside its component, "
|
|
f"using output index {source.output_index} of type '{io_type}', which cannot "
|
|
"cross the current component boundary. Current ComfyUI-Modal transport only "
|
|
"supports JSON-compatible values, bytes, and media or tensor-like outputs such "
|
|
"as VIDEO, AUDIO, IMAGE, MASK, LATENT, SIGMAS, INT, FLOAT, BOOLEAN, and STRING."
|
|
)
|
|
|
|
|
|
def validate_remote_component_transport_compatibility(
|
|
prompt: dict[str, Any],
|
|
components: list[RemoteComponentPlan],
|
|
nodes_module: Any,
|
|
) -> None:
|
|
"""Reject remote components whose true graph boundaries require unsupported transport."""
|
|
validation_errors: list[str] = []
|
|
consumers = _build_consumer_map(prompt)
|
|
logger.info(
|
|
"Validating %d remote components for transport compatibility.", len(components)
|
|
)
|
|
|
|
for component in components:
|
|
for boundary_input in component.boundary_inputs:
|
|
source_prompt_node = prompt.get(boundary_input.source.node_id)
|
|
if source_prompt_node is None:
|
|
continue
|
|
source_class_type = str(source_prompt_node["class_type"])
|
|
source_class = nodes_module.NODE_CLASS_MAPPINGS.get(source_class_type)
|
|
if source_class is None:
|
|
continue
|
|
|
|
source_output_types, _, _ = _normalize_output_metadata(source_class)
|
|
if boundary_input.source.output_index >= len(source_output_types):
|
|
continue
|
|
io_type = str(source_output_types[boundary_input.source.output_index])
|
|
if _is_transportable_output_type(io_type):
|
|
continue
|
|
|
|
for target in boundary_input.targets:
|
|
target_class_type = str(prompt[target.node_id]["class_type"])
|
|
validation_errors.append(
|
|
_describe_input_boundary_error(
|
|
component=component,
|
|
target=target,
|
|
target_class_type=target_class_type,
|
|
source=boundary_input.source,
|
|
source_class_type=source_class_type,
|
|
io_type=io_type,
|
|
)
|
|
)
|
|
|
|
for boundary_output in component.boundary_outputs:
|
|
if _is_transportable_output_type(boundary_output.io_type):
|
|
continue
|
|
|
|
source_class_type = str(
|
|
prompt[boundary_output.source.node_id]["class_type"]
|
|
)
|
|
for local_consumer in consumers.get(boundary_output.source, []):
|
|
if local_consumer.node_id in component.node_ids:
|
|
continue
|
|
local_consumer_class_type = str(
|
|
prompt[local_consumer.node_id]["class_type"]
|
|
)
|
|
validation_errors.append(
|
|
_describe_output_boundary_error(
|
|
component=component,
|
|
source=boundary_output.source,
|
|
source_class_type=source_class_type,
|
|
io_type=boundary_output.io_type,
|
|
local_consumer=local_consumer,
|
|
local_consumer_class_type=local_consumer_class_type,
|
|
)
|
|
)
|
|
|
|
if validation_errors:
|
|
raise ModalPromptValidationError("\n".join(validation_errors))
|
|
logger.info("Remote component transport validation passed.")
|
|
|