From 4539bf70908cff8f0531a0804d3dc215dd63ca25 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sat, 29 Aug 2026 12:48:42 -0700 Subject: [PATCH] Split cloud execution runtime responsibilities --- .gitignore | 4 + README.md | 2 +- cloud_execution_control.py | 61 +++ cloud_mapped_execution.py | 381 +++++++++++++ cloud_prompt_execution.py | 882 ++---------------------------- cloud_prompt_validation.py | 534 ++++++++++++++++++ cloud_runtime_logging.py | 99 ++++ cloud_streaming.py | 5 +- comfyui_modal_sync_cloud.py | 188 ++----- docs/LEARNINGS.md | 5 + tests/test_modal_executor_node.py | 24 +- 11 files changed, 1199 insertions(+), 986 deletions(-) create mode 100644 cloud_execution_control.py create mode 100644 cloud_mapped_execution.py create mode 100644 cloud_prompt_validation.py create mode 100644 cloud_runtime_logging.py diff --git a/.gitignore b/.gitignore index 3f4c40d..7c55883 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/README.md b/README.md index 4e8fec0..39a40ba 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/cloud_execution_control.py b/cloud_execution_control.py new file mode 100644 index 0000000..81a0d46 --- /dev/null +++ b/cloud_execution_control.py @@ -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) diff --git a/cloud_mapped_execution.py b/cloud_mapped_execution.py new file mode 100644 index 0000000..6e16080 --- /dev/null +++ b/cloud_mapped_execution.py @@ -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, + ) + diff --git a/cloud_prompt_execution.py b/cloud_prompt_execution.py index 3c3cf10..ed3be9e 100644 --- a/cloud_prompt_execution.py +++ b/cloud_prompt_execution.py @@ -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, - ) diff --git a/cloud_prompt_validation.py b/cloud_prompt_validation.py new file mode 100644 index 0000000..c39fe72 --- /dev/null +++ b/cloud_prompt_validation.py @@ -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 + + diff --git a/cloud_runtime_logging.py b/cloud_runtime_logging.py new file mode 100644 index 0000000..6415119 --- /dev/null +++ b/cloud_runtime_logging.py @@ -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, + ) diff --git a/cloud_streaming.py b/cloud_streaming.py index 5532ec1..bbdcf90 100644 --- a/cloud_streaming.py +++ b/cloud_streaming.py @@ -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"), ) - diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index b6593de..be2f768 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -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( diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index ab4ccdb..56bd08f 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -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. diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 4ea5063..6dcfb00 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -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(