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."""