nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor: reduce duplication and simplify arg passing across USDU modules

Browse Source 
- Pass UpscaleCoreArgs through static mode instead of 11 individual args
- Extract _post_with_retry helper from duplicate retry loops in worker_comms
- Remove 4 passthrough wrapper pairs in job_store
- Extract _process_image_tiles from duplicate master/worker tile loops
- Merge 4 override functions into 2 in prompt_transform (seed+value, collector+upscale)
- Split _async_collect_results into focused static/dynamic methods

Net reduction of ~256 lines across 10 files.
This commit is contained in:
Robert Wojciechowski
2026-03-08 22:05:29 +00:00
parent 5e724665a1
commit 9739f432a5
10 changed files with 342 additions and 598 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/orchestration/prompt_transform.py
+49 -106
View File
@@ -454,97 +454,47 @@ def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str
return job_map
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json):
"""Configure DistributedSeed nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
inputs["enabled_worker_ids"] = enabled_json
if is_master:
inputs["worker_id"] = ""
else:
inputs["worker_id"] = str(participant_id)
def _override_simple_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json, class_names):
"""Configure simple distributed nodes (DistributedSeed, DistributedValue)."""
for class_name in class_names:
for node_id in prompt_index.nodes_for_class(class_name):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
inputs["enabled_worker_ids"] = enabled_json
inputs["worker_id"] = "" if is_master else str(participant_id)
def _override_collector_nodes(
prompt_copy,
prompt_index,
is_master,
participant_id,
job_id_map,
master_url,
enabled_json,
delegate_master,
def _override_job_nodes(
prompt_copy, prompt_index, is_master, participant_id,
job_id_map, master_url, enabled_json, delegate_master,
class_names, skip_if_upstream=None,
):
"""Configure DistributedCollector nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
"""Configure job-aware distributed nodes (collector, upscale)."""
for class_name in class_names:
for node_id in prompt_index.nodes_for_class(class_name):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
node.setdefault("inputs", {})["pass_through"] = True
continue
if skip_if_upstream and prompt_index.has_upstream(node_id, skip_if_upstream):
node.setdefault("inputs", {})["pass_through"] = True
continue
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
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"] = participant_id
inputs["delegate_only"] = False
def _override_upscale_nodes(
prompt_copy,
prompt_index,
is_master,
participant_id,
job_id_map,
master_url,
enabled_json,
delegate_master=False,
):
"""Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
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
if is_master:
inputs.pop("master_url", None)
inputs.pop("worker_id", None)
inputs["delegate_only"] = bool(delegate_master)
else:
inputs["master_url"] = master_url
inputs["worker_id"] = participant_id
inputs["delegate_only"] = False
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json):
"""Configure DistributedValue nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedValue"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
inputs["enabled_worker_ids"] = enabled_json
if is_master:
inputs["worker_id"] = ""
else:
inputs["worker_id"] = str(participant_id)
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
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"] = participant_id
inputs["delegate_only"] = False
def _override_branch_nodes(
@@ -684,27 +634,20 @@ def apply_participant_overrides(
is_master = participant_id == "master"
enabled_json = json.dumps(enabled_worker_ids)
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json)
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json)
_override_collector_nodes(
prompt_copy,
prompt_index,
is_master,
participant_id,
job_id_map,
master_url,
enabled_json,
delegate_master,
_override_simple_nodes(
prompt_copy, prompt_index, is_master, participant_id, enabled_json,
["DistributedSeed", "DistributedValue"],
)
_override_upscale_nodes(
prompt_copy,
prompt_index,
is_master,
participant_id,
job_id_map,
master_url,
enabled_json,
delegate_master,
_override_job_nodes(
prompt_copy, prompt_index, is_master, participant_id,
job_id_map, master_url, enabled_json, delegate_master,
["DistributedCollector"],
skip_if_upstream="UltimateSDUpscaleDistributed",
)
_override_job_nodes(
prompt_copy, prompt_index, is_master, participant_id,
job_id_map, master_url, enabled_json, delegate_master,
["UltimateSDUpscaleDistributed"],
)
prompt_copy = _override_branch_nodes(
prompt_copy,
nodes/distributed_upscale.py
+4 -6
View File
@@ -276,10 +276,9 @@ class UltimateSDUpscaleDistributed(
)
# Static mode - enhanced with health monitoring and retry logic
return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
return self._process_worker_static_sync(upscaled_image, core_args,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
force_uniform_tiles, multi_job_id, master_url,
worker_id, enabled_workers)
def process_master(
@@ -352,10 +351,9 @@ class UltimateSDUpscaleDistributed(
)
# Static mode - enhanced with unified job management
return self._process_master_static_sync(upscaled_image, core_args.model, core_args.positive, core_args.negative, core_args.vae,
core_args.seed, core_args.steps, core_args.cfg, core_args.sampler_name, core_args.scheduler, core_args.denoise,
return self._process_master_static_sync(upscaled_image, core_args,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, core_args.tiled_decode, multi_job_id, enabled_workers,
force_uniform_tiles, multi_job_id, enabled_workers,
all_tiles, num_tiles_per_image)
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
tests/unit/upscale/test_job_store.py
+1 -1
View File
@@ -150,7 +150,7 @@ class JobStoreTests(unittest.IsolatedAsyncioTestCase):
}
)
drained = await job_store_module._drain_results_queue("job-drain")
drained = await job_store_module.drain_results_queue("job-drain")
self.assertEqual(drained, 3)
self.assertEqual(job_data.completed_tasks[5], "img-5")
tests/unit/upscale/test_result_collector.py
+22 -61
View File
@@ -107,24 +107,6 @@ class _Node(result_collector_module.ResultCollectorMixin):
return 1
class _WrapperNode(result_collector_module.ResultCollectorMixin):
def __init__(self):
self.calls = []
async def _async_collect_results(
self,
multi_job_id,
num_workers,
mode="static",
remaining_to_collect=None,
batch_size=None,
):
self.calls.append(
(multi_job_id, num_workers, mode, remaining_to_collect, batch_size)
)
return {"mode": mode}
class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.prompt_server = job_store_module.server.PromptServer.instance
@@ -156,10 +138,9 @@ class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
}
)
collected = await self.node._async_collect_results(
collected = await self.node._async_collect_worker_tiles(
"job-static",
num_workers=1,
mode="static",
)
self.assertEqual(list(collected.keys()), [9])
@@ -179,12 +160,12 @@ class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
}
)
completed = await self.node._async_collect_results(
completed = await self.node.collect_dynamic_images(
"job-dynamic",
num_workers=2,
mode="dynamic",
remaining_to_collect=1,
num_workers=2,
batch_size=3,
master_processed_count=0,
)
self.assertEqual(completed[2], "img-2")
@@ -198,11 +179,12 @@ class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
old_timeout = result_collector_module.get_worker_timeout_seconds
result_collector_module.get_worker_timeout_seconds = lambda: 0.01
try:
completed = await self.node._async_collect_results(
completed = await self.node.collect_dynamic_images(
"job-timeout",
remaining_to_collect=None,
num_workers=1,
mode="dynamic",
batch_size=4,
master_processed_count=0,
)
finally:
result_collector_module.get_worker_timeout_seconds = old_timeout
@@ -230,58 +212,37 @@ class ResultCollectorMixinTests(unittest.IsolatedAsyncioTestCase):
)
self.assertEqual(worker_ids, ["worker-a"])
async def test_mode_mismatch_and_unsupported_mode_raise(self):
async def test_static_mode_mismatch_raises(self):
self.prompt_server.distributed_pending_tile_jobs["job-mismatch"] = (
job_models_module.ImageJobState("job-mismatch")
)
with self.assertRaises(RuntimeError):
await self.node._async_collect_worker_tiles(
"job-mismatch",
num_workers=1,
)
async def test_dynamic_mode_mismatch_raises(self):
self.prompt_server.distributed_pending_tile_jobs["job-mismatch"] = (
job_models_module.TileJobState("job-mismatch")
)
with self.assertRaises(RuntimeError):
await self.node._async_collect_results(
await self.node.collect_dynamic_images(
"job-mismatch",
remaining_to_collect=1,
num_workers=1,
mode="dynamic",
batch_size=1,
)
self.prompt_server.distributed_pending_tile_jobs["job-unsupported"] = (
job_models_module.TileJobState("job-unsupported")
)
with self.assertRaises(RuntimeError):
await self.node._async_collect_results(
"job-unsupported",
num_workers=1,
mode="unsupported",
master_processed_count=0,
)
async def test_mark_image_completed_updates_job_state(self):
job_data = job_models_module.ImageJobState("job-mark")
self.prompt_server.distributed_pending_tile_jobs["job-mark"] = job_data
await self.node._mark_image_completed("job-mark", 7, "img-7")
await self.node.mark_image_completed("job-mark", 7, "img-7")
self.assertEqual(job_data.completed_images[7], "img-7")
async def test_wrapper_methods_delegate_to_async_collect_results(self):
node = _WrapperNode()
tile_result = await node._async_collect_worker_tiles("job-wrap", 3)
image_result = await node._async_collect_dynamic_images(
"job-wrap",
remaining_to_collect=2,
num_workers=4,
batch_size=8,
master_processed_count=0,
)
self.assertEqual(tile_result["mode"], "static")
self.assertEqual(image_result["mode"], "dynamic")
self.assertEqual(
node.calls,
[
("job-wrap", 3, "static", None, None),
("job-wrap", 4, "dynamic", 2, 8),
],
)
if __name__ == "__main__":
unittest.main()
tests/unit/upscale/test_static.py
-4
View File
@@ -133,10 +133,6 @@ def _load_static_mode_module():
job_store_module.cleanup_job = _noop
job_store_module.drain_results_queue = _noop
job_store_module.get_completed_count = _noop
job_store_module._mark_task_completed = _noop
job_store_module._cleanup_job = _noop
job_store_module._drain_results_queue = _noop
job_store_module._get_completed_count = _noop
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
upscale/job_store.py
+5 -25
View File
@@ -125,7 +125,7 @@ async def init_static_job_batched(
)
async def _drain_results_queue(multi_job_id: str) -> int:
async def drain_results_queue(multi_job_id: str) -> int:
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
@@ -163,7 +163,7 @@ async def _drain_results_queue(multi_job_id: str) -> int:
return collected
async def _get_completed_count(multi_job_id: str) -> int:
async def get_completed_count(multi_job_id: str) -> int:
"""Get count of completed tasks."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
@@ -173,7 +173,7 @@ async def _get_completed_count(multi_job_id: str) -> int:
return 0
async def _mark_task_completed(multi_job_id: str, task_id: int, result: Any) -> None:
async def mark_task_completed(multi_job_id: str, task_id: int, result: Any) -> None:
"""Mark a task as completed."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
@@ -182,7 +182,7 @@ async def _mark_task_completed(multi_job_id: str, task_id: int, result: Any) ->
job_data.completed_tasks[task_id] = result
async def _cleanup_job(multi_job_id: str) -> None:
async def cleanup_job(multi_job_id: str) -> None:
"""Cleanup the job data."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
@@ -195,25 +195,5 @@ async def init_job_queue(
multi_job_id: str,
config: JobQueueInitConfig,
) -> None:
"""Public wrapper for unified job queue initialization."""
"""Public entry point for unified job queue initialization."""
await _init_job_queue(multi_job_id, config)
async def drain_results_queue(multi_job_id: str) -> int:
"""Public wrapper for draining queued worker results."""
return await _drain_results_queue(multi_job_id)
async def get_completed_count(multi_job_id: str) -> int:
"""Public wrapper for counting completed tasks in a job."""
return await _get_completed_count(multi_job_id)
async def mark_task_completed(multi_job_id: str, task_id: int, result: Any) -> None:
"""Public wrapper for marking a task complete."""
await _mark_task_completed(multi_job_id, task_id, result)
async def cleanup_job(multi_job_id: str) -> None:
"""Public wrapper for deleting finished job state."""
await _cleanup_job(multi_job_id)
upscale/modes/dynamic.py
+62 -69
View File
@@ -67,6 +67,53 @@ class DynamicModeMixin:
time.sleep(JOB_POLL_INTERVAL)
return False
def _process_image_tiles(
self,
image_idx: int,
local_image: Image.Image,
single_tensor: torch.Tensor,
core_args: UpscaleCoreArgs,
all_tiles: list,
tile_width: int,
tile_height: int,
padding: int,
mask_blur: int,
width: int,
height: int,
force_uniform_tiles: bool,
context: DynamicModeContext,
per_tile_callback: Any | None = None,
) -> Image.Image:
"""Process all tiles for a single image. Returns the processed PIL image."""
positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
core_args.positive, core_args.negative, image_idx,
)
for tile_idx, pos in enumerate(all_tiles):
source_tensor = pil_to_tensor(local_image)
if single_tensor.is_cuda:
source_tensor = source_tensor.cuda()
local_image = context.tile_ops.process_and_blend_tile(
tile_idx, pos, source_tensor, local_image,
core_args.model,
positive_sliced,
negative_sliced,
core_args.vae,
core_args.seed,
core_args.steps,
core_args.cfg,
core_args.sampler_name,
core_args.scheduler,
core_args.denoise,
tile_width,
tile_height,
padding, mask_blur, width, height, force_uniform_tiles,
core_args.tiled_decode,
batch_idx=image_idx,
)
if per_tile_callback:
per_tile_callback()
return local_image
def _handle_master_dynamic_idle_state(
self,
multi_job_id,
@@ -231,42 +278,14 @@ class DynamicModeMixin:
# Process locally
single_tensor = upscaled_image[image_idx:image_idx+1]
local_image = result_images[image_idx]
image_seed = core_args.seed
# Pre-slice conditioning once per image (not per tile)
positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
core_args.positive,
core_args.negative,
image_idx,
local_image = self._process_image_tiles(
image_idx, local_image, single_tensor, core_args, all_tiles,
tile_width, tile_height, padding, mask_blur, width, height,
force_uniform_tiles, context,
per_tile_callback=lambda: run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1),
)
for tile_idx, pos in enumerate(all_tiles):
source_tensor = pil_to_tensor(local_image)
if single_tensor.is_cuda:
source_tensor = source_tensor.cuda()
local_image = context.tile_ops.process_and_blend_tile(
tile_idx, pos, source_tensor, local_image,
core_args.model,
positive_sliced,
negative_sliced,
core_args.vae,
image_seed,
core_args.steps,
core_args.cfg,
core_args.sampler_name,
core_args.scheduler,
core_args.denoise,
tile_width,
tile_height,
padding, mask_blur, width, height, force_uniform_tiles,
core_args.tiled_decode,
batch_idx=image_idx,
)
# Yield after each tile to minimize worker downtime
run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
# Note: No per-tile drain here – that's what makes this "per-image"
result_images[image_idx] = local_image
# Mark as completed
@@ -404,41 +423,15 @@ class DynamicModeMixin:
local_image = tensor_to_pil(single_tensor, 0).copy()
# Process all tiles for this image
image_seed = core_args.seed
# Pre-slice conditioning once per image (not per tile)
positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
core_args.positive,
core_args.negative,
image_idx,
)
for tile_idx, pos in enumerate(all_tiles):
source_tensor = pil_to_tensor(local_image)
if single_tensor.is_cuda:
source_tensor = source_tensor.cuda()
local_image = context.tile_ops.process_and_blend_tile(
tile_idx, pos, source_tensor, local_image,
core_args.model,
positive_sliced,
negative_sliced,
core_args.vae,
image_seed,
core_args.steps,
core_args.cfg,
core_args.sampler_name,
core_args.scheduler,
core_args.denoise,
tile_width,
tile_height,
padding, mask_blur, width, height, force_uniform_tiles,
core_args.tiled_decode,
batch_idx=image_idx,
)
run_async_in_server_loop(
local_image = self._process_image_tiles(
image_idx, local_image, single_tensor, core_args, all_tiles,
tile_width, tile_height, padding, mask_blur, width, height,
force_uniform_tiles, context,
per_tile_callback=lambda: run_async_in_server_loop(
context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
timeout=5.0
)
timeout=5.0,
),
)
# Send processed image back to master
try:
upscale/modes/static.py
+15 -130
View File
@@ -24,7 +24,6 @@ from ..mode_contexts import (
TileOpsCollaborator,
WorkerCommsCollaborator,
)
from ..processing_args import UpscaleCoreArgs
from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
@@ -79,36 +78,14 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
):
"""Extract one tile position for the whole batch and process it."""
tx, ty = all_tiles[tile_id]
tile_batch_args = TileBatchArgs(
core=UpscaleCoreArgs(
model=model,
positive=positive,
negative=negative,
vae=vae,
seed=seed,
steps=steps,
cfg=cfg,
sampler_name=sampler_name,
scheduler=scheduler,
denoise=denoise,
tiled_decode=tiled_decode,
),
core=core_args,
tile_width=tile_width,
tile_height=tile_height,
padding=padding,
@@ -178,17 +155,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
mode_context: StaticModeContext | None = None,
@@ -206,17 +173,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
)
@@ -243,10 +200,9 @@ class StaticModeMixin:
)
return processed_items
def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
def _process_worker_static_sync(self, upscaled_image, core_args,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
force_uniform_tiles, multi_job_id, master_url,
worker_id, enabled_workers,
mode_context: StaticModeContext | None = None):
"""Worker static mode processing with optional dynamic queue pulling."""
@@ -314,17 +270,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
)
@@ -495,17 +441,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
mode_context: StaticModeContext | None = None,
@@ -540,17 +476,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
mode_context=context,
@@ -573,17 +499,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
mode_context: StaticModeContext | None = None,
@@ -627,17 +543,7 @@ class StaticModeMixin:
tile_height,
padding,
force_uniform_tiles,
model,
positive,
negative,
vae,
seed,
steps,
cfg,
sampler_name,
scheduler,
denoise,
tiled_decode,
core_args,
width,
height,
mode_context=context,
@@ -697,10 +603,9 @@ class StaticModeMixin:
result_tensor = result_tensor.cuda()
return result_tensor
def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
def _process_master_static_sync(self, upscaled_image, core_args,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
force_uniform_tiles, multi_job_id, enabled_workers,
all_tiles, num_tiles_per_image,
mode_context: StaticModeContext | None = None):
"""Static mode master processing with optional dynamic queue pulling."""
@@ -734,17 +639,7 @@ class StaticModeMixin:
tile_height=tile_height,
padding=padding,
force_uniform_tiles=force_uniform_tiles,
model=model,
positive=positive,
negative=negative,
vae=vae,
seed=seed,
steps=steps,
cfg=cfg,
sampler_name=sampler_name,
scheduler=scheduler,
denoise=denoise,
tiled_decode=tiled_decode,
core_args=core_args,
width=width,
height=height,
mode_context=context,
@@ -764,17 +659,7 @@ class StaticModeMixin:
tile_height=tile_height,
padding=padding,
force_uniform_tiles=force_uniform_tiles,
model=model,
positive=positive,
negative=negative,
vae=vae,
seed=seed,
steps=steps,
cfg=cfg,
sampler_name=sampler_name,
scheduler=scheduler,
denoise=denoise,
tiled_decode=tiled_decode,
core_args=core_args,
width=width,
height=height,
mode_context=context,
upscale/result_collector.py
+142 -149
View File
@@ -37,180 +37,92 @@ class ResultCollectorMixin:
"""Default timeout requeue hook; override in host mixins when needed."""
return await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
async def _async_collect_results(self, multi_job_id, num_workers, mode='static',
remaining_to_collect=None, batch_size=None):
"""Unified async helper to collect results from workers (tiles or images)."""
# Get the already initialized queue
def _get_job_data_snapshot(self, prompt_server, multi_job_id):
"""Get a snapshot of job data for timeout logging (non-async helper)."""
current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if isinstance(current_job_data, BaseJobState):
return current_job_data
return None
async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
"""Collect tiles from workers in static mode."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if mode == 'dynamic':
if not isinstance(job_data, ImageJobState):
raise RuntimeError(
f"Mode mismatch: expected dynamic, got {getattr(job_data, 'mode', 'unknown')}"
)
q = job_data.queue
completed_images = job_data.completed_images
expected_count = remaining_to_collect or batch_size
elif mode == 'static':
if not isinstance(job_data, TileJobState):
raise RuntimeError(
f"Mode mismatch: expected static, got {getattr(job_data, 'mode', 'unknown')}"
)
q = job_data.queue
expected_count = len(job_data.completed_tasks) + job_data.pending_tasks.qsize()
else:
raise RuntimeError(f"Unsupported mode: {mode}")
item_type = "images" if mode == 'dynamic' else "tiles"
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
if not isinstance(job_data, TileJobState):
raise RuntimeError(
f"Mode mismatch: expected static, got {getattr(job_data, 'mode', 'unknown')}"
)
q = job_data.queue
expected_count = len(job_data.completed_tasks) + job_data.pending_tasks.qsize()
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} tiles from {num_workers} workers")
collected_results = {}
workers_done = set()
# Unify collector/upscaler wait behavior with the UI worker timeout
timeout = float(get_worker_timeout_seconds())
last_heartbeat_check = time.time()
wait_started_at = time.time()
collected_count = 0
while len(workers_done) < num_workers:
# Check for user interruption
if comfy.model_management.processing_interrupted():
log("Processing interrupted by user")
raise comfy.model_management.InterruptProcessingException()
# For dynamic mode with remaining_to_collect, check if we've collected enough
if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
break
job_data_snapshot = None
async with prompt_server.distributed_tile_jobs_lock:
current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if isinstance(current_job_data, BaseJobState):
job_data_snapshot = current_job_data
job_data_snapshot = self._get_job_data_snapshot(prompt_server, multi_job_id)
try:
# Shorter poll for dynamic mode, but never exceed the configured timeout
wait_timeout = (min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout) if mode == 'dynamic' else timeout)
result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
result = await asyncio.wait_for(q.get(), timeout=timeout)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
if mode == 'static':
# Handle tiles
tiles = result.get('tiles', [])
debug_log(
f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker "
f"'{worker_id}' (is_last={is_last})"
)
for tile_data in tiles:
if 'batch_idx' not in tile_data:
log("UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
continue
tiles = result.get('tiles', [])
debug_log(
f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker "
f"'{worker_id}' (is_last={is_last})"
)
for tile_data in tiles:
if 'batch_idx' not in tile_data:
log("UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
continue
tile_idx = tile_data['tile_idx']
key = tile_data.get('global_idx', tile_idx)
entry = {
**tile_data,
'tile_idx': tile_idx,
'worker_id': worker_id,
'global_idx': key,
}
collected_results[entry['global_idx']] = entry
tile_idx = tile_data['tile_idx']
key = tile_data.get('global_idx', tile_idx)
entry = {
**tile_data,
'tile_idx': tile_idx,
'worker_id': worker_id,
'global_idx': key,
}
_ = entry['tile_idx'], entry['worker_id'], entry['global_idx']
collected_results[entry['global_idx']] = entry
elif mode == 'dynamic':
# Handle full images
if 'image_idx' in result and 'image' in result:
image_idx = result['image_idx']
image_pil = result['image']
completed_images[image_idx] = image_pil
collected_results[image_idx] = image_pil
collected_count += 1
debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
if is_last:
workers_done.add(worker_id)
debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
except asyncio.TimeoutError:
current_time = time.time()
waiting_workers = type(self)._log_worker_timeout_status(
self,
job_data_snapshot,
current_time,
multi_job_id,
self, job_data_snapshot, current_time, multi_job_id,
)
if mode == 'dynamic':
# Check for worker timeouts periodically
if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
# Use the class hook (can be overridden by host mixins/tests).
requeued = await type(self)._check_and_requeue_timed_out_workers(
self,
multi_job_id,
batch_size,
)
if requeued > 0:
log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
last_heartbeat_check = current_time
# Check if we've been waiting too long overall
if current_time - wait_started_at > timeout:
elapsed = current_time - wait_started_at
log(
"UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
)
break
else:
elapsed = current_time - wait_started_at
log(
f"UltimateSDUpscale Master - Heartbeat timeout waiting for {item_type}; "
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
)
break
debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
# Clean up job queue
elapsed = current_time - wait_started_at
log(
f"UltimateSDUpscale Master - Heartbeat timeout waiting for tiles; "
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
)
break
debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} tiles from {len(workers_done)} workers")
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
return collected_results if mode == 'static' else completed_images
async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
"""Async helper to collect tiles from workers."""
return await type(self)._async_collect_results(self, multi_job_id, num_workers, mode='static')
async def _mark_image_completed(self, multi_job_id, image_idx, image_pil):
"""Mark an image as completed in the job data."""
# Mark the image as completed with the image data
await mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if isinstance(job_data, ImageJobState):
job_data.completed_images[image_idx] = image_pil
async def mark_image_completed(self, multi_job_id, image_idx, image_pil):
"""Public image-completion API."""
await self._mark_image_completed(multi_job_id, image_idx, image_pil)
async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
"""Collect remaining processed images from workers."""
_ = master_processed_count
return await type(self)._async_collect_results(
self,
multi_job_id,
num_workers,
mode='dynamic',
remaining_to_collect=remaining_to_collect,
batch_size=batch_size,
)
return collected_results
async def collect_dynamic_images(
self,
@@ -220,11 +132,92 @@ class ResultCollectorMixin:
batch_size,
master_processed_count,
):
"""Public dynamic image-collection API."""
return await self._async_collect_dynamic_images(
multi_job_id,
remaining_to_collect,
num_workers,
batch_size,
master_processed_count,
)
"""Collect remaining processed images from workers in dynamic mode."""
_ = master_processed_count
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if not isinstance(job_data, ImageJobState):
raise RuntimeError(
f"Mode mismatch: expected dynamic, got {getattr(job_data, 'mode', 'unknown')}"
)
q = job_data.queue
completed_images = job_data.completed_images
expected_count = remaining_to_collect or batch_size
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} images from {num_workers} workers")
workers_done = set()
timeout = float(get_worker_timeout_seconds())
last_heartbeat_check = time.time()
wait_started_at = time.time()
collected_count = 0
while len(workers_done) < num_workers:
if comfy.model_management.processing_interrupted():
log("Processing interrupted by user")
raise comfy.model_management.InterruptProcessingException()
if remaining_to_collect and collected_count >= remaining_to_collect:
break
job_data_snapshot = None
async with prompt_server.distributed_tile_jobs_lock:
job_data_snapshot = self._get_job_data_snapshot(prompt_server, multi_job_id)
try:
wait_timeout = min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout)
result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
if 'image_idx' in result and 'image' in result:
image_idx = result['image_idx']
image_pil = result['image']
completed_images[image_idx] = image_pil
collected_count += 1
debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
if is_last:
workers_done.add(worker_id)
debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
except asyncio.TimeoutError:
current_time = time.time()
waiting_workers = type(self)._log_worker_timeout_status(
self, job_data_snapshot, current_time, multi_job_id,
)
if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
requeued = await type(self)._check_and_requeue_timed_out_workers(
self, multi_job_id, batch_size,
)
if requeued > 0:
log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
last_heartbeat_check = current_time
if current_time - wait_started_at > timeout:
elapsed = current_time - wait_started_at
log(
"UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
)
break
debug_log(f"UltimateSDUpscale Master - Collection complete. Got {collected_count} images from {len(workers_done)} workers")
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
return completed_images
async def mark_image_completed(self, multi_job_id, image_idx, image_pil):
"""Mark an image as completed in the job data."""
await mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if isinstance(job_data, ImageJobState):
job_data.completed_images[image_idx] = image_pil
upscale/worker_comms.py
+42 -47
View File
@@ -30,6 +30,37 @@ class WorkerCommsMixin:
def _master_auth_headers() -> dict[str, str]:
return distributed_auth_headers(load_config())
async def _post_with_retry(
self,
url: str,
*,
build_form: callable,
max_retries: int = 5,
initial_delay: float = 0.5,
max_delay: float = 5.0,
error_context: str = "",
) -> None:
"""POST form data with exponential backoff. build_form is called fresh each attempt."""
retry_delay = initial_delay
for attempt in range(max_retries):
try:
session = await get_client_session()
async with session.post(
url,
data=build_form(),
headers=self._master_auth_headers(),
) as response:
response.raise_for_status()
return
except Exception as e:
if attempt < max_retries - 1:
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
await asyncio.sleep(retry_delay)
retry_delay = min(retry_delay * 2, max_delay)
else:
log(f"{error_context}: {e}")
raise
async def send_heartbeat(
self,
multi_job_id: str,
@@ -128,27 +159,11 @@ class WorkerCommsMixin:
)
return data
# Retry logic with exponential backoff
max_retries = 5
retry_delay = 0.5
for attempt in range(max_retries):
try:
session = await get_client_session()
url = f"{master_url}/distributed/submit_tiles"
async with session.post(
url,
data=_build_chunk_form(),
headers=self._master_auth_headers(),
) as response:
response.raise_for_status()
break
except Exception as e:
if attempt < max_retries - 1:
await asyncio.sleep(retry_delay)
retry_delay = min(retry_delay * 2, 5.0)
else:
log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
raise
await self._post_with_retry(
f"{master_url}/distributed/submit_tiles",
build_form=_build_chunk_form,
error_context=f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index}",
)
debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
chunk_index += 1
@@ -293,32 +308,12 @@ class WorkerCommsMixin:
)
return data
# Retry logic
max_retries = 5
retry_delay = 0.5
for attempt in range(max_retries):
try:
session = await get_client_session()
url = f"{master_url}/distributed/submit_image"
async with session.post(
url,
data=_build_image_form(),
headers=self._master_auth_headers(),
) as response:
response.raise_for_status()
debug_log(f"Successfully sent image {image_idx} to master")
return
except Exception as e:
if attempt < max_retries - 1:
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
await asyncio.sleep(retry_delay)
retry_delay *= 2
else:
log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
raise
await self._post_with_retry(
f"{master_url}/distributed/submit_image",
build_form=_build_image_form,
error_context=f"Failed to send image {image_idx}",
)
debug_log(f"Successfully sent image {image_idx} to master")
async def send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
"""Send completion signal to master in dynamic mode."""