From 9739f432a5966ec0b14fc72db9de0c124b22942b Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 8 Mar 2026 22:05:29 +0000 Subject: [PATCH] Refactor: reduce duplication and simplify arg passing across USDU modules - 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. --- api/orchestration/prompt_transform.py | 155 ++++------- nodes/distributed_upscale.py | 10 +- tests/unit/upscale/test_job_store.py | 2 +- tests/unit/upscale/test_result_collector.py | 83 ++---- tests/unit/upscale/test_static.py | 4 - upscale/job_store.py | 30 +- upscale/modes/dynamic.py | 131 +++++---- upscale/modes/static.py | 145 +--------- upscale/result_collector.py | 291 ++++++++++---------- upscale/worker_comms.py | 89 +++--- 10 files changed, 342 insertions(+), 598 deletions(-) diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 92050f4..9e18894 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -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, diff --git a/nodes/distributed_upscale.py b/nodes/distributed_upscale.py index 0319afb..862621f 100644 --- a/nodes/distributed_upscale.py +++ b/nodes/distributed_upscale.py @@ -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: diff --git a/tests/unit/upscale/test_job_store.py b/tests/unit/upscale/test_job_store.py index 91f69a6..d59a052 100644 --- a/tests/unit/upscale/test_job_store.py +++ b/tests/unit/upscale/test_job_store.py @@ -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") diff --git a/tests/unit/upscale/test_result_collector.py b/tests/unit/upscale/test_result_collector.py index f42666e..fa0f228 100644 --- a/tests/unit/upscale/test_result_collector.py +++ b/tests/unit/upscale/test_result_collector.py @@ -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() diff --git a/tests/unit/upscale/test_static.py b/tests/unit/upscale/test_static.py index f1b6e2c..f1e392b 100644 --- a/tests/unit/upscale/test_static.py +++ b/tests/unit/upscale/test_static.py @@ -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") diff --git a/upscale/job_store.py b/upscale/job_store.py index f049116..7051624 100644 --- a/upscale/job_store.py +++ b/upscale/job_store.py @@ -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) diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py index 9f35db5..018c420 100644 --- a/upscale/modes/dynamic.py +++ b/upscale/modes/dynamic.py @@ -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: diff --git a/upscale/modes/static.py b/upscale/modes/static.py index 4c6c69e..ba57265 100644 --- a/upscale/modes/static.py +++ b/upscale/modes/static.py @@ -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, diff --git a/upscale/result_collector.py b/upscale/result_collector.py index f5fdd66..d7868e3 100644 --- a/upscale/result_collector.py +++ b/upscale/result_collector.py @@ -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 diff --git a/upscale/worker_comms.py b/upscale/worker_comms.py index 09f4e31..e9482f9 100644 --- a/upscale/worker_comms.py +++ b/upscale/worker_comms.py @@ -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."""