diff --git a/upscale/job_models.py b/upscale/job_models.py new file mode 100644 index 0000000..970fff6 --- /dev/null +++ b/upscale/job_models.py @@ -0,0 +1,49 @@ +from dataclasses import dataclass, field +import asyncio +import time + + +class BaseJobState: + """Marker base class for typed USDU job state containers.""" + + +@dataclass +class TileJobState(BaseJobState): + """Typed state container for static (tile) USDU jobs.""" + + multi_job_id: str + mode: str = field(default="static", init=False) + queue: asyncio.Queue = field(default_factory=asyncio.Queue) + pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue) + completed_tasks: dict = field(default_factory=dict) + worker_status: dict = field(default_factory=dict) + assigned_to_workers: dict = field(default_factory=dict) + batch_size: int = 0 + num_tiles_per_image: int = 0 + batched_static: bool = False + created_at: float = field(default_factory=time.monotonic) + + +@dataclass +class ImageJobState(BaseJobState): + """Typed state container for dynamic (per-image) USDU jobs.""" + + 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) + batch_size: int = 0 + num_tiles_per_image: int = 0 + batched_static: bool = False + created_at: float = field(default_factory=time.monotonic) + + @property + def pending_tasks(self): + return self.pending_images + + @property + def completed_tasks(self): + return self.completed_images diff --git a/upscale/job_state.py b/upscale/job_state.py index 87aff71..6cdf09e 100644 --- a/upscale/job_state.py +++ b/upscale/job_state.py @@ -1,28 +1,27 @@ -import asyncio, time -import server +import asyncio + from ..utils.logging import debug_log from ..utils.usdu_managment import ( - ensure_tile_jobs_initialized, _check_and_requeue_timed_out_workers as _requeue_usdu, - JOB_COMPLETED_TASKS, JOB_PENDING_TASKS, + ensure_tile_jobs_initialized, ) +from .job_models import ImageJobState, TileJobState class JobStateMixin: async def _get_job_data(self, multi_job_id): - """Return a shallow snapshot of job data while holding the lock briefly.""" + """Return current job data reference while holding lock briefly.""" 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 not job_data: - return None - return dict(job_data) + return prompt_server.distributed_pending_tile_jobs.get(multi_job_id) async def _get_all_completed_tasks(self, multi_job_id): """Helper to retrieve all completed tasks from the job data.""" job_data = await self._get_job_data(multi_job_id) - if job_data and JOB_COMPLETED_TASKS in job_data: - return dict(job_data[JOB_COMPLETED_TASKS]) + if isinstance(job_data, TileJobState): + return dict(job_data.completed_tasks) + if isinstance(job_data, ImageJobState): + return dict(job_data.completed_images) return {} async def _get_next_image_index(self, multi_job_id): @@ -31,15 +30,14 @@ class JobStateMixin: pending_queue = None async with prompt_server.distributed_tile_jobs_lock: job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if job_data and 'pending_images' in job_data: - pending_queue = job_data['pending_images'] + if isinstance(job_data, ImageJobState): + pending_queue = job_data.pending_images if pending_queue is None: return None try: - image_idx = await asyncio.wait_for(pending_queue.get(), timeout=1.0) - return image_idx + return await asyncio.wait_for(pending_queue.get(), timeout=1.0) except asyncio.TimeoutError: return None @@ -49,15 +47,14 @@ class JobStateMixin: pending_queue = None async with prompt_server.distributed_tile_jobs_lock: job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if job_data and JOB_PENDING_TASKS in job_data: - pending_queue = job_data[JOB_PENDING_TASKS] + if isinstance(job_data, TileJobState): + pending_queue = job_data.pending_tasks if pending_queue is None: return None try: - tile_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) - return tile_idx + return await asyncio.wait_for(pending_queue.get(), timeout=0.1) except asyncio.TimeoutError: return None @@ -66,8 +63,10 @@ class JobStateMixin: 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 job_data and 'completed_images' in job_data: - return len(job_data['completed_images']) + if isinstance(job_data, ImageJobState): + return len(job_data.completed_images) + if isinstance(job_data, TileJobState): + return len(job_data.completed_tasks) return 0 async def _get_all_completed_images(self, multi_job_id): @@ -75,8 +74,8 @@ class JobStateMixin: 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 job_data and 'completed_images' in job_data: - return job_data['completed_images'].copy() + if isinstance(job_data, ImageJobState): + return job_data.completed_images.copy() return {} async def _get_pending_count(self, multi_job_id): @@ -84,18 +83,20 @@ class JobStateMixin: 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 job_data and 'pending_images' in job_data: - return job_data['pending_images'].qsize() + if isinstance(job_data, ImageJobState): + return job_data.pending_images.qsize() + if isinstance(job_data, TileJobState): + return job_data.pending_tasks.qsize() return 0 async def _drain_worker_results_queue(self, multi_job_id): - """Drain pending worker results from queue and update completed_images. Returns count of drained images.""" + """Drain pending worker results from queue and update completed images.""" prompt_server = ensure_tile_jobs_initialized() worker_queue = None async with prompt_server.distributed_tile_jobs_lock: job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if job_data and 'queue' in job_data and 'completed_images' in job_data: - worker_queue = job_data['queue'] + if isinstance(job_data, ImageJobState): + worker_queue = job_data.queue if worker_queue is None: return 0 @@ -113,17 +114,16 @@ class JobStateMixin: collected = 0 async with prompt_server.distributed_tile_jobs_lock: job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if not job_data or 'completed_images' not in job_data: + if not isinstance(job_data, ImageJobState): return 0 - completed_images = job_data['completed_images'] for result in drained_results: - worker_id = result.get('worker_id') - if 'image_idx' in result and 'image' in result: - image_idx = result['image_idx'] - image_pil = result['image'] - if image_idx not in completed_images: - completed_images[image_idx] = image_pil + worker_id = result.get("worker_id") + if "image_idx" in result and "image" in result: + image_idx = result["image_idx"] + image_pil = result["image"] + if image_idx not in job_data.completed_images: + job_data.completed_images[image_idx] = image_pil collected += 1 debug_log(f"Drained image {image_idx} from worker {worker_id}") @@ -133,6 +133,5 @@ class JobStateMixin: return collected async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): - """Check for timed out workers and requeue their assigned images. Returns count of requeued images.""" - # Use the original function from usdu_managment.py + """Check for timed out workers and requeue their assigned images.""" return await _requeue_usdu(multi_job_id, batch_size) diff --git a/upscale/modes/static.py b/upscale/modes/static.py index 0ad7d8d..54960a4 100644 --- a/upscale/modes/static.py +++ b/upscale/modes/static.py @@ -5,13 +5,13 @@ from ...utils.logging import debug_log, log from ...utils.image import tensor_to_pil, pil_to_tensor from ...utils.async_helpers import run_async_in_server_loop from ...utils.config import get_worker_timeout_seconds -from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH +from ...utils.constants import HEARTBEAT_INTERVAL, MAX_BATCH, TILE_SEND_TIMEOUT, TILE_WAIT_TIMEOUT from ...utils.usdu_managment import ( ensure_tile_jobs_initialized, init_static_job_batched, _mark_task_completed, _cleanup_job, _drain_results_queue, _send_heartbeat_to_master, _get_completed_count, - JOB_PENDING_TASKS, JOB_WORKER_STATUS, ) +from ..job_models import TileJobState class StaticModeMixin: @@ -67,12 +67,39 @@ class StaticModeMixin: ) return processed_batch, x1, y1, ew, eh - def _flush_tiles_to_master(self, processed_tiles, multi_job_id, master_url, padding, worker_id): + def _flush_tiles_to_master( + self, + processed_tiles, + multi_job_id, + master_url, + padding, + worker_id, + is_final_flush=False, + ): """Send accumulated tile payloads to master and return a fresh accumulator.""" if not processed_tiles: + if is_final_flush: + run_async_in_server_loop( + self.send_tiles_batch_to_master( + [], + multi_job_id, + master_url, + padding, + worker_id, + is_final_flush=True, + ), + timeout=TILE_SEND_TIMEOUT, + ) return processed_tiles run_async_in_server_loop( - self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + self.send_tiles_batch_to_master( + processed_tiles, + multi_job_id, + master_url, + padding, + worker_id, + is_final_flush=is_final_flush, + ), timeout=TILE_SEND_TIMEOUT ) return [] @@ -236,12 +263,12 @@ class StaticModeMixin: # Send tiles in batches within loop if len(processed_tiles) >= MAX_BATCH: processed_tiles = self._flush_tiles_to_master( - processed_tiles, multi_job_id, master_url, padding, worker_id + processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=False ) # Send any remaining tiles processed_tiles = self._flush_tiles_to_master( - processed_tiles, multi_job_id, master_url, padding, worker_id + processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=True ) debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") @@ -266,7 +293,7 @@ class StaticModeMixin: # Check and requeue timed-out workers periodically current_time = time.time() - if current_time - last_heartbeat_check >= 10.0: + if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL: requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total) if requeued_count > 0: log(f"Requeued {requeued_count} tasks from timed-out workers") @@ -289,9 +316,9 @@ class StaticModeMixin: 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 job_data: - pending_queue = job_data.get(JOB_PENDING_TASKS) - active_workers = list(job_data.get(JOB_WORKER_STATUS, {}).keys()) + if isinstance(job_data, TileJobState): + pending_queue = job_data.pending_tasks + active_workers = list(job_data.worker_status.keys()) if pending_queue and not pending_queue.empty() and len(active_workers) == 0: log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.") break diff --git a/upscale/result_collector.py b/upscale/result_collector.py index b87da93..fde2a5e 100644 --- a/upscale/result_collector.py +++ b/upscale/result_collector.py @@ -1,13 +1,14 @@ import asyncio, time import comfy.model_management import server +from ..utils.constants import HEARTBEAT_INTERVAL from ..utils.logging import debug_log, log from ..utils.config import get_worker_timeout_seconds from ..utils.usdu_managment import ( ensure_tile_jobs_initialized, _mark_task_completed, _check_and_requeue_timed_out_workers, - JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS, ) +from .job_models import BaseJobState, ImageJobState, TileJobState class ResultCollectorMixin: @@ -20,20 +21,23 @@ class ResultCollectorMixin: 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, dict) or 'mode' not in job_data: - raise RuntimeError("Invalid job data structure") - if job_data['mode'] != mode: - raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}") - q = job_data['queue'] - - # For dynamic mode, get reference to completed_images if mode == 'dynamic': - completed_images = job_data['completed_images'] - # Calculate expected count for logging + 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: - # Calculate expected count from job data for static mode - expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set())) + 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") @@ -43,6 +47,7 @@ class ResultCollectorMixin: # 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: @@ -65,42 +70,34 @@ class ResultCollectorMixin: if mode == 'static': # Handle tiles tiles = result.get('tiles', []) - if tiles: - # Batch mode - debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") - - for tile_data in tiles: - # Validate required fields - if 'batch_idx' not in tile_data: - log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") - continue - - tile_idx = tile_data['tile_idx'] - # Use global_idx as key if available (for batch processing) - key = tile_data.get('global_idx', tile_idx) - - # Store the full tile data including metadata; prefer PIL image if present - 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 - else: - # Single tile mode (backward compat) - tile_idx = result['tile_idx'] - collected_results[tile_idx] = result - debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + 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 @@ -117,10 +114,10 @@ class ResultCollectorMixin: debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed") except asyncio.TimeoutError: + current_time = time.time() if mode == 'dynamic': # Check for worker timeouts periodically - current_time = time.time() - if current_time - last_heartbeat_check >= 10.0: + 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: @@ -128,11 +125,39 @@ class ResultCollectorMixin: last_heartbeat_check = current_time # Check if we've been waiting too long overall - if current_time - last_heartbeat_check > timeout: - log(f"UltimateSDUpscale Master - Overall timeout waiting for images") + if current_time - wait_started_at > timeout: + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + worker_status = dict(job_data.worker_status) if isinstance(job_data, BaseJobState) else {} + waiting_workers = list(worker_status.keys()) + for worker_id, last_seen in worker_status.items(): + elapsed_for_worker = max(0.0, current_time - float(last_seen)) + log( + "UltimateSDUpscale Master - Heartbeat timeout: " + f"worker={worker_id}, elapsed={elapsed_for_worker:.1f}s" + ) + 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: - log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}") + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + worker_status = dict(job_data.worker_status) if isinstance(job_data, BaseJobState) else {} + waiting_workers = list(worker_status.keys()) + for worker_id, last_seen in worker_status.items(): + elapsed_for_worker = max(0.0, current_time - float(last_seen)) + log( + "UltimateSDUpscale Master - Heartbeat timeout: " + f"worker={worker_id}, elapsed={elapsed_for_worker:.1f}s" + ) + 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") @@ -155,8 +180,8 @@ class ResultCollectorMixin: 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 job_data and 'completed_images' in job_data: - job_data['completed_images'][image_idx] = image_pil + 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.""" diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py index f52f478..bee033e 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_managment.py @@ -4,14 +4,16 @@ import json import copy import os import io -from aiohttp import web, ClientTimeout +from aiohttp import web import server from PIL import Image # Import from other utilities from .logging import debug_log, log -from .network import handle_api_error, get_client_session +from .network import build_worker_url, handle_api_error, get_client_session, probe_worker +from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState # We avoid converting to tensors on the master for tiles; blending uses PIL +from typing import List, Optional # Configure maximum payload size (50MB default, configurable via environment variable) MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) @@ -98,28 +100,11 @@ def _parse_tiles_from_form(data): return tiles - -# Unified Job Data Structure Keys -JOB_QUEUE = 'queue' -JOB_MODE = 'mode' -JOB_COMPLETED_TASKS = 'completed_tasks' -JOB_WORKER_STATUS = 'worker_status' -JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers' -JOB_PENDING_TASKS = 'pending_tasks' -JOB_BATCH_SIZE = 'batch_size' # For dynamic -JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static - -# Task Types -TASK_TYPE_TILE = 'tile' -TASK_TYPE_IMAGE = 'image' - -from typing import List, Optional - async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None): """Initialize queue for dynamic mode (per-image), with collector fields. - - Creates JOB_PENDING_TASKS with image indices - - Adds 'completed_images' dict and 'pending_images' alias used by collectors + - Creates pending image queue entries + - Initializes dynamic completion/status tracking """ await _init_job_queue( multi_job_id, @@ -128,19 +113,13 @@ async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: all_indices=all_indices or list(range(batch_size)), enabled_workers=enabled_workers, ) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - job_data['completed_images'] = {} - job_data['pending_images'] = job_data[JOB_PENDING_TASKS] debug_log(f"Job {multi_job_id} initialized with {batch_size} images") async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]): """Initialize queue for static mode (batched-per-tile). - - Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1] + - Populates pending tile queue with tile ids [0..num_tiles_per_image-1] """ await _init_job_queue( multi_job_id, @@ -152,7 +131,15 @@ async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_ ) # Initialization handled by master; avoid duplicate init logs here -async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False): +async def _init_job_queue( + multi_job_id, + mode, + batch_size=None, + num_tiles_per_image=None, + all_indices=None, + enabled_workers=None, + batched_static: bool = False, +): """Unified initialization for job queues in static and dynamic modes.""" prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: @@ -160,49 +147,38 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima debug_log(f"Queue already exists for {multi_job_id}") return - job_data = { - JOB_QUEUE: asyncio.Queue(), - JOB_MODE: mode, - JOB_COMPLETED_TASKS: {}, - JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []}, - JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []}, - JOB_PENDING_TASKS: asyncio.Queue(), - } + if mode == 'dynamic': + job_data = ImageJobState(multi_job_id=multi_job_id) + elif mode == 'static': + job_data = TileJobState(multi_job_id=multi_job_id) + else: + raise ValueError(f"Unknown mode: {mode}") + + job_data.worker_status = {w: time.time() for w in enabled_workers or []} + job_data.assigned_to_workers = {w: [] for w in enabled_workers or []} if mode == 'dynamic': - job_data[JOB_BATCH_SIZE] = batch_size - pending_queue = job_data[JOB_PENDING_TASKS] - for i in (all_indices or range(batch_size)): + job_data.batch_size = int(batch_size or 0) + pending_queue = job_data.pending_images + for i in (all_indices or range(int(batch_size or 0))): await pending_queue.put(i) debug_log(f"Initialized image queue with {batch_size} pending items") elif mode == 'static': - job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image - job_data[JOB_BATCH_SIZE] = batch_size - job_data['batched_static'] = bool(batched_static) + job_data.num_tiles_per_image = int(num_tiles_per_image or 0) + job_data.batch_size = int(batch_size or 0) + job_data.batched_static = bool(batched_static) # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1] - pending_queue = job_data[JOB_PENDING_TASKS] + pending_queue = job_data.pending_tasks if batched_static and num_tiles_per_image is not None: for i in range(num_tiles_per_image): await pending_queue.put(i) else: - total_tiles = batch_size * num_tiles_per_image + total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0) for i in range(total_tiles): await pending_queue.put(i) - # Keep backward compatibility - if task assignments provided, still track them - if task_assignments and enabled_workers: - # task_assignments[0] is for master, 1+ are for workers - for i, worker_id in enumerate(enabled_workers): - if i + 1 < len(task_assignments): - job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1] - debug_log(f"Worker {worker_id} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)") - else: - raise ValueError(f"Unknown mode: {mode}") - prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data -# Note: legacy task distribution and queue pull helpers removed - async def _drain_results_queue(multi_job_id): """Drain pending results from queue and update completed_tasks. Returns count drained. @@ -211,10 +187,10 @@ async def _drain_results_queue(multi_job_id): 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 not job_data or JOB_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data: + if not isinstance(job_data, BaseJobState): return 0 - q = job_data[JOB_QUEUE] - completed_tasks = job_data[JOB_COMPLETED_TASKS] + q = job_data.queue + completed_tasks = job_data.completed_tasks collected = 0 while True: @@ -237,26 +213,10 @@ async def _drain_results_queue(multi_job_id): if task_id not in completed_tasks: completed_tasks[task_id] = tile_data collected += 1 - elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat - task_id = result.get('global_idx', result['tile_idx']) - if task_id not in completed_tasks: - completed_tasks[task_id] = { - 'tensor': result['tensor'], - 'tile_idx': result['tile_idx'], - 'x': result['x'], - 'y': result['y'], - 'extracted_width': result['extracted_width'], - 'extracted_height': result['extracted_height'], - 'padding': result['padding'], - 'batch_idx': result.get('batch_idx', 0), - 'global_idx': task_id - } - collected += 1 - if is_last: # Track worker completion - if worker_id in job_data[JOB_WORKER_STATUS]: - del job_data[JOB_WORKER_STATUS][worker_id] + if worker_id in job_data.worker_status: + del job_data.worker_status[worker_id] return collected @@ -269,27 +229,30 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): return 0 current_time = time.time() + if not isinstance(job_data, BaseJobState): + return 0 + requeued_count = 0 - completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {}) + completed_tasks = job_data.completed_tasks # Allow override via config setting 'worker_timeout_seconds' cfg = load_config() hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT)) - for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()): + for worker, last_heartbeat in list(job_data.worker_status.items()): age = current_time - last_heartbeat debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s") if age > hb_timeout: # Busy-only grace policy: require positive signal from worker (/prompt) # We also log assignment state for diagnostics but do not grace on it alone. - assigned = job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []) + assigned = job_data.assigned_to_workers.get(worker, []) incomplete_assigned = 0 try: if assigned: - batched_static = bool(job_data.get('batched_static', False)) + batched_static = bool(job_data.batched_static) if batched_static: - num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1) - batch_size = job_data.get(JOB_BATCH_SIZE, 1) + num_tiles_per_image = int(job_data.num_tiles_per_image or 1) + batch_size = int(job_data.batch_size or 1) for task_id in assigned: for b in range(batch_size): gidx = b * num_tiles_per_image + task_id @@ -305,44 +268,37 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): debug_log(f"Assigned diagnostics failed for worker {worker}: {e}") busy = False - probe_status = None probe_queue = None try: cfg_workers = load_config().get('workers', []) wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None) if wrec: - host = wrec.get('host') or 'localhost' - port = int(wrec.get('port', 8188)) - url = f"http://{host}:{port}/prompt" - debug_log(f"Probing worker {worker} at {url}") - session = await get_client_session() - async with session.get(url, timeout=ClientTimeout(total=2.0)) as resp: - probe_status = resp.status - if resp.status == 200: - try: - payload = await resp.json() - probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0)) - except Exception: - probe_queue = 0 - busy = probe_queue is not None and probe_queue > 0 + worker_url = build_worker_url(wrec) + debug_log(f"Probing worker {worker} at {worker_url}/prompt") + payload = await probe_worker(worker_url, timeout=2.0) + if payload is not None: + probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0)) + busy = probe_queue is not None and probe_queue > 0 except Exception as e: debug_log(f"Probe failed for worker {worker}: {e}") finally: - debug_log(f"Probe diagnostics: http_status={probe_status} queue_remaining={probe_queue}") + debug_log( + f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}" + ) if busy: - job_data[JOB_WORKER_STATUS][worker] = current_time + job_data.worker_status[worker] = current_time debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue") continue - log(f"Worker {worker} timed out") - for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []): + log(f"Worker {worker} heartbeat timed out after {age:.1f}s") + for task_id in job_data.assigned_to_workers.get(worker, []): # If batched_static, task_id is a tile_idx; consider it complete only if # all corresponding global_idx entries are present in completed_tasks. - batched_static = bool(job_data.get('batched_static', False)) + batched_static = bool(job_data.batched_static) if batched_static: - num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1) - batch_size = job_data.get(JOB_BATCH_SIZE, 1) + num_tiles_per_image = int(job_data.num_tiles_per_image or 1) + batch_size = int(job_data.batch_size or 1) # Check all global indices for this tile across the batch all_done = True for b in range(batch_size): @@ -351,17 +307,15 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): all_done = False break if not all_done: - await job_data[JOB_PENDING_TASKS].put(task_id) + await job_data.pending_tasks.put(task_id) requeued_count += 1 else: - # Legacy/global-idx mode: task_id is a global index key if task_id not in completed_tasks: - await job_data[JOB_PENDING_TASKS].put(task_id) + await job_data.pending_tasks.put(task_id) requeued_count += 1 - if JOB_WORKER_STATUS in job_data: - del job_data[JOB_WORKER_STATUS][worker] - if JOB_ASSIGNED_TO_WORKERS in job_data: - job_data[JOB_ASSIGNED_TO_WORKERS][worker] = [] + job_data.worker_status.pop(worker, None) + if worker in job_data.assigned_to_workers: + job_data.assigned_to_workers[worker] = [] return requeued_count @@ -370,8 +324,8 @@ async def _get_completed_count(multi_job_id): 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 job_data and JOB_COMPLETED_TASKS in job_data: - return len(job_data[JOB_COMPLETED_TASKS]) + if isinstance(job_data, BaseJobState): + return len(job_data.completed_tasks) return 0 async def _mark_task_completed(multi_job_id, task_id, result): @@ -379,8 +333,8 @@ async def _mark_task_completed(multi_job_id, task_id, result): 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 job_data and JOB_COMPLETED_TASKS in job_data: - job_data[JOB_COMPLETED_TASKS][task_id] = result + if isinstance(job_data, BaseJobState): + job_data.completed_tasks[task_id] = result async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id): """Send heartbeat to master.""" @@ -417,8 +371,8 @@ async def heartbeat_endpoint(request): async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_WORKER_STATUS in job_data: - job_data[JOB_WORKER_STATUS][worker_id] = time.time() + if isinstance(job_data, BaseJobState): + job_data.worker_status[worker_id] = time.time() debug_log(f"Heartbeat from worker {worker_id}") return web.json_response({"status": "success"}) else: @@ -454,16 +408,15 @@ async def submit_tiles_endpoint(request): async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + if not isinstance(job_data, TileJobState): return await handle_api_error(request, "Job not configured for tile submissions", 400) - if JOB_QUEUE in job_data: - await job_data[JOB_QUEUE].put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) + await job_data.queue.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) try: tiles = _parse_tiles_from_form(data) @@ -474,10 +427,10 @@ async def submit_tiles_endpoint(request): async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + if not isinstance(job_data, TileJobState): return await handle_api_error(request, "Job not configured for tile submissions", 400) - - q = job_data[JOB_QUEUE] + + q = job_data.queue if batch_size > 0 or len(tiles) > 0: await q.put({ 'worker_id': worker_id, @@ -527,32 +480,29 @@ async def submit_image_endpoint(request): async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + if not isinstance(job_data, ImageJobState): return await handle_api_error(request, "Job not configured for image submissions", 400) - if JOB_QUEUE in job_data: - await job_data[JOB_QUEUE].put({ - 'worker_id': worker_id, - 'image_idx': image_idx, - 'image': img, - 'is_last': is_last - }) - return web.json_response({"status": "success"}) + await job_data.queue.put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) # Handle completion signal (no image data) elif is_last: async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + if not isinstance(job_data, ImageJobState): return await handle_api_error(request, "Job not configured for image submissions", 400) - if JOB_QUEUE in job_data: - await job_data[JOB_QUEUE].put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] # For compatibility - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) + await job_data.queue.put({ + 'worker_id': worker_id, + 'is_last': True, + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) else: return await handle_api_error(request, "Missing image data or invalid request", 400) @@ -560,10 +510,6 @@ async def submit_image_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) -# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles. - - - # Helper functions for shallow copying conditioning without duplicating models def clone_control_chain(control, clone_hint=True): """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" @@ -606,12 +552,13 @@ def ensure_tile_jobs_initialized(): prompt_server.distributed_pending_tile_jobs = {} prompt_server.distributed_tile_jobs_lock = asyncio.Lock() else: - # Clean up any legacy queue structures that don't have the 'mode' field - # (Should be rare after fixes, but keep for safety) - to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() - if not isinstance(job_data, dict) or 'mode' not in job_data] - for job_id in to_remove: - debug_log(f"Removing legacy queue structure for job {job_id}") + invalid_job_ids = [ + job_id + for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() + if not isinstance(job_data, BaseJobState) + ] + for job_id in invalid_job_ids: + debug_log(f"Removing invalid job state for {job_id}") del prompt_server.distributed_pending_tile_jobs[job_id] return prompt_server @@ -631,27 +578,23 @@ async def request_image_endpoint(request): async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: + if not isinstance(job_data, BaseJobState): return await handle_api_error(request, "Invalid job data structure", 500) - - mode = job_data['mode'] - - # Handle both dynamic and static modes - if mode == 'dynamic' and 'pending_images' in job_data: - pending_queue = job_data['pending_images'] - elif mode == 'static' and JOB_PENDING_TASKS in job_data: - pending_queue = job_data[JOB_PENDING_TASKS] + + mode = job_data.mode + if isinstance(job_data, ImageJobState): + pending_queue = job_data.pending_images + elif isinstance(job_data, TileJobState): + pending_queue = job_data.pending_tasks else: return await handle_api_error(request, "Invalid job configuration", 400) try: task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) # Track assigned task - if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: - job_data['assigned_to_workers'][worker_id].append(task_idx) + job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx) # Update worker heartbeat - if 'worker_status' in job_data: - job_data['worker_status'][worker_id] = time.time() + job_data.worker_status[worker_id] = time.time() # Get estimated remaining count remaining = pending_queue.qsize() # Approximate @@ -661,7 +604,7 @@ async def request_image_endpoint(request): return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining}) else: # static debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}") - return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.get('batched_static', False)}) + return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.batched_static}) except asyncio.TimeoutError: if mode == 'dynamic': return web.json_response({"image_idx": None}) @@ -681,7 +624,5 @@ async def job_status_endpoint(request): 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) - ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) + ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None) return web.json_response({"ready": ready}) - -