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

Unify worker timeout via UI across collector/upscaler (90s); add interruptible sliced waits + cleanup to collector

Browse Source
This commit is contained in:
robertvoy
2025-09-13 10:26:03 +10:00
parent b69fc40fbd
commit 4b8836cd03
6 changed files with 166 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
distributed.py
+119 -105
View File
@@ -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())
distributed_upscale.py
+9 -5
View File
@@ -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
utils/config.py
+22 -1
View File
@@ -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")
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))
utils/constants.py
+1 -1
View File
@@ -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
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '90')) # Worker heartbeat timeout in seconds (default 90s)
web/main.js
+13 -2
View File
@@ -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) {
web/sidebarRenderer.js
+2 -2
View File
@@ -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;