diff --git a/api/orchestration/serverless_media_prompt.py b/api/orchestration/serverless_media_prompt.py index a86446e..de7ee32 100644 --- a/api/orchestration/serverless_media_prompt.py +++ b/api/orchestration/serverless_media_prompt.py @@ -16,6 +16,7 @@ from graphlib import CycleError, TopologicalSorter from ...serverless_bridge.contract import SOCKET_KINDS, VALUE_SOCKETS, identifier, validate_bindings from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, validate_options from .serverless_prompt import ( + COLLECTOR_OUTPUT_COUNT, _PRIMITIVE_NODES, _REMOVED_SCALAR_NODES, _is_link, @@ -99,7 +100,7 @@ def _validate_prompt(prepared_prompt, *, require_collectors=True): if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES: if slot != 0: raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}") - elif source_class == _COLLECTOR and slot >= len(COLLECTOR_SOCKETS): + elif source_class == _COLLECTOR and slot >= COLLECTOR_OUTPUT_COUNT: raise ValueError(f"invalid output slot {slot} for collector {source_id!r}") parents.add(source_id) dependencies[node_id] = parents diff --git a/api/orchestration/serverless_prompt.py b/api/orchestration/serverless_prompt.py index 1168a56..32be508 100644 --- a/api/orchestration/serverless_prompt.py +++ b/api/orchestration/serverless_prompt.py @@ -29,8 +29,12 @@ _PRIMITIVE_NODES = { "PrimitiveInt": 0, "PrimitiveString": "", } +# DistributedCollector outputs: images, audio, then the per-GPU previews +# (PER_GPU_OUTPUT_NAMES in nodes/collector.py). Previews are valid link +# sources in a workflow even though a pool only ever returns images and audio. +COLLECTOR_OUTPUT_COUNT = 2 + 9 _REMOVED_NODE_OUTPUTS = { - "DistributedCollector": {0, 1}, + "DistributedCollector": set(range(COLLECTOR_OUTPUT_COUNT)), "DistributedSeed": {0}, "DistributedValue": {0}, "DistributedModelName": {0}, diff --git a/nodes/collector.py b/nodes/collector.py index 7aef1e8..89001c7 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -21,6 +21,11 @@ from ..utils import transfer_codec prompt_server = _server.PromptServer.instance +# Per-GPU preview outputs follow the combined images/audio outputs. worker_N is +# the Nth worker in the config, so a slot keeps its GPU when others are toggled. +MAX_WORKER_OUTPUTS = 8 +PER_GPU_OUTPUT_NAMES = ("master",) + tuple(f"worker_{i + 1}" for i in range(MAX_WORKER_OUTPUTS)) + def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes): """Materialize every transfer chunk (runs in an executor).""" @@ -113,7 +118,12 @@ class DistributedCollectorNode: RETURN_NAMES = ("images", "audio") FUNCTION = "run" CATEGORY = "image" - + + def __init__(self): + # Image batch each participant produced in this run, keyed by + # PER_GPU_OUTPUT_NAMES slot. Runtime instances are per-execution. + self.participant_images = {} + @staticmethod def _unwrap_list_input(value): """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST.""" @@ -159,6 +169,9 @@ class DistributedCollectorNode: # 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)) + if len(normalized) == 1: + # torch.cat would copy a single batch for nothing. + return normalized[0] return ensure_contiguous(torch.cat(normalized, dim=0)) def _normalize_images_input(self, images): @@ -232,6 +245,9 @@ class DistributedCollectorNode: if not multi_job_id or pass_through: if pass_through: debug_log("Collector: pass-through mode enabled, returning images unchanged") + elif images is not None: + # Not distributed: everything was produced on this GPU. + self.participant_images = {"master": images} return (images, audio if audio is not None else empty_audio) # Use async helper to run in server loop @@ -655,6 +671,28 @@ class DistributedCollectorNode: worker_audio[worker_id] = self._normalize_audio_input(media["audio"]) return stored + @staticmethod + def _ordered_worker_frames(worker_images, worker_order): + """Yield (worker_id, image_index, tensor) in combined-batch order. + + Expected workers come in ``worker_order``; any unexpected worker ids + follow in sorted order. + """ + seen = set() + for worker_id_str in [str(worker_id) for worker_id in (worker_order or [])]: + seen.add(worker_id_str) + frames = worker_images.get(worker_id_str) + if not frames: + continue + for idx in sorted(frames): + yield worker_id_str, idx, frames[idx] + for worker_id_str in sorted(worker_images): + if worker_id_str in seen: + continue + frames = worker_images[worker_id_str] + for idx in sorted(frames): + yield worker_id_str, idx, frames[idx] + def _reorder_and_combine_tensors( self, worker_images: dict, @@ -672,23 +710,9 @@ class DistributedCollectorNode: 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() - for worker_id_str in ordered_worker_ids: - seen.add(worker_id_str) - if worker_id_str not in worker_images: - 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()): - if worker_id_str in seen: - 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}") + for worker_id_str, idx, tensor in self._ordered_worker_frames(worker_images, worker_order): + ordered_tensors.append(tensor) + sources.append(f"worker {worker_id_str!r} image {idx}") cpu_tensors = [] for t in ordered_tensors: @@ -703,6 +727,79 @@ class DistributedCollectorNode: return self._combine_image_tensors([fallback_images], ["master fallback"]) return None + @staticmethod + def _config_worker_id(participant_id): + """Config entry a participant belongs to. + + Serverless pools run as ``serverless::`` participants + (see queue_orchestration) and all preview on their pool's slot. + """ + if participant_id.startswith("serverless:"): + return participant_id[len("serverless:"):].rsplit(":", 1)[0] + return participant_id + + @staticmethod + def _worker_output_slots(participant_ids): + """Map participant ids to worker_N preview slots (many-to-one). + + Configured workers keep their config position; participants outside + the config take the free slots after the configured workers. + """ + try: + config_ids = [str(worker.get("id")) for worker in load_config().get("workers", [])] + except Exception as e: + debug_log(f"Collector - Could not read worker order for preview outputs: {e}") + config_ids = [] + slots = {} + overflow = {} + for participant_id in participant_ids: + config_id = DistributedCollectorNode._config_worker_id(participant_id) + if config_id in config_ids: + index = config_ids.index(config_id) + else: + index = overflow.setdefault(config_id, len(config_ids) + len(overflow)) + if index < MAX_WORKER_OUTPUTS: + slots[participant_id] = PER_GPU_OUTPUT_NAMES[index + 1] + else: + debug_log(f"Collector - No preview output for {participant_id}; only in the combined images") + return slots + + def _split_participant_images(self, worker_images, worker_order, master_images, combined, master_frames): + """Return each participant's own images keyed by preview slot. + + ``combined`` holds the master's ``master_frames`` frames and then the + worker frames in ``_ordered_worker_frames`` order, so worker previews + are views into it rather than a second copy of every frame. + """ + participants = {} + if master_images is not None: + participants["master"] = master_images + frames = list(self._ordered_worker_frames(worker_images, worker_order)) + slots = self._worker_output_slots(list(dict.fromkeys(worker_id for worker_id, _, _ in frames))) + segments = {} # slot -> [start, end) frame runs within combined + offset = master_frames + for worker_id, _idx, tensor in frames: + count = tensor.shape[0] + # Empty placeholders were dropped from the combined batch too. + if count == 0: + continue + slot = slots.get(worker_id) + if slot is not None: + runs = segments.setdefault(slot, []) + if runs and runs[-1][1] == offset: + runs[-1][1] = offset + count + else: + runs.append([offset, offset + count]) + offset += count + if offset != combined.shape[0]: + # Never attribute frames to the wrong GPU; the combined output is unaffected. + log(f"Master - Per-GPU previews skipped: expected {offset} combined frames, got {combined.shape[0]}") + return participants + for slot, runs in segments.items(): + views = [combined[start:end] for start, end in runs] + participants[slot] = views[0] if len(views) == 1 else torch.cat(views, dim=0) + return participants + 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): if is_worker: # Worker mode: send images and audio to master in a single batch @@ -781,6 +878,8 @@ class DistributedCollectorNode: exact_requested = False num_workers = len(expected_workers) if num_workers == 0: + if images is not None: + self.participant_images = {"master": images} return (images, audio if audio is not None else self.EMPTY_AUDIO) # Create the queue before any expensive local work to avoid job_complete race. @@ -1021,6 +1120,11 @@ class DistributedCollectorNode: else: debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total " f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") + master_frames = 0 if delegate_mode or images_on_cpu is None else master_batch_size + self.participant_images = self._split_participant_images( + worker_images, enabled_workers, images_on_cpu if master_frames else None, + combined, master_frames, + ) return (combined, combined_audio) except Exception as e: if serverless_results or exact_requested: diff --git a/nodes/v3.py b/nodes/v3.py index 05bbd97..a71c2e7 100644 --- a/nodes/v3.py +++ b/nodes/v3.py @@ -5,9 +5,10 @@ This avoids sharing mutable instance state through V3's sanitized class clones. """ import comfy.samplers from comfy_api.v0_0_2 import io +from comfy_execution.graph_utils import ExecutionBlocker from . import utilities as _utilities -from .collector import DistributedCollectorNode as _CollectorRuntime +from .collector import DistributedCollectorNode as _CollectorRuntime, PER_GPU_OUTPUT_NAMES from ..utils import transfer_codec as _transfer_codec from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue @@ -130,7 +131,15 @@ class DistributedCollector(io.ComfyNode): io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY, min=1, max=100, step=1, optional=True, tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)], - outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')], + # Per-GPU outputs are appended so saved links to images/audio keep their slots. + outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')] + [ + io.Image.Output(id=name, display_name=name, + tooltip='Images produced by this GPU only, for previewing. ' + 'worker_N is the Nth worker in the Distributed panel. ' + 'Nodes attached to a GPU that did not take part are skipped. ' + 'After Ultimate SD Upscale Distributed every GPU contributes ' + 'tiles, so only the combined images output is available.') + for name in PER_GPU_OUTPUT_NAMES], is_input_list=True, accept_all_inputs=True, ) @@ -140,12 +149,16 @@ class DistributedCollector(io.ComfyNode): worker_id='', pass_through=False, delegate_only=False, transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY): - return io.NodeOutput(*_CollectorRuntime().run( + runtime = _CollectorRuntime() + result = runtime.run( images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id, is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids, worker_batch_size=worker_batch_size, worker_id=worker_id, pass_through=pass_through, delegate_only=delegate_only, - transfer_format=transfer_format, transfer_quality=transfer_quality)) + transfer_format=transfer_format, transfer_quality=transfer_quality) + # A silent blocker skips previews wired to a GPU that produced nothing. + per_gpu = [runtime.participant_images.get(name, ExecutionBlocker(None)) for name in PER_GPU_OUTPUT_NAMES] + return io.NodeOutput(*result, *per_gpu) class UltimateSDUpscaleDistributed(io.ComfyNode): diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py index a75cc52..e2dc2f9 100644 --- a/tests/helpers/v3_runtime_check.py +++ b/tests/helpers/v3_runtime_check.py @@ -43,6 +43,11 @@ ADDED_OPTIONAL_INPUTS = { }, } +# Outputs appended after the V1 baseline outputs, so saved links keep their slots. +ADDED_OUTPUTS = { + 'DistributedCollector': ['master'] + [f'worker_{index + 1}' for index in range(8)], +} + # 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', 'DistributedImageListStart', 'DistributedImageListEnd', @@ -374,7 +379,8 @@ async def main(): assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id] assert 'DistributedMediaCollector' not in mapping collector_info = mapping['DistributedCollector'].GET_NODE_INFO_V1() - assert collector_info['output'] == ['IMAGE', 'AUDIO'], collector_info + assert collector_info['output'] == ['IMAGE', 'AUDIO'] + ['IMAGE'] * 9, collector_info + assert collector_info['output_name'][:2] == ['images', 'audio'], collector_info assert not any(collector_info['output_is_list']), collector_info assert list(collector_info['input']['optional']) == [ 'images', 'audio', 'transfer_format', 'transfer_quality'] @@ -393,6 +399,11 @@ async def main(): assert old['output'] == ['*'] * 10 old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10 new = json.loads(json.dumps(cls.GET_NODE_INFO_V1())) + added_outputs = ADDED_OUTPUTS.get(node_id, []) + if added_outputs: + assert new['output_name'][len(old['output']):] == added_outputs, (node_id, new['output_name']) + for key in ('output', 'output_name', 'output_is_list'): + new[key] = new[key][:len(old['output'])] for group in ('required', 'optional'): old_inputs = old['input'].get(group, {}) new_inputs = dict(new['input'].get(group, {})) diff --git a/tests/test_collector_image_transport.py b/tests/test_collector_image_transport.py index 6994167..b37967b 100644 --- a/tests/test_collector_image_transport.py +++ b/tests/test_collector_image_transport.py @@ -223,3 +223,128 @@ def test_empty_image_results_raise_instead_of_reaching_save_image(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)) + + +def test_master_collection_records_each_participant_in_its_config_slot(media_modules): + module, _ = media_modules + module.load_config = lambda: {"workers": [{"id": "worker-off"}, {"id": "worker-a"}]} + module.prompt_server.distributed_jobs_lock = asyncio.Lock() + queue = asyncio.Queue() + worker_frames = [torch.full((1, 8, 12, 3), value) for value in (0.25, 0.5)] + # Out-of-order delivery must still preview in image order. + for index in (1, 0): + queue.put_nowait({"worker_id": "worker-a", "image_index": index, + "tensor": worker_frames[index], "is_last": index == 0}) + queue.put_nowait({"worker_id": "serverless:pool", "image_index": 0, + "tensor": torch.full((1, 8, 12, 3), 0.75), "is_last": True}) + module.prompt_server.distributed_pending_jobs = {"image-test": queue} + master = torch.zeros(1, 8, 12, 3) + collector = module.DistributedCollectorNode() + + combined, _ = asyncio.run(collector.execute( + master, None, multi_job_id="image-test", enabled_worker_ids='["worker-a", "serverless:pool"]')) + + previews = collector.participant_images + assert combined.shape[0] == 4 + # worker-a is the second configured worker; the unconfigured serverless + # participant takes the first slot after the configured workers. + assert set(previews) == {"master", "worker_2", "worker_3"} + assert torch.equal(previews["master"], master) + assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.25, 0.5] + assert previews["worker_3"][:, 0, 0, 0].tolist() == [0.75] + + +def test_delegate_master_has_no_master_preview(media_modules): + module, _ = media_modules + module.load_config = lambda: {"workers": [{"id": "worker-a"}]} + module.prompt_server.distributed_jobs_lock = asyncio.Lock() + queue = asyncio.Queue() + worker = torch.full((1, 8, 12, 3), 0.5) + queue.put_nowait({"worker_id": "worker-a", "image_index": 0, "tensor": worker, "is_last": True}) + module.prompt_server.distributed_pending_jobs = {"image-test": queue} + collector = module.DistributedCollectorNode() + + asyncio.run(collector.execute( + torch.empty(0, 64, 64, 3), None, multi_job_id="image-test", + enabled_worker_ids='["worker-a"]', delegate_only=True)) + + assert list(collector.participant_images) == ["worker_1"] + assert torch.equal(collector.participant_images["worker_1"], worker) + + +def test_undistributed_run_previews_everything_as_master(media_modules): + module, _ = media_modules + images = torch.zeros(2, 8, 12, 3) + collector = module.DistributedCollectorNode() + collector.run(images=[images]) + assert list(collector.participant_images) == ["master"] + assert torch.equal(collector.participant_images["master"], images) + + pass_through = module.DistributedCollectorNode() + pass_through.run(images=[images], multi_job_id=["job"], pass_through=[True]) + assert pass_through.participant_images == {} + + +def test_empty_worker_batch_keeps_master_preview_and_combined_result(media_modules): + module, _ = media_modules + module.load_config = lambda: {"workers": [{"id": "worker-a"}]} + module.prompt_server.distributed_jobs_lock = asyncio.Lock() + queue = asyncio.Queue() + queue.put_nowait({"worker_id": "worker-a", "image_index": 0, + "tensor": torch.empty(0, 8, 12, 3), "is_last": True}) + module.prompt_server.distributed_pending_jobs = {"image-test": queue} + master = torch.zeros(1, 8, 12, 3) + collector = module.DistributedCollectorNode() + + combined, _ = asyncio.run(collector.execute( + master, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]')) + + assert combined.shape[0] == 1 + assert list(collector.participant_images) == ["master"] + + +def test_serverless_pool_jobs_share_their_pools_preview_slot(media_modules): + module, _ = media_modules + module.load_config = lambda: {"workers": [{"id": "gpu-a"}, {"id": "pool", "type": "serverless"}]} + module.prompt_server.distributed_jobs_lock = asyncio.Lock() + queue = asyncio.Queue() + for job, value in ((2, 0.75), (1, 0.5)): + queue.put_nowait({"worker_id": f"serverless:pool:{job}", "image_index": 0, + "tensor": torch.full((1, 8, 12, 3), value), "is_last": True}) + module.prompt_server.distributed_pending_jobs = {"image-test": queue} + collector = module.DistributedCollectorNode() + + asyncio.run(collector.execute( + torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test", + enabled_worker_ids='["serverless:pool:1", "serverless:pool:2"]')) + + previews = collector.participant_images + assert set(previews) == {"master", "worker_2"} + # Pool jobs appear in their enabled order, not delivery order. + assert previews["worker_2"][:, 0, 0, 0].tolist() == [0.5, 0.75] + + +def test_serverless_validators_accept_every_collector_output(media_modules): + module, _ = media_modules + legacy = _load_file("serverless_prompt_slot_check", ROOT / "api/orchestration/serverless_prompt.py") + assert legacy.COLLECTOR_OUTPUT_COUNT == 2 + len(module.PER_GPU_OUTPUT_NAMES) + assert legacy._REMOVED_NODE_OUTPUTS["DistributedCollector"] == set(range(legacy.COLLECTOR_OUTPUT_COUNT)) + + +def test_worker_previews_are_views_of_the_combined_batch(media_modules): + module, _ = media_modules + module.load_config = lambda: {"workers": [{"id": "worker-a"}]} + module.prompt_server.distributed_jobs_lock = asyncio.Lock() + queue = asyncio.Queue() + for index in range(2): + queue.put_nowait({"worker_id": "worker-a", "image_index": index, + "tensor": torch.full((1, 8, 12, 3), 0.5), "is_last": index == 1}) + module.prompt_server.distributed_pending_jobs = {"image-test": queue} + collector = module.DistributedCollectorNode() + + combined, _ = asyncio.run(collector.execute( + torch.zeros(1, 8, 12, 3), None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]')) + + preview = collector.participant_images["worker_1"] + assert preview.shape == (2, 8, 12, 3) and preview.is_contiguous() + assert preview.data_ptr() == combined[1:].data_ptr() diff --git a/tests/test_serverless_media_prompt.py b/tests/test_serverless_media_prompt.py index f17be44..b0a0dc8 100644 --- a/tests/test_serverless_media_prompt.py +++ b/tests/test_serverless_media_prompt.py @@ -424,3 +424,21 @@ assert graph[bindings[0]["node_id"]]["class_type"] == "MediaExportV1" ''' result = subprocess.run([sys.executable, "-c", script, str(ROOT)], text=True, capture_output=True) assert result.returncode == 0, result.stderr + + +def test_links_from_collector_preview_outputs_are_pruned_not_rejected(): + # Per-GPU preview sockets follow images/audio; a pool never returns them, + # but a workflow previewing them must still compile for typed pools. + prompt = _prompt() + prompt["preview:1"] = _node("PreviewImage", images=["collector:1", 2]) + prompt["preview:2"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT - 1]) + graph, bindings = _run(prompt) + assert [binding["socket"] for binding in bindings] == ["images"] + assert not any(node["class_type"] == "PreviewImage" for node in graph.values()) + + +def test_rejects_collector_output_slots_beyond_the_schema(): + prompt = _prompt() + prompt["preview:1"] = _node("PreviewImage", images=["collector:1", legacy.COLLECTOR_OUTPUT_COUNT]) + with pytest.raises(ValueError, match="invalid output slot"): + _run(prompt) diff --git a/tests/test_v3_collector_transfer.py b/tests/test_v3_collector_transfer.py index cd66fd8..352893c 100644 --- a/tests/test_v3_collector_transfer.py +++ b/tests/test_v3_collector_transfer.py @@ -13,6 +13,7 @@ import pytest ROOT = Path(__file__).resolve().parents[1] PKG = "dist_v3_transfer_testpkg" +PER_GPU_OUTPUT_NAMES = ("master", "worker_1", "worker_2") class _Field: @@ -51,10 +52,20 @@ def v3(monkeypatch): calls = [] class CollectorRuntime: + participant_images_for_run = {} + + def __init__(self): + self.participant_images = {} + def run(self, **kwargs): calls.append(kwargs) + self.participant_images = dict(self.participant_images_for_run) return (kwargs["images"], kwargs["audio"]) + class ExecutionBlocker: + def __init__(self, message): + self.message = message + def module(name, **attrs): mod = types.ModuleType(name) mod.__dict__.update(attrs) @@ -68,13 +79,16 @@ def v3(monkeypatch): module(f"{PKG}.utils", __path__=[]) module(f"{PKG}.nodes.utilities") module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object) - module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime) + module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime, + PER_GPU_OUTPUT_NAMES=PER_GPU_OUTPUT_NAMES) 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__=[]) module("comfy_api.v0_0_2", io=_fake_io()) + module("comfy_execution", __path__=[]) + module("comfy_execution.graph_utils", ExecutionBlocker=ExecutionBlocker) codec_spec = importlib.util.spec_from_file_location(f"{PKG}.utils.transfer_codec", ROOT / "utils" / "transfer_codec.py") codec = importlib.util.module_from_spec(codec_spec) @@ -85,6 +99,7 @@ def v3(monkeypatch): mod = importlib.util.module_from_spec(spec) monkeypatch.setitem(sys.modules, f"{PKG}.nodes.v3", mod) spec.loader.exec_module(mod) + mod._TestRuntime = CollectorRuntime return mod, codec, calls @@ -127,3 +142,24 @@ def test_registered_collector_defaults_to_legacy_transfer(v3): assert calls[-1]["transfer_format"] == "legacy_png" assert calls[-1]["transfer_quality"] == 90 + + +def test_registered_collector_appends_per_gpu_outputs_after_combined_outputs(v3): + mod, _, _ = v3 + schema = mod.DistributedCollector.define_schema() + + assert [field.kind for field in schema.outputs] == ["Image", "Audio", "Image", "Image", "Image"] + assert [field.kwargs["display_name"] for field in schema.outputs] == [ + "images", "audio", "master", "worker_1", "worker_2"] + + +def test_registered_collector_blocks_per_gpu_outputs_without_images(v3): + mod, _, _ = v3 + mod._TestRuntime.participant_images_for_run = {"master": "m", "worker_2": "w2"} + + _, values, _ = mod.DistributedCollector.execute(images=["img"], audio=["aud"]) + + assert values[:2] == (["img"], ["aud"]) + assert values[2] == "m" + assert type(values[3]).__name__ == "ExecutionBlocker" and values[3].message is None + assert values[4] == "w2" diff --git a/web/collectorOutputs.js b/web/collectorOutputs.js new file mode 100644 index 0000000..81b0a63 --- /dev/null +++ b/web/collectorOutputs.js @@ -0,0 +1,94 @@ +import { app } from "/scripts/app.js"; +import { ENDPOINTS, NODE_CLASSES } from "./constants.js"; + +// worker_N previews the Nth configured worker (see nodes/collector.py). +const WORKER_OUTPUT_PREFIX = "worker_"; +const WORKERS_CHANGED_EVENT = "distributed:workers-changed"; + +const trackedNodes = new Set(); +let workers = []; +let workersPromise = null; +let listenerAttached = false; + +async function fetchWorkers() { + try { + const resp = await fetch(ENDPOINTS.CONFIG); + if (!resp.ok) return []; + const config = await resp.json(); + return Array.isArray(config.workers) ? config.workers : []; + } catch { + return []; + } +} + +// Collectors created together (a workflow load) share one config fetch. +function loadWorkers() { + if (!workersPromise) { + workersPromise = fetchWorkers().then((list) => { + workers = list; + if (list.length === 0) workersPromise = null; + return list; + }); + } + return workersPromise; +} + +// Only the displayed label changes; slot names stay fixed so links survive. +// Unnamed slots drop the label so saved workflows carry no redundant copies. +function labelWorkerOutputs(node, configWorkers) { + for (const output of node.outputs || []) { + if (!output?.name?.startsWith(WORKER_OUTPUT_PREFIX)) continue; + const index = parseInt(output.name.slice(WORKER_OUTPUT_PREFIX.length), 10) - 1; + const worker = configWorkers[index]; + if (worker?.name) { + output.label = `${output.name} (${worker.name})`; + } else { + delete output.label; + } + } + node.setDirtyCanvas?.(true, true); +} + +function relabelAll() { + for (const node of trackedNodes) labelWorkerOutputs(node, workers); +} + +function attachListener() { + if (listenerAttached || typeof window === "undefined") return; + window.addEventListener(WORKERS_CHANGED_EVENT, (event) => { + if (Array.isArray(event?.detail?.workers)) { + workers = event.detail.workers; + relabelAll(); + } + }); + listenerAttached = true; +} + +app.registerExtension({ + name: "Distributed.CollectorOutputs", + async nodeCreated(node) { + if (node.comfyClass !== NODE_CLASSES.DISTRIBUTED_COLLECTOR) return; + try { + attachListener(); + trackedNodes.add(node); + + const originalConfigure = node.configure; + node.configure = function (data) { + const result = originalConfigure ? originalConfigure.call(this, data) : undefined; + labelWorkerOutputs(this, workers); + return result; + }; + + const originalOnRemoved = node.onRemoved; + node.onRemoved = function () { + trackedNodes.delete(this); + return originalOnRemoved ? originalOnRemoved.call(this) : undefined; + }; + + if (workers.length === 0) await loadWorkers(); + labelWorkerOutputs(node, workers); + } catch (error) { + console.error("Error in DistributedCollector outputs extension:", error); + } + }, +}); diff --git a/web/tests/collectorOutputs.test.js b/web/tests/collectorOutputs.test.js new file mode 100644 index 0000000..bdee2fe --- /dev/null +++ b/web/tests/collectorOutputs.test.js @@ -0,0 +1,62 @@ +import { readFileSync } from "node:fs"; +import vm from "node:vm"; +import { describe, expect, it } from "vitest"; + +const source = readFileSync(new URL("../collectorOutputs.js", import.meta.url), "utf8") + .replace(/^import .*;\r?\n/gm, ""); + +function setup(workers, onFetch = () => {}) { + let extension; + const listeners = new Map(); + vm.runInNewContext(source, { + app: { registerExtension: (ext) => { extension = ext; } }, + ENDPOINTS: { CONFIG: "/distributed/config" }, + NODE_CLASSES: { DISTRIBUTED_COLLECTOR: "DistributedCollector" }, + fetch: async () => { onFetch(); return { ok: true, json: async () => ({ workers }) }; }, + Promise, + window: { addEventListener: (name, fn) => listeners.set(name, fn) }, + parseInt, + Array, + console, + }); + const node = { + comfyClass: "DistributedCollector", + outputs: ["images", "audio", "master", "worker_1", "worker_2", "worker_3"].map((name) => ({ name })), + }; + return { extension, listeners, node }; +} + +describe("DistributedCollector per-GPU outputs", () => { + it("labels worker_N with the Nth configured worker and keeps slot names", async () => { + const { extension, node } = setup([{ id: "a", name: "GPU A", enabled: false }, { id: "b", name: "GPU B" }]); + await extension.nodeCreated(node); + expect(node.outputs.map((o) => o.name)).toEqual( + ["images", "audio", "master", "worker_1", "worker_2", "worker_3"]); + expect(node.outputs.slice(3).map((o) => o.label)).toEqual( + ["worker_1 (GPU A)", "worker_2 (GPU B)", undefined]); + expect(node.outputs[2].label).toBeUndefined(); + }); + + it("shares one config fetch across collectors created together", async () => { + let fetches = 0; + const { extension, node } = setup([{ id: "a", name: "GPU A" }], () => { fetches += 1; }); + const second = { ...node, outputs: node.outputs.map((o) => ({ ...o })) }; + await Promise.all([extension.nodeCreated(node), extension.nodeCreated(second)]); + expect(fetches).toBe(1); + expect(second.outputs[3].label).toBe("worker_1 (GPU A)"); + }); + + it("clears a stale saved label when the slot has no configured worker", async () => { + const { extension, node } = setup([]); + node.outputs[3].label = "worker_1 (Old PC)"; + await extension.nodeCreated(node); + expect("label" in node.outputs[3]).toBe(false); + }); + + it("relabels when the worker list changes", async () => { + const { extension, listeners, node } = setup([{ id: "a", name: "GPU A" }]); + await extension.nodeCreated(node); + listeners.get("distributed:workers-changed")({ detail: { workers: [{ id: "c", name: "Renamed" }] } }); + expect(node.outputs[3].label).toBe("worker_1 (Renamed)"); + }); +});