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