diff --git a/distributed.py b/distributed.py index 18d1353..c22014c 100644 --- a/distributed.py +++ b/distributed.py @@ -9,6 +9,7 @@ import aiohttp from aiohttp import web import io import server +import comfy.model_management import subprocess import platform import time @@ -22,7 +23,7 @@ from comfy.utils import ProgressBar # Import shared utilities from .utils.logging import debug_log, log -from .utils.config import CONFIG_FILE, get_default_config, load_config, save_config, ensure_config_exists +from .utils.config import CONFIG_FILE, get_default_config, load_config, save_config, ensure_config_exists, get_worker_timeout_seconds from .utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous from .utils.process import is_process_alive, terminate_process, get_python_executable from .utils.network import handle_api_error, get_server_port, get_server_loop, get_client_session, cleanup_client_session @@ -1912,8 +1913,10 @@ class DistributedCollectorNode: collected_count = 0 workers_done = set() - # Use a reasonable timeout for the first image - timeout = WORKER_JOB_TIMEOUT + # Use unified worker timeout from config/UI with simple sliced waits + base_timeout = float(get_worker_timeout_seconds()) + slice_timeout = min(0.5, base_timeout) # small per-wait slice to recheck interrupt + last_activity = time.time() # Get queue size before starting @@ -1924,116 +1927,127 @@ class DistributedCollectorNode: # NEW: Initialize progress bar for workers (total = num_workers) p = ProgressBar(num_workers) - while len(workers_done) < num_workers: - try: - # Get the queue again each time to ensure we have the right reference - async with prompt_server.distributed_jobs_lock: - q = prompt_server.distributed_pending_jobs[multi_job_id] - current_size = q.qsize() - - - result = await asyncio.wait_for(q.get(), timeout=timeout) - worker_id = result['worker_id'] - is_last = result.get('is_last', False) - - # Check if batch mode - tensors = result.get('tensors', []) - indices = result.get('indices', []) # Get the indices - if tensors: - # Batch mode - debug_log(f"Master - Got batch from worker {worker_id}, size={len(tensors)}, is_last={is_last}") + try: + while len(workers_done) < num_workers: + # Check for user interruption to abort collection promptly + comfy.model_management.throw_exception_if_processing_interrupted() + try: + # Get the queue again each time to ensure we have the right reference + async with prompt_server.distributed_jobs_lock: + q = prompt_server.distributed_pending_jobs[multi_job_id] + current_size = q.qsize() - if worker_id not in worker_images: - worker_images[worker_id] = {} + result = await asyncio.wait_for(q.get(), timeout=slice_timeout) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) - # Use actual indices if available, otherwise fall back to sequential - if indices: - for i, tensor in enumerate(tensors): - actual_idx = indices[i] - worker_images[worker_id][actual_idx] = tensor + # Check if batch mode + tensors = result.get('tensors', []) + indices = result.get('indices', []) # Get the indices + if tensors: + # Batch mode + debug_log(f"Master - Got batch from worker {worker_id}, size={len(tensors)}, is_last={is_last}") + + if worker_id not in worker_images: + worker_images[worker_id] = {} + + # Use actual indices if available, otherwise fall back to sequential + if indices: + for i, tensor in enumerate(tensors): + actual_idx = indices[i] + worker_images[worker_id][actual_idx] = tensor + else: + # Fallback for backward compatibility + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + + collected_count += len(tensors) else: - # Fallback for backward compatibility - for idx, tensor in enumerate(tensors): - worker_images[worker_id][idx] = tensor + # Single image mode (backward compat) + image_index = result['image_index'] + tensor = result['tensor'] - - collected_count += len(tensors) - else: - # Single image mode (backward compat) - image_index = result['image_index'] - tensor = result['tensor'] - - debug_log(f"Master - Got single result from worker {worker_id}, image {image_index}, is_last={is_last}") - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - - collected_count += 1 - - # Once we start receiving images, use shorter timeout - timeout = WORKER_JOB_TIMEOUT - - if is_last: - workers_done.add(worker_id) - p.update(1) # +1 per completed worker - else: - pass # Continue waiting for more results - - except asyncio.TimeoutError: - missing_workers = set(str(w) for w in enabled_workers) - workers_done - log(f"Master - Timeout. Still waiting for workers: {list(missing_workers)}") - - # Check queue size again with lock - async with prompt_server.distributed_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_jobs: - final_q = prompt_server.distributed_pending_jobs[multi_job_id] - final_size = final_q.qsize() + debug_log(f"Master - Got single result from worker {worker_id}, image {image_index}, is_last={is_last}") - # Try to drain any remaining items - remaining_items = [] - while not final_q.empty(): - try: - item = final_q.get_nowait() - remaining_items.append(item) - except asyncio.QueueEmpty: - break + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor - if remaining_items: - # Process them - for item in remaining_items: - worker_id = item['worker_id'] - is_last = item.get('is_last', False) - - # Check if batch mode - tensors = item.get('tensors', []) - if tensors: - # Batch mode - if worker_id not in worker_images: - worker_images[worker_id] = {} + collected_count += 1 + + # Record activity and refresh timeout baseline + last_activity = time.time() + base_timeout = float(get_worker_timeout_seconds()) + + if is_last: + workers_done.add(worker_id) + p.update(1) # +1 per completed worker + + except asyncio.TimeoutError: + # If we still have time, continue polling; otherwise handle timeout + if (time.time() - last_activity) < base_timeout: + comfy.model_management.throw_exception_if_processing_interrupted() + continue + # Re-check for user interruption after timeout expiry + comfy.model_management.throw_exception_if_processing_interrupted() + missing_workers = set(str(w) for w in enabled_workers) - workers_done + log(f"Master - Timeout. Still waiting for workers: {list(missing_workers)}") + + # Check queue size again with lock + async with prompt_server.distributed_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_jobs: + final_q = prompt_server.distributed_pending_jobs[multi_job_id] + final_size = final_q.qsize() + + # Try to drain any remaining items + remaining_items = [] + while not final_q.empty(): + try: + item = final_q.get_nowait() + remaining_items.append(item) + except asyncio.QueueEmpty: + break + + if remaining_items: + # Process them + for item in remaining_items: + worker_id = item['worker_id'] + is_last = item.get('is_last', False) - for idx, tensor in enumerate(tensors): - worker_images[worker_id][idx] = tensor + # Check if batch mode + tensors = item.get('tensors', []) + if tensors: + # Batch mode + if worker_id not in worker_images: + worker_images[worker_id] = {} + + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + + collected_count += len(tensors) + else: + # Single image mode + image_index = item['image_index'] + tensor = item['tensor'] + + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor + + collected_count += 1 - collected_count += len(tensors) - else: - # Single image mode - image_index = item['image_index'] - tensor = item['tensor'] - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - collected_count += 1 - - if is_last: - workers_done.add(worker_id) - p.update(1) # +1 here too - else: - log(f"Master - Queue {multi_job_id} no longer exists!") - break + if is_last: + workers_done.add(worker_id) + p.update(1) # +1 here too + else: + log(f"Master - Queue {multi_job_id} no longer exists!") + break + except comfy.model_management.InterruptProcessingException: + # Cleanup queue on interruption and re-raise to abort prompt cleanly + async with prompt_server.distributed_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_jobs: + del prompt_server.distributed_pending_jobs[multi_job_id] + raise total_collected = sum(len(imgs) for imgs in worker_images.values()) diff --git a/distributed_upscale.py b/distributed_upscale.py index e3c4c27..284742a 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -19,6 +19,7 @@ from .utils.logging import debug_log, log from .utils.image import tensor_to_pil, pil_to_tensor from .utils.network import get_client_session from .utils.async_helpers import run_async_in_server_loop +from .utils.config import get_worker_timeout_seconds from .utils.constants import ( TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH @@ -46,7 +47,7 @@ from .utils.usdu_managment import init_dynamic_job, init_static_job_batched # Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants # They can be overridden via environment variables: # - COMFYUI_MAX_BATCH (default: 20) -# - COMFYUI_HEARTBEAT_TIMEOUT (default: 120) +# - COMFYUI_HEARTBEAT_TIMEOUT (default: 90) # Sync wrapper decorator for async methods def sync_wrapper(async_func): @@ -329,7 +330,8 @@ class UltimateSDUpscaleDistributed: collected_results = {} workers_done = set() - timeout = 120.0 if mode == 'dynamic' else 60.0 # Longer timeout for full images + # Unify collector/upscaler wait behavior with the UI worker timeout + timeout = float(get_worker_timeout_seconds()) last_heartbeat_check = time.time() collected_count = 0 @@ -344,8 +346,8 @@ class UltimateSDUpscaleDistributed: break try: - # Shorter timeout for regular checks in dynamic mode - wait_timeout = 10.0 if mode == 'dynamic' else timeout + # Shorter poll for dynamic mode, but never exceed the configured timeout + wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout) result = await asyncio.wait_for(q.get(), timeout=wait_timeout) worker_id = result['worker_id'] is_last = result.get('is_last', False) @@ -1523,9 +1525,11 @@ class UltimateSDUpscaleDistributed: if remaining_to_collect > 0: debug_log(f"Waiting for {remaining_to_collect} more images from workers") + # Use the unified worker timeout for the collection phase + collection_timeout = float(get_worker_timeout_seconds()) collected_images = run_async_in_server_loop( self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count), - timeout=TILE_COLLECTION_TIMEOUT * 2 + timeout=collection_timeout ) # Merge collected with already completed diff --git a/utils/config.py b/utils/config.py index f1e8040..2569663 100644 --- a/utils/config.py +++ b/utils/config.py @@ -5,6 +5,9 @@ import os import json from .logging import log +# Import defaults for timeout fallbacks +from .constants import HEARTBEAT_TIMEOUT + CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json") def get_default_config(): @@ -47,4 +50,22 @@ def ensure_config_exists(): from .logging import debug_log debug_log("Created default config file") else: - log("Could not create default config file") \ No newline at end of file + log("Could not create default config file") + +def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int: + """Return the unified worker timeout (seconds). + + Priority: + 1) UI-configured setting `settings.worker_timeout_seconds` + 2) Fallback to provided `default` (defaults to HEARTBEAT_TIMEOUT which itself + can be overridden via the COMFYUI_HEARTBEAT_TIMEOUT env var) + + This value should be used anywhere we consider a worker "timed out" from the + master's perspective (e.g., collector waits, upscaler result collection). + """ + try: + cfg = load_config() + val = int(cfg.get('settings', {}).get('worker_timeout_seconds', default)) + return max(1, val) + except Exception: + return max(1, int(default)) diff --git a/utils/constants.py b/utils/constants.py index a4ca5bd..ad3de25 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -38,4 +38,4 @@ MEMORY_CLEAR_DELAY = 0.5 MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs) # Heartbeat monitoring -HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds \ No newline at end of file +HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '90')) # Worker heartbeat timeout in seconds (default 90s) diff --git a/web/main.js b/web/main.js index e1026d2..1f0b007 100644 --- a/web/main.js +++ b/web/main.js @@ -154,11 +154,22 @@ class DistributedExtension { try { await this.api.updateSetting(key, value); - + + const prettyKey = key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase()); + let detail; + if (key === 'worker_timeout_seconds') { + const secs = parseInt(value, 10); + detail = `Worker Timeout set to ${Number.isFinite(secs) ? secs : value}s`; + } else if (typeof value === 'boolean') { + detail = `${prettyKey} ${value ? 'enabled' : 'disabled'}`; + } else { + detail = `${prettyKey} set to ${value}`; + } + app.extensionManager.toast.add({ severity: "success", summary: "Setting Updated", - detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`, + detail, life: 2000 }); } catch (error) { diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js index a8b42a3..e184d1c 100644 --- a/web/sidebarRenderer.js +++ b/web/sidebarRenderer.js @@ -280,7 +280,7 @@ export async function renderSidebarContent(extension, el) { timeoutLabel.htmlFor = "setting-worker-timeout"; timeoutLabel.textContent = "Worker Timeout"; timeoutLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: default;"; - timeoutLabel.title = "Seconds without a heartbeat before a worker is considered timed out and its tasks are requeued. Default 60. For WAN, consider 300–600."; + timeoutLabel.title = "Seconds without a heartbeat before a worker is considered timed out and its tasks are requeued. Default 90. For WAN, consider 300–600."; const timeoutInput = document.createElement("input"); timeoutInput.type = "number"; @@ -288,7 +288,7 @@ export async function renderSidebarContent(extension, el) { timeoutInput.min = "10"; timeoutInput.step = "1"; timeoutInput.style.cssText = "width: 80px; padding: 2px 6px; background: #222; color: #ddd; border: 1px solid #333; border-radius: 3px;"; - timeoutInput.value = (extension.config?.settings?.worker_timeout_seconds ?? 60); + timeoutInput.value = (extension.config?.settings?.worker_timeout_seconds ?? 90); timeoutInput.onchange = (e) => { const v = parseInt(e.target.value, 10); if (!Number.isFinite(v) || v <= 0) return;