feat(workflow-sweep): sweep runner, USCGSweepCollect node and routes
workflow_sweep_runner feeds jobs through ComfyUI's own queue (windowed feeder that
mirrors server.py post_prompt in-process via ComfyBridge), tracks them in history,
pauses on interrupt/queue-clear/3 errors and resumes by fingerprint. USCGSweepCollect
saves each job's image into the USCG session and appends manifest items. Routes
/config_tester/sweep/{parse,dry_run,start,status,pause,resume,cancel}; the Generator
node starts a sweep for workflow configs and returns the dashboard without waiting.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Jason Hoku
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
e5312812d9
commit
9ef138bf72
+77
-1
@@ -14,6 +14,7 @@ from .config_builder_node import UltimateConfigBuilder
|
||||
from .json_text_node import SmartJSONTextNode
|
||||
from .metadata_packer import pack_metadata_into_image
|
||||
from .directory_scanner import scan_directory_for_images
|
||||
from .workflow_sweep_collect import USCGSweepCollect
|
||||
import sys
|
||||
|
||||
# Register a sys.modules alias for cross-companion imports.
|
||||
@@ -1120,6 +1121,79 @@ async def cancel_upscale(request):
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# API: WORKFLOW SWEEP (sweep any API-format workflow through ComfyUI's own executor)
|
||||
# JSON bodies only, no path parameters: workflow JSON is request-body data, never a filesystem path.
|
||||
# =============================================================================
|
||||
|
||||
@server.PromptServer.instance.routes.post("/config_tester/sweep/parse")
|
||||
async def sweep_parse(request):
|
||||
"""Describe an API-format workflow for the Builder (nodes, outputs, image sources, suggested bindings)."""
|
||||
try:
|
||||
data = await request.json()
|
||||
from .workflow_sweep import parse_for_ui
|
||||
return web.json_response(parse_for_ui(data.get("api_json")))
|
||||
except ValueError as e:
|
||||
return web.json_response({"ok": False, "errors": [str(e)], "error": str(e)}, status=400)
|
||||
except Exception as e:
|
||||
print(f"[ConfigTester] Error parsing workflow: {e}")
|
||||
return web.json_response({"ok": False, "errors": [str(e)], "error": str(e)}, status=500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/config_tester/sweep/dry_run")
|
||||
async def sweep_dry_run(request):
|
||||
"""First N concrete prompts (with collector) + validation of job 1. Nothing is queued."""
|
||||
try:
|
||||
data = await request.json()
|
||||
from .workflow_sweep_runner import dry_run
|
||||
return web.json_response(await dry_run(data))
|
||||
except ValueError as e:
|
||||
return web.json_response({"error": str(e)}, status=400)
|
||||
except Exception as e:
|
||||
print(f"[ConfigTester] Error in sweep dry run: {e}")
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/config_tester/sweep/start")
|
||||
async def sweep_start(request):
|
||||
"""Start a workflow sweep from the Builder (no Generator node needed; also the resume path)."""
|
||||
try:
|
||||
data = await request.json()
|
||||
from .workflow_sweep_runner import api_start
|
||||
status, body = await api_start(data)
|
||||
return web.json_response(body, status=status)
|
||||
except Exception as e:
|
||||
print(f"[ConfigTester] Error starting sweep: {e}")
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/config_tester/sweep/status")
|
||||
async def sweep_status(request):
|
||||
from .workflow_sweep_runner import get_status
|
||||
status = get_status(request.query.get("sweep_id") or None, request.query.get("session_name") or None)
|
||||
if status is None:
|
||||
return web.json_response({"error": "No such sweep"}, status=404)
|
||||
return web.json_response(status)
|
||||
|
||||
async def _sweep_control(request, action):
|
||||
try:
|
||||
data = await request.json()
|
||||
from . import workflow_sweep_runner as runner
|
||||
ok, state = getattr(runner, action)(data.get("sweep_id", ""))
|
||||
return web.json_response({"ok": ok, "state": state}, status=200 if ok else 409)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/config_tester/sweep/pause")
|
||||
async def sweep_pause(request):
|
||||
return await _sweep_control(request, "pause")
|
||||
|
||||
@server.PromptServer.instance.routes.post("/config_tester/sweep/resume")
|
||||
async def sweep_resume(request):
|
||||
return await _sweep_control(request, "resume")
|
||||
|
||||
@server.PromptServer.instance.routes.post("/config_tester/sweep/cancel")
|
||||
async def sweep_cancel(request):
|
||||
return await _sweep_control(request, "cancel")
|
||||
|
||||
|
||||
def _create_image_links(source_dir, link_dir, items):
|
||||
"""
|
||||
Create symlinks (or copies on Windows) from source images and videos into
|
||||
@@ -1173,7 +1247,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"UltimateSamplerGrid": SamplerGridTester,
|
||||
"UltimateGridDashboard": SamplerConfigDashboardViewer,
|
||||
"UltimateConfigBuilder": UltimateConfigBuilder,
|
||||
"SmartJSONText": SmartJSONTextNode
|
||||
"SmartJSONText": SmartJSONTextNode,
|
||||
"USCGSweepCollect": USCGSweepCollect,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -1181,6 +1256,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UltimateGridDashboard": "Ultimate Grid Dashboard (Viewer)",
|
||||
"UltimateConfigBuilder": "Ultimate Config Builder",
|
||||
"SmartJSONText": "Smart JSON Text",
|
||||
"USCGSweepCollect": "USCG Sweep Collect",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
+3
-3
@@ -5,6 +5,7 @@ Handles saving and merging manifest.json with user modifications
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
|
||||
# Thread-safe lock for manifest file writes.
|
||||
@@ -123,9 +124,8 @@ def load_existing_manifest(manifest_path):
|
||||
def _item_file_path(manifest_path, item):
|
||||
"""Local path of an item's image, derived from its /view URL, confined to <session>/images."""
|
||||
try:
|
||||
from urllib.parse import urlparse, parse_qs
|
||||
q = parse_qs(urlparse(item.get("file", "")).query)
|
||||
filename = os.path.basename((q.get("filename") or [""])[0])
|
||||
m = re.search(r"[?&]filename=([^&]+)", str(item.get("file", "")))
|
||||
filename = os.path.basename(m.group(1)) if m else ""
|
||||
if not filename:
|
||||
return None
|
||||
images_dir = os.path.join(os.path.dirname(os.path.abspath(manifest_path)), "images")
|
||||
|
||||
@@ -46,6 +46,20 @@ except ImportError:
|
||||
from .html_generator import get_html_template
|
||||
|
||||
|
||||
def _workflow_configs(configs_json):
|
||||
"""The configs list when configs_json holds Workflow Sweep entries (model_type "workflow"), else None."""
|
||||
if not isinstance(configs_json, str) or "workflow" not in configs_json:
|
||||
return None
|
||||
try:
|
||||
parsed = json.loads(configs_json)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
configs = parsed.get("configs") if isinstance(parsed, dict) else parsed
|
||||
if isinstance(configs, list) and any(isinstance(c, dict) and c.get("model_type") == "workflow" for c in configs):
|
||||
return configs
|
||||
return None
|
||||
|
||||
|
||||
class SamplerGridTester:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -96,6 +110,11 @@ class SamplerGridTester:
|
||||
"optional_positive", "optional_negative", "optional_latent"
|
||||
]
|
||||
|
||||
# Workflow Sweep configs: a re-queue must restart/resume the sweep (finished jobs are skipped by
|
||||
# fingerprint), never hit ComfyUI's node cache
|
||||
if _workflow_configs(kwargs.get("configs_json")):
|
||||
return float("NaN")
|
||||
|
||||
# Check if any optional input is connected (not None and present in kwargs)
|
||||
has_optional = any(
|
||||
kwargs.get(key) is not None for key in optional_keys
|
||||
@@ -327,6 +346,29 @@ class SamplerGridTester:
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
# Workflow Sweep: hand the template + bindings to the sweep runner and return the dashboard
|
||||
# immediately. The sweep prompts go through ComfyUI's own queue AFTER this prompt finishes, so
|
||||
# this node never waits on them (a node waiting on a prompt it queued would deadlock the executor).
|
||||
wf_configs = _workflow_configs(configs_json)
|
||||
if wf_configs is not None:
|
||||
if any(x is not None for x in (optional_model, optional_clip, optional_vae,
|
||||
optional_positive, optional_negative, optional_latent)):
|
||||
print("[GridTester] ⚠️ Workflow Sweep ignores the optional_* inputs (the workflow has its own loaders)")
|
||||
from .workflow_sweep_runner import start_sweep, SweepBusy
|
||||
from .workflow_sweep import read_manifest_items, session_paths
|
||||
settings = dict(session_settings or {})
|
||||
settings["session_name"] = session_name
|
||||
settings["overwrite_existing"] = overwrite_existing
|
||||
settings["add_random_seeds_to_gens"] = add_random_seeds_to_gens
|
||||
try:
|
||||
sweep = start_sweep(wf_configs, settings, origin_node=unique_id)
|
||||
except SweepBusy as e:
|
||||
raise RuntimeError(str(e))
|
||||
manifest = read_manifest_items(session_paths(sweep.session)["manifest"])
|
||||
manifest.setdefault("session_name", sweep.session)
|
||||
manifest.setdefault("meta", {})
|
||||
return (get_html_template(sweep.session, manifest, unique_id),)
|
||||
|
||||
return run_generation_loop(
|
||||
self,
|
||||
ckpt_name, positive_text, negative_text, seed, denoise, vae_batch_size,
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
"""USCGSweepCollect node: saves images + manifest items inside the benchmarks folder, notifies the runner."""
|
||||
import functools
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
PKG = os.path.basename(ROOT)
|
||||
sys.path.insert(0, ROOT)
|
||||
|
||||
import workflow_sweep as wfs # noqa: E402
|
||||
import workflow_sweep_collect as collect_mod # noqa: E402
|
||||
import workflow_sweep_runner as runner # noqa: E402
|
||||
|
||||
|
||||
def _load_real(name, private):
|
||||
spec = importlib.util.spec_from_file_location(f"{PKG}.{private}", os.path.join(ROOT, f"{name}.py"))
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
mod.__package__ = PKG
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
class FakeTensor:
|
||||
"""Quacks like a torch IMAGE batch [B,H,W,C] float 0..1 (cpu().numpy())."""
|
||||
|
||||
def __init__(self, arr):
|
||||
self.arr = arr
|
||||
|
||||
def __len__(self):
|
||||
return len(self.arr)
|
||||
|
||||
def __getitem__(self, i):
|
||||
return FakeTensor(self.arr[i])
|
||||
|
||||
def cpu(self):
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
return self.arr
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def env(tmp_path, monkeypatch):
|
||||
import folder_paths
|
||||
monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False)
|
||||
mu = _load_real("manifest_utils", "manifest_utils_real2")
|
||||
monkeypatch.setattr(collect_mod, "_append_items", mu.append_items)
|
||||
torch = sys.modules.get("torch")
|
||||
patched = torch is not None and not hasattr(torch, "inference_mode")
|
||||
if patched:
|
||||
torch.inference_mode = lambda *a, **k: (lambda f: f)
|
||||
try:
|
||||
cim = _load_real("image_generation", "image_generation_real2").create_image_metadata
|
||||
finally:
|
||||
if patched:
|
||||
del torch.inference_mode
|
||||
real_build = wfs.build_collect_item
|
||||
monkeypatch.setattr(wfs, "build_collect_item", functools.partial(real_build, create_image_metadata=cim))
|
||||
runner.reset_registry()
|
||||
yield tmp_path
|
||||
runner.reset_registry()
|
||||
|
||||
|
||||
def images(n=1, h=16, w=24):
|
||||
return FakeTensor(np.random.rand(n, h, w, 3).astype(np.float32))
|
||||
|
||||
|
||||
CONF = {"model_type": "workflow", "model": "m", "sampler": "euler", "scheduler": "simple", "steps": 20,
|
||||
"cfg": 4.0, "seed": 3, "denoise": 1.0, "positive": "a teapot", "negative": "", "wf_fingerprint": "ab" * 10}
|
||||
|
||||
|
||||
def job_json(**extra):
|
||||
d = {"conf": CONF, "sweep_id": "sw1", "job_index": 0, "gen_index": 7, "origin_node": "12", "image_format": "webp"}
|
||||
d.update(extra)
|
||||
return json.dumps(d)
|
||||
|
||||
|
||||
def test_node_definition_matches_plan():
|
||||
cls = collect_mod.USCGSweepCollect
|
||||
spec = cls.INPUT_TYPES()
|
||||
assert spec["required"]["images"] == ("IMAGE",) and "session_name" in spec["required"]
|
||||
assert spec["required"]["job_json"][1]["multiline"] is True
|
||||
assert spec["hidden"] == {"prompt": "PROMPT", "unique_id": "UNIQUE_ID"}
|
||||
assert cls.RETURN_TYPES == () and cls.OUTPUT_NODE is True and cls.FUNCTION == "collect"
|
||||
assert cls.IS_CHANGED() != cls.IS_CHANGED() # NaN: never cached
|
||||
assert collect_mod.NODE_CLASS_MAPPINGS["USCGSweepCollect"] is cls
|
||||
|
||||
|
||||
def test_collect_saves_image_manifest_item_and_ui_output(env):
|
||||
out = collect_mod.USCGSweepCollect().collect(images(2), "mysession", job_json(), prompt={"1": {}}, unique_id="40")
|
||||
base = env / "benchmarks" / "mysession"
|
||||
manifest = json.loads((base / "manifest.json").read_text())
|
||||
assert len(manifest["items"]) == 2 and len(out["ui"]["images"]) == 2
|
||||
item = manifest["items"][0]
|
||||
assert item["model_type"] == "workflow" and item["width"] == 24 and item["height"] == 16 # from the tensor
|
||||
assert item["gen_index"] == 7 and item["seed"] == 3 and item["denoise"] == 1
|
||||
assert item["wf_sweep_id"] == "sw1" and item["wf_job_index"] == 0
|
||||
assert item["favorited"] is False and item["rejected"] is False and item["note"] == ""
|
||||
assert item["file"].startswith("/view?filename=img_") and item["file"].endswith(
|
||||
"&type=output&subfolder=benchmarks/mysession/images")
|
||||
assert {m["batch_idx"] for m in manifest["items"]} == {0, 1}
|
||||
assert len({m["id"] for m in manifest["items"]}) == 2
|
||||
for f in out["ui"]["images"]:
|
||||
assert f["subfolder"] == "benchmarks/mysession/images" and f["type"] == "output"
|
||||
assert Image.open(base / "images" / f["filename"]).size == (24, 16)
|
||||
|
||||
|
||||
def test_collect_png_embeds_prompt_chunk_and_formats(env):
|
||||
prompt = {"3": {"class_type": "KSampler", "inputs": {}}}
|
||||
out = collect_mod.USCGSweepCollect().collect(images(), "fmt", job_json(image_format="png"), prompt=prompt)
|
||||
img = Image.open(env / "benchmarks" / "fmt" / "images" / out["ui"]["images"][0]["filename"])
|
||||
assert img.format == "PNG" and json.loads(img.text["prompt"]) == prompt
|
||||
out = collect_mod.USCGSweepCollect().collect(images(), "fmt", job_json(image_format="jpg"))
|
||||
assert out["ui"]["images"][0]["filename"].endswith(".jpg")
|
||||
out = collect_mod.USCGSweepCollect().collect(images(), "fmt", job_json(image_format="weird"))
|
||||
assert out["ui"]["images"][0]["filename"].endswith(".webp")
|
||||
|
||||
|
||||
def test_collect_sanitises_session_name_and_stays_in_benchmarks(env):
|
||||
collect_mod.USCGSweepCollect().collect(images(), "../../escape/..\\x", job_json())
|
||||
assert (env / "benchmarks" / "escapex" / "manifest.json").exists()
|
||||
assert not (env.parent / "escape").exists()
|
||||
written = [p for p in env.rglob("*") if p.is_file()]
|
||||
assert all(str(p).startswith(str(env / "benchmarks")) for p in written)
|
||||
|
||||
|
||||
def test_collect_works_without_a_sweep_manual_use(env):
|
||||
out = collect_mod.USCGSweepCollect().collect(images(), "manual", "{}")
|
||||
item = json.loads((env / "benchmarks" / "manual" / "manifest.json").read_text())["items"][0]
|
||||
assert item["model_type"] == "workflow" and "wf_sweep_id" not in item and len(out["ui"]["images"]) == 1
|
||||
collect_mod.USCGSweepCollect().collect(images(), "manual", "not json at all") # garbage job_json is tolerated
|
||||
|
||||
|
||||
def test_collect_appends_to_fresh_disk_manifest_preserving_favorites(env):
|
||||
d = env / "benchmarks" / "fav"
|
||||
d.mkdir(parents=True)
|
||||
(d / "manifest.json").write_text(json.dumps({"session_name": "fav", "meta": {"k": 1},
|
||||
"items": [{"id": 1, "favorited": True}]}))
|
||||
collect_mod.USCGSweepCollect().collect(images(), "fav", job_json())
|
||||
m = json.loads((d / "manifest.json").read_text())
|
||||
assert [i.get("id") for i in m["items"]][-1] == 1 and m["items"][-1]["favorited"] is True
|
||||
assert m["meta"] == {"k": 1} and len(m["items"]) == 2
|
||||
|
||||
|
||||
def test_collect_replace_item_ids_removes_old_item_and_file(env):
|
||||
d = env / "benchmarks" / "rep"
|
||||
(d / "images").mkdir(parents=True)
|
||||
(d / "images" / "img_old.webp").write_bytes(b"x")
|
||||
(d / "manifest.json").write_text(json.dumps({"items": [
|
||||
{"id": 99, "file": "/view?filename=img_old.webp&type=output&subfolder=benchmarks/rep/images"}]}))
|
||||
collect_mod.USCGSweepCollect().collect(images(), "rep", job_json(replace_item_ids=[99]))
|
||||
m = json.loads((d / "manifest.json").read_text())
|
||||
assert [i["id"] for i in m["items"]] != [99] and all(i["id"] != 99 for i in m["items"]) and len(m["items"]) == 1
|
||||
assert not (d / "images" / "img_old.webp").exists()
|
||||
|
||||
|
||||
def test_collect_notifies_runner_and_dashboard(env, monkeypatch):
|
||||
events = []
|
||||
import server
|
||||
ps = type("PS", (), {"send_sync": lambda self, e, d: events.append((e, d))})()
|
||||
monkeypatch.setattr(server.PromptServer, "instance", ps, raising=False)
|
||||
called = []
|
||||
monkeypatch.setattr(runner, "on_collected", lambda sid, idx, items: called.append((sid, idx, items)) or True)
|
||||
collect_mod.USCGSweepCollect().collect(images(), "notify", job_json())
|
||||
assert called[0][0] == "sw1" and called[0][1] == 0 and len(called[0][2]) == 1
|
||||
(event, data), = events
|
||||
assert event == "ultimate_grid.update" and data["node"] == "12" and data["session_name"] == "notify"
|
||||
assert len(data["new_items"]) == 1 and isinstance(data["meta"], dict)
|
||||
|
||||
|
||||
def test_collect_duration_uses_runner_start_time(env, monkeypatch):
|
||||
monkeypatch.setattr(runner, "job_started_at", lambda sid, idx: time.time() - 10)
|
||||
collect_mod.USCGSweepCollect().collect(images(), "dur", job_json())
|
||||
item = json.loads((env / "benchmarks" / "dur" / "manifest.json").read_text())["items"][0]
|
||||
assert 9.5 < item["duration"] < 12
|
||||
|
||||
|
||||
def test_package_registers_collector_node_and_sweep_routes():
|
||||
pkg = sys.modules.get(PKG)
|
||||
mappings = getattr(pkg, "NODE_CLASS_MAPPINGS", None)
|
||||
assert mappings is not None and "USCGSweepCollect" in mappings
|
||||
assert pkg.NODE_DISPLAY_NAME_MAPPINGS["USCGSweepCollect"] == "USCG Sweep Collect"
|
||||
for fn in ("sweep_parse", "sweep_dry_run", "sweep_start", "sweep_status", "sweep_pause", "sweep_resume", "sweep_cancel"):
|
||||
assert callable(getattr(pkg, fn)), fn
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Generator node in workflow mode: IS_CHANGED, hand-off to the sweep runner, never blocks."""
|
||||
import importlib
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
PKG = os.path.basename(ROOT)
|
||||
sys.path.insert(0, ROOT)
|
||||
|
||||
# the Generator imports the runner as <package>.workflow_sweep_runner; use that same module instance here
|
||||
runner = importlib.import_module(f"{PKG}.workflow_sweep_runner")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def sampler_node():
|
||||
lora_utils = sys.modules[f"{PKG}.lora_utils"]
|
||||
if not hasattr(lora_utils, "load_and_save_tags"):
|
||||
lora_utils.load_and_save_tags = lambda *a, **k: None
|
||||
sys.modules["comfy"].utils = sys.modules["comfy.utils"]
|
||||
sys.modules["comfy"].sd = sys.modules["comfy.sd"]
|
||||
# root conftest registers a stub for sampler_node; load the real file under a private name
|
||||
spec = importlib.util.spec_from_file_location(f"{PKG}.sampler_node_real", os.path.join(ROOT, "sampler_node.py"))
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
mod.__package__ = PKG
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
def wf_config():
|
||||
with open(os.path.join(ROOT, "tests", "fixtures", "image_moss_image_t2i.json"), encoding="utf-8") as f:
|
||||
wf = json.load(f)
|
||||
return {"model_type": "workflow", "name": "c", "workflow": wf,
|
||||
"workflow_output": {"node_id": "38", "output_index": 0},
|
||||
"workflow_bindings": [{"id": "b", "kind": "literal", "targets": [{"node_id": "32", "input": "seed"}],
|
||||
"values": [1, 2]}]}
|
||||
|
||||
|
||||
def test_workflow_configs_detection(sampler_node):
|
||||
detect = sampler_node._workflow_configs
|
||||
assert detect(json.dumps({"configs": [wf_config()]})) and detect(json.dumps([wf_config()]))
|
||||
assert detect(json.dumps({"configs": [{"model_type": "checkpoint", "model": "workflow.safetensors"}]})) is None
|
||||
assert detect("not json workflow") is None and detect(None) is None and detect("{}") is None
|
||||
|
||||
|
||||
def test_is_changed_is_nan_for_workflow_and_stable_for_standard(sampler_node):
|
||||
cls = sampler_node.SamplerGridTester
|
||||
wf = json.dumps({"configs": [wf_config()]})
|
||||
assert cls.IS_CHANGED(configs_json=wf) != cls.IS_CHANGED(configs_json=wf) # NaN
|
||||
std = json.dumps({"configs": [{"model_type": "checkpoint", "model": "a"}]})
|
||||
assert cls.IS_CHANGED(configs_json=std) == cls.IS_CHANGED(configs_json=std)
|
||||
|
||||
|
||||
class _FakeBridge:
|
||||
def __init__(self):
|
||||
self.enqueued = []
|
||||
|
||||
def run_on_loop(self, coro):
|
||||
import asyncio
|
||||
return asyncio.run(coro)
|
||||
|
||||
async def validate(self, prompt, partial):
|
||||
return True, None, {}
|
||||
|
||||
async def enqueue(self, prompt, partial, client_id):
|
||||
self.enqueued.append(prompt)
|
||||
return f"p{len(self.enqueued)}", None, {}
|
||||
|
||||
def queue_state(self):
|
||||
return set(), {f"p{i + 1}" for i in range(len(self.enqueued))}
|
||||
|
||||
def history_entry(self, pid):
|
||||
return None
|
||||
|
||||
def delete_queued(self, pid):
|
||||
return True
|
||||
|
||||
def interrupt(self, pid):
|
||||
return False
|
||||
|
||||
def send_sync(self, event, data):
|
||||
pass
|
||||
|
||||
|
||||
def test_run_tests_starts_sweep_and_returns_dashboard_without_waiting(sampler_node, tmp_path, monkeypatch):
|
||||
import folder_paths
|
||||
monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False)
|
||||
monkeypatch.setattr(folder_paths, "get_filename_list", lambda t: ["a.safetensors"], raising=False)
|
||||
runner.reset_registry()
|
||||
bridge = _FakeBridge()
|
||||
runner.set_bridge(bridge)
|
||||
captured = {}
|
||||
monkeypatch.setattr(sampler_node, "get_html_template",
|
||||
lambda title, manifest, node_id: captured.update(title=title, manifest=manifest, node=node_id) or "<html>")
|
||||
monkeypatch.setattr(runner, "TICK", 0.05)
|
||||
# class_info=None default arg of start_sweep would hit ComfyUI's registry; the stub `nodes` has none -> None
|
||||
cfg = json.dumps({"configs": [wf_config()], "_session_settings": {"session_name": "gen_wf", "start_at_job": 0}})
|
||||
try:
|
||||
out = sampler_node.SamplerGridTester().run_tests(cfg, unique_id="77")
|
||||
assert out == ("<html>",) and captured["title"] == "gen_wf" and captured["node"] == "77"
|
||||
assert captured["manifest"]["session_name"] == "gen_wf" and captured["manifest"]["items"] == []
|
||||
sw = runner._active_sweep()
|
||||
assert sw is not None and sw.origin_node == "77" and len(sw.jobs) == 2 # returned while jobs are pending
|
||||
# a second Generator run while the sweep is active is refused with a readable error
|
||||
with pytest.raises(RuntimeError, match="already active"):
|
||||
sampler_node.SamplerGridTester().run_tests(cfg, unique_id="77")
|
||||
finally:
|
||||
sw = runner._active_sweep()
|
||||
if sw is not None:
|
||||
runner.cancel(sw.id)
|
||||
sw.thread.join(5)
|
||||
runner.set_bridge(None)
|
||||
runner.reset_registry()
|
||||
|
||||
|
||||
def test_mixed_workflow_and_standard_configs_raise(sampler_node, tmp_path, monkeypatch):
|
||||
import folder_paths
|
||||
monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False)
|
||||
monkeypatch.setattr(folder_paths, "get_filename_list", lambda t: ["a.safetensors"], raising=False)
|
||||
runner.reset_registry()
|
||||
runner.set_bridge(_FakeBridge())
|
||||
cfg = json.dumps({"configs": [wf_config(), {"model_type": "checkpoint", "model": "x"}]})
|
||||
try:
|
||||
with pytest.raises(ValueError, match="cannot be mixed"):
|
||||
sampler_node.SamplerGridTester().run_tests(cfg, unique_id="1")
|
||||
finally:
|
||||
runner.set_bridge(None)
|
||||
runner.reset_registry()
|
||||
@@ -0,0 +1,628 @@
|
||||
"""Workflow Sweep runner: window, results, errors, breaker, interrupt/lost, cancel, resume, bridge mirror."""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, ROOT)
|
||||
|
||||
import workflow_sweep as wfs # noqa: E402
|
||||
import workflow_sweep_runner as runner # noqa: E402
|
||||
|
||||
|
||||
def load(name):
|
||||
with open(os.path.join(ROOT, "tests", "fixtures", name), "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def entry(values=(1, 2, 3, 4, 5, 6), **extra):
|
||||
e = {"model_type": "workflow", "name": "c", "workflow": load("image_moss_image_t2i.json"),
|
||||
"workflow_output": {"node_id": "38", "output_index": 0},
|
||||
"workflow_bindings": [{"id": "b", "kind": "literal", "targets": [{"node_id": "32", "input": "seed"}],
|
||||
"values": list(values)}]}
|
||||
e.update(extra)
|
||||
return e
|
||||
|
||||
|
||||
class FakeBridge:
|
||||
"""Synchronous stand-in for ComfyBridge with a hand-driven queue/history."""
|
||||
|
||||
def __init__(self):
|
||||
self.queue, self.running, self.history, self.meta = [], [], {}, {}
|
||||
self.enqueued, self.events, self.deleted, self.interrupted = [], [], [], []
|
||||
self.reject = None # callable(prompt) -> error string | None
|
||||
self.n = 0
|
||||
|
||||
def run_on_loop(self, coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
async def validate(self, prompt, partial):
|
||||
err = self.reject(prompt) if self.reject else None
|
||||
if err:
|
||||
return False, {"type": "prompt_outputs_failed_validation", "message": err}, \
|
||||
{"29": {"class_type": "UNETLoader", "errors": [{"message": err, "details": "unet_name"}]}}
|
||||
return True, None, {}
|
||||
|
||||
async def enqueue(self, prompt, partial, client_id):
|
||||
ok, error, node_errors = await self.validate(prompt, partial)
|
||||
if not ok:
|
||||
return None, error, node_errors
|
||||
self.n += 1
|
||||
pid = f"pid-{self.n}"
|
||||
collector = [n for n in prompt.values() if n["class_type"] == "USCGSweepCollect"][0]
|
||||
self.meta[pid] = json.loads(collector["inputs"]["job_json"])
|
||||
self.queue.append(pid)
|
||||
self.enqueued.append((pid, prompt, partial, client_id))
|
||||
return pid, None, {}
|
||||
|
||||
def queue_state(self):
|
||||
return set(self.running), set(self.queue)
|
||||
|
||||
def history_entry(self, pid):
|
||||
return self.history.get(pid)
|
||||
|
||||
def delete_queued(self, pid):
|
||||
if pid in self.queue:
|
||||
self.queue.remove(pid)
|
||||
self.deleted.append(pid)
|
||||
return True
|
||||
return False
|
||||
|
||||
def interrupt(self, pid):
|
||||
self.interrupted.append(pid)
|
||||
return pid in self.running
|
||||
|
||||
def send_sync(self, event, data):
|
||||
self.events.append((event, data))
|
||||
|
||||
# ---- test controls
|
||||
def start(self, pid):
|
||||
self.queue.remove(pid)
|
||||
self.running.append(pid)
|
||||
|
||||
def finish(self, pid, kind="success", collect=True, ts=(1000, 3500)):
|
||||
if pid in self.queue:
|
||||
self.queue.remove(pid)
|
||||
if pid in self.running:
|
||||
self.running.remove(pid)
|
||||
msgs = [["execution_start", {"prompt_id": pid, "timestamp": ts[0]}]]
|
||||
if kind == "success":
|
||||
msgs.append(["execution_success", {"prompt_id": pid, "timestamp": ts[1]}])
|
||||
status = "success"
|
||||
elif kind == "error":
|
||||
msgs.append(["execution_error", {"prompt_id": pid, "node_id": "32", "node_type": "KSampler",
|
||||
"exception_message": "CUDA out of memory", "timestamp": ts[1]}])
|
||||
status = "error"
|
||||
else:
|
||||
msgs.append(["execution_interrupted", {"prompt_id": pid, "timestamp": ts[1]}])
|
||||
status = "error"
|
||||
self.history[pid] = {"status": {"status_str": status, "completed": kind == "success", "messages": msgs},
|
||||
"outputs": {}}
|
||||
if kind == "success" and collect:
|
||||
meta = self.meta[pid]
|
||||
runner.on_collected(meta["sweep_id"], meta["job_index"], [{"id": 1}])
|
||||
|
||||
def progress_events(self):
|
||||
return [d for e, d in self.events if e == "ultimate_grid.progress"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def env(tmp_path, monkeypatch):
|
||||
import folder_paths
|
||||
monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False)
|
||||
runner.reset_registry()
|
||||
bridge = FakeBridge()
|
||||
yield types.SimpleNamespace(bridge=bridge, tmp=tmp_path)
|
||||
runner.reset_registry()
|
||||
|
||||
|
||||
def make(env, configs=None, settings=None, origin="node7", thread=False):
|
||||
settings = {"session_name": "wf_test", **(settings or {})}
|
||||
sw = runner.prepare_sweep(configs or [entry()], settings, origin, None, env.bridge, class_info=None)
|
||||
runner.launch(sw, start_thread=thread)
|
||||
return sw
|
||||
|
||||
|
||||
def drive(sw, bridge, n=1):
|
||||
for _ in range(n):
|
||||
if not runner.tick(sw):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def complete_all(sw, bridge, limit=100):
|
||||
for _ in range(limit):
|
||||
for pid in list(bridge.queue):
|
||||
bridge.start(pid)
|
||||
for pid in list(bridge.running):
|
||||
bridge.finish(pid)
|
||||
if not runner.tick(sw):
|
||||
return
|
||||
raise AssertionError("sweep did not finish")
|
||||
|
||||
|
||||
def test_window_respected_and_refilled(env):
|
||||
sw = make(env)
|
||||
drive(sw, env.bridge)
|
||||
assert len(env.bridge.queue) == runner.WINDOW == 2
|
||||
drive(sw, env.bridge, 3)
|
||||
assert len(env.bridge.enqueued) == 2 # nothing new while two are in flight
|
||||
pid = env.bridge.queue[0]
|
||||
env.bridge.start(pid)
|
||||
env.bridge.finish(pid)
|
||||
drive(sw, env.bridge)
|
||||
assert len(env.bridge.enqueued) == 3 # the freed slot is refilled
|
||||
assert len(sw.in_flight()) == 2
|
||||
|
||||
|
||||
def test_enqueued_prompt_shape_and_synthetic_client_id(env):
|
||||
sw = make(env)
|
||||
drive(sw, env.bridge)
|
||||
pid, prompt, partial, client_id = env.bridge.enqueued[0]
|
||||
collector_id = [k for k, v in prompt.items() if v["class_type"] == "USCGSweepCollect"][0]
|
||||
assert partial == [collector_id] # the template's SaveImage / PreviewAny are not executed
|
||||
assert client_id == f"uscg-sweep-{sw.id}"
|
||||
job_json = json.loads(prompt[collector_id]["inputs"]["job_json"])
|
||||
assert job_json["sweep_id"] == sw.id and job_json["job_index"] == 0
|
||||
assert "class_type" not in json.dumps(job_json) # job conf only, never the template
|
||||
assert prompt[collector_id]["inputs"]["session_name"] == "wf_test"
|
||||
assert prompt["32"]["inputs"]["seed"] == 1
|
||||
|
||||
|
||||
def test_run_other_outputs_executes_everything(env):
|
||||
sw = make(env, [entry(workflow_run_other_outputs=True)])
|
||||
drive(sw, env.bridge)
|
||||
assert env.bridge.enqueued[0][2] is None
|
||||
|
||||
|
||||
def test_success_path_progress_events_and_completion(env):
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
complete_all(sw, env.bridge)
|
||||
assert sw.state == "completed"
|
||||
assert [j["status"] for j in sw.jobs] == ["done"] * 3
|
||||
assert sw.jobs[0]["duration"] == pytest.approx(2.5) # from ComfyUI's execution_start/success timestamps
|
||||
events = env.bridge.progress_events()
|
||||
assert events[0]["current_job"] == 0 and events[0]["total_jobs"] == 3 and events[0]["phase"] == "sweep"
|
||||
assert all(e["node"] == "node7" and e["session_name"] == "wf_test" for e in events)
|
||||
last = events[-1]
|
||||
assert last["complete"] is True and last["total_generated"] == 3 and last["eta_str"] == "Done"
|
||||
assert any(e["current_job"] == 2 and not e.get("complete") for e in events)
|
||||
status = runner.get_status(sw.id)
|
||||
assert status["done"] == 3 and status["errors"] == 0 and status["state"] == "completed"
|
||||
assert runner.get_status(session="wf_test")["sweep_id"] == sw.id
|
||||
|
||||
|
||||
def test_spec_file_written_without_secrets(env):
|
||||
sw = make(env, [entry((1, 2))])
|
||||
complete_all(sw, env.bridge)
|
||||
path = env.tmp / "benchmarks" / "wf_test" / "workflow_sweeps" / f"{sw.id}.json"
|
||||
spec = json.loads(path.read_text())
|
||||
assert spec["state"] == "completed" and len(spec["jobs"]) == 2
|
||||
assert spec["jobs"][0]["subs"] == [["32", "seed", 1]] and spec["jobs"][0]["prompt_id"] == "pid-1"
|
||||
assert list(spec["templates"].values())[0]["output"] == ["38", 0]
|
||||
|
||||
|
||||
def test_manifest_created_with_meta(env):
|
||||
make(env)
|
||||
data = json.loads((env.tmp / "benchmarks" / "wf_test" / "manifest.json").read_text())
|
||||
assert data["session_name"] == "wf_test" and data["items"] == [] and data["meta"]["image_format"] == "webp"
|
||||
|
||||
|
||||
def test_job_error_recorded_and_sweep_continues(env):
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
drive(sw, env.bridge)
|
||||
p1, p2 = env.bridge.queue[:2]
|
||||
env.bridge.start(p1)
|
||||
env.bridge.finish(p1, "error")
|
||||
env.bridge.start(p2)
|
||||
env.bridge.finish(p2)
|
||||
complete_all(sw, env.bridge)
|
||||
assert [j["status"] for j in sw.jobs] == ["error", "done", "done"]
|
||||
assert "CUDA out of memory" in sw.jobs[0]["error"] and "KSampler" in sw.jobs[0]["error"]
|
||||
assert sw.state == "completed" and "1 job" in sw.reason
|
||||
status = runner.get_status(sw.id)
|
||||
assert status["errors"] == 1 and status["last_errors"][0]["index"] == 0
|
||||
assert sw.consecutive_errors == 0 # reset by the later success
|
||||
|
||||
|
||||
def test_three_consecutive_errors_pause(env):
|
||||
sw = make(env, [entry(range(1, 9))])
|
||||
for _ in range(3):
|
||||
drive(sw, env.bridge)
|
||||
for pid in list(env.bridge.queue):
|
||||
env.bridge.start(pid)
|
||||
for pid in list(env.bridge.running):
|
||||
env.bridge.finish(pid, "error")
|
||||
break
|
||||
drive(sw, env.bridge)
|
||||
if sw.state == "paused":
|
||||
break
|
||||
drive(sw, env.bridge)
|
||||
assert sw.state == "paused" and "in a row failed" in sw.reason and "CUDA" in sw.reason
|
||||
|
||||
|
||||
def test_invalid_job_continues_and_first_invalid_fails_start(env):
|
||||
env.bridge.reject = lambda prompt: "Value not in list: unet_name" if prompt["32"]["inputs"]["seed"] == 2 else None
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
complete_all(sw, env.bridge)
|
||||
assert [j["status"] for j in sw.jobs] == ["done", "invalid", "done"]
|
||||
assert "node 29 (UNETLoader" in sw.jobs[1]["error"] and "Value not in list" in sw.jobs[1]["error"]
|
||||
# job 1 invalid at start -> the sweep is refused with the node errors (route answers 400)
|
||||
env.bridge.reject = lambda prompt: "model missing"
|
||||
sw2 = runner.prepare_sweep([entry((1, 2))], {"session_name": "other"}, None, None, env.bridge, class_info=None)
|
||||
runner.reset_registry()
|
||||
with pytest.raises(ValueError, match="first job is not valid.*model missing"):
|
||||
asyncio.run(runner.validate_first(sw2))
|
||||
|
||||
|
||||
def test_interrupted_pauses_and_resume_requeues(env):
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
drive(sw, env.bridge)
|
||||
p1 = env.bridge.queue[0]
|
||||
env.bridge.start(p1)
|
||||
env.bridge.finish(p1, "interrupted")
|
||||
drive(sw, env.bridge)
|
||||
assert sw.state == "paused" and "interrupted" in sw.reason
|
||||
assert sw.jobs[0]["status"] == "interrupted"
|
||||
# the second (still queued) prompt is pulled while paused
|
||||
drive(sw, env.bridge)
|
||||
assert env.bridge.deleted and sw.jobs[1]["status"] == "pending"
|
||||
ok, state = runner.resume(sw.id)
|
||||
assert ok and state == "running" and sw.jobs[0]["status"] == "pending"
|
||||
complete_all(sw, env.bridge)
|
||||
assert [j["status"] for j in sw.jobs] == ["done"] * 3 and sw.state == "completed"
|
||||
|
||||
|
||||
def test_lost_prompt_pauses(env, monkeypatch):
|
||||
monkeypatch.setattr(runner, "LOST_GRACE", -1)
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
drive(sw, env.bridge)
|
||||
env.bridge.queue.clear() # user cleared ComfyUI's queue
|
||||
drive(sw, env.bridge)
|
||||
assert sw.state == "paused" and "removed from ComfyUI's queue" in sw.reason
|
||||
assert {j["status"] for j in sw.jobs[:2]} == {"lost"}
|
||||
runner.resume(sw.id)
|
||||
assert [j["status"] for j in sw.jobs[:2]] == ["pending", "pending"]
|
||||
|
||||
|
||||
def test_success_without_collected_item_is_an_error(env, monkeypatch):
|
||||
monkeypatch.setattr(runner, "COLLECT_GRACE", -1)
|
||||
sw = make(env, [entry((1,))])
|
||||
drive(sw, env.bridge)
|
||||
pid = env.bridge.queue[0]
|
||||
env.bridge.start(pid)
|
||||
env.bridge.finish(pid, collect=False)
|
||||
drive(sw, env.bridge, 2)
|
||||
assert sw.jobs[0]["status"] == "error" and "collector" in sw.jobs[0]["error"]
|
||||
|
||||
|
||||
def test_cancel_deletes_queued_interrupts_running_and_marks_rest(env):
|
||||
sw = make(env, [entry(range(1, 6))])
|
||||
drive(sw, env.bridge)
|
||||
p1, p2 = env.bridge.queue[:2]
|
||||
env.bridge.start(p1)
|
||||
drive(sw, env.bridge)
|
||||
assert sw.jobs[0]["status"] == "running"
|
||||
ok, state = runner.cancel(sw.id)
|
||||
assert ok and state == "cancelled"
|
||||
assert drive(sw, env.bridge) is False
|
||||
assert env.bridge.deleted == [p2] and env.bridge.interrupted == [p1]
|
||||
assert {j["status"] for j in sw.jobs} == {"cancelled"}
|
||||
assert runner.cancel(sw.id)[0] is False # already finished
|
||||
assert runner.get_status(sw.id)["state"] == "cancelled"
|
||||
|
||||
|
||||
def test_pause_pulls_queued_but_lets_running_finish_then_resume(env):
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
drive(sw, env.bridge)
|
||||
p1, p2 = env.bridge.queue[:2]
|
||||
env.bridge.start(p1)
|
||||
drive(sw, env.bridge)
|
||||
assert runner.pause(sw.id) == (True, "paused")
|
||||
drive(sw, env.bridge)
|
||||
assert env.bridge.deleted == [p2] and sw.jobs[0]["status"] == "running" and sw.jobs[1]["status"] == "pending"
|
||||
env.bridge.finish(p1)
|
||||
drive(sw, env.bridge)
|
||||
assert sw.jobs[0]["status"] == "done" and sw.state == "paused" and not env.bridge.queue
|
||||
assert runner.pause(sw.id)[0] is False
|
||||
runner.resume(sw.id)
|
||||
complete_all(sw, env.bridge)
|
||||
assert sw.state == "completed"
|
||||
|
||||
|
||||
def _write_manifest(env, items):
|
||||
d = env.tmp / "benchmarks" / "wf_test"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
(d / "manifest.json").write_text(json.dumps({"session_name": "wf_test", "items": items, "meta": {"x": 1}}))
|
||||
|
||||
|
||||
def test_resume_skips_fingerprints_in_manifest(env):
|
||||
plan = wfs.plan_workflow_configs([entry((1, 2, 3))])
|
||||
done_fps = [plan["jobs"][0]["wf_fingerprint"], plan["jobs"][2]["wf_fingerprint"]]
|
||||
_write_manifest(env, [{"id": 11, "wf_fingerprint": done_fps[0], "gen_index": 4},
|
||||
{"id": 12, "wf_fingerprint": done_fps[1]}, {"id": 13}])
|
||||
sw = make(env, [entry((1, 2, 3))])
|
||||
assert [j["status"] for j in sw.jobs] == ["skipped", "pending", "skipped"]
|
||||
assert sw.skipped_existing == 2 and sw.to_run == 1
|
||||
assert sw.jobs[1]["gen_index"] == 5 + 1 # continues after the highest existing gen_index
|
||||
complete_all(sw, env.bridge)
|
||||
assert len(env.bridge.enqueued) == 1 and env.bridge.enqueued[0][1]["32"]["inputs"]["seed"] == 2
|
||||
first_event = env.bridge.progress_events()[0]
|
||||
assert first_event["current_job"] == 2 and first_event["total_jobs"] == 3 # "2/3 already done"
|
||||
|
||||
|
||||
def test_overwrite_existing_reruns_and_carries_replace_ids(env):
|
||||
plan = wfs.plan_workflow_configs([entry((1, 2))])
|
||||
_write_manifest(env, [{"id": 11, "wf_fingerprint": plan["jobs"][0]["wf_fingerprint"]}])
|
||||
sw = make(env, [entry((1, 2))], {"overwrite_existing": True})
|
||||
assert [j["status"] for j in sw.jobs] == ["pending", "pending"] and sw.jobs[0]["replace_ids"] == [11]
|
||||
drive(sw, env.bridge)
|
||||
jj = env.bridge.meta[env.bridge.enqueued[0][0]]
|
||||
assert jj["replace_item_ids"] == [11]
|
||||
|
||||
|
||||
def test_start_at_job_skips_first_n(env):
|
||||
sw = make(env, [entry((1, 2, 3, 4))], {"start_at_job": 2})
|
||||
assert [j["status"] for j in sw.jobs] == ["skipped", "skipped", "pending", "pending"]
|
||||
assert sw.skipped_existing == 0
|
||||
|
||||
|
||||
def test_all_done_sweep_completes_immediately(env):
|
||||
plan = wfs.plan_workflow_configs([entry((1,))])
|
||||
_write_manifest(env, [{"id": 11, "wf_fingerprint": plan["jobs"][0]["wf_fingerprint"]}])
|
||||
sw = make(env, [entry((1,))])
|
||||
assert drive(sw, env.bridge) is False and sw.state == "completed"
|
||||
|
||||
|
||||
def test_one_active_sweep_and_409(env):
|
||||
sw = make(env)
|
||||
with pytest.raises(runner.SweepBusy):
|
||||
runner.prepare_sweep([entry()], {"session_name": "x"}, None, None, env.bridge, class_info=None)
|
||||
status, body = asyncio.run(runner.api_start(
|
||||
{"configs_json": json.dumps({"configs": [entry()], "_session_settings": {"session_name": "x"}})},
|
||||
bridge=env.bridge, class_info=None))
|
||||
assert status == 409 and body["sweep_id"] == sw.id
|
||||
runner.cancel(sw.id)
|
||||
drive(sw, env.bridge)
|
||||
runner.prepare_sweep([entry()], {"session_name": "x"}, None, None, env.bridge, class_info=None) # free again
|
||||
|
||||
|
||||
def test_api_start_ok_and_bad_input(env):
|
||||
body = {"configs_json": json.dumps({"configs": [entry((1, 2))], "_session_settings": {"session_name": "api"}}),
|
||||
"node_id": "5"}
|
||||
status, out = asyncio.run(runner.api_start(body, bridge=env.bridge, class_info=None))
|
||||
assert status == 200 and out["total"] == 2 and out["to_run"] == 2
|
||||
runner.cancel(out["sweep_id"])
|
||||
time.sleep(runner.TICK * 3)
|
||||
status, out = asyncio.run(runner.api_start({"configs_json": json.dumps({"configs": [{"model_type": "checkpoint"}]})},
|
||||
bridge=env.bridge, class_info=None))
|
||||
assert status == 400 and "cannot be mixed" in out["error"]
|
||||
status, out = asyncio.run(runner.api_start({}, bridge=env.bridge, class_info=None))
|
||||
assert status == 400
|
||||
|
||||
|
||||
def test_start_sweep_returns_while_jobs_pending_and_never_blocks(env):
|
||||
t0 = time.time()
|
||||
sw = runner.start_sweep([entry(range(1, 50))], {"session_name": "nb"}, "n1", None, env.bridge, class_info=None)
|
||||
assert time.time() - t0 < 2 and sw.state == "running"
|
||||
time.sleep(runner.TICK * 3) # the feeder thread fills the window on its own
|
||||
assert len(env.bridge.enqueued) == runner.WINDOW
|
||||
assert runner.cancel(sw.id)[0]
|
||||
sw.thread.join(5)
|
||||
assert not sw.thread.is_alive() and sw.state == "cancelled"
|
||||
|
||||
|
||||
def test_dry_run_shape_and_no_queueing(env):
|
||||
body = {"configs_json": json.dumps({"configs": [entry((1, 2, 3))], "_session_settings": {"session_name": "dry"}}),
|
||||
"limit": 2}
|
||||
out = asyncio.run(runner.dry_run(body, bridge=env.bridge, class_info=None))
|
||||
assert out["total"] == 3 and out["to_run"] == 3 and len(out["jobs"]) == 2
|
||||
assert out["validation"]["ok"] is True and out["axes"][0]["count"] == 3
|
||||
j = out["jobs"][0]
|
||||
assert j["prompt"]["32"]["inputs"]["seed"] == 1 and j["partial_execution_targets"]
|
||||
assert any(n["class_type"] == "USCGSweepCollect" for n in j["prompt"].values())
|
||||
assert env.bridge.enqueued == [] and runner._active_sweep() is None
|
||||
env.bridge.reject = lambda p: "bad model"
|
||||
out = asyncio.run(runner.dry_run(body, bridge=env.bridge, class_info=None))
|
||||
assert out["validation"]["ok"] is False and "bad model" in out["validation"]["error"]
|
||||
|
||||
|
||||
def test_dry_run_allowed_while_a_sweep_is_active(env):
|
||||
make(env)
|
||||
body = {"configs_json": json.dumps({"configs": [entry((1,))], "_session_settings": {"session_name": "d2"}})}
|
||||
assert asyncio.run(runner.dry_run(body, bridge=env.bridge, class_info=None))["total"] == 1
|
||||
|
||||
|
||||
def test_configs_from_body_uses_builder_state_transformer():
|
||||
state = {"session_name": "st", "config_arrays": [{
|
||||
"name": "WF", "kind": "workflow",
|
||||
"workflow": {"api_prompt": load("sdxl_default_api.json"), "output": {"node_id": "8", "output_index": 0},
|
||||
"bindings": []}}]}
|
||||
configs, settings = runner.configs_from_body({"state": state})
|
||||
assert configs[0]["model_type"] == "workflow" and settings["session_name"] == "st"
|
||||
with pytest.raises(ValueError):
|
||||
runner.configs_from_body({})
|
||||
|
||||
|
||||
def test_classify_history_entries():
|
||||
ok = {"status": {"status_str": "success", "messages": [["execution_start", {"timestamp": 1}],
|
||||
["execution_success", {"timestamp": 3}]]}}
|
||||
assert runner._classify(ok)[0] == "success"
|
||||
err = {"status": {"status_str": "error", "messages": [["execution_error", {"node_id": "3", "node_type": "K",
|
||||
"exception_message": "boom"}]]}}
|
||||
kind, detail = runner._classify(err)
|
||||
assert kind == "error" and "boom" in detail["error"] and "K #3" in detail["error"]
|
||||
assert runner._classify({"status": {"status_str": "error", "messages": [["execution_interrupted", {}]]}})[0] == "interrupted"
|
||||
assert runner._classify({"status": None})[0] == "unknown"
|
||||
assert runner._classify(None)[0] == "unknown"
|
||||
|
||||
|
||||
def test_summarize_errors():
|
||||
msg = runner.summarize_errors({"message": "Prompt outputs failed validation", "details": "x"},
|
||||
{"29": {"class_type": "UNETLoader", "errors": [{"message": "Value not in list",
|
||||
"details": "unet_name: 'a'"}]}},
|
||||
load("image_moss_image_t2i.json"))
|
||||
assert "Prompt outputs failed validation" in msg and "node 29 (UNETLoader / Load Diffusion Model)" in msg
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ ComfyBridge mirrors server.py post_prompt
|
||||
|
||||
class _Queue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
def _fake_comfy(monkeypatch, with_sensitive=True, validate_args=3, valid=True, handlers=None, replace=True):
|
||||
calls = {"validate": [], "replace": [], "handlers": []}
|
||||
|
||||
class PS:
|
||||
number = 5.0
|
||||
prompt_queue = _Queue()
|
||||
loop = None
|
||||
|
||||
def trigger_on_prompt(self, json_data):
|
||||
calls["handlers"].append(dict(json_data))
|
||||
if handlers:
|
||||
json_data = handlers(json_data)
|
||||
return json_data
|
||||
|
||||
class node_replace_manager: # noqa: N801
|
||||
@staticmethod
|
||||
def apply_replacements(prompt):
|
||||
calls["replace"].append(prompt)
|
||||
prompt["9999"] = {"class_type": "Replaced", "inputs": {}}
|
||||
|
||||
ps = PS()
|
||||
if not replace:
|
||||
del PS.node_replace_manager
|
||||
server_mod = sys.modules["server"]
|
||||
monkeypatch.setattr(server_mod, "PromptServer", types.SimpleNamespace(instance=ps), raising=False)
|
||||
monkeypatch.setattr(server_mod, "valid_workflow_metadata", lambda jd: jd.get("workflow_metadata"), raising=False)
|
||||
monkeypatch.setattr(server_mod, "workflow_metadata_from_prompt", lambda ed: None, raising=False)
|
||||
|
||||
ex = types.ModuleType("execution")
|
||||
if validate_args == 3:
|
||||
async def validate_prompt(prompt_id, prompt, partial_execution_list):
|
||||
calls["validate"].append((prompt_id, partial_execution_list))
|
||||
return (valid, None if valid else {"message": "bad"}, ["38x"], {} if valid else {"3": {"errors": []}})
|
||||
elif validate_args == 2:
|
||||
def validate_prompt(prompt_id, prompt):
|
||||
calls["validate"].append((prompt_id, None))
|
||||
return (valid, None, ["38x"], {})
|
||||
else:
|
||||
def validate_prompt(prompt):
|
||||
calls["validate"].append((None, None))
|
||||
return (valid, None, ["38x"], {})
|
||||
ex.validate_prompt = validate_prompt
|
||||
if with_sensitive:
|
||||
ex.SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org")
|
||||
monkeypatch.setitem(sys.modules, "execution", ex)
|
||||
return ps, calls
|
||||
|
||||
|
||||
def test_bridge_enqueue_mirrors_post_prompt(monkeypatch):
|
||||
ps, calls = _fake_comfy(monkeypatch)
|
||||
prompt = {"1": {"class_type": "X", "inputs": {}}}
|
||||
pid, err, node_errors = asyncio.run(runner.ComfyBridge().enqueue(prompt, ["1"], "cid"))
|
||||
assert err is None and node_errors == {}
|
||||
(item,) = ps.prompt_queue.items
|
||||
assert len(item) == 6 # (number, prompt_id, prompt, extra_data, outputs, sensitive)
|
||||
number, prompt_id, queued_prompt, extra, outputs, sensitive = item
|
||||
assert number == 5.0 and ps.number == 6.0 # number taken, then incremented
|
||||
assert prompt_id == pid and len(pid) == 36
|
||||
assert queued_prompt is prompt and "9999" in prompt # node replacements applied before validation
|
||||
assert extra["client_id"] == "cid" and isinstance(extra["create_time"], int)
|
||||
assert outputs == ["38x"] and sensitive == {}
|
||||
assert calls["validate"] == [(pid, ["1"])]
|
||||
assert calls["handlers"][0]["prompt"] is prompt and calls["handlers"][0]["partial_execution_targets"] == ["1"]
|
||||
|
||||
|
||||
def test_bridge_enqueue_honours_on_prompt_handler_and_strips_sensitive(monkeypatch):
|
||||
def handler(jd):
|
||||
jd["extra_data"] = {"api_key_comfy_org": "SECRET", "other": 1}
|
||||
jd["number"] = 2
|
||||
return jd
|
||||
ps, _ = _fake_comfy(monkeypatch, handlers=handler)
|
||||
asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c"))
|
||||
number, _, _, extra, _, sensitive = ps.prompt_queue.items[0]
|
||||
assert number == 2.0 and ps.number == 5.0 # an explicit number does not consume the counter
|
||||
assert sensitive == {"api_key_comfy_org": "SECRET"} and "api_key_comfy_org" not in extra and extra["other"] == 1
|
||||
|
||||
|
||||
def test_bridge_enqueue_older_server_shapes(monkeypatch):
|
||||
ps, calls = _fake_comfy(monkeypatch, with_sensitive=False, validate_args=2, replace=False)
|
||||
pid, err, _ = asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c"))
|
||||
assert err is None and len(ps.prompt_queue.items[0]) == 5 and calls["validate"][0][0] == pid
|
||||
ps, calls = _fake_comfy(monkeypatch, validate_args=1)
|
||||
asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c"))
|
||||
assert len(ps.prompt_queue.items[0]) == 6
|
||||
|
||||
|
||||
def test_bridge_enqueue_invalid_prompt_is_not_queued(monkeypatch):
|
||||
ps, _ = _fake_comfy(monkeypatch, valid=False)
|
||||
pid, err, node_errors = asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c"))
|
||||
assert pid is None and err == {"message": "bad"} and "3" in node_errors and ps.prompt_queue.items == []
|
||||
|
||||
|
||||
def test_bridge_validate_does_not_queue(monkeypatch):
|
||||
ps, calls = _fake_comfy(monkeypatch)
|
||||
ok, err, ne = asyncio.run(runner.ComfyBridge().validate({"1": {"class_type": "X", "inputs": {}}}, ["1"]))
|
||||
assert ok and ps.prompt_queue.items == [] and len(calls["validate"]) == 1
|
||||
|
||||
|
||||
def test_run_on_loop_refuses_the_loop_thread_and_works_from_others(monkeypatch):
|
||||
loop = asyncio.new_event_loop()
|
||||
t = threading.Thread(target=loop.run_forever, daemon=True)
|
||||
t.start()
|
||||
ps, _ = _fake_comfy(monkeypatch)
|
||||
ps.loop = loop
|
||||
bridge = runner.ComfyBridge()
|
||||
|
||||
async def double(x):
|
||||
return x * 2
|
||||
assert bridge.run_on_loop(double(21)) == 42
|
||||
|
||||
async def from_loop():
|
||||
c = double(1)
|
||||
with pytest.raises(RuntimeError, match="event loop thread"):
|
||||
bridge.run_on_loop(c)
|
||||
return "refused"
|
||||
assert asyncio.run_coroutine_threadsafe(from_loop(), loop).result(5) == "refused"
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
t.join(5)
|
||||
|
||||
|
||||
def test_bridge_queue_helpers(monkeypatch):
|
||||
removed = []
|
||||
|
||||
class Q:
|
||||
def get_current_queue_volatile(self):
|
||||
return [(0, "r1", {}, {}, [])], [(1, "q1", {}, {}, [], {}), (2, "q2", {}, {}, [], {})]
|
||||
|
||||
def get_history(self, pid, map_function=None):
|
||||
return {pid: map_function({"status": {"status_str": "success"}, "outputs": {}, "prompt": "big"})} \
|
||||
if pid == "h1" else {}
|
||||
|
||||
def delete_queue_item(self, fn):
|
||||
removed.append(fn((1, "q1")))
|
||||
return True
|
||||
|
||||
def interrupt_if_running(self, pid):
|
||||
return pid == "r1"
|
||||
ps, _ = _fake_comfy(monkeypatch)
|
||||
ps.prompt_queue = Q()
|
||||
b = runner.ComfyBridge()
|
||||
assert b.queue_state() == ({"r1"}, {"q1", "q2"})
|
||||
assert b.history_entry("h1") == {"status": {"status_str": "success"}, "outputs": {}} # no deep copy of prompt
|
||||
assert b.history_entry("nope") is None
|
||||
assert b.delete_queued("q1") is True and removed == [True]
|
||||
assert b.interrupt("r1") is True and b.interrupt("q1") is False
|
||||
@@ -950,3 +950,20 @@ def build_collect_item(conf, width, height, duration, batch_idx, gen_index, crea
|
||||
|
||||
def new_item_id():
|
||||
return int(time.time() * 100000) + random.randint(0, 1000)
|
||||
|
||||
|
||||
def parse_for_ui(api_json, class_info=comfy_class_info):
|
||||
"""Response body of POST /config_tester/sweep/parse. Raises ValueError for unusable input."""
|
||||
prompt = parse_api_workflow(api_json)
|
||||
sources = find_image_sources(prompt, class_info)
|
||||
default = sources[0] if sources else None
|
||||
suggestions = suggest_bindings(prompt, default, class_info) if default else []
|
||||
warnings = []
|
||||
nodes_ = describe_nodes(prompt, class_info)
|
||||
missing = sorted({n["class_type"] for n in nodes_ if n["missing_class"]})
|
||||
if missing:
|
||||
warnings.append("These node types are not installed in this ComfyUI: " + ", ".join(missing))
|
||||
if not sources:
|
||||
warnings.append("No IMAGE output was found; pick the image socket to collect manually.")
|
||||
return {"ok": True, "nodes": nodes_, "outputs": find_output_nodes(prompt, class_info),
|
||||
"image_sources": sources, "suggestions": suggestions, "warnings": warnings, "errors": []}
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
"""
|
||||
USCGSweepCollect: the output node a Workflow Sweep splices onto each job's image socket.
|
||||
|
||||
It saves the image(s) into the USCG session folder (benchmarks/<session>/images) and appends manifest
|
||||
items with the standard fields, so the Dashboard, favourites, filters and export keep working. It never
|
||||
waits on anything and never writes outside the benchmarks folder (the session name is sanitised with
|
||||
the same rule as every other route). `job_json` carries the job's conf only, never the workflow template.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
try:
|
||||
from . import workflow_sweep as wfs
|
||||
except ImportError: # top-level import (tests)
|
||||
import workflow_sweep as wfs
|
||||
|
||||
|
||||
def _append_items(*args, **kwargs):
|
||||
try:
|
||||
from .manifest_utils import append_items
|
||||
except ImportError:
|
||||
from manifest_utils import append_items
|
||||
return append_items(*args, **kwargs)
|
||||
|
||||
|
||||
def _runner():
|
||||
try:
|
||||
from . import workflow_sweep_runner as runner
|
||||
except ImportError:
|
||||
import workflow_sweep_runner as runner
|
||||
return runner
|
||||
|
||||
|
||||
def _tensor_to_array(images, b):
|
||||
"""images[b] -> uint8 HxWxC numpy array (SaveImage's conversion), duck-typed so no torch import is needed."""
|
||||
arr = images[b]
|
||||
if hasattr(arr, "cpu"):
|
||||
arr = arr.cpu()
|
||||
if hasattr(arr, "numpy"):
|
||||
arr = arr.numpy()
|
||||
return np.clip(255.0 * np.asarray(arr), 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
class USCGSweepCollect:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"session_name": ("STRING", {"default": "workflow_sweep"}),
|
||||
"job_json": ("STRING", {"multiline": True, "default": "{}"}),
|
||||
},
|
||||
"hidden": {"prompt": "PROMPT", "unique_id": "UNIQUE_ID"},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "collect"
|
||||
CATEGORY = "sampling/testing"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
# every job must run its collector, whatever the cache says about its ancestors
|
||||
return float("NaN")
|
||||
|
||||
def collect(self, images, session_name="workflow_sweep", job_json="{}", prompt=None, unique_id=None):
|
||||
paths = wfs.session_paths(session_name)
|
||||
session = paths["session"]
|
||||
os.makedirs(paths["images"], exist_ok=True)
|
||||
|
||||
try:
|
||||
job = json.loads(job_json or "{}")
|
||||
job = job if isinstance(job, dict) else {}
|
||||
except (TypeError, ValueError):
|
||||
job = {}
|
||||
conf = job.get("conf") if isinstance(job.get("conf"), dict) else {}
|
||||
sweep_id, job_index = job.get("sweep_id"), job.get("job_index")
|
||||
|
||||
fmt = str(job.get("image_format") or "webp").lower().strip()
|
||||
if fmt not in ("webp", "png", "jpg", "jpeg"):
|
||||
fmt = "webp"
|
||||
ext = "jpg" if fmt in ("jpg", "jpeg") else fmt
|
||||
|
||||
started = None
|
||||
if sweep_id is not None and job_index is not None:
|
||||
started = _runner().job_started_at(sweep_id, job_index)
|
||||
duration = round(time.time() - started, 2) if started else 0
|
||||
|
||||
gen_index = job.get("gen_index")
|
||||
items, files = [], []
|
||||
for b in range(len(images)):
|
||||
arr = _tensor_to_array(images, b)
|
||||
img = Image.fromarray(arr)
|
||||
item = wfs.build_collect_item(conf, arr.shape[1], arr.shape[0], duration, b, gen_index)
|
||||
item["id"] = wfs.new_item_id()
|
||||
filename = f"img_{item['id']}.{ext}"
|
||||
filepath = os.path.join(paths["images"], filename)
|
||||
if fmt == "png":
|
||||
info = PngInfo()
|
||||
if prompt is not None: # same text chunk SaveImage writes: the image can be dragged back into ComfyUI
|
||||
info.add_text("prompt", json.dumps(prompt))
|
||||
img.save(filepath, format="PNG", pnginfo=info)
|
||||
elif ext == "jpg":
|
||||
img.save(filepath, format="JPEG", quality=95)
|
||||
else:
|
||||
img.save(filepath, quality=80)
|
||||
item["file"] = f"/view?filename={filename}&type=output&subfolder=benchmarks/{session}/images"
|
||||
if sweep_id is not None:
|
||||
item["wf_sweep_id"] = sweep_id
|
||||
if job_index is not None:
|
||||
item["wf_job_index"] = job_index
|
||||
items.append(item)
|
||||
files.append({"filename": filename, "subfolder": f"benchmarks/{session}/images", "type": "output"})
|
||||
|
||||
manifest = _append_items(paths["manifest"], items, remove_ids=job.get("replace_item_ids") or (),
|
||||
session_name=session)
|
||||
try: # live dashboard update (same event flush_batch_with_vae sends)
|
||||
from server import PromptServer
|
||||
PromptServer.instance.send_sync("ultimate_grid.update", {
|
||||
"node": job.get("origin_node"), "session_name": session, "new_items": items,
|
||||
"meta": (manifest or {}).get("meta", {})})
|
||||
except Exception:
|
||||
pass
|
||||
if sweep_id is not None:
|
||||
try:
|
||||
_runner().on_collected(sweep_id, job_index, items)
|
||||
except Exception:
|
||||
pass
|
||||
print(f"[GridTester] 🧪 Sweep collected {len(items)} image(s) into '{session}'")
|
||||
return {"ui": {"images": files}}
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"USCGSweepCollect": USCGSweepCollect}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"USCGSweepCollect": "USCG Sweep Collect"}
|
||||
@@ -0,0 +1,912 @@
|
||||
"""
|
||||
Workflow Sweep runner: feeds a sweep's jobs to ComfyUI's own queue, one prompt per job.
|
||||
|
||||
A sweep is prepared synchronously (expand, fingerprint, skip jobs already in the manifest), then a
|
||||
daemon thread keeps at most WINDOW sweep prompts queued/running and follows each through ComfyUI's
|
||||
history. Every enqueue repeats server.py's post_prompt in-process (on_prompt handlers, node
|
||||
replacements, awaited validate_prompt, number, extra_data, queue tuple), so a queued sweep prompt is
|
||||
exactly what ComfyUI would have queued from its UI.
|
||||
|
||||
Hard rules honoured here:
|
||||
* no network code (the enqueue is an in-process call on the server event loop),
|
||||
* the executor is never blocked: nothing in this module (or the collector node) waits on a prompt
|
||||
it queued; the Generator node starts a sweep and returns immediately,
|
||||
* `run_on_loop` is never called from the loop thread, and no lock is held across it.
|
||||
|
||||
All loop-bound ComfyUI access goes through `ComfyBridge`, so the runner is testable with a fake.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
|
||||
try:
|
||||
from . import workflow_sweep as wfs
|
||||
except ImportError: # imported as a top-level module (tests)
|
||||
import workflow_sweep as wfs
|
||||
|
||||
|
||||
WINDOW = 2 # max sweep prompts queued/running at once
|
||||
TICK = 0.25 # seconds between runner polls
|
||||
LOST_GRACE = 5.0 # a prompt missing from queue, running set and history this long is "lost"
|
||||
COLLECT_GRACE = 5.0 # success without a collected item for this long is an error
|
||||
BREAKER = 3 # consecutive error/invalid jobs that pause the sweep
|
||||
SPEC_EVERY = 10 # spec file rewrite cadence (terminal jobs)
|
||||
|
||||
ACTIVE_STATES = ("preparing", "running", "paused")
|
||||
TERMINAL_JOB = ("done", "error", "invalid", "skipped", "cancelled")
|
||||
|
||||
|
||||
class SweepBusy(Exception):
|
||||
"""Another workflow sweep is still active (route -> 409)."""
|
||||
def __init__(self, sweep_id):
|
||||
super().__init__(f"A workflow sweep is already active ({sweep_id}). Cancel it or wait for it to finish.")
|
||||
self.sweep_id = sweep_id
|
||||
|
||||
|
||||
def summarize_errors(error, node_errors, prompt=None):
|
||||
"""One readable line from ComfyUI's validate_prompt error + node_errors."""
|
||||
parts = []
|
||||
if isinstance(error, dict):
|
||||
msg = error.get("message") or error.get("type") or "invalid prompt"
|
||||
if error.get("details"):
|
||||
msg += f" ({error['details']})"
|
||||
parts.append(msg)
|
||||
elif error:
|
||||
parts.append(str(error))
|
||||
for nid, ne in (node_errors or {}).items():
|
||||
ne = ne if isinstance(ne, dict) else {}
|
||||
ct = ne.get("class_type") or ((prompt or {}).get(str(nid)) or {}).get("class_type") or "?"
|
||||
title = wfs.node_title(prompt, nid) if prompt else ct
|
||||
label = ct if title == ct else f"{ct} / {title}"
|
||||
msgs = "; ".join(f"{e.get('message', '')} {e.get('details', '')}".strip() for e in (ne.get("errors") or []))
|
||||
parts.append(f"node {nid} ({label}): {msgs or 'invalid'}")
|
||||
return (" | ".join(parts) or "invalid prompt")[:1500]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- bridge to ComfyUI
|
||||
|
||||
class ComfyBridge:
|
||||
"""Thin wrapper over PromptServer.instance / execution / nodes (lazy imports, version guards)."""
|
||||
|
||||
def _ps(self):
|
||||
import server
|
||||
return server.PromptServer.instance
|
||||
|
||||
def _execution(self):
|
||||
import execution
|
||||
return execution
|
||||
|
||||
def run_on_loop(self, coro, timeout=300):
|
||||
"""Run a coroutine on the server event loop from a NON-loop thread and wait for it."""
|
||||
ps = self._ps()
|
||||
loop = ps.loop
|
||||
try:
|
||||
running = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
running = None
|
||||
if running is loop:
|
||||
coro.close()
|
||||
raise RuntimeError("run_on_loop was called from the server event loop thread (would deadlock)")
|
||||
return asyncio.run_coroutine_threadsafe(coro, loop).result(timeout)
|
||||
|
||||
async def _validate(self, prompt_id, prompt, partial):
|
||||
execution = self._execution()
|
||||
fn = execution.validate_prompt
|
||||
n = len(inspect.signature(fn).parameters)
|
||||
if n >= 3:
|
||||
result = fn(prompt_id, prompt, partial)
|
||||
elif n == 2:
|
||||
result = fn(prompt_id, prompt)
|
||||
else:
|
||||
result = fn(prompt)
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
return result # (valid, error, outputs, node_errors)
|
||||
|
||||
async def validate(self, prompt, partial_targets):
|
||||
"""Validate without queueing. Returns (ok, error, node_errors)."""
|
||||
valid = await self._validate(str(uuid.uuid4()), prompt, partial_targets)
|
||||
return bool(valid[0]), valid[1], valid[3] if len(valid) > 3 else {}
|
||||
|
||||
async def enqueue(self, prompt, partial_targets, client_id):
|
||||
"""Mirror of server.py post_prompt (ComfyUI 0.42.5). Returns (prompt_id|None, error, node_errors)."""
|
||||
ps = self._ps()
|
||||
execution = self._execution()
|
||||
json_data = {"prompt": prompt, "client_id": client_id}
|
||||
if partial_targets is not None:
|
||||
json_data["partial_execution_targets"] = partial_targets
|
||||
if hasattr(ps, "trigger_on_prompt"):
|
||||
json_data = ps.trigger_on_prompt(json_data)
|
||||
|
||||
if "number" in json_data:
|
||||
number = float(json_data["number"])
|
||||
else:
|
||||
number = ps.number
|
||||
if "front" in json_data and json_data["front"]:
|
||||
number = -number
|
||||
ps.number += 1
|
||||
|
||||
if "prompt" not in json_data:
|
||||
return None, {"type": "no_prompt", "message": "No prompt provided"}, {}
|
||||
prompt = json_data["prompt"]
|
||||
prompt_id = str(uuid.uuid4())
|
||||
partial = json_data.get("partial_execution_targets")
|
||||
|
||||
manager = getattr(ps, "node_replace_manager", None)
|
||||
if manager is not None and hasattr(manager, "apply_replacements"):
|
||||
manager.apply_replacements(prompt)
|
||||
|
||||
valid = await self._validate(prompt_id, prompt, partial)
|
||||
extra_data = json_data.get("extra_data", {})
|
||||
if not isinstance(extra_data, dict):
|
||||
extra_data = {}
|
||||
if "client_id" in json_data:
|
||||
extra_data["client_id"] = json_data["client_id"]
|
||||
|
||||
# workflow metadata for websocket messages (0.42.x); absent on older servers
|
||||
extra_data.pop("workflow_metadata", None)
|
||||
try:
|
||||
import server as _server
|
||||
valid_meta = getattr(_server, "valid_workflow_metadata", None)
|
||||
meta_from_prompt = getattr(_server, "workflow_metadata_from_prompt", None)
|
||||
metadata = valid_meta(json_data) if valid_meta else None
|
||||
if metadata is None and meta_from_prompt:
|
||||
metadata = meta_from_prompt(extra_data)
|
||||
if metadata is not None:
|
||||
extra_data["workflow_metadata"] = metadata
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not valid[0]:
|
||||
return None, valid[1], (valid[3] if len(valid) > 3 else {})
|
||||
|
||||
outputs_to_execute = valid[2]
|
||||
sensitive_keys = getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", None)
|
||||
extra_data["create_time"] = int(time.time() * 1000)
|
||||
if sensitive_keys is not None: # 6-tuple queue items (0.42.x)
|
||||
sensitive = {}
|
||||
for key in sensitive_keys:
|
||||
if key in extra_data:
|
||||
sensitive[key] = extra_data.pop(key)
|
||||
ps.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive))
|
||||
else:
|
||||
ps.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
|
||||
return prompt_id, None, (valid[3] if len(valid) > 3 else {})
|
||||
|
||||
def queue_state(self):
|
||||
"""(running prompt ids, queued prompt ids), one atomic snapshot."""
|
||||
running, queued = self._ps().prompt_queue.get_current_queue_volatile()
|
||||
return {item[1] for item in running}, {item[1] for item in queued}
|
||||
|
||||
def history_entry(self, prompt_id):
|
||||
"""{status:{status_str, messages}, outputs} or None when ComfyUI has no history for it."""
|
||||
q = self._ps().prompt_queue
|
||||
|
||||
def slim(p):
|
||||
return {"status": p.get("status"), "outputs": p.get("outputs")}
|
||||
try:
|
||||
got = q.get_history(prompt_id, map_function=slim)
|
||||
except TypeError: # older servers: no map_function
|
||||
got = q.get_history(prompt_id)
|
||||
return (got or {}).get(prompt_id)
|
||||
|
||||
def delete_queued(self, prompt_id):
|
||||
return bool(self._ps().prompt_queue.delete_queue_item(lambda item: item[1] == prompt_id))
|
||||
|
||||
def interrupt(self, prompt_id):
|
||||
q = self._ps().prompt_queue
|
||||
if hasattr(q, "interrupt_if_running"):
|
||||
return bool(q.interrupt_if_running(prompt_id))
|
||||
running, _ = self.queue_state()
|
||||
if prompt_id in running:
|
||||
import nodes
|
||||
nodes.interrupt_processing()
|
||||
return True
|
||||
return False
|
||||
|
||||
def send_sync(self, event, data):
|
||||
try:
|
||||
self._ps().send_sync(event, data)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
_bridge = None
|
||||
|
||||
|
||||
def get_bridge():
|
||||
global _bridge
|
||||
if _bridge is None:
|
||||
_bridge = ComfyBridge()
|
||||
return _bridge
|
||||
|
||||
|
||||
def set_bridge(bridge):
|
||||
"""Tests inject a FakeBridge."""
|
||||
global _bridge
|
||||
_bridge = bridge
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- sweep model
|
||||
|
||||
class Sweep:
|
||||
def __init__(self, sweep_id, session, jobs, templates, origin_node, client_id, settings, warnings, axes, bridge):
|
||||
self.id = sweep_id
|
||||
self.session = session
|
||||
self.jobs = jobs
|
||||
self.templates = templates
|
||||
self.origin_node = origin_node
|
||||
self.client_id = client_id
|
||||
self.settings = settings # {image_format, overwrite_existing, ...}
|
||||
self.warnings = warnings
|
||||
self.axes = axes
|
||||
self.bridge = bridge
|
||||
self.state = "preparing"
|
||||
self.reason = ""
|
||||
self.created_at = time.time()
|
||||
self.started_at = None
|
||||
self.finished_at = None
|
||||
self.lock = threading.RLock()
|
||||
self.wake = threading.Event()
|
||||
self.thread = None
|
||||
self.needs_cleanup = False
|
||||
self.consecutive_errors = 0
|
||||
self.durations = []
|
||||
self.last_errors = []
|
||||
self.terminal_since_spec = 0
|
||||
self.eta_start = None
|
||||
|
||||
# counters (computed from job states so they can never drift)
|
||||
def count(self, *statuses):
|
||||
return sum(1 for j in self.jobs if j["status"] in statuses)
|
||||
|
||||
@property
|
||||
def skipped_existing(self):
|
||||
return sum(1 for j in self.jobs if j["status"] == "skipped" and j.get("skip_reason") == "existing")
|
||||
|
||||
@property
|
||||
def to_run(self):
|
||||
return sum(1 for j in self.jobs if j["status"] != "skipped")
|
||||
|
||||
def in_flight(self):
|
||||
return [j for j in self.jobs if j["status"] in ("enqueued", "running")]
|
||||
|
||||
|
||||
_sweeps = {}
|
||||
_registry_lock = threading.RLock()
|
||||
|
||||
|
||||
def reset_registry():
|
||||
"""Tests only."""
|
||||
with _registry_lock:
|
||||
_sweeps.clear()
|
||||
|
||||
|
||||
def _active_sweep():
|
||||
with _registry_lock:
|
||||
for sw in _sweeps.values():
|
||||
if sw.state in ACTIVE_STATES:
|
||||
return sw
|
||||
return None
|
||||
|
||||
|
||||
def _find(sweep_id=None, session=None):
|
||||
with _registry_lock:
|
||||
if sweep_id:
|
||||
return _sweeps.get(str(sweep_id))
|
||||
if session:
|
||||
matches = [s for s in _sweeps.values() if s.session == session]
|
||||
return max(matches, key=lambda s: s.created_at) if matches else None
|
||||
return None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- prepare / validate / launch
|
||||
|
||||
def _manifest_state(manifest_path):
|
||||
"""({fingerprint: [item ids]}, next gen_index) from the manifest on disk."""
|
||||
data = wfs.read_manifest_items(manifest_path)
|
||||
by_fp, max_gen = {}, -1
|
||||
for it in data.get("items") or []:
|
||||
if not isinstance(it, dict):
|
||||
continue
|
||||
fp = it.get("wf_fingerprint")
|
||||
if fp and it.get("id") is not None:
|
||||
by_fp.setdefault(fp, []).append(it["id"])
|
||||
g = it.get("gen_index")
|
||||
if isinstance(g, int) and g > max_gen:
|
||||
max_gen = g
|
||||
return by_fp, max_gen + 1
|
||||
|
||||
|
||||
def prepare_sweep(configs, session_settings=None, origin_node=None, client_id=None, bridge=None,
|
||||
class_info=wfs.comfy_class_info, check_busy=True):
|
||||
"""
|
||||
Expand + fingerprint + skip finished jobs. Raises ValueError (readable) or SweepBusy.
|
||||
Pure apart from reading the session manifest.
|
||||
"""
|
||||
busy = _active_sweep() if check_busy else None
|
||||
if busy is not None:
|
||||
raise SweepBusy(busy.id)
|
||||
settings = dict(session_settings or {})
|
||||
if not isinstance(configs, list) or not configs:
|
||||
raise ValueError("There are no configs to run.")
|
||||
extra_seeds = []
|
||||
try:
|
||||
n_extra = int(settings.get("add_random_seeds_to_gens", 0) or 0)
|
||||
except (TypeError, ValueError):
|
||||
n_extra = 0
|
||||
if n_extra > 0:
|
||||
extra_seeds = [random.randint(0, 2**32 - 1) for _ in range(n_extra)]
|
||||
plan = wfs.plan_workflow_configs(configs, extra_seeds, class_info=class_info)
|
||||
|
||||
paths = wfs.session_paths(settings.get("session_name") or "workflow_sweep")
|
||||
overwrite = bool(settings.get("overwrite_existing", False))
|
||||
try:
|
||||
start_at = max(0, int(settings.get("start_at_job", 0) or 0))
|
||||
except (TypeError, ValueError):
|
||||
start_at = 0
|
||||
by_fp, gen_base = _manifest_state(paths["manifest"])
|
||||
|
||||
jobs = []
|
||||
for idx, conf in enumerate(plan["jobs"]):
|
||||
job = {"index": idx, "fingerprint": conf["wf_fingerprint"], "conf": conf, "status": "pending",
|
||||
"prompt_id": None, "error": None, "gen_index": gen_base + idx, "replace_ids": [],
|
||||
"enqueued_at": None, "started_at": None, "finished_at": None, "duration": None,
|
||||
"collected": False, "missing_since": None, "success_seen_at": None}
|
||||
if idx < start_at:
|
||||
job["status"], job["skip_reason"] = "skipped", "start_at_job"
|
||||
elif conf["wf_fingerprint"] in by_fp:
|
||||
if overwrite:
|
||||
job["replace_ids"] = list(by_fp[conf["wf_fingerprint"]])
|
||||
else:
|
||||
job["status"], job["skip_reason"] = "skipped", "existing"
|
||||
jobs.append(job)
|
||||
|
||||
sw = Sweep(uuid.uuid4().hex[:12], paths["session"], jobs, plan["templates"], origin_node,
|
||||
client_id or None, {"image_format": str(settings.get("image_format", "webp")),
|
||||
"overwrite_existing": overwrite, "start_at_job": start_at,
|
||||
"total_jobs": len(jobs)},
|
||||
list(plan["warnings"]), plan["axes"], bridge or get_bridge())
|
||||
if sw.client_id is None:
|
||||
sw.client_id = f"uscg-sweep-{sw.id}"
|
||||
return sw
|
||||
|
||||
|
||||
def build_job_request(sw, job):
|
||||
"""(prompt, partial_execution_targets) for one job."""
|
||||
conf = job["conf"]
|
||||
tpl = sw.templates[conf["wf_template_hash"]]
|
||||
job_json = {"conf": conf, "sweep_id": sw.id, "job_index": job["index"], "gen_index": job["gen_index"],
|
||||
"origin_node": sw.origin_node, "replace_item_ids": job["replace_ids"],
|
||||
"image_format": sw.settings.get("image_format", "webp")}
|
||||
prompt, cid = wfs.build_job_prompt(tpl["prompt"], conf["wf_subs"], tpl["output"],
|
||||
{"session_name": sw.session,
|
||||
"job_json": json.dumps(job_json, ensure_ascii=False)})
|
||||
partial = None if tpl["run_other_outputs"] else [cid]
|
||||
return prompt, partial
|
||||
|
||||
|
||||
async def validate_first(sw):
|
||||
"""Validate the first job that will run (on the event loop). Raises ValueError with node errors."""
|
||||
job = next((j for j in sw.jobs if j["status"] == "pending"), None)
|
||||
if job is None:
|
||||
return
|
||||
prompt, partial = build_job_request(sw, job)
|
||||
ok, error, node_errors = await sw.bridge.validate(prompt, partial)
|
||||
if not ok:
|
||||
raise ValueError("The first job is not valid: " + summarize_errors(error, node_errors, prompt))
|
||||
|
||||
|
||||
def _new_meta(sw):
|
||||
first = sw.jobs[0]["conf"] if sw.jobs else {}
|
||||
return {"positive": first.get("positive", ""), "negative": first.get("negative", ""),
|
||||
"model": first.get("model", ""), "seed": first.get("seed", 0) if isinstance(first.get("seed", 0), int) else 0,
|
||||
"vae_batch_size": 1, "configs_json": "", "resolutions_json": "[]",
|
||||
"image_format": sw.settings.get("image_format", "webp"), "workflow_sweep": True}
|
||||
|
||||
|
||||
def ensure_manifest(sw):
|
||||
"""Create the session manifest when missing (existing manifests are left alone)."""
|
||||
paths = wfs.session_paths(sw.session)
|
||||
os.makedirs(paths["images"], exist_ok=True)
|
||||
if not os.path.exists(paths["manifest"]):
|
||||
data = {"session_name": sw.session, "meta": _new_meta(sw), "items": []}
|
||||
with open(paths["manifest"], "w") as f:
|
||||
json.dump(data, f, indent=4)
|
||||
return data
|
||||
data = wfs.read_manifest_items(paths["manifest"])
|
||||
return data
|
||||
|
||||
|
||||
def launch(sw, start_thread=True):
|
||||
"""Register the sweep and start its feeder thread (tests drive tick() themselves with start_thread=False)."""
|
||||
with _registry_lock:
|
||||
busy = _active_sweep()
|
||||
if busy is not None and busy is not sw:
|
||||
raise SweepBusy(busy.id)
|
||||
_sweeps[sw.id] = sw
|
||||
ensure_manifest(sw)
|
||||
with sw.lock:
|
||||
sw.state = "running"
|
||||
sw.started_at = sw.eta_start = time.time()
|
||||
print(f"[GridTester] 🧪 Workflow sweep {sw.id}: {sw.to_run} job(s) to run, "
|
||||
f"{sw.skipped_existing} already in the manifest (session '{sw.session}')")
|
||||
for w in sw.warnings:
|
||||
print(f"[GridTester] ⚠️ {w}")
|
||||
_write_spec(sw)
|
||||
_progress(sw)
|
||||
if start_thread:
|
||||
sw.thread = threading.Thread(target=_run, args=(sw,), name=f"uscg-sweep-{sw.id}", daemon=True)
|
||||
sw.thread.start()
|
||||
return sw
|
||||
|
||||
|
||||
def start_sweep(configs, session_settings=None, origin_node=None, client_id=None, bridge=None,
|
||||
class_info=wfs.comfy_class_info):
|
||||
"""Prepare + validate job 1 on the loop + launch. For a NON-loop thread (the Generator node).
|
||||
Returns immediately after launch; never waits for any sweep prompt."""
|
||||
sw = prepare_sweep(configs, session_settings, origin_node, client_id, bridge, class_info)
|
||||
sw.bridge.run_on_loop(validate_first(sw))
|
||||
return launch(sw)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- feeder thread
|
||||
|
||||
def _run(sw):
|
||||
try:
|
||||
while True:
|
||||
alive = tick(sw)
|
||||
if not alive:
|
||||
break
|
||||
sw.wake.wait(TICK)
|
||||
sw.wake.clear()
|
||||
except Exception as e: # never let the feeder die silently
|
||||
traceback.print_exc()
|
||||
with sw.lock:
|
||||
sw.state, sw.reason = "failed", f"internal error: {e}"
|
||||
sw.finished_at = time.time()
|
||||
_write_spec(sw)
|
||||
_progress(sw)
|
||||
|
||||
|
||||
def _fail_job(sw, job, status, error, params=None):
|
||||
job["status"] = status
|
||||
job["error"] = error
|
||||
job["finished_at"] = time.time()
|
||||
sw.last_errors.append({"index": job["index"], "status": status, "error": error,
|
||||
"params": job["conf"].get("wf_params")})
|
||||
del sw.last_errors[:-20]
|
||||
sw.consecutive_errors += 1
|
||||
sw.terminal_since_spec += 1
|
||||
print(f"[GridTester] ❌ Sweep job {job['index'] + 1}/{len(sw.jobs)} {status}: {error}")
|
||||
if sw.consecutive_errors >= BREAKER and sw.state == "running":
|
||||
sw.state = "paused"
|
||||
sw.reason = f"{BREAKER} jobs in a row failed; last: {error}"
|
||||
|
||||
|
||||
def _classify(history):
|
||||
"""-> (kind, detail) with kind in success|error|interrupted|unknown, from a ComfyUI history entry."""
|
||||
status = (history or {}).get("status") or {}
|
||||
messages = status.get("messages") or []
|
||||
names = {}
|
||||
for m in messages:
|
||||
if isinstance(m, (list, tuple)) and len(m) == 2:
|
||||
names[m[0]] = m[1] if isinstance(m[1], dict) else {}
|
||||
if "execution_interrupted" in names:
|
||||
return "interrupted", names["execution_interrupted"]
|
||||
if "execution_error" in names or status.get("status_str") == "error":
|
||||
d = names.get("execution_error", {})
|
||||
text = d.get("exception_message") or "execution error"
|
||||
if d.get("node_type") or d.get("node_id"):
|
||||
text = f"{d.get('node_type', '?')} #{d.get('node_id', '?')}: {text}"
|
||||
return "error", {"error": text.strip()[:1500], "start": names.get("execution_start"),
|
||||
"end": names.get("execution_error")}
|
||||
if status.get("status_str") == "success" or "execution_success" in names:
|
||||
return "success", {"start": names.get("execution_start"), "end": names.get("execution_success")}
|
||||
return "unknown", {}
|
||||
|
||||
|
||||
def _poll(sw, inflight):
|
||||
b = sw.bridge
|
||||
now = time.time()
|
||||
running_ids, queued_ids = b.queue_state() # snapshot FIRST, then history (race-free ordering)
|
||||
for job in inflight:
|
||||
pid = job["prompt_id"]
|
||||
if pid in running_ids:
|
||||
job["missing_since"] = None
|
||||
if job["status"] == "enqueued":
|
||||
job["status"], job["started_at"] = "running", now
|
||||
continue
|
||||
if pid in queued_ids:
|
||||
job["missing_since"] = None
|
||||
continue
|
||||
history = b.history_entry(pid)
|
||||
if history is None:
|
||||
job["missing_since"] = job["missing_since"] or now
|
||||
if now - job["missing_since"] > LOST_GRACE:
|
||||
with sw.lock:
|
||||
job["status"] = "lost"
|
||||
job["error"] = "the prompt disappeared from ComfyUI's queue and history (queue cleared?)"
|
||||
if sw.state == "running":
|
||||
sw.state, sw.reason = "paused", "a sweep prompt was removed from ComfyUI's queue; resume to re-run it"
|
||||
print(f"[GridTester] ⚠️ Sweep job {job['index'] + 1} lost; sweep paused")
|
||||
continue
|
||||
kind, detail = _classify(history)
|
||||
with sw.lock:
|
||||
if kind == "success":
|
||||
job["success_seen_at"] = job["success_seen_at"] or now
|
||||
if job["collected"]:
|
||||
_finish_done(sw, job, detail, now)
|
||||
elif now - job["success_seen_at"] > COLLECT_GRACE:
|
||||
_fail_job(sw, job, "error", "the job ran but the collector produced no image")
|
||||
_after_terminal(sw)
|
||||
elif kind == "error":
|
||||
_fail_job(sw, job, "error", detail["error"])
|
||||
_after_terminal(sw)
|
||||
elif kind == "interrupted":
|
||||
job["status"], job["error"] = "interrupted", "interrupted in ComfyUI"
|
||||
if sw.state == "running":
|
||||
sw.state, sw.reason = "paused", "a job was interrupted in ComfyUI; resume to re-run it"
|
||||
print(f"[GridTester] ⏸️ Sweep job {job['index'] + 1} interrupted; sweep paused")
|
||||
else:
|
||||
job["missing_since"] = job["missing_since"] or now
|
||||
if now - job["missing_since"] > LOST_GRACE:
|
||||
_fail_job(sw, job, "error", "ComfyUI finished the prompt without a recognisable status")
|
||||
_after_terminal(sw)
|
||||
|
||||
|
||||
def _duration_from(detail, fallback):
|
||||
try:
|
||||
start, end = detail.get("start") or {}, detail.get("end") or {}
|
||||
d = (float(end["timestamp"]) - float(start["timestamp"])) / 1000.0
|
||||
if d >= 0:
|
||||
return d
|
||||
except (KeyError, TypeError, ValueError):
|
||||
pass
|
||||
return fallback
|
||||
|
||||
|
||||
def _finish_done(sw, job, detail, now):
|
||||
job["status"] = "done"
|
||||
job["finished_at"] = now
|
||||
fallback = (now - job["started_at"]) if job["started_at"] else 0.0
|
||||
job["duration"] = _duration_from(detail, fallback)
|
||||
sw.durations.append(job["duration"])
|
||||
sw.consecutive_errors = 0
|
||||
sw.terminal_since_spec += 1
|
||||
_after_terminal(sw)
|
||||
|
||||
|
||||
def _after_terminal(sw):
|
||||
_progress(sw)
|
||||
if sw.terminal_since_spec >= SPEC_EVERY:
|
||||
sw.terminal_since_spec = 0
|
||||
_write_spec(sw)
|
||||
|
||||
|
||||
def _enqueue_job(sw, job):
|
||||
try:
|
||||
prompt, partial = build_job_request(sw, job)
|
||||
except Exception as e:
|
||||
with sw.lock:
|
||||
_fail_job(sw, job, "invalid", f"could not build the prompt: {e}")
|
||||
_after_terminal(sw)
|
||||
return
|
||||
try:
|
||||
pid, error, node_errors = sw.bridge.run_on_loop(sw.bridge.enqueue(prompt, partial, sw.client_id))
|
||||
except Exception as e:
|
||||
with sw.lock:
|
||||
_fail_job(sw, job, "invalid", f"could not queue the prompt: {e}")
|
||||
_after_terminal(sw)
|
||||
return
|
||||
with sw.lock:
|
||||
if pid is None:
|
||||
_fail_job(sw, job, "invalid", summarize_errors(error, node_errors, prompt))
|
||||
_after_terminal(sw)
|
||||
else:
|
||||
job["status"], job["prompt_id"], job["enqueued_at"] = "enqueued", pid, time.time()
|
||||
job["missing_since"] = None
|
||||
|
||||
|
||||
def _do_cleanup(sw):
|
||||
"""Cancel: drop our queued prompts, interrupt the running one, mark the rest cancelled."""
|
||||
b = sw.bridge
|
||||
for job in list(sw.in_flight()):
|
||||
pid = job["prompt_id"]
|
||||
try:
|
||||
if job["status"] == "enqueued" and b.delete_queued(pid):
|
||||
pass
|
||||
else:
|
||||
b.interrupt(pid)
|
||||
except Exception as e:
|
||||
print(f"[GridTester] ⚠️ Sweep cancel could not remove {pid}: {e}")
|
||||
with sw.lock:
|
||||
for job in sw.jobs:
|
||||
if job["status"] in ("pending", "enqueued", "running", "interrupted", "lost"):
|
||||
job["status"] = "cancelled"
|
||||
sw.needs_cleanup = False
|
||||
sw.finished_at = time.time()
|
||||
_write_spec(sw)
|
||||
_progress(sw)
|
||||
|
||||
|
||||
def _finish(sw):
|
||||
with sw.lock:
|
||||
errors = sw.count("error", "invalid")
|
||||
sw.state = "completed"
|
||||
sw.reason = f"{errors} job(s) failed" if errors else ""
|
||||
sw.finished_at = time.time()
|
||||
print(f"[GridTester] 🎉 Workflow sweep {sw.id} complete: {sw.count('done')} image(s), {sw.count('error', 'invalid')} failed")
|
||||
_write_spec(sw)
|
||||
_progress(sw, complete=True)
|
||||
|
||||
|
||||
def tick(sw):
|
||||
"""One feeder iteration. Returns False when the sweep is over (thread should exit)."""
|
||||
b = sw.bridge
|
||||
with sw.lock:
|
||||
if sw.state in ("completed", "failed") or (sw.state == "cancelled" and not sw.needs_cleanup):
|
||||
return False
|
||||
cleanup = sw.state == "cancelled" and sw.needs_cleanup
|
||||
inflight = sw.in_flight()
|
||||
if cleanup:
|
||||
_do_cleanup(sw)
|
||||
return False
|
||||
|
||||
if inflight:
|
||||
_poll(sw, inflight)
|
||||
|
||||
with sw.lock:
|
||||
state = sw.state
|
||||
queued = [j for j in sw.in_flight() if j["status"] == "enqueued"]
|
||||
if state == "paused":
|
||||
for job in queued: # pull what has not started; a started job finishes
|
||||
try:
|
||||
pulled = b.delete_queued(job["prompt_id"])
|
||||
except Exception:
|
||||
pulled = False
|
||||
if pulled:
|
||||
with sw.lock:
|
||||
job.update(status="pending", prompt_id=None, enqueued_at=None, missing_since=None)
|
||||
return True
|
||||
|
||||
while True:
|
||||
with sw.lock:
|
||||
if sw.state != "running":
|
||||
break
|
||||
nxt = next((j for j in sw.jobs if j["status"] == "pending"), None)
|
||||
if nxt is None or len(sw.in_flight()) >= WINDOW:
|
||||
break
|
||||
_enqueue_job(sw, nxt)
|
||||
|
||||
with sw.lock:
|
||||
if sw.state != "running":
|
||||
return True
|
||||
if not sw.in_flight() and not any(j["status"] == "pending" for j in sw.jobs):
|
||||
finished = True
|
||||
else:
|
||||
finished = False
|
||||
if finished:
|
||||
_finish(sw)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- public controls
|
||||
|
||||
def pause(sweep_id):
|
||||
sw = _find(sweep_id)
|
||||
if sw is None:
|
||||
return False, "unknown sweep"
|
||||
with sw.lock:
|
||||
if sw.state != "running":
|
||||
return False, f"sweep is {sw.state}"
|
||||
sw.state, sw.reason = "paused", "paused by user"
|
||||
sw.wake.set()
|
||||
return True, "paused"
|
||||
|
||||
|
||||
def resume(sweep_id):
|
||||
sw = _find(sweep_id)
|
||||
if sw is None:
|
||||
return False, "unknown sweep"
|
||||
with sw.lock:
|
||||
if sw.state != "paused":
|
||||
return False, f"sweep is {sw.state}"
|
||||
for job in sw.jobs:
|
||||
if job["status"] in ("interrupted", "lost"): # failed jobs stay failed; these two never really ran
|
||||
job.update(status="pending", prompt_id=None, error=None, missing_since=None, success_seen_at=None)
|
||||
sw.consecutive_errors = 0
|
||||
sw.state, sw.reason = "running", ""
|
||||
sw.wake.set()
|
||||
return True, "running"
|
||||
|
||||
|
||||
def cancel(sweep_id):
|
||||
sw = _find(sweep_id)
|
||||
if sw is None:
|
||||
return False, "unknown sweep"
|
||||
with sw.lock:
|
||||
if sw.state not in ACTIVE_STATES:
|
||||
return False, f"sweep is {sw.state}"
|
||||
sw.state, sw.reason, sw.needs_cleanup = "cancelled", "cancelled by user", True
|
||||
sw.wake.set()
|
||||
return True, "cancelled"
|
||||
|
||||
|
||||
def on_collected(sweep_id, job_index, items):
|
||||
"""Called by the collector node after it saved a job's image(s). False when the sweep is unknown."""
|
||||
sw = _find(sweep_id)
|
||||
if sw is None:
|
||||
return False
|
||||
try:
|
||||
job = sw.jobs[int(job_index)]
|
||||
except (IndexError, TypeError, ValueError):
|
||||
return False
|
||||
with sw.lock:
|
||||
job["collected"] = True
|
||||
job["items"] = [it.get("id") for it in (items or []) if isinstance(it, dict)]
|
||||
return True
|
||||
|
||||
|
||||
def job_started_at(sweep_id, job_index):
|
||||
sw = _find(sweep_id)
|
||||
if sw is None:
|
||||
return None
|
||||
try:
|
||||
return sw.jobs[int(job_index)]["started_at"]
|
||||
except (IndexError, TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def get_status(sweep_id=None, session=None):
|
||||
sw = _find(sweep_id, session)
|
||||
if sw is None:
|
||||
return None
|
||||
with sw.lock:
|
||||
done = sw.count("done")
|
||||
errors = sw.count("error", "invalid")
|
||||
recent = sw.durations[-10:]
|
||||
avg = sum(recent) / len(recent) if recent else 0.0
|
||||
remaining = sw.count("pending", "enqueued", "running", "interrupted", "lost")
|
||||
return {"sweep_id": sw.id, "session_name": sw.session, "state": sw.state, "reason": sw.reason,
|
||||
"total": len(sw.jobs), "to_run": sw.to_run, "skipped_existing": sw.skipped_existing,
|
||||
"done": done, "errors": errors, "in_flight": len(sw.in_flight()), "pending": sw.count("pending"),
|
||||
"eta_str": _eta_str(avg * remaining) if recent and sw.state == "running" else "",
|
||||
"avg_duration": round(avg, 1), "warnings": list(sw.warnings),
|
||||
"last_errors": list(sw.last_errors)}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- progress + spec file
|
||||
|
||||
def _eta_str(seconds):
|
||||
seconds = int(max(0, seconds))
|
||||
h, m, s = seconds // 3600, (seconds % 3600) // 60, seconds % 60
|
||||
return f"{h}h {m}m {s}s" if h else (f"{m}m {s}s" if m else f"{s}s")
|
||||
|
||||
|
||||
def _progress(sw, complete=False):
|
||||
"""ultimate_grid.progress for the Dashboard's ETA bar (phase 'sweep')."""
|
||||
try:
|
||||
with sw.lock:
|
||||
total = len(sw.jobs)
|
||||
done = sw.count("done")
|
||||
errors = sw.count("error", "invalid")
|
||||
current = sw.skipped_existing + done
|
||||
remaining = sw.count("pending", "enqueued", "running", "interrupted", "lost")
|
||||
recent = sw.durations[-10:]
|
||||
avg = sum(recent) / len(recent) if recent else 0.0
|
||||
payload = {
|
||||
"node": sw.origin_node, "session_name": sw.session, "phase": "sweep", "sweep_id": sw.id,
|
||||
"state": sw.state, "errors": errors, "current_job": current, "total_jobs": total,
|
||||
"progress_pct": int((current + errors) / total * 100) if total else 100,
|
||||
"eta_str": _eta_str(avg * remaining) if recent and remaining else "—",
|
||||
"finish_time": time.strftime("%H:%M:%S", time.localtime(time.time() + avg * remaining)) if recent and remaining else "",
|
||||
"avg_duration": round(avg, 1), "last_duration": round(sw.durations[-1], 1) if sw.durations else 0,
|
||||
}
|
||||
if complete:
|
||||
elapsed = time.time() - (sw.eta_start or sw.created_at)
|
||||
payload.update({"complete": True, "eta_str": "Done", "progress_pct": 100,
|
||||
"total_elapsed": f"{int(elapsed // 3600)}h {int((elapsed % 3600) // 60)}m {int(elapsed % 60)}s",
|
||||
"total_generated": done})
|
||||
sw.bridge.send_sync("ultimate_grid.progress", payload)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _write_spec(sw):
|
||||
"""benchmarks/<session>/workflow_sweeps/<sweep_id>.json: templates, bindings and the job list, so any
|
||||
executed prompt can be reconstructed. Never served by any route."""
|
||||
try:
|
||||
paths = wfs.session_paths(sw.session)
|
||||
os.makedirs(paths["sweeps"], exist_ok=True)
|
||||
with sw.lock:
|
||||
spec = {
|
||||
"sweep_id": sw.id, "session_name": sw.session, "state": sw.state, "reason": sw.reason,
|
||||
"created_at": sw.created_at, "settings": sw.settings, "axes": [a for c in sw.axes for a in c["axes"]], "warnings": sw.warnings,
|
||||
"templates": {h: {"prompt": t["prompt"], "output": t["output"],
|
||||
"run_other_outputs": t["run_other_outputs"]} for h, t in sw.templates.items()},
|
||||
"jobs": [{"index": j["index"], "fingerprint": j["fingerprint"], "status": j["status"],
|
||||
"prompt_id": j["prompt_id"], "error": j["error"], "duration": j["duration"],
|
||||
"template": j["conf"].get("wf_template_hash"), "subs": j["conf"].get("wf_subs"),
|
||||
"params": j["conf"].get("wf_params")} for j in sw.jobs],
|
||||
}
|
||||
path = os.path.join(paths["sweeps"], f"{sw.id}.json")
|
||||
with open(path + ".tmp", "w") as f:
|
||||
json.dump(spec, f, indent=1)
|
||||
os.replace(path + ".tmp", path)
|
||||
except Exception as e:
|
||||
print(f"[GridTester] ⚠️ Could not write the sweep spec file: {e}")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- request-body helpers (routes)
|
||||
|
||||
def configs_from_body(body):
|
||||
"""(configs list, session_settings) from {state} (Builder state, through the single source of truth)
|
||||
or {configs_json}. Raises ValueError."""
|
||||
if not isinstance(body, dict):
|
||||
raise ValueError("The request body must be a JSON object.")
|
||||
if isinstance(body.get("state"), dict):
|
||||
try:
|
||||
from .config_builder_node import UltimateConfigBuilder
|
||||
except ImportError:
|
||||
from config_builder_node import UltimateConfigBuilder
|
||||
raw = UltimateConfigBuilder.state_to_configs_json(body["state"])
|
||||
elif body.get("configs_json") is not None:
|
||||
raw = body["configs_json"]
|
||||
else:
|
||||
raise ValueError("Send {state} (the Builder state) or {configs_json}.")
|
||||
parsed = json.loads(raw) if isinstance(raw, (str, bytes)) else raw
|
||||
if isinstance(parsed, dict):
|
||||
return parsed.get("configs") or [], parsed.get("_session_settings") or {}
|
||||
return parsed, {}
|
||||
|
||||
|
||||
async def dry_run(body, bridge=None, class_info=wfs.comfy_class_info):
|
||||
"""Concrete first prompts (with collector) + validation of the first job, nothing is queued."""
|
||||
configs, settings = configs_from_body(body)
|
||||
try:
|
||||
limit = max(1, min(20, int(body.get("limit", 3))))
|
||||
except (TypeError, ValueError):
|
||||
limit = 3
|
||||
sw = prepare_sweep(configs, settings, None, None, bridge, class_info, check_busy=False)
|
||||
jobs_out = []
|
||||
for job in sw.jobs[:limit]:
|
||||
prompt, partial = build_job_request(sw, job)
|
||||
jobs_out.append({"index": job["index"], "fingerprint": job["fingerprint"], "skip": job["status"] == "skipped",
|
||||
"params": job["conf"].get("wf_params"), "partial_execution_targets": partial,
|
||||
"prompt": prompt})
|
||||
validation = {"ok": True, "error": None, "node_errors": {}}
|
||||
first = next((j for j in sw.jobs if j["status"] == "pending"), None)
|
||||
if first is not None:
|
||||
prompt, partial = build_job_request(sw, first)
|
||||
try:
|
||||
ok, error, node_errors = await sw.bridge.validate(prompt, partial)
|
||||
validation = {"ok": ok, "error": summarize_errors(error, node_errors, prompt) if not ok else None,
|
||||
"node_errors": node_errors if not ok else {}}
|
||||
except Exception as e:
|
||||
validation = {"ok": False, "error": f"validation could not run: {e}", "node_errors": {}}
|
||||
return {"total": len(sw.jobs), "to_run": sw.to_run, "skipped_existing": sw.skipped_existing,
|
||||
"axes": [a for c in sw.axes for a in c["axes"]], "warnings": sw.warnings, "validation": validation, "jobs": jobs_out}
|
||||
|
||||
|
||||
async def api_start(body, bridge=None, class_info=wfs.comfy_class_info):
|
||||
"""POST /config_tester/sweep/start. Returns (http_status, json_body). Runs on the event loop."""
|
||||
try:
|
||||
configs, settings = configs_from_body(body)
|
||||
sw = prepare_sweep(configs, settings, body.get("node_id"), body.get("client_id"), bridge, class_info)
|
||||
await validate_first(sw)
|
||||
launch(sw)
|
||||
except SweepBusy as e:
|
||||
return 409, {"error": str(e), "sweep_id": e.sweep_id}
|
||||
except ValueError as e:
|
||||
return 400, {"error": str(e)}
|
||||
return 200, {"sweep_id": sw.id, "total": len(sw.jobs), "to_run": sw.to_run,
|
||||
"skipped_existing": sw.skipped_existing, "warnings": sw.warnings}
|
||||
Reference in New Issue
Block a user