import asyncio, time import comfy.model_management import server from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL from ..utils.logging import debug_log, log from ..utils.config import get_worker_timeout_seconds from .job_store import ensure_tile_jobs_initialized, _mark_task_completed from .job_timeout import _check_and_requeue_timed_out_workers from .job_models import BaseJobState, ImageJobState, TileJobState class ResultCollectorMixin: """ Mixin for master-side result collection in USDU distributed jobs. Expected co-mixins/attributes: - JobStateMixin methods for queue/task access. - `self._check_and_requeue_timed_out_workers(...)` coroutine. - `self._async_yield(...)` optional helper from WorkerCommsMixin. """ def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]: """Log timeout elapsed seconds for each tracked worker and return worker ids.""" if not isinstance(job_data, BaseJobState): return [] worker_status = dict(job_data.worker_status) for worker_id, last_seen in worker_status.items(): elapsed = max(0.0, current_time - float(last_seen)) log( "UltimateSDUpscale Master - Heartbeat timeout: " f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s" ) return list(worker_status.keys()) 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 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") 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 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) 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 tile_idx = tile_data['tile_idx'] key = tile_data.get('global_idx', tile_idx) entry = { 'tile_idx': tile_idx, 'x': tile_data['x'], 'y': tile_data['y'], 'extracted_width': tile_data['extracted_width'], 'extracted_height': tile_data['extracted_height'], 'padding': tile_data['padding'], 'worker_id': worker_id, 'batch_idx': tile_data.get('batch_idx', 0), 'global_idx': tile_data.get('global_idx', tile_idx), } if 'image' in tile_data: entry['image'] = tile_data['image'] elif 'tensor' in tile_data: entry['tensor'] = tile_data['tensor'] collected_results[key] = 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 = self._log_worker_timeout_status(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 method to check and requeue requeued = await self._check_and_requeue_timed_out_workers(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 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 self._async_collect_results(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 _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): """Collect remaining processed images from workers.""" return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', remaining_to_collect=remaining_to_collect, batch_size=batch_size)