113 lines
3.8 KiB
Python
113 lines
3.8 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
|
|
|
|
"""Coordinate one context-managed Comfy server from focused lifecycle owners."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import platform
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from tools.comfy_api import JsonObject, LoopbackComfyClient
|
|
|
|
from .artifacts import IntegrationArtifacts
|
|
from .loopback_port import is_loopback_port_available, reserve_loopback_port
|
|
from .readiness import wait_for_server
|
|
from .server_process import ComfyServerCommand, WindowsComfyProcess
|
|
|
|
MAX_LAUNCH_ATTEMPTS = 4
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RunningManagedComfy:
|
|
"""Expose ready server state without transferring lifecycle ownership."""
|
|
|
|
client: LoopbackComfyClient
|
|
process: WindowsComfyProcess
|
|
system_stats: JsonObject
|
|
port: int
|
|
|
|
|
|
class ManagedComfyServer:
|
|
"""Start, verify, expose, and always clean one exact Comfy process tree."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
comfy_root: Path,
|
|
artifacts: IntegrationArtifacts,
|
|
required_node_ids: frozenset[str],
|
|
readiness_timeout: float = 180.0,
|
|
launch_arguments: tuple[str, ...] = (),
|
|
) -> None:
|
|
"""Retain immutable launch inputs without starting external state."""
|
|
|
|
self._root = comfy_root.resolve()
|
|
self._artifacts = artifacts
|
|
self._required = required_node_ids
|
|
self._readiness_timeout = readiness_timeout
|
|
self._launch_arguments = launch_arguments
|
|
self._running: RunningManagedComfy | None = None
|
|
|
|
def __enter__(self) -> RunningManagedComfy:
|
|
"""Start and verify a server, retrying only a proven port collision."""
|
|
|
|
for attempt in range(1, MAX_LAUNCH_ATTEMPTS + 1):
|
|
reservation = reserve_loopback_port()
|
|
port = reservation.port
|
|
reservation.release()
|
|
try:
|
|
return self._start_on_port(port)
|
|
except Exception:
|
|
if attempt == MAX_LAUNCH_ATTEMPTS or is_loopback_port_available(port):
|
|
raise
|
|
raise AssertionError("Managed Comfy launch attempts were not exhausted.")
|
|
|
|
def _start_on_port(self, port: int) -> RunningManagedComfy:
|
|
"""Launch and verify one server attempt on an owned candidate port."""
|
|
|
|
command = ComfyServerCommand(
|
|
self._root,
|
|
self._root / "venv" / "Scripts" / "python.exe",
|
|
port,
|
|
self._launch_arguments,
|
|
)
|
|
process = WindowsComfyProcess.start(
|
|
command,
|
|
stdout_path=self._artifacts.stdout_path,
|
|
stderr_path=self._artifacts.stderr_path,
|
|
)
|
|
self._artifacts.record_started(
|
|
command=command.arguments(),
|
|
environment={
|
|
"comfy_root": str(command.comfy_root),
|
|
"python_executable": str(command.python_executable),
|
|
"platform": platform.platform(),
|
|
},
|
|
required_node_ids=self._required,
|
|
port=port,
|
|
pid=process.pid,
|
|
)
|
|
client = LoopbackComfyClient(f"http://127.0.0.1:{port}")
|
|
try:
|
|
stats = wait_for_server(
|
|
client,
|
|
process,
|
|
self._required,
|
|
timeout=self._readiness_timeout,
|
|
)
|
|
except BaseException:
|
|
process.stop()
|
|
raise
|
|
self._running = RunningManagedComfy(client, process, stats, port)
|
|
return self._running
|
|
|
|
def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
|
|
"""Stop only the process tree created by this context."""
|
|
|
|
del exc_type, exc, traceback
|
|
if self._running is not None:
|
|
self._running.process.stop()
|