215 lines
7.9 KiB
Python
215 lines
7.9 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Trace ordered Comfy node execution over one loopback WebSocket session."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from dataclasses import dataclass
|
|
from importlib import import_module
|
|
from typing import Any
|
|
|
|
from tools.comfy_api import JsonObject, LoopbackComfyClient
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ComfyNodeExecutionTiming:
|
|
"""Retain one node's observed start and duration until the next event."""
|
|
|
|
node_id: str
|
|
class_type: str
|
|
started_at_ns: int
|
|
elapsed_ms: float
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ComfyExecutionTraceResult:
|
|
"""Retain one completed prompt history and its ordered node timings."""
|
|
|
|
prompt_id: str
|
|
submission_started_at_ns: int
|
|
completed_at_ns: int
|
|
cached_node_ids: tuple[str, ...]
|
|
nodes: tuple[ComfyNodeExecutionTiming, ...]
|
|
history: JsonObject
|
|
|
|
@property
|
|
def elapsed_ms(self) -> float:
|
|
"""Return observed submission-to-terminal wall time."""
|
|
|
|
return (self.completed_at_ns - self.submission_started_at_ns) / 1_000_000.0
|
|
|
|
def class_totals_ms(self) -> dict[str, float]:
|
|
"""Aggregate ordered node durations by exact runtime class type."""
|
|
|
|
totals: dict[str, float] = {}
|
|
for node in self.nodes:
|
|
totals[node.class_type] = totals.get(node.class_type, 0.0) + node.elapsed_ms
|
|
return totals
|
|
|
|
|
|
class ComfyExecutionTraceAccumulator:
|
|
"""Convert WebSocket lifecycle messages into ordered node intervals."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
prompt_id: str,
|
|
prompt: dict[str, JsonObject],
|
|
submission_started_at_ns: int,
|
|
) -> None:
|
|
"""Retain immutable prompt identity and initialize empty trace state."""
|
|
|
|
if not prompt_id:
|
|
raise ValueError("Execution trace prompt id must be non-empty.")
|
|
if submission_started_at_ns < 1:
|
|
raise ValueError("Execution trace submission timestamp must be positive.")
|
|
self._prompt_id = prompt_id
|
|
self._prompt = prompt
|
|
self._submission_started_at_ns = submission_started_at_ns
|
|
self._cached_node_ids: tuple[str, ...] = ()
|
|
self._current: tuple[str, int] | None = None
|
|
self._nodes: list[ComfyNodeExecutionTiming] = []
|
|
self._completed_at_ns: int | None = None
|
|
|
|
def observe(self, message: object, *, observed_at_ns: int) -> bool:
|
|
"""Consume one JSON message and report terminal completion."""
|
|
|
|
if observed_at_ns < self._submission_started_at_ns:
|
|
raise ValueError("Execution trace observation precedes submission.")
|
|
if not isinstance(message, dict):
|
|
return False
|
|
event_type = message.get("type")
|
|
data = message.get("data")
|
|
if not isinstance(event_type, str) or not isinstance(data, dict):
|
|
return False
|
|
message_prompt_id = data.get("prompt_id")
|
|
if message_prompt_id != self._prompt_id:
|
|
return False
|
|
if event_type == "execution_cached":
|
|
self._cached_node_ids = self._cached_nodes(data.get("nodes"))
|
|
return False
|
|
if event_type == "executing":
|
|
node_id = data.get("node")
|
|
if not isinstance(node_id, str) or not node_id:
|
|
raise ValueError("Execution trace node id is invalid.")
|
|
self._close_current(observed_at_ns)
|
|
self._current = (node_id, observed_at_ns)
|
|
return False
|
|
if event_type == "execution_error":
|
|
raise RuntimeError(
|
|
"Comfy execution trace observed an error: "
|
|
f"{data.get('exception_type')}: {data.get('exception_message')}"
|
|
)
|
|
if event_type == "execution_success":
|
|
self._close_current(observed_at_ns)
|
|
self._completed_at_ns = observed_at_ns
|
|
return True
|
|
return False
|
|
|
|
def finish(self, history: JsonObject) -> ComfyExecutionTraceResult:
|
|
"""Return a completed trace or fail on missing lifecycle evidence."""
|
|
|
|
if self._completed_at_ns is None:
|
|
raise TimeoutError("Comfy execution trace did not reach terminal success.")
|
|
if not self._nodes:
|
|
raise ValueError("Comfy execution trace observed no node starts.")
|
|
return ComfyExecutionTraceResult(
|
|
prompt_id=self._prompt_id,
|
|
submission_started_at_ns=self._submission_started_at_ns,
|
|
completed_at_ns=self._completed_at_ns,
|
|
cached_node_ids=self._cached_node_ids,
|
|
nodes=tuple(self._nodes),
|
|
history=history,
|
|
)
|
|
|
|
def _close_current(self, completed_at_ns: int) -> None:
|
|
"""Close the current node interval against one observed boundary."""
|
|
|
|
if self._current is None:
|
|
return
|
|
node_id, started_at_ns = self._current
|
|
if completed_at_ns < started_at_ns:
|
|
raise ValueError("Execution trace node completion precedes its start.")
|
|
node = self._prompt.get(node_id)
|
|
class_type = node.get("class_type") if isinstance(node, dict) else None
|
|
if not isinstance(class_type, str) or not class_type:
|
|
raise ValueError(f"Execution trace node {node_id!r} has no class type.")
|
|
self._nodes.append(
|
|
ComfyNodeExecutionTiming(
|
|
node_id=node_id,
|
|
class_type=class_type,
|
|
started_at_ns=started_at_ns,
|
|
elapsed_ms=(completed_at_ns - started_at_ns) / 1_000_000.0,
|
|
)
|
|
)
|
|
self._current = None
|
|
|
|
@staticmethod
|
|
def _cached_nodes(value: object) -> tuple[str, ...]:
|
|
"""Return canonical cached node ids from one lifecycle message."""
|
|
|
|
if not isinstance(value, list) or any(
|
|
not isinstance(item, str) for item in value
|
|
):
|
|
raise ValueError("Execution trace cached node ids are invalid.")
|
|
return tuple(value)
|
|
|
|
|
|
class LoopbackExecutionTrace:
|
|
"""Submit one prompt and capture its loopback WebSocket node lifecycle."""
|
|
|
|
def execute(
|
|
self,
|
|
client: LoopbackComfyClient,
|
|
*,
|
|
prompt: dict[str, JsonObject],
|
|
timeout: float,
|
|
) -> ComfyExecutionTraceResult:
|
|
"""Return one synchronized trace and completed history."""
|
|
|
|
if timeout <= 0.0:
|
|
raise ValueError("Execution trace timeout must be positive.")
|
|
websocket: Any = import_module("websocket")
|
|
connection: Any = websocket.create_connection(
|
|
client.websocket_url,
|
|
timeout=min(timeout, 5.0),
|
|
)
|
|
try:
|
|
submission_started_at_ns = time.perf_counter_ns()
|
|
prompt_id = client.submit(prompt)
|
|
accumulator = ComfyExecutionTraceAccumulator(
|
|
prompt_id=prompt_id,
|
|
prompt=prompt,
|
|
submission_started_at_ns=submission_started_at_ns,
|
|
)
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
connection.settimeout(min(1.0, max(deadline - time.monotonic(), 0.01)))
|
|
try:
|
|
raw_message = connection.recv()
|
|
except Exception as error:
|
|
if isinstance(error, websocket.WebSocketTimeoutException):
|
|
continue
|
|
raise
|
|
if not isinstance(raw_message, str):
|
|
continue
|
|
decoded: object = json.loads(raw_message)
|
|
if accumulator.observe(
|
|
decoded,
|
|
observed_at_ns=time.perf_counter_ns(),
|
|
):
|
|
history = client.wait_for_history(prompt_id, timeout=timeout)
|
|
return accumulator.finish(history)
|
|
raise TimeoutError(
|
|
f"Comfy execution trace exceeded {timeout}s: {prompt_id}."
|
|
)
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
LOOPBACK_EXECUTION_TRACE = LoopbackExecutionTrace()
|