Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fef70b7d94 | ||
|
|
c64338e566 | ||
|
|
a399556ee0 | ||
|
|
5159cc7166 | ||
|
|
f149cbdf4d | ||
|
|
efa5dd624c | ||
|
|
0cc1631eec | ||
|
|
1940208035 | ||
|
|
0ec9f5d4bd | ||
|
|
213d491f06 |
+3
-1
@@ -136,7 +136,9 @@ def _decode_canonical_png_tensor(image_payload):
|
||||
|
||||
try:
|
||||
with Image.open(io.BytesIO(png_bytes)) as img:
|
||||
img = img.convert("RGB")
|
||||
# PNG callbacks must retain alpha just like local IMAGE tensors.
|
||||
has_alpha = "A" in img.getbands() or "transparency" in img.info
|
||||
img = img.convert("RGBA" if has_alpha else "RGB")
|
||||
tensor = pil_to_tensor(img)
|
||||
return ensure_contiguous(tensor)
|
||||
except Exception as exc:
|
||||
|
||||
+19
-3
@@ -14,6 +14,20 @@ DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
|
||||
_LIFECYCLE_MARKER = object()
|
||||
|
||||
|
||||
def _without_runtime_fingerprints(prompt: dict) -> dict:
|
||||
"""Compare submitted structure without ComfyUI's per-node cache annotation.
|
||||
|
||||
PromptQueue snapshots the graph before IsChangedCache adds ``is_changed``
|
||||
to executing nodes. Ignore only that top-level field, never inputs or nested
|
||||
metadata, and leave both the queued and executing graphs untouched.
|
||||
"""
|
||||
return {
|
||||
node_id: {key: value for key, value in node.items() if key != "is_changed"}
|
||||
if isinstance(node, dict) else node
|
||||
for node_id, node in prompt.items()
|
||||
}
|
||||
|
||||
|
||||
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.
|
||||
|
||||
@@ -23,8 +37,9 @@ def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt
|
||||
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:
|
||||
execution_prompt = _without_runtime_fingerprints(prompt)
|
||||
prepared, targets = prepare_image_list_execution(execution_prompt)
|
||||
if prepared != execution_prompt or not targets:
|
||||
return False
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
for name in ("get_current_queue_volatile", "get_current_queue"):
|
||||
@@ -35,7 +50,8 @@ def has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt
|
||||
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):
|
||||
and entry[1] == prompt_id and isinstance(entry[2], dict)
|
||||
and _without_runtime_fingerprints(entry[2]) == execution_prompt):
|
||||
roots = entry[4]
|
||||
return (isinstance(roots, (list, tuple))
|
||||
and all(isinstance(root, str) for root in roots)
|
||||
|
||||
+52
-5
@@ -121,6 +121,46 @@ class DistributedCollectorNode:
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _combine_image_tensors(tensors, sources):
|
||||
"""Validate BHWC batches, retaining alpha and deterministic input order."""
|
||||
nonempty = []
|
||||
reference_shape = None
|
||||
reference_source = None
|
||||
for tensor, source in zip(tensors, sources):
|
||||
if not isinstance(tensor, torch.Tensor):
|
||||
raise ValueError(f"Collector expected an IMAGE tensor from {source}")
|
||||
shape = tuple(tensor.shape)
|
||||
if tensor.ndim != 4 or shape[-1] not in (1, 3, 4) or min(shape[1:3]) <= 0:
|
||||
raise ValueError(
|
||||
f"Collector received invalid IMAGE shape {shape} from {source}; "
|
||||
"expected [B, H, W, C] with 1, 3, or 4 channels"
|
||||
)
|
||||
# Empty master placeholders must not determine the worker image grid.
|
||||
if shape[0] == 0:
|
||||
continue
|
||||
if reference_shape is None:
|
||||
reference_shape, reference_source = shape, source
|
||||
elif shape[1:3] != reference_shape[1:3]:
|
||||
raise ValueError(
|
||||
"Collector cannot combine images with different spatial dimensions: "
|
||||
f"{reference_source} has shape {reference_shape}; {source} has shape {shape}"
|
||||
)
|
||||
nonempty.append(tensor)
|
||||
if not nonempty:
|
||||
raise ValueError("Collector cannot return an empty image batch")
|
||||
|
||||
channels = max(tensor.shape[-1] for tensor in nonempty)
|
||||
normalized = []
|
||||
for tensor in nonempty:
|
||||
if tensor.shape[-1] == 1 and channels > 1:
|
||||
tensor = tensor.expand(-1, -1, -1, 3)
|
||||
if tensor.shape[-1] == 3 and channels == 4:
|
||||
# RGB is opaque; never discard alpha from another participant.
|
||||
tensor = torch.cat((tensor, torch.ones_like(tensor[..., :1])), dim=-1)
|
||||
normalized.append(ensure_contiguous(tensor))
|
||||
return ensure_contiguous(torch.cat(normalized, dim=0))
|
||||
|
||||
def _normalize_images_input(self, images):
|
||||
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
|
||||
if isinstance(images, (list, tuple)):
|
||||
@@ -130,7 +170,9 @@ class DistributedCollectorNode:
|
||||
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
|
||||
if len(images) == 1:
|
||||
return ensure_contiguous(images[0])
|
||||
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
|
||||
return self._combine_image_tensors(
|
||||
images, [f"IMAGE list item {index}" for index in range(len(images))]
|
||||
)
|
||||
return ensure_contiguous(images)
|
||||
|
||||
def _normalize_audio_input(self, audio):
|
||||
@@ -624,9 +666,11 @@ class DistributedCollectorNode:
|
||||
):
|
||||
"""Assemble final tensor, or return None when the job contains only audio."""
|
||||
ordered_tensors = []
|
||||
sources = []
|
||||
if not delegate_mode and images_on_cpu is not None:
|
||||
for i in range(master_batch_size):
|
||||
ordered_tensors.append(images_on_cpu[i:i+1])
|
||||
sources.append(f"master image {i}")
|
||||
|
||||
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
|
||||
seen = set()
|
||||
@@ -636,6 +680,7 @@ class DistributedCollectorNode:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
sources.append(f"worker {worker_id_str!r} image {idx}")
|
||||
|
||||
# Append any unexpected worker ids deterministically.
|
||||
for worker_id_str in sorted(worker_images.keys()):
|
||||
@@ -643,6 +688,7 @@ class DistributedCollectorNode:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
sources.append(f"worker {worker_id_str!r} image {idx}")
|
||||
|
||||
cpu_tensors = []
|
||||
for t in ordered_tensors:
|
||||
@@ -652,9 +698,9 @@ class DistributedCollectorNode:
|
||||
cpu_tensors.append(t)
|
||||
|
||||
if cpu_tensors:
|
||||
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
||||
return self._combine_image_tensors(cpu_tensors, sources)
|
||||
if fallback_images is not None:
|
||||
return ensure_contiguous(fallback_images)
|
||||
return self._combine_image_tensors([fallback_images], ["master fallback"])
|
||||
return None
|
||||
|
||||
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
|
||||
@@ -944,8 +990,9 @@ class DistributedCollectorNode:
|
||||
f"{sorted(worker_images)}: {e}"
|
||||
) from e
|
||||
log(f"Master - Error combining images: {e}")
|
||||
# Preserve collected audio for regular workers when image assembly fails.
|
||||
return (images, combined_audio)
|
||||
# Do not silently discard worker images or pass an empty
|
||||
# fallback batch downstream to SaveImage.
|
||||
raise
|
||||
except BaseException:
|
||||
# Remove the queue first, preventing cancelled work from delivering into
|
||||
# a collector that is already unwinding.
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.5.0"
|
||||
version = "1.6.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
"""CPU-only proof against real PromptQueue and IsChangedCache; no services."""
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
core = Path(sys.argv[1]).resolve()
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, str(core))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
import execution
|
||||
import nodes
|
||||
from comfy_execution.graph import DynamicPrompt
|
||||
sys.path.insert(0, str(root))
|
||||
from image_list.lifecycle import has_safe_execution_targets
|
||||
|
||||
|
||||
class Source:
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'seed': ('INT',)}}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, seed):
|
||||
return f'fingerprint-{seed}'
|
||||
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS['ProofSource'] = Source
|
||||
nodes.NODE_CLASS_MAPPINGS['ProofUtility'] = type('Utility', (), {'OUTPUT_NODE': True})
|
||||
prompt = {
|
||||
'source': {'class_type': 'ProofSource', 'inputs': {'seed': 7}},
|
||||
'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
|
||||
'op': {'class_type': 'ProofUtility', 'inputs': {'image': ['start', 0]}},
|
||||
'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['op', 0], 'context': ['start', 3]}},
|
||||
}
|
||||
server = SimpleNamespace(queue_updated=lambda: None)
|
||||
queue = execution.PromptQueue(server)
|
||||
server.prompt_queue = queue
|
||||
queue.put((0, 'owned', prompt, {}, ['end']))
|
||||
item, _ = queue.get(timeout=0)
|
||||
assert item[2] is prompt
|
||||
snapshot = queue.get_current_queue_volatile()[0][0][2]
|
||||
assert snapshot is not prompt
|
||||
assert has_safe_execution_targets(server, 'owned', prompt)
|
||||
cache = execution.IsChangedCache('owned', DynamicPrompt(prompt), None)
|
||||
assert asyncio.run(cache.get('source')) == ['fingerprint-7']
|
||||
assert prompt['source']['is_changed'] == ['fingerprint-7']
|
||||
assert 'is_changed' not in snapshot['source']
|
||||
assert has_safe_execution_targets(server, 'owned', prompt), 'core fingerprint invalidated safe roots'
|
||||
prompt['source']['inputs']['seed'] = 8
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt), 'changed input accepted'
|
||||
prompt['source']['inputs']['seed'] = 7
|
||||
assert not has_safe_execution_targets(server, 'not-owned', prompt)
|
||||
queue.get_current_queue_volatile()[0][0][4].append('op')
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt), 'unsafe output root accepted'
|
||||
print('FINGERPRINT_PROOF_OK')
|
||||
@@ -405,8 +405,9 @@ def test_exact_collector_rejects_old_worker_8bit_result(transport, monkeypatch,
|
||||
def test_exact_collector_shape_mismatch_raises_instead_of_returning_local_images(transport, monkeypatch):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, frames(count=1, channels=4))
|
||||
with pytest.raises(RuntimeError, match="shape|combin|mismatch"):
|
||||
# Channel mixing is supported; incompatible spatial sizes still fail.
|
||||
await send(ctx, frames(count=1, channels=4, width=10))
|
||||
with pytest.raises(RuntimeError, match="spatial dimensions"):
|
||||
await ctx.node.execute(
|
||||
frames(count=1, channels=3), None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
@@ -416,6 +417,33 @@ def test_exact_collector_shape_mismatch_raises_instead_of_returning_local_images
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_channels,worker_channels", [(3, 4), (4, 3)])
|
||||
def test_exact_collector_mixed_rgb_rgba_preserves_every_value(transport, monkeypatch, local_channels, worker_channels):
|
||||
local = frames(count=1, channels=local_channels)
|
||||
worker = frames(count=1, channels=worker_channels)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, worker)
|
||||
result, _audio = await ctx.node.execute(
|
||||
local, None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
)
|
||||
assert result.shape == (2, 7, 9, 4)
|
||||
assert result.is_contiguous()
|
||||
# RGB views into RGBA have a four-channel stride; materialize the
|
||||
# sliced comparison without changing the output-contiguity check.
|
||||
assert_exact(result[:1, ..., :local_channels].contiguous(), local)
|
||||
assert_exact(result[1:, ..., :worker_channels].contiguous(), worker)
|
||||
if local_channels == 3:
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
if worker_channels == 3:
|
||||
assert torch.all(result[1:, ..., 3] == 1)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_incomplete_worker_raises_and_cleans_queue(transport, monkeypatch):
|
||||
"""A timed-out worker's first frame is not a successful partial result."""
|
||||
envelope, parts = exact_chunk(transport, count=1)
|
||||
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Real PNG/tensor regressions; only server/config/HTTP plumbing is stubbed."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import io
|
||||
import base64
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_file(name, path):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def media_modules():
|
||||
# Existing harnesses provide ComfyUI/config doubles. Restore their module
|
||||
# replacements so this file can coexist with the rest of the test suite.
|
||||
with patch.dict(sys.modules):
|
||||
image_utils = _load_file("real_image_utils", ROOT / "utils/image.py")
|
||||
collector_tests = _load_file("collector_harness", ROOT / "tests/test_collector_list_inputs.py")
|
||||
collector = collector_tests._load_collector_module()
|
||||
sys.modules["aiohttp"].web = web
|
||||
route_tests = _load_file("route_harness", ROOT / "tests/api/test_distributed_queue.py")
|
||||
routes = route_tests._load_job_routes_module()
|
||||
for module in (collector, routes):
|
||||
module.pil_to_tensor = image_utils.pil_to_tensor
|
||||
module.ensure_contiguous = image_utils.ensure_contiguous
|
||||
collector.tensor_to_pil = image_utils.tensor_to_pil
|
||||
routes.prompt_server = collector.prompt_server
|
||||
yield collector, routes
|
||||
|
||||
|
||||
def _png_payload(image):
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["RGB", "RGBA", "L", "LA", "palette", "palette_alpha"])
|
||||
def test_png_decoder_preserves_alpha_and_normalizes_opaque_modes(media_modules, mode):
|
||||
_, routes = media_modules
|
||||
if mode.startswith("palette"):
|
||||
image = Image.new("P", (12, 8), 0)
|
||||
image.putpalette([16, 32, 64] + [0] * 765)
|
||||
if mode == "palette_alpha":
|
||||
image.info["transparency"] = 0
|
||||
else:
|
||||
color = {"RGB": (16, 32, 64), "RGBA": (16, 32, 64, 96), "L": 32, "LA": (32, 96)}[mode]
|
||||
image = Image.new(mode, (12, 8), color)
|
||||
alpha = mode in {"RGBA", "LA", "palette_alpha"}
|
||||
expected = np.array(image.convert("RGBA" if alpha else "RGB"), dtype=np.float32) / 255
|
||||
decoded = routes._decode_canonical_png_tensor(_png_payload(image))
|
||||
assert decoded.shape == (1, 8, 12, 4 if alpha else 3)
|
||||
assert decoded.is_contiguous()
|
||||
np.testing.assert_array_equal(decoded[0].numpy(), expected)
|
||||
|
||||
|
||||
async def _master_collect(module, local, worker_images, delegate=False):
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
for index, tensor in enumerate(worker_images):
|
||||
queue.put_nowait({"worker_id": "worker-a", "image_index": index,
|
||||
"tensor": tensor, "is_last": index == len(worker_images) - 1})
|
||||
module.prompt_server.distributed_pending_jobs = {"image-test": queue}
|
||||
return await module.DistributedCollectorNode().execute(
|
||||
local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=delegate,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_channels,worker_channels", [(3, 3), (4, 4), (3, 4), (4, 3), (None, 4)])
|
||||
def test_worker_png_callback_and_collector_preserve_pixels_and_alpha(
|
||||
media_modules, local_channels, worker_channels
|
||||
):
|
||||
module, routes = media_modules
|
||||
local = None if local_channels is None else torch.full((1, 8, 12, local_channels), 0.5)
|
||||
worker = torch.full((2, 8, 12, worker_channels), 0.25)
|
||||
if worker_channels == 4:
|
||||
worker[..., 3] = torch.linspace(0, 1, 12).reshape(1, 1, 12)
|
||||
responses = []
|
||||
|
||||
async def scenario():
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
module.prompt_server.distributed_pending_jobs = {"image-test": asyncio.Queue()}
|
||||
|
||||
class Response:
|
||||
def __init__(self, payload):
|
||||
self.payload = payload
|
||||
|
||||
async def __aenter__(self):
|
||||
async def request_json():
|
||||
return self.payload
|
||||
response = await routes.job_complete_endpoint(types.SimpleNamespace(json=request_json))
|
||||
responses.append(response.status)
|
||||
assert response.status == 200, response.text
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
class Session:
|
||||
def post(self, _url, json, timeout):
|
||||
return Response(json)
|
||||
|
||||
async def get_session():
|
||||
return Session()
|
||||
|
||||
module.get_client_session = get_session
|
||||
module.encode_audio_payload = lambda _audio: None
|
||||
collector = module.DistributedCollectorNode()
|
||||
await collector.send_batch_to_master(worker, None, "image-test", "http://master", "worker-a")
|
||||
result, _ = await collector.execute(
|
||||
local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=local is None,
|
||||
)
|
||||
assert "image-test" not in module.prompt_server.distributed_pending_jobs
|
||||
return result
|
||||
|
||||
result = asyncio.run(scenario())
|
||||
channels = max(local_channels or 3, worker_channels)
|
||||
offset = 0 if local is None else 1
|
||||
assert responses == [200, 200]
|
||||
assert result.shape == (offset + 2, 8, 12, channels)
|
||||
assert result.is_contiguous()
|
||||
if local is not None:
|
||||
assert torch.equal(result[:1, ..., :local_channels], local)
|
||||
if local_channels == 3 and channels == 4:
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
assert torch.allclose(result[offset:, ..., :worker_channels], worker, atol=1 / 255)
|
||||
if worker_channels == 3 and channels == 4:
|
||||
assert torch.all(result[offset:, ..., 3] == 1)
|
||||
|
||||
|
||||
def test_collector_mixed_rgb_rgba_list_promotes_rgb_to_opaque_rgba(media_modules):
|
||||
module, _ = media_modules
|
||||
rgb = torch.zeros(1, 8, 12, 3)
|
||||
rgba = torch.full((1, 8, 12, 4), 0.5)
|
||||
result, _ = module.DistributedCollectorNode().run(images=[rgb, rgba])
|
||||
assert result.shape == (2, 8, 12, 4)
|
||||
assert torch.equal(result[:1, ..., :3], rgb)
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
assert torch.equal(result[1:], rgba)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("other_channels", [1, 3, 4])
|
||||
def test_collector_preserves_grayscale_or_promotes_it_for_color_batches(media_modules, other_channels):
|
||||
module, _ = media_modules
|
||||
gray = torch.full((1, 8, 12, 1), 0.25)
|
||||
other = torch.full((1, 8, 12, other_channels), 0.5)
|
||||
result, _ = module.DistributedCollectorNode().run(images=[gray, other])
|
||||
assert result.shape == (2, 8, 12, other_channels)
|
||||
assert torch.all(result[:1, ..., :min(other_channels, 3)] == 0.25)
|
||||
if other_channels == 4:
|
||||
assert torch.all(result[:1, ..., 3] == 1)
|
||||
assert torch.equal(result[1:], other)
|
||||
|
||||
|
||||
def test_mixed_channel_merge_keeps_configured_worker_and_image_order(media_modules):
|
||||
module, _ = media_modules
|
||||
collector = module.DistributedCollectorNode()
|
||||
master = torch.zeros(1, 8, 12, 3)
|
||||
worker_a = torch.full((1, 8, 12, 4), 0.25)
|
||||
worker_b = torch.full((1, 8, 12, 3), 0.5)
|
||||
unexpected = torch.full((1, 8, 12, 3), 0.75)
|
||||
result = collector._reorder_and_combine_tensors(
|
||||
{"worker-a": {0: worker_a}, "worker-b": {2: worker_b, 0: master},
|
||||
"unexpected": {0: unexpected}},
|
||||
["worker-b", "worker-a"], 1, master, False, master,
|
||||
)
|
||||
assert result.shape == (5, 8, 12, 4)
|
||||
assert result[:, 0, 0, 0].tolist() == [0, 0, 0.5, 0.25, 0.75]
|
||||
assert result[:, 0, 0, 3].tolist() == [1, 1, 1, 0.25, 1]
|
||||
|
||||
|
||||
def test_collector_spatial_mismatch_raises_with_shapes_not_successful_fallback(media_modules):
|
||||
module, _ = media_modules
|
||||
local = torch.zeros(1, 8, 12, 3)
|
||||
worker = torch.zeros(1, 8, 13, 3)
|
||||
with pytest.raises(ValueError, match="spatial dimensions") as error:
|
||||
asyncio.run(_master_collect(module, local, [worker]))
|
||||
message = str(error.value)
|
||||
assert "master image 0" in message
|
||||
assert "worker 'worker-a' image 0" in message
|
||||
assert "(1, 8, 12, 3)" in message and "(1, 8, 13, 3)" in message
|
||||
assert "image-test" not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(1, 8, 12), (1, 8, 12, 5)])
|
||||
def test_collector_rejects_invalid_worker_image_shape(media_modules, shape):
|
||||
module, _ = media_modules
|
||||
with pytest.raises(ValueError, match="worker 'worker-a' image 0"):
|
||||
asyncio.run(_master_collect(module, None, [torch.zeros(shape)], delegate=True))
|
||||
|
||||
|
||||
def test_delegate_master_ignores_empty_placeholder_shape(media_modules):
|
||||
module, _ = media_modules
|
||||
placeholder = torch.empty(0, 64, 64, 3)
|
||||
worker = torch.full((1, 8, 12, 4), 0.5)
|
||||
result, _ = asyncio.run(_master_collect(module, placeholder, [worker], delegate=True))
|
||||
assert torch.equal(result, worker)
|
||||
|
||||
|
||||
def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules):
|
||||
module, _ = media_modules
|
||||
with pytest.raises(ValueError, match="empty image batch"):
|
||||
asyncio.run(_master_collect(module, torch.empty(0, 64, 64, 3),
|
||||
[torch.empty(0, 8, 12, 3)], delegate=True))
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Execution-target proof must survive core cache annotations, not graph edits."""
|
||||
import copy
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from image_list.lifecycle import has_safe_execution_targets
|
||||
from test_image_list_master_runtime import graph
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proof(monkeypatch):
|
||||
prompt = graph()
|
||||
queued = copy.deepcopy(prompt)
|
||||
entry = [0, 'owned', queued, {}, ['end']]
|
||||
server = SimpleNamespace(prompt_queue=SimpleNamespace(
|
||||
get_current_queue_volatile=lambda: ([entry], [])))
|
||||
monkeypatch.setitem(sys.modules, 'nodes', SimpleNamespace(NODE_CLASS_MAPPINGS={
|
||||
'Operation': type('Flagged', (), {'OUTPUT_NODE': True})}))
|
||||
return prompt, queued, entry, server
|
||||
|
||||
|
||||
@pytest.mark.parametrize('value', [['digest'], [None], float('nan'), {'nested': [1, 2]}])
|
||||
@pytest.mark.parametrize('side', ['executing', 'queued', 'both'])
|
||||
def test_top_level_fingerprints_do_not_invalidate_proof(proof, value, side):
|
||||
prompt, queued, _, server = proof
|
||||
if side in ('executing', 'both'):
|
||||
prompt['source']['is_changed'] = value
|
||||
if side in ('queued', 'both'):
|
||||
queued['source']['is_changed'] = ['different queued fingerprint']
|
||||
assert has_safe_execution_targets(server, 'owned', prompt)
|
||||
# Comparisons must not strip or rewrite the caller's data.
|
||||
if side in ('executing', 'both'):
|
||||
assert prompt['source']['is_changed'] is value
|
||||
if side in ('queued', 'both'):
|
||||
assert queued['source']['is_changed'] == ['different queued fingerprint']
|
||||
|
||||
|
||||
@pytest.mark.parametrize('change', ['input', 'link', 'class', 'node_id', 'nested_is_changed', 'other_metadata'])
|
||||
def test_fingerprint_tolerance_does_not_hide_graph_changes(proof, change):
|
||||
prompt, _, _, server = proof
|
||||
prompt['source']['is_changed'] = ['digest']
|
||||
if change == 'input':
|
||||
prompt['source']['inputs']['seed'] = 2
|
||||
elif change == 'link':
|
||||
prompt['op']['inputs']['image'] = ['start', 1]
|
||||
elif change == 'class':
|
||||
prompt['source']['class_type'] = 'OtherSource'
|
||||
elif change == 'node_id':
|
||||
prompt['new_source'] = prompt.pop('source')
|
||||
prompt['start']['inputs']['images'] = ['new_source', 0]
|
||||
elif change == 'nested_is_changed':
|
||||
prompt['source']['inputs']['is_changed'] = True
|
||||
else:
|
||||
prompt['source']['_meta'] = {'title': 'changed'}
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('roots', [['op', 'end'], [], ['end', 'end'], None, [True]])
|
||||
def test_fingerprints_never_grant_unscoped_or_invalid_roots(proof, roots):
|
||||
prompt, _, entry, server = proof
|
||||
prompt['source']['is_changed'] = ['digest']
|
||||
entry[4] = roots
|
||||
assert not has_safe_execution_targets(server, 'owned', prompt)
|
||||
|
||||
|
||||
def test_real_core_queue_snapshot_and_fingerprint():
|
||||
core = Path(os.environ.get('COMFYUI_SOURCE_ROOT', str(Path(sys.executable).parents[2])))
|
||||
if not (core / 'execution.py').exists():
|
||||
pytest.skip('real ComfyUI source not available beside this interpreter')
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(ROOT / 'tests/helpers/image_list_fingerprint_proof.py'), str(core)],
|
||||
text=True, capture_output=True, timeout=90)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
assert 'FINGERPRINT_PROOF_OK' in result.stdout
|
||||
@@ -56,7 +56,7 @@ def test_serverless_pending_waits_past_normal_worker_timeout():
|
||||
queue.put_nowait({
|
||||
"worker_id": participant,
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 2, 2, 2),
|
||||
"tensor": torch.ones(1, 2, 2, 3),
|
||||
"is_last": True,
|
||||
})
|
||||
|
||||
@@ -74,7 +74,7 @@ def test_serverless_pending_waits_past_normal_worker_timeout():
|
||||
|
||||
images, _audio = asyncio.run(run_collection())
|
||||
|
||||
assert tuple(images.shape) == (1, 2, 2, 2)
|
||||
assert tuple(images.shape) == (1, 2, 2, 3)
|
||||
assert manager.pending_calls
|
||||
assert probe_calls == []
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
@@ -162,12 +162,12 @@ def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
|
||||
queue.put_nowait({
|
||||
"worker_id": "serverless:pool:1",
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 3, 2, 2),
|
||||
"tensor": torch.ones(1, 3, 2, 3),
|
||||
"is_last": True,
|
||||
})
|
||||
manager = _FakeManager()
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
local_images = torch.zeros(1, 2, 2, 3)
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
|
||||
asyncio.run(collector.execute(
|
||||
@@ -188,7 +188,7 @@ def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager(
|
||||
collector = _set_up_collector(module, job_id, queue)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
local_images = torch.zeros(1, 2, 2, 3)
|
||||
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
|
||||
|
||||
images, _audio = asyncio.run(collector.execute(
|
||||
|
||||
Reference in New Issue
Block a user