1610 lines
59 KiB
Python
1610 lines
59 KiB
Python
"""Workflow mapping, transportability, and remote graph partitioning."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import defaultdict, deque
|
|
from typing import Any, Callable, Iterator, Mapping
|
|
|
|
if __package__:
|
|
from .execution_environments import ExecutionProvider, WorkflowExecutionPreferences
|
|
from .llm_profiles import get_llm_profile
|
|
from .modal_executor_node import (
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
MODAL_MAP_INPUT_NODE_ID,
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
)
|
|
from .remote_configuration_nodes import REMOTE_CONFIGURATION_NODE_IDS
|
|
from .remote_plan_types import (
|
|
InputTarget,
|
|
LinkedOutputRef,
|
|
RemoteExpansionReason,
|
|
RemoteNodeAnalysis,
|
|
)
|
|
from .settings import ModalSyncSettings, get_settings
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from execution_environments import ExecutionProvider, WorkflowExecutionPreferences
|
|
from llm_profiles import get_llm_profile
|
|
from modal_executor_node import (
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
MODAL_MAP_INPUT_NODE_ID,
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
)
|
|
from remote_configuration_nodes import REMOTE_CONFIGURATION_NODE_IDS
|
|
from remote_plan_types import (
|
|
InputTarget,
|
|
LinkedOutputRef,
|
|
RemoteExpansionReason,
|
|
RemoteNodeAnalysis,
|
|
)
|
|
from settings import ModalSyncSettings, get_settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_REMOTE_TOGGLE_WIDGET_NAME = "Run Remotely"
|
|
_LEGACY_REMOTE_TOGGLE_WIDGET_NAME = "Run on Modal"
|
|
_LOCAL_ONLY_REMOTE_CLASS_TYPES = frozenset(
|
|
{
|
|
"ModalEndpointChat",
|
|
"VastAILeaseConfiguration",
|
|
*REMOTE_CONFIGURATION_NODE_IDS,
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
MODAL_MAP_INPUT_NODE_ID,
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
}
|
|
)
|
|
_TRANSPORTABLE_OUTPUT_TYPES = frozenset(
|
|
{
|
|
"*",
|
|
"AUDIO",
|
|
"BOOLEAN",
|
|
"FLOAT",
|
|
"IMAGE",
|
|
"INT",
|
|
"LATENT",
|
|
"MASK",
|
|
"SIGMAS",
|
|
"STRING",
|
|
"VIDEO",
|
|
}
|
|
)
|
|
_INEXPENSIVE_REMOTE_BOUNDARY_TYPES = frozenset(
|
|
{
|
|
"BOOLEAN",
|
|
"FLOAT",
|
|
"INT",
|
|
"STRING",
|
|
}
|
|
)
|
|
|
|
|
|
def _get_nodes_module() -> Any:
|
|
"""Import the ComfyUI nodes module lazily."""
|
|
import nodes
|
|
|
|
return nodes
|
|
|
|
|
|
def _is_link(value: Any) -> bool:
|
|
"""Return whether a prompt input value is a ComfyUI link."""
|
|
return (
|
|
isinstance(value, list)
|
|
and len(value) == 2
|
|
and all(not isinstance(item, dict) for item in value)
|
|
)
|
|
|
|
|
|
def _prompt_node_required_provider(
|
|
prompt_node: Mapping[str, Any],
|
|
) -> ExecutionProvider | None:
|
|
"""Return the provider required by one curated node backend, when known."""
|
|
if str(prompt_node.get("class_type") or "") != "ModalLLM":
|
|
return None
|
|
inputs = prompt_node.get("inputs")
|
|
if not isinstance(inputs, Mapping):
|
|
return None
|
|
model_reference = inputs.get("model_profile")
|
|
if not isinstance(model_reference, str) or not model_reference.strip():
|
|
return None
|
|
|
|
try:
|
|
profile = get_llm_profile(model_reference.strip())
|
|
except ValueError:
|
|
return None
|
|
if profile.backend == "llama_cpp_server":
|
|
return ExecutionProvider.SSH_DOCKER
|
|
return None
|
|
|
|
|
|
def _iter_payload_input_strings(value: Any) -> Iterator[str]:
|
|
"""Yield string literals nested inside one serialized prompt input value."""
|
|
if isinstance(value, str):
|
|
yield value
|
|
return
|
|
if isinstance(value, list):
|
|
if len(value) == 2 and isinstance(value[0], str):
|
|
return
|
|
for item in value:
|
|
yield from _iter_payload_input_strings(item)
|
|
return
|
|
if isinstance(value, dict):
|
|
for nested_value in value.values():
|
|
yield from _iter_payload_input_strings(nested_value)
|
|
|
|
|
|
def _looks_like_workflow_node(fragment: dict[str, Any]) -> bool:
|
|
"""Return whether a JSON fragment resembles a saved ComfyUI workflow node."""
|
|
return "id" in fragment and "properties" in fragment
|
|
|
|
|
|
def _iter_workflow_nodes(
|
|
workflow_fragment: Any,
|
|
visited_object_ids: set[int] | None = None,
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""Yield workflow node dictionaries from a nested saved workflow fragment."""
|
|
if visited_object_ids is None:
|
|
visited_object_ids = set()
|
|
|
|
if isinstance(workflow_fragment, dict):
|
|
object_id = id(workflow_fragment)
|
|
if object_id in visited_object_ids:
|
|
return
|
|
visited_object_ids.add(object_id)
|
|
|
|
if _looks_like_workflow_node(workflow_fragment):
|
|
yield workflow_fragment
|
|
|
|
for value in workflow_fragment.values():
|
|
yield from _iter_workflow_nodes(value, visited_object_ids)
|
|
return
|
|
|
|
if isinstance(workflow_fragment, list):
|
|
object_id = id(workflow_fragment)
|
|
if object_id in visited_object_ids:
|
|
return
|
|
visited_object_ids.add(object_id)
|
|
|
|
for item in workflow_fragment:
|
|
yield from _iter_workflow_nodes(item, visited_object_ids)
|
|
|
|
|
|
def _workflow_subgraph_definitions(
|
|
workflow_fragment: Any,
|
|
) -> dict[str, dict[str, Any]]:
|
|
"""Return reusable ComfyUI subgraph definitions keyed by their node type id."""
|
|
if not isinstance(workflow_fragment, dict):
|
|
return {}
|
|
definitions = workflow_fragment.get("definitions")
|
|
if not isinstance(definitions, dict):
|
|
return {}
|
|
subgraphs = definitions.get("subgraphs")
|
|
if not isinstance(subgraphs, list):
|
|
return {}
|
|
|
|
return {
|
|
str(subgraph["id"]): subgraph
|
|
for subgraph in subgraphs
|
|
if isinstance(subgraph, dict) and subgraph.get("id") is not None
|
|
}
|
|
|
|
|
|
def _iter_workflow_nodes_with_ancestors(
|
|
workflow_fragment: Any,
|
|
ancestor_node_ids: tuple[str, ...] = (),
|
|
visited_fragments: set[tuple[int, tuple[str, ...], tuple[str, ...]]] | None = None,
|
|
subgraph_definitions: Mapping[str, dict[str, Any]] | None = None,
|
|
active_subgraph_definition_ids: tuple[str, ...] = (),
|
|
) -> Iterator[tuple[dict[str, Any], tuple[str, ...]]]:
|
|
"""Yield workflow nodes with instance paths across embedded and defined subgraphs."""
|
|
if visited_fragments is None:
|
|
visited_fragments = set()
|
|
if subgraph_definitions is None:
|
|
subgraph_definitions = _workflow_subgraph_definitions(workflow_fragment)
|
|
|
|
if isinstance(workflow_fragment, dict):
|
|
fragment_identity = (
|
|
id(workflow_fragment),
|
|
ancestor_node_ids,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
if fragment_identity in visited_fragments:
|
|
return
|
|
visited_fragments.add(fragment_identity)
|
|
|
|
next_ancestor_node_ids = ancestor_node_ids
|
|
if _looks_like_workflow_node(workflow_fragment):
|
|
node_id = str(workflow_fragment.get("id"))
|
|
yield workflow_fragment, ancestor_node_ids
|
|
next_ancestor_node_ids = ancestor_node_ids + (node_id,)
|
|
|
|
subgraph_definition_id = str(workflow_fragment.get("type") or "")
|
|
subgraph_definition = subgraph_definitions.get(subgraph_definition_id)
|
|
if (
|
|
subgraph_definition is not None
|
|
and subgraph_definition_id not in active_subgraph_definition_ids
|
|
):
|
|
yield from _iter_workflow_nodes_with_ancestors(
|
|
subgraph_definition.get("nodes", []),
|
|
next_ancestor_node_ids,
|
|
visited_fragments,
|
|
subgraph_definitions,
|
|
active_subgraph_definition_ids + (subgraph_definition_id,),
|
|
)
|
|
|
|
for key, value in workflow_fragment.items():
|
|
if key == "definitions":
|
|
continue
|
|
yield from _iter_workflow_nodes_with_ancestors(
|
|
value,
|
|
next_ancestor_node_ids,
|
|
visited_fragments,
|
|
subgraph_definitions,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
return
|
|
|
|
if isinstance(workflow_fragment, list):
|
|
fragment_identity = (
|
|
id(workflow_fragment),
|
|
ancestor_node_ids,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
if fragment_identity in visited_fragments:
|
|
return
|
|
visited_fragments.add(fragment_identity)
|
|
|
|
for item in workflow_fragment:
|
|
yield from _iter_workflow_nodes_with_ancestors(
|
|
item,
|
|
ancestor_node_ids,
|
|
visited_fragments,
|
|
subgraph_definitions,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
|
|
|
|
def _resolve_prompt_node_ids_for_workflow_node(
|
|
workflow_node_id: str,
|
|
ancestor_node_ids: tuple[str, ...],
|
|
prompt_node_ids: set[str],
|
|
) -> set[str]:
|
|
"""Resolve one saved workflow node id to matching queued prompt node ids."""
|
|
path_segments = ancestor_node_ids + (workflow_node_id,)
|
|
|
|
for index in range(len(path_segments)):
|
|
candidate = ":".join(path_segments[index:])
|
|
resolved_prompt_node_ids: set[str] = set()
|
|
if candidate in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(candidate)
|
|
descendant_prefix = f"{candidate}:"
|
|
resolved_prompt_node_ids.update(
|
|
prompt_node_id
|
|
for prompt_node_id in prompt_node_ids
|
|
if prompt_node_id.startswith(descendant_prefix)
|
|
)
|
|
if resolved_prompt_node_ids:
|
|
return resolved_prompt_node_ids
|
|
|
|
resolved_prompt_node_ids: set[str] = set()
|
|
|
|
for ancestor_node_id in reversed(ancestor_node_ids):
|
|
if ancestor_node_id in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(ancestor_node_id)
|
|
descendant_prefix = f"{ancestor_node_id}:"
|
|
resolved_prompt_node_ids.update(
|
|
prompt_node_id
|
|
for prompt_node_id in prompt_node_ids
|
|
if prompt_node_id.startswith(descendant_prefix)
|
|
)
|
|
break
|
|
|
|
return resolved_prompt_node_ids
|
|
|
|
|
|
def _workflow_node_path(
|
|
workflow_node_id: str, ancestor_node_ids: tuple[str, ...]
|
|
) -> str:
|
|
"""Return one workflow node's composed path, including subgraph ancestors."""
|
|
if not ancestor_node_ids:
|
|
return workflow_node_id
|
|
return ":".join((*ancestor_node_ids, workflow_node_id))
|
|
|
|
|
|
def _extract_marked_workflow_node_paths(
|
|
workflow: dict[str, Any] | None,
|
|
settings: ModalSyncSettings | None = None,
|
|
) -> set[str]:
|
|
"""Return composed workflow paths for nodes explicitly marked remote in metadata."""
|
|
if workflow is None:
|
|
return set()
|
|
|
|
marker = (settings or get_settings()).marker_property
|
|
marked_workflow_node_paths: set[str] = set()
|
|
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
|
properties = node.get("properties") or {}
|
|
if not properties.get(marker):
|
|
continue
|
|
marked_workflow_node_paths.add(
|
|
_workflow_node_path(str(node.get("id")), ancestor_node_ids)
|
|
)
|
|
return marked_workflow_node_paths
|
|
|
|
|
|
def _build_workflow_prompt_resolution_maps(
|
|
workflow: dict[str, Any] | None,
|
|
prompt_node_ids: set[str],
|
|
) -> tuple[dict[str, set[str]], dict[str, set[str]]]:
|
|
"""Return bidirectional mappings between workflow paths and queued prompt ids."""
|
|
workflow_path_to_prompt_node_ids: dict[str, set[str]] = {}
|
|
prompt_node_id_to_workflow_paths: dict[str, set[str]] = defaultdict(set)
|
|
|
|
if workflow is None:
|
|
for prompt_node_id in prompt_node_ids:
|
|
workflow_path_to_prompt_node_ids[prompt_node_id] = {prompt_node_id}
|
|
prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id)
|
|
return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths
|
|
|
|
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
|
workflow_node_id = str(node.get("id"))
|
|
workflow_node_path = _workflow_node_path(workflow_node_id, ancestor_node_ids)
|
|
resolved_prompt_node_ids = _resolve_prompt_node_ids_for_workflow_node(
|
|
workflow_node_id,
|
|
ancestor_node_ids,
|
|
prompt_node_ids,
|
|
)
|
|
if not ancestor_node_ids and workflow_node_id in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(workflow_node_id)
|
|
if workflow_node_path in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(workflow_node_path)
|
|
if not resolved_prompt_node_ids:
|
|
continue
|
|
|
|
workflow_path_to_prompt_node_ids[workflow_node_path] = resolved_prompt_node_ids
|
|
for prompt_node_id in resolved_prompt_node_ids:
|
|
prompt_node_id_to_workflow_paths[prompt_node_id].add(workflow_node_path)
|
|
|
|
for prompt_node_id in prompt_node_ids:
|
|
workflow_path_to_prompt_node_ids.setdefault(prompt_node_id, {prompt_node_id})
|
|
prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id)
|
|
|
|
return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths
|
|
|
|
|
|
def _resolve_requested_prompt_node_ids(
|
|
requested_workflow_node_paths: set[str],
|
|
prompt_node_ids: set[str],
|
|
workflow_path_to_prompt_node_ids: dict[str, set[str]],
|
|
) -> set[str]:
|
|
"""Resolve requested workflow-node paths to the prompt ids that queue-time rewrite sees."""
|
|
requested_prompt_node_ids: set[str] = set()
|
|
for requested_workflow_node_path in requested_workflow_node_paths:
|
|
if requested_workflow_node_path in workflow_path_to_prompt_node_ids:
|
|
requested_prompt_node_ids.update(
|
|
workflow_path_to_prompt_node_ids[requested_workflow_node_path]
|
|
)
|
|
continue
|
|
if requested_workflow_node_path in prompt_node_ids:
|
|
requested_prompt_node_ids.add(requested_workflow_node_path)
|
|
return requested_prompt_node_ids
|
|
|
|
|
|
def _best_workflow_path_for_prompt_node(
|
|
prompt_node_id: str,
|
|
prompt_node_id_to_workflow_paths: dict[str, set[str]],
|
|
) -> str:
|
|
"""Choose the most specific workflow path for one queued prompt node id."""
|
|
workflow_node_paths = prompt_node_id_to_workflow_paths.get(
|
|
prompt_node_id, {prompt_node_id}
|
|
)
|
|
return max(
|
|
workflow_node_paths,
|
|
key=lambda workflow_node_path: (
|
|
workflow_node_path.count(":"),
|
|
len(workflow_node_path),
|
|
workflow_node_path,
|
|
),
|
|
)
|
|
|
|
|
|
def _resolve_workflow_node_paths_for_prompt_nodes(
|
|
prompt_node_ids: set[str],
|
|
prompt_node_id_to_workflow_paths: dict[str, set[str]],
|
|
) -> set[str]:
|
|
"""Map queued prompt ids back to the workflow node paths the UI can mark remote."""
|
|
return {
|
|
_best_workflow_path_for_prompt_node(
|
|
prompt_node_id, prompt_node_id_to_workflow_paths
|
|
)
|
|
for prompt_node_id in prompt_node_ids
|
|
}
|
|
|
|
|
|
def extract_remote_node_ids(
|
|
workflow: dict[str, Any] | None,
|
|
settings: ModalSyncSettings | None = None,
|
|
prompt_node_ids: set[str] | None = None,
|
|
) -> set[str]:
|
|
"""Return the node ids marked for remote execution in the workflow metadata."""
|
|
if workflow is None:
|
|
return set()
|
|
|
|
marker = (settings or get_settings()).marker_property
|
|
remote_node_ids: set[str] = set()
|
|
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
|
if _workflow_node_remote_enabled(node, marker):
|
|
node_id = str(node.get("id"))
|
|
if prompt_node_ids is None:
|
|
remote_node_ids.add(node_id)
|
|
continue
|
|
|
|
resolved_prompt_node_ids = _resolve_prompt_node_ids_for_workflow_node(
|
|
node_id,
|
|
ancestor_node_ids,
|
|
prompt_node_ids,
|
|
)
|
|
if resolved_prompt_node_ids:
|
|
remote_node_ids.update(resolved_prompt_node_ids)
|
|
logger.info(
|
|
"Mapped workflow Modal marker from node %s with ancestors %s to prompt nodes %s.",
|
|
node_id,
|
|
list(ancestor_node_ids),
|
|
sorted(resolved_prompt_node_ids),
|
|
)
|
|
return remote_node_ids
|
|
|
|
|
|
def requested_remote_node_ids(
|
|
*,
|
|
prompt: Mapping[str, Any],
|
|
workflow: Mapping[str, Any] | None,
|
|
settings: ModalSyncSettings,
|
|
) -> set[str]:
|
|
"""Return explicit remote markers or every eligible node for auto placement."""
|
|
preferences = WorkflowExecutionPreferences.from_workflow(workflow)
|
|
if not preferences.auto_place:
|
|
return extract_remote_node_ids(
|
|
dict(workflow) if workflow is not None else None,
|
|
settings,
|
|
set(prompt),
|
|
)
|
|
selected_node_ids = {
|
|
str(node_id)
|
|
for node_id, prompt_node in prompt.items()
|
|
if isinstance(prompt_node, Mapping)
|
|
and str(prompt_node.get("class_type") or "")
|
|
not in _LOCAL_ONLY_REMOTE_CLASS_TYPES
|
|
and not str(prompt_node.get("class_type") or "").startswith(
|
|
"ModalUniversalExecutor"
|
|
)
|
|
}
|
|
logger.info(
|
|
"Automatic remote placement selected %d eligible prompt nodes.",
|
|
len(selected_node_ids),
|
|
)
|
|
return selected_node_ids
|
|
|
|
|
|
def _workflow_node_remote_enabled(node: Mapping[str, Any], marker: str) -> bool:
|
|
"""Return the visible remote toggle, falling back to the saved marker property."""
|
|
named_widget_values = node.get("widgets_values_named")
|
|
if isinstance(named_widget_values, Mapping):
|
|
visible_toggle_value = named_widget_values.get(_REMOTE_TOGGLE_WIDGET_NAME)
|
|
if not isinstance(visible_toggle_value, bool):
|
|
visible_toggle_value = named_widget_values.get(
|
|
_LEGACY_REMOTE_TOGGLE_WIDGET_NAME
|
|
)
|
|
if isinstance(visible_toggle_value, bool):
|
|
return visible_toggle_value
|
|
|
|
properties = node.get("properties")
|
|
return isinstance(properties, Mapping) and bool(properties.get(marker))
|
|
|
|
|
|
def _normalize_output_metadata(
|
|
node_class: type[Any],
|
|
) -> tuple[tuple[str, ...], tuple[str, ...], tuple[bool, ...]]:
|
|
"""Return normalized output metadata for a node class."""
|
|
if hasattr(node_class, "GET_SCHEMA"):
|
|
node_class.GET_SCHEMA()
|
|
|
|
output_types = tuple(getattr(node_class, "RETURN_TYPES", ("*",))) or ("*",)
|
|
default_names = tuple(f"output_{index}" for index, _ in enumerate(output_types))
|
|
output_names = tuple(getattr(node_class, "RETURN_NAMES", default_names))
|
|
output_is_list = tuple(
|
|
getattr(node_class, "OUTPUT_IS_LIST", (False,) * len(output_types))
|
|
)
|
|
|
|
if len(output_names) < len(output_types):
|
|
output_names = output_names + default_names[len(output_names) :]
|
|
if len(output_is_list) < len(output_types):
|
|
output_is_list = output_is_list + (False,) * (
|
|
len(output_types) - len(output_is_list)
|
|
)
|
|
|
|
return (
|
|
output_types,
|
|
output_names[: len(output_types)],
|
|
output_is_list[: len(output_types)],
|
|
)
|
|
|
|
|
|
def _is_transportable_output_type(io_type: str) -> bool:
|
|
"""Return whether a ComfyUI output type can cross the current transport."""
|
|
normalized_parts = [part.strip() for part in io_type.split(",") if part.strip()]
|
|
return bool(normalized_parts) and all(
|
|
part in _TRANSPORTABLE_OUTPUT_TYPES for part in normalized_parts
|
|
)
|
|
|
|
|
|
def _is_inexpensive_remote_boundary_type(io_type: str) -> bool:
|
|
"""Return whether a remote edge is cheap enough to remain a component boundary."""
|
|
normalized_parts = [part.strip() for part in io_type.split(",") if part.strip()]
|
|
return bool(normalized_parts) and all(
|
|
part in _INEXPENSIVE_REMOTE_BOUNDARY_TYPES for part in normalized_parts
|
|
)
|
|
|
|
|
|
def _build_consumer_map(
|
|
prompt: dict[str, Any]
|
|
) -> dict[LinkedOutputRef, list[InputTarget]]:
|
|
"""Build a reverse map from node outputs to downstream prompt inputs."""
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]] = defaultdict(list)
|
|
for node_id, prompt_node in prompt.items():
|
|
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])
|
|
)
|
|
consumers[source].append(
|
|
InputTarget(node_id=str(node_id), input_name=str(input_name))
|
|
)
|
|
return consumers
|
|
|
|
|
|
def _sandwiched_local_node_ids(
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
) -> set[str]:
|
|
"""Return local nodes on paths that leave and then re-enter remote execution."""
|
|
downstream_node_ids: dict[str, set[str]] = defaultdict(set)
|
|
upstream_node_ids: dict[str, set[str]] = defaultdict(set)
|
|
for target_node_id, prompt_node in prompt.items():
|
|
for input_value in (prompt_node.get("inputs") or {}).values():
|
|
if not _is_link(input_value):
|
|
continue
|
|
source_node_id = str(input_value[0])
|
|
safe_target_node_id = str(target_node_id)
|
|
downstream_node_ids[source_node_id].add(safe_target_node_id)
|
|
upstream_node_ids[safe_target_node_id].add(source_node_id)
|
|
|
|
def local_closure(adjacency: Mapping[str, set[str]]) -> set[str]:
|
|
"""Return local nodes reachable from the remote seed set in one direction."""
|
|
reachable_local_node_ids: set[str] = set()
|
|
pending_node_ids = list(sorted(remote_node_ids))
|
|
visited_node_ids = set(remote_node_ids)
|
|
while pending_node_ids:
|
|
current_node_id = pending_node_ids.pop()
|
|
for adjacent_node_id in adjacency.get(current_node_id, set()):
|
|
if adjacent_node_id in visited_node_ids:
|
|
continue
|
|
visited_node_ids.add(adjacent_node_id)
|
|
if adjacent_node_id in remote_node_ids:
|
|
continue
|
|
if adjacent_node_id not in prompt:
|
|
continue
|
|
reachable_local_node_ids.add(adjacent_node_id)
|
|
pending_node_ids.append(adjacent_node_id)
|
|
return reachable_local_node_ids
|
|
|
|
downstream_local_node_ids = local_closure(downstream_node_ids)
|
|
upstream_local_node_ids = local_closure(upstream_node_ids)
|
|
return downstream_local_node_ids & upstream_local_node_ids
|
|
|
|
|
|
def _node_output_refs(
|
|
prompt: dict[str, Any], node_id: str, nodes_module: Any
|
|
) -> list[LinkedOutputRef]:
|
|
"""Return declared output refs for one prompt node."""
|
|
prompt_node = prompt.get(node_id)
|
|
if prompt_node is None:
|
|
return []
|
|
node_class = nodes_module.NODE_CLASS_MAPPINGS.get(
|
|
str(prompt_node.get("class_type"))
|
|
)
|
|
if node_class is None:
|
|
return []
|
|
output_types, _, _ = _normalize_output_metadata(node_class)
|
|
return [
|
|
LinkedOutputRef(node_id=node_id, output_index=output_index)
|
|
for output_index, _io_type in enumerate(output_types)
|
|
]
|
|
|
|
|
|
def _downstream_node_ids_from_targets(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
seed_targets: list[InputTarget],
|
|
nodes_module: Any,
|
|
) -> set[str]:
|
|
"""Return every node reachable downstream from the supplied input targets."""
|
|
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 prompt:
|
|
continue
|
|
visited_node_ids.add(node_id)
|
|
for output_ref in _node_output_refs(prompt, node_id, nodes_module):
|
|
for downstream_target in consumers.get(output_ref, []):
|
|
downstream_node_id = str(downstream_target.node_id)
|
|
if downstream_node_id not in visited_node_ids:
|
|
pending_node_ids.append(downstream_node_id)
|
|
return visited_node_ids
|
|
|
|
|
|
def _is_non_returning_tap_terminal_node(
|
|
*,
|
|
prompt_node: dict[str, Any] | None,
|
|
nodes_module: Any,
|
|
) -> bool:
|
|
"""Return whether a local tap node should stay local instead of being absorbed."""
|
|
if prompt_node is None:
|
|
return False
|
|
class_type = str(prompt_node.get("class_type"))
|
|
if class_type == "PreviewImage":
|
|
return True
|
|
node_class = nodes_module.NODE_CLASS_MAPPINGS.get(class_type)
|
|
return bool(getattr(node_class, "OUTPUT_NODE", False))
|
|
|
|
|
|
def _non_returning_local_tap_node_ids(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
component_node_ids: set[str],
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
) -> tuple[set[str], set[str]]:
|
|
"""Return non-returning local tap branches without making local nodes remote."""
|
|
local_tap_node_ids: set[str] = set()
|
|
local_tap_terminal_node_ids: set[str] = set()
|
|
for component_node_id in sorted(component_node_ids):
|
|
for source in _node_output_refs(prompt, component_node_id, nodes_module):
|
|
output_consumers = consumers.get(source, [])
|
|
local_consumers = [
|
|
target
|
|
for target in output_consumers
|
|
if target.node_id not in component_node_ids
|
|
and target.node_id not in remote_node_ids
|
|
]
|
|
remote_consumers = [
|
|
target
|
|
for target in output_consumers
|
|
if target.node_id in remote_node_ids
|
|
]
|
|
if not local_consumers or not remote_consumers:
|
|
continue
|
|
for local_consumer in local_consumers:
|
|
branch_node_ids = _downstream_node_ids_from_targets(
|
|
prompt=prompt,
|
|
consumers=consumers,
|
|
seed_targets=[local_consumer],
|
|
nodes_module=nodes_module,
|
|
)
|
|
if branch_node_ids & remote_node_ids:
|
|
continue
|
|
local_tap_terminal_node_ids.update(branch_node_ids)
|
|
return local_tap_node_ids, local_tap_terminal_node_ids
|
|
|
|
|
|
def _non_returning_local_output_consumers(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
source: LinkedOutputRef,
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
) -> tuple[list[InputTarget], list[InputTarget]]:
|
|
"""Partition local consumers by whether their branch later returns remote."""
|
|
non_returning_consumers: list[InputTarget] = []
|
|
returning_consumers: list[InputTarget] = []
|
|
for target in consumers.get(source, []):
|
|
if target.node_id in remote_node_ids:
|
|
continue
|
|
branch_node_ids = _downstream_node_ids_from_targets(
|
|
prompt=prompt,
|
|
consumers=consumers,
|
|
seed_targets=[target],
|
|
nodes_module=nodes_module,
|
|
)
|
|
if branch_node_ids and not (branch_node_ids & remote_node_ids):
|
|
non_returning_consumers.append(target)
|
|
else:
|
|
returning_consumers.append(target)
|
|
return non_returning_consumers, returning_consumers
|
|
|
|
|
|
def _output_supports_parallel_local_materialization(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
source: LinkedOutputRef,
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
) -> bool:
|
|
"""Return whether every local consumer is independent of later remote work."""
|
|
(
|
|
non_returning_consumers,
|
|
returning_consumers,
|
|
) = _non_returning_local_output_consumers(
|
|
prompt=prompt,
|
|
source=source,
|
|
remote_node_ids=remote_node_ids,
|
|
consumers=consumers,
|
|
nodes_module=nodes_module,
|
|
)
|
|
return bool(non_returning_consumers) and not returning_consumers
|
|
|
|
|
|
def _expand_component_for_non_transportable_local_outputs(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
component_node_ids: set[str],
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
) -> tuple[set[str], set[str]]:
|
|
"""Absorb local consumers needed to avoid non-transportable component outputs."""
|
|
expanded_node_ids = set(component_node_ids)
|
|
added_node_ids: set[str] = set()
|
|
changed = True
|
|
while changed:
|
|
changed = False
|
|
(
|
|
upstream_expanded_node_ids,
|
|
_expansion_reasons,
|
|
) = _expand_remote_node_ids_for_non_transportable_inputs(
|
|
prompt=prompt,
|
|
remote_node_ids=expanded_node_ids,
|
|
nodes_module=nodes_module,
|
|
)
|
|
upstream_added_node_ids = upstream_expanded_node_ids - expanded_node_ids
|
|
if upstream_added_node_ids:
|
|
expanded_node_ids = upstream_expanded_node_ids
|
|
added_node_ids.update(upstream_added_node_ids)
|
|
changed = True
|
|
|
|
for node_id in sorted(expanded_node_ids):
|
|
prompt_node = prompt.get(node_id)
|
|
if prompt_node is None:
|
|
continue
|
|
node_class = nodes_module.NODE_CLASS_MAPPINGS.get(
|
|
str(prompt_node.get("class_type"))
|
|
)
|
|
if node_class is None:
|
|
continue
|
|
output_types, _, _ = _normalize_output_metadata(node_class)
|
|
for output_index, io_type in enumerate(output_types):
|
|
if _is_transportable_output_type(str(io_type)):
|
|
continue
|
|
source = LinkedOutputRef(node_id=node_id, output_index=output_index)
|
|
for target in consumers.get(source, []):
|
|
target_node_id = str(target.node_id)
|
|
if (
|
|
target_node_id in expanded_node_ids
|
|
or target_node_id in remote_node_ids
|
|
):
|
|
continue
|
|
if target_node_id not in prompt:
|
|
continue
|
|
if _is_non_returning_tap_terminal_node(
|
|
prompt_node=prompt.get(target_node_id),
|
|
nodes_module=nodes_module,
|
|
):
|
|
continue
|
|
expanded_node_ids.add(target_node_id)
|
|
added_node_ids.add(target_node_id)
|
|
changed = True
|
|
logger.info(
|
|
"Absorbing local node %s because it consumes non-transportable output %s:%d from a preview-tap remote component.",
|
|
target_node_id,
|
|
source.node_id,
|
|
source.output_index,
|
|
)
|
|
return expanded_node_ids, added_node_ids
|
|
|
|
|
|
def _remote_output_io_type(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
node_id: str,
|
|
output_index: int,
|
|
nodes_module: Any,
|
|
) -> str | None:
|
|
"""Return the declared output type for one prompt node output when available."""
|
|
prompt_node = prompt.get(node_id)
|
|
if prompt_node is None:
|
|
return None
|
|
|
|
class_type = str(prompt_node["class_type"])
|
|
node_class = nodes_module.NODE_CLASS_MAPPINGS.get(class_type)
|
|
if node_class is None:
|
|
return None
|
|
|
|
output_types, _, _ = _normalize_output_metadata(node_class)
|
|
if output_index < 0 or output_index >= len(output_types):
|
|
return None
|
|
return str(output_types[output_index])
|
|
|
|
|
|
def _remote_output_is_list(
|
|
*,
|
|
prompt: dict[str, Any],
|
|
node_id: str,
|
|
output_index: int,
|
|
nodes_module: Any,
|
|
) -> bool:
|
|
"""Return whether one prompt node output is declared as a list output."""
|
|
prompt_node = prompt.get(node_id)
|
|
if prompt_node is None:
|
|
return False
|
|
|
|
class_type = str(prompt_node["class_type"])
|
|
node_class = nodes_module.NODE_CLASS_MAPPINGS.get(class_type)
|
|
if node_class is None:
|
|
return False
|
|
|
|
output_types, _output_names, output_is_list = _normalize_output_metadata(node_class)
|
|
if output_index < 0 or output_index >= len(output_types):
|
|
return False
|
|
return bool(output_is_list[output_index])
|
|
|
|
|
|
def _remote_group_dependency_edges(
|
|
prompt: Mapping[str, Any],
|
|
remote_node_ids: set[str],
|
|
group_for_node: Callable[[str], str],
|
|
) -> dict[str, set[str]]:
|
|
"""Return current coarse dependency edges for a remote-node grouping."""
|
|
dependency_edges: dict[str, set[str]] = defaultdict(set)
|
|
for downstream_node_id in remote_node_ids:
|
|
downstream_node = prompt.get(downstream_node_id)
|
|
if not isinstance(downstream_node, Mapping):
|
|
continue
|
|
downstream_group = group_for_node(downstream_node_id)
|
|
for input_value in (downstream_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 remote_node_ids:
|
|
continue
|
|
upstream_group = group_for_node(upstream_node_id)
|
|
if upstream_group != downstream_group:
|
|
dependency_edges[upstream_group].add(downstream_group)
|
|
return dependency_edges
|
|
|
|
|
|
def _has_alternate_group_path_through_protected_group(
|
|
dependency_edges: Mapping[str, set[str]],
|
|
start_group: str,
|
|
target_group: str,
|
|
protected_groups: set[str],
|
|
) -> bool:
|
|
"""Return whether a non-direct path crosses a provider-constrained group."""
|
|
pending = [
|
|
(downstream_group, downstream_group in protected_groups)
|
|
for downstream_group in dependency_edges.get(start_group, set())
|
|
if downstream_group != target_group
|
|
]
|
|
visited: set[tuple[str, bool]] = set()
|
|
while pending:
|
|
current_group, crossed_protected_group = pending.pop()
|
|
if current_group == target_group and crossed_protected_group:
|
|
return True
|
|
state = (current_group, crossed_protected_group)
|
|
if state in visited:
|
|
continue
|
|
visited.add(state)
|
|
pending.extend(
|
|
(
|
|
downstream_group,
|
|
crossed_protected_group or downstream_group in protected_groups,
|
|
)
|
|
for downstream_group in dependency_edges.get(current_group, set())
|
|
)
|
|
return False
|
|
|
|
|
|
def _remote_component_partition_groups(
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
*,
|
|
preserve_nontransportable_affinity: bool = True,
|
|
) -> dict[str, set[str]]:
|
|
"""Return component groups after merging remote nodes across costly boundaries."""
|
|
parent: dict[str, str] = {node_id: node_id for node_id in remote_node_ids}
|
|
downstream_remote_node_ids_by_node_id: dict[str, set[str]] = defaultdict(set)
|
|
provider_constrained_node_ids = {
|
|
node_id
|
|
for node_id in remote_node_ids
|
|
if isinstance((prompt_node := prompt.get(node_id)), Mapping)
|
|
and _prompt_node_required_provider(prompt_node) is not None
|
|
}
|
|
|
|
def find(node_id: str) -> str:
|
|
"""Return the canonical union-find representative for one remote node."""
|
|
root = parent[node_id]
|
|
while root != parent[root]:
|
|
root = parent[root]
|
|
while node_id != root:
|
|
next_node_id = parent[node_id]
|
|
parent[node_id] = root
|
|
node_id = next_node_id
|
|
return root
|
|
|
|
def union(
|
|
left_node_id: str,
|
|
right_node_id: str,
|
|
*,
|
|
preserve_coarse_dag: bool | None = None,
|
|
) -> bool:
|
|
"""Merge two remote nodes unless doing so creates a coarse dependency cycle."""
|
|
left_root = find(left_node_id)
|
|
right_root = find(right_node_id)
|
|
if left_root == right_root:
|
|
return True
|
|
should_preserve_coarse_dag = (
|
|
bool(provider_constrained_node_ids)
|
|
if preserve_coarse_dag is None
|
|
else preserve_coarse_dag
|
|
)
|
|
if should_preserve_coarse_dag:
|
|
dependency_edges = _remote_group_dependency_edges(
|
|
prompt,
|
|
remote_node_ids,
|
|
find,
|
|
)
|
|
protected_groups = {
|
|
find(node_id) for node_id in provider_constrained_node_ids
|
|
}
|
|
creates_cycle = _has_alternate_group_path_through_protected_group(
|
|
dependency_edges,
|
|
left_root,
|
|
right_root,
|
|
protected_groups,
|
|
) or _has_alternate_group_path_through_protected_group(
|
|
dependency_edges,
|
|
right_root,
|
|
left_root,
|
|
protected_groups,
|
|
)
|
|
if creates_cycle:
|
|
logger.info(
|
|
"Splitting remote nodes %s and %s to preserve an acyclic "
|
|
"component graph.",
|
|
left_node_id,
|
|
right_node_id,
|
|
)
|
|
return False
|
|
canonical_root = min(left_root, right_root)
|
|
merged_root = max(left_root, right_root)
|
|
parent[merged_root] = canonical_root
|
|
return True
|
|
|
|
for node_id in sorted(remote_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 remote_node_ids:
|
|
continue
|
|
downstream_remote_node_ids_by_node_id[upstream_node_id].add(node_id)
|
|
upstream_prompt_node = prompt.get(upstream_node_id)
|
|
source = LinkedOutputRef(
|
|
node_id=upstream_node_id, output_index=int(input_value[1])
|
|
)
|
|
if (
|
|
upstream_prompt_node is not None
|
|
and str(upstream_prompt_node.get("class_type"))
|
|
== MODAL_MAP_INPUT_NODE_ID
|
|
):
|
|
union(node_id, upstream_node_id, preserve_coarse_dag=False)
|
|
continue
|
|
io_type = _remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=upstream_node_id,
|
|
output_index=source.output_index,
|
|
nodes_module=nodes_module,
|
|
)
|
|
if (
|
|
preserve_nontransportable_affinity
|
|
and io_type is not None
|
|
and not _is_transportable_output_type(io_type)
|
|
):
|
|
logger.info(
|
|
"Co-locating remote nodes %s -> %s because %s output %d "
|
|
"cannot cross a component boundary.",
|
|
upstream_node_id,
|
|
node_id,
|
|
io_type,
|
|
source.output_index,
|
|
)
|
|
union(node_id, upstream_node_id, preserve_coarse_dag=False)
|
|
continue
|
|
downstream_provider = _prompt_node_required_provider(prompt_node)
|
|
upstream_provider = (
|
|
_prompt_node_required_provider(upstream_prompt_node)
|
|
if isinstance(upstream_prompt_node, Mapping)
|
|
else None
|
|
)
|
|
if downstream_provider != upstream_provider and (
|
|
downstream_provider is not None or upstream_provider is not None
|
|
):
|
|
logger.info(
|
|
"Splitting remote nodes %s -> %s across provider boundary %s -> %s.",
|
|
upstream_node_id,
|
|
node_id,
|
|
upstream_provider.value if upstream_provider is not None else "any",
|
|
downstream_provider.value
|
|
if downstream_provider is not None
|
|
else "any",
|
|
)
|
|
continue
|
|
if _output_supports_parallel_local_materialization(
|
|
prompt=prompt,
|
|
source=source,
|
|
remote_node_ids=remote_node_ids,
|
|
consumers=consumers,
|
|
nodes_module=nodes_module,
|
|
):
|
|
logger.info(
|
|
"Splitting remote nodes %s -> %s because output %d also feeds only non-returning local work.",
|
|
upstream_node_id,
|
|
node_id,
|
|
source.output_index,
|
|
)
|
|
continue
|
|
if io_type is not None and _is_inexpensive_remote_boundary_type(io_type):
|
|
continue
|
|
if io_type is not None and _is_transportable_output_type(io_type):
|
|
logger.info(
|
|
"Co-locating remote nodes %s -> %s because %s output %d is "
|
|
"transportable but expensive.",
|
|
upstream_node_id,
|
|
node_id,
|
|
io_type,
|
|
source.output_index,
|
|
)
|
|
union(node_id, upstream_node_id)
|
|
|
|
for remote_node_id in sorted(remote_node_ids):
|
|
prompt_node = prompt.get(remote_node_id)
|
|
if (
|
|
prompt_node is None
|
|
or str(prompt_node.get("class_type")) != MODAL_MAP_INPUT_NODE_ID
|
|
):
|
|
continue
|
|
pending_node_ids = [remote_node_id]
|
|
visited_node_ids: set[str] = set()
|
|
while pending_node_ids:
|
|
current_node_id = pending_node_ids.pop()
|
|
if current_node_id in visited_node_ids:
|
|
continue
|
|
visited_node_ids.add(current_node_id)
|
|
union(
|
|
remote_node_id,
|
|
current_node_id,
|
|
preserve_coarse_dag=False,
|
|
)
|
|
for downstream_node_id in sorted(
|
|
downstream_remote_node_ids_by_node_id.get(current_node_id, set())
|
|
):
|
|
pending_node_ids.append(downstream_node_id)
|
|
|
|
groups: dict[str, set[str]] = defaultdict(set)
|
|
for node_id in sorted(remote_node_ids):
|
|
groups[find(node_id)].add(node_id)
|
|
return groups
|
|
|
|
|
|
def _component_topological_order(
|
|
prompt: dict[str, Any],
|
|
component_groups: dict[str, set[str]],
|
|
) -> list[list[str]]:
|
|
"""Return component node ids ordered from upstream to downstream."""
|
|
_, dependency_edges, indegree_by_component_id = _component_dependency_graph(
|
|
prompt, component_groups
|
|
)
|
|
merged_component_groups = _merge_cyclic_component_groups(
|
|
component_groups=component_groups,
|
|
dependency_edges=dependency_edges,
|
|
)
|
|
if merged_component_groups != component_groups:
|
|
return _component_topological_order(prompt, merged_component_groups)
|
|
|
|
ready_component_ids = deque(
|
|
sorted(
|
|
[
|
|
component_id
|
|
for component_id, indegree in indegree_by_component_id.items()
|
|
if indegree == 0
|
|
]
|
|
)
|
|
)
|
|
ordered_components: list[list[str]] = []
|
|
emitted_component_ids: set[str] = set()
|
|
|
|
while ready_component_ids:
|
|
component_id = ready_component_ids.popleft()
|
|
if component_id in emitted_component_ids:
|
|
continue
|
|
emitted_component_ids.add(component_id)
|
|
ordered_components.append(sorted(component_groups[component_id]))
|
|
for downstream_component_id in sorted(dependency_edges[component_id]):
|
|
indegree_by_component_id[downstream_component_id] -= 1
|
|
if indegree_by_component_id[downstream_component_id] == 0:
|
|
ready_component_ids.append(downstream_component_id)
|
|
|
|
if len(emitted_component_ids) == len(component_groups):
|
|
return ordered_components
|
|
|
|
logger.warning(
|
|
"Transport-aware component ordering encountered a cycle or unresolved dependency; falling back to stable component order."
|
|
)
|
|
for component_id in sorted(component_groups):
|
|
if component_id in emitted_component_ids:
|
|
continue
|
|
ordered_components.append(sorted(component_groups[component_id]))
|
|
return ordered_components
|
|
|
|
|
|
def _component_dependency_graph(
|
|
prompt: dict[str, Any],
|
|
component_groups: dict[str, set[str]],
|
|
) -> tuple[dict[str, str], dict[str, set[str]], dict[str, int]]:
|
|
"""Return component membership, downstream edges, and indegrees for the coarse component DAG."""
|
|
component_id_by_node_id: dict[str, str] = {}
|
|
for representative_node_id, component_node_ids in component_groups.items():
|
|
for node_id in component_node_ids:
|
|
component_id_by_node_id[node_id] = representative_node_id
|
|
|
|
dependency_edges: dict[str, set[str]] = {
|
|
representative_node_id: set() for representative_node_id in component_groups
|
|
}
|
|
indegree_by_component_id: dict[str, int] = {
|
|
representative_node_id: 0 for representative_node_id in component_groups
|
|
}
|
|
|
|
for node_id, representative_node_id in component_id_by_node_id.items():
|
|
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])
|
|
upstream_component_id = component_id_by_node_id.get(upstream_node_id)
|
|
if (
|
|
upstream_component_id is None
|
|
or upstream_component_id == representative_node_id
|
|
):
|
|
continue
|
|
if representative_node_id in dependency_edges[upstream_component_id]:
|
|
continue
|
|
dependency_edges[upstream_component_id].add(representative_node_id)
|
|
indegree_by_component_id[representative_node_id] += 1
|
|
|
|
return component_id_by_node_id, dependency_edges, indegree_by_component_id
|
|
|
|
|
|
def _component_execution_stages(
|
|
prompt: dict[str, Any],
|
|
component_groups: dict[str, set[str]],
|
|
) -> list[list[str]]:
|
|
"""Return one best-effort stage decomposition for concurrent remote component execution."""
|
|
_, dependency_edges, indegree_by_component_id = _component_dependency_graph(
|
|
prompt, component_groups
|
|
)
|
|
merged_component_groups = _merge_cyclic_component_groups(
|
|
component_groups=component_groups,
|
|
dependency_edges=dependency_edges,
|
|
)
|
|
if merged_component_groups != component_groups:
|
|
return _component_execution_stages(prompt, merged_component_groups)
|
|
|
|
remaining_indegrees = dict(indegree_by_component_id)
|
|
ready_component_ids = sorted(
|
|
[
|
|
component_id
|
|
for component_id, indegree in remaining_indegrees.items()
|
|
if indegree == 0
|
|
]
|
|
)
|
|
execution_stages: list[list[str]] = []
|
|
emitted_component_ids: set[str] = set()
|
|
|
|
while ready_component_ids:
|
|
current_stage = list(ready_component_ids)
|
|
execution_stages.append(current_stage)
|
|
next_ready_component_ids: set[str] = set()
|
|
for component_id in current_stage:
|
|
emitted_component_ids.add(component_id)
|
|
for downstream_component_id in sorted(
|
|
dependency_edges.get(component_id, set())
|
|
):
|
|
remaining_indegrees[downstream_component_id] -= 1
|
|
if remaining_indegrees[downstream_component_id] == 0:
|
|
next_ready_component_ids.add(downstream_component_id)
|
|
ready_component_ids = sorted(
|
|
component_id
|
|
for component_id in next_ready_component_ids
|
|
if component_id not in emitted_component_ids
|
|
)
|
|
|
|
if len(emitted_component_ids) == len(component_groups):
|
|
return execution_stages
|
|
|
|
fallback_stage = [
|
|
component_id
|
|
for component_id in sorted(component_groups)
|
|
if component_id not in emitted_component_ids
|
|
]
|
|
if fallback_stage:
|
|
logger.warning(
|
|
"Transport-aware execution-stage planning encountered a cycle or unresolved dependency; appending fallback stage %s.",
|
|
fallback_stage,
|
|
)
|
|
execution_stages.append(fallback_stage)
|
|
return execution_stages
|
|
|
|
|
|
def _estimated_stage_parallelism(
|
|
execution_stages: list[list[str]],
|
|
mapped_component_ids: set[str],
|
|
*,
|
|
mapped_component_weight: int,
|
|
max_parallelism_cap: int | None = None,
|
|
) -> int:
|
|
"""Return the weighted best-effort parallelism estimate over staged remote execution."""
|
|
stage_parallelism = 0
|
|
for stage in execution_stages:
|
|
current_stage_parallelism = sum(
|
|
mapped_component_weight if component_id in mapped_component_ids else 1
|
|
for component_id in stage
|
|
)
|
|
stage_parallelism = max(stage_parallelism, current_stage_parallelism)
|
|
if max_parallelism_cap is not None:
|
|
return min(stage_parallelism, max_parallelism_cap)
|
|
return stage_parallelism
|
|
|
|
|
|
def _merge_cyclic_component_groups(
|
|
*,
|
|
component_groups: dict[str, set[str]],
|
|
dependency_edges: dict[str, set[str]],
|
|
) -> dict[str, set[str]]:
|
|
"""Collapse cyclic coarse component groups into SCC-merged groups."""
|
|
component_ids = sorted(component_groups)
|
|
reverse_edges: dict[str, set[str]] = {
|
|
component_id: set() for component_id in component_ids
|
|
}
|
|
for upstream_component_id, downstream_component_ids in dependency_edges.items():
|
|
for downstream_component_id in downstream_component_ids:
|
|
reverse_edges.setdefault(downstream_component_id, set()).add(
|
|
upstream_component_id
|
|
)
|
|
|
|
visited_component_ids: set[str] = set()
|
|
finish_order: list[str] = []
|
|
|
|
def visit_forward(component_id: str) -> None:
|
|
"""Record reverse-topological finish order over the coarse graph."""
|
|
if component_id in visited_component_ids:
|
|
return
|
|
visited_component_ids.add(component_id)
|
|
for downstream_component_id in sorted(
|
|
dependency_edges.get(component_id, set())
|
|
):
|
|
visit_forward(downstream_component_id)
|
|
finish_order.append(component_id)
|
|
|
|
for component_id in component_ids:
|
|
visit_forward(component_id)
|
|
|
|
assigned_component_ids: set[str] = set()
|
|
merged_groups: dict[str, set[str]] = {}
|
|
merged_sccs = 0
|
|
|
|
def visit_reverse(component_id: str, scc_component_ids: set[str]) -> None:
|
|
"""Collect one SCC by walking reverse edges from the finish-order seed."""
|
|
if component_id in assigned_component_ids:
|
|
return
|
|
assigned_component_ids.add(component_id)
|
|
scc_component_ids.add(component_id)
|
|
for upstream_component_id in sorted(reverse_edges.get(component_id, set())):
|
|
visit_reverse(upstream_component_id, scc_component_ids)
|
|
|
|
for component_id in reversed(finish_order):
|
|
if component_id in assigned_component_ids:
|
|
continue
|
|
scc_component_ids: set[str] = set()
|
|
visit_reverse(component_id, scc_component_ids)
|
|
merged_node_ids: set[str] = set()
|
|
for scc_component_id in scc_component_ids:
|
|
merged_node_ids.update(component_groups[scc_component_id])
|
|
representative_node_id = min(merged_node_ids)
|
|
merged_groups[representative_node_id] = merged_node_ids
|
|
if len(scc_component_ids) > 1:
|
|
merged_sccs += 1
|
|
|
|
if merged_sccs:
|
|
logger.warning(
|
|
"Transport-aware coarse component graph contained %d cyclic SCC(s); merging them back into larger remote components.",
|
|
merged_sccs,
|
|
)
|
|
return merged_groups
|
|
|
|
|
|
def _build_remote_components(
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
) -> list[list[str]]:
|
|
"""Partition remote-marked nodes into transport-aware DAG components."""
|
|
component_groups = _remote_component_partition_groups(
|
|
prompt,
|
|
remote_node_ids,
|
|
consumers,
|
|
nodes_module,
|
|
)
|
|
components = _component_topological_order(prompt, component_groups)
|
|
logger.info(
|
|
"Partitioned %d remote nodes into %d transport-aware remote components: %s",
|
|
len(remote_node_ids),
|
|
len(components),
|
|
components,
|
|
)
|
|
return components
|
|
|
|
|
|
def _expand_remote_node_ids_for_non_transportable_inputs(
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
nodes_module: Any,
|
|
) -> tuple[set[str], list[RemoteExpansionReason]]:
|
|
"""Grow the remote set upstream until non-transportable inputs stay inside the remote island."""
|
|
expanded_remote_node_ids = set(remote_node_ids)
|
|
added_node_ids: set[str] = set()
|
|
reasons: list[RemoteExpansionReason] = []
|
|
|
|
changed = True
|
|
while changed:
|
|
changed = False
|
|
for node_id in sorted(expanded_remote_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 in expanded_remote_node_ids:
|
|
continue
|
|
|
|
upstream_prompt_node = prompt.get(upstream_node_id)
|
|
if upstream_prompt_node is None:
|
|
continue
|
|
|
|
upstream_class_type = str(upstream_prompt_node["class_type"])
|
|
upstream_class = nodes_module.NODE_CLASS_MAPPINGS.get(
|
|
upstream_class_type
|
|
)
|
|
if upstream_class is None:
|
|
continue
|
|
|
|
output_types, _, _ = _normalize_output_metadata(upstream_class)
|
|
output_index = int(input_value[1])
|
|
if output_index >= len(output_types):
|
|
continue
|
|
|
|
io_type = str(output_types[output_index])
|
|
if _is_transportable_output_type(io_type):
|
|
continue
|
|
|
|
expanded_remote_node_ids.add(upstream_node_id)
|
|
reason = RemoteExpansionReason(
|
|
node_id=upstream_node_id,
|
|
class_type=upstream_class_type,
|
|
required_by_node_id=node_id,
|
|
required_by_class_type=str(prompt_node["class_type"]),
|
|
output_index=output_index,
|
|
io_type=io_type,
|
|
)
|
|
reasons.append(reason)
|
|
added_node_ids.add(upstream_node_id)
|
|
changed = True
|
|
logger.info(
|
|
"Auto-expanded remote execution upstream: added node %s (%s) because node %s depends on non-transportable type '%s'.",
|
|
reason.node_id,
|
|
reason.class_type,
|
|
reason.required_by_node_id,
|
|
reason.io_type,
|
|
)
|
|
break
|
|
if changed:
|
|
break
|
|
|
|
if added_node_ids:
|
|
logger.info(
|
|
"Expanded remote node set from %d to %d nodes by absorbing upstream non-transportable dependencies: %s",
|
|
len(remote_node_ids),
|
|
len(expanded_remote_node_ids),
|
|
sorted(added_node_ids),
|
|
)
|
|
return expanded_remote_node_ids, reasons
|
|
|
|
|
|
def _terminal_remote_video_source(
|
|
*,
|
|
node_id: str,
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
consumers: dict[LinkedOutputRef, list[InputTarget]],
|
|
nodes_module: Any,
|
|
) -> LinkedOutputRef | None:
|
|
"""Return the remote VIDEO feeding one safe terminal SaveVideo artifact sink."""
|
|
prompt_node = prompt.get(node_id)
|
|
if prompt_node is None or str(prompt_node.get("class_type")) != "SaveVideo":
|
|
return None
|
|
inputs = prompt_node.get("inputs") or {}
|
|
video_input = inputs.get("video")
|
|
if not _is_link(video_input):
|
|
return None
|
|
video_source = LinkedOutputRef(str(video_input[0]), int(video_input[1]))
|
|
linked_source_node_ids = {
|
|
str(input_value[0]) for input_value in inputs.values() if _is_link(input_value)
|
|
}
|
|
if not linked_source_node_ids.issubset(remote_node_ids):
|
|
return None
|
|
if any(
|
|
consumers.get(output_ref)
|
|
for output_ref in _node_output_refs(prompt, node_id, nodes_module)
|
|
):
|
|
return None
|
|
io_type = _remote_output_io_type(
|
|
prompt=prompt,
|
|
node_id=video_source.node_id,
|
|
output_index=video_source.output_index,
|
|
nodes_module=nodes_module,
|
|
)
|
|
return video_source if io_type == "VIDEO" else None
|
|
|
|
|
|
def _expand_remote_node_ids_for_terminal_video_sinks(
|
|
prompt: dict[str, Any],
|
|
remote_node_ids: set[str],
|
|
nodes_module: Any,
|
|
) -> set[str]:
|
|
"""Run terminal SaveVideo consumers remotely to avoid exporting raw VIDEO tensors."""
|
|
expanded_remote_node_ids = set(remote_node_ids)
|
|
consumers = _build_consumer_map(prompt)
|
|
added_node_ids: set[str] = set()
|
|
for node_id in sorted(prompt):
|
|
if node_id in expanded_remote_node_ids:
|
|
continue
|
|
video_source = _terminal_remote_video_source(
|
|
node_id=node_id,
|
|
prompt=prompt,
|
|
remote_node_ids=expanded_remote_node_ids,
|
|
consumers=consumers,
|
|
nodes_module=nodes_module,
|
|
)
|
|
if video_source is None:
|
|
continue
|
|
expanded_remote_node_ids.add(node_id)
|
|
added_node_ids.add(node_id)
|
|
logger.info(
|
|
"Auto-expanded terminal SaveVideo node %s into remote execution because its "
|
|
"VIDEO input comes from remote node %s; the encoded output artifact will be "
|
|
"materialized locally without exporting raw frames.",
|
|
node_id,
|
|
video_source.node_id,
|
|
)
|
|
|
|
if added_node_ids:
|
|
logger.info(
|
|
"Expanded remote node set with terminal video artifact sinks: %s.",
|
|
sorted(added_node_ids),
|
|
)
|
|
return expanded_remote_node_ids
|
|
|
|
|
|
def analyze_remote_node_selection(
|
|
prompt: dict[str, Any],
|
|
workflow: dict[str, Any] | None,
|
|
seed_workflow_node_paths: list[str],
|
|
settings: ModalSyncSettings | None = None,
|
|
nodes_module: Any | None = None,
|
|
) -> RemoteNodeAnalysis:
|
|
"""Return the nodes the UI should mark remote for one context-menu expansion request."""
|
|
resolved_settings = settings or get_settings()
|
|
resolved_nodes_module = nodes_module or _get_nodes_module()
|
|
prompt_node_ids = {str(node_id) for node_id in prompt.keys()}
|
|
requested_workflow_node_paths = {
|
|
str(seed_workflow_node_path)
|
|
for seed_workflow_node_path in seed_workflow_node_paths
|
|
if str(seed_workflow_node_path)
|
|
}
|
|
current_remote_node_ids = extract_remote_node_ids(
|
|
workflow,
|
|
resolved_settings,
|
|
prompt_node_ids,
|
|
)
|
|
current_remote_workflow_node_paths = _extract_marked_workflow_node_paths(
|
|
workflow,
|
|
resolved_settings,
|
|
)
|
|
(
|
|
workflow_path_to_prompt_node_ids,
|
|
prompt_node_id_to_workflow_paths,
|
|
) = _build_workflow_prompt_resolution_maps(workflow, prompt_node_ids)
|
|
requested_node_ids = _resolve_requested_prompt_node_ids(
|
|
requested_workflow_node_paths,
|
|
prompt_node_ids,
|
|
workflow_path_to_prompt_node_ids,
|
|
)
|
|
initial_remote_node_ids = current_remote_node_ids | requested_node_ids
|
|
(
|
|
resolved_remote_node_ids,
|
|
reasons,
|
|
) = _expand_remote_node_ids_for_non_transportable_inputs(
|
|
prompt=prompt,
|
|
remote_node_ids=initial_remote_node_ids,
|
|
nodes_module=resolved_nodes_module,
|
|
)
|
|
resolved_remote_node_ids = _expand_remote_node_ids_for_terminal_video_sinks(
|
|
prompt=prompt,
|
|
remote_node_ids=resolved_remote_node_ids,
|
|
nodes_module=resolved_nodes_module,
|
|
)
|
|
sandwiched_local_node_ids = _sandwiched_local_node_ids(
|
|
prompt,
|
|
resolved_remote_node_ids,
|
|
)
|
|
resolved_workflow_node_paths = (
|
|
_resolve_workflow_node_paths_for_prompt_nodes(
|
|
resolved_remote_node_ids,
|
|
prompt_node_id_to_workflow_paths,
|
|
)
|
|
| current_remote_workflow_node_paths
|
|
)
|
|
added_node_ids = resolved_remote_node_ids - current_remote_node_ids
|
|
added_workflow_node_paths = (
|
|
resolved_workflow_node_paths - current_remote_workflow_node_paths
|
|
)
|
|
|
|
return RemoteNodeAnalysis(
|
|
requested_node_ids=sorted(requested_node_ids),
|
|
requested_workflow_node_paths=sorted(requested_workflow_node_paths),
|
|
current_remote_node_ids=sorted(current_remote_node_ids),
|
|
current_remote_workflow_node_paths=sorted(current_remote_workflow_node_paths),
|
|
resolved_remote_node_ids=sorted(resolved_remote_node_ids),
|
|
resolved_workflow_node_paths=sorted(resolved_workflow_node_paths),
|
|
added_node_ids=sorted(added_node_ids),
|
|
added_workflow_node_paths=sorted(added_workflow_node_paths),
|
|
sandwiched_local_node_ids=sorted(sandwiched_local_node_ids),
|
|
reasons=reasons,
|
|
)
|
|
|