diff --git a/.gitignore b/.gitignore index a90cf2f..8137ee1 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ !/cloud_app_guard.py !/cloud_durable_invocation.py !/cloud_image_env.py +!/cloud_prompt_server_shims.py !/cloud_runtime_context.py !/cloud_session_bridge.py !/comfyui_modal_sync_cloud.py diff --git a/README.md b/README.md index 628a32e..c0bbbf6 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`, and session-bridge storage, replay, and snapshot-profile caching live in `cloud_session_bridge.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`, 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`, and headless ComfyUI server compatibility plus progress tracing live in `cloud_prompt_server_shims.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. ### Vast.ai API Simulator diff --git a/cloud_prompt_server_shims.py b/cloud_prompt_server_shims.py new file mode 100644 index 0000000..7dea1d3 --- /dev/null +++ b/cloud_prompt_server_shims.py @@ -0,0 +1,520 @@ +"""Headless ComfyUI prompt-server shims and execution event tracing.""" + +from __future__ import annotations + +from dataclasses import dataclass +from io import BytesIO +import logging +import threading +import time +from typing import Any, Callable + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class CloudPromptServerHooks: + """Callbacks supplied by the cloud orchestration and execution owners.""" + + collapse_cache_slot: Callable[[Any, bool], Any] + emit_cloud_info: Callable[..., None] + meaningful_progress_values: Callable[[dict[str, Any]], tuple[float, float] | None] + + +_PROMPT_SERVER_HOOKS: CloudPromptServerHooks | None = None + + +def configure_cloud_prompt_server_hooks(hooks: CloudPromptServerHooks) -> None: + """Install prompt-server callbacks without importing upward into the entrypoint.""" + global _PROMPT_SERVER_HOOKS + _PROMPT_SERVER_HOOKS = hooks + + +def _prompt_server_hooks() -> CloudPromptServerHooks: + """Return configured callbacks or fail with a clear import-order error.""" + if _PROMPT_SERVER_HOOKS is None: + raise RuntimeError("Cloud prompt-server hooks have not been configured.") + return _PROMPT_SERVER_HOOKS + + +def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any: + """Delegate cache-slot normalization to prompt execution.""" + return _prompt_server_hooks().collapse_cache_slot(slot_values, is_list) + + +def _emit_cloud_info(message: str, *args: Any) -> None: + """Delegate timestamped cloud logging to the stable entrypoint.""" + _prompt_server_hooks().emit_cloud_info(message, *args) + + +def _meaningful_progress_values( + node_state: dict[str, Any], +) -> tuple[float, float] | None: + """Delegate progress-state filtering to the stable entrypoint.""" + return _prompt_server_hooks().meaningful_progress_values(node_state) + + +class _NullPromptServer: + """Minimal PromptExecutor server stub for headless subgraph execution.""" + + def __init__(self) -> None: + """Initialize the no-op prompt server state.""" + self.client_id: str | None = None + self.last_node_id: str | None = None + self.last_prompt_id: str | None = None + + def send_sync( + self, event: str, data: dict[str, Any], client_id: str | None + ) -> None: + """Discard PromptExecutor progress and status events.""" + logger.debug( + "Suppressed remote prompt event %s for client %s.", event, client_id + ) + + +class _HeadlessPromptQueue: + """In-memory PromptQueue compatibility surface for remote custom-node hooks.""" + + def __init__(self) -> None: + """Initialize the queue collections exposed by ComfyUI's PromptQueue.""" + self._lock = threading.RLock() + self.queue: list[Any] = [] + self.currently_running: dict[int, Any] = {} + self.history: dict[str, Any] = {} + self.flags: dict[str, Any] = {} + + def put(self, item: Any) -> None: + """Retain a custom-node requeue request for compatibility diagnostics.""" + with self._lock: + self.queue.append(item) + logger.warning( + "A remote custom node queued headless follow-up work; the request is retained " + "for compatibility but has no background queue consumer." + ) + + def get_current_queue(self) -> tuple[list[Any], list[Any]]: + """Return snapshots of running and pending compatibility queue items.""" + with self._lock: + return list(self.currently_running.values()), list(self.queue) + + def get_current_queue_volatile(self) -> tuple[list[Any], list[Any]]: + """Return shallow running and pending queue snapshots.""" + return self.get_current_queue() + + def get_tasks_remaining(self) -> int: + """Return the number of running and pending compatibility tasks.""" + with self._lock: + return len(self.currently_running) + len(self.queue) + + def set_flag(self, name: str, data: Any) -> None: + """Store a queue flag for custom nodes that inspect PromptQueue state.""" + with self._lock: + self.flags[name] = data + + def get_flags(self, reset: bool = True) -> dict[str, Any]: + """Return queue flags and optionally clear them.""" + with self._lock: + flags = dict(self.flags) + if reset: + self.flags.clear() + return flags + + +class _HeadlessPromptServerInstance: + """Minimal PromptServer.instance replacement for custom-node import side effects.""" + + def __init__(self, node_replace_manager: Any | None) -> None: + """Initialize route registration and no-op websocket state.""" + from aiohttp import web + + self.routes = web.RouteTableDef() + self.app = web.Application() + self.node_replace_manager = node_replace_manager + self.supports = ["custom_nodes_from_web"] + self.client_id: str | None = None + self.last_node_id: str | None = None + self.number = 0 + self.prompt_queue = _HeadlessPromptQueue() + self.on_prompt_handlers: list[Any] = [] + + def queue_updated(self) -> None: + """Accept PromptQueue update notifications without a websocket client.""" + logger.debug("Suppressed headless remote prompt queue update.") + + async def send( + self, event: str, data: dict[str, Any], sid: str | None = None + ) -> None: + """Discard async websocket sends from import-time custom-node helpers.""" + logger.debug( + "Suppressed headless remote prompt event %s for client %s.", event, sid + ) + + def send_sync( + self, event: str, data: dict[str, Any], sid: str | None = None + ) -> None: + """Discard sync websocket sends from import-time custom-node helpers.""" + logger.debug( + "Suppressed headless remote prompt event %s for client %s.", event, sid + ) + + def send_progress_text( + self, + text: bytes | bytearray | str, + node_id: str, + sid: str | None = None, + ) -> None: + """Discard node progress text emitted by headless remote node execution.""" + del text + logger.debug( + "Suppressed headless remote prompt text for node %s client %s.", + node_id, + sid, + ) + + def add_on_prompt_handler(self, handler: Any) -> None: + """Record prompt handlers registered by custom nodes during import.""" + self.on_prompt_handlers.append(handler) + + +class _TracingPromptServer(_NullPromptServer): + """PromptExecutor server stub that records coarse per-node execution timings.""" + + def __init__( + self, + prompt_id: str, + prompt: dict[str, Any], + status_callback: Callable[[dict[str, Any]], None] | None = None, + ) -> None: + """Initialize timing state for a specific prompt execution.""" + super().__init__() + self.prompt_id = prompt_id + self.prompt = prompt + self._status_callback = status_callback + self._active_node_id: str | None = None + self._active_node_started_at: float | None = None + self.last_prompt_id = prompt_id + self._boundary_outputs_by_node_id: dict[str, list[dict[str, Any]]] = {} + self._lookup_cache_entry: Callable[[str], Any | None] | None = None + self._published_boundary_outputs: set[tuple[str, int]] = set() + + def _classify_node_role(self, class_type: str) -> str: + """Return a coarse role name for a node class.""" + normalized = class_type.lower() + if "loader" in normalized or normalized in {"clipvisionencode"}: + return "model_load" + if "ksampler" in normalized or "sampler" in normalized: + return "sampling" + if "encode" in normalized: + return "conditioning" + return "node" + + def _log_node_finish(self, reason: str) -> None: + """Emit a timing line for the currently active node when one is running.""" + if self._active_node_id is None or self._active_node_started_at is None: + return + + node_id = self._active_node_id + node_info = self.prompt.get(node_id, {}) + class_type = str(node_info.get("class_type", "")) + role = self._classify_node_role(class_type) + elapsed_seconds = time.perf_counter() - self._active_node_started_at + _emit_cloud_info( + "Remote node %s class_type=%s role=%s finished in %.3fs reason=%s", + node_id, + class_type, + role, + elapsed_seconds, + reason, + ) + self._active_node_id = None + self._active_node_started_at = None + + def emit_preview_update( + self, + *, + node_id: str, + preview_image: Any, + ) -> None: + """Publish one preview image update through the status callback.""" + if self._status_callback is None: + return + + try: + image_type, image, max_size = preview_image + image_buffer = BytesIO() + save_kwargs: dict[str, Any] = {"format": image_type} + if image_type == "JPEG": + save_kwargs["quality"] = 95 + elif image_type == "PNG": + save_kwargs["compress_level"] = 1 + image.save(image_buffer, **save_kwargs) + except Exception: + logger.exception( + "Failed to serialize remote preview image for node %s.", node_id + ) + return + + try: + from comfy_execution.progress import get_progress_state + + registry = get_progress_state() + display_node_id = registry.dynprompt.get_display_node_id(node_id) + parent_node_id = registry.dynprompt.get_parent_node_id(node_id) + real_node_id = registry.dynprompt.get_real_node_id(node_id) + except Exception: + logger.exception( + "Failed to resolve preview metadata for remote node %s.", node_id + ) + display_node_id = node_id + parent_node_id = None + real_node_id = node_id + + self._status_callback( + { + "event_type": "preview", + "node_id": str(node_id), + "display_node_id": ( + str(display_node_id) if display_node_id is not None else None + ), + "parent_node_id": ( + str(parent_node_id) if parent_node_id is not None else None + ), + "real_node_id": str(real_node_id) if real_node_id is not None else None, + "image_type": str(image_type), + "image_bytes": image_buffer.getvalue(), + "max_size": int(max_size) if max_size is not None else None, + } + ) + + def configure_boundary_output_stream( + self, + *, + boundary_outputs: list[dict[str, Any]], + lookup_cache_entry: Callable[[str], Any | None], + ) -> None: + """Configure streamed remote boundary-output publication for this execution.""" + outputs_by_node_id: dict[str, list[dict[str, Any]]] = {} + for boundary_output in boundary_outputs: + preview_target_node_ids = [ + str(node_id) + for node_id in boundary_output.get("preview_target_node_ids", []) + if str(node_id) + ] + if not preview_target_node_ids: + continue + if str(boundary_output.get("io_type")) != "IMAGE": + continue + node_id = str(boundary_output["node_id"]) + outputs_by_node_id.setdefault(node_id, []).append(boundary_output) + + self._boundary_outputs_by_node_id = outputs_by_node_id + self._lookup_cache_entry = lookup_cache_entry + self._published_boundary_outputs.clear() + + def _emit_boundary_outputs_for_node(self, node_id: str | None) -> None: + """Publish configured boundary image outputs for one completed node once.""" + if ( + node_id is None + or self._status_callback is None + or self._lookup_cache_entry is None + ): + return + + boundary_outputs = self._boundary_outputs_by_node_id.get(str(node_id), []) + if not boundary_outputs: + return + + cache_entry = self._lookup_cache_entry(str(node_id)) + if cache_entry is None: + return + + cache_outputs = getattr(cache_entry, "outputs", None) + if not isinstance(cache_outputs, (list, tuple)): + return + + for boundary_output in boundary_outputs: + output_index = int(boundary_output["output_index"]) + publication_key = (str(node_id), output_index) + if publication_key in self._published_boundary_outputs: + continue + if output_index >= len(cache_outputs): + continue + + preview_target_node_ids = [ + str(target_node_id) + for target_node_id in boundary_output.get("preview_target_node_ids", []) + if str(target_node_id) + ] + if not preview_target_node_ids: + continue + + self._status_callback( + { + "event_type": "boundary_output", + "node_id": str(node_id), + "output_index": output_index, + "io_type": str(boundary_output.get("io_type", "")), + "is_list": bool(boundary_output.get("is_list", False)), + "preview_target_node_ids": preview_target_node_ids, + "value": _collapse_cache_slot( + slot_values=cache_outputs[output_index], + is_list=bool(boundary_output.get("is_list", False)), + ), + } + ) + self._published_boundary_outputs.add(publication_key) + + def send_sync( + self, event: str, data: dict[str, Any], client_id: str | None + ) -> None: + """Track per-node timing transitions from PromptExecutor progress events.""" + if event == "modal_llm_progress": + if self._status_callback is None: + return + node_id = str(data.get("node_id") or self._active_node_id or "") + if not node_id: + return + progress_event = { + "event_type": "node_progress", + "node_id": node_id, + "display_node_id": node_id, + "real_node_id": node_id, + "value": float(data.get("value", 0.0)), + "max": float(data.get("max", 1.0)), + "stage": str(data.get("stage") or ""), + "message": str(data.get("message") or ""), + "indeterminate": bool(data.get("indeterminate", False)), + } + for field_name in ( + "unit", + "elapsed_seconds", + "time_to_first_token_seconds", + "tokens_per_second", + ): + if data.get(field_name) is not None: + progress_event[field_name] = data[field_name] + self._status_callback(progress_event) + return + + if event == "executing": + next_node_id = data.get("node") + if next_node_id != self._active_node_id: + self._emit_boundary_outputs_for_node(self._active_node_id) + self._log_node_finish(reason="next_node") + if next_node_id is not None and next_node_id != self._active_node_id: + node_info = self.prompt.get(str(next_node_id), {}) + class_type = str(node_info.get("class_type", "")) + role = self._classify_node_role(class_type) + if self._status_callback is not None: + self._status_callback( + { + "phase": "executing", + "active_node_id": str(next_node_id), + "active_node_class_type": class_type, + "active_node_role": role, + } + ) + self._active_node_id = str(next_node_id) + self._active_node_started_at = time.perf_counter() + self.last_node_id = self._active_node_id + _emit_cloud_info( + "Remote node %s class_type=%s role=%s started", + self._active_node_id, + class_type, + role, + ) + return + + if event == "progress_state": + if self._status_callback is None: + return + + nodes_payload = data.get("nodes") + if not isinstance(nodes_payload, dict): + return + + tracked_node_id = self._active_node_id + tracked_node_state: dict[str, Any] | None = None + if tracked_node_id is not None: + candidate_state = nodes_payload.get(tracked_node_id) + if isinstance(candidate_state, dict): + tracked_node_state = candidate_state + + if tracked_node_state is None: + for node_state in nodes_payload.values(): + if ( + isinstance(node_state, dict) + and node_state.get("state") == "running" + ): + tracked_node_state = node_state + break + + if tracked_node_state is None: + return + + display_node_id = tracked_node_state.get("display_node_id") + real_node_id = tracked_node_state.get("real_node_id") + reported_node_id = ( + display_node_id or real_node_id or tracked_node_state.get("node_id") + ) + if reported_node_id is None: + return + progress_values = _meaningful_progress_values(tracked_node_state) + if progress_values is None: + logger.debug( + "Ignoring non-meaningful remote progress_state for node_id=%s state=%s value=%r max=%r.", + reported_node_id, + tracked_node_state.get("state"), + tracked_node_state.get("value"), + tracked_node_state.get("max"), + ) + return + progress_value, max_value = progress_values + + self._status_callback( + { + "event_type": "node_progress", + "node_id": str(reported_node_id), + "display_node_id": ( + str(display_node_id) if display_node_id is not None else None + ), + "real_node_id": str(real_node_id) + if real_node_id is not None + else None, + "value": progress_value, + "max": max_value, + } + ) + return + + if event == "executed": + executed_node_id = data.get("node") + if ( + executed_node_id is not None + and str(executed_node_id) == self._active_node_id + ): + self._log_node_finish(reason="executed") + if self._status_callback is not None and data.get("output") is not None: + self._status_callback( + { + "event_type": "executed", + "node_id": str(data.get("node")), + "display_node_id": ( + str(data["display_node"]) + if data.get("display_node") is not None + else None + ), + "output": data.get("output"), + } + ) + return + + if event in {"execution_error", "execution_interrupted", "execution_success"}: + self._emit_boundary_outputs_for_node(self._active_node_id) + self._log_node_finish(reason=event) + if self._status_callback is not None: + self._status_callback({"phase": event}) + return + + super().send_sync(event, data, client_id) + diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 5c3067b..646cd8e 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -118,6 +118,24 @@ from settings import ( # noqa: E402 - paths are bootstrapped above. get_settings, modal_deployment_app_name, ) +try: # noqa: E402 - support package and flat Modal-container imports. + from .cloud_prompt_server_shims import ( + CloudPromptServerHooks, + _HeadlessPromptQueue, + _HeadlessPromptServerInstance, + _NullPromptServer, + _TracingPromptServer, + configure_cloud_prompt_server_hooks, + ) +except ImportError: # pragma: no cover - exercised by flat cloud imports. + from cloud_prompt_server_shims import ( + CloudPromptServerHooks, + _HeadlessPromptQueue, + _HeadlessPromptServerInstance, + _NullPromptServer, + _TracingPromptServer, + configure_cloud_prompt_server_hooks, + ) try: # noqa: E402 - support package and flat Modal-container imports. from .cloud_session_bridge import ( CloudSessionBridgeHooks, @@ -600,469 +618,6 @@ def _maybe_schedule_container_termination_on_error( return True -class _NullPromptServer: - """Minimal PromptExecutor server stub for headless subgraph execution.""" - - def __init__(self) -> None: - """Initialize the no-op prompt server state.""" - self.client_id: str | None = None - self.last_node_id: str | None = None - self.last_prompt_id: str | None = None - - def send_sync( - self, event: str, data: dict[str, Any], client_id: str | None - ) -> None: - """Discard PromptExecutor progress and status events.""" - logger.debug( - "Suppressed remote prompt event %s for client %s.", event, client_id - ) - - -class _HeadlessPromptQueue: - """In-memory PromptQueue compatibility surface for remote custom-node hooks.""" - - def __init__(self) -> None: - """Initialize the queue collections exposed by ComfyUI's PromptQueue.""" - self._lock = threading.RLock() - self.queue: list[Any] = [] - self.currently_running: dict[int, Any] = {} - self.history: dict[str, Any] = {} - self.flags: dict[str, Any] = {} - - def put(self, item: Any) -> None: - """Retain a custom-node requeue request for compatibility diagnostics.""" - with self._lock: - self.queue.append(item) - logger.warning( - "A remote custom node queued headless follow-up work; the request is retained " - "for compatibility but has no background queue consumer." - ) - - def get_current_queue(self) -> tuple[list[Any], list[Any]]: - """Return snapshots of running and pending compatibility queue items.""" - with self._lock: - return list(self.currently_running.values()), list(self.queue) - - def get_current_queue_volatile(self) -> tuple[list[Any], list[Any]]: - """Return shallow running and pending queue snapshots.""" - return self.get_current_queue() - - def get_tasks_remaining(self) -> int: - """Return the number of running and pending compatibility tasks.""" - with self._lock: - return len(self.currently_running) + len(self.queue) - - def set_flag(self, name: str, data: Any) -> None: - """Store a queue flag for custom nodes that inspect PromptQueue state.""" - with self._lock: - self.flags[name] = data - - def get_flags(self, reset: bool = True) -> dict[str, Any]: - """Return queue flags and optionally clear them.""" - with self._lock: - flags = dict(self.flags) - if reset: - self.flags.clear() - return flags - - -class _HeadlessPromptServerInstance: - """Minimal PromptServer.instance replacement for custom-node import side effects.""" - - def __init__(self, node_replace_manager: Any | None) -> None: - """Initialize route registration and no-op websocket state.""" - from aiohttp import web - - self.routes = web.RouteTableDef() - self.app = web.Application() - self.node_replace_manager = node_replace_manager - self.supports = ["custom_nodes_from_web"] - self.client_id: str | None = None - self.last_node_id: str | None = None - self.number = 0 - self.prompt_queue = _HeadlessPromptQueue() - self.on_prompt_handlers: list[Any] = [] - - def queue_updated(self) -> None: - """Accept PromptQueue update notifications without a websocket client.""" - logger.debug("Suppressed headless remote prompt queue update.") - - async def send( - self, event: str, data: dict[str, Any], sid: str | None = None - ) -> None: - """Discard async websocket sends from import-time custom-node helpers.""" - logger.debug( - "Suppressed headless remote prompt event %s for client %s.", event, sid - ) - - def send_sync( - self, event: str, data: dict[str, Any], sid: str | None = None - ) -> None: - """Discard sync websocket sends from import-time custom-node helpers.""" - logger.debug( - "Suppressed headless remote prompt event %s for client %s.", event, sid - ) - - def send_progress_text( - self, - text: bytes | bytearray | str, - node_id: str, - sid: str | None = None, - ) -> None: - """Discard node progress text emitted by headless remote node execution.""" - del text - logger.debug( - "Suppressed headless remote prompt text for node %s client %s.", - node_id, - sid, - ) - - def add_on_prompt_handler(self, handler: Any) -> None: - """Record prompt handlers registered by custom nodes during import.""" - self.on_prompt_handlers.append(handler) - - -class _TracingPromptServer(_NullPromptServer): - """PromptExecutor server stub that records coarse per-node execution timings.""" - - def __init__( - self, - prompt_id: str, - prompt: dict[str, Any], - status_callback: Callable[[dict[str, Any]], None] | None = None, - ) -> None: - """Initialize timing state for a specific prompt execution.""" - super().__init__() - self.prompt_id = prompt_id - self.prompt = prompt - self._status_callback = status_callback - self._active_node_id: str | None = None - self._active_node_started_at: float | None = None - self.last_prompt_id = prompt_id - self._boundary_outputs_by_node_id: dict[str, list[dict[str, Any]]] = {} - self._lookup_cache_entry: Callable[[str], Any | None] | None = None - self._published_boundary_outputs: set[tuple[str, int]] = set() - - def _classify_node_role(self, class_type: str) -> str: - """Return a coarse role name for a node class.""" - normalized = class_type.lower() - if "loader" in normalized or normalized in {"clipvisionencode"}: - return "model_load" - if "ksampler" in normalized or "sampler" in normalized: - return "sampling" - if "encode" in normalized: - return "conditioning" - return "node" - - def _log_node_finish(self, reason: str) -> None: - """Emit a timing line for the currently active node when one is running.""" - if self._active_node_id is None or self._active_node_started_at is None: - return - - node_id = self._active_node_id - node_info = self.prompt.get(node_id, {}) - class_type = str(node_info.get("class_type", "")) - role = self._classify_node_role(class_type) - elapsed_seconds = time.perf_counter() - self._active_node_started_at - _emit_cloud_info( - "Remote node %s class_type=%s role=%s finished in %.3fs reason=%s", - node_id, - class_type, - role, - elapsed_seconds, - reason, - ) - self._active_node_id = None - self._active_node_started_at = None - - def emit_preview_update( - self, - *, - node_id: str, - preview_image: Any, - ) -> None: - """Publish one preview image update through the status callback.""" - if self._status_callback is None: - return - - try: - image_type, image, max_size = preview_image - image_buffer = BytesIO() - save_kwargs: dict[str, Any] = {"format": image_type} - if image_type == "JPEG": - save_kwargs["quality"] = 95 - elif image_type == "PNG": - save_kwargs["compress_level"] = 1 - image.save(image_buffer, **save_kwargs) - except Exception: - logger.exception( - "Failed to serialize remote preview image for node %s.", node_id - ) - return - - try: - from comfy_execution.progress import get_progress_state - - registry = get_progress_state() - display_node_id = registry.dynprompt.get_display_node_id(node_id) - parent_node_id = registry.dynprompt.get_parent_node_id(node_id) - real_node_id = registry.dynprompt.get_real_node_id(node_id) - except Exception: - logger.exception( - "Failed to resolve preview metadata for remote node %s.", node_id - ) - display_node_id = node_id - parent_node_id = None - real_node_id = node_id - - self._status_callback( - { - "event_type": "preview", - "node_id": str(node_id), - "display_node_id": ( - str(display_node_id) if display_node_id is not None else None - ), - "parent_node_id": ( - str(parent_node_id) if parent_node_id is not None else None - ), - "real_node_id": str(real_node_id) if real_node_id is not None else None, - "image_type": str(image_type), - "image_bytes": image_buffer.getvalue(), - "max_size": int(max_size) if max_size is not None else None, - } - ) - - def configure_boundary_output_stream( - self, - *, - boundary_outputs: list[dict[str, Any]], - lookup_cache_entry: Callable[[str], Any | None], - ) -> None: - """Configure streamed remote boundary-output publication for this execution.""" - outputs_by_node_id: dict[str, list[dict[str, Any]]] = {} - for boundary_output in boundary_outputs: - preview_target_node_ids = [ - str(node_id) - for node_id in boundary_output.get("preview_target_node_ids", []) - if str(node_id) - ] - if not preview_target_node_ids: - continue - if str(boundary_output.get("io_type")) != "IMAGE": - continue - node_id = str(boundary_output["node_id"]) - outputs_by_node_id.setdefault(node_id, []).append(boundary_output) - - self._boundary_outputs_by_node_id = outputs_by_node_id - self._lookup_cache_entry = lookup_cache_entry - self._published_boundary_outputs.clear() - - def _emit_boundary_outputs_for_node(self, node_id: str | None) -> None: - """Publish configured boundary image outputs for one completed node once.""" - if ( - node_id is None - or self._status_callback is None - or self._lookup_cache_entry is None - ): - return - - boundary_outputs = self._boundary_outputs_by_node_id.get(str(node_id), []) - if not boundary_outputs: - return - - cache_entry = self._lookup_cache_entry(str(node_id)) - if cache_entry is None: - return - - cache_outputs = getattr(cache_entry, "outputs", None) - if not isinstance(cache_outputs, (list, tuple)): - return - - for boundary_output in boundary_outputs: - output_index = int(boundary_output["output_index"]) - publication_key = (str(node_id), output_index) - if publication_key in self._published_boundary_outputs: - continue - if output_index >= len(cache_outputs): - continue - - preview_target_node_ids = [ - str(target_node_id) - for target_node_id in boundary_output.get("preview_target_node_ids", []) - if str(target_node_id) - ] - if not preview_target_node_ids: - continue - - self._status_callback( - { - "event_type": "boundary_output", - "node_id": str(node_id), - "output_index": output_index, - "io_type": str(boundary_output.get("io_type", "")), - "is_list": bool(boundary_output.get("is_list", False)), - "preview_target_node_ids": preview_target_node_ids, - "value": _collapse_cache_slot( - slot_values=cache_outputs[output_index], - is_list=bool(boundary_output.get("is_list", False)), - ), - } - ) - self._published_boundary_outputs.add(publication_key) - - def send_sync( - self, event: str, data: dict[str, Any], client_id: str | None - ) -> None: - """Track per-node timing transitions from PromptExecutor progress events.""" - if event == "modal_llm_progress": - if self._status_callback is None: - return - node_id = str(data.get("node_id") or self._active_node_id or "") - if not node_id: - return - progress_event = { - "event_type": "node_progress", - "node_id": node_id, - "display_node_id": node_id, - "real_node_id": node_id, - "value": float(data.get("value", 0.0)), - "max": float(data.get("max", 1.0)), - "stage": str(data.get("stage") or ""), - "message": str(data.get("message") or ""), - "indeterminate": bool(data.get("indeterminate", False)), - } - for field_name in ( - "unit", - "elapsed_seconds", - "time_to_first_token_seconds", - "tokens_per_second", - ): - if data.get(field_name) is not None: - progress_event[field_name] = data[field_name] - self._status_callback(progress_event) - return - - if event == "executing": - next_node_id = data.get("node") - if next_node_id != self._active_node_id: - self._emit_boundary_outputs_for_node(self._active_node_id) - self._log_node_finish(reason="next_node") - if next_node_id is not None and next_node_id != self._active_node_id: - node_info = self.prompt.get(str(next_node_id), {}) - class_type = str(node_info.get("class_type", "")) - role = self._classify_node_role(class_type) - if self._status_callback is not None: - self._status_callback( - { - "phase": "executing", - "active_node_id": str(next_node_id), - "active_node_class_type": class_type, - "active_node_role": role, - } - ) - self._active_node_id = str(next_node_id) - self._active_node_started_at = time.perf_counter() - self.last_node_id = self._active_node_id - _emit_cloud_info( - "Remote node %s class_type=%s role=%s started", - self._active_node_id, - class_type, - role, - ) - return - - if event == "progress_state": - if self._status_callback is None: - return - - nodes_payload = data.get("nodes") - if not isinstance(nodes_payload, dict): - return - - tracked_node_id = self._active_node_id - tracked_node_state: dict[str, Any] | None = None - if tracked_node_id is not None: - candidate_state = nodes_payload.get(tracked_node_id) - if isinstance(candidate_state, dict): - tracked_node_state = candidate_state - - if tracked_node_state is None: - for node_state in nodes_payload.values(): - if ( - isinstance(node_state, dict) - and node_state.get("state") == "running" - ): - tracked_node_state = node_state - break - - if tracked_node_state is None: - return - - display_node_id = tracked_node_state.get("display_node_id") - real_node_id = tracked_node_state.get("real_node_id") - reported_node_id = ( - display_node_id or real_node_id or tracked_node_state.get("node_id") - ) - if reported_node_id is None: - return - progress_values = _meaningful_progress_values(tracked_node_state) - if progress_values is None: - logger.debug( - "Ignoring non-meaningful remote progress_state for node_id=%s state=%s value=%r max=%r.", - reported_node_id, - tracked_node_state.get("state"), - tracked_node_state.get("value"), - tracked_node_state.get("max"), - ) - return - progress_value, max_value = progress_values - - self._status_callback( - { - "event_type": "node_progress", - "node_id": str(reported_node_id), - "display_node_id": ( - str(display_node_id) if display_node_id is not None else None - ), - "real_node_id": str(real_node_id) - if real_node_id is not None - else None, - "value": progress_value, - "max": max_value, - } - ) - return - - if event == "executed": - executed_node_id = data.get("node") - if ( - executed_node_id is not None - and str(executed_node_id) == self._active_node_id - ): - self._log_node_finish(reason="executed") - if self._status_callback is not None and data.get("output") is not None: - self._status_callback( - { - "event_type": "executed", - "node_id": str(data.get("node")), - "display_node_id": ( - str(data["display_node"]) - if data.get("display_node") is not None - else None - ), - "output": data.get("output"), - } - ) - return - - if event in {"execution_error", "execution_interrupted", "execution_success"}: - self._emit_boundary_outputs_for_node(self._active_node_id) - self._log_node_finish(reason=event) - if self._status_callback is not None: - self._status_callback({"phase": event}) - return - - super().send_sync(event, data, client_id) def _build_cloud_log_formatter() -> logging.Formatter: @@ -6096,6 +5651,14 @@ configure_cloud_session_bridge_hooks( ) ) +configure_cloud_prompt_server_hooks( + CloudPromptServerHooks( + collapse_cache_slot=_collapse_cache_slot, + emit_cloud_info=_emit_cloud_info, + meaningful_progress_values=_meaningful_progress_values, + ) +) + if modal is not None: # pragma: no branch - remote entrypoint configuration. settings = globals().get("__comfy_modal_settings_override__") or get_settings() diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 24b62f9..cde32ed 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -666,3 +666,4 @@ - Modal image policy is easier to verify when path filtering, environment construction, dependency layering, secrets, and class options move together. The stable entrypoint can keep constructing the actual `modal.Image` and decorators while importing these pure builders, preserving deployment shape without mixing it into runtime execution logic. - Durable invocation extraction has two stable-boundary concerns: mutable fallback stores must move with the lifecycle functions that use them, while public cloud exception classes must stay in the entrypoint. Inject the exception factories once after those classes are defined, and patch the extracted state owner in tests; compatibility re-exports are call surfaces, not mutable-state aliases. - Session-bridge replay crosses bootstrap and prompt-execution boundaries, but those upward calls do not justify importing the cloud entrypoint. Inject the narrow cloning, logging, graph, and execution callbacks after module initialization; keep bridge values, replay stacks, snapshot profiles, and their locks private to the session owner, and make tests patch that owner directly. +- Headless PromptServer compatibility is a distinct adapter boundary: custom-node imports need a queue and route surface, while PromptExecutor needs progress, preview, and boundary-output tracing. Move those shims together and inject only cache-slot normalization, progress filtering, and cloud logging so neither the shim module nor its tests depend on the monolithic entrypoint.