Extract headless cloud prompt server shims
This commit is contained in:
Ken Simpson
+26
-463
@@ -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", "<unknown>"))
|
||||
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", "<unknown>"))
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user