Add DistributedJoin and dynamic list-queue splitter mode
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
+75
-1
@@ -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."""
|
||||
|
||||
@@ -26,6 +26,7 @@ from .nodes import (
|
||||
DistributedListSplitter,
|
||||
DistributedListCollector,
|
||||
DistributedBranch,
|
||||
DistributedJoin,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
|
||||
@@ -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",
|
||||
|
||||
+269
@@ -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)
|
||||
+146
-1
@@ -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,)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
+168
-15
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user