diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index d07c666..75d2b68 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -228,15 +228,18 @@ def prune_prompt_for_worker(prompt_obj): """Prune worker prompt to distributed nodes and their upstream dependencies.""" collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector") list_collector_ids = find_nodes_by_class(prompt_obj, "DistributedListCollector") + join_ids = find_nodes_by_class(prompt_obj, "DistributedJoin") upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed") branch_ids = find_nodes_by_class(prompt_obj, "DistributedBranch") - distributed_ids = collector_ids + list_collector_ids + upscale_ids + branch_ids + distributed_ids = collector_ids + list_collector_ids + join_ids + upscale_ids + branch_ids if not distributed_ids: return prompt_obj connected = _find_upstream_nodes(prompt_obj, distributed_ids) if branch_ids: connected.update(_find_downstream_nodes(prompt_obj, branch_ids)) + if join_ids: + connected.update(_find_downstream_nodes(prompt_obj, join_ids)) pruned_prompt = {} for node_id in connected: @@ -310,7 +313,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids): "title": "Distributed Empty Image (auto-added)", }, } - collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0] + input_name = "images" + if collector_entry.get("class_type") == "DistributedJoin": + input_name = "input" + collector_entry.setdefault("inputs", {})[input_name] = [placeholder_id, 0] debug_log( f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt." ) @@ -323,7 +329,9 @@ def generate_job_id_map(prompt_index, prefix): job_map = {} distributed_nodes = ( prompt_index.nodes_for_class("DistributedCollector") + + prompt_index.nodes_for_class("DistributedListSplitter") + prompt_index.nodes_for_class("DistributedListCollector") + + prompt_index.nodes_for_class("DistributedJoin") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") + prompt_index.nodes_for_class("DistributedBranch") ) @@ -420,7 +428,16 @@ 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_list_splitter_nodes(prompt_copy, prompt_index, participant_id, enabled_worker_ids, delegate_master): +def _override_list_splitter_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + enabled_worker_ids, + job_id_map, + master_url, + delegate_master, +): """Configure participant chunk mapping for DistributedListSplitter nodes.""" participants = _resolve_participants(enabled_worker_ids, delegate_master) total_participants = max(1, len(participants)) @@ -438,6 +455,14 @@ def _override_list_splitter_nodes(prompt_copy, prompt_index, participant_id, ena inputs = node.setdefault("inputs", {}) inputs["participant_index"] = participant_index inputs["total_participants"] = total_participants + inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + inputs["is_worker"] = not is_master + if is_master: + inputs.pop("master_url", None) + inputs["worker_id"] = "master" + else: + inputs["master_url"] = master_url + inputs["worker_id"] = participant_id def _override_list_collector_nodes( @@ -524,6 +549,74 @@ def _override_branch_nodes( return prompt_copy +def _find_upstream_branch_node(prompt_copy, start_node_id): + """Locate the nearest upstream DistributedBranch node for the provided node.""" + visited = set() + stack = [str(start_node_id)] + while stack: + node_id = stack.pop() + if node_id in visited: + continue + visited.add(node_id) + node = prompt_copy.get(node_id) + if not isinstance(node, dict): + continue + + inputs = node.get("inputs", {}) + for input_value in inputs.values(): + if not (isinstance(input_value, list) and len(input_value) == 2): + continue + source_id = str(input_value[0]) + source_node = prompt_copy.get(source_id) + if not isinstance(source_node, dict): + continue + if source_node.get("class_type") == "DistributedBranch": + return source_id + stack.append(source_id) + return None + + +def _override_join_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + delegate_master, +): + """Configure DistributedJoin nodes for branch convergence.""" + for node_id in prompt_index.nodes_for_class("DistributedJoin"): + node = prompt_copy.get(node_id) + if not isinstance(node, dict): + continue + + upstream_branch_id = _find_upstream_branch_node(prompt_copy, node_id) + assigned_branch = -1 + if upstream_branch_id is not None: + branch_node = prompt_copy.get(upstream_branch_id, {}) + branch_inputs = branch_node.get("inputs", {}) if isinstance(branch_node, dict) else {} + try: + assigned_branch = int(branch_inputs.get("assigned_branch", -1)) + except (TypeError, ValueError): + assigned_branch = -1 + + inputs = node.setdefault("inputs", {}) + inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + inputs["is_worker"] = not is_master + inputs["enabled_worker_ids"] = enabled_json + inputs["assigned_branch"] = assigned_branch + if is_master: + inputs["delegate_only"] = bool(delegate_master) + inputs.pop("master_url", None) + inputs.pop("worker_id", None) + else: + inputs["master_url"] = master_url + inputs["worker_id"] = str(participant_id) + inputs["delegate_only"] = False + + def apply_participant_overrides( prompt_copy, participant_id, @@ -543,8 +636,11 @@ def apply_participant_overrides( _override_list_splitter_nodes( prompt_copy, prompt_index, + is_master, participant_id, enabled_worker_ids, + job_id_map, + master_url, delegate_master, ) _override_collector_nodes( @@ -585,5 +681,15 @@ def apply_participant_overrides( job_id_map, delegate_master, ) + _override_join_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + delegate_master, + ) return prompt_copy diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 3f10c77..1a32c62 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -326,6 +326,7 @@ async def orchestrate_distributed_execution( queue_node_ids = ( prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class("DistributedListCollector") + + prompt_index.nodes_for_class("DistributedJoin") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") ) for node_id in queue_node_ids: @@ -345,8 +346,10 @@ async def orchestrate_distributed_execution( ) if delegate_master: - collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector") + find_nodes_by_class( - master_prompt, "DistributedListCollector" + collector_ids = ( + find_nodes_by_class(master_prompt, "DistributedCollector") + + find_nodes_by_class(master_prompt, "DistributedListCollector") + + find_nodes_by_class(master_prompt, "DistributedJoin") ) upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") if upscale_nodes: @@ -355,7 +358,7 @@ async def orchestrate_distributed_execution( ) elif not collector_ids: debug_log( - "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master." + "Delegate-only master mode requested but no collector/join nodes found in master prompt. Running full prompt on master." ) else: master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids) diff --git a/api/usdu_routes.py b/api/usdu_routes.py index 2735ad4..04ca0d4 100644 --- a/api/usdu_routes.py +++ b/api/usdu_routes.py @@ -7,7 +7,7 @@ from PIL import Image import server from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState -from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized +from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized, init_dynamic_job from ..upscale.payload_parsers import _parse_tiles_from_form from ..utils.logging import debug_log from ..utils.network import handle_api_error @@ -215,6 +215,80 @@ async def request_image_endpoint(request): return await handle_api_error(request, e, 500) +@server.PromptServer.instance.routes.post("/distributed/init_list_queue") +async def init_list_queue_endpoint(request): + """Initialize a shared queue of list indices for dynamic DistributedListSplitter mode.""" + try: + data = await request.json() + multi_job_id = data.get("multi_job_id") + list_size = data.get("list_size", 0) + enabled_workers = data.get("enabled_workers", []) + + if not multi_job_id: + return await handle_api_error(request, "Missing multi_job_id", 400) + + try: + list_size = max(0, int(list_size)) + except (TypeError, ValueError): + return await handle_api_error(request, "list_size must be an integer", 400) + + worker_ids = [str(worker_id) for worker_id in (enabled_workers or []) if str(worker_id).strip()] + await init_dynamic_job( + multi_job_id=str(multi_job_id), + batch_size=list_size, + enabled_workers=worker_ids, + all_indices=list(range(list_size)), + ) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(str(multi_job_id)) + remaining = ( + job_data.pending_images.qsize() + if isinstance(job_data, ImageJobState) + else 0 + ) + + return web.json_response({"status": "success", "remaining": remaining}) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.post("/distributed/request_list_item") +async def request_list_item_endpoint(request): + """Allow workers to pull one list item index for dynamic list distribution.""" + try: + data = await request.json() + worker_id = data.get("worker_id") + multi_job_id = data.get("multi_job_id") + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(str(multi_job_id)) + if not isinstance(job_data, BaseJobState): + return await handle_api_error(request, "Job not found", 404) + + pending_queue = getattr(job_data, "pending_images", None) or getattr(job_data, "pending_tasks", None) + if pending_queue is None: + return await handle_api_error(request, "Job not configured for list item requests", 400) + + try: + item_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + except asyncio.TimeoutError: + return web.json_response({"item_idx": None}) + + worker_id = str(worker_id) + job_data.assigned_to_workers.setdefault(worker_id, []).append(item_idx) + job_data.worker_status[worker_id] = time.time() + remaining = pending_queue.qsize() + return web.json_response({"item_idx": int(item_idx), "estimated_remaining": remaining}) + except Exception as e: + return await handle_api_error(request, e, 500) + + @server.PromptServer.instance.routes.get("/distributed/job_status") async def job_status_endpoint(request): """Endpoint to check if a job is ready.""" diff --git a/distributed.py b/distributed.py index 2f77856..054e35e 100644 --- a/distributed.py +++ b/distributed.py @@ -26,6 +26,7 @@ from .nodes import ( DistributedListSplitter, DistributedListCollector, DistributedBranch, + DistributedJoin, AnyType, ByPassTypeTuple, any_type, diff --git a/nodes/__init__.py b/nodes/__init__.py index 33e10eb..b1b0904 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -13,12 +13,14 @@ from .collector import DistributedCollectorNode from .list_splitter import DistributedListSplitter from .list_collector import DistributedListCollector from .branch import DistributedBranch +from .join import DistributedJoin NODE_CLASS_MAPPINGS = { "DistributedCollector": DistributedCollectorNode, "DistributedListSplitter": DistributedListSplitter, "DistributedListCollector": DistributedListCollector, "DistributedBranch": DistributedBranch, + "DistributedJoin": DistributedJoin, "DistributedSeed": DistributedSeed, "DistributedModelName": DistributedModelName, "DistributedValue": DistributedValue, @@ -31,6 +33,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "DistributedListSplitter": "Distributed List Splitter", "DistributedListCollector": "Distributed List Collector", "DistributedBranch": "Distributed Branch", + "DistributedJoin": "Distributed Join", "DistributedSeed": "Distributed Seed", "DistributedModelName": "Distributed Model Name", "DistributedValue": "Distributed Value", diff --git a/nodes/join.py b/nodes/join.py new file mode 100644 index 0000000..c162bd3 --- /dev/null +++ b/nodes/join.py @@ -0,0 +1,269 @@ +import asyncio +import base64 +import io +import json +import time + +import aiohttp +import torch + +from ..utils.async_helpers import run_async_in_server_loop +from ..utils.config import get_worker_timeout_seconds, is_master_delegate_only +from ..utils.constants import HEARTBEAT_INTERVAL +from ..utils.image import ensure_contiguous, tensor_to_pil +from ..utils.logging import debug_log, log +from ..utils.network import get_client_session +from .utilities import any_type + + +MAX_BRANCH_OUTPUTS = 10 + + +def _get_prompt_server_instance(): + import server as _server + + return _server.PromptServer.instance + + +def _throw_if_processing_interrupted(): + try: + import comfy.model_management as model_management + + model_management.throw_exception_if_processing_interrupted() + except Exception: + return + + +class DistributedJoin: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input": (any_type,), + "num_branches": ("INT", {"default": 2, "min": 2, "max": 10, "step": 1}), + }, + "hidden": { + "multi_job_id": ("STRING", {"default": ""}), + "is_worker": ("BOOLEAN", {"default": False}), + "master_url": ("STRING", {"default": ""}), + "enabled_worker_ids": ("STRING", {"default": "[]"}), + "worker_id": ("STRING", {"default": ""}), + "assigned_branch": ("INT", {"default": -1, "min": -1, "max": 9}), + "delegate_only": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = tuple([any_type] * MAX_BRANCH_OUTPUTS) + RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(MAX_BRANCH_OUTPUTS)]) + FUNCTION = "run" + CATEGORY = "utils" + + def run( + self, + input, + num_branches=2, + multi_job_id="", + is_worker=False, + master_url="", + enabled_worker_ids="[]", + worker_id="", + assigned_branch=-1, + delegate_only=False, + ): + if not multi_job_id: + return self._build_outputs(input, assigned_branch) + + return run_async_in_server_loop( + self.execute( + input, + num_branches=num_branches, + multi_job_id=multi_job_id, + is_worker=is_worker, + master_url=master_url, + enabled_worker_ids=enabled_worker_ids, + worker_id=worker_id, + assigned_branch=assigned_branch, + delegate_only=delegate_only, + ) + ) + + def _build_outputs(self, value, assigned_branch): + outputs = [None] * MAX_BRANCH_OUTPUTS + try: + branch_idx = int(assigned_branch) + except (TypeError, ValueError): + branch_idx = -1 + + if 0 <= branch_idx < MAX_BRANCH_OUTPUTS: + outputs[branch_idx] = value + return tuple(outputs) + + def _parse_enabled_workers(self, enabled_worker_ids): + try: + raw = json.loads(enabled_worker_ids) + except Exception: + raw = [] + + workers = [] + seen = set() + for worker_id in raw if isinstance(raw, list) else []: + worker_id_str = str(worker_id) + if worker_id_str in seen: + continue + seen.add(worker_id_str) + workers.append(worker_id_str) + return workers + + def _expected_worker_ids_for_branches(self, enabled_workers, delegate_mode, num_branches): + participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers) + participant_count = len(participants) + if participant_count <= 0: + return [] + + expected = [] + for participant_index, participant_id in enumerate(participants): + # A participant receives at least one branch iff it appears in the first `num_branches` slots. + if participant_index >= num_branches: + continue + if participant_id == "master": + continue + expected.append(str(participant_id)) + return expected + + async def _send_branch_result_to_master(self, value, branch_idx, multi_job_id, master_url, worker_id): + if not isinstance(value, torch.Tensor): + raise ValueError("DistributedJoin currently supports torch.Tensor results for worker transfer.") + + image_batch = value + if image_batch.ndim == 3: + image_batch = image_batch.unsqueeze(0) + if image_batch.ndim != 4 or image_batch.shape[0] <= 0: + raise ValueError("DistributedJoin tensor result must be IMAGE-like with shape [B,H,W,C] or [H,W,C].") + + img = tensor_to_pil(image_batch, 0) + byte_io = io.BytesIO() + img.save(byte_io, format="PNG", compress_level=0) + encoded_image = base64.b64encode(byte_io.getvalue()).decode("utf-8") + + payload = { + "job_id": str(multi_job_id), + "worker_id": str(worker_id), + "batch_idx": int(branch_idx), + "image": f"data:image/png;base64,{encoded_image}", + "is_last": True, + } + + session = await get_client_session() + url = f"{master_url}/distributed/job_complete" + async with session.post( + url, + json=payload, + timeout=aiohttp.ClientTimeout(total=60), + ) as response: + response.raise_for_status() + + async def execute( + self, + input, + num_branches=2, + multi_job_id="", + is_worker=False, + master_url="", + enabled_worker_ids="[]", + worker_id="", + assigned_branch=-1, + delegate_only=False, + ): + try: + branch_idx = int(assigned_branch) + except (TypeError, ValueError): + branch_idx = -1 + + outputs = [None] * MAX_BRANCH_OUTPUTS + + if is_worker: + if 0 <= branch_idx < MAX_BRANCH_OUTPUTS: + try: + await self._send_branch_result_to_master( + input, + branch_idx, + multi_job_id, + master_url, + worker_id, + ) + outputs[branch_idx] = input + except Exception as exc: + log(f"Worker - DistributedJoin failed to send branch result: {exc}") + return tuple(outputs) + + delegate_mode = bool(delegate_only or is_master_delegate_only()) + try: + branch_count = int(num_branches) + except (TypeError, ValueError): + branch_count = 2 + branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS)) + + enabled_workers = self._parse_enabled_workers(enabled_worker_ids) + expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count) + expected_workers = set(expected_worker_ids) + + if not delegate_mode and 0 <= branch_idx < MAX_BRANCH_OUTPUTS: + outputs[branch_idx] = input + + if not expected_workers: + return tuple(outputs) + + prompt_server = _get_prompt_server_instance() + async with prompt_server.distributed_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_jobs: + prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() + + workers_done = set() + base_timeout = float(get_worker_timeout_seconds()) + slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout) + last_activity = time.time() + + try: + while len(workers_done) < len(expected_workers): + _throw_if_processing_interrupted() + try: + async with prompt_server.distributed_jobs_lock: + queue = prompt_server.distributed_pending_jobs[multi_job_id] + result = await asyncio.wait_for(queue.get(), timeout=slice_timeout) + except asyncio.TimeoutError: + if (time.time() - last_activity) < base_timeout: + continue + missing_workers = sorted(expected_workers - workers_done) + log( + "Master - DistributedJoin heartbeat timeout. " + f"Still waiting for workers: {missing_workers}" + ) + break + + worker_id_value = str(result.get("worker_id", "")) + image_index = result.get("image_index") + tensor = result.get("tensor") + is_last = bool(result.get("is_last", False)) + + try: + slot_idx = int(image_index) + except (TypeError, ValueError): + slot_idx = -1 + + if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and tensor is not None: + if isinstance(tensor, torch.Tensor): + if tensor.is_cuda: + tensor = tensor.cpu() + tensor = ensure_contiguous(tensor) + outputs[slot_idx] = tensor + + last_activity = time.time() + base_timeout = float(get_worker_timeout_seconds()) + if is_last and worker_id_value in expected_workers: + workers_done.add(worker_id_value) + + finally: + async with prompt_server.distributed_jobs_lock: + prompt_server.distributed_pending_jobs.pop(multi_job_id, None) + + return tuple(outputs) diff --git a/nodes/list_splitter.py b/nodes/list_splitter.py index aed4991..e5c15c6 100644 --- a/nodes/list_splitter.py +++ b/nodes/list_splitter.py @@ -1,3 +1,12 @@ +import asyncio + +import aiohttp + +from ..upscale.job_models import BaseJobState, ImageJobState +from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job +from ..utils.async_helpers import run_async_in_server_loop +from ..utils.logging import debug_log +from ..utils.network import get_client_session from .utilities import _chunk_bounds @@ -10,10 +19,15 @@ class DistributedListSplitter: return { "required": { "images": ("IMAGE",), + "mode": (["static", "dynamic"], {"default": "static"}), }, "hidden": { "participant_index": ("INT", {"default": 0, "min": 0, "max": 1024}), "total_participants": ("INT", {"default": 1, "min": 1, "max": 1024}), + "multi_job_id": ("STRING", {"default": ""}), + "is_worker": ("BOOLEAN", {"default": False}), + "master_url": ("STRING", {"default": ""}), + "worker_id": ("STRING", {"default": ""}), }, } @@ -22,12 +36,35 @@ class DistributedListSplitter: FUNCTION = "split" CATEGORY = "image" - def split(self, images, participant_index=0, total_participants=1): + def split( + self, + images, + mode="static", + participant_index=0, + total_participants=1, + multi_job_id="", + is_worker=False, + master_url="", + worker_id="", + ): if images is None: return ([],) image_list = images if isinstance(images, list) else [images] + if mode != "dynamic" or not multi_job_id: + return self._split_static(image_list, participant_index, total_participants) + return run_async_in_server_loop( + self._split_dynamic( + image_list, + multi_job_id, + bool(is_worker), + master_url, + worker_id, + ) + ) + + def _split_static(self, image_list, participant_index, total_participants): split_count = max(1, int(total_participants or 1)) if split_count <= 1: return (list(image_list),) @@ -41,3 +78,111 @@ class DistributedListSplitter: start, end = bounds[index] return (list(image_list[start:end]),) + + async def _ensure_local_dynamic_queue(self, multi_job_id, total_items): + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + existing = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if isinstance(existing, ImageJobState): + return + await init_dynamic_job( + multi_job_id, + batch_size=total_items, + enabled_workers=[], + all_indices=list(range(total_items)), + ) + + async def _pull_items_from_local_queue(self, multi_job_id, worker_id): + prompt_server = ensure_tile_jobs_initialized() + pulled = [] + worker_key = str(worker_id or "master") + + while True: + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not isinstance(job_data, BaseJobState): + break + pending_queue = getattr(job_data, "pending_images", None) or getattr(job_data, "pending_tasks", None) + if pending_queue is None: + break + + try: + item_idx = pending_queue.get_nowait() + except asyncio.QueueEmpty: + break + + job_data.assigned_to_workers.setdefault(worker_key, []).append(item_idx) + pulled.append(int(item_idx)) + + return pulled + + async def _request_item_from_master(self, multi_job_id, master_url, worker_id): + session = await get_client_session() + url = f"{master_url}/distributed/request_list_item" + try: + async with session.post( + url, + json={"multi_job_id": str(multi_job_id), "worker_id": str(worker_id or "")}, + timeout=aiohttp.ClientTimeout(total=10), + ) as response: + if response.status == 404: + return "not_ready", None + if response.status != 200: + return "error", None + payload = await response.json() + return "ok", payload.get("item_idx") + except Exception: + return "error", None + + async def _pull_items_from_master(self, multi_job_id, master_url, worker_id): + pulled = [] + consecutive_not_ready = 0 + max_not_ready_polls = 40 + + while True: + status, item_idx = await self._request_item_from_master(multi_job_id, master_url, worker_id) + + if status == "not_ready": + if consecutive_not_ready >= max_not_ready_polls: + break + consecutive_not_ready += 1 + await asyncio.sleep(0.25) + continue + + if status == "error": + if pulled: + break + await asyncio.sleep(0.25) + continue + + consecutive_not_ready = 0 + if item_idx is None: + break + + try: + pulled.append(int(item_idx)) + except (TypeError, ValueError): + continue + + return pulled + + async def _split_dynamic(self, image_list, multi_job_id, is_worker, master_url, worker_id): + if not image_list: + return ([],) + + if is_worker: + pulled_indices = await self._pull_items_from_master(multi_job_id, master_url, worker_id) + else: + await self._ensure_local_dynamic_queue(multi_job_id, len(image_list)) + pulled_indices = await self._pull_items_from_local_queue(multi_job_id, worker_id) + + selected_images = [ + image_list[index] + for index in pulled_indices + if isinstance(index, int) and 0 <= index < len(image_list) + ] + debug_log( + f"DistributedListSplitter dynamic mode: role={'worker' if is_worker else 'master'} " + f"pulled={len(selected_images)} items for job {multi_job_id}" + ) + return (selected_images,) diff --git a/tests/api/test_usdu_routes.py b/tests/api/test_usdu_routes.py index 4932663..1d2c3aa 100644 --- a/tests/api/test_usdu_routes.py +++ b/tests/api/test_usdu_routes.py @@ -103,6 +103,21 @@ def _load_usdu_routes_module(): job_store_module = types.ModuleType(f"{package_name}.upscale.job_store") job_store_module.MAX_PAYLOAD_SIZE = 1024 job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"] + + async def _init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None): + prompt_server = prompt_server_holder["value"] + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + return + job_data = ImageJobState(multi_job_id=multi_job_id) + job_data.worker_status = {str(worker_id): 0 for worker_id in (enabled_workers or [])} + job_data.assigned_to_workers = {str(worker_id): [] for worker_id in (enabled_workers or [])} + indices = all_indices if all_indices is not None else list(range(int(batch_size or 0))) + for idx in indices: + await job_data.pending_images.put(int(idx)) + prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data + + job_store_module.init_dynamic_job = _init_dynamic_job sys.modules[f"{package_name}.upscale.job_store"] = job_store_module job_models_module = types.ModuleType(f"{package_name}.upscale.job_models") @@ -242,6 +257,34 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase): self.assertTrue(response.payload.get("batched_static")) self.assertEqual(job_data.assigned_to_workers["worker-a"], [4]) + async def test_init_list_queue_initializes_pending_items(self): + request = _FakeRequest( + json_payload={ + "multi_job_id": "list-job-1", + "list_size": 3, + "enabled_workers": ["worker-a", "worker-b"], + } + ) + response = await usdu_routes.init_list_queue_endpoint(request) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("status"), "success") + self.assertEqual(response.payload.get("remaining"), 3) + + async def test_request_list_item_assigns_next_index(self): + prompt_server = usdu_routes._prompt_server_holder["value"] + job_data = usdu_routes._ImageJobState("list-job-2") + await job_data.pending_images.put(11) + prompt_server.distributed_pending_tile_jobs["list-job-2"] = job_data + + request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "list-job-2"}) + response = await usdu_routes.request_list_item_endpoint(request) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("item_idx"), 11) + self.assertEqual(response.payload.get("estimated_remaining"), 0) + self.assertEqual(job_data.assigned_to_workers["worker-a"], [11]) + async def test_submit_tiles_completion_signal_enqueues_last_marker(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._TileJobState("job-4") diff --git a/tests/test_join.py b/tests/test_join.py new file mode 100644 index 0000000..4595cb4 --- /dev/null +++ b/tests/test_join.py @@ -0,0 +1,204 @@ +import asyncio +import importlib.util +import json +import sys +import types +import unittest +from pathlib import Path + +import torch +from PIL import Image + + +def _bootstrap_package(package_name): + 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__ = [] + sys.modules[package_name] = root_pkg + + nodes_pkg = types.ModuleType(f"{package_name}.nodes") + nodes_pkg.__path__ = [] + sys.modules[f"{package_name}.nodes"] = nodes_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.get_worker_timeout_seconds = lambda: 0.2 + config_module.is_master_delegate_only = lambda: False + sys.modules[f"{package_name}.utils.config"] = config_module + + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.HEARTBEAT_INTERVAL = 1.0 + sys.modules[f"{package_name}.utils.constants"] = constants_module + + image_module = types.ModuleType(f"{package_name}.utils.image") + + def ensure_contiguous(tensor): + return tensor.contiguous() + + def tensor_to_pil(image_batch, index): + tensor = image_batch[index].detach().cpu().clamp(0, 1) + array = (tensor.numpy() * 255.0).astype("uint8") + return Image.fromarray(array) + + image_module.ensure_contiguous = ensure_contiguous + image_module.tensor_to_pil = tensor_to_pil + sys.modules[f"{package_name}.utils.image"] = image_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + + async def get_client_session(): + raise RuntimeError("test should patch get_client_session") + + network_module.get_client_session = get_client_session + sys.modules[f"{package_name}.utils.network"] = network_module + + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") + async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro) + sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module + + +def _load_module(package_name, module_rel_path, module_name): + module_path = Path(__file__).resolve().parents[1] / module_rel_path + spec = importlib.util.spec_from_file_location( + f"{package_name}.{module_name}", + module_path, + ) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_join_module(): + package_name = "dist_join_testpkg" + _bootstrap_package(package_name) + _load_module(package_name, "nodes/utilities.py", "nodes.utilities") + return _load_module(package_name, "nodes/join.py", "nodes.join") + + +join_module = _load_join_module() + + +class _FakeResponseCtx: + def __init__(self, recorder, payload, url): + self._recorder = recorder + self._payload = payload + self._url = url + + async def __aenter__(self): + self._recorder.append({"url": self._url, "payload": self._payload}) + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def raise_for_status(self): + return None + + +class _FakeSession: + def __init__(self): + self.calls = [] + + def post(self, url, json=None, timeout=None): + return _FakeResponseCtx(self.calls, json, url) + + +class _FakePromptServer: + def __init__(self): + self.distributed_pending_jobs = {} + self.distributed_jobs_lock = asyncio.Lock() + + +class DistributedJoinTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.node = join_module.DistributedJoin() + + async def test_worker_mode_sends_result_with_assigned_branch_index(self): + fake_session = _FakeSession() + + async def _get_session(): + return fake_session + + join_module.get_client_session = _get_session + + tensor = torch.ones((1, 4, 4, 3), dtype=torch.float32) + outputs = await self.node.execute( + tensor, + num_branches=2, + multi_job_id="join-job", + is_worker=True, + master_url="http://master.local:8188", + worker_id="worker-a", + assigned_branch=1, + ) + + self.assertEqual(len(fake_session.calls), 1) + payload = fake_session.calls[0]["payload"] + self.assertEqual(payload["batch_idx"], 1) + self.assertTrue(payload["is_last"]) + self.assertEqual(payload["worker_id"], "worker-a") + self.assertTrue(torch.equal(outputs[1], tensor)) + + async def test_master_mode_collects_worker_results_into_branch_slots(self): + fake_server = _FakePromptServer() + join_module._get_prompt_server_instance = lambda: fake_server + + queue = asyncio.Queue() + worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32) + worker_2 = torch.full((1, 4, 4, 3), 0.8, dtype=torch.float32) + await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True}) + await queue.put({"tensor": worker_2, "worker_id": "worker-b", "image_index": 2, "is_last": True}) + fake_server.distributed_pending_jobs["join-job"] = queue + + master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32) + outputs = await self.node.execute( + master_tensor, + num_branches=3, + multi_job_id="join-job", + is_worker=False, + enabled_worker_ids=json.dumps(["worker-a", "worker-b"]), + assigned_branch=0, + ) + + self.assertTrue(torch.equal(outputs[0], master_tensor)) + self.assertTrue(torch.equal(outputs[1], worker_1)) + self.assertTrue(torch.equal(outputs[2], worker_2)) + + async def test_unassigned_slots_remain_none(self): + fake_server = _FakePromptServer() + join_module._get_prompt_server_instance = lambda: fake_server + + queue = asyncio.Queue() + worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32) + await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True}) + fake_server.distributed_pending_jobs["join-job"] = queue + + master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32) + outputs = await self.node.execute( + master_tensor, + num_branches=2, + multi_job_id="join-job", + is_worker=False, + enabled_worker_ids=json.dumps(["worker-a"]), + assigned_branch=0, + ) + + self.assertIsNone(outputs[2]) + self.assertIsNone(outputs[9]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_list_splitter.py b/tests/test_list_splitter.py index f38a0c8..0fb69ce 100644 --- a/tests/test_list_splitter.py +++ b/tests/test_list_splitter.py @@ -1,7 +1,9 @@ +import asyncio import importlib.util import sys import types import unittest +from dataclasses import dataclass, field from pathlib import Path import torch @@ -24,11 +26,77 @@ def _bootstrap_package(package_name): utils_pkg.__path__ = [] sys.modules[f"{package_name}.utils"] = utils_pkg + upscale_pkg = types.ModuleType(f"{package_name}.upscale") + upscale_pkg.__path__ = [] + sys.modules[f"{package_name}.upscale"] = upscale_pkg + logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.debug_log = lambda *_args, **_kwargs: None logging_module.log = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.logging"] = logging_module + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") + async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro) + sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + + async def _get_client_session(): + raise RuntimeError("test should patch get_client_session") + + network_module.get_client_session = _get_client_session + sys.modules[f"{package_name}.utils.network"] = network_module + + job_models_module = types.ModuleType(f"{package_name}.upscale.job_models") + + class BaseJobState: + pass + + @dataclass + class ImageJobState(BaseJobState): + multi_job_id: str + mode: str = field(default="dynamic", init=False) + queue: asyncio.Queue = field(default_factory=asyncio.Queue) + pending_images: asyncio.Queue = field(default_factory=asyncio.Queue) + completed_images: dict = field(default_factory=dict) + worker_status: dict = field(default_factory=dict) + assigned_to_workers: dict = field(default_factory=dict) + + @property + def pending_tasks(self): + return self.pending_images + + job_models_module.BaseJobState = BaseJobState + job_models_module.ImageJobState = ImageJobState + sys.modules[f"{package_name}.upscale.job_models"] = job_models_module + + prompt_server_holder = { + "value": types.SimpleNamespace( + distributed_tile_jobs_lock=asyncio.Lock(), + distributed_pending_tile_jobs={}, + ) + } + + job_store_module = types.ModuleType(f"{package_name}.upscale.job_store") + + def ensure_tile_jobs_initialized(): + return prompt_server_holder["value"] + + async def init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None): + prompt_server = prompt_server_holder["value"] + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + return + job_data = ImageJobState(multi_job_id=multi_job_id) + for idx in (all_indices if all_indices is not None else list(range(batch_size))): + await job_data.pending_images.put(idx) + prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data + + job_store_module.ensure_tile_jobs_initialized = ensure_tile_jobs_initialized + job_store_module.init_dynamic_job = init_dynamic_job + job_store_module._prompt_server_holder = prompt_server_holder + sys.modules[f"{package_name}.upscale.job_store"] = job_store_module + def _load_module(package_name, module_rel_path, module_name): module_path = Path(__file__).resolve().parents[1] / module_rel_path @@ -47,29 +115,60 @@ def _load_list_splitter_module(): package_name = "dist_list_splitter_testpkg" _bootstrap_package(package_name) _load_module(package_name, "nodes/utilities.py", "nodes.utilities") - return _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter") + module = _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter") + return module splitter_module = _load_list_splitter_module() +class _FakeResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def json(self): + return self.payload + + +class _FakeSession: + def __init__(self, responses): + self.responses = list(responses) + self.calls = [] + + def post(self, url, json=None, timeout=None): + self.calls.append({"url": url, "json": json}) + if not self.responses: + return _FakeResponse({"item_idx": None}, status=200) + payload = self.responses.pop(0) + return _FakeResponse(payload, status=200) + + class DistributedListSplitterTests(unittest.TestCase): def setUp(self): self.node = splitter_module.DistributedListSplitter() + prompt_server = splitter_module.ensure_tile_jobs_initialized() + prompt_server.distributed_pending_tile_jobs = {} def _images(self, count): return [torch.full((1, 4, 4, 3), float(idx), dtype=torch.float32) for idx in range(count)] def test_single_participant_returns_full_list(self): images = self._images(5) - out = self.node.split(images, participant_index=0, total_participants=1)[0] + out = self.node.split(images, mode="static", participant_index=0, total_participants=1)[0] self.assertEqual(len(out), 5) self.assertEqual([int(t[0, 0, 0, 0].item()) for t in out], [0, 1, 2, 3, 4]) def test_two_participants_split_six_images_evenly(self): images = self._images(6) - p0 = self.node.split(images, participant_index=0, total_participants=2)[0] - p1 = self.node.split(images, participant_index=1, total_participants=2)[0] + p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0] + p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0] self.assertEqual(len(p0), 3) self.assertEqual(len(p1), 3) self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2]) @@ -77,17 +176,14 @@ class DistributedListSplitterTests(unittest.TestCase): def test_three_participants_split_seven_images(self): images = self._images(7) - p0 = self.node.split(images, participant_index=0, total_participants=3)[0] - p1 = self.node.split(images, participant_index=1, total_participants=3)[0] - p2 = self.node.split(images, participant_index=2, total_participants=3)[0] + p0 = self.node.split(images, mode="static", participant_index=0, total_participants=3)[0] + p1 = self.node.split(images, mode="static", participant_index=1, total_participants=3)[0] + p2 = self.node.split(images, mode="static", participant_index=2, total_participants=3)[0] self.assertEqual([len(p0), len(p1), len(p2)], [3, 2, 2]) - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2]) - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p1], [3, 4]) - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p2], [5, 6]) def test_more_participants_than_images_yields_empty_slices(self): images = self._images(2) - outputs = [self.node.split(images, participant_index=i, total_participants=5)[0] for i in range(5)] + outputs = [self.node.split(images, mode="static", participant_index=i, total_participants=5)[0] for i in range(5)] self.assertEqual([len(chunk) for chunk in outputs], [1, 1, 0, 0, 0]) def test_mixed_dimensions_are_preserved(self): @@ -96,13 +192,70 @@ class DistributedListSplitterTests(unittest.TestCase): torch.zeros((1, 8, 6, 3), dtype=torch.float32), torch.zeros((1, 3, 9, 3), dtype=torch.float32), ] - - p0 = self.node.split(images, participant_index=0, total_participants=2)[0] - p1 = self.node.split(images, participant_index=1, total_participants=2)[0] - + p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0] + p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0] self.assertEqual([tuple(t.shape) for t in p0], [(1, 4, 4, 3), (1, 8, 6, 3)]) self.assertEqual([tuple(t.shape) for t in p1], [(1, 3, 9, 3)]) + def test_dynamic_mode_master_pulls_from_local_queue(self): + images = self._images(4) + + pulled = self.node.split( + images, + mode="dynamic", + multi_job_id="job-dynamic", + is_worker=False, + worker_id="master", + )[0] + + self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [0, 1, 2, 3]) + + def test_dynamic_mode_worker_pulls_via_http(self): + images = self._images(5) + fake_session = _FakeSession([ + {"item_idx": 1, "estimated_remaining": 3}, + {"item_idx": 3, "estimated_remaining": 2}, + {"item_idx": None}, + ]) + + async def _get_session(): + return fake_session + + splitter_module.get_client_session = _get_session + + pulled = self.node.split( + images, + mode="dynamic", + multi_job_id="job-dynamic", + is_worker=True, + master_url="http://master.local:8188", + worker_id="worker-a", + )[0] + + self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [1, 3]) + self.assertGreaterEqual(len(fake_session.calls), 3) + + def test_dynamic_mode_exhausted_queue_returns_empty(self): + images = self._images(3) + + first_pull = self.node.split( + images, + mode="dynamic", + multi_job_id="job-dynamic", + is_worker=False, + worker_id="master", + )[0] + second_pull = self.node.split( + images, + mode="dynamic", + multi_job_id="job-dynamic", + is_worker=False, + worker_id="master", + )[0] + + self.assertEqual(len(first_pull), 3) + self.assertEqual(second_pull, []) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index cd6be7b..55c7e52 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -103,6 +103,18 @@ def _branch_prompt(): } +def _join_prompt(): + """1 → 2(DistributedBranch) -> 3(branch0) and 4(branch1) -> 5(DistributedJoin) -> 6(SaveImage).""" + return { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}}, + "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}}, + "4": {"class_type": "Sharpen", "inputs": {"image": ["2", 1]}}, + "5": {"class_type": "DistributedJoin", "inputs": {"input": ["3", 0], "num_branches": 2}}, + "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}}, + } + + def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False): if enabled_worker_ids is None: enabled_worker_ids = ["worker-a", "worker-b"] @@ -290,6 +302,12 @@ class PrunePromptForWorkerTests(unittest.TestCase): self.assertIn("5", result) self.assertIn("6", result) + def test_join_anchor_keeps_join_and_downstream(self): + prompt = _join_prompt() + result = pt.prune_prompt_for_worker(prompt) + self.assertIn("5", result) + self.assertIn("6", result) + # --------------------------------------------------------------------------- # prepare_delegate_master_prompt @@ -375,6 +393,22 @@ class GenerateJobIdMapTests(unittest.TestCase): job_map = pt.generate_job_id_map(idx, "run") self.assertEqual(job_map["8"], "run_8") + def test_maps_list_splitter_nodes(self): + prompt = { + "7": {"class_type": "DistributedListSplitter", "inputs": {}}, + } + idx = pt.PromptIndex(prompt) + job_map = pt.generate_job_id_map(idx, "run") + self.assertEqual(job_map["7"], "run_7") + + def test_maps_join_nodes(self): + prompt = { + "10": {"class_type": "DistributedJoin", "inputs": {}}, + } + idx = pt.PromptIndex(prompt) + job_map = pt.generate_job_id_map(idx, "run") + self.assertEqual(job_map["10"], "run_10") + def test_maps_branch_nodes(self): prompt = { "9": {"class_type": "DistributedBranch", "inputs": {"num_branches": 2}}, @@ -553,11 +587,18 @@ class ApplyOverridesListSplitterTests(unittest.TestCase): result = _apply(self._splitter_prompt(), "master", enabled_worker_ids=["worker-a", "worker-b"]) self.assertEqual(result["1"]["inputs"]["participant_index"], 0) self.assertEqual(result["1"]["inputs"]["total_participants"], 3) + self.assertFalse(result["1"]["inputs"]["is_worker"]) + self.assertNotIn("master_url", result["1"]["inputs"]) + self.assertEqual(result["1"]["inputs"]["worker_id"], "master") + self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1") def test_worker_gets_incremented_participant_index(self): result = _apply(self._splitter_prompt(), "worker-a", enabled_worker_ids=["worker-a", "worker-b"]) self.assertEqual(result["1"]["inputs"]["participant_index"], 1) self.assertEqual(result["1"]["inputs"]["total_participants"], 3) + self.assertTrue(result["1"]["inputs"]["is_worker"]) + self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a") + self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com") def test_delegate_mode_shifts_worker_to_index_zero(self): result = _apply( @@ -648,5 +689,33 @@ class ApplyOverridesBranchTests(unittest.TestCase): self.assertNotIn("5", worker_b_result) self.assertNotIn("6", worker_b_result) + +class ApplyOverridesJoinTests(unittest.TestCase): + def test_join_inherits_assigned_branch_from_upstream_branch_node(self): + master_prompt = { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}}, + "3": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 0], "num_branches": 2}}, + } + worker_prompt = { + "1": {"class_type": "KSampler", "inputs": {}}, + "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}}, + "3": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 1], "num_branches": 2}}, + } + + master_result = _apply(master_prompt, "master", enabled_worker_ids=["worker-a"]) + worker_result = _apply(worker_prompt, "worker-a", enabled_worker_ids=["worker-a"]) + + self.assertEqual(master_result["3"]["inputs"]["assigned_branch"], 0) + self.assertEqual(worker_result["3"]["inputs"]["assigned_branch"], 1) + + def test_join_sets_multi_job_id(self): + prompt = { + "1": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 0]}}, + "2": {"class_type": "KSampler", "inputs": {}}, + } + result = _apply(prompt, "master", enabled_worker_ids=["worker-a"]) + self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1") + if __name__ == "__main__": unittest.main() diff --git a/web/constants.js b/web/constants.js index 896f187..1ec83d7 100644 --- a/web/constants.js +++ b/web/constants.js @@ -227,6 +227,7 @@ export const NODE_CLASSES = { DISTRIBUTED_LIST_SPLITTER: 'DistributedListSplitter', DISTRIBUTED_LIST_COLLECTOR: 'DistributedListCollector', DISTRIBUTED_BRANCH: 'DistributedBranch', + DISTRIBUTED_JOIN: 'DistributedJoin', DISTRIBUTED_SEED: 'DistributedSeed', DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage', UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed', diff --git a/web/executionUtils.js b/web/executionUtils.js index c9438b6..64559df 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -10,9 +10,10 @@ export function setupInterceptor(extension) { const hasListSplitter = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_SPLITTER).length > 0; const hasListCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_COLLECTOR).length > 0; const hasBranch = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_BRANCH).length > 0; + const hasJoin = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_JOIN).length > 0; const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0; - if (hasCollector || hasListSplitter || hasListCollector || hasBranch || hasDistUpscale) { + if (hasCollector || hasListSplitter || hasListCollector || hasBranch || hasJoin || hasDistUpscale) { const result = await executeParallelDistributed(extension, prompt); // Immediate status check for instant feedback checkAllWorkerStatuses(extension);