Split cloud execution runtime responsibilities
This commit is contained in:
Ken Simpson
@@ -12,12 +12,16 @@
|
||||
!/cloud_app_guard.py
|
||||
!/cloud_comfy_bootstrap.py
|
||||
!/cloud_durable_invocation.py
|
||||
!/cloud_execution_control.py
|
||||
!/cloud_image_env.py
|
||||
!/cloud_mapped_execution.py
|
||||
!/cloud_node_output_cache.py
|
||||
!/cloud_prompt_execution.py
|
||||
!/cloud_prompt_validation.py
|
||||
!/cloud_prompt_server_shims.py
|
||||
!/cloud_prewarm.py
|
||||
!/cloud_runtime_context.py
|
||||
!/cloud_runtime_logging.py
|
||||
!/cloud_session_bridge.py
|
||||
!/cloud_streaming.py
|
||||
!/cloud_volume_reload.py
|
||||
|
||||
@@ -516,7 +516,7 @@ COMFYUI_ROOT=/tmp/comfyui-modal-test/ComfyUI \
|
||||
|
||||
For full local development including the remote and Apple-local extras, use `uv sync --extra remote --extra local-apple --group test`.
|
||||
|
||||
The prompt interception layer is organized as a compatibility aggregator over focused modules: planning values, graph analysis, component planning, prompt rewriting, execution scheduling, queue lifecycle, diagnostics, and HTTP route domains each have a dedicated owner. Route handlers share an immutable `RouteContext`, while `api_intercept.py` retains only registration and compatibility re-exports. Host-side remote execution follows the same ownership model: deployment loading, runtime compatibility checks, and auto-deploy state live in `remote/modal_deployment.py`; proactive warmup, snapshot profiles, and local-gap keepalive state live in `remote/modal_warmup.py`; mapped or implicit-batch lane scheduling lives in `remote/mapped_execution.py`; and CPU ModelStager state lives in `remote/modal_llm_profile_staging.py`. Adjacent invocation and staging callbacks are injected by the `remote/modal_app.py` orchestrator. Blocking Modal executors are constructed lazily on first use, and importing the host orchestrator no longer constructs a vestigial Modal app or volume. Cloud helpers receive tail-created Modal dictionaries and volumes through `cloud_runtime_context.py`, allowing focused cloud subsystems to move without reaching back into entry-module globals; pre-deploy app collision checks live in `cloud_app_guard.py`, image filters and worker construction policy live in `cloud_image_env.py`, durable invocation replay and live-canary coordination live in `cloud_durable_invocation.py`, session-bridge storage and replay live in `cloud_session_bridge.py`, headless ComfyUI server compatibility plus progress tracing live in `cloud_prompt_server_shims.py`, ComfyUI initialization, custom-node loading, asset path translation, compatibility patches, and warm loader caches live in `cloud_comfy_bootstrap.py`, distributed node-output key construction, serialization, restore, and write-back live in `cloud_node_output_cache.py`, prompt validation, reusable executor management, local node execution, and mapped/phased subgraph execution live in `cloud_prompt_execution.py`, bounded progress/result delivery plus abandoned invocation cleanup live in `cloud_streaming.py`, mounted-volume visibility and reload recovery live in `cloud_volume_reload.py`, and snapshot, loader, LLM, and request prewarming live in `cloud_prewarm.py`. Root modules support both package imports used by ComfyUI and flat imports used by the Modal container; `tests/test_import_shapes.py` guards both layouts and the stable `comfyui_modal_sync_cloud` surface.
|
||||
The prompt interception layer is organized as a compatibility aggregator over focused modules: planning values, graph analysis, component planning, prompt rewriting, execution scheduling, queue lifecycle, diagnostics, and HTTP route domains each have a dedicated owner. Route handlers share an immutable `RouteContext`, while `api_intercept.py` retains only registration and compatibility re-exports. Host-side remote execution follows the same ownership model: deployment loading, runtime compatibility checks, and auto-deploy state live in `remote/modal_deployment.py`; proactive warmup, snapshot profiles, and local-gap keepalive state live in `remote/modal_warmup.py`; mapped or implicit-batch lane scheduling lives in `remote/mapped_execution.py`; and CPU ModelStager state lives in `remote/modal_llm_profile_staging.py`. Adjacent invocation and staging callbacks are injected by the `remote/modal_app.py` orchestrator. Blocking Modal executors are constructed lazily on first use, and importing the host orchestrator no longer constructs a vestigial Modal app or volume. Cloud helpers receive tail-created Modal dictionaries and volumes through `cloud_runtime_context.py`, active request cleanup through `cloud_execution_control.py`, and timestamped phase reporting through `cloud_runtime_logging.py`, allowing focused cloud subsystems to move without reaching back into entry-module globals; pre-deploy app collision checks live in `cloud_app_guard.py`, image filters and worker construction policy live in `cloud_image_env.py`, durable invocation replay and live-canary coordination live in `cloud_durable_invocation.py`, session-bridge storage and replay live in `cloud_session_bridge.py`, headless ComfyUI server compatibility plus progress tracing live in `cloud_prompt_server_shims.py`, ComfyUI initialization, custom-node loading, asset path translation, compatibility patches, and warm loader caches live in `cloud_comfy_bootstrap.py`, distributed node-output key construction, serialization, restore, and write-back live in `cloud_node_output_cache.py`, prompt schema validation and normalization live in `cloud_prompt_validation.py`, reusable executor management and local subgraph execution live in `cloud_prompt_execution.py`, mapped/phased cloud execution lives in `cloud_mapped_execution.py`, bounded progress/result delivery plus abandoned invocation cleanup live in `cloud_streaming.py`, mounted-volume visibility and reload recovery live in `cloud_volume_reload.py`, and snapshot, loader, LLM, and request prewarming live in `cloud_prewarm.py`. Root modules support both package imports used by ComfyUI and flat imports used by the Modal container; `tests/test_import_shapes.py` guards both layouts and the stable `comfyui_modal_sync_cloud` surface.
|
||||
|
||||
Resident LLM inference uses the same ownership boundaries: `llm_types.py` defines backend-neutral values, `llm_inputs.py` normalizes multimodal and file inputs, `vllm_instrumentation.py` owns process-wide vLLM execution policy and Triton telemetry, and `llm_backend_transformers.py`, `llm_backend_llamacpp.py`, and `llm_backend_vllm.py` isolate backend implementations. `modal_llm_runtime.py` remains the model-residency coordinator, while `local_llm_runtime.py` consumes shared types and input preparation without importing private helpers from the Modal backend. CPU-side staging adapts the general cross-process ownership lock in `snapshot_lease.py` to LLM-specific progress events.
|
||||
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Active cloud execution registration and interrupt cleanup."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Iterator
|
||||
|
||||
try:
|
||||
import modal # type: ignore
|
||||
except ModuleNotFoundError: # pragma: no cover - remote entrypoint only.
|
||||
modal = None
|
||||
|
||||
try:
|
||||
from .cloud_runtime_context import interrupt_flag_store
|
||||
except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
from cloud_runtime_context import interrupt_flag_store
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RemoteExecutionControl:
|
||||
"""Track cancellation state for one active remote execution."""
|
||||
|
||||
cancellation_event: threading.Event
|
||||
interrupt_flag_key: str
|
||||
|
||||
|
||||
def _remote_execution_key(payload: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Return the registry key for one active remote execution."""
|
||||
prompt_id = str(
|
||||
payload.get("prompt_id") or payload.get("component_id") or "modal-subgraph"
|
||||
)
|
||||
component_id = str(payload.get("component_id") or "single-node")
|
||||
return prompt_id, component_id
|
||||
|
||||
|
||||
def _remote_interrupt_flag_key(prompt_id: str, component_id: str) -> str:
|
||||
"""Return the shared Modal interrupt-store key for one payload execution."""
|
||||
return f"{prompt_id}:{component_id}"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _registered_remote_execution(
|
||||
payload: dict[str, Any],
|
||||
) -> Iterator[_RemoteExecutionControl]:
|
||||
"""Prepare interruption state for one active remote execution."""
|
||||
prompt_id, component_id = _remote_execution_key(payload)
|
||||
control = _RemoteExecutionControl(
|
||||
cancellation_event=threading.Event(),
|
||||
interrupt_flag_key=_remote_interrupt_flag_key(prompt_id, component_id),
|
||||
)
|
||||
try:
|
||||
yield control
|
||||
finally:
|
||||
interrupt_store = interrupt_flag_store()
|
||||
if modal is not None and interrupt_store is not None:
|
||||
interrupt_store.pop(control.interrupt_flag_key, None)
|
||||
@@ -0,0 +1,381 @@
|
||||
"""Mapped and phased execution for the cloud worker runtime."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import threading
|
||||
from typing import Any, Callable
|
||||
|
||||
try:
|
||||
from . import cloud_prompt_execution as _prompt_execution
|
||||
from .serialization import (
|
||||
join_mapped_values_for_scheduler,
|
||||
split_mapped_value,
|
||||
)
|
||||
except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
import cloud_prompt_execution as _prompt_execution
|
||||
from serialization import (
|
||||
join_mapped_values_for_scheduler,
|
||||
split_mapped_value,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _execute_prompt_subgraph(
|
||||
payload: dict[str, Any],
|
||||
hydrated_inputs: dict[str, Any],
|
||||
custom_nodes_root: Path | None,
|
||||
status_callback: Callable[[dict[str, Any]], None] | None = None,
|
||||
cancellation_event: threading.Event | None = None,
|
||||
interrupt_store: Any | None = None,
|
||||
interrupt_flag_key: str | None = None,
|
||||
) -> tuple[Any, ...]:
|
||||
"""Execute one phase through the cloud prompt-execution owner."""
|
||||
return _prompt_execution._execute_subgraph_prompt(
|
||||
payload,
|
||||
hydrated_inputs,
|
||||
custom_nodes_root,
|
||||
status_callback=status_callback,
|
||||
cancellation_event=cancellation_event,
|
||||
interrupt_store=interrupt_store,
|
||||
interrupt_flag_key=interrupt_flag_key,
|
||||
)
|
||||
|
||||
|
||||
def _remote_subgraph_execution_error(message: str) -> RuntimeError:
|
||||
"""Construct the stable cloud entrypoint's configured subgraph error."""
|
||||
return _prompt_execution.RemoteSubgraphExecutionError(message)
|
||||
|
||||
|
||||
def _mapped_phase_definition(
|
||||
payload: dict[str, Any], phase_key: str
|
||||
) -> dict[str, Any] | None:
|
||||
"""Return one explicit mapped phase definition when queue-time planning provided it."""
|
||||
phase_payload = payload.get(phase_key)
|
||||
if isinstance(phase_payload, dict):
|
||||
return phase_payload
|
||||
return None
|
||||
|
||||
|
||||
def _shared_subgraph_payload_fields(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return the payload fields shared by every explicit mapped phase."""
|
||||
shared_fields = {
|
||||
"prompt_id": payload.get("prompt_id"),
|
||||
"extra_data": copy.deepcopy(payload.get("extra_data") or {}),
|
||||
"requires_volume_reload": bool(payload.get("requires_volume_reload", True)),
|
||||
"volume_reload_marker": payload.get("volume_reload_marker"),
|
||||
"uploaded_volume_paths": list(payload.get("uploaded_volume_paths", [])),
|
||||
"terminate_container_on_error": bool(
|
||||
payload.get("terminate_container_on_error", True)
|
||||
),
|
||||
"custom_nodes_bundle": payload.get("custom_nodes_bundle"),
|
||||
}
|
||||
snapshot_profile_key = payload.get("snapshot_profile_key")
|
||||
if isinstance(snapshot_profile_key, str) and snapshot_profile_key.strip():
|
||||
shared_fields["snapshot_profile_key"] = snapshot_profile_key.strip()
|
||||
remote_session = payload.get("remote_session")
|
||||
if remote_session is not None:
|
||||
shared_fields["remote_session"] = copy.deepcopy(remote_session)
|
||||
if bool(payload.get("clear_remote_session")):
|
||||
shared_fields["clear_remote_session"] = True
|
||||
return shared_fields
|
||||
|
||||
|
||||
def _build_phase_subgraph_payload(
|
||||
payload: dict[str, Any],
|
||||
phase_key: str,
|
||||
component_id: str,
|
||||
) -> 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}.")
|
||||
|
||||
return {
|
||||
"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)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
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 _aggregate_mapped_phase_outputs(
|
||||
per_item_outputs: list[tuple[Any, ...]],
|
||||
payload: dict[str, Any],
|
||||
) -> tuple[Any, ...]:
|
||||
"""Join ordered mapped-phase outputs back 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 _remote_subgraph_execution_error(
|
||||
"Mapped remote execution produced inconsistent output arity."
|
||||
)
|
||||
|
||||
aggregated_outputs: list[Any] = []
|
||||
boundary_outputs = list(payload.get("boundary_outputs", []))
|
||||
for output_index in range(output_count):
|
||||
boundary_output = (
|
||||
boundary_outputs[output_index]
|
||||
if output_index < len(boundary_outputs)
|
||||
else {}
|
||||
)
|
||||
aggregated_outputs.append(
|
||||
_merge_static_or_mapped_values(
|
||||
[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 _merge_static_or_mapped_values(
|
||||
values: list[Any],
|
||||
*,
|
||||
io_type: str,
|
||||
is_list: bool,
|
||||
scheduler_is_list: bool,
|
||||
) -> Any:
|
||||
"""Join mapped per-item outputs using the shared transport serializer rules."""
|
||||
return join_mapped_values_for_scheduler(
|
||||
values,
|
||||
io_type=io_type,
|
||||
is_list=is_list,
|
||||
scheduler_is_list=scheduler_is_list,
|
||||
)
|
||||
|
||||
|
||||
def _merge_static_and_mapped_outputs(
|
||||
*,
|
||||
static_outputs: tuple[Any, ...],
|
||||
mapped_outputs: tuple[Any, ...],
|
||||
payload: dict[str, Any],
|
||||
) -> tuple[Any, ...]:
|
||||
"""Reassemble one mapped component's static and mapped outputs in proxy order."""
|
||||
combined_outputs: list[Any] = []
|
||||
static_output_index = 0
|
||||
mapped_output_index = 0
|
||||
for boundary_output in payload.get("boundary_outputs", []):
|
||||
if bool(boundary_output.get("mapped_output")):
|
||||
if mapped_output_index >= len(mapped_outputs):
|
||||
raise _remote_subgraph_execution_error(
|
||||
"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 _remote_subgraph_execution_error(
|
||||
"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 _remote_subgraph_execution_error(
|
||||
"Mapped remote execution produced extra outputs that did not match the declared boundary outputs."
|
||||
)
|
||||
return tuple(combined_outputs)
|
||||
|
||||
|
||||
def _execute_mapped_subgraph_payload(
|
||||
payload: dict[str, Any],
|
||||
hydrated_inputs: dict[str, Any],
|
||||
custom_nodes_root: Path | None,
|
||||
status_callback: Callable[[dict[str, Any]], None] | None = None,
|
||||
cancellation_event: threading.Event | None = None,
|
||||
interrupt_store: Any | None = None,
|
||||
interrupt_flag_key: str | None = None,
|
||||
) -> tuple[Any, ...]:
|
||||
"""Execute one mapped payload inside a single remote runtime process."""
|
||||
mapped_input = payload.get("mapped_input") or {}
|
||||
mapped_input_name = str(mapped_input.get("proxy_input_name") or "")
|
||||
if not mapped_input_name:
|
||||
raise _remote_subgraph_execution_error(
|
||||
"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, ...] = ()
|
||||
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",
|
||||
)
|
||||
if static_phase_payload.get("execute_node_ids"):
|
||||
static_phase_outputs = _execute_prompt_subgraph(
|
||||
static_phase_payload,
|
||||
dict(broadcast_inputs),
|
||||
custom_nodes_root,
|
||||
status_callback=status_callback,
|
||||
cancellation_event=cancellation_event,
|
||||
interrupt_store=interrupt_store,
|
||||
interrupt_flag_key=interrupt_flag_key,
|
||||
)
|
||||
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)
|
||||
|
||||
if status_callback is not None:
|
||||
status_callback(
|
||||
{
|
||||
"event_type": "node_progress",
|
||||
"node_id": str(payload.get("component_id") or "modal-subgraph"),
|
||||
"display_node_id": str(payload.get("component_id") or "modal-subgraph"),
|
||||
"value": 0.0,
|
||||
"max": float(len(mapped_items)),
|
||||
"aggregate_only": True,
|
||||
}
|
||||
)
|
||||
|
||||
per_item_outputs: list[tuple[Any, ...]] = []
|
||||
for item_index, item_value in enumerate(mapped_items):
|
||||
last_lane_node_id: str | None = None
|
||||
lane_id = str(payload.get("mapped_progress_lane_id") or item_index)
|
||||
|
||||
def publish_item_status(progress_state: dict[str, Any]) -> None:
|
||||
"""Attach mapped-lane metadata to one per-item progress event."""
|
||||
nonlocal last_lane_node_id
|
||||
if status_callback is None:
|
||||
return
|
||||
event_type = str(progress_state.get("event_type", ""))
|
||||
if event_type == "node_progress":
|
||||
reported_node_id = progress_state.get(
|
||||
"real_node_id"
|
||||
) or progress_state.get("node_id")
|
||||
if reported_node_id is not None:
|
||||
last_lane_node_id = str(reported_node_id)
|
||||
status_callback(
|
||||
{
|
||||
**progress_state,
|
||||
"lane_id": lane_id,
|
||||
"item_index": item_index,
|
||||
}
|
||||
)
|
||||
return
|
||||
if event_type in {"executed", "preview", "boundary_output"}:
|
||||
status_callback({**progress_state, "item_index": item_index})
|
||||
|
||||
item_payload = _build_phase_subgraph_payload(
|
||||
payload,
|
||||
"mapped_phase",
|
||||
f"{payload.get('component_id', 'modal-subgraph')}::item:{item_index}",
|
||||
)
|
||||
item_inputs = dict(broadcast_inputs)
|
||||
item_inputs[mapped_input_name] = item_value
|
||||
per_item_outputs.append(
|
||||
_execute_prompt_subgraph(
|
||||
item_payload,
|
||||
item_inputs,
|
||||
custom_nodes_root,
|
||||
status_callback=publish_item_status,
|
||||
cancellation_event=cancellation_event,
|
||||
interrupt_store=interrupt_store,
|
||||
interrupt_flag_key=interrupt_flag_key,
|
||||
)
|
||||
)
|
||||
if status_callback is not None:
|
||||
status_callback(
|
||||
{
|
||||
"event_type": "node_progress",
|
||||
"node_id": last_lane_node_id
|
||||
or str(payload.get("component_id") or "modal-subgraph"),
|
||||
"display_node_id": last_lane_node_id
|
||||
or str(payload.get("component_id") or "modal-subgraph"),
|
||||
"value": 0.0,
|
||||
"max": 1.0,
|
||||
"lane_id": lane_id,
|
||||
"item_index": item_index,
|
||||
"clear": True,
|
||||
}
|
||||
)
|
||||
status_callback(
|
||||
{
|
||||
"event_type": "node_progress",
|
||||
"node_id": str(payload.get("component_id") or "modal-subgraph"),
|
||||
"display_node_id": str(
|
||||
payload.get("component_id") or "modal-subgraph"
|
||||
),
|
||||
"value": float(item_index + 1),
|
||||
"max": float(len(mapped_items)),
|
||||
"aggregate_only": True,
|
||||
}
|
||||
)
|
||||
|
||||
mapped_phase_payload = _build_phase_subgraph_payload(
|
||||
payload,
|
||||
"mapped_phase",
|
||||
f"{payload.get('component_id', 'modal-subgraph')}::mapped",
|
||||
)
|
||||
mapped_outputs = _aggregate_mapped_phase_outputs(
|
||||
per_item_outputs,
|
||||
{"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,
|
||||
)
|
||||
|
||||
+42
-840
@@ -36,6 +36,27 @@ try:
|
||||
_prompt_executor_cache_get_sync,
|
||||
)
|
||||
from .cloud_prompt_server_shims import _NullPromptServer, _TracingPromptServer
|
||||
from .cloud_prompt_validation import (
|
||||
_coerce_primitive_prompt_input_value,
|
||||
_coerce_prompt_primitive_input_values,
|
||||
_extract_prompt_executor_error,
|
||||
_extract_prompt_executor_error_payload,
|
||||
_format_prompt_executor_error_payload,
|
||||
_is_link,
|
||||
_log_prompt_executor_failure_details,
|
||||
_node_input_type_map,
|
||||
_node_input_types,
|
||||
_node_required_input_names,
|
||||
_normalize_link_output_index,
|
||||
_normalize_prompt_input_value,
|
||||
_normalize_subgraph_payload,
|
||||
_resolve_required_subgraph_nodes,
|
||||
_summarize_suspicious_prompt_inputs,
|
||||
_trim_subgraph_payload_to_required_nodes,
|
||||
_unwrap_wrapped_prompt_link,
|
||||
_validate_prompt_input_shapes,
|
||||
_validate_required_prompt_inputs,
|
||||
)
|
||||
from .cloud_session_bridge import (
|
||||
_RemoteSessionBridgeResolutionStats,
|
||||
_build_remote_session_bridge_record,
|
||||
@@ -48,7 +69,6 @@ try:
|
||||
)
|
||||
from .remote_protocol import (
|
||||
BOUNDARY_INPUT_SIGNATURES_KEY as _BOUNDARY_INPUT_SIGNATURES_KEY,
|
||||
PRIMITIVE_WIDGET_INPUT_TYPES as _PRIMITIVE_WIDGET_INPUT_TYPES,
|
||||
)
|
||||
from .serialization import (
|
||||
coerce_serialized_node_outputs,
|
||||
@@ -87,6 +107,27 @@ except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
_prompt_executor_cache_get_sync,
|
||||
)
|
||||
from cloud_prompt_server_shims import _NullPromptServer, _TracingPromptServer
|
||||
from cloud_prompt_validation import (
|
||||
_coerce_primitive_prompt_input_value,
|
||||
_coerce_prompt_primitive_input_values,
|
||||
_extract_prompt_executor_error,
|
||||
_extract_prompt_executor_error_payload,
|
||||
_format_prompt_executor_error_payload,
|
||||
_is_link,
|
||||
_log_prompt_executor_failure_details,
|
||||
_node_input_type_map,
|
||||
_node_input_types,
|
||||
_node_required_input_names,
|
||||
_normalize_link_output_index,
|
||||
_normalize_prompt_input_value,
|
||||
_normalize_subgraph_payload,
|
||||
_resolve_required_subgraph_nodes,
|
||||
_summarize_suspicious_prompt_inputs,
|
||||
_trim_subgraph_payload_to_required_nodes,
|
||||
_unwrap_wrapped_prompt_link,
|
||||
_validate_prompt_input_shapes,
|
||||
_validate_required_prompt_inputs,
|
||||
)
|
||||
from cloud_session_bridge import (
|
||||
_RemoteSessionBridgeResolutionStats,
|
||||
_build_remote_session_bridge_record,
|
||||
@@ -99,7 +140,6 @@ except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
)
|
||||
from remote_protocol import (
|
||||
BOUNDARY_INPUT_SIGNATURES_KEY as _BOUNDARY_INPUT_SIGNATURES_KEY,
|
||||
PRIMITIVE_WIDGET_INPUT_TYPES as _PRIMITIVE_WIDGET_INPUT_TYPES,
|
||||
)
|
||||
from serialization import (
|
||||
coerce_serialized_node_outputs,
|
||||
@@ -771,510 +811,6 @@ def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any:
|
||||
return slot_values
|
||||
|
||||
|
||||
def _extract_prompt_executor_error(executor: Any) -> str:
|
||||
"""Extract a useful failure message from a PromptExecutor run."""
|
||||
for event, data in reversed(executor.status_messages):
|
||||
if event == "execution_error":
|
||||
return _format_prompt_executor_error_payload(data)
|
||||
if event == "execution_interrupted":
|
||||
return "Remote subgraph execution was interrupted."
|
||||
return "Remote subgraph execution failed."
|
||||
|
||||
|
||||
def _format_prompt_executor_error_payload(data: Any) -> str:
|
||||
"""Return a richer human-readable PromptExecutor failure message when available."""
|
||||
if not isinstance(data, dict):
|
||||
return "Remote subgraph execution failed."
|
||||
|
||||
message = str(data.get("exception_message") or "Remote subgraph execution failed.")
|
||||
node_id = data.get("node_id")
|
||||
node_type = data.get("node_type")
|
||||
current_inputs = data.get("current_inputs")
|
||||
if node_id is None and node_type is None and not current_inputs:
|
||||
return message
|
||||
|
||||
details: list[str] = [message]
|
||||
if node_id is not None or node_type is not None:
|
||||
details.append(f"node_id={node_id!r} node_type={node_type!r}")
|
||||
if current_inputs:
|
||||
details.append(f"current_inputs={current_inputs!r}")
|
||||
return " | ".join(details)
|
||||
|
||||
|
||||
def _extract_prompt_executor_error_payload(executor: Any) -> dict[str, Any] | None:
|
||||
"""Return the most recent PromptExecutor execution_error payload when present."""
|
||||
for event, data in reversed(executor.status_messages):
|
||||
if event == "execution_error" and isinstance(data, dict):
|
||||
return data
|
||||
return None
|
||||
|
||||
|
||||
def _summarize_suspicious_prompt_inputs(prompt: dict[str, Any]) -> list[str]:
|
||||
"""Return compact descriptions of prompt inputs that still look list-wrapped."""
|
||||
findings: list[str] = []
|
||||
for node_id, node_info in sorted(prompt.items()):
|
||||
inputs = node_info.get("inputs") or {}
|
||||
for input_name, input_value in inputs.items():
|
||||
if isinstance(input_value, list) and len(input_value) == 1:
|
||||
findings.append(f"{node_id}.{input_name}={input_value!r}")
|
||||
continue
|
||||
if (
|
||||
isinstance(input_value, list)
|
||||
and len(input_value) == 2
|
||||
and isinstance(input_value[0], str)
|
||||
and isinstance(input_value[1], list)
|
||||
):
|
||||
findings.append(f"{node_id}.{input_name}={input_value!r}")
|
||||
return findings
|
||||
|
||||
|
||||
def _node_input_types(
|
||||
node_class: type[Any],
|
||||
live_inputs: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Return one node class's input schema finalized for the live prompt inputs."""
|
||||
input_types_callable = getattr(node_class, "INPUT_TYPES", None)
|
||||
if not callable(input_types_callable):
|
||||
return {}
|
||||
|
||||
raw_input_types = input_types_callable()
|
||||
if not isinstance(raw_input_types, dict):
|
||||
return {}
|
||||
|
||||
has_v3_dynamic_inputs = any(
|
||||
isinstance(input_config, tuple)
|
||||
and bool(input_config)
|
||||
and isinstance(input_config[0], str)
|
||||
and input_config[0].startswith("COMFY_")
|
||||
and input_config[0].endswith("_V3")
|
||||
for section_name in ("required", "optional")
|
||||
for input_config in (
|
||||
raw_input_types.get(section_name, {}).values()
|
||||
if isinstance(raw_input_types.get(section_name), dict)
|
||||
else ()
|
||||
)
|
||||
)
|
||||
if not has_v3_dynamic_inputs:
|
||||
return raw_input_types
|
||||
|
||||
from comfy_api.latest import _io
|
||||
|
||||
finalized_input_types, _, _ = _io.get_finalized_class_inputs(
|
||||
raw_input_types,
|
||||
live_inputs or {},
|
||||
)
|
||||
logger.debug(
|
||||
"Finalized V3 input schema for %s against prompt inputs %s.",
|
||||
node_class.__name__,
|
||||
sorted(str(input_name) for input_name in (live_inputs or {})),
|
||||
)
|
||||
return finalized_input_types
|
||||
|
||||
|
||||
def _node_input_type_map(
|
||||
node_class: type[Any],
|
||||
live_inputs: dict[str, Any] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Return one node class's finalized input types keyed by prompt input name."""
|
||||
input_types = _node_input_types(node_class, live_inputs)
|
||||
|
||||
input_type_map: dict[str, str] = {}
|
||||
for section_name in ("required", "optional", "hidden"):
|
||||
section = input_types.get(section_name)
|
||||
if not isinstance(section, dict):
|
||||
continue
|
||||
for input_name, input_config in section.items():
|
||||
if not isinstance(input_config, tuple) or not input_config:
|
||||
continue
|
||||
declared_type = input_config[0]
|
||||
if isinstance(declared_type, str):
|
||||
input_type_map[str(input_name)] = declared_type
|
||||
return input_type_map
|
||||
|
||||
|
||||
def _node_required_input_names(
|
||||
node_class: type[Any],
|
||||
live_inputs: dict[str, Any] | None = None,
|
||||
) -> set[str]:
|
||||
"""Return required input names after expanding V3 dynamic prompt inputs."""
|
||||
required_inputs = _node_input_types(node_class, live_inputs).get("required")
|
||||
if not isinstance(required_inputs, dict):
|
||||
return set()
|
||||
return {str(input_name) for input_name in required_inputs}
|
||||
|
||||
|
||||
def _coerce_primitive_prompt_input_value(
|
||||
*,
|
||||
node_id: str,
|
||||
class_type: str,
|
||||
input_name: str,
|
||||
declared_type: str,
|
||||
input_value: Any,
|
||||
) -> Any:
|
||||
"""Coerce one primitive prompt literal using ComfyUI's `validate_inputs` semantics."""
|
||||
literal_value = (
|
||||
input_value.get("__value__")
|
||||
if isinstance(input_value, dict) and "__value__" in input_value
|
||||
else input_value
|
||||
)
|
||||
if isinstance(literal_value, list):
|
||||
return input_value
|
||||
|
||||
try:
|
||||
if declared_type == "INT":
|
||||
return int(literal_value)
|
||||
if declared_type == "FLOAT":
|
||||
return float(literal_value)
|
||||
if declared_type == "STRING":
|
||||
return str(literal_value)
|
||||
if declared_type == "BOOLEAN":
|
||||
return bool(literal_value)
|
||||
except (OverflowError, TypeError, ValueError) as exc:
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"Remote subgraph input could not be coerced to the declared primitive socket type."
|
||||
f" node_id={node_id!r} node_type={class_type!r}"
|
||||
f" input_name={input_name!r} declared_type={declared_type!r}"
|
||||
f" received_value={literal_value!r}"
|
||||
) from exc
|
||||
|
||||
return input_value
|
||||
|
||||
|
||||
def _coerce_prompt_primitive_input_values(
|
||||
prompt: dict[str, Any],
|
||||
node_mapping: dict[str, type[Any]],
|
||||
) -> None:
|
||||
"""Mutate prompt literals in-place to match ComfyUI's primitive widget coercion."""
|
||||
for node_id, prompt_node in sorted(prompt.items()):
|
||||
class_type = str(prompt_node.get("class_type"))
|
||||
node_class = node_mapping.get(class_type)
|
||||
if node_class is None:
|
||||
continue
|
||||
inputs = prompt_node.get("inputs") or {}
|
||||
input_type_map = _node_input_type_map(node_class, inputs)
|
||||
if not input_type_map:
|
||||
continue
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
declared_type = input_type_map.get(str(input_name))
|
||||
if declared_type not in _PRIMITIVE_WIDGET_INPUT_TYPES:
|
||||
continue
|
||||
if (
|
||||
isinstance(input_value, list)
|
||||
and len(input_value) == 2
|
||||
and isinstance(input_value[0], str)
|
||||
):
|
||||
continue
|
||||
coerced_value = _coerce_primitive_prompt_input_value(
|
||||
node_id=str(node_id),
|
||||
class_type=class_type,
|
||||
input_name=str(input_name),
|
||||
declared_type=declared_type,
|
||||
input_value=input_value,
|
||||
)
|
||||
if coerced_value is not input_value:
|
||||
logger.debug(
|
||||
"Coerced remote primitive input %s.%s from %r to %r for type %s.",
|
||||
node_id,
|
||||
input_name,
|
||||
input_value,
|
||||
coerced_value,
|
||||
declared_type,
|
||||
)
|
||||
inputs[input_name] = coerced_value
|
||||
|
||||
|
||||
def _validate_prompt_input_shapes(
|
||||
prompt: dict[str, Any],
|
||||
node_mapping: dict[str, type[Any]],
|
||||
boundary_input_specs: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Reject prompt inputs that still look invalid for primitive widget sockets."""
|
||||
boundary_targets = {
|
||||
(str(target.get("node_id")), str(target.get("input_name")))
|
||||
for boundary_input in (boundary_input_specs or [])
|
||||
for target in boundary_input.get("targets", [])
|
||||
if target.get("node_id") is not None and target.get("input_name") is not None
|
||||
}
|
||||
for node_id, prompt_node in sorted(prompt.items()):
|
||||
class_type = str(prompt_node.get("class_type"))
|
||||
node_class = node_mapping.get(class_type)
|
||||
if node_class is None:
|
||||
continue
|
||||
inputs = prompt_node.get("inputs") or {}
|
||||
input_type_map = _node_input_type_map(node_class, inputs)
|
||||
if not input_type_map:
|
||||
continue
|
||||
for input_name, input_value in inputs.items():
|
||||
declared_type = input_type_map.get(str(input_name))
|
||||
if declared_type not in _PRIMITIVE_WIDGET_INPUT_TYPES:
|
||||
continue
|
||||
if (
|
||||
isinstance(input_value, list)
|
||||
and len(input_value) == 2
|
||||
and isinstance(input_value[0], str)
|
||||
):
|
||||
continue
|
||||
if (str(node_id), str(input_name)) in boundary_targets:
|
||||
continue
|
||||
literal_value = (
|
||||
input_value.get("__value__")
|
||||
if isinstance(input_value, dict) and "__value__" in input_value
|
||||
else input_value
|
||||
)
|
||||
if isinstance(literal_value, list):
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"Remote subgraph input has an invalid list value for a primitive socket."
|
||||
f" node_id={node_id!r} node_type={class_type!r}"
|
||||
f" input_name={input_name!r} declared_type={declared_type!r}"
|
||||
f" received_value={literal_value!r}"
|
||||
)
|
||||
|
||||
|
||||
def _validate_required_prompt_inputs(
|
||||
prompt: dict[str, Any],
|
||||
node_mapping: dict[str, type[Any]],
|
||||
) -> None:
|
||||
"""Reject remote prompt nodes that are missing declared required inputs."""
|
||||
for node_id, prompt_node in sorted(prompt.items()):
|
||||
class_type = str(prompt_node.get("class_type"))
|
||||
node_class = node_mapping.get(class_type)
|
||||
if node_class is None:
|
||||
continue
|
||||
inputs = prompt_node.get("inputs") or {}
|
||||
missing_inputs = sorted(
|
||||
_node_required_input_names(node_class, inputs) - set(inputs.keys())
|
||||
)
|
||||
if not missing_inputs:
|
||||
continue
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"Remote subgraph prompt is missing required node inputs before execution."
|
||||
f" node_id={node_id!r} node_type={class_type!r}"
|
||||
f" missing_inputs={missing_inputs!r} available_inputs={sorted(str(key) for key in inputs)!r}"
|
||||
)
|
||||
|
||||
|
||||
def _log_prompt_executor_failure_details(
|
||||
*,
|
||||
component_id: str,
|
||||
prompt: dict[str, Any],
|
||||
normalized_payload: dict[str, Any],
|
||||
executor: Any,
|
||||
) -> None:
|
||||
"""Emit high-signal diagnostics for one remote PromptExecutor failure."""
|
||||
error_payload = _extract_prompt_executor_error_payload(executor)
|
||||
suspicious_inputs = _summarize_suspicious_prompt_inputs(prompt)
|
||||
logger.error(
|
||||
"Remote PromptExecutor failed for component=%s execute_node_ids=%s boundary_outputs=%s suspicious_inputs=%s error_payload=%s",
|
||||
component_id,
|
||||
normalized_payload.get("execute_node_ids", []),
|
||||
[
|
||||
{
|
||||
"node_id": boundary_output.get("node_id"),
|
||||
"output_index": boundary_output.get("output_index"),
|
||||
}
|
||||
for boundary_output in normalized_payload.get("boundary_outputs", [])
|
||||
],
|
||||
suspicious_inputs,
|
||||
error_payload,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_required_subgraph_nodes(
|
||||
prompt: dict[str, Any],
|
||||
execute_node_ids: list[str],
|
||||
) -> list[str]:
|
||||
"""Return the dependency closure needed to execute the requested subgraph nodes."""
|
||||
required: set[str] = set()
|
||||
pending = list(execute_node_ids)
|
||||
logger.info(
|
||||
"Resolving dependency closure for remote execute targets: %s", execute_node_ids
|
||||
)
|
||||
while pending:
|
||||
node_id = str(pending.pop())
|
||||
if node_id not in prompt:
|
||||
logger.warning(
|
||||
"Skipping missing remote execute target %s while resolving dependency closure.",
|
||||
node_id,
|
||||
)
|
||||
continue
|
||||
if node_id in required:
|
||||
continue
|
||||
required.add(node_id)
|
||||
for input_value in (prompt[node_id].get("inputs") or {}).values():
|
||||
if _is_link(input_value):
|
||||
pending.append(str(input_value[0]))
|
||||
resolved = sorted(required)
|
||||
logger.info("Resolved remote dependency closure: %s", resolved)
|
||||
return resolved
|
||||
|
||||
|
||||
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 _normalize_link_output_index(value: Any) -> Any:
|
||||
"""Unwrap a singleton list around a prompt-link output index when present."""
|
||||
while isinstance(value, list) and len(value) == 1:
|
||||
value = value[0]
|
||||
return value
|
||||
|
||||
|
||||
def _unwrap_wrapped_prompt_link(value: Any) -> Any:
|
||||
"""Collapse nested singleton wrappers around one serialized prompt link when present."""
|
||||
candidate = value
|
||||
while isinstance(candidate, list) and len(candidate) == 1:
|
||||
candidate = candidate[0]
|
||||
if _is_link(candidate):
|
||||
return [candidate[0], _normalize_link_output_index(candidate[1])]
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_prompt_input_value(value: Any, io_type: str | None = None) -> Any:
|
||||
"""Unwrap transport-added singleton wrappers only for scalar-like prompt input values."""
|
||||
wrapped_link = _unwrap_wrapped_prompt_link(value)
|
||||
if wrapped_link is not value:
|
||||
return wrapped_link
|
||||
while (
|
||||
isinstance(value, list)
|
||||
and len(value) == 1
|
||||
and (
|
||||
io_type in _PRIMITIVE_WIDGET_INPUT_TYPES
|
||||
or value[0] is None
|
||||
or isinstance(value[0], bool | int | float | str)
|
||||
)
|
||||
):
|
||||
value = value[0]
|
||||
if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str):
|
||||
return [value[0], _normalize_link_output_index(value[1])]
|
||||
if value is None or isinstance(value, bool | int | float | str):
|
||||
return value
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_subgraph_payload(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return a subgraph payload with canonical prompt-link and output-index shapes."""
|
||||
normalized_payload = copy.deepcopy(payload)
|
||||
|
||||
for node_info in normalized_payload.get("subgraph_prompt", {}).values():
|
||||
inputs = node_info.get("inputs") or {}
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
inputs[input_name] = _normalize_prompt_input_value(input_value)
|
||||
|
||||
for boundary_output in normalized_payload.get("boundary_outputs", []):
|
||||
if "node_id" in boundary_output and isinstance(
|
||||
boundary_output["node_id"], list
|
||||
):
|
||||
boundary_output["node_id"] = _normalize_prompt_input_value(
|
||||
boundary_output["node_id"]
|
||||
)
|
||||
if "output_index" in boundary_output:
|
||||
boundary_output["output_index"] = _normalize_link_output_index(
|
||||
boundary_output["output_index"]
|
||||
)
|
||||
|
||||
normalized_payload["execute_node_ids"] = [
|
||||
_normalize_prompt_input_value(node_id)
|
||||
for node_id in normalized_payload.get("execute_node_ids", [])
|
||||
]
|
||||
|
||||
return normalized_payload
|
||||
|
||||
|
||||
def _trim_subgraph_payload_to_required_nodes(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Trim a subgraph payload down to the dependency closure of its execute targets."""
|
||||
trimmed_payload = copy.deepcopy(payload)
|
||||
prompt = trimmed_payload.get("subgraph_prompt", {})
|
||||
if not isinstance(prompt, dict):
|
||||
return trimmed_payload
|
||||
|
||||
prompt_node_ids = {str(node_id) for node_id in prompt}
|
||||
requested_execute_node_ids = [
|
||||
str(node_id) for node_id in trimmed_payload.get("execute_node_ids", [])
|
||||
]
|
||||
available_execute_node_ids = [
|
||||
node_id for node_id in requested_execute_node_ids if node_id in prompt_node_ids
|
||||
]
|
||||
dropped_execute_node_ids = [
|
||||
node_id
|
||||
for node_id in requested_execute_node_ids
|
||||
if node_id not in prompt_node_ids
|
||||
]
|
||||
if dropped_execute_node_ids:
|
||||
logger.warning(
|
||||
"Dropping remote execute targets absent from subgraph prompt for component=%s: %s",
|
||||
payload.get("component_id"),
|
||||
dropped_execute_node_ids,
|
||||
)
|
||||
|
||||
required_node_ids = set(
|
||||
_resolve_required_subgraph_nodes(
|
||||
prompt=prompt,
|
||||
execute_node_ids=available_execute_node_ids,
|
||||
)
|
||||
)
|
||||
if not required_node_ids:
|
||||
return trimmed_payload
|
||||
|
||||
original_node_ids = list(prompt.keys())
|
||||
trimmed_payload["subgraph_prompt"] = {
|
||||
str(node_id): prompt[node_id]
|
||||
for node_id in original_node_ids
|
||||
if str(node_id) in required_node_ids
|
||||
}
|
||||
trimmed_payload["boundary_inputs"] = [
|
||||
{
|
||||
**copy.deepcopy(boundary_input),
|
||||
"targets": [
|
||||
copy.deepcopy(target)
|
||||
for target in boundary_input.get("targets", [])
|
||||
if str(target.get("node_id")) in required_node_ids
|
||||
],
|
||||
}
|
||||
for boundary_input in trimmed_payload.get("boundary_inputs", [])
|
||||
if any(
|
||||
str(target.get("node_id")) in required_node_ids
|
||||
for target in boundary_input.get("targets", [])
|
||||
)
|
||||
]
|
||||
trimmed_payload["boundary_outputs"] = [
|
||||
copy.deepcopy(boundary_output)
|
||||
for boundary_output in trimmed_payload.get("boundary_outputs", [])
|
||||
if str(boundary_output.get("node_id")) in required_node_ids
|
||||
]
|
||||
trimmed_payload["component_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("component_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
trimmed_payload["execute_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("execute_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
trimmed_payload["mapped_execute_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("mapped_execute_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
trimmed_payload["static_execute_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("static_execute_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
logger.info(
|
||||
"Trimmed remote subgraph payload %s from %d prompt nodes to %d required nodes.",
|
||||
payload.get("component_id"),
|
||||
len(original_node_ids),
|
||||
len(trimmed_payload["subgraph_prompt"]),
|
||||
)
|
||||
return trimmed_payload
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _PreparedSubgraphExecution:
|
||||
"""Hold validated prompt state required by one PromptExecutor invocation."""
|
||||
@@ -1853,337 +1389,3 @@ def execute_subgraph_locally(
|
||||
remote_session_store().clear_session(session_handle)
|
||||
with _timed_phase("serialize_boundary_outputs", component=component_id):
|
||||
return serialize_node_outputs(outputs)
|
||||
|
||||
|
||||
def _mapped_phase_definition(
|
||||
payload: dict[str, Any], phase_key: str
|
||||
) -> dict[str, Any] | None:
|
||||
"""Return one explicit mapped phase definition when queue-time planning provided it."""
|
||||
phase_payload = payload.get(phase_key)
|
||||
if isinstance(phase_payload, dict):
|
||||
return phase_payload
|
||||
return None
|
||||
|
||||
|
||||
def _shared_subgraph_payload_fields(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return the payload fields shared by every explicit mapped phase."""
|
||||
shared_fields = {
|
||||
"prompt_id": payload.get("prompt_id"),
|
||||
"extra_data": copy.deepcopy(payload.get("extra_data") or {}),
|
||||
"requires_volume_reload": bool(payload.get("requires_volume_reload", True)),
|
||||
"volume_reload_marker": payload.get("volume_reload_marker"),
|
||||
"uploaded_volume_paths": list(payload.get("uploaded_volume_paths", [])),
|
||||
"terminate_container_on_error": bool(
|
||||
payload.get("terminate_container_on_error", True)
|
||||
),
|
||||
"custom_nodes_bundle": payload.get("custom_nodes_bundle"),
|
||||
}
|
||||
snapshot_profile_key = payload.get("snapshot_profile_key")
|
||||
if isinstance(snapshot_profile_key, str) and snapshot_profile_key.strip():
|
||||
shared_fields["snapshot_profile_key"] = snapshot_profile_key.strip()
|
||||
remote_session = payload.get("remote_session")
|
||||
if remote_session is not None:
|
||||
shared_fields["remote_session"] = copy.deepcopy(remote_session)
|
||||
if bool(payload.get("clear_remote_session")):
|
||||
shared_fields["clear_remote_session"] = True
|
||||
return shared_fields
|
||||
|
||||
|
||||
def _build_phase_subgraph_payload(
|
||||
payload: dict[str, Any],
|
||||
phase_key: str,
|
||||
component_id: str,
|
||||
) -> 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}.")
|
||||
|
||||
return {
|
||||
"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)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
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 _aggregate_mapped_phase_outputs(
|
||||
per_item_outputs: list[tuple[Any, ...]],
|
||||
payload: dict[str, Any],
|
||||
) -> tuple[Any, ...]:
|
||||
"""Join ordered mapped-phase outputs back 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."
|
||||
)
|
||||
|
||||
aggregated_outputs: list[Any] = []
|
||||
boundary_outputs = list(payload.get("boundary_outputs", []))
|
||||
for output_index in range(output_count):
|
||||
boundary_output = (
|
||||
boundary_outputs[output_index]
|
||||
if output_index < len(boundary_outputs)
|
||||
else {}
|
||||
)
|
||||
aggregated_outputs.append(
|
||||
_merge_static_or_mapped_values(
|
||||
[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 _merge_static_or_mapped_values(
|
||||
values: list[Any],
|
||||
*,
|
||||
io_type: str,
|
||||
is_list: bool,
|
||||
scheduler_is_list: bool,
|
||||
) -> Any:
|
||||
"""Join mapped per-item outputs using the shared transport serializer rules."""
|
||||
from serialization import join_mapped_values_for_scheduler
|
||||
|
||||
return join_mapped_values_for_scheduler(
|
||||
values,
|
||||
io_type=io_type,
|
||||
is_list=is_list,
|
||||
scheduler_is_list=scheduler_is_list,
|
||||
)
|
||||
|
||||
|
||||
def _merge_static_and_mapped_outputs(
|
||||
*,
|
||||
static_outputs: tuple[Any, ...],
|
||||
mapped_outputs: tuple[Any, ...],
|
||||
payload: dict[str, Any],
|
||||
) -> tuple[Any, ...]:
|
||||
"""Reassemble one mapped component's static and mapped outputs in proxy order."""
|
||||
combined_outputs: list[Any] = []
|
||||
static_output_index = 0
|
||||
mapped_output_index = 0
|
||||
for boundary_output in payload.get("boundary_outputs", []):
|
||||
if bool(boundary_output.get("mapped_output")):
|
||||
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 _execute_mapped_subgraph_payload(
|
||||
payload: dict[str, Any],
|
||||
hydrated_inputs: dict[str, Any],
|
||||
custom_nodes_root: Path | None,
|
||||
status_callback: Callable[[dict[str, Any]], None] | None = None,
|
||||
cancellation_event: threading.Event | None = None,
|
||||
interrupt_store: Any | None = None,
|
||||
interrupt_flag_key: str | None = None,
|
||||
) -> tuple[Any, ...]:
|
||||
"""Execute one mapped payload inside a single remote runtime process."""
|
||||
mapped_input = payload.get("mapped_input") or {}
|
||||
mapped_input_name = str(mapped_input.get("proxy_input_name") or "")
|
||||
if not mapped_input_name:
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"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."
|
||||
)
|
||||
|
||||
from serialization import split_mapped_value
|
||||
|
||||
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, ...] = ()
|
||||
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",
|
||||
)
|
||||
if static_phase_payload.get("execute_node_ids"):
|
||||
static_phase_outputs = _execute_subgraph_prompt(
|
||||
static_phase_payload,
|
||||
dict(broadcast_inputs),
|
||||
custom_nodes_root,
|
||||
status_callback=status_callback,
|
||||
cancellation_event=cancellation_event,
|
||||
interrupt_store=interrupt_store,
|
||||
interrupt_flag_key=interrupt_flag_key,
|
||||
)
|
||||
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)
|
||||
|
||||
if status_callback is not None:
|
||||
status_callback(
|
||||
{
|
||||
"event_type": "node_progress",
|
||||
"node_id": str(payload.get("component_id") or "modal-subgraph"),
|
||||
"display_node_id": str(payload.get("component_id") or "modal-subgraph"),
|
||||
"value": 0.0,
|
||||
"max": float(len(mapped_items)),
|
||||
"aggregate_only": True,
|
||||
}
|
||||
)
|
||||
|
||||
per_item_outputs: list[tuple[Any, ...]] = []
|
||||
for item_index, item_value in enumerate(mapped_items):
|
||||
last_lane_node_id: str | None = None
|
||||
lane_id = str(payload.get("mapped_progress_lane_id") or item_index)
|
||||
|
||||
def publish_item_status(progress_state: dict[str, Any]) -> None:
|
||||
"""Attach mapped-lane metadata to one per-item progress event."""
|
||||
nonlocal last_lane_node_id
|
||||
if status_callback is None:
|
||||
return
|
||||
event_type = str(progress_state.get("event_type", ""))
|
||||
if event_type == "node_progress":
|
||||
reported_node_id = progress_state.get(
|
||||
"real_node_id"
|
||||
) or progress_state.get("node_id")
|
||||
if reported_node_id is not None:
|
||||
last_lane_node_id = str(reported_node_id)
|
||||
status_callback(
|
||||
{
|
||||
**progress_state,
|
||||
"lane_id": lane_id,
|
||||
"item_index": item_index,
|
||||
}
|
||||
)
|
||||
return
|
||||
if event_type in {"executed", "preview", "boundary_output"}:
|
||||
status_callback({**progress_state, "item_index": item_index})
|
||||
|
||||
item_payload = _build_phase_subgraph_payload(
|
||||
payload,
|
||||
"mapped_phase",
|
||||
f"{payload.get('component_id', 'modal-subgraph')}::item:{item_index}",
|
||||
)
|
||||
item_inputs = dict(broadcast_inputs)
|
||||
item_inputs[mapped_input_name] = item_value
|
||||
per_item_outputs.append(
|
||||
_execute_subgraph_prompt(
|
||||
item_payload,
|
||||
item_inputs,
|
||||
custom_nodes_root,
|
||||
status_callback=publish_item_status,
|
||||
cancellation_event=cancellation_event,
|
||||
interrupt_store=interrupt_store,
|
||||
interrupt_flag_key=interrupt_flag_key,
|
||||
)
|
||||
)
|
||||
if status_callback is not None:
|
||||
status_callback(
|
||||
{
|
||||
"event_type": "node_progress",
|
||||
"node_id": last_lane_node_id
|
||||
or str(payload.get("component_id") or "modal-subgraph"),
|
||||
"display_node_id": last_lane_node_id
|
||||
or str(payload.get("component_id") or "modal-subgraph"),
|
||||
"value": 0.0,
|
||||
"max": 1.0,
|
||||
"lane_id": lane_id,
|
||||
"item_index": item_index,
|
||||
"clear": True,
|
||||
}
|
||||
)
|
||||
status_callback(
|
||||
{
|
||||
"event_type": "node_progress",
|
||||
"node_id": str(payload.get("component_id") or "modal-subgraph"),
|
||||
"display_node_id": str(
|
||||
payload.get("component_id") or "modal-subgraph"
|
||||
),
|
||||
"value": float(item_index + 1),
|
||||
"max": float(len(mapped_items)),
|
||||
"aggregate_only": True,
|
||||
}
|
||||
)
|
||||
|
||||
mapped_phase_payload = _build_phase_subgraph_payload(
|
||||
payload,
|
||||
"mapped_phase",
|
||||
f"{payload.get('component_id', 'modal-subgraph')}::mapped",
|
||||
)
|
||||
mapped_outputs = _aggregate_mapped_phase_outputs(
|
||||
per_item_outputs,
|
||||
{"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,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,534 @@
|
||||
"""Prompt normalization, schema validation, and failure diagnostics."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from .remote_protocol import (
|
||||
PRIMITIVE_WIDGET_INPUT_TYPES as _PRIMITIVE_WIDGET_INPUT_TYPES,
|
||||
)
|
||||
except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
from remote_protocol import (
|
||||
PRIMITIVE_WIDGET_INPUT_TYPES as _PRIMITIVE_WIDGET_INPUT_TYPES,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RemoteSubgraphExecutionError: type[RuntimeError] = RuntimeError
|
||||
|
||||
|
||||
def configure_cloud_prompt_validation_error(
|
||||
error_type: type[RuntimeError],
|
||||
) -> None:
|
||||
"""Install the stable cloud entrypoint's subgraph execution error type."""
|
||||
global RemoteSubgraphExecutionError
|
||||
RemoteSubgraphExecutionError = error_type
|
||||
|
||||
|
||||
def _extract_prompt_executor_error(executor: Any) -> str:
|
||||
"""Extract a useful failure message from a PromptExecutor run."""
|
||||
for event, data in reversed(executor.status_messages):
|
||||
if event == "execution_error":
|
||||
return _format_prompt_executor_error_payload(data)
|
||||
if event == "execution_interrupted":
|
||||
return "Remote subgraph execution was interrupted."
|
||||
return "Remote subgraph execution failed."
|
||||
|
||||
|
||||
def _format_prompt_executor_error_payload(data: Any) -> str:
|
||||
"""Return a richer human-readable PromptExecutor failure message when available."""
|
||||
if not isinstance(data, dict):
|
||||
return "Remote subgraph execution failed."
|
||||
|
||||
message = str(data.get("exception_message") or "Remote subgraph execution failed.")
|
||||
node_id = data.get("node_id")
|
||||
node_type = data.get("node_type")
|
||||
current_inputs = data.get("current_inputs")
|
||||
if node_id is None and node_type is None and not current_inputs:
|
||||
return message
|
||||
|
||||
details: list[str] = [message]
|
||||
if node_id is not None or node_type is not None:
|
||||
details.append(f"node_id={node_id!r} node_type={node_type!r}")
|
||||
if current_inputs:
|
||||
details.append(f"current_inputs={current_inputs!r}")
|
||||
return " | ".join(details)
|
||||
|
||||
|
||||
def _extract_prompt_executor_error_payload(executor: Any) -> dict[str, Any] | None:
|
||||
"""Return the most recent PromptExecutor execution_error payload when present."""
|
||||
for event, data in reversed(executor.status_messages):
|
||||
if event == "execution_error" and isinstance(data, dict):
|
||||
return data
|
||||
return None
|
||||
|
||||
|
||||
def _summarize_suspicious_prompt_inputs(prompt: dict[str, Any]) -> list[str]:
|
||||
"""Return compact descriptions of prompt inputs that still look list-wrapped."""
|
||||
findings: list[str] = []
|
||||
for node_id, node_info in sorted(prompt.items()):
|
||||
inputs = node_info.get("inputs") or {}
|
||||
for input_name, input_value in inputs.items():
|
||||
if isinstance(input_value, list) and len(input_value) == 1:
|
||||
findings.append(f"{node_id}.{input_name}={input_value!r}")
|
||||
continue
|
||||
if (
|
||||
isinstance(input_value, list)
|
||||
and len(input_value) == 2
|
||||
and isinstance(input_value[0], str)
|
||||
and isinstance(input_value[1], list)
|
||||
):
|
||||
findings.append(f"{node_id}.{input_name}={input_value!r}")
|
||||
return findings
|
||||
|
||||
|
||||
def _node_input_types(
|
||||
node_class: type[Any],
|
||||
live_inputs: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Return one node class's input schema finalized for the live prompt inputs."""
|
||||
input_types_callable = getattr(node_class, "INPUT_TYPES", None)
|
||||
if not callable(input_types_callable):
|
||||
return {}
|
||||
|
||||
raw_input_types = input_types_callable()
|
||||
if not isinstance(raw_input_types, dict):
|
||||
return {}
|
||||
|
||||
has_v3_dynamic_inputs = any(
|
||||
isinstance(input_config, tuple)
|
||||
and bool(input_config)
|
||||
and isinstance(input_config[0], str)
|
||||
and input_config[0].startswith("COMFY_")
|
||||
and input_config[0].endswith("_V3")
|
||||
for section_name in ("required", "optional")
|
||||
for input_config in (
|
||||
raw_input_types.get(section_name, {}).values()
|
||||
if isinstance(raw_input_types.get(section_name), dict)
|
||||
else ()
|
||||
)
|
||||
)
|
||||
if not has_v3_dynamic_inputs:
|
||||
return raw_input_types
|
||||
|
||||
from comfy_api.latest import _io
|
||||
|
||||
finalized_input_types, _, _ = _io.get_finalized_class_inputs(
|
||||
raw_input_types,
|
||||
live_inputs or {},
|
||||
)
|
||||
logger.debug(
|
||||
"Finalized V3 input schema for %s against prompt inputs %s.",
|
||||
node_class.__name__,
|
||||
sorted(str(input_name) for input_name in (live_inputs or {})),
|
||||
)
|
||||
return finalized_input_types
|
||||
|
||||
|
||||
def _node_input_type_map(
|
||||
node_class: type[Any],
|
||||
live_inputs: dict[str, Any] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Return one node class's finalized input types keyed by prompt input name."""
|
||||
input_types = _node_input_types(node_class, live_inputs)
|
||||
|
||||
input_type_map: dict[str, str] = {}
|
||||
for section_name in ("required", "optional", "hidden"):
|
||||
section = input_types.get(section_name)
|
||||
if not isinstance(section, dict):
|
||||
continue
|
||||
for input_name, input_config in section.items():
|
||||
if not isinstance(input_config, tuple) or not input_config:
|
||||
continue
|
||||
declared_type = input_config[0]
|
||||
if isinstance(declared_type, str):
|
||||
input_type_map[str(input_name)] = declared_type
|
||||
return input_type_map
|
||||
|
||||
|
||||
def _node_required_input_names(
|
||||
node_class: type[Any],
|
||||
live_inputs: dict[str, Any] | None = None,
|
||||
) -> set[str]:
|
||||
"""Return required input names after expanding V3 dynamic prompt inputs."""
|
||||
required_inputs = _node_input_types(node_class, live_inputs).get("required")
|
||||
if not isinstance(required_inputs, dict):
|
||||
return set()
|
||||
return {str(input_name) for input_name in required_inputs}
|
||||
|
||||
|
||||
def _coerce_primitive_prompt_input_value(
|
||||
*,
|
||||
node_id: str,
|
||||
class_type: str,
|
||||
input_name: str,
|
||||
declared_type: str,
|
||||
input_value: Any,
|
||||
) -> Any:
|
||||
"""Coerce one primitive prompt literal using ComfyUI's `validate_inputs` semantics."""
|
||||
literal_value = (
|
||||
input_value.get("__value__")
|
||||
if isinstance(input_value, dict) and "__value__" in input_value
|
||||
else input_value
|
||||
)
|
||||
if isinstance(literal_value, list):
|
||||
return input_value
|
||||
|
||||
try:
|
||||
if declared_type == "INT":
|
||||
return int(literal_value)
|
||||
if declared_type == "FLOAT":
|
||||
return float(literal_value)
|
||||
if declared_type == "STRING":
|
||||
return str(literal_value)
|
||||
if declared_type == "BOOLEAN":
|
||||
return bool(literal_value)
|
||||
except (OverflowError, TypeError, ValueError) as exc:
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"Remote subgraph input could not be coerced to the declared primitive socket type."
|
||||
f" node_id={node_id!r} node_type={class_type!r}"
|
||||
f" input_name={input_name!r} declared_type={declared_type!r}"
|
||||
f" received_value={literal_value!r}"
|
||||
) from exc
|
||||
|
||||
return input_value
|
||||
|
||||
|
||||
def _coerce_prompt_primitive_input_values(
|
||||
prompt: dict[str, Any],
|
||||
node_mapping: dict[str, type[Any]],
|
||||
) -> None:
|
||||
"""Mutate prompt literals in-place to match ComfyUI's primitive widget coercion."""
|
||||
for node_id, prompt_node in sorted(prompt.items()):
|
||||
class_type = str(prompt_node.get("class_type"))
|
||||
node_class = node_mapping.get(class_type)
|
||||
if node_class is None:
|
||||
continue
|
||||
inputs = prompt_node.get("inputs") or {}
|
||||
input_type_map = _node_input_type_map(node_class, inputs)
|
||||
if not input_type_map:
|
||||
continue
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
declared_type = input_type_map.get(str(input_name))
|
||||
if declared_type not in _PRIMITIVE_WIDGET_INPUT_TYPES:
|
||||
continue
|
||||
if (
|
||||
isinstance(input_value, list)
|
||||
and len(input_value) == 2
|
||||
and isinstance(input_value[0], str)
|
||||
):
|
||||
continue
|
||||
coerced_value = _coerce_primitive_prompt_input_value(
|
||||
node_id=str(node_id),
|
||||
class_type=class_type,
|
||||
input_name=str(input_name),
|
||||
declared_type=declared_type,
|
||||
input_value=input_value,
|
||||
)
|
||||
if coerced_value is not input_value:
|
||||
logger.debug(
|
||||
"Coerced remote primitive input %s.%s from %r to %r for type %s.",
|
||||
node_id,
|
||||
input_name,
|
||||
input_value,
|
||||
coerced_value,
|
||||
declared_type,
|
||||
)
|
||||
inputs[input_name] = coerced_value
|
||||
|
||||
|
||||
def _validate_prompt_input_shapes(
|
||||
prompt: dict[str, Any],
|
||||
node_mapping: dict[str, type[Any]],
|
||||
boundary_input_specs: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Reject prompt inputs that still look invalid for primitive widget sockets."""
|
||||
boundary_targets = {
|
||||
(str(target.get("node_id")), str(target.get("input_name")))
|
||||
for boundary_input in (boundary_input_specs or [])
|
||||
for target in boundary_input.get("targets", [])
|
||||
if target.get("node_id") is not None and target.get("input_name") is not None
|
||||
}
|
||||
for node_id, prompt_node in sorted(prompt.items()):
|
||||
class_type = str(prompt_node.get("class_type"))
|
||||
node_class = node_mapping.get(class_type)
|
||||
if node_class is None:
|
||||
continue
|
||||
inputs = prompt_node.get("inputs") or {}
|
||||
input_type_map = _node_input_type_map(node_class, inputs)
|
||||
if not input_type_map:
|
||||
continue
|
||||
for input_name, input_value in inputs.items():
|
||||
declared_type = input_type_map.get(str(input_name))
|
||||
if declared_type not in _PRIMITIVE_WIDGET_INPUT_TYPES:
|
||||
continue
|
||||
if (
|
||||
isinstance(input_value, list)
|
||||
and len(input_value) == 2
|
||||
and isinstance(input_value[0], str)
|
||||
):
|
||||
continue
|
||||
if (str(node_id), str(input_name)) in boundary_targets:
|
||||
continue
|
||||
literal_value = (
|
||||
input_value.get("__value__")
|
||||
if isinstance(input_value, dict) and "__value__" in input_value
|
||||
else input_value
|
||||
)
|
||||
if isinstance(literal_value, list):
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"Remote subgraph input has an invalid list value for a primitive socket."
|
||||
f" node_id={node_id!r} node_type={class_type!r}"
|
||||
f" input_name={input_name!r} declared_type={declared_type!r}"
|
||||
f" received_value={literal_value!r}"
|
||||
)
|
||||
|
||||
|
||||
def _validate_required_prompt_inputs(
|
||||
prompt: dict[str, Any],
|
||||
node_mapping: dict[str, type[Any]],
|
||||
) -> None:
|
||||
"""Reject remote prompt nodes that are missing declared required inputs."""
|
||||
for node_id, prompt_node in sorted(prompt.items()):
|
||||
class_type = str(prompt_node.get("class_type"))
|
||||
node_class = node_mapping.get(class_type)
|
||||
if node_class is None:
|
||||
continue
|
||||
inputs = prompt_node.get("inputs") or {}
|
||||
missing_inputs = sorted(
|
||||
_node_required_input_names(node_class, inputs) - set(inputs.keys())
|
||||
)
|
||||
if not missing_inputs:
|
||||
continue
|
||||
raise RemoteSubgraphExecutionError(
|
||||
"Remote subgraph prompt is missing required node inputs before execution."
|
||||
f" node_id={node_id!r} node_type={class_type!r}"
|
||||
f" missing_inputs={missing_inputs!r} available_inputs={sorted(str(key) for key in inputs)!r}"
|
||||
)
|
||||
|
||||
|
||||
def _log_prompt_executor_failure_details(
|
||||
*,
|
||||
component_id: str,
|
||||
prompt: dict[str, Any],
|
||||
normalized_payload: dict[str, Any],
|
||||
executor: Any,
|
||||
) -> None:
|
||||
"""Emit high-signal diagnostics for one remote PromptExecutor failure."""
|
||||
error_payload = _extract_prompt_executor_error_payload(executor)
|
||||
suspicious_inputs = _summarize_suspicious_prompt_inputs(prompt)
|
||||
logger.error(
|
||||
"Remote PromptExecutor failed for component=%s execute_node_ids=%s boundary_outputs=%s suspicious_inputs=%s error_payload=%s",
|
||||
component_id,
|
||||
normalized_payload.get("execute_node_ids", []),
|
||||
[
|
||||
{
|
||||
"node_id": boundary_output.get("node_id"),
|
||||
"output_index": boundary_output.get("output_index"),
|
||||
}
|
||||
for boundary_output in normalized_payload.get("boundary_outputs", [])
|
||||
],
|
||||
suspicious_inputs,
|
||||
error_payload,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_required_subgraph_nodes(
|
||||
prompt: dict[str, Any],
|
||||
execute_node_ids: list[str],
|
||||
) -> list[str]:
|
||||
"""Return the dependency closure needed to execute the requested subgraph nodes."""
|
||||
required: set[str] = set()
|
||||
pending = list(execute_node_ids)
|
||||
logger.info(
|
||||
"Resolving dependency closure for remote execute targets: %s", execute_node_ids
|
||||
)
|
||||
while pending:
|
||||
node_id = str(pending.pop())
|
||||
if node_id not in prompt:
|
||||
logger.warning(
|
||||
"Skipping missing remote execute target %s while resolving dependency closure.",
|
||||
node_id,
|
||||
)
|
||||
continue
|
||||
if node_id in required:
|
||||
continue
|
||||
required.add(node_id)
|
||||
for input_value in (prompt[node_id].get("inputs") or {}).values():
|
||||
if _is_link(input_value):
|
||||
pending.append(str(input_value[0]))
|
||||
resolved = sorted(required)
|
||||
logger.info("Resolved remote dependency closure: %s", resolved)
|
||||
return resolved
|
||||
|
||||
|
||||
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 _normalize_link_output_index(value: Any) -> Any:
|
||||
"""Unwrap a singleton list around a prompt-link output index when present."""
|
||||
while isinstance(value, list) and len(value) == 1:
|
||||
value = value[0]
|
||||
return value
|
||||
|
||||
|
||||
def _unwrap_wrapped_prompt_link(value: Any) -> Any:
|
||||
"""Collapse nested singleton wrappers around one serialized prompt link when present."""
|
||||
candidate = value
|
||||
while isinstance(candidate, list) and len(candidate) == 1:
|
||||
candidate = candidate[0]
|
||||
if _is_link(candidate):
|
||||
return [candidate[0], _normalize_link_output_index(candidate[1])]
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_prompt_input_value(value: Any, io_type: str | None = None) -> Any:
|
||||
"""Unwrap transport-added singleton wrappers only for scalar-like prompt input values."""
|
||||
wrapped_link = _unwrap_wrapped_prompt_link(value)
|
||||
if wrapped_link is not value:
|
||||
return wrapped_link
|
||||
while (
|
||||
isinstance(value, list)
|
||||
and len(value) == 1
|
||||
and (
|
||||
io_type in _PRIMITIVE_WIDGET_INPUT_TYPES
|
||||
or value[0] is None
|
||||
or isinstance(value[0], bool | int | float | str)
|
||||
)
|
||||
):
|
||||
value = value[0]
|
||||
if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str):
|
||||
return [value[0], _normalize_link_output_index(value[1])]
|
||||
if value is None or isinstance(value, bool | int | float | str):
|
||||
return value
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_subgraph_payload(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return a subgraph payload with canonical prompt-link and output-index shapes."""
|
||||
normalized_payload = copy.deepcopy(payload)
|
||||
|
||||
for node_info in normalized_payload.get("subgraph_prompt", {}).values():
|
||||
inputs = node_info.get("inputs") or {}
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
inputs[input_name] = _normalize_prompt_input_value(input_value)
|
||||
|
||||
for boundary_output in normalized_payload.get("boundary_outputs", []):
|
||||
if "node_id" in boundary_output and isinstance(
|
||||
boundary_output["node_id"], list
|
||||
):
|
||||
boundary_output["node_id"] = _normalize_prompt_input_value(
|
||||
boundary_output["node_id"]
|
||||
)
|
||||
if "output_index" in boundary_output:
|
||||
boundary_output["output_index"] = _normalize_link_output_index(
|
||||
boundary_output["output_index"]
|
||||
)
|
||||
|
||||
normalized_payload["execute_node_ids"] = [
|
||||
_normalize_prompt_input_value(node_id)
|
||||
for node_id in normalized_payload.get("execute_node_ids", [])
|
||||
]
|
||||
|
||||
return normalized_payload
|
||||
|
||||
|
||||
def _trim_subgraph_payload_to_required_nodes(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Trim a subgraph payload down to the dependency closure of its execute targets."""
|
||||
trimmed_payload = copy.deepcopy(payload)
|
||||
prompt = trimmed_payload.get("subgraph_prompt", {})
|
||||
if not isinstance(prompt, dict):
|
||||
return trimmed_payload
|
||||
|
||||
prompt_node_ids = {str(node_id) for node_id in prompt}
|
||||
requested_execute_node_ids = [
|
||||
str(node_id) for node_id in trimmed_payload.get("execute_node_ids", [])
|
||||
]
|
||||
available_execute_node_ids = [
|
||||
node_id for node_id in requested_execute_node_ids if node_id in prompt_node_ids
|
||||
]
|
||||
dropped_execute_node_ids = [
|
||||
node_id
|
||||
for node_id in requested_execute_node_ids
|
||||
if node_id not in prompt_node_ids
|
||||
]
|
||||
if dropped_execute_node_ids:
|
||||
logger.warning(
|
||||
"Dropping remote execute targets absent from subgraph prompt for component=%s: %s",
|
||||
payload.get("component_id"),
|
||||
dropped_execute_node_ids,
|
||||
)
|
||||
|
||||
required_node_ids = set(
|
||||
_resolve_required_subgraph_nodes(
|
||||
prompt=prompt,
|
||||
execute_node_ids=available_execute_node_ids,
|
||||
)
|
||||
)
|
||||
if not required_node_ids:
|
||||
return trimmed_payload
|
||||
|
||||
original_node_ids = list(prompt.keys())
|
||||
trimmed_payload["subgraph_prompt"] = {
|
||||
str(node_id): prompt[node_id]
|
||||
for node_id in original_node_ids
|
||||
if str(node_id) in required_node_ids
|
||||
}
|
||||
trimmed_payload["boundary_inputs"] = [
|
||||
{
|
||||
**copy.deepcopy(boundary_input),
|
||||
"targets": [
|
||||
copy.deepcopy(target)
|
||||
for target in boundary_input.get("targets", [])
|
||||
if str(target.get("node_id")) in required_node_ids
|
||||
],
|
||||
}
|
||||
for boundary_input in trimmed_payload.get("boundary_inputs", [])
|
||||
if any(
|
||||
str(target.get("node_id")) in required_node_ids
|
||||
for target in boundary_input.get("targets", [])
|
||||
)
|
||||
]
|
||||
trimmed_payload["boundary_outputs"] = [
|
||||
copy.deepcopy(boundary_output)
|
||||
for boundary_output in trimmed_payload.get("boundary_outputs", [])
|
||||
if str(boundary_output.get("node_id")) in required_node_ids
|
||||
]
|
||||
trimmed_payload["component_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("component_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
trimmed_payload["execute_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("execute_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
trimmed_payload["mapped_execute_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("mapped_execute_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
trimmed_payload["static_execute_node_ids"] = [
|
||||
str(node_id)
|
||||
for node_id in trimmed_payload.get("static_execute_node_ids", [])
|
||||
if str(node_id) in required_node_ids
|
||||
]
|
||||
logger.info(
|
||||
"Trimmed remote subgraph payload %s from %d prompt nodes to %d required nodes.",
|
||||
payload.get("component_id"),
|
||||
len(original_node_ids),
|
||||
len(trimmed_payload["subgraph_prompt"]),
|
||||
)
|
||||
return trimmed_payload
|
||||
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Timestamped logging and phase timing for the cloud worker runtime."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Any, Iterator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_TARGET_LOGGER: logging.Logger = logger
|
||||
_HANDLER_NAME = "comfyui-modal-sync-cloud-timestamped"
|
||||
|
||||
|
||||
def _build_cloud_log_formatter() -> logging.Formatter:
|
||||
"""Return the default formatter for remote Modal-Sync logs with timestamps."""
|
||||
return logging.Formatter(
|
||||
fmt="%(asctime)s.%(msecs)03d +%(relativeCreated)07.0fms %(levelname)s [%(name)s] %(message)s",
|
||||
datefmt="%H:%M:%S",
|
||||
)
|
||||
|
||||
|
||||
def configure_cloud_runtime_logging(
|
||||
target_logger: logging.Logger,
|
||||
handler_name: str,
|
||||
) -> logging.Logger:
|
||||
"""Configure and retain the stable entrypoint logger used by cloud helpers."""
|
||||
global _TARGET_LOGGER
|
||||
global _HANDLER_NAME
|
||||
_TARGET_LOGGER = target_logger
|
||||
_HANDLER_NAME = handler_name
|
||||
target_logger.setLevel(logging.INFO)
|
||||
for existing_handler in target_logger.handlers:
|
||||
if getattr(existing_handler, "name", "") == handler_name:
|
||||
return target_logger
|
||||
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.set_name(handler_name)
|
||||
handler.setLevel(logging.INFO)
|
||||
handler.setFormatter(_build_cloud_log_formatter())
|
||||
target_logger.addHandler(handler)
|
||||
target_logger.propagate = False
|
||||
return target_logger
|
||||
|
||||
|
||||
def _is_modal_container_runtime() -> bool:
|
||||
"""Return whether the current process is executing inside a Modal container."""
|
||||
return os.getenv("MODAL_IS_REMOTE") == "1" or bool(os.getenv("MODAL_TASK_ID"))
|
||||
|
||||
|
||||
def _cloud_formatter() -> logging.Formatter:
|
||||
"""Return the configured formatter used for cloud phase trace lines."""
|
||||
for existing_handler in _TARGET_LOGGER.handlers:
|
||||
if getattr(existing_handler, "name", "") == _HANDLER_NAME:
|
||||
formatter = existing_handler.formatter
|
||||
if formatter is not None:
|
||||
return formatter
|
||||
return _build_cloud_log_formatter()
|
||||
|
||||
|
||||
def _emit_cloud_info(message: str, *args: Any) -> None:
|
||||
"""Emit an info line and mirror it to stdout inside Modal containers."""
|
||||
if not _is_modal_container_runtime():
|
||||
_TARGET_LOGGER.info(message, *args)
|
||||
return
|
||||
|
||||
record = _TARGET_LOGGER.makeRecord(
|
||||
_TARGET_LOGGER.name,
|
||||
logging.INFO,
|
||||
__file__,
|
||||
0,
|
||||
message,
|
||||
args,
|
||||
exc_info=None,
|
||||
)
|
||||
print(_cloud_formatter().format(record), file=sys.stdout, flush=True)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _timed_phase(phase: str, **fields: Any) -> Iterator[None]:
|
||||
"""Log a start/finish pair with elapsed time for a named execution phase."""
|
||||
field_suffix = ""
|
||||
if fields:
|
||||
rendered_fields = " ".join(f"{key}={value}" for key, value in fields.items())
|
||||
field_suffix = f" {rendered_fields}"
|
||||
phase_started_at = time.perf_counter()
|
||||
_emit_cloud_info("Starting %s%s", phase, field_suffix)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_emit_cloud_info(
|
||||
"Finished %s in %.3fs%s",
|
||||
phase,
|
||||
time.perf_counter() - phase_started_at,
|
||||
field_suffix,
|
||||
)
|
||||
+2
-3
@@ -24,8 +24,8 @@ try:
|
||||
_execute_payload_with_output_capture,
|
||||
_fail_remote_invocation,
|
||||
)
|
||||
from .cloud_mapped_execution import _execute_mapped_subgraph_payload
|
||||
from .cloud_prompt_execution import (
|
||||
_execute_mapped_subgraph_payload,
|
||||
execute_node_locally,
|
||||
execute_subgraph_locally,
|
||||
)
|
||||
@@ -50,8 +50,8 @@ except ImportError: # pragma: no cover - flat Modal-container import.
|
||||
_execute_payload_with_output_capture,
|
||||
_fail_remote_invocation,
|
||||
)
|
||||
from cloud_mapped_execution import _execute_mapped_subgraph_payload
|
||||
from cloud_prompt_execution import (
|
||||
_execute_mapped_subgraph_payload,
|
||||
execute_node_locally,
|
||||
execute_subgraph_locally,
|
||||
)
|
||||
@@ -384,4 +384,3 @@ def _stream_remote_payload_events(
|
||||
event_buffer.dropped_progress_events,
|
||||
payload.get("component_id"),
|
||||
)
|
||||
|
||||
|
||||
+52
-136
@@ -18,8 +18,6 @@ import tempfile
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
@@ -278,18 +276,26 @@ except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
||||
configure_cloud_node_output_cache_hooks,
|
||||
)
|
||||
try: # noqa: E402 - support package and flat Modal-container imports.
|
||||
from .cloud_mapped_execution import (
|
||||
_aggregate_mapped_phase_outputs,
|
||||
_build_phase_subgraph_payload,
|
||||
_execute_mapped_subgraph_payload,
|
||||
_mapped_phase_definition,
|
||||
_merge_static_and_mapped_outputs,
|
||||
_merge_static_or_mapped_values,
|
||||
_shared_subgraph_payload_fields,
|
||||
_split_phase_outputs,
|
||||
)
|
||||
from .cloud_prompt_validation import configure_cloud_prompt_validation_error
|
||||
from .cloud_prompt_execution import (
|
||||
CloudPromptExecutionHooks,
|
||||
_ReusablePromptExecutorState,
|
||||
_aggregate_mapped_phase_outputs,
|
||||
_apply_boundary_inputs,
|
||||
_boundary_input_cache_signature,
|
||||
_build_phase_subgraph_payload,
|
||||
_coerce_primitive_prompt_input_value,
|
||||
_coerce_prompt_primitive_input_values,
|
||||
_collapse_cache_slot,
|
||||
_copy_json_safe_prompt_metadata,
|
||||
_execute_mapped_subgraph_payload,
|
||||
_execute_node_locally_raw,
|
||||
_execute_prompt_executor_compat,
|
||||
_execute_subgraph_prompt,
|
||||
@@ -301,9 +307,6 @@ try: # noqa: E402 - support package and flat Modal-container imports.
|
||||
_invoke_original_node,
|
||||
_is_link,
|
||||
_log_prompt_executor_failure_details,
|
||||
_mapped_phase_definition,
|
||||
_merge_static_and_mapped_outputs,
|
||||
_merge_static_or_mapped_values,
|
||||
_node_input_type_map,
|
||||
_node_input_types,
|
||||
_node_required_input_names,
|
||||
@@ -316,9 +319,7 @@ try: # noqa: E402 - support package and flat Modal-container imports.
|
||||
_reset_prompt_executor_request_state,
|
||||
_resolve_required_subgraph_nodes,
|
||||
_serialize_prompt_executor_cache_scope,
|
||||
_shared_subgraph_payload_fields,
|
||||
_short_circuit_restored_session_output_subgraph,
|
||||
_split_phase_outputs,
|
||||
_summarize_suspicious_prompt_inputs,
|
||||
_temporary_node_mapping,
|
||||
_temporary_progress_hook,
|
||||
@@ -333,18 +334,26 @@ try: # noqa: E402 - support package and flat Modal-container imports.
|
||||
execute_subgraph_locally,
|
||||
)
|
||||
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
||||
from cloud_mapped_execution import (
|
||||
_aggregate_mapped_phase_outputs,
|
||||
_build_phase_subgraph_payload,
|
||||
_execute_mapped_subgraph_payload,
|
||||
_mapped_phase_definition,
|
||||
_merge_static_and_mapped_outputs,
|
||||
_merge_static_or_mapped_values,
|
||||
_shared_subgraph_payload_fields,
|
||||
_split_phase_outputs,
|
||||
)
|
||||
from cloud_prompt_validation import configure_cloud_prompt_validation_error
|
||||
from cloud_prompt_execution import (
|
||||
CloudPromptExecutionHooks,
|
||||
_ReusablePromptExecutorState,
|
||||
_aggregate_mapped_phase_outputs,
|
||||
_apply_boundary_inputs,
|
||||
_boundary_input_cache_signature,
|
||||
_build_phase_subgraph_payload,
|
||||
_coerce_primitive_prompt_input_value,
|
||||
_coerce_prompt_primitive_input_values,
|
||||
_collapse_cache_slot,
|
||||
_copy_json_safe_prompt_metadata,
|
||||
_execute_mapped_subgraph_payload,
|
||||
_execute_node_locally_raw,
|
||||
_execute_prompt_executor_compat,
|
||||
_execute_subgraph_prompt,
|
||||
@@ -356,9 +365,6 @@ except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
||||
_invoke_original_node,
|
||||
_is_link,
|
||||
_log_prompt_executor_failure_details,
|
||||
_mapped_phase_definition,
|
||||
_merge_static_and_mapped_outputs,
|
||||
_merge_static_or_mapped_values,
|
||||
_node_input_type_map,
|
||||
_node_input_types,
|
||||
_node_required_input_names,
|
||||
@@ -371,9 +377,7 @@ except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
||||
_reset_prompt_executor_request_state,
|
||||
_resolve_required_subgraph_nodes,
|
||||
_serialize_prompt_executor_cache_scope,
|
||||
_shared_subgraph_payload_fields,
|
||||
_short_circuit_restored_session_output_subgraph,
|
||||
_split_phase_outputs,
|
||||
_summarize_suspicious_prompt_inputs,
|
||||
_temporary_node_mapping,
|
||||
_temporary_progress_hook,
|
||||
@@ -677,6 +681,20 @@ try: # noqa: E402 - support package and flat Modal-container imports.
|
||||
snapshot_profile_store,
|
||||
volume_store,
|
||||
)
|
||||
from .cloud_execution_control import (
|
||||
_RemoteExecutionControl,
|
||||
_registered_remote_execution,
|
||||
_remote_execution_key,
|
||||
_remote_interrupt_flag_key,
|
||||
)
|
||||
from .cloud_runtime_logging import (
|
||||
_build_cloud_log_formatter,
|
||||
_cloud_formatter,
|
||||
_emit_cloud_info,
|
||||
_is_modal_container_runtime,
|
||||
_timed_phase,
|
||||
configure_cloud_runtime_logging,
|
||||
)
|
||||
except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
||||
from cloud_runtime_context import (
|
||||
interrupt_flag_store,
|
||||
@@ -687,11 +705,26 @@ except ImportError: # pragma: no cover - exercised by flat cloud imports.
|
||||
snapshot_profile_store,
|
||||
volume_store,
|
||||
)
|
||||
from cloud_execution_control import (
|
||||
_RemoteExecutionControl,
|
||||
_registered_remote_execution,
|
||||
_remote_execution_key,
|
||||
_remote_interrupt_flag_key,
|
||||
)
|
||||
from cloud_runtime_logging import (
|
||||
_build_cloud_log_formatter,
|
||||
_cloud_formatter,
|
||||
_emit_cloud_info,
|
||||
_is_modal_container_runtime,
|
||||
_timed_phase,
|
||||
configure_cloud_runtime_logging,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Cloud app and ComfyUI bootstrap state.
|
||||
_CLOUD_HANDLER_NAME = "comfyui-modal-sync-cloud-timestamped"
|
||||
configure_cloud_runtime_logging(logger, _CLOUD_HANDLER_NAME)
|
||||
|
||||
# Poisoned-container retirement state.
|
||||
_CONTAINER_TERMINATION_LOCK = threading.Lock()
|
||||
@@ -757,14 +790,6 @@ def _guard_against_existing_modal_app(settings: Any, modal_module: Any) -> None:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RemoteExecutionControl:
|
||||
"""Track interruption state for one active remote payload execution."""
|
||||
|
||||
cancellation_event: threading.Event
|
||||
interrupt_flag_key: str
|
||||
|
||||
|
||||
def _meaningful_progress_values(
|
||||
node_state: dict[str, Any]
|
||||
) -> tuple[float, float] | None:
|
||||
@@ -924,72 +949,6 @@ def _maybe_schedule_container_termination_on_error(
|
||||
|
||||
|
||||
|
||||
def _build_cloud_log_formatter() -> logging.Formatter:
|
||||
"""Return the default formatter for remote Modal-Sync logs with timestamps."""
|
||||
return logging.Formatter(
|
||||
fmt="%(asctime)s.%(msecs)03d +%(relativeCreated)07.0fms %(levelname)s [%(name)s] %(message)s",
|
||||
datefmt="%H:%M:%S",
|
||||
)
|
||||
|
||||
|
||||
def _configure_cloud_logging() -> logging.Logger:
|
||||
"""Install a dedicated timestamped handler for the cloud runtime logger."""
|
||||
logger.setLevel(logging.INFO)
|
||||
for existing_handler in logger.handlers:
|
||||
if getattr(existing_handler, "name", "") == _CLOUD_HANDLER_NAME:
|
||||
return logger
|
||||
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.set_name(_CLOUD_HANDLER_NAME)
|
||||
handler.setLevel(logging.INFO)
|
||||
handler.setFormatter(_build_cloud_log_formatter())
|
||||
logger.addHandler(handler)
|
||||
logger.propagate = False
|
||||
return logger
|
||||
|
||||
|
||||
def _is_modal_container_runtime() -> bool:
|
||||
"""Return whether the current process is executing inside a Modal container."""
|
||||
return os.getenv("MODAL_IS_REMOTE") == "1" or bool(os.getenv("MODAL_TASK_ID"))
|
||||
|
||||
|
||||
def _cloud_formatter() -> logging.Formatter:
|
||||
"""Return the configured formatter used for cloud phase trace lines."""
|
||||
for existing_handler in logger.handlers:
|
||||
if getattr(existing_handler, "name", "") == _CLOUD_HANDLER_NAME:
|
||||
formatter = existing_handler.formatter
|
||||
if formatter is not None:
|
||||
return formatter
|
||||
return _build_cloud_log_formatter()
|
||||
|
||||
|
||||
def _emit_cloud_info(message: str, *args: Any) -> None:
|
||||
"""Emit an info line through logging and mirror it to stdout inside Modal containers."""
|
||||
if not _is_modal_container_runtime():
|
||||
logger.info(message, *args)
|
||||
return
|
||||
|
||||
record = logger.makeRecord(
|
||||
logger.name,
|
||||
logging.INFO,
|
||||
__file__,
|
||||
0,
|
||||
message,
|
||||
args,
|
||||
exc_info=None,
|
||||
)
|
||||
print(_cloud_formatter().format(record), file=sys.stdout, flush=True)
|
||||
|
||||
|
||||
def _remote_execution_key(payload: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Return the registry key for one active remote execution."""
|
||||
prompt_id = str(
|
||||
payload.get("prompt_id") or payload.get("component_id") or "modal-subgraph"
|
||||
)
|
||||
component_id = str(payload.get("component_id") or "single-node")
|
||||
return prompt_id, component_id
|
||||
|
||||
|
||||
def _observe_remote_workflow_for_llm_mode(payload: dict[str, Any]) -> None:
|
||||
"""Record real workflow arrivals for container-local vLLM auto promotion."""
|
||||
if payload.get("payload_kind") == "canary":
|
||||
@@ -1007,50 +966,6 @@ def _observe_remote_workflow_for_llm_mode(payload: dict[str, Any]) -> None:
|
||||
observe_modal_workflow_execution(normalized_prompt_id)
|
||||
|
||||
|
||||
def _remote_interrupt_flag_key(prompt_id: str, component_id: str) -> str:
|
||||
"""Return the shared Modal interrupt-store key for one payload execution."""
|
||||
return f"{prompt_id}:{component_id}"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _registered_remote_execution(
|
||||
payload: dict[str, Any],
|
||||
) -> Iterator[_RemoteExecutionControl]:
|
||||
"""Prepare interruption state for one active remote execution."""
|
||||
prompt_id, component_id = _remote_execution_key(payload)
|
||||
control = _RemoteExecutionControl(
|
||||
cancellation_event=threading.Event(),
|
||||
interrupt_flag_key=_remote_interrupt_flag_key(prompt_id, component_id),
|
||||
)
|
||||
try:
|
||||
yield control
|
||||
finally:
|
||||
interrupt_store = interrupt_flag_store()
|
||||
if modal is not None and interrupt_store is not None:
|
||||
interrupt_store.pop(control.interrupt_flag_key, None)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _timed_phase(phase: str, **fields: Any) -> Iterator[None]:
|
||||
"""Log a start/finish pair with elapsed time for a named execution phase."""
|
||||
field_suffix = ""
|
||||
if fields:
|
||||
rendered_fields = " ".join(f"{key}={value}" for key, value in fields.items())
|
||||
field_suffix = f" {rendered_fields}"
|
||||
phase_started_at = time.perf_counter()
|
||||
_emit_cloud_info("Starting %s%s", phase, field_suffix)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_emit_cloud_info(
|
||||
"Finished %s in %.3fs%s",
|
||||
phase,
|
||||
time.perf_counter() - phase_started_at,
|
||||
field_suffix,
|
||||
)
|
||||
|
||||
|
||||
_configure_cloud_logging()
|
||||
|
||||
|
||||
def _stream_remote_payload_events(
|
||||
@@ -1119,6 +1034,7 @@ configure_cloud_prompt_execution_hooks(
|
||||
remote_subgraph_error=RemoteSubgraphExecutionError,
|
||||
)
|
||||
)
|
||||
configure_cloud_prompt_validation_error(RemoteSubgraphExecutionError)
|
||||
|
||||
configure_cloud_prompt_server_hooks(
|
||||
CloudPromptServerHooks(
|
||||
|
||||
@@ -682,3 +682,8 @@
|
||||
- A heartbeating file lease needs no knowledge of Hugging Face profiles or staging event dataclasses. Let `snapshot_lease.py` own process identity, tokens, heartbeat timing, stale-owner reclamation, and atomic release; the staging layer can adapt its single wait message into provider-specific UI progress.
|
||||
- Dynamic proxy classes need both lower-level payload/router services and residual workflow-gate callbacks. Keep payload registries and provider routing below the factory, inject the few residual async hooks plus static node classes after `modal_executor_node.py` finishes defining them, and preserve the old module only as a compatibility surface.
|
||||
- Factory and registry extraction changes monkeypatch ownership twice: provider-client replacements belong to `remote_executor_router.py`, while warmup contexts and cache-friendly payload maps belong to `proxy_payloads.py`. Re-exported calls remain convenient, but re-exporting their mutable dictionaries would create misleading test seams.
|
||||
## Extracted cloud helpers must resolve configured error identities dynamically
|
||||
|
||||
The stable cloud exception classes are defined by `comfyui_modal_sync_cloud.py` and injected into lower-level owners after imports finish. A helper extracted from `cloud_prompt_execution.py` must not bind `RemoteSubgraphExecutionError` with `from ... import ...` during module import, because that captures the temporary fallback class. `cloud_mapped_execution.py` delegates exception construction and phase execution through the prompt-execution module at call time, preserving both the configured exception identity and the test patch boundary without introducing an import cycle.
|
||||
|
||||
The entrypoint logger has a similar identity constraint: tests and operators inspect the `comfyui_modal_sync_cloud` logger directly. `cloud_runtime_logging.py` therefore owns formatting and phase timing while configuring the entrypoint's logger object, rather than silently switching cloud output to a new logger name.
|
||||
|
||||
@@ -73,6 +73,16 @@ def _cloud_prompt_execution_owner() -> Any:
|
||||
return importlib.import_module("cloud_prompt_execution")
|
||||
|
||||
|
||||
def _cloud_mapped_execution_owner() -> Any:
|
||||
"""Return the module that owns cloud mapped execution helpers."""
|
||||
return importlib.import_module("cloud_mapped_execution")
|
||||
|
||||
|
||||
def _cloud_execution_control_owner() -> Any:
|
||||
"""Return the module that owns active cloud execution registration."""
|
||||
return importlib.import_module("cloud_execution_control")
|
||||
|
||||
|
||||
def _cloud_streaming_owner() -> Any:
|
||||
"""Return the module that owns cloud stream buffering and payload dispatch."""
|
||||
return importlib.import_module("cloud_streaming")
|
||||
@@ -11017,8 +11027,8 @@ def test_modal_cloud_execute_mapped_subgraph_payload_injects_static_bridges(
|
||||
return (f"mapped:{hydrated_inputs['remote_input_1']}:{hydrated_inputs['static_input_0']}",)
|
||||
|
||||
monkeypatch.setattr(
|
||||
_cloud_prompt_execution_owner(),
|
||||
"_execute_subgraph_prompt",
|
||||
_cloud_mapped_execution_owner(),
|
||||
"_execute_prompt_subgraph",
|
||||
fake_execute_subgraph_prompt,
|
||||
)
|
||||
|
||||
@@ -11164,8 +11174,8 @@ def test_modal_cloud_execute_mapped_subgraph_payload_preserves_assigned_lane_id(
|
||||
return (f"mapped:{hydrated_inputs['remote_input_1']}",)
|
||||
|
||||
monkeypatch.setattr(
|
||||
_cloud_prompt_execution_owner(),
|
||||
"_execute_subgraph_prompt",
|
||||
_cloud_mapped_execution_owner(),
|
||||
"_execute_prompt_subgraph",
|
||||
fake_execute_subgraph_prompt,
|
||||
)
|
||||
|
||||
@@ -15071,9 +15081,11 @@ def test_modal_cloud_registered_execution_preserves_pre_start_interrupt_flags(
|
||||
return None
|
||||
|
||||
interrupt_flags = FakeInterruptFlags()
|
||||
monkeypatch.setattr(modal_cloud_module, "modal", object())
|
||||
monkeypatch.setattr(_cloud_execution_control_owner(), "modal", object())
|
||||
monkeypatch.setattr(
|
||||
modal_cloud_module, "interrupt_flag_store", lambda: interrupt_flags
|
||||
_cloud_execution_control_owner(),
|
||||
"interrupt_flag_store",
|
||||
lambda: interrupt_flags,
|
||||
)
|
||||
|
||||
with modal_cloud_module._registered_remote_execution(
|
||||
|
||||
Reference in New Issue
Block a user