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

Refactor USDU job state to typed dataclasses

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-02-23 23:46:23 +00:00
parent 0f384f51d6
commit 7fc4f05abd
5 changed files with 320 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
upscale/job_models.py
+49
View File
@@ -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
upscale/job_state.py
+37 -38
View File
@@ -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)
upscale/modes/static.py
+37 -10
View File
@@ -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
upscale/result_collector.py
+80 -55
View File
@@ -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."""
utils/usdu_managment.py
+117 -176
View File
@@ -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})