Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6db5692d78 | ||
|
|
591ffc7396 | ||
|
|
5f9a2f0fc6 | ||
|
|
c30036ae31 | ||
|
|
41caf02481 | ||
|
|
494d5506a0 | ||
|
|
7ed60b93ca |
@@ -27,6 +27,11 @@
|
||||
- Intelligent distribution
|
||||
- Handles single images and videos
|
||||
|
||||
#### Dynamic Image-List Processing
|
||||
- Process a whole image list in parallel: the master and every worker take the next image when they free up, and results come back in input order
|
||||
- Use **Distributed Image List Start / End** around a per-image processing region
|
||||
- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
|
||||
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
@@ -43,7 +48,7 @@ ComfyUI Distributed supports four types of workers:
|
||||
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
|
||||
- **Remote Workers** - GPUs on other computers in your network
|
||||
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
|
||||
- **Serverless (Runpod)** - Queue-based Runpod endpoints for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
- **Serverless (Runpod)** - Queue-based endpoints for image Collector jobs; no Distributed install or inbound tunnel on the endpoint. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
|
||||
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
|
||||
|
||||
@@ -179,6 +184,8 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
|
||||
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
|
||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||
| **Distributed Image List Start** | Turns an image list or batch into one job per image, scheduled to the next available participant; mask and value inputs are shared across images or paired one per image |
|
||||
| **Distributed Image List End** | Closes the paired region and returns images, masks, and texts as lists in the original input order |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -4,3 +4,4 @@ from . import tunnel_routes # noqa: F401
|
||||
from . import worker_routes # noqa: F401
|
||||
from . import job_routes # noqa: F401
|
||||
from . import usdu_routes # noqa: F401
|
||||
from . import image_list_routes # noqa: F401
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Bounded, typed HTTP transport for dynamic image-list worker tasks."""
|
||||
import json
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..image_list.runtime import (
|
||||
MAX_REQUEST_BYTES,
|
||||
PROTOCOL_VERSION,
|
||||
decode_result,
|
||||
delete_worker_item,
|
||||
set_master_result,
|
||||
store_worker_item,
|
||||
)
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
async def _json_body(request):
|
||||
length = request.content_length
|
||||
if length is not None and length > MAX_REQUEST_BYTES:
|
||||
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
|
||||
raw = bytearray()
|
||||
async for chunk in request.content.iter_chunked(64 * 1024):
|
||||
raw.extend(chunk)
|
||||
if len(raw) > MAX_REQUEST_BYTES:
|
||||
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
|
||||
try:
|
||||
return json.loads(raw.decode("utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f"Invalid image-list JSON body: {exc}") from exc
|
||||
|
||||
|
||||
@prompt_server.routes.get("/distributed/image-list/capabilities")
|
||||
async def image_list_capabilities_endpoint(_request):
|
||||
return web.json_response({
|
||||
"supported": True,
|
||||
"execution_targets": True,
|
||||
"value_results": True,
|
||||
"protocol_version": PROTOCOL_VERSION,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd", "DistributedImageListGetValue"],
|
||||
"transport": "typed-float-tensor-json-v1",
|
||||
})
|
||||
|
||||
|
||||
@prompt_server.routes.post("/distributed/image-list/item")
|
||||
async def image_list_item_endpoint(request):
|
||||
try:
|
||||
data = await _json_body(request)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected an image-list item object")
|
||||
context_id = data.get("context_id")
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise ValueError("context_id must be a non-empty string")
|
||||
item = data.get("item")
|
||||
identity = store_worker_item(data.get("token"), item)
|
||||
return web.json_response({"status": "stored", "item_index": item["index"],
|
||||
"item_identity": identity})
|
||||
except (ValueError, TypeError) as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
except Exception as exc:
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
|
||||
|
||||
@prompt_server.routes.delete("/distributed/image-list/item/{token}")
|
||||
async def image_list_delete_item_endpoint(request):
|
||||
token = request.match_info.get("token", "")
|
||||
return web.json_response({"status": "deleted", "removed": delete_worker_item(token)})
|
||||
|
||||
|
||||
@prompt_server.routes.post("/distributed/image-list/result")
|
||||
async def image_list_result_endpoint(request):
|
||||
try:
|
||||
data = await _json_body(request)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected an image-list result object")
|
||||
context_id = data.get("context_id")
|
||||
token = data.get("token")
|
||||
item_index = data.get("item_index")
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise ValueError("context_id must be a non-empty string")
|
||||
if not isinstance(token, str) or not token:
|
||||
raise ValueError("token must be a non-empty string")
|
||||
if isinstance(item_index, bool) or not isinstance(item_index, int) or item_index < 0:
|
||||
raise ValueError("item_index must be a non-negative integer")
|
||||
result = decode_result(data.get("result"))
|
||||
set_master_result(context_id, item_index, token, result)
|
||||
return web.json_response({"status": "accepted", "item_index": item_index})
|
||||
except (ValueError, TypeError) as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
except Exception as exc:
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
+156
-22
@@ -6,6 +6,9 @@ import aiohttp
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session, probe_worker
|
||||
# Must match image_list.runtime.PROTOCOL_VERSION (kept local so this module
|
||||
# stays importable without the image-list runtime; a test pins them equal).
|
||||
IMAGE_LIST_PROTOCOL_VERSION = 2
|
||||
try:
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
except ImportError:
|
||||
@@ -74,30 +77,49 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
|
||||
"workflow": payload.get("workflow"),
|
||||
"client_id": client_id,
|
||||
}
|
||||
if "partial_execution_targets" in payload:
|
||||
ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
|
||||
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
|
||||
ws_url = f"{ws_url}/distributed/worker_ws"
|
||||
session = await get_client_session()
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + timeout
|
||||
ws = None
|
||||
|
||||
async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
|
||||
await ws.send_json(ws_payload)
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
data = json.loads(msg.data or "{}")
|
||||
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
|
||||
if data.get("ok"):
|
||||
return
|
||||
error_text = data.get("error") or "Worker rejected websocket dispatch."
|
||||
validation_error = data.get("validation_error")
|
||||
node_errors = data.get("node_errors")
|
||||
if validation_error:
|
||||
error_text = f"{error_text} | validation_error={validation_error}"
|
||||
if node_errors:
|
||||
error_text = f"{error_text} | node_errors={node_errors}"
|
||||
raise RuntimeError(error_text)
|
||||
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
|
||||
raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
|
||||
|
||||
raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
|
||||
try:
|
||||
async with asyncio.timeout_at(deadline):
|
||||
ws = await session.ws_connect(ws_url, heartbeat=20)
|
||||
await ws.send_json(ws_payload)
|
||||
while True:
|
||||
msg = await ws.receive()
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
data = json.loads(msg.data or "{}")
|
||||
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
|
||||
if data.get("ok"):
|
||||
return data.get("prompt_id")
|
||||
error_text = data.get("error") or "Worker rejected websocket dispatch."
|
||||
validation_error = data.get("validation_error")
|
||||
node_errors = data.get("node_errors")
|
||||
if validation_error:
|
||||
error_text = f"{error_text} | validation_error={validation_error}"
|
||||
if node_errors:
|
||||
error_text = f"{error_text} | node_errors={node_errors}"
|
||||
raise RuntimeError(error_text)
|
||||
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED):
|
||||
raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(
|
||||
"Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
|
||||
) from exc
|
||||
finally:
|
||||
if ws is not None:
|
||||
# Closing must not discard an already acknowledged prompt ID or
|
||||
# replace a definite rejection with an ambiguous timeout.
|
||||
try:
|
||||
async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)):
|
||||
await ws.close()
|
||||
except (asyncio.TimeoutError, aiohttp.ClientError):
|
||||
pass
|
||||
|
||||
|
||||
async def dispatch_worker_prompt(
|
||||
@@ -107,11 +129,18 @@ async def dispatch_worker_prompt(
|
||||
client_id=None,
|
||||
use_websocket=False,
|
||||
trace_execution_id=None,
|
||||
return_prompt_id=False,
|
||||
*,
|
||||
partial_execution_targets=None,
|
||||
):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
worker_url = build_worker_url(worker)
|
||||
url = build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
target_payload = {}
|
||||
if partial_execution_targets is not None:
|
||||
target_payload["partial_execution_targets"] = list(partial_execution_targets)
|
||||
payload.update(target_payload)
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
@@ -120,15 +149,18 @@ async def dispatch_worker_prompt(
|
||||
|
||||
if use_websocket:
|
||||
try:
|
||||
await _dispatch_via_websocket(
|
||||
prompt_id = await _dispatch_via_websocket(
|
||||
worker_url,
|
||||
{
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
**target_payload,
|
||||
},
|
||||
client_id,
|
||||
)
|
||||
return
|
||||
if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
|
||||
raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
|
||||
return prompt_id if return_prompt_id else None
|
||||
except Exception as exc:
|
||||
worker_id = worker.get("id")
|
||||
if trace_execution_id:
|
||||
@@ -144,6 +176,108 @@ async def dispatch_worker_prompt(
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
if not return_prompt_id:
|
||||
return None
|
||||
try:
|
||||
response_data = await resp.json()
|
||||
except Exception as exc:
|
||||
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
|
||||
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
|
||||
if not isinstance(prompt_id, str) or not prompt_id:
|
||||
raise RuntimeError("Worker /prompt response omitted prompt_id")
|
||||
return prompt_id if return_prompt_id else None
|
||||
|
||||
|
||||
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *, require_value_results=False):
|
||||
"""Require protocol/node support and report which direct workers are idle.
|
||||
|
||||
This performs read-only GETs only; it never submits a prompt or contacts a
|
||||
serverless provider. Unsupported worker releases fail before master inputs run.
|
||||
"""
|
||||
if not workers:
|
||||
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
|
||||
session = await get_client_session()
|
||||
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
|
||||
semaphore = asyncio.Semaphore(8)
|
||||
|
||||
async def inspect(worker):
|
||||
worker_id = str(worker.get("id") or "unknown")
|
||||
check = "capability"
|
||||
try:
|
||||
async with semaphore:
|
||||
cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
|
||||
async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
|
||||
if response.status != 200:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} does not expose the capability endpoint "
|
||||
f"(HTTP {response.status}); upgrade that worker before running this feature"
|
||||
)
|
||||
capability = await response.json()
|
||||
if not isinstance(capability, dict):
|
||||
raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
|
||||
supported = capability.get("supported")
|
||||
protocol_version = capability.get("protocol_version")
|
||||
nodes = capability.get("node_ids")
|
||||
if (
|
||||
not isinstance(supported, bool)
|
||||
or type(protocol_version) is not int
|
||||
or not isinstance(nodes, list)
|
||||
or any(not isinstance(node_id, str) for node_id in nodes)
|
||||
):
|
||||
raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
|
||||
if (
|
||||
supported is not True
|
||||
or protocol_version != IMAGE_LIST_PROTOCOL_VERSION
|
||||
or not required_nodes.issubset(nodes)
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} is incompatible: requires protocol "
|
||||
f"v{IMAGE_LIST_PROTOCOL_VERSION} and nodes "
|
||||
f"{sorted(required_nodes)}; no work was submitted"
|
||||
)
|
||||
|
||||
if capability.get("execution_targets") is not True:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks explicit execution_targets support; "
|
||||
"upgrade that worker before running this feature; no work was submitted"
|
||||
)
|
||||
|
||||
if require_value_results and capability.get("value_results") is not True:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks explicit value_results support; "
|
||||
"upgrade that worker before using End overrides; no work was submitted"
|
||||
)
|
||||
|
||||
check = "queue"
|
||||
queue_url = build_worker_url(worker, "/queue")
|
||||
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
|
||||
if response.status != 200:
|
||||
raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
|
||||
queue = await response.json()
|
||||
if (
|
||||
not isinstance(queue, dict)
|
||||
or not isinstance(queue.get("queue_running"), list)
|
||||
or not isinstance(queue.get("queue_pending"), list)
|
||||
):
|
||||
raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
|
||||
idle = not queue["queue_running"] and not queue["queue_pending"]
|
||||
return worker_id, idle
|
||||
except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc:
|
||||
log(
|
||||
f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; "
|
||||
f"skipping: {exc}"
|
||||
)
|
||||
return worker_id, None
|
||||
|
||||
statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
|
||||
idle_workers = {
|
||||
worker_id: True
|
||||
for worker_id, idle in statuses
|
||||
if idle is True
|
||||
}
|
||||
if require_idle and not idle_workers:
|
||||
raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
|
||||
return idle_workers
|
||||
|
||||
|
||||
async def select_active_workers(
|
||||
|
||||
@@ -29,6 +29,11 @@ def _normalize_media_reference(value):
|
||||
return None
|
||||
cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
|
||||
if MEDIA_FILE_RE.search(cleaned):
|
||||
# Output/temp annotations select a different master-side media root.
|
||||
# Preserve them through discovery/loading and as distinct rewrite keys.
|
||||
annotation = re.search(r"\s*\[(output|temp)\]$", value)
|
||||
if annotation:
|
||||
return f"{cleaned} [{annotation.group(1)}]"
|
||||
return cleaned
|
||||
return None
|
||||
|
||||
@@ -171,6 +176,12 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
|
||||
"""Upload one media file to worker iff missing or hash-mismatched."""
|
||||
session = await get_client_session()
|
||||
normalized = filename.replace("\\", "/")
|
||||
annotation = re.search(r"\s*\[(output|temp)\]$", normalized)
|
||||
if annotation:
|
||||
# Worker uploads live under input/. Do not collide across master roots,
|
||||
# or overwrite a different reference with the same basename.
|
||||
basename = normalized[:annotation.start()].rstrip().rsplit("/", 1)[-1]
|
||||
normalized = f"distributed-media/{annotation.group(1)}/{file_hash}/{basename}"
|
||||
|
||||
check_url = build_worker_url(worker, "/distributed/check_file")
|
||||
try:
|
||||
|
||||
@@ -458,11 +458,55 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def _resolve_direct_workers_for_image_list(config, requested_ids):
|
||||
"""Resolve enabled workers by default, or exact explicitly selected direct workers."""
|
||||
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
|
||||
configured_workers = (config or {}).get("workers", [])
|
||||
configured_ids = {
|
||||
str(worker.get("id") or "").strip()
|
||||
for worker in configured_workers
|
||||
if str(worker.get("id") or "").strip()
|
||||
}
|
||||
if requested is not None:
|
||||
unknown_ids = requested - configured_ids
|
||||
if unknown_ids:
|
||||
raise RuntimeError(
|
||||
"Unknown explicitly selected image-list worker ID(s): "
|
||||
+ ", ".join(sorted(unknown_ids))
|
||||
)
|
||||
direct = []
|
||||
for worker in configured_workers:
|
||||
worker_id = str(worker.get("id") or "").strip()
|
||||
if not worker_id:
|
||||
continue
|
||||
if requested is None:
|
||||
if not worker.get("enabled", False):
|
||||
continue
|
||||
elif worker_id not in requested:
|
||||
continue
|
||||
if str(worker.get("type") or "local").strip().lower() == "serverless":
|
||||
continue
|
||||
try:
|
||||
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
|
||||
direct.append({
|
||||
"id": worker_id,
|
||||
"name": str(worker.get("name") or worker_id),
|
||||
"host": worker.get("host"),
|
||||
"port": port,
|
||||
"type": worker.get("type", "local"),
|
||||
})
|
||||
return direct
|
||||
|
||||
|
||||
def generate_job_id_map(prompt_index, prefix):
|
||||
"""Create stable per-node job IDs for distributed nodes."""
|
||||
job_map = {}
|
||||
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
|
||||
"UltimateSDUpscaleDistributed"
|
||||
distributed_nodes = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
+ prompt_index.nodes_for_class("DistributedImageListStart")
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
@@ -557,6 +601,22 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
|
||||
job_id_map, master_url, delegate_master, workflow_meta=None):
|
||||
"""Attach protocol context to Start without changing its visible V3 sockets."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
|
||||
inputs["image_list_master_participates"] = participant_id == "master" and not delegate_master
|
||||
inputs["image_list_master_callback"] = str(master_url or "")
|
||||
inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
|
||||
if workflow_meta is not None:
|
||||
inputs["image_list_workflow_meta"] = workflow_meta
|
||||
|
||||
|
||||
def apply_participant_overrides(
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
@@ -565,6 +625,7 @@ def apply_participant_overrides(
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
workflow_meta=None,
|
||||
):
|
||||
"""Return a prompt copy with hidden inputs configured for master/worker."""
|
||||
is_master = participant_id == "master"
|
||||
@@ -592,5 +653,15 @@ def apply_participant_overrides(
|
||||
master_url,
|
||||
enabled_json,
|
||||
)
|
||||
_override_image_list_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
workflow_meta=workflow_meta,
|
||||
)
|
||||
|
||||
return prompt_copy
|
||||
|
||||
+77
-14
@@ -35,7 +35,10 @@ from .orchestration.prompt_transform import (
|
||||
prepare_delegate_master_prompt,
|
||||
prune_prompt_for_worker,
|
||||
WORKER_SINK_CLASS,
|
||||
_resolve_direct_workers_for_image_list,
|
||||
)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
from ..image_list.execution_plan import prepare_image_list_execution
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
@@ -252,6 +255,8 @@ async def orchestrate_distributed_execution(
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
|
||||
config = load_config()
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
|
||||
(
|
||||
@@ -261,8 +266,46 @@ async def orchestrate_distributed_execution(
|
||||
media_sync_timeout_seconds,
|
||||
) = _resolve_orchestration_limits(config)
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart"
|
||||
for node in prompt_obj.values())
|
||||
target_kwargs = {}
|
||||
if image_list_mode:
|
||||
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
|
||||
target_kwargs["partial_execution_targets"] = outer_targets
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
if image_list_mode:
|
||||
# Validate before probing or queueing anything so malformed boundaries and
|
||||
# old worker releases fail before any upstream image is materialized.
|
||||
validate_image_list_regions(prompt_obj, explicit_targets=True)
|
||||
other_distributed = (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
)
|
||||
if other_distributed:
|
||||
raise ValueError(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
|
||||
"Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
configured_serverless = [
|
||||
str(worker.get("id")) for worker in config.get("workers", [])
|
||||
if worker.get("enabled", False)
|
||||
and str(worker.get("type") or "local").lower() == "serverless"
|
||||
and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
|
||||
]
|
||||
if configured_serverless:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Image List v1 explicitly excludes unsupported serverless pools: "
|
||||
+ ", ".join(configured_serverless),
|
||||
)
|
||||
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
|
||||
if not workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker; "
|
||||
"serverless pools are unsupported"
|
||||
)
|
||||
else:
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
@@ -278,26 +321,42 @@ async def orchestrate_distributed_execution(
|
||||
),
|
||||
)
|
||||
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
|
||||
if not workers and delegate_master:
|
||||
# Keep the legacy Collector fallback, but never override image-list
|
||||
# orchestrator-only mode merely because its direct workers are unavailable.
|
||||
if not image_list_mode and not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
|
||||
)
|
||||
delegate_master = False
|
||||
|
||||
active_workers, delegate_master = await select_active_workers(
|
||||
active_workers, selected_delegate_master = await select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
False if image_list_mode else delegate_master,
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
if image_list_mode:
|
||||
from .orchestration.dispatch import preflight_image_list_workers
|
||||
if not active_workers and delegate_master:
|
||||
raise RuntimeError(
|
||||
"No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
|
||||
)
|
||||
from ..image_list.runtime import END_CLASS, end_override_names
|
||||
requires_values = any(end_override_names(node.get("inputs", {}))
|
||||
for node in prompt_obj.values() if node.get("class_type") == END_CLASS)
|
||||
value_preflight = {"require_value_results": True} if requires_values else {}
|
||||
# Persist candidates, not this submission-time idle snapshot: a preceding
|
||||
# master prompt may still be using them. The image-list runtime rechecks
|
||||
# capability and availability before assigning the first item.
|
||||
if active_workers:
|
||||
await preflight_image_list_workers(active_workers, require_idle=False,
|
||||
**value_preflight)
|
||||
else:
|
||||
delegate_master = selected_delegate_master
|
||||
|
||||
if load_balance_requested:
|
||||
if load_balance_requested and not image_list_mode:
|
||||
candidate_workers = list(active_workers)
|
||||
if not delegate_master:
|
||||
# Include master in load balancing only when master participation is enabled.
|
||||
@@ -360,6 +419,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
@@ -377,9 +437,10 @@ async def orchestrate_distributed_execution(
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
workflow_meta=workflow_meta,
|
||||
)
|
||||
|
||||
if delegate_master:
|
||||
if delegate_master and not image_list_mode:
|
||||
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
@@ -400,7 +461,7 @@ async def orchestrate_distributed_execution(
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
|
||||
)
|
||||
worker_payloads = []
|
||||
if active_workers:
|
||||
if active_workers and not image_list_mode:
|
||||
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
|
||||
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
|
||||
worker_payloads = await asyncio.gather(
|
||||
@@ -431,8 +492,9 @@ async def orchestrate_distributed_execution(
|
||||
manager.ensure_capacity(len(serverless_jobs))
|
||||
|
||||
# Do not leave collector queues behind when offline compilation/media validation fails.
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
if not image_list_mode:
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
if regular_payloads:
|
||||
await asyncio.gather(
|
||||
@@ -455,6 +517,7 @@ async def orchestrate_distributed_execution(
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
except Exception:
|
||||
if serverless_jobs:
|
||||
@@ -483,4 +546,4 @@ async def orchestrate_distributed_execution(
|
||||
execution_trace_id,
|
||||
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
|
||||
)
|
||||
return prompt_id, prompt_number, len(worker_payloads), node_errors
|
||||
return prompt_id, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
|
||||
|
||||
@@ -61,6 +61,13 @@ async def worker_ws_endpoint(request):
|
||||
})
|
||||
continue
|
||||
|
||||
if not isinstance(data, dict):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack", "request_id": None, "ok": False,
|
||||
"error": "Expected a websocket message object.",
|
||||
})
|
||||
continue
|
||||
|
||||
if data.get("type") != "dispatch_prompt":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
@@ -81,10 +88,16 @@ async def worker_ws_endpoint(request):
|
||||
continue
|
||||
|
||||
try:
|
||||
target_kwargs = {}
|
||||
if "partial_execution_targets" in data:
|
||||
from ..utils.async_helpers import validate_execution_targets
|
||||
target_kwargs["partial_execution_targets"] = validate_execution_targets(
|
||||
prompt, data["partial_execution_targets"])
|
||||
prompt_id = await queue_prompt_payload(
|
||||
prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=data.get("client_id"),
|
||||
**target_kwargs,
|
||||
)
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
# Dynamic image-list queue
|
||||
|
||||
Process a list of images across the master and direct local or remote ComfyUI workers without assigning fixed chunks to each participant.
|
||||
|
||||
The master evaluates the input list once. By default, the master and each selected idle worker receive one image job at a time; when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. The master restores input order before running downstream nodes.
|
||||
|
||||
The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. The master still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
|
||||
|
||||
## Wiring
|
||||
|
||||
```text
|
||||
Image list or IMAGE batch
|
||||
|
|
||||
Distributed Image List Start -------- context ---------+
|
||||
| image |
|
||||
v |
|
||||
Per-image processing |
|
||||
| IMAGE |
|
||||
v v
|
||||
Distributed Image List End
|
||||
| image_list
|
||||
v
|
||||
Preview / Save / next stage
|
||||
```
|
||||
|
||||
1. Install the same feature-capable version of Distributed on the master and every selected worker. Workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation.
|
||||
2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. The master must be reachable from remote workers at its configured callback address.
|
||||
3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
|
||||
4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
|
||||
5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
|
||||
6. Put preview/save nodes **after End**, then queue normally. API callers should submit through `/distributed/queue`, not directly through `/prompt`.
|
||||
|
||||
ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; for each assigned item, Start provides one frame and its aligned values. End returns ordered lists, not a concatenated tensor batch, so differently sized images can be preserved.
|
||||
|
||||
## Inputs and outputs
|
||||
|
||||
Start accepts:
|
||||
|
||||
- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames; each frame becomes one job.
|
||||
- Optional `mask`.
|
||||
- Optional `value_1`, `value_2`, … inputs. Only `value_1` is shown at first; connecting it adds `value_2`, and so on, up to `value_10`. Each accepts an IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN.
|
||||
|
||||
`mask` and every value follow one rule:
|
||||
|
||||
- **1 value** is shared by every image.
|
||||
- **One value per image** pairs them in order.
|
||||
- **Any other count** is an error. Values are never padded or repeated silently.
|
||||
|
||||
Only `images` is split into frames. A batch connected to a value input stays a single value, so a 2-image batch can travel with each item (for example, two reference images per job). To pass a different batch to each image, connect a list with one batch per image.
|
||||
|
||||
Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`, then one `value_N` output per connected `value_N` input, with the same type. Unused value outputs are hidden. Seeds and strengths are ordinary values: connect a single number to share it, a list for one per image, or derive it from `item_index`.
|
||||
|
||||
Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
|
||||
|
||||
End accepts `images`, optional `masks` / `texts`, and `context`. Its first three outputs remain `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional mask/text outputs contain `None` entries.
|
||||
|
||||
### Keep values through End without extra wires
|
||||
|
||||
Start's existing **context** connection now also lets End recover each item's original `value_1` … `value_10`. You do not need to wire those values individually into End. Context remains a lightweight execution reference, not a general-purpose data pipe.
|
||||
|
||||
End appends a **values** output (`DISTRIBUTED_IMAGE_LIST_VALUES`): one value bundle per image, in the same source order. Connect it to **Distributed Image List Get Value**, then choose `value_name` to extract the field you need. The getter preserves its type and ComfyUI maps it across the bundles, producing a list aligned with End's images. Missing fields fail clearly instead of producing empty placeholders.
|
||||
|
||||
For example, to preserve filenames:
|
||||
|
||||
```text
|
||||
filenames → Start.value_1
|
||||
Start.context ─────────────→ End.context
|
||||
Start.image → processing → End.images
|
||||
├─ image_list → Save.images
|
||||
└─ values → Get Value [value_1] → Save.filename_prefix
|
||||
```
|
||||
|
||||
**New or changed per-item values:** connect the calculated result to End's optional `value_N` input. It replaces the original field with that number, or adds it if Start did not supply it. Every connected override must produce exactly one supported value for each item. A tensor batch remains one value, just as it does at Start. Image, mask, string, integer, float and boolean values are supported; arbitrary objects/models/latents and non-finite numbers are not.
|
||||
|
||||
End shows override inputs progressively. Use **Add value override** to expose a later slot without connecting the earlier ones. Disconnecting an override never renumbers the other End fields: `value_3` always means `value_3`. Unlike Start's native Autogrow inputs, End overrides must keep these identities stable. Existing image/mask/text/context sockets and output indices are unchanged. If the frontend extension is unavailable, Python still declares all optional override inputs and the values output.
|
||||
|
||||
Original values are retained on the master; workers only return explicitly connected overrides, avoiding a redundant return trip for unchanged reference images or other values. Overrides use the existing exact value transport and size limits. Selected workers must advertise `value_results: true` when overrides are requested, checked before execution and again at dispatch. Original-value passthrough alone does not require that additional capability.
|
||||
|
||||
Values leaving the region must come from **End**, not directly from Start. To chain regions, extract an End value and connect it to the next Start's corresponding input. Values remain available downstream after End cleans up the execution context; they do not depend on a later lookup of expired job state.
|
||||
|
||||
Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. The master still validates the prompt: it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
|
||||
|
||||
Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes.
|
||||
|
||||
End's optional `texts` is a connection-only socket. For workflows saved before this socket became connection-only, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End.
|
||||
|
||||
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Their worker copies use separate content-addressed paths under `input/distributed-media/`; each worker's rewritten prompt is reused across items without reusing item tokens.
|
||||
|
||||
## Native cache reuse
|
||||
|
||||
Repeated items can reuse **ComfyUI's native processing cache** on the same participant. Start uses a deterministic identity of the exact image, optional mask, all supplied value slots and source index; tensor dtype, shape and bytes are preserved. Fresh job IDs, assignment tokens and callback addresses live at End, so they no longer invalidate otherwise unchanged upstream processing. End still validates the current assignment and delivers results in source order on every run.
|
||||
|
||||
- Processing settings, seeds, links and node fingerprints still participate in ComfyUI's normal signatures. Changed content invalidates reuse; nodes that declare themselves noncacheable still run.
|
||||
- The master's private executor retains one native LRU output cache between successful slots, capped at **128 entries when retained**. An active graph may temporarily exceed that target; oversized generations are not retained. Failed or cancelled local slots do not seed reuse, and `--cache-none` disables the local output cache. Direct workers use their own configured native caches.
|
||||
- Cache hits are best-effort, not guaranteed: eviction, restarts, changed ordering/index, changed settings and assignment to a different worker can cause misses. There is no worker affinity, disk cache or separate completed-result store.
|
||||
- Reuse retains tensor/model references and can increase RAM/VRAM use; the entry cap is **not a byte limit**. Exact identity calculation also adds CPU hashing and, for GPU-resident inputs, a CPU copy.
|
||||
- Remote items are still uploaded. Each upload returns a worker-computed identity; if a compatible older worker does not return a valid identity, that assignment uses the legacy graph without stable-input reuse. A worker downgrade between items also falls back safely. Assignment expiry and current-item validation apply even when processing hits the cache.
|
||||
- Keep side-effect nodes outside the region when they must run on every submission. A cached processing node follows ordinary ComfyUI semantics and does not repeat its save/preview side effects; End's delivery and downstream saving still run.
|
||||
|
||||
## Scope and failure behavior
|
||||
|
||||
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
|
||||
- Queue submission checks selected direct workers for reachability and feature compatibility, but does not freeze the worker list to that moment's idle snapshot. A selected worker busy with an earlier prompt remains a candidate for the queued run. Immediately before a region starts assigning images, its saved candidates are checked again: workers still busy or transiently unreachable are excluded; healthy idle candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers excluded by the initial reachability probe are not retained, and newly selected workers are not added to an already-queued run. Workers becoming free after the region starts are not added mid-region.
|
||||
- **Orchestrator Only** requires a selected, reachable, feature-compatible direct worker at submission, but can queue while all selected workers are busy. At region start it requires at least one idle compatible direct worker and fails if none are available; it never falls back to master processing. The queue response's `worker_count` counts saved remote candidates, not guaranteed execution-time participants.
|
||||
- Omitted worker IDs select enabled direct workers. Explicit IDs select the named configured direct workers, including disabled ones; an empty list stays empty, unknown IDs fail, and serverless workers are excluded. Local workers use the master's loopback callback; remote workers use its configured public address.
|
||||
- Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no direct workers are available. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
|
||||
- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
|
||||
- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
|
||||
- Queue preparation works on an execution copy, leaving the saved canvas unchanged. Unused per-item preview/save branches are pruned by dependency reachability, not node names. Required computational utilities marked `OUTPUT_NODE` stay in the graph; explicit `partial_execution_targets` prevent them from also becoming independent outer-master roots. End and legitimate post-End or unrelated outputs are retained. Bypass/rejoin branches still fail boundary validation rather than being silently removed.
|
||||
- Every selected worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation.
|
||||
- **Target selection is not a sandbox.** A required utility still runs its own save, preview or other side effects when it executes for an item. Keep intentional final saving after End.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and are not subject to those HTTP limits.
|
||||
- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
|
||||
- Worker errors fail the region. The worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue.
|
||||
- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are ambiguous, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
|
||||
- State is in memory. Restarting the master does not resume an interrupted region. Abandoned master contexts and expired worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
|
||||
|
||||
### Time and size limits
|
||||
|
||||
`settings.image_list_job_timeout_seconds` defaults to `0` (no arbitrary total-run cap). A positive finite number opts into an overall deadline with owned-work cleanup; negative, non-finite or malformed settings fail explicitly. There is no fixed one-hour region deadline. The setting is separate from the existing worker liveness timeout.
|
||||
|
||||
Uploads, media preparation and submission acknowledgements are independently bounded. WebSocket acknowledgement waiting includes connect/send/receive and is not extended by heartbeats or unrelated messages. A user interrupt allows up to one second for a late acknowledgement; cleanup requests use five-second bounds each.
|
||||
|
||||
**4K float32 RGB outputs cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/mask/value combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits.
|
||||
|
||||
Use trusted workers and trusted workflows, as with the existing ComfyUI API. Capability tokens protect assignment/result matching; they are not a replacement for network access control.
|
||||
|
||||
## Verification
|
||||
|
||||
The generic acceptance helper starts two isolated loopback workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. No managed ComfyUI service or model is started.
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
|
||||
bash scripts/test-web.sh
|
||||
```
|
||||
|
||||
Without `COMFYUI_SOURCE_ROOT`, real-framework tests are explicitly skipped. This coverage is not a substitute for checking your model workflow and connectivity on the actual GPU worker machines.
|
||||
@@ -3,7 +3,7 @@
|
||||
A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
|
||||
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
|
||||
normal ComfyUI API workflow per requested image and collects the returned images.
|
||||
The endpoint workers **do not need ComfyUI-Distributed installed**.
|
||||
The endpoint workers do not need ComfyUI-Distributed installed.
|
||||
|
||||
## Setup
|
||||
|
||||
@@ -93,7 +93,7 @@ incomplete serverless batch.
|
||||
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
|
||||
- Interrupting a collector or using the panel cancellation action requests
|
||||
cancellation only for jobs submitted by this master. It never purges the endpoint.
|
||||
- A network failure during submission is **not retried automatically**: the job
|
||||
- A network failure during submission is not retried automatically: the job
|
||||
may already be billable. Check the Runpod job list before queuing again.
|
||||
|
||||
## Persistence and recovery
|
||||
@@ -107,7 +107,7 @@ available on Windows.
|
||||
|
||||
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
|
||||
not image data, workflow payloads, or API keys. A restarted master cannot resume
|
||||
an interrupted collector. It does **not** automatically resubmit jobs. Retained
|
||||
an interrupted collector. It does not automatically resubmit jobs. Retained
|
||||
IDs are available through `/distributed/serverless/status/<worker_id>`; the worker's
|
||||
cancel action can request cancellation of known pre-restart jobs. If no job ID
|
||||
was received, check Runpod manually. Cancellation is best-effort and may race
|
||||
@@ -127,7 +127,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
|
||||
every unresolved record with its local `record_id`, provider `job_id` (if
|
||||
received), and state. These records remain visible if endpoint health fails.
|
||||
2. Request cancellation from the panel, or send
|
||||
`POST /distributed/serverless/reconcile` with the body below to **read** the
|
||||
`POST /distributed/serverless/reconcile` with the body below to read the
|
||||
retained provider IDs. Only a confirmed terminal state releases protection.
|
||||
Network errors, expired IDs, or ambiguous responses do not.
|
||||
|
||||
@@ -143,7 +143,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
|
||||
{"worker_id": "YOUR_WORKER_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
|
||||
```
|
||||
|
||||
This records `RESOLVED`: an explicit operator acknowledgment, **not** a
|
||||
This records `RESOLVED`: an explicit operator acknowledgment, not a
|
||||
provider-confirmed cancellation. It cannot resolve a live local task and does
|
||||
not stop remote compute. It releases the record's reservation and may let
|
||||
already-waiting jobs proceed. Never automate this acknowledgment as a retry
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Dynamic per-image image-list execution internals."""
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Prepare image-list execution without changing the saved/API prompt or registry."""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import sys
|
||||
|
||||
from .runtime import (
|
||||
END_CLASS,
|
||||
END_RESULT_INPUTS,
|
||||
_ancestors_for_inputs,
|
||||
_downstream_until,
|
||||
_image_list_pairs,
|
||||
_is_output_node,
|
||||
_prompt_links,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
def _reachable(roots: set[str], edges: dict[str, set[str]]) -> set[str]:
|
||||
visited = set()
|
||||
pending = list(roots)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
pending.extend(edges[node_id] - visited)
|
||||
return visited
|
||||
|
||||
|
||||
def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]:
|
||||
"""Return an execution copy and its outer-master output targets.
|
||||
|
||||
Disposable per-item side branches are removed before boundary validation.
|
||||
Required output-node utilities remain normal dependencies, but must not be
|
||||
independent outer roots. Callers must use these targets (and End-only targets
|
||||
for item execution); this does not suppress side effects inside required nodes.
|
||||
"""
|
||||
pairs = _image_list_pairs(prompt)
|
||||
prepared = copy.deepcopy(prompt)
|
||||
upstream = {node_id: set() for node_id in prepared}
|
||||
downstream = {node_id: set() for node_id in prepared}
|
||||
for node_id, node in prepared.items():
|
||||
for _, source, _ in _prompt_links(node, prepared):
|
||||
upstream[node_id].add(source)
|
||||
downstream[source].add(node_id)
|
||||
|
||||
# Never erase another region or its dependencies to make a bad graph pass.
|
||||
boundaries = {node_id for pair in pairs for node_id in pair}
|
||||
required = _reachable(boundaries, upstream)
|
||||
discarded = set()
|
||||
for start_id, end_id in pairs:
|
||||
region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, END_RESULT_INPUTS)
|
||||
outside = _downstream_until(prepared, start_id, end_id) - region_nodes
|
||||
# Preserve bypass joins and their whole chain into post-End computation.
|
||||
# They must fail boundary validation, not disappear with debug previews.
|
||||
post_end = _reachable({end_id}, downstream)
|
||||
protected = required | _reachable(post_end, upstream)
|
||||
discarded.update(outside - protected)
|
||||
for node_id in discarded:
|
||||
del prepared[node_id]
|
||||
|
||||
regions = validate_image_list_regions(prepared, explicit_targets=True)
|
||||
per_item = set().union(*(region["branch_ids"] for region in regions))
|
||||
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
|
||||
targets = [
|
||||
node_id for node_id, node in prepared.items()
|
||||
if node_id not in per_item and isinstance(node, dict)
|
||||
and (node.get("class_type") == END_CLASS or _is_output_node(node, registry))
|
||||
]
|
||||
return prepared, targets
|
||||
@@ -0,0 +1,174 @@
|
||||
"""Server-owned background cleanup for image-list runtime state."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, cast
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
DEFAULT_SWEEP_INTERVAL_SECONDS = 30.0
|
||||
DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
|
||||
_LIFECYCLE_MARKER = object()
|
||||
|
||||
|
||||
def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
|
||||
"""Prove outer roots from the running queue, not a forgeable prompt marker.
|
||||
|
||||
Plain /prompt remains strict when ComfyUI selected all outputs, including a
|
||||
flagged per-item dependency. No saved workflow field grants this exemption.
|
||||
"""
|
||||
if not prompt_id:
|
||||
return False
|
||||
from .execution_plan import prepare_image_list_execution
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
if prepared != prompt or not targets:
|
||||
return False
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
for name in ("get_current_queue_volatile", "get_current_queue"):
|
||||
reader = getattr(queue, name, None)
|
||||
if not callable(reader):
|
||||
continue
|
||||
try:
|
||||
running, _pending = reader()
|
||||
for entry in running:
|
||||
if (isinstance(entry, (tuple, list)) and len(entry) > 4
|
||||
and entry[1] == prompt_id and entry[2] == prompt):
|
||||
roots = entry[4]
|
||||
return (isinstance(roots, (list, tuple))
|
||||
and all(isinstance(root, str) for root in roots)
|
||||
and len(roots) == len(targets) and set(roots) == set(targets))
|
||||
return False
|
||||
except Exception:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _queue_prompt_ids(prompt_server: Any) -> set[str]:
|
||||
"""Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
readers = (
|
||||
getattr(queue, "get_current_queue_volatile", None),
|
||||
getattr(queue, "get_current_queue", None),
|
||||
)
|
||||
errors: list[Exception] = []
|
||||
for reader in readers:
|
||||
if not callable(reader):
|
||||
continue
|
||||
try:
|
||||
snapshot = reader()
|
||||
if not isinstance(snapshot, (tuple, list)) or len(snapshot) < 2:
|
||||
raise TypeError("prompt queue snapshot must contain running and pending entries")
|
||||
running, pending = snapshot[0], snapshot[1]
|
||||
prompt_ids: set[str] = set()
|
||||
for entries in (running, pending):
|
||||
if not isinstance(entries, Iterable):
|
||||
raise TypeError("prompt queue entries must be iterable")
|
||||
for entry in entries:
|
||||
if isinstance(entry, (tuple, list)) and len(entry) > 1:
|
||||
prompt_id = entry[1]
|
||||
if isinstance(prompt_id, str):
|
||||
prompt_ids.add(prompt_id)
|
||||
return prompt_ids
|
||||
except Exception as exc:
|
||||
errors.append(exc)
|
||||
|
||||
if errors:
|
||||
raise RuntimeError("unable to read the prompt queue") from errors[-1]
|
||||
raise RuntimeError("prompt queue has no supported queue reader")
|
||||
|
||||
|
||||
def _sweep_orphan_contexts(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
|
||||
queue_read_ok = True
|
||||
try:
|
||||
queued_prompt_ids = _queue_prompt_ids(prompt_server)
|
||||
except Exception:
|
||||
# A failed queue read must not make a submitted context look orphaned.
|
||||
LOGGER.warning("Could not read prompt queue during image-list lifecycle sweep", exc_info=True)
|
||||
queued_prompt_ids = set()
|
||||
queue_read_ok = False
|
||||
|
||||
now = time.monotonic()
|
||||
with runtime._LOCK:
|
||||
for context_id, state in list(runtime._MASTER_JOBS.items()):
|
||||
if getattr(state, "active", False):
|
||||
continue
|
||||
prompt_id = getattr(state, "prompt_id", None)
|
||||
if prompt_id is not None:
|
||||
if not queue_read_ok or prompt_id in queued_prompt_ids:
|
||||
continue
|
||||
created_at = getattr(state, "created_at", None)
|
||||
expired = (
|
||||
isinstance(created_at, (int, float))
|
||||
and not isinstance(created_at, bool)
|
||||
and math.isfinite(created_at)
|
||||
and now - created_at >= orphan_grace_seconds
|
||||
)
|
||||
if expired and runtime._MASTER_JOBS.get(context_id) is state:
|
||||
runtime.remove_master_job(context_id)
|
||||
|
||||
|
||||
def _sweep_once(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
|
||||
try:
|
||||
runtime._sweep_worker_items()
|
||||
except Exception:
|
||||
LOGGER.exception("Could not sweep expired image-list worker staging")
|
||||
_sweep_orphan_contexts(prompt_server, runtime, orphan_grace_seconds)
|
||||
|
||||
|
||||
async def _sweep_loop(
|
||||
prompt_server: Any,
|
||||
runtime: Any,
|
||||
interval_seconds: float,
|
||||
orphan_grace_seconds: float,
|
||||
) -> None:
|
||||
while True:
|
||||
try:
|
||||
_sweep_once(prompt_server, runtime, orphan_grace_seconds)
|
||||
except Exception:
|
||||
# Keep the server-owned task alive if one pass encounters unexpected state.
|
||||
LOGGER.exception("Image-list lifecycle sweep failed")
|
||||
await asyncio.sleep(interval_seconds)
|
||||
|
||||
|
||||
def install_lifecycle(
|
||||
prompt_server: Any,
|
||||
runtime: Any,
|
||||
*,
|
||||
interval_seconds: float = DEFAULT_SWEEP_INTERVAL_SECONDS,
|
||||
orphan_grace_seconds: float = DEFAULT_ORPHAN_GRACE_SECONDS,
|
||||
) -> None:
|
||||
"""Register one cancellable periodic sweep in the prompt server's app context."""
|
||||
interval_seconds = float(interval_seconds)
|
||||
orphan_grace_seconds = float(orphan_grace_seconds)
|
||||
if not math.isfinite(interval_seconds) or interval_seconds <= 0:
|
||||
raise ValueError("interval_seconds must be a positive finite value")
|
||||
if not math.isfinite(orphan_grace_seconds) or orphan_grace_seconds < 0:
|
||||
raise ValueError("orphan_grace_seconds must be a non-negative finite value")
|
||||
|
||||
app = getattr(prompt_server, "app", None)
|
||||
cleanup_ctx = getattr(app, "cleanup_ctx", None)
|
||||
if cleanup_ctx is None:
|
||||
raise TypeError("prompt_server must expose an app with cleanup_ctx")
|
||||
if any(getattr(factory, "_image_list_lifecycle_marker", None) is _LIFECYCLE_MARKER
|
||||
for factory in cleanup_ctx):
|
||||
return
|
||||
|
||||
async def lifecycle_context(_app):
|
||||
task = asyncio.create_task(
|
||||
_sweep_loop(prompt_server, runtime, interval_seconds, orphan_grace_seconds),
|
||||
name="image-list-lifecycle-sweeper",
|
||||
)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
cast(Any, lifecycle_context)._image_list_lifecycle_marker = _LIFECYCLE_MARKER
|
||||
cleanup_ctx.append(lifecycle_context)
|
||||
@@ -0,0 +1,223 @@
|
||||
"""A queue-free Comfy graph slot, only while the synchronous outer End is blocked.
|
||||
|
||||
Do NOT use PromptExecutor.execute[_async] here. Its prompt lifecycle clears the
|
||||
process interrupt, resets progress handlers, changes RAM-cache callbacks and
|
||||
starts/ends a model tracker. Those belong to the suspended parent executor.
|
||||
Instead use core's ExecutionList/execute (including lazy/async/expanded nodes)
|
||||
with private provider-free caches and a private server facade. Model outputs are
|
||||
registered with the existing parent tracker, whose lifetime includes this work.
|
||||
|
||||
Core progress and GraphBuilder defaults are process-global. Borrow them only
|
||||
inside the owned execution thread, restore them in finally, and join that thread
|
||||
before End returns. Cancellation never clears or raises the global interrupt:
|
||||
it stops locally at node/progress boundaries. A non-cooperative synchronous node
|
||||
must finish before cleanup can return; abandoning its thread is not safe.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
from .runtime import ImageListCancelledError
|
||||
|
||||
|
||||
_EXECUTION_LOCK = threading.Lock()
|
||||
_CACHE_LOCK = threading.Lock()
|
||||
# One process-local native output cache, not a second completed-result cache.
|
||||
# Core's LRU can temporarily exceed its target for the active graph; oversized
|
||||
# generations are discarded rather than retained after an item/run finishes.
|
||||
NATIVE_CACHE_MAX_ENTRIES = 128
|
||||
_RETAINED_OUTPUTS = None
|
||||
|
||||
|
||||
def _take_outputs(execution):
|
||||
global _RETAINED_OUTPUTS
|
||||
with _CACHE_LOCK:
|
||||
outputs, _RETAINED_OUTPUTS = _RETAINED_OUTPUTS, None
|
||||
if getattr(execution.args, "cache_none", False):
|
||||
return execution.NullCache()
|
||||
return outputs if outputs is not None else execution.LRUCache(
|
||||
execution.CacheKeySetInputSignature, max_size=NATIVE_CACHE_MAX_ENTRIES)
|
||||
|
||||
|
||||
def _retain_outputs(outputs):
|
||||
global _RETAINED_OUTPUTS
|
||||
# Only the bounded native LRU is eligible. NullCache and test doubles are not.
|
||||
if getattr(outputs, "max_size", None) != NATIVE_CACHE_MAX_ENTRIES:
|
||||
return
|
||||
outputs.clean_unused()
|
||||
if len(outputs.cache) > NATIVE_CACHE_MAX_ENTRIES:
|
||||
return
|
||||
with _CACHE_LOCK:
|
||||
_RETAINED_OUTPUTS = outputs
|
||||
|
||||
|
||||
class LocalExecutionSlot:
|
||||
def __init__(self, *, model_callback=None):
|
||||
self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-list-master")
|
||||
self._stop = threading.Event()
|
||||
self._future = None
|
||||
self._caches = None
|
||||
self._failed = False
|
||||
self._model_callback = model_callback
|
||||
self._closed = False
|
||||
self._loop = None
|
||||
self._task = None
|
||||
self._state_lock = threading.Lock()
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
|
||||
if self._closed or self._stop.is_set():
|
||||
raise ImageListCancelledError("Local image-list slot has stopped")
|
||||
if self._future is not None and not self._future.done():
|
||||
raise RuntimeError("Local image-list slot already has an in-flight item")
|
||||
self._future = self._pool.submit(self._run, prompt, prompt_id, end_id, workflow_meta)
|
||||
return self._future
|
||||
|
||||
def request_stop(self):
|
||||
self._stop.set()
|
||||
# Wakes a graph waiting on an async node as well as stopping synchronous
|
||||
# nodes at their next progress/between-node checkpoint.
|
||||
with self._state_lock:
|
||||
if self._loop is not None and self._task is not None:
|
||||
self._loop.call_soon_threadsafe(self._task.cancel)
|
||||
|
||||
async def close(self):
|
||||
if self._closed:
|
||||
return
|
||||
# A stop requested before close indicates cancellation, not normal drain.
|
||||
cancelled = self._stop.is_set()
|
||||
self.request_stop()
|
||||
try:
|
||||
if self._future is not None:
|
||||
# Do not cancel this Future: cancelling an asyncio wrapper cannot
|
||||
# kill its running thread, and falsely makes cleanup look done.
|
||||
waiter = asyncio.wrap_future(self._future)
|
||||
while not waiter.done():
|
||||
try:
|
||||
await asyncio.shield(waiter)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except Exception:
|
||||
break # The scheduler propagates execution errors.
|
||||
if waiter.done() and not waiter.cancelled():
|
||||
waiter.exception()
|
||||
finally:
|
||||
self._pool.shutdown(wait=True, cancel_futures=True)
|
||||
if not cancelled and not self._failed and self._caches is not None:
|
||||
_retain_outputs(self._caches.outputs)
|
||||
self._caches = None
|
||||
self._closed = True
|
||||
|
||||
def _run(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
# The normal Comfy prompt worker is single-threaded and blocked in End.
|
||||
# Fail closed if another local slot is nevertheless executing, rather
|
||||
# than interleave writes to core's process-global execution state.
|
||||
if not _EXECUTION_LOCK.acquire(blocking=False):
|
||||
raise RuntimeError("Another local image-list executor is already active")
|
||||
try:
|
||||
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
|
||||
except asyncio.CancelledError as exc:
|
||||
self._failed = True
|
||||
raise ImageListCancelledError("Local image-list execution cancelled") from exc
|
||||
except BaseException:
|
||||
self._failed = True
|
||||
raise
|
||||
finally:
|
||||
with self._state_lock:
|
||||
self._loop = self._task = None
|
||||
_EXECUTION_LOCK.release()
|
||||
|
||||
def _check_stopped(self):
|
||||
import comfy.model_management
|
||||
if self._stop.is_set() or comfy.model_management.processing_interrupted():
|
||||
raise ImageListCancelledError("Local image-list execution interrupted")
|
||||
|
||||
async def _execute(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
import execution
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
from comfy_execution import progress
|
||||
from comfy_execution.graph_utils import GraphBuilder
|
||||
|
||||
self._check_stopped()
|
||||
with self._state_lock:
|
||||
self._loop, self._task = asyncio.get_running_loop(), asyncio.current_task()
|
||||
dynamic_prompt = execution.DynamicPrompt(prompt)
|
||||
private_progress = progress.ProgressRegistry(prompt_id, dynamic_prompt)
|
||||
parent_progress = progress.global_progress_registry
|
||||
parent_hook = comfy.utils.PROGRESS_BAR_HOOK
|
||||
prefix_keys = ("_default_prefix_root", "_default_prefix_call_index", "_default_prefix_graph_index")
|
||||
parent_prefix = {key: getattr(GraphBuilder, key) for key in prefix_keys}
|
||||
progress.global_progress_registry = private_progress
|
||||
|
||||
def progress_hook(value, total, preview=None, prompt_id=None, node_id=None):
|
||||
self._check_stopped()
|
||||
# Never route a child progress/preview event through the parent's
|
||||
# client or reset its handlers. Keep the private progress accurate.
|
||||
if node_id is not None:
|
||||
private_progress.update_progress(node_id, value, total, preview)
|
||||
|
||||
comfy.utils.PROGRESS_BAR_HOOK = progress_hook
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
if self._caches is None:
|
||||
outputs = _take_outputs(execution)
|
||||
objects = execution.HierarchicalCache(execution.CacheKeySetID)
|
||||
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
|
||||
caches = self._caches
|
||||
changed = execution.IsChangedCache(prompt_id, dynamic_prompt, caches.outputs)
|
||||
for cache in caches.all:
|
||||
await cache.set_prompt(dynamic_prompt, prompt.keys(), changed)
|
||||
cache.clean_unused()
|
||||
execution_list = execution.ExecutionList(dynamic_prompt, caches.outputs, self._model_callback)
|
||||
execution_list.add_node(end_id)
|
||||
server = SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
|
||||
assets = execution.NoAssets(execution.args)
|
||||
extra = {"extra_pnginfo": {"workflow": workflow_meta}} if workflow_meta is not None else {}
|
||||
executed, pending_subgraphs, pending_async, ui = set(), {}, {}, {}
|
||||
while not execution_list.is_empty():
|
||||
self._check_stopped()
|
||||
node_id, error, exc = await execution_list.stage_node_execution()
|
||||
if error is not None:
|
||||
raise RuntimeError(f"Local image-list graph failed: {error}") from exc
|
||||
self._check_stopped()
|
||||
result, error, exc = await execution.execute(
|
||||
server, dynamic_prompt, caches, node_id, extra, executed, prompt_id,
|
||||
execution_list, pending_subgraphs, pending_async, ui, assets)
|
||||
if result == execution.ExecutionResult.FAILURE:
|
||||
if isinstance(exc, (comfy.model_management.InterruptProcessingException,
|
||||
ImageListCancelledError)):
|
||||
raise ImageListCancelledError("Local image-list execution interrupted") from exc
|
||||
raise RuntimeError(f"Local image-list node {node_id} failed: {error}") from exc
|
||||
if result == execution.ExecutionResult.PENDING:
|
||||
execution_list.unstage_node_execution()
|
||||
else:
|
||||
execution_list.complete_node_execution()
|
||||
self._check_stopped()
|
||||
finally:
|
||||
# Core execute creates async-node and unblock tasks. This is a private
|
||||
# event loop, so every other task is owned here. Drain/cancel them
|
||||
# before restoring globals, including when the graph failed early.
|
||||
with self._state_lock:
|
||||
self._task = None # Stop requests must not interrupt this drain.
|
||||
try:
|
||||
while True:
|
||||
pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
||||
if not pending:
|
||||
break
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
drain = asyncio.gather(*pending, return_exceptions=True)
|
||||
while not drain.done():
|
||||
try:
|
||||
await asyncio.shield(drain)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
finally:
|
||||
progress.global_progress_registry = parent_progress
|
||||
comfy.utils.PROGRESS_BAR_HOOK = parent_hook
|
||||
for key, value in parent_prefix.items():
|
||||
setattr(GraphBuilder, key, value)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,376 @@
|
||||
"""Public V3 Start/End nodes for dynamic per-image image-list execution."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..image_list.runtime import (
|
||||
MAX_VALUE_SLOTS,
|
||||
VALUE_NAMES,
|
||||
VALUES_INPUT,
|
||||
end_override_names,
|
||||
create_master_job,
|
||||
normalize_image_list_inputs,
|
||||
normalize_worker_result,
|
||||
run_image_list_job,
|
||||
encode_result,
|
||||
encode_request_payload,
|
||||
unwrap_single,
|
||||
is_valid_item_identity,
|
||||
)
|
||||
|
||||
|
||||
def _stable_worker_binding(prompt, start_id, identity):
|
||||
"""Resolve transient state only from the exact Start -> End context wire."""
|
||||
if not is_valid_item_identity(identity):
|
||||
raise ValueError("Invalid worker item identity")
|
||||
start = prompt.get(start_id) if isinstance(prompt, dict) else None
|
||||
if (not isinstance(start, dict) or start.get("class_type") != "DistributedImageListStart"
|
||||
or start.get("inputs", {}).get("worker_item_identity") != identity):
|
||||
raise RuntimeError("Worker item identity is missing its paired Start")
|
||||
candidates = []
|
||||
for end_id, node in prompt.items():
|
||||
if not isinstance(node, dict) or node.get("class_type") != "DistributedImageListEnd":
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
link = inputs.get("context")
|
||||
if isinstance(link, (list, tuple)) and list(link) == [start_id, 3]:
|
||||
candidates.append((end_id, inputs))
|
||||
if len(candidates) != 1:
|
||||
raise RuntimeError("Worker Start must have exactly one paired End context wire")
|
||||
end_id, inputs = candidates[0]
|
||||
if inputs.get("worker_item_identity") != identity:
|
||||
raise ValueError("Paired End worker item identity does not match Start")
|
||||
for name in ("worker_result_token", "worker_context_id"):
|
||||
if not isinstance(inputs.get(name), str) or not inputs[name]:
|
||||
raise RuntimeError(f"Worker paired End is missing {name}")
|
||||
index = inputs.get("worker_item_index")
|
||||
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
|
||||
raise ValueError("Worker paired End item index is invalid")
|
||||
return end_id, inputs
|
||||
|
||||
|
||||
# Types the transport can carry exactly. Anything else (models, latents, ...)
|
||||
# must be loaded inside the region or shared as a model dependency instead.
|
||||
_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
|
||||
assert len(VALUE_NAMES) == MAX_VALUE_SLOTS
|
||||
|
||||
|
||||
class DistributedImageListStart(io.ComfyNode):
|
||||
"""Materialize a batch once, then expose ordered per-item typed values."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DistributedImageListStart",
|
||||
display_name="Distributed Image List Start",
|
||||
category="image/distributed",
|
||||
description=("Starts one per-image processing region. Each image becomes one job. "
|
||||
"mask and value_N inputs take 1 value (shared by every image) or exactly "
|
||||
"one value per image, paired in order."),
|
||||
inputs=[
|
||||
io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one job."),
|
||||
io.Mask.Input("mask", optional=True,
|
||||
tooltip="1 mask shared by every image, or one mask per image."),
|
||||
io.Autogrow.Input(VALUES_INPUT, optional=True, template=io.Autogrow.TemplateNames(
|
||||
io.MultiType.Input("value", types=_VALUE_TYPES,
|
||||
tooltip=("Any image, mask, text, number or boolean. 1 value is shared by "
|
||||
"every image; one value per image pairs them in order. "
|
||||
"A batch stays one value.")),
|
||||
names=list(VALUE_NAMES), min=0)),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="image", is_output_list=True),
|
||||
io.Mask.Output(display_name="mask", is_output_list=True),
|
||||
io.Int.Output(display_name="item_index", is_output_list=True),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Output(display_name="context", is_output_list=True),
|
||||
# value_N outputs mirror value_N inputs. The web extension hides
|
||||
# unused ones and shows the connected input's type.
|
||||
*[io.AnyType.Output(id=name, display_name=name, is_output_list=True) for name in VALUE_NAMES],
|
||||
],
|
||||
is_input_list=True,
|
||||
accept_all_inputs=True,
|
||||
not_idempotent=True,
|
||||
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, mask=None, values=None,
|
||||
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
|
||||
image_list_master_callback="", image_list_execution_id="",
|
||||
image_list_workflow_meta=None, image_list_master_participates=None,
|
||||
image_list_local_execution=False, worker_item_identity=None):
|
||||
worker_item_token = unwrap_single(worker_item_token)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
worker_item_identity = unwrap_single(worker_item_identity)
|
||||
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
|
||||
if worker_item_identity is not None:
|
||||
prompt = getattr(cls.hidden, "prompt", None)
|
||||
start_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
_, delivery = _stable_worker_binding(prompt, start_id, worker_item_identity)
|
||||
# Only cache misses execute Start. Never put this token/context in
|
||||
# Start's literal inputs or returned context: both would poison reuse.
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
item = consume_worker_item(delivery["worker_result_token"],
|
||||
item_identity=worker_item_identity,
|
||||
item_index=delivery["worker_item_index"])
|
||||
item_values = item.get("values") or {}
|
||||
return io.NodeOutput(
|
||||
[item["image"]], [item.get("mask")], [item["index"]], [worker_item_identity],
|
||||
*[[item_values.get(name)] for name in VALUE_NAMES],
|
||||
)
|
||||
if worker_item_token:
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
item = consume_worker_item(worker_item_token)
|
||||
else:
|
||||
item = run_async_in_server_loop(_consume_worker_item(worker_item_token), timeout=10)
|
||||
context_id = worker_context_id
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
raise RuntimeError("Worker image-list prompt is missing its paired context")
|
||||
item_values = item.get("values") or {}
|
||||
return io.NodeOutput(
|
||||
[item["image"]], [item.get("mask")], [item["index"]], [context_id],
|
||||
*[[item_values.get(name)] for name in VALUE_NAMES],
|
||||
)
|
||||
|
||||
items = normalize_image_list_inputs(images, mask, values)
|
||||
prompt = getattr(cls.hidden, "prompt", None)
|
||||
start_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
if not prompt or not start_id:
|
||||
raise RuntimeError("Distributed Image List Start needs a submitted ComfyUI prompt context")
|
||||
try:
|
||||
worker_ids = json.loads(str(unwrap_single(image_list_enabled_worker_ids or "[]")))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise RuntimeError("Invalid preflighted image-list worker list") from exc
|
||||
if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
|
||||
raise RuntimeError("Invalid preflighted image-list worker list")
|
||||
from comfy_execution.utils import get_executing_context
|
||||
from ..image_list.lifecycle import has_safe_execution_targets
|
||||
import server
|
||||
executing = get_executing_context()
|
||||
explicit_targets = has_safe_execution_targets(
|
||||
getattr(server.PromptServer, "instance", None),
|
||||
executing.prompt_id if executing is not None else None, prompt)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
region = next((row for row in regions if row["start_id"] == start_id), None)
|
||||
if region is None:
|
||||
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
|
||||
master_participates = unwrap_single(image_list_master_participates)
|
||||
if master_participates is None:
|
||||
from ..utils.config import load_config
|
||||
master_participates = not bool((load_config().get("settings", {}) or {}).get("master_delegate_only", False))
|
||||
if not isinstance(master_participates, bool):
|
||||
raise ValueError("image_list_master_participates must be a boolean")
|
||||
context_id = create_master_job(
|
||||
start_id, region["end_id"], prompt, items, worker_ids,
|
||||
str(unwrap_single(image_list_master_callback or "")),
|
||||
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
|
||||
prompt_id=executing.prompt_id if executing is not None else None,
|
||||
master_participates=master_participates,
|
||||
explicit_targets=explicit_targets,
|
||||
)
|
||||
return io.NodeOutput(
|
||||
[item["image"] for item in items],
|
||||
[item.get("mask") for item in items],
|
||||
[item["index"] for item in items],
|
||||
[context_id for _ in items],
|
||||
*[[item["values"].get(name) for item in items] for name in VALUE_NAMES],
|
||||
)
|
||||
|
||||
|
||||
async def _consume_worker_item(token):
|
||||
# Item state is local to each worker process; source nodes consume it exactly once.
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
return consume_worker_item(token)
|
||||
|
||||
|
||||
class DistributedImageListEnd(io.ComfyNode):
|
||||
"""Collect branch results in source order without evaluating the branch on master."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DistributedImageListEnd",
|
||||
display_name="Distributed Image List End",
|
||||
category="image/distributed",
|
||||
description=("Closes a paired image-list region and returns ordered image, mask, text and values lists. "
|
||||
"Start values are retained; connected value_N inputs replace or add that key per item."),
|
||||
inputs=[
|
||||
io.Image.Input("images", lazy=True),
|
||||
io.Mask.Input("masks", optional=True, lazy=True),
|
||||
io.String.Input("texts", optional=True, lazy=True, force_input=True),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
|
||||
*[io.MultiType.Input(name, types=_VALUE_TYPES, optional=True, lazy=True,
|
||||
tooltip="Replace or add this value for each item; exactly one value per item.")
|
||||
for name in VALUE_NAMES],
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="image_list", is_output_list=True),
|
||||
io.Mask.Output(display_name="masks", is_output_list=True),
|
||||
io.String.Output(display_name="texts", is_output_list=True),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Output(
|
||||
id="values", display_name="values", is_output_list=True),
|
||||
],
|
||||
is_input_list=True,
|
||||
is_output_node=True,
|
||||
accept_all_inputs=True,
|
||||
not_idempotent=True,
|
||||
hidden=[io.Hidden.prompt, io.Hidden.unique_id, io.Hidden.execution_list],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(cls, images=None, masks=None, texts=None, context=None,
|
||||
worker_result_token="", **kwargs):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
if worker_result_token:
|
||||
# Decide which optional sockets are connected from the actual worker prompt;
|
||||
# absent optional inputs must not be demanded from Comfy's lazy executor.
|
||||
prompt = getattr(cls.hidden, "prompt", {}) or {}
|
||||
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
|
||||
if isinstance(node, dict)
|
||||
and node.get("class_type") == "DistributedImageListEnd"
|
||||
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
|
||||
required = []
|
||||
for name, value in (("images", images), ("masks", masks), ("texts", texts),
|
||||
*((name, kwargs.get(name)) for name in end_override_names(end_inputs))):
|
||||
connected = name == "images" or name in end_inputs
|
||||
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
|
||||
if connected and missing:
|
||||
required.append(name)
|
||||
return required
|
||||
# Master only needs Start's context. Never request branch outputs here.
|
||||
if context is None or (isinstance(context, (tuple, list)) and all(item is None for item in context)):
|
||||
return ["context"]
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, masks=None, texts=None, context=None,
|
||||
worker_result_token="", worker_item_index=-1, master_callback_url="",
|
||||
image_list_local_execution=False, worker_context_id="",
|
||||
worker_item_identity=None, **values):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
master_callback_url = unwrap_single(master_callback_url)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
worker_item_identity = unwrap_single(worker_item_identity)
|
||||
if worker_item_identity is not None and not worker_result_token:
|
||||
raise RuntimeError("Worker item identity is missing its delivery token")
|
||||
if worker_result_token:
|
||||
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
|
||||
if isinstance(item, str) and item]
|
||||
if not context_values or len(set(context_values)) != 1:
|
||||
raise RuntimeError("Worker Image List End is missing its paired context")
|
||||
prompt = getattr(cls.hidden, "prompt", {}) or {}
|
||||
context_id = context_values[0]
|
||||
if worker_item_identity is not None:
|
||||
end_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
node = prompt.get(end_id, {})
|
||||
link = node.get("inputs", {}).get("context")
|
||||
if (node.get("class_type") != "DistributedImageListEnd"
|
||||
or not isinstance(link, (list, tuple)) or len(link) != 2):
|
||||
raise RuntimeError("Worker End is missing its paired Start context wire")
|
||||
paired_id, end_inputs = _stable_worker_binding(prompt, link[0], worker_item_identity)
|
||||
if (paired_id != end_id or context_values != [worker_item_identity]
|
||||
or end_inputs["worker_context_id"] != worker_context_id
|
||||
or end_inputs["worker_result_token"] != worker_result_token
|
||||
or end_inputs["worker_item_index"] != unwrap_single(worker_item_index)
|
||||
or end_inputs.get("master_callback_url", "") != master_callback_url):
|
||||
raise RuntimeError("Worker End identity/context does not match its current delivery")
|
||||
from ..image_list.runtime import finish_worker_item
|
||||
# On a miss this consumes Start's receipt; on a hit it validates
|
||||
# and drops the current staged tensors. Both paths enforce TTL.
|
||||
finish_worker_item(worker_result_token, worker_item_identity,
|
||||
unwrap_single(worker_item_index))
|
||||
context_id = worker_context_id
|
||||
else:
|
||||
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
|
||||
if isinstance(node, dict) and node.get("class_type") == "DistributedImageListEnd"
|
||||
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
|
||||
expected = end_override_names(end_inputs)
|
||||
# Missing connected lazy values must fail instead of becoming an
|
||||
# implicit fallback to the original Start value.
|
||||
supplied = {name: values.get(name) for name in expected}
|
||||
supplied.update(values)
|
||||
result = normalize_worker_result(images, masks, texts, supplied)
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import set_master_result
|
||||
set_master_result(context_id, int(unwrap_single(worker_item_index)),
|
||||
str(worker_result_token), result)
|
||||
else:
|
||||
payload = encode_result(result)
|
||||
run_async_in_server_loop(
|
||||
submit_worker_result(str(master_callback_url), str(worker_result_token),
|
||||
context_id, int(unwrap_single(worker_item_index)), payload), timeout=60,
|
||||
)
|
||||
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")],
|
||||
[result.get("values", {})])
|
||||
|
||||
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
|
||||
if isinstance(item, str) and item]
|
||||
if not context_values or len(set(context_values)) != 1:
|
||||
raise RuntimeError("Distributed Image List End requires the paired Start context wire")
|
||||
context_id = context_values[0]
|
||||
from ..image_list.runtime import ImageListCancelledError
|
||||
try:
|
||||
# The runtime owns the configurable deadline and bounded remote cleanup.
|
||||
execution_list = getattr(cls.hidden, "execution_list", None)
|
||||
# Local model outputs belong to this still-running outer prompt.
|
||||
# Add them to its existing tracker; never start/end another tracker
|
||||
# that could unpin a model still needed by the parent.
|
||||
model_callback = getattr(execution_list, "output_link_callback", None)
|
||||
result_lists = run_async_in_server_loop(run_image_list_job(
|
||||
context_id, model_callback=model_callback, include_values=True))
|
||||
except ImageListCancelledError:
|
||||
import comfy.model_management
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
images_out, masks_out, texts_out, values_out = result_lists
|
||||
if len({len(output) for output in result_lists}) != 1:
|
||||
raise RuntimeError("Image-list collector received mismatched image, mask, text or values result counts")
|
||||
return io.NodeOutput(images_out, masks_out, texts_out, values_out)
|
||||
|
||||
|
||||
class DistributedImageListGetValue(io.ComfyNode):
|
||||
"""Ordinary scalar execution lets core map over End's ordered values list."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DistributedImageListGetValue",
|
||||
display_name="Distributed Image List Get Value",
|
||||
category="image/distributed",
|
||||
inputs=[io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Input("values"),
|
||||
io.Combo.Input("value_name", options=list(VALUE_NAMES), default="value_1")],
|
||||
outputs=[io.AnyType.Output(id="value", display_name="value")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, values, value_name="value_1"):
|
||||
if value_name not in VALUE_NAMES:
|
||||
raise ValueError(f"Unknown image-list value slot {value_name!r}")
|
||||
if not isinstance(values, dict):
|
||||
raise TypeError("Image List Get Value expects one values mapping from Image List End")
|
||||
if value_name not in values:
|
||||
raise ValueError(f"Image List Get Value: {value_name} is absent from this item's values")
|
||||
return io.NodeOutput(values[value_name])
|
||||
|
||||
|
||||
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
|
||||
if not callback_url.startswith(("http://", "https://")):
|
||||
raise RuntimeError("Worker image-list callback URL is missing or invalid")
|
||||
import aiohttp
|
||||
from ..utils.network import get_client_session
|
||||
session = await get_client_session()
|
||||
payload = {"context_id": context_id, "token": item_token,
|
||||
"item_index": item_index, "result": result}
|
||||
async with session.post(
|
||||
f"{callback_url.rstrip('/')}/distributed/image-list/result",
|
||||
data=encode_request_payload(payload),
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as response:
|
||||
if response.status >= 400:
|
||||
body = (await response.text())[:500]
|
||||
raise RuntimeError(f"Master rejected image-list result (HTTP {response.status}): {body}")
|
||||
+2
-1
@@ -10,6 +10,7 @@ from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from ..utils import transfer_codec as _transfer_codec
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@@ -211,4 +212,4 @@ class DistributedWorkerSink(io.ComfyNode):
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
|
||||
DistributedWorkerSink]
|
||||
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedImageListGetValue]
|
||||
|
||||
@@ -24,6 +24,9 @@ def initialize():
|
||||
|
||||
ensure_distributed_state(server.PromptServer.instance)
|
||||
ensure_tile_jobs_initialized()
|
||||
from ..image_list import runtime as image_list_runtime
|
||||
from ..image_list.lifecycle import install_lifecycle
|
||||
install_lifecycle(server.PromptServer.instance, image_list_runtime)
|
||||
|
||||
if not os.environ.get('COMFYUI_IS_WORKER'):
|
||||
atexit.register(sync_cleanup)
|
||||
|
||||
@@ -243,6 +243,55 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
|
||||
self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
|
||||
|
||||
|
||||
class AnnotatedMediaUploadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_output_temp_and_input_uploads_do_not_collide(self):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
class Response:
|
||||
status = 200
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
async def json(self):
|
||||
return self.data
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
class Form:
|
||||
def __init__(self):
|
||||
self.fields = {}
|
||||
def add_field(self, key, value, **kwargs):
|
||||
self.fields[key] = value
|
||||
if key == 'image':
|
||||
self.fields['name'] = kwargs['filename']
|
||||
|
||||
checked = []
|
||||
def post(url, *, json=None, data=None, **kwargs):
|
||||
if url.endswith('/check_file'):
|
||||
checked.append(json['filename'])
|
||||
return Response({'exists': False})
|
||||
self.assertEqual(data.fields['type'], 'input')
|
||||
self.assertEqual(data.fields['image'], b'real file bytes')
|
||||
return Response({'name': data.fields['name'], 'subfolder': data.fields['subfolder']})
|
||||
|
||||
refs = ['shared/ref.png', 'shared/ref.png [output]', 'shared/ref.png [temp]']
|
||||
prompt = {str(i): {'inputs': {'image': ref}} for i, ref in enumerate(refs)}
|
||||
self.assertEqual(ms._find_media_references(prompt), refs)
|
||||
with patch.object(ms, 'get_client_session', AsyncMock(return_value=types.SimpleNamespace(post=post))), \
|
||||
patch.object(ms, 'aiohttp', types.SimpleNamespace(FormData=Form, ClientTimeout=lambda **kw: kw)):
|
||||
paths = [await ms._upload_media_to_worker({}, ref, b'real file bytes', 'fixture-hash', 'image/png')
|
||||
for ref in refs]
|
||||
expected = ['shared/ref.png', 'distributed-media/output/fixture-hash/ref.png',
|
||||
'distributed-media/temp/fixture-hash/ref.png']
|
||||
self.assertEqual(checked, expected)
|
||||
self.assertEqual(paths, [(True, path) for path in expected])
|
||||
ms._rewrite_prompt_media_inputs(prompt, dict(zip(refs, expected)))
|
||||
self.assertEqual([node['inputs']['image'] for node in prompt.values()], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
@@ -271,6 +271,74 @@ def _load_worker_routes_module():
|
||||
worker_routes = _load_worker_routes_module()
|
||||
|
||||
|
||||
class WebsocketExecutionTargetTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
from test_async_helpers import _load_async_helpers_module
|
||||
self.modules = patch.dict(sys.modules)
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
self.helper, self.server = _load_async_helpers_module()
|
||||
sys.modules['dist_api_worker_testpkg.utils.async_helpers'] = self.helper
|
||||
self.prompt = {'9': {'class_type': 'Output'}, '10': {'class_type': 'Output'}}
|
||||
|
||||
async def receive(self, payload):
|
||||
import json
|
||||
sent = []
|
||||
class Socket:
|
||||
async def prepare(self, request):
|
||||
pass
|
||||
async def send_json(self, value):
|
||||
sent.append(value)
|
||||
def __aiter__(self):
|
||||
async def messages():
|
||||
yield types.SimpleNamespace(type=worker_routes.aiohttp.WSMsgType.TEXT,
|
||||
data=json.dumps(payload))
|
||||
return messages()
|
||||
with patch.object(worker_routes, 'web', types.SimpleNamespace(WebSocketResponse=lambda **kw: Socket())), \
|
||||
patch.object(worker_routes, 'queue_prompt_payload', self.helper.queue_prompt_payload), \
|
||||
patch.object(worker_routes, 'PromptValidationError', self.helper.PromptValidationError):
|
||||
await worker_routes.worker_ws_endpoint(_FakeRequest())
|
||||
return sent[0]
|
||||
|
||||
def payload(self, targets):
|
||||
return {'type': 'dispatch_prompt', 'request_id': 'req', 'prompt': self.prompt,
|
||||
'partial_execution_targets': targets}
|
||||
|
||||
async def test_explicit_targets_reach_queue_and_acknowledgement(self):
|
||||
ack = await self.receive(self.payload(['9']))
|
||||
self.assertTrue(ack['ok'], ack)
|
||||
self.assertEqual(ack['request_id'], 'req')
|
||||
self.assertEqual(self.server.prompt_queue.items[0][4], ['9'])
|
||||
|
||||
async def test_malformed_targets_and_messages_are_rejected(self):
|
||||
for targets in (None, [], '9', [True], [9], ['absent'], ['9', '9']):
|
||||
with self.subTest(targets=targets):
|
||||
ack = await self.receive(self.payload(targets))
|
||||
self.assertFalse(ack['ok'])
|
||||
self.assertIn('partial_execution_targets', ack['error'])
|
||||
self.assertFalse((await self.receive([]))['ok'])
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_validation_cannot_silently_queue_a_partial_subset(self):
|
||||
ack = await self.receive(self.payload(['9', '10']))
|
||||
self.assertFalse(ack['ok'])
|
||||
self.assertEqual(ack['validation_error']['type'], 'partial_execution_targets_failed_validation')
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_missing_target_field_keeps_legacy_queue_call(self):
|
||||
from unittest.mock import AsyncMock
|
||||
queued = AsyncMock(return_value='legacy')
|
||||
original = self.helper.queue_prompt_payload
|
||||
self.helper.queue_prompt_payload = queued
|
||||
try:
|
||||
payload = self.payload(['9'])
|
||||
payload.pop('partial_execution_targets')
|
||||
self.assertTrue((await self.receive(payload))['ok'])
|
||||
self.assertNotIn('partial_execution_targets', queued.call_args.kwargs)
|
||||
finally:
|
||||
self.helper.queue_prompt_payload = original
|
||||
|
||||
|
||||
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_network_info_reports_actual_master_and_available_ports(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Offline CPU proof against ComfyUI's real signatures and PromptExecutor.
|
||||
|
||||
Exercises one native worker executor and separate private master-slot lifetimes,
|
||||
including their shared stable Start / fresh End contract.
|
||||
"""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def main(core):
|
||||
sys.path.insert(0, str(Path(core).resolve()))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True # Must precede all model-management imports.
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
import torch
|
||||
from comfy_execution.caching import CacheKeySetInputSignature
|
||||
|
||||
name = '_native_cache_identity_probe'
|
||||
for suffix in ('', '.image_list', '.nodes', '.utils'):
|
||||
package = types.ModuleType(name + suffix)
|
||||
package.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = package
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('offline native proof attempted a server-loop bridge')
|
||||
|
||||
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
|
||||
rt = importlib.import_module(name + '.image_list.runtime')
|
||||
boundary = importlib.import_module(name + '.nodes.image_list')
|
||||
calls = []
|
||||
|
||||
class Operation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'mask': ('MASK',),
|
||||
'text': ('STRING', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
|
||||
def run(self, image, mask, text):
|
||||
calls.append(text)
|
||||
return image + 1, mask + 1, text + ':done'
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
|
||||
DistributedImageListEnd=boundary.DistributedImageListEnd,
|
||||
Operation=Operation)
|
||||
prompt = {
|
||||
'start': {'class_type': rt.START_CLASS, 'inputs': {'image_list_execution_id': 'outer-A'}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {
|
||||
'image': ['start', 0], 'mask': ['start', 1], 'text': ['start', 4]}},
|
||||
'end': {'class_type': rt.END_CLASS, 'inputs': {
|
||||
'images': ['op', 0], 'masks': ['op', 1], 'texts': ['op', 2],
|
||||
'context': ['start', 3], 'value_3': ['start', 2]}},
|
||||
}
|
||||
data = rt.normalize_image_list_inputs(torch.zeros(1, 2, 3, 3, dtype=torch.float64),
|
||||
torch.ones(1, 2, 3, dtype=torch.bfloat16),
|
||||
{'value_1': 'original', 'value_2': 3})[0]
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
|
||||
def prepare(data):
|
||||
token = rt.create_worker_task_token()
|
||||
context = rt.create_master_job('start', 'end', prompt, [data], ['worker'], '')
|
||||
rt.register_assignment(context, data['index'], 'worker', token)
|
||||
identity = rt.store_worker_item(token, rt.encode_item(data))
|
||||
result = rt.build_worker_prompt(prompt, region, token, context, data['index'], '', item_identity=identity)
|
||||
result['start']['inputs']['image_list_local_execution'] = True
|
||||
result['end']['inputs']['image_list_local_execution'] = True
|
||||
return token, context, result
|
||||
|
||||
async def signature(graph, executor):
|
||||
dynamic = execution.DynamicPrompt(graph)
|
||||
changed = execution.IsChangedCache('signature-only', dynamic, executor.caches.outputs)
|
||||
keys = CacheKeySetInputSignature(dynamic, graph.keys(), changed)
|
||||
await keys.add_keys(graph.keys())
|
||||
return {node: keys.get_data_key(node) for node in graph}
|
||||
|
||||
async def probe():
|
||||
server = types.SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
|
||||
executor = execution.PromptExecutor(server, cache_type=execution.CacheType.LRU,
|
||||
cache_args={'lru': 64, 'ram': 0, 'ram_inactive': 0})
|
||||
ta, ca, a = prepare(data)
|
||||
prompt['start']['inputs']['image_list_execution_id'] = 'outer-B'
|
||||
tb, cb, b = prepare(data)
|
||||
assert ta != tb and ca != cb
|
||||
ka, kb = await signature(a, executor), await signature(b, executor)
|
||||
assert ka['start'] == kb['start'] and ka['op'] == kb['op']
|
||||
assert ka['end'] != kb['end']
|
||||
assert a['start']['inputs'] == b['start']['inputs']
|
||||
# Exercise actual V3 list mapping, hidden.prompt/unique_id, lazy inputs,
|
||||
# native node signatures and cache hits, not a reimplementation of them.
|
||||
for number, (token, context, graph) in enumerate(((ta, ca, a), (tb, cb, b))):
|
||||
valid = await execution.validate_prompt(f'native-{number}', graph, ['end'])
|
||||
assert valid[0], valid
|
||||
await executor.execute_async(graph, f'native-{number}', {}, ['end'])
|
||||
assert executor.success, executor.status_messages
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
result = rt.get_master_job(context).results[0]
|
||||
assert torch.equal(result['image'], data['image'] + 1)
|
||||
assert torch.equal(result['mask'], data['mask'] + 1)
|
||||
assert result['text'] == 'original:done'
|
||||
assert result['values'] == {'value_3': 0}
|
||||
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
|
||||
if number:
|
||||
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
|
||||
assert calls == ['original'], calls
|
||||
for failure in ('expired', 'mismatched'):
|
||||
token, context, graph = prepare(data)
|
||||
expires, staged = rt._WORKER_ITEMS[token]
|
||||
if failure == 'expired':
|
||||
rt._WORKER_ITEMS[token] = (0, staged)
|
||||
else:
|
||||
staged['image'] = staged['image'] + 9
|
||||
await executor.execute_async(graph, f'rejected-{failure}', {}, ['end'])
|
||||
assert not executor.success, failure
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
|
||||
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
|
||||
assert calls == ['original'], calls
|
||||
for part in ('image', 'mask', 'value', 'index'):
|
||||
changed = copy.deepcopy(data)
|
||||
if part in ('image', 'mask'):
|
||||
changed[part].reshape(-1)[0] += 1
|
||||
elif part == 'value':
|
||||
changed['values']['value_2'] = 4 # Even an unconnected slot is content.
|
||||
else:
|
||||
changed['index'] = 1
|
||||
token, context, graph = prepare(changed)
|
||||
keys = await signature(graph, executor)
|
||||
assert keys['start'] != ka['start'] and keys['op'] != ka['op'], part
|
||||
await executor.execute_async(graph, f'changed-{part}', {}, ['end'])
|
||||
assert executor.success, executor.status_messages
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
assert rt.get_master_job(context).received_indexes == {changed['index']}
|
||||
assert len(calls) == 5, calls
|
||||
# Combine the real Start/End identity path with cache transfer between
|
||||
# private master slots, not just independent signature/retention probes.
|
||||
local = importlib.import_module(name + '.image_list.local_executor')
|
||||
before = len(calls)
|
||||
for number in range(2):
|
||||
token, context, graph = prepare(data)
|
||||
slot = local.LocalExecutionSlot()
|
||||
try:
|
||||
await asyncio.wrap_future(slot.submit(graph, f'local-{number}', 'end'))
|
||||
result = rt.get_master_job(context).results[0]
|
||||
assert torch.equal(result['image'], data['image'] + 1)
|
||||
assert result['values'] == {'value_3': 0}
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
finally:
|
||||
await slot.close()
|
||||
rt.remove_master_job(context)
|
||||
assert local._RETAINED_OUTPUTS is not None
|
||||
assert len(calls) == before + 1, calls
|
||||
assert not rt._WORKER_ITEMS
|
||||
assert not torch.cuda.is_initialized()
|
||||
print('NATIVE_CACHE_IDENTITY_OK stable Start/op keys; fresh End key; '
|
||||
'1 operation for 2 same-data jobs; 4 changed-content misses; '
|
||||
'expired/mismatched cached inputs rejected; current delivery and cleanup')
|
||||
|
||||
asyncio.run(probe())
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(sys.argv[1])
|
||||
@@ -0,0 +1,688 @@
|
||||
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
|
||||
|
||||
Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers.
|
||||
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
|
||||
No managed ComfyUI service, user workflow, model, or remote worker is touched.
|
||||
"""
|
||||
import asyncio
|
||||
from collections import Counter
|
||||
import copy
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
||||
WORKER = '--worker' in sys.argv
|
||||
DELAY = float(sys.argv[3]) if WORKER else 0
|
||||
sys.path.insert(0, str(COMFY_ROOT))
|
||||
os.environ['COMFYUI_IS_WORKER'] = '1'
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
from app.assets.manager import default_asset_manager
|
||||
import comfy.model_management
|
||||
import execution
|
||||
import nodes
|
||||
import server
|
||||
import torch
|
||||
from aiohttp import web
|
||||
import aiohttp
|
||||
|
||||
torch.set_num_threads(1)
|
||||
SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
|
||||
UTILITY_CALLS, ITEM_EXECUTIONS, VALUE_CAPTURES = [], [], []
|
||||
ALLOW_MASTER = False
|
||||
COUNT = 7
|
||||
|
||||
|
||||
def source_rows():
|
||||
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
|
||||
'reference_image': torch.full((2, 2, 3, 3), (i + 1) / 16),
|
||||
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
|
||||
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
|
||||
for i in range(COUNT)]
|
||||
|
||||
|
||||
class TestImageSource:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT')
|
||||
OUTPUT_IS_LIST = (True,) * 6
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
assert not WORKER, 'upstream image source leaked into worker prompt'
|
||||
SOURCE_CALLS.append(True)
|
||||
rows = source_rows()
|
||||
return tuple([row[key] for row in rows] for key in
|
||||
('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
|
||||
|
||||
|
||||
class TestSharedFlag:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('BOOLEAN',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
assert not WORKER, 'shared value source leaked into worker prompt'
|
||||
return (True,)
|
||||
|
||||
|
||||
class TestCachedModel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('MODEL',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
MODEL_CALLS.append(True)
|
||||
return (object(),)
|
||||
|
||||
|
||||
def assert_item_execution(prompt, index):
|
||||
"""Prove the operation belongs to an item, not merely a participating host."""
|
||||
from comfy_execution.utils import get_executing_context
|
||||
context = get_executing_context()
|
||||
assert context is not None
|
||||
assert prompt['4']['inputs'].get('worker_result_token'), 'per-item node ran in outer graph'
|
||||
assert (prompt['2']['inputs'].get('worker_item_identity')
|
||||
or prompt['2']['inputs'].get('worker_item_token')), 'per-item source binding missing'
|
||||
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
|
||||
assert len(running) == 1, running
|
||||
entry = running[0]
|
||||
if WORKER:
|
||||
assert entry[1] == context.prompt_id and entry[2] == prompt
|
||||
assert entry[4] == ['4'], entry[4]
|
||||
else:
|
||||
assert ALLOW_MASTER, 'per-item node ran on delegate-only master'
|
||||
assert entry[1] != context.prompt_id, 'per-item node ran in outer executor'
|
||||
assert set(entry[4]) == {'4', '5'}, entry[4]
|
||||
return {'index': index, 'prompt_id': context.prompt_id,
|
||||
'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]}
|
||||
|
||||
|
||||
class TestRequiredOutputUtility:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})},
|
||||
'hidden': {'prompt': 'PROMPT'}}
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, index, prompt):
|
||||
assert_item_execution(prompt, index)
|
||||
UTILITY_CALLS.append(index)
|
||||
return (image,)
|
||||
|
||||
|
||||
class TestDiscardedPreview:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',)}}
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, image):
|
||||
raise AssertionError('discarded per-item output preview executed')
|
||||
|
||||
|
||||
class TestImageOperation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'reference_image': ('IMAGE',), 'mask': ('MASK',),
|
||||
'text': ('STRING', {'forceInput': True}),
|
||||
'strength': ('FLOAT', {'forceInput': True}),
|
||||
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
|
||||
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)},
|
||||
'hidden': {'prompt': 'PROMPT'}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt):
|
||||
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
|
||||
ITEM_EXECUTIONS.append(assert_item_execution(prompt, index))
|
||||
BRANCH_CALLS.append(index)
|
||||
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
|
||||
until = time.monotonic() + DELAY
|
||||
while time.monotonic() < until:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
time.sleep(0.01)
|
||||
assert seed == 100 + index
|
||||
assert flag is True
|
||||
assert text == f'item-{index}'
|
||||
assert strength == (index + 1) / 16
|
||||
assert torch.all(reference_image == strength)
|
||||
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
|
||||
|
||||
|
||||
class TestValueOverride:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'text': ('STRING', {'forceInput': True}),
|
||||
'index': ('INT', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('STRING', 'INT')
|
||||
FUNCTION = 'run'
|
||||
def run(self, text, index):
|
||||
assert WORKER or ALLOW_MASTER, 'override branch executed on delegate-only master'
|
||||
return (text + ':override', 200 + index)
|
||||
|
||||
|
||||
class TestCapture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})},
|
||||
'optional': {'values': ('DISTRIBUTED_IMAGE_LIST_VALUES',),
|
||||
'picked': ('STRING', {'forceInput': True})}}
|
||||
INPUT_IS_LIST = True
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, images, masks, texts, values=None, picked=None):
|
||||
assert not WORKER, 'post-End capture leaked into a worker prompt'
|
||||
CAPTURED.append((images, masks, texts))
|
||||
VALUE_CAPTURES.append((values, picked))
|
||||
return ()
|
||||
|
||||
|
||||
async def setup():
|
||||
assets = default_asset_manager()
|
||||
ps = server.PromptServer(asyncio.get_running_loop(), assets)
|
||||
assert await nodes.load_custom_node(str(ROOT))
|
||||
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
||||
boundary = importlib.import_module(module.__name__ + '.nodes.image_list')
|
||||
encode_result = boundary.encode_result
|
||||
def checked_encode(result):
|
||||
encoded = encode_result(result)
|
||||
if 'values' in encoded:
|
||||
assert set(encoded['values']) == {f'value_{n}' for n in (2, 3, 4, 5, 7, 8, 9, 10)}
|
||||
return encoded
|
||||
boundary.encode_result = checked_encode
|
||||
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
|
||||
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
|
||||
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
|
||||
TestRequiredOutputUtility=TestRequiredOutputUtility,
|
||||
TestDiscardedPreview=TestDiscardedPreview, TestValueOverride=TestValueOverride)
|
||||
# Match the CORS-enabled direct-worker launch contract for canvas probes.
|
||||
app = web.Application(middlewares=[server.create_cors_middleware('*')])
|
||||
# Register the actual feature routes, not test replacements.
|
||||
app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')])
|
||||
return ps, assets, module, app
|
||||
|
||||
|
||||
async def listen(app):
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, '127.0.0.1', 0)
|
||||
await site.start()
|
||||
return runner, site._server.sockets[0].getsockname()[1]
|
||||
|
||||
|
||||
def finish_queue_item(ps, item_id, executor):
|
||||
ps.prompt_queue.task_done(item_id, executor.history_result,
|
||||
status=execution.PromptQueue.ExecutionStatus(
|
||||
status_str='success' if executor.success else 'error',
|
||||
completed=executor.success, messages=executor.status_messages),
|
||||
process_item=lambda item: item[:5] + item[6:])
|
||||
|
||||
|
||||
async def worker_main():
|
||||
ps, assets, module, app = await setup()
|
||||
wake = asyncio.Event()
|
||||
interrupted_ids, submissions = [], []
|
||||
core_handlers = {(route.method, route.path): route.handler for route in ps.routes}
|
||||
async def prompt(request):
|
||||
data = await request.json()
|
||||
# Use core's actual validation/enqueue/response schema, including
|
||||
# partial_execution_targets; only observe it and wake our CPU consumer.
|
||||
response = await core_handlers['POST', '/prompt'](request)
|
||||
if response.status == 200:
|
||||
accepted = json.loads(response.body)
|
||||
_, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
|
||||
submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'),
|
||||
'queued': entry[4], 'node_ids': sorted(entry[2]),
|
||||
'end_ids': [key for key, node in entry[2].items()
|
||||
if node['class_type'] == 'DistributedImageListEnd']})
|
||||
wake.set()
|
||||
return response
|
||||
async def interrupt(request):
|
||||
interrupted_ids.append((await request.json())['prompt_id'])
|
||||
return await core_handlers['POST', '/interrupt'](request)
|
||||
async def stats(request):
|
||||
running, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS,
|
||||
'item_executions': ITEM_EXECUTIONS, 'submissions': submissions,
|
||||
'running': [e[1] for e in running], 'pending': [e[1] for e in pending],
|
||||
'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()})
|
||||
async def consume():
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
while True:
|
||||
await wake.wait()
|
||||
while (queued := ps.prompt_queue.get(timeout=0)) is not None:
|
||||
entry, item_id = queued
|
||||
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
|
||||
finish_queue_item(ps, item_id, executor)
|
||||
wake.clear()
|
||||
app.add_routes([route for route in ps.routes if
|
||||
(route.path in ('/queue', '/history/{prompt_id}') or
|
||||
(route.path == '/prompt' and route.method == 'GET'))])
|
||||
app.router.add_post('/prompt', prompt)
|
||||
app.router.add_post('/interrupt', interrupt)
|
||||
app.router.add_get('/test-stats', stats)
|
||||
runner, port = await listen(app)
|
||||
consumer = asyncio.create_task(consume())
|
||||
print(f'READY:{port}', flush=True)
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
consumer.cancel()
|
||||
await asyncio.gather(consumer, return_exceptions=True)
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def graph(callback, worker_ids):
|
||||
return {
|
||||
'1': {'class_type': 'TestImageSource', 'inputs': {}},
|
||||
# Generic Autogrow value slots carry per-item reference images, text,
|
||||
# strength and seed (one value per image), plus a shared boolean.
|
||||
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': ['1', 0], 'mask': ['1', 2],
|
||||
'values.value_1': ['1', 1], 'values.value_2': ['1', 3],
|
||||
'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0],
|
||||
'values.value_6': ['1', 2],
|
||||
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
|
||||
'3': {'class_type': 'TestImageOperation', 'inputs': {
|
||||
key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2),
|
||||
('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
|
||||
| {'image': ['8', 0], 'model': ['6', 0]}},
|
||||
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3],
|
||||
'value_2': ['10', 0], 'value_3': ['2', 6], 'value_4': ['2', 7], 'value_5': ['2', 8],
|
||||
'value_7': ['3', 1], 'value_8': ['2', 4], 'value_9': ['2', 9], 'value_10': ['10', 1]}},
|
||||
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2],
|
||||
'values': ['4', 3], 'picked': ['get', 0]}},
|
||||
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
|
||||
'7': {'class_type': 'TestSharedFlag', 'inputs': {}},
|
||||
'8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}},
|
||||
'9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}},
|
||||
'10': {'class_type': 'TestValueOverride', 'inputs': {'text': ['2', 5], 'index': ['2', 2]}},
|
||||
'get': {'class_type': 'DistributedImageListGetValue', 'inputs': {'values': ['4', 3], 'value_name': 'value_2'}},
|
||||
}
|
||||
|
||||
|
||||
def check_carried_values():
|
||||
values, picked = VALUE_CAPTURES[-1]
|
||||
assert len(values) == len(picked) == COUNT
|
||||
for index, (actual, source) in enumerate(zip(values, source_rows())):
|
||||
assert set(actual) == {f'value_{n}' for n in range(1, 11)}, actual.keys()
|
||||
assert torch.equal(actual['value_1'], source['reference_image'])
|
||||
assert actual['value_2'] == picked[index] == f'item-{index}:override'
|
||||
assert actual['value_3'] == source['strength']
|
||||
assert actual['value_4'] == source['seed']
|
||||
assert actual['value_5'] is True
|
||||
assert torch.equal(actual['value_6'], source['mask'])
|
||||
assert torch.equal(actual['value_7'], source['mask'] * 0.5)
|
||||
assert torch.equal(actual['value_8'], source['reference_image'])
|
||||
assert actual['value_8'].shape[0] == 2, 'value tensor batch was split'
|
||||
assert torch.equal(actual['value_9'], source['mask'])
|
||||
assert actual['value_10'] == 200 + index
|
||||
|
||||
|
||||
async def coordinator():
|
||||
global ALLOW_MASTER
|
||||
ps, assets, module, app = await setup()
|
||||
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
|
||||
execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan')
|
||||
lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle')
|
||||
config_mod = importlib.import_module(module.__name__ + '.utils.config')
|
||||
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
|
||||
runner, port = await listen(app)
|
||||
ps.port = port # Match the listener used by local-worker callback resolution.
|
||||
processes, stderr_files = [], []
|
||||
# Old worker capability refusal is read-only and precedes input evaluation.
|
||||
legacy_app = web.Application()
|
||||
legacy_requests = []
|
||||
async def legacy(request):
|
||||
legacy_requests.append(request.method)
|
||||
return web.json_response({'error': 'not found'}, status=404)
|
||||
legacy_app.router.add_route('*', '/{path:.*}', legacy)
|
||||
legacy_runner, legacy_port = await listen(legacy_app)
|
||||
dispatch_module = importlib.import_module(module.__name__ + '.api.orchestration.dispatch')
|
||||
try:
|
||||
try:
|
||||
await dispatch_module.preflight_image_list_workers([
|
||||
{'id': 'legacy', 'type': 'local', 'host': '127.0.0.1', 'port': legacy_port}])
|
||||
raise AssertionError('legacy worker accepted')
|
||||
except RuntimeError as exc:
|
||||
assert 'upgrade' in str(exc), str(exc)
|
||||
assert legacy_requests == ['GET'] and not SOURCE_CALLS
|
||||
finally:
|
||||
await legacy_runner.cleanup()
|
||||
print('IMAGE_LIST_PREFLIGHT_OK legacy worker rejected before materialization or writes')
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix='image-list-network-', dir=scratch) as tmp:
|
||||
workers = []
|
||||
for name, delay in [('fast', 0.02), ('slow', 0.5)]:
|
||||
err = open(Path(tmp) / f'{name}.log', 'w+')
|
||||
stderr_files.append(err)
|
||||
proc = await asyncio.create_subprocess_exec(sys.executable, str(Path(__file__).resolve()),
|
||||
str(COMFY_ROOT), '--worker', str(delay), stdout=asyncio.subprocess.PIPE, stderr=err,
|
||||
env={**os.environ, 'TMPDIR': tmp})
|
||||
processes.append(proc)
|
||||
async def ready():
|
||||
while True:
|
||||
line = await proc.stdout.readline()
|
||||
if not line:
|
||||
err.seek(0)
|
||||
raise AssertionError(f'Worker exited: {err.read()}')
|
||||
if line.startswith(b'READY:'):
|
||||
return int(line.split(b':')[1])
|
||||
worker_port = await asyncio.wait_for(ready(), 60)
|
||||
workers.append({'id': name, 'type': 'local', 'host': '127.0.0.1',
|
||||
'port': worker_port, 'enabled': True})
|
||||
config = {'workers': workers, 'settings': {'websocket_orchestration': False, 'master_delegate_only': True},
|
||||
'master': {'host': f'http://127.0.0.1:{port}'}}
|
||||
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
|
||||
valid = await execution.validate_prompt('network-acceptance', prompt, None)
|
||||
assert valid[0], valid
|
||||
assert {'4', '5', '8', '9'} == set(valid[2]), valid
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
# Exercise production queue preparation; only master enqueue is redirected
|
||||
# to this isolated executor, not the managed server's live queue.
|
||||
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
|
||||
prompt['2']['inputs'].pop('image_list_master_callback')
|
||||
saved_prompt = copy.deepcopy(prompt)
|
||||
outer_executions = []
|
||||
async def enqueue(prepared, workflow, client_id, **kwargs):
|
||||
targets = kwargs['partial_execution_targets']
|
||||
planned, planned_targets = execution_plan.prepare_image_list_execution(prepared)
|
||||
assert planned == prepared and targets == planned_targets
|
||||
if '9' in saved_prompt and 'TestImageOperation' in {
|
||||
node['class_type'] for node in prepared.values()}:
|
||||
assert set(prepared) == set(saved_prompt) - {'9'}, prepared
|
||||
assert prepared['8'] == saved_prompt['8'], 'required utility changed'
|
||||
assert prepared['5'] == saved_prompt['5'], 'post-End output changed'
|
||||
assert set(targets) == {'4', '5'}, targets
|
||||
assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True
|
||||
assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True
|
||||
pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4())
|
||||
checked = await execution.validate_prompt(pid, prepared, targets)
|
||||
assert checked[0], checked
|
||||
assert not checked[3] and Counter(checked[2]) == Counter(targets), checked
|
||||
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
# Exactly the core queue tuple/put/get/task_done lifecycle. Start
|
||||
# must prove its roots from this real running record, not a mock.
|
||||
number = len(outer_executions)
|
||||
ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {}))
|
||||
entry, item_id = ps.prompt_queue.get(timeout=0)
|
||||
assert lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
before_captures = len(CAPTURED)
|
||||
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
|
||||
finish_queue_item(ps, item_id, executor)
|
||||
assert ps.prompt_queue.get_current_queue_volatile() == ([], [])
|
||||
history = ps.prompt_queue.get_history(pid)[pid]
|
||||
assert Counter(history['prompt'][4]) == Counter(targets)
|
||||
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
outer_executions.append({'prompt_id': pid, 'targets': targets,
|
||||
'captures': len(CAPTURED) - before_captures})
|
||||
if executor.success:
|
||||
assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once'
|
||||
return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]}
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', [w['id'] for w in workers]), 45)
|
||||
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
|
||||
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
|
||||
assert prompt == saved_prompt, 'queue preparation changed the saved/source graph'
|
||||
if not executor.success:
|
||||
for err in stderr_files:
|
||||
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS and not UTILITY_CALLS
|
||||
assert len(CAPTURED) == 1
|
||||
check_carried_values()
|
||||
images, masks, texts = CAPTURED[0]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (i, images[i])
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), i
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', texts
|
||||
async with aiohttp.ClientSession() as session:
|
||||
stats = []
|
||||
for worker in workers:
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
stats.append(await response.json())
|
||||
counts = [len(row['calls']) for row in stats]
|
||||
assert counts[0] > counts[1] > 0, counts
|
||||
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
|
||||
for row in stats:
|
||||
assert Counter(row['utility_calls']) == Counter(row['calls']), row
|
||||
assert len(row['submissions']) == len(row['calls'])
|
||||
assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls'])
|
||||
for submission in row['submissions']:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
assert not runtime._MASTER_JOBS
|
||||
# Three sequential regions on one worker, including optional outputs
|
||||
# left unconnected. Region N+1 must not race the prior worker End.
|
||||
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
|
||||
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
|
||||
chained['constant'] = {'class_type': 'TestSharedFlag', 'inputs': {}}
|
||||
previous = '1'
|
||||
value_source = 'constant'
|
||||
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
|
||||
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': [previous, 0], 'values.value_1': [value_source, 0]}}
|
||||
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': [start, 0], 'context': [start, 3]}}
|
||||
value_source = f'get-{end}'
|
||||
chained[value_source] = {'class_type': 'DistributedImageListGetValue', 'inputs': {
|
||||
'values': [end, 3], 'value_name': 'value_1'}}
|
||||
previous = end
|
||||
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
|
||||
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2], 'values': ['7', 3]}}
|
||||
saved_chained = copy.deepcopy(chained)
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
chained, None, 'test', ['slow']), 45)
|
||||
assert chained == saved_chained
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
|
||||
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
|
||||
assert CAPTURED[1][1:] == ([None, None], [None, None])
|
||||
assert VALUE_CAPTURES[1][0] == [{'value_1': True}, {'value_1': True}]
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
|
||||
# Same public queue path, now with master participation enabled. The
|
||||
# parent queue stays occupied while its local slot processes items.
|
||||
ALLOW_MASTER = True
|
||||
config['settings']['master_delegate_only'] = False
|
||||
participation_counts = {}
|
||||
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
|
||||
# Keep cold-run exact-once assertions independent of other modes.
|
||||
importlib.import_module(module.__name__ + '.image_list.local_executor')._RETAINED_OUTPUTS = None
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
|
||||
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
|
||||
async with aiohttp.ClientSession() as session:
|
||||
before_worker_stats = []
|
||||
for worker in workers:
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
before_worker_stats.append(await response.json())
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', selected), 45)
|
||||
assert prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
assert executor.success, (mode, executor.status_messages)
|
||||
assert queued[2] == len(selected)
|
||||
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
|
||||
local_calls = BRANCH_CALLS[before_local:]
|
||||
assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode
|
||||
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
|
||||
assert local_calls, f'{mode}: master did not process any items'
|
||||
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
|
||||
remote_calls = []
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for worker, before_stats in zip(workers, before_worker_stats):
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
current = await response.json()
|
||||
calls = current['calls'][len(before_stats['calls']):]
|
||||
utilities = current['utility_calls'][len(before_stats['utility_calls']):]
|
||||
submissions = current['submissions'][len(before_stats['submissions']):]
|
||||
events = current['item_executions'][len(before_stats['item_executions']):]
|
||||
assert Counter(utilities) == Counter(calls), (mode, current)
|
||||
assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events})
|
||||
for submission in submissions:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
remote_calls.extend(calls)
|
||||
assert not current['running'] and not current['pending'], current
|
||||
if mode == 'mixed':
|
||||
assert remote_calls, 'mixed execution failed to use direct workers'
|
||||
else:
|
||||
assert not remote_calls and len(local_calls) == COUNT
|
||||
assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
|
||||
check_carried_values()
|
||||
images, masks, texts = CAPTURED[-1]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
|
||||
# Repeat the completed master-only queue with a fresh outer executor.
|
||||
# Every item must be delivered again without recomputing its branch.
|
||||
before_sources, before_local = len(SOURCE_CALLS), len(BRANCH_CALLS)
|
||||
before_models, before_utilities = len(MODEL_CALLS), len(UTILITY_CALLS)
|
||||
before_captures = len(CAPTURED)
|
||||
previous_outputs = CAPTURED[-1]
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'load_config', return_value=config), \
|
||||
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
|
||||
prompt, None, 'test', []), 45)
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(SOURCE_CALLS) == before_sources + 1
|
||||
assert len(BRANCH_CALLS) == before_local
|
||||
assert len(MODEL_CALLS) == before_models
|
||||
assert len(UTILITY_CALLS) == before_utilities
|
||||
assert len(CAPTURED) == before_captures + 1
|
||||
for before, after in zip(previous_outputs[:2], CAPTURED[-1][:2]):
|
||||
assert len(before) == len(after) == COUNT
|
||||
assert all(torch.equal(a, b) for a, b in zip(before, after))
|
||||
assert previous_outputs[2] == CAPTURED[-1][2]
|
||||
check_carried_values()
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
assert prompt == saved_prompt
|
||||
print('IMAGE_LIST_CACHE_REQUEUE_OK unchanged items reused; fresh ordered delivery and cleanup')
|
||||
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
|
||||
print('IMAGE_LIST_VALUES_OK passthrough all six types; sparse override/add; Get; sequential Get-to-Start; direct/mixed/master-only')
|
||||
assert len(outer_executions) == 5 and all(row['captures'] == 1 for row in outer_executions)
|
||||
print('IMAGE_LIST_TARGETS_OK ' + json.dumps({
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
|
||||
'required_utility_once_per_item': True, 'preview_discarded': True,
|
||||
'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True}))
|
||||
ALLOW_MASTER = False
|
||||
config['settings']['master_delegate_only'] = True
|
||||
# Preserve the original direct-runtime cancellation/failure fixture:
|
||||
# it has no outer queue and must not claim target-control evidence.
|
||||
direct_prompt = copy.deepcopy(prompt)
|
||||
del direct_prompt['8'], direct_prompt['9']
|
||||
direct_prompt['3']['inputs']['image'] = ['2', 0]
|
||||
# Real Comfy user-interrupt flag while a real worker executor is active.
|
||||
rows = source_rows()
|
||||
items = runtime.normalize_image_list_inputs(
|
||||
[row['image'] for row in rows], [row['mask'] for row in rows],
|
||||
{f'value_{n}': [row[key] for row in rows]
|
||||
for n, key in ((1, 'reference_image'), (2, 'text'), (3, 'strength'), (4, 'seed'), (6, 'mask'))}
|
||||
| {'value_5': True})
|
||||
slow = workers[1]
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def slow_stats():
|
||||
async with session.get(f"http://127.0.0.1:{slow['port']}/test-stats") as response:
|
||||
return await response.json()
|
||||
async def until(predicate):
|
||||
for _ in range(150):
|
||||
current = await slow_stats()
|
||||
if predicate(current):
|
||||
return current
|
||||
await asyncio.sleep(0.02)
|
||||
raise AssertionError('worker state transition timed out')
|
||||
before = await until(lambda state: not state['running'] and not state['pending'])
|
||||
context = runtime.create_master_job('2', '4', direct_prompt, items, ['slow'],
|
||||
f'http://127.0.0.1:{port}')
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
task = asyncio.create_task(runtime.run_image_list_job(context, include_values=True))
|
||||
active = await until(lambda state: bool(state['running']))
|
||||
owned_id = active['running'][0]
|
||||
comfy.model_management.interrupt_current_processing(True)
|
||||
try:
|
||||
await asyncio.wait_for(task, 5)
|
||||
raise AssertionError('interrupt was ignored')
|
||||
except runtime.ImageListCancelledError:
|
||||
pass
|
||||
finally:
|
||||
comfy.model_management.interrupt_current_processing(False)
|
||||
after = await until(lambda state: not state['running'] and not state['pending'])
|
||||
assert after['interrupted'] == [owned_id], after
|
||||
assert len(after['calls']) <= len(before['calls']) + 1
|
||||
assert set(before['history']) <= set(after['history'])
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
|
||||
# A worker exception before End must fail, not hang or return a partial list.
|
||||
bad_rows = copy.deepcopy(items)
|
||||
bad_rows[0]['values']['value_4'] = -1
|
||||
context = runtime.create_master_job('2', '4', direct_prompt, bad_rows, ['fast'],
|
||||
f'http://127.0.0.1:{port}')
|
||||
with patch.object(config_mod, 'load_config', return_value=config), \
|
||||
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
|
||||
try:
|
||||
await asyncio.wait_for(runtime.run_image_list_job(context, include_values=True), 5)
|
||||
raise AssertionError('worker exception was ignored')
|
||||
except RuntimeError as exc:
|
||||
assert 'failed image-list item' in str(exc), str(exc)
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_FAILURE_OK worker exception surfaced; context removed')
|
||||
print('IMAGE_LIST_NETWORK_OK ' + json.dumps({'workers': 2, 'items': COUNT,
|
||||
'assignments': dict(zip(['fast', 'slow'], counts)), 'ordered_exact_outputs': True,
|
||||
'master_branch_calls': len(BRANCH_CALLS), 'source_calls': len(SOURCE_CALLS)}))
|
||||
finally:
|
||||
for proc in processes:
|
||||
if proc.returncode is None:
|
||||
proc.terminate()
|
||||
await asyncio.gather(*(proc.wait() for proc in processes))
|
||||
for err in stderr_files:
|
||||
err.close()
|
||||
await runner.cleanup()
|
||||
network = importlib.import_module(module.__name__ + '.utils.network')
|
||||
session = await network.get_client_session()
|
||||
await session.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(worker_main() if WORKER else coordinator())
|
||||
@@ -0,0 +1,114 @@
|
||||
"""CPU-only native Comfy execution across private-slot lifetimes."""
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
|
||||
pkg = '_local_retention_probe'
|
||||
for name, path in [(pkg, ROOT), (pkg + '.image_list', ROOT / 'image_list')]:
|
||||
module = types.ModuleType(name)
|
||||
module.__path__ = [str(path)]
|
||||
sys.modules[name] = module
|
||||
local = importlib.import_module(pkg + '.image_list.local_executor')
|
||||
|
||||
CALLS = []
|
||||
OUTPUTS = []
|
||||
|
||||
class Counted:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'value': ('INT',)}}
|
||||
RETURN_TYPES = ('INT',)
|
||||
FUNCTION = 'run'
|
||||
def run(self, value):
|
||||
if value < 0:
|
||||
raise ValueError('intentional retention failure')
|
||||
CALLS.append(value)
|
||||
return (value,)
|
||||
|
||||
class Uncached(Counted):
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float('nan')
|
||||
|
||||
class Capture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'value': ('INT',), 'delivery': ('STRING',)}}
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = 'run'
|
||||
def run(self, value, delivery):
|
||||
OUTPUTS.append((value, delivery))
|
||||
return ()
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(RetentionCounted=Counted, RetentionUncached=Uncached, RetentionCapture=Capture)
|
||||
|
||||
async def run(slot, value, delivery, uncached=False):
|
||||
prompt = {
|
||||
'source': {'class_type': 'RetentionUncached' if uncached else 'RetentionCounted', 'inputs': {'value': value}},
|
||||
'end': {'class_type': 'RetentionCapture', 'inputs': {'value': ['source', 0], 'delivery': delivery}},
|
||||
}
|
||||
await asyncio.wrap_future(slot.submit(prompt, delivery, 'end'))
|
||||
|
||||
async def main():
|
||||
first = local.LocalExecutionSlot()
|
||||
await run(first, 1, 'first-1')
|
||||
await run(first, 2, 'first-2')
|
||||
await first.close()
|
||||
assert local._RETAINED_OUTPUTS is not None
|
||||
second = local.LocalExecutionSlot()
|
||||
await run(second, 1, 'second-1')
|
||||
await run(second, 2, 'second-2')
|
||||
await run(second, 3, 'changed')
|
||||
await second.close()
|
||||
assert CALLS == [1, 2, 3], CALLS
|
||||
assert OUTPUTS == [(1, 'first-1'), (2, 'first-2'), (1, 'second-1'), (2, 'second-2'), (3, 'changed')], OUTPUTS
|
||||
third = local.LocalExecutionSlot()
|
||||
await run(third, 4, 'uncached-1', True)
|
||||
await run(third, 4, 'uncached-2', True)
|
||||
await third.close()
|
||||
assert CALLS[-2:] == [4, 4]
|
||||
comfy.cli_args.args.cache_none = True
|
||||
disabled = local.LocalExecutionSlot()
|
||||
await run(disabled, 1, 'disabled-1')
|
||||
await run(disabled, 1, 'disabled-2')
|
||||
await disabled.close()
|
||||
assert CALLS[-2:] == [1, 1]
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
comfy.cli_args.args.cache_none = False
|
||||
failed = local.LocalExecutionSlot()
|
||||
await run(failed, 7, 'before-failure')
|
||||
try:
|
||||
await run(failed, -1, 'failure')
|
||||
except RuntimeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError('failure fixture unexpectedly succeeded')
|
||||
await failed.close()
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
cancelled = local.LocalExecutionSlot()
|
||||
await run(cancelled, 8, 'cancelled')
|
||||
cancelled.request_stop()
|
||||
await cancelled.close()
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
oversized = execution.LRUCache(execution.CacheKeySetInputSignature, max_size=local.NATIVE_CACHE_MAX_ENTRIES)
|
||||
empty = execution.DynamicPrompt({})
|
||||
await oversized.set_prompt(empty, [], execution.IsChangedCache('bound', empty, oversized))
|
||||
oversized.min_generation = oversized.generation
|
||||
oversized.cache = {n: n for n in range(local.NATIVE_CACHE_MAX_ENTRIES + 1)}
|
||||
local._retain_outputs(oversized)
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
print('PASS: cross-run native cache reuse, fresh delivery, input change, noncacheable nodes, cache_none, cancellation, retention bound')
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -45,7 +45,8 @@ ADDED_OPTIONAL_INPUTS = {
|
||||
|
||||
# Nodes added after the V1 baseline; not part of saved-workflow contracts.
|
||||
# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
|
||||
ADDED_NODES = {'DistributedWorkerSink'}
|
||||
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd',
|
||||
'DistributedImageListGetValue'}
|
||||
|
||||
|
||||
def normalized_input(value):
|
||||
@@ -254,6 +255,109 @@ async def check_execution(mapping, module, prompt_server, asset_manager):
|
||||
print('V3_EXECUTION_OK nine nodes; upscale GPU boundary mocked; preview artifact verified; sink writes nothing')
|
||||
|
||||
|
||||
async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager):
|
||||
"""Execute real native V3 list/lazy flow; master must skip the processing branch."""
|
||||
branch_calls = []
|
||||
class ImageListProbe:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',)}}
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
FUNCTION = 'run'
|
||||
CATEGORY = 'tests'
|
||||
def run(self, image):
|
||||
branch_calls.append(image)
|
||||
return (image,)
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS['ImageListTestProbe'] = ImageListProbe
|
||||
runtime = sys.modules[module.__name__ + '.image_list.runtime']
|
||||
original_contexts = set(runtime._MASTER_JOBS)
|
||||
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
|
||||
workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
|
||||
async def fake_worker_results(_context_id, *, model_callback=None, include_values=False):
|
||||
state = runtime.get_master_job(_context_id)
|
||||
assert state.prompt_id == 'image-list-lazy', state.prompt_id
|
||||
assert state.workflow_meta == workflow_meta, state.workflow_meta
|
||||
assert len(prompt_server.app.cleanup_ctx) >= 1, 'lifecycle cleanup was not registered'
|
||||
return (
|
||||
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
|
||||
[None, None, None], [None, None, None],
|
||||
[{}, {}, {}],
|
||||
)
|
||||
|
||||
graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 3, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': ['1', 0], 'image_list_enabled_worker_ids': '["test-worker"]',
|
||||
'image_list_master_callback': 'http://127.0.0.1:1',
|
||||
'image_list_workflow_meta': workflow_meta,
|
||||
}},
|
||||
'3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}},
|
||||
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['3', 0], 'context': ['2', 3],
|
||||
}},
|
||||
'5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('image-list-lazy', graph, None)
|
||||
assert valid[0], valid
|
||||
validator = runtime.validate_image_list_regions
|
||||
assert 'texts' not in graph['4']['inputs']
|
||||
assert len(validator(graph)) == 1
|
||||
connected_text_graph = json.loads(json.dumps(graph))
|
||||
connected_text_graph['6'] = {
|
||||
'class_type': 'ImageListTextProbe', 'inputs': {'text': ['2', 4]},
|
||||
}
|
||||
connected_text_graph['4']['inputs']['texts'] = ['6', 0]
|
||||
assert len(validator(connected_text_graph)) == 1
|
||||
for literal_text in ('', 'per-item text'):
|
||||
literal_text_graph = json.loads(json.dumps(graph))
|
||||
literal_text_graph['4']['inputs']['texts'] = literal_text
|
||||
try:
|
||||
validator(literal_text_graph)
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError(f'End accepted unconnected literal text: {literal_text!r}')
|
||||
print('IMAGE_LIST_END_TEXT_INPUT_OK native forceInput; unconnected omitted; linked accepted; literals rejected')
|
||||
import folder_paths
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
try:
|
||||
# Resolve genuine Comfy output/temp annotations against the pinned core,
|
||||
# not only the controlled folder_paths seam in unit tests.
|
||||
media = sys.modules[module.__name__ + '.api.orchestration.media_sync']
|
||||
with tempfile.TemporaryDirectory(prefix='media-roots-', dir=scratch) as media_root:
|
||||
for kind in ('input', 'output', 'temp'):
|
||||
source = Path(media_root) / kind / 'shared/ref.png'
|
||||
source.parent.mkdir(parents=True)
|
||||
source.write_bytes(kind.encode())
|
||||
with patch.object(folder_paths, 'input_directory', str(Path(media_root) / 'input')), \
|
||||
patch.object(folder_paths, 'output_directory', str(Path(media_root) / 'output')), \
|
||||
patch.object(folder_paths, 'temp_directory', str(Path(media_root) / 'temp')):
|
||||
for kind in ('output', 'temp'):
|
||||
refs = media._find_media_references({'1': {'inputs': {'image': f'shared/ref.png [{kind}]'}}})
|
||||
assert len(refs) == 1
|
||||
assert media._load_media_file_sync(refs[0])[0] == kind.encode()
|
||||
with tempfile.TemporaryDirectory(prefix='image-list-lazy-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp), \
|
||||
patch.object(image_list_nodes, 'run_image_list_job', fake_worker_results):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, 'image-list-lazy', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
assert branch_calls == [], 'master evaluated the lazy per-image branch'
|
||||
records = executor.history_result['outputs']['5']['images']
|
||||
assert len(records) == 3, records
|
||||
for record in records:
|
||||
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
||||
with Image.open(preview) as image:
|
||||
assert image.size == (8, 8) and image.mode == 'RGB'
|
||||
finally:
|
||||
nodes.NODE_CLASS_MAPPINGS.pop('ImageListTestProbe', None)
|
||||
for context_id in set(runtime._MASTER_JOBS) - original_contexts:
|
||||
runtime.remove_master_job(context_id)
|
||||
print('IMAGE_LIST_LAZY_EXECUTION_OK real PromptExecutor; branch skipped on master; 3 previews verified')
|
||||
|
||||
|
||||
async def main():
|
||||
asset_manager = default_asset_manager()
|
||||
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
@@ -264,7 +368,7 @@ async def main():
|
||||
extension = await module.comfy_entrypoint()
|
||||
classes = await extension.get_node_list()
|
||||
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
|
||||
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 9
|
||||
assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 12
|
||||
assert set(mapping) == set(BASELINE) | ADDED_NODES
|
||||
for node_id in ADDED_NODES:
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
|
||||
@@ -315,9 +419,93 @@ async def main():
|
||||
if old['output_node']:
|
||||
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
||||
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
||||
start_schema = mapping['DistributedImageListStart'].GET_SCHEMA()
|
||||
end_schema = mapping['DistributedImageListEnd'].GET_SCHEMA()
|
||||
assert start_schema.node_id == 'DistributedImageListStart'
|
||||
assert start_schema.display_name == 'Distributed Image List Start'
|
||||
assert end_schema.node_id == 'DistributedImageListEnd'
|
||||
assert end_schema.display_name == 'Distributed Image List End'
|
||||
runtime = sys.modules[module.__name__ + '.image_list.runtime']
|
||||
assert [input.id for input in end_schema.inputs] == ['images', 'masks', 'texts', 'context'] + list(runtime.VALUE_NAMES)
|
||||
assert [output.io_type for output in end_schema.outputs] == [
|
||||
'IMAGE', 'MASK', 'STRING', 'DISTRIBUTED_IMAGE_LIST_VALUES']
|
||||
assert end_schema.outputs[3].id == end_schema.outputs[3].display_name == 'values'
|
||||
assert all(output.is_output_list for output in end_schema.outputs)
|
||||
for override in end_schema.inputs[4:]:
|
||||
assert override.optional and override.lazy
|
||||
assert override.get_io_type() == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN'
|
||||
getter = mapping['DistributedImageListGetValue']
|
||||
getter_schema = getter.GET_SCHEMA()
|
||||
assert getter_schema.display_name == 'Distributed Image List Get Value'
|
||||
assert not getter_schema.is_input_list
|
||||
assert getter_schema.outputs[0].id == getter_schema.outputs[0].display_name == 'value'
|
||||
assert getter_schema.outputs[0].io_type == '*'
|
||||
assert getter_schema.inputs[1].options == list(runtime.VALUE_NAMES)
|
||||
assert getter_schema.inputs[1].default == 'value_1'
|
||||
# Real core maps scalar mappings, never arbitrary tensor batches.
|
||||
tensor = torch.ones(3, 2, 2, 3)
|
||||
data, missing, hidden = execution.get_input_data({'value_name': 'value_1'}, getter, 'get-value')
|
||||
data['values'] = [{'value_1': tensor}, {'value_1': 17}, {'value_1': False}]
|
||||
mapped = await execution._async_map_node_over_list('get-value', 'get-value', getter, data,
|
||||
getter.FUNCTION, v3_data=hidden)
|
||||
mapped = await execution.resolve_map_node_over_list_results(mapped)
|
||||
assert len(mapped) == 3 and mapped[0].result[0] is tensor
|
||||
assert mapped[1].result == (17,) and mapped[2].result == (False,)
|
||||
try:
|
||||
getter.execute({}, 'value_1')
|
||||
raise AssertionError('missing Get Value key accepted')
|
||||
except ValueError as exc:
|
||||
assert 'absent' in str(exc)
|
||||
texts_input = next(input for input in end_schema.inputs if input.id == 'texts')
|
||||
texts_schema = texts_input.as_dict()
|
||||
assert texts_schema.get('forceInput') is True, texts_schema
|
||||
assert 'default' not in texts_schema, texts_schema
|
||||
assert [out.io_type for out in start_schema.outputs] == [
|
||||
'IMAGE', 'MASK', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT'] + ['*'] * 10
|
||||
assert [out.display_name for out in start_schema.outputs][4:] == [f'value_{i}' for i in range(1, 11)]
|
||||
assert all(out.is_output_list for out in start_schema.outputs)
|
||||
assert start_schema.is_input_list and start_schema.accept_all_inputs
|
||||
start_info = json.loads(json.dumps(mapping['DistributedImageListStart'].GET_NODE_INFO_V1()))
|
||||
assert list(start_info['input']['optional']) == ['images', 'mask', 'values'], start_info['input']
|
||||
values_spec = start_info['input']['optional']['values']
|
||||
assert values_spec[0] == 'COMFY_AUTOGROW_V3', values_spec
|
||||
template = values_spec[1]['template']
|
||||
assert template['names'] == [f'value_{i}' for i in range(1, 11)] and template['min'] == 0, template
|
||||
assert template['input']['required']['value'][0] == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN', template
|
||||
# Autogrow + INPUT_IS_LIST through the real parser/list mapper, values
|
||||
# still split 1-or-N, unconnected slots ignored.
|
||||
# get_input_data must accept Autogrow names; constants are wrapped once.
|
||||
parsed, missing, hidden = execution.get_input_data(
|
||||
{'values.value_1': 'shared', 'values.value_3': 4},
|
||||
mapping['DistributedImageListStart'], unique_id='autogrow-probe', extra_data={})
|
||||
assert not missing and parsed['values.value_1'] == ['shared'], parsed
|
||||
# Linked inputs arrive as the upstream output list (one element per value).
|
||||
prepared = {'images': [torch.zeros((3, 2, 2, 3))], 'values.value_1': ['shared'],
|
||||
'values.value_3': [4, 5, 6]}
|
||||
seen = {}
|
||||
runtime_mod = sys.modules[module.__name__ + '.image_list.runtime']
|
||||
def capture(_images, _mask, values):
|
||||
seen['values'] = values
|
||||
raise RuntimeError('captured')
|
||||
nodes_mod = sys.modules[module.__name__ + '.nodes.image_list']
|
||||
with patch.object(nodes_mod, 'normalize_image_list_inputs', capture):
|
||||
try:
|
||||
await asyncio.to_thread(lambda: asyncio.run(execution._async_map_node_over_list(
|
||||
'autogrow-probe', 'autogrow-probe', mapping['DistributedImageListStart'], prepared,
|
||||
'execute', v3_data=hidden)))
|
||||
except RuntimeError as exc:
|
||||
assert 'captured' in str(exc), exc
|
||||
assert seen['values'] == {'value_1': ['shared'], 'value_3': [4, 5, 6]}, seen
|
||||
rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values'])
|
||||
assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)]
|
||||
print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item')
|
||||
assert all(out.is_output_list for out in end_schema.outputs)
|
||||
assert end_schema.is_input_list and end_schema.is_output_node
|
||||
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
|
||||
print('SCHEMA_PARITY_OK', len(mapping))
|
||||
check_saved_workflows(mapping)
|
||||
await check_execution(mapping, module, prompt_server, asset_manager)
|
||||
await check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager)
|
||||
print('V3_ACCEPTANCE_OK')
|
||||
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
@@ -87,5 +88,59 @@ class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.module, self.server = _load_async_helpers_module()
|
||||
self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}}
|
||||
self.validate = AsyncMock(return_value=(True, None, ["9"], {}))
|
||||
self.module.execution.validate_prompt = self.validate
|
||||
|
||||
async def test_targets_survive_callback_omission_and_are_queued_exactly(self):
|
||||
seen = []
|
||||
def callback(payload):
|
||||
seen.append(payload["partial_execution_targets"][:])
|
||||
payload["partial_execution_targets"].append("10")
|
||||
return {"prompt": payload["prompt"]}
|
||||
self.server.trigger_on_prompt = callback
|
||||
targets = ["9"]
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
|
||||
self.assertEqual(seen, [["9"]])
|
||||
self.assertEqual(targets, ["9"])
|
||||
self.assertEqual(self.validate.call_args.args[2], ["9"])
|
||||
self.assertEqual(self.server.prompt_queue.items[0][4], ["9"])
|
||||
|
||||
async def test_invalid_requests_never_reach_validation_or_queue(self):
|
||||
for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}):
|
||||
with self.subTest(targets=targets):
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
|
||||
self.validate.assert_not_awaited()
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_removed_target_after_callback_is_rejected(self):
|
||||
self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}}
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"])
|
||||
self.validate.assert_not_awaited()
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self):
|
||||
for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}),
|
||||
(["9", "10"], {"9": {"errors": [{"message": "bad"}]}}),
|
||||
(["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})):
|
||||
with self.subTest(roots=roots, errors=errors):
|
||||
self.validate.return_value = (True, None, roots, errors)
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"])
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_none_preserves_legacy_payload_and_validation(self):
|
||||
seen = []
|
||||
self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload
|
||||
await self.module.queue_prompt_payload(self.prompt)
|
||||
self.assertNotIn("partial_execution_targets", seen[0])
|
||||
self.assertIsNone(self.validate.call_args.args[2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -4,7 +4,7 @@ import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
def _load_dispatch_module():
|
||||
@@ -55,7 +55,13 @@ def _load_dispatch_module():
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _ClientConnectorError(Exception):
|
||||
class _ClientError(Exception):
|
||||
pass
|
||||
|
||||
class _ClientConnectorError(_ClientError):
|
||||
pass
|
||||
|
||||
class _ClientConnectionError(_ClientError):
|
||||
pass
|
||||
|
||||
class _WSMsgType:
|
||||
@@ -76,7 +82,9 @@ def _load_dispatch_module():
|
||||
self.closed = True
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
setattr(aiohttp_module, "ClientError", _ClientError)
|
||||
aiohttp_module.ClientConnectorError = _ClientConnectorError
|
||||
setattr(aiohttp_module, "ClientConnectionError", _ClientConnectionError)
|
||||
aiohttp_module.WSMsgType = _WSMsgType
|
||||
aiohttp_module.TCPConnector = _TCPConnector
|
||||
aiohttp_module.ClientSession = _ClientSession
|
||||
@@ -100,6 +108,113 @@ def _load_dispatch_module():
|
||||
dispatch = _load_dispatch_module()
|
||||
|
||||
|
||||
def _load_prompt_transform_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
|
||||
package_name = "prompt_transform_worker_selection_testpkg"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".utils"):
|
||||
package = types.ModuleType(package_name + suffix)
|
||||
package.__path__ = []
|
||||
sys.modules[package_name + suffix] = package
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.prompt_transform",
|
||||
module_path,
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
prompt_transform = _load_prompt_transform_module()
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, status=200, payload=None, error=None):
|
||||
self.status = status
|
||||
self.payload = payload
|
||||
self.error = error
|
||||
|
||||
async def __aenter__(self):
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
async def json(self):
|
||||
return self.payload
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, responses):
|
||||
self.responses = responses
|
||||
|
||||
def get(self, url, timeout=None):
|
||||
del timeout
|
||||
worker_id = url.removeprefix("http://").split("/", 1)[0]
|
||||
endpoint = "capabilities" if url.endswith("/capabilities") else "queue"
|
||||
response = self.responses[(worker_id, endpoint)]
|
||||
if isinstance(response, BaseException):
|
||||
return _FakeResponse(error=response)
|
||||
status, payload = response
|
||||
return _FakeResponse(status=status, payload=payload)
|
||||
|
||||
|
||||
def _valid_capability():
|
||||
return {
|
||||
"supported": True,
|
||||
"execution_targets": True,
|
||||
"protocol_version": 2,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
|
||||
}
|
||||
|
||||
|
||||
def _valid_queue(running=None, pending=None):
|
||||
return {"queue_running": running or [], "queue_pending": pending or []}
|
||||
|
||||
|
||||
class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_http_preserves_targets_and_omits_default_field(self):
|
||||
response = types.SimpleNamespace(raise_for_status=lambda: None)
|
||||
class Response:
|
||||
async def __aenter__(self): return response
|
||||
async def __aexit__(self, *args): pass
|
||||
calls = []
|
||||
def post(url, **kwargs):
|
||||
calls.append(kwargs['json'])
|
||||
return Response()
|
||||
session = types.SimpleNamespace(post=post)
|
||||
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
|
||||
partial_execution_targets=['9'])
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None)
|
||||
self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
|
||||
self.assertNotIn('partial_execution_targets', calls[1])
|
||||
|
||||
async def test_websocket_sends_targets_all_the_way_to_wire(self):
|
||||
import json
|
||||
sent = []
|
||||
async def receive():
|
||||
return types.SimpleNamespace(type=dispatch.aiohttp.WSMsgType.TEXT, data=json.dumps({
|
||||
'type': 'dispatch_ack', 'request_id': sent[-1]['request_id'],
|
||||
'ok': True, 'prompt_id': 'worker-prompt'}))
|
||||
ws = types.SimpleNamespace(send_json=AsyncMock(side_effect=sent.append),
|
||||
receive=receive, close=AsyncMock())
|
||||
session = types.SimpleNamespace(ws_connect=AsyncMock(return_value=ws))
|
||||
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
|
||||
result = await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
|
||||
use_websocket=True, return_prompt_id=True, partial_execution_targets=['9'])
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, use_websocket=True)
|
||||
self.assertEqual(result, 'worker-prompt')
|
||||
self.assertEqual(sent[0]['partial_execution_targets'], ['9'])
|
||||
self.assertNotIn('partial_execution_targets', sent[1])
|
||||
|
||||
|
||||
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_select_active_workers_filters_offline(self):
|
||||
workers = [
|
||||
@@ -249,6 +364,231 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
self.assertIsNone(selected)
|
||||
|
||||
async def test_image_list_preflight_skips_unreachable_worker_and_keeps_healthy_worker(self):
|
||||
workers = [{"id": "healthy"}, {"id": "offline"}]
|
||||
responses = {
|
||||
(worker_id, "capabilities"): (200, _valid_capability())
|
||||
for worker_id in ("healthy", "offline")
|
||||
}
|
||||
responses.update({
|
||||
(worker_id, "queue"): (200, _valid_queue())
|
||||
for worker_id in ("healthy", "offline")
|
||||
})
|
||||
responses[("offline", "capabilities")] = dispatch.aiohttp.ClientConnectionError("connection refused")
|
||||
session = _FakeSession(responses)
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
), patch.object(dispatch, "log") as log_mock:
|
||||
result = await dispatch.preflight_image_list_workers(workers)
|
||||
|
||||
self.assertEqual(result, {"healthy": True})
|
||||
self.assertTrue(any("offline" in str(call) for call in log_mock.call_args_list))
|
||||
|
||||
async def test_image_list_preflight_skips_timeout_and_keeps_healthy_worker(self):
|
||||
workers = [{"id": "healthy"}, {"id": "slow"}]
|
||||
responses: dict = {
|
||||
(worker_id, "capabilities"): (200, _valid_capability())
|
||||
for worker_id in ("healthy", "slow")
|
||||
}
|
||||
responses.update({
|
||||
(worker_id, "queue"): (200, _valid_queue())
|
||||
for worker_id in ("healthy", "slow")
|
||||
})
|
||||
responses[("slow", "queue")] = asyncio.TimeoutError()
|
||||
session = _FakeSession(responses)
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
), patch.object(dispatch, "log") as log_mock:
|
||||
result = await dispatch.preflight_image_list_workers(workers)
|
||||
|
||||
self.assertEqual(result, {"healthy": True})
|
||||
self.assertTrue(any("slow" in str(call) for call in log_mock.call_args_list))
|
||||
|
||||
async def test_image_list_preflight_excludes_busy_workers(self):
|
||||
workers = [{"id": "idle"}, {"id": "busy"}]
|
||||
responses = {
|
||||
(worker_id, "capabilities"): (200, _valid_capability())
|
||||
for worker_id in ("idle", "busy")
|
||||
}
|
||||
responses.update({
|
||||
("idle", "queue"): (200, _valid_queue()),
|
||||
("busy", "queue"): (200, _valid_queue(running=[{"prompt": "existing"}])),
|
||||
})
|
||||
session = _FakeSession(responses)
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
result = await dispatch.preflight_image_list_workers(workers)
|
||||
|
||||
self.assertEqual(result, {"idle": True})
|
||||
|
||||
async def test_image_list_preflight_reports_no_idle_workers_when_all_are_unavailable(self):
|
||||
workers = [{"id": "offline-a"}, {"id": "offline-b"}]
|
||||
responses = {
|
||||
(worker_id, "capabilities"): dispatch.aiohttp.ClientConnectionError("unreachable")
|
||||
for worker_id in ("offline-a", "offline-b")
|
||||
}
|
||||
session = _FakeSession(responses)
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
self.assertEqual(
|
||||
await dispatch.preflight_image_list_workers(workers, require_idle=False),
|
||||
{},
|
||||
)
|
||||
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
|
||||
await dispatch.preflight_image_list_workers(workers, require_idle=True)
|
||||
|
||||
async def test_image_list_preflight_requires_exact_execution_targets_support(self):
|
||||
for value in (None, False, 1, "true", []):
|
||||
with self.subTest(value=value):
|
||||
capability = _valid_capability()
|
||||
if value is None:
|
||||
capability.pop("execution_targets")
|
||||
else:
|
||||
capability["execution_targets"] = value
|
||||
session = _FakeSession({("worker", "capabilities"): (200, capability)})
|
||||
async def fake_session():
|
||||
return session
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "execution_targets.*upgrade"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_value_results_requires_exact_true_only_for_overrides(self):
|
||||
for value in (None, False, 1, "true", [], True):
|
||||
capability = _valid_capability()
|
||||
if value is not None:
|
||||
capability['value_results'] = value
|
||||
session = _FakeSession({('worker', 'capabilities'): (200, capability),
|
||||
('worker', 'queue'): (200, _valid_queue())})
|
||||
async def fake_session():
|
||||
return session
|
||||
with patch.object(dispatch, 'get_client_session', side_effect=fake_session), patch.object(
|
||||
dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
self.assertEqual(await dispatch.preflight_image_list_workers([{'id': 'worker'}]), {'worker': True})
|
||||
if value is True:
|
||||
self.assertEqual(await dispatch.preflight_image_list_workers(
|
||||
[{'id': 'worker'}], require_value_results=True), {'worker': True})
|
||||
else:
|
||||
with self.assertRaisesRegex(RuntimeError, 'value_results.*upgrade'):
|
||||
await dispatch.preflight_image_list_workers([{'id': 'worker'}], require_value_results=True)
|
||||
|
||||
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
|
||||
malformed = {"supported": True, "protocol_version": 2, "node_ids": None}
|
||||
session = _FakeSession({("worker", "capabilities"): (200, malformed)})
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "malformed.*capabilit"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self):
|
||||
incompatible = _valid_capability()
|
||||
incompatible["protocol_version"] = 1
|
||||
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "incompatible"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_preflight_fails_closed_for_missing_required_node(self):
|
||||
incompatible = _valid_capability()
|
||||
incompatible["node_ids"] = ["DistributedImageListStart"]
|
||||
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "incompatible"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_preflight_fails_closed_for_malformed_queue(self):
|
||||
responses = {
|
||||
("worker", "capabilities"): (200, _valid_capability()),
|
||||
("worker", "queue"): (200, {"queue_running": [], "queue_pending": None}),
|
||||
}
|
||||
session = _FakeSession(responses)
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "malformed /queue"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_preflight_does_not_hide_unexpected_errors(self):
|
||||
session = _FakeSession({("worker", "capabilities"): ValueError("programming error")})
|
||||
|
||||
async def fake_session():
|
||||
return session
|
||||
|
||||
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
|
||||
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(ValueError, "programming error"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
|
||||
class ImageListDirectWorkerResolutionTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.config = {"workers": [
|
||||
{"id": "enabled", "enabled": True, "host": "enabled-host", "port": 8188},
|
||||
{"id": "disabled", "enabled": False, "host": "disabled-host", "port": 8189},
|
||||
{"id": "serverless", "enabled": True, "type": "serverless"},
|
||||
{"id": "other", "enabled": True, "host": "other-host", "port": 8190},
|
||||
]}
|
||||
|
||||
def test_implicit_selection_uses_only_enabled_direct_workers(self):
|
||||
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, None)
|
||||
self.assertEqual([worker["id"] for worker in workers], ["enabled", "other"])
|
||||
|
||||
def test_explicit_selection_includes_exact_disabled_direct_worker(self):
|
||||
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, ["disabled"])
|
||||
self.assertEqual([worker["id"] for worker in workers], ["disabled"])
|
||||
|
||||
def test_explicit_empty_selection_returns_no_workers(self):
|
||||
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, []), [])
|
||||
|
||||
def test_unknown_explicit_id_raises_instead_of_falling_back(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "missing"):
|
||||
prompt_transform._resolve_direct_workers_for_image_list(self.config, ["enabled", "missing"])
|
||||
|
||||
def test_explicit_serverless_selection_is_excluded(self):
|
||||
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, ["serverless"]), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,341 @@
|
||||
"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from image_list.runtime import (
|
||||
END_CLASS,
|
||||
IMAGE_OUT,
|
||||
CONTEXT_OUT,
|
||||
START_CLASS,
|
||||
VALUE_OUT_BASE,
|
||||
ImageListCancelledError,
|
||||
build_worker_prompt,
|
||||
create_master_job,
|
||||
decode_item,
|
||||
decode_result,
|
||||
dynamic_schedule,
|
||||
encode_item,
|
||||
encode_result,
|
||||
normalize_image_list_inputs,
|
||||
normalize_worker_result,
|
||||
remove_master_job,
|
||||
set_master_result,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
def start_node(image_link):
|
||||
return {"class_type": START_CLASS, "inputs": {"images": image_link}}
|
||||
|
||||
|
||||
def graph_one_region():
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
|
||||
"3": start_node(["2", 0]),
|
||||
"4": {"class_type": "GenericFrameOperation", "inputs": {
|
||||
"image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", VALUE_OUT_BASE],
|
||||
}},
|
||||
"5": {"class_type": END_CLASS, "inputs": {
|
||||
"images": ["4", 0], "context": ["3", CONTEXT_OUT],
|
||||
}},
|
||||
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
||||
}
|
||||
|
||||
|
||||
def graph_two_sequential_regions():
|
||||
return {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
|
||||
"2": start_node(["1", 0]),
|
||||
"3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
|
||||
"4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
|
||||
"5": start_node(["4", 0]),
|
||||
"6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
|
||||
"7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
|
||||
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
|
||||
}
|
||||
|
||||
|
||||
class ImageListNormalizationTests(unittest.TestCase):
|
||||
def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
|
||||
images = [
|
||||
torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
|
||||
torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
|
||||
]
|
||||
refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
|
||||
masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
|
||||
torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
|
||||
rows = normalize_image_list_inputs(images, masks, {
|
||||
"value_1": [refs[:1], refs[1:]], "value_2": ["first", "second"],
|
||||
"value_3": [0.2], "value_4": [37],
|
||||
})
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
|
||||
self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
|
||||
self.assertTrue(torch.equal(rows[0]["values"]["value_1"], refs[:1]))
|
||||
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
|
||||
self.assertEqual([row["index"] for row in rows], [0, 1])
|
||||
self.assertEqual([row["values"]["value_2"] for row in rows], ["first", "second"])
|
||||
self.assertEqual([row["values"]["value_3"] for row in rows], [0.2, 0.2])
|
||||
self.assertEqual([row["values"]["value_4"] for row in rows], [37, 37])
|
||||
|
||||
def test_one_value_is_shared_and_n_values_pair_in_order(self):
|
||||
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
|
||||
rows = normalize_image_list_inputs(image, values={
|
||||
"value_1": ["same"], "value_2": [0.4], "value_3": [9], "value_4": [True]})
|
||||
self.assertEqual([row["values"]["value_1"] for row in rows], ["same"] * 3)
|
||||
self.assertEqual([row["values"]["value_2"] for row in rows], [0.4] * 3)
|
||||
self.assertEqual([row["values"]["value_3"] for row in rows], [9] * 3)
|
||||
self.assertEqual([row["values"]["value_4"] for row in rows], [True] * 3)
|
||||
rows = normalize_image_list_inputs(image, values={"value_2": [1, 2, 3]})
|
||||
self.assertEqual([row["values"] for row in rows], [{"value_2": 1}, {"value_2": 2}, {"value_2": 3}])
|
||||
with self.assertRaisesRegex(ValueError, "value_1 has 2 values; connect exactly 1.*or 3"):
|
||||
normalize_image_list_inputs(image, values={"value_1": ["one", "two"]})
|
||||
with self.assertRaisesRegex(ValueError, "mask has 2 values"):
|
||||
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)))
|
||||
shared_mask = normalize_image_list_inputs(image, mask=torch.ones((1, 2, 4)))
|
||||
self.assertTrue(all(torch.equal(row["mask"], torch.ones((1, 2, 4))) for row in shared_mask))
|
||||
|
||||
def test_value_batches_stay_whole_and_only_images_split(self):
|
||||
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
||||
pair = torch.rand((2, 4, 4, 3)) # e.g. two references for one item
|
||||
rows = normalize_image_list_inputs(images, values={"value_1": [pair]})
|
||||
self.assertTrue(all(torch.equal(row["values"]["value_1"], pair) for row in rows))
|
||||
per_item = [torch.rand((2, 4, 4, 3)), torch.rand((2, 4, 4, 3))]
|
||||
rows = normalize_image_list_inputs(images, values={"value_1": per_item})
|
||||
self.assertTrue(torch.equal(rows[1]["values"]["value_1"], per_item[1]))
|
||||
|
||||
def test_rejects_unsupported_values_and_counts(self):
|
||||
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
||||
with self.assertRaisesRegex(ValueError, "value_1 has 3 values"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [1, 2, 3]})
|
||||
with self.assertRaisesRegex(TypeError, "unsupported type dict"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [{"samples": 1}]})
|
||||
with self.assertRaisesRegex(TypeError, "floating-point"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [torch.zeros(2, dtype=torch.int64)]})
|
||||
with self.assertRaisesRegex(TypeError, "finite"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [float("nan")]})
|
||||
with self.assertRaisesRegex(ValueError, "Unknown Start value slot"):
|
||||
normalize_image_list_inputs(images, values={"value_11": [1]})
|
||||
with self.assertRaisesRegex(ValueError, "at least one image"):
|
||||
normalize_image_list_inputs(None)
|
||||
# Unconnected Autogrow slots arrive as None and are ignored.
|
||||
rows = normalize_image_list_inputs(images, values={"value_1": [None], "value_2": None})
|
||||
self.assertEqual([row["values"] for row in rows], [{}, {}])
|
||||
|
||||
def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
|
||||
item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
|
||||
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
|
||||
"values": {"value_1": torch.ones((2, 1, 3, 4), dtype=torch.float16),
|
||||
"value_2": "safe text", "value_3": 0.625, "value_4": 42,
|
||||
"value_5": False},
|
||||
"index": 3}
|
||||
restored = decode_item(encode_item(item))
|
||||
self.assertTrue(torch.equal(restored["image"], item["image"]))
|
||||
self.assertEqual(restored["image"].dtype, torch.float64)
|
||||
self.assertTrue(torch.equal(restored["values"]["value_1"], item["values"]["value_1"]))
|
||||
self.assertEqual(restored["values"]["value_1"].dtype, torch.float16)
|
||||
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
|
||||
self.assertEqual(restored["values"]["value_2"], "safe text")
|
||||
self.assertEqual(restored["values"]["value_3"], 0.625)
|
||||
self.assertIs(restored["values"]["value_4"], 42)
|
||||
self.assertIs(restored["values"]["value_5"], False)
|
||||
with self.assertRaisesRegex(ValueError, "unknown slot"):
|
||||
decode_item({**encode_item(item), "values": {"strength": {"type": "FLOAT", "value": 1.0}}})
|
||||
with self.assertRaisesRegex(ValueError, "unsupported type"):
|
||||
decode_item({**encode_item(item), "values": {"value_1": {"type": "PATH", "value": "/x"}}})
|
||||
with self.assertRaisesRegex(TypeError, "floating-point tensor"):
|
||||
encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
|
||||
with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
|
||||
from image_list.runtime import encode_value
|
||||
encode_value("/untrusted/path", "PATH")
|
||||
|
||||
def test_result_transport_preserves_ragged_result_shapes(self):
|
||||
result = normalize_worker_result(
|
||||
[torch.ones((1, 3, 5, 4), dtype=torch.float32)],
|
||||
[torch.ones((1, 3, 5), dtype=torch.float16)],
|
||||
["done"],
|
||||
)
|
||||
encoded = encode_result(result)
|
||||
decoded = decode_result(encoded)
|
||||
self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
|
||||
self.assertEqual(decoded["image"].dtype, torch.float32)
|
||||
self.assertEqual(decoded["mask"].dtype, torch.float16)
|
||||
self.assertEqual(decoded["text"], "done")
|
||||
with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
|
||||
normalize_worker_result(torch.zeros((2, 3, 4, 3)))
|
||||
|
||||
|
||||
class ImageListGraphBoundaryTests(unittest.TestCase):
|
||||
def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
|
||||
prompt = graph_one_region()
|
||||
region = validate_image_list_regions(prompt)[0]
|
||||
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
|
||||
"ctx", 1, "http://master.invalid")
|
||||
self.assertEqual(set(worker), {"1", "3", "4", "5"})
|
||||
self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
|
||||
self.assertNotIn("images", worker["3"]["inputs"])
|
||||
|
||||
def test_worker_prompt_drops_per_item_value_links(self):
|
||||
prompt = graph_one_region()
|
||||
prompt["8"] = {"class_type": "StringList", "inputs": {}}
|
||||
prompt["3"]["inputs"].update({"values.value_1": ["8", 0], "values.value_2": ["8", 0]})
|
||||
import sys, types
|
||||
from unittest.mock import patch
|
||||
registry = {"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL",)})}
|
||||
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
|
||||
region = validate_image_list_regions(prompt)[0]
|
||||
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012", "ctx", 0, "http://m")
|
||||
self.assertFalse([name for name in worker["3"]["inputs"] if name.startswith("values.")])
|
||||
self.assertNotIn("8", worker)
|
||||
self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
|
||||
self.assertNotIn("6", worker)
|
||||
|
||||
def test_rejects_upstream_processing_and_leaf_image_bypasses(self):
|
||||
# Both a leaf source and a generated batch must cross Start exactly once.
|
||||
for generated in (False, True):
|
||||
with self.subTest(generated=generated):
|
||||
prompt = graph_one_region()
|
||||
if generated:
|
||||
prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
|
||||
prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
|
||||
prompt["4"]["inputs"]["reference"] = ["2", 0]
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
|
||||
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
|
||||
def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import patch
|
||||
prompt = graph_one_region()
|
||||
prompt["7"] = {"class_type": "SeedConstant", "inputs": {"value": 9}}
|
||||
prompt["2"]["inputs"].update({"model": ["1", 0], "seed": ["7", 0]})
|
||||
prompt["4"]["inputs"]["shared_seed"] = ["7", 0]
|
||||
registry = {
|
||||
"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL", "CLIP", "VAE")}),
|
||||
"SeedConstant": type("Constant", (), {"RETURN_TYPES": ("INT",)}),
|
||||
}
|
||||
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
|
||||
region = validate_image_list_regions(prompt)[0]
|
||||
worker = build_worker_prompt(prompt, region, "token", "ctx", 0, "http://master")
|
||||
self.assertNotIn("2", worker)
|
||||
self.assertTrue({"1", "7"}.issubset(worker))
|
||||
|
||||
def test_unknown_or_list_producing_shared_dependency_is_not_exempt(self):
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import patch
|
||||
for node_class in (None, type("ListLoader", (), {
|
||||
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,)})):
|
||||
prompt = graph_one_region()
|
||||
prompt["2"]["inputs"]["model"] = ["1", 0]
|
||||
registry = {} if node_class is None else {"CheckpointLoaderSimple": node_class}
|
||||
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
|
||||
with self.assertRaisesRegex(ValueError, "bypasses Start.*1"):
|
||||
validate_image_list_regions(prompt)
|
||||
|
||||
def test_two_sequential_regions_are_valid_and_isolated(self):
|
||||
regions = validate_image_list_regions(graph_two_sequential_regions())
|
||||
self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
|
||||
p = graph_two_sequential_regions()
|
||||
second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
|
||||
"http://master.invalid")
|
||||
self.assertEqual(set(second), {"5", "6", "7"})
|
||||
|
||||
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
|
||||
prompt = graph_one_region()
|
||||
prompt["5"]["inputs"]["context"] = ["2", 0]
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "typed context output"):
|
||||
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
nested = graph_one_region()
|
||||
nested["7"] = start_node(["4", 0])
|
||||
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
|
||||
nested["5"]["inputs"]["images"] = ["8", 0]
|
||||
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
|
||||
validate_image_list_regions(nested, explicit_targets=explicit_targets)
|
||||
outside = graph_one_region()
|
||||
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "cross-boundary"):
|
||||
validate_image_list_regions(outside, explicit_targets=explicit_targets)
|
||||
|
||||
def test_rejects_duplicate_or_missing_results(self):
|
||||
prompt = graph_one_region()
|
||||
prompt["5"]["inputs"].pop("images")
|
||||
with self.assertRaisesRegex(ValueError, "images must be connected"):
|
||||
validate_image_list_regions(prompt)
|
||||
|
||||
|
||||
class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_fast_workers_get_more_items_and_order_is_stable(self):
|
||||
assignments = {"fast": [], "slow": []}
|
||||
async def dispatch(worker, index):
|
||||
assignments[worker].append(index)
|
||||
delay = 0.002 if worker == "fast" else 0.025
|
||||
task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
|
||||
return {"index": index, "task": task}
|
||||
async def wait_next(inflight):
|
||||
done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
|
||||
task = next(iter(done))
|
||||
worker = next(w for w, handle in inflight.items() if handle["task"] is task)
|
||||
index, result = task.result()
|
||||
return worker, index, result
|
||||
async def cleanup(_worker, handle):
|
||||
handle["task"].cancel()
|
||||
ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
|
||||
self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
|
||||
self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
|
||||
self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
|
||||
|
||||
async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
|
||||
async def dispatch(worker, index):
|
||||
return {"index": index, "prompt_id": f"{worker}-{index}"}
|
||||
cancelled = []
|
||||
async def cleanup(worker, handle):
|
||||
cancelled.append((worker, handle["prompt_id"]))
|
||||
async def timeout(_inflight):
|
||||
raise TimeoutError("worker timed out")
|
||||
with self.assertRaisesRegex(TimeoutError, "timed out"):
|
||||
await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
|
||||
self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
|
||||
|
||||
cancelled.clear()
|
||||
async def wrong_index(inflight):
|
||||
worker, handle = next(iter(inflight.items()))
|
||||
return worker, handle["index"] + 99, "bad"
|
||||
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
|
||||
await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
|
||||
self.assertEqual(len(cancelled), 1)
|
||||
|
||||
cancelled.clear()
|
||||
with self.assertRaises(ImageListCancelledError):
|
||||
await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
|
||||
interrupted=lambda: True)
|
||||
self.assertEqual(len(cancelled), 0)
|
||||
|
||||
async def test_no_workers_and_duplicate_results_fail_closed(self):
|
||||
async def unused(*_args):
|
||||
self.fail("must not dispatch without workers")
|
||||
with self.assertRaisesRegex(RuntimeError, "No free"):
|
||||
await dynamic_schedule(2, [], unused, unused, unused)
|
||||
async def dispatch(worker, index):
|
||||
return {"index": index}
|
||||
async def duplicate(inflight):
|
||||
worker, _handle = next(iter(inflight.items()))
|
||||
return worker, 0, "same"
|
||||
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
|
||||
await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,245 @@
|
||||
"""CPU-only execution-copy planning; no ComfyUI imports or prompt execution."""
|
||||
import copy
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from image_list.execution_plan import prepare_image_list_execution
|
||||
from image_list.runtime import (
|
||||
CONTEXT_OUT,
|
||||
END_CLASS,
|
||||
IMAGE_OUT,
|
||||
START_CLASS,
|
||||
build_worker_prompt,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
class OutputNode:
|
||||
OUTPUT_NODE = True
|
||||
|
||||
|
||||
class SchemaOutputNode:
|
||||
OUTPUT_NODE = False # V3's schema, not this fallback, is authoritative.
|
||||
|
||||
@classmethod
|
||||
def GET_NODE_INFO_V1(cls):
|
||||
return {"output_node": True, "output": ("IMAGE",)}
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registry(monkeypatch):
|
||||
classes = {
|
||||
END_CLASS: SchemaOutputNode,
|
||||
"SaveImage": OutputNode,
|
||||
"PreviewImage": OutputNode,
|
||||
"OutputUtility": OutputNode,
|
||||
"SchemaUtility": SchemaOutputNode,
|
||||
"ModelLoader": ModelLoader,
|
||||
}
|
||||
monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(NODE_CLASS_MAPPINGS=classes))
|
||||
return classes
|
||||
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
"loader": node("ModelLoader"),
|
||||
"source": node("ImageSource"),
|
||||
"start": node(START_CLASS, images=["source", 0]),
|
||||
"process": node("FrameOperation", image=["start", IMAGE_OUT], model=["loader", 0]),
|
||||
"end": node(END_CLASS, images=["process", 0], context=["start", CONTEXT_OUT]),
|
||||
"save": node("SaveImage", images=["end", 0]),
|
||||
}
|
||||
|
||||
|
||||
def test_discards_multistep_previews_but_preserves_other_output_roots(registry):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"debug-transform": node("DebugTransform", images=["start", IMAGE_OUT]),
|
||||
"debug-preview": node("PreviewImage", images=["debug-transform", 0]),
|
||||
"process-preview": node("PreviewImage", images=["process", 0]),
|
||||
"unrelated-source": node("ImageSource"),
|
||||
"unrelated-preview": node("PreviewImage", images=["unrelated-source", 0]),
|
||||
"upstream-preview": node("PreviewImage", images=["source", 0]),
|
||||
"post-transform": node("FrameOperation", images=["end", 0]),
|
||||
"post-preview": node("PreviewImage", images=["post-transform", 0]),
|
||||
})
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError, match="cross-boundary"):
|
||||
validate_image_list_regions(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {"debug-transform", "debug-preview", "process-preview"}
|
||||
assert targets == ["end", "save", "unrelated-preview", "upstream-preview", "post-preview"]
|
||||
validate_image_list_regions(prepared, explicit_targets=True)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", ["OutputUtility", "SchemaUtility"])
|
||||
def test_required_output_utility_keeps_complete_dependency_chain(registry, class_type):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"utility-config": node("Constant", value=7),
|
||||
"utility-first": node(class_type, images=["process", 0], config=["utility-config", 0]),
|
||||
"utility-second": node(class_type, images=["utility-first", 0]),
|
||||
"result": node("FrameOperation", image=["utility-second", 0]),
|
||||
"discarded-preview": node("PreviewImage", images=["utility-first", 0]),
|
||||
})
|
||||
prompt["end"]["inputs"]["images"] = ["result", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
registered = dict(registry)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {"discarded-preview"}
|
||||
assert targets == ["end", "save"]
|
||||
with pytest.raises(ValueError, match="Output node"):
|
||||
validate_image_list_regions(prepared)
|
||||
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
|
||||
worker = build_worker_prompt(prepared, region, "token", "context", 0, "http://master.invalid")
|
||||
assert set(worker) == set(prepared) - {"source", "save"}
|
||||
assert worker["utility-second"]["inputs"] == {"images": ["utility-first", 0]}
|
||||
assert prompt == before
|
||||
assert registry == registered
|
||||
assert OutputNode.OUTPUT_NODE is True
|
||||
assert SchemaOutputNode.OUTPUT_NODE is False
|
||||
assert SchemaOutputNode.GET_NODE_INFO_V1()["output_node"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["start", "process"])
|
||||
@pytest.mark.parametrize("direct", [False, True])
|
||||
def test_internal_bypass_into_post_end_computation_is_rejected(registry, source, direct):
|
||||
prompt = graph()
|
||||
prompt["post-process"] = node("Composite", images=["end", 0])
|
||||
if not direct:
|
||||
prompt["bypass"] = node("FrameOperation", image=[source, 0])
|
||||
source = "bypass"
|
||||
prompt["post-process"]["inputs"]["reference"] = [source, 0]
|
||||
prompt["save"]["inputs"]["images"] = ["post-process", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError, match="cross-boundary"):
|
||||
prepare_image_list_execution(prompt)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
def test_sequential_namespaced_regions_and_intermediate_output_are_preserved(registry):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"next-start": node(START_CLASS, images=["end", 0]),
|
||||
"next-process": node("SchemaUtility", images=["next-start", IMAGE_OUT], model=["loader", 0]),
|
||||
"next-end": node(END_CLASS, images=["next-process", 0], context=["next-start", CONTEXT_OUT]),
|
||||
"next-save": node("SaveImage", images=["next-end", 0]),
|
||||
"first-debug": node("PreviewImage", images=["process", 0]),
|
||||
"second-debug": node("PreviewImage", images=["next-process", 0]),
|
||||
})
|
||||
renamed = {key: f"group:7/{key}" for key in prompt}
|
||||
for spec in prompt.values():
|
||||
for name, value in spec["inputs"].items():
|
||||
if isinstance(value, list):
|
||||
spec["inputs"][name] = [renamed[value[0]], value[1]]
|
||||
prompt = {renamed[key]: spec for key, spec in prompt.items()}
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {renamed["first-debug"], renamed["second-debug"]}
|
||||
assert targets == [renamed[name] for name in ("end", "save", "next-end", "next-save")]
|
||||
regions = validate_image_list_regions(prepared, explicit_targets=True)
|
||||
assert {(r["start_id"], r["end_id"]) for r in regions} == {
|
||||
(renamed["start"], renamed["end"]), (renamed["next-start"], renamed["next-end"]),
|
||||
}
|
||||
second = next(r for r in regions if r["start_id"] == renamed["next-start"])
|
||||
worker = build_worker_prompt(prepared, second, "token", "context", 0, "http://master.invalid")
|
||||
assert set(worker) == {renamed[name] for name in ("loader", "next-start", "next-process", "next-end")}
|
||||
|
||||
|
||||
def test_shared_loader_and_its_independent_output_root_are_preserved(registry):
|
||||
registry["ModelLoader"] = type("OutputLoader", (ModelLoader, OutputNode), {})
|
||||
prompt = graph()
|
||||
prompt["source"]["inputs"]["model"] = ["loader", 0]
|
||||
prompt["independent-save"] = node("SaveImage", model=["loader", 0])
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == prompt
|
||||
assert targets == ["loader", "end", "save", "independent-save"]
|
||||
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
|
||||
assert "loader" in region["node_ids"]
|
||||
assert "loader" not in region["branch_ids"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("unsafe_class", [None, type("ListLoader", (), {
|
||||
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,),
|
||||
})])
|
||||
def test_unknown_or_list_producing_shared_loader_still_fails_closed(registry, unsafe_class):
|
||||
if unsafe_class is None:
|
||||
registry.pop("ModelLoader")
|
||||
else:
|
||||
registry["ModelLoader"] = unsafe_class
|
||||
prompt = graph()
|
||||
prompt["source"]["inputs"]["model"] = ["loader", 0]
|
||||
prompt["debug"] = node("PreviewImage", images=["process", 0])
|
||||
with pytest.raises(ValueError, match="bypasses Start.*loader"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
def test_upstream_per_item_bypass_is_not_hidden_by_preview_pruning(registry):
|
||||
prompt = graph()
|
||||
prompt["process"]["inputs"]["reference"] = ["source", 0]
|
||||
prompt["debug"] = node("PreviewImage", images=["process", 0])
|
||||
with pytest.raises(ValueError, match="bypasses Start.*source"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nested", [False, True])
|
||||
def test_overlapping_or_nested_regions_cannot_be_discarded_as_side_branches(registry, nested):
|
||||
prompt = graph()
|
||||
prompt["other-start"] = node(START_CLASS, images=["process" if nested else "source", 0])
|
||||
prompt["other-process"] = node("FrameOperation", image=["other-start", IMAGE_OUT], reference=["process", 0])
|
||||
prompt["other-end"] = node(END_CLASS, images=["other-process", 0], context=["other-start", CONTEXT_OUT])
|
||||
with pytest.raises(ValueError, match="Nested/ambiguous|cross-boundary|ambiguous processing"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid", ["missing-end", "wrong-context", "missing-images", "unrelated-images"])
|
||||
def test_pairing_and_result_validation_cannot_be_pruned_away(registry, invalid):
|
||||
prompt = graph()
|
||||
if invalid == "missing-end":
|
||||
del prompt["end"]
|
||||
elif invalid == "wrong-context":
|
||||
prompt["end"]["inputs"]["context"] = ["source", 0]
|
||||
elif invalid == "missing-images":
|
||||
del prompt["end"]["inputs"]["images"]
|
||||
else:
|
||||
prompt["end"]["inputs"]["images"] = ["source", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError):
|
||||
prepare_image_list_execution(prompt)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
def test_execution_copy_is_deep_and_repeated_preparation_is_stable(registry):
|
||||
prompt = graph()
|
||||
prompt["process"]["_meta"] = {"title": "Original", "tags": ["saved"]}
|
||||
before = copy.deepcopy(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
prepared_again, targets_again = prepare_image_list_execution(prepared)
|
||||
assert prepared_again == prepared
|
||||
assert targets_again == targets
|
||||
prepared["start"]["inputs"]["images"][0] = "changed"
|
||||
prepared["process"]["_meta"]["tags"].append("execution-only")
|
||||
assert prompt == before
|
||||
assert prepared_again == before
|
||||
|
||||
|
||||
def test_v3_false_output_flag_overrides_legacy_true(registry):
|
||||
registry["NotOutput"] = type("NotOutput", (), {
|
||||
"OUTPUT_NODE": True,
|
||||
"GET_NODE_INFO_V1": classmethod(lambda cls: {"output_node": False}),
|
||||
})
|
||||
prompt = graph()
|
||||
prompt["not-output"] = node("NotOutput", images=["end", 0])
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == prompt
|
||||
assert targets == ["end", "save"]
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Regression coverage at the scheduler and real HTTP transport boundaries."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
from image_list import runtime
|
||||
from test_image_list import graph_two_sequential_regions, start_node
|
||||
|
||||
|
||||
def load_routes():
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
name = '_image_list_transport_tests'
|
||||
routes = web.RouteTableDef()
|
||||
modules = {
|
||||
name: types.ModuleType(name),
|
||||
name + '.api': types.ModuleType(name + '.api'),
|
||||
name + '.image_list': types.ModuleType(name + '.image_list'),
|
||||
name + '.image_list.runtime': runtime,
|
||||
'server': types.SimpleNamespace(PromptServer=types.SimpleNamespace(
|
||||
instance=types.SimpleNamespace(routes=routes))),
|
||||
}
|
||||
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', root / 'api/image_list_routes.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec.loader.exec_module(module)
|
||||
return module, routes
|
||||
|
||||
|
||||
class CapabilityTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_capability_advertises_target_control_without_tensor_protocol_bump(self):
|
||||
module, _routes = load_routes()
|
||||
response = await module.image_list_capabilities_endpoint(None)
|
||||
payload = json.loads(response.text)
|
||||
self.assertIs(payload['execution_targets'], True)
|
||||
self.assertEqual(payload['protocol_version'], 2)
|
||||
|
||||
|
||||
class TransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.module, routes = load_routes()
|
||||
app = web.Application()
|
||||
app.add_routes(routes)
|
||||
self.client = TestClient(TestServer(app))
|
||||
await self.client.start_server()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.client.close()
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
async def test_fragmented_chunked_item_upload_is_consumable(self):
|
||||
token = runtime.create_worker_task_token()
|
||||
item = {'image': torch.full((1, 2, 3, 3), 0.25), 'index': 2}
|
||||
data = json.dumps({'token': token, 'context_id': 'context', 'item': runtime.encode_item(item)}).encode()
|
||||
async def chunks():
|
||||
yield data[:8]
|
||||
await asyncio.sleep(0.02)
|
||||
yield data[8:]
|
||||
response = await self.client.post('/distributed/image-list/item', data=chunks())
|
||||
self.assertEqual(response.status, 200, await response.text())
|
||||
result = runtime.consume_worker_item(token)
|
||||
self.assertTrue(torch.equal(result['image'], item['image']))
|
||||
self.assertEqual(result['index'], 2)
|
||||
|
||||
async def test_complete_item_upload_and_replay_rejection(self):
|
||||
token = runtime.create_worker_task_token()
|
||||
data = {'token': token, 'context_id': 'context', 'item': runtime.encode_item(
|
||||
{'image': torch.zeros((1, 2, 3, 3)), 'index': 0})}
|
||||
response = await self.client.post('/distributed/image-list/item', json=data)
|
||||
self.assertEqual(response.status, 200, await response.text())
|
||||
response = await self.client.post('/distributed/image-list/item', json=data)
|
||||
self.assertEqual(response.status, 400)
|
||||
runtime.consume_worker_item(token)
|
||||
with self.assertRaisesRegex(RuntimeError, 'already consumed'):
|
||||
runtime.consume_worker_item(token)
|
||||
|
||||
async def test_chunked_body_limit_is_aggregate(self):
|
||||
async def chunks():
|
||||
yield b'{"context_id":"'
|
||||
await asyncio.sleep(0.01)
|
||||
yield b'x' * 60
|
||||
yield b'"}'
|
||||
with patch.object(self.module, 'MAX_REQUEST_BYTES', 64):
|
||||
response = await self.client.post('/distributed/image-list/item', data=chunks())
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn('exceeds', await response.text())
|
||||
|
||||
async def test_invalid_result_shapes_are_rejected(self):
|
||||
for kind, shape in [('image', (1, 2, 3)), ('mask', (1, 2, 3, 1))]:
|
||||
result = runtime.encode_result({'image': torch.zeros((1, 2, 3, 3)),
|
||||
'mask': torch.zeros((1, 2, 3))})
|
||||
result[kind]['shape'] = list(shape)
|
||||
with self.assertRaises(ValueError):
|
||||
runtime.decode_result(result)
|
||||
|
||||
|
||||
class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_user_interrupt_during_result_wait_cleans_up_promptly(self):
|
||||
entered = asyncio.Event()
|
||||
cancelled = asyncio.Event()
|
||||
wait_cancelled = asyncio.Event()
|
||||
dispatches, cleanups = [], []
|
||||
async def dispatch(worker, index):
|
||||
dispatches.append(index)
|
||||
return {'index': index}
|
||||
async def wait_next(_inflight):
|
||||
entered.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
wait_cancelled.set()
|
||||
async def cleanup(worker, handle):
|
||||
cleanups.append((worker, handle['index']))
|
||||
task = asyncio.create_task(runtime.dynamic_schedule(
|
||||
3, ['one'], dispatch, wait_next, cleanup, cancelled.is_set))
|
||||
await entered.wait()
|
||||
cancelled.set()
|
||||
with self.assertRaises(runtime.ImageListCancelledError):
|
||||
await asyncio.wait_for(task, timeout=1)
|
||||
self.assertTrue(wait_cancelled.is_set())
|
||||
self.assertEqual(dispatches, [0])
|
||||
self.assertEqual(cleanups, [('one', 0)])
|
||||
|
||||
async def test_already_interrupted_does_not_submit_any_work(self):
|
||||
async def unused(*args):
|
||||
self.fail('pre-cancelled job must not dispatch or wait')
|
||||
with self.assertRaises(runtime.ImageListCancelledError):
|
||||
await runtime.dynamic_schedule(2, ['one'], unused, unused, unused, lambda: True)
|
||||
|
||||
async def test_task_cancel_during_submit_waits_for_owned_handle(self):
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
cleanups = []
|
||||
async def dispatch(worker, index):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return {'index': index, 'prompt_id': 'owned'}
|
||||
async def unused(*args):
|
||||
self.fail('cancelled dispatch must not enter result wait')
|
||||
async def cleanup(worker, handle):
|
||||
cleanups.append(handle['prompt_id'])
|
||||
task = asyncio.create_task(runtime.dynamic_schedule(2, ['one'], dispatch, unused, cleanup))
|
||||
await entered.wait()
|
||||
task.cancel()
|
||||
release.set()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await task
|
||||
self.assertEqual(cleanups, ['owned'])
|
||||
|
||||
def test_output_nodes_inside_region_are_rejected(self):
|
||||
prompt = graph_two_sequential_regions()
|
||||
output_class = type('OutputOperation', (), {'OUTPUT_NODE': True})
|
||||
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'FirstFrameOperation': output_class})
|
||||
with patch.dict(sys.modules, {'nodes': registry}):
|
||||
with self.assertRaisesRegex(ValueError, 'must be placed after'):
|
||||
runtime.validate_image_list_regions(prompt)
|
||||
|
||||
def test_shared_scalars_are_validated_before_dispatch(self):
|
||||
for value in (float('nan'), object(), torch.zeros(2, dtype=torch.int32)):
|
||||
with self.assertRaises(TypeError):
|
||||
runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), values={'value_1': [value]})
|
||||
|
||||
def test_three_sequential_regions_can_share_worker_loader(self):
|
||||
prompt = graph_two_sequential_regions()
|
||||
prompt.pop('8')
|
||||
prompt.update({
|
||||
'8': start_node(['7', 0]),
|
||||
'9': {'class_type': 'ThirdOperation', 'inputs': {'image': ['8', 0]}},
|
||||
'10': {'class_type': runtime.END_CLASS,
|
||||
'inputs': {'images': ['9', 0], 'context': ['8', runtime.CONTEXT_OUT]}},
|
||||
'11': {'class_type': 'GenericModelLoader', 'inputs': {}},
|
||||
})
|
||||
for node_id in ['3', '6', '9']:
|
||||
prompt[node_id]['inputs']['model'] = ['11', 0]
|
||||
# Shared dependencies are verified against the registered output schema.
|
||||
loader = type('GenericModelLoader', (), {'RETURN_TYPES': ('MODEL',)})
|
||||
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'GenericModelLoader': loader})
|
||||
with patch.dict(sys.modules, {'nodes': registry}):
|
||||
regions = runtime.validate_image_list_regions(prompt)
|
||||
self.assertEqual(len(regions), 3)
|
||||
for region in regions:
|
||||
self.assertIn('11', region['node_ids'])
|
||||
@@ -0,0 +1,254 @@
|
||||
"""Tests for periodic image-list server lifecycle cleanup."""
|
||||
import asyncio
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from image_list import lifecycle, runtime
|
||||
|
||||
|
||||
class ExecutionTargetProofTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
from test_image_list_master_runtime import graph
|
||||
self.prompt = graph()
|
||||
self.registry = SimpleNamespace(NODE_CLASS_MAPPINGS={
|
||||
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})})
|
||||
self.modules = patch.dict('sys.modules', {'nodes': self.registry})
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
|
||||
def server(self, roots, *, prompt=None, running=True):
|
||||
item = (0, 'owned', self.prompt if prompt is None else prompt, {}, roots, {})
|
||||
return SimpleNamespace(prompt_queue=SimpleNamespace(
|
||||
get_current_queue_volatile=lambda: ([item], []) if running else ([], [item])))
|
||||
|
||||
def test_outer_roots_in_running_queue_are_authoritative(self):
|
||||
self.assertTrue(lifecycle.has_safe_execution_targets(self.server(['end']), 'owned', self.prompt))
|
||||
for roots in (['op', 'end'], [], ['end', 'end'], None, [True]):
|
||||
with self.subTest(roots=roots):
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(roots), 'owned', self.prompt))
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end']), 'other', self.prompt))
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], running=False), 'owned', self.prompt))
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['end'], prompt={}), 'owned', self.prompt))
|
||||
|
||||
def test_prompt_marker_is_not_an_execution_guarantee(self):
|
||||
self.prompt['start']['inputs']['image_list_explicit_targets'] = True
|
||||
self.assertFalse(lifecycle.has_safe_execution_targets(self.server(['op', 'end']), 'owned', self.prompt))
|
||||
with self.assertRaisesRegex(ValueError, 'must be placed after'):
|
||||
runtime.create_master_job('start', 'end', self.prompt, [], [], '', master_participates=True)
|
||||
self.assertEqual(runtime._MASTER_JOBS, {})
|
||||
|
||||
def test_start_node_uses_running_queue_proof(self):
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
import torch
|
||||
|
||||
root = Path(__file__).resolve().parents[1]
|
||||
package = '_image_list_start_targets'
|
||||
modules = {}
|
||||
for suffix in ('', '.nodes', '.image_list', '.utils'):
|
||||
module = ModuleType(package + suffix)
|
||||
module.__path__ = [str(root / suffix.lstrip('.'))]
|
||||
modules[package + suffix] = module
|
||||
modules[package + '.image_list.runtime'] = runtime
|
||||
modules[package + '.image_list.lifecycle'] = lifecycle
|
||||
modules[package + '.utils.async_helpers'] = SimpleNamespace(run_async_in_server_loop=None)
|
||||
io = SimpleNamespace(ComfyNode=type('ComfyNode', (), {}), NodeOutput=lambda *args: args,
|
||||
**{name: object() for name in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean')})
|
||||
modules['comfy_api.v0_0_2'] = SimpleNamespace(io=io)
|
||||
modules['comfy_execution.utils'] = SimpleNamespace(
|
||||
get_executing_context=lambda: SimpleNamespace(prompt_id='owned'))
|
||||
server = SimpleNamespace(PromptServer=SimpleNamespace(instance=self.server(['end'])))
|
||||
modules['server'] = server
|
||||
spec = importlib.util.spec_from_file_location(package + '.nodes.image_list', root / 'nodes/image_list.py')
|
||||
node_module = importlib.util.module_from_spec(spec)
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec.loader.exec_module(node_module)
|
||||
start = node_module.DistributedImageListStart
|
||||
start.hidden = SimpleNamespace(prompt=self.prompt, unique_id='start')
|
||||
result = start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
|
||||
context = result[3][0]
|
||||
self.assertIsNotNone(runtime.get_master_job(context))
|
||||
runtime.remove_master_job(context)
|
||||
server.PromptServer.instance = self.server(['op', 'end'])
|
||||
with self.assertRaisesRegex(ValueError, 'must be placed after'):
|
||||
start.execute(images=torch.zeros((1, 2, 2, 3)), image_list_master_participates=True)
|
||||
self.assertEqual(runtime._MASTER_JOBS, {})
|
||||
|
||||
def test_owned_target_controlled_start_can_create_job(self):
|
||||
context = runtime.create_master_job('start', 'end', self.prompt, [], [], '',
|
||||
master_participates=True, explicit_targets=True)
|
||||
self.addCleanup(runtime.remove_master_job, context)
|
||||
self.assertIsNotNone(runtime.get_master_job(context))
|
||||
|
||||
|
||||
class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
with runtime._LOCK:
|
||||
runtime._MASTER_JOBS.clear()
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
def tearDown(self):
|
||||
with runtime._LOCK:
|
||||
runtime._MASTER_JOBS.clear()
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
@staticmethod
|
||||
def _state(prompt_id=None, *, active=False, age=10_000.0):
|
||||
return SimpleNamespace(
|
||||
prompt_id=prompt_id,
|
||||
active=active,
|
||||
created_at=time.monotonic() - age,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _server(queue):
|
||||
return SimpleNamespace(app=web.Application(), prompt_queue=queue)
|
||||
|
||||
async def _start(self, server, *, interval=0.01, grace=0.01):
|
||||
lifecycle.install_lifecycle(
|
||||
server,
|
||||
runtime,
|
||||
interval_seconds=interval,
|
||||
orphan_grace_seconds=grace,
|
||||
)
|
||||
runner = web.AppRunner(server.app)
|
||||
await runner.setup()
|
||||
return runner
|
||||
|
||||
async def test_unattended_sweep_removes_orphan_and_expired_staging(self):
|
||||
context_id = "orphan"
|
||||
with runtime._LOCK:
|
||||
cast(Any, runtime._MASTER_JOBS)[context_id] = self._state()
|
||||
runtime._WORKER_ITEMS["expired-token"] = (time.monotonic() - 1, {"image": "stale"})
|
||||
|
||||
queue = SimpleNamespace(
|
||||
get_current_queue_volatile=lambda: ([], []),
|
||||
)
|
||||
server = self._server(queue)
|
||||
swept = asyncio.Event()
|
||||
sweep_worker_items = runtime._sweep_worker_items
|
||||
|
||||
def track_worker_sweep():
|
||||
swept.set()
|
||||
sweep_worker_items()
|
||||
|
||||
with patch.object(runtime, "_sweep_worker_items", side_effect=track_worker_sweep):
|
||||
runner = await self._start(server)
|
||||
try:
|
||||
await asyncio.wait_for(swept.wait(), timeout=0.5)
|
||||
self.assertNotIn(context_id, runtime._MASTER_JOBS)
|
||||
self.assertNotIn("expired-token", runtime._WORKER_ITEMS)
|
||||
with self.assertRaisesRegex(ValueError, "Unknown or expired"):
|
||||
runtime.set_master_result(context_id, 0, "late-token", {})
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
async def test_old_active_and_queued_contexts_survive_but_orphans_expire(self):
|
||||
states = {
|
||||
"active": self._state("active-prompt", active=True),
|
||||
"queued": self._state("queued-prompt"),
|
||||
"orphan": self._state("missing-prompt"),
|
||||
"unbound": self._state(None),
|
||||
}
|
||||
with runtime._LOCK:
|
||||
cast(Any, runtime._MASTER_JOBS).update(states)
|
||||
queue = SimpleNamespace(
|
||||
get_current_queue_volatile=lambda: (
|
||||
[[1, "active-prompt", {}]],
|
||||
[(2, "queued-prompt", {}, {}, None, None)],
|
||||
),
|
||||
)
|
||||
runner = await self._start(self._server(queue))
|
||||
try:
|
||||
await asyncio.sleep(0.04)
|
||||
self.assertIs(runtime._MASTER_JOBS.get("active"), states["active"])
|
||||
self.assertIs(runtime._MASTER_JOBS.get("queued"), states["queued"])
|
||||
self.assertNotIn("orphan", runtime._MASTER_JOBS)
|
||||
self.assertNotIn("unbound", runtime._MASTER_JOBS)
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
async def test_falls_back_to_nonvolatile_queue_reader(self):
|
||||
state = self._state("queued-prompt")
|
||||
with runtime._LOCK:
|
||||
cast(Any, runtime._MASTER_JOBS)["queued"] = state
|
||||
fallback_reads = 0
|
||||
|
||||
def volatile_failure():
|
||||
raise RuntimeError("volatile view unavailable")
|
||||
|
||||
def current_queue():
|
||||
nonlocal fallback_reads
|
||||
fallback_reads += 1
|
||||
return ([], [[1, "queued-prompt", {}]])
|
||||
|
||||
queue = SimpleNamespace(
|
||||
get_current_queue_volatile=volatile_failure,
|
||||
get_current_queue=current_queue,
|
||||
)
|
||||
runner = await self._start(self._server(queue))
|
||||
try:
|
||||
await asyncio.sleep(0.04)
|
||||
self.assertGreaterEqual(fallback_reads, 1)
|
||||
self.assertIs(runtime._MASTER_JOBS.get("queued"), state)
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
async def test_queue_read_error_protects_bound_contexts(self):
|
||||
with runtime._LOCK:
|
||||
cast(Any, runtime._MASTER_JOBS)["bound"] = self._state("bound-prompt")
|
||||
cast(Any, runtime._MASTER_JOBS)["unbound"] = self._state(None)
|
||||
|
||||
def fail_read():
|
||||
raise RuntimeError("queue unavailable")
|
||||
|
||||
queue = SimpleNamespace(
|
||||
get_current_queue_volatile=fail_read,
|
||||
get_current_queue=fail_read,
|
||||
)
|
||||
runner = await self._start(self._server(queue))
|
||||
try:
|
||||
await asyncio.sleep(0.04)
|
||||
self.assertIn("bound", runtime._MASTER_JOBS)
|
||||
self.assertNotIn("unbound", runtime._MASTER_JOBS)
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
async def test_install_and_removal_are_idempotent_and_task_stops_on_cleanup(self):
|
||||
queue_reads = 0
|
||||
|
||||
def read_queue():
|
||||
nonlocal queue_reads
|
||||
queue_reads += 1
|
||||
return ([], [])
|
||||
|
||||
server = self._server(SimpleNamespace(get_current_queue_volatile=read_queue))
|
||||
lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
|
||||
lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
|
||||
self.assertEqual(len(server.app.cleanup_ctx), 1)
|
||||
|
||||
with runtime._LOCK:
|
||||
cast(Any, runtime._MASTER_JOBS)["removed"] = self._state()
|
||||
runtime.remove_master_job("removed")
|
||||
runtime.remove_master_job("removed")
|
||||
self.assertNotIn("removed", runtime._MASTER_JOBS)
|
||||
|
||||
runner = web.AppRunner(server.app)
|
||||
await runner.setup()
|
||||
await asyncio.sleep(0.035)
|
||||
self.assertGreaterEqual(queue_reads, 1)
|
||||
await runner.cleanup()
|
||||
reads_after_cleanup = queue_reads
|
||||
await asyncio.sleep(0.035)
|
||||
self.assertEqual(queue_reads, reads_after_cleanup)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Run the real framework's CPU-only private-cache retention regression."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_native_local_cache_retention():
|
||||
core = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not core:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT for CPU-only real-framework validation')
|
||||
helper = Path(__file__).parent / 'helpers' / 'local_cache_retention_check.py'
|
||||
env = dict(os.environ, CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='')
|
||||
env.pop('PYTHONPATH', None)
|
||||
result = subprocess.run([sys.executable, str(helper), core], env=env,
|
||||
text=True, capture_output=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'PASS: cross-run native cache reuse' in result.stdout
|
||||
@@ -0,0 +1,506 @@
|
||||
"""Master slot regressions, plus an opt-in real CPU ComfyUI subprocess probe."""
|
||||
import asyncio
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
'source': {'class_type': 'Source', 'inputs': {}},
|
||||
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
|
||||
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['op', 0], 'context': ['start', 3]}},
|
||||
}
|
||||
|
||||
|
||||
class Response:
|
||||
status = 200
|
||||
|
||||
def __init__(self, data=None):
|
||||
self.data = data
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
async def json(self):
|
||||
return self.data
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
|
||||
class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
name = '_master_slot_tests'
|
||||
self.modules = patch.dict(sys.modules)
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = module
|
||||
self.rt = importlib.import_module(name + '.image_list.runtime')
|
||||
self.workers = []
|
||||
self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
|
||||
self.dispatched = []
|
||||
self.local_calls = []
|
||||
self.tokens = []
|
||||
self.closed = 0
|
||||
self.local_error = None
|
||||
self.local_delay = 0.005
|
||||
self.interrupted = False
|
||||
self.local_started = asyncio.Event()
|
||||
self.local_finished = asyncio.Event()
|
||||
self.deadline = 0
|
||||
self.slot_instances = []
|
||||
self.config = {'workers': self.workers, 'settings': {}}
|
||||
fixture = self
|
||||
|
||||
class FakeSlot:
|
||||
def __init__(self, model_callback=None):
|
||||
fixture.slot_instances.append(self)
|
||||
self.tasks = []
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
|
||||
async def execute():
|
||||
inputs = prompt[end_id]['inputs']
|
||||
index = inputs['worker_item_index']
|
||||
token = inputs['worker_result_token']
|
||||
fixture.local_calls.append(index)
|
||||
fixture.tokens.append(token)
|
||||
fixture.local_started.set()
|
||||
try:
|
||||
await asyncio.sleep(fixture.local_delay)
|
||||
if fixture.local_error:
|
||||
raise fixture.local_error
|
||||
item = fixture.rt.consume_worker_item(token)
|
||||
fixture.rt.set_master_result(fixture.context, index, token, {
|
||||
'image': item['image'], 'mask': item.get('mask'),
|
||||
'text': item['values'].get('value_1')})
|
||||
finally:
|
||||
fixture.local_finished.set()
|
||||
task = asyncio.create_task(execute())
|
||||
self.tasks.append(task)
|
||||
return task
|
||||
|
||||
def request_stop(self):
|
||||
pass
|
||||
|
||||
async def close(self):
|
||||
# Model a synchronous node that must finish before its thread can join.
|
||||
await asyncio.gather(*self.tasks, return_exceptions=True)
|
||||
fixture.closed += 1
|
||||
|
||||
dependencies = {
|
||||
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=FakeSlot),
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=self.preflight,
|
||||
dispatch_worker_prompt=AsyncMock(side_effect=self.dispatch)),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda config, ids: self.workers),
|
||||
'.api.orchestration.media_sync': types.SimpleNamespace(
|
||||
sync_worker_media=AsyncMock(), fetch_worker_path_separator=AsyncMock(return_value=None)),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
|
||||
get_worker_timeout_seconds=lambda: 60,
|
||||
get_image_list_job_timeout_seconds=lambda: self.deadline),
|
||||
'.utils.network': types.SimpleNamespace(
|
||||
build_master_callback_url=lambda *a, **kw: 'http://master',
|
||||
build_worker_url=lambda worker, path: 'http://remote' + path,
|
||||
get_client_session=AsyncMock(return_value=self)),
|
||||
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
sys.modules[name + suffix] = module
|
||||
self.network = dependencies['.utils.network']
|
||||
sys.modules['aiohttp'] = types.SimpleNamespace(ClientTimeout=lambda **kw: kw, ClientError=ConnectionError)
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
|
||||
sys.modules['comfy'] = comfy
|
||||
sys.modules['comfy.model_management'] = comfy.model_management
|
||||
|
||||
def create(self, participates=True, count=5):
|
||||
self.items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
|
||||
'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)},
|
||||
'index': i}
|
||||
for i in range(count)]
|
||||
self.context = self.rt.create_master_job('start', 'end', graph(), self.items,
|
||||
[w['id'] for w in self.workers], '', master_participates=participates)
|
||||
self.state = self.rt.get_master_job(self.context)
|
||||
return self.context
|
||||
|
||||
def post(self, url, *, data=None, **kwargs):
|
||||
if url.endswith('/distributed/image-list/item'):
|
||||
payload = json.loads(data)
|
||||
self.rt.store_worker_item(payload['token'], payload['item'])
|
||||
return Response({})
|
||||
|
||||
def delete(self, url, **kwargs):
|
||||
self.rt.delete_worker_item(url.rsplit('/', 1)[-1])
|
||||
return Response({})
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
pid = url.rsplit('/', 1)[-1]
|
||||
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
|
||||
|
||||
async def dispatch(self, worker, prompt, meta, **kwargs):
|
||||
inputs = prompt['end']['inputs']
|
||||
index, token = inputs['worker_item_index'], inputs['worker_result_token']
|
||||
self.dispatched.append(index)
|
||||
# A remote callback is ready immediately; the local slot is slower.
|
||||
item = self.rt.consume_worker_item(token)
|
||||
self.rt.set_master_result(self.context, index, token, {
|
||||
'image': item['image'], 'mask': item['mask'], 'text': item['values']['value_1']})
|
||||
return f'remote-{index}'
|
||||
|
||||
async def test_local_only_preserves_order_and_uses_no_http(self):
|
||||
result = await self.rt.run_image_list_job(self.create())
|
||||
self.assertEqual(result[2], [str(i) for i in range(5)])
|
||||
self.assertEqual(self.local_calls, list(range(5)))
|
||||
self.assertEqual(len(self.slot_instances), 1)
|
||||
self.assertEqual(len(set(self.tokens)), 5)
|
||||
self.network.get_client_session.assert_not_awaited()
|
||||
self.assertEqual(self.closed, 1)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
for i, (image, mask) in enumerate(zip(*result[:2])):
|
||||
self.assertTrue(torch.equal(image, self.items[i]['image']))
|
||||
self.assertTrue(torch.equal(mask, self.items[i]['mask']))
|
||||
|
||||
async def test_mixed_pool_claims_from_one_queue_and_relaxes_remote_idle(self):
|
||||
self.workers.append({'id': 'master', 'type': 'local'}) # Not the in-process master slot.
|
||||
self.local_delay = 0.04
|
||||
result = await self.rt.run_image_list_job(self.create(count=9))
|
||||
self.assertTrue(self.local_calls)
|
||||
self.assertTrue(self.dispatched)
|
||||
self.assertEqual(sorted(self.local_calls + self.dispatched), list(range(9)))
|
||||
self.assertEqual(result[2], [str(i) for i in range(9)])
|
||||
self.preflight.assert_awaited_once_with(self.workers, require_idle=False)
|
||||
|
||||
async def test_orchestrator_mode_never_instantiates_local_slot(self):
|
||||
self.workers.append({'id': 'worker', 'type': 'remote'})
|
||||
await self.rt.run_image_list_job(self.create(participates=False))
|
||||
self.assertFalse(self.slot_instances)
|
||||
self.preflight.assert_awaited_once_with(self.workers, require_idle=True)
|
||||
with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
|
||||
self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
|
||||
|
||||
async def test_values_reconstructed_in_source_order_without_transporting_originals(self):
|
||||
self.workers.append({'id': 'remote', 'type': 'remote'})
|
||||
context = self.create(count=6)
|
||||
tensor = torch.ones(4, 2, 2, 3, dtype=torch.float64)
|
||||
for index, item in enumerate(self.items):
|
||||
item['values'].update(value_2=tensor, value_3=True, value_4=index / 8)
|
||||
self.local_delay = 0.1 # Remote results arrive ahead of master item zero.
|
||||
with patch.object(self.rt, 'encode_result', side_effect=AssertionError('originals must not return')):
|
||||
result = await self.rt.run_image_list_job(context, include_values=True)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
self.assertEqual(len(result), 4)
|
||||
for index, row in enumerate(result[3]):
|
||||
self.assertEqual(row['value_1'], str(index))
|
||||
self.assertIs(row['value_2'], tensor)
|
||||
self.assertIs(row['value_3'], True)
|
||||
self.assertEqual(row['value_4'], index / 8)
|
||||
self.assertTrue(self.dispatched and self.local_calls)
|
||||
|
||||
async def test_overrides_merge_add_and_recheck_capability_before_every_item(self):
|
||||
self.workers.append({'id': 'remote', 'type': 'remote'})
|
||||
context = self.create(participates=False, count=3)
|
||||
self.state.original_prompt['end']['inputs'].update(value_1=['start', 4], value_10=['start', 2])
|
||||
original_set = self.rt.set_master_result
|
||||
def override(context, index, token, result):
|
||||
result['values'] = {'value_1': f'new-{index}', 'value_10': index}
|
||||
return original_set(context, index, token, result)
|
||||
with patch.object(self.rt, 'set_master_result', side_effect=override):
|
||||
result = await self.rt.run_image_list_job(context, include_values=True)
|
||||
self.assertEqual(result[3], [{'value_1': f'new-{index}', 'value_10': index} for index in range(3)])
|
||||
self.assertEqual(self.preflight.await_count, 4)
|
||||
self.assertTrue(all(call.kwargs['require_value_results'] is True
|
||||
for call in self.preflight.await_args_list))
|
||||
self.assertEqual(self.items[0]['values'], {'value_1': '0'})
|
||||
|
||||
async def test_override_capability_downgrade_stops_before_upload(self):
|
||||
self.workers.append({'id': 'remote', 'type': 'remote'})
|
||||
context = self.create(participates=False)
|
||||
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
|
||||
self.preflight.side_effect = [{'remote': True}, RuntimeError('value_results unsupported')]
|
||||
with patch.object(self, 'post', wraps=self.post) as upload:
|
||||
with self.assertRaisesRegex(RuntimeError, 'value_results'):
|
||||
await self.rt.run_image_list_job(context, include_values=True)
|
||||
upload.assert_not_called()
|
||||
self.assertFalse(self.dispatched or self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_missing_expected_override_fails_without_partial_output(self):
|
||||
self.workers.append({'id': 'remote', 'type': 'remote'})
|
||||
context = self.create(participates=False)
|
||||
self.state.original_prompt['end']['inputs']['value_1'] = ['start', 4]
|
||||
with self.assertRaisesRegex(ValueError, 'missing'):
|
||||
await self.rt.run_image_list_job(context, include_values=True)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_values_failure_and_cancellation_return_no_partial_lists(self):
|
||||
self.local_error = RuntimeError('branch failed')
|
||||
with self.assertRaisesRegex(RuntimeError, 'branch failed'):
|
||||
await self.rt.run_image_list_job(self.create(), include_values=True)
|
||||
self.local_error = None
|
||||
self.local_started.clear()
|
||||
self.local_delay = 0.1
|
||||
task = asyncio.create_task(self.rt.run_image_list_job(self.create(), include_values=True))
|
||||
await self.local_started.wait()
|
||||
task.cancel()
|
||||
with self.assertRaises(self.rt.ImageListCancelledError):
|
||||
await task
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_local_failure_surfaces_and_reclaims_staged_token(self):
|
||||
self.local_error = RuntimeError('local branch failed')
|
||||
with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
|
||||
await self.rt.run_image_list_job(self.create())
|
||||
self.assertEqual(self.closed, 1)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_deadline_joins_local_execution_before_removing_context(self):
|
||||
self.deadline, self.local_delay = 0.03, 0.1
|
||||
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
||||
await self.rt.run_image_list_job(self.create())
|
||||
self.assertTrue(self.local_finished.is_set())
|
||||
self.assertEqual(self.closed, 1)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
|
||||
async def test_cancelled_scheduler_joins_owned_execution_even_if_cancelled_again(self):
|
||||
self.local_delay = 0.12
|
||||
task = asyncio.create_task(self.rt.run_image_list_job(self.create()))
|
||||
await self.local_started.wait()
|
||||
task.cancel()
|
||||
await asyncio.sleep(0.02)
|
||||
task.cancel()
|
||||
with self.assertRaises(self.rt.ImageListCancelledError):
|
||||
await task
|
||||
self.assertTrue(self.local_finished.is_set())
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
|
||||
class RealFrameworkTests(unittest.TestCase):
|
||||
@unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for the real CPU framework probe')
|
||||
def test_real_framework_local_slot(self):
|
||||
result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--real-probe',
|
||||
os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=60)
|
||||
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
||||
self.assertIn('MASTER_SLOT_REAL_OK', result.stdout)
|
||||
|
||||
|
||||
def real_framework_probe(comfy_root):
|
||||
import threading
|
||||
import time
|
||||
sys.path.insert(0, str(Path(comfy_root).resolve()))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
import comfy.memory_management as memory
|
||||
import comfy.model_management as models
|
||||
import comfy.utils as utils
|
||||
from comfy_execution import progress
|
||||
from comfy_execution.graph_utils import GraphBuilder
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
name = '_real_master_slot'
|
||||
for suffix in ('', '.image_list', '.utils', '.nodes'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = module
|
||||
# No HTTP/service/extension bootstrap is needed by the local slot. Fail if
|
||||
# either boundary tries to use a server-loop bridge instead of direct memory.
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('local boundary used the server loop or /prompt')
|
||||
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
|
||||
config = {'settings': {}}
|
||||
sys.modules[name + '.utils.config'] = types.SimpleNamespace(load_config=lambda: config)
|
||||
rt = importlib.import_module(name + '.image_list.runtime')
|
||||
local = importlib.import_module(name + '.image_list.local_executor')
|
||||
boundary = importlib.import_module(name + '.nodes.image_list')
|
||||
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
|
||||
DistributedImageListEnd=boundary.DistributedImageListEnd)
|
||||
loads, calls = [], []
|
||||
entered, stopped = threading.Event(), threading.Event()
|
||||
mode = ['normal']
|
||||
model = object()
|
||||
|
||||
class Loader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('MODEL',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
loads.append(True)
|
||||
return (model,)
|
||||
|
||||
class Operation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'model': ('MODEL',),
|
||||
'mask': ('MASK',), 'text': ('STRING', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
async def run(self, image, model, mask, text):
|
||||
calls.append(text)
|
||||
entered.set()
|
||||
if mode[0] == 'error':
|
||||
raise ValueError('probe branch failure')
|
||||
if mode[0] in ('cancel', 'interrupt'):
|
||||
try:
|
||||
while True:
|
||||
utils.ProgressBar(1).update(1)
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
await asyncio.sleep(0.02)
|
||||
stopped.set()
|
||||
if mode[0] == 'sync':
|
||||
# No progress/interrupt checks: cleanup still must join, not
|
||||
# release the parent just because an asyncio Future is cancelled.
|
||||
time.sleep(0.08)
|
||||
stopped.set()
|
||||
await asyncio.sleep(0)
|
||||
return image + 1, mask + 1, text + ':done'
|
||||
|
||||
class Expanded:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return Operation.INPUT_TYPES()
|
||||
RETURN_TYPES = Operation.RETURN_TYPES
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, model, mask, text):
|
||||
builder = GraphBuilder()
|
||||
node = builder.node('Operation', image=image, model=model, mask=mask, text=text)
|
||||
return {'result': (node.out(0), node.out(1), node.out(2)), 'expand': builder.finalize()}
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(Loader=Loader, Operation=Operation, Expanded=Expanded)
|
||||
prompt = graph()
|
||||
prompt['loader'] = {'class_type': 'Loader', 'inputs': {}}
|
||||
prompt['op']['inputs'].update(model=['loader', 0], mask=['start', rt.MASK_OUT],
|
||||
text=['start', rt.VALUE_OUT_BASE])
|
||||
prompt['end']['inputs'].update(masks=['op', 1], texts=['op', 2])
|
||||
items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
|
||||
'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)}, 'index': i}
|
||||
for i in range(3)]
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
parent_progress = progress.ProgressRegistry('outer', execution.DynamicPrompt({}))
|
||||
parent_progress.nodes['outer'] = {'state': progress.NodeState.Running, 'value': 3, 'max': 7}
|
||||
progress.global_progress_registry = parent_progress
|
||||
parent_hook = lambda *a, **kw: None
|
||||
utils.PROGRESS_BAR_HOOK = parent_hook
|
||||
ram_callback = lambda *a, **kw: None
|
||||
memory.set_ram_cache_release_state(ram_callback, 123)
|
||||
GraphBuilder.set_default_prefix('outer-prefix', 7, 9)
|
||||
|
||||
def restored():
|
||||
assert progress.global_progress_registry is parent_progress
|
||||
assert parent_progress.nodes == {'outer': {'state': progress.NodeState.Running, 'value': 3, 'max': 7}}
|
||||
assert utils.PROGRESS_BAR_HOOK is parent_hook
|
||||
assert memory.extra_ram_release_callback is ram_callback and memory.RAM_CACHE_HEADROOM == 123
|
||||
assert (GraphBuilder._default_prefix_root, GraphBuilder._default_prefix_call_index,
|
||||
GraphBuilder._default_prefix_graph_index) == ('outer-prefix', 7, 9)
|
||||
assert not any(t.name.startswith('image-list-master') for t in threading.enumerate())
|
||||
|
||||
async def probe():
|
||||
# Verify actual V3 input parsing and list mapping, not only direct execute.
|
||||
for supplied, delegated, expected in [(None, False, True), (None, True, False),
|
||||
(False, False, False), (True, True, True)]:
|
||||
config['settings']['master_delegate_only'] = delegated
|
||||
inputs = {'image_list_enabled_worker_ids': '["remote"]'}
|
||||
if supplied is not None:
|
||||
inputs['image_list_master_participates'] = supplied
|
||||
cls = boundary.DistributedImageListStart
|
||||
data, _, v3 = execution.get_input_data(inputs, cls, 'start', dynprompt=execution.DynamicPrompt(prompt))
|
||||
data['images'] = [torch.zeros((3, 2, 2, 3))]
|
||||
result = await execution._async_map_node_over_list('outer', 'start', cls, data, cls.FUNCTION, v3_data=v3)
|
||||
result = await execution.resolve_map_node_over_list_results(result)
|
||||
context = result[0].result[rt.CONTEXT_OUT][0]
|
||||
assert rt.get_master_job(context).master_participates is expected
|
||||
rt.remove_master_job(context)
|
||||
|
||||
for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
|
||||
# Independent lifecycle cases start cold; retention is tested separately.
|
||||
local._RETAINED_OUTPUTS = None
|
||||
mode[0] = behavior
|
||||
prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
|
||||
loads.clear()
|
||||
calls.clear()
|
||||
entered.clear()
|
||||
stopped.clear()
|
||||
tracked = []
|
||||
slot = local.LocalExecutionSlot(model_callback=tracked.append)
|
||||
context = rt.create_master_job('start', 'end', prompt, items, [], '', master_participates=True)
|
||||
try:
|
||||
for index in range(3 if behavior in ('normal', 'expand') else 1):
|
||||
token = rt.create_worker_task_token()
|
||||
rt.register_assignment(context, index, 'local', token)
|
||||
rt._WORKER_ITEMS[token] = (time.monotonic() + 60, items[index])
|
||||
worker_prompt = rt.build_worker_prompt(prompt, region, token, context, index, '')
|
||||
for key in ('start', 'end'):
|
||||
worker_prompt[key]['inputs']['image_list_local_execution'] = True
|
||||
if behavior == 'already-interrupted':
|
||||
models.interrupt_current_processing(True)
|
||||
future = slot.submit(worker_prompt, f'local-{index}', 'end')
|
||||
if behavior in ('cancel', 'interrupt', 'sync'):
|
||||
assert await asyncio.to_thread(entered.wait, 5)
|
||||
if behavior == 'interrupt':
|
||||
models.interrupt_current_processing(True)
|
||||
else:
|
||||
slot.request_stop()
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.shield(asyncio.wrap_future(future)), 5)
|
||||
except rt.ImageListCancelledError:
|
||||
assert behavior in ('cancel', 'interrupt', 'sync', 'already-interrupted'), behavior
|
||||
except RuntimeError as exc:
|
||||
assert behavior == 'error' and 'probe branch failure' in str(exc), str(exc)
|
||||
else:
|
||||
assert behavior in ('normal', 'expand'), behavior
|
||||
result = rt.get_master_job(context).results[index]
|
||||
assert torch.equal(result['image'], items[index]['image'] + 1)
|
||||
assert torch.equal(result['mask'], items[index]['mask'] + 1)
|
||||
assert result['text'] == f'{index}:done'
|
||||
if behavior in ('normal', 'expand'):
|
||||
assert len(loads) == 1 and calls == ['0', '1', '2'], (loads, calls)
|
||||
assert any(model in outputs for outputs in tracked), tracked
|
||||
if behavior in ('cancel', 'interrupt', 'sync'):
|
||||
assert stopped.is_set(), behavior
|
||||
if behavior == 'already-interrupted':
|
||||
assert models.processing_interrupted() and not calls
|
||||
finally:
|
||||
await slot.close()
|
||||
for token in rt.get_master_job(context).expected_tokens.values():
|
||||
rt.delete_worker_item(token)
|
||||
rt.remove_master_job(context)
|
||||
models.interrupt_current_processing(False) # Only test setup/teardown.
|
||||
restored()
|
||||
assert not rt._MASTER_JOBS and not rt._WORKER_ITEMS
|
||||
asyncio.run(probe())
|
||||
print('MASTER_SLOT_REAL_OK V3 metadata; cache reuse; async/expanded graph; errors; cancellation; state restoration')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if len(sys.argv) > 1 and sys.argv[1] == '--real-probe':
|
||||
real_framework_probe(sys.argv[2])
|
||||
else:
|
||||
unittest.main()
|
||||
@@ -0,0 +1,286 @@
|
||||
"""Image-list participation uses the existing master orchestrator setting."""
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from test_serverless_integration import PKG, load_modules
|
||||
|
||||
|
||||
class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.q, modules, self.ps = load_modules()
|
||||
self.modules_patch = patch.dict(sys.modules, modules)
|
||||
self.modules_patch.start()
|
||||
self.addCleanup(self.modules_patch.stop)
|
||||
self.worker = {"id": "worker", "name": "Worker", "type": "local", "enabled": True}
|
||||
self.config = {"workers": [self.worker], "settings": {}}
|
||||
self.q.load_config = lambda: self.config
|
||||
# Model the legacy helper's fallback: image-list must not inherit it.
|
||||
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw:
|
||||
(workers, delegate if workers else False))
|
||||
self.preflight = AsyncMock(return_value={"worker": True})
|
||||
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
|
||||
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
|
||||
self.prompt = {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 3]}},
|
||||
}
|
||||
|
||||
async def queue(self, **kwargs):
|
||||
return await self.q.orchestrate_distributed_execution(self.prompt, None, "client", **kwargs)
|
||||
|
||||
def prepared_inputs(self):
|
||||
return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
|
||||
|
||||
async def test_prepared_master_prunes_side_outputs_without_mutating_input(self):
|
||||
import copy
|
||||
self.prompt['op'] = {'class_type': 'Flagged', 'inputs': {'image': ['2', 0]}}
|
||||
self.prompt['3']['inputs']['images'] = ['op', 0]
|
||||
self.prompt['preview'] = {'class_type': 'Flagged', 'inputs': {'image': ['op', 0]}}
|
||||
self.prompt['save'] = {'class_type': 'Flagged', 'inputs': {'image': ['3', 0]}}
|
||||
original = copy.deepcopy(self.prompt)
|
||||
registry = type('Registry', (), {'NODE_CLASS_MAPPINGS': {
|
||||
'Flagged': type('Flagged', (), {'OUTPUT_NODE': True})}})
|
||||
with patch.dict(sys.modules, {'nodes': registry}):
|
||||
await self.queue()
|
||||
queued = self.q.queue_prompt_payload.call_args
|
||||
self.assertEqual(queued.kwargs['partial_execution_targets'], ['3', 'save'])
|
||||
self.assertIn('op', queued.args[0])
|
||||
self.assertNotIn('preview', queued.args[0])
|
||||
self.assertEqual(self.prompt, original)
|
||||
|
||||
async def test_default_includes_master_and_direct_workers(self):
|
||||
result = await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
|
||||
self.assertEqual(result[2], 1) # Public worker_count still counts remote participants only.
|
||||
self.preflight.assert_awaited_once_with(
|
||||
[{"id": "worker", "name": "Worker", "type": "local", "host": None, "port": 8188}],
|
||||
require_idle=False)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
||||
|
||||
async def test_end_overrides_require_value_capability_before_master_queue(self):
|
||||
self.prompt['3']['inputs']['value_10'] = ['2', 2]
|
||||
await self.queue()
|
||||
self.assertIs(self.preflight.call_args.kwargs['require_value_results'], True)
|
||||
self.q.queue_prompt_payload.reset_mock()
|
||||
self.preflight.side_effect = RuntimeError('value_results unsupported')
|
||||
with self.assertRaisesRegex(RuntimeError, 'value_results'):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_orchestrator_setting_excludes_master_processing(self):
|
||||
self.config["settings"]["master_delegate_only"] = True
|
||||
await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
|
||||
self.assertIs(self.preflight.call_args.kwargs["require_idle"], False)
|
||||
# Upstream source stays on master even in orchestrator-only mode.
|
||||
self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
|
||||
|
||||
async def test_explicit_override_wins_over_config_in_both_directions(self):
|
||||
for configured, requested in [(True, False), (False, True)]:
|
||||
with self.subTest(configured=configured, requested=requested):
|
||||
self.config["settings"]["master_delegate_only"] = configured
|
||||
await self.queue(delegate_master=requested)
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], not requested)
|
||||
|
||||
async def test_master_only_with_explicit_empty_worker_selection(self):
|
||||
result = await self.queue(enabled_worker_ids=[])
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
||||
self.assertEqual(result[2], 0)
|
||||
self.preflight.assert_not_awaited()
|
||||
|
||||
async def test_master_only_when_no_workers_are_configured(self):
|
||||
self.config["workers"] = []
|
||||
await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
|
||||
async def test_unavailable_workers_do_not_prevent_enabled_master(self):
|
||||
self.q.select_active_workers.return_value = ([], False)
|
||||
self.q.select_active_workers.side_effect = None
|
||||
self.preflight.return_value = {}
|
||||
await self.queue()
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
|
||||
def use_real_preflight(self, queues):
|
||||
"""Exercise capability/queue parsing with read-only fake HTTP responses."""
|
||||
from test_dispatch_selection import dispatch, _FakeSession, _valid_capability
|
||||
|
||||
responses = {(worker_id, "capabilities"): (200, _valid_capability())
|
||||
for worker_id in queues}
|
||||
responses.update({(worker_id, "queue"): (200, queue)
|
||||
for worker_id, queue in queues.items()})
|
||||
session = _FakeSession(responses)
|
||||
for target, value in (
|
||||
("get_client_session", AsyncMock(return_value=session)),
|
||||
("build_worker_url", lambda worker, path: f"http://{worker['id']}{path}"),
|
||||
):
|
||||
patched = patch.object(dispatch, target, value)
|
||||
patched.start()
|
||||
self.addCleanup(patched.stop)
|
||||
self.preflight = AsyncMock(wraps=dispatch.preflight_image_list_workers)
|
||||
sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
|
||||
return responses
|
||||
|
||||
async def run_prepared_selection(self):
|
||||
"""Run real runtime selection up to the scheduler, without any item I/O."""
|
||||
rt = sys.modules[PKG + ".image_list.runtime"]
|
||||
inputs = self.prepared_inputs()
|
||||
image = object()
|
||||
context = rt.create_master_job(
|
||||
"2", "3", self.q.queue_prompt_payload.call_args.args[0],
|
||||
[{"image": image, "index": 0}],
|
||||
json.loads(inputs["image_list_enabled_worker_ids"]),
|
||||
inputs["image_list_master_callback"],
|
||||
master_participates=inputs["image_list_master_participates"],
|
||||
explicit_targets=True,
|
||||
)
|
||||
self.addCleanup(rt.remove_master_job, context)
|
||||
self.scheduler = AsyncMock(return_value=[{"image": image}])
|
||||
self.local_slot = Mock(return_value=types.SimpleNamespace(close=AsyncMock()))
|
||||
self.runtime_session = AsyncMock(return_value=object())
|
||||
config = sys.modules[PKG + ".utils.config"]
|
||||
network = sys.modules[PKG + ".utils.network"]
|
||||
with patch.object(config, "load_config", lambda: self.config), \
|
||||
patch.object(config, "get_worker_timeout_seconds", lambda: 60, create=True), \
|
||||
patch.object(config, "get_image_list_job_timeout_seconds", lambda: 0, create=True), \
|
||||
patch.object(network, "get_client_session", self.runtime_session, create=True), \
|
||||
patch.object(network, "build_worker_url", Mock(), create=True), \
|
||||
patch.object(rt, "dynamic_schedule", self.scheduler), \
|
||||
patch.dict(sys.modules, {PKG + ".image_list.local_executor": types.SimpleNamespace(
|
||||
LocalExecutionSlot=self.local_slot)}):
|
||||
await rt.run_image_list_job(context)
|
||||
self.assertIsNone(rt.get_master_job(context))
|
||||
return self.scheduler.call_args.args[1]
|
||||
|
||||
async def test_busy_worker_at_enqueue_is_available_to_scheduler_when_idle_at_start(self):
|
||||
responses = self.use_real_preflight({
|
||||
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
|
||||
})
|
||||
result = await self.queue()
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
|
||||
self.assertEqual(result[2], 1) # Candidates, not a promise of execution-time availability.
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
|
||||
# A newly enabled worker is not implicitly added to this queued selection.
|
||||
self.config["workers"].append({"id": "later", "enabled": True})
|
||||
slots = await self.run_prepared_selection()
|
||||
self.assertEqual([slot for slot in slots if isinstance(slot, str)], ["worker"])
|
||||
self.assertEqual(len(slots), 2)
|
||||
self.assertEqual(self.preflight.await_count, 2)
|
||||
self.assertEqual(self.preflight.await_args_list[0].args, self.preflight.await_args_list[1].args)
|
||||
|
||||
async def test_worker_still_busy_at_start_is_excluded_from_scheduler(self):
|
||||
self.use_real_preflight({
|
||||
"worker": {"queue_running": [], "queue_pending": [[1, "prompt-a"]]},
|
||||
})
|
||||
await self.queue()
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
|
||||
slots = await self.run_prepared_selection()
|
||||
self.assertEqual(len(slots), 1)
|
||||
self.assertNotIn("worker", slots)
|
||||
self.assertEqual(self.preflight.await_count, 2)
|
||||
self.runtime_session.assert_not_awaited()
|
||||
|
||||
async def test_explicit_empty_selection_stays_master_only_at_start(self):
|
||||
self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
|
||||
await self.queue(enabled_worker_ids=[])
|
||||
slots = await self.run_prepared_selection()
|
||||
self.assertEqual(len(slots), 1)
|
||||
self.assertNotIn("worker", slots)
|
||||
self.preflight.assert_not_awaited()
|
||||
self.runtime_session.assert_not_awaited()
|
||||
|
||||
async def test_busy_incompatible_worker_fails_before_master_queue(self):
|
||||
responses = self.use_real_preflight({
|
||||
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
|
||||
})
|
||||
capability = responses[("worker", "capabilities")][1]
|
||||
for field, invalid, message in (
|
||||
("protocol_version", 1, "incompatible"),
|
||||
("execution_targets", False, "execution_targets"),
|
||||
("value_results", False, "value_results"),
|
||||
):
|
||||
with self.subTest(field=field), patch.dict(capability, {field: invalid}):
|
||||
if field == "value_results":
|
||||
self.prompt["3"]["inputs"]["value_10"] = ["2", 2]
|
||||
with self.assertRaisesRegex(RuntimeError, message):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
async def test_orchestrator_with_all_workers_busy_queues_for_idle_start(self):
|
||||
responses = self.use_real_preflight({
|
||||
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
|
||||
})
|
||||
result = await self.queue(delegate_master=True)
|
||||
self.q.queue_prompt_payload.assert_awaited_once()
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
|
||||
self.assertEqual(result[2], 1)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
|
||||
self.assertEqual(await self.run_prepared_selection(), ["worker"])
|
||||
self.local_slot.assert_not_called()
|
||||
self.assertEqual(self.preflight.await_count, 2)
|
||||
self.assertIs(self.preflight.await_args_list[0].kwargs["require_idle"], False)
|
||||
self.assertIs(self.preflight.await_args_list[1].kwargs["require_idle"], True)
|
||||
|
||||
async def test_orchestrator_rechecks_availability_without_master_fallback(self):
|
||||
responses = self.use_real_preflight({"worker": {"queue_running": [], "queue_pending": []}})
|
||||
await self.queue(delegate_master=True)
|
||||
responses[("worker", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
|
||||
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
|
||||
await self.run_prepared_selection()
|
||||
self.assertEqual(self.preflight.await_count, 2)
|
||||
self.scheduler.assert_not_awaited()
|
||||
self.local_slot.assert_not_called()
|
||||
|
||||
async def test_orchestrator_retains_busy_candidates_when_another_worker_is_idle(self):
|
||||
self.config["workers"].append({"id": "idle", "enabled": True})
|
||||
responses = self.use_real_preflight({
|
||||
"worker": {"queue_running": [[1, "prompt-a"]], "queue_pending": []},
|
||||
"idle": {"queue_running": [], "queue_pending": []},
|
||||
})
|
||||
await self.queue(delegate_master=True)
|
||||
self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker", "idle"])
|
||||
responses[("worker", "queue")] = (200, {"queue_running": [], "queue_pending": []})
|
||||
responses[("idle", "queue")] = (200, {"queue_running": [[1, "other"]], "queue_pending": []})
|
||||
self.assertEqual(await self.run_prepared_selection(), ["worker"])
|
||||
self.local_slot.assert_not_called()
|
||||
|
||||
async def test_orchestrator_without_workers_fails_before_queue(self):
|
||||
self.config["settings"]["master_delegate_only"] = True
|
||||
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
|
||||
await self.queue(enabled_worker_ids=[])
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_orchestrator_never_inherits_legacy_offline_master_fallback(self):
|
||||
self.config["settings"]["master_delegate_only"] = True
|
||||
self.q.select_active_workers.side_effect = None
|
||||
self.q.select_active_workers.return_value = ([], False)
|
||||
self.preflight.return_value = {}
|
||||
with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_serverless_only_uses_master_without_provider_dispatch(self):
|
||||
self.config["workers"] = [{"id": "pool", "type": "serverless", "enabled": True}]
|
||||
result = await self.queue()
|
||||
self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
|
||||
self.assertEqual(result[2], 0)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
|
||||
self.preflight.side_effect = RuntimeError("upgrade worker")
|
||||
with self.assertRaisesRegex(RuntimeError, "upgrade worker"):
|
||||
await self.queue()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
@@ -0,0 +1,527 @@
|
||||
"""Stable content signatures with fresh, expiring delivery capabilities."""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from image_list import runtime as rt
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def item():
|
||||
return rt.normalize_image_list_inputs(
|
||||
torch.arange(18, dtype=torch.float64).reshape(1, 2, 3, 3),
|
||||
torch.zeros(1, 2, 3, dtype=torch.bfloat16),
|
||||
{'value_1': 'original', 'value_2': True, 'value_3': 7, 'value_4': 0.25,
|
||||
'value_10': torch.ones(2, 3, dtype=torch.float16)})[0]
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
'source': {'class_type': 'Source', 'inputs': {}},
|
||||
'start': {'class_type': rt.START_CLASS, 'inputs': {
|
||||
'images': ['source', 0], 'image_list_execution_id': 'transient',
|
||||
'image_list_workflow_meta': {'job': 'metadata'},
|
||||
'image_list_enabled_worker_ids': '["worker"]',
|
||||
'image_list_master_callback': 'http://old-master',
|
||||
'values.value_1': 'inherited', 'worker_context_id': 'old-context'}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
|
||||
'end': {'class_type': rt.END_CLASS, 'inputs': {
|
||||
'images': ['op', 0], 'context': ['start', 3], 'value_3': ['start', 2]}},
|
||||
}
|
||||
|
||||
|
||||
def worker_prompt(token, context, data, *, local=True):
|
||||
prompt = graph()
|
||||
result = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
|
||||
token, context, data['index'], 'http://master',
|
||||
item_identity=rt.item_content_identity(data))
|
||||
if local:
|
||||
result['start']['inputs']['image_list_local_execution'] = True
|
||||
result['end']['inputs']['image_list_local_execution'] = True
|
||||
return result
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_state():
|
||||
yield
|
||||
rt._WORKER_ITEMS.clear()
|
||||
rt._MASTER_JOBS.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def boundaries(monkeypatch):
|
||||
"""Load only the boundary module; restore every replaced module on exit."""
|
||||
name = '_native_cache_identity_tests'
|
||||
for suffix in ('', '.nodes', '.utils', '.image_list', '.api'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, name + '.image_list.runtime', rt)
|
||||
io = types.SimpleNamespace(ComfyNode=type('ComfyNode', (), {}),
|
||||
NodeOutput=lambda *args: args)
|
||||
for cls in ('Image', 'Mask', 'String', 'Int', 'Float', 'Boolean'):
|
||||
setattr(io, cls, object())
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api', types.ModuleType('comfy_api'))
|
||||
monkeypatch.setitem(sys.modules, 'comfy_api.v0_0_2', types.SimpleNamespace(io=io))
|
||||
monkeypatch.setitem(sys.modules, name + '.utils.async_helpers',
|
||||
types.SimpleNamespace(run_async_in_server_loop=lambda coro, **kw: asyncio.run(coro)))
|
||||
spec = importlib.util.spec_from_file_location(name + '.nodes.image_list', ROOT / 'nodes/image_list.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, module)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def start(boundaries, prompt):
|
||||
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
|
||||
return boundaries.DistributedImageListStart.execute(**prompt['start']['inputs'])
|
||||
|
||||
|
||||
def end(boundaries, prompt, cached):
|
||||
boundaries.DistributedImageListEnd.hidden = types.SimpleNamespace(prompt=prompt, unique_id='end')
|
||||
inputs = {k: v for k, v in prompt['end']['inputs'].items()
|
||||
if k not in ('images', 'context', 'value_3')}
|
||||
return boundaries.DistributedImageListEnd.execute(
|
||||
images=cached[0], context=cached[3], value_3=cached[2], **inputs)
|
||||
|
||||
|
||||
def stage(data):
|
||||
token = rt.create_worker_task_token()
|
||||
rt.store_worker_item(token, rt.encode_item(data))
|
||||
return token
|
||||
|
||||
|
||||
def master(data, token):
|
||||
context = rt.create_master_job('start', 'end', graph(), [data], ['worker'], '')
|
||||
rt.register_assignment(context, data['index'], 'worker', token)
|
||||
return context
|
||||
|
||||
|
||||
def test_identity_is_exact_typed_order_independent_and_not_transport_bounded(monkeypatch):
|
||||
original = item()
|
||||
identity = rt.item_content_identity(original)
|
||||
assert rt.is_valid_item_identity(identity)
|
||||
reordered = copy.deepcopy(original)
|
||||
reordered['values'] = dict(reversed(list(reordered['values'].items())))
|
||||
assert rt.item_content_identity(reordered) == identity
|
||||
strided = copy.deepcopy(original)
|
||||
strided['image'] = original['image'].transpose(1, 2).contiguous().transpose(1, 2)
|
||||
assert not strided['image'].is_contiguous()
|
||||
assert rt.item_content_identity(strided) == identity
|
||||
monkeypatch.setattr(rt, 'MAX_PAYLOAD_BYTES', 1)
|
||||
with pytest.raises(ValueError, match='transport limit'):
|
||||
rt.encode_item(original)
|
||||
assert rt.item_content_identity(original) == identity
|
||||
scalar_identities = set()
|
||||
for value in (True, 1, 1.0, '1', 0.0, -0.0):
|
||||
changed = copy.deepcopy(original)
|
||||
changed['values']['value_2'] = value
|
||||
scalar_identities.add(rt.item_content_identity(changed))
|
||||
assert len(scalar_identities) == 6
|
||||
|
||||
|
||||
@pytest.mark.parametrize('change', ['image', 'mask', 'value', 'index', 'dtype', 'shape', 'slot', 'missing-mask'])
|
||||
def test_every_item_component_invalidates_identity(change):
|
||||
original, changed = item(), item()
|
||||
if change in ('image', 'mask'):
|
||||
changed[change].reshape(-1)[0] = 1
|
||||
elif change == 'value':
|
||||
changed['values']['value_10'][0, 0] = 2
|
||||
elif change == 'index':
|
||||
changed['index'] = 1
|
||||
elif change == 'dtype':
|
||||
changed['image'] = changed['image'].float()
|
||||
elif change == 'shape':
|
||||
changed['values']['value_10'] = changed['values']['value_10'].reshape(3, 2)
|
||||
elif change == 'slot':
|
||||
changed['values']['value_9'] = changed['values'].pop('value_10')
|
||||
else:
|
||||
changed['mask'] = None
|
||||
assert rt.item_content_identity(original) != rt.item_content_identity(changed)
|
||||
|
||||
|
||||
def test_float64_low_bits_and_bfloat16_bytes_are_not_quantized():
|
||||
data, changed = item(), item()
|
||||
changed['image'].reshape(-1)[0] = torch.nextafter(torch.tensor(0., dtype=torch.float64),
|
||||
torch.tensor(1., dtype=torch.float64))
|
||||
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
|
||||
changed = item()
|
||||
changed['mask'].view(torch.int16).reshape(-1)[0] = 1
|
||||
assert rt.item_content_identity(data) != rt.item_content_identity(changed)
|
||||
|
||||
|
||||
def test_stable_start_replaced_wholesale_and_legacy_graph_preserved():
|
||||
prompt = graph()
|
||||
frozen = copy.deepcopy(prompt)
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
legacy = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a')
|
||||
assert legacy['start']['inputs']['worker_item_token'] == 'a'
|
||||
assert legacy['start']['inputs']['image_list_execution_id'] == 'transient'
|
||||
identity = rt.item_content_identity(item())
|
||||
a = rt.build_worker_prompt(prompt, region, 'a', 'ctx-a', 0, 'http://a', item_identity=identity)
|
||||
b = rt.build_worker_prompt(prompt, region, 'b', 'ctx-b', 0, 'http://b', item_identity=identity)
|
||||
assert a['start']['inputs'] == b['start']['inputs'] == {'worker_item_identity': identity}
|
||||
assert a['end']['inputs']['worker_context_id'] == 'ctx-a'
|
||||
assert b['end']['inputs']['worker_result_token'] == 'b'
|
||||
assert b['end']['inputs']['worker_item_identity'] == identity
|
||||
assert b['end']['inputs']['value_3'] == ['start', 2]
|
||||
assert prompt == frozen
|
||||
|
||||
|
||||
@pytest.mark.parametrize('bad', ['', 'sha256:' + 'z' * 64, 'a' * 64, False, ['bad']])
|
||||
def test_malformed_identity_rejected(bad):
|
||||
assert not rt.is_valid_item_identity(bad)
|
||||
prompt = graph()
|
||||
with pytest.raises((ValueError, TypeError), match='identity'):
|
||||
rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0],
|
||||
'token', 'context', 0, '', item_identity=bad)
|
||||
|
||||
|
||||
def test_cache_miss_then_hit_delivers_current_context_token_and_values(boundaries):
|
||||
data = item()
|
||||
a = stage(data)
|
||||
context_a = master(data, a)
|
||||
prompt_a = worker_prompt(a, context_a, data)
|
||||
cached = start(boundaries, prompt_a)
|
||||
assert cached[3] == [rt.item_content_identity(data)]
|
||||
assert 'image' not in rt._WORKER_ITEMS[a][1] # only a small, expiring receipt remains
|
||||
end(boundaries, prompt_a, cached)
|
||||
assert a not in rt._WORKER_ITEMS
|
||||
assert rt.get_master_job(context_a).results[0]['values'] == {'value_3': 0}
|
||||
b = stage(data)
|
||||
context_b = master(data, b)
|
||||
prompt_b = worker_prompt(b, context_b, data)
|
||||
assert prompt_a['start'] == prompt_b['start']
|
||||
# A native hit does NOT execute Start; End must consume this current upload.
|
||||
end(boundaries, prompt_b, cached)
|
||||
assert b not in rt._WORKER_ITEMS
|
||||
state = rt.get_master_job(context_b)
|
||||
assert state.expected_tokens[0] == b
|
||||
assert torch.equal(state.results[0]['image'], data['image'])
|
||||
assert rt._collect_output_lists([state.results[0]], state=state, include_values=True)[3][0]['value_3'] == 0
|
||||
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
|
||||
end(boundaries, prompt_b, cached)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('miss', [False, True])
|
||||
def test_expired_item_cannot_deliver_even_with_cached_start(boundaries, miss):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
context = master(data, token)
|
||||
prompt = worker_prompt(token, context, data)
|
||||
cached = start(boundaries, prompt) if miss else ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
|
||||
_, staged = rt._WORKER_ITEMS[token]
|
||||
rt._WORKER_ITEMS[token] = (0, staged)
|
||||
with pytest.raises(RuntimeError, match='missing|expired|consumed'):
|
||||
end(boundaries, prompt, cached)
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
|
||||
|
||||
def test_identity_mismatch_on_start_and_on_cache_hit_fails_closed(boundaries):
|
||||
data, different = item(), item()
|
||||
different['image'][0, 0, 0, 0] = 3
|
||||
for hit in (False, True):
|
||||
token = stage(different)
|
||||
context = master(different, token)
|
||||
prompt = worker_prompt(token, context, data)
|
||||
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
|
||||
with pytest.raises((ValueError, RuntimeError), match='identity'):
|
||||
end(boundaries, prompt, cached) if hit else start(boundaries, prompt)
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
rt.delete_worker_item(token)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('change', ['ambiguous', 'wrong-link', 'wrong-start-id', 'end-identity', 'end-index'])
|
||||
def test_start_requires_exact_paired_end_binding(boundaries, change):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
prompt = worker_prompt(token, 'context', data)
|
||||
if change == 'ambiguous':
|
||||
prompt['second-end'] = copy.deepcopy(prompt['end'])
|
||||
elif change == 'wrong-link':
|
||||
prompt['end']['inputs']['context'] = ['start', 2]
|
||||
elif change == 'wrong-start-id':
|
||||
prompt['renamed'] = prompt.pop('start')
|
||||
elif change == 'end-identity':
|
||||
prompt['end']['inputs']['worker_item_identity'] = 'sha256:' + '0' * 64
|
||||
else:
|
||||
prompt['end']['inputs']['worker_item_index'] = 1
|
||||
boundaries.DistributedImageListStart.hidden = types.SimpleNamespace(prompt=prompt, unique_id='start')
|
||||
with pytest.raises((RuntimeError, ValueError), match='paired|identity|index'):
|
||||
boundaries.DistributedImageListStart.execute(worker_item_identity=rt.item_content_identity(data),
|
||||
image_list_local_execution=True)
|
||||
|
||||
|
||||
def test_end_checks_stable_context_before_delivering(boundaries):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
context = master(data, token)
|
||||
prompt = worker_prompt(token, context, data)
|
||||
cached = list(start(boundaries, prompt))
|
||||
cached[3] = ['sha256:' + '0' * 64]
|
||||
with pytest.raises((RuntimeError, ValueError), match='identity|context'):
|
||||
end(boundaries, prompt, cached)
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
|
||||
|
||||
def test_legacy_worker_path_still_consumes_and_delivers_job_context(boundaries):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
context = master(data, token)
|
||||
prompt = graph()
|
||||
prompt = rt.build_worker_prompt(prompt, rt.validate_image_list_regions(prompt)[0], token, context, 0, '')
|
||||
prompt['start']['inputs']['image_list_local_execution'] = True
|
||||
prompt['end']['inputs']['image_list_local_execution'] = True
|
||||
cached = start(boundaries, prompt)
|
||||
assert cached[3] == [context]
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
end(boundaries, prompt, cached)
|
||||
assert rt.get_master_job(context).received_indexes == {0}
|
||||
|
||||
|
||||
def test_remote_cache_hit_submits_fresh_callback_metadata(boundaries, monkeypatch):
|
||||
data = item()
|
||||
token = stage(data)
|
||||
prompt = worker_prompt(token, 'fresh-context', data, local=False)
|
||||
submit = AsyncMock()
|
||||
monkeypatch.setattr(boundaries, 'submit_worker_result', submit)
|
||||
cached = ([data['image']], [data['mask']], [0], [rt.item_content_identity(data)])
|
||||
end(boundaries, prompt, cached)
|
||||
assert submit.await_args.args[:4] == ('http://master', token, 'fresh-context', 0)
|
||||
assert submit.await_args.args[4]['values']['value_3'] == {'type': 'INT', 'value': 0}
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
|
||||
|
||||
def test_item_route_returns_server_computed_identity(boundaries, monkeypatch):
|
||||
from aiohttp import web
|
||||
name = '_native_cache_identity_tests'
|
||||
server = types.SimpleNamespace(PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=web.RouteTableDef())))
|
||||
monkeypatch.setitem(sys.modules, 'server', server)
|
||||
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', ROOT / 'api/image_list_routes.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, module)
|
||||
spec.loader.exec_module(module)
|
||||
data = item()
|
||||
token = rt.create_worker_task_token()
|
||||
monkeypatch.setattr(module, '_json_body', AsyncMock(return_value={
|
||||
'token': token, 'context_id': 'context', 'item': rt.encode_item(data),
|
||||
'item_identity': 'sha256:' + '0' * 64}))
|
||||
response = asyncio.run(module.image_list_item_endpoint(object()))
|
||||
assert response.status == 200
|
||||
assert json.loads(response.body)['item_identity'] == rt.item_content_identity(data)
|
||||
assert len(rt._WORKER_ITEMS[token]) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize('responses', [
|
||||
('valid', 'valid'), ('valid', 'absent'), ('invalid', 'valid'),
|
||||
('bad-json', 'valid'), ('non-object', 'valid'),
|
||||
])
|
||||
def test_runtime_opts_in_per_upload_and_can_downgrade(monkeypatch, responses):
|
||||
name = '_native_cache_identity_runtime'
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
|
||||
runtime = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, runtime)
|
||||
spec.loader.exec_module(runtime)
|
||||
workers = [{'id': 'worker', 'type': 'local'}]
|
||||
dispatched, uploads = [], []
|
||||
preflight = AsyncMock(return_value={'worker': True})
|
||||
sync_media = AsyncMock()
|
||||
|
||||
class Response:
|
||||
status = 200
|
||||
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
pass
|
||||
|
||||
async def json(self):
|
||||
if self.data == 'bad-json':
|
||||
raise ValueError('old response is not JSON')
|
||||
return self.data
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
class Session:
|
||||
def post(self, url, *, data, **kw):
|
||||
assert url.endswith('/distributed/image-list/item')
|
||||
body = json.loads(data)
|
||||
identity = runtime.store_worker_item(body['token'], body['item'])
|
||||
mode = responses[len(uploads)]
|
||||
uploads.append(body)
|
||||
response = {'valid': {'item_identity': identity}, 'absent': {'status': 'stored'},
|
||||
'invalid': {'item_identity': 'bad'}, 'bad-json': 'bad-json', 'non-object': []}[mode]
|
||||
return Response(response)
|
||||
|
||||
def get(self, url, **kw):
|
||||
pid = url.rsplit('/', 1)[-1]
|
||||
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
|
||||
|
||||
def delete(self, url, **kw):
|
||||
runtime.delete_worker_item(url.rsplit('/', 1)[-1])
|
||||
return Response({})
|
||||
|
||||
async def dispatch(worker, prompt, meta, **kwargs):
|
||||
dispatched.append(copy.deepcopy(prompt))
|
||||
delivery = prompt['end']['inputs']
|
||||
token, index = delivery['worker_result_token'], delivery['worker_item_index']
|
||||
identity = prompt['start']['inputs'].get('worker_item_identity')
|
||||
options = {'item_identity': identity, 'item_index': index} if identity else {}
|
||||
data = runtime.consume_worker_item(token, **options)
|
||||
if identity:
|
||||
runtime.finish_worker_item(token, identity, index)
|
||||
runtime.set_master_result(context, index, token, {
|
||||
'image': data['image'], 'text': str(index), 'values': {'value_3': index}})
|
||||
assert kwargs['partial_execution_targets'] == ['end']
|
||||
return f'worker-{index}'
|
||||
|
||||
dependencies = {
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=preflight, dispatch_worker_prompt=dispatch),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda *args: workers),
|
||||
'.api.orchestration.media_sync': types.SimpleNamespace(sync_worker_media=sync_media),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
|
||||
get_worker_timeout_seconds=lambda: 60),
|
||||
'.utils.network': types.SimpleNamespace(build_master_callback_url=lambda *a, **kw: 'http://master',
|
||||
build_worker_url=lambda worker, path: 'http://worker' + path,
|
||||
get_client_session=AsyncMock(return_value=Session())),
|
||||
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
|
||||
monkeypatch.setitem(sys.modules, 'comfy', comfy)
|
||||
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
|
||||
items = [item(), item()]
|
||||
items[1]['index'] = 1
|
||||
context = runtime.create_master_job('start', 'end', graph(), items, ['worker'], '')
|
||||
state = runtime.get_master_job(context)
|
||||
try:
|
||||
result = asyncio.run(runtime._run_image_list_job(state, include_values=True))
|
||||
assert result[2] == ['0', '1']
|
||||
assert [row['value_3'] for row in result[3]] == [0, 1]
|
||||
assert len(uploads) == len(dispatched) == 2
|
||||
assert sync_media.await_count == 1
|
||||
for index, prompt in enumerate(dispatched):
|
||||
start_inputs = prompt['start']['inputs']
|
||||
if responses[index] == 'valid':
|
||||
assert start_inputs == {'worker_item_identity': runtime.item_content_identity(items[index])}
|
||||
assert prompt['end']['inputs']['worker_context_id'] == context
|
||||
else:
|
||||
assert 'worker_item_identity' not in start_inputs
|
||||
assert start_inputs['worker_item_token'] == uploads[index]['token']
|
||||
assert start_inputs['worker_context_id'] == context
|
||||
assert 'worker_item_identity' not in prompt['end']['inputs']
|
||||
assert prompt['end']['inputs']['worker_result_token'] == uploads[index]['token']
|
||||
assert not runtime._WORKER_ITEMS
|
||||
finally:
|
||||
runtime.remove_master_job(context)
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
|
||||
def test_runtime_local_path_hashes_without_transport(monkeypatch):
|
||||
name = '_native_cache_identity_local_runtime'
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
|
||||
runtime = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, runtime)
|
||||
spec.loader.exec_module(runtime)
|
||||
dispatched = []
|
||||
|
||||
class Slot:
|
||||
def __init__(self, **kw):
|
||||
pass
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
dispatched.append(prompt)
|
||||
delivery = prompt[end_id]['inputs']
|
||||
identity = prompt['start']['inputs']['worker_item_identity']
|
||||
data = runtime.consume_worker_item(delivery['worker_result_token'], item_identity=identity,
|
||||
item_index=delivery['worker_item_index'])
|
||||
runtime.finish_worker_item(delivery['worker_result_token'], identity, data['index'])
|
||||
runtime.set_master_result(delivery['worker_context_id'], data['index'], delivery['worker_result_token'],
|
||||
{'image': data['image'], 'values': {'value_3': data['index']}})
|
||||
future = asyncio.get_running_loop().create_future()
|
||||
future.set_result(None)
|
||||
return future
|
||||
|
||||
def request_stop(self):
|
||||
pass
|
||||
|
||||
async def close(self):
|
||||
pass
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('local path attempted HTTP/transport')
|
||||
|
||||
dependencies = {
|
||||
'.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=Slot),
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=forbidden, dispatch_worker_prompt=forbidden),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda *args: []),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: {'settings': {}},
|
||||
get_worker_timeout_seconds=lambda: 60),
|
||||
'.utils.network': types.SimpleNamespace(build_master_callback_url=forbidden,
|
||||
build_worker_url=forbidden, get_client_session=forbidden),
|
||||
'.utils.logging': types.SimpleNamespace(log=lambda *args: None),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
monkeypatch.setitem(sys.modules, name + suffix, module)
|
||||
monkeypatch.setitem(sys.modules, 'aiohttp', types.SimpleNamespace(
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError))
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: False)
|
||||
monkeypatch.setitem(sys.modules, 'comfy', comfy)
|
||||
monkeypatch.setitem(sys.modules, 'comfy.model_management', comfy.model_management)
|
||||
monkeypatch.setattr(runtime, 'MAX_PAYLOAD_BYTES', 1)
|
||||
monkeypatch.setattr(runtime, 'encode_item', forbidden)
|
||||
data = item()
|
||||
context = runtime.create_master_job('start', 'end', graph(), [data], [], '', master_participates=True)
|
||||
try:
|
||||
result = asyncio.run(runtime._run_image_list_job(runtime.get_master_job(context), include_values=True))
|
||||
assert dispatched[0]['start']['inputs'] == {'worker_item_identity': runtime.item_content_identity(data),
|
||||
'image_list_local_execution': True}
|
||||
assert result[3][0]['value_3'] == 0
|
||||
assert not runtime._WORKER_ITEMS
|
||||
finally:
|
||||
runtime.remove_master_job(context)
|
||||
runtime._WORKER_ITEMS.clear()
|
||||
|
||||
|
||||
def test_real_native_signatures_and_cached_delivery():
|
||||
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
|
||||
if not (core / 'comfy_execution/caching.py').exists():
|
||||
pytest.skip('real ComfyUI source not available beside this interpreter')
|
||||
result = subprocess.run([sys.executable, str(ROOT / 'tests/helpers/image_list_native_cache_identity.py'), str(core)],
|
||||
text=True, capture_output=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'NATIVE_CACHE_IDENTITY_OK' in result.stdout
|
||||
@@ -0,0 +1,358 @@
|
||||
"""Exercise production per-item orchestration through controlled provider seams."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from test_image_list import graph_one_region
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
class Response:
|
||||
def __init__(self, data=None, status=200):
|
||||
self.data, self.status = data, status
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
async def json(self):
|
||||
return self.data
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status >= 400:
|
||||
raise RuntimeError(f'HTTP {self.status}')
|
||||
|
||||
|
||||
class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
name = '_image_list_boundaries'
|
||||
self.modules = {}
|
||||
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
|
||||
module = types.ModuleType(name + suffix)
|
||||
module.__path__ = []
|
||||
self.modules[name + suffix] = module
|
||||
self.patch_modules = patch.dict(sys.modules, self.modules)
|
||||
self.patch_modules.start()
|
||||
self.addCleanup(self.patch_modules.stop)
|
||||
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
|
||||
self.rt = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = self.rt
|
||||
spec.loader.exec_module(self.rt)
|
||||
self.worker = {'id': 'worker', 'host': 'worker', 'type': 'remote', 'port': 8188}
|
||||
self.config = {'workers': [self.worker], 'settings': {}}
|
||||
self.prompts, self.uploads, self.deletions, self.cancel_requests = [], [], [], []
|
||||
self.reject = False
|
||||
self.interrupted = False
|
||||
self.completed = True
|
||||
self.queue = {'queue_running': [], 'queue_pending': []}
|
||||
self.callback = Mock(side_effect=lambda worker, **_kwargs:
|
||||
'http://127.0.0.1:8188' if worker['type'] == 'local' else 'http://public-master')
|
||||
self.sync = AsyncMock(side_effect=self.sync_media)
|
||||
self.dispatch = AsyncMock(side_effect=self.dispatch_prompt)
|
||||
self.logs = []
|
||||
dependencies = {
|
||||
'.api.orchestration.dispatch': types.SimpleNamespace(
|
||||
preflight_image_list_workers=AsyncMock(return_value={'worker': True}),
|
||||
dispatch_worker_prompt=self.dispatch),
|
||||
'.api.orchestration.prompt_transform': types.SimpleNamespace(
|
||||
_resolve_direct_workers_for_image_list=lambda config, ids: config['workers']),
|
||||
'.api.orchestration.media_sync': types.SimpleNamespace(
|
||||
sync_worker_media=self.sync, fetch_worker_path_separator=AsyncMock(return_value='/'),
|
||||
convert_paths_for_platform=lambda prompt, separator: prompt),
|
||||
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
|
||||
get_worker_timeout_seconds=lambda: 60,
|
||||
get_image_list_job_timeout_seconds=lambda config=None: 0),
|
||||
'.utils.network': types.SimpleNamespace(build_master_url=lambda **kwargs: 'http://public-master',
|
||||
build_master_callback_url=self.callback,
|
||||
build_worker_url=lambda worker, path='': f"http://{worker['id']}{path}",
|
||||
get_client_session=AsyncMock(return_value=self)),
|
||||
'.utils.logging': types.SimpleNamespace(log=self.logs.append),
|
||||
}
|
||||
for suffix, module in dependencies.items():
|
||||
sys.modules[name + suffix] = module
|
||||
self.config_module = dependencies['.utils.config']
|
||||
# Other legacy test modules replace aiohttp during collection. This
|
||||
# controlled provider fixture needs only its explicit timeout/error seam.
|
||||
sys.modules['aiohttp'] = types.SimpleNamespace(
|
||||
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError, ClientConnectionError=ConnectionError)
|
||||
comfy = types.ModuleType('comfy')
|
||||
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
|
||||
sys.modules['comfy'] = comfy
|
||||
sys.modules['comfy.model_management'] = comfy.model_management
|
||||
|
||||
async def sync_media(self, worker, prompt):
|
||||
prompt['7']['inputs']['image'] = 'shared/ref.png'
|
||||
|
||||
async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
|
||||
self.assertEqual(kwargs.get('partial_execution_targets'), ['5'])
|
||||
if self.reject:
|
||||
raise RuntimeError('prompt rejected')
|
||||
self.prompts.append((json.loads(json.dumps(prompt)), meta))
|
||||
inputs = prompt['5']['inputs']
|
||||
token = inputs['worker_result_token']
|
||||
item = self.rt.consume_worker_item(token)
|
||||
self.rt.set_master_result(self.context, inputs['worker_item_index'], token, {'image': item['image']})
|
||||
return f'prompt-{len(self.prompts)}'
|
||||
|
||||
def post(self, url, *, json=None, data=None, **kwargs):
|
||||
if data is not None:
|
||||
json = __import__('json').loads(data)
|
||||
if url.endswith('/distributed/image-list/item'):
|
||||
self.uploads.append(json['token'])
|
||||
self.rt.store_worker_item(json['token'], json['item'])
|
||||
else:
|
||||
self.cancel_requests.append((url, json))
|
||||
return Response({})
|
||||
|
||||
def delete(self, url, **kwargs):
|
||||
token = url.rsplit('/', 1)[-1]
|
||||
self.deletions.append(token)
|
||||
self.rt.delete_worker_item(token)
|
||||
return Response({})
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
if url.endswith('/queue'):
|
||||
return Response(self.queue)
|
||||
pid = url.rsplit('/', 1)[-1]
|
||||
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}}
|
||||
if self.completed else {})
|
||||
|
||||
def create(self, count=2):
|
||||
prompt = graph_one_region()
|
||||
prompt['7'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared/ref.png [output]'}}
|
||||
prompt['4']['inputs']['reference'] = ['7', 0]
|
||||
self.metadata = {'nodes': [{'id': 42}], 'extra': {'label': 'generic fixture'}}
|
||||
self.context = self.rt.create_master_job('3', '5', prompt,
|
||||
[{'image': torch.zeros((1, 2, 2, 3)), 'index': i} for i in range(count)],
|
||||
['worker'], 'http://public-master', self.metadata)
|
||||
return self.context
|
||||
|
||||
async def test_runtime_preflight_rejects_old_target_capability_before_upload(self):
|
||||
from test_dispatch_selection import _load_dispatch_module, _FakeSession, _valid_capability
|
||||
real_dispatch = _load_dispatch_module()
|
||||
capability = _valid_capability()
|
||||
capability.pop('execution_targets')
|
||||
session = _FakeSession({('worker', 'capabilities'): (200, capability)})
|
||||
module = sys.modules['_image_list_boundaries.api.orchestration.dispatch']
|
||||
module.preflight_image_list_workers = real_dispatch.preflight_image_list_workers
|
||||
with patch.object(real_dispatch, 'get_client_session', AsyncMock(return_value=session)), patch.object(
|
||||
real_dispatch, 'build_worker_url', side_effect=lambda worker, path: f"http://{worker['id']}{path}"
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, 'execution_targets.*upgrade'):
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertFalse(self.uploads)
|
||||
self.dispatch.assert_not_awaited()
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
|
||||
async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
|
||||
context = self.create()
|
||||
await self.rt.run_image_list_job(context)
|
||||
self.assertEqual(self.sync.await_count, 1)
|
||||
self.assertEqual([p['7']['inputs']['image'] for p, _ in self.prompts],
|
||||
['shared/ref.png', 'shared/ref.png'])
|
||||
tokens = [p['5']['inputs']['worker_result_token'] for p, _ in self.prompts]
|
||||
self.assertEqual(len(set(tokens)), 2)
|
||||
self.assertEqual([p['5']['inputs']['worker_item_index'] for p, _ in self.prompts], [0, 1])
|
||||
self.assertEqual([meta for _, meta in self.prompts], [self.metadata, self.metadata])
|
||||
|
||||
async def test_annotated_media_roots_survive_windows_prep_and_repeated_items(self):
|
||||
import tempfile
|
||||
|
||||
context = self.create(2)
|
||||
state = self.rt.get_master_job(context)
|
||||
state.original_prompt['8'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared\\ref.png [temp]'}}
|
||||
state.original_prompt['4']['inputs']['second_reference'] = ['8', 0]
|
||||
name = '_image_list_boundaries.api.orchestration.media_sync'
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT / 'api/orchestration/media_sync.py')
|
||||
assert spec is not None and spec.loader is not None
|
||||
media = importlib.util.module_from_spec(spec)
|
||||
sys.modules['_image_list_boundaries.utils.logging'].debug_log = self.logs.append
|
||||
sys.modules['_image_list_boundaries.utils.trace_logger'] = types.SimpleNamespace(
|
||||
trace_debug=lambda *args: None, trace_info=lambda *args: None)
|
||||
spec.loader.exec_module(media)
|
||||
media.fetch_worker_path_separator = AsyncMock(return_value='\\')
|
||||
converted = Mock(wraps=media.convert_paths_for_platform)
|
||||
media.convert_paths_for_platform = converted
|
||||
uploads = []
|
||||
|
||||
async def upload(worker, filename, content, digest, mime, **kwargs):
|
||||
uploads.append((filename, content))
|
||||
return True, f'worker-input/{len(uploads)}.png'
|
||||
|
||||
media._upload_media_to_worker = AsyncMock(side_effect=upload)
|
||||
media.sync_worker_media = AsyncMock(wraps=media.sync_worker_media)
|
||||
with tempfile.TemporaryDirectory(prefix='image-list-media-') as directory:
|
||||
root = Path(directory)
|
||||
for kind in ('input', 'output', 'temp'):
|
||||
path = root / kind / 'shared/ref.png'
|
||||
path.parent.mkdir(parents=True)
|
||||
path.write_bytes(kind.encode())
|
||||
|
||||
def resolve_annotated(filename):
|
||||
for kind in ('output', 'temp', 'input'):
|
||||
suffix = f' [{kind}]'
|
||||
if filename.endswith(suffix):
|
||||
return str(root / kind / filename[:-len(suffix)])
|
||||
return str(root / 'input' / filename)
|
||||
|
||||
with patch.dict(sys.modules, {name: media, 'folder_paths': types.SimpleNamespace(
|
||||
get_annotated_filepath=resolve_annotated)}):
|
||||
await self.rt.run_image_list_job(context)
|
||||
self.assertEqual(uploads, [('shared/ref.png [output]', b'output'),
|
||||
('shared/ref.png [temp]', b'temp')])
|
||||
media.sync_worker_media.assert_awaited_once()
|
||||
media.fetch_worker_path_separator.assert_awaited_once()
|
||||
converted.assert_called_once()
|
||||
self.assertEqual(converted.call_args.args[1], '\\')
|
||||
for prompt, _ in self.prompts:
|
||||
self.assertEqual(prompt['7']['inputs']['image'], 'worker-input/1.png')
|
||||
self.assertEqual(prompt['8']['inputs']['image'], 'worker-input/2.png')
|
||||
self.assertEqual(state.original_prompt['7']['inputs']['image'], 'shared/ref.png [output]')
|
||||
|
||||
async def test_local_worker_callback_does_not_use_public_master_address(self):
|
||||
self.worker['type'] = 'local'
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://127.0.0.1:8188')
|
||||
self.callback.assert_called()
|
||||
|
||||
async def test_remote_worker_uses_topology_callback_helper(self):
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://public-master')
|
||||
self.callback.assert_called()
|
||||
|
||||
def test_metadata_is_forwarded_into_master_start_prompt(self):
|
||||
from test_prompt_transform import pt
|
||||
prompt = graph_one_region()
|
||||
metadata = {'nodes': [{'id': 1}]}
|
||||
index = pt.PromptIndex(prompt)
|
||||
transformed = pt.apply_participant_overrides(
|
||||
prompt, 'master', ['worker'], {'3': 'run'}, 'http://public-master', False,
|
||||
index, workflow_meta=metadata)
|
||||
self.assertEqual(transformed['3']['inputs']['image_list_workflow_meta'], metadata)
|
||||
|
||||
async def test_aggregate_input_limit_fails_before_upload(self):
|
||||
self.rt.MAX_REQUEST_BYTES = 64
|
||||
with self.assertRaisesRegex(ValueError, 'complete.*request'):
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertFalse(self.uploads)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
|
||||
def test_tensor_and_combined_request_limits_preserve_dtype(self):
|
||||
item = {'image': torch.zeros((1, 4, 4, 3)), 'index': 0}
|
||||
encoded = self.rt.encode_item(item)
|
||||
payload = {'context_id': 'example', 'token': 't' * 48, 'item': encoded}
|
||||
single = self.rt.encode_request_payload(payload)
|
||||
self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32')
|
||||
self.rt.MAX_REQUEST_BYTES = len(single)
|
||||
self.assertEqual(self.rt.encode_request_payload(payload), single)
|
||||
item['values'] = {'value_1': item['image']}
|
||||
item['mask'] = torch.zeros((1, 4, 4))
|
||||
payload['item'] = self.rt.encode_item(item)
|
||||
with self.assertRaisesRegex(ValueError, 'complete.*request'):
|
||||
self.rt.encode_request_payload(payload)
|
||||
# Size is checked before allocating/copying a 4K tensor onto the CPU.
|
||||
with self.assertRaisesRegex(ValueError, 'transport limit'):
|
||||
self.rt.encode_value(torch.empty((1, 2160, 3840, 3), device='meta'), 'IMAGE')
|
||||
|
||||
async def test_rejected_submission_reclaims_uploaded_token_without_prompt_id(self):
|
||||
self.reject = True
|
||||
with self.assertRaisesRegex(RuntimeError, 'prompt rejected'):
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertEqual(self.deletions, self.uploads)
|
||||
self.assertFalse(self.rt._WORKER_ITEMS)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
self.assertFalse(self.cancel_requests)
|
||||
|
||||
async def test_user_interrupt_during_pending_submission_is_bounded(self):
|
||||
entered = asyncio.Event()
|
||||
async def never_ack(*args, **kwargs):
|
||||
entered.set()
|
||||
await asyncio.Event().wait()
|
||||
self.dispatch.side_effect = never_ack
|
||||
self.rt.SUBMISSION_CANCEL_GRACE_SECONDS = 0.02
|
||||
task = asyncio.create_task(self.rt.run_image_list_job(self.create(2)))
|
||||
await asyncio.wait_for(entered.wait(), 1)
|
||||
self.interrupted = True
|
||||
with self.assertRaises(self.rt.ImageListCancelledError):
|
||||
await asyncio.wait_for(task, 0.8)
|
||||
self.assertEqual(len(self.uploads), 1)
|
||||
self.assertEqual(self.deletions, self.uploads)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
self.assertTrue(any('ambiguous' in line for line in self.logs), self.logs)
|
||||
|
||||
async def test_owned_running_prompt_survives_cold_load_and_old_hour_cap(self):
|
||||
self.completed = False
|
||||
context = self.create(2)
|
||||
clock = [0.0]
|
||||
calls = [0]
|
||||
original_get = self.get
|
||||
def get(url, **kwargs):
|
||||
if '/history/' in url:
|
||||
calls[0] += 1
|
||||
clock[0] += 3700
|
||||
self.completed = calls[0] >= 3
|
||||
if url.endswith('/queue'):
|
||||
return Response({'queue_running': [[0, 'prompt-1']] if calls[0] == 2 else [],
|
||||
'queue_pending': [[0, 'prompt-1']] if calls[0] == 1 else []})
|
||||
return original_get(url, **kwargs)
|
||||
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
|
||||
with patch.object(self, 'get', side_effect=get):
|
||||
result = await self.rt.run_image_list_job(context)
|
||||
self.assertEqual(len(result[0]), 2)
|
||||
self.assertGreaterEqual(clock[0], 7200)
|
||||
self.assertFalse(self.cancel_requests)
|
||||
|
||||
async def test_unrelated_busy_prompt_does_not_refresh_owned_liveness(self):
|
||||
self.completed = False
|
||||
clock = [0.0]
|
||||
original_get = self.get
|
||||
def get(url, **kwargs):
|
||||
if '/history/' in url:
|
||||
clock[0] += 70
|
||||
if url.endswith('/queue'):
|
||||
return Response({'queue_running': [[0, 'not-our-prompt']], 'queue_pending': []})
|
||||
return original_get(url, **kwargs)
|
||||
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
|
||||
with patch.object(self, 'get', side_effect=get):
|
||||
with self.assertRaisesRegex(TimeoutError, 'liveness'):
|
||||
await self.rt.run_image_list_job(self.create(1))
|
||||
self.assertEqual(self.cancel_requests, [
|
||||
('http://worker/queue', {'delete': ['prompt-1']}),
|
||||
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
|
||||
|
||||
async def test_failed_queue_cleanup_still_interrupts_and_reclaims_staging(self):
|
||||
self.completed = False
|
||||
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
|
||||
original_post = self.post
|
||||
def post(url, **kwargs):
|
||||
result = original_post(url, **kwargs)
|
||||
return Response(status=500) if url.endswith('/queue') else result
|
||||
with patch.object(self, 'post', side_effect=post):
|
||||
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
||||
await self.rt.run_image_list_job(self.create(2))
|
||||
self.assertEqual(self.cancel_requests, [
|
||||
('http://worker/queue', {'delete': ['prompt-1']}),
|
||||
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
|
||||
self.assertEqual(self.deletions, self.uploads)
|
||||
self.assertTrue(any('cleanup unconfirmed' in line for line in self.logs), self.logs)
|
||||
self.assertEqual(self.dispatch.await_count, 1)
|
||||
|
||||
async def test_configured_total_deadline_cleans_owned_work(self):
|
||||
self.completed = False
|
||||
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
|
||||
self.queue = {'queue_running': [[0, 'prompt-1']], 'queue_pending': []}
|
||||
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
|
||||
await asyncio.wait_for(self.rt.run_image_list_job(self.create(2)), 1)
|
||||
self.assertEqual(len(self.uploads), 1)
|
||||
self.assertEqual(self.deletions, self.uploads)
|
||||
self.assertFalse(self.rt._MASTER_JOBS)
|
||||
@@ -0,0 +1,106 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_config_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
|
||||
package_name = "image_list_timeout_config_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = [str(module_path.parent)]
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.logging"] = logging_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
sys.modules[f"{package_name}.constants"] = constants_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
config_module = _load_config_module()
|
||||
|
||||
|
||||
SETTING_KEY = "image_list_job_timeout_seconds"
|
||||
|
||||
|
||||
def _config_with_timeout(value):
|
||||
return {"settings": {SETTING_KEY: value}}
|
||||
|
||||
|
||||
def test_default_config_sets_image_list_job_timeout_to_zero():
|
||||
assert config_module.get_default_config()["settings"][SETTING_KEY] == 0.0
|
||||
|
||||
|
||||
def test_missing_value_defaults_to_disabled_timeout():
|
||||
assert config_module.get_image_list_job_timeout_seconds({}) == 0.0
|
||||
assert config_module.get_image_list_job_timeout_seconds({"settings": {}}) == 0.0
|
||||
|
||||
|
||||
def test_none_config_loads_current_config(monkeypatch):
|
||||
cfg = _config_with_timeout("12.5")
|
||||
monkeypatch.setattr(config_module, "load_config", lambda: cfg)
|
||||
|
||||
assert config_module.get_image_list_job_timeout_seconds() == 12.5
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[
|
||||
(0, 0.0),
|
||||
(0.0, 0.0),
|
||||
("0", 0.0),
|
||||
(" 0.0 ", 0.0),
|
||||
(1, 1.0),
|
||||
(2.5, 2.5),
|
||||
("3", 3.0),
|
||||
(" 4.25 ", 4.25),
|
||||
],
|
||||
)
|
||||
def test_accepts_zero_and_positive_numeric_values(value, expected):
|
||||
result = config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
|
||||
|
||||
assert result == expected
|
||||
assert isinstance(result, float)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
[
|
||||
True,
|
||||
False,
|
||||
-1,
|
||||
-0.01,
|
||||
float("nan"),
|
||||
float("inf"),
|
||||
float("-inf"),
|
||||
"not-a-number",
|
||||
"",
|
||||
" ",
|
||||
"-1",
|
||||
"nan",
|
||||
"inf",
|
||||
None,
|
||||
[],
|
||||
{},
|
||||
],
|
||||
)
|
||||
def test_rejects_invalid_values_with_setting_name(value):
|
||||
with pytest.raises(ValueError, match=SETTING_KEY):
|
||||
config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Context-carried values and sparse, per-item End overrides."""
|
||||
import copy
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from image_list import runtime as rt
|
||||
from image_list.execution_plan import prepare_image_list_execution
|
||||
|
||||
|
||||
def value_graph():
|
||||
# Use a minimal, closed region with an override-only ancestor.
|
||||
return {
|
||||
'source': {'class_type': 'Source', 'inputs': {}},
|
||||
'start': {'class_type': rt.START_CLASS, 'inputs': {'images': ['source', 0]}},
|
||||
'override': {'class_type': 'Transform', 'inputs': {'value': ['start', 4]}},
|
||||
'end': {'class_type': rt.END_CLASS, 'inputs': {
|
||||
'images': ['start', 0], 'context': ['start', 3], 'value_1': ['override', 0],
|
||||
'value_10': ['start', 5]}},
|
||||
}
|
||||
|
||||
|
||||
def values():
|
||||
return {'value_1': torch.arange(72, dtype=torch.float64).reshape(2, 3, 4, 3),
|
||||
'value_2': torch.ones(3, 2, 2, dtype=torch.bfloat16), 'value_3': 'name',
|
||||
'value_4': 42, 'value_5': 0.125, 'value_6': True}
|
||||
|
||||
|
||||
def assert_values(actual, expected):
|
||||
assert actual.keys() == expected.keys()
|
||||
for key, value in expected.items():
|
||||
assert type(actual[key]) is type(value)
|
||||
if isinstance(value, torch.Tensor):
|
||||
assert actual[key].dtype == value.dtype
|
||||
assert torch.equal(actual[key], value)
|
||||
else:
|
||||
assert actual[key] == value
|
||||
|
||||
|
||||
def test_normalize_and_roundtrip_all_override_kinds_keeps_batches_whole():
|
||||
expected = values()
|
||||
result = rt.normalize_worker_result(torch.zeros(1, 2, 2, 3),
|
||||
values={key: [value] for key, value in expected.items()})
|
||||
assert_values(result['values'], expected)
|
||||
wire = rt.encode_result(result)
|
||||
assert_values(rt.decode_result(wire)['values'], expected)
|
||||
assert wire['values']['value_1']['shape'][0] == 2
|
||||
assert wire['values']['value_2']['shape'][0] == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize('bad', [None, [], [None], [1, 2], [[1]], object(),
|
||||
float('nan'), float('inf'), torch.tensor([float('nan')]),
|
||||
torch.tensor([float('inf')]), torch.ones(1, dtype=torch.int32)])
|
||||
def test_invalid_override_rejected(bad):
|
||||
with pytest.raises((ValueError, TypeError)):
|
||||
rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={'value_1': bad})
|
||||
|
||||
|
||||
def test_unknown_override_names_and_wire_types_rejected():
|
||||
for name in ('value_0', 'value_11', 'values.value_1'):
|
||||
with pytest.raises(ValueError, match='slot'):
|
||||
rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={name: 1})
|
||||
base = rt.encode_result({'image': torch.zeros(1, 2, 2, 3)})
|
||||
for payload in (None, [], {'value_0': {'type': 'INT', 'value': 1}},
|
||||
{'value_1': {'type': 'LATENT', 'value': {}}},
|
||||
{'value_1': {'type': 'FLOAT', 'value': float('nan')}},
|
||||
{'value_1': {'type': 'STRING', 'value': None}},
|
||||
{'value_1': rt.encode_value(torch.tensor([float('inf')]), 'TENSOR')}):
|
||||
with pytest.raises((ValueError, TypeError)):
|
||||
rt.decode_result({**base, 'values': payload})
|
||||
|
||||
|
||||
def test_no_overrides_keeps_legacy_wire_and_does_not_encode_original_tensors():
|
||||
image = torch.zeros(1, 2, 2, 3)
|
||||
result = rt.normalize_worker_result(image)
|
||||
assert set(result) == {'image', 'mask', 'text'}
|
||||
with patch.object(rt, 'encode_value', wraps=rt.encode_value) as encode:
|
||||
wire = rt.encode_result(result)
|
||||
assert set(wire) == {'image', 'mask', 'text'}
|
||||
assert [call.args[1] for call in encode.call_args_list] == ['IMAGE', 'MASK', 'STRING']
|
||||
assert set(rt.decode_result(wire)) == {'image', 'mask', 'text'}
|
||||
|
||||
|
||||
def test_override_ancestors_survive_plan_and_worker_prompt():
|
||||
prompt = value_graph()
|
||||
frozen = copy.deepcopy(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == frozen and prompt == frozen
|
||||
assert targets == ['end']
|
||||
region = rt.validate_image_list_regions(prepared)[0]
|
||||
assert 'override' in region['branch_ids']
|
||||
worker = rt.build_worker_prompt(prepared, region, 'token', 'context', 0, 'http://master')
|
||||
assert worker['end']['inputs']['value_1'] == ['override', 0]
|
||||
assert worker['end']['inputs']['value_10'] == ['start', 5]
|
||||
assert 'source' not in worker
|
||||
for bad in ('value_0', 'value_11', 'values.value_1'):
|
||||
invalid = copy.deepcopy(prompt)
|
||||
invalid['end']['inputs'][bad] = ['start', 4]
|
||||
with pytest.raises(ValueError, match='slot'):
|
||||
rt.validate_image_list_regions(invalid)
|
||||
|
||||
|
||||
def test_missing_or_unexpected_result_override_fails_before_recording():
|
||||
prompt = value_graph()
|
||||
items = rt.normalize_image_list_inputs(torch.zeros(1, 2, 2, 3), values={'value_1': 'original'})
|
||||
context = rt.create_master_job('start', 'end', prompt, items, ['worker'], '')
|
||||
try:
|
||||
rt.register_assignment(context, 0, 'worker', 'token')
|
||||
for result_values in ({}, {'value_1': 'updated'},
|
||||
{'value_1': 'updated', 'value_10': 10, 'value_9': False}):
|
||||
with pytest.raises(ValueError, match='override'):
|
||||
rt.set_master_result(context, 0, 'token', {'image': items[0]['image'], 'values': result_values})
|
||||
assert rt.get_master_job(context).received_indexes == set()
|
||||
rt.set_master_result(context, 0, 'token', {'image': items[0]['image'],
|
||||
'values': {'value_1': 'updated', 'value_10': 10}})
|
||||
assert rt.get_master_job(context).received_indexes == {0}
|
||||
finally:
|
||||
rt.remove_master_job(context)
|
||||
|
||||
|
||||
def test_override_return_request_budget_is_aggregate():
|
||||
result = rt.normalize_worker_result(torch.zeros(1, 2, 2, 3), values={'value_1': ['x' * 500]})
|
||||
with patch.object(rt, 'MAX_REQUEST_BYTES', 400):
|
||||
with pytest.raises(ValueError, match='request'):
|
||||
rt.encode_request_payload({'result': rt.encode_result(result)})
|
||||
|
||||
|
||||
class TestResultHTTP:
|
||||
def test_missing_and_invalid_overrides_fail_through_real_http(self):
|
||||
import asyncio
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
from test_image_list_integration_gaps import load_routes
|
||||
async def check():
|
||||
_module, routes = load_routes()
|
||||
app = web.Application()
|
||||
app.add_routes(routes)
|
||||
prompt = value_graph()
|
||||
items = rt.normalize_image_list_inputs(torch.zeros(1, 2, 2, 3), values={'value_1': 'original'})
|
||||
context = rt.create_master_job('start', 'end', prompt, items, ['worker'], '')
|
||||
token = rt.create_worker_task_token()
|
||||
rt.register_assignment(context, 0, 'worker', token)
|
||||
try:
|
||||
async with TestClient(TestServer(app)) as client:
|
||||
caps = await (await client.get('/distributed/image-list/capabilities')).json()
|
||||
assert caps['value_results'] is True and caps['protocol_version'] == 2
|
||||
result = rt.encode_result({'image': items[0]['image']})
|
||||
envelope = {'context_id': context, 'token': token, 'item_index': 0, 'result': result}
|
||||
response = await client.post('/distributed/image-list/result', json=envelope)
|
||||
assert response.status == 400
|
||||
assert 'missing' in (await response.json())['error']
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
result['values'] = {'value_1': {'type': 'INT', 'value': [1, 2]},
|
||||
'value_10': {'type': 'INT', 'value': 10}}
|
||||
response = await client.post('/distributed/image-list/result', json=envelope)
|
||||
assert response.status == 400
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
result['values']['value_1'] = {'type': 'STRING', 'value': 'new'}
|
||||
response = await client.post('/distributed/image-list/result', json=envelope)
|
||||
assert response.status == 200
|
||||
assert rt.get_master_job(context).results[0]['values'] == {'value_1': 'new', 'value_10': 10}
|
||||
finally:
|
||||
rt.remove_master_job(context)
|
||||
asyncio.run(check())
|
||||
@@ -0,0 +1,212 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
|
||||
def _load_dispatch_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
|
||||
package_name = f"image_list_ws_deadline_testpkg_{uuid.uuid4().hex}"
|
||||
package_modules = {}
|
||||
|
||||
for suffix in ("", ".api", ".api.orchestration", ".utils"):
|
||||
package = types.ModuleType(package_name + suffix)
|
||||
package.__path__ = []
|
||||
package_modules[package_name + suffix] = package
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
|
||||
setattr(logging_module, "log", lambda *_args, **_kwargs: None)
|
||||
package_modules[logging_module.__name__] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
setattr(network_module, "build_worker_url", lambda *_args, **_kwargs: "http://example.invalid")
|
||||
|
||||
async def _unused_probe_worker(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
async def _unused_get_client_session(*_args, **_kwargs):
|
||||
raise RuntimeError("get_client_session must be patched by this test")
|
||||
|
||||
setattr(network_module, "probe_worker", _unused_probe_worker)
|
||||
setattr(network_module, "get_client_session", _unused_get_client_session)
|
||||
package_modules[network_module.__name__] = network_module
|
||||
|
||||
module_name = f"{package_name}.api.orchestration.dispatch"
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
package_modules[module_name] = module
|
||||
|
||||
previous_modules = {name: sys.modules.get(name) for name in package_modules}
|
||||
try:
|
||||
sys.modules.update(package_modules)
|
||||
spec.loader.exec_module(module)
|
||||
finally:
|
||||
for name, previous in previous_modules.items():
|
||||
if previous is None:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = previous
|
||||
return module
|
||||
|
||||
|
||||
dispatch = _load_dispatch_module()
|
||||
|
||||
|
||||
class WebsocketDispatchDeadlineTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.mode = None
|
||||
self.received_payload = None
|
||||
self.dispatch_received = asyncio.Event()
|
||||
self.peer_finished = asyncio.Event()
|
||||
self.irrelevant_messages_sent = 0
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/distributed/worker_ws", self._handle_websocket)
|
||||
self.runner = web.AppRunner(app)
|
||||
await self.runner.setup()
|
||||
self.site = web.TCPSite(self.runner, "127.0.0.1", 0)
|
||||
await self.site.start()
|
||||
server = cast(Any, self.site._server)
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
self.worker_url = f"http://127.0.0.1:{port}"
|
||||
self.session = aiohttp.ClientSession()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.session.close()
|
||||
await self.runner.cleanup()
|
||||
|
||||
async def _get_client_session(self):
|
||||
return self.session
|
||||
|
||||
async def _handle_websocket(self, request):
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
try:
|
||||
payload = await ws.receive_json()
|
||||
self.received_payload = payload
|
||||
self.dispatch_received.set()
|
||||
|
||||
if self.mode in ("success", "slow-close"):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": "unrelated-request",
|
||||
"ok": True,
|
||||
"prompt_id": "wrong-prompt",
|
||||
})
|
||||
await ws.send_json({"type": "heartbeat"})
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": payload["request_id"],
|
||||
"ok": True,
|
||||
"prompt_id": "prompt-123",
|
||||
})
|
||||
if self.mode == "slow-close":
|
||||
await asyncio.sleep(0.15)
|
||||
elif self.mode == "rejection":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": payload["request_id"],
|
||||
"ok": False,
|
||||
"error": "Invalid prompt",
|
||||
"validation_error": {"node": "bad"},
|
||||
"node_errors": {"bad": "missing input"},
|
||||
})
|
||||
elif self.mode in ("no-ack", "wait-for-cancel"):
|
||||
async def send_irrelevant_messages():
|
||||
while True:
|
||||
try:
|
||||
await ws.send_json({"type": "heartbeat"})
|
||||
self.irrelevant_messages_sent += 1
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": "unrelated-request",
|
||||
"ok": True,
|
||||
"prompt_id": "not-this-request",
|
||||
})
|
||||
self.irrelevant_messages_sent += 1
|
||||
except (ConnectionError, RuntimeError):
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
sender = asyncio.create_task(send_irrelevant_messages())
|
||||
try:
|
||||
async for _message in ws:
|
||||
pass
|
||||
finally:
|
||||
sender.cancel()
|
||||
try:
|
||||
await sender
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
finally:
|
||||
self.peer_finished.set()
|
||||
return ws
|
||||
|
||||
async def _dispatch(self, timeout):
|
||||
with patch.object(dispatch, "get_client_session", self._get_client_session):
|
||||
return await dispatch._dispatch_via_websocket(
|
||||
self.worker_url,
|
||||
{"prompt": {"node": {"class_type": "Example"}}, "workflow": {"name": "test"}},
|
||||
"client-456",
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def test_returns_matching_acknowledgement_prompt_id(self):
|
||||
self.mode = "success"
|
||||
|
||||
prompt_id = await self._dispatch(timeout=0.5)
|
||||
|
||||
self.assertEqual(prompt_id, "prompt-123")
|
||||
received_payload = self.received_payload
|
||||
assert received_payload is not None
|
||||
self.assertEqual(received_payload["type"], "dispatch_prompt")
|
||||
self.assertTrue(received_payload["request_id"])
|
||||
self.assertEqual(received_payload["client_id"], "client-456")
|
||||
self.assertEqual(received_payload["prompt"], {"node": {"class_type": "Example"}})
|
||||
self.assertEqual(received_payload["workflow"], {"name": "test"})
|
||||
|
||||
async def test_acknowledged_prompt_id_survives_slow_close(self):
|
||||
self.mode = "slow-close"
|
||||
self.assertEqual(await self._dispatch(timeout=0.03), "prompt-123")
|
||||
|
||||
async def test_matching_rejection_preserves_worker_error_details(self):
|
||||
self.mode = "rejection"
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "Invalid prompt.*validation_error=.*node_errors="):
|
||||
await self._dispatch(timeout=0.5)
|
||||
|
||||
async def test_repeated_irrelevant_messages_do_not_reset_ack_deadline(self):
|
||||
self.mode = "no-ack"
|
||||
started = asyncio.get_running_loop().time()
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "submission status is ambiguous"):
|
||||
await asyncio.wait_for(self._dispatch(timeout=0.12), timeout=0.5)
|
||||
|
||||
elapsed = asyncio.get_running_loop().time() - started
|
||||
self.assertLess(elapsed, 0.5)
|
||||
self.assertGreaterEqual(self.irrelevant_messages_sent, 4)
|
||||
await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
|
||||
|
||||
async def test_outer_cancellation_propagates_and_closes_peer(self):
|
||||
self.mode = "wait-for-cancel"
|
||||
task = asyncio.create_task(self._dispatch(timeout=2.0))
|
||||
await asyncio.wait_for(self.dispatch_received.wait(), timeout=0.5)
|
||||
|
||||
task.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(task, timeout=0.5)
|
||||
await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -69,6 +69,8 @@ def v3(monkeypatch):
|
||||
module(f"{PKG}.nodes.utilities")
|
||||
module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
|
||||
module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
|
||||
module(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object,
|
||||
DistributedImageListGetValue=object)
|
||||
module("comfy", __path__=[])
|
||||
module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
|
||||
module("comfy_api", __path__=[])
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
@@ -19,3 +20,32 @@ def test_native_v3_runtime():
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'V3_ACCEPTANCE_OK' in result.stdout
|
||||
|
||||
|
||||
def test_image_list_network_runtime():
|
||||
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not comfy_root:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT for model-free multi-process HTTP acceptance')
|
||||
helper = Path(__file__).parent / 'helpers' / 'image_list_network_check.py'
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
|
||||
capture_output=True, text=True, timeout=180,
|
||||
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_VALUES_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_MASTER_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_PREFLIGHT_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_SEQUENTIAL_OK' in result.stdout
|
||||
assert 'IMAGE_LIST_FAILURE_OK' in result.stdout
|
||||
reports = [json.loads(line.removeprefix('IMAGE_LIST_TARGETS_OK '))
|
||||
for line in result.stdout.splitlines() if line.startswith('IMAGE_LIST_TARGETS_OK ')]
|
||||
assert len(reports) == 1, result.stdout
|
||||
assert reports[0] == {
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': ['4', '5'], 'worker_targets': ['4'],
|
||||
'required_utility_once_per_item': True, 'preview_discarded': True,
|
||||
'saved_prompt_unchanged': True, 'capture_once_per_outer_prompt': True,
|
||||
}
|
||||
|
||||
+41
-2
@@ -105,24 +105,63 @@ class PromptValidationError(RuntimeError):
|
||||
super().__init__(f"Invalid prompt: {merged}")
|
||||
|
||||
|
||||
def validate_execution_targets(prompt, targets):
|
||||
"""Reject malformed explicit roots instead of falling back to all outputs."""
|
||||
if (not isinstance(targets, list) or not targets
|
||||
or any(not isinstance(node_id, str) or not node_id or node_id not in prompt
|
||||
for node_id in targets)
|
||||
or len(set(targets)) != len(targets)):
|
||||
raise PromptValidationError({
|
||||
"type": "invalid_partial_execution_targets",
|
||||
"message": "partial_execution_targets must be a non-empty list of unique existing node IDs",
|
||||
"details": "", "extra_info": {},
|
||||
})
|
||||
return list(targets)
|
||||
|
||||
|
||||
async def queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta=None,
|
||||
client_id=None,
|
||||
include_queue_metadata=False,
|
||||
*,
|
||||
partial_execution_targets=None,
|
||||
):
|
||||
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
||||
"""Validate and queue a prompt, preserving any explicitly requested roots."""
|
||||
targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
|
||||
if partial_execution_targets is not None else None)
|
||||
payload = {"prompt": prompt_obj}
|
||||
if targets is not None:
|
||||
payload["partial_execution_targets"] = list(targets)
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
prompt = payload["prompt"]
|
||||
if targets is not None:
|
||||
# Hooks may replace the envelope or mutate its target list. The caller's
|
||||
# requested roots remain authoritative, but must still exist afterwards.
|
||||
validate_execution_targets(prompt, targets)
|
||||
|
||||
prompt_id = str(uuid.uuid4())
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, targets)
|
||||
if not valid[0]:
|
||||
error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
raise PromptValidationError(error_payload, node_errors)
|
||||
|
||||
if targets is not None:
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
target_errors = {
|
||||
node_id: entry for node_id, entry in node_errors.items()
|
||||
if node_id in targets or (isinstance(entry, dict)
|
||||
and set(entry.get("dependent_outputs", [])) & set(targets))
|
||||
}
|
||||
if set(valid[2]) != set(targets) or target_errors:
|
||||
raise PromptValidationError({
|
||||
"type": "partial_execution_targets_failed_validation",
|
||||
"message": "Not all requested execution targets passed validation",
|
||||
"details": "Refusing to queue a subset or additional output roots",
|
||||
"extra_info": {"requested_targets": targets},
|
||||
}, node_errors)
|
||||
|
||||
extra_data = {"create_time": int(time.time() * 1000)}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
|
||||
+34
-1
@@ -5,6 +5,7 @@ import asyncio
|
||||
import copy
|
||||
import os
|
||||
import json
|
||||
import math
|
||||
from contextlib import asynccontextmanager
|
||||
from .logging import log
|
||||
from .secret_fields import is_secret_field_name as _is_secret_field_name
|
||||
@@ -81,7 +82,8 @@ def get_default_config():
|
||||
"worker_probe_concurrency": 8,
|
||||
"worker_prep_concurrency": 4,
|
||||
"media_sync_concurrency": 2,
|
||||
"media_sync_timeout_seconds": 120
|
||||
"media_sync_timeout_seconds": 120,
|
||||
"image_list_job_timeout_seconds": 0.0
|
||||
},
|
||||
"tunnel": {
|
||||
"status": "stopped",
|
||||
@@ -215,6 +217,37 @@ def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
|
||||
return max(1, int(default))
|
||||
|
||||
|
||||
def get_image_list_job_timeout_seconds(config=None) -> float:
|
||||
"""Return the validated image-list total-job timeout in seconds.
|
||||
|
||||
A value of zero disables the total-job deadline. Missing settings default
|
||||
to zero; configured values must be finite, non-negative numbers.
|
||||
"""
|
||||
setting_key = "settings.image_list_job_timeout_seconds"
|
||||
cfg = load_config() if config is None else config
|
||||
if not isinstance(cfg, dict):
|
||||
raise ValueError(f"{setting_key} must be a finite non-negative number")
|
||||
|
||||
settings = cfg.get("settings", {})
|
||||
if not isinstance(settings, dict):
|
||||
raise ValueError(f"{setting_key} must be a finite non-negative number")
|
||||
|
||||
value = settings.get("image_list_job_timeout_seconds", 0.0)
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
||||
raise ValueError(f"{setting_key} must be a finite non-negative number")
|
||||
if isinstance(value, str) and not value.strip():
|
||||
raise ValueError(f"{setting_key} must be a finite non-negative number")
|
||||
|
||||
try:
|
||||
timeout = float(value)
|
||||
except (OverflowError, ValueError):
|
||||
raise ValueError(f"{setting_key} must be a finite non-negative number") from None
|
||||
|
||||
if not math.isfinite(timeout) or timeout < 0:
|
||||
raise ValueError(f"{setting_key} must be a finite non-negative number")
|
||||
return timeout
|
||||
|
||||
|
||||
def is_master_delegate_only() -> bool:
|
||||
"""Returns True when master should skip local workload and act as orchestrator only."""
|
||||
try:
|
||||
|
||||
@@ -225,6 +225,7 @@ export const ENDPOINTS = {
|
||||
|
||||
export const NODE_CLASSES = {
|
||||
DISTRIBUTED_COLLECTOR: 'DistributedCollector',
|
||||
IMAGE_LIST_START: 'DistributedImageListStart',
|
||||
DISTRIBUTED_SEED: 'DistributedSeed',
|
||||
DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
|
||||
UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',
|
||||
|
||||
+13
-3
@@ -9,7 +9,9 @@ export function setupInterceptor(extension) {
|
||||
const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
|
||||
const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
|
||||
|
||||
if (hasCollector || hasDistUpscale) {
|
||||
const hasImageList = findNodesByClass(prompt.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
|
||||
|
||||
if (hasCollector || hasDistUpscale || hasImageList) {
|
||||
const result = await executeParallelDistributed(extension, prompt);
|
||||
// Immediate status check for instant feedback
|
||||
checkAllWorkerStatuses(extension);
|
||||
@@ -25,14 +27,22 @@ export function setupInterceptor(extension) {
|
||||
export async function executeParallelDistributed(extension, promptWrapper) {
|
||||
const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
|
||||
try {
|
||||
const enabledWorkers = extension.enabledWorkers;
|
||||
const hasImageList = findNodesByClass(promptWrapper.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
|
||||
// Image-list v1 uses direct workers only; do not contact paid provider pools.
|
||||
const enabledWorkers = hasImageList
|
||||
? extension.enabledWorkers.filter(worker => worker.type !== "serverless")
|
||||
: extension.enabledWorkers;
|
||||
extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
|
||||
|
||||
// Pre-flight health check on all enabled workers
|
||||
const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
|
||||
|
||||
if (hasImageList && activeWorkers.length === 0 && extension.config?.settings?.master_delegate_only) {
|
||||
throw new Error("Distributed Image List in orchestrator-only mode requires an active direct worker.");
|
||||
}
|
||||
|
||||
// Case: Enabled workers but all offline
|
||||
if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
|
||||
if (!hasImageList && activeWorkers.length === 0 && enabledWorkers.length > 0) {
|
||||
extension.log("No active workers found. All enabled workers are offline.");
|
||||
if (extension.ui?.showToast) {
|
||||
extension.ui.showToast(extension.app, "error", "All Workers Offline",
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
|
||||
// Distributed Image List Start declares value_1..value_10 outputs in Python so
|
||||
// saved links and prompt output indexes stay valid. ComfyUI grows the matching
|
||||
// `values.value_N` inputs natively (Autogrow) but has no growable outputs, so
|
||||
// this extension only presents them: it shows outputs for connected inputs,
|
||||
// copies each input's type onto its output, and moves output links when
|
||||
// Autogrow compacts inputs after a middle input is disconnected.
|
||||
// If this script fails, the node still runs; all outputs simply stay visible.
|
||||
|
||||
const START_CLASS = "DistributedImageListStart";
|
||||
const FIXED_OUTPUTS = 4; // image, mask, item_index, context
|
||||
const MAX_VALUES = 10;
|
||||
const INPUT_PREFIX = "values.value_";
|
||||
|
||||
function valueOrdinal(name) {
|
||||
if (typeof name !== "string" || !name.startsWith(INPUT_PREFIX)) return null;
|
||||
const ordinal = Number.parseInt(name.slice(INPUT_PREFIX.length), 10);
|
||||
return ordinal >= 1 && ordinal <= MAX_VALUES ? ordinal : null;
|
||||
}
|
||||
|
||||
function getLink(graph, id) {
|
||||
if (id == null || !graph) return undefined;
|
||||
if (typeof graph.getLink === "function") return graph.getLink(id) ?? undefined;
|
||||
const links = graph.links;
|
||||
return (typeof links?.get === "function" ? links.get(id) : links?.[id]) ?? undefined;
|
||||
}
|
||||
|
||||
function connectedValues(node) {
|
||||
const result = new Map();
|
||||
for (const input of node.inputs ?? []) {
|
||||
const ordinal = valueOrdinal(input.name);
|
||||
if (!ordinal) continue;
|
||||
const link = getLink(node.graph, input.link);
|
||||
if (!link) continue;
|
||||
const origin = node.graph.getNodeById?.(link.origin_id);
|
||||
const originType = origin?.outputs?.[link.origin_slot]?.type;
|
||||
const type = originType && originType !== "*" ? originType : (link.type ?? "*");
|
||||
result.set(ordinal, { type, key: `${link.origin_id}:${link.origin_slot}` });
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
function outputLinks(node, ordinal) {
|
||||
return node.outputs?.[FIXED_OUTPUTS + ordinal - 1]?.links ?? [];
|
||||
}
|
||||
|
||||
function typesCompatible(outputType, inputType) {
|
||||
if (!outputType || !inputType || outputType === "*" || inputType === "*") return true;
|
||||
const accepted = String(inputType).split(",");
|
||||
return String(outputType).split(",").some(type => accepted.includes(type));
|
||||
}
|
||||
|
||||
function moveOutputLinks(node, from, to) {
|
||||
const graph = node.graph;
|
||||
const targets = outputLinks(node, from)
|
||||
.map(id => getLink(graph, id))
|
||||
.filter(Boolean)
|
||||
.map(link => [graph.getNodeById?.(link.target_id), link.target_slot])
|
||||
.filter(([target]) => target);
|
||||
// The moved links carry the same value, so give the destination output its
|
||||
// type first; LiteGraph refuses connections to a mismatched output type.
|
||||
node.outputs[FIXED_OUTPUTS + to - 1].type = node.outputs[FIXED_OUTPUTS + from - 1].type;
|
||||
// Connecting a target input replaces its existing link from `from`.
|
||||
for (const [target, slot] of targets) node.connect(FIXED_OUTPUTS + to - 1, target, slot);
|
||||
}
|
||||
|
||||
function disconnectValueOutput(node, ordinal) {
|
||||
if (outputLinks(node, ordinal).length) node.disconnectOutput(FIXED_OUTPUTS + ordinal - 1);
|
||||
}
|
||||
|
||||
// Autogrow removes a disconnected middle input and shifts later input links
|
||||
// down by one. Shift the matching outputs' links the same way, and drop the
|
||||
// removed value's downstream links, so no consumer silently receives a
|
||||
// different value.
|
||||
function followCompactedInputs(node, current) {
|
||||
const previous = node._distributedValueSources;
|
||||
if (!previous) return;
|
||||
const now = [];
|
||||
for (const [ordinal, info] of current) now[ordinal - 1] = info.key;
|
||||
const before = previous;
|
||||
if (now.length !== before.length - 1) return;
|
||||
let removed = 0;
|
||||
while (removed < now.length && now[removed] === before[removed]) removed++;
|
||||
// Removing the last value shifts nothing; keep its links so a temporary
|
||||
// rewire does not silently drop downstream connections.
|
||||
if (removed >= now.length) return;
|
||||
for (let i = removed; i < now.length; i++) {
|
||||
if (now[i] !== before[i + 1]) return; // Not a single Autogrow compaction.
|
||||
}
|
||||
disconnectValueOutput(node, removed + 1);
|
||||
for (let i = removed + 1; i < before.length; i++) moveOutputLinks(node, i + 1, i);
|
||||
}
|
||||
|
||||
function syncImageListStartOutputs(node) {
|
||||
if (!node?.graph || !Array.isArray(node.outputs) || node.outputs.length < FIXED_OUTPUTS) return;
|
||||
const current = connectedValues(node);
|
||||
followCompactedInputs(node, current);
|
||||
const sources = [];
|
||||
for (const [ordinal, info] of current) sources[ordinal - 1] = info.key;
|
||||
node._distributedValueSources = Array.from(sources);
|
||||
|
||||
let visible = Math.max(0, ...current.keys());
|
||||
for (let ordinal = MAX_VALUES; ordinal > visible; ordinal--) {
|
||||
// Never drop a wired output; its links keep it visible until removed.
|
||||
if (outputLinks(node, ordinal).length) visible = ordinal;
|
||||
}
|
||||
const listShape = node.outputs[0]?.shape;
|
||||
while (node.outputs.length < FIXED_OUTPUTS + visible) {
|
||||
const ordinal = node.outputs.length - FIXED_OUTPUTS + 1;
|
||||
node.addOutput(`value_${ordinal}`, "*", listShape === undefined ? undefined : { shape: listShape });
|
||||
}
|
||||
while (node.outputs.length > FIXED_OUTPUTS + visible) node.removeOutput(node.outputs.length - 1);
|
||||
|
||||
for (let ordinal = 1; ordinal <= visible; ordinal++) {
|
||||
const output = node.outputs[FIXED_OUTPUTS + ordinal - 1];
|
||||
const type = current.get(ordinal)?.type ?? "*";
|
||||
output.name = `value_${ordinal}`;
|
||||
if (output.type !== type) {
|
||||
output.type = type;
|
||||
for (const id of [...(output.links ?? [])]) {
|
||||
const link = getLink(node.graph, id);
|
||||
if (!link) continue;
|
||||
const target = node.graph.getNodeById?.(link.target_id);
|
||||
const targetType = target?.inputs?.[link.target_slot]?.type;
|
||||
if (target && !typesCompatible(type, targetType)) target.disconnectInput(link.target_slot);
|
||||
else link.type = type;
|
||||
}
|
||||
}
|
||||
}
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function scheduleSync(node) {
|
||||
if (node._distributedValueSyncPending) return;
|
||||
node._distributedValueSyncPending = true;
|
||||
const run = () => {
|
||||
node._distributedValueSyncPending = false;
|
||||
try {
|
||||
syncImageListStartOutputs(node);
|
||||
} catch (error) {
|
||||
console.error("[Distributed] Image List Start output sync failed:", error);
|
||||
}
|
||||
};
|
||||
// Run after Autogrow's own animation-frame input compaction.
|
||||
if (typeof requestAnimationFrame === "function") requestAnimationFrame(() => setTimeout(run, 0));
|
||||
else setTimeout(run, 0);
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "Distributed.ImageListStartValues",
|
||||
nodeCreated(node) {
|
||||
if (node.comfyClass !== START_CLASS) return;
|
||||
const onConnectionsChange = node.onConnectionsChange;
|
||||
node.onConnectionsChange = function (...args) {
|
||||
const result = onConnectionsChange?.apply(this, args);
|
||||
scheduleSync(this);
|
||||
return result;
|
||||
};
|
||||
const onConfigure = node.onConfigure;
|
||||
node.onConfigure = function (...args) {
|
||||
const result = onConfigure?.apply(this, args);
|
||||
// A loaded graph is the source of truth for prior wiring.
|
||||
this._distributedValueSources = undefined;
|
||||
scheduleSync(this);
|
||||
return result;
|
||||
};
|
||||
scheduleSync(node);
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,174 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
|
||||
// End's Python schema owns all ten optional override slots. Only presentation
|
||||
// grows here: unlike Start's Autogrow inputs, End keys must never compact, since
|
||||
// value_3 always overrides source value_3, even if value_2 is disconnected.
|
||||
const START = "DistributedImageListStart";
|
||||
const END = "DistributedImageListEnd";
|
||||
const GET = "DistributedImageListGetValue";
|
||||
const PIPE = "DISTRIBUTED_IMAGE_LIST_VALUES";
|
||||
const VALUE_TYPES = "IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN";
|
||||
const MAX_VALUES = 10;
|
||||
const pendingGraphs = new WeakSet();
|
||||
|
||||
function ordinal(name) {
|
||||
const match = /^value_([1-9]|10)$/.exec(name ?? "");
|
||||
return match ? Number(match[1]) : null;
|
||||
}
|
||||
|
||||
function getLink(graph, id) {
|
||||
if (id == null) return undefined;
|
||||
if (typeof graph?.getLink === "function") return graph.getLink(id);
|
||||
return graph?.links?.get?.(id) ?? graph?.links?.[id];
|
||||
}
|
||||
|
||||
function sourceOf(node, inputName) {
|
||||
const input = node.inputs?.find(slot => slot.name === inputName);
|
||||
const link = getLink(node.graph, input?.link);
|
||||
const source = link && node.graph?.getNodeById?.(link.origin_id);
|
||||
return source ? { node: source, slot: link.origin_slot, link } : null;
|
||||
}
|
||||
|
||||
function isReroute(node) {
|
||||
return [node?.comfyClass, node?.type].includes("Reroute");
|
||||
}
|
||||
|
||||
function throughReroutes(source) {
|
||||
const seen = new Set();
|
||||
while (source && isReroute(source.node)) {
|
||||
if (seen.has(source.node)) return null;
|
||||
seen.add(source.node);
|
||||
source = sourceOf(source.node, source.node.inputs?.[0]?.name);
|
||||
}
|
||||
return source;
|
||||
}
|
||||
|
||||
function sourceType(source, seen) {
|
||||
source = throughReroutes(source);
|
||||
if (!source) return "*";
|
||||
const { node, slot, link } = source;
|
||||
if (node.comfyClass === GET) return getterType(node, seen);
|
||||
const type = node.outputs?.[slot]?.type;
|
||||
return type && type !== "*" ? type : (link.type || "*");
|
||||
}
|
||||
|
||||
function getterType(node, seen = new Set()) {
|
||||
if (seen.has(node)) return "*";
|
||||
seen.add(node);
|
||||
const name = node.widgets?.find(widget => widget.name === "value_name")?.value;
|
||||
if (!ordinal(name)) return "*";
|
||||
const endSource = throughReroutes(sourceOf(node, "values"));
|
||||
if (endSource?.node.comfyClass !== END || endSource.slot !== 3) return "*";
|
||||
const end = endSource.node;
|
||||
const override = sourceOf(end, name);
|
||||
if (override) return sourceType(override, seen);
|
||||
const startSource = throughReroutes(sourceOf(end, "context"));
|
||||
if (startSource?.node.comfyClass !== START || startSource.slot !== 3) return "*";
|
||||
return sourceType(sourceOf(startSource.node, `values.${name}`), seen);
|
||||
}
|
||||
|
||||
function compatible(outputType, inputType) {
|
||||
if (!outputType || !inputType || outputType === "*" || inputType === "*") return true;
|
||||
return String(outputType).split(",").some(type => String(inputType).split(",").includes(type));
|
||||
}
|
||||
|
||||
function syncGetter(node) {
|
||||
const output = node.outputs?.[0];
|
||||
if (!output) return;
|
||||
const type = getterType(node);
|
||||
if (output.type === type) return;
|
||||
output.type = type;
|
||||
for (const id of [...(output.links ?? [])]) {
|
||||
const link = getLink(node.graph, id);
|
||||
if (!link) continue;
|
||||
const target = node.graph?.getNodeById?.(link.target_id);
|
||||
if (target && !compatible(type, target.inputs?.[link.target_slot]?.type)) {
|
||||
target.disconnectInput(link.target_slot);
|
||||
} else {
|
||||
link.type = type;
|
||||
}
|
||||
}
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function syncEnd(node) {
|
||||
if (!Array.isArray(node.inputs) || !Array.isArray(node.outputs)) return;
|
||||
// Old saved End nodes have only the three original outputs. Append rather
|
||||
// than replacing any socket, preserving every existing output index/link.
|
||||
if (node.outputs.length === 3) {
|
||||
const shape = node.outputs[0]?.shape;
|
||||
node.addOutput("values", PIPE, shape === undefined ? undefined : { shape });
|
||||
}
|
||||
const connected = node.inputs.filter(input => ordinal(input.name) && input.link != null);
|
||||
const highest = Math.max(0, ...connected.map(input => ordinal(input.name)));
|
||||
const visible = Math.min(MAX_VALUES, Math.max(highest + 1, node._distributedVisibleOverrides ?? 1));
|
||||
for (let i = node.inputs.length - 1; i >= 0; i--) {
|
||||
const input = node.inputs[i];
|
||||
if (ordinal(input.name) > visible && input.link == null) node.removeInput(i);
|
||||
}
|
||||
for (let i = 1; i <= visible; i++) {
|
||||
if (!node.inputs.some(input => input.name === `value_${i}`)) {
|
||||
node.addInput(`value_${i}`, VALUE_TYPES);
|
||||
}
|
||||
}
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function schedule(graph) {
|
||||
if (!graph || pendingGraphs.has(graph)) return;
|
||||
pendingGraphs.add(graph);
|
||||
const run = () => {
|
||||
// Keep the guard while syncing, since disconnectInput can re-enter the
|
||||
// connection hook. A fresh user event schedules a fresh traversal.
|
||||
try {
|
||||
for (const node of graph._nodes ?? []) if (node.comfyClass === END) syncEnd(node);
|
||||
for (const node of graph._nodes ?? []) if (node.comfyClass === GET) syncGetter(node);
|
||||
} catch (error) {
|
||||
console.error("[Distributed] Image List values presentation failed:", error);
|
||||
} finally {
|
||||
pendingGraphs.delete(graph);
|
||||
}
|
||||
};
|
||||
// Let native Start Autogrow and its existing output-sync finish first.
|
||||
if (typeof requestAnimationFrame === "function") requestAnimationFrame(() => setTimeout(run, 0));
|
||||
else setTimeout(run, 0);
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "Distributed.ImageListEndValues",
|
||||
nodeCreated(node) {
|
||||
if (![START, END, GET].includes(node.comfyClass) && !isReroute(node)) return;
|
||||
if (node.comfyClass === END) {
|
||||
node.addWidget("button", "Add value override", null, () => {
|
||||
const shown = Math.max(0, ...(node.inputs ?? []).map(input => ordinal(input.name) ?? 0));
|
||||
node._distributedVisibleOverrides = Math.min(MAX_VALUES, shown + 1);
|
||||
schedule(node.graph);
|
||||
}, { serialize: false });
|
||||
}
|
||||
for (const hook of ["onConnectionsChange", "onConfigure", "onAdded", "onWidgetChanged"]) {
|
||||
const original = node[hook];
|
||||
node[hook] = function (...args) {
|
||||
const result = original?.apply(this, args);
|
||||
schedule(this.graph);
|
||||
return result;
|
||||
};
|
||||
}
|
||||
// Widget callbacks are also used by frontend versions that do not call
|
||||
// the node's onWidgetChanged hook for a combo selection.
|
||||
if (node.comfyClass === GET) {
|
||||
const widget = node.widgets?.find(widget => widget.name === "value_name");
|
||||
if (widget) {
|
||||
const callback = widget.callback;
|
||||
widget.callback = function (...args) {
|
||||
const result = callback?.apply(this, args);
|
||||
schedule(node.graph);
|
||||
return result;
|
||||
};
|
||||
}
|
||||
}
|
||||
schedule(node.graph);
|
||||
},
|
||||
loadedGraphNode(node) {
|
||||
schedule(node.graph);
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,83 @@
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest';
|
||||
vi.mock('../../scripts/api.js', () => ({ api: { clientId: 'test-client' } }), { virtual: true });
|
||||
vi.mock('../workerLifecycle.js', () => ({
|
||||
checkAllWorkerStatuses: vi.fn(), getWorkerUrl: () => 'http://worker.invalid',
|
||||
}));
|
||||
vi.mock('../workerUtils.js', () => ({
|
||||
applyProbeResultToWorkerDot: vi.fn(),
|
||||
findNodesByClass: (prompt, cls) => Object.values(prompt).filter(node => node.class_type === cls),
|
||||
}));
|
||||
import { api } from '../../scripts/api.js';
|
||||
import { setupInterceptor } from '../executionUtils.js';
|
||||
|
||||
const prompt = { output: { '1': { class_type: 'DistributedImageListStart', inputs: {} } } };
|
||||
function extension(workers, online = true) {
|
||||
return {
|
||||
isEnabled: true, enabledWorkers: workers, log: vi.fn(), config: { settings: {} },
|
||||
originalQueuePrompt: vi.fn().mockResolvedValue({ prompt_id: 'local' }),
|
||||
api: { probeWorker: vi.fn().mockResolvedValue({ ok: online }),
|
||||
getServerlessStatus: vi.fn(),
|
||||
queueDistributed: vi.fn().mockResolvedValue({ prompt_id: 'distributed' }) },
|
||||
ui: { showToast: vi.fn() },
|
||||
};
|
||||
}
|
||||
afterEach(() => vi.useRealTimers());
|
||||
describe('image-list queue interception', () => {
|
||||
it('routes image-list graphs through backend orchestration', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: ['one'] }));
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('does not fall back when orchestrator-only workers are offline', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }], false);
|
||||
ext.config.settings.master_delegate_only = true;
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
|
||||
});
|
||||
it('does not probe serverless providers for an image-list job', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'pool', type: 'serverless' }]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: [], delegate_master: false }));
|
||||
expect(ext.api.getServerlessStatus).not.toHaveBeenCalled();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it.each([{ workers: [] }, { workers: [{ id: 'one', type: 'local' }] }])('queues master-only through orchestration when no workers are online: $workers', async ({ workers }) => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension(workers, false);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: [], delegate_master: false }));
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('passes orchestrator-only mode to backend when workers are online', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([{ id: 'one', type: 'local' }]);
|
||||
ext.config.settings.master_delegate_only = true;
|
||||
setupInterceptor(ext);
|
||||
await api.queuePrompt(0, prompt);
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ delegate_master: true }));
|
||||
});
|
||||
it('does not bypass backend errors with local execution', async () => {
|
||||
vi.useFakeTimers();
|
||||
const ext = extension([]);
|
||||
ext.api.queueDistributed.mockRejectedValue(new Error('invalid region'));
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, prompt)).toBeNull();
|
||||
expect(ext.api.queueDistributed).toHaveBeenCalled();
|
||||
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
|
||||
});
|
||||
it('preserves ordinary prompt handling', async () => {
|
||||
const ext = extension([]);
|
||||
setupInterceptor(ext);
|
||||
expect(await api.queuePrompt(0, { output: {} })).toEqual({ prompt_id: 'local' });
|
||||
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,174 @@
|
||||
import { readFileSync } from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
const source = readFileSync(new URL("../imageListStart.js", import.meta.url), "utf8")
|
||||
.replace(/^import .*;\r?\n/gm, "");
|
||||
|
||||
// Minimal LiteGraph-like graph: enough link bookkeeping to exercise output
|
||||
// add/remove, retyping, and link moves the extension performs.
|
||||
function setup() {
|
||||
let extension;
|
||||
const tasks = [];
|
||||
const links = new Map();
|
||||
let nextLink = 1;
|
||||
const nodes = new Map();
|
||||
const graph = {
|
||||
links,
|
||||
getLink: id => links.get(id),
|
||||
getNodeById: id => nodes.get(id),
|
||||
};
|
||||
const makeSource = (id, type) => {
|
||||
const node = { id, outputs: [{ type, links: [] }] };
|
||||
nodes.set(id, node);
|
||||
return node;
|
||||
};
|
||||
const makeTarget = (id, type) => {
|
||||
const node = {
|
||||
id, inputs: [{ type, link: null }],
|
||||
disconnectInput(slot) {
|
||||
const link = links.get(this.inputs[slot].link);
|
||||
if (!link) return;
|
||||
const origin = nodes.get(link.origin_id);
|
||||
origin.outputs[link.origin_slot].links =
|
||||
origin.outputs[link.origin_slot].links.filter(l => l !== link.id);
|
||||
links.delete(link.id);
|
||||
this.inputs[slot].link = null;
|
||||
},
|
||||
};
|
||||
nodes.set(id, node);
|
||||
return node;
|
||||
};
|
||||
const fixed = ["image", "mask", "item_index", "context"].map(name => ({ name, type: "IMAGE", links: [], shape: 6 }));
|
||||
const start = {
|
||||
id: 1, comfyClass: "DistributedImageListStart", graph,
|
||||
inputs: [{ name: "images", link: null }, { name: "mask", link: null }, { name: "values.value_1", link: null }],
|
||||
outputs: [...fixed, ...Array.from({ length: 10 }, (_, i) => ({ name: `value_${i + 1}`, type: "*", links: [] }))],
|
||||
addOutput(name, type, extra) { this.outputs.push({ name, type, links: [], ...(extra ?? {}) }); },
|
||||
removeOutput(slot) {
|
||||
if (this.outputs[slot].links.length) throw new Error("removed a wired output");
|
||||
this.outputs.splice(slot, 1);
|
||||
},
|
||||
connect(slot, target, targetSlot) {
|
||||
// LiteGraph refuses a link whose output type does not fit the input.
|
||||
const outType = this.outputs[slot].type;
|
||||
const inType = target.inputs[targetSlot].type;
|
||||
if (outType !== "*" && inType !== "*" && !String(inType).split(",").includes(outType)) return null;
|
||||
target.disconnectInput(targetSlot);
|
||||
const id = nextLink++;
|
||||
links.set(id, { id, origin_id: this.id, origin_slot: slot, target_id: target.id,
|
||||
target_slot: targetSlot, type: this.outputs[slot].type });
|
||||
this.outputs[slot].links.push(id);
|
||||
target.inputs[targetSlot].link = id;
|
||||
return links.get(id);
|
||||
},
|
||||
disconnectOutput(slot) {
|
||||
for (const id of [...this.outputs[slot].links]) {
|
||||
const link = links.get(id);
|
||||
nodes.get(link.target_id).disconnectInput(link.target_slot);
|
||||
}
|
||||
},
|
||||
setDirtyCanvas() {},
|
||||
};
|
||||
nodes.set(1, start);
|
||||
vm.runInNewContext(source, {
|
||||
app: { registerExtension: ext => { extension = ext; } },
|
||||
setTimeout: callback => { tasks.push(callback); },
|
||||
console,
|
||||
}, { filename: "imageListStart.js" });
|
||||
extension.nodeCreated(start);
|
||||
const flush = () => { while (tasks.length) tasks.shift()(); };
|
||||
// Wire a source output into Start's value input, as Autogrow would.
|
||||
const connectValue = (ordinal, sourceNode) => {
|
||||
const name = `values.value_${ordinal}`;
|
||||
let input = start.inputs.find(i => i.name === name);
|
||||
if (!input) { input = { name, link: null }; start.inputs.push(input); }
|
||||
const id = nextLink++;
|
||||
links.set(id, { id, origin_id: sourceNode.id, origin_slot: 0, target_id: 1,
|
||||
target_slot: start.inputs.indexOf(input), type: sourceNode.outputs[0].type });
|
||||
sourceNode.outputs[0].links.push(id);
|
||||
input.link = id;
|
||||
if (!start.inputs.some(i => i.name === `values.value_${ordinal + 1}`)) {
|
||||
start.inputs.push({ name: `values.value_${ordinal + 1}`, link: null });
|
||||
}
|
||||
start.onConnectionsChange(1, start.inputs.indexOf(input), true);
|
||||
flush();
|
||||
};
|
||||
// Mimic Autogrow's compaction after disconnecting a middle input.
|
||||
const disconnectValueAndCompact = ordinal => {
|
||||
const valueInputs = start.inputs.filter(i => i.name.startsWith("values."));
|
||||
const linksBefore = valueInputs.map(i => i.link);
|
||||
links.delete(linksBefore[ordinal - 1]);
|
||||
const remaining = linksBefore.filter((l, i) => i !== ordinal - 1 && l != null);
|
||||
start.inputs = start.inputs.filter(i => !i.name.startsWith("values."));
|
||||
remaining.forEach((l, i) => start.inputs.push({ name: `values.value_${i + 1}`, link: l }));
|
||||
start.inputs.push({ name: `values.value_${remaining.length + 1}`, link: null });
|
||||
start.onConnectionsChange(1, 0, false);
|
||||
flush();
|
||||
};
|
||||
flush();
|
||||
return { start, makeSource, makeTarget, connectValue, disconnectValueAndCompact, links, flush };
|
||||
}
|
||||
|
||||
const valueOutputs = start => start.outputs.slice(4).map(o => `${o.name}:${o.type}`);
|
||||
|
||||
describe("Distributed Image List Start value outputs", () => {
|
||||
it("shows no value outputs until a value input is connected", () => {
|
||||
const h = setup();
|
||||
expect(h.start.outputs.map(o => o.name)).toEqual(["image", "mask", "item_index", "context"]);
|
||||
});
|
||||
|
||||
it("adds one typed output per connected value input", () => {
|
||||
const h = setup();
|
||||
h.connectValue(1, h.makeSource(10, "IMAGE"));
|
||||
h.connectValue(2, h.makeSource(11, "STRING"));
|
||||
expect(valueOutputs(h.start)).toEqual(["value_1:IMAGE", "value_2:STRING"]);
|
||||
expect(h.start.outputs[4].shape).toBe(6);
|
||||
});
|
||||
|
||||
it("keeps downstream wiring on the same value when a middle input is removed", () => {
|
||||
const h = setup();
|
||||
h.connectValue(1, h.makeSource(10, "IMAGE"));
|
||||
h.connectValue(2, h.makeSource(11, "STRING"));
|
||||
h.connectValue(3, h.makeSource(12, "INT"));
|
||||
const imageUser = h.makeTarget(20, "IMAGE");
|
||||
const stringUser = h.makeTarget(21, "STRING");
|
||||
const intUser = h.makeTarget(22, "INT");
|
||||
h.start.connect(4, imageUser, 0);
|
||||
h.start.connect(5, stringUser, 0);
|
||||
h.start.connect(6, intUser, 0);
|
||||
|
||||
h.disconnectValueAndCompact(2);
|
||||
|
||||
expect(valueOutputs(h.start)).toEqual(["value_1:IMAGE", "value_2:INT"]);
|
||||
expect(h.links.get(imageUser.inputs[0].link).origin_slot).toBe(4);
|
||||
expect(h.links.get(intUser.inputs[0].link).origin_slot).toBe(5);
|
||||
expect(stringUser.inputs[0].link).toBe(null);
|
||||
});
|
||||
|
||||
it("disconnects a consumer whose type no longer matches the input", () => {
|
||||
const h = setup();
|
||||
h.connectValue(1, h.makeSource(10, "IMAGE"));
|
||||
const imageUser = h.makeTarget(20, "IMAGE");
|
||||
h.start.connect(4, imageUser, 0);
|
||||
// Replace value_1's source with a STRING source.
|
||||
const input = h.start.inputs.find(i => i.name === "values.value_1");
|
||||
h.links.delete(input.link);
|
||||
input.link = null;
|
||||
h.connectValue(1, h.makeSource(11, "STRING"));
|
||||
expect(valueOutputs(h.start)).toEqual(["value_1:STRING"]);
|
||||
expect(imageUser.inputs[0].link).toBe(null);
|
||||
});
|
||||
|
||||
it("never removes an output that still has links", () => {
|
||||
const h = setup();
|
||||
h.connectValue(1, h.makeSource(10, "IMAGE"));
|
||||
h.start.connect(4, h.makeTarget(20, "*"), 0);
|
||||
const input = h.start.inputs.find(i => i.name === "values.value_1");
|
||||
h.links.delete(input.link);
|
||||
input.link = null;
|
||||
h.start.onConnectionsChange(1, 2, false);
|
||||
h.flush();
|
||||
expect(valueOutputs(h.start)).toEqual(["value_1:*"]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,190 @@
|
||||
import { readFileSync } from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
const source = readFileSync(new URL("../imageListValues.js", import.meta.url), "utf8")
|
||||
.replace(/^import .*;\r?\n/gm, "");
|
||||
const PIPE = "DISTRIBUTED_IMAGE_LIST_VALUES";
|
||||
|
||||
function setup() {
|
||||
let extension;
|
||||
let nextId = 1;
|
||||
const tasks = [];
|
||||
const graph = {
|
||||
_nodes: [], links: new Map(),
|
||||
getNodeById(id) { return this._nodes.find(node => node.id === id); },
|
||||
getLink(id) { return this.links.get(id); },
|
||||
};
|
||||
vm.runInNewContext(source, {
|
||||
app: { registerExtension: ext => { extension = ext; } },
|
||||
setTimeout: callback => tasks.push(callback), console,
|
||||
});
|
||||
function node(comfyClass, inputs, outputs, widgets = []) {
|
||||
const n = {
|
||||
id: nextId++, comfyClass, graph, widgets,
|
||||
inputs: inputs.map(name => ({ name, type: "*", link: null })),
|
||||
outputs: outputs.map(([name, type]) => ({ name, type, links: [], shape: 6 })),
|
||||
addWidget(type, name, value, callback, options) { this.widgets.push({ type, name, value, callback, options }); },
|
||||
addInput(name, type) { this.inputs.push({ name, type, link: null }); },
|
||||
addOutput(name, type, extra) { this.outputs.push({ name, type, links: [], ...extra }); },
|
||||
removeInput(slot) {
|
||||
if (this.inputs[slot].link != null) throw new Error("Removed wired input");
|
||||
this.inputs.splice(slot, 1);
|
||||
this.inputs.forEach((input, index) => {
|
||||
if (input.link != null) graph.links.get(input.link).target_slot = index;
|
||||
});
|
||||
},
|
||||
disconnectInput(slot) {
|
||||
const input = this.inputs[slot];
|
||||
const link = graph.links.get(input.link);
|
||||
if (link) {
|
||||
const out = graph.getNodeById(link.origin_id).outputs[link.origin_slot];
|
||||
out.links = out.links.filter(id => id !== link.id);
|
||||
graph.links.delete(link.id);
|
||||
}
|
||||
input.link = null;
|
||||
this.onConnectionsChange?.();
|
||||
},
|
||||
setDirtyCanvas() {},
|
||||
};
|
||||
graph._nodes.push(n);
|
||||
extension.nodeCreated(n);
|
||||
return n;
|
||||
}
|
||||
const flush = () => { while (tasks.length) tasks.shift()(); };
|
||||
const connect = (from, slot, to, name) => {
|
||||
const targetSlot = to.inputs.findIndex(input => input.name === name);
|
||||
if (targetSlot < 0) throw new Error(`Missing input ${name}`);
|
||||
to.disconnectInput(targetSlot);
|
||||
const id = nextId++;
|
||||
graph.links.set(id, { id, origin_id: from.id, origin_slot: slot, target_id: to.id,
|
||||
target_slot: targetSlot, type: from.outputs[slot].type });
|
||||
from.outputs[slot].links.push(id);
|
||||
to.inputs[targetSlot].link = id;
|
||||
from.onConnectionsChange?.();
|
||||
to.onConnectionsChange?.();
|
||||
return id;
|
||||
};
|
||||
const start = node("DistributedImageListStart", ["images", "mask", "values.value_1", "values.value_2", "values.value_3"],
|
||||
[["image", "IMAGE"], ["mask", "MASK"], ["item_index", "INT"], ["context", "DISTRIBUTED_IMAGE_LIST_CONTEXT"]]);
|
||||
const end = node("DistributedImageListEnd", ["images", "masks", "texts", "context", ...Array.from({ length: 10 }, (_, i) => `value_${i + 1}`)],
|
||||
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"], ["values", PIPE]]);
|
||||
const get = node("DistributedImageListGetValue", ["values"], [["value", "*"]], [{ name: "value_name", value: "value_1" }]);
|
||||
connect(start, 3, end, "context");
|
||||
connect(end, 3, get, "values");
|
||||
flush();
|
||||
return { node, graph, start, end, get, flush, connect, extension };
|
||||
}
|
||||
|
||||
const valueNames = node => node.inputs.filter(input => input.name.startsWith("value_")).map(input => input.name);
|
||||
const addOverride = h => { h.end.widgets.find(widget => widget.name === "Add value override").callback(); h.flush(); };
|
||||
|
||||
describe("Image List End context values", () => {
|
||||
it("keeps original sockets and shows only the first optional override initially", () => {
|
||||
const h = setup();
|
||||
expect(h.end.inputs.map(input => input.name)).toEqual(["images", "masks", "texts", "context", "value_1"]);
|
||||
expect(h.end.outputs.map(output => output.name)).toEqual(["image_list", "masks", "texts", "values"]);
|
||||
expect(h.end.widgets[0].options.serialize).toBe(false);
|
||||
});
|
||||
|
||||
it("appends the values output when loading an old three-output End", () => {
|
||||
const h = setup();
|
||||
const old = h.node("DistributedImageListEnd", ["images", "masks", "texts", "context"],
|
||||
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"]]);
|
||||
h.flush();
|
||||
expect(old.outputs[3]).toMatchObject({ name: "values", type: PIPE, shape: 6 });
|
||||
expect(valueNames(old)).toEqual(["value_1"]);
|
||||
});
|
||||
|
||||
it("grows overrides without compacting field identities or existing links", () => {
|
||||
const h = setup();
|
||||
addOverride(h);
|
||||
addOverride(h);
|
||||
const src = h.node("PrimitiveInt", [], [["INT", "INT"]]);
|
||||
const thirdLink = h.connect(src, 0, h.end, "value_3");
|
||||
h.connect(src, 0, h.end, "value_2");
|
||||
h.flush();
|
||||
h.end.disconnectInput(h.end.inputs.findIndex(input => input.name === "value_2"));
|
||||
h.flush();
|
||||
expect(valueNames(h.end)).toEqual(["value_1", "value_2", "value_3", "value_4"]);
|
||||
const link = h.graph.links.get(thirdLink);
|
||||
expect(h.end.inputs[link.target_slot].name).toBe("value_3");
|
||||
for (let i = 0; i < 20; i++) addOverride(h);
|
||||
expect(valueNames(h.end)).toHaveLength(10);
|
||||
});
|
||||
|
||||
it.each(["IMAGE", "MASK", "STRING", "INT", "FLOAT", "BOOLEAN"])("infers %s from the original Start value through context", type => {
|
||||
const h = setup();
|
||||
const source = h.node("Source", [], [["value", type]]);
|
||||
h.connect(source, 0, h.start, "values.value_1");
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe(type);
|
||||
});
|
||||
|
||||
it("prefers the End override type, then returns to the original after disconnect", () => {
|
||||
const h = setup();
|
||||
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_1");
|
||||
h.connect(h.node("Source", [], [["value", "STRING"]]), 0, h.end, "value_1");
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe("STRING");
|
||||
h.end.disconnectInput(h.end.inputs.findIndex(input => input.name === "value_1"));
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe("INT");
|
||||
});
|
||||
|
||||
it("updates a selected field and disconnects incompatible consumers", () => {
|
||||
const h = setup();
|
||||
h.connect(h.node("Source", [], [["value", "STRING"]]), 0, h.start, "values.value_1");
|
||||
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_2");
|
||||
h.flush();
|
||||
const consumer = h.node("Save", ["filename"], []);
|
||||
consumer.inputs[0].type = "STRING";
|
||||
h.connect(h.get, 0, consumer, "filename");
|
||||
h.get.widgets[0].value = "value_2";
|
||||
h.get.widgets[0].callback();
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe("INT");
|
||||
expect(consumer.inputs[0].link).toBe(null);
|
||||
});
|
||||
|
||||
it("keeps compatible consumer links and uses wildcard for unavailable graph metadata", () => {
|
||||
const h = setup();
|
||||
h.connect(h.node("Source", [], [["value", "INT"]]), 0, h.start, "values.value_1");
|
||||
h.flush();
|
||||
const consumer = h.node("Consumer", ["value"], []);
|
||||
consumer.inputs[0].type = "INT,FLOAT";
|
||||
const id = h.connect(h.get, 0, consumer, "value");
|
||||
h.get.disconnectInput(0);
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe("*");
|
||||
expect(h.graph.links.has(id)).toBe(true);
|
||||
});
|
||||
|
||||
it("follows reroutes for context and values, and terminates on a cycle", () => {
|
||||
const h = setup();
|
||||
h.connect(h.node("Source", [], [["value", "FLOAT"]]), 0, h.start, "values.value_1");
|
||||
const reroute = h.node("Reroute", ["input"], [["output", "*"]]);
|
||||
h.connect(h.end, 3, reroute, "input");
|
||||
h.connect(reroute, 0, h.get, "values");
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe("FLOAT");
|
||||
h.connect(reroute, 0, reroute, "input");
|
||||
h.flush();
|
||||
expect(h.get.outputs[0].type).toBe("*");
|
||||
});
|
||||
|
||||
it("recomputes types across sequential regions", () => {
|
||||
const h = setup();
|
||||
const secondStart = h.node("DistributedImageListStart", ["values.value_1"],
|
||||
[["image", "IMAGE"], ["mask", "MASK"], ["item_index", "INT"], ["context", "DISTRIBUTED_IMAGE_LIST_CONTEXT"]]);
|
||||
const secondEnd = h.node("DistributedImageListEnd", ["images", "masks", "texts", "context"],
|
||||
[["image_list", "IMAGE"], ["masks", "MASK"], ["texts", "STRING"], ["values", PIPE]]);
|
||||
const secondGet = h.node("DistributedImageListGetValue", ["values"], [["value", "*"]], [{ name: "value_name", value: "value_1" }]);
|
||||
h.connect(h.node("Source", [], [["value", "BOOLEAN"]]), 0, h.start, "values.value_1");
|
||||
h.connect(h.get, 0, secondStart, "values.value_1");
|
||||
h.connect(secondStart, 3, secondEnd, "context");
|
||||
h.connect(secondEnd, 3, secondGet, "values");
|
||||
h.flush();
|
||||
expect(secondGet.outputs[0].type).toBe("BOOLEAN");
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user