Compare commits

..
Author SHA1 Message Date
Robert Wojciechowski fef70b7d94 Update pyproject.toml 2026-10-14 21:51:03 +11:00
Robert Wojciechowski c64338e566 Update pyproject.toml 2026-10-14 21:50:38 +11:00
Robert Wojciechowski a399556ee0 Merge pull request #92 from robertvoy/fix/issue-91-preserve-image-alpha
fix: preserve alpha and validate distributed image batches
2026-10-14 03:47:36 -07:00
Robert Wojciechowski 5159cc7166 fix: integrate alpha preservation with current collector transports
Merge main while retaining exact transfer, serverless cleanup and image assembly error propagation. Update RGB fixtures and verify mixed RGB/RGBA exact HTTP transport without value loss.
2026-10-14 10:46:16 +00:00
Robert Wojciechowski f149cbdf4d fix: preserve image-list target proof after cache fingerprinting 2026-10-14 04:37:56 +00:00
Robert Wojciechowski efa5dd624c Merge pull request #102 from robertvoy/feat/image-list-serverless
feat: schedule image-list items on bounded serverless slots
2026-10-13 21:18:54 -07:00
Robert Wojciechowski 0cc1631eec Merge pull request #101 from robertvoy/feat/serverless-multimodal
feat: add versioned lossless serverless media transport
2026-10-13 21:18:39 -07:00
Robert Wojciechowski 1940208035 Merge pull request #100 from robertvoy/feat/image-list-execution-targets
feat: preserve image-list execution targets and context values
2026-10-13 21:18:23 -07:00
Robert Wojciechowski 0ec9f5d4bd Merge pull request #104 from robertvoy/fix/image-list-native-cache-reuse
fix: preserve native image-list processing cache reuse
2026-10-13 21:18:09 -07:00
Robert Wojciechowski 213d491f06 fix: preserve image alpha in distributed collection 2026-10-10 05:59:52 +00:00
9 changed files with 476 additions and 17 deletions
+3 -1
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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')
+30 -2
View File
@@ -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)
+225
View File
@@ -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
+5 -5
View File
@@ -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(