335 lines
12 KiB
Python
335 lines
12 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
|
|
|
|
"""Execute one resumable P10.3 corpus through bounded managed-Comfy segments."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from tools.attention_coupling_benchmark.manifest_types import BenchmarkManifest
|
|
from tools.attention_coupling_benchmark.results import parse_completed_outputs
|
|
from tools.comfy_api import LoopbackComfyClient
|
|
from tools.comfy_integration.loopback_port import is_loopback_port_available
|
|
from tools.comfy_integration.managed_server import ManagedComfyServer
|
|
|
|
from .blind import BlindCorpusStager
|
|
from .corpus import VisualCorpusRecorder
|
|
from .input_artifacts import VisualPositionInputs
|
|
from .matrix import SourcePosition, VisualPosition, source_positions, visual_positions
|
|
from .run_artifacts import VisualBenchmarkRunArtifacts
|
|
from .workflow import BuiltVisualWorkflow, VisualBenchmarkWorkflowBuilder
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
DEFAULT_SEGMENT_SIZE = 10
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PendingCorpusWork:
|
|
"""Describe exact unrecorded sources and scored outputs in frozen order."""
|
|
|
|
sources: tuple[SourcePosition, ...]
|
|
visuals: tuple[VisualPosition, ...]
|
|
|
|
@property
|
|
def count(self) -> int:
|
|
"""Return the number of positions still requiring execution."""
|
|
|
|
return len(self.sources) + len(self.visuals)
|
|
|
|
|
|
class VisualCorpusExecutor:
|
|
"""Coordinate durable corpus work across bounded clean Comfy lifetimes."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
comfy_root: Path,
|
|
run_artifacts: VisualBenchmarkRunArtifacts,
|
|
manifest: BenchmarkManifest,
|
|
readiness_timeout: float,
|
|
prompt_timeout: float,
|
|
segment_size: int = DEFAULT_SEGMENT_SIZE,
|
|
) -> None:
|
|
"""Retain validated execution policy and focused collaborators."""
|
|
|
|
if segment_size < 1:
|
|
raise ValueError("P10.3 segment size must be at least one.")
|
|
self._comfy_root = comfy_root.resolve()
|
|
self._run = run_artifacts
|
|
self._manifest = manifest
|
|
self._readiness_timeout = readiness_timeout
|
|
self._prompt_timeout = prompt_timeout
|
|
self._segment_size = segment_size
|
|
self._builder = VisualBenchmarkWorkflowBuilder()
|
|
self._recorder = VisualCorpusRecorder(run_artifacts.root, manifest)
|
|
self._sources = source_positions(manifest)
|
|
self._visuals = visual_positions(manifest)
|
|
self._cases = {case.case_id: case for case in manifest.cases}
|
|
self._required = required_node_ids(
|
|
manifest,
|
|
self._builder,
|
|
self._sources,
|
|
self._visuals,
|
|
run_artifacts.run_id,
|
|
)
|
|
|
|
def execute(
|
|
self,
|
|
*,
|
|
max_sources: int | None = None,
|
|
max_visuals: int | None = None,
|
|
) -> Path:
|
|
"""Execute bounded pending segments and publish only a complete corpus."""
|
|
|
|
pending = self.pending_work(
|
|
max_sources=max_sources,
|
|
max_visuals=max_visuals,
|
|
)
|
|
source_index = 0
|
|
visual_index = 0
|
|
while source_index < len(pending.sources) or visual_index < len(
|
|
pending.visuals
|
|
):
|
|
segment_sources = pending.sources[
|
|
source_index : source_index + self._segment_size
|
|
]
|
|
remaining_capacity = self._segment_size - len(segment_sources)
|
|
segment_visuals = pending.visuals[
|
|
visual_index : visual_index + remaining_capacity
|
|
]
|
|
self._execute_segment(segment_sources, segment_visuals)
|
|
source_index += len(segment_sources)
|
|
visual_index += len(segment_visuals)
|
|
|
|
if max_sources is not None or max_visuals is not None:
|
|
return self._recorder.journal_path
|
|
corpus_path = self._recorder.finalize(
|
|
system_stats=self._run.system_stats,
|
|
cleanup_verified=True,
|
|
)
|
|
packet_path = BlindCorpusStager(self._manifest).stage(corpus_path).packet_path
|
|
self._run.record_completed(packet_path)
|
|
return packet_path
|
|
|
|
def pending_work(
|
|
self,
|
|
*,
|
|
max_sources: int | None = None,
|
|
max_visuals: int | None = None,
|
|
) -> PendingCorpusWork:
|
|
"""Select frozen ordered work while preserving exact resume semantics."""
|
|
|
|
bounded_sources = (
|
|
self._sources[:max_sources] if max_sources is not None else self._sources
|
|
)
|
|
bounded_visuals = (
|
|
self._visuals[:max_visuals] if max_visuals is not None else self._visuals
|
|
)
|
|
pending_sources = tuple(
|
|
source
|
|
for source in bounded_sources
|
|
if source.source_id not in self._recorder.completed_source_ids
|
|
)
|
|
available_sources = self._recorder.completed_source_ids | {
|
|
source.source_id for source in bounded_sources
|
|
}
|
|
pending_visuals: list[VisualPosition] = []
|
|
for visual in bounded_visuals:
|
|
if visual.artifact_id in self._recorder.completed_artifact_ids:
|
|
continue
|
|
if (
|
|
visual.source_id is not None
|
|
and visual.source_id not in available_sources
|
|
):
|
|
if max_sources is not None:
|
|
break
|
|
raise RuntimeError(
|
|
f"P10.3 required neutral source is unavailable: {visual.source_id}"
|
|
)
|
|
pending_visuals.append(visual)
|
|
return PendingCorpusWork(pending_sources, tuple(pending_visuals))
|
|
|
|
def _execute_segment(
|
|
self,
|
|
sources: tuple[SourcePosition, ...],
|
|
visuals: tuple[VisualPosition, ...],
|
|
) -> None:
|
|
"""Execute at most one bounded batch in a fresh owned Comfy process."""
|
|
|
|
segment = self._run.create_segment()
|
|
completed = 0
|
|
all_input_cleanup = True
|
|
try:
|
|
with ManagedComfyServer(
|
|
comfy_root=self._comfy_root,
|
|
artifacts=segment,
|
|
required_node_ids=self._required,
|
|
readiness_timeout=self._readiness_timeout,
|
|
) as running:
|
|
for source in sources:
|
|
LOGGER.info("P10.3 segment source: %s", source.source_id)
|
|
workflow = self._builder.build_source(
|
|
self._manifest,
|
|
source,
|
|
run_id=self._run.run_id,
|
|
)
|
|
_execute_source(
|
|
running.client,
|
|
self._recorder,
|
|
source,
|
|
workflow,
|
|
prompt_timeout=self._prompt_timeout,
|
|
)
|
|
completed += 1
|
|
for visual in visuals:
|
|
LOGGER.info("P10.3 segment visual: %s", visual.artifact_id)
|
|
cleanup = self._execute_visual(running.client, visual)
|
|
all_input_cleanup = all_input_cleanup and cleanup
|
|
completed += 1
|
|
system_stats = running.system_stats
|
|
port = running.port
|
|
process = running.process
|
|
process_running = process.is_running
|
|
port_available = is_loopback_port_available(port)
|
|
segment.record_cleanup(
|
|
process_running=process_running,
|
|
port_available=port_available,
|
|
)
|
|
if not all_input_cleanup:
|
|
raise RuntimeError("P10.3 segment input cleanup failed.")
|
|
self._run.record_segment_completed(
|
|
segment,
|
|
completed_positions=completed,
|
|
system_stats=system_stats,
|
|
)
|
|
except BaseException as error:
|
|
segment.record_failure(error)
|
|
self._run.record_segment_failed(segment, error)
|
|
raise
|
|
|
|
def _execute_visual(
|
|
self,
|
|
client: LoopbackComfyClient,
|
|
position: VisualPosition,
|
|
) -> bool:
|
|
"""Execute one visual with exact temporary input cleanup."""
|
|
|
|
source_path = (
|
|
self._recorder.source_path(position.source_id)
|
|
if position.source_id is not None
|
|
else None
|
|
)
|
|
inputs = VisualPositionInputs(
|
|
self._comfy_root / "input",
|
|
self._run.run_id,
|
|
position,
|
|
self._cases[position.case_id],
|
|
source_path=source_path,
|
|
)
|
|
with inputs:
|
|
workflow = self._builder.build_visual(
|
|
self._manifest,
|
|
position,
|
|
run_id=self._run.run_id,
|
|
mask_names=inputs.mask_names,
|
|
source_name=inputs.source_name,
|
|
)
|
|
_execute_visual(
|
|
client,
|
|
self._recorder,
|
|
position,
|
|
workflow,
|
|
mask_sha256s=inputs.mask_sha256s,
|
|
prompt_timeout=self._prompt_timeout,
|
|
)
|
|
return inputs.cleanup_verified
|
|
|
|
|
|
def _execute_source(
|
|
client: LoopbackComfyClient,
|
|
recorder: VisualCorpusRecorder,
|
|
source: SourcePosition,
|
|
workflow: BuiltVisualWorkflow,
|
|
*,
|
|
prompt_timeout: float,
|
|
) -> None:
|
|
"""Execute and record one neutral source through the shared history parser."""
|
|
|
|
prompt_id = client.submit(workflow.prompt)
|
|
history = client.wait_for_history(prompt_id, timeout=prompt_timeout)
|
|
outputs = parse_completed_outputs(
|
|
history,
|
|
metrics_node_id=workflow.metrics_node_id,
|
|
save_node_id=workflow.save_node_id,
|
|
)
|
|
recorder.record_source(
|
|
source,
|
|
workflow,
|
|
outputs,
|
|
history=history,
|
|
prompt_id=prompt_id,
|
|
image_bytes=client.download_image(outputs.image),
|
|
)
|
|
|
|
|
|
def _execute_visual(
|
|
client: LoopbackComfyClient,
|
|
recorder: VisualCorpusRecorder,
|
|
position: VisualPosition,
|
|
workflow: BuiltVisualWorkflow,
|
|
*,
|
|
mask_sha256s: tuple[str, ...],
|
|
prompt_timeout: float,
|
|
) -> None:
|
|
"""Execute and record one exact scored-output position."""
|
|
|
|
prompt_id = client.submit(workflow.prompt)
|
|
history = client.wait_for_history(prompt_id, timeout=prompt_timeout)
|
|
outputs = parse_completed_outputs(
|
|
history,
|
|
metrics_node_id=workflow.metrics_node_id,
|
|
save_node_id=workflow.save_node_id,
|
|
)
|
|
recorder.record_visual(
|
|
position,
|
|
workflow,
|
|
outputs,
|
|
history=history,
|
|
prompt_id=prompt_id,
|
|
image_bytes=client.download_image(outputs.image),
|
|
mask_sha256s=mask_sha256s,
|
|
)
|
|
|
|
|
|
def required_node_ids(
|
|
manifest: BenchmarkManifest,
|
|
builder: VisualBenchmarkWorkflowBuilder,
|
|
sources: tuple[SourcePosition, ...],
|
|
positions: tuple[VisualPosition, ...],
|
|
run_id: str,
|
|
) -> frozenset[str]:
|
|
"""Build representative graphs and return every exact required node ID."""
|
|
|
|
required = set(
|
|
builder.build_source(manifest, sources[0], run_id=run_id).required_node_ids
|
|
)
|
|
case = next(
|
|
value for value in manifest.cases if value.case_id == positions[0].case_id
|
|
)
|
|
mask_names = tuple(f"mask-{index}.png" for index in range(len(case.masks)))
|
|
for position in positions[:10]:
|
|
required.update(
|
|
builder.build_visual(
|
|
manifest,
|
|
position,
|
|
run_id=run_id,
|
|
mask_names=mask_names,
|
|
source_name="source.png" if position.is_refinement else None,
|
|
).required_node_ids
|
|
)
|
|
return frozenset(required)
|